From 58844d3bda3c74ba35c3a571de0a8d2fcf2b6a79 Mon Sep 17 00:00:00 2001 From: mateo-berri Date: Thu, 20 Aug 2026 02:36:47 -0700 Subject: [PATCH] refactor(realtime): inject the vertex access token resolver Take the resolver and its timeout as parameters of the bounded helper and bind the vertex one once at module level, so the timeout tests drive an injected fake instead of patching a shared singleton. --- litellm/realtime_api/main.py | 15 ++- litellm/types/llms/vertex_ai.py | 13 ++- tests/test_litellm/realtime_api/test_main.py | 99 +++++++++++++------- 3 files changed, 86 insertions(+), 41 deletions(-) diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 8fde7cb75c5..56b3931711e 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -15,7 +15,7 @@ from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.xai.common_utils import XAIModelInfo from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES +from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES, VertexAccessTokenResolver from litellm.types.realtime import ( RealtimeClientSecretRequest, RealtimeExpiresAfter, @@ -43,6 +43,7 @@ openai_realtime: Final = OpenAIRealtime() bedrock_realtime: Final = BedrockRealtime() xai_realtime: Final = XAIRealtime() vertex_llm_base: Final = VertexBase() +vertex_access_token_resolver: Final[VertexAccessTokenResolver] = vertex_llm_base._ensure_access_token_async base_llm_http_handler = BaseLLMHTTPHandler() @@ -290,20 +291,22 @@ async def arealtime_calls( async def _resolve_vertex_access_token_bounded( credentials: VERTEX_CREDENTIALS_TYPES | None, project_id: str | None, + resolver: VertexAccessTokenResolver, + timeout_seconds: float, ) -> tuple[str, str]: try: return await asyncio.wait_for( - vertex_llm_base._ensure_access_token_async( + resolver( credentials=credentials, project_id=project_id, custom_llm_provider="vertex_ai", ), - timeout=REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS, + timeout=timeout_seconds, ) except asyncio.TimeoutError as e: raise ValueError( "Vertex AI realtime: timed out fetching Google OAuth access token after " - f"{REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS}s; check network egress from the proxy " + f"{timeout_seconds}s; check network egress from the proxy " "to the OAuth token endpoint (oauth2.googleapis.com)" ) from e @@ -508,6 +511,8 @@ async def _arealtime( ) = await _resolve_vertex_access_token_bounded( credentials=vertex_credentials, project_id=vertex_project, + resolver=vertex_access_token_resolver, + timeout_seconds=REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS, ) vertex_realtime_config: Final = VertexAIRealtimeConfig( @@ -588,6 +593,8 @@ async def _realtime_health_check( ) = await _resolve_vertex_access_token_bounded( credentials=VertexBase.safe_get_vertex_ai_credentials(vertex_model_params), project_id=VertexBase.safe_get_vertex_ai_project(vertex_model_params), + resolver=vertex_access_token_resolver, + timeout_seconds=REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS, ) vertex_realtime_config: Final = VertexAIRealtimeConfig( access_token=access_token, diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index b750563432e..3b95b786631 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -1,5 +1,5 @@ from enum import Enum -from typing import Any, Final, Literal +from typing import Any, Final, Literal, Protocol from typing_extensions import ( Required, @@ -747,6 +747,17 @@ class VertexVideoGenerationResponse(TypedDict, total=False): VERTEX_CREDENTIALS_TYPES = str | dict[str, str] +class VertexAccessTokenResolver(Protocol): + """Resolves a Google OAuth access token and the project id it belongs to.""" + + async def __call__( + self, + credentials: VERTEX_CREDENTIALS_TYPES | None, + project_id: str | None, + custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"], + ) -> tuple[str, str]: ... + + class VertexPartnerProvider(str, Enum): mistralai = "mistralai" llama = "llama" diff --git a/tests/test_litellm/realtime_api/test_main.py b/tests/test_litellm/realtime_api/test_main.py index 8ed7fb06e84..9f48d4d427b 100644 --- a/tests/test_litellm/realtime_api/test_main.py +++ b/tests/test_litellm/realtime_api/test_main.py @@ -93,12 +93,70 @@ def test_client_secret_session_model_takes_priority_over_top_level(monkeypatch): assert captured["request_data"]["session"]["model"] == "gpt-realtime-session" +async def _hanging_resolver(credentials, project_id, custom_llm_provider) -> tuple[str, str]: + await asyncio.sleep(30) + return "", "" + + +async def _thread_offloaded_hanging_resolver(credentials, project_id, custom_llm_provider) -> tuple[str, str]: + from litellm.litellm_core_utils.asyncify import asyncify + + await asyncify(time.sleep)(30) + return "", "" + + +async def _instant_resolver(credentials, project_id, custom_llm_provider) -> tuple[str, str]: + return "token-abc", "resolved-project" + + @pytest.mark.asyncio -async def test_arealtime_vertex_hung_credential_resolution_raises_promptly(monkeypatch): +async def test_vertex_credential_resolution_returns_the_resolved_token_and_project(): + assert await realtime_main._resolve_vertex_access_token_bounded( + credentials="fake-credentials", + project_id="fake-project", + resolver=_instant_resolver, + timeout_seconds=5, + ) == ("token-abc", "resolved-project") + + +@pytest.mark.asyncio +async def test_vertex_credential_resolution_times_out_instead_of_hanging(): """Regression for the realtime accept-then-silence hang: a stalled Google - OAuth token refresh used to block _arealtime's vertex branch unbounded - (minutes of zero frames for the client). It must instead raise a clear, - prompt error naming the credential-resolution timeout.""" + OAuth token refresh used to block the vertex branch unbounded (minutes of + zero frames for the client). It must raise promptly and name the timeout.""" + start = time.monotonic() + with pytest.raises(ValueError, match="timed out fetching Google OAuth access token"): + await realtime_main._resolve_vertex_access_token_bounded( + credentials="fake-credentials", + project_id="fake-project", + resolver=_hanging_resolver, + timeout_seconds=0.05, + ) + assert time.monotonic() - start < 5 + + +@pytest.mark.asyncio +async def test_vertex_credential_resolution_bounds_a_thread_offloaded_refresh(): + """The real stall is a blocking google-auth refresh that runs in a worker + thread via asyncify, not a plain awaitable sleep. A timeout that only bounds + cancellable awaits would leave that shape hanging, so bound the shape the + proxy actually runs.""" + start = time.monotonic() + with pytest.raises(ValueError, match="timed out fetching Google OAuth access token"): + await realtime_main._resolve_vertex_access_token_bounded( + credentials="fake-credentials", + project_id="fake-project", + resolver=_thread_offloaded_hanging_resolver, + timeout_seconds=0.05, + ) + assert time.monotonic() - start < 5 + + +@pytest.mark.asyncio +async def test_arealtime_vertex_branch_resolves_credentials_under_a_bound(monkeypatch): + """The wiring half of the regression: the vertex branch of _arealtime must + go through the bounded resolver, so a hung token refresh surfaces as a + prompt error there rather than as an accepted-then-silent websocket.""" async def hanging_token_refresh(**kwargs): await asyncio.sleep(30) @@ -107,38 +165,7 @@ async def test_arealtime_vertex_hung_credential_resolution_raises_promptly(monke return model, "vertex_ai", None, api_base monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider) - monkeypatch.setattr(realtime_main.vertex_llm_base, "_ensure_access_token_async", hanging_token_refresh) - monkeypatch.setattr(realtime_main, "REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS", 0.05) - - start = time.monotonic() - with pytest.raises(ValueError, match="timed out fetching Google OAuth access token"): - await realtime_main._arealtime.__wrapped__( - model="gemini-live-2.5-flash", - websocket=MagicMock(), - litellm_logging_obj=FakeLogging(), - vertex_credentials="fake-credentials", - vertex_project="fake-project", - vertex_location="us-central1", - ) - assert time.monotonic() - start < 5 - - -@pytest.mark.asyncio -async def test_arealtime_vertex_credential_timeout_survives_thread_offloaded_refresh(monkeypatch): - """The real stall is a blocking google-auth refresh that runs in a worker - thread via asyncify, not a plain awaitable sleep. A timeout that only bounds - cancellable awaits would leave that shape hanging, so bound the shape the - proxy actually runs.""" - from litellm.litellm_core_utils.asyncify import asyncify - - async def thread_offloaded_hanging_refresh(**kwargs): - return await asyncify(time.sleep)(30) - - def mock_get_llm_provider(model, api_base, api_key): - return model, "vertex_ai", None, api_base - - monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider) - monkeypatch.setattr(realtime_main.vertex_llm_base, "_ensure_access_token_async", thread_offloaded_hanging_refresh) + monkeypatch.setattr(realtime_main, "vertex_access_token_resolver", hanging_token_refresh) monkeypatch.setattr(realtime_main, "REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS", 0.05) start = time.monotonic()