resume long running requst from new client

Author: ihrpr

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<void> {
@@ -240,7 +246,7 @@ export class StreamableHTTPClientTransport implements Transport {
     }, delay);
   }
 
-  private _handleSseStream(stream: ReadableStream<Uint8Array> | null): void {
+  private _handleSseStream(stream: ReadableStream<Uint8Array> | 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<void> {
+  async send(message: JSONRPCMessage | JSONRPCMessage[], options?: { lastEventId?: string, onLastEventIdUpdate?: (event: string) => void }): Promise<void> {
     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;
+  }
 }

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<void> {
   console.log('MCP Interactive Client');
@@ -186,7 +188,10 @@ async function connect(url?: string): Promise<void> {
     }
 
     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<void> {
 
     // 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<string, unknown>): Promise<vo
     };
 
     console.log(`Calling tool '${name}' with args:`, args);
-    const result = await client.request(request, CallToolResultSchema);
+    const onLastEventIdUpdate = (event: string) => {
+      notificationsToolLastEventId = event;
+    };
+    const result = await client.request(request, CallToolResultSchema, {
+      lastEventId: notificationsToolLastEventId, onLastEventIdUpdate
+    });
 
     console.log('Tool result:');
     result.content.forEach(item => {

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);
     }

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<string, string> = {

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<z.infer<T>> {
-    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);
       });

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<void>;
+  send(message: JSONRPCMessage, options?: { relatedRequestId?: RequestId, lastEventId?: string, onLastEventIdUpdate?: (event: string) => void }): Promise<void>;
 
   /**
    * Closes the connection.