Repository navigation
feat(security): per-token call rate limit for agent tokens - #1746
Conversation
Cap how many upstream tool calls one agent token may make per minute, so a
runaway agent loop cannot burn through an upstream's quota. Off by default.
The limit is a token bucket per agent token, enforced at the managed-client
admission seam, so it covers call_tool_*, direct routing mode, code_execution
and the REST tool-call path. Only agent-token callers are limited; the admin
API key and unauthenticated callers are not.
- agent_token_rate_limit {per_minute, burst} sets the default for every
agent token and is hot-reloadable.
- mcpproxy token create --rate-limit/--rate-burst and rate_limit on
POST /api/v1/tokens give one token its own budget; per_minute 0 exempts it.
- A refused call is a typed rate_limited LimitError: an MCP error result,
or HTTP 429 with Retry-After over REST. It never reaches the upstream and
takes no concurrency slot.
- The refusal is recorded once as a rejected tool_call naming the token, and
counted in mcpproxy_tool_calls_rejected_total{reason="rate_limited",
scope="token"}. The audit-line reason vocabulary gains rate_limited.
Refs smart-mcp-proxy#1699
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
Only oas/docs.go conflicted (generated); regenerated with make swagger.
Dumbris
left a comment
There was a problem hiding this comment.
Thank you, James. This is what we agreed on the issue and it is tidy: one bucket per token at the admission seam, a typed 429 with Retry-After, one rejected tool_call record per refusal, and the Shed() predicate replacing the hard-coded queue_full || queue_timeout checks, so the REST and MCP paths could not drift. Approving.
What I changed on your branch
oas/docs.go conflicted with main (it is generated). I merged main into the branch (82d13fe) and regenerated it with make swagger. The resulting OAS diff against main is only your new config type and the two 429 descriptions. Nothing else changed, so the PR is mergeable again.
What I verified
- Build,
go build -tags server, andgo test -racepass forinternal/{upstream,auth,config,httpapi,audit,contracts,observability,runtime,server,storage}andcmd/mcpproxy(theinternal/serverrun used theunit-tests.yml-skipregex)../scripts/test-api-e2e.shpasses (70/70) on the merged tree. - Both golangci-lint v2 runs (bare and
--build-tags server) have no findings in any file this PR touches. The ones reported are in other files. - A real
mcpproxy servewithagent_token_rate_limit {per_minute:60, burst:3}and theeverythingserver:- the default-budget token gets 3×200, then 429 with
Retry-After: 1and your message text - a token created with
rate_limit {per_minute:6, burst:2}gets 2×200, then 429 withRetry-After: 10 - a token with
rate_limit {per_minute:0}gets 8×200, and the admin key gets 8×200 rate_limitis echoed on create and inGET /api/v1/tokens;per_minute:-1is a 400 naming the field- over MCP (
/mcp, bearer token),call_tool_readreturns anisErrorresult with the same text code_executionwithburst:4: a script making 7call_toolcalls gets 4 ok, then 3 rate-limit errors, so each call a script makes counts- the activity log has
status: rejected,metadata.rejection_reason: rate_limited,rejection_scope: tokenandtoken_nameon each refusal, one row per refusal
- the default-budget token gets 3×200, then 429 with
- No Gemini or other bot review comments are on this PR (only Codecov).
- One cross-model round with opencode
github-copilot/gpt-6.1-sol, in two chunks (limiter and seam; configuration, storage and reporting). I checked each finding and none needs a change before merge. Neither chunk found a bypass of the admission seam. Both chunks confirmed that theAuthContextis built by one constructor, thatrate_limit: {per_minute: 0}survives regenerate, and that no caller can set its own budget.- Sweep before a budget change: if the burst is raised on a hot reload, a long-idle (already full) bucket is swept and re-created at the new, larger burst instead of refilling to the old size first. The extra calls are bounded by the new burst the operator chose, so I treat this as acceptable.
- A deleted token and a new token with the same name share a bucket until it refills, so the replacement can be refused for up to one refill period. This errs on the strict side and I agree with keying by owner and name.
- Rejection records carry the token name, prefix and client id but not the owner's
user_idor the effective profile, and a refusedcode_executionsub-call has noparent_id. Server-edition users therefore do not see these refusals in their own activity view, and a profile filter or drill-down misses them. The concurrency sheds have the same gap today, so I did not ask for it here; it is a small follow-up.
For Algis (the choices in the PR body):
- A token's own budget replaces the default (0 exempts it).
rate_limitedcounts inRejectedCount.- There is no route to change the budget of an existing token.
I have not merged.
|
Merged. Thank you, @james-crowley! A looping agent with a write-tier token could fire hundreds of calls before anyone looked; now each token has a per-minute budget with a proper 429 and Retry-After, and admin and exempt tokens are untouched. This is also the defence in depth behind the scan-default decision in #1713. Contributions like this are what keep MCPProxy moving, and they matter a lot to the project. One small ask, only if you feel like it: a short public note about your change (a tweet, a LinkedIn post, a line in a community you are part of) helps more people discover MCPProxy. No pressure at all, and I would be glad to repost it. |
Description
Closes #1699. Adds a per-token call rate limit for agent tokens, built on the plan in the maintainer's comment on the issue.
acquireAdmission, before the concurrency registry), so it coverscall_tool_*, direct routing mode,code_execution(each call a script makes counts) and the REST tool-call path with one check.ListToolsandPingare not counted, same as the concurrency limits.per_minute: 0). The limiter is in the repo (internal/upstream/limiter/token_rate.go, injectable clock), no new dependency. State is in memory, keyed by owner + token name (names are only unique per owner); idle buckets are swept. A restart refills every bucket.AuthTypeAgentcallers with a token name. The admin API key, the Unix socket and unauthenticated callers go through unchanged.agent_token_rate_limit: {per_minute, burst}is the default for every agent token. Hot-reloadable (goes throughapplyConcurrencyLimits; added toundiffedHotConfigFields). Negative values, values over 100000 and aburstwithout aper_minuteare validation errors that name the field.mcpproxy token create --rate-limit N [--rate-burst M]andrate_limitonPOST /api/v1/tokensstore a budget on the token. It replaces the config default for that token;per_minute: 0exempts it.GET /api/v1/tokens[/{name}]andtoken showreturn it. The record gets one new omitempty field, so no storage migration.LimitErrorreasonrate_limitedwith scopetokengoes through the existing shed plumbing instead of a parallel path (aShed()predicate replaces the hard-codedqueue_full || queue_timeoutchecks):isErrortool result.429withRetry-After(seconds, rounded up, minimum 1).tool_callin the activity log with reasonrate_limited, written by the same observer as the concurrency sheds, so there is exactly one row per refusal and no second authz line. It carries the token name, prefix and client id (the observer now passes the attribution through). The audit-linereasonvocabulary gainsrate_limited, in both schema copies.mcpproxy_tool_calls_rejected_total{reason="rate_limited",scope="token"}(existing counter, new label values).docs/features/agent-tokens.md(Call Rate Limit),docs/configuration.md,docs/configuration/config-file.md,docs/features/audit-log.md, the spec 107 audit-line contract, CHANGELOG. OAS regenerated withmake swagger(the diff is only the new config type and the two 429 descriptions).Choices I made that you may want different:
0is an explicit exemption. That gives a "no limit for this one batch token" without turning the default off for everyone. Say if you want the stricter of the two instead.burstdefaults toper_minute, soper_minute: 60alone allows 60 calls at once, then one a second.reason: rate_limitedis added to the sharedRejectedCountin the activity summary, so it counts alongside the concurrency sheds. Say if you want it separate.scope="token"without the name; the name is in the activity record).Testing
I have tested these changes locally
I have added/updated tests that prove my fix is effective or my feature works
All existing tests pass (see below)
internal/upstream/limiter: bucket math with a fake clock (burst allowed, then refused with the exactRetryAfter, refill, per-key isolation, per-token override and exemption, changed budget, sweep, concurrent callers), error text andShed().internal/upstream/managed: enforcement at the seam: burst then refusal withRetry-After, refill, per-token isolation, off by default, admin and unauthenticated callers unaffected, a refused call takes no concurrency slot.internal/server: end to end against a counting upstream throughcall_tool_read,code_execution, direct mode and the REST typed error. The upstream is never reached after the budget is spent.internal/runtime: the rejected activity record names the token; a concurrency shed stays unattributed; a real runtime with a live HTTP upstream, the metric sink and a config reload;IssueTokenvalidation and storage.internal/httpapi:rate_limiton create / get / list, validation, exemption, and the 429 +Retry-Aftermapping (1.5s becomes 2).internal/config,internal/auth,internal/audit,internal/upstream,cmd/mcpproxy: validation, JSON round trip, schema vocabulary, manager default follows reloads, CLI flags and output.go build ./...,go build -tags server ./...,go vet ./...andgo vet -tags server ./...are clean.golangci-lintv2.9.0 with.github/.golangci.ymlreports 0 issues, bare and with--build-tags server.go test -race -skip "E2E|Binary|MCPProtocol|TestInfoEndpoint|TestGracefulShutdownNoPanic|TestSocketInfoEndpoint" ./...passes, as does the server-tag CI command for./internal/{serveredition,config,oauth,server,httpapi,storage}/...and-shuffle=onon the packages above.