fix(realtime): bound Vertex credential resolution and make realtime failures loud - #37604
Conversation
…ailures loud A /v1/realtime connection to a Vertex AI Live model accepted the WebSocket upgrade and then went silent: a stalled Google OAuth token fetch blocked the handler before any session event, and the eventual failure closed the socket with a bare 1011 and no error event, so callers saw an open socket, no frames, and no reason. Bound the pre-session token fetch with REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS (20s default) and, on any realtime failure, send an OpenAI-style error event before closing with a reason that names the failure. Close reasons are truncated by bytes, not characters, since an over-long reason makes the close frame itself fail.
|
Mateo seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account. You have signed the CLA already but the status is still pending? Let us recheck it. |
Greptile SummaryBounds Vertex AI realtime credential resolution and improves WebSocket failure reporting
Confidence Score: 5/5The PR appears safe to merge No blocking failure remains
|
| Filename | Overview |
|---|---|
| litellm/realtime_api/main.py | Adds bounded Vertex credential resolution to realtime sessions and health checks without leaving an eligible blocking issue |
| litellm/litellm_core_utils/realtime_errors.py | Introduces shared error-event serialization and RFC-compliant close-reason truncation |
| litellm/llms/custom_httpx/llm_http_handler.py | Sends a structured error event and reasoned close when realtime backend setup fails |
| litellm/proxy/proxy_server.py | Makes phase-two realtime routing failures visible while tolerating an already-closed client socket |
| litellm/types/llms/vertex_ai.py | Adds the typed asynchronous access-token resolver protocol used by bounded credential resolution |
| litellm/types/realtime.py | Defines typed OpenAI-style realtime error event payloads |
| tests/test_litellm/realtime_api/test_main.py | Covers successful, stalled, thread-offloaded, and wired Vertex credential resolution |
| tests/test_litellm/proxy/realtime_endpoints/test_realtime_webrtc_endpoints.py | Verifies phase-two failures produce an error event and safe close behavior |
Reviews (3): Last reviewed commit: "fix(realtime): resolve the vertex token ..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
![veria-ai[bot]](https://avatars.githubusercontent.com/in/1723776?s=60&v=4){kind=small}
| await websocket.close(code=1011, reason="Internal server error") | ||
| redacted_error: Final = _redact_string(str(e)) | ||
| try: | ||
| await websocket.send_text(realtime_error_event(redacted_error, error_type="server_error")) |
There was a problem hiding this comment.
Low: Internal exception disclosure
When realtime routing or backend setup fails, an authenticated caller now receives the raw exception text. _redact_string removes recognized credential patterns, but leaves details such as GCP project IDs, private upstream hostnames, and filesystem paths; the same issue occurs in litellm/llms/custom_httpx/llm_http_handler.py:5982. Return a fixed public error message for unexpected exceptions, whitelist specific safe messages such as the credential timeout, and retain the full exception only in server logs.
There was a problem hiding this comment.
llm_http_handler already sent this redacted text before the PR. Clamping the proxy to a fixed string restores the opaque failure this fix removes
There was a problem hiding this comment.
Thanks for the detail. I can't automatically confirm this is safe to dismiss, so I'm leaving the thread open for a maintainer to make the call.
PR overviewThis pull request bounds Vertex credential resolution and changes realtime routing and backend setup failures to surface errors more clearly. It updates realtime proxy and custom HTTP handling paths. One low-impact information disclosure remains open: an authenticated caller who triggers certain realtime failures may receive raw exception details, including project IDs, internal hostnames, or filesystem paths. No issues have yet been addressed, so unexpected exceptions should still be replaced with a fixed public message while full details remain in server logs. Open issues (1)
Fixed/addressed: 0 · PR risk: 3/10 |
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
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.
Binding the bound method at import froze the module-level VertexBase instance, so callers that swap it no longer reached their replacement.
|
bugbot run |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using high effort and found 1 potential issue.
Autofix Details
Bugbot Autofix prepared a fix for the issue found in the latest run.
- ✅ Fixed: Timeout cancels Vertex refresh lock
- Scheduled the resolver as a Task, added a module-level strong-ref registry, and shielded it inside
wait_forso the inflight refresh keeps VertexBase's per-key async lock held past timeout, forcing retries to queue on it rather than spawn new asyncify(refresh_auth) worker threads that would starve anyio's thread limiter.
- Scheduled the resolver as a Task, added a module-level strong-ref registry, and shielded it inside
Or push these changes by commenting:
@cursor push 7d4980c7e6
Preview (7d4980c7e6)
diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py
--- a/litellm/realtime_api/main.py
+++ b/litellm/realtime_api/main.py
@@ -299,21 +299,33 @@
)
+_INFLIGHT_VERTEX_TOKEN_TASKS: Final[set[asyncio.Task[tuple[str, str]]]] = set() # mutable-ok: strong-ref registry so bounded resolvers keep running past timeout and keep VertexBase's per-key async refresh lock held; without it retries would spawn new asyncify(refresh_auth) threads that queue on the instance-wide _sync_refresh_lock and starve anyio's thread limiter
+
+
+def _drop_inflight_vertex_token_task(task: asyncio.Task[tuple[str, str]]) -> None:
+ _INFLIGHT_VERTEX_TOKEN_TASKS.discard(task)
+ if not task.cancelled():
+ # Consume the exception so tasks that outlived their bounded caller don't emit "exception was never retrieved" warnings.
+ _ = task.exception()
+
+
async def _resolve_vertex_access_token_bounded(
credentials: VERTEX_CREDENTIALS_TYPES | None,
project_id: str | None,
resolver: VertexAccessTokenResolver,
timeout_seconds: float,
) -> tuple[str, str]:
+ task: Final = asyncio.ensure_future(
+ resolver(
+ credentials=credentials,
+ project_id=project_id,
+ custom_llm_provider="vertex_ai",
+ )
+ )
+ _INFLIGHT_VERTEX_TOKEN_TASKS.add(task)
+ task.add_done_callback(_drop_inflight_vertex_token_task)
try:
- return await asyncio.wait_for(
- resolver(
- credentials=credentials,
- project_id=project_id,
- custom_llm_provider="vertex_ai",
- ),
- timeout=timeout_seconds,
- )
+ return await asyncio.wait_for(asyncio.shield(task), timeout=timeout_seconds)
except asyncio.TimeoutError as e:
raise ValueError(
"Vertex AI realtime: timed out fetching Google OAuth access token after "
diff --git a/tests/test_litellm/realtime_api/test_main.py b/tests/test_litellm/realtime_api/test_main.py
--- a/tests/test_litellm/realtime_api/test_main.py
+++ b/tests/test_litellm/realtime_api/test_main.py
@@ -153,6 +153,52 @@
@pytest.mark.asyncio
+async def test_bounded_resolver_keeps_inflight_refresh_alive_on_timeout():
+ """Regression: `asyncio.wait_for` used to cancel the token resolver on
+ timeout, which released VertexBase's per-key async refresh lock while the
+ underlying google-auth refresh thread kept running holding the instance-wide
+ `_sync_refresh_lock`. Router retries would then spawn additional refresh
+ threads that queued on that sync lock and starved anyio's thread limiter.
+ The bounded resolver must let the inflight resolution keep running past the
+ caller's bound so the per-key async lock stays held and later callers
+ (retries or unrelated realtime traffic) queue on the async lock instead of
+ spawning new blocking refresh threads."""
+ started = asyncio.Event()
+
+ async def slow_resolver(credentials, project_id, custom_llm_provider) -> tuple[str, str]:
+ started.set()
+ await asyncio.sleep(30)
+ return "", ""
+
+ before = set(realtime_main._INFLIGHT_VERTEX_TOKEN_TASKS)
+ 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=slow_resolver,
+ timeout_seconds=0.05,
+ )
+ assert started.is_set()
+ new_tasks = realtime_main._INFLIGHT_VERTEX_TOKEN_TASKS - before
+ assert len(new_tasks) == 1
+ inflight = next(iter(new_tasks))
+ try:
+ assert not inflight.done(), (
+ "Bounded resolver ended the inflight token refresh on timeout: "
+ "VertexBase's per-key async lock would be released while the refresh "
+ "thread still holds _sync_refresh_lock, letting retries spawn new "
+ "refresh threads that queue on the sync lock"
+ )
+ assert not inflight.cancelled()
+ finally:
+ inflight.cancel()
+ try:
+ await inflight
+ except BaseException:
+ pass
+
+
+@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 aYou can send follow-ups to the cloud agent here.
|
bugbot run |
There was a problem hiding this comment.
✅ Bugbot reviewed your changes and found no new issues!
Comment @cursor review or bugbot run to trigger another review on this PR
Reviewed by Cursor Bugbot for commit d643136. Configure here.
3 participants
TLDR
Problem this solves:
/v1/realtimeaccepts the upgrade, then goes silent for minutesHow it solves it:
errorevent before closing the socketUser Flow
Before: a developer pointing a voice app at a Vertex AI Live model gets a socket that connects and then never says anything
wss://litellm-domain/v1/realtime?model=gemini-live-2.5-flash-native-audiowith their virtual keysession.createdarrives, noerrorevent arrives, and the socket sits open with zero frames on itInternal server error, having delivered nothinggpt-realtime-2025-08-28on the same gateway connects and streams normally, so nothing on the client side looks brokenAfter: the same connection fails within about a minute and says exactly what went wrong
wss://litellm-domain/v1/realtime?model=gemini-live-2.5-flash-native-audiowith their virtual keyerrorevent of typeserver_errorwhose message says the Google OAuth token fetch timed out after 20s and names egress tooauth2.googleapis.comas the thing to checksession.createdin under a second, andgpt-realtime-2025-08-28is unchangedRelevant issues
Linear ticket
Resolves LIT-5867
Pre-Submission checklist
uv run pytest tests/test_litellm/<your_test_file>.py -v. Leave the suites (make test-unit-*,make test-unit) to CI: it finishes in ~15 minutes where a laptop takes an hour or more@greptileaito re-request a review after pushing changes)Screenshots / Proof of Fix
Shared setup, identical on both sides. One proxy against real Vertex AI and real OpenAI, on a random high port, with three Vertex AI Live deployments registered through
/model/new: one shaped the way a customer's config is (tags: [],guardrails: [],vector_store_ids: []), one clean, and one whose service-account JSON has itstoken_uripointed at an unroutable address so the OAuth token fetch hangs the way blocked egress does in the customer's cluster.Before (5290150)
Vertex AI Live, OAuth token endpoint unreachable
python ws_probe.py 41207 gemini-live-badtoken 300Vertex AI Live, working credentials
python ws_probe.py 41207 gemini-live-2.5-flash-native-audio 20andpython ws_probe.py 41207 gemini-live-clean 20OpenAI realtime control
python ws_probe.py 41207 gpt-realtime-2025-08-28 20After (d643136)
Vertex AI Live, OAuth token endpoint unreachable
python ws_probe.py 39177 gemini-live-badtoken 300Vertex AI Live, working credentials
python ws_probe.py 39177 gemini-live-2.5-flash-native-audio 20andpython ws_probe.py 39177 gemini-live-clean 20OpenAI realtime control
python ws_probe.py 39177 gpt-realtime-2025-08-28 20Type
🐛 Bug Fix
Caveats (if any)
REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDStunes it per deploymentFinal Attestation
Note
Cursor Bugbot is generating a summary for commit d643136. Configure here.