fix: reuse sessions correctly in streamable HTTP transport (#615)

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

Author: ezynda3

server/streamable_http.go

@@ -309,12 +309,23 @@ func (s *StreamableHTTPServer) handlePost(w http.ResponseWriter, r *http.Request
 		}
 	}
 
+	// For non-initialize requests, try to reuse existing registered session
+	var session *streamableHttpSession
+	if !isInitializeRequest {
+		if sessionValue, ok := s.server.sessions.Load(sessionID); ok {
+			if existingSession, ok := sessionValue.(*streamableHttpSession); ok {
+				session = existingSession
+			}
+		}
+	}
+
 	// Check if a persistent session exists (for sampling support), otherwise create ephemeral session
 	// Persistent sessions are created by GET (continuous listening) connections
-	var session *streamableHttpSession
-	if sessionInterface, exists := s.activeSessions.Load(sessionID); exists {
-		if persistentSession, ok := sessionInterface.(*streamableHttpSession); ok {
-			session = persistentSession
+	if session == nil {
+		if sessionInterface, exists := s.activeSessions.Load(sessionID); exists {
+			if persistentSession, ok := sessionInterface.(*streamableHttpSession); ok {
+				session = persistentSession
+			}
 		}
 	}
 
@@ -417,6 +428,21 @@ func (s *StreamableHTTPServer) handlePost(w http.ResponseWriter, r *http.Request
 			s.logger.Errorf("Failed to write response: %v", err)
 		}
 	}
+
+	// Register session after successful initialization
+	// Only register if not already registered (e.g., by a GET connection)
+	if isInitializeRequest && sessionID != "" {
+		if _, exists := s.server.sessions.Load(sessionID); !exists {
+			// Store in activeSessions to prevent duplicate registration from GET
+			s.activeSessions.Store(sessionID, session)
+			// Register the session with the MCPServer for notification support
+			if err := s.server.RegisterSession(ctx, session); err != nil {
+				s.logger.Errorf("Failed to register POST session: %v", err)
+				s.activeSessions.Delete(sessionID)
+				// Don't fail the request, just log the error
+			}
+		}
+	}
 }
 
 func (s *StreamableHTTPServer) handleGet(w http.ResponseWriter, r *http.Request) {

server/streamable_http_test.go

@@ -1314,3 +1314,170 @@ func TestInsecureStatefulSessionIdManager(t *testing.T) {
 		}
 	})
 }
