diff --git a/.changeset/chat-consumer-stop-onabort.md b/.changeset/chat-consumer-stop-onabort.md new file mode 100644 index 0000000000..78b5f9185e --- /dev/null +++ b/.changeset/chat-consumer-stop-onabort.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai': patch +--- + +A `chat()` consumer that stops iterating (`break` or `return()`) now runs the middleware `onAbort` hook, so every run gets exactly one terminal hook. Stopping at a client-tool or approval pause's `RUN_FINISHED` chunk runs `onFinish`, as a drained pause does, so the paused state (such as a `withSandbox` sandbox) is kept. diff --git a/docs/advanced/middleware-advanced.md b/docs/advanced/middleware-advanced.md index 34dceb5264..8b1457972f 100644 --- a/docs/advanced/middleware-advanced.md +++ b/docs/advanced/middleware-advanced.md @@ -785,10 +785,10 @@ Exactly **one** terminal hook runs per `chat()` call: | Hook | When it runs | | --- | --- | | `onFinish` | The run completed normally | -| `onAbort` | The run was aborted: by `ctx.abort()`, by an external `AbortSignal`, or by an `abort` decision from `onBeforeToolCall` | +| `onAbort` | The run was aborted: by `ctx.abort()`, by an external `AbortSignal`, by an `abort` decision from `onBeforeToolCall`, or by a consumer that stopped iterating (`break` or `return()`) | | `onError` | An error was not handled | -`onFinish` also runs when the call pauses for a client tool or an approval. The run is not complete yet. To tell a pause from a completion, look in `onChunk` for a `RUN_FINISHED` chunk with `outcome.type === 'interrupt'`. +`onFinish` also runs when the call pauses for a client tool or an approval. The run is not complete yet. To tell a pause from a completion, look in `onChunk` for a `RUN_FINISHED` chunk with `outcome.type === 'interrupt'`. A consumer that stops iterating at that pause's `RUN_FINISHED` chunk still gets `onFinish`. Stopping at any other chunk, including the final `RUN_FINISHED` of a completed run, runs `onAbort`. `onAbort` gets `info.cancelRequested`. It is `true` only when an explicit cancel aborted the run, never when the client only disconnected. diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index 2cfdcfe546..c2563139e5 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -1549,6 +1549,35 @@ class TextEngine< } async *run(): AsyncGenerator { + const steps = this.runSteps() + for await (const chunk of steps) { + let resumed = false + try { + yield chunk + resumed = true + } finally { + // The consumer stopped iterating (`break` or `return()`). + if (!resumed) { + if ( + chunk.type === EventType.RUN_FINISHED && + chunk.outcome?.type === 'interrupt' + ) { + // A pause ends this invocation. Finish it unseen so it reports + // `onFinish`, as a drained run does, and keeps the paused state. + for await (const _chunk of steps) { + // discard + } + } else { + // Abort before the loop closes `runSteps`, so its `finally` + // reports `onAbort`. + this.middlewareAbortController?.abort() + } + } + } + } + } + + private async *runSteps(): AsyncGenerator { this.beforeRun() this.logger.agentLoop('run started', { threadId: this.middlewareCtx.threadId, @@ -6143,22 +6172,28 @@ class TextEngine< }, ) - while (!settled) { + try { + while (!settled) { + yield* this.drainMiddlewareCustomQueue() + if (settled) break + await Promise.race([ + done, + new Promise((resolve) => { + if (this.middlewareCustomQueue.length > 0) { + resolve() + return + } + this.middlewareCustomWaiters.push(resolve) + }), + ]) + } + yield* this.drainMiddlewareCustomQueue() - if (settled) break - await Promise.race([ - done, - new Promise((resolve) => { - if (this.middlewareCustomQueue.length > 0) { - resolve() - return - } - this.middlewareCustomWaiters.push(resolve) - }), - ]) + } finally { + // A consumer that stops at a custom event closes this generator while + // the work runs. Wait for it, so no hook work outlives `onAbort`. + await done } - - yield* this.drainMiddlewareCustomQueue() if (error !== undefined) throw error return result as T } diff --git a/packages/ai/tests/consumer-stop-terminal-hook.test.ts b/packages/ai/tests/consumer-stop-terminal-hook.test.ts new file mode 100644 index 0000000000..05506f81cd --- /dev/null +++ b/packages/ai/tests/consumer-stop-terminal-hook.test.ts @@ -0,0 +1,153 @@ +import { describe, expect, it } from 'vitest' +import { chat } from '../src/activities/chat/index' +import { fakeText } from '../src/testing' +import { EventType } from '../src/types' +import { clientTool, createMockAdapter, ev, serverTool } from './test-utils' +import type { ChatMiddleware } from '../src/activities/chat/middleware/types' +import type { StreamChunk, Tool } from '../src/types' + +function terminalRecorder() { + const hooks: Array = [] + const middleware: ChatMiddleware = { + name: 'record-terminal', + onFinish: () => { + hooks.push('onFinish') + }, + onAbort: () => { + hooks.push('onAbort') + }, + onError: () => { + hooks.push('onError') + }, + } + return { hooks, middleware } +} + +function run(middleware: ChatMiddleware) { + const adapter = fakeText() + adapter.setResponses([{ text: 'a long answer in many chunks' }]) + return chat({ + threadId: 'thread-1', + adapter, + messages: [{ role: 'user', content: 'Hi' }], + middleware: [middleware], + }) as AsyncIterable +} + +function runWithTool(middleware: ChatMiddleware, tool: Tool) { + const { adapter, calls } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.toolStart('call-1', tool.name), + ev.toolArgs('call-1', '{}'), + ev.runFinished('tool_calls'), + ], + [ + ev.runStarted('run-2'), + ev.textStart(), + ev.textContent('done'), + ev.textEnd(), + ev.runFinished('stop', 'run-2'), + ], + ], + }) + const stream = chat({ + threadId: 'thread-1', + adapter, + messages: [{ role: 'user', content: 'Use the tool' }], + tools: [tool], + middleware: [middleware], + }) as AsyncIterable + return { stream, calls } +} + +describe('a consumer that stops iterating chat() runs one terminal hook', () => { + it('break after the first TEXT_MESSAGE_CONTENT runs onAbort', async () => { + const { hooks, middleware } = terminalRecorder() + for await (const chunk of run(middleware)) { + if (chunk.type === EventType.TEXT_MESSAGE_CONTENT) break + } + expect(hooks).toEqual(['onAbort']) + }) + + it('break after the final RUN_FINISHED runs onAbort', async () => { + const { hooks, middleware } = terminalRecorder() + for await (const chunk of run(middleware)) { + if (chunk.type === EventType.RUN_FINISHED) break + } + expect(hooks).toEqual(['onAbort']) + }) + + it('break after a tool turn RUN_FINISHED runs onAbort and no next turn', async () => { + const { hooks, middleware } = terminalRecorder() + const { stream, calls } = runWithTool( + middleware, + serverTool('lookup', () => 'found'), + ) + for await (const chunk of stream) { + if (chunk.type === EventType.RUN_FINISHED) break + } + expect(hooks).toEqual(['onAbort']) + expect(calls).toHaveLength(1) + }) + + it.each([ + ['a client tool', clientTool('ask_client')], + [ + 'an approval', + { ...serverTool('delete_file', () => 'deleted'), needsApproval: true }, + ], + ])('break after %s pause RUN_FINISHED runs onFinish', async (_kind, tool) => { + const { hooks, middleware } = terminalRecorder() + const { stream, calls } = runWithTool(middleware, tool) + let stoppedAt: StreamChunk | undefined + for await (const chunk of stream) { + if (chunk.type === EventType.RUN_FINISHED) { + stoppedAt = chunk + break + } + } + expect(stoppedAt).toMatchObject({ outcome: { type: 'interrupt' } }) + expect(hooks).toEqual(['onFinish']) + expect(calls).toHaveLength(1) + }) + + it('break at a custom event from setup runs onAbort after setup ends', async () => { + const events: Array = [] + const middleware: ChatMiddleware = { + name: 'slow-setup', + async setup(ctx) { + ctx.emitCustomEvent('setup:progress', {}) + await new Promise((resolve) => setTimeout(resolve, 20)) + events.push('setup done') + }, + onAbort: () => { + events.push('onAbort') + }, + } + for await (const chunk of run(middleware)) { + if (chunk.type === EventType.CUSTOM) break + } + expect(events).toEqual(['setup done', 'onAbort']) + }) + + it('return() after the first TEXT_MESSAGE_CONTENT runs onAbort', async () => { + const { hooks, middleware } = terminalRecorder() + const iterator = run(middleware)[Symbol.asyncIterator]() + let next = await iterator.next() + while (!next.done && next.value.type !== EventType.TEXT_MESSAGE_CONTENT) { + next = await iterator.next() + } + await iterator.return?.() + expect(hooks).toEqual(['onAbort']) + }) + + it('a full drain runs onFinish once', async () => { + const { hooks, middleware } = terminalRecorder() + for await (const _ of run(middleware)) { + // drain + } + expect(hooks).toEqual(['onFinish']) + }) +}) diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index 31c7217dd5..f0866e4936 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -79,6 +79,7 @@ import { Route as ApiByteplusSeedance1080pWireRouteImport } from './routes/api.b import { Route as ApiCallOwnershipWireRouteImport } from './routes/api.call-ownership-wire' import { Route as ApiChatRouteImport } from './routes/api.chat' import { Route as ApiChatCompletionsIncompleteStreamRouteImport } from './routes/api.chat-completions-incomplete-stream' +import { Route as ApiChatConsumerStopRouteImport } from './routes/api.chat-consumer-stop' import { Route as ApiClientToolResumeStoreRouteImport } from './routes/api.client-tool-resume-store' import { Route as ApiCloudflareBindingHeadersWireRouteImport } from './routes/api.cloudflare-binding-headers-wire' import { Route as ApiCloudflareBindingWireRouteImport } from './routes/api.cloudflare-binding-wire' @@ -588,6 +589,11 @@ const ApiChatCompletionsIncompleteStreamRoute = path: '/api/chat-completions-incomplete-stream', getParentRoute: () => rootRouteImport, } as any) +const ApiChatConsumerStopRoute = ApiChatConsumerStopRouteImport.update({ + id: '/api/chat-consumer-stop', + path: '/api/chat-consumer-stop', + getParentRoute: () => rootRouteImport, +} as any) const ApiClientToolResumeStoreRoute = ApiClientToolResumeStoreRouteImport.update({ id: '/api/client-tool-resume-store', @@ -1393,6 +1399,7 @@ export interface FileRoutesByFullPath { '/api/call-ownership-wire': typeof ApiCallOwnershipWireRoute '/api/chat': typeof ApiChatRoute '/api/chat-completions-incomplete-stream': typeof ApiChatCompletionsIncompleteStreamRoute + '/api/chat-consumer-stop': typeof ApiChatConsumerStopRoute '/api/client-tool-resume-store': typeof ApiClientToolResumeStoreRoute '/api/cloudflare-binding-headers-wire': typeof ApiCloudflareBindingHeadersWireRoute '/api/cloudflare-binding-wire': typeof ApiCloudflareBindingWireRoute @@ -1605,6 +1612,7 @@ export interface FileRoutesByTo { '/api/call-ownership-wire': typeof ApiCallOwnershipWireRoute '/api/chat': typeof ApiChatRoute '/api/chat-completions-incomplete-stream': typeof ApiChatCompletionsIncompleteStreamRoute + '/api/chat-consumer-stop': typeof ApiChatConsumerStopRoute '/api/client-tool-resume-store': typeof ApiClientToolResumeStoreRoute '/api/cloudflare-binding-headers-wire': typeof ApiCloudflareBindingHeadersWireRoute '/api/cloudflare-binding-wire': typeof ApiCloudflareBindingWireRoute @@ -1818,6 +1826,7 @@ export interface FileRoutesById { '/api/call-ownership-wire': typeof ApiCallOwnershipWireRoute '/api/chat': typeof ApiChatRoute '/api/chat-completions-incomplete-stream': typeof ApiChatCompletionsIncompleteStreamRoute + '/api/chat-consumer-stop': typeof ApiChatConsumerStopRoute '/api/client-tool-resume-store': typeof ApiClientToolResumeStoreRoute '/api/cloudflare-binding-headers-wire': typeof ApiCloudflareBindingHeadersWireRoute '/api/cloudflare-binding-wire': typeof ApiCloudflareBindingWireRoute @@ -2032,6 +2041,7 @@ export interface FileRouteTypes { | '/api/call-ownership-wire' | '/api/chat' | '/api/chat-completions-incomplete-stream' + | '/api/chat-consumer-stop' | '/api/client-tool-resume-store' | '/api/cloudflare-binding-headers-wire' | '/api/cloudflare-binding-wire' @@ -2244,6 +2254,7 @@ export interface FileRouteTypes { | '/api/call-ownership-wire' | '/api/chat' | '/api/chat-completions-incomplete-stream' + | '/api/chat-consumer-stop' | '/api/client-tool-resume-store' | '/api/cloudflare-binding-headers-wire' | '/api/cloudflare-binding-wire' @@ -2456,6 +2467,7 @@ export interface FileRouteTypes { | '/api/call-ownership-wire' | '/api/chat' | '/api/chat-completions-incomplete-stream' + | '/api/chat-consumer-stop' | '/api/client-tool-resume-store' | '/api/cloudflare-binding-headers-wire' | '/api/cloudflare-binding-wire' @@ -2669,6 +2681,7 @@ export interface RootRouteChildren { ApiCallOwnershipWireRoute: typeof ApiCallOwnershipWireRoute ApiChatRoute: typeof ApiChatRoute ApiChatCompletionsIncompleteStreamRoute: typeof ApiChatCompletionsIncompleteStreamRoute + ApiChatConsumerStopRoute: typeof ApiChatConsumerStopRoute ApiClientToolResumeStoreRoute: typeof ApiClientToolResumeStoreRoute ApiCloudflareBindingHeadersWireRoute: typeof ApiCloudflareBindingHeadersWireRoute ApiCloudflareBindingWireRoute: typeof ApiCloudflareBindingWireRoute @@ -3299,6 +3312,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiChatCompletionsIncompleteStreamRouteImport parentRoute: typeof rootRouteImport } + '/api/chat-consumer-stop': { + id: '/api/chat-consumer-stop' + path: '/api/chat-consumer-stop' + fullPath: '/api/chat-consumer-stop' + preLoaderRoute: typeof ApiChatConsumerStopRouteImport + parentRoute: typeof rootRouteImport + } '/api/client-tool-resume-store': { id: '/api/client-tool-resume-store' path: '/api/client-tool-resume-store' @@ -4415,6 +4435,7 @@ const rootRouteChildren: RootRouteChildren = { ApiChatRoute: ApiChatRoute, ApiChatCompletionsIncompleteStreamRoute: ApiChatCompletionsIncompleteStreamRoute, + ApiChatConsumerStopRoute: ApiChatConsumerStopRoute, ApiClientToolResumeStoreRoute: ApiClientToolResumeStoreRoute, ApiCloudflareBindingHeadersWireRoute: ApiCloudflareBindingHeadersWireRoute, ApiCloudflareBindingWireRoute: ApiCloudflareBindingWireRoute, diff --git a/testing/e2e/src/routes/api.chat-consumer-stop.ts b/testing/e2e/src/routes/api.chat-consumer-stop.ts new file mode 100644 index 0000000000..3391428608 --- /dev/null +++ b/testing/e2e/src/routes/api.chat-consumer-stop.ts @@ -0,0 +1,43 @@ +import { createFileRoute } from '@tanstack/react-router' +import { chat } from '@tanstack/ai' +import { createTextAdapter } from '@/lib/providers' +import type { Provider } from '@/lib/types' + +/** + * Iterates `chat()` and stops with `break` at the first text chunk, with no + * abort. Returns the terminal middleware hooks that ran. + */ +export const Route = createFileRoute('/api/chat-consumer-stop')({ + server: { + handlers: { + POST: async ({ request }) => { + const url = new URL(request.url) + const provider = url.searchParams.get('provider') as Provider + const testId = url.searchParams.get('testId') ?? undefined + const hooks: Array = [] + for await (const event of chat({ + ...createTextAdapter(provider, undefined, undefined, testId), + threadId: 'chat-consumer-stop', + messages: [{ role: 'user', content: '[chat] recommend a guitar' }], + middleware: [ + { + name: 'record-terminal-hooks', + onFinish: () => { + hooks.push('onFinish') + }, + onAbort: () => { + hooks.push('onAbort') + }, + onError: () => { + hooks.push('onError') + }, + }, + ], + })) { + if (event.type === 'TEXT_MESSAGE_CONTENT') break + } + return Response.json({ hooks }) + }, + }, + }, +}) diff --git a/testing/e2e/tests/chat.spec.ts b/testing/e2e/tests/chat.spec.ts index a50795ae49..0536bb73a7 100644 --- a/testing/e2e/tests/chat.spec.ts +++ b/testing/e2e/tests/chat.spec.ts @@ -57,6 +57,17 @@ for (const provider of providersFor('chat')) { }) } +test('a consumer that breaks out of chat() runs onAbort once', async ({ + request, + testId, +}) => { + const response = await request.post( + `/api/chat-consumer-stop?provider=openai&testId=${encodeURIComponent(testId)}`, + ) + expect(response.ok()).toBe(true) + expect(await response.json()).toEqual({ hooks: ['onAbort'] }) +}) + test('preserves UI message IDs at the server conversion boundary', async ({ request, }) => {