Merge 4af9c78747 into 13335266f5

litellm/cost_calculator.py

@@ -3113,14 +3113,19 @@ class ResponsesWebSocketTokenUsageProcessor(BaseTokenUsageProcessor):
         return response_ids[-1] if response_ids else None
 
     @staticmethod
-    def create_logging_object(results: Sequence[Mapping[str, object]]) -> LiteLLMRealtimeStreamLoggingObject:
+    def create_logging_object(
+        results: Sequence[Mapping[str, object]], litellm_model_name: str | None = None
+    ) -> LiteLLMRealtimeStreamLoggingObject:
         tier_partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier(results)
-        return LiteLLMRealtimeStreamLoggingObject(
+        logging_object: Final = LiteLLMRealtimeStreamLoggingObject(
             id=ResponsesWebSocketTokenUsageProcessor.response_id_from_responses_ws_results(results),
             usage=ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results(results),
             results=cast(OpenAIRealtimeStreamList, list(results)),  # cast-ok: raw WS frames, unvalidated field
             service_tier=next(iter(tier_partition)) if len(tier_partition) == 1 else None,
         )
+        if litellm_model_name is not None:
+            logging_object.hidden_params = {"litellm_model_name": litellm_model_name}
+        return logging_object
 
 
 _TRANSCRIPTION_COMPLETED_EVENT_TYPE: Final = "conversation.item.input_audio_transcription.completed"

litellm/litellm_core_utils/litellm_logging.py

@@ -986,6 +986,16 @@ class Logging(LiteLLMLoggingBaseClass):
                 return model_id
         return None
 
+    def _routed_deployment_model(self) -> str | None:
+        litellm_params: Final = self._model_call_details_view.get("litellm_params")
+        metadata: Final = (
+            StandardLoggingPayloadSetup.merge_litellm_metadata(_LITELLM_PARAMS.validate_python(litellm_params))
+            if isinstance(litellm_params, dict)
+            else {}
+        )
+        deployment: Final = metadata.get("deployment")
+        return deployment if isinstance(deployment, str) else None
+
     def get_deployment_model_for_cost(self) -> str | None:
         """The provider-qualified model to price against.
 
@@ -2253,6 +2263,7 @@ class Logging(LiteLLMLoggingBaseClass):
             and litellm_params.get(CallTypes.agenerate_content.value, False) is not True
             and litellm_params.get(CallTypes.agenerate_content_stream.value, False) is not True
             and litellm_params.get(CallTypes.arealtime.value, False) is not True
+            and litellm_params.get(CallTypes.aresponses_websocket.value, False) is not True
         )
 
     _is_sync_litellm_request = is_sync_litellm_request
@@ -2508,7 +2519,8 @@ class Logging(LiteLLMLoggingBaseClass):
 
         elif self.call_type == CallTypes.aresponses_websocket.value and isinstance(result, list):  # pyright: ignore[reportUnknownMemberType]  # Logging.call_type is untyped
             logging_result = ResponsesWebSocketTokenUsageProcessor.create_logging_object(
-                results=result  # pyright: ignore[reportUnknownArgumentType]  # raw event dicts from the WS stream
+                results=result,  # pyright: ignore[reportUnknownArgumentType]  # raw event dicts from the WS stream
+                litellm_model_name=self._routed_deployment_model(),
             )
 
         elif self.relays_llm_passthrough and isinstance(result, Response):

litellm/llms/custom_httpx/llm_http_handler.py

@@ -6910,6 +6910,7 @@ class BaseLLMHTTPHandler:
         custom_llm_provider: str | None = None,
         first_message: str | None = None,
         request_defaults: ResponsesWebSocketRequestDefaults | None = None,
+        turn_admission: Callable[[], Awaitable[None]] | None = None,
         frame_guard: "ResponsesWsFrameGuard | None" = None,
         **kwargs: object,
     ) -> Exception | None:
@@ -6972,12 +6973,12 @@ class BaseLLMHTTPHandler:
                 custom_llm_provider=custom_llm_provider,
                 first_message=first_message,
                 quota_callbacks=_ws_quota_callbacks,
+                turn_admission=turn_admission,
                 guardrail_callbacks=_ws_guardrail_callbacks,  # pyright: ignore[reportUnknownArgumentType]  # duck-typed callbacks from litellm.callbacks
                 frame_guard=frame_guard,
                 **kwargs,
             )
-            await handler.run()
-            return None
+            return await handler.run()
 
         import websockets
         from websockets.asyncio.client import ClientConnection
@@ -7068,6 +7069,7 @@ class BaseLLMHTTPHandler:
                     authorized_model=model,
                     custom_llm_provider=custom_llm_provider,
                     request_defaults=request_defaults,
+                    turn_admission=turn_admission,
                     frame_guard=frame_guard,
                 )
                 return await streaming.bidirectional_forward()

litellm/proxy/hooks/proxy_track_cost_callback.py

@@ -2,8 +2,11 @@ import asyncio
 import traceback
 from collections.abc import Callable, Mapping, Sequence
 from datetime import datetime
+from types import MappingProxyType
 from typing import TYPE_CHECKING, Final, Protocol, cast
 
+from pydantic import TypeAdapter, ValidationError
+
 import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.batches.batch_utils import batch_cost_is_final
@@ -17,6 +20,7 @@ from litellm.litellm_core_utils.core_helpers import (
 )
 from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
 from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost
+from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.auth.auth_checks import (
     get_key_object,
@@ -100,6 +104,37 @@ def _proxy_stamped_used_client_oauth_token(
     return stamped if isinstance(stamped, bool) else None
 
 
+def _str_keyed_fields(value: object) -> Mapping[str, object]:
+    return value if is_str_keyed_mapping(value) else {}
+
+
+def _is_model_group(value: object) -> bool:
+    return isinstance(value, str) and bool(value)
+
+
+_TAG_LIST: Final = TypeAdapter(list[str])
+
+
+def _is_tag_list(value: object) -> bool:
+    try:
+        return bool(_TAG_LIST.validate_python(value, strict=True))
+    except ValidationError:
+        return False
+
+
+_ROUTED_METADATA_FIELD_CHECKS: Final = MappingProxyType({"model_group": _is_model_group, "tags": _is_tag_list})
+
+
+def _routed_metadata_fallback(request_data: Mapping[str, object]) -> Mapping[str, object]:
+    caller_metadata: Final = _str_keyed_fields(request_data.get("metadata"))
+    routed_metadata: Final = _str_keyed_fields(request_data.get("litellm_metadata"))
+    return {
+        field: routed_metadata[field]
+        for field, is_valid in _ROUTED_METADATA_FIELD_CHECKS.items()
+        if not caller_metadata.get(field) and is_valid(routed_metadata.get(field))
+    }
+
+
 def _failure_snapshot_with_guardrails_recorded_after_it(
     failure_snapshot: object, recorded_guardrail_information: object
 ) -> Mapping[str, object] | None:
@@ -220,10 +255,11 @@ class _ProxyDBLogger(CustomLogger):
             metadata=_metadata, original_exception=original_exception
         )
 
-        _metadata["used_client_oauth_token"] = _proxy_stamped_used_client_oauth_token(request_data, request_route)
+        request_fields: Final = _str_keyed_fields(request_data)
+        _metadata["used_client_oauth_token"] = _proxy_stamped_used_client_oauth_token(request_fields, request_route)
 
         existing_metadata: Final[dict] = request_data.get("metadata", None) or {}
-        existing_metadata.update(_metadata)
+        existing_metadata.update({**_routed_metadata_fallback(request_fields), **_metadata})
 
         litellm_metadata_bucket: Final = request_data.get("litellm_metadata")
         if (
@@ -471,11 +507,7 @@ class _ProxyDBLogger(CustomLogger):
                 await _release_budget_reservation(budget_reservation=budget_reservation)
                 # Non-model call types (health checks, afile_delete) have no model or standard_logging_object.
                 # Use .get() for "stream" to avoid KeyError on health checks.
-                # WS session wrappers (_aresponses_websocket, _arealtime) also reach here with
-                # result=None; their per-turn costs are tracked on the inner aresponses/realtime calls.
-                if sl_object is None and (
-                    not kwargs.get("model") or kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime")
-                ):
+                if sl_object is None and (not kwargs.get("model") or kwargs.get("call_type") == "_arealtime"):
                     verbose_proxy_logger.warning(
                         "Cost tracking - skipping, no standard_logging_object for call_type=%s",
                         kwargs.get("call_type", "unknown"),

litellm/proxy/response_api_endpoints/endpoints.py

@@ -78,6 +78,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio
     request_tags_from_request_data,
 )
 from litellm.proxy.route_llm_request import raise_if_required_body_param_missing
+from litellm.proxy.utils import handle_exception_on_proxy
 from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWsFrameGuard
 from litellm.router import Router
 from litellm.types.llms.base import LiteLLMBaseModel
@@ -1861,12 +1862,14 @@ def _routing_hints_from_first_ws_frame(first_message: str) -> Mapping[str, objec
 
 
 def _responses_ws_failure_frame(failure: Exception) -> str:
-    raw_status: Final = getattr(failure, "status_code", None)
-    status: Final = raw_status if isinstance(raw_status, int) and not isinstance(raw_status, bool) else 500
+    proxy_exception: Final = handle_exception_on_proxy(failure)
+    status: Final = int(proxy_exception.code) if str(proxy_exception.code).isdigit() else 500
     error_type: Final = (
         "rate_limit_exceeded" if status == 429 else "invalid_request_error" if 400 <= status < 500 else "server_error"
     )
-    return json.dumps({"type": "error", "status": status, "error": {"type": error_type, "message": str(failure)}})
+    return json.dumps(
+        {"type": "error", "status": status, "error": {"type": error_type, "message": proxy_exception.message}}
+    )
 
 
 async def _enforce_responses_ws_first_frame_model_auth(
@@ -1949,7 +1952,7 @@ def _responses_ws_frame_guard(
     }
     model: Final = admitted.get("model")
 
-    # a blocked frame is billed as its own failed request; only the first turn settles the admission reservation
+    # blocked and refused frames bill as their own failed requests; only the first turn settles the reservation
     rejected_frame_auth: Final = user_api_key_dict.model_copy(update={"budget_reservation": None})
 
     async def guard(frame: str) -> GuardedResponsesWsFrame:
@@ -1997,6 +2000,7 @@ def _responses_ws_frame_guard(
             frame=json.dumps(rewritten),
             guardrail_information=_recorded_guardrail_information(guarded_data.get(metadata_key)),
             pii_tokens=_recorded_pii_tokens(guarded_data.get("metadata")),
+            bill_refusal=partial(_bill_rejected_frame, rejected_frame_auth, frame_data=guarded_data),
         )
 
     return guard
@@ -2167,7 +2171,6 @@ async def _responses_websocket_session(
                 request_data=routed_data,
             )
     except Exception as e:
-        verbose_proxy_logger.exception("Responses WebSocket error")
         with contextlib.suppress(Exception):
             await websocket.send_text(_responses_ws_failure_frame(e))
         await proxy_logging_obj.post_call_failure_hook(

litellm/responses/main.py

@@ -1,7 +1,7 @@
 import asyncio
 import contextvars
 import json
-from collections.abc import Coroutine, Generator, Iterable, Mapping, Sequence
+from collections.abc import Awaitable, Callable, Coroutine, Generator, Iterable, Mapping, Sequence
 from contextlib import contextmanager
 from dataclasses import dataclass
 from functools import partial
@@ -73,7 +73,7 @@ from litellm.llms.openai.data_residency import infer_openai_data_residency
 from litellm.types.responses.main import *
 from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults
 from litellm.types.router import GenericLiteLLMParams
-from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD, all_litellm_params
+from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD, CallTypes, all_litellm_params
 from litellm.utils import (
     ProviderConfigManager,
     client,
@@ -2609,6 +2609,7 @@ async def _aresponses_websocket(
     api_base: str | None = None,
     api_key: str | None = None,
     timeout: float | None = None,
+    turn_admission: Callable[[], Awaitable[None]] | None = None,
     responses_ws_frame_guard: "ResponsesWsFrameGuard | None" = None,
     **kwargs,
 ) -> Exception | None:
@@ -2624,7 +2625,7 @@ async def _aresponses_websocket(
     litellm_logging_obj: Final[LiteLLMLoggingObj] = kwargs.get("litellm_logging_obj")
     user: Final = kwargs.get("user", None)
     litellm_params: Final = GenericLiteLLMParams(**kwargs)
-    litellm_params_dict: Final = get_litellm_params(**kwargs)
+    litellm_params_dict: Final = {**get_litellm_params(**kwargs), CallTypes.aresponses_websocket.value: True}
 
     (
         provider_model,
@@ -2702,6 +2703,7 @@ async def _aresponses_websocket(
         litellm_metadata=_build_litellm_metadata_for_ws(kwargs),
         custom_llm_provider=_custom_llm_provider,
         request_defaults=_build_responses_websocket_request_defaults(deployment_kwargs),
+        turn_admission=turn_admission,
         frame_guard=responses_ws_frame_guard,
         **remaining_kwargs,
     )

litellm/responses/streaming_iterator.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import copy
 import json
 import time
@@ -1929,6 +1930,38 @@ async def _enforce_frame_project_quota(
         )
 
 
+async def _admit_frame_turn(turn_admission: Callable[[], Awaitable[None]] | None, raw_message: str) -> Exception | None:
+    frame: Final = _json_object_or_none(raw_message)
+    if turn_admission is None or frame is None or frame.get("type") != "response.create":
+        return None
+    try:
+        await turn_admission()
+    except RateLimitError as refusal:
+        return refusal
+    except Exception as e:
+        verbose_logger.exception("Responses WS turn admission failed: %s", e)
+        return e
+    return None
+
+
+def _turn_refusal_frame(refusal: RateLimitError) -> str:
+    return json.dumps(
+        {"type": "error", "status": 429, "error": {"type": "rate_limit_exceeded", "message": str(refusal)}}
+    )
+
+
+def _turn_admission_error_frame(failure: Exception) -> str:
+    return json.dumps({"type": "error", "status": 500, "error": {"type": "server_error", "message": str(failure)}})
+
+
+def _turn_rejection_frame(rejection: Exception) -> str:
+    match rejection:
+        case RateLimitError():
+            return _turn_refusal_frame(rejection)
+        case _:
+            return _turn_admission_error_frame(rejection)
+
+
 RESPONSES_WS_LOGGED_EVENT_TYPES: Final = [
     "response.created",
     "response.completed",
@@ -2001,16 +2034,28 @@ def _unmask_responses_event(response_str: str, pii_tokens: Mapping[str, str], un
 @dataclass(frozen=True, slots=True)
 class GuardedResponsesWsFrame:
     """A ``response.create`` frame after the proxy's pre_call guardrails ran on it: the frame to forward (guardrails
-    may rewrite ``input`` and ``instructions``), the guardrail records that bill against this turn only, and the
-    Presidio ``pii_tokens`` the walk masked with, which the handler needs to unmask the provider's reply."""
+    may rewrite ``input`` and ``instructions``), the guardrail records that bill against this turn only, the
+    Presidio ``pii_tokens`` the walk masked with, which the handler needs to unmask the provider's reply, and how
+    to bill those guardrail records when the deployment then refuses the turn."""
 
     frame: str
     guardrail_information: tuple[Mapping[str, object], ...] = ()
     pii_tokens: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({}))
+    bill_refusal: Callable[[Exception], Awaitable[None]] | None = None
 
 
 ResponsesWsFrameGuard: TypeAlias = Callable[[str], Awaitable[GuardedResponsesWsFrame]]
 
+
+async def _reject_guarded_turn(
+    websocket: ResponsesClientWebSocket, guarded: GuardedResponsesWsFrame, rejection: Exception
+) -> None:
+    with contextlib.suppress(Exception):
+        await websocket.send_text(_turn_rejection_frame(rejection))
+    if guarded.guardrail_information and guarded.bill_refusal is not None:
+        await guarded.bill_refusal(rejection)
+
+
 _RESPONSES_WS_TERMINAL_EVENT_TYPES: Final = _HISTORY_TERMINAL_EVENT_TYPES | _RESPONSES_WS_FAILURE_EVENT_TYPES
 
 
@@ -2194,6 +2239,7 @@ class ResponsesWebSocketStreaming:
         authorized_model: str | None = None,
         custom_llm_provider: str | None = None,
         request_defaults: ResponsesWebSocketRequestDefaults | None = None,
+        turn_admission: Callable[[], Awaitable[None]] | None = None,
         frame_guard: ResponsesWsFrameGuard | None = None,
     ):
         self.websocket = websocket
@@ -2215,6 +2261,7 @@ class ResponsesWebSocketStreaming:
         # response.create frame to prevent deployment-substitution attacks.
         self.authorized_model: str | None = authorized_model
         self.request_defaults: ResponsesWebSocketRequestDefaults | None = request_defaults
+        self.turn_admission: Callable[[], Awaitable[None]] | None = turn_admission
 
     def _start_turn(
         self, client_frame: str | None, guardrail_information: tuple[Mapping[str, object], ...] = ()
@@ -2266,9 +2313,11 @@ class ResponsesWebSocketStreaming:
             await self._send_error_event("invalid_request_error", str(e))
             return None
 
-    async def _forward_guarded_frame(self, message: str) -> None:
+    async def _forward_guarded_frame(
+        self, message: str, turn_admission: Callable[[], Awaitable[None]] | None = None
+    ) -> None:
         guarded: Final = await self._guard_frame(message)
-        if guarded is None:
+        if guarded is None or not await self._admit_or_reject_turn(guarded, message, turn_admission):
             return
         if guarded.pii_tokens:
             self._remember_pii_tokens(guarded.pii_tokens)
@@ -2662,10 +2711,27 @@ class ResponsesWebSocketStreaming:
                 self.quota_callbacks, self.user_api_key_dict, self.authorized_model, message
             )
         except RateLimitError as e:
-            await self._send_error_event("rate_limit_exceeded", str(e))
+            await self._send_turn_refusal(e)
             return False
         return True
 
+    async def _admit_or_reject_turn(
+        self, guarded: GuardedResponsesWsFrame, message: str, turn_admission: Callable[[], Awaitable[None]] | None
+    ) -> bool:
+        rejection: Final = await _admit_frame_turn(turn_admission, message)
+        if rejection is None:
+            return True
+        if self._failure is None:
+            self._failure = rejection
+        await _reject_guarded_turn(self.websocket, guarded, rejection)
+        return False
+
+    async def _send_turn_refusal(self, refusal: RateLimitError) -> None:
+        try:
+            await self.websocket.send_text(_turn_refusal_frame(refusal))
+        except Exception:  # noqa: BLE001, S110  # client may already be gone
+            pass
+
     async def _send_error_event(self, error_type: str, message: str) -> None:
         try:
             await self.websocket.send_text(
@@ -2684,7 +2750,7 @@ class ResponsesWebSocketStreaming:
                 message = await self.websocket.receive_text()
                 if not await self._enforce_or_reject_frame(message):
                     continue
-                await self._forward_guarded_frame(message)
+                await self._forward_guarded_frame(message, self.turn_admission)
 
         except Exception as e:
             verbose_logger.debug("Responses WS client_to_backend ended: %s", e)
@@ -2761,6 +2827,7 @@ class ManagedResponsesWebSocketHandler:
         custom_llm_provider: str | None = None,
         first_message: str | None = None,
         quota_callbacks: Sequence[ProjectQuotaCallback] | None = None,
+        turn_admission: Callable[[], Awaitable[None]] | None = None,
         frame_guard: ResponsesWsFrameGuard | None = None,
         guardrail_callbacks: Sequence[PresidioGuardrailCallback] | None = None,
         **kwargs: object,
@@ -2784,6 +2851,8 @@ class ManagedResponsesWebSocketHandler:
         self._connection_provider = self._resolve_provider(model) or custom_llm_provider
         self.first_message = first_message
         self.quota_callbacks: tuple[ProjectQuotaCallback, ...] = tuple(quota_callbacks) if quota_callbacks else ()
+        self.turn_admission: Callable[[], Awaitable[None]] | None = turn_admission
+        self._failure: Exception | None = None
         # Carry through safe pass-through kwargs (e.g. extra_headers)
         self.extra_kwargs: dict[str, object] = {k: v for k, v in kwargs.items() if k not in _MANAGED_WS_SKIP_KWARGS}
         # In-memory session history: response_id → full accumulated message list.
@@ -2812,12 +2881,11 @@ class ManagedResponsesWebSocketHandler:
             return None
 
     async def _send_error(self, message: str, error_type: str = "server_error") -> None:
-        try:
-            await self.websocket.send_text(
-                json.dumps({"type": "error", "error": {"type": error_type, "message": message}})
-            )
-        except Exception:
-            pass
+        await self._send_frame(json.dumps({"type": "error", "error": {"type": error_type, "message": message}}))
+
+    async def _send_frame(self, frame: str) -> None:
+        with contextlib.suppress(Exception):
+            await self.websocket.send_text(frame)
 
     def _get_history_messages(self, previous_response_id: str) -> list[dict[str, object]]:
         """
