fix(sdk): stop chat.agent replay blocking the event loop on long turns

Recovering a long unfinished `chat.agent` response on a continuation run
could block the worker's event loop for minutes, long enough to miss
heartbeats and fail the run. Capturing the response at the end of a long
turn had the same cost.

Both paths rebuilt the message by draining `readUIMessageStream` and
keeping its last snapshot. That reducer emits a `structuredClone` of the
whole message after almost every chunk, and every chunk was enqueued
into one web stream up front, where each dequeue is linear in the queue
length on Node. Both costs are quadratic in turn length.

The reductions now go through one helper:

```ts
const message = await reduceUIMessageChunks(chunks, { message: original });
```

- It runs the same AI SDK reducer through `createUIMessageStream`'s
`onFinish`, which mutates one state and hands back only the final
message. That message is cloned once so it shares nothing with the
recorded chunks.
- Chunks are fed one at a time, each after the previous one comes out,
and the loop yields to the event loop every 10ms, so heartbeats keep
running.
- The result matches the last `readUIMessageStream` snapshot, including
`undefined` when no chunk would emit one. A malformed chunk is
identified exactly, and the chunks before it are reduced again on the
same path, so content before it is kept without a slow fallback.

A 320,000-chunk text response now reduces in under a second. Tests cover
equivalence with `readUIMessageStream` across finished, unfinished,
resumed and continued messages, a malformed chunk after a long prefix,
detachment from the input chunks, and bounded timer gaps during large
reductions.

Mono-RevId: 4ce17e79a4750a7c949bbc3448881ff0a98dc705

Author: ericallam

.changeset/chat-replay-final-snapshot.md

@@ -0,0 +1,5 @@
+---
+"@trigger.dev/sdk": patch
+---
+
+Recovering a long unfinished `chat.agent` response on a continuation run no longer blocks the worker. Rebuilding the response from its streamed chunks now takes time linear in its length and yields to the event loop as it goes, so heartbeats keep firing and the run is not killed mid-replay. Capturing the response at the end of a long turn gets the same speedup.

packages/trigger-sdk/src/imports/ai-runtime-cjs.cts

@@ -7,6 +7,8 @@ const ai = require("ai");
 // @ts-ignore
 module.exports.convertToModelMessages = ai.convertToModelMessages;
 // @ts-ignore
+module.exports.createUIMessageStream = ai.createUIMessageStream;
+// @ts-ignore
 module.exports.dynamicTool = ai.dynamicTool;
 // @ts-ignore
 module.exports.generateId = ai.generateId;

packages/trigger-sdk/src/imports/ai-runtime.ts

