feat: Add server hooks:OnRequestInitialization #164
Conversation
WalkthroughThis change introduces a new hook mechanism, Changes
Possibly related PRs
Tip ⚡💬 Agentic Chat (Pro Plan, General Availability)
📜 Recent review detailsConfiguration used: CodeRabbit UI ⛔ Files ignored due to path filters (2)
📒 Files selected for processing (4)
🚧 Files skipped from review as they are similar to previous changes (4)
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. 🪧 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: 2
🧹 Nitpick comments (3)
server/hooks.go (2)
56-58: Improve hook documentation for clearer usageThe documentation for
OnBeforeHandleRequestFunccould be more descriptive and include grammatical corrections.-// OnBeforeHandleRequestFunc is a function that called before handle diff request method -// Should any errors arise during func execution, the service will promptly return the corresponding error message. +// OnBeforeHandleRequestFunc is a function that is called before handling a request. +// It can be used for validation, authorization, or other preprocessing tasks. +// If the function returns an error, request processing will be halted and an error +// response will be returned to the client with the error message.
235-246: Implementation looks good with proper error handlingThe
beforeHandleRequestmethod correctly implements the sequential execution of hooks with early return on error.Consider adding a brief comment explaining that the first hook that returns an error will stop the execution chain:
func (c *Hooks) beforeHandleRequest(ctx context.Context, id any, message any) error { if c == nil { return nil } + // Execute hooks sequentially, returning early if any hook returns an error for _, hook := range c.OnBeforeHandleRequest { err := hook(ctx, id, message) if err != nil { return err } } return nil }examples/everything/main.go (1)
47-51: Enhance example with authorization demonstrationThe hook example is simple and clear, but could be improved by showing a real-world use case of authorization verification.
Consider enhancing the example to demonstrate authorization:
hooks.AddBeforeHandleRequest(func(ctx context.Context, id any, message any) error { fmt.Printf("beforeHandleRequest: %v, %v\n", id, message) // authorization verification and other preprocessing tasks are performed. + + // Example of authorization check based on request properties + // If this was a real application, you might check authentication tokens, + // permissions, rate limits, etc. + /* + if shouldRejectRequest(message) { + return fmt.Errorf("request rejected: unauthorized or invalid") + } + */ + return nil })
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
examples/everything/main.go(1 hunks)server/hooks.go(3 hunks)server/request_handler.go(1 hunks)server/server_test.go(3 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (1)
server/request_handler.go (1)
mcp/types.go (1)
PARSE_ERROR(222-222)
🔇 Additional comments (4)
server/hooks.go (2)
92-92: LGTM - Field addition looks goodThe field addition in the Hooks struct is appropriately positioned with other similar hook fields.
231-233: LGTM - AddBeforeHandleRequest method is consistentThe implementation follows the same pattern as other hook registration methods.
server/server_test.go (2)
1208-1209: LGTM - Counter addition for the new hookThe counter variable addition is correctly aligned with other hook counters.
1343-1344: LGTM - Hook invocation count verificationThe assertion for the hook invocation count is correct and ensures the hook is called exactly once for each request.
server/server_test.go
Outdated
| hooks.AddBeforeHandleRequest(func(ctx context.Context, id any, message any) error { | ||
| beforeHandleRequestCount++ | ||
| return nil | ||
| }) |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Add error path test for hook functionality
The hook implementation in the test looks good but only tests the success path where the hook returns nil. Consider adding a test that verifies the error handling path.
Add a test that verifies error handling:
hooks.AddBeforeHandleRequest(func(ctx context.Context, id any, message any) error {
beforeHandleRequestCount++
+
+ // To test error handling, you could return an error for a specific request
+ // For example, when calling a specific method:
+ /*
+ jsonReq := make(map[string]interface{})
+ if err := json.Unmarshal(message.(json.RawMessage), &jsonReq); err == nil {
+ if method, ok := jsonReq["method"].(string); ok && method == "tools/call" {
+ return errors.New("hook rejection test")
+ }
+ }
+ */
+
return nil
})Also, add a test case to verify the error response:
// Test case for hook error
errorMessage := server.HandleMessage(context.Background(), []byte(`{
"jsonrpc": "2.0",
"id": 5,
"method": "tools/call",
"params": {
"name": "test-hook-error"
}
}`))
// Verify error response has the correct error message
errorResp, ok := errorMessage.(mcp.JSONRPCError)
assert.True(t, ok)
assert.Equal(t, mcp.PARSE_ERROR, errorResp.Error.Code)
assert.Contains(t, errorResp.Error.Message, "hook rejection test")
use mcp.INVALID_REQUEST replace mcp.PARSE_ERROR
server/hooks.go
Outdated
| @@ -86,6 +89,7 @@ type Hooks struct { | |||
| OnBeforeAny []BeforeAnyHookFunc | |||
| OnSuccess []OnSuccessHookFunc | |||
| OnError []OnErrorHookFunc | |||
| OnBeforeHandleRequest []OnBeforeHandleRequestFunc | |||
There was a problem hiding this comment.
two things:
- if there is a
beforexx, there will always aafterxx beforexxmethod body will always invokebeforeAny(...)first
There was a problem hiding this comment.
This hook primarily serves as a pre-processing step for requests. The current naming convention "beforeXX" is somewhat ambiguous, so I should consider renaming it.
How about OnRequestInitialization
rename OnBeforeHandleRequest to OnRequestInitialization
AlexNiny
changed the title
| handleErr := s.hooks.onRequestInitialization(ctx, baseMessage.ID, message) | ||
| if handleErr != nil { | ||
| return createErrorResponse( | ||
| baseMessage.ID, | ||
| mcp.INVALID_REQUEST, | ||
| handleErr.Error(), | ||
| ) | ||
| } |
There was a problem hiding this comment.
FYI - this and server/hooks.go are a generated files, you'll have to update the template (see comment at the top)
There was a problem hiding this comment.
Thank you for your reminder. The revised amendments have been submitted.
update tmpl file generate hooks.go/request_handler.go by "go generate" fix example function name error
update test message
…reHandleRequest # Conflicts: # server/hooks.go # server/internal/gen/hooks.go.tmpl # server/internal/gen/request_handler.go.tmpl # server/request_handler.go
* add hooks:beforeHandleRequest * Update request_handler.go use mcp.INVALID_REQUEST replace mcp.PARSE_ERROR * rename hook name rename OnBeforeHandleRequest to OnRequestInitialization * update tmpl file update tmpl file generate hooks.go/request_handler.go by "go generate" fix example function name error * update test update test message --------- Co-authored-by: nine <[email protected]>
New Features
Introducing a new Hook method: AddOnRequestInitialization
Enabling users to perform validation or preprocessing operations before the actual request handling—functioning similarly to middleware. Execution proceeds only upon successful validation; otherwise, an error is returned.
Usage
Summary by CodeRabbit
Summary by CodeRabbit