@@ -3179,7 +3247,9 @@ class ManagedResponsesWebSocketHandler:
     # Core request handler
     # ------------------------------------------------------------------
 
-    async def _process_response_create(self, raw_message: str) -> None:
+    async def _process_response_create(
+        self, raw_message: str, turn_admission: Callable[[], Awaitable[None]] | None = None
+    ) -> None:
         """
         Parse one ``response.create`` event, call ``litellm.aresponses(stream=True)``,
         and forward every streaming event to the client.
@@ -3220,12 +3290,17 @@ class ManagedResponsesWebSocketHandler:
                 self.quota_callbacks, self.user_api_key_dict, self.model_group or self.model, raw_message
             )
         except RateLimitError as e:
-            await self._send_error(str(e), error_type="rate_limit_exceeded")
+            await self._send_frame(_turn_refusal_frame(e))
             return
-
         guarded: Final = await self._guard_frame(raw_message)
         if guarded is None:
             return
+        rejection: Final = await _admit_frame_turn(turn_admission, raw_message)
+        if rejection is not None:
+            if self._failure is None:
+                self._failure = rejection
+            await _reject_guarded_turn(self.websocket, guarded, rejection)
+            return
         if guarded.pii_tokens:
             self._pii_tokens = {**self._pii_tokens, **guarded.pii_tokens}
 
@@ -3269,7 +3344,7 @@ class ManagedResponsesWebSocketHandler:
     # Main entry point
     # ------------------------------------------------------------------
 
-    async def run(self) -> None:
+    async def run(self) -> Exception | None:
         """
         Main loop: accept ``response.create`` events sequentially and handle
         each one before waiting for the next message.
@@ -3285,8 +3360,9 @@ class ManagedResponsesWebSocketHandler:
                     verbose_logger.debug("ManagedResponsesWS: client disconnected: %s", exc)
                     break
 
-                await self._process_response_create(message)
+                await self._process_response_create(message, self.turn_admission)
 
         except Exception as exc:
             verbose_logger.exception("ManagedResponsesWS: unexpected error: %s", exc)
             await self._send_error(f"Internal server error: {exc}")
+        return self._failure

litellm/router.py

@@ -289,6 +289,7 @@ from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import (
     deployment_has_io_token_limits,
     refund_stale_reservation_before_retry,
     set_io_token_rate_limit_request_kwargs,
+    without_io_token_limits,
 )
 from litellm.router_utils.pre_call_checks.model_rate_limit_check import (
     ModelRateLimitingCheck,
@@ -5962,6 +5963,7 @@ class Router:
             )
 
             custom_llm_provider: Final = provider_for_generic_call(data)
+            generic_function_name: Final[str] = getattr(original_generic_function, "__name__", "")
 
             response_kwargs: Final = {
                 **data,
@@ -5969,6 +5971,11 @@ class Router:
                 **kwargs,
                 "model": model_name,
                 **_with_router_resolved_session_model(kwargs.get("session"), model_name),
+                **(
+                    {"turn_admission": partial(self.async_admit_deployment_turn, model, deployment, kwargs)}
+                    if generic_function_name == "_aresponses_websocket"
+                    else {}
+                ),
             }
             # Only set custom_llm_provider if it's not None
             if custom_llm_provider is not None:
@@ -5978,7 +5985,7 @@ class Router:
                 self,
                 deployment,
                 response_kwargs,
-                surface_for_call(getattr(original_generic_function, "__name__", "")),
+                surface_for_call(generic_function_name),
             )
             async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span):
                 response = await original_generic_function(**compacted_input)
