From 8bd2509bc259e3b554bc459dfe7debb16a26aee9 Mon Sep 17 00:00:00 2001 From: prosolis <5590409+prosolis@users.noreply.github.com> Date: Fri, 26 Jun 2026 14:16:07 -0700 Subject: [PATCH] Checkpoint: raise output cap to 4096 + salvage truncated JSON The grammar checkpoint capped num_predict at 1024, but qwen3.5:9b ignores the prompt's "up to 5 issues" and emits ~17-20 suggestions (~2000 tokens) on a 300+ word doc. The output hit the cap mid-array (done_reason=length), the JSON never closed, and ParseCheckpoint found no parseable object -> a 502 in ~20s on every long doc (1024 tok @ ~50 tok/s, not a network timeout). The repeated failures made the writing-assist helper look permanently asleep. Fix: - Raise the cap to checkpointMaxTokens=4096. It is a ceiling, not a target -- the model stops at its JSON close, so shorter docs are unaffected; only genuinely long outputs use the headroom. - Make ParseCheckpoint salvage the completed {...} suggestion objects from a truncated array (refactor extractJSONObject onto a shared firstBalancedObject scanner), so an over-long doc degrades to partial feedback instead of a hard 502. Verified live on millenia: 'The Missing Key' (382 words) now returns 200 with 17 suggestions in ~38s, previously 502 every time. Claude-Session: https://claude.ai/code/session_016Yr6jELuRc7hyzYLccQKZd --- internal/llm/checkpoint.go | 91 ++++++++++++++++++++++++++++----- internal/llm/checkpoint_test.go | 15 ++++++ 2 files changed, 93 insertions(+), 13 deletions(-) diff --git a/internal/llm/checkpoint.go b/internal/llm/checkpoint.go index a769e4d..f1a622c 100644 --- a/internal/llm/checkpoint.go +++ b/internal/llm/checkpoint.go @@ -14,6 +14,16 @@ import ( // that protects the inference endpoint from rapid repeat checks. const CheckpointInterval = 30 * time.Second +// checkpointMaxTokens caps the checkpoint's generated output. Each suggestion +// echoes its original sentence, a replacement, and a friendly explanation, so +// the full JSON for the prompt's "up to 5 issues" runs ~2,000 tokens on a long +// document (measured ~2,007 for a 380-word doc). The old 1,024 cap truncated +// the JSON mid-array — done_reason "length" — leaving an object that never +// closed, so ParseCheckpoint found no parseable JSON and the pass 502'd on +// every long doc. This is a ceiling, not a target: the model stops at its JSON +// close well before it, so the headroom costs nothing on shorter docs. +const checkpointMaxTokens = 4096 + // RawSuggestion is one item as the model emits it. Positions are resolved // server-side from Original; the frontend re-anchors by string at render time. type RawSuggestion struct { @@ -34,7 +44,7 @@ type checkpointResponse struct { func RunCheckpoint(ctx context.Context, client LLMClient, contentText, tone string) ([]RawSuggestion, error) { raw, err := client.Complete(ctx, CompletionRequest{ Messages: CheckpointMessages(TruncateDoc(contentText), tone), - MaxTokens: 1024, + MaxTokens: checkpointMaxTokens, Temperature: 0.3, RepetitionPenalty: 1.15, TopP: 0.9, @@ -48,20 +58,32 @@ func RunCheckpoint(ctx context.Context, client LLMClient, contentText, tone stri // ParseCheckpoint extracts the suggestions array from a model response. Smaller // models sometimes wrap JSON in prose or markdown fences despite instructions, -// so we salvage the outermost {...} object before decoding. +// so we salvage the outermost {...} object before decoding. When the model +// truncates mid-array (it ran out of output budget) the outer object never +// closes and won't parse — we then recover the suggestion objects that *did* +// complete, so a long document still gets feedback instead of a hard failure. func ParseCheckpoint(raw string) ([]RawSuggestion, error) { - jsonText := extractJSONObject(raw) - if jsonText == "" { - return nil, fmt.Errorf("checkpoint: no JSON object in model output: %q", truncateForError(raw)) + var suggestions []RawSuggestion + parsed := false + if jsonText := extractJSONObject(raw); jsonText != "" { + var resp checkpointResponse + if err := json.Unmarshal([]byte(jsonText), &resp); err == nil { + suggestions = resp.Suggestions + parsed = true + } } - var parsed checkpointResponse - if err := json.Unmarshal([]byte(jsonText), &parsed); err != nil { - return nil, fmt.Errorf("checkpoint: parse JSON: %w (got %q)", err, truncateForError(jsonText)) + if !parsed { + // No closeable object, or it failed to decode (most often a truncated + // array). Salvage the individual suggestion objects that completed. + suggestions = salvageSuggestions(raw) + if suggestions == nil { + return nil, fmt.Errorf("checkpoint: no JSON object in model output: %q", truncateForError(raw)) + } } // Drop items the model returned with an empty original — they can't be // anchored — and normalize whitespace the model may have echoed. - out := parsed.Suggestions[:0] - for _, s := range parsed.Suggestions { + out := suggestions[:0] + for _, s := range suggestions { s.Original = strings.TrimSpace(s.Original) s.Replacement = strings.TrimSpace(s.Replacement) if s.Original == "" { @@ -72,12 +94,55 @@ func ParseCheckpoint(raw string) ([]RawSuggestion, error) { return out, nil } +// salvageSuggestions recovers as many complete suggestion objects as it can from +// a response whose top-level JSON didn't parse — typically one truncated mid +// "suggestions" array. It scans the array region for brace-balanced {...} +// objects and keeps each that decodes into a RawSuggestion with a non-empty +// original; the final, cut-off object simply never balances and is dropped. +// Returns nil when there is no salvageable array (so the caller can error). +func salvageSuggestions(raw string) []RawSuggestion { + key := strings.Index(raw, `"suggestions"`) + if key < 0 { + return nil + } + lb := strings.IndexByte(raw[key:], '[') + if lb < 0 { + return nil + } + rest := raw[key+lb+1:] + out := []RawSuggestion{} + for { + obj, end := firstBalancedObject(rest) + if end < 0 { + break + } + var s RawSuggestion + if err := json.Unmarshal([]byte(obj), &s); err == nil && strings.TrimSpace(s.Original) != "" { + out = append(out, s) + } + rest = rest[end:] + } + if len(out) == 0 { + return nil + } + return out +} + // extractJSONObject returns the substring from the first '{' to its matching // closing '}', or "" if none. Tolerates fences/preamble around the object. func extractJSONObject(s string) string { + obj, _ := firstBalancedObject(s) + return obj +} + +// firstBalancedObject returns the first brace-balanced {...} substring in s and +// the index in s just past its closing brace, or ("", -1) if no object closes +// (none present, or the only one is truncated). String contents — including +// braces inside quoted values — are skipped so they can't unbalance the count. +func firstBalancedObject(s string) (string, int) { start := strings.IndexByte(s, '{') if start < 0 { - return "" + return "", -1 } depth := 0 inString := false @@ -98,11 +163,11 @@ func extractJSONObject(s string) string { case c == '}': depth-- if depth == 0 { - return s[start : i+1] + return s[start : i+1], i + 1 } } } - return "" + return "", -1 } func truncateForError(s string) string { diff --git a/internal/llm/checkpoint_test.go b/internal/llm/checkpoint_test.go index b3c8272..b42c7ba 100644 --- a/internal/llm/checkpoint_test.go +++ b/internal/llm/checkpoint_test.go @@ -42,6 +42,21 @@ func TestParseCheckpoint(t *testing.T) { raw: `{"suggestions":[{"original":"use {x}","replacement":"use x","explanation":"drop the braces","type":"clarity"}]}`, want: 1, }, + { + // The bug that 502'd every long doc: the model hit num_predict and the + // JSON never closed. Two suggestions completed; the third is cut off. + // Salvage keeps the two that completed instead of failing outright. + name: "truncated mid-array salvages completed objects", + raw: `{"suggestions":[{"original":"I has","replacement":"I have","explanation":"subject-verb agreement","type":"grammar"},{"original":"a apple","replacement":"an apple","explanation":"use an before a vowel","type":"grammar"},{"original":"she relised","replacement":"she real`, + want: 2, + }, + { + // A truncated array where even the first object is incomplete has + // nothing to salvage → still an error (the caller releases + retries). + name: "truncated before any object closes", + raw: `{"suggestions":[{"original":"It was a thursdy morning when Clara`, + wantErr: true, + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) {