Fix: salvage truncated labeling replies and account for labeling token cost (#1690, #1694)

Two related fixes in the community-labeling path:

#1690 (thanks @vdgbcrypto): a truncated or slightly malformed reply no
longer discards the whole batch with "Expecting value: line 1 column 6".
`_parse_label_response` now salvages the complete `"id": "name"` pairs from
a reply that failed a strict `json.loads` (e.g. one truncated mid-object),
raising only when no pairs can be recovered. The per-batch token budget was
also raised (256 + 48*n, was 64 + 24*n) so models that prepend a short
preamble have headroom to finish the JSON. The exact provider truncation
could not be reproduced without a live key; the parser and budget address
the mechanism.

#1694 (thanks @sub4biz): cluster-only mode reported a hardcoded
`0 input * 0 output` token cost because the labeling LLM calls were never
accounted for. `_call_llm` now accumulates per-response usage into an
optional accumulator threaded through the labeling path and surfaced in
GRAPH_REPORT.md. Backends that do not return usage (the Claude Code CLI)
contribute nothing, which is honest rather than estimated.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

CHANGELOG.md

@@ -7,6 +7,8 @@ Full release notes with details on each version: [GitHub Releases](https://githu
 - Fix: `--update`-style section writes to `CLAUDE.md`/`AGENTS.md` no longer corrupt or drop content (#1688, thanks @bdfinst). `_replace_or_append_section` located its managed block by substring (`marker in content`) and `next(... if marker in line)`, so a heading that appeared as a substring of another line (or duplicate headings) matched the wrong offset and the rewrite could truncate the file. It now matches the section heading exactly (`line.strip() == marker`), appends when absent, and prefers the last exact match when several exist, so unrelated content is preserved.
 - Fix: token estimation no longer crashes on files containing tiktoken special-token text like `<|endoftext|>` (#1685, thanks @Kyzcreig). `_TOKENIZER.encode(content)` raises `ValueError` by default when the text contains a special token, which aborted packing on docs/corpora that merely mention these strings. Both `encode` sites now pass `disallowed_special=()` so such text is tokenized as ordinary bytes.
 - Fix: the Ollama backend no longer multiplies a hang by the retry count (#1686, thanks @Kyzcreig). A stalled local model would wedge for `timeout * (max_retries + 1)`, which with the default 6 retries turned one long stall into a very long one. Ollama now defaults to zero client-side retries (a local model that stalls will not un-stall on retry); set `GRAPHIFY_MAX_RETRIES` to opt back in. Other backends are unchanged. Note: the underlying stall is non-deterministic and driven by the model server, so this bounds the wait rather than eliminating the hang.
+- Fix: a truncated or slightly malformed community-labeling reply no longer discards the whole batch (#1690, thanks @vdgbcrypto). `_parse_label_response` now salvages the complete `"id": "name"` pairs from a reply that failed a strict `json.loads` (e.g. a reply truncated mid-object), raising only when no pairs can be recovered. The per-batch token budget was also raised (`256 + 48*n`, was `64 + 24*n`) to give models that prepend a short preamble enough headroom to finish the JSON. The exact provider truncation in the report could not be reproduced without a live key; the parser and budget fixes address the mechanism.
+- Fix: cluster-only mode now reports the real token cost of community labeling instead of a hardcoded zero (#1694, thanks @sub4biz). The labeling LLM calls were never accounted for, so `GRAPH_REPORT.md`'s "Token cost" line always read `0 input · 0 output` in cluster-only runs. `_call_llm` now accumulates per-response usage into an optional accumulator that is threaded through the labeling path and surfaced in the report. Backends that do not return usage (the Claude Code CLI) still contribute nothing, which is honest rather than estimated.
 
 ## 0.9.7 (2026-07-06)
 

graphify/__main__.py

@@ -3560,6 +3560,10 @@ def main() -> None:
                 }
             except Exception:
                 existing_labels = {}
+        # Accumulate token usage from the labeling LLM calls so cluster-only mode
+        # reports real cost instead of a hardcoded zero (#1694). Stays {0, 0} on
+        # the reuse / no-label paths, which make no LLM calls.
+        label_token_usage = {"input": 0, "output": 0}
         if labels_path.exists() and not force_relabel:
             # Reuse saved labels, but don't blindly trust them: the graph may have
             # been re-scoped/re-clustered since labeling, in which case a cid now
@@ -3643,6 +3647,7 @@ def main() -> None:
             generated_labels, _ = generate_community_labels(
                 G, label_communities_input, backend=label_backend, model=label_model, gods=gods,
                 max_concurrency=label_max_concurrency, batch_size=label_batch_size,
+                usage_out=label_token_usage,
             )
             # Only let the LLM OVERRIDE where it produced a real name — its no-backend
             # fallback returns "Community {cid}" placeholders, which must not clobber
@@ -3653,7 +3658,7 @@ def main() -> None:
             })
         stages.mark("label")
         questions = suggest_questions(G, communities, labels)
-        tokens = {"input": 0, "output": 0}
+        tokens = label_token_usage
         from graphify.export import _git_head as _gh
         _commit = _gh()
         from graphify.report import load_learning_for_report as _llfr

graphify/llm.py

@@ -1949,9 +1949,15 @@ def _call_llm(
     backend: str,
     max_tokens: int = 200,
     model: str | None = None,
+    usage_out: dict | None = None,
 ) -> str:
     """Send a plain-text prompt to `backend` and return the model's text reply.
 
+    When ``usage_out`` is provided it is accumulated in place with ``input`` and
+    ``output`` token counts from the response, so callers (community labeling)
+    can total the cost of otherwise-uninstrumented LLM calls (#1694). Existing
+    callers that omit it are unaffected.
+
     Used by lightweight callers (e.g. `graphify.dedup` LLM tiebreaker) that
     don't need the full extraction prompt or JSON-shaped output. Mirrors the
     backend dispatch logic of `extract_files_direct` but skips the
@@ -1975,6 +1981,11 @@ def _call_llm(
         )
     mdl = model or _default_model_for_backend(backend)
 
+    def _rec(inp, out) -> None:
+        if usage_out is not None:
+            usage_out["input"] = usage_out.get("input", 0) + int(inp or 0)
+            usage_out["output"] = usage_out.get("output", 0) + int(out or 0)
+
     if backend == "claude":
         try:
             import anthropic
@@ -1986,6 +1997,9 @@ def _call_llm(
             max_tokens=max_tokens,
             messages=[{"role": "user", "content": prompt}],
         )
+        u = getattr(resp, "usage", None)
+        if u is not None:
+            _rec(getattr(u, "input_tokens", 0), getattr(u, "output_tokens", 0))
         return resp.content[0].text if resp.content else ""
 
     if backend == "claude-cli":
@@ -2019,6 +2033,14 @@ def _call_llm(
         if proc.returncode != 0:
             raise RuntimeError(f"claude -p exited {proc.returncode}: {proc.stderr.strip()[:500]}")
         envelope = _claude_cli_envelope(proc.stdout)
+        cli_usage = envelope.get("usage") or {}
+        if cli_usage:
+            _rec(
+                (cli_usage.get("input_tokens", 0) or 0)
+                + (cli_usage.get("cache_read_input_tokens", 0) or 0)
+                + (cli_usage.get("cache_creation_input_tokens", 0) or 0),
+                cli_usage.get("output_tokens", 0),
+            )
         return envelope.get("result", "")
 
 
@@ -2036,6 +2058,9 @@ def _call_llm(
             messages=[{"role": "user", "content": [{"text": prompt}]}],
             inferenceConfig=_bedrock_inference_config(max_tokens, mdl),
         )
+        bu = resp.get("usage") or {}
+        if bu:
+            _rec(bu.get("inputTokens", 0), bu.get("outputTokens", 0))
         return resp.get("output", {}).get("message", {}).get("content", [{}])[0].get("text", "")
 
     if backend == "azure":
@@ -2056,6 +2081,9 @@ def _call_llm(
         resp = azure_client.chat.completions.create(**azure_kwargs)
         if not resp.choices or resp.choices[0].message is None:
             raise ValueError("Azure OpenAI returned empty or filtered response")
+        au = getattr(resp, "usage", None)
+        if au is not None:
+            _rec(getattr(au, "prompt_tokens", 0), getattr(au, "completion_tokens", 0))
         return resp.choices[0].message.content or ""
 
     # OpenAI-compatible (kimi, openai, gemini, ollama)
@@ -2088,6 +2116,9 @@ def _call_llm(
     resp = client.chat.completions.create(**kwargs)
     if not resp.choices or resp.choices[0].message is None:
         raise ValueError("LLM returned empty or filtered response")
+    ou = getattr(resp, "usage", None)
+    if ou is not None:
+        _rec(getattr(ou, "prompt_tokens", 0), getattr(ou, "completion_tokens", 0))
     return resp.choices[0].message.content or ""
 
 
@@ -2257,9 +2288,25 @@ def _parse_label_response(text: str, labeled_cids: list[int]) -> dict[int, str]:
         start, end = cleaned.find("{"), cleaned.rfind("}")
         if start != -1 and end > start:
             cleaned = cleaned[start:end + 1]
-    data = json.loads(cleaned)
-    if not isinstance(data, dict):
-        raise ValueError("label response is not a JSON object")
+    data: dict | None = None
+    try:
+        parsed = json.loads(cleaned)
+        if isinstance(parsed, dict):
+            data = parsed
+    except (json.JSONDecodeError, ValueError):
+        data = None
+    if data is None:
+        # Salvage: pull the complete "<cid>": "<name>" pairs directly. A model
+        # can truncate its reply mid-object (a stingy token budget or a preamble
+        # eating the completion), which used to hard-fail the whole batch with
+        # e.g. `Expecting value: line 1 column 6` on a `{"0":` fragment (#1690).
+        # Recovering the pairs that DID arrive labels those communities instead
+        # of dropping the entire batch to placeholders.
+        pairs = re.findall(r'"?(-?\d+)"?\s*:\s*"([^"\\]*(?:\\.[^"\\]*)*)"', cleaned)
+        if pairs:
+            data = {k: v for k, v in pairs}
+        else:
+            raise ValueError(f"label response is not parseable JSON: {text[:120]!r}")
     out: dict[int, str] = {}
     for cid in labeled_cids:
         name = data.get(str(cid))
@@ -2278,6 +2325,7 @@ def _label_batch_with_retry(
     model: str | None,
     depth: int = 0,
     max_depth: int = 3,
+    usage_out: dict | None = None,
 ) -> dict[int, str]:
     """Label a batch of communities, splitting in half and retrying on parse failure.
 
@@ -2301,10 +2349,18 @@ def _label_batch_with_retry(
         "Respond ONLY with a JSON object mapping the community id (as a string) to "
         "its name - no prose, no markdown fences.\n\n" + "\n".join(batch_lines)
     )
-    max_tokens = _resolve_max_tokens(min(64 + 24 * len(batch_cids), 8192))
+    # Budget generously: a 2-5 word name is ~10 tokens, but models (notably
+    # gemini) often prepend a short preamble or reasoning that eats the
+    # completion and truncates the JSON mid-object, which used to fail the whole
+    # batch (#1690). The old 64 + 24*n floor left no headroom.
+    max_tokens = _resolve_max_tokens(min(256 + 48 * len(batch_cids), 8192))
     call_kwargs: dict = {"backend": backend, "max_tokens": max_tokens}
     if model is not None:
         call_kwargs["model"] = model
+    # Only forward usage_out when the caller wants accounting, so existing
+    # callers (and their test doubles) see the unchanged _call_llm signature.
+    if usage_out is not None:
+        call_kwargs["usage_out"] = usage_out
 
     try:
         text = _call_llm(prompt, **call_kwargs)
@@ -2325,10 +2381,12 @@ def _label_batch_with_retry(
         left = _label_batch_with_retry(
             batch_cids[:mid], batch_lines[:mid],
             backend=backend, model=model, depth=depth + 1, max_depth=max_depth,
+            usage_out=usage_out,
         )
         right = _label_batch_with_retry(
             batch_cids[mid:], batch_lines[mid:],
             backend=backend, model=model, depth=depth + 1, max_depth=max_depth,
+            usage_out=usage_out,
         )
         return left | right
 
@@ -2344,6 +2402,7 @@ def label_communities(
     top_k: int = _LABEL_TOP_K,
     batch_size: int = _LABEL_BATCH_SIZE,
     max_concurrency: int = 4,
+    usage_out: dict | None = None,
 ) -> dict[int, str]:
     """Return a complete ``{cid: name}`` map using ``backend`` for naming.
 
@@ -2386,19 +2445,30 @@ def label_communities(
     def _run_batch(batch_idx: int):
         start = batch_idx * batch_size
         end = min(start + batch_size, len(labeled_cids))
+        # Accumulate token usage into a per-batch dict so concurrent workers
+        # never race on the shared accumulator; it is merged on the main thread
+        # in _merge (#1694).
+        batch_usage: dict = {} if usage_out is not None else None
+        batch_kwargs = {"usage_out": batch_usage} if usage_out is not None else {}
         try:
             parsed = _label_batch_with_retry(
                 labeled_cids[start:end], lines[start:end], backend=backend, model=model,
+                **batch_kwargs,
             )
-            return batch_idx, parsed, None
+            return batch_idx, parsed, None, batch_usage
         except Exception as exc:  # noqa: BLE001 - reported per-batch; surfaced below
-            return batch_idx, None, exc
+            return batch_idx, None, exc, batch_usage
 
     written = 0
     errors: dict[int, Exception] = {}
 
-    def _merge(batch_idx: int, parsed, exc) -> None:
+    def _merge(batch_idx: int, parsed, exc, batch_usage=None) -> None:
         nonlocal written
+        # Count tokens even for a failed batch: the LLM call was billed whether
+        # or not the reply parsed.
+        if usage_out is not None and batch_usage:
+            usage_out["input"] = usage_out.get("input", 0) + batch_usage.get("input", 0)
+            usage_out["output"] = usage_out.get("output", 0) + batch_usage.get("output", 0)
         if exc is not None:
             errors[batch_idx] = exc
             start = batch_idx * batch_size
@@ -2440,6 +2510,7 @@ def generate_community_labels(
     quiet: bool = False,
     max_concurrency: int = 4,
     batch_size: int = _LABEL_BATCH_SIZE,
+    usage_out: dict | None = None,
 ) -> tuple[dict[int, str], str]:
     """CLI entry point: resolve a backend, name communities, and degrade to
     ``Community N`` placeholders on any failure (no backend, API error, malformed
@@ -2462,6 +2533,7 @@ def generate_community_labels(
         labels = label_communities(
             G, communities, backend=backend, model=model, gods=gods,
             max_concurrency=max_concurrency, batch_size=batch_size,
+            usage_out=usage_out,
         )
         return labels, "llm"
     except Exception as exc:

tests/test_labeling.py

@@ -82,7 +82,7 @@ def test_label_cli_passes_model_override(tmp_path, monkeypatch):
     captured = {}
 
     def fake_generate(G, communities, *, backend=None, model=None, gods=None,
-                      quiet=False, max_concurrency=4, batch_size=100):
+                      quiet=False, max_concurrency=4, batch_size=100, usage_out=None):
         captured["backend"] = backend
         captured["model"] = model
         captured["max_concurrency"] = max_concurrency
@@ -143,7 +143,7 @@ def test_label_cli_missing_only_preserves_existing_labels(tmp_path, monkeypatch)
     captured = {}
 
     def fake_generate(G, communities, *, backend=None, model=None, gods=None,
-                      quiet=False, max_concurrency=4, batch_size=100):
+                      quiet=False, max_concurrency=4, batch_size=100, usage_out=None):
         captured["communities"] = dict(communities)
         return {1: "Payment Flow"}, "llm"
 
@@ -410,3 +410,67 @@ def test_label_communities_forces_serial_for_ollama(monkeypatch):
     monkeypatch.delenv("GRAPHIFY_OLLAMA_PARALLEL", raising=False)
     label_communities(G, communities, backend="ollama", batch_size=1, max_concurrency=8)
     assert state["peak"] == 1, "ollama must be forced serial"
+
+
+def test_label_communities_salvages_truncated_reply(monkeypatch):
+    # #1690: a reply truncated mid-object (a stingy token budget or model
+    # preamble) used to hard-fail the whole batch with `Expecting value: line 1
+    # column 6`. The complete pairs that arrived are now salvaged.
+    G, communities = _graph()
+    monkeypatch.setattr(
+        "graphify.llm._call_llm",
+        lambda p, *, backend, max_tokens=200: '{"0": "Order Management", "1":',
+    )
+    labels = label_communities(G, communities, backend="gemini")
+    assert labels[0] == "Order Management"   # salvaged
+    assert labels[1] == "Community 1"         # truncated cid falls back to placeholder
+
+
+def test_label_communities_accumulates_token_usage(monkeypatch):
+    # #1694: cluster-only mode reported zero labeling cost because token usage
+    # from the naming LLM calls was never accumulated. label_communities now
+    # fills a caller-supplied usage_out accumulator, summed across all batches.
+    G, communities = _many_communities(6)
+
+    def fake_call(prompt, *, backend, max_tokens=200, usage_out=None):
+        if usage_out is not None:
+            usage_out["input"] = usage_out.get("input", 0) + 100
+            usage_out["output"] = usage_out.get("output", 0) + 10
+        # one name per community id present in this batch
+        cids = [int(line.split()[1].rstrip(":")) for line in prompt.splitlines()
+                if line.startswith("Community ")]
+        return json.dumps({str(c): f"Name {c}" for c in cids})
+
+    monkeypatch.setattr("graphify.llm._call_llm", fake_call)
+    usage = {"input": 0, "output": 0}
+    # batch_size=2 -> 3 batches, run serially so the count is deterministic
+    labels = label_communities(
+        G, communities, backend="gemini", batch_size=2, max_concurrency=1,
+        usage_out=usage,
+    )
+    assert len(labels) == 6
+    assert usage == {"input": 300, "output": 30}  # 3 batches * (100, 10)
+
+
+def test_label_communities_counts_tokens_for_failed_batch(monkeypatch):
+    # A batch whose reply can't be parsed was still billed by the provider, so
+    # its tokens must be counted even though it contributes no label (#1694).
+    G, communities = _graph()
+
+    def fake_call(prompt, *, backend, max_tokens=200, usage_out=None):
+        if usage_out is not None:
+            usage_out["input"] = usage_out.get("input", 0) + 50
+            usage_out["output"] = usage_out.get("output", 0) + 5
+        return "not json at all"
+
+    monkeypatch.setattr("graphify.llm._call_llm", fake_call)
+    usage = {"input": 0, "output": 0}
+    # single community -> no split retry; the only batch fails to parse, so
+    # label_communities re-raises (every batch failed) after counting tokens.
+    G2 = nx.Graph()
+    G2.add_node("a", label="alpha")
+    with pytest.raises((ValueError, json.JSONDecodeError)):
+        label_communities(
+            G2, {0: ["a"]}, backend="gemini", usage_out=usage,
+        )
+    assert usage == {"input": 50, "output": 5}