fix(sdk): dedup newUIMessages against pre-hydration accumulator

Author: ericallam

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

@@ -6158,6 +6158,20 @@ function chatAgent<
                   }
 
                   if (hydrateMessages) {
+                    // Snapshot the ids the accumulator knew BEFORE this
+                    // turn ran — used below to decide whether an
+                    // incoming wire message is genuinely new or just a
+                    // state advance on an existing entry. We can't use
+                    // the post-`hydrateMessages` array for this because
+                    // the canonical hook pattern pushes the incoming
+                    // user message into the persisted chain and
+                    // returns it.
+                    const previouslyKnownMessageIds = new Set(
+                      accumulatedUIMessages
+                        .map((m) => m.id)
+                        .filter((id): id is string => typeof id === "string")
+                    );
+
                     // Backend hydration: load the full message history from the user's
                     // backend, replacing the built-in accumulator entirely. With slim
                     // wire, `incomingMessages` is consistently 0-or-1-length — what
@@ -6217,18 +6231,25 @@ function chatAgent<
 
                     // Track new messages for onTurnComplete.newUIMessages.
                     // Only push for genuinely new ids — HITL continuations
-                    // whose incoming wire id matches an existing hydrated
-                    // entry are state advances on an old message, not new
-                    // messages. The non-hydrate branch below has the same
-                    // semantic (push only on append, not on merge).
+                    // whose incoming wire id matches an existing entry are
+                    // state advances on an old message, not new messages.
+                    // We compare against `previouslyKnownMessageIds`
+                    // captured BEFORE hydration, not against `hydrated`:
+                    // the canonical hydrate pattern pushes the incoming
+                    // user message into the persisted chain and returns
+                    // it, so the new id IS in `hydrated`, which would
+                    // wrongly drop every fresh user turn from
+                    // `newUIMessages`. The non-hydrate branch below has
+                    // the same "push only on append" semantic via its
+                    // own append-vs-replace path.
                     if (
                       currentWirePayload.trigger === "submit-message" &&
                       cleanedUIMessages.length > 0
                     ) {
                       const lastUI = cleanedUIMessages[cleanedUIMessages.length - 1]!;
                       const matchedExisting =
                         lastUI.id !== undefined &&
-                        hydrated.some((m) => m.id === lastUI.id);
+                        previouslyKnownMessageIds.has(lastUI.id);
                       if (!matchedExisting) {
                         turnNewUIMessages.push(lastUI);
                         const lastModel = (await toModelMessages([lastUI]))[0];

packages/trigger-sdk/test/mockChatAgent.test.ts

@@ -179,6 +179,64 @@ describe("mockChatAgent", () => {
     }
   });
 
+  it("hydrate path: fresh user message lands in onTurnComplete.newUIMessages", async () => {
+    // The dedup that suppresses HITL-continuation pushes to
+    // `turnNewUIMessages` must compare against the PRE-hydration
+    // accumulator, not the post-hydration chain. A `hydrateMessages`
+    // hook that pushes the incoming user message into its persisted
+    // chain (the canonical pattern documented in
+    // /ai-chat/lifecycle-hooks#hydratemessages and the
+    // `upsertIncomingMessage` helper) would otherwise see the new
+    // user message in the returned `hydrated` array for every fresh
+    // turn, causing the dedup to wrongly fire and drop the user
+    // message from `newUIMessages` / `newMessages`.
+    const stored: any[] = [];
+
+    const model = new MockLanguageModelV3({
+      doStream: async () => ({ stream: textStream("hi") }),
+    });
+
+    let capturedNewUIMessages: any[] | undefined;
+    const agent = chat.agent({
+      id: "mockChatAgent.hydrate-newui-fresh-user",
+      hydrateMessages: async ({ trigger, incomingMessages }) => {
+        // Canonical pattern: push the incoming user message into the
+        // persisted chain and return the chain. Mirrors what
+        // `upsertIncomingMessage` does.
+        if (trigger === "submit-message" && incomingMessages.length > 0) {
+          const newMsg = incomingMessages[incomingMessages.length - 1]!;
+          const exists = newMsg.id
+            ? stored.some((m) => m.id === newMsg.id)
+            : false;
+          if (!exists) stored.push(newMsg);
+        }
+        return [...stored];
+      },
+      onTurnComplete: async ({ newUIMessages }) => {
+        capturedNewUIMessages = newUIMessages;
+      },
+      run: async ({ messages, signal }) => {
+        return streamText({ model, messages, abortSignal: signal });
+      },
+    });
+
+    const harness = mockChatAgent(agent, { chatId: "test-hydrate-newui-fresh-user" });
+    try {
+      await harness.sendMessage(userMessage("hello", "u-fresh"));
+      await new Promise((r) => setTimeout(r, 50));
+
+      // `newUIMessages` for a fresh user turn must contain BOTH the
+      // user message and the assistant response. The bug surfaces as
+      // length=1 (assistant only — user dropped by the wrong dedup).
+      expect(capturedNewUIMessages).toBeDefined();
+      const roles = capturedNewUIMessages!.map((m: any) => m.role);
+      expect(roles).toEqual(["user", "assistant"]);
+      expect(capturedNewUIMessages![0]!.id).toBe("u-fresh");
+    } finally {
+      await harness.close();
+    }
+  });
+
   it("merges HITL tool answer onto head assistant when AI SDK regenerates the id", async () => {
     // Regression for TRI-9137: customers (Arena AI) report that the AI SDK
     // intermittently mints a fresh id on `addToolOutput` resume, breaking