fix(proxy): check the model in every OpenAI websocket passthrough frame and price the session

litellm/proxy/auth/auth_checks.py

@@ -5291,79 +5291,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,
@@ -5677,6 +5604,115 @@ async def can_key_call_resolved_model(
             )
 
 
+async def can_personal_user_call_model(
+    model: str,
+    valid_token: UserAPIKeyAuth,
+    llm_router: Router | None,
+) -> None:
+    """Apply the personal user's model allowlist, the check ``common_checks`` runs for a key with no team."""
+    from litellm.proxy.proxy_server import (
+        prisma_client,
+        proxy_logging_obj,
+        user_api_key_cache,
+    )
+
+    if valid_token.team_id is not None or valid_token.user_id is None or prisma_client is None:
+        return
+    try:
+        user_object: Final = 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,
+            parent_otel_span=valid_token.parent_otel_span,
+            proxy_logging_obj=proxy_logging_obj,
+        )
+    except UserNotFoundError:
+        return
+    if user_object is None:
+        return
+    await can_user_call_model(
+        model=model,
+        llm_router=llm_router,
+        user_object=user_object,
+        key_model_aliases=key_model_aliases_for_auth_check(valid_token),
+    )
+
+
+async def _matching_team_object(
+    valid_token: UserAPIKeyAuth,
+    prisma_client: PrismaClient,
+    user_api_key_cache: UserApiKeyCache,
+    proxy_logging_obj: ProxyLogging,
+) -> LiteLLM_TeamTableCachedObj | None:
+    if valid_token.team_id is None:
+        return None
+    try:
+        return await get_team_object(
+            team_id=valid_token.team_id,
+            prisma_client=prisma_client,
+            user_api_key_cache=user_api_key_cache,
+            parent_otel_span=valid_token.parent_otel_span,
+            proxy_logging_obj=proxy_logging_obj,
+        )
+    except Exception as e:  # noqa: BLE001  # fail-safe: a team we cannot read grants no group, it never breaks auth
+        verbose_proxy_logger.debug("model access group attribution could not read the team: %s", e)
+        return None
+
+
+async def enforce_model_access_group_budget(
+    model: str,
+    valid_token: UserAPIKeyAuth,
+    llm_router: Router | None,
+) -> None:
+    """Attribute this model to the budgeted access groups that authorize it and refuse an exhausted one.
+
+    The pair ``common_checks`` runs for every HTTP request, for callers that reach a model only
+    through a group: the stamp is what the spend writer charges afterwards, and the check is what
+    stops a group whose budget is already gone.
+
+    Raises:
+        BudgetExceededError if a group that authorized this model is over its max budget.
+    """
+    from litellm.proxy.proxy_server import (
+        prisma_client,
+        proxy_logging_obj,
+        user_api_key_cache,
+    )
+
+    if prisma_client is None:
+        return
+    matched: Final = await stamp_matched_model_access_groups(
+        model=model,
+        valid_token=valid_token,
+        team_object=await _matching_team_object(
+            valid_token=valid_token,
+            prisma_client=prisma_client,
+            user_api_key_cache=user_api_key_cache,
+            proxy_logging_obj=proxy_logging_obj,
+        ),
+        project_object=(
+            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,
+            )
+        ),
+        llm_router=llm_router,
+        prisma_client=prisma_client,
+        user_api_key_cache=user_api_key_cache,
+        proxy_logging_obj=proxy_logging_obj,
+    )
+    await _model_access_group_max_budget_check(
+        matched_model_access_groups=matched,
+        prisma_client=prisma_client,
+        user_api_key_cache=user_api_key_cache,
+    )
+
+
 def can_org_access_model(
     model: str,
     org_object: LiteLLM_OrganizationTable | None,

litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py

@@ -15,7 +15,7 @@ import os
 import posixpath
 import re
 import sys
-from collections.abc import AsyncGenerator, Callable, Mapping, Sequence
+from collections.abc import AsyncGenerator, Callable, Mapping
 from dataclasses import dataclass
 from functools import partial
 from types import MappingProxyType
@@ -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,12 @@ from litellm.llms.oss_decision import OssDecisionProvider, oss_connection, valid
 from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
 from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error
 from litellm.proxy._types import *
-from litellm.proxy.auth.auth_checks import enforced_model_allowlists
+from litellm.proxy.auth.auth_checks import (
+    can_agent_call_model,
+    can_key_call_resolved_model,
+    can_personal_user_call_model,
+    enforce_model_access_group_budget,
+)
 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
@@ -110,6 +116,7 @@ from litellm.proxy.pass_through_endpoints.injected_stream_usage import (
 )
 from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
     HttpPassThroughEndpointHelpers,
+    WebsocketFrameModelGate,
     create_pass_through_route,
     create_websocket_passthrough_route,
     websocket_passthrough_request,
@@ -173,7 +180,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
 
@@ -203,9 +210,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:
@@ -3261,19 +3267,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
@@ -3290,6 +3283,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=(
@@ -3298,14 +3296,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")
@@ -3314,19 +3304,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
 
 
@@ -3341,6 +3321,7 @@ class _WebsocketRelay(Protocol):
         forward_headers: bool,
         endpoint: str,
         accept_websocket: bool,
+        client_frame_model_gate: WebsocketFrameModelGate | None = None,
     ) -> None: ...
 
 
@@ -3348,18 +3329,42 @@ 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_frame_model_gate() -> WebsocketFrameModelGate:
+    from litellm.proxy.proxy_server import (
+        llm_model_list,
+        llm_router,
+        prisma_client,
+        proxy_logging_obj,
+        user_api_key_cache,
+    )
 
-    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,
-        )
+    async def gate(model: str, valid_token: UserAPIKeyAuth, /) -> str | None:
+        try:
+            await can_key_call_resolved_model(
+                model=model, llm_model_list=llm_model_list, valid_token=valid_token, llm_router=llm_router
+            )
+            await can_agent_call_model(
+                model=model,
+                valid_token=valid_token,
+                llm_router=llm_router,
+                prisma_client=prisma_client,
+                user_api_key_cache=user_api_key_cache,
+                proxy_logging_obj=proxy_logging_obj,
+            )
+            await can_personal_user_call_model(model=model, valid_token=valid_token, llm_router=llm_router)
+            await enforce_model_access_group_budget(model=model, valid_token=valid_token, llm_router=llm_router)
+        except ProxyException as denial:
+            return denial.message
+        except HTTPException as denial:
+            return str(denial.detail)
+        except litellm.BudgetExceededError as exhausted:
+            return exhausted.message
+        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 _OPENAI_WS_UNAVAILABLE_AUTH_REFUSAL
+        return None
 
