fix(gemini): add oauth loopback timeout and normalize code assist model handling

litellm/llms/gemini/authenticator.py

@@ -53,6 +53,15 @@ class GeminiAuthenticator:
             )
         return client_id, client_secret
 
+    @staticmethod
+    def _get_access_token_or_raise(creds: Dict[str, Any], source: str) -> str:
+        token = creds.get("access_token")
+        if not token or not isinstance(token, str):
+            raise Exception(
+                f"OAuth login failed: missing access_token in {source} response."
+            )
+        return token
+
     def get_token(self) -> str:
         """
         Get the OAuth token, refreshing if necessary.
@@ -88,7 +97,7 @@ class GeminiAuthenticator:
         verbose_logger.info("Starting Gemini OAuth login flow...")
         creds = self._login()
         self._write_oauth_creds(creds)
-        return creds.get("access_token")
+        return self._get_access_token_or_raise(creds, "authorization_code")
 
     def _refresh_token(self, refresh_token: str) -> str:
         """Refresh the access token using the refresh token."""
@@ -118,7 +127,7 @@ class GeminiAuthenticator:
 
         self._write_oauth_creds(creds)
 
-        return creds.get("access_token")
+        return self._get_access_token_or_raise(creds, "refresh_token")
 
     def _ensure_token_dir(self) -> None:
         """Ensure the token directory exists."""
@@ -234,9 +243,23 @@ class GeminiAuthenticator:
         webbrowser.open(auth_url)
 
         # Wait for callback; browsers may hit non-callback paths first (e.g. /favicon.ico).
-        while auth_code is None and error is None:
-            server.handle_request()
-        server.server_close()
+        # Use a bounded loop so headless/CI environments fail fast instead of hanging forever.
+        loopback_timeout_s = float(
+            os.getenv("GEMINI_OAUTH_LOOPBACK_TIMEOUT_SECONDS", "120")
+        )
+        loopback_timeout_s = max(loopback_timeout_s, 1.0)
+        deadline = time.monotonic() + loopback_timeout_s
+        server.timeout = 1.0
+        try:
+            while auth_code is None and error is None:
+                server.handle_request()
+                if time.monotonic() >= deadline:
+                    raise Exception(
+                        "OAuth login timed out. Re-run `litellm-proxy gemini login` "
+                        f"and complete the browser flow within {int(loopback_timeout_s)} seconds."
+                    )
+        finally:
+            server.server_close()
 
         if error:
             raise Exception(f"OAuth login failed: {error}")

litellm/llms/google_code_assist/transformation.py

@@ -97,6 +97,7 @@ class GoogleCodeAssistConfig(VertexGeminiConfig):
 
         # Create a copy to avoid mutating the original list
         messages_copy = copy.deepcopy(messages)
+        model_name = model.split("/")[-1]
 
         # Separate system instruction
         system_instruction, filtered_messages = _transform_system_message(
@@ -105,14 +106,14 @@ class GoogleCodeAssistConfig(VertexGeminiConfig):
 
         # Convert the rest of messages
         contents = _gemini_convert_messages_with_history(
-            messages=filtered_messages, model=model
+            messages=filtered_messages, model=model_name
         )
 
         # 2. Build vertex-style nested request
         generation_config = {}
         # Handle parameter mapping
         base_params = self.map_openai_params(
-            {}, optional_params.copy(), model, messages
+            {}, optional_params.copy(), model_name, messages
         )
 
         for key in ["temperature", "topP", "topK", "maxOutputTokens", "stopSequences"]:
@@ -126,12 +127,6 @@ class GoogleCodeAssistConfig(VertexGeminiConfig):
             generation_config["thinkingConfig"] = {
                 "includeThoughts": base_params.pop("include_thoughts")
             }
-        elif "thinkingConfig" in optional_params:
-            generation_config["thinkingConfig"] = optional_params["thinkingConfig"]
-        elif "include_thoughts" in optional_params:
-            generation_config["thinkingConfig"] = {
-                "includeThoughts": optional_params["include_thoughts"]
-            }
 
         vertex_request = {
             "contents": contents,
@@ -149,7 +144,6 @@ class GoogleCodeAssistConfig(VertexGeminiConfig):
 
         # 3. Wrap in Code Assist envelope (matches verified gemini-cli structure)
         user_prompt_id = f"litellm-{uuid.uuid4()}"
-        model_name = model.split("/")[-1]
 
         ca_request = {
             "model": model_name,

litellm/main.py

@@ -3703,7 +3703,7 @@ def completion(  # type: ignore # noqa: PLR0915
                     logger_fn=logger_fn,
                 )
             else:
-                model_response = google_code_assist_chat.completion(
+                response = google_code_assist_chat.completion(
                     model=model,
                     messages=messages,
                     model_response=model_response,
@@ -3713,7 +3713,6 @@ def completion(  # type: ignore # noqa: PLR0915
                     logging_obj=logging,
                     logger_fn=logger_fn,
                 )
-                response = model_response
 
         elif custom_llm_provider == "predibase":
             tenant_id = (

tests/test_litellm/llms/gemini/test_gemini_authenticator.py

@@ -0,0 +1,68 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from litellm.llms.gemini.authenticator import GeminiAuthenticator
+
+
+def test_get_token_raises_if_login_response_missing_access_token(tmp_path):
+    auth = GeminiAuthenticator()
+    auth.oauth_creds_file = str(tmp_path / "oauth_creds.json")
+
+    with (
+        patch.object(auth, "_login", return_value={"refresh_token": "abc"}),
+        patch.object(auth, "_write_oauth_creds"),
+    ):
+        with pytest.raises(Exception, match="missing access_token"):
+            auth.get_token()
+
+
+def test_refresh_token_raises_if_response_missing_access_token(tmp_path):
+    auth = GeminiAuthenticator()
+    auth.oauth_creds_file = str(tmp_path / "oauth_creds.json")
+
+    mock_resp = MagicMock()
+    mock_resp.raise_for_status.return_value = None
+    mock_resp.json.return_value = {"expires_in": 3600}
+
+    with (
+        patch.object(
+            GeminiAuthenticator,
+            "_get_oauth_client_credentials",
+            return_value=("client-id", "client-secret"),
+        ),
+        patch("litellm.llms.gemini.authenticator.httpx.post", return_value=mock_resp),
+        patch.object(auth, "_write_oauth_creds"),
+    ):
+        with pytest.raises(Exception, match="missing access_token"):
+            auth._refresh_token("refresh-token")
+
+
+def test_login_times_out_without_oauth_callback(monkeypatch):
+    auth = GeminiAuthenticator()
+    monkeypatch.setenv("GEMINI_OAUTH_LOOPBACK_TIMEOUT_SECONDS", "1")
+
+    class _NoRequestServer:
+        def __init__(self, *args, **kwargs):
+            self.server_port = 12345
+            self.timeout = None
+
+        def handle_request(self):
+            return None
+
+        def server_close(self):
+            return None
+
+    with (
+        patch.object(
+            GeminiAuthenticator,
+            "_get_oauth_client_credentials",
+            return_value=("client-id", "client-secret"),
+        ),
+        patch(
+            "litellm.llms.gemini.authenticator.http.server.HTTPServer", _NoRequestServer
+        ),
+        patch("litellm.llms.gemini.authenticator.webbrowser.open", return_value=True),
+    ):
+        with pytest.raises(Exception, match="OAuth login timed out"):
+            auth._login()

tests/test_litellm/llms/google_code_assist/test_google_code_assist_transformation.py

@@ -0,0 +1,42 @@
+from unittest.mock import patch
+
+from litellm.llms.google_code_assist.transformation import GoogleCodeAssistConfig
+
+
+def test_transform_request_uses_base_model_name_for_internal_gemini_helpers():
+    class _CaptureConfig(GoogleCodeAssistConfig):
+        def __init__(self):
+            super().__init__()
+            self.model_seen_by_map = None
+
+        def map_openai_params(
+            self,
+            non_default_params: dict,
+            optional_params: dict,
+            model: str,
+            messages: list,
+        ) -> dict:
+            self.model_seen_by_map = model
+            return {"thinkingConfig": {"includeThoughts": True}}
+
+    config = _CaptureConfig()
+    with (
+        patch(
+            "litellm.llms.vertex_ai.gemini.transformation._transform_system_message",
+            return_value=(None, [{"role": "user", "content": "hi"}]),
+        ),
+        patch(
+            "litellm.llms.vertex_ai.gemini.transformation._gemini_convert_messages_with_history",
+            return_value=[{"role": "user", "parts": [{"text": "hi"}]}],
+        ) as mock_convert,
+    ):
+        result = config.transform_request(
+            model="google_code_assist/gemini-2.5-pro",
+            messages=[{"role": "user", "content": "hi"}],
+            optional_params={"thinkingConfig": {"includeThoughts": False}},
+            litellm_params={},
+        )
+
+    assert config.model_seen_by_map == "gemini-2.5-pro"
+    assert mock_convert.call_args.kwargs["model"] == "gemini-2.5-pro"
+    assert result["model"] == "gemini-2.5-pro"