tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-03 02:20:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(rate-limit): stop v3 limiter from leaking internal stash to provider body (#27913)

Browse Source 
* 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 <mateo-berri@users.noreply.github.com>

* 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 <cursoragent@cursor.com>
Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>
This commit is contained in:
Mateo Wang
2026-05-14 10:53:04 -07:00
committed by GitHub
co-authored by Cursor Agent Mateo Wang
parent 867470fd68
commit 7a462a4220
3 changed files with 221 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/hooks/parallel_request_limiter_v3.py
+90 -44
View File
@@ -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 []
)
tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py
+118
View File
@@ -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}"
tests/test_litellm/proxy/hooks/test_tpm_concurrent.py
+13 -13
View File
@@ -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"),