fix: notification break the client tool call #642
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Contributor
WalkthroughAdds a per-session SSE upgrade invocation before sending notifications and introduces a regression test that verifies adding a tool during an in-progress tool call does not break a streamable HTTP client response. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes
Possibly related PRs
Suggested labels
Suggested reviewers
Pre-merge checks and finishing touches❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Signed-off-by: hai.yue <hai.yue@ingka.com>
Contributor
|
@yuehaii would you mind adding a regression test for this? Something like the following? // TestStreamableHTTP_AddToolDuringToolCall tests that adding a tool while a tool call
// is in progress doesn't break the client's response.
// This is a regression test for issue #638 where notifications sent via
// sendNotificationToAllClients during an in-progress request would cause
// the response to fail with "unexpected nil response".
func TestStreamableHTTP_AddToolDuringToolCall(t *testing.T) {
mcpServer := NewMCPServer("test-mcp-server", "1.0",
WithToolCapabilities(true), // Enable tool list change notifications
)
// Add a tool that takes some time to complete
mcpServer.AddTool(mcp.NewTool("slow_tool",
mcp.WithDescription("A tool that takes time to complete"),
), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
// Simulate work that takes some time
time.Sleep(100 * time.Millisecond)
return mcp.NewToolResultText("done"), nil
})
server := NewTestStreamableHTTPServer(mcpServer, WithStateful(true))
defer server.Close()
// Initialize to get session
resp, err := postJSON(server.URL, initRequest)
if err != nil {
t.Fatalf("Failed to initialize: %v", err)
}
sessionID := resp.Header.Get(HeaderKeySessionID)
resp.Body.Close()
if sessionID == "" {
t.Fatal("Expected session ID in response header")
}
// Start the tool call in a goroutine
resultChan := make(chan struct {
statusCode int
body string
err error
})
go func() {
toolRequest := map[string]any{
"jsonrpc": "2.0",
"id": 1,
"method": "tools/call",
"params": map[string]any{
"name": "slow_tool",
},
}
toolBody, _ := json.Marshal(toolRequest)
req, _ := http.NewRequest("POST", server.URL, bytes.NewReader(toolBody))
req.Header.Set("Content-Type", "application/json")
req.Header.Set(HeaderKeySessionID, sessionID)
resp, err := server.Client().Do(req)
if err != nil {
resultChan <- struct {
statusCode int
body string
err error
}{0, "", err}
return
}
defer resp.Body.Close()
body, _ := io.ReadAll(resp.Body)
resultChan <- struct {
statusCode int
body string
err error
}{resp.StatusCode, string(body), nil}
}()
// Wait a bit then add a new tool while the slow_tool is executing
// This triggers sendNotificationToAllClients
time.Sleep(50 * time.Millisecond)
mcpServer.AddTool(mcp.NewTool("new_tool",
mcp.WithDescription("A new tool added during execution"),
), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
return mcp.NewToolResultText("new tool result"), nil
})
// Wait for the tool call to complete
result := <-resultChan
if result.err != nil {
t.Fatalf("Tool call failed with error: %v", result.err)
}
if result.statusCode != http.StatusOK {
t.Errorf("Expected status 200, got %d. Body: %s", result.statusCode, result.body)
}
// The response should contain the tool result
// It may be SSE format (text/event-stream) due to the notification upgrade
if !strings.Contains(result.body, "done") {
t.Errorf("Expected response to contain 'done', got: %s", result.body)
}
} |
Signed-off-by: hai.yue <hai.yue@ingka.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
server/streamable_http_test.go (1)
2255-2341: Regression scenario is well-covered; consider adding a timeout to avoid hangsThe test correctly models the bug scenario (stateful session,
WithToolCapabilities(true), slow tool, concurrentAddTool, and asserting HTTP 200 +"done"in the body). That should reliably catch regressions around the nil-response behavior.One small robustness improvement: right now, if the server path ever deadlocks or blocks indefinitely, the goroutine doing
server.Client().Do(req)and the main goroutine waiting onresultChancould hang the test. You can keep the behavior the same but bound the failure mode by giving the request a context timeout:- go func() { + go func() { toolRequest := map[string]any{ "jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": map[string]any{ "name": "slow_tool", }, } - toolBody, _ := json.Marshal(toolRequest) - req, _ := http.NewRequest("POST", server.URL, bytes.NewReader(toolBody)) + toolBody, _ := json.Marshal(toolRequest) + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + req, _ := http.NewRequestWithContext(ctx, http.MethodPost, server.URL, bytes.NewReader(toolBody)) req.Header.Set("Content-Type", "application/json") req.Header.Set(HeaderKeySessionID, sessionID) resp, err := server.Client().Do(req)With this, a regression that causes the call to stall will still fail the test (via a non-nil
result.err) but won’t hang the test suite. Optionally, if you ever see flakes under heavy load, you could also increase the handler sleep or coordinate start-of-work via a channel instead of relying purely on the 50ms/100ms timing window, but that’s likely overkill here.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
server/streamable_http_test.go(1 hunks)
🧰 Additional context used
📓 Path-based instructions (2)
**/*.go
📄 CodeRabbit inference engine (AGENTS.md)
**/*.go: Order imports: standard library first, then third-party, then local packages (goimports enforces this)
Follow Go naming conventions: exported identifiers in PascalCase; unexported in camelCase; acronyms uppercase (HTTP, JSON, MCP)
Error handling: return sentinel errors, wrap with fmt.Errorf("context: %w", err), and check with errors.Is/As
Prefer explicit types and strongly-typed structs; avoid using any except where protocol flexibility is required (e.g., Arguments any)
All exported types and functions must have GoDoc comments starting with the identifier name; avoid inline comments unless necessary
Functions that are handlers or long-running must accept context.Context as the first parameter
Ensure thread safety for shared state using sync.Mutex and document thread-safety requirements in comments
For JSON: use json struct tags with omitempty for optional fields; use json.RawMessage for flexible/deferred parsing
Files:
server/streamable_http_test.go
**/*_test.go
📄 CodeRabbit inference engine (AGENTS.md)
**/*_test.go: Testing: use testify/assert and testify/require
Write table-driven tests using a tests := []struct{ name, ... } pattern
Go test files must end with _test.go
Files:
server/streamable_http_test.go
🧠 Learnings (2)
📚 Learning: 2025-06-23T11:10:42.948Z
Learnt from: floatingIce91
Repo: mark3labs/mcp-go PR: 401
File: server/server.go:1082-1092
Timestamp: 2025-06-23T11:10:42.948Z
Learning: In Go MCP server, ServerTool.Tool field is only used for tool listing and indexing, not for tool execution or middleware. During handleToolCall, only the Handler field is used, so dynamic tools don't need the Tool field populated.
Applied to files:
server/streamable_http_test.go
📚 Learning: 2025-03-04T06:59:43.882Z
Learnt from: xinwo
Repo: mark3labs/mcp-go PR: 35
File: mcp/tools.go:107-137
Timestamp: 2025-03-04T06:59:43.882Z
Learning: Tool responses from the MCP server shouldn't contain RawInputSchema, which is why the UnmarshalJSON method for the Tool struct is implemented to handle only the structured InputSchema format.
Applied to files:
server/streamable_http_test.go
🧬 Code graph analysis (1)
server/streamable_http_test.go (4)
mcp/utils.go (1)
NewToolResultText(271-280)server/streamable_http.go (2)
NewTestStreamableHTTPServer(1381-1385)WithStateful(84-90)client/transport/constants.go (1)
HeaderKeySessionID(5-5)server/constants.go (1)
HeaderKeySessionID(5-5)
Contributor
Author
I have added this regression test. Thanks for the assist, @ezynda3 |
sd2k
pushed a commit
to grafana/mcp-go
that referenced
this pull request
Nov 27, 2025
* feat: client roots feature * feat: finish client roots, pass unit and integration test * client roots http sample code * client roots for stdio and pass integration test * update roots stio client example * add godoc and const of rootlist * update godoc and data format * update examples for client roots * add fallback for demonstration * adjust roots path and signals of examples * update roots http client example * samples: fix unit test and refactor with lint * examples: refactor to adapt windows os and nitpick comments * update for nitpick comments * refactor for nitpick comments * fix: notifications breaking the tool call Signed-off-by: hai.yue <hai.yue@ingka.com> * add a regression test mark3labs#642 (comment) Signed-off-by: hai.yue <hai.yue@ingka.com>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
If there is adding/removing tool of the server, it will trigger the notification. And if the client try to call tool at this time, it will get failure.
Fixes #638
Type of Change
Checklist
Additional Information
The root cause is that the notification already upgrade client communication to SSE. But session.upgradeToSSE.Load() is always false and go into incorrect http stream processing logic.
I have verify the solution locally and will add a test case for such special scenario later.
Summary by CodeRabbit
Improvements
Tests
✏️ Tip: You can customize this high-level summary in your review settings.