diff --git a/litellm/llms/gemini/authenticator.py b/litellm/llms/gemini/authenticator.py index bd7d6d0b8dc..ef67116da69 100644 --- a/litellm/llms/gemini/authenticator.py +++ b/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}") diff --git a/litellm/llms/google_code_assist/transformation.py b/litellm/llms/google_code_assist/transformation.py index b58c7ed780d..273b53d7ba4 100644 --- a/litellm/llms/google_code_assist/transformation.py +++ b/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, diff --git a/litellm/main.py b/litellm/main.py index b59ea1bc801..40807265949 100644 --- a/litellm/main.py +++ b/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 = ( diff --git a/tests/test_litellm/llms/gemini/test_gemini_authenticator.py b/tests/test_litellm/llms/gemini/test_gemini_authenticator.py new file mode 100644 index 00000000000..889ff537c60 --- /dev/null +++ b/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() diff --git a/tests/test_litellm/llms/google_code_assist/test_google_code_assist_transformation.py b/tests/test_litellm/llms/google_code_assist/test_google_code_assist_transformation.py new file mode 100644 index 00000000000..ccaebc5b2d6 --- /dev/null +++ b/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"