PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2609.35760v1cs.LGcs.AIcs.SE

TL;DR

Agent token consumption is difficult to forecast because actions and context growth make later-call costs depend on the execution trajectory. TokenCast composes segment costs and updates forecasts from observed evidence without additional LLM calls. Across four benchmarks and six models it reduces MAE by an average of 14.5% over 96 comparisons, and its budget-control replay uses 21.3% fewer tokens at matched trace completion.

  • Problem

    Agent actions and accumulated context make later-call inputs and total task consumption variable, complicating forecasts before execution is complete.

    Source: §1 · pp. 1–2

  • Method

    TokenCast represents each execution segment with a composable cost representation and updates direct and compositional forecasts from execution evidence without additional LLM calls.

    Source: §1 · p. 3

  • Results

    Across four benchmarks and six agent models, TokenCast’s MAE reduction averages 14.5% over 96 comparisons with the strongest comparator; budget-control replay uses 21.3% fewer tokens than a fixed-budget policy at matched trace completion.

    Source: §5 · p. 10

  • Takeaways & Limitations

    In sequential agent execution, carried context couples remaining cost to future call counts and input lengths, and forecasts can be updated without querying an LLM for the estimate.

    Source: §1 · p. 2

    The offline budget-control replay does not measure task resolution.

    Source: §4.3 · p. 9

Abstract

from arXiv · show

When a large language model (LLM) agent executes the same task, token consumption can vary by over an order of magnitude across runs. The agent chooses its next steps based on tool feedback and intermediate results, while the growing context steadily inflates the input size of every subsequent call. The total consumption of a task is therefore hard to predict before execution and the prediction must be revised as the run unfolds. In this paper, we propose TokenCast, which learns a composable cost representation for each execution segment, recording its own consumption and the context growth it introduces. Composing adjacent segments yields a cumulative estimate that captures the extra input cost incurred when context from earlier segments is re-read by every later call. As execution unfolds, newly observed evidence refreshes the forecast, requiring no additional LLM calls and incurring a mean cumulative prediction time of 32.8 ms per run on SWE-bench Verified. Across 4 task suites and 6 agent models, TokenCast's mean absolute error reduction against the strongest comparator averages 14.5% over 96 evaluated combinations. In offline budget-control replay, TokenCast uses 21.3% fewer tokens on average than a fixed-budget policy at matched trace completion. The code is available at https://github.com/DEFENSE-SEU/TokenCast.

1 Introduction

Agent token costs are difficult to forecast because tool-driven actions and accumulated context make later calls and total consumption variable. TokenCast composes segment costs to account for repeated input consumption, then refreshes forecasts from execution evidence; it reduces error across evaluated configurations and saves tokens in budget-control replay.

1 Introduction

  • Same-task executions can differ in token consumption by up to 30× because agent actions vary and earlier context growth increases input costs for subsequent calls.The accumulated dialogue and tool outputs become context for later requests.Source: §1 · pp. 1–2
  • Open-ended agents depend on tool feedback and intermediate results, so their behavior cannot be fully forecast from a workflow structure available before execution.TokenCast focuses on remaining provider-accounted input and output consumption in sequential agent execution.Source: §1 · p. 2
  • TokenCast factorizes segment costs and composes adjacent segments to account for repeated input consumption as earlier context is re-read by later calls.The staged prefix–suffix predictor combines direct and compositional forecasts and updates from execution evidence without additional LLM calls.Source: §1 · p. 3
  • Across four benchmarks and six agent models, TokenCast’s MAE reduction averages 14.5% over 96 comparisons with the strongest comparator, while budget-control replay uses 21.3% fewer tokens at matched trace completion.The replay comparison is against a fixed-budget policy.Source: §1 · p. 3

2 Related Work

Prior work predicts single-response length or task-level consumption, but these approaches rely on known prompts, model internals, or workflow structure. TokenCast addresses carried-context costs in open-ended execution, while prior trace studies describe consumption retrospectively.

Request-level output length prediction

  • Request-level predictors estimate one response from a known prompt and may refine the estimate during decoding, but their assumptions do not hold when later agent requests have not yet been created.These methods may use prompt features or model internals such as intermediate states and entropy statistics.Source: §2 · p. 3

