Predict query result sizes instead of re-serializing per candidate

src/docpull/context_query.py

@@ -71,6 +71,15 @@ def _measure(payload: dict[str, Any]) -> int:
         payload["budget"]["used_tokens"] = size
 
 
+def _settle(size_without_used: int, used: int) -> int:
+    # Integer mirror of _measure, including its dependence on the starting used_tokens.
+    while True:
+        size = size_without_used + len(str(used))
+        if size == used:
+            return size
+        used = size
+
+
 def _spans(content: str) -> Iterator[tuple[int, int]]:
     """Keep paragraphs and fenced code blocks whole, with exact character offsets."""
     start = offset = 0
@@ -255,17 +264,34 @@ def query_context(
     ranked = sorted(groups.values(), key=lambda g: (-g["score"], g["text"]))
     deferred: list[tuple[dict[str, Any], list[dict[str, Any]]]] = []
     payload["omitted_passages"] = len(ranked)
+    used = payload["budget"]
+    rest = len(serialize_query_result(payload).encode("utf-8")) - len(str(used["used_tokens"]))
     for group in ranked:
         citations = sorted(group["citations"].values(), key=lambda c: c["id"])
         result = {"text": group["text"], "citations": citations[:1], "omitted_citations": len(citations) - 1}
         if len(payload["results"]) >= limit:
             break
+        # Predict the exact size without re-serializing the payload; _measure confirms fits.
+        omitted = payload["omitted_passages"]
+        candidate = (
+            rest
+            + len(serialize_query_result(result).encode("utf-8"))
+            + (1 if payload["results"] else 0)
+            + len(str(omitted - 1))
+            - len(str(omitted))
+        )
+        predicted = _settle(candidate, used["used_tokens"])
+        if predicted > max_tokens:
+            used["used_tokens"] = predicted  # the value a rejected _measure leaves behind
+            continue
         payload["results"].append(result)
         payload["omitted_passages"] -= 1
-        if _measure(payload) > max_tokens:
+        size = _measure(payload)
+        if size > max_tokens:
             payload["results"].pop()
             payload["omitted_passages"] += 1
         else:
+            rest = size - len(str(size))
             deferred.append((result, citations[1:]))
     # Spend remaining space on duplicate provenance after selecting distinct passages.
     for result, citations in deferred:

tests/test_context_query.py

@@ -11,7 +11,7 @@
 
 import pytest
 
-from docpull import ContextQueryError, query_context
+from docpull import ContextQueryError, context_query, query_context
 from docpull import project as project_module
 from docpull.cli import main
 from docpull.context_query import serialize_query_result
@@ -280,3 +280,17 @@ async def test_context_query_stdio_contract(tmp_path: Path, monkeypatch: pytest.
             "context_query", {"query": "retry", "project_roots": [str(tmp_path / "missing")]}
         )
         assert missing.isError is True
+
+
+def test_broad_query_matches_full_measurement(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    documents: list[tuple[str, str | None]] = [
+        (f"# Doc {i}\n\nRetry backoff café {i}: {'wait ' * (i * 37 % 90)}\n", None) for i in range(200)
+    ]
+    root = _project(tmp_path, monkeypatch, "broad", documents)
+    args = {"project_roots": [root], "max_tokens": 3000, "limit": 100}
+    fast = query_context("retry backoff", **args)
+    assert 0 < len(fast["results"]) < 100
+    assert fast["omitted_passages"] > 0
+    # Forcing every prediction to "fits" sends each candidate through the original _measure path.
+    monkeypatch.setattr(context_query, "_settle", lambda size, used: 0)
+    assert serialize_query_result(query_context("retry backoff", **args)) == serialize_query_result(fast)