fix(chat): prevent useChat resume from hanging on completed turns

Add isStreaming flag to session state — set true when streaming starts,
  false on turn-complete. reconnectToStream returns null immediately when
  isStreaming is false, so resume: true is safe to pass unconditionally.
This commit is contained in:
Eric Allam
2026-04-15 17:42:38 +01:00
parent 49df95b924
commit 5bb7ae6823
2 changed files with 183 additions and 10 deletions
+160 -3
View File
@@ -502,6 +502,145 @@ describe("TriggerChatTransport", () => {
expect(receivedChunks.length).toBeGreaterThan(0);
});
it("should return null when session exists but isStreaming is false (TRI-8557)", async () => {
// Simulate a session restored from DB after a completed turn
const transport = new TriggerChatTransport({
task: "my-task",
accessToken: "token",
sessions: {
"chat-completed": {
runId: "run_completed",
publicAccessToken: "pub_token",
lastEventId: "42",
isStreaming: false,
},
},
});
// reconnectToStream should return null immediately — no hanging
const result = await transport.reconnectToStream({
chatId: "chat-completed",
});
expect(result).toBeNull();
});
it("should reconnect when session exists and isStreaming is true", async () => {
global.fetch = vi.fn().mockImplementation(async (url: string | URL) => {
const urlStr = typeof url === "string" ? url : url.toString();
if (urlStr.includes("/realtime/v1/streams/")) {
const chunks: UIMessageChunk[] = [
{ type: "text-start", id: "part-1" },
{ type: "text-delta", id: "part-1", delta: "Resumed!" },
{ type: "text-end", id: "part-1" },
];
return new Response(createSSEStream(sseEncode(chunks)), {
status: 200,
headers: {
"content-type": "text/event-stream",
"X-Stream-Version": "v1",
},
});
}
throw new Error(`Unexpected fetch URL: ${urlStr}`);
});
const transport = new TriggerChatTransport({
task: "my-task",
accessToken: "token",
baseURL: "https://api.test.trigger.dev",
sessions: {
"chat-streaming": {
runId: "run_streaming",
publicAccessToken: "pub_token",
lastEventId: "10",
isStreaming: true,
},
},
});
const stream = await transport.reconnectToStream({
chatId: "chat-streaming",
});
expect(stream).toBeInstanceOf(ReadableStream);
});
it("should set isStreaming to false via onSessionChange when turn completes", async () => {
const sessionChanges: Array<{
chatId: string;
session: { isStreaming?: boolean } | null;
}> = [];
global.fetch = vi.fn().mockImplementation(async (url: string | URL) => {
const urlStr = typeof url === "string" ? url : url.toString();
if (urlStr.includes("/trigger")) {
return new Response(JSON.stringify({ id: "run_streaming_flag" }), {
status: 200,
headers: {
"content-type": "application/json",
"x-trigger-jwt": "pub_token",
},
});
}
if (urlStr.includes("/realtime/v1/streams/")) {
const chunks = [
{ type: "text-start", id: "part-1" },
{ type: "text-delta", id: "part-1", delta: "Hi" },
{ type: "text-end", id: "part-1" },
{ type: "trigger:turn-complete", publicAccessToken: "refreshed_token" },
];
return new Response(createSSEStream(sseEncode(chunks)), {
status: 200,
headers: {
"content-type": "text/event-stream",
"X-Stream-Version": "v1",
},
});
}
throw new Error(`Unexpected fetch URL: ${urlStr}`);
});
const transport = new TriggerChatTransport({
task: "my-task",
accessToken: "token",
baseURL: "https://api.test.trigger.dev",
onSessionChange: (chatId, session) => {
sessionChanges.push({ chatId, session });
},
});
const stream = await transport.sendMessages({
trigger: "submit-message",
chatId: "chat-flag-test",
messageId: undefined,
messages: [createUserMessage("Hello")],
abortSignal: undefined,
});
// Drain the stream
const reader = stream.getReader();
while (true) {
const { done } = await reader.read();
if (done) break;
}
// Find the session changes for this chat
const changes = sessionChanges.filter((c) => c.chatId === "chat-flag-test");
// First change: session created with isStreaming: true
expect(changes[0]?.session?.isStreaming).toBe(true);
// Last change: turn completed, isStreaming: false
const lastChange = changes[changes.length - 1];
expect(lastChange?.session?.isStreaming).toBe(false);
});
});
describe("renewRunAccessToken", () => {
@@ -651,6 +790,15 @@ describe("TriggerChatTransport", () => {
if (done) break;
}
// Simulate mid-stream state (isStreaming must be true for reconnect to attempt)
const session = transport.getSession("chat-fail-renew");
transport.setOnSessionChange(undefined); // prevent side-effects
// Re-seed with isStreaming: true to simulate reconnect during an active turn
(transport as any).sessions.set("chat-fail-renew", {
...session,
isStreaming: true,
});
const stream = await transport.reconnectToStream({ chatId: "chat-fail-renew" });
const reader = stream!.getReader();
await expect(reader.read()).rejects.toMatchObject({ status: 401 });
@@ -1013,13 +1161,21 @@ describe("TriggerChatTransport", () => {
const r2 = s2.getReader();
while (!(await r2.read()).done) {}
// Both sessions should be independently reconnectable
// Both sessions should exist but not be reconnectable (turns completed)
const sessionA = transport.getSession("session-a");
const sessionB = transport.getSession("session-b");
expect(sessionA).toBeDefined();
expect(sessionB).toBeDefined();
expect(sessionA!.isStreaming).toBe(false);
expect(sessionB!.isStreaming).toBe(false);
// Completed turns return null on reconnect (TRI-8557 fix)
const streamA = await transport.reconnectToStream({ chatId: "session-a" });
const streamB = await transport.reconnectToStream({ chatId: "session-b" });
const streamC = await transport.reconnectToStream({ chatId: "nonexistent" });
expect(streamA).toBeInstanceOf(ReadableStream);
expect(streamB).toBeInstanceOf(ReadableStream);
expect(streamA).toBeNull();
expect(streamB).toBeNull();
expect(streamC).toBeNull();
});
});
@@ -2060,6 +2216,7 @@ describe("TriggerChatTransport", () => {
runId: triggerRunId,
publicAccessToken: publicToken,
lastEventId: undefined,
isStreaming: true,
});
// Consume stream
+23 -7
View File
@@ -173,7 +173,7 @@ type TriggerChatTransportOptionsBase<TClientData = unknown> = {
* });
* ```
*/
sessions?: Record<string, { runId: string; publicAccessToken: string; lastEventId?: string }>;
sessions?: Record<string, { runId: string; publicAccessToken: string; lastEventId?: string; isStreaming?: boolean }>;
/**
* Called whenever a chat session's state changes.
@@ -203,7 +203,7 @@ type TriggerChatTransportOptionsBase<TClientData = unknown> = {
*/
onSessionChange?: (
chatId: string,
session: { runId: string; publicAccessToken: string; lastEventId?: string } | null
session: { runId: string; publicAccessToken: string; lastEventId?: string; isStreaming?: boolean } | null
) => void;
/**
@@ -336,6 +336,8 @@ type ChatSessionState = {
lastEventId?: string;
/** Set when the stream was aborted mid-turn (stop). On reconnect, skip chunks until trigger:turn-complete. */
skipToTurnComplete?: boolean;
/** Whether the agent is currently streaming a response. Set on first chunk, cleared on turn-complete. */
isStreaming?: boolean;
};
/**
@@ -385,7 +387,7 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
private _onSessionChange:
| ((
chatId: string,
session: { runId: string; publicAccessToken: string; lastEventId?: string } | null
session: { runId: string; publicAccessToken: string; lastEventId?: string; isStreaming?: boolean } | null
) => void)
| undefined;
@@ -428,6 +430,7 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
runId: session.runId,
publicAccessToken: session.publicAccessToken,
lastEventId: session.lastEventId,
isStreaming: session.isStreaming,
});
}
}
@@ -515,6 +518,9 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
this.activeStreams.delete(chatId);
}
currentSession.isStreaming = true;
this.notifySessionChange(chatId, currentSession);
return this.subscribeToStream(
currentSession.runId,
currentSession.publicAccessToken,
@@ -534,7 +540,7 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
const { runId, publicAccessToken } = await this.triggerNewRun(chatId, triggerPayload, "trigger");
const newSession: ChatSessionState = { runId, publicAccessToken };
const newSession: ChatSessionState = { runId, publicAccessToken, isStreaming: true };
this.sessions.set(chatId, newSession);
this.notifySessionChange(chatId, newSession);
return this.subscribeToStream(runId, publicAccessToken, abortSignal, chatId, {
@@ -617,6 +623,12 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
return null;
}
// No active stream — the last turn completed before the page refreshed.
// Return null so useChat settles into "ready" state instead of hanging.
if (!session.isStreaming) {
return null;
}
// Deduplicate: if there's already an active stream for this chatId,
// return null so the second caller no-ops.
if (this.activeStreams.has(options.chatId)) {
@@ -790,13 +802,14 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
*/
getSession = (
chatId: string
): { runId: string; publicAccessToken: string; lastEventId?: string } | undefined => {
): { runId: string; publicAccessToken: string; lastEventId?: string; isStreaming?: boolean } | undefined => {
const session = this.sessions.get(chatId);
if (!session) return undefined;
return {
runId: session.runId,
publicAccessToken: session.publicAccessToken,
lastEventId: session.lastEventId,
isStreaming: session.isStreaming,
};
};
@@ -808,7 +821,7 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
callback:
| ((
chatId: string,
session: { runId: string; publicAccessToken: string; lastEventId?: string } | null
session: { runId: string; publicAccessToken: string; lastEventId?: string; isStreaming?: boolean } | null
) => void)
| undefined
): void {
@@ -966,6 +979,7 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
runId: session.runId,
publicAccessToken: session.publicAccessToken,
lastEventId: session.lastEventId,
isStreaming: session.isStreaming,
});
} else {
this._onSessionChange(chatId, null);
@@ -1212,8 +1226,10 @@ export class TriggerChatTransport implements ChatTransport<UIMessage> {
if (session && typeof chunk.publicAccessToken === "string") {
session.publicAccessToken = chunk.publicAccessToken;
}
// Notify with updated session (including refreshed token)
// Mark streaming as complete so reconnectToStream doesn't
// hang on page refresh when no turn is in-flight.
if (session) {
session.isStreaming = false;
this.notifySessionChange(chatId, session);
}