diff --git a/.changeset/nested-subagent-resume.md b/.changeset/nested-subagent-resume.md new file mode 100644 index 0000000000..a948686be9 --- /dev/null +++ b/.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`. diff --git a/docs/chat/subagents.md b/docs/chat/subagents.md index 327aa42b7a..b2a7585cdc 100644 --- a/docs/chat/subagents.md +++ b/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. diff --git a/packages/ai-harness/tests/subagent-waits.test.ts b/packages/ai-harness/tests/subagent-waits.test.ts index ba3b095f43..c0e62a31b1 100644 --- a/packages/ai-harness/tests/subagent-waits.test.ts +++ b/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([ diff --git a/packages/ai-persistence/src/middleware.ts b/packages/ai-persistence/src/middleware.ts index 9886cd2c84..38fbae500d 100644 --- a/packages/ai-persistence/src/middleware.ts +++ b/packages/ai-persistence/src/middleware.ts @@ -2332,14 +2332,15 @@ export function withPersistence( } // 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 diff --git a/packages/ai-persistence/src/subagent-runs.ts b/packages/ai-persistence/src/subagent-runs.ts index 8bb109413d..5916109567 100644 --- a/packages/ai-persistence/src/subagent-runs.ts +++ b/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, + }), } } diff --git a/packages/ai-persistence/tests/subagent-persistence-cards.test.ts b/packages/ai-persistence/tests/subagent-persistence-cards.test.ts index a8cccaf85c..b789f29a1c 100644 --- a/packages/ai-persistence/tests/subagent-persistence-cards.test.ts +++ b/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 + 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) => + 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() diff --git a/packages/ai/src/activities/chat/agents/spawn.ts b/packages/ai/src/activities/chat/agents/spawn.ts index 4ed7807e28..6ed92991d2 100644 --- a/packages/ai/src/activities/chat/agents/spawn.ts +++ b/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) { + 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, + cards: ReadonlyArray, +): Array { + 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, + pending: ReadonlyArray, + parentSubagentRunId: string, + load: LoadChild, +): Promise> { + 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 => { + 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 && { diff --git a/packages/ai/src/activities/chat/agents/turn.ts b/packages/ai/src/activities/chat/agents/turn.ts index 1acd7f2da0..affe3f2d8b 100644 --- a/packages/ai/src/activities/chat/agents/turn.ts +++ b/packages/ai/src/activities/chat/agents/turn.ts @@ -83,6 +83,50 @@ function groupInterruptIds(group: SubagentWireGroup): Array { ] } +/** + * The subagent cards in `messages`: UI subagent parts, then wire groups. + * Nested cards stay inside their parent card. + */ +export function subagentCards(messages: ReadonlyArray): Array< + Omit & { + /** The interrupt ids of the child and its nested children. */ + ids: Set + plan: unknown + } +> { + const parts = messages.flatMap((message) => + 'parts' in message + ? message.parts.filter((part) => part.type === 'subagent') + : [], + ) + return [ + ...parts.map(({ subagent }) => ({ + subagentRunId: subagent.id, + name: subagent.name, + status: subagent.status, + ...(subagent.parentToolCallId !== undefined && { + parentToolCallId: subagent.parentToolCallId, + }), + messages: subagent.messages, + text: uiText(subagent.messages), + ids: new Set(partInterruptIds({ type: 'subagent', subagent })), + plan: planOf(subagent.metadata), + })), + ...splitSubagentWire(messages).groups.map((group) => ({ + subagentRunId: group.id, + name: group.info.name, + status: group.info.status, + ...(group.info.parentToolCallId !== undefined && { + parentToolCallId: group.info.parentToolCallId, + }), + messages: group.messages, + text: subagentWireText(group.messages), + ids: new Set(groupInterruptIds(group)), + plan: planOf(group.info.metadata), + })), + ] +} + /** * Find the children of the trailing assistant turn that `resume` answers. * Returns undefined when no child owns a resume entry. @@ -97,46 +141,11 @@ export function readSubagentTurn( message.role === 'user' && wireSubagentRunId(message) === undefined, ) const before = messages.slice(0, lastUser + 1) - const turn = messages.slice(lastUser + 1) - - const found: Array & { ids: Set }> = - [] - let plan: unknown - for (const message of turn) { - if (!('parts' in message)) continue - for (const part of message.parts) { - if (part.type !== 'subagent') continue - plan ??= planOf(part.subagent.metadata) - found.push({ - subagentRunId: part.subagent.id, - name: part.subagent.name, - status: part.subagent.status, - ...(part.subagent.parentToolCallId !== undefined && { - parentToolCallId: part.subagent.parentToolCallId, - }), - messages: part.subagent.messages, - text: uiText(part.subagent.messages), - ids: new Set(partInterruptIds(part)), - }) - } - } - for (const group of splitSubagentWire(turn).groups) { - plan ??= planOf(group.info.metadata) - found.push({ - subagentRunId: group.id, - name: group.info.name, - status: group.info.status, - ...(group.info.parentToolCallId !== undefined && { - parentToolCallId: group.info.parentToolCallId, - }), - messages: group.messages, - text: subagentWireText(group.messages), - ids: new Set(groupInterruptIds(group)), - }) - } + const found = subagentCards(messages.slice(lastUser + 1)) + const plan = found.find((card) => card.plan !== undefined)?.plan const owned = new Set() - const children = found.map(({ ids, ...child }) => { + const children = found.map(({ ids, plan: _plan, ...child }) => { const entries = resume.filter((entry) => ids.has(entry.interruptId)) for (const entry of entries) owned.add(entry.interruptId) return { ...child, resume: entries } diff --git a/packages/ai/src/activities/chat/middleware/load-child.ts b/packages/ai/src/activities/chat/middleware/load-child.ts index 9c29f827d0..762897919b 100644 --- a/packages/ai/src/activities/chat/middleware/load-child.ts +++ b/packages/ai/src/activities/chat/middleware/load-child.ts @@ -8,8 +8,16 @@ import type { ModelMessage, UIMessage } from '../../../types' * child by `sessionId`. When `agent` is set, only that agent can continue * the child. A middleware that provides it must list `LoadChildCapability` * in `provides`. Only then do the agent tools take `sessionId`. + * + * With `parentSubagentRunId`, it loads a nested child: a child that the child + * `parentSubagentRunId` started. It returns `undefined` when that child did + * not start it. A loader that ignores the option fails closed for nested + * children. */ -export type LoadChild = (subagentRunId: string) => Promise< +export type LoadChild = ( + subagentRunId: string, + options?: { parentSubagentRunId?: string }, +) => Promise< | { messages: Array /** The name of the agent the child ran under. */ @@ -21,6 +29,10 @@ export type LoadChild = (subagentRunId: string) => Promise< * tool does not resume the child. */ pending?: Array + /** While the child waits: the ids of the interrupts it waits on. */ + interruptIds?: Array + /** The tool call of the parent that started the child, if any. */ + parentToolCallId?: string } | undefined >