fix(sdk): chat streams reconnect when the body ends mid-turn

Author: kathiekiwi

.changeset/chat-stream-mid-turn-reconnect.md

@@ -0,0 +1,6 @@
+---
+"@trigger.dev/core": patch
+"@trigger.dev/sdk": patch
+---
+
+Chat streams now reconnect when the connection drops mid-turn, instead of leaving the reply stuck as if it were still generating.

packages/core/src/v3/apiClient/runStream.ts

@@ -206,6 +206,13 @@ export class SSEStreamSubscription implements StreamSubscription {
   private cancelledByConsumer = false;
   private completeNotified = false;
 
+  /**
+   * True when the most recent response carried `X-Session-Settled: true` —
+   * the server has no more records coming, so a clean end of the body is
+   * terminal rather than the end of a long-poll window.
+   */
+  sessionSettled = false;
+
   constructor(
     private url: string,
     private options: {
@@ -414,6 +421,7 @@ export class SSEStreamSubscription implements StreamSubscription {
       }
 
       const streamVersion = response.headers.get("X-Stream-Version") ?? "v1";
+      this.sessionSettled = response.headers.get("X-Session-Settled") === "true";
       this.retryCount = 0; // reset on success
       armStall();
 

packages/trigger-sdk/src/v3/chat.test.ts

@@ -1054,6 +1054,90 @@ describe("TriggerChatTransport", () => {
     });
   });
 
+  describe("stream body ends mid-turn", () => {
+    it("resubscribes from the last event id when the close was not settled", async () => {
+      const subscribeHeaders: Headers[] = [];
+      global.fetch = vi.fn().mockImplementation(async (url: string | URL, init?: RequestInit) => {
+        const urlStr = typeof url === "string" ? url : url.toString();
+        if (isSessionStreamAppendUrl(urlStr)) return defaultAppendResponse();
+        if (isSessionOutSubscribeUrl(urlStr)) {
+          subscribeHeaders.push(new Headers(init?.headers));
+          // First connection ends mid-turn: one chunk, no turn-complete,
+          // no `X-Session-Settled`.
+          return subscribeHeaders.length === 1
+            ? defaultSseResponse([{ type: "text-start", id: "part-1" }])
+            : defaultSseResponse([
+                { type: "text-delta", id: "part-1", delta: "resumed" },
+                { type: "trigger:turn-complete" },
+              ]);
+        }
+        throw new Error(`Unexpected URL: ${urlStr}`);
+      });
+
+      const transport = new TriggerChatTransport({
+        task: "my-chat-task",
+        accessToken: () => "pat",
+        sessions: { "chat-eof": { publicAccessToken: "p" } },
+      });
+
+      const stream = await transport.sendMessages({
+        trigger: "submit-message",
+        chatId: "chat-eof",
+        messageId: undefined,
+        messages: [createUserMessage("hi")],
+        abortSignal: undefined,
+      });
+      const chunks = await drainChunks(stream);
+
+      expect(subscribeHeaders).toHaveLength(2);
+      expect(subscribeHeaders[1]?.get("Last-Event-ID")).toBe("1");
+      expect(chunks).toEqual([
+        { type: "text-start", id: "part-1" },
+        { type: "text-delta", id: "part-1", delta: "resumed" },
+      ]);
+      expect(transport.getSession("chat-eof")?.isStreaming).toBe(false);
+    });
+
+    it("stops and clears isStreaming when the close was settled", async () => {
+      let subscribeCount = 0;
+      global.fetch = vi.fn().mockImplementation(async (url: string | URL) => {
+        const urlStr = typeof url === "string" ? url : url.toString();
+        if (isSessionStreamAppendUrl(urlStr)) return defaultAppendResponse();
+        if (isSessionOutSubscribeUrl(urlStr)) {
+          subscribeCount++;
+          const response = defaultSseResponse([{ type: "text-start", id: "part-1" }]);
+          const headers = new Headers(response.headers);
+          headers.set("X-Session-Settled", "true");
+          return new Response(response.body, { status: 200, headers });
+        }
+        throw new Error(`Unexpected URL: ${urlStr}`);
+      });
+
+      const onSessionChange = vi.fn();
+      const transport = new TriggerChatTransport({
+        task: "my-chat-task",
+        accessToken: () => "pat",
+        onSessionChange,
+        sessions: { "chat-settled": { publicAccessToken: "p" } },
+      });
+
+      const stream = await transport.sendMessages({
+        trigger: "submit-message",
+        chatId: "chat-settled",
+        messageId: undefined,
+        messages: [createUserMessage("hi")],
+        abortSignal: undefined,
+      });
+      await drainChunks(stream);
+
+      expect(subscribeCount).toBe(1);
+      expect(transport.getSession("chat-settled")?.isStreaming).toBe(false);
+      expect(
+        onSessionChange.mock.calls.some(([, session]) => session && session.isStreaming === false)
+      ).toBe(true);
+    });
+  });
+
   describe("multi-tab coordination", () => {
     it("isReadOnly defaults to false when multiTab is disabled", () => {
       const transport = new TriggerChatTransport({

packages/trigger-sdk/src/v3/chat.ts

@@ -1759,6 +1759,50 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
           }
         };
 
+        const openWithAuthRetry = async () => {
+          try {
+            return await connectSseOnce(state.publicAccessToken);
+          } catch (e) {
+            if (!isAuthError(e)) throw e;
+            const fresh = await this.resolveAccessToken({ chatId });
+            state.publicAccessToken = fresh;
+            this.notifySessionChange(chatId, state);
+            return await connectSseOnce(fresh);
+          }
+        };
+
+        // A body that ends without a turn-complete is only terminal when the
+        // server says the session settled — otherwise the turn is still
+        // running and we lost the connection (long-poll window closed, proxy
+        // restarted). Resubscribe from `state.lastEventId`, bounded so a
+        // permanently empty stream can't spin.
+        const MAX_EOF_RESUBSCRIBES = 5;
+        let eofResubscribes = 0;
+
+        const resumeAfterEof = async () => {
+          while (
+            state.isStreaming &&
+            !currentSubscription?.sessionSettled &&
+            !combinedSignal.aborted &&
+            eofResubscribes < MAX_EOF_RESUBSCRIBES
+          ) {
+            eofResubscribes++;
+            await new Promise((resolve) =>
+              setTimeout(resolve, Math.min(100 * 2 ** (eofResubscribes - 1), 5_000))
+            );
+            const opened = await openWithAuthRetry();
+            if (opened) return opened;
+          }
+
+          // Settled close, or the turn is gone — tell the UI instead of
+          // leaving it spinning on a stream nobody will finish.
+          if (state.isStreaming) {
+            state.isStreaming = false;
+            this.notifySessionChange(chatId, state);
+          }
+          return null;
+        };
+
         try {
           let reader: ReadableStreamDefaultReader<{
             id: string;
@@ -1767,30 +1811,13 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
           }>;
           let primed: { id: string; chunk: unknown; timestamp: number } | undefined;
 
-          try {
-            const opened = await connectSseOnce(state.publicAccessToken);
-            if (opened === null) {
-              controller.close();
-              return;
-            }
-            reader = opened.reader;
-            primed = opened.primed;
-          } catch (e) {
-            if (isAuthError(e)) {
-              const fresh = await this.resolveAccessToken({ chatId });
-              state.publicAccessToken = fresh;
-              this.notifySessionChange(chatId, state);
-              const opened = await connectSseOnce(fresh);
-              if (opened === null) {
-                controller.close();
-                return;
-              }
-              reader = opened.reader;
-              primed = opened.primed;
-            } else {
-              throw e;
-            }
+          const opened = (await openWithAuthRetry()) ?? (await resumeAfterEof());
+          if (opened === null) {
+            controller.close();
+            return;
           }
+          reader = opened.reader;
+          primed = opened.primed;
 
           this.emitEvent({
             type: "stream-connected",
@@ -1814,10 +1841,19 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
             } else {
               const next = await reader.read();
               if (next.done) {
-                controller.close();
-                return;
+                const resumed = await resumeAfterEof();
+                if (resumed === null) {
+                  controller.close();
+                  return;
+                }
+                reader = resumed.reader;
+                primed = resumed.primed;
+                continue;
               }
               value = next.value;
+              // A productive connection re-earns the resubscribe budget, so a
+              // long turn spanning many long-poll windows keeps streaming.
+              eofResubscribes = 0;
             }
 
             if (combinedSignal.aborted) {