improve event store interface

ihrpr · ihrpr · commit 67a4595ae7ba · 2025-04-15T11:03:44.000+01:00
diff --git a/src/client/streamableHttp.test.ts b/src/client/streamableHttp.test.ts
@@ -164,7 +164,7 @@ describe("StreamableHTTPClientTransport", () => {
     // We expect the 405 error to be caught and handled gracefully
     // This should not throw an error that breaks the transport
     await transport.start();
-    await expect(transport["_startOrAuthStandaloneSSE"]()).resolves.not.toThrow("Failed to open SSE stream: Method Not Allowed");
+    await expect(transport["_startOrAuthStandaloneSSE"]({})).resolves.not.toThrow("Failed to open SSE stream: Method Not Allowed");
     // Check that GET was attempted
     expect(global.fetch).toHaveBeenCalledWith(
       expect.anything(),
@@ -208,7 +208,7 @@ describe("StreamableHTTPClientTransport", () => {
     transport.onmessage = messageSpy;
 
     await transport.start();
-    await transport["_startOrAuthStandaloneSSE"]();
+    await transport["_startOrAuthStandaloneSSE"]({});
 
     // Give time for the SSE event to be processed
     await new Promise(resolve => setTimeout(resolve, 50));
@@ -289,10 +289,10 @@ describe("StreamableHTTPClientTransport", () => {
     // Verify options were set correctly (checking implementation details)
     // Access private properties for testing
     const transportInstance = transport as unknown as {
-      _defaultReconnectionOptions: StreamableHTTPReconnectionOptions;
+      _reconnectionOptions: StreamableHTTPReconnectionOptions;
     };
-    expect(transportInstance._defaultReconnectionOptions.initialReconnectionDelay).toBe(1000);
-    expect(transportInstance._defaultReconnectionOptions.maxRetries).toBe(2);
+    expect(transportInstance._reconnectionOptions.initialReconnectionDelay).toBe(500);
+    expect(transportInstance._reconnectionOptions.maxRetries).toBe(5);
   });
 
   it("should pass lastEventId when reconnecting", async () => {
@@ -313,9 +313,9 @@ describe("StreamableHTTPClientTransport", () => {
     await transport.start();
     // Type assertion to access private method
     const transportWithPrivateMethods = transport as unknown as {
-      _startOrAuthStandaloneSSE: (lastEventId?: string) => Promise<void>
+      _startOrAuthStandaloneSSE: (options: { lastEventId?: string }) => Promise<void>
     };
-    await transportWithPrivateMethods._startOrAuthStandaloneSSE("test-event-id");
+    await transportWithPrivateMethods._startOrAuthStandaloneSSE({ lastEventId: "test-event-id" });
 
     // Verify fetch was called with the lastEventId header
     expect(fetchSpy).toHaveBeenCalled();
@@ -382,7 +382,7 @@ describe("StreamableHTTPClientTransport", () => {
 
     await transport.start();
 
-    await transport["_startOrAuthStandaloneSSE"]();
+    await transport["_startOrAuthStandaloneSSE"]({});
     expect((actualReqInit.headers as Headers).get("x-custom-header")).toBe("CustomValue");
 
     requestInit.headers["X-Custom-Header"] = "SecondCustomValue";
diff --git a/src/examples/server/simpleStreamableHttp.ts b/src/examples/server/simpleStreamableHttp.ts
@@ -12,18 +12,18 @@ class InMemoryEventStore implements EventStore {
   /**
    * Generates a unique event ID for a given stream ID
    */
-  generateEventId(streamId: string): string {
+  private generateEventId(streamId: string): string {
     return `${streamId}_${Date.now()}_${Math.random().toString(36).substring(2, 10)}`;
   }
 
-  getStreamIdFromEventId(eventId: string): string {
+  private getStreamIdFromEventId(eventId: string): string {
     const parts = eventId.split('_');
     return parts.length > 0 ? parts[0] : '';
   }
 
-
   /**
    * Stores an event with a generated event ID
+   * Implements EventStore.storeEvent
    */
   async storeEvent(streamId: string, message: JSONRPCMessage): Promise<string> {
     const eventId = this.generateEventId(streamId);
@@ -33,17 +33,26 @@ class InMemoryEventStore implements EventStore {
   }
 
   /**
-   * Retrieves events that occurred after a specific event
+   * Replays events that occurred after a specific event ID
+   * Implements EventStore.replayEventsAfter
    */
-  async getEventsAfter(lastEventId: string): Promise<Array<{ eventId: string, message: JSONRPCMessage }>> {
+  async replayEventsAfter(lastEventId: string, 
+    { send }: { send: (eventId: string, message: JSONRPCMessage) => Promise<void> }
+  ): Promise<string> {
     if (!lastEventId || !this.events.has(lastEventId)) {
-      return [];
+      console.log(`No events found for lastEventId: ${lastEventId}`);
+      return '';
     }
 
     // Extract the stream ID from the event ID
-    const streamId = lastEventId.split('_')[0];
-    const result: Array<{ eventId: string, message: JSONRPCMessage }> = [];
+    const streamId = this.getStreamIdFromEventId(lastEventId);
+    if (!streamId) {
+      console.log(`Could not extract streamId from lastEventId: ${lastEventId}`);
+      return '';
+    }
+
     let foundLastEvent = false;
+    let eventCount = 0;
 
     // Sort events by eventId for chronological ordering
     const sortedEvents = [...this.events.entries()].sort((a, b) => a[0].localeCompare(b[0]));
@@ -54,21 +63,21 @@ class InMemoryEventStore implements EventStore {
         continue;
       }
 
-      // Start collecting events after we find the lastEventId
+      // Start sending events after we find the lastEventId
       if (eventId === lastEventId) {
         foundLastEvent = true;
         continue;
       }
 
       if (foundLastEvent) {
-        result.push({ eventId, message });
+        await send(eventId, message);
+        eventCount++;
       }
     }
 
-    console.log(`Found ${result.length} events after ${lastEventId} for replay`);
-    return result;
+    console.log(`Replayed ${eventCount} events after ${lastEventId} for stream ${streamId}`);
+    return streamId;
   }
-
 }
 
 // Create an MCP server with implementation details
diff --git a/src/server/streamableHttp.test.ts b/src/server/streamableHttp.test.ts
@@ -1,7 +1,7 @@
 import { createServer, type Server, IncomingMessage, ServerResponse } from "node:http";
 import { AddressInfo } from "node:net";
 import { randomUUID } from "node:crypto";
-import { EventStore, StreamableHTTPServerTransport } from "./streamableHttp.js";
+import { EventStore, StreamableHTTPServerTransport, EventId, StreamId } from "./streamableHttp.js";
 import { McpServer } from "./mcp.js";
 import { CallToolResult, JSONRPCMessage } from "../types.js";
 import { z } from "zod";
@@ -912,31 +912,25 @@ describe("StreamableHTTPServerTransport with resumability", () => {
 
   // Simple implementation of EventStore
   const eventStore: EventStore = {
-    generateEventId(streamId: string): string {
-      return `${streamId}_${randomUUID()}`;
-    },
-    getStreamIdFromEventId(eventId: string): string {
-      return eventId.split('_')[0]; // Extract stream ID from the event ID
-    },
+
     async storeEvent(streamId: string, message: JSONRPCMessage): Promise<string> {
-      const eventId = this.generateEventId(streamId);
+      const eventId = `${streamId}_${randomUUID()}`;
       storedEvents.set(eventId, { eventId, message });
       return eventId;
     },
 
-    async getEventsAfter(lastEventId: string): Promise<Array<{ eventId: string, message: JSONRPCMessage }>> {
-      const streamId = lastEventId.split('_')[0]; // Extract stream ID from the event ID
-      const result: Array<{ eventId: string, message: JSONRPCMessage }> = [];
-
+    async replayEventsAfter(lastEventId: EventId, { send }: {
+      send: (eventId: EventId, message: JSONRPCMessage) => Promise<void>
+    }): Promise<StreamId> {
+      const streamId = lastEventId.split('_')[0];
+      // Extract stream ID from the event ID
       // For test simplicity, just return all events with matching streamId that aren't the lastEventId
-      // This avoids issues with event ordering in tests
       for (const [eventId, { message }] of storedEvents.entries()) {
         if (eventId.startsWith(streamId) && eventId !== lastEventId) {
-          result.push({ eventId, message });
+          await send(eventId, message);
         }
       }
-
-      return result;
+      return streamId;
     },
   };
 
diff --git a/src/server/streamableHttp.ts b/src/server/streamableHttp.ts
@@ -1,45 +1,30 @@
 import { IncomingMessage, ServerResponse } from "node:http";
 import { Transport } from "../shared/transport.js";
-import { isJSONRPCNotification, isJSONRPCRequest, isJSONRPCResponse, JSONRPCMessage, JSONRPCMessageSchema, RequestId } from "../types.js";
+import { isJSONRPCRequest, isJSONRPCResponse, JSONRPCMessage, JSONRPCMessageSchema, RequestId } from "../types.js";
 import getRawBody from "raw-body";
 import contentType from "content-type";
 import { randomUUID } from "node:crypto";
 
 const MAXIMUM_MESSAGE_SIZE = "4mb";
 
+export type StreamId = string;
+export type EventId = string;
+
 /**
  * Interface for resumability support via event storage
  */
 export interface EventStore {
-  /**
-   * Generates a unique event ID for a given stream ID
-   * @param streamId The stream ID to include in the event ID
-   * @returns A unique event ID that includes the stream ID
-   */
-  generateEventId(streamId: string): string;
-
   /**
    * Stores an event for later retrieval
    * @param streamId ID of the stream the event belongs to
    * @param message The JSON-RPC message to store
    * @returns The generated event ID for the stored event
    */
-  storeEvent(streamId: string, message: JSONRPCMessage): Promise<string>;
+  storeEvent(streamId: StreamId, message: JSONRPCMessage): Promise<EventId>;
 
-  /**
-   * Retrieves events for a stream starting from a given event ID
-   * @param lastEventId The event ID to start from
-   * @returns Array of stored events with their event IDs
-   */
-  getEventsAfter(lastEventId: string): Promise<Array<{ eventId: string, message: JSONRPCMessage }>>;
-
-  /**
-   * Extracts the stream ID from an event ID
-   * This is necessary for resumability to identify which stream the event belongs to
-   * @param eventId The event ID to extract stream ID from
-   * @returns The stream ID portion of the event ID
-   */
-  getStreamIdFromEventId(eventId: string): string;
+  replayEventsAfter(lastEventId: EventId, { send }: {
+    send: (eventId: EventId, message: JSONRPCMessage) => Promise<void>
+  }): Promise<StreamId>;
 }
 
 /**
@@ -232,12 +217,7 @@ export class StreamableHTTPServerTransport implements Transport {
     if (!this._eventStore) {
       return;
     }
-
     try {
-      const events = await this._eventStore.getEventsAfter(lastEventId);
-      const streamId = this._eventStore.getStreamIdFromEventId(lastEventId);
-
-      this._streamMapping.set(streamId, res);
       const headers: Record<string, string> = {
         "Content-Type": "text/event-stream",
         "Cache-Control": "no-cache, no-transform",
@@ -248,9 +228,16 @@ export class StreamableHTTPServerTransport implements Transport {
         headers["mcp-session-id"] = this.sessionId;
       }
       res.writeHead(200, headers).flushHeaders();
-      for (const { eventId, message } of events) {
-        this.writeSSEEvent(res, message, eventId);
-      }
+
+      const streamId = await this._eventStore?.replayEventsAfter(lastEventId, {
+        send: async (eventId: string, message: JSONRPCMessage) => {
+          if (!this.writeSSEEvent(res, message, eventId)) {
+            this.onerror?.(new Error("Failed replay events"));
+            res.end();
+          }
+        }
+      });
+      this._streamMapping.set(streamId, res);
     } catch (error) {
       this.onerror?.(error as Error);
     }
