Track exact SSE event IDs for chat reconnect

Co-authored-by: Eric Allam <eric@trigger.dev>

Author: cursoragent

packages/ai/src/chatTransport.test.ts

@@ -419,7 +419,7 @@ describe("TriggerChatTransport", function () {
         }
 
         const state = runStore.get("chat-2");
-        return Boolean(state && state.lastEventId === "0");
+        return Boolean(state && state.lastEventId === "1-0");
       });
 
       const reconnectStream = await transport.reconnectToStream({
@@ -429,7 +429,7 @@ describe("TriggerChatTransport", function () {
       expect(reconnectStream).not.toBeNull();
 
       const reconnectChunks = await readChunks(reconnectStream!);
-      expect(reconnectLastEventId).toBe("0");
+      expect(reconnectLastEventId).toBe("1-0");
       expect(reconnectChunks).toHaveLength(2);
       expect(reconnectChunks[0]).toMatchObject({
         chunk: { type: "text-delta", id: "msg_2", delta: "world" },

packages/ai/src/chatTransport.ts

@@ -2,6 +2,8 @@ import {
   ApiClient,
   ApiRequestOptions,
   makeIdempotencyKey,
+  SSEStreamPart,
+  SSEStreamSubscription,
   stringifyIO,
   TriggerOptions,
 } from "@trigger.dev/core/v3";
@@ -205,47 +207,61 @@ export class TriggerChatTransport<
     runState: TriggerChatRunState,
     abortSignal: AbortSignal | undefined,
     lastEventId?: string
-  ): Promise<ReadableStream<UIMessageChunk>> {
+  ): Promise<ReadableStream<SSEStreamPart<InferUIMessageChunk<UI_MESSAGE>>>> {
     const streamClient = new ApiClient(
       this.baseURL,
       runState.publicAccessToken,
       this.previewBranch,
       this.requestOptions
     );
 
-    const stream = await streamClient.fetchStream<InferUIMessageChunk<UI_MESSAGE>>(
-      runState.runId,
-      runState.streamKey,
+    const subscription = new SSEStreamSubscription(
+      this.createStreamUrl(runState.runId, runState.streamKey),
       {
+        headers: streamClient.getHeaders(),
         signal: abortSignal,
         timeoutInSeconds: this.timeoutInSeconds,
         lastEventId,
       }
     );
 
-    return stream as unknown as ReadableStream<UIMessageChunk>;
+    return (await subscription.subscribe()) as ReadableStream<
+      SSEStreamPart<InferUIMessageChunk<UI_MESSAGE>>
+    >;
   }
 
-  private createTrackedStream(chatId: string, stream: ReadableStream<UIMessageChunk>) {
+  private createTrackedStream(
+    chatId: string,
+    stream: ReadableStream<SSEStreamPart<InferUIMessageChunk<UI_MESSAGE>>>
+  ) {
     const teeStreams = stream.tee();
     const trackingStream = teeStreams[0];
     const consumerStream = teeStreams[1];
 
     this.consumeTrackingStream(chatId, trackingStream);
 
-    return consumerStream;
+    return consumerStream.pipeThrough(
+      new TransformStream<SSEStreamPart<InferUIMessageChunk<UI_MESSAGE>>, UIMessageChunk>({
+        transform(part, controller) {
+          controller.enqueue(part.chunk as UIMessageChunk);
+        },
+      })
+    );
   }
 
-  private async consumeTrackingStream(chatId: string, stream: ReadableStream<UIMessageChunk>) {
+  private async consumeTrackingStream(
+    chatId: string,
+    stream: ReadableStream<SSEStreamPart<InferUIMessageChunk<UI_MESSAGE>>>
+  ) {
     try {
-      for await (const _chunk of stream) {
+      for await (const part of stream) {
         const runState = await this.runStore.get(chatId);
 
         if (!runState) {
           return;
         }
 
-        runState.lastEventId = incrementLastEventId(runState.lastEventId);
+        runState.lastEventId = part.id;
         await this.runStore.set(runState);
       }
 
@@ -274,6 +290,14 @@ export class TriggerChatTransport<
 
     return handle as TriggerTaskResponse;
   }
+
+  private createStreamUrl(runId: string, streamKey: string): string {
+    const normalizedBaseUrl = this.baseURL.replace(/\/$/, "");
+    const encodedRunId = encodeURIComponent(runId);
+    const encodedStreamKey = encodeURIComponent(streamKey);
+
+    return `${normalizedBaseUrl}/realtime/v1/streams/${encodedRunId}/${encodedStreamKey}`;
+  }
 }
 
 export function createTriggerChatTransport<
@@ -426,17 +450,4 @@ async function createTriggerTaskOptions(
   };
 }
 
-function incrementLastEventId(lastEventId: string | undefined): string {
-  if (!lastEventId) {
-    return "0";
-  }
-
-  const numberValue = Number.parseInt(lastEventId, 10);
-  if (Number.isNaN(numberValue)) {
-    return "0";
-  }
-
-  return String(numberValue + 1);
-}
-
 export type { TriggerChatTaskContext };