From c33d88f0ea4c4ef3151aaf4e7644fcb3cf8214dd Mon Sep 17 00:00:00 2001
From: sax3l <simon@siax.io>
Date: Mon, 2 Nov 2026 17:18:34 +0100
Subject: [PATCH] fix(workbench): tolerate flat or fenced JSON from a small
 local model

The decision-server model sometimes returns {field: value} or wraps the JSON
in a markdown fence. Accept both, while still requiring every value to be an
exact, unique substring of the source before it can be validated.
---
 src/document-workbench/structure-document.ts  | 45 ++++++++++++++-----
 .../structure-document.test.ts                | 14 ++++++
 2 files changed, 49 insertions(+), 10 deletions(-)

diff --git a/src/document-workbench/structure-document.ts b/src/document-workbench/structure-document.ts
index 6095a46..84ed35e 100644
--- a/src/document-workbench/structure-document.ts
+++ b/src/document-workbench/structure-document.ts
@@ -133,31 +133,54 @@ export async function structureDocument(input: {
   try {
     parsed = JSON.parse(content);
   } catch {
-    return { status: 'needs_review', issues: [{ field: '', code: 'invalid_model_json' }], model };
+    // Small models often wrap JSON in a markdown fence or add prose.
+    const fenced = /```(?:json)?\s*([\s\S]*?)```/i.exec(content);
+    let inner = fenced ? fenced[1]!.trim() : '';
+    if (!inner) {
+      const start = content.indexOf('{');
+      const end = content.lastIndexOf('}');
+      inner = start >= 0 && end > start ? content.slice(start, end + 1) : '';
+    }
+    try {
+      parsed = inner ? JSON.parse(inner) : undefined;
+    } catch {
+      parsed = undefined;
+    }
+    if (parsed === undefined) {
+      return { status: 'needs_review', issues: [{ field: '', code: 'invalid_model_json' }], model };
+    }
   }
   if (!parsed || typeof parsed !== 'object' || Array.isArray(parsed)) {
     return { status: 'needs_review', issues: [{ field: '', code: 'invalid_model_json' }], model };
   }
-  const candidate = parsed as { value?: unknown; evidence?: unknown };
-  const asObject =
-    candidate.value && typeof candidate.value === 'object' && !Array.isArray(candidate.value)
-      ? (candidate.value as Record<string, unknown>)
-      : undefined;
+  const candidate = parsed as Record<string, unknown>;
+  const fieldNames = new Set(snapshot.fields.map((field) => field.name));
+  // Accept the documented {value, evidence} envelope, a flat object of the
+  // requested fields, or {value: {...}} with extra prose keys.
+  let asObject: Record<string, unknown> | undefined;
+  if (candidate.value && typeof candidate.value === 'object' && !Array.isArray(candidate.value)) {
+    asObject = candidate.value as Record<string, unknown>;
+  } else if (
+    Object.keys(candidate).length > 0 &&
+    Object.keys(candidate).every((key) => fieldNames.has(key))
+  ) {
+    asObject = candidate;
+  }
   if (!asObject) return { status: 'needs_review', issues: [{ field: '', code: 'invalid_output' }], model };
 
   // 1. Trust model-supplied evidence only if it validates as-is.
-  if (Array.isArray(candidate.evidence)) {
+  if (Array.isArray((candidate as { evidence?: unknown }).evidence)) {
     const supplied = validateStructuredOutput({
       fields: snapshot.fields,
       sources: snapshot.sources,
       value: asObject,
-      evidence: candidate.evidence as FieldEvidence[],
+      evidence: (candidate as { evidence?: unknown }).evidence as FieldEvidence[],
     });
     if (supplied.status === 'mechanically_validated') {
       return {
         status: 'mechanically_validated',
         value: asObject,
-        evidence: candidate.evidence as FieldEvidence[],
+        evidence: (candidate as { evidence?: unknown }).evidence as FieldEvidence[],
         model,
       };
     }
@@ -169,7 +192,9 @@ export async function structureDocument(input: {
       fields: snapshot.fields,
       sources: snapshot.sources,
       value: asObject,
-      evidence: Array.isArray(candidate.evidence) ? (candidate.evidence as FieldEvidence[]) : [],
+      evidence: Array.isArray((candidate as { evidence?: unknown }).evidence)
+        ? ((candidate as { evidence?: unknown }).evidence as FieldEvidence[])
+        : [],
     });
     return { status: 'needs_review', issues: checked.issues, model };
   }
diff --git a/tests/document-workbench/structure-document.test.ts b/tests/document-workbench/structure-document.test.ts
index 66725e8..4a6480b 100644
--- a/tests/document-workbench/structure-document.test.ts
+++ b/tests/document-workbench/structure-document.test.ts
@@ -122,6 +122,20 @@ describe('structure document workflow', () => {
       { field: 'name', sourceId: 'page1', start: 6, end: 10, transform: 'exact' },
     ]);
   });
+  it('accepts a flat object or a fenced JSON object from a small local model', async () => {
+    const flat = await structureDocument({
+      ...base,
+      gateway: { complete: vi.fn(async () => ({ ...response, content: JSON.stringify({ name: 'Acme' }) })) },
+    });
+    expect(flat).toMatchObject({ status: 'mechanically_validated', value: { name: 'Acme' } });
+    const fenced = await structureDocument({
+      ...base,
+      gateway: {
+        complete: vi.fn(async () => ({ ...response, content: '```json\n{"value":{"name":"Acme"}}\n```' })),
+      },
+    });
+    expect(fenced).toMatchObject({ status: 'mechanically_validated', value: { name: 'Acme' } });
+  });
   it('refuses ambiguous or absent values when computing evidence itself', async () => {
     const twice = await structureDocument({
       ...base,
-- 
2.54.0