@@ -13,6 +13,7 @@
 // @ts-ignore
 import {
   convertToModelMessages,
+  createUIMessageStream,
   dynamicTool,
   generateId,
   getToolName,
@@ -28,6 +29,7 @@ import {
 // @ts-ignore
 export {
   convertToModelMessages,
+  createUIMessageStream,
   dynamicTool,
   generateId,
   getToolName,

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

@@ -75,7 +75,6 @@ import {
   getToolName,
   isToolUIPart,
   jsonSchema,
-  readUIMessageStream,
   streamText as aiStreamText,
   zodSchema,
 } from "../imports/ai-runtime.js";
@@ -97,6 +96,7 @@ import {
 import { responseAfterCompaction } from "./compactionResponse.js";
 import { withToolResultsInCallOrder } from "./toolResultOrder.js";
 import { ManagedChatResponse, createOrderedChatWriter } from "./managedChatResponse.js";
+import { reduceUIMessageChunks } from "./uiMessageChunks.js";
 import {
   convertSteeredMessages,
   retainStepMessages,
@@ -439,9 +439,9 @@ export function __setReplaySessionOutTailImplForTests(
  *   2. Filter out the agent's control chunks (`type: "trigger:*"`) — they
  *      ride on the same stream as the user-visible UIMessageChunks.
  *   3. Split chunks at `start`/`finish` boundaries so each segment is a
- *      single message, then feed each segment through the AI SDK's
- *      `readUIMessageStream` reducer (the same one `useChat` uses on the
- *      browser side) and grab the final emitted snapshot.
+ *      single message, then reduce each segment with the AI SDK's reducer
+ *      (the same one `useChat` uses on the browser side) via
+ *      {@link reduceUIMessageChunks}, which builds only the final state.
  *   4. The trailing message — if it never received a `finish` chunk —
  *      goes through `cleanupAbortedParts` so partial in-flight parts
  *      don't leak into the next turn's accumulator. Drop it entirely
@@ -497,8 +497,8 @@ async function replaySessionOutTail<TUIMessage extends UIMessage>(
     }
     if (!current) {
       // Chunk arrived before any `start`. Synthesize a segment so the reducer
-      // has something to work with — `readUIMessageStream` tolerates a missing
-      // `start` because we pass `message: undefined`.
+      // has something to work with. The reducer tolerates a missing `start`
+      // and leaves the message ID empty.
       current = { chunks: [], closed: false };
       segments.push(current);
     }
@@ -515,17 +515,9 @@ async function replaySessionOutTail<TUIMessage extends UIMessage>(
   for (let i = 0; i < segments.length; i++) {
     const seg = segments[i]!;
     const isTrailing = i === segments.length - 1 && !seg.closed;
-    const segmentStream = new ReadableStream<UIMessageChunk>({
-      start(controller) {
-        for (const c of seg.chunks) controller.enqueue(c);
-        controller.close();
-      },
-    });
     let last: UIMessage | undefined;
     try {
-      for await (const snapshot of readUIMessageStream({ stream: segmentStream })) {
-        last = snapshot;
-      }
+      last = await reduceUIMessageChunks(seg.chunks);
     } catch (error) {
       // Reducer error — the segment is malformed. Skip it and keep going so a
       // single corrupt chunk doesn't sink the entire replay.
@@ -555,15 +547,15 @@ async function replaySessionOutTail<TUIMessage extends UIMessage>(
 
 /**
  * Test-only entry point that bypasses `__setReplaySessionOutTailImplForTests`
- * and reaches the real `apiClient.subscribeToSessionStream` + chunk-segment
- * splitter + `readUIMessageStream` reducer. Pairs with the snapshot
+ * and reaches the real `apiClient.readSessionStreamRecords` + chunk-segment
+ * splitter + {@link reduceUIMessageChunks} reducer. Pairs with the snapshot
  * production-path wrappers above. Lets `replay-session-out.test.ts` drive
  * synthetic chunk sequences through the real reducer to lock down chunk-
  * stream → `UIMessage[]` correctness — if the AI SDK's chunk semantics
  * shift in a future version, the test catches it before customers do.
  *
- * Tests should mock `apiClient.subscribeToSessionStream` (e.g. via
- * `vi.spyOn(apiClient, ...)`) to feed a `ReadableStream<UIMessageChunk>`.
+ * Tests should stub `apiClient.readSessionStreamRecords` to return the
+ * recorded chunks as session stream records.
  *
  * Not part of the public API.
  * @internal
@@ -11689,7 +11681,7 @@ function tapUIMessageChunks(
  * Reconstruct a partial assistant `UIMessage` from the raw chunks that
  * streamed before a failure — the fallback for {@link pipeChatAndCapture}
  * when a transport error abandons the stream before `onFinish` runs. Uses the
- * same `readUIMessageStream` reducer as the boot-time replay path. Returns
+ * same {@link reduceUIMessageChunks} reducer as the boot-time replay path. Returns
  * `undefined` if there's nothing to assemble or the reducer throws.
  */
 async function assemblePartialFromChunks(chunks: UIMessageChunk[]): Promise<UIMessage | undefined> {
@@ -11699,17 +11691,7 @@ async function assemblePartialFromChunks(chunks: UIMessageChunk[]): Promise<UIMe
   });
   if (relevant.length === 0) return undefined;
   try {
-    const stream = new ReadableStream<UIMessageChunk>({
-      start(controller) {
-        for (const chunk of relevant) controller.enqueue(chunk);
-        controller.close();
-      },
-    });
-    let last: UIMessage | undefined;
-    for await (const message of readUIMessageStream({ stream })) {
-      last = message;
-    }
-    return last;
+    return await reduceUIMessageChunks(relevant);
   } catch {
     return undefined;
   }

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

@@ -1,5 +1,6 @@
 import type { UIMessage, UIMessageChunk } from "ai";
-import { generateId, readUIMessageStream } from "../imports/ai-runtime.js";
+import { generateId } from "../imports/ai-runtime.js";
+import { reduceUIMessageChunks } from "./uiMessageChunks.js";
 
 /** One ordered output channel for a turn's managed model and data chunks. */
 export class ManagedChatResponse {
@@ -147,21 +148,10 @@ export class ManagedChatResponse {
   async snapshot(options?: { message: UIMessage; from: number }): Promise<UIMessage | undefined> {
     if (!this.chunks.length) return options?.message;
     const chunks = this.chunks.slice(options?.from ?? 0);
-    let message: UIMessage | undefined = options?.message;
     const original = options?.message ?? this.original;
-    const stream = new ReadableStream<UIMessageChunk>({
-      start(controller) {
-        for (const chunk of chunks) controller.enqueue(chunk);
-        controller.close();
-      },
-    });
-    for await (const update of readUIMessageStream({
-      stream,
-      ...(original ? { message: structuredClone(original) } : {}),
-      // An error chunk must leave previously streamed content recoverable.
-      terminateOnError: false,
-    }))
-      message = update;
+    const message =
+      (await reduceUIMessageChunks(chunks, original ? { message: original } : undefined)) ??
+      options?.message;
     return message ? { ...message, id: message.id || this.fallbackId } : undefined;
   }
 

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

@@ -16,6 +16,7 @@ import {
   type ChatSnapshotV1,
 } from "../ai.js";
 import { createTestSessionHandle, type TestSessionOutState } from "./test-session-handle.js";
+import { reduceUIMessageChunks } from "../uiMessageChunks.js";
 
 /** Pre-seed locals before the agent's `run()` starts. */
 type SetupLocals = (locals: { set<T>(key: LocalsKey<T>, value: T): void }) => void | Promise<void>;
@@ -898,22 +899,17 @@ export function mockChatAgent(
 /**
  * Reduce a synthetic UIMessageChunk[] sequence into the UIMessage[] that
  * the runtime's `replaySessionOutTail` would produce. Splits chunks at
- * `start` boundaries and feeds each segment through AI SDK's
- * `readUIMessageStream`. The trailing un-finished segment goes through
+ * `start` boundaries and reduces each segment with
+ * `reduceUIMessageChunks`. The trailing un-finished segment goes through
  * `cleanupAbortedParts`. Mirrors the production reducer used in
  * `ai.ts:replaySessionOutTail`.
  */
 async function reduceChunksToMessages(chunks: UIMessageChunk[]): Promise<UIMessage[]> {
   if (chunks.length === 0) return [];
   const aiModule = (await import("ai")) as {
-    readUIMessageStream?: (args: {
-      stream: ReadableStream<UIMessageChunk>;
-    }) => AsyncIterable<UIMessage>;
     cleanupAbortedParts?: (msg: UIMessage) => UIMessage;
   };
-  const readUIMessageStream = aiModule.readUIMessageStream;
   const cleanupAbortedParts = aiModule.cleanupAbortedParts;
-  if (!readUIMessageStream) return [];
 
   type Segment = { chunks: UIMessageChunk[]; closed: boolean };
   const segments: Segment[] = [];
@@ -939,17 +935,9 @@ async function reduceChunksToMessages(chunks: UIMessageChunk[]): Promise<UIMessa
   for (let i = 0; i < segments.length; i++) {
     const seg = segments[i]!;
     const isTrailing = i === segments.length - 1 && !seg.closed;
-    const segmentStream = new ReadableStream<UIMessageChunk>({
-      start(controller) {
-        for (const c of seg.chunks) controller.enqueue(c);
-        controller.close();
-      },
-    });
     let last: UIMessage | undefined;
     try {
-      for await (const snapshot of readUIMessageStream({ stream: segmentStream })) {
-        last = snapshot;
-      }
+      last = await reduceUIMessageChunks(seg.chunks);
     } catch {
       // Skip malformed segment — tests can assert by inspecting what makes it through.
       continue;