+
+func TestStreamableHTTP_SendNotificationToSpecificClient(t *testing.T) {
+	t.Run("POST session registration enables SendNotificationToSpecificClient", func(t *testing.T) {
+		hooks := &Hooks{}
+		var registeredSessionID string
+		var mu sync.Mutex
+		var sessionRegistered sync.WaitGroup
+		sessionRegistered.Add(1)
+
+		hooks.AddOnRegisterSession(func(ctx context.Context, session ClientSession) {
+			mu.Lock()
+			registeredSessionID = session.SessionID()
+			mu.Unlock()
+			sessionRegistered.Done()
+		})
+
+		mcpServer := NewMCPServer("test", "1.0.0", WithHooks(hooks))
+		testServer := NewTestStreamableHTTPServer(mcpServer)
+		defer testServer.Close()
+
+		// Send initialize request to register session
+		resp, err := postJSON(testServer.URL, initRequest)
+		if err != nil {
+			t.Fatalf("Failed to send initialize request: %v", err)
+		}
+		defer resp.Body.Close()
+
+		if resp.StatusCode != http.StatusOK {
+			t.Fatalf("Expected status 200, got %d", resp.StatusCode)
+		}
+
+		// Get session ID from response header
+		sessionID := resp.Header.Get(HeaderKeySessionID)
+		if sessionID == "" {
+			t.Fatal("Expected session ID in response header")
+		}
+
+		// Wait for session registration
+		done := make(chan struct{})
+		go func() {
+			sessionRegistered.Wait()
+			close(done)
+		}()
+
+		select {
+		case <-done:
+			// Session registered successfully
+		case <-time.After(2 * time.Second):
+			t.Fatal("Timeout waiting for session registration")
+		}
+
+		mu.Lock()
+		if registeredSessionID != sessionID {
+			t.Errorf("Expected registered session ID %s, got %s", sessionID, registeredSessionID)
+		}
+		mu.Unlock()
+
+		// Now test SendNotificationToSpecificClient
+		err = mcpServer.SendNotificationToSpecificClient(sessionID, "test/notification", map[string]any{
+			"message": "test notification",
+		})
+		if err != nil {
+			t.Errorf("SendNotificationToSpecificClient failed: %v", err)
+		}
+	})
+
+	t.Run("Session reuse for non-initialize requests", func(t *testing.T) {
+		mcpServer := NewMCPServer("test", "1.0.0")
+
+		// Add a tool that sends a notification
+		mcpServer.AddTool(mcp.NewTool("notify_tool"), func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+			session := ClientSessionFromContext(ctx)
+			if session == nil {
+				return mcp.NewToolResultError("no session in context"), nil
+			}
+
+			// Try to send notification to specific client
+			server := ServerFromContext(ctx)
+			err := server.SendNotificationToSpecificClient(session.SessionID(), "tool/notification", map[string]any{
+				"from": "tool",
+			})
+			if err != nil {
+				return mcp.NewToolResultError(fmt.Sprintf("notification failed: %v", err)), nil
+			}
+
+			return mcp.NewToolResultText("notification sent"), nil
+		})
+
+		testServer := NewTestStreamableHTTPServer(mcpServer)
+		defer testServer.Close()
+
+		// Initialize session
+		resp, err := postJSON(testServer.URL, initRequest)
+		if err != nil {
+			t.Fatalf("Failed to send initialize request: %v", err)
+		}
+		sessionID := resp.Header.Get(HeaderKeySessionID)
+		resp.Body.Close()
+
+		if sessionID == "" {
+			t.Fatal("Expected session ID in response header")
+		}
+
+		// Give time for registration to complete
+		time.Sleep(100 * time.Millisecond)
+
+		// Call tool with the session ID
+		toolCallRequest := map[string]any{
+			"jsonrpc": "2.0",
+			"id":      2,
+			"method":  "tools/call",
+			"params": map[string]any{
+				"name": "notify_tool",
+			},
+		}
+
+		jsonBody, _ := json.Marshal(toolCallRequest)
+		req, _ := http.NewRequest(http.MethodPost, testServer.URL, bytes.NewBuffer(jsonBody))
+		req.Header.Set("Content-Type", "application/json")
+		req.Header.Set(HeaderKeySessionID, sessionID)
+
+		resp, err = http.DefaultClient.Do(req)
+		if err != nil {
+			t.Fatalf("Failed to call tool: %v", err)
+		}
+		defer resp.Body.Close()
+
+		bodyBytes, _ := io.ReadAll(resp.Body)
+		bodyStr := string(bodyBytes)
+
+		// Response might be SSE format if notification was sent
+		var toolResponse jsonRPCResponse
+		if strings.HasPrefix(bodyStr, "event: message") {
+			// Parse SSE format
+			lines := strings.Split(bodyStr, "\n")
+			for _, line := range lines {
+				if strings.HasPrefix(line, "data: ") {
+					jsonData := strings.TrimPrefix(line, "data: ")
+					if err := json.Unmarshal([]byte(jsonData), &toolResponse); err == nil {
+						break
+					}
+				}
+			}
+		} else {
+			if err := json.Unmarshal(bodyBytes, &toolResponse); err != nil {
+				t.Fatalf("Failed to unmarshal response: %v. Body: %s", err, bodyStr)
+			}
+		}
+
+		if toolResponse.Error != nil {
+			t.Errorf("Tool call failed: %v", toolResponse.Error)
+		}
+
+		// Verify the tool result indicates success
+		if result, ok := toolResponse.Result["content"].([]any); ok {
+			if len(result) > 0 {
+				if content, ok := result[0].(map[string]any); ok {
+					if text, ok := content["text"].(string); ok {
+						if text != "notification sent" {
+							t.Errorf("Expected 'notification sent', got %s", text)
+						}
+					}
+				}
+			}
+		}
+	})
+}