@tanstack/ai 0.41.0 → 0.43.0

package/src/activities/chat/mcp/manager.ts

@@ -1,4 +1,4 @@
-import type { ServerTool } from '../tools/tool-definition'
+import type { AnyServerTool } from '../tools/tool-definition'
 import type { ChatMCPOptions, MCPToolSource } from './types'
 
 /**
@@ -18,7 +18,7 @@ import type { ChatMCPOptions, MCPToolSource } from './types'
  * drains. If discovery results were ever cached and reused across runs, this
  * would bind a closure over an already-closed source — bind onto a copy then.
  */
-function bindReadResource(tool: ServerTool, source: MCPToolSource): void {
+function bindReadResource(tool: AnyServerTool, source: MCPToolSource): void {
   if (!source.readResource) return
   const meta = (
     tool.metadata as { mcp?: { uiResourceUri?: string } } | undefined
@@ -71,13 +71,13 @@ export class MCPManager {
    * (no `onDiscoveryError`, or it re-threw) or a duplicate tool name; in that
    * case it first closes any connected sources when the policy is 'close'.
    */
-  async discover(): Promise<Array<ServerTool>> {
+  async discover(): Promise<Array<AnyServerTool>> {
     if (this.#sources.length === 0) return []
     try {
       const settled = await Promise.allSettled(
         this.#sources.map((s) => s.tools({ lazy: this.#lazyTools })),
       )
-      const tools: Array<ServerTool> = []
+      const tools: Array<AnyServerTool> = []
       const zipped = this.#sources.map(
         (source, i) => [source, settled[i]] as const,
       )

package/src/activities/chat/mcp/types.ts

@@ -1,4 +1,4 @@
-import type { ServerTool } from '../tools/tool-definition'
+import type { AnyServerTool } from '../tools/tool-definition'
 
 /**
  * The shape `readResource` resolves to — a structural subset of MCP's
@@ -26,7 +26,7 @@ export interface MCPToolSource {
   // Keep the options shape in sync with ai-mcp's `ToolsOptions` — extra
   // optional fields added there still match structurally, but chat() only
   // forwards what is declared here.
-  tools: (options?: { lazy?: boolean }) => Promise<Array<ServerTool>>
+  tools: (options?: { lazy?: boolean }) => Promise<Array<AnyServerTool>>
   close: () => Promise<void>
   /**
    * Reads an MCP resource by URI. Used by the chat manager to eagerly fetch

package/src/activities/chat/messages.ts

@@ -211,6 +211,7 @@ function isToolCallIncluded(part: ToolCallPart): boolean {
   return (
     part.state === 'input-complete' ||
     part.state === 'complete' ||
+    part.state === 'approval-requested' ||
     part.state === 'approval-responded' ||
     part.state === 'error' ||
     part.output !== undefined
@@ -614,12 +615,13 @@ export function modelMessagesToUIMessages(
         })
       } else {
         // No assistant message to merge into, create a standalone one
-        const toolResultUIMessage = modelMessageToUIMessage(msg)
+        const toolResultUIMessage = modelMessageToUIMessage(msg, msg.id)
         uiMessages.push(toolResultUIMessage)
       }
     } else {
-      // Regular message
-      const uiMessage = modelMessageToUIMessage(msg)
+      // Regular message. Preserve a persisted stable id so a hydrated message
+      // keeps the same identity as its live stream (enables in-place resume).
+      const uiMessage = modelMessageToUIMessage(msg, msg.id)
       uiMessages.push(uiMessage)
 
       // Track assistant messages for potential tool result merging

package/src/activities/chat/middleware/builder.ts

@@ -104,6 +104,6 @@ export function createChatMiddleware(): ChatMiddlewareBuilder<
   // object reused across `.use()` calls, but the type accumulates `TProvided`
   // and `TList` per call — TypeScript cannot derive that from runtime values, so
   // a structural `as` is impossible and the double assertion is irreducible.
-  // eslint-disable-next-line no-restricted-syntax -- irreducible: type-level accumulation cannot be expressed from a single runtime object
+  // oxlint-disable-next-line eslint-js/no-restricted-syntax -- irreducible: type-level accumulation cannot be expressed from a single runtime object
   return builder as unknown as ChatMiddlewareBuilder<readonly [], never>
 }

package/src/activities/chat/middleware/compose.ts

@@ -1,5 +1,5 @@
 import { aiEventClient } from '@tanstack/ai-event-client'
-import type { StreamChunk } from '../../../types'
+import type { AgentLoopState, StreamChunk } from '../../../types'
 import type { InternalLogger } from '../../../logger/internal-logger'
 import type {
   AbortInfo,
@@ -18,6 +18,12 @@ import type {
   UsageInfo,
 } from './types'
 
+/** One middleware's terminal-hook throw, captured instead of propagated. */
+interface HookFailure {
+  middleware: string
+  error: unknown
+}
+
 /** Check if a middleware should be skipped for instrumentation events. */
 function shouldSkipInstrumentation(mw: ChatMiddleware<any>): boolean {
   return mw.name === 'devtools' || mw.name === 'strip-to-spec'
@@ -474,18 +480,106 @@ export class MiddlewareRunner<TContext = unknown> {
     }
   }
 
+  /**
+   * Await ONE terminal hook and RETURN its throw instead of letting it escape
+   * the caller's loop, logging it on the `errors` channel first so the failure
+   * is never invisible. `undefined` means the hook completed.
+   *
+   * Capturing (rather than swallowing at this level) is what lets isolation and
+   * reporting coexist: every caller gives every middleware its turn, and then
+   * each decides on its own whether the collected failures are worth telling the
+   * caller about. See {@link runOnFinish} vs {@link runOnAbort} /
+   * {@link runOnError}.
+   */
+  private async captureTerminalHook(
+    mw: ChatMiddleware<TContext>,
+    hookName: 'onFinish' | 'onAbort' | 'onError',
+    invoke: () => void | Promise<void>,
+  ): Promise<HookFailure | undefined> {
+    try {
+      await invoke()
+      return undefined
+    } catch (error) {
+      this.logger.errors(`middleware ${hookName} hook failed`, {
+        middleware: mw.name ?? 'unnamed',
+        hook: hookName,
+        error,
+      })
+      return { middleware: mw.name ?? 'unnamed', error }
+    }
+  }
+
   /**
    * Run onFinish on all middleware in order.
+   *
+   * ISOLATED **and** REPORTED. `onFinish` is the only terminal fan-out on the
+   * SUCCESS path, and it is where `withPersistence.onFinish` writes the
+   * assistant turn through the store. So the two properties are needed together
+   * and neither may be traded for the other:
+   *
+   * - ISOLATION: every middleware's hook runs even if an earlier one threw, so a
+   *   transient store error cannot skip a later middleware's own bookkeeping.
+   *   Each failure is captured by {@link captureTerminalHook}, not propagated
+   *   mid-loop.
+   * - REPORTING: after the loop, the failures are rethrown. `chat()`'s catch
+   *   treats what we throw as a genuine error (it is not a
+   *   `MiddlewareAbortError`, and `structuralInterruptFailure` does not match
+   *   it) and rethrows it out of the generator.
+   *
+   * What that rethrow can and cannot achieve depends on the transport, because
+   * this fan-out is awaited AFTER the adapter's `RUN_FINISHED` has already been
+   * yielded (`chat()` yields terminal chunks while streaming, then awaits this
+   * hook on its way out). The success terminal is therefore already gone; the
+   * rethrow can only append to what the consumer saw, never retract it:
+   *
+   * - NON-DURABLE transport: the throw escapes the generator mid-response, and
+   *   the SSE / HTTP-stream encoder turns it into a TRAILING `RUN_ERROR` on the
+   *   wire carrying the store's own message and `code`. `ai-client` surfaces
+   *   that as an error status, so the user is not told the turn was saved when
+   *   it was not.
+   * - DURABLE transport: the throw reaches the durability sink instead. The
+   *   terminal was already persisted AND forwarded, so the sink deliberately
+   *   does NOT append a second, contradictory terminal, and `terminalForwarded`
+   *   (see `stream-to-response.ts`) suppresses the rethrow to the live consumer.
+   *   The `RUN_FINISHED` stands and the failure is RECORDED SERVER-SIDE on the
+   *   sink's `errors` channel. That is the intended outcome, not a gap: the save
+   *   failed, not the run — the consumer did receive the complete stream, so
+   *   telling it the run errored would be the lie. What the rethrow buys here is
+   *   that the sink sees the failure at all; while this loop swallowed, the only
+   *   trace anywhere was {@link captureTerminalHook}'s log line.
+   *
+   * Either way, swallowing is the one option ruled out: a failed
+   * `messages.append` would otherwise leave a `completed` run record with the
+   * assistant turn missing from storage and nothing beyond a middleware log
+   * line, and the client would go on to send a history the server has no record
+   * of.
+   *
+   * A single failure is rethrown AS-IS so the store's own error — its message,
+   * `cause`, `code` and `instanceof` identity — is what reaches the caller and
+   * the wire; wrapping the common case would bury it. Two or more become an
+   * `AggregateError` (never a `MiddlewareAbortError`, so it cannot be mistaken
+   * for an abort) rather than picking a winner and dropping the rest.
    */
   async runOnFinish(
     ctx: ChatMiddlewareContext<TContext>,
     info: FinishInfo,
   ): Promise<void> {
+    const failures: Array<HookFailure> = []
+    let firstFailure: HookFailure | undefined
+
     for (const mw of this.middlewares) {
-      if (mw.onFinish) {
+      const hook = mw.onFinish
+      if (hook) {
         const skip = shouldSkipInstrumentation(mw)
         const start = Date.now()
-        await mw.onFinish(ctx, info)
+        const failure = await this.captureTerminalHook(mw, 'onFinish', () =>
+          hook.call(mw, ctx, info),
+        )
+        if (failure !== undefined) {
+          firstFailure ??= failure
+          failures.push(failure)
+          continue
+        }
         if (!skip) {
           this.logger.middleware(
             `hook=onFinish middleware=${mw.name ?? 'unnamed'}`,
@@ -502,21 +596,49 @@ export class MiddlewareRunner<TContext = unknown> {
         }
       }
     }
+
+    if (firstFailure !== undefined) {
+      throw failures.length === 1
+        ? firstFailure.error
+        : new AggregateError(
+            failures.map((f) => f.error),
+            `${failures.length} middleware onFinish hooks failed: ` +
+              failures.map((f) => f.middleware).join(', '),
+          )
+    }
   }
 
   /**
    * Run onAbort on all middleware in order.
+   *
+   * ISOLATED and DELIBERATELY SWALLOWED. `onAbort` is a pure teardown fan-out
+   * released from `chat()`'s `finally`, on a path where the outcome is already
+   * decided: the run stopped, and the caller is being told why. A throw here has
+   * nothing better to report than the abort reason it would DISPLACE — the
+   * `finally` would surface a flaky store's error in place of "client
+   * disconnected" — so failures are logged on the `errors` channel and go no
+   * further. That is not a silent failure; it is refusing to let teardown
+   * rewrite an outcome it did not produce.
+   *
+   * Isolation matters independently: these hooks release PER-MIDDLEWARE
+   * resources (`withSandbox.onAbort` detaches or destroys the sandbox and stamps
+   * `detachedSince`; `withPersistence.onAbort` records the run status through the
+   * store), so an unguarded loop turns one transient store error into a
+   * permanently leaked sandbox for every middleware ordered after it.
    */
   async runOnAbort(
     ctx: ChatMiddlewareContext<TContext>,
     info: AbortInfo,
   ): Promise<void> {
     for (const mw of this.middlewares) {
-      if (mw.onAbort) {
+      const hook = mw.onAbort
+      if (hook) {
         const skip = shouldSkipInstrumentation(mw)
         const start = Date.now()
-        await mw.onAbort(ctx, info)
-        if (!skip) {
+        const failure = await this.captureTerminalHook(mw, 'onAbort', () =>
+          hook.call(mw, ctx, info),
+        )
+        if (failure === undefined && !skip) {
           this.logger.middleware(
             `hook=onAbort middleware=${mw.name ?? 'unnamed'}`,
             { middleware: mw.name ?? 'unnamed', hook: 'onAbort' },
@@ -536,17 +658,32 @@ export class MiddlewareRunner<TContext = unknown> {
 
   /**
    * Run onError on all middleware in order.
+   *
+   * ISOLATED and DELIBERATELY SWALLOWED, for the same reason as
+   * {@link runOnAbort} and NOT merely because it is teardown: the run has
+   * already failed, `info.error` IS that failure, and `chat()` rethrows it to the
+   * caller the moment this fan-out returns. A propagated hook throw could only
+   * REPLACE the run's real error with a teardown artifact — strictly less
+   * information for the caller, who is already learning the run failed. Reporting
+   * would buy nothing and cost the diagnosis, so failures are logged on the
+   * `errors` channel and stop there.
+   *
+   * Contrast {@link runOnFinish}, where nothing else is telling the caller
+   * anything is wrong — which is why that one reports.
    */
   async runOnError(
     ctx: ChatMiddlewareContext<TContext>,
     info: ErrorInfo,
   ): Promise<void> {
     for (const mw of this.middlewares) {
-      if (mw.onError) {
+      const hook = mw.onError
+      if (hook) {
         const skip = shouldSkipInstrumentation(mw)
         const start = Date.now()
-        await mw.onError(ctx, info)
-        if (!skip) {
+        const failure = await this.captureTerminalHook(mw, 'onError', () =>
+          hook.call(mw, ctx, info),
+        )
+        if (failure === undefined && !skip) {
           this.logger.middleware(
             `hook=onError middleware=${mw.name ?? 'unnamed'}`,
             { middleware: mw.name ?? 'unnamed', hook: 'onError' },
@@ -595,6 +732,46 @@ export class MiddlewareRunner<TContext = unknown> {
     }
   }
 
+  /**
+   * Run onShouldContinue through middleware in order (AND semantics).
+   * Any explicit `false` stops further iterations; `true` / void / undefined pass.
+   * Called after `agentLoopStrategy` has already approved continuation.
+   */
+  async runOnShouldContinue(
+    ctx: ChatMiddlewareContext<TContext>,
+    state: AgentLoopState,
+  ): Promise<boolean> {
+    for (const mw of this.middlewares) {
+      if (mw.onShouldContinue) {
+        const skip = shouldSkipInstrumentation(mw)
+        const start = Date.now()
+        const result = await mw.onShouldContinue(ctx, state)
+        if (!skip) {
+          this.logger.middleware(
+            `hook=onShouldContinue middleware=${mw.name ?? 'unnamed'}`,
+            {
+              middleware: mw.name ?? 'unnamed',
+              hook: 'onShouldContinue',
+              result,
+            },
+          )
+          aiEventClient.emit('middleware:hook:executed', {
+            ...instrumentCtx(ctx),
+            middlewareName: mw.name || 'unnamed',
+            hookName: 'onShouldContinue',
+            iteration: ctx.iteration,
+            duration: Date.now() - start,
+            hasTransform: result === false,
+          })
+        }
+        if (result === false) {
+          return false
+        }
+      }
+    }
+    return true
+  }
+
   /**
    * Run onToolPhaseComplete on all middleware in order.
    * Called after all tool calls in an iteration have been processed.

package/src/activities/chat/middleware/index.ts

@@ -3,6 +3,8 @@ export type {
   ChatMiddlewareContext,
   ChatMiddlewarePhase,
   ChatMiddlewareConfig,
+  ChatResumeToolState,
+  ChatResumeGenericResolution,
   StructuredOutputMiddlewareConfig,
   ToolCallHookContext,
   BeforeToolCallDecision,
@@ -39,3 +41,27 @@ export type {
 } from './builder'
 export { validateCapabilities } from './validate'
 export type { AnyChatMiddleware } from './types'
+
+export {
+  LocksCapability,
+  getLocks,
+  provideLocks,
+  InMemoryLockStore,
+  withLocks,
+  defineLock,
+} from './locks'
+export type { LockStore } from './locks'
+
+export {
+  isRunStatus,
+  isTerminalRunStatus,
+  defineRunStore,
+  InMemoryRunStore,
+} from './run-store'
+export type {
+  RunStatus,
+  TerminalRunStatus,
+  RunRecord,
+  RunError,
+  RunStore,
+} from './run-store'

package/src/activities/chat/middleware/locks.ts

@@ -0,0 +1,102 @@
+/**
+ * Distributed-mutex primitive — the neutral home for the `'locks'` capability.
+ *
+ * Capability identity is by object reference (see `createCapability`). Any
+ * middleware may PROVIDE a {@link LockStore} via {@link withLocks} /
+ * {@link provideLocks}; consumers (notably `@tanstack/ai-sandbox` `ensure`)
+ * read it with {@link getLocks}. Coordination, not state persistence.
+ */
+import { createCapability } from './capabilities'
+import { defineChatMiddleware } from './define'
+import type { ChatMiddleware, ChatMiddlewareContext } from './types'
+
+/**
+ * Mutual exclusion around a critical section keyed by `key`. A distributed
+ * backend (e.g. a Cloudflare Durable Object) is the only kind safe across
+ * instances; the in-memory default is correct within a single process only.
+ * Lease-backed implementations abort `signal` as soon as ownership can no longer
+ * be guaranteed; the callback must stop externally visible mutations when it
+ * aborts. Callbacks that ignore `signal` (e.g. the sandbox `ensure` critical
+ * section) remain valid — a `() => Promise<T>` is assignable to the
+ * signal-taking parameter.
+ */
+export interface LockStore {
+  withLock: <T>(
+    key: string,
+    fn: (signal: AbortSignal) => Promise<T>,
+  ) => Promise<T>
+}
+
+/**
+ * Type a {@link LockStore} implementation inline: pass the object and get
+ * autocomplete + contract checking, with no separate `: LockStore` annotation.
+ * Hand the result to {@link withLocks}.
+ */
+export function defineLock(lock: LockStore): LockStore {
+  return lock
+}
+
+/**
+ * The lock capability. Provided by {@link withLocks} or any middleware that
+ * calls {@link provideLocks}.
+ */
+export const LocksCapability = createCapability<LockStore>()('locks')
+
+/** Destructured accessors: `getLocks(ctx)` / `provideLocks(ctx, store)`. */
+export const [getLocks, provideLocks] = LocksCapability
+
+/**
+ * In-memory {@link LockStore} — a per-key promise chain. Correct within a single
+ * process; multi-instance correctness needs a distributed lock backend.
+ */
+export class InMemoryLockStore implements LockStore {
+  private readonly chains = new Map<string, Promise<unknown>>()
+
+  withLock<T>(
+    key: string,
+    fn: (signal: AbortSignal) => Promise<T>,
+  ): Promise<T> {
+    const prior = this.chains.get(key) ?? Promise.resolve()
+    const runCriticalSection = () => fn(new AbortController().signal)
+    // Chain after the prior holder regardless of how it settled.
+    const run = prior.then(runCriticalSection, runCriticalSection)
+    // Swallow rejections so one failure doesn't poison the lock, then drop the
+    // chain entry once this tail is still the latest — otherwise long-lived
+    // processes accumulate settled promises for every distinct key forever.
+    const settled = run.then(
+      () => undefined,
+      () => undefined,
+    )
+    this.chains.set(key, settled)
+    void settled.then(() => {
+      if (this.chains.get(key) === settled) {
+        this.chains.delete(key)
+      }
+    })
+    return run
+  }
+}
+
+/**
+ * Provide a {@link LockStore} on the chat middleware capability bus.
+ *
+ * Coordination only — independent of chat state persistence. A lock provided
+ * here reaches any later middleware that reads {@link LocksCapability}
+ * (including `withSandbox`).
+ *
+ * ```ts
+ * middleware: [
+ *   withLocks(distributedLocks),
+ *   withSandbox(sandbox),
+ * ]
+ * ```
+ */
+export function withLocks(locks: LockStore): ChatMiddleware {
+  return defineChatMiddleware({
+    name: 'locks',
+    provides: [LocksCapability],
+    setup(ctx: ChatMiddlewareContext) {
+      provideLocks(ctx, locks)
+    },
+  })
+}

package/src/activities/chat/middleware/pending-turn.ts

@@ -0,0 +1,47 @@
+/**
+ * Internal seam that lets a SLOW middleware ask the persistence layer to store
+ * the user's pending turn NOW, before that middleware starts its slow work.
+ *
+ * WHY THIS EXISTS. Chat persistence stores the pending turn from `onStart`, which
+ * is the earliest hook that holds the merged message list. `onStart` runs after
+ * EVERY middleware `setup`, and that is normally a few milliseconds. `withSandbox`
+ * breaks the assumption: its `setup` creates a sandbox and clones a repository,
+ * which takes minutes. For that whole window the thread holds nothing, so a reload
+ * — or a second device — asks the server for the conversation and is told it is
+ * empty. The user sees no sign of the message they just sent.
+ *
+ * The provider owns the rule for WHAT to store. A caller must not rebuild that
+ * rule: `saveThread` replaces the whole thread, so a caller that stored only the
+ * newly-sent message would delete the history. Asking the owner to store the turn
+ * keeps the merge in one place.
+ *
+ * OPT-IN, so nothing changes for a fast run. Persistence offers this seam on every
+ * durable run; only a middleware that is about to be slow calls it. A run with no
+ * such middleware never calls it, and the `onStart` store stays the only one.
+ *
+ * A SEAM RATHER THAN A NEW MIDDLEWARE HOOK, deliberately. A lifecycle hook that
+ * runs before `setup` would be public API, and it would have to explain itself to
+ * every middleware author. This concern has one provider (`withPersistence`) and
+ * one caller (`withSandbox`), and it reaches across packages through the same
+ * internal channel the sandbox layer already uses for its runtime.
+ */
+import { createCapability } from './capabilities'
+
+export interface PendingTurnSnapshot {
+  /**
+   * Store the user's pending turn now.
+   *
+   * Idempotent: the later `onStart` store replaces the thread with the same or a
+   * more complete list, so calling this changes what is visible EARLIER without
+   * changing what is visible at the end.
+   *
+   * Rejects only if the store itself fails. Callers treat that as non-fatal — a
+   * run that cannot pre-store its turn is still a run worth doing.
+   */
+  snapshot: () => Promise<void>
+}
+
+export const PendingTurnCapability =
+  createCapability<PendingTurnSnapshot>()('pending-turn')
+
+export const [getPendingTurn, providePendingTurn] = PendingTurnCapability

package/src/activities/chat/middleware/run-disconnect.ts

@@ -0,0 +1,62 @@
+/**
+ * Internal seam the chat engine PROVIDES so middleware can learn that the
+ * DELIVERY socket closed while the run is still going — and, crucially, learn it
+ * WITHOUT the run being cancelled.
+ *
+ * WHY THIS EXISTS. Before it, the only way a disconnect reached middleware was for
+ * the application to mirror `request.signal` into `chat()`'s `abortController`,
+ * which ABORTS THE RUN. For a durable run that is precisely wrong, and wrong in
+ * the most expensive direction: `chat()` returns at its `isCancelled()` check
+ * immediately after middleware `setup`, so the harness adapter's `chatStream` is
+ * never called and the agent in the sandbox that `setup` just spent minutes
+ * creating is NEVER LAUNCHED. The user switched away during "starting the
+ * sandbox", came back, and found an empty log belonging to a run that had done
+ * nothing — with no takeover able to recover it, because an agent that never
+ * started wrote no journal to replay. Applications were forced to choose between
+ * "the middleware learns about the disconnect" and "the run survives it".
+ *
+ * So a disconnect is delivered as a NOTIFICATION. `withSandbox` uses it to stamp
+ * `detachedSince`/`sandboxKey` and publish the detach verdict while the run keeps
+ * producing into its still-open durable log — which is exactly what a re-attaching
+ * client tails to catch up.
+ *
+ * A SUBSCRIPTION RATHER THAN A MIDDLEWARE HOOK, deliberately. `ChatMiddleware` is
+ * public API and a new lifecycle hook there is a permanent commitment — including
+ * the obligation to explain that it is the one hook that is NOT terminal. This
+ * concern has exactly one consumer in the tree (`withSandbox`) and it reaches it
+ * through the same internal channel the sandbox layer already uses for its runtime
+ * (`SandboxRuntimeCapability`), so it ships with no public surface at all.
+ *
+ * DISPATCHED FROM THE TRANSPORT, NOT FROM THE RUN'S UNWINDING. That is what makes
+ * it prompt. A run suspended inside a minutes-wide `setup` cannot dispatch anything
+ * from its own `finally`, because the `finally` is reached only once the generator
+ * unwinds — which is what made `detachedSince` land three minutes late.
+ */
+import { createCapability } from './capabilities'
+
+/**
+ * Registry of disconnect listeners for one run.
+ *
+ * Listeners must do BOOKKEEPING ONLY. The run is still executing, so releasing
+ * anything it depends on — stopping a file watcher, destroying a sandbox — breaks
+ * a healthy run. Teardown belongs in the terminal hooks, which still run exactly
+ * once afterwards.
+ *
+ * A listener may return a promise; the engine awaits all of them before the run
+ * finishes, so bookkeeping cannot be lost to a race with the run's own completion.
+ */
+export interface RunDisconnect {
+  /**
+   * Register `listener`, called at most once per run when the delivery socket
+   * closes. Registering after the socket has ALREADY closed calls `listener`
+   * immediately — otherwise a middleware whose `setup` was still running during
+   * the disconnect would silently never hear about it, which is the exact window
+   * the common disconnect lands in.
+   */
+  subscribe: (listener: () => void | Promise<void>) => void
+}
+
+export const RunDisconnectCapability =
+  createCapability<RunDisconnect>()('run-disconnect')
+
+export const [getRunDisconnect, provideRunDisconnect] = RunDisconnectCapability