@tanstack/ai/src/activities/chat/middleware/types.ts

Type-safe TypeScript AI SDK for streaming chat, tool calling, agents, structured outputs, and multimodal generation.

import type {
  JSONSchema,
  ModelMessage,
  StreamChunk,
  TokenUsage,
  Tool,
  ToolCall,
} from '../../../types'
import type { SystemPrompt } from '../../../system-prompts'
import type {
  Capability,
  CapabilityHandle,
  CapabilityRegistry,
} from './capabilities'

// ===========================
// Middleware Context
// ===========================

/**
 * Phase of the chat middleware lifecycle.
 * - 'init': Initial config transform before the chat engine starts
 * - 'beforeModel': Before each adapter chatStream call (per agent iteration)
 * - 'modelStream': During model streaming
 * - 'beforeTools': Before tool execution phase
 * - 'afterTools': After tool execution phase
 * - 'structuredOutput': During the final structured-output adapter call (set
 *   for chunks from adapter.structuredOutputStream or the synthesized fallback)
 */
export type ChatMiddlewarePhase =
  | 'init'
  | 'beforeModel'
  | 'modelStream'
  | 'beforeTools'
  | 'afterTools'
  | 'structuredOutput'

/**
 * Stable context object passed to all middleware hooks.
 * Created once per chat() invocation and shared across all hooks.
 */
export interface ChatMiddlewareContext<TContext = unknown> {
  /** Unique identifier for this chat request */
  requestId: string
  /** Unique identifier for this stream */
  streamId: string
  /** AG-UI run identifier for correlating client and server events */
  runId: string
  /**
   * AG-UI thread identifier — a stable per-conversation ID used to
   * correlate client and server devtools events. Resolves to the
   * caller-provided `threadId` (or legacy `conversationId`), or an
   * auto-generated value when neither is supplied.
   */
  threadId: string
  /**
   * @deprecated Use `threadId` instead. Retained as an alias of
   * `threadId` so middleware written before the AG-UI rename keeps
   * working unchanged. Will be removed in a future major release.
   */
  conversationId?: string
  /** Current lifecycle phase */
  phase: ChatMiddlewarePhase
  /** Current agent loop iteration (0-indexed) */
  iteration: number
  /** Running count of chunks yielded so far */
  chunkIndex: number
  /** Abort signal from the chat request */
  signal?: AbortSignal
  /** Abort the chat run with a reason */
  abort: (reason?: string) => void
  /** Runtime context provided by chat() options */
  context: TContext
  /**
   * Defer a non-blocking side-effect promise.
   * Deferred promises do not block streaming and are awaited
   * after the terminal hook (onFinish/onAbort/onError).
   */
  defer: (promise: Promise<unknown>) => void

  // --- Provider / adapter info (immutable for the lifetime of the request) ---

  /** Provider name (e.g., 'openai', 'anthropic') */
  provider: string
  /** Model identifier (e.g., 'gpt-4o') */
  model: string
  /** Source of the chat invocation — always 'server' for server-side chat */
  source: 'client' | 'server'
  /** Whether the chat is streaming */
  streaming: boolean

  // --- Config-derived info (may update per-iteration via onConfig) ---

  /** System prompts configured for this chat */
  systemPrompts: Array<SystemPrompt>
  /** Names of configured tools, if any */
  toolNames?: Array<string>
  /** Flattened generation options (metadata) */
  options?: Record<string, unknown> | undefined
  /** Provider-specific model options */
  modelOptions?: Record<string, unknown> | undefined

  // --- Computed info ---

  /** Number of messages at the start of the request */
  messageCount: number
  /** Whether tools are configured */
  hasTools: boolean

  // --- Mutable per-iteration state ---

  /** Current assistant message ID (changes per iteration) */
  currentMessageId: string | null
  /** Accumulated text content for the current iteration */
  accumulatedContent: string

  // --- References ---

  /** Current messages array (read-only view) */
  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
}

// ===========================
// Config passed to onConfig
// ===========================

/**
 * Chat configuration that middleware can observe or transform.
 * This is a subset of the chat engine's effective configuration
 * that middleware is allowed to modify.
 */
export interface ChatMiddlewareConfig {
  messages: Array<ModelMessage>
  systemPrompts: Array<SystemPrompt>
  tools: Array<Tool>
  metadata?: Record<string, unknown> | undefined
  modelOptions?: Record<string, unknown> | undefined
}

/**
 * Config passed to onStructuredOutputConfig.
 *
 * Mirrors ChatMiddlewareConfig minus `tools` (the final structured-output call
 * is a single typed-response request, not an agentic loop — tools cannot be
 * forwarded to it), plus the `outputSchema` being sent to the provider.
 * Middleware may transform the schema (e.g., inject $defs, strip
 * vendor-incompatible keywords) by returning a partial that includes
 * `outputSchema`.
 */
