Feature: Request Hooks

[tylergannon]

MCPServer is upgraded with the capacity to notify before and/or after requests run. This will be really helpful:

• Allows collecting statistics on various types of actions, and this allows to avoid having to add analytics code into every tool definition.
• This opens up the possibility of knowing how often resources or prompts are used.
• Also enables the MCP implementation to report on which agent has invoked it, and possibly behave differently in each case (Cursor vs Claude Desktop, etc).

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.

Summary by CodeRabbit

• New Features
  • Introduced an extensible hooks system for enhanced monitoring, logging, and error reporting during request processing.
  • Improved handling of incoming messages, resulting in a more robust server response mechanism.
• Documentation
  • Added a new “Extras” section to the user guide outlining how to configure and leverage request hooks.
• Tests
  • Expanded test coverage to ensure reliable hook execution and to validate the improved error handling mechanisms.

[coderabbitai[bot]]

Walkthrough

This 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 WithHooks function, and enhanced test cases to validate the new functionality. Additionally, dependency management in go.mod is updated to reflect the direct usage of an external package.

Changes

File(s)	Change Summary
README.md	Added a new "Extras" section with a "Request Hooks" subsection that documents how to create and integrate a Hooks object for observability and logging.
go.mod	Updated dependency handling by changing github.com/yosida95/uritemplate/v3 v3.0.2 from indirect to a direct requirement.
mcp/types.go	Introduced the MCPMethod type and corresponding constants (e.g., MethodInitialize, MethodPing, etc.) to improve type safety in referencing MCP methods.
examples/everything/main.go, server/hooks.go	Added a hooks system with various callback types (OnBeforeAny, OnAfterAny, OnError, etc.) and demonstrated its usage by initializing and integrating a hooks object within server creation.
server/request_handler.go, server/server.go	Refactored JSON-RPC message processing by adding structured error types (e.g., UnparseableMessageError), updating handler function signatures, and integrating hooks via WithHooks.
server/server_test.go	Modified tests to validate the updated error handling, adjusted type assertions (from pointer to value), and added tests (e.g., TestMCPServer_WithHooks) for the new hooks system.

Possibly related PRs

• mark3labs/mcp-go#36: The changes in the main PR, which introduce a hooks system for request processing, are related to the modifications in the retrieved PR that also involve changes to the server/server.go file, specifically in how capabilities are managed and processed.
• mark3labs/mcp-go#38: The changes in the main PR, which introduce a hooks system for request handling, are related to the modifications in the retrieved PR that also involve the HandleMessage method in the server.go file, as both PRs focus on enhancing the server's request processing capabilities.

Suggested reviewers

• ezynda3

---

📜 Recent review details

Configuration used: CodeRabbit UI Review profile: CHILL Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 0f061096eb811aed2ed5cb2a3f3a8a2783c1d453 and 252237ef65321ad9b3af5f7c5ac8cd3ce67b0ac5.

⛔ Files ignored due to path filters (1)

