feat: wire AGUIMock into config, suite, exports, and CLI

Author: jpr5

src/cli.ts

@@ -6,6 +6,7 @@ import { createServer } from "./server.js";
 import { loadFixtureFile, loadFixturesFromDir, validateFixtures } from "./fixture-loader.js";
 import { Logger, type LogLevel } from "./logger.js";
 import { watchFixtures } from "./watcher.js";
+import { AGUIMock } from "./agui-mock.js";
 import type { ChaosConfig, RecordConfig } from "./types.js";
 
 const HELP = `
@@ -32,6 +33,9 @@ Options:
       --provider-azure <url>      Upstream URL for Azure OpenAI
       --provider-ollama <url>     Upstream URL for Ollama
       --provider-cohere <url>     Upstream URL for Cohere
+      --agui-record              Enable AG-UI recording (proxy unmatched AG-UI requests)
+      --agui-upstream <url>      Upstream AG-UI agent URL (used with --agui-record)
+      --agui-proxy-only          AG-UI proxy mode: forward without saving
       --chaos-drop <rate>   Probability (0-1) of dropping requests with 500
       --chaos-malformed <rate>  Probability (0-1) of returning malformed JSON
       --chaos-disconnect <rate> Probability (0-1) of destroying connection
@@ -60,6 +64,9 @@ const { values } = parseArgs({
     "provider-azure": { type: "string" },
     "provider-ollama": { type: "string" },
     "provider-cohere": { type: "string" },
+    "agui-record": { type: "boolean", default: false },
+    "agui-upstream": { type: "string" },
+    "agui-proxy-only": { type: "boolean", default: false },
     "chaos-drop": { type: "string" },
     "chaos-malformed": { type: "string" },
     "chaos-disconnect": { type: "string" },
@@ -168,6 +175,22 @@ if (values.record || values["proxy-only"]) {
   };
 }
 
+// Parse AG-UI record/proxy config from CLI flags
+let aguiMount: { path: string; handler: AGUIMock } | undefined;
+if (values["agui-record"] || values["agui-proxy-only"]) {
+  if (!values["agui-upstream"]) {
+    console.error("Error: --agui-record/--agui-proxy-only requires --agui-upstream");
+    process.exit(1);
+  }
+  const agui = new AGUIMock();
+  agui.enableRecording({
+    upstream: values["agui-upstream"],
+    fixturePath: resolve(fixturePath, "agui-recorded"),
+    proxyOnly: values["agui-proxy-only"],
+  });
+  aguiMount = { path: "/agui", handler: agui };
+}
+
 async function main() {
   // Load fixtures from path (detect file vs directory)
   let isDir: boolean;
@@ -219,17 +242,23 @@ async function main() {
     }
   }
 
-  const instance = await createServer(fixtures, {
-    port,
-    host,
-    latency,
-    chunkSize,
-    logLevel,
-    chaos,
-    metrics: values.metrics,
-    record,
-    strict: values.strict,
-  });
+  const mounts = aguiMount ? [aguiMount] : undefined;
+
+  const instance = await createServer(
+    fixtures,
+    {
+      port,
+      host,
+      latency,
+      chunkSize,
+      logLevel,
+      chaos,
+      metrics: values.metrics,
+      record,
+      strict: values.strict,
+    },
+    mounts,
+  );
 
   logger.info(`aimock server listening on ${instance.url}`);
 

src/config-loader.ts

@@ -3,9 +3,11 @@ import * as path from "node:path";
 import { LLMock } from "./llmock.js";
 import { MCPMock } from "./mcp-mock.js";
 import { A2AMock } from "./a2a-mock.js";
+import { AGUIMock } from "./agui-mock.js";
 import type { ChaosConfig, RecordConfig } from "./types.js";
 import type { MCPToolDefinition, MCPPromptDefinition } from "./mcp-types.js";
 import type { A2AAgentDefinition, A2APart, A2AArtifact, A2AStreamEvent } from "./a2a-types.js";
+import type { AGUIEvent } from "./agui-types.js";
 import { VectorMock } from "./vector-mock.js";
 import type { QueryResult } from "./vector-types.js";
 import { Logger } from "./logger.js";
@@ -56,6 +58,18 @@ export interface A2AConfig {
   agents?: A2AConfigAgent[];
 }
 
+export interface AGUIConfigFixture {
+  match: { message?: string; toolName?: string; stateKey?: string };
+  text?: string; // shorthand: uses buildTextResponse
+  events?: AGUIEvent[]; // raw events
+  delayMs?: number;
+}
+
+export interface AGUIConfig {
+  path?: string; // mount path, default "/agui"
+  fixtures?: AGUIConfigFixture[];
+}
+
 export interface VectorConfigCollection {
   name: string;
   dimension: number;
@@ -80,6 +94,7 @@ export interface AimockConfig {
   };
   mcp?: MCPConfig;
   a2a?: A2AConfig;
+  agui?: AGUIConfig;
   vector?: VectorConfig;
   services?: { search?: boolean; rerank?: boolean; moderate?: boolean };
   metrics?: boolean;
@@ -198,6 +213,38 @@ export async function startFromConfig(
     logger.info(`A2AMock mounted at ${a2aPath}`);
   }
 
+  // AG-UI
+  if (config.agui) {
+    const aguiConfig = config.agui;
+    const agui = new AGUIMock();
+
+    if (aguiConfig.fixtures) {
+      for (const f of aguiConfig.fixtures) {
+        if (f.text) {
+          agui.onMessage(f.match.message ?? /.*/, f.text, { delayMs: f.delayMs });
+        } else if (f.events) {
+          agui.addFixture({
+            match: {
+              message: f.match.message,
+              toolName: f.match.toolName,
+              stateKey: f.match.stateKey,
+            },
+            events: f.events,
+            delayMs: f.delayMs,
+          });
+        } else {
+          logger.warn(
+            `AG-UI fixture has neither text nor events — it will be skipped (match: ${JSON.stringify(f.match)})`,
+          );
+        }
+      }
+    }
+
+    const aguiPath = aguiConfig.path ?? "/agui";
+    llmock.mount(aguiPath, agui);
+    logger.info(`AGUIMock mounted at ${aguiPath}`);
+  }
+
   // Vector
   if (config.vector) {
     const vectorConfig = config.vector;

src/index.ts

@@ -156,6 +156,55 @@ export type {
   A2ATaskState,
 } from "./a2a-types.js";
 
+// AG-UI
+export { AGUIMock } from "./agui-mock.js";
+export { proxyAndRecordAGUI } from "./agui-recorder.js";
+export type {
+  AGUIMockOptions,
+  AGUIRunAgentInput,
+  AGUIMessage,
+  AGUIToolDefinition,
+  AGUIToolCall,
+  AGUIEvent,
+  AGUIEventType,
+  AGUIFixture,
+  AGUIFixtureMatch,
+  AGUIRecordConfig,
+  // Key individual event types
+  AGUIRunStartedEvent,
+  AGUIRunFinishedEvent,
+  AGUIRunErrorEvent,
+  AGUITextMessageStartEvent,
+  AGUITextMessageContentEvent,
+  AGUITextMessageEndEvent,
+  AGUITextMessageChunkEvent,
+  AGUIToolCallStartEvent,
+  AGUIToolCallArgsEvent,
+  AGUIToolCallEndEvent,
+  AGUIToolCallResultEvent,
+  AGUIStateSnapshotEvent,
+  AGUIStateDeltaEvent,
+  AGUIMessagesSnapshotEvent,
+  AGUIActivitySnapshotEvent,
+  AGUIActivityDeltaEvent,
+} from "./agui-types.js";
+export {
+  buildTextResponse as buildAGUITextResponse,
+  buildTextChunkResponse as buildAGUITextChunkResponse,
+  buildToolCallResponse as buildAGUIToolCallResponse,
+  buildStateUpdate as buildAGUIStateUpdate,
+  buildStateDelta as buildAGUIStateDelta,
+  buildMessagesSnapshot as buildAGUIMessagesSnapshot,
+  buildReasoningResponse as buildAGUIReasoningResponse,
+  buildActivityResponse as buildAGUIActivityResponse,
+  buildErrorResponse as buildAGUIErrorResponse,
+  buildStepWithText as buildAGUIStepWithText,
+  buildCompositeResponse as buildAGUICompositeResponse,
+  extractLastUserMessage as extractAGUILastUserMessage,
+  findFixture as findAGUIFixture,
+  writeAGUIEventStream,
+} from "./agui-handler.js";
+
 // JSON-RPC
 export { createJsonRpcDispatcher } from "./jsonrpc.js";
 export type { JsonRpcResponse, MethodHandler, JsonRpcDispatcherOptions } from "./jsonrpc.js";

src/suite.ts

@@ -2,23 +2,27 @@ import { LLMock } from "./llmock.js";
 import { MCPMock } from "./mcp-mock.js";
 import { A2AMock } from "./a2a-mock.js";
 import { VectorMock } from "./vector-mock.js";
+import { AGUIMock } from "./agui-mock.js";
 import type { MockServerOptions } from "./types.js";
 import type { MCPMockOptions } from "./mcp-types.js";
 import type { A2AMockOptions } from "./a2a-types.js";
 import type { VectorMockOptions } from "./vector-types.js";
+import type { AGUIMockOptions } from "./agui-types.js";
 
 export interface MockSuiteOptions {
   llm?: MockServerOptions;
   mcp?: MCPMockOptions;
   a2a?: A2AMockOptions;
   vector?: VectorMockOptions;
+  agui?: AGUIMockOptions;
 }
 
 export interface MockSuite {
   llm: LLMock;
   mcp?: MCPMock;
   a2a?: A2AMock;
   vector?: VectorMock;
+  agui?: AGUIMock;
   start(): Promise<void>;
   stop(): Promise<void>;
   reset(): void;
@@ -29,6 +33,7 @@ export async function createMockSuite(options: MockSuiteOptions = {}): Promise<M
   let mcp: MCPMock | undefined;
   let a2a: A2AMock | undefined;
   let vector: VectorMock | undefined;
+  let agui: AGUIMock | undefined;
 
   if (options.mcp) {
     mcp = new MCPMock(options.mcp);
@@ -45,11 +50,17 @@ export async function createMockSuite(options: MockSuiteOptions = {}): Promise<M
     llm.mount("/vector", vector);
   }
 
+  if (options.agui) {
+    agui = new AGUIMock(options.agui);
+    llm.mount("/agui", agui);
+  }
+
   return {
     llm,
     mcp,
     a2a,
     vector,
+    agui,
     async start() {
       await llm.start();
     },
@@ -61,6 +72,7 @@ export async function createMockSuite(options: MockSuiteOptions = {}): Promise<M
       if (mcp) mcp.reset();
       if (a2a) a2a.reset();
       if (vector) vector.reset();
+      if (agui) agui.reset();
     },
   };
 }