From b6b0e2f1c55be1db1f473ef01843dfdb3043ed14 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 18:13:09 -0700 Subject: [PATCH 01/17] test(proxy): cover model access on the Vertex AI Live websocket passthrough (internal copy of #40078) Co-authored-by: Marty Sullivan --- .../test_pass_through_endpoints.py | 483 +++++++++++++++++- 1 file changed, 481 insertions(+), 2 deletions(-) diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index f0baceecdaa..953aefe7f7c 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -35,7 +35,7 @@ from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._lazy_features import LazyFeature, attach_lazy_features -from litellm.proxy._types import ProxyException, UserAPIKeyAuth +from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import router as llm_passthrough_router from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( @@ -7175,6 +7175,39 @@ class RecordingUpstreamWebSocket: raise StopAsyncIteration +class ClosableUpstreamWebSocket: + """An upstream whose ``recv`` ends when ``close`` does, and whose close keeps yielding after that. + + A real close handshake outlives the ``recv`` it unblocks, so the reader task finishes while the closing + task is still suspended. That is the interleaving that leaves a refusal stranded in a cancelled task + """ + + CLOSE_HANDSHAKE_YIELDS = 5 + + def __init__(self): + self._closing = asyncio.Event() + self.send = AsyncMock() + self.close = AsyncMock(side_effect=self._close) + + async def _close(self): + self._closing.set() + for _ in range(self.CLOSE_HANDSHAKE_YIELDS): + await asyncio.sleep(0) + + async def recv(self, decode: bool = True): + from websockets.exceptions import ConnectionClosedOK + from websockets.frames import Close + + await self._closing.wait() + raise ConnectionClosedOK(rcvd=Close(1000, ""), sent=Close(1000, ""), rcvd_then_sent=True) + + def __aiter__(self): + return self + + async def __anext__(self): + raise StopAsyncIteration + + def _client_websocket(receive): from starlette.websockets import WebSocketState @@ -7210,7 +7243,7 @@ def _patched_websocket_passthrough_environment(upstream_ws): mock_worker.ensure_initialized_and_enqueue = MagicMock( side_effect=lambda async_coroutine: async_coroutine.close() ) - yield + yield SimpleNamespace(proxy_logging=mock_proxy_logging, logging_worker=mock_worker) async def _pending_receive(): @@ -7336,6 +7369,452 @@ async def test_websocket_passthrough_rewrites_gateway_alias_setup_model(): assert sent_setup["model"] == "projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash" +VERTEX_LIVE_TARGET = "wss://aiplatform.googleapis.com/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent" + + +VERTEX_LIVE_GROUP = "gemini-live-2.5-flash-native-audio" + +VERTEX_LIVE_ALIAS_GROUP = "gemini-live-native-audio" + +VERTEX_LIVE_CLIENT_ADDRESSINGS = ( + "{model}", + "models/{model}", + "vertex_ai/{model}", + "publishers/google/models/{model}", + "projects/proj-db/locations/global/publishers/google/models/{model}", +) + + +def _vertex_live_router(): + return litellm.Router( + model_list=[ + { + "model_name": "gemini-live-2.5-flash", + "litellm_params": {"model": "vertex_ai/gemini-live-2.5-flash"}, + }, + { + "model_name": VERTEX_LIVE_GROUP, + "litellm_params": {"model": f"vertex_ai/{VERTEX_LIVE_GROUP}"}, + }, + { + "model_name": VERTEX_LIVE_ALIAS_GROUP, + "litellm_params": {"model": f"vertex_ai/{VERTEX_LIVE_GROUP}"}, + }, + { + "model_name": "gemini-3-pro-live", + "litellm_params": {"model": "vertex_ai/gemini-3-pro-live"}, + }, + ] + ) + + +async def _run_vertex_live_gated_passthrough(client_frames, valid_token): + """Drive /vertex_ai/live with the real frame gate, returning the client and upstream sockets""" + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _build_vertex_live_client_frame_gate, + _build_vertex_live_setup_model_rewriter, + ) + + llm_router = _vertex_live_router() + upstream_ws = RecordingUpstreamWebSocket() + websocket = _client_websocket( + AsyncMock(side_effect=[*client_frames, {"type": "websocket.disconnect"}]), + ) + + with _patched_websocket_passthrough_environment(upstream_ws): + await websocket_passthrough_request( + websocket=websocket, + target=VERTEX_LIVE_TARGET, + custom_headers={"Authorization": "Bearer token"}, + user_api_key_dict=valid_token, + forward_headers=False, + endpoint="/vertex_ai/live", + accept_websocket=False, + setup_model_rewriter=_build_vertex_live_setup_model_rewriter( + vertex_project="proj-db", + vertex_location="global", + llm_router=llm_router, + ), + client_frame_gate=_build_vertex_live_client_frame_gate( + llm_model_list=llm_router.get_model_list(), + llm_router=llm_router, + ), + ) + + return websocket, upstream_ws + + +def _setup_frame(model: str) -> dict: + return { + "type": "websocket.receive", + "text": json.dumps({"setup": {"model": model, "generationConfig": {"responseModalities": ["TEXT"]}}}), + } + + +def _policy_close_reason(websocket) -> str | None: + return next( + (call.kwargs.get("reason") for call in websocket.close.await_args_list if call.kwargs.get("code") == 1008), + None, + ) + + +@pytest.mark.asyncio +async def test_vertex_live_setup_frame_naming_a_model_the_key_cannot_call_is_refused(): + """ + ?model= is optional on this route, so a key that names an unlisted model only in the setup frame used to + reach Vertex with the gateway's own service-account credentials + """ + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame("totally-unlisted-model-abc")], + UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), + ) + + upstream_ws.send.assert_not_awaited() + reason = _policy_close_reason(websocket) + assert reason is not None and "totally-unlisted-model-abc" in reason + upstream_ws.close.assert_awaited() + + +@pytest.mark.asyncio +async def test_vertex_live_setup_frame_naming_a_permitted_model_still_reaches_vertex(): + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame("gemini-live-2.5-flash")], + UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), + ) + + upstream_ws.send.assert_awaited_once() + sent_setup = json.loads(upstream_ws.send.await_args.args[0])["setup"] + assert sent_setup["model"] == "projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash" + assert _policy_close_reason(websocket) is None + + +@pytest.mark.asyncio +async def test_vertex_live_setup_frame_from_an_unrestricted_key_still_reaches_vertex(): + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame("gemini-3-pro-live")], + UserAPIKeyAuth(token="hashed", models=[]), + ) + + upstream_ws.send.assert_awaited_once() + assert _policy_close_reason(websocket) is None + + +@pytest.mark.asyncio +async def test_vertex_live_setup_frame_is_checked_against_the_team_for_an_all_team_models_key(): + """ + ``all-team-models`` skips the key allowlist, so the team's own restriction is the only thing left to apply + """ + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame("gemini-3-pro-live")], + UserAPIKeyAuth( + token="hashed", + models=["all-team-models"], + team_id="team-1", + team_models=["gemini-live-2.5-flash"], + ), + ) + + upstream_ws.send.assert_not_awaited() + reason = _policy_close_reason(websocket) + assert reason is not None and "gemini-3-pro-live" in reason + + +@pytest.mark.asyncio +async def test_vertex_live_setup_frame_is_gated_when_it_arrives_as_binary(): + """A client that sends its setup frame as bytes must not skip the check the text path applies""" + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [ + { + "type": "websocket.receive", + "bytes": json.dumps({"setup": {"model": "totally-unlisted-model-abc"}}).encode(), + } + ], + UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), + ) + + upstream_ws.send.assert_not_awaited() + assert _policy_close_reason(websocket) is not None + + +@pytest.mark.asyncio +async def test_vertex_live_setup_frame_with_an_escaped_key_is_still_gated(): + """Large audio frames skip the parse on a substring test, so a JSON-escaped key must not slip past it""" + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [{"type": "websocket.receive", "text": '{"\\u0073etup": {"model": "totally-unlisted-model-abc"}}'}], + UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), + ) + + upstream_ws.send.assert_not_awaited() + assert _policy_close_reason(websocket) is not None + + +@pytest.mark.asyncio +async def test_vertex_live_non_setup_frames_are_forwarded_without_a_model_check(): + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [ + _setup_frame("gemini-live-2.5-flash"), + {"type": "websocket.receive", "text": json.dumps({"realtimeInput": {"audio": {"data": "AAAA"}}})}, + ], + UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), + ) + + assert upstream_ws.send.await_count == 2 + assert _policy_close_reason(websocket) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS) +@pytest.mark.parametrize("group", [VERTEX_LIVE_GROUP, VERTEX_LIVE_ALIAS_GROUP]) +async def test_vertex_live_setup_frame_authorizes_a_permitted_group_in_every_addressing(addressing, group): + """ + The Live SDK wraps the caller's model as ``models/`` and the docs tell callers to send the full + resource path, so a key holding the group has to work whichever addressing its client actually sends + """ + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame(addressing.format(model=group))], + UserAPIKeyAuth(token="hashed", models=[group]), + ) + + upstream_ws.send.assert_awaited_once() + sent_setup = json.loads(upstream_ws.send.await_args.args[0])["setup"] + assert sent_setup["model"].startswith("projects/") + assert _policy_close_reason(websocket) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS) +async def test_vertex_live_setup_frame_refuses_an_unpermitted_group_in_every_addressing(addressing): + """Normalizing the addressing must not become a way to reach a group the key does not hold""" + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame(addressing.format(model="gemini-3-pro-live"))], + UserAPIKeyAuth(token="hashed", models=[VERTEX_LIVE_GROUP]), + ) + + upstream_ws.send.assert_not_awaited() + reason = _policy_close_reason(websocket) + assert reason is not None and "gemini-3-pro-live" in reason + + +@pytest.mark.asyncio +@pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS) +async def test_vertex_live_setup_frame_naming_no_group_at_all_is_refused_in_every_addressing(addressing): + """A name that resolves to no group is authorized as it arrived, so an unknown model cannot ride a prefix in""" + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame(addressing.format(model="totally-unlisted-model-abc"))], + UserAPIKeyAuth(token="hashed", models=[VERTEX_LIVE_GROUP]), + ) + + upstream_ws.send.assert_not_awaited() + assert _policy_close_reason(websocket) is not None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS[1:]) +async def test_vertex_live_setup_frame_does_not_strip_a_prefix_off_a_model_no_group_serves(addressing): + """ + Stripping the addressing off a name the router does not serve would let a wildcard key reach an arbitrary + publisher model, and a full path would carry a project of the caller's choosing with it + """ + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame(addressing.format(model="gemini-4-pro-live-unserved"))], + UserAPIKeyAuth(token="hashed", models=["gemini-*"]), + ) + + upstream_ws.send.assert_not_awaited() + assert _policy_close_reason(websocket) is not None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS) +async def test_vertex_live_setup_frame_is_checked_against_the_team_in_every_addressing(addressing): + """The team allowlist is keyed by group too, so the addressing must not decide whether the team is checked""" + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame(addressing.format(model="gemini-3-pro-live"))], + UserAPIKeyAuth( + token="hashed", + models=["all-team-models"], + team_id="team-1", + team_models=[VERTEX_LIVE_GROUP], + ), + ) + + upstream_ws.send.assert_not_awaited() + reason = _policy_close_reason(websocket) + assert reason is not None and "gemini-3-pro-live" in reason + + +@pytest.mark.asyncio +@pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS) +async def test_vertex_live_setup_frame_resolves_an_access_group_grant_in_every_addressing(addressing): + """ + ``get_model_access_groups`` is keyed by group as well, and returns nothing for a prefixed name, so a key + entitled through an access group rather than a model name needs the same normalization + """ + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _build_vertex_live_client_frame_gate, + ) + + llm_router = litellm.Router( + model_list=[ + { + "model_name": VERTEX_LIVE_GROUP, + "litellm_params": {"model": f"vertex_ai/{VERTEX_LIVE_GROUP}"}, + "model_info": {"access_groups": ["live-team"]}, + }, + ] + ) + gate = _build_vertex_live_client_frame_gate( + llm_model_list=llm_router.get_model_list(), + llm_router=llm_router, + ) + + frame = json.dumps({"setup": {"model": addressing.format(model=VERTEX_LIVE_GROUP)}}) + denial = await gate(frame, UserAPIKeyAuth(token="hashed", models=["live-team"])) + + assert denial is None + + +async def _run_vertex_live_route(client_frames, valid_token): + """Drive the /vertex_ai/live route, so the gate under test is whichever one the route itself wires""" + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + vertex_ai_live_websocket_passthrough, + vertex_llm_base, + ) + + upstream_ws = RecordingUpstreamWebSocket() + websocket = _client_websocket( + AsyncMock(side_effect=[*client_frames, {"type": "websocket.disconnect"}]), + ) + + with ( + _patched_websocket_passthrough_environment(upstream_ws) as environment, + patch.object( # test-quality-ok: an OAuth exchange with Google, so this is the boundary being faked + vertex_llm_base, + "_ensure_access_token_async", + AsyncMock(return_value=("token-abc", "proj-db")), + ), + ): + await vertex_ai_live_websocket_passthrough( + websocket=websocket, + model="gemini-live-2.5-flash", + vertex_project="proj-db", + vertex_location="global", + user_api_key_dict=valid_token, + ) + + return websocket, upstream_ws, environment + + +@pytest.mark.asyncio +async def test_vertex_live_route_gates_the_setup_frame_even_when_the_query_model_is_permitted(): + """ + ?model= is all connect-time auth sees, so a key that passes it and then names a different model in its setup + frame reaches Vertex on the gateway's own credentials unless the route gates the frames as well + """ + websocket, upstream_ws, environment = await _run_vertex_live_route( + [_setup_frame("gemini-3-pro-live")], + UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), + ) + + upstream_ws.send.assert_not_awaited() + reason = _policy_close_reason(websocket) + assert reason is not None and "gemini-3-pro-live" in reason + refusal = environment.proxy_logging.post_call_failure_hook.await_args.kwargs["original_exception"] + assert refusal.code == "403" + assert refusal.type == ProxyErrorTypes.key_model_access_denied + + +@pytest.mark.asyncio +async def test_vertex_live_route_keeps_the_key_allowlist_out_of_the_close_reason_and_in_the_proxy_log(caplog): + """The client learns which model was refused, while only the proxy log names the models the key may reach""" + with caplog.at_level("WARNING", logger="LiteLLM Proxy"): + websocket, _, _ = await _run_vertex_live_route( + [_setup_frame("gemini-3-pro-live")], + UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), + ) + + reason = _policy_close_reason(websocket) + assert reason is not None and "gemini-live-2.5-flash" not in reason + assert any( + "gemini-3-pro-live" in record.getMessage() and "gemini-live-2.5-flash" in record.getMessage() + for record in caplog.records + ) + + +@pytest.mark.asyncio +async def test_vertex_live_route_forwards_a_setup_frame_the_key_is_allowed_to_run(): + websocket, upstream_ws, _ = await _run_vertex_live_route( + [_setup_frame("gemini-live-2.5-flash")], + UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), + ) + + upstream_ws.send.assert_awaited_once() + assert _policy_close_reason(websocket) is None + + +async def _run_passthrough_with_frame_gate(gate): + upstream_ws = ClosableUpstreamWebSocket() + websocket = _client_websocket( + AsyncMock( + side_effect=[ + {"type": "websocket.receive", "text": json.dumps({"setup": {"model": "gemini-live-2.5-flash"}})}, + {"type": "websocket.disconnect"}, + ] + ), + ) + + with _patched_websocket_passthrough_environment(upstream_ws) as environment: + await websocket_passthrough_request( + websocket=websocket, + target=VERTEX_LIVE_TARGET, + custom_headers={"Authorization": "Bearer token"}, + user_api_key_dict=UserAPIKeyAuth(token="hashed"), + forward_headers=False, + endpoint="/vertex_ai/live", + accept_websocket=False, + client_frame_gate=gate, + ) + + return websocket, upstream_ws, environment + + +@pytest.mark.asyncio +async def test_websocket_passthrough_reports_a_refused_client_frame_as_a_failure(): + """A denied session logged through the success hooks tells spend tracking and audit that it succeeded""" + denial = ProxyException( + message="gemini-live-2.5-flash: key not allowed to access model", + type=ProxyErrorTypes.key_model_access_denied, + param="model", + code=403, + ) + + async def refusing_gate(frame_data, valid_token): + return denial + + websocket, upstream_ws, environment = await _run_passthrough_with_frame_gate(refusing_gate) + + assert _policy_close_reason(websocket) == "gemini-live-2.5-flash: key not allowed to access model" + upstream_ws.close.assert_awaited() + environment.proxy_logging.post_call_success_hook.assert_not_awaited() + environment.logging_worker.ensure_initialized_and_enqueue.assert_not_called() + environment.proxy_logging.post_call_failure_hook.assert_awaited_once() + failure_kwargs = environment.proxy_logging.post_call_failure_hook.await_args.kwargs + assert failure_kwargs["original_exception"] is denial + assert failure_kwargs["request_data"]["litellm_logging_obj"].litellm_call_id + + +@pytest.mark.asyncio +async def test_websocket_passthrough_still_reports_an_allowed_session_as_a_success(): + async def permitting_gate(frame_data, valid_token): + return None + + _, _, environment = await _run_passthrough_with_frame_gate(permitting_gate) + + environment.proxy_logging.post_call_failure_hook.assert_not_awaited() + environment.proxy_logging.post_call_success_hook.assert_awaited_once() + environment.logging_worker.ensure_initialized_and_enqueue.assert_called_once() + + @pytest.mark.parametrize( "setup_model", ["gemini-live-2.5-flash", "models/gemini-live-2.5-flash", "publishers/google/models/gemini-live-2.5-flash"], From 408db64293a5425a2530bf34ceb3583ec1b5bb2a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 19:01:04 -0700 Subject: [PATCH 02/17] fix(proxy): enforce model access on the Vertex AI Live websocket passthrough The Live route takes its model from the client's setup frame after connect-time auth, so a key limited to one model could open a session on any Vertex Live model. Each setup frame's model is now checked against the key, team and project allowlists, and a refused frame closes the socket with 1008 before anything reaches Vertex. Vertex also reads setup frames that strict JSON cannot parse (single quotes, trailing commas, bare keys), so such a frame is refused instead of being forwarded unchecked. Co-authored-by: Marty Sullivan --- .../llm_passthrough_endpoints.py | 139 ++++++++++++++++-- .../pass_through_endpoints.py | 123 ++++++++++------ .../test_pass_through_endpoints.py | 36 ++++- 3 files changed, 237 insertions(+), 61 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 16fbd1de376..6afd8c53974 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -72,7 +72,8 @@ from litellm.llms.oss_decision import OssDecisionProvider, oss_connection, valid from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error from litellm.proxy._types import * -from litellm.proxy.auth.auth_checks import enforced_model_allowlists +from litellm.proxy.auth.auth_checks import can_key_call_resolved_model, enforced_model_allowlists +from litellm.proxy.auth.auth_utils import log_model_access_denial from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports @@ -110,6 +111,7 @@ from litellm.proxy.pass_through_endpoints.injected_stream_usage import ( ) from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( HttpPassThroughEndpointHelpers, + WebsocketClientFrameGate, create_pass_through_route, create_websocket_passthrough_route, websocket_passthrough_request, @@ -145,6 +147,7 @@ from .passthrough_endpoint_router import PassthroughEndpointRouter if TYPE_CHECKING: from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig from litellm.router import Router + from litellm.types.router import DeploymentTypedDict ProxyConfig = _ProxyConfig # rebind-ok: conditional type alias else: @@ -3803,6 +3806,12 @@ VERTEX_LIVE_UNCONFIGURED_CLOSE_REASON: Final = ( "Vertex AI auth failed: set a use_in_pass_through vertex model, default_vertex_config, or DEFAULT_VERTEXAI_* env" ) +VERTEX_LIVE_UNVERIFIED_MODEL_CLOSE_REASON: Final = ( + "Could not check your key's access to the model this setup frame names; ask a proxy admin to read the logs" +) + +VERTEX_LIVE_UNREADABLE_SETUP_CLOSE_REASON: Final = "Send setup frames as strict JSON so the proxy can check their model" + VERTEX_PUBLISHER_MODEL_PREFIX: Final = "publishers/google/models/" VERTEX_PUBLISHERS_SEGMENT: Final = "publishers/" @@ -3827,6 +3836,120 @@ def _get_llm_router() -> Router | None: return llm_router +@dataclass(frozen=True, slots=True) +class _UnreadableVertexLiveFrame: + pass + + +class _VertexLiveSetup(TypedDict, total=False): + model: ReadOnly[str] + + +class _VertexLiveClientFrame(TypedDict, total=False): + setup: ReadOnly[_VertexLiveSetup] + + +_VERTEX_LIVE_CLIENT_FRAME: Final = TypeAdapter(_VertexLiveClientFrame) + + +def _may_carry_vertex_live_setup(frame_data: str | bytes) -> bool: + """ + Whether a frame is worth parsing: audio and video frames stream continuously and are large, and one that + cannot spell ``setup`` cannot carry one. Vertex also reads single-quoted and bare keys, so the bare word is + the test, and a backslash means the key may be escaped, which only a parse settles + """ + if isinstance(frame_data, bytes): + return b"setup" in frame_data or b"\\" in frame_data + return "setup" in frame_data or "\\" in frame_data + + +def _vertex_live_setup_model(frame_data: str | bytes) -> str | _UnreadableVertexLiveFrame | None: + """ + The model a Vertex AI Live ``setup`` frame names, which is the only place BidiGenerateContent carries one. + + Vertex parses trailing commas, single quotes and bare keys that ``json.loads`` refuses, so a frame that may + carry a setup but fails a strict parse is unreadable rather than setup-free + """ + if not _may_carry_vertex_live_setup(frame_data): + return None + try: + frame: Final = _VERTEX_LIVE_CLIENT_FRAME.validate_python(json.loads(frame_data)) + except (json.JSONDecodeError, UnicodeDecodeError): + return _UnreadableVertexLiveFrame() + except ValidationError: + return None + setup: Final = frame.get("setup") + return None if setup is None else setup.get("model") or None + + +def _vertex_live_deployment(setup_model: str, llm_router: Router | None) -> DeploymentTypedDict | None: + """ + The router deployment whose group a ``setup`` frame's model names, whichever addressing the client used. + + The Live SDK wraps whatever the caller typed as ``models/``, and clients also send LiteLLM ids and + full Vertex resource paths, so the group is read from the last segment as well as from the whole string + """ + if llm_router is None: + return None + candidates: Final = (setup_model, setup_model.rsplit("/", 1)[-1]) + return next( + ( + deployment + for deployment in (llm_router.get_model_list() or ()) + if deployment.get("model_name") in candidates + ), + None, + ) + + +def _build_vertex_live_client_frame_gate(llm_router: Router | None) -> WebsocketClientFrameGate: + """ + Apply the key's, team's and project's model access to the model a client's ``setup`` frame names. + + ``?model=`` is optional on this route and the Live protocol carries the real model in the first client + frame, so connect-time auth has no model to check and every frame that names one has to be authorized + here instead. + + The allowlists and the access groups behind them are keyed by model group, and every addressing a client + sends carries that group inside a prefix, so the group is what gets authorized. A frame naming no group + is still checked as it arrived, which is the strictest reading available for a name nothing recognises + """ + + async def gate(frame_data: str | bytes, valid_token: UserAPIKeyAuth, /) -> ProxyException | None: + model: Final = _vertex_live_setup_model(frame_data) + if model is None: + return None + if isinstance(model, _UnreadableVertexLiveFrame): + return ProxyException( + message=VERTEX_LIVE_UNREADABLE_SETUP_CLOSE_REASON, + type=ProxyErrorTypes.bad_request_error, + param="setup", + code=400, + ) + deployment: Final = _vertex_live_deployment(model, llm_router) + try: + await can_key_call_resolved_model( + model=model if deployment is None else deployment["model_name"], + llm_model_list=None if llm_router is None else llm_router.get_model_list(), + valid_token=valid_token, + llm_router=llm_router, + ) + except ProxyException as denial: + log_model_access_denial(denial) + return denial + except Exception: # noqa: BLE001 # fail closed: a gate that cannot answer refuses, and says so in band + verbose_proxy_logger.exception("Vertex AI Live passthrough: model access check failed for %s", model) + return ProxyException( + message=VERTEX_LIVE_UNVERIFIED_MODEL_CLOSE_REASON, + type=ProxyErrorTypes.internal_server_error, + param="model", + code=500, + ) + return None + + return gate + + def _resolve_vertex_live_credentials( vertex_project: str | None, vertex_location: str | None, @@ -3885,17 +4008,8 @@ def _resolve_alias_to_upstream_model(setup_model: str, llm_router: Router | None """ 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"].get("model") - for deployment in (llm_router.get_model_list() or ()) - if deployment.get("model_name") in candidates - ), - None, - ) + deployment: Final = _vertex_live_deployment(setup_model, llm_router) + upstream: Final = None if deployment is None else deployment["litellm_params"].get("model") if upstream is None: return setup_model try: @@ -4009,6 +4123,7 @@ async def vertex_ai_live_websocket_passthrough( vertex_location=resolved_location, llm_router=_get_llm_router(), ), + client_frame_gate=_build_vertex_live_client_frame_gate(llm_router=_get_llm_router()), ) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index fe7ac30e434..6eda40e369f 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -10,7 +10,7 @@ from dataclasses import dataclass from datetime import datetime from itertools import chain, count, groupby from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast from urllib.parse import urlencode, urlparse import httpx @@ -30,7 +30,7 @@ from pydantic import TypeAdapter from starlette.datastructures import UploadFile as StarletteUploadFile from starlette.routing import BaseRoute, Route from starlette.websockets import WebSocketState -from websockets.asyncio.client import connect +from websockets.asyncio.client import ClientConnection, connect from websockets.exceptions import ( ConnectionClosedError, ConnectionClosedOK, @@ -2454,6 +2454,39 @@ def _with_trace_context(headers: Mapping[str, str], parent_span: object) -> dict return inject_trace_context(headers, parent_span=parent_span) +class WebsocketClientFrameGate(Protocol): + """Authorize one client frame, returning the denial to report or ``None`` to let it through. + + A websocket passthrough authenticates before any frame arrives, so a protocol that names its model + inside a frame has no model to authorize at connect time. The frame, not the extracted model, is the + argument because where a model sits is part of each provider's wire protocol + """ + + async def __call__(self, frame_data: str | bytes, valid_token: UserAPIKeyAuth, /) -> ProxyException | None: ... + + +async def _relay_client_frame( + websocket: WebSocket, + upstream_ws: ClientConnection, + checked_frame: str | bytes, + forwarded_frame: str | bytes, + gate: WebsocketClientFrameGate | None, + valid_token: UserAPIKeyAuth, +) -> ProxyException | None: + """Forward one client frame upstream, or refuse it and hand back the denial that ended the session. + + The denial travels out so the caller reports a failed session instead of leaving the success hooks to + record a refusal as a success, and tearing the upstream down is the caller's job: awaiting that close + here lets the upstream reader finish first, which cancels this task mid-close and loses the denial + """ + refusal: Final = None if gate is None else await gate(checked_frame, valid_token) + if refusal is None: + await upstream_ws.send(forwarded_frame) + return None + await websocket.close(code=CloseCode.POLICY_VIOLATION, reason=_truncated_close_reason(refusal.message)) + return refusal + + async def websocket_passthrough_request( websocket: WebSocket, target: str, @@ -2464,6 +2497,7 @@ async def websocket_passthrough_request( cost_per_request: float | None = None, accept_websocket: bool = True, setup_model_rewriter: Callable[[str], str] | None = None, + client_frame_gate: WebsocketClientFrameGate | None = None, ): """ WebSocket passthrough request handler. @@ -2477,6 +2511,7 @@ async def websocket_passthrough_request( endpoint: The endpoint path (for logging purposes) cost_per_request: Optional field - cost per request to the target endpoint setup_model_rewriter: Optional rewrite of the setup frame's model before it reaches the upstream + client_frame_gate: Optional per-frame authorization, applied before a frame reaches the upstream """ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.proxy_server import proxy_config, proxy_logging_obj @@ -2572,6 +2607,14 @@ async def websocket_passthrough_request( ) logging_obj.model_call_details["litellm_call_id"] = litellm_call_id + async def log_failure(exception: Exception, traceback_str: str | None = None) -> None: + await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=exception, + request_data={**kwargs, "litellm_logging_obj": logging_obj}, + traceback_str=traceback_str, + ) + # Pre-call logging logging_obj.pre_call( input=[{"role": "user", "content": "WebSocket connection"}], @@ -2603,8 +2646,8 @@ async def websocket_passthrough_request( "WebSocket passthrough (%s): Upstream connection established successfully", endpoint ) - async def forward_client_to_upstream() -> None: - """Forward messages from client to upstream WebSocket""" + async def forward_client_to_upstream() -> ProxyException | None: + """Forward messages from client to upstream WebSocket, returning any denial that ended it""" try: while True: message = await websocket.receive() @@ -2673,9 +2716,27 @@ async def websocket_passthrough_request( ) # Not a JSON message or doesn't contain setup data - await upstream_ws.send(_rewrite_vertex_live_setup_model(text_data, setup_model_rewriter)) + refusal = await _relay_client_frame( + websocket, + upstream_ws, + text_data, + _rewrite_vertex_live_setup_model(text_data, setup_model_rewriter), + client_frame_gate, + user_api_key_dict, + ) + if refusal is not None: + return refusal elif bytes_data is not None: - await upstream_ws.send(bytes_data) + refusal = await _relay_client_frame( + websocket, + upstream_ws, + bytes_data, + bytes_data, + client_frame_gate, + user_api_key_dict, + ) + if refusal is not None: + return refusal except asyncio.CancelledError: raise except Exception: @@ -2748,13 +2809,25 @@ async def websocket_passthrough_request( except asyncio.CancelledError: pass + task_results: Final = tuple(task.result() for task in done if task.exception() is None) + + # A refusal outranks any exception it caused: closing the client socket is what makes the + # upstream reader's next send fail, so raising that instead would report the wrong cause + client_frame_refusal: Final = next( + (result for result in task_results if isinstance(result, ProxyException)), None + ) + if client_frame_refusal is not None: + await upstream_ws.close() + await log_failure(client_frame_refusal) + return + # Check for exceptions in completed tasks for task in done: exception = task.exception() if exception is not None: raise exception - upstream_close: Final = _upstream_close_to_relay(task.result() for task in done) + upstream_close: Final = _upstream_close_to_relay(task_results) if upstream_close is not None and _client_socket_is_open(websocket): await websocket.close( code=upstream_close.code, @@ -2829,23 +2902,7 @@ async def websocket_passthrough_request( except InvalidStatus as exc: verbose_proxy_logger.exception("WebSocket passthrough (%s): upstream rejected WebSocket connection", endpoint) - # Prepare request payload for logging - request_payload = {} - if kwargs: - for key, value in kwargs.items(): - request_payload[key] = value - if logging_obj is not None: - request_payload["litellm_logging_obj"] = logging_obj - - # Log the connection failure using the same pattern as HTTP - await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=user_api_key_dict, - original_exception=exc, - request_data=request_payload, - traceback_str=traceback.format_exc( - limit=MAXIMUM_TRACEBACK_LINES_TO_LOG, - ), - ) + await log_failure(exc, traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG)) if _client_socket_is_open(websocket): await websocket.close( @@ -2857,23 +2914,7 @@ async def websocket_passthrough_request( "WebSocket passthrough (%s): unexpected error while proxying WebSocket", endpoint ) - # Prepare request payload for logging - request_payload = {} - if kwargs: - for key, value in kwargs.items(): - request_payload[key] = value - if logging_obj is not None: - request_payload["litellm_logging_obj"] = logging_obj - - # Log the unexpected error using the same pattern as HTTP - await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=user_api_key_dict, - original_exception=e, - request_data=request_payload, - traceback_str=traceback.format_exc( - limit=MAXIMUM_TRACEBACK_LINES_TO_LOG, - ), - ) + await log_failure(e, traceback.format_exc(limit=MAXIMUM_TRACEBACK_LINES_TO_LOG)) if _client_socket_is_open(websocket): await websocket.close(code=1011, reason="WebSocket passthrough error") diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 953aefe7f7c..1fac949fe46 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -7435,10 +7435,7 @@ async def _run_vertex_live_gated_passthrough(client_frames, valid_token): vertex_location="global", llm_router=llm_router, ), - client_frame_gate=_build_vertex_live_client_frame_gate( - llm_model_list=llm_router.get_model_list(), - llm_router=llm_router, - ), + client_frame_gate=_build_vertex_live_client_frame_gate(llm_router=llm_router), ) return websocket, upstream_ws @@ -7548,6 +7545,32 @@ async def test_vertex_live_setup_frame_with_an_escaped_key_is_still_gated(): assert _policy_close_reason(websocket) is not None +@pytest.mark.asyncio +@pytest.mark.parametrize("frame_kind", ["text", "bytes"]) +@pytest.mark.parametrize( + "lenient_setup", + [ + '{"setup": {"model": "totally-unlisted-model-abc",},}', + "{'setup': {'model': 'totally-unlisted-model-abc'}}", + '{setup: {model: "totally-unlisted-model-abc"}}', + ], +) +async def test_vertex_live_setup_frame_only_vertex_can_parse_is_refused(lenient_setup, frame_kind): + """Vertex reads trailing commas, single quotes and bare keys, which strict JSON refuses to parse""" + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [ + { + "type": "websocket.receive", + frame_kind: lenient_setup if frame_kind == "text" else lenient_setup.encode(), + } + ], + UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), + ) + + upstream_ws.send.assert_not_awaited() + assert _policy_close_reason(websocket) is not None + + @pytest.mark.asyncio async def test_vertex_live_non_setup_frames_are_forwarded_without_a_model_check(): websocket, upstream_ws = await _run_vertex_live_gated_passthrough( @@ -7663,10 +7686,7 @@ async def test_vertex_live_setup_frame_resolves_an_access_group_grant_in_every_a }, ] ) - gate = _build_vertex_live_client_frame_gate( - llm_model_list=llm_router.get_model_list(), - llm_router=llm_router, - ) + gate = _build_vertex_live_client_frame_gate(llm_router=llm_router) frame = json.dumps({"setup": {"model": addressing.format(model=VERTEX_LIVE_GROUP)}}) denial = await gate(frame, UserAPIKeyAuth(token="hashed", models=["live-team"])) From 42de0e6c893587f19c24cbe231e1a119549308d5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 21:06:39 -0700 Subject: [PATCH 03/17] fix(proxy): bill turns before a refused Vertex Live frame and authorize publisher ids Turns served before a refused setup frame are logged and billed as a success at the model they ran on, and the frame is relayed or refused before its model is recorded. A name no model group serves is checked by its publisher id, bare or vertex_ai/ prefixed, and an exact group match wins over the last path segment. Frames the proxy cannot read (lenient JSON, a repeated key, or the wrong shape) are refused only for callers with a model limit. The gate now also runs the user and agent model checks, and the close reason is cut at the last whole sentence that fits. --- litellm/proxy/auth/auth_checks.py | 39 +++ .../llm_passthrough_endpoints.py | 143 ++++++++--- .../pass_through_endpoints.py | 33 +-- .../test_pass_through_endpoints.py | 228 ++++++++++++++++-- 4 files changed, 375 insertions(+), 68 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 8887d86570a..adfb4bc14f7 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5677,6 +5677,45 @@ async def can_key_call_resolved_model( ) +async def can_caller_call_resolved_model( + model: str, + valid_token: UserAPIKeyAuth, + llm_router: litellm.Router | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, +) -> None: + """``can_key_call_resolved_model`` plus the agent grant and a personal key owner's own models, which + ``common_checks`` applies to a model named at auth time and a model named after auth needs as well""" + await can_key_call_resolved_model( + model=model, + llm_model_list=None if llm_router is None else llm_router.get_model_list(), + valid_token=valid_token, + llm_router=llm_router, + ) + await can_agent_call_model( + model=model, + valid_token=valid_token, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if valid_token.team_id is not None or prisma_client is None: + return + await can_user_call_model( + model=model, + llm_router=llm_router, + user_object=await _user_object_or_none( + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ), + key_model_aliases=key_model_aliases_for_auth_check(valid_token), + ) + + def can_org_access_model( model: str, org_object: LiteLLM_OrganizationTable | None, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 6afd8c53974..5acceb54ce3 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -72,7 +72,7 @@ from litellm.llms.oss_decision import OssDecisionProvider, oss_connection, valid from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error from litellm.proxy._types import * -from litellm.proxy.auth.auth_checks import can_key_call_resolved_model, enforced_model_allowlists +from litellm.proxy.auth.auth_checks import can_caller_call_resolved_model, enforced_model_allowlists from litellm.proxy.auth.auth_utils import log_model_access_denial from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.route_checks import RouteChecks @@ -3812,6 +3812,8 @@ VERTEX_LIVE_UNVERIFIED_MODEL_CLOSE_REASON: Final = ( VERTEX_LIVE_UNREADABLE_SETUP_CLOSE_REASON: Final = "Send setup frames as strict JSON so the proxy can check their model" +VERTEX_LIVE_NAME_NO_ALLOWLIST_HOLDS: Final = "litellm-probe/a-setup-frame-model-the-proxy-could-not-read" + VERTEX_PUBLISHER_MODEL_PREFIX: Final = "publishers/google/models/" VERTEX_PUBLISHERS_SEGMENT: Final = "publishers/" @@ -3863,80 +3865,146 @@ def _may_carry_vertex_live_setup(frame_data: str | bytes) -> bool: return "setup" in frame_data or "\\" in frame_data +def _object_with_unique_keys(pairs: Sequence[tuple[str, object]]) -> Mapping[str, object]: + unique: Final = MappingProxyType(dict(pairs)) + if len(unique) < len(pairs): + raise ValueError("a JSON object repeats a key") + return unique + + def _vertex_live_setup_model(frame_data: str | bytes) -> str | _UnreadableVertexLiveFrame | None: """ The model a Vertex AI Live ``setup`` frame names, which is the only place BidiGenerateContent carries one. - Vertex parses trailing commas, single quotes and bare keys that ``json.loads`` refuses, so a frame that may - carry a setup but fails a strict parse is unreadable rather than setup-free + Vertex parses trailing commas, single quotes and bare keys that ``json.loads`` refuses, and keeps the first of + two keys when the second is null, so a frame that may carry a setup but fails a strict parse, repeats a key, + or holds a setup of the wrong shape is unreadable rather than setup-free """ if not _may_carry_vertex_live_setup(frame_data): return None try: - frame: Final = _VERTEX_LIVE_CLIENT_FRAME.validate_python(json.loads(frame_data)) - except (json.JSONDecodeError, UnicodeDecodeError): + frame: Final = _VERTEX_LIVE_CLIENT_FRAME.validate_python( + json.loads(frame_data, object_pairs_hook=_object_with_unique_keys) + ) + except ValueError: return _UnreadableVertexLiveFrame() - except ValidationError: - return None setup: Final = frame.get("setup") return None if setup is None else setup.get("model") or None +def _deployment_named(model_name: str, deployments: Sequence[DeploymentTypedDict]) -> DeploymentTypedDict | None: + return next((deployment for deployment in deployments if deployment.get("model_name") == model_name), None) + + def _vertex_live_deployment(setup_model: str, llm_router: Router | None) -> DeploymentTypedDict | None: """ The router deployment whose group a ``setup`` frame's model names, whichever addressing the client used. The Live SDK wraps whatever the caller typed as ``models/``, and clients also send LiteLLM ids and - full Vertex resource paths, so the group is read from the last segment as well as from the whole string + full Vertex resource paths, so the group is read from the last segment when the whole string names none """ if llm_router is None: return None - candidates: Final = (setup_model, setup_model.rsplit("/", 1)[-1]) - return next( - ( - deployment - for deployment in (llm_router.get_model_list() or ()) - if deployment.get("model_name") in candidates - ), - None, + deployments: Final = tuple(llm_router.get_model_list() or ()) + exact: Final = _deployment_named(setup_model, deployments) + return exact if exact is not None else _deployment_named(setup_model.rsplit("/", 1)[-1], deployments) + + +def _vertex_live_model_names(setup_model: str, llm_router: Router | None) -> tuple[str, ...]: + """ + The names a key can hold for the model a ``setup`` frame runs: the router group the setup rewriter sends it + through, else the publisher model id Vertex runs, bare or under LiteLLM's ``vertex_ai/`` prefix. + + A full resource path skips the rewriter and reaches Vertex as sent, so only its publisher model id counts, and + a wildcard is matched against that id alone so addressing in front of it cannot satisfy the pattern + """ + deployment: Final = ( + None if setup_model.startswith("projects/") else _vertex_live_deployment(setup_model, llm_router) ) + if deployment is not None: + return (deployment["model_name"],) + publisher_model: Final = setup_model.rsplit("/", 1)[-1] + return (publisher_model, f"vertex_ai/{publisher_model}") + + +async def _model_access_denial( + model: str, valid_token: UserAPIKeyAuth, llm_router: Router | None +) -> ProxyException | None: + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + try: + await can_caller_call_resolved_model( + model=model, + valid_token=valid_token, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + except ProxyException as denial: + return denial + except HTTPException as refusal: + return ProxyException( + message=str(refusal.detail), + type=ProxyErrorTypes.auth_error, + param="model", + code=refusal.status_code, + ) + return None + + +async def _any_name_allowed_denial( + model_names: tuple[str, ...], valid_token: UserAPIKeyAuth, llm_router: Router | None +) -> ProxyException | None: + denials: Final = tuple([await _model_access_denial(name, valid_token, llm_router) for name in model_names]) + return None if None in denials else denials[0] + + +async def _vertex_live_setup_denial( + setup_model: str | _UnreadableVertexLiveFrame, valid_token: UserAPIKeyAuth, llm_router: Router | None +) -> ProxyException | None: + match setup_model: + case str(): + denial: Final = await _any_name_allowed_denial( + _vertex_live_model_names(setup_model, llm_router), valid_token, llm_router + ) + if denial is not None: + log_model_access_denial(denial) + return denial + case _UnreadableVertexLiveFrame(): + unlimited: Final = await _any_name_allowed_denial( + (VERTEX_LIVE_NAME_NO_ALLOWLIST_HOLDS,), valid_token, llm_router + ) + return ( + None + if unlimited is None + else ProxyException( + message=VERTEX_LIVE_UNREADABLE_SETUP_CLOSE_REASON, + type=ProxyErrorTypes.bad_request_error, + param="setup", + code=400, + ) + ) def _build_vertex_live_client_frame_gate(llm_router: Router | None) -> WebsocketClientFrameGate: """ - Apply the key's, team's and project's model access to the model a client's ``setup`` frame names. + Apply the checks a model named at auth time gets to the model a client's ``setup`` frame names. ``?model=`` is optional on this route and the Live protocol carries the real model in the first client frame, so connect-time auth has no model to check and every frame that names one has to be authorized here instead. - The allowlists and the access groups behind them are keyed by model group, and every addressing a client - sends carries that group inside a prefix, so the group is what gets authorized. A frame naming no group - is still checked as it arrived, which is the strictest reading available for a name nothing recognises + A frame the proxy cannot read the way Vertex does only reaches Vertex for a caller who may call any model, + which is a caller whose checks accept a name no allowlist holds """ async def gate(frame_data: str | bytes, valid_token: UserAPIKeyAuth, /) -> ProxyException | None: model: Final = _vertex_live_setup_model(frame_data) if model is None: return None - if isinstance(model, _UnreadableVertexLiveFrame): - return ProxyException( - message=VERTEX_LIVE_UNREADABLE_SETUP_CLOSE_REASON, - type=ProxyErrorTypes.bad_request_error, - param="setup", - code=400, - ) - deployment: Final = _vertex_live_deployment(model, llm_router) try: - await can_key_call_resolved_model( - model=model if deployment is None else deployment["model_name"], - llm_model_list=None if llm_router is None else llm_router.get_model_list(), - valid_token=valid_token, - llm_router=llm_router, - ) - except ProxyException as denial: - log_model_access_denial(denial) - return denial + return await _vertex_live_setup_denial(model, valid_token, llm_router) except Exception: # noqa: BLE001 # fail closed: a gate that cannot answer refuses, and says so in band verbose_proxy_logger.exception("Vertex AI Live passthrough: model access check failed for %s", model) return ProxyException( @@ -3945,7 +4013,6 @@ def _build_vertex_live_client_frame_gate(llm_router: Router | None) -> Websocket param="model", code=500, ) - return None return gate diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 6eda40e369f..9fdb0b74455 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2403,12 +2403,15 @@ def _json_object_frame(frame: str | bytes) -> dict[str, object] | None: def _truncated_close_reason(reason: str) -> str: """ - Fit a close reason inside the byte budget a WebSocket close frame allows, without splitting a character + Fit a close reason inside the byte budget a WebSocket close frame allows, ending on the last whole sentence + that fits when there is one and never splitting a character """ encoded: Final = reason.encode("utf-8") if len(encoded) <= WEBSOCKET_CLOSE_REASON_MAX_BYTES: return reason - return encoded[:WEBSOCKET_CLOSE_REASON_MAX_BYTES].decode("utf-8", errors="ignore") + clipped: Final = encoded[:WEBSOCKET_CLOSE_REASON_MAX_BYTES].decode("utf-8", errors="ignore") + sentence_end: Final = clipped.rfind(". ") + return clipped if sentence_end == -1 else clipped[: sentence_end + 1] SENDABLE_CLOSE_CODES: Final = frozenset(CloseCode) - frozenset( @@ -2660,6 +2663,16 @@ async def websocket_passthrough_request( bytes_data: bytes | None = message.get("bytes") if text_data is not None: + refusal = await _relay_client_frame( + websocket, + upstream_ws, + text_data, + _rewrite_vertex_live_setup_model(text_data, setup_model_rewriter), + client_frame_gate, + user_api_key_dict, + ) + if refusal is not None: + return refusal # Try to extract model from client setup message for Vertex AI Live if endpoint and "/vertex_ai/live" in endpoint: verbose_proxy_logger.debug( @@ -2716,16 +2729,6 @@ async def websocket_passthrough_request( ) # Not a JSON message or doesn't contain setup data - refusal = await _relay_client_frame( - websocket, - upstream_ws, - text_data, - _rewrite_vertex_live_setup_model(text_data, setup_model_rewriter), - client_frame_gate, - user_api_key_dict, - ) - if refusal is not None: - return refusal elif bytes_data is not None: refusal = await _relay_client_frame( websocket, @@ -2812,19 +2815,21 @@ async def websocket_passthrough_request( task_results: Final = tuple(task.result() for task in done if task.exception() is None) # A refusal outranks any exception it caused: closing the client socket is what makes the - # upstream reader's next send fail, so raising that instead would report the wrong cause + # upstream reader's next send fail, so raising that instead would report the wrong cause. + # Turns that completed before a refusal were served, so they are logged and billed as a success client_frame_refusal: Final = next( (result for result in task_results if isinstance(result, ProxyException)), None ) if client_frame_refusal is not None: await upstream_ws.close() + if client_frame_refusal is not None and not websocket_messages: await log_failure(client_frame_refusal) return # Check for exceptions in completed tasks for task in done: exception = task.exception() - if exception is not None: + if exception is not None and client_frame_refusal is None: raise exception upstream_close: Final = _upstream_close_to_relay(task_results) diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 1fac949fe46..c7163b22666 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -7240,10 +7240,16 @@ def _patched_websocket_passthrough_environment(upstream_ws): mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_success_hook = AsyncMock() mock_proxy_logging.post_call_failure_hook = AsyncMock() + enqueued_handler = MagicMock() mock_worker.ensure_initialized_and_enqueue = MagicMock( - side_effect=lambda async_coroutine: async_coroutine.close() + side_effect=lambda async_coroutine: ( + enqueued_handler(**async_coroutine.cr_frame.f_locals), + async_coroutine.close(), + ) + ) + yield SimpleNamespace( + proxy_logging=mock_proxy_logging, logging_worker=mock_worker, enqueued_handler=enqueued_handler ) - yield SimpleNamespace(proxy_logging=mock_proxy_logging, logging_worker=mock_worker) async def _pending_receive(): @@ -7408,20 +7414,27 @@ def _vertex_live_router(): ) -async def _run_vertex_live_gated_passthrough(client_frames, valid_token): +async def _run_vertex_live_gated_passthrough(client_frames, valid_token, llm_router=None): """Drive /vertex_ai/live with the real frame gate, returning the client and upstream sockets""" + upstream_ws = RecordingUpstreamWebSocket() + websocket, _ = await _run_vertex_live_gated_session( + AsyncMock(side_effect=[*client_frames, {"type": "websocket.disconnect"}]), + upstream_ws, + valid_token, + llm_router or _vertex_live_router(), + ) + return websocket, upstream_ws + + +async def _run_vertex_live_gated_session(receive, upstream_ws, valid_token, llm_router): from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( _build_vertex_live_client_frame_gate, _build_vertex_live_setup_model_rewriter, ) - llm_router = _vertex_live_router() - upstream_ws = RecordingUpstreamWebSocket() - websocket = _client_websocket( - AsyncMock(side_effect=[*client_frames, {"type": "websocket.disconnect"}]), - ) + websocket = _client_websocket(receive) - with _patched_websocket_passthrough_environment(upstream_ws): + with _patched_websocket_passthrough_environment(upstream_ws) as environment: await websocket_passthrough_request( websocket=websocket, target=VERTEX_LIVE_TARGET, @@ -7438,7 +7451,7 @@ async def _run_vertex_live_gated_passthrough(client_frames, valid_token): client_frame_gate=_build_vertex_live_client_frame_gate(llm_router=llm_router), ) - return websocket, upstream_ws + return websocket, environment def _setup_frame(model: str) -> dict: @@ -7469,6 +7482,7 @@ async def test_vertex_live_setup_frame_naming_a_model_the_key_cannot_call_is_ref upstream_ws.send.assert_not_awaited() reason = _policy_close_reason(websocket) assert reason is not None and "totally-unlisted-model-abc" in reason + assert reason.endswith(".") upstream_ws.close.assert_awaited() @@ -7571,6 +7585,69 @@ async def test_vertex_live_setup_frame_only_vertex_can_parse_is_refused(lenient_ assert _policy_close_reason(websocket) is not None +@pytest.mark.asyncio +@pytest.mark.parametrize( + "lenient_setup", + [ + '{"setup": {"model": "gemini-3-pro-live",},}', + "{'setup': {'model': 'gemini-3-pro-live'}}", + '{setup: {model: "gemini-3-pro-live"}}', + ], +) +async def test_vertex_live_setup_frame_only_vertex_can_parse_still_reaches_vertex_for_an_unrestricted_key( + lenient_setup, +): + """A key that may call any model loses nothing the proxy cannot read, so its lenient frames still go through""" + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [{"type": "websocket.receive", "text": lenient_setup}], + UserAPIKeyAuth(token="hashed", models=[]), + ) + + upstream_ws.send.assert_awaited_once() + assert _policy_close_reason(websocket) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "repeated_key_setup", + [ + '{"setup": {"model": "gemini-3-pro-live", "model": null}}', + '{"setup": {"model": "gemini-3-pro-live"}, "setup": null}', + ], +) +async def test_vertex_live_setup_frame_repeating_a_key_as_null_is_refused(repeated_key_setup): + """ + Vertex keeps the first of two keys when the second is null, while ``json.loads`` keeps the last, so the + proxy would read no model where Vertex runs the first one + """ + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [{"type": "websocket.receive", "text": repeated_key_setup}], + UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), + ) + + upstream_ws.send.assert_not_awaited() + assert _policy_close_reason(websocket) is not None + + +@pytest.mark.asyncio +async def test_vertex_live_client_content_that_mentions_setup_is_forwarded(): + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [ + _setup_frame("gemini-live-2.5-flash"), + { + "type": "websocket.receive", + "text": json.dumps( + {"clientContent": {"turns": [{"role": "user", "parts": [{"text": "help me setup a meeting"}]}]}} + ), + }, + ], + UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), + ) + + assert upstream_ws.send.await_count == 2 + assert _policy_close_reason(websocket) is None + + @pytest.mark.asyncio async def test_vertex_live_non_setup_frames_are_forwarded_without_a_model_check(): websocket, upstream_ws = await _run_vertex_live_gated_passthrough( @@ -7632,19 +7709,79 @@ async def test_vertex_live_setup_frame_naming_no_group_at_all_is_refused_in_ever @pytest.mark.asyncio -@pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS[1:]) -async def test_vertex_live_setup_frame_does_not_strip_a_prefix_off_a_model_no_group_serves(addressing): +@pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS) +@pytest.mark.parametrize("key_model", ["gemini-4-pro-live-unserved", "vertex_ai/gemini-4-pro-live-unserved"]) +async def test_vertex_live_setup_frame_authorizes_a_model_no_group_serves_by_its_publisher_id(addressing, key_model): """ - Stripping the addressing off a name the router does not serve would let a wildcard key reach an arbitrary - publisher model, and a full path would carry a project of the caller's choosing with it + With no group to resolve, the setup rewriter sends Vertex the last path segment, so a key holding that + publisher model, bare or as a LiteLLM ``vertex_ai/`` id, can open it whichever addressing its client sends """ websocket, upstream_ws = await _run_vertex_live_gated_passthrough( [_setup_frame(addressing.format(model="gemini-4-pro-live-unserved"))], - UserAPIKeyAuth(token="hashed", models=["gemini-*"]), + UserAPIKeyAuth(token="hashed", models=[key_model]), + ) + + upstream_ws.send.assert_awaited_once() + sent_setup = json.loads(upstream_ws.send.await_args.args[0])["setup"] + assert sent_setup["model"].endswith("/models/gemini-4-pro-live-unserved") + assert _policy_close_reason(websocket) is None + + +@pytest.mark.asyncio +async def test_vertex_live_setup_frame_addressing_cannot_satisfy_a_wildcard_for_another_publisher_model(): + """Vertex runs only the last segment, so the segments ahead of it must not be what matches the key's wildcard""" + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame("vertex_ai/gemini-anything/forbidden-model")], + UserAPIKeyAuth(token="hashed", models=["vertex_ai/gemini-*"]), ) upstream_ws.send.assert_not_awaited() - assert _policy_close_reason(websocket) is not None + reason = _policy_close_reason(websocket) + assert reason is not None and "forbidden-model" in reason + + +@pytest.mark.asyncio +async def test_vertex_live_setup_frame_prefers_the_group_its_whole_name_matches(): + """A group named ``vertex_ai/`` is the one the rewriter sends, so it is the one the key must hold""" + llm_router = litellm.Router( + model_list=[ + {"model_name": "gemini-3-pro-live", "litellm_params": {"model": "vertex_ai/gemini-3-pro-live"}}, + {"model_name": "vertex_ai/gemini-3-pro-live", "litellm_params": {"model": "vertex_ai/gemini-3-pro-live"}}, + ] + ) + + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame("vertex_ai/gemini-3-pro-live")], + UserAPIKeyAuth(token="hashed", models=["vertex_ai/gemini-3-pro-live"]), + llm_router=llm_router, + ) + + upstream_ws.send.assert_awaited_once() + assert _policy_close_reason(websocket) is None + + +@pytest.mark.asyncio +async def test_vertex_live_setup_frame_is_checked_against_the_owner_of_a_personal_key(): + """A personal key's owner can be limited to fewer models than the key lists, and auth enforces that limit too""" + from litellm.proxy.proxy_server import user_api_key_cache + + user_id = f"live-user-{uuid.uuid4().hex}" + await user_api_key_cache.async_set_cache( + key=user_id, + value={"user_id": user_id, "user_email": "live@example.com", "models": ["gemini-live-2.5-flash"]}, + ) + + with patch( # test-quality-ok: read at call time, and the user row comes from the cache seeded above + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ): + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame("gemini-3-pro-live")], + UserAPIKeyAuth(token="hashed", user_id=user_id, models=[]), + ) + + upstream_ws.send.assert_not_awaited() + reason = _policy_close_reason(websocket) + assert reason is not None and "gemini-3-pro-live" in reason @pytest.mark.asyncio @@ -7772,6 +7909,65 @@ async def test_vertex_live_route_forwards_a_setup_frame_the_key_is_allowed_to_ru assert _policy_close_reason(websocket) is None +class ScriptedUpstreamWebSocket: + """An upstream that sends the given frames in order, then stays open with nothing more to say""" + + def __init__(self, frames): + self._frames = iter(frames) + self.sent_every_frame = asyncio.Event() + self.send = AsyncMock() + self.close = AsyncMock() + + async def recv(self, decode: bool = True): + frame = next(self._frames, None) + if frame is not None: + return frame + self.sent_every_frame.set() + await asyncio.Event().wait() + + def __aiter__(self): + return self + + async def __anext__(self): + raise StopAsyncIteration + + +@pytest.mark.asyncio +async def test_vertex_live_turns_served_before_a_refused_frame_are_still_billed(): + """ + A refusal used to log the whole session as a failure, which records no spend, so a key could run allowed + turns and then send a refused frame to have them go unbilled. The refused frame's model must not become + the model those turns are priced at either + """ + upstream_ws = ScriptedUpstreamWebSocket( + ( + json.dumps({"setupComplete": {}}), + json.dumps({"serverContent": {"turnComplete": True}, "usageMetadata": {"totalTokenCount": 60}}), + ) + ) + + async def client_frames(): + yield _setup_frame("gemini-live-2.5-flash") + await upstream_ws.sent_every_frame.wait() + yield _setup_frame("gemini-3-pro-live") + + websocket, environment = await _run_vertex_live_gated_session( + client_frames().__anext__, + upstream_ws, + UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), + _vertex_live_router(), + ) + + reason = _policy_close_reason(websocket) + assert reason is not None and "gemini-3-pro-live" in reason + upstream_ws.send.assert_awaited_once() + upstream_ws.close.assert_awaited() + environment.logging_worker.ensure_initialized_and_enqueue.assert_called_once() + assert environment.enqueued_handler.call_args.kwargs["logging_obj"].model == "gemini-live-2.5-flash" + environment.proxy_logging.post_call_success_hook.assert_awaited_once() + environment.proxy_logging.post_call_failure_hook.assert_not_awaited() + + async def _run_passthrough_with_frame_gate(gate): upstream_ws = ClosableUpstreamWebSocket() websocket = _client_websocket( From 3ee80f6acabe3c1bc75a58c5486de449a43a4ee0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 22:33:40 -0700 Subject: [PATCH 04/17] fix(proxy): check the Vertex Live model the setup frame forwards, not an alias of it A key, team, router or global alias name no longer grants a Live setup frame, since the frame reaches Vertex under the name sent. A full projects/ path is allowed for any group serving that publisher model. Frames the proxy cannot read are refused for any caller with a model limit, and upstream close reasons keep the byte cut. --- .../llm_passthrough_endpoints.py | 115 +++++++--- .../pass_through_endpoints.py | 58 +++-- .../test_pass_through_endpoints.py | 217 +++++++++++++++--- 3 files changed, 311 insertions(+), 79 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 5acceb54ce3..f57d1a1af63 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -72,9 +72,16 @@ from litellm.llms.oss_decision import OssDecisionProvider, oss_connection, valid from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error from litellm.proxy._types import * -from litellm.proxy.auth.auth_checks import can_caller_call_resolved_model, enforced_model_allowlists +from litellm.proxy.auth.auth_checks import ( + can_caller_call_resolved_model, + enforced_model_allowlists, + key_model_aliases_for_auth_check, + personal_key_team_model_access_applies, + team_model_aliases_for_auth_check, +) from litellm.proxy.auth.auth_utils import log_model_access_denial from litellm.proxy.auth.handle_jwt import JWTHandler +from litellm.proxy.auth.model_access_denied import model_access_denied_client_message from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports _get_bearer_token, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -3812,8 +3819,6 @@ VERTEX_LIVE_UNVERIFIED_MODEL_CLOSE_REASON: Final = ( VERTEX_LIVE_UNREADABLE_SETUP_CLOSE_REASON: Final = "Send setup frames as strict JSON so the proxy can check their model" -VERTEX_LIVE_NAME_NO_ALLOWLIST_HOLDS: Final = "litellm-probe/a-setup-frame-model-the-proxy-could-not-read" - VERTEX_PUBLISHER_MODEL_PREFIX: Final = "publishers/google/models/" VERTEX_PUBLISHERS_SEGMENT: Final = "publishers/" @@ -3910,12 +3915,24 @@ def _vertex_live_deployment(setup_model: str, llm_router: Router | None) -> Depl return exact if exact is not None else _deployment_named(setup_model.rsplit("/", 1)[-1], deployments) +def _deployment_upstream_model(deployment: DeploymentTypedDict) -> str | None: + upstream: Final = deployment["litellm_params"].get("model") + if upstream is None: + return None + try: + _, provider, _, _ = litellm.get_llm_provider(model=upstream) + except litellm.exceptions.BadRequestError: + return upstream + return upstream.removeprefix(f"{provider}/") + + def _vertex_live_model_names(setup_model: str, llm_router: Router | None) -> tuple[str, ...]: """ The names a key can hold for the model a ``setup`` frame runs: the router group the setup rewriter sends it - through, else the publisher model id Vertex runs, bare or under LiteLLM's ``vertex_ai/`` prefix. + through, else every group serving the publisher model Vertex runs, and that model's id bare or under LiteLLM's + ``vertex_ai/`` prefix unless a group holds that name for another model. - A full resource path skips the rewriter and reaches Vertex as sent, so only its publisher model id counts, and + A full resource path skips the rewriter and reaches Vertex as sent, so only its publisher model counts, and a wildcard is matched against that id alone so addressing in front of it cannot satisfy the pattern """ deployment: Final = ( @@ -3923,8 +3940,30 @@ def _vertex_live_model_names(setup_model: str, llm_router: Router | None) -> tup ) if deployment is not None: return (deployment["model_name"],) - publisher_model: Final = setup_model.rsplit("/", 1)[-1] - return (publisher_model, f"vertex_ai/{publisher_model}") + deployments: Final = () if llm_router is None else tuple(llm_router.get_model_list() or ()) + publisher_model: Final = _vertex_publisher_model_suffix(setup_model) + serving_groups: Final = tuple( + dict.fromkeys( + candidate["model_name"] + for candidate in deployments + if _vertex_publisher_model_suffix(_deployment_upstream_model(candidate) or "") == publisher_model + ) + ) + group_names: Final = frozenset(candidate["model_name"] for candidate in deployments) + publisher_model_id: Final = setup_model.rsplit("/", 1)[-1] + return ( + *serving_groups, + *(name for name in (publisher_model_id, f"vertex_ai/{publisher_model_id}") if name not in group_names), + ) + + +def _is_model_alias(name: str, valid_token: UserAPIKeyAuth, llm_router: Router | None) -> bool: + return ( + name in (key_model_aliases_for_auth_check(valid_token) or ()) + or name in (team_model_aliases_for_auth_check(valid_token) or ()) + or name in litellm.model_alias_map + or (llm_router is not None and name in llm_router.model_group_alias) + ) async def _model_access_denial( @@ -3953,11 +3992,42 @@ async def _model_access_denial( return None -async def _any_name_allowed_denial( - model_names: tuple[str, ...], valid_token: UserAPIKeyAuth, llm_router: Router | None +async def _vertex_live_model_denial( + setup_model: str, valid_token: UserAPIKeyAuth, llm_router: Router | None ) -> ProxyException | None: + """ + The setup rewriter forwards a name as sent, so a name an alias maps to another model grants nothing here + """ + model_names: Final = tuple( + name + for name in _vertex_live_model_names(setup_model, llm_router) + if not _is_model_alias(name, valid_token, llm_router) + ) denials: Final = tuple([await _model_access_denial(name, valid_token, llm_router) for name in model_names]) - return None if None in denials else denials[0] + if None in denials: + return None + return ( + denials[0] + if denials + else ModelAccessDeniedProxyException( + message=model_access_denied_client_message(model=setup_model), + internal_message=f"Vertex AI Live setup model {setup_model} is only reachable through a model alias", + type=ProxyErrorTypes.key_model_access_denied, + param="model", + code=403, + ) + ) + + +async def _caller_has_model_limits(valid_token: UserAPIKeyAuth) -> bool: + return ( + valid_token.agent_id is not None + or valid_token.managed_agent_policy is not None + or valid_token.agent_caller is not None + or valid_token.end_user_id is not None + or personal_key_team_model_access_applies(valid_token) + or _has_model_restrictions(await _proxy_model_allowlists()(valid_token)) + ) async def _vertex_live_setup_denial( @@ -3965,25 +4035,20 @@ async def _vertex_live_setup_denial( ) -> ProxyException | None: match setup_model: case str(): - denial: Final = await _any_name_allowed_denial( - _vertex_live_model_names(setup_model, llm_router), valid_token, llm_router - ) + denial: Final = await _vertex_live_model_denial(setup_model, valid_token, llm_router) if denial is not None: log_model_access_denial(denial) return denial case _UnreadableVertexLiveFrame(): - unlimited: Final = await _any_name_allowed_denial( - (VERTEX_LIVE_NAME_NO_ALLOWLIST_HOLDS,), valid_token, llm_router - ) return ( - None - if unlimited is None - else ProxyException( + ProxyException( message=VERTEX_LIVE_UNREADABLE_SETUP_CLOSE_REASON, type=ProxyErrorTypes.bad_request_error, param="setup", code=400, ) + if await _caller_has_model_limits(valid_token) + else None ) @@ -3996,7 +4061,7 @@ def _build_vertex_live_client_frame_gate(llm_router: Router | None) -> Websocket here instead. A frame the proxy cannot read the way Vertex does only reaches Vertex for a caller who may call any model, - which is a caller whose checks accept a name no allowlist holds + which is a caller with no model allowlist and no agent, end user or team grant standing in for one """ async def gate(frame_data: str | bytes, valid_token: UserAPIKeyAuth, /) -> ProxyException | None: @@ -4076,14 +4141,8 @@ def _resolve_alias_to_upstream_model(setup_model: str, llm_router: Router | None The Live SDK wraps whatever the caller typed as ``models/``, so a gateway alias arrives prefixed """ deployment: Final = _vertex_live_deployment(setup_model, llm_router) - upstream: Final = None if deployment is None else deployment["litellm_params"].get("model") - if upstream is None: - return setup_model - try: - _, provider, _, _ = litellm.get_llm_provider(model=upstream) - except litellm.exceptions.BadRequestError: - return upstream - return upstream.removeprefix(f"{provider}/") + upstream: Final = None if deployment is None else _deployment_upstream_model(deployment) + return setup_model if upstream is None else upstream async def vertex_ai_live_websocket_passthrough( diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 9fdb0b74455..c8949dc3a01 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2403,15 +2403,21 @@ def _json_object_frame(frame: str | bytes) -> dict[str, object] | None: def _truncated_close_reason(reason: str) -> str: """ - Fit a close reason inside the byte budget a WebSocket close frame allows, ending on the last whole sentence - that fits when there is one and never splitting a character + Fit a close reason inside the byte budget a WebSocket close frame allows, without splitting a character """ encoded: Final = reason.encode("utf-8") if len(encoded) <= WEBSOCKET_CLOSE_REASON_MAX_BYTES: return reason - clipped: Final = encoded[:WEBSOCKET_CLOSE_REASON_MAX_BYTES].decode("utf-8", errors="ignore") + return encoded[:WEBSOCKET_CLOSE_REASON_MAX_BYTES].decode("utf-8", errors="ignore") + + +def _refusal_close_reason(message: str) -> str: + """ + The proxy's own refusal message as a close reason, ending on the last whole sentence that fits + """ + clipped: Final = _truncated_close_reason(message) sentence_end: Final = clipped.rfind(". ") - return clipped if sentence_end == -1 else clipped[: sentence_end + 1] + return clipped if clipped == message or sentence_end == -1 else clipped[: sentence_end + 1] SENDABLE_CLOSE_CODES: Final = frozenset(CloseCode) - frozenset( @@ -2486,7 +2492,7 @@ async def _relay_client_frame( if refusal is None: await upstream_ws.send(forwarded_frame) return None - await websocket.close(code=CloseCode.POLICY_VIOLATION, reason=_truncated_close_reason(refusal.message)) + await websocket.close(code=CloseCode.POLICY_VIOLATION, reason=_refusal_close_reason(refusal.message)) return refusal @@ -2663,16 +2669,17 @@ async def websocket_passthrough_request( bytes_data: bytes | None = message.get("bytes") if text_data is not None: - refusal = await _relay_client_frame( - websocket, - upstream_ws, - text_data, - _rewrite_vertex_live_setup_model(text_data, setup_model_rewriter), - client_frame_gate, - user_api_key_dict, - ) - if refusal is not None: - return refusal + if ( + text_refusal := await _relay_client_frame( + websocket, + upstream_ws, + text_data, + _rewrite_vertex_live_setup_model(text_data, setup_model_rewriter), + client_frame_gate, + user_api_key_dict, + ) + ) is not None: + return text_refusal # Try to extract model from client setup message for Vertex AI Live if endpoint and "/vertex_ai/live" in endpoint: verbose_proxy_logger.debug( @@ -2730,16 +2737,17 @@ async def websocket_passthrough_request( # Not a JSON message or doesn't contain setup data elif bytes_data is not None: - refusal = await _relay_client_frame( - websocket, - upstream_ws, - bytes_data, - bytes_data, - client_frame_gate, - user_api_key_dict, - ) - if refusal is not None: - return refusal + if ( + bytes_refusal := await _relay_client_frame( + websocket, + upstream_ws, + bytes_data, + bytes_data, + client_frame_gate, + user_api_key_dict, + ) + ) is not None: + return bytes_refusal except asyncio.CancelledError: raise except Exception: diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index c7163b22666..e95a132eae2 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -7,7 +7,7 @@ import os import sys import uuid import zlib -from collections.abc import AsyncIterator, Callable, Mapping +from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Mapping, Sequence from contextlib import ExitStack, contextmanager from dataclasses import dataclass from datetime import datetime @@ -30,7 +30,7 @@ from starlette.datastructures import UploadFile as StarletteUploadFile import litellm from litellm._logging import verbose_proxy_logger -from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS +from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS, WEBSOCKET_CLOSE_REASON_MAX_BYTES from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -53,6 +53,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( initialize_pass_through_endpoints, pass_through_request, resolve_llm_passthrough_timeout, + WebsocketClientFrameGate, resolve_pass_through_request_timeout, set_env_variables_in_header, websocket_passthrough_request, @@ -7184,27 +7185,27 @@ class ClosableUpstreamWebSocket: CLOSE_HANDSHAKE_YIELDS = 5 - def __init__(self): + def __init__(self) -> None: self._closing = asyncio.Event() self.send = AsyncMock() self.close = AsyncMock(side_effect=self._close) - async def _close(self): + async def _close(self) -> None: self._closing.set() for _ in range(self.CLOSE_HANDSHAKE_YIELDS): await asyncio.sleep(0) - async def recv(self, decode: bool = True): + async def recv(self, decode: bool = True) -> str: from websockets.exceptions import ConnectionClosedOK from websockets.frames import Close await self._closing.wait() raise ConnectionClosedOK(rcvd=Close(1000, ""), sent=Close(1000, ""), rcvd_then_sent=True) - def __aiter__(self): + def __aiter__(self) -> "ClosableUpstreamWebSocket": return self - async def __anext__(self): + async def __anext__(self) -> str: raise StopAsyncIteration @@ -7285,6 +7286,36 @@ async def test_websocket_passthrough_relays_upstream_policy_close_to_client(): websocket.close.assert_awaited_once_with(code=1008, reason=upstream_reason) +@pytest.mark.asyncio +async def test_websocket_passthrough_relays_as_much_of_a_long_upstream_close_reason_as_fits(): + from websockets.exceptions import ConnectionClosedError + from websockets.frames import Close + + upstream_reason = ( + "Quota exceeded for concurrent_sessions. Request more quota for aiplatform.googleapis.com " + "BidiGenerateContent sessions in us-central1 from the Google Cloud console quotas page" + ) + upstream_ws = ClosingUpstreamWebSocket( + ConnectionClosedError(rcvd=Close(1011, upstream_reason), sent=Close(1011, ""), rcvd_then_sent=True) + ) + websocket = _client_websocket(_pending_receive) + + with _patched_websocket_passthrough_environment(upstream_ws): + 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, + ) + + relayed_reason = websocket.close.await_args.kwargs["reason"] + assert upstream_reason.startswith(relayed_reason) + assert len(relayed_reason.encode("utf-8")) == WEBSOCKET_CLOSE_REASON_MAX_BYTES + + @pytest.mark.asyncio async def test_websocket_passthrough_keeps_normal_upstream_close_normal(): from websockets.exceptions import ConnectionClosedOK @@ -7382,16 +7413,19 @@ VERTEX_LIVE_GROUP = "gemini-live-2.5-flash-native-audio" VERTEX_LIVE_ALIAS_GROUP = "gemini-live-native-audio" -VERTEX_LIVE_CLIENT_ADDRESSINGS = ( +VERTEX_LIVE_GROUP_ADDRESSINGS = ( "{model}", "models/{model}", "vertex_ai/{model}", "publishers/google/models/{model}", - "projects/proj-db/locations/global/publishers/google/models/{model}", ) +VERTEX_LIVE_FULL_PATH = "projects/proj-db/locations/global/publishers/google/models/{model}" -def _vertex_live_router(): +VERTEX_LIVE_CLIENT_ADDRESSINGS = (*VERTEX_LIVE_GROUP_ADDRESSINGS, VERTEX_LIVE_FULL_PATH) + + +def _vertex_live_router() -> litellm.Router: return litellm.Router( model_list=[ { @@ -7414,7 +7448,11 @@ def _vertex_live_router(): ) -async def _run_vertex_live_gated_passthrough(client_frames, valid_token, llm_router=None): +async def _run_vertex_live_gated_passthrough( + client_frames: Sequence[Mapping[str, str | bytes]], + valid_token: UserAPIKeyAuth, + llm_router: litellm.Router | None = None, +) -> tuple[MagicMock, RecordingUpstreamWebSocket]: """Drive /vertex_ai/live with the real frame gate, returning the client and upstream sockets""" upstream_ws = RecordingUpstreamWebSocket() websocket, _ = await _run_vertex_live_gated_session( @@ -7426,7 +7464,12 @@ async def _run_vertex_live_gated_passthrough(client_frames, valid_token, llm_rou return websocket, upstream_ws -async def _run_vertex_live_gated_session(receive, upstream_ws, valid_token, llm_router): +async def _run_vertex_live_gated_session( + receive: Callable[[], Awaitable[Mapping[str, str | bytes]]], + upstream_ws: "RecordingUpstreamWebSocket | ScriptedUpstreamWebSocket", + valid_token: UserAPIKeyAuth, + llm_router: litellm.Router, +) -> tuple[MagicMock, SimpleNamespace]: from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( _build_vertex_live_client_frame_gate, _build_vertex_live_setup_model_rewriter, @@ -7454,14 +7497,14 @@ async def _run_vertex_live_gated_session(receive, upstream_ws, valid_token, llm_ return websocket, environment -def _setup_frame(model: str) -> dict: +def _setup_frame(model: str) -> dict[str, str]: return { "type": "websocket.receive", "text": json.dumps({"setup": {"model": model, "generationConfig": {"responseModalities": ["TEXT"]}}}), } -def _policy_close_reason(websocket) -> str | None: +def _policy_close_reason(websocket: MagicMock) -> str | None: return next( (call.kwargs.get("reason") for call in websocket.close.await_args_list if call.kwargs.get("code") == 1008), None, @@ -7629,6 +7672,37 @@ async def test_vertex_live_setup_frame_repeating_a_key_as_null_is_refused(repeat assert _policy_close_reason(websocket) is not None +@pytest.mark.asyncio +@pytest.mark.parametrize( + "limited_token", + [ + UserAPIKeyAuth(token="hashed", models=["litellm-*"]), + UserAPIKeyAuth(token="hashed", models=["vertex_ai/gemini-*"]), + UserAPIKeyAuth(token="hashed", models=[], team_id="team-1", team_models=["gemini-live-2.5-flash"]), + UserAPIKeyAuth(token="hashed", models=[], agent_id="agent-1"), + UserAPIKeyAuth(token="hashed", models=[], end_user_id="end-user-1"), + ], + ids=["litellm-wildcard", "vertex-wildcard", "team", "agent", "end-user"], +) +async def test_vertex_live_setup_frame_only_vertex_can_parse_is_refused_for_any_limited_caller(limited_token): + """A wildcard can match any name the proxy might test the caller with, so only a caller with no model limit + anywhere may send a frame the proxy cannot read""" + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [ + { + "type": "websocket.receive", + "text": "{'setup': {'model': '" + + VERTEX_LIVE_FULL_PATH.format(model="gemini-3-pro-live") + + "',},}", + } + ], + limited_token, + ) + + upstream_ws.send.assert_not_awaited() + assert _policy_close_reason(websocket) is not None + + @pytest.mark.asyncio async def test_vertex_live_client_content_that_mentions_setup_is_forwarded(): websocket, upstream_ws = await _run_vertex_live_gated_passthrough( @@ -7663,12 +7737,12 @@ async def test_vertex_live_non_setup_frames_are_forwarded_without_a_model_check( @pytest.mark.asyncio -@pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS) +@pytest.mark.parametrize("addressing", VERTEX_LIVE_GROUP_ADDRESSINGS) @pytest.mark.parametrize("group", [VERTEX_LIVE_GROUP, VERTEX_LIVE_ALIAS_GROUP]) async def test_vertex_live_setup_frame_authorizes_a_permitted_group_in_every_addressing(addressing, group): """ - The Live SDK wraps the caller's model as ``models/`` and the docs tell callers to send the full - resource path, so a key holding the group has to work whichever addressing its client actually sends + The Live SDK wraps the caller's model as ``models/``, so a key holding the group has to work + whichever addressing its client actually sends """ websocket, upstream_ws = await _run_vertex_live_gated_passthrough( [_setup_frame(addressing.format(model=group))], @@ -7760,6 +7834,93 @@ async def test_vertex_live_setup_frame_prefers_the_group_its_whole_name_matches( assert _policy_close_reason(websocket) is None +@pytest.mark.asyncio +@pytest.mark.parametrize("group", [VERTEX_LIVE_GROUP, VERTEX_LIVE_ALIAS_GROUP]) +async def test_vertex_live_full_path_is_authorized_by_any_group_serving_its_publisher_model(group): + """The docs tell callers to send the full resource path, which Vertex runs as sent, so every group serving + that publisher model grants it, whatever the group is called""" + full_path = VERTEX_LIVE_FULL_PATH.format(model=VERTEX_LIVE_GROUP) + + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame(full_path)], + UserAPIKeyAuth(token="hashed", models=[group]), + ) + + upstream_ws.send.assert_awaited_once() + assert json.loads(upstream_ws.send.await_args.args[0])["setup"]["model"] == full_path + assert _policy_close_reason(websocket) is None + + +@pytest.mark.asyncio +async def test_vertex_live_full_path_is_not_granted_by_a_group_named_after_it_that_serves_another_model(): + llm_router = litellm.Router( + model_list=[ + {"model_name": "gemini-3-pro-live", "litellm_params": {"model": "vertex_ai/gemini-live-2.5-flash"}}, + ] + ) + + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame(VERTEX_LIVE_FULL_PATH.format(model="gemini-3-pro-live"))], + UserAPIKeyAuth(token="hashed", models=["gemini-3-pro-live"]), + llm_router=llm_router, + ) + + upstream_ws.send.assert_not_awaited() + reason = _policy_close_reason(websocket) + assert reason is not None and "gemini-3-pro-live" in reason + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "valid_token", + [ + UserAPIKeyAuth( + token="hashed", + models=["gemini-live-2.5-flash"], + aliases={"gemini-3-pro-live": "gemini-live-2.5-flash"}, + ), + UserAPIKeyAuth( + token="hashed", + models=["all-team-models"], + team_id="team-1", + team_models=["gemini-live-2.5-flash"], + team_model_aliases={"gemini-3-pro-live": "gemini-live-2.5-flash"}, + ), + ], + ids=["key-alias", "team-alias"], +) +async def test_vertex_live_setup_frame_naming_an_alias_cannot_open_the_model_the_alias_shadows(valid_token): + """ + An alias passes the access check by its target, while the setup rewriter sends the name as written, so a + key that could set an alias would open the model the alias name belongs to + """ + websocket, upstream_ws = await _run_vertex_live_gated_passthrough([_setup_frame("gemini-3-pro-live")], valid_token) + + upstream_ws.send.assert_not_awaited() + reason = _policy_close_reason(websocket) + assert reason is not None and "gemini-3-pro-live" in reason + + +@pytest.mark.asyncio +async def test_vertex_live_setup_frame_naming_a_router_alias_cannot_open_the_model_the_alias_shadows(): + llm_router = litellm.Router( + model_list=[ + {"model_name": "gemini-live-2.5-flash", "litellm_params": {"model": "vertex_ai/gemini-live-2.5-flash"}}, + ], + model_group_alias={"gemini-3-pro-live": "gemini-live-2.5-flash"}, + ) + + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame("gemini-3-pro-live")], + UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), + llm_router=llm_router, + ) + + upstream_ws.send.assert_not_awaited() + reason = _policy_close_reason(websocket) + assert reason is not None and "gemini-3-pro-live" in reason + + @pytest.mark.asyncio async def test_vertex_live_setup_frame_is_checked_against_the_owner_of_a_personal_key(): """A personal key's owner can be limited to fewer models than the key lists, and auth enforces that limit too""" @@ -7831,7 +7992,9 @@ async def test_vertex_live_setup_frame_resolves_an_access_group_grant_in_every_a assert denial is None -async def _run_vertex_live_route(client_frames, valid_token): +async def _run_vertex_live_route( + client_frames: Sequence[Mapping[str, str | bytes]], valid_token: UserAPIKeyAuth +) -> tuple[MagicMock, RecordingUpstreamWebSocket, SimpleNamespace]: """Drive the /vertex_ai/live route, so the gate under test is whichever one the route itself wires""" from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( vertex_ai_live_websocket_passthrough, @@ -7912,23 +8075,23 @@ async def test_vertex_live_route_forwards_a_setup_frame_the_key_is_allowed_to_ru class ScriptedUpstreamWebSocket: """An upstream that sends the given frames in order, then stays open with nothing more to say""" - def __init__(self, frames): + def __init__(self, frames: Iterable[str]) -> None: self._frames = iter(frames) self.sent_every_frame = asyncio.Event() self.send = AsyncMock() self.close = AsyncMock() - async def recv(self, decode: bool = True): + async def recv(self, decode: bool = True) -> str | None: frame = next(self._frames, None) if frame is not None: return frame self.sent_every_frame.set() await asyncio.Event().wait() - def __aiter__(self): + def __aiter__(self) -> "ScriptedUpstreamWebSocket": return self - async def __anext__(self): + async def __anext__(self) -> str: raise StopAsyncIteration @@ -7946,7 +8109,7 @@ async def test_vertex_live_turns_served_before_a_refused_frame_are_still_billed( ) ) - async def client_frames(): + async def client_frames() -> AsyncIterator[dict[str, str]]: yield _setup_frame("gemini-live-2.5-flash") await upstream_ws.sent_every_frame.wait() yield _setup_frame("gemini-3-pro-live") @@ -7968,7 +8131,9 @@ async def test_vertex_live_turns_served_before_a_refused_frame_are_still_billed( environment.proxy_logging.post_call_failure_hook.assert_not_awaited() -async def _run_passthrough_with_frame_gate(gate): +async def _run_passthrough_with_frame_gate( + gate: WebsocketClientFrameGate, +) -> tuple[MagicMock, ClosableUpstreamWebSocket, SimpleNamespace]: upstream_ws = ClosableUpstreamWebSocket() websocket = _client_websocket( AsyncMock( @@ -8004,7 +8169,7 @@ async def test_websocket_passthrough_reports_a_refused_client_frame_as_a_failure code=403, ) - async def refusing_gate(frame_data, valid_token): + async def refusing_gate(frame_data: str | bytes, valid_token: UserAPIKeyAuth, /) -> ProxyException | None: return denial websocket, upstream_ws, environment = await _run_passthrough_with_frame_gate(refusing_gate) @@ -8021,7 +8186,7 @@ async def test_websocket_passthrough_reports_a_refused_client_frame_as_a_failure @pytest.mark.asyncio async def test_websocket_passthrough_still_reports_an_allowed_session_as_a_success(): - async def permitting_gate(frame_data, valid_token): + async def permitting_gate(frame_data: str | bytes, valid_token: UserAPIKeyAuth, /) -> ProxyException | None: return None _, _, environment = await _run_passthrough_with_frame_gate(permitting_gate) From bb14db22c768c86092b4fe3d319503fc0b50c977 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:54:09 -0700 Subject: [PATCH 05/17] fix(proxy): count only Vertex AI deployments for Vertex AI Live model access A group another provider serves no longer grants a Vertex full path or gets its name rewritten for Vertex, and a caller with no model limit keeps working when every candidate name is a group serving another model. --- .../llm_passthrough_endpoints.py | 81 ++++++++++------- .../test_pass_through_endpoints.py | 87 +++++++++++++++++++ 2 files changed, 135 insertions(+), 33 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index e51b58a6c7b..2a536b9ffd1 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -3914,31 +3914,45 @@ def _deployment_named(model_name: str, deployments: Sequence[DeploymentTypedDict return next((deployment for deployment in deployments if deployment.get("model_name") == model_name), None) +VERTEX_LIVE_PROVIDERS: Final = frozenset({LlmProviders.VERTEX_AI.value, LlmProviders.VERTEX_AI_BETA.value}) + + +def _vertex_upstream_model(deployment: DeploymentTypedDict) -> str | None: + """ + The publisher model a Vertex AI deployment runs, or ``None`` for any other provider's deployment, since this + route runs every session on the proxy's Vertex credentials whatever provider a group was configured with + """ + litellm_params: Final = deployment["litellm_params"] + if "model" not in litellm_params: + return None + upstream: Final = litellm_params["model"] + try: + _, provider, _, _ = litellm.get_llm_provider( + model=upstream, custom_llm_provider=litellm_params.get("custom_llm_provider") + ) + except litellm.exceptions.BadRequestError: + return upstream + return upstream.removeprefix(f"{provider}/") if provider in VERTEX_LIVE_PROVIDERS else None + + +def _vertex_deployments(llm_router: Router | None) -> tuple[DeploymentTypedDict, ...]: + if llm_router is None: + return () + return tuple(deployment for deployment in (llm_router.get_model_list() or ()) if _vertex_upstream_model(deployment)) + + def _vertex_live_deployment(setup_model: str, llm_router: Router | None) -> DeploymentTypedDict | None: """ - The router deployment whose group a ``setup`` frame's model names, whichever addressing the client used. + The Vertex AI deployment whose group a ``setup`` frame's model names, whichever addressing the client used. The Live SDK wraps whatever the caller typed as ``models/``, and clients also send LiteLLM ids and full Vertex resource paths, so the group is read from the last segment when the whole string names none """ - if llm_router is None: - return None - deployments: Final = tuple(llm_router.get_model_list() or ()) + deployments: Final = _vertex_deployments(llm_router) exact: Final = _deployment_named(setup_model, deployments) return exact if exact is not None else _deployment_named(setup_model.rsplit("/", 1)[-1], deployments) -def _deployment_upstream_model(deployment: DeploymentTypedDict) -> str | None: - upstream: Final = deployment["litellm_params"].get("model") - if upstream is None: - return None - try: - _, provider, _, _ = litellm.get_llm_provider(model=upstream) - except litellm.exceptions.BadRequestError: - return upstream - return upstream.removeprefix(f"{provider}/") - - def _vertex_live_model_names(setup_model: str, llm_router: Router | None) -> tuple[str, ...]: """ The names a key can hold for the model a ``setup`` frame runs: the router group the setup rewriter sends it @@ -3953,16 +3967,17 @@ def _vertex_live_model_names(setup_model: str, llm_router: Router | None) -> tup ) if deployment is not None: return (deployment["model_name"],) - deployments: Final = () if llm_router is None else tuple(llm_router.get_model_list() or ()) publisher_model: Final = _vertex_publisher_model_suffix(setup_model) serving_groups: Final = tuple( dict.fromkeys( candidate["model_name"] - for candidate in deployments - if _vertex_publisher_model_suffix(_deployment_upstream_model(candidate) or "") == publisher_model + for candidate in _vertex_deployments(llm_router) + if _vertex_publisher_model_suffix(_vertex_upstream_model(candidate) or "") == publisher_model ) ) - group_names: Final = frozenset(candidate["model_name"] for candidate in deployments) + group_names: Final = frozenset( + candidate["model_name"] for candidate in (() if llm_router is None else llm_router.get_model_list() or ()) + ) publisher_model_id: Final = setup_model.rsplit("/", 1)[-1] return ( *serving_groups, @@ -4016,20 +4031,20 @@ async def _vertex_live_model_denial( for name in _vertex_live_model_names(setup_model, llm_router) if not _is_model_alias(name, valid_token, llm_router) ) - denials: Final = tuple([await _model_access_denial(name, valid_token, llm_router) for name in model_names]) - if None in denials: - return None - return ( - denials[0] - if denials - else ModelAccessDeniedProxyException( - message=model_access_denied_client_message(model=setup_model), - internal_message=f"Vertex AI Live setup model {setup_model} is only reachable through a model alias", - type=ProxyErrorTypes.key_model_access_denied, - param="model", - code=403, + if not model_names: + return ( + ModelAccessDeniedProxyException( + message=model_access_denied_client_message(model=setup_model), + internal_message=f"Vertex AI Live setup model {setup_model} has no name a model limit can grant", + type=ProxyErrorTypes.key_model_access_denied, + param="model", + code=403, + ) + if await _caller_has_model_limits(valid_token) + else None ) - ) + denials: Final = tuple([await _model_access_denial(name, valid_token, llm_router) for name in model_names]) + return None if None in denials else denials[0] async def _caller_has_model_limits(valid_token: UserAPIKeyAuth) -> bool: @@ -4154,7 +4169,7 @@ def _resolve_alias_to_upstream_model(setup_model: str, llm_router: Router | None The Live SDK wraps whatever the caller typed as ``models/``, so a gateway alias arrives prefixed """ deployment: Final = _vertex_live_deployment(setup_model, llm_router) - upstream: Final = None if deployment is None else _deployment_upstream_model(deployment) + upstream: Final = None if deployment is None else _vertex_upstream_model(deployment) return setup_model if upstream is None else upstream diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 4002816f2e3..6452b6c7ccf 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -7968,6 +7968,93 @@ async def test_vertex_live_full_path_is_not_granted_by_a_group_named_after_it_th assert reason is not None and "gemini-3-pro-live" in reason +@pytest.mark.asyncio +@pytest.mark.parametrize( + "setup_model", + [VERTEX_LIVE_FULL_PATH.format(model=VERTEX_LIVE_GROUP), "openai-native-audio"], + ids=["full-path", "group-name"], +) +async def test_vertex_live_setup_frame_is_not_granted_by_a_group_another_provider_serves(setup_model: str): + """Every Live session runs on the proxy's Vertex credentials, so a group another provider serves under the + same model id must not open the Vertex model""" + llm_router = litellm.Router( + model_list=[ + { + "model_name": "openai-native-audio", + "litellm_params": {"model": f"openai/{VERTEX_LIVE_GROUP}", "api_base": "https://openai-compatible"}, + }, + ] + ) + + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame(setup_model)], + UserAPIKeyAuth(token="hashed", models=["openai-native-audio"]), + llm_router=llm_router, + ) + + upstream_ws.send.assert_not_awaited() + assert _policy_close_reason(websocket) is not None + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("custom_llm_provider", "reaches_vertex"), [("vertex_ai", True), ("openai", False)]) +async def test_vertex_live_full_path_is_authorized_only_by_a_group_whose_custom_provider_is_vertex( + custom_llm_provider: str, reaches_vertex: bool +): + llm_router = litellm.Router( + model_list=[ + { + "model_name": "native-audio", + "litellm_params": { + "model": VERTEX_LIVE_GROUP, + "custom_llm_provider": custom_llm_provider, + "api_base": "https://openai-compatible", + }, + }, + ] + ) + + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame(VERTEX_LIVE_FULL_PATH.format(model=VERTEX_LIVE_GROUP))], + UserAPIKeyAuth(token="hashed", models=["native-audio"]), + llm_router=llm_router, + ) + + assert upstream_ws.send.await_count == (1 if reaches_vertex else 0) + assert (_policy_close_reason(websocket) is None) is reaches_vertex + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("valid_token", "reaches_vertex"), + [ + (UserAPIKeyAuth(token="hashed", models=[]), True), + (UserAPIKeyAuth(token="hashed", models=["gemini-3-pro-live"]), False), + ], + ids=["unrestricted-key", "limited-key"], +) +async def test_vertex_live_full_path_whose_every_name_is_taken_by_another_model_refuses_only_limited_keys( + valid_token: UserAPIKeyAuth, reaches_vertex: bool +): + """Groups holding both the bare and the ``vertex_ai/`` id for another model leave no name a key could hold, + which says nothing about a key that may call any model""" + llm_router = litellm.Router( + model_list=[ + {"model_name": "gemini-3-pro-live", "litellm_params": {"model": "vertex_ai/gemini-live-2.5-flash"}}, + {"model_name": "vertex_ai/gemini-3-pro-live", "litellm_params": {"model": "vertex_ai/gemini-live-2.5-flash"}}, + ] + ) + + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame(VERTEX_LIVE_FULL_PATH.format(model="gemini-3-pro-live"))], + valid_token, + llm_router=llm_router, + ) + + assert upstream_ws.send.await_count == (1 if reaches_vertex else 0) + assert (_policy_close_reason(websocket) is None) is reaches_vertex + + @pytest.mark.asyncio @pytest.mark.parametrize( "valid_token", From c5f7118770937275f235c040bd890073612004be Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 19:07:37 -0700 Subject: [PATCH 06/17] fix(proxy): skip OAuth provider logins when scanning deployments for Vertex AI Live --- .../llm_passthrough_endpoints.py | 3 +- .../test_pass_through_endpoints.py | 30 +++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 2a536b9ffd1..3731abe3134 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -50,6 +50,7 @@ from litellm.constants import ( from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.integrations.custom_guardrail import guardrail_request_data_with_streaming, is_guardrail_intervention from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix +from litellm.litellm_core_utils.get_llm_provider_logic import resolve_model_and_provider_for_metadata from litellm.llms.anthropic.common_utils import AnthropicModelInfo, merge_anthropic_beta_headers from litellm.llms.azure.passthrough.transformation import ( foreign_azure_deployment, @@ -3927,7 +3928,7 @@ def _vertex_upstream_model(deployment: DeploymentTypedDict) -> str | None: return None upstream: Final = litellm_params["model"] try: - _, provider, _, _ = litellm.get_llm_provider( + _, provider = resolve_model_and_provider_for_metadata( model=upstream, custom_llm_provider=litellm_params.get("custom_llm_provider") ) except litellm.exceptions.BadRequestError: diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 6452b6c7ccf..8bd2106e8cd 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -12,6 +12,7 @@ from contextlib import ExitStack, contextmanager from dataclasses import dataclass from datetime import datetime from io import BytesIO +from pathlib import Path from types import MappingProxyType, ModuleType, SimpleNamespace from typing import Final, Optional, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -8055,6 +8056,35 @@ async def test_vertex_live_full_path_whose_every_name_is_taken_by_another_model_ assert (_policy_close_reason(websocket) is None) is reaches_vertex +@pytest.mark.asyncio +async def test_vertex_live_setup_frame_never_refreshes_the_token_of_a_chatgpt_group_on_the_same_proxy( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + (tmp_path / "auth.json").write_text( + json.dumps({"access_token": "expired", "refresh_token": "refresh", "expires_at": 0}) + ) + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + monkeypatch.delenv("CHATGPT_AUTH_FILE", raising=False) + llm_router: Final = litellm.Router( + model_list=[ + {"model_name": "codex", "litellm_params": {"model": "chatgpt/gpt-5.4"}}, + {"model_name": "native-audio", "litellm_params": {"model": f"vertex_ai/{VERTEX_LIVE_GROUP}"}}, + ] + ) + + with respx.mock(assert_all_called=False) as transport: + outbound: Final = transport.route().respond(500) + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame("native-audio")], + UserAPIKeyAuth(token="hashed", models=["native-audio"]), + llm_router=llm_router, + ) + + assert not outbound.called + upstream_ws.send.assert_awaited_once() + assert _policy_close_reason(websocket) is None + + @pytest.mark.asyncio @pytest.mark.parametrize( "valid_token", From d48e964b6e406512ffb2e3c3c9c369804bf78416 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 19:57:08 -0700 Subject: [PATCH 07/17] fix(proxy): close the Vertex AI Live sockets only after the refusal is kept Closing the client socket inside the frame relay let the upstream reader finish first, which cancelled the relay mid-close, lost the refusal, and logged a session that served turns as a failure. The relay now returns the refusal and the handler closes both sockets. --- .../pass_through_endpoints.py | 12 ++-- .../test_pass_through_endpoints.py | 68 +++++++++++++++++-- 2 files changed, 68 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index bd24be1b808..d09ea8818fd 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2519,7 +2519,6 @@ class WebsocketClientFrameGate(Protocol): async def _relay_client_frame( - websocket: WebSocket, upstream_ws: ClientConnection, checked_frame: str | bytes, forwarded_frame: str | bytes, @@ -2529,14 +2528,12 @@ async def _relay_client_frame( """Forward one client frame upstream, or refuse it and hand back the denial that ended the session. The denial travels out so the caller reports a failed session instead of leaving the success hooks to - record a refusal as a success, and tearing the upstream down is the caller's job: awaiting that close - here lets the upstream reader finish first, which cancels this task mid-close and loses the denial + record a refusal as a success, and closing either socket is the caller's job: awaiting a close here lets + the upstream reader finish first, which cancels this task mid-close and loses the denial """ refusal: Final = None if gate is None else await gate(checked_frame, valid_token) if refusal is None: await upstream_ws.send(forwarded_frame) - return None - await websocket.close(code=CloseCode.POLICY_VIOLATION, reason=_refusal_close_reason(refusal.message)) return refusal @@ -2715,7 +2712,6 @@ async def websocket_passthrough_request( if text_data is not None: if ( text_refusal := await _relay_client_frame( - websocket, upstream_ws, text_data, _rewrite_vertex_live_setup_model(text_data, setup_model_rewriter), @@ -2783,7 +2779,6 @@ async def websocket_passthrough_request( elif bytes_data is not None: if ( bytes_refusal := await _relay_client_frame( - websocket, upstream_ws, bytes_data, bytes_data, @@ -2873,6 +2868,9 @@ async def websocket_passthrough_request( (result for result in task_results if isinstance(result, ProxyException)), None ) if client_frame_refusal is not None: + await websocket.close( + code=CloseCode.POLICY_VIOLATION, reason=_refusal_close_reason(client_frame_refusal.message) + ) await upstream_ws.close() if client_frame_refusal is not None and not websocket_messages: await log_failure(client_frame_refusal) diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 8bd2106e8cd..e388c658200 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -7555,7 +7555,7 @@ async def _run_vertex_live_gated_passthrough( """Drive /vertex_ai/live with the real frame gate, returning the client and upstream sockets""" upstream_ws = RecordingUpstreamWebSocket() websocket, _ = await _run_vertex_live_gated_session( - AsyncMock(side_effect=[*client_frames, {"type": "websocket.disconnect"}]), + _client_websocket(AsyncMock(side_effect=[*client_frames, {"type": "websocket.disconnect"}])), upstream_ws, valid_token, llm_router or _vertex_live_router(), @@ -7564,7 +7564,7 @@ async def _run_vertex_live_gated_passthrough( async def _run_vertex_live_gated_session( - receive: Callable[[], Awaitable[Mapping[str, str | bytes]]], + websocket: MagicMock, upstream_ws: "RecordingUpstreamWebSocket | ScriptedUpstreamWebSocket", valid_token: UserAPIKeyAuth, llm_router: litellm.Router, @@ -7574,8 +7574,6 @@ async def _run_vertex_live_gated_session( _build_vertex_live_setup_model_rewriter, ) - websocket = _client_websocket(receive) - with _patched_websocket_passthrough_environment(upstream_ws) as environment: await websocket_passthrough_request( websocket=websocket, @@ -8330,7 +8328,7 @@ async def test_vertex_live_turns_served_before_a_refused_frame_are_still_billed( yield _setup_frame("gemini-3-pro-live") websocket, environment = await _run_vertex_live_gated_session( - client_frames().__anext__, + _client_websocket(client_frames().__anext__), upstream_ws, UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), _vertex_live_router(), @@ -8346,6 +8344,66 @@ async def test_vertex_live_turns_served_before_a_refused_frame_are_still_billed( environment.proxy_logging.post_call_failure_hook.assert_not_awaited() +class StreamingUpstreamWebSocket(ScriptedUpstreamWebSocket): + """An upstream that sends the given frames in order, then keeps streaming model output until it is closed""" + + async def recv(self, decode: bool = True) -> str | None: + frame: Final = next(self._frames, None) + if frame is not None: + return frame + self.sent_every_frame.set() + await asyncio.sleep(0) + return json.dumps({"serverContent": {"modelTurn": {"parts": [{"text": "still talking"}]}}}) + + +def _client_websocket_with_a_close_handshake(receive) -> MagicMock: + """A client socket that, like Starlette's, refuses sends once its close starts and yields while closing""" + from starlette.websockets import WebSocketState + + websocket: Final = _client_websocket(receive) + + async def close(*args: object, **kwargs: object) -> None: + websocket.application_state = WebSocketState.DISCONNECTED + for _ in range(ClosableUpstreamWebSocket.CLOSE_HANDSHAKE_YIELDS): + await asyncio.sleep(0) + + async def send_text(data: str) -> None: + if websocket.application_state == WebSocketState.DISCONNECTED: + raise RuntimeError('Cannot call "send" once a close message has been sent.') + + websocket.close = AsyncMock(side_effect=close) + websocket.send_text = AsyncMock(side_effect=send_text) + return websocket + + +@pytest.mark.asyncio +async def test_vertex_live_turns_served_before_a_frame_refused_mid_stream_are_still_billed(): + upstream_ws: Final = StreamingUpstreamWebSocket( + ( + json.dumps({"setupComplete": {}}), + json.dumps({"serverContent": {"turnComplete": True}, "usageMetadata": {"totalTokenCount": 60}}), + ) + ) + + async def client_frames() -> AsyncIterator[dict[str, str]]: + yield _setup_frame("gemini-live-2.5-flash") + await upstream_ws.sent_every_frame.wait() + yield _setup_frame("gemini-3-pro-live") + + websocket, environment = await _run_vertex_live_gated_session( + _client_websocket_with_a_close_handshake(client_frames().__anext__), + upstream_ws, + UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), + _vertex_live_router(), + ) + + reason: Final = _policy_close_reason(websocket) + assert reason is not None and "gemini-3-pro-live" in reason + upstream_ws.close.assert_awaited() + environment.proxy_logging.post_call_success_hook.assert_awaited_once() + environment.proxy_logging.post_call_failure_hook.assert_not_awaited() + + async def _run_passthrough_with_frame_gate( gate: WebsocketClientFrameGate, ) -> tuple[MagicMock, ClosableUpstreamWebSocket, SimpleNamespace]: From 7fa65b3c5a2c9da5ce7dff340d2ea868cbfab48f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 20:22:45 -0700 Subject: [PATCH 08/17] fix(proxy): name a deployment's declared provider without resolving it in the Vertex AI Live scan Resolving a declared provider reads its API key and base, which a hosted secret manager answers with a network call for every deployment on every Live setup frame. The scan now takes the provider from custom_llm_provider or a known model prefix and resolves only bare names. --- .../get_llm_provider_logic.py | 16 +++++ .../llm_passthrough_endpoints.py | 10 ++-- .../test_pass_through_endpoints.py | 58 +++++++++++++++++++ 3 files changed, 78 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 19c1d5c414c..6ace8c42505 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -158,6 +158,22 @@ def declared_authenticating_provider(model: str | None, custom_llm_provider: str return declared if declared in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO else None +def declared_provider(model: str, custom_llm_provider: str | None = None) -> str | None: + """The provider this pair names outright, or None when only get_llm_provider can tell. + + get_llm_provider reads a named provider's API key and base on the way to returning it, and a hosted + secret manager answers each read with a network call, so a caller that needs only the name reads it here. + """ + if custom_llm_provider: + return custom_llm_provider + prefix, slash, _ = model.partition("/") + if not slash: + return None + if prefix in litellm.provider_list or JSONProviderRegistry.exists(prefix) or is_registered_custom_provider(prefix): + return prefix + return None + + def resolve_model_and_provider_for_metadata(model: str, custom_llm_provider: str | None = None) -> tuple[str, str]: declared: Final = declared_authenticating_provider(model, custom_llm_provider) if declared is not None: diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 3731abe3134..68c11688c9c 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -50,7 +50,7 @@ from litellm.constants import ( from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.integrations.custom_guardrail import guardrail_request_data_with_streaming, is_guardrail_intervention from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix -from litellm.litellm_core_utils.get_llm_provider_logic import resolve_model_and_provider_for_metadata +from litellm.litellm_core_utils.get_llm_provider_logic import declared_provider, inferred_provider from litellm.llms.anthropic.common_utils import AnthropicModelInfo, merge_anthropic_beta_headers from litellm.llms.azure.passthrough.transformation import ( foreign_azure_deployment, @@ -3927,11 +3927,9 @@ def _vertex_upstream_model(deployment: DeploymentTypedDict) -> str | None: if "model" not in litellm_params: return None upstream: Final = litellm_params["model"] - try: - _, provider = resolve_model_and_provider_for_metadata( - model=upstream, custom_llm_provider=litellm_params.get("custom_llm_provider") - ) - except litellm.exceptions.BadRequestError: + declared: Final = declared_provider(upstream, litellm_params.get("custom_llm_provider")) + provider: Final = declared or inferred_provider(upstream) + if provider is None: return upstream return upstream.removeprefix(f"{provider}/") if provider in VERTEX_LIVE_PROVIDERS else None diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index e388c658200..659be2666bc 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -34,6 +34,8 @@ from litellm._logging import verbose_proxy_logger from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS, WEBSOCKET_CLOSE_REASON_MAX_BYTES from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.custom_secret_manager import CustomSecretManager +from litellm.llms.custom_llm import CustomLLM from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._lazy_features import LazyFeature, attach_lazy_features from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth @@ -65,6 +67,7 @@ from litellm.proxy.pass_through_endpoints.success_handler import ( from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.types import utils as types_utils from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, @@ -8083,6 +8086,61 @@ async def test_vertex_live_setup_frame_never_refreshes_the_token_of_a_chatgpt_gr assert _policy_close_reason(websocket) is None +class SecretReadRecorder(CustomSecretManager): + """A hosted secret manager that answers every read and records the name it was asked for""" + + def __init__(self) -> None: + self.names_read = MagicMock() + + def sync_read_secret( + self, secret_name: str, optional_params: dict | None = None, timeout: float | httpx.Timeout | None = None + ) -> str: + self.names_read(secret_name) + return "a-value" + + async def async_read_secret( + self, secret_name: str, optional_params: dict | None = None, timeout: float | httpx.Timeout | None = None + ) -> str: + return self.sync_read_secret(secret_name) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "other_group", + [ + {"model": "groq/llama-4-scout"}, + {"model": "llama-4-scout", "custom_llm_provider": "groq"}, + {"model": "veniceai/venice-uncensored"}, + {"model": "in-house/assistant"}, + ], + ids=["provider-prefix", "custom-llm-provider", "json-provider", "registered-custom-provider"], +) +async def test_vertex_live_setup_frame_never_reads_the_secret_manager_for_another_providers_group( + monkeypatch: pytest.MonkeyPatch, other_group: dict[str, str] +): + monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "in-house", "custom_handler": CustomLLM()}]) + llm_router: Final = litellm.Router( + model_list=[ + {"model_name": "other", "litellm_params": other_group}, + {"model_name": "native-audio", "litellm_params": {"model": f"vertex_ai/{VERTEX_LIVE_GROUP}"}}, + ] + ) + secret_manager: Final = SecretReadRecorder() + monkeypatch.setattr(litellm, "secret_manager_client", secret_manager) + monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM) + monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only")) + + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame("native-audio")], + UserAPIKeyAuth(token="hashed", models=["native-audio"]), + llm_router=llm_router, + ) + + secret_manager.names_read.assert_not_called() + upstream_ws.send.assert_awaited_once() + assert _policy_close_reason(websocket) is None + + @pytest.mark.asyncio @pytest.mark.parametrize( "valid_token", From f0c8bc293e1019dbec9668e6ccd908e927d74c79 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 23:32:53 -0700 Subject: [PATCH 09/17] fix(proxy): authorize the Vertex Live setup model like a connect-time model and keep main's group rewriting Re-run full websocket auth for the setup model, so JWT scopes, model budgets and the pricing block apply, and release its budget reservation. Resolve a listed router alias to its target group, and keep sending another provider's group to its Vertex model as main does --- .../llm_passthrough_endpoints.py | 144 ++++++++++++------ .../pass_through_endpoints.py | 3 - .../test_get_llm_provider_logic.py | 21 +++ tests/unit/proxy/auth/test_auth_checks.py | 34 ++++- .../test_pass_through_endpoints.py | 137 +++++++++++++++-- 5 files changed, 282 insertions(+), 57 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 8703eeb22b1..1cfab1f928e 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -15,11 +15,11 @@ import os import posixpath import re import sys -from collections.abc import AsyncGenerator, Callable, Mapping, Sequence +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from functools import partial from types import MappingProxyType -from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast +from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias, cast import httpx import openai @@ -90,6 +90,7 @@ from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy modul get_bearer_token, is_no_auth_dev_mode, user_api_key_auth, + user_api_key_auth_accepted_websocket_for_model, user_api_key_auth_websocket, user_api_key_auth_websocket_for_model, ) @@ -3933,20 +3934,26 @@ def _deployment_named(model_name: str, deployments: Sequence[DeploymentTypedDict VERTEX_LIVE_PROVIDERS: Final = frozenset({LlmProviders.VERTEX_AI.value, LlmProviders.VERTEX_AI_BETA.value}) -def _vertex_upstream_model(deployment: DeploymentTypedDict) -> str | None: - """ - The publisher model a Vertex AI deployment runs, or ``None`` for any other provider's deployment, since this - route runs every session on the proxy's Vertex credentials whatever provider a group was configured with - """ +def _upstream_model_and_provider(deployment: DeploymentTypedDict) -> tuple[str, str | None] | None: litellm_params: Final = deployment["litellm_params"] if "model" not in litellm_params: return None upstream: Final = litellm_params["model"] declared: Final = declared_provider(upstream, litellm_params.get("custom_llm_provider")) provider: Final = declared or inferred_provider(upstream) - if provider is None: - return upstream - return upstream.removeprefix(f"{provider}/") if provider in VERTEX_LIVE_PROVIDERS else None + return (upstream if provider is None else upstream.removeprefix(f"{provider}/"), provider) + + +def _vertex_upstream_model(deployment: DeploymentTypedDict) -> str | None: + """ + The publisher model a Vertex AI deployment runs, or ``None`` for any other provider's deployment, since a + group another provider serves grants nothing on the proxy's Vertex credentials + """ + resolved: Final = _upstream_model_and_provider(deployment) + if resolved is None: + return None + upstream, provider = resolved + return upstream if provider is None or provider in VERTEX_LIVE_PROVIDERS else None def _vertex_deployments(llm_router: Router | None) -> tuple[DeploymentTypedDict, ...]: @@ -3955,33 +3962,35 @@ def _vertex_deployments(llm_router: Router | None) -> tuple[DeploymentTypedDict, return tuple(deployment for deployment in (llm_router.get_model_list() or ()) if _vertex_upstream_model(deployment)) -def _vertex_live_deployment(setup_model: str, llm_router: Router | None) -> DeploymentTypedDict | None: +def _live_deployment(setup_model: str, llm_router: Router | None) -> DeploymentTypedDict | None: """ - The Vertex AI deployment whose group a ``setup`` frame's model names, whichever addressing the client used. + The deployment whose group a ``setup`` frame's model names, whichever addressing the client used. The Live SDK wraps whatever the caller typed as ``models/``, and clients also send LiteLLM ids and full Vertex resource paths, so the group is read from the last segment when the whole string names none """ - deployments: Final = _vertex_deployments(llm_router) + if setup_model.startswith("projects/") or llm_router is None: + return None + deployments: Final = tuple(llm_router.get_model_list() or ()) exact: Final = _deployment_named(setup_model, deployments) return exact if exact is not None else _deployment_named(setup_model.rsplit("/", 1)[-1], deployments) def _vertex_live_model_names(setup_model: str, llm_router: Router | None) -> tuple[str, ...]: """ - The names a key can hold for the model a ``setup`` frame runs: the router group the setup rewriter sends it - through, else every group serving the publisher model Vertex runs, and that model's id bare or under LiteLLM's - ``vertex_ai/`` prefix unless a group holds that name for another model. + The names a key can hold for the model a ``setup`` frame runs: the Vertex AI group the setup rewriter sends it + through, else every Vertex AI group serving the publisher model Vertex runs, and that model's id bare or under + LiteLLM's ``vertex_ai/`` prefix unless a group holds that name for another model. A full resource path skips the rewriter and reaches Vertex as sent, so only its publisher model counts, and a wildcard is matched against that id alone so addressing in front of it cannot satisfy the pattern """ - deployment: Final = ( - None if setup_model.startswith("projects/") else _vertex_live_deployment(setup_model, llm_router) - ) - if deployment is not None: - return (deployment["model_name"],) - publisher_model: Final = _vertex_publisher_model_suffix(setup_model) + deployment: Final = _live_deployment(setup_model, llm_router) + if deployment is not None and _vertex_upstream_model(deployment) is not None: + alias_target: Final = None if llm_router is None else llm_router.get_model_from_alias(deployment["model_name"]) + return (alias_target or deployment["model_name"],) + opened_model: Final = _resolve_alias_to_upstream_model(setup_model, llm_router) + publisher_model: Final = _vertex_publisher_model_suffix(opened_model) serving_groups: Final = tuple( dict.fromkeys( candidate["model_name"] @@ -3992,7 +4001,7 @@ def _vertex_live_model_names(setup_model: str, llm_router: Router | None) -> tup group_names: Final = frozenset( candidate["model_name"] for candidate in (() if llm_router is None else llm_router.get_model_list() or ()) ) - publisher_model_id: Final = setup_model.rsplit("/", 1)[-1] + publisher_model_id: Final = opened_model.rsplit("/", 1)[-1] return ( *serving_groups, *(name for name in (publisher_model_id, f"vertex_ai/{publisher_model_id}") if name not in group_names), @@ -4008,6 +4017,19 @@ def _is_model_alias(name: str, valid_token: UserAPIKeyAuth, llm_router: Router | ) +def _as_model_denial(refusal: ProxyException | HTTPException) -> ProxyException: + match refusal: + case ProxyException(): + return refusal + case HTTPException(): + return ProxyException( + message=str(refusal.detail), + type=ProxyErrorTypes.auth_error, + param="model", + code=refusal.status_code, + ) + + async def _model_access_denial( model: str, valid_token: UserAPIKeyAuth, llm_router: Router | None ) -> ProxyException | None: @@ -4022,23 +4044,47 @@ async def _model_access_denial( user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) - except ProxyException as denial: - return denial - except HTTPException as refusal: - return ProxyException( - message=str(refusal.detail), - type=ProxyErrorTypes.auth_error, - param="model", - code=refusal.status_code, - ) + except (ProxyException, HTTPException) as refusal: + return _as_model_denial(refusal) return None +VertexLiveModelAuthorizer: TypeAlias = Callable[[str], Awaitable[UserAPIKeyAuth]] + + +async def _model_reauthorization_denial( + model: str, authorize_model: VertexLiveModelAuthorizer +) -> ProxyException | None: + from litellm.proxy.spend_tracking.budget_reservation import release_or_invalidate_budget_reservation + + try: + model_auth: Final = await authorize_model(model) + except (ProxyException, HTTPException) as refusal: + return _as_model_denial(refusal) + await release_or_invalidate_budget_reservation(budget_reservation=model_auth.budget_reservation) + return None + + +async def _first_reauthorization_denial( + model_names: tuple[str, ...], authorize_model: VertexLiveModelAuthorizer +) -> ProxyException | None: + denial: Final = await _model_reauthorization_denial(model_names[0], authorize_model) + if denial is None or len(model_names) == 1: + return denial + later_denial: Final = await _first_reauthorization_denial(model_names[1:], authorize_model) + return None if later_denial is None else denial + + async def _vertex_live_model_denial( - setup_model: str, valid_token: UserAPIKeyAuth, llm_router: Router | None + setup_model: str, + valid_token: UserAPIKeyAuth, + llm_router: Router | None, + authorize_model: VertexLiveModelAuthorizer, ) -> ProxyException | None: """ - The setup rewriter forwards a name as sent, so a name an alias maps to another model grants nothing here + The setup rewriter forwards an alias name as sent unless it is a listed router alias, which already resolved to + its target group, so an alias name grants nothing here. A name the key may call is then authorized the way a + request naming it at connect time is """ model_names: Final = tuple( name @@ -4058,7 +4104,10 @@ async def _vertex_live_model_denial( else None ) denials: Final = tuple([await _model_access_denial(name, valid_token, llm_router) for name in model_names]) - return None if None in denials else denials[0] + callable_names: Final = tuple(name for name, denial in zip(model_names, denials, strict=True) if denial is None) + if not callable_names: + return denials[0] + return await _first_reauthorization_denial(callable_names, authorize_model) async def _caller_has_model_limits(valid_token: UserAPIKeyAuth) -> bool: @@ -4073,11 +4122,14 @@ async def _caller_has_model_limits(valid_token: UserAPIKeyAuth) -> bool: async def _vertex_live_setup_denial( - setup_model: str | _UnreadableVertexLiveFrame, valid_token: UserAPIKeyAuth, llm_router: Router | None + setup_model: str | _UnreadableVertexLiveFrame, + valid_token: UserAPIKeyAuth, + llm_router: Router | None, + authorize_model: VertexLiveModelAuthorizer, ) -> ProxyException | None: match setup_model: case str(): - denial: Final = await _vertex_live_model_denial(setup_model, valid_token, llm_router) + denial: Final = await _vertex_live_model_denial(setup_model, valid_token, llm_router, authorize_model) if denial is not None: log_model_access_denial(denial) return denial @@ -4094,7 +4146,9 @@ async def _vertex_live_setup_denial( ) -def _build_vertex_live_client_frame_gate(llm_router: Router | None) -> WebsocketClientFrameGate: +def _build_vertex_live_client_frame_gate( + llm_router: Router | None, authorize_model: VertexLiveModelAuthorizer +) -> WebsocketClientFrameGate: """ Apply the checks a model named at auth time gets to the model a client's ``setup`` frame names. @@ -4111,7 +4165,7 @@ def _build_vertex_live_client_frame_gate(llm_router: Router | None) -> Websocket if model is None: return None try: - return await _vertex_live_setup_denial(model, valid_token, llm_router) + return await _vertex_live_setup_denial(model, valid_token, llm_router, authorize_model) except Exception: # noqa: BLE001 # fail closed: a gate that cannot answer refuses, and says so in band verbose_proxy_logger.exception("Vertex AI Live passthrough: model access check failed for %s", model) return ProxyException( @@ -4182,9 +4236,9 @@ def _resolve_alias_to_upstream_model(setup_model: str, llm_router: Router | None """ The Live SDK wraps whatever the caller typed as ``models/``, so a gateway alias arrives prefixed """ - deployment: Final = _vertex_live_deployment(setup_model, llm_router) - upstream: Final = None if deployment is None else _vertex_upstream_model(deployment) - return setup_model if upstream is None else upstream + deployment: Final = _live_deployment(setup_model, llm_router) + resolved: Final = None if deployment is None else _upstream_model_and_provider(deployment) + return setup_model if resolved is None else resolved[0] async def vertex_ai_live_websocket_passthrough( @@ -4193,6 +4247,7 @@ async def vertex_ai_live_websocket_passthrough( vertex_project: str | None = None, vertex_location: str | None = None, user_api_key_dict: UserAPIKeyAuth | None = None, + authorize_model: VertexLiveModelAuthorizer | None = None, ): """ Vertex AI Live API WebSocket Pass-through Function @@ -4291,7 +4346,10 @@ async def vertex_ai_live_websocket_passthrough( vertex_location=resolved_location, llm_router=_get_llm_router(), ), - client_frame_gate=_build_vertex_live_client_frame_gate(llm_router=_get_llm_router()), + client_frame_gate=_build_vertex_live_client_frame_gate( + llm_router=_get_llm_router(), + authorize_model=authorize_model or partial(user_api_key_auth_accepted_websocket_for_model, websocket), + ), ) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index d09ea8818fd..c1f733c738c 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2456,9 +2456,6 @@ def _truncated_close_reason(reason: str) -> str: def _refusal_close_reason(message: str) -> str: - """ - The proxy's own refusal message as a close reason, ending on the last whole sentence that fits - """ clipped: Final = _truncated_close_reason(message) sentence_end: Final = clipped.rfind(". ") return clipped if clipped == message or sentence_end == -1 else clipped[: sentence_end + 1] diff --git a/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py b/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py index 25de5ef60d9..529a2bf9e32 100644 --- a/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py +++ b/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py @@ -11,6 +11,7 @@ import litellm from litellm import CustomLLM from litellm.litellm_core_utils import get_llm_provider_logic from litellm.litellm_core_utils.get_llm_provider_logic import ( + declared_provider, get_llm_provider, resolve_model_and_provider_for_metadata, get_provider_affinity_header, @@ -70,6 +71,26 @@ def test_is_registered_custom_provider(registered_custom_provider: str, candidat assert is_registered_custom_provider(candidate) is expected +@pytest.mark.parametrize( + ("model", "custom_llm_provider", "expected"), + [ + ("vertex_ai/gemini-live-2.5-flash", None, "vertex_ai"), + ("llama-4-scout", "groq", "groq"), + ("gemini/gemini-2.0-flash-live-001", "vertex_ai", "vertex_ai"), + ("veniceai/venice-uncensored", None, "veniceai"), + (f"{CUSTOM_PROVIDER}/my-model", None, CUSTOM_PROVIDER), + ("gemini-live-2.5-flash", None, None), + ("models/gemini-live-2.5-flash", None, None), + ], + ids=["provider-prefix", "custom-llm-provider", "custom-llm-provider-wins", "json-provider", "custom-provider", + "bare-name", "unknown-prefix"], +) +def test_declared_provider_names_only_a_provider_the_pair_spells_out( + registered_custom_provider: str, model: str, custom_llm_provider: str | None, expected: str | None +) -> None: + assert declared_provider(model, custom_llm_provider) == expected + + def test_get_llm_provider_leaves_fal_ai_api_base_unset_for_global_fallback() -> None: _, provider, _, api_base = get_llm_provider(model="fal_ai/fal-ai/flux/schnell") assert provider == "fal_ai" diff --git a/tests/unit/proxy/auth/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py index e783f01769e..9a080284efe 100644 --- a/tests/unit/proxy/auth/test_auth_checks.py +++ b/tests/unit/proxy/auth/test_auth_checks.py @@ -1,7 +1,7 @@ # What is this? ## Tests if 'get_end_user_object' works as expected -import sys, os, asyncio, time, random, uuid +import sys, os, asyncio, time, random, re, uuid import hashlib import json import traceback @@ -39,6 +39,7 @@ from litellm.proxy.utils import PrismaClient import litellm.proxy.common_utils.auth_cache_invalidation_pubsub as pubsub_module from litellm.caching.in_memory_cache import InMemoryCache from litellm.proxy.auth.auth_checks import ( + can_caller_call_resolved_model, can_team_access_model, delete_cache_key_objects, is_model_cost_zero, @@ -2425,3 +2426,34 @@ async def test_bulk_eviction_clears_the_redis_copies_before_waiting_on_a_wedged_ assert refilled == [] assert redis.entries == {} + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("team_id", "refused"), [(None, True), ("live-team", False)], ids=["personal-key", "team-key"]) +async def test_can_caller_call_resolved_model_holds_a_personal_key_to_its_owners_models( + team_id: str | None, refused: bool +) -> None: + owner_id: Final = f"live-owner-{uuid.uuid4()}" + cache: Final = UserApiKeyCache() + await cache.async_set_cache(key=owner_id, value=LiteLLM_UserTable(user_id=owner_id, models=["gemini-3.8-live"])) + router: Final = litellm.Router( + model_list=[ + {"model_name": "gemini-3.8-live", "litellm_params": {"model": "vertex_ai/gemini-3.8-live"}}, + {"model_name": "gemini-live-2.5-flash", "litellm_params": {"model": "vertex_ai/gemini-live-2.5-flash"}}, + ] + ) + + call: Final = can_caller_call_resolved_model( + model="gemini-live-2.5-flash", + valid_token=UserAPIKeyAuth(api_key="sk-live", user_id=owner_id, team_id=team_id), + llm_router=router, + prisma_client=MagicMock(), + user_api_key_cache=cache, + proxy_logging_obj=MagicMock(), + ) + + if refused: + with pytest.raises(ProxyException, match=re.escape("gemini-live-2.5-flash")): + await call + else: + await call diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 659be2666bc..22c109acb63 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -7550,10 +7550,18 @@ def _vertex_live_router() -> litellm.Router: ) +def _reauthorized_as(valid_token: UserAPIKeyAuth) -> Callable[[str], Awaitable[UserAPIKeyAuth]]: + async def authorize_model(model: str) -> UserAPIKeyAuth: + return valid_token + + return authorize_model + + async def _run_vertex_live_gated_passthrough( client_frames: Sequence[Mapping[str, str | bytes]], valid_token: UserAPIKeyAuth, llm_router: litellm.Router | None = None, + authorize_model: Callable[[str], Awaitable[UserAPIKeyAuth]] | None = None, ) -> tuple[MagicMock, RecordingUpstreamWebSocket]: """Drive /vertex_ai/live with the real frame gate, returning the client and upstream sockets""" upstream_ws = RecordingUpstreamWebSocket() @@ -7562,6 +7570,7 @@ async def _run_vertex_live_gated_passthrough( upstream_ws, valid_token, llm_router or _vertex_live_router(), + authorize_model, ) return websocket, upstream_ws @@ -7571,6 +7580,7 @@ async def _run_vertex_live_gated_session( upstream_ws: "RecordingUpstreamWebSocket | ScriptedUpstreamWebSocket", valid_token: UserAPIKeyAuth, llm_router: litellm.Router, + authorize_model: Callable[[str], Awaitable[UserAPIKeyAuth]] | None = None, ) -> tuple[MagicMock, SimpleNamespace]: from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( _build_vertex_live_client_frame_gate, @@ -7591,7 +7601,10 @@ async def _run_vertex_live_gated_session( vertex_location="global", llm_router=llm_router, ), - client_frame_gate=_build_vertex_live_client_frame_gate(llm_router=llm_router), + client_frame_gate=_build_vertex_live_client_frame_gate( + llm_router=llm_router, + authorize_model=authorize_model or _reauthorized_as(valid_token), + ), ) return websocket, environment @@ -8173,23 +8186,26 @@ async def test_vertex_live_setup_frame_naming_an_alias_cannot_open_the_model_the @pytest.mark.asyncio -async def test_vertex_live_setup_frame_naming_a_router_alias_cannot_open_the_model_the_alias_shadows(): +@pytest.mark.parametrize("hidden", [False, True], ids=["listed-alias", "hidden-alias"]) +async def test_vertex_live_setup_frame_naming_a_router_alias_cannot_open_the_model_the_alias_shadows(hidden): + """A listed router alias runs as its target group, while the setup rewriter cannot see a hidden one and sends it + to Vertex as the model whose id it reuses, which a key holding only the target must not open""" llm_router = litellm.Router( model_list=[ {"model_name": "gemini-live-2.5-flash", "litellm_params": {"model": "vertex_ai/gemini-live-2.5-flash"}}, ], - model_group_alias={"gemini-3-pro-live": "gemini-live-2.5-flash"}, + model_group_alias={"gemini-3-pro-live": {"model": "gemini-live-2.5-flash", "hidden": hidden}}, ) - websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + _, upstream_ws = await _run_vertex_live_gated_passthrough( [_setup_frame("gemini-3-pro-live")], UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), llm_router=llm_router, ) - upstream_ws.send.assert_not_awaited() - reason = _policy_close_reason(websocket) - assert reason is not None and "gemini-3-pro-live" in reason + sent_models = [json.loads(call.args[0])["setup"]["model"] for call in upstream_ws.send.await_args_list] + assert not any(model.endswith("/gemini-3-pro-live") for model in sent_models) + assert sent_models == ([] if hidden else ["projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash"]) @pytest.mark.asyncio @@ -8235,6 +8251,69 @@ async def test_vertex_live_setup_frame_is_checked_against_the_team_in_every_addr assert reason is not None and "gemini-3-pro-live" in reason +def _another_providers_live_router() -> litellm.Router: + return litellm.Router( + model_list=[ + {"model_name": "studio-live", "litellm_params": {"model": "gemini/gemini-2.0-flash-live-001"}}, + {"model_name": "vertex-flash-live", "litellm_params": {"model": "vertex_ai/gemini-2.0-flash-live-001"}}, + ] + ) + + +@pytest.mark.asyncio +async def test_vertex_live_setup_frame_naming_another_providers_group_runs_that_groups_model_on_vertex(): + """A caller who may call any model kept reaching Vertex through a group another provider serves before the + gate existed, so the setup rewriter still sends that group's model id rather than the group name""" + _, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame("models/studio-live")], + UserAPIKeyAuth(token="hashed"), + llm_router=_another_providers_live_router(), + ) + + sent_setup = json.loads(upstream_ws.send.await_args.args[0])["setup"] + assert sent_setup["model"] == "projects/proj-db/locations/global/publishers/google/models/gemini-2.0-flash-live-001" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("models", "reaches_vertex"), + [(["studio-live"], False), (["vertex-flash-live"], True), (["vertex_ai/gemini-2.0-flash-live-001"], True)], + ids=["other-provider-group", "vertex-group-serving-it", "vertex-model-id"], +) +async def test_vertex_live_setup_frame_naming_another_providers_group_is_authorized_by_the_vertex_model_it_runs( + models, reaches_vertex +): + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame("studio-live")], + UserAPIKeyAuth(token="hashed", models=models), + llm_router=_another_providers_live_router(), + ) + + assert upstream_ws.send.await_count == (1 if reaches_vertex else 0) + reason = _policy_close_reason(websocket) + assert (reason is None) is reaches_vertex + assert reaches_vertex or ("vertex-flash-live" in reason and "vertex_ai/studio-live" not in reason) + + +@pytest.mark.asyncio +async def test_vertex_live_setup_frame_naming_a_router_alias_is_authorized_by_its_target_group(): + """The setup rewriter runs a router alias as its target group, so the key holding that group may open it""" + llm_router = litellm.Router( + model_list=[{"model_name": "vertex-live", "litellm_params": {"model": f"vertex_ai/{VERTEX_LIVE_GROUP}"}}], + model_group_alias={"gemini-live-alias": "vertex-live"}, + ) + + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame("models/gemini-live-alias")], + UserAPIKeyAuth(token="hashed", models=["vertex-live"]), + llm_router=llm_router, + ) + + assert _policy_close_reason(websocket) is None + sent_setup = json.loads(upstream_ws.send.await_args.args[0])["setup"] + assert sent_setup["model"] == f"projects/proj-db/locations/global/publishers/google/models/{VERTEX_LIVE_GROUP}" + + @pytest.mark.asyncio @pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS) async def test_vertex_live_setup_frame_resolves_an_access_group_grant_in_every_addressing(addressing): @@ -8255,16 +8334,19 @@ async def test_vertex_live_setup_frame_resolves_an_access_group_grant_in_every_a }, ] ) - gate = _build_vertex_live_client_frame_gate(llm_router=llm_router) + token = UserAPIKeyAuth(token="hashed", models=["live-team"]) + gate = _build_vertex_live_client_frame_gate(llm_router=llm_router, authorize_model=_reauthorized_as(token)) frame = json.dumps({"setup": {"model": addressing.format(model=VERTEX_LIVE_GROUP)}}) - denial = await gate(frame, UserAPIKeyAuth(token="hashed", models=["live-team"])) + denial = await gate(frame, token) assert denial is None async def _run_vertex_live_route( - client_frames: Sequence[Mapping[str, str | bytes]], valid_token: UserAPIKeyAuth + client_frames: Sequence[Mapping[str, str | bytes]], + valid_token: UserAPIKeyAuth, + reauthorize_through_user_api_key_auth: bool = False, ) -> tuple[MagicMock, RecordingUpstreamWebSocket, SimpleNamespace]: """Drive the /vertex_ai/live route, so the gate under test is whichever one the route itself wires""" from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( @@ -8276,6 +8358,12 @@ async def _run_vertex_live_route( websocket = _client_websocket( AsyncMock(side_effect=[*client_frames, {"type": "websocket.disconnect"}]), ) + websocket.headers = {"authorization": "Bearer sk-live-client"} + websocket.scope = { + "type": "websocket", + "path": "/vertex_ai/live", + "headers": [(b"authorization", b"Bearer sk-live-client")], + } with ( _patched_websocket_passthrough_environment(upstream_ws) as environment, @@ -8291,6 +8379,7 @@ async def _run_vertex_live_route( vertex_project="proj-db", vertex_location="global", user_api_key_dict=valid_token, + authorize_model=None if reauthorize_through_user_api_key_auth else _reauthorized_as(valid_token), ) return websocket, upstream_ws, environment @@ -8315,6 +8404,34 @@ async def test_vertex_live_route_gates_the_setup_frame_even_when_the_query_model assert refusal.type == ProxyErrorTypes.key_model_access_denied +@pytest.mark.asyncio +async def test_vertex_live_route_authorizes_the_setup_model_the_way_connect_time_auth_checks_a_query_model(): + """ + Connect-time auth applies JWT scopes, per-model budgets and the unpriced-model block to ``?model=``, so a model + the key's list allows still has to pass them when the setup frame is the first place it is named + """ + budget_denial = ProxyException( + message="Budget has been exceeded for model gemini-live-2.5-flash.", + type=ProxyErrorTypes.budget_exceeded, + param="model", + code=400, + ) + with patch( + "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True, side_effect=budget_denial + ) as reauthorization: + websocket, upstream_ws, _ = await _run_vertex_live_route( + [_setup_frame("models/gemini-live-2.5-flash")], + UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), + reauthorize_through_user_api_key_auth=True, + ) + + upstream_ws.send.assert_not_awaited() + assert _policy_close_reason(websocket) == budget_denial.message + reauthorized_request = reauthorization.call_args.kwargs["request"] + assert json.loads(await reauthorized_request.body()) == {"model": "gemini-live-2.5-flash"} + assert reauthorization.call_args.kwargs["api_key"] == "Bearer sk-live-client" + + @pytest.mark.asyncio async def test_vertex_live_route_keeps_the_key_allowlist_out_of_the_close_reason_and_in_the_proxy_log(caplog): """The client learns which model was refused, while only the proxy log names the models the key may reach""" From 98492c716fcaa89e8720299b28a42ee27686decc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 03:00:53 -0700 Subject: [PATCH 10/17] fix(proxy): log a refused Vertex Live setup model once, rewrite binary setup frames, and check a group an alias shadows as itself A setup model refused by its re-authorization was logged by the auth handler and again by the passthrough, so it now carries a marker that it was already logged. Binary setup frames were checked as if rewritten but forwarded as sent, so they go through the same rewriter. A router alias named like a real group made the gate check the alias target while the rewriter sent the group's own model, so a name that is a real group is checked as that group. A regression test pins the release of the budget the re-authorization reserves. --- .../llm_passthrough_endpoints.py | 23 +++-- .../pass_through_endpoints.py | 69 +++++++++++---- .../test_pass_through_endpoints.py | 86 +++++++++++++++---- 3 files changed, 135 insertions(+), 43 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 1cfab1f928e..de5d3047989 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -127,7 +127,9 @@ from litellm.proxy.pass_through_endpoints.injected_stream_usage import ( without_injected_usage_frames, ) from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + ClientFrameRefusal, HttpPassThroughEndpointHelpers, + LoggedClientFrameRefusal, WebsocketClientFrameGate, create_pass_through_route, create_websocket_passthrough_route, @@ -3987,8 +3989,11 @@ def _vertex_live_model_names(setup_model: str, llm_router: Router | None) -> tup """ deployment: Final = _live_deployment(setup_model, llm_router) if deployment is not None and _vertex_upstream_model(deployment) is not None: - alias_target: Final = None if llm_router is None else llm_router.get_model_from_alias(deployment["model_name"]) - return (alias_target or deployment["model_name"],) + group: Final = deployment["model_name"] + alias_target: Final = ( + None if llm_router is None or group in llm_router.model_names else llm_router.get_model_from_alias(group) + ) + return (alias_target or group,) opened_model: Final = _resolve_alias_to_upstream_model(setup_model, llm_router) publisher_model: Final = _vertex_publisher_model_suffix(opened_model) serving_groups: Final = tuple( @@ -4054,20 +4059,20 @@ VertexLiveModelAuthorizer: TypeAlias = Callable[[str], Awaitable[UserAPIKeyAuth] async def _model_reauthorization_denial( model: str, authorize_model: VertexLiveModelAuthorizer -) -> ProxyException | None: +) -> LoggedClientFrameRefusal | None: from litellm.proxy.spend_tracking.budget_reservation import release_or_invalidate_budget_reservation try: model_auth: Final = await authorize_model(model) except (ProxyException, HTTPException) as refusal: - return _as_model_denial(refusal) + return LoggedClientFrameRefusal(_as_model_denial(refusal)) await release_or_invalidate_budget_reservation(budget_reservation=model_auth.budget_reservation) return None async def _first_reauthorization_denial( model_names: tuple[str, ...], authorize_model: VertexLiveModelAuthorizer -) -> ProxyException | None: +) -> LoggedClientFrameRefusal | None: denial: Final = await _model_reauthorization_denial(model_names[0], authorize_model) if denial is None or len(model_names) == 1: return denial @@ -4080,7 +4085,7 @@ async def _vertex_live_model_denial( valid_token: UserAPIKeyAuth, llm_router: Router | None, authorize_model: VertexLiveModelAuthorizer, -) -> ProxyException | None: +) -> ClientFrameRefusal | None: """ The setup rewriter forwards an alias name as sent unless it is a listed router alias, which already resolved to its target group, so an alias name grants nothing here. A name the key may call is then authorized the way a @@ -4126,11 +4131,11 @@ async def _vertex_live_setup_denial( valid_token: UserAPIKeyAuth, llm_router: Router | None, authorize_model: VertexLiveModelAuthorizer, -) -> ProxyException | None: +) -> ClientFrameRefusal | None: match setup_model: case str(): denial: Final = await _vertex_live_model_denial(setup_model, valid_token, llm_router, authorize_model) - if denial is not None: + if isinstance(denial, ProxyException): log_model_access_denial(denial) return denial case _UnreadableVertexLiveFrame(): @@ -4160,7 +4165,7 @@ def _build_vertex_live_client_frame_gate( which is a caller with no model allowlist and no agent, end user or team grant standing in for one """ - async def gate(frame_data: str | bytes, valid_token: UserAPIKeyAuth, /) -> ProxyException | None: + async def gate(frame_data: str | bytes, valid_token: UserAPIKeyAuth, /) -> ClientFrameRefusal | None: model: Final = _vertex_live_setup_model(frame_data) if model is None: return None diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index c1f733c738c..e2d6ce4a63e 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -10,7 +10,7 @@ from dataclasses import dataclass from datetime import datetime from itertools import chain, count, groupby from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, TypedDict, cast, overload from urllib.parse import urlencode, urlparse import httpx @@ -2395,28 +2395,39 @@ def create_websocket_passthrough_route( return websocket_endpoint_func -def _rewrite_vertex_live_setup_model(text_data: str, setup_model_rewriter: Callable[[str], str] | None) -> str: +@overload +def _rewrite_vertex_live_setup_model(frame_data: str, setup_model_rewriter: Callable[[str], str] | None) -> str: ... + + +@overload +def _rewrite_vertex_live_setup_model(frame_data: bytes, setup_model_rewriter: Callable[[str], str] | None) -> bytes: ... + + +def _rewrite_vertex_live_setup_model( + frame_data: str | bytes, setup_model_rewriter: Callable[[str], str] | None +) -> str | bytes: """ Rewrite the model of a Vertex AI Live ``setup`` frame, leaving every other frame byte-identical """ if setup_model_rewriter is None: - return text_data + return frame_data try: - message: Final[object] = _DECODED_JSON.validate_python(json.loads(text_data)) - except json.JSONDecodeError: - return text_data + message: Final[object] = _DECODED_JSON.validate_python(json.loads(frame_data)) + except (json.JSONDecodeError, UnicodeDecodeError): + return frame_data if not isinstance(message, dict): - return text_data + return frame_data setup: Final = message.get("setup") if not isinstance(setup, dict): - return text_data + return frame_data setup_model: Final = setup.get("model") if not isinstance(setup_model, str): - return text_data + return frame_data rewritten_model: Final = setup_model_rewriter(setup_model) if rewritten_model == setup_model: - return text_data - return json.dumps({**message, "setup": {**setup, "model": rewritten_model}}) + return frame_data + rewritten_frame: Final = json.dumps({**message, "setup": {**setup, "model": rewritten_model}}) + return rewritten_frame if isinstance(frame_data, str) else rewritten_frame.encode("utf-8") def _resolved_vertex_live_setup( @@ -2504,15 +2515,32 @@ def _with_trace_context(headers: Mapping[str, str], parent_span: object) -> dict return inject_trace_context(headers, parent_span=parent_span) +@dataclass(frozen=True, slots=True) +class LoggedClientFrameRefusal: + denial: ProxyException + + +ClientFrameRefusal: TypeAlias = ProxyException | LoggedClientFrameRefusal + + +def _client_frame_denial(refusal: ClientFrameRefusal) -> ProxyException: + match refusal: + case LoggedClientFrameRefusal(denial=denial): + return denial + case ProxyException(): + return refusal + + class WebsocketClientFrameGate(Protocol): - """Authorize one client frame, returning the denial to report or ``None`` to let it through. + """Authorize one client frame, returning the denial to report, one its check already reported, or ``None`` to + let it through. A websocket passthrough authenticates before any frame arrives, so a protocol that names its model inside a frame has no model to authorize at connect time. The frame, not the extracted model, is the argument because where a model sits is part of each provider's wire protocol """ - async def __call__(self, frame_data: str | bytes, valid_token: UserAPIKeyAuth, /) -> ProxyException | None: ... + async def __call__(self, frame_data: str | bytes, valid_token: UserAPIKeyAuth, /) -> ClientFrameRefusal | None: ... async def _relay_client_frame( @@ -2521,7 +2549,7 @@ async def _relay_client_frame( forwarded_frame: str | bytes, gate: WebsocketClientFrameGate | None, valid_token: UserAPIKeyAuth, -) -> ProxyException | None: +) -> ClientFrameRefusal | None: """Forward one client frame upstream, or refuse it and hand back the denial that ended the session. The denial travels out so the caller reports a failed session instead of leaving the success hooks to @@ -2693,7 +2721,7 @@ async def websocket_passthrough_request( "WebSocket passthrough (%s): Upstream connection established successfully", endpoint ) - async def forward_client_to_upstream() -> ProxyException | None: + async def forward_client_to_upstream() -> ClientFrameRefusal | None: """Forward messages from client to upstream WebSocket, returning any denial that ended it""" try: while True: @@ -2778,7 +2806,7 @@ async def websocket_passthrough_request( bytes_refusal := await _relay_client_frame( upstream_ws, bytes_data, - bytes_data, + _rewrite_vertex_live_setup_model(bytes_data, setup_model_rewriter), client_frame_gate, user_api_key_dict, ) @@ -2862,15 +2890,18 @@ async def websocket_passthrough_request( # upstream reader's next send fail, so raising that instead would report the wrong cause. # Turns that completed before a refusal were served, so they are logged and billed as a success client_frame_refusal: Final = next( - (result for result in task_results if isinstance(result, ProxyException)), None + (result for result in task_results if isinstance(result, (ProxyException, LoggedClientFrameRefusal))), + None, ) if client_frame_refusal is not None: await websocket.close( - code=CloseCode.POLICY_VIOLATION, reason=_refusal_close_reason(client_frame_refusal.message) + code=CloseCode.POLICY_VIOLATION, + reason=_refusal_close_reason(_client_frame_denial(client_frame_refusal).message), ) await upstream_ws.close() - if client_frame_refusal is not None and not websocket_messages: + if isinstance(client_frame_refusal, ProxyException) and not websocket_messages: await log_failure(client_frame_refusal) + if client_frame_refusal is not None and not websocket_messages: return # Check for exceptions in completed tasks diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 22c109acb63..dc93a294fb4 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -7342,7 +7342,7 @@ def _patched_websocket_passthrough_environment(upstream_ws): ): mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_success_hook = AsyncMock() - mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None) enqueued_handler = MagicMock() mock_worker.ensure_initialized_and_enqueue = MagicMock( side_effect=lambda async_coroutine: ( @@ -7703,6 +7703,19 @@ async def test_vertex_live_setup_frame_is_gated_when_it_arrives_as_binary(): assert _policy_close_reason(websocket) is not None +@pytest.mark.asyncio +async def test_vertex_live_binary_setup_frame_reaches_vertex_as_the_model_the_gate_authorized() -> None: + binary_setup: Final = json.dumps({"setup": {"model": f"publishers/google/models/{VERTEX_LIVE_ALIAS_GROUP}"}}) + _, upstream_ws = await _run_vertex_live_gated_passthrough( + [{"type": "websocket.receive", "bytes": binary_setup.encode()}], + UserAPIKeyAuth(token="hashed", models=[VERTEX_LIVE_ALIAS_GROUP]), + ) + + sent_frame: Final = upstream_ws.send.await_args.args[0] + assert isinstance(sent_frame, bytes) + assert json.loads(sent_frame)["setup"]["model"] == VERTEX_LIVE_FULL_PATH.format(model=VERTEX_LIVE_GROUP) + + @pytest.mark.asyncio async def test_vertex_live_setup_frame_with_an_escaped_key_is_still_gated(): """Large audio frames skip the parse on a substring test, so a JSON-escaped key must not slip past it""" @@ -8186,13 +8199,22 @@ async def test_vertex_live_setup_frame_naming_an_alias_cannot_open_the_model_the @pytest.mark.asyncio -@pytest.mark.parametrize("hidden", [False, True], ids=["listed-alias", "hidden-alias"]) -async def test_vertex_live_setup_frame_naming_a_router_alias_cannot_open_the_model_the_alias_shadows(hidden): - """A listed router alias runs as its target group, while the setup rewriter cannot see a hidden one and sends it - to Vertex as the model whose id it reuses, which a key holding only the target must not open""" - llm_router = litellm.Router( +@pytest.mark.parametrize( + ("hidden", "shadows_a_group"), + [(False, False), (True, False), (False, True), (True, True)], + ids=["listed-alias", "hidden-alias", "listed-alias-named-like-a-group", "hidden-alias-named-like-a-group"], +) +async def test_vertex_live_setup_frame_naming_a_router_alias_cannot_open_the_model_the_alias_shadows( + hidden: bool, shadows_a_group: bool +) -> None: + """A listed router alias runs as its target group, while the setup rewriter cannot see a hidden one, or one named + like a group, and sends it to Vertex as the model whose id it reuses, which a key holding only the target must + not open""" + shadowed_group: Final = {"model_name": "gemini-3-pro-live", "litellm_params": {"model": "vertex_ai/gemini-3-pro-live"}} + llm_router: Final = litellm.Router( model_list=[ {"model_name": "gemini-live-2.5-flash", "litellm_params": {"model": "vertex_ai/gemini-live-2.5-flash"}}, + *([shadowed_group] if shadows_a_group else []), ], model_group_alias={"gemini-3-pro-live": {"model": "gemini-live-2.5-flash", "hidden": hidden}}, ) @@ -8205,7 +8227,11 @@ async def test_vertex_live_setup_frame_naming_a_router_alias_cannot_open_the_mod sent_models = [json.loads(call.args[0])["setup"]["model"] for call in upstream_ws.send.await_args_list] assert not any(model.endswith("/gemini-3-pro-live") for model in sent_models) - assert sent_models == ([] if hidden else ["projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash"]) + assert sent_models == ( + [] + if hidden or shadows_a_group + else ["projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash"] + ) @pytest.mark.asyncio @@ -8416,20 +8442,50 @@ async def test_vertex_live_route_authorizes_the_setup_model_the_way_connect_time param="model", code=400, ) - with patch( - "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True, side_effect=budget_denial - ) as reauthorization: - websocket, upstream_ws, _ = await _run_vertex_live_route( + valid_token: Final = UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]) + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.user_api_key_auth_builder", autospec=True, return_value=valid_token + ) as reauthentication, + patch( + "litellm.proxy.auth.user_api_key_auth.run_centralized_common_checks", + autospec=True, + side_effect=budget_denial, + ), + ): + websocket, upstream_ws, environment = await _run_vertex_live_route( [_setup_frame("models/gemini-live-2.5-flash")], - UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), + valid_token, reauthorize_through_user_api_key_auth=True, ) upstream_ws.send.assert_not_awaited() assert _policy_close_reason(websocket) == budget_denial.message - reauthorized_request = reauthorization.call_args.kwargs["request"] - assert json.loads(await reauthorized_request.body()) == {"model": "gemini-live-2.5-flash"} - assert reauthorization.call_args.kwargs["api_key"] == "Bearer sk-live-client" + reauthenticated_request: Final = reauthentication.call_args.kwargs["request"] + assert json.loads(await reauthenticated_request.body()) == {"model": "gemini-live-2.5-flash"} + assert reauthentication.call_args.kwargs["api_key"] == "Bearer sk-live-client" + environment.proxy_logging.post_call_failure_hook.assert_awaited_once() + reported: Final = environment.proxy_logging.post_call_failure_hook.await_args.kwargs + assert reported["original_exception"].message == budget_denial.message + assert reported["request_data"]["model"] == "gemini-live-2.5-flash" + + +@pytest.mark.asyncio +async def test_vertex_live_route_releases_the_budget_it_reserved_to_authorize_an_allowed_setup_model(): + """ + Re-authorizing the setup model reserves the request's worst-case cost the way connect-time auth does, and the + session's spend is settled on the connect-time token, so an unreleased reservation pins the key's spend counter + and throttles its later sessions until the counter expires + """ + valid_token: Final = UserAPIKeyAuth( + token="hashed", + models=["gemini-live-2.5-flash"], + budget_reservation={"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False}, + ) + _, upstream_ws, _ = await _run_vertex_live_route([_setup_frame("models/gemini-live-2.5-flash")], valid_token) + + upstream_ws.send.assert_awaited_once() + assert valid_token.budget_reservation is not None and valid_token.budget_reservation["finalized"] is True @pytest.mark.asyncio From 0c9ce85e40febb4fc9f17d894763b54a32298f4a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 03:28:24 -0700 Subject: [PATCH 11/17] fix(proxy): price Vertex Live sessions opened with a binary setup frame at their setup model A binary setup frame is gated and rewritten like a text one, so the session it opens is served, but only text frames fed the model the session is priced at, and a binary-opened session logged model unknown with no spend. Both frame kinds now go through one setup-model reader, which reuses the JSON object frame parser the upstream side already uses --- .../pass_through_endpoints.py | 83 ++++++------------- .../test_pass_through_endpoints.py | 30 +++++++ 2 files changed, 56 insertions(+), 57 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index e2d6ce4a63e..59fb2e842a3 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2430,6 +2430,16 @@ def _rewrite_vertex_live_setup_model( return rewritten_frame if isinstance(frame_data, str) else rewritten_frame.encode("utf-8") +def _vertex_live_client_setup_model( + frame_data: str | bytes, setup_model_rewriter: Callable[[str], str] | None +) -> str | None: + message: Final = _json_object_frame(frame_data) + setup: Final = None if message is None else message.get("setup") + if not isinstance(setup, dict) or "model" not in setup: + return None + return _extract_model_from_vertex_ai_setup(_resolved_vertex_live_setup(setup, setup_model_rewriter)) + + def _resolved_vertex_live_setup( setup_data: Mapping[str, object], setup_model_rewriter: Callable[[str], str] | None ) -> Mapping[str, object]: @@ -2721,6 +2731,20 @@ async def websocket_passthrough_request( "WebSocket passthrough (%s): Upstream connection established successfully", endpoint ) + is_vertex_live: Final = bool(endpoint and "/vertex_ai/live" in endpoint) + + def record_client_setup_model(frame_data: str | bytes) -> None: + if not is_vertex_live: + return + extracted_model: Final = _vertex_live_client_setup_model(frame_data, setup_model_rewriter) + if not extracted_model: + return + kwargs["model"] = extracted_model + kwargs["custom_llm_provider"] = "vertex_ai-language-models" + logging_obj.model = extracted_model + logging_obj.model_call_details["model"] = extracted_model + logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai" + async def forward_client_to_upstream() -> ClientFrameRefusal | None: """Forward messages from client to upstream WebSocket, returning any denial that ended it""" try: @@ -2745,62 +2769,7 @@ async def websocket_passthrough_request( ) ) is not None: return text_refusal - # Try to extract model from client setup message for Vertex AI Live - if endpoint and "/vertex_ai/live" in endpoint: - verbose_proxy_logger.debug( - "WebSocket passthrough (%s): Processing client message for model extraction", - endpoint, - ) - try: - client_message = _DECODED_JSON.validate_python(json.loads(text_data)) - if isinstance(client_message, dict) and "setup" in client_message: - setup_data = client_message["setup"] - verbose_proxy_logger.debug( - "WebSocket passthrough (%s): Found setup data in client message: %s", - endpoint, - setup_data, - ) - if isinstance(setup_data, dict) and "model" in setup_data: - extracted_model = _extract_model_from_vertex_ai_setup( - _resolved_vertex_live_setup(setup_data, setup_model_rewriter) - ) - if extracted_model: - kwargs["model"] = extracted_model - kwargs["custom_llm_provider"] = "vertex_ai-language-models" - # Update logging object with correct model - logging_obj.model = extracted_model - logging_obj.model_call_details["model"] = extracted_model - logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai" - verbose_proxy_logger.info( - "WebSocket passthrough (%s): Successfully extracted model '%s' and set provider to 'vertex_ai' from client setup message", - endpoint, - extracted_model, - ) - else: - verbose_proxy_logger.warning( - "WebSocket passthrough (%s): Failed to extract model from client setup data: %s", - endpoint, - setup_data, - ) - else: - verbose_proxy_logger.debug( - "WebSocket passthrough (%s): Setup data does not contain model field: %s", - endpoint, - setup_data, - ) - else: - verbose_proxy_logger.debug( - "WebSocket passthrough (%s): Client message does not contain setup data", - endpoint, - ) - except (json.JSONDecodeError, KeyError, TypeError) as e: - verbose_proxy_logger.debug( - "WebSocket passthrough (%s): Client message is not a valid setup message: %s", - endpoint, - e, - ) - # Not a JSON message or doesn't contain setup data - + record_client_setup_model(text_data) elif bytes_data is not None: if ( bytes_refusal := await _relay_client_frame( @@ -2812,6 +2781,7 @@ async def websocket_passthrough_request( ) ) is not None: return bytes_refusal + record_client_setup_model(bytes_data) except asyncio.CancelledError: raise except Exception: @@ -2835,7 +2805,6 @@ async def websocket_passthrough_request( logging_obj.model_call_details["model"] = extracted_model logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai_language_models" - is_vertex_live: Final = bool(endpoint and "/vertex_ai/live" in endpoint) json_frame_ordinal: Final = count() async def relay_upstream_frame(upstream_message: str | bytes) -> None: diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index dc93a294fb4..2d7264c1021 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -8575,6 +8575,36 @@ async def test_vertex_live_turns_served_before_a_refused_frame_are_still_billed( environment.proxy_logging.post_call_failure_hook.assert_not_awaited() +@pytest.mark.asyncio +async def test_vertex_live_session_opened_with_a_binary_setup_frame_is_billed_at_its_setup_model(): + """ + A setup frame sent as binary is gated and relayed like a text one, so the session it opens must be priced at + the setup model too instead of logging as model unknown with no spend + """ + upstream_ws = ScriptedUpstreamWebSocket( + ( + json.dumps({"setupComplete": {}}), + json.dumps({"serverContent": {"turnComplete": True}, "usageMetadata": {"totalTokenCount": 60}}), + ) + ) + + async def client_frames() -> AsyncIterator[dict[str, str | bytes]]: + yield {"type": "websocket.receive", "bytes": json.dumps({"setup": {"model": "gemini-live-2.5-flash"}}).encode()} + await upstream_ws.sent_every_frame.wait() + yield {"type": "websocket.disconnect"} + + _, environment = await _run_vertex_live_gated_session( + _client_websocket(client_frames().__anext__), + upstream_ws, + UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), + _vertex_live_router(), + ) + + upstream_ws.send.assert_awaited_once() + environment.logging_worker.ensure_initialized_and_enqueue.assert_called_once() + assert environment.enqueued_handler.call_args.kwargs["logging_obj"].model == "gemini-live-2.5-flash" + + class StreamingUpstreamWebSocket(ScriptedUpstreamWebSocket): """An upstream that sends the given frames in order, then keeps streaming model output until it is closed""" From b4f69f1590e641795d40a6abf1e1d63407b11143 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 04:57:48 -0700 Subject: [PATCH 12/17] fix(proxy): authorize Vertex Live router aliases and vertex_ai ids like HTTP requests --- .../llm_passthrough_endpoints.py | 94 ++++++++++++------- .../pass_through_endpoints.py | 4 +- .../test_pass_through_endpoints.py | 65 +++++++++++-- 3 files changed, 123 insertions(+), 40 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index de5d3047989..23650ded311 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -3923,7 +3923,7 @@ def _vertex_live_setup_model(frame_data: str | bytes) -> str | _UnreadableVertex frame: Final = _VERTEX_LIVE_CLIENT_FRAME.validate_python( json.loads(frame_data, object_pairs_hook=_object_with_unique_keys) ) - except ValueError: + except (ValueError, RecursionError): return _UnreadableVertexLiveFrame() setup: Final = frame.get("setup") return None if setup is None else setup.get("model") or None @@ -3958,10 +3958,23 @@ def _vertex_upstream_model(deployment: DeploymentTypedDict) -> str | None: return upstream if provider is None or provider in VERTEX_LIVE_PROVIDERS else None +def _routed_deployments(llm_router: Router) -> tuple[DeploymentTypedDict, ...]: + """Every deployment under the name a request reaches it by, which leaves out a group a router alias shares a name + with, since the alias wins""" + return ( + *llm_router.get_model_list_from_model_alias(), + *( + deployment + for deployment in llm_router.get_model_list() or () + if deployment["model_name"] not in llm_router.model_group_alias + ), + ) + + def _vertex_deployments(llm_router: Router | None) -> tuple[DeploymentTypedDict, ...]: if llm_router is None: return () - return tuple(deployment for deployment in (llm_router.get_model_list() or ()) if _vertex_upstream_model(deployment)) + return tuple(deployment for deployment in _routed_deployments(llm_router) if _vertex_upstream_model(deployment)) def _live_deployment(setup_model: str, llm_router: Router | None) -> DeploymentTypedDict | None: @@ -3981,19 +3994,15 @@ def _live_deployment(setup_model: str, llm_router: Router | None) -> DeploymentT def _vertex_live_model_names(setup_model: str, llm_router: Router | None) -> tuple[str, ...]: """ The names a key can hold for the model a ``setup`` frame runs: the Vertex AI group the setup rewriter sends it - through, else every Vertex AI group serving the publisher model Vertex runs, and that model's id bare or under - LiteLLM's ``vertex_ai/`` prefix unless a group holds that name for another model. + through, else every Vertex AI group serving the publisher model Vertex runs, and either way that model's id bare + or under LiteLLM's ``vertex_ai/`` prefix unless a group holds that name. A full resource path skips the rewriter and reaches Vertex as sent, so only its publisher model counts, and a wildcard is matched against that id alone so addressing in front of it cannot satisfy the pattern """ deployment: Final = _live_deployment(setup_model, llm_router) - if deployment is not None and _vertex_upstream_model(deployment) is not None: - group: Final = deployment["model_name"] - alias_target: Final = ( - None if llm_router is None or group in llm_router.model_names else llm_router.get_model_from_alias(group) - ) - return (alias_target or group,) + if deployment is not None and (upstream := _vertex_upstream_model(deployment)) is not None: + return (deployment["model_name"], *_publisher_model_id_names(upstream, llm_router)) opened_model: Final = _resolve_alias_to_upstream_model(setup_model, llm_router) publisher_model: Final = _vertex_publisher_model_suffix(opened_model) serving_groups: Final = tuple( @@ -4003,22 +4012,31 @@ def _vertex_live_model_names(setup_model: str, llm_router: Router | None) -> tup if _vertex_publisher_model_suffix(_vertex_upstream_model(candidate) or "") == publisher_model ) ) + return (*serving_groups, *_publisher_model_id_names(opened_model, llm_router)) + + +def _publisher_model_id_names(model: str, llm_router: Router | None) -> tuple[str, ...]: group_names: Final = frozenset( candidate["model_name"] for candidate in (() if llm_router is None else llm_router.get_model_list() or ()) ) - publisher_model_id: Final = opened_model.rsplit("/", 1)[-1] - return ( - *serving_groups, - *(name for name in (publisher_model_id, f"vertex_ai/{publisher_model_id}") if name not in group_names), - ) + publisher_model_id: Final = model.rsplit("/", 1)[-1] + return tuple(name for name in (publisher_model_id, f"vertex_ai/{publisher_model_id}") if name not in group_names) -def _is_model_alias(name: str, valid_token: UserAPIKeyAuth, llm_router: Router | None) -> bool: +def _is_listed_router_alias(name: str, llm_router: Router) -> bool: + return any(deployment["model_name"] == name for deployment in llm_router.get_model_list_from_model_alias()) + + +def _is_alias_the_rewriter_ignores(name: str, valid_token: UserAPIKeyAuth, llm_router: Router | None) -> bool: return ( name in (key_model_aliases_for_auth_check(valid_token) or ()) or name in (team_model_aliases_for_auth_check(valid_token) or ()) or name in litellm.model_alias_map - or (llm_router is not None and name in llm_router.model_group_alias) + or ( + llm_router is not None + and name in llm_router.model_group_alias + and not _is_listed_router_alias(name, llm_router) + ) ) @@ -4070,14 +4088,30 @@ async def _model_reauthorization_denial( return None -async def _first_reauthorization_denial( - model_names: tuple[str, ...], authorize_model: VertexLiveModelAuthorizer -) -> LoggedClientFrameRefusal | None: - denial: Final = await _model_reauthorization_denial(model_names[0], authorize_model) +async def _first_model_denial( + model_names: tuple[str, ...], + valid_token: UserAPIKeyAuth, + llm_router: Router | None, + authorize_model: VertexLiveModelAuthorizer, +) -> ClientFrameRefusal | None: + """ + The first name the key may call that also passes re-authorization opens the session, and no later name is + checked. Otherwise the refusal is the first re-authorization denial, which says more than an access denial + """ + name: Final = model_names[0] + denial: Final = await _model_access_denial(name, valid_token, llm_router) or await _model_reauthorization_denial( + name, authorize_model + ) if denial is None or len(model_names) == 1: return denial - later_denial: Final = await _first_reauthorization_denial(model_names[1:], authorize_model) - return None if later_denial is None else denial + later_denial: Final = await _first_model_denial(model_names[1:], valid_token, llm_router, authorize_model) + if later_denial is None: + return None + return ( + later_denial + if isinstance(later_denial, LoggedClientFrameRefusal) and isinstance(denial, ProxyException) + else denial + ) async def _vertex_live_model_denial( @@ -4087,14 +4121,14 @@ async def _vertex_live_model_denial( authorize_model: VertexLiveModelAuthorizer, ) -> ClientFrameRefusal | None: """ - The setup rewriter forwards an alias name as sent unless it is a listed router alias, which already resolved to - its target group, so an alias name grants nothing here. A name the key may call is then authorized the way a - request naming it at connect time is + The setup rewriter runs a listed router alias as its target group, the way a request naming it does, and sends + any other alias name to Vertex as written, so only a listed router alias grants anything here. A name the key may + call is then authorized the way a request naming it at connect time is """ model_names: Final = tuple( name for name in _vertex_live_model_names(setup_model, llm_router) - if not _is_model_alias(name, valid_token, llm_router) + if not _is_alias_the_rewriter_ignores(name, valid_token, llm_router) ) if not model_names: return ( @@ -4108,11 +4142,7 @@ async def _vertex_live_model_denial( if await _caller_has_model_limits(valid_token) else None ) - denials: Final = tuple([await _model_access_denial(name, valid_token, llm_router) for name in model_names]) - callable_names: Final = tuple(name for name, denial in zip(model_names, denials, strict=True) if denial is None) - if not callable_names: - return denials[0] - return await _first_reauthorization_denial(callable_names, authorize_model) + return await _first_model_denial(model_names, valid_token, llm_router, authorize_model) async def _caller_has_model_limits(valid_token: UserAPIKeyAuth) -> bool: diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 59fb2e842a3..872e9b8e9cc 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2413,7 +2413,7 @@ def _rewrite_vertex_live_setup_model( return frame_data try: message: Final[object] = _DECODED_JSON.validate_python(json.loads(frame_data)) - except (json.JSONDecodeError, UnicodeDecodeError): + except (json.JSONDecodeError, UnicodeDecodeError, RecursionError): return frame_data if not isinstance(message, dict): return frame_data @@ -2461,7 +2461,7 @@ def _json_object_frame(frame: str | bytes) -> dict[str, object] | None: decoded: Final[object] = _DECODED_JSON.validate_python( json.loads(frame if isinstance(frame, str) else frame.decode("utf-8")) ) - except (json.JSONDecodeError, UnicodeDecodeError): + except (json.JSONDecodeError, UnicodeDecodeError, RecursionError): return None return decoded if isinstance(decoded, dict) else None diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 2d7264c1021..2b71e6dfc96 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -7754,6 +7754,18 @@ async def test_vertex_live_setup_frame_only_vertex_can_parse_is_refused(lenient_ assert _policy_close_reason(websocket) is not None +@pytest.mark.asyncio +async def test_vertex_live_setup_frame_nested_too_deep_to_parse_is_refused() -> None: + nested_setup: Final = '{"setup": ' + "[" * 100_000 + "]" * 100_000 + "}" + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [{"type": "websocket.receive", "text": nested_setup}], + UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), + ) + + upstream_ws.send.assert_not_awaited() + assert _policy_close_reason(websocket) is not None + + @pytest.mark.asyncio @pytest.mark.parametrize( "lenient_setup", @@ -8207,9 +8219,9 @@ async def test_vertex_live_setup_frame_naming_an_alias_cannot_open_the_model_the async def test_vertex_live_setup_frame_naming_a_router_alias_cannot_open_the_model_the_alias_shadows( hidden: bool, shadows_a_group: bool ) -> None: - """A listed router alias runs as its target group, while the setup rewriter cannot see a hidden one, or one named - like a group, and sends it to Vertex as the model whose id it reuses, which a key holding only the target must - not open""" + """A listed router alias runs as its target group even when a group shares its name, while the setup rewriter + cannot see a hidden one and sends it to Vertex as the model whose id it reuses, which a key holding only the + target must not open""" shadowed_group: Final = {"model_name": "gemini-3-pro-live", "litellm_params": {"model": "vertex_ai/gemini-3-pro-live"}} llm_router: Final = litellm.Router( model_list=[ @@ -8228,12 +8240,53 @@ async def test_vertex_live_setup_frame_naming_a_router_alias_cannot_open_the_mod sent_models = [json.loads(call.args[0])["setup"]["model"] for call in upstream_ws.send.await_args_list] assert not any(model.endswith("/gemini-3-pro-live") for model in sent_models) assert sent_models == ( - [] - if hidden or shadows_a_group - else ["projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash"] + [] if hidden else ["projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash"] ) +@pytest.mark.asyncio +@pytest.mark.parametrize("held_model", ["gemini-3-pro-live", "gemini-live-2.5-flash"], ids=["alias", "alias-target"]) +async def test_vertex_live_setup_frame_naming_a_router_alias_named_like_a_group_is_authorized_like_a_request( + held_model: str, +) -> None: + """A request naming a router alias reaches its target even when a group shares the alias name, so a key holding + the alias or its target runs the target on Live too, and still cannot reach the group the alias shadows""" + llm_router: Final = litellm.Router( + model_list=[ + {"model_name": "gemini-live-2.5-flash", "litellm_params": {"model": "vertex_ai/gemini-live-2.5-flash"}}, + {"model_name": "gemini-3-pro-live", "litellm_params": {"model": "vertex_ai/gemini-3-pro-live"}}, + ], + model_group_alias={"gemini-3-pro-live": "gemini-live-2.5-flash"}, + ) + key: Final = UserAPIKeyAuth(token="hashed", models=[held_model]) + + _, alias_upstream = await _run_vertex_live_gated_passthrough( + [_setup_frame("models/gemini-3-pro-live")], key, llm_router=llm_router + ) + _, shadowed_upstream = await _run_vertex_live_gated_passthrough( + [_setup_frame(VERTEX_LIVE_FULL_PATH.format(model="gemini-3-pro-live"))], key, llm_router=llm_router + ) + + sent_models = [json.loads(call.args[0])["setup"]["model"] for call in alias_upstream.send.await_args_list] + assert sent_models == [VERTEX_LIVE_FULL_PATH.format(model="gemini-live-2.5-flash")] + shadowed_upstream.send.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS) +async def test_vertex_live_setup_frame_is_authorized_by_the_litellm_model_id_in_every_addressing(addressing): + """A request naming a deployment's ``vertex_ai/`` id reaches that deployment, so a key holding the id runs the + model it names on Live whichever addressing the setup frame uses""" + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame(addressing.format(model=VERTEX_LIVE_GROUP))], + UserAPIKeyAuth(token="hashed", models=[f"vertex_ai/{VERTEX_LIVE_GROUP}"]), + ) + + assert _policy_close_reason(websocket) is None + sent_setup = json.loads(upstream_ws.send.await_args.args[0])["setup"] + assert sent_setup["model"] == VERTEX_LIVE_FULL_PATH.format(model=VERTEX_LIVE_GROUP) + + @pytest.mark.asyncio async def test_vertex_live_setup_frame_is_checked_against_the_owner_of_a_personal_key(): """A personal key's owner can be limited to fewer models than the key lists, and auth enforces that limit too""" From 00fcc96002b5c43b91081d88b8d168ef4240cd3b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 06:11:52 -0700 Subject: [PATCH 13/17] fix(proxy): keep a Vertex Live setup group's budget and pricing answer when the key may call it Every name a setup frame matches runs the same deployment, so the first name the key may call is re-authorized and its denial stands instead of falling through to the bare model id, which no group's budget or pricing check covers --- .../llm_passthrough_endpoints.py | 22 ++++++--------- .../test_pass_through_endpoints.py | 27 +++++++++++++++++++ 2 files changed, 35 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 23650ded311..8ffc6f04c1b 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -4095,23 +4095,17 @@ async def _first_model_denial( authorize_model: VertexLiveModelAuthorizer, ) -> ClientFrameRefusal | None: """ - The first name the key may call that also passes re-authorization opens the session, and no later name is - checked. Otherwise the refusal is the first re-authorization denial, which says more than an access denial + Every name runs the same deployment, so the first name the key may call is re-authorized and that answer + stands, budget and pricing checks included. Only a key that may call none of them gets the first access denial """ name: Final = model_names[0] - denial: Final = await _model_access_denial(name, valid_token, llm_router) or await _model_reauthorization_denial( - name, authorize_model - ) - if denial is None or len(model_names) == 1: - return denial + access_denial: Final = await _model_access_denial(name, valid_token, llm_router) + if access_denial is None: + return await _model_reauthorization_denial(name, authorize_model) + if len(model_names) == 1: + return access_denial later_denial: Final = await _first_model_denial(model_names[1:], valid_token, llm_router, authorize_model) - if later_denial is None: - return None - return ( - later_denial - if isinstance(later_denial, LoggedClientFrameRefusal) and isinstance(denial, ProxyException) - else denial - ) + return access_denial if isinstance(later_denial, ProxyException) else later_denial async def _vertex_live_model_denial( diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 2b71e6dfc96..0db2500bec0 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -8541,6 +8541,33 @@ async def test_vertex_live_route_releases_the_budget_it_reserved_to_authorize_an assert valid_token.budget_reservation is not None and valid_token.budget_reservation["finalized"] is True +@pytest.mark.asyncio +async def test_vertex_live_setup_group_the_key_may_call_but_fails_budget_is_not_opened_under_its_bare_model_id(): + """ + A setup frame naming a group runs that group's deployment whichever name the key holds, so when the key may call + the group its budget and pricing checks decide, and the group's bare model id cannot open the session instead + """ + group_denial: Final = ProxyException( + message=f"Budget has been exceeded for model {VERTEX_LIVE_GROUP}.", + type=ProxyErrorTypes.budget_exceeded, + param="model", + code=400, + ) + valid_token: Final = UserAPIKeyAuth(token="hashed") + + async def authorize_model(model: str) -> UserAPIKeyAuth: + if model == VERTEX_LIVE_GROUP: + raise group_denial + return valid_token + + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame(VERTEX_LIVE_GROUP)], valid_token, authorize_model=authorize_model + ) + + upstream_ws.send.assert_not_awaited() + assert _policy_close_reason(websocket) == group_denial.message + + @pytest.mark.asyncio async def test_vertex_live_route_keeps_the_key_allowlist_out_of_the_close_reason_and_in_the_proxy_log(caplog): """The client learns which model was refused, while only the proxy log names the models the key may reach""" From 3e8a1795752d98c761875522294e543d1ca469e0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 06:57:20 -0700 Subject: [PATCH 14/17] test(proxy): type the Vertex Live gate tests' parameters --- .../test_pass_through_endpoints.py | 28 +++++++++---------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 0db2500bec0..3aec2ac5d31 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -7738,7 +7738,7 @@ async def test_vertex_live_setup_frame_with_an_escaped_key_is_still_gated(): '{setup: {model: "totally-unlisted-model-abc"}}', ], ) -async def test_vertex_live_setup_frame_only_vertex_can_parse_is_refused(lenient_setup, frame_kind): +async def test_vertex_live_setup_frame_only_vertex_can_parse_is_refused(lenient_setup: str, frame_kind: str): """Vertex reads trailing commas, single quotes and bare keys, which strict JSON refuses to parse""" websocket, upstream_ws = await _run_vertex_live_gated_passthrough( [ @@ -7796,7 +7796,7 @@ async def test_vertex_live_setup_frame_only_vertex_can_parse_still_reaches_verte '{"setup": {"model": "gemini-3-pro-live"}, "setup": null}', ], ) -async def test_vertex_live_setup_frame_repeating_a_key_as_null_is_refused(repeated_key_setup): +async def test_vertex_live_setup_frame_repeating_a_key_as_null_is_refused(repeated_key_setup: str): """ Vertex keeps the first of two keys when the second is null, while ``json.loads`` keeps the last, so the proxy would read no model where Vertex runs the first one @@ -7822,7 +7822,7 @@ async def test_vertex_live_setup_frame_repeating_a_key_as_null_is_refused(repeat ], ids=["litellm-wildcard", "vertex-wildcard", "team", "agent", "end-user"], ) -async def test_vertex_live_setup_frame_only_vertex_can_parse_is_refused_for_any_limited_caller(limited_token): +async def test_vertex_live_setup_frame_only_vertex_can_parse_is_refused_for_any_limited_caller(limited_token: UserAPIKeyAuth): """A wildcard can match any name the proxy might test the caller with, so only a caller with no model limit anywhere may send a frame the proxy cannot read""" websocket, upstream_ws = await _run_vertex_live_gated_passthrough( @@ -7877,7 +7877,7 @@ async def test_vertex_live_non_setup_frames_are_forwarded_without_a_model_check( @pytest.mark.asyncio @pytest.mark.parametrize("addressing", VERTEX_LIVE_GROUP_ADDRESSINGS) @pytest.mark.parametrize("group", [VERTEX_LIVE_GROUP, VERTEX_LIVE_ALIAS_GROUP]) -async def test_vertex_live_setup_frame_authorizes_a_permitted_group_in_every_addressing(addressing, group): +async def test_vertex_live_setup_frame_authorizes_a_permitted_group_in_every_addressing(addressing: str, group: str): """ The Live SDK wraps the caller's model as ``models/``, so a key holding the group has to work whichever addressing its client actually sends @@ -7895,7 +7895,7 @@ async def test_vertex_live_setup_frame_authorizes_a_permitted_group_in_every_add @pytest.mark.asyncio @pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS) -async def test_vertex_live_setup_frame_refuses_an_unpermitted_group_in_every_addressing(addressing): +async def test_vertex_live_setup_frame_refuses_an_unpermitted_group_in_every_addressing(addressing: str): """Normalizing the addressing must not become a way to reach a group the key does not hold""" websocket, upstream_ws = await _run_vertex_live_gated_passthrough( [_setup_frame(addressing.format(model="gemini-3-pro-live"))], @@ -7909,7 +7909,7 @@ async def test_vertex_live_setup_frame_refuses_an_unpermitted_group_in_every_add @pytest.mark.asyncio @pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS) -async def test_vertex_live_setup_frame_naming_no_group_at_all_is_refused_in_every_addressing(addressing): +async def test_vertex_live_setup_frame_naming_no_group_at_all_is_refused_in_every_addressing(addressing: str): """A name that resolves to no group is authorized as it arrived, so an unknown model cannot ride a prefix in""" websocket, upstream_ws = await _run_vertex_live_gated_passthrough( [_setup_frame(addressing.format(model="totally-unlisted-model-abc"))], @@ -7923,7 +7923,7 @@ async def test_vertex_live_setup_frame_naming_no_group_at_all_is_refused_in_ever @pytest.mark.asyncio @pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS) @pytest.mark.parametrize("key_model", ["gemini-4-pro-live-unserved", "vertex_ai/gemini-4-pro-live-unserved"]) -async def test_vertex_live_setup_frame_authorizes_a_model_no_group_serves_by_its_publisher_id(addressing, key_model): +async def test_vertex_live_setup_frame_authorizes_a_model_no_group_serves_by_its_publisher_id(addressing: str, key_model: str): """ With no group to resolve, the setup rewriter sends Vertex the last path segment, so a key holding that publisher model, bare or as a LiteLLM ``vertex_ai/`` id, can open it whichever addressing its client sends @@ -7974,7 +7974,7 @@ async def test_vertex_live_setup_frame_prefers_the_group_its_whole_name_matches( @pytest.mark.asyncio @pytest.mark.parametrize("group", [VERTEX_LIVE_GROUP, VERTEX_LIVE_ALIAS_GROUP]) -async def test_vertex_live_full_path_is_authorized_by_any_group_serving_its_publisher_model(group): +async def test_vertex_live_full_path_is_authorized_by_any_group_serving_its_publisher_model(group: str): """The docs tell callers to send the full resource path, which Vertex runs as sent, so every group serving that publisher model grants it, whatever the group is called""" full_path = VERTEX_LIVE_FULL_PATH.format(model=VERTEX_LIVE_GROUP) @@ -8198,7 +8198,7 @@ async def test_vertex_live_setup_frame_never_reads_the_secret_manager_for_anothe ], ids=["key-alias", "team-alias"], ) -async def test_vertex_live_setup_frame_naming_an_alias_cannot_open_the_model_the_alias_shadows(valid_token): +async def test_vertex_live_setup_frame_naming_an_alias_cannot_open_the_model_the_alias_shadows(valid_token: UserAPIKeyAuth): """ An alias passes the access check by its target, while the setup rewriter sends the name as written, so a key that could set an alias would open the model the alias name belongs to @@ -8274,7 +8274,7 @@ async def test_vertex_live_setup_frame_naming_a_router_alias_named_like_a_group_ @pytest.mark.asyncio @pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS) -async def test_vertex_live_setup_frame_is_authorized_by_the_litellm_model_id_in_every_addressing(addressing): +async def test_vertex_live_setup_frame_is_authorized_by_the_litellm_model_id_in_every_addressing(addressing: str): """A request naming a deployment's ``vertex_ai/`` id reaches that deployment, so a key holding the id runs the model it names on Live whichever addressing the setup frame uses""" websocket, upstream_ws = await _run_vertex_live_gated_passthrough( @@ -8313,7 +8313,7 @@ async def test_vertex_live_setup_frame_is_checked_against_the_owner_of_a_persona @pytest.mark.asyncio @pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS) -async def test_vertex_live_setup_frame_is_checked_against_the_team_in_every_addressing(addressing): +async def test_vertex_live_setup_frame_is_checked_against_the_team_in_every_addressing(addressing: str): """The team allowlist is keyed by group too, so the addressing must not decide whether the team is checked""" websocket, upstream_ws = await _run_vertex_live_gated_passthrough( [_setup_frame(addressing.format(model="gemini-3-pro-live"))], @@ -8395,7 +8395,7 @@ async def test_vertex_live_setup_frame_naming_a_router_alias_is_authorized_by_it @pytest.mark.asyncio @pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS) -async def test_vertex_live_setup_frame_resolves_an_access_group_grant_in_every_addressing(addressing): +async def test_vertex_live_setup_frame_resolves_an_access_group_grant_in_every_addressing(addressing: str): """ ``get_model_access_groups`` is keyed by group as well, and returns nothing for a prefixed name, so a key entitled through an access group rather than a model name needs the same normalization @@ -8569,7 +8569,7 @@ async def test_vertex_live_setup_group_the_key_may_call_but_fails_budget_is_not_ @pytest.mark.asyncio -async def test_vertex_live_route_keeps_the_key_allowlist_out_of_the_close_reason_and_in_the_proxy_log(caplog): +async def test_vertex_live_route_keeps_the_key_allowlist_out_of_the_close_reason_and_in_the_proxy_log(caplog: pytest.LogCaptureFixture): """The client learns which model was refused, while only the proxy log names the models the key may reach""" with caplog.at_level("WARNING", logger="LiteLLM Proxy"): websocket, _, _ = await _run_vertex_live_route( @@ -8697,7 +8697,7 @@ class StreamingUpstreamWebSocket(ScriptedUpstreamWebSocket): return json.dumps({"serverContent": {"modelTurn": {"parts": [{"text": "still talking"}]}}}) -def _client_websocket_with_a_close_handshake(receive) -> MagicMock: +def _client_websocket_with_a_close_handshake(receive: Callable[[], Awaitable[Mapping[str, str | bytes]]]) -> MagicMock: """A client socket that, like Starlette's, refuses sends once its close starts and yields while closing""" from starlette.websockets import WebSocketState From 85cdf88ea635ffcc53951943e3fef31adac5cf89 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 07:21:33 -0700 Subject: [PATCH 15/17] test(proxy): mark the Vertex Live gate tests' locals Final --- .../test_pass_through_endpoints.py | 100 +++++++++--------- 1 file changed, 50 insertions(+), 50 deletions(-) diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 3aec2ac5d31..07eff6c3d08 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -7343,7 +7343,7 @@ def _patched_websocket_passthrough_environment(upstream_ws): mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) mock_proxy_logging.post_call_success_hook = AsyncMock() mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None) - enqueued_handler = MagicMock() + enqueued_handler: Final = MagicMock() mock_worker.ensure_initialized_and_enqueue = MagicMock( side_effect=lambda async_coroutine: ( enqueued_handler(**async_coroutine.cr_frame.f_locals), @@ -7393,14 +7393,14 @@ async def test_websocket_passthrough_relays_as_much_of_a_long_upstream_close_rea from websockets.exceptions import ConnectionClosedError from websockets.frames import Close - upstream_reason = ( + upstream_reason: Final = ( "Quota exceeded for concurrent_sessions. Request more quota for aiplatform.googleapis.com " "BidiGenerateContent sessions in us-central1 from the Google Cloud console quotas page" ) - upstream_ws = ClosingUpstreamWebSocket( + upstream_ws: Final = ClosingUpstreamWebSocket( ConnectionClosedError(rcvd=Close(1011, upstream_reason), sent=Close(1011, ""), rcvd_then_sent=True) ) - websocket = _client_websocket(_pending_receive) + websocket: Final = _client_websocket(_pending_receive) with _patched_websocket_passthrough_environment(upstream_ws): await websocket_passthrough_request( @@ -7413,7 +7413,7 @@ async def test_websocket_passthrough_relays_as_much_of_a_long_upstream_close_rea accept_websocket=False, ) - relayed_reason = websocket.close.await_args.kwargs["reason"] + relayed_reason: Final = websocket.close.await_args.kwargs["reason"] assert upstream_reason.startswith(relayed_reason) assert len(relayed_reason.encode("utf-8")) == WEBSOCKET_CLOSE_REASON_MAX_BYTES @@ -7564,7 +7564,7 @@ async def _run_vertex_live_gated_passthrough( authorize_model: Callable[[str], Awaitable[UserAPIKeyAuth]] | None = None, ) -> tuple[MagicMock, RecordingUpstreamWebSocket]: """Drive /vertex_ai/live with the real frame gate, returning the client and upstream sockets""" - upstream_ws = RecordingUpstreamWebSocket() + upstream_ws: Final = RecordingUpstreamWebSocket() websocket, _ = await _run_vertex_live_gated_session( _client_websocket(AsyncMock(side_effect=[*client_frames, {"type": "websocket.disconnect"}])), upstream_ws, @@ -7636,7 +7636,7 @@ async def test_vertex_live_setup_frame_naming_a_model_the_key_cannot_call_is_ref ) upstream_ws.send.assert_not_awaited() - reason = _policy_close_reason(websocket) + reason: Final = _policy_close_reason(websocket) assert reason is not None and "totally-unlisted-model-abc" in reason assert reason.endswith(".") upstream_ws.close.assert_awaited() @@ -7650,7 +7650,7 @@ async def test_vertex_live_setup_frame_naming_a_permitted_model_still_reaches_ve ) upstream_ws.send.assert_awaited_once() - sent_setup = json.loads(upstream_ws.send.await_args.args[0])["setup"] + sent_setup: Final = json.loads(upstream_ws.send.await_args.args[0])["setup"] assert sent_setup["model"] == "projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash" assert _policy_close_reason(websocket) is None @@ -7682,7 +7682,7 @@ async def test_vertex_live_setup_frame_is_checked_against_the_team_for_an_all_te ) upstream_ws.send.assert_not_awaited() - reason = _policy_close_reason(websocket) + reason: Final = _policy_close_reason(websocket) assert reason is not None and "gemini-3-pro-live" in reason @@ -7888,7 +7888,7 @@ async def test_vertex_live_setup_frame_authorizes_a_permitted_group_in_every_add ) upstream_ws.send.assert_awaited_once() - sent_setup = json.loads(upstream_ws.send.await_args.args[0])["setup"] + sent_setup: Final = json.loads(upstream_ws.send.await_args.args[0])["setup"] assert sent_setup["model"].startswith("projects/") assert _policy_close_reason(websocket) is None @@ -7903,7 +7903,7 @@ async def test_vertex_live_setup_frame_refuses_an_unpermitted_group_in_every_add ) upstream_ws.send.assert_not_awaited() - reason = _policy_close_reason(websocket) + reason: Final = _policy_close_reason(websocket) assert reason is not None and "gemini-3-pro-live" in reason @@ -7934,7 +7934,7 @@ async def test_vertex_live_setup_frame_authorizes_a_model_no_group_serves_by_its ) upstream_ws.send.assert_awaited_once() - sent_setup = json.loads(upstream_ws.send.await_args.args[0])["setup"] + sent_setup: Final = json.loads(upstream_ws.send.await_args.args[0])["setup"] assert sent_setup["model"].endswith("/models/gemini-4-pro-live-unserved") assert _policy_close_reason(websocket) is None @@ -7948,14 +7948,14 @@ async def test_vertex_live_setup_frame_addressing_cannot_satisfy_a_wildcard_for_ ) upstream_ws.send.assert_not_awaited() - reason = _policy_close_reason(websocket) + reason: Final = _policy_close_reason(websocket) assert reason is not None and "forbidden-model" in reason @pytest.mark.asyncio async def test_vertex_live_setup_frame_prefers_the_group_its_whole_name_matches(): """A group named ``vertex_ai/`` is the one the rewriter sends, so it is the one the key must hold""" - llm_router = litellm.Router( + llm_router: Final = litellm.Router( model_list=[ {"model_name": "gemini-3-pro-live", "litellm_params": {"model": "vertex_ai/gemini-3-pro-live"}}, {"model_name": "vertex_ai/gemini-3-pro-live", "litellm_params": {"model": "vertex_ai/gemini-3-pro-live"}}, @@ -7977,7 +7977,7 @@ async def test_vertex_live_setup_frame_prefers_the_group_its_whole_name_matches( async def test_vertex_live_full_path_is_authorized_by_any_group_serving_its_publisher_model(group: str): """The docs tell callers to send the full resource path, which Vertex runs as sent, so every group serving that publisher model grants it, whatever the group is called""" - full_path = VERTEX_LIVE_FULL_PATH.format(model=VERTEX_LIVE_GROUP) + full_path: Final = VERTEX_LIVE_FULL_PATH.format(model=VERTEX_LIVE_GROUP) websocket, upstream_ws = await _run_vertex_live_gated_passthrough( [_setup_frame(full_path)], @@ -7991,7 +7991,7 @@ async def test_vertex_live_full_path_is_authorized_by_any_group_serving_its_publ @pytest.mark.asyncio async def test_vertex_live_full_path_is_not_granted_by_a_group_named_after_it_that_serves_another_model(): - llm_router = litellm.Router( + llm_router: Final = litellm.Router( model_list=[ {"model_name": "gemini-3-pro-live", "litellm_params": {"model": "vertex_ai/gemini-live-2.5-flash"}}, ] @@ -8004,7 +8004,7 @@ async def test_vertex_live_full_path_is_not_granted_by_a_group_named_after_it_th ) upstream_ws.send.assert_not_awaited() - reason = _policy_close_reason(websocket) + reason: Final = _policy_close_reason(websocket) assert reason is not None and "gemini-3-pro-live" in reason @@ -8017,7 +8017,7 @@ async def test_vertex_live_full_path_is_not_granted_by_a_group_named_after_it_th async def test_vertex_live_setup_frame_is_not_granted_by_a_group_another_provider_serves(setup_model: str): """Every Live session runs on the proxy's Vertex credentials, so a group another provider serves under the same model id must not open the Vertex model""" - llm_router = litellm.Router( + llm_router: Final = litellm.Router( model_list=[ { "model_name": "openai-native-audio", @@ -8041,7 +8041,7 @@ async def test_vertex_live_setup_frame_is_not_granted_by_a_group_another_provide async def test_vertex_live_full_path_is_authorized_only_by_a_group_whose_custom_provider_is_vertex( custom_llm_provider: str, reaches_vertex: bool ): - llm_router = litellm.Router( + llm_router: Final = litellm.Router( model_list=[ { "model_name": "native-audio", @@ -8078,7 +8078,7 @@ async def test_vertex_live_full_path_whose_every_name_is_taken_by_another_model_ ): """Groups holding both the bare and the ``vertex_ai/`` id for another model leave no name a key could hold, which says nothing about a key that may call any model""" - llm_router = litellm.Router( + llm_router: Final = litellm.Router( model_list=[ {"model_name": "gemini-3-pro-live", "litellm_params": {"model": "vertex_ai/gemini-live-2.5-flash"}}, {"model_name": "vertex_ai/gemini-3-pro-live", "litellm_params": {"model": "vertex_ai/gemini-live-2.5-flash"}}, @@ -8206,7 +8206,7 @@ async def test_vertex_live_setup_frame_naming_an_alias_cannot_open_the_model_the websocket, upstream_ws = await _run_vertex_live_gated_passthrough([_setup_frame("gemini-3-pro-live")], valid_token) upstream_ws.send.assert_not_awaited() - reason = _policy_close_reason(websocket) + reason: Final = _policy_close_reason(websocket) assert reason is not None and "gemini-3-pro-live" in reason @@ -8237,7 +8237,7 @@ async def test_vertex_live_setup_frame_naming_a_router_alias_cannot_open_the_mod llm_router=llm_router, ) - sent_models = [json.loads(call.args[0])["setup"]["model"] for call in upstream_ws.send.await_args_list] + sent_models: Final = [json.loads(call.args[0])["setup"]["model"] for call in upstream_ws.send.await_args_list] assert not any(model.endswith("/gemini-3-pro-live") for model in sent_models) assert sent_models == ( [] if hidden else ["projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash"] @@ -8267,7 +8267,7 @@ async def test_vertex_live_setup_frame_naming_a_router_alias_named_like_a_group_ [_setup_frame(VERTEX_LIVE_FULL_PATH.format(model="gemini-3-pro-live"))], key, llm_router=llm_router ) - sent_models = [json.loads(call.args[0])["setup"]["model"] for call in alias_upstream.send.await_args_list] + sent_models: Final = [json.loads(call.args[0])["setup"]["model"] for call in alias_upstream.send.await_args_list] assert sent_models == [VERTEX_LIVE_FULL_PATH.format(model="gemini-live-2.5-flash")] shadowed_upstream.send.assert_not_awaited() @@ -8283,7 +8283,7 @@ async def test_vertex_live_setup_frame_is_authorized_by_the_litellm_model_id_in_ ) assert _policy_close_reason(websocket) is None - sent_setup = json.loads(upstream_ws.send.await_args.args[0])["setup"] + sent_setup: Final = json.loads(upstream_ws.send.await_args.args[0])["setup"] assert sent_setup["model"] == VERTEX_LIVE_FULL_PATH.format(model=VERTEX_LIVE_GROUP) @@ -8292,7 +8292,7 @@ async def test_vertex_live_setup_frame_is_checked_against_the_owner_of_a_persona """A personal key's owner can be limited to fewer models than the key lists, and auth enforces that limit too""" from litellm.proxy.proxy_server import user_api_key_cache - user_id = f"live-user-{uuid.uuid4().hex}" + user_id: Final = f"live-user-{uuid.uuid4().hex}" await user_api_key_cache.async_set_cache( key=user_id, value={"user_id": user_id, "user_email": "live@example.com", "models": ["gemini-live-2.5-flash"]}, @@ -8307,7 +8307,7 @@ async def test_vertex_live_setup_frame_is_checked_against_the_owner_of_a_persona ) upstream_ws.send.assert_not_awaited() - reason = _policy_close_reason(websocket) + reason: Final = _policy_close_reason(websocket) assert reason is not None and "gemini-3-pro-live" in reason @@ -8326,7 +8326,7 @@ async def test_vertex_live_setup_frame_is_checked_against_the_team_in_every_addr ) upstream_ws.send.assert_not_awaited() - reason = _policy_close_reason(websocket) + reason: Final = _policy_close_reason(websocket) assert reason is not None and "gemini-3-pro-live" in reason @@ -8349,7 +8349,7 @@ async def test_vertex_live_setup_frame_naming_another_providers_group_runs_that_ llm_router=_another_providers_live_router(), ) - sent_setup = json.loads(upstream_ws.send.await_args.args[0])["setup"] + sent_setup: Final = json.loads(upstream_ws.send.await_args.args[0])["setup"] assert sent_setup["model"] == "projects/proj-db/locations/global/publishers/google/models/gemini-2.0-flash-live-001" @@ -8369,7 +8369,7 @@ async def test_vertex_live_setup_frame_naming_another_providers_group_is_authori ) assert upstream_ws.send.await_count == (1 if reaches_vertex else 0) - reason = _policy_close_reason(websocket) + reason: Final = _policy_close_reason(websocket) assert (reason is None) is reaches_vertex assert reaches_vertex or ("vertex-flash-live" in reason and "vertex_ai/studio-live" not in reason) @@ -8377,7 +8377,7 @@ async def test_vertex_live_setup_frame_naming_another_providers_group_is_authori @pytest.mark.asyncio async def test_vertex_live_setup_frame_naming_a_router_alias_is_authorized_by_its_target_group(): """The setup rewriter runs a router alias as its target group, so the key holding that group may open it""" - llm_router = litellm.Router( + llm_router: Final = litellm.Router( model_list=[{"model_name": "vertex-live", "litellm_params": {"model": f"vertex_ai/{VERTEX_LIVE_GROUP}"}}], model_group_alias={"gemini-live-alias": "vertex-live"}, ) @@ -8389,7 +8389,7 @@ async def test_vertex_live_setup_frame_naming_a_router_alias_is_authorized_by_it ) assert _policy_close_reason(websocket) is None - sent_setup = json.loads(upstream_ws.send.await_args.args[0])["setup"] + sent_setup: Final = json.loads(upstream_ws.send.await_args.args[0])["setup"] assert sent_setup["model"] == f"projects/proj-db/locations/global/publishers/google/models/{VERTEX_LIVE_GROUP}" @@ -8404,7 +8404,7 @@ async def test_vertex_live_setup_frame_resolves_an_access_group_grant_in_every_a _build_vertex_live_client_frame_gate, ) - llm_router = litellm.Router( + llm_router: Final = litellm.Router( model_list=[ { "model_name": VERTEX_LIVE_GROUP, @@ -8413,11 +8413,11 @@ async def test_vertex_live_setup_frame_resolves_an_access_group_grant_in_every_a }, ] ) - token = UserAPIKeyAuth(token="hashed", models=["live-team"]) - gate = _build_vertex_live_client_frame_gate(llm_router=llm_router, authorize_model=_reauthorized_as(token)) + token: Final = UserAPIKeyAuth(token="hashed", models=["live-team"]) + gate: Final = _build_vertex_live_client_frame_gate(llm_router=llm_router, authorize_model=_reauthorized_as(token)) - frame = json.dumps({"setup": {"model": addressing.format(model=VERTEX_LIVE_GROUP)}}) - denial = await gate(frame, token) + frame: Final = json.dumps({"setup": {"model": addressing.format(model=VERTEX_LIVE_GROUP)}}) + denial: Final = await gate(frame, token) assert denial is None @@ -8433,8 +8433,8 @@ async def _run_vertex_live_route( vertex_llm_base, ) - upstream_ws = RecordingUpstreamWebSocket() - websocket = _client_websocket( + upstream_ws: Final = RecordingUpstreamWebSocket() + websocket: Final = _client_websocket( AsyncMock(side_effect=[*client_frames, {"type": "websocket.disconnect"}]), ) websocket.headers = {"authorization": "Bearer sk-live-client"} @@ -8476,9 +8476,9 @@ async def test_vertex_live_route_gates_the_setup_frame_even_when_the_query_model ) upstream_ws.send.assert_not_awaited() - reason = _policy_close_reason(websocket) + reason: Final = _policy_close_reason(websocket) assert reason is not None and "gemini-3-pro-live" in reason - refusal = environment.proxy_logging.post_call_failure_hook.await_args.kwargs["original_exception"] + refusal: Final = environment.proxy_logging.post_call_failure_hook.await_args.kwargs["original_exception"] assert refusal.code == "403" assert refusal.type == ProxyErrorTypes.key_model_access_denied @@ -8489,7 +8489,7 @@ async def test_vertex_live_route_authorizes_the_setup_model_the_way_connect_time Connect-time auth applies JWT scopes, per-model budgets and the unpriced-model block to ``?model=``, so a model the key's list allows still has to pass them when the setup frame is the first place it is named """ - budget_denial = ProxyException( + budget_denial: Final = ProxyException( message="Budget has been exceeded for model gemini-live-2.5-flash.", type=ProxyErrorTypes.budget_exceeded, param="model", @@ -8577,7 +8577,7 @@ async def test_vertex_live_route_keeps_the_key_allowlist_out_of_the_close_reason UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), ) - reason = _policy_close_reason(websocket) + reason: Final = _policy_close_reason(websocket) assert reason is not None and "gemini-live-2.5-flash" not in reason assert any( "gemini-3-pro-live" in record.getMessage() and "gemini-live-2.5-flash" in record.getMessage() @@ -8606,7 +8606,7 @@ class ScriptedUpstreamWebSocket: self.close = AsyncMock() async def recv(self, decode: bool = True) -> str | None: - frame = next(self._frames, None) + frame: Final = next(self._frames, None) if frame is not None: return frame self.sent_every_frame.set() @@ -8626,7 +8626,7 @@ async def test_vertex_live_turns_served_before_a_refused_frame_are_still_billed( turns and then send a refused frame to have them go unbilled. The refused frame's model must not become the model those turns are priced at either """ - upstream_ws = ScriptedUpstreamWebSocket( + upstream_ws: Final = ScriptedUpstreamWebSocket( ( json.dumps({"setupComplete": {}}), json.dumps({"serverContent": {"turnComplete": True}, "usageMetadata": {"totalTokenCount": 60}}), @@ -8645,7 +8645,7 @@ async def test_vertex_live_turns_served_before_a_refused_frame_are_still_billed( _vertex_live_router(), ) - reason = _policy_close_reason(websocket) + reason: Final = _policy_close_reason(websocket) assert reason is not None and "gemini-3-pro-live" in reason upstream_ws.send.assert_awaited_once() upstream_ws.close.assert_awaited() @@ -8661,7 +8661,7 @@ async def test_vertex_live_session_opened_with_a_binary_setup_frame_is_billed_at A setup frame sent as binary is gated and relayed like a text one, so the session it opens must be priced at the setup model too instead of logging as model unknown with no spend """ - upstream_ws = ScriptedUpstreamWebSocket( + upstream_ws: Final = ScriptedUpstreamWebSocket( ( json.dumps({"setupComplete": {}}), json.dumps({"serverContent": {"turnComplete": True}, "usageMetadata": {"totalTokenCount": 60}}), @@ -8748,8 +8748,8 @@ async def test_vertex_live_turns_served_before_a_frame_refused_mid_stream_are_st async def _run_passthrough_with_frame_gate( gate: WebsocketClientFrameGate, ) -> tuple[MagicMock, ClosableUpstreamWebSocket, SimpleNamespace]: - upstream_ws = ClosableUpstreamWebSocket() - websocket = _client_websocket( + upstream_ws: Final = ClosableUpstreamWebSocket() + websocket: Final = _client_websocket( AsyncMock( side_effect=[ {"type": "websocket.receive", "text": json.dumps({"setup": {"model": "gemini-live-2.5-flash"}})}, @@ -8776,7 +8776,7 @@ async def _run_passthrough_with_frame_gate( @pytest.mark.asyncio async def test_websocket_passthrough_reports_a_refused_client_frame_as_a_failure(): """A denied session logged through the success hooks tells spend tracking and audit that it succeeded""" - denial = ProxyException( + denial: Final = ProxyException( message="gemini-live-2.5-flash: key not allowed to access model", type=ProxyErrorTypes.key_model_access_denied, param="model", @@ -8793,7 +8793,7 @@ async def test_websocket_passthrough_reports_a_refused_client_frame_as_a_failure environment.proxy_logging.post_call_success_hook.assert_not_awaited() environment.logging_worker.ensure_initialized_and_enqueue.assert_not_called() environment.proxy_logging.post_call_failure_hook.assert_awaited_once() - failure_kwargs = environment.proxy_logging.post_call_failure_hook.await_args.kwargs + failure_kwargs: Final = environment.proxy_logging.post_call_failure_hook.await_args.kwargs assert failure_kwargs["original_exception"] is denial assert failure_kwargs["request_data"]["litellm_logging_obj"].litellm_call_id From a3fc0fcf84aad9b21bf46751fe4893546d6c266b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 08:10:53 -0700 Subject: [PATCH 16/17] test(proxy): give the Vertex Live gate tests return types --- .../test_pass_through_endpoints.py | 90 +++++++++---------- 1 file changed, 45 insertions(+), 45 deletions(-) diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 07eff6c3d08..0f305c07e71 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -7389,7 +7389,7 @@ async def test_websocket_passthrough_relays_upstream_policy_close_to_client(): @pytest.mark.asyncio -async def test_websocket_passthrough_relays_as_much_of_a_long_upstream_close_reason_as_fits(): +async def test_websocket_passthrough_relays_as_much_of_a_long_upstream_close_reason_as_fits() -> None: from websockets.exceptions import ConnectionClosedError from websockets.frames import Close @@ -7625,7 +7625,7 @@ def _policy_close_reason(websocket: MagicMock) -> str | None: @pytest.mark.asyncio -async def test_vertex_live_setup_frame_naming_a_model_the_key_cannot_call_is_refused(): +async def test_vertex_live_setup_frame_naming_a_model_the_key_cannot_call_is_refused() -> None: """ ?model= is optional on this route, so a key that names an unlisted model only in the setup frame used to reach Vertex with the gateway's own service-account credentials @@ -7643,7 +7643,7 @@ async def test_vertex_live_setup_frame_naming_a_model_the_key_cannot_call_is_ref @pytest.mark.asyncio -async def test_vertex_live_setup_frame_naming_a_permitted_model_still_reaches_vertex(): +async def test_vertex_live_setup_frame_naming_a_permitted_model_still_reaches_vertex() -> None: websocket, upstream_ws = await _run_vertex_live_gated_passthrough( [_setup_frame("gemini-live-2.5-flash")], UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), @@ -7656,7 +7656,7 @@ async def test_vertex_live_setup_frame_naming_a_permitted_model_still_reaches_ve @pytest.mark.asyncio -async def test_vertex_live_setup_frame_from_an_unrestricted_key_still_reaches_vertex(): +async def test_vertex_live_setup_frame_from_an_unrestricted_key_still_reaches_vertex() -> None: websocket, upstream_ws = await _run_vertex_live_gated_passthrough( [_setup_frame("gemini-3-pro-live")], UserAPIKeyAuth(token="hashed", models=[]), @@ -7667,7 +7667,7 @@ async def test_vertex_live_setup_frame_from_an_unrestricted_key_still_reaches_ve @pytest.mark.asyncio -async def test_vertex_live_setup_frame_is_checked_against_the_team_for_an_all_team_models_key(): +async def test_vertex_live_setup_frame_is_checked_against_the_team_for_an_all_team_models_key() -> None: """ ``all-team-models`` skips the key allowlist, so the team's own restriction is the only thing left to apply """ @@ -7687,7 +7687,7 @@ async def test_vertex_live_setup_frame_is_checked_against_the_team_for_an_all_te @pytest.mark.asyncio -async def test_vertex_live_setup_frame_is_gated_when_it_arrives_as_binary(): +async def test_vertex_live_setup_frame_is_gated_when_it_arrives_as_binary() -> None: """A client that sends its setup frame as bytes must not skip the check the text path applies""" websocket, upstream_ws = await _run_vertex_live_gated_passthrough( [ @@ -7717,7 +7717,7 @@ async def test_vertex_live_binary_setup_frame_reaches_vertex_as_the_model_the_ga @pytest.mark.asyncio -async def test_vertex_live_setup_frame_with_an_escaped_key_is_still_gated(): +async def test_vertex_live_setup_frame_with_an_escaped_key_is_still_gated() -> None: """Large audio frames skip the parse on a substring test, so a JSON-escaped key must not slip past it""" websocket, upstream_ws = await _run_vertex_live_gated_passthrough( [{"type": "websocket.receive", "text": '{"\\u0073etup": {"model": "totally-unlisted-model-abc"}}'}], @@ -7738,7 +7738,7 @@ async def test_vertex_live_setup_frame_with_an_escaped_key_is_still_gated(): '{setup: {model: "totally-unlisted-model-abc"}}', ], ) -async def test_vertex_live_setup_frame_only_vertex_can_parse_is_refused(lenient_setup: str, frame_kind: str): +async def test_vertex_live_setup_frame_only_vertex_can_parse_is_refused(lenient_setup: str, frame_kind: str) -> None: """Vertex reads trailing commas, single quotes and bare keys, which strict JSON refuses to parse""" websocket, upstream_ws = await _run_vertex_live_gated_passthrough( [ @@ -7777,7 +7777,7 @@ async def test_vertex_live_setup_frame_nested_too_deep_to_parse_is_refused() -> ) async def test_vertex_live_setup_frame_only_vertex_can_parse_still_reaches_vertex_for_an_unrestricted_key( lenient_setup, -): +) -> None: """A key that may call any model loses nothing the proxy cannot read, so its lenient frames still go through""" websocket, upstream_ws = await _run_vertex_live_gated_passthrough( [{"type": "websocket.receive", "text": lenient_setup}], @@ -7796,7 +7796,7 @@ async def test_vertex_live_setup_frame_only_vertex_can_parse_still_reaches_verte '{"setup": {"model": "gemini-3-pro-live"}, "setup": null}', ], ) -async def test_vertex_live_setup_frame_repeating_a_key_as_null_is_refused(repeated_key_setup: str): +async def test_vertex_live_setup_frame_repeating_a_key_as_null_is_refused(repeated_key_setup: str) -> None: """ Vertex keeps the first of two keys when the second is null, while ``json.loads`` keeps the last, so the proxy would read no model where Vertex runs the first one @@ -7822,7 +7822,7 @@ async def test_vertex_live_setup_frame_repeating_a_key_as_null_is_refused(repeat ], ids=["litellm-wildcard", "vertex-wildcard", "team", "agent", "end-user"], ) -async def test_vertex_live_setup_frame_only_vertex_can_parse_is_refused_for_any_limited_caller(limited_token: UserAPIKeyAuth): +async def test_vertex_live_setup_frame_only_vertex_can_parse_is_refused_for_any_limited_caller(limited_token: UserAPIKeyAuth) -> None: """A wildcard can match any name the proxy might test the caller with, so only a caller with no model limit anywhere may send a frame the proxy cannot read""" websocket, upstream_ws = await _run_vertex_live_gated_passthrough( @@ -7842,7 +7842,7 @@ async def test_vertex_live_setup_frame_only_vertex_can_parse_is_refused_for_any_ @pytest.mark.asyncio -async def test_vertex_live_client_content_that_mentions_setup_is_forwarded(): +async def test_vertex_live_client_content_that_mentions_setup_is_forwarded() -> None: websocket, upstream_ws = await _run_vertex_live_gated_passthrough( [ _setup_frame("gemini-live-2.5-flash"), @@ -7861,7 +7861,7 @@ async def test_vertex_live_client_content_that_mentions_setup_is_forwarded(): @pytest.mark.asyncio -async def test_vertex_live_non_setup_frames_are_forwarded_without_a_model_check(): +async def test_vertex_live_non_setup_frames_are_forwarded_without_a_model_check() -> None: websocket, upstream_ws = await _run_vertex_live_gated_passthrough( [ _setup_frame("gemini-live-2.5-flash"), @@ -7877,7 +7877,7 @@ async def test_vertex_live_non_setup_frames_are_forwarded_without_a_model_check( @pytest.mark.asyncio @pytest.mark.parametrize("addressing", VERTEX_LIVE_GROUP_ADDRESSINGS) @pytest.mark.parametrize("group", [VERTEX_LIVE_GROUP, VERTEX_LIVE_ALIAS_GROUP]) -async def test_vertex_live_setup_frame_authorizes_a_permitted_group_in_every_addressing(addressing: str, group: str): +async def test_vertex_live_setup_frame_authorizes_a_permitted_group_in_every_addressing(addressing: str, group: str) -> None: """ The Live SDK wraps the caller's model as ``models/``, so a key holding the group has to work whichever addressing its client actually sends @@ -7895,7 +7895,7 @@ async def test_vertex_live_setup_frame_authorizes_a_permitted_group_in_every_add @pytest.mark.asyncio @pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS) -async def test_vertex_live_setup_frame_refuses_an_unpermitted_group_in_every_addressing(addressing: str): +async def test_vertex_live_setup_frame_refuses_an_unpermitted_group_in_every_addressing(addressing: str) -> None: """Normalizing the addressing must not become a way to reach a group the key does not hold""" websocket, upstream_ws = await _run_vertex_live_gated_passthrough( [_setup_frame(addressing.format(model="gemini-3-pro-live"))], @@ -7909,7 +7909,7 @@ async def test_vertex_live_setup_frame_refuses_an_unpermitted_group_in_every_add @pytest.mark.asyncio @pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS) -async def test_vertex_live_setup_frame_naming_no_group_at_all_is_refused_in_every_addressing(addressing: str): +async def test_vertex_live_setup_frame_naming_no_group_at_all_is_refused_in_every_addressing(addressing: str) -> None: """A name that resolves to no group is authorized as it arrived, so an unknown model cannot ride a prefix in""" websocket, upstream_ws = await _run_vertex_live_gated_passthrough( [_setup_frame(addressing.format(model="totally-unlisted-model-abc"))], @@ -7923,7 +7923,7 @@ async def test_vertex_live_setup_frame_naming_no_group_at_all_is_refused_in_ever @pytest.mark.asyncio @pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS) @pytest.mark.parametrize("key_model", ["gemini-4-pro-live-unserved", "vertex_ai/gemini-4-pro-live-unserved"]) -async def test_vertex_live_setup_frame_authorizes_a_model_no_group_serves_by_its_publisher_id(addressing: str, key_model: str): +async def test_vertex_live_setup_frame_authorizes_a_model_no_group_serves_by_its_publisher_id(addressing: str, key_model: str) -> None: """ With no group to resolve, the setup rewriter sends Vertex the last path segment, so a key holding that publisher model, bare or as a LiteLLM ``vertex_ai/`` id, can open it whichever addressing its client sends @@ -7940,7 +7940,7 @@ async def test_vertex_live_setup_frame_authorizes_a_model_no_group_serves_by_its @pytest.mark.asyncio -async def test_vertex_live_setup_frame_addressing_cannot_satisfy_a_wildcard_for_another_publisher_model(): +async def test_vertex_live_setup_frame_addressing_cannot_satisfy_a_wildcard_for_another_publisher_model() -> None: """Vertex runs only the last segment, so the segments ahead of it must not be what matches the key's wildcard""" websocket, upstream_ws = await _run_vertex_live_gated_passthrough( [_setup_frame("vertex_ai/gemini-anything/forbidden-model")], @@ -7953,7 +7953,7 @@ async def test_vertex_live_setup_frame_addressing_cannot_satisfy_a_wildcard_for_ @pytest.mark.asyncio -async def test_vertex_live_setup_frame_prefers_the_group_its_whole_name_matches(): +async def test_vertex_live_setup_frame_prefers_the_group_its_whole_name_matches() -> None: """A group named ``vertex_ai/`` is the one the rewriter sends, so it is the one the key must hold""" llm_router: Final = litellm.Router( model_list=[ @@ -7974,7 +7974,7 @@ async def test_vertex_live_setup_frame_prefers_the_group_its_whole_name_matches( @pytest.mark.asyncio @pytest.mark.parametrize("group", [VERTEX_LIVE_GROUP, VERTEX_LIVE_ALIAS_GROUP]) -async def test_vertex_live_full_path_is_authorized_by_any_group_serving_its_publisher_model(group: str): +async def test_vertex_live_full_path_is_authorized_by_any_group_serving_its_publisher_model(group: str) -> None: """The docs tell callers to send the full resource path, which Vertex runs as sent, so every group serving that publisher model grants it, whatever the group is called""" full_path: Final = VERTEX_LIVE_FULL_PATH.format(model=VERTEX_LIVE_GROUP) @@ -7990,7 +7990,7 @@ async def test_vertex_live_full_path_is_authorized_by_any_group_serving_its_publ @pytest.mark.asyncio -async def test_vertex_live_full_path_is_not_granted_by_a_group_named_after_it_that_serves_another_model(): +async def test_vertex_live_full_path_is_not_granted_by_a_group_named_after_it_that_serves_another_model() -> None: llm_router: Final = litellm.Router( model_list=[ {"model_name": "gemini-3-pro-live", "litellm_params": {"model": "vertex_ai/gemini-live-2.5-flash"}}, @@ -8014,7 +8014,7 @@ async def test_vertex_live_full_path_is_not_granted_by_a_group_named_after_it_th [VERTEX_LIVE_FULL_PATH.format(model=VERTEX_LIVE_GROUP), "openai-native-audio"], ids=["full-path", "group-name"], ) -async def test_vertex_live_setup_frame_is_not_granted_by_a_group_another_provider_serves(setup_model: str): +async def test_vertex_live_setup_frame_is_not_granted_by_a_group_another_provider_serves(setup_model: str) -> None: """Every Live session runs on the proxy's Vertex credentials, so a group another provider serves under the same model id must not open the Vertex model""" llm_router: Final = litellm.Router( @@ -8040,7 +8040,7 @@ async def test_vertex_live_setup_frame_is_not_granted_by_a_group_another_provide @pytest.mark.parametrize(("custom_llm_provider", "reaches_vertex"), [("vertex_ai", True), ("openai", False)]) async def test_vertex_live_full_path_is_authorized_only_by_a_group_whose_custom_provider_is_vertex( custom_llm_provider: str, reaches_vertex: bool -): +) -> None: llm_router: Final = litellm.Router( model_list=[ { @@ -8075,7 +8075,7 @@ async def test_vertex_live_full_path_is_authorized_only_by_a_group_whose_custom_ ) async def test_vertex_live_full_path_whose_every_name_is_taken_by_another_model_refuses_only_limited_keys( valid_token: UserAPIKeyAuth, reaches_vertex: bool -): +) -> None: """Groups holding both the bare and the ``vertex_ai/`` id for another model leave no name a key could hold, which says nothing about a key that may call any model""" llm_router: Final = litellm.Router( @@ -8098,7 +8098,7 @@ async def test_vertex_live_full_path_whose_every_name_is_taken_by_another_model_ @pytest.mark.asyncio async def test_vertex_live_setup_frame_never_refreshes_the_token_of_a_chatgpt_group_on_the_same_proxy( monkeypatch: pytest.MonkeyPatch, tmp_path: Path -): +) -> None: (tmp_path / "auth.json").write_text( json.dumps({"access_token": "expired", "refresh_token": "refresh", "expires_at": 0}) ) @@ -8155,7 +8155,7 @@ class SecretReadRecorder(CustomSecretManager): ) async def test_vertex_live_setup_frame_never_reads_the_secret_manager_for_another_providers_group( monkeypatch: pytest.MonkeyPatch, other_group: dict[str, str] -): +) -> None: monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "in-house", "custom_handler": CustomLLM()}]) llm_router: Final = litellm.Router( model_list=[ @@ -8198,7 +8198,7 @@ async def test_vertex_live_setup_frame_never_reads_the_secret_manager_for_anothe ], ids=["key-alias", "team-alias"], ) -async def test_vertex_live_setup_frame_naming_an_alias_cannot_open_the_model_the_alias_shadows(valid_token: UserAPIKeyAuth): +async def test_vertex_live_setup_frame_naming_an_alias_cannot_open_the_model_the_alias_shadows(valid_token: UserAPIKeyAuth) -> None: """ An alias passes the access check by its target, while the setup rewriter sends the name as written, so a key that could set an alias would open the model the alias name belongs to @@ -8274,7 +8274,7 @@ async def test_vertex_live_setup_frame_naming_a_router_alias_named_like_a_group_ @pytest.mark.asyncio @pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS) -async def test_vertex_live_setup_frame_is_authorized_by_the_litellm_model_id_in_every_addressing(addressing: str): +async def test_vertex_live_setup_frame_is_authorized_by_the_litellm_model_id_in_every_addressing(addressing: str) -> None: """A request naming a deployment's ``vertex_ai/`` id reaches that deployment, so a key holding the id runs the model it names on Live whichever addressing the setup frame uses""" websocket, upstream_ws = await _run_vertex_live_gated_passthrough( @@ -8288,7 +8288,7 @@ async def test_vertex_live_setup_frame_is_authorized_by_the_litellm_model_id_in_ @pytest.mark.asyncio -async def test_vertex_live_setup_frame_is_checked_against_the_owner_of_a_personal_key(): +async def test_vertex_live_setup_frame_is_checked_against_the_owner_of_a_personal_key() -> None: """A personal key's owner can be limited to fewer models than the key lists, and auth enforces that limit too""" from litellm.proxy.proxy_server import user_api_key_cache @@ -8313,7 +8313,7 @@ async def test_vertex_live_setup_frame_is_checked_against_the_owner_of_a_persona @pytest.mark.asyncio @pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS) -async def test_vertex_live_setup_frame_is_checked_against_the_team_in_every_addressing(addressing: str): +async def test_vertex_live_setup_frame_is_checked_against_the_team_in_every_addressing(addressing: str) -> None: """The team allowlist is keyed by group too, so the addressing must not decide whether the team is checked""" websocket, upstream_ws = await _run_vertex_live_gated_passthrough( [_setup_frame(addressing.format(model="gemini-3-pro-live"))], @@ -8340,7 +8340,7 @@ def _another_providers_live_router() -> litellm.Router: @pytest.mark.asyncio -async def test_vertex_live_setup_frame_naming_another_providers_group_runs_that_groups_model_on_vertex(): +async def test_vertex_live_setup_frame_naming_another_providers_group_runs_that_groups_model_on_vertex() -> None: """A caller who may call any model kept reaching Vertex through a group another provider serves before the gate existed, so the setup rewriter still sends that group's model id rather than the group name""" _, upstream_ws = await _run_vertex_live_gated_passthrough( @@ -8361,7 +8361,7 @@ async def test_vertex_live_setup_frame_naming_another_providers_group_runs_that_ ) async def test_vertex_live_setup_frame_naming_another_providers_group_is_authorized_by_the_vertex_model_it_runs( models, reaches_vertex -): +) -> None: websocket, upstream_ws = await _run_vertex_live_gated_passthrough( [_setup_frame("studio-live")], UserAPIKeyAuth(token="hashed", models=models), @@ -8375,7 +8375,7 @@ async def test_vertex_live_setup_frame_naming_another_providers_group_is_authori @pytest.mark.asyncio -async def test_vertex_live_setup_frame_naming_a_router_alias_is_authorized_by_its_target_group(): +async def test_vertex_live_setup_frame_naming_a_router_alias_is_authorized_by_its_target_group() -> None: """The setup rewriter runs a router alias as its target group, so the key holding that group may open it""" llm_router: Final = litellm.Router( model_list=[{"model_name": "vertex-live", "litellm_params": {"model": f"vertex_ai/{VERTEX_LIVE_GROUP}"}}], @@ -8395,7 +8395,7 @@ async def test_vertex_live_setup_frame_naming_a_router_alias_is_authorized_by_it @pytest.mark.asyncio @pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS) -async def test_vertex_live_setup_frame_resolves_an_access_group_grant_in_every_addressing(addressing: str): +async def test_vertex_live_setup_frame_resolves_an_access_group_grant_in_every_addressing(addressing: str) -> None: """ ``get_model_access_groups`` is keyed by group as well, and returns nothing for a prefixed name, so a key entitled through an access group rather than a model name needs the same normalization @@ -8465,7 +8465,7 @@ async def _run_vertex_live_route( @pytest.mark.asyncio -async def test_vertex_live_route_gates_the_setup_frame_even_when_the_query_model_is_permitted(): +async def test_vertex_live_route_gates_the_setup_frame_even_when_the_query_model_is_permitted() -> None: """ ?model= is all connect-time auth sees, so a key that passes it and then names a different model in its setup frame reaches Vertex on the gateway's own credentials unless the route gates the frames as well @@ -8484,7 +8484,7 @@ async def test_vertex_live_route_gates_the_setup_frame_even_when_the_query_model @pytest.mark.asyncio -async def test_vertex_live_route_authorizes_the_setup_model_the_way_connect_time_auth_checks_a_query_model(): +async def test_vertex_live_route_authorizes_the_setup_model_the_way_connect_time_auth_checks_a_query_model() -> None: """ Connect-time auth applies JWT scopes, per-model budgets and the unpriced-model block to ``?model=``, so a model the key's list allows still has to pass them when the setup frame is the first place it is named @@ -8524,7 +8524,7 @@ async def test_vertex_live_route_authorizes_the_setup_model_the_way_connect_time @pytest.mark.asyncio -async def test_vertex_live_route_releases_the_budget_it_reserved_to_authorize_an_allowed_setup_model(): +async def test_vertex_live_route_releases_the_budget_it_reserved_to_authorize_an_allowed_setup_model() -> None: """ Re-authorizing the setup model reserves the request's worst-case cost the way connect-time auth does, and the session's spend is settled on the connect-time token, so an unreleased reservation pins the key's spend counter @@ -8542,7 +8542,7 @@ async def test_vertex_live_route_releases_the_budget_it_reserved_to_authorize_an @pytest.mark.asyncio -async def test_vertex_live_setup_group_the_key_may_call_but_fails_budget_is_not_opened_under_its_bare_model_id(): +async def test_vertex_live_setup_group_the_key_may_call_but_fails_budget_is_not_opened_under_its_bare_model_id() -> None: """ A setup frame naming a group runs that group's deployment whichever name the key holds, so when the key may call the group its budget and pricing checks decide, and the group's bare model id cannot open the session instead @@ -8569,7 +8569,7 @@ async def test_vertex_live_setup_group_the_key_may_call_but_fails_budget_is_not_ @pytest.mark.asyncio -async def test_vertex_live_route_keeps_the_key_allowlist_out_of_the_close_reason_and_in_the_proxy_log(caplog: pytest.LogCaptureFixture): +async def test_vertex_live_route_keeps_the_key_allowlist_out_of_the_close_reason_and_in_the_proxy_log(caplog: pytest.LogCaptureFixture) -> None: """The client learns which model was refused, while only the proxy log names the models the key may reach""" with caplog.at_level("WARNING", logger="LiteLLM Proxy"): websocket, _, _ = await _run_vertex_live_route( @@ -8586,7 +8586,7 @@ async def test_vertex_live_route_keeps_the_key_allowlist_out_of_the_close_reason @pytest.mark.asyncio -async def test_vertex_live_route_forwards_a_setup_frame_the_key_is_allowed_to_run(): +async def test_vertex_live_route_forwards_a_setup_frame_the_key_is_allowed_to_run() -> None: websocket, upstream_ws, _ = await _run_vertex_live_route( [_setup_frame("gemini-live-2.5-flash")], UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), @@ -8620,7 +8620,7 @@ class ScriptedUpstreamWebSocket: @pytest.mark.asyncio -async def test_vertex_live_turns_served_before_a_refused_frame_are_still_billed(): +async def test_vertex_live_turns_served_before_a_refused_frame_are_still_billed() -> None: """ A refusal used to log the whole session as a failure, which records no spend, so a key could run allowed turns and then send a refused frame to have them go unbilled. The refused frame's model must not become @@ -8656,7 +8656,7 @@ async def test_vertex_live_turns_served_before_a_refused_frame_are_still_billed( @pytest.mark.asyncio -async def test_vertex_live_session_opened_with_a_binary_setup_frame_is_billed_at_its_setup_model(): +async def test_vertex_live_session_opened_with_a_binary_setup_frame_is_billed_at_its_setup_model() -> None: """ A setup frame sent as binary is gated and relayed like a text one, so the session it opens must be priced at the setup model too instead of logging as model unknown with no spend @@ -8718,7 +8718,7 @@ def _client_websocket_with_a_close_handshake(receive: Callable[[], Awaitable[Map @pytest.mark.asyncio -async def test_vertex_live_turns_served_before_a_frame_refused_mid_stream_are_still_billed(): +async def test_vertex_live_turns_served_before_a_frame_refused_mid_stream_are_still_billed() -> None: upstream_ws: Final = StreamingUpstreamWebSocket( ( json.dumps({"setupComplete": {}}), @@ -8774,7 +8774,7 @@ async def _run_passthrough_with_frame_gate( @pytest.mark.asyncio -async def test_websocket_passthrough_reports_a_refused_client_frame_as_a_failure(): +async def test_websocket_passthrough_reports_a_refused_client_frame_as_a_failure() -> None: """A denied session logged through the success hooks tells spend tracking and audit that it succeeded""" denial: Final = ProxyException( message="gemini-live-2.5-flash: key not allowed to access model", @@ -8799,7 +8799,7 @@ async def test_websocket_passthrough_reports_a_refused_client_frame_as_a_failure @pytest.mark.asyncio -async def test_websocket_passthrough_still_reports_an_allowed_session_as_a_success(): +async def test_websocket_passthrough_still_reports_an_allowed_session_as_a_success() -> None: async def permitting_gate(frame_data: str | bytes, valid_token: UserAPIKeyAuth, /) -> ProxyException | None: return None From 725de174579d18b506b142e8f0c1531214d0624e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 08:48:28 -0700 Subject: [PATCH 17/17] test(proxy): type the last two parametrized Vertex Live test parameters --- .../pass_through_endpoints/test_pass_through_endpoints.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 0f305c07e71..7bc4ff177dc 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -7776,7 +7776,7 @@ async def test_vertex_live_setup_frame_nested_too_deep_to_parse_is_refused() -> ], ) async def test_vertex_live_setup_frame_only_vertex_can_parse_still_reaches_vertex_for_an_unrestricted_key( - lenient_setup, + lenient_setup: str, ) -> None: """A key that may call any model loses nothing the proxy cannot read, so its lenient frames still go through""" websocket, upstream_ws = await _run_vertex_live_gated_passthrough( @@ -8360,7 +8360,7 @@ async def test_vertex_live_setup_frame_naming_another_providers_group_runs_that_ ids=["other-provider-group", "vertex-group-serving-it", "vertex-model-id"], ) async def test_vertex_live_setup_frame_naming_another_providers_group_is_authorized_by_the_vertex_model_it_runs( - models, reaches_vertex + models: list[str], reaches_vertex: bool ) -> None: websocket, upstream_ws = await _run_vertex_live_gated_passthrough( [_setup_frame("studio-live")],