Skip to content

v1 rollout token limits double-count the shared prefix on forked traces #2079

Description

@dumko2001

When a trace forks, RolloutLimits double-counts the shared prefix and stops the rollout too early.

RolloutLimits.reached() checks trace.num_input_tokens / num_total_tokens, and those sum each branch's counts (trace.py). Forked branches share the same prefix nodes, so the prefix is counted once per branch. A two-branch rollout sees ~2x its real input on the shared part, and trips the cap before it actually should.

Minimal case (shared prompt prefix, two sampled leaves):

branch[0] input=2   branch[1] input=2
unique graph input = 2
trace.num_input_tokens (summed) = 4
max_input_tokens = 3  ->  reached() refuses (max_input_tokens); real usage is within budget

The check runs before each turn, so once a fork exists mid-rollout (subagents, compaction, retokenization drift) later turns get refused early.

The reason I'm filing instead of just sending a patch: there's no obviously-correct number.

  • summing branches over-counts the shared prefix (current behavior)
  • max over branches under-counts genuinely parallel work (N subagents each 50k -> reports 50k)
  • provider usage can't recover the unique input footprint either — each branch's prompt_tokens re-includes the prefix

So it comes down to what the caps are meant to bound. Roughly:

  1. unique context footprint (dedupe shared graph nodes; token-id based, training-only), or
  2. provider cost (re-sent prefix is real work and should count repeatedly), or
  3. keep usage-based but attribute incremental input per unique node/call.

I have a repro and a prototype for (1). Happy to send a PR once you decide which semantics you want — didn't want to bake in a behavior change unilaterally.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions