fix(llm): let the task schema disambiguate JSON candidates in model prose

internal/llm/adapter.go

@@ -267,26 +267,44 @@ func RunStructuredWithSessionResume[T any](ctx context.Context, adapter Adapter,
 	return StructuredResult[T]{Value: retryValue, Response: retryResponse, SessionID: retrySessionID, ValidationAttempts: attempts, AcceptedOutput: retryAcceptedOutput}, nil
 }
 
-// decodeStructuredAccepted strict-decodes data, then on failure recovers a
-// response that wraps exactly one balanced top-level JSON object in surrounding
-// prose by decoding the extracted object with the same schema decoder. When the
-// extracted object also fails the schema, that error is returned because it
-// describes the real schema violation; otherwise the strict error stands.
+// decodeStructuredAccepted strict-decodes data, then on failure decodes each
+// balanced top-level JSON object found in the surrounding prose with the same
+// schema decoder. Exactly one candidate passing the schema is accepted; zero or
+// several keep the strict error, except a lone failing candidate reports its
+// own schema error because that describes the real violation.
 func decodeStructuredAccepted[T any](decode Decoder[T], data []byte) (T, []byte, error) {
 	value, err := decode(data)
 	if err == nil {
 		return value, data, nil
 	}
 	var zero T
-	extracted, ok := extractSingleJSONObject(data)
-	if !ok || bytes.Equal(extracted, data) {
+	candidates := extractJSONObjects(data)
+	if len(candidates) == 0 || (len(candidates) == 1 && bytes.Equal(candidates[0], data)) {
 		return zero, nil, err
 	}
-	extractedValue, extractedErr := decode(extracted)
-	if extractedErr != nil {
-		return zero, nil, extractedErr
+	var (
+		accepted       T
+		acceptedOutput []byte
+		candidateErr   error
+		passed         int
+	)
+	for _, candidate := range candidates {
+		candidateValue, decodeErr := decode(candidate)
+		if decodeErr != nil {
+			candidateErr = decodeErr
+			continue
+		}
+		accepted, acceptedOutput = candidateValue, candidate
+		passed++
+	}
+	switch {
+	case passed == 1:
+		return accepted, acceptedOutput, nil
+	case len(candidates) == 1:
+		return zero, nil, candidateErr
+	default:
+		return zero, nil, err
 	}
-	return extractedValue, extracted, nil
 }
 
 // runOnceWithSession runs a single attempt and retries transient provider

internal/llm/adapter_test.go

@@ -488,8 +488,8 @@ func TestRunStructuredProseRecovery(t *testing.T) {
 			decodeCalls++
 			return "", errors.New("invalid")
 		})
-		if decodeCalls != 2 {
-			t.Fatalf("decode calls = %d, want strict decode only per attempt (no extracted candidate)", decodeCalls)
+		if decodeCalls != 6 {
+			t.Fatalf("decode calls = %d, want 1 strict + 2 candidate decodes per attempt", decodeCalls)
 		}
 		if !errors.Is(err, ErrStructuredOutputInvalidAfterRetry) {
 			t.Fatalf("RunStructured error = %v, want %v", err, ErrStructuredOutputInvalidAfterRetry)
@@ -538,6 +538,72 @@ func TestRunStructuredProseRecovery(t *testing.T) {
 			t.Fatalf("requests = %d, want exactly one retry", requests)
 		}
 	})
+
+	t.Run("recovers the schema-valid object when the preamble drafts another", func(t *testing.T) {
+		adapter := &FakeAdapter{}
+		adapter.Queue(FakeResult{Response: Response{StructuredOutput: []byte(`<think>{"ok":false}</think>{"ok":true}`)}})
+
+		result, err := RunStructuredWithSessionResume(context.Background(), adapter, "", Request{Prompt: "prompt"}, decodeProbe)
+		if err != nil {
+			t.Fatalf("RunStructured: %v", err)
+		}
+		if !result.Value.OK {
+			t.Fatalf("value = %#v, want recovered object", result.Value)
+		}
+		if requests := len(adapter.Requests()); requests != 1 {
+			t.Fatalf("requests = %d, want recovery without retry", requests)
+		}
+	})
+
+	t.Run("rejects two schema-valid objects", func(t *testing.T) {
+		adapter := &FakeAdapter{}
+		adapter.Queue(FakeResult{Response: Response{StructuredOutput: []byte(`{"ok":true} and {"ok":true}`)}})
+		adapter.Queue(FakeResult{Response: Response{StructuredOutput: []byte(`{"ok":true} and {"ok":true}`)}})
+
+		_, err := RunStructuredWithSessionResume(context.Background(), adapter, "", Request{Prompt: "prompt"}, decodeProbe)
+		if !errors.Is(err, ErrStructuredOutputInvalidAfterRetry) {
+			t.Fatalf("RunStructured error = %v, want %v", err, ErrStructuredOutputInvalidAfterRetry)
+		}
+		if requests := len(adapter.Requests()); requests != 2 {
+			t.Fatalf("requests = %d, want retry path preserved", requests)
+		}
+	})
+
+	t.Run("multiple failing candidates keep the strict error", func(t *testing.T) {
+		adapter := &FakeAdapter{}
+		adapter.Queue(FakeResult{Response: Response{StructuredOutput: []byte(`<think>{"ok":false}</think>{"ok":false}`)}})
+		adapter.Queue(FakeResult{Response: Response{StructuredOutput: []byte(`{"ok":true}`)}})
+
+		result, err := RunStructuredWithSessionResume(context.Background(), adapter, "", Request{Prompt: "prompt"}, decodeProbe)
+		if err != nil {
+			t.Fatalf("RunStructured: %v", err)
+		}
+		if !result.Value.OK {
+			t.Fatalf("value = %#v, want retry value", result.Value)
+		}
+		requests := adapter.Requests()
+		if len(requests) != 2 {
+			t.Fatalf("requests = %d, want one retry", len(requests))
+		}
+		if !strings.Contains(requests[1].Prompt, "invalid character '<'") {
+			t.Fatalf("retry prompt = %q, want strict decode error", requests[1].Prompt)
+		}
+	})
+
+	t.Run("bare object failing schema decodes once", func(t *testing.T) {
+		adapter := &FakeAdapter{}
+		adapter.Queue(FakeResult{Response: Response{StructuredOutput: []byte(`{"ok":false}`)}})
+		adapter.Queue(FakeResult{Response: Response{StructuredOutput: []byte(`{"ok":false}`)}})
+
+		decodeCalls := 0
+		_, _ = RunStructuredWithSessionResume(context.Background(), adapter, "", Request{Prompt: "prompt"}, func([]byte) (probe, error) {
+			decodeCalls++
+			return probe{}, errors.New("invalid")
+		})
+		if decodeCalls != 2 {
+			t.Fatalf("decode calls = %d, want one strict decode per attempt", decodeCalls)
+		}
+	})
 }
 
 func TestFakeAdapterQuotaAndResume(t *testing.T) {

internal/llm/extract.go

@@ -2,36 +2,25 @@ package llm
 
 import "encoding/json"
 
-// extractSingleJSONObject returns the single balanced top-level JSON object
-// found in data, if exactly one syntactically valid candidate exists. Zero or
-// multiple valid candidates return ok=false: ambiguous output must not be
-// recovered. Invalid spans are deliberately re-scanned one byte past their
+// extractJSONObjects returns every balanced, syntactically valid top-level JSON
+// object in data, in order. Invalid spans are re-scanned one byte past their
 // opening brace so prose braces wrapping a valid object cannot hide it; the
-// resulting worst case is quadratic in nesting depth, which is irrelevant at
-// LLM response sizes.
-func extractSingleJSONObject(data []byte) ([]byte, bool) {
-	var (
-		candidate []byte
-		found     bool
-	)
+// worst case is quadratic in nesting depth, which is irrelevant at LLM
+// response sizes.
+func extractJSONObjects(data []byte) [][]byte {
+	var candidates [][]byte
 	for i := 0; i < len(data); i++ {
 		if data[i] != '{' {
 			continue
 		}
 		span, end, balanced := scanBalancedObject(data, i)
 		if !balanced || !json.Valid(span) {
-			// Keep scanning inside invalid spans: prose braces wrapping a
-			// valid object must not hide it.
 			continue
 		}
-		if found {
-			return nil, false
-		}
-		candidate = span
-		found = true
+		candidates = append(candidates, span)
 		i = end
 	}
-	return candidate, found
+	return candidates
 }
 
 // scanBalancedObject scans a {...} span starting at data[start], tracking

internal/llm/extract_test.go

@@ -2,36 +2,40 @@ package llm
 
 import "testing"
 
-func TestExtractSingleJSONObject(t *testing.T) {
+func TestExtractJSONObjects(t *testing.T) {
 	cases := []struct {
-		name   string
-		input  string
-		want   string
-		wantOK bool
+		name  string
+		input string
+		want  []string
 	}{
-		{"bare object", `{"a":1}`, `{"a":1}`, true},
-		{"leading and trailing prose", `Sure! Here it is: {"a":1} Hope that helps.`, `{"a":1}`, true},
-		{"nested objects and arrays", `prose {"a":{"b":1},"c":[{"d":2}]} prose`, `{"a":{"b":1},"c":[{"d":2}]}`, true},
-		{"braces inside strings", `note {"msg":"use { and } freely"} end`, `{"msg":"use { and } freely"}`, true},
-		{"escaped quotes inside strings", `{"msg":"she said \"hi\" {ok}"}`, `{"msg":"she said \"hi\" {ok}"}`, true},
-		{"prose braces alongside one valid object", `Here is {the thing}: {"a":1}`, `{"a":1}`, true},
-		{"valid object nested in invalid prose braces", `before {note {"ok":true}} after`, `{"ok":true}`, true},
-		{"markdown fenced object", "```json\n{\"a\":1}\n```", `{"a":1}`, true},
-		{"single object inside top-level array", `[{"a":1}]`, `{"a":1}`, true},
-		{"multiple objects inside top-level array", `[{"a":1},{"a":2}]`, "", false},
-		{"zero objects", `no json here`, "", false},
-		{"two objects", `{"a":1} and {"a":2}`, "", false},
-		{"unbalanced brace", `broken {"a":1`, "", false},
-		{"empty input", ``, "", false},
+		{"bare object", `{"a":1}`, []string{`{"a":1}`}},
+		{"leading and trailing prose", `Sure! Here it is: {"a":1} Hope that helps.`, []string{`{"a":1}`}},
+		{"nested objects and arrays", `prose {"a":{"b":1},"c":[{"d":2}]} prose`, []string{`{"a":{"b":1},"c":[{"d":2}]}`}},
+		{"braces inside strings", `note {"msg":"use { and } freely"} end`, []string{`{"msg":"use { and } freely"}`}},
+		{"escaped quotes inside strings", `{"msg":"she said \"hi\" {ok}"}`, []string{`{"msg":"she said \"hi\" {ok}"}`}},
+		{"prose braces alongside one valid object", `Here is {the thing}: {"a":1}`, []string{`{"a":1}`}},
+		{"valid object nested in invalid prose braces", `before {note {"ok":true}} after`, []string{`{"ok":true}`}},
+		{"markdown fenced object", "```json\n{\"a\":1}\n```", []string{`{"a":1}`}},
+		{"single object inside top-level array", `[{"a":1}]`, []string{`{"a":1}`}},
+		{"multiple objects inside top-level array", `[{"a":1},{"a":2}]`, []string{`{"a":1}`, `{"a":2}`}},
+		{"two objects", `{"a":1} and {"a":2}`, []string{`{"a":1}`, `{"a":2}`}},
+		{"tag preamble drafting an object", `<think>{"a":0}</think>{"a":1}`, []string{`{"a":0}`, `{"a":1}`}},
+		{"empty braces in preamble", `<think>not {}.</think>{"a":1}`, []string{`{}`, `{"a":1}`}},
+		{"zero objects", `no json here`, nil},
+		{"unbalanced brace", `broken {"a":1`, nil},
+		{"truncated after preamble", `<think>x</think>{"a":`, nil},
+		{"empty input", ``, nil},
 	}
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
-			got, ok := extractSingleJSONObject([]byte(tc.input))
-			if ok != tc.wantOK {
-				t.Fatalf("extractSingleJSONObject(%q) ok = %v, want %v", tc.input, ok, tc.wantOK)
+			got := extractJSONObjects([]byte(tc.input))
+			if len(got) != len(tc.want) {
+				t.Fatalf("extractJSONObjects(%q) = %q, want %q", tc.input, got, tc.want)
 			}
-			if string(got) != tc.want {
-				t.Fatalf("extractSingleJSONObject(%q) = %q, want %q", tc.input, got, tc.want)
+			for i := range got {
+				if string(got[i]) != tc.want[i] {
+					t.Fatalf("extractJSONObjects(%q)[%d] = %q, want %q", tc.input, i, got[i], tc.want[i])
+				}
 			}
 		})
 	}