-    return resolve
+    return gate
 
 
 def _negotiated_websocket_subprotocol(websocket: WebSocket) -> str | None:
@@ -3379,12 +3384,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] = {
@@ -3438,6 +3443,7 @@ async def openai_websocket_proxy_route(
         forward_headers=False,
         endpoint=websocket.url.path,
         accept_websocket=False,
+        client_frame_model_gate=frame_model_gate,
     )
 
 
@@ -3683,7 +3689,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:

litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py

@@ -0,0 +1,207 @@
+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.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_COMPLETED: Final = "response.completed"
+_MODEL_NAMING_EVENTS: Final = frozenset(
+    {_SESSION_CREATED, _RESPONSE_CREATED, _REALTIME_RESPONSE_DONE, _RESPONSES_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 _ModelNamingEvent(BaseModel):
+    model_config = ConfigDict(extra="ignore", frozen=True)
+
+    session: _NamedModel | None = None
+    response: _NamedModel | None = None
+
+
+class _CompletedResponse(BaseModel):
+    model_config = ConfigDict(extra="ignore", frozen=True)
+
+    model: str | None = None
+    usage: ResponseAPIUsage
+
+
+class _ResponsesCompletedEvent(BaseModel):
+    model_config = ConfigDict(extra="ignore", frozen=True)
+
+    type: Literal["response.completed"]
+    response: _CompletedResponse
+
+
+_MODEL_NAMING_EVENT: Final = TypeAdapter(_ModelNamingEvent)
+_RESPONSES_COMPLETED_EVENT: Final = TypeAdapter(_ResponsesCompletedEvent)
+
+
+@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 _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") != _RESPONSES_COMPLETED:
+        return None
+    try:
+        completed: Final = _RESPONSES_COMPLETED_EVENT.validate_python(event)
+    except ValidationError:
+        verbose_proxy_logger.warning(
+            "OpenAI websocket passthrough: a response.completed frame carried no readable usage, so it is not billed"
+        )
+        return None
+    return _CompletedTurn(
+        model=completed.response.model or session_model,
+        usage=ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(  # pyright: ignore[reportPrivateUsage]  # the one shared Responses usage converter
+            completed.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

litellm/proxy/pass_through_endpoints/pass_through_endpoints.py

@@ -5,12 +5,13 @@ import json
 import posixpath
 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, TypedDict, cast
 from urllib.parse import urlencode, urlparse
 
 import httpx
@@ -23,6 +24,7 @@ from fastapi import (
     Response,
     UploadFile,
     WebSocket,
+    WebSocketDisconnect,
     status,
 )
 from fastapi.responses import StreamingResponse
@@ -30,7 +32,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,
@@ -55,6 +57,7 @@ from litellm.integrations.custom_guardrail import (
 )
 from litellm.integrations.custom_logger import CustomLogger
 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,
@@ -64,6 +67,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
@@ -2454,6 +2458,106 @@ def _with_trace_context(headers: Mapping[str, str], parent_span: object) -> dict
     return inject_trace_context(headers, parent_span=parent_span)
 
 
+class WebsocketFrameModelGate(Protocol):
+    async def __call__(self, model: str, valid_token: UserAPIKeyAuth, /) -> str | None: ...
+
+
+@dataclass(frozen=True, slots=True)
+class _ForwardClientFrame:
+    pass
+
+
+@dataclass(frozen=True, slots=True)
+class _RefuseClientFrame:
+    message: str
+
+
+_PLAIN_MODEL_KEY: Final = '"model"'
+_JSON_STRING_ESCAPE: Final = "\\"
+_CLIENT_FRAME_MODEL_PLACEMENTS: Final = (
+    ("model",),
+    ("response", "model"),
+    ("session", "model"),
+    ("session", "input_audio_transcription", "model"),
+    ("session", "audio", "input", "transcription", "model"),
+)
+
+
+def _may_name_a_model(text_data: str) -> bool:
+    return _PLAIN_MODEL_KEY in text_data or _JSON_STRING_ESCAPE in text_data
+
+
+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_model(frame: Mapping[str, object], placement: Sequence[str]) -> str | None:
+    *parents, leaf = placement
+    container: Final = reduce(_nested_frame, parents, frame)
+    if container is None:
+        return None
+    value: Final = container.get(leaf)
+    return value if isinstance(value, str) else None
+
+
+def _client_frame_models(frame_data: str | bytes) -> tuple[str, ...]:
+    if isinstance(frame_data, str) and not _may_name_a_model(frame_data):
+        return ()
+    frame: Final = _json_object_frame(frame_data)
+    if frame is None:
+        return ()
+    placed: Final = (_placed_model(frame, placement) for placement in _CLIENT_FRAME_MODEL_PLACEMENTS)
+    return tuple(dict.fromkeys(model for model in placed if model is not None))
+
+
+async def _client_frame_verdict(
+    frame_data: str | bytes, gate: WebsocketFrameModelGate | None, valid_token: UserAPIKeyAuth
+) -> _ForwardClientFrame | _RefuseClientFrame:
+    if gate is None:
+        return _ForwardClientFrame()
+    for model in _client_frame_models(frame_data):
+        if (refusal := await gate(model, valid_token)) is not None:
+            return _RefuseClientFrame(message=refusal)
+    return _ForwardClientFrame()
+
+
+async def _refuse_client_frame(websocket: WebSocket, upstream_ws: ClientConnection, message: str) -> None:
+    await websocket.send_text(realtime_error_event(message, error_type="invalid_request_error"))
+    await websocket.close(code=CloseCode.POLICY_VIOLATION, reason=_truncated_close_reason(message))
+    await upstream_ws.close()
+
+
+def _stamp_billed_access_groups(kwargs: Mapping[str, object], valid_token: UserAPIKeyAuth) -> None:
+    """Refresh the session's logging metadata with the access groups its frames were authorized through.
+
+    The gate only learns them once a frame names a model, which is long after the relay stamped the
+    key's groups onto its metadata, so the spend writer would otherwise charge no group for the session.
+    """
+    groups: Final = valid_token.matched_model_access_groups
+    litellm_params: Final = as_str_mapping(kwargs.get("litellm_params"))
+    metadata: Final = None if litellm_params is None else litellm_params.get("metadata")
+    if groups and isinstance(metadata, MutableMapping):
+        metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = groups
+
+
+async def _relay_client_frame(
+    websocket: WebSocket,
+    upstream_ws: ClientConnection,
+    checked_frame: str | bytes,
+    forwarded_frame: str | bytes,
+    gate: WebsocketFrameModelGate | None,
+    valid_token: UserAPIKeyAuth,
+) -> bool:
+    """Forward one client frame upstream, or refuse it and report that the session is over."""
+    match await _client_frame_verdict(checked_frame, gate, valid_token):
+        case _RefuseClientFrame(message=refusal_message):
+            await _refuse_client_frame(websocket, upstream_ws, refusal_message)
+            return False
+        case _ForwardClientFrame():
+            await upstream_ws.send(forwarded_frame)
+            return True
+
+
 async def websocket_passthrough_request(
     websocket: WebSocket,
     target: str,
@@ -2464,6 +2568,7 @@ async def websocket_passthrough_request(
     cost_per_request: float | None = None,
     accept_websocket: bool = True,
     setup_model_rewriter: Callable[[str], str] | None = None,
+    client_frame_model_gate: WebsocketFrameModelGate | None = None,
 ):
     """
     WebSocket passthrough request handler.
@@ -2673,9 +2778,25 @@ 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 not await _relay_client_frame(
+                                websocket,
+                                upstream_ws,
+                                text_data,
+                                _rewrite_vertex_live_setup_model(text_data, setup_model_rewriter),
+                                client_frame_model_gate,
+                                user_api_key_dict,
+                            ):
+                                break
                         elif bytes_data is not None:
-                            await upstream_ws.send(bytes_data)
+                            if not await _relay_client_frame(
+                                websocket,
+                                upstream_ws,
+                                bytes_data,
+                                bytes_data,
+                                client_frame_model_gate,
+                                user_api_key_dict,
+                            ):
+                                break
                 except asyncio.CancelledError:
                     raise
                 except Exception:
@@ -2702,11 +2823,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
@@ -2715,6 +2832,13 @@ async def websocket_passthrough_request(
                     return
                 websocket_messages.append(message_data)
 
+            async def relay_upstream_frame(upstream_message: str | bytes) -> None:
+                record_upstream_frame(upstream_message)
+                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:
                     while True:
@@ -2722,6 +2846,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
@@ -2767,6 +2896,8 @@ async def websocket_passthrough_request(
             passthrough_logging_payload["response_body"] = websocket_messages
             passthrough_logging_payload["end_time"] = end_time
 
+            _stamp_billed_access_groups(kwargs, user_api_key_dict)
+
             # Remove logging_obj from kwargs to avoid duplicate keyword argument
             success_kwargs: Final = kwargs.copy()
             success_kwargs.pop("logging_obj", None)

litellm/proxy/pass_through_endpoints/success_handler.py

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