@@ -8933,8 +8940,7 @@ class Router:
         Track remaining tpm/rpm quota for model in model_list
         """
         try:
-            # WS session wrappers fire with result=None; per-turn costs tracked by inner calls.
-            if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"):
+            if kwargs.get("call_type") == "_arealtime":
                 return
             if is_batch_retrieve_call_type(kwargs.get("call_type")):
                 return
@@ -9602,6 +9608,71 @@ class Router:
                         )
                     raise e
 
+    def _deployment_as_configured_now(self, deployment: Mapping[str, object]) -> Mapping[str, object]:
+        model_info: Final = deployment.get("model_info")
+        model_id: Final = (
+            _MODEL_INFO_ADAPTER.validate_python(model_info).get("id") if isinstance(model_info, dict) else None
+        )
+        current: Final = self.get_deployment(model_id) if isinstance(model_id, str) else None
+        return deployment if current is None else current.model_dump(exclude_none=True)
+
+    async def async_admit_deployment_turn(
+        self, model_group: str, opened_on: Mapping[str, object], request_kwargs: Mapping[str, object]
+    ) -> None:
+        deployment: Final = self._deployment_as_configured_now(opened_on)
+        budget_limiter: Final = self.get_router_deployment_budget_limiter()
+        if budget_limiter is not None:
+            try:
+                await budget_limiter.async_filter_deployments(  # pyright: ignore[reportUnknownMemberType]  # legacy signature takes untyped deployment dicts
+                    model=model_group,
+                    healthy_deployments=[dict(deployment)],
+                    messages=None,
+                    request_kwargs=dict(request_kwargs),
+                )
+            except ValueError as e:
+                raise litellm.RateLimitError(
+                    message=str(e),
+                    llm_provider="",
+                    model=model_group,
+                    category=litellm.RateLimitErrorCategory.LITELLM_RATE_LIMIT,
+                ) from e
+        await self._admit_turn_through_routing_strategy(model_group, deployment, request_kwargs)
+        await self.async_routing_strategy_pre_call_checks(  # pyright: ignore[reportUnknownMemberType]  # legacy signature takes an untyped deployment dict
+            deployment=dict(without_io_token_limits(deployment)), parent_otel_span=None
+        )
+
+    async def _admit_turn_through_routing_strategy(
+        self, model_group: str, deployment: Mapping[str, object], request_kwargs: Mapping[str, object]
+    ) -> None:
+        routing_kwargs: Final = {key: value for key, value in request_kwargs.items() if key != "litellm_logging_obj"}
+        strategy, selector = self._get_routing_context(  # pyright: ignore[reportUnknownMemberType]  # legacy signature takes an untyped request dict
+            model=model_group, request_kwargs=routing_kwargs
+        )
+        if strategy == "simple-shuffle" or selector is None:
+            return
+        selected: Final = await self._select_deployment_async(  # pyright: ignore[reportUnknownMemberType]  # legacy signature takes untyped deployment dicts
+            strategy=strategy,
+            selector=selector,
+            model=model_group,
+            healthy_deployments=[dict(deployment)],
+            messages=None,
+            input=None,
+            request_kwargs=routing_kwargs,
+        )
+        if selected is None:
+            no_deployment: Final = await async_raise_no_deployment_exception(
+                litellm_router_instance=self, model=model_group, parent_otel_span=None
+            )
+            raise litellm.RateLimitError(
+                message=str(no_deployment),
+                llm_provider="",
+                model=model_group,
+                category=litellm.RateLimitErrorCategory.LITELLM_RATE_LIMIT,
+            ) from no_deployment
+        await self._async_override_selector_pre_call_check(  # pyright: ignore[reportUnknownMemberType]  # legacy signature takes an untyped deployment dict
+            strategy, selector, dict(deployment), None
+        )
+
     @contextlib.asynccontextmanager
     async def _deployment_slot(
         self, deployment: dict, kwargs: Mapping[str, object], parent_otel_span: Span | None

litellm/router_strategy/budget_limiter.py

@@ -505,8 +505,7 @@ class RouterBudgetLimiting(CustomLogger):
     async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
         """Original method now uses helper functions"""
         verbose_router_logger.debug("in RouterBudgetLimiting.async_log_success_event")
-        # WS session wrappers fire with result=None; per-turn costs tracked by inner calls.
-        if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"):
+        if kwargs.get("call_type") == "_arealtime":
             return
         standard_logging_payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None)
         if standard_logging_payload is None:

litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py

@@ -23,6 +23,7 @@ from litellm._internal_context import with_service_target
 from litellm._logging import verbose_router_logger
 from litellm.caching.dual_cache import DualCache
 from litellm.litellm_core_utils.token_counter import offload_token_count
+from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping
 from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET
 from litellm.types.router import RouterCacheEnum, RouterErrors
 from litellm.utils import get_utc_datetime
@@ -96,6 +97,24 @@ def deployment_has_io_token_limits(deployment: dict) -> bool:
     return itpm is not None or otpm is not None
 
 
+_IO_TOKEN_LIMIT_KEYS: Final = frozenset({"itpm", "otpm"})
+_IO_TOKEN_LIMIT_SECTIONS: Final = frozenset({"litellm_params", "model_info"})
+
+
+def without_io_token_limits(deployment: Mapping[str, object]) -> Mapping[str, object]:
+    return {
+        key: _without_io_token_limit_keys(value) if key in _IO_TOKEN_LIMIT_SECTIONS else value
+        for key, value in deployment.items()
+        if key not in _IO_TOKEN_LIMIT_KEYS
+    }
+
+
+def _without_io_token_limit_keys(section: object) -> object:
+    if not is_str_keyed_mapping(section):
+        return section
+    return {key: value for key, value in section.items() if key not in _IO_TOKEN_LIMIT_KEYS}
+
+
 def _get_cache_keys(deployment: dict, current_minute: str) -> tuple[str, str] | None:
     model_id: Final = deployment.get("model_info", {}).get("id")
     deployment_name: Final = deployment.get("litellm_params", {}).get("model")

tests/code_coverage_tests/router_code_coverage.py

@@ -134,6 +134,9 @@ ignored_function_names = [
     "_provider_model_id",  # Tested through get_router_model_info in test_router_model_cost_isolation.py
     "_names_route_model",  # Tested through get_router_model_info in test_router_model_cost_isolation.py
     "_declared_base_model",  # Tested through get_router_model_info in test_router_model_cost_isolation.py
+    "async_admit_deployment_turn",  # Bound as the WebSocket turn_admission partial, tested through it in test_router_helper_utils.py
+    "_deployment_as_configured_now",  # Read by async_admit_deployment_turn, tested through turn_admission in test_router_helper_utils.py
+    "_admit_turn_through_routing_strategy",  # Run by async_admit_deployment_turn, tested through turn_admission in test_router_helper_utils.py
 ]
 
 

tests/integration/routing/test_responses_websocket_turn_admission.py

@@ -0,0 +1,629 @@
+import asyncio
+import hashlib
+import os
+import threading
+import uuid
+from collections.abc import AsyncIterator, Callable, Iterator, Mapping
+from contextlib import asynccontextmanager, contextmanager
+from dataclasses import dataclass
+from functools import partial
+from datetime import UTC, datetime
+from pathlib import Path
+from types import MappingProxyType
+from typing import Final
+
+import httpx
+import pytest
+import websockets
+import yaml
+from integration._support.client import (
+    JSON_OBJECT,
+    Gateway,
+    Scenario,
+    eventually,
+    gateway_from_environment,
+    held,
+    list_value,
+    object_value,
+    string_value,
+)
+from integration._support.database import read_rows
+from integration._support.process import graceful_stop_seconds, owned_proxy
+from integration._support.responses_ws_peer import (
+    PROVIDER_MODEL,
+    USAGE,
+    Peer,
+    Script,
+    completed,
+    create,
+    drain,
+    events,
+    responses_peer,
+    turn,
+    until_terminal,
+)
+from integration._support.tls import write_self_signed_cert
+from pydantic import JsonValue
+from redis import Redis
+
+pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 240)
+
+PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml"
+INPUT_RATE: Final = 0.001
+OUTPUT_RATE: Final = 0.002
+TURN_SPEND: Final = USAGE["input_tokens"] * INPUT_RATE + USAGE["output_tokens"] * OUTPUT_RATE
+TURN_TOKENS: Final = USAGE["total_tokens"]
+MANAGED_PROVIDER: Final = "mistral"
+SPEND_ROWS_FOR_KEY: Final = (
+    'SELECT request_id, spend, status, model_group, metadata, request_tags FROM "LiteLLM_SpendLogs" WHERE api_key = %s'
+)
+TURN_CHECK_FAULT: Final = "turn_check_fault"
+FAULTING_TURN_CHECK: Final = f"""
+from litellm.integrations.custom_logger import CustomLogger
+
+
+class FaultingTurnCheck(CustomLogger):
+    async def async_pre_call_check(self, deployment, parent_otel_span):
+        fault = (deployment.get("model_info") or {{}}).get("{TURN_CHECK_FAULT}")
+        if fault:
+            raise RuntimeError(fault)
+        return None
+
+
+instance = FaultingTurnCheck()
+"""
+FrameReader = Callable[[], int]
+
+
+@dataclass(frozen=True, slots=True)
+class Deployment:
+    name: str
+    identity: str
+    forwarded_turns: FrameReader
+
+
+def _key_hash(key: str) -> str:
+    return hashlib.sha256(key.encode()).hexdigest()
+
+
+def _redis() -> Redis:
+    return Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]), decode_responses=True)
+
+
+def _config(directory: Path, router_settings: Mapping[str, JsonValue], callbacks: tuple[str, ...]) -> Path:
+    base: Final = JSON_OBJECT.validate_python(yaml.safe_load(PROXY_CONFIG.read_text()))
+    path: Final = directory / "turn-admission.yaml"
+    path.write_text(
+        yaml.safe_dump(
+            {
+                **base,
+                "litellm_settings": {
+                    **object_value(base["litellm_settings"]),
+                    **({"callbacks": list(callbacks)} if callbacks else {}),
+                },
+                "router_settings": {
+                    **object_value(base["router_settings"]),
+                    "redis_host": os.environ["REDIS_HOST"],
+                    "redis_port": int(os.environ["REDIS_PORT"]),
+                    **router_settings,
+                },
+            }
+        )
+    )
+    return path
+
+
+@contextmanager
+def _admitting_proxy(
+    tmp_path_factory: pytest.TempPathFactory,
+    cert: tuple[Path, Path],
+    router_settings: Mapping[str, JsonValue],
+    callback_sources: Mapping[str, str] = MappingProxyType({}),
+) -> Iterator[Gateway]:
+    directory: Final = tmp_path_factory.mktemp("responses-ws-turn-admission")
+    for module, source in callback_sources.items():
+        (directory / f"{module}.py").write_text(source)
+    config: Final = _config(directory, router_settings, tuple(f"{module}.instance" for module in callback_sources))
+    with gateway_from_environment() as base:
+        with owned_proxy(base, directory, {"SSL_CERT_FILE": str(cert[0])}, config=config) as proxy:
+            yield proxy
+
+
+@pytest.fixture(scope="module")
+def cert(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]:
+    return write_self_signed_cert(tmp_path_factory.mktemp("responses-ws-turn-admission-cert"))
+
+
+@pytest.fixture(scope="module")
+def limited(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) -> Iterator[Gateway]:
+    with _admitting_proxy(tmp_path_factory, cert, {"optional_pre_call_checks": ["enforce_model_rate_limits"]}) as proxy:
+        yield proxy
+
+
+@pytest.fixture(scope="module")
+def faulting(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) -> Iterator[Gateway]:
+    with _admitting_proxy(tmp_path_factory, cert, {}, {"faulting_turn_check": FAULTING_TURN_CHECK}) as proxy:
+        yield proxy
+
+
+@pytest.fixture(scope="module")
+def usage_routed(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) -> Iterator[Gateway]:
+    with _admitting_proxy(tmp_path_factory, cert, {"routing_strategy": "usage-based-routing-v2"}) as proxy:
+        yield proxy
+
+
+def _native(scenario: Scenario, peer: Peer, **limits: JsonValue) -> Deployment:
+    identity: Final = f"ws-admission-{uuid.uuid4().hex}"
+    name: Final = scenario.model(
+        model_info={"id": identity},
+        model=f"openai/{PROVIDER_MODEL}",
+        api_base=peer.url,
+        input_cost_per_token=INPUT_RATE,
+        output_cost_per_token=OUTPUT_RATE,
+        **limits,
+    )
+    return Deployment(
+        name, identity, lambda: sum(frame.get("type") == "response.create" for frame in drain(peer.frames))
+    )
+
+
+def _upstream_chats(gateway: Gateway, provider_model: str) -> int:
+    with httpx.Client(base_url=gateway.upstream_url, timeout=15, trust_env=False) as upstream:
+        observed: Final = JSON_OBJECT.validate_json(upstream.get("/__observations").content)["requests"]
+    assert isinstance(observed, list)
+    return sum(object_value(object_value(entry)["body"]).get("model") == provider_model for entry in observed)
+
+
+def _managed(scenario: Scenario, **limits: JsonValue) -> Deployment:
+    identity: Final = f"ws-admission-{uuid.uuid4().hex}"
+    provider_model: Final = f"ws-managed-{uuid.uuid4().hex}"
+    name: Final = scenario.model(
+        model_info={"id": identity},
+        model=f"{MANAGED_PROVIDER}/{provider_model}",
+        input_cost_per_token=INPUT_RATE,
+        output_cost_per_token=OUTPUT_RATE,
+        **limits,
+    )
+    return Deployment(name, identity, lambda: _upstream_chats(scenario.gateway, provider_model))
+
+
+def _socket_url(gateway: Gateway, model: str) -> str:
+    return f"{str(gateway.client.base_url).rstrip('/').replace('http://', 'ws://')}/v1/responses?model={model}"
+
+
+@asynccontextmanager
+async def _socket(
+    gateway: Gateway, key: str, model: str, headers: Mapping[str, str] | None = None
+) -> AsyncIterator[websockets.ClientConnection]:
+    async with websockets.connect(
+        _socket_url(gateway, model),
+        additional_headers={"Authorization": f"Bearer {key}", **(headers or {})},
+        open_timeout=10,
+    ) as connection:
+        yield connection
+
+
+async def _say(connection: websockets.ClientConnection, model: str, text: str = "") -> tuple[dict[str, JsonValue], ...]:
+    return await turn(connection, create(model, text or f"turn {uuid.uuid4().hex}"))
+
+
+def _assert_refused(received: tuple[dict[str, JsonValue], ...]) -> None:
+    assert [event.get("type") for event in received] == ["error"], received
+    assert received[0]["status"] == 429, received
+    assert object_value(received[0]["error"])["type"] == "rate_limit_exceeded", received
+
+
+def _fresh_minute() -> str:
+    eventually(lambda: datetime.now(UTC), lambda now: now.second < 30, seconds=65)
+    return datetime.now(UTC).strftime("%H-%M")
+
+
+def _deployment_spend(deployment: Deployment) -> float:
+    with _redis() as redis:
+        return float(redis.get(f"deployment_spend:{deployment.identity}:1d") or 0)
+
+
+def _configured(gateway: Gateway, deployment: Deployment, section: str) -> dict[str, JsonValue]:
+    entries: Final = list_value(gateway.get("/model/info", {"litellm_model_id": deployment.identity})["data"])
+    return object_value(object_value(entries[0])[section])
+
+
+def _configured_tpm(gateway: Gateway, deployment: Deployment) -> JsonValue:
+    return _configured(gateway, deployment, "litellm_params").get("tpm")
+
+
+def _tokens_counted(identity: str, minute: str) -> int:
+    with _redis() as redis:
+        return sum(int(redis.get(key) or 0) for key in redis.scan_iter(f"{identity}:*:tpm:{minute}"))
+
+
+def _settled_count(read: Callable[[], int]) -> int:
+    reached: Final = eventually(read, lambda tokens: tokens > 0, seconds=15)
+    return held(read, lambda tokens: tokens == reached, holding=3, seconds=20)
+
+
+async def _turns_against_tpm(
+    gateway: Gateway, key: str, deployment: Deployment
+) -> tuple[tuple[dict[str, JsonValue], ...], tuple[dict[str, JsonValue], ...], tuple[dict[str, JsonValue], ...], int]:
+    minute: Final = await asyncio.to_thread(_fresh_minute)
+    tokens: Final = partial(_tokens_counted, deployment.identity, minute)
+    async with _socket(gateway, key, deployment.name) as connection:
+        first: Final = await _say(connection, deployment.name)
+        counted_after_first: Final = await asyncio.to_thread(_settled_count, tokens)
+        second: Final = await _say(connection, deployment.name)
+        if second[-1].get("type") == "response.completed":
+            await _until(tokens, lambda counted: counted > counted_after_first)
+        return first, second, await _say(connection, deployment.name), counted_after_first
+
+
+async def _until(read: Callable[[], float], satisfied: Callable[[float], bool]) -> float:
+    return await asyncio.to_thread(eventually, read, satisfied, 15)
+
+
+def _failure_rows(key: str) -> list[dict[str, JsonValue]]:
+    rows: Final = eventually(
+        lambda: read_rows(SPEND_ROWS_FOR_KEY, (_key_hash(key),)),
+        lambda found: any(row["status"] == "failure" for row in found),
+        seconds=70,
+    )
+    return [row for row in rows if row["status"] == "failure"]
+
+
+@pytest.mark.parametrize("relay", ("native", "managed"))
+def test_a_turn_past_the_deployment_budget_is_refused_on_the_open_socket_and_logged_once_as_its_failure(
+    limited: Gateway, cert: tuple[Path, Path], relay: str
+) -> None:
+    with responses_peer(cert) as peer, limited.scenario() as scenario:
+        budget: Final = {"max_budget": TURN_SPEND / 2, "budget_duration": "1d"}
+        deployment: Final = _native(scenario, peer, **budget) if relay == "native" else _managed(scenario, **budget)
+        key: Final = scenario.key(models=[deployment.name])
+        tag: Final = f"ws-turn-{uuid.uuid4().hex}"
+
+        async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]:
+            async with _socket(limited, key, deployment.name, {"x-litellm-tags": tag}) as connection:
+                first: Final = await _say(connection, deployment.name)
+                await _until(lambda: _deployment_spend(deployment), lambda spent: spent > TURN_SPEND / 2)
+                return (first, await _say(connection, deployment.name), await _say(connection, deployment.name))
+
+        first, second, third = asyncio.run(session())
+
+        completed(first)
+        _assert_refused(second)
+        _assert_refused(third)
+        assert deployment.forwarded_turns() == 1
+        failures: Final = _failure_rows(key)
+        assert len(failures) == 1, failures
+        assert failures[0]["model_group"] == deployment.name, failures
+        assert tag in list_value(failures[0]["request_tags"]), failures
+        assert float(failures[0]["spend"]) == 0.0, failures
+        error: Final = object_value(object_value(failures[0]["metadata"])["error_information"])
+        assert str(error["error_code"]) == "429", error
+
+
+def _set_turn_check_fault(gateway: Gateway, deployment: Deployment, fault: str) -> None:
+    gateway.request(
+        "PATCH",
+        f"/model/{deployment.identity}/update",
+        {"model_info": {"id": deployment.identity, TURN_CHECK_FAULT: fault}},
+    ).raise_for_status()
+    eventually(
+        lambda: _configured(gateway, deployment, "model_info").get(TURN_CHECK_FAULT) or "",
+        lambda configured: configured == fault,
+        seconds=15,
+    )
+
+
+def _assert_turn_check_failed(received: tuple[dict[str, JsonValue], ...], fault: str) -> None:
+    assert [event.get("type") for event in received] == ["error"], received
+    assert received[0]["status"] == 500, received
+    error: Final = object_value(received[0]["error"])
+    assert error["type"] == "server_error", received
+    assert fault in string_value(error["message"]), received
+
+
+def _first_turn_held_until(release: threading.Event) -> Script:
+    def script(response_id: str) -> tuple[dict[str, JsonValue], ...]:
+        if response_id == "resp_peer_1":
+            release.wait(timeout=60)
+        return events(f"resp_{uuid.uuid4().hex}", "seven")
+
+    return script
+
+
+@pytest.mark.parametrize("relay", ("native", "managed"))
+def test_a_turn_whose_admission_check_fails_gets_a_500_frame_and_the_socket_serves_the_next_turn(
+    faulting: Gateway, cert: tuple[Path, Path], relay: str
+) -> None:
+    with responses_peer(cert) as peer, faulting.scenario() as scenario:
+        deployment: Final = _native(scenario, peer) if relay == "native" else _managed(scenario)
+        key: Final = scenario.key(models=[deployment.name])
+        fault: Final = f"turn check unavailable {uuid.uuid4().hex}"
+
+        async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]:
+            async with _socket(faulting, key, deployment.name) as connection:
+                first: Final = await _say(connection, deployment.name)
+                await asyncio.to_thread(_set_turn_check_fault, faulting, deployment, fault)
+                second: Final = await _say(connection, deployment.name)
+                third: Final = await _say(connection, deployment.name)
+                await asyncio.to_thread(_set_turn_check_fault, faulting, deployment, "")
+                return first, second, third, await _say(connection, deployment.name)
+
+        first, second, third, fourth = asyncio.run(session())
+
+        completed(first)
+        _assert_turn_check_failed(second, fault)
+        _assert_turn_check_failed(third, fault)
+        completed(fourth)
+        assert deployment.forwarded_turns() == 2
+        failures: Final = _failure_rows(key)
+        assert len(failures) == 1, failures
+        assert failures[0]["model_group"] == deployment.name, failures
+        recorded: Final = object_value(object_value(failures[0]["metadata"])["error_information"])
+        assert fault in string_value(recorded["error_message"]), recorded
+
+
+def test_a_turn_still_streaming_when_the_next_turns_check_fails_completes_and_is_billed(
+    faulting: Gateway, cert: tuple[Path, Path]
+) -> None:
+    release: Final = threading.Event()
+    with responses_peer(cert, _first_turn_held_until(release)) as peer, faulting.scenario() as scenario:
+        deployment: Final = _native(scenario, peer)
+        key: Final = scenario.key(models=[deployment.name])
+        fault: Final = f"turn check unavailable {uuid.uuid4().hex}"
+
+        async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]:
+            async with _socket(faulting, key, deployment.name) as connection:
+                try:
+                    await connection.send(create(deployment.name, f"streaming {uuid.uuid4().hex}"))
+                    await asyncio.to_thread(_set_turn_check_fault, faulting, deployment, fault)
+                    rejected: Final = await _say(connection, deployment.name)
+                finally:
+                    release.set()
+                return rejected, await until_terminal(connection, ())
+
+        rejected, streamed = asyncio.run(session())
+
+        _assert_turn_check_failed(rejected, fault)
+        billed: Final = completed(streamed)
+        assert deployment.forwarded_turns() == 1
+        rows: Final = eventually(
+            lambda: read_rows(SPEND_ROWS_FOR_KEY, (_key_hash(key),)),
+            lambda found: any(row["request_id"] == billed["id"] for row in found),
+            seconds=70,
+        )
+        row: Final = next(row for row in rows if row["request_id"] == billed["id"])
+        assert row["status"] == "success", row
+        assert float(row["spend"]) == pytest.approx(TURN_SPEND), row
+
+
+def test_a_socket_opened_after_websocket_turns_spent_the_budget_is_refused_and_so_is_http(
+    limited: Gateway, cert: tuple[Path, Path]
+) -> None:
+    with responses_peer(cert) as peer, limited.scenario() as scenario:
+        deployment: Final = _native(scenario, peer, max_budget=TURN_SPEND / 2, budget_duration="1d")
+        key: Final = scenario.key(models=[deployment.name])
+
+        async def spend_then_reopen() -> tuple[tuple[dict[str, JsonValue], ...], int | None]:
+            async with _socket(limited, key, deployment.name) as spending:
+                completed(await _say(spending, deployment.name))
+            await _until(lambda: _deployment_spend(deployment), lambda spent: spent > TURN_SPEND / 2)
+            async with _socket(limited, key, deployment.name) as refused:
+                received: Final = await _say(refused, deployment.name)
+                await asyncio.wait_for(refused.wait_closed(), timeout=20)
+                return received, refused.close_code
+
+        received, close_code = asyncio.run(spend_then_reopen())
+
+        _assert_refused(received)
+        assert close_code == 1011, close_code
+        assert peer.paths.qsize() == 1
+        assert deployment.forwarded_turns() == 1
+        http: Final = limited.request(
+            "POST", "/v1/responses", {"model": deployment.name, "input": f"after {uuid.uuid4().hex}"}, key=key
+        )
+        assert http.status_code == 429, http.text
+        assert peer.paths.qsize() == 1
+
+
+def test_a_socket_the_router_refuses_to_open_gets_the_status_and_message_http_returns(
+    limited: Gateway, cert: tuple[Path, Path]
+) -> None:
+    with responses_peer(cert) as peer, limited.scenario() as scenario:
+        name: Final = scenario.model(
+            model=f"openai/{PROVIDER_MODEL}",
+            api_base=peer.url,
+            tags=[f"granted-{uuid.uuid4().hex}"],
+            model_info={"enable_tag_filtering": True},
+        )
+        key: Final = scenario.key(models=[name])
+        foreign: Final = {"x-litellm-tags": f"foreign-{uuid.uuid4().hex}"}
+        http: Final = limited.request(
+            "POST", "/v1/responses", {"model": name, "input": "tag routing"}, key=key, headers=foreign
+        )
+        assert http.status_code != 200, http.text
+        http_error: Final = object_value(JSON_OBJECT.validate_json(http.content)["error"])
+
+        async def refused_open() -> tuple[dict[str, JsonValue], ...]:
+            async with _socket(limited, key, name, foreign) as connection:
+                return await _say(connection, name)
+
+        received: Final = asyncio.run(refused_open())
+
+        assert [event.get("type") for event in received] == ["error"], received
+        assert received[0]["status"] == http.status_code, (received, http.text)
+        assert object_value(received[0]["error"])["message"] == string_value(http_error["message"]), (
+            received,
+            http.text,
+        )
+        assert peer.paths.qsize() == 0
+
+
+def test_each_turn_counts_as_one_request_toward_the_deployment_rpm(limited: Gateway, cert: tuple[Path, Path]) -> None:
+    with responses_peer(cert) as peer, limited.scenario() as scenario:
+        deployment: Final = _native(scenario, peer, rpm=1)
+        key: Final = scenario.key(models=[deployment.name])
+        _fresh_minute()
+
+        async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]:
+            async with _socket(limited, key, deployment.name) as connection:
+                return (await _say(connection, deployment.name), await _say(connection, deployment.name))
+
+        first, second = asyncio.run(session())
+
+        completed(first)
+        _assert_refused(second)
+        assert deployment.forwarded_turns() == 1
+
+
+def test_a_deployment_tpm_refuses_a_turn_only_once_earlier_turns_used_it_up(
+    limited: Gateway, cert: tuple[Path, Path]
+) -> None:
+    with responses_peer(cert) as peer, limited.scenario() as scenario:
+        deployment: Final = _native(scenario, peer, tpm=TURN_TOKENS + 2)
+        key: Final = scenario.key(models=[deployment.name])
+
+        first, second, third, counted_after_first = asyncio.run(_turns_against_tpm(limited, key, deployment))
+
+        completed(first)
+        completed(second)
+        _assert_refused(third)
+        assert counted_after_first == TURN_TOKENS
+        assert deployment.forwarded_turns() == 2
+
+
+def test_a_tpm_limit_lowered_while_the_socket_is_open_applies_to_its_next_turn(
+    limited: Gateway, cert: tuple[Path, Path]
+) -> None:
+    with responses_peer(cert) as peer, limited.scenario() as scenario:
+        deployment: Final = _native(scenario, peer, tpm=100 * TURN_TOKENS)
+        key: Final = scenario.key(models=[deployment.name])
+        minute: Final = _fresh_minute()
+
+        def lower_tpm() -> None:
+            limited.request(
+                "PATCH", f"/model/{deployment.identity}/update", {"litellm_params": {"tpm": TURN_TOKENS - 2}}
+            ).raise_for_status()
+            eventually(lambda: _configured_tpm(limited, deployment), lambda tpm: tpm == TURN_TOKENS - 2, seconds=15)
+
+        async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]:
+            async with _socket(limited, key, deployment.name) as connection:
+                first: Final = await _say(connection, deployment.name)
+                await _until(partial(_tokens_counted, deployment.identity, minute), lambda tokens: tokens > 0)
+                await asyncio.to_thread(lower_tpm)
+                return (first, await _say(connection, deployment.name))
+
+        first, second = asyncio.run(session())
+
+        completed(first)
+        _assert_refused(second)
+        assert deployment.forwarded_turns() == 1
+
+
+def test_a_socket_on_an_otpm_limited_deployment_reserves_once_and_serves_every_turn(
+    limited: Gateway, cert: tuple[Path, Path]
+) -> None:
+    with responses_peer(cert) as peer, limited.scenario() as scenario:
+        deployment: Final = _native(scenario, peer, otpm=10_000)
+        key: Final = scenario.key(models=[deployment.name])
+        _fresh_minute()
+
+        async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]:
+            async with _socket(limited, key, deployment.name) as connection:
+                return tuple([await _say(connection, deployment.name) for _ in range(5)])
+
+        turns: Final = asyncio.run(session())
+
+        assert [completed(received)["status"] for received in turns] == ["completed"] * 5
+        assert deployment.forwarded_turns() == 5
+
+
+def test_usage_based_routing_refuses_a_turn_only_once_earlier_turns_used_up_the_deployment_tpm(
+    usage_routed: Gateway, cert: tuple[Path, Path]
+) -> None:
+    with responses_peer(cert) as peer, usage_routed.scenario() as scenario:
+        deployment: Final = _native(scenario, peer, tpm=TURN_TOKENS + 2)
+        key: Final = scenario.key(models=[deployment.name])
+
+        first, second, third, counted_after_first = asyncio.run(_turns_against_tpm(usage_routed, key, deployment))
+
+        completed(first)
+        completed(second)
+        _assert_refused(third)
+        assert counted_after_first == TURN_TOKENS
+        assert deployment.forwarded_turns() == 2
+
+
+def _assert_http_chat_counted_once(gateway: Gateway) -> None:
+    with gateway.scenario() as scenario:
+        identity: Final = f"ws-admission-{uuid.uuid4().hex}"
+        name: Final = scenario.model(model_info={"id": identity}, tpm=1_000_000)
+        key: Final = scenario.key(models=[name])
+        minute: Final = _fresh_minute()
+        answered: Final = gateway.request(
+            "POST",
+            "/v1/chat/completions",
+            {"model": name, "messages": [{"role": "user", "content": f"http turn {uuid.uuid4().hex}"}]},
+            key=key,
+        )
+        assert answered.status_code == 200, answered.text
+        tokens: Final = object_value(JSON_OBJECT.validate_json(answered.content)["usage"])["total_tokens"]
+        assert isinstance(tokens, int) and tokens > 0, answered.text
+
+        assert _settled_count(partial(_tokens_counted, identity, minute)) == tokens
+
+
+def test_an_http_call_counts_its_tokens_once_toward_the_deployment_tpm_under_model_rate_limits(
+    limited: Gateway,
+) -> None:
+    _assert_http_chat_counted_once(limited)
+
+
+def test_an_http_call_counts_its_tokens_once_toward_the_deployment_tpm_under_usage_based_routing(
+    usage_routed: Gateway,
+) -> None:
+    _assert_http_chat_counted_once(usage_routed)
+
+
+def test_usage_based_routing_serves_every_turn_while_the_deployment_has_headroom(
+    usage_routed: Gateway, cert: tuple[Path, Path]
+) -> None:
+    with responses_peer(cert) as peer, usage_routed.scenario() as scenario:
+        deployment: Final = _native(scenario, peer, tpm=1_000_000, rpm=1_000)
+        key: Final = scenario.key(models=[deployment.name])
+
+        async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]:
+            async with _socket(usage_routed, key, deployment.name) as connection:
+                return tuple([await _say(connection, deployment.name) for _ in range(3)])
+
+        turns: Final = asyncio.run(session())
+
+        assert [completed(received)["status"] for received in turns] == ["completed"] * 3
+        assert deployment.forwarded_turns() == 3
+
+
+def test_a_frame_over_the_project_input_quota_gets_a_429_error_frame_and_never_reaches_the_provider(
+    limited: Gateway,
+) -> None:
+    with limited.scenario() as scenario:
+        deployment: Final = _managed(scenario)
+        team: Final = scenario.team()
+        key: Final = scenario.key(
+            models=[deployment.name],
+            team_id=team,
+            project_id=scenario.project(team, model_itpm_limit={deployment.name: 100}),
+        )
+
+        async def session() -> tuple[tuple[dict[str, JsonValue], ...], ...]:
+            async with _socket(limited, key, deployment.name) as connection:
+                return (
+                    await _say(connection, deployment.name, "short"),
+                    await _say(connection, deployment.name, "long " * 400),
+                    await _say(connection, deployment.name, "long " * 400),
+                )
+
+        first, second, third = asyncio.run(session())
+
+        completed(first)
+        _assert_refused(third)
+        answered: Final = sum(received[-1].get("type") == "response.completed" for received in (first, second, third))
+        assert deployment.forwarded_turns() == answered

tests/integration/spend/test_failed_dispatch_tokens.py

@@ -1,11 +1,13 @@
 import json
 import uuid
+from collections.abc import Callable, Mapping
 from typing import Final
 
 import pytest
-from integration._support.client import Gateway, eventually
+from integration._support.client import Gateway, eventually, list_value
 from integration._support.database import read_rows
 from integration._support.wire import Reply, Request, wire_server
+from pydantic import JsonValue
 
 
 @pytest.mark.covers("spend.failed_dispatch.failure_row_records_estimated_input_tokens")
@@ -48,3 +50,79 @@ def test_provider_500_after_dispatch_records_estimated_prompt_tokens_on_failure_
         assert row["status"] == "failure" and float(row["spend"]) == 0 and row["completion_tokens"] == 0, row
         assert row["prompt_tokens"] > 0, f"failure row lost the dispatched input tokens: {row}"
         assert row["total_tokens"] == row["prompt_tokens"], row
+
+
+def _chat_body(model: str, text: str) -> Mapping[str, JsonValue]:
+    return {"model": model, "messages": [{"role": "user", "content": text}]}
+
+
+def _responses_body(model: str, text: str) -> Mapping[str, JsonValue]:
+    return {"model": model, "input": text}
+
+
+def _messages_body(model: str, text: str) -> Mapping[str, JsonValue]:
+    return {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": text}]}
+
+
+FAILURE_ROUTES: Final = pytest.mark.parametrize(
+    ("route", "body"),
+    (
+        ("/v1/chat/completions", _chat_body),
+        ("/v1/responses", _responses_body),
+        ("/v1/messages", _messages_body),
+    ),
+    ids=("chat", "responses", "messages"),
+)
+
+
+def _provider_500_failure_row(
+    gateway: Gateway,
+    route: str,
+    body: Callable[[str, str], Mapping[str, JsonValue]],
+    headers: Mapping[str, str] | None = None,
+) -> tuple[str, dict[str, JsonValue]]:
+    def provider(request: Request) -> Reply:
+        assert request.method == "POST"
+        return Reply(
+            status=500,
+            body=b'{"error":{"message":"synthetic provider outage","type":"server_error","code":"500"}}',
+        )
+
+    with wire_server(provider) as wire, gateway.scenario() as scenario:
+        model: Final = scenario.model(api_base=wire.url + "/v1", num_retries=0)
+        key: Final = scenario.key(models=[model])
+        failed: Final = gateway.request(
+            "POST", route, body(model, f"failed dispatch group {uuid.uuid4().hex}"), key=key, headers=headers
+        )
+        assert failed.status_code == 500, failed.text
+        assert wire.drain()
+        rows: Final = eventually(
+            lambda: read_rows(
+                'SELECT status, model_group, request_tags FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
+                (failed.headers["x-litellm-call-id"],),
+            ),
+            lambda values: len(values) == 1,
+            seconds=70,
+        )
+        assert rows[0]["status"] == "failure", rows
+        return model, rows[0]
+
+
+@FAILURE_ROUTES
+def test_a_provider_500_failure_row_names_the_model_group_the_caller_asked_for(
+    gateway: Gateway, route: str, body: Callable[[str, str], Mapping[str, JsonValue]]
+) -> None:
+    model, row = _provider_500_failure_row(gateway, route, body)
+
+    assert row["model_group"] == model, row
+
+
+@FAILURE_ROUTES
+def test_a_provider_500_failure_row_keeps_the_tags_the_caller_sent(
+    gateway: Gateway, route: str, body: Callable[[str, str], Mapping[str, JsonValue]]
+) -> None:
+    tag: Final = f"failed-dispatch-{uuid.uuid4().hex}"
+
+    _, row = _provider_500_failure_row(gateway, route, body, {"x-litellm-tags": tag})
+
+    assert tag in list_value(row["request_tags"]), row

tests/unit/litellm_core_utils/test_litellm_logging.py

@@ -2066,6 +2066,56 @@ async def test_arealtime_marks_litellm_params_async(monkeypatch):
     assert LitellmLogging.is_sync_litellm_request(captured["litellm_params"]) is False
 
 
+@pytest.mark.asyncio
+async def test_aresponses_websocket_turn_reaches_a_custom_logger_through_the_async_hook_only(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    from litellm.litellm_core_utils import thread_pool_executor
+    from litellm.responses.main import base_llm_http_handler
+
+    class InlineExecutor:
+        def submit(self, fn, *args, **kwargs):
+            fn(*args, **kwargs)
+
+    sessions: list[LitellmLogging] = []
+
+    async def open_session(**kwargs):
+        sessions.append(kwargs["logging_obj"])
+
+    async_logged = asyncio.Event()
+    sync_call_ids: list[str] = []
+    async_call_ids: list[str] = []
+
+    class CaptureLogger(CustomLogger):
+        def log_success_event(self, kwargs, response_obj, start_time, end_time):
+            sync_call_ids.append(kwargs["litellm_call_id"])
+
+        async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+            async_call_ids.append(kwargs["litellm_call_id"])
+            async_logged.set()
+
+    logger = CaptureLogger()
+    monkeypatch.setattr(litellm, "callbacks", [logger])
+    monkeypatch.setattr(litellm, "success_callback", ["cache", logger])
+    monkeypatch.setattr(litellm, "_async_success_callback", [logger])
+    monkeypatch.setattr(thread_pool_executor, "executor", InlineExecutor())
+    with patch.object(  # test-quality-ok: the provider socket is the seam; how a turn of the session logs is under test
+        base_llm_http_handler, "async_responses_websocket", AsyncMock(side_effect=open_session)
+    ):
+        await litellm._aresponses_websocket(model="openai/gpt-4o", websocket=MagicMock(), api_key="sk-test")
+    turn = sessions[0].copy_for_new_call("ws-turn-2", dt_object.now())
+
+    await turn.dispatch_success_handlers(
+        litellm.ResponsesAPIResponse(
+            id="resp_turn_2", created_at=0, output=[], parallel_tool_calls=False, tool_choice="auto", tools=[]
+        ),
+        prefer_async_handlers=True,
+    )
+    await asyncio.wait_for(async_logged.wait(), timeout=10)
+
+    assert (async_call_ids, sync_call_ids) == (["ws-turn-2"], [])
+
+
 @pytest.mark.asyncio
 async def test_aresponses_websocket_hands_back_the_provider_failure_without_a_success_log(
     monkeypatch: pytest.MonkeyPatch,

tests/unit/proxy/hooks/test_proxy_track_cost_callback.py

@@ -4,7 +4,7 @@ import logging
 from datetime import datetime, timedelta, timezone
 from types import SimpleNamespace
 from typing import Final
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, create_autospec, patch
 
 import httpx
 import pytest
@@ -136,6 +136,77 @@ async def test_async_post_call_failure_hook_carries_guardrail_info_from_litellm_
         assert metadata["original_key"] == "original_value"
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("metadata_buckets", "expected_model_group"),
+    [
+        ({"metadata": {}, "litellm_metadata": {"model_group": "ws-budget-sock"}}, "ws-budget-sock"),
+        ({"metadata": {"model_group": "caller-group"}, "litellm_metadata": {"model_group": "routed"}}, "caller-group"),
+        ({"metadata": {}, "litellm_metadata": {}}, ""),
+    ],
+)
+async def test_async_post_call_failure_hook_logs_the_routed_model_group_on_the_failure_row(
+    metadata_buckets: dict[str, dict[str, str]], expected_model_group: str
+):
+    request_data: Final = {
+        "model": "ws-budget-sock",
+        "input": "Reply with exactly the word two",
+        "proxy_server_request": {"request_id": "test_request_id"},
+        **metadata_buckets,
+    }
+
+    spend_writer: Final = create_autospec(DBSpendUpdateWriter, instance=True)
+    await ProxyDBLogger(spend_writer=lambda: spend_writer).async_post_call_failure_hook(
+        request_data=request_data,
+        original_exception=litellm.RateLimitError(message="crossed budget", llm_provider="", model="ws-budget-sock"),
+        user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key"),
+    )
+
+    row: Final = get_logging_payload(
+        kwargs=spend_writer.update_database.call_args.kwargs["kwargs"],
+        response_obj={},
+        start_time=datetime.now(),
+        end_time=datetime.now(),
+    )
+    assert row["model_group"] == expected_model_group
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("metadata_buckets", "expected_tags"),
+    [
+        ({"metadata": {}, "litellm_metadata": {"tags": ["ws-tag"]}}, ["ws-tag"]),
+        ({"metadata": {"tags": ["caller-tag"]}, "litellm_metadata": {"tags": ["routed-tag"]}}, ["caller-tag"]),
+        ({"metadata": {}, "litellm_metadata": {"tags": "not-a-list"}}, []),
+        ({"metadata": {}, "litellm_metadata": {"tags": ["ws-tag", 1, {"a": 1}]}}, []),
+    ],
+)
+async def test_async_post_call_failure_hook_logs_the_routed_tags_on_the_failure_row(
+    metadata_buckets: dict[str, dict[str, object]], expected_tags: list[str]
+):
+    request_data: Final = {
+        "model": "ws-budget-sock",
+        "input": "Reply with exactly the word two",
+        "proxy_server_request": {"request_id": "test_request_id"},
+        **metadata_buckets,
+    }
+
+    spend_writer: Final = create_autospec(DBSpendUpdateWriter, instance=True)
+    await ProxyDBLogger(spend_writer=lambda: spend_writer).async_post_call_failure_hook(
+        request_data=request_data,
+        original_exception=litellm.RateLimitError(message="crossed budget", llm_provider="", model="ws-budget-sock"),
+        user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key"),
+    )
+
+    row: Final = get_logging_payload(
+        kwargs=spend_writer.update_database.call_args.kwargs["kwargs"],
+        response_obj={},
+        start_time=datetime.now(),
+        end_time=datetime.now(),
+    )
+    assert json.loads(row["request_tags"]) == expected_tags
+
+
 @pytest.mark.asyncio
 async def test_async_post_call_failure_hook_does_not_clobber_guardrail_info_in_metadata():
     logger = ProxyDBLogger()
@@ -1105,6 +1176,38 @@ async def test_track_cost_callback_skips_when_no_standard_logging_object():
         mock_proxy_logging.failed_tracking_alert.assert_not_called()
 
 
+class _FailedTrackingAlerts:
+    def __init__(self) -> None:
+        self.messages: tuple[str, ...] = ()
+
+    async def failed_tracking_alert(self, error_message: str, failing_model: str) -> None:
+        self.messages = (*self.messages, error_message)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(("call_type", "alerted"), [("_aresponses_websocket", True), ("_arealtime", False)])
+async def test_track_cost_callback_alerts_on_a_websocket_turn_with_no_logging_object(call_type: str, alerted: bool):
+    kwargs: Final = {
+        "call_type": call_type,
+        "model": "gpt-5-mini",
+        "litellm_call_id": "turn-call-id",
+        "litellm_params": {},
+        "stream": False,
+    }
+    alerts: Final = _FailedTrackingAlerts()
+
+    with patch("litellm.proxy.proxy_server.proxy_logging_obj", alerts):
+        await ProxyDBLogger()._PROXY_track_cost_callback(
+            kwargs=kwargs,
+            completion_response=None,
+            start_time=datetime.now(),
+            end_time=datetime.now(),
+        )
+        await asyncio.sleep(0)
+
+    assert any(f"call_type: {call_type}" in message for message in alerts.messages) is alerted
+
+
 @pytest.mark.asyncio
 async def test_track_cost_callback_defers_in_progress_background_interaction():  # test-quality-ok: writing no spend row and raising no alert is the whole observable contract of the deferral path
     """

