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

docs/source/concepts/gateway-and-trajectories.md

@@ -78,6 +78,13 @@ One session may contain multiple model turns. Tool observations are encoded as c
 
 Concurrent requests may create multiple chains within one session. Chains sharing a message prefix reuse the same encoded context where possible, then materialize as separate trajectories during finalization.
 
+When a client rewrites only the most recent Assistant message, the Gateway rolls
+the matching chain back to the start of that Assistant turn and re-encodes the
+replacement suffix. This preserves token, mask, and rollout-log-probability
+alignment without materializing a redundant trajectory. The behavior is enabled
+by default and can be disabled with
+`actor_rollout_ref.rollout.custom.agent_framework.enable_last_assistant_rollback=false`.
+
 ## Reward Flow
 
 The built-in Task Runner posts:
@@ -142,6 +149,9 @@ actor_rollout_ref.rollout.custom.agent_framework
 Important knobs include:
 
 - `gateway_count`: Gateway actor pool size.
+- `enable_last_assistant_rollback`: reuses a chain when only its latest Assistant
+  message is rewritten. Defaults to `true`; set it to `false` to preserve the
+  previous split-on-rewrite behavior.
 - `agent_runners`: Runner import paths and arguments.
 - `dispatch_mode`: inline async execution or Ray tasks.
 - `max_concurrent_sessions`: per-Runner concurrency limit.

examples/blackbox_recipes/sandbox_client.py

@@ -84,8 +84,7 @@ def init(cls) -> None:
         token = os.getenv("OPENYUANRONG_TOKEN")
         if not server or not token:
             raise ValueError(
-                "OPENYUANRONG_SERVER_ADDRESS and OPENYUANRONG_TOKEN "
-                "environment variables must be set for sandbox"
+                "OPENYUANRONG_SERVER_ADDRESS and OPENYUANRONG_TOKEN environment variables must be set for sandbox"
             )
         # Reverse tunnel TLS verify
         os.environ["TUNNEL_SSL_VERIFY"] = os.getenv("OPENYUANRONG_TUNNEL_SSL_VERIFY", "0")
@@ -242,4 +241,4 @@ async def cleanup(self) -> None:
                     logger.info("sandbox %s already stopped", sandbox_id)
             except Exception as e:
                 logger.warning("Failed to kill sandbox %s: %s", sandbox_id, e)
-            self._sandbox = None
+            self._sandbox = None

tests/uni_agent/framework/test_generate_sequences_on_cpu.py

@@ -112,15 +112,22 @@ async def _build_framework_with_agent_runners(
 
 
 @pytest.mark.parametrize(
-    ("data_config", "expected_chat_template_kwargs"),
+    ("data_config", "rollback_config", "expected_rollback", "expected_chat_template_kwargs"),
     [
-        ({}, {}),
-        ({"apply_chat_template_kwargs": {"thinking": True}}, {"thinking": True}),
+        ({}, {}, True, {}),
+        (
+            {"apply_chat_template_kwargs": {"thinking": True}},
+            {"enable_last_assistant_rollback": False},
+            False,
+            {"thinking": True},
+        ),
     ],
 )
 def test_build_gateway_manager_wires_gateway_config_defaults(
     monkeypatch,
     data_config,
+    rollback_config,
+    expected_rollback,
     expected_chat_template_kwargs,
 ):
     from omegaconf import OmegaConf
@@ -152,7 +159,12 @@ def __init__(self, *, llm_client, gateway_count, gateway_actor_config):
                     "prompt_length": 128,
                     "response_length": 64,
                     "multi_turn": {"format": "hermes"},
-                    "custom": {"agent_framework": {"gateway_count": 2}},
+                    "custom": {
+                        "agent_framework": {
+                            "gateway_count": 2,
+                            **rollback_config,
+                        }
+                    },
                 },
             },
         }
@@ -166,6 +178,7 @@ def __init__(self, *, llm_client, gateway_count, gateway_actor_config):
     assert captured["gateway_actor_config"].prompt_length == 128
     assert captured["gateway_actor_config"].response_length == 64
     assert captured["gateway_actor_config"].tool_parser_name == "hermes"
+    assert captured["gateway_actor_config"].enable_last_assistant_rollback is expected_rollback
     assert isinstance(captured["gateway_actor_config"].apply_chat_template_kwargs, dict)
     assert captured["gateway_actor_config"].apply_chat_template_kwargs == expected_chat_template_kwargs
 

tests/uni_agent/gateway/test_gateway_actor_on_cpu.py

@@ -49,6 +49,43 @@ def test_gateway_actor_config_rejects_non_positive_response_length(response_leng
         GatewayActorConfig(tokenizer=FakeTokenizer(), response_length=response_length)
 
 
+@pytest.mark.parametrize("value", ["true", 1, None])
+def test_gateway_actor_config_rejects_non_bool_last_assistant_rollback(value):
+    from uni_agent.gateway.config import GatewayActorConfig
+
+    with pytest.raises(ValueError, match="enable_last_assistant_rollback must be a bool"):
+        GatewayActorConfig(tokenizer=FakeTokenizer(), enable_last_assistant_rollback=value)
+
+
+def test_gateway_actor_config_enables_last_assistant_rollback_by_default():
+    from uni_agent.gateway.config import GatewayActorConfig
+
+    assert GatewayActorConfig(tokenizer=FakeTokenizer()).enable_last_assistant_rollback is True
+
+
+@pytest.mark.asyncio
+async def test_gateway_actor_forwards_last_assistant_rollback_to_session():
+    from uni_agent.gateway.config import GatewayActorConfig
+    from uni_agent.gateway.gateway import _GatewayActor
+
+    actor = _GatewayActor(
+        GatewayActorConfig(tokenizer=FakeTokenizer()),
+        SequencedBackend(["BAD", "FIXED"]),
+    )
+    actor._server_base_url = "http://test"
+    await actor.create_session("rollback-enabled")
+    prompt = [{"role": "user", "content": "run"}]
+    await actor._handle_openai_chat_completions("rollback-enabled", {"messages": prompt})
+    await actor._handle_openai_chat_completions(
+        "rollback-enabled",
+        {"messages": [*prompt, {"role": "user", "content": "user_error"}]},
+    )
+
+    state = await actor.get_session_state("rollback-enabled")
+    assert state["active_chain_ids"] == [1]
+    assert state["rollback_count"] == 1
+
+
 @pytest.mark.asyncio
 async def test_gateway_actor_max_tokens_clamped_to_remaining_response_budget():
     """Continuation requests clamp ``max_tokens`` to the selected chain budget."""
@@ -773,7 +810,7 @@ async def test_gateway_actor_continuation_with_tool_returned_image_appends_media
     import uni_agent.gateway.session.codec as codec_mod
     from uni_agent.gateway.config import GatewayActorConfig
     from uni_agent.gateway.gateway import _GatewayActor
-    from verl.utils.chat_template import apply_chat_template, initialize_system_prompt
+    from verl.utils.tokenizer.chat_template import apply_chat_template, initialize_system_prompt
 
     monkeypatch.setattr(codec_mod, "_extract_tool_calls_with_sglang_or_vllm", fake_tool_call_dispatch)
     processor = FakeProcessor()