fix(sdk): deliver the head-start handover partial when hydrateMessages is registered

ericallam · ericallam · commit f4e7fb274775 · 2026-06-11T15:03:02.000+01:00
The turn-0 handover splice only ran on the default accumulation path, so
agents registering hydrateMessages lost the warm route's step-1 response:
pure-text turns fired onTurnComplete with no assistant message, tool-call
turns re-ran step 1 from scratch under a fresh messageId, and the
head-start user message never reached the hydrate hook. The first-turn
history now reaches hydrateMessages as incoming messages, and the splice
runs after both accumulation branches, deduped by the handover messageId.
diff --git a/.changeset/chat-headstart-hydrate.md b/.changeset/chat-headstart-hydrate.md
@@ -0,0 +1,5 @@
+---
+"@trigger.dev/sdk": patch
+---
+
+Fix `chat.headStart` when `hydrateMessages` is registered. The warm route's step-1 partial now reaches the agent's accumulator on the hydrate path, so `onTurnComplete` carries the full first turn (the head-start user message included), tool-call handovers resume from step 2 instead of re-running step 1, and the assistant `messageId` stays stable across the handover.
diff --git a/packages/trigger-sdk/src/v3/ai.ts b/packages/trigger-sdk/src/v3/ai.ts
@@ -6362,6 +6362,21 @@ function chatAgent<
 
                   let cleanedUIMessages: TUIMessage[] = cleanedIncomingMessages;
 
+                  // Turn-0 head-start with hydrateMessages: the boot seeding from
+                  // `payload.headStartMessages` is non-hydrate-only, so ship the
+                  // route handler's first-turn history to the hydrate hook as
+                  // incoming messages instead (gated on the pending handover).
+                  if (
+                    turn === 0 &&
+                    hydrateMessages &&
+                    cleanedUIMessages.length === 0 &&
+                    (locals.get(chatHandoverPartialKey)?.length ?? 0) > 0 &&
+                    Array.isArray(payload.headStartMessages) &&
+                    payload.headStartMessages.length > 0
+                  ) {
+                    cleanedUIMessages = payload.headStartMessages as TUIMessage[];
+                  }
+
                   // Validate/transform UIMessages before conversion — catches malformed
                   // messages from storage or untrusted input before they reach the model.
                   // Slim wire: triggers like `regenerate-message` carry no incoming
@@ -6568,46 +6583,54 @@ function chatAgent<
                     // `preload` / `close` / `handover-prepare` and submits
                     // with no incoming message fall through with the boot-
                     // seeded accumulator unchanged.
+                  }
 
