@tanstack/ai 0.28.0 → 0.31.0

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

@@ -0,0 +1,162 @@
+/** Options accepted by a capability getter. */
+export interface CapabilityGetOptions {
+  /** When true, return undefined instead of throwing if the capability is absent. */
+  optional?: boolean
+}
+
+/**
+ * The minimal context shape a capability accessor needs. The full
+ * `ChatMiddlewareContext` satisfies this (it has `capabilities`), so accessors
+ * accept any middleware context without referencing `any`.
+ */
+export interface CapabilityContext {
+  capabilities: CapabilityRegistry
+}
+
+/** Reads a capability value off a context. Overloaded so the flag narrows the return. */
+export interface CapabilityGetter<TValue> {
+  (ctx: CapabilityContext): TValue
+  (ctx: CapabilityContext, opts: { optional: true }): TValue | undefined
+}
+
+/** Writes a capability value onto a context. */
+export type CapabilityProvider<TValue> = (
+  ctx: CapabilityContext,
+  value: TValue,
+) => void
+
+/**
+ * A capability handle. It is BOTH a `[get, provide]` tuple (array-destructurable)
+ * AND the identity used in middleware `requires`/`provides` declarations.
+ *
+ * Runtime identity is this object's reference. The `capabilityName` literal is
+ * used for diagnostics and COMPILE-TIME tracking only — capability names MUST be
+ * unique across an app or the type-level coverage check conflates them.
+ */
+export type Capability<
+  TValue = unknown,
+  TName extends string = string,
+> = readonly [
+  get: CapabilityGetter<TValue>,
+  provide: CapabilityProvider<TValue>,
+] & {
+  readonly capabilityName: TName
+  /** @internal Presence check for the post-setup assertion. */
+  has: (ctx: CapabilityContext) => boolean
+}
+
+/**
+ * A capability handle with permissive value/name — for use as a constraint in
+ * `requires`/`provides` arrays. Concentrates `any` in one named alias (same
+ * convention as `AnyTextAdapter`/`AnyTool`); needed so `Capability<SpecificT>`
+ * is assignable to the handle-array element type.
+ */
+export type CapabilityHandle = Capability<any, string>
+
+/**
+ * Per-request bookkeeping: which capabilities were provided, plus the
+ * duplicate-provide notification. Capability VALUES live in per-capability
+ * WeakMaps (see `createCapability`), not here — this only tracks presence.
+ */
+export class CapabilityRegistry {
+  private readonly provided = new Set<CapabilityHandle>()
+  private onDuplicate?: (name: string) => void
+
+  /** Register a callback fired when a handle is provided more than once. */
+  setOnDuplicate(cb: (name: string) => void): void {
+    this.onDuplicate = cb
+  }
+
+  /** Record that `handle` was provided; fire the duplicate callback on repeats. */
+  markProvided(handle: CapabilityHandle): void {
+    if (this.provided.has(handle)) this.onDuplicate?.(handle.capabilityName)
+    this.provided.add(handle)
+  }
+
+  has(handle: CapabilityHandle): boolean {
+    return this.provided.has(handle)
+  }
+}
+
+/**
+ * Create a capability. Returns a hybrid handle that destructures to
+ * `[get, provide]` and is itself the identity for `requires`/`provides`.
+ *
+ * Curried so the value type is supplied explicitly while the name literal is
+ * INFERRED from the argument: `createCapability<T>()('name')`. (A single call
+ * `createCapability<T>('name')` cannot work — supplying `T` explicitly stops
+ * TypeScript inferring the name, collapsing it to `string` and defeating the
+ * compile-time coverage check that keys on the literal name.)
+ *
+ * @example Provider + consumer middleware
+ * ```ts
+ * const counterCapability = createCapability<{ value: number }>()('counter')
+ * const [getCounter, provideCounter] = counterCapability
+ *
+ * const withCounter = defineChatMiddleware({
+ *   name: 'counter',
+ *   provides: [counterCapability],
+ *   setup(ctx) { provideCounter(ctx, { value: 0 }) },
+ * })
+ *
+ * const readsCounter = defineChatMiddleware({
+ *   name: 'reads-counter',
+ *   requires: [counterCapability],
+ *   onChunk(ctx) { getCounter(ctx).value++ },
+ * })
+ *
+ * chat({ adapter, messages, middleware: [withCounter, readsCounter] })
+ * ```
+ *
+ * @remarks Capability `name`s must be unique across your app: compile-time
+ * coverage tracking keys on the name literal (runtime keys on reference).
+ */
+export function createCapability<TValue = unknown>(): <
+  const TName extends string,
+>(
+  name: TName,
+) => Capability<TValue, TName> {
+  return <const TName extends string>(
+    name: TName,
+  ): Capability<TValue, TName> => {
+    // Each capability owns a typed WeakMap keyed by the context object. Because
+    // the value type is TValue, reads are typed with no assertion.
+    const values = new WeakMap<CapabilityContext, TValue>()
+
+    function get(ctx: CapabilityContext): TValue
+    function get(
+      ctx: CapabilityContext,
+      opts: { optional: true },
+    ): TValue | undefined
+    function get(
+      ctx: CapabilityContext,
+      opts?: CapabilityGetOptions,
+    ): TValue | undefined {
+      if (!values.has(ctx)) {
+        if (opts?.optional) return undefined
+        throw new Error(
+          `Capability "${name}" was requested but never provided. Ensure a ` +
+            `middleware provides it in setup(), ordered before this consumer.`,
+        )
+      }
+      return values.get(ctx)
+    }
+
+    const provide: CapabilityProvider<TValue> = (ctx, value) => {
+      values.set(ctx, value)
+      ctx.capabilities.markProvided(handle)
+    }
+
+    const pair: readonly [
+      CapabilityGetter<TValue>,
+      CapabilityProvider<TValue>,
+    ] = [get, provide]
+    // Object.assign's return type is the intersection of the tuple and the
+    // props, which IS Capability<TValue, TName> — no cast needed.
+    const handle = Object.assign(pair, {
+      capabilityName: name,
+      has: (ctx: CapabilityContext) => values.has(ctx),
+    })
+    return handle
+  }
+}

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

