From c2ca9ee2e9da121353a06be0b31e50a6e26baeb8 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Thu, 20 Aug 2026 15:55:19 -0400 Subject: [PATCH] fix(rate-limiting): charge token/dollar usage against the window admission checked Success accounting recomputed a fresh timestamp instead of reusing admission's own, so a call slow enough to cross a period_seconds boundary got admitted against one window's counter but charged into the next window's fresh, empty one. Stash admission's timestamp on the request's model_call_details and reuse it at success time so both stages agree on the same bucket. --- litellm/proxy/hooks/tag_rate_limiter.py | 32 ++++++++++- .../proxy/hooks/test_tag_rate_limiter.py | 55 +++++++++++++++++++ 2 files changed, 86 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/hooks/tag_rate_limiter.py b/litellm/proxy/hooks/tag_rate_limiter.py index 18f89018235..f36b872f5ee 100644 --- a/litellm/proxy/hooks/tag_rate_limiter.py +++ b/litellm/proxy/hooks/tag_rate_limiter.py @@ -539,6 +539,19 @@ _CONCURRENCY_MIN_SAFETY_TTL_SECONDS: Final = 3600 # can't be forged or guessed. _PENDING_CONCURRENCY_KEYS_FIELD: Final[str] = "_tag_rate_limiter_pending_concurrency_keys" +# The admission-time timestamp a hop's token/dollar checks classified their +# bucket against, stashed on the same model_call_details object so success +# accounting recomputes the identical bucket_id (int(now) // period_seconds) +# instead of a fresh one. A completion can take long enough for a fresh +# timestamp at success time to land in the *next* window than the one +# admission actually checked, letting a burst of calls admitted against one +# (still-under-limit) window get charged entirely into the next window's +# fresh, unrelated counter -- silently bypassing the limit right around each +# rollover. Overwritten by each hop's own admission (last-write-wins), which +# is correct: success only ever fires for whichever hop actually served the +# request, so its own most recent admission timestamp is the right one. +_ADMISSION_TIME_FIELD: Final[str] = "_tag_rate_limiter_admission_time" + class _TagRateLimitIndex: """Rebuilds the limits index when `llm_router.model_list` changes, or at @@ -766,6 +779,22 @@ def _queue_pending_concurrency_reservations( pending.extend(reservations) # mutable-ok: see comment above +def _record_admission_time(request_kwargs: Mapping[str, object], now: float) -> None: + """Stash this hop's admission timestamp -- see `_ADMISSION_TIME_FIELD`'s + docstring for why. Silently a no-op without a real logging object + (defensive only; every real request has one): success accounting falls + back to its own fresh timestamp, same as before this fix existed.""" + logging_obj: Final = request_kwargs.get("litellm_logging_obj") + model_call_details: Final = getattr(logging_obj, "model_call_details", None) + if isinstance(model_call_details, dict): + model_call_details[_ADMISSION_TIME_FIELD] = now + + +def _admission_time_or(kwargs: Mapping[str, object], fallback: float) -> float: + recorded: Final = kwargs.get(_ADMISSION_TIME_FIELD) + return recorded if isinstance(recorded, float) else fallback + + @dataclass(frozen=True, slots=True) class _CachePartition: internal_usage_cache: InternalUsageCache @@ -996,6 +1025,7 @@ class _PROXY_TagRateLimiter( # pyright: ignore[reportUnusedClass] # only refer ) now: Final = self._time_provider().timestamp() + _record_admission_time(resolved_request_kwargs, now) classified: Final = tuple( check for configured_limit in configured @@ -1326,7 +1356,7 @@ class _PROXY_TagRateLimiter( # pyright: ignore[reportUnusedClass] # only refer if not tags: return - now: Final = self._time_provider().timestamp() + now: Final = _admission_time_or(kwargs, fallback=self._time_provider().timestamp()) increment_by_unit: Final[Mapping[_LimitUnit, float]] = MappingProxyType( { "tokens": float(standard_logging_object.get("total_tokens") or 0), diff --git a/tests/test_litellm/proxy/hooks/test_tag_rate_limiter.py b/tests/test_litellm/proxy/hooks/test_tag_rate_limiter.py index fcf0e065571..369f0477016 100644 --- a/tests/test_litellm/proxy/hooks/test_tag_rate_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_tag_rate_limiter.py @@ -1039,6 +1039,61 @@ async def test_log_success_event_accounts_against_the_key_hash_admission_checked assert await limiter.internal_usage_cache.async_get_cache(key=unkeyed_bucket, litellm_parent_otel_span=None) is None +@pytest.mark.asyncio +async def test_log_success_event_charges_the_window_admission_checked_not_a_later_one(time_controller): + """ + Admission classifies its bucket as int(now) // period_seconds at filter + time; success accounting used to recompute a fresh now of its own, so a + call slow enough to cross a period_seconds boundary between admission and + completion got admitted against one window's (still-open) counter but + charged into the next window's fresh, empty one -- silently bypassing the + limit for calls straddling each rollover. Success must charge the exact + window admission classified against, not whatever window happens to be + current when the response finishes. + """ + token_limits = { + "token_limits": {"limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 500, "period_seconds": 60}]} + } + router = litellm.Router(model_list=[_deployment("grp", "dep-1", token_limits)]) + limiter = _make_limiter(time_controller) + limiter.update_variables(llm_router=router) + healthy = router.model_list + request_kwargs, kwargs = _call_context(["end_user_id:u1"]) + + admission_time = time_controller.now().timestamp() + result = await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + assert result == healthy + + # The response takes long enough to cross into the next 60s window before + # completing. + time_controller.advance(61) + kwargs["standard_logging_object"] = { + "model_group": "grp", + "model_id": "dep-1", + "total_tokens": 42, + "response_cost": 0.01, + } + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + admitted_window_bucket = _expected_bucket_key("grp", "tokens", "per_minute", "end_user_id", "u1", 60, admission_time) + later_window_bucket = _expected_bucket_key( + "grp", "tokens", "per_minute", "end_user_id", "u1", 60, time_controller.now().timestamp() + ) + assert ( + float( + await limiter.internal_usage_cache.async_get_cache(key=admitted_window_bucket, litellm_parent_otel_span=None) + ) + == 42.0 + ) + assert ( + await limiter.internal_usage_cache.async_get_cache(key=later_window_bucket, litellm_parent_otel_span=None) + is None + ) + + # --------------------------------------------------------------------------- # concurrency limits -- reserve at admission, release on success/failure # ---------------------------------------------------------------------------