From 4f04e59ca036860ed1a83eb6169b4bdd20b31bb1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 20 Aug 2026 03:05:48 -0700 Subject: [PATCH] fix: harden vertex live passthrough against client model forms and dict credentials - accept the Live SDK's models/ and LiteLLM's vertex_ai/ when rewriting the setup model - keep a dict service account intact instead of stringifying it - treat same-target deployments holding different credentials as ambiguous - guard both websocket states before every close so a second close cannot raise - build the sendable close codes from the public CloseCode enum --- .../llm_passthrough_endpoints.py | 40 +++++-- .../pass_through_endpoints.py | 32 ++++-- .../passthrough_endpoint_router.py | 28 ++++- .../test_llm_pass_through_endpoints.py | 107 ++++++++++++++++++ .../test_pass_through_endpoints.py | 32 ++++++ .../test_passthrough_endpoint_router.py | 53 +++++++++ 6 files changed, 266 insertions(+), 26 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index c2b221b1f7a..7ce41c1d5b6 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -2384,6 +2384,21 @@ VERTEX_LIVE_UNCONFIGURED_CLOSE_REASON: Final = ( VERTEX_PUBLISHER_MODEL_PREFIX: Final = "publishers/google/models/" +VERTEX_PUBLISHERS_SEGMENT: Final = "publishers/" + + +def _vertex_publisher_model_suffix(model: str) -> str: + """ + Turn whatever the client named into the ``publishers//models/`` tail of a Vertex resource name. + + Clients send bare ids, LiteLLM ids (``vertex_ai/gemini-live-2.5-flash``), and the Live SDK's ``models/``, + and a publisher model id never contains a slash, so anything ahead of the last one is addressing, not identity + """ + publishers_at: Final = model.find(VERTEX_PUBLISHERS_SEGMENT) + if publishers_at != -1: + return model[publishers_at:] + return f"{VERTEX_PUBLISHER_MODEL_PREFIX}{model.rsplit('/', 1)[-1]}" + def _get_llm_router() -> "Router | None": from litellm.proxy.proxy_server import llm_router @@ -2397,8 +2412,12 @@ def _resolve_vertex_live_credentials( model: str | None, ) -> VertexPassThroughCredentials | None: """ - Resolution order: an explicit project/location registration or ``default_vertex_config``, then any DB model - entry flagged ``use_in_pass_through``, then the ``DEFAULT_VERTEXAI_*`` env vars + Resolution order: an explicit project/location registration, then ``default_vertex_config`` (which the proxy + fills from the ``DEFAULT_VERTEXAI_*`` env vars whenever the yaml leaves it out), then any DB model entry + flagged ``use_in_pass_through``. + + DB entries come last on purpose: an operator who set a global default already said which project + pass-through traffic should bill to, and this route silently ignoring that would be the worse surprise """ keyed: Final = passthrough_endpoint_router.get_vertex_credentials( project_id=vertex_project, @@ -2436,22 +2455,23 @@ def _build_vertex_live_setup_model_rewriter( if setup_model.startswith("projects/"): return setup_model aliased: Final = _resolve_alias_to_upstream_model(setup_model, llm_router) - return ( - f"projects/{vertex_project}/locations/{vertex_location}/" - f"{VERTEX_PUBLISHER_MODEL_PREFIX}{aliased.removeprefix(VERTEX_PUBLISHER_MODEL_PREFIX)}" - ) + return f"projects/{vertex_project}/locations/{vertex_location}/{_vertex_publisher_model_suffix(aliased)}" return rewrite def _resolve_alias_to_upstream_model(setup_model: str, llm_router: "Router | None") -> str: + """ + The Live SDK wraps whatever the caller typed as ``models/``, so a gateway alias arrives prefixed + """ if llm_router is None: return setup_model + candidates: Final = (setup_model, setup_model.rsplit("/", 1)[-1]) upstream: Final = next( ( - deployment["litellm_params"]["model"] + deployment["litellm_params"].get("model") for deployment in (llm_router.get_model_list() or ()) - if deployment.get("model_name") == setup_model + if deployment.get("model_name") in candidates ), None, ) @@ -2500,9 +2520,7 @@ async def vertex_ai_live_websocket_passthrough( vertex_credentials_config.vertex_location if vertex_credentials_config is not None else None ) credentials_value: Final = ( - str(vertex_credentials_config.vertex_credentials) - if vertex_credentials_config is not None and vertex_credentials_config.vertex_credentials is not None - else None + vertex_credentials_config.vertex_credentials if vertex_credentials_config is not None else None ) try: diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 5608e8b384d..d45421489e7 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -32,7 +32,7 @@ from websockets.exceptions import ( ConnectionClosedOK, InvalidStatus, ) -from websockets.frames import EXTERNAL_CLOSE_CODES, Close +from websockets.frames import Close, CloseCode import litellm from litellm._logging import verbose_proxy_logger @@ -1928,11 +1928,26 @@ def _truncated_close_reason(reason: str) -> str: return encoded[:WEBSOCKET_CLOSE_REASON_MAX_BYTES].decode("utf-8", errors="ignore") +SENDABLE_CLOSE_CODES: Final = frozenset(CloseCode) - frozenset( + {CloseCode.NO_STATUS_RCVD, CloseCode.ABNORMAL_CLOSURE, CloseCode.TLS_HANDSHAKE} +) + + +def _client_socket_is_open(websocket: WebSocket) -> bool: + """ + Starlette tracks the two halves separately and raises on a second close, so both have to still be live + """ + return ( + websocket.client_state != WebSocketState.DISCONNECTED + and websocket.application_state != WebSocketState.DISCONNECTED + ) + + def _upstream_close_to_relay(task_results: Iterable[object]) -> Close | None: """ The upstream close worth telling the client about: anything other than a plain, reasonless normal close. - Codes outside ``EXTERNAL_CLOSE_CODES`` and the private range never travel on the wire (1006 for a socket that + Codes outside ``SENDABLE_CLOSE_CODES`` and the private range never travel on the wire (1006 for a socket that died without a close frame, 1005 for one that sent no code), so relaying them would build an invalid frame """ upstream_close: Final = next((result for result in task_results if isinstance(result, Close)), None) @@ -1940,7 +1955,7 @@ def _upstream_close_to_relay(task_results: Iterable[object]) -> Close | None: return None if upstream_close.code == 1000 and upstream_close.reason == "": return None - if upstream_close.code not in EXTERNAL_CLOSE_CODES and not 3000 <= upstream_close.code < 5000: + if upstream_close.code not in SENDABLE_CLOSE_CODES and not 3000 <= upstream_close.code < 5000: return None return upstream_close @@ -2268,7 +2283,7 @@ async def websocket_passthrough_request( raise exception upstream_close: Final = _upstream_close_to_relay(task.result() for task in done) - if upstream_close is not None and websocket.application_state != WebSocketState.DISCONNECTED: + if upstream_close is not None and _client_socket_is_open(websocket): await websocket.close( code=upstream_close.code, reason=_truncated_close_reason(upstream_close.reason), @@ -2359,7 +2374,7 @@ async def websocket_passthrough_request( ), ) - if websocket.client_state != WebSocketState.DISCONNECTED: + if _client_socket_is_open(websocket): await websocket.close( code=getattr(exc, "status_code", 1011), reason="Upstream connection rejected", @@ -2387,13 +2402,10 @@ async def websocket_passthrough_request( ), ) - if websocket.client_state != WebSocketState.DISCONNECTED: + if _client_socket_is_open(websocket): await websocket.close(code=1011, reason="WebSocket passthrough error") finally: - if ( - websocket.client_state != WebSocketState.DISCONNECTED - and websocket.application_state != WebSocketState.DISCONNECTED - ): + if _client_socket_is_open(websocket): await websocket.close() diff --git a/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py b/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py index 28067c842cd..7fd607fc4d0 100644 --- a/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py +++ b/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py @@ -1,3 +1,4 @@ +import json from collections.abc import Callable from typing import TYPE_CHECKING, Final @@ -27,6 +28,15 @@ def _get_str_value(values: dict[str, object] | None, key: str) -> str | None: return value if isinstance(value, str) else None +def _credential_identity(credentials: VERTEX_CREDENTIALS_TYPES | None) -> str | None: + """ + A hashable stand-in for a credential, so two deployments can be compared for holding the same one + """ + if isinstance(credentials, dict): + return json.dumps(credentials, sort_keys=True) + return credentials + + class PassthroughEndpointRouter: """ Use this class to Get credentials for pass-through endpoints @@ -127,8 +137,8 @@ class PassthroughEndpointRouter: ``deployment_key_to_vertex_credentials`` is only reachable when the caller names a project and location, which WebSocket clients never do, so DB-stored deployments need this lookup to be usable at all. - With no model to go on, only deployments that agree on a project and location answer: guessing between - two Vertex projects would mint a token for one and later send the other one's model name + With no model to go on, only deployments that agree on a project, a location, and a credential answer: + guessing between two Vertex projects would mint a token for one and later send the other one's model name """ llm_router: Final = self.llm_router_getter() if llm_router is None: @@ -149,7 +159,12 @@ class PassthroughEndpointRouter: if matched is not None: return matched targets: Final = frozenset( - (credentials.vertex_project, credentials.vertex_location) for _, credentials in resolved + ( + credentials.vertex_project, + credentials.vertex_location, + _credential_identity(credentials.vertex_credentials), + ) + for _, credentials in resolved ) if len(targets) != 1: return None @@ -172,9 +187,12 @@ class PassthroughEndpointRouter: vertex_location: Final = _get_str_value(credential_values, "vertex_location") or litellm_params.get( "vertex_location" ) - vertex_credentials: Final = _get_str_value(credential_values, "vertex_credentials") or litellm_params.get( - "vertex_credentials" + stored_credentials: Final = ( + credential_values.get("vertex_credentials") if credential_values is not None else None ) + vertex_credentials: Final = ( + stored_credentials if isinstance(stored_credentials, (str, dict)) else None + ) or litellm_params.get("vertex_credentials") if vertex_project is None or vertex_location is None: return None return VertexPassThroughCredentials( diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 5ed974c7a47..f994fba371b 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -3889,6 +3889,9 @@ class TestComprehendMedicalProxyRoute: assert exc_info.value.status_code == 400 +LIVE_RESOURCE_PATH = "projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash" + + class TestVertexAILiveWebsocketPassthrough: def _websocket(self): from starlette.websockets import WebSocketState @@ -3959,6 +3962,110 @@ class TestVertexAILiveWebsocketPassthrough: ) websocket.close.assert_not_awaited() + @pytest.mark.parametrize( + "setup_model, expected", + [ + ("gemini-live-2.5-flash", LIVE_RESOURCE_PATH), + ("models/gemini-live-2.5-flash", LIVE_RESOURCE_PATH), + ("vertex_ai/gemini-live-2.5-flash", LIVE_RESOURCE_PATH), + ("gemini-live", LIVE_RESOURCE_PATH), + ("models/gemini-live", LIVE_RESOURCE_PATH), + ( + "publishers/meta/models/llama-3.3-70b-instruct-maas", + "projects/proj-db/locations/global/publishers/meta/models/llama-3.3-70b-instruct-maas", + ), + ( + "projects/other/locations/us-central1/publishers/google/models/gemini-2.0-flash", + "projects/other/locations/us-central1/publishers/google/models/gemini-2.0-flash", + ), + ], + ) + def test_setup_model_rewriter_normalises_the_forms_clients_send(self, setup_model, expected): + from litellm.proxy.pass_through_endpoints import ( + llm_passthrough_endpoints as passthrough_module, + ) + + llm_router = litellm.Router( + model_list=[ + { + "model_name": "gemini-live", + "litellm_params": { + "model": "vertex_ai/gemini-live-2.5-flash", + "use_in_pass_through": True, + "vertex_project": "proj-db", + "vertex_location": "global", + }, + } + ] + ) + + rewriter = passthrough_module._build_vertex_live_setup_model_rewriter( + vertex_project="proj-db", + vertex_location="global", + llm_router=llm_router, + ) + + assert rewriter is not None + assert rewriter(setup_model) == expected + + @pytest.mark.asyncio + async def test_default_vertex_config_outranks_db_deployment(self, monkeypatch): + from litellm.proxy.pass_through_endpoints import ( + llm_passthrough_endpoints as passthrough_module, + ) + from litellm.types.passthrough_endpoints.vertex_ai import ( + VertexPassThroughCredentials, + ) + + llm_router = litellm.Router( + model_list=[ + { + "model_name": "gemini-live", + "litellm_params": { + "model": "vertex_ai/gemini-live-2.5-flash", + "use_in_pass_through": True, + "vertex_project": "proj-db", + "vertex_location": "global", + "vertex_credentials": '{"type": "db_account"}', + }, + } + ] + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + monkeypatch.setattr( + passthrough_module.passthrough_endpoint_router, + "default_vertex_config", + VertexPassThroughCredentials( + vertex_project="proj-env", + vertex_location="global", + vertex_credentials='{"type": "env_account"}', + ), + ) + self._clear_vertex_env(monkeypatch) + websocket = self._websocket() + ensure_token = AsyncMock(return_value=("token-abc", "proj-env")) + ws_passthrough = AsyncMock() + + with ( + patch.object(passthrough_module.vertex_llm_base, "_ensure_access_token_async", ensure_token), + patch.object(passthrough_module, "websocket_passthrough_request", ws_passthrough), + ): + await passthrough_module.vertex_ai_live_websocket_passthrough( + websocket=websocket, + model="gemini-live", + user_api_key_dict=UserAPIKeyAuth(), + ) + + ensure_token.assert_awaited_once_with( + credentials='{"type": "env_account"}', + project_id="proj-env", + custom_llm_provider="vertex_ai_beta", + ) + rewriter = ws_passthrough.await_args.kwargs["setup_model_rewriter"] + assert rewriter("gemini-live") == ( + "projects/proj-env/locations/global/publishers/google/models/gemini-live-2.5-flash" + ) + @pytest.mark.asyncio async def test_credential_failure_close_names_configuration_options(self, monkeypatch): from litellm.proxy.pass_through_endpoints import ( diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 2663396bf53..00097166c13 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -5243,6 +5243,38 @@ async def test_websocket_passthrough_leaves_full_resource_setup_model_untouched( ) +@pytest.mark.asyncio +async def test_websocket_passthrough_does_not_close_twice_when_success_logging_fails(): + from websockets.exceptions import ConnectionClosedError + from websockets.frames import Close + + upstream_reason = "Publisher Model `projects/p/locations/global/publishers/google/models/nope` was not found" + upstream_ws = ClosingUpstreamWebSocket( + ConnectionClosedError(rcvd=Close(1008, upstream_reason), sent=Close(1008, ""), rcvd_then_sent=True) + ) + websocket = _client_websocket(_pending_receive) + + with ( + _patched_websocket_passthrough_environment(upstream_ws), + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints." + "GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue", + side_effect=RuntimeError("logging worker down"), + ), + ): + await websocket_passthrough_request( + websocket=websocket, + target="wss://aiplatform.googleapis.com/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent", + custom_headers={"Authorization": "Bearer token"}, + user_api_key_dict=UserAPIKeyAuth(), + forward_headers=False, + endpoint="/vertex_ai/live", + accept_websocket=False, + ) + + websocket.close.assert_awaited_once_with(code=1008, reason=upstream_reason) + + def _passthrough_kwargs_for_reservation( user_api_key_dict: UserAPIKeyAuth, parsed_body: Optional[dict] = None ) -> dict: diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoint_router.py b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoint_router.py index f86bfb8bb1b..e3cbc2d507f 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoint_router.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoint_router.py @@ -293,6 +293,59 @@ def test_vertex_without_hint_falls_back_when_deployments_share_a_target(): assert resolved is not None and resolved.vertex_project == "proj-one" +def test_vertex_without_hint_refuses_to_guess_between_service_accounts(): + llm_router = litellm.Router( + model_list=[ + _vertex_deployment( + "gemini-flash", + "vertex_ai/gemini-2.5-flash", + vertex_project="proj-one", + vertex_location="global", + vertex_credentials='{"client_email": "flash@proj-one.iam"}', + ), + _vertex_deployment( + "gemini-live", + "vertex_ai/gemini-live-2.5-flash", + vertex_project="proj-one", + vertex_location="global", + vertex_credentials='{"client_email": "live@proj-one.iam"}', + ), + ] + ) + passthrough_router = _passthrough_router(llm_router) + + assert passthrough_router.get_vertex_credentials_from_router_deployments(model=None) is None + + +def test_vertex_named_credential_keeps_dict_service_account(): + service_account = {"type": "service_account", "client_email": "live@proj-db.iam"} + CredentialAccessor.upsert_credentials( + [ + _vertex_credential( + "cred_gcp_dict", + { + "vertex_project": "proj-db", + "vertex_location": "global", + "vertex_credentials": service_account, + }, + ) + ] + ) + llm_router = litellm.Router( + model_list=[ + _vertex_deployment( + "gemini-live", "vertex_ai/gemini-live-2.5-flash", litellm_credential_name="cred_gcp_dict" + ) + ] + ) + passthrough_router = _passthrough_router(llm_router) + + resolved = passthrough_router.get_vertex_credentials_from_router_deployments(model=None) + + assert resolved is not None + assert resolved.vertex_credentials == service_account + + def test_no_flagged_vertex_deployment_returns_none(): llm_router = litellm.Router( model_list=[