@tanstack/ai 0.29.0 → 0.32.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/generateImage/adapter.ts

@@ -1,4 +1,8 @@
-import type { ImageGenerationOptions, ImageGenerationResult } from '../../types'
+import type {
+  ImageGenerationOptions,
+  ImageGenerationResult,
+  ModelInputModalitiesByName,
+} from '../../types'
 
 /**
  * Resolve the size type for a model from the model-size map.
@@ -29,6 +33,8 @@ export interface ImageAdapterConfig {
  * - TProviderOptions: Base provider-specific options (already resolved)
  * - TModelProviderOptionsByName: Map from model name to its specific provider options
  * - TModelSizeByName: Map from model name to its supported sizes
+ * - TModelInputModalitiesByName: Map from model name to the non-text prompt
+ *   modalities it accepts (constrains the `prompt` part types at compile time)
  */
 export interface ImageAdapter<
   TModel extends string = string,
@@ -38,6 +44,8 @@ export interface ImageAdapter<
     string,
     string
   >,
+  TModelInputModalitiesByName extends ModelInputModalitiesByName =
+    ModelInputModalitiesByName,
 > {
   /** Discriminator for adapter kind - used by generate() to determine API shape */
   readonly kind: 'image'
@@ -53,6 +61,7 @@ export interface ImageAdapter<
     providerOptions: TProviderOptions
     modelProviderOptionsByName: TModelProviderOptionsByName
     modelSizeByName: TModelSizeByName
+    modelInputModalitiesByName: TModelInputModalitiesByName
   }
 
   /**
@@ -67,7 +76,7 @@ export interface ImageAdapter<
  * An ImageAdapter with any/unknown type parameters.
  * Useful as a constraint in generic functions and interfaces.
  */
-export type AnyImageAdapter = ImageAdapter<any, any, any, any>
+export type AnyImageAdapter = ImageAdapter<any, any, any, any, any>
 
 /**
  * Abstract base class for image generation adapters.
@@ -83,11 +92,14 @@ export abstract class BaseImageAdapter<
     string,
     string
   >,
+  TModelInputModalitiesByName extends ModelInputModalitiesByName =
+    ModelInputModalitiesByName,
 > implements ImageAdapter<
   TModel,
   TProviderOptions,
   TModelProviderOptionsByName,
-  TModelSizeByName
+  TModelSizeByName,
+  TModelInputModalitiesByName
 > {
   readonly kind = 'image' as const
   abstract readonly name: string
@@ -98,6 +110,7 @@ export abstract class BaseImageAdapter<
     providerOptions: TProviderOptions
     modelProviderOptionsByName: TModelProviderOptionsByName
     modelSizeByName: TModelSizeByName
+    modelInputModalitiesByName: TModelInputModalitiesByName
   }
 
   protected config: ImageAdapterConfig

package/src/activities/generateImage/index.ts

@@ -8,10 +8,16 @@
 import { aiEventClient } from '@tanstack/ai-event-client'
 import { streamGenerationResult } from '../stream-generation-result.js'
 import { resolveDebugOption } from '../../logger/resolve'
+import { resolveMediaPrompt } from '../../utilities/media-prompt'
 import type { InternalLogger } from '../../logger/internal-logger'
 import type { DebugOption } from '../../logger/types'
 import type { ImageAdapter } from './adapter'
-import type { ImageGenerationResult, StreamChunk } from '../../types'
+import type {
+  ImageGenerationResult,
+  MediaPrompt,
+  MediaPromptFor,
+  StreamChunk,
+} from '../../types'
 
 // ===========================
 // Activity Kind
@@ -55,6 +61,23 @@ export type ImageSizeForModel<TAdapter, TModel extends string> =
         : string
     : string
 
+/**
+ * Extract the prompt type a model accepts from an ImageAdapter via ~types.
+ * Adapters declare a per-model input-modality map; models in the map get a
+ * `prompt` narrowed to text + their supported part types (text-only models
+ * accept `string | Array<TextPart>`), so unsupported media parts fail at
+ * compile time. Adapters without a map fall back to the full MediaPrompt.
+ */
+export type ImagePromptForModel<TAdapter, TModel extends string> =
+  TAdapter extends ImageAdapter<any, any, any, any, infer ModsByName>
+    ? string extends keyof ModsByName
+      ? // No explicit map - accept the full union
+        MediaPrompt
+      : TModel extends keyof ModsByName
+        ? MediaPromptFor<ModsByName[TModel][number]>
+        : MediaPrompt
+    : MediaPrompt
+
 // ===========================
 // Activity Options Type
 // ===========================
@@ -72,8 +95,16 @@ export type ImageActivityOptions<
 > = {
   /** The image adapter to use (must be created with a model) */
   adapter: TAdapter & { kind: typeof kind }
-  /** Text description of the desired image(s) */
-  prompt: string
+  /**
+   * Description of the desired image(s). Either a plain string, or — for
+   * models that support image-conditioned generation — an ordered array of
+   * content parts interleaving text with image inputs (image-to-image,
+   * reference-guided, edit, multi-reference). Media parts may carry
+   * `metadata.role` (`'reference' | 'mask' | 'control' | 'character'`) to
+   * disambiguate intent. The accepted part types are narrowed per model via
+   * the adapter's input-modality map.
+   */
+  prompt: ImagePromptForModel<TAdapter, TAdapter['model']>
   /** Number of images to generate (default: 1) */
   numberOfImages?: number
   /** Image size in WIDTHxHEIGHT format (e.g., "1024x1024") */
@@ -203,13 +234,26 @@ async function runGenerateImage<
   const startTime = Date.now()
   const logger: InternalLogger = resolveDebugOption(options.debug)
 
+  // Devtools events carry the flattened prompt text plus media-part counts —
+  // the wire payload stays `prompt: string` regardless of the prompt shape.
+  const resolved = resolveMediaPrompt(rest.prompt)
+
   aiEventClient.emit('image:request:started', {
     requestId,
     provider: adapter.name,
     model,
-    prompt: rest.prompt,
+    prompt: resolved.text,
     numberOfImages: rest.numberOfImages,
     size: rest.size,
+    ...(resolved.images.length > 0 && {
+      imageInputCount: resolved.images.length,
+    }),
+    ...(resolved.videos.length > 0 && {
+      videoInputCount: resolved.videos.length,
+    }),
+    ...(resolved.audios.length > 0 && {
+      audioInputCount: resolved.audios.length,
+    }),
     modelOptions: rest.modelOptions,
     timestamp: startTime,
   })