Add pre-run token usage estimation

[Aaron Aspinwall (AaronAspinwall123)]

Summary

• estimate model token usage locally before a run, with a likely range and per-stage breakdown
• deliberately bias point estimates high so planning is more likely to overestimate than underestimate actual usage
• mirror runtime resume, force, artifact-cache, multi-turn, and tool-loop behavior while excluding opaque target internals
• record provider usage coverage and actual-versus-estimated accuracy in metrics.json only when the run and telemetry are complete
• show the estimate before viewer submission, then show compact estimate-versus-reported usage in the CLI, run viewer, and standalone HTML export

How token estimation works

This is a local, deterministic, stage-aware forecast, not an LLM prediction or a dry-run against the target. We tokenize what is already known and use conservative, schema-shaped projections for content that has not been generated yet. The CLI and viewer use the same implementation in assert_ai/core/token_estimator.py, with read-only run planning in assert_ai/runner.py.

1. Plan only the work that will run

Resolve the config and overrides, then apply the runtime's stage selection, artifact-cache, resume, and forced-stage rules. Reusable generation artifacts and compatible completed inference/judgments are excluded; forced or invalidated work is included. Existing test cases, taxonomies, and transcripts are used when applicable. For a first run, representative generated structures stand in for artifacts that do not exist yet. Estimation does not call a model, execute tools, or create run artifacts.

2. Count known inputs and project stage outputs

Known messages, system prompts, tool schemas, and structured-response schemas are counted locally with LiteLLM's model-aware tokenizer. Provider/model names are mapped to supported tokenizer families; unsupported names or tokenizer failures use a roughly four-characters-per-token fallback. This is token accounting, not dollar-cost estimation.

Stage	Projection method
Behavior categories (systematize)	Build the real two-step systematization/conversion prompts; tokenize representative structured outputs sized to the requested category count, subject to output limits.
Test cases (test_set)	Build generation jobs and prompts with the runtime helpers, requested sample counts, behavior categories, variations, and tool source. Include model-generated variation levels when needed. Estimate structured case outputs per job, capped by generation limits.
Execute target (inference)	Count a prompt request or simulate the configured conversation up to max_turns, including tester requests, target requests, and tool round trips. Project response lengths using the rules below.
Judge (judge)	Build the real scoring prompt and response contract around existing or projected transcript XML. Output is the larger of 512 tokens or the tokenized representative judgment, capped by the judge limit. Multiply requests and usage by judge.n.

3. Use high-side completion assumptions, not the full budget

When a target output limit B is present, final-answer projections are:

cap(B) = max(1, round(0.875 * B))
prompt output = min(cap(B), max(512, min(768, round(0.75 * B))))
scenario output per turn = min(cap(B), 384)

Without a target limit, the prompt/scenario baselines are 512/384 tokens. Tester messages use 55 tokens, capped by the tester limit; simulated tool results use 90 tokens. For example, a target limit of 256 projects a 224-token final answer, while a limit of 900 projects a 675-token prompt answer. These are deliberately conservative heuristics, calibrated against the saved measurements below, not learned predictions of the exact answer length.

4. Carry projected content into later requests

Each projected tester/target message is appended to conversation history before estimating the next request. Earlier output therefore also becomes input on later turns and in the judge transcript; usage is not just a single-call estimate multiplied by turn count.

Tool-enabled targets assume one tool round trip per target turn: count a schema-shaped tool-call payload separately from one final answer, add the tool result, and retain that history in later target requests and judge input. Simulated tools add their own model call, using the same prompt builder as execution, including the conversation and previous tool results. Projected judge XML uses the runtime formatter and its message-length truncation.

5. Aggregate and report uncertainty

For larger case/job lists, count up to 24 evenly spaced profiles/jobs per group and scale their mean usage to the full count. Sum projected input and output tokens across pending calls and stages. The usual displayed range is 65%-135% of the point estimate: a heuristic planning band, not a statistically fitted confidence interval or guaranteed upper bound. High-side bias comes from the output assumptions, not a blanket multiplier on every known input.

