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

Merge branch 'litellm_openai_ws_frame_model_gate_rescue' of https://github.com/BerriAI/litellm into litellm_openai_ws_frame_model_gate

Browse source
This commit is contained in:
mateo-berri 2026-10-28 14:29:30 -07:00
commit 99123c6426
11 changed files with 1211 additions and 480 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -5305,79 +5305,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,

136
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
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
@ -50,6 +50,7 @@ from litellm.constants import (
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
from litellm.integrations.custom_guardrail import guardrail_request_data_with_streaming, is_guardrail_intervention
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.llms.anthropic.common_utils import AnthropicModelInfo, merge_anthropic_beta_headers
from litellm.llms.azure.passthrough.transformation import (
foreign_azure_deployment,
@ -72,7 +73,7 @@ from litellm.llms.oss_decision import OssDecisionProvider, oss_connection, valid
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error
from litellm.proxy._types import *
from litellm.proxy.auth.auth_checks import enforced_model_allowlists
from litellm.proxy.auth.auth_checks import can_agent_call_model, can_key_call_resolved_model
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
@ -80,6 +81,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,
)
@ -115,8 +117,16 @@ 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 (
HttpPassThroughEndpointHelpers,
WebsocketFrameModelAllowed,
WebsocketFrameModelGate,
WebsocketFrameModelRefused,
WebsocketFrameModelVerdict,
WebsocketUpstreamFrameClassifier,
create_pass_through_route,
create_websocket_passthrough_route,
websocket_passthrough_request,
@ -180,7 +190,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
@ -210,9 +220,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:
@ -3274,19 +3283,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
@ -3303,6 +3299,11 @@ 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_DISABLED_REFUSAL: Final = _OpenAIWebsocketRefusal(
close_reason="OpenAI websocket passthrough is disabled",
message=(
@ -3311,14 +3312,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")
@ -3327,19 +3320,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
@ -3354,6 +3337,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: ...
@ -3361,18 +3346,62 @@ 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,
)
return resolve
FrameModelAuthorizer: TypeAlias = Callable[[WebSocket, str], Awaitable[UserAPIKeyAuth]]
def _frame_model_authorizer() -> FrameModelAuthorizer:
return user_api_key_auth_accepted_websocket_for_model
def _proxy_frame_model_gate(
authorize_frame_model: Annotated[FrameModelAuthorizer, Depends(_frame_model_authorizer)],
) -> WebsocketFrameModelGate:
from litellm.proxy.proxy_server import (
llm_model_list,
llm_router,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
from litellm.proxy.spend_tracking.budget_reservation import release_or_invalidate_budget_reservation
async def authorize(websocket: WebSocket, model: str) -> WebsocketFrameModelAllowed:
frame_auth: Final = await authorize_frame_model(websocket, model)
access_groups: Final = tuple(frame_auth.matched_model_access_groups or ())
try:
await can_key_call_resolved_model(
model=model, llm_model_list=llm_model_list, valid_token=frame_auth, llm_router=llm_router
)
await can_agent_call_model(
model=model,
valid_token=frame_auth,
llm_router=llm_router,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
finally:
await release_or_invalidate_budget_reservation(budget_reservation=frame_auth.budget_reservation)
return WebsocketFrameModelAllowed(access_groups=access_groups)
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:
@ -3392,12 +3421,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] = {
@ -3451,6 +3480,8 @@ async def openai_websocket_proxy_route(
forward_headers=False,
endpoint=websocket.url.path,
accept_websocket=False,
client_frame_model_gate=frame_model_gate,
upstream_frame_kind=openai_websocket_frame_kind,
)
@ -3696,7 +3727,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:

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

@ -0,0 +1,218 @@
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.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, ResponseAPIUsage
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"
_RESPONSE_CREATED: Final = "response.created"
_REALTIME_RESPONSE_DONE: Final = "response.done"
_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, _RESPONSE_CREATED}
_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 _ModelNamingEvent(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
session: _NamedModel | None = None
response: _NamedModel | None = None
class _EndedResponse(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
model: str | None = None
usage: ResponseAPIUsage | None = None
class _ResponsesTurnEndEvent(BaseModel):
model_config = ConfigDict(extra="ignore", frozen=True)
type: Literal["response.completed", "response.incomplete", "response.failed"]
response: _EndedResponse
_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
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
return next(
(
container.model
for container in (naming.session, naming.response)
if container is not None and container.model
),
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"
return "relayed"
def _realtime_billing(messages: WebsocketMessages, model: str, logging_obj: LiteLLMLoggingObj) -> _Billing | None:
if not any(event.get("type") == _REALTIME_RESPONSE_DONE 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,
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) -> _Billing | None:
turns: Final = tuple(turn for turn in (_responses_turn(event, model) for event in messages) if turn is not None)
if not turns:
return None
usage: Final = RealtimeAPITokenUsageProcessor.combine_usage_objects([turn.usage for turn in turns])
return _Billing(usage=usage, cost=sum(_responses_cost(turn.usage, turn.model) for turn in turns))
def _responses_cost(usage: Usage, model: str) -> float:
try:
return litellm.completion_cost(
completion_response=ModelResponse(model=model, usage=usage),
model=model,
custom_llm_provider=OPENAI_WEBSOCKET_PROVIDER,
)
except Exception: # noqa: BLE001 # completion_cost raises a bare Exception for an unpriced model, and the turn must still log
verbose_proxy_logger.warning(
"OpenAI websocket passthrough: could not price a Responses turn on %s, logging it at zero", model
)
return 0.0
def _priced_result(model: str, billing: _Billing, start_time: datetime) -> ModelResponse:
result: Final = ModelResponse(
id=f"openai-websocket-{start_time.timestamp()}",
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
)
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),
"kwargs": {
**kwargs,
"model": model,
"custom_llm_provider": OPENAI_WEBSOCKET_PROVIDER,
"response_cost": billing.cost,
},
}
return priced

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

@ -3,14 +3,16 @@ import asyncio
import copy
import json
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 typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, TypedDict, cast
from urllib.parse import urlencode, urlparse
import httpx
@ -23,6 +25,7 @@ from fastapi import (
Response,
UploadFile,
WebSocket,
WebSocketDisconnect,
status,
)
from fastapi.responses import StreamingResponse
@ -30,7 +33,7 @@ from pydantic import TypeAdapter
from starlette.datastructures import UploadFile as StarletteUploadFile
from starlette.routing import BaseRoute, Route
from starlette.websockets import WebSocketState
from websockets.asyncio.client import connect
from websockets.asyncio.client import ClientConnection, connect
from websockets.exceptions import (
ConnectionClosedError,
ConnectionClosedOK,
@ -56,6 +59,7 @@ from litellm.integrations.custom_guardrail import (
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.otel.model.semconv import PASSTHROUGH_OPERATION_KEY, GenAIOperation
from litellm.litellm_core_utils.core_helpers import (
as_str_mapping,
bind_budget_reservation_to_callbacks,
get_metadata_variable_name_from_kwargs,
get_or_create_metadata_bucket,
@ -65,6 +69,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
@ -141,6 +146,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
EndpointType,
PassthroughStandardLoggingPayload,
WebsocketUpstreamFrameKind,
)
from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, Usage
@ -2486,6 +2492,178 @@ 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 _ForwardClientFrame:
turn_access_groups: tuple[str, ...] | None
@dataclass(frozen=True, slots=True)
class _RefuseClientFrame:
refusal: WebsocketFrameModelRefused
@dataclass(frozen=True, slots=True)
class _ClientFrameRequests:
turn_model: str | None
models: tuple[str, ...]
names_stored_prompt: bool
_NO_CLIENT_FRAME_REQUESTS: Final = _ClientFrameRequests(turn_model=None, models=(), names_stored_prompt=False)
_FRAME_TEXT_WORTH_PARSING: Final = ('"model"', '"prompt"', "\\")
_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"))
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."
)
_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 _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 _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)
prompts: Final = (_placed_value(frame, placement) for placement in _STORED_PROMPT_PLACEMENTS)
return _ClientFrameRequests(
turn_model=next(iter(turn_models), None),
models=tuple(dict.fromkeys((*turn_models, *_placed_models(frame, _TRANSCRIPTION_MODEL_PLACEMENTS)))),
names_stored_prompt=any(isinstance(prompt, Mapping) for prompt in prompts),
)
def _stored_prompt_refusal(probe: WebsocketFrameModelRefused) -> WebsocketFrameModelRefused:
if probe.error_type.endswith("_model_access_denied"):
return WebsocketFrameModelRefused(message=_STORED_PROMPT_REFUSAL, 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 _ForwardClientFrame(turn_access_groups=None)
requests: Final = _client_frame_requests(frame_data)
verdicts: Final = tuple([await gate(websocket, model) for model in requests.models])
refusal: Final = next((verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelRefused)), None)
if refusal is not None:
return _RefuseClientFrame(refusal=refusal)
if requests.names_stored_prompt and isinstance(
probe := await gate(websocket, STORED_PROMPT_PROBE_MODEL), WebsocketFrameModelRefused
):
return _RefuseClientFrame(refusal=_stored_prompt_refusal(probe))
allowed: Final = tuple(verdict for verdict in verdicts if isinstance(verdict, WebsocketFrameModelAllowed))
if requests.turn_model is None or not allowed:
return _ForwardClientFrame(turn_access_groups=None)
return _ForwardClientFrame(turn_access_groups=allowed[0].access_groups)
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_client_frame(
websocket: WebSocket, upstream_ws: ClientConnection, 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))
await upstream_ws.close()
def _stamp_billed_access_groups(kwargs: Mapping[str, object], access_groups: tuple[str, ...]) -> None:
"""Charge a turn to the access groups that authorized its model, and to none when no group did."""
litellm_params: Final = as_str_mapping(kwargs.get("litellm_params"))
metadata: Final = None if litellm_params is None else litellm_params.get("metadata")
if isinstance(metadata, MutableMapping):
metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = list(access_groups)
async def _relay_client_frame(
websocket: WebSocket,
upstream_ws: ClientConnection,
checked_frame: str | bytes,
forwarded_frame: str | bytes,
gate: WebsocketFrameModelGate | None,
) -> _ForwardClientFrame | _RefuseClientFrame:
"""Forward one client frame upstream, or refuse it and end the session."""
verdict: Final = await _client_frame_verdict(websocket, checked_frame, gate)
match verdict:
case _RefuseClientFrame(refusal=refusal):
await _refuse_client_frame(websocket, upstream_ws, refusal)
case _ForwardClientFrame():
await upstream_ws.send(forwarded_frame)
return verdict
_EXPLICIT_SUCCESS_KWARGS: Final = frozenset({"logging_obj", "passthrough_logging_payload", "custom_llm_provider"})
class _WebsocketTurnLedger:
"""The groups that authorized the turn in flight, when it started, and how many turns the socket has logged."""
def __init__(self, access_groups: tuple[str, ...], started: datetime) -> None:
self.access_groups = access_groups
self.started = started
self.logged_turns = 0
def authorize(self, access_groups: tuple[str, ...] | None) -> None:
if access_groups is not None:
self.access_groups = access_groups
def close_turn(self, end_time: datetime) -> None:
self.logged_turns += 1
self.started = end_time
async def websocket_passthrough_request(
websocket: WebSocket,
target: str,
@ -2496,6 +2674,8 @@ 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,
):
"""
WebSocket passthrough request handler.
@ -2509,6 +2689,9 @@ 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
"""
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.proxy.proxy_server import proxy_config, proxy_logging_obj
@ -2544,26 +2727,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
@ -2585,35 +2773,117 @@ 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(
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,
) -> None:
frames: Final = list(websocket_messages)
websocket_messages[:] = (
()
if upstream_frame_kind is None
else (frame for frame in frames if 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_access_groups(call_kwargs, ledger.access_groups)
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,
**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() -> None:
end_time: Final = datetime.now()
if ledger.logged_turns == 0:
log_frames(logging_obj, passthrough_logging_payload, kwargs, end_time)
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.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)
ledger.close_turn(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)
@ -2705,9 +2975,27 @@ async def websocket_passthrough_request(
)
# Not a JSON message or doesn't contain setup data
await upstream_ws.send(_rewrite_vertex_live_setup_model(text_data, setup_model_rewriter))
if isinstance(
relayed_text := await _relay_client_frame(
websocket,
upstream_ws,
text_data,
_rewrite_vertex_live_setup_model(text_data, setup_model_rewriter),
client_frame_model_gate,
),
_RefuseClientFrame,
):
break
ledger.authorize(relayed_text.turn_access_groups)
elif bytes_data is not None:
await upstream_ws.send(bytes_data)
if isinstance(
relayed_bytes := await _relay_client_frame(
websocket, upstream_ws, bytes_data, bytes_data, client_frame_model_gate
),
_RefuseClientFrame,
):
break
ledger.authorize(relayed_bytes.turn_access_groups)
except asyncio.CancelledError:
raise
except Exception:
@ -2734,11 +3022,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
@ -2746,6 +3030,17 @@ async def websocket_passthrough_request(
_extract_vertex_live_model_from_setup_response(message_data)
return
websocket_messages.append(message_data)
if upstream_frame_kind is not None and upstream_frame_kind(message_data) == "ends_turn":
log_turn()
async def relay_upstream_frame(upstream_message: str | bytes) -> None:
record_upstream_frame(upstream_message)
if not _client_socket_is_open(websocket):
return
if isinstance(upstream_message, bytes):
await websocket.send_bytes(upstream_message)
else:
await websocket.send_text(upstream_message)
async def forward_upstream_to_client() -> Close | None:
try:
@ -2754,6 +3049,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
@ -2793,62 +3093,8 @@ 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)
# 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,
)
)
bind_budget_reservation_to_callbacks(logging_obj.litellm_params)
if ledger.logged_turns == 0:
log_turn()
# Call the proxy logging success hook
if proxy_logging_obj:

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

@ -1,4 +1,5 @@
import json
from collections.abc import Mapping, Sequence
from datetime import datetime
from types import MappingProxyType
from typing import Any, Final
@ -8,7 +9,7 @@ import httpx
from litellm.constants import AZURE_SPEECH_CUSTOM_LLM_PROVIDER
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,
)
@ -64,6 +65,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,
@ -114,6 +123,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:
@ -385,21 +395,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({}),
@ -415,6 +422,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,
@ -584,6 +621,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

@ -422,7 +422,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,8 +1,10 @@
from collections.abc import Mapping, Sequence
from datetime import datetime
from enum import Enum
from typing import Final
from typing import Final, Literal, TypeAlias
from pydantic import Field
from typing_extensions import TypedDict
from typing_extensions import ReadOnly, TypedDict
from litellm.types.llms.base import LiteLLMPydanticObjectBase
@ -27,6 +29,8 @@ LITELLM_PASS_THROUGH_ENDPOINT_MARKER: Final = "__litellm_pass_through_endpoint__
SETTLES_BACKGROUND_RESPONSE_KWARG: Final = "settles_background_response"
WebsocketUpstreamFrameKind: TypeAlias = Literal["ends_turn", "opens_session", "relayed"]
class BackgroundResponseSettlement(str, Enum):
SETTLES = "settles"
@ -62,7 +66,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
"""
@ -74,6 +78,8 @@ class PassthroughStandardLoggingPayload(TypedDict, total=False):
Optional field, we use this for cost tracking only if it's set.
"""
end_time: ReadOnly[datetime]
class PassThroughGuardrailSettings(LiteLLMPydanticObjectBase):
"""

63
tests/unit/proxy/auth/test_model_access_group_budgets.py
View file

@ -27,7 +27,6 @@ from litellm.proxy.auth.auth_checks import (
_model_access_group_max_budget_check,
collect_matched_model_access_groups,
common_checks,
enforce_model_access_group_budget,
stamp_matched_model_access_groups,
)
from litellm.proxy.common_utils.reset_budget_job import _model_access_group_counter_key
@ -552,65 +551,3 @@ async def test_common_checks_blocks_a_request_whose_group_is_over_budget():
async def test_free_model_routes_skip_the_model_access_group_budget_check():
"""skip_budget_checks is how free models stay free; it has to cover this budget too."""
assert await _common_checks_with_over_budget_group(skip_budget_checks=True) is True
async def _enforce_for_a_route_without_a_request_body(
*rows: _MagBudgetRow,
key_models: list[str],
spend_by_counter_key: dict[str, float],
) -> UserAPIKeyAuth:
"""Run the pair the websocket frame gate runs, which reaches the proxy for what it needs."""
valid_token = UserAPIKeyAuth(api_key="hashed", models=key_models, user_id=USER_ID)
cache = await _cache()
read, _ = _spend_reader(spend_by_counter_key)
with (
patch("litellm.proxy.proxy_server.prisma_client", _RecordingPrismaClient(*rows)), # test-quality-ok: read lazily inside the helper, which takes no client
patch("litellm.proxy.proxy_server.user_api_key_cache", cache), # test-quality-ok: same lazy import
patch("litellm.proxy.proxy_server.proxy_logging_obj", ProxyLogging(user_api_key_cache=cache)), # test-quality-ok: same lazy import
patch("litellm.proxy.proxy_server.get_current_spend", read), # test-quality-ok: get_current_spend is lazily imported inside the budget check and has no injection point
):
await enforce_model_access_group_budget(
model="gpt-4o",
valid_token=valid_token,
llm_router=Router(model_list=MODEL_LIST),
)
return valid_token
@pytest.mark.asyncio
async def test_caller_granted_the_model_only_through_an_exhausted_group_is_refused():
"""Routes with no request body, the websocket passthrough among them, run this instead of
``common_checks``, so the group budget has to stop them the same way."""
with pytest.raises(litellm.BudgetExceededError) as exc_info:
await _enforce_for_a_route_without_a_request_body(
_MagBudgetRow("tier-a", max_budget=1.0),
key_models=["tier-a"],
spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 1.5},
)
assert exc_info.value.entity_id == "tier-a"
assert exc_info.value.current_cost == 1.5
@pytest.mark.asyncio
async def test_group_that_authorized_the_call_is_stamped_for_the_spend_writer():
"""Without the stamp the call runs on the group's budget and never pays into it, so the pool
can never catch up with what it authorized."""
valid_token = await _enforce_for_a_route_without_a_request_body(
_MagBudgetRow("tier-a", max_budget=10.0),
key_models=["tier-a"],
spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 1.0},
)
assert valid_token.matched_model_access_groups == ["tier-a"]
@pytest.mark.asyncio
async def test_model_granted_directly_charges_no_group():
valid_token = await _enforce_for_a_route_without_a_request_body(
_MagBudgetRow("tier-a", max_budget=1.0),
key_models=["gpt-4o"],
spend_by_counter_key={MODEL_ACCESS_GROUP_COUNTER_KEY: 1.5},
)
assert valid_token.matched_model_access_groups is None

51
tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py
View file

@ -8,6 +8,7 @@ import pytest
import litellm
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket_passthrough_logging_handler import (
OpenAIWebsocketPassthroughLoggingHandler,
openai_websocket_frame_kind,
)
REALTIME_MODEL: Final = "gpt-realtime-2.1"
@ -42,9 +43,11 @@ def _realtime_response_done(text_in: int, audio_in: int, text_out: int, audio_ou
}
def _responses_completed(model: str, input_tokens: int, output_tokens: int) -> dict:
def _responses_completed(
model: str, input_tokens: int, output_tokens: int, event: str = "response.completed"
) -> dict:
return {
"type": "response.completed",
"type": event,
"response": {
"id": "resp_fake",
"model": model,
@ -288,3 +291,47 @@ def test_a_turn_the_caller_walked_out_of_still_names_the_model_it_ran():
assert handled["result"] is None
assert handled["kwargs"] == {"model": REASONING_MODEL, "custom_llm_provider": "openai"}
assert logging_obj.model_call_details["model"] == REASONING_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 = _handle(
[_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5, event=event)],
url_route="/openai_passthrough/v1/responses",
)
prices = litellm.model_cost[RESPONSES_MODEL]
expected_cost = 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 = _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(
("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", "relayed"),
("response.output_text.delta", "relayed"),
("error", "relayed"),
],
)
def test_every_frame_that_settles_a_turn_closes_it_and_nothing_else_does(frame_type, kind):
assert openai_websocket_frame_kind({"type": frame_type}) == kind

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

@ -42,9 +42,13 @@ from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import route
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS,
LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY,
STORED_PROMPT_PROBE_MODEL,
HttpPassThroughEndpointHelpers,
InitPassThroughEndpointHelpers,
SafeRouteAdder,
WebsocketFrameModelAllowed,
WebsocketFrameModelRefused,
WebsocketFrameModelVerdict,
_registered_pass_through_routes,
_truncate_upstream_error_body,
_update_metadata_with_tags_in_header,
@ -58,6 +62,9 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
set_env_variables_in_header,
websocket_passthrough_request,
)
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.success_handler import (
PassThroughEndpointLogging,
)
@ -7309,7 +7316,7 @@ def _patched_websocket_passthrough_environment(upstream_ws):
mock_worker.ensure_initialized_and_enqueue = MagicMock(
side_effect=lambda async_coroutine: async_coroutine.close()
)
yield
yield mock_proxy_logging
async def _pending_receive():
@ -7620,14 +7627,28 @@ SUCCESS_HANDLER: Final = (
)
KEY_MODEL_REFUSAL: Final = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']."
KEY_MODEL_CLOSE_REASON: Final = "key not allowed to access model."
class RecordingFrameModelGate:
def __init__(self, refusals: Mapping[str, str] = MappingProxyType({})):
def __init__(
self,
refusals: Mapping[str, str] = MappingProxyType({}),
groups: Mapping[str, tuple[str, ...]] = MappingProxyType({}),
error_type: str = "key_model_access_denied",
):
self.refusals = refusals
self.groups = groups
self.error_type = error_type
self.checked: list[str] = []
async def __call__(self, model: str, valid_token: UserAPIKeyAuth, /) -> Optional[str]:
async def __call__(self, websocket: object, model: str, /) -> WebsocketFrameModelVerdict:
self.checked.append(model)
return self.refusals.get(model)
refusal = self.refusals.get(model)
if refusal is not None:
return WebsocketFrameModelRefused(message=refusal, error_type=self.error_type)
return WebsocketFrameModelAllowed(access_groups=self.groups.get(model, ()))
def _received_frame(frame):
@ -7692,8 +7713,7 @@ async def test_websocket_passthrough_checks_the_model_a_client_frame_asks_for(fr
@pytest.mark.asyncio
async def test_websocket_passthrough_refuses_a_frame_naming_a_model_the_key_cannot_call():
refusal = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']."
gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": refusal}))
gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL}))
frame = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}})
relayed = await _relay_client_frames([frame], gate)
@ -7701,16 +7721,34 @@ async def test_websocket_passthrough_refuses_a_frame_naming_a_model_the_key_cann
assert relayed.upstream.send.await_args_list == []
assert json.loads(relayed.client.send_text.await_args.args[0]) == {
"type": "error",
"error": {"type": "invalid_request_error", "message": refusal},
"error": {"type": "key_model_access_denied", "message": KEY_MODEL_REFUSAL},
}
assert relayed.client.close.await_args_list == [call(code=1008, reason=refusal)]
assert relayed.client.close.await_args_list == [call(code=1008, reason=KEY_MODEL_CLOSE_REASON)]
relayed.upstream.close.assert_awaited()
@pytest.mark.asyncio
async def test_websocket_passthrough_closes_with_the_whole_first_sentence_of_the_refusal():
"""
The close frame only has room for 123 bytes, so a refusal cut at that limit
reads as a truncated number. The first sentence names the reason whole, and
the error event before the close still carries every word with its type.
"""
refusal = "Budget has been exceeded! Key=" + "f" * 64 + " Spend=1.23456789, Budget=1.0"
gate = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": refusal}), error_type="budget_exceeded")
relayed = await _relay_client_frames([json.dumps({"type": "response.create", "model": "gpt-5.4"})], gate)
assert relayed.client.close.await_args_list == [call(code=1008, reason="Budget has been exceeded!")]
assert json.loads(relayed.client.send_text.await_args.args[0])["error"] == {
"type": "budget_exceeded",
"message": refusal,
}
@pytest.mark.asyncio
async def test_websocket_passthrough_refuses_the_whole_frame_when_any_named_model_is_denied():
refusal = "key not allowed to access model. This key can only access models=['gpt-realtime-2.1']."
gate = RecordingFrameModelGate(MappingProxyType({"gpt-4o-transcribe": refusal}))
gate = RecordingFrameModelGate(MappingProxyType({"gpt-4o-transcribe": KEY_MODEL_REFUSAL}))
frame = json.dumps(
{
"type": "session.update",
@ -7722,13 +7760,12 @@ async def test_websocket_passthrough_refuses_the_whole_frame_when_any_named_mode
assert gate.checked == ["gpt-realtime-2.1", "gpt-4o-transcribe"]
assert relayed.upstream.send.await_args_list == []
assert relayed.client.close.await_args_list == [call(code=1008, reason=refusal)]
assert relayed.client.close.await_args_list == [call(code=1008, reason=KEY_MODEL_CLOSE_REASON)]
@pytest.mark.asyncio
async def test_websocket_passthrough_stops_relaying_after_a_refused_frame():
refusal = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']."
gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": refusal}))
gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL}))
frames = [
json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}),
json.dumps({"type": "response.create"}),
@ -7795,27 +7832,79 @@ async def test_websocket_passthrough_checks_a_model_named_with_escaped_json(fram
"\\u006dodel" asks for a model just as plainly as one carrying "model" and
has to face the same check.
"""
refusal = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']."
gate = RecordingFrameModelGate(MappingProxyType({expected_model: refusal}))
gate = RecordingFrameModelGate(MappingProxyType({expected_model: KEY_MODEL_REFUSAL}))
relayed = await _relay_client_frames([frame], gate)
assert gate.checked == [expected_model]
assert relayed.upstream.send.await_args_list == []
assert relayed.client.close.await_args_list == [call(code=1008, reason=refusal)]
assert relayed.client.close.await_args_list == [call(code=1008, reason=KEY_MODEL_CLOSE_REASON)]
@pytest.mark.asyncio
async def test_websocket_passthrough_checks_the_model_a_binary_frame_asks_for():
refusal = "key not allowed to access model. This key can only access models=['gpt-5.4-mini']."
gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": refusal}))
gate = RecordingFrameModelGate(MappingProxyType({"gpt-realtime-2.1": KEY_MODEL_REFUSAL}))
frame = json.dumps({"type": "session.update", "session": {"model": "gpt-realtime-2.1"}}).encode()
relayed = await _relay_client_frames([frame], gate)
assert gate.checked == ["gpt-realtime-2.1"]
assert relayed.upstream.send.await_args_list == []
assert relayed.client.close.await_args_list == [call(code=1008, reason=refusal)]
assert relayed.client.close.await_args_list == [call(code=1008, reason=KEY_MODEL_CLOSE_REASON)]
STORED_PROMPT_FRAMES: Final = (
{"type": "response.create", "prompt": {"id": "pmpt_fake", "version": "2"}},
{"type": "response.create", "response": {"prompt": {"id": "pmpt_fake"}}},
{"type": "session.update", "session": {"prompt": {"id": "pmpt_fake"}}},
)
@pytest.mark.asyncio
@pytest.mark.parametrize("frame", STORED_PROMPT_FRAMES)
async def test_websocket_passthrough_refuses_a_stored_prompt_for_a_key_limited_to_some_models(frame):
"""
A stored prompt carries its own model on the provider side, so a key held to
a model list could reach any model through one. The gateway cannot see that
model, so a key that is not allowed every model cannot run a stored prompt.
"""
gate = RecordingFrameModelGate(MappingProxyType({STORED_PROMPT_PROBE_MODEL: KEY_MODEL_REFUSAL}))
relayed = await _relay_client_frames([json.dumps(frame)], gate)
error_message: Final = json.loads(relayed.client.send_text.await_args.args[0])["error"]["message"]
close_reason: Final = relayed.client.close.await_args.kwargs["reason"]
assert gate.checked == [STORED_PROMPT_PROBE_MODEL]
assert relayed.upstream.send.await_args_list == []
assert relayed.client.close.await_args.kwargs["code"] == 1008
assert "stored prompt" in close_reason
assert close_reason.endswith(".") and error_message.startswith(close_reason)
@pytest.mark.asyncio
@pytest.mark.parametrize("frame", STORED_PROMPT_FRAMES)
async def test_websocket_passthrough_forwards_a_stored_prompt_for_a_key_allowed_every_model(frame):
gate = RecordingFrameModelGate()
relayed = await _relay_client_frames([json.dumps(frame)], gate)
assert [call_.args[0] for call_ in relayed.upstream.send.await_args_list] == [json.dumps(frame)]
assert relayed.client.close.await_args_list == [call()]
@pytest.mark.asyncio
async def test_websocket_passthrough_keeps_the_budget_reason_when_a_stored_prompt_is_refused_for_spend():
refusal = "Budget has been exceeded! Key=hashed Spend=2.0, Budget=1.0"
gate = RecordingFrameModelGate(
MappingProxyType({STORED_PROMPT_PROBE_MODEL: refusal}), error_type="budget_exceeded"
)
relayed = await _relay_client_frames([json.dumps(STORED_PROMPT_FRAMES[0])], gate)
assert json.loads(relayed.client.send_text.await_args.args[0])["error"] == {
"type": "budget_exceeded",
"message": refusal,
}
@pytest.mark.asyncio
@ -7889,6 +7978,196 @@ async def test_websocket_passthrough_still_bills_a_session_the_client_walked_out
]
class ScriptedUpstreamWebSocket:
"""Answers each client frame with its scripted reply, and flushes what was still in flight when it is closed"""
def __init__(self, *opening: str, replies: Sequence[Sequence[str]] = (), in_flight: Sequence[str] = ()):
self._frames: asyncio.Queue[str | None] = asyncio.Queue()
for frame in opening:
self._frames.put_nowait(frame)
self._replies = iter(replies)
self._in_flight = in_flight
self._closed = False
self.send = AsyncMock(side_effect=self._reply)
self.close = AsyncMock(side_effect=self._flush_and_close)
async def _reply(self, frame: str | bytes) -> None:
for reply in next(self._replies, ()):
self._frames.put_nowait(reply)
async def _flush_and_close(self, *args: object, **kwargs: object) -> None:
if self._closed:
return
self._closed = True
for frame in self._in_flight:
self._frames.put_nowait(frame)
self._frames.put_nowait(None)
while not self._frames.empty():
await asyncio.sleep(0)
async def recv(self, decode: bool | None = None) -> str:
from websockets.exceptions import ConnectionClosedOK
from websockets.frames import Close
frame = await self._frames.get()
if frame is None:
raise ConnectionClosedOK(rcvd=Close(1000, ""), sent=Close(1000, ""), rcvd_then_sent=True)
return frame
def _turn_taking_client(frames: Sequence[str], answered_by: str) -> MagicMock:
"""A client that sends its next frame only once the answer to the last one reached it, and that Starlette
refuses to send to once it is closed"""
from starlette.websockets import WebSocketState
delivered: list[str] = []
script = iter(frames)
receives = iter(range(len(frames) + 1))
async def receive():
answers_due = next(receives)
while sum(answered_by in text for text in delivered) < answers_due:
await asyncio.sleep(0)
frame = next(script, None)
return {"type": "websocket.disconnect"} if frame is None else _received_frame(frame)
websocket = _client_websocket(receive)
async def send_text(text: str) -> None:
if websocket.application_state == WebSocketState.DISCONNECTED:
raise RuntimeError('Cannot call "send" once a close message has been sent.')
delivered.append(text)
websocket.send_text = AsyncMock(side_effect=send_text)
return websocket
async def _relay_turns(client, upstream_ws, gate, endpoint="/openai_passthrough/v1/responses"):
with (
_patched_websocket_passthrough_environment(upstream_ws) as proxy_logging,
patch(SUCCESS_HANDLER, new_callable=AsyncMock) as success_handler,
):
await asyncio.wait_for(
websocket_passthrough_request(
websocket=client,
target="wss://api.openai.com/v1/responses",
custom_headers={"Authorization": "Bearer sk-provider"},
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
forward_headers=False,
endpoint=endpoint,
accept_websocket=False,
client_frame_model_gate=gate,
upstream_frame_kind=openai_websocket_frame_kind,
),
timeout=10,
)
return SimpleNamespace(success_handler=success_handler, proxy_logging=proxy_logging)
def _responses_turn(model: str) -> list[str]:
return [
json.dumps({"type": "response.created", "response": {"model": model}}),
json.dumps(
{
"type": "response.completed",
"response": {"model": model, "usage": {"input_tokens": 9, "output_tokens": 3, "total_tokens": 12}},
}
),
]
@pytest.mark.asyncio
async def test_websocket_passthrough_logs_each_turn_on_its_own_row_charged_to_the_groups_that_allowed_it():
"""
One socket can run a model from one access group and then a model from
another. Each turn has to reach spend as it finishes, under its own call,
and be charged to the group that let its model through, so a group budget
sees its share during the session and never pays for the other group's turn.
"""
gate = RecordingFrameModelGate(
groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",), "gpt-5.4": ("tier-b",)}),
)
client = _turn_taking_client(
[
json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}),
json.dumps({"type": "response.create", "model": "gpt-5.4"}),
],
answered_by='"response.completed"',
)
upstream_ws = ScriptedUpstreamWebSocket(replies=[_responses_turn("gpt-5.4-mini"), _responses_turn("gpt-5.4")])
relayed = await _relay_turns(client, upstream_ws, gate)
rows = [logged.kwargs for logged in relayed.success_handler.call_args_list]
assert [[frame["response"]["model"] for frame in row["response_body"]] for row in rows] == [
["gpt-5.4-mini", "gpt-5.4-mini"],
["gpt-5.4", "gpt-5.4"],
]
assert [row["litellm_params"]["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] for row in rows] == [
["tier-a"],
["tier-b"],
]
assert [row["litellm_params"]["metadata"]["user_api_key"] for row in rows] == ["hashed-key", "hashed-key"]
assert rows[0]["litellm_call_id"] != rows[1]["litellm_call_id"]
assert rows[1]["start_time"] == rows[0]["end_time"]
relayed.proxy_logging.post_call_failure_hook.assert_not_awaited()
@pytest.mark.asyncio
async def test_websocket_passthrough_carries_the_session_the_provider_opened_into_every_later_turn():
"""
A realtime provider names the session's model once, when it opens the
session, and never again. Each later turn logs on its own, so it has to
carry that opening frame or it would be priced as a model nobody named.
"""
client = _turn_taking_client(
[json.dumps({"type": "response.create"}), json.dumps({"type": "response.create"})],
answered_by='"response.done"',
)
upstream_ws = ScriptedUpstreamWebSocket(
REALTIME_SESSION_CREATED, replies=[[REALTIME_RESPONSE_DONE], [REALTIME_RESPONSE_DONE]]
)
relayed = await _relay_turns(client, upstream_ws, RecordingFrameModelGate(), "/openai/v1/realtime")
assert [
[frame["type"] for frame in logged.kwargs["response_body"]]
for logged in relayed.success_handler.call_args_list
] == [["session.created", "response.done"], ["session.created", "response.done"]]
@pytest.mark.asyncio
async def test_websocket_passthrough_still_bills_the_answer_in_flight_when_a_later_frame_is_refused():
"""
A refused frame closes the client while the provider may still be sending the
answer to an earlier, allowed frame. That answer cannot go to a closed client,
but it was paid for, so it has to reach spend instead of turning the whole
session into a failure that bills nothing.
"""
gate = RecordingFrameModelGate(MappingProxyType({"gpt-5.4": KEY_MODEL_REFUSAL}))
client = _turn_taking_client(
[
json.dumps({"type": "response.create", "model": "gpt-5.4-mini"}),
json.dumps({"type": "response.create", "model": "gpt-5.4"}),
],
answered_by='"response.created"',
)
created, completed = _responses_turn("gpt-5.4-mini")
upstream_ws = ScriptedUpstreamWebSocket(replies=[[created]], in_flight=[completed])
relayed = await _relay_turns(client, upstream_ws, gate)
assert [json.loads(text.args[0])["type"] for text in client.send_text.await_args_list] == [
"response.created",
"error",
]
assert [
[frame["type"] for frame in logged.kwargs["response_body"]]
for logged in relayed.success_handler.call_args_list
] == [["response.created", "response.completed"]]
relayed.proxy_logging.post_call_failure_hook.assert_not_awaited()
DEEPGRAM_LISTEN_TARGET = "wss://api.deepgram.com/v1/listen?model=nova-3&encoding=linear16&sample_rate=16000"
DEEPGRAM_INTERIM_FRAME = json.dumps(
{
@ -8816,15 +9095,6 @@ def test_passthrough_client_cannot_forge_session_id_omission(client_metadata_key
)
class StampingFrameModelGate:
def __init__(self, groups: list[str]):
self.groups = groups
async def __call__(self, model: str, valid_token: UserAPIKeyAuth, /) -> Optional[str]:
valid_token.matched_model_access_groups = self.groups
return None
async def _billed_metadata(gate) -> dict:
served = await _relay_client_frames(
[json.dumps({"type": "response.create", "model": "gpt-5.4-mini"})],
@ -8838,7 +9108,7 @@ async def _billed_metadata(gate) -> dict:
async def test_websocket_session_is_charged_to_the_group_its_frames_were_authorized_through():
"""A websocket names its model in a frame, long after the relay built its logging metadata, so the
groups the gate matched have to be picked up at the end or the group pays for nothing it allowed."""
metadata = await _billed_metadata(StampingFrameModelGate(["tier-a"]))
metadata = await _billed_metadata(RecordingFrameModelGate(groups=MappingProxyType({"gpt-5.4-mini": ("tier-a",)})))
assert metadata[MODEL_ACCESS_GROUP_METADATA_KEY] == ["tier-a"]
assert metadata["user_api_key"] == "hashed-key"

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

@ -1,21 +1,22 @@
"""OpenAI passthrough WebSocket route: registration, opt-in gating, refusals, and per-frame model checks."""
import json
from collections.abc import Iterator, Mapping
from contextlib import contextmanager
from collections.abc import Mapping
from dataclasses import dataclass
from types import MappingProxyType, SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock, 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 Litellm_EntityType, LiteLLM_UserTable, UserAPIKeyAuth
from litellm.proxy._types import Litellm_EntityType, ProxyErrorTypes, ProxyException, UserAPIKeyAuth
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
_OPENAI_WS_DISABLED_REFUSAL,
_OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL,
@ -24,7 +25,16 @@ from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
openai_websocket_proxy_route,
router,
)
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import WebsocketFrameModelGate
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,7 @@ class _RelayCall:
endpoint: str
accept_websocket: bool
client_frame_model_gate: WebsocketFrameModelGate | None
upstream_frame_kind: WebsocketUpstreamFrameClassifier | None
class _FakeRelay:
@ -93,6 +104,7 @@ class _FakeRelay:
endpoint: str,
accept_websocket: bool,
client_frame_model_gate: WebsocketFrameModelGate | None = None,
upstream_frame_kind: WebsocketUpstreamFrameClassifier | None = None,
) -> None:
self.calls.append(
_RelayCall(
@ -102,18 +114,14 @@ class _FakeRelay:
endpoint=endpoint,
accept_websocket=accept_websocket,
client_frame_model_gate=client_frame_model_gate,
upstream_frame_kind=upstream_frame_kind,
)
)
class _FakeFrameModelGate:
def __init__(self, refusal: str | None = None) -> None:
self.refusal = refusal
self.checked: list[str] = []
async def __call__(self, model: str, valid_token: UserAPIKeyAuth, /) -> str | None:
self.checked.append(model)
return self.refusal
async def __call__(self, websocket: object, model: str, /) -> WebsocketFrameModelVerdict:
return WebsocketFrameModelAllowed(access_groups=())
@dataclass(frozen=True, slots=True)
@ -157,6 +165,7 @@ async def test_openai_websocket_forwards_query_and_keeps_provider_auth(prefix, m
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]
@ -310,26 +319,53 @@ async def test_openai_websocket_connects_model_restricted_keys(key_models):
assert websocket.closed is None
@pytest.mark.asyncio
async def test_proxy_frame_model_gate_allows_a_model_the_key_owns():
gate: Final = _proxy_frame_model_gate()
token: Final = UserAPIKeyAuth(models=["gpt-realtime-2.1", "gpt-5.4-mini"])
class _FrameAuthorizer:
"""Stands in for the full key, team, user and budget auth the gate re-runs for every model a frame names."""
with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the gate needs no database
assert await gate("gpt-realtime-2.1", token) is None
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):
gate: Final = _proxy_frame_model_gate()
token: Final = UserAPIKeyAuth(models=["gpt-realtime-2.1", "gpt-5.4-mini"])
authorizer: Final = _FrameAuthorizer(UserAPIKeyAuth(models=["gpt-realtime-2.1", "gpt-5.4-mini"]))
with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the gate needs no database
refusal = await gate(requested_model, token)
refusal: Final = _refusal(await _gate_verdict(authorizer, requested_model))
assert refusal is not None
assert requested_model in refusal
assert requested_model in refusal.message
assert refusal.error_type == ProxyErrorTypes.key_model_access_denied.value
@pytest.mark.asyncio
@ -341,7 +377,6 @@ async def test_proxy_frame_model_gate_refuses_a_model_the_key_lacks(requested_mo
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
gate: Final = _proxy_frame_model_gate()
token: Final = UserAPIKeyAuth(models=[], agent_id="agent-fake")
token.managed_agent_policy = AgentResponse(
agent_id="agent-fake",
@ -350,132 +385,105 @@ async def test_proxy_frame_model_gate_refuses_a_model_the_agent_was_not_granted(
object_permission={"object_permission_id": "grant-fake", "models": granted_models},
)
with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the gate needs no database
refusal = await gate("gpt-5.4", token)
refusal: Final = _refusal(await _gate_verdict(_FrameAuthorizer(token), "gpt-5.4"))
assert refusal is not None
assert refusal_names in refusal
GET_USER_OBJECT: Final = "litellm.proxy.auth.auth_checks.get_user_object"
def _personal_user(models: list[str]) -> LiteLLM_UserTable:
return LiteLLM_UserTable(user_id="user-fake", models=models, max_budget=None, spend=0.0)
@contextmanager
def _personal_user_allowlist(models: list[str]) -> Iterator[AsyncMock]:
"""Stand up the database-backed owner lookup the allowlist check runs, and report every call it made."""
lookup: Final = AsyncMock(return_value=_personal_user(models))
with (
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: read at call time
patch(GET_USER_OBJECT, lookup),
):
yield lookup
@pytest.mark.asyncio
async def test_proxy_frame_model_gate_refuses_a_model_the_person_behind_the_key_lacks():
"""
A key with no team of its own carries its owner's model allowlist, the same
one every HTTP route applies, so a frame naming a model that person cannot
use is refused even when the key itself names no models at all.
"""
gate: Final = _proxy_frame_model_gate()
token: Final = UserAPIKeyAuth(models=[], user_id="user-fake")
with _personal_user_allowlist(["gpt-5.4-mini"]):
refusal = await gate("gpt-5.4", token)
assert refusal is not None
assert "gpt-5.4" in refusal
@pytest.mark.asyncio
async def test_proxy_frame_model_gate_allows_a_model_the_person_behind_the_key_owns():
gate: Final = _proxy_frame_model_gate()
token: Final = UserAPIKeyAuth(models=[], user_id="user-fake")
with _personal_user_allowlist(["gpt-5.4-mini", "gpt-realtime-2.1"]):
assert await gate("gpt-realtime-2.1", token) is None
@pytest.mark.asyncio
async def test_proxy_frame_model_gate_leaves_a_team_key_to_the_teams_allowlist():
"""
A key that belongs to a team is governed by the team's models, not by the
personal allowlist of whoever created it, so the owner's own restrictions
must not close a session the team is entitled to.
"""
gate: Final = _proxy_frame_model_gate()
token: Final = UserAPIKeyAuth(models=[], user_id="user-fake", team_id="team-fake")
with _personal_user_allowlist(["gpt-5.4-mini"]) as owner_lookup:
assert await gate("gpt-5.4", token) is None
assert owner_lookup.await_args_list == []
assert refusal_names in refusal.message
@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):
gate: Final = _proxy_frame_model_gate()
token: Final = UserAPIKeyAuth(models=unrestricted_models)
verdict: Final = await _gate_verdict(_FrameAuthorizer(UserAPIKeyAuth(models=unrestricted_models)), "gpt-5.4")
with patch("litellm.proxy.proxy_server.prisma_client", None): # test-quality-ok: the gate needs no database
assert await gate("gpt-realtime-2.1", token) is None
assert verdict == WebsocketFrameModelAllowed(access_groups=())
ENFORCE_GROUP_BUDGET: Final = (
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.enforce_model_access_group_budget"
@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_refuses_a_model_whose_access_group_is_out_of_budget():
"""
A model that sits in a budgeted access group is refused on every other route once
that group has spent its budget, so a frame naming it closes the session instead
of quietly running on an empty pool. The key is allowed the model outright here,
so nothing but the group's budget can produce the refusal.
"""
exhausted: Final = 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",
)
gate: Final = _proxy_frame_model_gate()
token: Final = UserAPIKeyAuth(models=["gpt-5.4-mini"])
@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"))
with (
patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: the model check needs no database
patch(ENFORCE_GROUP_BUDGET, AsyncMock(side_effect=exhausted)),
):
refusal = await gate("gpt-5.4-mini", token)
assert refusal is not None
assert "Budget has been exceeded" in refusal
assert "Model access group=tier-a" in refusal
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_it_cannot_reach_the_budgets():
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.
"""
gate: Final = _proxy_frame_model_gate()
token: Final = UserAPIKeyAuth(models=["gpt-5.4-mini"])
refusal: Final = _refusal(
await _gate_verdict(_FrameAuthorizer(denial=ConnectionError("connection pool is exhausted")), "gpt-5.4-mini")
)
with (
patch("litellm.proxy.proxy_server.prisma_client", None), # test-quality-ok: the model check needs no database
patch(ENFORCE_GROUP_BUDGET, AsyncMock(side_effect=ConnectionError("connection pool is exhausted"))),
):
refusal = await gate("gpt-5.4-mini", token)
assert refusal == _OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL
assert "connection pool" not in refusal
assert refusal == WebsocketFrameModelRefused(message=_OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL, error_type="server_error")