From b31983f5b37b164e23b436776fabe71ca356e077 Mon Sep 17 00:00:00 2001 From: Aayush <21987529+aayushbtw@users.noreply.github.com> Date: Sun, 8 Nov 2026 13:42:40 +0400 Subject: [PATCH] fix(ai-effect): run middleware onAbort when a chat Stream is interrupted toStream now aborts the run's signal before it calls the iterator's return(), and awaits that teardown, so core chat() runs onAbort before the interrupt completes. Closes #2091 --- .changeset/ai-effect-interrupt-onabort.md | 5 + packages/ai-effect/src/lift.ts | 25 +++-- packages/ai-effect/tests/chat.test.ts | 111 +++++++++++++++++++++- testing/e2e/src/routes/api.effect-chat.ts | 32 +++++++ testing/e2e/tests/effect-chat.spec.ts | 15 ++- 5 files changed, 177 insertions(+), 11 deletions(-) create mode 100644 .changeset/ai-effect-interrupt-onabort.md diff --git a/.changeset/ai-effect-interrupt-onabort.md b/.changeset/ai-effect-interrupt-onabort.md new file mode 100644 index 0000000000..84fee3902f --- /dev/null +++ b/.changeset/ai-effect-interrupt-onabort.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-effect': patch +--- + +Interrupting a `chat` Stream, with `Stream.take` or `Fiber.interrupt`, now runs the middleware `onAbort` hook, and the interrupt completes after the run's teardown. diff --git a/packages/ai-effect/src/lift.ts b/packages/ai-effect/src/lift.ts index 85dcfe7f6c..d369cc4c42 100644 --- a/packages/ai-effect/src/lift.ts +++ b/packages/ai-effect/src/lift.ts @@ -39,14 +39,23 @@ export const toEffect = (run: (signal: AbortSignal) => Promise) => export const toStream = (run: (signal: AbortSignal) => AsyncIterable) => Stream.unwrap( Effect.gen(function* () { - const controller = yield* Effect.acquireRelease( - Effect.sync(() => new AbortController()), - (created) => Effect.sync(() => created.abort()), + const controller = new AbortController() + // Abort before return() so the run ends as aborted. fromAsyncIterable + // calls return() first and does not await it mid-pull, so it gets no return(). + const iterator = yield* Effect.acquireRelease( + Effect.try({ + try: () => run(controller.signal)[Symbol.asyncIterator](), + catch: toAIError, + }), + (created) => + Effect.promise(async () => { + controller.abort() + await created.return?.() + }), + ) + return Stream.fromAsyncIterable( + { [Symbol.asyncIterator]: () => ({ next: () => iterator.next() }) }, + toAIError, ) - const iterable = yield* Effect.try({ - try: () => run(controller.signal), - catch: toAIError, - }) - return Stream.fromAsyncIterable(iterable, toAIError) }), ) diff --git a/packages/ai-effect/tests/chat.test.ts b/packages/ai-effect/tests/chat.test.ts index 614a64be8a..70b9ba4dac 100644 --- a/packages/ai-effect/tests/chat.test.ts +++ b/packages/ai-effect/tests/chat.test.ts @@ -3,7 +3,13 @@ import { Context, Effect, Fiber, Layer, Schema, Stream } from 'effect' import { toolDefinition } from '@tanstack/ai' import { fakeText } from '@tanstack/ai/testing' import { chat, serverTool, toolSchema } from '../src/chat' -import type { ChatResult, ModelMessage, StreamChunk } from '@tanstack/ai' +import type { + AbortInfo, + ChatMiddleware, + ChatResult, + ModelMessage, + StreamChunk, +} from '@tanstack/ai' import type { FakeResponseStep } from '@tanstack/ai/testing' import type { AIError } from '../src/lift' @@ -212,6 +218,109 @@ describe('chat', () => { }) }) +describe('chat abort reaches middleware onAbort', () => { + function abortRecorder() { + const aborts: Array = [] + const middleware: ChatMiddleware = { + name: 'record-abort', + onAbort: async (_ctx, info) => { + // Stands in for a hook that awaits I/O, such as marking the run aborted in a store. + await new Promise((resolve) => setTimeout(resolve, 10)) + aborts.push(info) + }, + } + return { middleware, aborts } + } + + it('Stream.take stops the run through onAbort', async () => { + const { step, seen } = recordRequest('a long answer in many chunks') + const adapter = fakeWith([step]) + const { middleware, aborts } = abortRecorder() + + await Effect.runPromise( + Stream.runCollect( + chat({ + adapter, + messages, + threadId: 't1', + middleware: [middleware], + }).pipe( + Stream.filter((chunk) => chunk.type === 'TEXT_MESSAGE_CONTENT'), + Stream.take(1), + ), + ), + ) + + expect(seen.signal?.aborted).toBe(true) + expect(aborts).toHaveLength(1) + }) + + it('interrupting while the consumer handles a chunk runs onAbort', async () => { + const { step, seen } = recordRequest('a long answer in many chunks') + const adapter = fakeWith([step]) + const { middleware, aborts } = abortRecorder() + let markFirstText = () => {} + const firstText = new Promise((resolve) => { + markFirstText = resolve + }) + + const fiber = Effect.runFork( + Stream.runDrain( + chat({ + adapter, + messages, + threadId: 't1', + middleware: [middleware], + }).pipe( + Stream.tap((chunk) => + chunk.type === 'TEXT_MESSAGE_CONTENT' + ? Effect.sync(markFirstText).pipe(Effect.andThen(Effect.never)) + : Effect.void, + ), + ), + ), + ) + await firstText + await Effect.runPromise(Fiber.interrupt(fiber)) + + expect(seen.signal?.aborted).toBe(true) + expect(aborts).toHaveLength(1) + }) + + it('interrupting while a chunk is pending runs onAbort before the interrupt completes', async () => { + const { step, seen } = recordRequest('a long answer in many chunks') + const adapter = fakeText({ tokensPerSecond: 50 }) + adapter.setResponses([step]) + const { middleware, aborts } = abortRecorder() + let markFirstText = () => {} + const firstText = new Promise((resolve) => { + markFirstText = resolve + }) + + const fiber = Effect.runFork( + Stream.runDrain( + chat({ + adapter, + messages, + threadId: 't1', + middleware: [middleware], + }).pipe( + Stream.tap((chunk) => + Effect.sync(() => { + if (chunk.type === 'TEXT_MESSAGE_CONTENT') markFirstText() + }), + ), + ), + ), + ) + await firstText + await Effect.runPromise(Fiber.interrupt(fiber)) + + expect(seen.signal?.aborted).toBe(true) + expect(aborts).toHaveLength(1) + }) +}) + describe('toolSchema', () => { const input = toolSchema( Schema.Struct({ diff --git a/testing/e2e/src/routes/api.effect-chat.ts b/testing/e2e/src/routes/api.effect-chat.ts index 5aac689838..1008dd302b 100644 --- a/testing/e2e/src/routes/api.effect-chat.ts +++ b/testing/e2e/src/routes/api.effect-chat.ts @@ -11,6 +11,8 @@ import { createImageAdapter } from '@/lib/media-providers' * - Default: `AI.chat` with an Effect tool. The tool reads a `Weather` * service that a `Layer` provides. The run streams back as SSE. * - `?kind=image`: `AI.generateImage` runs as an Effect and returns JSON. + * - `?kind=stop`: `Stream.take` stops the chat after its first text chunk, + * and the JSON counts the middleware `onAbort` calls. */ class Weather extends Context.Service< Weather, @@ -55,6 +57,36 @@ export const Route = createFileRoute('/api/effect-chat')({ return Response.json(result) } + if (url.searchParams.get('kind') === 'stop') { + let aborts = 0 + await Effect.runPromise( + Stream.runDrain( + AI.chat({ + ...createTextAdapter('openai', undefined, aimockPort, testId), + threadId: 'effect-stop', + messages: [ + { + role: 'user', + content: '[abort-test] tell me a long story', + }, + ], + middleware: [ + { + name: 'count-aborts', + onAbort: () => { + aborts++ + }, + }, + ], + }).pipe( + Stream.filter((chunk) => chunk.type === 'TEXT_MESSAGE_CONTENT'), + Stream.take(1), + ), + ), + ) + return Response.json({ aborts }) + } + const chunks = AI.chat({ ...createTextAdapter('openai', undefined, aimockPort, testId), threadId: 'effect-chat', diff --git a/testing/e2e/tests/effect-chat.spec.ts b/testing/e2e/tests/effect-chat.spec.ts index df42846b8e..3df589c403 100644 --- a/testing/e2e/tests/effect-chat.spec.ts +++ b/testing/e2e/tests/effect-chat.spec.ts @@ -3,7 +3,8 @@ import { test, expect } from './fixtures' /** * `@tanstack/ai-effect` against aimock. The route runs `AI.chat` with an * Effect tool that reads a `Weather` service from a `Layer`, and - * `AI.generateImage` as an Effect. Fixtures: `fixtures/effect-chat/basic.json`. + * `AI.generateImage` as an Effect. Fixtures: `fixtures/effect-chat/basic.json`, + * and the slow story in `fixtures/abort/basic.json`. */ const PIXEL_PNG = @@ -22,7 +23,7 @@ function eventsOf(sse: string): Array { .map((line) => JSON.parse(line.slice('data: '.length))) } -function routeUrl(testId: string, aimockPort: number, kind?: 'image') { +function routeUrl(testId: string, aimockPort: number, kind?: 'image' | 'stop') { const params = new URLSearchParams({ testId, aimockPort: String(aimockPort), @@ -67,4 +68,14 @@ test.describe('ai-effect', () => { expect.objectContaining({ b64Json: PIXEL_PNG }), ]) }) + + test('stopping the chat Stream early runs middleware onAbort', async ({ + request, + testId, + aimockPort, + }) => { + const res = await request.post(routeUrl(testId, aimockPort, 'stop')) + expect(res.ok()).toBe(true) + expect(await res.json()).toEqual({ aborts: 1 }) + }) })