Task-level and multi-step prediction

  • Task-level predictors use task descriptions or model cost assessments, while workflow systems rely on structures or dependencies unavailable in open-ended agents before execution.Self-Prediction estimates consumption by stage; open-ended agent behavior depends on tool feedback, environment state, and intermediate results.Source: §2 · p. 4
  • Chimera predicts remaining workflow output with a quantile random forest, and Pythia infers likely paths and role-level output lengths from historical traces.TokenCast instead addresses repeated input costs from carried context and uses no additional LLM calls.Source: §2 · p. 4

Trace studies

  • Deployed-agent trace studies characterize iterative review loops, long contexts, and prefix reuse after the fact rather than forecasting consumption during a running task.Source: §2 · p. 4

3 Method

The method forecasts token use at call and task levels, representing contiguous execution segments so their costs can be composed as context grows. It combines direct and compositional task forecasts, refreshes them as evidence arrives, and trains separate predictors for the forecasting components.

3.1 Problem Formulation

  • A task consists of interleaved LLM calls and tool interactions; the method predicts call consumption or total/remaining task consumption at specified execution points.The four settings are Task Start, Call Start, In-call Update, and Task Update.Source: §3.1 · p. 4

3.2 Call–Task Forecasting

  • A segment is represented by call count, net input-length change, and residual cost beyond its starting-input baseline; composing adjacent segments adds n_B g_A to rebase the later segment’s cost.Provider-billed reasoning tokens contribute to residual cost, but contribute to context change only when they persist in the conversation.Source: §3.2 · pp. 4–5
  • Task-level forecasts combine a direct remaining-cost prediction with a prefix–suffix composition path conditioned on predicted context change and boundary state.The suffix input baseline is set by predicted context change; the suffix representation is converted to remaining cost using the composition identity.Source: §3.2 · p. 5
  • As calls complete, TokenCast refreshes forecasts using confirmed cumulative cost and completed tool outcomes, while the next request’s input length remains predicted until assembly.At Call Start the assembled request length is known; during generation, the committed output prefix and streaming timing add features.Source: §3.2 · p. 5

3.3 Compositional Learning

  • Separate LightGBM models fit direct, prefix, and suffix predictions from completed traces, with local loss weights motivated by how context-change and call-count errors affect composed cost.Cross-fitting trains the suffix model on predicted prefix boundaries, and a correction model is trained on out-of-fold outputs of the full pipeline.Source: §3.3 · p. 5
  • Separate quantile models estimate 0.05 and 0.95 interval endpoints, which are widened using held-out calibration residuals and bounded below by confirmed consumption.Source: §3.3 · pp. 5–6

4 Experiments

Across the experiments, TokenCast improves token-consumption forecasts and supports lower-token budget control at matched trace completion. Reliability, generalization, and ablation results show how performance varies with prediction point and design choices.

4.1 Experimental Setup

  • Evaluation uses 11,712 traces from 240 tasks across four suites and six agent LLMs, with adapted baselines evaluated at four prediction points.Tasks cover software engineering, retrieval-based question answering, knowledge-based reasoning, and long-context understanding.Source: §4.1 · pp. 6–7

4.2.1 Main Results

  • TokenCast’s MAE reduction against the lowest comparator MAE averages 14.5% across 96 benchmark–model–prediction-point combinations.At Task Start, the average reduction is −2.2%; it is 30.4% at In-call Update and 27.8% at Task Update.Source: §4.2.1 · p. 7

4.2.2 Prediction Reliability

  • Repeated GPT-5.4 runs show substantial consumption spread across four benchmarks; calibrated intervals reduce MIS by 32.0% on average versus Self-Prediction’s native intervals.At Task Start, TokenCast intervals cover 82.0% of outcomes against a nominal 90% level, versus 52.7% for Self-Prediction.Source: §4.2.2 · p. 7

4.2.3 Generalization

  • On unseen task types, MAE ratios versus Self-Prediction fall from 1.31/1.47 zero-shot to 0.82/0.85 with 20 tasks; held-out-agent Call Start reaches 0.95× history-median after 20 tasks.The task-type ratios correspond to Call Start and Task Update; the held-out-agent result is at Call Start.Source: §4.2.3 · pp. 7–8