Retries, extra tool calls, provider-added web-search context, and hidden provider overhead are not included. Model calls inside opaque callable, connector, endpoint, or sandbox targets are explicitly excluded; owned tester/judge calls can still be estimated. Notes surface these limitations and any unavailable stage estimate.

6. Compare with provider usage after completion

Keep the pre-run forecast and aggregate provider-reported input/output usage separately. CLI and viewer show the signed difference (actual - estimate) / estimate only when execution and tracked-call telemetry are complete; incomplete totals are labeled reported/partial rather than presented as a reliable accuracy result. The paired calibration section below explains how the latest algorithm was replayed against the same saved usage without fresh model calls.

Viewer

The Summary & submit step now estimates usage before the evaluation starts. The compact row shows the conservative total, likely range, and tracked call count. It uses the same normalized payload as submission, makes no provider calls, and does not reserve a run directory. Submission waits while the estimate is loading; superseded or abandoned requests terminate their estimator process and clean up the temporary config.

Completed runs now use a single compact Token usage row instead of the previous three-card layout. It shows:

• estimated total and likely range
• provider-reported total, input/output split, call coverage, and cached-input rate
• signed actual-versus-estimate difference and an in-range/out-of-range badge when telemetry is complete
• stage estimates and estimator notes under collapsed Details
• partial provider totals labeled Reported, with the incomplete-coverage explanation kept visible

Older token metrics without the newer accuracy status field remain readable. Runs without token telemetry show a small unavailable state rather than a large empty card.

CLI

assert-ai estimate --config PATH exposes the same read-only estimate without starting a run; --output json is used by the viewer endpoint. assert-ai run prints the estimate before uncached stages execute and reports provider usage plus estimate accuracy after completion. Zero tracked-token estimates still emit explicit text and preserve notes for opaque callable, endpoint, connector, or sandbox usage.

Validation

• full Python suite: 1,516 passed and 30 skipped across the main run plus isolated reruns for dotenv and installed-package metadata tests
• post-merge model/token/cache/viewer regression set: 197 passed, 6 skipped
• npm run check (0 errors; 6 existing warnings in unrelated files)
• npm run build
• live /api/runs/estimate check: 83,242 estimated tokens, likely 54,107-112,377, 28 tracked calls, and no suite directory created
• cancellation regression verifies an estimator that ignores the first termination signal is forcibly closed before temporary-config cleanup

Accuracy review: estimate-only replay

The revised estimator is closer on all 12 original scenarios, while keeping 12/12 point estimates above saved actual usage and 12/12 actual totals inside the displayed likely range. Both CLI and viewer receive these changes through the shared estimator.

Calculation changes

• Count a schema-shaped tool-call payload separately from the final answer instead of allocating two full answers. Retain tool calls/results in subsequent target requests and projected judge transcripts.
• Use the runtime simulator prompt builder, including conversation and previous tool results, and the runtime XML formatter, including message truncation.
• Recalibrate heuristics: small target completion budgets use 87.5% rather than assuming 100% utilization; reduce the judge floor from 640 to 512 tokens while retaining contract-size scaling and output-limit caps. Larger prompt projections retain the existing high-side baseline and budget scaling. No scenario-specific rules or global input-token scaling were added.

Paired results

Method: estimate-only replay against saved provider usage; no fresh model calls. Used the same 12 azure/gpt-5.4-mini scenarios, fixture files, and original per-case overrides. The previous estimator reproduced every original total exactly before changes; fixture SHA-256 fingerprints remained unchanged. Ground truth is provider-reported token usage, not evaluation scores. Prediction runs created no run artifacts; calibration data stays local and uncommitted.

