feat(server): add WithRateLimit per-session rate-limit middleware - #910
feat(server): add WithRateLimit per-session rate-limit middleware#910ultramcu wants to merge 3 commits into
Conversation
|
Connected to Huly®: MCP_G-474 |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Run ID: ⛔ Files ignored due to path filters (2)
📒 Files selected for processing (5)
✅ Files skipped from review due to trivial changes (1)
🚧 Files skipped from review as they are similar to previous changes (4)
WalkthroughAdds ChangesRate Limit Middleware
Estimated code review effort: 4 (Complex) | ~45 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
server/ratelimit_extra_test.go (1)
13-22: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMake the invalid-options coverage a named table test.
This already exercises multiple bad configurations, so a
tests := []struct{ name string; opts RateLimitOpts }+t.Runshape would align better with the repo’s test conventions and make it easier to add the missing negative-StaleTTLcase once validation is tightened. As per coding guidelines,**/*_test.go: "implement table-driven tests withtests := []struct{ name, ... }; test files must end in_test.go".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@server/ratelimit_extra_test.go` around lines 13 - 22, Refactor TestRateLimit_PanicsOnBadOpts into a named table-driven test using a tests slice with name and RateLimitOpts fields, and wrap each case in t.Run so the invalid-option coverage matches repo test conventions. Keep the existing panic assertion around WithRateLimit, and include the current bad RPS/Burst/GlobalBurst cases as named entries so it is easy to add the missing negative StaleTTL case later.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@server/ratelimit.go`:
- Around line 76-79: Reject negative StaleTTL values in the rate limit config:
when reaping is enabled, a negative ttl makes the sweep cutoff land in the
future and clears all limiters every tick. Add validation in the option setup
path for WithRateLimit and the reaper setup around ReapInterval/StaleTTL so that
StaleTTL < 0 panics or fails fast, matching the existing bad-options behavior.
Use the StaleTTL field and the reaper/cutoff logic in the rate limiter
implementation to locate the checks.
- Around line 227-230: The global-deny path in the rate limiter is creating
per-key state via st.touch(key), which can populate entries[key] for requests
that never reach the per-key allow path. Update the logic in the rate-limiting
flow around st.touch, Allow, and the per-key limiter creation so global denials
only refresh lastSeen for an already-existing key and do not allocate a new
entry. Keep limiter creation confined to the per-key allow path, and apply the
same adjustment wherever this global-deny handling is duplicated.
- Around line 81-82: The denied-call warning in the rate limiter is missing the
documented session ID, so the `Logger` behavior in `server/ratelimit.go` no
longer matches the comment. Update the logging path in the denied-call handling
used by `RateLimit`/its helper so the Warn entry includes the session ID
alongside the derived key and tool name, or adjust the `Logger` field docs if
that identity is intentionally unavailable after `KeyFunc` transforms it.
---
Nitpick comments:
In `@server/ratelimit_extra_test.go`:
- Around line 13-22: Refactor TestRateLimit_PanicsOnBadOpts into a named
table-driven test using a tests slice with name and RateLimitOpts fields, and
wrap each case in t.Run so the invalid-option coverage matches repo test
conventions. Keep the existing panic assertion around WithRateLimit, and include
the current bad RPS/Burst/GlobalBurst cases as named entries so it is easy to
add the missing negative StaleTTL case later.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 8ea0b267-0324-4f67-bda8-3a2174ac3b6e
⛔ Files ignored due to path filters (2)
go.sumis excluded by!**/*.sumotel/go.sumis excluded by!**/*.sum
📒 Files selected for processing (4)
go.modserver/ratelimit.goserver/ratelimit_extra_test.goserver/ratelimit_test.go
| // StaleTTL is how long a key may be idle before its limiter is evicted by | ||
| // the reaper. 0 selects the default of 10m. Ignored when the reaper is | ||
| // disabled (ReapInterval < 0). | ||
| StaleTTL time.Duration |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Reject negative StaleTTL values.
With ttl < 0, cutoff := now.Add(-ttl) lands in the future, so every sweep deletes every key. That effectively resets exhausted buckets on each reap tick instead of only evicting idle limiters. Since WithRateLimit already fails fast on bad opts, this should also panic when reaping is enabled and StaleTTL < 0.
Suggested fix
func WithRateLimit(opts RateLimitOpts) ServerOption {
if opts.RPS <= 0 {
panic(fmt.Sprintf("server.WithRateLimit: RPS must be > 0, got %v", opts.RPS))
}
if opts.Burst < 1 {
panic(fmt.Sprintf("server.WithRateLimit: Burst must be >= 1, got %d", opts.Burst))
}
if opts.GlobalRPS > 0 && opts.GlobalBurst < 1 {
panic(fmt.Sprintf("server.WithRateLimit: GlobalBurst must be >= 1 when GlobalRPS > 0, got %d", opts.GlobalBurst))
}
+ if opts.ReapInterval >= 0 && opts.StaleTTL < 0 {
+ panic(fmt.Sprintf("server.WithRateLimit: StaleTTL must be >= 0 when reaping is enabled, got %s", opts.StaleTTL))
+ }
store := newRateLimitStore(opts)Also applies to: 99-111, 288-300
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@server/ratelimit.go` around lines 76 - 79, Reject negative StaleTTL values in
the rate limit config: when reaping is enabled, a negative ttl makes the sweep
cutoff land in the future and clears all limiters every tick. Add validation in
the option setup path for WithRateLimit and the reaper setup around
ReapInterval/StaleTTL so that StaleTTL < 0 panics or fails fast, matching the
existing bad-options behavior. Use the StaleTTL field and the reaper/cutoff
logic in the rate limiter implementation to locate the checks.
| // Logger, when non-nil, receives a Warn line per denied call with the key, | ||
| // tool name, and session ID. When nil, denials are not logged. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Denied-call logging is missing the documented session ID.
Right now the warning line only includes the derived key and tool name. When KeyFunc collapses or rewrites session identities, that loses the caller identity the Logger field promises. Either add the session ID attribute here or relax the field docs.
Also applies to: 209-213
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@server/ratelimit.go` around lines 81 - 82, The denied-call warning in the
rate limiter is missing the documented session ID, so the `Logger` behavior in
`server/ratelimit.go` no longer matches the comment. Update the logging path in
the denied-call handling used by `RateLimit`/its helper so the Warn entry
includes the session ID alongside the derived key and tool name, or adjust the
`Logger` field docs if that identity is intentionally unavailable after
`KeyFunc` transforms it.
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (3)
server/ratelimit.go (2)
263-269: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPass
context.Contextinto the long-running reaper.
reapis a long-running goroutine but does not acceptcontext.Contextas its first parameter.Proposed fix
- go st.reap(st.reapInterval, st.staleTTL, st.reaperStop) + go st.reap(context.Background(), st.reapInterval, st.staleTTL, st.reaperStop) } @@ -func (st *rateLimitStore) reap(interval, ttl time.Duration, stop <-chan struct{}) { +func (st *rateLimitStore) reap(ctx context.Context, interval, ttl time.Duration, stop <-chan struct{}) { ticker := time.NewTicker(interval) defer ticker.Stop() for { select { + case <-ctx.Done(): + return case <-stop: returnAs per coding guidelines, “Always accept
context.Contextas the first parameter in handlers and long-running functions.”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@server/ratelimit.go` around lines 263 - 269, The long-running reaper should follow the context-first guideline: update rateLimitStore.reap to accept context.Context as its first parameter, and thread that context through the goroutine startup and any blocking loop/stop checks so the reaper can be canceled cleanly. Locate the call site that starts the reaper and the reap method itself, and replace the custom stop-channel-only flow with context-based cancellation while preserving the existing idle-eviction behavior.Source: Coding guidelines
144-148: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument what the mutex protects.
rateLimitStorehas shared mutable state guarded bymu; add a short comment documenting the protected fields and thread-safety boundary.Proposed fix
- mu sync.Mutex + mu sync.Mutex // guards entries, reaperOn, and reaperStop entries map[string]*rateLimitEntryAs per coding guidelines, “Use
sync.Mutexfor shared state and document thread-safety requirements in comments.”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@server/ratelimit.go` around lines 144 - 148, The rateLimitStore mutex needs a thread-safety comment documenting what shared state it guards. Add a short comment on the mu field in rateLimitStore explaining that it protects entries, reaperOn, and reaperStop, and that now is the injectable clock rather than protected mutable state. Keep the wording aligned with the existing struct fields so the boundary is clear for future changes.Source: Coding guidelines
server/ratelimit_test.go (1)
365-378: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert that reaping restores the full burst, not just one token.
This loop stops after the first allowed call, so an implementation that recreates the limiter with
Burst: 1would still pass. Afterregained, spend the remainingburst-1calls and assert the next call is denied again to lock down the “fresh bucket” contract from the PR objective.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@server/ratelimit_test.go` around lines 365 - 378, The idle-key reaper test currently only verifies that a request becomes allowed again, which can miss a limiter recreated with only one token. Update the test around the regained check in ratelimit_test.go to use the existing rlCall and denied helpers to consume the remaining burst after the first allowed request, then assert the next call is denied so the fresh bucket behavior is enforced.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@server/ratelimit_test.go`:
- Around line 47-71: The helper rlCall in server/ratelimit_test.go should not
use require.Truef for the type assertion because it can stop only the worker
goroutine and leave TestRateLimit_Concurrency hanging. Update rlCall to return
an error when the result is not a *mcp.CallToolResult, and move the assertion to
the caller after wg.Wait() so failures are handled in the main test flow.
In `@server/ratelimit.go`:
- Around line 81-83: Update the Logger field contract in the rate limit
middleware so the godoc matches actual behavior: the Logger comment currently
promises key, tool name, and session ID, but the logging path only includes the
key and tool name. Either revise the comment attached to Logger in ratelimit.go
to remove the session ID mention, or, if intended, update the denied-call
logging code in the rate-limit middleware to actually emit the session ID as
well; keep the contract and implementation consistent in both referenced
locations.
- Around line 55-58: The default deny message documentation in the OnDeny godoc
is out of sync with the actual behavior in the rate limiter. Update the comment
near OnDeny to match the exact string returned by the default deny path in the
rate limit handler, or change the default deny branch in the rate-limit logic to
return the documented message. Use the OnDeny documentation and the default deny
implementation in server/ratelimit.go as the source of truth.
- Around line 102-111: In WithRateLimit, add validation for
RateLimitOpts.StaleTTL so negative values are rejected up front; this
misconfiguration causes sweep to calculate a future cutoff and can evict active
entries repeatedly. Mirror the existing RPS/Burst checks by panicking with a
clear message when opts.StaleTTL < 0, and keep the validation alongside the
other RateLimitOpts guardrails in WithRateLimit.
---
Nitpick comments:
In `@server/ratelimit_test.go`:
- Around line 365-378: The idle-key reaper test currently only verifies that a
request becomes allowed again, which can miss a limiter recreated with only one
token. Update the test around the regained check in ratelimit_test.go to use the
existing rlCall and denied helpers to consume the remaining burst after the
first allowed request, then assert the next call is denied so the fresh bucket
behavior is enforced.
In `@server/ratelimit.go`:
- Around line 263-269: The long-running reaper should follow the context-first
guideline: update rateLimitStore.reap to accept context.Context as its first
parameter, and thread that context through the goroutine startup and any
blocking loop/stop checks so the reaper can be canceled cleanly. Locate the call
site that starts the reaper and the reap method itself, and replace the custom
stop-channel-only flow with context-based cancellation while preserving the
existing idle-eviction behavior.
- Around line 144-148: The rateLimitStore mutex needs a thread-safety comment
documenting what shared state it guards. Add a short comment on the mu field in
rateLimitStore explaining that it protects entries, reaperOn, and reaperStop,
and that now is the injectable clock rather than protected mutable state. Keep
the wording aligned with the existing struct fields so the boundary is clear for
future changes.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 8ea0b267-0324-4f67-bda8-3a2174ac3b6e
⛔ Files ignored due to path filters (2)
go.sumis excluded by!**/*.sumotel/go.sumis excluded by!**/*.sum
📒 Files selected for processing (4)
go.modserver/ratelimit.goserver/ratelimit_extra_test.goserver/ratelimit_test.go
| func rlCall(t *testing.T, s *MCPServer, sessionID, toolName string) (*mcp.CallToolResult, error) { | ||
| t.Helper() | ||
| ctx := s.WithContext(context.Background(), fakeSession{ | ||
| sessionID: sessionID, | ||
| notificationChannel: make(chan mcp.JSONRPCNotification, 1), | ||
| initialized: true, | ||
| }) | ||
| req := mcp.CallToolRequest{ | ||
| Params: mcp.CallToolParams{ | ||
| Name: toolName, | ||
| Arguments: map[string]any{}, | ||
| }, | ||
| } | ||
| res, reqErr := s.handleToolCall(ctx, 1, req) | ||
| if reqErr != nil { | ||
| return nil, reqErr.err | ||
| } | ||
| // handleToolCall returns `any` on success; for tools/call it is a | ||
| // *mcp.CallToolResult. | ||
| if res == nil { | ||
| return nil, nil | ||
| } | ||
| ctr, ok := res.(*mcp.CallToolResult) | ||
| require.Truef(t, ok, "expected *mcp.CallToolResult, got %T", res) | ||
| return ctr, nil |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '1,180p' server/ratelimit_test.go
printf '\n---\n'
sed -n '180,420p' server/ratelimit_test.goRepository: mark3labs/mcp-go
Length of output: 1920
🏁 Script executed:
#!/bin/bash
set -euo pipefail
git ls-files server/ratelimit_test.go
printf '\n---\n'
sed -n '1,120p' server/ratelimit_test.go
printf '\n---\n'
sed -n '320,420p' server/ratelimit_test.goRepository: mark3labs/mcp-go
Length of output: 1920
Avoid require.Truef in rlCall
rlCall is used from worker goroutines in TestRateLimit_Concurrency; require calls FailNow, which can stop only that goroutine and leave the test hanging if the type assertion ever fails. Return an error here and assert after wg.Wait() instead.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@server/ratelimit_test.go` around lines 47 - 71, The helper rlCall in
server/ratelimit_test.go should not use require.Truef for the type assertion
because it can stop only the worker goroutine and leave
TestRateLimit_Concurrency hanging. Update rlCall to return an error when the
result is not a *mcp.CallToolResult, and move the assertion to the caller after
wg.Wait() so failures are handled in the main test flow.
| // OnDeny produces the response returned to the caller when a call is denied | ||
| // (the wrapped handler is not invoked). When nil, the default returns a | ||
| // tool execution error result (mcp.NewToolResultError) carrying the message | ||
| // "rate limit exceeded". |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Keep the default deny message documentation in sync.
Line 57 says the default message is "rate limit exceeded", but Line 196 returns "rate limit exceeded: too many tool calls, please retry later". Please update the godoc to the exact message or make the implementation match the documented contract.
Also applies to: 195-196
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@server/ratelimit.go` around lines 55 - 58, The default deny message
documentation in the OnDeny godoc is out of sync with the actual behavior in the
rate limiter. Update the comment near OnDeny to match the exact string returned
by the default deny path in the rate limit handler, or change the default deny
branch in the rate-limit logic to return the documented message. Use the OnDeny
documentation and the default deny implementation in server/ratelimit.go as the
source of truth.
| // Logger, when non-nil, receives a Warn line per denied call with the key, | ||
| // tool name, and session ID. When nil, denials are not logged. | ||
| Logger *slog.Logger |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Fix the logger field contract.
The godoc promises key, tool name, and session ID, but the middleware logs only the key and tool name. Prefer updating the comment to avoid promising extra session-ID logging unless you intentionally want that additional identifier in logs.
Also applies to: 209-213
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@server/ratelimit.go` around lines 81 - 83, Update the Logger field contract
in the rate limit middleware so the godoc matches actual behavior: the Logger
comment currently promises key, tool name, and session ID, but the logging path
only includes the key and tool name. Either revise the comment attached to
Logger in ratelimit.go to remove the session ID mention, or, if intended, update
the denied-call logging code in the rate-limit middleware to actually emit the
session ID as well; keep the contract and implementation consistent in both
referenced locations.
| func WithRateLimit(opts RateLimitOpts) ServerOption { | ||
| if opts.RPS <= 0 { | ||
| panic(fmt.Sprintf("server.WithRateLimit: RPS must be > 0, got %v", opts.RPS)) | ||
| } | ||
| if opts.Burst < 1 { | ||
| panic(fmt.Sprintf("server.WithRateLimit: Burst must be >= 1, got %d", opts.Burst)) | ||
| } | ||
| if opts.GlobalRPS > 0 && opts.GlobalBurst < 1 { | ||
| panic(fmt.Sprintf("server.WithRateLimit: GlobalBurst must be >= 1 when GlobalRPS > 0, got %d", opts.GlobalBurst)) | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Reject negative StaleTTL values.
With StaleTTL < 0, sweep computes a future cutoff, so active entries are evicted on every reap cycle and their buckets reset. That weakens rate limiting under misconfiguration.
Proposed fix
if opts.GlobalRPS > 0 && opts.GlobalBurst < 1 {
panic(fmt.Sprintf("server.WithRateLimit: GlobalBurst must be >= 1 when GlobalRPS > 0, got %d", opts.GlobalBurst))
}
+if opts.StaleTTL < 0 {
+ panic(fmt.Sprintf("server.WithRateLimit: StaleTTL must be >= 0, got %v", opts.StaleTTL))
+}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| func WithRateLimit(opts RateLimitOpts) ServerOption { | |
| if opts.RPS <= 0 { | |
| panic(fmt.Sprintf("server.WithRateLimit: RPS must be > 0, got %v", opts.RPS)) | |
| } | |
| if opts.Burst < 1 { | |
| panic(fmt.Sprintf("server.WithRateLimit: Burst must be >= 1, got %d", opts.Burst)) | |
| } | |
| if opts.GlobalRPS > 0 && opts.GlobalBurst < 1 { | |
| panic(fmt.Sprintf("server.WithRateLimit: GlobalBurst must be >= 1 when GlobalRPS > 0, got %d", opts.GlobalBurst)) | |
| } | |
| func WithRateLimit(opts RateLimitOpts) ServerOption { | |
| if opts.RPS <= 0 { | |
| panic(fmt.Sprintf("server.WithRateLimit: RPS must be > 0, got %v", opts.RPS)) | |
| } | |
| if opts.Burst < 1 { | |
| panic(fmt.Sprintf("server.WithRateLimit: Burst must be >= 1, got %d", opts.Burst)) | |
| } | |
| if opts.GlobalRPS > 0 && opts.GlobalBurst < 1 { | |
| panic(fmt.Sprintf("server.WithRateLimit: GlobalBurst must be >= 1 when GlobalRPS > 0, got %d", opts.GlobalBurst)) | |
| } | |
| if opts.StaleTTL < 0 { | |
| panic(fmt.Sprintf("server.WithRateLimit: StaleTTL must be >= 0, got %v", opts.StaleTTL)) | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@server/ratelimit.go` around lines 102 - 111, In WithRateLimit, add validation
for RateLimitOpts.StaleTTL so negative values are rejected up front; this
misconfiguration causes sweep to calculate a future cutoff and can evict active
entries repeatedly. Mirror the existing RPS/Burst checks by panicking with a
clear message when opts.StaleTTL < 0, and keep the validation alongside the
other RateLimitOpts guardrails in WithRateLimit.
|
@ultramcu first time reviewing a PR like this, so take it with that in mind. |
|
Thanks for the careful read, @tarunbtw — and good eye. You're right: I've removed it (8463b74) — |
Adds a built-in per-session token-bucket rate limiter as a tool-handler middleware (golang.org/x/time/rate), with an optional global ceiling, a pluggable KeyFunc and OnDeny, and a lazy self-terminating reaper for idle limiters (no background goroutine on an idle server). Closes mark3labs#896.
The reaper already self-terminates when the store drains and restarts on new traffic, and nothing closed reaperStop (the MCPServer has no shutdown lifecycle to wire it to). Remove the speculative stop channel; reap now loops on the ticker directly.
CI lint/test failures were: (1) the otel submodule (replace => ..) needed golang.org/x/time added to its go.mod for the new dependency; (2) the rate-limit tests tripped the repo's modernize + usetesting linters. Fixes: add the otel go.mod require line, use range-over-int, atomic.Int64 for the flagged counters, drop the redundant loop-var copy, and use t.Context() in tests. No behavior change.
8463b74 to
2257324
Compare
|
Rebased onto latest
Verified locally: |
2 participants
Description
Adds a built-in per-session rate-limit tool-handler middleware, as proposed and discussed in #896.
WithRateLimit(opts)installs a token-bucket limiter (viagolang.org/x/time/rate) in front of the tool-handler chain, mirroring the existingWithLogger/WithRecoveryoption pattern:ClientSessionFromContext(a no-session caller falls back to a single shared bucket, so it can't bypass the limit). Override withKeyFuncfor tenant/claim-based keying.GlobalRPS/GlobalBurst) checked before the per-key bucket.OnDenyfor the denied-request response.ReapInterval < 0disables it.Fixes #896.
Type of Change
Checklist
Additional Information
On the denial response / JSON-RPC code. The issue sketched a JSON-RPC
-32029error for denials, but aToolHandlerFunccan't surface a custom JSON-RPC code today — the dispatch path maps any handler/middleware error toINTERNAL_ERROR(-32603), and therequestErrortype that carries a code is unexported. So the defaultOnDenyreturnsmcp.NewToolResultError("rate limit exceeded …")(anIsErrorCallToolResult), consistent with how this package surfaces validation failures. Callers who want different semantics can supply their ownOnDeny. Happy to revisit if a typed protocol-level error becomes returnable from handlers.Tests.
server/ratelimit_test.gocovers deny-after-burst, per-session isolation, the global ceiling, customKeyFunc/OnDeny, concurrency (under-race), and reaper eviction; plus construction-validation and the no-session fallback. The full suite (incl. theotel/submodule, whosego.sumis updated for the new dep) andgo test -racepass.Follow-ups. Logging/metrics hooks can build on the shipped
WithLogger(#892) and the pendingWithMeter(#893) — @tarunbtw kindly offered to pair on that integration.Summary by CodeRabbit