4.3 Case Study

  • On 288 runs from 144 SWE-bench Verified tasks, TokenCast used 21.3% fewer tokens across seven budgets at matched trace completion, including prediction costs.Trace completion means reaching the recorded terminal state; replay does not measure task resolution. Example forecasts rise after a verification failure and after post-answer file operations.Source: §4.3 · pp. 8–9, Figure 4 · p. 9

4.4 Ablation Study

  • Request length leads at Task Start and Call Start, committed prefix length during generation; all-call updating takes 32.8 ms cumulative prediction time.After a call completes, last-call input tokens, remaining plan items, and calls without progress have similar relative importance.Source: §4.4 · pp. 9–10, Figure 5 · p. 10
  • Removing segment composition raises normalized average MAE from 0.69 to 0.78; at Task Update, the full correction model reaches MAE 0.62.Replacing the segment triple with total token count, input length, and output length raises task-level and call-level normalized MAEs from 0.71 to 0.88 and 0.68 to 0.82, respectively.Source: §4.4 · p. 9
  • Dropping cross-fitting and cost weighting together raises normalized average MAE from 0.69 to 0.78 and lowers 90% interval coverage from 90.6% to 87.4%.Dropping cross-fitting alone gives 0.73 MAE, while dropping cost weighting alone gives 0.72; LightGBM has the lowest error and highest coverage among six base predictors, with 0.8 ms model time.Source: §4.4 · pp. 9–10

5 Conclusion

TokenCast forecasts agent token consumption by composing segment costs so context growth is reflected in later calls. Across four benchmarks and six agent LLMs, it reduces MAE against the strongest comparator by an average of 14.5% across 96 combinations; offline budget-control replay uses 21.3% fewer tokens at matched trace completion.

5 Conclusion

  • 14.5% average MAE reduction against the strongest comparator across 96 combinations, with 21.3% fewer tokens than a fixed-budget policy at matched trace completion in offline replay.The mean cumulative task-level prediction time is 32.8 ms per run on SWE-bench Verified.Source: §5 · p. 10
  • TokenCast composes segment cost representations to account for context growth that increases downstream call costs.The representation separates the starting-input baseline from incremental consumption.Source: §5 · p. 10
  • Transfer to a new task domain or agent LLM improves with a small set of target tasks from that setting.The predictors are learned from recorded traces.Source: §5 · p. 10
  • Integrating forecasts into a runtime decision framework that actively manages execution under budget constraints remains future work.Source: §5 · p. 10

Ethics Statement

The study uses public benchmarks and recorded agent trajectories, with no human subjects or private user data collected. Models and datasets are used under their respective licenses and terms.

Ethics Statement

A Comparison with Related Methods

TokenCast forecasts call and remaining-task consumption from execution features and a segment-cost representation. Its execution evidence includes task, history, action, generated-prefix, recall, and timing features.

A Comparison with Related Methods

B.1 Execution Evidence

  • The observation record uses only information available at the prediction point; provider records supply token usage, and a logical call includes its initial request and retries.Future request input length remains predicted until the request is assembled.Source: §B.1 · p. 16
  • Task features encode word unigrams and bigrams plus character 3- to 5-grams, with 12,000 terms for each representation.A benchmark difficulty label is prepended to the task statement when available.Source: §B.1 · p. 16
  • Action labels distinguish reads, edits, tracking, tests, and other runs; a call’s stage is determined by its most recent action.Tool errors or nonzero shell exits mark an action as failed.Source: §B.1 · p. 16
  • Execution history tracks plan progress and recent consumption statistics, with Task Update also recording failures, retries, and calls without progress.Progress is defined as an edit, test, tracking update, or completed plan item; unavailable window statistics are masked.Source: §B.1 · p. 16
  • Recent-action slots record action type, status, result size, elapsed calls, and reported result, while the action window is selected from 3, 5, or 8 using validation tasks.Source: §B.1 · pp. 16–17
  • Generated-prefix features hash words and word pairs and track structural cues such as JSON validity, nesting, open strings, and code fences.Source: §B.1 · p. 17
  • Prefix recall compares the current generated prefix with earlier requests and checkpoints using cosine similarity, retaining associated final and remaining lengths.Source: §B.1 · p. 17

