diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 38242291b71..6bb300fc0a0 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/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 7d379a6f042..d9281a056bd 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5744,6 +5744,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.registry().deployments(), + 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 33b18c3d939..fd1d0801203 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 @@ -52,6 +52,7 @@ 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.fallback_generalizations import match_routing_generalization +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, @@ -76,14 +77,23 @@ from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.llms.vllm.common_utils import endpoint_under_api_base 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_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 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, ) @@ -121,7 +131,10 @@ 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, YieldsToConfiguredPassThroughRoute, create_pass_through_route, create_websocket_passthrough_route, @@ -159,6 +172,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: @@ -3940,6 +3954,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/" @@ -3964,6 +3984,339 @@ 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 _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, 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, object_pairs_hook=_object_with_unique_keys) + ) + except (ValueError, RecursionError): + return _UnreadableVertexLiveFrame() + 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) + + +VERTEX_LIVE_PROVIDERS: Final = frozenset({LlmProviders.VERTEX_AI.value, LlmProviders.VERTEX_AI_BETA.value}) + + +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) + 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 _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.registry().deployments(source="alias"), + *( + deployment + for deployment in llm_router.registry().deployments() or () + if deployment["model_name"] not in llm_router.registry().aliases() + ), + ) + + +def _vertex_deployments(llm_router: Router | None) -> tuple[DeploymentTypedDict, ...]: + if llm_router is None: + return () + 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: + """ + 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 + """ + if setup_model.startswith("projects/") or llm_router is None: + return None + deployments: Final = tuple(llm_router.registry().deployments() 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 Vertex AI group the setup rewriter sends it + 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 (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( + dict.fromkeys( + candidate["model_name"] + for candidate in _vertex_deployments(llm_router) + 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.registry().deployments() or ()) + ) + 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_listed_router_alias(name: str, llm_router: Router) -> bool: + return any(deployment["model_name"] == name for deployment in llm_router.registry().deployments(source="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.registry().aliases() + and not _is_listed_router_alias(name, llm_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: + 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, 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 +) -> 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 LoggedClientFrameRefusal(_as_model_denial(refusal)) + await release_or_invalidate_budget_reservation(budget_reservation=model_auth.budget_reservation) + return None + + +async def _first_model_denial( + model_names: tuple[str, ...], + valid_token: UserAPIKeyAuth, + llm_router: Router | None, + authorize_model: VertexLiveModelAuthorizer, +) -> ClientFrameRefusal | None: + """ + 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] + 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) + return access_denial if isinstance(later_denial, ProxyException) else later_denial + + +async def _vertex_live_model_denial( + setup_model: str, + valid_token: UserAPIKeyAuth, + llm_router: Router | None, + authorize_model: VertexLiveModelAuthorizer, +) -> ClientFrameRefusal | None: + """ + 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_alias_the_rewriter_ignores(name, valid_token, llm_router) + ) + 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 + ) + return await _first_model_denial(model_names, valid_token, llm_router, authorize_model) + + +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( + setup_model: str | _UnreadableVertexLiveFrame, + valid_token: UserAPIKeyAuth, + llm_router: Router | None, + authorize_model: VertexLiveModelAuthorizer, +) -> ClientFrameRefusal | None: + match setup_model: + case str(): + denial: Final = await _vertex_live_model_denial(setup_model, valid_token, llm_router, authorize_model) + if isinstance(denial, ProxyException): + log_model_access_denial(denial) + return denial + case _UnreadableVertexLiveFrame(): + return ( + 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 + ) + + +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. + + ``?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. + + 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 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, /) -> ClientFrameRefusal | None: + model: Final = _vertex_live_setup_model(frame_data) + if model is None: + return None + try: + 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( + message=VERTEX_LIVE_UNVERIFIED_MODEL_CLOSE_REASON, + type=ProxyErrorTypes.internal_server_error, + param="model", + code=500, + ) + + return gate + + def _resolve_vertex_live_credentials( vertex_project: str | None, vertex_location: str | None, @@ -4022,24 +4375,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 """ - 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.registry().deployments() or ()) - if deployment.get("model_name") in candidates - ), - None, - ) - 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}/") + 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( @@ -4048,6 +4386,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 @@ -4153,6 +4492,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(), + 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 cfac28db1a9..8c32b6ef0dc 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -11,7 +11,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, TypeAlias, TypedDict, cast, overload from urllib.parse import urlencode, urlparse import httpx @@ -34,7 +34,7 @@ from starlette.routing import BaseRoute, Match, Route from starlette.types import Scope from starlette.websockets import WebSocketState from typing_extensions import ReadOnly -from websockets.asyncio.client import connect +from websockets.asyncio.client import ClientConnection, connect from websockets.exceptions import ( ConnectionClosedError, ConnectionClosedOK, @@ -2492,28 +2492,49 @@ 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, RecursionError): + 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 _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( @@ -2537,7 +2558,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 @@ -2552,6 +2573,12 @@ def _truncated_close_reason(reason: str) -> str: return encoded[:WEBSOCKET_CLOSE_REASON_MAX_BYTES].decode("utf-8", errors="ignore") +def _refusal_close_reason(message: str) -> str: + 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] + + SENDABLE_CLOSE_CODES: Final = frozenset(CloseCode) - frozenset( {CloseCode.NO_STATUS_RCVD, CloseCode.ABNORMAL_CLOSURE, CloseCode.TLS_HANDSHAKE} ) @@ -2595,6 +2622,53 @@ 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, 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, /) -> ClientFrameRefusal | None: ... + + +async def _relay_client_frame( + upstream_ws: ClientConnection, + checked_frame: str | bytes, + forwarded_frame: str | bytes, + gate: WebsocketClientFrameGate | None, + valid_token: UserAPIKeyAuth, +) -> 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 + 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 refusal + + async def websocket_passthrough_request( websocket: WebSocket, target: str, @@ -2605,6 +2679,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, builtin_provider_route: bool = False, model: str | None = None, ): @@ -2620,6 +2695,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 builtin_provider_route: Whether a built-in provider route serves the connection, so its URL names the provider model: The model the caller named before connecting, logged on a failure before the setup frame arrives """ @@ -2717,6 +2793,21 @@ 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=_build_passthrough_failure_request_payload( + parsed_body=None if model is None else {"model": model}, + kwargs=kwargs, + logging_obj=logging_obj, + custom_llm_provider=None, + url=httpx.URL(target), + builtin_provider_route=builtin_provider_route, + ), + traceback_str=traceback_str, + ) + # Pre-call logging logging_obj.pre_call( input=[{"role": "user", "content": "WebSocket connection"}], @@ -2748,8 +2839,22 @@ 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""" + 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: while True: message = await websocket.receive() @@ -2762,65 +2867,29 @@ async def websocket_passthrough_request( bytes_data: bytes | None = message.get("bytes") if text_data is not None: - # 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, + if ( + text_refusal := await _relay_client_frame( + upstream_ws, + text_data, + _rewrite_vertex_live_setup_model(text_data, setup_model_rewriter), + client_frame_gate, + user_api_key_dict, ) - 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 - - await upstream_ws.send(_rewrite_vertex_live_setup_model(text_data, setup_model_rewriter)) + ) is not None: + return text_refusal + record_client_setup_model(text_data) elif bytes_data is not None: - await upstream_ws.send(bytes_data) + if ( + bytes_refusal := await _relay_client_frame( + upstream_ws, + bytes_data, + _rewrite_vertex_live_setup_model(bytes_data, setup_model_rewriter), + client_frame_gate, + user_api_key_dict, + ) + ) is not None: + return bytes_refusal + record_client_setup_model(bytes_data) except asyncio.CancelledError: raise except Exception: @@ -2844,7 +2913,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: @@ -2893,13 +2961,33 @@ 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. + # 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, LoggedClientFrameRefusal))), + None, + ) + if client_frame_refusal is not None: + await websocket.close( + code=CloseCode.POLICY_VIOLATION, + reason=_refusal_close_reason(_client_frame_denial(client_frame_refusal).message), + ) + await upstream_ws.close() + 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 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.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, @@ -2974,21 +3062,7 @@ async def websocket_passthrough_request( except InvalidStatus as exc: verbose_proxy_logger.exception("WebSocket passthrough (%s): upstream rejected WebSocket connection", endpoint) - await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=user_api_key_dict, - original_exception=exc, - request_data=_build_passthrough_failure_request_payload( - parsed_body=None if model is None else {"model": model}, - kwargs=kwargs, - logging_obj=logging_obj, - custom_llm_provider=None, - url=httpx.URL(target), - builtin_provider_route=builtin_provider_route, - ), - 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( @@ -3000,21 +3074,7 @@ async def websocket_passthrough_request( "WebSocket passthrough (%s): unexpected error while proxying WebSocket", endpoint ) - await proxy_logging_obj.post_call_failure_hook( - user_api_key_dict=user_api_key_dict, - original_exception=e, - request_data=_build_passthrough_failure_request_payload( - parsed_body=None if model is None else {"model": model}, - kwargs=kwargs, - logging_obj=logging_obj, - custom_llm_provider=None, - url=httpx.URL(target), - builtin_provider_route=builtin_provider_route, - ), - 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/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 9d4fe67614e..a4365477784 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, @@ -2427,6 +2428,37 @@ async def test_bulk_eviction_clears_the_redis_copies_before_waiting_on_a_wedged_ 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 + + class _SignallingPublishRedisClient(_SlowPublishRedisClient): def __init__(self) -> None: super().__init__() 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 a9098ac3d59..19dfba701ec 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,11 +7,12 @@ 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 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 @@ -31,13 +32,15 @@ from starlette.datastructures import UploadFile as StarletteUploadFile import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache -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.integrations.custom_secret_manager import CustomSecretManager from litellm.integrations.prometheus import PrometheusLogger +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 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.llm_provider_handlers.vertex_passthrough_logging_handler import ( @@ -59,6 +62,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, @@ -71,6 +75,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id from litellm.proxy.utils import ProxyLogging 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.utils import CallTypes from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY, @@ -7294,6 +7299,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) -> None: + self._closing = asyncio.Event() + self.send = AsyncMock() + self.close = AsyncMock(side_effect=self._close) + + 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) -> 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) -> "ClosableUpstreamWebSocket": + return self + + async def __anext__(self) -> str: + raise StopAsyncIteration + + def _client_websocket(receive): from starlette.websockets import WebSocketState @@ -7325,11 +7363,17 @@ 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: Final = 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 async def _pending_receive(): @@ -7365,6 +7409,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() -> None: + from websockets.exceptions import ConnectionClosedError + from websockets.frames import Close + + 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: Final = ClosingUpstreamWebSocket( + ConnectionClosedError(rcvd=Close(1011, upstream_reason), sent=Close(1011, ""), rcvd_then_sent=True) + ) + websocket: Final = _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: 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 + + @pytest.mark.asyncio async def test_websocket_passthrough_keeps_normal_upstream_close_normal(): from websockets.exceptions import ConnectionClosedOK @@ -7455,6 +7529,1308 @@ 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_GROUP_ADDRESSINGS = ( + "{model}", + "models/{model}", + "vertex_ai/{model}", + "publishers/google/models/{model}", +) + +VERTEX_LIVE_FULL_PATH = "projects/proj-db/locations/global/publishers/google/models/{model}" + +VERTEX_LIVE_CLIENT_ADDRESSINGS = (*VERTEX_LIVE_GROUP_ADDRESSINGS, VERTEX_LIVE_FULL_PATH) + + +def _vertex_live_router() -> litellm.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"}, + }, + ] + ) + + +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: Final = RecordingUpstreamWebSocket() + websocket, _ = await _run_vertex_live_gated_session( + _client_websocket(AsyncMock(side_effect=[*client_frames, {"type": "websocket.disconnect"}])), + upstream_ws, + valid_token, + llm_router or _vertex_live_router(), + authorize_model, + ) + return websocket, upstream_ws + + +async def _run_vertex_live_gated_session( + websocket: MagicMock, + 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, + _build_vertex_live_setup_model_rewriter, + ) + + 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=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_router=llm_router, + authorize_model=authorize_model or _reauthorized_as(valid_token), + ), + ) + + return websocket, environment + + +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: MagicMock) -> 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() -> 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 + """ + 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: 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() + + +@pytest.mark.asyncio +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"]), + ) + + upstream_ws.send.assert_awaited_once() + 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 + + +@pytest.mark.asyncio +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=[]), + ) + + 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() -> None: + """ + ``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: Final = _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() -> 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( + [ + { + "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_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() -> 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"}}'}], + 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("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: 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( + [ + { + "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_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", + [ + '{"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: 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( + [{"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: 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 + """ + 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 +@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: 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( + [ + { + "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() -> None: + 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() -> None: + 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_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) -> 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 + """ + 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: Final = 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: 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"))], + UserAPIKeyAuth(token="hashed", models=[VERTEX_LIVE_GROUP]), + ) + + upstream_ws.send.assert_not_awaited() + reason: Final = _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: 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"))], + 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) +@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) -> 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 + """ + websocket, upstream_ws = await _run_vertex_live_gated_passthrough( + [_setup_frame(addressing.format(model="gemini-4-pro-live-unserved"))], + UserAPIKeyAuth(token="hashed", models=[key_model]), + ) + + upstream_ws.send.assert_awaited_once() + 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 + + +@pytest.mark.asyncio +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")], + UserAPIKeyAuth(token="hashed", models=["vertex_ai/gemini-*"]), + ) + + upstream_ws.send.assert_not_awaited() + 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() -> 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=[ + {"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 +@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) -> 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) + + 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() -> None: + llm_router: Final = 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: Final = _policy_close_reason(websocket) + 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) -> 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( + 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 +) -> None: + llm_router: Final = 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 +) -> 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( + 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 +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}) + ) + 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 + + +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] +) -> None: + 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", + [ + 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: 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 + """ + websocket, upstream_ws = await _run_vertex_live_gated_passthrough([_setup_frame("gemini-3-pro-live")], valid_token) + + upstream_ws.send.assert_not_awaited() + reason: Final = _policy_close_reason(websocket) + assert reason is not None and "gemini-3-pro-live" in reason + + +@pytest.mark.asyncio +@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 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=[ + {"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}}, + ) + + _, 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, + ) + + 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"] + ) + + +@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: 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() + + +@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) -> 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( + [_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: Final = 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() -> 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 + + 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"]}, + ) + + 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: Final = _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_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"))], + UserAPIKeyAuth( + token="hashed", + models=["all-team-models"], + team_id="team-1", + team_models=[VERTEX_LIVE_GROUP], + ), + ) + + upstream_ws.send.assert_not_awaited() + reason: Final = _policy_close_reason(websocket) + 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() -> 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( + [_setup_frame("models/studio-live")], + UserAPIKeyAuth(token="hashed"), + llm_router=_another_providers_live_router(), + ) + + 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" + + +@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: list[str], reaches_vertex: bool +) -> None: + 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: 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) + + +@pytest.mark.asyncio +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}"}}], + 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: 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}" + + +@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) -> 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 + """ + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _build_vertex_live_client_frame_gate, + ) + + llm_router: Final = litellm.Router( + model_list=[ + { + "model_name": VERTEX_LIVE_GROUP, + "litellm_params": {"model": f"vertex_ai/{VERTEX_LIVE_GROUP}"}, + "model_info": {"access_groups": ["live-team"]}, + }, + ] + ) + 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: Final = json.dumps({"setup": {"model": addressing.format(model=VERTEX_LIVE_GROUP)}}) + denial: Final = await gate(frame, token) + + assert denial is None + + +async def _run_vertex_live_route( + 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 ( + vertex_ai_live_websocket_passthrough, + vertex_llm_base, + ) + + upstream_ws: Final = RecordingUpstreamWebSocket() + websocket: Final = _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, + 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, + authorize_model=None if reauthorize_through_user_api_key_auth else _reauthorized_as(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() -> 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 + """ + 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: Final = _policy_close_reason(websocket) + assert reason is not None and "gemini-3-pro-live" in reason + 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 + + +@pytest.mark.asyncio +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 + """ + budget_denial: Final = ProxyException( + message="Budget has been exceeded for model gemini-live-2.5-flash.", + type=ProxyErrorTypes.budget_exceeded, + param="model", + code=400, + ) + 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")], + valid_token, + reauthorize_through_user_api_key_auth=True, + ) + + upstream_ws.send.assert_not_awaited() + assert _policy_close_reason(websocket) == budget_denial.message + 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() -> 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 + 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 +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 + """ + 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: 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( + [_setup_frame("gemini-3-pro-live")], + UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]), + ) + + 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() + for record in caplog.records + ) + + +@pytest.mark.asyncio +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"]), + ) + + upstream_ws.send.assert_awaited_once() + 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: 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) -> str | None: + frame: Final = next(self._frames, None) + if frame is not None: + return frame + self.sent_every_frame.set() + await asyncio.Event().wait() + + def __aiter__(self) -> "ScriptedUpstreamWebSocket": + return self + + async def __anext__(self) -> str: + raise StopAsyncIteration + + +@pytest.mark.asyncio +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 + the model those turns are priced at either + """ + upstream_ws: Final = ScriptedUpstreamWebSocket( + ( + 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(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.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() + + +@pytest.mark.asyncio +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 + """ + upstream_ws: Final = 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""" + + 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: 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 + + 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() -> None: + 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]: + upstream_ws: Final = ClosableUpstreamWebSocket() + websocket: Final = _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() -> 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", + type=ProxyErrorTypes.key_model_access_denied, + param="model", + code=403, + ) + + 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) + + 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: 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 + + +@pytest.mark.asyncio +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 + + _, _, 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"],