feat(budget): measure the token margin against the provider instead of guessing 50% high

[arhyneRWU]

Summary

Item 3 of the plan you set out in #629: "For hosted models, use provider-native counting where it exists, falling back to the current margin where it doesn't." Items 1 and 2 landed as #648/#719 and #774.

DEFAULT_TOKEN_SAFETY_MARGIN is 1.5 because that covers the densest content anyone measured — a currency-dense budget table diverging ~45% from cl100k_base. Applied to everything, a hosted model answering prose plans against a third less window than it has, and crosses the #632 routing threshold early. Anthropic and Google publish exact counts, so where one is available the margin is now measured instead of assumed.

It is a ratio, not a substituted total, and the margin does not become 1.0. I said 1.0 in the first draft of #816 and that was wrong — three things in the existing code prevent it:

1. The margin is applied per component (estimate_input_tokens:795-806), and BudgetPlan.total_input_tokens:567 is a computed property summing those parts. There is no slot for a total that isn't the sum of its components.
2. The planner recounts on mutated text inside its trim loop (:981-1052). A network counter can't live there, and a count of the original text is void once the text is cut.
3. A pre-plan count doesn't measure what's sent — the wire payload is assembled after planning, which is what REQUEST_SCAFFOLD_TOKENS exists to absorb.

So: count the assembled prompt, count the same payload locally, and use max(1.0, provider / local) as that request's margin. A ratio stays approximately valid as content is trimmed where a total does not.

Changes

backend/app/services/native_token_count.py (new) — fetches a provider's count, or reports plainly that it couldn't. Gated on protocol before any I/O, since OpenAIModel.count_tokens raises and every OpenAI-compatible model would otherwise decrypt a key and build a client per turn to learn that. Five-second timeout, deliberately not the request's own settings — build_thinking_model_settings sets 120s and the shared client's RateLimitRetryTransport retries a 429 six times honouring Retry-After up to 60s, which would put a two-minute stall in front of every message. Blanket except by necessity: Anthropic raises ModelHTTPError, Google's genai errors are unwrapped, a misconfigured protocol reaches OpenAIModel; there's no complete and stable union of those. CancelledError passes through, being the caller leaving rather than a failure.

backend/app/services/context_budget.py — NativeCount plus a rung between the tiktoken-is-exact check and the 1.5 default, so it only ever displaces the guess. Rejected on model-name mismatch, zero baseline, non-positive count, or non-finite ratio, each falling through to the ladder. Clamped at 1.0: a provider counting below tiktoken may reduce over-inflation, never reclaim window. The "estimated, not exact" warning stays quiet only when a measurement actually supplied the margin. No new imports — the module stays a dependency leaf.

backend/app/services/model_routing.py — _sized_for re-derived the current model's margin to restate a request in a candidate's units. Handed a natively-counted number it would still answer 1.5 and divide by a factor never applied, understating the request by a third exactly where the router decides whether a candidate can hold it — the #648 defect at the routing boundary, which that function's docstring says it exists to prevent. Callers now state the margin their number was measured with.

backend/app/services/chat_service.py — the wiring. _build_chat_prompt is the existing assembly extracted verbatim, so the pre-flight counts the request chat is about to make rather than a reconstruction. The payload counted is the whole prompt including reference documents: in document chat the documents are the payload, and they're the digit-dense content the 1.5 was sized for — sampling question and history alone would measure prose and apply the answer to a budget table.

Three things I'd want you to push back on

The pre-flight is unconditional. One serialized round trip ahead of the first token on every eligible turn. Gating it on "the estimate is near a boundary where the answer could change" would be cheaper, but it makes is this count exact? depend on request size, which is hard to reason about later. I have no latency data for either count endpoint — if you do, that's the number that should decide this.

The reach is narrower than "Claude and Gemini". detect_api_protocol:372 routes a bare claude-* name to the OpenAI-compatible path, and the OpenRouter / Insight-AI branches produce OpenAIModel, which has no count endpoint. This fires only for models an admin explicitly configured api_protocol: anthropic or google with a direct provider key — which the ModelEditor presets do create. I deliberately did not add a symmetric auto-route for claude-*: that would change which client every existing Claude deployment uses, a behaviour change well beyond token counting.

Google's count omits the system instruction. pydantic-ai attaches system_instruction only for providers other than google-gla, and an api-key GoogleProvider is google-gla. Rather than fudge it, the system prompt is excluded from both sides of the ratio. Counting it on one side only would depress the ratio for a reason unrelated to tokenizer divergence — and a depressed ratio clamps to margin 1.0 over a tiktoken figure that itself under-counts these models, which is the hard-fail direction.

Also deliberately untouched: stored_count_margin. It answers "how safe is a count someone else took earlier, when I don't have the text", and a provider count requires the text — so find_oversize_documents and find_context_overflow keep the 1.5 allowance.

Test Plan

