fix: reuse sessions correctly in streamable HTTP transport #615
Conversation
Fixes session management in the streamable HTTP server to properly reuse registered sessions for POST requests instead of always creating ephemeral sessions. This enables SendNotificationToSpecificClient and session-aware features to work correctly with POST-based interactions. Changes: - Check s.server.sessions for existing sessions before creating ephemeral ones - Register sessions after successful initialization from POST requests - Store sessions in both s.server.sessions and s.activeSessions for consistency - Add comprehensive tests for session reuse and notification delivery Fixes #614
WalkthroughUpdates handlePost in streamable_http to reuse existing sessions for non-initialize requests, defer session registration until after a successful initialize response, and manage activeSessions with rollback on registration failure. Adds tests validating per-client notifications and session reuse during tool invocations. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
Merging this branch will increase overall coverage
Coverage by fileChanged files (no unit tests)
Please note that the "Total", "Covered", and "Missed" counts above refer to code statements instead of lines of code. The value in brackets refers to the test coverage of that file in the old version of the code. Changed unit test files
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
server/streamable_http.go (2)
125-127: Update outdated documentation.The comment states that POST handlers will not trigger session registration, but the new code on lines 432-445 now registers sessions from POST initialize requests. This documentation is now incorrect.
Apply this diff to update the comment:
-// Notice: -// Except for the GET handlers(listening), the POST handlers(request/notification) will -// not trigger the session registration. So the methods like `SendNotificationToSpecificClient` -// or `hooks.onRegisterSession` will not be triggered for POST messages. +// Notice: +// Session registration is triggered by GET handlers (listening) and POST initialize requests. +// For non-initialize POST requests, sessions are reused if already registered. +// Methods like `SendNotificationToSpecificClient` and `hooks.onRegisterSession` will be +// triggered for both GET connections and POST initialize requests.
291-293: Update misleading comment about ephemeral sessions.The comment states "The session is ephemeral. Its life is the same as the request." However, for initialize requests, the session is now registered (lines 432-445) and persists across requests, so it's no longer ephemeral.
Apply this diff to clarify the comment:
// Prepare the session for the mcp server -// The session is ephemeral. Its life is the same as the request. It's only created -// for interaction with the mcp server. +// For initialize requests, the session will be registered and persist across requests. +// For non-initialize requests, the session is reused if registered, or an ephemeral +// session is created with a lifetime equal to the request.
🧹 Nitpick comments (1)
server/streamable_http_test.go (1)
1383-1482: Replace sleep with robust synchronization.The test uses
time.Sleep(100 * time.Millisecond)on line 1421 to wait for registration to complete. This can lead to flaky tests on slow systems or under load.Consider using a WaitGroup similar to the first subtest:
+ hooks := &Hooks{} + var sessionRegistered sync.WaitGroup + sessionRegistered.Add(1) + + hooks.AddOnRegisterSession(func(ctx context.Context, session ClientSession) { + sessionRegistered.Done() + }) + - mcpServer := NewMCPServer("test", "1.0.0") + mcpServer := NewMCPServer("test", "1.0.0", WithHooks(hooks)) // ... initialize session code ... - // Give time for registration to complete - time.Sleep(100 * time.Millisecond) + // Wait for registration to complete + done := make(chan struct{}) + go func() { + sessionRegistered.Wait() + close(done) + }() + + select { + case <-done: + // Registration complete + case <-time.After(2 * time.Second): + t.Fatal("Timeout waiting for session registration") + } // Call tool with the session ID
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
server/streamable_http.go(2 hunks)server/streamable_http_test.go(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
server/streamable_http_test.go (7)
server/hooks.go (1)
Hooks(94-121)server/session.go (2)
ClientSession(11-20)ClientSessionFromContext(82-87)server/server.go (3)
NewMCPServer(335-363)WithHooks(287-291)ServerFromContext(80-85)server/streamable_http.go (1)
NewTestStreamableHTTPServer(1089-1093)server/constants.go (1)
HeaderKeySessionID(5-5)mcp/tools.go (3)
NewTool(684-706)CallToolRequest(54-58)CallToolResult(40-51)mcp/utils.go (2)
NewToolResultError(393-403)NewToolResultText(271-280)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: test
- GitHub Check: coverage
🔇 Additional comments (5)
server/streamable_http.go (3)
312-320: LGTM! Session reuse logic correctly implements the fix.The code correctly checks
s.server.sessionsfor existing registered sessions before creating ephemeral ones, enabling session-aware features likeSendNotificationToSpecificClientfor POST requests. The type assertion provides safety against incorrect session types.
324-329: LGTM! Correct fallback to activeSessions for sampling support.The modified logic correctly checks
activeSessionsonly when no registered session was found in the previous step. This maintains support for persistent sessions created by GET connections while respecting the new session reuse logic.
432-445: LGTM! Registration logic correctly handles race conditions.The code properly:
- Checks for existing registration to prevent duplicates (e.g., from concurrent GET connections)
- Stores in
activeSessionsbefore callingRegisterSessionto coordinate with GET handler- Rolls back
activeSessionson registration failure without failing the request- Uses appropriate error logging
The defensive check against
s.server.sessionscombined with the atomicLoadOrStorein the GET handler (line 465) effectively prevents duplicate registration races.server/streamable_http_test.go (2)
1318-1381: LGTM! Comprehensive test with proper synchronization.The test correctly:
- Uses a WaitGroup with timeout to wait for async registration
- Verifies the registered session ID matches the initialize response
- Tests
SendNotificationToSpecificClientfunctionalityThe synchronization approach is robust and avoids flaky timing issues.
1447-1464: LGTM! Test correctly handles both response formats.The test properly handles both SSE format (when notifications are sent) and JSON format (when no notifications occur), making it resilient to the conditional response format logic in the production code.
1 participant
Description
Fixes session management in the streamable HTTP transport to properly reuse registered sessions for POST requests instead of always creating ephemeral sessions. This enables
SendNotificationToSpecificClientand other session-aware features to work correctly with POST-based interactions.Changes:
s.server.sessionsbefore creating ephemeral sessionss.server.sessions(viaRegisterSession) ands.activeSessionsfor consistencySendNotificationToSpecificClientworks with streamable HTTP sessionsFixes #614
Type of Change
Checklist
MCP Spec Compliance
N/A - This is a bug fix for existing functionality, not a new spec implementation.
Additional Information
Testing:
TestStreamableHTTP_SendNotificationToSpecificClientwith two subtestsgo test ./... -race)Backward Compatibility:
This fix restores intended behavior and maintains backward compatibility by preserving the ephemeral session fallback. No breaking changes.
Summary by CodeRabbit
New Features
Bug Fixes
Tests