scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

Merge 8b47acaef0 into eddfb5fb20

Browse source
This commit is contained in:
Marty Sullivan 2026-09-12 09:58:13 -07:00 • committed by GitHub
commit 474ffc83ee
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 654 additions and 21 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

126
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -39,7 +39,7 @@ from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.passthrough.main import AsyncPassthroughStreamingResponse
from litellm.proxy._types import *
from litellm.proxy.auth.auth_checks import enforced_model_allowlists
from litellm.proxy.auth.auth_checks import can_key_call_resolved_model, enforced_model_allowlists
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.auth.user_api_key_auth import (
@ -62,6 +62,7 @@ from litellm.proxy.common_utils.sse_keepalive import (
from litellm.proxy.pass_through_endpoints.common_utils import get_litellm_virtual_key
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
HttpPassThroughEndpointHelpers,
WebsocketClientFrameGate,
create_pass_through_route,
create_websocket_passthrough_route,
websocket_passthrough_request,
@ -90,6 +91,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:
@ -2717,6 +2719,10 @@ 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_PUBLISHER_MODEL_PREFIX: Final = "publishers/google/models/"
VERTEX_PUBLISHERS_SEGMENT: Final = "publishers/"
@ -2741,6 +2747,107 @@ def _get_llm_router() -> Router | None:
return llm_router
def _get_llm_model_list() -> list | None: # mutable-ok: mirrors proxy_server.llm_model_list's own type
from litellm.proxy.proxy_server import llm_model_list
return llm_model_list
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. A backslash means the key may be escaped, which only a parse settles
"""
if isinstance(frame_data, bytes):
return b'"setup"' in frame_data or b"\\" in frame_data
return '"setup"' in frame_data or "\\" in frame_data
def _vertex_live_setup_model(frame_data: str | bytes) -> str | None:
"""
The model a Vertex AI Live ``setup`` frame names, which is the only place BidiGenerateContent carries one
"""
if not _may_carry_vertex_live_setup(frame_data):
return None
try:
frame: Final = json.loads(frame_data)
except (json.JSONDecodeError, UnicodeDecodeError):
return None
setup: Final = frame.get("setup") if isinstance(frame, dict) else None
model: Final = setup.get("model") if isinstance(setup, dict) else None
return model if isinstance(model, str) and model else None
def _vertex_live_deployment(setup_model: str, llm_router: Router | None) -> DeploymentTypedDict | None:
"""
The router deployment whose group a ``setup`` frame's model names, whichever addressing the client used.
The Live SDK wraps whatever the caller typed as ``models/<name>``, and clients also send LiteLLM ids and
full Vertex resource paths, so the group is read from the last segment as well as from the whole string
"""
if llm_router is None:
return None
candidates: Final = (setup_model, setup_model.rsplit("/", 1)[-1])
return next(
(
deployment
for deployment in (llm_router.get_model_list() or ())
if deployment.get("model_name") in candidates
),
None,
)
def _build_vertex_live_client_frame_gate(
llm_model_list: list | None, # mutable-ok: matches can_key_call_resolved_model's own type
llm_router: Router | None,
) -> WebsocketClientFrameGate:
"""
Apply the key's, team's and project's model access to the model a client's ``setup`` frame names.
``?model=`` is optional on this route and the Live protocol carries the real model in the first client
frame, so connect-time auth has no model to check and every frame that names one has to be authorized
here instead.
The allowlists and the access groups behind them are keyed by model group, and every addressing a client
sends carries that group inside a prefix, so the group is what gets authorized. A frame naming no group
is still checked as it arrived, which is the strictest reading available for a name nothing recognises
"""
async def gate(frame_data: str | bytes, valid_token: UserAPIKeyAuth, /) -> ProxyException | None:
model: Final = _vertex_live_setup_model(frame_data)
if model is None:
return None
deployment: Final = _vertex_live_deployment(model, llm_router)
try:
await can_key_call_resolved_model(
model=model if deployment is None else deployment["model_name"],
llm_model_list=llm_model_list,
valid_token=valid_token,
llm_router=llm_router,
)
except ProxyException as denial:
# A close frame carries 123 bytes, so the model leads: the tail of the allowlist dump is what a
# caller can most afford to lose
return ProxyException(
message=f"{model}: {denial.message}",
type=denial.type,
param=denial.param,
code=denial.code,
)
except Exception: # noqa: BLE001 # fail closed: a gate that cannot answer refuses, and says so in band
verbose_proxy_logger.exception("Vertex AI Live passthrough: model access check failed for %s", model)
return ProxyException(
message=VERTEX_LIVE_UNVERIFIED_MODEL_CLOSE_REASON,
type=ProxyErrorTypes.internal_server_error,
param="model",
code=500,
)
return None
return gate
def _resolve_vertex_live_credentials(
vertex_project: str | None,
vertex_location: str | None,
@ -2799,17 +2906,8 @@ def _resolve_alias_to_upstream_model(setup_model: str, llm_router: Router | None
"""
The Live SDK wraps whatever the caller typed as ``models/<name>``, so a gateway alias arrives prefixed
"""
if llm_router is None:
return setup_model
candidates: Final = (setup_model, setup_model.rsplit("/", 1)[-1])
upstream: Final = next(
(
deployment["litellm_params"].get("model")
for deployment in (llm_router.get_model_list() or ())
if deployment.get("model_name") in candidates
),
None,
)
deployment: Final = _vertex_live_deployment(setup_model, llm_router)
upstream: Final = None if deployment is None else deployment["litellm_params"].get("model")
if upstream is None:
return setup_model
try:
@ -2923,6 +3021,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_model_list=_get_llm_model_list(),
llm_router=_get_llm_router(),
),
)

83
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -9,7 +9,7 @@ from collections.abc import AsyncGenerator, Callable, Iterable, Mapping, Sequenc
from dataclasses import dataclass
from datetime import datetime
from itertools import groupby
from typing import TYPE_CHECKING, Any, Final, TypedDict, cast
from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast
from urllib.parse import urlencode, urlparse
import httpx
@ -27,7 +27,7 @@ from fastapi import (
from fastapi.responses import StreamingResponse
from starlette.datastructures import UploadFile as StarletteUploadFile
from starlette.websockets import WebSocketState
from websockets.asyncio.client import connect
from websockets.asyncio.client import ClientConnection, connect
from websockets.exceptions import (
ConnectionClosedError,
ConnectionClosedOK,
@ -2132,6 +2132,39 @@ def _upstream_close_to_relay(task_results: Iterable[object]) -> Close | None:
return upstream_close
class WebsocketClientFrameGate(Protocol):
"""Authorize one client frame, returning the denial to report or ``None`` to let it through.
A websocket passthrough authenticates before any frame arrives, so a protocol that names its model
inside a frame has no model to authorize at connect time. The frame, not the extracted model, is the
argument because where a model sits is part of each provider's wire protocol
"""
async def __call__(self, frame_data: str | bytes, valid_token: UserAPIKeyAuth, /) -> ProxyException | None: ...
async def _relay_client_frame(
websocket: WebSocket,
upstream_ws: ClientConnection,
checked_frame: str | bytes,
forwarded_frame: str | bytes,
gate: WebsocketClientFrameGate | None,
valid_token: UserAPIKeyAuth,
) -> ProxyException | None:
"""Forward one client frame upstream, or refuse it and hand back the denial that ended the session.
The denial travels out so the caller reports a failed session instead of leaving the success hooks to
record a refusal as a success, and tearing the upstream down is the caller's job: awaiting that close
here lets the upstream reader finish first, which cancels this task mid-close and loses the denial
"""
refusal: Final = None if gate is None else await gate(checked_frame, valid_token)
if refusal is None:
await upstream_ws.send(forwarded_frame)
return None
await websocket.close(code=CloseCode.POLICY_VIOLATION, reason=_truncated_close_reason(refusal.message))
return refusal
async def websocket_passthrough_request(
websocket: WebSocket,
target: str,
@ -2142,6 +2175,7 @@ async def websocket_passthrough_request(
cost_per_request: float | None = None,
accept_websocket: bool = True,
setup_model_rewriter: Callable[[str], str] | None = None,
client_frame_gate: WebsocketClientFrameGate | None = None,
):
"""
WebSocket passthrough request handler.
@ -2155,6 +2189,7 @@ async def websocket_passthrough_request(
endpoint: The endpoint path (for logging purposes)
cost_per_request: Optional field - cost per request to the target endpoint
setup_model_rewriter: Optional rewrite of the setup frame's model before it reaches the upstream
client_frame_gate: Optional per-frame authorization, applied before a frame reaches the upstream
"""
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.proxy.proxy_server import proxy_config, proxy_logging_obj
@ -2284,8 +2319,8 @@ async def websocket_passthrough_request(
"WebSocket passthrough (%s): Upstream connection established successfully", endpoint
)
async def forward_client_to_upstream() -> None:
"""Forward messages from client to upstream WebSocket"""
async def forward_client_to_upstream() -> ProxyException | None:
"""Forward messages from client to upstream WebSocket, returning any denial that ended it"""
try:
while True:
message = await websocket.receive()
@ -2352,9 +2387,27 @@ async def websocket_passthrough_request(
)
# Not a JSON message or doesn't contain setup data
await upstream_ws.send(_rewrite_vertex_live_setup_model(text_data, setup_model_rewriter))
refusal = await _relay_client_frame(
websocket,
upstream_ws,
text_data,
_rewrite_vertex_live_setup_model(text_data, setup_model_rewriter),
client_frame_gate,
user_api_key_dict,
)
if refusal is not None:
return refusal
elif bytes_data is not None:
await upstream_ws.send(bytes_data)
refusal = await _relay_client_frame(
websocket,
upstream_ws,
bytes_data,
bytes_data,
client_frame_gate,
user_api_key_dict,
)
if refusal is not None:
return refusal
except asyncio.CancelledError:
raise
except Exception:
@ -2456,13 +2509,29 @@ async def websocket_passthrough_request(
except asyncio.CancelledError:
pass
task_results: Final = tuple(task.result() for task in done if task.exception() is None)
# A refusal outranks any exception it caused: closing the client socket is what makes the
# upstream reader's next send fail, so raising that instead would report the wrong cause
client_frame_refusal: Final = next(
(result for result in task_results if isinstance(result, ProxyException)), None
)
if client_frame_refusal is not None:
await upstream_ws.close()
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=client_frame_refusal,
request_data={**kwargs, "litellm_logging_obj": logging_obj}, # mutable-ok: the hook pops keys
)
return
# Check for exceptions in completed tasks
for task in done:
exception = task.exception()
if exception is not None:
raise exception
upstream_close: Final = _upstream_close_to_relay(task.result() for task in done)
upstream_close: Final = _upstream_close_to_relay(task_results)
if upstream_close is not None and _client_socket_is_open(websocket):
await websocket.close(
code=upstream_close.code,

466
tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py
View file

@ -32,7 +32,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
)
@ -4869,6 +4869,39 @@ class RecordingUpstreamWebSocket:
raise StopAsyncIteration
class ClosableUpstreamWebSocket:
"""An upstream whose ``recv`` ends when ``close`` does, and whose close keeps yielding after that.
A real close handshake outlives the ``recv`` it unblocks, so the reader task finishes while the closing
task is still suspended. That is the interleaving that leaves a refusal stranded in a cancelled task
"""
CLOSE_HANDSHAKE_YIELDS = 5
def __init__(self):
self._closing = asyncio.Event()
self.send = AsyncMock()
self.close = AsyncMock(side_effect=self._close)
async def _close(self):
self._closing.set()
for _ in range(self.CLOSE_HANDSHAKE_YIELDS):
await asyncio.sleep(0)
async def recv(self, decode: bool = True):
from websockets.exceptions import ConnectionClosedOK
from websockets.frames import Close
await self._closing.wait()
raise ConnectionClosedOK(rcvd=Close(1000, ""), sent=Close(1000, ""), rcvd_then_sent=True)
def __aiter__(self):
return self
async def __anext__(self):
raise StopAsyncIteration
def _client_websocket(receive):
from starlette.websockets import WebSocketState
@ -4904,7 +4937,7 @@ def _patched_websocket_passthrough_environment(upstream_ws):
mock_worker.ensure_initialized_and_enqueue = MagicMock(
side_effect=lambda async_coroutine: async_coroutine.close()
)
yield
yield SimpleNamespace(proxy_logging=mock_proxy_logging, logging_worker=mock_worker)
async def _pending_receive():
@ -5030,6 +5063,435 @@ async def test_websocket_passthrough_rewrites_gateway_alias_setup_model():
assert sent_setup["model"] == "projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash"
VERTEX_LIVE_TARGET = "wss://aiplatform.googleapis.com/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent"
VERTEX_LIVE_GROUP = "gemini-live-2.5-flash-native-audio"
VERTEX_LIVE_ALIAS_GROUP = "gemini-live-native-audio"
VERTEX_LIVE_CLIENT_ADDRESSINGS = (
"{model}",
"models/{model}",
"vertex_ai/{model}",
"publishers/google/models/{model}",
"projects/proj-db/locations/global/publishers/google/models/{model}",
)
def _vertex_live_router():
return litellm.Router(
model_list=[
{
"model_name": "gemini-live-2.5-flash",
"litellm_params": {"model": "vertex_ai/gemini-live-2.5-flash"},
},
{
"model_name": VERTEX_LIVE_GROUP,
"litellm_params": {"model": f"vertex_ai/{VERTEX_LIVE_GROUP}"},
},
{
"model_name": VERTEX_LIVE_ALIAS_GROUP,
"litellm_params": {"model": f"vertex_ai/{VERTEX_LIVE_GROUP}"},
},
{
"model_name": "gemini-3-pro-live",
"litellm_params": {"model": "vertex_ai/gemini-3-pro-live"},
},
]
)
async def _run_vertex_live_gated_passthrough(client_frames, valid_token):
"""Drive /vertex_ai/live with the real frame gate, returning the client and upstream sockets"""
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
_build_vertex_live_client_frame_gate,
_build_vertex_live_setup_model_rewriter,
)
llm_router = _vertex_live_router()
upstream_ws = RecordingUpstreamWebSocket()
websocket = _client_websocket(
AsyncMock(side_effect=[*client_frames, {"type": "websocket.disconnect"}]),
)
with _patched_websocket_passthrough_environment(upstream_ws):
await websocket_passthrough_request(
websocket=websocket,
target=VERTEX_LIVE_TARGET,
custom_headers={"Authorization": "Bearer token"},
user_api_key_dict=valid_token,
forward_headers=False,
endpoint="/vertex_ai/live",
accept_websocket=False,
setup_model_rewriter=_build_vertex_live_setup_model_rewriter(
vertex_project="proj-db",
vertex_location="global",
llm_router=llm_router,
),
client_frame_gate=_build_vertex_live_client_frame_gate(
llm_model_list=llm_router.get_model_list(),
llm_router=llm_router,
),
)
return websocket, upstream_ws
def _setup_frame(model: str) -> dict:
return {
"type": "websocket.receive",
"text": json.dumps({"setup": {"model": model, "generationConfig": {"responseModalities": ["TEXT"]}}}),
}
def _policy_close_reason(websocket) -> str | None:
return next(
(call.kwargs.get("reason") for call in websocket.close.await_args_list if call.kwargs.get("code") == 1008),
None,
)
@pytest.mark.asyncio
async def test_vertex_live_setup_frame_naming_a_model_the_key_cannot_call_is_refused():
"""
?model= is optional on this route, so a key that names an unlisted model only in the setup frame used to
reach Vertex with the gateway's own service-account credentials
"""
websocket, upstream_ws = await _run_vertex_live_gated_passthrough(
[_setup_frame("totally-unlisted-model-abc")],
UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]),
)
upstream_ws.send.assert_not_awaited()
reason = _policy_close_reason(websocket)
assert reason is not None and "totally-unlisted-model-abc" in reason
upstream_ws.close.assert_awaited()
@pytest.mark.asyncio
async def test_vertex_live_setup_frame_naming_a_permitted_model_still_reaches_vertex():
websocket, upstream_ws = await _run_vertex_live_gated_passthrough(
[_setup_frame("gemini-live-2.5-flash")],
UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]),
)
upstream_ws.send.assert_awaited_once()
sent_setup = json.loads(upstream_ws.send.await_args.args[0])["setup"]
assert sent_setup["model"] == "projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash"
assert _policy_close_reason(websocket) is None
@pytest.mark.asyncio
async def test_vertex_live_setup_frame_from_an_unrestricted_key_still_reaches_vertex():
websocket, upstream_ws = await _run_vertex_live_gated_passthrough(
[_setup_frame("gemini-3-pro-live")],
UserAPIKeyAuth(token="hashed", models=[]),
)
upstream_ws.send.assert_awaited_once()
assert _policy_close_reason(websocket) is None
@pytest.mark.asyncio
async def test_vertex_live_setup_frame_is_checked_against_the_team_for_an_all_team_models_key():
"""
``all-team-models`` skips the key allowlist, so the team's own restriction is the only thing left to apply
"""
websocket, upstream_ws = await _run_vertex_live_gated_passthrough(
[_setup_frame("gemini-3-pro-live")],
UserAPIKeyAuth(
token="hashed",
models=["all-team-models"],
team_id="team-1",
team_models=["gemini-live-2.5-flash"],
),
)
upstream_ws.send.assert_not_awaited()
reason = _policy_close_reason(websocket)
assert reason is not None and "team not allowed to access model" in reason
@pytest.mark.asyncio
async def test_vertex_live_setup_frame_is_gated_when_it_arrives_as_binary():
"""A client that sends its setup frame as bytes must not skip the check the text path applies"""
websocket, upstream_ws = await _run_vertex_live_gated_passthrough(
[
{
"type": "websocket.receive",
"bytes": json.dumps({"setup": {"model": "totally-unlisted-model-abc"}}).encode(),
}
],
UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]),
)
upstream_ws.send.assert_not_awaited()
assert _policy_close_reason(websocket) is not None
@pytest.mark.asyncio
async def test_vertex_live_setup_frame_with_an_escaped_key_is_still_gated():
"""Large audio frames skip the parse on a substring test, so a JSON-escaped key must not slip past it"""
websocket, upstream_ws = await _run_vertex_live_gated_passthrough(
[{"type": "websocket.receive", "text": '{"\\u0073etup": {"model": "totally-unlisted-model-abc"}}'}],
UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]),
)
upstream_ws.send.assert_not_awaited()
assert _policy_close_reason(websocket) is not None
@pytest.mark.asyncio
async def test_vertex_live_non_setup_frames_are_forwarded_without_a_model_check():
websocket, upstream_ws = await _run_vertex_live_gated_passthrough(
[
_setup_frame("gemini-live-2.5-flash"),
{"type": "websocket.receive", "text": json.dumps({"realtimeInput": {"audio": {"data": "AAAA"}}})},
],
UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]),
)
assert upstream_ws.send.await_count == 2
assert _policy_close_reason(websocket) is None
@pytest.mark.asyncio
@pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS)
@pytest.mark.parametrize("group", [VERTEX_LIVE_GROUP, VERTEX_LIVE_ALIAS_GROUP])
async def test_vertex_live_setup_frame_authorizes_a_permitted_group_in_every_addressing(addressing, group):
"""
The Live SDK wraps the caller's model as ``models/<name>`` and the docs tell callers to send the full
resource path, so a key holding the group has to work whichever addressing its client actually sends
"""
websocket, upstream_ws = await _run_vertex_live_gated_passthrough(
[_setup_frame(addressing.format(model=group))],
UserAPIKeyAuth(token="hashed", models=[group]),
)
upstream_ws.send.assert_awaited_once()
sent_setup = json.loads(upstream_ws.send.await_args.args[0])["setup"]
assert sent_setup["model"].startswith("projects/")
assert _policy_close_reason(websocket) is None
@pytest.mark.asyncio
@pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS)
async def test_vertex_live_setup_frame_refuses_an_unpermitted_group_in_every_addressing(addressing):
"""Normalizing the addressing must not become a way to reach a group the key does not hold"""
websocket, upstream_ws = await _run_vertex_live_gated_passthrough(
[_setup_frame(addressing.format(model="gemini-3-pro-live"))],
UserAPIKeyAuth(token="hashed", models=[VERTEX_LIVE_GROUP]),
)
upstream_ws.send.assert_not_awaited()
reason = _policy_close_reason(websocket)
assert reason is not None and "gemini-3-pro-live" in reason
@pytest.mark.asyncio
@pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS)
async def test_vertex_live_setup_frame_naming_no_group_at_all_is_refused_in_every_addressing(addressing):
"""A name that resolves to no group is authorized as it arrived, so an unknown model cannot ride a prefix in"""
websocket, upstream_ws = await _run_vertex_live_gated_passthrough(
[_setup_frame(addressing.format(model="totally-unlisted-model-abc"))],
UserAPIKeyAuth(token="hashed", models=[VERTEX_LIVE_GROUP]),
)
upstream_ws.send.assert_not_awaited()
assert _policy_close_reason(websocket) is not None
@pytest.mark.asyncio
@pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS[1:])
async def test_vertex_live_setup_frame_does_not_strip_a_prefix_off_a_model_no_group_serves(addressing):
"""
Stripping the addressing off a name the router does not serve would let a wildcard key reach an arbitrary
publisher model, and a full path would carry a project of the caller's choosing with it
"""
websocket, upstream_ws = await _run_vertex_live_gated_passthrough(
[_setup_frame(addressing.format(model="gemini-4-pro-live-unserved"))],
UserAPIKeyAuth(token="hashed", models=["gemini-*"]),
)
upstream_ws.send.assert_not_awaited()
assert _policy_close_reason(websocket) is not None
@pytest.mark.asyncio
@pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS)
async def test_vertex_live_setup_frame_is_checked_against_the_team_in_every_addressing(addressing):
"""The team allowlist is keyed by group too, so the addressing must not decide whether the team is checked"""
websocket, upstream_ws = await _run_vertex_live_gated_passthrough(
[_setup_frame(addressing.format(model="gemini-3-pro-live"))],
UserAPIKeyAuth(
token="hashed",
models=["all-team-models"],
team_id="team-1",
team_models=[VERTEX_LIVE_GROUP],
),
)
upstream_ws.send.assert_not_awaited()
reason = _policy_close_reason(websocket)
assert reason is not None and "team not allowed to access model" in reason
@pytest.mark.asyncio
@pytest.mark.parametrize("addressing", VERTEX_LIVE_CLIENT_ADDRESSINGS)
async def test_vertex_live_setup_frame_resolves_an_access_group_grant_in_every_addressing(addressing):
"""
``get_model_access_groups`` is keyed by group as well, and returns nothing for a prefixed name, so a key
entitled through an access group rather than a model name needs the same normalization
"""
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
_build_vertex_live_client_frame_gate,
)
llm_router = litellm.Router(
model_list=[
{
"model_name": VERTEX_LIVE_GROUP,
"litellm_params": {"model": f"vertex_ai/{VERTEX_LIVE_GROUP}"},
"model_info": {"access_groups": ["live-team"]},
},
]
)
gate = _build_vertex_live_client_frame_gate(
llm_model_list=llm_router.get_model_list(),
llm_router=llm_router,
)
frame = json.dumps({"setup": {"model": addressing.format(model=VERTEX_LIVE_GROUP)}})
denial = await gate(frame, UserAPIKeyAuth(token="hashed", models=["live-team"]))
assert denial is None
async def _run_vertex_live_route(client_frames, valid_token):
"""Drive the /vertex_ai/live route, so the gate under test is whichever one the route itself wires"""
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
vertex_ai_live_websocket_passthrough,
vertex_llm_base,
)
upstream_ws = RecordingUpstreamWebSocket()
websocket = _client_websocket(
AsyncMock(side_effect=[*client_frames, {"type": "websocket.disconnect"}]),
)
with (
_patched_websocket_passthrough_environment(upstream_ws) as environment,
patch.object( # test-quality-ok: an OAuth exchange with Google, so this is the boundary being faked
vertex_llm_base,
"_ensure_access_token_async",
AsyncMock(return_value=("token-abc", "proj-db")),
),
):
await vertex_ai_live_websocket_passthrough(
websocket=websocket,
model="gemini-live-2.5-flash",
vertex_project="proj-db",
vertex_location="global",
user_api_key_dict=valid_token,
)
return websocket, upstream_ws, environment
@pytest.mark.asyncio
async def test_vertex_live_route_gates_the_setup_frame_even_when_the_query_model_is_permitted():
"""
?model= is all connect-time auth sees, so a key that passes it and then names a different model in its setup
frame reaches Vertex on the gateway's own credentials unless the route gates the frames as well
"""
websocket, upstream_ws, environment = await _run_vertex_live_route(
[_setup_frame("gemini-3-pro-live")],
UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]),
)
upstream_ws.send.assert_not_awaited()
reason = _policy_close_reason(websocket)
assert reason is not None and "gemini-3-pro-live" in reason
refusal = environment.proxy_logging.post_call_failure_hook.await_args.kwargs["original_exception"]
assert refusal.code == "403"
assert refusal.type == ProxyErrorTypes.key_model_access_denied
@pytest.mark.asyncio
async def test_vertex_live_route_forwards_a_setup_frame_the_key_is_allowed_to_run():
websocket, upstream_ws, _ = await _run_vertex_live_route(
[_setup_frame("gemini-live-2.5-flash")],
UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]),
)
upstream_ws.send.assert_awaited_once()
assert _policy_close_reason(websocket) is None
async def _run_passthrough_with_frame_gate(gate):
upstream_ws = ClosableUpstreamWebSocket()
websocket = _client_websocket(
AsyncMock(
side_effect=[
{"type": "websocket.receive", "text": json.dumps({"setup": {"model": "gemini-live-2.5-flash"}})},
{"type": "websocket.disconnect"},
]
),
)
with _patched_websocket_passthrough_environment(upstream_ws) as environment:
await websocket_passthrough_request(
websocket=websocket,
target=VERTEX_LIVE_TARGET,
custom_headers={"Authorization": "Bearer token"},
user_api_key_dict=UserAPIKeyAuth(token="hashed"),
forward_headers=False,
endpoint="/vertex_ai/live",
accept_websocket=False,
client_frame_gate=gate,
)
return websocket, upstream_ws, environment
@pytest.mark.asyncio
async def test_websocket_passthrough_reports_a_refused_client_frame_as_a_failure():
"""A denied session logged through the success hooks tells spend tracking and audit that it succeeded"""
denial = ProxyException(
message="gemini-live-2.5-flash: key not allowed to access model",
type=ProxyErrorTypes.key_model_access_denied,
param="model",
code=403,
)
async def refusing_gate(frame_data, valid_token):
return denial
websocket, upstream_ws, environment = await _run_passthrough_with_frame_gate(refusing_gate)
assert _policy_close_reason(websocket) == "gemini-live-2.5-flash: key not allowed to access model"
upstream_ws.close.assert_awaited()
environment.proxy_logging.post_call_success_hook.assert_not_awaited()
environment.logging_worker.ensure_initialized_and_enqueue.assert_not_called()
environment.proxy_logging.post_call_failure_hook.assert_awaited_once()
failure_kwargs = environment.proxy_logging.post_call_failure_hook.await_args.kwargs
assert failure_kwargs["original_exception"] is denial
assert failure_kwargs["request_data"]["litellm_logging_obj"].litellm_call_id
@pytest.mark.asyncio
async def test_websocket_passthrough_still_reports_an_allowed_session_as_a_success():
async def permitting_gate(frame_data, valid_token):
return None
_, _, environment = await _run_passthrough_with_frame_gate(permitting_gate)
environment.proxy_logging.post_call_failure_hook.assert_not_awaited()
environment.proxy_logging.post_call_success_hook.assert_awaited_once()
environment.logging_worker.ensure_initialized_and_enqueue.assert_called_once()
@pytest.mark.asyncio
@pytest.mark.parametrize("rcvd_close", [None, "abnormal", "no_status"])
async def test_websocket_passthrough_does_not_relay_unsendable_upstream_close(rcvd_close):