Metric	Previous	Updated
Mean absolute percentage error, actual denominator	26.6%	20.5%
Mean absolute percentage error, estimate denominator (original PR metric)	20.0%	16.1%
Weighted absolute percentage error, actual denominator	23.8%	18.4%
Conservative point estimates	12/12	12/12
Actual totals inside displayed range	12/12	12/12
Estimated call counts matching saved actual counts	12/12	12/12

Scenario	Previous estimate	Updated estimate	Saved actual	Updated overestimate
prompt-tiny	3,119	2,975	2,477	+20.1%
prompt-short	3,515	3,323	2,542	+30.7%
prompt-medium	3,695	3,503	2,954	+18.6%
prompt-long	4,237	4,045	3,951	+2.4%
prompt-long-output	4,453	4,325	3,964	+9.1%
prompt-batch	14,034	13,266	10,678	+24.2%
scenario-one-turn	4,948	4,756	3,781	+25.8%
scenario-two-turn	7,666	7,346	5,508	+33.4%
scenario-four-turn	15,118	14,350	10,066	+42.6%
scenario-long	11,381	10,869	10,517	+3.3%
tool	5,160	4,698	3,613	+30.0%
full-pipeline	19,928	19,544	18,489	+5.7%

Updated overestimate is (estimate - actual) / actual. Standard MAPE uses abs(estimate - actual) / actual; the original PR metric instead divides by the estimate. Weighted error is sum(abs(estimate - actual)) / sum(actual). Standard MAPE improved by about 23% relative.

This is the same small calibration set, not independent evidence of accuracy on unseen workloads. Point estimates remain planning heuristics, not guaranteed upper bounds. Additional tool calls, retries, unusually long outputs, hidden provider overhead, and opaque target internals remain limitations.

Regression results

264 passed, 6 skipped. Re-ran the same model-client, estimator, runner-usage, stage-filter, artifact-cache, CLI, viewer-spawn, and viewer-token-usage selectors, plus existing runtime-mode/inference-stage coverage for the shared simulator prompt builder. Added deterministic regressions for completion-budget headroom, judge contract scaling/caps, tool payload accounting, accumulated tool history, simulator substitutions, and runtime XML escaping/truncation.

Original Azure calibration (before accuracy review)

Ran 12 fresh azure/gpt-5.4-mini evaluations after the conservative tuning, covering prompt length, requested output length, batching, one-to-four-turn conversations, long conversational context, a simulated tool loop, and a generated end-to-end pipeline. Generated calibration artifacts remain local and are not committed.

Scenario	Estimated	Likely range	Actual	Actual vs estimate
Tiny prompt	3,119	2,027-4,211	2,477	-20.6%
Short prompt	3,515	2,285-4,745	2,542	-27.7%
Medium prompt	3,695	2,402-4,988	2,954	-20.1%
Long context	4,237	2,754-5,720	3,951	-6.8%
Long output	4,453	2,894-6,012	3,964	-11.0%
Four-prompt batch	14,034	9,122-18,946	10,678	-23.9%
One-turn scenario	4,948	3,216-6,680	3,781	-23.6%
Two-turn scenario	7,666	4,983-10,349	5,508	-28.2%
Four-turn scenario	15,118	9,827-20,409	10,066	-33.4%
Long-context scenario	11,381	7,398-15,364	10,517	-7.6%
Simulated tool loop	5,160	3,354-6,966	3,613	-30.0%
Generated full pipeline	19,928	12,953-26,903	18,489	-7.2%

Difference is (actual - estimate) / estimate; negative values mean the estimate was conservative.

• Conservative point estimates: 12/12 runs (100%)
• Range coverage: 12/12 runs (100%)
• Complete provider usage: 12/12 runs (100%)
• Mean absolute error: 20.0%
• Median absolute error: 22.1%
• Maximum absolute error: 33.4%
• Bias: actual usage averaged 20.0% below estimate