B.2 Predictors and Training

TokenCast uses separate predictors for four call- and task-level prediction settings, with task forecasts combining direct and compositional paths. Segment composition propagates context growth, and separately calibrated intervals are bounded below by consumption already confirmed within the prediction scope.

B.2 Predictors and Training

  • Call Start and In-call Update predict complete current-call consumption, while Task Start and Task Update combine direct remaining-task prediction with a prefix–boundary–suffix compositional forecast.The task-level pipeline applies a correction model after computing the compositional forecast.Source: §B.2 · p. 17
  • Training labels come from completed traces, and task-level instances pair remaining-consumption targets with prefix representations, ending boundaries, and suffix representations.Inputs are restricted to evidence available at the corresponding prediction point.Source: §B.2 · p. 17
  • Task-level fitting cross-fits prefix boundaries for suffix training, recomputes suffix residual labels against predicted input baselines, and uses out-of-fold forecasts to fit the correction model.Source: §B.2 · p. 17

B.3 Theoretical Analysis

  • Composing adjacent segment representations adds the downstream cross-term for earlier context growth, and the composition is associative.For a complete trace, the resulting residual is invariant to how the trace is partitioned.Source: §B.3 · pp. 17–18
  • Suffix cost error depends on errors in prefix context growth, suffix call count, and suffix residual; boundary prediction can affect both the learned residual and suffix-model inputs.Cross-fitting exposes suffix training to out-of-fold boundary errors.Source: §B.3 · p. 18

B.4 Interval Calibration

  • Separate quantile models produce raw interval endpoints, which are calibrated for each prediction setting on held-out tasks and bounded below by confirmed consumption.Calibration parameters remain fixed during inference.Source: §B.4 · p. 18

C.1 Trace Collection

Repeated executions of the same tasks show wide token-consumption spread, while calibrated intervals improve interval scores over native Self-Prediction intervals. The section also describes trace collection, evaluation splits, baselines, and metrics.

C.1 Trace Collection

  • The collection spans four benchmarks, task-specific workspaces, two agent harnesses, six agent models, repeated runs, and provider-recorded per-request token usage.Runs are capped by time, calls, retries, output length, and total tokens; Table 6 reports median tokens and calls and correctness by benchmark and model.Source: §C.1 · pp. 18–19

C.2 Data Splits

  • The evaluation uses task-weighted out-of-fold predictions across five folds and three random seeds, with all runs of a task kept in the same fold.The next fold calibrates intervals, the following fold selects hyperparameters, and the remaining two train.Source: §C.2 · p. 20

C.3 Baselines

  • TRAIL, EGTP, and TIE are adapted output-length baselines, while Self-Prediction queries the agent LLM for consumption estimates and intervals at each prediction point.Learned baselines use the same training, validation, and calibration partitions.Source: §C.3 · pp. 20–21

C.4 Metrics

  • Metrics include task-weighted MAE, WAPE, interval coverage, interval width, and interval score; Norm. Avg. averages normalized MAE across four prediction points.The history-median predictor supplies the normalization reference for each benchmark, agent LLM, and prediction point.Source: §C.4 · p. 21

D.1 Run-to-Run Variation and Interval Reliability

  • Repeated GPT-5.4 executions show wide token-consumption spread across all four benchmarks, and TokenCast’s wider calibrated intervals improve MIS over Self-Prediction’s native uncalibrated intervals.The repeated-run analysis covers 48 anchor tasks; MIS penalizes interval width and outcomes falling outside the interval.Source: §D.1 · pp. 21–22, Figure 6 · p. 22

D.2 Results Across Benchmarks and Agent LLMs

This section reports MAE on SWE-bench Verified and defines the table’s normalized average relative to a history-median predictor.

