refactor: improve async handling and response management in StreamableHTTPServerTransport

Author: chrisribe

src/transports/streamable-http.ts

@@ -2,7 +2,7 @@
  * StreamableHTTPServerTransport - MCP transport over HTTP
  * 
  * Implements the MCP Server transport interface using HTTP instead of stdio.
- * Supports multiple clients connecting via HTTP POST requests.
+ * Properly handles the async nature of MCP over synchronous HTTP.
  */
 
 import { createServer, IncomingMessage, ServerResponse } from 'http';
@@ -12,81 +12,32 @@ import { debugLog } from '../utils/debug.js';
 interface HTTPTransportOptions {
   port?: number;
   host?: string;
-}
-
-/**
- * Simple request-scoped response buffer for storing responses
- * Each HTTP request gets its own buffer to collect responses
- */
-class ResponseBuffer {
-  private responses: JSONRPCMessage[] = [];
-
-  add(message: JSONRPCMessage): void {
-    this.responses.push(message);
-  }
-
-  getAll(): JSONRPCMessage[] {
-    return this.responses;
-  }
-}
-
-// Thread-local style buffer (request-scoped)
-let currentBuffer: ResponseBuffer | null = null;
-
-/**
- * Helper functions for HTTP responses
- */
-function sendJSON(res: ServerResponse, statusCode: number, data: any): void {
-  res.writeHead(statusCode, { 'Content-Type': 'application/json' });
-  res.end(JSON.stringify(data));
-}
-
-function sendError(res: ServerResponse, statusCode: number, code: number, message: string, id: any = null): void {
-  sendJSON(res, statusCode, {
-    jsonrpc: '2.0',
-    error: { code, message },
-    id
-  });
-}
-
-function sendSuccess(res: ServerResponse, data: any): void {
-  sendJSON(res, 200, data);
-}
-
-function sendNotFound(res: ServerResponse): void {
-  sendJSON(res, 404, { error: 'Not found' });
-}
-
-function sendPayloadTooLarge(res: ServerResponse): void {
-  sendJSON(res, 413, {
-    error: 'Payload too large',
-    code: -32600
-  });
-}
-
-function sendAccepted(res: ServerResponse): void {
-  res.writeHead(202);
-  res.end();
+  timeout?: number; // Request timeout in ms
 }
 
 export class StreamableHTTPServerTransport {
   private port: number;
   private host: string;
+  private timeout: number;
   private httpServer: any = null;
-
-  // Transport interface properties (standard MCP interface)
+  
+  // Map request IDs to their response handlers
+  private pendingResponses = new Map<string | number, {
+    resolve: (message: JSONRPCMessage) => void;
+    timer: NodeJS.Timeout;
+  }>();
+
+  // Transport interface properties
   onmessage?: (message: JSONRPCMessage) => void;
   onclose?: () => void;
   onerror?: (error: Error) => void;
 
   constructor(options: HTTPTransportOptions = {}) {
     this.port = options.port || 3000;
     this.host = options.host || 'localhost';
+    this.timeout = options.timeout || 30000; // 30 second default
   }
 
-  /**
-   * Start the HTTP server
-   */
   async start(): Promise<void> {
     if (this.httpServer) {
       throw new Error('Server already started');
@@ -113,10 +64,18 @@ export class StreamableHTTPServerTransport {
     });
   }
 
-  /**
-   * Close the HTTP server
-   */
   async close(): Promise<void> {
+    // Clear all pending responses
+    for (const [id, pending] of this.pendingResponses) {
+      clearTimeout(pending.timer);
+      pending.resolve({
+        jsonrpc: '2.0',
+        error: { code: -32603, message: 'Server shutting down' },
+        id
+      });
+    }
+    this.pendingResponses.clear();
+
     return new Promise((resolve) => {
       if (this.httpServer) {
         this.httpServer.close(() => {
@@ -130,23 +89,25 @@ export class StreamableHTTPServerTransport {
   }
 
   /**
-   * Send a message to the client
-   * In HTTP mode, responses are buffered and sent in the response body
+   * Send a message (response) back to the waiting HTTP request
    */
   async send(message: JSONRPCMessage): Promise<void> {
     debugLog('[HTTP Transport] Sending message:', message);
 
-    // If we're in a request context, buffer the response
-    if (currentBuffer) {
-      currentBuffer.add(message);
-    } else {
-      debugLog('[HTTP Transport] Warning: send() called outside of request context');
+    // Check if this is a response to a pending request
+    if ('id' in message && message.id !== null) {
+      const pending = this.pendingResponses.get(message.id);
+      if (pending) {
+        clearTimeout(pending.timer);
+        this.pendingResponses.delete(message.id);
+        pending.resolve(message);
+      } else {
+        debugLog('[HTTP Transport] No pending request for response:', message.id);
+      }
     }
+    // Notifications don't need handling in HTTP context
   }
 
-  /**
-   * Handle incoming HTTP requests
-   */
   private async handleRequest(req: IncomingMessage, res: ServerResponse): Promise<void> {
     // Add CORS headers
     res.setHeader('Access-Control-Allow-Origin', '*');
@@ -162,7 +123,7 @@ export class StreamableHTTPServerTransport {
 
     // Health check endpoint
     if (req.url === '/health' && req.method === 'GET') {
-      sendSuccess(res, { status: 'ok', transport: 'http' });
+      this.sendJSON(res, 200, { status: 'ok', transport: 'http' });
       return;
     }
 
@@ -173,80 +134,113 @@ export class StreamableHTTPServerTransport {
     }
 
     // Unknown endpoint
-    sendNotFound(res);
+    this.sendJSON(res, 404, { error: 'Not found' });
   }
 
-  /**
-   * Handle MCP protocol POST requests
-   */
   private async handleMCPRequest(req: IncomingMessage, res: ServerResponse): Promise<void> {
-    let body = '';
+    try {
+      const body = await this.readBody(req);
+      const request = this.parseRequest(body);
+      
+      if (!request) {
+        this.sendError(res, 400, -32700, 'Parse error');
+        return;
+      }
 
-    // Read request body
-    req.on('data', (chunk: Buffer) => {
-        body += chunk.toString('utf-8');
+      debugLog('[HTTP Transport] Received request:', request);
 
-        // Prevent body from growing too large
-        if (body.length > 1024 * 1024) {
-          req.removeAllListeners('data');
-          sendPayloadTooLarge(res);
+      // Handle based on whether it's a request or notification
+      if ('id' in request && request.id !== null) {
+        // Request - expects a response
+        const response = await this.waitForResponse(request);
+        this.sendJSON(res, 200, response);
+      } else {
+        // Notification - no response expected
+        if (this.onmessage) {
+          this.onmessage(request);
         }
-      });
+        res.writeHead(202);
+        res.end();
+      }
+    } catch (error: any) {
+      debugLog('[HTTP Transport] Request error:', error);
+      if (!res.headersSent) {
+        this.sendError(res, 500, -32603, error.message || 'Internal server error');
+      }
+    }
+  }
+
+  private async waitForResponse(request: JSONRPCMessage): Promise<JSONRPCMessage> {
+    return new Promise((resolve) => {
+      // Type narrowing - we know this request has an id because we checked before calling this method
+      if (!('id' in request) || request.id === null || request.id === undefined) {
+        resolve({
+          jsonrpc: '2.0',
+          error: { code: -32600, message: 'Invalid request - missing id' }
+        } as JSONRPCMessage);
+        return;
+      }
+      
+      const id = request.id;
+      
+      // Set up timeout
+      const timer = setTimeout(() => {
+        this.pendingResponses.delete(id);
+        resolve({
+          jsonrpc: '2.0',
+          error: { code: -32603, message: 'Request timeout' },
+          id
+        } as JSONRPCMessage);
+      }, this.timeout);
+
+      // Store the response handler
+      this.pendingResponses.set(id, { resolve, timer });
+
+      // Process the request
+      if (this.onmessage) {
+        this.onmessage(request);
+      }
+    });
+  }
 
-      req.on('end', async () => {
-        try {
-          // Parse JSON-RPC request
-          let request: JSONRPCMessage;
-
-          try {
-            request = JSON.parse(body);
-            debugLog('[HTTP Transport] Received request:', request);
-          } catch (parseError) {
-            debugLog('[HTTP Transport] Parse error:', parseError);
-            sendError(res, 400, -32700, 'Parse error');
-            return;
-          }
-
-          // Create response buffer for this request
-          const buffer = new ResponseBuffer();
-          currentBuffer = buffer;
-
-          try {
-            // Call the message handler (set by MCP Server)
-            if (this.onmessage) {
-              this.onmessage(request);
-              
-              // Wait for handler to process and buffer responses
-              // TODO: Replace with proper async handling
-              await new Promise(resolve => setTimeout(resolve, 100));
-            }
-
-            // Send buffered responses
-            const responses = buffer.getAll();
-
-            if (responses.length > 0) {
-              sendJSON(res, 200, responses.length === 1 ? responses[0] : responses);
-            } else {
-              // No response (notification or async processing)
-              sendAccepted(res);
-            }
-          } finally {
-            currentBuffer = null;
-          }
-        } catch (error: any) {
-          debugLog('[HTTP Transport] Request error:', error);
-
-          if (!res.headersSent) {
-            sendError(res, 500, -32603, error.message || 'Internal server error');
-          }
+  private async readBody(req: IncomingMessage): Promise<string> {
+    return new Promise((resolve, reject) => {
+      let body = '';
+      
+      req.on('data', (chunk: Buffer) => {
+        body += chunk.toString('utf-8');
+        
+        // Prevent body from growing too large (1MB limit)
+        if (body.length > 1024 * 1024) {
+          req.removeAllListeners();
+          reject(new Error('Payload too large'));
         }
       });
 
-    req.on('error', (error: Error) => {
-      debugLog('[HTTP Transport] Request error:', error);
-      if (!res.headersSent) {
-        sendError(res, 400, -32600, 'Invalid request');
-      }
+      req.on('end', () => resolve(body));
+      req.on('error', reject);
+    });
+  }
+
+  private parseRequest(body: string): JSONRPCMessage | null {
+    try {
+      return JSON.parse(body);
+    } catch {
+      return null;
+    }
+  }
+
+  // Helper methods
+  private sendJSON(res: ServerResponse, statusCode: number, data: any): void {
+    res.writeHead(statusCode, { 'Content-Type': 'application/json' });
+    res.end(JSON.stringify(data));
+  }
+
+  private sendError(res: ServerResponse, statusCode: number, code: number, message: string, id: any = null): void {
+    this.sendJSON(res, statusCode, {
+      jsonrpc: '2.0',
+      error: { code, message },
+      id
     });
   }
 }