@tanstack/ai 0.29.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/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'
 

package/src/middlewares/otel.ts

@@ -20,6 +20,7 @@ import type {
   ChatMiddleware,
   ChatMiddlewareContext,
 } from '../activities/chat/middleware/types'
+import type { TokenUsage } from '../types'
 
 /**
  * Scope (role) of an OTel span emitted by this middleware.
@@ -179,6 +180,59 @@ function firstNumber(...candidates: Array<unknown>): number | undefined {
   return undefined
 }
 
+/**
+ * Build the full set of `gen_ai.usage.*` span attributes from a `TokenUsage`.
+ *
+ * Beyond input/output tokens, this emits provider-reported cost, total tokens,
+ * cache and reasoning breakdowns, and duration-based billing — every field is
+ * guarded so spans stay clean when a provider doesn't report it. Cache and
+ * reasoning use the official GenAI semconv names; `gen_ai.usage.cost` and
+ * `gen_ai.usage.total_tokens` are de-facto extensions consumed by backends
+ * like PostHog (which otherwise re-derive cost from their own price tables,
+ * losing cache discounts and gateway markup). Fields with no semconv or
+ * de-facto convention (`costDetails`, `durationSeconds`) are
+ * TanStack-namespaced. Deliberately not emitted: `unitsBilled`,
+ * `providerUsageDetails`, and the per-modality token breakdowns — those are
+ * media-oriented; media-activity observability is tracked in #720.
+ */
+function usageAttributes(usage: TokenUsage): Record<string, AttributeValue> {
+  const attrs: Record<string, AttributeValue> = {
+    'gen_ai.usage.input_tokens': usage.promptTokens,
+    'gen_ai.usage.output_tokens': usage.completionTokens,
+  }
+  const optional: Array<[key: string, value: unknown]> = [
+    ['gen_ai.usage.total_tokens', usage.totalTokens],
+    ['gen_ai.usage.cost', usage.cost],
+    [
+      'gen_ai.usage.cache_read.input_tokens',
+      usage.promptTokensDetails?.cachedTokens,
+    ],
+    [
+      'gen_ai.usage.cache_creation.input_tokens',
+      usage.promptTokensDetails?.cacheWriteTokens,
+    ],
+    [
+      'gen_ai.usage.reasoning.output_tokens',
+      usage.completionTokensDetails?.reasoningTokens,
+    ],
+    ['tanstack.ai.usage.duration_seconds', usage.durationSeconds],
+    ['tanstack.ai.usage.upstream_cost', usage.costDetails?.upstreamCost],
+    [
+      'tanstack.ai.usage.upstream_input_cost',
+      usage.costDetails?.upstreamInputCost,
+    ],
+    [
+      'tanstack.ai.usage.upstream_output_cost',
+      usage.costDetails?.upstreamOutputCost,
+    ],
+  ]
+  for (const [key, value] of optional) {
+    const num = firstNumber(value)
+    if (num !== undefined) attrs[key] = num
+  }
+  return attrs
+}
+
 function errorMessage(err: unknown): string | undefined {
   if (err instanceof Error) return err.message
   if (typeof err === 'string') return err
@@ -524,10 +578,7 @@ export function otelMiddleware(options: OtelMiddlewareOptions): ChatMiddleware {
         // `runOnUsage` when `chunk.usage` is present, and `onUsage` is the
         // canonical place for the metric. Recording in both would double-count.
         if (chunk.usage) {
-          span.setAttributes({
-            'gen_ai.usage.input_tokens': chunk.usage.promptTokens,
-            'gen_ai.usage.output_tokens': chunk.usage.completionTokens,
-          })
+          span.setAttributes(usageAttributes(chunk.usage))
         }
 
         if (captureContent && state.assistantTextBuffer.length > 0) {
@@ -584,10 +635,7 @@ export function otelMiddleware(options: OtelMiddlewareOptions): ChatMiddleware {
         }
 
         const span = state.currentIterationSpan ?? state.rootSpan
-        span.setAttributes({
-          'gen_ai.usage.input_tokens': usage.promptTokens,
-          'gen_ai.usage.output_tokens': usage.completionTokens,
-        })
+        span.setAttributes(usageAttributes(usage))
       })
     },
 
@@ -905,10 +953,7 @@ export function otelMiddleware(options: OtelMiddlewareOptions): ChatMiddleware {
         }
 
         if (info.usage) {
-          state.rootSpan.setAttributes({
-            'gen_ai.usage.input_tokens': info.usage.promptTokens,
-            'gen_ai.usage.output_tokens': info.usage.completionTokens,
-          })
+          state.rootSpan.setAttributes(usageAttributes(info.usage))
         }
         if (info.finishReason) {
           state.rootSpan.setAttribute('gen_ai.response.finish_reasons', [