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