def retry_on_session_expiration(func):
    """
    Decorator to retry a method upon session expiration errors with exponential backoff.
    """
    max_attempts = UCAI_DATABRICKS_SESSION_RETRY_MAX_ATTEMPTS.get()

    @functools.wraps(func)
    def wrapper(self, *args, **kwargs):
        for attempt in range(1, max_attempts + 1):
            try:
                result = func(self, *args, **kwargs)
                if hasattr(result, "error") and result.error and _is_session_expired(result.error):
                    raise SessionExpirationException(result.error)
                _logger.info("Successfully re-acquired connection to a serverless instance.")
                return result
            except SessionExpirationException as e:
                if not hasattr(self, "_is_default_client") or not self._is_default_client:
                    refresh_message = (
                        f"Failed to execute {func.__name__} due to session expiration. "
                        "Unable to automatically refresh session when using a custom client. "
                        "Recreate the DatabricksFunctionClient with a new client to recreate "
                        "the custom client session."
                    )
                    raise RuntimeError(refresh_message) from e

                if attempt < max_attempts:
                    delay = min(
                        SESSION_RETRY_BASE_DELAY * (2 ** (attempt - 1)), SESSION_RETRY_MAX_DELAY
                    )
                    _logger.warning(
                        f"Session expired. Retrying attempt {attempt} of {max_attempts}. "
                        f"Refreshing session and retrying after {delay} seconds..."
                    )
                    if hasattr(self, "refresh_client_and_session"):
                        self.refresh_client_and_session()
                    time.sleep(delay)
                    continue
                else:
                    refresh_failure_message = (
                        f"Failed to execute {func.__name__} after {max_attempts} attempts due to session expiration. "
                        "Generate a new session_id by detaching and reattaching your compute."
                    )
                    raise RuntimeError(refresh_failure_message) from e

    return wrapper

    Decorator to retry a method upon session expiration errors with exponential backoff.
    
ai/core/src/unitycatalog/ai/core/utils/retry_utils.py