Merge branch 'ihrpr/fix-streamable-connection-close' into ihrpr/examples

src/server/streamableHttp.test.ts

@@ -842,8 +842,8 @@ describe("StreamableHTTPServerTransport", () => {
         },
         body: JSON.stringify(initMessage),
       });
-
-      await transport.handleRequest(initReq, mockResponse);
+      const initResponse = createMockResponse();
+      await transport.handleRequest(initReq, initResponse);
       mockResponse.writeHead.mockClear();
     });
 
@@ -934,17 +934,17 @@ describe("StreamableHTTPServerTransport", () => {
       // Now stream should be closed
       expect(mockResponse.end).toHaveBeenCalled();
     });
-    
+
     it("should keep stream open when multiple requests share the same connection", async () => {
       // Create a fresh response for this test  
       const sharedResponse = createMockResponse();
-      
+
       // Send two requests in a batch that will share the same connection
       const batchRequests: JSONRPCMessage[] = [
         { jsonrpc: "2.0", method: "method1", params: {}, id: "req1" },
         { jsonrpc: "2.0", method: "method2", params: {}, id: "req2" }
       ];
-      
+
       const req = createMockRequest({
         method: "POST",
         headers: {
@@ -954,40 +954,40 @@ describe("StreamableHTTPServerTransport", () => {
         },
         body: JSON.stringify(batchRequests)
       });
-      
+
       await transport.handleRequest(req, sharedResponse);
-      
+
       // Respond to first request
       const response1: JSONRPCMessage = {
-        jsonrpc: "2.0", 
+        jsonrpc: "2.0",
         result: { value: "result1" },
         id: "req1"
       };
-      
+
       await transport.send(response1);
-      
+
       // Connection should remain open because req2 is still pending
       expect(sharedResponse.write).toHaveBeenCalledWith(
         expect.stringContaining(`event: message\ndata: ${JSON.stringify(response1)}\n\n`)
       );
       expect(sharedResponse.end).not.toHaveBeenCalled();
-      
+
       // Respond to second request
       const response2: JSONRPCMessage = {
         jsonrpc: "2.0",
         result: { value: "result2" },
         id: "req2"
       };
-      
+
       await transport.send(response2);
-      
+
       // Now connection should close as all requests are complete
       expect(sharedResponse.write).toHaveBeenCalledWith(
-        expect.stringContaining(`event: message\ndata: ${JSON.stringify(response2)}\n\n`)  
+        expect.stringContaining(`event: message\ndata: ${JSON.stringify(response2)}\n\n`)
       );
       expect(sharedResponse.end).toHaveBeenCalled();
     });
-    
+
     it("should clean up connection tracking when a response is sent", async () => {
       const req = createMockRequest({
         method: "POST",
@@ -997,32 +997,32 @@ describe("StreamableHTTPServerTransport", () => {
           "mcp-session-id": transport.sessionId
         },
         body: JSON.stringify({
-          jsonrpc: "2.0", 
-          method: "test", 
-          params: {}, 
+          jsonrpc: "2.0",
+          method: "test",
+          params: {},
           id: "cleanup-test"
         })
       });
-      
+
       const response = createMockResponse();
       await transport.handleRequest(req, response);
-      
+
       // Verify that the request is tracked in the SSE map
-      expect(transport["_sseResponseMapping"].size).toBe(1);
+      expect(transport["_sseResponseMapping"].size).toBe(2);
       expect(transport["_sseResponseMapping"].has("cleanup-test")).toBe(true);
-      
+
       // Send a response
       await transport.send({
         jsonrpc: "2.0",
         result: {},
         id: "cleanup-test"
       });
-      
+
       // Verify that the mapping was cleaned up
-      expect(transport["_sseResponseMapping"].size).toBe(0);
+      expect(transport["_sseResponseMapping"].size).toBe(1);
       expect(transport["_sseResponseMapping"].has("cleanup-test")).toBe(false);
     });
-    
+
     it("should clean up connection tracking when client disconnects", async () => {
       // Setup two requests that share a connection
       const req = createMockRequest({
@@ -1034,12 +1034,12 @@ describe("StreamableHTTPServerTransport", () => {
         },
         body: JSON.stringify([
           { jsonrpc: "2.0", method: "longRunning1", params: {}, id: "req1" },
-          { jsonrpc: "2.0", method: "longRunning2", params: {}, id: "req2" } 
+          { jsonrpc: "2.0", method: "longRunning2", params: {}, id: "req2" }
         ])
       });
-      
+
       const response = createMockResponse();
-      
+
       // We need to manually store the callback to trigger it later
       let closeCallback: (() => void) | undefined;
       response.on.mockImplementation((event, callback: () => void) => {
@@ -1048,19 +1048,19 @@ describe("StreamableHTTPServerTransport", () => {
         }
         return response;
       });
-      
+
       await transport.handleRequest(req, response);
-      
+
       // Both requests should be mapped to the same response
-      expect(transport["_sseResponseMapping"].size).toBe(2);
+      expect(transport["_sseResponseMapping"].size).toBe(3);
       expect(transport["_sseResponseMapping"].get("req1")).toBe(response);
       expect(transport["_sseResponseMapping"].get("req2")).toBe(response);
-      
+
       // Simulate client disconnect by triggering the stored callback
       if (closeCallback) closeCallback();
-      
+
       // All entries using this response should be removed
-      expect(transport["_sseResponseMapping"].size).toBe(0);
+      expect(transport["_sseResponseMapping"].size).toBe(1);
       expect(transport["_sseResponseMapping"].has("req1")).toBe(false);
       expect(transport["_sseResponseMapping"].has("req2")).toBe(false);
     });