feat: add custom handler support for all MCP server methods #580
Conversation
- 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
WalkthroughAdds optional per-method handler hooks to Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested reviewers
Pre-merge checks (3 passed)✅ Passed checks (3 passed)
Tip 👮 Agentic pre-merge checks are now available in preview!Pro plan users can now enable pre-merge checks in their settings to enforce checklists before merging PRs.
Please see the documentation for more information. Example: reviews:
pre_merge_checks:
custom_checks:
- name: "Undocumented Breaking Changes"
mode: "warning"
instructions: |
Pass/fail criteria: All breaking changes to public APIs, CLI flags, environment variables, configuration keys, database schemas, or HTTP/GraphQL endpoints must be documented in the "Breaking Change" section of the PR description and in CHANGELOG.md. Exclude purely internal or private changes (e.g., code not exported from package entry points or explicitly marked as internal).Please share your feedback with us on this Discord post. ✨ Finishing touches
🧪 Generate 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. 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
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