Merge fix/checkpoint-token-truncation: stop 502s on long-doc checkpoints

Claude-Session: https://claude.ai/code/session_016Yr6jELuRc7hyzYLccQKZd

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)
+	var suggestions []RawSuggestion
-	if jsonText == "" {
+	parsed := false
+	if jsonText := extractJSONObject(raw); jsonText != "" {
+		var resp checkpointResponse
+		if err := json.Unmarshal([]byte(jsonText), &resp); err == nil {
+			suggestions = resp.Suggestions
+			parsed = true
+		}
+	}
+	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))
 		}
-	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))
 	}
 	// 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]
+	out := suggestions[:0]
-	for _, s := range parsed.Suggestions {
+	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 {

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) {