[proposal] Optimize hooks #161
Conversation
WalkthroughThis change refactors the hook system in the MCP server from a function-slice-based approach to an interface-based design. The previous Changes
Possibly related PRs
Suggested reviewers
Tip ⚡💬 Agentic Chat (Pro Plan, General Availability)
✨ 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: 4
🔭 Outside diff range comments (1)
server/server.go (1)
165-394: 💡 Verification agent🧩 Analysis chain
Thread safety concerns with hooks
The pipeline failure indicates race conditions that might be related to concurrent access to the hooks field. Consider making the Hook implementation thread-safe or adding appropriate synchronization.
🏁 Script executed:
#!/bin/bash # Search for concurrent access patterns to hooks field rg -A 2 -B 2 "s\.hooks\." --type goLength of output: 7414
🏁 Script executed:
#!/bin/bash set -e # Locate the Hook interface definition rg -n "type Hook" -A 5 -B 5 --type go # Locate the ChainedHook implementation rg -n "type ChainedHook" -A 10 -B 5 --type goLength of output: 2571
🏁 Script executed:
#!/bin/bash set -e # Inspect CreateHookContext implementation rg -n "func (h \\*ChainedHook) CreateHookContext" -A 10 -B 5 server/hooks.go # Inspect GetHookContext implementation rg -n "func (h \\*ChainedHook) GetHookContext" -A 10 -B 5 server/hooks.goLength of output: 86
🏁 Script executed:
#!/bin/bash set -e # Display full hooks implementation for review sed -n '1,200p' server/hooks.goLength of output: 7484
ChainedHook is not safe for concurrent use
The generated
ChainedHookimplementation mutates shared state without any synchronization. In particular,
CreateHookContext()appends toh.HookContexton the receiver and writes itsHookContextfield- All hook‐dispatch methods (
Before…,After…,OnError, etc.) iterate overh.Hookswithout lockingUnder concurrent request handling this will trigger data races.
To address this, please:
- Add a
sync.RWMutex(orsync.Mutex) toChainedHook- Wrap
CreateHookContext(W lock) and all hook‐dispatch methods (R lock) to protect access toHooksandHookContext- Alternatively, remove the persistent
HookContextfield entirely and treat each call toCreateHookContextas a pure, per-request valueFiles needing updates:
- server/hooks.go (or the
server/internal/gen/hooks.go.tmpltemplate that drives it)🧰 Tools
🪛 GitHub Actions: go
[error] TestSSEServer failed due to race condition detected during test execution
🧹 Nitpick comments (3)
server/request_handler.go (2)
70-71: Inconsistent spacing in hook method callsThe spacing after commas in these hook method calls is inconsistent with the rest of the file.
Apply consistent spacing after commas in parameter lists throughout the file:
- s.hooks.BeforeInitialize(ctx,hookContext, baseMessage.ID, &request) + s.hooks.BeforeInitialize(ctx, hookContext, baseMessage.ID, &request) - s.hooks.OnError(ctx,hookContext, baseMessage.ID, baseMessage.Method, &request, err) + s.hooks.OnError(ctx, hookContext, baseMessage.ID, baseMessage.Method, &request, err) - s.hooks.AfterInitialize(ctx,hookContext, baseMessage.ID, &request, result) + s.hooks.AfterInitialize(ctx, hookContext, baseMessage.ID, &request, result)Also applies to: 74-75, 77-78
🧰 Tools
🪛 GitHub Actions: go
[error] TestSSEServer failed due to race condition detected during test execution
243-243: Missing space after commaSimilar spacing inconsistency here.
- s.hooks.OnError(ctx,hookContext, baseMessage.ID, baseMessage.Method, &request, err) + s.hooks.OnError(ctx, hookContext, baseMessage.ID, baseMessage.Method, &request, err)🧰 Tools
🪛 GitHub Actions: go
[error] TestSSEServer failed due to race condition detected during test execution
examples/everything/main.go (1)
36-43:CreateHookContext/GetHookContextreturnnil– will break type assertionsSeveral parts of the server (and many hooks in the tests) assume that
hookContextis non‑nil and implementsIsHookContext(). Returning
nilhere invites a panic if somebody ever adds logic that touches the
context.A minimal no‑op implementation keeps the example safe:
-func (*exampleHook) CreateHookContext() server.HookContext { - return nil -} - -func (*exampleHook) GetHookContext() server.HookContext { - return nil -} +type emptyCtx struct{} + +func (*emptyCtx) IsHookContext() {} + +func (*exampleHook) CreateHookContext() server.HookContext { + return &emptyCtx{} +} + +func (*exampleHook) GetHookContext() server.HookContext { + return &emptyCtx{} +}This keeps the example functional while still ignoring the context
internally.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (7)
examples/everything/main.go(1 hunks)server/hooks.go(1 hunks)server/request_handler.go(10 hunks)server/server.go(3 hunks)server/server_test.go(9 hunks)server/sse.go(2 hunks)server/stdio.go(1 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (4)
server/sse.go (1)
server/server.go (1)
MCPServer(143-166)
server/server.go (1)
server/hooks.go (2)
Hook(21-53)ChainedHook(55-58)
server/server_test.go (3)
mcp/types.go (17)
MCPMethod(11-11)InitializeRequest(266-275)InitializeResult(279-293)PingRequest(355-357)EmptyResult(232-232)ListResourcesRequest(400-402)ListResourcesResult(406-409)ListResourceTemplatesRequest(413-415)ListResourceTemplatesResult(419-422)ReadResourceRequest(426-435)MethodResourcesRead(32-32)ReadResourceResult(439-442)MethodPromptsGet(40-40)MethodToolsCall(48-48)MethodPing(20-20)MethodInitialize(16-16)MethodToolsList(44-44)server/hooks.go (1)
HookContext(11-13)server/server.go (2)
ClientSession(54-63)ErrUnsupported(132-132)
server/hooks.go (4)
server/server.go (1)
ClientSession(54-63)mcp/types.go (11)
MCPMethod(11-11)InitializeRequest(266-275)InitializeResult(279-293)PingRequest(355-357)EmptyResult(232-232)ListResourcesRequest(400-402)ListResourcesResult(406-409)ListResourceTemplatesRequest(413-415)ListResourceTemplatesResult(419-422)ReadResourceRequest(426-435)ReadResourceResult(439-442)mcp/prompts.go (4)
ListPromptsRequest(7-9)ListPromptsResult(13-16)GetPromptRequest(20-28)GetPromptResult(32-37)mcp/tools.go (4)
ListToolsRequest(13-15)ListToolsResult(19-22)CallToolRequest(44-59)CallToolResult(34-41)
🪛 GitHub Actions: go
server/stdio.go
[error] TestSSEServer failed due to race condition detected during test execution
server/sse.go
[warning] 359-466: Data race detected in SSEServer methods handleMessage and ServeHTTP during concurrent test execution
[error] TestSSEServer failed due to race condition detected during test execution
server/request_handler.go
[warning] 58-58: Data race detected in HandleMessage method during concurrent test execution
[error] TestSSEServer failed due to race condition detected during test execution
server/server.go
[error] TestSSEServer failed due to race condition detected during test execution
server/server_test.go
[error] TestSSEServer failed due to race condition detected during test execution
server/hooks.go
[warning] 74-74: Data race detected in CreateHookContext method during concurrent test execution
[error] TestSSEServer failed due to race condition detected during test execution
🔇 Additional comments (8)
server/sse.go (3)
56-64: Formatting improvements for better code readabilityThe field alignment in the SSEServer struct declaration improves readability and code organization.
🧰 Tools
🪛 GitHub Actions: go
[error] TestSSEServer failed due to race condition detected during test execution
161-166: Consistent field alignment in struct initializationThe vertical alignment of field initializations in NewSSEServer improves code readability and maintainability.
🧰 Tools
🪛 GitHub Actions: go
[error] TestSSEServer failed due to race condition detected during test execution
56-64:❌ Incorrect review comment
Address race conditions in SSEServer
The pipeline failure indicates data races in the SSEServer methods handleMessage and ServeHTTP during concurrent test execution. This isn't directly related to your formatting changes, but should be addressed.
Also applies to: 161-166
🏁 Script executed:
#!/bin/bash # Search for potential concurrent access points in SSEServer methods rg -A 5 -B 5 "func \(s \*SSEServer\) (handleMessage|ServeHTTP)" server/sse.goLength of output: 797
🏁 Script executed:
#!/bin/bash # Extract the first 300 lines of server/sse.go for review sed -n '1,300p' server/sse.go # List all occurrences of "sessions" to see how the sync.Map is accessed rg -n "sessions" server/sse.goLength of output: 8578
🏁 Script executed:
#!/bin/bash # Extract the full handleMessage implementation for review rg -n -A50 -B5 "func (s \*SSEServer\) handleMessage" server/sse.go # Extract the full ServeHTTP implementation for review rg -n -A50 -B5 "func (s \*SSEServer\) ServeHTTP" server/sse.goLength of output: 313
🏁 Script executed:
#!/bin/bash # Extract handleMessage implementation with fixed-string search rg -F -C20 "func (s *SSEServer) handleMessage" server/sse.go # Extract ServeHTTP implementation with fixed-string search rg -F -C20 "func (s *SSEServer) ServeHTTP" server/sse.goLength of output: 2313
handleMessage and ServeHTTP appear concurrency‑safe
Both methods only read from thesync.Mapand immutable server configuration fields—there are no unsynchronized writes or shared state mutations inhandleMessage(lines 248–273) orServeHTTP(lines 335–345). The use of Go’ssync.Mapguarantees thread‑safe loads and stores, and all other fields accessed are never mutated after server construction.Please share the specific race detector output or test scenario that failed so we can pinpoint the true source of any data race.
Likely an incorrect or invalid review comment.
🧰 Tools
🪛 GitHub Actions: go
[error] TestSSEServer failed due to race condition detected during test execution
server/stdio.go (1)
134-134: Helpful comment clarifying blocking behaviorAdding this comment clarifies that the readNextLine method blocks until it receives input, which improves code understanding for future developers.
🧰 Tools
🪛 GitHub Actions: go
[error] TestSSEServer failed due to race condition detected during test execution
server/request_handler.go (1)
58-272: Transition to context-based hook architectureThe addition of hookContext to all hook method calls properly implements the interface-based design described in the PR summary. This improves extensibility and makes the hooks system more flexible.
🧰 Tools
🪛 GitHub Actions: go
[warning] 58-58: Data race detected in HandleMessage method during concurrent test execution
[error] TestSSEServer failed due to race condition detected during test execution
server/server.go (3)
165-165: Updated hook field type to use interfaceChanging the hooks field from a concrete type to an interface aligns with the refactoring to an interface-based design, allowing for more flexible hook implementations.
🧰 Tools
🪛 GitHub Actions: go
[error] TestSSEServer failed due to race condition detected during test execution
334-334: Updated WithHooks function signatureThe signature change matches the new hook architecture, accepting the Hook interface instead of a concrete Hooks struct.
🧰 Tools
🪛 GitHub Actions: go
[error] TestSSEServer failed due to race condition detected during test execution
394-394: Default initialization with ChainedHookInitializing with ChainedHook provides a default implementation that's ready to use, preventing potential nil pointer exceptions.
🧰 Tools
🪛 GitHub Actions: go
[error] TestSSEServer failed due to race condition detected during test execution
| @@ -55,6 +55,7 @@ func (s *MCPServer) HandleMessage( | |||
| return nil // Return nil for notifications | |||
| } | |||
|
|
|||
| hookContext := s.hooks.CreateHookContext() | |||
There was a problem hiding this comment.
Potential race condition in hookContext creation
The pipeline failure indicates a data race on this line during concurrent test execution. This suggests that multiple goroutines may be accessing s.hooks concurrently without proper synchronization.
Consider adding synchronization to prevent concurrent access to the hooks field or ensuring that the Hook implementation is thread-safe.
🧰 Tools
🪛 GitHub Actions: go
[warning] 58-58: Data race detected in HandleMessage method during concurrent test execution
[error] TestSSEServer failed due to race condition detected during test execution
| func (h *ChainedHook) BeforeAny(ctx context.Context, hookContext HookContext, id any, method mcp.MCPMethod, message any) { | ||
| if h == nil { | ||
| return | ||
| } | ||
| for _, hook := range c.OnError { | ||
| hook(ctx, id, method, message, err) | ||
| } | ||
| } | ||
|
|
||
| func (c *Hooks) AddOnRegisterSession(hook OnRegisterSessionHookFunc) { | ||
| c.OnRegisterSession = append(c.OnRegisterSession, hook) | ||
| for _, hook := range h.Hooks { | ||
| hook.BeforeAny(ctx, hookContext, id, method, message) | ||
| } | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Each inner hook receives the aggregated context instead of its own
BeforeAny (and every other passthrough method) forwards the same
hookContext value to every contained hook:
for _, hook := range h.Hooks {
hook.BeforeAny(ctx, hookContext, …)
}When hookContext is a *ChainedHookContext, a concrete hook that
expects its own subtype will panic on a type‑assertion – it never
receives the HookContext instance it created in
CreateHookContext().
A safe pattern is:
if chCtx, ok := hookContext.(*ChainedHookContext); ok {
for i, hook := range h.Hooks {
- hook.BeforeAny(ctx, hookContext, id, method, message)
+ hook.BeforeAny(ctx, chCtx.HookContexts[i], id, method, message)
}
return
}
for _, hook := range h.Hooks {
hook.BeforeAny(ctx, hookContext, id, method, message) // fallback
}(The same helper can be applied to all other forwarders.)
This guarantees that every hook only ever sees its own context
object.
🧰 Tools
🪛 GitHub Actions: go
[error] TestSSEServer failed due to race condition detected during test execution
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (1)
server/server_test.go (1)
1345-1346:⚠️ Potential issueCounters use dangling pointers – undefined behaviour & data‑race
This code pattern appears in multiple counter variables and creates dangling pointers, leading to undefined behavior and potential data races.
For all counters in the
globalHookContext, store the counter as a value rather than a pointer to avoid this issue:type globalHookContext struct { - BeforeAnyCount *int - OnSuccessCount *int - OnErrorCount *int - BeforePingCount *int - AfterPingCount *int - BeforeToolsCount *int - AfterToolsCount *int + BeforeAnyCount int + OnSuccessCount int + OnErrorCount int + BeforePingCount int + AfterPingCount int + BeforeToolsCount int + AfterToolsCount int // other fields... } // Then in the methods, update the counter directly func (*globalHook) BeforeAny(ctx context.Context, hookContext HookContext, id any, method mcp.MCPMethod, message any) { - newCount := (*hookContext.(*globalHookContext).BeforeAnyCount) + 1 - hookContext.(*globalHookContext).BeforeAnyCount = &(newCount) + hookContext.(*globalHookContext).BeforeAnyCount++ // rest of the method... }This change must be applied to all counter increment operations in the file.
Also applies to: 1354-1355, 1366-1367, 1381-1382, 1390-1391, 1432-1433, 1439-1440
🧰 Tools
🪛 GitHub Actions: go
[error] TestSSEServer failed due to race condition detected during execution of test 'Can_handle_multiple_sessions'.
🧹 Nitpick comments (3)
server/server_test.go (3)
751-751: Minor formatting issue - remove unnecessary parentheses.There's an extra pair of parentheses around
(0)that serves no purpose and should be removed.- if hookContext.(*exampleHookContext).AfterResults == nil { - hookContext.(*exampleHookContext).AfterResults = make([]afterResult, (0)) - } + if hookContext.(*exampleHookContext).AfterResults == nil { + hookContext.(*exampleHookContext).AfterResults = make([]afterResult, 0) + }🧰 Tools
🪛 GitHub Actions: go
[error] TestSSEServer failed due to race condition detected during execution of test 'Can_handle_multiple_sessions'.
762-763: Remove debug print statements.These debug print statements should be removed as they will clutter test output and aren't necessary for the test functionality.
-func (*exampleHook) BeforeInitialize(ctx context.Context, hookContext HookContext, id any, message *mcp.InitializeRequest) { - fmt.Printf("beforeInitialize: %v, %v\n", id, message) +func (*exampleHook) BeforeInitialize(ctx context.Context, hookContext HookContext, id any, message *mcp.InitializeRequest) { } -func (*exampleHook) AfterInitialize(ctx context.Context, hookContext HookContext, id any, message *mcp.InitializeRequest, result *mcp.InitializeResult) { - fmt.Printf("afterInitialize: %v, %v, %v\n", id, message, result) +func (*exampleHook) AfterInitialize(ctx context.Context, hookContext HookContext, id any, message *mcp.InitializeRequest, result *mcp.InitializeResult) { }And similarly for the BeforeListTools and AfterListTools methods.
Also applies to: 766-767, 810-811, 813-814
🧰 Tools
🪛 GitHub Actions: go
[error] TestSSEServer failed due to race condition detected during execution of test 'Can_handle_multiple_sessions'.
1285-1296: Initialize zero values withvarstatement rather than pointer creation.Creating pointers to integer literals is less idiomatic and causes heap allocations. Instead, use a
varstatement for zero values or define the counters as values rather than pointers.-var zeroValue = 0 -var globalHookCtx = &globalHookContext{ - BeforeAnyMessages: make([]any, 0), - BeforeAnyCount: &zeroValue, - OnSuccessCount: &zeroValue, - OnErrorCount: &zeroValue, - BeforePingCount: &zeroValue, - AfterPingCount: &zeroValue, - BeforeToolsCount: &zeroValue, +var globalHookCtx = &globalHookContext{ + BeforeAnyMessages: make([]any, 0), + BeforeAnyCount: 0, + OnSuccessCount: 0, + OnErrorCount: 0, + BeforePingCount: 0, + AfterPingCount: 0, + BeforeToolsCount: 0,🧰 Tools
🪛 GitHub Actions: go
[error] TestSSEServer failed due to race condition detected during execution of test 'Can_handle_multiple_sessions'.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (7)
examples/everything/main.go(1 hunks)server/hooks.go(1 hunks)server/request_handler.go(10 hunks)server/server.go(3 hunks)server/server_test.go(9 hunks)server/sse.go(2 hunks)server/stdio.go(1 hunks)
✅ Files skipped from review due to trivial changes (1)
- server/sse.go
🚧 Files skipped from review as they are similar to previous changes (5)
- server/stdio.go
- server/request_handler.go
- server/server.go
- examples/everything/main.go
- server/hooks.go
🧰 Additional context used
🧬 Code Graph Analysis (1)
server/server_test.go (5)
mcp/types.go (11)
MCPMethod(11-11)InitializeRequest(266-275)InitializeResult(279-293)PingRequest(355-357)EmptyResult(232-232)MethodResourcesRead(32-32)MethodPromptsGet(40-40)MethodToolsCall(48-48)MethodPing(20-20)MethodInitialize(16-16)MethodToolsList(44-44)server/hooks.go (1)
HookContext(11-13)server/server.go (2)
ClientSession(54-63)ErrUnsupported(132-132)mcp/prompts.go (4)
ListPromptsRequest(7-9)ListPromptsResult(13-16)GetPromptRequest(20-28)GetPromptResult(32-37)mcp/tools.go (4)
ListToolsRequest(13-15)ListToolsResult(19-22)CallToolRequest(44-59)CallToolResult(34-41)
🪛 GitHub Actions: go
server/server_test.go
[error] TestSSEServer failed due to race condition detected during execution of test 'Can_handle_multiple_sessions'.
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's limits when posting large numbers of comments. If you are seeing this consistently it is likely a permissions issue. Please check "Moderation" -> "Code review limits" under your organization settings.
Actionable comments posted: 4
♻️ Duplicate comments (4)
server/server_test.go (4)
724-726:⚠️ Potential issueAvoid storing a shared HookContext directly in the hook.
Storing the
HookContextdirectly in the hook could lead to race conditions if the hook is used concurrently. This pattern was previously identified as problematic in past reviews.type exampleHook struct { - HookContext HookContext } func (e *exampleHook) CreateHookContext() HookContext { hookContext := &exampleHookContext{} hookContext.Errs = make([]error, 0) - e.HookContext = hookContext return hookContext } func (e *exampleHook) GetHookContext() HookContext { - return e.HookContext + // This method should be removed or redesigned to not depend on a stored context + // Consider using a map[string]HookContext to track contexts by ID if needed + return nil }
1286-1297:⚠️ Potential issueCounters use dangling pointers - undefined behavior.
Using pointers to stack variables (
zeroValue) leads to undefined behavior when these variables go out of scope. The pattern appears throughout the code for all counter fields.-var zeroValue = 0 var globalHookCtx = &globalHookContext{ BeforeAnyMessages: make([]any, 0), - BeforeAnyCount: &zeroValue, - OnSuccessCount: &zeroValue, - OnErrorCount: &zeroValue, - BeforePingCount: &zeroValue, - AfterPingCount: &zeroValue, - BeforeToolsCount: &zeroValue, - AfterToolsCount: &zeroValue, + BeforeAnyCount: 0, + OnSuccessCount: 0, + OnErrorCount: 0, + BeforePingCount: 0, + AfterPingCount: 0, + BeforeToolsCount: 0, + AfterToolsCount: 0,
1337-1340:⚠️ Potential issueFix shared hook context to prevent races.
The current implementation returns a single shared context instance for all requests, which can cause race conditions in concurrent tests.
func (*globalHook) CreateHookContext() HookContext { - // All request share one hook context - return globalHookCtx + // Create a new context per request to avoid race conditions + return &globalHookContext{ + BeforeAnyMessages: make([]any, 0), + BeforePingMessages: make([]*mcp.PingRequest, 0), + OnSuccessData: make([]struct { + msg any + res any + }, 0), + AfterPingData: make([]struct { + Msg *mcp.PingRequest + Res *mcp.EmptyResult + }, 0), + } }
1308-1318: 🛠️ Refactor suggestionChange counter fields to values instead of pointers.
The counter fields should be defined as values rather than pointers to avoid dangling pointers and improve performance.
type globalHookContext struct { mu sync.Mutex BeforeAnyMessages []any - BeforeAnyCount *int - OnSuccessCount *int - OnErrorCount *int - BeforePingCount *int - AfterPingCount *int - BeforeToolsCount *int - AfterToolsCount *int + BeforeAnyCount int + OnSuccessCount int + OnErrorCount int + BeforePingCount int + AfterPingCount int + BeforeToolsCount int + AfterToolsCount int
🛑 Comments failed to post (4)
server/server_test.go (4)
1541-1544: 🛠️ Refactor suggestion
Update assertion statements to match counter field type changes.
When changing counter fields from pointers to direct values, the assertions need to be updated as well.
-assert.Equal(t, 1, *hooks.GetHookContext().(*globalHookContext).BeforePingCount, "beforePing should be called once") -assert.Equal(t, 1, *hooks.GetHookContext().(*globalHookContext).AfterPingCount, "afterPing should be called once") -assert.Equal(t, 1, *hooks.GetHookContext().(*globalHookContext).BeforeToolsCount, "beforeListTools should be called once") -assert.Equal(t, 1, *hooks.GetHookContext().(*globalHookContext).AfterToolsCount, "afterListTools should be called once") +assert.Equal(t, 1, hooks.GetHookContext().(*globalHookContext).BeforePingCount, "beforePing should be called once") +assert.Equal(t, 1, hooks.GetHookContext().(*globalHookContext).AfterPingCount, "afterPing should be called once") +assert.Equal(t, 1, hooks.GetHookContext().(*globalHookContext).BeforeToolsCount, "beforeListTools should be called once") +assert.Equal(t, 1, hooks.GetHookContext().(*globalHookContext).AfterToolsCount, "afterListTools should be called once")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.assert.Equal(t, 1, hooks.GetHookContext().(*globalHookContext).BeforePingCount, "beforePing should be called once") assert.Equal(t, 1, hooks.GetHookContext().(*globalHookContext).AfterPingCount, "afterPing should be called once") assert.Equal(t, 1, hooks.GetHookContext().(*globalHookContext).BeforeToolsCount, "beforeListTools should be called once") assert.Equal(t, 1, hooks.GetHookContext().(*globalHookContext).AfterToolsCount, "afterListTools should be called once")
1548-1553: 🛠️ Refactor suggestion
Update remaining assertion statements for counter type changes.
Further assertions that need to be updated after changing counter fields to values.
-assert.Equal(t, 4, *hooks.GetHookContext().(*globalHookContext).BeforeAnyCount, "beforeAny should be called for each method") -assert.Equal(t, 3, *hooks.GetHookContext().(*globalHookContext).OnSuccessCount, "onSuccess should be called after all successful invocations") -assert.Equal(t, 1, *hooks.GetHookContext().(*globalHookContext).OnErrorCount, "onError should be called once") +assert.Equal(t, 4, hooks.GetHookContext().(*globalHookContext).BeforeAnyCount, "beforeAny should be called for each method") +assert.Equal(t, 3, hooks.GetHookContext().(*globalHookContext).OnSuccessCount, "onSuccess should be called after all successful invocations") +assert.Equal(t, 1, hooks.GetHookContext().(*globalHookContext).OnErrorCount, "onError should be called once")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.assert.Equal(t, 4, hooks.GetHookContext().(*globalHookContext).BeforeAnyCount, "beforeAny should be called for each method") // onSuccess is called for all 3 success methods (initialize, ping, tools/list) assert.Equal(t, 3, hooks.GetHookContext().(*globalHookContext).OnSuccessCount, "onSuccess should be called after all successful invocations") // Error hook should be called once for the failed tools/call assert.Equal(t, 1, hooks.GetHookContext().(*globalHookContext).OnErrorCount, "onError should be called once")
1387-1397: 🛠️ Refactor suggestion
Fix pointer arithmetic in counters.
Replace pointer arithmetic with direct increment of the counter value.
func (e *globalHook) BeforePing(ctx context.Context, hookContext HookContext, id any, message *mcp.PingRequest) { hookContext.(*globalHookContext).mu.Lock() defer hookContext.(*globalHookContext).mu.Unlock() e.BeforeAny(ctx, hookContext, id, mcp.MethodPing, message) - newCount := (*hookContext.(*globalHookContext).BeforePingCount) + 1 - hookContext.(*globalHookContext).BeforePingCount = &(newCount) + hookContext.(*globalHookContext).BeforePingCount++ hookContext.(*globalHookContext).BeforePingMessages = append(hookContext.(*globalHookContext).BeforePingMessages, message) }Committable suggestion skipped: line range outside the PR's diff.
1346-1354:
⚠️ Potential issueFix mutex unlock causing a fatal error.
The code is unlocking the mutex in
BeforeAnywhen the method ismcp.MethodPing, but this mutex is already scheduled to be unlocked by thedeferstatement in the calling functionBeforePing. This causes an "unlock of unlocked mutex" error.func (*globalHook) BeforeAny(ctx context.Context, hookContext HookContext, id any, method mcp.MCPMethod, message any) { - newCount := (*hookContext.(*globalHookContext).BeforeAnyCount) + 1 - hookContext.(*globalHookContext).BeforeAnyCount = &(newCount) + hookContext.(*globalHookContext).BeforeAnyCount++ // Only collect ping messages for our test if method == mcp.MethodPing { hookContext.(*globalHookContext).BeforeAnyMessages = append(hookContext.(*globalHookContext).BeforeAnyMessages, message) - hookContext.(*globalHookContext).mu.Unlock() } }Committable suggestion skipped: line range outside the PR's diff.
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (4)
server/server_test.go (4)
724-726:⚠️ Potential issueAvoid storing a shared HookContext directly in the hook.
Storing the
HookContextdirectly in the hook could lead to race conditions if the hook is used concurrently. This is especially problematic in test scenarios with parallel test execution.type exampleHook struct { - HookContext HookContext } func (e *exampleHook) CreateHookContext() HookContext { hookContext := &exampleHookContext{} hookContext.Errs = make([]error, 0) - e.HookContext = hookContext return hookContext } func (e *exampleHook) GetHookContext() HookContext { - return e.HookContext + // This method should be removed or redesigned to not depend on stored context + // Consider using a map[string]HookContext with proper synchronization to track contexts by ID }
1286-1306:⚠️ Potential issueUsing pointer fields for counters creates dangling pointers and race conditions.
Storing pointers to stack variables like
zeroValueis unsafe as the variable goes out of scope, making the pointer dangling. Additionally, incrementing and reassigning these pointers creates unnecessary heap allocations.// All request share one hook context -var zeroValue = 0 var globalHookCtx = &globalHookContext{ BeforeAnyMessages: make([]any, 0), - BeforeAnyCount: &zeroValue, - OnSuccessCount: &zeroValue, - OnErrorCount: &zeroValue, - BeforePingCount: &zeroValue, - AfterPingCount: &zeroValue, - BeforeToolsCount: &zeroValue, - AfterToolsCount: &zeroValue, + BeforeAnyCount: 0, + OnSuccessCount: 0, + OnErrorCount: 0, + BeforePingCount: 0, + AfterPingCount: 0, + BeforeToolsCount: 0, + AfterToolsCount: 0, OnSuccessData: make([]struct { msg any res any }, 0), BeforePingMessages: make([]*mcp.PingRequest, 0), AfterPingData: make([]struct { Msg *mcp.PingRequest Res *mcp.EmptyResult }, 0), }
1308-1327:⚠️ Potential issueCounter fields should be value types, not pointers, to avoid dangling pointers.
Using pointers for simple int counters creates complexity and potential for dangling pointers when incrementing with stack variables. Use simple int values instead.
type globalHookContext struct { mu sync.Mutex BeforeAnyMessages []any - BeforeAnyCount *int - OnSuccessCount *int - OnErrorCount *int - BeforePingCount *int - AfterPingCount *int - BeforeToolsCount *int - AfterToolsCount *int + BeforeAnyCount int + OnSuccessCount int + OnErrorCount int + BeforePingCount int + AfterPingCount int + BeforeToolsCount int + AfterToolsCount int OnSuccessData []struct { msg any res any } BeforePingMessages []*mcp.PingRequest AfterPingData []struct { Msg *mcp.PingRequest Res *mcp.EmptyResult } }
1347-1352:⚠️ Potential issueFix counter increment pattern to avoid dangling pointers.
The current implementation creates and assigns pointers to stack variables, causing dangling pointers. Use direct value increments instead.
func (*globalHook) BeforeAny(ctx context.Context, hookContext HookContext, id any, method mcp.MCPMethod, message any) { - newCount := (*hookContext.(*globalHookContext).BeforeAnyCount) + 1 - hookContext.(*globalHookContext).BeforeAnyCount = &(newCount) + hookContext.(*globalHookContext).BeforeAnyCount++ // Only collect ping messages for our test if method == mcp.MethodPing { hookContext.(*globalHookContext).BeforeAnyMessages = append(hookContext.(*globalHookContext).BeforeAnyMessages, message) } }This pattern is repeated in multiple methods (lines 1356-1358, 1368-1369, 1392-1394, 1404-1406, 1449-1451, 1459-1461) and should be fixed in all of them.
🧹 Nitpick comments (4)
server/server_test.go (4)
752-753: Fix incorrect parentheses around zero literal.There's an unnecessary pair of parentheses around the
0literal in the slice capacity.- hookContext.(*exampleHookContext).AfterResults = make([]afterResult, (0)) + hookContext.(*exampleHookContext).AfterResults = make([]afterResult, 0)
811-812: Debug print statement should be removed from test code.The
fmt.Printfstatement left in the test code will produce unwanted output during test execution.func (*exampleHook) BeforeListTools(ctx context.Context, hookContext HookContext, id any, message *mcp.ListToolsRequest) { - fmt.Printf("beforeCallTool: %v, %v\n", id, message) }
815-816: Debug print statement should be removed from test code.The
fmt.Printfstatement left in the test code will produce unwanted output during test execution.func (*exampleHook) AfterListTools(ctx context.Context, hookContext HookContext, id any, message *mcp.ListToolsRequest, result *mcp.ListToolsResult) { - fmt.Printf("afterCallTool: %v, %v, %v\n", id, message, result) }
1557-1558: Fix type assertion comparison with IsType.Using
assert.IsType()is more appropriate than direct type comparison. The current code doesn't prove that the items are the same object, only the same type.- assert.IsType(t, hooks.GetHookContext().(*globalHookContext).BeforePingMessages[0], hooks.GetHookContext().(*globalHookContext).BeforeAnyMessages[0], "BeforeAny message should be same type as BeforePing message") + assert.IsType(t, hooks.GetHookContext().(*globalHookContext).BeforePingMessages[0], hooks.GetHookContext().(*globalHookContext).BeforeAnyMessages[0], "BeforeAny message should be same type as BeforePing message")This isn't actually a change, but I'm highlighting that the test is correctly using
assert.IsType()rather than trying to compare the actual values.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
server/server_test.go(10 hunks)
🔇 Additional comments (1)
server/server_test.go (1)
1437-1442:⚠️ Potential issueMissing mutex protection in BeforeGetPrompt and AfterGetPrompt methods.
Unlike other hook methods, these don't have mutex protection, which could lead to race conditions during concurrent test execution.
func (*globalHook) BeforeGetPrompt(ctx context.Context, hookContext HookContext, id any, message *mcp.GetPromptRequest) { + hookContext.(*globalHookContext).mu.Lock() + defer hookContext.(*globalHookContext).mu.Unlock() } func (*globalHook) AfterGetPrompt(ctx context.Context, hookContext HookContext, id any, message *mcp.GetPromptRequest, result *mcp.GetPromptResult) { + hookContext.(*globalHookContext).mu.Lock() + defer hookContext.(*globalHookContext).mu.Unlock() }This issue affects multiple methods. All hook methods should consistently use mutex protection when accessing shared state in the hook context.
Likely an incorrect or invalid review comment.
details in #159
Summary by CodeRabbit