Align Assistent debug patch extraction with string generate flags and unfenced JSON.

Co-authored-by: Cursor <cursoragent@cursor.com>

src/gpu_rent/debug_assistent.py

@@ -52,6 +52,7 @@ _PATCH_KEYS = frozenset(
         "vision_from",
         "vision_slots",
         "variants",
+        "checkpoint",
     }
 )
 _FENCE_RE = re.compile(r"```(?:json)?\s*([\s\S]*?)```", re.IGNORECASE)
@@ -343,12 +344,39 @@ def _compact_api(name: str, data: Any) -> Any:
     return {"keys": list(data.keys())[:20]}
 
 
+def _generate_flag_on(obj: dict[str, Any] | None) -> bool:
+    if not isinstance(obj, dict):
+        return False
+    g = obj.get("generate")
+    if g is True or g == 1:
+        return True
+    if isinstance(g, str) and g.strip().lower() in {"true", "1", "yes", "on"}:
+        return True
+    acts = obj.get("actions")
+    return isinstance(acts, list) and "generate" in [str(a) for a in acts]
+
+
+def _normalize_extracted_patch(obj: dict[str, Any]) -> dict[str, Any]:
+    patch = dict(obj)
+    if _generate_flag_on(patch):
+        patch["generate"] = True
+    if isinstance(patch.get("ask"), str):
+        one = patch["ask"].strip()
+        if one:
+            patch["ask"] = [one]
+        else:
+            patch.pop("ask", None)
+    return patch
+
+
 def extract_assistent_patch(text: str | None) -> dict[str, Any]:
     """Pull last fenced JSON patch from Assistent reply (mirrors SA.extractPatch)."""
     if not text:
         return {"prose": text or "", "patch": None}
-    last_patch: dict[str, Any] | None = None
+    last_any: dict[str, Any] | None = None
-    prose = text
+    last_any_span: tuple[int, int] | None = None
+    last_term: dict[str, Any] | None = None
+    last_term_span: tuple[int, int] | None = None
     for match in _FENCE_RE.finditer(text):
         try:
             obj = json.loads(match.group(1).strip())
@@ -358,15 +386,29 @@ def extract_assistent_patch(text: str | None) -> dict[str, Any]:
             continue
         if not any(k in obj and obj[k] is not None for k in _PATCH_KEYS):
             continue
-        patch = dict(obj)
+        patch = _normalize_extracted_patch(obj)
-        acts = [str(a) for a in patch["actions"]] if isinstance(patch.get("actions"), list) else []
+        last_any = patch
-        if patch.get("generate") is True or "generate" in acts:
+        last_any_span = (match.start(), match.end())
-            patch["generate"] = True
+        if _generate_flag_on(patch) or patch.get("look_at") is not None or patch.get("ask"):
-        if isinstance(patch.get("ask"), str):
+            last_term = patch
-            patch["ask"] = [patch["ask"]]
+            last_term_span = last_any_span
-        last_patch = patch
+        elif isinstance(patch.get("prompt"), str) and len(patch["prompt"].strip()) >= 48:
-        prose = (text[: match.start()] + text[match.end() :]).strip()
+            last_term = patch
-    return {"prose": prose, "patch": last_patch}
+            last_term_span = last_any_span
+    chosen = last_term or last_any
+    span = last_term_span or last_any_span
+    if chosen and span:
+        prose = (text[: span[0]] + text[span[1] :]).strip()
+        return {"prose": prose, "patch": chosen}
+    brace = text.rfind("{")
+    if brace >= 0:
+        try:
+            obj = json.loads(text[brace:].strip())
+        except (json.JSONDecodeError, TypeError, ValueError):
+            obj = None
+        if isinstance(obj, dict) and any(k in obj and obj[k] is not None for k in _PATCH_KEYS):
+            return {"prose": text[:brace].strip(), "patch": _normalize_extracted_patch(obj)}
+    return {"prose": text, "patch": None}
 
 
 def _summarize_patch(patch: dict[str, Any] | None) -> dict[str, Any] | None:

src/gpu_rent/debug_assistent_session.py

@@ -36,6 +36,7 @@ from gpu_rent.debug_assistent import (
     _swarm_base,
     collect_assistent_deep,
     extract_assistent_patch,
+    _generate_flag_on,
 )
 
 SESSION_TTL_SEC = 45 * 60
@@ -245,14 +246,7 @@ def analyze_exact_merge(
 ) -> dict[str, Any]:
     """Simulate client mergeExactParamsForGenerate for diagnostics."""
     patch = patch if isinstance(patch, dict) else None
-    acts = patch.get("actions") if patch else None
+    wants_gen = _generate_flag_on(patch)
-    wants_gen = bool(
-        patch
-        and (
-            patch.get("generate") is True
-            or (isinstance(acts, list) and "generate" in [str(a) for a in acts])
-        )
-    )
     defaults = resolve_exact_profile_defaults(exact, profile_name=krea_profile)
     if recommended and isinstance(recommended, dict):
         for k in EXACT_GENERATE_PARAM_KEYS:

tests/test_debug_api.py

@@ -190,6 +190,17 @@ def test_extract_assistent_patch():
     assert got2["patch"]["steps"] == 8
     assert got2["patch"]["cfg"] == 2
 
+    str_flag = extract_assistent_patch(
+        'ok\n```json\n{"prompt":"a slender redhead in leather","generate":"true"}\n```'
+    )
+    assert str_flag["patch"]["generate"] is True
+
+    raw = extract_assistent_patch(
+        'done\n{"prompt":"a fiery redhead, street light, 85mm","generate":true}'
+    )
+    assert raw["patch"]["generate"] is True
+    assert raw["prose"] == "done"
+
 
 def test_chat_eval_mocked_http(monkeypatch, tmp_path):
     from gpu_rent import debug_assistent