fix(proxy): keep per-turn WebSocket limits, request bodies and collector reservations right

Free a rejected request's parallel slot through the request it runs in, and skip
native Responses WebSocket connections, whose disconnect frees the slot. Count
WebSocket turns toward the connection's tag TPM. In collector mode, claim the
shared budget reservation before the spend event is built, so only the first
turn settles it. Store each turn's own frame as its request body, and send an
id-less error to the oldest waiting turn unless an earlier turn is streaming.

litellm/litellm_core_utils/litellm_logging.py

@@ -541,6 +541,17 @@ _PER_CALL_LITELLM_PARAM_DICTS: Final = frozenset({"metadata", "litellm_metadata"
 _LITELLM_PARAMS: Final = TypeAdapter(dict[str, object])
 
 
+def _proxy_server_request_with_body(
+    litellm_params: Mapping[str, object], request_body: Mapping[str, object]
+) -> Mapping[str, object]:
+    proxy_server_request: Final = litellm_params.get("proxy_server_request")
+    return (
+        {"proxy_server_request": {**_LITELLM_PARAMS.validate_python(proxy_server_request), "body": dict(request_body)}}
+        if isinstance(proxy_server_request, Mapping)
+        else {}
+    )
+
+
 def _provider_response_id(source: object) -> str | None:
     candidate: Final = source.get("id") if isinstance(source, dict) else getattr(source, "id", None)
     return candidate if isinstance(candidate, str) and candidate else None
@@ -767,7 +778,9 @@ class Logging(LiteLLMLoggingBaseClass):
         self.on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None
         self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None
 
-    def copy_for_new_call(self, litellm_call_id: str, start_time: datetime.datetime) -> "Logging":
+    def copy_for_new_call(
+        self, litellm_call_id: str, start_time: datetime.datetime, request_body: Mapping[str, object]
+    ) -> "Logging":
         connection_params: Final = _LITELLM_PARAMS.validate_python(self.litellm_params)
         litellm_params: Final[dict[str, object]] = {
             **connection_params,
@@ -776,6 +789,7 @@ class Logging(LiteLLMLoggingBaseClass):
                 for key, value in connection_params.items()
                 if key in _PER_CALL_LITELLM_PARAM_DICTS and isinstance(value, dict)
             },
+            **_proxy_server_request_with_body(connection_params, request_body),
             "litellm_call_id": litellm_call_id,
         }
         new_call: Final = copy.copy(self)

litellm/proxy/hooks/parallel_request_limiter_v3.py

@@ -742,15 +742,18 @@ def get_or_create_request_stash() -> RequestRateLimiterStash:
     return stash
 
 
-def _request_owner_call_id(data: Mapping[str, object]) -> str | None:
-    from litellm.proxy.common_request_processing import request_litellm_call_id
-
-    return request_litellm_call_id(data)
+def _tpm_limited_tags(call_stash: RequestRateLimiterStash | None) -> frozenset[str]:
+    if call_stash is not None:
+        return call_stash.tpm_limited_tags
+    connection: Final = get_request_stash()
+    return connection.tpm_limited_tags if connection is not None and connection.websocket_connection else frozenset()
 
 
 def claim_request_stash_for_data(data: dict) -> RequestRateLimiterStash:
+    from litellm.proxy.common_request_processing import request_litellm_call_id
+
     stash: Final = get_or_create_request_stash()
-    owner_call_id: Final = _request_owner_call_id(_REQUEST_RATE_LIMIT_DATA.validate_python(data))
+    owner_call_id: Final = request_litellm_call_id(_REQUEST_RATE_LIMIT_DATA.validate_python(data))
     if owner_call_id is not None:
         stash.owner_litellm_call_id = owner_call_id
     return stash
@@ -5067,7 +5070,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
         targets: Final = self._collect_tpm_scope_targets(
             standard_logging_metadata=standard_logging_metadata,
             kwargs=kwargs,
-            tpm_limited_tags=stash.tpm_limited_tags if stash is not None else frozenset(),
+            tpm_limited_tags=_tpm_limited_tags(stash),
             model_group=reconcile_model.group if reconcile_model is not None else None,
         )
         targets.extend(
@@ -5458,12 +5461,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
         both this hook and async_log_failure_event end up running in the same
         flow, only the first release/refund applies. A mid-stream failure
         relayed here with recovered partial usage settles the reservation at
-        that usage instead of refunding it.
+        that usage instead of refunding it. A native Responses WebSocket
+        connection keeps its slot and reservation through a failed turn, and
+        its disconnect releases them.
         """
         try:
-            failed_request: Final = _REQUEST_RATE_LIMIT_DATA.validate_python(request_data)
-            stash: Final = get_request_stash_for_call(_request_owner_call_id(failed_request))
-            if stash is None:
+            stash: Final = get_request_stash()
+            if stash is None or stash.websocket_connection:
                 return
             await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span)
 

litellm/proxy/hooks/proxy_track_cost_callback.py

@@ -139,6 +139,12 @@ class _ProxyDBLogger(CustomLogger):
         if self.spend_event_producer is None or not is_offloadable_success(response_obj):
             await self._PROXY_track_cost_callback(kwargs, response_obj, start_time, end_time)
             return
+        _claim_budget_reservation(
+            budget_reservation=_get_budget_reservation_from_metadata(
+                metadata=get_litellm_metadata_from_kwargs(kwargs=kwargs)
+            ),
+            litellm_call_id=_string_call_id(_CALLBACK_DATA.validate_python(kwargs)),
+        )
         event: Final = build_spend_event(
             kwargs,
             response_obj,
@@ -170,7 +176,7 @@ class _ProxyDBLogger(CustomLogger):
     ) -> None:
         budget_reservation: Final = _claim_budget_reservation(
             budget_reservation=user_api_key_dict.budget_reservation,
-            litellm_call_id=_request_call_id(_CALLBACK_DATA.validate_python(request_data)),
+            litellm_call_id=_string_call_id(_CALLBACK_DATA.validate_python(request_data)),
         )
         try:
             await _release_budget_reservation(budget_reservation=budget_reservation)
@@ -353,7 +359,7 @@ class _ProxyDBLogger(CustomLogger):
                 _write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata)
             budget_reservation: Final = _claim_budget_reservation(
                 budget_reservation=_get_budget_reservation_from_metadata(metadata=metadata),
-                litellm_call_id=_logged_call_id(_CALLBACK_DATA.validate_python(kwargs)),
+                litellm_call_id=_string_call_id(_CALLBACK_DATA.validate_python(kwargs)),
             )
             if (
                 isinstance(completion_response, LiteLLMBatch)
@@ -709,14 +715,8 @@ def _get_budget_reservation_from_metadata(metadata: Mapping[str, object]) -> dic
 _CALLBACK_DATA: Final = TypeAdapter(Mapping[str, object])
 
 
-def _request_call_id(request_data: Mapping[str, object]) -> str | None:
-    from litellm.proxy.common_request_processing import request_litellm_call_id
-
-    return request_litellm_call_id(request_data)
-
-
-def _logged_call_id(kwargs: Mapping[str, object]) -> str | None:
-    call_id: Final = kwargs.get("litellm_call_id")
+def _string_call_id(callback_data: Mapping[str, object]) -> str | None:
+    call_id: Final = callback_data.get("litellm_call_id")
     return call_id if isinstance(call_id, str) else None
 
 

litellm/responses/streaming_iterator.py

@@ -2007,7 +2007,9 @@ class ResponsesWebSocketStreaming:
     def _start_turn(self, client_frame: str | None) -> None:
         turn_start: Final = datetime.now()  # noqa: DTZ005  # logging pipeline uses naive datetimes
         turn_logging_obj: Final = self.logging_obj.copy_for_new_call(
-            litellm_call_id=str(uuid.uuid4()), start_time=turn_start
+            litellm_call_id=str(uuid.uuid4()),
+            start_time=turn_start,
+            request_body=_json_object_or_none(client_frame) or {},
         )
         turn_logging_obj.model_call_details["messages"] = list(_logged_user_messages(client_frame))
         turn_logging_obj.pre_call(input=client_frame, api_key="")
@@ -2036,10 +2038,11 @@ class ResponsesWebSocketStreaming:
     def _turn_index_for_event(self, event_obj: Mapping[str, object]) -> int:
         """
         An event naming a response id belongs to the turn whose ``response.created``
-        carried it. The backend answers accepted turns strictly in order, so an
-        id-less event goes to the oldest turn still waiting for its ``response.created``,
-        except an id-less ``error``: that is the backend rejecting the newest pending
-        ``response.create`` at parse time while an earlier turn is still streaming.
+        carried it. The backend answers turns strictly in order, so an id-less event,
+        a rejected turn's ``error`` included, goes to the oldest turn still waiting for
+        its ``response.created``. The one exception is an id-less ``error`` while a turn
+        is streaming: the backend rejects a ``response.create`` it cannot parse as soon
+        as it reads it, so that error belongs to the newest pending turn.
         """
         response_id: Final = _ws_event_response_id(event_obj)
         turn_ids: Final = tuple(turn.response_id for turn in self._open_turns)
@@ -2048,7 +2051,7 @@ class ResponsesWebSocketStreaming:
         pending: Final = tuple(index for index, turn_id in enumerate(turn_ids) if turn_id is None)
         if not pending:
             return 0
-        if response_id is None and event_obj.get("type") == "error":
+        if response_id is None and event_obj.get("type") == "error" and len(pending) < len(turn_ids):
             return pending[-1]
         return pending[0]
 

tests/unit/litellm_core_utils/test_litellm_logging.py

@@ -8267,10 +8267,12 @@ def test_copy_for_new_call_gives_each_call_its_own_id_clock_and_metadata():
     reservation = {"reserved_cost": 0.5, "finalized": False}
     connection = _responses_ws_logging_obj()
     connection_start = connection.start_time
+    connection_request = {"url": "ws://proxy/v1/responses", "body": {"first_message": "turn 1 frame"}}
     connection.litellm_params = {
         "metadata": {"user_api_key_alias": "codex", "budget_reservation": reservation},
         "litellm_metadata": {"model_info": {"id": "dep-1"}},
         "litellm_call_id": connection.litellm_call_id,
+        "proxy_server_request": connection_request,
     }
     connection.model_call_details = {
         **connection.model_call_details,
@@ -8279,7 +8281,8 @@ def test_copy_for_new_call_gives_each_call_its_own_id_clock_and_metadata():
     }
     turn_start = datetime.datetime(2026, 10, 9, 17, 24, 47)
 
-    turn = connection.copy_for_new_call(litellm_call_id="turn-1-call", start_time=turn_start)
+    turn_frame = {"type": "response.create", "input": "turn 2 frame"}
+    turn = connection.copy_for_new_call(litellm_call_id="turn-1-call", start_time=turn_start, request_body=turn_frame)
     turn.litellm_params["metadata"]["status"] = "failure"
     turn.litellm_params["litellm_metadata"]["model_group"] = "gpt-4o"
     turn.model_call_details["standard_logging_object"] = {"id": "resp_1"}
@@ -8290,6 +8293,14 @@ def test_copy_for_new_call_gives_each_call_its_own_id_clock_and_metadata():
     assert turn.model_call_details["litellm_params"]["litellm_call_id"] == "turn-1-call"
     assert turn.model_call_details["first_api_call_start_time"] is None
     assert turn.litellm_params["metadata"]["budget_reservation"] is reservation
+    assert turn.model_call_details["litellm_params"]["proxy_server_request"] == {
+        "url": "ws://proxy/v1/responses",
+        "body": turn_frame,
+    }
+    assert connection.litellm_params["proxy_server_request"] == {
+        "url": "ws://proxy/v1/responses",
+        "body": {"first_message": "turn 1 frame"},
+    }
     assert (connection.litellm_call_id, connection.start_time) == ("responses-ws-usage-test", connection_start)
     assert connection.litellm_params["litellm_call_id"] == "responses-ws-usage-test"
     assert "status" not in connection.litellm_params["metadata"]

tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py

@@ -4383,9 +4383,11 @@ async def test_post_call_failure_hook_for_another_call_keeps_the_owner_slot_v3()
     handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
     user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1)
     counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests"
-    connection_data: Dict[str, Any] = {"model": "gpt-4o-mini", "litellm_call_id": "socket-call"}
     await handler.async_pre_call_hook(
-        user_api_key_dict=user_api_key_dict, cache=local_cache, data=connection_data, call_type=""
+        user_api_key_dict=user_api_key_dict,
+        cache=local_cache,
+        data={"model": "gpt-4o-mini", "litellm_call_id": "socket-call"},
+        call_type=CallTypes.aresponses_websocket.value,
     )
 
     await handler.async_post_call_failure_hook(
@@ -4393,14 +4395,10 @@ async def test_post_call_failure_hook_for_another_call_keeps_the_owner_slot_v3()
         original_exception=Exception("previous_response_not_found"),
         user_api_key_dict=user_api_key_dict,
     )
-    in_flight_after_turn_failure = handler._gauge_in_flight_from_cache_value(
+    in_flight_after_turn_failure: Final = handler._gauge_in_flight_from_cache_value(
         await local_cache.async_get_cache(key=counter_key)
     )
-    await handler.async_post_call_failure_hook(
-        request_data=connection_data,
-        original_exception=Exception("socket closed"),
-        user_api_key_dict=user_api_key_dict,
-    )
+    await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict)
 
     assert in_flight_after_turn_failure == 1
     assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0
@@ -4436,15 +4434,16 @@ async def _admit_pass_through_with_client_call_id(
 
 
 @pytest.mark.asyncio
-async def test_post_call_failure_hook_releases_the_slot_the_proxy_call_id_owns_v3():
+async def test_post_call_failure_hook_releases_the_slot_the_proxy_call_id_owns_v3(monkeypatch: pytest.MonkeyPatch):
     _api_key = hash_token("sk-passthrough-client-id-failure")
     local_cache = DualCache()
     handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
     user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1)
     counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests"
     logging_obj = await _admit_pass_through_with_client_call_id(handler, local_cache, user_api_key_dict)
+    monkeypatch.setattr(litellm, "callbacks", [handler])
 
-    await handler.async_post_call_failure_hook(
+    await ProxyLogging(user_api_key_cache=UserApiKeyCache()).post_call_failure_hook(
         request_data={"litellm_call_id": "client-chosen-id", "litellm_logging_obj": logging_obj},
         original_exception=Exception("rejected by a pre-call hook"),
         user_api_key_dict=user_api_key_dict,
@@ -5224,6 +5223,41 @@ async def test_tag_429_names_the_tag_that_is_over_its_limit(monkeypatch):
     assert "cell-blocked" in str(exc_info.value.detail)
 
 
+@pytest.mark.asyncio
+async def test_websocket_turns_count_toward_the_connection_tag_tpm_v3(monkeypatch):
+    monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60")
+    monkeypatch.setenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", "false")
+    _request_stash.set(None)
+    resolver, _ = _static_tag_limits({"cell-1": TagRateLimit(rpm_limit=None, tpm_limit=1_000)})
+    local_cache = DualCache()
+    handler = _PROXY_MaxParallelRequestsHandler(
+        internal_usage_cache=InternalUsageCache(local_cache), tag_rate_limit_resolver=resolver
+    )
+    user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("sk-ws-tag-tpm"))
+    await handler.async_pre_call_hook(
+        user_api_key_dict=user_api_key_dict,
+        cache=local_cache,
+        data={"model": "gpt-4o-mini", "metadata": {"tags": ["cell-1"]}, "litellm_call_id": "socket-call"},
+        call_type=CallTypes.aresponses_websocket.value,
+    )
+    tokens_before_turns: Final = int(await local_cache.async_get_cache("{tag:cell-1}:tokens") or 0)
+
+    for turn_call_id in ("turn-1-call", "turn-2-call"):
+        await handler.async_log_success_event(
+            kwargs={
+                "litellm_call_id": turn_call_id,
+                "standard_logging_object": {"metadata": {"user_api_key_hash": user_api_key_dict.api_key}},
+                "model": "gpt-4o-mini",
+            },
+            response_obj=ModelResponse(usage=Usage(prompt_tokens=60, completion_tokens=60, total_tokens=120)),
+            start_time=datetime.now(),
+            end_time=datetime.now(),
+        )
+    await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict)
+
+    assert await local_cache.async_get_cache("{tag:cell-1}:tokens") == tokens_before_turns + 240
+
+
 @pytest.mark.asyncio
 async def test_tag_object_tpm_limit_enforced_v3(monkeypatch):
     monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60")

tests/unit/proxy/hooks/test_proxy_track_cost_callback.py

@@ -2718,6 +2718,44 @@ async def test_async_log_success_event_hands_the_sidecar_a_compact_event_and_ski
     assert event.response_cost == 0.0125
 
 
+@pytest.mark.asyncio
+async def test_collector_events_let_only_the_first_turn_settle_a_shared_reservation(tmp_path):
+    from litellm.litellm_core_utils.core_helpers import (
+        budget_reservation_from_metadata,
+        get_litellm_metadata_from_kwargs,
+    )
+    from litellm.proxy.spend_tracking.budget_reservation import claim_budget_reservation
+    from litellm.proxy.spend_tracking.spend_event import spend_event_callback_args
+
+    handler = _RecordingHandler()
+    consumer = SpendEventConsumer(handler)
+    address = UnixAddress(path=str(tmp_path / "spend.sock"))
+    server = await consumer.serve(address)
+    producer = SpendEventProducer(
+        address=address, on_unavailable="fallback", buffer_size=10, connect_timeout=1.0, fallback=_no_fallback
+    )
+    logger = ProxyDBLogger(producer)
+    connection_kwargs = _offload_kwargs()
+
+    for turn_call_id in ("turn-1", "turn-2", "turn-3"):
+        await logger.async_log_success_event(
+            {**connection_kwargs, "litellm_call_id": turn_call_id}, _offload_response(), datetime.now(), datetime.now()
+        )
+    await producer.close(drain_timeout=5.0)
+    server.close()
+    assert await consumer.drain(timeout=5.0) == 0
+
+    sidecar_kwargs = [spend_event_callback_args(decode_spend_event(line)).kwargs for line in handler.lines]
+    assert [
+        claim_budget_reservation(
+            budget_reservation=budget_reservation_from_metadata(get_litellm_metadata_from_kwargs(kwargs)),
+            litellm_call_id=kwargs["litellm_call_id"],
+        )
+        is not None
+        for kwargs in sidecar_kwargs
+    ] == [True, False, False]
+
+
 @pytest.mark.asyncio
 async def test_async_log_success_event_keeps_batch_retrieves_in_process():
     producer = SpendEventProducer(

tests/unit/responses/test_responses_websocket_all_providers.py

@@ -1907,12 +1907,13 @@ class _TurnLoggingRecorder:
         self.connection = MagicMock()
         self.connection.copy_for_new_call = MagicMock(side_effect=self._copy_for_new_call)
 
-    def _copy_for_new_call(self, litellm_call_id: str, start_time: object) -> MagicMock:
+    def _copy_for_new_call(self, litellm_call_id: str, start_time: object, request_body: object) -> MagicMock:
         from unittest.mock import AsyncMock
 
         turn = MagicMock()
         turn.litellm_call_id = litellm_call_id
         turn.start_time = start_time
+        turn.request_body = request_body
         turn.model_call_details = {}
         turn.dispatch_success_handlers = AsyncMock()
         turn.dispatch_failure_handlers = AsyncMock()
@@ -3336,6 +3337,9 @@ class TestNativeWebSocketPerTurnLogging:
             [{"role": "user", "content": "two"}],
             [{"role": "user", "content": "three"}],
         ]
+        assert [turn.request_body for turn in recorder.turns] == [
+            json.loads(_response_create(text)) for text in ("one", "two", "three")
+        ]
         for turn in recorder.turns:
             turn.dispatch_failure_handlers.assert_not_awaited()
 
@@ -3402,6 +3406,47 @@ class TestNativeWebSocketPerTurnLogging:
         assert "max_output_tokens" in str(second.dispatch_failure_handlers.await_args[0][0])
         assert second.record_partial_usage_for_failure.call_args[0][0].total_tokens == 0
 
+    @pytest.mark.asyncio
+    async def test_a_rejected_turn_answered_in_order_fails_and_the_next_turn_keeps_its_response(self):
+        """OpenAI answers a response.create it parses but rejects (previous_response_not_found) in order,
+        before it starts the next frame's response. With both turns still waiting, the id-less error is the
+        older turn's failure and the next response.created belongs to the newer turn."""
+        import websockets.exceptions  # noqa: F401  (lazy submodule must be importable)
+
+        recorder = _TurnLoggingRecorder()
+        handler = _make_streaming(
+            websocket=MagicMock(),
+            backend_ws=MagicMock(),
+            logging_obj=recorder.connection,
+            request_data={},
+            authorized_model="gpt-5.4-mini",
+            custom_llm_provider="openai",
+        )
+        created, _delta, completed = _completed_turn("resp_2", 40, 4)
+        rejection = {
+            "type": "error",
+            "error": {
+                "type": "invalid_request_error",
+                "code": "previous_response_not_found",
+                "message": "Previous response with id 'resp_missing' not found.",
+            },
+        }
+
+        handler._record_client_frame(_response_create("first"))
+        handler._record_client_frame(_response_create("second"))
+        for event in (rejection, created, completed):
+            handler._record_backend_event(json.dumps(event))
+        assert await _until(lambda: recorder.logged_turns() == 2)
+
+        first, second = recorder.turns
+        first.dispatch_success_handlers.assert_not_awaited()
+        assert "resp_missing" in str(first.dispatch_failure_handlers.await_args[0][0])
+        second.dispatch_failure_handlers.assert_not_awaited()
+        assert [event["type"] for event in second.dispatch_success_handlers.await_args[0][0]] == [
+            "response.created",
+            "response.completed",
+        ]
+
     @pytest.mark.asyncio
     async def test_a_failed_turn_is_booked_alone_and_leaves_the_other_turns_successful(self):
         from unittest.mock import AsyncMock