From 6714590c9c8e41b67e7e923497d24ea15c13b997 Mon Sep 17 00:00:00 2001 From: balazss Date: Tue, 17 Mar 2026 22:37:55 -0700 Subject: [PATCH] fix(code-assist): offload async oauth lookup and validate token --- litellm/llms/gemini/fallback_handler.py | 6 +- litellm/llms/google_code_assist/chat.py | 13 +++- litellm/main.py | 7 +- .../test_google_code_assist.py | 76 +++++++++++++++++++ 4 files changed, 96 insertions(+), 6 deletions(-) diff --git a/litellm/llms/gemini/fallback_handler.py b/litellm/llms/gemini/fallback_handler.py index a59497657e6..08cccf222c7 100644 --- a/litellm/llms/gemini/fallback_handler.py +++ b/litellm/llms/gemini/fallback_handler.py @@ -4,6 +4,8 @@ from litellm._logging import verbose_logger from litellm.llms.gemini.common_utils import should_fallback_to_google_code_assist from litellm.llms.google_code_assist.chat import GoogleCodeAssistChat +_google_code_assist_chat = GoogleCodeAssistChat() + async def run_gemini_acompletion_with_code_assist_fallback( primary_call: Awaitable[Any], @@ -27,7 +29,7 @@ async def run_gemini_acompletion_with_code_assist_fallback( "Gemini request failed with ACCESS_TOKEN_SCOPE_INSUFFICIENT. " "Falling back to google_code_assist." ) - return await GoogleCodeAssistChat().acompletion(**fallback_kwargs) + return await _google_code_assist_chat.acompletion(**fallback_kwargs) def run_gemini_completion_with_code_assist_fallback( @@ -52,4 +54,4 @@ def run_gemini_completion_with_code_assist_fallback( "Gemini request failed with ACCESS_TOKEN_SCOPE_INSUFFICIENT. " "Falling back to google_code_assist." ) - return GoogleCodeAssistChat().completion(**fallback_kwargs) + return _google_code_assist_chat.completion(**fallback_kwargs) diff --git a/litellm/llms/google_code_assist/chat.py b/litellm/llms/google_code_assist/chat.py index 7bb74cb4a54..beea6ca997e 100644 --- a/litellm/llms/google_code_assist/chat.py +++ b/litellm/llms/google_code_assist/chat.py @@ -1,3 +1,4 @@ +import asyncio import httpx from typing import Any, Optional @@ -43,6 +44,11 @@ class GoogleCodeAssistChat: token = gemini_auth_data.get("token") initial_project_id = gemini_auth_data.get("project_id") + if not token: + raise GoogleCodeAssistError( + status_code=401, + message="Missing Gemini OAuth token value. Re-run 'litellm-proxy gemini login'.", + ) client = _get_httpx_client() @@ -102,7 +108,7 @@ class GoogleCodeAssistChat: try: from litellm.llms.gemini.common_utils import get_gemini_oauth_token - gemini_auth_data = get_gemini_oauth_token() + gemini_auth_data = await asyncio.to_thread(get_gemini_oauth_token) if not gemini_auth_data: raise GoogleCodeAssistError( status_code=401, @@ -111,6 +117,11 @@ class GoogleCodeAssistChat: token = gemini_auth_data.get("token") initial_project_id = gemini_auth_data.get("project_id") + if not token: + raise GoogleCodeAssistError( + status_code=401, + message="Missing Gemini OAuth token value. Re-run 'litellm-proxy gemini login'.", + ) async_handler = AsyncHTTPHandler() try: diff --git a/litellm/main.py b/litellm/main.py index 40807265949..26eaced03f6 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -161,6 +161,8 @@ from litellm.utils import ( validate_openai_optional_params, ) +_google_code_assist_chat = GoogleCodeAssistChat() + from ._logging import verbose_logger from .caching.caching import disable_cache, enable_cache, update_cache from .litellm_core_utils.core_helpers import safe_deep_copy @@ -3690,9 +3692,8 @@ def completion( # type: ignore # noqa: PLR0915 response = model_response elif custom_llm_provider == "google_code_assist": - google_code_assist_chat = GoogleCodeAssistChat() if acompletion is True: - response = google_code_assist_chat.acompletion( + response = _google_code_assist_chat.acompletion( model=model, messages=messages, model_response=model_response, @@ -3703,7 +3704,7 @@ def completion( # type: ignore # noqa: PLR0915 logger_fn=logger_fn, ) else: - response = google_code_assist_chat.completion( + response = _google_code_assist_chat.completion( model=model, messages=messages, model_response=model_response, diff --git a/tests/test_litellm/llms/google_code_assist/test_google_code_assist.py b/tests/test_litellm/llms/google_code_assist/test_google_code_assist.py index 5340443775e..c760e8d6ead 100644 --- a/tests/test_litellm/llms/google_code_assist/test_google_code_assist.py +++ b/tests/test_litellm/llms/google_code_assist/test_google_code_assist.py @@ -70,6 +70,22 @@ class TestGoogleCodeAssist: assert response.choices[0].message.content == "Hello world" assert response.usage.total_tokens == 7 + @patch("litellm.llms.gemini.common_utils.get_gemini_oauth_token") + def test_completion_raises_when_token_missing(self, mock_get_token): + mock_get_token.return_value = {"token": None} + + handler = GoogleCodeAssistChat() + with pytest.raises(Exception, match="Missing Gemini OAuth token value"): + handler.completion( + model="google_code_assist/gemini-1.5-flash", + messages=[{"role": "user", "content": "hi"}], + model_response=ModelResponse(), + print_verbose=False, + logging_obj=MagicMock(), + optional_params={}, + litellm_params={}, + ) + @pytest.mark.asyncio @patch("litellm.llms.google_code_assist.chat.AsyncHTTPHandler.post") @patch( @@ -132,3 +148,63 @@ class TestGoogleCodeAssist: assert response.choices[0].message.content == "Async success" mock_async_close.assert_awaited_once() + + @pytest.mark.asyncio + @patch( + "litellm.llms.google_code_assist.chat.asyncio.to_thread", new_callable=AsyncMock + ) + async def test_acompletion_uses_thread_for_oauth_lookup(self, mock_to_thread): + mock_to_thread.return_value = {"token": "test-token"} + + with ( + patch( + "litellm.llms.google_code_assist.chat.GoogleCodeAssistChat._ahandle_handshake", + new_callable=AsyncMock, + return_value="final-project", + ), + patch( + "litellm.llms.google_code_assist.chat.AsyncHTTPHandler.post", + new_callable=AsyncMock, + ) as mock_async_post, + patch( + "litellm.llms.google_code_assist.chat.AsyncHTTPHandler.close", + new_callable=AsyncMock, + ), + ): + completion_data = { + "response": { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "Async success"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 2, + "totalTokenCount": 7, + }, + } + } + mock_completion_resp = httpx.Response( + status_code=200, + content=json.dumps(completion_data).encode(), + request=httpx.Request("POST", "https://completion"), + ) + mock_async_post.return_value = mock_completion_resp + + handler = GoogleCodeAssistChat() + await handler.acompletion( + model="google_code_assist/gemini-1.5-flash", + messages=[{"role": "user", "content": "hi"}], + model_response=ModelResponse(), + print_verbose=False, + logging_obj=MagicMock(), + optional_params={}, + litellm_params={}, + ) + + mock_to_thread.assert_awaited_once()