From 5bb7ae6823273ebb385a93a510344f9df045fcdf Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Wed, 15 Apr 2026 17:42:38 +0100 Subject: [PATCH] fix(chat): prevent useChat resume from hanging on completed turns MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- packages/trigger-sdk/src/v3/chat.test.ts | 163 ++++++++++++++++++++++- packages/trigger-sdk/src/v3/chat.ts | 30 ++++- 2 files changed, 183 insertions(+), 10 deletions(-) diff --git a/packages/trigger-sdk/src/v3/chat.test.ts b/packages/trigger-sdk/src/v3/chat.test.ts index 92636f62b..25ed39e62 100644 --- a/packages/trigger-sdk/src/v3/chat.test.ts +++ b/packages/trigger-sdk/src/v3/chat.test.ts @@ -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 diff --git a/packages/trigger-sdk/src/v3/chat.ts b/packages/trigger-sdk/src/v3/chat.ts index 8a66542b2..db6ef2e40 100644 --- a/packages/trigger-sdk/src/v3/chat.ts +++ b/packages/trigger-sdk/src/v3/chat.ts @@ -173,7 +173,7 @@ type TriggerChatTransportOptionsBase = { * }); * ``` */ - sessions?: Record; + sessions?: Record; /** * Called whenever a chat session's state changes. @@ -203,7 +203,7 @@ type TriggerChatTransportOptionsBase = { */ 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 { 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 { runId: session.runId, publicAccessToken: session.publicAccessToken, lastEventId: session.lastEventId, + isStreaming: session.isStreaming, }); } } @@ -515,6 +518,9 @@ export class TriggerChatTransport implements ChatTransport { 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 { 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 { 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 { */ 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 { 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 { 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 { 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); }