D.2 Results Across Benchmarks and Agent LLMs

  • Norm. Avg. averages normalized MAE across prediction points, using the corresponding history-median predictor as the reference.The history median is the median target from training tasks for the same benchmark, agent LLM, and prediction point.Source: §D.2 · p. 23
  • Table 7 reports MAE on SWE-bench Verified in tokens, with k denoting thousands of tokens.It also reports Norm. Avg., the macro-average of MAE normalized at each prediction point by the history-median MAE.Source: Table 7 · p. 23, §D.2 · p. 23
  • The supplied section passages specify the table’s metrics and scope but do not provide cell values or identify a winning method.Table marks denote best and second-best results using unrounded values.Source: Table 7 · p. 23, §D.2 · p. 23

D.3 Error Relative to Consumption

The section examines error relative to consumption, transfer to unseen tasks and models, length and reasoning shifts, and the overhead of online forecast updates.

D.3 Error Relative to Consumption

  • WAPE expresses weighted MAE as a percentage of mean target consumption, and Table 11 reports WAPE and mean target consumption across four benchmarks and four prediction points.The supplied passages do not state the table’s WAPE cell values.Source: §D.3 · pp. 23, 27, Table 11 · p. 27

D.4.1 Unseen Task Types

  • With 10 held-out-domain tasks, MAE relative to Self-Prediction is 0.86 at Call Start and 0.91 at Task Update; with 20 tasks, the ratios are 0.82 and 0.85.Zero-shot ratios are 1.31 at Call Start and 1.47 at Task Update.Source: §D.4.1 · p. 28

D.4.2 Unseen Agent LLMs

  • With 3–10 target-model tasks, transfer to an unseen agent LLM achieves lower Call Start MAE than training on those target-model tasks alone.The transferred predictor’s Call Start MAE is 1.02 times the target history median with no target tasks and 0.95 with 20 tasks.Source: §D.4.2 · p. 28, Figure 7 · p. 28

D.4.3 Cross-Harness Transfer

  • Table 12 reports TokenCast MAE relative to the history median under each test harness and includes transfer from DeepSeek Harness to OpenHands on held-out tasks.The supplied passages specify the table’s scope but not its results.Source: §D.4.3 · p. 28, Table 12 · p. 29

D.4.4 Length Extrapolation

  • 44.9% lower MAE at In-call Update and 26.8% lower at Task Update than Self-Prediction on the longest-task band; MAE rises by at most 7.8% versus matched random splits.The longest-task evaluation trains on the four shorter call-count bands and tests on the longest band.Source: §D.4.4 · pp. 28–29, Table 13 · p. 29

D.4.5 Reasoning Configuration Shift

  • Under a GPT-5.4 reasoning-configuration shift, median run consumption falls from 273k to 238k tokens; within the off setting, TokenCast reduces MAE versus Self-Prediction by 50.8% at In-call Update.The within-off comparison also reports reductions of 20.7% at Call Start and 32.2% at Task Update.Source: §D.4.5 · p. 29, Table 14 · p. 29

D.5 Online Update Frequency

  • Every-call updating averages 19.7 forecasts and 32.8 ± 30.5 ms per run, while three-call and five-call schedules average 11.9 ± 10.4 ms and 7.7 ± 6.4 ms.Every-call updating remains a small fraction of the 129 s median run time.Source: §D.5 · pp. 29–30, Table 15 · p. 30

D.6 Ablation Study

The ablations examine predictor and composition choices; the section also reports offline budget-control results and illustrates forecast changes during code repair.

D.6 Ablation Study

  • LightGBM has the lowest normalized average MAE, 0.69, and highest 90% interval coverage, 90.6%, among the compared base predictors.Its inference time is 0.8 ms, compared with 3.9 ms for XGBoost and 5.2 ms for CatBoost.Source: §D.6 · p. 30, Table 16 · p. 30
  • Combining direct and compositional forecasts with boundary variables reduces normalized MAE to 0.80 at Task Start and 0.62 at Task Update, with 90.6% coverage.The compositional path alone is worse than direct regression at Task Start but better at Task Update.Source: §D.6 · p. 30, Table 17 · p. 31
  • Removing composition raises normalized average MAE from the full segment representation, and removing g raises task-level MAE from 0.71 to 0.79.Removing the segment triple’s composition gives normalized average MAE of 0.78; removing g also changes call-level predictions and does not isolate the cross-term.Source: §D.6 · pp. 30–31, Table 18 · p. 31
  • Removing cross-fitting or the cost-weighted loss raises normalized average MAE by 0.03–0.04; removing both raises MAE to 0.78 and lowers coverage to 87.4%.Removing the correction model raises MAE to 0.74 and lowers coverage to 89.0%.Source: §D.6 · p. 31, Table 19 · p. 31

