fix(sdk): tolerate AI SDK assistant-id regen on HITL addToolOutput resume

When the AI SDK regenerates the assistant message id on an
addToolOutput-driven HITL continuation, our id-merge in hydrateMessages
fails to attach the tool answer to the existing head — duplicating the
assistant in the accumulator. Reported by Arena AI, who maintains a
content-match workaround in their hydrateMessages keyed on tool_call_id.

Add a run-scoped `toolCallId -> head messageId` map populated whenever
an assistant message containing tool parts lands in the accumulator.
The submit-message id-merge now falls back to this map when id-match
fails: walk the incoming message's tool parts, look up by toolCallId,
rewrite the incoming id back to the recorded head, then retry id-match.

Could not reproduce on current AI SDK 6.0.116 (id is preserved through
addToolOutput) but ship the mapping so the merge stays robust against
older versions and edge cases we haven't observed. Customer-side
content-match workarounds become unnecessary.

Closes TRI-9137.

Author: ericallam

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

@@ -1272,6 +1272,56 @@ const chatUIStreamPerTurnKey = locals.create<ChatUIMessageStreamOptions<UIMessag
   "chat.uiMessageStreamOptions.perTurn"
 );
 
+/**
+ * Run-scoped `toolCallId → assistant messageId` map. Records the head
+ * assistant id whenever the accumulator absorbs an assistant message
+ * containing tool parts. Used as a fallback in the id-merge for
+ * incoming tool-answer messages — if the AI SDK regenerates the
+ * assistant id on a HITL `addToolOutput` resume, we look up the
+ * original head id by `toolCallId` and rewrite it before the merge.
+ *
+ * Customer-side workaround for the same case is documented in Arena
+ * AI's chat-agent task; lifting it into the SDK so customers don't
+ * have to. See TRI-9137.
+ * @internal
+ */
+const chatToolCallToMessageIdKey = locals.create<Map<string, string>>(
+  "chat.toolCallToMessageId"
+);
+
+function recordToolCallIdsFromMessage(message: { id?: string; role?: string; parts?: unknown[] } | undefined) {
+  if (!message || message.role !== "assistant" || !message.id) return;
+  let map = locals.get(chatToolCallToMessageIdKey);
+  if (!map) {
+    map = new Map();
+    locals.set(chatToolCallToMessageIdKey, map);
+  }
+  for (const part of message.parts ?? []) {
+    if (typeof part !== "object" || part == null) continue;
+    const toolCallId = (part as { toolCallId?: unknown }).toolCallId;
+    if (typeof toolCallId === "string" && toolCallId.length > 0) {
+      map.set(toolCallId, message.id);
+    }
+  }
+}
+
+function rewriteIncomingIdViaToolCallMap<T extends { id?: string; parts?: unknown[] }>(
+  incoming: T
+): T {
+  const map = locals.get(chatToolCallToMessageIdKey);
+  if (!map || map.size === 0) return incoming;
+  for (const part of incoming.parts ?? []) {
+    if (typeof part !== "object" || part == null) continue;
+    const toolCallId = (part as { toolCallId?: unknown }).toolCallId;
+    if (typeof toolCallId !== "string" || toolCallId.length === 0) continue;
+    const headId = map.get(toolCallId);
+    if (headId && headId !== incoming.id) {
+      return { ...incoming, id: headId };
+    }
+  }
+  return incoming;
+}
+
 // ---------------------------------------------------------------------------
 // Token usage helpers (internal)
 // ---------------------------------------------------------------------------
@@ -4539,18 +4589,36 @@ function chatAgent<
                       // IDs match because we always pass generateMessageId + originalMessages
                       // to toUIMessageStream, so the backend's start chunk carries the same
                       // messageId that the frontend uses.
+                      //
+                      // Fallback for HITL `addToolOutput` continuations where the AI SDK
+                      // regenerates the assistant id (Arena AI report, TRI-9137): if the
+                      // id-match fails, look up the head messageId via toolCallId and
+                      // rewrite the incoming id before retrying. The mapping is
+                      // populated whenever an assistant containing tool parts lands in
+                      // the accumulator.
                       let replaced = false;
