diff --git a/.changeset/approval-raw-edits-resume.md b/.changeset/approval-raw-edits-resume.md new file mode 100644 index 0000000000..305d549292 --- /dev/null +++ b/.changeset/approval-raw-edits-resume.md @@ -0,0 +1,9 @@ +--- +'@tanstack/ai': minor +'@tanstack/ai-persistence': patch +--- + +Approved edits to tool arguments stay raw until the final input check. + +- `@tanstack/ai`: a resume no longer rejects approved `editedArgs` that do not fit the tool's `inputSchema`. The edit goes through `onBeforeToolCall` and the final input check, like the model's arguments. If it still does not fit, the approval is used, the tool does not run, and the model gets the input error as the tool result. The transcript keeps the approved edit as it was sent, but only when the call passed the check. +- `@tanstack/ai-persistence`: `withPersistence` binds the client-tool phase of an approved tool call to its approval. After a restart, the client result resumes with the approval that came before it, and a stale or missing approval fails the resume. diff --git a/docs/config.json b/docs/config.json index dd6239b947..aac1ab8435 100644 --- a/docs/config.json +++ b/docs/config.json @@ -503,7 +503,7 @@ "label": "Tool Approval", "to": "interrupts/tool-approval", "addedAt": "2026-08-04", - "updatedAt": "2026-08-14" + "updatedAt": "2026-10-13" }, { "label": "Multiple Interrupts", diff --git a/docs/interrupts/tool-approval.md b/docs/interrupts/tool-approval.md index 0c724bd274..3a60952f87 100644 --- a/docs/interrupts/tool-approval.md +++ b/docs/interrupts/tool-approval.md @@ -194,7 +194,7 @@ Now the decision carries a payload, and approval can also replace the arguments: interrupt.resolveInterrupt(true, { payload: { note: 'Reviewed' } }) // Approve, but replace the arguments first. editedArgs is a full replacement, -// not a merge, and is validated against the tool's inputSchema. +// not a merge. The tool's inputSchema checks it just before the tool runs. interrupt.resolveInterrupt(true, { editedArgs: { amount: 12, recipient: 'Ada' }, payload: { note: 'Capped to policy' }, @@ -206,7 +206,17 @@ interrupt.resolveInterrupt(false, { payload: { reason: 'Too large' } }) Only approval accepts `editedArgs`. Without an `approvalSchema` the boolean shorthand `resolveInterrupt(true)` / `resolveInterrupt(false)` is all you need. -The server re-validates the whole decision before it runs the tool. +The server checks the whole decision again when it resumes. + +The tool's `inputSchema` checks `editedArgs` after middleware, just before the +tool runs, the same as the model's arguments. On the server, edited arguments +that do not fit do not fail the resume. The approval is used, the tool does not +run, and the model gets the input error as the result of that call. See +[Server Tools](../tools/server-tools). + +The browser client also checks `editedArgs` against the tool's `inputSchema` +when it has the schema, and it refuses an edit that does not fit before it +sends the decision. ## Consume the decision on the server diff --git a/packages/ai-harness/tests/interrupt-restart.test.ts b/packages/ai-harness/tests/interrupt-restart.test.ts index ed840ddb6f..6dd2e1b786 100644 --- a/packages/ai-harness/tests/interrupt-restart.test.ts +++ b/packages/ai-harness/tests/interrupt-restart.test.ts @@ -329,6 +329,122 @@ describe.each([false, true])( }, ) + it('consumes an approval after a final input error without executing its tool', async () => { + const stores = memoryPersistence().stores + const persistence = durable + ? { + stores: { + log: memoryLogStore(), + runs: stores.runs, + metadata: stores.metadata, + interrupts: stores.interrupts, + }, + } + : { stores } + const remove = removeTool() + const before = vi.fn() + const main = mockAdapter([ + () => toolCall('remove', { path: 'a.txt' }, 'invalid-final-call'), + () => text('Input rejected.'), + ]) + const config = defineHarness({ + name: 'invalid-final-approval', + adapter: main.adapter, + tools: [remove.tool], + middleware: [{ name: 'observe', onBeforeToolCall: before }], + }) + const host = createHarnessHost({ persistence }) + const session = await host.open(config, { threadId: THREAD }) + const initial = await session.send('remove') + const approval = initial.interrupts?.[0] + if (!approval) throw new Error('The turn did not stop for approval.') + const receipt = await session.resolve([ + { + interruptId: approval.id, + status: 'resolved', + payload: { approved: true, editedArgs: { path: {} } }, + }, + ]) + await expect( + session.operation(receipt.operationId ?? ''), + ).resolves.toEqual({ text: 'Input rejected.', media: [] }) + expect(before).toHaveBeenCalledTimes(1) + expect(remove.execute).not.toHaveBeenCalled() + expect(session.snapshot().pendingInterrupts).toEqual([]) + expect((await stores.interrupts.get(approval.id))?.status).toBe( + 'resolved', + ) + await host.close() + }) + + it('keeps approval after schema drift and retries with the original schema', async () => { + const stores = memoryPersistence().stores + const persistence = durable + ? { + stores: { + log: memoryLogStore(), + runs: stores.runs, + metadata: stores.metadata, + interrupts: stores.interrupts, + }, + } + : { stores } + const original = removeTool() + const main = mockAdapter([ + () => toolCall('remove', { path: 'a.txt' }, 'drift-call'), + () => text('Done.'), + ]) + const config = defineHarness({ + name: 'schema-drift', + adapter: main.adapter, + tools: [original.tool], + }) + const first = createHarnessHost({ persistence }) + const session = await first.open(config, { threadId: THREAD }) + const initial = await session.send('remove') + const approval = initial.interrupts?.[0] + if (!approval) throw new Error('The turn did not stop for approval.') + await first.close() + const changedExecute = vi.fn(async () => 'must not run') + const changedTool = toolDefinition({ + name: 'remove', + description: 'Remove', + needsApproval: true, + inputSchema: z.object({ path: z.number() }), + }).server(changedExecute) + const changed = createHarnessHost({ persistence }) + const changedSession = await changed.open( + defineHarness({ + name: 'schema-drift', + adapter: main.adapter, + tools: [changedTool], + }), + { threadId: THREAD }, + ) + const rejected = await changedSession.resolve([ + { interruptId: approval.id, status: 'resolved', payload: true }, + ]) + await expect( + changedSession.operation(rejected.operationId ?? ''), + ).rejects.toThrow() + expect(changedExecute).not.toHaveBeenCalled() + expect(original.execute).not.toHaveBeenCalled() + expect( + changedSession.snapshot().pendingInterrupts.map((item) => item.id), + ).toEqual([approval.id]) + await changed.close() + const next = createHarnessHost({ persistence }) + const retrySession = await next.open(config, { threadId: THREAD }) + const retry = await retrySession.resolve([ + { interruptId: approval.id, status: 'resolved', payload: true }, + ]) + await expect( + retrySession.operation(retry.operationId ?? ''), + ).resolves.toEqual({ text: 'Done.', media: [] }) + expect(original.execute).toHaveBeenCalledTimes(1) + await next.close() + }) + it.each(['cancel', 'close'])( 'keeps approval when %s stops an admitted resolve before execution', async (kind) => { @@ -521,6 +637,92 @@ describe.each([false, true])( }, ) + it('recovers the new client phase after metadata replacement fails', async () => { + const stores = memoryPersistence().stores + const persistence = durable + ? { + stores: { + log: memoryLogStore(), + runs: stores.runs, + metadata: stores.metadata, + interrupts: stores.interrupts, + }, + } + : { stores } + const tool = toolDefinition({ + name: 'check', + description: 'Check', + needsApproval: true, + inputSchema: z.object({ count: z.number() }), + outputSchema: z.object({ ok: z.boolean() }), + }).client() + const main = mockAdapter([ + () => toolCall('check', { count: 1 }, 'phase-call'), + () => text('Checked.'), + ]) + const config = defineHarness({ + name: 'phase-recovery', + adapter: main.adapter, + tools: [tool], + }) + const first = createHarnessHost({ persistence }) + const session = await first.open(config, { threadId: THREAD }) + const initial = await session.send('check') + const approval = initial.interrupts?.[0] + if (!approval) throw new Error('The turn did not stop for approval.') + const metadata = vi + .spyOn(stores.metadata, 'set') + .mockRejectedValueOnce(new Error('metadata replacement failed')) + const accepted = await session.resolve([ + { + interruptId: approval.id, + status: 'resolved', + payload: { approved: true, editedArgs: { count: '2' } }, + }, + ]) + await session.operation(accepted.operationId ?? '') + metadata.mockRestore() + const pending = session.snapshot().pendingInterrupts + expect(pending).toHaveLength(1) + const client = pending[0] + if (!client) throw new Error('The client phase is missing.') + expect(client.id).not.toBe(approval.id) + expect(client.metadata).toMatchObject({ + kind: 'client_tool', + input: { count: 2 }, + }) + expect((await stores.interrupts.get(approval.id))?.status).toBe( + 'resolved', + ) + await first.close() + const next = createHarnessHost({ persistence }) + const reopened = await next.open(config, { threadId: THREAD }) + expect( + reopened.snapshot().pendingInterrupts.map((item) => item.id), + ).toEqual([client.id]) + const wrong = await reopened.resolve([ + { + interruptId: client.id, + status: 'resolved', + payload: { ok: 'wrong' }, + }, + ]) + await expect( + reopened.operation(wrong.operationId ?? ''), + ).rejects.toThrow() + expect( + reopened.snapshot().pendingInterrupts.map((item) => item.id), + ).toEqual([client.id]) + const valid = await reopened.resolve([ + { interruptId: client.id, status: 'resolved', payload: { ok: true } }, + ]) + await expect( + reopened.operation(valid.operationId ?? ''), + ).resolves.toEqual({ text: 'Checked.', media: [] }) + expect(reopened.snapshot().pendingInterrupts).toEqual([]) + await next.close() + }) + it('keeps foreign pending descriptors outside the harness batch', async () => { const stores = memoryPersistence().stores const persistence = durable diff --git a/packages/ai-harness/tests/replay-adapter-parity.test.ts b/packages/ai-harness/tests/replay-adapter-parity.test.ts new file mode 100644 index 0000000000..ed746a0738 --- /dev/null +++ b/packages/ai-harness/tests/replay-adapter-parity.test.ts @@ -0,0 +1,372 @@ +import { describe, expect, it, vi } from 'vitest' +import { EventType, toolDefinition } from '@tanstack/ai' +import { memoryLogStore, memoryPersistence } from '@tanstack/ai-persistence' +import { z } from 'zod' +import { createHarnessHost, defineHarness } from '../src' +import { gate, mockAdapter, text, toolCall } from './helpers' +import type { MessageSource, StreamChunk } from '@tanstack/ai' + +const firstSource: MessageSource = { + provider: 'first-provider', + api: 'first-api', + model: 'first-requested', +} +const secondSource: MessageSource = { + provider: 'second-provider', + api: 'second-api', + model: 'second-requested', +} + +function sourced( + chunks: Array, + source: MessageSource, + responseId: string, +): Array { + return chunks.map((chunk) => ({ + ...chunk, + metadata: { + ...chunk.metadata, + tanstack: { + ...chunk.metadata?.tanstack, + source, + ...(chunk.type === EventType.RUN_FINISHED + ? { responseId, model: 'reported-model' } + : {}), + }, + }, + })) +} + +describe('saved harness replay across turn adapters', () => { + it('uses the override with saved source, genuine identity and paired raw tool arguments', async () => { + const persistence = memoryPersistence() + const host = createHarnessHost({ persistence }) + const execute = vi.fn(async () => 'found') + const lookup = toolDefinition({ + name: 'lookup', + description: 'Look up a fact', + inputSchema: z.object({ q: z.string() }), + }).server(execute) + const first = mockAdapter([ + () => + sourced( + toolCall('lookup', { q: 'first' }, 'call-original'), + firstSource, + 'generation-tool', + ), + () => sourced(text('First answer'), firstSource, 'generation-first'), + ]) + const second = mockAdapter([ + () => sourced(text('Second answer'), secondSource, 'generation-second'), + ]) + const session = await host.open( + defineHarness({ + name: 'test/replay-parity', + adapter: first.adapter, + tools: [lookup], + }), + { threadId: 'replay' }, + ) + try { + await session.send('first') + const before = JSON.parse(JSON.stringify(await session.transcript())) + await session.send('second', { overrides: { adapter: second.adapter } }) + expect(execute).toHaveBeenCalledTimes(1) + expect(first.calls).toHaveLength(2) + expect(second.calls).toHaveLength(1) + const received = second.calls[0].messages + const call = received.find( + (message: { toolCalls?: ReadonlyArray }) => + message.toolCalls?.length, + ) + expect(call.toolCalls).toEqual([ + { + id: 'call-original', + type: 'function', + function: { name: 'lookup', arguments: '{"q":"first"}' }, + metadata: { tanstack: { source: firstSource } }, + }, + ]) + expect( + received.find( + (message: { toolCallId?: string }) => + message.toolCallId === 'call-original', + ), + ).toMatchObject({ role: 'tool', content: 'found' }) + expect(call.metadata?.tanstack?.source).toEqual(firstSource) + expect( + received.find( + (message: { content?: unknown }) => + message.content === 'First answer', + ).metadata.tanstack, + ).toMatchObject({ source: firstSource, responseId: 'generation-first' }) + expect( + JSON.parse( + JSON.stringify((await session.transcript()).slice(0, before.length)), + ), + ).toEqual(before) + const saved = await persistence.stores.messages.loadThread('replay') + expect(saved.at(-1)?.metadata?.tanstack).toMatchObject({ + source: secondSource, + responseId: 'generation-second', + }) + } finally { + await host.close() + } + }) + + it.each([ + { label: 'valid', raw: '{"count":2}', count: 2, executes: 1 }, + { label: 'coercible', raw: '{"count":"3"}', count: 3, executes: 1 }, + { label: 'rejected', raw: '{"count":{}}', count: undefined, executes: 0 }, + { label: 'empty', raw: '', count: undefined, executes: 0 }, + ])( + 'checks $label raw input before dispatch and preserves it in saved history', + async ({ raw, count, executes }) => { + const execute = vi.fn(async (input: { count: number }) => input.count) + const tool = toolDefinition({ + name: 'measure', + description: 'Measure a count', + inputSchema: z.object({ count: z.number().finite() }), + }).server(execute) + const model = mockAdapter([ + () => + toolCall('measure', {}, 'raw-call').map((chunk) => + chunk.type === EventType.TOOL_CALL_ARGS + ? { ...chunk, delta: raw } + : chunk, + ), + () => text('Checked'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/raw-replay', + adapter: model.adapter, + tools: [tool], + }), + { threadId: 'raw' }, + ) + try { + await session.send('measure') + expect(execute).toHaveBeenCalledTimes(executes) + if (executes) expect(execute.mock.calls[0]?.[0]).toEqual({ count }) + const saved = await session.transcript() + expect( + saved.find((message) => message.toolCalls?.length)?.toolCalls?.[0] + ?.function.arguments, + ).toBe(raw) + expect( + model.calls[1].messages.find( + (message: { toolCalls?: ReadonlyArray }) => + message.toolCalls?.length, + ).toolCalls[0].function.arguments, + ).toBe(raw) + } finally { + await host.close() + } + }, + ) +}) + +describe.each([false, true])( + 'approval replay after a host restart (durable: %s)', + (durable) => { + it('runs approved raw edits through middleware and one authored transform', async () => { + const memory = memoryPersistence() + const { runs, metadata, interrupts } = memory.stores + const persistence = durable + ? { stores: { log: memoryLogStore(), runs, metadata, interrupts } } + : memory + const seen: Array = [] + let transforms = 0 + const execute = vi.fn(async (input: { count: number }) => input) + const tool = toolDefinition({ + name: 'approve_checked', + description: 'Approve a checked count', + needsApproval: true, + inputSchema: z.object({ + count: z.number().transform((value) => { + transforms++ + return value + 1 + }), + }), + }).server(execute) + const model = mockAdapter([ + () => toolCall('approve_checked', { count: 1 }, 'approved-call'), + () => text('Resumed'), + ]) + const config = defineHarness({ + name: 'test/approval-replay', + adapter: model.adapter, + tools: [tool], + middleware: [ + { + name: 'raw-edit', + onBeforeToolCall(_ctx, hook) { + seen.push(hook.args) + return { type: 'transformArgs', args: { count: 4 } } + }, + }, + ], + }) + const first = createHarnessHost({ persistence }) + let approvalId = '' + try { + const session = await first.open(config, { threadId: 'approved' }) + const stopped = await session.send('approve') + approvalId = stopped.interrupts?.[0]?.id ?? '' + expect(approvalId).not.toBe('') + expect(seen).toEqual([]) + expect(execute).not.toHaveBeenCalled() + } finally { + await first.close() + } + const next = createHarnessHost({ persistence }) + try { + const session = await next.open(config, { threadId: 'approved' }) + const wrong = await session.resolve([ + { + interruptId: 'wrong-correlation', + status: 'resolved', + payload: true, + }, + ]) + expect(wrong.status).toBe('accepted') + await expect( + session.operation(wrong.operationId ?? ''), + ).rejects.toThrow('Missing resume entry') + expect(seen).toEqual([]) + expect(execute).not.toHaveBeenCalled() + const before = transforms + const receipt = await session.resolve([ + { + interruptId: approvalId, + status: 'resolved', + payload: { approved: true, editedArgs: { count: '2' } }, + }, + ]) + expect(receipt.status).toBe('accepted') + await session.operation(receipt.operationId ?? '') + expect(seen).toEqual([{ count: '2' }]) + expect(transforms - before).toBe(1) + expect(execute).toHaveBeenCalledTimes(1) + expect(execute.mock.calls[0]?.[0]).toEqual({ count: 5 }) + expect( + (await session.transcript()).find( + (message) => message.toolCalls?.length, + )?.toolCalls?.[0]?.function.arguments, + ).toBe('{"count":"2"}') + } finally { + await next.close() + } + }) + + it('runs no middleware or execution for a denied approval', async () => { + const memory = memoryPersistence() + const { runs, metadata, interrupts } = memory.stores + const persistence = durable + ? { stores: { log: memoryLogStore(), runs, metadata, interrupts } } + : memory + const execute = vi.fn(async () => 'Should not run') + const hook = vi.fn() + const tool = toolDefinition({ + name: 'deny_checked', + description: 'Check approval', + needsApproval: true, + inputSchema: z.object({}), + }).server(execute) + const model = mockAdapter([ + () => toolCall('deny_checked', {}, 'denied-call'), + () => text('Denied'), + ]) + const host = createHarnessHost({ persistence }) + try { + const session = await host.open( + defineHarness({ + name: 'test/denied-replay', + adapter: model.adapter, + tools: [tool], + middleware: [{ name: 'count-hooks', onBeforeToolCall: hook }], + }), + { threadId: 'denied' }, + ) + const stopped = await session.send('deny') + const id = stopped.interrupts?.[0]?.id + if (!id) throw new Error('Missing approval') + const receipt = await session.resolve([ + { interruptId: id, status: 'resolved', payload: false }, + ]) + await session.operation(receipt.operationId ?? '') + expect(hook).not.toHaveBeenCalled() + expect(execute).not.toHaveBeenCalled() + } finally { + await host.close() + } + }) + }, +) + +it('detaches a local event reader without aborting the model turn', async () => { + const hold = gate() + const stores = memoryPersistence().stores + const persistence = { + stores: { + log: memoryLogStore(), + runs: stores.runs, + metadata: stores.metadata, + }, + } + const main = mockAdapter([ + () => + (async function* () { + yield { + type: EventType.RUN_STARTED, + runId: 'provider-run', + threadId: 'detach', + timestamp: Date.now(), + } + await hold.opened + yield* sourced( + text('Finished after detach.'), + firstSource, + 'generation-detached', + ) + })(), + ]) + const host = createHarnessHost({ persistence }) + const session = await host.open( + defineHarness({ name: 'detach-replay', adapter: main.adapter }), + { threadId: 'detach' }, + ) + const turn = session.send('continue') + const local = new AbortController() + const detached = (async () => { + for await (const entry of turn.events({ + from: '0', + signal: local.signal, + })) { + if (entry.event.type === EventType.RUN_STARTED) { + local.abort() + break + } + } + })() + await detached + expect(turn.status()).toBe('running') + hold.open() + await expect(turn).resolves.toEqual({ + text: 'Finished after detach.', + media: [], + }) + const stored = await session.transcript() + const assistant = stored.find((message) => message.role === 'assistant') + expect(assistant?.error).toBeUndefined() + expect(assistant?.metadata?.tanstack?.stopReason).not.toBe('aborted') + expect(assistant?.metadata?.tanstack?.responseId).toBe('generation-detached') + const events = [] + for await (const entry of turn.events({ from: '0' })) events.push(entry.event) + expect(events.some((event) => event.type === EventType.RUN_ERROR)).toBe(false) + expect(main.calls).toHaveLength(1) + await host.close() +}) diff --git a/packages/ai-persistence/src/middleware.ts b/packages/ai-persistence/src/middleware.ts index 7291857964..7f3dbdc723 100644 --- a/packages/ai-persistence/src/middleware.ts +++ b/packages/ai-persistence/src/middleware.ts @@ -78,6 +78,7 @@ import type { ChatTranscriptStores, InterruptCommitEntry, InterruptRecord, + InterruptStore, MessagePage, RunStore, } from './types' @@ -717,11 +718,134 @@ function durableGenericFailure( ]) } +const APPROVAL_CONTEXT = 'tanstack:approvalContext' + +type ApprovalContext = { + v: 1 + approvalInterruptId: string + approvalRunId: string + generation: number + toolCallId: string + toolName: string + inputSchemaHash: string + approvalSchemaHash: string + responseSchemaHash: string +} + +function readApprovalContext(value: unknown): ApprovalContext | undefined { + const record = objectValue(value) + const fields = [ + 'v', + 'approvalInterruptId', + 'approvalRunId', + 'generation', + 'toolCallId', + 'toolName', + 'inputSchemaHash', + 'approvalSchemaHash', + 'responseSchemaHash', + ] + if ( + !record || + Array.isArray(value) || + Object.keys(record).length !== fields.length || + fields.some((key) => !Object.hasOwn(record, key)) + ) + return undefined + const { + v, + approvalInterruptId, + approvalRunId, + generation, + toolCallId, + toolName, + inputSchemaHash, + approvalSchemaHash, + responseSchemaHash, + } = record + if ( + v !== 1 || + typeof generation !== 'number' || + !Number.isSafeInteger(generation) || + generation < 0 || + typeof approvalInterruptId !== 'string' || + !approvalInterruptId || + typeof approvalRunId !== 'string' || + !approvalRunId || + typeof toolCallId !== 'string' || + !toolCallId || + typeof toolName !== 'string' || + !toolName || + typeof inputSchemaHash !== 'string' || + !inputSchemaHash || + typeof approvalSchemaHash !== 'string' || + !approvalSchemaHash || + typeof responseSchemaHash !== 'string' || + !responseSchemaHash + ) + return undefined + return { + v, + approvalInterruptId, + approvalRunId, + generation, + toolCallId, + toolName, + inputSchemaHash, + approvalSchemaHash, + responseSchemaHash, + } +} + +function clientInterruptPayload( + interrupt: unknown, + state: RunStateEntry | undefined, +): Record { + const payload = interruptPayload(interrupt) + if (!isPersistedInterruptDescriptor(payload)) return payload + const client = readInterruptBinding(payload) + if (client?.kind !== 'client-tool-execution') return payload + const metadata = { ...objectValue(payload.metadata) } + delete metadata[APPROVAL_CONTEXT] + for (const previous of state?.pendingResumes?.pending ?? []) { + if (!isPersistedInterruptDescriptor(previous.payload)) continue + const binding = readInterruptBinding(previous.payload) + const resume = state?.pendingResumes?.resumeByInterruptId.get( + previous.interruptId, + ) + const decision = objectValue(resume?.payload) + if ( + binding?.kind !== 'tool-approval' || + binding.toolCallId !== client.toolCallId || + binding.toolName !== client.toolName || + binding.generation !== client.generation || + resume?.status !== 'resolved' || + (resume.payload !== true && decision?.approved !== true) + ) + continue + const context: ApprovalContext = { + v: 1, + approvalInterruptId: previous.interruptId, + approvalRunId: previous.runId, + generation: binding.generation, + toolCallId: binding.toolCallId, + toolName: binding.toolName, + inputSchemaHash: binding.inputSchemaHash, + approvalSchemaHash: binding.approvalSchemaHash, + responseSchemaHash: binding.responseSchemaHash, + } + metadata[APPROVAL_CONTEXT] = context + break + } + return { ...payload, metadata } +} + async function durableGenericResumeState( ctx: ChatMiddlewareContext, pending: Array, resume: ReadonlyArray, tools: Array, + interrupts: InterruptStore, ): Promise { const registry = getGenericInterruptDefinitionRegistry(ctx, { optional: true, @@ -859,6 +983,112 @@ async function durableGenericResumeState( if (validated.errors.length > 0 || !validated.resumeToolState) { throw new InterruptResumeValidationError(validated.errors) } + const approvals = new Map(validated.resumeToolState.approvals) + for (const record of records) { + if (record.binding.kind !== 'client-tool-execution') continue + const toolName = record.binding.toolName + const tool = tools.find((candidate) => candidate.name === toolName) + if (!tool?.needsApproval) continue + const metadata = objectValue(objectValue(record.payload)?.metadata) + const hasContext = !!metadata && Object.hasOwn(metadata, APPROVAL_CONTEXT) + const context = hasContext + ? readApprovalContext(metadata?.[APPROVAL_CONTEXT]) + : undefined + const previous = await interrupts.get( + context?.approvalInterruptId ?? `approval_${record.binding.toolCallId}`, + ) + const isMatch = + !!previous && + previous.status === 'resolved' && + previous.threadId === ctx.threadId && + (hasContext + ? !!context && context.approvalRunId === previous.runId + : previous.runId === record.binding.interruptedRunId) && + isPersistedInterruptDescriptor(previous.payload) + // A record with no approval context (written before this binding, or by + // a path that does not add it) resumes as it did before. + if (!hasContext && !isMatch) continue + if ( + !previous || + !isMatch || + !isPersistedInterruptDescriptor(previous.payload) + ) { + const persisted = pending.find( + (entry) => entry.interruptId === record.interruptId, + ) + if (!persisted) + throw new Error( + `Missing persisted client interrupt ${record.interruptId}`, + ) + throw durableGenericFailure( + ctx, + persisted, + `Client tool ${record.binding.toolName} has no matching resolved approval.`, + ) + } + const binding = readInterruptBinding(previous.payload) + if ( + !binding || + binding.kind !== 'tool-approval' || + binding.toolCallId !== record.binding.toolCallId || + binding.toolName !== record.binding.toolName || + previous.payload.id !== previous.interruptId || + binding.interruptedRunId !== previous.runId || + binding.generation !== record.binding.generation || + binding.interruptId !== previous.interruptId || + (hasContext && + (!context || + context.approvalInterruptId !== previous.interruptId || + context.toolCallId !== binding.toolCallId || + context.toolName !== binding.toolName || + context.generation !== binding.generation || + context.inputSchemaHash !== binding.inputSchemaHash || + context.approvalSchemaHash !== binding.approvalSchemaHash || + context.responseSchemaHash !== binding.responseSchemaHash)) + ) { + throw durableGenericFailure( + ctx, + previous, + `Client tool ${record.binding.toolName} has a stale resolved approval.`, + ) + } + const restored = await validateInterruptResumeBatch({ + threadId: ctx.threadId, + interruptedRunId: binding.interruptedRunId, + generation: binding.generation, + pending: [ + { + interruptId: previous.interruptId, + payload: previous.payload, + binding, + }, + ], + resume: [ + { + interruptId: previous.interruptId, + status: 'resolved', + payload: previous.response, + }, + ], + tools, + ...(previous.resolvedAt !== undefined + ? { now: previous.resolvedAt } + : {}), + }) + if (restored.errors.length > 0 || !restored.resumeToolState) + throw new InterruptResumeValidationError(restored.errors) + const decision = restored.resumeToolState.approvals?.get(binding.toolCallId) + if ( + decision !== true && + (typeof decision !== 'object' || !decision.approved) + ) + throw durableGenericFailure( + ctx, + previous, + `Client tool ${record.binding.toolName} was not approved.`, + ) + approvals.set(binding.toolCallId, decision) + } type GenericRecord = PendingInterruptResumeRecord & { binding: Extract< PendingInterruptResumeRecord['binding'], @@ -900,6 +1130,7 @@ async function durableGenericResumeState( genericRecords.sort((left, right) => left.batchIndex - right.batchIndex) return { ...validated.resumeToolState, + approvals, genericInterruptRequests: new Map( genericRecords.flatMap(({ record }) => record.genericRequest @@ -2268,6 +2499,7 @@ export function withPersistence( ownedPending, config.resume ?? [], config.tools, + persistence.stores.interrupts, ) patch.resume = [] if (resumeToolState || genericResumeState) { @@ -2464,14 +2696,20 @@ export function withPersistence( if (wantsInterrupts && persistence.stores.interrupts) { // The run reached a new interrupt boundary, so the resumes it consumed // are committed before the fresh interrupts are recorded. + // Read the approvals of this run before the commit clears them, so + // a client-tool phase keeps the approval it came from. + const payloads = chunk.outcome.interrupts.map((interrupt) => ({ + interrupt, + payload: clientInterruptPayload(interrupt, state), + })) await commitPendingResumes(state, persistence.stores.interrupts) - for (const interrupt of chunk.outcome.interrupts) { + for (const { interrupt, payload } of payloads) { await persistence.stores.interrupts.create({ interruptId: interrupt.id, runId: ctx.runId, threadId: ctx.threadId, requestedAt: Date.now(), - payload: interruptPayload(interrupt), + payload, }) } } diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index 5fead436a2..b9c427f13b 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -2636,7 +2636,9 @@ class TextEngine< ) // Consume the async generator, yielding custom events and collecting the return value + const rawEdits = this.captureApprovedArgumentEdits(pendingToolCalls) const executionResult = yield* this.drainToolCallGenerator(generator) + this.retainApprovedArgumentEdits(executionResult, rawEdits) // Check if middleware aborted during pending tool execution if (this.isMiddlewareAborted()) { @@ -2870,7 +2872,9 @@ class TextEngine< ) // Consume the async generator, yielding custom events and collecting the return value + const rawEdits = this.captureApprovedArgumentEdits(toolCalls) const executionResult = yield* this.drainToolCallGenerator(generator) + this.retainApprovedArgumentEdits(executionResult, rawEdits) this.middlewareCtx.phase = 'afterTools' @@ -3896,8 +3900,7 @@ class TextEngine< } } const actionable = - executed ?? - (yield* this.drainToolCallGenerator(this.boundaryToolRequests(toolCalls))) + executed ?? (yield* this.drainBoundaryToolRequests(toolCalls)) for (const chunk of this.buildToolResultChunks( actionable.results, finishEvent, @@ -3915,6 +3918,17 @@ class TextEngine< return true } + private async *drainBoundaryToolRequests( + toolCalls: ReadonlyArray, + ): AsyncGenerator { + const rawEdits = this.captureApprovedArgumentEdits(toolCalls) + const result = yield* this.drainToolCallGenerator( + this.boundaryToolRequests(toolCalls), + ) + this.retainApprovedArgumentEdits(result, rawEdits) + return result + } + /** * The approval previews and client requests a boundary pauses with. They get * the same middleware hook and final input check as the tool phase, and a @@ -5599,6 +5613,164 @@ class TextEngine< this.middlewareCtx.modelOptions = config.modelOptions } + /** Copy only lossless JSON raw edits. Opaque execution inputs stay untouched. */ + private captureApprovedArgumentEdits( + calls: ReadonlyArray, + ): Map { + const unsupported = Symbol('unsupported') + const copy = (value: unknown, ancestors = new Set()): unknown => { + if ( + value === null || + typeof value === 'string' || + typeof value === 'boolean' + ) + return value + if (typeof value === 'number') + return Number.isFinite(value) && !Object.is(value, -0) + ? value + : unsupported + if (typeof value !== 'object' || ancestors.has(value)) return unsupported + const array = Array.isArray(value) + const prototype = Object.getPrototypeOf(value) + if ( + array + ? prototype !== Array.prototype + : prototype !== Object.prototype && prototype !== null + ) + return unsupported + const properties = Object.getOwnPropertyDescriptors(value) + const keys = Reflect.ownKeys(properties) + if (array && keys.length !== properties.length?.value + 1) + return unsupported + ancestors.add(value) + const result: Record | Array = array + ? [] + : Object.create(null) + for (const key of keys) { + if (array && key === 'length') continue + const property = typeof key === 'string' ? properties[key] : undefined + if ( + !property || + !property.enumerable || + !('value' in property) || + (array && !/^(0|[1-9]\d*)$/.test(String(key))) + ) + return unsupported + const copied = copy(property.value, ancestors) + if (copied === unsupported) return unsupported + Object.defineProperty(result, key, { + value: copied, + enumerable: true, + configurable: true, + writable: true, + }) + } + ancestors.delete(value) + return result + } + const snapshots = new Map() + for (const call of calls) { + const resolution = + this.resumeApprovals.get(call.id) ?? + this.resumeApprovals.get(`approval_${call.id}`) ?? + this.initialApprovals.get(call.id) ?? + this.initialApprovals.get(`approval_${call.id}`) + if ( + typeof resolution !== 'object' || + !resolution.approved || + // ponytail: `== null` matches `editedArgs ?? input` in executeToolCalls. + resolution.editedArgs == null + ) + continue + try { + const snapshot = copy(resolution.editedArgs) + if (snapshot !== unsupported) { + const historyCalls = this.messages.flatMap((message) => + message.role === 'assistant' ? (message.toolCalls ?? []) : [], + ) + const selected = historyCalls.includes(call) + ? call + : historyCalls + .slice() + .reverse() + .find( + (candidate) => + candidate.id === call.id && + candidate.function.name === call.function.name && + candidate.function.arguments === call.function.arguments, + ) + if (selected) snapshots.set(selected, snapshot) + } + } catch { + // An opaque value can execute without changing its JSON history. + } + } + return snapshots + } + + /** + * Keep checked raw edits in the selected history occurrence. A call passed + * the check when it waits for the client or ran with its input. A middleware + * skip has no input, and a failed check is an output error. + */ + private retainApprovedArgumentEdits( + executed: ExecuteToolCallsResult, + snapshots: Map, + ): void { + if (snapshots.size === 0) return + const selectedIds = new Set([ + ...executed.results + .filter( + (entry) => + entry.input !== undefined && entry.state !== 'output-error', + ) + .map((entry) => entry.toolCallId), + ...executed.needsClientExecution.map((entry) => entry.toolCallId), + ]) + const historyCalls = this.messages.flatMap((message) => + message.role === 'assistant' ? (message.toolCalls ?? []) : [], + ) + const update = (messages: Array): Array => { + const selected = new Map() + const calls = messages.flatMap((message) => + message.role === 'assistant' ? (message.toolCalls ?? []) : [], + ) + for (const [call, snapshot] of snapshots) { + if (!selectedIds.has(call.id)) continue + if (calls.includes(call)) { + selected.set(call, snapshot) + continue + } + const sameCall = (candidate: ToolCall) => + candidate.id === call.id && + candidate.function.name === call.function.name && + candidate.function.arguments === call.function.arguments + const occurrence = historyCalls.filter(sameCall).indexOf(call) + const equivalent = calls.filter(sameCall)[occurrence] + if (equivalent) selected.set(equivalent, snapshot) + } + return messages.map((message) => { + if (message.role !== 'assistant' || !message.toolCalls) return message + let changed = false + const toolCalls = message.toolCalls.map((call) => { + if (!selected.has(call)) return call + const argumentsJson = JSON.stringify(selected.get(call)) + if (argumentsJson === undefined) return call + if (argumentsJson === call.function.arguments) return call + changed = true + return { + ...call, + function: { ...call.function, arguments: argumentsJson }, + } + }) + return changed ? { ...message, toolCalls } : message + }) + } + this.messages = update(this.messages) + this.providerMessages = update(this.providerMessages) + this.middlewareCtx.messages = this.messages + } + private setToolPhase(phase: ToolPhaseResult): void { this.toolPhase = phase } diff --git a/packages/ai/src/interrupt-resume.ts b/packages/ai/src/interrupt-resume.ts index e78029a381..11a13380e8 100644 --- a/packages/ai/src/interrupt-resume.ts +++ b/packages/ai/src/interrupt-resume.ts @@ -636,17 +636,8 @@ export async function validateInterruptResumeBatch( `Approval ${record.interruptId} cannot edit arguments without an input schema.`, ), ) - } else { - await pushSchemaIssues({ - request: input, - errors, - interruptId: record.interruptId, - schema: tool.inputSchema, - value: envelope.editedArgs, - code: 'invalid-edited-args', - label: `Approval ${record.interruptId} edited arguments are invalid`, - }) } + // Edited input stays raw until middleware and the final execution check. } const branch = approved ? approval.branches.approve diff --git a/packages/ai/tests/chat.test.ts b/packages/ai/tests/chat.test.ts index eac41978be..2193827923 100644 --- a/packages/ai/tests/chat.test.ts +++ b/packages/ai/tests/chat.test.ts @@ -1514,19 +1514,19 @@ describe('chat()', () => { }) it('validates an entire ephemeral batch before executing any tool', async () => { + const beforeTool = vi.fn() const firstExecute = vi.fn().mockReturnValue({ ok: true }) const secondExecute = vi.fn().mockReturnValue({ ok: true }) const { adapter, calls } = createMockAdapter({ iterations: [[ev.runStarted(), ev.runFinished('stop')]], }) - // A Standard Schema input so the library validates the edited args. Raw - // JSON Schema is intentionally not validated by the library (the app owns - // that), so the invalid `editedArgs` below is caught via this path. + // Reject the malformed decision before tool middleware or execution. const inputSchema = z.object({ action: z.string() }) const chunks = await collectChunks( chat({ adapter, + middleware: [{ name: 'batch-check', onBeforeToolCall: beforeTool }], threadId: 'thread-1', runId: 'continuation-run', parentRunId: 'interrupted-run', @@ -1573,21 +1573,22 @@ describe('chat()', () => { status: 'resolved', payload: { approved: true, - editedArgs: { action: 1 }, + editedArgs: { action: 'first-edit' }, }, }, { interruptId: 'approval_call_2', status: 'resolved', payload: { - approved: true, - editedArgs: { action: 2 }, + approved: 'yes', + editedArgs: { action: 'second-edit' }, }, }, ], }) as AsyncIterable, ) + expect(beforeTool).not.toHaveBeenCalled() expect(firstExecute).not.toHaveBeenCalled() expect(secondExecute).not.toHaveBeenCalled() expect(calls).toHaveLength(0) @@ -1597,13 +1598,9 @@ describe('chat()', () => { metadata: expect.objectContaining({ tanstack: expect.objectContaining({ interruptErrors: expect.arrayContaining([ - expect.objectContaining({ - interruptId: 'approval_call_1', - code: 'invalid-edited-args', - }), expect.objectContaining({ interruptId: 'approval_call_2', - code: 'invalid-edited-args', + code: 'invalid-payload', }), expect.objectContaining({ scope: 'batch', diff --git a/packages/ai/tests/tool-input-validation.test.ts b/packages/ai/tests/tool-input-validation.test.ts index 14b5dec60a..4c5ac6f6e3 100644 --- a/packages/ai/tests/tool-input-validation.test.ts +++ b/packages/ai/tests/tool-input-validation.test.ts @@ -1525,3 +1525,57 @@ describe('repairToolInput', () => { ]) }) }) + +describe('an approved edit in the transcript', () => { + const run = async ( + editedArgs: unknown, + middleware: Array = [], + ) => { + const { adapter, calls } = createMockAdapter({ + iterations: [[ev.runStarted(), ev.runFinished()]], + }) + const execute = vi.fn(async (_input: unknown) => 'ok') + await collectChunks( + chat({ + adapter, + parentRunId: 'old', + messages: history({ count: 1 }), + tools: [ + { ...clientCheck(z.object({ count: z.number() }), true), execute }, + ], + resume: [ + { + interruptId: 'approval_call', + status: 'resolved', + payload: { approved: true, editedArgs }, + }, + ], + middleware, + }), + ) + const sent = calls[0]?.messages.flatMap((message) => + message.role === 'assistant' ? (message.toolCalls ?? []) : [], + ) + return { execute, args: sent?.[0]?.function.arguments } + } + + it('keeps an edit that the tool ran with', async () => { + const { execute, args } = await run({ count: 2 }) + expect(execute.mock.calls[0]?.[0]).toEqual({ count: 2 }) + expect(args).toBe('{"count":2}') + }) + + it('keeps the model input for a null edit, which the tool ignores', async () => { + const { execute, args } = await run(null) + expect(execute.mock.calls[0]?.[0]).toEqual({ count: 1 }) + expect(args).toBe('{"count":1}') + }) + + it('keeps the model input when middleware skips the call', async () => { + const { execute, args } = await run({ count: 2 }, [ + { onBeforeToolCall: () => ({ type: 'skip', result: 'skipped' }) }, + ]) + expect(execute).not.toHaveBeenCalled() + expect(args).toBe('{"count":1}') + }) +})