D.7 Budget Control

  • The replay controller stops at Task Update when confirmed consumption plus a selected quantile of predicted remaining consumption exceeds the budget.It charges stopped runs at their stopping-checkpoint consumption and includes prediction processing and time in replay accounting.Source: §D.7 · p. 32
  • Only TokenCast matches fixed-budget trace completion at every tested budget in the offline replay.The replay uses 288 GPT-5.4 runs from 144 SWE-bench Verified tasks.Source: §D.7 · pp. 31–32, Table 20 · p. 33

E.1 Code Repair

  • In the code-repair run, the task-total forecast rises to 281,746 tokens after verification failure, then falls to 252,381 after a compatibility workaround passes assertions.The run ultimately consumes 243,371 tokens; subsequent inspection, artifact preparation, and submission account for 37.9% of the final total.Source: §E.1 · pp. 34–35, Table 22 · p. 35
  • During call 12, the 90% interval contracts from 911 tokens at Call Start to 65 tokens at 1,450 generated bytes while the forecast stays close to the recorded 13,420-token call cost.Figure 9 tracks confirmed consumption over wall time, and Figure 10 shows task-level and current-call forecast updates.Source: §E.1 · p. 34, Figure 9 · p. 34, Table 22 · p. 35

E.2 Long-Context QA

The Long-Context QA section illustrates how repeated context makes extra calls costly and describes a self-prediction prompt for estimating consumption at four points in execution. It covers prediction scope, retained context, uncertainty, output routing, and estimation overhead using the evidence available at each point.

E.2 Long-Context QA

  • After a failed first write, replacing the placeholder at call 3 was followed by 395,353 more tokens; the Task Update forecast rose from 526k to 835k.The file operations extended the execution to six calls; Figure 11 and Table 23 show task- and call-level forecasts.Source: Figure 11 · p. 35, §E.2 · pp. 35–36, Table 23 · p. 36
  • 787,465 tokens total included 785,146 input tokens (99.7%), with repeated input contributing 778,284 tokens.The recorded execution is decomposed as T = 6 × 129,714 + 9,181 = 787,465.Source: Table 23 · p. 36, §E.2 · p. 36

F Self-Prediction Prompt

  • The self-prediction protocol adapts Bai et al. [2026] to four prediction points, using initial environment inspection at Task Start and observed execution evidence at later points without advancing the task.At later points, unavailable fields remain unknown rather than being treated as zero.Source: F Self-Prediction Prompt · pp. 36–37
  • Task Start forecasts a complete run, Task Update forecasts only future consumption, and call-level forecasts cover the current call, including retries and any generated prefix.A subsequent request after tool feedback belongs to a later call, outside the current call-level estimate.Source: F Self-Prediction Prompt · pp. 36–37
  • Task-level estimates divide the expected continuation into phases and account for both added calls and retained content reread in later requests.Input estimates use supplied counts where available and otherwise reflect current context, expected additions, and supported retention behavior.Source: F Self-Prediction Prompt · pp. 37–38
  • The prompt anchors estimates in confirmed usage without double-counting, accounts for supported retries, and requests median, 5th-percentile, and 95th-percentile totals.The lower and upper estimates define an interval targeting 90% coverage.Source: F Self-Prediction Prompt · p. 38
  • The prompt routes JSON estimates through a finish tool when available, counts estimation-only calls and inspections as overhead, and defines predicted totals by prediction scope.The total represents complete-task consumption at Task Start, complete-call consumption at both call-level points, and remaining-task consumption at Task Update.Source: F Self-Prediction Prompt · p. 39

Related papers

Loading 2609.35760v1…