From e957b3b065ec1bf5021430946f3efb1e5740350b Mon Sep 17 00:00:00 2001 From: Jason Kim Date: Sun, 26 Jul 2026 23:44:42 +0900 Subject: [PATCH 1/3] fix(ai): handle client tool waits as invocation boundaries MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Client-tool and approval waits end the current server invocation without ending the overall generation. Treating this boundary inconsistently left middleware lifecycle hooks—and therefore telemetry spans and duration metrics—unfinished. Treat the wait as terminal for middleware exactly once, while keeping structured-output finalization suspended until the client result resumes generation in a subsequent invocation. Add unit and E2E regression coverage for both sides of this lifecycle boundary. --- .changeset/fuzzy-tools-wait.md | 5 + packages/ai/src/activities/chat/index.ts | 83 ++--- .../ai/tests/chat-client-tool-wait.test.ts | 288 ++++++++++++++++++ .../structured-client-tool-wait.json | 27 ++ testing/e2e/src/lib/middleware-test-tools.ts | 9 + testing/e2e/src/routes/api.middleware-test.ts | 16 +- testing/e2e/src/routes/middleware-test.tsx | 57 +++- testing/e2e/tests/client-tool-wait.spec.ts | 102 +++++++ 8 files changed, 538 insertions(+), 49 deletions(-) create mode 100644 .changeset/fuzzy-tools-wait.md create mode 100644 packages/ai/tests/chat-client-tool-wait.test.ts create mode 100644 testing/e2e/fixtures/middleware-test/structured-client-tool-wait.json create mode 100644 testing/e2e/src/lib/middleware-test-tools.ts create mode 100644 testing/e2e/tests/client-tool-wait.spec.ts diff --git a/.changeset/fuzzy-tools-wait.md b/.changeset/fuzzy-tools-wait.md new file mode 100644 index 0000000000..0fe0ec3498 --- /dev/null +++ b/.changeset/fuzzy-tools-wait.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai': patch +--- + +Finish middleware lifecycle hooks while waiting for client tools or approvals, and defer structured-output finalization until their results arrive in a later chat invocation. diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index 91f84d3490..236d1037a8 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -1121,6 +1121,43 @@ class TextEngine< return this.finalizationError } + private async runTerminalHook(): Promise { + if (this.terminalHookCalled || this.isCancelled()) return + + this.terminalHookCalled = true + + if (this.finalizationError) { + const errForHook = new Error( + this.finalizationError.message, + this.finalizationError.cause !== undefined + ? { cause: this.finalizationError.cause } + : undefined, + ) + if (this.finalizationError.code !== undefined) { + Object.defineProperty(errForHook, 'code', { + value: this.finalizationError.code, + enumerable: true, + }) + } + await this.middlewareRunner.runOnError(this.middlewareCtx, { + error: errForHook, + duration: Date.now() - this.streamStartTime, + }) + return + } + + this.addTerminalAssistantMessages() + await this.middlewareRunner.runOnFinish(this.middlewareCtx, { + finishReason: this.lastFinishReason, + duration: Date.now() - this.streamStartTime, + content: this.accumulatedContent, + usage: rebuildTokenUsage( + this.finishedEvent?.usage, + tanstackMetadata(this.finishedEvent ?? undefined)?.usage, + ), + }) + } + async *run(): AsyncGenerator { this.beforeRun() this.logger.agentLoop('run started', { @@ -1164,6 +1201,7 @@ class TextEngine< const pendingPhase = yield* this.checkForPendingToolCalls() if (pendingPhase === 'wait') { + await this.runTerminalHook() return } @@ -1271,46 +1309,11 @@ class TextEngine< } } - // Call terminal hook (skip when waiting for client — stream is paused, not finished). - // Priority: finalizationError → onError; otherwise normal onFinish. - // Skip on cancellation — the finally block routes aborts to onAbort. - if ( - !this.terminalHookCalled && - this.toolPhase !== 'wait' && - !this.isCancelled() - ) { - if (this.finalizationError) { - this.terminalHookCalled = true - const errForHook = new Error( - this.finalizationError.message, - this.finalizationError.cause !== undefined - ? { cause: this.finalizationError.cause } - : undefined, - ) - if (this.finalizationError.code !== undefined) { - Object.defineProperty(errForHook, 'code', { - value: this.finalizationError.code, - enumerable: true, - }) - } - await this.middlewareRunner.runOnError(this.middlewareCtx, { - error: errForHook, - duration: Date.now() - this.streamStartTime, - }) - } else { - this.addTerminalAssistantMessages() - this.terminalHookCalled = true - await this.middlewareRunner.runOnFinish(this.middlewareCtx, { - finishReason: this.lastFinishReason, - duration: Date.now() - this.streamStartTime, - content: this.accumulatedContent, - usage: rebuildTokenUsage( - this.finishedEvent?.usage, - tanstackMetadata(this.finishedEvent ?? undefined)?.usage, - ), - }) - } - } + // A client-tool or approval wait ends this server-side chat invocation. + // Structured-output finalization stays paused until the caller submits + // the result in a new invocation, but middleware must still observe one + // terminal hook for the current invocation. + await this.runTerminalHook() } catch (error: unknown) { if ( error instanceof Error && diff --git a/packages/ai/tests/chat-client-tool-wait.test.ts b/packages/ai/tests/chat-client-tool-wait.test.ts new file mode 100644 index 0000000000..582bd13de7 --- /dev/null +++ b/packages/ai/tests/chat-client-tool-wait.test.ts @@ -0,0 +1,288 @@ +import { describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { chat } from '../src/activities/chat/index' +import { otelMiddleware } from '../src/middlewares/otel' +import { EventType } from '../src/types' +import { clientTool, collectChunks, createMockAdapter, ev } from './test-utils' +import { createFakeMeter, createFakeTracer } from './middlewares/fake-otel' +import type { ChatMiddleware } from '../src/activities/chat/middleware/types' +import type { StreamChunk } from '../src/types' + +const ResultSchema = z.object({ value: z.string() }) +const usage = { promptTokens: 5, completionTokens: 3, totalTokens: 8 } + +function clientToolTurn(): Array { + return [ + ev.runStarted(), + ev.toolStart('call-1', 'ask_client'), + ev.toolArgs('call-1', '{}'), + ev.runFinished('tool_calls', 'run-1', usage), + ] +} + +function structuredTextTurn(value: string): Array { + return [ + ev.runStarted(), + ev.textStart(), + ev.textContent(JSON.stringify({ value })), + ev.textEnd(), + ev.runFinished('stop'), + ] +} + +function createTerminalSpy() { + const onFinish = vi.fn() + const onAbort = vi.fn() + const onError = vi.fn() + const middleware: ChatMiddleware = { + name: 'terminal-spy', + onFinish, + onAbort, + onError, + } + + return { middleware, onFinish, onAbort, onError } +} + +function expectOnlyFinish({ + onFinish, + onAbort, + onError, +}: ReturnType) { + expect(onFinish).toHaveBeenCalledOnce() + expect(onAbort).not.toHaveBeenCalled() + expect(onError).not.toHaveBeenCalled() +} + +function expectNoStructuredFinalization(chunks: Array) { + const structuredChunks = chunks.filter( + (chunk) => + chunk.type === EventType.CUSTOM && + (chunk.name === 'structured-output.start' || + chunk.name === 'structured-output.complete'), + ) + expect(structuredChunks).toHaveLength(0) + expect(chunks.some((chunk) => chunk.type === EventType.RUN_ERROR)).toBe(false) +} + +function expectClientToolWait(chunks: Array) { + expect( + chunks.some( + (chunk) => + chunk.type === EventType.RUN_FINISHED && + chunk.outcome?.type === 'interrupt' && + chunk.outcome.interrupts.some( + (interrupt) => + interrupt.reason === 'tanstack:client_tool_execution' && + interrupt.toolCallId === 'call-1', + ), + ), + ).toBe(true) +} + +describe('client-tool wait lifecycle', () => { + it('calls onFinish exactly once when a live client tool waits', async () => { + const { adapter } = createMockAdapter({ iterations: [clientToolTurn()] }) + const terminal = createTerminalSpy() + + const chunks = await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Ask the client' }], + tools: [clientTool('ask_client')], + middleware: [terminal.middleware], + }), + ) + + expectClientToolWait(chunks) + expectOnlyFinish(terminal) + expect(terminal.onFinish.mock.calls[0]?.[1]).toMatchObject({ + finishReason: 'tool_calls', + content: '', + usage, + }) + }) + + it('calls onFinish exactly once for a pending client tool early return', async () => { + const { adapter, calls } = createMockAdapter({ iterations: [] }) + const terminal = createTerminalSpy() + + const chunks = await collectChunks( + chat({ + adapter, + messages: [ + { role: 'user', content: 'Ask the client' }, + { + role: 'assistant', + content: '', + toolCalls: [ + { + id: 'call-1', + type: 'function' as const, + function: { name: 'ask_client', arguments: '{}' }, + }, + ], + }, + ], + tools: [clientTool('ask_client')], + middleware: [terminal.middleware], + }), + ) + + expect(calls).toHaveLength(0) + expectClientToolWait(chunks) + expectOnlyFinish(terminal) + expect(terminal.onFinish.mock.calls[0]?.[1]).toMatchObject({ + finishReason: null, + content: '', + usage: undefined, + }) + }) + + it('ends OpenTelemetry spans and records duration while waiting', async () => { + const { adapter } = createMockAdapter({ iterations: [clientToolTurn()] }) + const fakeTracer = createFakeTracer() + const fakeMeter = createFakeMeter() + + await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Ask the client' }], + tools: [clientTool('ask_client')], + middleware: [ + otelMiddleware({ + tracer: fakeTracer.tracer, + meter: fakeMeter.meter, + }), + ], + }), + ) + + expect(fakeTracer.spans).toHaveLength(2) + expect(fakeTracer.spans.every((span) => span.ended)).toBe(true) + expect( + fakeMeter.records.filter( + (record) => record.name === 'gen_ai.client.operation.duration', + ), + ).toHaveLength(1) + expect( + fakeMeter.records.filter( + (record) => record.name === 'gen_ai.client.token.usage', + ), + ).toHaveLength(2) + }) +}) + +describe('client-tool wait with structured output', () => { + it('does not harvest native-combined output before the client result', async () => { + const structuredOutput = vi.fn(async () => ({ + data: { value: 'premature' }, + rawText: '{"value":"premature"}', + })) + const { adapter } = createMockAdapter({ + iterations: [clientToolTurn()], + structuredOutput, + supportsCombinedToolsAndSchema: true, + }) + const terminal = createTerminalSpy() + + const chunks = await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Ask the client' }], + tools: [clientTool('ask_client')], + outputSchema: ResultSchema, + stream: true, + middleware: [terminal.middleware], + }), + ) + + expect(structuredOutput).not.toHaveBeenCalled() + expectNoStructuredFinalization(chunks) + expectOnlyFinish(terminal) + }) + + it('does not call fallback finalization before the client result', async () => { + const structuredOutput = vi.fn(async () => ({ + data: { value: 'premature' }, + rawText: '{"value":"premature"}', + })) + const { adapter } = createMockAdapter({ + iterations: [clientToolTurn()], + structuredOutput, + }) + const terminal = createTerminalSpy() + + const chunks = await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'Ask the client' }], + tools: [clientTool('ask_client')], + outputSchema: ResultSchema, + stream: true, + middleware: [terminal.middleware], + }), + ) + + expect(structuredOutput).not.toHaveBeenCalled() + expectNoStructuredFinalization(chunks) + expectOnlyFinish(terminal) + }) + + it('completes structured output in the next invocation with the client result', async () => { + const { adapter, calls } = createMockAdapter({ + iterations: [structuredTextTurn('client-result')], + supportsCombinedToolsAndSchema: true, + }) + + const chunks = await collectChunks( + chat({ + adapter, + messages: [ + { role: 'user', content: 'Ask the client' }, + { + role: 'assistant', + content: '', + toolCalls: [ + { + id: 'call-1', + type: 'function' as const, + function: { name: 'ask_client', arguments: '{}' }, + }, + ], + }, + { + role: 'tool', + content: '{"value":"client-result"}', + toolCallId: 'call-1', + }, + ], + tools: [clientTool('ask_client')], + outputSchema: ResultSchema, + stream: true, + }), + ) + + expect(calls).toHaveLength(1) + expect(calls[0]?.messages.some((message) => message.role === 'tool')).toBe( + true, + ) + + const complete = chunks.find( + (chunk) => + chunk.type === EventType.CUSTOM && + chunk.name === 'structured-output.complete', + ) + expect(complete).toBeDefined() + if ( + complete?.type !== EventType.CUSTOM || + complete.name !== 'structured-output.complete' + ) { + throw new Error('Expected structured-output.complete') + } + expect(complete.value.object).toEqual({ value: 'client-result' }) + expect(chunks.some((chunk) => chunk.type === EventType.RUN_ERROR)).toBe( + false, + ) + }) +}) diff --git a/testing/e2e/fixtures/middleware-test/structured-client-tool-wait.json b/testing/e2e/fixtures/middleware-test/structured-client-tool-wait.json new file mode 100644 index 0000000000..68fd9cd324 --- /dev/null +++ b/testing/e2e/fixtures/middleware-test/structured-client-tool-wait.json @@ -0,0 +1,27 @@ +{ + "fixtures": [ + { + "match": { + "userMessage": "[structured-client-tool-wait] run test", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "get_client_context", + "arguments": "{}" + } + ] + } + }, + { + "match": { + "userMessage": "[structured-client-tool-wait] run test", + "sequenceIndex": 1 + }, + "response": { + "content": "{\"name\":\"Client Context Guitar\",\"price\":999,\"reason\":\"Recommended using client-result\",\"rating\":5}" + } + } + ] +} diff --git a/testing/e2e/src/lib/middleware-test-tools.ts b/testing/e2e/src/lib/middleware-test-tools.ts new file mode 100644 index 0000000000..2e2a110e8e --- /dev/null +++ b/testing/e2e/src/lib/middleware-test-tools.ts @@ -0,0 +1,9 @@ +import { toolDefinition } from '@tanstack/ai' +import { z } from 'zod' + +export const clientContextToolDefinition = toolDefinition({ + name: 'get_client_context', + description: 'Get context that is only available in the browser', + inputSchema: z.object({}), + outputSchema: z.object({ context: z.string() }), +}) diff --git a/testing/e2e/src/routes/api.middleware-test.ts b/testing/e2e/src/routes/api.middleware-test.ts index 41f14857f2..6f293332f2 100644 --- a/testing/e2e/src/routes/api.middleware-test.ts +++ b/testing/e2e/src/routes/api.middleware-test.ts @@ -11,6 +11,7 @@ import { import { otelMiddleware } from '@tanstack/ai/middlewares/otel' import { memoryMiddleware } from '@tanstack/ai-memory' import type { MemoryAdapter } from '@tanstack/ai-memory' +import { clientContextToolDefinition } from '@/lib/middleware-test-tools' import { getMemoryCapture, recordMemoryConfig, @@ -645,20 +646,23 @@ export const Route = createFileRoute('/api/middleware-test')({ ? genericTools(testId, genericScenario) : scenario === 'with-tool' ? [weatherTool] - : [] + : scenario === 'structured-client-tool-wait' + ? [clientContextToolDefinition] + : [] - // The two `structured-output*` scenarios both bind the same - // guitar schema; they differ only in what the spec asserts (phases - // observed vs RUN_STARTED/RUN_FINISHED uniqueness). A single - // outputSchema branch keeps the route narrow. + // Structured scenarios bind the same guitar schema. The client-tool + // variant also passes an isomorphic tool definition so the first + // server invocation can stop at the browser-execution boundary. const isStructured = scenario === 'structured-output' || - scenario === 'structured-output-stream' + scenario === 'structured-output-stream' || + scenario === 'structured-client-tool-wait' const rawStream = isStructured ? chat({ ...adapterOptions, messages: params.messages, + tools, middleware, threadId: params.threadId, runId: params.runId, diff --git a/testing/e2e/src/routes/middleware-test.tsx b/testing/e2e/src/routes/middleware-test.tsx index c55a1cfada..398b3ba775 100644 --- a/testing/e2e/src/routes/middleware-test.tsx +++ b/testing/e2e/src/routes/middleware-test.tsx @@ -1,4 +1,4 @@ -import { useMemo, useState } from 'react' +import { useMemo, useRef, useState } from 'react' import { createFileRoute } from '@tanstack/react-router' import { useChat, @@ -12,6 +12,7 @@ import { renderReviewTool, reviewPlan, } from '@/lib/generic-middleware-interrupts' +import { clientContextToolDefinition } from '@/lib/middleware-test-tools' const MIDDLEWARE_MODES = [ { id: 'none', label: 'No Middleware' }, @@ -111,6 +112,19 @@ function MiddlewareTestPage() { saveCount: number }>({ configs: [], saveCount: 0 }) const [clientToolExecutions, setClientToolExecutions] = useState(0) + const [clientToolWaiting, setClientToolWaiting] = useState(false) + const clientToolResolver = useRef< + ((value: { context: string }) => void) | null + >(null) + const [clientContextTool] = useState(() => + clientContextToolDefinition.client( + () => + new Promise<{ context: string }>((resolve) => { + clientToolResolver.current = resolve + setClientToolWaiting(true) + }), + ), + ) const clientToolList = useMemo( () => @@ -120,11 +134,12 @@ function MiddlewareTestPage() { setClientToolExecutions((count) => count + 1) return { rendered: true, reviewId } }), + clientContextTool, ), - [], + [clientContextTool], ) - const { messages, sendMessage, isLoading, interrupts } = useChat< + const { messages, sendMessage, isLoading, error, interrupts } = useChat< typeof clientToolList, undefined, unknown, @@ -181,9 +196,19 @@ function MiddlewareTestPage() { setTestComplete(false) setPhaseCapture(EMPTY_PHASE_CAPTURE) setClientToolExecutions(0) + setClientToolWaiting(false) + clientToolResolver.current = null sendMessage(`[${scenario}] run test`) } + const handleResolveClientTool = () => { + const resolve = clientToolResolver.current + if (!resolve) return + clientToolResolver.current = null + setClientToolWaiting(false) + resolve({ context: 'client-result' }) + } + type ActiveInterrupt = (typeof interrupts)[number] const reviewInterrupts = interrupts.filter( ( @@ -232,6 +257,9 @@ function MiddlewareTestPage() { +