feat: quick return tool-call request, send response via SSE in goroutine #163
Conversation
WalkthroughThe update refactors the message handling flow within the SSE server implementation. The Changes
Possibly related PRs
✨ Finishing Touches
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: 1
🧹 Nitpick comments (2)
server/sse.go (2)
358-358: Consider adding documentation about the changed behavior.This line explains the purpose, but consider updating function-level comments at line 324 to document the changed behavior that all responses are now sent via SSE rather than in the HTTP response.
375-377: Consider adding backpressure mechanism for queue full scenarios.Currently, if the event queue is full, the message is silently dropped. Consider adding a backpressure mechanism or at minimum logging this situation to help diagnose issues.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
server/sse.go(3 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (1)
server/sse.go (1)
server/server.go (1)
MCPServer(143-166)
🔇 Additional comments (3)
server/sse.go (3)
56-64: Code formatting improvement looks good.The alignment of struct fields improves readability and maintains consistent style.
161-166: Struct initialization formatting improvement looks good.The alignment of initialization fields maintains better consistency with the struct definition.
361-379:✅ Verification successful
Verify that goroutine cleanup is properly handled.
The goroutine spawned here will continue running until it processes the message and sends any response. Ensure that it will properly terminate even if the session is closed during processing.
🏁 Script executed:
#!/bin/bash # Check if there's proper handling of session termination in other goroutines grep -A 10 "case <-session.done" server/sse.goLength of output: 1046
🏁 Script executed:
#!/bin/bash # Locate definition(s) of HandleMessage to inspect its signature rg -n "HandleMessage" -C 2Length of output: 7050
🏁 Script executed:
#!/bin/bash # Display the surrounding context of the goroutine to see where ctx comes from sed -n '300,420p' server/sse.goLength of output: 3549
No goroutine leak – single‐run worker auto‑terminates
The anonymous goroutine in handleMessage only processes one request then returns:
- It calls
s.server.HandleMessage(ctx, rawMessage)once.- After that, the
selectonsession.eventQueueincludes both<-session.done(skips queuing) and adefaultcase, so it never blocks indefinitely.- Once that select completes, the function exits and the goroutine is freed.
Because it’s not a long‐lived loop, there’s no need for additional cleanup.
|
@ezynda3 Okay, done! |
server/sse.go
Outdated
| if response != nil { | ||
| eventData, err := json.Marshal(response) | ||
| if err != nil { | ||
| s.writeJSONRPCError(w, nil, mcp.INTERNAL_ERROR, "Fail to marshal response") |
There was a problem hiding this comment.
Better to just log this as we have already sent a response and doesn't really make sense write anymore to the response...
There was a problem hiding this comment.
@ezynda3 Maybe I should have to return this error to sse, otherwise the client might not know if the execution was successful?
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
server/sse.go (1)
373-373: Minor: Unnecessary use of fmt.Sprintf for static string.The static analysis tool flagged that using fmt.Sprintf here is unnecessary since there are no format specifiers.
- message = fmt.Sprintf("event: message\ndata: {\"error\": \"internal error\",\"jsonrpc\": \"2.0\", \"id\": null}\n\n") + message = "event: message\ndata: {\"error\": \"internal error\",\"jsonrpc\": \"2.0\", \"id\": null}\n\n"🧰 Tools
🪛 golangci-lint (1.64.8)
373-373: S1039: unnecessary use of fmt.Sprintf
(gosimple)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
server/sse.go(5 hunks)server/sse_test.go(6 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (1)
server/sse.go (2)
server/server.go (1)
MCPServer(143-166)examples/custom_context/main.go (1)
MCPServer(100-102)
🪛 golangci-lint (1.64.8)
server/sse_test.go
197-197: testinggoroutine: call to (*testing.T).Fatalf from a non-test goroutine
(govet)
server/sse.go
373-373: S1039: unnecessary use of fmt.Sprintf
(gosimple)
🔇 Additional comments (7)
server/sse.go (4)
7-7: Good: Added log package import for error handling.Adding the log package import is necessary for the new error logging implemented in the handleMessage function.
57-65: Improved struct field alignment for better readability.The consistent alignment of struct fields makes the code more readable.
162-167: Improved constructor field alignment consistent with struct definition.The alignment in the constructor now matches the struct definition style.
385-386: Good: Added logging for queue full condition.Properly logging the condition where the event queue is full helps with monitoring and debugging.
server/sse_test.go (3)
751-758: Good: Helper function to reduce code duplication.The new
readSeeEventhelper function improves code maintainability by extracting the common pattern of reading SSE events.
625-634: Good: Updated test to handle asynchronous SSE responses.The code correctly reads and parses the SSE event data to validate the response, aligning with the changes in the main code.
672-682: Good: Consistent handling of SSE events across tests.The code applies the same pattern for reading and parsing SSE events, maintaining consistency across the test suite.
server/sse.go
Outdated
| // quick return request, send 202 Accepted with no body, then deal the message and sent response via SSE | ||
| w.WriteHeader(http.StatusAccepted) | ||
|
|
||
| // Only send response if there is one (not for notifications) | ||
| if response != nil { | ||
| eventData, _ := json.Marshal(response) | ||
| go func() { | ||
| // Process message through MCPServer | ||
| response := s.server.HandleMessage(ctx, rawMessage) | ||
|
|
||
| // Queue the event for sending via SSE | ||
| select { | ||
| case session.eventQueue <- fmt.Sprintf("event: message\ndata: %s\n\n", eventData): | ||
| // Event queued successfully | ||
| case <-session.done: | ||
| // Session is closed, don't try to queue | ||
| default: | ||
| // Queue is full, could log this | ||
| } | ||
| // Only send response if there is one (not for notifications) | ||
| if response != nil { | ||
| var message string | ||
| eventData, err := json.Marshal(response) | ||
| if err != nil { | ||
| // If there is an error marshalling the response, send a generic error response | ||
| log.Printf("failed to marshal response: %v", err) | ||
| message = fmt.Sprintf("event: message\ndata: {\"error\": \"internal error\",\"jsonrpc\": \"2.0\", \"id\": null}\n\n") | ||
| return | ||
| } | ||
| message = fmt.Sprintf("event: message\ndata: %s\n\n", eventData) | ||
|
|
||
| // Send HTTP response | ||
| w.Header().Set("Content-Type", "application/json") | ||
| w.WriteHeader(http.StatusAccepted) | ||
| json.NewEncoder(w).Encode(response) | ||
| } else { | ||
| // For notifications, just send 202 Accepted with no body | ||
| w.WriteHeader(http.StatusAccepted) | ||
| } | ||
| // Queue the event for sending via SSE | ||
| select { | ||
| case session.eventQueue <- message: | ||
| // Event queued successfully | ||
| case <-session.done: | ||
| // Session is closed, don't try to queue | ||
| default: | ||
| // Queue is full, log this situation | ||
| log.Printf("Event queue full for session %s", sessionID) | ||
| } | ||
| } | ||
| }() |
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Good implementation of asynchronous message processing and quick response.
This change effectively decouples the HTTP response from the message processing cycle, improving responsiveness by immediately returning 202 Accepted and handling the message asynchronously via SSE. The implementation follows a common pattern for async processing in API handlers.
However, there's an issue with the error handling for JSON marshaling:
- message = fmt.Sprintf("event: message\ndata: {\"error\": \"internal error\",\"jsonrpc\": \"2.0\", \"id\": null}\n\n")
- return
+ log.Printf("failed to marshal response: %v", err)
+ returnThe error message string is created but never used since the function returns immediately after logging the error.
🏁 Script executed:
#!/bin/bash
# Check if there are any other places where JSON marshaling errors are handled consistently
grep -A 5 "json\.Marshal" server/sse.goLength of output: 726
Ensure error SSE frames are queued before returning
In server/sse.go around the json.Marshal(response) error block (lines ~365–371), you build an error message but return immediately, so the client never receives it. Update the error path to enqueue that message just like the success case:
if err != nil {
log.Printf("failed to marshal response: %v", err)
- message = fmt.Sprintf("event: message\ndata: {\"error\": \"internal error\",\"jsonrpc\": \"2.0\", \"id\": null}\n\n")
- return
+ message = fmt.Sprintf("event: message\ndata: {\"error\": \"internal error\",\"jsonrpc\": \"2.0\", \"id\": null}\n\n")
+ select {
+ case session.eventQueue <- message:
+ // Error event queued
+ case <-session.done:
+ // Session closed
+ default:
+ log.Printf("Event queue full for session %s", sessionID)
+ }
+ return
}This guarantees that clients get an SSE error frame when marshaling fails.
📝 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.
| // quick return request, send 202 Accepted with no body, then deal the message and sent response via SSE | |
| w.WriteHeader(http.StatusAccepted) | |
| // Only send response if there is one (not for notifications) | |
| if response != nil { | |
| eventData, _ := json.Marshal(response) | |
| go func() { | |
| // Process message through MCPServer | |
| response := s.server.HandleMessage(ctx, rawMessage) | |
| // Queue the event for sending via SSE | |
| select { | |
| case session.eventQueue <- fmt.Sprintf("event: message\ndata: %s\n\n", eventData): | |
| // Event queued successfully | |
| case <-session.done: | |
| // Session is closed, don't try to queue | |
| default: | |
| // Queue is full, could log this | |
| } | |
| // Only send response if there is one (not for notifications) | |
| if response != nil { | |
| var message string | |
| eventData, err := json.Marshal(response) | |
| if err != nil { | |
| // If there is an error marshalling the response, send a generic error response | |
| log.Printf("failed to marshal response: %v", err) | |
| message = fmt.Sprintf("event: message\ndata: {\"error\": \"internal error\",\"jsonrpc\": \"2.0\", \"id\": null}\n\n") | |
| return | |
| } | |
| message = fmt.Sprintf("event: message\ndata: %s\n\n", eventData) | |
| // Send HTTP response | |
| w.Header().Set("Content-Type", "application/json") | |
| w.WriteHeader(http.StatusAccepted) | |
| json.NewEncoder(w).Encode(response) | |
| } else { | |
| // For notifications, just send 202 Accepted with no body | |
| w.WriteHeader(http.StatusAccepted) | |
| } | |
| // Queue the event for sending via SSE | |
| select { | |
| case session.eventQueue <- message: | |
| // Event queued successfully | |
| case <-session.done: | |
| // Session is closed, don't try to queue | |
| default: | |
| // Queue is full, log this situation | |
| log.Printf("Event queue full for session %s", sessionID) | |
| } | |
| } | |
| }() | |
| // quick return request, send 202 Accepted with no body, then deal the message and send response via SSE | |
| w.WriteHeader(http.StatusAccepted) | |
| go func() { | |
| // Process message through MCPServer | |
| response := s.server.HandleMessage(ctx, rawMessage) | |
| // Only send response if there is one (not for notifications) | |
| if response != nil { | |
| var message string | |
| eventData, err := json.Marshal(response) | |
| if err != nil { | |
| // If there is an error marshalling the response, send a generic error response | |
| log.Printf("failed to marshal response: %v", err) | |
| message = fmt.Sprintf("event: message\ndata: {\"error\": \"internal error\",\"jsonrpc\": \"2.0\", \"id\": null}\n\n") | |
| select { | |
| case session.eventQueue <- message: | |
| // Error event queued | |
| case <-session.done: | |
| // Session closed | |
| default: | |
| log.Printf("Event queue full for session %s", sessionID) | |
| } | |
| return | |
| } | |
| message = fmt.Sprintf("event: message\ndata: %s\n\n", eventData) | |
| // Queue the event for sending via SSE | |
| select { | |
| case session.eventQueue <- message: | |
| // Event queued successfully | |
| case <-session.done: | |
| // Session is closed, don't try to queue | |
| default: | |
| // Queue is full, log this situation | |
| log.Printf("Event queue full for session %s", sessionID) | |
| } | |
| } | |
| }() |
🧰 Tools
🪛 golangci-lint (1.64.8)
373-373: S1039: unnecessary use of fmt.Sprintf
(gosimple)
| endpointEvent, err = readSeeEvent(sseResp) | ||
| if err != nil { | ||
| t.Fatalf("Failed to read SSE response: %v", err) | ||
| } | ||
| respFromSee := strings.TrimSpace( | ||
| strings.Split(strings.Split(endpointEvent, "data: ")[1], "\n")[0], | ||
| ) | ||
|
|
||
| fmt.Printf("========> %v", respFromSee) | ||
| var response map[string]interface{} | ||
| if err := json.NewDecoder(resp.Body).Decode(&response); err != nil { | ||
| if err := json.NewDecoder(strings.NewReader(respFromSee)).Decode(&response); err != nil { |
There was a problem hiding this comment.
Fix: Testing function called from goroutine.
The static analysis tool flagged that t.Fatalf is being called from a non-test goroutine, which can lead to race conditions or panics in the test framework.
- t.Fatalf("Failed to read SSE response: %v", err)
+ t.Errorf("Session %d: Failed to read SSE response: %v", sessionNum, err)
+ returnCommittable suggestion skipped: line range outside the PR's diff.
🧰 Tools
🪛 golangci-lint (1.64.8)
197-197: testinggoroutine: call to (*testing.T).Fatalf from a non-test goroutine
(govet)
CeerDecy
force-pushed
the
feature/request-timeout
branch
from
cde2b2a to
73d36a9
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
server/sse.go (1)
359-389:⚠️ Potential issueCritical error: Error message isn't being sent to the client when JSON marshaling fails.
The implementation successfully decouples the HTTP response from message processing, improving responsiveness by returning 202 Accepted immediately and processing messages asynchronously. However, there's a critical issue in the error handling path that needs to be fixed.
When JSON marshaling fails, you're creating an error message but returning from the goroutine immediately without sending it to the client. This means clients won't receive any notification when this error occurs.
if eventData, err := json.Marshal(response); err != nil { // If there is an error marshalling the response, send a generic error response log.Printf("failed to marshal response: %v", err) message = fmt.Sprintf("event: message\ndata: {\"error\": \"internal error\",\"jsonrpc\": \"2.0\", \"id\": null}\n\n") - return + // Queue the error event for sending via SSE + select { + case session.eventQueue <- message: + // Error event queued successfully + case <-session.done: + // Session is closed, don't try to queue + default: + // Queue is full, log this situation + log.Printf("Event queue full for session %s", sessionID) + } + return } else { message = fmt.Sprintf("event: message\ndata: %s\n\n", eventData) }Also, as noted by the static analysis tool, the error message can be simplified:
-message = fmt.Sprintf("event: message\ndata: {\"error\": \"internal error\",\"jsonrpc\": \"2.0\", \"id\": null}\n\n") +message = "event: message\ndata: {\"error\": \"internal error\",\"jsonrpc\": \"2.0\", \"id\": null}\n\n"🧰 Tools
🪛 golangci-lint (1.64.8)
372-372: S1039: unnecessary use of fmt.Sprintf
(gosimple)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
server/sse.go(5 hunks)server/sse_test.go(6 hunks)
🧰 Additional context used
🪛 golangci-lint (1.64.8)
server/sse_test.go
197-197: testinggoroutine: call to (*testing.T).Fatalf from a non-test goroutine
(govet)
server/sse.go
372-372: S1039: unnecessary use of fmt.Sprintf
(gosimple)
🔇 Additional comments (4)
server/sse.go (3)
7-7: Appropriate addition of the log package.The log package is correctly imported to support the new logging functionality introduced in the handleMessage function.
57-65: Improved struct field alignment.The struct fields have been reorganized with consistent alignment, improving code readability.
162-167: Consistent struct field formatting.The struct field initialization formatting is now consistent with the struct declaration, enhancing maintainability.
server/sse_test.go (1)
751-758: Good implementation of helper function for reading SSE events.The
readSeeEventfunction is a clean abstraction that centralizes the logic for reading from SSE response bodies, reducing code duplication across tests.
| @@ -59,13 +59,10 @@ func TestSSEServer(t *testing.T) { | |||
| defer sseResp.Body.Close() | |||
|
|
|||
| // Read the endpoint event | |||
| buf := make([]byte, 1024) | |||
| n, err := sseResp.Body.Read(buf) | |||
| endpointEvent, err := readSeeEvent(sseResp) | |||
There was a problem hiding this comment.
Fix test code to accommodate asynchronous response delivery.
The test code has been properly adapted to read and process responses from SSE events instead of directly from HTTP response bodies, which aligns with the server-side changes. However, there's a critical issue with using t.Fatalf in a goroutine.
Using t.Fatalf from a non-test goroutine (line 197) can cause race conditions or panics in the test framework. This was previously flagged in code reviews but hasn't been fixed.
endpointEvent, err = readSeeEvent(sseResp)
if err != nil {
- t.Fatalf("Failed to read SSE response: %v", err)
+ t.Errorf("Session %d: Failed to read SSE response: %v", sessionNum, err)
+ return
}Apply this fix to all instances of t.Fatalf used in goroutines in this file.
Also applies to: 195-204, 582-632, 672-682
to solve the #162
the ToolCall request should return an HTTP 202 (Accepted) status code immediately, without waiting for
handleMessageto complete.like https://github.com/modelcontextprotocol/python-sdk src/mcp/server/sse.py 174
Summary by CodeRabbit
Summary by CodeRabbit