feat: Tool Handler Middleware #123
Conversation
WalkthroughThe pull request introduces a middleware mechanism for tool handlers in the MCP server. It updates the documentation to describe two new server options: one for adding tool handler middleware and another for recovery from panics during tool calls. In the implementation, a new middleware type and a middleware slice in the MCPServer have been added, along with functions to register middleware and recovery behavior. The middleware is applied in reverse order, and tests have been added to verify that the server recovers correctly from panics in tool handler execution. Changes
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (2)
🧰 Additional context used🧬 Code Graph Analysis (1)server/server_test.go (3)
🔇 Additional comments (1)
✨ Finishing Touches
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
server/server.go (2)
42-44: Fix the type name to avoid confusion
The type name appears to have a minor misspelling ("Middlware" instead of "Middleware"), which could lead to confusion or inconsistency in the codebase.Apply this diff to correct the spelling:
-type ToolHandlerMiddlware func(ToolHandlerFunc) ToolHandlerFunc +type ToolHandlerMiddleware func(ToolHandlerFunc) ToolHandlerFunc
298-310: Consider capturing stack traces for better debugging
While panic recovery is implemented correctly, adding a stack trace can assist in diagnosing issues that lead to panics.Example diff:
import ( + "runtime/debug" ... ) func WithRecovery() ServerOption { return WithToolHandlerMiddleware(func(next ToolHandlerFunc) ToolHandlerFunc { return func(ctx context.Context, request mcp.CallToolRequest) (result *mcp.CallToolResult, err error) { defer func() { if r := recover(); r != nil { - err = fmt.Errorf("panic recovered in %s tool handler: %v", request.Params.Name, r) + stack := debug.Stack() + err = fmt.Errorf("panic recovered in %s tool handler: %v\nStack trace:\n%s", request.Params.Name, r, stack) } }() return next(ctx, request) } }) }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
README.md(2 hunks)server/server.go(4 hunks)server/server_test.go(1 hunks)
🧰 Additional context used
🧬 Code Definitions (1)
server/server_test.go (3)
mcp/tools.go (3)
CallToolRequest(44-59)CallToolResult(34-41)NewTool(127-142)server/server.go (2)
NewMCPServer(356-382)WithRecovery(299-310)mcp/types.go (2)
JSONRPCError(205-218)INTERNAL_ERROR(226-226)
🔇 Additional comments (5)
server/server.go (2)
142-155: Initialization of new fields looks good
The added fields, including the new RWMutex and the additional slice for tool handler middlewares, appear consistent with the rest of the struct.
288-296: Straightforward middleware registration
The functional option pattern for adding tool handler middleware is clear and aligns well with existing patterns.README.md (2)
125-125: Recovery middleware reference
Including theserver.WithRecovery()call in the server initialization is a clear illustration of the new feature.
527-531: Documentation of tool handler middleware is succinct and clear
The new section effectively explains how to add and use tool handler middleware, including panic recovery.server/server_test.go (1)
1347-1378: Good test coverage for panic recovery
The new test effectively validates that panics in tool handlers are recovered and translated into the correct JSON-RPC error response.
wimspaargaren
force-pushed
the
mcp-go-tool-handler-middleware
branch
from
857743c to
5766c8d
Compare
wimspaargaren
force-pushed
the
mcp-go-tool-handler-middleware
branch
from
5766c8d to
f5a3593
Compare
While trying out mcp host with model
--model ollama:llama3.2:latestand the mcp calculator server example, I noticed the tool call sometimes gets stuck. After some investigation, I figured out that the values of x and y are sometimes received asstringinstead offloat64.Because of this, the
toolHandlerFuncin the server will panic (as the type cast is not being checked in the example) and the client will hang until the context times out.From a client perspective it's hard to figure out what is going on when a toolHandlerFunc panics, as it waits until it receives a message where an ID is set.
I think this issue might be related as well: mark3labs/mcphost#11
In order to address this, this PR adds a
ToolHandlerMiddlwarewhich can be added via theServerOption.This PR includes:
WithToolHandlerMiddlewareoption to add toolHandler middlewaresWithRecoveryoption to let the server recover from panics in toolHandler function callsSummary by CodeRabbit