From 57b3d64d42ef1faa18475a7b0e56bd8d48d91705 Mon Sep 17 00:00:00 2001
From: admin-raintree <277948009+admin-raintree@users.noreply.github.com>
Date: Mon, 12 Oct 2026 13:22:29 -0700
Subject: [PATCH] Predict query result sizes instead of re-serializing per
 candidate

Once the token budget fills, query_context measured the whole payload for
every remaining ranked group, making selection O(groups x payload size).
Each candidate is now serialized once and its exact size predicted from a
running total; _measure only confirms candidates that fit. Rejected
candidates still leave the same used_tokens value the old loop did, so the
payload stays byte-identical at power-of-ten boundaries.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/docpull/context_query.py | 28 +++++++++++++++++++++++++++-
 tests/test_context_query.py  | 16 +++++++++++++++-
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/src/docpull/context_query.py b/src/docpull/context_query.py
index b2ab3f6..34a0f60 100644
--- a/src/docpull/context_query.py
+++ b/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:
diff --git a/tests/test_context_query.py b/tests/test_context_query.py
index 84d5db5..802b488 100644
--- a/tests/test_context_query.py
+++ b/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)