From 6ccd11643f387ef7b576224c59ed9920ee92d39d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 15 Oct 2026 19:44:50 -0700 Subject: [PATCH] fix(proxy): keep the legacy limiter's WebSocket slot until the socket closes --- .../proxy/hooks/parallel_request_limiter.py | 77 +++++++++-- litellm/proxy/utils.py | 3 + .../hooks/test_parallel_request_limiter.py | 130 +++++++++++++++++- 3 files changed, 193 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index 19251b53f13..54f4fbf562a 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -1,6 +1,8 @@ import asyncio import sys from collections.abc import Callable +from contextvars import ContextVar +from dataclasses import dataclass from datetime import datetime, timedelta from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn @@ -23,7 +25,7 @@ from litellm.proxy.auth.budget_throttle import throttled_limit from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup -from litellm.types.utils import Usage +from litellm.types.utils import CallTypes, Usage if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -48,6 +50,17 @@ def _response_total_tokens(response_obj: object) -> int: return response_usage.total_tokens if isinstance(response_usage, Usage) else 0 +@dataclass(frozen=True, slots=True) +class _WebSocketConnectionSlot: + request_count_key: str | None + holds_global_slot: bool + + +_websocket_connection_slot: Final[ContextVar[_WebSocketConnectionSlot | None]] = ContextVar( + "litellm_legacy_rate_limiter_websocket_connection_slot", default=None +) + + class CacheObject(TypedDict): current_global_requests: dict | None request_count_api_key: dict | None @@ -63,7 +76,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): self.internal_usage_cache = internal_usage_cache self._clock: Final = clock - def print_verbose(self, print_statement) -> None: + def print_verbose(self, print_statement: object) -> None: try: verbose_proxy_logger.debug(print_statement) if litellm.set_verbose: @@ -499,6 +512,41 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): litellm_parent_otel_span=user_api_key_dict.parent_otel_span, ) # don't block execution for cache updates ) + if call_type == CallTypes.aresponses_websocket.value: + _websocket_connection_slot.set( + _WebSocketConnectionSlot( + request_count_key=(f"{api_key}::{precise_minute}::request_count" if api_key is not None else None), + holds_global_slot=global_max_parallel_requests is not None, + ) + ) + + @with_service_target("rate_limits") + async def async_release_websocket_connection_slot(self) -> None: + connection: Final = _websocket_connection_slot.get() + if connection is None: + return + _websocket_connection_slot.set(None) + if connection.holds_global_slot: + await self.internal_usage_cache.async_increment_cache( + key="global_max_parallel_requests", + value=-1, + local_only=True, + litellm_parent_otel_span=None, + ) + if connection.request_count_key is None: + return + current: Final = await self.internal_usage_cache.async_get_cache( + key=connection.request_count_key, + litellm_parent_otel_span=None, + ) + if current is None: + return + await self.internal_usage_cache.async_set_cache( + connection.request_count_key, + {**current, "current_requests": max(current["current_requests"] - 1, 0)}, + ttl=60, + litellm_parent_otel_span=None, + ) @with_service_target("rate_limits") async def async_log_success_event(self, kwargs, response_obj: object, start_time, end_time) -> None: @@ -520,6 +568,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): "user_api_key_model_max_budget", None ) user_api_key_end_user_id: Final = kwargs.get("user") + requests_released: Final = 0 if _websocket_connection_slot.get() is not None else 1 user_api_key_metadata: Final = kwargs["litellm_params"]["metadata"].get("user_api_key_metadata", {}) or {} user_api_key_team_metadata = kwargs["litellm_params"]["metadata"].get("user_api_key_team_metadata", None) @@ -540,7 +589,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): # decrement await self.internal_usage_cache.async_increment_cache( key=_key, - value=-1, + value=-requests_released, local_only=True, litellm_parent_otel_span=litellm_parent_otel_span, ) @@ -562,13 +611,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): key=request_count_api_key, litellm_parent_otel_span=litellm_parent_otel_span, ) or { - "current_requests": 1, + "current_requests": requests_released, "current_tpm": 0, "current_rpm": 0, } new_val = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - requests_released, 0), "current_tpm": current["current_tpm"] + total_tokens, "current_rpm": current["current_rpm"], } @@ -603,13 +652,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): key=request_count_api_key, litellm_parent_otel_span=litellm_parent_otel_span, ) or { - "current_requests": 1, + "current_requests": requests_released, "current_tpm": 0, "current_rpm": 0, } new_val = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - requests_released, 0), "current_tpm": current["current_tpm"] + total_tokens, "current_rpm": current["current_rpm"], } @@ -629,13 +678,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): key=request_count_api_key, litellm_parent_otel_span=litellm_parent_otel_span, ) or { - "current_requests": 1, + "current_requests": requests_released, "current_tpm": total_tokens, "current_rpm": 1, } new_val = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - requests_released, 0), "current_tpm": current["current_tpm"] + total_tokens, "current_rpm": current["current_rpm"], } @@ -655,13 +704,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): key=request_count_api_key, litellm_parent_otel_span=litellm_parent_otel_span, ) or { - "current_requests": 1, + "current_requests": requests_released, "current_tpm": total_tokens, "current_rpm": 1, } new_val = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - requests_released, 0), "current_tpm": current["current_tpm"] + total_tokens, "current_rpm": current["current_rpm"], } @@ -681,13 +730,13 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): key=request_count_api_key, litellm_parent_otel_span=litellm_parent_otel_span, ) or { - "current_requests": 1, + "current_requests": requests_released, "current_tpm": total_tokens, "current_rpm": 1, } new_val = { - "current_requests": max(current["current_requests"] - 1, 0), + "current_requests": max(current["current_requests"] - requests_released, 0), "current_tpm": current["current_tpm"] + total_tokens, "current_rpm": current["current_rpm"], } @@ -712,7 +761,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): global_max_parallel_requests: Final = _metadata.get("global_max_parallel_requests", None) user_api_key: Final = _metadata.get("user_api_key", None) self.print_verbose(f"user_api_key: [set={user_api_key is not None}]") - if user_api_key is None: + if user_api_key is None or _websocket_connection_slot.get() is not None: return ## decrement call count if call failed diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9788c4635c9..bc70604a5b2 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4405,6 +4405,9 @@ class ProxyLogging: double-decrement under the limiter's in-memory fallback. """ limiter: Final = self.get_proxy_hook("parallel_request_limiter") + if isinstance(limiter, PROXY_MaxParallelRequestsHandler): + await limiter.async_release_websocket_connection_slot() + return if not isinstance(limiter, PROXY_MaxParallelRequestsHandler_v3): return await limiter.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter.py b/tests/unit/proxy/hooks/test_parallel_request_limiter.py index 89cb6bd1ef2..3f5eef4593f 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter.py @@ -2,8 +2,9 @@ Unit Tests for the max parallel request limiter v1 for the proxy """ +import asyncio import itertools -from collections.abc import Callable, Iterator, Mapping +from collections.abc import Awaitable, Callable, Iterator, Mapping from datetime import datetime from typing import Final @@ -12,11 +13,12 @@ import pytest from litellm.caching.caching import DualCache 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 from litellm.proxy.hooks.parallel_request_limiter import ( PROXY_MaxParallelRequestsHandler, ) -from litellm.proxy.utils import InternalUsageCache, hash_token -from litellm.types.utils import EmbeddingResponse, ModelResponse, TextCompletionResponse, Usage +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) LAST_MICROSECOND_OF_JANUARY: Final = datetime(2026, 1, 31, 23, 59, 59, 999999) @@ -323,3 +325,125 @@ async def test_retry_after_is_the_seconds_until_the_next_minute_of_the_injected_ ) assert exc_info.value.headers["retry-after"] == "14.5" + + +async def _log_websocket_turn_success( + handler: PROXY_MaxParallelRequestsHandler, metadata: Mapping[str, object] +) -> None: + await handler.async_log_success_event( + kwargs={"litellm_params": {"metadata": dict(metadata)}}, + response_obj=None, + start_time=FROZEN_INSTANT, + end_time=FROZEN_INSTANT, + ) + + +async def _log_websocket_turn_failure( + handler: PROXY_MaxParallelRequestsHandler, metadata: Mapping[str, object] +) -> None: + await handler.async_log_failure_event( + kwargs={"litellm_params": {"metadata": dict(metadata)}, "exception": Exception("upstream turn failed")}, + response_obj=None, + start_time=FROZEN_INSTANT, + end_time=FROZEN_INSTANT, + ) + + +def _proxy_logging_with(handler: PROXY_MaxParallelRequestsHandler) -> ProxyLogging: + proxy_logging: Final = ProxyLogging(user_api_key_cache=UserApiKeyCache()) + proxy_logging.proxy_hook_mapping["parallel_request_limiter"] = handler + return proxy_logging + + +async def _requests_admitted( + handler: PROXY_MaxParallelRequestsHandler, session: UserAPIKeyAuth, metadata: Mapping[str, object] +) -> int: + for admitted in range(3): + try: + await handler.async_pre_call_hook( + user_api_key_dict=session, + cache=DualCache(), + data={"metadata": dict(metadata)}, + call_type=CallTypes.acompletion.value, + ) + except ProxyRateLimitError: + return admitted + await asyncio.sleep(0) + return 3 + + +@pytest.mark.parametrize("log_turn", [_log_websocket_turn_success, _log_websocket_turn_failure]) +@pytest.mark.parametrize( + "max_parallel_requests,limit_metadata,admitted_while_open", + [(1, {}, 0), (None, {"global_max_parallel_requests": 2}, 1)], + ids=["key_limit", "global_limit"], +) +@pytest.mark.asyncio +async def test_a_websocket_connection_holds_its_parallel_slot_through_its_turns_until_it_closes( + log_turn: Callable[[PROXY_MaxParallelRequestsHandler, Mapping[str, object]], Awaitable[None]], + max_parallel_requests: int | None, + limit_metadata: Mapping[str, object], + admitted_while_open: int, +): + handler: Final = PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()), clock=_frozen_clock + ) + api_key: Final = hash_token("sk-websocket-slot") + session: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=max_parallel_requests) + turn_logged: Final = asyncio.Event() + socket_closed: Final = asyncio.Event() + + async def websocket_connection() -> None: + await handler.async_pre_call_hook( + user_api_key_dict=session, + cache=DualCache(), + data={"metadata": dict(limit_metadata)}, + call_type=CallTypes.aresponses_websocket.value, + ) + await asyncio.sleep(0) + await log_turn(handler, {"user_api_key": api_key, "user_api_key_model_max_budget": {}, **limit_metadata}) + turn_logged.set() + await socket_closed.wait() + await _proxy_logging_with(handler).arelease_max_parallel_requests_on_disconnect(session) + + connection: Final = asyncio.create_task(websocket_connection()) + await turn_logged.wait() + assert await _requests_admitted(handler, session, limit_metadata) == admitted_while_open + + socket_closed.set() + await connection + assert await _requests_admitted(handler, session, limit_metadata) == 1 + + +@pytest.mark.asyncio +async def test_a_websocket_turn_logged_in_a_later_minute_leaves_no_slot_held_after_the_socket_closes(): + handler: Final = PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()), clock=_clock_rolling_over_after_first_read() + ) + api_key: Final = hash_token("sk-websocket-rollover") + session: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + + async def websocket_connection() -> None: + await handler.async_pre_call_hook( + user_api_key_dict=session, cache=DualCache(), data={}, call_type=CallTypes.aresponses_websocket.value + ) + await asyncio.sleep(0) + await _log_websocket_turn_success(handler, {"user_api_key": api_key, "user_api_key_model_max_budget": {}}) + await handler.async_release_websocket_connection_slot() + + await asyncio.create_task(websocket_connection()) + + assert await _requests_admitted(handler, session, {}) == 1 + + +@pytest.mark.asyncio +async def test_a_disconnect_release_outside_a_websocket_connection_leaves_the_request_slot_held(): + handler: Final = PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()), clock=_frozen_clock + ) + session: Final = UserAPIKeyAuth(api_key=hash_token("sk-http-slot"), max_parallel_requests=1) + + assert await _requests_admitted(handler, session, {}) == 1 + await _proxy_logging_with(handler).arelease_max_parallel_requests_on_disconnect(session) + + assert await _requests_admitted(handler, session, {}) == 0