tests/unit/proxy/response_api_endpoints/test_endpoints.py

@@ -14,12 +14,15 @@ from unittest.mock import AsyncMock, MagicMock, patch
 import httpx
 import pytest
 import respx
+from fastapi import HTTPException
 from fastapi.testclient import TestClient
 from httpx import Response
 
 import litellm
+from litellm.proxy._types import ProxyException
 from litellm.proxy.proxy_server import app
 from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
+from litellm.types.router import RouterErrors
 from tests._master_key import MASTER_KEY
 
 
@@ -1648,7 +1651,7 @@ class TestManagedResponsesWSFirstMessage:
 
         processed: list = []
 
-        async def fake_process(msg: str) -> None:
+        async def fake_process(msg: str, turn_admission: object = None) -> None:
             processed.append(msg)
 
         handler = ManagedResponsesWebSocketHandler(
@@ -1693,30 +1696,32 @@ class TestResponsesWSStreamingFirstMessage:
         backend_ws.send.assert_awaited_once_with(first)
 
 
+@pytest.mark.parametrize(
+    ("failure", "expected_status", "expected_type"),
+    [
+        (ValueError("No deployments available - crossed budget: Exceeded budget for deployment"), 429, "rate_limit_exceeded"),
+        (litellm.BadRequestError(message="bad input", model="m", llm_provider="openai"), 400, "invalid_request_error"),
+        (RuntimeError("boom"), 500, "server_error"),
+        (ValueError(f"{RouterErrors.no_deployments_with_tag_routing.value}=['free']"), 401, "invalid_request_error"),
+        (ProxyException(message="Authentication Error", type="auth_error", param=None, code=401), 401, "invalid_request_error"),
+        (HTTPException(status_code=403, detail="key not allowed"), 403, "invalid_request_error"),
+    ],
+    ids=["crossed-budget", "bad-request", "unexpected", "tag-routing", "proxy-auth", "http-exception"],
+)
+def test_responses_ws_failure_frame_uses_the_http_status_for_the_same_failure(
+    failure: Exception, expected_status: int, expected_type: str
+):
+    from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_failure_frame
+    from litellm.proxy.utils import handle_exception_on_proxy
+
+    frame: Final = json.loads(_responses_ws_failure_frame(failure))
+
+    assert frame["status"] == expected_status
+    assert frame["error"]["type"] == expected_type
+    assert frame["error"]["message"] == handle_exception_on_proxy(failure).message
+
+
 class TestWSSessionCostTracking:
-    @pytest.mark.asyncio
-    async def test_router_budget_limiter_skips_aresponses_websocket_call_type(self):
-        """
-        RouterBudgetLimiting.async_log_success_event must not raise when
-        call_type='_aresponses_websocket', even when standard_logging_object is None.
-        Per-turn costs are tracked by individual aresponses calls inside the session;
-        the outer session wrapper fires with result=None.
-        """
-        from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
-
-        limiter = RouterBudgetLimiting.__new__(RouterBudgetLimiting)
-        kwargs = {
-            "call_type": "_aresponses_websocket",
-            "standard_logging_object": None,
-            "litellm_params": {"custom_llm_provider": "vertex_ai"},
-        }
-        await limiter.async_log_success_event(
-            kwargs=kwargs,
-            response_obj=None,
-            start_time=None,
-            end_time=None,
-        )
-
     @pytest.mark.asyncio
     async def test_router_budget_limiter_skips_arealtime_call_type(self):
         """Same guard applies to _arealtime WS session wrappers."""
@@ -2049,6 +2054,44 @@ class TestResponsesWSFirstFrameModelAuth:
         assert failure["request_data"]["input"] == "hi"
         assert len(failure["request_data"]["litellm_call_id"]) == 36
 
+    @pytest.mark.asyncio
+    async def test_frame_guard_bills_a_passed_frame_the_deployment_refuses_with_its_guardrail_records(self):
+        from litellm.exceptions import RateLimitError
+        from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger
+        from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard
+
+        async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only):
+            data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}]
+            return {**data, "input": data["input"] + " [guarded]"}
+
+        proxy_logging_obj = MagicMock()
+        proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook)
+        spend_logger = MagicMock(spec=ProxyDBLogger)
+        spend_logger.async_post_call_failure_hook = AsyncMock()
+        rate_limiter = MagicMock()
+        rate_limiter.async_post_call_failure_hook = AsyncMock()
+        user_api_key_dict = MagicMock()
+        refusal = RateLimitError(message="crossed budget", llm_provider="", model="")
+        guard = _responses_ws_frame_guard(
+            {"model": "gpt-4o-mini", "litellm_metadata": {}}, user_api_key_dict, proxy_logging_obj
+        )
+
+        guarded = await guard(json.dumps({"type": "response.create", "input": "hi"}))
+        assert guarded.bill_refusal is not None
+        with patch.object(litellm, "callbacks", [rate_limiter, spend_logger]):
+            await guarded.bill_refusal(refusal)
+
+        rate_limiter.async_post_call_failure_hook.assert_not_awaited()
+        spend_logger.async_post_call_failure_hook.assert_awaited_once()
+        failure = spend_logger.async_post_call_failure_hook.await_args.kwargs
+        user_api_key_dict.model_copy.assert_called_once_with(update={"budget_reservation": None})
+        assert failure["user_api_key_dict"] is user_api_key_dict.model_copy.return_value
+        assert failure["original_exception"] is refusal
+        assert failure["request_data"]["litellm_metadata"]["standard_logging_guardrail_information"] == [
+            {"guardrail_cost": 0.5}
+        ]
+        assert failure["request_data"]["input"] == "hi [guarded]"
+
     @pytest.mark.asyncio
     async def test_frame_guard_guards_the_nested_response_body_and_copies_back_tools(self):
         from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard

