From 7a462a4220cd10d7486d4ac956e9fd68d4396bbc Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 14 May 2026 10:53:04 -0700 Subject: [PATCH] fix(rate-limit): stop v3 limiter from leaking internal stash to provider body (#27913) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(rate-limit): stop v3 limiter from leaking internal stash to provider body PR #27001 (atomic TPM rate limit) introduced a reservation flow that writes four LiteLLM-internal keys onto the request data dict: _litellm_rate_limit_descriptors _litellm_tpm_reserved_tokens _litellm_tpm_reserved_model _litellm_tpm_reserved_scopes _litellm_tpm_reservation_released These keys are forwarded as request body params to the upstream provider, which rejects them as unknown fields: OpenAI -> 400 'Unknown parameter: _litellm_rate_limit_descriptors' (mapped by litellm to RateLimitError / 429, hiding the bug behind a misleading 'throttling_error' code) Anthropic -> 400 '_litellm_rate_limit_descriptors: Extra inputs are not permitted' Net effect: every chat completion against any real provider fails the moment a virtual key has any tpm_limit / rpm_limit set — i.e. v3-enforced key-level TPM/RPM limits are broken end-to-end. The v3 RPM/TPM check itself still runs (raises 429 on over-limit), but the success path poisons the upstream body. Reproduced on litellm_internal_staging HEAD (410ce761dc) against gpt-4o-mini and claude-haiku-4-5 with a 1-RPM/1-TPM key — first request fails with the provider's unknown-field error. Fix: the stash is metadata only. - Add RATE_LIMIT_DESCRIPTORS_KEY constant and a _LITELLM_STASH_KEYS registry so we have a single source of truth for stash keys. - New helper _stash_value_in_metadata_channels writes to data['metadata'] / data['litellm_metadata'] without touching the top level. - _stash_reservation_in_data and the descriptor stash now route through that helper. _mark_reservation_released stops writing top-level. - _lookup_stashed_value also checks kwargs['metadata'] / kwargs['litellm_metadata'] (raw request_data shape) in addition to kwargs['litellm_params']['metadata'] (completion kwargs shape). - async_post_call_failure_hook now reads descriptors via the unified metadata lookup instead of request_data.get(top-level). - Defense in depth: async_pre_call_hook strips any stash key that somehow surfaced at the top level (stale cache, future refactor, test fixture) before returning. Tests: - New regression test asserts no _litellm_* stash key is present at the top level of data after async_pre_call_hook, and that the metadata channel still carries the reservation + descriptors so success / failure reconciliation works. - Existing test_tpm_concurrent.py tests that asserted top-level presence are updated to read from data['metadata'] — the location is an implementation detail; the spec is that post-call callbacks can resolve the stash. Verified end-to-end against OpenAI gpt-4o-mini and Anthropic claude-haiku-4-5 via /v1/chat/completions on a low-rpm key: - With limits not exceeded: HTTP 200, valid completion response, no leaked fields in body. - With RPM exceeded: HTTP 429 from v3 enforcement ('Rate limit exceeded ... Limit type: requests'). - With TPM exceeded: HTTP 429 from v3 enforcement ('Rate limit exceeded ... Limit type: tokens'). Full v3 hook test suite passes (171 tests). Co-authored-by: Mateo Wang * chore(rate-limit): use RATE_LIMIT_DESCRIPTORS_KEY constant in test, trim noisy comments Address greptile P2: test fixture now uses the imported constant. Drop comments that re-explain what well-named identifiers already convey. * fix(rate-limit): reject caller-supplied stash values to prevent TPM-refund abuse Strip _LITELLM_STASH_KEYS from data top-level and both metadata channels at the start of async_pre_call_hook. Without this, an authenticated caller can inject _litellm_rate_limit_descriptors plus _litellm_tpm_reserved_tokens in body metadata, trigger a proxy-side rejection, and cause async_post_call_failure_hook to refund TPM counters against attacker-named scopes (e.g. another tenant's api_key). --------- Co-authored-by: Cursor Agent Co-authored-by: Mateo Wang --- .../hooks/parallel_request_limiter_v3.py | 134 ++++++++++++------ .../hooks/test_parallel_request_limiter_v3.py | 118 +++++++++++++++ .../proxy/hooks/test_tpm_concurrent.py | 26 ++-- 3 files changed, 221 insertions(+), 57 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index cd797483b2..283a3d8d10 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -224,6 +224,17 @@ TPM_RESERVED_SCOPES_KEY = "_litellm_tpm_reserved_scopes" # (e.g. async_log_failure_event firing after async_post_call_failure_hook) # does not double-refund. TPM_RESERVATION_RELEASED_KEY = "_litellm_tpm_reservation_released" +RATE_LIMIT_DESCRIPTORS_KEY = "_litellm_rate_limit_descriptors" +# Stash keys live ONLY in metadata channels — never at the top level of the +# request body. Top-level keys are forwarded as body params to upstream +# providers, which reject unknown fields with 400/429 errors. +_LITELLM_STASH_KEYS: Tuple[str, ...] = ( + TPM_RESERVED_TOKENS_KEY, + TPM_RESERVED_MODEL_KEY, + TPM_RESERVED_SCOPES_KEY, + TPM_RESERVATION_RELEASED_KEY, + RATE_LIMIT_DESCRIPTORS_KEY, +) class RateLimitDescriptorRateLimitObject(TypedDict, total=False): @@ -1892,6 +1903,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ verbose_proxy_logger.debug("Inside Rate Limit Pre-Call Hook") + # Reject caller-supplied stash values before any read/write. Otherwise + # a client can inject ``_litellm_rate_limit_descriptors`` / + # ``_litellm_tpm_reserved_tokens`` in body ``metadata`` and have + # ``async_post_call_failure_hook`` refund TPM counters against scopes + # they name (e.g. another tenant's api_key). + self._strip_stash_keys_from_all_channels(data) + ######################################################### # Check if the call type has a specific rate limiter # eg. for Batch APIs we need to use the batch rate limiter to read the input file and count the tokens and requests @@ -2024,7 +2042,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): descriptors=descriptors, ) else: - data["_litellm_rate_limit_descriptors"] = descriptors + self._stash_value_in_metadata_channels( + data=data, + key=RATE_LIMIT_DESCRIPTORS_KEY, + value=descriptors, + ) # Capture the exact (key, value) scopes the reservation # incremented so post-call reconciliation only applies # the (actual - reserved) delta to those — unreserved @@ -2059,6 +2081,29 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): f"TPM tokens reserved: {estimated_tokens} for model {requested_model}" ) + # Defense-in-depth: scrub any stash key that escaped onto data + # top-level (stale cache hit, router pass, test fixture) before the + # body is forwarded to the provider. + self._strip_stash_keys_from_top_level(data) + + @staticmethod + def _strip_stash_keys_from_top_level(data: Any) -> None: + if not isinstance(data, dict): + return + for stash_key in _LITELLM_STASH_KEYS: + data.pop(stash_key, None) + + @classmethod + def _strip_stash_keys_from_all_channels(cls, data: Any) -> None: + if not isinstance(data, dict): + return + cls._strip_stash_keys_from_top_level(data) + for channel in ("metadata", "litellm_metadata"): + channel_dict = data.get(channel) + if isinstance(channel_dict, dict): + for stash_key in _LITELLM_STASH_KEYS: + channel_dict.pop(stash_key, None) + def _create_pipeline_operations( self, key: str, @@ -2233,18 +2278,29 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return specified_rate_limit_type @staticmethod + def _stash_value_in_metadata_channels( + data: Dict[str, Any], + key: str, + value: Any, + ) -> None: + for channel in ("metadata", "litellm_metadata"): + existing = data.get(channel) + if isinstance(existing, dict): + existing[key] = value + elif channel == "metadata": + # ``litellm_metadata`` is owned by the router; don't conjure + # it here. + data[channel] = {key: value} + + @classmethod def _stash_reservation_in_data( + cls, data: Dict[str, Any], estimated_tokens: int, reserved_model: Optional[str], reserved_scopes: Optional[List[Tuple[str, str]]] = None, ) -> None: """ - Persist the reservation amount, model, and reserved scopes into every - channel a callback might read from: top-level kwargs (via ``**data``), - request metadata, and litellm_metadata. Keeps reservation and - reconciliation in sync. - ``reserved_scopes`` is serialized as a list of [key, value] pairs so it round-trips through JSON-based metadata transports. """ @@ -2252,30 +2308,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): [[k, v] for k, v in reserved_scopes] if reserved_scopes else None ) - data[TPM_RESERVED_TOKENS_KEY] = estimated_tokens + cls._stash_value_in_metadata_channels( + data=data, key=TPM_RESERVED_TOKENS_KEY, value=estimated_tokens + ) if reserved_model: - data[TPM_RESERVED_MODEL_KEY] = reserved_model + cls._stash_value_in_metadata_channels( + data=data, key=TPM_RESERVED_MODEL_KEY, value=reserved_model + ) if scopes_payload is not None: - data[TPM_RESERVED_SCOPES_KEY] = scopes_payload - - for channel in ("metadata", "litellm_metadata"): - existing = data.get(channel) - if isinstance(existing, dict): - existing[TPM_RESERVED_TOKENS_KEY] = estimated_tokens - if reserved_model: - existing[TPM_RESERVED_MODEL_KEY] = reserved_model - if scopes_payload is not None: - existing[TPM_RESERVED_SCOPES_KEY] = scopes_payload - elif channel == "metadata": - # Only auto-create ``metadata`` (preserves prior behavior); - # ``litellm_metadata`` is set by the router and shouldn't be - # conjured here. - stash: Dict[str, Any] = {TPM_RESERVED_TOKENS_KEY: estimated_tokens} - if reserved_model: - stash[TPM_RESERVED_MODEL_KEY] = reserved_model - if scopes_payload is not None: - stash[TPM_RESERVED_SCOPES_KEY] = scopes_payload - data[channel] = stash + cls._stash_value_in_metadata_channels( + data=data, key=TPM_RESERVED_SCOPES_KEY, value=scopes_payload + ) @staticmethod def _lookup_stashed_value( @@ -2284,19 +2327,19 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): key: str, ) -> Any: """ - Resolve a stashed value from any of the channels the request data can - flow through to a callback. - - Checks (in priority order): - 1. kwargs (top-level data fields propagate via **data) - 2. kwargs["litellm_params"]["metadata"] (request metadata channel) - 3. standard_logging_metadata (covers tests that mock the SLO directly) + Resolve a stashed value from any metadata channel the request data + can flow through to a callback. Top-level ``kwargs`` is not checked + because stash keys must never live there. """ - candidate = kwargs.get(key) if isinstance(kwargs, dict) else None - if candidate is None: - litellm_params = ( - kwargs.get("litellm_params") if isinstance(kwargs, dict) else None - ) + candidate: Any = None + if isinstance(kwargs, dict): + for channel in ("metadata", "litellm_metadata"): + channel_dict = kwargs.get(channel) + if isinstance(channel_dict, dict) and key in channel_dict: + candidate = channel_dict.get(key) + if candidate is not None: + return candidate + litellm_params = kwargs.get("litellm_params") if isinstance(litellm_params, dict): lp_metadata = litellm_params.get("metadata") if isinstance(lp_metadata, dict): @@ -2390,7 +2433,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ if not isinstance(data, dict): return - data[TPM_RESERVATION_RELEASED_KEY] = True for channel in ("metadata", "litellm_metadata"): existing = data.get(channel) if isinstance(existing, dict): @@ -2811,9 +2853,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return # Refund directly against the descriptors we reserved against — - # the pre-call hook stashes them on the request data before - # success/failure callbacks run. - stashed = request_data.get("_litellm_rate_limit_descriptors") + # the pre-call hook stashes them in the request-data metadata + # channels before success/failure callbacks run. + stashed = self._lookup_stashed_value( + kwargs=request_data, + standard_logging_metadata=None, + key=RATE_LIMIT_DESCRIPTORS_KEY, + ) descriptors: List[RateLimitDescriptor] = ( stashed if isinstance(stashed, list) else [] ) 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 e9ac1794ac..3e2eb4b02c 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 @@ -2775,3 +2775,121 @@ async def test_project_model_rate_limits_not_triggered_for_other_model_v3(): assert ( "model_per_project" not in descriptor_keys ), f"model_per_project should not be added for unrelated model, got: {descriptor_keys}" + + +@pytest.mark.asyncio +async def test_pre_call_hook_does_not_leak_internal_stash_to_request_body(): + """Regression for #27001: stash keys must stay in metadata, never on + the top level of ``data`` (which gets forwarded as the provider body).""" + from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _LITELLM_STASH_KEYS, + RATE_LIMIT_DESCRIPTORS_KEY, + TPM_RESERVED_TOKENS_KEY, + ) + + _api_key = hash_token("sk-leak-regression") + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + tpm_limit=1000, + rpm_limit=5, + ) + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache), + ) + + async def mock_should_rate_limit(descriptors, **kwargs): + return {"overall_code": "OK", "statuses": []} + + async def mock_reserve_tpm_tokens(descriptors, estimated_tokens, **kwargs): + return { + "overall_code": "OK", + "statuses": [ + { + "code": "OK", + "current_limit": 1000, + "limit_remaining": 1000 - estimated_tokens, + "descriptor_key": d["key"], + "descriptor_value": d["value"], + "rate_limit_type": "tokens", + } + for d in descriptors + ], + } + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + parallel_request_handler.reserve_tpm_tokens = mock_reserve_tpm_tokens + + data: Dict[str, Any] = { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 10, + } + + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=data, + call_type="completion", + ) + + leaked = [k for k in _LITELLM_STASH_KEYS if k in data] + assert not leaked, f"stash keys leaked to top level: {leaked}" + + metadata = data.get("metadata") or {} + assert metadata.get(TPM_RESERVED_TOKENS_KEY) + assert isinstance(metadata.get(RATE_LIMIT_DESCRIPTORS_KEY), list) + + +@pytest.mark.asyncio +async def test_pre_call_hook_rejects_caller_supplied_stash_values(): + """Caller cannot pre-populate stash keys in body metadata to drive a + later TPM refund against an arbitrary scope.""" + from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _LITELLM_STASH_KEYS, + RATE_LIMIT_DESCRIPTORS_KEY, + TPM_RESERVED_TOKENS_KEY, + ) + + user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("sk-no-limits")) + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache), + ) + + victim_descriptors = [ + { + "key": "api_key", + "value": "victim-key-hash", + "rate_limit": {"tokens_per_unit": 10000, "window_size": 60}, + } + ] + data: Dict[str, Any] = { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + TPM_RESERVED_TOKENS_KEY: 9999, + RATE_LIMIT_DESCRIPTORS_KEY: victim_descriptors, + "metadata": { + TPM_RESERVED_TOKENS_KEY: 9999, + RATE_LIMIT_DESCRIPTORS_KEY: victim_descriptors, + }, + "litellm_metadata": { + TPM_RESERVED_TOKENS_KEY: 9999, + RATE_LIMIT_DESCRIPTORS_KEY: victim_descriptors, + }, + } + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=data, + call_type="completion", + ) + + for channel in ( + data, + data.get("metadata") or {}, + data.get("litellm_metadata") or {}, + ): + leaked = [k for k in _LITELLM_STASH_KEYS if k in channel] + assert not leaked, f"caller-supplied stash survived in {channel!r}: {leaked}" diff --git a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py index 297d18d1ab..e294d1471d 100644 --- a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py +++ b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py @@ -23,6 +23,7 @@ import pytest from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + RATE_LIMIT_DESCRIPTORS_KEY, TPM_RESERVATION_RELEASED_KEY, TPM_RESERVED_MODEL_KEY, TPM_RESERVED_SCOPES_KEY, @@ -606,9 +607,9 @@ async def test_contentless_request_reserves_minimum(rate_limiter): data=data, call_type="", ) - assert ( - data.get(TPM_RESERVED_TOKENS_KEY) == 1 - ), "Contentless request should reserve the floor of 1 token" + assert (data.get("metadata") or {}).get( + TPM_RESERVED_TOKENS_KEY + ) == 1, "Contentless request should reserve the floor of 1 token" counter_after_two = int( await cache.async_get_cache(key=counter_key, local_only=True) or 0 @@ -701,7 +702,7 @@ async def test_reservation_released_on_proxy_rejection(rate_limiter): data=data, call_type="", ) - reserved = data[TPM_RESERVED_TOKENS_KEY] + reserved = (data.get("metadata") or {})[TPM_RESERVED_TOKENS_KEY] assert reserved > 0 counter_key = handler.create_rate_limit_keys( @@ -726,9 +727,9 @@ async def test_reservation_released_on_proxy_rejection(rate_limiter): f"Reservation leaked: counter={counter_after_release} after " f"proxy-level rejection refund (expected 0)." ) - assert data.get(TPM_RESERVATION_RELEASED_KEY) is True, ( - "Released marker must be stamped to prevent async_log_failure_event " - "from double-refunding." + assert (data.get("metadata") or {}).get(TPM_RESERVATION_RELEASED_KEY) is True, ( + "Released marker must be stamped to prevent " + "async_log_failure_event from double-refunding." ) @@ -760,12 +761,7 @@ async def test_reservation_release_idempotent(rate_limiter): shared_metadata = { "user_api_key_hash": api_key, TPM_RESERVED_TOKENS_KEY: 100, - } - - request_data = { - "metadata": shared_metadata, - TPM_RESERVED_TOKENS_KEY: 100, - "_litellm_rate_limit_descriptors": [ + RATE_LIMIT_DESCRIPTORS_KEY: [ { "key": "api_key", "value": api_key, @@ -774,6 +770,10 @@ async def test_reservation_release_idempotent(rate_limiter): ], } + request_data = { + "metadata": shared_metadata, + } + await handler.async_post_call_failure_hook( request_data=request_data, original_exception=Exception("rejected"),