-                      for (const incoming of cleanedUIMessages) {
-                        const idx = accumulatedUIMessages.findIndex(
+                      for (const raw of cleanedUIMessages) {
+                        let incoming = raw;
+                        let idx = accumulatedUIMessages.findIndex(
                           (m) => m.id === incoming.id
                         );
+                        if (idx === -1) {
+                          const rewritten = rewriteIncomingIdViaToolCallMap(incoming);
+                          if (rewritten.id !== incoming.id) {
+                            incoming = rewritten as typeof raw;
+                            idx = accumulatedUIMessages.findIndex(
+                              (m) => m.id === incoming.id
+                            );
+                          }
+                        }
                         if (idx !== -1) {
                           accumulatedUIMessages[idx] = incoming as TUIMessage;
                           replaced = true;
                         } else {
                           accumulatedUIMessages.push(incoming as TUIMessage);
                           turnNewUIMessages.push(incoming as TUIMessage);
                         }
+                        recordToolCallIdsFromMessage(incoming);
                       }
                       if (replaced) {
                         // Reconvert all model messages since a replacement changes the structure
@@ -5006,6 +5074,12 @@ function chatAgent<
                     }
                     turnNewUIMessages.push(capturedResponseMessage);
                     locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
+                    // Record toolCallId → head messageId so a HITL
+                    // continuation next turn can recover the head id
+                    // even if the AI SDK regenerates it. See
+                    // `chatToolCallToMessageIdKey` for the full
+                    // rationale (TRI-9137).
+                    recordToolCallIdsFromMessage(capturedResponseMessage);
                     try {
                       const responseModelMessages = await toModelMessages([
                         stripProviderMetadata(capturedResponseMessage),

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

@@ -178,6 +178,135 @@ describe("mockChatAgent", () => {
     }
   });
 
+  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
+    // id-based dedup. Our SDK records `toolCallId → head messageId` whenever
+    // an assistant with tool parts lands in the accumulator and uses that
+    // map as a fallback in the merge so a fresh-id incoming still attaches
+    // to the right head.
+    const { z } = await import("zod");
+    const { tool } = await import("ai");
+
+    const askUserTool = tool({
+      description: "Ask the user a question.",
+      inputSchema: z.object({ question: z.string() }),
+      // No execute — HITL round-trip via addToolOutput.
+    });
+
+    const HEAD_TOOL_CALL_ID = "tc_regression_9137";
+
+    // Turn 1: model emits a tool-call for askUser. No text, no finish-reason
+    // logic beyond `tool-calls`. Agent's response will carry a tool-input-
+    // available part with HEAD_TOOL_CALL_ID.
+    const turn1Stream = simulateReadableStream({
+      chunks: [
+        { type: "tool-input-start", id: HEAD_TOOL_CALL_ID, toolName: "askUser" },
+        {
+          type: "tool-input-delta",
+          id: HEAD_TOOL_CALL_ID,
+          delta: JSON.stringify({ question: "what color?" }),
+        },
+        { type: "tool-input-end", id: HEAD_TOOL_CALL_ID },
+        {
+          type: "tool-call",
+          toolCallId: HEAD_TOOL_CALL_ID,
+          toolName: "askUser",
+          input: JSON.stringify({ question: "what color?" }),
+        },
+        {
+          type: "finish",
+          finishReason: { unified: "tool-calls", raw: "tool_calls" },
+          usage: {
+            inputTokens: { total: 10, noCache: 10, cacheRead: undefined, cacheWrite: undefined },
+            outputTokens: { total: 10, text: 0, reasoning: undefined },
+          },
+        },
+      ] as LanguageModelV3StreamPart[],
+    });
+
+    // Turn 2: model produces a final text response — exercises the post-HITL
+    // continuation streamText after the tool answer is merged in.
+    const turn2Stream = textStream("blue is great");
+
+    let callIdx = 0;
+    const model = new MockLanguageModelV3({
+      doStream: async () => ({ stream: callIdx++ === 0 ? turn1Stream : turn2Stream }),
+    });
+
+    const turnsSeen: { turn: number; uiMessages: any[] }[] = [];
+
+    const agent = chat.agent({
+      id: "mockChatAgent.hitl-id-regen",
+      tools: { askUser: askUserTool },
+      onTurnComplete: async ({ turn, uiMessages }) => {
+        turnsSeen.push({
+          turn,
+          uiMessages: uiMessages.map((m) => ({
+            id: m.id,
+            role: m.role,
+            toolStates: (m.parts ?? [])
+              .filter((p: any) => typeof p?.toolCallId === "string")
+              .map((p: any) => ({ toolCallId: p.toolCallId, state: p.state })),
+          })),
+        });
+      },
+      run: async ({ messages, signal }) => {
+        return streamText({ model, messages, tools: { askUser: askUserTool }, abortSignal: signal });
+      },
+    });
+
+    const harness = mockChatAgent(agent, { chatId: "test-hitl-id-regen" });
+    try {
+      // Turn 1: user message → agent emits tool-input-available for askUser
+      await harness.sendMessage(userMessage("hi"));
+      await new Promise((r) => setTimeout(r, 50));
+
+      // Capture the head assistant id the agent produced.
+      const turn1 = turnsSeen.at(-1);
+      const headAssistant = turn1?.uiMessages.find(
+        (m) => m.role === "assistant" && m.toolStates.length > 0
+      );
+      expect(headAssistant?.id).toBeTruthy();
+      const HEAD_ID = headAssistant!.id as string;
+
+      // Turn 2: simulate AI SDK regenerating the assistant id on
+      // addToolOutput resume — fresh id, but the same toolCallId in
+      // tool-output-available state.
+      const FRESH_ID = "regenerated-by-ai-sdk-" + Math.random().toString(36).slice(2);
+      const toolAnswerMessage = {
+        id: FRESH_ID,
+        role: "assistant" as const,
+        parts: [
+          {
+            type: "tool-askUser",
+            toolCallId: HEAD_TOOL_CALL_ID,
+            state: "output-available" as const,
+            input: { question: "what color?" },
+            output: { color: "blue" },
+          },
+        ],
+      };
+      await harness.sendMessage(toolAnswerMessage as any);
+      await new Promise((r) => setTimeout(r, 50));
+
+      // The merge must rewrite FRESH_ID back to HEAD_ID via the toolCallId
+      // map, attaching the tool answer to the existing head — no duplicate.
+      const turn2 = turnsSeen.at(-1);
+      expect(turn2).toBeTruthy();
+      const assistantsWithToolCall = turn2!.uiMessages.filter(
+        (m) =>
+          m.role === "assistant" &&
+          m.toolStates.some((t: any) => t.toolCallId === HEAD_TOOL_CALL_ID)
+      );
+      expect(assistantsWithToolCall).toHaveLength(1);
+      expect(assistantsWithToolCall[0]!.id).toBe(HEAD_ID);
+      expect(turn2!.uiMessages.find((m) => m.id === FRESH_ID)).toBeUndefined();
+    } finally {
+      await harness.close();
+    }
+  });
+
   it("routes custom actions through actionSchema + onAction", async () => {
     const model = new MockLanguageModelV3({
       doStream: async () => ({ stream: textStream("ok") }),