From f507a118af66eda884026b7d4cfeaec9fcb805cd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 30 Jul 2026 14:54:28 -0700 Subject: [PATCH] fix(rate-limits): pin the request stash to its owning litellm_call_id so nested calls cannot release it --- .../proxy/hooks/dynamic_rate_limiter_v3.py | 2 + .../hooks/parallel_request_limiter_v3.py | 45 +++++-- .../hooks/test_parallel_request_limiter_v3.py | 115 ++++++++++++++++++ 3 files changed, 155 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 1aaeda6ba95..932146800e2 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -23,6 +23,7 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( RateLimitDescriptorRateLimitObject, RateLimitResponse, _PROXY_MaxParallelRequestsHandler_v3, + claim_request_stash_for_data, get_or_create_request_stash, ) from litellm.proxy.hooks.rate_limiter_utils import ( @@ -601,6 +602,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): if "model" not in data: return None + claim_request_stash_for_data(data) model = data["model"] priority = self._get_priority_from_user_api_key_dict(user_api_key_dict=user_api_key_dict) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 9d7423166ad..b04ef5f7087 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -354,8 +354,17 @@ class RequestRateLimiterStash: ``reservation_released`` flag and ``parallel_slot`` clearing effective across sibling callbacks: the first release wins, later callbacks observe the cleared state. + + Because the stash is context-inherited, nested LiteLLM calls made inside + the request (LLM-judge guardrails, silent experiments) would also see it + from their own logging callbacks. ``owner_litellm_call_id`` pins the stash + to the proxy request's ``litellm_call_id`` so those callbacks can tell the + owning request's events apart from a nested call's: router retries and + fallbacks reuse the request's call id and keep access, while nested calls + mint fresh ids and are ignored. """ + owner_litellm_call_id: Optional[str] = None rate_limit_response: Optional[RateLimitResponse] = None parallel_slot: Optional[ParallelSlotAcquisition] = None reserved_tokens: int = 0 @@ -381,6 +390,30 @@ def get_or_create_request_stash() -> RequestRateLimiterStash: return stash +def claim_request_stash_for_data(data: dict) -> RequestRateLimiterStash: + stash = get_or_create_request_stash() + owner_call_id = data.get("litellm_call_id") + if isinstance(owner_call_id, str): + stash.owner_litellm_call_id = owner_call_id + return stash + + +def get_request_stash_for_call(litellm_call_id: Optional[str]) -> Optional[RequestRateLimiterStash]: + stash = _request_stash.get() + if stash is None: + return None + if stash.owner_litellm_call_id is None or litellm_call_id is None: + return stash + return stash if litellm_call_id == stash.owner_litellm_call_id else None + + +def _call_id_from_callback_kwargs(kwargs: object) -> Optional[str]: + if not isinstance(kwargs, dict): + return None + call_id = kwargs.get("litellm_call_id") + return call_id if isinstance(call_id, str) else None + + class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def __init__( self, @@ -2342,7 +2375,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ verbose_proxy_logger.debug("Inside Rate Limit Pre-Call Hook") - stash = get_or_create_request_stash() + stash = claim_request_stash_for_data(data) ######################################################### # Check if the call type has a specific rate limiter @@ -2536,8 +2569,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): stored_response = stash.rate_limit_response if stored_response is not None: stored_response["statuses"].extend(tpm_response["statuses"]) - elif tpm_response["statuses"]: - stash.rate_limit_response = tpm_response verbose_proxy_logger.debug(f"TPM tokens reserved: {estimated_tokens} for model {requested_model}") @@ -2886,7 +2917,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if total_tokens == 0: total_tokens = self._aggregate_only_total_tokens(usage=_usage) - stash = get_request_stash() + stash = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs)) reserved_tokens = stash.reserved_tokens if stash is not None else 0 reserved_model = stash.reserved_model if stash is not None else None reserved_scopes: FrozenSet[Tuple[str, str]] = stash.reserved_scopes if stash is not None else frozenset() @@ -2945,7 +2976,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): try: verbose_proxy_logger.debug("INSIDE parallel request limiter ASYNC SUCCESS LOGGING") - stash = get_request_stash() + stash = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs)) acquisition = stash.parallel_slot if stash is not None else None if stash is not None and acquisition is not None: await self._release_parallel_request_slots( @@ -3002,7 +3033,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if not isinstance(kwargs, dict): return - stash = get_request_stash() + stash = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs)) rate_limit_response = stash.rate_limit_response if stash is not None else None statuses = rate_limit_response["statuses"] if rate_limit_response is not None else [] if not statuses: @@ -3044,7 +3075,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): pipeline_operations: List[RedisPipelineIncrementOperation] = [] - stash = get_request_stash() + stash = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs)) acquisition = stash.parallel_slot if stash is not None else None if stash is not None and acquisition is not None: await self._release_parallel_request_slots( diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index d546b629f0f..56bfd1829b5 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -20,6 +20,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.parallel_request_limiter_v3 import ( PARALLEL_REQUEST_SLOT_TTL_SECONDS, ParallelSlotAcquisition, + RequestRateLimiterStash, _request_stash, get_or_create_request_stash, get_request_stash, @@ -3345,6 +3346,120 @@ async def test_pre_call_hook_ignores_caller_supplied_stash_values(): assert stash.reserved_tokens == 0 +@pytest.mark.asyncio +async def test_log_events_from_nested_calls_leave_owner_stash_alone(monkeypatch): + """ + A nested LiteLLM call made inside the request (LLM-judge guardrail, + silent experiment) inherits the request context and fires the same global + logging callbacks with a fresh ``litellm_call_id``. Those callbacks must + not release the owning request's parallel slot or refund its TPM + reservation; only events carrying the owner's call id may. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + _api_key = hash_token("sk-nested-guard") + 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" + ) + parallel_key = f"{{api_key:{_api_key}}}:max_parallel_requests" + + 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 + assert stash.owner_litellm_call_id == "owner-call-id" + reserved = stash.reserved_tokens + assert reserved > 0 + + nested_kwargs = { + "litellm_call_id": "nested-guardrail-call", + "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, + } + await handler.async_log_success_event( + kwargs=nested_kwargs, response_obj=None, start_time=None, end_time=None + ) + await handler.async_log_failure_event( + kwargs=nested_kwargs, response_obj=None, start_time=None, end_time=None + ) + + assert stash.parallel_slot is not None + assert stash.reservation_released is False + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) == 1 + assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == reserved + + owner_kwargs = { + "litellm_call_id": "owner-call-id", + "standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}, + } + await handler.async_log_failure_event( + kwargs=owner_kwargs, response_obj=None, start_time=None, end_time=None + ) + + assert stash.parallel_slot is None + assert stash.reservation_released is True + assert handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) == 0 + assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 0 + + +@pytest.mark.asyncio +async def test_stash_applies_when_owner_or_callback_call_id_missing(): + """ + The owner guard only rejects a positive mismatch. A stash never claimed + by a pre-call hook (no owner id) must stay visible to any callback, and a + claimed stash must stay visible to callbacks whose kwargs carry no call + id — otherwise reservations and slots would strand on request paths that + do not thread ``litellm_call_id`` into their logging kwargs. + """ + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + unclaimed = get_or_create_request_stash() + unclaimed.reserved_tokens = 42 + await handler.async_log_failure_event( + kwargs={"litellm_call_id": "any-id", "standard_logging_object": {}}, + response_obj=None, + start_time=None, + end_time=None, + ) + assert unclaimed.reservation_released is True + + claimed = RequestRateLimiterStash( + owner_litellm_call_id="owner-1", reserved_tokens=42 + ) + _request_stash.set(claimed) + await handler.async_log_failure_event( + kwargs={"standard_logging_object": {}}, + response_obj=None, + start_time=None, + end_time=None, + ) + assert claimed.reservation_released is True + + # ----------------------- Per-MCP-server rate limiting (v3) -----------------------