@@ -163,6 +163,57 @@ export class MiddlewareRunner<TContext = unknown> {
     return current
   }
 
+  /**
+   * Run all `setup` hooks in array order, then assert every declared `provides`
+   * capability was actually provided. Wires the last-wins duplicate-provide
+   * warning into the registry. Runs before init `onConfig`.
+   *
+   * Takes the full `ChatMiddlewareContext` — the same stable context the engine
+   * threads through every other hook — because it both forwards `ctx` to each
+   * `setup` hook and emits instrumentation events from it.
+   */
+  async runSetup(ctx: ChatMiddlewareContext<TContext>): Promise<void> {
+    ctx.capabilities.setOnDuplicate((name) => {
+      this.logger.warn(
+        `capability "${name}" was provided more than once; last provider wins`,
+        { capability: name },
+      )
+    })
+
+    for (const mw of this.middlewares) {
+      if (mw.setup) {
+        const skip = shouldSkipInstrumentation(mw)
+        const start = Date.now()
+        await mw.setup(ctx)
+        if (!skip) {
+          this.logger.middleware(
+            `hook=setup middleware=${mw.name ?? 'unnamed'}`,
+            { middleware: mw.name ?? 'unnamed', hook: 'setup' },
+          )
+          aiEventClient.emit('middleware:hook:executed', {
+            ...instrumentCtx(ctx),
+            middlewareName: mw.name || 'unnamed',
+            hookName: 'setup',
+            iteration: ctx.iteration,
+            duration: Date.now() - start,
+            hasTransform: false,
+          })
+        }
+      }
+    }
+
+    for (const mw of this.middlewares) {
+      for (const handle of mw.provides ?? []) {
+        if (!ctx.capabilities.has(handle)) {
+          throw new Error(
+            `Middleware "${mw.name ?? 'unnamed'}" declares it provides ` +
+              `"${handle.capabilityName}" but never called provide() in setup().`,
+          )
+        }
+      }
+    }
+  }
+
   /**
    * Call onStart on all middleware in order.
    */

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