The high-side bias is intentional: compared with the earlier tuning, mean absolute error increases from 14.0% to 20.0%, while underestimates drop from 2/12 runs to 0/12. A separate CLI/viewer demo estimated 4,453 tokens (likely 2,894-6,012) and reported 3,913 actual tokens, 12.1% below the estimate.

[Chang Liu (changliu2)]

Requesting changes on exact head 23d1080 for four token-accounting correctness issues:

1. A no-op resume erases the completed run's existing token telemetry. The new stage_usage or token_estimate condition rewrites metrics.json even when no stage recorded new usage; a zero-pending-call estimate therefore replaces prior totals and stage data with zeros. Preserve the existing metrics when no new usage was recorded, or store invocation and cumulative totals separately.
2. Viewer estimation and execution use different results roots when ARTIFACTS_ROOT is customized. The temporary estimate config receives artifacts_root/results_dir, but writeRunConfigFiles() serializes the original config and the spawned runner receives no equivalent overrides. The viewer reserves and monitors the custom root while ASSERT writes under <MEASUREMENTS_ROOT>\artifacts\results. Apply the same resolved roots to both estimation and submission.
3. pipeline.inference.max_tool_calls has no effect on estimates. The simulated-tools example produces the identical inference estimate for limits 1 and 10 - 675 calls and 1,152,200 tokens - while HostedSession continues target/resolver calls up to the configured cap. Model the configured loop, accumulated context, simulator calls, and resulting judge transcript, or make the upper bound reflect that cap.
4. Opaque-target caveats disappear when tracked usage is zero. _log_token_estimate() returns before printing notes at zero tokens, and viewer normalization drops a zero-valued estimate entirely, hiding the fact that callable/endpoint/sandbox target-internal usage was excluded. Preserve and render note-only estimates.

[Chang Liu (changliu2)]

Re-reviewed exact head 45afbc41. The new projection work improves representative tool payloads, context accumulation, and calibration, but none of the four prior blockers is resolved:

1. A no-op resume still rewrites metrics.json whenever the zero-valued token_estimate dictionary exists, replacing completed-run telemetry with zero usage (assert_ai/runner.py:1186-1198).
2. Viewer estimation still injects custom artifacts_root / results_dir only into its temporary estimate config; writeRunConfigFiles() persists the original config for execution, so custom-root estimates and spawned runs can target different trees (viewer/src/lib/server/run-spawn.ts:699,1031-1033).
3. max_tool_calls still does not affect projections. Prompt and scenario estimates explicitly assume one tool round trip and exclude additional calls, while runtime may loop to the configured cap (assert_ai/core/token_estimator.py:1193-1196,1263-1267).
4. Note-only opaque-target caveats still disappear: the runner returns before logging notes at zero tokens, completed-run normalization drops a zero-valued estimate before reading notes, and the new-run preview neither types nor renders notes (assert_ai/runner.py:548-551; viewer/src/lib/server/data.ts:174-188; viewer/src/routes/new/+page.svelte:47-53,1672-1676).

All checks are green, but these cases remain uncovered. Keeping changes requested.

[Chang Liu (changliu2)]

Re-reviewed exact head 45afbc4. The new commit improves per-round tool payload, history, and transcript accounting, but none of the four requested blockers is resolved:

1. runner.py is unchanged: a no-op resume still rewrites an existing metrics.json from a truthy zero-valued estimate and empty stage_usage, replacing prior telemetry with zero totals.
2. viewer/src/lib/server/run-spawn.ts is unchanged: custom ARTIFACTS_ROOT values are still injected only into the estimate config, not the submitted run config, so the viewer and runner can use different result trees.
3. pipeline.inference.max_tool_calls still has no effect. On the current head, the simulated-tools example returns the same inference estimate for limits 1 and 10: 675 calls and 1,369,075 tokens. The projection still models one round trip while runtime can continue to the configured cap.
4. Zero-token opaque-target caveats are still suppressed: run logging returns before notes, and viewer normalization drops the zero-valued estimate.

Please keep this in changes-requested state until all four are addressed.