tests/unit/responses/test_responses_websocket_all_providers.py

@@ -7,6 +7,7 @@ Tests that:
 3. Providers without native websocket support use ManagedResponsesWebSocketHandler
 """
 
+import asyncio
 import json
 from collections.abc import Mapping
 from typing import Final
@@ -1331,7 +1332,7 @@ class TestWebSocketProjectQuotaEnforcement:
     callback, not just the first one."""
 
     @pytest.mark.asyncio
-    async def test_managed_handler_blocks_frame_rejected_by_quota_callback(self, monkeypatch):
+    async def test_managed_handler_blocks_frame_rejected_by_quota_callback(self, monkeypatch: pytest.MonkeyPatch):
         from unittest.mock import AsyncMock, MagicMock
 
         import litellm
@@ -1377,11 +1378,11 @@ class TestWebSocketProjectQuotaEnforcement:
         quota_callback.enforce_project_io_token_quota_for_frame.assert_awaited_once()
         assert aresponses_called is False
         mock_websocket.send_text.assert_called_once()
-        error_event = mock_websocket.send_text.call_args[0][0]
-        assert "rate_limit_exceeded" in error_event
+        error_event = json.loads(mock_websocket.send_text.call_args[0][0])
+        assert (error_event["status"], error_event["error"]["type"]) == (429, "rate_limit_exceeded")
 
     @pytest.mark.asyncio