@@ -0,0 +1,34 @@
+import type { CapabilityHandle } from './capabilities'
+import type { ChatMiddleware } from './types'
+
+/**
+ * A middleware whose `requires`/`provides` tuple types are captured precisely
+ * (via `const` inference) for the array coverage check and the builder.
+ */
+export interface DefinedChatMiddleware<
+  TContext,
+  TRequires extends ReadonlyArray<CapabilityHandle>,
+  TProvides extends ReadonlyArray<CapabilityHandle>,
+> extends ChatMiddleware<TContext> {
+  requires?: TRequires
+  provides?: TProvides
+}
+
+/**
+ * Identity helper for authoring middleware with precise capability inference.
+ * Returns the middleware unchanged at runtime; only sharpens its type so the
+ * `chat()` array coverage check and `createChatMiddleware` builder can read the
+ * exact `requires`/`provides`.
+ */
+export function defineChatMiddleware<
+  TContext = unknown,
+  const TRequires extends ReadonlyArray<CapabilityHandle> = readonly [],
+  const TProvides extends ReadonlyArray<CapabilityHandle> = readonly [],
+>(
+  middleware: ChatMiddleware<TContext> & {
+    requires?: TRequires
+    provides?: TProvides
+  },
+): DefinedChatMiddleware<TContext, TRequires, TProvides> {
+  return middleware
+}

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

@@ -16,3 +16,23 @@ export type {
 } from './types'
 
 export { MiddlewareRunner } from './compose'
+
+export { createCapability, CapabilityRegistry } from './capabilities'
+export type {
+  Capability,
+  CapabilityHandle,
+  CapabilityContext,
+  CapabilityGetter,
+  CapabilityProvider,
+  CapabilityGetOptions,
+} from './capabilities'
+export { defineChatMiddleware } from './define'
+export type { DefinedChatMiddleware } from './define'
+export { createChatMiddleware } from './builder'
+export type {
+  ChatMiddlewareBuilder,
+  MissingCapabilities,
+  NamesOf,
+} from './builder'
+export { validateCapabilities } from './validate'
+export type { AnyChatMiddleware } from './types'

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

@@ -7,6 +7,11 @@ import type {
   ToolCall,
 } from '../../../types'
 import type { SystemPrompt } from '../../../system-prompts'
+import type {
+  Capability,
+  CapabilityHandle,
+  CapabilityRegistry,
+} from './capabilities'
 
 // ===========================
 // Middleware Context
@@ -115,6 +120,28 @@ export interface ChatMiddlewareContext<TContext = unknown> {
   messages: ReadonlyArray<ModelMessage>
   /** Generate a unique ID with the given prefix */
   createId: (prefix: string) => string
+  /**
+   * Capability bookkeeping for this request. Populated by middleware `setup`
+   * hooks (via `provide` accessors) and read by later middleware (via `get`
+   * accessors). Prefer the accessors returned by `createCapability` over using
+   * this directly. Orthogonal to `context` (the user runtime context).
+   */
+  capabilities: CapabilityRegistry
+  /**
+   * Read a provided capability by its handle. Equivalent to the handle's own
+   * `get` accessor (`getX(ctx)`); throws if the capability was never provided.
+   */
+  get: <TValue>(capability: Capability<TValue>) => TValue
+  /**
+   * Read a capability by its handle, returning `undefined` if it was never
+   * provided (never throws).
+   */
+  getOptional: <TValue>(capability: Capability<TValue>) => TValue | undefined
+  /**
+   * Provide a capability value. Equivalent to the handle's own `provide`
+   * accessor (`provideX(ctx, value)`). Typically called from `setup`.
+   */
+  provide: <TValue>(capability: Capability<TValue>, value: TValue) => void
 }
 
 // ===========================
