feat(sdk): upsertIncomingMessage helper for hydrateMessages

Author: ericallam

.changeset/chat-slim-wire-merge.md

@@ -9,4 +9,23 @@ Fix `chat.agent` HITL continuations on reasoning-heavy turns. Two changes that w
 
 Note: `onValidateMessages` receives the slim wire on HITL turns. If you call `validateUIMessages` from `ai` against the full `messages` array it will reject the slim assistant; filter to user messages (or skip on HITL turns) — see the updated docstring on `onValidateMessages` for the recommended pattern.
 
+For `hydrateMessages` hooks that persist the chain, this release also adds a small helper to the `@trigger.dev/sdk/ai` surface:
+
+```ts
+import { chat, upsertIncomingMessage } from "@trigger.dev/sdk/ai";
+
+chat.agent({
+  hydrateMessages: async ({ chatId, trigger, incomingMessages }) => {
+    const record = await db.chat.findUnique({ where: { id: chatId } });
+    const stored = record?.messages ?? [];
+    if (upsertIncomingMessage(stored, { trigger, incomingMessages })) {
+      await db.chat.update({ where: { id: chatId }, data: { messages: stored } });
+    }
+    return stored;
+  },
+});
+```
+
+It pushes fresh user messages by id, no-ops on HITL continuations (the incoming shares an id with the existing assistant — the runtime overlays the new tool-state advance), and skips on non-`submit-message` triggers. Returns `true` if it mutated `stored` so the caller knows whether to persist.
+
 Net effect: `chat.addToolOutput(...)` / `chat.addToolApproveResponse(...)` on multi-step reasoning agents (OpenAI Responses with `store: false`, Anthropic extended thinking, etc.) no longer blows the cap and no longer corrupts the LLM input.

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

@@ -199,6 +199,63 @@ export type InferChatUIMessage<TTask extends AnyTask> = TTask extends Task<
   ? TUIM
   : UIMessage;
 
+/**
+ * Upsert an incoming wire message into the customer's DB-backed chain
+ * inside a `hydrateMessages` hook. Returns `true` iff the chain was
+ * mutated (the caller should persist).
+ *
+ * Handles the three cases that matter:
+ *
+ *  - **Non-submit-message trigger** (`regenerate-message` / `action`,
+ *    or `submit-message` with no incoming): no-op. Returns `false`.
+ *  - **Incoming id already in `stored`** (HITL `addToolOutput` /
+ *    `addToolApproveResponse` continuation — the wire carries the
+ *    existing assistant's id with a slim resolution payload): no-op.
+ *    The runtime's per-turn merge overlays the new tool-state advance
+ *    onto the existing entry; pushing again would duplicate the row
+ *    in the chain you return, and the duplicate slim copy would hit
+ *    `toModelMessages` with no `input`. Returns `false`.
+ *  - **Incoming id not in `stored`** (typically a fresh user message
+ *    on a new turn): push. Returns `true`.
+ *
+ * Mutates `stored` in place. The caller persists `stored`, not the
+ * return value.
+ *
+ * @example
+ * ```ts
+ * import { chat, upsertIncomingMessage } from "@trigger.dev/sdk/ai";
+ *
+ * chat.agent({
+ *   hydrateMessages: async ({ chatId, trigger, incomingMessages }) => {
+ *     const record = await db.chat.findUnique({ where: { id: chatId } });
+ *     const stored = record?.messages ?? [];
+ *     if (upsertIncomingMessage(stored, { trigger, incomingMessages })) {
+ *       await db.chat.update({ where: { id: chatId }, data: { messages: stored } });
+ *     }
+ *     return stored;
+ *   },
+ * });
+ * ```
+ */
+export function upsertIncomingMessage<TMsg extends UIMessage = UIMessage>(
+  stored: TMsg[],
+  event: {
+    trigger: "submit-message" | "regenerate-message" | "action";
+    incomingMessages: TMsg[];
+  }
+): boolean {
+  if (event.trigger !== "submit-message") return false;
+  if (event.incomingMessages.length === 0) return false;
+  const newMsg = event.incomingMessages[event.incomingMessages.length - 1];
+  if (!newMsg) return false;
+  if (newMsg.id) {
+    const existingIdx = stored.findIndex((m) => m.id === newMsg.id);
+    if (existingIdx !== -1) return false;
+  }
+  stored.push(newMsg);
+  return true;
+}
+
 /**
  * Tool-part states that the client advances and ships back over the wire.
  * Covers HITL `addToolOutput` (output-available / output-error) and the

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

@@ -2571,7 +2571,7 @@ export type PendingMessagesOptions<TUIM extends UIMessage = UIMessage> = {
 // React hooks (`@trigger.dev/sdk/chat/react`) can import it without
 // dragging `ai.ts` into the browser graph. Re-exported here so
 // `@trigger.dev/sdk/ai` consumers still see it.
-export { PENDING_MESSAGE_INJECTED_TYPE } from "./ai-shared.js";
+export { PENDING_MESSAGE_INJECTED_TYPE, upsertIncomingMessage } from "./ai-shared.js";
 import { PENDING_MESSAGE_INJECTED_TYPE } from "./ai-shared.js";
 
 /** @internal */

