Predict query result sizes instead of re-serializing per candidate

[admin-raintree]

Broad context queries now check most candidates against the token budget without re-serializing the whole payload, and the output is unchanged.

Problem

Once the budget filled before limit was reached, query_context called _measure(payload) for every remaining ranked group. Each call ran json.dumps on the whole payload, so selection took O(groups × payload size). On a 3,000-document project, the query "retry backoff" took about 2.6s, roughly half of it in json.dumps.

Change

• Each candidate result is serialized once. Its exact size is the running total plus the result's bytes, the comma, and any change in the number of digits in omitted_passages.
• _settle repeats _measure's fixed-point loop for used_tokens using integers only.
• A candidate that won't fit is skipped without a full serialization. A candidate that will fit is still confirmed by _measure.
• A rejected candidate still leaves the same used_tokens value the old loop left behind. That value can affect the next measurement near a power of ten, such as 999 versus 1001, so keeping it keeps the payload byte-identical.
• The canonical serialization, the duplicate-citation loop and the final budget check are unchanged.

Checks

• uv run pytest tests/test_context_query.py -q: 18 passed.
• New test test_broad_query_matches_full_measurement: the budget fills before limit on a 200-passage query. The test forces every candidate through the original _measure path and asserts that the serialized payloads are identical.
• Local equivalence sweep, not committed: on 3,000 documents, old and new outputs were identical across 42 budgets, including 999, 1000, 1001, 9999, 10000 and 10001. On that synthetic project, the query took 0.21s with the old loop and 0.08s with the new one.
• ruff check, ruff format --check and mypy pass on the changed files. bun run validate:full was not run.

🤖 Generated with Claude Code