From 02fac7d9ccca6f3efb22f74a9a4894c93920a6dc Mon Sep 17 00:00:00 2001
From: Synvoya <16019863+Synvoya@users.noreply.github.com>
Date: Wed, 22 Jul 2026 22:12:36 +1000
Subject: [PATCH] fix(llm): keep label retries above token floor

---
 graphify/llm.py           |  5 +++--
 tests/test_label_retry.py | 25 +++++++++++++++++++++++++
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/graphify/llm.py b/graphify/llm.py
index 6c0602358..a945e41b9 100644
--- a/graphify/llm.py
+++ b/graphify/llm.py
@@ -2800,8 +2800,9 @@ def _label_batch_with_retry(
     # 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))
+    # batch (#1690). Keep at least 512 tokens as adaptive retries shrink the
+    # batch; otherwise the recovery path starves its own base case (#2086).
+    max_tokens = _resolve_max_tokens(max(512, 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
diff --git a/tests/test_label_retry.py b/tests/test_label_retry.py
index 8baae81f7..8a86ef10b 100644
--- a/tests/test_label_retry.py
+++ b/tests/test_label_retry.py
@@ -44,3 +44,28 @@ def fake_call_llm(prompt: str, **_kwargs) -> str:
 
     assert result == {42: "Label 42", 99: "Label 99", 137: "Label 137", 201: "Label 201"}
     assert call_count["n"] >= 2
+
+
+def test_label_batch_retry_keeps_a_safe_output_token_floor(monkeypatch):
+    """Split retries must not starve small batches of output tokens (#2086)."""
+    token_budgets: list[int] = []
+    monkeypatch.delenv("GRAPHIFY_MAX_OUTPUT_TOKENS", raising=False)
+
+    def fake_call_llm(prompt: str, **kwargs) -> str:
+        token_budgets.append(kwargs["max_tokens"])
+        cids = [int(m) for m in re.findall(r"Community (\d+):", prompt)]
+        if len(cids) > 1:
+            return "not json"
+        return json.dumps({str(cids[0]): f"Label {cids[0]}"})
+
+    monkeypatch.setattr(llm_mod, "_call_llm", fake_call_llm)
+
+    result = llm_mod._label_batch_with_retry(
+        [1, 2],
+        ["Community 1: auth", "Community 2: billing"],
+        backend="gemini",
+        model=None,
+    )
+
+    assert result == {1: "Label 1", 2: "Label 2"}
+    assert token_budgets == [512, 512, 512]