From 43cb015844a92c6a08bb5914aee78c2615d4a9cd Mon Sep 17 00:00:00 2001 From: ihrpr Date: Sun, 13 Apr 2025 21:40:15 +0100 Subject: [PATCH] resume long running requst from new client --- src/client/streamableHttp.ts | 25 +++++++++++++++--- src/examples/client/simpleStreamableHttp.ts | 16 ++++++++++-- src/examples/server/simpleStreamableHttp.ts | 20 +++++++++------ src/server/streamableHttp.ts | 28 +++++++++------------ src/shared/protocol.ts | 18 ++++++++++--- src/shared/transport.ts | 2 +- 6 files changed, 76 insertions(+), 33 deletions(-) diff --git a/src/client/streamableHttp.ts b/src/client/streamableHttp.ts index 9451e8d9..77acc9c4 100644 --- a/src/client/streamableHttp.ts +++ b/src/client/streamableHttp.ts @@ -71,6 +71,11 @@ export type StreamableHTTPClientTransportOptions = { * Options to configure the reconnection behavior. */ reconnectionOptions?: StreamableHTTPReconnectionOptions; + /** + * Session ID for the connection. This is used to identify the session on the server. + * When not provided and connecting to a server that supports session IDs, the server will generate a new session ID. + */ + sessionId?: string; }; /** @@ -98,6 +103,7 @@ export class StreamableHTTPClientTransport implements Transport { this._requestInit = opts?.requestInit; this._authProvider = opts?.authProvider; this._reconnectionOptions = opts?.reconnectionOptions || this._defaultReconnectionOptions; + this._sessionId = opts?.sessionId; } private async _authThenStart(): Promise { @@ -240,7 +246,7 @@ export class StreamableHTTPClientTransport implements Transport { }, delay); } - private _handleSseStream(stream: ReadableStream | null): void { + private _handleSseStream(stream: ReadableStream | null, onLastEventIdUpdate?: (event: string) => void): void { if (!stream) { return; } @@ -266,6 +272,7 @@ export class StreamableHTTPClientTransport implements Transport { // Update last event ID if provided if (event.id) { lastEventId = event.id; + onLastEventIdUpdate?.(lastEventId); } if (!event.event || event.event === "message") { @@ -330,8 +337,16 @@ export class StreamableHTTPClientTransport implements Transport { this.onclose?.(); } - async send(message: JSONRPCMessage | JSONRPCMessage[]): Promise { + async send(message: JSONRPCMessage | JSONRPCMessage[], options?: { lastEventId?: string, onLastEventIdUpdate?: (event: string) => void }): Promise { try { + // If client passes in a lastEventId in the request options, we need to reconnect the SSE stream + const { lastEventId, onLastEventIdUpdate } = options ?? {}; + if (lastEventId) { + // If we have at last event ID, we need to reconnect the SSE stream + this._startOrAuthStandaloneSSE(lastEventId).catch(err => this.onerror?.(err)); + return; + } + const headers = await this._commonHeaders(); headers.set("content-type", "application/json"); headers.set("accept", "application/json, text/event-stream"); @@ -393,7 +408,7 @@ export class StreamableHTTPClientTransport implements Transport { // Handle SSE stream responses for requests // We use the same handler as standalone streams, which now supports // reconnection with the last event ID - this._handleSseStream(response.body); + this._handleSseStream(response.body, onLastEventIdUpdate); } else if (contentType?.includes("application/json")) { // For non-streaming servers, we might get direct JSON responses const data = await response.json(); @@ -416,4 +431,8 @@ export class StreamableHTTPClientTransport implements Transport { throw error; } } + + get sessionId(): string | undefined { + return this._sessionId; + } } diff --git a/src/examples/client/simpleStreamableHttp.ts b/src/examples/client/simpleStreamableHttp.ts index 923ffbc2..922aa0ce 100644 --- a/src/examples/client/simpleStreamableHttp.ts +++ b/src/examples/client/simpleStreamableHttp.ts @@ -29,6 +29,8 @@ let notificationCount = 0; let client: Client | null = null; let transport: StreamableHTTPClientTransport | null = null; let serverUrl = 'http://localhost:3000/mcp'; +let notificationsToolLastEventId: string | undefined = undefined; +let sessionId: string | undefined = undefined; async function main(): Promise { console.log('MCP Interactive Client'); @@ -186,7 +188,10 @@ async function connect(url?: string): Promise { } transport = new StreamableHTTPClientTransport( - new URL(serverUrl) + new URL(serverUrl), + { + sessionId: sessionId + } ); // Set up notification handlers @@ -218,6 +223,8 @@ async function connect(url?: string): Promise { // Connect the client await client.connect(transport); + sessionId = transport.sessionId + console.log('Transport created with session ID:', sessionId); console.log('Connected to MCP server'); } catch (error) { console.error('Failed to connect:', error); @@ -291,7 +298,12 @@ async function callTool(name: string, args: Record): Promise { + notificationsToolLastEventId = event; + }; + const result = await client.request(request, CallToolResultSchema, { + lastEventId: notificationsToolLastEventId, onLastEventIdUpdate + }); console.log('Tool result:'); result.content.forEach(item => { diff --git a/src/examples/server/simpleStreamableHttp.ts b/src/examples/server/simpleStreamableHttp.ts index a45b3158..dc03f1db 100644 --- a/src/examples/server/simpleStreamableHttp.ts +++ b/src/examples/server/simpleStreamableHttp.ts @@ -172,14 +172,18 @@ server.tool( while (count === 0 || counter < count) { counter++; - await sendNotification({ - method: "notifications/message", - params: { - level: "info", - data: `Periodic notification #${counter} at ${new Date().toISOString()}` - } - }); - + try { + await sendNotification({ + method: "notifications/message", + params: { + level: "info", + data: `Periodic notification #${counter} at ${new Date().toISOString()}` + } + }); + } + catch (error) { + console.error("Error sending notification:", error); + } // Wait for the specified interval await sleep(interval); } diff --git a/src/server/streamableHttp.ts b/src/server/streamableHttp.ts index b8be6d15..7740a08f 100644 --- a/src/server/streamableHttp.ts +++ b/src/server/streamableHttp.ts @@ -217,7 +217,6 @@ export class StreamableHTTPServerTransport implements Transport { // Assign the response to the standalone SSE stream this._streamMapping.set(this._standaloneSseStreamId, res); - // Set up close handler for client disconnects res.on("close", () => { this._streamMapping.delete(this._standaloneSseStreamId); @@ -345,8 +344,10 @@ export class StreamableHTTPServerTransport implements Transport { const isInitializationRequest = messages.some( msg => 'method' in msg && msg.method === 'initialize' ); + const mcpSessionId = req.headers["mcp-session-id"] as string | undefined; if (isInitializationRequest) { - if (this._initialized) { + // if generateSessionId is not set, the server does not support session management + if (this._initialized && this.sessionId !== undefined && mcpSessionId !== this.sessionId) { res.writeHead(400).end(JSON.stringify({ jsonrpc: "2.0", error: { @@ -368,7 +369,7 @@ export class StreamableHTTPServerTransport implements Transport { })); return; } - this.sessionId = this.sessionIdGenerator(); + this.sessionId = mcpSessionId ?? this.sessionIdGenerator(); this._initialized = true; } @@ -419,17 +420,8 @@ export class StreamableHTTPServerTransport implements Transport { this._requestToStreamMapping.set(message.id, streamId); } } - // Set up close handler for client disconnects res.on("close", () => { - // find a stream ID for this response - // Remove all entries that reference this response - for (const [id, stream] of this._requestToStreamMapping.entries()) { - if (streamId === stream) { - this._requestToStreamMapping.delete(id); - this._requestResponseMap.delete(id); - } - } this._streamMapping.delete(streamId); }); @@ -577,7 +569,7 @@ export class StreamableHTTPServerTransport implements Transport { // Get the response for this request const streamId = this._requestToStreamMapping.get(requestId); const response = this._streamMapping.get(streamId!); - if (!streamId || !response) { + if (!streamId) { throw new Error(`No connection established for request ID: ${String(requestId)}`); } @@ -588,9 +580,10 @@ export class StreamableHTTPServerTransport implements Transport { if (this._eventStore) { eventId = await this._eventStore.storeEvent(streamId, message); } - - // Write the event to the response stream - this.writeSSEEvent(response, message, eventId); + if (response) { + // Write the event to the response stream + this.writeSSEEvent(response, message, eventId); + } } if (isJSONRPCResponse(message)) { @@ -603,6 +596,9 @@ export class StreamableHTTPServerTransport implements Transport { const allResponsesReady = relatedIds.every(id => this._requestResponseMap.has(id)); if (allResponsesReady) { + if (!response) { + throw new Error(`No connection established for request ID: ${String(requestId)}`); + } if (this._enableJsonResponse) { // All responses ready, send as JSON const headers: Record = { diff --git a/src/shared/protocol.ts b/src/shared/protocol.ts index 07bfc02c..6193f5f1 100644 --- a/src/shared/protocol.ts +++ b/src/shared/protocol.ts @@ -87,6 +87,18 @@ export type RequestOptions = { * May be used to indicate to the transport which incoming request to associate this outgoing request with. */ relatedRequestId?: RequestId; + + /** + * May be used to indicate to the transport which last event ID to associate this outgoing request with. + * This is used to resume a long-running requests that may have been interrupted and a new instance of a client is being created. + */ + lastEventId?: string; + + /** + * A callback that is invoked when the last event ID is updated. + * This is used to notidy the client that the last event ID has changed, so that client can update its state accordingly. + */ + onLastEventIdUpdate?: (event: string) => void; }; /** @@ -501,7 +513,7 @@ export abstract class Protocol< resultSchema: T, options?: RequestOptions, ): Promise> { - const { relatedRequestId } = options ?? {}; + const { relatedRequestId, lastEventId, onLastEventIdUpdate } = options ?? {}; return new Promise((resolve, reject) => { if (!this._transport) { @@ -543,7 +555,7 @@ export abstract class Protocol< requestId: messageId, reason: String(reason), }, - }, { relatedRequestId }) + }, { relatedRequestId, lastEventId, onLastEventIdUpdate }) .catch((error) => this._onerror(new Error(`Failed to send cancellation: ${error}`)), ); @@ -581,7 +593,7 @@ export abstract class Protocol< this._setupTimeout(messageId, timeout, options?.maxTotalTimeout, timeoutHandler, options?.resetTimeoutOnProgress ?? false); - this._transport.send(jsonrpcRequest, { relatedRequestId }).catch((error) => { + this._transport.send(jsonrpcRequest, { relatedRequestId, lastEventId, onLastEventIdUpdate }).catch((error) => { this._cleanupTimeout(messageId); reject(error); }); diff --git a/src/shared/transport.ts b/src/shared/transport.ts index e464653b..cfbe8461 100644 --- a/src/shared/transport.ts +++ b/src/shared/transport.ts @@ -18,7 +18,7 @@ export interface Transport { * * If present, `relatedRequestId` is used to indicate to the transport which incoming request to associate this outgoing message with. */ - send(message: JSONRPCMessage, options?: { relatedRequestId?: RequestId }): Promise; + send(message: JSONRPCMessage, options?: { relatedRequestId?: RequestId, lastEventId?: string, onLastEventIdUpdate?: (event: string) => void }): Promise; /** * Closes the connection.