fix(chat): make a response streamed from onAction part of the conversation

Returning a StreamTextResult from onAction piped it to the browser and stopped
there. The accumulator never saw it, no snapshot recorded it, and actions fire no
onTurnComplete — so the user read a good answer that the model had no memory of,
and the next turn carried on from the answer regenerate had just replaced. The
disagreement between the screen and the conversation was invisible until that
next turn contradicted it.

The action branch now captures what it pipes, using the pipeChatAndCapture that
already existed for exactly this, and appends the message to the accumulator.
Persistence beyond the snapshot is still the app's job, since an action fires no
turn hook — pipeAndCapture hands back the same message for that.

Also folds the snapshot write added for rolled-back history into one helper used
by both action paths, so a regenerate that both rolls back and answers writes
once rather than twice, and the cursor-preservation rule lives in one place.

The two fixes needed each other: with the rollback persisted but the response
dropped, a regenerate left the snapshot empty rather than stale — still wrong,
just differently.

Author: ericallam

.changeset/action-stream-into-conversation.md

@@ -0,0 +1,5 @@
+---
+"@trigger.dev/sdk": patch
+---
+
+A response streamed back from `onAction` is now part of the conversation. Returning a `StreamTextResult` from an action sent it to the browser and nowhere else, so a regenerate showed the user a new answer that the model had no memory of — the next turn carried on from the answer that had just been replaced.

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

@@ -6561,6 +6561,56 @@ function chatAgent<
        * keeps an action's write cursor-neutral.
        */
       let lastSnapshotOutEventId: string | undefined;
