feat(gemini): implement auto-retry on 401 Unauthorized errors

- Add proactive token refresh and retry logic for 401 responses
- Support both streaming and non-streaming request retries
- Prevent infinite retry loops with _auth_retry_attempted flag
- Add comprehensive unit and behavior tests for auth retry scenarios

Author: Mateusz

src/connectors/gemini_base/connector.py

@@ -1398,6 +1398,7 @@ async def _chat_completions_code_assist(
         processed_messages: list[Any],
         effective_model: str,
         _in_graceful_degradation: bool = False,
+        _auth_retry_attempted: bool = False,
         **kwargs: Any,
     ) -> ResponseEnvelope | StreamingResponseEnvelope:
         """Handle chat completions using the Code Assist API.
@@ -1486,6 +1487,44 @@ async def _chat_completions_code_assist(
             )
 
         except AuthenticationError as e:
+            # Handle 401 authentication errors with token refresh and retry
+            if not _auth_retry_attempted:
+                logger.info(
+                    "Received 401 Unauthorized in non-streaming request, attempting token refresh and retry..."
+                )
+                try:
+                    # Use 30s timeout for refresh, leaving room for retry request
+                    AUTH_RETRY_TIMEOUT = 30.0
+                    refreshed = await asyncio.wait_for(
+                        self._refresh_token_if_needed(),
+                        timeout=AUTH_RETRY_TIMEOUT,
+                    )
+                    if refreshed:
+                        logger.info(
+                            "Token refresh successful, retrying non-streaming request..."
+                        )
+                        return await self._chat_completions_code_assist(
+                            request_data=request_data,
+                            processed_messages=processed_messages,
+                            effective_model=effective_model,
+                            _in_graceful_degradation=_in_graceful_degradation,
+                            _auth_retry_attempted=True,  # Prevent infinite retry loops
+                            **kwargs,
+                        )
+                    else:
+                        logger.warning(
+                            "Token refresh failed; will raise 401 error to caller"
+                        )
+                except asyncio.TimeoutError:
+                    logger.warning(
+                        f"Token refresh timed out after {AUTH_RETRY_TIMEOUT}s; raising 401 to caller"
+                    )
+                except Exception as refresh_error:
+                    logger.error(
+                        f"Error during token refresh attempt: {refresh_error}",
+                        exc_info=True,
+                    )
+            # If we reach here, refresh failed or already retried - raise original error
             logger.error(f"Authentication error during API call: {e}", exc_info=True)
             raise
         except BackendError as e:
@@ -1561,6 +1600,7 @@ async def stream_generator(
                 *,
                 _allow_tool_retry: bool = True,
                 without_tools: bool = False,
+                _auth_retry_attempted: bool = False,
             ) -> AsyncGenerator[ProcessedResponse, None]:
                 import json
 
@@ -1727,9 +1767,54 @@ def _build_error_chunk(
                                 code = "quota_exceeded"
                             elif response.status_code == 429:
                                 code = "rate_limit_exceeded"
+                            elif response.status_code == 401:
+                                code = "auth_error"
                         elif isinstance(error_detail, str) and error_detail.strip():
                             error_message = error_detail
 
+                        # Handle 401 authentication errors with token refresh and retry
+                        if response.status_code == 401 and not _auth_retry_attempted:
+                            logger.info(
+                                "Received 401 Unauthorized from backend, attempting token refresh and retry..."
+                            )
+                            with contextlib.suppress(Exception):
+                                response.close()
+
+                            # Trigger proactive token refresh with timeout
+                            # Use 30s timeout for refresh, leaving room for retry request
+                            AUTH_RETRY_TIMEOUT = 30.0
+                            try:
+                                refreshed = await asyncio.wait_for(
+                                    self._refresh_token_if_needed(),
+                                    timeout=AUTH_RETRY_TIMEOUT,
+                                )
+                                if refreshed:
+                                    logger.info(
+                                        "Token refresh successful, retrying streaming request..."
+                                    )
+                                    # Recursively call stream_generator with retry flag set
+                                    async for retry_chunk in stream_generator(
+                                        _allow_tool_retry=_allow_tool_retry,
+                                        without_tools=without_tools,
+                                        _auth_retry_attempted=True,  # Prevent infinite retry loops
+                                    ):
+                                        yield retry_chunk
+                                    return  # Successfully handled via retry
+                                else:
+                                    logger.warning(
+                                        "Token refresh failed; will return 401 error to client"
+                                    )
+                            except asyncio.TimeoutError:
+                                logger.warning(
+                                    f"Token refresh timed out after {AUTH_RETRY_TIMEOUT}s; returning 401 to client"
+                                )
+                            except Exception as refresh_error:
+                                logger.error(
+                                    f"Error during token refresh attempt: {refresh_error}",
+                                    exc_info=True,
+                                )
+                            # If we reach here, refresh failed - continue to raise error below
+
                         # Attach retry-after hint when available
                         retry_delay = None
                         if response.status_code == 429:

src/connectors/gemini_base/graceful_degradation.py

@@ -24,8 +24,8 @@
 logger = logging.getLogger(__name__)
 
 
-def is_rate_limit_like_error(error: BackendError) -> bool:
-    """Determine whether an error should trigger graceful degradation retries.
+def is_rate_limit_like_error(error: BackendError) -> bool:
+    """Determine whether an error should trigger graceful degradation retries.
 
     Args:
         error: The BackendError to check.
@@ -113,8 +113,8 @@ def calculate_retry_delay(
     return max(min_delay, base_delay + jitter)
 
 
-class GracefulDegradationManager:
-    """Manages graceful degradation state for a connector."""
+class GracefulDegradationManager:
+    """Manages graceful degradation state for a connector."""
 
     def __init__(
         self,
@@ -132,9 +132,9 @@ def __init__(
         self.model_retry_states: dict[str, ModelRetryState] = {}
         self.permanently_failed = False
 
-    def is_rate_limit_like_error(self, error: BackendError) -> bool:
-        """Determine whether an error should trigger graceful degradation."""
-        return is_rate_limit_like_error(error)
+    def is_rate_limit_like_error(self, error: BackendError) -> bool:
+        """Determine whether an error should trigger graceful degradation."""
+        return is_rate_limit_like_error(error)
 
     def is_in_cooldown(self, model: str) -> bool:
         """Check if a model is currently in cooldown."""
@@ -163,9 +163,22 @@ def get_or_create_state(self, model: str) -> ModelRetryState:
             self.model_retry_states[model] = ModelRetryState()
         return self.model_retry_states[model]
 
-    def get_models_to_try(self, original_model: str, disable_fallback: bool = False) -> list[str]:
-        """Return only the original model; fallbacks are handled upstream."""
-        return [original_model]
+    def get_models_to_try(
+        self, original_model: str, disable_fallback: bool = False
+    ) -> list[str]:
+        """Return only the original model; fallbacks are handled upstream.
+
+        Args:
+            original_model: The model to use.
+            disable_fallback: Reserved for API compatibility; fallbacks handled upstream.
+
+        Returns:
+            List containing only the original model.
+        """
+        # disable_fallback is intentionally unused - fallbacks are handled upstream
+        # by the connector layer. This parameter is kept for API compatibility.
+        _ = disable_fallback  # Explicitly acknowledge the parameter
+        return [original_model]
 
     def record_attempt(self) -> None:
         """Record an attempt in metrics."""
@@ -201,10 +214,10 @@ def get_metrics(self) -> dict[str, Any]:
         return self.metrics.as_dict()
 
 
-__all__ = [
-    "GracefulDegradationManager",
-    "calculate_retry_delay",
-    "is_model_in_cooldown",
-    "is_rate_limit_like_error",
-    "set_model_cooldown",
-]
+__all__ = [
+    "GracefulDegradationManager",
+    "calculate_retry_delay",
+    "is_model_in_cooldown",
+    "is_rate_limit_like_error",
+    "set_model_cooldown",
+]