-                    if (turn === 0) {
-                      // Head-start handover splice (turn 0 only): the
-                      // `chat.handover` route handler signalled a mid-turn
-                      // handover, so splice its partial assistant response
-                      // (text + pending tool-calls + the synthesized
-                      // tool-approval round) onto the accumulator.
-                      // `streamText` then hits AI SDK's initial-tool-
-                      // execution branch, runs the agent-side tool executes,
-                      // and resumes from step 2 — skipping the first model
-                      // call (already done by the handler).
-                      //
-                      // We also synthesize a UIMessage form of the partial
-                      // assistant and push it to `accumulatedUIMessages` so
-                      // AI SDK's `processUIMessageStream` (invoked when the
-                      // run loop calls `runResult.toUIMessageStream({
-                      // onFinish })`) can initialize `state.message` from
-                      // the trailing assistant in `originalMessages`. Without
-                      // that, the `tool-output-available` chunks emitted by
-                      // the initial-tool-execution branch can't find their
-                      // matching tool-call in state and AI SDK throws
-                      // `UIMessageStreamError: No tool invocation found`.
-                      const pendingHandoverPartial = locals.get(chatHandoverPartialKey);
-                      if (pendingHandoverPartial && pendingHandoverPartial.length > 0) {
+                  if (turn === 0) {
+                    // Head-start handover splice (turn 0 only, BOTH
+                    // accumulation branches — hydrate and default): the
+                    // `chat.handover` route handler signalled a mid-turn
+                    // handover, so splice its partial assistant response
+                    // (text + pending tool-calls + the synthesized
+                    // tool-approval round) onto the accumulator.
+                    // `streamText` then hits AI SDK's initial-tool-
+                    // execution branch, runs the agent-side tool executes,
+                    // and resumes from step 2 — skipping the first model
+                    // call (already done by the handler).
+                    //
+                    // We also synthesize a UIMessage form of the partial
+                    // assistant and push it to `accumulatedUIMessages` so
+                    // AI SDK's `processUIMessageStream` (invoked when the
+                    // run loop calls `runResult.toUIMessageStream({
+                    // onFinish })`) can initialize `state.message` from
+                    // the trailing assistant in `originalMessages`. Without
+                    // that, the `tool-output-available` chunks emitted by
+                    // the initial-tool-execution branch can't find their
+                    // matching tool-call in state and AI SDK throws
+                    // `UIMessageStreamError: No tool invocation found`.
+                    const pendingHandoverPartial = locals.get(chatHandoverPartialKey);
+                    if (pendingHandoverPartial && pendingHandoverPartial.length > 0) {
+                      const handoverMessageId = locals.get(chatHandoverMessageIdKey);
+                      // Skip if the hydrated chain already persisted the
+                      // partial under the handover messageId.
+                      const alreadyInChain =
+                        handoverMessageId !== undefined &&
+                        accumulatedUIMessages.some((m) => m.id === handoverMessageId);
+                      if (!alreadyInChain) {
                         accumulatedMessages.push(...pendingHandoverPartial);
-                        const handoverMessageId = locals.get(chatHandoverMessageIdKey);
                         const partialUI = synthesizeHandoverUIMessage(
                           pendingHandoverPartial,
                           handoverMessageId
                         );
                         if (partialUI) {
                           accumulatedUIMessages.push(partialUI as TUIMessage);
                         }
-                        locals.set(chatHandoverPartialKey, []); // consume once
                       }
+                      locals.set(chatHandoverPartialKey, []); // consume once
                     }
-
-                    locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
                   }
 
+                  locals.set(chatCurrentUIMessagesKey, accumulatedUIMessages);
+
                   } // end if (trigger !== "action")
 
                   // ── Action result handling ──────────────────────────────
diff --git a/packages/trigger-sdk/src/v3/test/mock-chat-agent.ts b/packages/trigger-sdk/src/v3/test/mock-chat-agent.ts
@@ -92,6 +92,12 @@ export type MockChatAgentOptions = {
    *   `sendHandover()` / `sendHandoverSkip()` to dispatch the handover signal.
    */
   mode?: "preload" | "submit-message" | "handover-prepare" | "continuation";
+  /**
+   * First-turn UIMessage history shipped on the BOOT payload. Only
+   * meaningful with `mode: "handover-prepare"` — mirrors the
+   * `chat.headStart` route handler's `basePayload.headStartMessages`.
+   */
+  headStartMessages?: UIMessage[];
   /**
    * Pre-seed the snapshot the agent reads at run boot. The runtime's
    * snapshot read is replaced with one that returns this snapshot
@@ -501,6 +507,7 @@ export function mockChatAgent(
         metadata: clientData,
         ...(!isContinuationMode && options.continuation ? { continuation: true } : {}),
         ...(options.previousRunId ? { previousRunId: options.previousRunId } : {}),
+        ...(options.headStartMessages ? { headStartMessages: options.headStartMessages } : {}),
       };
 
       sendSessionInput = drivers.sessions.in.send;
diff --git a/packages/trigger-sdk/test/chatHandover.test.ts b/packages/trigger-sdk/test/chatHandover.test.ts
@@ -260,6 +260,202 @@ describe("chat.handover", () => {
     }
   });
 
+  it("pure-text head-start (isFinal: true) with hydrateMessages persists the partial (TRI-10715)", async () => {
+    // Same as the pure-text case above, but the customer registers
+    // `hydrateMessages` (the documented DB-as-source-of-truth pattern).
+    // The head-start user message must reach the hydrate hook as
+    // `incomingMessages`, and the warm route's partial must land in the
+    // accumulator so `onTurnComplete` carries the full first turn.
+    const runFn = vi.fn();
+    const stored: { id: string; role: string; parts: unknown[] }[] = [];
+    const hydrateIncomingRoles: string[] = [];
+    let captured:
+      | { responseId?: string; responseText?: string; roles?: string[] }
+      | undefined;
+
+    const agent = chat.agent({
+      id: "chat.handover.hydrate-pure-text",
+      hydrateMessages: async ({ incomingMessages }) => {
+        hydrateIncomingRoles.push(...incomingMessages.map((m) => m.role));
+        for (const m of incomingMessages) {
+          if (!stored.some((s) => s.id === m.id)) stored.push(m as (typeof stored)[number]);
+        }
+        return [...stored] as never;
+      },
+      onTurnComplete: ({ responseMessage, uiMessages }) => {
+        captured = {
+          responseId: responseMessage?.id,
+          responseText: (responseMessage?.parts ?? [])
+            .filter((p) => p.type === "text")
+            .map((p) => (p as { text?: string }).text || "")
+            .join(""),
+          roles: uiMessages.map((m) => m.role),
+        };
+      },
+      run: async ({ messages, signal }) => {
+        runFn();
+        return streamText({
+          model: new MockLanguageModelV3({
+            doStream: async () => ({ stream: textStream("should-not-run") }),
+          }),
+          messages,
+          abortSignal: signal,
+        });
+      },
+    });
+
+    const harness = mockChatAgent(agent, {
+      chatId: "test-handover-hydrate-final",
+      mode: "handover-prepare",
+      headStartMessages: [
+        { id: "hs-user-1", role: "user", parts: [{ type: "text", text: "say hi" }] },
+      ],
+    });
+
+    try {
+      await harness.sendHandover({
+        partialAssistantMessage: [
+          {
+            role: "assistant",
+            content: [{ type: "text", text: "Hi there, hope you're well." }],
+          },
+        ],
+        messageId: "asst-hydrate-1",
+        isFinal: true,
+      });
+      await new Promise((r) => setTimeout(r, 30));
+
+      // isFinal — the agent never calls the user's run().
+      expect(runFn).not.toHaveBeenCalled();
+
+      // The head-start user message reached the hydrate hook as incoming.
+      expect(hydrateIncomingRoles).toContain("user");
+
+      // onTurnComplete carries the full first turn: user + the warm
+      // route's assistant, under the handover messageId.
+      expect(captured).toBeDefined();
+      expect(captured!.roles).toEqual(["user", "assistant"]);
+      expect(captured!.responseId).toBe("asst-hydrate-1");
+      expect(captured!.responseText).toBe("Hi there, hope you're well.");
+    } finally {
+      await harness.close();
+    }
+  });
+
+  it("tool-call handover (isFinal: false) with hydrateMessages resumes from step 2 (TRI-10715)", async () => {
+    // Hydrate variant of the schema-only tool-call case: the spliced
+    // partial (assistant + approval round) must reach the agent's
+    // streamText so AI SDK executes the pending tool instead of
+    // re-running step 1 from scratch against an empty/short prompt.
+    const toolExecute = vi.fn(async ({ city }: { city: string }) => ({ city, temp: 22 }));
+    const weatherTool = tool({
+      description: "Look up weather",
+      inputSchema: z.object({ city: z.string() }),
+      execute: toolExecute,
+    });
+
+    const stored: { id: string; role: string; parts: unknown[] }[] = [];
+    let runMessageRoles: string[] | undefined;
+    let captured: { roles?: string[]; assistantIds?: (string | undefined)[] } | undefined;
+
+    const agent = chat.agent({
+      id: "chat.handover.hydrate-schema-only-tool",
+      hydrateMessages: async ({ incomingMessages }) => {
+        for (const m of incomingMessages) {
+          if (!stored.some((s) => s.id === m.id)) stored.push(m as (typeof stored)[number]);
+        }
+        return [...stored] as never;
+      },
+      onTurnComplete: ({ uiMessages }) => {
+        captured = {
+          roles: uiMessages.map((m) => m.role),
+          assistantIds: uiMessages.filter((m) => m.role === "assistant").map((m) => m.id),
+        };
+      },
+      run: async ({ messages, signal }) => {
+        runMessageRoles = messages.map((m) => m.role);
+        return streamText({
+          model: new MockLanguageModelV3({
+            doStream: async () => ({ stream: textStream("the weather in tokyo is 22°C") }),
+          }),
+          messages,
+          tools: { weather: weatherTool },
+          abortSignal: signal,
+        });
+      },
+    });
+
+    const harness = mockChatAgent(agent, {
+      chatId: "test-handover-hydrate-tool",
+      mode: "handover-prepare",
+      headStartMessages: [
+        { id: "hs-user-2", role: "user", parts: [{ type: "text", text: "weather in tokyo?" }] },
+      ],
+    });
+
+    try {
+      const turn = await harness.sendHandover({
+        isFinal: false,
+        messageId: "asst-hydrate-2",
+        partialAssistantMessage: [
+          {
+            role: "assistant",
+            content: [
+              { type: "text", text: "let me check the weather" },
+              {
+                type: "tool-call",
+                toolCallId: "tc-h1",
+                toolName: "weather",
+                input: { city: "tokyo" },
+              },
+              {
+                type: "tool-approval-request",
+                approvalId: "handover-approval-h1",
+                toolCallId: "tc-h1",
+              },
+            ],
+          },
+          {
+            role: "tool",
+            content: [
+              {
+                type: "tool-approval-response",
+                approvalId: "handover-approval-h1",
+                approved: true,
+              },
+            ],
+          },
+        ],
+      });
+      await new Promise((r) => setTimeout(r, 30));
+
+      // The resume prompt contained the full splice: user + partial
+      // assistant + approval round — NOT an empty/user-only prompt.
+      expect(runMessageRoles).toEqual(["user", "assistant", "tool"]);
+
+      // AI SDK's initial-tool-execution branch ran the agent-side
+      // execute (no step-1 re-run).
+      expect(toolExecute).toHaveBeenCalledWith(
+        expect.objectContaining({ city: "tokyo" }),
+        expect.anything()
+      );
+
+      // Step-2 text streamed through session.out.
+      const text = turn.chunks
+        .filter((c) => c.type === "text-delta")
+        .map((c) => (c as { delta: string }).delta)
+        .join("");
+      expect(text).toContain("tokyo");
+
+      // One assistant in the final chain, under the handover messageId.
+      expect(captured).toBeDefined();
+      expect(captured!.roles).toEqual(["user", "assistant"]);
+      expect(captured!.assistantIds).toEqual(["asst-hydrate-2"]);
+    } finally {
+      await harness.close();
+    }
+  });
+
   it("onTurnStart fires after the handover signal arrives (lazy)", async () => {
     // Hooks should not fire during the wait — only once handover lands
     // and a real turn begins. Verifies the order so customers can

-Original file line number
+Diff line change
@@ @@ -0,0 +1,5 @@ @@
 +---
 +"@trigger.dev/sdk": patch
 +---
++
 +Fix `chat.headStart` when `hydrateMessages` is registered. The warm route's step-1 partial now reaches the agent's accumulator on the hydrate path, so `onTurnComplete` carries the full first turn (the head-start user message included), tool-call handovers resume from step 2 instead of re-running step 1, and the assistant `messageId` stays stable across the handover.