From 8b631a31e0137350155b9844e6fdc191ee06e1e6 Mon Sep 17 00:00:00 2001 From: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 00:54:47 -0700 Subject: [PATCH 1/4] perf(rate-limiter): charge v3 token counters only for scopes admission rate limited The v3 success path wrote a {scope}:tokens counter for every key, user, team, team member, end user, organization and per-model scope on the request, including scopes with no rate limit that no request ever reads. The pre-call hook now records every (scope, value) it built a rate limit descriptor for, plus the model group it checked, on the request stash. The success path keeps only those targets, keeps the per-model targets when a different model group served the request (fallback), and keeps every target when no admission was recorded for the call. Resolves LIT-9875 --- .../hooks/parallel_request_limiter_v3.py | 52 ++- .../spend/test_rate_limit_token_counters.py | 422 ++++++++++++++++++ .../hooks/test_parallel_request_limiter_v3.py | 183 +++++++- 3 files changed, 654 insertions(+), 3 deletions(-) create mode 100644 tests/integration/spend/test_rate_limit_token_counters.py diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index d858eff3348..972e944f729 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -515,6 +515,9 @@ _AUDIO_BYTES_PER_TOKEN: Final = 1600 PROJECT_ITPM_DESCRIPTOR_KEY: Final = "model_per_project_itpm" PROJECT_OTPM_DESCRIPTOR_KEY: Final = "model_per_project_otpm" PTU_TEAM_DESCRIPTOR_KEY: Final = "model_per_team_ptu" +_MODEL_SCOPED_TPM_TARGETS: Final = frozenset( + {"model_per_key", "model_per_team", "model_per_organization", "model_per_project"} +) # How long an acquired slot counts toward the in-flight total before it is # considered leaked (worker crashed without any release callback firing) and # pruned. Also the longest request duration the gauge can track: a request @@ -665,6 +668,14 @@ class CallTypeRateLimiter(Protocol): ) -> Exception | str | dict[str, object] | None: ... +@dataclass(frozen=True, slots=True) +class AdmittedRateLimitScopes: + """Every ``(scope, value)`` the pre-call hook built a rate limit descriptor for, and the model group it checked.""" + + scopes: frozenset[tuple[str, str]] + model_group: str | None + + @dataclass(slots=True) class RequestRateLimiterStash: """ @@ -712,6 +723,7 @@ class RequestRateLimiterStash: batch_tpd_refund_ops: tuple[ReservationAwareIncrementOperation, ...] = () reservation_released: bool = False tpm_limited_tags: frozenset[str] = field(default_factory=frozenset) + admitted_scopes: AdmittedRateLimitScopes | None = None @dataclass(frozen=True, slots=True) @@ -4197,6 +4209,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): for d in descriptors if d["key"] == "tag" and d["rate_limit"] is not None and d["rate_limit"].get("tokens_per_unit") is not None ) + stash.admitted_scopes = self._admitted_scopes(descriptors, requested_model) # Only check rate limits if we have descriptors with actual limits if descriptors: @@ -5017,6 +5030,36 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): targets.extend(("tag", tag) for tag in sorted(tpm_limited_tags)) return targets + def _admitted_scopes( + self, descriptors: Sequence[RateLimitDescriptor], requested_model: str | None + ) -> AdmittedRateLimitScopes: + model: Final = self._rate_limited_model(requested_model) + return AdmittedRateLimitScopes( + scopes=frozenset((d["key"], d["value"]) for d in descriptors), + model_group=model.group if model is not None else None, + ) + + @staticmethod + def _rate_limited_targets( + targets: Sequence[tuple[str, str]], + admitted: AdmittedRateLimitScopes | None, + model_group: str | None, + ) -> Sequence[tuple[str, str]]: + """ + Keep the TPM targets admission built a rate limit descriptor for, plus + the per-model targets when tokens are settled against a model group + admission did not check. Every target stays when no admission was + recorded for the call. + """ + if admitted is None: + return targets + served_admitted_group: Final = model_group == admitted.model_group + return tuple( + target + for target in targets + if target in admitted.scopes or (not served_admitted_group and target[0] in _MODEL_SCOPED_TPM_TARGETS) + ) + def _build_reservation_aware_tpm_ops( self, targets: Sequence[tuple[str, str]], @@ -5197,10 +5240,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): for scope in sorted(reserved_scopes) if scope[0] in ("agent", "agent_session") and scope not in targets ) + limited_targets: Final = self._rate_limited_targets( + targets=targets, + admitted=stash.admitted_scopes if stash is not None else None, + model_group=reconcile_model.group if reconcile_model is not None else None, + ) charged_targets: Final = ( - [target for target in targets if target[0] != "model_per_team"] + [target for target in limited_targets if target[0] != "model_per_team"] if self._key_owns_model_tpm_limit_from_request_metadata(request_metadata, reconcile_model) - else targets + else limited_targets ) if reserved_tokens > 0 and total_tokens < reserved_tokens: verbose_proxy_logger.debug( diff --git a/tests/integration/spend/test_rate_limit_token_counters.py b/tests/integration/spend/test_rate_limit_token_counters.py new file mode 100644 index 00000000000..d51dbe6b428 --- /dev/null +++ b/tests/integration/spend/test_rate_limit_token_counters.py @@ -0,0 +1,422 @@ +import asyncio +import json +import os +import uuid +from collections.abc import Callable, Iterator, Mapping, Set +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from dataclasses import dataclass, replace +from hashlib import sha256 +from typing import Final + +import anthropic +import httpx +import openai +import pytest +from integration._support.client import Gateway, Scenario, eventually, held +from integration._support.openai_wire import chat_reply, responses_reply +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue +from redis import Redis + +USAGE: Final = {"prompt_tokens": 12, "completion_tokens": 8, "total_tokens": 20} +WIRE_TOTAL_TOKENS: Final = 8 +Fields = Callable[[str], Mapping[str, JsonValue]] + + +def completion(request: Request) -> Reply: + assert request.target == "/v1/chat/completions" + body: Final = json.loads(request.body) + return Reply( + body=json.dumps( + { + "id": "chatcmpl-" + uuid.uuid4().hex, + "object": "chat.completion", + "created": 1, + "model": body["model"], + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": USAGE, + } + ).encode() + ) + + +def outage(request: Request) -> Reply: + assert request.target == "/v1/chat/completions" + return Reply(status=500, body=json.dumps({"error": {"message": "synthetic outage"}}).encode()) + + +def openai_endpoints(request: Request) -> Reply: + body: Final = json.loads(request.body) + stream: Final = body.get("stream") is True + if request.target == "/v1/chat/completions": + return chat_reply("chatcmpl-" + uuid.uuid4().hex, body["model"], "ok", stream=stream) + assert request.target == "/v1/responses", request.target + return responses_reply("resp_" + uuid.uuid4().hex, body["model"], "ok", stream=stream) + + +def counter(scope: str, value: str) -> str: + return f"{{{scope}:{value}}}:tokens" + + +def read_counters(cache: Redis, keys: Mapping[str, str]) -> Mapping[str, float | None]: + values: Final = cache.mget(list(keys.values())) + return {name: None if raw is None else float(raw) for name, raw in zip(keys, values, strict=True)} + + +def no_fields(_model: str) -> Mapping[str, JsonValue]: + return {} + + +@dataclass(frozen=True, slots=True) +class Caller: + key: str + user: str + team: str + scopes: Mapping[str, str] + body: Mapping[str, JsonValue] + + +@dataclass(frozen=True, slots=True) +class Cell: + name: str + charged: str + also_charged: frozenset[str] = frozenset() + key: Fields = no_fields + user: Fields = no_fields + team: Fields = no_fields + organization: Fields | None = None + project: Fields | None = None + customer: Fields | None = None + tag: Fields | None = None + + +def caller(scenario: Scenario, cell: Cell, model: str) -> Caller: + gateway: Final = scenario.gateway + organization: Final = None if cell.organization is None else scenario.organization(**cell.organization(model)) + team: Final = scenario.team( + **cell.team(model), **({} if organization is None else {"organization_id": organization}) + ) + user: Final = scenario.user(user_role="internal_user", **cell.user(model)) + gateway.post("/team/member_add", {"team_id": team, "member": {"role": "user", "user_id": user}}) + project: Final = None if cell.project is None else scenario.project(team, **cell.project(model)) + key: Final = scenario.key( + user_id=user, team_id=team, **({} if project is None else {"project_id": project}), **cell.key(model) + ) + hashed: Final = sha256(key.encode()).hexdigest() + scopes: Final = { + "api_key": counter("api_key", hashed), + "user": counter("user", user), + "team": counter("team", team), + "team_member": counter("team_member", f"{team}:{user}"), + "model_per_key": counter("model_per_key", f"{hashed}:{model}"), + "model_per_team": counter("model_per_team", f"{team}:{model}"), + } + body: Final[dict[str, JsonValue]] = {} + if organization is not None: + scopes["organization"] = counter("organization", organization) + scopes["model_per_organization"] = counter("model_per_organization", f"{organization}:{model}") + if project is not None: + scopes["model_per_project"] = counter("model_per_project", f"{project}:{model}") + if cell.customer is not None: + customer: Final = f"integration-customer-{uuid.uuid4().hex}" + gateway.post("/customer/new", {"user_id": customer, **cell.customer(model)}) + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [customer]}) + scopes["end_user"] = counter("end_user", customer) + body["user"] = customer + if cell.tag is not None: + tag: Final = f"integration-tag-{uuid.uuid4().hex}" + gateway.post("/tag/new", {"name": tag, **cell.tag(model)}) + scenario.cleanups.callback(gateway.request, "POST", "/tag/delete", {"name": tag}) + scopes["tag"] = counter("tag", tag) + body["metadata"] = {"tags": [tag]} + return Caller(key, user, team, scopes, body) + + +def key_rpm_only(_model: str) -> Mapping[str, JsonValue]: + return {"rpm_limit": 1000} + + +CELLS: Final = ( + Cell("organization_rpm_only", charged="organization", organization=lambda _model: {"rpm_limit": 1000}), + Cell( + "organization_model_rpm_only", + charged="model_per_organization", + organization=lambda model: {"model_rpm_limit": {model: 1000}}, + ), + Cell("key_rpm_only", charged="api_key", key=key_rpm_only), + Cell("user_rpm_only", charged="user", user=lambda _model: {"rpm_limit": 1000}), + Cell("team_rpm_only", charged="team", team=lambda _model: {"rpm_limit": 1000}), + Cell("team_member_rpm_only", charged="team_member", team=lambda _model: {"team_member_rpm_limit": 1000}), + Cell("key_model_rpm_only", charged="model_per_key", key=lambda model: {"model_rpm_limit": {model: 1000}}), + Cell( + "team_model_rpm_only", + charged="model_per_team", + also_charged=frozenset({"model_per_key"}), + team=lambda model: {"model_rpm_limit": {model: 1000}}, + ), + Cell( + "project_model_rpm_only", charged="model_per_project", project=lambda model: {"model_rpm_limit": {model: 1000}} + ), + Cell("end_user_rpm_only", charged="end_user", customer=lambda _model: {"rpm_limit": 1000}), + Cell("tag_tpm", charged="tag", tag=lambda _model: {"tpm_limit": 100000}), + Cell("key_tpm_reserved", charged="api_key", key=lambda _model: {"tpm_limit": 100000}), + Cell( + "key_dynamic_tpm", + charged="api_key", + key=lambda _model: {"tpm_limit": 100000, "tpm_limit_type": "dynamic"}, + ), +) + + +@contextmanager +def redis_counters() -> Iterator[Redis]: + with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache: + yield cache + + +def assert_only_charged(cache: Redis, scopes: Mapping[str, str], charged: Set[str], tokens: float) -> None: + counters: Final = eventually( + lambda: read_counters(cache, scopes), + lambda values: all(values[name] == tokens for name in charged), + seconds=30, + ) + assert {name: value for name, value in counters.items() if name not in charged} == { + name: None for name in scopes if name not in charged + } + + +def chat(gateway: Gateway, identity: Caller, model: str, prompt: str, **extra: JsonValue) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": prompt}], **identity.body, **extra}, + key=identity.key, + ) + + +@pytest.mark.parametrize("cell", CELLS, ids=[cell.name for cell in CELLS]) +def test_success_charges_token_counters_only_for_rate_limited_scopes(gateway: Gateway, cell: Cell) -> None: + prompt: Final = f"token counters {cell.name} {uuid.uuid4().hex}" + with wire_server(completion) as upstream, gateway.scenario() as scenario, redis_counters() as cache: + model: Final = scenario.model(api_base=upstream.url + "/v1") + identity: Final = caller(scenario, cell, model) + response: Final = chat(gateway, identity, model, prompt, max_tokens=500) + assert response.status_code == 200, response.text + assert response.json()["usage"]["total_tokens"] == USAGE["total_tokens"] + assert [json.loads(seen.body)["messages"] for seen in upstream.drain()] == [ + [{"role": "user", "content": prompt}] + ] + + assert_only_charged(cache, identity.scopes, {cell.charged, *cell.also_charged}, USAGE["total_tokens"]) + + +def gateway_url(gateway: Gateway) -> str: + return str(gateway.client.base_url).rstrip("/") + + +def openai_chat_sync(gateway: Gateway, key: str, model: str, prompt: str) -> int | None: + with ( + httpx.Client(timeout=30, trust_env=False) as transport, + openai.OpenAI( + api_key=key, base_url=gateway_url(gateway) + "/v1", max_retries=0, http_client=transport + ) as client, + ): + response: Final = client.chat.completions.create(model=model, messages=[{"role": "user", "content": prompt}]) + return None if response.usage is None else response.usage.total_tokens + + +async def openai_chat_stream_async(gateway: Gateway, key: str, model: str, prompt: str) -> int | None: + async with ( + httpx.AsyncClient(timeout=30, trust_env=False) as transport, + openai.AsyncOpenAI( + api_key=key, base_url=gateway_url(gateway) + "/v1", max_retries=0, http_client=transport + ) as client, + ): + stream: Final = await client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": prompt}], + stream=True, + stream_options={"include_usage": True}, + ) + usage: Final = [chunk.usage.total_tokens async for chunk in stream if chunk.usage is not None] + return usage[-1] if usage else None + + +def anthropic_messages_sync(gateway: Gateway, key: str, model: str, prompt: str) -> int | None: + with ( + httpx.Client(timeout=30, trust_env=False) as transport, + anthropic.Anthropic(api_key=key, base_url=gateway_url(gateway), max_retries=0, http_client=transport) as client, + ): + message: Final = client.messages.create( + model=model, max_tokens=50, messages=[{"role": "user", "content": prompt}] + ) + return message.usage.input_tokens + message.usage.output_tokens + + +async def anthropic_messages_stream_async(gateway: Gateway, key: str, model: str, prompt: str) -> int | None: + async with ( + httpx.AsyncClient(timeout=30, trust_env=False) as transport, + anthropic.AsyncAnthropic( + api_key=key, base_url=gateway_url(gateway), max_retries=0, http_client=transport + ) as client, + client.messages.stream(model=model, max_tokens=50, messages=[{"role": "user", "content": prompt}]) as stream, + ): + message: Final = await stream.get_final_message() + return message.usage.input_tokens + message.usage.output_tokens + + +def openai_responses_sync(gateway: Gateway, key: str, model: str, prompt: str) -> int | None: + with ( + httpx.Client(timeout=30, trust_env=False) as transport, + openai.OpenAI( + api_key=key, base_url=gateway_url(gateway) + "/v1", max_retries=0, http_client=transport + ) as client, + ): + response: Final = client.responses.create(model=model, input=prompt) + return None if response.usage is None else response.usage.total_tokens + + +async def openai_responses_stream_async(gateway: Gateway, key: str, model: str, prompt: str) -> int | None: + async with ( + httpx.AsyncClient(timeout=30, trust_env=False) as transport, + openai.AsyncOpenAI( + api_key=key, base_url=gateway_url(gateway) + "/v1", max_retries=0, http_client=transport + ) as client, + ): + stream: Final = await client.responses.create(model=model, input=prompt, stream=True) + usage: Final = [ + event.response.usage.total_tokens + async for event in stream + if event.type == "response.completed" and event.response.usage is not None + ] + return usage[-1] if usage else None + + +ENDPOINTS: Final[Mapping[str, Callable[[Gateway, str, str, str], int | None]]] = { + "chat_sync": openai_chat_sync, + "chat_stream_async": lambda *args: asyncio.run(openai_chat_stream_async(*args)), + "messages_sync": anthropic_messages_sync, + "messages_stream_async": lambda *args: asyncio.run(anthropic_messages_stream_async(*args)), + "responses_sync": openai_responses_sync, + "responses_stream_async": lambda *args: asyncio.run(openai_responses_stream_async(*args)), +} + + +@pytest.mark.parametrize("endpoint", ENDPOINTS) +def test_every_endpoint_charges_only_the_rate_limited_scope(gateway: Gateway, endpoint: str) -> None: + prompt: Final = f"token counters {endpoint} {uuid.uuid4().hex}" + with wire_server(openai_endpoints) as upstream, gateway.scenario() as scenario, redis_counters() as cache: + model: Final = scenario.model(api_base=upstream.url + "/v1") + identity: Final = caller(scenario, Cell(endpoint, charged="api_key", key=key_rpm_only), model) + + assert ENDPOINTS[endpoint](gateway, identity.key, model, prompt) == WIRE_TOTAL_TOKENS + assert len(upstream.drain()) == 1 + + assert_only_charged(cache, identity.scopes, {"api_key"}, WIRE_TOTAL_TOKENS) + + +def test_fallback_hop_charges_the_rate_limited_model_it_served(gateway: Gateway) -> None: + prompt: Final = f"token counters fallback {uuid.uuid4().hex}" + with ( + wire_server(outage) as primary_upstream, + wire_server(completion) as fallback_upstream, + gateway.scenario() as scenario, + redis_counters() as cache, + ): + primary: Final = scenario.model(api_base=primary_upstream.url + "/v1") + fallback: Final = scenario.model(api_base=fallback_upstream.url + "/v1") + cell: Final = Cell( + "fallback", + charged="model_per_key", + key=lambda _model: {"model_tpm_limit": {fallback: 100000}}, + ) + identity: Final = caller(scenario, cell, fallback) + response: Final = chat(gateway, identity, primary, prompt, fallbacks=[fallback]) + assert response.status_code == 200, response.text + assert response.headers["x-litellm-model-group"] == fallback + assert len(primary_upstream.drain()) == 1 + assert len(fallback_upstream.drain()) == 1 + + hashed: Final = sha256(identity.key.encode()).hexdigest() + scopes: Final = { + name: identity.scopes[name] for name in ("api_key", "user", "team", "team_member", "model_per_key") + } + assert_only_charged( + cache, + {**scopes, "primary_model_per_key": counter("model_per_key", f"{hashed}:{primary}")}, + {"model_per_key"}, + USAGE["total_tokens"], + ) + + +def same_key(_scenario: Scenario, identity: Caller) -> Caller: + return identity + + +def another_key_of_the_same_user(scenario: Scenario, identity: Caller) -> Caller: + return replace(identity, key=scenario.key(user_id=identity.user, team_id=identity.team)) + + +@pytest.mark.parametrize( + "cell, next_caller", + [ + (Cell("key_tpm", charged="api_key", key=lambda _model: {"tpm_limit": 21}), same_key), + ( + Cell("user_tpm_across_keys", charged="user", user=lambda _model: {"tpm_limit": 21}), + another_key_of_the_same_user, + ), + ], + ids=["key_tpm", "user_tpm_across_keys"], +) +def test_rate_limited_scope_rejects_once_charged_tokens_reach_its_limit( + gateway: Gateway, cell: Cell, next_caller: Callable[[Scenario, Caller], Caller] +) -> None: + with wire_server(completion) as upstream, gateway.scenario() as scenario, redis_counters() as cache: + model: Final = scenario.model(api_base=upstream.url + "/v1") + identity: Final = caller(scenario, cell, model) + first: Final = chat(gateway, identity, model, "hi", max_tokens=1) + assert first.status_code == 200, first.text + eventually( + lambda: read_counters(cache, identity.scopes)[cell.charged], + lambda tokens: tokens == USAGE["total_tokens"], + seconds=30, + ) + + rejected: Final = chat(gateway, next_caller(scenario, identity), model, "hi again", max_tokens=1) + assert rejected.status_code == 429, rejected.text + assert len(upstream.drain()) == 1 + assert_only_charged(cache, identity.scopes, {cell.charged}, USAGE["total_tokens"]) + + +def test_concurrent_burst_charges_every_request_to_the_rate_limited_scope_only(gateway: Gateway) -> None: + requests: Final = 20 + with wire_server(completion) as upstream, gateway.scenario() as scenario, redis_counters() as cache: + model: Final = scenario.model(api_base=upstream.url + "/v1") + identity: Final = caller(scenario, Cell("burst", charged="api_key", key=key_rpm_only), model) + with ThreadPoolExecutor(max_workers=requests) as pool: + statuses: Final = list( + pool.map( + lambda index: chat(gateway, identity, model, f"token counters burst {index} {uuid.uuid4().hex}"), + range(requests), + ) + ) + assert [response.status_code for response in statuses] == [200] * requests + assert len(upstream.drain()) == requests + + assert_only_charged(cache, identity.scopes, {"api_key"}, requests * USAGE["total_tokens"]) + + +def test_upstream_outage_charges_no_token_counter(gateway: Gateway) -> None: + with wire_server(outage) as upstream, gateway.scenario() as scenario, redis_counters() as cache: + model: Final = scenario.model(api_base=upstream.url + "/v1") + identity: Final = caller(scenario, Cell("outage", charged="api_key", key=key_rpm_only), model) + response: Final = chat(gateway, identity, model, f"token counters outage {uuid.uuid4().hex}") + assert response.status_code == 500, response.text + assert len(upstream.drain()) == 1 + + held( + lambda: read_counters(cache, identity.scopes), + lambda values: set(values.values()) == {None}, + holding=3, + ) diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 7af747fc26d..82b5bab8362 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -6,7 +6,7 @@ import asyncio import logging import sys import time -from collections.abc import Iterator, Sequence +from collections.abc import Awaitable, Callable, Iterator, Sequence from contextlib import contextmanager from datetime import datetime, timedelta, timezone from typing import Any, Dict, Final, List, Optional @@ -8686,3 +8686,184 @@ async def test_a2a_url_target_owns_invocation_fee_and_request_limit( await _rpm_request(limiter, cache, auth, "a2a/cheap") assert denied.value.status_code == 429 assert "expensive" in str(denied.value.detail) + + +_SCOPE_TOKEN_KEYS: Final = frozenset( + { + "{api_key:key-hash}:tokens", + "{user:u}:tokens", + "{team:t}:tokens", + "{team_member:t:u}:tokens", + "{model_per_key:key-hash:gpt-4o-mini}:tokens", + "{model_per_team:t:gpt-4o-mini}:tokens", + } +) + + +def _scoped_success_kwargs( + model_group: str, key_metadata: Dict[str, object] | None = None, call_id: str | None = None +) -> Dict[str, object]: + return { + **({} if call_id is None else {"litellm_call_id": call_id}), + "standard_logging_object": { + "metadata": {"user_api_key_hash": "key-hash", "user_api_key_user_id": "u", "user_api_key_team_id": "t"} + }, + "litellm_params": { + "metadata": { + "model_group": model_group, + "user_api_key_metadata": key_metadata or {}, + "user_api_key_team_metadata": {}, + } + }, + "model": model_group, + } + + +async def _admit_scoped_request( + handler: _PROXY_MaxParallelRequestsHandler, auth: UserAPIKeyAuth, model: str, call_id: str | None = None +) -> None: + await handler.async_pre_call_hook( + user_api_key_dict=auth, + cache=handler.internal_usage_cache.dual_cache, + data={ + "model": model, + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 20, + **({} if call_id is None else {"litellm_call_id": call_id}), + }, + call_type="acompletion", + ) + + +def _charged_token_keys(handler: _PROXY_MaxParallelRequestsHandler, kwargs: Dict[str, object]) -> frozenset[str]: + operations: Final = handler._build_success_event_pipeline_operations( + kwargs=kwargs, + response_obj=ModelResponse(usage=Usage(prompt_tokens=2, completion_tokens=1, total_tokens=3)), + rate_limit_type="total", + ) + return frozenset(operation["key"] for operation in operations) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "limits, charged", + [ + ({}, frozenset()), + ({"rpm_limit": 100}, frozenset({"{api_key:key-hash}:tokens"})), + ({"max_parallel_requests": 10}, frozenset({"{api_key:key-hash}:tokens"})), + ({"tpm_limit": 1000}, frozenset({"{api_key:key-hash}:tokens"})), + ({"tpm_limit": 1000, "metadata": {"tpm_limit_type": "dynamic"}}, frozenset({"{api_key:key-hash}:tokens"})), + ({"user_rpm_limit": 100}, frozenset({"{user:u}:tokens"})), + ({"team_tpm_limit": 1000}, frozenset({"{team:t}:tokens"})), + ({"team_member_rpm_limit": 100}, frozenset({"{team_member:t:u}:tokens"})), + ( + {"metadata": {"model_rpm_limit": {"gpt-4o-mini": 100}}}, + frozenset({"{model_per_key:key-hash:gpt-4o-mini}:tokens"}), + ), + ], + ids=[ + "no_limits", + "key_rpm", + "key_max_parallel", + "key_tpm_reserved", + "key_dynamic_tpm", + "user_rpm", + "team_tpm_reserved", + "team_member_rpm", + "key_model_rpm", + ], +) +async def test_success_charges_tokens_only_to_scopes_admission_rate_limited( + limits: Dict[str, object], charged: frozenset[str] +) -> None: + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + await _admit_scoped_request( + handler, UserAPIKeyAuth(api_key="key-hash", user_id="u", team_id="t", **limits), "gpt-4o-mini" + ) + + assert _charged_token_keys(handler, _scoped_success_kwargs("gpt-4o-mini")) == charged + + +@pytest.mark.asyncio +async def test_fallback_success_charges_the_model_scopes_of_the_group_that_served_it() -> None: + key_metadata: Final = {"model_tpm_limit": {"fallback-model": 1000}} + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + await _admit_scoped_request( + handler, UserAPIKeyAuth(api_key="key-hash", user_id="u", team_id="t", metadata=key_metadata), "primary-model" + ) + + assert _charged_token_keys(handler, _scoped_success_kwargs("fallback-model", key_metadata)) == { + "{model_per_key:key-hash:fallback-model}:tokens" + } + + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "success_end_user, charged", + [ + (None, frozenset({"{end_user:admitted-customer}:tokens"})), + ("admitted-customer", frozenset({"{end_user:admitted-customer}:tokens"})), + ("other-customer", frozenset()), + ], + ids=["metadata_end_user", "same_body_user", "different_body_user"], +) +async def test_success_charges_the_end_user_admission_checked_and_no_other( + success_end_user: str | None, charged: frozenset[str] +) -> None: + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + await _admit_scoped_request( + handler, + UserAPIKeyAuth(api_key="key-hash", end_user_id="admitted-customer", end_user_rpm_limit=100), + "gpt-4o-mini", + ) + kwargs: Final = _scoped_success_kwargs("gpt-4o-mini") + kwargs["standard_logging_object"] = { + "metadata": {"user_api_key_hash": "key-hash", "user_api_key_end_user_id": "admitted-customer"} + } + if success_end_user is not None: + kwargs["user"] = success_end_user + + assert _charged_token_keys(handler, kwargs) == charged + + +@pytest.mark.asyncio +async def test_success_on_an_aliased_model_charges_the_model_limit_of_the_resolved_group() -> None: + key_metadata: Final = {"model_rpm_limit": {"target": 100}} + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()), model_group_resolver=_resolve_alias_to_target + ) + await _admit_scoped_request( + handler, UserAPIKeyAuth(api_key="key-hash", user_id="u", team_id="t", metadata=key_metadata), "alias" + ) + + assert _charged_token_keys(handler, _scoped_success_kwargs("target", key_metadata)) == { + "{model_per_key:key-hash:target}:tokens" + } + +async def _no_stash(handler: _PROXY_MaxParallelRequestsHandler) -> str | None: + _request_stash.set(None) + return None + + +async def _nested_call(handler: _PROXY_MaxParallelRequestsHandler) -> str | None: + await _admit_scoped_request( + handler, UserAPIKeyAuth(api_key="key-hash", user_id="u", team_id="t"), "gpt-4o-mini", "proxy-request" + ) + return "nested-call" + + +async def _stash_without_admission(handler: _PROXY_MaxParallelRequestsHandler) -> str | None: + _request_stash.set(RequestRateLimiterStash()) + return None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("arrange", [_no_stash, _nested_call, _stash_without_admission]) +async def test_success_charges_every_scope_when_the_call_was_not_admitted( + arrange: Callable[[_PROXY_MaxParallelRequestsHandler], Awaitable[str | None]], +) -> None: + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + call_id: Final = await arrange(handler) + + assert _charged_token_keys(handler, _scoped_success_kwargs("gpt-4o-mini", call_id=call_id)) == _SCOPE_TOKEN_KEYS From 8f9ba9b987fb305364924980a33716e7aa3e6fd4 Mon Sep 17 00:00:00 2001 From: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 18:20:36 -0700 Subject: [PATCH 2/4] perf(rate-limiter): write v3 token counters only into open rate limit windows The previous revision charged tokens only to the scopes admission checked, which stopped master key, JWT, header-mapped and numeric `user` traffic from counting toward limits that other requests of the same scope enforce. Every unreserved success charge now carries its scope's window key and goes through OPEN_WINDOW_TOKEN_INCREMENT_SCRIPT, which increments the counter only while that window exists. A scope with no limit never opens a window, so it gets no counter; a counter written outside a window was deleted by the next admission anyway, so every limit sees the same token totals as before. A hash-tag group that mixes plain and gated charges goes out as one script call, so the group lands or falls back as a whole. Reserved scopes keep their reconciliation deltas, and the no-Redis path still increments in memory as before. Resolves LIT-9875 --- .../hooks/parallel_request_limiter_v3.py | 172 +++++++++++------- .../spend/test_rate_limit_token_counters.py | 115 +++++++++++- ...st_redis_ttl_preserving_token_increment.py | 41 +++++ .../test_request_redis_batch_post_call.py | 63 +++++++ .../hooks/test_parallel_request_limiter_v3.py | 146 +++------------ 5 files changed, 346 insertions(+), 191 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 972e944f729..4c5da039511 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -468,6 +468,25 @@ end return results """ +OPEN_WINDOW_TOKEN_INCREMENT_SCRIPT: Final = """ +-- KEYS: (window_key, counter_key) pairs; ARGV: (increment, ttl_seconds) pairs at the same indexes. +-- A counter is incremented only while its scope's rate limit window is open. A pair whose window key is +-- its own counter key has no window to wait for and is always incremented. +local applied = 0 +for i = 1, #KEYS, 2 do + local counter_key = KEYS[i + 1] + if KEYS[i] == counter_key or redis.call('EXISTS', KEYS[i]) == 1 then + local ttl_seconds = tonumber(ARGV[i + 1]) + redis.call('INCRBYFLOAT', counter_key, ARGV[i]) + if ttl_seconds > 0 and redis.call('TTL', counter_key) == -1 then + redis.call('EXPIRE', counter_key, ttl_seconds) + end + applied = applied + 1 + end +end +return applied +""" + # Redis cluster slot count REDIS_CLUSTER_SLOTS: Final = 16384 REDIS_NODE_HASHTAG_NAME: Final = "all_keys" @@ -515,9 +534,6 @@ _AUDIO_BYTES_PER_TOKEN: Final = 1600 PROJECT_ITPM_DESCRIPTOR_KEY: Final = "model_per_project_itpm" PROJECT_OTPM_DESCRIPTOR_KEY: Final = "model_per_project_otpm" PTU_TEAM_DESCRIPTOR_KEY: Final = "model_per_team_ptu" -_MODEL_SCOPED_TPM_TARGETS: Final = frozenset( - {"model_per_key", "model_per_team", "model_per_organization", "model_per_project"} -) # How long an acquired slot counts toward the in-flight total before it is # considered leaked (worker crashed without any release callback firing) and # pruned. Also the longest request duration the gauge can track: a request @@ -620,6 +636,12 @@ class ReservationAwareIncrementOperation(RedisPipelineIncrementOperation): reservation_backend: NotRequired[Literal["redis", "local"]] +class OpenWindowIncrementOperation(RedisPipelineIncrementOperation): + """A token increment Redis applies only while ``open_window_key`` holds its scope's rate limit window.""" + + open_window_key: str + + class RateLimitResponseWithDescriptors(TypedDict): descriptors: list[RateLimitDescriptor] response: RateLimitResponse @@ -668,14 +690,6 @@ class CallTypeRateLimiter(Protocol): ) -> Exception | str | dict[str, object] | None: ... -@dataclass(frozen=True, slots=True) -class AdmittedRateLimitScopes: - """Every ``(scope, value)`` the pre-call hook built a rate limit descriptor for, and the model group it checked.""" - - scopes: frozenset[tuple[str, str]] - model_group: str | None - - @dataclass(slots=True) class RequestRateLimiterStash: """ @@ -723,7 +737,6 @@ class RequestRateLimiterStash: batch_tpd_refund_ops: tuple[ReservationAwareIncrementOperation, ...] = () reservation_released: bool = False tpm_limited_tags: frozenset[str] = field(default_factory=frozenset) - admitted_scopes: AdmittedRateLimitScopes | None = None @dataclass(frozen=True, slots=True) @@ -823,6 +836,11 @@ def _as_window_guarded(operation: RedisPipelineIncrementOperation) -> Reservatio ) +def _open_window_key(operation: RedisPipelineIncrementOperation) -> str | None: + open_window_key: Final = operation.get("open_window_key") + return open_window_key if isinstance(open_window_key, str) else None + + def _call_id_from_callback_kwargs(kwargs: object) -> str | None: if not isinstance(kwargs, dict): return None @@ -843,6 +861,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): batch_rate_limiter_script: _AsyncLuaScript | None batch_counter_read_script: _AsyncLuaScript | None token_increment_script: _AsyncLuaScript | None + open_window_token_increment_script: _AsyncLuaScript | None check_and_increment_by_n_script: _AsyncLuaScript | None window_guarded_token_increment_script: _AsyncLuaScript | None parallel_acquire_script: _AsyncLuaScript | None @@ -878,6 +897,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self.token_increment_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( TOKEN_INCREMENT_SCRIPT ) + self.open_window_token_increment_script = ( + self.internal_usage_cache.dual_cache.redis_cache.async_register_script( + OPEN_WINDOW_TOKEN_INCREMENT_SCRIPT + ) + ) self.check_and_increment_by_n_script = ( self.internal_usage_cache.dual_cache.redis_cache.async_register_script(CHECK_AND_INCREMENT_BY_N_SCRIPT) ) @@ -899,6 +923,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self.batch_rate_limiter_script = None self.batch_counter_read_script = None self.token_increment_script = None + self.open_window_token_increment_script = None self.check_and_increment_by_n_script = None self.window_guarded_token_increment_script = None self.parallel_acquire_script = None @@ -4209,7 +4234,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): for d in descriptors if d["key"] == "tag" and d["rate_limit"] is not None and d["rate_limit"].get("tokens_per_unit") is not None ) - stash.admitted_scopes = self._admitted_scopes(descriptors, requested_model) # Only check rate limits if we have descriptors with actual limits if descriptors: @@ -4545,8 +4569,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ Execute token increment script grouped by hash tag for cluster compatibility. """ - if self.token_increment_script is None: + token_increment_script: Final = self.token_increment_script + if token_increment_script is None: return + open_window_script: Final = self.open_window_token_increment_script # Group operations by hash tag for Redis cluster compatibility operation_keys: Final = [op["key"] for op in pipeline_operations] @@ -4555,6 +4581,26 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): for _hash_tag, group_keys in key_groups.items(): # Get operations for this hash tag group group_operations = [op for op in pipeline_operations if op["key"] in group_keys] + window_keys = [_open_window_key(op) for op in group_operations] + if open_window_script is not None and any(window_key is not None for window_key in window_keys): + # One script for the whole group, so its plain and window-gated increments land together + await self._run_token_increment_script( + OPEN_WINDOW_TOKEN_INCREMENT_SCRIPT, + open_window_script, + tuple( + itertools.chain.from_iterable( + (window_key or op["key"], op["key"]) + for op, window_key in zip(group_operations, window_keys, strict=True) + ) + ), + tuple( + itertools.chain.from_iterable( + (op["increment_value"], op["ttl"] or 0) for op in group_operations + ) + ), + group_operations, + ) + continue keys = [] args = [] @@ -4572,25 +4618,34 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): keys.append(op["key"]) args.extend([op["increment_value"], ttl_value]) - if self._defer_token_increment_script(keys, args, group_operations): - continue - await self.token_increment_script( - keys=keys, - args=args, + await self._run_token_increment_script( + TOKEN_INCREMENT_SCRIPT, token_increment_script, keys, args, group_operations ) + async def _run_token_increment_script( + self, + source: str, + script: _AsyncLuaScript, + keys: Sequence[str], + args: Sequence[float], + group_operations: Sequence[RedisPipelineIncrementOperation], + ) -> None: + if not self._defer_token_increment_script(source, script, keys, args, group_operations): + await script(keys=keys, args=args) + def _defer_token_increment_script( self, - keys: list[str], - args: list[int], - group_operations: list["RedisPipelineIncrementOperation"], + source: str, + script: _AsyncLuaScript, + keys: Sequence[str], + args: Sequence[float], + group_operations: Sequence[RedisPipelineIncrementOperation], ) -> bool: """Declared into the request's post-call pipeline instead of its own EVALSHA round trip; a failed script falls back to the plain increment pipeline for its own group, as the direct path does.""" redis_cache: Final = self.internal_usage_cache.dual_cache.redis_cache - script: Final = self.token_increment_script batch: Final = None if redis_cache is None else active_post_call_redis_batch(redis_cache) - if batch is None or script is None: + if batch is None: return False async def fall_back(future: asyncio.Future[object]) -> None: @@ -4603,10 +4658,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): future.exception(), ) await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( - increment_list=group_operations, + increment_list=[*group_operations], ) - batch.script(TOKEN_INCREMENT_SCRIPT, script, keys, args).on_settled(fall_back) + batch.script(source, script, keys, args).on_settled(fall_back) return True @with_service_target("rate_limits") @@ -5030,36 +5085,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): targets.extend(("tag", tag) for tag in sorted(tpm_limited_tags)) return targets - def _admitted_scopes( - self, descriptors: Sequence[RateLimitDescriptor], requested_model: str | None - ) -> AdmittedRateLimitScopes: - model: Final = self._rate_limited_model(requested_model) - return AdmittedRateLimitScopes( - scopes=frozenset((d["key"], d["value"]) for d in descriptors), - model_group=model.group if model is not None else None, - ) - - @staticmethod - def _rate_limited_targets( - targets: Sequence[tuple[str, str]], - admitted: AdmittedRateLimitScopes | None, - model_group: str | None, - ) -> Sequence[tuple[str, str]]: - """ - Keep the TPM targets admission built a rate limit descriptor for, plus - the per-model targets when tokens are settled against a model group - admission did not check. Every target stays when no admission was - recorded for the call. - """ - if admitted is None: - return targets - served_admitted_group: Final = model_group == admitted.model_group - return tuple( - target - for target in targets - if target in admitted.scopes or (not served_admitted_group and target[0] in _MODEL_SCOPED_TPM_TARGETS) - ) - def _build_reservation_aware_tpm_ops( self, targets: Sequence[tuple[str, str]], @@ -5096,6 +5121,24 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) return ops + def _build_open_window_tpm_ops( + self, + targets: Sequence[tuple[str, str]], + tokens: int, + ) -> list[OpenWindowIncrementOperation]: + """Charge ``tokens`` to each target's counter while that scope has an open rate limit window.""" + if tokens == 0: + return [] + return [ + OpenWindowIncrementOperation( + key=self.create_rate_limit_keys(scope_key, scope_value, "tokens"), + increment_value=tokens, + ttl=self.window_size, + open_window_key=f"{{{scope_key}:{scope_value}}}:window", + ) + for scope_key, scope_value in targets + ] + def _build_project_reservation_op( self, scope: tuple[str, str], @@ -5240,15 +5283,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): for scope in sorted(reserved_scopes) if scope[0] in ("agent", "agent_session") and scope not in targets ) - limited_targets: Final = self._rate_limited_targets( - targets=targets, - admitted=stash.admitted_scopes if stash is not None else None, - model_group=reconcile_model.group if reconcile_model is not None else None, - ) charged_targets: Final = ( - [target for target in limited_targets if target[0] != "model_per_team"] + [target for target in targets if target[0] != "model_per_team"] if self._key_owns_model_tpm_limit_from_request_metadata(request_metadata, reconcile_model) - else limited_targets + else targets ) if reserved_tokens > 0 and total_tokens < reserved_tokens: verbose_proxy_logger.debug( @@ -5259,12 +5297,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) pipeline_operations.extend( self._build_reservation_aware_tpm_ops( - targets=charged_targets, + targets=[target for target in charged_targets if target in reserved_scopes], reserved_scopes=reserved_scopes, actual_tokens=total_tokens, reserved_tokens=reserved_tokens, ) ) + pipeline_operations.extend( + self._build_open_window_tpm_ops( + targets=[target for target in charged_targets if target not in reserved_scopes], + tokens=total_tokens, + ) + ) pipeline_operations.extend( self._build_team_ptu_tpm_ops( standard_logging_metadata, # pyright: ignore[reportUnknownArgumentType] # untyped logging metadata diff --git a/tests/integration/spend/test_rate_limit_token_counters.py b/tests/integration/spend/test_rate_limit_token_counters.py index d51dbe6b428..dc0d133900d 100644 --- a/tests/integration/spend/test_rate_limit_token_counters.py +++ b/tests/integration/spend/test_rate_limit_token_counters.py @@ -7,18 +7,23 @@ from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass, replace from hashlib import sha256 +from pathlib import Path from typing import Final import anthropic import httpx import openai import pytest -from integration._support.client import Gateway, Scenario, eventually, held +import yaml +from integration._support.client import Gateway, Scenario, eventually, held, object_value from integration._support.openai_wire import chat_reply, responses_reply +from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue from redis import Redis +from litellm.constants import LITELLM_PROXY_ADMIN_NAME + USAGE: Final = {"prompt_tokens": 12, "completion_tokens": 8, "total_tokens": 20} WIRE_TOTAL_TOKENS: Final = 8 Fields = Callable[[str], Mapping[str, JsonValue]] @@ -316,7 +321,7 @@ def test_every_endpoint_charges_only_the_rate_limited_scope(gateway: Gateway, en assert_only_charged(cache, identity.scopes, {"api_key"}, WIRE_TOTAL_TOKENS) -def test_fallback_hop_charges_the_rate_limited_model_it_served(gateway: Gateway) -> None: +def test_fallback_hop_charges_the_open_window_of_the_rate_limited_model_it_served(gateway: Gateway) -> None: prompt: Final = f"token counters fallback {uuid.uuid4().hex}" with ( wire_server(outage) as primary_upstream, @@ -332,11 +337,13 @@ def test_fallback_hop_charges_the_rate_limited_model_it_served(gateway: Gateway) key=lambda _model: {"model_tpm_limit": {fallback: 100000}}, ) identity: Final = caller(scenario, cell, fallback) + direct: Final = chat(gateway, identity, fallback, prompt + " direct") + assert direct.status_code == 200, direct.text response: Final = chat(gateway, identity, primary, prompt, fallbacks=[fallback]) assert response.status_code == 200, response.text assert response.headers["x-litellm-model-group"] == fallback assert len(primary_upstream.drain()) == 1 - assert len(fallback_upstream.drain()) == 1 + assert len(fallback_upstream.drain()) == 2 hashed: Final = sha256(identity.key.encode()).hexdigest() scopes: Final = { @@ -346,7 +353,7 @@ def test_fallback_hop_charges_the_rate_limited_model_it_served(gateway: Gateway) cache, {**scopes, "primary_model_per_key": counter("model_per_key", f"{hashed}:{primary}")}, {"model_per_key"}, - USAGE["total_tokens"], + 2 * USAGE["total_tokens"], ) @@ -420,3 +427,103 @@ def test_upstream_outage_charges_no_token_counter(gateway: Gateway) -> None: lambda values: set(values.values()) == {None}, holding=3, ) + + +def test_numeric_body_user_is_charged_to_its_rate_limited_end_user_counter(gateway: Gateway) -> None: + customer: Final = uuid.uuid4().int % 10**15 + with wire_server(completion) as upstream, gateway.scenario() as scenario, redis_counters() as cache: + model: Final = scenario.model(api_base=upstream.url + "/v1") + gateway.post("/customer/new", {"user_id": str(customer), "rpm_limit": 1000}) + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [str(customer)]}) + identity: Final = replace( + caller(scenario, Cell("numeric_end_user", charged="end_user"), model), body={"user": customer} + ) + response: Final = chat(gateway, identity, model, f"token counters numeric user {uuid.uuid4().hex}") + assert response.status_code == 200, response.text + assert [json.loads(seen.body)["user"] for seen in upstream.drain()] == [customer] + + assert_only_charged( + cache, + {**identity.scopes, "end_user": counter("end_user", str(customer))}, + {"end_user"}, + USAGE["total_tokens"], + ) + + +USER_HEADER: Final = "x-integration-openwebui-user" + + +def proxy_admin(scenario: Scenario, tpm_limit: int) -> str: + return scenario.user(user_id=LITELLM_PROXY_ADMIN_NAME, user_role="proxy_admin", tpm_limit=tpm_limit) + + +def internal_user(scenario: Scenario, tpm_limit: int) -> str: + return scenario.user(user_role="internal_user", tpm_limit=tpm_limit) + + +def master_key_request(scenario: Scenario, _user: str, body: Mapping[str, JsonValue]) -> httpx.Response: + return scenario.gateway.request("POST", "/v1/chat/completions", body) + + +def header_mapped_service_key_request( + scenario: Scenario, user: str, body: Mapping[str, JsonValue] +) -> httpx.Response: + return scenario.gateway.request( + "POST", "/v1/chat/completions", body, key=scenario.key(), headers={USER_HEADER: user} + ) + + +def config_with_user_header_mapping(directory: Path) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["general_settings"] = { + **object_value(config["general_settings"]), + "user_header_mappings": [{"header_name": USER_HEADER, "litellm_user_role": "internal_user"}], + } + path: Final = directory / "config.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.mark.parametrize( + "limited_user, unchecked_request", + [(proxy_admin, master_key_request), (internal_user, header_mapped_service_key_request)], + ids=["master_key", "user_header_mapping"], +) +def test_unchecked_traffic_counts_toward_the_user_limit_another_request_opened( + gateway: Gateway, + tmp_path: Path, + limited_user: Callable[[Scenario, int], str], + unchecked_request: Callable[[Scenario, str, Mapping[str, JsonValue]], httpx.Response], +) -> None: + redis_state: Final = {"REDIS_HOST": os.environ["REDIS_HOST"], "REDIS_PORT": os.environ["REDIS_PORT"]} + with ( + wire_server(completion) as upstream, + owned_proxy(gateway, tmp_path, redis_state, config=config_with_user_header_mapping(tmp_path)) as candidate, + candidate.scenario() as scenario, + redis_counters() as cache, + ): + model: Final = scenario.model(api_base=upstream.url + "/v1") + user: Final = limited_user(scenario, 100000) + user_counter: Final = {"user": counter("user", user)} + cache.delete(counter("user", user), f"{{user:{user}}}:window") + + checked: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"checked {uuid.uuid4().hex}"}]}, + key=scenario.key(user_id=user), + ) + assert checked.status_code == 200, checked.text + eventually(lambda: read_counters(cache, user_counter)["user"], lambda tokens: tokens == USAGE["total_tokens"]) + + unchecked: Final = unchecked_request( + scenario, user, {"model": model, "messages": [{"role": "user", "content": f"unchecked {uuid.uuid4().hex}"}]} + ) + assert unchecked.status_code == 200, unchecked.text + assert len(upstream.drain()) == 2 + + eventually( + lambda: read_counters(cache, user_counter)["user"], + lambda tokens: tokens == 2 * USAGE["total_tokens"], + seconds=30, + ) diff --git a/tests/integration/spend/test_redis_ttl_preserving_token_increment.py b/tests/integration/spend/test_redis_ttl_preserving_token_increment.py index 8b7a21f73d5..a35602bdc98 100644 --- a/tests/integration/spend/test_redis_ttl_preserving_token_increment.py +++ b/tests/integration/spend/test_redis_ttl_preserving_token_increment.py @@ -7,6 +7,7 @@ from redis import Redis from litellm.caching.caching import DualCache from litellm.caching.redis_cache import RedisCache +from litellm.proxy.hooks.parallel_request_limiter_v3 import OpenWindowIncrementOperation from litellm.proxy.hooks.parallel_request_limiter_v3 import ( PROXY_MaxParallelRequestsHandler_v3 as _PROXY_MaxParallelRequestsHandler, ) @@ -63,3 +64,43 @@ async def test_async_increment_tokens_with_ttl_preservation() -> None: finally: await redis_cache.async_delete_cache(key_with_ttl) await redis_cache.async_delete_cache(key_without_ttl) + + +@pytest.mark.asyncio +async def test_open_window_token_increment_charges_only_scopes_whose_window_is_open() -> None: + redis_host: Final = os.environ["REDIS_HOST"] + redis_port: Final = int(os.environ["REDIS_PORT"]) + redis_cache: Final = RedisCache(host=redis_host, port=redis_port) + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache(redis_cache=redis_cache)) + ) + suffix: Final = uuid.uuid4().hex[:8] + limited: Final = f"{{user:limited-{suffix}}}" + unlimited: Final = f"{{user:unlimited-{suffix}}}" + ungated: Final = f"{{team:ungated-{suffix}}}" + + with Redis(host=redis_host, port=redis_port) as raw: + try: + assert raw.set(f"{limited}:window", "1", ex=60) is True + assert raw.set(f"{limited}:tokens", "5", ex=60) is True + await handler.async_increment_tokens_with_ttl_preservation( + pipeline_operations=[ + RedisPipelineIncrementOperation(key=f"{ungated}:tokens", increment_value=20, ttl=60), + *( + OpenWindowIncrementOperation( + key=f"{scope}:tokens", increment_value=20, ttl=60, open_window_key=f"{scope}:window" + ) + for scope in (limited, unlimited) + ), + ] + ) + + assert raw.mget([f"{limited}:tokens", f"{unlimited}:tokens", f"{unlimited}:window"]) == [b"25", None, None] + assert raw.get(f"{ungated}:tokens") == b"20" + assert 0 < raw.ttl(f"{limited}:tokens") <= 60 + assert 0 < raw.ttl(f"{ungated}:tokens") <= 60 + finally: + raw.delete( + *(f"{scope}:{suffix_key}" for scope in (limited, unlimited) for suffix_key in ("window", "tokens")), + f"{ungated}:tokens", + ) diff --git a/tests/unit/caching/test_request_redis_batch_post_call.py b/tests/unit/caching/test_request_redis_batch_post_call.py index 5de785e44db..c23f88fa891 100644 --- a/tests/unit/caching/test_request_redis_batch_post_call.py +++ b/tests/unit/caching/test_request_redis_batch_post_call.py @@ -30,8 +30,10 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.proxy.auth.auth_object_prefetch import AUTH_OBJECTS_TARGET from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + OPEN_WINDOW_TOKEN_INCREMENT_SCRIPT, PARALLEL_RELEASE_SCRIPT, TOKEN_INCREMENT_SCRIPT, + OpenWindowIncrementOperation, ParallelSlotAcquisition, RequestRateLimiterStash, PROXY_MaxParallelRequestsHandler_v3, @@ -231,6 +233,67 @@ async def test_a_failed_operation_fails_only_its_owner_and_the_owner_applies_its assert await limiter.internal_usage_cache.dual_cache.in_memory_cache.async_get_cache("{team:t1}:tokens") is None +def _open_window_op(scope: str) -> OpenWindowIncrementOperation: + return OpenWindowIncrementOperation( + key=scope + ":tokens", increment_value=10, ttl=60, open_window_key=scope + ":window" + ) + + +@pytest.mark.asyncio +async def test_a_group_with_open_window_charges_rides_the_post_call_pipeline_as_one_script(): + client = FakeClient(_ok_replies) + redis_cache = PostCallFakeRedisCache(client) + limiter = _limiter(redis_cache) + + with request_redis_batch_scope(): + await limiter.async_increment_tokens_with_ttl_preservation( + [*_token_ops("{team:t1}:tokens"), _open_window_op("{api_key:k1}")] + ) + await limiter.async_increment_tokens_with_ttl_preservation(_token_ops("{user:u1}:tokens")) + await flush_post_call_redis_batches() + + assert len(client.pipelines) == 1 + assert [command for command in client.pipelines[0].commands if command[0] == "EVALSHA"] == [ + ( + "EVALSHA", + sha_of(OPEN_WINDOW_TOKEN_INCREMENT_SCRIPT), + 4, + "{team:t1}:tokens", + "{team:t1}:tokens", + "{api_key:k1}:window", + "{api_key:k1}:tokens", + 10, + 60, + 10, + 60, + ), + ("EVALSHA", sha_of(TOKEN_INCREMENT_SCRIPT), 1, "{user:u1}:tokens", 10, 60), + ] + assert redis_cache.alone == [] + + +@pytest.mark.asyncio +async def test_a_failed_open_window_script_falls_back_to_the_plain_increment(): + def replies(command: tuple[object, ...]) -> object: + if command[0] == "EVALSHA" and command[1] == sha_of(OPEN_WINDOW_TOKEN_INCREMENT_SCRIPT): + return Exception("ERR Lua") + return _ok_replies(command) + + client = FakeClient(replies) + redis_cache = PostCallFakeRedisCache(client) + limiter = _limiter(redis_cache) + + with request_redis_batch_scope(): + await limiter.async_increment_tokens_with_ttl_preservation( + [*_token_ops("{team:t1}:tokens"), _open_window_op("{user:u1}")] + ) + await flush_post_call_redis_batches() + + assert redis_cache.alone == [("INCRBYFLOAT", "{team:t1}:tokens", 10), ("INCRBYFLOAT", "{user:u1}:tokens", 10)] + assert await limiter.internal_usage_cache.dual_cache.in_memory_cache.async_get_cache("{team:t1}:tokens") == 10 + assert await limiter.internal_usage_cache.dual_cache.in_memory_cache.async_get_cache("{user:u1}:tokens") == 10 + + @pytest.mark.asyncio async def test_a_failed_slot_release_script_releases_the_slot_in_memory(): def replies(command: tuple[object, ...]) -> object: diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 82b5bab8362..f4c9371f8b7 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -6,7 +6,7 @@ import asyncio import logging import sys import time -from collections.abc import Awaitable, Callable, Iterator, Sequence +from collections.abc import Iterator, Sequence from contextlib import contextmanager from datetime import datetime, timedelta, timezone from typing import Any, Dict, Final, List, Optional @@ -8688,6 +8688,7 @@ async def test_a2a_url_target_owns_invocation_fee_and_request_limit( assert "expensive" in str(denied.value.detail) + _SCOPE_TOKEN_KEYS: Final = frozenset( { "{api_key:key-hash}:tokens", @@ -8700,170 +8701,69 @@ _SCOPE_TOKEN_KEYS: Final = frozenset( ) -def _scoped_success_kwargs( - model_group: str, key_metadata: Dict[str, object] | None = None, call_id: str | None = None -) -> Dict[str, object]: +def _scoped_success_kwargs(**body: object) -> dict[str, object]: return { - **({} if call_id is None else {"litellm_call_id": call_id}), "standard_logging_object": { "metadata": {"user_api_key_hash": "key-hash", "user_api_key_user_id": "u", "user_api_key_team_id": "t"} }, "litellm_params": { "metadata": { - "model_group": model_group, - "user_api_key_metadata": key_metadata or {}, + "model_group": "gpt-4o-mini", + "user_api_key_metadata": {}, "user_api_key_team_metadata": {}, } }, - "model": model_group, + "model": "gpt-4o-mini", + **body, } -async def _admit_scoped_request( - handler: _PROXY_MaxParallelRequestsHandler, auth: UserAPIKeyAuth, model: str, call_id: str | None = None -) -> None: +async def _admit_scoped_request(handler: _PROXY_MaxParallelRequestsHandler, auth: UserAPIKeyAuth) -> None: await handler.async_pre_call_hook( user_api_key_dict=auth, cache=handler.internal_usage_cache.dual_cache, - data={ - "model": model, - "messages": [{"role": "user", "content": "hello"}], - "max_tokens": 20, - **({} if call_id is None else {"litellm_call_id": call_id}), - }, + data={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 20}, call_type="acompletion", ) -def _charged_token_keys(handler: _PROXY_MaxParallelRequestsHandler, kwargs: Dict[str, object]) -> frozenset[str]: +def _charged_windows(handler: _PROXY_MaxParallelRequestsHandler, kwargs: dict[str, object]) -> dict[str, object]: operations: Final = handler._build_success_event_pipeline_operations( kwargs=kwargs, response_obj=ModelResponse(usage=Usage(prompt_tokens=2, completion_tokens=1, total_tokens=3)), rate_limit_type="total", ) - return frozenset(operation["key"] for operation in operations) + return {operation["key"]: operation.get("open_window_key") for operation in operations} @pytest.mark.asyncio @pytest.mark.parametrize( - "limits, charged", + "limits, reserved", [ ({}, frozenset()), - ({"rpm_limit": 100}, frozenset({"{api_key:key-hash}:tokens"})), - ({"max_parallel_requests": 10}, frozenset({"{api_key:key-hash}:tokens"})), + ({"rpm_limit": 100}, frozenset()), + ({"user_rpm_limit": 100}, frozenset()), ({"tpm_limit": 1000}, frozenset({"{api_key:key-hash}:tokens"})), - ({"tpm_limit": 1000, "metadata": {"tpm_limit_type": "dynamic"}}, frozenset({"{api_key:key-hash}:tokens"})), - ({"user_rpm_limit": 100}, frozenset({"{user:u}:tokens"})), ({"team_tpm_limit": 1000}, frozenset({"{team:t}:tokens"})), - ({"team_member_rpm_limit": 100}, frozenset({"{team_member:t:u}:tokens"})), - ( - {"metadata": {"model_rpm_limit": {"gpt-4o-mini": 100}}}, - frozenset({"{model_per_key:key-hash:gpt-4o-mini}:tokens"}), - ), - ], - ids=[ - "no_limits", - "key_rpm", - "key_max_parallel", - "key_tpm_reserved", - "key_dynamic_tpm", - "user_rpm", - "team_tpm_reserved", - "team_member_rpm", - "key_model_rpm", ], + ids=["no_limits", "key_rpm", "user_rpm", "key_tpm_reserved", "team_tpm_reserved"], ) -async def test_success_charges_tokens_only_to_scopes_admission_rate_limited( - limits: Dict[str, object], charged: frozenset[str] +async def test_success_charges_every_unreserved_scope_only_while_its_rate_limit_window_is_open( + limits: dict[str, object], reserved: frozenset[str] ) -> None: handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) - await _admit_scoped_request( - handler, UserAPIKeyAuth(api_key="key-hash", user_id="u", team_id="t", **limits), "gpt-4o-mini" - ) + await _admit_scoped_request(handler, UserAPIKeyAuth(api_key="key-hash", user_id="u", team_id="t", **limits)) - assert _charged_token_keys(handler, _scoped_success_kwargs("gpt-4o-mini")) == charged - - -@pytest.mark.asyncio -async def test_fallback_success_charges_the_model_scopes_of_the_group_that_served_it() -> None: - key_metadata: Final = {"model_tpm_limit": {"fallback-model": 1000}} - handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) - await _admit_scoped_request( - handler, UserAPIKeyAuth(api_key="key-hash", user_id="u", team_id="t", metadata=key_metadata), "primary-model" - ) - - assert _charged_token_keys(handler, _scoped_success_kwargs("fallback-model", key_metadata)) == { - "{model_per_key:key-hash:fallback-model}:tokens" + assert _charged_windows(handler, _scoped_success_kwargs()) == { + key: None if key in reserved else key.removesuffix(":tokens") + ":window" for key in _SCOPE_TOKEN_KEYS } - @pytest.mark.asyncio -@pytest.mark.parametrize( - "success_end_user, charged", - [ - (None, frozenset({"{end_user:admitted-customer}:tokens"})), - ("admitted-customer", frozenset({"{end_user:admitted-customer}:tokens"})), - ("other-customer", frozenset()), - ], - ids=["metadata_end_user", "same_body_user", "different_body_user"], -) -async def test_success_charges_the_end_user_admission_checked_and_no_other( - success_end_user: str | None, charged: frozenset[str] -) -> None: +async def test_success_charges_a_numeric_body_user_to_the_window_its_admission_opened() -> None: handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) await _admit_scoped_request( - handler, - UserAPIKeyAuth(api_key="key-hash", end_user_id="admitted-customer", end_user_rpm_limit=100), - "gpt-4o-mini", - ) - kwargs: Final = _scoped_success_kwargs("gpt-4o-mini") - kwargs["standard_logging_object"] = { - "metadata": {"user_api_key_hash": "key-hash", "user_api_key_end_user_id": "admitted-customer"} - } - if success_end_user is not None: - kwargs["user"] = success_end_user - - assert _charged_token_keys(handler, kwargs) == charged - - -@pytest.mark.asyncio -async def test_success_on_an_aliased_model_charges_the_model_limit_of_the_resolved_group() -> None: - key_metadata: Final = {"model_rpm_limit": {"target": 100}} - handler: Final = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()), model_group_resolver=_resolve_alias_to_target - ) - await _admit_scoped_request( - handler, UserAPIKeyAuth(api_key="key-hash", user_id="u", team_id="t", metadata=key_metadata), "alias" + handler, UserAPIKeyAuth(api_key="key-hash", end_user_id="123", end_user_rpm_limit=100) ) - assert _charged_token_keys(handler, _scoped_success_kwargs("target", key_metadata)) == { - "{model_per_key:key-hash:target}:tokens" - } - -async def _no_stash(handler: _PROXY_MaxParallelRequestsHandler) -> str | None: - _request_stash.set(None) - return None - - -async def _nested_call(handler: _PROXY_MaxParallelRequestsHandler) -> str | None: - await _admit_scoped_request( - handler, UserAPIKeyAuth(api_key="key-hash", user_id="u", team_id="t"), "gpt-4o-mini", "proxy-request" - ) - return "nested-call" - - -async def _stash_without_admission(handler: _PROXY_MaxParallelRequestsHandler) -> str | None: - _request_stash.set(RequestRateLimiterStash()) - return None - - -@pytest.mark.asyncio -@pytest.mark.parametrize("arrange", [_no_stash, _nested_call, _stash_without_admission]) -async def test_success_charges_every_scope_when_the_call_was_not_admitted( - arrange: Callable[[_PROXY_MaxParallelRequestsHandler], Awaitable[str | None]], -) -> None: - handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) - call_id: Final = await arrange(handler) - - assert _charged_token_keys(handler, _scoped_success_kwargs("gpt-4o-mini", call_id=call_id)) == _SCOPE_TOKEN_KEYS + assert _charged_windows(handler, _scoped_success_kwargs(user=123))["{end_user:123}:tokens"] == "{end_user:123}:window" From 5b94aee8b3dd7debd93b9b525e5f2a84812f8555 Mon Sep 17 00:00:00 2001 From: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 18:40:43 -0700 Subject: [PATCH 3/4] test(rate-limiter): isolate the master key cell's Redis and pin healthy dynamic TPM as unchecked The master key cell shared the worker Redis with the main gateway, whose own master key traffic caches a default_user_id user object without the cell's TPM limit, so the owned proxy could admit the cell's key without the user limit. The cell now runs the owned proxy on a Redis of its own. A dynamic key TPM limit is not checked while the model is healthy, so it opens no window; the cell now expects only the user's RPM-opened counter. --- .../spend/test_rate_limit_token_counters.py | 22 +++++++++++-------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/tests/integration/spend/test_rate_limit_token_counters.py b/tests/integration/spend/test_rate_limit_token_counters.py index dc0d133900d..1936b6dc260 100644 --- a/tests/integration/spend/test_rate_limit_token_counters.py +++ b/tests/integration/spend/test_rate_limit_token_counters.py @@ -18,6 +18,7 @@ import yaml from integration._support.client import Gateway, Scenario, eventually, held, object_value from integration._support.openai_wire import chat_reply, responses_reply from integration._support.process import owned_proxy +from integration._support.redis_process import owned_redis from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue from redis import Redis @@ -167,9 +168,10 @@ CELLS: Final = ( Cell("tag_tpm", charged="tag", tag=lambda _model: {"tpm_limit": 100000}), Cell("key_tpm_reserved", charged="api_key", key=lambda _model: {"tpm_limit": 100000}), Cell( - "key_dynamic_tpm", - charged="api_key", + "key_dynamic_tpm_on_a_healthy_model", + charged="user", key=lambda _model: {"tpm_limit": 100000, "tpm_limit_type": "dynamic"}, + user=lambda _model: {"rpm_limit": 1000}, ), ) @@ -465,9 +467,7 @@ def master_key_request(scenario: Scenario, _user: str, body: Mapping[str, JsonVa return scenario.gateway.request("POST", "/v1/chat/completions", body) -def header_mapped_service_key_request( - scenario: Scenario, user: str, body: Mapping[str, JsonValue] -) -> httpx.Response: +def header_mapped_service_key_request(scenario: Scenario, user: str, body: Mapping[str, JsonValue]) -> httpx.Response: return scenario.gateway.request( "POST", "/v1/chat/completions", body, key=scenario.key(), headers={USER_HEADER: user} ) @@ -495,17 +495,21 @@ def test_unchecked_traffic_counts_toward_the_user_limit_another_request_opened( limited_user: Callable[[Scenario, int], str], unchecked_request: Callable[[Scenario, str, Mapping[str, JsonValue]], httpx.Response], ) -> None: - redis_state: Final = {"REDIS_HOST": os.environ["REDIS_HOST"], "REDIS_PORT": os.environ["REDIS_PORT"]} with ( wire_server(completion) as upstream, - owned_proxy(gateway, tmp_path, redis_state, config=config_with_user_header_mapping(tmp_path)) as candidate, + owned_redis(tmp_path / "redis") as private_redis, + owned_proxy( + gateway, + tmp_path, + {"REDIS_HOST": private_redis.host, "REDIS_PORT": str(private_redis.port)}, + config=config_with_user_header_mapping(tmp_path), + ) as candidate, candidate.scenario() as scenario, - redis_counters() as cache, + Redis(host=private_redis.host, port=private_redis.port) as cache, ): model: Final = scenario.model(api_base=upstream.url + "/v1") user: Final = limited_user(scenario, 100000) user_counter: Final = {"user": counter("user", user)} - cache.delete(counter("user", user), f"{{user:{user}}}:window") checked: Final = candidate.request( "POST", From acdcad9769531a59d034e9e0a5d700657bf9bbae Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 26 Oct 2026 18:11:37 -0700 Subject: [PATCH 4/4] fix(rate-limiter): charge token counters of calls that skipped admission without waiting for a window A proxy sub-call, such as the compaction summary or the safeguards classifier, never passes admission, so the scopes it charges have no open window and the sub-call precheck reads their counters directly. Gating those charges on a window left a key's summary-model TPM limit unenforced; they are now charged as before. --- .../hooks/parallel_request_limiter_v3.py | 13 ++++++++--- .../test_safeguard_results_compaction.py | 14 +++++++++++ .../hooks/test_parallel_request_limiter_v3.py | 23 +++++++++++++++++++ 3 files changed, 47 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index ee3aac84d4a..4bc67202e14 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -5325,10 +5325,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): reserved_tokens=reserved_tokens, ) ) + unreserved_targets: Final = [target for target in charged_targets if target not in reserved_scopes] + # A call that never passed admission, such as a proxy sub-call, opened no window, and the sub-call + # precheck reads its counters directly, so it charges them unconditionally pipeline_operations.extend( - self._build_open_window_tpm_ops( - targets=[target for target in charged_targets if target not in reserved_scopes], - tokens=total_tokens, + self._build_open_window_tpm_ops(targets=unreserved_targets, tokens=total_tokens) + if stash is not None + else self._build_reservation_aware_tpm_ops( + targets=unreserved_targets, + reserved_scopes=reserved_scopes, + actual_tokens=total_tokens, + reserved_tokens=0, ) ) pipeline_operations.extend( diff --git a/tests/integration/messages_endpoint/safeguards/test_safeguard_results_compaction.py b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_compaction.py index 38aef23c630..587fef4e2f5 100644 --- a/tests/integration/messages_endpoint/safeguards/test_safeguard_results_compaction.py +++ b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_compaction.py @@ -140,6 +140,20 @@ def test_caller_without_the_summary_model_sees_the_edit_error(compacting: Rig) - assert len(compacting.backend_calls(received)) == 1 +def test_caller_over_the_summary_models_tpm_limit_sees_the_edit_error(compacting: Rig) -> None: + with compacting.gateway.scenario() as scenario: + key: Final = scenario.key(model_tpm_limit={CLASSIFIER_MODEL: 1}) + within_nonce: Final = nonce() + compacting.register(within_nonce, chat_backend(tool_id(within_nonce), stream=False), summarizing_or_classifying) + within: Final = applied_edit(json_body(compacting.post(long_request(within_nonce), key=key))) + assert within["type"] == "compact_20260112" and "error" not in within, within + over_nonce: Final = nonce() + compacting.register(over_nonce, chat_backend(tool_id(over_nonce), stream=False), summarizing_or_classifying) + over: Final = json_body(compacting.post(long_request(over_nonce), key=key)) + assert applied_edit(over)["error"] == "summary_model_rate_limit_exceeded", over + assert compacting.classifier_calls(compacting.received(over_nonce)) == () + + def assert_summarized_and_classified_on_their_models(rig: Rig, request_nonce: str, body: dict[str, JsonValue]) -> str: received: Final = rig.received(request_nonce) targets: Final = [request.target for request in received] diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 4fba76555ee..4c29d460521 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -8842,3 +8842,26 @@ async def test_success_charges_a_numeric_body_user_to_the_window_its_admission_o ) assert _charged_windows(handler, _scoped_success_kwargs(user=123))["{end_user:123}:tokens"] == "{end_user:123}:window" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("parent_admitted", [True, False], ids=["sub_call_of_an_admitted_request", "no_admission"]) +async def test_success_of_a_call_that_skipped_admission_charges_every_scope_without_waiting_for_a_window( + parent_admitted: bool, +) -> None: + handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + if parent_admitted: + await handler.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="key-hash", user_id="u", team_id="t"), + cache=handler.internal_usage_cache.dual_cache, + data={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}], + "litellm_call_id": "parent", + }, + call_type="acompletion", + ) + + assert _charged_windows(handler, _scoped_success_kwargs(litellm_call_id="sub-call")) == dict.fromkeys( + _SCOPE_TOKEN_KEYS + )