From 33ed9c2fad0fd0ea451affb277a351c2bf5db8a3 Mon Sep 17 00:00:00 2001
From: Scott Severance <scott@scottseverance.net>
Date: Wed, 26 Aug 2026 13:52:48 -0500
Subject: [PATCH] fix(ai-improve): tolerate prose-wrapped JSON from the model

Reproduced 3/3 on Data-Science-Toolkit: the selection call got a valid
response, but Haiku sometimes reasons in prose ("Looking at this
codebase, I need to identify...") instead of emitting bare JSON, despite
the existing "ONLY valid JSON" instruction. ask_json() only tried
json.loads() on the whole trimmed response, so any surrounding prose
made it discard an otherwise-usable answer and skip the run.

- Strengthen SELECT_PROMPT and WRITE_PROMPT: explicitly forbid
  reasoning/commentary, since a caller parses the response
  programmatically
- Add a fallback in ask_json(): if the whole response isn't valid JSON,
  scan for the first balanced {...} object anywhere in it and parse
  that instead of giving up
- Print more of the raw response on a genuine parse failure (800 -> 2000
  chars) for easier debugging next time this happens
---
 scripts/ai_improve.py | 36 +++++++++++++++++++++++++++++++-----
 1 file changed, 31 insertions(+), 5 deletions(-)

diff --git a/scripts/ai_improve.py b/scripts/ai_improve.py
index d2d5252..550f40d 100644
--- a/scripts/ai_improve.py
+++ b/scripts/ai_improve.py
@@ -166,7 +166,11 @@ def rank(path: str) -> int:
 genuinely different area of the codebase
 - Do NOT pick README.md, *.yml/*.yaml workflow files, or lock files
 
-Respond with ONLY valid JSON — no markdown fences, no prose outside the object:
+Do not think out loud, explain your reasoning, or walk through your analysis of \
+the codebase — a caller is parsing your response programmatically and any text \
+outside the object below will break it. Respond with ONLY the JSON object, \
+starting with `{{` and ending with `}}` — no markdown fences, no preamble, no \
+commentary before or after it:
 {{"file_path":"relative/path/to/file","plan":"one sentence naming the concrete change"}}
 """
 
@@ -193,7 +197,10 @@ def rank(path: str) -> int:
 - Commit message: conventional commits format, imperative mood, ≤72 chars \
 (e.g. "fix: handle empty feature store on first run")
 
-Respond with ONLY valid JSON — no markdown fences, no prose outside the object:
+Do not think out loud, explain your reasoning, or walk through your analysis — a \
+caller is parsing your response programmatically and any text outside the object \
+below will break it. Respond with ONLY the JSON object, starting with `{{` and \
+ending with `}}` — no markdown fences, no preamble, no commentary before or after it:
 {{"file_content":"complete file content here",\
 "commit_message":"type(scope): description","pr_title":"Short PR title (≤60 chars)",\
 "pr_body":"## What\\nOne sentence.\\n\\n## Why\\nOne sentence."}}
@@ -226,9 +233,28 @@ def ask_json(client, prompt: str, max_tokens: int) -> dict | None:
 
     try:
         return json.loads(raw)
-    except json.JSONDecodeError as exc:
-        print(f"JSON parse error: {exc}\nRaw response:\n{raw[:800]}")
-        return None
+    except json.JSONDecodeError:
+        pass
+
+    # Despite the instruction, the model sometimes reasons in prose before (or
+    # instead of) emitting bare JSON. Rather than discard a usable answer,
+    # look for the first balanced {...} object anywhere in the response.
+    start = raw.find("{")
+    if start != -1:
+        depth = 0
+        for i, ch in enumerate(raw[start:], start):
+            if ch == "{":
+                depth += 1
+            elif ch == "}":
+                depth -= 1
+                if depth == 0:
+                    try:
+                        return json.loads(raw[start:i + 1])
+                    except json.JSONDecodeError:
+                        break
+
+    print(f"JSON parse error; no valid JSON object found.\nRaw response:\n{raw[:2000]}")
+    return None
 
 
 def safe_target(raw_path: str, repo_root: Path) -> Path | None: