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

This commit is contained in:
Katia Bulatova
2026-08-09 22:48:13 +00:00
parent 08871bf37d
commit df64bec62d
4 changed files with 159 additions and 25 deletions
@@ -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.
@@ -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();
+84
View File
@@ -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({
+61 -25
View File
@@ -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) {