packages/trigger-sdk/test/wire-shape.test.ts

@@ -11,7 +11,7 @@ import "../src/v3/test/index.js";
 import type { UIMessage } from "ai";
 import { describe, expect, expectTypeOf, it } from "vitest";
 import type { ChatInputChunk, ChatTaskWirePayload } from "../src/v3/ai-shared.js";
-import { slimSubmitMessageForWire } from "../src/v3/ai-shared.js";
+import { slimSubmitMessageForWire, upsertIncomingMessage } from "../src/v3/ai-shared.js";
 
 describe("ChatTaskWirePayload (slim wire shape)", () => {
   it("encodes and decodes a submit-message payload through JSON", () => {
@@ -158,6 +158,119 @@ describe("ChatTaskWirePayload (slim wire shape)", () => {
   });
 });
 
+describe("upsertIncomingMessage", () => {
+  const userMsg = (id: string, text: string): UIMessage => ({
+    id,
+    role: "user",
+    parts: [{ type: "text", text }],
+  });
+
+  it("pushes a fresh user message and returns true", () => {
+    const stored: UIMessage[] = [userMsg("u-1", "first")];
+    const mutated = upsertIncomingMessage(stored, {
+      trigger: "submit-message",
+      incomingMessages: [userMsg("u-2", "second")],
+    });
+    expect(mutated).toBe(true);
+    expect(stored).toHaveLength(2);
+    expect(stored[1]!.id).toBe("u-2");
+  });
+
+  it("no-ops when incoming id is already in stored (HITL continuation)", () => {
+    const head = {
+      id: "asst-1",
+      role: "assistant" as const,
+      parts: [{ type: "tool-search", toolCallId: "tc-1", state: "input-available", input: {} } as never],
+    };
+    const stored: UIMessage[] = [userMsg("u-1", "hi"), head];
+    const slim = {
+      id: "asst-1",
+      role: "assistant" as const,
+      parts: [{ type: "tool-search", toolCallId: "tc-1", state: "output-available", output: {} } as never],
+    };
+    const mutated = upsertIncomingMessage(stored, {
+      trigger: "submit-message",
+      incomingMessages: [slim],
+    });
+    expect(mutated).toBe(false);
+    expect(stored).toHaveLength(2);
+    // The original head is untouched — the runtime's per-turn merge
+    // overlays the resolution; the customer's stored array is just
+    // the pre-merge snapshot.
+    expect(stored[1]).toBe(head);
+  });
+
+  it("no-ops on regenerate-message trigger", () => {
+    const stored: UIMessage[] = [userMsg("u-1", "hi")];
+    const mutated = upsertIncomingMessage(stored, {
+      trigger: "regenerate-message",
+      incomingMessages: [userMsg("u-2", "ignored")],
+    });
+    expect(mutated).toBe(false);
+    expect(stored).toHaveLength(1);
+  });
+
+  it("no-ops on action trigger", () => {
+    const stored: UIMessage[] = [userMsg("u-1", "hi")];
+    const mutated = upsertIncomingMessage(stored, {
+      trigger: "action",
+      incomingMessages: [],
+    });
+    expect(mutated).toBe(false);
+    expect(stored).toHaveLength(1);
+  });
+
+  it("no-ops on empty incomingMessages", () => {
+    const stored: UIMessage[] = [userMsg("u-1", "hi")];
+    const mutated = upsertIncomingMessage(stored, {
+      trigger: "submit-message",
+      incomingMessages: [],
+    });
+    expect(mutated).toBe(false);
+    expect(stored).toHaveLength(1);
+  });
+
+  it("only inspects the last incoming message (slim wire ships at most one)", () => {
+    const stored: UIMessage[] = [userMsg("u-1", "hi")];
+    const mutated = upsertIncomingMessage(stored, {
+      trigger: "submit-message",
+      incomingMessages: [userMsg("ignored", "ignored"), userMsg("u-3", "new")],
+    });
+    expect(mutated).toBe(true);
+    expect(stored).toHaveLength(2);
+    expect(stored[1]!.id).toBe("u-3");
+  });
+
+  it("pushes when newMsg has no id (no dedup possible)", () => {
+    const stored: UIMessage[] = [userMsg("u-1", "hi")];
+    const incoming = { role: "user", parts: [{ type: "text", text: "no id" }] } as unknown as UIMessage;
+    const mutated = upsertIncomingMessage(stored, {
+      trigger: "submit-message",
+      incomingMessages: [incoming],
+    });
+    expect(mutated).toBe(true);
+    expect(stored).toHaveLength(2);
+  });
+
+  it("accepts the full hydrateMessages event without re-packaging", () => {
+    // Customers can pass the destructured event directly — the helper
+    // only reads `trigger` + `incomingMessages` but ignores any other
+    // fields the event happens to carry.
+    const stored: UIMessage[] = [];
+    const event = {
+      chatId: "chat-1",
+      turn: 0,
+      trigger: "submit-message" as const,
+      incomingMessages: [userMsg("u-1", "hi")],
+      previousMessages: [],
+      continuation: false,
+    };
+    const mutated = upsertIncomingMessage(stored, event);
+    expect(mutated).toBe(true);
+    expect(stored).toHaveLength(1);
+  });
+});
+
 describe("slimSubmitMessageForWire", () => {
   it("passes user messages through unchanged", () => {
     const userMsg: UIMessage = {

references/ai-chat/src/trigger/chat.ts

@@ -1,4 +1,4 @@
-import { chat, type ChatTaskWirePayload } from "@trigger.dev/sdk/ai";
+import { chat, upsertIncomingMessage, type ChatTaskWirePayload } from "@trigger.dev/sdk/ai";
 import { logger, prompts, skills } from "@trigger.dev/sdk";
 
 import {
@@ -887,30 +887,19 @@ export const aiChatHydrated = chat
 
     // Load message history from the database on every turn.
     // The frontend's accumulated messages are ignored — the DB is the
-    // single source of truth. New user messages arrive in `incomingMessages`
-    // and are upserted by id + persisted before returning.
-    //
-    // HITL continuations (tool-output answers from `addToolOutput`) arrive
-    // as a slim assistant message whose id matches the existing assistant
-    // already in the chain — those land via the SDK's per-turn merge, so
-    // they're skipped here. Only brand-new ids (typically the user's next
-    // message) get appended.
+    // single source of truth. `upsertIncomingMessage` handles HITL
+    // continuations (slim wire sharing an id with the existing
+    // assistant — no-op so the runtime overlays the new state) and
+    // fresh user messages (push + persist).
     hydrateMessages: async ({ chatId, trigger, incomingMessages }) => {
       const record = await prisma.chat.findUnique({ where: { id: chatId } });
       const stored = (record?.messages as unknown as UIMessage[]) ?? [];
 
-      if (trigger === "submit-message" && incomingMessages.length > 0) {
-        const newMsg = incomingMessages[incomingMessages.length - 1]!;
-        const existingIdx = newMsg.id
-          ? stored.findIndex((m) => m.id === newMsg.id)
-          : -1;
-        if (existingIdx === -1) {
-          stored.push(newMsg);
-          await prisma.chat.update({
-            where: { id: chatId },
-            data: { messages: stored as unknown as ChatMessagesForWrite },
-          });
-        }
+      if (upsertIncomingMessage(stored, { trigger, incomingMessages })) {
+        await prisma.chat.update({
+          where: { id: chatId },
+          data: { messages: stored as unknown as ChatMessagesForWrite },
+        });
       }
 
       return stored;