@@ -344,6 +371,36 @@ export interface ChatMiddleware<TContext = unknown> {
   /** Optional name for debugging and identification */
   name?: string
 
+  /**
+   * Capabilities this middleware requires. `chat()` validates that some
+   * middleware (or the adapter) provides each one; unsatisfied requirements are
+   * a compile-time error (array coverage / builder) and a runtime error before
+   * the adapter runs.
+   */
+  requires?: ReadonlyArray<CapabilityHandle>
+
+  /**
+   * Capabilities this middleware provides. Each declared capability MUST be
+   * provided (via its `provide` accessor) inside `setup`, or `chat()` throws
+   * after the setup phase.
+   */
+  provides?: ReadonlyArray<CapabilityHandle>
+
+  /**
+   * Capabilities this middleware uses if present but does not require.
+   * Non-gating: never causes a validation error. Read with
+   * `getX(ctx, { optional: true })`.
+   */
+  optionalRequires?: ReadonlyArray<CapabilityHandle>
+
+  /**
+   * Provisioning hook. Runs FIRST — before `onConfig` (init) — across all
+   * middleware in array order. Use it to call `provide` accessors so later
+   * middleware (`onConfig` onward) can consume the capabilities. Receives the
+   * stable context; does NOT receive the mutable config.
+   */
+  setup?: (ctx: ChatMiddlewareContext<TContext>) => void | Promise<void>
+
   /**
    * Called to observe or transform the chat configuration.
    * Called at init and at the beginning of each agent iteration.
@@ -476,3 +533,6 @@ export interface ChatMiddleware<TContext = unknown> {
     info: ErrorInfo,
   ) => void | Promise<void>
 }
+
+/** A `ChatMiddleware` with a permissive context — for use as a constraint. */
+export type AnyChatMiddleware = ChatMiddleware<any>

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

@@ -0,0 +1,55 @@
+import type { CapabilityHandle } from './capabilities'
+import type { AnyChatMiddleware } from './types'
+
+/** Minimal adapter shape needed for capability validation. */
+interface CapabilityRequiringAdapter {
+  name: string
+  requires?: ReadonlyArray<CapabilityHandle>
+}
+
+/**
+ * Runtime validation: every required capability (from middleware `requires` and
+ * the adapter's `requires`) must be provided by some middleware's `provides`.
+ * `optionalRequires` is never gating. Throws a clear error otherwise.
+ *
+ * Presence only — ORDER is not validated here (a provider may appear after its
+ * consumer in the array and still pass). Use the `createChatMiddleware()`
+ * builder for compile-time order enforcement; at runtime, a consumer that reads
+ * a not-yet-provided capability during `setup` fails loud via its getter.
+ */
+export function validateCapabilities(
+  middlewares: ReadonlyArray<AnyChatMiddleware>,
+  adapter: CapabilityRequiringAdapter,
+): void {
+  const provided = new Set<CapabilityHandle>()
+  for (const mw of middlewares) {
+    for (const handle of mw.provides ?? []) provided.add(handle)
+  }
+
+  const providedNames = (): string => {
+    const names = [...provided].map((h) => h.capabilityName)
+    return names.length ? names.join(', ') : 'none'
+  }
+
+  for (const handle of adapter.requires ?? []) {
+    if (!provided.has(handle)) {
+      throw new Error(
+        `Adapter "${adapter.name}" requires capability "${handle.capabilityName}". ` +
+          `Provided capabilities: ${providedNames()}. ` +
+          `Add a middleware that provides "${handle.capabilityName}".`,
+      )
+    }
+  }
+
+  for (const mw of middlewares) {
+    for (const handle of mw.requires ?? []) {
+      if (!provided.has(handle)) {
+        throw new Error(
+          `Middleware "${mw.name ?? 'unnamed'}" requires capability ` +
+            `"${handle.capabilityName}". Provided capabilities: ${providedNames()}. ` +
+            `Add a middleware that provides "${handle.capabilityName}".`,
+        )
+      }
+    }
+  }
+}

package/src/activities/chat/stream/message-updaters.ts

@@ -226,7 +226,7 @@ export function updateToolCallWithOutput(
       parts[index] = {
         ...toolCallPart,
         output: errorText ? { error: errorText } : output,
-        state: state ?? (errorText ? 'input-complete' : 'complete'),
+        state: state ?? (errorText ? 'error' : 'complete'),
       }
     }
 

package/src/activities/chat/stream/processor.ts

@@ -317,7 +317,7 @@ export class StreamProcessor {
       this.messages,
       toolCallId,
       output,
-      error ? 'input-complete' : undefined,
+      error ? 'error' : undefined,
       error,
     )
 
@@ -1184,7 +1184,7 @@ export class StreamProcessor {
         this.messages,
         chunk.toolCallId,
         output,
-        chunk.state === 'output-error' ? 'input-complete' : undefined,
+        chunk.state === 'output-error' ? 'error' : undefined,
       )
 
       // Step 2: Create/update the tool-result part (for LLM conversation history)
@@ -1240,7 +1240,7 @@ export class StreamProcessor {
       this.messages,
       chunk.toolCallId,
       output,
-      chunk.state === 'output-error' ? 'input-complete' : undefined,
+      chunk.state === 'output-error' ? 'error' : undefined,
     )
 
     // Step 2: Create/update the tool-result part
@@ -1690,11 +1690,22 @@ export class StreamProcessor {
     _index: number,
     toolCall: InternalToolCallState,
   ): void {
+    // Finalize the internal bookkeeping: the call's input arguments ARE
+    // complete regardless of whether execution later failed, so the call still
+    // counts as a completed tool call in getCompletedToolCalls()/getState().
     toolCall.state = 'input-complete'
 
     // Try final parse
     toolCall.parsedArguments = this.jsonParser.parse(toolCall.arguments)
 
+    // Don't downgrade the rendered part of a call that already reached the
+    // terminal 'error' state (e.g. an output-error TOOL_CALL_RESULT arrived
+    // without a preceding TOOL_CALL_END). The RUN_FINISHED / finalizeStream
+    // safety net must not clobber a failed call back to 'input-complete'.
+    if (this.isToolCallPartErrored(toolCall.id)) {
+      return
+    }
+
     // Update UIMessage
     this.messages = updateToolCallPart(this.messages, messageId, {
       id: toolCall.id,
@@ -1714,6 +1725,22 @@ export class StreamProcessor {
     )
   }
 
+  /**
+   * Whether the rendered tool-call part for the given id has reached the
+   * terminal 'error' state. Used to prevent the completion safety net from
+   * downgrading a failed call back to 'input-complete'.
+   */
+  private isToolCallPartErrored(toolCallId: string): boolean {
+    return this.messages.some((msg) =>
+      msg.parts.some(
+        (part) =>
+          part.type === 'tool-call' &&
+          part.id === toolCallId &&
+          part.state === 'error',
+      ),
+    )
+  }
+
   /**
    * Emit pending text update for a specific message.
    *

package/src/activities/generateVideo/index.ts

@@ -15,6 +15,7 @@ import type { DebugOption } from '../../logger/types'
 import type { VideoAdapter } from './adapter'
 import type {
   StreamChunk,
+  TokenUsage,
   VideoJobResult,
   VideoStatusResult,
   VideoUrlResult,
@@ -380,6 +381,7 @@ async function* runStreamingVideoGeneration<
             status: 'completed',
             url: urlResult.url,
             expiresAt: urlResult.expiresAt,
+            ...(urlResult.usage ? { usage: urlResult.usage } : {}),
           },
           timestamp: Date.now(),
         } as StreamChunk
@@ -454,6 +456,7 @@ export async function getVideoJobStatus<
   progress?: number
   url?: string
   error?: string
+  usage?: TokenUsage
 }> {
   const { adapter, jobId } = options
   const requestId = createId('video-status')
@@ -487,10 +490,19 @@ export async function getVideoJobStatus<
         duration: Date.now() - startTime,
         timestamp: Date.now(),
       })
+      if (urlResult.usage) {
+        aiEventClient.emit('video:usage', {
+          requestId,
+          model: adapter.model,
+          usage: urlResult.usage,
+          timestamp: Date.now(),
+        })
+      }
       return {
         status: statusResult.status,
         progress: statusResult.progress,
         url: urlResult.url,
+        ...(urlResult.usage ? { usage: urlResult.usage } : {}),
       }
     } catch (error) {
       const errorMessage =

package/src/index.ts

@@ -118,6 +118,20 @@ export type {
   ErrorInfo,
 } from './activities/chat/middleware/index'
 
+// Capability primitives + middleware builder
+export {
+  createCapability,
+  defineChatMiddleware,
+  createChatMiddleware,
+} from './activities/chat/middleware/index'
+export type {
+  Capability,
+  CapabilityHandle,
+  CapabilityContext,
+  CapabilityGetter,
+  CapabilityProvider,
+} from './activities/chat/middleware/index'
+
 // All types
 export * from './types'