From 2d860f05d4ae0c8b9d52b5424f6eaf3467fdaf53 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 01:17:07 -0700 Subject: [PATCH] fix(proxy): bill pipelined websocket turns to their own frame and close the review gaps --- litellm/litellm_core_utils/core_helpers.py | 6 +- litellm/proxy/common_utils/realtime_utils.py | 3 +- .../llm_passthrough_endpoints.py | 77 +++++-- ...i_websocket_passthrough_logging_handler.py | 29 ++- .../pass_through_endpoints.py | 188 +++++++++++----- .../proxy/common_utils/test_realtime_cache.py | 7 + ...i_websocket_passthrough_logging_handler.py | 45 +++- .../test_pass_through_endpoints.py | 212 ++++++++++++++++-- .../test_success_handler.py | 2 +- .../test_openai_ws_passthrough_routes.py | 55 +++++ 10 files changed, 530 insertions(+), 94 deletions(-) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 55025508e40..58412198064 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -3,7 +3,7 @@ import copy import logging import re -from collections.abc import Collection, Iterable, Iterator, Mapping +from collections.abc import Collection, Iterable, Iterator, Mapping, Sequence from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast, runtime_checkable @@ -59,6 +59,10 @@ def as_str_mapping(value: object) -> Mapping[str, object] | None: return value if isinstance(value, Mapping) else None +def as_object_list(value: object) -> Sequence[object] | None: + return value if isinstance(value, list) else None + + def decoded_json_object(value: object) -> Mapping[str, object]: mapping: Final = as_str_mapping(value) if mapping is not None: diff --git a/litellm/proxy/common_utils/realtime_utils.py b/litellm/proxy/common_utils/realtime_utils.py index c641748843e..e96685bdeae 100644 --- a/litellm/proxy/common_utils/realtime_utils.py +++ b/litellm/proxy/common_utils/realtime_utils.py @@ -1,3 +1,4 @@ +import json from functools import lru_cache from typing import Final @@ -10,7 +11,7 @@ def realtime_request_body(model: str | None) -> bytes: Generate the realtime websocket request body. Cached with LRU semantics to avoid repeated string formatting work while keeping memory usage bounded. """ - return f'{{"model": "{model or ""}"}}'.encode() + return json.dumps({"model": model or ""}).encode() _realtime_request_body: Final = realtime_request_body diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 2e51bc3465b..57c0e78c462 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -15,7 +15,7 @@ import os import posixpath import re import sys -from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from functools import partial from types import MappingProxyType @@ -50,7 +50,7 @@ from litellm.constants import ( 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.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping, attribute_or_none from litellm.llms.anthropic.common_utils import AnthropicModelInfo, merge_anthropic_beta_headers from litellm.llms.azure.passthrough.transformation import ( foreign_azure_deployment, @@ -73,7 +73,12 @@ from litellm.llms.oss_decision import OssDecisionProvider, oss_connection, valid from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error from litellm.proxy._types import * -from litellm.proxy.auth.auth_checks import can_agent_call_model, can_key_call_resolved_model +from litellm.proxy.auth.auth_checks import ( + can_agent_call_model, + can_key_call_resolved_model, + key_model_aliases_for_auth_check, + team_model_aliases_for_auth_check, +) 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 ( # noqa: F401 # legacy module exports @@ -102,6 +107,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg safe_get_request_headers, safe_set_request_parsed_body, ) +from litellm.proxy.common_utils.model_listing_utils import alias_map from litellm.proxy.common_utils.openai_error_payload import LITELLM_CALL_ID_HEADER from litellm.proxy.common_utils.resource_ownership import is_proxy_admin from litellm.proxy.common_utils.sse_keepalive import ( @@ -121,6 +127,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket openai_websocket_frame_kind, ) from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + STORED_PROMPT_PROBE_MODEL, HttpPassThroughEndpointHelpers, WebsocketFrameModelAllowed, WebsocketFrameModelGate, @@ -132,7 +139,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( websocket_passthrough_request, ) from litellm.proxy.utils import ProxyLogging as ProxyLoggingType -from litellm.proxy.utils import is_known_model +from litellm.proxy.utils import is_known_model, normalize_route_for_root_path from litellm.proxy.vector_store_endpoints.utils import ( assert_proxy_admin_for_vector_store_index_management, assert_user_can_access_vector_store, @@ -3304,6 +3311,11 @@ _OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL: Final = ( "and ask a proxy admin to check the gateway logs if it keeps happening." ) +_OPENAI_WS_ALIASED_MODEL_REFUSAL: Final = ( + "This key cannot use a model alias over OpenAI websocket passthrough. The passthrough sends each frame to " + "OpenAI unchanged, so name the model this key may use directly." +) + _OPENAI_WS_DISABLED_REFUSAL: Final = _OpenAIWebsocketRefusal( close_reason="OpenAI websocket passthrough is disabled", message=( @@ -3357,36 +3369,65 @@ def _frame_model_authorizer() -> FrameModelAuthorizer: return user_api_key_auth_accepted_websocket_for_model +def _proxy_model_list() -> Sequence[object] | None: + from litellm.proxy import proxy_server + + return as_object_list(attribute_or_none(proxy_server, "llm_model_list")) + + +def _gateway_remaps_model(model: str, frame_auth: UserAPIKeyAuth, llm_router: Router | None) -> bool: + alias_maps: Final = ( + key_model_aliases_for_auth_check(frame_auth) or {}, + team_model_aliases_for_auth_check(frame_auth) or {}, + alias_map(litellm.model_alias_map), + ) + targets: Final = ( + *(aliases.get(model) for aliases in alias_maps), + llm_router.get_model_from_alias(model) if llm_router is not None else None, + ) + return any(target not in (None, model) for target in targets) + + +def _aliased_model_verdict( + allowed: WebsocketFrameModelAllowed, probe: WebsocketFrameModelVerdict +) -> WebsocketFrameModelVerdict: + match probe: + case WebsocketFrameModelRefused(error_type=error_type) if error_type.endswith("_model_access_denied"): + return WebsocketFrameModelRefused(message=_OPENAI_WS_ALIASED_MODEL_REFUSAL, error_type=error_type) + case WebsocketFrameModelRefused(): + return probe + case WebsocketFrameModelAllowed(): + return allowed + + def _proxy_frame_model_gate( authorize_frame_model: Annotated[FrameModelAuthorizer, Depends(_frame_model_authorizer)], ) -> WebsocketFrameModelGate: - from litellm.proxy.proxy_server import ( - llm_model_list, - llm_router, - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) + from litellm.proxy import proxy_server from litellm.proxy.spend_tracking.budget_reservation import release_or_invalidate_budget_reservation - async def authorize(websocket: WebSocket, model: str) -> WebsocketFrameModelAllowed: + async def authorize(websocket: WebSocket, model: str) -> WebsocketFrameModelVerdict: frame_auth: Final = await authorize_frame_model(websocket, model) access_groups: Final = tuple(frame_auth.matched_model_access_groups or ()) + llm_router: Final = proxy_server.llm_router try: await can_key_call_resolved_model( - model=model, llm_model_list=llm_model_list, valid_token=frame_auth, llm_router=llm_router + model=model, llm_model_list=_proxy_model_list(), valid_token=frame_auth, llm_router=llm_router ) await can_agent_call_model( model=model, valid_token=frame_auth, llm_router=llm_router, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, + prisma_client=proxy_server.prisma_client, + user_api_key_cache=proxy_server.user_api_key_cache, + proxy_logging_obj=proxy_server.proxy_logging_obj, ) finally: await release_or_invalidate_budget_reservation(budget_reservation=frame_auth.budget_reservation) - return WebsocketFrameModelAllowed(access_groups=access_groups) + allowed: Final = WebsocketFrameModelAllowed(access_groups=access_groups) + if not _gateway_remaps_model(model, frame_auth, llm_router): + return allowed + return _aliased_model_verdict(allowed, await gate(websocket, STORED_PROMPT_PROBE_MODEL)) async def gate(websocket: WebSocket, model: str, /) -> WebsocketFrameModelVerdict: try: @@ -3478,7 +3519,7 @@ async def openai_websocket_proxy_route( custom_headers=custom_headers, user_api_key_dict=user_api_key_dict, forward_headers=False, - endpoint=websocket.url.path, + endpoint=normalize_route_for_root_path(websocket.url.path) or websocket.url.path, accept_websocket=False, client_frame_model_gate=frame_model_gate, upstream_frame_kind=openai_websocket_frame_kind, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py index c2f4e85005c..951764f74b0 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py @@ -9,6 +9,7 @@ from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger from litellm.cost_calculator import RealtimeAPITokenUsageProcessor, handle_realtime_stream_cost_calculation +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.responses.utils import ResponseAPILoggingUtils @@ -139,15 +140,18 @@ def _responses_turn(event: Mapping[str, object], session_model: str) -> _Complet ) -def _responses_billing(messages: WebsocketMessages, model: str) -> _Billing | None: +def _responses_billing(messages: WebsocketMessages, model: str, requested_model: str | None) -> _Billing | None: turns: Final = tuple(turn for turn in (_responses_turn(event, model) for event in messages) if turn is not None) if not turns: return None usage: Final = RealtimeAPITokenUsageProcessor.combine_usage_objects([turn.usage for turn in turns]) - return _Billing(usage=usage, cost=sum(_responses_cost(turn.usage, turn.model) for turn in turns)) + return _Billing( + usage=usage, + cost=sum(_responses_cost(turn.usage, (turn.model, requested_model)) for turn in turns), + ) -def _responses_cost(usage: Usage, model: str) -> float: +def _model_cost(usage: Usage, model: str) -> float | None: try: return litellm.completion_cost( completion_response=ModelResponse(model=model, usage=usage), @@ -155,10 +159,25 @@ def _responses_cost(usage: Usage, model: str) -> float: custom_llm_provider=OPENAI_WEBSOCKET_PROVIDER, ) except Exception: # noqa: BLE001 # completion_cost raises a bare Exception for an unpriced model, and the turn must still log + return None + + +def _responses_cost(usage: Usage, price_as: Sequence[str | None]) -> float: + costs: Final = (_model_cost(usage, model) for model in dict.fromkeys(price_as) if model is not None) + cost: Final = next((cost for cost in costs if cost), None) + if cost is None: verbose_proxy_logger.warning( - "OpenAI websocket passthrough: could not price a Responses turn on %s, logging it at zero", model + "OpenAI websocket passthrough: could not price a Responses turn on %s, logging it at zero", price_as ) return 0.0 + return cost + + +def _requested_model(kwargs: Mapping[str, object]) -> str | None: + litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) + metadata: Final = None if litellm_params is None else as_str_mapping(litellm_params.get("metadata")) + model_group: Final = None if metadata is None else metadata.get("model_group") + return model_group if isinstance(model_group, str) else None def _priced_result(model: str, billing: _Billing, start_time: datetime) -> ModelResponse: @@ -195,7 +214,7 @@ class OpenAIWebsocketPassthroughLoggingHandler: ) billing: Final = _realtime_billing(websocket_messages, model, logging_obj) or _responses_billing( - websocket_messages, model + websocket_messages, model, _requested_model(kwargs) ) if billing is None: verbose_proxy_logger.debug("OpenAI websocket passthrough (%s): no usage reported, skipping cost", url_route) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index befdc821734..af67a451c96 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -13,9 +13,10 @@ from functools import reduce from itertools import chain, count, groupby from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, TypedDict, cast -from urllib.parse import urlencode, urlparse +from urllib.parse import parse_qs, urlencode, urlparse import httpx +import orjson from fastapi import ( APIRouter, Depends, @@ -60,6 +61,7 @@ from litellm.integrations.custom_guardrail import ( from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.otel.model.semconv import PASSTHROUGH_OPERATION_KEY, GenAIOperation from litellm.litellm_core_utils.core_helpers import ( + as_object_list, as_str_mapping, bind_budget_reservation_to_callbacks, get_metadata_variable_name_from_kwargs, @@ -2441,11 +2443,16 @@ def _resolved_vertex_live_setup( return {**setup_data, "model": setup_model_rewriter(setup_model)} +def _decoded_frame(frame: str | bytes) -> object: + try: + return _DECODED_JSON.validate_python(orjson.loads(frame)) + except orjson.JSONDecodeError: + return _DECODED_JSON.validate_python(json.loads(frame if isinstance(frame, str) else frame.decode("utf-8"))) + + def _json_object_frame(frame: str | bytes) -> dict[str, object] | None: try: - decoded: Final[object] = _DECODED_JSON.validate_python( - json.loads(frame if isinstance(frame, str) else frame.decode("utf-8")) - ) + decoded: Final = _decoded_frame(frame) except (json.JSONDecodeError, UnicodeDecodeError): return None return decoded if isinstance(decoded, dict) else None @@ -2525,9 +2532,16 @@ class WebsocketFrameModelGate(Protocol): WebsocketUpstreamFrameClassifier: TypeAlias = Callable[[Mapping[str, object]], WebsocketUpstreamFrameKind] +@dataclass(frozen=True, slots=True) +class _TurnAuthorization: + requested_model: str | None + access_groups: tuple[str, ...] + + @dataclass(frozen=True, slots=True) class _ForwardClientFrame: - turn_access_groups: tuple[str, ...] | None + starts_turn: bool + turn: _TurnAuthorization | None @dataclass(frozen=True, slots=True) @@ -2537,24 +2551,35 @@ class _RefuseClientFrame: @dataclass(frozen=True, slots=True) class _ClientFrameRequests: + starts_turn: bool turn_model: str | None models: tuple[str, ...] - names_stored_prompt: bool + unverifiable_model_refusal: str | None -_NO_CLIENT_FRAME_REQUESTS: Final = _ClientFrameRequests(turn_model=None, models=(), names_stored_prompt=False) -_FRAME_TEXT_WORTH_PARSING: Final = ('"model"', '"prompt"', "\\") +_NO_CLIENT_FRAME_REQUESTS: Final = _ClientFrameRequests( + starts_turn=False, turn_model=None, models=(), unverifiable_model_refusal=None +) +_FORWARD_WITHOUT_TURN: Final = _ForwardClientFrame(starts_turn=False, turn=None) +_TURN_STARTING_FRAME: Final = "response.create" +_FRAME_TEXT_WORTH_PARSING: Final = ('"model"', '"prompt"', '"tools"', "\\") _TURN_MODEL_PLACEMENTS: Final = (("model",), ("response", "model"), ("session", "model")) _TRANSCRIPTION_MODEL_PLACEMENTS: Final = ( ("session", "input_audio_transcription", "model"), ("session", "audio", "input", "transcription", "model"), ) _STORED_PROMPT_PLACEMENTS: Final = (("prompt",), ("response", "prompt"), ("session", "prompt")) +_TOOL_LIST_PLACEMENTS: Final = (("tools",), ("response", "tools"), ("session", "tools")) +_TOOL_THAT_RUNS_A_MODEL: Final = "image_generation" STORED_PROMPT_PROBE_MODEL: Final = "litellm-unverifiable-stored-prompt-model" _STORED_PROMPT_REFUSAL: Final = ( "This key is limited to specific models, so it cannot run a stored prompt. " "The gateway cannot check which model a stored prompt runs." ) +_UNNAMED_TOOL_MODEL_REFUSAL: Final = ( + "This key is limited to specific models, so an image_generation tool has to name its model. " + "The gateway cannot check which model an unnamed one runs." +) _FIRST_SENTENCE: Final = re.compile(r".*?[.!?](?=\s|$)") @@ -2577,6 +2602,21 @@ def _placed_models(frame: Mapping[str, object], placements: Iterable[Sequence[st return tuple(value for value in placed if isinstance(value, str)) +def _placed_tools(frame: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + tool_lists: Final = (_placed_value(frame, placement) for placement in _TOOL_LIST_PLACEMENTS) + tools: Final = chain.from_iterable(as_object_list(tool_list) or () for tool_list in tool_lists) + return tuple(tool for tool in map(as_str_mapping, tools) if tool is not None) + + +def _unverifiable_model_refusal(frame: Mapping[str, object], tools: Sequence[Mapping[str, object]]) -> str | None: + prompts: Final = (_placed_value(frame, placement) for placement in _STORED_PROMPT_PLACEMENTS) + if any(isinstance(prompt, Mapping) for prompt in prompts): + return _STORED_PROMPT_REFUSAL + if any(tool.get("type") == _TOOL_THAT_RUNS_A_MODEL and not isinstance(tool.get("model"), str) for tool in tools): + return _UNNAMED_TOOL_MODEL_REFUSAL + return None + + def _client_frame_requests(frame_data: str | bytes) -> _ClientFrameRequests: if isinstance(frame_data, str) and not _may_name_a_model(frame_data): return _NO_CLIENT_FRAME_REQUESTS @@ -2584,17 +2624,21 @@ def _client_frame_requests(frame_data: str | bytes) -> _ClientFrameRequests: if frame is None: return _NO_CLIENT_FRAME_REQUESTS turn_models: Final = _placed_models(frame, _TURN_MODEL_PLACEMENTS) - prompts: Final = (_placed_value(frame, placement) for placement in _STORED_PROMPT_PLACEMENTS) + tools: Final = _placed_tools(frame) + tool_models: Final = (model for model in (tool.get("model") for tool in tools) if isinstance(model, str)) return _ClientFrameRequests( + starts_turn=frame.get("type") == _TURN_STARTING_FRAME, turn_model=next(iter(turn_models), None), - models=tuple(dict.fromkeys((*turn_models, *_placed_models(frame, _TRANSCRIPTION_MODEL_PLACEMENTS)))), - names_stored_prompt=any(isinstance(prompt, Mapping) for prompt in prompts), + models=tuple( + dict.fromkeys((*turn_models, *_placed_models(frame, _TRANSCRIPTION_MODEL_PLACEMENTS), *tool_models)) + ), + unverifiable_model_refusal=_unverifiable_model_refusal(frame, tools), ) -def _stored_prompt_refusal(probe: WebsocketFrameModelRefused) -> WebsocketFrameModelRefused: +def _unverifiable_model_probe_refusal(probe: WebsocketFrameModelRefused, message: str) -> WebsocketFrameModelRefused: if probe.error_type.endswith("_model_access_denied"): - return WebsocketFrameModelRefused(message=_STORED_PROMPT_REFUSAL, error_type=probe.error_type) + return WebsocketFrameModelRefused(message=message, error_type=probe.error_type) return probe @@ -2602,20 +2646,25 @@ async def _client_frame_verdict( websocket: WebSocket, frame_data: str | bytes, gate: WebsocketFrameModelGate | None ) -> _ForwardClientFrame | _RefuseClientFrame: if gate is None: - return _ForwardClientFrame(turn_access_groups=None) + return _FORWARD_WITHOUT_TURN requests: Final = _client_frame_requests(frame_data) verdicts: Final = tuple([await gate(websocket, model) for model in requests.models]) refusal: Final = next((verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelRefused)), None) if refusal is not None: return _RefuseClientFrame(refusal=refusal) - if requests.names_stored_prompt and isinstance( + if requests.unverifiable_model_refusal is not None and isinstance( probe := await gate(websocket, STORED_PROMPT_PROBE_MODEL), WebsocketFrameModelRefused ): - return _RefuseClientFrame(refusal=_stored_prompt_refusal(probe)) + return _RefuseClientFrame(refusal=_unverifiable_model_probe_refusal(probe, requests.unverifiable_model_refusal)) allowed: Final = tuple(verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelAllowed)) - if requests.turn_model is None or not allowed: - return _ForwardClientFrame(turn_access_groups=None) - return _ForwardClientFrame(turn_access_groups=allowed[0].access_groups) + return _ForwardClientFrame( + starts_turn=requests.starts_turn, + turn=( + _TurnAuthorization(requested_model=requests.turn_model, access_groups=allowed[0].access_groups) + if requests.turn_model is not None and allowed + else None + ), + ) def _refusal_close_reason(message: str) -> str: @@ -2631,12 +2680,54 @@ async def _refuse_client_frame( await upstream_ws.close() -def _stamp_billed_access_groups(kwargs: Mapping[str, object], access_groups: tuple[str, ...]) -> None: - """Charge a turn to the access groups that authorized its model, and to none when no group did.""" +def _stamp_billed_turn(kwargs: Mapping[str, object], turn: _TurnAuthorization) -> None: + """Charge a turn to the access groups that authorized its model, and its model budget to the model it asked for, + since the provider may answer under a dated snapshot name the budget was never set on.""" litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) metadata: Final = None if litellm_params is None else litellm_params.get("metadata") - if isinstance(metadata, MutableMapping): - metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = list(access_groups) + if not isinstance(metadata, MutableMapping): + return + metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = list(turn.access_groups) + if turn.requested_model is not None: + metadata["model_group"] = turn.requested_model + + +def _opening_model(target: str) -> str | None: + return next(iter(parse_qs(urlparse(target).query).get("model", ())), None) + + +class _WebsocketTurnLedger: + """The groups that authorized each turn not yet logged, when the next row starts, and how many rows were logged. + + The provider runs a socket's response.create frames one at a time in the order they were sent, so each row is + charged to the groups that allowed the oldest turn still unlogged. A realtime session keeps the model it opened on, + so its rows stay charged to the groups that opened it. + """ + + def __init__(self, opening: _TurnAuthorization, started: datetime) -> None: + self.opening = opening + self.queued_turns: tuple[_TurnAuthorization | None, ...] = () + self.session_opened = False + self.started = started + self.logged_turns = 0 + + @property + def current(self) -> _TurnAuthorization: + oldest: Final = self.queued_turns[0] if self.queued_turns else None + return self.opening if oldest is None else oldest + + def queue_turn(self, turn: _TurnAuthorization | None) -> None: + if not self.session_opened: + self.queued_turns = (*self.queued_turns, turn) + + def open_session(self) -> None: + self.session_opened = True + self.queued_turns = () + + def close_turn(self, end_time: datetime) -> None: + self.logged_turns += 1 + self.started = end_time + self.queued_turns = self.queued_turns[1:] async def _relay_client_frame( @@ -2645,12 +2736,16 @@ async def _relay_client_frame( checked_frame: str | bytes, forwarded_frame: str | bytes, gate: WebsocketFrameModelGate | None, + ledger: _WebsocketTurnLedger, ) -> _ForwardClientFrame | _RefuseClientFrame: """Forward one client frame upstream, or refuse it and end the session.""" verdict: Final = await _client_frame_verdict(websocket, checked_frame, gate) match verdict: case _RefuseClientFrame(refusal=refusal): await _refuse_client_frame(websocket, upstream_ws, refusal) + case _ForwardClientFrame(starts_turn=True, turn=turn): + ledger.queue_turn(turn) + await upstream_ws.send(forwarded_frame) case _ForwardClientFrame(): await upstream_ws.send(forwarded_frame) return verdict @@ -2659,23 +2754,6 @@ async def _relay_client_frame( _EXPLICIT_SUCCESS_KWARGS: Final = frozenset({"logging_obj", "passthrough_logging_payload", "custom_llm_provider"}) -class _WebsocketTurnLedger: - """The groups that authorized the turn in flight, when it started, and how many turns the socket has logged.""" - - def __init__(self, access_groups: tuple[str, ...], started: datetime) -> None: - self.access_groups = access_groups - self.started = started - self.logged_turns = 0 - - def authorize(self, access_groups: tuple[str, ...] | None) -> None: - if access_groups is not None: - self.access_groups = access_groups - - def close_turn(self, end_time: datetime) -> None: - self.logged_turns += 1 - self.started = end_time - - async def websocket_passthrough_request( websocket: WebSocket, target: str, @@ -2835,7 +2913,11 @@ async def websocket_passthrough_request( self.request = MockWebSocketRequest(target_url) ledger: Final = _WebsocketTurnLedger( - access_groups=tuple(user_api_key_dict.matched_model_access_groups or ()), started=start_time + opening=_TurnAuthorization( + requested_model=_opening_model(target), + access_groups=tuple(user_api_key_dict.matched_model_access_groups or ()), + ), + started=start_time, ) def log_frames( @@ -2856,7 +2938,7 @@ async def websocket_passthrough_request( "end_time": end_time, } if client_frame_model_gate is not None: - _stamp_billed_access_groups(call_kwargs, ledger.access_groups) + _stamp_billed_turn(call_kwargs, ledger.current) success_kwargs: Final = { key: value for key, value in call_kwargs.items() if key not in _EXPLICIT_SUCCESS_KWARGS } @@ -2988,26 +3070,25 @@ async def websocket_passthrough_request( # Not a JSON message or doesn't contain setup data if isinstance( - relayed_text := await _relay_client_frame( + await _relay_client_frame( websocket, upstream_ws, text_data, _rewrite_vertex_live_setup_model(text_data, setup_model_rewriter), client_frame_model_gate, + ledger, ), _RefuseClientFrame, ): break - ledger.authorize(relayed_text.turn_access_groups) elif bytes_data is not None: if isinstance( - relayed_bytes := await _relay_client_frame( - websocket, upstream_ws, bytes_data, bytes_data, client_frame_model_gate + await _relay_client_frame( + websocket, upstream_ws, bytes_data, bytes_data, client_frame_model_gate, ledger ), _RefuseClientFrame, ): break - ledger.authorize(relayed_bytes.turn_access_groups) except asyncio.CancelledError: raise except Exception: @@ -3042,8 +3123,13 @@ async def websocket_passthrough_request( _extract_vertex_live_model_from_setup_response(message_data) return websocket_messages.append(message_data) - if upstream_frame_kind is not None and upstream_frame_kind(message_data) == "ends_turn": - log_turn() + match None if upstream_frame_kind is None else upstream_frame_kind(message_data): + case "opens_session": + ledger.open_session() + case "ends_turn": + log_turn() + case _: + pass async def relay_upstream_frame(upstream_message: str | bytes) -> None: record_upstream_frame(upstream_message) @@ -3105,7 +3191,9 @@ async def websocket_passthrough_request( reason=_truncated_close_reason(upstream_close.reason), ) - if ledger.logged_turns == 0: + if upstream_frame_kind is None or any( + upstream_frame_kind(frame) != "opens_session" for frame in websocket_messages + ): log_turn() # Call the proxy logging success hook diff --git a/tests/unit/proxy/common_utils/test_realtime_cache.py b/tests/unit/proxy/common_utils/test_realtime_cache.py index 6f5751fa923..6177a00d0aa 100644 --- a/tests/unit/proxy/common_utils/test_realtime_cache.py +++ b/tests/unit/proxy/common_utils/test_realtime_cache.py @@ -1,3 +1,4 @@ +import json from typing import Any, cast import pytest @@ -40,6 +41,12 @@ def test_realtime_request_body_caches_each_model_separately(): assert gpt4o_body_first is not gpt4o_mini_body +def test_realtime_request_body_keeps_a_model_name_with_quotes_inside_its_own_field(): + body = realtime_request_body('gpt-4o", "user": "someone-else') + + assert json.loads(body) == {"model": 'gpt-4o", "user": "someone-else'} + + def test_realtime_query_params_template_caches_each_pair_separately(): params_with_intent_first = _realtime_query_params_template("gpt-4o", "intent-a") params_with_intent_second = _realtime_query_params_template("gpt-4o", "intent-a") diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py index 253634760fa..c259d89c07f 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping, Sequence from datetime import datetime from types import MappingProxyType from typing import Final @@ -22,11 +23,11 @@ def _logging_obj() -> MagicMock: return logging_obj -def _session_created(model: str) -> dict: +def _session_created(model: str) -> dict[str, object]: return {"type": "session.created", "session": {"id": "sess_fake", "model": model}} -def _realtime_response_done(text_in: int, audio_in: int, text_out: int, audio_out: int) -> dict: +def _realtime_response_done(text_in: int, audio_in: int, text_out: int, audio_out: int) -> dict[str, object]: return { "type": "response.done", "response": { @@ -45,7 +46,7 @@ def _realtime_response_done(text_in: int, audio_in: int, text_out: int, audio_ou def _responses_completed( model: str, input_tokens: int, output_tokens: int, event: str = "response.completed" -) -> dict: +) -> dict[str, object]: return { "type": event, "response": { @@ -63,7 +64,12 @@ def _responses_completed( } -def _handle(messages, kwargs=MappingProxyType({}), logging_obj=None, url_route="/openai/v1/realtime"): +def _handle( + messages: Sequence[Mapping[str, object]], + kwargs: Mapping[str, object] = MappingProxyType({}), + logging_obj: MagicMock | None = None, + url_route: str = "/openai/v1/realtime", +): return OpenAIWebsocketPassthroughLoggingHandler().openai_websocket_passthrough_handler( websocket_messages=messages, logging_obj=logging_obj if logging_obj is not None else _logging_obj(), @@ -320,6 +326,37 @@ def test_a_failed_responses_turn_that_reported_no_usage_is_not_billed(): assert handled["kwargs"] == {"model": RESPONSES_MODEL, "custom_llm_provider": "openai"} +@pytest.mark.parametrize( + ("turn_model", "requested_model", "priced_as"), + [ + ("gpt-unreleased-qa-snapshot", RESPONSES_MODEL, RESPONSES_MODEL), + ("gpt-5.4-mini-qa-preview", RESPONSES_MODEL, RESPONSES_MODEL), + (RESPONSES_MODEL, REASONING_MODEL, RESPONSES_MODEL), + ], + ids=["unmapped_model", "model_the_cost_map_prices_at_zero", "priced_turn_model"], +) +def test_a_responses_turn_falls_back_to_the_requested_model_only_when_its_own_has_no_price( + turn_model: str, requested_model: str, priced_as: str +): + """ + The provider can answer under a model name the cost map does not price + yet. That turn still ran the model the key was allowed to ask for, so it is + billed at that model's price instead of $0, while a turn whose own model + has a price keeps it. + """ + handled = _handle( + [_responses_completed(turn_model, input_tokens=13, output_tokens=5)], + kwargs=MappingProxyType({"litellm_params": {"metadata": {"model_group": requested_model}}}), + url_route="/openai_passthrough/v1/responses", + ) + + prices = litellm.model_cost[priced_as] + expected_cost = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"] + assert expected_cost > 0 + assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost) + assert handled["kwargs"]["model"] == turn_model + + @pytest.mark.parametrize( ("frame_type", "kind"), [ 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 3d53a8d7b67..ac42491c389 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 @@ -47,6 +47,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( InitPassThroughEndpointHelpers, SafeRouteAdder, WebsocketFrameModelAllowed, + WebsocketFrameModelGate, WebsocketFrameModelRefused, WebsocketFrameModelVerdict, _registered_pass_through_routes, @@ -7651,13 +7652,17 @@ class RecordingFrameModelGate: return WebsocketFrameModelAllowed(access_groups=self.groups.get(model, ())) -def _received_frame(frame): +def _received_frame(frame: str | bytes) -> Mapping[str, str | bytes]: if isinstance(frame, bytes): return {"type": "websocket.receive", "bytes": frame} return {"type": "websocket.receive", "text": frame} -async def _relay_client_frames(frames, gate, user_api_key_dict=None): +async def _relay_client_frames( + frames: Sequence[str | bytes], + gate: WebsocketFrameModelGate | None, + user_api_key_dict: UserAPIKeyAuth | None = None, +) -> SimpleNamespace: upstream_ws = RecordingUpstreamWebSocket() websocket = _client_websocket( AsyncMock(side_effect=[_received_frame(frame) for frame in frames] + [{"type": "websocket.disconnect"}]) @@ -7907,6 +7912,53 @@ async def test_websocket_passthrough_keeps_the_budget_reason_when_a_stored_promp } +@pytest.mark.asyncio +async def test_websocket_passthrough_checks_the_model_a_hosted_tool_names(): + """A hosted tool such as image generation runs its own model, so that model gets the same check as the turn's""" + gate = RecordingFrameModelGate(MappingProxyType({"gpt-image-2": KEY_MODEL_REFUSAL})) + frame = { + "type": "response.create", + "model": "gpt-5.4-mini", + "tools": [{"type": "web_search"}, {"type": "image_generation", "model": "gpt-image-2"}], + } + + relayed = await _relay_client_frames([json.dumps(frame)], gate) + + assert gate.checked == ["gpt-5.4-mini", "gpt-image-2"] + assert relayed.upstream.send.await_args_list == [] + assert relayed.client.close.await_args.kwargs["code"] == 1008 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "frame", + [ + {"type": "response.create", "model": "gpt-5.4-mini", "tools": [{"type": "image_generation"}]}, + {"type": "response.create", "response": {"tools": [{"type": "image_generation", "model": None}]}}, + ], +) +async def test_websocket_passthrough_refuses_an_image_tool_without_a_model_for_a_key_limited_to_some_models(frame): + gate = RecordingFrameModelGate(MappingProxyType({STORED_PROMPT_PROBE_MODEL: KEY_MODEL_REFUSAL})) + + relayed = await _relay_client_frames([json.dumps(frame)], gate) + + assert gate.checked[-1] == STORED_PROMPT_PROBE_MODEL + assert relayed.upstream.send.await_args_list == [] + assert "image_generation tool has to name its model" in relayed.client.close.await_args.kwargs["reason"] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_checks_a_frame_only_the_standard_json_parser_reads(): + """A lone surrogate escape is valid JSON syntax the fast parser refuses, so the frame is still checked""" + gate = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL})) + frame = '{"type": "response.create", "model": "gpt-5.4", "metadata": {"note": "\\ud800"}}' + + relayed = await _relay_client_frames([frame], gate) + + assert gate.checked == ["gpt-5.4"] + assert relayed.upstream.send.await_args_list == [] + + @pytest.mark.asyncio async def test_websocket_passthrough_forwards_binary_audio_untouched(): gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": "denied"})) @@ -8015,14 +8067,16 @@ class ScriptedUpstreamWebSocket: return frame -def _turn_taking_client(frames: Sequence[str], answered_by: str) -> MagicMock: - """A client that sends its next frame only once the answer to the last one reached it, and that Starlette - refuses to send to once it is closed""" +def _turn_taking_client( + frames: Sequence[str], answered_by: str, answers_due: Sequence[int] | None = None +) -> MagicMock: + """A client that sends each frame, and then leaves, once the answers due by then reached it (by default the + answer to every earlier frame), and that Starlette refuses to send to once it is closed""" from starlette.websockets import WebSocketState delivered: list[str] = [] script = iter(frames) - receives = iter(range(len(frames) + 1)) + receives = iter(range(len(frames) + 1) if answers_due is None else answers_due) async def receive(): answers_due = next(receives) @@ -8042,7 +8096,14 @@ def _turn_taking_client(frames: Sequence[str], answered_by: str) -> MagicMock: return websocket -async def _relay_turns(client, upstream_ws, gate, endpoint="/openai_passthrough/v1/responses"): +async def _relay_turns( + client: MagicMock, + upstream_ws: ScriptedUpstreamWebSocket, + gate: WebsocketFrameModelGate, + endpoint: str = "/openai_passthrough/v1/responses", + target: str = "wss://api.openai.com/v1/responses", + opening_access_groups: Sequence[str] = (), +) -> SimpleNamespace: with ( _patched_websocket_passthrough_environment(upstream_ws) as proxy_logging, patch(SUCCESS_HANDLER, new_callable=AsyncMock) as success_handler, @@ -8050,9 +8111,11 @@ async def _relay_turns(client, upstream_ws, gate, endpoint="/openai_passthrough/ await asyncio.wait_for( websocket_passthrough_request( websocket=client, - target="wss://api.openai.com/v1/responses", + target=target, custom_headers={"Authorization": "Bearer sk-provider"}, - user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + user_api_key_dict=UserAPIKeyAuth( + api_key="hashed-key", matched_model_access_groups=list(opening_access_groups) + ), forward_headers=False, endpoint=endpoint, accept_websocket=False, @@ -8065,15 +8128,23 @@ async def _relay_turns(client, upstream_ws, gate, endpoint="/openai_passthrough/ def _responses_turn(model: str) -> list[str]: - return [ - json.dumps({"type": "response.created", "response": {"model": model}}), + return [_responses_turn_started(model), _responses_turn_ended(model)] + + +def _responses_turn_started(model: str) -> str: + return json.dumps({"type": "response.created", "response": {"model": model}}) + + +def _responses_turn_ended(model: str) -> str: + return ( + json.dumps( { "type": "response.completed", "response": {"model": model, "usage": {"input_tokens": 9, "output_tokens": 3, "total_tokens": 12}}, } - ), - ] + ) + ) @pytest.mark.asyncio @@ -8113,6 +8184,119 @@ async def test_websocket_passthrough_logs_each_turn_on_its_own_row_charged_to_th relayed.proxy_logging.post_call_failure_hook.assert_not_awaited() +@pytest.mark.asyncio +async def test_websocket_passthrough_charges_pipelined_turns_to_the_frame_that_asked_for_each(): + """ + A client can send its next response.create before the provider finished the + last one, and the provider then runs them one after the other. Each turn + still has to reach spend under the groups that let its own frame through, + and under the model that frame asked for, even when the provider answers + under a dated snapshot name a per-model budget was never set on. + """ + gate = RecordingFrameModelGate( + groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}), + ) + client = _turn_taking_client( + [ + json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}), + json.dumps({"type": "response.create", "model": "gpt-5.4"}), + ], + answered_by='"response.completed"', + answers_due=(0, 0, 2), + ) + upstream_ws = ScriptedUpstreamWebSocket( + replies=[_responses_turn("gpt-5.4-mini-2099-01-01"), _responses_turn("gpt-5.4-2099-01-01")] + ) + + relayed = await _relay_turns(client, upstream_ws, gate) + + rows = [logged.kwargs for logged in relayed.success_handler.call_args_list] + assert [[frame["response"]["model"] for frame in row["response_body"]] for row in rows] == [ + ["gpt-5.4-mini-2099-01-01", "gpt-5.4-mini-2099-01-01"], + ["gpt-5.4-2099-01-01", "gpt-5.4-2099-01-01"], + ] + assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [ + ["tier-a"], + ["tier-b"], + ] + assert [row["litellm_params"]["metadata"]["model_group"] for row in rows] == ["gpt-5.4-mini", "gpt-5.4"] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_charges_realtime_turns_to_the_session_the_key_opened(): + """ + A realtime session runs the model it was opened on, so its turns stay + charged to the groups and the model the key opened it with, whatever model + a later session.update or response.create names. + """ + gate = RecordingFrameModelGate(groups=MappingProxyType({"gpt-realtime-2.1-mini": ("tier-b",)})) + client = _turn_taking_client( + [ + json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1-mini"}}), + json.dumps({"type": "response.create", "response": {"model": "gpt-realtime-2.1-mini"}}), + ], + answered_by='"type"', + answers_due=(1, 2, 4), + ) + upstream_ws = ScriptedUpstreamWebSocket( + REALTIME_SESSION_CREATED, + replies=[ + [json.dumps({"type": "session.updated"})], + [json.dumps({"type": "response.created", "response": {"status": "in_progress"}}), REALTIME_RESPONSE_DONE], + ], + ) + + relayed = await _relay_turns( + client, + upstream_ws, + gate, + "/openai/v1/realtime", + target="wss://api.openai.com/v1/realtime?model=gpt-realtime-2.1", + opening_access_groups=("tier-a",), + ) + + assert gate.checked == ["gpt-realtime-2.1-mini", "gpt-realtime-2.1-mini"] + rows = [logged.kwargs for logged in relayed.success_handler.call_args_list] + assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [["tier-a"]] + assert [row["litellm_params"]["metadata"]["model_group"] for row in rows] == ["gpt-realtime-2.1"] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_logs_no_spend_row_for_a_session_refused_before_any_answer(): + """A session refused at its first frame ran nothing, so its only record is the refusal, not an extra $0 row""" + gate = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL})) + client = _turn_taking_client([json.dumps({"type": "response.create", "model": "gpt-5.4"})], answered_by='"error"') + + relayed = await _relay_turns(client, ScriptedUpstreamWebSocket(), gate) + + assert relayed.success_handler.call_args_list == [] + + +@pytest.mark.asyncio +async def test_websocket_passthrough_logs_the_unfinished_turn_a_client_walked_out_of(): + """ + A client can leave while a later turn is still running. That turn never + ends, but it was started and may be billed, so it still gets its own row. + """ + client = _turn_taking_client( + [ + json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}), + json.dumps({"type": "response.create", "model": "gpt-5.4"}), + ], + answered_by='"response.created"', + ) + upstream_ws = ScriptedUpstreamWebSocket( + replies=[_responses_turn("gpt-5.4-mini"), [_responses_turn_started("gpt-5.4")]] + ) + + relayed = await _relay_turns(client, upstream_ws, RecordingFrameModelGate()) + + assert [ + [json.loads(frame)["type"] if isinstance(frame, str) else frame["type"] for frame in logged.kwargs["response_body"]] + for logged in relayed.success_handler.call_args_list + ] == [["response.created", "response.completed"], ["response.created"]] + + @pytest.mark.asyncio async def test_websocket_passthrough_carries_the_session_the_provider_opened_into_every_later_turn(): """ @@ -9095,7 +9279,7 @@ def test_passthrough_client_cannot_forge_session_id_omission(client_metadata_key ) -async def _billed_metadata(gate) -> dict: +async def _billed_metadata(gate: WebsocketFrameModelGate) -> Mapping[str, object]: served = await _relay_client_frames( [json.dumps({"type": "response.create", "model": "gpt-5.4-mini"})], gate, diff --git a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py index 5b3f337fb5e..28ad1795611 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py @@ -667,7 +667,7 @@ REALTIME_FRAMES: Final = [ ] -def _normalize(url_route: str, response_body): +def _normalize(url_route: str, response_body: dict[str, object] | list[dict[str, object]]) -> dict: logging_obj = MagicMock() logging_obj.model_call_details = {} return PassThroughEndpointLogging().normalize_llm_passthrough_logging_payload( diff --git a/tests/unit/proxy/test_openai_ws_passthrough_routes.py b/tests/unit/proxy/test_openai_ws_passthrough_routes.py index 836104e3342..89478a852f3 100644 --- a/tests/unit/proxy/test_openai_ws_passthrough_routes.py +++ b/tests/unit/proxy/test_openai_ws_passthrough_routes.py @@ -18,6 +18,7 @@ import litellm from litellm.llms.openai.workload_identity import _workload_identity_auth from litellm.proxy._types import Litellm_EntityType, ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _OPENAI_WS_ALIASED_MODEL_REFUSAL, _OPENAI_WS_DISABLED_REFUSAL, _OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL, _openai_websocket_refusal, @@ -173,6 +174,18 @@ async def test_openai_websocket_forwards_query_and_keeps_provider_auth(prefix, m assert websocket.closed is None +@pytest.mark.asyncio +async def test_openai_websocket_reports_its_route_without_the_proxy_root_path(monkeypatch): + """Route allowlists and spend rows key on the route the proxy serves, which a root path prefix is not part of""" + monkeypatch.setenv("SERVER_ROOT_PATH", "/litellm") + websocket = _FakeWebSocket("/litellm/openai/v1/responses", "") + + with patch(GET_CREDENTIALS, return_value="sk-provider"): + served = await _serve(websocket, "v1/responses", UserAPIKeyAuth(), ENABLED) + + assert [call.endpoint for call in served.relay.calls] == ["/openai/v1/responses"] + + @pytest.mark.asyncio async def test_openai_websocket_accepts_first_client_subprotocol(): websocket = _FakeWebSocket( @@ -390,6 +403,48 @@ async def test_proxy_frame_model_gate_refuses_a_model_the_agent_was_not_granted( assert refusal_names in refusal.message +@pytest.mark.asyncio +@pytest.mark.parametrize( + "token", + [ + UserAPIKeyAuth(models=["gpt-5.4-mini"], aliases={"gpt-5.4": "gpt-5.4-mini"}), + UserAPIKeyAuth(models=["gpt-5.4-mini"], team_model_aliases={"gpt-5.4": "gpt-5.4-mini"}), + ], + ids=["key_alias", "team_alias"], +) +async def test_proxy_frame_model_gate_refuses_an_alias_that_would_reach_openai_unchanged(token): + """ + The passthrough sends each frame to OpenAI as written, so OpenAI runs the + alias name itself. An alias pointing at an allowed model must not let a key + limited to that model run the model the alias is named after. + """ + refusal: Final = _refusal(await _gate_verdict(_FrameAuthorizer(token), "gpt-5.4")) + + assert refusal == WebsocketFrameModelRefused( + message=_OPENAI_WS_ALIASED_MODEL_REFUSAL, error_type=ProxyErrorTypes.key_model_access_denied.value + ) + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_refuses_a_proxy_wide_alias_for_a_key_limited_to_its_target(monkeypatch): + monkeypatch.setattr(litellm, "model_alias_map", {"gpt-5.4": "gpt-5.4-mini"}) + + refusal: Final = _refusal( + await _gate_verdict(_FrameAuthorizer(UserAPIKeyAuth(models=["gpt-5.4-mini"])), "gpt-5.4") + ) + + assert refusal.message == _OPENAI_WS_ALIASED_MODEL_REFUSAL + + +@pytest.mark.asyncio +async def test_proxy_frame_model_gate_lets_a_key_that_may_use_any_model_send_an_alias(): + token: Final = UserAPIKeyAuth(models=[], aliases={"fast": "gpt-5.4-mini"}) + + verdict: Final = await _gate_verdict(_FrameAuthorizer(token), "fast") + + assert verdict == WebsocketFrameModelAllowed(access_groups=()) + + @pytest.mark.asyncio @pytest.mark.parametrize("unrestricted_models", [[], ["all-proxy-models"], ["*"]]) async def test_proxy_frame_model_gate_allows_every_model_for_unrestricted_keys(unrestricted_models):