-    async def test_managed_handler_forwards_frame_allowed_by_quota_callback(self, monkeypatch):
+    async def test_managed_handler_forwards_frame_allowed_by_quota_callback(self, monkeypatch: pytest.MonkeyPatch):
         from unittest.mock import AsyncMock, MagicMock
 
         import litellm
@@ -1462,7 +1463,8 @@ class TestWebSocketProjectQuotaEnforcement:
         assert allowed is False
         mock_backend_ws.send.assert_not_called()
         mock_websocket.send_text.assert_called_once()
-        assert "rate_limit_exceeded" in mock_websocket.send_text.call_args[0][0]
+        error_event = json.loads(mock_websocket.send_text.call_args[0][0])
+        assert (error_event["status"], error_event["error"]["type"]) == (429, "rate_limit_exceeded")
 
     @pytest.mark.asyncio
     async def test_native_handler_runs_the_frame_guard_on_every_turn_and_bills_each_turn_its_own_guardrail_records(
@@ -1580,6 +1582,425 @@ class TestWebSocketProjectQuotaEnforcement:
         quota_callback.enforce_project_io_token_quota_for_frame.assert_awaited_once()
 
 
+def _turn_frames() -> tuple[str, str, str]:
+    return tuple(json.dumps({"type": "response.create", "input": word}) for word in ("one", "two", "three"))
+
+
+def _admission_refusing_the_next_turn():
+    from unittest.mock import AsyncMock
+
+    from litellm.exceptions import RateLimitError
+
+    return AsyncMock(side_effect=[RateLimitError(message="crossed budget", llm_provider="", model=""), None])
+
+
+class _ClientSocket:
+    """A client socket that, like Starlette's, stops receiving once the server closes it."""
+
+    def __init__(self, *frames: str) -> None:
+        self._frames: Final = iter(frames)
+        self.sent: Final[list[dict[str, object]]] = []
+        self.close_code: int | None = None
+
+    async def send_text(self, data: str) -> None:
+        self.sent.append(json.loads(data))
+
+    async def receive_text(self) -> str:
+        frame: Final = next(self._frames, None)
+        if self.close_code is not None or frame is None:
+            raise RuntimeError("WebSocket is not connected")
+        return frame
+
+    async def close(self, code: int = 1000, reason: str | None = None) -> None:
+        self.close_code = code
+
+
+def _admission_crashing_the_next_turn():
+    from unittest.mock import AsyncMock
+
+    return AsyncMock(side_effect=[ValueError("custom check failed"), None])
+
+
+def _guard_blocking(word: str):
+    from litellm.responses.streaming_iterator import GuardedResponsesWsFrame
+
+    async def guard(frame: str) -> GuardedResponsesWsFrame:
+        if word in json.loads(frame)["input"]:
+            raise ValueError("blocked by guardrail")
+        return GuardedResponsesWsFrame(frame=frame)
+
+    return guard
+
+
+def _guard_recording(guardrail_information: tuple[Mapping[str, object], ...], bill_refusal: object):
+    from litellm.responses.streaming_iterator import GuardedResponsesWsFrame
+
+    async def guard(frame: str) -> GuardedResponsesWsFrame:
+        return GuardedResponsesWsFrame(
+            frame=frame, guardrail_information=guardrail_information, bill_refusal=bill_refusal
+        )
+
+    return guard
+
+
+def _errors(websocket: _ClientSocket) -> list[tuple[object, object]]:
+    return [(frame.get("status"), frame["error"]) for frame in websocket.sent if frame.get("type") == "error"]
+
+
+class TestWebSocketTurnAdmission:
+    """A turn after the connection's first is admitted by the deployment again before it reaches the provider."""
+
+    @pytest.mark.asyncio
+    async def test_native_relay_drops_a_later_turn_the_deployment_refuses(self):
+        from unittest.mock import AsyncMock
+
+        first, second, third = _turn_frames()
+        cancel: Final = json.dumps({"type": "response.cancel"})
+        turn_admission: Final = _admission_refusing_the_next_turn()
+        backend_ws: Final = MagicMock()
+        backend_ws.send = AsyncMock()
+        websocket: Final = MagicMock()
+        websocket.send_text = AsyncMock()
+        websocket.receive_text = AsyncMock(side_effect=[second, cancel, third, Exception("stop")])
+        handler: Final = _make_streaming(
+            websocket=websocket,
+            backend_ws=backend_ws,
+            request_data={},
+            first_message=first,
+            turn_admission=turn_admission,
+        )
+
+        await handler.client_to_backend()
+
+        assert [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] == [
+            json.loads(first),
+            json.loads(cancel),
+            json.loads(third),
+        ]
+        assert turn_admission.await_count == 2
+        refusal: Final = json.loads(websocket.send_text.await_args.args[0])
+        assert refusal["status"] == 429
+        assert refusal["error"]["type"] == "rate_limit_exceeded"
+        assert "crossed budget" in refusal["error"]["message"]
+
+    @pytest.mark.asyncio
+    async def test_native_relay_reports_a_refused_turn_as_the_socket_failure(self):
+        from unittest.mock import AsyncMock
+
+        from litellm.exceptions import RateLimitError
+
+        first, second, _ = _turn_frames()
+        backend_ws: Final = MagicMock()
+        backend_ws.send = AsyncMock()
+        backend_ws.close = AsyncMock()
+        backend_ws.recv = AsyncMock(side_effect=Exception("stop"))
+        websocket: Final = MagicMock()
+        websocket.send_text = AsyncMock()
+        websocket.receive_text = AsyncMock(side_effect=[second, Exception("client gone")])
+        handler: Final = _make_streaming(
+            websocket=websocket,
+            backend_ws=backend_ws,
+            request_data={},
+            first_message=first,
+            turn_admission=_admission_refusing_the_next_turn(),
+        )
+
+        failure: Final = await handler.bidirectional_forward()
+
+        assert isinstance(failure, RateLimitError)
+        assert "crossed budget" in str(failure)
+
+    @pytest.mark.asyncio
+    async def test_managed_relay_skips_a_later_turn_the_deployment_refuses(self, monkeypatch: pytest.MonkeyPatch):
+        from unittest.mock import AsyncMock
+
+        import litellm
+        from litellm.exceptions import RateLimitError
+        from litellm.litellm_core_utils.litellm_logging import Logging
+        from litellm.responses.streaming_iterator import ManagedResponsesWebSocketHandler
+
+        answered: Final[list[object]] = []
+
+        async def fake_aresponses(*args: object, **kwargs: object):
+            answered.append(kwargs.get("input"))
+
+            async def _empty():
+                return
+                yield
+
+            return _empty()
+
+        monkeypatch.setattr(litellm, "aresponses", fake_aresponses)
+        first, second, third = _turn_frames()
+        turn_admission: Final = _admission_refusing_the_next_turn()
+        websocket: Final = MagicMock()
+        websocket.send_text = AsyncMock()
+        websocket.receive_text = AsyncMock(side_effect=[second, third, Exception("stop")])
+        handler: Final = ManagedResponsesWebSocketHandler(
+            websocket=websocket,
+            model="test-model",
+            logging_obj=Logging(
+                model="test-model",
+                messages=[],
+                stream=True,
+                call_type="aresponses",
+                start_time=0,
+                litellm_call_id="test-id",
+                function_id="test-func",
+            ),
+            first_message=first,
+            turn_admission=turn_admission,
+        )
+
+        failure: Final = await handler.run()
+
+        assert len(answered) == 2
+        assert "two" not in json.dumps(answered)
+        assert turn_admission.await_count == 2
+        assert isinstance(failure, RateLimitError)
+        assert "crossed budget" in str(failure)
+        refusals: Final = [
+            json.loads(call.args[0])
+            for call in websocket.send_text.await_args_list
+            if json.loads(call.args[0]).get("type") == "error"
+        ]
+        assert [(refusal["status"], refusal["error"]["type"]) for refusal in refusals] == [(429, "rate_limit_exceeded")]
+
+
+    @pytest.mark.asyncio
+    async def test_native_relay_runs_the_guardrail_before_the_deployment_counts_a_turn(self):
+        from unittest.mock import AsyncMock
+
+        first, second, third = _turn_frames()
+        turn_admission: Final = _admission_refusing_the_next_turn()
+        backend_ws: Final = MagicMock()
+        backend_ws.send = AsyncMock()
+        websocket: Final = _ClientSocket(second, third)
+        handler: Final = _make_streaming(
+            websocket=websocket,
+            backend_ws=backend_ws,
+            request_data={},
+            first_message=first,
+            frame_guard=_guard_blocking("two"),
+            turn_admission=turn_admission,
+        )
+
+        await handler.client_to_backend()
+
+        assert [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] == [json.loads(first)]
+        assert turn_admission.await_count == 1
+        assert [(status, error["type"]) for status, error in _errors(websocket)] == [
+            (None, "invalid_request_error"),
+            (429, "rate_limit_exceeded"),
+        ]
+
+    @pytest.mark.asyncio
+    async def test_managed_relay_runs_the_guardrail_before_the_deployment_counts_a_turn(
+        self, monkeypatch: pytest.MonkeyPatch
+    ):
+        import litellm
+        from litellm.litellm_core_utils.litellm_logging import Logging
+        from litellm.responses.streaming_iterator import ManagedResponsesWebSocketHandler
+
+        answered: Final[list[object]] = []
+
+        async def fake_aresponses(*args: object, **kwargs: object):
+            answered.append(kwargs.get("input"))
+
+            async def _empty():
+                return
+                yield
+
+            return _empty()
+
+        monkeypatch.setattr(litellm, "aresponses", fake_aresponses)
+        first, second, third = _turn_frames()
+        turn_admission: Final = _admission_refusing_the_next_turn()
+        websocket: Final = _ClientSocket(second, third)
+        handler: Final = ManagedResponsesWebSocketHandler(
+            websocket=websocket,
+            model="test-model",
+            logging_obj=Logging(
+                model="test-model",
+                messages=[],
+                stream=True,
+                call_type="aresponses",
+                start_time=0,
+                litellm_call_id="test-id",
+                function_id="test-func",
+            ),
+            first_message=first,
+            frame_guard=_guard_blocking("two"),
+            turn_admission=turn_admission,
+        )
+
+        await handler.run()
+
+        assert answered == ["one"]
+        assert turn_admission.await_count == 1
+        assert [(status, error["type"]) for status, error in _errors(websocket)] == [
+            (None, "invalid_request_error"),
+            (429, "rate_limit_exceeded"),
+        ]
+
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        ("guardrail_information", "billed"), [(({"guardrail_cost": 0.5},), True), ((), False)]
+    )
+    async def test_native_relay_bills_the_guardrails_of_a_turn_the_deployment_refuses(
+        self, guardrail_information: tuple[Mapping[str, object], ...], billed: bool
+    ):
+        from unittest.mock import AsyncMock, call
+
+        from litellm.exceptions import RateLimitError
+
+        first, second, third = _turn_frames()
+        refusal: Final = RateLimitError(message="crossed budget", llm_provider="", model="")
+        bill_refusal: Final = AsyncMock()
+        backend_ws: Final = MagicMock()
+        backend_ws.send = AsyncMock()
+        handler: Final = _make_streaming(
+            websocket=_ClientSocket(second, third),
+            backend_ws=backend_ws,
+            request_data={},
+            first_message=first,
+            frame_guard=_guard_recording(guardrail_information, bill_refusal),
+            turn_admission=AsyncMock(side_effect=[refusal, None]),
+        )
+
+        await handler.client_to_backend()
+
+        assert [json.loads(sent.args[0])["input"] for sent in backend_ws.send.await_args_list] == ["one", "three"]
+        assert bill_refusal.await_args_list == ([call(refusal)] if billed else [])
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        ("guardrail_information", "billed"), [(({"guardrail_cost": 0.5},), True), ((), False)]
+    )
+    async def test_managed_relay_bills_the_guardrails_of_a_turn_the_deployment_refuses(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        guardrail_information: tuple[Mapping[str, object], ...],
+        billed: bool,
+    ):
+        from unittest.mock import AsyncMock, call
+
+        import litellm
+        from litellm.exceptions import RateLimitError
+        from litellm.litellm_core_utils.litellm_logging import Logging
+        from litellm.responses.streaming_iterator import ManagedResponsesWebSocketHandler
+
+        answered: Final[list[object]] = []
+
+        async def fake_aresponses(*args: object, **kwargs: object):
+            answered.append(kwargs.get("input"))
+
+            async def _empty():
+                return
+                yield
+
+            return _empty()
+
+        monkeypatch.setattr(litellm, "aresponses", fake_aresponses)
+        first, second, third = _turn_frames()
+        refusal: Final = RateLimitError(message="crossed budget", llm_provider="", model="")
+        bill_refusal: Final = AsyncMock()
+        handler: Final = ManagedResponsesWebSocketHandler(
+            websocket=_ClientSocket(second, third),
+            model="test-model",
+            logging_obj=Logging(
+                model="test-model",
+                messages=[],
+                stream=True,
+                call_type="aresponses",
+                start_time=0,
+                litellm_call_id="test-id",
+                function_id="test-func",
+            ),
+            first_message=first,
+            frame_guard=_guard_recording(guardrail_information, bill_refusal),
+            turn_admission=AsyncMock(side_effect=[refusal, None]),
+        )
+
+        await handler.run()
+
+        assert answered == ["one", "three"]
+        assert bill_refusal.await_args_list == ([call(refusal)] if billed else [])
+
+    @pytest.mark.asyncio
+    async def test_native_relay_answers_a_crashed_turn_check_with_a_500_and_keeps_serving(self):
+        from unittest.mock import AsyncMock
+
+        first, second, third = _turn_frames()
+        backend_ws: Final = MagicMock()
+        backend_ws.send = AsyncMock()
+        backend_ws.close = AsyncMock()
+        backend_ws.recv = AsyncMock(side_effect=asyncio.Event().wait)
+        websocket: Final = _ClientSocket(second, third)
+        handler: Final = _make_streaming(
+            websocket=websocket,
+            backend_ws=backend_ws,
+            request_data={},
+            first_message=first,
+            turn_admission=_admission_crashing_the_next_turn(),
+        )
+
+        failure: Final = await handler.bidirectional_forward()
+
+        assert isinstance(failure, ValueError)
+        assert [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] == [
+            json.loads(first),
+            json.loads(third),
+        ]
+        assert _errors(websocket) == [(500, {"type": "server_error", "message": "custom check failed"})]
+        assert websocket.close_code is None
+
+    @pytest.mark.asyncio
+    async def test_managed_relay_answers_a_crashed_turn_check_with_a_500_and_keeps_serving(
+        self, monkeypatch: pytest.MonkeyPatch
+    ):
+        import litellm
+        from litellm.litellm_core_utils.litellm_logging import Logging
+        from litellm.responses.streaming_iterator import ManagedResponsesWebSocketHandler
+
+        answered: Final[list[object]] = []
+
+        async def fake_aresponses(*args: object, **kwargs: object):
+            answered.append(kwargs.get("input"))
+
+            async def _empty():
+                return
+                yield
+
+            return _empty()
+
+        monkeypatch.setattr(litellm, "aresponses", fake_aresponses)
+        first, second, third = _turn_frames()
+        websocket: Final = _ClientSocket(second, third)
+        handler: Final = ManagedResponsesWebSocketHandler(
+            websocket=websocket,
+            model="test-model",
+            logging_obj=Logging(
+                model="test-model",
+                messages=[],
+                stream=True,
+                call_type="aresponses",
+                start_time=0,
+                litellm_call_id="test-id",
+                function_id="test-func",
+            ),
+            first_message=first,
+            turn_admission=_admission_crashing_the_next_turn(),
+        )
+
+        failure: Final = await handler.run()
+
+        assert isinstance(failure, ValueError)
+        assert answered == ["one", "three"]
+        assert _errors(websocket) == [(500, {"type": "server_error", "message": "custom check failed"})]
+        assert websocket.close_code is None
+
 def _deployment_defaults():
     from types import MappingProxyType
 

