From 67a4595ae7bacd7a0a52891d0892d0a9e0feee69 Mon Sep 17 00:00:00 2001 From: ihrpr Date: Tue, 15 Apr 2025 11:03:44 +0100 Subject: [PATCH] improve event store interface --- src/client/streamableHttp.test.ts | 16 +++---- src/examples/server/simpleStreamableHttp.ts | 35 +++++++++------ src/server/streamableHttp.test.ts | 26 +++++------ src/server/streamableHttp.ts | 49 ++++++++------------- 4 files changed, 58 insertions(+), 68 deletions(-) diff --git a/src/client/streamableHttp.test.ts b/src/client/streamableHttp.test.ts index a73927ac..073cc9ac 100644 --- 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 + _startOrAuthStandaloneSSE: (options: { lastEventId?: string }) => Promise }; - 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 index a45b3158..153e35b7 100644 --- 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 { 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> { + async replayEventsAfter(lastEventId: string, + { send }: { send: (eventId: string, message: JSONRPCMessage) => Promise } + ): Promise { 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 index a03a457d..7cb062bc 100644 --- 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 { - const eventId = this.generateEventId(streamId); + const eventId = `${streamId}_${randomUUID()}`; storedEvents.set(eventId, { eventId, message }); return eventId; }, - async getEventsAfter(lastEventId: string): Promise> { - 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 + }): Promise { + 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 index 37dc2479..58d0fe14 100644 --- 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; + storeEvent(streamId: StreamId, message: JSONRPCMessage): Promise; - /** - * 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>; - - /** - * 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 + }): Promise; } /** @@ -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 = { "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); }