wip allowing for resumability

Author: mattzcarey

packages/agents/src/mcp/worker-transport.ts

@@ -44,6 +44,7 @@ export interface TransportState {
   sessionId?: string;
   initialized: boolean;
   protocolVersion?: ProtocolVersion;
+  requestMappings?: Record<string, string>;
 }
 
 export interface WorkerTransportOptions {
@@ -76,6 +77,7 @@ export class WorkerTransport implements Transport {
   private protocolVersion?: ProtocolVersion;
   private storage?: MCPStorageApi;
   private stateRestored = false;
+  private eventCounters = new Map<string, number>();
 
   sessionId?: string;
   onclose?: () => void;
@@ -90,6 +92,51 @@ export class WorkerTransport implements Transport {
     this.storage = options?.storage;
   }
 
+  /**
+   * Generate a unique event ID for SSE spec compliance.
+   * Counter is ephemeral (not persisted) - validation only checks streamId existence.
+   */
+  private generateEventId(streamId: string): string {
+    if (!this.sessionIdGenerator) {
+      return `session-${streamId}-${Date.now()}`;
+    }
+
+    const currentCounter = this.eventCounters.get(streamId) || 0;
+    const nextCounter = currentCounter + 1;
+    this.eventCounters.set(streamId, nextCounter);
+
+    return `${this.sessionId || "session"}-${streamId}-${nextCounter}`;
+  }
+
+  /**
+   * Parse an event ID to extract streamId and counter.
+   * Event ID format: "{sessionId}-{streamId}-{counter}"
+   */
+  private parseEventId(
+    eventId: string
+  ): { streamId: string; counter: number } | null {
+    const sessionPrefix = this.sessionId ? `${this.sessionId}-` : "session-";
+    if (!eventId.startsWith(sessionPrefix)) {
+      return null;
+    }
+
+    const remainder = eventId.slice(sessionPrefix.length);
+
+    const lastDash = remainder.lastIndexOf("-");
+    if (lastDash === -1) {
+      return null;
+    }
+
+    const counterStr = remainder.slice(lastDash + 1);
+    const counter = parseInt(counterStr, 10);
+    if (isNaN(counter)) {
+      return null;
+    }
+
+    const streamId = remainder.slice(0, lastDash);
+    return { streamId, counter };
+  }
+
   /**
    * Restore transport state from persistent storage.
    * This is automatically called on start.
@@ -107,11 +154,33 @@ export class WorkerTransport implements Transport {
       this.protocolVersion = state.protocolVersion;
     }
 
+    if (state?.requestMappings) {
+      for (const [requestId, streamId] of Object.entries(
+        state.requestMappings
+      )) {
+        this.requestToStreamMapping.set(requestId, streamId);
+      }
+    }
+
     this.stateRestored = true;
   }
 
+  /**
+   * Convert a Map to a plain Record for JSON serialization.
+   */
+  private mapToRecord<K extends string | number, V>(
+    map: Map<K, V>
+  ): Record<string, V> {
+    const record: Record<string, V> = {};
+    for (const [key, value] of map.entries()) {
+      record[String(key)] = value;
+    }
+    return record;
+  }
+
   /**
    * Persist current transport state to storage.
+   * Called on initialization and POST completion only.
    */
   private async saveState() {
     if (!this.storage) {
@@ -121,7 +190,8 @@ export class WorkerTransport implements Transport {
     const state: TransportState = {
       sessionId: this.sessionId,
       initialized: this.initialized,
-      protocolVersion: this.protocolVersion
+      protocolVersion: this.protocolVersion,
+      requestMappings: this.mapToRecord(this.requestToStreamMapping)
     };
 
     await Promise.resolve(this.storage.set(state));
@@ -289,15 +359,73 @@ export class WorkerTransport implements Transport {
       return sessionError;
     }
 
-    // Validate protocol version on subsequent requests
     const versionError = this.validateProtocolVersion(request);
     if (versionError) {
       return versionError;
     }
 
-    const streamId = this.standaloneSseStreamId;
+    const lastEventId = request.headers.get("Last-Event-ID");
+    let streamId: string;
+    let isResumingStream = false;
+
+    if (lastEventId) {
+      const parsed = this.parseEventId(lastEventId);
+      if (!parsed) {
+        return new Response(
+          JSON.stringify({
+            jsonrpc: "2.0",
+            error: {
+              code: -32000,
+              message: `Cannot resume: invalid event ID format "${lastEventId}".`
+            },
+            id: null
+          }),
+          {
+            status: 404,
+            headers: {
+              "Content-Type": "application/json",
+              ...this.getHeaders()
+            }
+          }
+        );
+      }
+
+      // Validate streamId exists by deriving from requestMappings + GET stream
+      // POST streams are valid if they have pending requests
+      // GET stream is always valid (persistent notification channel)
+      const validStreamIds = new Set(
+        Array.from(this.requestToStreamMapping.values())
+      );
+      validStreamIds.add(this.standaloneSseStreamId);
 
-    if (this.streamMapping.get(streamId) !== undefined) {
+      if (!validStreamIds.has(parsed.streamId)) {
+        return new Response(
+          JSON.stringify({
+            jsonrpc: "2.0",
+            error: {
+              code: -32000,
+              message: `Cannot resume: stream "${parsed.streamId}" not found. POST streams expire after responses complete.`
+            },
+            id: null
+          }),
+          {
+            status: 404,
+            headers: {
+              "Content-Type": "application/json",
+              ...this.getHeaders()
+            }
+          }
+        );
+      }
+
+      streamId = parsed.streamId;
+      isResumingStream = true;
+    } else {
+      streamId = this.standaloneSseStreamId;
+    }
+
+    const existingStream = this.streamMapping.get(streamId);
+    if (existingStream !== undefined && !isResumingStream) {
       return new Response(
         JSON.stringify({
           jsonrpc: "2.0",
@@ -340,6 +468,9 @@ export class WorkerTransport implements Transport {
       }
     }, 30000);
 
+    if (isResumingStream && existingStream) {
+      existingStream.cleanup();
+    }
     this.streamMapping.set(streamId, {
       writer,
       encoder,
@@ -350,6 +481,10 @@ export class WorkerTransport implements Transport {
       }
     });
 
+    const primingEventId = this.generateEventId(streamId);
+    const primingEvent = `id: ${primingEventId}\ndata: \n\n`;
+    writer.write(encoder.encode(primingEvent)).catch(() => {});
+
     return new Response(readable, { headers });
   }
 
@@ -608,6 +743,10 @@ export class WorkerTransport implements Transport {
       }
     }
 
+    const primingEventId = this.generateEventId(streamId);
+    const primingEvent = `id: ${primingEventId}\ndata: \n\n`;
+    writer.write(encoder.encode(primingEvent)).catch(() => {});
+
     for (const message of messages) {
       this.onmessage?.(message, { requestInfo });
     }
@@ -738,6 +877,8 @@ export class WorkerTransport implements Transport {
 
     this.streamMapping.clear();
     this.requestResponseMap.clear();
+    this.requestToStreamMapping.clear();
+    this.eventCounters.clear();
     this.onclose?.();
   }
 
@@ -766,7 +907,8 @@ export class WorkerTransport implements Transport {
       }
 
       if (standaloneSse.writer && standaloneSse.encoder) {
-        const data = `event: message\ndata: ${JSON.stringify(message)}\n\n`;
+        const eventId = this.generateEventId(this.standaloneSseStreamId);
+        const data = `id: ${eventId}\nevent: message\ndata: ${JSON.stringify(message)}\n\n`;
         await standaloneSse.writer.write(standaloneSse.encoder.encode(data));
       }
       return;
@@ -788,7 +930,8 @@ export class WorkerTransport implements Transport {
 
     if (!this.enableJsonResponse) {
       if (response.writer && response.encoder) {
-        const data = `event: message\ndata: ${JSON.stringify(message)}\n\n`;
+        const eventId = this.generateEventId(streamId);
+        const data = `id: ${eventId}\nevent: message\ndata: ${JSON.stringify(message)}\n\n`;
         await response.writer.write(response.encoder.encode(data));
       }
     }
@@ -829,6 +972,8 @@ export class WorkerTransport implements Transport {
           this.requestResponseMap.delete(id);
           this.requestToStreamMapping.delete(id);
         }
+
+        this.saveState();
       }
     }
   }