tests/unit/router_strategy/test_budget_limiter.py

@@ -60,7 +60,9 @@ async def _log_success(limiter: RouterBudgetLimiting, kwargs: dict[str, object])
 
 
 @pytest.mark.asyncio
-@pytest.mark.parametrize("call_type", ["aresponses", "anthropic_messages", "aembedding", "arerank"])
+@pytest.mark.parametrize(
+    "call_type", ["aresponses", "_aresponses_websocket", "anthropic_messages", "aembedding", "arerank"]
+)
 async def test_provider_spend_tracked_when_litellm_params_omits_provider(disable_budget_sync, call_type):
     """Non-chat surfaces carry the provider only on the standard logging payload."""
     limiter = RouterBudgetLimiting(

tests/unit/test_router/test_router_helper_utils.py

@@ -1,6 +1,7 @@
 import json
 import os
 from datetime import datetime, timezone
+from collections.abc import Mapping
 from typing import Final
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -13,6 +14,7 @@ from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_
 from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
 from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo
 from litellm.types.utils import (
+    CallTypes,
     ModelResponse,
     StandardLoggingHiddenParams,
     StandardLoggingMetadata,
@@ -482,6 +484,227 @@ async def test_deployment_callback_on_success(sync_mode):
     assert tpm_key is not None
 
 
+@pytest.mark.asyncio
+async def test_deployment_callback_on_success_counts_a_responses_websocket_turn():
+    import time
+
+    model_list = [
+        {
+            "model_name": "gpt-5-mini",
+            "litellm_params": {"model": "gpt-5-mini", "api_key": "sk-fake", "rpm": 100},
+            "model_info": {"id": "ws-deployment"},
+        }
+    ]
+    router = Router(model_list=model_list)
+    standard_logging_payload = create_standard_logging_payload()
+    standard_logging_payload["total_tokens"] = 17
+
+    def turn_kwargs(call_type: str) -> dict:
+        return {
+            "call_type": call_type,
+            "litellm_params": {
+                "metadata": {"model_group": "gpt-5-mini"},
+                "model_info": {"id": "ws-deployment"},
+            },
+            "standard_logging_object": standard_logging_payload,
+        }
+
+    tpm_key = await router.deployment_callback_on_success(
+        kwargs=turn_kwargs("_aresponses_websocket"),
+        completion_response=None,
+        start_time=time.time(),
+        end_time=time.time(),
+    )
+    assert tpm_key is not None
+    assert await router.cache.async_get_cache(key=tpm_key) == 17
+
+    realtime_key = await router.deployment_callback_on_success(
+        kwargs=turn_kwargs("_arealtime"),
+        completion_response=None,
+        start_time=time.time(),
+        end_time=time.time(),
+    )
+    assert realtime_key is None
+    assert await router.cache.async_get_cache(key=tpm_key) == 17
+
+
+async def _kwargs_routed_to(router: Router, model: str, function_name: str) -> dict[str, object]:
+    routed: Final[dict[str, object]] = {}
+
+    async def call(**kwargs: object) -> None:
+        routed.update(kwargs)
+
+    call.__name__ = function_name
+    await router._ageneric_api_call_with_fallbacks_helper(model=model, original_generic_function=call)
+    return routed
+
+
+async def _finish_35_token_websocket_turn(routed: Mapping[str, object], call_id: str) -> None:
+    import time
+
+    from litellm.litellm_core_utils.litellm_logging import Logging
+
+    turn: Final = Logging(
+        model="gpt-5-mini",
+        messages=[],
+        stream=False,
+        call_type=CallTypes.aresponses_websocket.value,
+        start_time=time.time(),
+        litellm_call_id=call_id,
+        function_id=call_id,
+    )
+    turn.update_environment_variables(litellm_params={"litellm_metadata": routed["litellm_metadata"]}, optional_params={})
+    await turn.dispatch_success_handlers(
+        [{"type": "response.completed", "response": {"usage": {"input_tokens": 20, "output_tokens": 15, "total_tokens": 35}}}],
+        prefer_async_handlers=True,
+    )
+
+
+@pytest.mark.asyncio
+async def test_responses_websocket_turn_is_refused_once_its_deployment_crosses_budget():
+    router = Router(
+        model_list=[
+            {
+                "model_name": "ws-budget",
+                "litellm_params": {
+                    "model": "openai/gpt-5-mini",
+                    "api_key": "sk-fake",
+                    "max_budget": 0.001,
+                    "budget_duration": "1d",
+                },
+                "model_info": {"id": "ws-budget-dep"},
+            }
+        ]
+    )
+    assert "turn_admission" not in await _kwargs_routed_to(router, "ws-budget", "aresponses")
+    admit_turn: Final = (await _kwargs_routed_to(router, "ws-budget", "_aresponses_websocket"))["turn_admission"]
+    assert callable(admit_turn)
+
+    await admit_turn()
+    await router.cache.async_set_cache("deployment_spend:ws-budget-dep:1d", 0.002)
+
+    with pytest.raises(litellm.RateLimitError, match="crossed budget"):
+        await admit_turn()
+
+
+@pytest.mark.asyncio
+async def test_responses_websocket_turn_after_the_first_counts_toward_deployment_rpm():
+    router = Router(
+        model_list=[
+            {
+                "model_name": "ws-rpm",
+                "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake", "rpm": 1},
+                "model_info": {"id": "ws-rpm-dep"},
+            }
+        ],
+        optional_pre_call_checks=["enforce_model_rate_limits"],
+    )
+    admit_turn: Final = (await _kwargs_routed_to(router, "ws-rpm", "_aresponses_websocket"))["turn_admission"]
+    assert callable(admit_turn)
+
+    with pytest.raises(litellm.RateLimitError, match="RPM limit=1"):
+        await admit_turn()
+
+
+@pytest.mark.asyncio
+async def test_responses_websocket_turn_is_checked_against_limits_changed_after_the_socket_opened():
+    router = Router(
+        model_list=[
+            {
+                "model_name": "ws-rpm-lowered",
+                "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake", "rpm": 100},
+                "model_info": {"id": "ws-rpm-lowered-dep"},
+            }
+        ],
+        optional_pre_call_checks=["enforce_model_rate_limits"],
+    )
+    admit_turn: Final = (await _kwargs_routed_to(router, "ws-rpm-lowered", "_aresponses_websocket"))["turn_admission"]
+    assert callable(admit_turn)
+    await admit_turn()
+
+    router.upsert_deployment(
+        Deployment(
+            model_name="ws-rpm-lowered",
+            litellm_params=LiteLLM_Params(model="openai/gpt-5-mini", api_key="sk-fake", rpm=1),
+            model_info=ModelInfo(id="ws-rpm-lowered-dep"),
+        )
+    )
+
+    with pytest.raises(litellm.RateLimitError, match="RPM limit=1"):
+        await admit_turn()
+
+
+@pytest.mark.asyncio
+async def test_responses_websocket_turn_counts_toward_deployment_tpm(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(litellm, "callbacks", [])
+    router = Router(
+        model_list=[
+            {
+                "model_name": "ws-tpm",
+                "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake", "tpm": 30},
+                "model_info": {"id": "ws-tpm-dep"},
+            }
+        ],
+        optional_pre_call_checks=["enforce_model_rate_limits"],
+    )
+    monkeypatch.setattr(litellm, "_async_success_callback", list(litellm.callbacks))
+    routed: Final = await _kwargs_routed_to(router, "ws-tpm", "_aresponses_websocket")
+    admit_turn: Final = routed["turn_admission"]
+    assert callable(admit_turn)
+    await admit_turn()
+
+    await _finish_35_token_websocket_turn(routed, "ws-tpm-turn")
+
+    with pytest.raises(litellm.RateLimitError, match="TPM limit=30, current usage=35"):
+        await admit_turn()
+
+
+@pytest.mark.asyncio
+async def test_responses_websocket_turn_is_refused_by_usage_based_routing_once_deployment_tpm_is_used(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    monkeypatch.setattr(litellm, "callbacks", [])
+    router = Router(
+        model_list=[
+            {
+                "model_name": "ws-tpm-v2",
+                "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake", "tpm": 30},
+                "model_info": {"id": "ws-tpm-v2-dep"},
+            }
+        ],
+        routing_strategy="usage-based-routing-v2",
+    )
+    monkeypatch.setattr(litellm, "_async_success_callback", list(litellm.callbacks))
+    routed: Final = await _kwargs_routed_to(router, "ws-tpm-v2", "_aresponses_websocket")
+    admit_turn: Final = routed["turn_admission"]
+    assert callable(admit_turn)
+    await admit_turn()
+
+    await _finish_35_token_websocket_turn(routed, "ws-tpm-v2-turn")
+
+    with pytest.raises(litellm.RateLimitError):
+        await admit_turn()
+
+
+@pytest.mark.asyncio
+async def test_responses_websocket_turns_do_not_reserve_deployment_otpm():
+    router = Router(
+        model_list=[
+            {
+                "model_name": "ws-otpm",
+                "litellm_params": {"model": "openai/ws-otpm-unpriced", "api_key": "sk-fake", "otpm": 5000},
+                "model_info": {"id": "ws-otpm-dep"},
+            }
+        ],
+        optional_pre_call_checks=["enforce_model_rate_limits"],
+    )
+    admit_turn: Final = (await _kwargs_routed_to(router, "ws-otpm", "_aresponses_websocket"))["turn_admission"]
+    assert callable(admit_turn)
+
+    for _ in range(3):
+        await admit_turn()
+
+
 @pytest.mark.asyncio
 async def test_deployment_callback_on_success_tracks_tpm_for_io_deployment():
     """