Skip to content

feat(security): per-token call rate limit for agent tokens - #1746

Merged
Dumbris merged 2 commits into
smart-mcp-proxy:mainfrom
james-crowley:feat/token-rate-limit
Oct 25, 2026
Merged

Dumbris merged 2 commits into
smart-mcp-proxy:mainfrom
james-crowley:feat/token-rate-limit

Conversation

@james-crowley

Copy link
Copy Markdown
Contributor

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.

  • Where it is enforced. At the managed-client admission seam (acquireAdmission, before the concurrency registry), so it covers call_tool_*, direct routing mode, code_execution (each call a script makes counts) and the REST tool-call path with one check. ListTools and Ping are not counted, same as the concurrency limits.
  • The limit. A token bucket per agent token: N calls per minute plus a burst. Off by default (absent or 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.
  • Agent tokens only. The check runs only for AuthTypeAgent callers with a token name. The admin API key, the Unix socket and unauthenticated callers go through unchanged.
  • Config. agent_token_rate_limit: {per_minute, burst} is the default for every agent token. Hot-reloadable (goes through applyConcurrencyLimits; added to undiffedHotConfigFields). Negative values, values over 100000 and a burst without a per_minute are validation errors that name the field.
  • Per token. mcpproxy token create --rate-limit N [--rate-burst M] and rate_limit on POST /api/v1/tokens store a budget on the token. It replaces the config default for that token; per_minute: 0 exempts it. GET /api/v1/tokens[/{name}] and token show return it. The record gets one new omitempty field, so no storage migration.
  • Refusal. A new LimitError reason rate_limited with scope token goes through the existing shed plumbing instead of a parallel path (a Shed() predicate replaces the hard-coded queue_full || queue_timeout checks):
    • MCP: an isError tool result.
    • REST: 429 with Retry-After (seconds, rounded up, minimum 1).
    • Message: "Rate limit reached: this agent token may make N tool calls per minute (bursts of up to M) and has used its allowance. Retry in S second(s)."
    • A refused call takes no concurrency slot and never reaches the upstream.
  • Record. The refusal is one rejected tool_call in the activity log with reason rate_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-line reason vocabulary gains rate_limited, in both schema copies.
  • Metric. mcpproxy_tool_calls_rejected_total{reason="rate_limited",scope="token"} (existing counter, new label values).
  • Docs: 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 with make swagger (the diff is only the new config type and the two 429 descriptions).

Choices I made that you may want different:

  • A token's own budget replaces the config default rather than stacking on it, and 0 is 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.
  • There is no way to change a budget on an existing token (there is no token update route today). Re-create the token, or change the config default.
  • Changing a budget keeps the calls a token has already spent (clamped to the new size); it does not refund them.
  • burst defaults to per_minute, so per_minute: 60 alone allows 60 calls at once, then one a second.
  • reason: rate_limited is added to the shared RejectedCount in the activity summary, so it counts alongside the concurrency sheds. Say if you want it separate.
  • Out of scope, as agreed on the issue: a tighter write/destructive cap, a daily budget, and a per-token metric label (token names would be unbounded cardinality, so the metric stays 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 exact RetryAfter, refill, per-key isolation, per-token override and exemption, changed budget, sweep, concurrent callers), error text and Shed().

  • internal/upstream/managed: enforcement at the seam: burst then refusal with Retry-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 through call_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; IssueToken validation and storage.

  • internal/httpapi: rate_limit on create / get / list, validation, exemption, and the 429 + Retry-After mapping (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 ./... and go vet -tags server ./... are clean. golangci-lint v2.9.0 with .github/.golangci.yml reports 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=on on the packages above.

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-commenter

codecov-commenter commented Oct 23, 2026 •

Copy link
Copy Markdown

⚠️ Please install the 'codecov app svg image' to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 91.00719% with 25 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
internal/upstream/limiter/token_rate.go 82.71% 10 Missing and 4 partials ⚠️
internal/upstream/managed/admission.go 84.37% 2 Missing and 3 partials ⚠️
cmd/mcpproxy/token_cmd.go 95.83% 1 Missing and 1 partial ⚠️
internal/upstream/limiter/errors.go 92.30% 1 Missing and 1 partial ⚠️
internal/runtime/runtime.go 0.00% 0 Missing and 1 partial ⚠️
internal/server/server.go 0.00% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

Only oas/docs.go conflicted (generated); regenerated with make swagger.

@Dumbris Dumbris left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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, and go test -race pass for internal/{upstream,auth,config,httpapi,audit,contracts,observability,runtime,server,storage} and cmd/mcpproxy (the internal/server run used the unit-tests.yml -skip regex). ./scripts/test-api-e2e.sh passes (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 serve with agent_token_rate_limit {per_minute:60, burst:3} and the everything server:
    • the default-budget token gets 3×200, then 429 with Retry-After: 1 and your message text
    • a token created with rate_limit {per_minute:6, burst:2} gets 2×200, then 429 with Retry-After: 10
    • a token with rate_limit {per_minute:0} gets 8×200, and the admin key gets 8×200
    • rate_limit is echoed on create and in GET /api/v1/tokens; per_minute:-1 is a 400 naming the field
    • over MCP (/mcp, bearer token), call_tool_read returns an isError result with the same text
    • code_execution with burst:4: a script making 7 call_tool calls 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: token and token_name on each refusal, one row per refusal
  • 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 the AuthContext is built by one constructor, that rate_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_id or the effective profile, and a refused code_execution sub-call has no parent_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_limited counts in RejectedCount.
  • There is no route to change the budget of an existing token.

I have not merged.

@Dumbris
Dumbris merged commit f13cb22 into smart-mcp-proxy:main Oct 25, 2026
50 checks passed
@Dumbris

Dumbris commented Oct 25, 2026

Copy link
Copy Markdown
Member

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature]: per-token call rate limit / budget for agent tokens

3 participants