diff --git a/.changeset/cancelled-client-tool-resume.md b/.changeset/cancelled-client-tool-resume.md new file mode 100644 index 0000000000..db79010ffe --- /dev/null +++ b/.changeset/cancelled-client-tool-resume.md @@ -0,0 +1,8 @@ +--- +'@tanstack/ai-persistence': patch +--- + +`withPersistence` now maps cancelled client-tool and approval resume entries +into `cancelledToolCallIds`. A resume batch that is only cancellations still +produces a `resumeToolState`, so the engine can complete the turn instead of +emitting another `client_tool_*` interrupt. diff --git a/packages/ai-persistence/src/middleware.ts b/packages/ai-persistence/src/middleware.ts index 5416130637..99df9d5411 100644 --- a/packages/ai-persistence/src/middleware.ts +++ b/packages/ai-persistence/src/middleware.ts @@ -396,6 +396,7 @@ function resumeToolStateFromPending( ): ChatResumeToolState | undefined { const approvals = new Map() const clientToolResults = new Map() + const cancelledToolCallIds = new Set() for (const interrupt of pending) { const entry = resumeByInterruptId.get(interrupt.interruptId) @@ -405,6 +406,10 @@ function resumeToolStateFromPending( const reason = stringField(interrupt.payload, 'reason') const toolCallId = stringField(interrupt.payload, 'toolCallId') + if (entry.status === 'cancelled' && toolCallId) { + cancelledToolCallIds.add(toolCallId) + } + if (kind === 'approval' || reason === 'approval_required') { approvals.set(interrupt.interruptId, resolvedApprovalDecision(entry)) continue @@ -419,8 +424,14 @@ function resumeToolStateFromPending( } } - if (approvals.size === 0 && clientToolResults.size === 0) return undefined - return { approvals, clientToolResults } + if ( + approvals.size === 0 && + clientToolResults.size === 0 && + cancelledToolCallIds.size === 0 + ) { + return undefined + } + return { approvals, clientToolResults, cancelledToolCallIds } } /** diff --git a/packages/ai-persistence/tests/interrupts.test.ts b/packages/ai-persistence/tests/interrupts.test.ts index 752bf4eab8..c848f58d3c 100644 --- a/packages/ai-persistence/tests/interrupts.test.ts +++ b/packages/ai-persistence/tests/interrupts.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it, vi } from 'vitest' -import { EventType, chat } from '@tanstack/ai' +import { EventType, chat, defineChatMiddleware } from '@tanstack/ai' import type { AnyTextAdapter, StreamChunk, Tool } from '@tanstack/ai' import { memoryPersistence } from '../src/memory' import { withPersistence } from '../src/middleware' @@ -87,6 +87,39 @@ const runFinished = (runId = 'r1'): StreamChunk => ({ timestamp: 1, }) +const toolCallFinished = (runId = 'r1'): StreamChunk => ({ + type: EventType.RUN_FINISHED, + runId, + threadId: 't1', + finishReason: 'tool_calls', + timestamp: 1, +}) + +const toolCallChunks = () => [ + runStarted(), + toolStart(), + toolArgs(), + toolCallFinished(), +] + +async function persistClientToolTurn( + persistence: ReturnType, + tools: Array, +) { + const first = mockAdapter([toolCallChunks()]) + await collect( + chat({ + adapter: first.adapter, + messages: [{ role: 'user', content: 'hi' }], + tools, + runId: 'r1', + threadId: 't1', + middleware: [withPersistence(persistence)], + }) as AsyncIterable, + ) + return first +} + const clientTool = (name: string): Tool => ({ name, description: `${name} client tool`, @@ -144,20 +177,7 @@ describe('interrupt persistence', () => { it('does not persist duplicate records before terminal interrupt outcome', async () => { const persistence = memoryPersistence() const create = vi.spyOn(persistence.stores.interrupts!, 'create') - const { adapter } = mockAdapter([ - [ - runStarted(), - toolStart(), - toolArgs(), - { - type: EventType.RUN_FINISHED, - runId: 'r1', - threadId: 't1', - finishReason: 'tool_calls', - timestamp: 1, - }, - ], - ]) + const { adapter } = mockAdapter([toolCallChunks()]) await collect( chat({ @@ -260,29 +280,9 @@ describe('interrupt persistence', () => { // interrupt. Feeding the client output then drives exactly one model call. it('applies persisted approval and client-tool resume decisions with empty client messages', async () => { const persistence = memoryPersistence() - const toolCallChunks = () => [ - runStarted(), - toolStart(), - toolArgs(), - { - type: EventType.RUN_FINISHED, - runId: 'r1', - threadId: 't1', - finishReason: 'tool_calls', - timestamp: 1, - } as StreamChunk, - ] - const first = mockAdapter([toolCallChunks()]) - await collect( - chat({ - adapter: first.adapter, - messages: [{ role: 'user', content: 'hi' }], - tools: [approvalClientTool('clientSearch')], - runId: 'r1', - threadId: 't1', - middleware: [withPersistence(persistence)], - }) as AsyncIterable, - ) + await persistClientToolTurn(persistence, [ + approvalClientTool('clientSearch'), + ]) const approvalInterrupt = await persistence.stores.interrupts!.get( 'approval_tool-call-1', @@ -373,6 +373,65 @@ describe('interrupt persistence', () => { expect(await persistence.stores.interrupts!.listPending('t1')).toEqual([]) }) + // Issue #1088: cancelling a hydrated client-tool interrupt under + // withPersistence must complete the turn. Persistence clears `config.resume` + // and must therefore put the cancelled toolCallId on `cancelledToolCallIds`. + // Otherwise the engine treats the stored tool call as unhandled and emits + // another `client_tool_*` interrupt instead of an output-error. + it('completes a cancelled client-tool resume from persisted state with empty client messages', async () => { + const persistence = memoryPersistence() + await persistClientToolTurn(persistence, [clientTool('clientSearch')]) + + const pending = await persistence.stores.interrupts!.get( + 'client_tool_tool-call-1', + ) + expect(pending?.status).toBe('pending') + + const afterCancel = mockAdapter([ + [runStarted(), text('cancelled-and-done'), runFinished('r1')], + ]) + const chunks = await collect( + chat({ + adapter: afterCancel.adapter, + messages: [], + tools: [clientTool('clientSearch')], + runId: 'r1', + threadId: 't1', + resume: [ + { + interruptId: 'client_tool_tool-call-1', + status: 'cancelled', + }, + ], + middleware: [withPersistence(persistence)], + }) as AsyncIterable, + ) + + expect(afterCancel.calls).toHaveLength(1) + expect(chunks).toContainEqual( + expect.objectContaining({ + type: EventType.TOOL_CALL_RESULT, + toolCallId: 'tool-call-1', + content: JSON.stringify({ error: 'Tool execution cancelled' }), + }), + ) + expect( + chunks.find( + (chunk) => + chunk.type === EventType.RUN_FINISHED && + chunk.outcome?.type === 'interrupt', + ), + ).toBeUndefined() + expect(chunks).toContainEqual( + expect.objectContaining({ delta: 'cancelled-and-done' }), + ) + expect( + (await persistence.stores.interrupts!.get('client_tool_tool-call-1')) + ?.status, + ).toBe('cancelled') + expect(await persistence.stores.interrupts!.listPending('t1')).toEqual([]) + }) + it('rejects invalid resume entries against pending interrupts', async () => { const persistence = memoryPersistence() const first = mockAdapter([[runStarted(), interruptFinished()]]) @@ -775,6 +834,13 @@ describe('interrupt persistence', () => { }) const run = mockAdapter([[runStarted(), text('ok'), runFinished('r1')]]) + const resumeStates: Array | undefined> = [] + const observeResumeState = defineChatMiddleware({ + name: 'observe-resume-state', + onConfig(_ctx, config) { + resumeStates.push(config.resumeToolState?.cancelledToolCallIds) + }, + }) await collect( chat({ adapter: run.adapter, @@ -782,7 +848,7 @@ describe('interrupt persistence', () => { runId: 'r1', threadId: 't1', resume: [{ interruptId: 'approval-1', status: 'cancelled' }], - middleware: [withPersistence(persistence)], + middleware: [withPersistence(persistence), observeResumeState], }) as AsyncIterable, ) @@ -790,6 +856,7 @@ describe('interrupt persistence', () => { run.calls[0] as { approvals?: ReadonlyMap } ).approvals expect(approvals?.get('approval-1')).toBe(false) + expect(resumeStates[0]?.has('tc1')).toBe(true) expect( (await persistence.stores.interrupts!.get('approval-1'))?.status, ).toBe('cancelled') @@ -806,6 +873,13 @@ describe('interrupt persistence', () => { }) const run = mockAdapter([[runStarted(), text('ok'), runFinished('r1')]]) + const resumeStates: Array | undefined> = [] + const observeResumeState = defineChatMiddleware({ + name: 'observe-resume-state', + onConfig(_ctx, config) { + resumeStates.push(config.resumeToolState?.cancelledToolCallIds) + }, + }) await collect( chat({ adapter: run.adapter, @@ -819,7 +893,7 @@ describe('interrupt persistence', () => { payload: { answer: 99 }, }, ], - middleware: [withPersistence(persistence)], + middleware: [withPersistence(persistence), observeResumeState], }) as AsyncIterable, ) @@ -829,6 +903,7 @@ describe('interrupt persistence', () => { run.calls[0] as { clientToolResults?: ReadonlyMap } ).clientToolResults expect(clientToolResults?.get('tc1')).toBeUndefined() + expect(resumeStates[0]?.has('tc1')).toBe(true) expect((await persistence.stores.interrupts!.get('client-1'))?.status).toBe( 'cancelled', )