From 1bbe4108c708e0dc452264e464d8665bf3b682cc Mon Sep 17 00:00:00 2001 From: zackcxb Date: Wed, 22 Jul 2026 14:09:51 +0000 Subject: [PATCH 1/2] feat(gateway): rollback rewritten last assistant --- .../concepts/gateway-and-trajectories.md | 10 + .../test_generate_sequences_on_cpu.py | 21 +- .../gateway/test_gateway_actor_on_cpu.py | 39 ++- .../test_session_multiple_chains_on_cpu.py | 307 ++++++++++++++++++ uni_agent/framework/entry.py | 1 + uni_agent/gateway/config.py | 8 + uni_agent/gateway/gateway.py | 2 + uni_agent/gateway/session/session.py | 301 ++++++++++++++--- 8 files changed, 630 insertions(+), 59 deletions(-) diff --git a/docs/source/concepts/gateway-and-trajectories.md b/docs/source/concepts/gateway-and-trajectories.md index ef770909..e543c67e 100644 --- a/docs/source/concepts/gateway-and-trajectories.md +++ b/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. diff --git a/tests/uni_agent/framework/test_generate_sequences_on_cpu.py b/tests/uni_agent/framework/test_generate_sequences_on_cpu.py index 42ce5bdc..fad5b3ec 100644 --- a/tests/uni_agent/framework/test_generate_sequences_on_cpu.py +++ b/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 diff --git a/tests/uni_agent/gateway/test_gateway_actor_on_cpu.py b/tests/uni_agent/gateway/test_gateway_actor_on_cpu.py index d8b7e49b..6435e4d6 100644 --- a/tests/uni_agent/gateway/test_gateway_actor_on_cpu.py +++ b/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() diff --git a/tests/uni_agent/gateway/test_session_multiple_chains_on_cpu.py b/tests/uni_agent/gateway/test_session_multiple_chains_on_cpu.py index e92058a1..c7f52676 100644 --- a/tests/uni_agent/gateway/test_session_multiple_chains_on_cpu.py +++ b/tests/uni_agent/gateway/test_session_multiple_chains_on_cpu.py @@ -2,6 +2,7 @@ from types import SimpleNamespace import pytest +import torch from fastapi import HTTPException from tests.uni_agent.support import FakeProcessor, FakeTokenizer, SequencedBackend, fake_vision_info_extractor @@ -33,6 +34,7 @@ def _session( *, response_length: int | None = None, sampling_params: dict | None = None, + enable_last_assistant_rollback: bool = False, processor=None, vision_info_extractor=None, tool_parser_name: str | None = None, @@ -47,6 +49,7 @@ def _session( ), response_length=response_length, sampling_params=sampling_params, + enable_last_assistant_rollback=enable_last_assistant_rollback, ) @@ -57,6 +60,15 @@ def test_gateway_session_rejects_non_positive_response_length(response_length): _session("invalid-response-length", response_length=response_length) +def test_gateway_session_enables_last_assistant_rollback_by_default(): + session = GatewaySession( + SessionHandle(session_id="rollback-default"), + MessageCodec(FakeTokenizer()), + ) + + assert session._enable_last_assistant_rollback is True + + async def _run(session: GatewaySession, backend: SequencedBackend, messages: list[dict], **payload_extra): request = openai_to_internal( {"model": "dummy-model", "messages": messages, **payload_extra}, @@ -114,6 +126,20 @@ def release_call(self, index: int) -> None: self.calls[index]["release"].set() +class _ExpandedImageTokenProcessor(FakeProcessor): + """Mirror vision processors that expand one image into multiple model tokens.""" + + def __call__(self, **kwargs): + output = super().__call__(**kwargs) + if kwargs.get("images"): + expanded_ids = [] + for token_id in output["input_ids"][0].tolist(): + expanded_ids.extend([token_id, token_id] if token_id == self.image_token_id else [token_id]) + output["input_ids"] = torch.tensor([expanded_ids], dtype=torch.long) + output["attention_mask"] = torch.ones_like(output["input_ids"]) + return output + + def _image_message(url: str, text: str) -> dict: return { "role": "user", @@ -228,6 +254,287 @@ async def test_multiple_chains_context_compaction_starts_new_chain(): assert all(t.response_mask == [1] * len(t.response_ids) for t in trajectories) +@pytest.mark.asyncio +async def test_last_assistant_rollback_reencodes_replacement_suffix_as_masked_context(): + """Drop an abandoned assistant and retain its replacement prompt as context.""" + session = _session( + "rollback-token-truth", + sampling_params={"logprobs": True}, + enable_last_assistant_rollback=True, + ) + first_messages = [{"role": "user", "content": "run mini-swe"}] + rewrite_messages = [ + *first_messages, + {"role": "user", "content": "user_error: missing import"}, + ] + expected_suffix = "user:user_error: missing import\nassistant:" + + await _run(session, _LogprobBackend([("FORMAT_ERROR", "full")]), first_messages) + await _run(session, _LogprobBackend([("FIXED", "full")]), rewrite_messages) + + state = session.snapshot_state() + assert state["active_chain_ids"] == [1] + assert state["rollback_count"] == 1 + assert state["rollback_dropped_trainable_tokens_total"] == len("FORMAT_ERROR") + [chain] = session.active_chains + assert [message["role"] for message in chain.message_history] == ["user", "user", "assistant"] + assert _decode_response_ids(chain.buffer.response_ids) == expected_suffix + "FIXED" + assert chain.buffer.response_mask == [0] * len(expected_suffix) + [1] * len("FIXED") + assert chain.buffer.response_logprobs == [0.0] * len(expected_suffix) + [-0.1] * len("FIXED") + assert len(chain.buffer.response_logprobs) == len(chain.buffer.response_ids) + + +@pytest.mark.asyncio +async def test_last_assistant_rollback_splits_when_history_changes_before_boundary(): + """Split when request drift starts before the latest assistant boundary.""" + session = _session("rollback-split-before-boundary", enable_last_assistant_rollback=True) + backend = SequencedBackend(["A1", "A2", "A3"]) + first_messages = [{"role": "user", "content": "start"}] + second_messages = [ + *first_messages, + {"role": "assistant", "content": "A1"}, + {"role": "user", "content": "original tool result"}, + ] + edited_messages = [ + *first_messages, + {"role": "assistant", "content": "A1"}, + {"role": "user", "content": "edited tool result"}, + {"role": "assistant", "content": "A2"}, + {"role": "user", "content": "continue"}, + ] + + await _run(session, backend, first_messages) + await _run(session, backend, second_messages) + await _run(session, backend, edited_messages) + + state = session.snapshot_state() + assert state["active_chain_ids"] == [1, 2] + assert state["rollback_count"] == 0 + chains_by_id = {chain.chain_id: chain for chain in session.active_chains} + assert chains_by_id[1].message_history[2]["content"] == "original tool result" + assert chains_by_id[2].message_history[2]["content"] == "edited tool result" + assert _decode_response_ids(chains_by_id[1].buffer.response_ids).endswith("A2") + assert _decode_response_ids(chains_by_id[2].buffer.response_ids) == "A3" + + +@pytest.mark.asyncio +async def test_last_assistant_rollback_ambiguous_deepest_boundary_splits_over_shallower_exact(): + """Do not fall back to a shallower exact chain when the deepest rewrite is ambiguous.""" + session = _session("rollback-ambiguous-deepest", enable_last_assistant_rollback=True) + prompt = [{"role": "user", "content": "same prompt"}] + continuation = [ + *prompt, + {"role": "assistant", "content": "A1"}, + {"role": "user", "content": "second turn"}, + ] + rewrite = [*continuation, {"role": "user", "content": "user_error"}] + backend = SequencedBackend(["A1", "A1", "A1", "A2A", "A2B", "FIXED"]) + + await _run(session, backend, prompt) + await _run(session, backend, prompt) + await _run(session, backend, prompt) + await _run(session, backend, continuation) + await _run(session, backend, continuation) + await _run(session, backend, rewrite) + + state = session.snapshot_state() + assert state["active_chain_ids"] == [1, 2, 3, 4] + assert state["rollback_count"] == 0 + chains_by_id = {chain.chain_id: chain for chain in session.active_chains} + assert _decode_response_ids(chains_by_id[1].buffer.response_ids) == "A1" + assert _decode_response_ids(chains_by_id[4].buffer.response_ids) == "FIXED" + + +@pytest.mark.asyncio +async def test_last_assistant_rollback_ignores_shallower_candidate_for_unique_deepest_boundary(): + """Rollback the unique deepest match even when a shallower rewrite also matches.""" + session = _session("rollback-deepest", enable_last_assistant_rollback=True) + prompt = [{"role": "user", "content": "same prompt"}] + continuation = [ + *prompt, + {"role": "assistant", "content": "A1"}, + {"role": "user", "content": "second turn"}, + ] + rewrite = [*continuation, {"role": "user", "content": "user_error"}] + + await _run(session, SequencedBackend(["SHALLOW"]), prompt) + await _run(session, SequencedBackend(["A1"]), prompt) + await _run(session, SequencedBackend(["A2"]), continuation) + await _run(session, SequencedBackend(["FIXED"]), rewrite) + + state = session.snapshot_state() + assert state["active_chain_ids"] == [1, 2] + assert state["rollback_count"] == 1 + chains_by_id = {chain.chain_id: chain for chain in session.active_chains} + assert _decode_response_ids(chains_by_id[1].buffer.response_ids) == "SHALLOW" + deep_text = _decode_response_ids(chains_by_id[2].buffer.response_ids) + assert deep_text.startswith("A1") + assert deep_text.endswith("FIXED") + assert "A2" not in deep_text + + +@pytest.mark.asyncio +async def test_last_assistant_rollback_excludes_reserved_candidate(): + """Do not select an in-flight chain as a rollback target.""" + session = _session("rollback-reserved", enable_last_assistant_rollback=True) + prompt = [{"role": "user", "content": "same prompt"}] + rewrite_messages = [ + *prompt, + {"role": "user", "content": "user_error while busy"}, + ] + + await _run(session, SequencedBackend(["BAD"]), prompt) + session.reserved_chain_ids.add(1) + await _run(session, SequencedBackend(["FIXED"]), rewrite_messages) + + state = session.snapshot_state() + assert state["active_chain_ids"] == [1, 2] + assert state["rollback_count"] == 0 + chains_by_id = {chain.chain_id: chain for chain in session.active_chains} + assert _decode_response_ids(chains_by_id[1].buffer.response_ids) == "BAD" + assert _decode_response_ids(chains_by_id[2].buffer.response_ids) == "FIXED" + + +@pytest.mark.asyncio +async def test_last_assistant_rollback_prefers_longer_service_chain_over_exact_short_chain(): + """Rollback the deeper live chain instead of continuing its exact prefix sibling.""" + session = _session("rollback-longest-service", enable_last_assistant_rollback=True) + prompt = [{"role": "user", "content": "same prompt"}] + continuation = [ + *prompt, + {"role": "assistant", "content": "A1"}, + {"role": "user", "content": "second turn"}, + ] + rewrite_messages = [ + *continuation, + {"role": "user", "content": "user_error replaces old assistant"}, + ] + + await _run(session, SequencedBackend(["A1"]), prompt) + await _run(session, SequencedBackend(["A1"]), prompt) + await _run(session, SequencedBackend(["A2"]), continuation) + await _run(session, SequencedBackend(["FIXED"]), rewrite_messages) + + state = session.snapshot_state() + assert state["active_chain_ids"] == [1, 2] + assert state["rollback_count"] == 1 + chains_by_id = {chain.chain_id: chain for chain in session.active_chains} + assert _decode_response_ids(chains_by_id[1].buffer.response_ids) == "A1" + long_chain_text = _decode_response_ids(chains_by_id[2].buffer.response_ids) + assert long_chain_text.startswith("A1") + assert long_chain_text.endswith("FIXED") + assert "A2" not in long_chain_text + + +@pytest.mark.asyncio +async def test_last_assistant_rollback_tie_prefers_exact_chain_without_drop(): + """Continue the exact chain when its service value ties a rollback candidate.""" + session = _session("rollback-exact-tie", enable_last_assistant_rollback=True) + prompt = [{"role": "user", "content": "same prompt"}] + echoed_assistant = [*prompt, {"role": "assistant", "content": "A1"}] + rewrite_messages = [ + *echoed_assistant, + {"role": "user", "content": "replace the newer assistant"}, + ] + + await _run(session, SequencedBackend(["A1"]), prompt) + await _run(session, SequencedBackend(["A1"]), prompt) + await _run(session, SequencedBackend(["A2"]), echoed_assistant) + await _run(session, SequencedBackend(["FIXED"]), rewrite_messages) + + state = session.snapshot_state() + assert state["active_chain_ids"] == [1, 2] + assert state["rollback_count"] == 0 + chains_by_id = {chain.chain_id: chain for chain in session.active_chains} + assert _decode_response_ids(chains_by_id[1].buffer.response_ids).endswith("FIXED") + assert _decode_response_ids(chains_by_id[2].buffer.response_ids).endswith("A2") + + +def test_chain_prefix_hash_match_accepts_empty_history_for_any_request(): + """Document the current empty-history wildcard behavior without changing it.""" + session = _session("empty-history-prefix") + incoming_hashes = session._extend_message_prefix_hashes( + [], + [{"role": "user", "content": "any request"}], + ) + empty_chain = SimpleNamespace(message_history=[], message_tip_hash="unused") + + assert session._is_chain_prefix_hash_match( + chain=empty_chain, + incoming_message_prefix_hashes=incoming_hashes, + ) + + +@pytest.mark.asyncio +async def test_last_assistant_rollback_rejects_misaligned_stored_logprobs(): + """Fail loudly instead of slicing a chain whose token truth is already corrupt.""" + session = _session( + "rollback-logprob-assert", + sampling_params={"logprobs": True}, + enable_last_assistant_rollback=True, + ) + prompt = [{"role": "user", "content": "run"}] + rewrite_messages = [*prompt, {"role": "user", "content": "user_error"}] + + await _run(session, _LogprobBackend([("BAD", "full")]), prompt) + session.active_chains[0].buffer.response_logprobs.pop() + + with pytest.raises(AssertionError, match="response_logprobs must be empty or aligned"): + await _run(session, _LogprobBackend([("FIXED", "full")]), rewrite_messages) + + +@pytest.mark.asyncio +async def test_last_assistant_rollback_multimodal_suffix_reencodes_expanded_media_tokens(): + """Keep rollback media aligned when one image expands into multiple tokens.""" + session = _session( + "rollback-multimodal", + processor=_ExpandedImageTokenProcessor(), + vision_info_extractor=fake_vision_info_extractor, + enable_last_assistant_rollback=True, + ) + first_messages = [_image_message("image://old.png", "inspect old")] + replacement_message = _image_message("image://error.png", "user_error image") + rewrite_messages = [*first_messages, replacement_message] + backend = SequencedBackend(["BAD_IMAGE", "FIXED_IMAGE"]) + expected_suffix_ids = session._codec.encode_incremental( + [replacement_message], + image_data=["image://error.png"], + ) + + await _run(session, backend, first_messages) + await _run(session, backend, rewrite_messages) + + state = session.snapshot_state() + assert state["active_chain_ids"] == [1] + assert state["rollback_count"] == 1 + assert state["rollback_dropped_trainable_tokens_total"] == len("BAD_IMAGE") + assert backend.calls[1]["image_data"] == ["image://old.png", "image://error.png"] + [chain] = session.active_chains + assert chain.image_data == ["image://old.png", "image://error.png"] + assert chain.video_data is None + assert chain.buffer.response_ids[: len(expected_suffix_ids)] == expected_suffix_ids + assert chain.buffer.response_mask == [0] * len(expected_suffix_ids) + [1] * len("FIXED_IMAGE") + + +@pytest.mark.asyncio +async def test_last_assistant_rollback_disabled_splits_rewritten_assistant(): + """Allow explicit opt-out so rewritten assistants still split trajectories.""" + session = _session("rollback-disabled", enable_last_assistant_rollback=False) + prompt = [{"role": "user", "content": "run mini-swe"}] + rewrite = [*prompt, {"role": "user", "content": "user_error"}] + backend = SequencedBackend(["BAD", "FIXED"]) + + await _run(session, backend, prompt) + await _run(session, backend, rewrite) + + state = session.snapshot_state() + assert state["active_chain_ids"] == [1, 2] + assert state["rollback_count"] == 0 + chains_by_id = {chain.chain_id: chain for chain in session.active_chains} + assert _decode_response_ids(chains_by_id[1].buffer.response_ids) == "BAD" + assert _decode_response_ids(chains_by_id[2].buffer.response_ids) == "FIXED" + + @pytest.mark.asyncio async def test_multiple_chains_repeated_same_prompt_creates_siblings_and_continues_latest(): """Create siblings for repeated prompts and continue the most recently updated one.""" diff --git a/uni_agent/framework/entry.py b/uni_agent/framework/entry.py index 48e1ac29..523572e7 100644 --- a/uni_agent/framework/entry.py +++ b/uni_agent/framework/entry.py @@ -45,6 +45,7 @@ def build_gateway_manager(*, config, llm_client) -> GatewayManager: apply_chat_template_kwargs=dict(apply_chat_template_kwargs), prompt_length=config.actor_rollout_ref.rollout.prompt_length, response_length=config.actor_rollout_ref.rollout.response_length, + enable_last_assistant_rollback=af_cfg.get("enable_last_assistant_rollback", True), ) return GatewayManager( diff --git a/uni_agent/gateway/config.py b/uni_agent/gateway/config.py index bba2dac6..ca955262 100644 --- a/uni_agent/gateway/config.py +++ b/uni_agent/gateway/config.py @@ -29,6 +29,8 @@ class GatewayActorConfig: vision_info_extractor_kwargs: Static kwargs forwarded to the extractor. prompt_length: Optional prompt-token budget stored on gateway sessions. response_length: Optional response-token budget stored on gateway sessions. + enable_last_assistant_rollback: Whether latest-assistant rewrites may + rollback and reuse an existing chain. Enabled by default. """ tokenizer: Any @@ -41,7 +43,13 @@ class GatewayActorConfig: vision_info_extractor_kwargs: dict[str, Any] | None = None prompt_length: int | None = None response_length: int | None = None + enable_last_assistant_rollback: bool = True def __post_init__(self) -> None: + if type(self.enable_last_assistant_rollback) is not bool: + raise ValueError( + "enable_last_assistant_rollback must be a bool, " + f"got {type(self.enable_last_assistant_rollback).__name__}" + ) if self.response_length is not None and self.response_length <= 0: raise ValueError(f"response_length must be positive when set, got {self.response_length}") diff --git a/uni_agent/gateway/gateway.py b/uni_agent/gateway/gateway.py index 2954395a..43213338 100644 --- a/uni_agent/gateway/gateway.py +++ b/uni_agent/gateway/gateway.py @@ -80,6 +80,7 @@ def __init__(self, config: GatewayActorConfig, backend): ) self._prompt_length = config.prompt_length self._response_length = config.response_length + self._enable_last_assistant_rollback = config.enable_last_assistant_rollback self._sessions: dict[str, GatewaySession] = {} self._app = FastAPI() self._server_port: int | None = None @@ -253,6 +254,7 @@ async def create_session( prompt_length=self._prompt_length, response_length=self._response_length, sampling_params=sampling_params, + enable_last_assistant_rollback=self._enable_last_assistant_rollback, ) return handle diff --git a/uni_agent/gateway/session/session.py b/uni_agent/gateway/session/session.py index 94c61e8f..b7d6441c 100644 --- a/uni_agent/gateway/session/session.py +++ b/uni_agent/gateway/session/session.py @@ -56,6 +56,19 @@ class TrajectoryBuffer: routed_experts: Any | None = None +@dataclass +class LastAssistantStart: + """Stable chain lengths captured immediately before its latest assistant.""" + + response_ids_len: int + response_mask_len: int + response_logprobs_len: int + message_history_len: int + image_data_len: int + video_data_len: int + tip_hash: str + + @dataclass class ChainState: """One active linear trajectory chain in a gateway session.""" @@ -67,6 +80,7 @@ class ChainState: buffer: TrajectoryBuffer image_data: list[Any] | None video_data: list[Any] | None + last_assistant_start: LastAssistantStart updated_seq: int @@ -101,6 +115,12 @@ class EncodedData: a new chain. incoming_message_prefix_hashes: Stable prefix hashes for the normalized request history. + last_assistant_start: Lengths immediately before this generation's + assistant output is appended. + rollback_applied: Whether prepare removed the selected chain's latest + assistant before re-encoding the incoming suffix. + rollback_dropped_trainable_tokens: Number of mask=1 tokens removed by + that rollback. """ buffer: TrajectoryBuffer @@ -113,6 +133,9 @@ class EncodedData: length_exhausted_trajectory: Trajectory | None chain_id: int | None incoming_message_prefix_hashes: list[str] = field(default_factory=list) + last_assistant_start: LastAssistantStart | None = None + rollback_applied: bool = False + rollback_dropped_trainable_tokens: int = 0 @dataclass @@ -153,6 +176,7 @@ def __init__( prompt_length: int | None = None, response_length: int | None = None, sampling_params: dict[str, Any] | None = None, + enable_last_assistant_rollback: bool = True, ): """Create an active session bound to a handle and model codec.""" if response_length is not None and response_length <= 0: @@ -165,11 +189,14 @@ def __init__( self._prompt_length = prompt_length self._response_length = response_length self._sampling_params = dict(sampling_params or {}) + self._enable_last_assistant_rollback = enable_last_assistant_rollback self.active_chains: list[ChainState] = [] self.materialized_chains: list[MaterializedChain] = [] self.reserved_chain_ids: set[int] = set() self._next_chain_id = 1 self._order_seq = 0 + self._rollback_count = 0 + self._rollback_dropped_trainable_tokens_total = 0 self.reward_info: dict[str, Any] = {} self.phase = SessionPhase.ACTIVE self.created_at = time.time() @@ -245,6 +272,7 @@ async def run_generation(self, request: InternalGenerationRequest, backend) -> G f"got {len(log_probs)} logprobs for {len(response_ids)} tokens" ) encoded.buffer.response_logprobs.extend(log_probs) + self._assert_response_logprob_alignment(encoded.buffer) # R3 router replay: the backend returns routing for the full context # it just prefilled (prompt + response so far + new tokens), so keep @@ -335,6 +363,8 @@ def snapshot_state(self) -> dict[str, Any]: "num_active_chains": len(self.active_chains), "active_chain_ids": [chain.chain_id for chain in self.active_chains], "active_chain_tip_hashes": {chain.chain_id: chain.message_tip_hash for chain in self.active_chains}, + "rollback_count": self._rollback_count, + "rollback_dropped_trainable_tokens_total": self._rollback_dropped_trainable_tokens_total, } async def _prepare_generation_inputs( @@ -349,6 +379,8 @@ async def _prepare_generation_inputs( tools=tools, incoming_message_prefix_hashes=incoming_message_prefix_hashes, ) + rollback_applied = False + rollback_dropped_trainable_tokens = 0 if selected_chain is None: image_data, video_data = await self._codec.extract_multi_modal_data(messages) @@ -362,54 +394,134 @@ async def _prepare_generation_inputs( chain_id = None else: buffer = self._copy_trajectory_buffer(selected_chain.buffer) + self._assert_response_logprob_alignment(buffer) image_data, video_data = self._copy_chain_media(selected_chain) chain_id = selected_chain.chain_id - incremental_messages = messages[len(selected_chain.message_history) :] - new_image_data = None - new_video_data = None - incremental_ids = [] - already_exhausted = self._response_length is not None and len(buffer.response_mask) >= self._response_length - if incremental_messages and not already_exhausted: - new_image_data, new_video_data = await self._codec.extract_multi_modal_data(incremental_messages) - incremental_ids = self._codec.encode_incremental( - incremental_messages, - image_data=new_image_data, - video_data=new_video_data, - ) + rollback_to_last_assistant = not self._is_chain_prefix_hash_match( + chain=selected_chain, + incoming_message_prefix_hashes=incoming_message_prefix_hashes, + ) + if rollback_to_last_assistant: + last_assistant_start = selected_chain.last_assistant_start + assert last_assistant_start.response_ids_len <= len(buffer.response_ids) + assert last_assistant_start.response_mask_len <= len(buffer.response_mask) + assert last_assistant_start.response_logprobs_len <= len(buffer.response_logprobs) + rollback_dropped_trainable_tokens = sum(buffer.response_mask[last_assistant_start.response_mask_len :]) + del buffer.response_ids[last_assistant_start.response_ids_len :] + del buffer.response_mask[last_assistant_start.response_mask_len :] + del buffer.response_logprobs[last_assistant_start.response_logprobs_len :] + self._assert_response_logprob_alignment(buffer) + + stored_image_data = list(selected_chain.image_data or []) + stored_video_data = list(selected_chain.video_data or []) + assert last_assistant_start.image_data_len <= len(stored_image_data) + assert last_assistant_start.video_data_len <= len(stored_video_data) + image_data = stored_image_data[: last_assistant_start.image_data_len] or None + video_data = stored_video_data[: last_assistant_start.video_data_len] or None + suffix_messages = messages[last_assistant_start.message_history_len :] + suffix_ids: list[int] = [] + new_image_data = None + new_video_data = None + if suffix_messages: + new_image_data, new_video_data = await self._codec.extract_multi_modal_data(suffix_messages) + suffix_ids = self._codec.encode_incremental( + suffix_messages, + image_data=new_image_data, + video_data=new_video_data, + ) - if already_exhausted or ( - self._response_length is not None - and len(buffer.response_mask) + len(incremental_ids) >= self._response_length - ): - context_ids = buffer.prompt_ids + buffer.response_ids - return EncodedData( - buffer=buffer, - context_ids=context_ids, - sampling_params={}, - messages=list(messages), - tools=tools, - image_data=image_data, - video_data=video_data, - length_exhausted_trajectory=self._build_materialized_trajectory( - chain=selected_chain, - extra_fields={"materialization_reason": "max_response_length"}, - ), - chain_id=selected_chain.chain_id, - incoming_message_prefix_hashes=list(incoming_message_prefix_hashes), + buffer.response_ids.extend(suffix_ids) + buffer.response_mask.extend([0] * len(suffix_ids)) + if sampling_params.get("logprobs", False): + buffer.response_logprobs.extend([0.0] * len(suffix_ids)) + self._assert_response_logprob_alignment(buffer) + if new_image_data: + if image_data is None: + image_data = [] + image_data.extend(new_image_data) + if new_video_data: + if video_data is None: + video_data = [] + video_data.extend(new_video_data) + rollback_applied = True + + if self._response_length is not None and len(buffer.response_mask) >= self._response_length: + context_ids = buffer.prompt_ids + buffer.response_ids + working_chain = replace( + selected_chain, + message_history=list(messages), + message_tip_hash=incoming_message_prefix_hashes[-1], + buffer=buffer, + image_data=self._copy_media_list(image_data), + video_data=self._copy_media_list(video_data), + ) + return EncodedData( + buffer=buffer, + context_ids=context_ids, + sampling_params={}, + messages=list(messages), + tools=tools, + image_data=image_data, + video_data=video_data, + length_exhausted_trajectory=self._build_materialized_trajectory( + chain=working_chain, + extra_fields={"materialization_reason": "max_response_length"}, + ), + chain_id=selected_chain.chain_id, + incoming_message_prefix_hashes=list(incoming_message_prefix_hashes), + rollback_applied=True, + rollback_dropped_trainable_tokens=rollback_dropped_trainable_tokens, + ) + else: + incremental_messages = messages[len(selected_chain.message_history) :] + new_image_data = None + new_video_data = None + incremental_ids = [] + already_exhausted = ( + self._response_length is not None and len(buffer.response_mask) >= self._response_length ) + if incremental_messages and not already_exhausted: + new_image_data, new_video_data = await self._codec.extract_multi_modal_data(incremental_messages) + incremental_ids = self._codec.encode_incremental( + incremental_messages, + image_data=new_image_data, + video_data=new_video_data, + ) - buffer.response_ids.extend(incremental_ids) - buffer.response_mask.extend([0] * len(incremental_ids)) - if sampling_params.get("logprobs", False): - buffer.response_logprobs.extend([0.0] * len(incremental_ids)) - if new_image_data: - if image_data is None: - image_data = [] - image_data.extend(new_image_data) - if new_video_data: - if video_data is None: - video_data = [] - video_data.extend(new_video_data) + if already_exhausted or ( + self._response_length is not None + and len(buffer.response_mask) + len(incremental_ids) >= self._response_length + ): + context_ids = buffer.prompt_ids + buffer.response_ids + return EncodedData( + buffer=buffer, + context_ids=context_ids, + sampling_params={}, + messages=list(messages), + tools=tools, + image_data=image_data, + video_data=video_data, + length_exhausted_trajectory=self._build_materialized_trajectory( + chain=selected_chain, + extra_fields={"materialization_reason": "max_response_length"}, + ), + chain_id=selected_chain.chain_id, + incoming_message_prefix_hashes=list(incoming_message_prefix_hashes), + ) + + buffer.response_ids.extend(incremental_ids) + buffer.response_mask.extend([0] * len(incremental_ids)) + if sampling_params.get("logprobs", False): + buffer.response_logprobs.extend([0.0] * len(incremental_ids)) + self._assert_response_logprob_alignment(buffer) + if new_image_data: + if image_data is None: + image_data = [] + image_data.extend(new_image_data) + if new_video_data: + if video_data is None: + video_data = [] + video_data.extend(new_video_data) context_ids = buffer.prompt_ids + buffer.response_ids remaining_response_budget = ( @@ -420,6 +532,13 @@ async def _prepare_generation_inputs( sampling_params.get("max_tokens", remaining_response_budget), remaining_response_budget, ) + last_assistant_start = self._snapshot_last_assistant_start( + buffer=buffer, + message_history_len=len(messages), + image_data=image_data, + video_data=video_data, + tip_hash=incoming_message_prefix_hashes[-1], + ) return EncodedData( buffer=buffer, context_ids=context_ids, @@ -431,6 +550,9 @@ async def _prepare_generation_inputs( length_exhausted_trajectory=None, chain_id=chain_id, incoming_message_prefix_hashes=list(incoming_message_prefix_hashes), + last_assistant_start=last_assistant_start, + rollback_applied=rollback_applied, + rollback_dropped_trainable_tokens=rollback_dropped_trainable_tokens, ) def _select_chain( @@ -439,19 +561,52 @@ def _select_chain( tools: list[dict[str, Any]] | None, incoming_message_prefix_hashes: list[str], ) -> ChainState | None: - candidates = [ - chain - for chain in self.active_chains - if chain.chain_id not in self.reserved_chain_ids - and chain.active_tool_schemas == tools - and self._is_chain_prefix_hash_match( + ranked_candidates = [] + deepest_rollback_candidates = [] + deepest_rollback_service_value = -1 + for chain in self.active_chains: + if chain.chain_id in self.reserved_chain_ids or chain.active_tool_schemas != tools: + continue + assistant_start = chain.last_assistant_start + assistant_start_len = assistant_start.message_history_len + # A request ending exactly at the boundary is a fresh sample from + # the same prompt, not a rewrite of the abandoned assistant. + if assistant_start_len >= len(incoming_message_prefix_hashes): + continue + if incoming_message_prefix_hashes[assistant_start_len - 1] != assistant_start.tip_hash: + continue + if self._is_chain_prefix_hash_match( chain=chain, incoming_message_prefix_hashes=incoming_message_prefix_hashes, - ) - ] - if not candidates: + ): + ranked_candidates.append((chain, len(chain.message_history), True)) + continue + if self._enable_last_assistant_rollback: + if assistant_start_len > deepest_rollback_service_value: + deepest_rollback_candidates = [chain] + deepest_rollback_service_value = assistant_start_len + elif assistant_start_len == deepest_rollback_service_value: + deepest_rollback_candidates.append(chain) + + if len(deepest_rollback_candidates) == 1: + rollback_chain = deepest_rollback_candidates[0] + ranked_candidates.append((rollback_chain, deepest_rollback_service_value, False)) + elif deepest_rollback_candidates and deepest_rollback_service_value > max( + (candidate[1] for candidate in ranked_candidates), + default=-1, + ): + return None + if not ranked_candidates: return None - return max(candidates, key=lambda chain: (len(chain.message_history), chain.updated_seq, chain.chain_id)) + return max( + ranked_candidates, + key=lambda candidate: ( + candidate[1], + candidate[2], + candidate[0].updated_seq, + candidate[0].chain_id, + ), + )[0] def _is_chain_prefix_hash_match( self, @@ -511,6 +666,37 @@ def _copy_media_list(self, media: list[Any] | None) -> list[Any] | None: # Copy only the container; media payloads may not be deepcopyable. return list(media) if media is not None else None + def _assert_response_logprob_alignment(self, buffer: TrajectoryBuffer) -> None: + assert len(buffer.response_logprobs) in { + 0, + len(buffer.response_ids), + }, "response_logprobs must be empty or aligned with response_ids" + + def _snapshot_last_assistant_start( + self, + *, + buffer: TrajectoryBuffer, + message_history_len: int, + image_data: list[Any] | None, + video_data: list[Any] | None, + tip_hash: str, + ) -> LastAssistantStart: + return LastAssistantStart( + response_ids_len=len(buffer.response_ids), + response_mask_len=len(buffer.response_mask), + response_logprobs_len=len(buffer.response_logprobs), + message_history_len=message_history_len, + image_data_len=len(image_data or []), + video_data_len=len(video_data or []), + tip_hash=tip_hash, + ) + + def _record_rollback_stats(self, encoded: EncodedData) -> None: + if not encoded.rollback_applied: + return + self._rollback_count += 1 + self._rollback_dropped_trainable_tokens_total += encoded.rollback_dropped_trainable_tokens + def _commit_generation_to_chain(self, encoded: EncodedData, assistant_msg: dict[str, Any]) -> None: message_history = list(encoded.messages) + [assistant_msg] message_prefix_hashes = self._extend_message_prefix_hashes( @@ -518,6 +704,9 @@ def _commit_generation_to_chain(self, encoded: EncodedData, assistant_msg: dict[ [assistant_msg], ) assert len(message_prefix_hashes) == len(message_history) + if encoded.last_assistant_start is None: + raise RuntimeError("last assistant start is missing") + self._record_rollback_stats(encoded) if encoded.chain_id is None: order_seq = self._next_order_seq() chain_id = self._allocate_chain_id() @@ -530,6 +719,7 @@ def _commit_generation_to_chain(self, encoded: EncodedData, assistant_msg: dict[ buffer=encoded.buffer, image_data=self._copy_media_list(encoded.image_data), video_data=self._copy_media_list(encoded.video_data), + last_assistant_start=encoded.last_assistant_start, updated_seq=order_seq, ) ) @@ -545,6 +735,7 @@ def _commit_generation_to_chain(self, encoded: EncodedData, assistant_msg: dict[ buffer=encoded.buffer, image_data=self._copy_media_list(encoded.image_data), video_data=self._copy_media_list(encoded.video_data), + last_assistant_start=encoded.last_assistant_start, updated_seq=order_seq, ) @@ -559,6 +750,7 @@ def _close_length_exhausted_chain(self, encoded: EncodedData) -> None: order_seq=order_seq, ) ) + self._record_rollback_stats(encoded) del self.active_chains[chain_index] def _find_active_chain(self, chain_id: int) -> tuple[int, ChainState]: @@ -596,8 +788,9 @@ def _build_materialized_trajectory( chain: ChainState, extra_fields: dict[str, Any] | None = None, ) -> Trajectory: + self._assert_response_logprob_alignment(chain.buffer) response_logprobs = None - if chain.buffer.response_logprobs and len(chain.buffer.response_logprobs) == len(chain.buffer.response_ids): + if chain.buffer.response_logprobs: response_logprobs = list(chain.buffer.response_logprobs) return Trajectory( prompt_ids=list(chain.buffer.prompt_ids), From bcaabdbf9a5dc157c001b0db237101529a833219 Mon Sep 17 00:00:00 2001 From: zackcxb Date: Thu, 23 Jul 2026 02:53:42 +0000 Subject: [PATCH 2/2] chore(examples): format sandbox client --- examples/blackbox_recipes/sandbox_client.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/examples/blackbox_recipes/sandbox_client.py b/examples/blackbox_recipes/sandbox_client.py index c2f3e0f8..7fa05f2a 100644 --- a/examples/blackbox_recipes/sandbox_client.py +++ b/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 \ No newline at end of file + self._sandbox = None