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 {
@@ -127,17 +192,33 @@ func NewRateLimiter(interval time.Duration) *RateLimiter {
 }
 
 // Allow reports whether a checkpoint may run for docID now. When allowed it
-// records the time and returns (true, 0); when throttled it returns
+// records the time and returns (true, 0, recordedAt); when throttled it returns
-// (false, retryAfter) where retryAfter is the wait until the next allowed run.
+// (false, retryAfter, zero) where retryAfter is the wait until the next allowed
-func (rl *RateLimiter) Allow(docID string) (bool, time.Duration) {
+// run. recordedAt lets a caller whose pass fails hand the exact timestamp to
+// Release so it rolls back only its own slot.
+func (rl *RateLimiter) Allow(docID string) (bool, time.Duration, time.Time) {
 	rl.mu.Lock()
 	defer rl.mu.Unlock()
 	now := time.Now()
 	if last, ok := rl.last[docID]; ok {
 		if elapsed := now.Sub(last); elapsed < rl.interval {
-			return false, rl.interval - elapsed
+			return false, rl.interval - elapsed, time.Time{}
 		}
 	}
 	rl.last[docID] = now
-	return true, 0
+	return true, 0, now
+}
+
+// Release rolls back the slot a prior Allow recorded for docID, but only if no
+// newer Allow has since claimed it. Called when an LLM pass fails: Allow runs
+// before the model call, so without this a failed checkpoint would hold the
+// per-document slot for the full interval and a retry (or the frontend's
+// auto-retry) would hit the throttle path and get the stale/empty set back
+// instead of re-running. `at` is the time the failed Allow returned.
+func (rl *RateLimiter) Release(docID string, at time.Time) {
+	rl.mu.Lock()
+	defer rl.mu.Unlock()
+	if last, ok := rl.last[docID]; ok && last.Equal(at) {
+		delete(rl.last, docID)
+	}
 }

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) {
@@ -65,16 +80,31 @@ func TestParseCheckpoint(t *testing.T) {
 func TestRateLimiter(t *testing.T) {
 	rl := NewRateLimiter(30 * time.Second)
 
-	if ok, _ := rl.Allow("doc1"); !ok {
+	ok, _, at := rl.Allow("doc1")
+	if !ok {
 		t.Fatal("first call should be allowed")
 	}
-	if ok, retry := rl.Allow("doc1"); ok || retry <= 0 {
+	if ok, retry, _ := rl.Allow("doc1"); ok || retry <= 0 {
 		t.Fatalf("immediate second call should be throttled, got ok=%v retry=%v", ok, retry)
 	}
 	// A different document is independent.
-	if ok, _ := rl.Allow("doc2"); !ok {
+	if ok, _, _ := rl.Allow("doc2"); !ok {
 		t.Fatal("different doc should be allowed")
 	}
+
+	// Releasing doc1's slot (as a failed pass does) lets the next check re-run
+	// immediately instead of waiting out the interval.
+	rl.Release("doc1", at)
+	if ok, _, _ := rl.Allow("doc1"); !ok {
+		t.Fatal("after Release the next call should be allowed again")
+	}
+	// Release only rolls back its own slot: a stale timestamp must not evict the
+	// slot a newer Allow holds (so a late failure can't cancel a fresh in-flight
+	// check). doc2 still holds its slot from above; a zero-time Release is a no-op.
+	rl.Release("doc2", time.Time{})
+	if ok, _, _ := rl.Allow("doc2"); ok {
+		t.Fatal("stale Release must not free doc2's current slot")
+	}
 }
 
 func TestTruncateDoc(t *testing.T) {

internal/suggestions/handlers.go

@@ -102,7 +102,8 @@ func (h *Handler) runPass(w http.ResponseWriter, r *http.Request, limiter *llm.R
 		return
 	}
 
-	if ok, _ := limiter.Allow(docID); !ok {
+	ok, _, slotAt := limiter.Allow(docID)
+	if !ok {
 		// Throttled: return the existing pending set unchanged rather than an
 		// error, so the frontend keeps showing current suggestions.
 		existing, err := h.fetchPending(docID)
@@ -116,6 +117,10 @@ func (h *Handler) runPass(w http.ResponseWriter, r *http.Request, limiter *llm.R
 
 	raw, err := run(r.Context(), h.Client, contentText, tone)
 	if err != nil {
+		// Allow ran before the model call, so a failed pass would otherwise hold
+		// the per-document slot for the full interval — stranding the frontend's
+		// auto-retry on the throttle path. Release it so a retry can re-run.
+		limiter.Release(docID, slotAt)
 		errorJSON(w, http.StatusBadGateway, "llm pass failed: "+err.Error())
 		return
 	}

web/src/hooks/useCheckpoint.ts

@@ -20,17 +20,29 @@ export function useCheckpoint(docId: string | null) {
   const [llmDown, setLlmDown] = useState(false)
 
   const debounceRef = useRef<ReturnType<typeof setTimeout>>(undefined)
+  // Pending auto-retry timer for a failed checkpoint (see runCheck).
+  const retryRef = useRef<ReturnType<typeof setTimeout>>(undefined)
   const docIdRef = useRef(docId)
   docIdRef.current = docId
 
   // Token to discard responses from a doc we've since navigated away from.
   const runRef = useRef(0)
 
-  const runCheck = useCallback(async () => {
+  // Backoff schedule for re-running a checkpoint that failed. A paste fires
+  // exactly ONE checkpoint, and nothing re-fires until the next keystroke — so a
+  // single transient failure (a momentary stall on the shared inference box)
+  // would otherwise strand the user on "resting" indefinitely after a paste.
+  // These delays clear the server's 30s per-doc floor by the later attempts, and
+  // a failed pass now releases its slot server-side so a retry can truly re-run.
+  const RETRY_DELAYS_MS = [3000, 12000, 35000]
+
+  const runCheck = useCallback(async (attempt = 0) => {
     const id = docIdRef.current
     if (!id) return
+    clearTimeout(retryRef.current)
     const run = ++runRef.current
     setChecking(true)
+    let retrying = false
     try {
       const fresh = await api.checkDoc(id)
       if (run === runRef.current && id === docIdRef.current) {
@@ -39,9 +51,18 @@ export function useCheckpoint(docId: string | null) {
       }
     } catch (err) {
       console.error('checkpoint failed', err)
-      if (run === runRef.current) setLlmDown(true)
+      if (run !== runRef.current) return
+      if (attempt < RETRY_DELAYS_MS.length) {
+        // Keep trying quietly — don't flag "resting" until retries are exhausted.
+        retrying = true
+        retryRef.current = setTimeout(() => void runCheck(attempt + 1), RETRY_DELAYS_MS[attempt])
+      } else {
+        setLlmDown(true)
+      }
     } finally {
-      if (run === runRef.current) setChecking(false)
+      // Stay in the "checking" state while a retry is queued so the breathing dot
+      // keeps reassuring rather than flickering off between attempts.
+      if (run === runRef.current && !retrying) setChecking(false)
     }
   }, [])
 
@@ -51,6 +72,7 @@ export function useCheckpoint(docId: string | null) {
   const runVoice = useCallback(async () => {
     const id = docIdRef.current
     if (!id) return
+    clearTimeout(retryRef.current) // a voice pass supersedes a queued grammar retry
     const run = ++runRef.current
     setVoicing(true)
     try {
@@ -70,13 +92,15 @@ export function useCheckpoint(docId: string | null) {
   // Call on every edit; schedules a check 4s after typing settles.
   const schedule = useCallback(() => {
     clearTimeout(debounceRef.current)
-    debounceRef.current = setTimeout(runCheck, DEBOUNCE_MS)
+    clearTimeout(retryRef.current) // a fresh edit supersedes any queued retry
+    debounceRef.current = setTimeout(() => void runCheck(), DEBOUNCE_MS)
   }, [runCheck])
 
   // Load existing pending suggestions whenever the document changes, and cancel
   // any in-flight debounce from the previous doc.
   useEffect(() => {
     clearTimeout(debounceRef.current)
+    clearTimeout(retryRef.current)
     runRef.current++
     setSuggestions([])
     setChecking(false)
@@ -97,7 +121,13 @@ export function useCheckpoint(docId: string | null) {
     }
   }, [docId])
 
-  useEffect(() => () => clearTimeout(debounceRef.current), [])
+  useEffect(
+    () => () => {
+      clearTimeout(debounceRef.current)
+      clearTimeout(retryRef.current)
+    },
+    [],
+  )
 
   // Drop one suggestion locally (after accept/dismiss) without a refetch.
   const removeSuggestion = useCallback((id: string) => {