diff --git a/.changeset/ai-after-tool-hook-error.md b/.changeset/ai-after-tool-hook-error.md new file mode 100644 index 0000000000..c10c60ce99 --- /dev/null +++ b/.changeset/ai-after-tool-hook-error.md @@ -0,0 +1,7 @@ +--- +'@tanstack/ai': patch +--- + +Treat a failed `onAfterToolCall` middleware hook as a chat run error instead of recording a second, contradictory tool result. + +Deliver and persist that error for durable streams even after an earlier tool iteration emitted `RUN_FINISHED`. diff --git a/docs/advanced/middleware-advanced.md b/docs/advanced/middleware-advanced.md index c69c4038fe..00d27771d2 100644 --- a/docs/advanced/middleware-advanced.md +++ b/docs/advanced/middleware-advanced.md @@ -637,7 +637,13 @@ A guard cannot change the arguments. To change them, use `onBeforeToolCall`. ### onAfterToolCall -`onAfterToolCall` runs after each tool runs (or is skipped). Every middleware runs, in array order. +`onAfterToolCall` runs after each tool runs (or is skipped). Hooks run in order. +If one hook throws, later hooks do not run. +The chat run ends through `onError`. Catch errors inside the hook when its work +is optional. + +Durable response helpers also deliver and store the `RUN_ERROR`, even after an +earlier tool iteration finished. ```typescript import { type ChatMiddleware } from "@tanstack/ai"; diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index 89c7c4da44..b9bf6de17c 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -123,6 +123,7 @@ import { MiddlewareRunner } from './middleware/compose' import { getRunDetached } from './middleware/run-store' import { LoadChildCapability } from './middleware/load-child' import { publishRunDetachedSignal } from '../../delivery-detach' +import { publishRunCompletionSignal } from '../../delivery-completion' import { publishRunDisconnectHandler } from '../../delivery-disconnect' import { provideSandboxRuntime } from './middleware/sandbox-runtime' import { provideRunDisconnect } from './middleware/run-disconnect' @@ -1163,6 +1164,7 @@ class TextEngine< // observe both cancellation sources via ctx.abortSignal. private readonly toolAbortSignal?: AbortSignal private terminalHookCalled = false + private runEnded = false /** * Latched the first time the delivery socket closes; see `notifyDisconnected`. * Also read by `subscribe` so a listener registered AFTER the disconnect (a @@ -1492,6 +1494,7 @@ class TextEngine< } private async runTerminalHook(): Promise { + this.runEnded = true if (this.terminalHookCalled || this.isCancelled()) return this.terminalHookCalled = true @@ -1573,6 +1576,7 @@ class TextEngine< if (this.earlyTermination) { yield* this.emitSuccessfulEarlyTermination() + this.runEnded = true if (!this.terminalHookCalled) { this.terminalHookCalled = true await this.middlewareRunner.runOnFinish(this.middlewareCtx, { @@ -1734,6 +1738,7 @@ class TextEngine< 'continuationRunId' in error && typeof error.continuationRunId === 'string' ) { + this.runEnded = true this.terminalHookCalled = true yield* this.pipeThroughMiddleware({ type: EventType.RUN_FINISHED, @@ -1823,6 +1828,10 @@ class TextEngine< this.middlewareCtx.toolNames = this.eventToolNames } + hasRunEnded(): boolean { + return this.runEnded + } + private async beginCycle(): Promise { if (this.cyclePhase === 'processText') { await this.beginIteration() @@ -3980,6 +3989,7 @@ class TextEngine< }) } for (const output of terminalOutputs) { + this.runEnded = true yield this.publicInterruptTerminal(output) } return true @@ -6318,19 +6328,20 @@ function validateInterruptDefinitions( /** * The slice of the engine that the durable delivery sink reaches back into, in - * BOTH directions: it reads the detach verdict (`wasDetached`) and pushes the - * socket-closed fact in (`notifyDisconnected`). Filled by the generator body as - * soon as its engine exists. + * BOTH directions: it reads run completion and detach state, and pushes the + * socket-closed fact in (`notifyDisconnected`). Filled by the generator body + * as soon as its engine exists. */ interface DeliveryEngineRef { current?: { wasDetached: () => boolean notifyDisconnected: () => void + hasRunEnded: () => boolean } } /** - * Publish both delivery-side seams for `stream`. + * Publish delivery-side state and callbacks for `stream`. * * Shared by the two streaming paths so they cannot drift apart — the * structured-output path having been wired for one seam and not the other is @@ -6341,6 +6352,7 @@ function publishDeliverySeams( stream: object, engineRef: DeliveryEngineRef, ): void { + publishRunCompletionSignal(stream, () => engineRef.current?.hasRunEnded()) // A thunk, evaluated on the sink's teardown path: the engine does not exist // yet, and the verdict it will report is only written during `onAbort`. publishRunDetachedSignal( diff --git a/packages/ai/src/activities/chat/tools/tool-calls.ts b/packages/ai/src/activities/chat/tools/tool-calls.ts index fa984b6d52..992fbfec70 100644 --- a/packages/ai/src/activities/chat/tools/tool-calls.ts +++ b/packages/ai/src/activities/chat/tools/tool-calls.ts @@ -786,6 +786,7 @@ export async function* executeServerTool( subagentInterrupts?: Array, ): AsyncGenerator { const startTime = Date.now() + let resultRecorded = false try { if (!tool.execute) { throw new Error(`Tool ${toolName} has no execute() implementation`) @@ -821,6 +822,7 @@ export async function* executeServerTool( ? outcome.text : compactForModel(outcome.result), } + resultRecorded = true await pushResultAndRunAfterHook( results, { @@ -869,6 +871,7 @@ export async function* executeServerTool( const finalResult = typeof result === 'string' ? safeJsonParse(result) : (result ?? null) + resultRecorded = true await pushResultAndRunAfterHook( results, { @@ -891,6 +894,10 @@ export async function* executeServerTool( middlewareHooks, ) } catch (error: unknown) { + // A post-execution hook failed after the tool result was recorded. Let the + // run handle that failure; reporting it as a second tool result would give + // the model contradictory outcomes and call the hook a second time. + if (resultRecorded) throw error const duration = Date.now() - startTime // Flush remaining events diff --git a/packages/ai/src/delivery-completion.ts b/packages/ai/src/delivery-completion.ts new file mode 100644 index 0000000000..4db74b491a --- /dev/null +++ b/packages/ai/src/delivery-completion.ts @@ -0,0 +1,15 @@ +/** Chat streams can emit RUN_FINISHED before the agent loop actually ends. */ +const completionSignals = new WeakMap boolean | undefined>() + +/** Publish final-run state before terminal hooks and cleanup can fail. */ +export function publishRunCompletionSignal( + stream: object, + completed: () => boolean | undefined, +): void { + completionSignals.set(stream, completed) +} + +/** Undefined preserves terminal-event handling for user-provided iterables. */ +export function wasRunCompleted(stream: object): boolean | undefined { + return completionSignals.get(stream)?.() +} diff --git a/packages/ai/src/stream-to-response.ts b/packages/ai/src/stream-to-response.ts index 165dd1f94f..e7f0f0b147 100644 --- a/packages/ai/src/stream-to-response.ts +++ b/packages/ai/src/stream-to-response.ts @@ -5,6 +5,7 @@ import { isTerminalRunStatus, } from './activities/chat/middleware/run-store' import { wasRunDetached } from './delivery-detach' +import { wasRunCompleted } from './delivery-completion' import { notifyRunDisconnected } from './delivery-disconnect' import { resolveResumeRunId } from './stream-durability' import { EventType } from './types' @@ -715,14 +716,12 @@ export function durableStreamSource( async function* produce(): AsyncIterable { let batch: Array = [] - let terminalPersisted = false - // Whether a terminal event was actually delivered LIVE to the consumer (as - // opposed to only appended to the log). Distinguishes "the run already ended - // on the wire" from "the log has a terminal but the consumer never saw one", - // which governs whether a late durability-cleanup failure may be rethrown. - // Only ever assigned inside the nested flush() closure, which TS's - // control-flow analysis can't observe (see the disable at the read site). - let terminalForwarded = false + // A persisted terminal need not have reached the live consumer. Chat + // iteration terminals are provisional until the engine ends the run. + const terminals: { + persisted?: 'RUN_FINISHED' | 'RUN_ERROR' + forwarded?: 'RUN_FINISHED' | 'RUN_ERROR' + } = {} let failure: RecordedFailure | undefined let terminalCause: unknown let hasTerminalCause = false @@ -756,17 +755,14 @@ export function durableStreamSource( validateOffset(offset) idByChunk.set(chunk, offset) }) - if ( - toForward.some( - (chunk) => - chunk.type === 'RUN_FINISHED' || chunk.type === 'RUN_ERROR', - ) - ) { - terminalPersisted = true + for (const chunk of toForward) { + if (chunk.type === 'RUN_FINISHED' || chunk.type === 'RUN_ERROR') { + terminals.persisted = chunk.type + } } for (const chunk of toForward) { if (chunk.type === 'RUN_FINISHED' || chunk.type === 'RUN_ERROR') { - terminalForwarded = true + terminals.forwarded = chunk.type } yield chunk } @@ -908,7 +904,7 @@ export function durableStreamSource( // What is ALREADY IN THE LOG is deliberately NOT a conjunct. An agent-loop // run emits one `RUN_FINISHED` PER ITERATION — the intermediate // `finishReason: 'tool_calls'` terminal is flushed at its boundary - // mid-run — so `terminalPersisted` means "some terminal is in the log", + // mid-run — so `terminals.persisted` means "some terminal is in the log", // never "the run ended". Gating on it terminalized the log of a healthy, // still-running agent for every tool-calling run. Nor can the sink // distinguish a final terminal from an intermediate one by its @@ -922,6 +918,11 @@ export function durableStreamSource( !hasTerminalCause && wasRunDetached(stream) + const completed = wasRunCompleted(stream) !== false + const terminalPersisted = + terminals.persisted === 'RUN_ERROR' || + (terminals.persisted === 'RUN_FINISHED' && completed) + if ( !detached && needsTerminalPersistence(terminalPersisted, cancelled, hasTerminalCause) @@ -933,7 +934,6 @@ export function durableStreamSource( const cause = hasTerminalCause ? terminalCause : { name: 'AbortError' } try { await durability.append([runErrorChunk(cause)]) - terminalPersisted = true } catch (terminalError) { // Rethrown to the live consumer below, but a joiner replaying the log // only ever sees a generic incomplete error — so record the real @@ -982,7 +982,9 @@ export function durableStreamSource( // provider that throws AFTER emitting its own terminal, whose error is // otherwise neither delivered nor logged.) if (failure !== undefined) { - // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition -- terminalForwarded is set only inside the flush() closure, which TS CFA narrows away here + const terminalForwarded = + terminals.forwarded === 'RUN_ERROR' || + (terminals.forwarded === 'RUN_FINISHED' && completed) if (!terminalForwarded) { // eslint-disable-next-line no-unsafe-finally throw failure.error diff --git a/packages/ai/tests/middleware.test.ts b/packages/ai/tests/middleware.test.ts index 532fcc7419..93b88f9d3a 100644 --- a/packages/ai/tests/middleware.test.ts +++ b/packages/ai/tests/middleware.test.ts @@ -1384,6 +1384,49 @@ describe('chat() middleware', () => { // onAfterToolCall error path // ========================================================================== describe('onAfterToolCall error handling', () => { + it('does not turn a failed after-hook into a second tool result', async () => { + const execute = vi.fn(() => ({ sent: true })) + const onAfterToolCall = vi.fn() + const onError = vi.fn() + const { adapter } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.toolStart('tc-1', 'sendEmail'), + ev.toolArgs('tc-1', '{}'), + ev.toolEnd('tc-1', 'sendEmail', { input: {} }), + ev.runFinished('tool_calls'), + ], + ], + }) + + const stream = chat({ + adapter, + threadId: 'thread-1', + messages: [{ role: 'user', content: 'Send it' }], + tools: [serverTool('sendEmail', execute)], + middleware: [ + { + name: 'failing-after-hook', + onAfterToolCall: (...args) => { + onAfterToolCall(...args) + if (onAfterToolCall.mock.calls.length === 1) { + throw new Error('post-hook failed') + } + }, + onError: (_ctx, info) => onError(info), + }, + ], + }) + + await expect( + collectChunks(stream as AsyncIterable), + ).rejects.toThrow('post-hook failed') + expect(execute).toHaveBeenCalledOnce() + expect(onAfterToolCall).toHaveBeenCalledOnce() + expect(onError).toHaveBeenCalledOnce() + }) + it('should report tool execution errors in onAfterToolCall', async () => { const afterCalls: Array<{ ok: boolean; error?: unknown }> = [] diff --git a/packages/ai/tests/stream-to-response-durability.test.ts b/packages/ai/tests/stream-to-response-durability.test.ts index 18979584ba..56c93b8b68 100644 --- a/packages/ai/tests/stream-to-response-durability.test.ts +++ b/packages/ai/tests/stream-to-response-durability.test.ts @@ -101,6 +101,72 @@ function fixedOffsetDurability( } describe('toServerSentEventsResponse with durability', () => { + it.each(['SSE', 'NDJSON'] as const)( + 'delivers and persists a later after-tool error over %s', + async (format) => { + const { adapter } = createMockAdapter({ + iterations: ['first', 'second'].map((id) => [ + ev.runStarted(), + ev.toolStart(id, 'lookup'), + ev.toolArgs(id, '{}'), + ev.toolEnd(id), + ev.runFinished('tool_calls'), + ]), + }) + const execute = vi.fn(() => 'found') + const onError = vi.fn() + const onAfterToolCall = vi.fn(() => { + if (onAfterToolCall.mock.calls.length === 2) { + throw new Error('second hook failed') + } + }) + const stream = chat({ + adapter, + threadId: 'thread-1', + messages: [{ role: 'user' as const, content: 'Look twice' }], + tools: [ + toolDefinition({ + name: 'lookup', + description: 'Return a value', + }).server(execute), + ], + middleware: [{ name: 'fail-second-hook', onAfterToolCall, onError }], + }) + const durability = memoryStream( + new Request(`https://example.test/chat?runId=after-hook-${format}`, { + method: 'POST', + }), + ) + const response = ( + format === 'SSE' ? toServerSentEventsResponse : toHttpResponse + )(stream, { durability: { adapter: durability } }) + const body = await readBody(response) + const events = + format === 'SSE' ? parseSseEvents(body) : parseNdjsonEvents(body) + const types = events.map((event) => field(event, 'type')) + expect(types).toContain(EventType.RUN_FINISHED) + expect(types.at(-1)).toBe(EventType.RUN_ERROR) + expect(types.filter((type) => type === EventType.RUN_ERROR)).toHaveLength( + 1, + ) + expect(field(events.at(-1)!, 'message')).toBe('second hook failed') + expect( + events.filter( + (event) => field(event, 'type') === EventType.TOOL_CALL_RESULT, + ), + ).toHaveLength(1) + expect(execute).toHaveBeenCalledTimes(2) + expect(onAfterToolCall).toHaveBeenCalledTimes(2) + expect(onError).toHaveBeenCalledOnce() + + const stored = (await durability.snapshot()).map((entry) => entry.chunk) + expect(stored.at(-1)?.type).toBe(EventType.RUN_ERROR) + expect( + stored.filter((chunk) => chunk.type === EventType.RUN_ERROR), + ).toHaveLength(1) + }, + ) + it('keeps persisting while a connected client does not read', async () => { const durability = memoryStream( new Request('https://example.test/api/chat?runId=unread-connected', { diff --git a/testing/e2e/fixtures/middleware-test/after-tool-error.json b/testing/e2e/fixtures/middleware-test/after-tool-error.json new file mode 100644 index 0000000000..69895e1656 --- /dev/null +++ b/testing/e2e/fixtures/middleware-test/after-tool-error.json @@ -0,0 +1,24 @@ +{ + "fixtures": [ + { + "match": { "userMessage": "[with-tool] run twice", "sequenceIndex": 0 }, + "response": { + "toolCalls": [ + { "name": "get_weather", "arguments": "{\"city\":\"NYC\"}" } + ] + } + }, + { + "match": { "userMessage": "[with-tool] run twice", "sequenceIndex": 1 }, + "response": { + "toolCalls": [ + { "name": "get_weather", "arguments": "{\"city\":\"London\"}" } + ] + } + }, + { + "match": { "userMessage": "[with-tool] run twice", "sequenceIndex": 2 }, + "response": { "content": "Both cities are sunny." } + } + ] +} diff --git a/testing/e2e/src/routes/api.middleware-test.ts b/testing/e2e/src/routes/api.middleware-test.ts index cccdabc1d5..25f1d3dfa7 100644 --- a/testing/e2e/src/routes/api.middleware-test.ts +++ b/testing/e2e/src/routes/api.middleware-test.ts @@ -5,6 +5,7 @@ import { createCapability, EventType, maxIterations, + memoryStream, toServerSentEventsResponse, toolDefinition, } from '@tanstack/ai' @@ -614,6 +615,21 @@ export const Route = createFileRoute('/api/middleware-test')({ middleware.push(chunkTransformMiddleware) if (middlewareMode === 'tool-skip') middleware.push(toolSkipMiddleware) + if ( + middlewareMode === 'after-tool-error' || + middlewareMode === 'after-tool-error-durable' + ) { + let calls = 0 + const failingCall = + middlewareMode === 'after-tool-error-durable' ? 2 : 1 + middleware.push({ + name: 'after-tool-error', + onAfterToolCall() { + // Throw once so retrying the hook would hide the failure. + if (++calls === failingCall) throw new Error('post-hook failed') + }, + }) + } if (middlewareMode === 'tool-replace') middleware.push(toolReplaceMiddleware) if (middlewareMode === 'tool-bad-args') @@ -771,7 +787,12 @@ export const Route = createFileRoute('/api/middleware-test')({ ? teeForPhaseCapture(rawStream, testId) : rawStream - return toServerSentEventsResponse(stream, { abortController }) + return toServerSentEventsResponse(stream, { + abortController, + ...(middlewareMode === 'after-tool-error-durable' && { + durability: { adapter: memoryStream(request) }, + }), + }) } catch (error: any) { console.error('[api.middleware-test] Error:', error.message) if (error.name === 'AbortError' || abortController.signal.aborted) { diff --git a/testing/e2e/tests/middleware.spec.ts b/testing/e2e/tests/middleware.spec.ts index e711be3b92..8a27cd907c 100644 --- a/testing/e2e/tests/middleware.spec.ts +++ b/testing/e2e/tests/middleware.spec.ts @@ -34,6 +34,55 @@ async function fetchPhaseCapture( } test.describe('Middleware Lifecycle', () => { + for (const durable of [false, true]) { + test(`a failed after-tool hook ends the run (durable=${durable})`, async ({ + page, + testId, + aimockPort, + }) => { + const response = await page.request.post('/api/middleware-test', { + data: { + threadId: testId, + runId: `${testId}-after-tool-error`, + messages: [ + { + id: 'user-1', + role: 'user', + content: durable + ? '[with-tool] run twice' + : '[with-tool] run test', + }, + ], + tools: [], + context: [], + forwardedProps: { + scenario: 'with-tool', + middlewareMode: durable + ? 'after-tool-error-durable' + : 'after-tool-error', + testId, + aimockPort, + }, + }, + }) + const events = (await response.text()) + .split('\n') + .filter((line) => line.startsWith('data: ')) + .map((line) => JSON.parse(line.slice(6)) as { type: string }) + + expect(events.some((event) => event.type === 'RUN_ERROR')).toBe(true) + expect(events.at(-1)?.type).toBe('RUN_ERROR') + expect(events.filter((event) => event.type === 'RUN_ERROR')).toHaveLength( + 1, + ) + expect( + events.filter((event) => event.type === 'TOOL_CALL_RESULT'), + ).toHaveLength(durable ? 1 : 0) + if (durable) + expect(events.some((event) => event.type === 'RUN_FINISHED')).toBe(true) + }) + } + test('adapter RUN_ERROR calls onError once', async ({ page, testId,