fix(sdk): extract pendingToolCalls from raw partial before cleanupAbortedParts strips them

Author: ericallam

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

@@ -510,6 +510,13 @@ type ReplaySessionOutTailResult<TUIMessage extends UIMessage> = {
    * the tail ended cleanly (every segment closed).
    */
   partial: TUIMessage | undefined;
+  /**
+   * The trailing assistant message BEFORE `cleanupAbortedParts` ran. Same
+   * `undefined` semantics as `partial`. Use this when you need to inspect
+   * tool parts the cleanup would strip (e.g. `input-available` /
+   * `input-streaming` orphans surfaced via `pendingToolCalls`).
+   */
+  partialRaw: TUIMessage | undefined;
 };
 
 type ReplaySessionOutTailImpl = <TUIMessage extends UIMessage>(
@@ -581,7 +588,7 @@ async function replaySessionOutTail<TUIMessage extends UIMessage>(
     if (type.startsWith("trigger:")) continue;
     collected.push(chunk as UIMessageChunk);
   }
-  if (collected.length === 0) return { settled: [], partial: undefined };
+  if (collected.length === 0) return { settled: [], partial: undefined, partialRaw: undefined };
 
   // Split chunks into per-message segments. A `start` chunk demarcates the
   // beginning of an assistant message; chunks before any `start` (rare —
@@ -612,6 +619,7 @@ async function replaySessionOutTail<TUIMessage extends UIMessage>(
 
   const settled: TUIMessage[] = [];
   let partial: TUIMessage | undefined;
+  let partialRaw: TUIMessage | undefined;
   for (let i = 0; i < segments.length; i++) {
     const seg = segments[i]!;
     const isTrailing = i === segments.length - 1 && !seg.closed;
@@ -641,11 +649,16 @@ async function replaySessionOutTail<TUIMessage extends UIMessage>(
       const cleaned = cleanupAbortedParts(last as TUIMessage);
       if (cleaned.parts.length === 0) continue;
       partial = cleaned;
+      // Keep the raw pre-cleanup message too — recovery boot extracts
+      // `pendingToolCalls` from it, since `cleanupAbortedParts` strips
+      // exactly the input-streaming / input-available tool parts that
+      // we want to surface.
+      partialRaw = last as TUIMessage;
     } else {
       settled.push(last as TUIMessage);
     }
   }
-  return { settled, partial };
+  return { settled, partial, partialRaw };
 }
 
 /**
@@ -4972,6 +4985,7 @@ function chatAgent<
       let bootSnapshot: ChatSnapshotV1<TUIMessage> | undefined;
       let replayedSettled: TUIMessage[] = [];
       let replayedPartial: TUIMessage | undefined;
+      let replayedPartialRaw: TUIMessage | undefined;
       let replayedInTail: { message: TUIMessage; metadata: unknown; seqNum: number }[] = [];
       // Wire payloads to dispatch as turns before the regular session.in
       // pump kicks in. Populated by `onRecoveryBoot.recoveredTurns` (or its
@@ -5036,6 +5050,7 @@ function chatAgent<
               );
               replayedSettled = replayResult.settled;
               replayedPartial = replayResult.partial;
+              replayedPartialRaw = replayResult.partialRaw;
             } catch (error) {
               logger.warn(
                 "chat.agent: session.out replay failed; using snapshot only",
@@ -5160,7 +5175,11 @@ function chatAgent<
         let hookRecoveredTurns: TUIMessage[] | undefined;
         let hookBeforeBoot: (() => Promise<void>) | undefined;
         if (couldHavePriorState && hasRecoveredState && onRecoveryBoot) {
-          const pendingToolCalls = extractPendingToolCallsFromPartial(partialAssistant);
+          // Extract from the RAW partial (pre-cleanup). `cleanupAbortedParts`
+          // strips exactly the input-streaming / input-available tool parts
+          // we want to surface here, so the cleaned `partialAssistant` would
+          // always report zero pending tool calls.
+          const pendingToolCalls = extractPendingToolCallsFromPartial(replayedPartialRaw);
           const previousRunIdForHook = previousRunId ?? "";
           let hookResult: RecoveryBootResult<TUIMessage> | void = undefined;
           const { writer: hookWriter, flush: hookFlush } = createLazyChatWriter();

packages/trigger-sdk/src/v3/test/mock-chat-agent.ts

@@ -416,7 +416,14 @@ export function mockChatAgent(
       seededReplayChunks.length === 0
         ? []
         : ((await reduceChunksToMessages(seededReplayChunks)) as unknown[]);
-    return { settled, partial: seededReplayPartial } as never;
+    // For the mock harness, `partialRaw` is the same as `partial` — we
+    // don't model cleanupAbortedParts separately. Recovery tests that
+    // need a partialRaw distinct from partial install their own stub.
+    return {
+      settled,
+      partial: seededReplayPartial,
+      partialRaw: seededReplayPartial,
+    } as never;
   });
 
   // session.in tail override: each seeded UIMessage becomes a

packages/trigger-sdk/test/recovery-boot.test.ts

@@ -5,6 +5,7 @@ import { mockChatAgent } from "../src/v3/test/index.js";
 import { describe, expect, it, vi } from "vitest";
 import { chat } from "../src/v3/ai.js";
 import type { RecoveryBootEvent, RecoveryBootResult } from "../src/v3/ai.js";
+import { __setReplaySessionOutTailImplForTests } from "../src/v3/ai.js";
 import { simulateReadableStream, streamText } from "ai";
 import { MockLanguageModelV3 } from "ai/test";
 import type { LanguageModelV3StreamPart } from "@ai-sdk/provider";
@@ -125,6 +126,74 @@ describe("onRecoveryBoot — chat.agent recovery hook", () => {
     }
   });
 
+  it("pendingToolCalls is extracted from the RAW partial (pre-cleanupAbortedParts)", async () => {
+    // Real-world scenario: cancel-mid-tool-call. Session.out has tool-call
+    // chunks but the tool never returned. cleanupAbortedParts strips the
+    // input-available tool part from the partial used for the chain (you
+    // don't want orphan tool calls poisoning the model context), but
+    // `pendingToolCalls` should still surface what was happening.
+    const cleanedPartial = {
+      id: "a-orphan",
+      role: "assistant" as const,
+      parts: [{ type: "text" as const, text: "Let me look that up" }],
+    };
+    const rawPartial = {
+      id: "a-orphan",
+      role: "assistant" as const,
+      parts: [
+        { type: "text" as const, text: "Let me look that up" },
+        {
+          type: "tool-search" as const,
+          toolCallId: "tc-pending",
+          state: "input-available" as const,
+          input: { q: "vietnamese pho" },
+        },
+      ],
+    } as unknown as typeof cleanedPartial;
+
+    const captured: { event?: RecoveryBootEvent } = {};
+    const model = new MockLanguageModelV3({
+      doStream: async () => ({ stream: textStream("ok") }),
+    });
+    const u1 = userMessage("buffered", "u-1");
+    const agent = chat.agent({
+      id: "recovery-boot.pending-tool-from-raw",
+      onRecoveryBoot: async (event) => {
+        captured.event = event;
+        return {};
+      },
+      run: async ({ messages, signal }) =>
+        streamText({ model, messages, abortSignal: signal }),
+    });
+    const harness = mockChatAgent(agent, {
+      chatId: "pending-tool-from-raw",
+      continuation: true,
+      previousRunId: "run_prior",
+    });
+    harness.seedSessionInTail([u1 as never]);
+    // Install AFTER mockChatAgent — its constructor sets its own default
+    // override that we want to replace for this test.
+    __setReplaySessionOutTailImplForTests(async () =>
+      ({
+        settled: [],
+        partial: cleanedPartial,
+        partialRaw: rawPartial,
+      }) as never
+    );
+    try {
+      await new Promise((r) => setTimeout(r, 50));
+      expect(captured.event).toBeDefined();
+      // Cleaned partial → chain (no input-available tool part)
+      expect(captured.event!.partialAssistant?.parts).toHaveLength(1);
+      // pendingToolCalls → from raw (input-available tool part visible)
+      expect(captured.event!.pendingToolCalls).toHaveLength(1);
+      expect(captured.event!.pendingToolCalls[0]!.toolCallId).toBe("tc-pending");
+      expect(captured.event!.pendingToolCalls[0]!.toolName).toBe("search");
+    } finally {
+      await harness.close();
+    }
+  });
+
   it("fires when there are in-flight users (no partial)", async () => {
     const captured: { event?: RecoveryBootEvent } = {};
     const model = new MockLanguageModelV3({