From 73d081a130d6b2832d8c96010deca9b1d5f52e0d Mon Sep 17 00:00:00 2001 From: Aayush <21987529+aayushbtw@users.noreply.github.com> Date: Sun, 8 Nov 2026 17:40:19 +0400 Subject: [PATCH] fix(ai): run onAbort when a chat() consumer stops iterating The engine aborts the run when the consumer stops at a yield (break or return()), so runSteps()'s finally reports onAbort and every run gets exactly one terminal hook. Closes #2097 --- .changeset/chat-consumer-stop-onabort.md | 5 ++ docs/advanced/middleware-advanced.md | 2 +- packages/ai/src/activities/chat/index.ts | 14 ++++ .../tests/consumer-stop-terminal-hook.test.ts | 71 +++++++++++++++++++ testing/e2e/src/routeTree.gen.ts | 21 ++++++ .../e2e/src/routes/api.chat-consumer-stop.ts | 43 +++++++++++ testing/e2e/tests/chat.spec.ts | 11 +++ 7 files changed, 166 insertions(+), 1 deletion(-) create mode 100644 .changeset/chat-consumer-stop-onabort.md create mode 100644 packages/ai/tests/consumer-stop-terminal-hook.test.ts create mode 100644 testing/e2e/src/routes/api.chat-consumer-stop.ts diff --git a/.changeset/chat-consumer-stop-onabort.md b/.changeset/chat-consumer-stop-onabort.md new file mode 100644 index 0000000000..36249adfca --- /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. diff --git a/docs/advanced/middleware-advanced.md b/docs/advanced/middleware-advanced.md index 34dceb5264..3ba461b4d1 100644 --- a/docs/advanced/middleware-advanced.md +++ b/docs/advanced/middleware-advanced.md @@ -785,7 +785,7 @@ 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'`. diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index 2cfdcfe546..cfdb6e62e6 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -1549,6 +1549,20 @@ class TextEngine< } async *run(): AsyncGenerator { + for await (const chunk of this.runSteps()) { + let resumed = false + try { + yield chunk + resumed = true + } finally { + // The consumer stopped iterating (`break` or `return()`). Abort before + // the loop closes `runSteps`, so its `finally` reports `onAbort`. + if (!resumed) this.middlewareAbortController?.abort() + } + } + } + + private async *runSteps(): AsyncGenerator { this.beforeRun() this.logger.agentLoop('run started', { threadId: this.middlewareCtx.threadId, 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..d95e213ad4 --- /dev/null +++ b/packages/ai/tests/consumer-stop-terminal-hook.test.ts @@ -0,0 +1,71 @@ +import { describe, expect, it } from 'vitest' +import { chat } from '../src/activities/chat/index' +import { fakeText } from '../src/testing' +import { EventType } from '../src/types' +import type { ChatMiddleware } from '../src/activities/chat/middleware/types' +import type { StreamChunk } 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 +} + +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 RUN_FINISHED runs one terminal hook', async () => { + const { hooks, middleware } = terminalRecorder() + for await (const chunk of run(middleware)) { + if (chunk.type === EventType.RUN_FINISHED) break + } + expect(hooks).toHaveLength(1) + }) + + 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, }) => {