feat(llm): split and retry chunks that hit max_completion_tokens truncation

Token-budget chunking cuts the truncation rate but doesn't eliminate
it. Output token cost scales with extractable concept density rather
than input tokens — a chunk that lands on a directory of dense design
docs can pack under the input budget while needing more than
`max_completion_tokens=8192` to express every named concept, so the
response is truncated mid-string and `_parse_llm_json` returns an
empty fragment.

Pre-tuning chunk size to be conservative enough that this never
happens leaves throughput on the table for the common case. Adding a
hard `max_files_per_chunk` cap on top of `token_budget` reintroduces
the "tune a static constant" problem the previous commit set out to
fix.

The fix uses the API's own truncation signal:

1. `_call_openai_compat` and `_call_claude` now expose `finish_reason`
   on the result dict (Anthropic's `stop_reason == "max_tokens"` is
   normalised to `"length"`).
2. `_extract_with_adaptive_retry` checks it: when truncated, splits
   the chunk in half and recurses on each half. Recursion is bounded
   by `max_retry_depth` (default 3 → at most 8x fanout per top-level
   chunk).
3. Single-file chunks that truncate can't recover and surface a
   warning rather than infinite-loop.
4. `extract_corpus_parallel` routes every chunk through the retry
   wrapper. The `on_chunk_done` callback fires once per top-level
   chunk with the merged result — recursive splits are invisible to
   callers.

graphify/llm.py

