chat(): a consumer that stops iterating (break / return()) runs no terminal middleware hook #2097

TanStack AI version

@tanstack/ai 1.2.0. Reproduced on main at c32f641.

Framework/Library version

Node.js 24. No framework is needed.

Describe the bug and the steps to reproduce it

When a chat() consumer stops iterating without aborting, no terminal middleware hook runs: no onFinish, onAbort or onError. The consumer can stop with break in for await or by calling the iterator's return(). A middleware's run record, sandbox or spans are then never closed. The docs say exactly one terminal hook runs per chat() call (docs/advanced/middleware.md:98, middleware-advanced.md:783).

Cause. In packages/ai/src/activities/chat/index.ts:1802-1812, the run's finally calls onAbort only when isCancelled() is true. A bare return() at a yield skips runTerminalHook() (:1740) and the catch, and the run is not cancelled, so nothing runs.

This does not look like an intended detach:

• A detach is reported through onAbort itself: ai-sandbox calls provideRunDetached inside onAbort (ai-sandbox/src/middleware.ts:315).
• Core's own detach sink never calls return(). It keeps pulling (stream-to-response.ts:273-276, 325-328).
• Core's own cancel path aborts first and only then calls return() (:330-334).

So a bare return() is just a stop that no hook records.

Related: this is the core cause of the Stream.take case in #2091. An ai-effect Stream.take(1) probe, with ai-effect untouched, fails against main's @tanstack/ai and passes with the fix below. The other case in #2091, an interrupt during a pending read, is separate: Effect does not await return() while a next() is pending, and #2094 handles that. Closed #1028 fixed an earlier gap of the same kind, for client-tool waits.

Expected: stopping iteration runs exactly one terminal hook (onAbort).

Actual: no terminal hook runs.

Your Minimal, Reproducible Example - (Sandbox Highly Recommended)

A new file, packages/ai/tests/consumer-stop-terminal-hook.test.ts. On main:

 ❯ |@tanstack/ai| tests/consumer-stop-terminal-hook.test.ts (4 tests | 3 failed)
     × break after the first TEXT_MESSAGE_CONTENT
     × break after the first RUN_FINISHED
     × return() after the first TEXT_MESSAGE_CONTENT
AssertionError: expected [] to have a length of 1 but got +0
      Tests  3 failed | 1 passed (4)

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<string> = []
  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<StreamChunk>
}

describe('a consumer that stops iterating chat() runs one terminal hook', () => {
  it.each([EventType.TEXT_MESSAGE_CONTENT, EventType.RUN_FINISHED])(
    'break after the first %s',
    async (stopAt) => {
      const { hooks, middleware } = terminalRecorder()
      for await (const chunk of run(middleware)) {
        if (chunk.type === stopAt) break
      }
      expect(hooks).toHaveLength(1)
    },
  )

  it('return() after the first TEXT_MESSAGE_CONTENT', 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).toHaveLength(1)
  })

  it('a full drain runs onFinish once', async () => {
    const { hooks, middleware } = terminalRecorder()
    for await (const _ of run(middleware)) {
      // drain
    }
    expect(hooks).toEqual(['onFinish'])
  })
})

Proposed fix

In the engine, so both streaming callers of engine.run() (:6527, :7525) get it. When the consumer stops while the run is live, abort and then return, in the same order core's own sink uses. runSteps's finally then reports onAbort:

   async *run(): AsyncGenerator<StreamChunk> {
+    const steps = this.runSteps()
+    let atYield = false
+    try {
+      for (let next = await steps.next(); !next.done; next = await steps.next()) {
+        atYield = true
+        yield next.value
+        atYield = false
+      }
+    } finally {
+      // The consumer stopped iterating while the run was live: stop it, so
+      // runSteps's finally reports onAbort.
+      if (atYield) {
+        this.middlewareAbortController?.abort()
+        await steps.return(undefined)
+      }
+    }
+  }
+
+  private async *runSteps(): AsyncGenerator<StreamChunk> {
     this.beforeRun()

With it, all 4 cases pass, and each stop case records exactly ['onAbort']. The packages/ai suite passes (158 files, 2696 tests), and tsc is clean.

I deliberately did not use the rule "any finally without a hook means onAbort". The boundary-interrupt returns at :1672 and :1687 exit with no hook and no cancel, and that rule would wrongly give them onAbort. From reading the code (not run), those pauses seem to run no terminal hook either, though the docs say onFinish runs on pauses. That may be worth its own look.

Open question. With this fix, a break right at the final RUN_FINISHED gives onAbort, even though all output was delivered. Core yields the final RUN_FINISHED before runTerminalHook() (:1695 vs :1740). Making that case onFinish would need a "final terminal already yielded" flag. I left that out to keep the change small. Happy to go either way.

Do you intend to try to help solve this bug with your own PR?

Yes, I think I know how to fix it and will discuss it in the comments of this issue

Terms & Code of Conduct

• I agree to follow this project's Code of Conduct
• I understand that if my bug cannot be reliable reproduced in a debuggable environment, it will probably not be fixed and this issue may even be closed.