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.
This commit is contained in:
yucheng 2026-10-26 18:11:37 -07:00
parent 5b94aee8b3
commit acdcad9769
3 changed files with 47 additions and 3 deletions

View file

@ -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(

View file

@ -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]

View file

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