fix: enhance auditor's JSON parsing and error handling

Added robust error handling for auditor responses, including a retry mechanism for JSON parsing and a fallback to extract verdicts from partial outputs. Improved the extractJSON utility to handle cases where JSON objects may be emitted without braces.

lib/agent/nodes/auditor.ts

@@ -6,6 +6,19 @@ import { chatComplete } from "../llm";
 import { extractJSON, loadPolicy } from "../utils";
 import { ResearchState, AuditResult, appendReasoning } from "../state";
 
+const MAX_AUDITOR_ATTEMPTS = 3;
+
+const VERDICT_REGEX = /"verdict"\s*:\s*"(approved|rejected|needs_revision)"/;
+
+/**
+ * Format Zod errors for inclusion in the next prompt so the LLM can self-correct.
+ */
+function formatZodErrors(issues: { path: unknown[]; message: string }[]): string {
+  return issues
+    .map((i) => `  - ${i.path.map((p) => String(p)).join(".")}: ${i.message}`)
+    .join("\n");
+}
+
 function normalizeStringArray(value: unknown): string[] {
   if (!Array.isArray(value)) return [];
   return value.map((item) => {
@@ -20,6 +33,34 @@ function normalizeStringArray(value: unknown): string[] {
   });
 }
 
+function parseAuditResult(parsed: unknown) {
+  const base = parsed as Record<string, unknown>;
+  return AuditResult.safeParse({
+    ...base,
+    policyViolations: normalizeStringArray(base.policyViolations),
+    suggestions: normalizeStringArray(base.suggestions),
+    auditedAt: new Date().toISOString(),
+  });
+}
+
+/**
+ * Last resort when JSON and retries fail: recover verdict from a substring fragment.
+ */
+function auditFromVerdictRegex(rawThought: string): AuditResult | null {
+  const m = rawThought.match(VERDICT_REGEX);
+  if (!m) return null;
+  const verdict = m[1];
+  const result = AuditResult.safeParse({
+    verdict,
+    policyViolations: [
+      "Auditor output could not be parsed as full JSON; verdict was recovered from partial output.",
+    ],
+    suggestions: [],
+    auditedAt: new Date().toISOString(),
+  });
+  return result.success ? result.data : null;
+}
+
 /**
  * Compares the current plan against POLICY.md.
  * Sets `auditResult` with verdict + structured feedback.
@@ -61,7 +102,7 @@ Global rules:
 - Do not include markdown fences or any prose outside the JSON object.
   `.trim();
 
-  const userMessage = `
+  const userMessageBase = `
 POLICY:
 ${policy}
 
@@ -69,23 +110,52 @@ PLAN TO AUDIT:
 ${JSON.stringify(state.plan, null, 2)}
   `.trim();
 
-  const rawThought = await chatComplete(system, userMessage);
+  let lastRawThought = "";
+  let lastParseError: string | null = null;
+
+  for (let attempt = 1; attempt <= MAX_AUDITOR_ATTEMPTS; attempt++) {
+    const parseFeedback =
+      lastParseError &&
+      `\n\nYour previous response had errors. Fix them and return ONLY valid JSON:\n${lastParseError}`;
+
+    const rawThought = await chatComplete(
+      system,
+      userMessageBase + (parseFeedback ?? "")
+    );
+    lastRawThought = rawThought;
+
+    let parsed: unknown;
+    try {
+      parsed = extractJSON(rawThought);
+    } catch {
+      lastParseError = `Could not parse as JSON. Output started with: ${rawThought.slice(0, 200)}`;
+      continue;
+    }
+
+    const zodResult = parseAuditResult(parsed);
+    if (zodResult.success) {
+      return buildAuditorState(state, zodResult.data, rawThought);
+    }
 
-  let parsed: unknown;
-  try {
-    parsed = extractJSON(rawThought);
-  } catch {
-    throw new Error(`Auditor produced non-JSON output: ${rawThought.slice(0, 300)}`);
+    lastParseError = formatZodErrors(zodResult.error.issues);
   }
 
-  const base = parsed as Record<string, unknown>;
-  const auditResult = AuditResult.parse({
-    ...base,
-    policyViolations: normalizeStringArray(base.policyViolations),
-    suggestions: normalizeStringArray(base.suggestions),
-    auditedAt: new Date().toISOString(),
-  });
+  const regexAudit = auditFromVerdictRegex(lastRawThought);
+  if (regexAudit) {
+    return buildAuditorState(state, regexAudit, lastRawThought);
+  }
+
+  throw new Error(
+    `Auditor failed after ${MAX_AUDITOR_ATTEMPTS} attempts. The model did not return valid JSON matching the audit schema. ` +
+      `Try a larger model or lower sampling temperature. Last output (excerpt): ${lastRawThought.slice(0, 400)}`
+  );
+}
 
+function buildAuditorState(
+  state: ResearchState,
+  auditResult: AuditResult,
+  rawThought: string
+): Partial<ResearchState> {
   const isRejected = auditResult.verdict !== "approved";
 
   const rejectionFeedback: string | null = isRejected
@@ -121,4 +191,3 @@ ${JSON.stringify(state.plan, null, 2)}
     updatedAt: new Date().toISOString(),
   };
 }
-

lib/agent/utils/extract-json.ts

@@ -2,6 +2,28 @@
  * JSON extraction helper — handles models that add extra text
  */
 
+/**
+ * Small models often emit JSON object *contents* without `{` / `}`.
+ * Strip a trailing comma before wrapping so `{"a":1,}` is not produced.
+ */
+function tryParseBraceWrappedObjectBody(text: string): unknown | null {
+  const trimmed = text.trim();
+  let candidate = trimmed;
+  if (!candidate.startsWith('"')) {
+    const q = candidate.indexOf('"');
+    if (q === -1) return null;
+    candidate = candidate.slice(q);
+  }
+  if (!candidate.includes(":")) return null;
+  const inner = candidate.replace(/,\s*$/, "").trim();
+  if (!inner.startsWith('"')) return null;
+  try {
+    return JSON.parse(`{${inner}}`);
+  } catch {
+    return null;
+  }
+}
+
 /**
  * Extracts JSON from model output that may contain extra text,
  * markdown fences, or other non-JSON content.
@@ -37,6 +59,12 @@ export function extractJSON(text: string): unknown {
     }
   }
 
+  // Braceless object body, e.g. `"verdict": "needs_revision"`
+  const fromCleaned = tryParseBraceWrappedObjectBody(cleaned);
+  if (fromCleaned !== null) return fromCleaned;
+
+  const fromOriginal = tryParseBraceWrappedObjectBody(text);
+  if (fromOriginal !== null) return fromOriginal;
+
   throw new Error(`Could not extract valid JSON from: ${text.slice(0, 200)}...`);
 }
-