fix(proxy): charge successful Responses WebSocket turns to TPM and ITPM/OTPM

A native Responses WebSocket turn reaches the rate limiter's success
callback with its raw event list as the response and no
combined_usage_object, so the turn's tokens never reached the key's TPM
window or the project's input and output token limits. Read the usage
from the turn's events when the call type is a Responses WebSocket turn.

litellm/proxy/hooks/parallel_request_limiter_v3.py

@@ -42,6 +42,7 @@ from litellm.caching.redis_batch import (
 )
 from litellm.caching.redis_cache import log_redis_failure
 from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY
+from litellm.cost_calculator import ResponsesWebSocketTokenUsageProcessor
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
     get_str_from_messages,
@@ -767,6 +768,20 @@ def _call_id_from_callback_kwargs(kwargs: object) -> str | None:
     return call_id if isinstance(call_id, str) else None
 
 
+_RESPONSES_WS_TURN_EVENTS: Final = TypeAdapter(tuple[Mapping[str, object], ...])
+
+
+def _unmodeled_response_usage(kwargs: Mapping[str, object], response_obj: object) -> Usage | None:
+    combined_usage: Final = kwargs.get("combined_usage_object")
+    if isinstance(combined_usage, Usage):
+        return combined_usage
+    if kwargs.get("call_type") != CallTypes.aresponses_websocket.value or not isinstance(response_obj, list):
+        return None
+    return ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results(
+        _RESPONSES_WS_TURN_EVENTS.validate_python(response_obj)
+    )
+
+
 def _parse_output_cap_value(raw_value: object) -> int | None:
     if isinstance(raw_value, bool) or not isinstance(raw_value, (int, float, str)):
         return None
@@ -4793,7 +4808,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
         response_usage: Final = self._resolve_io_token_reconcile_usage(response_obj)
         if response_usage[2]:
             return response_usage
-        return self._resolve_io_token_reconcile_usage(kwargs.get("combined_usage_object"))
+        return self._resolve_io_token_reconcile_usage(_unmodeled_response_usage(kwargs, response_obj))
 
     def _build_websocket_turn_io_ops(
         self,
@@ -5036,7 +5051,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
         # Get total tokens from response. Responses LiteLLM does not model
         # (e.g. pass-through, whose usage is reported by the upstream rather
         # than parsed out of the body) carry their usage in
-        # ``combined_usage_object`` instead, and would otherwise never charge
+        # ``combined_usage_object`` instead, and a native Responses WebSocket
+        # turn in its raw event list; either would otherwise never charge
         # the TPM window.
         _usage: Usage | dict | None = None
         if isinstance(
@@ -5050,9 +5066,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
         ):
             _usage = getattr(response_obj, "usage", None)
         else:
-            _combined_usage: Final = kwargs.get("combined_usage_object")
-            if isinstance(_combined_usage, Usage):
-                _usage = _combined_usage
+            _usage = _unmodeled_response_usage(kwargs, response_obj)
         total_tokens = self._get_total_tokens_from_usage(usage=_usage, rate_limit_type=rate_limit_type)
         if total_tokens == 0:
             total_tokens = self._aggregate_only_total_tokens(usage=_usage)

tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py

@@ -8063,6 +8063,45 @@ async def test_websocket_turn_failure_counts_its_usage_toward_tpm_v3(monkeypatch
     assert get_request_stash().reservation_released is False
 
 
+@pytest.mark.asyncio
+async def test_websocket_turn_success_counts_its_usage_toward_tpm_v3(monkeypatch):
+    """
+    A successful native Responses WebSocket turn is logged with the raw event list
+    as its response, not a modeled response object. Its usage must still reach the
+    key's TPM window right away, while the socket stays open.
+    """
+    monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False)
+    _api_key = hash_token("sk-ws-successful-turn")
+    local_cache = DualCache()
+    handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
+    user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=10_000, max_parallel_requests=2)
+    tokens_key = handler.create_rate_limit_keys(key="api_key", value=_api_key, rate_limit_type="tokens")
+    reserved = await _admit_with_tpm_reservation(
+        handler, local_cache, user_api_key_dict, CallTypes.aresponses_websocket.value
+    )
+    turn_kwargs = {
+        key: value
+        for key, value in _websocket_turn_kwargs("turn-1", _api_key, Usage()).items()
+        if key != "combined_usage_object"
+    }
+
+    await handler.async_log_success_event(
+        kwargs={**turn_kwargs, "call_type": CallTypes.aresponses_websocket.value},
+        response_obj=[
+            {"type": "response.created", "response": {"id": "resp_turn_1"}},
+            {
+                "type": "response.completed",
+                "response": {"id": "resp_turn_1", "usage": {"input_tokens": 12, "output_tokens": 5, "total_tokens": 17}},
+            },
+        ],
+        start_time=datetime.now(),
+        end_time=datetime.now(),
+    )
+
+    assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == reserved + 17
+    assert get_request_stash().reservation_released is False
+
+
 @pytest.mark.asyncio
 async def test_failure_event_without_a_stash_charges_nothing_outside_a_websocket_connection_v3(monkeypatch):
     """

tests/unit/proxy/hooks/test_tpm_concurrent.py

@@ -3768,13 +3768,14 @@ async def test_websocket_turn_usage_counts_toward_project_itpm_and_otpm(rate_lim
     assert await counters() == (stash.itpm_reserved_tokens, stash.otpm_reserved_tokens)
 
     await handler.async_log_success_event(
-        kwargs=turn_kwargs("turn-1"),
-        response_obj=ResponsesAPIResponse(
-            id="resp_1",
-            created_at=int(datetime.now().timestamp()),
-            output=[],
-            usage=ResponseAPIUsage(input_tokens=30, output_tokens=7, total_tokens=37),
-        ),
+        kwargs={**turn_kwargs("turn-1"), "call_type": CallTypes.aresponses_websocket.value},
+        response_obj=[
+            {"type": "response.created", "response": {"id": "resp_1"}},
+            {
+                "type": "response.completed",
+                "response": {"id": "resp_1", "usage": {"input_tokens": 30, "output_tokens": 7, "total_tokens": 37}},
+            },
+        ],
         start_time=datetime.now(),
         end_time=datetime.now(),
     )