diff --git a/src/client/streamableHttp.test.ts b/src/client/streamableHttp.test.ts index 2dfa5869..22958582 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["_startOrAuthSse"]()).resolves.not.toThrow("Failed to open SSE stream: Method Not Allowed"); + await expect(transport["_startOrAuthSse"]({})).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["_startOrAuthSse"](); + await transport["_startOrAuthSse"]({}); // 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 { - _startOrAuthSse: (lastEventId?: string) => Promise + _startOrAuthSse: (options: { lastEventId?: string }) => Promise }; - await transportWithPrivateMethods._startOrAuthSse("test-event-id"); + await transportWithPrivateMethods._startOrAuthSse({ 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["_startOrAuthSse"](); + await transport["_startOrAuthSse"]({}); expect((actualReqInit.headers as Headers).get("x-custom-header")).toBe("CustomValue"); requestInit.headers["X-Custom-Header"] = "SecondCustomValue"; diff --git a/src/client/streamableHttp.ts b/src/client/streamableHttp.ts index 43030239..63677d55 100644 --- a/src/client/streamableHttp.ts +++ b/src/client/streamableHttp.ts @@ -1,9 +1,16 @@ -import { log } from "node:console"; import { Transport } from "../shared/transport.js"; import { isJSONRPCNotification, JSONRPCMessage, JSONRPCMessageSchema } from "../types.js"; import { auth, AuthResult, OAuthClientProvider, UnauthorizedError } from "./auth.js"; import { EventSourceParserStream } from "eventsource-parser/stream"; +// Default reconnection options for StreamableHTTP connections +const DEFAULT_STREAMABLE_HTTP_RECONNECTION_OPTIONS: StreamableHTTPReconnectionOptions = { + initialReconnectionDelay: 1000, + maxReconnectionDelay: 30000, + reconnectionDelayGrowFactor: 1.5, + maxRetries: 2, +}; + export class StreamableHTTPError extends Error { constructor( public readonly code: number | undefined, @@ -13,6 +20,16 @@ export class StreamableHTTPError extends Error { } } +/** + * Options for starting or authenticating an SSE connection + */ +export interface StartSSEOptions { + /** + * The ID of the last received event, used for resuming a disconnected stream + */ + lastEventId?: string; +} + /** * Configuration options for reconnection behavior of the StreamableHTTPClientTransport. */ @@ -37,7 +54,7 @@ export interface StreamableHTTPReconnectionOptions { /** * Maximum number of reconnection attempts before giving up. - * Default is 0 (unlimited). + * Default is 2. */ maxRetries: number; } @@ -102,8 +119,8 @@ export class StreamableHTTPClientTransport implements Transport { this._url = url; this._requestInit = opts?.requestInit; this._authProvider = opts?.authProvider; - this._reconnectionOptions = opts?.reconnectionOptions || this._defaultReconnectionOptions; this._sessionId = opts?.sessionId; + this._reconnectionOptions = opts?.reconnectionOptions ?? DEFAULT_STREAMABLE_HTTP_RECONNECTION_OPTIONS; } private async _authThenStart(): Promise { @@ -123,7 +140,7 @@ export class StreamableHTTPClientTransport implements Transport { throw new UnauthorizedError(); } - return await this._startOrAuthSse(); + return await this._startOrAuthSse({ lastEventId: undefined }); } private async _commonHeaders(): Promise { @@ -144,7 +161,9 @@ export class StreamableHTTPClientTransport implements Transport { ); } - private async _startOrAuthSse(lastEventId?: string): Promise { + + private async _startOrAuthSse(options: StartSSEOptions): Promise { + const { lastEventId } = options; try { // Try to open an initial SSE stream with GET to listen for server messages // This is optional according to the spec - server may not support it @@ -187,19 +206,9 @@ export class StreamableHTTPClientTransport implements Transport { } } - // Default reconnection options - private readonly _defaultReconnectionOptions: StreamableHTTPReconnectionOptions = { - initialReconnectionDelay: 1000, - maxReconnectionDelay: 30000, - reconnectionDelayGrowFactor: 1.5, - maxRetries: 2, - }; - - // We no longer need global reconnection state as it will be maintained per stream /** - * Calculates the next reconnection delay using exponential backoff algorithm - * with jitter for more effective reconnections in high load scenarios. + * Calculates the next reconnection delay using backoff algorithm * * @param attempt Current reconnection attempt count for the specific stream * @returns Time to wait in milliseconds before next reconnection attempt @@ -233,12 +242,11 @@ export class StreamableHTTPClientTransport implements Transport { // Calculate next delay based on current attempt count const delay = this._getNextReconnectionDelay(attemptCount); - log(`Reconnection attempt ${attemptCount + 1} in ${delay}ms...`); // Schedule the reconnection setTimeout(() => { // Use the last event ID to resume where we left off - this._startOrAuthSse(lastEventId).catch(error => { + this._startOrAuthSse({ lastEventId }).catch(error => { this.onerror?.(new Error(`Failed to reconnect SSE stream: ${error instanceof Error ? error.message : String(error)}`)); // Schedule another attempt if this one failed, incrementing the attempt counter this._scheduleReconnection(lastEventId, attemptCount + 1); @@ -253,7 +261,7 @@ export class StreamableHTTPClientTransport implements Transport { let lastEventId: string | undefined; const processStream = async () => { - // this is the closest we can get to trying to cath network errors + // this is the closest we can get to trying to catch network errors // if something happens reader will throw try { // Create a pipeline: binary stream -> text decoder -> SSE parser @@ -286,7 +294,7 @@ export class StreamableHTTPClientTransport implements Transport { } } catch (error) { // Handle stream errors - likely a network disconnect - this.onerror?.(new Error(`SSE stream disconnected: ${error instanceof Error ? error.message : String(error)}`)); + this.onerror?.(new Error(`SSE stream disconnected: ${error}`)); // Attempt to reconnect if the stream disconnects unexpectedly and we aren't closing if (this._abortController && !this._abortController.signal.aborted) { @@ -343,7 +351,7 @@ export class StreamableHTTPClientTransport implements Transport { const { lastEventId, onLastEventIdUpdate } = options ?? {}; if (lastEventId) { // If we have at last event ID, we need to reconnect the SSE stream - this._startOrAuthSse(lastEventId).catch(err => this.onerror?.(err)); + this._startOrAuthSse({ lastEventId }).catch(err => this.onerror?.(err)); return; } @@ -390,7 +398,7 @@ export class StreamableHTTPClientTransport implements Transport { // if it's supported by the server if (isJSONRPCNotification(message) && message.method === "notifications/initialized") { // Start without a lastEventId since this is a fresh connection - this._startOrAuthSse().catch(err => this.onerror?.(err)); + this._startOrAuthSse({ lastEventId: undefined }).catch(err => this.onerror?.(err)); } return; } diff --git a/src/examples/server/simpleStreamableHttp.ts b/src/examples/server/simpleStreamableHttp.ts index dc03f1db..98333730 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/integration-tests/taskResumability.test.ts b/src/integration-tests/taskResumability.test.ts index c63d3f04..cd0042ae 100644 --- a/src/integration-tests/taskResumability.test.ts +++ b/src/integration-tests/taskResumability.test.ts @@ -14,11 +14,11 @@ import { z } from 'zod'; class InMemoryEventStore implements EventStore { private events: Map = new Map(); - 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] : ''; } @@ -29,14 +29,18 @@ class InMemoryEventStore implements EventStore { return eventId; } - async getEventsAfter(lastEventId: string): Promise> { + async replayEventsAfter(lastEventId: string, + { send }: { send: (eventId: string, message: JSONRPCMessage) => Promise } + ): Promise { if (!lastEventId || !this.events.has(lastEventId)) { - return []; + return ''; } // Extract the stream ID from the event ID const streamId = this.getStreamIdFromEventId(lastEventId); - const result: Array<{ eventId: string, message: JSONRPCMessage }> = []; + if (!streamId) { + return ''; + } let foundLastEvent = false; // Sort events by eventId for chronological ordering @@ -55,11 +59,11 @@ class InMemoryEventStore implements EventStore { } if (foundLastEvent) { - result.push({ eventId, message }); + await send(eventId, message); } } - return result; + return streamId; } } 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 0487e1af..012458f2 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; } /** @@ -231,12 +216,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", @@ -247,9 +227,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); } @@ -384,10 +371,8 @@ export class StreamableHTTPServerTransport implements Transport { // check if it contains requests const hasRequests = messages.some(isJSONRPCRequest); - const hasOnlyNotificationsOrResponses = messages.every(msg => - isJSONRPCNotification(msg) || isJSONRPCResponse(msg)); - if (hasOnlyNotificationsOrResponses) { + if (!hasRequests) { // if it only contains notifications or responses, return 202 res.writeHead(202).end();