fix(proxy): hold a websocket session's parallel slot until the socket closes

Each finished turn logged its spend row under the session's own call id, so the first turn's success callback freed the key's max_parallel_requests slot while the socket was still open, and a key limited to one session could open several at once. The session now claims its slot under its own call id, every turn row logs under a call id of its own, and the slot is released when the session ends, as before

Priced turn rows also take their row's call id as their id, so two turns that start in the same instant no longer share a spend log id and drop one row

litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_websocket_passthrough_logging_handler.py

@@ -200,9 +200,9 @@ def _requested_model(kwargs: Mapping[str, object]) -> str | None:
     return model_group if isinstance(model_group, str) else None
 
 
-def _priced_result(model: str, billing: _Billing, start_time: datetime) -> ModelResponse:
+def _priced_result(model: str, billing: _Billing, start_time: datetime, call_id: str) -> ModelResponse:
     result: Final = ModelResponse(
-        id=f"openai-websocket-{start_time.timestamp()}",
+        id=call_id,
         object="chat.completion",
         created=int(start_time.timestamp()),
         model=model,
@@ -246,7 +246,7 @@ class OpenAIWebsocketPassthroughLoggingHandler:
 
         logging_obj.model_call_details.update(MappingProxyType({"response_cost": billing.cost}))
         priced: Final[PassThroughEndpointLoggingTypedDict] = {
-            "result": _priced_result(model, billing, start_time),
+            "result": _priced_result(model, billing, start_time, logging_obj.litellm_call_id),
             "kwargs": {
                 **kwargs,
                 "model": model,

litellm/proxy/pass_through_endpoints/pass_through_endpoints.py

@@ -3124,14 +3124,16 @@ async def websocket_passthrough_request(
 
     def log_turn(lane: str | None, in_row: Callable[[Mapping[str, object]], bool]) -> None:
         end_time: Final = datetime.now()
-        if ledger.logged_turns == 0:
+        if upstream_frame_kind is None:
             log_frames(logging_obj, passthrough_logging_payload, kwargs, end_time, lane, in_row)
         else:
             turn_call_id: Final = str(uuid.uuid4())
             turn_logging: Final = new_call_logging(ledger.started, turn_call_id)
             turn_payload: Final = new_call_payload()
             turn_kwargs: Final = init_call_kwargs(
-                user_api_key_dict.model_copy(update={"budget_reservation": None}),
+                user_api_key_dict
+                if ledger.logged_turns == 0
+                else user_api_key_dict.model_copy(update={"budget_reservation": None}),
                 turn_payload,
                 turn_call_id,
                 turn_logging,
@@ -3141,7 +3143,9 @@ async def websocket_passthrough_request(
         ledger.close_turn(lane, end_time)
 
     ### CALL HOOKS ### - modify incoming data / reject request before calling the model
-    websocket_hook_data: Final = guardrail_request_data_with_streaming(MappingProxyType({}), is_streaming=True)
+    websocket_hook_data: Final = guardrail_request_data_with_streaming(
+        MappingProxyType({"litellm_call_id": litellm_call_id}), is_streaming=True
+    )
     await proxy_logging_obj.pre_call_hook(
         user_api_key_dict=user_api_key_dict,
         data=websocket_hook_data,

tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_websocket_passthrough_logging_handler.py

@@ -1,3 +1,4 @@
+import uuid
 from collections.abc import Mapping, Sequence
 from datetime import datetime
 from types import MappingProxyType
@@ -15,6 +16,8 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_websocket
     OpenAIWebsocketPassthroughLoggingHandler,
     openai_websocket_frame_kind,
 )
+from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id
+from litellm.types.utils import CallTypes
 
 REALTIME_MODEL: Final = "gpt-realtime-2.1"
 RESPONSES_MODEL: Final = "gpt-5.4-mini"
@@ -24,7 +27,7 @@ SESSION_START: Final = datetime(2026, 1, 1, 12, 0, 0)
 
 
 def _logging_obj() -> MagicMock:
-    logging_obj: Final = MagicMock(model_call_details={})
+    logging_obj: Final = MagicMock(model_call_details={}, litellm_call_id=str(uuid.uuid4()))
     return logging_obj
 
 
@@ -290,6 +293,27 @@ def test_the_price_of_a_multi_model_connection_reaches_the_spend_row():
     assert handled["kwargs"]["model"] == RESPONSES_MODEL
 
 
+def test_turns_that_start_in_the_same_instant_each_keep_their_own_spend_row():
+    """
+    Spend logs key each row by the id its result carries and skip a repeated id
+    without a word, so two turns that start in the same instant, on one socket
+    or on two, each need an id of their own or one of them is never billed.
+    """
+    turns: Final = [
+        _handle(
+            [_responses_completed(RESPONSES_MODEL, input_tokens=13, output_tokens=5)],
+            url_route="/openai_passthrough/v1/responses",
+        )
+        for _ in range(2)
+    ]
+
+    spend_log_ids: Final = {
+        get_spend_logs_id(CallTypes.pass_through.value, turn["result"].model_dump(), turn["kwargs"]) for turn in turns
+    }
+    assert len(spend_log_ids) == 2
+    assert None not in spend_log_ids
+
+
 def test_a_turn_the_caller_walked_out_of_still_names_the_model_it_ran():
     """
     A client that closes before response.completed leaves no usage behind to

tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py

@@ -1,5 +1,6 @@
 import asyncio
 import base64
+import contextvars
 import gzip
 import json
 import logging
@@ -39,6 +40,8 @@ from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.proxy._lazy_features import LazyFeature, attach_lazy_features
 from litellm.proxy._types import ProxyException, UserAPIKeyAuth
+from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
+from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3
 from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
 from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import router as llm_passthrough_router
 from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import (
@@ -77,7 +80,7 @@ from litellm.proxy.pass_through_endpoints.success_handler import (
 )
 from litellm.proxy.route_llm_request import ProxyModelNotFoundError
 from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id
-from litellm.proxy.utils import ProxyLogging
+from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token
 from litellm.types import utils as types_utils
 from litellm.types.guardrails import GuardrailEventHooks
 from litellm.types.utils import CallTypes
@@ -8204,6 +8207,89 @@ async def test_websocket_passthrough_logs_each_turn_on_its_own_row_charged_to_th
     relayed.proxy_logging.post_call_failure_hook.assert_not_awaited()
 
 
+@pytest.mark.asyncio
+async def test_websocket_passthrough_holds_the_keys_parallel_slot_until_the_session_ends():
+    """
+    A key allowed one parallel request holds that slot for as long as its
+    socket stays open. A turn that finishes and reaches spend mid-session must
+    not free it, or a second session gets in while the first is still running.
+    """
+    limiter_cache: Final = DualCache()
+    limiter: Final = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(limiter_cache))
+    key: Final = UserAPIKeyAuth(api_key=hash_token("sk-one-socket"), max_parallel_requests=1)
+    logged_rows: Final[list[asyncio.Task[None]]] = []
+
+    async def admit(user_api_key_dict: UserAPIKeyAuth, data: dict[str, object], call_type: str) -> None:
+        await limiter.async_pre_call_hook(
+            user_api_key_dict=user_api_key_dict, cache=limiter_cache, data=data, call_type=call_type
+        )
+
+    async def release(data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: object) -> None:
+        await limiter.async_post_call_success_hook(data=data, user_api_key_dict=user_api_key_dict, response=response)
+
+    async def log_row(**row: object) -> None:
+        await limiter.async_log_success_event(
+            kwargs=row, response_obj=None, start_time=row["start_time"], end_time=row["end_time"]
+        )
+
+    async def opens_another_session() -> bool:
+        await asyncio.gather(*logged_rows)
+        try:
+            await asyncio.create_task(
+                admit(key, {"litellm_call_id": str(uuid.uuid4())}, "pass_through_endpoint"),
+                context=contextvars.Context(),
+            )
+        except ProxyRateLimitError:
+            return False
+        return True
+
+    client: Final = _turn_taking_client(
+        [json.dumps({"type": "response.create", "model": "gpt-5.4-mini"})], answered_by='"response.completed"'
+    )
+    take_frame: Final = client.receive
+    while_open: Final[list[bool]] = []
+
+    async def receive() -> Mapping[str, str | bytes]:
+        message: Final = await take_frame()
+        if message["type"] == "websocket.disconnect":
+            while_open.append(await opens_another_session())
+        return message
+
+    client.receive = receive
+
+    with (
+        _patched_websocket_passthrough_environment(
+            ScriptedUpstreamWebSocket(replies=[_responses_turn("gpt-5.4-mini")])
+        ) as proxy_logging,
+        patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") as worker,
+        patch(SUCCESS_HANDLER, new_callable=AsyncMock, side_effect=log_row),
+    ):
+        proxy_logging.pre_call_hook = AsyncMock(side_effect=admit)
+        proxy_logging.post_call_success_hook = AsyncMock(side_effect=release)
+        worker.ensure_initialized_and_enqueue = MagicMock(
+            side_effect=lambda async_coroutine: logged_rows.append(asyncio.create_task(async_coroutine))
+        )
+        await asyncio.wait_for(
+            websocket_passthrough_request(
+                websocket=client,
+                target="wss://api.openai.com/v1/responses",
+                custom_headers={"Authorization": "Bearer sk-provider"},
+                user_api_key_dict=key,
+                forward_headers=False,
+                endpoint="/openai_passthrough/v1/responses",
+                accept_websocket=False,
+                client_frame_model_gate=RecordingFrameModelGate(),
+                upstream_frame_kind=openai_websocket_frame_kind,
+                walkout_turn_drain_seconds=0.05,
+            ),
+            timeout=10,
+        )
+        after_close: Final = await opens_another_session()
+
+    assert len(logged_rows) == 1
+    assert (while_open, after_close) == ([False], True)
+
+
 @pytest.mark.asyncio
 async def test_websocket_passthrough_charges_pipelined_turns_to_the_frame_that_asked_for_each():
     """

tests/unit/proxy/pass_through_endpoints/test_success_handler.py

@@ -668,7 +668,7 @@ REALTIME_FRAMES: Final = [
 
 
 def _normalize(url_route: str, response_body: dict[str, object] | list[dict[str, object]]) -> dict:
-    logging_obj: Final = MagicMock(model_call_details={})
+    logging_obj: Final = MagicMock(model_call_details={}, litellm_call_id="call-fake")
     return PassThroughEndpointLogging().normalize_llm_passthrough_logging_payload(
         httpx_response=MagicMock(spec=httpx.Response),
         response_body=response_body,