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] 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)