feat: add Use() method for attaching tool middleware at runtime

[majiayu000]

Description

Fixes #729

Adds Use() to MCPServer so middleware can be attached after server construction. The existing WithToolHandlerMiddleware option only works at init time, which doesn't fit use cases where middleware needs to be registered dynamically (e.g., per-tenant auth, conditional logging).

Two changes:

1. Use(mw ...ToolHandlerMiddleware) on MCPServer — acquires the write lock, appends to toolHandlerMiddlewares, done. Variadic so you can chain multiple in one call.
2. executeRegularToolAsTask was bypassing middleware entirely for tools with TaskSupportOptional. Applied the same middleware chain that handleToolCall already uses, so hybrid-mode task execution isn't a silent gap.

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• MCP spec compatibility implementation
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update
• Code refactoring (no functional changes)
• Performance improvement
• Tests only (no functional changes)
• Other (please describe):

Checklist

• My code follows the code style of this project
• I have performed a self-review of my own code
• I have added tests that prove my fix is effective or that my feature works
• I have updated the documentation accordingly

Additional Information

Tests added:

• TestMCPServer_Use_AppliesMiddlewareToToolCall — middleware runs on a normal tool call
• TestMCPServer_Use_MultipleMiddlewaresExecuteInOrder — ordering matches net/http convention (outermost first)
• TestMCPServer_Use_MiddlewareAppliedToRegularToolAsTask — covers the executeRegularToolAsTask path
• TestMCPServer_Use_ConcurrentSafe — pass -race, no data races

Summary by CodeRabbit

• New Features
  • Added a tool-handler middleware system with ordered chaining, concurrency-safe dynamic registration, and middleware applied to task-routed regular tool execution.
• Documentation
  • New middleware docs with examples, chaining semantics, recovery option, and nav entry for Middleware.
• Tests
  • Tests added for middleware ordering, concurrency safety, and task-path application.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 6313a18c-9993-4314-a6fa-e06bc81532db

📥 Commits

Reviewing files that changed from the base of the PR and between 9079bfe and 2ba4e2f.

📒 Files selected for processing (1)

• www/docs/pages/servers/middleware.mdx

✅ Files skipped from review due to trivial changes (1)

• www/docs/pages/servers/middleware.mdx

---

Walkthrough

Adds a runtime middleware registration method MCPServer.Use(...) and applies registered tool-handler middlewares to regular tools executed via the task path; includes new tests validating middleware ordering, concurrency safety, and task-path application, plus documentation and a sidebar entry.

Changes

Cohort / File(s)	Summary
Server implementation server/server.go	Added func (s *MCPServer) Use(mw ...ToolHandlerMiddleware) (mutex-protected) and updated executeRegularToolAsTask to wrap the tool handler with the registered middleware chain (middlewares applied outermost-first) before invocation.
Tests server/server_test.go	Added tests verifying middleware invocation on tools/call, middleware registration order (outermost-first), concurrent Use() safety, and middleware application when a regular tool is executed via the task path.
Docs & site config www/docs/pages/servers/middleware.mdx, www/vocs.config.ts	New middleware documentation page describing ToolHandlerMiddleware, Use() and constructor-time options, chaining semantics and examples; added sidebar entry for the new docs page.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

