Skip to content

Cache trace branches by node count#1799

Open
xeophon wants to merge 1 commit into
mainfrom
codex/reuse-rollout-token-counts
Open

Cache trace branches by node count#1799
xeophon wants to merge 1 commit into
mainfrom
codex/reuse-rollout-token-counts

Conversation

@xeophon

@xeophon xeophon commented Jun 21, 2026

Copy link
Copy Markdown
Member

Overview

Cache the derived Trace.branches view by len(trace.nodes) so repeated trace properties reuse the same branch materialization. Token-limit evaluation keeps using the existing Trace API and preserves its current semantics.

What changed

  • Add a private branch cache and cached node count to Trace.
  • Return a shallow copy of the cached branch list while the node count is unchanged, preserving caller isolation.
  • Rebuild the view after nodes are appended.
  • Let num_branches reuse the cached branch count when warm, while retaining its cheap leaf scan on cold count-only access.

Why

prompt_len, completion_len, and total_tokens each read Trace.branches. Token-limit checks can therefore reconstruct the same graph paths three times. Caching the derived view at its owner removes those repeated graph walks for every caller without adding topology-specific counting logic to the interception server.

The public property still returns an independent outer list, so caller operations such as sorting, popping, or clearing do not mutate the internal cache.

Performance

Median wall time with all three token limits enabled:

Workload Before Cache miss Cache hit
200k-node linear, 1 token/node 131.509 ms 66.529 ms 31.925 ms
2k nodes / 1k branches, 1 token/node 256.453 ms 146.099 ms 88.978 ms
2k-node linear, 32 tokens/node 0.824 ms 0.479 ms 0.286 ms
2k nodes / 10 branches, 32 tokens/node 5.705 ms 3.631 ms 2.534 ms
10k nodes / 10 branches, 32 tokens/node 28.442 ms 17.964 ms 12.643 ms

A cache miss represents the normal case after the trace grows; a hit represents repeated reads at the same node count.

Warm num_branches reads improve from 346.4 µs to 0.659 µs for a 10k-node linear trace and from 100.2 µs to 0.672 µs for a shared-prefix trace with 1k leaves. Cold count-only reads continue using graph.leaves and avoid materializing full branch paths.


Note

Low Risk
Localized performance optimization on a derived view; semantics stay the same as long as branches only change when nodes are appended (existing rollout model).

Overview
Adds a private branch cache on Trace keyed by len(nodes) so the expensive root→leaf branch materialization runs once per graph size instead of on every access.

branches now returns a cached copy when the node count is unchanged; otherwise it rebuilds, stores the result, and returns a copy (same outward behavior as before). num_branches reads len(_branch_cache) when the cache is warm and still falls back to graph.leaves(self) for cold count-only access.

This mainly speeds repeated reads of prompt_len, completion_len, and total_tokens (each iterates branches), including token-limit checks that previously walked the graph multiple times at the same node count.

Reviewed by Cursor Bugbot for commit ec04697. Bugbot is set up for automated code reviews on this repo. Configure here.

Note

Cache trace branches by node count in Trace

Adds memoization to Trace.branches and Trace.num_branches in trace.py to avoid recomputing branches on every access. The cache is keyed by node count and invalidated automatically when the graph grows. Both properties return results from _branch_cache when the node count matches, and fall back to full recomputation otherwise.

Macroscope summarized ec04697.

@macroscopeapp

macroscopeapp Bot commented Jun 21, 2026

Copy link
Copy Markdown

Approvability

Verdict: Approved

This PR adds a straightforward cache for the branches property based on node count. The implementation correctly returns defensive copies on both code paths, addressing the concern raised in the review comment. No functional behavior changes beyond improved performance.

You can customize Macroscope's approvability policy. Learn more.

@xeophon xeophon changed the base branch from feat/nano-as-v1 to main June 23, 2026 04:10
@xeophon xeophon force-pushed the codex/reuse-rollout-token-counts branch from af6db01 to 5f30b53 Compare June 23, 2026 04:17
@xeophon xeophon requested a review from mikasenghaas June 23, 2026 04:25

@mikasenghaas mikasenghaas left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this seems p minor?

@xeophon xeophon force-pushed the codex/reuse-rollout-token-counts branch from 5f30b53 to d6aef07 Compare June 23, 2026 15:07
@xeophon xeophon changed the title Reuse rollout token counts across limit checks Cache trace branches by node count Jun 23, 2026
@xeophon xeophon force-pushed the codex/reuse-rollout-token-counts branch from d6aef07 to 9f93543 Compare June 23, 2026 15:09

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 9f935437d8

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread verifiers/v1/trace.py Outdated
@xeophon xeophon force-pushed the codex/reuse-rollout-token-counts branch from 9f93543 to c0fea64 Compare June 23, 2026 15:12
@xeophon xeophon force-pushed the codex/reuse-rollout-token-counts branch from c0fea64 to ec04697 Compare June 23, 2026 15:31

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: ec04697fd7

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread verifiers/v1/trace.py
subagents). Branching falls out of walking each leaf's parents back to its root."""
node_count = len(self.nodes)
if self._branch_cache_node_count == node_count:
return self._branch_cache.copy()

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Return defensive Branch copies from the cache

Fresh evidence after the prior review: this returns only a shallow copy of the cached list, so callers still receive the exact cached Branch instances. If any post-processing mutates a branch object or its public nodes list, such as trace.branches[0].nodes.pop(), later reads while len(self.nodes) is unchanged reuse the corrupted branch and can make prompt_len, completion_len, or tool_messages wrong; before this change each access rebuilt new Branch objects. Cache immutable data or return defensive Branch copies.

Useful? React with 👍 / 👎.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants