feat(token-id-capture): chain a rollout's calls into one response

[ananthsub]

Builds one trainable Responses trajectory from a rollout's unordered token-capture records.

Reconstruction flow

flowchart TD
    E[Unordered TokenEntry records] --> Z[Exclude calls with no generated tokens]
    Z --> O[Order calls by prompt length]
    O --> P[Find the earlier call whose complete tokens are the longest strict prompt prefix]
    P --> A{Several candidates share the longest prefix?}
    A -->|yes| Q[Quarantine the ambiguous call]
    A -->|no| L[Attach the inferred predecessor]
    Q --> F[Build roots and chains]
    L --> F
    F --> S{Exactly one root and one trainable chain?}
    S -->|no| M[mask_sample = true]
    S -->|yes| R[Project contiguous Responses output]
    R --> C[Assert prefix continuity and report metrics]

Loading

The predecessor relationship is inferred only during reconstruction: call B follows call A when A's complete token sequence is a strict prefix of B's prompt.

Summary

• Chains calls through strict longest-prefix inference and quarantines indistinguishable candidates instead of guessing.
• Excludes empty generations so a filtered call cannot become the predecessor of its retry.
• Fails closed for incomplete capture, no trainable generation, multiple roots or chains, unresolved records, or unsafe projection.
• Reports roots, chains, quarantine, delivered-token fraction, empty generations, and parent-link failures.
• Keeps the low-level per_request builder for multi-trajectory consumers, while single-response delivery rejects it explicitly.
• Covers trajectories_from_source success, incomplete state, source failures, and unsupported single-response modes.

Depends on #2124. Consumed by #2126.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[ananthsub]

/claude review

[claude]

SHIP WITH CARE — trajectory builder + consumer for stitching captured token records into contiguous training trajectories. The correctness-critical parts hold up well.

What I checked hardest (this feeds RLHF; wrong tokens/masks corrupt training silently):

• Masking errs safe. Every uncertain condition — unresolved final-call retries, >1 root, incomplete capture, log-prob/token mismatch, ambiguous parents — resolves to mask_sample=True or quarantine rather than a guess. _assemble wraps build+project+assert in a try/except that degrades one rollout instead of killing a batch, which matches the stated caller (rollout-collection / trainer step loops).
• Projection invariant is self-checked. assert_prefix_contiguity runs on every assembled response, so a stitching bug fails the rollout rather than shipping a discontiguous prompt/generation sequence. Usage counts read from token-bearing items, so a leading content-only item won't KeyError or skew totals.
• Async correctness is clean. No httpx, no ray.get. Sync co-located path uses the new seal_now; the trainer path awaits source.seal. The _seal→seal_now rename has no external callers and seal() still delegates to it correctly.
• clear_token_captures_for_rollouts correctly addresses the append-mode id-reuse hazard (rerun of a deterministic rollout id would otherwise stitch two attempts into one chain).

Two non-blocking items posted inline:

• NOTE (consumer.py): per_request builder produces roots=0, so mask_sample masks 100% of its output — the strategy is effectively unusable through this consumer. Default prefix_merging unaffected.
• NOTE (consumer.py): trajectories_from_source and clear_token_captures_for_rollouts are new public API with no direct tests; the trainer-transport path has seal/incomplete handling a builder-only test doesn't cover.

Neither is a merge blocker. No BLOCKER or RISK findings — masking-on-uncertainty and the contiguity assertion make the failure mode "dropped sample," not "corrupted score."

[ananthsub]

/claude review

[claude]

Review incomplete: the PR revision changed during review. The captured base ref (f480877826ff9dbda89a44411ca608a5443ab9b3) no longer matches the current base (6aca3bf814baf9fc839609dd6577d886450dd4c7), so check-pr-revision.sh fails the pre-publish revision check. To avoid posting findings against a stale diff, no inline comments were posted. Please re-run the review against the current head.

[ananthsub]

/claude review

[cmunley1]

we should support returning the full tree / forest, or at least segments in a list with broadcasted reward.

Also would like to have a optional filter for certain prompts like title calls

[cmunley1]

Should this be asyncio.to_thread(_assemble, ...)

[cmunley1]

With rollouts with 100s to 1000 turns, _infer_parent scans every prior node and compares full cumulative prefixes, which becomes O(turns³) as context grows, and the recursive walk() below might hit Python’s recursion limit 1k turn scale.

Should we make traversal iterative? Would a trie help here? We should test 1000 turn