scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-31 06:56:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 245

Merge 0d80ba329f into a392a39627

Browse source
This commit is contained in:
zoveki 2026-10-31 06:44:01 +00:00 • committed by GitHub
commit 4ee287c040
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 3047 additions and 497 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/litellm_core_utils/core_helpers.py
View file

@ -4,7 +4,7 @@ import copy
import logging
import math
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
@ -104,6 +104,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:

73
litellm/proxy/auth/auth_checks.py
View file

@ -5359,79 +5359,6 @@ async def _granted_model_lists(
)
async def _user_object_or_none(
valid_token: UserAPIKeyAuth,
prisma_client: PrismaClient,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
) -> LiteLLM_UserTable | None:
try:
return await get_user_object(
user_id=valid_token.user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
user_id_upsert=False,
proxy_logging_obj=proxy_logging_obj,
)
except UserNotFoundError:
return None
async def enforced_model_allowlists(
valid_token: UserAPIKeyAuth,
prisma_client: PrismaClient | None,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
) -> tuple[Sequence[str], ...]:
"""One model allowlist per level that ``common_checks`` enforces on a request from this identity."""
key_models: Final = _resolve_key_models_for_auth_check(valid_token=valid_token)
if prisma_client is None:
return (key_models, tuple(valid_token.team_models or ()))
team_object: Final = (
None
if valid_token.team_id is None
else await get_team_object(
team_id=valid_token.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
)
user_object: Final = (
None
if team_object is not None
else await _user_object_or_none(
valid_token=valid_token,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
)
project_object: Final = (
None
if valid_token.project_id is None
else await get_project_object(
project_id=valid_token.project_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
)
return (
key_models,
team_object.models if team_object is not None else (),
await _team_member_granted_models(
valid_token=valid_token,
team_object=team_object,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
),
user_object.models if user_object is not None else (),
project_object.models if project_object is not None else (),
)
async def collect_matched_model_access_groups(
model: str | Sequence[str] | None,
valid_token: UserAPIKeyAuth | None,

3
litellm/proxy/common_utils/realtime_utils.py
View file

@ -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

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

@ -15,11 +15,11 @@ import os
import posixpath
import re
import sys
from collections.abc import AsyncGenerator, Callable, Mapping, Sequence
from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from functools import partial
from types import MappingProxyType
from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, cast
from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias, cast
import httpx
import openai
@ -51,6 +51,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_object_list, as_str_mapping, attribute_or_none
from litellm.litellm_core_utils.fallback_generalizations import match_routing_generalization
from litellm.llms.anthropic.common_utils import AnthropicModelInfo, merge_anthropic_beta_headers
from litellm.llms.azure.passthrough.transformation import (
@ -76,7 +77,12 @@ from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.llms.vllm.common_utils import endpoint_under_api_base
from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error
from litellm.proxy._types import *
from litellm.proxy.auth.auth_checks import enforced_model_allowlists
from litellm.proxy.auth.auth_checks import (
can_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
@ -84,6 +90,7 @@ from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy modul
get_bearer_token,
is_no_auth_dev_mode,
user_api_key_auth,
user_api_key_auth_accepted_websocket_for_model,
user_api_key_auth_websocket,
user_api_key_auth_websocket_for_model,
)
@ -105,6 +112,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 (
@ -120,15 +128,25 @@ from litellm.proxy.pass_through_endpoints.injected_stream_usage import (
stream_usage_relay,
without_injected_usage_frames,
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import (
openai_websocket_frame_kind,
)
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
STORED_PROMPT_PROBE_MODEL,
HttpPassThroughEndpointHelpers,
WebsocketFrameModelAllowed,
WebsocketFrameModelGate,
WebsocketFrameModelRefused,
WebsocketFrameModelVerdict,
WebsocketUpstreamFrameClassifier,
YieldsToConfiguredPassThroughRoute,
create_pass_through_route,
create_websocket_passthrough_route,
refuse_websocket,
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,
@ -187,7 +205,7 @@ def is_passthrough_request_using_router_model(request_body: dict, llm_router: li
"""
try:
model: Final = request_body.get("model")
return is_known_model(model, llm_router)
return isinstance(model, str) and is_known_model(model, llm_router)
except Exception:
return False
@ -217,9 +235,8 @@ def is_passthrough_request_streaming(request_body: object) -> bool:
A JSON body need not be an object, so a list or scalar can reach here; it
carries no streaming flag.
"""
if not isinstance(request_body, dict):
return False
return bool(request_body.get("stream", False))
body: Final = as_str_mapping(request_body)
return body is not None and bool(body.get("stream", False))
def _optional_str(value: object) -> str | None:
@ -3376,19 +3393,6 @@ def _join_url_paths(base_url: httpx.URL, path: str, custom_llm_provider: litellm
return joined_path_str
_OPENAI_WS_ALL_MODEL_ACCESS: Final = frozenset(
{
SpecialModelNames.all_proxy_models.value,
SpecialModelNames.all_team_models.value,
"*",
}
)
def _has_model_restrictions(model_allowlists: tuple[Sequence[str], ...]) -> bool:
return any(str(model) not in _OPENAI_WS_ALL_MODEL_ACCESS for allowlist in model_allowlists for model in allowlist)
@dataclass(frozen=True, slots=True)
class _OpenAIWebsocketRefusal:
close_reason: str
@ -3405,6 +3409,16 @@ class _OpenAIWebsocketErrorFrame(TypedDict):
error: ReadOnly[_OpenAIWebsocketErrorDetail]
_OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL: Final = (
"This gateway could not authorize the model named in this frame, so the session was closed. Try again, "
"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=(
@ -3413,14 +3427,6 @@ _OPENAI_WS_DISABLED_REFUSAL: Final = _OpenAIWebsocketRefusal(
),
)
_OPENAI_WS_MODEL_RESTRICTED_REFUSAL: Final = _OpenAIWebsocketRefusal(
close_reason="Keys with model restrictions cannot use OpenAI websocket passthrough",
message=(
"Keys with model restrictions cannot use OpenAI websocket passthrough, because this route "
"relays frames to the provider without reading which model they ask for."
),
)
def _is_openai_websocket_passthrough_enabled(general_settings: Mapping[str, object]) -> bool:
setting: Final = general_settings.get("enable_openai_websocket_passthrough")
@ -3429,19 +3435,9 @@ def _is_openai_websocket_passthrough_enabled(general_settings: Mapping[str, obje
return setting is True
class _OpenAIWebsocketModelAllowlists(Protocol):
async def __call__(self, valid_token: UserAPIKeyAuth, /) -> tuple[Sequence[str], ...]: ...
async def _openai_websocket_refusal(
user_api_key_dict: UserAPIKeyAuth,
general_settings: Mapping[str, object],
model_allowlists: _OpenAIWebsocketModelAllowlists,
) -> _OpenAIWebsocketRefusal | None:
def _openai_websocket_refusal(general_settings: Mapping[str, object]) -> _OpenAIWebsocketRefusal | None:
if not _is_openai_websocket_passthrough_enabled(general_settings):
return _OPENAI_WS_DISABLED_REFUSAL
if _has_model_restrictions(await model_allowlists(user_api_key_dict)):
return _OPENAI_WS_MODEL_RESTRICTED_REFUSAL
return None
@ -3456,6 +3452,8 @@ class _WebsocketRelay(Protocol):
forward_headers: bool,
endpoint: str,
accept_websocket: bool,
client_frame_model_gate: WebsocketFrameModelGate | None = None,
upstream_frame_kind: WebsocketUpstreamFrameClassifier | None = None,
) -> None: ...
@ -3463,18 +3461,95 @@ def _websocket_relay() -> _WebsocketRelay:
return websocket_passthrough_request
def _proxy_model_allowlists() -> _OpenAIWebsocketModelAllowlists:
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
def _proxy_error_type(denial: ProxyException) -> str:
return denial.type.value if isinstance(denial.type, ProxyErrorTypes) else denial.type
async def resolve(valid_token: UserAPIKeyAuth, /) -> tuple[Sequence[str], ...]:
return await enforced_model_allowlists(
valid_token=valid_token,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
FrameModelAuthorizer: TypeAlias = Callable[[WebSocket, str], Awaitable[UserAPIKeyAuth]]
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.registry().resolve(model, kind="alias") 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 import proxy_server
from litellm.proxy.spend_tracking.budget_reservation import release_or_invalidate_budget_reservation
async def authorize(websocket: WebSocket, model: str) -> WebsocketFrameModelVerdict:
connect_reservation: Final = attribute_or_none(websocket.state, "budget_reservation")
await release_or_invalidate_budget_reservation(
budget_reservation=connect_reservation if isinstance(connect_reservation, dict) else None
)
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=_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=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)
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))
return resolve
async def gate(websocket: WebSocket, model: str, /) -> WebsocketFrameModelVerdict:
try:
return await authorize(websocket, model)
except ProxyException as denial:
return WebsocketFrameModelRefused(message=denial.message, error_type=_proxy_error_type(denial))
except HTTPException as denial:
return WebsocketFrameModelRefused(message=str(denial.detail), error_type="invalid_request_error")
except litellm.BudgetExceededError as exhausted:
return WebsocketFrameModelRefused(message=exhausted.message, error_type="budget_exceeded")
except Exception: # noqa: BLE001 # fail closed: a gate that cannot answer refuses, and says so in band
verbose_proxy_logger.exception("OpenAI websocket passthrough: the frame model gate failed on %s", model)
return WebsocketFrameModelRefused(message=_OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL, error_type="server_error")
return gate
def _negotiated_websocket_subprotocol(websocket: WebSocket) -> str | None:
@ -3486,6 +3561,14 @@ def _negotiated_websocket_subprotocol(websocket: WebSocket) -> str | None:
return requested_subprotocols[0] if requested_subprotocols else None
async def _opening_model_refusal(
websocket: WebSocket, query_string: str, frame_model_gate: WebsocketFrameModelGate
) -> WebsocketFrameModelRefused | None:
opening_models: Final = dict.fromkeys(httpx.QueryParams(query_string).get_list("model"))
verdicts: Final = tuple([await frame_model_gate(websocket, model) for model in opening_models])
return next((verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelRefused)), None)
@router.websocket("/openai_passthrough/{endpoint:path}")
@router.websocket("/openai/{endpoint:path}")
async def openai_websocket_proxy_route(
@ -3494,12 +3577,12 @@ async def openai_websocket_proxy_route(
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth_websocket)],
general_settings: Annotated[Mapping[str, object], Depends(_proxy_general_settings)],
relay: Annotated[_WebsocketRelay, Depends(_websocket_relay)],
model_allowlists: Annotated[_OpenAIWebsocketModelAllowlists, Depends(_proxy_model_allowlists)],
frame_model_gate: Annotated[WebsocketFrameModelGate, Depends(_proxy_frame_model_gate)],
) -> None:
"""WebSocket passthrough for OpenAI prefixes (realtime / responses.connect)."""
negotiated_subprotocol: Final = _negotiated_websocket_subprotocol(websocket)
refusal: Final = await _openai_websocket_refusal(user_api_key_dict, general_settings, model_allowlists)
refusal: Final = _openai_websocket_refusal(general_settings)
if refusal is not None:
await websocket.accept(subprotocol=negotiated_subprotocol)
error_frame: Final[_OpenAIWebsocketErrorFrame] = {
@ -3545,14 +3628,21 @@ async def openai_websocket_proxy_route(
await websocket.accept(subprotocol=negotiated_subprotocol)
opening_refusal: Final = await _opening_model_refusal(websocket, query_string, frame_model_gate)
if opening_refusal is not None:
await refuse_websocket(websocket, opening_refusal)
return
await relay(
websocket=websocket,
target=wss_target,
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,
)
@ -3820,7 +3910,8 @@ async def _tinyfish_json_object_field_names(request: Request) -> frozenset[str]
parsed: Final = _DECODED_JSON.validate_python(json.loads(raw_body))
except (json.JSONDecodeError, UnicodeDecodeError):
return None
return frozenset(parsed) if isinstance(parsed, dict) else None
fields: Final = as_str_mapping(parsed)
return None if fields is None else frozenset(fields)
def _tinyfish_route_timeout() -> float | None:

257
litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py Normal file
View file

@ -0,0 +1,257 @@
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from types import MappingProxyType
from typing import Final, Literal, TypeAlias, cast # noqa: TID251 # cost helpers type raw frames as a TypedDict list
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
from litellm.types.llms.openai import OpenAIRealtimeStreamList, ResponsesAPIResponse
from litellm.types.passthrough_endpoints.pass_through_endpoints import WebsocketUpstreamFrameKind
from litellm.types.utils import LlmProviders, ModelResponse, Usage
WebsocketMessages: TypeAlias = Sequence[Mapping[str, object]]
OPENAI_WEBSOCKET_PROVIDER: Final = LlmProviders.OPENAI.value
_SESSION_CREATED: Final = "session.created"
_SESSION_UPDATED: Final = "session.updated"
_RESPONSE_CREATED: Final = "response.created"
_ERROR: Final = "error"
_REALTIME_RESPONSE_DONE: Final = "response.done"
_TRANSCRIPTION_COMPLETED: Final = "conversation.item.input_audio_transcription.completed"
_RESPONSES_TURN_ENDS: Final = frozenset({"response.completed", "response.incomplete", "response.failed"})
_TURN_ENDS: Final = _RESPONSES_TURN_ENDS | {_REALTIME_RESPONSE_DONE}
_MODEL_NAMING_EVENTS: Final = _TURN_ENDS | {_SESSION_CREATED, _SESSION_UPDATED, _RESPONSE_CREATED}
_REALTIME_BILLED_EVENTS: Final = frozenset({_REALTIME_RESPONSE_DONE, _TRANSCRIPTION_COMPLETED})
_RESPONSE_COST_HEADER: Final = "llm_provider-x-litellm-response-cost"
class _NamedModel(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
model: str | None = None
class _AudioInput(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
transcription: _NamedModel | None = None
class _SessionAudio(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
input: _AudioInput | None = None
class _NamedSession(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
model: str | None = None
audio: _SessionAudio | None = None
class _ModelNamingEvent(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
session: _NamedSession | None = None
response: _NamedModel | None = None
class _ResponsesTurnEndEvent(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
type: Literal["response.completed", "response.incomplete", "response.failed"]
response: ResponsesAPIResponse
_MODEL_NAMING_EVENT: Final = TypeAdapter(_ModelNamingEvent)
_RESPONSES_TURN_END_EVENT: Final = TypeAdapter(_ResponsesTurnEndEvent)
@dataclass(frozen=True, slots=True)
class _Billing:
usage: Usage
cost: float
@dataclass(frozen=True, slots=True)
class _CompletedTurn:
model: str
response: ResponsesAPIResponse
usage: Usage
def _event_model(event: Mapping[str, object]) -> str | None:
if event.get("type") not in _MODEL_NAMING_EVENTS:
return None
try:
naming: Final = _MODEL_NAMING_EVENT.validate_python(event)
except ValidationError:
return None
session: Final = naming.session or _NamedSession()
response: Final = naming.response or _NamedModel()
audio_input: Final = (session.audio or _SessionAudio()).input or _AudioInput()
transcription: Final = audio_input.transcription or _NamedModel()
return session.model or response.model or transcription.model or None
def observed_websocket_model(messages: WebsocketMessages) -> str | None:
return next((model for model in map(_event_model, messages) if model is not None), None)
def openai_websocket_frame_kind(frame: Mapping[str, object]) -> WebsocketUpstreamFrameKind:
frame_type: Final = frame.get("type")
if frame_type in _TURN_ENDS:
return "ends_turn"
if frame_type == _SESSION_CREATED:
return "opens_session"
if frame_type == _RESPONSE_CREATED:
return "starts_turn"
if frame_type == _ERROR:
return "refuses_turn"
return "relayed"
def _realtime_billing(messages: WebsocketMessages, model: str, logging_obj: LiteLLMLoggingObj) -> _Billing | None:
if not any(event.get("type") in _REALTIME_BILLED_EVENTS for event in messages):
return None
results: Final = cast( # cast-ok: the realtime cost helpers read raw provider frames through their TypedDict view
OpenAIRealtimeStreamList, tuple(event for event in messages if "type" in event)
)
usage: Final = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(results)
cost: Final = handle_realtime_stream_cost_calculation(
results=results,
combined_usage_object=usage,
custom_llm_provider=OPENAI_WEBSOCKET_PROVIDER,
litellm_model_name=model,
litellm_logging_obj=logging_obj,
)
return _Billing(usage=usage, cost=cost)
def _responses_turn(event: Mapping[str, object], session_model: str) -> _CompletedTurn | None:
if event.get("type") not in _RESPONSES_TURN_ENDS:
return None
try:
ended: Final = _RESPONSES_TURN_END_EVENT.validate_python(event)
except ValidationError:
verbose_proxy_logger.warning(
"OpenAI websocket passthrough: a %s frame carried no readable usage, so it is not billed", event.get("type")
)
return None
if ended.response.usage is None:
return None
return _CompletedTurn(
model=ended.response.model or session_model,
response=ended.response,
usage=ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # the one shared Responses usage converter
ended.response.usage
),
)
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.response, (turn.model, requested_model)) for turn in turns),
)
def _model_cost(response: ResponsesAPIResponse, model: str) -> float | None:
try:
return litellm.completion_cost(
completion_response=response,
model=model,
custom_llm_provider=OPENAI_WEBSOCKET_PROVIDER,
call_type="responses",
)
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(response: ResponsesAPIResponse, price_as: Sequence[str | None]) -> float:
costs: Final = (_model_cost(response, 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", 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, call_id: str) -> ModelResponse:
result: Final = ModelResponse(
id=call_id,
object="chat.completion",
created=int(start_time.timestamp()),
model=model,
usage=billing.usage,
)
hidden: Final = {"additional_headers": MappingProxyType({_RESPONSE_COST_HEADER: billing.cost})}
result._hidden_params = hidden # pyright: ignore[reportPrivateUsage] # the one settled-cost channel
return result
class OpenAIWebsocketPassthroughLoggingHandler:
def openai_websocket_passthrough_handler(
self,
websocket_messages: WebsocketMessages,
logging_obj: LiteLLMLoggingObj,
url_route: str,
start_time: datetime,
kwargs: Mapping[str, object],
) -> PassThroughEndpointLoggingTypedDict:
model: Final = observed_websocket_model(websocket_messages)
if model is None:
verbose_proxy_logger.debug("OpenAI websocket passthrough (%s): no model named, skipping cost", url_route)
unpriced: Final[PassThroughEndpointLoggingTypedDict] = {"result": None, "kwargs": {**kwargs}}
return unpriced
logging_obj.model = model # rebind-ok: the spend logger reads model and cost off the shared logging object
logging_obj.model_call_details.update(
MappingProxyType({"model": model, "custom_llm_provider": OPENAI_WEBSOCKET_PROVIDER})
)
billing: Final = _realtime_billing(websocket_messages, model, logging_obj) or _responses_billing(
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)
unbilled: Final[PassThroughEndpointLoggingTypedDict] = {
"result": None,
"kwargs": {**kwargs, "model": model, "custom_llm_provider": OPENAI_WEBSOCKET_PROVIDER},
}
return unbilled
logging_obj.model_call_details.update(MappingProxyType({"response_cost": billing.cost}))
priced: Final[PassThroughEndpointLoggingTypedDict] = {
"result": _priced_result(model, billing, start_time, logging_obj.litellm_call_id),
"kwargs": {
**kwargs,
"model": model,
"custom_llm_provider": OPENAI_WEBSOCKET_PROVIDER,
"response_cost": billing.cost,
},
}
return priced

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

@ -6,15 +6,17 @@ import posixpath
import re
import traceback
from base64 import b64encode
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, Sequence
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, MutableMapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from functools import reduce
from itertools import chain, count, groupby
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, TypedDict, cast
from urllib.parse import urlencode, urlparse
from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, TypedDict, cast
from urllib.parse import parse_qs, urlencode, urlparse
import httpx
import orjson
from fastapi import (
APIRouter,
Depends,
@ -24,6 +26,7 @@ from fastapi import (
Response,
UploadFile,
WebSocket,
WebSocketDisconnect,
status,
)
from fastapi.responses import StreamingResponse
@ -34,7 +37,7 @@ from starlette.routing import BaseRoute, Match, Route
from starlette.types import Scope
from starlette.websockets import WebSocketState
from typing_extensions import ReadOnly
from websockets.asyncio.client import connect
from websockets.asyncio.client import ClientConnection, connect
from websockets.exceptions import (
ConnectionClosedError,
ConnectionClosedOK,
@ -60,6 +63,8 @@ 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,
get_or_create_metadata_bucket,
@ -69,6 +74,7 @@ from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.litellm_logging import get_masked_values
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.litellm_core_utils.realtime_errors import realtime_error_event
from litellm.litellm_core_utils.redact_messages import should_redact_message_logging
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.llms.anthropic.common_utils import is_anthropic_messages_url
@ -150,6 +156,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
PassthroughStandardLoggingPayload,
WebSocketSessionMeter,
WebSocketSessionUsage,
WebsocketUpstreamFrameKind,
)
from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, CallTypes, Usage
@ -2534,11 +2541,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
@ -2597,6 +2609,307 @@ 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 WebsocketFrameModelAllowed:
access_groups: tuple[str, ...]
@dataclass(frozen=True, slots=True)
class WebsocketFrameModelRefused:
message: str
error_type: str
WebsocketFrameModelVerdict: TypeAlias = WebsocketFrameModelAllowed | WebsocketFrameModelRefused
class WebsocketFrameModelGate(Protocol):
async def __call__(self, websocket: WebSocket, model: str, /) -> WebsocketFrameModelVerdict: ...
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:
starts_turn: bool
turn: _TurnAuthorization | None
lane: str | None
@dataclass(frozen=True, slots=True)
class _RefuseClientFrame:
refusal: WebsocketFrameModelRefused
@dataclass(frozen=True, slots=True)
class _ClientFrameRequests:
starts_turn: bool
lane: str | None
turn_model: str | None
models: tuple[str, ...]
unverifiable_model_refusal: str | None
_NO_CLIENT_FRAME_REQUESTS: Final = _ClientFrameRequests(
starts_turn=False, lane=None, turn_model=None, models=(), unverifiable_model_refusal=None
)
_FORWARD_WITHOUT_TURN: Final = _ForwardClientFrame(starts_turn=False, turn=None, lane=None)
_TURN_STARTING_FRAME: Final = "response.create"
_FRAME_TEXT_WORTH_PARSING: Final = (f'"{_TURN_STARTING_FRAME}"', '"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|$)")
def _may_name_a_model(text_data: str) -> bool:
return any(marker in text_data for marker in _FRAME_TEXT_WORTH_PARSING)
def _frame_lane(frame: Mapping[str, object]) -> str | None:
stream_id: Final = frame.get("stream_id")
return stream_id if isinstance(stream_id, str) else None
def _frames_of_lane(lane: str | None) -> Callable[[Mapping[str, object]], bool]:
def in_lane(frame: Mapping[str, object]) -> bool:
return _frame_lane(frame) == lane
return in_lane
def _every_frame(_frame: Mapping[str, object]) -> bool:
return True
def _nested_frame(frame: Mapping[str, object] | None, key: str) -> Mapping[str, object] | None:
return None if frame is None else as_str_mapping(frame.get(key))
def _placed_value(frame: Mapping[str, object], placement: Sequence[str]) -> object:
*parents, leaf = placement
container: Final = reduce(_nested_frame, parents, frame)
return None if container is None else container.get(leaf)
def _placed_models(frame: Mapping[str, object], placements: Iterable[Sequence[str]]) -> tuple[str, ...]:
placed: Final = (_placed_value(frame, placement) for placement in placements)
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
frame: Final = _json_object_frame(frame_data)
if frame is None:
return _NO_CLIENT_FRAME_REQUESTS
turn_models: Final = _placed_models(frame, _TURN_MODEL_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,
lane=_frame_lane(frame),
turn_model=next(iter(turn_models), None),
models=tuple(
dict.fromkeys((*turn_models, *_placed_models(frame, _TRANSCRIPTION_MODEL_PLACEMENTS), *tool_models))
),
unverifiable_model_refusal=_unverifiable_model_refusal(frame, tools),
)
def _unverifiable_model_probe_refusal(probe: WebsocketFrameModelRefused, message: str) -> WebsocketFrameModelRefused:
if probe.error_type.endswith("_model_access_denied"):
return WebsocketFrameModelRefused(message=message, error_type=probe.error_type)
return probe
async def _client_frame_verdict(
websocket: WebSocket, frame_data: str | bytes, gate: WebsocketFrameModelGate | None
) -> _ForwardClientFrame | _RefuseClientFrame:
if gate is 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.unverifiable_model_refusal is not None and isinstance(
probe := await gate(websocket, STORED_PROMPT_PROBE_MODEL), WebsocketFrameModelRefused
):
return _RefuseClientFrame(refusal=_unverifiable_model_probe_refusal(probe, requests.unverifiable_model_refusal))
allowed: Final = tuple(verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelAllowed))
return _ForwardClientFrame(
starts_turn=requests.starts_turn,
lane=requests.lane,
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:
first_sentence: Final = _FIRST_SENTENCE.match(message)
return _truncated_close_reason(message if first_sentence is None else first_sentence.group(0))
async def refuse_websocket(websocket: WebSocket, refusal: WebsocketFrameModelRefused) -> None:
await websocket.send_text(realtime_error_event(refusal.message, error_type=refusal.error_type))
await websocket.close(code=CloseCode.POLICY_VIOLATION, reason=_refusal_close_reason(refusal.message))
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 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 in each lane, when the next row starts, and how many rows
were logged.
The provider runs the response.create frames of one lane (one stream_id, or none) one at a time in the order they
were sent, and runs separate lanes side by side, so each row is charged to the groups that allowed the oldest turn
still unlogged in its lane. A turn the provider refuses with an error event never runs, so it leaves its lane
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: Mapping[str | None, tuple[_TurnAuthorization | None, ...]] = MappingProxyType({})
self.session_opened = False
self.started = started
self.logged_turns = 0
def current(self, lane: str | None) -> _TurnAuthorization:
oldest: Final = next(iter(self.queued_turns.get(lane, ())), None)
return self.opening if oldest is None else oldest
def queue_turn(self, lane: str | None, turn: _TurnAuthorization | None) -> None:
if not self.session_opened:
self.queued_turns = MappingProxyType({**self.queued_turns, lane: (*self.queued_turns.get(lane, ()), turn)})
def open_session(self) -> None:
self.session_opened = True
self.queued_turns = MappingProxyType({})
def close_turn(self, lane: str | None, end_time: datetime) -> None:
self.logged_turns += 1
self.started = end_time
self._forget_turn(lane, 0)
def drop_refused_turn(self, lane: str | None, lane_is_running_a_turn: bool) -> None:
self._forget_turn(lane, 1 if lane_is_running_a_turn else 0)
def _forget_turn(self, lane: str | None, position: int) -> None:
turns: Final = self.queued_turns.get(lane, ())
self.queued_turns = MappingProxyType({**self.queued_turns, lane: (*turns[:position], *turns[position + 1 :])})
def _session_end_needs_a_row(
upstream_frame_kind: WebsocketUpstreamFrameClassifier | None,
unlogged_frames: Sequence[Mapping[str, object]],
ledger: _WebsocketTurnLedger,
) -> bool:
"""A session logs one row per turn as each ends, so at its end it only owes a row for a turn still running, or
for whatever the provider sent when no turn ever ended"""
if upstream_frame_kind is None:
return True
kinds: Final = frozenset(map(upstream_frame_kind, unlogged_frames))
return "starts_turn" in kinds or (ledger.logged_turns == 0 and bool(kinds - {"opens_session"}))
def _turn_in_flight(
upstream_frame_kind: WebsocketUpstreamFrameClassifier | None, unlogged_frames: Sequence[Mapping[str, object]]
) -> bool:
return upstream_frame_kind is not None and any(
upstream_frame_kind(frame) == "starts_turn" for frame in unlogged_frames
)
def _lanes_in_flight(
upstream_frame_kind: WebsocketUpstreamFrameClassifier, unlogged_frames: Sequence[Mapping[str, object]]
) -> tuple[str | None, ...]:
started: Final = (frame for frame in unlogged_frames if upstream_frame_kind(frame) == "starts_turn")
return tuple(dict.fromkeys(map(_frame_lane, started)))
async def _relay_client_frame(
websocket: WebSocket,
upstream_ws: ClientConnection,
checked_frame: str | bytes,
forwarded_frame: str | bytes,
gate: WebsocketFrameModelGate | None,
ledger: _WebsocketTurnLedger,
) -> _ForwardClientFrame | _RefuseClientFrame:
"""Forward one client frame upstream, or refuse it and close the client."""
verdict: Final = await _client_frame_verdict(websocket, checked_frame, gate)
match verdict:
case _RefuseClientFrame(refusal=refusal):
await refuse_websocket(websocket, refusal)
case _ForwardClientFrame(starts_turn=True, turn=turn, lane=lane):
ledger.queue_turn(lane, turn)
await upstream_ws.send(forwarded_frame)
case _ForwardClientFrame():
await upstream_ws.send(forwarded_frame)
return verdict
_EXPLICIT_SUCCESS_KWARGS: Final = frozenset({"logging_obj", "passthrough_logging_payload", "custom_llm_provider"})
WALKOUT_TURN_DRAIN_SECONDS: Final = 600.0
def _metered_client_text(meter: WebSocketSessionMeter | None, client_text: str) -> int:
return 0 if meter is None else meter.measure(client_text)
def _log_session_usage(logging_obj: LiteLLMLoggingObj, usage: WebSocketSessionUsage) -> None:
logging_obj.model = usage.model # rebind-ok: the spend logger reads model and cost off the shared logging object
logging_obj.model_call_details["model"] = usage.model # rebind-ok: same shared logging object
@ -2614,6 +2927,9 @@ 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_model_gate: WebsocketFrameModelGate | None = None,
upstream_frame_kind: WebsocketUpstreamFrameClassifier | None = None,
walkout_turn_drain_seconds: float = WALKOUT_TURN_DRAIN_SECONDS,
builtin_provider_route: bool = False,
model: str | None = None,
meter: WebSocketSessionMeter | None = None,
@ -2630,6 +2946,11 @@ 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_model_gate: Optional per-frame check of every model a client frame names
upstream_frame_kind: Optional sort of upstream frames into the one that ends a billable turn and the ones
that open the session, so each turn logs its own spend and keeps the session's provider-named model
walkout_turn_drain_seconds: How long to keep reading a turn the provider is still running after the client
left or was refused, so the provider's usage for it still reaches spend
builtin_provider_route: Whether a built-in provider route serves the connection, so its URL names the provider
model: The model the caller named before connecting, logged on a failure before the setup frame arrives
meter: Prices the session from the client text frames delivered upstream, for an upstream that reports no usage
@ -2669,26 +2990,31 @@ async def websocket_passthrough_request(
proxy_config=proxy_config,
route_description="websocket_passthrough",
)
logging_obj: Final = Logging(
model="unknown",
messages=[{"role": "user", "content": "WebSocket connection"}],
stream=True, # WebSockets are inherently streaming
call_type="pass_through_endpoint",
start_time=start_time,
litellm_call_id=litellm_call_id,
function_id="websocket_passthrough",
dynamic_success_callbacks=team_callbacks.success_callbacks,
dynamic_failure_callbacks=team_callbacks.failure_callbacks,
kwargs=team_callbacks.logging_kwargs,
)
# Create passthrough logging payload
passthrough_logging_payload: Final = PassthroughStandardLoggingPayload(
url=target,
request_body={}, # WebSocket doesn't have a traditional request body
request_method="WEBSOCKET",
cost_per_request=cost_per_request,
)
def new_call_logging(call_start: datetime, call_id: str) -> LiteLLMLoggingObj:
return Logging(
model="unknown",
messages=[{"role": "user", "content": "WebSocket connection"}],
stream=True,
call_type="pass_through_endpoint",
start_time=call_start,
litellm_call_id=call_id,
function_id="websocket_passthrough",
dynamic_success_callbacks=team_callbacks.success_callbacks,
dynamic_failure_callbacks=team_callbacks.failure_callbacks,
kwargs=team_callbacks.logging_kwargs,
)
def new_call_payload() -> PassthroughStandardLoggingPayload:
return PassthroughStandardLoggingPayload(
url=target,
request_body={},
request_method="WEBSOCKET",
cost_per_request=cost_per_request,
)
logging_obj: Final = new_call_logging(start_time, litellm_call_id)
passthrough_logging_payload: Final = new_call_payload()
# Create a dummy request object for WebSocket connections to maintain compatibility
# with the existing _init_kwargs_for_pass_through_endpoint function
@ -2710,38 +3036,132 @@ async def websocket_passthrough_request(
)
# Initialize kwargs for logging using the same pattern as HTTP passthrough
kwargs: Final = HttpPassThroughEndpointHelpers.init_kwargs_for_pass_through_endpoint(
user_api_key_dict=user_api_key_dict,
_parsed_body={}, # WebSocket doesn't have a traditional request body
passthrough_logging_payload=passthrough_logging_payload,
litellm_call_id=litellm_call_id,
request=dummy_request,
logging_obj=logging_obj,
def init_call_kwargs(
token: UserAPIKeyAuth,
call_payload: PassthroughStandardLoggingPayload,
call_id: str,
call_logging: LiteLLMLoggingObj,
) -> dict[str, object]: # mutable-ok: the Vertex relay stamps the model from setup frames into these kwargs
return HttpPassThroughEndpointHelpers.init_kwargs_for_pass_through_endpoint(
user_api_key_dict=token,
_parsed_body={}, # WebSocket doesn't have a traditional request body
passthrough_logging_payload=call_payload,
litellm_call_id=call_id,
request=dummy_request,
logging_obj=call_logging,
)
kwargs: Final = init_call_kwargs(user_api_key_dict, passthrough_logging_payload, litellm_call_id, logging_obj)
def start_call_logging(call_logging: LiteLLMLoggingObj, call_kwargs: Mapping[str, object]) -> None:
call_logging.update_environment_variables(
model="unknown",
user="unknown",
optional_params={},
litellm_params=dict(as_str_mapping(call_kwargs.get("litellm_params")) or {}),
call_type="pass_through_endpoint",
)
call_logging.pre_call(
input=[{"role": "user", "content": "WebSocket connection"}],
api_key="",
additional_args={
"complete_input_dict": {},
"api_base": target,
"headers": upstream_headers,
},
)
start_call_logging(logging_obj, kwargs)
class MockWebSocketRequest:
def __init__(self, target_url: str):
self.method = "WEBSOCKET"
self.url = target_url
class MockWebSocketResponse:
def __init__(self, target_url: str):
self.status_code = 200
self.text = "WebSocket connection successful"
self.headers: dict[str, str] = {}
self.request = MockWebSocketRequest(target_url)
ledger: Final = _WebsocketTurnLedger(
opening=_TurnAuthorization(
requested_model=_opening_model(target),
access_groups=tuple(user_api_key_dict.matched_model_access_groups or ()),
),
started=start_time,
)
# Update logging environment variables
logging_obj.update_environment_variables(
model="unknown",
user="unknown",
optional_params={},
litellm_params=dict(kwargs.get("litellm_params", {})),
call_type="pass_through_endpoint",
)
logging_obj.model_call_details["litellm_call_id"] = litellm_call_id
def log_frames(
call_logging: LiteLLMLoggingObj,
call_payload: PassthroughStandardLoggingPayload,
call_kwargs: Mapping[str, object],
end_time: datetime,
lane: str | None,
in_row: Callable[[Mapping[str, object]], bool],
) -> None:
all_frames: Final = list(websocket_messages)
frames: Final = [frame for frame in all_frames if in_row(frame)]
websocket_messages[:] = (
()
if upstream_frame_kind is None
else (frame for frame in all_frames if not in_row(frame) or upstream_frame_kind(frame) == "opens_session")
)
logged_payload: Final[PassthroughStandardLoggingPayload] = {
**call_payload,
"response_body": frames,
"end_time": end_time,
}
if client_frame_model_gate is not None:
_stamp_billed_turn(call_kwargs, ledger.current(lane))
success_kwargs: Final = {
key: value for key, value in call_kwargs.items() if key not in _EXPLICIT_SUCCESS_KWARGS
}
custom_llm_provider: Final = call_kwargs.get("custom_llm_provider")
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
async_coroutine=pass_through_endpoint_logging.pass_through_async_success_handler(
httpx_response=MockWebSocketResponse(target),
response_body=frames,
url_route=endpoint or "",
result="websocket_connection_successful",
start_time=ledger.started,
end_time=end_time,
logging_obj=call_logging,
cache_hit=False,
request_body={},
passthrough_logging_payload=logged_payload,
custom_llm_provider=custom_llm_provider if isinstance(custom_llm_provider, str) else None,
finalize_before_dispatch=False,
**success_kwargs,
)
)
bind_budget_reservation_to_callbacks(call_logging.litellm_params)
# Pre-call logging
logging_obj.pre_call(
input=[{"role": "user", "content": "WebSocket connection"}],
api_key="",
additional_args={
"complete_input_dict": {},
"api_base": target,
"headers": upstream_headers,
},
)
def log_turn(lane: str | None, in_row: Callable[[Mapping[str, object]], bool]) -> None:
end_time: Final = datetime.now()
if upstream_frame_kind is None:
log_frames(logging_obj, passthrough_logging_payload, kwargs, end_time, lane, in_row)
else:
turn_call_id: Final = str(uuid.uuid4())
turn_logging: Final = new_call_logging(ledger.started, turn_call_id)
turn_payload: Final = new_call_payload()
turn_kwargs: Final = init_call_kwargs(
user_api_key_dict
if ledger.logged_turns == 0
else user_api_key_dict.model_copy(update={"budget_reservation": None}),
turn_payload,
turn_call_id,
turn_logging,
)
start_call_logging(turn_logging, turn_kwargs)
log_frames(turn_logging, turn_payload, turn_kwargs, end_time, lane, in_row)
ledger.close_turn(lane, end_time)
### CALL HOOKS ### - modify incoming data / reject request before calling the model
websocket_hook_data: Final = guardrail_request_data_with_streaming(MappingProxyType({}), is_streaming=True)
websocket_hook_data: Final = guardrail_request_data_with_streaming(
MappingProxyType({"litellm_call_id": litellm_call_id}), is_streaming=True
)
await proxy_logging_obj.pre_call_hook(
user_api_key_dict=user_api_key_dict,
data=websocket_hook_data,
@ -2768,7 +3188,6 @@ async def websocket_passthrough_request(
message = await websocket.receive()
message_type = message.get("type")
if message_type == "websocket.disconnect":
await upstream_ws.close()
break
text_data: str | None = message.get("text")
@ -2831,11 +3250,29 @@ 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))
if meter is not None:
metered_total += meter.measure(text_data)
if isinstance(
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
metered_total += _metered_client_text(meter, text_data)
elif bytes_data is not None:
await upstream_ws.send(bytes_data)
if isinstance(
await _relay_client_frame(
websocket, upstream_ws, bytes_data, bytes_data, client_frame_model_gate, ledger
),
_RefuseClientFrame,
):
break
if not _turn_in_flight(upstream_frame_kind, websocket_messages):
await upstream_ws.close()
except asyncio.CancelledError:
raise
except Exception:
@ -2862,11 +3299,7 @@ async def websocket_passthrough_request(
is_vertex_live: Final = bool(endpoint and "/vertex_ai/live" in endpoint)
json_frame_ordinal: Final = count()
async def relay_upstream_frame(upstream_message: str | bytes) -> None:
if isinstance(upstream_message, bytes):
await websocket.send_bytes(upstream_message)
else:
await websocket.send_text(upstream_message)
def record_upstream_frame(upstream_message: str | bytes) -> None:
message_data: Final = _json_object_frame(upstream_message)
if message_data is None:
return
@ -2874,6 +3307,33 @@ async def websocket_passthrough_request(
_extract_vertex_live_model_from_setup_response(message_data)
return
websocket_messages.append(message_data)
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(_frame_lane(message_data), _frames_of_lane(_frame_lane(message_data)))
case "refuses_turn" if upstream_frame_kind is not None:
lane: Final = _frame_lane(message_data)
ledger.drop_refused_turn(
lane, lane in _lanes_in_flight(upstream_frame_kind, websocket_messages)
)
case _:
pass
async def relay_upstream_frame(upstream_message: str | bytes) -> None:
record_upstream_frame(upstream_message)
if not _client_socket_is_open(websocket):
if not _turn_in_flight(upstream_frame_kind, websocket_messages):
await upstream_ws.close()
return
try:
if isinstance(upstream_message, bytes):
await websocket.send_bytes(upstream_message)
else:
await websocket.send_text(upstream_message)
except WebSocketDisconnect:
if not _turn_in_flight(upstream_frame_kind, websocket_messages):
raise
async def forward_upstream_to_client() -> Close | None:
try:
@ -2882,6 +3342,11 @@ async def websocket_passthrough_request(
except (ConnectionClosedOK, ConnectionClosedError) as e:
verbose_proxy_logger.debug("Upstream WebSocket connection closed: %s", e)
return e.rcvd
except WebSocketDisconnect:
verbose_proxy_logger.debug(
"WebSocket passthrough (%s): client left while upstream was still sending", endpoint
)
return None
except asyncio.CancelledError:
verbose_proxy_logger.debug("asyncio.CancelledError in forward_upstream_to_client")
raise
@ -2893,12 +3358,15 @@ async def websocket_passthrough_request(
raise
# Create tasks for bidirectional message forwarding
tasks: Final = [
asyncio.create_task(forward_client_to_upstream()),
asyncio.create_task(forward_upstream_to_client()),
]
client_task: Final = asyncio.create_task(forward_client_to_upstream())
upstream_task: Final = asyncio.create_task(forward_upstream_to_client())
tasks: Final = (client_task, upstream_task)
done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
if client_task.done() and _turn_in_flight(upstream_frame_kind, websocket_messages):
await asyncio.wait((upstream_task,), timeout=walkout_turn_drain_seconds)
done: Final = tuple(task for task in tasks if task.done())
pending: Final = tuple(task for task in tasks if not task.done())
# Cancel remaining tasks
for task in pending:
@ -2921,64 +3389,17 @@ async def websocket_passthrough_request(
reason=_truncated_close_reason(upstream_close.reason),
)
end_time: Final = datetime.now()
# Update passthrough logging payload with response data
passthrough_logging_payload["response_body"] = websocket_messages
passthrough_logging_payload["end_time"] = end_time
# Remove logging_obj from kwargs to avoid duplicate keyword argument
success_kwargs: Final = kwargs.copy()
success_kwargs.pop("logging_obj", None)
# # Add user authentication context for database logging
# if user_api_key_dict:
# success_kwargs.setdefault('litellm_params', {})
# success_kwargs['litellm_params'].update({
# 'proxy_server_request': {
# 'body': {
# 'user': user_api_key_dict.user_id,
# 'team_id': user_api_key_dict.team_id,
# 'end_user_id': user_api_key_dict.end_user_id,
# }
# }
# })
# # Also add the user_api_key for direct access
# success_kwargs['user_api_key'] = user_api_key_dict.api_key
# Create a dummy httpx.Response for WebSocket connections
class MockWebSocketResponse:
def __init__(self, target_url: str):
self.status_code = 200
self.text = "WebSocket connection successful"
self.headers: dict[str, str] = {}
self.request = MockWebSocketRequest(target_url)
class MockWebSocketRequest:
def __init__(self, target_url: str):
self.method = "WEBSOCKET"
self.url = target_url
mock_response: Final = MockWebSocketResponse(target)
if meter is not None:
_log_session_usage(logging_obj, meter.price(metered_total))
# Use the same success handler as HTTP passthrough endpoints
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
async_coroutine=pass_through_endpoint_logging.pass_through_async_success_handler(
httpx_response=mock_response,
response_body=websocket_messages,
url_route=endpoint or "",
result="websocket_connection_successful",
start_time=start_time,
end_time=end_time,
logging_obj=logging_obj,
cache_hit=False,
request_body={},
**success_kwargs,
if _session_end_needs_a_row(upstream_frame_kind, websocket_messages, ledger):
lanes_in_flight: Final = (
() if upstream_frame_kind is None else _lanes_in_flight(upstream_frame_kind, websocket_messages)
)
)
bind_budget_reservation_to_callbacks(logging_obj.litellm_params)
*earlier_lanes, last_lane = lanes_in_flight or (None,)
for lane in earlier_lanes:
log_turn(lane, _frames_of_lane(lane))
log_turn(last_lane, _every_frame)
# Call the proxy logging success hook
if proxy_logging_obj:

70
litellm/proxy/pass_through_endpoints/success_handler.py
View file

@ -1,5 +1,6 @@
import asyncio
import json
from collections.abc import Mapping, Sequence
from contextlib import suppress
from datetime import datetime
from types import MappingProxyType
@ -13,7 +14,7 @@ from litellm._logging import verbose_proxy_logger
from litellm.constants import AZURE_SPEECH_CUSTOM_LLM_PROVIDER
from litellm.litellm_core_utils.core_helpers import bind_budget_reservation_to_callbacks
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import PassThroughEndpointLoggingResultValues
from litellm.proxy._types import PassThroughEndpointLoggingResultValues, PassThroughEndpointLoggingTypedDict
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
PassthroughStandardLoggingPayload,
)
@ -70,6 +71,14 @@ def _safe_response_text(httpx_response: httpx.Response) -> str:
return ""
def _websocket_frames(
response_body: Mapping[str, object] | Sequence[Mapping[str, object]] | None,
) -> tuple[Mapping[str, object], ...]:
return tuple(
message for message in (response_body if isinstance(response_body, list) else ()) if isinstance(message, dict)
)
class PassThroughEndpointLogging:
def __init__(
self,
@ -120,6 +129,7 @@ class PassThroughEndpointLogging:
# Vertex AI Live API WebSocket
self.TRACKED_VERTEX_AI_LIVE_ROUTES = ["/vertex_ai/live"]
self.TRACKED_OPENAI_WEBSOCKET_ROUTES: Final = ("/openai/", "/openai_passthrough/")
@property
def _log_dispatch(self) -> PassThroughLogDispatch:
@ -422,21 +432,18 @@ class PassThroughEndpointLogging:
standard_logging_response_object = vertex_ai_live_handler_result["result"]
kwargs = vertex_ai_live_handler_result["kwargs"]
elif DeepgramListenPassthroughLoggingHandler.is_deepgram_listen_route(url_route):
deepgram_handler_result: Final = (
DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler(
websocket_messages=tuple(
message
for message in (response_body if isinstance(response_body, list) else ())
if isinstance(message, dict)
),
logging_obj=logging_obj,
upstream_url=str(httpx_response.request.url),
kwargs=kwargs,
)
elif (
websocket_relay_result := self._websocket_relay_logging_result(
url_route=url_route,
websocket_frames=_websocket_frames(response_body),
logging_obj=logging_obj,
httpx_response=httpx_response,
start_time=start_time,
kwargs=kwargs,
)
standard_logging_response_object = deepgram_handler_result["result"] # rebind-ok: elif-chain
kwargs = deepgram_handler_result["kwargs"] # rebind-ok: elif-chain contract
) is not None:
standard_logging_response_object = websocket_relay_result["result"] # rebind-ok: elif-chain
kwargs = websocket_relay_result["kwargs"] # rebind-ok: elif-chain contract
elif FalAIPassthroughLoggingHandler.is_fal_ai_route(url_route, custom_llm_provider):
fal_ai_handler_result: Final = FalAIPassthroughLoggingHandler().fal_ai_passthrough_handler(
response_body=response_body if isinstance(response_body, dict) else MappingProxyType({}),
@ -452,6 +459,36 @@ class PassThroughEndpointLogging:
return_dict["kwargs"] = kwargs
return return_dict
def _websocket_relay_logging_result(
self,
url_route: str,
websocket_frames: tuple[Mapping[str, object], ...],
logging_obj: LiteLLMLoggingObj,
httpx_response: httpx.Response,
start_time: datetime,
kwargs: Mapping[str, object],
) -> PassThroughEndpointLoggingTypedDict | None:
if self.is_openai_websocket_route(url_route):
from .llm_provider_handlers.openai_websocket_passthrough_logging_handler import (
OpenAIWebsocketPassthroughLoggingHandler,
)
return OpenAIWebsocketPassthroughLoggingHandler().openai_websocket_passthrough_handler(
websocket_messages=websocket_frames,
logging_obj=logging_obj,
url_route=url_route,
start_time=start_time,
kwargs=kwargs,
)
if DeepgramListenPassthroughLoggingHandler.is_deepgram_listen_route(url_route):
return DeepgramListenPassthroughLoggingHandler().deepgram_listen_passthrough_handler(
websocket_messages=websocket_frames,
logging_obj=logging_obj,
upstream_url=str(httpx_response.request.url),
kwargs=kwargs,
)
return None
async def pass_through_async_success_handler(
self,
httpx_response: httpx.Response,
@ -625,6 +662,9 @@ class PassThroughEndpointLogging:
return True
return False
def is_openai_websocket_route(self, url_route: str) -> bool:
return url_route.startswith(self.TRACKED_OPENAI_WEBSOCKET_ROUTES)
def is_cursor_route(self, url_route: str, custom_llm_provider: str | None = None):
"""Check if the URL route is a Cursor Cloud Agents API route."""
if custom_llm_provider == "cursor":

2
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -432,7 +432,7 @@ async def invalidate_budget_reservation_counters(
async def release_or_invalidate_budget_reservation(
budget_reservation: dict | None, # mutable-ok: stamps finalized on the caller's shared reservation dict
budget_reservation: dict[str, object] | None, # mutable-ok: stamps finalized on the caller's reservation
) -> None:
"""Reconcile a still-open reservation on a terminal path that settles no cost.

12
litellm/types/passthrough_endpoints/pass_through_endpoints.py
View file

@ -1,9 +1,11 @@
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from typing import Final, Protocol
from typing import Final, Literal, Protocol, TypeAlias
from pydantic import Field
from typing_extensions import TypedDict
from typing_extensions import ReadOnly, TypedDict
from litellm.types.llms.base import LiteLLMPydanticObjectBase
@ -28,6 +30,8 @@ LITELLM_PASS_THROUGH_ENDPOINT_MARKER: Final = "__litellm_pass_through_endpoint__
SETTLES_BACKGROUND_RESPONSE_KWARG: Final = "settles_background_response"
WebsocketUpstreamFrameKind: TypeAlias = Literal["starts_turn", "ends_turn", "opens_session", "refuses_turn", "relayed"]
class BackgroundResponseSettlement(str, Enum):
SETTLES = "settles"
@ -63,7 +67,7 @@ class PassthroughStandardLoggingPayload(TypedDict, total=False):
"""
The body of the request
"""
response_body: dict | None # only tracked for non-streaming responses
response_body: dict | Sequence[Mapping[str, object]] | None # only tracked for non-streaming responses
"""
The body of the response
"""
@ -75,6 +79,8 @@ class PassthroughStandardLoggingPayload(TypedDict, total=False):
Optional field, we use this for cost tracking only if it's set.
"""
end_time: ReadOnly[datetime]
@dataclass(frozen=True, slots=True)
class WebSocketSessionUsage:

133
tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py
View file

@ -10998,139 +10998,6 @@ async def test_key_budget_error_keeps_the_masked_key_name(key_name):
assert f"Key=prod-key ({key_name}) Current cost" in message
class _UntouchedPrisma:
def __getattr__(self, name: str) -> object:
raise AssertionError(f"database reached through {name}")
class _MissingUserPrisma:
class db:
class litellm_usertable:
@staticmethod
async def find_unique(where: dict[str, str], include: dict[str, bool]) -> None:
return None
@pytest.mark.asyncio
async def test_enforced_model_allowlists_treats_a_missing_user_row_as_unrestricted():
from litellm.proxy.auth.auth_checks import enforced_model_allowlists
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.utils import ProxyLogging
cache = UserApiKeyCache()
scopes = await enforced_model_allowlists(
valid_token=UserAPIKeyAuth(token="hashed-fake", user_id="default_user_id"),
prisma_client=_MissingUserPrisma(),
user_api_key_cache=cache,
proxy_logging_obj=ProxyLogging(user_api_key_cache=cache),
)
assert [list(scope) for scope in scopes] == [[], [], [], [], []]
class _UnreachableUserPrisma:
class db:
class litellm_usertable:
@staticmethod
async def find_unique(where: dict[str, str], include: dict[str, bool]) -> None:
raise RuntimeError("database gone")
@pytest.mark.asyncio
async def test_enforced_model_allowlists_surfaces_a_failed_user_lookup():
from litellm.proxy.auth.auth_checks import enforced_model_allowlists
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.utils import ProxyLogging
cache = UserApiKeyCache()
with pytest.raises(ValueError, match="database gone"):
await enforced_model_allowlists(
valid_token=UserAPIKeyAuth(token="hashed-fake", user_id="user-fake"),
prisma_client=_UnreachableUserPrisma(),
user_api_key_cache=cache,
proxy_logging_obj=ProxyLogging(user_api_key_cache=cache),
)
@pytest.mark.asyncio
async def test_enforced_model_allowlists_reads_every_level_from_cache():
from litellm.proxy._types import (
LiteLLM_BudgetTable,
LiteLLM_ProjectTableCachedObj,
LiteLLM_TeamMembership,
LiteLLM_TeamTableCachedObj,
LiteLLM_UserTable,
)
from litellm.proxy.auth.auth_checks import enforced_model_allowlists
from litellm.proxy.common_utils.user_api_key_cache import (
UserApiKeyCache,
team_membership_reservation_cache_key,
)
from litellm.proxy.utils import ProxyLogging
cache = UserApiKeyCache()
proxy_logging_obj = ProxyLogging(user_api_key_cache=cache)
await cache.async_set_cache(
key="team_id:team-fake", value=LiteLLM_TeamTableCachedObj(team_id="team-fake", models=["gpt-4o"])
)
await cache.async_set_cache(
key=team_membership_reservation_cache_key(user_id="user-fake", team_id="team-fake"),
value=LiteLLM_TeamMembership(
user_id="user-fake",
team_id="team-fake",
litellm_budget_table=LiteLLM_BudgetTable(allowed_models=["gpt-4o-mini"]),
),
)
await cache.async_set_cache(
key="project_id:project-fake",
value=LiteLLM_ProjectTableCachedObj(project_id="project-fake", models=["gpt-4.1"]),
)
await cache.async_set_cache(key="user-fake", value=LiteLLM_UserTable(user_id="user-fake", models=["o3"]))
prisma_client = _UntouchedPrisma()
team_scoped = await enforced_model_allowlists(
valid_token=UserAPIKeyAuth(
token="hashed-fake",
models=["all-team-models"],
team_models=["gpt-4o", "gpt-4o-mini"],
user_id="user-fake",
team_id="team-fake",
project_id="project-fake",
),
prisma_client=prisma_client,
user_api_key_cache=cache,
proxy_logging_obj=proxy_logging_obj,
)
personal = await enforced_model_allowlists(
valid_token=UserAPIKeyAuth(token="hashed-fake", user_id="user-fake"),
prisma_client=prisma_client,
user_api_key_cache=cache,
proxy_logging_obj=proxy_logging_obj,
)
without_database = await enforced_model_allowlists(
valid_token=UserAPIKeyAuth(
token="hashed-fake",
models=["gpt-4o"],
team_models=["gpt-4o-mini"],
user_id="user-fake",
team_id="team-fake",
),
prisma_client=None,
user_api_key_cache=cache,
proxy_logging_obj=proxy_logging_obj,
)
assert [list(scope) for scope in team_scoped] == [
["gpt-4o", "gpt-4o-mini"],
["gpt-4o"],
["gpt-4o-mini"],
[],
["gpt-4.1"],
]
assert [list(scope) for scope in personal] == [[], [], [], ["o3"], []]
assert [list(scope) for scope in without_database] == [["gpt-4o"], ["gpt-4o-mini"]]
@pytest.mark.asyncio
@pytest.mark.parametrize("channel", ["team", "key"])
async def test_access_group_model_fallback_uses_the_injected_database(channel: str) -> None:

9
tests/unit/proxy/common_utils/test_realtime_cache.py
View file

@ -1,4 +1,5 @@
from typing import Any, cast
import json
from typing import Any, Final, 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: Final = 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")

518
tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py Normal file
View file

@ -0,0 +1,518 @@
import uuid
from collections.abc import Mapping, Sequence
from datetime import datetime
from types import MappingProxyType
from typing import Final
from unittest.mock import MagicMock
import httpx
import pytest
import litellm
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import (
OpenAIPassthroughLoggingHandler,
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import (
OpenAIWebsocketPassthroughLoggingHandler,
openai_websocket_frame_kind,
)
from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id
from litellm.types.utils import CallTypes
REALTIME_MODEL: Final = "gpt-realtime-2.1"
RESPONSES_MODEL: Final = "gpt-5.4-mini"
REASONING_MODEL: Final = "gpt-5.4"
TRANSCRIPTION_MODEL: Final = "gpt-4o-mini-transcribe"
SESSION_START: Final = datetime(2026, 1, 1, 12, 0, 0)
def _logging_obj() -> MagicMock:
logging_obj: Final = MagicMock(model_call_details={}, litellm_call_id=str(uuid.uuid4()))
return logging_obj
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[str, object]:
return {
"type": "response.done",
"response": {
"id": "resp_fake",
"status": "completed",
"usage": {
"input_tokens": text_in + audio_in,
"output_tokens": text_out + audio_out,
"total_tokens": text_in + audio_in + text_out + audio_out,
"input_token_details": {"text_tokens": text_in, "audio_tokens": audio_in, "cached_tokens": 0},
"output_token_details": {"text_tokens": text_out, "audio_tokens": audio_out},
},
},
}
def _responses_completed(
model: str,
input_tokens: int,
output_tokens: int,
event: str = "response.completed",
served: Mapping[str, object] = MappingProxyType({}),
) -> dict[str, object]:
return {
"type": event,
"response": {
"id": "resp_fake",
"object": "response",
"created_at": 1761700000,
"model": model,
"status": "completed",
"output": [],
**served,
"usage": {
"input_tokens": input_tokens,
"output_tokens": output_tokens,
"total_tokens": input_tokens + output_tokens,
"input_tokens_details": {"cached_tokens": 0},
"output_tokens_details": {"reasoning_tokens": 0},
},
},
}
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(),
url_route=url_route,
start_time=SESSION_START,
kwargs=kwargs,
)
def _realtime_cost(model: str, text_in: int, audio_in: int, text_out: int, audio_out: int) -> float:
prices: Final = litellm.model_cost[model]
return (
text_in * prices["input_cost_per_token"]
+ audio_in * prices["input_cost_per_audio_token"]
+ text_out * prices["output_cost_per_token"]
+ audio_out * prices["output_cost_per_audio_token"]
)
def test_realtime_session_is_priced_on_the_model_the_session_reported():
logging_obj: Final = _logging_obj()
messages: Final = [
_session_created(REALTIME_MODEL),
{"type": "response.created", "response": {"id": "resp_fake"}},
_realtime_response_done(text_in=33, audio_in=120, text_out=16, audio_out=240),
]
handled: Final = _handle(messages, logging_obj=logging_obj)
expected_cost: Final = _realtime_cost(REALTIME_MODEL, 33, 120, 16, 240)
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
assert handled["kwargs"]["model"] == REALTIME_MODEL
assert handled["kwargs"]["custom_llm_provider"] == "openai"
assert handled["result"].model == REALTIME_MODEL
assert handled["result"].usage.total_tokens == 409
assert logging_obj.model_call_details["response_cost"] == pytest.approx(expected_cost)
assert logging_obj.model_call_details["model"] == REALTIME_MODEL
def test_realtime_cost_covers_every_turn_of_the_session():
one_turn: Final = _handle([_session_created(REALTIME_MODEL), _realtime_response_done(10, 0, 5, 0)])
two_turns: Final = _handle(
[
_session_created(REALTIME_MODEL),
_realtime_response_done(10, 0, 5, 0),
_realtime_response_done(10, 0, 5, 0),
]
)
assert two_turns["kwargs"]["response_cost"] == pytest.approx(2 * one_turn["kwargs"]["response_cost"])
assert two_turns["result"].usage.total_tokens == 30
def test_responses_turn_is_priced_on_the_model_the_response_reported():
handled: Final = _handle(
[_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5)],
url_route="/openai_passthrough/v1/responses",
)
prices: Final = litellm.model_cost[RESPONSES_MODEL]
expected_cost: Final = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"]
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
assert handled["kwargs"]["model"] == RESPONSES_MODEL
assert handled["result"].usage.prompt_tokens == 13
assert handled["result"].usage.completion_tokens == 5
def test_responses_cost_covers_every_turn_of_the_connection():
handled: Final = _handle(
[
_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5),
_responses_completed(RESPONSES_MODEL, input_tokens=20, output_tokens=7),
],
url_route="/openai_passthrough/v1/responses",
)
prices: Final = litellm.model_cost[RESPONSES_MODEL]
expected_cost: Final = 33 * prices["input_cost_per_token"] + 12 * prices["output_cost_per_token"]
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
assert handled["result"].usage.total_tokens == 45
def test_a_model_the_caller_asked_for_never_prices_a_session_the_provider_never_named():
"""
The caller picks what to ask for, and the provider picks what to run. Only
the provider's own frames may set the price of a session, so a model the
caller put in the request cannot stand in for one the provider never named.
"""
logging_obj: Final = _logging_obj()
handled: Final = _handle(
[_realtime_response_done(text_in=10, audio_in=0, text_out=5, audio_out=0)],
kwargs=MappingProxyType({"model": REALTIME_MODEL}),
logging_obj=logging_obj,
)
assert handled["result"] is None
assert handled["kwargs"] == {"model": REALTIME_MODEL}
assert logging_obj.model_call_details == {}
def test_each_responses_turn_is_priced_on_the_model_that_turn_ran():
"""
One socket can run several models one after another, so a connection that
starts on a big model and finishes on a small one has to pay the big
model's price for the turn that used it.
"""
handled: Final = _handle(
[
_responses_completed(REASONING_MODEL, input_tokens=13, output_tokens=5),
_responses_completed(RESPONSES_MODEL, input_tokens=20, output_tokens=7),
],
url_route="/openai_passthrough/v1/responses",
)
first: Final = litellm.model_cost[REASONING_MODEL]
second: Final = litellm.model_cost[RESPONSES_MODEL]
expected_cost: Final = (
13 * first["input_cost_per_token"]
+ 5 * first["output_cost_per_token"]
+ 20 * second["input_cost_per_token"]
+ 7 * second["output_cost_per_token"]
)
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
assert handled["result"].usage.total_tokens == 45
def test_a_responses_turn_that_names_no_model_is_priced_on_the_session_model():
handled: Final = _handle(
[
_session_created(RESPONSES_MODEL),
_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, served={"model": None}),
],
url_route="/openai_passthrough/v1/responses",
)
prices: Final = litellm.model_cost[RESPONSES_MODEL]
expected_cost: Final = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"]
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
def test_a_session_nobody_named_a_model_for_is_left_unpriced():
logging_obj: Final = _logging_obj()
handled: Final = _handle(
[{"type": "response.done", "response": {"usage": {"input_tokens": 10, "output_tokens": 5}}}],
kwargs=MappingProxyType({"litellm_call_id": "call-fake"}),
logging_obj=logging_obj,
)
assert handled["result"] is None
assert handled["kwargs"] == {"litellm_call_id": "call-fake"}
assert logging_obj.model_call_details == {}
def test_a_session_that_reported_no_usage_still_records_its_model():
logging_obj: Final = _logging_obj()
handled: Final = _handle([_session_created(REALTIME_MODEL)], logging_obj=logging_obj)
assert handled["result"] is None
assert handled["kwargs"] == {"model": REALTIME_MODEL, "custom_llm_provider": "openai"}
assert logging_obj.model_call_details["model"] == REALTIME_MODEL
assert logging_obj.model_call_details["custom_llm_provider"] == "openai"
def test_a_responses_turn_without_readable_usage_is_not_billed():
handled: Final = _handle(
[{"type": "response.completed", "response": {"model": RESPONSES_MODEL, "status": "completed"}}],
url_route="/openai_passthrough/v1/responses",
)
assert handled["result"] is None
assert handled["kwargs"] == {"model": RESPONSES_MODEL, "custom_llm_provider": "openai"}
def test_the_price_of_a_multi_model_connection_reaches_the_spend_row():
"""
The logging layer prices whatever result it is handed, and one result can
only name one model, so a connection that ran two models has to carry the
price it already worked out or the cheaper model pays for both turns. The
row it lands in names the model the connection opened on.
"""
handled: Final = _handle(
[
_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5),
_responses_completed(REASONING_MODEL, input_tokens=20, output_tokens=7),
],
url_route="/openai_passthrough/v1/responses",
)
logged_cost: Final = litellm.response_cost_calculator(
response_object=handled["result"],
model=handled["kwargs"]["model"],
custom_llm_provider="openai",
call_type="completion",
optional_params={},
)
session_prices: Final = litellm.model_cost[RESPONSES_MODEL]
priced_on_the_session_model: Final = (
33 * session_prices["input_cost_per_token"] + 12 * session_prices["output_cost_per_token"]
)
assert logged_cost == pytest.approx(handled["kwargs"]["response_cost"])
assert logged_cost > priced_on_the_session_model
assert handled["kwargs"]["model"] == RESPONSES_MODEL
def test_turns_that_start_in_the_same_instant_each_keep_their_own_spend_row():
"""
Spend logs key each row by the id its result carries and skip a repeated id
without a word, so two turns that start in the same instant, on one socket
or on two, each need an id of their own or one of them is never billed.
"""
turns: Final = [
_handle(
[_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5)],
url_route="/openai_passthrough/v1/responses",
)
for _ in range(2)
]
spend_log_ids: Final = {
get_spend_logs_id(CallTypes.pass_through.value, turn["result"].model_dump(), turn["kwargs"]) for turn in turns
}
assert len(spend_log_ids) == 2
assert None not in spend_log_ids
def test_a_turn_the_caller_walked_out_of_still_names_the_model_it_ran():
"""
A client that closes before response.completed leaves no usage behind to
bill, but response.created already named the model, so the row says which
model answered instead of falling back to unknown.
"""
logging_obj: Final = _logging_obj()
handled: Final = _handle(
[{"type": "response.created", "response": {"id": "resp_fake", "model": REASONING_MODEL}}],
logging_obj=logging_obj,
url_route="/openai_passthrough/v1/responses",
)
assert handled["result"] is None
assert handled["kwargs"] == {"model": REASONING_MODEL, "custom_llm_provider": "openai"}
assert logging_obj.model_call_details["model"] == REASONING_MODEL
def _transcription_session(model: str) -> tuple[dict[str, object], ...]:
return (
{"type": "session.created", "session": {"type": "transcription", "audio": {"input": {"transcription": None}}}},
{
"type": "session.updated",
"session": {"type": "transcription", "audio": {"input": {"transcription": {"model": model}}}},
},
)
def _transcription_completed(usage: Mapping[str, object]) -> dict[str, object]:
return {
"type": "conversation.item.input_audio_transcription.completed",
"item_id": "item_fake",
"transcript": "hello",
"usage": usage,
}
@pytest.mark.parametrize(
("usage", "billed_units"),
[
({"type": "duration", "seconds": 3}, ((3, "input_cost_per_second"),)),
(
{
"type": "tokens",
"input_tokens": 30,
"output_tokens": 11,
"input_token_details": {"text_tokens": 0, "audio_tokens": 30},
},
((30, "input_cost_per_audio_token"), (11, "output_cost_per_token")),
),
],
)
def test_a_transcription_session_is_billed_for_every_clip_it_transcribed(usage, billed_units):
handled: Final = _handle(
[
*_transcription_session(TRANSCRIPTION_MODEL),
_transcription_completed(usage),
_transcription_completed(usage),
],
url_route="/openai_passthrough/v1/realtime",
)
logged_cost: Final = litellm.response_cost_calculator(
response_object=handled["result"],
model=handled["kwargs"]["model"],
custom_llm_provider="openai",
call_type="completion",
optional_params={},
)
prices: Final = litellm.model_cost[TRANSCRIPTION_MODEL]
expected_cost: Final = 2 * sum(units * prices[rate] for units, rate in billed_units)
assert expected_cost > 0
assert handled["kwargs"]["model"] == TRANSCRIPTION_MODEL
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
assert logged_cost == pytest.approx(expected_cost)
def test_a_conversation_that_also_transcribes_the_caller_is_logged_on_its_own_model():
handled: Final = _handle(
[
{
"type": "session.created",
"session": {
"model": REALTIME_MODEL,
"audio": {"input": {"transcription": {"model": TRANSCRIPTION_MODEL}}},
},
},
_realtime_response_done(text_in=10, audio_in=0, text_out=5, audio_out=0),
]
)
assert handled["kwargs"]["model"] == REALTIME_MODEL
@pytest.mark.parametrize("event", ["response.incomplete", "response.failed"])
def test_a_responses_turn_that_ended_early_is_billed_for_the_tokens_it_used(event):
"""
A turn cut short by max_output_tokens or a provider failure still reports
the tokens it burned, and OpenAI bills them, so the row has to as well.
"""
handled: Final = _handle(
[_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, event=event)],
url_route="/openai_passthrough/v1/responses",
)
prices: Final = litellm.model_cost[RESPONSES_MODEL]
expected_cost: Final = 13 * prices["input_cost_per_token"] + 5 * prices["output_cost_per_token"]
assert handled["kwargs"]["response_cost"] == pytest.approx(expected_cost)
assert handled["result"].usage.total_tokens == 18
def test_a_failed_responses_turn_that_reported_no_usage_is_not_billed():
handled: Final = _handle(
[{"type": "response.failed", "response": {"model": RESPONSES_MODEL, "status": "failed", "usage": None}}],
url_route="/openai_passthrough/v1/responses",
)
assert handled["result"] is None
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: Final = _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: Final = litellm.model_cost[priced_as]
expected_cost: Final = 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
def _http_passthrough_cost(response: Mapping[str, object]) -> float:
http_response: Final = httpx.Response(
200, json=response, request=httpx.Request("POST", "https://api.openai.com/v1/responses")
)
_, cost = OpenAIPassthroughLoggingHandler._build_responses_api_response_and_cost(
RESPONSES_MODEL, http_response, _logging_obj(), "openai"
)
return cost
@pytest.mark.parametrize(
"served",
[
{"service_tier": "priority"},
{"output": [{"type": "web_search_call", "id": "ws_fake", "status": "completed", "action": {"type": "search"}}]},
],
ids=["priority_tier", "hosted_web_search"],
)
def test_a_responses_turn_costs_what_the_http_passthrough_charges_for_the_same_response(served):
"""
The provider bills a turn by the tier that served it and the hosted tools it
called, not only its tokens, so a websocket turn has to cost exactly what the
same response costs over HTTP, never the cheaper token-only price.
"""
turn: Final = _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, served=served)
plain_turn: Final = _responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5)
handled: Final = _handle([turn], url_route="/openai_passthrough/v1/responses")
assert handled["kwargs"]["response_cost"] == pytest.approx(_http_passthrough_cost(turn["response"]))
assert handled["kwargs"]["response_cost"] > _http_passthrough_cost(plain_turn["response"])
@pytest.mark.parametrize(
("frame_type", "kind"),
[
("response.completed", "ends_turn"),
("response.incomplete", "ends_turn"),
("response.failed", "ends_turn"),
("response.done", "ends_turn"),
("session.created", "opens_session"),
("response.created", "starts_turn"),
("response.output_text.delta", "relayed"),
("error", "refuses_turn"),
],
)
def test_each_frame_is_classified_by_what_it_does_to_a_turn(frame_type, kind):
assert openai_websocket_frame_kind({"type": frame_type}) == kind

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

File diff suppressed because it is too large Load diff

53
tests/unit/proxy/pass_through_endpoints/test_success_handler.py
View file

@ -652,6 +652,59 @@ def test_vertex_ai_live_route_without_usage_frames_yields_no_logging_response(
assert normalized["kwargs"] == {"model": _LIVE_MODEL}
REALTIME_FRAMES: Final = [
{"type": "session.created", "session": {"id": "sess_fake", "model": "gpt-realtime-2.1"}},
{
"type": "response.done",
"response": {
"usage": {
"input_tokens": 33,
"output_tokens": 16,
"total_tokens": 49,
"input_token_details": {"text_tokens": 33, "audio_tokens": 0, "cached_tokens": 0},
"output_token_details": {"text_tokens": 16, "audio_tokens": 0},
}
},
},
]
def _normalize(url_route: str, response_body: dict[str, object] | list[dict[str, object]]) -> dict:
logging_obj: Final = MagicMock(model_call_details={}, litellm_call_id="call-fake")
return PassThroughEndpointLogging().normalize_llm_passthrough_logging_payload(
httpx_response=MagicMock(spec=httpx.Response),
response_body=response_body,
request_body={},
logging_obj=logging_obj,
url_route=url_route,
result="websocket_connection_successful",
start_time=_START,
end_time=_END,
cache_hit=False,
litellm_call_id="call-fake",
)
@pytest.mark.parametrize("url_route", ["/openai/v1/realtime", "/openai_passthrough/v1/realtime"])
def test_openai_websocket_frames_reach_the_websocket_cost_handler(url_route):
normalized: Final = _normalize(url_route, REALTIME_FRAMES)
assert normalized["kwargs"]["model"] == "gpt-realtime-2.1"
assert normalized["kwargs"]["custom_llm_provider"] == "openai"
prices: Final = litellm.model_cost["gpt-realtime-2.1"]
assert normalized["kwargs"]["response_cost"] == pytest.approx(
33 * prices["input_cost_per_token"] + 16 * prices["output_cost_per_token"]
)
assert normalized["standard_logging_response_object"].usage.total_tokens == 49
def test_openai_http_passthrough_bodies_are_left_to_the_http_handlers():
normalized: Final = _normalize("/openai/v1/responses", {"model": "gpt-5.4-mini", "usage": {"input_tokens": 13}})
assert normalized["standard_logging_response_object"] is None
assert normalized["kwargs"] == {"litellm_call_id": "call-fake"}
@pytest.mark.asyncio
@pytest.mark.parametrize("cancel_mode", ["once", "twice", "scope", "none"])
@pytest.mark.parametrize("cancel_phase", ["allocation", "lock"])

439
tests/unit/proxy/test_openai_ws_passthrough_routes.py
View file

@ -1,31 +1,41 @@
"""OpenAI passthrough WebSocket route: registration, opt-in gating, and refusals."""
"""OpenAI passthrough WebSocket route: registration, opt-in gating, refusals, and per-frame model checks."""
import json
from collections.abc import Mapping, Sequence
from collections.abc import Mapping
from dataclasses import dataclass
from types import MappingProxyType, SimpleNamespace
from typing import Final
from unittest.mock import patch
from unittest.mock import AsyncMock, patch
import httpx
import pytest
import respx
from fastapi import HTTPException
from starlette.routing import WebSocketRoute
from starlette.websockets import WebSocket
import litellm
from litellm.llms.openai.workload_identity import _workload_identity_auth
from litellm.proxy._types import UserAPIKeyAuth
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_MODEL_RESTRICTED_REFUSAL,
_has_model_restrictions,
_OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL,
_openai_websocket_refusal,
_proxy_model_allowlists,
_proxy_frame_model_gate,
openai_websocket_proxy_route,
router,
)
Scopes = tuple[Sequence[str], ...]
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import (
openai_websocket_frame_kind,
)
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
WebsocketFrameModelAllowed,
WebsocketFrameModelGate,
WebsocketFrameModelRefused,
WebsocketFrameModelVerdict,
WebsocketUpstreamFrameClassifier,
)
ENABLED: Final = MappingProxyType({"enable_openai_websocket_passthrough": True})
DISABLED_SETTINGS: Final = (
@ -76,6 +86,8 @@ class _RelayCall:
forward_headers: bool
endpoint: str
accept_websocket: bool
client_frame_model_gate: WebsocketFrameModelGate | None
upstream_frame_kind: WebsocketUpstreamFrameClassifier | None
class _FakeRelay:
@ -92,6 +104,8 @@ class _FakeRelay:
forward_headers: bool,
endpoint: str,
accept_websocket: bool,
client_frame_model_gate: WebsocketFrameModelGate | None = None,
upstream_frame_kind: WebsocketUpstreamFrameClassifier | None = None,
) -> None:
self.calls.append(
_RelayCall(
@ -100,24 +114,26 @@ class _FakeRelay:
forward_headers=forward_headers,
endpoint=endpoint,
accept_websocket=accept_websocket,
client_frame_model_gate=client_frame_model_gate,
upstream_frame_kind=upstream_frame_kind,
)
)
class _FakeModelAllowlists:
def __init__(self, scopes: Scopes) -> None:
self.scopes = scopes
self.calls: list[UserAPIKeyAuth] = []
class _FakeFrameModelGate:
def __init__(self, refusals: Mapping[str, WebsocketFrameModelRefused] = MappingProxyType({})) -> None:
self.refusals = refusals
self.checked: list[str] = []
async def __call__(self, valid_token: UserAPIKeyAuth, /) -> Scopes:
self.calls.append(valid_token)
return self.scopes
async def __call__(self, websocket: object, model: str, /) -> WebsocketFrameModelVerdict:
self.checked.append(model)
return self.refusals.get(model) or WebsocketFrameModelAllowed(access_groups=())
@dataclass(frozen=True, slots=True)
class _Served:
relay: _FakeRelay
allowlists: _FakeModelAllowlists
frame_model_gate: _FakeFrameModelGate
async def _serve(
@ -125,16 +141,16 @@ async def _serve(
endpoint: str,
user_api_key_dict: UserAPIKeyAuth,
general_settings: Mapping[str, object],
scopes: Scopes = (),
frame_model_gate: _FakeFrameModelGate | None = None,
) -> _Served:
served = _Served(relay=_FakeRelay(), allowlists=_FakeModelAllowlists(scopes))
served: Final = _Served(relay=_FakeRelay(), frame_model_gate=frame_model_gate or _FakeFrameModelGate())
await openai_websocket_proxy_route(
websocket=websocket,
endpoint=endpoint,
user_api_key_dict=user_api_key_dict,
general_settings=general_settings,
relay=served.relay,
model_allowlists=served.allowlists,
frame_model_gate=served.frame_model_gate,
)
return served
@ -143,18 +159,20 @@ async def _serve(
@pytest.mark.parametrize("prefix", ["openai", "openai_passthrough"])
async def test_openai_websocket_forwards_query_and_keeps_provider_auth(prefix, monkeypatch):
monkeypatch.delenv("OPENAI_API_BASE", raising=False)
websocket = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-4o-realtime-preview")
websocket: Final = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-realtime-2.1")
with patch(GET_CREDENTIALS, return_value="sk-provider"):
served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED)
assert served.relay.calls == [
_RelayCall(
target="wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview",
target="wss://api.openai.com/v1/realtime?model=gpt-realtime-2.1",
custom_headers=MappingProxyType({"Authorization": "Bearer sk-provider"}),
forward_headers=False,
endpoint=f"/{prefix}/v1/realtime",
accept_websocket=False,
client_frame_model_gate=served.frame_model_gate,
upstream_frame_kind=openai_websocket_frame_kind,
)
]
assert websocket.accepts == [None]
@ -162,11 +180,23 @@ 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: Final = _FakeWebSocket("/litellm/openai/v1/responses", "")
with patch(GET_CREDENTIALS, return_value="sk-provider"):
served: Final = 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(
"/openai/v1/realtime",
"model=gpt-4o-realtime-preview",
"model=gpt-realtime-2.1",
subprotocols="realtime, openai-insecure-api-key.sk-abc, openai-beta.realtime-v1",
)
@ -239,7 +269,7 @@ async def test_openai_websocket_closes_cleanly_when_workload_identity_exchange_f
@pytest.mark.asyncio
async def test_openai_websocket_closes_cleanly_when_provider_credentials_missing():
websocket = _FakeWebSocket("/openai/v1/realtime", "model=gpt-4o-realtime-preview")
websocket: Final = _FakeWebSocket("/openai/v1/realtime", "model=gpt-realtime-2.1")
with patch(GET_CREDENTIALS, return_value=None):
served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), ENABLED)
@ -255,7 +285,7 @@ async def test_openai_websocket_closes_cleanly_when_provider_credentials_missing
@pytest.mark.parametrize("prefix", ["openai", "openai_passthrough"])
@pytest.mark.parametrize("general_settings", DISABLED_SETTINGS)
async def test_openai_websocket_refused_unless_explicitly_enabled(prefix, general_settings):
websocket = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-4o-realtime-preview")
websocket: Final = _FakeWebSocket(f"/{prefix}/v1/realtime", "model=gpt-realtime-2.1")
served = await _serve(websocket, "v1/realtime", UserAPIKeyAuth(), general_settings)
@ -265,25 +295,21 @@ async def test_openai_websocket_refused_unless_explicitly_enabled(prefix, genera
assert served.relay.calls == []
@pytest.mark.asyncio
@pytest.mark.parametrize("general_settings", DISABLED_SETTINGS)
async def test_openai_websocket_refusal_is_disabled_for_falsy_settings(general_settings):
refusal = await _openai_websocket_refusal(UserAPIKeyAuth(), general_settings, _FakeModelAllowlists(()))
assert refusal is _OPENAI_WS_DISABLED_REFUSAL
def test_openai_websocket_refusal_is_disabled_for_falsy_settings(general_settings):
assert _openai_websocket_refusal(general_settings) is _OPENAI_WS_DISABLED_REFUSAL
@pytest.mark.asyncio
@pytest.mark.parametrize("value", [True, "true", "True"])
async def test_openai_websocket_refusal_is_none_for_truthy_settings(value):
settings = MappingProxyType({"enable_openai_websocket_passthrough": value})
assert await _openai_websocket_refusal(UserAPIKeyAuth(), settings, _FakeModelAllowlists(())) is None
def test_openai_websocket_refusal_is_none_for_truthy_settings(value):
assert _openai_websocket_refusal(MappingProxyType({"enable_openai_websocket_passthrough": value})) is None
@pytest.mark.asyncio
async def test_openai_websocket_refusal_echoes_requested_subprotocol():
websocket = _FakeWebSocket(
"/openai_passthrough/v1/realtime",
"model=gpt-4o-realtime-preview",
"model=gpt-realtime-2.1",
subprotocols="realtime, openai-beta.realtime-v1",
)
@ -294,68 +320,303 @@ async def test_openai_websocket_refusal_echoes_requested_subprotocol():
assert served.relay.calls == []
RESTRICTED_SCOPES: Final[tuple[Scopes, ...]] = (
(("gpt-4o",),),
((), ("gpt-4o-realtime-preview",)),
(("all-team-models",), ("gpt-4o",)),
((), ("all-proxy-models",), ("gpt-4o",)),
((), (), (), ("gpt-4o",)),
(("*",), (), (), (), ("gpt-4o",)),
@pytest.mark.asyncio
@pytest.mark.parametrize(
"key_models",
[["gpt-5.4-mini"], ["gpt-realtime-2.1", "gpt-5.4"], []],
)
UNRESTRICTED_SCOPES: Final[tuple[Scopes, ...]] = (
(),
((),),
(("all-proxy-models",),),
(("*",),),
(("all-team-models",), ("all-proxy-models",)),
((), (), (), (), ()),
(("*",), ("all-proxy-models",), ("all-team-models",), (), ()),
)
@pytest.mark.asyncio
@pytest.mark.parametrize("scopes", RESTRICTED_SCOPES)
async def test_openai_websocket_rejects_model_restricted_identities(scopes):
websocket = _FakeWebSocket("/openai/v1/realtime", "model=gpt-4o-realtime-preview")
user_api_key_dict = UserAPIKeyAuth(token="hashed-fake", user_id="user-fake", team_id="team-fake")
served = await _serve(websocket, "v1/realtime", user_api_key_dict, ENABLED, scopes)
assert "model restrictions" in websocket.error_message()
assert websocket.closed == (1008, _OPENAI_WS_MODEL_RESTRICTED_REFUSAL.close_reason)
assert served.relay.calls == []
assert served.allowlists.calls == [user_api_key_dict]
@pytest.mark.asyncio
@pytest.mark.parametrize("scopes", RESTRICTED_SCOPES)
async def test_openai_websocket_disabled_refusal_skips_allowlist_lookups(scopes):
allowlists = _FakeModelAllowlists(scopes)
refusal = await _openai_websocket_refusal(UserAPIKeyAuth(), MappingProxyType({}), allowlists)
assert refusal is _OPENAI_WS_DISABLED_REFUSAL
assert allowlists.calls == []
@pytest.mark.asyncio
@pytest.mark.parametrize("scopes", UNRESTRICTED_SCOPES)
async def test_openai_websocket_allows_unrestricted_identities(scopes):
websocket = _FakeWebSocket("/openai/v1/responses", "")
async def test_openai_websocket_connects_model_restricted_keys(key_models):
websocket: Final = _FakeWebSocket("/openai/v1/responses", "")
user_api_key_dict: Final = UserAPIKeyAuth(models=key_models, token="hashed-fake", user_id="user-fake")
with patch(GET_CREDENTIALS, return_value="sk-provider"):
served = await _serve(websocket, "v1/responses", UserAPIKeyAuth(), ENABLED, scopes)
served: Final = await _serve(websocket, "v1/responses", user_api_key_dict, ENABLED)
assert len(served.relay.calls) == 1
assert served.relay.calls[0].client_frame_model_gate is served.frame_model_gate
assert websocket.sent == []
assert websocket.closed is None
@pytest.mark.asyncio
async def test_proxy_model_allowlists_reads_the_token_scopes_without_a_database():
token: Final = UserAPIKeyAuth(models=[], team_id="team-fake", team_models=["gpt-4o"])
with patch("litellm.proxy.proxy_server.prisma_client", None):
scopes = await _proxy_model_allowlists()(token)
ALIAS_REFUSAL: Final = WebsocketFrameModelRefused(
message=_OPENAI_WS_ALIASED_MODEL_REFUSAL, error_type=ProxyErrorTypes.key_model_access_denied.value
)
assert tuple(tuple(scope) for scope in scopes) == ((), ("gpt-4o",))
assert _has_model_restrictions(scopes)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"query",
[
"model=gpt-realtime-2.1",
"model=gpt-realtime-2.1&model=gpt-realtime-2.1-mini",
"model=gpt-realtime-2.1-mini&model=gpt-realtime-2.1",
],
ids=["one_model", "refused_model_first", "refused_model_last"],
)
async def test_openai_websocket_refuses_an_opening_model_the_frame_gate_refuses(query):
"""
The relay sends the connection URL to OpenAI unchanged, so every model it names has to pass the
same gate as a frame, or an alias the upgrade resolves to an allowed model opens the session on the
alias name itself
"""
websocket: Final = _FakeWebSocket("/openai/v1/realtime", query)
gate: Final = _FakeFrameModelGate(MappingProxyType({"gpt-realtime-2.1": ALIAS_REFUSAL}))
with patch(GET_CREDENTIALS, return_value="sk-provider"):
served: Final = await _serve(
websocket, "v1/realtime", UserAPIKeyAuth(models=["gpt-realtime-2.1-mini"]), ENABLED, gate
)
assert websocket.error_message() == _OPENAI_WS_ALIASED_MODEL_REFUSAL
assert websocket.closed == (1008, "This key cannot use a model alias over OpenAI websocket passthrough.")
assert served.relay.calls == []
assert "gpt-realtime-2.1" in gate.checked
@pytest.mark.asyncio
async def test_openai_websocket_relays_a_session_once_every_opening_model_passes_the_frame_gate():
websocket: Final = _FakeWebSocket("/openai/v1/realtime", "model=gpt-realtime-2.1-mini")
with patch(GET_CREDENTIALS, return_value="sk-provider"):
served: Final = await _serve(
websocket, "v1/realtime", UserAPIKeyAuth(models=["gpt-realtime-2.1-mini"]), ENABLED
)
assert served.frame_model_gate.checked == ["gpt-realtime-2.1-mini"]
assert len(served.relay.calls) == 1
assert websocket.closed is None
class _FrameAuthorizer:
"""Stands in for the full key, team, user and budget auth the gate re-runs for every model a frame names."""
def __init__(self, token: UserAPIKeyAuth | None = None, denial: Exception | None = None) -> None:
self.token = token or UserAPIKeyAuth()
self.denial = denial
self.authorized: list[str] = []
async def __call__(self, websocket: object, model: str) -> UserAPIKeyAuth:
self.authorized.append(model)
if self.denial is not None:
raise self.denial
return self.token
SOCKET: Final = WebSocket({"type": "websocket"}, receive=AsyncMock(), send=AsyncMock())
async def _gate_verdict(authorizer: _FrameAuthorizer, model: str) -> WebsocketFrameModelVerdict:
with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the model checks need no database
return await _proxy_frame_model_gate(authorizer)(SOCKET, model)
def _refusal(verdict: WebsocketFrameModelVerdict) -> WebsocketFrameModelRefused:
assert isinstance(verdict, WebsocketFrameModelRefused)
return verdict
@pytest.mark.asyncio
async def test_proxy_frame_model_gate_runs_full_auth_for_the_model_the_frame_names():
authorizer: Final = _FrameAuthorizer(UserAPIKeyAuth(models=["gpt-realtime-2.1", "gpt-5.4-mini"]))
verdict: Final = await _gate_verdict(authorizer, "gpt-realtime-2.1")
assert verdict == WebsocketFrameModelAllowed(access_groups=())
assert authorizer.authorized == ["gpt-realtime-2.1"]
@pytest.mark.asyncio
@pytest.mark.parametrize("requested_model", ["gpt-5.4", "gpt-realtime-2.1-mini"])
async def test_proxy_frame_model_gate_refuses_a_model_the_key_lacks(requested_model):
authorizer: Final = _FrameAuthorizer(UserAPIKeyAuth(models=["gpt-realtime-2.1", "gpt-5.4-mini"]))
refusal: Final = _refusal(await _gate_verdict(authorizer, requested_model))
assert requested_model in refusal.message
assert refusal.error_type == ProxyErrorTypes.key_model_access_denied.value
@pytest.mark.asyncio
@pytest.mark.parametrize(
"granted_models, refusal_names",
[(["gpt-5.4-mini"], "gpt-5.4"), ([], "no model grants")],
ids=["model_outside_the_grant", "no_grants"],
)
async def test_proxy_frame_model_gate_refuses_a_model_the_agent_was_not_granted(granted_models, refusal_names):
from litellm.types.agents import AgentResponse
token: Final = UserAPIKeyAuth(models=[], agent_id="agent-fake")
token.managed_agent_policy = AgentResponse(
agent_id="agent-fake",
agent_name="Agent",
agent_card_params={},
object_permission={"object_permission_id": "grant-fake", "models": granted_models},
)
refusal: Final = _refusal(await _gate_verdict(_FrameAuthorizer(token), "gpt-5.4"))
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):
verdict: Final = await _gate_verdict(_FrameAuthorizer(UserAPIKeyAuth(models=unrestricted_models)), "gpt-5.4")
assert verdict == WebsocketFrameModelAllowed(access_groups=())
@pytest.mark.asyncio
async def test_proxy_frame_model_gate_hands_back_the_access_groups_that_authorized_the_model():
"""The turn a frame starts is billed to the groups that let its model through, so the gate has to report them."""
token: Final = UserAPIKeyAuth(models=["gpt-5.4-mini"], matched_model_access_groups=["tier-a"])
verdict: Final = await _gate_verdict(_FrameAuthorizer(token), "gpt-5.4-mini")
assert verdict == WebsocketFrameModelAllowed(access_groups=("tier-a",))
@pytest.mark.asyncio
@pytest.mark.parametrize(
"requested_model",
["gpt-5.4-mini", "gpt-5.4"],
ids=["allowed", "refused_by_the_model_check"],
)
async def test_proxy_frame_model_gate_releases_the_budget_it_reserved_for_the_check(requested_model):
"""
Full auth reserves budget for the request it authorizes, but a frame check is
not a request, so a reservation the gate leaves open would hold the key's
spend counter above its real spend and refuse later calls on every route.
"""
reservation: Final[dict[str, object]] = {"reserved_cost": 0.5, "entries": []}
token: Final = UserAPIKeyAuth(models=["gpt-5.4-mini"], budget_reservation=reservation)
await _gate_verdict(_FrameAuthorizer(token), requested_model)
assert token.budget_reservation == {**reservation, "finalized": True}
@pytest.mark.asyncio
async def test_proxy_frame_model_gate_does_not_count_the_sessions_own_reservation_as_key_spend():
"""
The upgrade reserves up to the key's remaining budget for the session, and the key budget check
counts open reservations as spend, so a frame checked while that reservation is open is refused
on a key with budget left, where an HTTP call with the same key goes through.
"""
connect_reservation: Final[dict[str, object]] = {"reserved_cost": 0.0001, "entries": []}
socket: Final = WebSocket(
{"type": "websocket", "state": {"budget_reservation": connect_reservation}}, receive=AsyncMock(), send=AsyncMock()
)
async def key_budget_check(websocket: object, model: str) -> UserAPIKeyAuth:
if connect_reservation.get("finalized") is not True:
raise litellm.BudgetExceededError(current_cost=0.0001, max_budget=0.0001)
return UserAPIKeyAuth(models=["gpt-5.4-mini"])
with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the model checks need no database
verdict: Final = await _proxy_frame_model_gate(key_budget_check)(socket, "gpt-5.4-mini")
assert verdict == WebsocketFrameModelAllowed(access_groups=())
@pytest.mark.asyncio
@pytest.mark.parametrize(
"denial, error_type, message",
[
(
ProxyException(
message="Budget has been exceeded! Key=hashed Spend=2.0, Budget=1.0",
type=ProxyErrorTypes.budget_exceeded,
param=None,
code=400,
),
"budget_exceeded",
"Budget has been exceeded! Key=hashed",
),
(
litellm.BudgetExceededError(
current_cost=2.0,
max_budget=1.0,
message="Budget has been exceeded! Model access group=tier-a Current cost: 2.0, Max budget: 1.0",
entity_type=Litellm_EntityType.MODEL_ACCESS_GROUP.value,
entity_id="tier-a",
),
"budget_exceeded",
"Model access group=tier-a",
),
(
ProxyException(
message="team not allowed to access model. This team can only access models=['gpt-5.4-mini']",
type=ProxyErrorTypes.team_model_access_denied,
param="model",
code=401,
),
"team_model_access_denied",
"team not allowed to access model",
),
(HTTPException(status_code=403, detail="Invalid Authorization header format"), "invalid_request_error", ""),
],
ids=["key_budget", "group_budget", "team_allowlist", "malformed_auth"],
)
async def test_proxy_frame_model_gate_refuses_what_full_auth_denies_with_its_reason(denial, error_type, message):
"""A frame is denied for the same reasons a request is, and the caller reads which reason in the error event."""
refusal: Final = _refusal(await _gate_verdict(_FrameAuthorizer(denial=denial), "gpt-5.4-mini"))
assert refusal.error_type == error_type
assert message in refusal.message
@pytest.mark.asyncio
async def test_proxy_frame_model_gate_refuses_out_loud_when_auth_cannot_answer():
"""
A gate that cannot answer has to close the session with a reason the caller can read.
Letting the failure escape drops the socket with no frame at all, which the caller
cannot tell apart from the network going away.
"""
refusal: Final = _refusal(
await _gate_verdict(_FrameAuthorizer(denial=ConnectionError("connection pool is exhausted")), "gpt-5.4-mini")
)
assert refusal == WebsocketFrameModelRefused(message=_OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL, error_type="server_error")