fix(llm): pin claude-cli output to a JSON schema when supported (#2076)

graphify/llm.py

@@ -1337,6 +1337,62 @@ def _claude_cli_envelope(stdout: str) -> dict:
     return envelope
 
 
+# A JSON Schema pinning the top-level shape graphify consumes. Passed to
+# `claude -p --json-schema` (structured output) so the CLI CONSTRAINS the model
+# to emit the object directly instead of relying on it CHOOSING to honour a
+# "raw JSON only" instruction in the prompt. Item internals stay loose so a
+# valid extraction is never rejected; the `result` envelope field still carries
+# the JSON string, so the parse path is unchanged. See #2076.
+_EXTRACTION_JSON_SCHEMA = json.dumps(
+    {
+        "type": "object",
+        "properties": {
+            "nodes": {"type": "array", "items": {"type": "object"}},
+            "edges": {"type": "array", "items": {"type": "object"}},
+            "hyperedges": {"type": "array", "items": {"type": "object"}},
+        },
+        "required": ["nodes", "edges"],
+    }
+)
+
+# Cache the `--json-schema` capability probe per resolved claude command so it
+# runs at most once per process (extract fans a chunk out per file/slice).
+_JSON_SCHEMA_SUPPORT: dict[str, bool] = {}
+
+
+def _claude_cli_supports_json_schema(claude_cmd: str) -> bool:
+    """Return True if this Claude Code CLI accepts ``--json-schema``.
+
+    Structured output (``--json-schema``) landed in newer Claude Code releases.
+    Probing ``claude --help`` for the flag is a direct capability check — more
+    reliable than guessing a version boundary — so graphify uses structured
+    output where it exists and falls back to the user-turn prompt on older CLIs
+    that predate it. Any probe failure is treated as "unsupported" (safe
+    fallback). Result is cached per resolved command.
+    """
+    import subprocess
+
+    cached = _JSON_SCHEMA_SUPPORT.get(claude_cmd)
+    if cached is not None:
+        return cached
+    try:
+        proc = subprocess.run(
+            [claude_cmd, "--help"],
+            capture_output=True,
+            text=True,
+            encoding="utf-8",
+            errors="replace",
+            timeout=30,
+            check=False,
+            **_no_window_kwargs(),
+        )
+        supported = "--json-schema" in (proc.stdout or "")
+    except (OSError, subprocess.SubprocessError):
+        supported = False
+    _JSON_SCHEMA_SUPPORT[claude_cmd] = supported
+    return supported
+
+
 def _call_claude_cli(user_message: str, max_tokens: int = 8192, *, deep_mode: bool = False, images: list[_ImageRef] | None = None) -> dict:
     """Call Claude via the locally-installed Claude Code CLI (`claude -p`).
 
@@ -1425,6 +1481,16 @@ def _call_claude_cli(user_message: str, max_tokens: int = 8192, *, deep_mode: bo
     cli_model = os.environ.get("GRAPHIFY_CLAUDE_CLI_MODEL", "").strip()
     if cli_model:
         cli_args.extend(["--model", cli_model])
+    # Constrain the output shape structurally where the CLI supports it. Newer
+    # Claude Code releases increasingly treat a bare file-dump prompt as an
+    # agentic task and REPORT the extraction in prose ("Knowledge graph
+    # extracted — 21 nodes, 20 edges…") instead of returning it; that parses to
+    # zero nodes, reads as truncation, and gets bisected without ever
+    # converging (#2076). --json-schema pins the object shape regardless of
+    # that framing; the user-turn prompt above stays as the fallback for older
+    # CLIs that predate the flag.
+    if _claude_cli_supports_json_schema(claude_cmd):
+        cli_args.extend(["--json-schema", _EXTRACTION_JSON_SCHEMA])
     proc = subprocess.run(
         cli_args,
         input=combined_message,

tests/test_claude_cli_backend.py

@@ -154,6 +154,63 @@ def test_user_turn_preserves_untrusted_source_guardrails(fake_claude):
     assert "untrusted_source" in sent
 
 
+# ---------- structured output via --json-schema (#2076) ----------
+# Newer Claude Code CLIs treat a bare file-dump prompt as an agentic task and
+# REPORT the extraction in prose instead of returning JSON, so the graph comes
+# out empty and adaptive-retry bisects forever. When the CLI supports
+# `--json-schema`, graphify constrains the output shape structurally so the
+# model must emit the object regardless of framing. Older CLIs that predate the
+# flag fall back to the user-turn prompt, unchanged.
+
+
+def test_json_schema_flag_added_when_cli_supports_it(monkeypatch, fake_claude):
+    """When the CLI advertises --json-schema, it is passed with a schema that
+    pins the top-level {nodes, edges} shape graphify parses."""
+    monkeypatch.setattr(llm, "_claude_cli_supports_json_schema", lambda cmd: True)
+    llm._call_claude_cli("dummy source", max_tokens=8192)
+    argv = fake_claude.call_args.args[0]
+    assert "--json-schema" in argv
+    schema = json.loads(argv[argv.index("--json-schema") + 1])
+    assert schema["type"] == "object"
+    assert set(schema["required"]) == {"nodes", "edges"}
+
+
+def test_json_schema_flag_absent_when_cli_lacks_it(monkeypatch, fake_claude):
+    """Older CLIs without --json-schema must not receive the flag (it would be
+    an unknown option) — extraction falls back to the user-turn prompt."""
+    monkeypatch.setattr(llm, "_claude_cli_supports_json_schema", lambda cmd: False)
+    result = llm._call_claude_cli("dummy source", max_tokens=8192)
+    argv = fake_claude.call_args.args[0]
+    assert "--json-schema" not in argv
+    assert len(result["nodes"]) == 2  # result envelope still carries the JSON
+
+
+def test_supports_json_schema_detects_flag_in_help():
+    llm._JSON_SCHEMA_SUPPORT.clear()
+    help_text = "Options:\n  --json-schema <schema>  JSON Schema for structured output\n"
+    completed = MagicMock(returncode=0, stdout=help_text, stderr="")
+    with patch("subprocess.run", return_value=completed):
+        assert llm._claude_cli_supports_json_schema("/fake/claude-new") is True
+
+
+def test_supports_json_schema_false_when_flag_absent():
+    llm._JSON_SCHEMA_SUPPORT.clear()
+    help_text = "Options:\n  --output-format <format>  text|json|stream-json\n"
+    completed = MagicMock(returncode=0, stdout=help_text, stderr="")
+    with patch("subprocess.run", return_value=completed):
+        assert llm._claude_cli_supports_json_schema("/fake/claude-old") is False
+
+
+def test_supports_json_schema_false_and_cached_on_probe_error():
+    """A probe that fails to run is treated as unsupported (safe fallback) and
+    cached so it is not re-probed for every chunk."""
+    llm._JSON_SCHEMA_SUPPORT.clear()
+    with patch("subprocess.run", side_effect=OSError("boom")) as run:
+        assert llm._claude_cli_supports_json_schema("/fake/claude-broken") is False
+        assert llm._claude_cli_supports_json_schema("/fake/claude-broken") is False
+    assert run.call_count == 1
+
+
 # ---------- Windows path resolution (#1072) ----------