fix(agent): realign from divergence using previous response length

[YeonwooSung]

Background

_SampleBuilder heals TITO / chat-template token drift by either REALIGN (overwrite the drifted tail in-place) or FORK (close the builder and start a new Sample). REALIGN is supposed to be a cheap repair of a short previous response.

Fixes #2338.

Problem

Two bugs in slime/agent/trajectory.py, both visible in the reporter's 3-turn example (default fork_threshold=1024): a 61-token transition turn REALIGN'd a 28k-token previous response, and 184–383 already-trained tokens before the BPE split were zeroed with it.

1. REALIGN gated on this turn's output_ids, not the previous response

classify_token_drift used len(turn.output_ids) to decide REALIGN vs FORK. REALIGN rewrites the previous response (it zeros that span and splices the new prompt suffix). A short follow-up therefore realigned — and discarded — a long trained span.

In the report: new_resp_len=61 < 1024 → REALIGN, zeroed_resp_tokens=28404.

2. _align_to_prompt overwrote from last_response_start_idx, not the divergence

Even when REALIGN is correct, zeroing from the start of the previous response also wipes tokens that still match the new prompt (the "innocent prefix" before the BPE split). The reporter measured 184 / 237 / 332 / 383 trained tokens killed before realign_at across four REALIGNs, with killed_masked_lm0 == 0 (the whole zeroed range was loss_mask=1 model output, not tool/history).

Solution

• classify_token_drift: REALIGN only when the divergence sits inside the most-recent response and both sacrificed_len = len(tokens) - last_response_start_idx (previous response) and len(turn.output_ids) are < fork_threshold. Otherwise FORK. _try_merge_assistant_rewrite is unchanged (different contract: merging a short abandoned assistant rewrite).
• _align_to_prompt(prompt_ids, realign_at): overwrite tokens / loss_mask / logprobs from the divergence only. Tokens before realign_at keep their existing mask and logprobs.
• append_turn passes the same _common_prefix_len cut that classify already computed (recomputed at append time; tokens are unchanged between the two calls).

Tests

tests/test_agent/test_trajectory_manager_branching.py:

• test_4_6_drift_B1_threshold_boundary now pins the new gate: REALIGN iff both sides are < threshold; a long previous response forks even if this turn is short (the old threshold=3, r1=4, r2=2 case, which used to REALIGN). After REALIGN, the innocent prefix keeps loss_mask=1 and only the suffix from realign_at is zeroed.
• test_2_4_drift_case_B1_short_replaces updated so the trained token before the last-token split stays loss_mask=1.

python -m pytest tests/test_agent/test_trajectory_manager_branching.py -q --tb=short
# 35 passed