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