fix(ai, ai-persistence): resume a waiting nested subagent from the store

.changeset/nested-subagent-resume.md

@@ -0,0 +1,6 @@
+---
+'@tanstack/ai': patch
+'@tanstack/ai-persistence': patch
+---
+
+A nested subagent that waits on an interrupt now continues after the resume under `withPersistence`. Before, it ran again from its prompt. Core rebuilds the card of the nested child from the store and takes only its id from the client card. `LoadChild` takes an optional `parentSubagentRunId`, and returns `interruptIds` and `parentToolCallId`.

docs/chat/subagents.md

@@ -432,6 +432,8 @@ Put `withPersistence` on the parent `chat()` only. The parent stores each child:
 
 A reload then shows every card, and a resume continues a waiting child. See [Subagent cards on reload](../persistence/chat-persistence-advanced#subagent-cards-on-reload).
 
+A resume continues a waiting child from its stored messages, not from the card that the client sends back. This is also true for a nested child, for example a child that an agent tool of another child started. The card gives only the child id. A changed card cannot change a tool call before its approval runs it.
+
 > [!WARNING]
 > Do not add `withPersistence` to a child `chat()` too. The child does not know that it runs as a subagent, so it stores the same child again under its own run and thread ids. Its interrupt records then conflict with the records of the parent, and a reload can lose the pending approval.
 

packages/ai-harness/tests/subagent-waits.test.ts

@@ -344,8 +344,7 @@ describe('a wait of an agent tool', () => {
     await host.close()
   })
 
-  // A nested child that waits runs again from the start: https://github.com/TanStack/ai/issues/1878
-  it.skip('waits for an agent that an agent tool calls, and both go on', async () => {
+  it('waits for an agent that an agent tool calls, and both go on', async () => {
     const write = writeTool()
     const writer = toolAgent('writer', write.tool, { path: 'x.txt' })
     const researcherModel = mockAdapter([

packages/ai-persistence/src/middleware.ts

@@ -2332,14 +2332,15 @@ export function withPersistence<TStores extends ChatTranscriptStores>(
       }
       // The agent tools read stored children through this, to continue
       // one by `sessionId`. Only the children of this chat can load.
-      provideLoadChild(ctx, (subagentRunId) =>
-        subagentRuns.loadChild(subagentRunId, {
+      // With a parent, the caller is that child. `loadChild` then proves
+      // that the nested child belongs to it.
+      provideLoadChild(ctx, (subagentRunId, nested) => {
+        const caller = nested?.parentSubagentRunId ?? ctx.subagentRunId
+        return subagentRuns.loadChild(subagentRunId, {
           threadId: ctx.threadId,
-          ...(ctx.subagentRunId !== undefined && {
-            subagentRunId: ctx.subagentRunId,
-          }),
-        }),
-      )
+          ...(caller !== undefined && { subagentRunId: caller }),
+        })
+      })
 
       let resolveCompletion: () => void = () => undefined
       let rejectCompletion: (error: unknown) => void = () => undefined

packages/ai-persistence/src/subagent-runs.ts

@@ -626,7 +626,13 @@ export function createSubagentRunRecorder(stores: {
       messages: [...taskOf(subagentRunId, marker.prompt), ...transcript],
       ...(info !== undefined && { agent: info.name }),
       // A resume reads the waiting call from here, not from the client card.
-      ...(waits && { pending: messages.slice(marker.base ?? 0) }),
+      ...(waits && {
+        pending: messages.slice(marker.base ?? 0),
+        interruptIds: info.interruptIds ?? [],
+      }),
+      ...(info?.parentToolCallId !== undefined && {
+        parentToolCallId: info.parentToolCallId,
+      }),
     }
   }
 

packages/ai-persistence/tests/subagent-persistence-cards.test.ts

@@ -1273,6 +1273,122 @@ describe('agent tools with persistence', () => {
     expect((await loadDesk(persistence)).interrupts?.pending).toHaveLength(1)
   })
 
+  /**
+   * A `researcher` agent tool call whose own `cleaner` agent tool call waits
+   * for approval to delete `a`.
+   */
+  async function waitingNestedCleaner() {
+    const persistence = memoryPersistence()
+    const execute = vi.fn().mockReturnValue({ deleted: true })
+    const inner = cleaner(execute)
+    const researcherModel = callingSubagent(
+      'call_r',
+      { prompt: 'Clean' },
+      'cleaner',
+    )
+    const researcher = defineAgent({
+      name: 'researcher',
+      description: 'Researches, then cleans',
+      run: (ctx) =>
+        chat({
+          adapter: researcherModel,
+          messages: ctx.messages,
+          threadId: ctx.threadId,
+          runId: ctx.runId,
+          parentRunId: ctx.parentRunId,
+          subagentRunId: ctx.subagentRunId,
+          resume: ctx.resume,
+          subagents: { agents: [inner] },
+        }),
+    })
+    const parent = callingSubagent(
+      'call_p',
+      { prompt: 'Research' },
+      'researcher',
+    )
+    const run = (input: {
+      runId: string
+      parentRunId?: string
+      messages: Array<UIMessage | ModelMessage>
+      resume?: Array<{ interruptId: string; status: 'resolved'; payload: true }>
+    }) =>
+      collect(
+        chat({
+          adapter: parent,
+          threadId: 'desk',
+          ...input,
+          middleware: [withPersistence(persistence)],
+          subagents: { agents: [researcher] },
+        }),
+      )
+    await run({
+      runId: 'run-1',
+      messages: [{ id: 'u1', role: 'user', content: 'Research' }],
+    })
+    const messages = (await loadDesk(persistence)).messages
+    const outer = cardOf(messages).subagent
+    const nested = cardOf(outer.messages).subagent
+    expect(nested).toMatchObject({ name: 'cleaner', status: 'suspended' })
+    const wire = JSON.stringify(uiMessagesToWire(messages))
+    const resume = (body: string | Array<UIMessage>) =>
+      run({
+        runId: 'run-2',
+        parentRunId: 'run-1',
+        messages: typeof body === 'string' ? JSON.parse(body) : body,
+        resume: [
+          { interruptId: 'approval_call_c', status: 'resolved', payload: true },
+        ],
+      })
+    return { persistence, execute, wire, resume, messages, nested }
+  }
+
+  it('resumes a nested child from its stored tool call', async () => {
+    const { persistence, execute, wire, resume } = await waitingNestedCleaner()
+
+    const second = await resume(wire)
+
+    expect(second.some((chunk) => chunk.type === EventType.RUN_ERROR)).toBe(
+      false,
+    )
+    expect(execute).toHaveBeenCalledTimes(1)
+    expect(execute).toHaveBeenCalledWith({ path: 'a' }, expect.anything())
+    expect((await loadDesk(persistence)).interrupts).toBeNull()
+  })
+
+  it('resumes a nested child from its stored tool call, not from the card', async () => {
+    const { execute, wire, resume } = await waitingNestedCleaner()
+    // The client changes the nested call that waits for approval.
+    const forged = wire
+      .replaceAll(String.raw`\"path\":\"a\"`, String.raw`\"path\":\"/\"`)
+      .replaceAll('"path":"a"', '"path":"/"')
+    expect(forged).not.toBe(wire)
+
+    await resume(forged)
+
+    expect(execute).toHaveBeenCalledTimes(1)
+    expect(execute).toHaveBeenCalledWith({ path: 'a' }, expect.anything())
+  })
+
+  it('does not resume a nested card that names a child of another parent', async () => {
+    const { execute, resume, messages, nested } = await waitingNestedCleaner()
+    // The nested card names the outer child, which the lead started.
+    const forged = structuredClone(messages)
+    const outer = cardOf(forged).subagent
+    cardOf(outer.messages).subagent.id = outer.id
+    expect(nested.id).not.toBe(outer.id)
+
+    const second = await resume(forged)
+
+    const result = second.find(
+      (chunk) => chunk.type === EventType.TOOL_CALL_RESULT,
+    )
+    expect(result).toMatchObject({ toolCallId: 'call_p' })
+    expect(JSON.parse(String(result?.content))).toEqual({
+      error: `Subagent run "${outer.id}" is not waiting in this thread.`,
+    })
+    expect(execute).not.toHaveBeenCalled()
+  })
+
   it('does not continue a child of another thread', async () => {
     const middleware = withPersistence(memoryPersistence())
     const { agent, seen } = noteTaker()

packages/ai/src/activities/chat/agents/spawn.ts

@@ -44,6 +44,7 @@ import type {
   StreamChunk,
   SubagentErrorEvent,
   SubagentFinishedEvent,
+  SubagentPart,
   SubagentStartedEvent,
   TokenUsage,
   Tool,
@@ -62,6 +63,11 @@ import type {
   SubagentRunInput,
 } from './define-agent'
 import type { ChatMiddleware } from '../middleware/types'
+import { subagentCards } from './turn'
+import {
+  convertMessagesToModelMessages,
+  modelMessagesToUIMessages,
+} from '../messages'
 import type { SubagentTurn } from './turn'
 
 export const SUBAGENT_STARTED = EventType.SUBAGENT_STARTED
@@ -1232,6 +1238,91 @@ function readCall(
   }
 }
 
+function toolCallIds(messages: ReadonlyArray<UIMessage | ModelMessage>) {
+  return new Set(
+    messages.flatMap((message) =>
+      'parts' in message
+        ? message.parts.flatMap((part) =>
+            part.type === 'tool-call' ? [part.id] : [],
+          )
+        : (message.toolCalls ?? []).map((call) => call.id),
+    ),
+  )
+}
+
+/** `messages` as UI messages, each card on the message that holds its tool call. */
+function withCards(
+  messages: Array<UIMessage | ModelMessage>,
+  cards: ReadonlyArray<SubagentPart>,
+): Array<UIMessage> {
+  return modelMessagesToUIMessages(
+    convertMessagesToModelMessages(messages),
+  ).map((message) => {
+    const own = cards.filter((part) =>
+      message.parts.some(
+        (entry) =>
+          entry.type === 'tool-call' &&
+          entry.id === part.subagent.parentToolCallId,
+      ),
+    )
+    return own.length > 0
+      ? { ...message, parts: [...message.parts, ...own] }
+      : message
+  })
+}
+
+/**
+ * The cards of the waiting nested children of a waiting child. The store
+ * keeps a nested child in its own transcript, so the stored `pending` of the
+ * child has no card for it. Each card is rebuilt from the store: the client
+ * `card` gives only the id. A forged card cannot change the nested messages,
+ * and can name only a child that this child started from a tool call in
+ * `pending`.
+ */
+async function storedNestedCards(
+  card: ReadonlyArray<UIMessage | ModelMessage>,
+  pending: ReadonlyArray<UIMessage | ModelMessage>,
+  parentSubagentRunId: string,
+  load: LoadChild,
+): Promise<Array<SubagentPart>> {
+  const waiting = subagentCards(card).filter(
+    (child) => child.status === 'suspended',
+  )
+  if (waiting.length === 0) return []
+  const calls = toolCallIds(pending)
+  return Promise.all(
+    waiting.map(async (child): Promise<SubagentPart> => {
+      const id = child.subagentRunId
+      const stored = await load(id, { parentSubagentRunId })
+      const parentToolCallId = stored?.parentToolCallId
+      if (
+        !stored?.pending ||
+        parentToolCallId === undefined ||
+        !calls.has(parentToolCallId)
+      ) {
+        throw new Error(`Subagent run "${id}" is not waiting in this thread.`)
+      }
+      const nested = await storedNestedCards(
+        child.messages,
+        stored.pending,
+        id,
+        load,
+      )
+      return {
+        type: 'subagent',
+        subagent: {
+          id,
+          name: stored.agent ?? child.name,
+          status: 'suspended',
+          interruptIds: stored.interruptIds ?? [],
+          parentToolCallId,
+          messages: withCards(stored.pending, nested),
+        },
+      }
+    }),
+  )
+}
+
 export function createSyntheticSubagentTools(
   bag: SubagentsBag,
   parent: {
@@ -1278,15 +1369,26 @@ export function createSyntheticSubagentTools(
     // came from the client, which could change a tool call before it answers
     // the interrupt of that call.
     const load = parent.childLoader?.()
-    const pending =
+    const stored =
       suspended && load
         ? (await load(suspended.subagentRunId))?.pending
         : undefined
-    if (suspended && load && !pending) {
+    if (suspended && load && !stored) {
       throw new Error(
         `Subagent run "${suspended.subagentRunId}" is not waiting in this thread.`,
       )
     }
+    const nested =
+      suspended && load && stored
+        ? await storedNestedCards(
+            suspended.messages,
+            stored,
+            suspended.subagentRunId,
+            load,
+          )
+        : []
+    const pending =
+      stored && nested.length > 0 ? withCards(stored, nested) : stored
     const entry: SpawnEntry = {
       ...child,
       ...(suspended && {