fix: harden AG-UI recorder, mock, and handler error handling

src/agui-mock.ts

@@ -19,7 +19,7 @@ import {
 } from "./agui-handler.js";
 import { flattenHeaders, readBody } from "./helpers.js";
 import { proxyAndRecordAGUI } from "./agui-recorder.js";
-import { Logger } from "./logger.js";
+import { Logger, type LogLevel } from "./logger.js";
 
 export class AGUIMock implements Mountable {
   private fixtures: AGUIFixture[] = [];
@@ -33,7 +33,7 @@ export class AGUIMock implements Mountable {
 
   constructor(options?: AGUIMockOptions) {
     this.options = options ?? {};
-    this.logger = new Logger("silent");
+    this.logger = new Logger((options?.logLevel as LogLevel) ?? "warn");
   }
 
   // ---- Fluent registration API ----
@@ -138,7 +138,16 @@ export class AGUIMock implements Mountable {
       this.registry.incrementCounter("aimock_agui_requests_total", { method: "POST" });
     }
 
-    const body = await readBody(req);
+    let body: string;
+    try {
+      body = await readBody(req);
+    } catch (err) {
+      res.writeHead(400, { "Content-Type": "application/json" });
+      const detail = err instanceof Error ? err.message : "body read failed";
+      res.end(JSON.stringify({ error: `Failed to read request body: ${detail}` }));
+      this.journalRequest(req, pathname, 400);
+      return true;
+    }
 
     let input: AGUIRunAgentInput;
     try {

src/agui-recorder.ts

@@ -115,23 +115,35 @@ function teeUpstreamStream(
       (upstreamRes) => {
         const upstreamStatus = upstreamRes.statusCode ?? 200;
 
-        // Set SSE headers on the client response
+        // Set appropriate headers on the client response
         if (!clientRes.headersSent) {
-          clientRes.writeHead(upstreamStatus, {
-            "Content-Type": "text/event-stream",
-            "Cache-Control": "no-cache",
-            Connection: "keep-alive",
-          });
+          if (upstreamStatus >= 200 && upstreamStatus < 300) {
+            clientRes.writeHead(upstreamStatus, {
+              "Content-Type": "text/event-stream",
+              "Cache-Control": "no-cache",
+              Connection: "keep-alive",
+            });
+          } else {
+            const ct = upstreamRes.headers["content-type"] || "application/json";
+            clientRes.writeHead(upstreamStatus, { "Content-Type": ct });
+          }
         }
 
         const chunks: Buffer[] = [];
+        let clientWriteFailed = false;
 
         upstreamRes.on("data", (chunk: Buffer) => {
           // Relay to client in real time
           try {
             clientRes.write(chunk);
-          } catch {
-            // Client connection may have closed — continue buffering for recording
+          } catch (err) {
+            if (!clientWriteFailed) {
+              clientWriteFailed = true;
+              logger?.warn(
+                "Client write failed during proxy relay:",
+                err instanceof Error ? err.message : String(err),
+              );
+            }
           }
           // Buffer for fixture construction
           chunks.push(chunk);
@@ -247,8 +259,10 @@ function parseSSEEvents(text: string, logger?: Logger): AGUIEvent[] {
         try {
           const parsed = JSON.parse(payload) as AGUIEvent;
           events.push(parsed);
-        } catch {
-          logger?.warn(`Skipping unparseable SSE data line: ${payload.slice(0, 200)}`);
+        } catch (err) {
+          const msg = err instanceof Error ? err.message : String(err);
+          if (logger) logger.warn(`Skipping unparseable SSE data line: ${payload.slice(0, 200)}`);
+          else console.warn(`Skipping unparseable SSE data line: ${msg}`);
         }
       }
     }

src/agui-types.ts

@@ -399,6 +399,7 @@ export interface AGUIFixture {
 export interface AGUIMockOptions {
   port?: number;
   host?: string;
+  logLevel?: string;
 }
 
 export interface AGUIRecordConfig {

src/cli.ts

@@ -97,8 +97,8 @@ const watchMode = values.watch!;
 const validateOnLoad = values["validate-on-load"]!;
 const logLevelStr = values["log-level"]!;
 
-if (!["silent", "info", "debug"].includes(logLevelStr)) {
-  console.error(`Invalid log-level: ${logLevelStr} (must be silent, info, or debug)`);
+if (!["silent", "warn", "info", "debug"].includes(logLevelStr)) {
+  console.error(`Invalid log-level: ${logLevelStr} (must be silent, warn, info, or debug)`);
   process.exit(1);
 }
 const logLevel = logLevelStr as LogLevel;

src/logger.ts

@@ -1,9 +1,10 @@
-export type LogLevel = "silent" | "info" | "debug";
+export type LogLevel = "silent" | "warn" | "info" | "debug";
 
 const LEVELS: Record<LogLevel, number> = {
   silent: 0,
-  info: 1,
-  debug: 2,
+  warn: 1,
+  info: 2,
+  debug: 3,
 };
 
 export class Logger {
@@ -26,10 +27,14 @@ export class Logger {
   }
 
   warn(...args: unknown[]): void {
-    console.warn("[aimock]", ...args);
+    if (this.level >= LEVELS.warn) {
+      console.warn("[aimock]", ...args);
+    }
   }
 
   error(...args: unknown[]): void {
-    console.error("[aimock]", ...args);
+    if (this.level >= LEVELS.warn) {
+      console.error("[aimock]", ...args);
+    }
   }
 }

src/types.ts

@@ -396,7 +396,7 @@ export interface MockServerOptions {
   latency?: number;
   chunkSize?: number;
   /** Log verbosity. CLI default is "info"; programmatic default (when omitted) is "silent". */
-  logLevel?: "silent" | "info" | "debug";
+  logLevel?: "silent" | "warn" | "info" | "debug";
   chaos?: ChaosConfig;
   /** Enable Prometheus-compatible /metrics endpoint. */
   metrics?: boolean;