fix: Add missing session cleanup to the StreamableHTTPServer DELETE handler #667
Conversation
WalkthroughThe pull request adds session cleanup logic to the StreamableHTTPServer's DELETE handler, which now removes the session from the activeSessions map and unregisters it from the MCP server. A corresponding test validates this cleanup behavior and hook invocation. Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes
Possibly related PRs
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 |
![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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
server/streamable_http.go (1)
657-682: Guard against empty session ID.If a client sends
DELETEwithoutMcp-Session-Id, the handler will delete session-scoped stores for""(Line 659). Consider returning400 Bad RequestwhensessionID == ""to prevent unintended cleanup of non-existent sessions.func (s *StreamableHTTPServer) handleDelete(w http.ResponseWriter, r *http.Request) { // delete request terminate the session sessionID := r.Header.Get(HeaderKeySessionID) + if sessionID == "" { + http.Error(w, "Missing session ID", http.StatusBadRequest) + return + } sessionIdManager := s.sessionIdManagerResolver.ResolveSessionIdManager(r)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
server/streamable_http.go(1 hunks)server/streamable_http_test.go(2 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.goserver/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-10-13T09:35:20.180Z
Learnt from: CR
Repo: mark3labs/mcp-go PR: 0
File: AGENTS.md:0-0
Timestamp: 2025-10-13T09:35:20.180Z
Learning: Applies to **/*_test.go : Testing: use testify/assert and testify/require
Applied to files:
server/streamable_http_test.go
📚 Learning: 2025-03-04T07:00:57.111Z
Learnt from: xinwo
Repo: mark3labs/mcp-go PR: 35
File: mcp/tools.go:0-0
Timestamp: 2025-03-04T07:00:57.111Z
Learning: The Tool struct in the mark3labs/mcp-go project should handle both InputSchema and RawInputSchema consistently between MarshalJSON and UnmarshalJSON methods, even though the tools response from MCP server typically doesn't contain rawInputSchema.
Applied to files:
server/streamable_http_test.go
🧬 Code graph analysis (1)
server/streamable_http_test.go (4)
server/session.go (1)
ClientSession(12-21)server/hooks.go (1)
Hooks(94-121)client/transport/constants.go (1)
HeaderKeySessionID(5-5)server/constants.go (1)
HeaderKeySessionID(5-5)
🔇 Additional comments (1)
server/streamable_http_test.go (1)
18-20: Good move switching new assertions to testify.
Matches AGENTS.md guidance. Based on learnings, ...
| func TestStreamableHTTP_Delete(t *testing.T) { | ||
| var hookCalled bool | ||
| var hookSession ClientSession | ||
|
|
||
| hooks := &Hooks{} | ||
| hooks.AddOnUnregisterSession(func(ctx context.Context, session ClientSession) { | ||
| hookCalled = true | ||
| hookSession = session | ||
| }) | ||
|
|
||
| mcpServer := NewMCPServer("test-mcp-server", "1.0", WithHooks(hooks)) | ||
| sseServer := NewStreamableHTTPServer(mcpServer, WithStateful(true)) | ||
| testServer := httptest.NewServer(sseServer) | ||
| defer testServer.Close() | ||
|
|
||
| resp, err := postJSON(testServer.URL, initRequest) | ||
| require.NoError(t, err) | ||
| resp.Body.Close() | ||
| sessionID := resp.Header.Get(HeaderKeySessionID) | ||
|
|
||
| req, _ := http.NewRequest(http.MethodDelete, testServer.URL, nil) | ||
| req.Header.Set(HeaderKeySessionID, sessionID) | ||
|
|
||
| resp, err = testServer.Client().Do(req) | ||
| require.NoError(t, err) | ||
| resp.Body.Close() | ||
|
|
||
| require.Equal(t, http.StatusOK, resp.StatusCode) | ||
|
|
||
| _, activeSessionExists := sseServer.activeSessions.Load(sessionID) | ||
| assert.False(t, activeSessionExists) | ||
|
|
||
| _, serverSessionExists := mcpServer.sessions.Load(sessionID) | ||
| assert.False(t, serverSessionExists) | ||
|
|
||
| assert.True(t, hookCalled) | ||
| assert.Equal(t, sessionID, hookSession.SessionID()) | ||
| } |
There was a problem hiding this comment.
Fix potential race/panic in the unregister hook assertions.
hookCalled / hookSession are written in the hook and read later without synchronization; because the hook currently sets hookCalled = true before hookSession = session, the final hookSession.SessionID() assertion can panic or flake if hook execution is async or reordered by the scheduler.
Prefer capturing the session via a channel (and also assert sessionID is non-empty):
func TestStreamableHTTP_Delete(t *testing.T) {
- var hookCalled bool
- var hookSession ClientSession
+ hookSessionCh := make(chan ClientSession, 1)
hooks := &Hooks{}
hooks.AddOnUnregisterSession(func(ctx context.Context, session ClientSession) {
- hookCalled = true
- hookSession = session
+ select {
+ case hookSessionCh <- session:
+ default:
+ }
})
@@
resp, err := postJSON(testServer.URL, initRequest)
require.NoError(t, err)
resp.Body.Close()
sessionID := resp.Header.Get(HeaderKeySessionID)
+ require.NotEmpty(t, sessionID)
@@
require.Equal(t, http.StatusOK, resp.StatusCode)
@@
- assert.True(t, hookCalled)
- assert.Equal(t, sessionID, hookSession.SessionID())
+ var hookSession ClientSession
+ select {
+ case hookSession = <-hookSessionCh:
+ case <-time.After(1 * time.Second):
+ t.Fatal("timeout waiting for OnUnregisterSession hook")
+ }
+ require.NotNil(t, hookSession)
+ assert.Equal(t, sessionID, hookSession.SessionID())
}📝 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.
| func TestStreamableHTTP_Delete(t *testing.T) { | |
| var hookCalled bool | |
| var hookSession ClientSession | |
| hooks := &Hooks{} | |
| hooks.AddOnUnregisterSession(func(ctx context.Context, session ClientSession) { | |
| hookCalled = true | |
| hookSession = session | |
| }) | |
| mcpServer := NewMCPServer("test-mcp-server", "1.0", WithHooks(hooks)) | |
| sseServer := NewStreamableHTTPServer(mcpServer, WithStateful(true)) | |
| testServer := httptest.NewServer(sseServer) | |
| defer testServer.Close() | |
| resp, err := postJSON(testServer.URL, initRequest) | |
| require.NoError(t, err) | |
| resp.Body.Close() | |
| sessionID := resp.Header.Get(HeaderKeySessionID) | |
| req, _ := http.NewRequest(http.MethodDelete, testServer.URL, nil) | |
| req.Header.Set(HeaderKeySessionID, sessionID) | |
| resp, err = testServer.Client().Do(req) | |
| require.NoError(t, err) | |
| resp.Body.Close() | |
| require.Equal(t, http.StatusOK, resp.StatusCode) | |
| _, activeSessionExists := sseServer.activeSessions.Load(sessionID) | |
| assert.False(t, activeSessionExists) | |
| _, serverSessionExists := mcpServer.sessions.Load(sessionID) | |
| assert.False(t, serverSessionExists) | |
| assert.True(t, hookCalled) | |
| assert.Equal(t, sessionID, hookSession.SessionID()) | |
| } | |
| func TestStreamableHTTP_Delete(t *testing.T) { | |
| hookSessionCh := make(chan ClientSession, 1) | |
| hooks := &Hooks{} | |
| hooks.AddOnUnregisterSession(func(ctx context.Context, session ClientSession) { | |
| select { | |
| case hookSessionCh <- session: | |
| default: | |
| } | |
| }) | |
| mcpServer := NewMCPServer("test-mcp-server", "1.0", WithHooks(hooks)) | |
| sseServer := NewStreamableHTTPServer(mcpServer, WithStateful(true)) | |
| testServer := httptest.NewServer(sseServer) | |
| defer testServer.Close() | |
| resp, err := postJSON(testServer.URL, initRequest) | |
| require.NoError(t, err) | |
| resp.Body.Close() | |
| sessionID := resp.Header.Get(HeaderKeySessionID) | |
| require.NotEmpty(t, sessionID) | |
| req, _ := http.NewRequest(http.MethodDelete, testServer.URL, nil) | |
| req.Header.Set(HeaderKeySessionID, sessionID) | |
| resp, err = testServer.Client().Do(req) | |
| require.NoError(t, err) | |
| resp.Body.Close() | |
| require.Equal(t, http.StatusOK, resp.StatusCode) | |
| _, activeSessionExists := sseServer.activeSessions.Load(sessionID) | |
| assert.False(t, activeSessionExists) | |
| _, serverSessionExists := mcpServer.sessions.Load(sessionID) | |
| assert.False(t, serverSessionExists) | |
| var hookSession ClientSession | |
| select { | |
| case hookSession = <-hookSessionCh: | |
| case <-time.After(1 * time.Second): | |
| t.Fatal("timeout waiting for OnUnregisterSession hook") | |
| } | |
| require.NotNil(t, hookSession) | |
| assert.Equal(t, sessionID, hookSession.SessionID()) | |
| } |
🤖 Prompt for AI Agents
In server/streamable_http_test.go around lines 2490 to 2527, the test reads
hookCalled/hookSession without synchronization which can race/panic; change the
hook to send the session (or sessionID) over a channel and in the test wait to
receive it (with a timeout) before making assertions, assert the received
sessionID is non-empty, then compare it to the request sessionID; this ensures
ordering and avoids read/write races on shared variables.
2 participants
Description
This PR adds two missing session cleanup calls to the streamable http transport's handleDelete method.
UnregisterSession()methodFixes #651
Type of Change
Checklist
Additional Information
When writing my test, I noticed an inconsistency between the existing assertion style and AGENTS.md.
AGENTS.md tells agents to use assert/require. The existing tests use
t.Error().Summary by CodeRabbit
Release Notes
Bug Fixes
Tests
✏️ Tip: You can customize this high-level summary in your review settings.