feat: delta-native message-graph trajectory

[mikasenghaas]

Summary

Move the v1 trajectory from a flat trajectory: list[Turn] — where every turn restated the whole prompt (and its tokens), so size was quadratic in turns — to a delta-native message graph. Trace.nodes (one MessageNode per distinct message, each linked to its predecessor) is the only stored model; branches and the conversation accessors are computed views over it. Each message is stored once, so in-memory / on-disk / wire size is linear; branching falls out of walking the graph; a training sample is a cheap concat along a branch path.

Changes

• verifiers/v1/graph.py (new) — MessageNode (parent, message, token_ids, mask, logprobs, finish_reason); message_hash (message equality); add_turn (reuse prefix nodes by (parent, hash), attribute each new message's tokens from the renderer's per-message spans, append the assistant node = generation-prompt scaffold + sampled completion); branches_from_nodes / leaves (the per-branch token views live on Branch).
• trace.py — Trace.nodes is the ground truth. Turn and Trace.trajectory are removed. Branch holds nodes and is self-describing: messages + the aligned token views token_ids/sampled_mask/logprobs, plus num_turns/prompt_len/completion_len/total_tokens; assistant_messages/tool_messages/has_response/is_truncated read the graph directly.
• interception/server.py, legacy.py build via graph.add_turn; clients/renderer.py threads the per-message token spans (message_token_spans()) it previously discarded.
• branching.py deleted — the graph supersedes it (message_hash is the message-equality).
• Token attribution is exact: concat(node.token_ids along a path) == prompt_ids + completion_ids the model saw (the previous assistant's closing scaffold anchors to the cumulative stored prefix, not message spans).

Data structure

A trace is a flat list of MessageNodes — one per distinct message — each pointing at its predecessor (parent, an index into nodes). The conversation is a root→leaf path; branches (compaction / subagents / resampling) are simply multiple leaves, so branching falls out of walking the graph. Each node stores only the tokens it adds (token_ids), so concatenating a path reproduces the full sequence — storage is linear, not quadratic.

nodes (flat list; parent = index into nodes)        branches = leaf→root paths
────────────────────────────────────────────        ──────────────────────────
 [0] system     parent=–  (root)                      0─1─2─3─4   → one branch
 [1] user       parent=0                              (S)(U)(A0)(U1)(A1)
 [2] assistant  parent=1   ◄ token_ids + mask + logprobs (sampled)
 [3] user       parent=2                              a resample/compaction forks it:
 [4] assistant  parent=3                                0─1─2─3─4      (leaf 4)
                                                              └─5─6    (leaf 6)  → two branches

On disk (results.jsonl — one trace per line)

{
  "id": "…", "task": { … },
  "nodes": [
    { "message": {"role":"system","content":"You sort."},          // root: "parent" omitted (= None)
      "token_ids":[1,2,3], "mask":[false,false,false], "logprobs":[] },
    { "parent":0, "message": {"role":"user","content":"cba?"},
      "token_ids":[4,5], "mask":[false,false], "logprobs":[] },
    { "parent":1, "message": {"role":"assistant","content":"abc"},  // a turn (assistant node)
      "token_ids":[9,8], "mask":[true,true], "logprobs":[-0.2,-0.2], "finish_reason":"stop" }
  ],
  "rewards": { … }, "reward": …, "metrics": { … },
  "is_completed": true, "stop_condition": "…", "is_truncated": false,
  "timing": { … }, "errors": []
}

• parent — index of the predecessor node; omitted for a root (None).
• token_ids — this message's delta contribution (leading template scaffold + its own tokens; for an assistant, the generation-prompt scaffold + the sampled completion).
• mask[i] — True ⇒ a model-sampled (trainable) token; False ⇒ scaffold / input. Only assistant nodes have Trues.
• logprobs — sampling logprobs, aligned to the True entries.
• finish_reason — assistant nodes only (truncation detection).
• branches / num_turns / assistant_messages / … are not stored — they're computed views over nodes.

Loading + building a training sample

import json, verifiers.v1 as vf

trace = vf.Trace[vf.Task].model_validate(json.loads(line))   # from results.jsonl or the wire

# branches — one training sample per branch; a view over the graph (no stored copies)
for b in trace.branches:
    conversation = b.messages          # list[Message] for this root→leaf path

# each branch is first-class and carries its own flat training sequence:
for branch in trace.branches:
    ids, mask, logprobs = branch.token_ids, branch.sampled_mask, branch.logprobs
    first = mask.index(True)                       # prompt | completion split = first sampled token
    sample = dict(
        prompt_ids          = ids[:first],
        completion_ids      = ids[first:],
        completion_mask     = mask[first:],        # True where the model generated (trainable)
        completion_logprobs = logprobs[first:],
    )

This is exactly what prime-rl's trace_to_samples does — one TrainingSample per branch.

Verification

• Size (linear vs quadratic) — same trace, new (graph) results.jsonl vs the old per-turn-full-prompt + expanded-branches layout:

  • real terminal-bench-2 (bash agent), same trace both ways:

    turns	new (graph)	old (trajectory)	reduction
    8	47 KB	0.40 MB	8.7×
    16	74 KB	1.35 MB	18.7×
    32	236 KB	7.70 MB	33.5×
    64	~0.5 MB	~31 MB	~67× (projected)
    128	~0.9 MB	~123 MB	~134× (projected)
    256	~1.8 MB	~493 MB	~268× (projected)

    8/16/32 are measured (same trace, both layouts); 64/128/256 are projected (old ∝ turns², new ∝ turns). The reduction tracks turn count — at 256 turns the old layout would be ~0.5 GB/trace vs ~1.8 MB.

  • synthetic 5→40 turns: 2.6× → 17.6×.

• Training samples correct end-to-end (with prime-rl #2763): RL on alphabet-sort trains cleanly through graph → trace_to_samples → trainer at 0% error on both v1 (native renderer path, 3.2 turns, 432/512 trainable) and v0 (legacy bridge, 2 turns, 368/512).
• Existing trace / serve / env tests green; ruff clean.

Notes

• Renderer-only training: per-message attribution comes from the renderer client; the chat/MITO client carries no spans → eval-only (non-trainable). Paired with prime-rl #2763 (graph-walk trace_to_samples).
• Multimodal is left out (not in this base, feat/nano-as-v1); MessageNode is trivially mm-extendable when mm lands.

---

Note

High Risk
Breaking change to v1 trace shape and serialization (trajectory/Turn removed); training and downstream consumers must read nodes/Branch instead of per-turn lists.

Overview
Replaces the v1 trace’s flat trajectory: list[Turn] (full prompt restated every turn → quadratic storage) with a delta-native message graph: Trace.nodes is the only persisted model; Turn and branching.py are removed.

verifiers/v1/graph.py adds MessageNode (parent link, per-message token_ids / mask / logprobs), add_turn (prefix dedup via message_hash, renderer span attribution), and graph walks for branches. Branch now holds nodes and exposes training-ready token_ids, sampled_mask, and logprobs by concatenating the path. The interception server and legacy v0 bridge record turns via graph.add_turn; the renderer client threads message_spans on TurnTokens (transient, exclude=True on persist).

Docs and the eval dashboard are updated for graph-based branching and token display (b.prompt_len / b.completion_len).

^{Reviewed by Cursor Bugbot for commit dae561f. Bugbot is set up for automated code reviews on this repo. Configure here.}

Note

Replace flat trajectory with a delta-native message graph in trace representation

• Introduces MessageNode in graph.py as the atomic unit of the trace, storing token_ids, mask, logprobs, and finish_reason per message.
• Replaces Trace.trajectory: list[Turn] with Trace.nodes: list[MessageNode]; branches, turn counts, and token aggregations are now derived from root→leaf paths through this graph.
• graph.add_turn reuses shared prompt prefixes via a hash-keyed head index, avoiding duplication of repeated messages across turns.
• Updates interception/server.py and legacy.py to write into the message graph rather than appending Turn objects.
• Risk: Turn and Trace.trajectory are removed; any code referencing these will break.

^{Macroscope summarized dae561f.}

[macroscopeapp]

Approvability

Verdict: Needs human review

This PR replaces core data structures (flat trajectory list → message graph) and changes how traces are stored, branches are computed, and training samples are built. Major refactors that restructure shared abstractions warrant human review.

^{You can customize Macroscope's approvability policy. Learn more.}