• server/internal/gen/hooks.go.tmpl is excluded by !**/gen/**

📒 Files selected for processing (4)

• examples/everything/main.go (1 hunks)
• server/hooks.go (1 hunks)
• server/server.go (12 hunks)
• server/server_test.go (17 hunks)

🧰 Additional context used

🧬 Code Definitions (4)

server/server_test.go (5)

server/request_handler.go (12) (12)

• result (61-61)
• result (80-80)
• err (20-20)
• 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)

mcp/types.go (6) (6)

• Result (173-177)
• MCPMethod (11-11)
• t (55-57)
• t (59-70)
• JSONRPCResponse (198-202)
• JSONRPCError (205-218)

mcp/prompts.go (1) (1)

• GetPromptResult (32-37)

server/hooks.go (1) (1)

• Hooks (79-101)

server/server.go (3) (3)

• NewMCPServer (324-350)
• UnparseableMessageError (70-74)
• ServerOption (29-29)

examples/everything/main.go (4)

server/hooks.go (1) (1)

• Hooks (79-101)

mcp/types.go (2) (2)

• InitializeRequest (266-275)
• InitializeResult (279-293)

mcp/tools.go (2) (2)

• CallToolRequest (44-59)
• CallToolResult (34-41)

server/server.go (2) (2)

• NewMCPServer (324-350)
• WithHooks (283-287)

server/hooks.go (1)

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.go (5)

mcp/types.go (9) (9)

• InitializeResult (279-293)
• EmptyResult (232-232)
• ListResourcesRequest (400-402)
• ListResourcesResult (406-409)
• ListResourceTemplatesRequest (413-415)
• ListResourceTemplatesResult (419-422)
• ReadResourceRequest (426-435)
• ReadResourceResult (439-442)
• Params (104-104)

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/prompts.go (3) (3)

• ListPromptsResult (13-16)
• GetPromptRequest (20-28)
• GetPromptResult (32-37)

mcp/tools.go (4) (4)

• ListToolsRequest (13-15)
• ListToolsResult (19-22)
• CallToolRequest (44-59)
• CallToolResult (34-41)

🔇 Additional comments (11)

server/hooks.go (1)

1-443: Well-structured hooks system implementation.

This file introduces a complete hooks system for the MCP server with consistent patterns for all request types. The approach of having both general hooks (BeforeAny, OnSuccess, OnError) and method-specific hooks provides excellent flexibility. The error handling documentation is particularly thorough with practical examples.

The hooks implementation pattern is consistent throughout the file, making it maintainable and extensible. Each hook follows the same null-checking pattern and appropriate type signatures.

server/server_test.go (3)

1151-1307: Thorough test coverage for the hooks system.

The test function provides comprehensive verification of the hooks functionality by testing:

1. Multiple hook types (before/after, general and method-specific)
2. Hook execution counts
3. Parameter passing between hooks
4. Error handling through hooks

Consider adding a comment at the beginning of this function explaining its purpose and test strategy for better readability.

---

667-730: Properly modified test cases for type assertion changes.

The test cases have been updated to use value type assertions (mcp.GetPromptResult) instead of pointer type assertions (*mcp.GetPromptResult). This is consistent with the changes in the handler function return types.

---

668-702: Enhanced error validation in test cases.

Test cases now include validation of the specific error types using errors.Is and errors.As, which is a more robust approach to error handling than simple string matching.

examples/everything/main.go (2)

33-55: Great practical example of the hooks system.

This code effectively demonstrates how to use the hooks system for logging request lifecycle events. Both general-purpose hooks and method-specific hooks are showcased.

Consider adding brief comments above each hook registration to explain the specific purpose or benefit of that hook type for better documentation.

---

63-63: Proper integration of hooks with server creation.

The hooks are correctly passed to the server via the WithHooks option, showing the intended usage pattern.

server/server.go (5)

68-121: Improved error handling with structured error types.

The new error types UnparseableMessageError and enhanced requestError provide better error handling capabilities:

1. Both implement proper error interfaces including Unwrap() for error chain inspection
2. UnparseableMessageError includes the raw message and method for better debugging
3. The ToJSONRPCError method provides clean conversion to the protocol's error format

This structured approach allows for more precise error handling using Go's errors.Is and errors.As patterns.

---

123-128: Well-defined standard error values.

The standard error values (ErrUnsupported, ErrResourceNotFound, etc.) enable consistent error checking with errors.Is() across the codebase. This is a good practice that improves error handling consistency.

---

149-149: Added hooks support to MCPServer.

The addition of the hooks field to the MCPServer struct enables the integration of the hooks system, allowing for extending server functionality at various request lifecycle points.

---

280-287: Clean implementation of WithHooks option.

The WithHooks function follows the established pattern for server options, making the hooks system integration consistent with other server configuration options.

---

451-503: Improved handler function signature.

The handler functions now return specific result types and error objects instead of generic JSONRPCMessage objects. This provides better type safety and allows for more precise error handling.

This pattern is consistently applied across all handler functions, which improves code maintainability.

✨ Finishing Touches

• [ ] 📝 Generate Docstrings

---

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Generate unit testing code for this file.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai generate unit testing code for this file.
  • @coderabbitai modularize this function.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read src/utils.ts and generate unit testing code.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
  • @coderabbitai help me debug CodeRabbit configuration file.

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)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai full review to do a full review from scratch and review all the files again.
• @coderabbitai summary to regenerate the summary of the PR.
• @coderabbitai generate docstrings to generate docstrings for this PR.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
• @coderabbitai help to get help.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

• You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
• Please see the configuration documentation for more information.
• If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.

[ezynda3]

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 Hooks instead of Callbacks? Just a preference and think it matches with the naming I've seen in other projects as well.

[tylergannon]

@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 server

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

Error Types and Usage

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.

Code Examples

The 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.