From 863f7a02ec17bc3cd44247a310d136477e341d01 Mon Sep 17 00:00:00 2001 From: Aaron Wong <6979793+zzwong@users.noreply.github.com> Date: Sat, 12 Sep 2026 22:09:33 -0500 Subject: [PATCH] fix(llm): let the task schema disambiguate JSON candidates in model prose Structured output already recovers a JSON object wrapped in prose, but only when exactly one syntactically valid top-level object is present. A model that drafts its answer inside a reasoning preamble emits a second valid object, and even the sentence "return JSON, not {}" is enough, so recovery gave up and a validation attempt was spent on text that could never parse. Extraction now returns every balanced candidate and the task's own decoder chooses among them: exactly one passing the schema is accepted, zero or several keep the existing error. A candidate that fails DisallowUnknownFields, or carries the wrong thread or schema version, was never a possible answer, so this changes which bytes reach the decoder without changing what the decoder accepts. --- internal/llm/adapter.go | 40 +++++++++++++++------ internal/llm/adapter_test.go | 70 ++++++++++++++++++++++++++++++++++-- internal/llm/extract.go | 27 +++++--------- internal/llm/extract_test.go | 52 ++++++++++++++------------- 4 files changed, 133 insertions(+), 56 deletions(-) diff --git a/internal/llm/adapter.go b/internal/llm/adapter.go index 0395d229..4bd244b8 100644 --- a/internal/llm/adapter.go +++ b/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 diff --git a/internal/llm/adapter_test.go b/internal/llm/adapter_test.go index 8e9d0fed..4842302f 100644 --- a/internal/llm/adapter_test.go +++ b/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(`{"ok":false}{"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(`{"ok":false}{"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) { diff --git a/internal/llm/extract.go b/internal/llm/extract.go index 1abe135d..4b12559b 100644 --- a/internal/llm/extract.go +++ b/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 diff --git a/internal/llm/extract_test.go b/internal/llm/extract_test.go index afb28266..f2a6044c 100644 --- a/internal/llm/extract_test.go +++ b/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", `{"a":0}{"a":1}`, []string{`{"a":0}`, `{"a":1}`}}, + {"empty braces in preamble", `not {}.{"a":1}`, []string{`{}`, `{"a":1}`}}, + {"zero objects", `no json here`, nil}, + {"unbalanced brace", `broken {"a":1`, nil}, + {"truncated after preamble", `x{"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]) + } } }) }