feat(rollout_gateway): linear-history mode via generate-time prefix h…

[hellodanylo]

feat(rollout_gateway): linear-history mode via generate-time prefix healing

Summary

Adds an opt-in history_mode="linear" to the rollout gateway that keeps a strictly linear, append-only agent history as one training sample per session instead of letting re-tokenization drift shatter it into many forked samples. It is implemented as a small LinearHealer (rollout_gateway/linear.py) that heals drift at generation time; the existing TrajectoryManager is unchanged and stays permanently CLEAN. Default behavior (history_mode="tree") is untouched.

Problem

The gateway owns tokenization, so each turn it re-renders the whole conversation the agent replays. Chat templates are deterministic, so system/user/tool messages re-render to identical ids — but a prior assistant message rarely does: its served tokens came from generation, while on the next turn it is re-rendered from a parsed message dict (text + tool_calls) back through the template, yielding near-identical ids that differ in whitespace, tool-call JSON spacing, or reasoning re-keying. TrajectoryManager detects that drift and, to stay safe for harnesses that may branch or edit history, resolves it by FORK (split the rollout into a second sample) or REALIGN (drop a turn's signal).

For a strictly linear, multi-turn tool-calling agent (the target use case here — a SWE agent whose history is append-only and where cosmetic drift is out-of-distribution-insignificant) that safety is pure overhead. On a real multi-turn SWE-agent run (Qwen3-Coder-30B, fork_threshold_tokens=0) the large majority of sessions forked into several records each, even though every rollout was linear. Each fork re-emits the shared prefix as loss_mask=0 context and splits the trajectory, so it is a large, avoidable training-efficiency loss.

Approach

Because the gateway controls the tokens handed to the backend, it can splice the exact ids it already served for prior turns over the client's re-rendered version before generation, so the model always generates on the canonical, drift-free context:

fed_{N+1} = served_prefix + close_N + delta_tail
O_{N+1}   = backend.generate(fed_{N+1})
served_prefix = fed_{N+1} + O_{N+1}    # extend the canonical sequence

where served_prefix is the canonical ids through the last generated turn, close_N is the template's between-message glue, and delta_tail is the genuinely-new tokens (new observation messages + generation prompt). O_{N+1} is then sampled on the canonical context, TrajectoryManager sees an append-only sequence, appends the tail, and never forks — one session, one TraceRecord, every generated turn trained on its true multi-turn context.

Linearity is decided in message space, at the same altitude TrajectoryManager matches replayed history: dict equality (which is order-independent), not raw token identity. The stored prev_messages (whose last element is the gateway's own parsed assistant message) must be a prefix of the incoming messages under ==, and the tools schema must be unchanged. The genuinely-new messages are then new = messages[len(prev_messages):], and delta_tail is rendered from the gateway's own stored messages, so no hand-rolled delimiter logic is needed and the client's re-serialization of prior turns never enters the fed tokens:

assert messages[:len(prev_messages)] == prev_messages and tools == prev_tools  # linearity check
new        = messages[len(prev_messages):]
r_prev     = render(prev_messages,       tools=prev_tools, add_generation_prompt=False)
r_ext      = render(prev_messages + new, tools=prev_tools, add_generation_prompt=True)
delta_tail = r_ext[len(r_prev):]                                               # new messages + generation prompt

Both renders share the literal prev_messages prefix and the same tools, so r_ext starts with r_prev by construction — the prior turns are re-rendered from the gateway's canonical dicts, not the client's replay. Deciding linearity at the message level (rather than by a raw token-prefix check on the client's render) is what keeps the healer in agreement with TrajectoryManager: a harness that replays a prior tool call with its arguments re-keyed (the adapter emits the wire form sort_keys=True while the manager keeps the model's emission order) or its JSON re-spaced produces a message that is still == to what the gateway stored — so the manager, matching by dict equality, keeps it linear — but it renders to different tokens under an order-sensitive template. A token-prefix check would wrongly flag that cosmetic re-serialization as non-linear and reset; matching at dict equality keeps it healed and reserves the non-linear path for a genuine edit/branch. This relies only on chat templates being prefix-consistent with add_generation_prompt=False (rendering [m0..mk] is a prefix of [m0..mk, mk+1]), which holds for turn-delimited templates such as Qwen and Llama-3.

What's in this PR

• rollout_gateway/linear.py — LinearHealer, the per-sid healer. Tokenizer-free and torch-free; it only calls the injected Renderer. heal returns the ids to feed the backend; commit advances per-sid state after a turn is actually recorded. heal/commit are split because commit needs the served output_ids (which only exist after generation) and must fire only for turns that are truly recorded — a failed or retried generation must not advance the canonical prefix — keeping healer state in lockstep with TrajectoryManager.
• rollout_gateway/gateway.py — RolloutGateway.__init__ gains history_mode and linear_on_nonlinear. In linear mode it builds one shared LinearHealer and injects it into every co-mounted adapter (like the shared TrajectoryManager), so a session's canonical state is coherent regardless of which wire protocol its turns arrive on. fork_threshold_tokens is ignored in linear mode (forking is disabled by construction) and a warning is logged if both are set.
• rollout_gateway/adapters/common.py — BaseAdapter heals between render and backend.generate in _run_turn, feeds the healed ids to both generate and the TurnRecord, and calls commit before record_turn. Session teardown drops per-sid healer state.
• rollout_gateway/__init__.py — exports LinearHealer.
• tests/rollout_gateway/test_linear_healer.py — new correctness suite (below).

Config surface

• history_mode: "tree" (default, current behavior) | "linear". Gateway-global, plumbed like fork_threshold_tokens; linear and tree sessions do not coexist within a run.
• linear_on_nonlinear: "reset" (default) | "error" | "passthrough" — behavior when a turn breaks the append-only assumption.

When the assumption breaks

The message-level linearity check fails when the incoming history is not an append-only extension of what the gateway stored: a dropped/edited/branched prior message (a genuine dict change, not merely a re-serialization), context compaction, a changed tools schema, or an LLM-client retry that re-issues the same prompt and produces a second generation from the same prefix. It also falls back when the closer probe cannot isolate the between-message glue (a message-type-dependent template — the close_unresolved counter). Behavior is controlled by linear_on_nonlinear:

• "reset" (default) — drop per-sid state and re-anchor to the current render, treating the jump as the start of a fresh linear segment. The turns before and after each stay drift-free; only the single jump turn conditions on the client's render. Correct for benign, agent-controlled jumps; increments a counter so a run that resets constantly is visible.
• "error" — raise and fail the rollout, for runs that want a hard guarantee the assumption holds.
• "passthrough" — stop healing this session and route the rest through the standard tree/FORK path (today's behavior) for the remainder of the session.

Validation

Unit tests (tests/rollout_gateway/test_linear_healer.py) drive the real LinearHealer + TrajectoryManager through the same heal → generate → commit pipeline BaseAdapter._run_turn uses, with a template-shaped fake renderer that reproduces assistant-only drift: renderer-driven drift forks without healing and collapses to one sample with it (loss mask covering every generated turn, served logprobs preserved); multi-turn drift stays one sample; the closer probe isolates the glue for both text and tool-call turns and does not double the closer when the served output already ends with the stop token; and the three linear_on_nonlinear paths behave as specified. The template-level assumptions (generation-prompt tokens equal the assistant open; prefix-consistency turn to turn; the closer probe returns <|im_end|>\n for both text and tool-call turns and is a genuine suffix of a tool-call r_prev; served ids end with the stop token) were confirmed against the live Qwen3-Coder-30B tokenizer/template.

Real-trace reconstruction. We reconstructed, at the token level, the per-turn prompt streams from real recorded multi-turn SWE-agent sessions (Qwen3-Coder-30B) and replayed them through the real TrajectoryManager two ways: the drifted prompts as they actually happened, and the canonical served prefix + drift-free new-observation tail that healing feeds. The raw replay reproduced the recorded fork counts exactly (fidelity check). Healing collapsed the multi-record baseline sessions to a single record each while leaving the trained (loss_mask=1) token count identical — so the whole difference is redundant re-emitted context that healing removes. Training latency scales ~linearly with the sequence length forwarded, not the record count, so the total-token reduction is the latency-relevant figure, and it is smaller than the record reduction because early forks carry short prefixes while the redundancy is dominated by late, large-context forks. A small number of sessions did not collapse to a single record; the observed cases were benign harness-side history rewrites — an LLM-client retry re-issuing the same prompt, or a harness recovering a hallucinated tool name into a different tool with restructured arguments — which the message-level linearity check correctly treats as non-linear and on_nonlinear="reset" re-anchors into separate linear segments.

Observability

LinearHealer tracks per turn: healed_turns and healed_prefix_tokens (turns healed and canonical prefix tokens spliced); nonlinear (message-level linearity failures — a genuine history edit/branch or tools change, not cosmetic re-serialization); close_unresolved (turns where the closer probe couldn't isolate the glue and fell back). A high nonlinear rate means the "linear" assumption is wrong for that harness and tree mode is the better fit — what happened to those jumps is the (known) linear_on_nonlinear mode, so there is no separate reset counter.

These are surfaced two ways. LinearHealer.counters is the run-cumulative total across every session the shared healer has seen. Per session, finish_session pops that session's own counters (LinearHealer.pop_stats, a fixed schema with zeros for counters that never fired) and attaches them to every record's metadata under linear_healer, riding the existing extra_metadata → TraceRecord.metadata seam — so a downstream consumer sees each rollout's healing stats without touching the healer. The per-session counters survive reset/passthrough jumps within a session (only per-sid healing state is cleared on a jump), so the final record still reports the full nonlinear total for that session.

Non-goals

• Message-level rewrites where a prior message's dict changes semantically (not just its key order or JSON spacing): treated as non-linear and re-anchored/failed per linear_on_nonlinear. Cosmetic re-serialization that keeps the dict == is handled — the linearity check matches at dict equality — but an actual content change (e.g. a harness that recovers a hallucinated tool name into a different tool with restructured arguments) is out of scope for healing.
• Sub-agent / parallel-branch capture: that is what tree mode is for; linear mode re-anchors or falls back rather than trying to represent it.

[luyuzhe111]

I think committing the design doc is a good idea and maybe we should do this in future PRs too.

For the long-term design tho, I had codex compare this design with the session tree design from Miles TITO v2. Its recommended long term design is as follows:

Recommended Long-Term Architecture

The strongest long-term design is to store exact token checkpoints directly in the trajectory tree instead of combining:

LinearHealer + message tree + post-hoc token reconciler

Each generated assistant node should own:

MessageNode
├── canonical messages
├── exact prompt token IDs
├── exact generated token IDs
├── generated-token logprobs
├── loss-mask ownership
└── parent checkpoint

For each incoming request:

1. Match it to the deepest compatible canonical message checkpoint.
2. Reuse that node's exact token sequence.
3. Incrementally tokenize only the appended message suffix.
4. Apply a model-specific boundary merge.
5. Generate from the assembled exact-token prompt.
6. Attach the new exact-token checkpoint as a child.

This gives:

• A single chain for linear rollouts.
• Natural branches for sub-agents and genuine rewrites.
• Explicit rollback or deduplication for retries.
• No tokenization-artifact forks.
• No need for post-hoc REALIGN.
• One shared identity model for messages and tokens.

This architecture resembles the newer Miles tree-based TITO session design and fits the repository's existing MessageNode tree well.

It makes good sense to me. Wondering what do you think of this? Do you think there is a path towards this long-term design with the current PR? should we incorporate this long-term vision to this design doc too?

[hellodanylo]

Yeah, the tree version of the healer would be roughly what Miles calls v2. When we have non-linear harnesses ready to test, we can implement the tree version of the healer as well. It's not a lot of complexity, but we need a good harness to test on.

[Linbo-Liu]

To get the token delta, this blog proposes to render a dummy message header + new message, instead of rendering the original messages, to increase efficiency.