export interface StructuredOutputMiddlewareConfig extends Omit<
  ChatMiddlewareConfig,
  'tools'
> {
  /** JSON Schema being sent to the provider for structured output. */
  outputSchema: JSONSchema
}

// ===========================
// Tool Call Hook Context
// ===========================

/**
 * Context provided to tool call hooks (onBeforeToolCall / onAfterToolCall).
 */
export interface ToolCallHookContext {
  /** The tool call being executed */
  toolCall: ToolCall
  /** The resolved tool definition, if found */
  tool: Tool | undefined
  /** Parsed arguments for the tool call */
  args: unknown
  /** Name of the tool */
  toolName: string
  /** ID of the tool call */
  toolCallId: string
}

/**
 * Decision returned from onBeforeToolCall.
 * - undefined/void: continue with normal execution
 * - { type: 'transformArgs', args }: replace args used for execution
 * - { type: 'skip', result }: skip execution, use provided result
 * - { type: 'abort', reason }: abort the entire chat run
 */
export type BeforeToolCallDecision =
  | void
  | undefined
  | null
  | { type: 'transformArgs'; args: unknown }
  | { type: 'skip'; result: unknown }
  | { type: 'abort'; reason?: string }

/**
 * Outcome information provided to onAfterToolCall.
 */
export interface AfterToolCallInfo {
  /** The tool call that was executed */
  toolCall: ToolCall
  /** The resolved tool definition */
  tool: Tool | undefined
  /** Name of the tool */
  toolName: string
  /** ID of the tool call */
  toolCallId: string
  /** Whether the execution succeeded */
  ok: boolean
  /** Duration of tool execution in milliseconds */
  duration: number
  /** The result (if ok) or error (if not ok) */
  result?: unknown
  error?: unknown
}

// ===========================
// Iteration Info
// ===========================

/**
 * Information passed to onIteration at the start of each agent loop iteration.
 */
export interface IterationInfo {
  /** 0-based iteration index */
  iteration: number
  /** The assistant message ID created for this iteration */
  messageId: string
}

// ===========================
// Tool Phase Complete Info
// ===========================

/**
 * Aggregate information passed to onToolPhaseComplete after all tool calls
 * in an iteration have been processed.
 */
export interface ToolPhaseCompleteInfo {
  /** Tool calls that were assigned to the assistant message */
  toolCalls: Array<ToolCall>
  /** Completed tool results */
  results: Array<{
    toolCallId: string
    toolName: string
    result: unknown
    duration?: number
  }>
  /** Tools that need user approval */
  needsApproval: Array<{
    toolCallId: string
    toolName: string
    input: unknown
    approvalId: string
  }>
  /** Tools that need client-side execution */
  needsClientExecution: Array<{
    toolCallId: string
    toolName: string
    input: unknown
  }>
}

// ===========================
// Usage Info
// ===========================

/**
 * Token usage statistics passed to the onUsage hook.
 * Extracted from the RUN_FINISHED chunk when usage data is present.
 *
 * Includes optional provider-reported `cost`/`costDetails` (see {@link TokenUsage}).
 * Kept as an interface extending `TokenUsage` to preserve declaration merging for
 * this publicly exported type.
 */
export interface UsageInfo extends TokenUsage {}

// ===========================
// Terminal Hook Info
// ===========================

/**
 * Information passed to onFinish.
 */
export interface FinishInfo {
  /** The finish reason from the last model response */
  finishReason: string | null
  /** Total duration of the chat run in milliseconds */
  duration: number
  /** Final accumulated text content */
  content: string
  /** Final usage totals, if available (optionally including provider-reported cost) */
  usage?: TokenUsage | undefined
}

/**
 * Information passed to onAbort.
 */
export interface AbortInfo {
  /** The reason for the abort, if provided */
  reason?: string
  /** Duration until abort in milliseconds */
  duration: number
}

/**
 * Information passed to onError.
 */
export interface ErrorInfo {
  /** The error that caused the failure */
  error: unknown
  /** Duration until error in milliseconds */
  duration: number
}

// ===========================
// Middleware Interface
// ===========================

