fix(proxy): pin the limiter slot owner to the proxy call id and refund unsettled reservations on disconnect

The v3 limiter stashed whichever litellm_call_id the request body carried as
the slot owner, so a client-supplied id on a pass-through body left the
proxy's own success and failure events unable to release the slot. The owner
is now the proxy's logging call id. A WebSocket connection whose turns log
under their own ids never settles the TPM, ITPM, or OTPM estimate the
connection reserved, so the disconnect release now refunds any reservation
no callback settled, once.

litellm/proxy/hooks/parallel_request_limiter_v3.py

@@ -740,9 +740,11 @@ def get_or_create_request_stash() -> RequestRateLimiterStash:
 
 
 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 = data.get("litellm_call_id")
-    if isinstance(owner_call_id, str):
+    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
 
@@ -5092,6 +5094,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
                         pipeline_operations=io_token_operations,
                         parent_otel_span=litellm_parent_otel_span,
                     )
+            if stash is not None:
+                stash.reservation_released = True
 
         except Exception as e:
             verbose_proxy_logger.exception("Error in rate limit success event: %s", e)
@@ -5293,8 +5297,29 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
         TTL prunes it. The stashed acquisition's presence (not the key
         object's current max_parallel_requests configuration, which can
         change mid-request) decides whether there is anything to release.
+        A TPM, ITPM or OTPM reservation no callback settled is refunded the
+        same way, so a native Responses WebSocket connection, whose turns log
+        under their own call ids, hands back the estimate it reserved.
         """
-        await self._release_stashed_parallel_slot(get_request_stash(), None)
+        stash: Final = get_request_stash()
+        await self._release_stashed_parallel_slot(stash, None)
+        await self._refund_unsettled_reservation(stash)
+
+    async def _refund_unsettled_reservation(self, stash: RequestRateLimiterStash | None) -> None:
+        if stash is None or stash.reservation_released:
+            return
+        stash.reservation_released = True
+        await self._refund_reserved_tokens(scopes=tuple(stash.reserved_scopes), amount=stash.reserved_tokens)
+        await self._refund_reserved_tokens(
+            scopes=tuple(stash.itpm_reserved_scopes),
+            amount=stash.itpm_reserved_tokens,
+            reservation_windows=stash.itpm_reserved_window_identities,
+        )
+        await self._refund_reserved_tokens(
+            scopes=tuple(stash.otpm_reserved_scopes),
+            amount=stash.otpm_reserved_tokens,
+            reservation_windows=stash.otpm_reserved_window_identities,
+        )
 
     @with_service_target("rate_limits")
     async def async_post_call_success_hook(self, data: dict, user_api_key_dict: UserAPIKeyAuth, response) -> None:

tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py

@@ -20,6 +20,7 @@ from litellm import Router
 from litellm.caching.caching import DualCache
 from litellm.caching.in_memory_cache import InMemoryCache
 from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
+from litellm.litellm_core_utils.litellm_logging import Logging
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
 from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
@@ -4235,6 +4236,77 @@ async def test_release_on_disconnect_works_when_key_config_changed_v3():
     ) == 0
 
 
+async def _admit_with_tpm_reservation(
+    handler: _PROXY_MaxParallelRequestsHandler, local_cache: DualCache, user_api_key_dict: UserAPIKeyAuth
+) -> int:
+    await handler.async_pre_call_hook(
+        user_api_key_dict=user_api_key_dict,
+        cache=local_cache,
+        data={
+            "model": "gpt-4o-mini",
+            "messages": [{"role": "user", "content": "hello"}],
+            "max_tokens": 50,
+            "litellm_call_id": "owner-call-id",
+        },
+        call_type="completion",
+    )
+    stash = get_request_stash()
+    assert stash is not None and stash.reserved_tokens > 0
+    return stash.reserved_tokens
+
+
+@pytest.mark.asyncio
+async def test_release_on_disconnect_refunds_an_unsettled_tpm_reservation_v3(monkeypatch):
+    """
+    A native Responses WebSocket connection reserves tokens at pre-call and logs its
+    turns under their own call ids, so no callback ever settles the connection's
+    reservation. The disconnect release must hand the estimate back.
+    """
+    monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False)
+    _api_key = hash_token("sk-ws-unsettled-reservation")
+    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)
+    assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == reserved
+
+    await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict)
+
+    assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 0
+    assert get_request_stash().reservation_released is True
+
+
+@pytest.mark.asyncio
+async def test_release_on_disconnect_keeps_a_settled_tpm_reservation_v3(monkeypatch):
+    """
+    When the owner's success event already settled the reservation to the actual
+    usage, the disconnect release must not refund the estimate a second time.
+    """
+    monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False)
+    _api_key = hash_token("sk-ws-settled-reservation")
+    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")
+    await _admit_with_tpm_reservation(handler, local_cache, user_api_key_dict)
+
+    await handler.async_log_success_event(
+        kwargs={
+            "litellm_call_id": "owner-call-id",
+            "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}},
+        },
+        response_obj=ModelResponse(usage=Usage(prompt_tokens=7, completion_tokens=3, total_tokens=10)),
+        start_time=datetime.now(),
+        end_time=datetime.now(),
+    )
+    assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 10
+
+    await handler.async_release_max_parallel_requests_on_disconnect(user_api_key_dict)
+
+    assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 10
+
+
 @pytest.mark.asyncio
 async def test_post_call_failure_hook_releases_parallel_slot_v3():
     """
