diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index b7f57d93a9b..a44f1d10fda 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1906,7 +1906,7 @@ async def _responses_websocket_session( ) await websocket.close(code=1011, reason="Internal server error") finally: - await proxy_logging_obj.arelease_max_parallel_requests_on_disconnect(user_api_key_dict) + await proxy_logging_obj.arelease_responses_websocket_slot(user_api_key_dict) @router.websocket("/v1/responses") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 5d2ebe73cac..be6532ba36b 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4405,12 +4405,20 @@ class ProxyLogging: double-decrement under the limiter's in-memory fallback. """ limiter: Final = self.get_proxy_hook("parallel_request_limiter") - if not isinstance(limiter, (PROXY_MaxParallelRequestsHandler_v3, _PROXY_MaxParallelRequestsHandler)): + if not isinstance(limiter, PROXY_MaxParallelRequestsHandler_v3): return await limiter.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) _arelease_max_parallel_requests_on_disconnect = arelease_max_parallel_requests_on_disconnect + async def arelease_responses_websocket_slot(self, user_api_key_dict: UserAPIKeyAuth) -> None: + """Release the slot a Responses WebSocket took at admission. Only call this once that admission succeeded: + the legacy limiter keeps no record of its acquisitions and decrements unconditionally.""" + limiter: Final = self.get_proxy_hook("parallel_request_limiter") + if not isinstance(limiter, (PROXY_MaxParallelRequestsHandler_v3, _PROXY_MaxParallelRequestsHandler)): + return + await limiter.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) + async def enforce_mcp_server_rate_limits( self, user_api_key_dict: UserAPIKeyAuth | None, diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter.py b/tests/unit/proxy/hooks/test_parallel_request_limiter.py index f17cad56751..08840978169 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter.py @@ -15,7 +15,7 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitErro from litellm.proxy.hooks.parallel_request_limiter import ( PROXY_MaxParallelRequestsHandler, ) -from litellm.proxy.utils import InternalUsageCache, hash_token +from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token from litellm.types.utils import CallTypes, EmbeddingResponse, ModelResponse, TextCompletionResponse, Usage FROZEN_INSTANT: Final = datetime(2026, 1, 31, 23, 59, 30) @@ -353,3 +353,26 @@ async def test_websocket_turn_callbacks_keep_the_sockets_slot_until_the_disconne await handler.async_release_max_parallel_requests_on_disconnect(UserAPIKeyAuth(api_key=api_key)) after_disconnect = await handler.internal_usage_cache.async_get_cache(key=bucket, litellm_parent_otel_span=None) assert after_disconnect["current_requests"] == 0, after_disconnect + + +@pytest.mark.asyncio +async def test_a_rejected_realtime_attempt_does_not_free_a_slot_it_never_took(): + api_key = hash_token("sk-at-limit") + handler = PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache()), clock=_frozen_clock) + bucket = f"{api_key}::{FROZEN_INSTANT.strftime('%Y-%m-%d-%H-%M')}::request_count" + await handler.internal_usage_cache.async_set_cache( + key=bucket, + value={"current_requests": 1, "current_tpm": 0, "current_rpm": 1}, + litellm_parent_otel_span=None, + ) + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = handler + auth = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + + await proxy_logging_obj.arelease_max_parallel_requests_on_disconnect(auth) + after_rejection = await handler.internal_usage_cache.async_get_cache(key=bucket, litellm_parent_otel_span=None) + assert after_rejection["current_requests"] == 1, after_rejection + + await proxy_logging_obj.arelease_responses_websocket_slot(auth) + after_socket_close = await handler.internal_usage_cache.async_get_cache(key=bucket, litellm_parent_otel_span=None) + assert after_socket_close["current_requests"] == 0, after_socket_close diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 0eaf619ba70..b047c29243a 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -1952,7 +1952,7 @@ class TestResponsesWSFirstFrameModelAuth: proxy_logging_obj = MagicMock() proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) - proxy_logging_obj.arelease_max_parallel_requests_on_disconnect = AsyncMock(return_value=None) + proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(return_value=None) user_api_key_dict = MagicMock() with ( @@ -2017,7 +2017,7 @@ class TestResponsesWSFirstFrameModelAuth: ) proxy_logging_obj = MagicMock() proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) - proxy_logging_obj.arelease_max_parallel_requests_on_disconnect = AsyncMock(return_value=None) + proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(return_value=None) user_api_key_dict = MagicMock() with (