fix(ai-persistence): preserve cancelled tool resumes (#1090)

* fix(ai-persistence): preserve cancelled tool resumes

* test(ai-persistence): cover cancelled client-tool hydrate resume

* chore: add changeset for cancelled client-tool resume

---------

Co-authored-by: Alem Tuzlak <t.zlak@hotmail.com>

Author: mikemikimike

.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.

packages/ai-persistence/src/middleware.ts

@@ -396,6 +396,7 @@ function resumeToolStateFromPending(
 ): ChatResumeToolState | undefined {
   const approvals = new Map<string, ToolApprovalResolution>()
   const clientToolResults = new Map<string, unknown>()
+  const cancelledToolCallIds = new Set<string>()
 
   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 }
 }
 
 /**

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<typeof memoryPersistence>,
+  tools: Array<Tool>,
+) {
+  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<StreamChunk>,
+  )
+  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<StreamChunk>,
-    )
+    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<StreamChunk>,
+    )
+
+    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,21 +834,29 @@ describe('interrupt persistence', () => {
     })
 
     const run = mockAdapter([[runStarted(), text('ok'), runFinished('r1')]])
+    const resumeStates: Array<ReadonlySet<string> | undefined> = []
+    const observeResumeState = defineChatMiddleware({
+      name: 'observe-resume-state',
+      onConfig(_ctx, config) {
+        resumeStates.push(config.resumeToolState?.cancelledToolCallIds)
+      },
+    })
     await collect(
       chat({
         adapter: run.adapter,
         messages: [],
         runId: 'r1',
         threadId: 't1',
         resume: [{ interruptId: 'approval-1', status: 'cancelled' }],
-        middleware: [withPersistence(persistence)],
+        middleware: [withPersistence(persistence), observeResumeState],
       }) as AsyncIterable<StreamChunk>,
     )
 
     const approvals = (
       run.calls[0] as { approvals?: ReadonlyMap<string, boolean> }
     ).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<ReadonlySet<string> | 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<StreamChunk>,
     )
 
@@ -829,6 +903,7 @@ describe('interrupt persistence', () => {
       run.calls[0] as { clientToolResults?: ReadonlyMap<string, unknown> }
     ).clientToolResults
     expect(clientToolResults?.get('tc1')).toBeUndefined()
+    expect(resumeStates[0]?.has('tc1')).toBe(true)
     expect((await persistence.stores.interrupts!.get('client-1'))?.status).toBe(
       'cancelled',
     )