fix(rate-limits): pin the request stash to its owning litellm_call_id so nested calls cannot release it

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)
 

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(

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) -----------------------