feat: add custom handler support for all MCP server methods #580
Conversation
WalkthroughAdds optional per-method handler hooks to Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested reviewers
Pre-merge checks and finishing touches✅ Passed checks (3 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: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
server/server.go (1)
1026-1036: Resource template path skips resource middlewares (panic recovery, logging, etc.).Direct resource handlers are wrapped, but template handlers are not. This can bypass WithResourceRecovery and other middlewares and lead to process crashes on panics.
Apply this diff to wrap the matched template handler with the same middleware chain:
if matched { - contents, err := matchedHandler(ctx, request) + finalHandler := matchedHandler + s.resourceMiddlewareMu.RLock() + mw := s.resourceHandlerMiddlewares + for i := len(mw) - 1; i >= 0; i-- { + finalHandler = mw[i](finalHandler) + } + s.resourceMiddlewareMu.RUnlock() + contents, err := finalHandler(ctx, request) if err != nil { return nil, &requestError{ id: id, code: mcp.INTERNAL_ERROR, err: err, } } return &mcp.ReadResourceResult{Contents: contents}, nil }
🧹 Nitpick comments (5)
server/server.go (5)
174-186: Group and document the new handler fields; prefer options over mutable fields.
- Define named types for each handler (e.g., InitializeHandlerFunc) and add brief doc comments for public API clarity.
- Consider ServerOptions (e.g., WithInitializeHandler(...)) to set these once at construction time and avoid racy mutation after the server starts handling requests.
666-676: Custom handler errors are always mapped to INTERNAL_ERROR; allow richer error signaling.Right now, custom handlers cannot surface INVALID_PARAMS, NOT_FOUND, etc. Consider supporting a lightweight interface (e.g., type JSONRPCErrorCarrier interface{ ToJSONRPCError() mcp.JSONRPCError }) or exported helpers so handlers can control the code. Fallback to INTERNAL_ERROR when not implemented.
Also applies to: 762-771, 781-790, 872-881, 922-931, 970-979, 1060-1069, 1110-1119, 1151-1160, 1252-1261
1321-1324: Global NotificationHandler shadows per-method handlers; chain both for additive behavior.Calling the global handler and then returning prevents previously registered method-specific handlers from running. Prefer chaining to preserve backward compatibility and allow global observability.
Apply this diff to call both:
- s.NotificationHandler(ctx, notification) - return nil + s.NotificationHandler(ctx, notification) + // fall through to per-method handlers as well
761-771: Guard against nil results from custom PingHandler for consistent JSON-RPC “result” shape.Returning null is valid JSON-RPC, but elsewhere EmptyResult{} is used. Normalize to avoid client surprises.
result, err := s.PingHandler(ctx, request) if err != nil { return nil, &requestError{ id: id, code: mcp.INTERNAL_ERROR, err: err, } } - return result, nil + if result == nil { + result = &mcp.EmptyResult{} + } + return result, nil
780-790: Same here: enforce non-nil result from custom SetLevelHandler.result, err := s.SetLevelHandler(ctx, request) if err != nil { return nil, &requestError{ id: id, code: mcp.INTERNAL_ERROR, err: err, } } - return result, nil + if result == nil { + result = &mcp.EmptyResult{} + } + return result, nil
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
server/server.go(12 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-06-23T11:10:42.948Z
Learnt from: floatingIce91
PR: mark3labs/mcp-go#401
File: server/server.go:1082-1092
Timestamp: 2025-06-23T11:10:42.948Z
Learning: In Go MCP server, ServerTool.Tool field is only used for tool listing and indexing, not for tool execution or middleware. During handleToolCall, only the Handler field is used, so dynamic tools don't need the Tool field populated.
Applied to files:
server/server.go
🧬 Code graph analysis (1)
server/server.go (3)
mcp/types.go (13)
InitializeRequest(414-418)InitializeResult(430-444)PingRequest(508-511)EmptyResult(378-378)ListResourcesRequest(561-564)ListResourcesResult(568-571)ListResourceTemplatesRequest(575-578)ListResourceTemplatesResult(582-585)ReadResourceRequest(589-593)ReadResourceResult(605-608)SetLevelRequest(749-753)JSONRPCNotification(333-336)INTERNAL_ERROR(367-367)mcp/prompts.go (4)
ListPromptsRequest(9-12)ListPromptsResult(16-19)GetPromptRequest(23-27)GetPromptResult(38-43)mcp/tools.go (4)
ListToolsRequest(18-21)ListToolsResult(25-28)CallToolRequest(54-58)CallToolResult(40-51)
🔇 Additional comments (2)
server/server.go (2)
1294-1304: Correct: tool execution uses the Handler, not the Tool metadata.This aligns with prior guidance that ServerTool.Tool is for listing/indexing only. Good consistency.
661-676: Resolved — handleInitialize call sites pass the id.
Invocations found at server/request_handler.go:96, server/sampling_test.go:133,146, and server/session_test.go:1132 — all supply the id argument.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (3)
server/sse.go (3)
476-482: Prefer net/url to construct the query (more robust, safer encoding).The fix works, but using url.Parse + url.Values avoids delimiter edge cases and ensures proper encoding if future ids change format.
Apply:
- if strings.Contains(endpointPath, "?") { - endpointPath += "&" - } else { - endpointPath += "?" - } - return fmt.Sprintf("%ssessionId=%s", endpointPath, sessionID) + if u, err := url.Parse(endpointPath); err == nil { + q := u.Query() + q.Set("sessionId", sessionID) + u.RawQuery = q.Encode() + return u.String() + } + // Fallback to preserve current behavior if parse fails. + sep := "?" + if strings.Contains(endpointPath, "?") { + sep = "&" + } + return endpointPath + sep + "sessionId=" + url.QueryEscape(sessionID)
441-444: Merge client query params via url.Values instead of string concatenation.Prevents malformed URLs and guarantees proper encoding when appending the initial request’s query.
- endpoint := s.GetMessageEndpointForClient(r, sessionID) - if s.appendQueryToMessageEndpoint && len(r.URL.RawQuery) > 0 { - endpoint += "&" + r.URL.RawQuery - } + endpoint := s.GetMessageEndpointForClient(r, sessionID) + if s.appendQueryToMessageEndpoint && r.URL.RawQuery != "" { + if u, err := url.Parse(endpoint); err == nil { + q := u.Query() + for k, vals := range r.URL.Query() { + for _, v := range vals { + q.Add(k, v) + } + } + u.RawQuery = q.Encode() + endpoint = u.String() + } else { + // Fallback to current behavior on parse error. + endpoint += "&" + r.URL.RawQuery + } + }
423-423: Normalize SSE frame formatting (consistent “data: ” and LF).Align with other events for consistency and readability.
- pingMsg := fmt.Sprintf("event: message\ndata:%s\n\n", messageBytes) + pingMsg := fmt.Sprintf("event: message\ndata: %s\n\n", messageBytes)- fmt.Fprintf(w, "event: endpoint\ndata: %s\r\n\r\n", endpoint) + fmt.Fprintf(w, "event: endpoint\ndata: %s\n\n", endpoint)Also applies to: 444-444
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
server/sse.go(1 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-04-28T00:14:49.263Z
Learnt from: robert-jackson-glean
PR: mark3labs/mcp-go#214
File: server/sse.go:0-0
Timestamp: 2025-04-28T00:14:49.263Z
Learning: The SSE server in mcp-go implements path sanitization within the `WithDynamicBasePath` function that ensures the dynamic base path starts with "/" and has no trailing "/" to prevent double slashes in URL construction.
Applied to files:
server/sse.go
- Add custom handler fields to MCPServer struct for all basic MCP methods - Implement custom handler logic in all handle* methods with proper error handling - Support custom handlers for: Initialize, Ping, SetLevel, ListResources, ListResourceTemplates, ReadResource, ListPrompts, GetPrompt, ListTools, CallTool, and Notification methods - Maintain backward compatibility by falling back to default behavior when custom handlers are not set - Enable more flexible server customization and middleware integration
lbbniu
force-pushed
the
custom_handler
branch
from
51ffe02 to
6b9c30f
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
server/server.go(12 hunks)server/sse.go(1 hunks)
🧰 Additional context used
📓 Path-based instructions (1)
**/*.go
📄 CodeRabbit inference engine (AGENTS.md)
**/*.go: Order imports: standard library first, then third-party, then local packages (goimports enforces this)
Follow Go naming conventions: exported identifiers in PascalCase; unexported in camelCase; acronyms uppercase (HTTP, JSON, MCP)
Error handling: return sentinel errors, wrap with fmt.Errorf("context: %w", err), and check with errors.Is/As
Prefer explicit types and strongly-typed structs; avoid using any except where protocol flexibility is required (e.g., Arguments any)
All exported types and functions must have GoDoc comments starting with the identifier name; avoid inline comments unless necessary
Functions that are handlers or long-running must accept context.Context as the first parameter
Ensure thread safety for shared state using sync.Mutex and document thread-safety requirements in comments
For JSON: use json struct tags with omitempty for optional fields; use json.RawMessage for flexible/deferred parsing
Files:
server/server.goserver/sse.go
🧠 Learnings (1)
📚 Learning: 2025-06-23T11:10:42.948Z
Learnt from: floatingIce91
PR: mark3labs/mcp-go#401
File: server/server.go:1082-1092
Timestamp: 2025-06-23T11:10:42.948Z
Learning: In Go MCP server, ServerTool.Tool field is only used for tool listing and indexing, not for tool execution or middleware. During handleToolCall, only the Handler field is used, so dynamic tools don't need the Tool field populated.
Applied to files:
server/server.go
🧬 Code graph analysis (1)
server/server.go (3)
mcp/types.go (13)
InitializeRequest(433-437)InitializeResult(449-463)PingRequest(531-534)EmptyResult(397-397)ListResourcesRequest(584-587)ListResourcesResult(591-594)ListResourceTemplatesRequest(598-601)ListResourceTemplatesResult(605-608)ReadResourceRequest(612-616)ReadResourceResult(628-631)SetLevelRequest(774-778)JSONRPCNotification(335-338)INTERNAL_ERROR(382-382)mcp/prompts.go (4)
ListPromptsRequest(9-12)ListPromptsResult(16-19)GetPromptRequest(23-27)GetPromptResult(38-43)mcp/tools.go (4)
ListToolsRequest(18-21)ListToolsResult(25-28)CallToolRequest(54-58)CallToolResult(40-51)
🔇 Additional comments (5)
server/sse.go (1)
499-504: LGTM! Correct fix for query parameter handling.This change properly handles existing query parameters by conditionally choosing the delimiter. Previously, appending
"?"unconditionally would produce malformed URLs like/path?x=1?sessionId=123. Now it correctly produces/path?x=1&sessionId=123.server/server.go (4)
878-888: Document that custom ListResourcesHandler bypasses session-specific resources.When
ListResourcesHandleris provided, it short-circuits the default logic that merges session-specific resources (lines 897-908). Similarly,ReadResourceHandlerbypasses resource middleware (lines 1036-1042) and template matching. Ensure this behavior is documented in the handler field's GoDoc comment so users understand they must implement these features if needed.
1208-1218: Document that custom tool handlers bypass middleware and filters.
ListToolsHandlerbypasses session-specific tools (lines 1239-1269) and all registered tool filters (lines 1272-1278).CallToolHandlerbypasses tool middleware (lines 1354-1361). Users implementing custom handlers must be aware they're responsible for these features if needed.Based on learnings.
664-667: LGTM! Parameter names clarified.Replacing blank identifiers (
_) with actual parameter names (id,request) improves readability and is necessary for passing these values to custom handlers.
787-797: ****The custom
SetLevelHandlerdesign intentionally allows handlers to bypass default validation—this is part of the optional pluggable callback API architecture. The validation responsibility falls to the custom handler implementer, not the base implementation. Since no custom handler implementations are introduced or used in this PR, no action is required.
|
|
||
| // custom handlers for basic methods | ||
| InitializeHandler func(ctx context.Context, request mcp.InitializeRequest) (*mcp.InitializeResult, error) | ||
| PingHandler func(ctx context.Context, request mcp.PingRequest) (*mcp.EmptyResult, error) | ||
| ListResourcesHandler func(ctx context.Context, request mcp.ListResourcesRequest) (*mcp.ListResourcesResult, error) | ||
| ListResourceTemplatesHandler func(ctx context.Context, request mcp.ListResourceTemplatesRequest) (*mcp.ListResourceTemplatesResult, error) | ||
| ReadResourceHandler func(ctx context.Context, request mcp.ReadResourceRequest) (*mcp.ReadResourceResult, error) | ||
| ListPromptsHandler func(ctx context.Context, request mcp.ListPromptsRequest) (*mcp.ListPromptsResult, error) | ||
| GetPromptHandler func(ctx context.Context, request mcp.GetPromptRequest) (*mcp.GetPromptResult, error) | ||
| ListToolsHandler func(ctx context.Context, request mcp.ListToolsRequest) (*mcp.ListToolsResult, error) | ||
| CallToolHandler func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) | ||
| SetLevelHandler func(ctx context.Context, request mcp.SetLevelRequest) (*mcp.EmptyResult, error) | ||
| NotificationHandler func(ctx context.Context, notification mcp.JSONRPCNotification) | ||
| } |
There was a problem hiding this comment.
Add GoDoc comments for exported handler fields.
All 11 exported handler fields lack documentation. Per coding guidelines, exported identifiers must have GoDoc comments starting with the identifier name. Users need to understand when these handlers are invoked, how errors are handled, and that they bypass default behavior (including middleware, session-specific resources/tools, and filters).
Example documentation pattern:
+ // InitializeHandler, if set, replaces the default initialization logic.
+ // It is invoked during the initialize request and bypasses all default capability
+ // negotiation. Errors are wrapped as INTERNAL_ERROR responses.
InitializeHandler func(ctx context.Context, request mcp.InitializeRequest) (*mcp.InitializeResult, error)
+ // PingHandler, if set, replaces the default ping logic. Errors are wrapped as INTERNAL_ERROR responses.
PingHandler func(ctx context.Context, request mcp.PingRequest) (*mcp.EmptyResult, error)As per coding guidelines.
Committable suggestion skipped: line range outside the PR's diff.
2 participants
Description
This PR adds comprehensive custom handler support to the MCP server, allowing developers to override the default behavior of all basic MCP methods. This enhancement provides greater flexibility for server customization and enables more sophisticated middleware integration patterns.
Type of Change
Checklist
MCP Spec Compliance
Additional Information
Key Changes
MCPServerstruct for all basic MCP methodshandle*methods with proper error handlingSupported Methods
InitializeHandler- Custom initialization logicPingHandler- Custom ping handlingSetLevelHandler- Custom logging level settingListResourcesHandler- Custom resource listingListResourceTemplatesHandler- Custom resource template listingReadResourceHandler- Custom resource readingListPromptsHandler- Custom prompt listingGetPromptHandler- Custom prompt retrievalListToolsHandler- Custom tool listingCallToolHandler- Custom tool executionNotificationHandler- Custom notification handlingUsage Example
Benefits
Summary by CodeRabbit
New Features
Improvements
Bug Fixes