fix(ai-effect): run middleware onAbort when a chat Stream is interrupted

.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.

packages/ai-effect/src/lift.ts

@@ -39,14 +39,23 @@ export const toEffect = <A>(run: (signal: AbortSignal) => Promise<A>) =>
 export const toStream = <C>(run: (signal: AbortSignal) => AsyncIterable<C>) =>
   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)
     }),
   )

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<AbortInfo> = []
+    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<void>((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<void>((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({

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',

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<WireEvent> {
     .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 })
+  })
 })