• feat: Tool Handler Middleware #123 — Adds server-level ToolHandlerMiddleware and constructor-time middleware options (WithToolHandlerMiddleware, WithRecovery) and applies middleware in tool handling paths.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 28.57% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely summarizes the main change: adding a Use() method for runtime middleware registration on MCPServer.
Description check	✅ Passed	The description fully adheres to the template structure with a clear issue reference (#729), correctly selected feature type, and all checklist items marked as completed with detailed implementation notes.
Linked Issues check	✅ Passed	The PR fully addresses issue #729 by implementing the Use() method for dynamic middleware registration, applying middleware to both direct and task-based tool execution, with comprehensive tests and documentation.
Out of Scope Changes check	✅ Passed	All changes are directly scoped to implementing the middleware feature: the Use() method, middleware application in task execution, comprehensive tests, and documentation—no unrelated modifications detected.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[ezynda3]

@majiayu000 Can you add documentation in the www dir that demonstrates this?

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (3)

www/docs/pages/servers/middleware.mdx (3)

162-170: Consider adding a brief explanation of WithRecovery().

The section mentions server.WithRecovery() but doesn't explain what it does. Adding a brief note would help readers understand its purpose:

-MCP-Go includes built-in recovery middleware that catches panics in tool handlers:
+MCP-Go includes built-in recovery middleware that catches panics in tool handlers and returns an error response instead of crashing the server:

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@www/docs/pages/servers/middleware.mdx` around lines 162 - 170, Add a
one-sentence explanation of server.WithRecovery() next to the example so readers
know its purpose: state that WithRecovery() installs built-in recovery
middleware on the MCP-Go server to catch panics in tool handlers and convert
them into safe error responses (preventing process crashes and logging the
error); place this sentence near the server.NewMCPServer(...) snippet
referencing WithRecovery() so the intent and behavior are clear.

---

103-103: Optional: Consider hyphenating compound adjective.

The static analysis tool suggests using "Rate-Limiting Middleware" (hyphenated) when "Rate Limiting" functions as a compound adjective modifying "Middleware." This follows standard English grammar conventions for compound adjectives.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@www/docs/pages/servers/middleware.mdx` at line 103, Update the header text
"Rate Limiting Middleware" to use a hyphenated compound adjective by changing it
to "Rate-Limiting Middleware" in the document; locate the header string "###
Rate Limiting Middleware" and replace it with "### Rate-Limiting Middleware".

---

47-132: Consider adding import statements and helper definitions for clarity.

The middleware examples are technically correct but could be enhanced:

• The rate limiting example (lines 108-132) uses server.GetSessionID(ctx) and rate.NewLimiter without showing imports (golang.org/x/time/rate)
• The auth example (lines 85-100) references extractToken, validateToken, and userKey without definitions

While these are acceptable for documentation, adding brief comments or a note about required imports would improve clarity for readers copying the examples.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@www/docs/pages/servers/middleware.mdx` around lines 47 - 132, Add missing
import notes and small helper definitions referenced by the examples: mention
the required imports (e.g., "context", "log", "fmt", and
"golang.org/x/time/rate") and clarify helper symbols used in the auth and
rate-limit examples; specifically document or stub extractToken, validateToken,
and userKey referenced by authMiddleware, and call out server.GetSessionID and
rate.NewLimiter used by rateLimiter/newRateLimiter/middleware so readers know to
provide or import those helpers before copying the examples.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@www/docs/pages/servers/middleware.mdx`:
- Around line 147-160: Clarify middleware scope: update the docs to state that
middleware registered via Use() is applied to calls handled by handleCallTool
and executeRegularToolAsTask (i.e., regular tools invoked via the task path) but
is NOT applied to native task tools handled by executeTaskTool or
TaskToolHandlerFunc; explicitly mention this distinction in the paragraph around
executeRegularToolAsTask and replace the blanket "applies to both regular and
task-routed calls" line with a note that only regular tools routed through the
task executor receive middleware, while native task tools do not.

---

Nitpick comments:
In `@www/docs/pages/servers/middleware.mdx`:
- Around line 162-170: Add a one-sentence explanation of server.WithRecovery()
next to the example so readers know its purpose: state that WithRecovery()
installs built-in recovery middleware on the MCP-Go server to catch panics in
tool handlers and convert them into safe error responses (preventing process
crashes and logging the error); place this sentence near the
server.NewMCPServer(...) snippet referencing WithRecovery() so the intent and
behavior are clear.
- Line 103: Update the header text "Rate Limiting Middleware" to use a
hyphenated compound adjective by changing it to "Rate-Limiting Middleware" in
the document; locate the header string "### Rate Limiting Middleware" and
replace it with "### Rate-Limiting Middleware".
- Around line 47-132: Add missing import notes and small helper definitions
referenced by the examples: mention the required imports (e.g., "context",
"log", "fmt", and "golang.org/x/time/rate") and clarify helper symbols used in
the auth and rate-limit examples; specifically document or stub extractToken,
validateToken, and userKey referenced by authMiddleware, and call out
server.GetSessionID and rate.NewLimiter used by
rateLimiter/newRateLimiter/middleware so readers know to provide or import those
helpers before copying the examples.

🪄 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: 16ade903-c597-4753-897a-095a2e6b919f

📥 Commits

Reviewing files that changed from the base of the PR and between 79583e2 and 9079bfe.

📒 Files selected for processing (2)

• www/docs/pages/servers/middleware.mdx
• www/vocs.config.ts

✅ Files skipped from review due to trivial changes (1)

• www/vocs.config.ts