fix(agent): restore a re-rendered assistant echo before rendering the prompt

[pollyloop]

Issue

In coding-agent RL, an assistant turn that calls Edit or Write can end up contributing zero trained tokens. The run reports nothing: no exception, no warning, and the same number of emitted Samples. The only trace is a smaller loss_mask sum, which nothing watches.

Measured across a 50-step GRPO run on SWE-rebench problems with the Claude Code harness (8xH20 per node; 4224 trained samples, 171,539 assistant turns, 38.8M assistant tokens): 44.3% of all turns containing an Edit or Write should have been trained and were not. Bash and Read turns were dropped less often than average (0.6x), so the model kept learning to read code and run tests while losing supervision on the step where it applies the fix.

Reproduce on a clean checkout — no GPU, sandbox, or model needed (repro_edit_write_turn_dropped.py for the demonstration, verify.py for the measurements in Validation below): https://gist.github.com/pollyloop/f7cafc2dd0854587e82acb4e31281e1b

scenario                            samples   mask   Edit   Write   Bash
------------------------------------------------------------------------
byte-exact echo (ideal)                   1      6      1       1      1
harness replayed it re-rendered           1      2      0       0      1

Cause

A harness may replay a prior assistant tool call re-rendered rather than byte-identical — the term _try_merge_assistant_rewrite already uses for this. The replay is not just formatted differently; the argument values themselves are altered. With the Claude Code CLI, two things happen to the echo:

1. it fills in a tool-schema default the model omitted — replace_all: False appears on an Edit that never carried that key;
2. string arguments are right-stripped per line, so trailing whitespace inside a code payload disappears.

Neither changes what the tool does, but both break ==.

That matters twice over, because a rewritten echo reaches the manager through two channels at once. _run_turn builds both from the same translated list:

translated, tools_schema = self._translate(body)                    # the harness's bytes
prompt_ids = _render_token_ids(translated, ...)                     # -> layer 2
...
self.manager.record_turn(sid, turn=turn, prompt_messages=translated, ...)   # -> layer 1

The test suite already names these layers: layer 1 is the routing tree, keyed on message identity; layer 2 is linearization, keyed on the token prefix. The rewrite breaks both, and each one on its own is enough to lose the turn:

Layer 1 — the routing tree. _find_mount_point matches history by strict dict equality (child.message == msg, trajectory.py:361), so the walk stops one level above the assistant node. _try_merge_assistant_rewrite finds exactly one short assistant leaf there and absorbs it — rewritten_node.turn = None (:423) — and _chain_to_samples filters turn is None out of asst_nodes. The generated tokens survive only as loss_mask=0 context.

Layer 2 — the token prefix. prompt_ids was rendered from the harness's arguments, so it no longer extends the tokens the builder holds. classify_token_drift sees the divergence inside the most recent response span and returns REALIGN, and _align_to_prompt overwrites that span with the prompt tail at loss_mask=0 (:216-224). Same outcome, different route.

This is why the fix has to act on translated itself rather than on the matching predicate: both channels are downstream of it.

This is why the file-writing tools are hit: they carry at least one of those two things, and either is enough on its own. Edit has both. Write has no schema default at all (injection rate 0%), yet its multi-line content is right-stripped just the same — and it is dropped at the same rate. Bash/Read arguments pass through untouched.

Looking past this class of difference is already the intent here. tool_call_dict stores arguments as a dict rather than a JSON string specifically so that key order cannot break the comparison — its docstring says "the trajectory manager matches history by dict equality, so a sampled leaf and its replayed echo compare equal regardless of key order." A filled-in default and per-line trailing whitespace are the same kind of difference — no effect on what the tool does, but fatal to == — and are currently not absorbed.

Fix

TrajectoryManager.restore_generated_messages(sid, messages) — swap a recognized echo back to the message we generated. The adapter calls it in _run_turn, between _translate and _render_token_ids:

translated, tools_schema = self._translate(body)
translated = self.manager.restore_generated_messages(sid, translated)
prompt_ids = _render_token_ids(translated, tok, tools=tools_schema, add_generation_prompt=True)

That placement is what makes it fix both layers at once. It is after _translate, so there is something to compare against the tree; and before _render_token_ids, so prompt_ids is rendered from our own bytes rather than the harness's. The turn then classifies CLEAN and keeps its gradient.

Fixing only layer 1 would not be enough: prompt_ids would still come from the harness's bytes, layer 2 would still return REALIGN, and the same span would still be zeroed (Validation item 1, third row).

