fix: Add missing session cleanup to the StreamableHTTPServer DELETE handler (#667)

* fix: Add missing session cleanup to the StreamableHTTPServer DELETE handler

* Fix merge typo

---------

Co-authored-by: Connor Zanin <[email protected]>

Author: cnnrznn

server/streamable_http.go

@@ -698,6 +698,8 @@ func (s *StreamableHTTPServer) handleDelete(w http.ResponseWriter, r *http.Reque
 	s.sessionLogLevels.delete(sessionID)
 	// remove current session's requstID information
 	s.sessionRequestIDs.Delete(sessionID)
+	s.activeSessions.Delete(sessionID)
+	s.server.UnregisterSession(r.Context(), sessionID)
 
 	w.WriteHeader(http.StatusOK)
 }

server/streamable_http_test.go

@@ -15,6 +15,8 @@ import (
 	"time"
 
 	"github.com/mark3labs/mcp-go/mcp"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 type jsonRPCResponse struct {
@@ -2485,6 +2487,45 @@ func TestStreamableHTTP_GET_NonFlusherReturns405(t *testing.T) {
 	})
 }
 
+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_DrainNotifications(t *testing.T) {
 	t.Run("drain pending notifications after response is computed", func(t *testing.T) {
 		mcpServer := NewMCPServer("test-mcp-server", "1.0")