[core] feat: support rollback for rewritten last assistant messages

[zackcxb]

What does this PR do?

This PR makes Gateway sessions reuse an existing chain when a client rewrites
only that chain's latest Assistant message.

The Gateway records the token/message/media boundary immediately before each
Assistant turn. When the next request preserves the earlier prefix but replaces
that Assistant turn, it rolls the chain back to the recorded boundary,
re-encodes the replacement suffix, and continues generation on the same chain.
This avoids materializing a redundant trajectory while preserving token IDs,
response masks, rollout log probabilities, and multimodal inputs.

Last-Assistant rollback is enabled by default. Set
actor_rollout_ref.rollout.custom.agent_framework.enable_last_assistant_rollback=false
to retain the previous split-on-rewrite behavior.

Checklist Before Starting

• Search for similar PRs or issues and paste at least one relevant link here: [BREAKING][core] refactor: unify agent runtime and training interfaces #83, [BREAKING][core] refactor: centralize distributed session logging #85, and [example]fix: unify openYuanrong sandbox name and fix retry error #87 above
• Format the PR title as [{modules}] {type}: {description} (checked by CI)

Test

Automated validation on upstream/main at ecdddea (#87):

PYTHONPATH=/tmp/uni-agent-python-isolation:/tmp/uni-agent-verl-78bba31 \
python -m pytest \
  tests/uni_agent/gateway \
  tests/uni_agent/framework/test_generate_sequences_on_cpu.py -q

166 passed, 4 warnings in 110.78s

All hooks pass for files changed by this PR:

pre-commit run --from-ref upstream/main --to-ref HEAD --show-diff-on-failure

ruff: pass
ruff-format: pass
mypy: pass
compileall: pass

pre-commit run --all-files additionally reformats
examples/blackbox_recipes/sandbox_client.py, an unchanged file introduced by
#87. That unrelated baseline formatting change is intentionally excluded from
this focused PR.

CC rollout experiment

Full rollout configuration:

model: Qwen/Qwen3.5-4B
samples: 8
rollouts per sample: 4
sessions: 32
max concurrent sessions: 4
max agent turns: 100
prompt length: 90112
response length: 49152
tensor parallelism: 4
GPU memory utilization: 0.90

Results:

Metric	Result
Successful sessions	32 / 32
Failed session UIDs	0
Sessions that triggered rollback	32 / 32 (100%)
Total rollbacks	229
Mean rollbacks per session	7.16
Rollback count P25 / median / P75 / max	4 / 5.5 / 7 / 27
Dropped trainable tokens	165,733
Mean dropped trainable tokens per rollback	723.7
Output trajectories	50
Extra trajectories	18
max_response_length materializations	18
Non-length split trajectories	0
Resolved SWE-bench samples	3 / 8 (37.5%)

All 50 trajectory records satisfied:

len(response_ids) == len(response_mask) == len(response_logprobs)

The full 8x4 run was collected on the #83-rebased rollback implementation.
After rebasing onto #85/#87, the rollback implementation in session.py is
byte-for-byte identical except that the feature default changes from false to
true in this PR.

A fresh #87 smoke run used MAX_SAMPLES=2, N=1, and the same model/token
budgets. It completed 2/2 sessions, produced three trajectories (one
max_response_length materialization), and all response ID/mask/log-probability
arrays remained aligned.

The CC recipe currently needs a separate post-#83 reward-wiring compatibility
fix. That fix was applied only in the local rollout workspace and is not part of
this PR.

Observed task-level warnings included agents reaching the configured 100-turn
limit, one malformed Tool-call parse, and one 143,813-token prompt exceeding the
140,288-token model context. These were not rollback or token-alignment failures;
the full experiment still finalized all 32 sessions.

API and Usage Example

Rollback is enabled when the option is omitted:

actor_rollout_ref:
  rollout:
    custom:
      agent_framework:
        gateway_count: 1

To preserve the previous behavior and split a rewritten latest Assistant into a
new chain:

actor_rollout_ref:
  rollout:
    custom:
      agent_framework:
        enable_last_assistant_rollback: false

Design & Code Changes

• Capture a stable LastAssistantStart boundary before every Assistant
  generation, including response IDs, response masks, log probabilities,
  message-history length, media lengths, and the prefix hash.
• Rank exact-prefix and rollback candidates by reusable service depth, preferring
  exact continuation on ties and refusing ambiguous deepest rollback matches.
• On rollback, truncate the selected chain to its recorded boundary and
  re-encode only the incoming replacement suffix.
• Preserve response-log-probability alignment and fail loudly if stored token
  truth is already inconsistent.
• Track rollback_count and
  rollback_dropped_trainable_tokens_total in the session snapshot.
• Forward the configuration through framework entry -> Gateway actor -> Gateway
  session. The default is true; explicit false remains supported.
• Cover exact matches, ambiguous candidates, reserved chains, multimodal suffix
  re-encoding, response-budget materialization, log-probability alignment, and
  explicit opt-out behavior.

This PR intentionally does not include the independent CC reward compatibility
fix or rollout snapshot logging used to collect the experiment metrics.

Checklist Before Submitting

• Read the Contribute Guide
• Run pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always
  • All PR files pass; the all-files command finds the unrelated [example]fix: unify openYuanrong sandbox name and fix retry error #87 baseline formatting issue documented above.
• Add or update docs/examples for user-facing changes
• Add tests or explain why tests are not practical
• Confirm the PR title matches the required format
• Confirm the placeholder text in this template has been replaced with real content

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[yyDing1]

LGTM