Repository navigation
fix(gateway): enforce the token rate limit instead of only accepting it - #3570
Conversation
rate_limit_tpm was accepted on a key, loaded into the key store, and read by
nobody. global_tpm came the same way from the control plane, was copied to
rc.Metadata["org_rate_limit_tpm"], and was read by nobody either. The rate
limit plugin only ever looked at RPM. An operator who set a token ceiling for
cost control had no ceiling at all, and nothing said so.
Tokens cannot be counted before a call: the cost is known once the provider
answers. So the window is read before the request and charged after it, which
lets one request cross the limit before the next is refused. That suits a
limit meant to bound spend over a minute, and the refusal carries its own code
rather than the request-limit one:
429 token_rate_limit_exceeded
Precedence matches RPM — the key's own limit, then the organization's. The two
budgets are independent: a token limit is now checked even when no request
limit is set, which it was not before, and the counters are in separate
windows.
The token budget lives in the in-memory limiter, so with the Redis limiter it
counts per pod and the effective limit is the configured one times the replica
count. The alternative was a second Lua sliding window counting a number only
the answering pod knows, in a file whose existing script no test in CI can
run. Per-pod matches how providers.max_concurrent already behaves, and is
written down rather than left to be discovered.
A limiter that cannot hold a token budget now says so once, at WARN, instead
of leaving the limit silently inert — which is the whole defect being fixed
here, and not worth reintroducing one interface lower.
Tanmaycode1
left a comment
Tanmaycode1
left a comment
There was a problem hiding this comment.
Confirmed one blocking correctness finding on this head: cached responses are charged to the token budget again, causing false quota exhaustion and incorrect 429s. The inline comment includes the reproduction and suggested fix. Verified through the actual pipeline engine, cache plugin, rate-limit plugin, key store, and in-memory limiter; the regression test passes on base 03fe69d and fails on this head. Existing gateway tests and scoped race checks passed, but do not cover this scenario.
| if !ok { | ||
| return pipeline.ResultContinue() | ||
| } | ||
| limiter.AddTokens(tokenWindowKey(keyID), tokensUsed(rc)) |
There was a problem hiding this comment.
[P1 — blocks merge] Do not charge cached response usage as new upstream token consumption (§PB; introduced by this PR).
RunPostPlugins also calls this method when the cache short-circuits the request. The cached response retains the usage from its original provider call, so this unconditional AddTokens charges those tokens again without any new provider work.
Confirmed reproduction using the actual pipeline engine, cache plugin, rate-limit plugin, key store, and in-memory limiter:
- Set a key's TPM limit to 100.
- Send a request that calls the provider and returns 60 total tokens.
- Repeat the identical request:
cache_status=hit_exact, with the provider-call count still 1. - Send a different request: this PR rejects it with HTTP 429 /
token_rate_limit_exceeded, even though only 60 upstream tokens have been used.
The same regression test passes on base 03fe69df9; on this head it fails because the cache hit raises the counter to 120. This is an incorrect new rejection, not the intended activation of a genuinely exhausted token limit.
Please skip accounting when rc.Flags.CacheHit is true and add a pipeline integration regression for this sequence. Cover semantic cache hits too: implementing only ShouldSkipOnCacheHit would currently skip exact hits, because the engine's interface check is restricted to cache_status == "hit_exact".
Review found this PR charged cached responses as new consumption. The engine runs post-plugins on a cache short-circuit too, and a cached response keeps the usage of the provider call that filled the cache, so the window was charged again for work nobody did. Repeated identical requests exhausted the budget and the next real request was refused with 429. ProcessResponse now returns early on rc.Flags.CacheHit. The flag is what to read rather than the pipeline's SkipOnCacheHit interface: the engine applies that only when cache_status is "hit_exact", so a conditional or semantic hit would still have been charged. Test covers all three hit kinds and asserts the remaining budget afterwards. Against this PR's previous head it fails on the conditional hit, refusing a request with 60 of 100 tokens actually spent.
Tanmaycode1
left a comment
Tanmaycode1
left a comment
There was a problem hiding this comment.
Re-reviewed the latest push. Approving: the reported cache double-charging blocker is fixed, and no remaining blockers were found.
Replayed the original real pipeline/cache/rate-limit sequence: the first provider call consumes 60 of 100 tokens, the exact cache hit adds no consumption, and the next fresh request reaches the provider. The same reproduction fails on the previous implementation. Exact, conditional and semantic cache-hit paths set rc.Flags.CacheHit, so the new guard covers all three.
Validation: the full gateway suite passed across 82 packages, with 23 Redis/Google-credential-dependent skips. Targeted ratelimit/cache/pipeline/redisstate tests passed under -race -count=2; touched-package vet, formatting and contract checks passed. The E2E coverage gate passes with the declared harness-gap exemption.
The documented post-response accounting and per-pod token-budget limitations remain intentional scope; they are not unresolved review blockers.
rate_limit_tpm and an org's global_tpm were accepted, stored and read by nobody, so a configured token ceiling did nothing. The window is read before the request and charged after it, because a request's cost is known only once the provider answers. A cache hit is not charged. Refusal carries its own code, token_rate_limit_exceeded. Cherry-picked from future-agi#3570 on dev (d24818e).
rate_limit_tpmwas accepted on a key, loaded into the key store, and read by nobody.global_tpmarrived the same way from the control plane, was copied intorc.Metadata["org_rate_limit_tpm"], and was read by nobody either. The rate limit plugin only ever looked at RPM.So an operator who set a token ceiling for cost control had no ceiling, and nothing anywhere said so. It is in the cross-service contract on both sides —
agentcc/contracts/gateway_admin.pyand the gateway's own generated contract — so removing the field would mean changing a contract to match a bug.The change
Tokens cannot be counted before a call: the cost is known once the provider answers. So the window is read in
ProcessRequestand charged inProcessResponse, which lets one request cross the limit before the next is refused. That suits a limit meant to bound spend over a minute, and it is stated rather than hidden.The refusal carries its own code, so a client can tell which budget it hit:
Precedence matches RPM: the key's own limit, then the organization's. The two budgets are independent — a token limit is now checked even when no request limit is set, which it was not, because the token check sat after the
rpm == 0early return. Their counters are in separate windows.ratelimit_tokens_limitandratelimit_tokens_remaininggo on the request context beside the existing RPM metadata.The per-pod limitation, deliberately
The token budget lives in the in-memory limiter. With the Redis limiter configured,
AllowTokensandAddTokensdelegate to its local fallback, so a token limit counts per pod and the effective limit is the configured one times the replica count.The alternative was a second Lua sliding window, counting a number only the answering pod knows, in a package whose existing script no test in CI can run —
internal/redisstatetests the Redis path only through its nil-client fallback. Shipping an untested Lua counter to close a dead-config finding is a worse trade than a documented approximation. Per-pod is also howproviders.max_concurrentalready behaves.A limiter that cannot hold a token budget now says so once at WARN rather than leaving the limit silently inert — that silence is the defect being fixed, and not worth reintroducing one interface lower.
Behaviour change to know about
An operator who already set
rate_limit_tpmorglobal_tpmwill start getting 429s they never got before. That is the setting working for the first time. It is worth a line in the release notes, because from the outside it looks like a new limit appearing.Tests
9 new cases. The enforcement is new API, so there is no
origin/devto fail against; instead each case was re-run with the enforcement stubbed out, and all four behaviour cases fail there:token_rate_limit_exceeded— stubbed:the second request was allowed with the window spentAllowTokenscharges nothing, so reading the window cannot bill a token per requestRateLimitersatisfiesTokenLimiterandPluginsatisfiespipeline.PluginFull gateway suite green, 82 packages.
go vetandgofmtclean on everything touched.Docs for the new limit are not in this PR; the gateway docs PR in flight covers the same page and will carry it.
E2E: exempt (harness-gap the e2e stack cannot configure auth keys or rate limits)
e2e/lib/managed-mock.tsadmits only thelogging,providersandserversections, so neitherauth.keys(whererate_limit_tpmlives) nor therate_limitingsection can be set in the e2e gateway config. Without a configured per-key limit there is no budget to exhaust. Covered by Go tests through the real plugin, key store and limiter, including the cache-hit path the review found.