scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 280

fix(proxy): log a refused Vertex Live setup model once, rewrite binary setup frames, and check a group an alias shadows as itself

Browse source 
A setup model refused by its re-authorization was logged by the auth handler and again by the passthrough, so it now carries a marker that it was already logged. Binary setup frames were checked as if rewritten but forwarded as sent, so they go through the same rewriter. A router alias named like a real group made the gate check the alias target while the rewriter sent the group's own model, so a name that is a real group is checked as that group. A regression test pins the release of the budget the re-authorization reserves.
This commit is contained in:
mateo-berri 2026-10-29 03:00:53 -07:00
parent 6d6bbfc7bb
commit 98492c716f
3 changed files with 135 additions and 43 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -127,7 +127,9 @@ from litellm.proxy.pass_through_endpoints.injected_stream_usage import (
without_injected_usage_frames,
)
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
ClientFrameRefusal,
HttpPassThroughEndpointHelpers,
LoggedClientFrameRefusal,
WebsocketClientFrameGate,
create_pass_through_route,
create_websocket_passthrough_route,
@ -3987,8 +3989,11 @@ def _vertex_live_model_names(setup_model: str, llm_router: Router | None) -> tup
"""
deployment: Final = _live_deployment(setup_model, llm_router)
if deployment is not None and _vertex_upstream_model(deployment) is not None:
alias_target: Final = None if llm_router is None else llm_router.get_model_from_alias(deployment["model_name"])
return (alias_target or deployment["model_name"],)
group: Final = deployment["model_name"]
alias_target: Final = (
None if llm_router is None or group in llm_router.model_names else llm_router.get_model_from_alias(group)
)
return (alias_target or group,)
opened_model: Final = _resolve_alias_to_upstream_model(setup_model, llm_router)
publisher_model: Final = _vertex_publisher_model_suffix(opened_model)
serving_groups: Final = tuple(
@ -4054,20 +4059,20 @@ VertexLiveModelAuthorizer: TypeAlias = Callable[[str], Awaitable[UserAPIKeyAuth]
async def _model_reauthorization_denial(
model: str, authorize_model: VertexLiveModelAuthorizer
) -> ProxyException | None:
) -> LoggedClientFrameRefusal | None:
from litellm.proxy.spend_tracking.budget_reservation import release_or_invalidate_budget_reservation
try:
model_auth: Final = await authorize_model(model)
except (ProxyException, HTTPException) as refusal:
return _as_model_denial(refusal)
return LoggedClientFrameRefusal(_as_model_denial(refusal))
await release_or_invalidate_budget_reservation(budget_reservation=model_auth.budget_reservation)
return None
async def _first_reauthorization_denial(
model_names: tuple[str, ...], authorize_model: VertexLiveModelAuthorizer
) -> ProxyException | None:
) -> LoggedClientFrameRefusal | None:
denial: Final = await _model_reauthorization_denial(model_names[0], authorize_model)
if denial is None or len(model_names) == 1:
return denial
@ -4080,7 +4085,7 @@ async def _vertex_live_model_denial(
valid_token: UserAPIKeyAuth,
llm_router: Router | None,
authorize_model: VertexLiveModelAuthorizer,
) -> ProxyException | None:
) -> ClientFrameRefusal | None:
"""
The setup rewriter forwards an alias name as sent unless it is a listed router alias, which already resolved to
its target group, so an alias name grants nothing here. A name the key may call is then authorized the way a
@ -4126,11 +4131,11 @@ async def _vertex_live_setup_denial(
valid_token: UserAPIKeyAuth,
llm_router: Router | None,
authorize_model: VertexLiveModelAuthorizer,
) -> ProxyException | None:
) -> ClientFrameRefusal | None:
match setup_model:
case str():
denial: Final = await _vertex_live_model_denial(setup_model, valid_token, llm_router, authorize_model)
if denial is not None:
if isinstance(denial, ProxyException):
log_model_access_denial(denial)
return denial
case _UnreadableVertexLiveFrame():
@ -4160,7 +4165,7 @@ def _build_vertex_live_client_frame_gate(
which is a caller with no model allowlist and no agent, end user or team grant standing in for one
"""
async def gate(frame_data: str | bytes, valid_token: UserAPIKeyAuth, /) -> ProxyException | None:
async def gate(frame_data: str | bytes, valid_token: UserAPIKeyAuth, /) -> ClientFrameRefusal | None:
model: Final = _vertex_live_setup_model(frame_data)
if model is None:
return None

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

@ -10,7 +10,7 @@ from dataclasses import dataclass
from datetime import datetime
from itertools import chain, count, groupby
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast
from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, TypedDict, cast, overload
from urllib.parse import urlencode, urlparse
import httpx
@ -2395,28 +2395,39 @@ def create_websocket_passthrough_route(
return websocket_endpoint_func
def _rewrite_vertex_live_setup_model(text_data: str, setup_model_rewriter: Callable[[str], str] | None) -> str:
@overload
def _rewrite_vertex_live_setup_model(frame_data: str, setup_model_rewriter: Callable[[str], str] | None) -> str: ...
@overload
def _rewrite_vertex_live_setup_model(frame_data: bytes, setup_model_rewriter: Callable[[str], str] | None) -> bytes: ...
def _rewrite_vertex_live_setup_model(
frame_data: str | bytes, setup_model_rewriter: Callable[[str], str] | None
) -> str | bytes:
"""
Rewrite the model of a Vertex AI Live ``setup`` frame, leaving every other frame byte-identical
"""
if setup_model_rewriter is None:
return text_data
return frame_data
try:
message: Final[object] = _DECODED_JSON.validate_python(json.loads(text_data))
except json.JSONDecodeError:
return text_data
message: Final[object] = _DECODED_JSON.validate_python(json.loads(frame_data))
except (json.JSONDecodeError, UnicodeDecodeError):
return frame_data
if not isinstance(message, dict):
return text_data
return frame_data
setup: Final = message.get("setup")
if not isinstance(setup, dict):
return text_data
return frame_data
setup_model: Final = setup.get("model")
if not isinstance(setup_model, str):
return text_data
return frame_data
rewritten_model: Final = setup_model_rewriter(setup_model)
if rewritten_model == setup_model:
return text_data
return json.dumps({**message, "setup": {**setup, "model": rewritten_model}})
return frame_data
rewritten_frame: Final = json.dumps({**message, "setup": {**setup, "model": rewritten_model}})
return rewritten_frame if isinstance(frame_data, str) else rewritten_frame.encode("utf-8")
def _resolved_vertex_live_setup(
@ -2504,15 +2515,32 @@ def _with_trace_context(headers: Mapping[str, str], parent_span: object) -> dict
return inject_trace_context(headers, parent_span=parent_span)
@dataclass(frozen=True, slots=True)
class LoggedClientFrameRefusal:
denial: ProxyException
ClientFrameRefusal: TypeAlias = ProxyException | LoggedClientFrameRefusal
def _client_frame_denial(refusal: ClientFrameRefusal) -> ProxyException:
match refusal:
case LoggedClientFrameRefusal(denial=denial):
return denial
case ProxyException():
return refusal
class WebsocketClientFrameGate(Protocol):
"""Authorize one client frame, returning the denial to report or ``None`` to let it through.
"""Authorize one client frame, returning the denial to report, one its check already reported, or ``None`` to
let it through.
A websocket passthrough authenticates before any frame arrives, so a protocol that names its model
inside a frame has no model to authorize at connect time. The frame, not the extracted model, is the
argument because where a model sits is part of each provider's wire protocol
"""
async def __call__(self, frame_data: str | bytes, valid_token: UserAPIKeyAuth, /) -> ProxyException | None: ...
async def __call__(self, frame_data: str | bytes, valid_token: UserAPIKeyAuth, /) -> ClientFrameRefusal | None: ...
async def _relay_client_frame(
@ -2521,7 +2549,7 @@ async def _relay_client_frame(
forwarded_frame: str | bytes,
gate: WebsocketClientFrameGate | None,
valid_token: UserAPIKeyAuth,
) -> ProxyException | None:
) -> ClientFrameRefusal | None:
"""Forward one client frame upstream, or refuse it and hand back the denial that ended the session.
The denial travels out so the caller reports a failed session instead of leaving the success hooks to
@ -2693,7 +2721,7 @@ async def websocket_passthrough_request(
"WebSocket passthrough (%s): Upstream connection established successfully", endpoint
)
async def forward_client_to_upstream() -> ProxyException | None:
async def forward_client_to_upstream() -> ClientFrameRefusal | None:
"""Forward messages from client to upstream WebSocket, returning any denial that ended it"""
try:
while True:
@ -2778,7 +2806,7 @@ async def websocket_passthrough_request(
bytes_refusal := await _relay_client_frame(
upstream_ws,
bytes_data,
bytes_data,
_rewrite_vertex_live_setup_model(bytes_data, setup_model_rewriter),
client_frame_gate,
user_api_key_dict,
)
@ -2862,15 +2890,18 @@ async def websocket_passthrough_request(
# upstream reader's next send fail, so raising that instead would report the wrong cause.
# Turns that completed before a refusal were served, so they are logged and billed as a success
client_frame_refusal: Final = next(
(result for result in task_results if isinstance(result, ProxyException)), None
(result for result in task_results if isinstance(result, (ProxyException, LoggedClientFrameRefusal))),
None,
)
if client_frame_refusal is not None:
await websocket.close(
code=CloseCode.POLICY_VIOLATION, reason=_refusal_close_reason(client_frame_refusal.message)
code=CloseCode.POLICY_VIOLATION,
reason=_refusal_close_reason(_client_frame_denial(client_frame_refusal).message),
)
await upstream_ws.close()
if client_frame_refusal is not None and not websocket_messages:
if isinstance(client_frame_refusal, ProxyException) and not websocket_messages:
await log_failure(client_frame_refusal)
if client_frame_refusal is not None and not websocket_messages:
return
# Check for exceptions in completed tasks

86
tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py
View file

@ -7342,7 +7342,7 @@ def _patched_websocket_passthrough_environment(upstream_ws):
):
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={})
mock_proxy_logging.post_call_success_hook = AsyncMock()
mock_proxy_logging.post_call_failure_hook = AsyncMock()
mock_proxy_logging.post_call_failure_hook = AsyncMock(return_value=None)
enqueued_handler = MagicMock()
mock_worker.ensure_initialized_and_enqueue = MagicMock(
side_effect=lambda async_coroutine: (
@ -7703,6 +7703,19 @@ async def test_vertex_live_setup_frame_is_gated_when_it_arrives_as_binary():
assert _policy_close_reason(websocket) is not None
@pytest.mark.asyncio
async def test_vertex_live_binary_setup_frame_reaches_vertex_as_the_model_the_gate_authorized() -> None:
binary_setup: Final = json.dumps({"setup": {"model": f"publishers/google/models/{VERTEX_LIVE_ALIAS_GROUP}"}})
_, upstream_ws = await _run_vertex_live_gated_passthrough(
[{"type": "websocket.receive", "bytes": binary_setup.encode()}],
UserAPIKeyAuth(token="hashed", models=[VERTEX_LIVE_ALIAS_GROUP]),
)
sent_frame: Final = upstream_ws.send.await_args.args[0]
assert isinstance(sent_frame, bytes)
assert json.loads(sent_frame)["setup"]["model"] == VERTEX_LIVE_FULL_PATH.format(model=VERTEX_LIVE_GROUP)
@pytest.mark.asyncio
async def test_vertex_live_setup_frame_with_an_escaped_key_is_still_gated():
"""Large audio frames skip the parse on a substring test, so a JSON-escaped key must not slip past it"""
@ -8186,13 +8199,22 @@ async def test_vertex_live_setup_frame_naming_an_alias_cannot_open_the_model_the
@pytest.mark.asyncio
@pytest.mark.parametrize("hidden", [False, True], ids=["listed-alias", "hidden-alias"])
async def test_vertex_live_setup_frame_naming_a_router_alias_cannot_open_the_model_the_alias_shadows(hidden):
"""A listed router alias runs as its target group, while the setup rewriter cannot see a hidden one and sends it
to Vertex as the model whose id it reuses, which a key holding only the target must not open"""
llm_router = litellm.Router(
@pytest.mark.parametrize(
("hidden", "shadows_a_group"),
[(False, False), (True, False), (False, True), (True, True)],
ids=["listed-alias", "hidden-alias", "listed-alias-named-like-a-group", "hidden-alias-named-like-a-group"],
)
async def test_vertex_live_setup_frame_naming_a_router_alias_cannot_open_the_model_the_alias_shadows(
hidden: bool, shadows_a_group: bool
) -> None:
"""A listed router alias runs as its target group, while the setup rewriter cannot see a hidden one, or one named
like a group, and sends it to Vertex as the model whose id it reuses, which a key holding only the target must
not open"""
shadowed_group: Final = {"model_name": "gemini-3-pro-live", "litellm_params": {"model": "vertex_ai/gemini-3-pro-live"}}
llm_router: Final = litellm.Router(
model_list=[
{"model_name": "gemini-live-2.5-flash", "litellm_params": {"model": "vertex_ai/gemini-live-2.5-flash"}},
*([shadowed_group] if shadows_a_group else []),
],
model_group_alias={"gemini-3-pro-live": {"model": "gemini-live-2.5-flash", "hidden": hidden}},
)
@ -8205,7 +8227,11 @@ async def test_vertex_live_setup_frame_naming_a_router_alias_cannot_open_the_mod
sent_models = [json.loads(call.args[0])["setup"]["model"] for call in upstream_ws.send.await_args_list]
assert not any(model.endswith("/gemini-3-pro-live") for model in sent_models)
assert sent_models == ([] if hidden else ["projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash"])
assert sent_models == (
[]
if hidden or shadows_a_group
else ["projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash"]
)
@pytest.mark.asyncio
@ -8416,20 +8442,50 @@ async def test_vertex_live_route_authorizes_the_setup_model_the_way_connect_time
param="model",
code=400,
)
with patch(
"litellm.proxy.auth.user_api_key_auth.user_api_key_auth", autospec=True, side_effect=budget_denial
) as reauthorization:
websocket, upstream_ws, _ = await _run_vertex_live_route(
valid_token: Final = UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"])
with (
patch(
"litellm.proxy.auth.user_api_key_auth.user_api_key_auth_builder", autospec=True, return_value=valid_token
) as reauthentication,
patch(
"litellm.proxy.auth.user_api_key_auth.run_centralized_common_checks",
autospec=True,
side_effect=budget_denial,
),
):
websocket, upstream_ws, environment = await _run_vertex_live_route(
[_setup_frame("models/gemini-live-2.5-flash")],
UserAPIKeyAuth(token="hashed", models=["gemini-live-2.5-flash"]),
valid_token,
reauthorize_through_user_api_key_auth=True,
)
upstream_ws.send.assert_not_awaited()
assert _policy_close_reason(websocket) == budget_denial.message
reauthorized_request = reauthorization.call_args.kwargs["request"]
assert json.loads(await reauthorized_request.body()) == {"model": "gemini-live-2.5-flash"}
assert reauthorization.call_args.kwargs["api_key"] == "Bearer sk-live-client"
reauthenticated_request: Final = reauthentication.call_args.kwargs["request"]
assert json.loads(await reauthenticated_request.body()) == {"model": "gemini-live-2.5-flash"}
assert reauthentication.call_args.kwargs["api_key"] == "Bearer sk-live-client"
environment.proxy_logging.post_call_failure_hook.assert_awaited_once()
reported: Final = environment.proxy_logging.post_call_failure_hook.await_args.kwargs
assert reported["original_exception"].message == budget_denial.message
assert reported["request_data"]["model"] == "gemini-live-2.5-flash"
@pytest.mark.asyncio
async def test_vertex_live_route_releases_the_budget_it_reserved_to_authorize_an_allowed_setup_model():
"""
Re-authorizing the setup model reserves the request's worst-case cost the way connect-time auth does, and the
session's spend is settled on the connect-time token, so an unreleased reservation pins the key's spend counter
and throttles its later sessions until the counter expires
"""
valid_token: Final = UserAPIKeyAuth(
token="hashed",
models=["gemini-live-2.5-flash"],
budget_reservation={"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False},
)
_, upstream_ws, _ = await _run_vertex_live_route([_setup_frame("models/gemini-live-2.5-flash")], valid_token)
upstream_ws.send.assert_awaited_once()
assert valid_token.budget_reservation is not None and valid_token.budget_reservation["finalized"] is True
@pytest.mark.asyncio