@@ -139,6 +139,10 @@ def _call_openai_compat(
     result["input_tokens"] = resp.usage.prompt_tokens if resp.usage else 0
     result["output_tokens"] = resp.usage.completion_tokens if resp.usage else 0
     result["model"] = model
+    # `finish_reason == "length"` means the model hit max_completion_tokens
+    # mid-generation. The JSON we got back is truncated; callers should
+    # treat this as a signal to retry with smaller input.
+    result["finish_reason"] = resp.choices[0].finish_reason
     return result
 
 
@@ -163,6 +167,10 @@ def _call_claude(api_key: str, model: str, user_message: str) -> dict:
     result["input_tokens"] = resp.usage.input_tokens if resp.usage else 0
     result["output_tokens"] = resp.usage.output_tokens if resp.usage else 0
     result["model"] = model
+    # Normalise Anthropic's `stop_reason` to the OpenAI-compat `finish_reason`
+    # vocabulary so the adaptive-retry layer doesn't have to know which
+    # backend produced the result.
+    result["finish_reason"] = "length" if resp.stop_reason == "max_tokens" else "stop"
     return result
 
 
@@ -261,6 +269,83 @@ def _pack_chunks_by_tokens(
     return chunks
 
 
+def _extract_with_adaptive_retry(
+    chunk: list[Path],
+    backend: str,
+    api_key: str | None,
+    model: str | None,
+    root: Path,
+    max_depth: int,
+    _depth: int = 0,
+) -> dict:
+    """Extract a chunk; if the response is truncated (`finish_reason="length"`),
+    split the chunk in half and recurse.
+
+    The signal driving the retry is the API's own `finish_reason` — `"length"`
+    means the model hit `max_completion_tokens` mid-output. The truncated JSON
+    has nothing useful in it (parse fails partway through a string or array),
+    so we discard it and re-extract on smaller inputs that produce shorter
+    outputs.
+
+    Recursion is capped at `max_depth` to bound worst-case cost. A chunk of N
+    files can split into up to 2**max_depth pieces — at depth=3 that's 8x. If
+    still truncated at the cap, we surface the (likely empty) result with a
+    warning rather than infinite-loop.
+
+    A single-file chunk that truncates is unrecoverable here — we can't make
+    one file smaller than itself, so we return what we got and warn.
+    """
+    result = extract_files_direct(
+        chunk, backend=backend, api_key=api_key, model=model, root=root
+    )
+
+    if result.get("finish_reason") != "length":
+        return result
+
+    if len(chunk) <= 1:
+        print(
+            f"[graphify] single-file chunk {chunk[0]} truncated at "
+            f"max_completion_tokens — partial result kept",
+            file=sys.stderr,
+        )
+        return result
+
+    if _depth >= max_depth:
+        print(
+            f"[graphify] chunk of {len(chunk)} still truncated at recursion "
+            f"depth {_depth} (max {max_depth}) — partial result kept",
+            file=sys.stderr,
+        )
+        return result
+
+    print(
+        f"[graphify] chunk of {len(chunk)} truncated at depth {_depth}, "
+        f"splitting into halves of {len(chunk) // 2} and "
+        f"{len(chunk) - len(chunk) // 2}",
+        file=sys.stderr,
+    )
+    mid = len(chunk) // 2
+    left = _extract_with_adaptive_retry(
+        chunk[:mid], backend, api_key, model, root, max_depth, _depth + 1
+    )
+    right = _extract_with_adaptive_retry(
+        chunk[mid:], backend, api_key, model, root, max_depth, _depth + 1
+    )
+
+    return {
+        "nodes": left.get("nodes", []) + right.get("nodes", []),
+        "edges": left.get("edges", []) + right.get("edges", []),
+        "hyperedges": left.get("hyperedges", []) + right.get("hyperedges", []),
+        "input_tokens": left.get("input_tokens", 0) + right.get("input_tokens", 0),
+        "output_tokens": left.get("output_tokens", 0) + right.get("output_tokens", 0),
+        "model": result.get("model"),
+        # Both halves either succeeded or have already surfaced their own
+        # truncation warning; the merged result is no longer truncated as a
+        # logical unit.
+        "finish_reason": "stop",
+    }
+
+
 def extract_corpus_parallel(
     files: list[Path],
     backend: str = "kimi",
@@ -271,6 +356,7 @@ def extract_corpus_parallel(
     on_chunk_done: Callable | None = None,
     token_budget: int | None = 60_000,
     max_concurrency: int = 4,
+    max_retry_depth: int = 3,
 ) -> dict:
     """Extract a corpus in chunks, merging results.
 
@@ -287,9 +373,20 @@ def extract_corpus_parallel(
           (default 4 — conservative to stay under provider rate limits).
         - Set `max_concurrency=1` to force sequential execution.
 
+    Adaptive retry on truncation:
+        - When the LLM returns `finish_reason="length"` (output truncated at
+          `max_completion_tokens`), the chunk is split in half and each half
+          re-extracted recursively, up to `max_retry_depth` levels deep
+          (default 3 → max 8x expansion of one chunk).
+        - This is signal-driven: chunks too dense to fit in one response
+          self-heal by splitting until they do, while well-sized chunks pay
+          no extra cost. Set `max_retry_depth=0` to disable retries.
+
     `on_chunk_done(idx, total, chunk_result)` fires once per chunk as it
     completes (in completion order, not submission order). `idx` is the
-    chunk's submission index so callers can correlate progress.
+    chunk's submission index so callers can correlate progress. The
+    callback fires once per top-level chunk; recursive splits are merged
+    transparently before the callback is invoked.
 
     Returns merged dict with nodes, edges, hyperedges, input_tokens,
     output_tokens. Failed chunks are logged to stderr and skipped — one bad
@@ -306,7 +403,14 @@ def extract_corpus_parallel(
     def _run_one(idx: int, chunk: list[Path]) -> tuple[int, dict | None, Exception | None]:
         t0 = time.time()
         try:
-            result = extract_files_direct(chunk, backend=backend, api_key=api_key, model=model, root=root)
+            result = _extract_with_adaptive_retry(
+                chunk,
+                backend=backend,
+                api_key=api_key,
+                model=model,
+                root=root,
+                max_depth=max_retry_depth,
+            )
             result["elapsed_seconds"] = round(time.time() - t0, 2)
             return idx, result, None
         except Exception as exc:  # noqa: BLE001 — caller-facing surface, log + continue

tests/test_chunking.py

@@ -275,3 +275,178 @@ def test_corpus_parallel_token_budget_default_packs_files(tmp_path):
     # 50 tiny files at default 60k token budget should pack into 1 chunk
     assert len(chunks_seen) == 1
     assert chunks_seen[0] == 50
+
+
+# ---- Adaptive retry on truncation -------------------------------------------
+
+def _stub_with_finish(file_count: int, finish_reason: str = "stop") -> dict:
+    """Build a stub extraction result with a controllable finish_reason."""
+    return {
+        "nodes": [{"id": f"n_{i}"} for i in range(file_count)],
+        "edges": [],
+        "hyperedges": [],
+        "input_tokens": 100 * file_count,
+        "output_tokens": 50 * file_count,
+        "finish_reason": finish_reason,
+    }
+
+
+def test_adaptive_retry_returns_directly_when_not_truncated(tmp_path):
+    """No retry when finish_reason='stop' — single call, result passes through."""
+    from graphify.llm import _extract_with_adaptive_retry
+
+    files = [tmp_path / f"f{i}.py" for i in range(4)]
+    for f in files:
+        f.write_text("x")
+
+    calls = []
+
+    def stub(chunk, **kwargs):
+        calls.append(len(chunk))
+        return _stub_with_finish(len(chunk), finish_reason="stop")
+
+    with patch("graphify.llm.extract_files_direct", side_effect=stub):
+        result = _extract_with_adaptive_retry(
+            files, backend="kimi", api_key=None, model=None, root=tmp_path, max_depth=3
+        )
+
+    assert calls == [4], f"expected 1 call of 4 files, got {calls}"
+    assert len(result["nodes"]) == 4
+
+
+def test_adaptive_retry_splits_when_finish_reason_length(tmp_path):
+    """finish_reason='length' triggers split-in-half. Both halves succeed
+    on the second try (mocked) and results merge."""
+    from graphify.llm import _extract_with_adaptive_retry
+
+    files = [tmp_path / f"f{i}.py" for i in range(4)]
+    for f in files:
+        f.write_text("x")
+
+    calls = []
+
+    def stub(chunk, **kwargs):
+        calls.append(len(chunk))
+        finish = "length" if len(chunk) == 4 else "stop"
+        return _stub_with_finish(len(chunk), finish_reason=finish)
+
+    with patch("graphify.llm.extract_files_direct", side_effect=stub):
+        result = _extract_with_adaptive_retry(
+            files, backend="kimi", api_key=None, model=None, root=tmp_path, max_depth=3
+        )
+
+    assert calls == [4, 2, 2], f"expected [4, 2, 2], got {calls}"
+    assert len(result["nodes"]) == 4
+    assert result["finish_reason"] == "stop"
+
+
+def test_adaptive_retry_recurses_for_persistent_truncation(tmp_path):
+    """When even the half-chunk truncates, split again. With 8 files and a
+    truncation cutoff at >2 files, splits 8 → 4 → 2 (4 leaves of 2)."""
+    from graphify.llm import _extract_with_adaptive_retry
+
+    files = [tmp_path / f"f{i}.py" for i in range(8)]
+    for f in files:
+        f.write_text("x")
+
+    calls = []
+
+    def stub(chunk, **kwargs):
+        calls.append(len(chunk))
+        finish = "length" if len(chunk) > 2 else "stop"
+        return _stub_with_finish(len(chunk), finish_reason=finish)
+
+    with patch("graphify.llm.extract_files_direct", side_effect=stub):
+        result = _extract_with_adaptive_retry(
+            files, backend="kimi", api_key=None, model=None, root=tmp_path, max_depth=3
+        )
+
+    # Tree: 8 (trunc) → 4 + 4 (both trunc) → 2+2+2+2 (all stop)
+    # Total calls: 1 + 2 + 4 = 7
+    assert sorted(calls) == [2, 2, 2, 2, 4, 4, 8]
+    assert len(result["nodes"]) == 8
+
+
+def test_adaptive_retry_caps_at_max_depth(tmp_path, capsys):
+    """If everything truncates, retries stop at max_depth — partial result
+    kept with a warning, no infinite loop."""
+    from graphify.llm import _extract_with_adaptive_retry
+
+    files = [tmp_path / f"f{i}.py" for i in range(8)]
+    for f in files:
+        f.write_text("x")
+
+    calls = []
+
+    def always_truncate(chunk, **kwargs):
+        calls.append(len(chunk))
+        return _stub_with_finish(len(chunk), finish_reason="length")
+
+    with patch("graphify.llm.extract_files_direct", side_effect=always_truncate):
+        _extract_with_adaptive_retry(
+            files, backend="kimi", api_key=None, model=None, root=tmp_path, max_depth=2
+        )
+
+    # max_depth=2 bounds the tree: root + 2 + 4 = 7 calls maximum
+    assert len(calls) <= 7, f"recursion not bounded — {len(calls)} calls"
+    err = capsys.readouterr().err
+    assert "still truncated" in err
+
+
+def test_adaptive_retry_single_file_truncation_does_not_recurse(tmp_path, capsys):
+    """A single file that truncates can't be split further — surface a
+    warning and return what we got. No infinite loop."""
+    from graphify.llm import _extract_with_adaptive_retry
+
+    f = tmp_path / "huge.py"; f.write_text("x")
+
+    calls = []
+
+    def stub(chunk, **kwargs):
+        calls.append(len(chunk))
+        return _stub_with_finish(len(chunk), finish_reason="length")
+
+    with patch("graphify.llm.extract_files_direct", side_effect=stub):
+        _extract_with_adaptive_retry(
+            [f], backend="kimi", api_key=None, model=None, root=tmp_path, max_depth=3
+        )
+
+    assert calls == [1], f"single-file chunk recursed; calls = {calls}"
+    err = capsys.readouterr().err
+    assert "single-file chunk" in err and "truncated" in err
+
+
+def test_corpus_parallel_uses_adaptive_retry(tmp_path):
+    """End-to-end: extract_corpus_parallel routes through adaptive retry,
+    so a chunk that truncates gets split and merged transparently before
+    on_chunk_done fires."""
+    from graphify.llm import extract_corpus_parallel
+
+    files = [tmp_path / f"f{i}.py" for i in range(4)]
+    for f in files:
+        f.write_text("x")
+
+    calls = []
+
+    def stub(chunk, **kwargs):
+        calls.append(len(chunk))
+        finish = "length" if len(chunk) == 4 else "stop"
+        return _stub_with_finish(len(chunk), finish_reason=finish)
+
+    chunk_done_args = []
+    with patch("graphify.llm.extract_files_direct", side_effect=stub):
+        result = extract_corpus_parallel(
+            files,
+            backend="kimi",
+            token_budget=None,
+            chunk_size=4,
+            max_concurrency=1,
+            on_chunk_done=lambda i, t, r: chunk_done_args.append((i, t, len(r["nodes"]))),
+        )
+
+    # Adaptive retry runs INSIDE _run_one: 4 → 2 + 2 = 3 underlying API calls
+    assert calls == [4, 2, 2]
+    # User-visible: 1 chunk completion (the merged result)
+    assert len(chunk_done_args) == 1
+    assert chunk_done_args[0] == (0, 1, 4)
+    assert len(result["nodes"]) == 4