/**
 * Chat middleware interface.
 *
 * All hooks are optional. Middleware is composed in array order:
 * - `onConfig`: config piped through middlewares in order (first transform influences later)
 * - `onChunk`: each output chunk is fed into the next middleware in order
 *
 * @example Logging middleware
 * ```ts
 * const loggingMiddleware: ChatMiddleware = {
 *   name: 'logging',
 *   onStart(ctx) { console.log('Chat started', ctx.requestId) },
 *   onChunk(ctx, chunk) { console.log('Chunk:', chunk.type) },
 *   onFinish(ctx, info) { console.log('Done:', info.duration, 'ms') },
 * }
 * ```
 *
 * @example Redaction middleware
 * ```ts
 * const redactionMiddleware: ChatMiddleware = {
 *   name: 'redaction',
 *   onChunk(ctx, chunk) {
 *     if (chunk.type === 'TEXT_MESSAGE_CONTENT') {
 *       return { ...chunk, delta: redact(chunk.delta) }
 *     }
 *   },
 * }
 * ```
 */
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.
   *
   * Return a partial config to merge with the current config, or void to pass through.
   * Only the fields you return are overwritten — everything else is preserved.
   */
  onConfig?: (
    ctx: ChatMiddlewareContext<TContext>,
    config: ChatMiddlewareConfig,
  ) =>
    | void
    | null
    | Partial<ChatMiddlewareConfig>
    | Promise<void | null | Partial<ChatMiddlewareConfig>>

  /**
   * Called at the start of the final structured-output call (when the chat
   * was invoked with outputSchema). Pipes through middleware in order, like
   * onConfig, but with access to the JSON Schema being sent to the provider.
   *
   * Return a partial to shallow-merge into the current config, or void to
   * pass through.
   *
   * Fires BEFORE onConfig at the structured-output boundary. onConfig also
   * re-fires at the same boundary with ctx.phase === 'structuredOutput',
   * receiving the post-onStructuredOutputConfig view of the config (minus
   * outputSchema). Use onConfig for general-purpose transforms that apply
   * to every adapter call; use this hook when you need to transform the
   * outputSchema or apply structured-output-specific behavior.
   */
  onStructuredOutputConfig?: (
    ctx: ChatMiddlewareContext<TContext>,
    config: StructuredOutputMiddlewareConfig,
  ) =>
    | void
    | null
    | Partial<StructuredOutputMiddlewareConfig>
    | Promise<void | null | Partial<StructuredOutputMiddlewareConfig>>

  /**
   * Called when the chat run starts (after initial onConfig).
   */
  onStart?: (ctx: ChatMiddlewareContext<TContext>) => void | Promise<void>

  /**
   * Called at the start of each agent loop iteration, after a new assistant message ID
   * is created. Use this to observe iteration boundaries.
   */
  onIteration?: (
    ctx: ChatMiddlewareContext<TContext>,
    info: IterationInfo,
  ) => void | Promise<void>

  /**
   * Called for every chunk yielded by chat().
   * Can observe, transform, expand, or drop chunks.
   *
   * @returns void (pass through), chunk (replace), chunk[] (expand), null (drop)
   */
  onChunk?: (
    ctx: ChatMiddlewareContext<TContext>,
    chunk: StreamChunk,
  ) =>
    | void
    | StreamChunk
    | Array<StreamChunk>
    | null
    | Promise<void | StreamChunk | Array<StreamChunk> | null>

  /**
   * Called before a tool is executed.
   * Can observe, transform args, skip execution, or abort the run.
   */
  onBeforeToolCall?: (
    ctx: ChatMiddlewareContext<TContext>,
    hookCtx: ToolCallHookContext,
  ) => BeforeToolCallDecision | Promise<BeforeToolCallDecision>

  /**
   * Called after a tool execution completes (success or failure).
   */
  onAfterToolCall?: (
    ctx: ChatMiddlewareContext<TContext>,
    info: AfterToolCallInfo,
  ) => void | Promise<void>

  /**
   * Called after all tool calls in an iteration have been processed.
   * Provides aggregate data about tool execution results, approvals, and client tools.
   */
  onToolPhaseComplete?: (
    ctx: ChatMiddlewareContext<TContext>,
    info: ToolPhaseCompleteInfo,
  ) => void | Promise<void>

  /**
   * Called when usage data is available from a RUN_FINISHED chunk.
   * Called once per model iteration that reports usage.
   */
  onUsage?: (
    ctx: ChatMiddlewareContext<TContext>,
    usage: UsageInfo,
  ) => void | Promise<void>

  /**
   * Called when the chat run completes normally.
   * Exactly one of onFinish/onAbort/onError will be called per run.
   */
  onFinish?: (
    ctx: ChatMiddlewareContext<TContext>,
    info: FinishInfo,
  ) => void | Promise<void>

  /**
   * Called when the chat run is aborted.
   * Exactly one of onFinish/onAbort/onError will be called per run.
   */
  onAbort?: (
    ctx: ChatMiddlewareContext<TContext>,
    info: AbortInfo,
  ) => void | Promise<void>

  /**
   * Called when the chat run encounters an unhandled error.
   * Exactly one of onFinish/onAbort/onError will be called per run.
   */
  onError?: (
    ctx: ChatMiddlewareContext<TContext>,
    info: ErrorInfo,
  ) => void | Promise<void>
}

/** A `ChatMiddleware` with a permissive context — for use as a constraint. */
export type AnyChatMiddleware = ChatMiddleware<any>