TokenCast: Forecasting Token Consumption During LLM Agent Execution
Chaoqian Ouyang, Ling Yue, Libin Zheng, Huanghui Guo, Shengxiang Xu, YiShu Wang, Ran Li, Jian Yin, Shaowu Pan, Shimin Di
cs.LG, cs.AI, cs.SE
2026-09-29
TokenCast predicts agent token use from composable segment costs, cutting MAE 14.5% vs the strongest baseline over 96 settings and saving 21.3% tokens in budget replay.
Run the same software-engineering task twice and an LLM agent can spend an order of magnitude more tokens on one of the two runs; prior work has reported spreads up to 30×. The next action follows tool feedback, so path length is unknown up front. Dialogue and tool output from earlier turns are stuffed into every later request, and each extra stretch of context raises the input bill of every subsequent hop. Total cost is future call count times how much history each hop rereads. Neither factor is known at the start.
Single-response length predictors such as TRAIL, EGTP, and TIE are enough for scheduling and memory reservation. Self-Prediction asks the agent to inspect the workspace and quote a total, which costs an extra LLM round. DSPy and SGLang need a program graph that open-ended agents do not expose. After a run starts, remaining cost still needs a forecast that updates from evidence already on the trace, without another model call.
A contiguous block of calls is a triple (n, g, b): call count, net change in input length, and a residual equal to total tokens minus starting input length times n. The residual holds generation cost and within-segment context growth. Reasoning tokens that are billed but not written back into the conversation go into b, not g.
Every call in the following segment rereads the extra context the prefix introduced, so the combined residual picks up nB × gA. That identity follows from the definitions. Any partition of a trace composes back to the same task total.
Task-level prediction runs a direct remaining-cost regressor beside a prefix-suffix compositional path, then a correction model. The compositional path predicts the prefix triple and its ending boundary, feeds the predicted input baseline into the suffix model, and adds the two with the identity. The correction model sees both point forecasts, their difference, and the boundary variables. Suffix training uses cross-fitted prefix boundaries so it never trains on the true boundary that would have been unknown at prediction time. Errors in g and n propagate through the identity, so those losses are weighted by nB and LB.
Four prediction points use only then-visible evidence. At task start the input is the statement and seed workspace. At call start the request is assembled, so the true input length is known. During generation a checkpoint is taken every 128 bytes of committed prefix, with streaming timings. After a call finishes, confirmed spend and tool outcomes refresh remaining cost. The base models are LightGBM. Separate quantile models give 5th and 95th percentiles, then held-out calibration widens the endpoints. No extra LLM calls.
Evaluation covers 4 suites, 6 agent models, 240 tasks, and 11,712 traces: SWE-bench Verified, Search-R1, MMLU-Pro, LongBench-v2. Across 96 suite-model-point cells, MAE versus the strongest comparator in that cell falls by 14.5% on average. Split by point: -2.2% at task start, 1.9% at call start, 30.4% during generation, 27.8% at task update. All 24 losses sit at task start (15) or call start (9). In-call and task-update never lose.
GPT-5.4 on SWE-bench Verified:
| Point | TokenCast | Best comparator | Method |
| Task start | 144.0k | 152.0k | Self-Prediction |
| Call start | 64.5 | 69.5 | TIE |
| In-call update | 38.9 | 74.6 | EGTP |
| Task update | 80.0k | 115.0k | TRAIL |
In-call MAE drops 47.9% versus EGTP (74.6 to 38.9). Task-update drops 30.4% versus TRAIL (115k to 80k). Task start is only 5.3% better than Self-Prediction on this GPT-5.4 cell. On Claude Opus 4.6 the same point is 150.7k against Self-Prediction's 138.9k; on DeepSeek-V4-Pro it is 212.9k against 194.5k. Before a run starts, asking the agent is often the better guess.
WAPE at GPT-5.4 task start is 4.2% on LongBench-v2 (mean target 771.2k) and 15.8% on MMLU-Pro (mean 50.6k). Shorter tasks have larger relative error. On anchor-task 90% intervals, mean interval score is 32.0% lower than Self-Prediction's native percentiles. Coverage at task start is 82.0% versus 52.7%, against a 90% nominal level.
Budget control is an offline replay on 144 SWE-bench tasks, 288 GPT-5.4 runs, and seven caps. At matched trace completion TokenCast spends 21.3% fewer tokens on average. Completion means reaching the recorded terminal state, not passing tests. Refreshing after every call costs 19.7 forecasts and 32.8 ms per run against a 129 s median wall time.
Replacing the segment triple with total, input, and output length raises normalized task MAE from 0.71 to 0.88. Dropping composition raises the overall average from 0.69 to 0.78. LightGBM is the lowest-error of six base predictors. Zero-shot transfer to independently released LiveClawBench traces yields Call Start and Task Update MAE 1.31× and 1.47× Self-Prediction. With 20 target-domain tasks those ratios fall to 0.82 and 0.85. Held-out agent transfer at Call Start is 1.02× the target history median with no target-model data.
For budget caps, queues, and user-facing cost quotes, this is a consumption layer that does not spend extra LLM calls. It needs in-distribution traces. The gap opens after execution has started: in-call and task-update MAE drop by about 30% versus the best baseline. Before the first call, Self-Prediction sometimes wins.
The 21.3% figure says a follow-along remaining-cost estimate can stop a run earlier than a fixed cap at the same trace-completion rate. The experiment replays recorded traces and does not change the agent's plan. Wiring these forecasts into a runtime that actually steers execution is left as future work. The paper aims at whole-trajectory cost when history is re-billed as input. That matches GPT-5.4 on SWE-bench, where input is 99% of tokens.
The authors state two limits: predictors are learned from recorded traces, so a new domain or agent needs some target tasks; the forecasts are not yet plugged into a policy that changes execution.
The 21.3% figure is not live savings. Replay only checks whether a trace reached its logged end state. Success rate and patch quality are unmeasured, and early stops could cut paths that would have fixed the bug. Task start is on average slightly worse than the strongest comparator. Interval coverage is 82% against a 90% nominal. For API agents, TRAIL and EGTP encode text with Llama-3.2-3B as a proxy for hidden states, so those comparisons are not symmetric.
The 30× same-task spread is cited from Bai et al. 2026. Figure 6 only shows a wide distribution over 48 GPT-5.4 anchor tasks and does not report a 30× statistic of its own. Zero-shot cross-domain transfer is still worse than asking the agent, so production use needs target-domain traces.