Custom Middleware

​

Core Concepts

• Strategy / Provider: Each middleware implements an IAgentMiddlewareStrategy, marked with @AgentMiddlewareStrategy('<provider>') and registered to AgentMiddlewareRegistry (plugin-sdk).
• Middleware Instance: The AgentMiddleware object returned by createMiddleware, containing state Schema, context Schema, optional tools, and lifecycle hooks.
• Node-based Integration: Connect WorkflowNodeTypeEnum.MIDDLEWARE nodes to Agents in the workflow graph. At runtime, they are loaded and executed in connection order via getAgentMiddlewares.
• Configuration & Metadata: TAgentMiddlewareMeta describes name, i18n labels, icon, description, and configuration Schema. The UI renders configuration panels accordingly.

​

Lifecycle Hooks & Capabilities

• stateSchema: Persistable middleware state (Zod object/optional/with defaults).
• contextSchema: Runtime-only readable context, not persisted.
• tools: Dynamically injected DynamicStructuredTool list.
• JumpToTarget: Return jumpTo in hooks to control jumps (e.g., model, tools, end).

​

StateSchema Deep Dive

​

What is StateSchema?

​

Zod Introduction

import { z } from 'zod'

// Define schema
const schema = z.object({
  fieldName: z.type().method()
})

// Validate data
const validData = schema.parse(unknownData)

​

StateSchema Example: SkillsMiddleware

readonly stateSchema = z.object({
  selectedSkillIds: z.array(z.string()).optional(),
  selectedSkillWorkspaceId: z.string().optional()
})

1. Runtime Dynamic Skill Selection: During Agent execution, dynamically select which skills to enable based on user input or context
2. Workspace Isolation: Support loading skills from different workspaces, enabling multi-tenancy and workspace isolation
3. State Persistence: Maintain user skill selection history for later queries and recovery

private resolveRuntimeSkillSelection(state: Record<string, unknown>): {
  skillIds: string[]
  workspaceId: string
  hasRuntimeSelection: boolean
}

1. Read selectedSkillIds and selectedSkillWorkspaceId from current state and initial state
2. Sanitize and normalize data using sanitizeSkillIds() and sanitizeWorkspaceId()
3. Prioritize current state values, fallback to initial state
4. Return resolution results and whether runtime selection exists

​

StateSchema Definition Best Practices

1. Use optional(): For optional fields, explicitly mark them as optional
   Copy

   fieldName: z.string().optional()
   
2. Type Matching: Ensure defined types match actual data types in use
3. Validation & Cleanup: Validate and clean state data before use
4. Documentation: Add detailed comments explaining each state field

​

Why StateSchema Matters

• Type Safety: Check state data types at compile time
• Runtime Validation: Validate state data completeness and correctness at execution
• Auto Transformation: Automatically transform and clean state data
• IDE Support: Better code completion and type hints

• Maintain middleware runtime configuration
• Share state across middleware chain
• Implement stateful complex logic
• Support conditional execution and dynamic behavior

​

Writing a Middleware

1. Define Strategy & Metadata

@Injectable()
@AgentMiddlewareStrategy('rateLimitMiddleware')
export class RateLimitMiddleware implements IAgentMiddlewareStrategy<RateLimitOptions> {
  readonly meta: TAgentMiddlewareMeta = {
    name: 'rateLimitMiddleware',
    label: { en_US: 'Rate Limit Middleware', zh_Hans: '限流中间件' },
    description: { en_US: 'Protect LLM calls with quotas', zh_Hans: '为模型调用增加配额保护' },
    configSchema: { /* JSON Schema for frontend form rendering */ }
  };

2. Implement createMiddleware

  createMiddleware(options: RateLimitOptions, ctx: IAgentMiddlewareContext): AgentMiddleware {
    const quota = options.quota ?? 100;
    return {
      name: 'rateLimitMiddleware',
      stateSchema: z.object({ used: z.number().default(0) }),
      beforeModel: async (state, runtime) => {
        if ((state.used ?? 0) >= quota) {
          return { jumpTo: 'end', messages: state.messages };
        }
        return { used: (state.used ?? 0) + 1 };
      },
      wrapToolCall: async (request, handler) => handler(request),
    };
  }
}

3. Register as Plugin Module

@XpertServerPlugin({
  imports: [CqrsModule],
  providers: [RateLimitMiddleware],
})
export class MyAgentMiddlewareModule {}

4. Runtime Integration

• During development/deployment, add the plugin to plugin environment variables.
• In the workflow editor, add middleware nodes to the Agent and configure order (execution order can be adjusted by arrangement).

​

Middleware Lifecycle

​

Initialization Phase

Constructor → Dependency Injection → createMiddleware() → Return AgentMiddleware

​

Execution Phase

Request with State → wrapModelCall() → [Custom Logic] → Model Call → Response

​

Related Type Definitions

interface IAgentMiddlewareStrategy<TOptions> {
  readonly meta: TAgentMiddlewareMeta
  readonly stateSchema: ZodSchema
  createMiddleware(options: TOptions, context: IAgentMiddlewareContext): Promise<AgentMiddleware>
}

interface AgentMiddleware {
  name: string
  stateSchema: ZodSchema
  tools?: Tool[]
  wrapModelCall?: (request: any, handler: (req: any) => Promise<any>) => Promise<any>
}

​

Built-in Examples

• SummarizationMiddleware
  Detects conversation length in beforeModel, triggers compression and replaces historical messages; supports triggering by token/message count or context ratio, and records compression to execution trace via WrapWorkflowNodeExecutionCommand.
• todoListMiddleware Injects write_todos tool, and appends system prompts in wrapModelCall to guide LLM in planning complex tasks. Tool returns Command to update Agent state.

​

Best Practices

• Implement only necessary hooks, keep them idempotent, avoid heavy blocking operations within hooks.
• Use stateSchema to strictly declare persistent data, preventing state interference between different middleware/Agents.
• In wrapModelCall/wrapToolCall, prioritize calling the passed handler to ensure default chain works, then add custom logic.
• For ratio-triggered logic, rely on model profile.maxInputTokens; fallback to absolute token limits when unavailable (see Summarization example).
• Following LangChain Middleware approach: decompose cross-cutting concerns like logging, auditing, caching, rate limiting, and gradual model switching into independent middleware, composing them by connection order.

​

Reference Resources

• Zod Documentation
• Agent Middleware Overview
• SkillsMiddleware Implementation