_is_same_tool_call_echo decides what counts as an echo. It does not try to reverse the harness's edits — nothing is recovered from the echo, and no attempt is made to guess how much whitespace was stripped. Instead it compares the two messages through a coarser view that both sides map onto:

• an argument equal to the default its own tool schema declares is dropped, since omitting an argument and passing its default are the same call. The schema is already in hand at the call site (_translate returns it one line earlier), so restore_generated_messages takes it and _tool_defaults reads properties.*.default out of it;
• strings are right-stripped per line.

If both sides agree there, the echo is the same call. Everything outside tool_calls — role, content, tool name, call count, any other key — must still match exactly.

With no schema passed, only the whitespace projection applies — the comparison gets stricter, which costs a restore rather than making a wrong one.

Recovering the original bytes is then trivial and lossless, because it does not come from the echo at all: out[depth] = echoes[0].message puts back the message we stored when the model produced it, whitespace and all.

_find_mount_point is therefore left alone, keeping strict dict equality: once the restore has run, record_turn receives the generated message and the exact-equality walk succeeds on its own (measured — Validation item 4).

Three files, +472 −0; no existing line is modified.

Validation

Items 1, 3 and 4 are a script in the gist — python3 verify.py inside any checkout reproduces those three tables, and it detects whether the checkout is patched (on a pristine one it shows the bug and skips the fix-dependent parts). Items 2 and 5 are the test suite, so CI runs them here. Item 6 is the one thing you cannot re-run: it came from a cluster run, and the data is not public — treat it as corroboration of scale, not as evidence the fix is correct. That case rests on 1–5.

1. Is the restore necessary, and is it sufficient? (verify.py, item 1.) Four configurations. Each layer is observed from the outside — a read-only wrapper around _try_merge_assistant_rewrite (layer 1) and classify_token_drift (layer 2) that records which branch is taken and forwards unchanged, so what runs is the shipped code:

configuration	mask	Edit	Write	Bash	layer 1	layer 2
byte-exact echo (ideal)	6	1	1	1	no merge	clean
harness rewrote the echo (main today)	2	0	0	1	2 merges fired	not reached
+ recognize echo, no restore	2	0	0	1	no merge	realign
+ restore before render	6	1	1	1	no merge	clean

Row 2 is main: layer 1 swallows the turn, which is why layer 2 never sees it. Row 3 is the layer-1-only fix — the merge now correctly declines, so the turn does reach layer 2, which diverges instead and zeroes the same span. The numbers do not budge. Row 4 fixes both layers and matches the ideal. Bash is unchanged throughout, as it must be.

2. Does it ever restore something it shouldn't? This is the risk the fix introduces, so it gets its own test (test_3_13). A harness that self-corrects or retries produces exactly these:

incoming assistant message	restored?
the real echo	yes
a different old_string	no
a different file	no
a different tool (Write, not Edit)	no
plain text, no tool call	no
leading whitespace differs	no
MultiEdit whose nested payload differs	no

The default-dropping deserves its own matrix, over a tool whose default is False (Edit.replace_all) and one whose default is True (Fetch.follow_redirects below) — test_3_11 covers both:

The observed edit is always the same shape — the model omitted the argument and the harness filled one in — so what matters is whether the filled-in value is the declared default:

declared default	harness filled in	same call?
False	False	yes
False	True	no
True	True	yes
True	False	no

Rows 1 and 3 are the same case — the harness filled in exactly what the schema declares, whichever value that is — and both are recognized because the declared value is read rather than assumed. Rows 2 and 4 fill in something the schema does not declare, which is a genuinely different call, so they keep taking the existing path.

Worth noting for the bound on all of this: restore_generated_messages only rewrites prompt_messages, the context for the next turn. Training uses turn.output_ids — what the model actually sampled — which the function never touches.

3. Does the comparison have to look past both kinds of edit? Yes. Narrowing the view so it looks past only one of them (verify.py, item 3) leaves turns untrained:

comparison looks past	Edit	Write
both (this PR)	1	1
trailing whitespace only	0	1
filled-in defaults only	0	0

Write only needs trailing whitespace looked past, since its schema declares no default to fill in. Edit needs both, because the harness applies both to it at once — look past one and its old_string still compares unequal, so the echo goes unrecognized and the turn is lost exactly as on main.

4. Would relaxing _find_mount_point help too? No (verify.py, item 4). Instrumenting a session where the predicate is relaxed as well, every echo it recognizes was already handled by the restore: 0 of its decisions came from _find_mount_point. Once the adapter restores before rendering, record_turn receives the generated message and exact equality succeeds on its own, so record_turn's core path is left untouched.

