fix(sdk): chat streams reconnect when the body ends mid-turn
This commit is contained in:
@@ -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();
|
||||
|
||||
|
||||
@@ -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({
|
||||
|
||||
@@ -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) {
|
||||
|
||||
Reference in New Issue
Block a user