Middleware System Overview

/**
 * Create typed middleware for extending agent behavior
 * @param config - Middleware configuration with hooks and schemas
 * @returns Typed middleware instance
 */
function createMiddleware<
  TSchema = any,
  TContextSchema = any,
  TTools extends readonly Tool[] = readonly Tool[]
>(
  config: MiddlewareConfig<TSchema, TContextSchema, TTools>
): AgentMiddleware<TSchema, TContextSchema, TTools>;

interface MiddlewareConfig<TSchema, TContextSchema, TTools> {
  /**
   * Middleware name for identification
   */
  name: string;

  /**
   * State schema (persisted across invocations)
   */
  stateSchema?: TSchema;

  /**
   * Context schema (read-only, not persisted)
   */
  contextSchema?: TContextSchema;

  /**
   * Additional tools provided by middleware
   */
  tools?: TTools;

  /**
   * Wrap tool execution
   */
  wrapToolCall?: WrapToolCallHook<TSchema, TContextSchema>;

  /**
   * Wrap model invocation
   */
  wrapModelCall?: WrapModelCallHook<TSchema, TContextSchema>;

  /**
   * Run before agent starts
   */
  beforeAgent?: BeforeAgentHook<TSchema, TContextSchema>;

  /**
   * Run before model call
   */
  beforeModel?: BeforeModelHook<TSchema, TContextSchema>;

  /**
   * Run after model call
   */
  afterModel?: AfterModelHook<TSchema, TContextSchema>;

  /**
   * Run after agent completes
   */
  afterAgent?: AfterAgentHook<TSchema, TContextSchema>;
}

/**
 * Wrap tool call execution
 */
type WrapToolCallHook<TSchema, TContext> = (
  request: ToolCallRequest,
  handler: ToolCallHandler<TSchema, TContext>,
  runtime: Runtime<TContext>
) => Promise<ToolResult>;

/**
 * Wrap model invocation
 */
type WrapModelCallHook<TSchema, TContext> = (
  state: State<TSchema>,
  handler: WrapModelCallHandler<TSchema, TContext>,
  runtime: Runtime<TContext>
) => Promise<State<TSchema>>;

/**
 * Run before agent starts
 */
type BeforeAgentHook<TSchema, TContext> = (
  state: State<TSchema>,
  runtime: Runtime<TContext>
) => Promise<State<TSchema>> | State<TSchema>;

/**
 * Run before model call
 */
type BeforeModelHook<TSchema, TContext> = (
  state: State<TSchema>,
  runtime: Runtime<TContext>
) => Promise<State<TSchema>> | State<TSchema>;

/**
 * Run after model call
 */
type AfterModelHook<TSchema, TContext> = (
  state: State<TSchema>,
  runtime: Runtime<TContext>
) => Promise<State<TSchema>> | State<TSchema>;

/**
 * Run after agent completes
 */
type AfterAgentHook<TSchema, TContext> = (
  state: State<TSchema>,
  runtime: Runtime<TContext>
) => Promise<State<TSchema>> | State<TSchema>;

/**
 * Tool call request information
 */
interface ToolCallRequest {
  toolName: string;
  args: Record<string, any>;
  toolCallId: string;
}

/**
 * Tool call handler function
 */
type ToolCallHandler<TSchema, TContext> = (
  request: ToolCallRequest
) => Promise<ToolResult>;

/**
 * Model call handler function
 */
type WrapModelCallHandler<TSchema, TContext> = (
  state: State<TSchema>
) => Promise<State<TSchema>>;

/**
 * Tool result
 */
interface ToolResult {
  content: string;
  error?: string;
}

/**
 * Base middleware interface
 */
interface AgentMiddleware<
  TSchema = any,
  TContextSchema = any,
  TFullContext = any,
  TTools extends readonly Tool[] = readonly Tool[]
> {
  [MIDDLEWARE_BRAND]: true;
  name: string;
  stateSchema?: TSchema;
  contextSchema?: TContextSchema;
  tools?: TTools;
  hooks: {
    wrapToolCall?: WrapToolCallHook<TSchema, TFullContext>;
    wrapModelCall?: WrapModelCallHook<TSchema, TFullContext>;
    beforeAgent?: BeforeAgentHook<TSchema, TFullContext>;
    beforeModel?: BeforeModelHook<TSchema, TFullContext>;
    afterModel?: AfterModelHook<TSchema, TFullContext>;
    afterAgent?: AfterAgentHook<TSchema, TFullContext>;
  };
}

/**
 * Unique symbol for middleware branding
 */
const MIDDLEWARE_BRAND: unique symbol;

/**
 * Middleware return type (state update)
 */
type MiddlewareResult<TState> = TState | Partial<TState> | void;

import {
  type InferMiddlewareState,
  type InferMiddlewareStates,
  type InferMiddlewareContext,
  type InferMiddlewareContexts,
  type InferMiddlewareTools,
} from "langchain";

// Infer single middleware state
type State = InferMiddlewareState<typeof middleware>;

// Infer combined middleware states
type CombinedStates = InferMiddlewareStates<typeof middlewareArray>;

// Infer middleware context
type Context = InferMiddlewareContext<typeof middleware>;

// Infer combined middleware contexts
type CombinedContexts = InferMiddlewareContexts<typeof middlewareArray>;

// Infer tools from middleware
type Tools = InferMiddlewareTools<typeof middlewareArray>;