From a094648b1d16a6289ab5e616c2196757de7cd4e8 Mon Sep 17 00:00:00 2001 From: kolaworld Date: Sun, 16 Aug 2026 15:02:17 -0400 Subject: [PATCH] fix(ai-client): preserve client-tool continuation ownership Keep native resume batches authoritative while interrupt descriptors remain, then allow legacy client-tool continuation after the native resume settles. Fixes #1106 --- .../fix-sequential-client-tool-resumes.md | 5 + docs/structured-outputs/with-tools.md | 2 +- packages/ai-client/src/chat-client.ts | 33 +-- .../tests/chat-client-resume.test.ts | 240 ++++++++++++++++++ testing/e2e/src/routes/api.tools-test.ts | 2 + .../tests/tools-test/race-conditions.spec.ts | 21 ++ 6 files changed, 281 insertions(+), 22 deletions(-) create mode 100644 .changeset/fix-sequential-client-tool-resumes.md diff --git a/.changeset/fix-sequential-client-tool-resumes.md b/.changeset/fix-sequential-client-tool-resumes.md new file mode 100644 index 0000000000..2b2693f38d --- /dev/null +++ b/.changeset/fix-sequential-client-tool-resumes.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-client': patch +--- + +Keep native interrupt ownership across sequential client-tool resumes. diff --git a/docs/structured-outputs/with-tools.md b/docs/structured-outputs/with-tools.md index 69e43d7be4..ed26096ad3 100644 --- a/docs/structured-outputs/with-tools.md +++ b/docs/structured-outputs/with-tools.md @@ -128,7 +128,7 @@ The full server-tool approval pattern lives in [Tool Approval Flow](../tools/too ## Client tools mid-run -Client tools — defined with `.client((input) => ...)` on the tool definition — execute automatically when the model calls them. The runtime sees the queued `tool-input-available` custom event, looks up the registered `.client()` implementation, runs it, and posts the result back. The agent loop continues to the structured-output stream once every client tool resolves. There's no `onToolCall` option to wire up on the hook side. +Client tools — defined with `.client((input) => ...)` on the tool definition — execute automatically when the model calls them. The server ends the current run with an internal `client-tool-execution` interrupt that does not appear in the public `interrupts` array. The client runs the registered `.client()` implementation and submits its output in a resume batch. Once every client tool resolves, the agent loop continues into the structured-output stream. There's no `onToolCall` option to wire up on the hook side. ```tsx import { toolDefinition } from "@tanstack/ai"; diff --git a/packages/ai-client/src/chat-client.ts b/packages/ai-client/src/chat-client.ts index 5f1dcb06ff..00cc05045c 100644 --- a/packages/ai-client/src/chat-client.ts +++ b/packages/ai-client/src/chat-client.ts @@ -1852,28 +1852,17 @@ export class ChatClient< } } - /** - * True when the client still has user-actionable interrupts (or is mid - * resume submission). Staged/submitting items that are already being - * continued do not block a later turn once the resume stream has cleared - * resume state. - */ + /** True while interrupt descriptors still own continuation. */ + private hasPendingInterrupts(): boolean { + return this.interruptManager.getDescriptors().length > 0 + } + + /** True while an interrupt batch owns the next user turn. */ private hasBlockingInterrupts(): boolean { - if (!this.lastResume && !this.activeInterruptSubmission) { - return false - } - if (this.activeInterruptSubmission) { - return true - } - return this.interruptManager - .getInterrupts() - .some( - (item) => - item.status === 'pending' || - item.status === 'validating' || - item.status === 'error' || - item.status === 'staged', - ) + return ( + this.activeInterruptSubmission !== undefined || + this.hasPendingInterrupts() + ) } /** True while a stream is active, a send is claiming the client, or the queue is draining. */ @@ -2583,6 +2572,8 @@ export class ChatClient< * Check if we should continue the flow and do so if needed */ private async checkForContinuation(): Promise { + if (this.hasPendingInterrupts()) return + // Prevent duplicate continuation attempts if (this.continuationPending || this.isLoading) { this.continuationSkipped = true diff --git a/packages/ai-client/tests/chat-client-resume.test.ts b/packages/ai-client/tests/chat-client-resume.test.ts index 47169442ab..7a4563dc07 100644 --- a/packages/ai-client/tests/chat-client-resume.test.ts +++ b/packages/ai-client/tests/chat-client-resume.test.ts @@ -1201,4 +1201,244 @@ describe('ChatClient resume', () => { ]) expect(client.getInterruptState().interruptErrors).toEqual([]) }) + + it('continues a legacy client tool emitted by a native resume', async () => { + const lookup = toolDefinition({ + name: 'lookup', + description: 'Look up', + inputSchema: z.object({ query: z.string() }), + outputSchema: z.object({ answer: z.number() }), + }).client(async () => ({ answer: 42 })) + const { adapter, contexts } = recordingAdapter([ + // The initial run pauses on a native interrupt. + (ctx) => [ + { + type: EventType.RUN_STARTED, + runId: ctx?.runId ?? 'interrupted-run', + threadId: ctx?.threadId ?? 'thread-1', + timestamp: Date.now(), + }, + { + type: EventType.RUN_FINISHED, + runId: ctx?.runId ?? 'interrupted-run', + threadId: ctx?.threadId ?? 'thread-1', + timestamp: Date.now(), + outcome: { + type: 'interrupt', + interrupts: [ + { + id: 'interrupt-1', + reason: 'approval_required', + metadata: { + kind: 'approval', + toolName: 'confirm', + input: {}, + }, + }, + ], + }, + }, + ], + // The native resume emits a legacy client tool, not another interrupt. + (ctx) => [ + { + type: EventType.RUN_STARTED, + runId: ctx?.runId ?? 'resume-run', + threadId: ctx?.threadId ?? 'thread-1', + timestamp: Date.now(), + }, + { + type: EventType.TOOL_CALL_START, + toolCallId: 'legacy-tool-call', + toolCallName: 'lookup', + toolName: 'lookup', + timestamp: Date.now(), + }, + { + type: EventType.TOOL_CALL_ARGS, + toolCallId: 'legacy-tool-call', + delta: JSON.stringify({ query: 'answer' }), + timestamp: Date.now(), + }, + { + type: EventType.CUSTOM, + name: 'tool-input-available', + value: { + toolCallId: 'legacy-tool-call', + toolName: 'lookup', + input: { query: 'answer' }, + }, + timestamp: Date.now(), + }, + { + type: EventType.RUN_FINISHED, + runId: ctx?.runId ?? 'resume-run', + threadId: ctx?.threadId ?? 'thread-1', + finishReason: 'tool_calls', + timestamp: Date.now(), + }, + ], + // The legacy tool result continues through an ordinary request. + (ctx) => [ + { + type: EventType.RUN_STARTED, + runId: ctx?.runId ?? 'final-run', + threadId: ctx?.threadId ?? 'thread-1', + timestamp: Date.now(), + }, + text('done'), + { + type: EventType.RUN_FINISHED, + runId: ctx?.runId ?? 'final-run', + threadId: ctx?.threadId ?? 'thread-1', + finishReason: 'stop', + timestamp: Date.now(), + }, + ], + ]) + const client = new ChatClient({ + connection: adapter, + threadId: 'thread-1', + tools: [lookup], + }) + + await client.sendMessage('hi') + resolveGenericInterrupt(client) + + await vi.waitFor(() => { + expect(contexts).toHaveLength(3) + expect( + client + .getMessages() + .some((message) => + message.parts.some( + (part) => part.type === 'text' && part.content === 'done', + ), + ), + ).toBe(true) + }) + expect(contexts[1]?.resume).toEqual([ + { + interruptId: 'interrupt-1', + status: 'resolved', + payload: { answer: 'continue' }, + }, + ]) + expect(contexts[2]?.resume).toBeUndefined() + expect(contexts[2]?.parentRunId).toBeUndefined() + }) + + it('keeps native interrupt ownership when a sequential client tool resume fails', async () => { + const outputSchema = z.object({ answer: z.number() }) + const lookup = toolDefinition({ + name: 'lookup', + description: 'Look up', + inputSchema: z.object({ query: z.string() }), + outputSchema, + }).client(async ({ query }) => ({ answer: query === 'first' ? 42 : 43 })) + const outputSchemaHash = hashSchemaInput(outputSchema) + const responseSchema = convertSchemaToJsonSchema(outputSchema) ?? {} + const responseSchemaHash = digestInterruptJson( + canonicalInterruptJson(responseSchema), + ) + const interrupt = + (toolCallId: string, query: string): Script => + (ctx) => { + const runId = ctx?.runId ?? `run-${toolCallId}` + const threadId = ctx?.threadId ?? 'thread-1' + return [ + { + type: EventType.RUN_STARTED, + runId, + threadId, + timestamp: Date.now(), + }, + { + type: EventType.TOOL_CALL_START, + toolCallId, + toolCallName: 'lookup', + toolName: 'lookup', + timestamp: Date.now(), + }, + { + type: EventType.TOOL_CALL_ARGS, + toolCallId, + delta: JSON.stringify({ query }), + timestamp: Date.now(), + }, + { + type: EventType.RUN_FINISHED, + runId, + threadId, + timestamp: Date.now(), + outcome: { + type: 'interrupt', + interrupts: [ + { + id: `client_tool_${toolCallId}`, + reason: 'tanstack:client_tool_execution', + toolCallId, + responseSchema, + metadata: { + kind: 'client_tool', + toolName: 'lookup', + input: { query }, + 'tanstack:interruptBinding': { + kind: 'client-tool-execution', + interruptId: `client_tool_${toolCallId}`, + interruptedRunId: runId, + generation: 0, + toolName: 'lookup', + toolCallId, + outputSchemaHash, + responseSchemaHash, + }, + }, + }, + ], + }, + }, + ] + } + + const { adapter, contexts } = recordingAdapter([ + interrupt('tool-call-1', 'first'), + interrupt('tool-call-2', 'second'), + { + chunks: [], + error: new Error('resume failed'), + }, + ]) + const client = new ChatClient({ + connection: adapter, + threadId: 'thread-1', + tools: [lookup], + }) + + await client.sendMessage('hi') + await vi.waitFor(() => { + expect(contexts).toHaveLength(3) + expect(client.getInterruptState().interruptErrors[0]?.code).toBe( + 'transport', + ) + }) + + expect(contexts[0]?.resume).toBeUndefined() + expect(contexts[1]?.parentRunId).toBe(contexts[0]?.runId) + expect(contexts[1]?.resume).toEqual([ + { + interruptId: 'client_tool_tool-call-1', + status: 'resolved', + payload: { answer: 42 }, + }, + ]) + expect(contexts[2]?.parentRunId).toBe(contexts[1]?.runId) + expect(contexts[2]?.resume).toEqual([ + { + interruptId: 'client_tool_tool-call-2', + status: 'resolved', + payload: { answer: 43 }, + }, + ]) + }) }) diff --git a/testing/e2e/src/routes/api.tools-test.ts b/testing/e2e/src/routes/api.tools-test.ts index d350281f78..1bb6f6e16b 100644 --- a/testing/e2e/src/routes/api.tools-test.ts +++ b/testing/e2e/src/routes/api.tools-test.ts @@ -226,6 +226,8 @@ export const Route = createFileRoute('/api/tools-test')({ context: runtimeContext, threadId: params.threadId, runId: params.runId, + ...(params.parentRunId ? { parentRunId: params.parentRunId } : {}), + ...(params.resume ? { resume: params.resume } : {}), agentLoopStrategy: maxIterations(20), abortController, }) diff --git a/testing/e2e/tests/tools-test/race-conditions.spec.ts b/testing/e2e/tests/tools-test/race-conditions.spec.ts index 6c341caa0c..dd1c1db60a 100644 --- a/testing/e2e/tests/tools-test/race-conditions.spec.ts +++ b/testing/e2e/tests/tools-test/race-conditions.spec.ts @@ -33,6 +33,17 @@ test.describe('Race Condition Tests', () => { testId, aimockPort, }) => { + const requestBodies: Array = [] + page.on('request', (request) => { + if ( + request.url().includes('/api/tools-test') && + request.method() === 'POST' + ) { + const body = request.postDataJSON() + if (body) requestBodies.push(body) + } + }) + await selectScenario(page, 'sequential-client-tools', testId, aimockPort) const startTime = Date.now() @@ -74,6 +85,16 @@ test.describe('Race Condition Tests', () => { expect(executionEvents[2]?.type).toBe('execution-start') expect(executionEvents[3]?.type).toBe('execution-complete') + await expect(page.locator('#messages-json-content')).toContainText( + 'Both notifications have been shown.', + ) + expect(requestBodies).toHaveLength(3) + expect(requestBodies[0]?.resume).toBeUndefined() + expect(requestBodies[1]?.parentRunId).toBe(requestBodies[0]?.runId) + expect(requestBodies[1]?.resume).toHaveLength(1) + expect(requestBodies[2]?.parentRunId).toBe(requestBodies[1]?.runId) + expect(requestBodies[2]?.resume).toHaveLength(1) + // If it takes too long (e.g., > 10 seconds), it might indicate blocking // (each tool takes ~50ms, so total should be well under 5 seconds) expect(duration).toBeLessThan(10000)