5. Tests. Six new cases in tests/test_agent/test_trajectory_manager_branching.py, in the existing 3.x style:

• test_3_9 — an Edit echo (both kinds of edit at once) is restored, the caller's list is not mutated, and the turn keeps its TurnRecord and trains;
• test_3_10 — a Write echo differing only by the right-stripping, so each kind of edit is covered on its own;
• test_3_11 — a tool whose declared default is True, where filling in True is the same call and filling in False is not: the reverse of Edit.replace_all;
• test_3_12 — with no schema passed, a filled-in default is no longer recognized and the echo is left alone, while whitespace-only echoes still restore;
• test_3_13 — the rejection table above;
• test_3_14 — an end-to-end 5-turn session along the adapter's path with the harness rewriting every echo.

All six fail on main (AttributeError: no attribute restore_generated_messages) and pass with this change. Run the way CI runs them (python tests/test_agent/<file>.py, CPU torch, pip install -e . --no-deps):

	main	this PR
test_trajectory_manager_branching.py	35 passed	41 passed
test_adapters.py	14 passed, 1 skipped	14 passed, 1 skipped
test_harness.py	9 passed	9 passed
test_agent_rollout_cpu.py	4 passed	4 passed
pre-commit run	—	clean

test_agent_rollout_cpu.py covers the adapter change as well: it drives the real _run_turn over HTTP loopback, where restore_generated_messages is called 4 times. That is also why the repro is a gist rather than a new file under tests/ — agent-test's file list is enumerated in pr-test.yml.j2, so a new file would not actually run without editing the generated workflow.

6. A real rollout, before and after. Two runs over the same 16 SWE-rebench problems with the Claude Code harness. Same model, same sampling config, same problems; the only difference is whether the restore is applied. The two runs were checked rollout by rollout to confirm they saw the same problems in the same order.

	restore off	restore on
Edit turns with loss_mask=0	59.4%	3.7%
Write turns with loss_mask=0	50.0%	0%
assistant tokens with loss_mask=0	30.6%	6.5%
samples emitted	96	86 (−10.4%)
truncated sessions	4/16	4/16 (unchanged)

Of the 285 rewritten tool calls in these runs, all 285 were recognized and restored. The fall in sample count is independent corroboration rather than a side effect: those extra samples were dead-end forks, which is what the rewrite-merge existed to avoid in the first place.

What this does not show: 16 problems cannot measure a change in model quality, and I make no claim there. This is a correctness fix — generated tokens were being excluded from the loss.

Notes

Anything counting drift classifications will report more of them. Every turn that survives to linearization gets classified by classify_token_drift as CLEAN, REALIGN or FORK. A turn the merge swallows never gets that far — turn = None drops it from asst_nodes before _split_chain_into_builders runs — so today it is absent from those counts entirely. This PR keeps those turns, so they are classified for the first time and the totals rise (1 classification before vs 2 after, on a 3-turn session). No turn drifts that did not drift before; turns that were invisible are now visible.

Two cases this does not fix. An argument filled in inside a nested object is not recognized, since only top-level arguments are compared against the schema; and a harness whose tool schemas omit default altogether leaves only the whitespace comparison, so a filled-in default is not recognized either. In both, the echo is left as the harness sent it and the turn is lost exactly as it is on main today — no improvement, but no new harm. The failure direction is the safe one: an echo that goes unrecognized costs that turn its gradient, whereas misrecognizing one would put the wrong message into the next prompt.

The merge mechanism and its threshold are untouched. This PR does not change _try_merge_assistant_rewrite or fork_threshold_tokens; it removes the drift that makes the merge fire, so a re-rendered echo no longer looks like a rewrite and there is no candidate to absorb. Worth noting separately, though: that function's docstring says the merge is "purely a cleanup" and that "forking is always safe". Structurally true, but the merge branch destroys a generated TurnRecord while the fork branch keeps it — so under any harness that re-renders its history, the choice between them is not signal-neutral, and the 1024 default silently decides how much of a turn survives. That assumption may deserve a look on its own; this PR just stops relying on it.

#2174 is open against the same file and will likely conflict. Happy to rebase on top of it, or to hold this until it lands.

Acknowledgement

Special thanks to the person who got this coding-agent RL pipeline running in our internal environment, which is where these measurements come from.

Analysis assisted by Claude Code.