Feature: Request Hooks #62
Conversation
WalkthroughThis pull request introduces a hooks system to the MCP server to facilitate customized logging, telemetry, and error handling during request processing. It adds comprehensive documentation in the README and new type definitions for MCP methods. The changes span modifications in server handling logic with structured error types, integration of hooks via a new Changes
Possibly related PRs
Suggested reviewers
📜 Recent review detailsConfiguration used: CodeRabbit UI ⛔ Files ignored due to path filters (1)
📒 Files selected for processing (4)
🧰 Additional context used🧬 Code Definitions (4)server/server_test.go (5)
examples/everything/main.go (4)
server/hooks.go (1)
server/server.go (5)
🔇 Additional comments (11)
✨ 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: 3
🧹 Nitpick comments (9)
README.md (1)
519-519: Spelling/grammar nitpick: remove the hyphen in 'improperly-formatted' for standard usage.“improperly” is an adverb, so the hyphen isn't needed. This is a purely stylistic point.
-... ability to count improperly-formatted requests ... +... ability to count improperly formatted requests ...🧰 Tools
🪛 LanguageTool
[uncategorized] ~519-~519: Although a hyphen is possible, it is not necessary in a compound modifier in which the first word is an adverb that ends in ‘ly’.
Context: ...facts, for example the ability to count improperly-formatted requests, or to log the agent identity ...(HYPHENATED_LY_ADVERB_ADJECTIVE)
server/request_handler.go (1)
11-278: Consider modularizing repeated callback pre- and post- calls.Each case block pattern (unmarshal → callback.before* → handle* → error check → callback.after*) is repeated. A small helper could reduce repetition and unify logic, though code generation partially addresses DRY concerns.
server/server_test.go (3)
1066-1080: Check for potential concurrency testing or parallel calls.
While this block correctly sets up callback counters and registers them, consider adding additional tests to ensure callback ordering and counts remain correct when multiple requests are processed in parallel. This will provide stronger guarantees of thread-safety and correctness in concurrent scenarios.
1081-1093: Log message structure refinement.
The callbacks use simplefmt.Printfstatements, which may not be ideal in production-grade logging. Consider integrating a structured logger (e.g.,logrus,zap) to capture callback context and severity levels consistently, facilitating easier analysis and troubleshooting.
1094-1110: Validate method-specific callback coverage.
The callbacks here cover only a few method-specific cases (Ping, ListTools). If you anticipate more methods in the future, consider adding further method-specific callbacks for thorough coverage or a generic fallback for unhandled methods.examples/everything/main.go (1)
31-55: Guard against sensitive data leaks in log messages.
The callback functions utilize%vinfmt.Printf, which may unintentionally log sensitive data included within request parameters (e.g., user PII). In a production environment, consider sanitizing or filtering sensitive details before printing to logs.Also applies to: 63-63
server/server.go (1)
85-85: Handle nil callbacks gracefully.
Although this pointer can remain nil if no callbacks are set, consider adding a quick safety check (e.g.,if s.callbacks != nil) before calling any callback methods to prevent potential nil-pointer dereferences when referencings.callbacks.server/callbacks.go (2)
7-9: Clarify the doc comment for OnBeforeAnyCallbackFunc.
The current wording suggests this callback is invoked "after the request is parsed but before the method is called," which is correct. However, consider prefacing the comment with the function name (e.g., "OnBeforeAnyCallbackFunc is...") to align with standard Go doc conventions, ensuring clarity when generated documentation is viewed.
116-350: Repeated boilerplate for typed request callbacks.
Many methods (e.g.,beforeInitialize,afterInitialize,beforePing,afterPing, etc.) follow a similar pattern. Although this is functional, it can become cumbersome to maintain. Consider a code generation approach or more generic abstraction to reduce duplication, improve readability, and centralize updates for typed callbacks.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (5)
server/internal/gen/README.mdis excluded by!**/gen/**server/internal/gen/callbacks.go.tmplis excluded by!**/gen/**server/internal/gen/data.gois excluded by!**/gen/**server/internal/gen/main.gois excluded by!**/gen/**server/internal/gen/request_handler.go.tmplis excluded by!**/gen/**
📒 Files selected for processing (8)
README.md(1 hunks)examples/everything/main.go(1 hunks)go.mod(1 hunks)mcp/types.go(1 hunks)server/callbacks.go(1 hunks)server/request_handler.go(1 hunks)server/server.go(10 hunks)server/server_test.go(3 hunks)
🧰 Additional context used
🧬 Code Definitions (4)
server/request_handler.go (2)
server/server.go (20) (20)
s(100-105)s(108-116)s(119-123)s(126-150)s(153-179)s(289-302)s(305-318)s(321-329)s(332-334)s(337-352)s(355-360)s(363-375)s(378-385)s(387-439)s(441-447)s(449-468)MCPServer(71-86)requestError(639-643)createErrorResponse(715-732)createResponse(707-713)mcp/types.go (13) (13)
JSONRPCMessage(49-49)MCPMethod(11-11)JSONRPC_VERSION(55-55)INVALID_REQUEST(197-197)MethodInitialize(14-14)InitializeRequest(240-249)InitializeResult(253-267)PingRequest(329-331)EmptyResult(206-206)ListResourcesRequest(374-376)ListResourcesResult(380-383)METHOD_NOT_FOUND(198-198)MethodResourcesRead(18-18)
server/server_test.go (6)
server/request_handler.go (12) (12)
result(59-59)result(78-78)err(18-18)request(58-58)request(77-77)request(96-96)request(121-121)request(146-146)request(171-171)request(196-196)request(221-221)request(246-246)mcp/types.go (4) (4)
Result(147-151)t(29-31)t(33-44)JSONRPCResponse(172-176)mcp/prompts.go (1) (1)
GetPromptResult(32-37)mcp/tools.go (3) (3)
t(82-104)ListToolsResult(19-22)CallToolResult(34-41)server/callbacks.go (1) (1)
Callbacks(46-68)server/server.go (1) (1)
NewMCPServer(260-286)
server/server.go (7)
server/callbacks.go (1) (1)
Callbacks(46-68)mcp/types.go (9) (9)
InitializeResult(253-267)EmptyResult(206-206)ListResourcesRequest(374-376)ListResourcesResult(380-383)ListResourceTemplatesRequest(387-389)ListResourceTemplatesResult(393-396)ReadResourceRequest(400-409)ReadResourceResult(413-416)Params(78-78)client/stdio.go (6) (6)
result(251-251)result(301-301)result(323-323)result(369-369)result(398-398)result(435-435)mcp/utils.go (2) (2)
result(468-468)result(546-546)server/request_handler.go (12) (12)
result(59-59)result(78-78)request(58-58)request(77-77)request(96-96)request(121-121)request(146-146)request(171-171)request(196-196)request(221-221)request(246-246)err(18-18)mcp/prompts.go (2) (2)
ListPromptsResult(13-16)GetPromptRequest(20-28)mcp/tools.go (4) (4)
ListToolsRequest(13-15)ListToolsResult(19-22)CallToolRequest(44-59)CallToolResult(34-41)
server/callbacks.go (1)
mcp/types.go (20) (20)
MCPMethod(11-11)InitializeRequest(240-249)InitializeResult(253-267)PingRequest(329-331)EmptyResult(206-206)ListResourcesRequest(374-376)ListResourcesResult(380-383)ListResourceTemplatesRequest(387-389)ListResourceTemplatesResult(393-396)ReadResourceRequest(400-409)ReadResourceResult(413-416)MethodInitialize(14-14)MethodPing(15-15)MethodResourcesList(16-16)MethodResourcesTemplatesList(17-17)MethodResourcesRead(18-18)MethodPromptsList(19-19)MethodPromptsGet(20-20)MethodToolsList(21-21)MethodToolsCall(22-22)
🪛 LanguageTool
README.md
[uncategorized] ~519-~519: Although a hyphen is possible, it is not necessary in a compound modifier in which the first word is an adverb that ends in ‘ly’.
Context: ...facts, for example the ability to count improperly-formatted requests, or to log the agent identity ...
(HYPHENATED_LY_ADVERB_ADJECTIVE)
🔇 Additional comments (8)
go.mod (1)
8-8:✅ Verification successful
Pinning the dependency version looks good, but verify security and compatibility.
This new direct dependency on
github.com/yosida95/uritemplate/v3is appropriate for the callback feature. However, it's best practice to confirm the chosen version for security advisories and ensure it meets all use cases.Below is a script to check for any known vulnerabilities or advisories regarding this module:
🏁 Script executed:
#!/bin/bash # Description: Check for security advisories for yosida95/uritemplate and verify the latest available version on pkg.go.dev. # Retrieve the module's information from pkg.go.dev curl -s "https://pkg.go.dev/github.com/yosida95/uritemplate/v3?tab=versions" | grep "Version" echo "Check if there are any GH security advisories for uritemplate..." gh api graphql -f query=' { securityVulnerabilities(first: 5, ecosystem: GO, package: "github.com/yosida95/uritemplate/v3") { nodes { advisory { summary severity publishedAt } vulnerableVersionRange firstPatchedVersion { identifier } } } }'Length of output: 8681
Dependency Version Verified: Security & Compatibility Confirmed
The pinned dependency
github.com/yosida95/uritemplate/v3 v3.0.2has been verified via the security advisory check and the package information available on pkg.go.dev. No vulnerabilities were reported by the GH API query, and the module appears to fully support the use case for the callback feature.README.md (1)
512-523: Great addition: “Extras” and “Request Callbacks” documentation.This new section is concise and well-defined, showing how to integrate callbacks into the server. It clearly describes the purpose of callbacks (telemetry, logging, error handling) and how to attach them using
server.WithCallbacks, helping developers adopt the feature without confusion.🧰 Tools
🪛 LanguageTool
[uncategorized] ~519-~519: Although a hyphen is possible, it is not necessary in a compound modifier in which the first word is an adverb that ends in ‘ly’.
Context: ...facts, for example the ability to count improperly-formatted requests, or to log the agent identity ...(HYPHENATED_LY_ADVERB_ADJECTIVE)
mcp/types.go (1)
11-23: String-based method constants enhance clarity.Defining
MCPMethodas astringtype and enumerating all recognized method constants improves type safety and readability. This approach avoids “magic strings” and facilitates future maintenance.server/request_handler.go (3)
11-54: Notification handling is correctly segregated from request/response flow.The check for
baseMessage.ID == niland subsequent unmarshal into a notification struct is a clean, JSON-RPC-compliant approach. Returningnilfor notifications avoids sending unnecessary responses.
56-70: Initialize request handling logic looks consistent.Unmarshaling the request and invoking
handleInitializeunder typed callbacks fits well with the new callback design. The structured error handling withrequestErrorensures clear feedback if the request is invalid.
96-119: Graceful fallback when resources are unavailable.Returning
METHOD_NOT_FOUNDifresourcescapabilities are absent is sensible. This switch-case extension ensures the contract is respected while providing explicit feedback to clients about unsupported features.server/server_test.go (1)
1134-1183: Ensure error callback logic aligns with the success callback flow.
In lines 1177–1179, theafterAnycallback is not called for failed methods, while theonErrorcallback is invoked exactly once for the invalid tool. This approach is valid, but confirm whether you need the afterAny callback to run even for failed requests (e.g., for cleanup or final logging).server/server.go (1)
216-224: Approve addition of WithCallbacks configuration.
This new option neatly integrates callback functionality into the server’s initialization flow, providing a clean, composable design for hooking into request lifecycle events. No issues noted.
tylergannon
force-pushed
the
server-callbacks
branch
5 times, most recently
from
2e9e3d7 to
35b3345
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
README.md (1)
517-520: Remove the hyphen to follow standard usage for adverb + adjective combinations.
Since "improperly" is an adverb ending in -ly, the hyphen before "formatted" is unnecessary:- to count improperly-formatted requests, or to log the agent identity + to count improperly formatted requests, or to log the agent identity🧰 Tools
🪛 LanguageTool
[uncategorized] ~519-~519: Although a hyphen is possible, it is not necessary in a compound modifier in which the first word is an adverb that ends in ‘ly’.
Context: ...facts, for example the ability to count improperly-formatted requests, or to log the agent identity ...(HYPHENATED_LY_ADVERB_ADJECTIVE)
server/server.go (1)
585-586: Remove debug print statement.This
fmt.Printlnstatement appears to be debugging code that should be removed or replaced with proper logging before merging to production.-fmt.Println("request.Params.Name", request.Params.Name)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (5)
server/internal/gen/README.mdis excluded by!**/gen/**server/internal/gen/callbacks.go.tmplis excluded by!**/gen/**server/internal/gen/data.gois excluded by!**/gen/**server/internal/gen/main.gois excluded by!**/gen/**server/internal/gen/request_handler.go.tmplis excluded by!**/gen/**
📒 Files selected for processing (8)
README.md(1 hunks)examples/everything/main.go(1 hunks)go.mod(1 hunks)mcp/types.go(1 hunks)server/callbacks.go(1 hunks)server/request_handler.go(1 hunks)server/server.go(10 hunks)server/server_test.go(3 hunks)
🚧 Files skipped from review as they are similar to previous changes (4)
- go.mod
- mcp/types.go
- server/server_test.go
- server/request_handler.go
🧰 Additional context used
🧬 Code Definitions (1)
server/server.go (3)
server/callbacks.go (1) (1)
Callbacks(48-70)mcp/types.go (12) (12)
InitializeResult(253-267)PingRequest(329-331)EmptyResult(206-206)ListResourcesRequest(374-376)ListResourcesResult(380-383)ListResourceTemplatesRequest(387-389)ListResourceTemplatesResult(393-396)ReadResourceRequest(400-409)ReadResourceResult(413-416)Params(78-78)JSONRPCError(179-192)JSONRPC_VERSION(55-55)server/request_handler.go (12) (12)
result(61-61)result(80-80)request(60-60)request(79-79)request(98-98)request(123-123)request(148-148)request(173-173)request(198-198)request(223-223)request(248-248)err(20-20)
🪛 LanguageTool
README.md
[uncategorized] ~519-~519: Although a hyphen is possible, it is not necessary in a compound modifier in which the first word is an adverb that ends in ‘ly’.
Context: ...facts, for example the ability to count improperly-formatted requests, or to log the agent identity ...
(HYPHENATED_LY_ADVERB_ADJECTIVE)
🔇 Additional comments (7)
server/callbacks.go (2)
48-70: Consider concurrency safety when managing callback slices.
This mirrors a previous comment stating that if callbacks can be registered or modified at runtime, you may face data races when reading them concurrently. Apply a concurrency-safe approach such as using read-write locks or another synchronization mechanism if runtime modifications are anticipated.
88-99: Avoid passing&messageto callbacks expecting a direct value.
This code still passes&messageinstead ofmessage, repeating the existing pointer-mismatch concern. The callback signatures expectmessage any, not*any.examples/everything/main.go (1)
35-55:❓ Verification inconclusive
Verify safe handling of potentially sensitive request data in logs.
You're logging all messages and results throughfmt.Printf. These may contain user data, raising privacy and security concerns. Please confirm whether this logging is appropriate in production environments.Below is a script to locate all
fmt.Printfcalls throughout the repository so you can review them for sensitive data leakage:
🏁 Script executed:
#!/bin/bash rg "fmt\.Printf" -A 2Length of output: 2767
Attention: Potential Sensitive Data Exposure in Logging
The code registers several callbacks that use
fmt.Printfto log dynamic data (e.g.,message,id,result). Our repository-wide search confirms similar patterns across multiple files (notably inexamples/everything/main.goand other modules). These log entries may inadvertently expose sensitive user data if they remain active in production.
- Location to Review:
examples/everything/main.go(lines 35–55)- Other occurrences in files like
client/stdio.goandclient/sse.goPlease verify whether this logging behavior is intended for production use. If these logs are only meant for debugging purposes, consider conditioning them on an environment flag or migrating to a structured logging framework that supports log sanitization and redaction.
server/server.go (4)
85-85: LGTM! The addition of callbacks field to MCPServer struct.The new
callbacksfield allows the server to support the callbacks feature described in the PR objectives. This design enables notifying users before/after requests are executed, facilitating statistics collection and agent identification.
216-223: Well-designed callback configuration function.The
WithCallbacksfunction follows the established pattern of other server configuration options and has clear documentation. It integrates seamlessly with the existing server option architecture.
387-439: Excellent refactoring of handler return types.The handler functions now return specific result types and a
requestErrorinstead of the genericmcp.JSONRPCMessage. This refactoring:
- Improves type safety by using concrete types
- Makes the code more maintainable
- Enables the callbacks feature by allowing typed results to be passed to callbacks
- Creates a consistent error handling pattern across all handlers
This change directly supports the PR objective of modifying request handler prototypes to return their specific result type and error.
Also applies to: 441-447, 449-468, 470-489, 491-548, 555-574, 576-605, 607-637, 664-691
639-662: Well-designed error handling type.The
requestErrortype effectively encapsulates the error information needed for JSON-RPC responses. The implementation includes:
- Fields for ID, error code, and error message
- Implementation of the standard Go
errorinterface- A method to convert to the JSON-RPC error format
This structured approach to error handling is cleaner than the previous direct creation of error responses.
|
Hi @tylergannon I really like this idea and think it's a killer feature to have. With the risk of being nitpicky though, could we rename this to |
|
@ezynda3 Done 👍🏽 Also added additional documentation comments and a more structured error messaging to allow the OnError callback better responses: Add error type handling and documentation for MCP serverThis commit enhances the MCP server's error handling capabilities by:
Error Types and UsageWe've added better support for typed errors that propagate through the request handling chain:
These errors can be interrogated using Code ExamplesThe documentation now includes examples like: hooks.AddOnError(func(id any, method mcp.MCPMethod, message any, err error) {
// Check for specific error types using errors.Is
if errors.Is(err, ErrUnsupported) {
// Handle capability not supported errors
log.Printf("Capability not supported: %v", err)
}
// Use errors.As to get specific error types
var parseErr = &UnparseableMessageError{}
if errors.As(err, &parseErr) {
// Access specific methods/fields of the error type
log.Printf("Failed to parse message for method %s: %v",
parseErr.GetMethod(), parseErr.Unwrap())
// Access the raw message that failed to parse
rawMsg := parseErr.GetMessage()
}
// Check for specific resource/prompt/tool errors
switch {
case errors.Is(err, ErrResourceNotFound):
log.Printf("Resource not found: %v", err)
case errors.Is(err, ErrPromptNotFound):
log.Printf("Prompt not found: %v", err)
case errors.Is(err, ErrToolNotFound):
log.Printf("Tool not found: %v", err)
}
})We've also added examples for testing scenarios: // Create a channel to receive errors for testing
errChan := make(chan error, 1)
// Register hook to capture and inspect errors
hooks := &Hooks{}
hooks.AddOnError(func(id any, method mcp.MCPMethod, message any, err error) {
// For capability-related errors
if errors.Is(err, ErrUnsupported) {
// Handle capability not supported
errChan <- err
return
}
// For parsing errors
var parseErr = &UnparseableMessageError{}
if errors.As(err, &parseErr) {
// Handle unparseable message errors
fmt.Printf("Failed to parse %s request: %v\n",
parseErr.GetMethod(), parseErr.Unwrap())
errChan <- parseErr
return
}
// For resource/prompt/tool not found errors
if errors.Is(err, ErrResourceNotFound) ||
errors.Is(err, ErrPromptNotFound) ||
errors.Is(err, ErrToolNotFound) {
// Handle not found errors
errChan <- err
return
}
// For other errors
errChan <- err
})These improvements make the MCP server more robust and user-friendly by providing clear error patterns and detailed documentation. |
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (8)
server/hooks.go (3)
9-11: Consider using consistent naming for hook function parameters.
Currently, these types use a genericid any, method mcp.MCPMethod, message anysignature for the different hooks. Although flexible, it may become harder to maintain if types vary for different hooks. Consider introducing stronger types or documented interfaces for clarity and safer refactoring.
178-199: Ensure adequate error context is provided.
TheonErrormethod nicely forwards the error to all hooks. However, consider including additional context (e.g., server name, timestamp, user info) for more comprehensive diagnostics if the error is to be stored or escalated.
200-442: Recognize repetitive pattern across before and after methods.
Most of these hooks follow a repetitive pattern of callingbeforeAny/afterAnyplus a local slice iteration. Consider abstracting common functionality into a helper to reduce boilerplate and simplify future maintenance.README.md (1)
516-520: Remove the hyphen from “improperly-formatted requests.”
According to style guides, an adverb ending in “-ly” (e.g. “improperly”) does not typically require a hyphen before an adjective.- to count improperly-formatted requests + to count improperly formatted requests🧰 Tools
🪛 LanguageTool
[uncategorized] ~519-~519: Although a hyphen is possible, it is not necessary in a compound modifier in which the first word is an adverb that ends in ‘ly’.
Context: ...facts, for example the ability to count improperly-formatted requests, or to log the agent identity ...(HYPHENATED_LY_ADVERB_ADJECTIVE)
examples/everything/main.go (2)
33-43: Validate log output vs. production readiness.
These hooks usefmt.Printfto print diagnostic messages. Consider using structured logging or a centralized logger for consistent formatting, severity levels, or log forwarding in production environments.
63-63: Suggest forming dedicated logging hooks in a separate function.
All hooks are defined inline inNewMCPServer. Extracting them into a dedicated helper (e.g.,func newLoggingHooks() *server.Hooks) may improve readability and testability.server/server.go (2)
649-650: Remove debug print statement.The
fmt.Printlnstatement should be removed as it appears to be debugging code that was accidentally left in.- fmt.Println("request.Params.Name", request.Params.Name) -
731-770: Consider updating notification handler for consistency.While most handler functions have been updated to return typed results, the
handleNotificationfunction still returns a generalmcp.JSONRPCMessage. For consistency, consider updating this handler to follow the same pattern as the other handlers.- func (s *MCPServer) handleNotification( - ctx context.Context, - notification mcp.JSONRPCNotification, - ) mcp.JSONRPCMessage { + func (s *MCPServer) handleNotification( + ctx context.Context, + notification mcp.JSONRPCNotification, + ) (*mcp.EmptyResult, *requestError) { s.mu.RLock() handler, ok := s.notificationHandlers[notification.Method] s.mu.RUnlock() if ok { handler(ctx, notification) } - return nil + return &mcp.EmptyResult{}, nil }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (5)
server/internal/gen/README.mdis excluded by!**/gen/**server/internal/gen/data.gois excluded by!**/gen/**server/internal/gen/hooks.go.tmplis excluded by!**/gen/**server/internal/gen/main.gois excluded by!**/gen/**server/internal/gen/request_handler.go.tmplis excluded by!**/gen/**
📒 Files selected for processing (7)
README.md(1 hunks)examples/everything/main.go(1 hunks)mcp/types.go(1 hunks)server/hooks.go(1 hunks)server/request_handler.go(1 hunks)server/server.go(12 hunks)server/server_test.go(16 hunks)
🚧 Files skipped from review as they are similar to previous changes (3)
- server/request_handler.go
- server/server_test.go
- mcp/types.go
🧰 Additional context used
🧬 Code Definitions (2)
server/server.go (4)
mcp/types.go (2) (2)
JSONRPCError(205-218)Params(104-104)server/hooks.go (1) (1)
Hooks(79-101)mcp/prompts.go (1) (1)
GetPromptRequest(20-28)mcp/tools.go (1) (1)
CallToolRequest(44-59)
server/hooks.go (2)
server/request_handler.go (10) (10)
result(61-61)result(80-80)result(99-99)result(124-124)result(149-149)result(174-174)result(199-199)result(224-224)result(249-249)err(20-20)server/server_test.go (2) (2)
err(919-919)_(1141-1141)
🪛 LanguageTool
README.md
[uncategorized] ~519-~519: Although a hyphen is possible, it is not necessary in a compound modifier in which the first word is an adverb that ends in ‘ly’.
Context: ...facts, for example the ability to count improperly-formatted requests, or to log the agent identity ...
(HYPHENATED_LY_ADVERB_ADJECTIVE)
🔇 Additional comments (18)
server/hooks.go (1)
79-101: Validate potential nil usage ofHooks.
TheHooksstruct organizes all hooks in slices, but some methods and internal calls rely onc == nilchecks. Ensure you consistently handle or document how an uninitializedHookspointer is used to avoid unexpectednilpointer dereferences.examples/everything/main.go (2)
44-55: Confirm initialization ordering.
ThebeforeInitializeandafterInitializehooks are defined but ensure that hooks requiring data from initialization run strictly after the server instance is fully set up. If additional dependencies need initialization, consider clarifying the order in documentation.
50-52: Proactively verify potential concurrency.
Hooks likeAddAfterCallToolandAddBeforeCallToolmight be invoked concurrently if multiple requests arrive. Confirm that the code within these hook callbacks is safe for concurrent execution, or add synchronization where needed.Also applies to: 53-55
server/server.go (15)
68-90: Good addition of structured error handling for unparseable messages.The
UnparseableMessageErrortype provides a robust way to handle JSON unmarshalling failures while preserving the original message and method information. This is a good improvement that makes error handling more explicit and structured.
92-121: Improved error handling with proper error wrapping.Updating
requestErrorto include an error field instead of a string is a good practice. The implementation ofError(),ToJSONRPCError(), andUnwrap()methods follows Go's error handling conventions well and enables error inspection viaerrors.Isanderrors.As.
123-128: Well-defined error constants.Defining common error sentinels like
ErrResourceNotFoundmakes error handling more consistent throughout the codebase and enables callers to check for specific error conditions usingerrors.Is.
149-149: Good addition of hooks field to MCPServer struct.The
hooksfield enables the callback functionality described in the PR objectives, which will allow for telemetry, logging, and other cross-cutting concerns.
280-287: Well-implemented hooks configuration function.The
WithHooksfunction follows the established pattern for server options in this codebase, making it consistent with other configuration functions. The comment provides clear documentation on its purpose.
451-503: Handler signature refactored to support hooks.The
handleInitializemethod now returns a typed result (*mcp.InitializeResult) and error (*requestError) instead of a genericmcp.JSONRPCMessage. This change supports the hooks system by providing typed objects that can be passed to callbacks.
505-511: Simplified ping handler with proper return type.The ping handler now returns a properly typed empty result, consistent with the other handler changes.
513-532: List resources handler properly returns typed result.The handler correctly implements the new return signature pattern, returning a typed result that can be used by hooks.
534-553: List resource templates handler follows new pattern.This handler has been updated consistently with the other handlers to return a typed result.
555-612: Read resource handler with improved error handling.The handler now returns specific error information using the
requestErrorstruct, which provides better context about the nature of failures. The error messages include appropriate wrapping of sentinel errors.
619-638: List prompts handler consistently updated.This handler follows the same pattern as the other handlers, returning a typed result.
640-669: Get prompt handler with better error handling.Aside from the debug print statement, this handler has been properly updated to the new pattern with improved error handling.
671-701: List tools handler follows consistent pattern.This handler has been updated to return a typed result in line with the other handlers.
702-729: Tool call handler with proper error propagation.The handler now returns appropriate error information when a tool is not found or when the tool handler encounters an error.
4-14: Added error package import.The addition of the
errorsimport is necessary for the new error handling features like error wrapping and sentinel errors.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
server/hooks.go (2)
9-11: Consider strongly typed parameters
OnBeforeAnyHookFuncusesanyfor the message; while flexible, typed parameters can improve clarity and reduce casting overhead in hook implementations.
79-101: Large struct might benefit from grouping
While auto-generated, theHooksstruct is sizable, with many slices. Consider grouping method-specific hooks in related sub-structures (e.g., resourceHooks, promptHooks) for improved readability.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (1)
server/internal/gen/hooks.go.tmplis excluded by!**/gen/**
📒 Files selected for processing (2)
server/hooks.go(1 hunks)server/server_test.go(12 hunks)
🧰 Additional context used
🧬 Code Definitions (1)
server/server_test.go (6)
mcp/utils.go (16) (16)
result(468-468)result(546-546)_(9-9)_(10-10)_(11-11)_(12-12)_(13-13)_(14-14)_(15-15)_(16-16)_(17-17)_(18-18)_(19-19)_(20-20)_(23-23)_(24-24)mcp/types.go (5) (5)
Result(173-177)MCPMethod(11-11)t(55-57)t(59-70)METHOD_NOT_FOUND(224-224)mcp/prompts.go (1) (1)
GetPromptResult(32-37)server/hooks.go (1) (1)
Hooks(79-101)server/server.go (4) (4)
NewMCPServer(324-350)WithHooks(283-287)ServerOption(29-29)ErrUnsupported(124-124)mcp/tools.go (4) (4)
t(82-104)ListToolsResult(19-22)CallToolRequest(44-59)CallToolResult(34-41)
🔇 Additional comments (16)
server/hooks.go (7)
1-2: Generated disclaimers look good
This is the standard go:generate disclaimer plus reference to the source template.
13-15: Symmetry with OnBeforeAnyHookFunc is consistent
OnAfterAnyHookFuncfollows the same signature pattern, preserving clarity for a post-execution hook. No issues spotted.
17-50: Descriptive documentation for OnErrorHookFunc
The docstring comprehensively demonstrates how to handle various errors and is a strong reference for implementers. Good job ensuring standard usage oferrors.Isanderrors.As.
160-167: beforeAny invocation logic is straightforward
Executes each registered “beforeAny” hook in sequence. No obvious concurrency or error-handling issues.
169-176: afterAny invocation logic is consistent
Executes “afterAny” hooks similarly. Looks good and mirrors the beforeAny approach.
201-443: Method-specific hook implementations align with code generation
All typed hooks (Initialize, Ping, ReadResource, etc.) cleanly wrap around the genericbeforeAny/afterAnycalls. The repetitive pattern is acceptable given code generation.
178-200:✅ Verification successful
Double-check usage of afterAny call for errors
Callingc.afterAny(id, method, message, err)on error paths is a valid design choice but slightly unconventional (passing an error object as "result"). Verify that any “afterAny” hooks correctly handle this scenario.
🏁 Script executed:
#!/bin/bash # Script to search afterAny references and check if they handle error types gracefully. rg -A 5 "afterAny"Length of output: 6286
Confirmed handling of error paths in afterAny hooks
- The tests in
server/server_test.goconfirm that theafterAnyinvocation with an error (i.e., passingerras the result) for error cases works as intended.- All registered hooks are correctly called, and type handling aligns with the design expectations.
- Consider adding a clarifying comment or documentation note to highlight this unconventional, yet validated, design decision.
server/server_test.go (9)
6-6: Minor import change
Importing"errors"is fine. Appears consistent with usage in the rest of the test file.
616-616: Type assertion from interface
Castingresp.Resulttomcp.GetPromptResultis safe given the prior validation. No issues spotted.
646-646: Matching pattern for the missing-argument path
Again casting tomcp.GetPromptResultis consistent with the “get prompts” flow. Looks correct.
668-675: Introducing error hooks in invalid messages test
Storing all errors inerrsviaAddOnErroris a good approach to verify multiple error scenarios. Code is clear and well-scoped for test usage.
680-680: Flexible error validation function
Adding avalidateErrcallback in the test table fosters targeted checks for custom error types. Nice design for extensibility.
696-702: Proper usage of errors.As to detect UnparseableMessageError
The test accurately checks for the parsed method and raw message. Great demonstration of advanced error handling.
713-713: Resetting errs before each test
Ensures no leakage of errors between sub-tests. Good practice for clean test state.
734-812: Handling capabilities absence with hooks
This block tests methods missing their capabilities. It correctly checks forMETHOD_NOT_FOUNDand ensuresErrUnsupportedsurfaces, verifying hook-based error capturing. Nicely done.
1025-1182: Comprehensive “WithHooks” test
Validates that all hooks (BeforeAny, AfterAny, OnError, and method-specific) trigger as intended. This thorough test raises confidence in the new callback system.
MCPServer is upgraded with the capacity to notify before and/or after requests run. Developers can choose Before / After events, and can submit callbacks that are notified during EVERY request lifecycle, or else specific request methods. The specific request method callbacks receive typed message and/or result objects. A special OnError callback is also provided. To accomplish this with a minimal margin for error, a bit of a refactoring was done: - The `MCPServer.HandleMessage` method was moved from server.go into its own file. - The individual request handlers' prototypes were changed. rather than returning `mcp.JSONRPCMessage`, they return a tuple of their specific result type and error. - A code generation schema was introduced, to facilitate adding functionality to each clause in the switch statement within `HandleMessage()`. - Now that HandleMessage() now receives a typed result object and possible error, it is able to pass those results to any configured callbacks. I considered several options including a strategy pattern to make HandleMessage() more DRY, but opted to stick with the minimal reorganization of the code, and to keep the algorithm as flat and obvious as possible. I think the codegen mitigates the repetitiveness of that section of code.
Reinstate test
This commit enhances the MCP server's error handling capabilities by:
1. Improving error propagation through the OnError hook system
2. Adding comprehensive documentation with usage examples
3. Ensuring error types can be properly inspected with Go's standard error handling patterns
We've added better support for typed errors that propagate through the request handling chain:
- `ErrUnsupported`: Used when a capability is not enabled on the server
- `UnparseableMessageError`: Used when parsing a message fails
- `ErrResourceNotFound`: Used when a requested resource doesn't exist
- `ErrPromptNotFound`: Used when a requested prompt doesn't exist
- `ErrToolNotFound`: Used when a requested tool doesn't exist
These errors can be interrogated using `errors.Is` and `errors.As` to provide more targeted error handling.
The documentation now includes examples like:
```go
hooks.AddOnError(func(id any, method mcp.MCPMethod, message any, err error) {
// Check for specific error types using errors.Is
if errors.Is(err, ErrUnsupported) {
// Handle capability not supported errors
log.Printf("Capability not supported: %v", err)
}
// Use errors.As to get specific error types
var parseErr = &UnparseableMessageError{}
if errors.As(err, &parseErr) {
// Access specific methods/fields of the error type
log.Printf("Failed to parse message for method %s: %v",
parseErr.GetMethod(), parseErr.Unwrap())
// Access the raw message that failed to parse
rawMsg := parseErr.GetMessage()
}
// Check for specific resource/prompt/tool errors
switch {
case errors.Is(err, ErrResourceNotFound):
log.Printf("Resource not found: %v", err)
case errors.Is(err, ErrPromptNotFound):
log.Printf("Prompt not found: %v", err)
case errors.Is(err, ErrToolNotFound):
log.Printf("Tool not found: %v", err)
}
})
```
We've also added examples for testing scenarios:
```go
// Create a channel to receive errors for testing
errChan := make(chan error, 1)
// Register hook to capture and inspect errors
hooks := &Hooks{}
hooks.AddOnError(func(id any, method mcp.MCPMethod, message any, err error) {
// For capability-related errors
if errors.Is(err, ErrUnsupported) {
// Handle capability not supported
errChan <- err
return
}
// For parsing errors
var parseErr = &UnparseableMessageError{}
if errors.As(err, &parseErr) {
// Handle unparseable message errors
fmt.Printf("Failed to parse %s request: %v\n",
parseErr.GetMethod(), parseErr.Unwrap())
errChan <- parseErr
return
}
// For resource/prompt/tool not found errors
if errors.Is(err, ErrResourceNotFound) ||
errors.Is(err, ErrPromptNotFound) ||
errors.Is(err, ErrToolNotFound) {
// Handle not found errors
errChan <- err
return
}
// For other errors
errChan <- err
})
```
These improvements make the MCP server more robust and user-friendly by providing clear error patterns and detailed documentation.
Added tests to verify that messages are propagated to the BeforeAny and AfterAny hooks
tylergannon
force-pushed
the
server-callbacks
branch
from
6c6275c to
0f06109
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
README.md (1)
519-519: Remove the hyphen in “improperly-formatted.”According to style conventions, adverbs ending in “-ly” typically don’t require a hyphen before adjectives.
Here’s a proposed change:
- to count improperly-formatted requests + to count improperly formatted requests🧰 Tools
🪛 LanguageTool
[uncategorized] ~519-~519: Although a hyphen is possible, it is not necessary in a compound modifier in which the first word is an adverb that ends in ‘ly’.
Context: ...facts, for example the ability to count improperly-formatted requests, or to log the agent identity ...(HYPHENATED_LY_ADVERB_ADJECTIVE)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (5)
server/internal/gen/README.mdis excluded by!**/gen/**server/internal/gen/data.gois excluded by!**/gen/**server/internal/gen/hooks.go.tmplis excluded by!**/gen/**server/internal/gen/main.gois excluded by!**/gen/**server/internal/gen/request_handler.go.tmplis excluded by!**/gen/**
📒 Files selected for processing (8)
README.md(1 hunks)examples/everything/main.go(1 hunks)go.mod(1 hunks)mcp/types.go(1 hunks)server/hooks.go(1 hunks)server/request_handler.go(1 hunks)server/server.go(12 hunks)server/server_test.go(17 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- go.mod
- mcp/types.go
🧰 Additional context used
🧬 Code Definitions (4)
server/request_handler.go (2)
server/server.go (18) (18)
s(164-169)s(172-180)s(183-187)s(190-214)s(217-243)s(353-366)s(369-382)s(385-393)s(396-398)s(401-416)s(419-424)s(427-439)s(442-449)s(451-503)s(505-511)s(513-532)MCPServer(135-150)requestError(94-98)mcp/types.go (14) (14)
MCPMethod(11-11)PARSE_ERROR(222-222)JSONRPC_VERSION(81-81)INVALID_REQUEST(223-223)JSONRPCNotification(192-195)InitializeRequest(266-275)InitializeResult(279-293)PingRequest(355-357)EmptyResult(232-232)ListResourcesRequest(400-402)ListResourcesResult(406-409)METHOD_NOT_FOUND(224-224)ReadResourceRequest(426-435)ReadResourceResult(439-442)
server/server_test.go (4)
mcp/types.go (4) (4)
Result(173-177)MCPMethod(11-11)t(55-57)t(59-70)mcp/prompts.go (1) (1)
GetPromptResult(32-37)server/hooks.go (1) (1)
Hooks(79-101)server/server.go (8) (8)
NewMCPServer(324-350)WithHooks(283-287)UnparseableMessageError(70-74)ErrToolNotFound(127-127)ErrPromptNotFound(126-126)ErrResourceNotFound(125-125)ServerOption(29-29)ErrUnsupported(124-124)
server/hooks.go (1)
server/request_handler.go (9) (9)
result(61-61)result(80-80)result(99-99)result(124-124)result(149-149)result(174-174)result(199-199)result(224-224)result(249-249)
server/server.go (6)
mcp/types.go (6) (6)
JSONRPCError(205-218)InitializeResult(279-293)EmptyResult(232-232)ListResourcesRequest(400-402)ListResourceTemplatesRequest(413-415)ReadResourceRequest(426-435)server/request_handler.go (12) (12)
err(20-20)result(61-61)result(80-80)request(60-60)request(79-79)request(98-98)request(123-123)request(148-148)request(173-173)request(198-198)request(223-223)request(248-248)server/hooks.go (1) (1)
Hooks(79-101)mcp/utils.go (2) (2)
result(468-468)result(546-546)mcp/prompts.go (2) (2)
ListPromptsResult(13-16)GetPromptRequest(20-28)mcp/tools.go (4) (4)
ListToolsRequest(13-15)ListToolsResult(19-22)CallToolRequest(44-59)CallToolResult(34-41)
🪛 LanguageTool
README.md
[uncategorized] ~519-~519: Although a hyphen is possible, it is not necessary in a compound modifier in which the first word is an adverb that ends in ‘ly’.
Context: ...facts, for example the ability to count improperly-formatted requests, or to log the agent identity ...
(HYPHENATED_LY_ADVERB_ADJECTIVE)
🔇 Additional comments (21)
server/hooks.go (1)
79-101: Consider concurrency safety for appending hooks.Currently, each “Add” method (e.g.,
AddBeforeAny) appends to a slice in theHooksstruct without synchronization. If these methods are called while the server is actively processing requests, data races could occur. If the intended usage is strictly single-threaded (i.e., the hooks are only added before the server starts handling requests), then this is acceptable. Otherwise, it might be prudent to document or ensure concurrency safety.Would you like me to generate a script to check for concurrent modifications of these hooks within the repository?
server/request_handler.go (2)
28-34: Robust handling of JSON parsing errors.Nicely done returning a JSON-RPC parse error response when
json.Unmarshalfails, ensuring the client is informed of malformed messages.
58-278: Commendable typed branching approach.Using a switch statement to parse and route each specific request type ensures strong type safety and clear separation of logic for each method. This approach, coupled with hooks, improves maintainability and extensibility.
server/server_test.go (6)
619-619: Type assertion changed from pointer to value type.The type assertion for the result field has been changed from a pointer type
*mcp.GetPromptResultto a value typemcp.GetPromptResult. This change aligns with the modified function signatures inserver/request_handler.gowhich now return value types instead of pointer types.
646-646: Type assertion changed from pointer to value type.Similar to the previous comment, this is a consistent change to use value types instead of pointer types for the result. This refactoring aligns with the changes made to function signatures throughout the codebase.
667-729: Improved error handling with structured error types.The test now validates that errors are properly propagated through the hooks system and that they match the expected error types. This change is part of the broader efforts to improve error handling by using Go's standard error handling patterns (
errors.Isanderrors.As).
733-867: Enhanced test coverage for error types and hooks.The changes to this test include additional validation of callbacks and error types. This provides better test coverage for the new error handling system, ensuring that specific error types (like
ErrToolNotFound,ErrPromptNotFound, andErrResourceNotFound) are properly propagated through the system.
871-950: Added error string validation for unsupported capabilities.The test now includes an
errStringfield and validation to ensure that the correct error messages are returned when handlers are called without the necessary capabilities. This enhances the test coverage for the error handling system.
1164-1320: Added comprehensive test for hooks system.This new test function validates the hook functionality by:
- Creating counters to verify each type of hook is called the expected number of times
- Setting up various hooks including generic ones (BeforeAny, AfterAny, OnError) and method-specific ones
- Making requests that should trigger different hooks
- Verifying type matching between generic and specific hooks
This comprehensive test ensures the hooks system works as expected and provides good coverage for the new feature.
examples/everything/main.go (2)
33-56: Well-implemented hooks for logging and debugging.The implementation demonstrates how to use the new hooks system for logging server operations. Each hook type (BeforeAny, AfterAny, OnError, etc.) is registered with appropriate logging functions that print relevant information about the request lifecycle.
This provides an excellent real-world example of how the hooks can be used for monitoring, debugging, and telemetry purposes. The logging format is consistent and includes all relevant information (method, ID, message, result, error) for each hook type.
63-63: Hooks configuration added to server options.The
server.WithHooks(hooks)option is correctly added to the server initialization, demonstrating the proper way to configure hooks when creating a new MCP server.server/server.go (10)
68-90: Well-designed error type for unparseable messages.The
UnparseableMessageErrortype provides a structured way to handle JSON parsing errors with access to the raw message, method, and underlying error. The implementation includes:
- A clear error message format
- Support for error unwrapping with
Unwrap()- Accessor methods for the raw message and method
This structure allows for detailed error reporting and inspection, which will be valuable for debugging and error handling.
92-121: Improved request error handling with error unwrapping.The redesigned
requestErrortype now properly implements the error interface with:
- Support for error unwrapping with
Unwrap()- Conversion to a JSON-RPC error response
- A clear error message format
This change enables more robust error handling and allows errors to be inspected with
errors.Isanderrors.As, aligning with Go's standard error handling patterns.
123-128: Standardized error variables for common error conditions.These standardized error variables provide consistent error types that can be checked using
errors.Is(). This approach is a best practice in Go error handling and improves error handling throughout the codebase.
149-149: Added hooks field to MCPServer struct.This field will store the hooks configuration, enabling the hook functionality throughout the server's operation. This is a clean way to add the hooks system to the existing server structure.
280-287: Well-documented WithHooks option function.The
WithHooksfunction is well-documented and follows the established pattern for server options. The function clearly explains its purpose and how it integrates with the request lifecycle.
451-503: Function signature updated to return typed result and error.The handler now returns a specific result type (
*mcp.InitializeResult) and an error (*requestError) instead of a genericmcp.JSONRPCMessage. This change improves type safety and error handling.
505-511: Simplified ping handler with improved return types.The ping handler now returns a more specific type (
*mcp.EmptyResult) and error (*requestError), following the same pattern as other handlers for consistency.
607-612: Improved error return with error wrapping.The error now includes the underlying
ErrResourceNotFounderror usingfmt.Errorfwith the%wverb. This allows for error inspection usingerrors.Is()and provides a more detailed error message.
652-657: Improved error handling for prompt not found.The error now wraps the
ErrPromptNotFounderror, allowing for error inspection usingerrors.Is(). This is consistent with the error handling approach in other handlers.
711-717: Consistent error handling for tool not found.The error handling follows the same pattern as other handlers, wrapping the
ErrToolNotFounderror to allow for error inspection usingerrors.Is().
MCPServer is upgraded with the capacity to notify before and/or after requests run. This will be really helpful:
Developers can choose Before / After events, and can submit callbacks that are notified during EVERY request lifecycle, or else specific request methods. The specific request method callbacks receive typed message and/or result objects.
A special OnError callback is also provided.
To accomplish this with a minimal margin for error, a bit of a refactoring was done:
MCPServer.HandleMessagemethod was moved from server.go into its own file.mcp.JSONRPCMessage, they return a tuple of their specific result type and error.HandleMessage().I considered several options including a strategy pattern to make HandleMessage() more DRY, but opted to stick with the minimal reorganization of the code, and to keep the algorithm as flat and obvious as possible. I think the codegen mitigates the repetitiveness of that section of code.
Summary by CodeRabbit