+
+      /**
+       * Persist the accumulator outside a turn.
+       *
+       * An action is not a turn, so it never reaches the turn-complete path where
+       * the snapshot is normally written — but it can change the conversation in
+       * two ways: a `chat.history` mutation, and a response streamed back from
+       * `onAction`. Both have to survive, and one write at the end of the action
+       * covers both rather than writing twice for a regenerate that does both.
+       *
+       * Cursor-neutral: an action has no turn cursor of its own, and writing
+       * `undefined` would drop the resume point the last turn established and make
+       * the next boot replay from further back.
+       */
+      const writeSnapshotOutsideTurn = async (reason: string) => {
+        if (hydrateMessages) return;
+        try {
+          await tracer.startActiveSpan(
+            "snapshot.write",
+            async () => {
+              const snapshotInCursor = chatInputRouter().resumeFloor();
+              await writeChatSnapshot<TUIMessage>(sessionIdForSnapshot, {
+                version: 1,
+                savedAt: Date.now(),
+                messages: accumulatedUIMessages,
+                lastOutEventId: lastSnapshotOutEventId,
+                lastInEventId:
+                  snapshotInCursor !== undefined ? String(snapshotInCursor) : undefined,
+              });
+            },
+            {
+              attributes: {
+                [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onStart",
+                [SemanticInternalAttributes.COLLAPSED]: true,
+                "chat.snapshot.reason": reason,
+                "chat.messages.count": accumulatedUIMessages.length,
+              },
+            }
+          );
+        } catch (error) {
+          logger.warn(
+            "chat.agent: snapshot write outside a turn failed; the change may not survive a continuation",
+            {
+              error: error instanceof Error ? error.message : String(error),
+              sessionId: sessionIdForSnapshot,
+              reason,
+            }
+          );
+        }
+      };
       let replayedSettled: TUIMessage[] = [];
       let replayedPartial: TUIMessage | undefined;
       let replayedPartialRaw: TUIMessage | undefined;
@@ -7548,6 +7598,13 @@ function chatAgent<
                 // string, or UIMessage from `onAction`. Turn counter
                 // does not advance.
                 let actionStreamResult: unknown = undefined;
+                /**
+                 * Whether this action changed the conversation, by rolling history
+                 * back or by streaming a response. Drives the single snapshot write
+                 * at the end — an action never reaches the turn-complete path that
+                 * normally does it.
+                 */
+                let actionChangedHistory = false;
                 if (isAction) {
                   // Parse and validate the action payload
                   const parsedAction = parseAction
@@ -7620,62 +7677,7 @@ function chatAgent<
                       accumulatedMessages = await toModelMessages(actionOverride);
                       locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
 
-                      /**
-                       * Persist it. An action is not a turn, so it never reaches the
-                       * turn-complete path below where the snapshot is normally
-                       * written — and `onAction` is exactly where `chat.history`
-                       * rollbacks are meant to happen.
-                       *
-                       * Without this the rollback lives only in this worker's
-                       * memory: undo works while the worker stays warm, and the
-                       * next continuation boots from a snapshot that still holds
-                       * the undone exchange, so the undo silently reverts minutes
-                       * later with no error. Awaited for the same reason as the
-                       * turn-complete write — the agent may suspend immediately
-                       * after, and in-flight promises do not reliably survive that.
-                       */
-                      if (!hydrateMessages) {
-                        try {
-                          await tracer.startActiveSpan(
-                            "snapshot.write",
-                            async () => {
-                              // The resume floor, not the dispatched high-water. An
-                              // action can run while records are still queued, and the
-                              // floor is held back below the earliest of those — writing
-                              // the high-water instead would advance the cursor past
-                              // records a replay still has to recover, losing them.
-                              const snapshotInCursor = chatInputRouter().resumeFloor();
-                              await writeChatSnapshot<TUIMessage>(sessionIdForSnapshot, {
-                                version: 1,
-                                savedAt: Date.now(),
-                                messages: accumulatedUIMessages,
-                                lastOutEventId: lastSnapshotOutEventId,
-                                lastInEventId:
-                                  snapshotInCursor !== undefined
-                                    ? String(snapshotInCursor)
-                                    : undefined,
-                              });
-                            },
-                            {
-                              attributes: {
-                                [SemanticInternalAttributes.STYLE_ICON]: "task-hook-onStart",
-                                [SemanticInternalAttributes.COLLAPSED]: true,
-                                "chat.id": currentWirePayload.chatId,
-                                "chat.snapshot.reason": "action",
-                                "chat.messages.count": accumulatedUIMessages.length,
-                              },
-                            }
-                          );
-                        } catch (error) {
-                          logger.warn(
-                            "chat.agent: snapshot write after action failed; the mutation may not survive a continuation",
-                            {
-                              error: error instanceof Error ? error.message : String(error),
-                              sessionId: sessionIdForSnapshot,
-                            }
-                          );
-                        }
-                      }
+                      actionChangedHistory = true;
                     }
                   } else {
                     warnMissingOnActionOnce();
@@ -7959,17 +7961,37 @@ function chatAgent<
                     isUIMessageStreamable(actionStreamResult)
                   ) {
                     try {
-                      const resolvedOptions = resolveUIMessageStreamOptions();
-                      const uiStream = (
-                        actionStreamResult as UIMessageStreamable
-                      ).toUIMessageStream({
-                        ...resolvedOptions,
-                        generateMessageId: resolvedOptions.generateMessageId ?? generateMessageId,
-                      });
-                      await pipeChat(uiStream, {
-                        signal: combinedSignal,
-                        spanName: "stream response",
-                      });
+                      /**
+                       * Captured, not just piped. The stream reaching the browser was
+                       * never the problem — the problem was that it stopped there, so
+                       * the user read an answer the accumulator had no record of and
+                       * the next turn contradicted the screen. Worst on regenerate,
+                       * which removes the old answer and used to leave nothing in its
+                       * place.
+                       *
+                       * Persistence beyond the snapshot is still the app's job: an
+                       * action fires no `onTurnComplete`, so an app owning its own
+                       * store has to write the row itself — `chat.pipeAndCapture`
+                       * hands back the same message for that.
+                       */
+                      const { message: actionResponse } = await pipeChatAndCapture(
+                        actionStreamResult as UIMessageStreamable,
+                        { signal: combinedSignal, spanName: "stream response" }
+                      );
+
+                      if (actionResponse) {
+                        const existingIdx = actionResponse.id
+                          ? accumulatedUIMessages.findIndex((m) => m.id === actionResponse.id)
+                          : -1;
+                        if (existingIdx !== -1) {
+                          accumulatedUIMessages[existingIdx] = actionResponse as TUIMessage;
+                        } else {
+                          accumulatedUIMessages.push(actionResponse as TUIMessage);
+                        }
+                        accumulatedMessages = await toModelMessages(accumulatedUIMessages);
+                        locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
+                        actionChangedHistory = true;
+                      }
                     } catch (error) {
                       if (
                         error instanceof Error &&
@@ -7982,6 +8004,10 @@ function chatAgent<
                     }
                   }
 
+                  if (actionChangedHistory) {
+                    await writeSnapshotOutsideTurn("action");
+                  }
+
                   await writeTurnCompleteChunk(currentWirePayload.chatId);
 
                   // Don't consume a turn iteration — actions aren't turns.

packages/trigger-sdk/test/action-stream-accumulator.test.ts

@@ -0,0 +1,98 @@
+// Import the test harness FIRST — installs the resource catalog so
+// `chat.agent()` below registers its task functions correctly.
+import { mockChatAgent } from "../src/v3/test/index.js";
+
+import { describe, expect, it } from "vitest";
+import { chat } from "../src/v3/ai.js";
+import { simulateReadableStream, streamText } from "ai";
+import type { UIMessage } from "ai";
+import { MockLanguageModelV3 } from "ai/test";
+import type { LanguageModelV3StreamPart } from "@ai-sdk/provider";
+import { z } from "zod";
+
+function textStream(text: string): ReadableStream<LanguageModelV3StreamPart> {
+  return simulateReadableStream({
+    chunks: [
+      { type: "text-start", id: "t1" },
+      { type: "text-delta", id: "t1", delta: text },
+      { type: "text-end", id: "t1" },
+      {
+        type: "finish",
+        finishReason: { unified: "stop", raw: "stop" },
+        usage: {
+          inputTokens: { total: 10, noCache: 10, cacheRead: undefined, cacheWrite: undefined },
+          outputTokens: { total: 10, text: 10, reasoning: undefined },
+        },
+      },
+    ],
+  });
+}
+
+function textOf(message: UIMessage): string {
+  return message.parts.map((part) => (part.type === "text" ? part.text : "")).join("");
+}
+
+describe("a StreamTextResult returned from onAction (TRI-13378)", () => {
+  it("becomes part of the conversation, not just something the browser saw", async () => {
+    const model = new MockLanguageModelV3({
+      doStream: async () => ({ stream: textStream("regenerated answer") }),
+    });
+
+    const agent = chat.agent({
+      id: "action-stream-accumulator",
+      actionSchema: z.discriminatedUnion("type", [z.object({ type: z.literal("regenerate") })]),
+
+      /**
+       * The bare shape the docs show: return the stream and let the runtime pipe
+       * it. The alternative — consuming it with `chat.pipeAndCapture` — is the
+       * workaround, so testing that instead would prove nothing about this path.
+       */
+      onAction: async ({ action, messages }) => {
+        if (action.type !== "regenerate") return;
+        chat.history.slice(0, -1);
+        return streamText({ model, messages });
+      },
+
+      run: async ({ messages, signal }) =>
+        streamText({
+          model: new MockLanguageModelV3({
+            doStream: async () => ({ stream: textStream("first answer") }),
+          }),
+          messages,
+          abortSignal: signal,
+        }),
+    });
+
+    const harness = mockChatAgent(agent, { chatId: "action-stream-accumulator" });
+
+    try {
+      await harness.sendMessage({
+        id: "u1",
+        role: "user",
+        parts: [{ type: "text", text: "ask" }],
+      });
+      await new Promise((r) => setTimeout(r, 30));
+
+      const turn = await harness.sendAction({ type: "regenerate" });
+      await new Promise((r) => setTimeout(r, 50));
+
+      // The browser did see it — that part was never broken.
+      const streamed = turn.chunks
+        .filter((c) => c.type === "text-delta")
+        .map((c) => (c as { delta: string }).delta)
+        .join("");
+      expect(streamed).toBe("regenerated answer");
+
+      /**
+       * And the conversation agrees with the screen. Before the fix the response
+       * was piped and dropped: absent from the accumulator, absent from the
+       * snapshot, so the next turn's model context contained the question and the
+       * *old* answer that regenerate had just removed.
+       */
+      const snapshot = harness.getSnapshot();
+      expect(snapshot?.messages.map(textOf)).toEqual(["ask", "regenerated answer"]);
+    } finally {
+      await harness.close();
+    }
+  });
+});