@@ -4328,6 +4400,75 @@ async def test_post_call_failure_hook_for_another_call_keeps_the_owner_slot_v3()
     assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0
 
 
+def _pass_through_logging_obj(litellm_call_id: str) -> Logging:
+    return Logging(
+        model="gpt-4o-mini",
+        messages=[],
+        stream=False,
+        call_type="pass_through_endpoint",
+        start_time=datetime.now(),
+        litellm_call_id=litellm_call_id,
+        function_id="pass-through",
+    )
+
+
+async def _admit_pass_through_with_client_call_id(
+    handler: _PROXY_MaxParallelRequestsHandler, local_cache: DualCache, user_api_key_dict: UserAPIKeyAuth
+) -> Logging:
+    """
+    A pass-through body may carry its own litellm_call_id. The proxy books the request under the id of the
+    logging object it put in the same data, so that id, not the client's, must own the parallel slot.
+    """
+    logging_obj = _pass_through_logging_obj("proxy-call")
+    await handler.async_pre_call_hook(
+        user_api_key_dict=user_api_key_dict,
+        cache=local_cache,
+        data={"litellm_call_id": "client-chosen-id", "litellm_logging_obj": logging_obj},
+        call_type="pass_through_endpoint",
+    )
+    return logging_obj
+
+
+@pytest.mark.asyncio
+async def test_post_call_failure_hook_releases_the_slot_the_proxy_call_id_owns_v3():
+    _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)
+
+    await handler.async_post_call_failure_hook(
+        request_data={"litellm_call_id": "proxy-call", "litellm_logging_obj": logging_obj},
+        original_exception=Exception("upstream unreachable"),
+        user_api_key_dict=user_api_key_dict,
+    )
+
+    assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0
+
+
+@pytest.mark.asyncio
+async def test_success_event_releases_the_slot_the_proxy_call_id_owns_v3():
+    _api_key = hash_token("sk-passthrough-client-id-success")
+    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"
+    await _admit_pass_through_with_client_call_id(handler, local_cache, user_api_key_dict)
+
+    await handler.async_log_success_event(
+        kwargs={
+            "litellm_call_id": "proxy-call",
+            "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}},
+        },
+        response_obj=ModelResponse(usage=Usage(prompt_tokens=5, completion_tokens=5, total_tokens=10)),
+        start_time=datetime.now(),
+        end_time=datetime.now(),
+    )
+
+    assert handler._gauge_in_flight_from_cache_value(await local_cache.async_get_cache(key=counter_key)) == 0
+
+
 @pytest.mark.asyncio
 async def test_success_event_releases_parallel_slot_v3(monkeypatch):
     """

tests/unit/test_router/test_router_helper_utils.py

@@ -482,6 +482,52 @@ 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():
+    """Each native Responses WebSocket turn logs under its own call, so the deployment's
+    tpm and rpm counters move once per turn, while a realtime session still counts nothing."""
+    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
+
+
 @pytest.mark.asyncio
 async def test_deployment_callback_on_success_tracks_tpm_for_io_deployment():
     """