• Shared checks pass (make ci) — make backend-ci locally: 4438 passed, 161 skipped, coverage 64.38% against the 50% gate, tier-1 integration 14 passed. Baseline on this branch point was 4410 passed; +28 accounts exactly for the tests added. I could not run make frontend-ci locally (npm is not on my PATH) and said so when opening this; CI's Frontend job has since run and passed, along with the other ten checks. No frontend/ file is touched by this PR.
• Release check (make release-check) — not applicable. No packaging or deployment path changed.
• Manually tested the affected feature(s) — not possible here, and this is the PR's main limitation. This deployment counts every model exactly from a local tokenizer.json, so nothing is on the guessed rung and there is no Claude or Gemini model to exercise. No live provider count endpoint was called at any point. What is covered, against stubbed providers: a good Anthropic count, a good Google count, NotImplementedError from an OpenAIModel, timeout, ModelHTTPError/429, a malformed usage object, and every ineligible protocol asserting that no client is constructed at all. Live verification needs a deployment holding direct provider keys.
• Updated CHANGELOG.md for user-facing or operator-facing changes
• Updated deploy/release docs — not applicable. No install or release path changed.

Related Issues

Closes #816.

Adjacent, filed separately and not addressed here: #817 — update_model replaces the whole model dict, so token_safety_margin, tokenizer_path and tokenizer_cache_root are deleted by the next save from the admin form. It constrains this work (any per-model opt-in would be wiped), which is why the pre-flight is gated on api_protocol — an existing field the form does manage — rather than on a new one.

🤖 Generated with Claude Code

https://claude.ai/code/session_01HucxUJ7gaA6iKveF8aruW2

[ViaJables]

Reviewed and pushed one commit. This is careful work — the fail-closed design is genuinely thorough, the blast radius is narrow (chat_service is the only caller of estimate_input_tokens/plan_and_compact_context, so extraction and workflows take an identical path), and the routing change is right: _sized_for re-deriving 1.5 for a natively-counted number would understate the request by a third at exactly the boundary where the router decides whether a candidate can hold it. I verified the pydantic-ai signatures against 1.71.0, including the google-gla system-instruction omission.

What I added: a circuit breaker. The module coalesces the failure log but not the failure work. A deployment whose egress to the count endpoint is slow, or whose org is rate-limited there, pays the full 5s timeout on every chat turn ahead of the first token, permanently — and after the first occurrence says so only at DEBUG. No metric, no way out short of a restart.

Three consecutive failures now pause a model for five minutes. The check sits after the protocol gate (so an ineligible model never enters the circuit) and before any model construction, key decryption or request (so an open circuit costs nothing). One request goes through when the cooldown elapses; any success clears the record, so an outage cannot disable counting until the next deploy. Per model name, monotonic clock.

Skipping is strictly better than blocking here — the estimate path it falls back to is the one that shipped before this feature, so a paused model gets the same budget it would have got anyway, without the wait.

Two changelog corrections, both things the entry asserted that the code does not do:

• "fires only for models an admin explicitly configured" is not true of Gemini. detect_api_protocol routes any name containing gemini to google on its own (llm_service.py:377), so a deployment already running one gets the pre-flight on upgrade without changing anything. Issue [Feature] Provider-native token counting for hosted models — Claude and Gemini fall to the 1.5 guessed margin #816 stated this correctly — the entry lost it.
• "one unconditional round trip" is no longer true, which is the point of the commit.

Your existing failure-logging tests needed updating rather than working around: _CIRCUIT is per-process state exactly like _UNAVAILABLE_LOGGED, so it joins that autouse reset — otherwise whether a test sees an open circuit depends on collection order. I also extended the repeated-failure test from 4 turns to 50, so it pins the bound rather than an arithmetic coincidence.

Two things I did not change

The payload is uploaded twice per turn. _build_chat_prompt is passed the pre-compaction segments, so a 150k-token document chat POSTs ~600KB to count_tokens and then the same ~600KB to /v1/messages. That is correct for measurement — the documents are the digit-dense content the 1.5 was sized for — but it doubles egress and pre-request latency on exactly the largest requests, which is also where the 5s budget is most likely to blow. Worth considering whether a sampled or capped payload would measure well enough.

A merge-forward hazard for major/agentic-chat. _ask_provider passes a bare ModelRequestParameters(), justified by "the chat agent registers none" — true on main, false on the v5 branch, where create_agentic_chat_agent registers the chat_tools registry. Merged forward, the provider would count a payload missing thousands of tokens of billable tool schemas, the ratio deflates toward 1.0, and the 1.5 cushion incidentally covering those schemas disappears — in the under-estimate direction. Nothing in the code or tests would flag it at merge time. Worth a guard or at least a comment before that merge.

[ViaJables]

Merged after review. Two things the review found that were not blocking, recorded here as follow-ups:

1. The pre-flight counts the whole uncompacted payload on every eligible turn (chat_service.py, the _measure_native_count call around the prompt build). That is a second full upload of every document segment ahead of the first token, and for a payload already above the model's window the count endpoint may reject it — three such turns open the circuit for five minutes for every user of that model, which is exactly the large-document case the margin matters for. Worth skipping the pre-flight when the local estimate already fits with headroom or exceeds the window outright, or capping the counted payload and applying the ratio.
2. native_token_count._ask_provider builds a fresh provider client (and decrypts the key) per count; passing the model instance the chat is about to build would halve that.

Committed on the branch before merge: the boolean-margin warning in model_routing._measured_margin now logs the offending value instead of None, and the CHANGELOG entry was moved back under [Unreleased] (git had auto-merged it into the released v4.12.0 section).