fix(ai): propagate errors from onAfterToolCall

.changeset/ai-after-tool-hook-error.md

@@ -0,0 +1,7 @@
+---
+'@tanstack/ai': patch
+---
+
+Treat a failed `onAfterToolCall` middleware hook as a chat run error instead of recording a second, contradictory tool result.
+
+Deliver and persist that error for durable streams even after an earlier tool iteration emitted `RUN_FINISHED`.

docs/advanced/middleware-advanced.md

@@ -637,7 +637,13 @@ A guard cannot change the arguments. To change them, use `onBeforeToolCall`.
 
 ### onAfterToolCall
 
-`onAfterToolCall` runs after each tool runs (or is skipped). Every middleware runs, in array order.
+`onAfterToolCall` runs after each tool runs (or is skipped). Hooks run in order.
+If one hook throws, later hooks do not run.
+The chat run ends through `onError`. Catch errors inside the hook when its work
+is optional.
+
+Durable response helpers also deliver and store the `RUN_ERROR`, even after an
+earlier tool iteration finished.
 
 ```typescript
 import { type ChatMiddleware } from "@tanstack/ai";

packages/ai/src/activities/chat/index.ts

@@ -123,6 +123,7 @@ import { MiddlewareRunner } from './middleware/compose'
 import { getRunDetached } from './middleware/run-store'
 import { LoadChildCapability } from './middleware/load-child'
 import { publishRunDetachedSignal } from '../../delivery-detach'
+import { publishRunCompletionSignal } from '../../delivery-completion'
 import { publishRunDisconnectHandler } from '../../delivery-disconnect'
 import { provideSandboxRuntime } from './middleware/sandbox-runtime'
 import { provideRunDisconnect } from './middleware/run-disconnect'
@@ -1163,6 +1164,7 @@ class TextEngine<
   // observe both cancellation sources via ctx.abortSignal.
   private readonly toolAbortSignal?: AbortSignal
   private terminalHookCalled = false
+  private runEnded = false
   /**
    * Latched the first time the delivery socket closes; see `notifyDisconnected`.
    * Also read by `subscribe` so a listener registered AFTER the disconnect (a
@@ -1492,6 +1494,7 @@ class TextEngine<
   }
 
   private async runTerminalHook(): Promise<void> {
+    this.runEnded = true
     if (this.terminalHookCalled || this.isCancelled()) return
 
     this.terminalHookCalled = true
@@ -1573,6 +1576,7 @@ class TextEngine<
 
       if (this.earlyTermination) {
         yield* this.emitSuccessfulEarlyTermination()
+        this.runEnded = true
         if (!this.terminalHookCalled) {
           this.terminalHookCalled = true
           await this.middlewareRunner.runOnFinish(this.middlewareCtx, {
@@ -1734,6 +1738,7 @@ class TextEngine<
         'continuationRunId' in error &&
         typeof error.continuationRunId === 'string'
       ) {
+        this.runEnded = true
         this.terminalHookCalled = true
         yield* this.pipeThroughMiddleware({
           type: EventType.RUN_FINISHED,
@@ -1823,6 +1828,10 @@ class TextEngine<
     this.middlewareCtx.toolNames = this.eventToolNames
   }
 
+  hasRunEnded(): boolean {
+    return this.runEnded
+  }
+
   private async beginCycle(): Promise<void> {
     if (this.cyclePhase === 'processText') {
       await this.beginIteration()
@@ -3980,6 +3989,7 @@ class TextEngine<
       })
     }
     for (const output of terminalOutputs) {
+      this.runEnded = true
       yield this.publicInterruptTerminal(output)
     }
     return true
@@ -6318,19 +6328,20 @@ function validateInterruptDefinitions(
 
 /**
  * The slice of the engine that the durable delivery sink reaches back into, in
- * BOTH directions: it reads the detach verdict (`wasDetached`) and pushes the
- * socket-closed fact in (`notifyDisconnected`). Filled by the generator body as
- * soon as its engine exists.
+ * BOTH directions: it reads run completion and detach state, and pushes the
+ * socket-closed fact in (`notifyDisconnected`). Filled by the generator body
+ * as soon as its engine exists.
  */
 interface DeliveryEngineRef {
   current?: {
     wasDetached: () => boolean
     notifyDisconnected: () => void
+    hasRunEnded: () => boolean
   }
 }
 
 /**
- * Publish both delivery-side seams for `stream`.
+ * Publish delivery-side state and callbacks for `stream`.
  *
  * Shared by the two streaming paths so they cannot drift apart — the
  * structured-output path having been wired for one seam and not the other is
@@ -6341,6 +6352,7 @@ function publishDeliverySeams(
   stream: object,
   engineRef: DeliveryEngineRef,
 ): void {
+  publishRunCompletionSignal(stream, () => engineRef.current?.hasRunEnded())
   // A thunk, evaluated on the sink's teardown path: the engine does not exist
   // yet, and the verdict it will report is only written during `onAbort`.
   publishRunDetachedSignal(

packages/ai/src/activities/chat/tools/tool-calls.ts

@@ -786,6 +786,7 @@ export async function* executeServerTool<TContext = unknown>(
   subagentInterrupts?: Array<Interrupt>,
 ): AsyncGenerator<CustomEvent | StreamChunk, void, void> {
   const startTime = Date.now()
+  let resultRecorded = false
   try {
     if (!tool.execute) {
       throw new Error(`Tool ${toolName} has no execute() implementation`)
@@ -821,6 +822,7 @@ export async function* executeServerTool<TContext = unknown>(
                 ? outcome.text
                 : compactForModel(outcome.result),
           }
+      resultRecorded = true
       await pushResultAndRunAfterHook(
         results,
         {
@@ -869,6 +871,7 @@ export async function* executeServerTool<TContext = unknown>(
     const finalResult =
       typeof result === 'string' ? safeJsonParse(result) : (result ?? null)
 
+    resultRecorded = true
     await pushResultAndRunAfterHook(
       results,
       {
@@ -891,6 +894,10 @@ export async function* executeServerTool<TContext = unknown>(
       middlewareHooks,
     )
   } catch (error: unknown) {
+    // A post-execution hook failed after the tool result was recorded. Let the
+    // run handle that failure; reporting it as a second tool result would give
+    // the model contradictory outcomes and call the hook a second time.
+    if (resultRecorded) throw error
     const duration = Date.now() - startTime
 
     // Flush remaining events

packages/ai/src/delivery-completion.ts

@@ -0,0 +1,15 @@
+/** Chat streams can emit RUN_FINISHED before the agent loop actually ends. */
+const completionSignals = new WeakMap<object, () => boolean | undefined>()
+
+/** Publish final-run state before terminal hooks and cleanup can fail. */
+export function publishRunCompletionSignal(
+  stream: object,
+  completed: () => boolean | undefined,
+): void {
+  completionSignals.set(stream, completed)
+}
+
+/** Undefined preserves terminal-event handling for user-provided iterables. */
+export function wasRunCompleted(stream: object): boolean | undefined {
+  return completionSignals.get(stream)?.()
+}

packages/ai/src/stream-to-response.ts

@@ -5,6 +5,7 @@ import {
   isTerminalRunStatus,
 } from './activities/chat/middleware/run-store'
 import { wasRunDetached } from './delivery-detach'
+import { wasRunCompleted } from './delivery-completion'
 import { notifyRunDisconnected } from './delivery-disconnect'
 import { resolveResumeRunId } from './stream-durability'
 import { EventType } from './types'
@@ -715,14 +716,12 @@ export function durableStreamSource<TOffset extends string>(
 
   async function* produce(): AsyncIterable<StreamChunk> {
     let batch: Array<StreamChunk> = []
-    let terminalPersisted = false
-    // Whether a terminal event was actually delivered LIVE to the consumer (as
-    // opposed to only appended to the log). Distinguishes "the run already ended
-    // on the wire" from "the log has a terminal but the consumer never saw one",
-    // which governs whether a late durability-cleanup failure may be rethrown.
-    // Only ever assigned inside the nested flush() closure, which TS's
-    // control-flow analysis can't observe (see the disable at the read site).
-    let terminalForwarded = false
+    // A persisted terminal need not have reached the live consumer. Chat
+    // iteration terminals are provisional until the engine ends the run.
+    const terminals: {
+      persisted?: 'RUN_FINISHED' | 'RUN_ERROR'
+      forwarded?: 'RUN_FINISHED' | 'RUN_ERROR'
+    } = {}
     let failure: RecordedFailure | undefined
     let terminalCause: unknown
     let hasTerminalCause = false
@@ -756,17 +755,14 @@ export function durableStreamSource<TOffset extends string>(
         validateOffset(offset)
         idByChunk.set(chunk, offset)
       })
-      if (
-        toForward.some(
-          (chunk) =>
-            chunk.type === 'RUN_FINISHED' || chunk.type === 'RUN_ERROR',
-        )
-      ) {
-        terminalPersisted = true
+      for (const chunk of toForward) {
+        if (chunk.type === 'RUN_FINISHED' || chunk.type === 'RUN_ERROR') {
+          terminals.persisted = chunk.type
+        }
       }
       for (const chunk of toForward) {
         if (chunk.type === 'RUN_FINISHED' || chunk.type === 'RUN_ERROR') {
-          terminalForwarded = true
+          terminals.forwarded = chunk.type
         }
         yield chunk
       }
@@ -908,7 +904,7 @@ export function durableStreamSource<TOffset extends string>(
       // What is ALREADY IN THE LOG is deliberately NOT a conjunct. An agent-loop
       // run emits one `RUN_FINISHED` PER ITERATION — the intermediate
       // `finishReason: 'tool_calls'` terminal is flushed at its boundary
-      // mid-run — so `terminalPersisted` means "some terminal is in the log",
+      // mid-run — so `terminals.persisted` means "some terminal is in the log",
       // never "the run ended". Gating on it terminalized the log of a healthy,
       // still-running agent for every tool-calling run. Nor can the sink
       // distinguish a final terminal from an intermediate one by its
@@ -922,6 +918,11 @@ export function durableStreamSource<TOffset extends string>(
         !hasTerminalCause &&
         wasRunDetached(stream)
 
+      const completed = wasRunCompleted(stream) !== false
+      const terminalPersisted =
+        terminals.persisted === 'RUN_ERROR' ||
+        (terminals.persisted === 'RUN_FINISHED' && completed)
+
       if (
         !detached &&
         needsTerminalPersistence(terminalPersisted, cancelled, hasTerminalCause)
@@ -933,7 +934,6 @@ export function durableStreamSource<TOffset extends string>(
         const cause = hasTerminalCause ? terminalCause : { name: 'AbortError' }
         try {
           await durability.append([runErrorChunk(cause)])
-          terminalPersisted = true
         } catch (terminalError) {
           // Rethrown to the live consumer below, but a joiner replaying the log
           // only ever sees a generic incomplete error — so record the real
@@ -982,7 +982,9 @@ export function durableStreamSource<TOffset extends string>(
       // provider that throws AFTER emitting its own terminal, whose error is
       // otherwise neither delivered nor logged.)
       if (failure !== undefined) {
-        // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition -- terminalForwarded is set only inside the flush() closure, which TS CFA narrows away here
+        const terminalForwarded =
+          terminals.forwarded === 'RUN_ERROR' ||
+          (terminals.forwarded === 'RUN_FINISHED' && completed)
         if (!terminalForwarded) {
           // eslint-disable-next-line no-unsafe-finally
           throw failure.error

packages/ai/tests/middleware.test.ts

@@ -1384,6 +1384,49 @@ describe('chat() middleware', () => {
   // onAfterToolCall error path
   // ==========================================================================
   describe('onAfterToolCall error handling', () => {
+    it('does not turn a failed after-hook into a second tool result', async () => {
+      const execute = vi.fn(() => ({ sent: true }))
+      const onAfterToolCall = vi.fn()
+      const onError = vi.fn()
+      const { adapter } = createMockAdapter({
+        iterations: [
+          [
+            ev.runStarted(),
+            ev.toolStart('tc-1', 'sendEmail'),
+            ev.toolArgs('tc-1', '{}'),
+            ev.toolEnd('tc-1', 'sendEmail', { input: {} }),
+            ev.runFinished('tool_calls'),
+          ],
+        ],
+      })
+
+      const stream = chat({
+        adapter,
+        threadId: 'thread-1',
+        messages: [{ role: 'user', content: 'Send it' }],
+        tools: [serverTool('sendEmail', execute)],
+        middleware: [
+          {
+            name: 'failing-after-hook',
+            onAfterToolCall: (...args) => {
+              onAfterToolCall(...args)
+              if (onAfterToolCall.mock.calls.length === 1) {
+                throw new Error('post-hook failed')
+              }
+            },
+            onError: (_ctx, info) => onError(info),
+          },
+        ],
+      })
+
+      await expect(
+        collectChunks(stream as AsyncIterable<StreamChunk>),
+      ).rejects.toThrow('post-hook failed')
+      expect(execute).toHaveBeenCalledOnce()
+      expect(onAfterToolCall).toHaveBeenCalledOnce()
+      expect(onError).toHaveBeenCalledOnce()
+    })
+
     it('should report tool execution errors in onAfterToolCall', async () => {
       const afterCalls: Array<{ ok: boolean; error?: unknown }> = []
 

packages/ai/tests/stream-to-response-durability.test.ts

@@ -101,6 +101,72 @@ function fixedOffsetDurability(
 }
 
 describe('toServerSentEventsResponse with durability', () => {
+  it.each(['SSE', 'NDJSON'] as const)(
+    'delivers and persists a later after-tool error over %s',
+    async (format) => {
+      const { adapter } = createMockAdapter({
+        iterations: ['first', 'second'].map((id) => [
+          ev.runStarted(),
+          ev.toolStart(id, 'lookup'),
+          ev.toolArgs(id, '{}'),
+          ev.toolEnd(id),
+          ev.runFinished('tool_calls'),
+        ]),
+      })
+      const execute = vi.fn(() => 'found')
+      const onError = vi.fn()
+      const onAfterToolCall = vi.fn(() => {
+        if (onAfterToolCall.mock.calls.length === 2) {
+          throw new Error('second hook failed')
+        }
+      })
+      const stream = chat({
+        adapter,
+        threadId: 'thread-1',
+        messages: [{ role: 'user' as const, content: 'Look twice' }],
+        tools: [
+          toolDefinition({
+            name: 'lookup',
+            description: 'Return a value',
+          }).server(execute),
+        ],
+        middleware: [{ name: 'fail-second-hook', onAfterToolCall, onError }],
+      })
+      const durability = memoryStream(
+        new Request(`https://example.test/chat?runId=after-hook-${format}`, {
+          method: 'POST',
+        }),
+      )
+      const response = (
+        format === 'SSE' ? toServerSentEventsResponse : toHttpResponse
+      )(stream, { durability: { adapter: durability } })
+      const body = await readBody(response)
+      const events =
+        format === 'SSE' ? parseSseEvents(body) : parseNdjsonEvents(body)
+      const types = events.map((event) => field(event, 'type'))
+      expect(types).toContain(EventType.RUN_FINISHED)
+      expect(types.at(-1)).toBe(EventType.RUN_ERROR)
+      expect(types.filter((type) => type === EventType.RUN_ERROR)).toHaveLength(
+        1,
+      )
+      expect(field(events.at(-1)!, 'message')).toBe('second hook failed')
+      expect(
+        events.filter(
+          (event) => field(event, 'type') === EventType.TOOL_CALL_RESULT,
+        ),
+      ).toHaveLength(1)
+      expect(execute).toHaveBeenCalledTimes(2)
+      expect(onAfterToolCall).toHaveBeenCalledTimes(2)
+      expect(onError).toHaveBeenCalledOnce()
+
+      const stored = (await durability.snapshot()).map((entry) => entry.chunk)
+      expect(stored.at(-1)?.type).toBe(EventType.RUN_ERROR)
+      expect(
+        stored.filter((chunk) => chunk.type === EventType.RUN_ERROR),
+      ).toHaveLength(1)
+    },
+  )
+
   it('keeps persisting while a connected client does not read', async () => {
     const durability = memoryStream(
       new Request('https://example.test/api/chat?runId=unread-connected', {