From 8646a8d9e6043d3d40ebbc6e51d54b83f5b9592f Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Mon, 1 Jun 2026 16:10:48 +0000 Subject: [PATCH 01/28] refactor(agent): replace segment-based trajectory with turn-node TrajectoryTree Replace slime/agent/trajectory.py (manual subagent/wipe/final segment bookkeeping) with slime/agent/trajectory_manager.py, which folds each turn into a per-session turn-node tree routed by text prefix. Sub-agent and compaction patterns now split into independent leaves automatically. Update Anthropic/OpenAI adapters and common helpers to the new record_turn / export_token_segments API, and point the coding_agent_rl example at slime.agent.trajectory_manager. --- examples/coding_agent_rl/README.md | 22 +- examples/coding_agent_rl/generate.py | 6 +- slime/agent/adapters/anthropic.py | 207 +++------- slime/agent/adapters/common.py | 211 +++++++++- slime/agent/adapters/openai.py | 119 +++--- slime/agent/trajectory.py | 208 ---------- slime/agent/trajectory_manager.py | 552 +++++++++++++++++++++++++++ 7 files changed, 846 insertions(+), 479 deletions(-) delete mode 100644 slime/agent/trajectory.py create mode 100644 slime/agent/trajectory_manager.py diff --git a/examples/coding_agent_rl/README.md b/examples/coding_agent_rl/README.md index 38e106a20e..4ff8d5bdff 100644 --- a/examples/coding_agent_rl/README.md +++ b/examples/coding_agent_rl/README.md @@ -145,16 +145,18 @@ The Anthropic adapter therefore follows a **string in, token out** contract: Multi-turn agents still force the adapter to tokenize later message histories, because tool observations and claude-code's own compacted messages -arrive as strings. `slime.agent.trajectory.merge_turns` stitches those later -prompts against the saved token stream: - -- New prompt suffixes that are tool/user/environment context are appended with - `loss_mask=0`. -- Fresh model outputs from SGLang are appended with `loss_mask=1`. -- If a later prompt no longer token-matches an earlier sampled output, the - unmatched suffix is dropped. If the drift cuts through the middle of a - previous model output, the retained prefix of that whole output turn is also - assigned `loss_mask=0`. +arrive as strings. `slime.agent.trajectory_manager` models each session as a +turn-node text-prefix tree and folds every request into it: + +- Each request is split into strictly-alternating `(prompt, response)` turns; + the current turn's prompt segment is taken by a segment render diff and its + response comes from SGLang, while historical responses are served from a + per-session truth cache so they keep their exact sampled token ids. +- New prompt suffixes that are tool/user/environment context export with + `loss_mask=0`; fresh model outputs export with `loss_mask=1`. +- Text divergence (compaction / history rewrite / sub-agent) branches the tree; + tail-turn token drift (TITO) is repaired in place and masked. `export` + deduplicates shared fork prefixes so they are trained at most once. That last case is the important correctness guard. A re-tokenization mismatch can make a string-level conversation look continuous while token-level diff --git a/examples/coding_agent_rl/generate.py b/examples/coding_agent_rl/generate.py index 64fb0d6660..2ce0aae815 100644 --- a/examples/coding_agent_rl/generate.py +++ b/examples/coding_agent_rl/generate.py @@ -10,10 +10,10 @@ 2. ``sandbox.git_diff`` captures the model-produced patch. 3. ``sandbox.evaluate`` scores that patch in a second clean sandbox. 4. ``_merge_samples`` combines reward + adapter ``TokenSegment``s, - delegating segment-to-``Sample`` fan-out to ``slime.agent.trajectory``. + delegating segment-to-``Sample`` fan-out to ``slime.agent.trajectory_manager``. All sandbox-side details live in ``sandbox.py``; the LLM plumbing -(Anthropic <-> SGLang /generate, token capture, 3-kind segment split) uses +(Anthropic <-> SGLang /generate, token capture, turn-node trajectory tree) uses ``slime.agent.adapters.AnthropicAdapter``. Dataset row ``metadata`` schema:: @@ -53,7 +53,7 @@ from typing import Any from slime.agent.adapters import AnthropicAdapter -from slime.agent.trajectory import TokenSegment, fan_out_sample_segments +from slime.agent.trajectory_manager import TokenSegment, fan_out_sample_segments from slime.utils.misc import SingletonMeta from slime.utils.processing_utils import load_tokenizer from slime.utils.types import Sample diff --git a/slime/agent/adapters/anthropic.py b/slime/agent/adapters/anthropic.py index e2d0a8300a..3a4f52a098 100644 --- a/slime/agent/adapters/anthropic.py +++ b/slime/agent/adapters/anthropic.py @@ -2,13 +2,13 @@ The adapter exposes ``/v1/messages`` and ``/v1/messages/count_tokens``. It renders each Anthropic message history with the served model's chat template, -calls SGLang ``/generate`` with ``input_ids``, and records the exact sampled -token ids/logprobs as ``TurnRecord`` objects. New code should use -``AnthropicAdapter`` and call ``finish_session()`` at trajectory end to drain -trainable ``TokenSegment`` objects. +calls SGLang ``/generate`` with ``input_ids``, and folds the turn into a +per-session turn-node :class:`~slime.agent.trajectory_manager.TrajectoryTree`. -It also handles Claude Code sub-agent and compaction patterns by splitting one -session into ``subagent``, ``wipe``, and ``final`` segments. +The tree routes everything by text prefix, so Claude Code sub-agent and +compaction patterns split into independent leaves automatically -- no manual +``active_sub`` / ``wipe`` bookkeeping. Call ``finish_session()`` at trajectory +end to drain trainable ``TokenSegment`` objects. """ from __future__ import annotations @@ -22,35 +22,31 @@ from aiohttp import web -from slime.agent.adapters.common import ADAPTER_KEY, REASONING_PARSER_KEY, TOKENIZER_KEY, TOOL_PARSER_KEY -from slime.agent.adapters.common import AdapterChain as Chain from slime.agent.adapters.common import ( + ADAPTER_KEY, + REASONING_PARSER_KEY, + TOKENIZER_KEY, + TOOL_PARSER_KEY, BaseAdapter, + SessionTrajectory, + assemble_turns, call_sglang_generate, ok_response, - render_token_ids, + render_prompt, request_session_id, ) -from slime.agent.adapters.common import stable_hash as _hash from slime.agent.parsing import parse_model_output -from slime.agent.trajectory import TokenSegment, TurnRecord, TurnSegment, make_turn_segment, merge_turn_segments +from slime.agent.trajectory_manager import TokenSegment, export_token_segments, record_turn logger = logging.getLogger(__name__) -# Tool names claude-code uses to dispatch a sub-agent. -_SUBAGENT_TOOLS = {"Task", "Agent"} - - @dataclasses.dataclass class Session: - main: Chain = dataclasses.field(default_factory=Chain) - active_sub: Chain | None = None # at most one sub-agent at a time - pending_dispatch_id: str = "" # tool_use_id we're waiting to close + traj: SessionTrajectory = dataclasses.field(default_factory=SessionTrajectory) sampling_defaults: dict = dataclasses.field(default_factory=dict) max_context_tokens: int = 0 lock: asyncio.Lock = dataclasses.field(default_factory=asyncio.Lock) - segments: list[TurnSegment] = dataclasses.field(default_factory=list) # frozen output class AnthropicAdapter(BaseAdapter): @@ -75,86 +71,14 @@ async def finish_session(self, sid: str, *, wait_timeout: float = 5.0) -> list[T s = self.store.pop(sid, None) if s is None: return [] - if s.active_sub is not None and s.active_sub.turns: - s.segments.append(make_turn_segment(s.active_sub.turns, kind="subagent")) - if s.main.turns: - s.segments.append(make_turn_segment(s.main.turns, kind="final")) - - return merge_turn_segments(s.segments) + return export_token_segments(s.traj.tree) # ============================================================================= -# 2. Per-turn stages +# Translation (Anthropic wire <-> chat-template messages) -- unchanged # ============================================================================= -def _select_chain(s: Session, body: dict) -> tuple[Chain, bool, str]: - """Decide which chain this turn operates on. - - 1. fingerprint body.messages and body.system into hashes - 2. if main now contains the tool_result for a pending sub dispatch, - snapshot the sub chain into s.segments and clear s.active_sub - 3. pick main vs s.active_sub based on whether request continues main's prefix - 4. classify as 'new' | 'append' | 'wipe' against the chosen target; - a wipe also snapshots the target's current state into s.segments - - Returns (target_chain, is_sub, kind). - """ - all_msgs = body.get("messages") or [] - msg_hashes = [_hash(m) for m in all_msgs] - req_system_hash = _hash(body.get("system")) if "system" in body else s.main.system_hash - - # Close active sub-agent if its dispatch tool_result has landed on main. - if s.pending_dispatch_id and s.active_sub is not None: - tu_id = s.pending_dispatch_id - for m in all_msgs: - if not isinstance(m, dict) or m.get("role") != "user": - continue - content = m.get("content") - if not isinstance(content, list): - continue - done = any( - isinstance(b, dict) and b.get("type") == "tool_result" and b.get("tool_use_id") == tu_id - for b in content - ) - if done: - if s.active_sub.turns: - s.segments.append(make_turn_segment(s.active_sub.turns, kind="subagent")) - s.active_sub = None - s.pending_dispatch_id = "" - break - - # Route: main iff request continues main's prefix. Sub system_hash can be - # "" (armed before sub dialled in), so never route by sub equality alone. - if s.active_sub is None: - target, is_sub = s.main, False - else: - main_continues = ( - req_system_hash == s.main.system_hash - and len(msg_hashes) >= s.main.seen_msgs - and msg_hashes[: s.main.seen_msgs] == s.main.msg_hashes[: s.main.seen_msgs] - ) - target, is_sub = (s.main, False) if main_continues else (s.active_sub, True) - - # Classify; snapshot a "wipe" segment first if we're discarding work. - if target.seen_msgs == 0: - kind = "new" - else: - is_append = ( - req_system_hash == target.system_hash - and len(msg_hashes) >= target.seen_msgs - and msg_hashes[: target.seen_msgs] == target.msg_hashes[: target.seen_msgs] - ) - if is_append: - kind = "append" - else: - if target.turns: - s.segments.append(make_turn_segment(target.turns, kind="wipe")) - kind = "wipe" - - return target, is_sub, kind - - def _flatten(c: Any) -> str: """Recursive Anthropic content flattener: text/tool_result(content) joined by newline, images replaced with a placeholder.""" @@ -241,47 +165,13 @@ def _anthropic_tools_to_chat_tools(anth_tools: list[dict] | None) -> list[dict] return ts or None -def _replace_chat_messages(target: Chain, body: dict) -> None: - """new/wipe: full reset of chat state and turn log.""" - all_msgs = body.get("messages") or [] - target.chat_messages = _translate_anthropic(all_msgs, body.get("system")) - if "system" in body: - target.system_hash = _hash(body.get("system")) - target.turns.clear() - target.seen_msgs = len(all_msgs) - target.msg_hashes = [_hash(m) for m in all_msgs] - if target.tools_schema is None: - target.tools_schema = _anthropic_tools_to_chat_tools(body.get("tools")) - - -def _extend_chat_messages(target: Chain, body: dict) -> None: - """append: translate only the new tail.""" - all_msgs = body.get("messages") or [] - translated = _translate_anthropic(all_msgs[target.seen_msgs :], None) - target.chat_messages.extend(translated) - - target.seen_msgs = len(all_msgs) - target.msg_hashes = [_hash(m) for m in all_msgs] - if target.tools_schema is None: - target.tools_schema = _anthropic_tools_to_chat_tools(body.get("tools")) - - -def _build_prompt(target: Chain, body: dict, kind: str, tok) -> list[int]: - """Replace/extend chat_messages and render input ids for sglang.""" - (_extend_chat_messages if kind == "append" else _replace_chat_messages)(target, body) - return render_token_ids(target, tok) - +# ============================================================================= +# Reply building (raw output text -> Anthropic content blocks) -- unchanged shape +# ============================================================================= -async def _generate( - prompt_ids: list[int], s: Session, body: dict, app, *, session_id: str | None = None -) -> TurnRecord: - """Call sglang and return a TurnRecord. - 1. build sampling_params (session defaults overlaid with body overrides) - 2. POST sglang /generate; on cancel/error fire /abort_request - 3. keep the exact prompt/output token ids; trajectory merge later compares - later prompt tokens with earlier outputs to build the loss mask - """ +async def _generate(prompt_ids: list[int], s: Session, body: dict, app, *, session_id: str | None = None): + """Call sglang and return a GenResult (output ids/logprobs/text).""" return await call_sglang_generate( prompt_ids, s, @@ -295,22 +185,19 @@ async def _generate( ) -def _build_reply(target: Chain, output_ids: list[int], finish: str, app) -> tuple[list[dict], str, str]: - """Turn the model's raw output ids into the reply we send back to claude-code. +def _build_reply(output_text: str, finish: str, tools_schema: list[dict] | None, app) -> tuple[list[dict], str, str]: + """Turn the model's raw output text into the reply we send back to claude-code. 1. parse decoded text -> (thinking, visible, tool_uses) via sglang parsers - 2. pack into Anthropic content blocks; tag dispatch_id when a tool_use - names Task/Agent (sub-agent trigger) + 2. pack into Anthropic content blocks 3. derive stop_reason: 'tool_use' | 'max_tokens' | 'end_turn' - Returns (blocks, stop_reason, dispatch_id). + Returns (blocks, stop_reason, dispatch_id). dispatch_id is retained for wire + compatibility but no longer drives routing (the tree splits sub-agents). """ - tok = app[TOKENIZER_KEY] - - raw_output = tok.decode(output_ids, skip_special_tokens=False) if output_ids else "" parsed = parse_model_output( - raw_output, - tools_schema=target.tools_schema, + output_text or "", + tools_schema=tools_schema, tool_parser_name=app[TOOL_PARSER_KEY], reasoning_parser_name=app[REASONING_PARSER_KEY], ) @@ -329,8 +216,6 @@ def _anthropic_blocks(thinking: str, visible: str, tool_uses: list[dict]) -> tup for tu in tool_uses: tu_id = f"toolu_{secrets.token_hex(8)}" blocks.append({"type": "tool_use", "id": tu_id, "name": tu["name"], "input": tu["input"]}) - if tu["name"] in _SUBAGENT_TOOLS: - dispatch_id = tu_id if not blocks: blocks.append({"type": "text", "text": ""}) return blocks, dispatch_id @@ -344,17 +229,8 @@ def _stop_reason(tool_uses: list[dict], finish: str) -> str: return "end_turn" -def _start_sub_chain(s: Session, dispatch_id: str) -> None: - """Start a fresh sub chain on this session and remember the tool_use_id - we'll watch for on main to know when this sub is done. The matching - 'sub done' step lives inside _select_chain.""" - s.pending_dispatch_id = dispatch_id - if s.active_sub is None: - s.active_sub = Chain() - - # ============================================================================= -# 3. Request handling -- one full turn + SSE wrap +# Request handling -- one full turn + SSE wrap # ============================================================================= @@ -369,19 +245,26 @@ async def _handle_request(request: web.Request) -> web.StreamResponse: if sid in adapter.closed: # session drained; refuse stragglers return web.Response(status=503, text="session closed") app = request.app + tok = app[TOKENIZER_KEY] s = adapter.store.setdefault(sid, Session()) task = asyncio.current_task() adapter.inflight.setdefault(sid, set()).add(task) try: async with s.lock: # same sid -> serialized - target, is_sub, kind = _select_chain(s, body) - ideal_ids = _build_prompt(target, body, kind, app[TOKENIZER_KEY]) - turn = await _generate(ideal_ids, s, body, app, session_id=sid) - blocks, stop, did = _build_reply(target, turn.output_ids, turn.finish_reason, app) - target.turns.append(turn) - if did and not is_sub: # sub doesn't nest - _start_sub_chain(s, did) - in_tok, out_tok = len(ideal_ids), len(turn.output_ids) + translated = _translate_anthropic(body.get("messages") or [], body.get("system")) + tools_schema = _anthropic_tools_to_chat_tools(body.get("tools")) + full_prompt_ids = render_prompt(s.traj, translated, tok, tools_schema) + gen = await _generate(full_prompt_ids, s, body, app, session_id=sid) + turns, pending_key = assemble_turns(s.traj, translated, tok, tools_schema, gen, full_prompt_ids) + node = record_turn(s.traj.tree, turns) + if node is not None: + s.traj.resp_truth[pending_key] = ( + list(gen.output_ids), + list(gen.output_log_probs), + gen.output_text, + ) + blocks, stop, _did = _build_reply(gen.output_text, gen.finish_reason, tools_schema, app) + in_tok, out_tok = len(full_prompt_ids), len(gen.output_ids) if body.get("stream") is True or "text/event-stream" in request.headers.get("Accept", ""): return await _stream_response(request, blocks, stop, in_tok, out_tok) return web.json_response(_message_response(body, blocks, stop, in_tok, out_tok)) diff --git a/slime/agent/adapters/common.py b/slime/agent/adapters/common.py index ed5d2e06a4..72c74191cd 100644 --- a/slime/agent/adapters/common.py +++ b/slime/agent/adapters/common.py @@ -1,4 +1,21 @@ -"""Shared adapter primitives for token-capturing agent rollouts.""" +"""Shared adapter primitives for token-capturing agent rollouts. + +Each HTTP request carries the full conversation history. We render it with the +served model's chat template, call SGLang ``/generate`` with ``input_ids``, and +fold the turn into a per-session :class:`~slime.agent.trajectory_manager.TrajectoryTree`. + +The tree does all routing (sub-agent / compaction / history-rewrite branch +automatically by text prefix), so there is no manual new/append/wipe logic. Two +things make this faithful under TITO (text-in-token-out) drift: + +* the **current turn's** incremental prompt segment is taken by *segment render + diff* (``render(through this prompt) - render(through prev response)``), which + is immune to drift because it compares two re-renders rather than a cache vs a + re-render (design §5.1); and +* **historical turns'** response tokens come from a per-session *truth cache* + keyed by the cumulative prompt that produced them, so they equal the tree + node tokens exactly (no re-tokenization). +""" from __future__ import annotations @@ -14,7 +31,7 @@ import aiohttp from aiohttp import web -from slime.agent.trajectory import TokenSegment, TurnRecord +from slime.agent.trajectory_manager import PromptSeg, RespSeg, TokenSegment, TrajectoryTree ADAPTER_KEY = web.AppKey("adapter", object) @@ -25,15 +42,39 @@ @dataclasses.dataclass -class AdapterChain: - """Protocol-neutral chat chain state used by HTTP adapters.""" +class GenResult: + """One assistant generation from the rollout engine (sglang ``/generate``). - system_hash: str = "" - chat_messages: list[dict] = dataclasses.field(default_factory=list) - tools_schema: list[dict] | None = None - seen_msgs: int = 0 - msg_hashes: list[str] = dataclasses.field(default_factory=list) - turns: list[TurnRecord] = dataclasses.field(default_factory=list) + ``output_text`` is ``decode(output_ids)`` cached once so reply builders and + the trajectory's response-segment text share a single detokenization. + """ + + output_ids: list[int] + output_log_probs: list[float] + finish_reason: str + output_text: str = "" + + +@dataclasses.dataclass +class SessionTrajectory: + """Per-session trajectory state: the turn-node tree plus the truth cache. + + ``resp_truth`` maps the cumulative prompt token tuple that *preceded* a + response to ``(output_ids, output_log_probs, output_text)``. The same turn, + seen as history in a later request, renders to the same cumulative prompt and + so retrieves its exact sampled tokens (a Case 2 append, never a false drift). + Keying on rendered prompt ids (a deterministic function of the messages), + rather than on echoed wire text, avoids the parse->serialize round-trip drift + that would otherwise make a text key miss. + + ``render_memo`` caches ``apply_chat_template`` results within a session so a + growing trajectory renders each distinct message-prefix once (O(n) renders + over the whole trajectory instead of O(n^2)). + """ + + tree: TrajectoryTree = dataclasses.field(default_factory=TrajectoryTree) + resp_truth: dict[tuple[int, ...], tuple[list[int], list[float], str]] = dataclasses.field(default_factory=dict) + render_memo: dict[tuple[str, bool], list[int]] = dataclasses.field(default_factory=dict) class BaseAdapter: @@ -95,15 +136,141 @@ def json_arguments(value: Any) -> str: return json.dumps(value, ensure_ascii=False) -def render_token_ids(chain: AdapterChain, tokenizer) -> list[int]: +def _extract_ids(enc: Any) -> list[int]: + ids = enc["input_ids"] if hasattr(enc, "__getitem__") and "input_ids" in enc else enc + return list(ids) + + +def render_token_ids( + messages: list[dict], tokenizer, *, tools: list[dict] | None = None, add_generation_prompt: bool = True +) -> list[int]: + """Render a chat-template message list to token ids.""" enc = tokenizer.apply_chat_template( - chain.chat_messages, - tools=chain.tools_schema, + messages, + tools=tools, tokenize=True, - add_generation_prompt=True, + add_generation_prompt=add_generation_prompt, ) - ids = enc["input_ids"] if hasattr(enc, "__getitem__") and "input_ids" in enc else enc - return list(ids) + return _extract_ids(enc) + + +# ============================================================================= +# Turn assembly: chat messages -> strictly-alternating (PromptSeg, RespSeg) turns +# ============================================================================= + + +def split_turns(chat_messages: list[dict]) -> list[tuple[list[dict], dict | None]]: + """Split a chat-template message list into ``(prompt_msgs, assistant_msg)`` + turns at every ``assistant`` boundary. + + ``prompt_msgs`` is the run of non-assistant messages since the previous + assistant. A trailing ``(prompt_msgs, None)`` turn is always appended: it is + the current turn whose response is about to be generated. + """ + specs: list[tuple[list[dict], dict | None]] = [] + buf: list[dict] = [] + for m in chat_messages: + if isinstance(m, dict) and m.get("role") == "assistant": + specs.append((buf, m)) + buf = [] + else: + buf.append(m) + specs.append((buf, None)) + return specs + + +def _render_memo( + traj: SessionTrajectory | None, + messages: list[dict], + tokenizer, + tools: list[dict] | None, + add_generation_prompt: bool, +) -> list[int]: + """Render with a per-session memo so each distinct message prefix is rendered + once across a growing trajectory (turns 1..n cost O(n) renders total, not + O(n^2)).""" + if traj is None: + return render_token_ids(messages, tokenizer, tools=tools, add_generation_prompt=add_generation_prompt) + key = (stable_hash([messages, tools]), add_generation_prompt) + cached = traj.render_memo.get(key) + if cached is None: + cached = render_token_ids(messages, tokenizer, tools=tools, add_generation_prompt=add_generation_prompt) + traj.render_memo[key] = cached + return list(cached) + + +def render_prompt(traj: SessionTrajectory, messages: list[dict], tokenizer, tools: list[dict] | None) -> list[int]: + """Render the full prompt (``add_generation_prompt=True``) to send to sglang, + memoized on the session so a replayed prefix is not re-rendered next turn.""" + return _render_memo(traj, messages, tokenizer, tools, True) + + +def assemble_turns( + traj: SessionTrajectory, + chat_messages: list[dict], + tokenizer, + tools: list[dict] | None, + gen: GenResult, + full_prompt_ids: list[int], +) -> tuple[list[tuple[PromptSeg, RespSeg]], tuple[int, ...]]: + """Build ``record_turn``'s strictly-alternating ``turns`` for this request. + + Every turn's incremental prompt segment comes from a *segment render diff* + (design §5.1): ``render(through this prompt) - render(through previous + response)``, comparing two re-renders so it is immune to TITO drift. The + previous render uses ``add_generation_prompt=False`` (it ends at a response), + so the diff is a clean length-based suffix. + + History response tokens come from the truth cache keyed by the cumulative + prompt that produced them (exact sampled ids; a Case 2 append, never a false + drift). On a cache miss they fall back to re-tokenization, which the tree + then handles as Case 3/4 (replace + mask). The current turn's response comes + from ``gen`` and reuses ``full_prompt_ids`` (the prompt already sent to + sglang) as its prompt segment. + + Does not write the cache -- the caller writes it only after ``record_turn`` + succeeds (no dangling cache on a sglang error). Returns ``(turns, + pending_key)`` where ``pending_key`` is the cumulative-prompt key under which + the current turn's response should be cached. + """ + specs = split_turns(chat_messages) + turns: list[tuple[PromptSeg, RespSeg]] = [] + prev_cum: list[dict] = [] + prev_ids: list[int] = [] + pending_key: tuple[int, ...] = tuple(full_prompt_ids) + + for prompt_msgs, assistant_msg in specs: + cum_prompt = prev_cum + prompt_msgs + if assistant_msg is None: + # Current turn: reuse the exact prompt already sent to sglang. + cur_ids = list(full_prompt_ids) + else: + cur_ids = _render_memo(traj, cum_prompt, tokenizer, tools, True) + delta = cur_ids[len(prev_ids) :] + p_text = tokenizer.decode(delta, skip_special_tokens=False) if delta else "" + p_seg = PromptSeg(p_text, list(delta)) + + if assistant_msg is None: + r_seg = RespSeg(gen.output_text, list(gen.output_ids), list(gen.output_log_probs)) + pending_key = tuple(cur_ids) + turns.append((p_seg, r_seg)) + break + + cached = traj.resp_truth.get(tuple(cur_ids)) + if cached is not None: + r_ids, r_lp, r_text = cached + r_seg = RespSeg(r_text, list(r_ids), list(r_lp)) + else: + full = _render_memo(traj, cum_prompt + [assistant_msg], tokenizer, tools, False) + r_ids = full[len(cur_ids) :] + r_text = tokenizer.decode(r_ids, skip_special_tokens=False) if r_ids else "" + r_seg = RespSeg(r_text, list(r_ids), [0.0] * len(r_ids)) + turns.append((p_seg, r_seg)) + + prev_cum = cum_prompt + [assistant_msg] + prev_ids = _render_memo(traj, prev_cum, tokenizer, tools, False) + + return turns, pending_key def request_session_id( @@ -185,7 +352,7 @@ async def call_sglang_generate( log_prefix: str, logger: logging.Logger, session_id: str | None = None, -) -> TurnRecord: +) -> GenResult: sp = _sampling_params(session, body, max_token_keys=max_token_keys, stop_keys=stop_keys) if session.max_context_tokens > 0: @@ -197,7 +364,7 @@ async def call_sglang_generate( len(prompt_ids), session.max_context_tokens, ) - return TurnRecord(prompt_ids=list(prompt_ids), output_ids=[], finish_reason="length") + return GenResult(output_ids=[], output_log_probs=[], finish_reason="length", output_text="") sp["max_new_tokens"] = min(int(sp.get("max_new_tokens", remaining_context)), remaining_context) sglang_url = app[SGLANG_URL_KEY] @@ -232,11 +399,13 @@ async def call_sglang_generate( pass raise - return TurnRecord( - prompt_ids=list(prompt_ids), + tok = app[TOKENIZER_KEY] if TOKENIZER_KEY in app else None + output_text = tok.decode(output_ids, skip_special_tokens=False) if (tok is not None and output_ids) else "" + return GenResult( output_ids=output_ids, - finish_reason=finish, output_log_probs=output_log_probs, + finish_reason=finish, + output_text=output_text, ) diff --git a/slime/agent/adapters/openai.py b/slime/agent/adapters/openai.py index f9d91d8ff5..cb8ed605fb 100644 --- a/slime/agent/adapters/openai.py +++ b/slime/agent/adapters/openai.py @@ -2,9 +2,10 @@ The adapter exposes ``/v1/chat/completions`` and ``/v1/responses``. Both endpoints render incoming messages with the served model's chat template, call -SGLang ``/generate`` with ``input_ids``, and record the exact sampled token -ids/logprobs as ``TurnRecord`` objects. New code should use ``OpenAIAdapter`` -and call ``finish_session()`` at trajectory end to drain trainable +SGLang ``/generate`` with ``input_ids``, and fold the turn into a per-session +turn-node :class:`~slime.agent.trajectory_manager.TrajectoryTree`. The tree +routes everything by text prefix, so there is no manual new/append/wipe +bookkeeping. Call ``finish_session()`` at trajectory end to drain trainable ``TokenSegment`` objects. """ @@ -20,25 +21,31 @@ from aiohttp import web -from slime.agent.adapters.common import ADAPTER_KEY, REASONING_PARSER_KEY, TOKENIZER_KEY, TOOL_PARSER_KEY -from slime.agent.adapters.common import AdapterChain as Chain -from slime.agent.adapters.common import BaseAdapter, call_sglang_generate +from slime.agent.adapters.common import ( + ADAPTER_KEY, + REASONING_PARSER_KEY, + TOKENIZER_KEY, + TOOL_PARSER_KEY, + BaseAdapter, + GenResult, + SessionTrajectory, + assemble_turns, + call_sglang_generate, +) from slime.agent.adapters.common import json_arguments as _json_arguments -from slime.agent.adapters.common import ok_response, render_token_ids, request_session_id -from slime.agent.adapters.common import stable_hash as _hash +from slime.agent.adapters.common import ok_response, render_prompt, request_session_id from slime.agent.parsing import ParsedModelOutput, parse_model_output -from slime.agent.trajectory import TokenSegment, TurnRecord, TurnSegment, make_turn_segment, merge_turn_segments +from slime.agent.trajectory_manager import TokenSegment, export_token_segments, record_turn logger = logging.getLogger(__name__) @dataclasses.dataclass class Session: - main: Chain = dataclasses.field(default_factory=Chain) + traj: SessionTrajectory = dataclasses.field(default_factory=SessionTrajectory) sampling_defaults: dict = dataclasses.field(default_factory=dict) max_context_tokens: int = 0 lock: asyncio.Lock = dataclasses.field(default_factory=asyncio.Lock) - segments: list[TurnSegment] = dataclasses.field(default_factory=list) class OpenAIAdapter(BaseAdapter): @@ -63,9 +70,7 @@ async def finish_session(self, sid: str, *, wait_timeout: float = 5.0) -> list[T s = self.store.pop(sid, None) if s is None: return [] - if s.main.turns: - s.segments.append(make_turn_segment(s.main.turns, kind="final")) - return merge_turn_segments(s.segments) + return export_token_segments(s.traj.tree) def _flatten_content(content: Any) -> str: @@ -239,48 +244,7 @@ def _responses_input_to_messages(input_value: Any, instructions: Any = None) -> return messages -def _select_kind(s: Session, messages: list[dict]) -> str: - target = s.main - msg_hashes = [_hash(m) for m in messages] - if target.seen_msgs == 0: - kind = "new" - else: - is_append = len(msg_hashes) >= target.seen_msgs and msg_hashes[: target.seen_msgs] == target.msg_hashes - if is_append: - kind = "append" - else: - if target.turns: - s.segments.append(make_turn_segment(target.turns, kind="wipe")) - kind = "wipe" - return kind - - -def _replace_chat_messages(target: Chain, messages: list[dict], tools_schema: list[dict] | None) -> None: - target.chat_messages = _translate_chat_messages(messages) - target.turns.clear() - target.seen_msgs = len(messages) - target.msg_hashes = [_hash(m) for m in messages] - if tools_schema is not None: - target.tools_schema = tools_schema - - -def _extend_chat_messages(target: Chain, messages: list[dict], tools_schema: list[dict] | None) -> None: - translated = _translate_chat_messages(messages[target.seen_msgs :]) - target.chat_messages.extend(translated) - target.seen_msgs = len(messages) - target.msg_hashes = [_hash(m) for m in messages] - if tools_schema is not None: - target.tools_schema = tools_schema - - -def _build_prompt(target: Chain, messages: list[dict], tools_schema: list[dict] | None, kind: str, tok) -> list[int]: - (_extend_chat_messages if kind == "append" else _replace_chat_messages)(target, messages, tools_schema) - return render_token_ids(target, tok) - - -async def _generate( - prompt_ids: list[int], s: Session, body: dict, app, *, session_id: str | None = None -) -> TurnRecord: +async def _generate(prompt_ids: list[int], s: Session, body: dict, app, *, session_id: str | None = None): return await call_sglang_generate( prompt_ids, s, @@ -294,12 +258,10 @@ async def _generate( ) -def _parse_turn(target: Chain, turn: TurnRecord, app) -> ParsedModelOutput: - tok = app[TOKENIZER_KEY] - raw_output = tok.decode(turn.output_ids, skip_special_tokens=False) if turn.output_ids else "" +def _parse_output(output_text: str, tools_schema: list[dict] | None, app) -> ParsedModelOutput: return parse_model_output( - raw_output, - tools_schema=target.tools_schema, + output_text or "", + tools_schema=tools_schema, tool_parser_name=app[TOOL_PARSER_KEY], reasoning_parser_name=app[REASONING_PARSER_KEY], ) @@ -365,25 +327,32 @@ def _request_session_id(request: web.Request, body: dict) -> str: async def _run_turn( request: web.Request, body: dict, messages: list[dict] -) -> tuple[TurnRecord, ParsedModelOutput, int, int]: +) -> tuple[GenResult, ParsedModelOutput, int, int]: sid = _request_session_id(request, body) adapter = request.app[ADAPTER_KEY] if sid in adapter.closed: raise web.HTTPServiceUnavailable(text="session closed") app = request.app + tok = app[TOKENIZER_KEY] s = adapter.store.setdefault(sid, Session()) task = asyncio.current_task() adapter.inflight.setdefault(sid, set()).add(task) try: async with s.lock: - target = s.main + translated = _translate_chat_messages(messages) tools_schema = _normalize_tools(body.get("tools")) - kind = _select_kind(s, messages) - prompt_ids = _build_prompt(target, messages, tools_schema, kind, app[TOKENIZER_KEY]) - turn = await _generate(prompt_ids, s, body, app, session_id=sid) - parsed = _parse_turn(target, turn, app) - target.turns.append(turn) - return turn, parsed, len(prompt_ids), len(turn.output_ids) + full_prompt_ids = render_prompt(s.traj, translated, tok, tools_schema) + gen = await _generate(full_prompt_ids, s, body, app, session_id=sid) + turns, pending_key = assemble_turns(s.traj, translated, tok, tools_schema, gen, full_prompt_ids) + node = record_turn(s.traj.tree, turns) + if node is not None: + s.traj.resp_truth[pending_key] = ( + list(gen.output_ids), + list(gen.output_log_probs), + gen.output_text, + ) + parsed = _parse_output(gen.output_text, tools_schema, app) + return gen, parsed, len(full_prompt_ids), len(gen.output_ids) finally: adapter.inflight.get(sid, set()).discard(task) @@ -393,10 +362,10 @@ async def _handle_chat_completions(request: web.Request) -> web.StreamResponse: messages = body.get("messages") or [] if not isinstance(messages, list): raise web.HTTPBadRequest(text="messages must be a list") - turn, parsed, in_tok, out_tok = await _run_turn(request, body, messages) + gen, parsed, in_tok, out_tok = await _run_turn(request, body, messages) if body.get("stream"): - return await _stream_chat_completion(request, body, parsed, turn.finish_reason, in_tok, out_tok) - return web.json_response(_chat_completion_response(body, parsed, turn.finish_reason, in_tok, out_tok)) + return await _stream_chat_completion(request, body, parsed, gen.finish_reason, in_tok, out_tok) + return web.json_response(_chat_completion_response(body, parsed, gen.finish_reason, in_tok, out_tok)) def _chat_completion_response( @@ -469,10 +438,10 @@ async def emit(choice_delta: dict[str, Any], finish_reason: str | None = None, u async def _handle_responses(request: web.Request) -> web.StreamResponse: body = await request.json() messages = _responses_input_to_messages(body.get("input", ""), body.get("instructions")) - turn, parsed, in_tok, out_tok = await _run_turn(request, body, messages) + gen, parsed, in_tok, out_tok = await _run_turn(request, body, messages) if body.get("stream"): - return await _stream_response(request, body, parsed, turn.finish_reason, in_tok, out_tok) - return web.json_response(_response_response(body, parsed, turn.finish_reason, in_tok, out_tok)) + return await _stream_response(request, body, parsed, gen.finish_reason, in_tok, out_tok) + return web.json_response(_response_response(body, parsed, gen.finish_reason, in_tok, out_tok)) def _response_output(parsed: ParsedModelOutput) -> list[dict[str, Any]]: diff --git a/slime/agent/trajectory.py b/slime/agent/trajectory.py deleted file mode 100644 index 51db112fde..0000000000 --- a/slime/agent/trajectory.py +++ /dev/null @@ -1,208 +0,0 @@ -"""Token-level trajectory helpers for agent rollouts.""" - -from __future__ import annotations - -import copy -import dataclasses -import logging -from typing import Any - -from slime.utils.types import Sample - - -logger = logging.getLogger(__name__) - - -@dataclasses.dataclass(frozen=True) -class TurnRecord: - """Exact token snapshot for one assistant generation. - - ``prompt_ids`` is the full tokenized prompt sent to the generator for that - turn. ``output_ids`` is the raw generated output, and - ``output_log_probs`` is aligned with it when the rollout engine returns - per-token log probabilities. - """ - - prompt_ids: list[int] - output_ids: list[int] - finish_reason: str - output_log_probs: list[float] = dataclasses.field(default_factory=list) - - -@dataclasses.dataclass(frozen=True) -class TokenSegment: - """One training segment assembled from an agent trajectory.""" - - prompt_ids: list[int] - response_ids: list[int] - loss_mask: list[int] - rollout_log_probs: list[float] = dataclasses.field(default_factory=list) - metadata: dict[str, Any] = dataclasses.field(default_factory=dict) - - -@dataclasses.dataclass(frozen=True) -class TurnSegment: - """A frozen group of turns before token-level merge.""" - - turns: list[TurnRecord] - metadata: dict[str, Any] = dataclasses.field(default_factory=dict) - - -def make_turn_segment( - turns: list[TurnRecord], - *, - kind: str = "", - metadata: dict[str, Any] | None = None, -) -> TurnSegment: - """Freeze turns and attach conventional segment metadata.""" - frozen_turns = list(turns) - segment_metadata = dict(metadata or {}) - if kind: - segment_metadata.setdefault("segment_kind", kind) - segment_metadata.setdefault("finish_reason", frozen_turns[-1].finish_reason if frozen_turns else "") - return TurnSegment(turns=frozen_turns, metadata=segment_metadata) - - -def _common_prefix_len(a: list[int], b: list[int]) -> int: - n = min(len(a), len(b)) - i = 0 - while i < n and a[i] == b[i]: - i += 1 - return i - - -def _output_log_probs(turn: TurnRecord) -> list[float]: - if len(turn.output_log_probs) == len(turn.output_ids): - return list(turn.output_log_probs) - logger.warning( - "[trajectory] turn logprob length mismatch; zeroing output logprobs (%d ids, %d logprobs)", - len(turn.output_ids), - len(turn.output_log_probs), - ) - return [0.0] * len(turn.output_ids) - - -def merge_turns(turns: list[TurnRecord], *, metadata: dict[str, Any] | None = None) -> TokenSegment | None: - """Replay turn records into one linear training segment. - - The first turn's prompt becomes the segment prompt. Later turn prompts are - stitched against ``prompt + response_so_far``. Any new prompt suffix is - non-model context and receives loss mask 0. If a later prompt diverges - inside a previous model output, the retained prefix of that whole output - turn is also masked out, because partial token matches are not a faithful - training target for that turn. - """ - if not turns: - return None - - prompt_ids = list(turns[0].prompt_ids) - response_ids: list[int] = [] - loss_mask: list[int] = [] - rollout_log_probs: list[float] = [] - output_spans: list[tuple[int, int]] = [] - - for i, turn in enumerate(turns): - if i > 0: - if turn.prompt_ids[: len(prompt_ids)] != prompt_ids: - logger.warning("[trajectory] merge prompt base changed; starting segment from drifted prompt") - prompt_ids = list(turn.prompt_ids) - response_ids = [] - loss_mask = [] - rollout_log_probs = [] - output_spans = [] - else: - prompt_suffix = turn.prompt_ids[len(prompt_ids) :] - matched_len = _common_prefix_len(response_ids, prompt_suffix) - if matched_len < len(response_ids): - logger.warning( - "[trajectory] merge prefix drift; truncating %d unstitched response tokens", - len(response_ids) - matched_len, - ) - for start, end in output_spans: - if start < matched_len < end: - loss_mask[start:matched_len] = [0] * (matched_len - start) - rollout_log_probs[start:matched_len] = [0.0] * (matched_len - start) - response_ids = response_ids[:matched_len] - loss_mask = loss_mask[:matched_len] - rollout_log_probs = rollout_log_probs[:matched_len] - output_spans = [ - (start, min(end, matched_len)) for start, end in output_spans if start < matched_len - ] - - context_tail = prompt_suffix[matched_len:] - response_ids.extend(context_tail) - loss_mask.extend([0] * len(context_tail)) - rollout_log_probs.extend([0.0] * len(context_tail)) - - output_start = len(response_ids) - response_ids.extend(turn.output_ids) - loss_mask.extend([1] * len(turn.output_ids)) - rollout_log_probs.extend(_output_log_probs(turn)) - output_spans.append((output_start, len(response_ids))) - - rollout_log_probs = [logprob if mask else 0.0 for logprob, mask in zip(rollout_log_probs, loss_mask, strict=True)] - - return TokenSegment( - prompt_ids=prompt_ids, - response_ids=response_ids, - loss_mask=loss_mask, - rollout_log_probs=rollout_log_probs, - metadata=dict(metadata or {}), - ) - - -def merge_turn_segments(segments: list[TurnSegment]) -> list[TokenSegment]: - """Merge frozen turn segments and keep every non-empty output.""" - out: list[TokenSegment] = [] - for turn_segment in segments: - token_segment = merge_turns(turn_segment.turns, metadata=turn_segment.metadata) - if token_segment is None: - continue - if token_segment.response_ids: - out.append(token_segment) - return out - - -def write_segment_to_sample(sample: Sample, segment: TokenSegment, reward: float, tokenizer) -> None: - """Populate token, mask, response, reward, and status fields from a segment.""" - sample.tokens = list(segment.prompt_ids) + list(segment.response_ids) - sample.response_length = len(segment.response_ids) - sample.loss_mask = list(segment.loss_mask) - sample.rollout_log_probs = list(segment.rollout_log_probs) - sample.response = tokenizer.decode(segment.response_ids, skip_special_tokens=False) - sample.reward = float(reward) - sample.status = Sample.Status.COMPLETED - - -def fan_out_sample_segments( - sample: Sample, - segments: list[TokenSegment], - reward: float, - tokenizer, - *, - metadata: dict[str, Any] | None = None, - rollout_id: int | None = None, -) -> list[Sample]: - """Emit one Sample per segment, splitting reward uniformly across them. - - Sibling samples share ``rollout_id`` so reducers that average by rollout do - not over-count trajectories split by compaction or sub-agent dispatch. - """ - k = len(segments) - per_segment_reward = float(reward) / max(1, k) - shared_rollout_id = getattr(sample, "index", None) if rollout_id is None else rollout_id - base_metadata = {**(sample.metadata or {}), **(metadata or {})} - - out: list[Sample] = [] - for i, segment in enumerate(segments): - sub = sample if i == 0 else copy.copy(sample) - write_segment_to_sample(sub, segment, per_segment_reward, tokenizer) - sub.rollout_id = shared_rollout_id - sub.metadata = { - **base_metadata, - **(segment.metadata or {}), - "segment_idx": i, - "num_segments": k, - } - out.append(sub) - return out diff --git a/slime/agent/trajectory_manager.py b/slime/agent/trajectory_manager.py new file mode 100644 index 0000000000..2c4e251e75 --- /dev/null +++ b/slime/agent/trajectory_manager.py @@ -0,0 +1,552 @@ +"""Turn-node trajectory manager for agent rollouts. + +A session's trajectory is modeled as a **text-prefix tree**: each ``Node`` is one +turn (incremental prompt segment + response segment). Text-prefix matching is the +primary signal; token-id comparison is secondary and only handles TITO +(text-in-token-out) drift after the text prefix already matches. ``export`` walks +the tree once and yields full-length ``(tokens, masks, logprobs)`` triples, one +sample per leaf. + +The core algorithm (``PromptSeg``/``RespSeg``/``Node``/``TrajectoryTree``/ +``MatchResult`` and ``match_prefix``/``attach_turn``/``record_turn``/``export``) +is the standalone turn-node design; see +``0601-Trajectory-manager/02-turn-node/02-turn-node-design.md``. + +The slime glue below the core (``TokenSegment``/``export_token_segments``/ +``fan_out_sample_segments``/``write_segment_to_sample``) is the only intentional +deviation from the design's zero-dependency constraint: it imports ``Sample`` and +adapts the full-length export into slime's prompt/response training segments. +""" + +from __future__ import annotations + +import copy +import dataclasses +import logging +import time +import uuid +from dataclasses import dataclass, field +from typing import Any + +from slime.utils.types import Sample + +logger = logging.getLogger(__name__) + + +# ============================================================================= +# Core turn-node tree (standalone design; zero training-side dependency) +# ============================================================================= + + +@dataclass +class PromptSeg: + text: str + tokens: list[int] + + +@dataclass +class RespSeg: + text: str + tokens: list[int] + logprobs: list[float] + + +@dataclass +class Node: + id: str + prompt_delta_text: str + prompt_delta_tokens: list[int] + resp_text: str + resp_tokens: list[int] + resp_logprobs: list[float] + parent: Node | None + children: list[Node] = field(default_factory=list) + loss_mask: bool = False + created_at: float = 0.0 + updated_at: float = 0.0 + replaced_count: int = 0 + warning: str | None = None + kind_is_root: bool = False + + @staticmethod + def new_root() -> Node: + now = time.time() + return Node( + id=uuid.uuid4().hex[:12], + prompt_delta_text="", + prompt_delta_tokens=[], + resp_text="", + resp_tokens=[], + resp_logprobs=[], + parent=None, + children=[], + created_at=now, + updated_at=now, + kind_is_root=True, + ) + + @staticmethod + def new_turn(prompt: PromptSeg, resp: RespSeg, parent: Node) -> Node: + now = time.time() + node = Node( + id=uuid.uuid4().hex[:12], + prompt_delta_text=prompt.text, + prompt_delta_tokens=list(prompt.tokens), + resp_text=resp.text, + resp_tokens=list(resp.tokens), + resp_logprobs=list(resp.logprobs), + parent=parent, + children=[], + created_at=now, + updated_at=now, + ) + parent.children.append(node) + return node + + @property + def node_text(self) -> str: + return self.prompt_delta_text + self.resp_text + + +@dataclass +class TrajectoryTree: + root: Node = field(default_factory=Node.new_root) + + +@dataclass +class MatchResult: + case: str + anchor: Node | None = None + text_matched_len: int = 0 + drift_nodes: list[Node] = field(default_factory=list) + lca: Node | None = None + fork_node: Node | None = None + residual_count: int = 0 + + +# ---- helpers ---- + + +def _text_lcp(a: str, b: str) -> int: + n = min(len(a), len(b)) + i = 0 + while i < n and a[i] == b[i]: + i += 1 + return i + + +def _chain_nodes(node: Node) -> list[Node]: + """root->node path, non-root nodes, ordered root to leaf.""" + out = [] + cur = node + while cur is not None and not cur.kind_is_root: + out.append(cur) + cur = cur.parent + out.reverse() + return out + + +def _chain_text(nodes: list[Node]) -> str: + return "".join(n.node_text for n in nodes) + + +def _collect_leaves(root: Node) -> list[Node]: + """Pre-order collect of all leaves (excluding root; empty if root childless).""" + leaves = [] + + def walk(n: Node): + if not n.children: + if not n.kind_is_root: + leaves.append(n) + return + for c in n.children: + walk(c) + + walk(root) + return leaves + + +def _log_replace(case, reason, node, before_text, after_text, before_tokens, after_tokens): + logger.warning( + "trajectory replace | case=%s reason=%s node=%s\n" " text : %r -> %r\n" " token: %s -> %s", + case, + reason, + node.id, + before_text, + after_text, + before_tokens, + after_tokens, + ) + + +# ---- match_prefix ---- + + +def _starting_chains(root: Node) -> list[Node]: + """ROOT's direct children = the head of each starting chain.""" + return list(root.children) + + +def _leaf_of_chain(head: Node) -> Node: + """Walk children[0..] to the leaf (single-chain helper).""" + cur = head + while cur.children: + cur = cur.children[0] + return cur + + +def match_prefix(tree: TrajectoryTree, turns) -> MatchResult: + incoming_text = "".join(p.text + r.text for (p, r) in turns) + + # For each root->leaf chain take the longest text prefix; pick the longest, + # ties broken by the latest created_at leaf. + best_leaf = None + best_len = -1 + for leaf in _collect_leaves(tree.root): + nodes = _chain_nodes(leaf) + ctext = _chain_text(nodes) + lcp = _text_lcp(ctext, incoming_text) + if lcp > best_len or (lcp == best_len and best_leaf is not None and leaf.created_at > best_leaf.created_at): + best_leaf = leaf + best_len = lcp + + # Case 1: no text overlap at all + if best_leaf is None or best_len == 0: + return MatchResult(case="case1", anchor=tree.root, text_matched_len=0) + + nodes = _chain_nodes(best_leaf) + ctext = _chain_text(nodes) + + # §7.1: incoming text <= chain text and fully covered (no new tail, incl. equal) + if best_len == len(incoming_text) and len(incoming_text) <= len(ctext): + return MatchResult(case="substring", anchor=best_leaf, text_matched_len=best_len) + + # Text reached leaf and incoming has a new tail -> TITO decision (Case 2/3/4) + if best_len == len(ctext) and len(incoming_text) > len(ctext): + drift = [] + for i, node in enumerate(nodes): + p, r = turns[i] + if list(p.tokens) != node.prompt_delta_tokens or list(r.tokens) != node.resp_tokens: + drift.append(node) + if not drift: + return MatchResult(case="case2", anchor=best_leaf, text_matched_len=best_len) + only_tail = len(drift) == 1 and drift[0] is nodes[-1] + return MatchResult( + case="case3" if only_tail else "case4", + anchor=best_leaf, + text_matched_len=best_len, + drift_nodes=drift, + ) + + # Text stops inside the chain (0 < best_len < len(ctext)) -> Case 5/6 + cum = 0 + fork_node = None + fork_idx = 0 + for idx, node in enumerate(nodes): + seg_len = len(node.node_text) + if cum + seg_len > best_len: + fork_node = node + fork_idx = idx + break + cum += seg_len + if fork_node is None: + fork_node = nodes[-1] + fork_idx = len(nodes) - 1 + residual_count = len(nodes) - fork_idx # X..leaf + lca = fork_node.parent + return MatchResult( + case="case5" if residual_count <= 1 else "case6", + anchor=best_leaf, + text_matched_len=best_len, + lca=lca, + fork_node=fork_node, + residual_count=residual_count, + ) + + +# ---- attach_turn ---- + + +def attach_turn(tree: TrajectoryTree, match: MatchResult, turns) -> Node | None: + last_p, last_r = turns[-1] + + if match.case == "substring": + logger.warning( + "trajectory drop | reason=substring-prefix redelivery; " + "incoming is a true prefix of an existing chain, dropped" + ) + return None + + if match.case == "case1": + return Node.new_turn(last_p, last_r, parent=tree.root) + + if match.case == "case2": + return Node.new_turn(last_p, last_r, parent=match.anchor) + + if match.case in ("case3", "case4"): + nodes = _chain_nodes(match.anchor) + # Replace each drifted node's segment in place (same-index turn segment). + idx_by_node = {id(n): i for i, n in enumerate(nodes)} + for node in match.drift_nodes: + i = idx_by_node[id(node)] + p, r = turns[i] + before_text, before_tokens = node.node_text, list(node.resp_tokens) + node.prompt_delta_text = p.text + node.prompt_delta_tokens = list(p.tokens) + node.resp_text = r.text + node.resp_tokens = list(r.tokens) + node.resp_logprobs = [0.0] * len(r.tokens) # §2.4 placeholder + node.replaced_count += 1 + node.updated_at = time.time() + reason = ( + "tokenizer TITO drift at tail turn" + if match.case == "case3" + else "multi-turn token drift (template changed?)" + ) + _log_replace(match.case, reason, node, before_text, node.node_text, before_tokens, node.resp_tokens) + # mask interval + if match.case == "case3": + match.drift_nodes[0].loss_mask = True + else: + first = match.drift_nodes[0] + first.warning = "multi-turn token drift" + start = idx_by_node[id(first)] + for node in nodes[start:]: + node.loss_mask = True + # append the new turn at the tail + return Node.new_turn(last_p, last_r, parent=nodes[-1]) + + if match.case == "case5": + nodes = _chain_nodes(match.anchor) + x = match.fork_node + xi = nodes.index(x) + p, r = turns[xi] + before_text, before_tokens = x.node_text, list(x.resp_tokens) + x.prompt_delta_text = p.text + x.prompt_delta_tokens = list(p.tokens) + x.resp_text = r.text + x.resp_tokens = list(r.tokens) + x.resp_logprobs = [0.0] * len(r.tokens) + x.loss_mask = True + x.replaced_count += 1 + x.updated_at = time.time() + _log_replace( + "case5", "upstream response format rewrite", x, before_text, x.node_text, before_tokens, x.resp_tokens + ) + # if turns has newer turns after x, append them in order + parent = x + for j in range(xi + 1, len(turns)): + pj, rj = turns[j] + parent = Node.new_turn(pj, rj, parent=parent) + return parent + + if match.case == "case6": + nodes = _chain_nodes(match.anchor) + x = match.fork_node + xi = nodes.index(x) + lca = match.lca # = x.parent + parent = lca + last = None + for j in range(xi, len(turns)): + pj, rj = turns[j] + last = Node.new_turn(pj, rj, parent=parent) + parent = last + return last + + raise ValueError(f"unknown match case: {match.case}") + + +# ---- record_turn ---- + + +def record_turn(tree: TrajectoryTree, turns) -> Node | None: + match = match_prefix(tree, turns) + return attach_turn(tree, match, turns) + + +# ---- export ---- + + +def _subtree_has_trainable(node: Node) -> bool: + """Whether any node in the subtree (rooted at ``node``) carries a trainable + response segment (``loss_mask`` False and non-empty ``resp_tokens``). Used by + owner selection so a fork's owner branch is not later dropped by the slice + layer for being all-mask.""" + stack = [node] + while stack: + cur = stack.pop() + if not cur.kind_is_root and not cur.loss_mask and cur.resp_tokens: + return True + stack.extend(cur.children) + return False + + +def _decide_fork_owners(root: Node) -> dict: + owners = {} + + def walk(n: Node): + if len(n.children) >= 2: + # Decision E: prefer the earliest child whose subtree still has a + # trainable leaf, so the shared prefix it owns survives slicing. + # Fall back to the earliest child when every branch is all-mask + # (then the shared prefix has nowhere trainable to go anyway). + trainable = [c for c in n.children if _subtree_has_trainable(c)] + pool = trainable or list(n.children) + owner = min(pool, key=lambda c: c.created_at) + owners[id(n)] = owner + for c in n.children: + walk(c) + + walk(root) + return owners + + +def _masked_by_fork(node: Node, leaf: Node, fork_owner: dict) -> bool: + """Whether ``node`` is force-masked because some fork ancestor ``F`` on the + leaf's path chose a non-owner branch toward ``leaf``: i.e. ``node`` is at or + before ``F`` (incl. ``F`` itself) and at ``F`` the leaf's branch != owner.""" + # leaf's root->leaf path (incl. root) + path = [] + cur = leaf + while cur is not None: + path.append(cur) + cur = cur.parent + path.reverse() # root ... leaf + pos = {id(n): i for i, n in enumerate(path)} + node_pos = pos.get(id(node)) + if node_pos is None: + return False + # for every fork point F on path (present in fork_owner) + for i, F in enumerate(path): + if id(F) not in fork_owner: + continue + # leaf's chosen child at F = path[i+1] + if i + 1 >= len(path): + continue + chosen_child = path[i + 1] + owner = fork_owner[id(F)] + if chosen_child is not owner: + # non-owner branch: mask nodes at or before F (root..F) + if node_pos <= i: + return True + return False + + +def export(tree: TrajectoryTree): + samples, masks, logprobs = [], [], [] + fork_owner = _decide_fork_owners(tree.root) + for leaf in _collect_leaves(tree.root): + nodes = _chain_nodes(leaf) + toks, mask, lp = [], [], [] + for node in nodes: + toks += node.prompt_delta_tokens + mask += [0] * len(node.prompt_delta_tokens) + lp += [0.0] * len(node.prompt_delta_tokens) + toks += node.resp_tokens + if _masked_by_fork(node, leaf, fork_owner): + bit = 0 + elif node.loss_mask: + bit = 0 + else: + bit = 1 + mask += [bit] * len(node.resp_tokens) + lp += list(node.resp_logprobs) + samples.append(toks) + masks.append(mask) + logprobs.append(lp) + return samples, masks, logprobs + + +# ============================================================================= +# slime glue: full-length export -> prompt/response TokenSegment + Sample fan-out +# ============================================================================= + + +@dataclasses.dataclass(frozen=True) +class TokenSegment: + """One training segment assembled from an agent trajectory. + + slime training invariants: ``tokens = prompt_ids + response_ids``, + ``response_length = len(response_ids)`` and + ``len(loss_mask) == len(rollout_log_probs) == response_length``. + """ + + prompt_ids: list[int] + response_ids: list[int] + loss_mask: list[int] + rollout_log_probs: list[float] = dataclasses.field(default_factory=list) + metadata: dict[str, Any] = dataclasses.field(default_factory=dict) + + +def export_token_segments(tree: TrajectoryTree, *, metadata: dict[str, Any] | None = None) -> list[TokenSegment]: + """Turn the tree's full-length export into slime ``TokenSegment``s. + + Each leaf chain is split at the first trainable token (design §6, + ``mask.index(1)``): tokens before the cut are the prompt, tokens from the cut + on are the response. Leaves whose whole chain is masked (no trainable token, + e.g. fork non-owner or Case4/5 fully-masked) are dropped -- they have no + trainable response and would break slime's ``response_length`` contract. + """ + out: list[TokenSegment] = [] + tokens_list, masks_list, logprobs_list = export(tree) + for tokens, mask, logprobs in zip(tokens_list, masks_list, logprobs_list, strict=True): + if 1 not in mask: + continue # all-mask chain: nothing trainable, drop + cut = mask.index(1) + response_ids = tokens[cut:] + if not response_ids: + continue + segment = TokenSegment( + prompt_ids=list(tokens[:cut]), + response_ids=list(response_ids), + loss_mask=list(mask[cut:]), + rollout_log_probs=list(logprobs[cut:]), + metadata={**(metadata or {}), "segment_kind": "leaf"}, + ) + assert len(segment.loss_mask) == len(segment.rollout_log_probs) == len(segment.response_ids) + out.append(segment) + return out + + +def write_segment_to_sample(sample: Sample, segment: TokenSegment, reward: float, tokenizer) -> None: + """Populate token, mask, response, reward, and status fields from a segment.""" + sample.tokens = list(segment.prompt_ids) + list(segment.response_ids) + sample.response_length = len(segment.response_ids) + sample.loss_mask = list(segment.loss_mask) + sample.rollout_log_probs = list(segment.rollout_log_probs) + sample.response = tokenizer.decode(segment.response_ids, skip_special_tokens=False) + sample.reward = float(reward) + sample.status = Sample.Status.COMPLETED + + +def fan_out_sample_segments( + sample: Sample, + segments: list[TokenSegment], + reward: float, + tokenizer, + *, + metadata: dict[str, Any] | None = None, +) -> list[Sample]: + """Emit one Sample per segment, splitting reward uniformly across them. + + Sibling samples share ``group_id`` so reducers that average by group do + not over-count trajectories split by compaction or sub-agent dispatch. + """ + k = len(segments) + per_segment_reward = float(reward) / max(1, k) + shared_group_id = sample.group_id if sample.group_id is not None else sample.index + base_metadata = {**(sample.metadata or {}), **(metadata or {})} + + out: list[Sample] = [] + for i, segment in enumerate(segments): + sub = sample if i == 0 else copy.copy(sample) + write_segment_to_sample(sub, segment, per_segment_reward, tokenizer) + sub.group_id = shared_group_id + sub.metadata = { + **base_metadata, + **(segment.metadata or {}), + "segment_idx": i, + "num_segments": k, + } + out.append(sub) + return out From f10552f571c5f5c006b673586f483b6d7d83b1e1 Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Tue, 2 Jun 2026 02:43:03 +0000 Subject: [PATCH 02/28] refactor(agent): drop dead code left by trajectory-manager refactor Remove vestigial bookkeeping the turn-node TrajectoryTree made redundant: * anthropic adapter: the always-empty dispatch_id plumbing in _anthropic_blocks / _build_reply (routing is now done by the tree, not by tool_use ids). * hoist the byte-identical Session dataclass and finish_session method from both adapters into common.BaseAdapter (shared session_cls + export_token_segments drain). * trajectory_manager: delete the unreferenced _starting_chains / _leaf_of_chain helpers. No behavior change; agent adapter and trajectory tests pass. --- slime/agent/adapters/anthropic.py | 38 ++++++++----------------------- slime/agent/adapters/common.py | 29 +++++++++++++++++++---- slime/agent/adapters/openai.py | 22 ++---------------- slime/agent/trajectory_manager.py | 13 ----------- 4 files changed, 36 insertions(+), 66 deletions(-) diff --git a/slime/agent/adapters/anthropic.py b/slime/agent/adapters/anthropic.py index 3a4f52a098..5b6d39fc5f 100644 --- a/slime/agent/adapters/anthropic.py +++ b/slime/agent/adapters/anthropic.py @@ -14,7 +14,6 @@ from __future__ import annotations import asyncio -import dataclasses import json import logging import secrets @@ -28,7 +27,7 @@ TOKENIZER_KEY, TOOL_PARSER_KEY, BaseAdapter, - SessionTrajectory, + Session, assemble_turns, call_sglang_generate, ok_response, @@ -36,24 +35,14 @@ request_session_id, ) from slime.agent.parsing import parse_model_output -from slime.agent.trajectory_manager import TokenSegment, export_token_segments, record_turn +from slime.agent.trajectory_manager import record_turn logger = logging.getLogger(__name__) -@dataclasses.dataclass -class Session: - traj: SessionTrajectory = dataclasses.field(default_factory=SessionTrajectory) - sampling_defaults: dict = dataclasses.field(default_factory=dict) - max_context_tokens: int = 0 - lock: asyncio.Lock = dataclasses.field(default_factory=asyncio.Lock) - - class AnthropicAdapter(BaseAdapter): """Anthropic Messages-compatible HTTP adapter with session lifecycle helpers.""" - session_cls = Session - def __init__(self, *, tokenizer, sglang_url, tool_parser=None, reasoning_parser=None) -> None: super().__init__( tokenizer=tokenizer, @@ -66,13 +55,6 @@ def __init__(self, *, tokenizer, sglang_url, tool_parser=None, reasoning_parser= self.app.router.add_get("/healthz", _ok) self.app.router.add_get("/v1/models", _ok) - async def finish_session(self, sid: str, *, wait_timeout: float = 5.0) -> list[TokenSegment]: - await self.shutdown_session(sid, wait_timeout=wait_timeout) - s = self.store.pop(sid, None) - if s is None: - return [] - return export_token_segments(s.traj.tree) - # ============================================================================= # Translation (Anthropic wire <-> chat-template messages) -- unchanged @@ -185,15 +167,14 @@ async def _generate(prompt_ids: list[int], s: Session, body: dict, app, *, sessi ) -def _build_reply(output_text: str, finish: str, tools_schema: list[dict] | None, app) -> tuple[list[dict], str, str]: +def _build_reply(output_text: str, finish: str, tools_schema: list[dict] | None, app) -> tuple[list[dict], str]: """Turn the model's raw output text into the reply we send back to claude-code. 1. parse decoded text -> (thinking, visible, tool_uses) via sglang parsers 2. pack into Anthropic content blocks 3. derive stop_reason: 'tool_use' | 'max_tokens' | 'end_turn' - Returns (blocks, stop_reason, dispatch_id). dispatch_id is retained for wire - compatibility but no longer drives routing (the tree splits sub-agents). + Returns (blocks, stop_reason). """ parsed = parse_model_output( output_text or "", @@ -201,24 +182,23 @@ def _build_reply(output_text: str, finish: str, tools_schema: list[dict] | None, tool_parser_name=app[TOOL_PARSER_KEY], reasoning_parser_name=app[REASONING_PARSER_KEY], ) - blocks, dispatch_id = _anthropic_blocks(parsed.reasoning, parsed.text, parsed.tool_uses) - return blocks, _stop_reason(parsed.tool_uses, finish), dispatch_id + blocks = _anthropic_blocks(parsed.reasoning, parsed.text, parsed.tool_uses) + return blocks, _stop_reason(parsed.tool_uses, finish) -def _anthropic_blocks(thinking: str, visible: str, tool_uses: list[dict]) -> tuple[list[dict], str]: +def _anthropic_blocks(thinking: str, visible: str, tool_uses: list[dict]) -> list[dict]: """Pack parsed model output into Anthropic content blocks.""" blocks: list[dict] = [] if thinking: blocks.append({"type": "thinking", "thinking": thinking}) if visible: blocks.append({"type": "text", "text": visible}) - dispatch_id = "" for tu in tool_uses: tu_id = f"toolu_{secrets.token_hex(8)}" blocks.append({"type": "tool_use", "id": tu_id, "name": tu["name"], "input": tu["input"]}) if not blocks: blocks.append({"type": "text", "text": ""}) - return blocks, dispatch_id + return blocks def _stop_reason(tool_uses: list[dict], finish: str) -> str: @@ -263,7 +243,7 @@ async def _handle_request(request: web.Request) -> web.StreamResponse: list(gen.output_log_probs), gen.output_text, ) - blocks, stop, _did = _build_reply(gen.output_text, gen.finish_reason, tools_schema, app) + blocks, stop = _build_reply(gen.output_text, gen.finish_reason, tools_schema, app) in_tok, out_tok = len(full_prompt_ids), len(gen.output_ids) if body.get("stream") is True or "text/event-stream" in request.headers.get("Accept", ""): return await _stream_response(request, blocks, stop, in_tok, out_tok) diff --git a/slime/agent/adapters/common.py b/slime/agent/adapters/common.py index 72c74191cd..c3705cf942 100644 --- a/slime/agent/adapters/common.py +++ b/slime/agent/adapters/common.py @@ -31,8 +31,7 @@ import aiohttp from aiohttp import web -from slime.agent.trajectory_manager import PromptSeg, RespSeg, TokenSegment, TrajectoryTree - +from slime.agent.trajectory_manager import PromptSeg, RespSeg, TokenSegment, TrajectoryTree, export_token_segments ADAPTER_KEY = web.AppKey("adapter", object) TOKENIZER_KEY = web.AppKey("tokenizer", object) @@ -77,10 +76,25 @@ class SessionTrajectory: render_memo: dict[tuple[str, bool], list[int]] = dataclasses.field(default_factory=dict) +@dataclasses.dataclass +class Session: + """Per-session adapter state shared by the OpenAI and Anthropic adapters. + + Holds the trajectory (turn-node tree + truth cache), the per-session sampling + defaults / context budget injected at ``open_session``, and a lock that + serializes concurrent requests carrying the same session id. + """ + + traj: SessionTrajectory = dataclasses.field(default_factory=SessionTrajectory) + sampling_defaults: dict = dataclasses.field(default_factory=dict) + max_context_tokens: int = 0 + lock: asyncio.Lock = dataclasses.field(default_factory=asyncio.Lock) + + class BaseAdapter: """Base HTTP adapter with per-instance session lifecycle state.""" - session_cls: type + session_cls: type = Session def __init__(self, *, tokenizer, sglang_url, tool_parser=None, reasoning_parser=None) -> None: self.store: dict[str, Any] = {} @@ -112,7 +126,14 @@ async def shutdown_session(self, sid: str, *, wait_timeout: float = 5.0) -> None await shutdown_session_tasks(sid, self.closed, self.inflight, wait_timeout=wait_timeout) async def finish_session(self, sid: str, *, wait_timeout: float = 5.0) -> list[TokenSegment]: - raise NotImplementedError + """Drain a session: wait out in-flight requests, then export trainable + ``TokenSegment``s from its trajectory tree. Idempotent -- a second call + for an already-popped sid returns ``[]``.""" + await self.shutdown_session(sid, wait_timeout=wait_timeout) + s = self.store.pop(sid, None) + if s is None: + return [] + return export_token_segments(s.traj.tree) def strip_cache_control(obj: Any) -> Any: diff --git a/slime/agent/adapters/openai.py b/slime/agent/adapters/openai.py index cb8ed605fb..79a1b0b824 100644 --- a/slime/agent/adapters/openai.py +++ b/slime/agent/adapters/openai.py @@ -12,7 +12,6 @@ from __future__ import annotations import asyncio -import dataclasses import json import logging import secrets @@ -28,31 +27,21 @@ TOOL_PARSER_KEY, BaseAdapter, GenResult, - SessionTrajectory, + Session, assemble_turns, call_sglang_generate, ) from slime.agent.adapters.common import json_arguments as _json_arguments from slime.agent.adapters.common import ok_response, render_prompt, request_session_id from slime.agent.parsing import ParsedModelOutput, parse_model_output -from slime.agent.trajectory_manager import TokenSegment, export_token_segments, record_turn +from slime.agent.trajectory_manager import record_turn logger = logging.getLogger(__name__) -@dataclasses.dataclass -class Session: - traj: SessionTrajectory = dataclasses.field(default_factory=SessionTrajectory) - sampling_defaults: dict = dataclasses.field(default_factory=dict) - max_context_tokens: int = 0 - lock: asyncio.Lock = dataclasses.field(default_factory=asyncio.Lock) - - class OpenAIAdapter(BaseAdapter): """OpenAI-compatible HTTP adapter with session lifecycle helpers.""" - session_cls = Session - def __init__(self, *, tokenizer, sglang_url, tool_parser=None, reasoning_parser=None) -> None: super().__init__( tokenizer=tokenizer, @@ -65,13 +54,6 @@ def __init__(self, *, tokenizer, sglang_url, tool_parser=None, reasoning_parser= self.app.router.add_get("/healthz", _ok) self.app.router.add_get("/v1/models", _ok) - async def finish_session(self, sid: str, *, wait_timeout: float = 5.0) -> list[TokenSegment]: - await self.shutdown_session(sid, wait_timeout=wait_timeout) - s = self.store.pop(sid, None) - if s is None: - return [] - return export_token_segments(s.traj.tree) - def _flatten_content(content: Any) -> str: """Flatten OpenAI text/content parts into a chat-template string.""" diff --git a/slime/agent/trajectory_manager.py b/slime/agent/trajectory_manager.py index 2c4e251e75..17333601e2 100644 --- a/slime/agent/trajectory_manager.py +++ b/slime/agent/trajectory_manager.py @@ -182,19 +182,6 @@ def _log_replace(case, reason, node, before_text, after_text, before_tokens, aft # ---- match_prefix ---- -def _starting_chains(root: Node) -> list[Node]: - """ROOT's direct children = the head of each starting chain.""" - return list(root.children) - - -def _leaf_of_chain(head: Node) -> Node: - """Walk children[0..] to the leaf (single-chain helper).""" - cur = head - while cur.children: - cur = cur.children[0] - return cur - - def match_prefix(tree: TrajectoryTree, turns) -> MatchResult: incoming_text = "".join(p.text + r.text for (p, r) in turns) From ba557aee2324e3a152577ac28321180b48c970b9 Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Thu, 4 Jun 2026 12:50:44 +0000 Subject: [PATCH 03/28] refactor(agent): port anthropic + trajectory_manager from trajectory-manager-migration-v2 Bring over the four wire/manager files from trajectory-manager-migration-v2 to land the same TrajectoryManager-based anthropic adapter on this branch: - examples/coding_agent_rl/{README,generate}.py: switch generate() to the list[Sample] return shape from adapter.finish_session, document the env knob SLIME_TITO_SNAPSHOT_MIN_LOSS_TOKENS. - slime/agent/adapters/anthropic.py: absorb the wire-side scrub / mid-list system fold / per-sid turn cap / cc title-gen skip, route through TrajectoryManager. - slime/agent/adapters/common.py: slim to the shared primitives still used by the anthropic path (TurnRecord, BaseAdapter, call_sglang_generate, shutdown_session_tasks, ok_response). - slime/agent/trajectory_manager.py: replace the segment-based path with the DFS routing + LCP alignment + TITO snapshot rescue implementation. openai.py is intentionally left untouched; adapters/__init__.py drops the OpenAIAdapter export so the package still imports under the slimmed common.py. The OpenAI adapter and its tests do not work under this commit and will be cleaned up in a follow-up. --- examples/coding_agent_rl/README.md | 22 +- examples/coding_agent_rl/generate.py | 76 ++- slime/agent/adapters/__init__.py | 3 +- slime/agent/adapters/anthropic.py | 515 +++++++++++++-- slime/agent/adapters/common.py | 259 +------- slime/agent/trajectory_manager.py | 947 +++++++++++++-------------- 6 files changed, 979 insertions(+), 843 deletions(-) diff --git a/examples/coding_agent_rl/README.md b/examples/coding_agent_rl/README.md index 4ff8d5bdff..38e106a20e 100644 --- a/examples/coding_agent_rl/README.md +++ b/examples/coding_agent_rl/README.md @@ -145,18 +145,16 @@ The Anthropic adapter therefore follows a **string in, token out** contract: Multi-turn agents still force the adapter to tokenize later message histories, because tool observations and claude-code's own compacted messages -arrive as strings. `slime.agent.trajectory_manager` models each session as a -turn-node text-prefix tree and folds every request into it: - -- Each request is split into strictly-alternating `(prompt, response)` turns; - the current turn's prompt segment is taken by a segment render diff and its - response comes from SGLang, while historical responses are served from a - per-session truth cache so they keep their exact sampled token ids. -- New prompt suffixes that are tool/user/environment context export with - `loss_mask=0`; fresh model outputs export with `loss_mask=1`. -- Text divergence (compaction / history rewrite / sub-agent) branches the tree; - tail-turn token drift (TITO) is repaired in place and masked. `export` - deduplicates shared fork prefixes so they are trained at most once. +arrive as strings. `slime.agent.trajectory.merge_turns` stitches those later +prompts against the saved token stream: + +- New prompt suffixes that are tool/user/environment context are appended with + `loss_mask=0`. +- Fresh model outputs from SGLang are appended with `loss_mask=1`. +- If a later prompt no longer token-matches an earlier sampled output, the + unmatched suffix is dropped. If the drift cuts through the middle of a + previous model output, the retained prefix of that whole output turn is also + assigned `loss_mask=0`. That last case is the important correctness guard. A re-tokenization mismatch can make a string-level conversation look continuous while token-level diff --git a/examples/coding_agent_rl/generate.py b/examples/coding_agent_rl/generate.py index 2ce0aae815..a858eda060 100644 --- a/examples/coding_agent_rl/generate.py +++ b/examples/coding_agent_rl/generate.py @@ -9,11 +9,12 @@ 1. ``sandbox.run_claude_code`` prepares the agent sandbox and runs claude-code. 2. ``sandbox.git_diff`` captures the model-produced patch. 3. ``sandbox.evaluate`` scores that patch in a second clean sandbox. - 4. ``_merge_samples`` combines reward + adapter ``TokenSegment``s, - delegating segment-to-``Sample`` fan-out to ``slime.agent.trajectory_manager``. + 4. ``_merge_samples`` combines reward + the ``list[Sample]`` returned by + ``adapter.finish_session(sid)`` (which drains the per-sid trajectory + tree inside ``TrajectoryManager``). All sandbox-side details live in ``sandbox.py``; the LLM plumbing -(Anthropic <-> SGLang /generate, token capture, turn-node trajectory tree) uses +(Anthropic <-> SGLang /generate, token capture, 3-kind segment split) uses ``slime.agent.adapters.AnthropicAdapter``. Dataset row ``metadata`` schema:: @@ -53,7 +54,6 @@ from typing import Any from slime.agent.adapters import AnthropicAdapter -from slime.agent.trajectory_manager import TokenSegment, fan_out_sample_segments from slime.utils.misc import SingletonMeta from slime.utils.processing_utils import load_tokenizer from slime.utils.types import Sample @@ -97,11 +97,22 @@ def __init__(self, args) -> None: "Without it the sandbox cannot dial back and the rollout will " "silently abort." ) + # Snapshot threshold: 0 disables; absent or malformed env => default 1000. + _snap_env = os.environ.get("SLIME_TITO_SNAPSHOT_MIN_LOSS_TOKENS") + try: + _snap_threshold = int(_snap_env) if _snap_env is not None else 1000 + except ValueError: + logger.warning( + "SLIME_TITO_SNAPSHOT_MIN_LOSS_TOKENS=%r is not an int; using 1000", + _snap_env, + ) + _snap_threshold = 1000 self.adapter = AnthropicAdapter( tokenizer=self.tokenizer, sglang_url=sglang_url, tool_parser=self.tool_parser, reasoning_parser=self.reasoning_parser, + tito_snapshot_min_loss_tokens=_snap_threshold, ) # handler_cancellation=True so a client disconnect cancels the handler # coroutine, arming the fire-and-forget /abort_request inside the @@ -128,9 +139,9 @@ def __init__(self, args) -> None: # --------------------------------------------------------------------------- # Trajectory -> Sample conversion -# adapter.finish_session() returns TokenSegments. One trajectory yields >=1 -# segments because the agent may compact + reset mid-run; trajectory.py handles -# the mechanical segment -> Sample fan-out. +# adapter.finish_session(sid) drains the per-sid tree in TrajectoryManager and +# returns a list[Sample]. One trajectory yields >=1 samples because the agent +# may compact + reset mid-run, forking sub-trees that each become a sample. # --------------------------------------------------------------------------- @dataclass(frozen=True) class RewardResult: @@ -168,33 +179,42 @@ def _merge_samples( *, sample: Sample, state: _State, - segments: list[TokenSegment], + samples: list[Sample], reward_result: RewardResult, elapsed_sec: float, instance_id: str, -): - if not segments: +) -> list[Sample]: + """Decorate per-leaf Samples returned by TrajectoryManager.get_trajectory. + + The manager already filled tokens / loss_mask / rollout_log_probs / + response_length / reward (reward / N). We add per-trajectory metadata + (is_solved / applied_cleanly / elapsed_sec / segment_idx) and decode + ``sample.response`` from the response tokens slice -- slime's training + logging path reads this string. + """ + if not samples: return _abort_result(sample, "adapter_session_empty") trajectory_metadata = { - **(sample.metadata or {}), "instance_id": instance_id, "is_solved": reward_result.is_solved, "applied_cleanly": reward_result.applied_cleanly, "elapsed_sec": elapsed_sec, } - # All K samples share rollout_id so the loss reducer counts this - # trajectory once. - fanned = fan_out_sample_segments( - sample, - segments, - reward_result.reward, - state.tokenizer, - metadata=trajectory_metadata, - ) - if not fanned: - raise ValueError("fan-out produced no samples") + k = len(samples) + for i, s in enumerate(samples): + s.metadata = { + **(s.metadata or {}), + **trajectory_metadata, + "segment_idx": i, + "num_segments": k, + } + rlen = int(s.response_length or 0) + if rlen and s.tokens: + s.response = state.tokenizer.decode(s.tokens[-rlen:], skip_special_tokens=False) + else: + s.response = "" logger.info( "[coding_agent_rl] %s: reward=%.2f solved=%s applied=%s elapsed=%.1fs segments=%d", @@ -203,9 +223,9 @@ def _merge_samples( reward_result.is_solved, reward_result.applied_cleanly, elapsed_sec, - len(fanned), + k, ) - return fanned + return samples # --------------------------------------------------------------------------- @@ -254,11 +274,15 @@ async def generate(args, sample: Sample, sampling_params: dict[str, Any]): is_solved=bool(is_solved), applied_cleanly=bool(applied_cleanly), ) - segments = await state.adapter.finish_session(session_id) + samples = await state.adapter.finish_session( + session_id, + base_sample=sample, + reward=float(reward_result.reward), + ) return _merge_samples( sample=sample, state=state, - segments=segments, + samples=samples, reward_result=reward_result, elapsed_sec=time.time() - t0, instance_id=instance_id, diff --git a/slime/agent/adapters/__init__.py b/slime/agent/adapters/__init__.py index d5bb725ea9..6ab89963df 100644 --- a/slime/agent/adapters/__init__.py +++ b/slime/agent/adapters/__init__.py @@ -2,6 +2,5 @@ from slime.agent.adapters.anthropic import AnthropicAdapter from slime.agent.adapters.common import BaseAdapter -from slime.agent.adapters.openai import OpenAIAdapter -__all__ = ["AnthropicAdapter", "BaseAdapter", "OpenAIAdapter"] +__all__ = ["AnthropicAdapter", "BaseAdapter"] diff --git a/slime/agent/adapters/anthropic.py b/slime/agent/adapters/anthropic.py index 5b6d39fc5f..6a36abf5fa 100644 --- a/slime/agent/adapters/anthropic.py +++ b/slime/agent/adapters/anthropic.py @@ -2,21 +2,25 @@ The adapter exposes ``/v1/messages`` and ``/v1/messages/count_tokens``. It renders each Anthropic message history with the served model's chat template, -calls SGLang ``/generate`` with ``input_ids``, and folds the turn into a -per-session turn-node :class:`~slime.agent.trajectory_manager.TrajectoryTree`. - -The tree routes everything by text prefix, so Claude Code sub-agent and -compaction patterns split into independent leaves automatically -- no manual -``active_sub`` / ``wipe`` bookkeeping. Call ``finish_session()`` at trajectory -end to drain trainable ``TokenSegment`` objects. +calls SGLang ``/generate`` with ``input_ids``, and feeds the turn into a +shared :class:`~slime.agent.trajectory_manager.TrajectoryManager` keyed by +session id. ``finish_session(sid)`` drains a session's trajectory into a list +of :class:`~slime.utils.types.Sample`. + +The per-sid tree inside TrajectoryManager handles sub-agent and compaction +patterns automatically (any divergence in the prompt prefix forks into a new +leaf), so we no longer track ``main`` / ``active_sub`` chains here. """ from __future__ import annotations import asyncio +import dataclasses import json import logging +import re import secrets +from collections.abc import Callable from typing import Any from aiohttp import web @@ -27,40 +31,220 @@ TOKENIZER_KEY, TOOL_PARSER_KEY, BaseAdapter, - Session, - assemble_turns, call_sglang_generate, ok_response, - render_prompt, request_session_id, ) -from slime.agent.parsing import parse_model_output -from slime.agent.trajectory_manager import record_turn +from slime.agent.parsing import ParsedModelOutput, parse_model_output +from slime.agent.trajectory_manager import TrajectoryManager +from slime.utils.types import Sample logger = logging.getLogger(__name__) +@dataclasses.dataclass +class Session: + """Per-sid adapter state: sampling defaults, context budget, request lock. + + Trajectory state lives in ``AnthropicAdapter.manager`` (one shared + TrajectoryManager keyed by sid across all sessions). + """ + + sampling_defaults: dict = dataclasses.field(default_factory=dict) + max_context_tokens: int = 0 + lock: asyncio.Lock = dataclasses.field(default_factory=asyncio.Lock) + + class AnthropicAdapter(BaseAdapter): """Anthropic Messages-compatible HTTP adapter with session lifecycle helpers.""" - def __init__(self, *, tokenizer, sglang_url, tool_parser=None, reasoning_parser=None) -> None: + session_cls = Session + + def __init__( + self, + *, + tokenizer, + sglang_url, + tool_parser=None, + reasoning_parser=None, + tito_snapshot_min_loss_tokens: int | None = None, + max_turns_per_sid: int | None = None, + on_turn_appended: Callable[..., None] | None = None, + ) -> None: super().__init__( tokenizer=tokenizer, sglang_url=sglang_url, tool_parser=tool_parser, reasoning_parser=reasoning_parser, ) + # ONE manager shared across all sids; per-sid trees live inside. + self.manager = TrajectoryManager( + tokenizer=tokenizer, + tito_snapshot_min_loss_tokens=tito_snapshot_min_loss_tokens, + ) + # Optional debug hook invoked after each successful append_turn. + # Signature: (sid, prompt_messages, tools, response_message, + # prompt_ids, response_ids, finish_reason) -> None. + # Exceptions are swallowed; never block the SSE response. + self.on_turn_appended: Callable[..., None] | None = on_turn_appended + # Per-sid turn cap; None disables. When set, /v1/messages returns 429 + # once a sid has made this many turns. Prevents runaway agents from + # burning the whole budget. + self.max_turns_per_sid: int | None = max_turns_per_sid + self._sid_turn_count: dict[str, int] = {} self.app.router.add_post("/v1/messages", _handle_request) self.app.router.add_post("/v1/messages/count_tokens", _count_tokens) self.app.router.add_get("/healthz", _ok) self.app.router.add_get("/v1/models", _ok) + async def finish_session( + self, + sid: str, + *, + base_sample: Sample | None = None, + reward: float = 0.0, + extra_metadata: dict[str, Any] | None = None, + wait_timeout: float = 5.0, + ) -> list[Sample]: + """Drain a session's trajectory into Sample objects. + + Waits out in-flight requests for ``sid``, then linearises the + per-sid tree via ``TrajectoryManager.get_trajectory``. Idempotent -- + a second call for an already-popped sid returns ``[]``. + """ + await self.shutdown_session(sid, wait_timeout=wait_timeout) + # Drop the per-sid adapter Session; the trajectory itself is in + # manager._trees and will be popped by get_trajectory(drop=True). + self.store.pop(sid, None) + return self.manager.get_trajectory( + sid, + base_sample=base_sample, + reward=reward, + extra_metadata=extra_metadata, + ) + # ============================================================================= -# Translation (Anthropic wire <-> chat-template messages) -- unchanged +# Translation (Anthropic wire <-> chat-template messages) # ============================================================================= +# Claude Code CLI leaks ``x-anthropic-billing-header: ...cch=;`` as a text +# block at the top of the system prompt. The cch hash changes per request, so +# without stripping it the rendered system tokens differ every turn and the +# manager tree can't chain consecutive turns together. +_CLAUDE_CODE_BILLING_HEADER_RE = re.compile( + r"^\s*x-anthropic-billing-header:[^\n]*\n?", + re.IGNORECASE, +) + + +def _scrub_claude_code_billing_header_in_body(body_obj: dict) -> bool: + """Strip Claude Code's billing-header sidechannel from ``body['system']``. + + Handles both Anthropic shapes (``system: str`` and + ``system: list[{type:"text",text:"..."}]``). Mutates ``body_obj`` in + place; returns True iff anything changed. + """ + sysm = body_obj.get("system") + changed = False + if isinstance(sysm, str): + cleaned = _CLAUDE_CODE_BILLING_HEADER_RE.sub("", sysm) + if cleaned != sysm: + body_obj["system"] = cleaned if cleaned.strip() else "" + changed = True + elif isinstance(sysm, list): + new_blocks: list = [] + for block in sysm: + if not isinstance(block, dict) or block.get("type") != "text": + new_blocks.append(block) + continue + txt = block.get("text") or "" + cleaned = _CLAUDE_CODE_BILLING_HEADER_RE.sub("", txt) + if not cleaned.strip(): + # Whole block was the sidechannel — drop it. + changed = True + continue + if cleaned != txt: + new_block = dict(block) + new_block["text"] = cleaned + new_blocks.append(new_block) + changed = True + else: + new_blocks.append(block) + if changed: + body_obj["system"] = new_blocks + return changed + + +_MID_SYSTEM_WRAP_PREFIX = "\n" +_MID_SYSTEM_WRAP_SUFFIX = "\n\n" + + +def _fold_mid_list_system_into_user(body_obj: dict) -> bool: + """Fold non-leading ``role: system`` messages into a neighbouring user + message as a ```` text block. Mutates ``body_obj`` in + place; returns True iff any fold happened. + + Claude Code CLI >= 2.1.161 inserts ``{"role":"system","content":""}`` in the middle of ``messages``. Qwen3-style chat templates reject + any system message past index 0 with ``System message must be at the + beginning.`` This wrap mirrors the older claude-code (<= 2.1.143) + behaviour by attaching the wrapped reminder to the preceding user message + (or the next one if no prior user message exists). + """ + msgs = body_obj.get("messages") + if not isinstance(msgs, list) or not msgs: + return False + + system_idx = [i for i, m in enumerate(msgs) if isinstance(m, dict) and m.get("role") == "system" and i > 0] + if not system_idx: + return False + + def _promote_to_list(msg: dict) -> list: + c = msg.get("content") + if isinstance(c, list): + return c + msg["content"] = [{"type": "text", "text": c if isinstance(c, str) else ""}] + return msg["content"] + + def _wrap(text: str) -> dict: + return { + "type": "text", + "text": _MID_SYSTEM_WRAP_PREFIX + text + _MID_SYSTEM_WRAP_SUFFIX, + } + + changed = False + TOMBSTONE: dict = {"__folded__": True} + for i in system_idx: + sys_msg = msgs[i] + wrapped = _wrap(_flatten(sys_msg.get("content"))) + target = None + for j in range(i - 1, -1, -1): + cand = msgs[j] + if isinstance(cand, dict) and cand.get("role") == "user": + target = cand + _promote_to_list(target).append(wrapped) + break + if target is None: + for j in range(i + 1, len(msgs)): + cand = msgs[j] + if isinstance(cand, dict) and cand.get("role") == "user": + target = cand + _promote_to_list(target).insert(0, wrapped) + break + if target is None: + msgs[i] = {"role": "user", "content": [wrapped]} + changed = True + continue + msgs[i] = TOMBSTONE + changed = True + + if changed: + body_obj["messages"] = [m for m in msgs if m is not TOMBSTONE] + return changed + + def _flatten(c: Any) -> str: """Recursive Anthropic content flattener: text/tool_result(content) joined by newline, images replaced with a placeholder.""" @@ -85,6 +269,48 @@ def _flatten(c: Any) -> str: return "\n".join(p for p in parts if p) +# Marker string Claude Code embeds in the system prompt of its per-session +# title-generation request (a meta request that asks the LLM to produce a +# short conversation title). Title-gen requests should NOT enter the RL +# trajectory — they aren't agent work. See spec +# docs/superpowers/specs/2026-06-04-skip-cc-title-gen-from-trajectory-design.md. +_CC_TITLE_GEN_MARKER = "Generate a concise, sentence-case title" + + +def _is_cc_title_generation_request( + translated: list[dict], + tools_schema: list[dict] | None, +) -> bool: + """Return True iff this is a Claude Code per-session title-generation request. + + Detection is AND-conjunction: + (1) ``tools_schema`` is falsy (cc sends tools=[]; converter returns None). + (2) one of the leading ``role=system`` messages' content contains + ``_CC_TITLE_GEN_MARKER``. + + Scanning stops at the first non-system message — title-gen system blocks + always sit at the head of the request. + """ + if tools_schema: + return False + for msg in translated: + if msg.get("role") != "system": + break + content = msg.get("content") + if isinstance(content, str): + if _CC_TITLE_GEN_MARKER in content: + return True + elif isinstance(content, list): + for block in content: + if ( + isinstance(block, dict) + and isinstance(block.get("text"), str) + and _CC_TITLE_GEN_MARKER in block["text"] + ): + return True + return False + + def _translate_anthropic(msgs: list[dict], system: Any) -> list[dict]: """Anthropic messages + system -> chat-template messages. Pure function.""" translated: list[dict] = [] @@ -114,7 +340,28 @@ def _translate_anthropic(msgs: list[dict], system: Any) -> list[dict]: elif b.get("type") == "thinking": thinkings.append(b.get("thinking", "")) elif b.get("type") == "tool_use": - tcs.append({"function": {"name": b.get("name", "tool"), "arguments": b.get("input") or {}}}) + # Match the canonical shape produced by + # _build_blocks_and_response_message for sampled leaves so + # that node_match_key (json.dumps sort_keys) hashes a + # replayed assistant identically to its leaf. We drop the + # wire-only "id" — see the matching note on the leaf side. + # NB: arguments stays a dict here (NOT a JSON string). + # The translated list is also fed to the chat template + # via apply_chat_template; Qwen3's template calls + # `arguments | items` which requires a mapping. A JSON + # string would raise "Can only get item pairs from a + # mapping." mid-render. node_match_key's json.dumps + # sort_keys=True recursively sorts dict keys so two + # equivalent dicts still hash identically. + tcs.append( + { + "type": "function", + "function": { + "name": b.get("name", "tool"), + "arguments": b.get("input") or {}, + }, + } + ) mo: dict[str, Any] = {"role": "assistant", "content": "".join(texts)} if thinkings: mo["reasoning_content"] = "".join(thinkings) @@ -148,65 +395,103 @@ def _anthropic_tools_to_chat_tools(anth_tools: list[dict] | None) -> list[dict] # ============================================================================= -# Reply building (raw output text -> Anthropic content blocks) -- unchanged shape +# Local chat-template render helper. +# +# The Anthropic adapter renders directly from a message list -- no chain +# bookkeeping needed because TrajectoryManager is the routing authority. # ============================================================================= -async def _generate(prompt_ids: list[int], s: Session, body: dict, app, *, session_id: str | None = None): - """Call sglang and return a GenResult (output ids/logprobs/text).""" - return await call_sglang_generate( - prompt_ids, - s, - body, - app, - max_token_keys=("max_tokens",), - stop_keys=("stop_sequences",), - log_prefix="anthropic_adapter", - logger=logger, - session_id=session_id, +def _render_token_ids( + messages: list[dict], + tokenizer, + *, + tools: list[dict] | None, + add_generation_prompt: bool = True, +) -> list[int]: + enc = tokenizer.apply_chat_template( + messages, + tools=tools, + tokenize=True, + add_generation_prompt=add_generation_prompt, ) + ids = enc["input_ids"] if hasattr(enc, "__getitem__") and "input_ids" in enc else enc + return list(ids) -def _build_reply(output_text: str, finish: str, tools_schema: list[dict] | None, app) -> tuple[list[dict], str]: - """Turn the model's raw output text into the reply we send back to claude-code. - - 1. parse decoded text -> (thinking, visible, tool_uses) via sglang parsers - 2. pack into Anthropic content blocks - 3. derive stop_reason: 'tool_use' | 'max_tokens' | 'end_turn' +# ============================================================================= +# Reply building: parsed output -> Anthropic blocks + OpenAI-shape response_message +# ============================================================================= - Returns (blocks, stop_reason). - """ - parsed = parse_model_output( - output_text or "", - tools_schema=tools_schema, - tool_parser_name=app[TOOL_PARSER_KEY], - reasoning_parser_name=app[REASONING_PARSER_KEY], - ) - blocks = _anthropic_blocks(parsed.reasoning, parsed.text, parsed.tool_uses) - return blocks, _stop_reason(parsed.tool_uses, finish) +def _build_blocks_and_response_message( + parsed: ParsedModelOutput, + finish: str, +) -> tuple[list[dict], str, dict[str, Any]]: + """Pack parsed model output into: + - Anthropic content blocks (sent over the wire), + - stop_reason, + - response_message (OpenAI shape) for TrajectoryManager.append_turn. -def _anthropic_blocks(thinking: str, visible: str, tool_uses: list[dict]) -> list[dict]: - """Pack parsed model output into Anthropic content blocks.""" + The tool_calls inside response_message use canonical JSON args (sorted + keys, str-encoded) so the node_match_key the manager computes for this + assistant turn matches the same turn replayed as history on the next + /v1/messages request. + """ blocks: list[dict] = [] - if thinking: - blocks.append({"type": "thinking", "thinking": thinking}) - if visible: - blocks.append({"type": "text", "text": visible}) - for tu in tool_uses: + if parsed.reasoning: + blocks.append({"type": "thinking", "thinking": parsed.reasoning}) + if parsed.text: + blocks.append({"type": "text", "text": parsed.text}) + + response_tcs: list[dict] = [] + for tu in parsed.tool_uses: tu_id = f"toolu_{secrets.token_hex(8)}" blocks.append({"type": "tool_use", "id": tu_id, "name": tu["name"], "input": tu["input"]}) + # NB: do NOT include tu_id here. The id is wire-only (clients use it + # to correlate tool_result blocks). When cc echoes this assistant on + # the next /v1/messages, it sends the original id; slime regenerates + # a fresh id each call. Including the id in response_message would + # make node_match_key differ between leaf and echo, breaking the + # leaf-vs-replay merge — see trajectory_manager DFS Step 1. + # + # arguments stays a dict to mirror _translate_anthropic. The + # trajectory_manager node_match_key uses json.dumps(sort_keys=True) + # which is invariant to dict key order, so two equivalent dicts + # hash identically. + response_tcs.append( + { + "type": "function", + "function": { + "name": tu["name"], + "arguments": tu.get("input") or {}, + }, + } + ) + if not blocks: blocks.append({"type": "text", "text": ""}) - return blocks + if parsed.tool_uses: + stop_reason = "tool_use" + elif finish == "length": + stop_reason = "max_tokens" + else: + stop_reason = "end_turn" + + response_message: dict[str, Any] = {"role": "assistant", "content": parsed.text or ""} + if parsed.reasoning: + response_message["reasoning_content"] = parsed.reasoning + if response_tcs: + response_message["tool_calls"] = response_tcs + + return blocks, stop_reason, response_message -def _stop_reason(tool_uses: list[dict], finish: str) -> str: + +def _finish_reason_for_manager(finish: str, tool_uses: list[dict]) -> str: if tool_uses: - return "tool_use" - if finish == "length": - return "max_tokens" - return "end_turn" + return "tool_calls" + return finish or "stop" # ============================================================================= @@ -224,6 +509,33 @@ async def _handle_request(request: web.Request) -> web.StreamResponse: adapter = request.app[ADAPTER_KEY] if sid in adapter.closed: # session drained; refuse stragglers return web.Response(status=503, text="session closed") + + # Per-sid turn cap (HTTP 429). When the adapter was constructed with a + # ``max_turns_per_sid`` ceiling, refuse further /v1/messages calls past + # that count so a runaway agent in a sandbox exits cleanly instead of + # burning the whole token budget. ``None`` (default) disables. + cap = adapter.max_turns_per_sid + if cap is not None: + prior = adapter._sid_turn_count.get(sid, 0) + if prior >= cap: + return web.json_response( + { + "error": { + "type": "rate_limit_error", + "message": (f"adapter: sid {sid!r} exceeded max_turns_per_sid={cap}; killing run"), + } + }, + status=429, + ) + adapter._sid_turn_count[sid] = prior + 1 + + # Strip Claude Code's per-request billing-header sidechannel BEFORE the + # adapter renders prompt_ids. Also fold mid-list ``role: system`` messages + # into a neighbouring user message so Qwen3 chat templates accept them. + # Both are no-ops when the relevant patterns aren't present. + _scrub_claude_code_billing_header_in_body(body) + _fold_mid_list_system_into_user(body) + app = request.app tok = app[TOKENIZER_KEY] s = adapter.store.setdefault(sid, Session()) @@ -233,21 +545,83 @@ async def _handle_request(request: web.Request) -> web.StreamResponse: async with s.lock: # same sid -> serialized translated = _translate_anthropic(body.get("messages") or [], body.get("system")) tools_schema = _anthropic_tools_to_chat_tools(body.get("tools")) - full_prompt_ids = render_prompt(s.traj, translated, tok, tools_schema) - gen = await _generate(full_prompt_ids, s, body, app, session_id=sid) - turns, pending_key = assemble_turns(s.traj, translated, tok, tools_schema, gen, full_prompt_ids) - node = record_turn(s.traj.tree, turns) - if node is not None: - s.traj.resp_truth[pending_key] = ( - list(gen.output_ids), - list(gen.output_log_probs), - gen.output_text, + prompt_ids = _render_token_ids(translated, tok, tools=tools_schema, add_generation_prompt=True) + + turn = await call_sglang_generate( + prompt_ids, + s, + body, + app, + max_token_keys=("max_tokens",), + stop_keys=("stop_sequences",), + log_prefix="anthropic_adapter", + logger=logger, + session_id=sid, + ) + + raw_output = tok.decode(turn.output_ids, skip_special_tokens=False) if turn.output_ids else "" + parsed = parse_model_output( + raw_output, + tools_schema=tools_schema, + tool_parser_name=app[TOOL_PARSER_KEY], + reasoning_parser_name=app[REASONING_PARSER_KEY], + ) + blocks, stop_reason, response_message = _build_blocks_and_response_message(parsed, turn.finish_reason) + + output_ids = list(turn.output_ids) + finish_reason = _finish_reason_for_manager(turn.finish_reason, parsed.tool_uses) + + if _is_cc_title_generation_request(translated, tools_schema): + # Claude Code meta request (per-session title generation). + # Skip the trajectory so it doesn't pollute the tree / become + # an RL sample. The on_turn_appended hook below still fires, + # so per-turn dumps (request, sse, openai.json) keep landing + # on disk for debugging. See spec + # docs/superpowers/specs/2026-06-04-skip-cc-title-gen-from-trajectory-design.md. + logger.info( + "skipping append_turn for cc title-generation request (sid=%s)", + sid, ) - blocks, stop = _build_reply(gen.output_text, gen.finish_reason, tools_schema, app) - in_tok, out_tok = len(full_prompt_ids), len(gen.output_ids) + else: + try: + adapter.manager.append_turn( + sid, + prompt_messages=translated, + tools=tools_schema, + prompt_ids=prompt_ids, + response_ids=output_ids, + response_logprobs=( + list(turn.output_log_probs) + if turn.output_log_probs and len(turn.output_log_probs) == len(turn.output_ids) + else None + ), + response_message=response_message, + finish_reason=finish_reason, + metadata={"sid": sid}, + ) + except Exception: + logger.exception("append_turn(sid=%s) failed", sid) + + hook = adapter.on_turn_appended + if hook is not None: + try: + hook( + sid, + translated, + tools_schema, + response_message, + prompt_ids, + output_ids, + finish_reason, + ) + except Exception: + logger.exception("on_turn_appended hook failed (sid=%s)", sid) + + in_tok, out_tok = len(prompt_ids), len(turn.output_ids) + if body.get("stream") is True or "text/event-stream" in request.headers.get("Accept", ""): - return await _stream_response(request, blocks, stop, in_tok, out_tok) - return web.json_response(_message_response(body, blocks, stop, in_tok, out_tok)) + return await _stream_response(request, blocks, stop_reason, in_tok, out_tok) + return web.json_response(_message_response(body, blocks, stop_reason, in_tok, out_tok)) finally: adapter.inflight.get(sid, set()).discard(task) @@ -279,7 +653,6 @@ async def _stream_response(request, blocks, stop_reason, in_tok, out_tok) -> web ) await out.prepare(request) - # message_start ms_data = { "type": "message_start", "message": { diff --git a/slime/agent/adapters/common.py b/slime/agent/adapters/common.py index c3705cf942..f34daf4bad 100644 --- a/slime/agent/adapters/common.py +++ b/slime/agent/adapters/common.py @@ -1,28 +1,9 @@ -"""Shared adapter primitives for token-capturing agent rollouts. - -Each HTTP request carries the full conversation history. We render it with the -served model's chat template, call SGLang ``/generate`` with ``input_ids``, and -fold the turn into a per-session :class:`~slime.agent.trajectory_manager.TrajectoryTree`. - -The tree does all routing (sub-agent / compaction / history-rewrite branch -automatically by text prefix), so there is no manual new/append/wipe logic. Two -things make this faithful under TITO (text-in-token-out) drift: - -* the **current turn's** incremental prompt segment is taken by *segment render - diff* (``render(through this prompt) - render(through prev response)``), which - is immune to drift because it compares two re-renders rather than a cache vs a - re-render (design §5.1); and -* **historical turns'** response tokens come from a per-session *truth cache* - keyed by the cumulative prompt that produced them, so they equal the tree - node tokens exactly (no re-tokenization). -""" +"""Shared adapter primitives for token-capturing agent rollouts.""" from __future__ import annotations import asyncio import dataclasses -import hashlib -import json import logging import uuid from collections.abc import Callable @@ -31,7 +12,6 @@ import aiohttp from aiohttp import web -from slime.agent.trajectory_manager import PromptSeg, RespSeg, TokenSegment, TrajectoryTree, export_token_segments ADAPTER_KEY = web.AppKey("adapter", object) TOKENIZER_KEY = web.AppKey("tokenizer", object) @@ -40,61 +20,27 @@ REASONING_PARSER_KEY = web.AppKey("reasoning_parser", object) -@dataclasses.dataclass -class GenResult: - """One assistant generation from the rollout engine (sglang ``/generate``). +@dataclasses.dataclass(frozen=True) +class TurnRecord: + """Exact token snapshot for one assistant generation, returned by + :func:`call_sglang_generate`. - ``output_text`` is ``decode(output_ids)`` cached once so reply builders and - the trajectory's response-segment text share a single detokenization. + ``prompt_ids`` is the full tokenized prompt sent to the generator for that + turn. ``output_ids`` is the raw generated output, and + ``output_log_probs`` is aligned with it when the rollout engine returns + per-token log probabilities. """ + prompt_ids: list[int] output_ids: list[int] - output_log_probs: list[float] finish_reason: str - output_text: str = "" - - -@dataclasses.dataclass -class SessionTrajectory: - """Per-session trajectory state: the turn-node tree plus the truth cache. - - ``resp_truth`` maps the cumulative prompt token tuple that *preceded* a - response to ``(output_ids, output_log_probs, output_text)``. The same turn, - seen as history in a later request, renders to the same cumulative prompt and - so retrieves its exact sampled tokens (a Case 2 append, never a false drift). - Keying on rendered prompt ids (a deterministic function of the messages), - rather than on echoed wire text, avoids the parse->serialize round-trip drift - that would otherwise make a text key miss. - - ``render_memo`` caches ``apply_chat_template`` results within a session so a - growing trajectory renders each distinct message-prefix once (O(n) renders - over the whole trajectory instead of O(n^2)). - """ - - tree: TrajectoryTree = dataclasses.field(default_factory=TrajectoryTree) - resp_truth: dict[tuple[int, ...], tuple[list[int], list[float], str]] = dataclasses.field(default_factory=dict) - render_memo: dict[tuple[str, bool], list[int]] = dataclasses.field(default_factory=dict) - - -@dataclasses.dataclass -class Session: - """Per-session adapter state shared by the OpenAI and Anthropic adapters. - - Holds the trajectory (turn-node tree + truth cache), the per-session sampling - defaults / context budget injected at ``open_session``, and a lock that - serializes concurrent requests carrying the same session id. - """ - - traj: SessionTrajectory = dataclasses.field(default_factory=SessionTrajectory) - sampling_defaults: dict = dataclasses.field(default_factory=dict) - max_context_tokens: int = 0 - lock: asyncio.Lock = dataclasses.field(default_factory=asyncio.Lock) + output_log_probs: list[float] = dataclasses.field(default_factory=list) class BaseAdapter: """Base HTTP adapter with per-instance session lifecycle state.""" - session_cls: type = Session + session_cls: type def __init__(self, *, tokenizer, sglang_url, tool_parser=None, reasoning_parser=None) -> None: self.store: dict[str, Any] = {} @@ -125,173 +71,8 @@ def open_session( async def shutdown_session(self, sid: str, *, wait_timeout: float = 5.0) -> None: await shutdown_session_tasks(sid, self.closed, self.inflight, wait_timeout=wait_timeout) - async def finish_session(self, sid: str, *, wait_timeout: float = 5.0) -> list[TokenSegment]: - """Drain a session: wait out in-flight requests, then export trainable - ``TokenSegment``s from its trajectory tree. Idempotent -- a second call - for an already-popped sid returns ``[]``.""" - await self.shutdown_session(sid, wait_timeout=wait_timeout) - s = self.store.pop(sid, None) - if s is None: - return [] - return export_token_segments(s.traj.tree) - - -def strip_cache_control(obj: Any) -> Any: - if isinstance(obj, dict): - return {k: strip_cache_control(v) for k, v in obj.items() if k != "cache_control"} - if isinstance(obj, list): - return [strip_cache_control(x) for x in obj] - return obj - - -def stable_hash(obj: Any) -> str: - payload = json.dumps(strip_cache_control(obj), sort_keys=True, ensure_ascii=False, default=str).encode("utf-8") - return hashlib.sha1(payload).hexdigest()[:12] - - -def json_arguments(value: Any) -> str: - if value is None: - return "{}" - if isinstance(value, str): - return value - return json.dumps(value, ensure_ascii=False) - - -def _extract_ids(enc: Any) -> list[int]: - ids = enc["input_ids"] if hasattr(enc, "__getitem__") and "input_ids" in enc else enc - return list(ids) - - -def render_token_ids( - messages: list[dict], tokenizer, *, tools: list[dict] | None = None, add_generation_prompt: bool = True -) -> list[int]: - """Render a chat-template message list to token ids.""" - enc = tokenizer.apply_chat_template( - messages, - tools=tools, - tokenize=True, - add_generation_prompt=add_generation_prompt, - ) - return _extract_ids(enc) - - -# ============================================================================= -# Turn assembly: chat messages -> strictly-alternating (PromptSeg, RespSeg) turns -# ============================================================================= - - -def split_turns(chat_messages: list[dict]) -> list[tuple[list[dict], dict | None]]: - """Split a chat-template message list into ``(prompt_msgs, assistant_msg)`` - turns at every ``assistant`` boundary. - - ``prompt_msgs`` is the run of non-assistant messages since the previous - assistant. A trailing ``(prompt_msgs, None)`` turn is always appended: it is - the current turn whose response is about to be generated. - """ - specs: list[tuple[list[dict], dict | None]] = [] - buf: list[dict] = [] - for m in chat_messages: - if isinstance(m, dict) and m.get("role") == "assistant": - specs.append((buf, m)) - buf = [] - else: - buf.append(m) - specs.append((buf, None)) - return specs - - -def _render_memo( - traj: SessionTrajectory | None, - messages: list[dict], - tokenizer, - tools: list[dict] | None, - add_generation_prompt: bool, -) -> list[int]: - """Render with a per-session memo so each distinct message prefix is rendered - once across a growing trajectory (turns 1..n cost O(n) renders total, not - O(n^2)).""" - if traj is None: - return render_token_ids(messages, tokenizer, tools=tools, add_generation_prompt=add_generation_prompt) - key = (stable_hash([messages, tools]), add_generation_prompt) - cached = traj.render_memo.get(key) - if cached is None: - cached = render_token_ids(messages, tokenizer, tools=tools, add_generation_prompt=add_generation_prompt) - traj.render_memo[key] = cached - return list(cached) - - -def render_prompt(traj: SessionTrajectory, messages: list[dict], tokenizer, tools: list[dict] | None) -> list[int]: - """Render the full prompt (``add_generation_prompt=True``) to send to sglang, - memoized on the session so a replayed prefix is not re-rendered next turn.""" - return _render_memo(traj, messages, tokenizer, tools, True) - - -def assemble_turns( - traj: SessionTrajectory, - chat_messages: list[dict], - tokenizer, - tools: list[dict] | None, - gen: GenResult, - full_prompt_ids: list[int], -) -> tuple[list[tuple[PromptSeg, RespSeg]], tuple[int, ...]]: - """Build ``record_turn``'s strictly-alternating ``turns`` for this request. - - Every turn's incremental prompt segment comes from a *segment render diff* - (design §5.1): ``render(through this prompt) - render(through previous - response)``, comparing two re-renders so it is immune to TITO drift. The - previous render uses ``add_generation_prompt=False`` (it ends at a response), - so the diff is a clean length-based suffix. - - History response tokens come from the truth cache keyed by the cumulative - prompt that produced them (exact sampled ids; a Case 2 append, never a false - drift). On a cache miss they fall back to re-tokenization, which the tree - then handles as Case 3/4 (replace + mask). The current turn's response comes - from ``gen`` and reuses ``full_prompt_ids`` (the prompt already sent to - sglang) as its prompt segment. - - Does not write the cache -- the caller writes it only after ``record_turn`` - succeeds (no dangling cache on a sglang error). Returns ``(turns, - pending_key)`` where ``pending_key`` is the cumulative-prompt key under which - the current turn's response should be cached. - """ - specs = split_turns(chat_messages) - turns: list[tuple[PromptSeg, RespSeg]] = [] - prev_cum: list[dict] = [] - prev_ids: list[int] = [] - pending_key: tuple[int, ...] = tuple(full_prompt_ids) - - for prompt_msgs, assistant_msg in specs: - cum_prompt = prev_cum + prompt_msgs - if assistant_msg is None: - # Current turn: reuse the exact prompt already sent to sglang. - cur_ids = list(full_prompt_ids) - else: - cur_ids = _render_memo(traj, cum_prompt, tokenizer, tools, True) - delta = cur_ids[len(prev_ids) :] - p_text = tokenizer.decode(delta, skip_special_tokens=False) if delta else "" - p_seg = PromptSeg(p_text, list(delta)) - - if assistant_msg is None: - r_seg = RespSeg(gen.output_text, list(gen.output_ids), list(gen.output_log_probs)) - pending_key = tuple(cur_ids) - turns.append((p_seg, r_seg)) - break - - cached = traj.resp_truth.get(tuple(cur_ids)) - if cached is not None: - r_ids, r_lp, r_text = cached - r_seg = RespSeg(r_text, list(r_ids), list(r_lp)) - else: - full = _render_memo(traj, cum_prompt + [assistant_msg], tokenizer, tools, False) - r_ids = full[len(cur_ids) :] - r_text = tokenizer.decode(r_ids, skip_special_tokens=False) if r_ids else "" - r_seg = RespSeg(r_text, list(r_ids), [0.0] * len(r_ids)) - turns.append((p_seg, r_seg)) - - prev_cum = cum_prompt + [assistant_msg] - prev_ids = _render_memo(traj, prev_cum, tokenizer, tools, False) - - return turns, pending_key + async def finish_session(self, sid: str, *, wait_timeout: float = 5.0) -> list: + raise NotImplementedError def request_session_id( @@ -373,7 +154,7 @@ async def call_sglang_generate( log_prefix: str, logger: logging.Logger, session_id: str | None = None, -) -> GenResult: +) -> TurnRecord: sp = _sampling_params(session, body, max_token_keys=max_token_keys, stop_keys=stop_keys) if session.max_context_tokens > 0: @@ -385,7 +166,7 @@ async def call_sglang_generate( len(prompt_ids), session.max_context_tokens, ) - return GenResult(output_ids=[], output_log_probs=[], finish_reason="length", output_text="") + return TurnRecord(prompt_ids=list(prompt_ids), output_ids=[], finish_reason="length") sp["max_new_tokens"] = min(int(sp.get("max_new_tokens", remaining_context)), remaining_context) sglang_url = app[SGLANG_URL_KEY] @@ -420,13 +201,11 @@ async def call_sglang_generate( pass raise - tok = app[TOKENIZER_KEY] if TOKENIZER_KEY in app else None - output_text = tok.decode(output_ids, skip_special_tokens=False) if (tok is not None and output_ids) else "" - return GenResult( + return TurnRecord( + prompt_ids=list(prompt_ids), output_ids=output_ids, - output_log_probs=output_log_probs, finish_reason=finish, - output_text=output_text, + output_log_probs=output_log_probs, ) diff --git a/slime/agent/trajectory_manager.py b/slime/agent/trajectory_manager.py index 17333601e2..8a5937176a 100644 --- a/slime/agent/trajectory_manager.py +++ b/slime/agent/trajectory_manager.py @@ -1,30 +1,58 @@ -"""Turn-node trajectory manager for agent rollouts. - -A session's trajectory is modeled as a **text-prefix tree**: each ``Node`` is one -turn (incremental prompt segment + response segment). Text-prefix matching is the -primary signal; token-id comparison is secondary and only handles TITO -(text-in-token-out) drift after the text prefix already matches. ``export`` walks -the tree once and yields full-length ``(tokens, masks, logprobs)`` triples, one -sample per leaf. - -The core algorithm (``PromptSeg``/``RespSeg``/``Node``/``TrajectoryTree``/ -``MatchResult`` and ``match_prefix``/``attach_turn``/``record_turn``/``export``) -is the standalone turn-node design; see -``0601-Trajectory-manager/02-turn-node/02-turn-node-design.md``. - -The slime glue below the core (``TokenSegment``/``export_token_segments``/ -``fan_out_sample_segments``/``write_segment_to_sample``) is the only intentional -deviation from the design's zero-dependency constraint: it imports ``Sample`` and -adapts the full-length export into slime's prompt/response training segments. +"""Per-role chunk-merging trajectory tree manager (C-plan: token-faithful). + +Design (Plan C, 2026-06-03): + +* The tree is a router only. DFS merge keys on ``(role, node_match_key)`` + alone — no prompt_ids prefix check. Same conversation prefix in + ``messages`` space always lands on the same path, regardless of any + chat_template re-tokenization drift across turns. + +* Each assistant leaf stores the THIS-TURN sglang snapshot: + ``turn_prompt_ids`` / ``turn_response_ids`` / ``turn_response_logprobs`` + / ``turn_finish_reason`` / ``turn_index``. Non-assistant nodes carry no + token attribution at all. + +* ``get_trajectory`` linearizes each leaf turn-by-turn using LCP-aligned + drop-and-replace: the cumulative tokens emitted so far are clamped to + the longest common prefix with the next turn's prompt; any prior tokens + past that LCP (the TITO drift suffix, including the previous turn's + response if it lands in the drift region) are DROPPED along with their + logprobs, then ``prompt[LCP:]`` is appended as loss_mask=0 (chat + template's authoritative re-rendering wins), then the current turn's + ``response`` is appended as loss_mask=1 with real logprobs. + +* Trade-off: previous-turn response tokens that fall inside the drift + region lose their training signal. In exchange, the final tokens + sequence matches what the live model actually conditioned on for every + later turn — logprobs stay coherent, no duplicated-content forks, no + reliance on chat_template being position-invariant. + +* Snapshot rescue (opt-in via ``tito_snapshot_min_loss_tokens``): when a + drift would drop >= N loss_mask=1 tokens, emit an extra "snapshot" + Sample alongside the main leaf. Snapshot tokens = cumulative pre-drop; + snapshot loss_mask is COMPLEMENTARY — 1 only at positions that the + main leaf is about to drop, 0 elsewhere. Snapshot reward = main-leaf + share; snapshot rollout_id = main-leaf rollout_id. Snapshot ∪ main on + loss_mask=1 tokens never overlap and their union equals the virtual + no-drift trajectory. The snapshotted drift is NOT counted in the main + sample's ``tito_dropped_*`` (it wasn't truly lost). + +* On drift, ``Sample.metadata`` records: + ``tito_dropped_tokens`` — total tokens dropped (NOT including + drifts that produced a snapshot) + ``tito_dropped_turns`` — number of turns that triggered a drop + ``tito_snapshots_emitted`` — set on main leaf when >=1 snapshot + sibling was emitted for the same leaf + ``tito_snapshot`` — True on a snapshot Sample + ``tito_snapshot_at_turn`` — turn index whose drift triggered it + ``tito_snapshot_loss_tokens`` — count of loss_mask=1 tokens in snapshot """ from __future__ import annotations -import copy -import dataclasses +import json import logging -import time -import uuid +from collections.abc import Iterator from dataclasses import dataclass, field from typing import Any @@ -33,101 +61,127 @@ logger = logging.getLogger(__name__) -# ============================================================================= -# Core turn-node tree (standalone design; zero training-side dependency) -# ============================================================================= +# =========================================================================== +# Node +# =========================================================================== -@dataclass -class PromptSeg: - text: str - tokens: list[int] - - -@dataclass -class RespSeg: - text: str - tokens: list[int] - logprobs: list[float] - - -@dataclass class Node: - id: str - prompt_delta_text: str - prompt_delta_tokens: list[int] - resp_text: str - resp_tokens: list[int] - resp_logprobs: list[float] - parent: Node | None - children: list[Node] = field(default_factory=list) - loss_mask: bool = False - created_at: float = 0.0 - updated_at: float = 0.0 - replaced_count: int = 0 - warning: str | None = None - kind_is_root: bool = False + """One node in the trajectory tree. + + Routing fields (every node): + role, messages, parent, children, metadata + + Per-turn snapshot fields (assistant leaves only — None on non-assistant + and on internal assistant nodes that aren't a turn's own leaf): + turn_prompt_ids: list[int] sglang prompt as fed to /generate + turn_response_ids: list[int] sglang output ids + turn_response_logprobs: list[float] + turn_finish_reason: str | None + turn_index: int 1-based, monotonic per session + """ - @staticmethod - def new_root() -> Node: - now = time.time() - return Node( - id=uuid.uuid4().hex[:12], - prompt_delta_text="", - prompt_delta_tokens=[], - resp_text="", - resp_tokens=[], - resp_logprobs=[], - parent=None, - children=[], - created_at=now, - updated_at=now, - kind_is_root=True, - ) + __slots__ = ( + # routing + "role", + "messages", + "metadata", + "parent", + "children", + # per-turn snapshot (assistant leaves) + "turn_prompt_ids", + "turn_response_ids", + "turn_response_logprobs", + "turn_finish_reason", + "turn_index", + ) - @staticmethod - def new_turn(prompt: PromptSeg, resp: RespSeg, parent: Node) -> Node: - now = time.time() - node = Node( - id=uuid.uuid4().hex[:12], - prompt_delta_text=prompt.text, - prompt_delta_tokens=list(prompt.tokens), - resp_text=resp.text, - resp_tokens=list(resp.tokens), - resp_logprobs=list(resp.logprobs), - parent=parent, - children=[], - created_at=now, - updated_at=now, - ) - parent.children.append(node) - return node + def __init__( + self, + *, + role: str | None = None, + messages: list[dict[str, Any]] | None = None, + metadata: dict[str, Any] | None = None, + parent: Node | None = None, + ) -> None: + self.role = role + self.messages = list(messages or []) + self.metadata = dict(metadata or {}) + self.parent: Node | None = parent + self.children: list[Node] = [] + # per-turn snapshot + self.turn_prompt_ids: list[int] | None = None + self.turn_response_ids: list[int] | None = None + self.turn_response_logprobs: list[float] | None = None + self.turn_finish_reason: str | None = None + self.turn_index: int | None = None @property - def node_text(self) -> str: - return self.prompt_delta_text + self.resp_text + def is_root(self) -> bool: + return self.parent is None + + def add_child(self, child: Node) -> Node: + child.parent = self + self.children.append(child) + return child + + def path_from_root(self) -> list[Node]: + """Ordered list of nodes from the first non-root ancestor down to self.""" + chain: list[Node] = [] + cur: Node | None = self + while cur is not None and not cur.is_root: + chain.append(cur) + cur = cur.parent + chain.reverse() + return chain + + def leaves(self) -> Iterator[Node]: + if not self.children: + yield self + return + for c in self.children: + yield from c.leaves() -@dataclass -class TrajectoryTree: - root: Node = field(default_factory=Node.new_root) +# =========================================================================== +# node_match_key + role-grouping helpers +# =========================================================================== -@dataclass -class MatchResult: - case: str - anchor: Node | None = None - text_matched_len: int = 0 - drift_nodes: list[Node] = field(default_factory=list) - lca: Node | None = None - fork_node: Node | None = None - residual_count: int = 0 +def node_match_key(messages: list[dict[str, Any]]) -> str: + """Identity key for a node's message list. + + json.dumps(sort_keys=True) sorts dict-internal keys recursively; list + element order is preserved (which is what we want: message order and + tool_calls order are both semantically significant). + """ + return json.dumps(messages, sort_keys=True, ensure_ascii=False) -# ---- helpers ---- +@dataclass +class _PromptGroup: + role: str + messages: list[dict[str, Any]] = field(default_factory=list) + + +def _group_messages_by_role( + messages: list[dict[str, Any]], +) -> list[_PromptGroup]: + groups: list[_PromptGroup] = [] + for m in messages: + role = m.get("role") + if not isinstance(role, str): + logger.warning("skipping message without string role: %r", m) + continue + if groups and groups[-1].role == role: + groups[-1].messages.append(m) + else: + groups.append(_PromptGroup(role=role, messages=[m])) + return groups -def _text_lcp(a: str, b: str) -> int: +def _lcp_len(a: list[int], b: list[int]) -> int: + """Length of the longest common prefix between two int lists.""" n = min(len(a), len(b)) i = 0 while i < n and a[i] == b[i]: @@ -135,405 +189,314 @@ def _text_lcp(a: str, b: str) -> int: return i -def _chain_nodes(node: Node) -> list[Node]: - """root->node path, non-root nodes, ordered root to leaf.""" - out = [] - cur = node - while cur is not None and not cur.kind_is_root: - out.append(cur) - cur = cur.parent - out.reverse() - return out - +# =========================================================================== +# TrajectoryManager +# =========================================================================== -def _chain_text(nodes: list[Node]) -> str: - return "".join(n.node_text for n in nodes) +class TrajectoryManager: + """Per-sid trajectory tree manager. -def _collect_leaves(root: Node) -> list[Node]: - """Pre-order collect of all leaves (excluding root; empty if root childless).""" - leaves = [] - - def walk(n: Node): - if not n.children: - if not n.kind_is_root: - leaves.append(n) - return - for c in n.children: - walk(c) - - walk(root) - return leaves - - -def _log_replace(case, reason, node, before_text, after_text, before_tokens, after_tokens): - logger.warning( - "trajectory replace | case=%s reason=%s node=%s\n" " text : %r -> %r\n" " token: %s -> %s", - case, - reason, - node.id, - before_text, - after_text, - before_tokens, - after_tokens, - ) - + See module docstring for the C-plan invariants. Each ``append_turn`` + mounts >=0 prompt nodes (under the deepest matching ancestor) + exactly + 1 assistant leaf carrying that turn's sglang snapshot. + """ -# ---- match_prefix ---- - - -def match_prefix(tree: TrajectoryTree, turns) -> MatchResult: - incoming_text = "".join(p.text + r.text for (p, r) in turns) - - # For each root->leaf chain take the longest text prefix; pick the longest, - # ties broken by the latest created_at leaf. - best_leaf = None - best_len = -1 - for leaf in _collect_leaves(tree.root): - nodes = _chain_nodes(leaf) - ctext = _chain_text(nodes) - lcp = _text_lcp(ctext, incoming_text) - if lcp > best_len or (lcp == best_len and best_leaf is not None and leaf.created_at > best_leaf.created_at): - best_leaf = leaf - best_len = lcp - - # Case 1: no text overlap at all - if best_leaf is None or best_len == 0: - return MatchResult(case="case1", anchor=tree.root, text_matched_len=0) - - nodes = _chain_nodes(best_leaf) - ctext = _chain_text(nodes) - - # §7.1: incoming text <= chain text and fully covered (no new tail, incl. equal) - if best_len == len(incoming_text) and len(incoming_text) <= len(ctext): - return MatchResult(case="substring", anchor=best_leaf, text_matched_len=best_len) - - # Text reached leaf and incoming has a new tail -> TITO decision (Case 2/3/4) - if best_len == len(ctext) and len(incoming_text) > len(ctext): - drift = [] - for i, node in enumerate(nodes): - p, r = turns[i] - if list(p.tokens) != node.prompt_delta_tokens or list(r.tokens) != node.resp_tokens: - drift.append(node) - if not drift: - return MatchResult(case="case2", anchor=best_leaf, text_matched_len=best_len) - only_tail = len(drift) == 1 and drift[0] is nodes[-1] - return MatchResult( - case="case3" if only_tail else "case4", - anchor=best_leaf, - text_matched_len=best_len, - drift_nodes=drift, + def __init__( + self, + *, + tokenizer=None, + chat_template_kwargs: dict[str, Any] | None = None, + end_of_turn_token_id: int | None = None, + tito_snapshot_min_loss_tokens: int | None = None, + ) -> None: + # tokenizer / chat_template_kwargs are no longer load-bearing under + # plan C, but the constructor signature is kept for callsite + # compatibility. _tokenizer is retained for forward-compat (callers + # constructing TrajectoryManager(tokenizer=tok) shouldn't break). + self._tokenizer = tokenizer + self._ct_kwargs: dict[str, Any] = dict(chat_template_kwargs or {}) + self._end_of_turn_token_id = end_of_turn_token_id + # Drift-snapshot threshold (loss_mask=1 token count inside drift suffix). + # None or <= 0 disables; behavior then matches the pre-feature output. + self._snap_threshold: int | None = ( + tito_snapshot_min_loss_tokens + if (tito_snapshot_min_loss_tokens is not None and tito_snapshot_min_loss_tokens > 0) + else None ) + self._trees: dict[str, Node] = {} + self._turn_count: dict[str, int] = {} + + # -------------------- public ------------------------------------------ + + def has_session(self, sid: str) -> bool: + return sid in self._trees + + def turn_count(self, sid: str) -> int: + return self._turn_count.get(sid, 0) + + def append_turn( + self, + sid: str, + *, + prompt_messages: list[dict[str, Any]], + tools: list[dict[str, Any]] | None, + prompt_ids: list[int], + response_ids: list[int], + response_logprobs: list[float] | None, + response_message: dict[str, Any] | None, + finish_reason: str | None, + metadata: dict[str, Any] | None = None, + ) -> None: + if not prompt_messages: + logger.warning("append_turn(sid=%s): empty prompt_messages; skipping", sid) + return + if response_logprobs is not None and len(response_logprobs) != len(response_ids): + raise ValueError( + f"response_logprobs length {len(response_logprobs)} != " f"response_ids length {len(response_ids)}" + ) - # Text stops inside the chain (0 < best_len < len(ctext)) -> Case 5/6 - cum = 0 - fork_node = None - fork_idx = 0 - for idx, node in enumerate(nodes): - seg_len = len(node.node_text) - if cum + seg_len > best_len: - fork_node = node - fork_idx = idx - break - cum += seg_len - if fork_node is None: - fork_node = nodes[-1] - fork_idx = len(nodes) - 1 - residual_count = len(nodes) - fork_idx # X..leaf - lca = fork_node.parent - return MatchResult( - case="case5" if residual_count <= 1 else "case6", - anchor=best_leaf, - text_matched_len=best_len, - lca=lca, - fork_node=fork_node, - residual_count=residual_count, - ) - - -# ---- attach_turn ---- - - -def attach_turn(tree: TrajectoryTree, match: MatchResult, turns) -> Node | None: - last_p, last_r = turns[-1] - - if match.case == "substring": - logger.warning( - "trajectory drop | reason=substring-prefix redelivery; " - "incoming is a true prefix of an existing chain, dropped" + root = self._trees.get(sid) + if root is None: + root = Node() + self._trees[sid] = root + + groups = _group_messages_by_role(prompt_messages) + + # Step 1: DFS by (role, node_match_key) ONLY. No prompt_ids check. + cur = root + i = 0 + while i < len(groups): + g_key = node_match_key(groups[i].messages) + match: Node | None = None + for child in cur.children: + if child.role == groups[i].role and node_match_key(child.messages) == g_key: + match = child + break + if match is None: + break + cur = match + i += 1 + + # Step 2: mount remaining prompt groups as plain routing nodes. + # Token attribution happens at get_trajectory time, not here. + for g in groups[i:]: + md: dict[str, Any] = {} + if g.role == "system" and tools is not None and not self._first_system_already_set(cur): + md["tools"] = list(tools) + cur = cur.add_child(Node(role=g.role, messages=list(g.messages), metadata=md)) + + # Step 3: assistant leaf with this turn's sglang snapshot. + asst_messages = [response_message] if response_message is not None else [] + asst = Node( + role="assistant", + messages=asst_messages, + metadata=dict(metadata or {}), ) - return None - - if match.case == "case1": - return Node.new_turn(last_p, last_r, parent=tree.root) - - if match.case == "case2": - return Node.new_turn(last_p, last_r, parent=match.anchor) - - if match.case in ("case3", "case4"): - nodes = _chain_nodes(match.anchor) - # Replace each drifted node's segment in place (same-index turn segment). - idx_by_node = {id(n): i for i, n in enumerate(nodes)} - for node in match.drift_nodes: - i = idx_by_node[id(node)] - p, r = turns[i] - before_text, before_tokens = node.node_text, list(node.resp_tokens) - node.prompt_delta_text = p.text - node.prompt_delta_tokens = list(p.tokens) - node.resp_text = r.text - node.resp_tokens = list(r.tokens) - node.resp_logprobs = [0.0] * len(r.tokens) # §2.4 placeholder - node.replaced_count += 1 - node.updated_at = time.time() - reason = ( - "tokenizer TITO drift at tail turn" - if match.case == "case3" - else "multi-turn token drift (template changed?)" + asst.turn_prompt_ids = list(prompt_ids) + asst.turn_response_ids = list(response_ids) + asst.turn_response_logprobs = list(response_logprobs) if response_logprobs is not None else None + asst.turn_finish_reason = finish_reason + asst.turn_index = self._turn_count.get(sid, 0) + 1 + cur.add_child(asst) + + self._turn_count[sid] = asst.turn_index + + def get_trajectory( + self, + sid: str, + *, + base_sample=None, + reward: float = 0.0, + extra_metadata: dict[str, Any] | None = None, + drop: bool = True, + ) -> list: + """Linearize each leaf into a slime ``Sample`` using LCP drop-and-replace. + + For each leaf, walk root→leaf collecting assistant nodes in order. + Start with tokens=[]. For each assistant turn k (1-based): + + 1. ``p = asst.turn_prompt_ids``, ``r = asst.turn_response_ids``. + 2. If k == 1: emit all of ``p`` as loss_mask=0 (plus 0.0 logprobs), + then ``r`` as loss_mask=1 with real logprobs. + 3. If k >= 2: compute ``L = LCP(tokens, p)``. Truncate tokens / + loss_mask / logprobs to length L (DROP everything past L — + that includes the previous turn's response tokens that fall + in the drift region; logging tells you how much was dropped). + Then append ``p[L:]`` (loss_mask=0) and ``r`` (loss_mask=1). + + When drift fires on at least one turn, the returned Sample's + ``metadata`` gains ``tito_dropped_tokens`` (total tokens dropped + across the leaf) and ``tito_dropped_turns`` (how many turns + triggered a drop). Both keys are absent when no drift occurs. + + When ``tito_snapshot_min_loss_tokens`` was passed to the constructor + and a drift would drop >= that many loss_mask=1 tokens, an extra + snapshot Sample is emitted before the main-leaf Sample carrying just + the to-be-lost tokens (complementary mask). See module docstring. + + See module docstring for the rationale. + """ + if base_sample is None: + base_sample = Sample(index=0, prompt="") + + root = self._trees.get(sid) + if root is None: + return [] + leaves = [leaf for leaf in root.leaves() if not leaf.is_root] + samples: list[Sample] = [] + for leaf in leaves: + chain = leaf.path_from_root() + # Only assistant leaves carrying this turn's sglang snapshot + # participate in TITO accumulation. Routing assistant nodes mounted + # from prior-turn replay (turn_prompt_ids is None) carry no token + # signal and would otherwise be misread as a full-trajectory drift. + asst_chain = [n for n in chain if n.role == "assistant" and n.turn_prompt_ids is not None] + + tokens: list[int] = [] + loss_mask: list[int] = [] + logprobs: list[float] = [] + total_dropped = 0 + dropped_turns = 0 + snapshots: list[tuple[list[int], list[int], list[float], int, int]] = [] + + for k, asst in enumerate(asst_chain, start=1): + p = list(asst.turn_prompt_ids or []) + r = list(asst.turn_response_ids or []) + lp = list(asst.turn_response_logprobs) if asst.turn_response_logprobs is not None else None + + if k == 1: + emit_prompt = p + else: + L = _lcp_len(tokens, p) + drift = len(tokens) - L + if drift > 0: + drift_loss_tokens = sum(loss_mask[L:]) + snap_emitted = False + if self._snap_threshold is not None and drift_loss_tokens >= self._snap_threshold: + snap_tokens = list(tokens) + snap_mask = [0] * L + list(loss_mask[L:]) + snap_lp = [0.0] * L + [ + (logprobs[i] if loss_mask[i] == 1 else 0.0) for i in range(L, len(tokens)) + ] + snapshots.append((snap_tokens, snap_mask, snap_lp, asst.turn_index, k - 1)) + snap_emitted = True + logger.warning( + "get_trajectory(sid=%s leaf turn=%s): TITO drift detected, " + "dropping %d prior tokens (incl. previous-turn response) to " + "realign with this turn's prompt%s", + sid, + asst.turn_index, + drift, + f"; snapshotted {drift_loss_tokens} loss tokens" if snap_emitted else "", + ) + if not snap_emitted: + total_dropped += drift + dropped_turns += 1 + tokens = tokens[:L] + loss_mask = loss_mask[:L] + logprobs = logprobs[:L] + emit_prompt = p[L:] + + tokens.extend(emit_prompt) + loss_mask.extend([0] * len(emit_prompt)) + logprobs.extend([0.0] * len(emit_prompt)) + + tokens.extend(r) + loss_mask.extend([1] * len(r)) + if lp is not None: + logprobs.extend(lp) + else: + logprobs.extend([0.0] * len(r)) + + last_asst = asst_chain[-1] if asst_chain else None + first_sys = next((n for n in chain if n.role == "system"), None) + tools_meta = first_sys.metadata.get("tools") if first_sys else None + base_md: dict[str, Any] = { + **(base_sample.metadata or {}), + **(extra_metadata or {}), + "tools": tools_meta, + } + per_leaf_reward = (reward / len(leaves)) if leaves else 0.0 + + # Emit snapshot sample(s) first, then the main-leaf sample. + for snap_tokens, snap_mask, snap_lp, drift_turn, cur_chain_idx in snapshots: + snap_finish = None + prev_idx = cur_chain_idx - 1 # asst_chain index of the previous (prefix's last) turn + if 0 <= prev_idx < len(asst_chain): + snap_finish = asst_chain[prev_idx].turn_finish_reason + snap_md = { + **base_md, + "finish_reason": snap_finish, + "tito_snapshot": True, + "tito_snapshot_at_turn": drift_turn, + "tito_snapshot_loss_tokens": sum(snap_mask), + } + samples.append( + Sample( + index=base_sample.index, + rollout_id=( + base_sample.rollout_id if base_sample.rollout_id is not None else base_sample.index + ), + prompt=base_sample.prompt, + label=base_sample.label, + tokens=snap_tokens, + response_length=sum(1 for m in snap_mask if m == 1), + loss_mask=snap_mask, + rollout_log_probs=snap_lp, + reward=per_leaf_reward, + status=Sample.Status.COMPLETED, + metadata=snap_md, + ) + ) + + response_length = sum(1 for m in loss_mask if m == 1) + main_md: dict[str, Any] = { + **base_md, + "finish_reason": last_asst.turn_finish_reason if last_asst else None, + } + if total_dropped > 0: + main_md["tito_dropped_tokens"] = total_dropped + main_md["tito_dropped_turns"] = dropped_turns + if snapshots: + main_md["tito_snapshots_emitted"] = len(snapshots) + samples.append( + Sample( + index=base_sample.index, + rollout_id=(base_sample.rollout_id if base_sample.rollout_id is not None else base_sample.index), + prompt=base_sample.prompt, + label=base_sample.label, + tokens=tokens, + response_length=response_length, + loss_mask=loss_mask, + rollout_log_probs=logprobs, + reward=per_leaf_reward, + status=Sample.Status.COMPLETED, + metadata=main_md, + ) ) - _log_replace(match.case, reason, node, before_text, node.node_text, before_tokens, node.resp_tokens) - # mask interval - if match.case == "case3": - match.drift_nodes[0].loss_mask = True - else: - first = match.drift_nodes[0] - first.warning = "multi-turn token drift" - start = idx_by_node[id(first)] - for node in nodes[start:]: - node.loss_mask = True - # append the new turn at the tail - return Node.new_turn(last_p, last_r, parent=nodes[-1]) - - if match.case == "case5": - nodes = _chain_nodes(match.anchor) - x = match.fork_node - xi = nodes.index(x) - p, r = turns[xi] - before_text, before_tokens = x.node_text, list(x.resp_tokens) - x.prompt_delta_text = p.text - x.prompt_delta_tokens = list(p.tokens) - x.resp_text = r.text - x.resp_tokens = list(r.tokens) - x.resp_logprobs = [0.0] * len(r.tokens) - x.loss_mask = True - x.replaced_count += 1 - x.updated_at = time.time() - _log_replace( - "case5", "upstream response format rewrite", x, before_text, x.node_text, before_tokens, x.resp_tokens - ) - # if turns has newer turns after x, append them in order - parent = x - for j in range(xi + 1, len(turns)): - pj, rj = turns[j] - parent = Node.new_turn(pj, rj, parent=parent) - return parent - - if match.case == "case6": - nodes = _chain_nodes(match.anchor) - x = match.fork_node - xi = nodes.index(x) - lca = match.lca # = x.parent - parent = lca - last = None - for j in range(xi, len(turns)): - pj, rj = turns[j] - last = Node.new_turn(pj, rj, parent=parent) - parent = last - return last - - raise ValueError(f"unknown match case: {match.case}") - - -# ---- record_turn ---- - - -def record_turn(tree: TrajectoryTree, turns) -> Node | None: - match = match_prefix(tree, turns) - return attach_turn(tree, match, turns) - - -# ---- export ---- - - -def _subtree_has_trainable(node: Node) -> bool: - """Whether any node in the subtree (rooted at ``node``) carries a trainable - response segment (``loss_mask`` False and non-empty ``resp_tokens``). Used by - owner selection so a fork's owner branch is not later dropped by the slice - layer for being all-mask.""" - stack = [node] - while stack: - cur = stack.pop() - if not cur.kind_is_root and not cur.loss_mask and cur.resp_tokens: - return True - stack.extend(cur.children) - return False - - -def _decide_fork_owners(root: Node) -> dict: - owners = {} - - def walk(n: Node): - if len(n.children) >= 2: - # Decision E: prefer the earliest child whose subtree still has a - # trainable leaf, so the shared prefix it owns survives slicing. - # Fall back to the earliest child when every branch is all-mask - # (then the shared prefix has nowhere trainable to go anyway). - trainable = [c for c in n.children if _subtree_has_trainable(c)] - pool = trainable or list(n.children) - owner = min(pool, key=lambda c: c.created_at) - owners[id(n)] = owner - for c in n.children: - walk(c) - - walk(root) - return owners - - -def _masked_by_fork(node: Node, leaf: Node, fork_owner: dict) -> bool: - """Whether ``node`` is force-masked because some fork ancestor ``F`` on the - leaf's path chose a non-owner branch toward ``leaf``: i.e. ``node`` is at or - before ``F`` (incl. ``F`` itself) and at ``F`` the leaf's branch != owner.""" - # leaf's root->leaf path (incl. root) - path = [] - cur = leaf - while cur is not None: - path.append(cur) - cur = cur.parent - path.reverse() # root ... leaf - pos = {id(n): i for i, n in enumerate(path)} - node_pos = pos.get(id(node)) - if node_pos is None: - return False - # for every fork point F on path (present in fork_owner) - for i, F in enumerate(path): - if id(F) not in fork_owner: - continue - # leaf's chosen child at F = path[i+1] - if i + 1 >= len(path): - continue - chosen_child = path[i + 1] - owner = fork_owner[id(F)] - if chosen_child is not owner: - # non-owner branch: mask nodes at or before F (root..F) - if node_pos <= i: - return True - return False - - -def export(tree: TrajectoryTree): - samples, masks, logprobs = [], [], [] - fork_owner = _decide_fork_owners(tree.root) - for leaf in _collect_leaves(tree.root): - nodes = _chain_nodes(leaf) - toks, mask, lp = [], [], [] - for node in nodes: - toks += node.prompt_delta_tokens - mask += [0] * len(node.prompt_delta_tokens) - lp += [0.0] * len(node.prompt_delta_tokens) - toks += node.resp_tokens - if _masked_by_fork(node, leaf, fork_owner): - bit = 0 - elif node.loss_mask: - bit = 0 - else: - bit = 1 - mask += [bit] * len(node.resp_tokens) - lp += list(node.resp_logprobs) - samples.append(toks) - masks.append(mask) - logprobs.append(lp) - return samples, masks, logprobs - - -# ============================================================================= -# slime glue: full-length export -> prompt/response TokenSegment + Sample fan-out -# ============================================================================= - - -@dataclasses.dataclass(frozen=True) -class TokenSegment: - """One training segment assembled from an agent trajectory. - - slime training invariants: ``tokens = prompt_ids + response_ids``, - ``response_length = len(response_ids)`` and - ``len(loss_mask) == len(rollout_log_probs) == response_length``. - """ + if drop: + self._trees.pop(sid, None) + self._turn_count.pop(sid, None) + return samples - prompt_ids: list[int] - response_ids: list[int] - loss_mask: list[int] - rollout_log_probs: list[float] = dataclasses.field(default_factory=list) - metadata: dict[str, Any] = dataclasses.field(default_factory=dict) + # -------------------- internals ---------------------------------------- + @staticmethod + def _first_system_already_set(start: Node) -> bool: + """Walk start->root looking for a system node already carrying tools.""" + cur: Node | None = start + while cur is not None and not cur.is_root: + if cur.role == "system" and cur.metadata.get("tools") is not None: + return True + cur = cur.parent + return False -def export_token_segments(tree: TrajectoryTree, *, metadata: dict[str, Any] | None = None) -> list[TokenSegment]: - """Turn the tree's full-length export into slime ``TokenSegment``s. - Each leaf chain is split at the first trainable token (design §6, - ``mask.index(1)``): tokens before the cut are the prompt, tokens from the cut - on are the response. Leaves whose whole chain is masked (no trainable token, - e.g. fork non-owner or Case4/5 fully-masked) are dropped -- they have no - trainable response and would break slime's ``response_length`` contract. - """ - out: list[TokenSegment] = [] - tokens_list, masks_list, logprobs_list = export(tree) - for tokens, mask, logprobs in zip(tokens_list, masks_list, logprobs_list, strict=True): - if 1 not in mask: - continue # all-mask chain: nothing trainable, drop - cut = mask.index(1) - response_ids = tokens[cut:] - if not response_ids: - continue - segment = TokenSegment( - prompt_ids=list(tokens[:cut]), - response_ids=list(response_ids), - loss_mask=list(mask[cut:]), - rollout_log_probs=list(logprobs[cut:]), - metadata={**(metadata or {}), "segment_kind": "leaf"}, - ) - assert len(segment.loss_mask) == len(segment.rollout_log_probs) == len(segment.response_ids) - out.append(segment) - return out - - -def write_segment_to_sample(sample: Sample, segment: TokenSegment, reward: float, tokenizer) -> None: - """Populate token, mask, response, reward, and status fields from a segment.""" - sample.tokens = list(segment.prompt_ids) + list(segment.response_ids) - sample.response_length = len(segment.response_ids) - sample.loss_mask = list(segment.loss_mask) - sample.rollout_log_probs = list(segment.rollout_log_probs) - sample.response = tokenizer.decode(segment.response_ids, skip_special_tokens=False) - sample.reward = float(reward) - sample.status = Sample.Status.COMPLETED - - -def fan_out_sample_segments( - sample: Sample, - segments: list[TokenSegment], - reward: float, - tokenizer, - *, - metadata: dict[str, Any] | None = None, -) -> list[Sample]: - """Emit one Sample per segment, splitting reward uniformly across them. - - Sibling samples share ``group_id`` so reducers that average by group do - not over-count trajectories split by compaction or sub-agent dispatch. - """ - k = len(segments) - per_segment_reward = float(reward) / max(1, k) - shared_group_id = sample.group_id if sample.group_id is not None else sample.index - base_metadata = {**(sample.metadata or {}), **(metadata or {})} - - out: list[Sample] = [] - for i, segment in enumerate(segments): - sub = sample if i == 0 else copy.copy(sample) - write_segment_to_sample(sub, segment, per_segment_reward, tokenizer) - sub.group_id = shared_group_id - sub.metadata = { - **base_metadata, - **(segment.metadata or {}), - "segment_idx": i, - "num_segments": k, - } - out.append(sub) - return out +__all__ = [ + "Node", + "TrajectoryManager", + "node_match_key", + "_group_messages_by_role", + "_lcp_len", +] From b8a5cd4e0902bb116fd53ce43d5b0902d0adaf69 Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Thu, 4 Jun 2026 15:50:01 +0000 Subject: [PATCH 04/28] refactor(agent): rewrite openai.py for Codex CLI + TrajectoryManager Rewrite slime/agent/adapters/openai.py on top of the new TrajectoryManager-based architecture so the Codex CLI (wire_api="chat", v0.30.0) running inside an e2b sandbox can drive the slime SGLang backend the same way anthropic.py drives Claude Code. Key wire-format alignments for Codex 0.30.0 (encoded in _build_oai_response / _stream_chat_completion): * Emit all parallel tool_calls in a single SSE chunk -- Codex 0.30 accumulates per-index arguments fragments across chunks and would otherwise merge them into one tool_call with concatenated args. * wire_message.tool_calls is truncated to the first call -- Codex silently drops the rest on echo, which would fork node_match_key. * When tool_calls are present, wire_message.content=None and manager_message.content="" -- Codex splits a single assistant-with-text-and-tool_calls into two echoed messages, so we suppress the text on the wire side to keep the echo single-shaped. * manager_message intentionally omits reasoning_content -- Codex strips it on echo; reasoning token ids stay in response_ids so loss is unaffected. Also revert Sample.rollout_id -> Sample.group_id in trajectory_manager.py to match the upstream Sample field rename (rollout_id is now write-only deprecated and raises on read), which is hit at finish_session time and is a prerequisite for the openai e2e path to run. Verified: pytest smoke (1 SWE instance, e2b sandbox + Codex CLI -> OpenAIAdapter -> local sglang:30000) -> rc=0, forks=0, leaves=1, turns=39 over 5.8M tokens with 32 tokens of expected TITO drift (reasoning text not echoed back). --- slime/agent/adapters/openai.py | 758 +++++++++++++++++------------- slime/agent/trajectory_manager.py | 8 +- 2 files changed, 434 insertions(+), 332 deletions(-) diff --git a/slime/agent/adapters/openai.py b/slime/agent/adapters/openai.py index 79a1b0b824..8697c09639 100644 --- a/slime/agent/adapters/openai.py +++ b/slime/agent/adapters/openai.py @@ -1,21 +1,32 @@ -"""OpenAI-compatible adapters for agent rollouts. - -The adapter exposes ``/v1/chat/completions`` and ``/v1/responses``. Both -endpoints render incoming messages with the served model's chat template, call -SGLang ``/generate`` with ``input_ids``, and fold the turn into a per-session -turn-node :class:`~slime.agent.trajectory_manager.TrajectoryTree`. The tree -routes everything by text prefix, so there is no manual new/append/wipe -bookkeeping. Call ``finish_session()`` at trajectory end to drain trainable -``TokenSegment`` objects. +"""OpenAI Chat-Completions adapter for agent rollouts. + +Mirrors :mod:`slime.agent.adapters.anthropic` but speaks the OpenAI +``/v1/chat/completions`` wire protocol so that the Codex CLI (and any other +OpenAI-compatible client) can drive the slime SGLang server. Each incoming +request is rendered with the served model's chat template, sent to SGLang +``/generate`` as ``input_ids``, parsed, and folded into a shared +:class:`~slime.agent.trajectory_manager.TrajectoryManager` keyed by session id. +``finish_session(sid)`` drains a session's trajectory into a list of +:class:`~slime.utils.types.Sample`. + +The per-sid tree inside TrajectoryManager handles sub-agent and compaction +patterns automatically (any divergence in the prompt prefix forks into a new +leaf), so we do not track explicit chains here. + +Only ``/v1/chat/completions`` is implemented; the older Responses API +(``/v1/responses``) is intentionally out of scope -- Codex 0.30.0 uses +``wire_api = "chat"``. """ from __future__ import annotations import asyncio +import dataclasses import json import logging import secrets import time +from collections.abc import Callable from typing import Any from aiohttp import web @@ -26,34 +37,100 @@ TOKENIZER_KEY, TOOL_PARSER_KEY, BaseAdapter, - GenResult, - Session, - assemble_turns, call_sglang_generate, + ok_response, + request_session_id, ) -from slime.agent.adapters.common import json_arguments as _json_arguments -from slime.agent.adapters.common import ok_response, render_prompt, request_session_id from slime.agent.parsing import ParsedModelOutput, parse_model_output -from slime.agent.trajectory_manager import record_turn +from slime.agent.trajectory_manager import TrajectoryManager +from slime.utils.types import Sample logger = logging.getLogger(__name__) -class OpenAIAdapter(BaseAdapter): - """OpenAI-compatible HTTP adapter with session lifecycle helpers.""" +@dataclasses.dataclass +class Session: + """Per-sid adapter state: sampling defaults, context budget, request lock. + + Trajectory state lives in ``OpenAIAdapter.manager`` (one shared + TrajectoryManager keyed by sid across all sessions). + """ + + sampling_defaults: dict = dataclasses.field(default_factory=dict) + max_context_tokens: int = 0 + lock: asyncio.Lock = dataclasses.field(default_factory=asyncio.Lock) + - def __init__(self, *, tokenizer, sglang_url, tool_parser=None, reasoning_parser=None) -> None: +class OpenAIAdapter(BaseAdapter): + """OpenAI Chat-Completions-compatible HTTP adapter with session lifecycle helpers.""" + + session_cls = Session + + def __init__( + self, + *, + tokenizer, + sglang_url, + tool_parser=None, + reasoning_parser=None, + tito_snapshot_min_loss_tokens: int | None = None, + max_turns_per_sid: int | None = None, + on_turn_appended: Callable[..., None] | None = None, + ) -> None: super().__init__( tokenizer=tokenizer, sglang_url=sglang_url, tool_parser=tool_parser, reasoning_parser=reasoning_parser, ) + # ONE manager shared across all sids; per-sid trees live inside. + self.manager = TrajectoryManager( + tokenizer=tokenizer, + tito_snapshot_min_loss_tokens=tito_snapshot_min_loss_tokens, + ) + # Optional debug hook invoked after each successful append_turn. + # Signature mirrors AnthropicAdapter.on_turn_appended: + # (sid, prompt_messages, tools, response_message, + # prompt_ids, response_ids, finish_reason) -> None. + # Exceptions are swallowed; never block the HTTP response. + self.on_turn_appended: Callable[..., None] | None = on_turn_appended + # Per-sid turn cap; None disables. When set, /v1/chat/completions + # returns 429 once a sid has made this many turns. + self.max_turns_per_sid: int | None = max_turns_per_sid + self._sid_turn_count: dict[str, int] = {} self.app.router.add_post("/v1/chat/completions", _handle_chat_completions) - self.app.router.add_post("/v1/responses", _handle_responses) self.app.router.add_get("/healthz", _ok) self.app.router.add_get("/v1/models", _ok) + async def finish_session( + self, + sid: str, + *, + base_sample: Sample | None = None, + reward: float = 0.0, + extra_metadata: dict[str, Any] | None = None, + wait_timeout: float = 5.0, + ) -> list[Sample]: + """Drain a session's trajectory into Sample objects. + + Waits out in-flight requests for ``sid``, then linearises the + per-sid tree via ``TrajectoryManager.get_trajectory``. Idempotent -- + a second call for an already-popped sid returns ``[]``. + """ + await self.shutdown_session(sid, wait_timeout=wait_timeout) + self.store.pop(sid, None) + return self.manager.get_trajectory( + sid, + base_sample=base_sample, + reward=reward, + extra_metadata=extra_metadata, + ) + + +# ============================================================================= +# Translation (OpenAI wire <-> chat-template messages) +# ============================================================================= + def _flatten_content(content: Any) -> str: """Flatten OpenAI text/content parts into a chat-template string.""" @@ -84,207 +161,255 @@ def _flatten_content(content: Any) -> str: return "\n".join(p for p in parts if p) -def _normalize_tool_call(call: dict[str, Any]) -> dict[str, Any]: - function = call.get("function") or {} - name = function.get("name") or call.get("name") or "tool" - arguments = function.get("arguments", call.get("arguments", {})) - out = { - "type": "function", - "function": { - "name": name, - "arguments": _json_arguments(arguments), - }, - } - if call.get("id"): - out["id"] = call["id"] - return out - +def _arguments_as_dict(arguments: Any) -> dict[str, Any]: + """Coerce wire-shape ``tool_calls[].function.arguments`` into a dict. -def _translate_chat_messages(messages: list[dict]) -> list[dict]: - """OpenAI chat messages -> tokenizer chat-template messages.""" + OpenAI sends ``arguments`` as a JSON-encoded string; the chat template and + ``trajectory_manager.node_match_key`` both expect a mapping. ``json.loads`` + is tried first; malformed payloads fall back to ``{"_raw_arguments": s}`` + (mirrors :func:`slime.agent.parsing.parse_tool_uses`). + """ + if isinstance(arguments, dict): + return arguments + if arguments is None: + return {} + if isinstance(arguments, str): + s = arguments.strip() + if not s: + return {} + try: + parsed = json.loads(s) + except json.JSONDecodeError: + return {"_raw_arguments": arguments} + return parsed if isinstance(parsed, dict) else {"_raw_arguments": arguments} + return {"_raw_arguments": str(arguments)} + + +def _translate_openai_chat(messages: list[dict]) -> list[dict]: + """OpenAI chat messages -> tokenizer chat-template messages. + + Mirrors :func:`slime.agent.adapters.anthropic._translate_anthropic` so that + a replayed assistant turn hashes identically (via + ``trajectory_manager.node_match_key``) to the leaf the manager appended on + the previous request. Two invariants must hold: + + * ``tool_calls[i].function.arguments`` is a ``dict`` (NOT a JSON string). + Qwen3-style chat templates call ``arguments | items`` which requires + a mapping; ``node_match_key`` uses ``json.dumps(sort_keys=True)`` so + equivalent dicts hash the same regardless of key order. + * Wire-only correlation ids are DROPPED: ``tool_call_id`` on ``role: + "tool"`` history messages and ``tool_calls[i].id`` on echoed assistant + messages. The adapter mints fresh ids on each response, so keeping the + wire ids would diverge the replay hash from the original leaf. + """ translated: list[dict] = [] for msg in messages: if not isinstance(msg, dict): continue role = msg.get("role") content = msg.get("content") - if role == "developer": + if role == "developer": # OpenAI Responses API alias role = "system" if role in {"system", "user"}: translated.append({"role": role, "content": _flatten_content(content)}) elif role == "tool": - tool_msg = {"role": "tool", "content": _flatten_content(content)} - if msg.get("tool_call_id"): - tool_msg["tool_call_id"] = msg["tool_call_id"] - translated.append(tool_msg) + # DROP tool_call_id -- wire-only correlation field; see docstring. + translated.append({"role": "tool", "content": _flatten_content(content)}) elif role == "assistant": - assistant: dict[str, Any] = {"role": "assistant", "content": _flatten_content(content)} - if msg.get("reasoning_content"): - assistant["reasoning_content"] = msg["reasoning_content"] + assistant: dict[str, Any] = { + "role": "assistant", + "content": _flatten_content(content), + } + reasoning = msg.get("reasoning_content") + if reasoning: + assistant["reasoning_content"] = reasoning tool_calls = msg.get("tool_calls") or [] - if tool_calls: - assistant["tool_calls"] = [_normalize_tool_call(c) for c in tool_calls if isinstance(c, dict)] + normalized: list[dict[str, Any]] = [] + for call in tool_calls: + if not isinstance(call, dict): + continue + function = call.get("function") or {} + name = function.get("name") or call.get("name") or "tool" + arguments = function.get("arguments") + if arguments is None: + arguments = call.get("arguments", {}) + # NB: arguments stays a dict (NOT a JSON string), and we DROP + # the wire-only ``id``. See docstring above. + normalized.append( + { + "type": "function", + "function": { + "name": name, + "arguments": _arguments_as_dict(arguments), + }, + } + ) + if normalized: + assistant["tool_calls"] = normalized translated.append(assistant) + # Unknown roles are silently dropped. return translated -def _normalize_tool(tool: dict[str, Any]) -> dict[str, Any] | None: - if not isinstance(tool, dict): +def _openai_tools_to_chat_tools(tools: list[dict] | None) -> list[dict] | None: + """Convert OpenAI tools list to tokenizer chat-template tool schema.""" + if not tools: return None - if tool.get("type") != "function": - return None - if isinstance(tool.get("function"), dict): - function = tool["function"] - name = function.get("name") - if not name: - return None - return { - "type": "function", - "function": { - "name": name, - "description": function.get("description", ""), - "parameters": function.get("parameters") or {"type": "object", "properties": {}}, - }, - } - name = tool.get("name") - if not name: - return None - return { - "type": "function", - "function": { - "name": name, - "description": tool.get("description", ""), - "parameters": tool.get("parameters") or {"type": "object", "properties": {}}, - }, - } - - -def _normalize_tools(tools: list[dict] | None) -> list[dict] | None: - normalized = [_normalize_tool(t) for t in tools or []] - return [t for t in normalized if t is not None] or None - - -def _responses_input_to_messages(input_value: Any, instructions: Any = None) -> list[dict]: - """Responses API input -> OpenAI chat message list. - - This intentionally covers the common message/function-call shapes used by - agent SDKs. Unknown input items are preserved as user text where possible. - """ - messages: list[dict] = [] - if instructions: - messages.append({"role": "system", "content": _flatten_content(instructions)}) - - if isinstance(input_value, str): - messages.append({"role": "user", "content": input_value}) - return messages - - if not isinstance(input_value, list): - messages.append({"role": "user", "content": _flatten_content(input_value)}) - return messages - - for item in input_value: - if isinstance(item, str): - messages.append({"role": "user", "content": item}) + normalized: list[dict] = [] + for tool in tools: + if not isinstance(tool, dict): continue - if not isinstance(item, dict): - messages.append({"role": "user", "content": str(item)}) + if tool.get("type") and tool.get("type") != "function": continue - - typ = item.get("type") - if typ == "function_call_output": - messages.append( + function = tool.get("function") if isinstance(tool.get("function"), dict) else None + if function is not None: + name = function.get("name") + if not name: + continue + normalized.append( { - "role": "tool", - "tool_call_id": item.get("call_id") or item.get("id") or "", - "content": item.get("output", ""), + "type": "function", + "function": { + "name": name, + "description": function.get("description", ""), + "parameters": function.get("parameters") or {"type": "object", "properties": {}}, + }, } ) - elif typ == "function_call": - messages.append( + else: + name = tool.get("name") + if not name: + continue + normalized.append( { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": item.get("call_id") or item.get("id") or f"call_{secrets.token_hex(8)}", - "type": "function", - "function": { - "name": item.get("name", "tool"), - "arguments": item.get("arguments", "{}"), - }, - } - ], + "type": "function", + "function": { + "name": name, + "description": tool.get("description", ""), + "parameters": tool.get("parameters") or {"type": "object", "properties": {}}, + }, } ) - elif item.get("role"): - messages.append({"role": item.get("role"), "content": item.get("content", "")}) - elif typ == "message": - messages.append({"role": item.get("role", "user"), "content": item.get("content", "")}) - else: - messages.append({"role": "user", "content": _flatten_content(item)}) - return messages - - -async def _generate(prompt_ids: list[int], s: Session, body: dict, app, *, session_id: str | None = None): - return await call_sglang_generate( - prompt_ids, - s, - body, - app, - max_token_keys=("max_output_tokens", "max_completion_tokens", "max_tokens"), - stop_keys=("stop",), - log_prefix="openai_adapter", - logger=logger, - session_id=session_id, + return normalized or None + + +# ============================================================================= +# Local chat-template render helper. Mirrors anthropic.py:_render_token_ids. +# ============================================================================= + + +def _render_token_ids( + messages: list[dict], + tokenizer, + *, + tools: list[dict] | None, + add_generation_prompt: bool = True, +) -> list[int]: + enc = tokenizer.apply_chat_template( + messages, + tools=tools, + tokenize=True, + add_generation_prompt=add_generation_prompt, ) + ids = enc["input_ids"] if hasattr(enc, "__getitem__") and "input_ids" in enc else enc + return list(ids) -def _parse_output(output_text: str, tools_schema: list[dict] | None, app) -> ParsedModelOutput: - return parse_model_output( - output_text or "", - tools_schema=tools_schema, - tool_parser_name=app[TOOL_PARSER_KEY], - reasoning_parser_name=app[REASONING_PARSER_KEY], - ) +# ============================================================================= +# Reply building: parsed output -> OpenAI wire message + manager response_message +# ============================================================================= + +def _build_oai_response(parsed: ParsedModelOutput, finish: str) -> tuple[dict[str, Any], dict[str, Any], str]: + """Return ``(wire_message, manager_message, finish_reason)``. -def _openai_tool_calls(tool_uses: list[dict[str, Any]]) -> list[dict[str, Any]]: - calls: list[dict[str, Any]] = [] - for tool_use in tool_uses: + ``wire_message`` follows OpenAI Chat-Completions spec: ``tool_calls[].id`` + is a unique correlation id, and ``tool_calls[].function.arguments`` is a + JSON-encoded **string** (clients depend on this). + + ``manager_message`` is the shape fed to ``TrajectoryManager.append_turn``: + ``tool_calls[].function.arguments`` is a **dict** so chat-template replay + (Qwen3 etc.) succeeds and ``node_match_key`` hashes match the echo on the + next turn. The wire-only ``id`` is omitted (the next turn's echo will not + include it; matching anthropic.py:444-470). + """ + wire_tool_calls: list[dict[str, Any]] = [] + manager_tool_calls: list[dict[str, Any]] = [] + for tu in parsed.tool_uses: + name = tu.get("name", "tool") + args_dict = tu.get("input") or {} + if not isinstance(args_dict, dict): + args_dict = {"_raw_arguments": str(args_dict)} call_id = f"call_{secrets.token_hex(12)}" - calls.append( + wire_tool_calls.append( { "id": call_id, "type": "function", "function": { - "name": tool_use.get("name", "tool"), - "arguments": _json_arguments(tool_use.get("input") or {}), + "name": name, + "arguments": json.dumps(args_dict, ensure_ascii=False, sort_keys=True), + }, + } + ) + manager_tool_calls.append( + { + "type": "function", + "function": { + "name": name, + "arguments": args_dict, }, } ) - return calls + wire_message: dict[str, Any] = { + "role": "assistant", + # OpenAI spec allows null content when tool_calls are present. The + # Codex CLI 0.30.0 splits a single assistant turn containing both + # text and tool_calls into TWO echoed messages on the next request + # (a tool_calls-only one followed by a text-only one), which breaks + # node_match_key against our leaf -- so when we have tool_calls we + # send content=null so the echo is a single tool_calls-only message. + "content": None if wire_tool_calls else (parsed.text or None), + } + # ``manager_message`` must match what the OAI client will echo on the + # next request, otherwise ``node_match_key`` diverges and every turn + # forks. Three empirically necessary differences vs ``wire_message``: + # + # * NO ``reasoning_content`` -- the Codex CLI strips it on echo, so we + # must not store it in the leaf either. The reasoning token ids are + # preserved in ``response_ids`` (used for loss), only the rendered + # text is dropped. + # * Only the FIRST ``tool_call`` -- the Codex CLI silently drops any + # additional parallel tool_calls on echo (validates them serially and + # aborts after the first response), so the leaf can only hold one. + # * When ``tool_calls`` is present, ``content`` is empty -- the wire + # side also sends ``content=null`` (see above) so the echo stays + # a single tool_calls-only message that matches the leaf. + manager_message: dict[str, Any] = { + "role": "assistant", + "content": "" if wire_tool_calls else (parsed.text or ""), + } + if parsed.reasoning: + wire_message["reasoning_content"] = parsed.reasoning + if wire_tool_calls: + wire_message["tool_calls"] = wire_tool_calls[:1] + manager_message["tool_calls"] = manager_tool_calls[:1] -def _finish_reason(parsed: ParsedModelOutput, finish: str) -> str: if parsed.tool_uses: - return "tool_calls" - if finish == "length": - return "length" - return "stop" + wire_finish = "tool_calls" + elif finish == "length": + wire_finish = "length" + else: + wire_finish = "stop" + return wire_message, manager_message, wire_finish -def _chat_message(parsed: ParsedModelOutput) -> dict[str, Any]: - tool_calls = _openai_tool_calls(parsed.tool_uses) - message: dict[str, Any] = { - "role": "assistant", - "content": parsed.text if parsed.text else None, - } - if parsed.reasoning: - message["reasoning_content"] = parsed.reasoning - if tool_calls: - message["tool_calls"] = tool_calls - return message + +def _finish_reason_for_manager(finish: str, tool_uses: list[dict]) -> str: + if tool_uses: + return "tool_calls" + return finish or "stop" def _usage(in_tok: int, out_tok: int) -> dict[str, int]: @@ -295,65 +420,130 @@ def _usage(in_tok: int, out_tok: int) -> dict[str, int]: } -def _responses_usage(in_tok: int, out_tok: int) -> dict[str, int]: - return { - "input_tokens": in_tok, - "output_tokens": out_tok, - "total_tokens": in_tok + out_tok, - } +# ============================================================================= +# Request handling -- one full turn + JSON or SSE response +# ============================================================================= def _request_session_id(request: web.Request, body: dict) -> str: + """Resolve sid from request. + + Tries ``Authorization: Bearer `` first (the canonical OpenAI auth + header; Codex CLI propagates ``OPENAI_API_KEY`` here), then falls back + to body-level hints (``metadata.session_id`` / ``user``). + """ return request_session_id(request, body=body) -async def _run_turn( - request: web.Request, body: dict, messages: list[dict] -) -> tuple[GenResult, ParsedModelOutput, int, int]: +async def _handle_chat_completions(request: web.Request) -> web.StreamResponse: + body = await request.json() + messages = body.get("messages") or [] + if not isinstance(messages, list): + raise web.HTTPBadRequest(text="messages must be a list") + sid = _request_session_id(request, body) adapter = request.app[ADAPTER_KEY] if sid in adapter.closed: - raise web.HTTPServiceUnavailable(text="session closed") + return web.Response(status=503, text="session closed") + + # Per-sid turn cap (HTTP 429). Mirrors anthropic.py:517-530. + cap = adapter.max_turns_per_sid + if cap is not None: + prior = adapter._sid_turn_count.get(sid, 0) + if prior >= cap: + return web.json_response( + { + "error": { + "type": "rate_limit_error", + "message": (f"adapter: sid {sid!r} exceeded max_turns_per_sid={cap}; killing run"), + } + }, + status=429, + ) + adapter._sid_turn_count[sid] = prior + 1 + app = request.app tok = app[TOKENIZER_KEY] s = adapter.store.setdefault(sid, Session()) task = asyncio.current_task() adapter.inflight.setdefault(sid, set()).add(task) try: - async with s.lock: - translated = _translate_chat_messages(messages) - tools_schema = _normalize_tools(body.get("tools")) - full_prompt_ids = render_prompt(s.traj, translated, tok, tools_schema) - gen = await _generate(full_prompt_ids, s, body, app, session_id=sid) - turns, pending_key = assemble_turns(s.traj, translated, tok, tools_schema, gen, full_prompt_ids) - node = record_turn(s.traj.tree, turns) - if node is not None: - s.traj.resp_truth[pending_key] = ( - list(gen.output_ids), - list(gen.output_log_probs), - gen.output_text, + async with s.lock: # same sid -> serialized + translated = _translate_openai_chat(messages) + tools_schema = _openai_tools_to_chat_tools(body.get("tools")) + prompt_ids = _render_token_ids(translated, tok, tools=tools_schema, add_generation_prompt=True) + + turn = await call_sglang_generate( + prompt_ids, + s, + body, + app, + max_token_keys=("max_completion_tokens", "max_tokens", "max_output_tokens"), + stop_keys=("stop",), + log_prefix="openai_adapter", + logger=logger, + session_id=sid, + ) + + raw_output = tok.decode(turn.output_ids, skip_special_tokens=False) if turn.output_ids else "" + parsed = parse_model_output( + raw_output, + tools_schema=tools_schema, + tool_parser_name=app[TOOL_PARSER_KEY], + reasoning_parser_name=app[REASONING_PARSER_KEY], + ) + wire_message, manager_message, wire_finish = _build_oai_response(parsed, turn.finish_reason) + + output_ids = list(turn.output_ids) + finish_reason_mgr = _finish_reason_for_manager(turn.finish_reason, parsed.tool_uses) + + try: + adapter.manager.append_turn( + sid, + prompt_messages=translated, + tools=tools_schema, + prompt_ids=prompt_ids, + response_ids=output_ids, + response_logprobs=( + list(turn.output_log_probs) + if turn.output_log_probs and len(turn.output_log_probs) == len(turn.output_ids) + else None + ), + response_message=manager_message, + finish_reason=finish_reason_mgr, + metadata={"sid": sid}, ) - parsed = _parse_output(gen.output_text, tools_schema, app) - return gen, parsed, len(full_prompt_ids), len(gen.output_ids) + except Exception: + logger.exception("append_turn(sid=%s) failed", sid) + + hook = adapter.on_turn_appended + if hook is not None: + try: + hook( + sid, + translated, + tools_schema, + manager_message, + prompt_ids, + output_ids, + finish_reason_mgr, + ) + except Exception: + logger.exception("on_turn_appended hook failed (sid=%s)", sid) + + in_tok, out_tok = len(prompt_ids), len(turn.output_ids) + + if body.get("stream") is True or "text/event-stream" in request.headers.get("Accept", ""): + return await _stream_chat_completion(request, body, wire_message, wire_finish, in_tok, out_tok) + return web.json_response(_chat_completion_response(body, wire_message, wire_finish, in_tok, out_tok)) finally: adapter.inflight.get(sid, set()).discard(task) -async def _handle_chat_completions(request: web.Request) -> web.StreamResponse: - body = await request.json() - messages = body.get("messages") or [] - if not isinstance(messages, list): - raise web.HTTPBadRequest(text="messages must be a list") - gen, parsed, in_tok, out_tok = await _run_turn(request, body, messages) - if body.get("stream"): - return await _stream_chat_completion(request, body, parsed, gen.finish_reason, in_tok, out_tok) - return web.json_response(_chat_completion_response(body, parsed, gen.finish_reason, in_tok, out_tok)) - - def _chat_completion_response( body: dict, - parsed: ParsedModelOutput, - finish: str, + wire_message: dict[str, Any], + wire_finish: str, in_tok: int, out_tok: int, ) -> dict[str, Any]: @@ -365,8 +555,8 @@ def _chat_completion_response( "choices": [ { "index": 0, - "message": _chat_message(parsed), - "finish_reason": _finish_reason(parsed, finish), + "message": wire_message, + "finish_reason": wire_finish, } ], "usage": _usage(in_tok, out_tok), @@ -376,11 +566,18 @@ def _chat_completion_response( async def _stream_chat_completion( request: web.Request, body: dict, - parsed: ParsedModelOutput, - finish: str, + wire_message: dict[str, Any], + wire_finish: str, in_tok: int, out_tok: int, ) -> web.StreamResponse: + """Emit the OpenAI Chat-Completions SSE stream. + + Each chunk shape: ``data: {chatcmpl ...}\n\n`` ending with ``data: [DONE]``. + The whole turn is fully realised on the server before we start streaming + (we don't have token-level deltas from SGLang here), so we emit one role + chunk, then content / reasoning / tool_calls in single delta chunks each. + """ out = web.StreamResponse( status=200, headers={ @@ -393,131 +590,38 @@ async def _stream_chat_completion( completion_id = f"chatcmpl_{secrets.token_hex(12)}" created = int(time.time()) - async def emit(choice_delta: dict[str, Any], finish_reason: str | None = None, usage: dict | None = None) -> None: + async def emit(delta: dict[str, Any], finish_reason: str | None = None, usage: dict | None = None) -> None: chunk = { "id": completion_id, "object": "chat.completion.chunk", "created": created, "model": body.get("model", "slime-actor"), - "choices": [{"index": 0, "delta": choice_delta, "finish_reason": finish_reason}], + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], } if usage is not None: chunk["usage"] = usage await out.write(f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n".encode()) await emit({"role": "assistant"}) - if parsed.reasoning: - await emit({"reasoning_content": parsed.reasoning}) - if parsed.text: - await emit({"content": parsed.text}) - for idx, call in enumerate(_openai_tool_calls(parsed.tool_uses)): - await emit({"tool_calls": [{**call, "index": idx}]}) - await emit({}, finish_reason=_finish_reason(parsed, finish), usage=_usage(in_tok, out_tok)) + reasoning = wire_message.get("reasoning_content") + if reasoning: + await emit({"reasoning_content": reasoning}) + content = wire_message.get("content") + if content: + await emit({"content": content}) + # NB: emit ALL tool_calls in a single chunk. Codex CLI 0.30.0 incorrectly + # accumulates per-index ``arguments`` fragments across chunks, causing N + # parallel tool_calls to collapse into a single call with a concatenated + # arguments string (``{"command": "..."}{"command": "..."}``) -- this then + # fails Codex's own arguments parser and the tool result becomes a parse + # error, breaking node_match_key alignment on the next turn. + tool_calls = wire_message.get("tool_calls") or [] + if tool_calls: + await emit({"tool_calls": [{**call, "index": idx} for idx, call in enumerate(tool_calls)]}) + await emit({}, finish_reason=wire_finish, usage=_usage(in_tok, out_tok)) await out.write(b"data: [DONE]\n\n") return out -async def _handle_responses(request: web.Request) -> web.StreamResponse: - body = await request.json() - messages = _responses_input_to_messages(body.get("input", ""), body.get("instructions")) - gen, parsed, in_tok, out_tok = await _run_turn(request, body, messages) - if body.get("stream"): - return await _stream_response(request, body, parsed, gen.finish_reason, in_tok, out_tok) - return web.json_response(_response_response(body, parsed, gen.finish_reason, in_tok, out_tok)) - - -def _response_output(parsed: ParsedModelOutput) -> list[dict[str, Any]]: - output: list[dict[str, Any]] = [] - if parsed.reasoning: - output.append( - { - "id": f"rs_{secrets.token_hex(12)}", - "type": "reasoning", - "summary": [{"type": "summary_text", "text": parsed.reasoning}], - } - ) - if parsed.text: - output.append( - { - "id": f"msg_{secrets.token_hex(12)}", - "type": "message", - "role": "assistant", - "status": "completed", - "content": [{"type": "output_text", "text": parsed.text, "annotations": []}], - } - ) - for call in _openai_tool_calls(parsed.tool_uses): - output.append( - { - "id": f"fc_{secrets.token_hex(12)}", - "type": "function_call", - "status": "completed", - "call_id": call["id"], - "name": call["function"]["name"], - "arguments": call["function"]["arguments"], - } - ) - if not output: - output.append( - { - "id": f"msg_{secrets.token_hex(12)}", - "type": "message", - "role": "assistant", - "status": "completed", - "content": [{"type": "output_text", "text": "", "annotations": []}], - } - ) - return output - - -def _response_response( - body: dict, - parsed: ParsedModelOutput, - finish: str, - in_tok: int, - out_tok: int, -) -> dict[str, Any]: - status = "incomplete" if finish == "length" else "completed" - return { - "id": f"resp_{secrets.token_hex(12)}", - "object": "response", - "created_at": int(time.time()), - "status": status, - "model": body.get("model", "slime-actor"), - "output": _response_output(parsed), - "usage": _responses_usage(in_tok, out_tok), - } - - -async def _stream_response( - request: web.Request, - body: dict, - parsed: ParsedModelOutput, - finish: str, - in_tok: int, - out_tok: int, -) -> web.StreamResponse: - out = web.StreamResponse( - status=200, - headers={ - "Content-Type": "text/event-stream", - "Cache-Control": "no-cache", - "Connection": "keep-alive", - }, - ) - await out.prepare(request) - response = _response_response(body, parsed, finish, in_tok, out_tok) - created = {"type": "response.created", "response": response} - await out.write(f"event: response.created\ndata: {json.dumps(created, ensure_ascii=False)}\n\n".encode()) - if parsed.text: - delta = {"type": "response.output_text.delta", "delta": parsed.text} - await out.write( - f"event: response.output_text.delta\ndata: {json.dumps(delta, ensure_ascii=False)}\n\n".encode() - ) - completed = {"type": "response.completed", "response": response} - await out.write(f"event: response.completed\ndata: {json.dumps(completed, ensure_ascii=False)}\n\n".encode()) - return out - - async def _ok(request: web.Request) -> web.Response: return await ok_response(request) diff --git a/slime/agent/trajectory_manager.py b/slime/agent/trajectory_manager.py index 8a5937176a..fd9f50b029 100644 --- a/slime/agent/trajectory_manager.py +++ b/slime/agent/trajectory_manager.py @@ -32,7 +32,7 @@ Sample alongside the main leaf. Snapshot tokens = cumulative pre-drop; snapshot loss_mask is COMPLEMENTARY — 1 only at positions that the main leaf is about to drop, 0 elsewhere. Snapshot reward = main-leaf - share; snapshot rollout_id = main-leaf rollout_id. Snapshot ∪ main on + share; snapshot group_id = main-leaf group_id. Snapshot ∪ main on loss_mask=1 tokens never overlap and their union equals the virtual no-drift trajectory. The snapshotted drift is NOT counted in the main sample's ``tito_dropped_*`` (it wasn't truly lost). @@ -435,9 +435,7 @@ def get_trajectory( samples.append( Sample( index=base_sample.index, - rollout_id=( - base_sample.rollout_id if base_sample.rollout_id is not None else base_sample.index - ), + group_id=(base_sample.group_id if base_sample.group_id is not None else base_sample.index), prompt=base_sample.prompt, label=base_sample.label, tokens=snap_tokens, @@ -463,7 +461,7 @@ def get_trajectory( samples.append( Sample( index=base_sample.index, - rollout_id=(base_sample.rollout_id if base_sample.rollout_id is not None else base_sample.index), + group_id=(base_sample.group_id if base_sample.group_id is not None else base_sample.index), prompt=base_sample.prompt, label=base_sample.label, tokens=tokens, From 44e7999dfa9beea3000ea0bc62ccbf236f9573df Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Fri, 5 Jun 2026 03:02:55 +0000 Subject: [PATCH 05/28] refactor(agent): centralize snapshot-threshold default + filter access log MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * TrajectoryManager owns the snapshot threshold default (1024) — drop None-passthrough from AnthropicAdapter and the hardcoded 1000 in examples/coding_agent_rl/generate.py so the single source of truth holds. * TrajectoryManager.__init__: remove dead kwargs (tokenizer, chat_template_kwargs, end_of_turn_token_id) — none were read since plan C. * FilteredAccessLogger drops HEAD heartbeats and only emits when status != 200 or elapsed > 120s — kills the web_log.py:232 spam without silencing real errors / slow handlers. --- examples/coding_agent_rl/aiohttp_threaded.py | 12 ++++++ examples/coding_agent_rl/generate.py | 31 ++++++++++----- slime/agent/adapters/anthropic.py | 10 +++-- slime/agent/trajectory_manager.py | 42 +++++++++++--------- 4 files changed, 62 insertions(+), 33 deletions(-) diff --git a/examples/coding_agent_rl/aiohttp_threaded.py b/examples/coding_agent_rl/aiohttp_threaded.py index a5a17652d4..4fe80e863a 100644 --- a/examples/coding_agent_rl/aiohttp_threaded.py +++ b/examples/coding_agent_rl/aiohttp_threaded.py @@ -8,6 +8,18 @@ from typing import Any from aiohttp import web +from aiohttp.web_log import AccessLogger + + +class FilteredAccessLogger(AccessLogger): + SLOW_THRESHOLD_SEC = 120.0 + + def log(self, request, response, time): + if request.method == "HEAD": + return + if response.status == 200 and time <= self.SLOW_THRESHOLD_SEC: + return + super().log(request, response, time) @dataclass diff --git a/examples/coding_agent_rl/generate.py b/examples/coding_agent_rl/generate.py index a858eda060..6d7ad5bab8 100644 --- a/examples/coding_agent_rl/generate.py +++ b/examples/coding_agent_rl/generate.py @@ -59,7 +59,7 @@ from slime.utils.types import Sample from . import sandbox -from .aiohttp_threaded import run_app_in_thread +from .aiohttp_threaded import FilteredAccessLogger, run_app_in_thread logger = logging.getLogger(__name__) @@ -97,16 +97,22 @@ def __init__(self, args) -> None: "Without it the sandbox cannot dial back and the rollout will " "silently abort." ) - # Snapshot threshold: 0 disables; absent or malformed env => default 1000. + # Snapshot threshold: env override only. Absent => let TrajectoryManager + # take its built-in default. ``0`` disables, malformed values are + # ignored with a warning. _snap_env = os.environ.get("SLIME_TITO_SNAPSHOT_MIN_LOSS_TOKENS") - try: - _snap_threshold = int(_snap_env) if _snap_env is not None else 1000 - except ValueError: - logger.warning( - "SLIME_TITO_SNAPSHOT_MIN_LOSS_TOKENS=%r is not an int; using 1000", - _snap_env, - ) - _snap_threshold = 1000 + _snap_threshold: int | None + if _snap_env is None: + _snap_threshold = None + else: + try: + _snap_threshold = int(_snap_env) + except ValueError: + logger.warning( + "SLIME_TITO_SNAPSHOT_MIN_LOSS_TOKENS=%r is not an int; falling back to TrajectoryManager default", + _snap_env, + ) + _snap_threshold = None self.adapter = AnthropicAdapter( tokenizer=self.tokenizer, sglang_url=sglang_url, @@ -124,7 +130,10 @@ def __init__(self, args) -> None: host=SHIM_BIND_HOST, port=SHIM_PORT, thread_name="anthropic-adapter", - runner_kwargs={"handler_cancellation": True}, + runner_kwargs={ + "handler_cancellation": True, + "access_log_class": FilteredAccessLogger, + }, ) self.adapter_url = f"http://{public_host}:{self.app_handle.port}" logger.info( diff --git a/slime/agent/adapters/anthropic.py b/slime/agent/adapters/anthropic.py index 6a36abf5fa..cb5f20916c 100644 --- a/slime/agent/adapters/anthropic.py +++ b/slime/agent/adapters/anthropic.py @@ -78,10 +78,12 @@ def __init__( reasoning_parser=reasoning_parser, ) # ONE manager shared across all sids; per-sid trees live inside. - self.manager = TrajectoryManager( - tokenizer=tokenizer, - tito_snapshot_min_loss_tokens=tito_snapshot_min_loss_tokens, - ) + # ``None`` here means "caller did not specify" → let TrajectoryManager's + # own default take over. Pass an int (incl. 0 to disable) to override. + mgr_kwargs: dict[str, int] = {} + if tito_snapshot_min_loss_tokens is not None: + mgr_kwargs["tito_snapshot_min_loss_tokens"] = tito_snapshot_min_loss_tokens + self.manager = TrajectoryManager(**mgr_kwargs) # Optional debug hook invoked after each successful append_turn. # Signature: (sid, prompt_messages, tools, response_message, # prompt_ids, response_ids, finish_reason) -> None. diff --git a/slime/agent/trajectory_manager.py b/slime/agent/trajectory_manager.py index fd9f50b029..8019a64290 100644 --- a/slime/agent/trajectory_manager.py +++ b/slime/agent/trajectory_manager.py @@ -205,20 +205,12 @@ class TrajectoryManager: def __init__( self, *, - tokenizer=None, - chat_template_kwargs: dict[str, Any] | None = None, - end_of_turn_token_id: int | None = None, - tito_snapshot_min_loss_tokens: int | None = None, + tito_snapshot_min_loss_tokens: int | None = 1024, ) -> None: - # tokenizer / chat_template_kwargs are no longer load-bearing under - # plan C, but the constructor signature is kept for callsite - # compatibility. _tokenizer is retained for forward-compat (callers - # constructing TrajectoryManager(tokenizer=tok) shouldn't break). - self._tokenizer = tokenizer - self._ct_kwargs: dict[str, Any] = dict(chat_template_kwargs or {}) - self._end_of_turn_token_id = end_of_turn_token_id # Drift-snapshot threshold (loss_mask=1 token count inside drift suffix). # None or <= 0 disables; behavior then matches the pre-feature output. + # Default 1024 trades a small per-trajectory snapshot overhead for not + # silently dropping >1k loss tokens when TITO drift hits. self._snap_threshold: int | None = ( tito_snapshot_min_loss_tokens if (tito_snapshot_min_loss_tokens is not None and tito_snapshot_min_loss_tokens > 0) @@ -419,18 +411,29 @@ def get_trajectory( } per_leaf_reward = (reward / len(leaves)) if leaves else 0.0 + # slime contract (see slime/backends/megatron_utils/data.py:139, + # slime/ray/rollout.py:695): ``loss_mask`` and ``rollout_log_probs`` + # cover only the response region — i.e. tokens AFTER the initial + # prompt — and ``response_length == len(loss_mask)``. Strip the + # first turn's prompt prefix here so all downstream consumers see + # tokens/loss_mask/logprobs in their canonical alignment. + first_prompt_len = len(asst_chain[0].turn_prompt_ids or []) if asst_chain else 0 + # Emit snapshot sample(s) first, then the main-leaf sample. for snap_tokens, snap_mask, snap_lp, drift_turn, cur_chain_idx in snapshots: snap_finish = None prev_idx = cur_chain_idx - 1 # asst_chain index of the previous (prefix's last) turn if 0 <= prev_idx < len(asst_chain): snap_finish = asst_chain[prev_idx].turn_finish_reason + snap_strip = min(first_prompt_len, len(snap_mask)) + snap_mask_resp = snap_mask[snap_strip:] + snap_lp_resp = snap_lp[snap_strip:] snap_md = { **base_md, "finish_reason": snap_finish, "tito_snapshot": True, "tito_snapshot_at_turn": drift_turn, - "tito_snapshot_loss_tokens": sum(snap_mask), + "tito_snapshot_loss_tokens": sum(snap_mask_resp), } samples.append( Sample( @@ -439,16 +442,19 @@ def get_trajectory( prompt=base_sample.prompt, label=base_sample.label, tokens=snap_tokens, - response_length=sum(1 for m in snap_mask if m == 1), - loss_mask=snap_mask, - rollout_log_probs=snap_lp, + response_length=len(snap_mask_resp), + loss_mask=snap_mask_resp, + rollout_log_probs=snap_lp_resp, reward=per_leaf_reward, status=Sample.Status.COMPLETED, metadata=snap_md, ) ) - response_length = sum(1 for m in loss_mask if m == 1) + main_strip = min(first_prompt_len, len(loss_mask)) + loss_mask_resp = loss_mask[main_strip:] + logprobs_resp = logprobs[main_strip:] + response_length = len(loss_mask_resp) main_md: dict[str, Any] = { **base_md, "finish_reason": last_asst.turn_finish_reason if last_asst else None, @@ -466,8 +472,8 @@ def get_trajectory( label=base_sample.label, tokens=tokens, response_length=response_length, - loss_mask=loss_mask, - rollout_log_probs=logprobs, + loss_mask=loss_mask_resp, + rollout_log_probs=logprobs_resp, reward=per_leaf_reward, status=Sample.Status.COMPLETED, metadata=main_md, From 5e59c254e1ef6805db68a5bc239c226956d1d123 Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Fri, 5 Jun 2026 07:50:19 +0000 Subject: [PATCH 06/28] feat(agent): add fork-merge rescue for short assistant rewrites When claude-code replays a session and reformats a prior assistant message (tool_call arg ordering, whitespace), the DFS breaks at that assistant group and every reformat would spawn a new sibling subtree. Opt-in via fork_merge_max_response_tokens: if exactly one leaf assistant sibling has turn_response_ids length < threshold, collapse onto it and mark it loss_mask=0 at linearization. Sample metadata records fork_merge_masked_tokens / fork_merge_turns; a warning logs each merge. - TrajectoryManager: __init__ kwarg, Step 1.5 in append_turn, mask=0 emit in get_trajectory; revert tito_snapshot_min_loss_tokens default back to None to keep the opt-in contract. - AnthropicAdapter / OpenAIAdapter: pass-through kwarg (only forwarded when non-None); fix OpenAIAdapter erroneously passing tokenizer= to TrajectoryManager. - examples/coding_agent_rl/generate.py: parse SLIME_FORK_MERGE_MAX_RESPONSE_TOKENS env var. E2E on 20 SWE tasks with threshold=1024: 5 rewrites merged (3164 masked tokens), asst-role forks 15->6 vs no-rescue baseline. --- examples/coding_agent_rl/generate.py | 18 +++++ slime/agent/adapters/anthropic.py | 3 + slime/agent/adapters/openai.py | 13 +++- slime/agent/trajectory_manager.py | 105 ++++++++++++++++++++++++++- 4 files changed, 131 insertions(+), 8 deletions(-) diff --git a/examples/coding_agent_rl/generate.py b/examples/coding_agent_rl/generate.py index 6d7ad5bab8..e4dddad705 100644 --- a/examples/coding_agent_rl/generate.py +++ b/examples/coding_agent_rl/generate.py @@ -113,12 +113,30 @@ def __init__(self, args) -> None: _snap_env, ) _snap_threshold = None + # Fork-merge threshold: same env-override convention. Triggers the + # rescue path in trajectory_manager when claude-code reformats a prior + # assistant message and the existing sibling's response is short + # enough to mask out instead of forking the tree. + _fork_merge_env = os.environ.get("SLIME_FORK_MERGE_MAX_RESPONSE_TOKENS") + _fork_merge_threshold: int | None + if _fork_merge_env is None: + _fork_merge_threshold = None + else: + try: + _fork_merge_threshold = int(_fork_merge_env) + except ValueError: + logger.warning( + "SLIME_FORK_MERGE_MAX_RESPONSE_TOKENS=%r is not an int; falling back to TrajectoryManager default", + _fork_merge_env, + ) + _fork_merge_threshold = None self.adapter = AnthropicAdapter( tokenizer=self.tokenizer, sglang_url=sglang_url, tool_parser=self.tool_parser, reasoning_parser=self.reasoning_parser, tito_snapshot_min_loss_tokens=_snap_threshold, + fork_merge_max_response_tokens=_fork_merge_threshold, ) # handler_cancellation=True so a client disconnect cancels the handler # coroutine, arming the fire-and-forget /abort_request inside the diff --git a/slime/agent/adapters/anthropic.py b/slime/agent/adapters/anthropic.py index cb5f20916c..995fc5b7d5 100644 --- a/slime/agent/adapters/anthropic.py +++ b/slime/agent/adapters/anthropic.py @@ -68,6 +68,7 @@ def __init__( tool_parser=None, reasoning_parser=None, tito_snapshot_min_loss_tokens: int | None = None, + fork_merge_max_response_tokens: int | None = None, max_turns_per_sid: int | None = None, on_turn_appended: Callable[..., None] | None = None, ) -> None: @@ -83,6 +84,8 @@ def __init__( mgr_kwargs: dict[str, int] = {} if tito_snapshot_min_loss_tokens is not None: mgr_kwargs["tito_snapshot_min_loss_tokens"] = tito_snapshot_min_loss_tokens + if fork_merge_max_response_tokens is not None: + mgr_kwargs["fork_merge_max_response_tokens"] = fork_merge_max_response_tokens self.manager = TrajectoryManager(**mgr_kwargs) # Optional debug hook invoked after each successful append_turn. # Signature: (sid, prompt_messages, tools, response_message, diff --git a/slime/agent/adapters/openai.py b/slime/agent/adapters/openai.py index 8697c09639..cea3e5d858 100644 --- a/slime/agent/adapters/openai.py +++ b/slime/agent/adapters/openai.py @@ -74,6 +74,7 @@ def __init__( tool_parser=None, reasoning_parser=None, tito_snapshot_min_loss_tokens: int | None = None, + fork_merge_max_response_tokens: int | None = None, max_turns_per_sid: int | None = None, on_turn_appended: Callable[..., None] | None = None, ) -> None: @@ -84,10 +85,14 @@ def __init__( reasoning_parser=reasoning_parser, ) # ONE manager shared across all sids; per-sid trees live inside. - self.manager = TrajectoryManager( - tokenizer=tokenizer, - tito_snapshot_min_loss_tokens=tito_snapshot_min_loss_tokens, - ) + # Mirror AnthropicAdapter: only forward kwargs the caller actually + # specified, so TrajectoryManager's own defaults stay authoritative. + mgr_kwargs: dict[str, int] = {} + if tito_snapshot_min_loss_tokens is not None: + mgr_kwargs["tito_snapshot_min_loss_tokens"] = tito_snapshot_min_loss_tokens + if fork_merge_max_response_tokens is not None: + mgr_kwargs["fork_merge_max_response_tokens"] = fork_merge_max_response_tokens + self.manager = TrajectoryManager(**mgr_kwargs) # Optional debug hook invoked after each successful append_turn. # Signature mirrors AnthropicAdapter.on_turn_appended: # (sid, prompt_messages, tools, response_message, diff --git a/slime/agent/trajectory_manager.py b/slime/agent/trajectory_manager.py index 8019a64290..87c229ad30 100644 --- a/slime/agent/trajectory_manager.py +++ b/slime/agent/trajectory_manager.py @@ -189,6 +189,25 @@ def _lcp_len(a: list[int], b: list[int]) -> int: return i +def _short_text_preview(messages: list[dict[str, Any]], *, limit: int) -> str: + """Compact text preview of an assistant message block for log lines. + + Extracts string ``content`` and any ``{"text": "..."}`` content blocks + (anthropic-style); falls back to the empty string if nothing parseable. + """ + parts: list[str] = [] + for m in messages: + c = m.get("content") + if isinstance(c, str): + parts.append(c) + elif isinstance(c, list): + for blk in c: + if isinstance(blk, dict) and isinstance(blk.get("text"), str): + parts.append(blk["text"]) + s = " ".join(parts).strip() + return s[:limit] + ("…" if len(s) > limit else "") + + # =========================================================================== # TrajectoryManager # =========================================================================== @@ -205,17 +224,26 @@ class TrajectoryManager: def __init__( self, *, - tito_snapshot_min_loss_tokens: int | None = 1024, + tito_snapshot_min_loss_tokens: int | None = None, + fork_merge_max_response_tokens: int | None = None, ) -> None: # Drift-snapshot threshold (loss_mask=1 token count inside drift suffix). # None or <= 0 disables; behavior then matches the pre-feature output. - # Default 1024 trades a small per-trajectory snapshot overhead for not - # silently dropping >1k loss tokens when TITO drift hits. self._snap_threshold: int | None = ( tito_snapshot_min_loss_tokens if (tito_snapshot_min_loss_tokens is not None and tito_snapshot_min_loss_tokens > 0) else None ) + # Fork-merge threshold: when DFS would break at an assistant group and + # exactly one non-leaf assistant sibling has turn_response_ids length + # STRICTLY LESS than this value, collapse the would-be fork onto that + # sibling (its response then enters trajectories with loss_mask=0). + # None or <= 0 disables; behavior matches pre-feature output. + self._fork_merge_threshold: int | None = ( + fork_merge_max_response_tokens + if (fork_merge_max_response_tokens is not None and fork_merge_max_response_tokens > 0) + else None + ) self._trees: dict[str, Node] = {} self._turn_count: dict[str, int] = {} @@ -270,6 +298,62 @@ def append_turn( cur = match i += 1 + # Step 1.5: assistant fork-merge rescue (opt-in via + # fork_merge_max_response_tokens). The typical claude-code pattern is: + # a later replay reformats an earlier assistant message (e.g. tool_call + # arg ordering, whitespace), which breaks the DFS at that assistant + # group. Without rescue, every such reformat spawns a new sibling + # subtree; with rescue, when the existing sibling's per-turn response + # is short enough that masking it out is the cheaper trade-off, we + # collapse onto that sibling and mark it for mask=0 at linearization. + if i < len(groups) and self._fork_merge_threshold is not None and groups[i].role == "assistant": + candidates = [ + c + for c in cur.children + if c.role == "assistant" + # Real rewrite footprint = the original turn's leaf assistant + # node: it was inserted by Step-3 of a prior turn (so + # turn_response_ids is populated), and no later turn has + # extended it (so it is still a leaf). Once a rewrite collapses + # onto it via this rescue, the new turn's user/tool + asst leaf + # are appended underneath as children — so the merge target + # MUST be a leaf at decision time, otherwise it has already + # diverged into mixed subchains and merging would tangle them. + and not c.children + and c.turn_response_ids is not None + and len(c.turn_response_ids) < self._fork_merge_threshold + ] + if len(candidates) == 1: + sib = candidates[0] + masked = len(sib.turn_response_ids or []) + preview = _short_text_preview(sib.messages, limit=160) + logger.warning( + "append_turn(sid=%s turn=%s): fork-merging assistant rewrite " + "into existing sibling (turn_index=%s, masked_response_tokens=%d, " + "sibling_response_preview=%r)", + sid, + self._turn_count.get(sid, 0) + 1, + sib.turn_index, + masked, + preview, + ) + sib.metadata["fork_merged"] = True + sib.metadata["fork_merge_masked_tokens"] = masked + cur = sib + i += 1 + elif len(candidates) >= 2: + # Legacy / pathological state: fork_merge wasn't on during + # earlier rewrites, or threshold was widened mid-session. + # Don't pick arbitrarily — fork as usual and surface a hint. + logger.warning( + "append_turn(sid=%s turn=%s): multiple eligible fork-merge " + "candidates (%d), refusing to merge — likely a legacy state " + "from a prior run without fork_merge enabled; forking instead.", + sid, + self._turn_count.get(sid, 0) + 1, + len(candidates), + ) + # Step 2: mount remaining prompt groups as plain routing nodes. # Token attribution happens at get_trajectory time, not here. for g in groups[i:]: @@ -350,12 +434,15 @@ def get_trajectory( logprobs: list[float] = [] total_dropped = 0 dropped_turns = 0 + fork_merge_masked_total = 0 + fork_merge_turns_total = 0 snapshots: list[tuple[list[int], list[int], list[float], int, int]] = [] for k, asst in enumerate(asst_chain, start=1): p = list(asst.turn_prompt_ids or []) r = list(asst.turn_response_ids or []) lp = list(asst.turn_response_logprobs) if asst.turn_response_logprobs is not None else None + is_merged = bool(asst.metadata.get("fork_merged")) if k == 1: emit_prompt = p @@ -395,12 +482,19 @@ def get_trajectory( logprobs.extend([0.0] * len(emit_prompt)) tokens.extend(r) - loss_mask.extend([1] * len(r)) + # fork-merged sibling: its response is "stale" — present in the + # tree only as a routing placeholder for the rewrites that + # collapsed onto it; mask it out of training. + loss_mask.extend([0 if is_merged else 1] * len(r)) if lp is not None: logprobs.extend(lp) else: logprobs.extend([0.0] * len(r)) + if is_merged: + fork_merge_masked_total += len(r) + fork_merge_turns_total += 1 + last_asst = asst_chain[-1] if asst_chain else None first_sys = next((n for n in chain if n.role == "system"), None) tools_meta = first_sys.metadata.get("tools") if first_sys else None @@ -464,6 +558,9 @@ def get_trajectory( main_md["tito_dropped_turns"] = dropped_turns if snapshots: main_md["tito_snapshots_emitted"] = len(snapshots) + if fork_merge_masked_total > 0: + main_md["fork_merge_masked_tokens"] = fork_merge_masked_total + main_md["fork_merge_turns"] = fork_merge_turns_total samples.append( Sample( index=base_sample.index, From 1d459bde6879d22d9406265003b713b7b37628cd Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Fri, 5 Jun 2026 09:12:16 +0000 Subject: [PATCH 07/28] fix(agent): replace sib.messages on fork-merge rescue MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rescue branch was merging the rewritten turn into the sibling node's metadata but leaving sib.messages as the pre-rewrite payload. The subsequent turn replays the rewritten payload in its prompt history, DFS-fails to match the (unchanged) sibling, falls through Step 1.5 (sibling is no longer a leaf since the new turn child attached), and forks anyway — defeating the rescue. Update sib.messages to the rewritten version at rescue time. The per-turn sglang snapshot (turn_response_ids/logprobs/turn_index) stays on the original node, and get_trajectory still emits it with loss_mask=0 via the fork_merged flag. Validated end-to-end on a 20-instance SWE batch: tool→2×assistant forks dropped 6 → 0; total forks 27 → 18. --- slime/agent/trajectory_manager.py | 1 + 1 file changed, 1 insertion(+) diff --git a/slime/agent/trajectory_manager.py b/slime/agent/trajectory_manager.py index 87c229ad30..bf8714c2f1 100644 --- a/slime/agent/trajectory_manager.py +++ b/slime/agent/trajectory_manager.py @@ -339,6 +339,7 @@ def append_turn( ) sib.metadata["fork_merged"] = True sib.metadata["fork_merge_masked_tokens"] = masked + sib.messages = list(groups[i].messages) cur = sib i += 1 elif len(candidates) >= 2: From 8d6fdc939aa06c21bb662e93198dd5179c9e7278 Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Fri, 5 Jun 2026 09:30:23 +0000 Subject: [PATCH 08/28] refactor(agent): drop billing-header scrub now cc emits no header CLAUDE_CODE_ATTRIBUTION_HEADER=0 (set in examples/coding_agent_rl/sandbox.py and the e2e test runner) tells claude-code to suppress the ``x-anthropic-billing-header: cc_version=...; cch=...;`` block it otherwise prepends to the system prompt. Verified on a 56-turn e2e batch: zero requests contained the header, no scrub mutations fired. Remove _scrub_claude_code_billing_header_in_body, its regex, the call site, and the now-unused `re` import. --- slime/agent/adapters/anthropic.py | 56 ++----------------------------- 1 file changed, 2 insertions(+), 54 deletions(-) diff --git a/slime/agent/adapters/anthropic.py b/slime/agent/adapters/anthropic.py index 995fc5b7d5..48968a500b 100644 --- a/slime/agent/adapters/anthropic.py +++ b/slime/agent/adapters/anthropic.py @@ -18,7 +18,6 @@ import dataclasses import json import logging -import re import secrets from collections.abc import Callable from typing import Any @@ -134,54 +133,6 @@ async def finish_session( # ============================================================================= -# Claude Code CLI leaks ``x-anthropic-billing-header: ...cch=;`` as a text -# block at the top of the system prompt. The cch hash changes per request, so -# without stripping it the rendered system tokens differ every turn and the -# manager tree can't chain consecutive turns together. -_CLAUDE_CODE_BILLING_HEADER_RE = re.compile( - r"^\s*x-anthropic-billing-header:[^\n]*\n?", - re.IGNORECASE, -) - - -def _scrub_claude_code_billing_header_in_body(body_obj: dict) -> bool: - """Strip Claude Code's billing-header sidechannel from ``body['system']``. - - Handles both Anthropic shapes (``system: str`` and - ``system: list[{type:"text",text:"..."}]``). Mutates ``body_obj`` in - place; returns True iff anything changed. - """ - sysm = body_obj.get("system") - changed = False - if isinstance(sysm, str): - cleaned = _CLAUDE_CODE_BILLING_HEADER_RE.sub("", sysm) - if cleaned != sysm: - body_obj["system"] = cleaned if cleaned.strip() else "" - changed = True - elif isinstance(sysm, list): - new_blocks: list = [] - for block in sysm: - if not isinstance(block, dict) or block.get("type") != "text": - new_blocks.append(block) - continue - txt = block.get("text") or "" - cleaned = _CLAUDE_CODE_BILLING_HEADER_RE.sub("", txt) - if not cleaned.strip(): - # Whole block was the sidechannel — drop it. - changed = True - continue - if cleaned != txt: - new_block = dict(block) - new_block["text"] = cleaned - new_blocks.append(new_block) - changed = True - else: - new_blocks.append(block) - if changed: - body_obj["system"] = new_blocks - return changed - - _MID_SYSTEM_WRAP_PREFIX = "\n" _MID_SYSTEM_WRAP_SUFFIX = "\n\n" @@ -534,11 +485,8 @@ async def _handle_request(request: web.Request) -> web.StreamResponse: ) adapter._sid_turn_count[sid] = prior + 1 - # Strip Claude Code's per-request billing-header sidechannel BEFORE the - # adapter renders prompt_ids. Also fold mid-list ``role: system`` messages - # into a neighbouring user message so Qwen3 chat templates accept them. - # Both are no-ops when the relevant patterns aren't present. - _scrub_claude_code_billing_header_in_body(body) + # Fold mid-list ``role: system`` messages into a neighbouring user message + # so Qwen3 chat templates accept them. No-op when no mid-list system msgs. _fold_mid_list_system_into_user(body) app = request.app From 2b4efc42580718906327138a0518be93e5d37d43 Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Mon, 8 Jun 2026 06:13:12 +0000 Subject: [PATCH 09/28] refactor(agent): migrate TrajectoryManager and adapters (v4) --- examples/coding_agent_rl/generate.py | 41 +- slime/agent/adapters/anthropic.py | 20 +- slime/agent/adapters/openai.py | 18 +- slime/agent/trajectory_manager.py | 746 ++++++++++++++++----------- 4 files changed, 455 insertions(+), 370 deletions(-) diff --git a/examples/coding_agent_rl/generate.py b/examples/coding_agent_rl/generate.py index e4dddad705..d7dd6a1616 100644 --- a/examples/coding_agent_rl/generate.py +++ b/examples/coding_agent_rl/generate.py @@ -97,46 +97,17 @@ def __init__(self, args) -> None: "Without it the sandbox cannot dial back and the rollout will " "silently abort." ) - # Snapshot threshold: env override only. Absent => let TrajectoryManager - # take its built-in default. ``0`` disables, malformed values are - # ignored with a warning. - _snap_env = os.environ.get("SLIME_TITO_SNAPSHOT_MIN_LOSS_TOKENS") - _snap_threshold: int | None - if _snap_env is None: - _snap_threshold = None - else: - try: - _snap_threshold = int(_snap_env) - except ValueError: - logger.warning( - "SLIME_TITO_SNAPSHOT_MIN_LOSS_TOKENS=%r is not an int; falling back to TrajectoryManager default", - _snap_env, - ) - _snap_threshold = None - # Fork-merge threshold: same env-override convention. Triggers the - # rescue path in trajectory_manager when claude-code reformats a prior - # assistant message and the existing sibling's response is short - # enough to mask out instead of forking the tree. - _fork_merge_env = os.environ.get("SLIME_FORK_MERGE_MAX_RESPONSE_TOKENS") - _fork_merge_threshold: int | None - if _fork_merge_env is None: - _fork_merge_threshold = None - else: - try: - _fork_merge_threshold = int(_fork_merge_env) - except ValueError: - logger.warning( - "SLIME_FORK_MERGE_MAX_RESPONSE_TOKENS=%r is not an int; falling back to TrajectoryManager default", - _fork_merge_env, - ) - _fork_merge_threshold = None + drift_fork_threshold = os.environ.get("SLIME_DRIFT_FORK_MIN_LOSS_TOKENS") + drift_fork_threshold = int(drift_fork_threshold) if drift_fork_threshold else None + fork_merge_threshold = os.environ.get("SLIME_FORK_MERGE_MAX_RESPONSE_TOKENS") + fork_merge_threshold = int(fork_merge_threshold) if fork_merge_threshold else None self.adapter = AnthropicAdapter( tokenizer=self.tokenizer, sglang_url=sglang_url, tool_parser=self.tool_parser, reasoning_parser=self.reasoning_parser, - tito_snapshot_min_loss_tokens=_snap_threshold, - fork_merge_max_response_tokens=_fork_merge_threshold, + drift_fork_min_loss_tokens=drift_fork_threshold, + fork_merge_max_response_tokens=fork_merge_threshold, ) # handler_cancellation=True so a client disconnect cancels the handler # coroutine, arming the fire-and-forget /abort_request inside the diff --git a/slime/agent/adapters/anthropic.py b/slime/agent/adapters/anthropic.py index 48968a500b..b521ca29e3 100644 --- a/slime/agent/adapters/anthropic.py +++ b/slime/agent/adapters/anthropic.py @@ -66,7 +66,7 @@ def __init__( sglang_url, tool_parser=None, reasoning_parser=None, - tito_snapshot_min_loss_tokens: int | None = None, + drift_fork_min_loss_tokens: int | None = None, fork_merge_max_response_tokens: int | None = None, max_turns_per_sid: int | None = None, on_turn_appended: Callable[..., None] | None = None, @@ -81,8 +81,8 @@ def __init__( # ``None`` here means "caller did not specify" → let TrajectoryManager's # own default take over. Pass an int (incl. 0 to disable) to override. mgr_kwargs: dict[str, int] = {} - if tito_snapshot_min_loss_tokens is not None: - mgr_kwargs["tito_snapshot_min_loss_tokens"] = tito_snapshot_min_loss_tokens + if drift_fork_min_loss_tokens is not None: + mgr_kwargs["drift_fork_min_loss_tokens"] = drift_fork_min_loss_tokens if fork_merge_max_response_tokens is not None: mgr_kwargs["fork_merge_max_response_tokens"] = fork_merge_max_response_tokens self.manager = TrajectoryManager(**mgr_kwargs) @@ -485,8 +485,6 @@ async def _handle_request(request: web.Request) -> web.StreamResponse: ) adapter._sid_turn_count[sid] = prior + 1 - # Fold mid-list ``role: system`` messages into a neighbouring user message - # so Qwen3 chat templates accept them. No-op when no mid-list system msgs. _fold_mid_list_system_into_user(body) app = request.app @@ -521,8 +519,9 @@ async def _handle_request(request: web.Request) -> web.StreamResponse: ) blocks, stop_reason, response_message = _build_blocks_and_response_message(parsed, turn.finish_reason) - output_ids = list(turn.output_ids) finish_reason = _finish_reason_for_manager(turn.finish_reason, parsed.tool_uses) + turn = dataclasses.replace(turn, finish_reason=finish_reason) + output_ids = list(turn.output_ids) if _is_cc_title_generation_request(translated, tools_schema): # Claude Code meta request (per-session title generation). @@ -539,17 +538,10 @@ async def _handle_request(request: web.Request) -> web.StreamResponse: try: adapter.manager.append_turn( sid, + turn=turn, prompt_messages=translated, tools=tools_schema, - prompt_ids=prompt_ids, - response_ids=output_ids, - response_logprobs=( - list(turn.output_log_probs) - if turn.output_log_probs and len(turn.output_log_probs) == len(turn.output_ids) - else None - ), response_message=response_message, - finish_reason=finish_reason, metadata={"sid": sid}, ) except Exception: diff --git a/slime/agent/adapters/openai.py b/slime/agent/adapters/openai.py index cea3e5d858..15236f2edc 100644 --- a/slime/agent/adapters/openai.py +++ b/slime/agent/adapters/openai.py @@ -73,7 +73,7 @@ def __init__( sglang_url, tool_parser=None, reasoning_parser=None, - tito_snapshot_min_loss_tokens: int | None = None, + drift_fork_min_loss_tokens: int | None = None, fork_merge_max_response_tokens: int | None = None, max_turns_per_sid: int | None = None, on_turn_appended: Callable[..., None] | None = None, @@ -88,8 +88,8 @@ def __init__( # Mirror AnthropicAdapter: only forward kwargs the caller actually # specified, so TrajectoryManager's own defaults stay authoritative. mgr_kwargs: dict[str, int] = {} - if tito_snapshot_min_loss_tokens is not None: - mgr_kwargs["tito_snapshot_min_loss_tokens"] = tito_snapshot_min_loss_tokens + if drift_fork_min_loss_tokens is not None: + mgr_kwargs["drift_fork_min_loss_tokens"] = drift_fork_min_loss_tokens if fork_merge_max_response_tokens is not None: mgr_kwargs["fork_merge_max_response_tokens"] = fork_merge_max_response_tokens self.manager = TrajectoryManager(**mgr_kwargs) @@ -499,23 +499,17 @@ async def _handle_chat_completions(request: web.Request) -> web.StreamResponse: ) wire_message, manager_message, wire_finish = _build_oai_response(parsed, turn.finish_reason) - output_ids = list(turn.output_ids) finish_reason_mgr = _finish_reason_for_manager(turn.finish_reason, parsed.tool_uses) + turn = dataclasses.replace(turn, finish_reason=finish_reason_mgr) + output_ids = list(turn.output_ids) try: adapter.manager.append_turn( sid, + turn=turn, prompt_messages=translated, tools=tools_schema, - prompt_ids=prompt_ids, - response_ids=output_ids, - response_logprobs=( - list(turn.output_log_probs) - if turn.output_log_probs and len(turn.output_log_probs) == len(turn.output_ids) - else None - ), response_message=manager_message, - finish_reason=finish_reason_mgr, metadata={"sid": sid}, ) except Exception: diff --git a/slime/agent/trajectory_manager.py b/slime/agent/trajectory_manager.py index bf8714c2f1..9fbad298d1 100644 --- a/slime/agent/trajectory_manager.py +++ b/slime/agent/trajectory_manager.py @@ -27,25 +27,42 @@ later turn — logprobs stay coherent, no duplicated-content forks, no reliance on chat_template being position-invariant. -* Snapshot rescue (opt-in via ``tito_snapshot_min_loss_tokens``): when a - drift would drop >= N loss_mask=1 tokens, emit an extra "snapshot" - Sample alongside the main leaf. Snapshot tokens = cumulative pre-drop; - snapshot loss_mask is COMPLEMENTARY — 1 only at positions that the - main leaf is about to drop, 0 elsewhere. Snapshot reward = main-leaf - share; snapshot group_id = main-leaf group_id. Snapshot ∪ main on - loss_mask=1 tokens never overlap and their union equals the virtual - no-drift trajectory. The snapshotted drift is NOT counted in the main - sample's ``tito_dropped_*`` (it wasn't truly lost). +* Drift fork (gated by ``drift_fork_min_loss_tokens``, default 1024 — ON + by default): when a drift would drop >= N loss_mask=1 tokens, the leaf + FORKS. This is the primary drift path. We emit an extra synthetic + "drift_fork" Sample at drain time alongside the main Sample. + Fork tokens = cumulative pre-drop; fork loss_mask is COMPLEMENTARY — 1 + only at positions that the main leaf is about to drop, 0 elsewhere. Fork + reward = main-leaf share; fork group_id = main-leaf group_id. Fork ∪ main + on loss_mask=1 tokens never overlap and their union equals the virtual + no-drift trajectory. The forked drift is NOT counted in the main sample's + ``tito_dropped_*`` (it wasn't truly lost). + + When the drift would drop < N loss_mask=1 tokens (or is a pure-prompt + drift losing 0 loss tokens), the secondary DROP path applies instead: + drop-and-replace on the main leaf only, accounted in ``tito_dropped_*``. * On drift, ``Sample.metadata`` records: - ``tito_dropped_tokens`` — total tokens dropped (NOT including - drifts that produced a snapshot) - ``tito_dropped_turns`` — number of turns that triggered a drop - ``tito_snapshots_emitted`` — set on main leaf when >=1 snapshot - sibling was emitted for the same leaf - ``tito_snapshot`` — True on a snapshot Sample - ``tito_snapshot_at_turn`` — turn index whose drift triggered it - ``tito_snapshot_loss_tokens`` — count of loss_mask=1 tokens in snapshot + ``tito_dropped_tokens`` — total tokens dropped (NOT including + drifts that produced a fork) + ``tito_dropped_turns`` — number of turns that triggered a drop + ``tito_drift_forks_emitted`` — set on main leaf when >=1 drift-fork + sibling was emitted for the same leaf + ``tito_drift_fork`` — True on a drift-fork Sample + ``tito_drift_fork_at_turn`` — turn index whose drift triggered it + ``tito_drift_fork_loss_tokens`` — count of loss_mask=1 tokens in fork + +* Fork-merge rescue (gated by ``fork_merge_max_response_tokens``, default + 1024 — ON by default; set <=0 to disable): a routing-time mechanism, + independent of the + linearization-time drift_fork/drop above. When DFS breaks at an assistant + group (a later replay reformats an earlier assistant message, e.g. + tool_call arg ordering or whitespace) and exactly one leaf sibling carries + a per-turn response shorter than the threshold, the rewrite is collapsed + onto that sibling instead of spawning a new sibling subtree. The merged + sibling's stale response then enters trajectories with loss_mask=0, + recorded in ``fork_merge_masked_tokens`` / ``fork_merge_turns``. drift_fork + handles token alignment; fork-merge prevents a routing fork from forming. """ from __future__ import annotations @@ -56,6 +73,7 @@ from dataclasses import dataclass, field from typing import Any +from slime.agent.adapters.common import TurnRecord from slime.utils.types import Sample logger = logging.getLogger(__name__) @@ -164,6 +182,37 @@ class _PromptGroup: messages: list[dict[str, Any]] = field(default_factory=list) +@dataclass +class _DriftFork: + """Pre-drop fork collected when TITO drift would lose >= threshold loss tokens.""" + + tokens: list[int] + # Complementary mask: 0 at positions main leaf keeps, 1 at positions main leaf will drop. + loss_mask: list[int] + logprobs: list[float] + # asst.turn_index whose prompt triggered the drift. + drift_turn_index: int | None + # finish_reason of the prior assistant turn — describes "what the prefix + # looked like" before drift. Captured at fork creation so the build + # path doesn't have to look back into the chain. + prev_finish_reason: str | None + + +@dataclass +class _LeafAccum: + """Result of walking one leaf's assistant chain.""" + + tokens: list[int] = field(default_factory=list) + loss_mask: list[int] = field(default_factory=list) + logprobs: list[float] = field(default_factory=list) + drift_forks: list[_DriftFork] = field(default_factory=list) + # accounting (emitted to main sample's metadata when > 0) + dropped_tokens: int = 0 + dropped_turns: int = 0 + fork_merge_masked_tokens: int = 0 + fork_merge_turns: int = 0 + + def _group_messages_by_role( messages: list[dict[str, Any]], ) -> list[_PromptGroup]: @@ -189,25 +238,6 @@ def _lcp_len(a: list[int], b: list[int]) -> int: return i -def _short_text_preview(messages: list[dict[str, Any]], *, limit: int) -> str: - """Compact text preview of an assistant message block for log lines. - - Extracts string ``content`` and any ``{"text": "..."}`` content blocks - (anthropic-style); falls back to the empty string if nothing parseable. - """ - parts: list[str] = [] - for m in messages: - c = m.get("content") - if isinstance(c, str): - parts.append(c) - elif isinstance(c, list): - for blk in c: - if isinstance(blk, dict) and isinstance(blk.get("text"), str): - parts.append(blk["text"]) - s = " ".join(parts).strip() - return s[:limit] + ("…" if len(s) > limit else "") - - # =========================================================================== # TrajectoryManager # =========================================================================== @@ -224,26 +254,21 @@ class TrajectoryManager: def __init__( self, *, - tito_snapshot_min_loss_tokens: int | None = None, - fork_merge_max_response_tokens: int | None = None, + drift_fork_min_loss_tokens: int = 1024, + fork_merge_max_response_tokens: int = 1024, ) -> None: - # Drift-snapshot threshold (loss_mask=1 token count inside drift suffix). - # None or <= 0 disables; behavior then matches the pre-feature output. - self._snap_threshold: int | None = ( - tito_snapshot_min_loss_tokens - if (tito_snapshot_min_loss_tokens is not None and tito_snapshot_min_loss_tokens > 0) - else None - ) + # Drift-fork threshold (loss_mask=1 token count inside drift suffix). + # When a drift would drop >= this many loss tokens, fork instead of + # dropping. Always an int (default 1024 — drift-fork ON by default). + # Set <=0 to effectively disable (combined with the drift_loss_tokens>0 + # guard, only true drifts above the threshold ever fork). + self._fork_threshold: int = drift_fork_min_loss_tokens # Fork-merge threshold: when DFS would break at an assistant group and # exactly one non-leaf assistant sibling has turn_response_ids length # STRICTLY LESS than this value, collapse the would-be fork onto that # sibling (its response then enters trajectories with loss_mask=0). - # None or <= 0 disables; behavior matches pre-feature output. - self._fork_merge_threshold: int | None = ( - fork_merge_max_response_tokens - if (fork_merge_max_response_tokens is not None and fork_merge_max_response_tokens > 0) - else None - ) + # Default 1024 — fork-merge ON by default; set <=0 to disable. + self._fork_merge_threshold: int = fork_merge_max_response_tokens self._trees: dict[str, Node] = {} self._turn_count: dict[str, int] = {} @@ -259,90 +284,85 @@ def append_turn( self, sid: str, *, + turn: TurnRecord, prompt_messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None, - prompt_ids: list[int], - response_ids: list[int], - response_logprobs: list[float] | None, response_message: dict[str, Any] | None, - finish_reason: str | None, metadata: dict[str, Any] | None = None, ) -> None: if not prompt_messages: logger.warning("append_turn(sid=%s): empty prompt_messages; skipping", sid) return - if response_logprobs is not None and len(response_logprobs) != len(response_ids): + if turn.output_log_probs and len(turn.output_log_probs) != len(turn.output_ids): raise ValueError( - f"response_logprobs length {len(response_logprobs)} != " f"response_ids length {len(response_ids)}" + f"turn.output_log_probs length {len(turn.output_log_probs)} != " + f"turn.output_ids length {len(turn.output_ids)}" ) - root = self._trees.get(sid) - if root is None: - root = Node() - self._trees[sid] = root - + root = self._trees.setdefault(sid, Node()) groups = _group_messages_by_role(prompt_messages) - # Step 1: DFS by (role, node_match_key) ONLY. No prompt_ids check. + cur, i = self._find_mount_point(root, groups) + cur, i = self._try_fork_merge_assistant(sid, cur, groups, i) + cur = self._mount_prompt_groups(cur, groups[i:], tools) + self._attach_assistant_leaf(sid, cur, turn=turn, response_message=response_message, metadata=metadata) + + def _find_mount_point(self, root: Node, groups: list[_PromptGroup]) -> tuple[Node, int]: + """DFS down the existing tree by ``(role, node_match_key)``. + + Returns ``(cur, i)``: ``cur`` is the deepest node whose path matches + ``groups[:i]`` exactly; ``i`` is the index into ``groups`` of the first + group that diverges from anything mounted so far (i.e., where this + turn's new content begins). + """ cur = root i = 0 while i < len(groups): g_key = node_match_key(groups[i].messages) - match: Node | None = None - for child in cur.children: - if child.role == groups[i].role and node_match_key(child.messages) == g_key: - match = child - break + match = next( + (c for c in cur.children if c.role == groups[i].role and node_match_key(c.messages) == g_key), + None, + ) if match is None: break cur = match i += 1 - - # Step 1.5: assistant fork-merge rescue (opt-in via - # fork_merge_max_response_tokens). The typical claude-code pattern is: - # a later replay reformats an earlier assistant message (e.g. tool_call - # arg ordering, whitespace), which breaks the DFS at that assistant - # group. Without rescue, every such reformat spawns a new sibling - # subtree; with rescue, when the existing sibling's per-turn response - # is short enough that masking it out is the cheaper trade-off, we - # collapse onto that sibling and mark it for mask=0 at linearization. - if i < len(groups) and self._fork_merge_threshold is not None and groups[i].role == "assistant": - candidates = [ - c - for c in cur.children - if c.role == "assistant" - # Real rewrite footprint = the original turn's leaf assistant - # node: it was inserted by Step-3 of a prior turn (so - # turn_response_ids is populated), and no later turn has - # extended it (so it is still a leaf). Once a rewrite collapses - # onto it via this rescue, the new turn's user/tool + asst leaf - # are appended underneath as children — so the merge target - # MUST be a leaf at decision time, otherwise it has already - # diverged into mixed subchains and merging would tangle them. - and not c.children - and c.turn_response_ids is not None - and len(c.turn_response_ids) < self._fork_merge_threshold - ] - if len(candidates) == 1: - sib = candidates[0] - masked = len(sib.turn_response_ids or []) - preview = _short_text_preview(sib.messages, limit=160) - logger.warning( - "append_turn(sid=%s turn=%s): fork-merging assistant rewrite " - "into existing sibling (turn_index=%s, masked_response_tokens=%d, " - "sibling_response_preview=%r)", - sid, - self._turn_count.get(sid, 0) + 1, - sib.turn_index, - masked, - preview, - ) - sib.metadata["fork_merged"] = True - sib.metadata["fork_merge_masked_tokens"] = masked - sib.messages = list(groups[i].messages) - cur = sib - i += 1 - elif len(candidates) >= 2: + return cur, i + + def _try_fork_merge_assistant(self, sid: str, cur: Node, groups: list[_PromptGroup], i: int) -> tuple[Node, int]: + """Optionally collapse an assistant-rewrite onto a single short leaf sibling. + + The typical claude-code pattern is: a later replay reformats an earlier + assistant message (e.g. tool_call arg ordering, whitespace), which + breaks DFS at that assistant group. Without rescue, every such reformat + spawns a new sibling subtree; with rescue, when the existing sibling's + per-turn response is short enough that masking it out is the cheaper + trade-off, we collapse onto that sibling and mark it for mask=0 at + linearization. + """ + if self._fork_merge_threshold <= 0: + return cur, i # feature off + if i >= len(groups) or groups[i].role != "assistant": + return cur, i # feature on, but this turn isn't an asst rewrite + + candidates = [ + c + for c in cur.children + if c.role == "assistant" + # Real rewrite footprint = the original turn's leaf assistant + # node: it was inserted by step 3 of a prior turn (so + # turn_response_ids is populated), and no later turn has + # extended it (so it is still a leaf). Once a rewrite collapses + # onto it via this rescue, the new turn's user/tool + asst leaf + # are appended underneath as children — so the merge target + # MUST be a leaf at decision time, otherwise it has already + # diverged into mixed subchains and merging would tangle them. + and not c.children + and c.turn_response_ids is not None + and len(c.turn_response_ids) < self._fork_merge_threshold + ] + if len(candidates) != 1: + if len(candidates) >= 2: # Legacy / pathological state: fork_merge wasn't on during # earlier rewrites, or threshold was widened mid-session. # Don't pick arbitrarily — fork as usual and surface a hint. @@ -354,29 +374,66 @@ def append_turn( self._turn_count.get(sid, 0) + 1, len(candidates), ) + return cur, i + + sib = candidates[0] + masked = len(sib.turn_response_ids or []) + logger.warning( + "append_turn(sid=%s turn=%s): fork-merging assistant rewrite " + "into existing sibling (turn_index=%s, masked_response_tokens=%d)", + sid, + self._turn_count.get(sid, 0) + 1, + sib.turn_index, + masked, + ) + sib.metadata["fork_merged"] = True + sib.metadata["fork_merge_masked_tokens"] = masked + sib.messages = list(groups[i].messages) + return sib, i + 1 - # Step 2: mount remaining prompt groups as plain routing nodes. - # Token attribution happens at get_trajectory time, not here. - for g in groups[i:]: + def _mount_prompt_groups( + self, + cur: Node, + remaining_groups: list[_PromptGroup], + tools: list[dict[str, Any]] | None, + ) -> Node: + """Attach each remaining prompt group as a routing node under ``cur``. + + Token attribution happens at get_trajectory time, not here. The tools + metadata is placed only on the FIRST system node on the path — + ``_first_system_already_set(cur)`` walks ``cur → root`` looking for a + system ancestor that already carries it, and ``cur`` here is the + deepest node from descent (+ optional merge), so the walk sees every + ancestor that's already mounted. + """ + for g in remaining_groups: md: dict[str, Any] = {} if g.role == "system" and tools is not None and not self._first_system_already_set(cur): md["tools"] = list(tools) cur = cur.add_child(Node(role=g.role, messages=list(g.messages), metadata=md)) + return cur - # Step 3: assistant leaf with this turn's sglang snapshot. - asst_messages = [response_message] if response_message is not None else [] + def _attach_assistant_leaf( + self, + sid: str, + cur: Node, + *, + turn: TurnRecord, + response_message: dict[str, Any] | None, + metadata: dict[str, Any] | None, + ) -> None: + """Attach this turn's assistant leaf carrying the sglang snapshot.""" asst = Node( role="assistant", - messages=asst_messages, + messages=[response_message] if response_message is not None else [], metadata=dict(metadata or {}), ) - asst.turn_prompt_ids = list(prompt_ids) - asst.turn_response_ids = list(response_ids) - asst.turn_response_logprobs = list(response_logprobs) if response_logprobs is not None else None - asst.turn_finish_reason = finish_reason + asst.turn_prompt_ids = list(turn.prompt_ids) + asst.turn_response_ids = list(turn.output_ids) + asst.turn_response_logprobs = list(turn.output_log_probs) if turn.output_log_probs else None + asst.turn_finish_reason = turn.finish_reason asst.turn_index = self._turn_count.get(sid, 0) + 1 cur.add_child(asst) - self._turn_count[sid] = asst.turn_index def get_trajectory( @@ -386,33 +443,14 @@ def get_trajectory( base_sample=None, reward: float = 0.0, extra_metadata: dict[str, Any] | None = None, - drop: bool = True, ) -> list: - """Linearize each leaf into a slime ``Sample`` using LCP drop-and-replace. - - For each leaf, walk root→leaf collecting assistant nodes in order. - Start with tokens=[]. For each assistant turn k (1-based): - - 1. ``p = asst.turn_prompt_ids``, ``r = asst.turn_response_ids``. - 2. If k == 1: emit all of ``p`` as loss_mask=0 (plus 0.0 logprobs), - then ``r`` as loss_mask=1 with real logprobs. - 3. If k >= 2: compute ``L = LCP(tokens, p)``. Truncate tokens / - loss_mask / logprobs to length L (DROP everything past L — - that includes the previous turn's response tokens that fall - in the drift region; logging tells you how much was dropped). - Then append ``p[L:]`` (loss_mask=0) and ``r`` (loss_mask=1). - - When drift fires on at least one turn, the returned Sample's - ``metadata`` gains ``tito_dropped_tokens`` (total tokens dropped - across the leaf) and ``tito_dropped_turns`` (how many turns - triggered a drop). Both keys are absent when no drift occurs. - - When ``tito_snapshot_min_loss_tokens`` was passed to the constructor - and a drift would drop >= that many loss_mask=1 tokens, an extra - snapshot Sample is emitted before the main-leaf Sample carrying just - the to-be-lost tokens (complementary mask). See module docstring. - - See module docstring for the rationale. + """Drain a sid into slime ``Sample`` objects, then drop the session. + + ``get_trajectory`` is the lifecycle boundary where the message routing + tree is linearized into token-normalized ``Sample`` objects. Each + routing leaf yields one main Sample plus one extra Sample per drift + fork. ``reward`` is split evenly across all emitted samples. The sid is + consumed: a second call for the same sid returns ``[]``. """ if base_sample is None: base_sample = Sample(index=0, prompt="") @@ -420,170 +458,260 @@ def get_trajectory( root = self._trees.get(sid) if root is None: return [] - leaves = [leaf for leaf in root.leaves() if not leaf.is_root] - samples: list[Sample] = [] - for leaf in leaves: - chain = leaf.path_from_root() - # Only assistant leaves carrying this turn's sglang snapshot - # participate in TITO accumulation. Routing assistant nodes mounted - # from prior-turn replay (turn_prompt_ids is None) carry no token - # signal and would otherwise be misread as a full-trajectory drift. - asst_chain = [n for n in chain if n.role == "assistant" and n.turn_prompt_ids is not None] - - tokens: list[int] = [] - loss_mask: list[int] = [] - logprobs: list[float] = [] - total_dropped = 0 - dropped_turns = 0 - fork_merge_masked_total = 0 - fork_merge_turns_total = 0 - snapshots: list[tuple[list[int], list[int], list[float], int, int]] = [] - - for k, asst in enumerate(asst_chain, start=1): - p = list(asst.turn_prompt_ids or []) - r = list(asst.turn_response_ids or []) - lp = list(asst.turn_response_logprobs) if asst.turn_response_logprobs is not None else None - is_merged = bool(asst.metadata.get("fork_merged")) - - if k == 1: - emit_prompt = p - else: - L = _lcp_len(tokens, p) - drift = len(tokens) - L - if drift > 0: - drift_loss_tokens = sum(loss_mask[L:]) - snap_emitted = False - if self._snap_threshold is not None and drift_loss_tokens >= self._snap_threshold: - snap_tokens = list(tokens) - snap_mask = [0] * L + list(loss_mask[L:]) - snap_lp = [0.0] * L + [ - (logprobs[i] if loss_mask[i] == 1 else 0.0) for i in range(L, len(tokens)) - ] - snapshots.append((snap_tokens, snap_mask, snap_lp, asst.turn_index, k - 1)) - snap_emitted = True - logger.warning( - "get_trajectory(sid=%s leaf turn=%s): TITO drift detected, " - "dropping %d prior tokens (incl. previous-turn response) to " - "realign with this turn's prompt%s", - sid, - asst.turn_index, - drift, - f"; snapshotted {drift_loss_tokens} loss tokens" if snap_emitted else "", - ) - if not snap_emitted: - total_dropped += drift - dropped_turns += 1 - tokens = tokens[:L] - loss_mask = loss_mask[:L] - logprobs = logprobs[:L] - emit_prompt = p[L:] - - tokens.extend(emit_prompt) - loss_mask.extend([0] * len(emit_prompt)) - logprobs.extend([0.0] * len(emit_prompt)) - - tokens.extend(r) - # fork-merged sibling: its response is "stale" — present in the - # tree only as a routing placeholder for the rewrites that - # collapsed onto it; mask it out of training. - loss_mask.extend([0 if is_merged else 1] * len(r)) - if lp is not None: - logprobs.extend(lp) - else: - logprobs.extend([0.0] * len(r)) - - if is_merged: - fork_merge_masked_total += len(r) - fork_merge_turns_total += 1 - - last_asst = asst_chain[-1] if asst_chain else None - first_sys = next((n for n in chain if n.role == "system"), None) - tools_meta = first_sys.metadata.get("tools") if first_sys else None - base_md: dict[str, Any] = { - **(base_sample.metadata or {}), - **(extra_metadata or {}), - "tools": tools_meta, - } - per_leaf_reward = (reward / len(leaves)) if leaves else 0.0 - - # slime contract (see slime/backends/megatron_utils/data.py:139, - # slime/ray/rollout.py:695): ``loss_mask`` and ``rollout_log_probs`` - # cover only the response region — i.e. tokens AFTER the initial - # prompt — and ``response_length == len(loss_mask)``. Strip the - # first turn's prompt prefix here so all downstream consumers see - # tokens/loss_mask/logprobs in their canonical alignment. - first_prompt_len = len(asst_chain[0].turn_prompt_ids or []) if asst_chain else 0 - - # Emit snapshot sample(s) first, then the main-leaf sample. - for snap_tokens, snap_mask, snap_lp, drift_turn, cur_chain_idx in snapshots: - snap_finish = None - prev_idx = cur_chain_idx - 1 # asst_chain index of the previous (prefix's last) turn - if 0 <= prev_idx < len(asst_chain): - snap_finish = asst_chain[prev_idx].turn_finish_reason - snap_strip = min(first_prompt_len, len(snap_mask)) - snap_mask_resp = snap_mask[snap_strip:] - snap_lp_resp = snap_lp[snap_strip:] - snap_md = { - **base_md, - "finish_reason": snap_finish, - "tito_snapshot": True, - "tito_snapshot_at_turn": drift_turn, - "tito_snapshot_loss_tokens": sum(snap_mask_resp), - } - samples.append( - Sample( - index=base_sample.index, - group_id=(base_sample.group_id if base_sample.group_id is not None else base_sample.index), - prompt=base_sample.prompt, - label=base_sample.label, - tokens=snap_tokens, - response_length=len(snap_mask_resp), - loss_mask=snap_mask_resp, - rollout_log_probs=snap_lp_resp, - reward=per_leaf_reward, - status=Sample.Status.COMPLETED, - metadata=snap_md, - ) - ) - main_strip = min(first_prompt_len, len(loss_mask)) - loss_mask_resp = loss_mask[main_strip:] - logprobs_resp = logprobs[main_strip:] - response_length = len(loss_mask_resp) - main_md: dict[str, Any] = { - **base_md, - "finish_reason": last_asst.turn_finish_reason if last_asst else None, - } - if total_dropped > 0: - main_md["tito_dropped_tokens"] = total_dropped - main_md["tito_dropped_turns"] = dropped_turns - if snapshots: - main_md["tito_snapshots_emitted"] = len(snapshots) - if fork_merge_masked_total > 0: - main_md["fork_merge_masked_tokens"] = fork_merge_masked_total - main_md["fork_merge_turns"] = fork_merge_turns_total - samples.append( - Sample( - index=base_sample.index, - group_id=(base_sample.group_id if base_sample.group_id is not None else base_sample.index), - prompt=base_sample.prompt, - label=base_sample.label, - tokens=tokens, - response_length=response_length, - loss_mask=loss_mask_resp, - rollout_log_probs=logprobs_resp, - reward=per_leaf_reward, - status=Sample.Status.COMPLETED, - metadata=main_md, - ) + samples: list[Sample] = [] + for routing_leaf in root.leaves(): + if routing_leaf.is_root: + continue + samples.extend( + self._normalize_routing_leaf(sid, routing_leaf, base_sample=base_sample, extra_metadata=extra_metadata) ) - if drop: - self._trees.pop(sid, None) - self._turn_count.pop(sid, None) + + # Reward is split evenly across every emitted sample (main + forks); + # the token-weighted reducer downstream then gives each loss token the + # trajectory's full R. Assigned after the fact so the per-leaf builder + # stays reward-agnostic. + per_sample_reward = (reward / len(samples)) if samples else 0.0 + for s in samples: + s.reward = per_sample_reward + + self._trees.pop(sid, None) + self._turn_count.pop(sid, None) return samples # -------------------- internals ---------------------------------------- + def _normalize_routing_leaf( + self, + sid: str, + leaf: Node, + *, + base_sample: Sample, + extra_metadata: dict[str, Any] | None, + ) -> list[Sample]: + """Linearize one routing leaf into its main Sample (+ drift-fork Samples). + + Drift forks come first, then the main sample, matching the original + drain order. Reward is left at 0.0 here and assigned by the caller. + """ + chain = leaf.path_from_root() + # Only assistant leaves carrying this turn's sglang snapshot + # participate in TITO accumulation. Routing assistant nodes mounted + # from prior-turn replay (turn_prompt_ids is None) carry no token + # signal and would otherwise be misread as a full-trajectory drift. + asst_chain = [n for n in chain if n.role == "assistant" and n.turn_prompt_ids is not None] + accum = self._accumulate_chain(sid, asst_chain) + first_sys = next((n for n in chain if n.role == "system"), None) + base_md = {"tools": first_sys.metadata.get("tools") if first_sys else None} + first_prompt_len = len(asst_chain[0].turn_prompt_ids or []) if asst_chain else 0 + + # Build one Sample from a linearized segment (a _DriftFork or the accum + # itself — both expose tokens/loss_mask/logprobs). The base/extra/clamp + # args are constant across segments, so this closure carries them. + def build(seg: _DriftFork | _LeafAccum, leaf_md: dict[str, Any]) -> Sample: + return self._build_leaf_sample( + base_sample=base_sample, + extra_metadata=extra_metadata, + leaf_metadata={**base_md, **leaf_md}, + tokens=seg.tokens, + loss_mask=seg.loss_mask, + logprobs=seg.logprobs, + first_prompt_len=first_prompt_len, + ) + + # Drift forks first, then the main sample — matching the original drain + # order. A fork's loss_mask is complementary (0 on [0:L], L >= + # first_prompt_len), so the response-region clamp never touches a loss=1 + # position and sum(loss_mask) is the final loss-token count. + samples = [ + build( + fork, + { + "finish_reason": fork.prev_finish_reason, + "tito_drift_fork": True, + "tito_drift_fork_at_turn": fork.drift_turn_index, + "tito_drift_fork_loss_tokens": sum(fork.loss_mask), + }, + ) + for fork in accum.drift_forks + ] + samples.append(build(accum, self._main_leaf_metadata(accum, asst_chain))) + return samples + + @staticmethod + def _main_leaf_metadata(accum: _LeafAccum, asst_chain: list[Node]) -> dict[str, Any]: + """Assemble the main sample's leaf metadata (conditional drift/merge keys).""" + last_asst = asst_chain[-1] if asst_chain else None + md: dict[str, Any] = {"finish_reason": last_asst.turn_finish_reason if last_asst else None} + if accum.dropped_tokens > 0: + md["tito_dropped_tokens"] = accum.dropped_tokens + md["tito_dropped_turns"] = accum.dropped_turns + if accum.drift_forks: + md["tito_drift_forks_emitted"] = len(accum.drift_forks) + if accum.fork_merge_masked_tokens > 0: + md["fork_merge_masked_tokens"] = accum.fork_merge_masked_tokens + md["fork_merge_turns"] = accum.fork_merge_turns + return md + + def _build_leaf_sample( + self, + *, + base_sample: Sample, + extra_metadata: dict[str, Any] | None, + leaf_metadata: dict[str, Any], + tokens: list[int], + loss_mask: list[int], + logprobs: list[float], + first_prompt_len: int, + ) -> Sample: + """Build one Sample from a linearized token segment. + + ``loss_mask`` / ``logprobs`` are clamped to the response region (the + leading first-turn prompt prefix is stripped) per the slime contract: + ``response_length == len(loss_mask)`` and loss_mask/logprobs cover only + the response region. ``reward`` is left at 0.0; the caller assigns the + per-sample share. + """ + loss_resp, lp_resp = self._response_region(loss_mask, logprobs, first_prompt_len) + metadata = { + **(base_sample.metadata or {}), + **(extra_metadata or {}), + **leaf_metadata, + } + return Sample( + index=base_sample.index, + group_id=base_sample.group_id if base_sample.group_id is not None else base_sample.index, + prompt=base_sample.prompt, + label=base_sample.label, + tokens=list(tokens), + response_length=len(loss_resp), + loss_mask=loss_resp, + rollout_log_probs=lp_resp, + reward=0.0, + status=Sample.Status.COMPLETED, + metadata=metadata, + ) + + def _accumulate_chain(self, sid: str, asst_chain: list[Node]) -> _LeafAccum: + """Apply LCP drop-and-replace to an assistant chain in turn order. + + Mutates and returns ``accum``. See the module docstring for the + algorithm, drift-fork contract, and fork-merge masking rule. + """ + accum = _LeafAccum() + for k, asst in enumerate(asst_chain, start=1): + # Read-only views: only fed to _lcp_len / .extend(), never mutated + # in place, so no defensive copy is needed here (the leaf already + # owns isolated copies from _attach_assistant_leaf). + prompt_ids = asst.turn_prompt_ids or [] + response_ids = asst.turn_response_ids or [] + response_logprobs = asst.turn_response_logprobs + is_merged = bool(asst.metadata.get("fork_merged")) + + # k == 1 falls out of the general case: LCP([], prompt) == 0 so + # _fork_or_drop_drift trivially returns the full prompt with no drop. + prev_finish_reason = asst_chain[k - 2].turn_finish_reason if k >= 2 else None + emit_prompt = self._fork_or_drop_drift( + sid, + accum, + prev_finish_reason=prev_finish_reason, + drift_turn_index=asst.turn_index, + prompt=prompt_ids, + ) + + accum.tokens.extend(emit_prompt) + accum.loss_mask.extend([0] * len(emit_prompt)) + accum.logprobs.extend([0.0] * len(emit_prompt)) + + accum.tokens.extend(response_ids) + # fork-merged sibling: its response is "stale" — present in the tree + # only as a routing placeholder for the rewrites that collapsed onto + # it; mask it out of training. + accum.loss_mask.extend([0 if is_merged else 1] * len(response_ids)) + accum.logprobs.extend(response_logprobs if response_logprobs is not None else [0.0] * len(response_ids)) + + if is_merged: + accum.fork_merge_masked_tokens += len(response_ids) + accum.fork_merge_turns += 1 + return accum + + def _fork_or_drop_drift( + self, + sid: str, + accum: _LeafAccum, + *, + prev_finish_reason: str | None, + drift_turn_index: int | None, + prompt: list[int], + ) -> list[int]: + """Resolve a TITO drift between cumulative tokens and the next prompt. + + Compute LCP ``L``; if there's no drift suffix, just return the prompt + tail. Otherwise decide FORK vs DROP (fork is the primary path), then + truncate ``accum.{tokens,loss_mask,logprobs}`` to ``L`` and return + ``prompt[L:]`` for the caller to emit. + + FORK (drift loses >= ``self._fork_threshold`` loss_mask=1 tokens): + append a complementary-mask ``_DriftFork`` to ``accum`` so the + dropped training signal survives as a sibling output leaf. The + drift is NOT counted toward ``accum.dropped_*``. + DROP (below threshold, or a pure-prompt drift losing 0 loss tokens): + drop-and-replace only, counted toward ``accum.dropped_*``. + """ + L = _lcp_len(accum.tokens, prompt) + drift = len(accum.tokens) - L + if drift == 0: + return prompt[L:] + + drift_loss_tokens = sum(accum.loss_mask[L:]) + # PRIMARY: fork — losing >= threshold loss tokens is worth a + # complementary-mask fork leaf so the dropped signal survives. The + # drift_loss_tokens>0 guard keeps pure-prompt drift on the DROP path + # even if the threshold is set to <=0. + forked = drift_loss_tokens > 0 and drift_loss_tokens >= self._fork_threshold + if forked: + accum.drift_forks.append( + _DriftFork( + tokens=list(accum.tokens), + loss_mask=[0] * L + list(accum.loss_mask[L:]), + logprobs=[0.0] * L + + [(accum.logprobs[i] if accum.loss_mask[i] == 1 else 0.0) for i in range(L, len(accum.tokens))], + drift_turn_index=drift_turn_index, + prev_finish_reason=prev_finish_reason, + ) + ) + else: + # SECONDARY: drop-and-replace only. + accum.dropped_tokens += drift + accum.dropped_turns += 1 + + logger.warning( + "get_trajectory(sid=%s leaf turn=%s): TITO drift detected, " + "dropping %d prior tokens (incl. previous-turn response) to " + "realign with this turn's prompt%s", + sid, + drift_turn_index, + drift, + f"; forked {drift_loss_tokens} loss tokens" if forked else "", + ) + + del accum.tokens[L:] + del accum.loss_mask[L:] + del accum.logprobs[L:] + return prompt[L:] + + @staticmethod + def _response_region( + loss_mask: list[int], + logprobs: list[float], + first_prompt_len: int, + ) -> tuple[list[int], list[float]]: + """Strip the leading first-turn prompt prefix from loss_mask / logprobs + so what remains is the response-region view slime expects (see + slime/backends/megatron_utils/data.py:139, slime/ray/rollout.py:695).""" + strip = min(first_prompt_len, len(loss_mask)) + return loss_mask[strip:], logprobs[strip:] + @staticmethod def _first_system_already_set(start: Node) -> bool: """Walk start->root looking for a system node already carrying tools.""" From 6f95f188969c5b8cb30c282b4d2213a974330357 Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Mon, 8 Jun 2026 07:24:05 +0000 Subject: [PATCH 10/28] refactor(agent): drop drift fork/merge params, strict exact-prefix linearization TrajectoryManager now uses strict exact-prefix linearization and raises on TITO drift, so the drift_fork_min_loss_tokens / fork_merge_max_response_tokens knobs are removed from both adapters. generate.py warns loudly if the corresponding env vars are still set, and stops attaching per-trajectory metadata to merged samples (revisit when dump/analysis needs it). --- examples/coding_agent_rl/generate.py | 40 +- slime/agent/adapters/anthropic.py | 11 +- slime/agent/adapters/openai.py | 11 +- slime/agent/trajectory_manager.py | 622 ++++++++------------------- 4 files changed, 202 insertions(+), 482 deletions(-) diff --git a/examples/coding_agent_rl/generate.py b/examples/coding_agent_rl/generate.py index d7dd6a1616..09a63a776c 100644 --- a/examples/coding_agent_rl/generate.py +++ b/examples/coding_agent_rl/generate.py @@ -97,17 +97,23 @@ def __init__(self, args) -> None: "Without it the sandbox cannot dial back and the rollout will " "silently abort." ) + # Kept for the staged re-add of drift tolerance: the strict + # exact-prefix TrajectoryManager no longer honors these, so warn loudly + # if an operator set them expecting fork/merge behavior. drift_fork_threshold = os.environ.get("SLIME_DRIFT_FORK_MIN_LOSS_TOKENS") - drift_fork_threshold = int(drift_fork_threshold) if drift_fork_threshold else None fork_merge_threshold = os.environ.get("SLIME_FORK_MERGE_MAX_RESPONSE_TOKENS") - fork_merge_threshold = int(fork_merge_threshold) if fork_merge_threshold else None + if drift_fork_threshold or fork_merge_threshold: + logger.warning( + "[coding_agent_rl] SLIME_DRIFT_FORK_MIN_LOSS_TOKENS / " + "SLIME_FORK_MERGE_MAX_RESPONSE_TOKENS are set but currently " + "ignored: TrajectoryManager uses strict exact-prefix " + "linearization and raises on TITO drift." + ) self.adapter = AnthropicAdapter( tokenizer=self.tokenizer, sglang_url=sglang_url, tool_parser=self.tool_parser, reasoning_parser=self.reasoning_parser, - drift_fork_min_loss_tokens=drift_fork_threshold, - fork_merge_max_response_tokens=fork_merge_threshold, ) # handler_cancellation=True so a client disconnect cancels the handler # coroutine, arming the fire-and-forget /abort_request inside the @@ -185,29 +191,15 @@ def _merge_samples( """Decorate per-leaf Samples returned by TrajectoryManager.get_trajectory. The manager already filled tokens / loss_mask / rollout_log_probs / - response_length / reward (reward / N). We add per-trajectory metadata - (is_solved / applied_cleanly / elapsed_sec / segment_idx) and decode - ``sample.response`` from the response tokens slice -- slime's training - logging path reads this string. + response_length / reward (reward / N). We decode ``sample.response`` from + the response tokens slice -- slime's training logging path reads this + string. Per-trajectory metadata is intentionally NOT attached to the + samples (kept empty); revisit when dump/analysis needs it. """ if not samples: return _abort_result(sample, "adapter_session_empty") - trajectory_metadata = { - "instance_id": instance_id, - "is_solved": reward_result.is_solved, - "applied_cleanly": reward_result.applied_cleanly, - "elapsed_sec": elapsed_sec, - } - - k = len(samples) - for i, s in enumerate(samples): - s.metadata = { - **(s.metadata or {}), - **trajectory_metadata, - "segment_idx": i, - "num_segments": k, - } + for s in samples: rlen = int(s.response_length or 0) if rlen and s.tokens: s.response = state.tokenizer.decode(s.tokens[-rlen:], skip_special_tokens=False) @@ -221,7 +213,7 @@ def _merge_samples( reward_result.is_solved, reward_result.applied_cleanly, elapsed_sec, - k, + len(samples), ) return samples diff --git a/slime/agent/adapters/anthropic.py b/slime/agent/adapters/anthropic.py index b521ca29e3..122e683c6c 100644 --- a/slime/agent/adapters/anthropic.py +++ b/slime/agent/adapters/anthropic.py @@ -66,8 +66,6 @@ def __init__( sglang_url, tool_parser=None, reasoning_parser=None, - drift_fork_min_loss_tokens: int | None = None, - fork_merge_max_response_tokens: int | None = None, max_turns_per_sid: int | None = None, on_turn_appended: Callable[..., None] | None = None, ) -> None: @@ -78,14 +76,7 @@ def __init__( reasoning_parser=reasoning_parser, ) # ONE manager shared across all sids; per-sid trees live inside. - # ``None`` here means "caller did not specify" → let TrajectoryManager's - # own default take over. Pass an int (incl. 0 to disable) to override. - mgr_kwargs: dict[str, int] = {} - if drift_fork_min_loss_tokens is not None: - mgr_kwargs["drift_fork_min_loss_tokens"] = drift_fork_min_loss_tokens - if fork_merge_max_response_tokens is not None: - mgr_kwargs["fork_merge_max_response_tokens"] = fork_merge_max_response_tokens - self.manager = TrajectoryManager(**mgr_kwargs) + self.manager = TrajectoryManager() # Optional debug hook invoked after each successful append_turn. # Signature: (sid, prompt_messages, tools, response_message, # prompt_ids, response_ids, finish_reason) -> None. diff --git a/slime/agent/adapters/openai.py b/slime/agent/adapters/openai.py index 15236f2edc..c1e71d1f84 100644 --- a/slime/agent/adapters/openai.py +++ b/slime/agent/adapters/openai.py @@ -73,8 +73,6 @@ def __init__( sglang_url, tool_parser=None, reasoning_parser=None, - drift_fork_min_loss_tokens: int | None = None, - fork_merge_max_response_tokens: int | None = None, max_turns_per_sid: int | None = None, on_turn_appended: Callable[..., None] | None = None, ) -> None: @@ -85,14 +83,7 @@ def __init__( reasoning_parser=reasoning_parser, ) # ONE manager shared across all sids; per-sid trees live inside. - # Mirror AnthropicAdapter: only forward kwargs the caller actually - # specified, so TrajectoryManager's own defaults stay authoritative. - mgr_kwargs: dict[str, int] = {} - if drift_fork_min_loss_tokens is not None: - mgr_kwargs["drift_fork_min_loss_tokens"] = drift_fork_min_loss_tokens - if fork_merge_max_response_tokens is not None: - mgr_kwargs["fork_merge_max_response_tokens"] = fork_merge_max_response_tokens - self.manager = TrajectoryManager(**mgr_kwargs) + self.manager = TrajectoryManager() # Optional debug hook invoked after each successful append_turn. # Signature mirrors AnthropicAdapter.on_turn_appended: # (sid, prompt_messages, tools, response_message, diff --git a/slime/agent/trajectory_manager.py b/slime/agent/trajectory_manager.py index 9fbad298d1..9ba327e17e 100644 --- a/slime/agent/trajectory_manager.py +++ b/slime/agent/trajectory_manager.py @@ -1,68 +1,38 @@ -"""Per-role chunk-merging trajectory tree manager (C-plan: token-faithful). +"""Per-message trajectory tree manager (C-plan: token-faithful). -Design (Plan C, 2026-06-03): +Design (Plan C, 2026-06-03; strict exact-prefix rewrite 2026-06-08; +per-message routing 2026-06-08): * The tree is a router only. DFS merge keys on ``(role, node_match_key)`` - alone — no prompt_ids prefix check. Same conversation prefix in - ``messages`` space always lands on the same path, regardless of any - chat_template re-tokenization drift across turns. + alone, one tree node per message — no prompt_ids prefix check and no + same-role grouping. Same conversation prefix in ``messages`` space always + lands on the same path, regardless of any chat_template re-tokenization + drift across turns. * Each assistant leaf stores the THIS-TURN sglang snapshot: ``turn_prompt_ids`` / ``turn_response_ids`` / ``turn_response_logprobs`` / ``turn_finish_reason`` / ``turn_index``. Non-assistant nodes carry no token attribution at all. -* ``get_trajectory`` linearizes each leaf turn-by-turn using LCP-aligned - drop-and-replace: the cumulative tokens emitted so far are clamped to - the longest common prefix with the next turn's prompt; any prior tokens - past that LCP (the TITO drift suffix, including the previous turn's - response if it lands in the drift region) are DROPPED along with their - logprobs, then ``prompt[LCP:]`` is appended as loss_mask=0 (chat - template's authoritative re-rendering wins), then the current turn's - ``response`` is appended as loss_mask=1 with real logprobs. - -* Trade-off: previous-turn response tokens that fall inside the drift - region lose their training signal. In exchange, the final tokens - sequence matches what the live model actually conditioned on for every - later turn — logprobs stay coherent, no duplicated-content forks, no - reliance on chat_template being position-invariant. - -* Drift fork (gated by ``drift_fork_min_loss_tokens``, default 1024 — ON - by default): when a drift would drop >= N loss_mask=1 tokens, the leaf - FORKS. This is the primary drift path. We emit an extra synthetic - "drift_fork" Sample at drain time alongside the main Sample. - Fork tokens = cumulative pre-drop; fork loss_mask is COMPLEMENTARY — 1 - only at positions that the main leaf is about to drop, 0 elsewhere. Fork - reward = main-leaf share; fork group_id = main-leaf group_id. Fork ∪ main - on loss_mask=1 tokens never overlap and their union equals the virtual - no-drift trajectory. The forked drift is NOT counted in the main sample's - ``tito_dropped_*`` (it wasn't truly lost). - - When the drift would drop < N loss_mask=1 tokens (or is a pure-prompt - drift losing 0 loss tokens), the secondary DROP path applies instead: - drop-and-replace on the main leaf only, accounted in ``tito_dropped_*``. - -* On drift, ``Sample.metadata`` records: - ``tito_dropped_tokens`` — total tokens dropped (NOT including - drifts that produced a fork) - ``tito_dropped_turns`` — number of turns that triggered a drop - ``tito_drift_forks_emitted`` — set on main leaf when >=1 drift-fork - sibling was emitted for the same leaf - ``tito_drift_fork`` — True on a drift-fork Sample - ``tito_drift_fork_at_turn`` — turn index whose drift triggered it - ``tito_drift_fork_loss_tokens`` — count of loss_mask=1 tokens in fork - -* Fork-merge rescue (gated by ``fork_merge_max_response_tokens``, default - 1024 — ON by default; set <=0 to disable): a routing-time mechanism, - independent of the - linearization-time drift_fork/drop above. When DFS breaks at an assistant - group (a later replay reformats an earlier assistant message, e.g. - tool_call arg ordering or whitespace) and exactly one leaf sibling carries - a per-turn response shorter than the threshold, the rewrite is collapsed - onto that sibling instead of spawning a new sibling subtree. The merged - sibling's stale response then enters trajectories with loss_mask=0, - recorded in ``fork_merge_masked_tokens`` / ``fork_merge_turns``. drift_fork - handles token alignment; fork-merge prevents a routing fork from forming. +* ``get_trajectory`` linearizes each leaf turn-by-turn with a STRICT + exact-prefix contract: walking the leaf's assistant chain root→leaf, the + cumulative ``(prompt + response)`` tokens emitted so far MUST be an exact + prefix of the next turn's ``turn_prompt_ids``. When it is, the new prompt + tail ``prompt[len(cumulative):]`` is appended as loss_mask=0, then the + turn's ``response`` is appended as loss_mask=1 with real logprobs. + +* When the prefix does NOT match, the upstream tokenization drifted (the + same history re-tokenized differently across turns). That is a bug to + surface, not to paper over: ``get_trajectory`` raises ``ValueError`` with + the sid, turn_index, common-prefix length, and drift size so the + offending turn is locatable. Note a drift introduced at an early turn can + surface several turns later — the prefix check catches it whenever the + re-rendered early region first diverges from the accumulated tokens. + + (History: an earlier design tolerated drift via LCP drop-and-replace plus + optional drift-fork / drop-accounting / fork-merge. That machinery is + removed here in favor of failing loudly; it can be re-added as an explicit + layer later if real drift turns out to be unavoidable.) """ from __future__ import annotations @@ -70,7 +40,6 @@ import json import logging from collections.abc import Iterator -from dataclasses import dataclass, field from typing import Any from slime.agent.adapters.common import TurnRecord @@ -106,6 +75,7 @@ class Node: "metadata", "parent", "children", + "match_key", # per-turn snapshot (assistant leaves) "turn_prompt_ids", "turn_response_ids", @@ -127,6 +97,10 @@ def __init__( self.metadata = dict(metadata or {}) self.parent: Node | None = parent self.children: list[Node] = [] + # messages is immutable after construction (only children are appended, + # never the message list itself), so the routing key is computed once + # here and reused on every descent instead of re-serializing per turn. + self.match_key = node_match_key(self.messages) # per-turn snapshot self.turn_prompt_ids: list[int] | None = None self.turn_response_ids: list[int] | None = None @@ -162,7 +136,7 @@ def leaves(self) -> Iterator[Node]: # =========================================================================== -# node_match_key + role-grouping helpers +# node_match_key + message helpers # =========================================================================== @@ -176,59 +150,6 @@ def node_match_key(messages: list[dict[str, Any]]) -> str: return json.dumps(messages, sort_keys=True, ensure_ascii=False) -@dataclass -class _PromptGroup: - role: str - messages: list[dict[str, Any]] = field(default_factory=list) - - -@dataclass -class _DriftFork: - """Pre-drop fork collected when TITO drift would lose >= threshold loss tokens.""" - - tokens: list[int] - # Complementary mask: 0 at positions main leaf keeps, 1 at positions main leaf will drop. - loss_mask: list[int] - logprobs: list[float] - # asst.turn_index whose prompt triggered the drift. - drift_turn_index: int | None - # finish_reason of the prior assistant turn — describes "what the prefix - # looked like" before drift. Captured at fork creation so the build - # path doesn't have to look back into the chain. - prev_finish_reason: str | None - - -@dataclass -class _LeafAccum: - """Result of walking one leaf's assistant chain.""" - - tokens: list[int] = field(default_factory=list) - loss_mask: list[int] = field(default_factory=list) - logprobs: list[float] = field(default_factory=list) - drift_forks: list[_DriftFork] = field(default_factory=list) - # accounting (emitted to main sample's metadata when > 0) - dropped_tokens: int = 0 - dropped_turns: int = 0 - fork_merge_masked_tokens: int = 0 - fork_merge_turns: int = 0 - - -def _group_messages_by_role( - messages: list[dict[str, Any]], -) -> list[_PromptGroup]: - groups: list[_PromptGroup] = [] - for m in messages: - role = m.get("role") - if not isinstance(role, str): - logger.warning("skipping message without string role: %r", m) - continue - if groups and groups[-1].role == role: - groups[-1].messages.append(m) - else: - groups.append(_PromptGroup(role=role, messages=[m])) - return groups - - def _lcp_len(a: list[int], b: list[int]) -> int: """Length of the longest common prefix between two int lists.""" n = min(len(a), len(b)) @@ -247,28 +168,11 @@ class TrajectoryManager: """Per-sid trajectory tree manager. See module docstring for the C-plan invariants. Each ``append_turn`` - mounts >=0 prompt nodes (under the deepest matching ancestor) + exactly - 1 assistant leaf carrying that turn's sglang snapshot. + mounts >=0 prompt nodes (one per message, under the deepest matching + ancestor) + exactly 1 assistant leaf carrying that turn's sglang snapshot. """ - def __init__( - self, - *, - drift_fork_min_loss_tokens: int = 1024, - fork_merge_max_response_tokens: int = 1024, - ) -> None: - # Drift-fork threshold (loss_mask=1 token count inside drift suffix). - # When a drift would drop >= this many loss tokens, fork instead of - # dropping. Always an int (default 1024 — drift-fork ON by default). - # Set <=0 to effectively disable (combined with the drift_loss_tokens>0 - # guard, only true drifts above the threshold ever fork). - self._fork_threshold: int = drift_fork_min_loss_tokens - # Fork-merge threshold: when DFS would break at an assistant group and - # exactly one non-leaf assistant sibling has turn_response_ids length - # STRICTLY LESS than this value, collapse the would-be fork onto that - # sibling (its response then enters trajectories with loss_mask=0). - # Default 1024 — fork-merge ON by default; set <=0 to disable. - self._fork_merge_threshold: int = fork_merge_max_response_tokens + def __init__(self) -> None: self._trees: dict[str, Node] = {} self._turn_count: dict[str, int] = {} @@ -300,27 +204,70 @@ def append_turn( ) root = self._trees.setdefault(sid, Node()) - groups = _group_messages_by_role(prompt_messages) - cur, i = self._find_mount_point(root, groups) - cur, i = self._try_fork_merge_assistant(sid, cur, groups, i) - cur = self._mount_prompt_groups(cur, groups[i:], tools) + cur, i = self._find_mount_point(root, prompt_messages) + cur = self._mount_prompt_messages(cur, prompt_messages[i:], tools) self._attach_assistant_leaf(sid, cur, turn=turn, response_message=response_message, metadata=metadata) - def _find_mount_point(self, root: Node, groups: list[_PromptGroup]) -> tuple[Node, int]: - """DFS down the existing tree by ``(role, node_match_key)``. + def get_trajectory( + self, + sid: str, + *, + base_sample=None, + reward: float = 0.0, + extra_metadata: dict[str, Any] | None = None, + ) -> list: + """Drain a sid into slime ``Sample`` objects, then drop the session. + + ``get_trajectory`` is the lifecycle boundary where the message routing + tree is linearized into token-normalized ``Sample`` objects. Each + routing leaf yields exactly one Sample. ``reward`` is split evenly + across all emitted samples. The sid is consumed: a second call for the + same sid returns ``[]``. + """ + if base_sample is None: + base_sample = Sample(index=0, prompt="") + + root = self._trees.get(sid) + if root is None: + return [] + + samples: list[Sample] = [] + for routing_leaf in root.leaves(): + if routing_leaf.is_root: + continue + chain = routing_leaf.path_from_root() + samples.extend(self._chain_to_sample(sid, chain, base_sample=base_sample, extra_metadata=extra_metadata)) + + # Reward is split evenly across every emitted sample (one per leaf); the + # token-weighted reducer downstream then gives each loss token the + # trajectory's full R. Assigned after the fact so the per-leaf builder + # stays reward-agnostic. + per_sample_reward = (reward / len(samples)) if samples else 0.0 + for s in samples: + s.reward = per_sample_reward + + self._trees.pop(sid, None) + self._turn_count.pop(sid, None) + return samples + + # -------------------- internals ---------------------------------------- + + def _find_mount_point(self, root: Node, messages: list[dict[str, Any]]) -> tuple[Node, int]: + """DFS down the existing tree by ``(role, node_match_key)``, per message. Returns ``(cur, i)``: ``cur`` is the deepest node whose path matches - ``groups[:i]`` exactly; ``i`` is the index into ``groups`` of the first - group that diverges from anything mounted so far (i.e., where this - turn's new content begins). + ``messages[:i]`` exactly; ``i`` is the index into ``messages`` of the + first message that diverges from anything mounted so far (i.e., where + this turn's new content begins). """ cur = root i = 0 - while i < len(groups): - g_key = node_match_key(groups[i].messages) + while i < len(messages): + m = messages[i] + m_key = node_match_key([m]) match = next( - (c for c in cur.children if c.role == groups[i].role and node_match_key(c.messages) == g_key), + (c for c in cur.children if c.role == m.get("role") and c.match_key == m_key), None, ) if match is None: @@ -329,88 +276,28 @@ def _find_mount_point(self, root: Node, groups: list[_PromptGroup]) -> tuple[Nod i += 1 return cur, i - def _try_fork_merge_assistant(self, sid: str, cur: Node, groups: list[_PromptGroup], i: int) -> tuple[Node, int]: - """Optionally collapse an assistant-rewrite onto a single short leaf sibling. - - The typical claude-code pattern is: a later replay reformats an earlier - assistant message (e.g. tool_call arg ordering, whitespace), which - breaks DFS at that assistant group. Without rescue, every such reformat - spawns a new sibling subtree; with rescue, when the existing sibling's - per-turn response is short enough that masking it out is the cheaper - trade-off, we collapse onto that sibling and mark it for mask=0 at - linearization. - """ - if self._fork_merge_threshold <= 0: - return cur, i # feature off - if i >= len(groups) or groups[i].role != "assistant": - return cur, i # feature on, but this turn isn't an asst rewrite - - candidates = [ - c - for c in cur.children - if c.role == "assistant" - # Real rewrite footprint = the original turn's leaf assistant - # node: it was inserted by step 3 of a prior turn (so - # turn_response_ids is populated), and no later turn has - # extended it (so it is still a leaf). Once a rewrite collapses - # onto it via this rescue, the new turn's user/tool + asst leaf - # are appended underneath as children — so the merge target - # MUST be a leaf at decision time, otherwise it has already - # diverged into mixed subchains and merging would tangle them. - and not c.children - and c.turn_response_ids is not None - and len(c.turn_response_ids) < self._fork_merge_threshold - ] - if len(candidates) != 1: - if len(candidates) >= 2: - # Legacy / pathological state: fork_merge wasn't on during - # earlier rewrites, or threshold was widened mid-session. - # Don't pick arbitrarily — fork as usual and surface a hint. - logger.warning( - "append_turn(sid=%s turn=%s): multiple eligible fork-merge " - "candidates (%d), refusing to merge — likely a legacy state " - "from a prior run without fork_merge enabled; forking instead.", - sid, - self._turn_count.get(sid, 0) + 1, - len(candidates), - ) - return cur, i - - sib = candidates[0] - masked = len(sib.turn_response_ids or []) - logger.warning( - "append_turn(sid=%s turn=%s): fork-merging assistant rewrite " - "into existing sibling (turn_index=%s, masked_response_tokens=%d)", - sid, - self._turn_count.get(sid, 0) + 1, - sib.turn_index, - masked, - ) - sib.metadata["fork_merged"] = True - sib.metadata["fork_merge_masked_tokens"] = masked - sib.messages = list(groups[i].messages) - return sib, i + 1 - - def _mount_prompt_groups( + def _mount_prompt_messages( self, cur: Node, - remaining_groups: list[_PromptGroup], + remaining_messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None, ) -> Node: - """Attach each remaining prompt group as a routing node under ``cur``. + """Attach each remaining prompt message as a routing node under ``cur``. - Token attribution happens at get_trajectory time, not here. The tools + One node per message (``node.messages`` is a singleton list). Token + attribution happens at get_trajectory time, not here. The tools metadata is placed only on the FIRST system node on the path — ``_first_system_already_set(cur)`` walks ``cur → root`` looking for a system ancestor that already carries it, and ``cur`` here is the deepest node from descent (+ optional merge), so the walk sees every ancestor that's already mounted. """ - for g in remaining_groups: + for m in remaining_messages: + role = m.get("role") md: dict[str, Any] = {} - if g.role == "system" and tools is not None and not self._first_system_already_set(cur): + if role == "system" and tools is not None and not self._first_system_already_set(cur): md["tools"] = list(tools) - cur = cur.add_child(Node(role=g.role, messages=list(g.messages), metadata=md)) + cur = cur.add_child(Node(role=role, messages=[m], metadata=md)) return cur def _attach_assistant_leaf( @@ -436,129 +323,103 @@ def _attach_assistant_leaf( cur.add_child(asst) self._turn_count[sid] = asst.turn_index - def get_trajectory( + def _chain_to_sample( self, sid: str, + chain: list[Node], *, - base_sample=None, - reward: float = 0.0, - extra_metadata: dict[str, Any] | None = None, - ) -> list: - """Drain a sid into slime ``Sample`` objects, then drop the session. + base_sample: Sample, + extra_metadata: dict[str, Any] | None, + ) -> list[Sample]: + """Linearize one root→leaf chain into a single Sample (strict exact-prefix). - ``get_trajectory`` is the lifecycle boundary where the message routing - tree is linearized into token-normalized ``Sample`` objects. Each - routing leaf yields one main Sample plus one extra Sample per drift - fork. ``reward`` is split evenly across all emitted samples. The sid is - consumed: a second call for the same sid returns ``[]``. + Walk the chain's assistant nodes root→leaf, accumulating tokens. Each + turn's cumulative ``(prompt + response)`` so far MUST be an exact prefix + of the next turn's ``turn_prompt_ids``; otherwise the upstream + tokenization drifted and we raise (see module docstring). Reward is left + at 0.0 here and assigned by the caller. """ - if base_sample is None: - base_sample = Sample(index=0, prompt="") - - root = self._trees.get(sid) - if root is None: - return [] - - samples: list[Sample] = [] - for routing_leaf in root.leaves(): - if routing_leaf.is_root: - continue - samples.extend( - self._normalize_routing_leaf(sid, routing_leaf, base_sample=base_sample, extra_metadata=extra_metadata) - ) + # Only assistant leaves carrying this turn's sglang snapshot participate. + # Routing assistant nodes mounted from prior-turn replay (turn_prompt_ids + # is None) carry no token signal and are skipped. + asst_chain = [n for n in chain if n.role == "assistant" and n.turn_prompt_ids is not None] - # Reward is split evenly across every emitted sample (main + forks); - # the token-weighted reducer downstream then gives each loss token the - # trajectory's full R. Assigned after the fact so the per-leaf builder - # stays reward-agnostic. - per_sample_reward = (reward / len(samples)) if samples else 0.0 - for s in samples: - s.reward = per_sample_reward + tokens: list[int] = [] + loss_mask: list[int] = [] + logprobs: list[float] = [] + for asst in asst_chain: + prompt = asst.turn_prompt_ids or [] + response = asst.turn_response_ids or [] + response_logprobs = asst.turn_response_logprobs - self._trees.pop(sid, None) - self._turn_count.pop(sid, None) - return samples + # Strict prefix check: tokens accumulated so far must be an exact + # prefix of this turn's prompt. For the first turn `tokens` is empty + # so this trivially holds. + n = len(tokens) + if prompt[:n] != tokens: + self._raise_prefix_drift(sid, asst_chain, asst, tokens, prompt) - # -------------------- internals ---------------------------------------- + new_prompt = prompt[n:] + tokens.extend(new_prompt) + loss_mask.extend([0] * len(new_prompt)) + logprobs.extend([0.0] * len(new_prompt)) - def _normalize_routing_leaf( - self, - sid: str, - leaf: Node, - *, - base_sample: Sample, - extra_metadata: dict[str, Any] | None, - ) -> list[Sample]: - """Linearize one routing leaf into its main Sample (+ drift-fork Samples). + tokens.extend(response) + loss_mask.extend([1] * len(response)) + logprobs.extend(response_logprobs if response_logprobs is not None else [0.0] * len(response)) - Drift forks come first, then the main sample, matching the original - drain order. Reward is left at 0.0 here and assigned by the caller. - """ - chain = leaf.path_from_root() - # Only assistant leaves carrying this turn's sglang snapshot - # participate in TITO accumulation. Routing assistant nodes mounted - # from prior-turn replay (turn_prompt_ids is None) carry no token - # signal and would otherwise be misread as a full-trajectory drift. - asst_chain = [n for n in chain if n.role == "assistant" and n.turn_prompt_ids is not None] - accum = self._accumulate_chain(sid, asst_chain) - first_sys = next((n for n in chain if n.role == "system"), None) - base_md = {"tools": first_sys.metadata.get("tools") if first_sys else None} first_prompt_len = len(asst_chain[0].turn_prompt_ids or []) if asst_chain else 0 - - # Build one Sample from a linearized segment (a _DriftFork or the accum - # itself — both expose tokens/loss_mask/logprobs). The base/extra/clamp - # args are constant across segments, so this closure carries them. - def build(seg: _DriftFork | _LeafAccum, leaf_md: dict[str, Any]) -> Sample: - return self._build_leaf_sample( + return [ + self._build_leaf_sample( base_sample=base_sample, extra_metadata=extra_metadata, - leaf_metadata={**base_md, **leaf_md}, - tokens=seg.tokens, - loss_mask=seg.loss_mask, - logprobs=seg.logprobs, + tokens=tokens, + loss_mask=loss_mask, + logprobs=logprobs, first_prompt_len=first_prompt_len, ) - - # Drift forks first, then the main sample — matching the original drain - # order. A fork's loss_mask is complementary (0 on [0:L], L >= - # first_prompt_len), so the response-region clamp never touches a loss=1 - # position and sum(loss_mask) is the final loss-token count. - samples = [ - build( - fork, - { - "finish_reason": fork.prev_finish_reason, - "tito_drift_fork": True, - "tito_drift_fork_at_turn": fork.drift_turn_index, - "tito_drift_fork_loss_tokens": sum(fork.loss_mask), - }, - ) - for fork in accum.drift_forks ] - samples.append(build(accum, self._main_leaf_metadata(accum, asst_chain))) - return samples @staticmethod - def _main_leaf_metadata(accum: _LeafAccum, asst_chain: list[Node]) -> dict[str, Any]: - """Assemble the main sample's leaf metadata (conditional drift/merge keys).""" - last_asst = asst_chain[-1] if asst_chain else None - md: dict[str, Any] = {"finish_reason": last_asst.turn_finish_reason if last_asst else None} - if accum.dropped_tokens > 0: - md["tito_dropped_tokens"] = accum.dropped_tokens - md["tito_dropped_turns"] = accum.dropped_turns - if accum.drift_forks: - md["tito_drift_forks_emitted"] = len(accum.drift_forks) - if accum.fork_merge_masked_tokens > 0: - md["fork_merge_masked_tokens"] = accum.fork_merge_masked_tokens - md["fork_merge_turns"] = accum.fork_merge_turns - return md + def _raise_prefix_drift( + sid: str, + asst_chain: list[Node], + asst: Node, + tokens: list[int], + prompt: list[int], + ) -> None: + """Raise on a TITO drift: accumulated tokens are not a prefix of prompt. + + Reports the common-prefix length and drift size, plus which earlier + assistant turn's prompt region the divergence falls in — a drift + introduced at an early turn can surface only when a later turn + re-renders that early region differently. + """ + L = _lcp_len(tokens, prompt) + drift = len(tokens) - L + # Locate which turn's prompt region L lands in, so an early-turn drift + # that surfaces several turns later is still attributable. + drift_in_turn = None + for prior in asst_chain: + if prior is asst: + break + if L < len(prior.turn_prompt_ids or []): + drift_in_turn = prior.turn_index + break + raise ValueError( + f"get_trajectory(sid={sid} turn={asst.turn_index}): TITO drift — " + f"accumulated tokens are not a prefix of this turn's prompt " + f"(common_prefix_len={L}, drift={drift} tokens; divergence falls in " + f"turn {drift_in_turn}'s prompt region). The same history " + f"re-tokenized differently across turns; refusing to silently " + f"drop/realign." + ) def _build_leaf_sample( self, *, base_sample: Sample, extra_metadata: dict[str, Any] | None, - leaf_metadata: dict[str, Any], tokens: list[int], loss_mask: list[int], logprobs: list[float], @@ -571,13 +432,19 @@ def _build_leaf_sample( ``response_length == len(loss_mask)`` and loss_mask/logprobs cover only the response region. ``reward`` is left at 0.0; the caller assigns the per-sample share. + + Sample metadata carries only ``extra_metadata`` (empty on the + production path): the per-row dataset metadata and per-turn tool / + finish_reason snapshot are intentionally NOT propagated onto the leaf + Sample. Dump/analysis tooling reads those off the tree nodes instead. """ - loss_resp, lp_resp = self._response_region(loss_mask, logprobs, first_prompt_len) - metadata = { - **(base_sample.metadata or {}), - **(extra_metadata or {}), - **leaf_metadata, - } + # Clamp loss_mask/logprobs to the response region: strip the leading + # first-turn prompt prefix so response_length == len(loss_mask), per the + # slime contract (see backends/megatron_utils/data.py:139, + # ray/rollout.py:695). + strip = min(first_prompt_len, len(loss_mask)) + loss_resp, lp_resp = loss_mask[strip:], logprobs[strip:] + metadata = dict(extra_metadata or {}) return Sample( index=base_sample.index, group_id=base_sample.group_id if base_sample.group_id is not None else base_sample.index, @@ -592,126 +459,6 @@ def _build_leaf_sample( metadata=metadata, ) - def _accumulate_chain(self, sid: str, asst_chain: list[Node]) -> _LeafAccum: - """Apply LCP drop-and-replace to an assistant chain in turn order. - - Mutates and returns ``accum``. See the module docstring for the - algorithm, drift-fork contract, and fork-merge masking rule. - """ - accum = _LeafAccum() - for k, asst in enumerate(asst_chain, start=1): - # Read-only views: only fed to _lcp_len / .extend(), never mutated - # in place, so no defensive copy is needed here (the leaf already - # owns isolated copies from _attach_assistant_leaf). - prompt_ids = asst.turn_prompt_ids or [] - response_ids = asst.turn_response_ids or [] - response_logprobs = asst.turn_response_logprobs - is_merged = bool(asst.metadata.get("fork_merged")) - - # k == 1 falls out of the general case: LCP([], prompt) == 0 so - # _fork_or_drop_drift trivially returns the full prompt with no drop. - prev_finish_reason = asst_chain[k - 2].turn_finish_reason if k >= 2 else None - emit_prompt = self._fork_or_drop_drift( - sid, - accum, - prev_finish_reason=prev_finish_reason, - drift_turn_index=asst.turn_index, - prompt=prompt_ids, - ) - - accum.tokens.extend(emit_prompt) - accum.loss_mask.extend([0] * len(emit_prompt)) - accum.logprobs.extend([0.0] * len(emit_prompt)) - - accum.tokens.extend(response_ids) - # fork-merged sibling: its response is "stale" — present in the tree - # only as a routing placeholder for the rewrites that collapsed onto - # it; mask it out of training. - accum.loss_mask.extend([0 if is_merged else 1] * len(response_ids)) - accum.logprobs.extend(response_logprobs if response_logprobs is not None else [0.0] * len(response_ids)) - - if is_merged: - accum.fork_merge_masked_tokens += len(response_ids) - accum.fork_merge_turns += 1 - return accum - - def _fork_or_drop_drift( - self, - sid: str, - accum: _LeafAccum, - *, - prev_finish_reason: str | None, - drift_turn_index: int | None, - prompt: list[int], - ) -> list[int]: - """Resolve a TITO drift between cumulative tokens and the next prompt. - - Compute LCP ``L``; if there's no drift suffix, just return the prompt - tail. Otherwise decide FORK vs DROP (fork is the primary path), then - truncate ``accum.{tokens,loss_mask,logprobs}`` to ``L`` and return - ``prompt[L:]`` for the caller to emit. - - FORK (drift loses >= ``self._fork_threshold`` loss_mask=1 tokens): - append a complementary-mask ``_DriftFork`` to ``accum`` so the - dropped training signal survives as a sibling output leaf. The - drift is NOT counted toward ``accum.dropped_*``. - DROP (below threshold, or a pure-prompt drift losing 0 loss tokens): - drop-and-replace only, counted toward ``accum.dropped_*``. - """ - L = _lcp_len(accum.tokens, prompt) - drift = len(accum.tokens) - L - if drift == 0: - return prompt[L:] - - drift_loss_tokens = sum(accum.loss_mask[L:]) - # PRIMARY: fork — losing >= threshold loss tokens is worth a - # complementary-mask fork leaf so the dropped signal survives. The - # drift_loss_tokens>0 guard keeps pure-prompt drift on the DROP path - # even if the threshold is set to <=0. - forked = drift_loss_tokens > 0 and drift_loss_tokens >= self._fork_threshold - if forked: - accum.drift_forks.append( - _DriftFork( - tokens=list(accum.tokens), - loss_mask=[0] * L + list(accum.loss_mask[L:]), - logprobs=[0.0] * L - + [(accum.logprobs[i] if accum.loss_mask[i] == 1 else 0.0) for i in range(L, len(accum.tokens))], - drift_turn_index=drift_turn_index, - prev_finish_reason=prev_finish_reason, - ) - ) - else: - # SECONDARY: drop-and-replace only. - accum.dropped_tokens += drift - accum.dropped_turns += 1 - - logger.warning( - "get_trajectory(sid=%s leaf turn=%s): TITO drift detected, " - "dropping %d prior tokens (incl. previous-turn response) to " - "realign with this turn's prompt%s", - sid, - drift_turn_index, - drift, - f"; forked {drift_loss_tokens} loss tokens" if forked else "", - ) - - del accum.tokens[L:] - del accum.loss_mask[L:] - del accum.logprobs[L:] - return prompt[L:] - - @staticmethod - def _response_region( - loss_mask: list[int], - logprobs: list[float], - first_prompt_len: int, - ) -> tuple[list[int], list[float]]: - """Strip the leading first-turn prompt prefix from loss_mask / logprobs - so what remains is the response-region view slime expects (see - slime/backends/megatron_utils/data.py:139, slime/ray/rollout.py:695).""" - strip = min(first_prompt_len, len(loss_mask)) - return loss_mask[strip:], logprobs[strip:] - @staticmethod def _first_system_already_set(start: Node) -> bool: """Walk start->root looking for a system node already carrying tools.""" @@ -727,6 +474,5 @@ def _first_system_already_set(start: Node) -> bool: "Node", "TrajectoryManager", "node_match_key", - "_group_messages_by_role", "_lcp_len", ] From 4fcbb24133af33682613d4db653e835fff602f8a Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Mon, 8 Jun 2026 11:00:09 +0000 Subject: [PATCH 11/28] feat(agent): assistant-rewrite merge to de-dilute reward Add the single tolerated exception to the strict exact-prefix TrajectoryManager contract: when cc re-renders a short prior assistant message (tool_call arg order / whitespace), DFS forks at that assistant and leaves the original short turn as a standalone stub leaf -> its own Sample, diluting the trajectory's evenly-split reward. _try_merge_assistant_rewrite absorbs such a rewrite onto the existing leaf when its response is short enough (fork_merge_max_response_tokens, default 1024), demoting that node to routing-only so it contributes 0 training tokens. Wire the threshold through Anthropic/OpenAI adapters and the coding_agent_rl generate entrypoint (env SLIME_FORK_MERGE_MAX_RESPONSE_TOKENS). --- examples/coding_agent_rl/generate.py | 18 ++--- slime/agent/adapters/anthropic.py | 8 ++- slime/agent/adapters/openai.py | 8 ++- slime/agent/trajectory_manager.py | 100 ++++++++++++++++++++++++++- 4 files changed, 123 insertions(+), 11 deletions(-) diff --git a/examples/coding_agent_rl/generate.py b/examples/coding_agent_rl/generate.py index 09a63a776c..a007bf40e9 100644 --- a/examples/coding_agent_rl/generate.py +++ b/examples/coding_agent_rl/generate.py @@ -97,16 +97,17 @@ def __init__(self, args) -> None: "Without it the sandbox cannot dial back and the rollout will " "silently abort." ) - # Kept for the staged re-add of drift tolerance: the strict - # exact-prefix TrajectoryManager no longer honors these, so warn loudly - # if an operator set them expecting fork/merge behavior. - drift_fork_threshold = os.environ.get("SLIME_DRIFT_FORK_MIN_LOSS_TOKENS") + # Assistant-rewrite merge threshold (see TrajectoryManager): when cc + # re-renders a short prior assistant, absorb it onto the existing leaf + # instead of forking a reward-diluting stub Sample. None -> manager + # default (1024); <=0 disables. fork_merge_threshold = os.environ.get("SLIME_FORK_MERGE_MAX_RESPONSE_TOKENS") - if drift_fork_threshold or fork_merge_threshold: + fork_merge_threshold = int(fork_merge_threshold) if fork_merge_threshold else None + # drift-fork remains unimplemented in the strict core; warn if set. + if os.environ.get("SLIME_DRIFT_FORK_MIN_LOSS_TOKENS"): logger.warning( - "[coding_agent_rl] SLIME_DRIFT_FORK_MIN_LOSS_TOKENS / " - "SLIME_FORK_MERGE_MAX_RESPONSE_TOKENS are set but currently " - "ignored: TrajectoryManager uses strict exact-prefix " + "[coding_agent_rl] SLIME_DRIFT_FORK_MIN_LOSS_TOKENS is set but " + "currently ignored: TrajectoryManager uses strict exact-prefix " "linearization and raises on TITO drift." ) self.adapter = AnthropicAdapter( @@ -114,6 +115,7 @@ def __init__(self, args) -> None: sglang_url=sglang_url, tool_parser=self.tool_parser, reasoning_parser=self.reasoning_parser, + fork_merge_max_response_tokens=fork_merge_threshold, ) # handler_cancellation=True so a client disconnect cancels the handler # coroutine, arming the fire-and-forget /abort_request inside the diff --git a/slime/agent/adapters/anthropic.py b/slime/agent/adapters/anthropic.py index 122e683c6c..bc563682d4 100644 --- a/slime/agent/adapters/anthropic.py +++ b/slime/agent/adapters/anthropic.py @@ -67,6 +67,7 @@ def __init__( tool_parser=None, reasoning_parser=None, max_turns_per_sid: int | None = None, + fork_merge_max_response_tokens: int | None = None, on_turn_appended: Callable[..., None] | None = None, ) -> None: super().__init__( @@ -76,7 +77,12 @@ def __init__( reasoning_parser=reasoning_parser, ) # ONE manager shared across all sids; per-sid trees live inside. - self.manager = TrajectoryManager() + # ``None`` means "caller did not specify" -> let TrajectoryManager use + # its own default for the assistant-rewrite merge threshold. + mgr_kwargs: dict[str, int] = {} + if fork_merge_max_response_tokens is not None: + mgr_kwargs["fork_merge_max_response_tokens"] = fork_merge_max_response_tokens + self.manager = TrajectoryManager(**mgr_kwargs) # Optional debug hook invoked after each successful append_turn. # Signature: (sid, prompt_messages, tools, response_message, # prompt_ids, response_ids, finish_reason) -> None. diff --git a/slime/agent/adapters/openai.py b/slime/agent/adapters/openai.py index c1e71d1f84..f051908a8f 100644 --- a/slime/agent/adapters/openai.py +++ b/slime/agent/adapters/openai.py @@ -74,6 +74,7 @@ def __init__( tool_parser=None, reasoning_parser=None, max_turns_per_sid: int | None = None, + fork_merge_max_response_tokens: int | None = None, on_turn_appended: Callable[..., None] | None = None, ) -> None: super().__init__( @@ -83,7 +84,12 @@ def __init__( reasoning_parser=reasoning_parser, ) # ONE manager shared across all sids; per-sid trees live inside. - self.manager = TrajectoryManager() + # ``None`` means "caller did not specify" -> let TrajectoryManager use + # its own default for the assistant-rewrite merge threshold. + mgr_kwargs: dict[str, int] = {} + if fork_merge_max_response_tokens is not None: + mgr_kwargs["fork_merge_max_response_tokens"] = fork_merge_max_response_tokens + self.manager = TrajectoryManager(**mgr_kwargs) # Optional debug hook invoked after each successful append_turn. # Signature mirrors AnthropicAdapter.on_turn_appended: # (sid, prompt_messages, tools, response_message, diff --git a/slime/agent/trajectory_manager.py b/slime/agent/trajectory_manager.py index 9ba327e17e..a10bfac6be 100644 --- a/slime/agent/trajectory_manager.py +++ b/slime/agent/trajectory_manager.py @@ -33,6 +33,19 @@ optional drift-fork / drop-accounting / fork-merge. That machinery is removed here in favor of failing loudly; it can be re-added as an explicit layer later if real drift turns out to be unavoidable.) + +* ONE tolerated exception to the strict contract: an assistant-rewrite merge. + cc sometimes re-renders a previously-recorded assistant message when feeding + it back as prompt (tool_call arg order, whitespace). The message no longer + matches, so DFS forks at that assistant — which does NOT raise (the rewrite + mounts as a routing-only node, skipped at linearization) but leaves the + original short turn as a standalone stub leaf -> its own Sample, diluting the + trajectory's evenly-split reward. ``_try_merge_assistant_rewrite`` absorbs + such a rewrite onto the existing leaf when its response is short enough + (``fork_merge_max_response_tokens``), demoting that node to routing-only so it + contributes 0 training tokens. Non-assistant mismatches, long responses, and + ambiguous cases are left to fork as usual; same-message-different-token drift + still raises. """ from __future__ import annotations @@ -172,7 +185,14 @@ class TrajectoryManager: ancestor) + exactly 1 assistant leaf carrying that turn's sglang snapshot. """ - def __init__(self) -> None: + def __init__(self, *, fork_merge_max_response_tokens: int = 1024) -> None: + # Rewrite-merge threshold: when DFS breaks at an assistant message and + # exactly one eligible short-response *leaf* sibling exists, absorb the + # rewrite onto it (demoted to routing-only -> 0 training tokens), + # compared against the abandoned turn's own turn_response_ids length. + # Default 1024 (ON); set <=0 to disable. This is the single tolerated + # exception to the strict exact-prefix contract; see module docstring. + self._fork_merge_threshold = fork_merge_max_response_tokens self._trees: dict[str, Node] = {} self._turn_count: dict[str, int] = {} @@ -206,6 +226,7 @@ def append_turn( root = self._trees.setdefault(sid, Node()) cur, i = self._find_mount_point(root, prompt_messages) + cur, i = self._try_merge_assistant_rewrite(sid, cur, prompt_messages, i) cur = self._mount_prompt_messages(cur, prompt_messages[i:], tools) self._attach_assistant_leaf(sid, cur, turn=turn, response_message=response_message, metadata=metadata) @@ -276,6 +297,83 @@ def _find_mount_point(self, root: Node, messages: list[dict[str, Any]]) -> tuple i += 1 return cur, i + def _try_merge_assistant_rewrite( + self, + sid: str, + cur: Node, + prompt_messages: list[dict[str, Any]], + i: int, + ) -> tuple[Node, int]: + """Absorb a short assistant-rewrite onto its existing node instead of forking. + + cc sometimes re-renders a previously-recorded assistant message when + feeding it back as prompt (tool_call arg order, whitespace). That breaks + DFS at the assistant and forks a fresh subtree, leaving the original + short turn as a standalone stub leaf -> its own Sample, diluting the + trajectory's evenly-split reward. Forking does NOT raise (the rewritten + message mounts as a routing-only node and is already skipped at + linearization); this merge is purely a reward-hygiene / de-fragmentation + optimization. + + When the diverging message is an assistant and exactly one eligible + *short-response leaf* sibling exists, adopt the rewritten message onto + that node and DEMOTE it to routing-only (clear its turn snapshot), so it + contributes 0 training tokens -- handled by the existing + ``turn_prompt_ids is not None`` filter in ``_chain_to_sample`` (no change + needed there). Any other mismatch (non-assistant message, long response, + non-leaf or ambiguous candidates) is left to fork as usual. + """ + if self._fork_merge_threshold <= 0: + return cur, i # feature off + if i >= len(prompt_messages) or prompt_messages[i].get("role") != "assistant": + return cur, i # genuine non-assistant history fork -> leave it + + candidates = [ + c + for c in cur.children + if c.role == "assistant" + # Leaf == rewrite of the immediately-previous assistant: no later + # turn has extended it yet. A non-leaf assistant has already grown a + # subchain; merging onto it would tangle that history. + and not c.children + # A real turn leaf carrying this turn's snapshot, not an already- + # demoted routing node (turn_prompt_ids cleared by a prior merge). + and c.turn_prompt_ids is not None and len(c.turn_response_ids or []) < self._fork_merge_threshold + ] + if len(candidates) != 1: + if len(candidates) >= 2: + # Ambiguous: don't pick arbitrarily — fork as usual and hint. + logger.warning( + "append_turn(sid=%s turn=%s): %d eligible rewrite-merge " + "candidates; forking instead (ambiguous mixed state).", + sid, + self._turn_count.get(sid, 0) + 1, + len(candidates), + ) + return cur, i + + sib = candidates[0] + # Observability breadcrumb only (NOT read by linearization). + sib.metadata["merged_rewrite"] = { + "abandoned_turn_index": sib.turn_index, + "abandoned_response_tokens": len(sib.turn_response_ids or []), + } + # Demote to routing-only: snapshot cleared -> skipped by the existing + # ``turn_prompt_ids is not None`` filter at linearization. Clearing + # turn_prompt_ids also prevents this node from being re-selected as a + # merge candidate on a later turn. + sib.turn_prompt_ids = None + sib.turn_response_ids = None + sib.turn_response_logprobs = None + sib.turn_finish_reason = None + sib.turn_index = None + # Adopt the rewritten message; the match_key cache MUST follow messages + # so a later turn's DFS descends through this (now rewritten) node + # instead of forking again. + sib.messages = [prompt_messages[i]] + sib.match_key = node_match_key(sib.messages) + return sib, i + 1 + def _mount_prompt_messages( self, cur: Node, From 0f1c5aa7b758d71bf8b4092741569d249ada0f21 Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Mon, 8 Jun 2026 11:11:12 +0000 Subject: [PATCH 12/28] feat(agent): TrajectoryManager re-accepts fork_merge_max_response_tokens --- slime/agent/trajectory_manager.py | 22 +- tests/test_agent/test_trajectory_manager.py | 1013 +++++++++++++++++++ 2 files changed, 1025 insertions(+), 10 deletions(-) create mode 100644 tests/test_agent/test_trajectory_manager.py diff --git a/slime/agent/trajectory_manager.py b/slime/agent/trajectory_manager.py index a10bfac6be..4053883f67 100644 --- a/slime/agent/trajectory_manager.py +++ b/slime/agent/trajectory_manager.py @@ -185,14 +185,16 @@ class TrajectoryManager: ancestor) + exactly 1 assistant leaf carrying that turn's sglang snapshot. """ - def __init__(self, *, fork_merge_max_response_tokens: int = 1024) -> None: - # Rewrite-merge threshold: when DFS breaks at an assistant message and - # exactly one eligible short-response *leaf* sibling exists, absorb the - # rewrite onto it (demoted to routing-only -> 0 training tokens), - # compared against the abandoned turn's own turn_response_ids length. - # Default 1024 (ON); set <=0 to disable. This is the single tolerated - # exception to the strict exact-prefix contract; see module docstring. - self._fork_merge_threshold = fork_merge_max_response_tokens + def __init__(self, *, fork_merge_max_response_tokens: int | None = None) -> None: + # Drift fork/replace threshold (see module docstring + spec). Only a + # case-B1 drift (divergence inside the immediately-previous turn's + # response region) compares its drift length against this value: + # drift < threshold -> replace (truncate + realign, dropped tail + # counted in tito_dropped_*); drift >= threshold -> fork. case A + # (prompt region) and case B2 (drift in an earlier turn's response) + # always fork regardless. <=0 forces B1 to fork too (max fidelity). + # ``None`` from the caller means "use the default". + self._fork_threshold: int = 1024 if fork_merge_max_response_tokens is None else fork_merge_max_response_tokens self._trees: dict[str, Node] = {} self._turn_count: dict[str, int] = {} @@ -323,7 +325,7 @@ def _try_merge_assistant_rewrite( needed there). Any other mismatch (non-assistant message, long response, non-leaf or ambiguous candidates) is left to fork as usual. """ - if self._fork_merge_threshold <= 0: + if self._fork_threshold <= 0: return cur, i # feature off if i >= len(prompt_messages) or prompt_messages[i].get("role") != "assistant": return cur, i # genuine non-assistant history fork -> leave it @@ -338,7 +340,7 @@ def _try_merge_assistant_rewrite( and not c.children # A real turn leaf carrying this turn's snapshot, not an already- # demoted routing node (turn_prompt_ids cleared by a prior merge). - and c.turn_prompt_ids is not None and len(c.turn_response_ids or []) < self._fork_merge_threshold + and c.turn_prompt_ids is not None and len(c.turn_response_ids or []) < self._fork_threshold ] if len(candidates) != 1: if len(candidates) >= 2: diff --git a/tests/test_agent/test_trajectory_manager.py b/tests/test_agent/test_trajectory_manager.py new file mode 100644 index 0000000000..3948f179a5 --- /dev/null +++ b/tests/test_agent/test_trajectory_manager.py @@ -0,0 +1,1013 @@ +"""Unit tests for src_v2.trajectory_manager (Plan C: token-faithful). + +What we test: + (1) DFS merge only on (role, node_match_key): same prefix in messages + space always lands on the same path regardless of prompt_ids drift. + (2) get_trajectory linearization: turn 1 = full prompt + response; + turn k>=2 = strict exact-prefix append; tokens / loss_mask / + logprobs all stay in sync. + (3) TITO drift handling: when turn k+1.prompt diverges mid-stream from + cumulative tokens, get_trajectory RAISES (no silent drop/realign). +""" + +from __future__ import annotations + +import json + +from slime.agent.adapters.common import TurnRecord # noqa: E402 +from slime.agent.trajectory_manager import TrajectoryManager, _lcp_len, node_match_key # noqa: E402 +from slime.utils.types import Sample # noqa: E402 + + +def _turn(prompt_ids, response_ids, *, finish_reason, logprobs=None): + """Helper: build the TurnRecord the way call_sglang_generate would. + + ``logprobs=None`` maps to an empty ``output_log_probs`` (the dataclass + default) so the manager treats this turn as carrying no logprob signal. + Pass an explicit list to attach per-token logprobs. + """ + return TurnRecord( + prompt_ids=list(prompt_ids), + output_ids=list(response_ids), + finish_reason=finish_reason, + output_log_probs=list(logprobs) if logprobs is not None else [], + ) + + +# --------------------------------------------------------------------------- +# Helper-level tests +# --------------------------------------------------------------------------- + + +def test_node_match_key_is_dict_internal_sort_only(): + a = [{"role": "u", "content": "x"}] + b = [{"content": "x", "role": "u"}] + assert node_match_key(a) == node_match_key(b) + + c = [{"role": "u", "content": "x"}, {"role": "u", "content": "y"}] + d = [{"role": "u", "content": "y"}, {"role": "u", "content": "x"}] + assert node_match_key(c) != node_match_key(d) + + e = [{"role": "assistant", "tool_calls": [{"id": "1", "type": "function"}]}] + f = [{"role": "assistant", "tool_calls": [{"type": "function", "id": "1"}]}] + assert node_match_key(e) == node_match_key(f) + print("PASS test_node_match_key_is_dict_internal_sort_only") + + +def test_lcp_len(): + assert _lcp_len([], []) == 0 + assert _lcp_len([1, 2, 3], []) == 0 + assert _lcp_len([], [1, 2, 3]) == 0 + assert _lcp_len([1, 2, 3], [1, 2, 3]) == 3 + assert _lcp_len([1, 2, 3], [1, 2, 4]) == 2 + assert _lcp_len([1, 2, 3, 4, 5], [1, 2, 3]) == 3 + print("PASS test_lcp_len") + + +def test_manager_accepts_fork_threshold(): + # default 1024 when unspecified + m_default = TrajectoryManager() + assert m_default._fork_threshold == 1024 + # explicit value honored + m_explicit = TrajectoryManager(fork_merge_max_response_tokens=256) + assert m_explicit._fork_threshold == 256 + # None -> default + m_none = TrajectoryManager(fork_merge_max_response_tokens=None) + assert m_none._fork_threshold == 1024 + print("PASS test_manager_accepts_fork_threshold") + + +# --------------------------------------------------------------------------- +# Fake tokenizer (kept only as a shape-matching prompt/response generator; +# trajectory_manager doesn't invoke it under plan C). +# --------------------------------------------------------------------------- + + +class FakeTokenizer: + ROLE_START = {"system": 9001, "user": 9002, "assistant": 9003, "tool": 9004} + ROLE_END = {"system": 9101, "user": 9102, "assistant": 9103, "tool": 9104} + + def apply_chat_template(self, messages, *, tools=None, add_generation_prompt=False, **kwargs): + out: list[int] = [] + for m in messages: + role = m["role"] + content = m.get("content") or "" + if not isinstance(content, str): + content = json.dumps(content, ensure_ascii=False) + out.append(self.ROLE_START[role]) + out.extend(ord(c) for c in content) + out.append(self.ROLE_END[role]) + if add_generation_prompt: + out.append(self.ROLE_START["assistant"]) + return out + + +def _render_prompt(messages, tools=None, tokenizer=None): + tok = tokenizer or FakeTokenizer() + return tok.apply_chat_template(messages, tools=tools, add_generation_prompt=True) + + +def _render_response(content_str, tokenizer=None): + tok = tokenizer or FakeTokenizer() + return [ord(c) for c in content_str] + [tok.ROLE_END["assistant"]] + + +# --------------------------------------------------------------------------- +# Plan-C semantics tests +# --------------------------------------------------------------------------- + + +SYSTEM_MSG = "You are a python coding agent." +TOOLS_OPENAI = [ + { + "type": "function", + "function": { + "name": "run_python", + "description": "Run python code.", + "parameters": {"type": "object", "properties": {"code": {"type": "string"}}}, + }, + }, +] + + +def _three_turn_session(tok): + """3-turn linear session via append_turn. Returns mgr, sid, per-turn (p,r).""" + mgr = TrajectoryManager() + sid = "three-turn" + + sys_msg = {"role": "system", "content": SYSTEM_MSG} + user1 = {"role": "user", "content": "Compute 2+2."} + asst1 = {"role": "assistant", "content": "Computing."} + tool1 = {"role": "tool", "content": "4"} + asst2 = {"role": "assistant", "content": "Answer is 4."} + + p1 = _render_prompt([sys_msg, user1], tokenizer=tok) + r1 = _render_response("Computing.", tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p1, r1, finish_reason="tool_calls", logprobs=[-0.5] * len(r1)), + prompt_messages=[sys_msg, user1], + tools=TOOLS_OPENAI, + response_message=asst1, + ) + + p2 = _render_prompt([sys_msg, user1, asst1, tool1], tokenizer=tok) + r2 = _render_response("Answer is 4.", tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p2, r2, finish_reason="stop", logprobs=[-0.4] * len(r2)), + prompt_messages=[sys_msg, user1, asst1, tool1], + tools=TOOLS_OPENAI, + response_message=asst2, + ) + return mgr, sid, [(p1, r1), (p2, r2)] + + +def test_append_single_turn_shapes_tree(): + tok = FakeTokenizer() + mgr = TrajectoryManager() + sid = "single" + sys_msg = {"role": "system", "content": "S"} + user1 = {"role": "user", "content": "u"} + p = _render_prompt([sys_msg, user1], tokenizer=tok) + r = _render_response("a", tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p, r, finish_reason="stop"), + prompt_messages=[sys_msg, user1], + tools=None, + response_message={"role": "assistant", "content": "a"}, + ) + chain = list(mgr._trees[sid].leaves())[0].path_from_root() + roles = [n.role for n in chain] + assert roles == ["system", "user", "assistant"], roles + asst = chain[-1] + assert asst.turn_index == 1 + assert asst.turn_prompt_ids == p + assert asst.turn_response_ids == r + assert asst.turn_finish_reason == "stop" + print("PASS test_append_single_turn_shapes_tree") + + +def test_append_three_turn_chain_no_fork(): + """3-turn session with consistent prompts -> exactly 1 leaf.""" + tok = FakeTokenizer() + mgr, sid, _ = _three_turn_session(tok) + leaves = [leaf for leaf in mgr._trees[sid].leaves() if not leaf.is_root] + assert len(leaves) == 1 + chain = leaves[0].path_from_root() + roles = [n.role for n in chain] + assert roles == ["system", "user", "assistant", "tool", "assistant"], roles + assert mgr.turn_count(sid) == 2 + print("PASS test_append_three_turn_chain_no_fork") + + +def test_fork_on_text_diff(): + """Different user content under shared sys -> 2 leaves, sys shared.""" + tok = FakeTokenizer() + mgr = TrajectoryManager() + sid = "fork-text" + sys_msg = {"role": "system", "content": "S"} + + for content in ["uA", "uB"]: + user = {"role": "user", "content": content} + p = _render_prompt([sys_msg, user], tokenizer=tok) + r = _render_response(content[-1], tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p, r, finish_reason="stop"), + prompt_messages=[sys_msg, user], + tools=None, + response_message={"role": "assistant", "content": content[-1]}, + ) + + root = mgr._trees[sid] + assert len(root.children) == 1, "sys node must be shared" + sys_node = root.children[0] + assert len(sys_node.children) == 2, "user level must fork" + leaves = [leaf for leaf in root.leaves() if not leaf.is_root] + assert len(leaves) == 2 + print("PASS test_fork_on_text_diff") + + +def test_no_fork_on_token_only_diff(): + """Plan C: same text but tampered prompt_ids -> NO fork (DFS ignores tokens). + + This is the load-bearing behavior change vs the old prefix-match design. + """ + tok = FakeTokenizer() + mgr = TrajectoryManager() + sid = "tokens-diff-only" + sys_msg = {"role": "system", "content": "S"} + user1 = {"role": "user", "content": "u"} + pa = _render_prompt([sys_msg, user1], tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(pa, _render_response("a", tokenizer=tok), finish_reason="stop"), + prompt_messages=[sys_msg, user1], + tools=None, + response_message={"role": "assistant", "content": "a"}, + ) + tampered = list(pa) + tampered[1] = tampered[1] ^ 1 + mgr.append_turn( + sid, + turn=_turn(tampered, _render_response("b", tokenizer=tok), finish_reason="stop"), + prompt_messages=[sys_msg, user1], + tools=None, + response_message={"role": "assistant", "content": "b"}, + ) + root = mgr._trees[sid] + # Same (sys, user) path -> shared, but two different assistant turns + # produce two assistant leaves under the same user node. + assert len(root.children) == 1 + sys_node = root.children[0] + assert len(sys_node.children) == 1 + user_node = sys_node.children[0] + assert len(user_node.children) == 2, "two distinct assistant turns hang off shared user" + leaves = [leaf for leaf in root.leaves() if not leaf.is_root] + assert len(leaves) == 2 + print("PASS test_no_fork_on_token_only_diff") + + +def test_cross_sid_isolation(): + tok = FakeTokenizer() + mgr = TrajectoryManager() + sys_msg = {"role": "system", "content": "S"} + for sid, content in [("sid-a", "uA"), ("sid-b", "uB")]: + user = {"role": "user", "content": content} + p = _render_prompt([sys_msg, user], tokenizer=tok) + r = _render_response(content[-1], tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p, r, finish_reason="stop"), + prompt_messages=[sys_msg, user], + tools=None, + response_message={"role": "assistant", "content": content[-1]}, + ) + assert len(list(mgr._trees["sid-a"].leaves())) == 1 + assert len(list(mgr._trees["sid-b"].leaves())) == 1 + print("PASS test_cross_sid_isolation") + + +def test_role_tool_in_chain(): + tok = FakeTokenizer() + mgr = TrajectoryManager() + sid = "tool-chain" + sys_msg = {"role": "system", "content": "S"} + user1 = {"role": "user", "content": "u"} + asst1 = {"role": "assistant", "content": "a1"} + tool_a = {"role": "tool", "content": "tA"} + tool_b = {"role": "tool", "content": "tB"} + asst2 = {"role": "assistant", "content": "a2"} + + p1 = _render_prompt([sys_msg, user1], tokenizer=tok) + r1 = _render_response("a1", tokenizer=tok) + p2 = _render_prompt([sys_msg, user1, asst1, tool_a, tool_b], tokenizer=tok) + r2 = _render_response("a2", tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p1, r1, finish_reason="stop"), + prompt_messages=[sys_msg, user1], + tools=None, + response_message=asst1, + ) + mgr.append_turn( + sid, + turn=_turn(p2, r2, finish_reason="stop"), + prompt_messages=[sys_msg, user1, asst1, tool_a, tool_b], + tools=None, + response_message=asst2, + ) + + chain = list(mgr._trees[sid].leaves())[0].path_from_root() + roles = [n.role for n in chain] + # Per-message routing: the two tool_results mount as two separate nodes. + assert roles == ["system", "user", "assistant", "tool", "tool", "assistant"], roles + assert chain[3].messages == [tool_a] + assert chain[4].messages == [tool_b] + print("PASS test_role_tool_in_chain") + + +def test_response_logprobs_length_mismatch_raises(): + tok = FakeTokenizer() + mgr = TrajectoryManager() + sys_msg = {"role": "system", "content": "S"} + user1 = {"role": "user", "content": "u"} + p = _render_prompt([sys_msg, user1], tokenizer=tok) + bad_turn = TurnRecord( + prompt_ids=p, + output_ids=[1, 2, 3], + finish_reason="stop", + output_log_probs=[-0.1, -0.2], + ) + try: + mgr.append_turn( + "x", + turn=bad_turn, + prompt_messages=[sys_msg, user1], + tools=None, + response_message={"role": "assistant", "content": ""}, + ) + except ValueError as e: + assert "output_log_probs" in str(e) + print("PASS test_response_logprobs_length_mismatch_raises") + return + raise AssertionError("expected ValueError") + + +def test_response_ids_empty_ok(): + tok = FakeTokenizer() + mgr = TrajectoryManager() + sys_msg = {"role": "system", "content": "S"} + user1 = {"role": "user", "content": "u"} + p = _render_prompt([sys_msg, user1], tokenizer=tok) + mgr.append_turn( + "x", + turn=_turn(p, [], finish_reason="stop"), + prompt_messages=[sys_msg, user1], + tools=None, + response_message=None, + ) + chain = list(mgr._trees["x"].leaves())[0].path_from_root() + asst = chain[-1] + assert asst.role == "assistant" + assert asst.turn_response_ids == [] + assert asst.turn_prompt_ids == p + assert asst.messages == [] + print("PASS test_response_ids_empty_ok") + + +# --------------------------------------------------------------------------- +# get_trajectory linearization (Plan C heart of the matter) +# --------------------------------------------------------------------------- + + +def test_get_trajectory_single_turn(): + tok = FakeTokenizer() + mgr = TrajectoryManager() + sid = "g1" + sys_msg = {"role": "system", "content": "S"} + user = {"role": "user", "content": "u"} + p = _render_prompt([sys_msg, user], tokenizer=tok) + r = _render_response("a", tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p, r, finish_reason="stop", logprobs=[-0.5] * len(r)), + prompt_messages=[sys_msg, user], + tools=TOOLS_OPENAI, + response_message={"role": "assistant", "content": "a"}, + ) + samples = mgr.get_trajectory(sid, base_sample=Sample(index=7, prompt="hi"), reward=1.0) + assert len(samples) == 1 + s = samples[0] + assert s.tokens == p + r + # slime contract: loss_mask / rollout_log_probs cover only the response + # region (tokens after the initial prompt) and len(loss_mask) == + # response_length. The first turn's prompt prefix is stripped. + assert s.loss_mask == [1] * len(r) + assert s.rollout_log_probs == [-0.5] * len(r) + assert s.response_length == len(r) + assert s.reward == 1.0 + # Leaf Sample metadata is intentionally empty: no dataset-row passthrough, + # no per-turn tools / finish_reason snapshot (those live on the tree nodes). + assert s.metadata == {} + print("PASS test_get_trajectory_single_turn") + + +def test_get_trajectory_clean_multiturn(): + """Clean 2-turn session (no drift) linearizes as turn1 prompt+resp then + turn2 (prompt - LCP) + resp, with full coherent loss_mask / logprobs.""" + tok = FakeTokenizer() + mgr, sid, turns = _three_turn_session(tok) + samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + s = samples[0] + + (p1, r1), (p2, r2) = turns + # LCP(p1+r1, p2) should equal len(p1)+len(r1) for our clean fake tokenizer + # (p2 starts exactly with p1 contents + asst response + tool block + new gen prompt) + L = _lcp_len(p1 + r1, p2) + assert L == len(p1) + len(r1), f"clean session LCP should equal cumulative, got {L}" + expected_tokens = p1 + r1 + p2[L:] + r2 + # loss_mask / rollout_log_probs are response-only (slime contract): + # turn1 prompt is stripped; turn1 response keeps mask=1, then the + # extra prompt slice (p2[L:]) is mask=0, then turn2 response is mask=1. + expected_loss = [1] * len(r1) + [0] * (len(p2) - L) + [1] * len(r2) + expected_logp = [-0.5] * len(r1) + [0.0] * (len(p2) - L) + [-0.4] * len(r2) + assert s.tokens == expected_tokens + assert s.loss_mask == expected_loss + assert s.rollout_log_probs == expected_logp + assert s.response_length == len(r1) + (len(p2) - L) + len(r2) + assert s.metadata == {} + print("PASS test_get_trajectory_clean_multiturn") + + +def test_get_trajectory_tito_drift_raises(): + """Strict exact-prefix: turn 2 prompt diverges mid-stream from cumulative + tokens. The accumulated turn-1 (prompt + response) is no longer a prefix of + turn 2's prompt, so get_trajectory must RAISE instead of dropping/realigning. + """ + tok = FakeTokenizer() + mgr = TrajectoryManager() + sid = "tito" + sys_msg = {"role": "system", "content": "S"} + user = {"role": "user", "content": "u"} + asst1 = {"role": "assistant", "content": "a1"} + tool = {"role": "tool", "content": "t"} + asst2 = {"role": "assistant", "content": "a2"} + + p1 = _render_prompt([sys_msg, user], tokenizer=tok) + r1 = _render_response("a1", tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p1, r1, finish_reason="tool_calls", logprobs=[-0.5] * len(r1)), + prompt_messages=[sys_msg, user], + tools=None, + response_message=asst1, + ) + # Build turn 2 prompt the "honest" way, then INJECT a synthetic divergence + # inside the assistant response region — simulating chat-template drift. We + # splice 3 fake tokens past the LCP so cumulative (p1+r1) is no longer a + # prefix of p2. + p2_honest = _render_prompt([sys_msg, user, asst1, tool], tokenizer=tok) + drift_at = len(p1) + 1 # inside r1 + p2 = list(p2_honest) + p2 = p2[:drift_at] + [77777, 77778, 77779] + p2[drift_at:] + r2 = _render_response("a2", tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p2, r2, finish_reason="stop", logprobs=[-0.4] * len(r2)), + prompt_messages=[sys_msg, user, asst1, tool], + tools=None, + response_message=asst2, + ) + + try: + mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + except ValueError as e: + msg = str(e) + assert "TITO drift" in msg, msg + assert f"turn={2}" in msg, msg # the turn whose prompt failed the check + assert "common_prefix_len" in msg, msg + print("PASS test_get_trajectory_tito_drift_raises") + return + raise AssertionError("expected ValueError on TITO drift") + + +def test_get_trajectory_tito_drift_late_surfacing_attributes_early_turn(): + """A drift introduced in an early turn's region but only re-rendered at a + later turn must still raise, and the message should attribute the divergence + to the early turn's prompt region (not just the failing turn). + """ + tok = FakeTokenizer() + mgr = TrajectoryManager() + sid = "tito-late" + sys_msg = {"role": "system", "content": "S"} + user = {"role": "user", "content": "u"} + asst1 = {"role": "assistant", "content": "a1"} + tool1 = {"role": "tool", "content": "t1"} + asst2 = {"role": "assistant", "content": "a2"} + tool2 = {"role": "tool", "content": "t2"} + asst3 = {"role": "assistant", "content": "a3"} + + # turn 1 + turn 2: clean continuation. + p1 = _render_prompt([sys_msg, user], tokenizer=tok) + r1 = _render_response("a1", tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p1, r1, finish_reason="tool_calls"), + prompt_messages=[sys_msg, user], + tools=None, + response_message=asst1, + ) + p2 = _render_prompt([sys_msg, user, asst1, tool1], tokenizer=tok) + r2 = _render_response("a2", tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p2, r2, finish_reason="tool_calls"), + prompt_messages=[sys_msg, user, asst1, tool1], + tools=None, + response_message=asst2, + ) + # turn 3: re-render the turn-1 region differently (splice a phantom token + # inside p1's range). Cumulative now diverges from p3 deep inside turn 1. + p3_honest = _render_prompt([sys_msg, user, asst1, tool1, asst2, tool2], tokenizer=tok) + drift_at = len(p1) - 1 # inside turn 1's prompt region + p3 = list(p3_honest) + p3 = p3[:drift_at] + [99999] + p3[drift_at:] + r3 = _render_response("a3", tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p3, r3, finish_reason="stop"), + prompt_messages=[sys_msg, user, asst1, tool1, asst2, tool2], + tools=None, + response_message=asst3, + ) + + try: + mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt="")) + except ValueError as e: + msg = str(e) + assert "TITO drift" in msg, msg + # The failing turn is turn 3, but the divergence falls in turn 1's region. + assert "turn=3" in msg, msg + assert "turn 1's prompt region" in msg, msg + print("PASS test_get_trajectory_tito_drift_late_surfacing_attributes_early_turn") + return + raise AssertionError("expected ValueError on late-surfacing TITO drift") + + +def test_get_trajectory_two_leaves_share_reward(): + """Forked tree (2 leaves) -> reward split evenly.""" + tok = FakeTokenizer() + mgr = TrajectoryManager() + sid = "split" + sys_msg = {"role": "system", "content": "S"} + for content in ["uA", "uB"]: + user = {"role": "user", "content": content} + p = _render_prompt([sys_msg, user], tokenizer=tok) + r = _render_response(content[-1], tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p, r, finish_reason="stop"), + prompt_messages=[sys_msg, user], + tools=None, + response_message={"role": "assistant", "content": content[-1]}, + ) + samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + assert len(samples) == 2 + assert all(s.reward == 1.0 for s in samples) + print("PASS test_get_trajectory_two_leaves_share_reward") + + +def test_drop_clears_sid(): + tok = FakeTokenizer() + mgr, sid, _ = _three_turn_session(tok) + assert mgr.has_session(sid) + mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt="")) + assert not mgr.has_session(sid) + assert mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt="")) == [] + print("PASS test_drop_clears_sid") + + +def test_debug_dump_shape(): + from tests.test_agent.test_claude_code_agent._dump_helpers import dump_tree_json, dump_tree_txt + + tok = FakeTokenizer() + mgr, sid, _ = _three_turn_session(tok) + txt = dump_tree_txt(mgr, sid) + assert isinstance(txt, str) and txt + for needle in ("session=", "[system]", "[user]", "[assistant]", "[tool]", "turns=2"): + assert needle in txt, f"missing {needle!r}" + # Plan C: assistant rows show turn= / prompt_ids= / response_ids= + assert "turn=1" in txt + assert "turn=2" in txt + assert "prompt_ids=" in txt + assert "response_ids=" in txt + + j = dump_tree_json(mgr, sid) + assert j["found"] is True and j["sid"] == sid and j["turns"] == 2 + assert j["nodes_total"] == 5 + + miss = dump_tree_txt(mgr, "no-such") + assert miss == "" + miss_j = dump_tree_json(mgr, "no-such") + assert miss_j == {"sid": "no-such", "found": False} + print("PASS test_debug_dump_shape") + + +def test_get_trajectory_skips_routing_assistant_in_drift_loop(): + """cc replays a foreign assistant the manager never recorded as a leaf; + per-message routing mounts it as a routing-only assistant that must be + filtered out of the strict-prefix walk. + + With per-message routing each prompt message mounts its own node, so the + previously-recorded ``asst2`` leaf (a single-message node) still matches + by ``(role, node_match_key)`` and the chain descends through it — turn 2 + stays on the main path and re-enters the strict-prefix check. Only the + extra ``foreign`` assistant, which the manager never saw via append_turn, + has no matching leaf and mounts as a routing-only assistant + (``turn_prompt_ids`` / ``turn_index`` both None). + + Such routing assistants must be filtered out of ``asst_chain`` (they + carry no per-turn snapshot). If one leaked into the strict prefix walk it + would look like a turn with an empty prompt and trip the exact-prefix + check — raising spuriously. This guards that the filter keeps the + otherwise-clean trajectory from raising, while turns 1/2/3 all stay in + the single linearized chain. + + Regression for the 20260604-120030 batch where 6 instances surfaced + routing assistants at depth 23-39, exact pattern: cc replays a prior + assistant message that the manager never recorded as its own leaf. + """ + tok = FakeTokenizer() + mgr = TrajectoryManager() + sid = "routing-asst" + sys_msg = {"role": "system", "content": "S"} + user1 = {"role": "user", "content": "u1"} + asst1 = {"role": "assistant", "content": "real-a1"} + tool1 = {"role": "tool", "content": "t1"} + asst2 = {"role": "assistant", "content": "real-a2"} + + # turn 1 + p1 = _render_prompt([sys_msg, user1], tokenizer=tok) + r1 = _render_response("real-a1", tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p1, r1, finish_reason="tool_calls"), + prompt_messages=[sys_msg, user1], + tools=None, + response_message=asst1, + ) + # turn 2 — clean continuation + p2 = _render_prompt([sys_msg, user1, asst1, tool1], tokenizer=tok) + r2 = _render_response("real-a2", tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p2, r2, finish_reason="tool_calls"), + prompt_messages=[sys_msg, user1, asst1, tool1], + tools=None, + response_message=asst2, + ) + # turn 3 — cc replays an extra prior asst that manager never saw via + # add_turn. Per-message routing matches asst2's own leaf and descends, so + # only the unmatched `foreign` message mounts as a routing-only assistant. + foreign = {"role": "assistant", "content": "foreign-msg"} + tool2 = {"role": "tool", "content": "t2"} + asst3 = {"role": "assistant", "content": "real-a3"} + p3 = _render_prompt([sys_msg, user1, asst1, tool1, asst2, foreign, tool2], tokenizer=tok) + r3 = _render_response("real-a3", tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p3, r3, finish_reason="stop"), + prompt_messages=[sys_msg, user1, asst1, tool1, asst2, foreign, tool2], + tools=None, + response_message=asst3, + ) + + # Per-message routing keeps everything on one path: asst2 matched, so no + # spurious fork. + leaves = [leaf for leaf in mgr._trees[sid].leaves() if not leaf.is_root] + assert len(leaves) == 1, [n.messages for n in leaves] + leaf3 = leaves[0] + assert leaf3.messages[0].get("content") == "real-a3" + chain = leaf3.path_from_root() + asst_nodes = [n for n in chain if n.role == "assistant"] + + # Exactly one routing assistant (the foreign replay); turns 1/2/3 keep + # their snapshots and stay in asst_chain. + routing_asst = [n for n in asst_nodes if n.turn_prompt_ids is None] + assert len(routing_asst) == 1, ( + f"expected exactly one routing assistant; got " + f"{[(n.turn_index, n.turn_prompt_ids is not None) for n in asst_nodes]}" + ) + assert routing_asst[0].turn_index is None + assert routing_asst[0].messages[0].get("content") == "foreign-msg" + snapshot_turns = [n.turn_index for n in asst_nodes if n.turn_prompt_ids is not None] + assert snapshot_turns == [1, 2, 3], snapshot_turns + + # The routing assistant is filtered out, so the strict prefix walk sees a + # clean chain (turns 1/2/3) and does NOT raise. + samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt="")) + assert len(samples) == 1 + print("PASS test_get_trajectory_skips_routing_assistant_in_drift_loop") + + +# --------------------------------------------------------------------------- +# Assistant-rewrite merge (single tolerated exception to strict exact-prefix) +# --------------------------------------------------------------------------- + + +def test_rewrite_merge_absorbs_short_assistant(): + """cc re-renders a short prior assistant; the manager absorbs the rewrite + onto the existing leaf (demoted to routing-only) instead of forking a + reward-diluting stub. One leaf, one Sample, original response not trained. + """ + tok = FakeTokenizer() + mgr = TrajectoryManager() # default threshold 1024 -> merge ON + sid = "rw-merge" + sys_msg = {"role": "system", "content": "S"} + user1 = {"role": "user", "content": "u"} + asst1 = {"role": "assistant", "content": "ok"} # short raw output + asst1_rw = {"role": "assistant", "content": "ok "} # cc-rewritten (whitespace) + tool1 = {"role": "tool", "content": "t"} + asst2 = {"role": "assistant", "content": "done"} + + p1 = _render_prompt([sys_msg, user1], tokenizer=tok) + r1 = _render_response("ok", tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p1, r1, finish_reason="tool_calls", logprobs=[-0.5] * len(r1)), + prompt_messages=[sys_msg, user1], + tools=None, + response_message=asst1, + ) + p2 = _render_prompt([sys_msg, user1, asst1_rw, tool1], tokenizer=tok) + r2 = _render_response("done", tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p2, r2, finish_reason="stop", logprobs=[-0.4] * len(r2)), + prompt_messages=[sys_msg, user1, asst1_rw, tool1], + tools=None, + response_message=asst2, + ) + + # Single chain (no fork): the rewrite was absorbed. + leaves = [leaf for leaf in mgr._trees[sid].leaves() if not leaf.is_root] + assert len(leaves) == 1, [n.messages for n in leaves] + chain = leaves[0].path_from_root() + assert [n.role for n in chain] == ["system", "user", "assistant", "tool", "assistant"] + + merged = chain[2] + # Demoted to routing-only: snapshot cleared, adopted the rewritten message. + assert merged.turn_prompt_ids is None + assert merged.turn_index is None + assert merged.messages == [asst1_rw] + assert merged.metadata["merged_rewrite"]["abandoned_turn_index"] == 1 + assert merged.metadata["merged_rewrite"]["abandoned_response_tokens"] == len(r1) + + # Linearization: only turn 2 participates; the abandoned turn-1 response is + # NOT trained. tokens == p2 + r2; loss covers only r2. + samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + s = samples[0] + assert s.tokens == p2 + r2 + assert s.loss_mask == [1] * len(r2) + assert s.rollout_log_probs == [-0.4] * len(r2) + assert s.reward == 1.0 + print("PASS test_rewrite_merge_absorbs_short_assistant") + + +def test_rewrite_merge_skips_long_assistant(): + """A long abandoned response (>= threshold) is NOT absorbed: it forks into + its own Sample (carrying enough real signal to train standalone). Forking + does not raise. + """ + tok = FakeTokenizer() + mgr = TrajectoryManager(fork_merge_max_response_tokens=2) # r1 (3 tok) >= 2 + sid = "rw-long" + sys_msg = {"role": "system", "content": "S"} + user1 = {"role": "user", "content": "u"} + asst1 = {"role": "assistant", "content": "ok"} + asst1_rw = {"role": "assistant", "content": "ok "} + tool1 = {"role": "tool", "content": "t"} + asst2 = {"role": "assistant", "content": "done"} + + p1 = _render_prompt([sys_msg, user1], tokenizer=tok) + r1 = _render_response("ok", tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p1, r1, finish_reason="tool_calls"), + prompt_messages=[sys_msg, user1], + tools=None, + response_message=asst1, + ) + p2 = _render_prompt([sys_msg, user1, asst1_rw, tool1], tokenizer=tok) + r2 = _render_response("done", tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p2, r2, finish_reason="stop"), + prompt_messages=[sys_msg, user1, asst1_rw, tool1], + tools=None, + response_message=asst2, + ) + + leaves = [leaf for leaf in mgr._trees[sid].leaves() if not leaf.is_root] + assert len(leaves) == 2, "long rewrite must fork, not merge" + samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + assert len(samples) == 2 # no raise + print("PASS test_rewrite_merge_skips_long_assistant") + + +def test_rewrite_merge_disabled_by_zero_threshold(): + """fork_merge_max_response_tokens=0 disables merge: every rewrite forks.""" + tok = FakeTokenizer() + mgr = TrajectoryManager(fork_merge_max_response_tokens=0) + sid = "rw-off" + sys_msg = {"role": "system", "content": "S"} + user1 = {"role": "user", "content": "u"} + asst1 = {"role": "assistant", "content": "ok"} + asst1_rw = {"role": "assistant", "content": "ok "} + tool1 = {"role": "tool", "content": "t"} + + p1 = _render_prompt([sys_msg, user1], tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p1, _render_response("ok", tokenizer=tok), finish_reason="tool_calls"), + prompt_messages=[sys_msg, user1], + tools=None, + response_message=asst1, + ) + p2 = _render_prompt([sys_msg, user1, asst1_rw, tool1], tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p2, _render_response("done", tokenizer=tok), finish_reason="stop"), + prompt_messages=[sys_msg, user1, asst1_rw, tool1], + tools=None, + response_message={"role": "assistant", "content": "done"}, + ) + leaves = [leaf for leaf in mgr._trees[sid].leaves() if not leaf.is_root] + assert len(leaves) == 2, "merge disabled -> rewrite forks" + print("PASS test_rewrite_merge_disabled_by_zero_threshold") + + +def test_rewrite_merge_ambiguous_candidates_fork(): + """Two eligible short-leaf assistant siblings -> ambiguous -> fork (no + arbitrary merge). + """ + tok = FakeTokenizer() + mgr = TrajectoryManager() + sid = "rw-ambig" + sys_msg = {"role": "system", "content": "S"} + user1 = {"role": "user", "content": "u"} + + # Two turns sharing the (sys, user) prefix produce two assistant leaves. + for content in ["a", "b"]: + p = _render_prompt([sys_msg, user1], tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p, _render_response(content, tokenizer=tok), finish_reason="stop"), + prompt_messages=[sys_msg, user1], + tools=None, + response_message={"role": "assistant", "content": content}, + ) + user_node = mgr._trees[sid].children[0].children[0] + assert len(user_node.children) == 2 + + # A third turn rewrites at the assistant slot -> two merge candidates. + asst_c = {"role": "assistant", "content": "c"} + tool1 = {"role": "tool", "content": "t"} + p3 = _render_prompt([sys_msg, user1, asst_c, tool1], tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p3, _render_response("d", tokenizer=tok), finish_reason="stop"), + prompt_messages=[sys_msg, user1, asst_c, tool1], + tools=None, + response_message={"role": "assistant", "content": "d"}, + ) + leaves = [leaf for leaf in mgr._trees[sid].leaves() if not leaf.is_root] + assert len(leaves) == 3, "ambiguous candidates must fork, not merge" + print("PASS test_rewrite_merge_ambiguous_candidates_fork") + + +def test_rewrite_merge_non_assistant_mismatch_forks(): + """A non-assistant divergence (different user) is left to fork; the merge + hook does not touch it even with merge enabled. + """ + tok = FakeTokenizer() + mgr = TrajectoryManager() + sid = "rw-nonasst" + sys_msg = {"role": "system", "content": "S"} + for content in ["uA", "uB"]: + user = {"role": "user", "content": content} + p = _render_prompt([sys_msg, user], tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p, _render_response(content[-1], tokenizer=tok), finish_reason="stop"), + prompt_messages=[sys_msg, user], + tools=None, + response_message={"role": "assistant", "content": content[-1]}, + ) + sys_node = mgr._trees[sid].children[0] + assert len(sys_node.children) == 2, "user-level divergence still forks" + print("PASS test_rewrite_merge_non_assistant_mismatch_forks") + + +def test_rewrite_merge_match_key_updated_so_next_turn_descends(): + """Regression: after merge, the node's cached match_key must follow the + adopted (rewritten) message so a LATER turn's DFS descends through it + instead of forking again. Also exercises clean strict-prefix continuation + across the merged node. + """ + tok = FakeTokenizer() + mgr = TrajectoryManager() + sid = "rw-matchkey" + sys_msg = {"role": "system", "content": "S"} + user1 = {"role": "user", "content": "u"} + asst1 = {"role": "assistant", "content": "ok"} + asst1_rw = {"role": "assistant", "content": "ok "} + tool1 = {"role": "tool", "content": "t1"} + asst2 = {"role": "assistant", "content": "second"} + tool2 = {"role": "tool", "content": "t2"} + asst3 = {"role": "assistant", "content": "third"} + + # turn 1: short assistant. + p1 = _render_prompt([sys_msg, user1], tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p1, _render_response("ok", tokenizer=tok), finish_reason="tool_calls"), + prompt_messages=[sys_msg, user1], + tools=None, + response_message=asst1, + ) + # turn 2: rewrite asst1 -> merge. + p2 = _render_prompt([sys_msg, user1, asst1_rw, tool1], tokenizer=tok) + r2 = _render_response("second", tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p2, r2, finish_reason="tool_calls", logprobs=[-0.4] * len(r2)), + prompt_messages=[sys_msg, user1, asst1_rw, tool1], + tools=None, + response_message=asst2, + ) + # turn 3: prompt carries the rewritten asst1_rw again; DFS must descend + # through the merged node (match_key updated) and not fork. + p3 = _render_prompt([sys_msg, user1, asst1_rw, tool1, asst2, tool2], tokenizer=tok) + r3 = _render_response("third", tokenizer=tok) + mgr.append_turn( + sid, + turn=_turn(p3, r3, finish_reason="stop", logprobs=[-0.3] * len(r3)), + prompt_messages=[sys_msg, user1, asst1_rw, tool1, asst2, tool2], + tools=None, + response_message=asst3, + ) + + leaves = [leaf for leaf in mgr._trees[sid].leaves() if not leaf.is_root] + assert len(leaves) == 1, ("match_key not updated -> spurious fork", [n.messages for n in leaves]) + + # Strict prefix holds across the merged node: turns 2 and 3 linearize into + # one clean Sample (no raise). The demoted turn-1 node is filtered out. + samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt="")) + assert len(samples) == 1 + s = samples[0] + L = _lcp_len(p2 + r2, p3) + assert s.tokens == p2 + r2 + p3[L:] + r3 + print("PASS test_rewrite_merge_match_key_updated_so_next_turn_descends") + + +# --------------------------------------------------------------------------- +# main +# --------------------------------------------------------------------------- + + +def main() -> None: + test_node_match_key_is_dict_internal_sort_only() + test_lcp_len() + test_append_single_turn_shapes_tree() + test_append_three_turn_chain_no_fork() + test_fork_on_text_diff() + test_no_fork_on_token_only_diff() + test_cross_sid_isolation() + test_role_tool_in_chain() + test_response_logprobs_length_mismatch_raises() + test_response_ids_empty_ok() + test_get_trajectory_single_turn() + test_get_trajectory_clean_multiturn() + test_get_trajectory_tito_drift_raises() + test_get_trajectory_tito_drift_late_surfacing_attributes_early_turn() + test_get_trajectory_two_leaves_share_reward() + test_drop_clears_sid() + test_debug_dump_shape() + test_get_trajectory_skips_routing_assistant_in_drift_loop() + test_rewrite_merge_absorbs_short_assistant() + test_rewrite_merge_skips_long_assistant() + test_rewrite_merge_disabled_by_zero_threshold() + test_rewrite_merge_ambiguous_candidates_fork() + test_rewrite_merge_non_assistant_mismatch_forks() + test_rewrite_merge_match_key_updated_so_next_turn_descends() + print("\nALL PLAN-C TESTS PASSED.") + + +if __name__ == "__main__": + main() From 473270256f0c50c0f2133c0c7a239163e1a73e02 Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Mon, 8 Jun 2026 13:57:44 +0000 Subject: [PATCH 13/28] docs(test): spec for TrajectoryManager e2e test script --- ...-08-trajectory-manager-e2e-tests-design.md | 164 ++++++++++++++++++ 1 file changed, 164 insertions(+) create mode 100644 docs/superpowers/specs/2026-06-08-trajectory-manager-e2e-tests-design.md diff --git a/docs/superpowers/specs/2026-06-08-trajectory-manager-e2e-tests-design.md b/docs/superpowers/specs/2026-06-08-trajectory-manager-e2e-tests-design.md new file mode 100644 index 0000000000..0321ba7ab6 --- /dev/null +++ b/docs/superpowers/specs/2026-06-08-trajectory-manager-e2e-tests-design.md @@ -0,0 +1,164 @@ +# TrajectoryManager 端到端测试脚本 — 设计文档 + +日期：2026-06-08 +状态：已批准设计，待实现 + +## 目标 + +写一个独立的端到端测试脚本，通过 `TrajectoryManager` 的两个公共接口 +`append_turn` 和 `get_trajectory`，从**数据结构角度**全面覆盖各种分叉情形： +prompt 分叉、assistant 分叉、token-ID 分叉，以及它们的组合。测试数据要 +**方便人类阅读**（语义化 token ID + 反查表），运行时既做严格 assertion， +又能打印可读的 tree / 线性化结果供人眼审查。 + +## 背景 + +`TrajectoryManager`（`slime/agent/trajectory_manager.py`）维护一棵 per-sid 的 +逐 message 路由树，并在 `get_trajectory` 时把每个 leaf 链线性化成 slime +`Sample`。它有两个正交的层： + +- **路由树层**（`append_turn`）：按 `(role, node_match_key)` 匹配，**只看 + message 身份**，与 token ID 无关。相同 message 前缀总落在同一路径。 +- **线性化层**（`get_trajectory`）：按 **token-ID 前缀**匹配累积 token，按 + 漂移位置走 case A / B1 / B2 路由（fork / replace），并做 cross-leaf dedup + 和 reward 均分。 + +现有 `tests/test_agent/test_trajectory_manager.py` 已有 28 个测试覆盖这些机制， +但用 `ord(char)` 映射 token、断言不易一眼看懂分叉点，且「两层同时发生」的 +组合格子覆盖较薄。本脚本是**独立新增**，与现有文件并存、互不依赖。 + +## 文件 + +`tests/test_agent/test_trajectory_manager_e2e.py` + +不改动 `trajectory_manager.py`；不依赖 sglang / 网络 / 真实 tokenizer；不碰 +现有 `test_trajectory_manager.py`。 + +## §1 基础设施 + +### 语义化 token 词表 + +token ID 用带语义段的小整数，看 assertion 一眼知道分叉在哪。每个 message +渲染成 `[START, ...body, END]`： + +``` +system : START=1000, END=1009, body 1001..1008 +user : START=2000, END=2009, body 2001..2008 +assistant: START=9000, END=9009, body 9001..9008 +tool : START=3000, END=3009, body 3001..3008 +gen-prompt 起手符 (add_generation_prompt): 9000 +漂移哨兵段: 7000..7099（不属于任何 role，dump 里一眼认出是人为漂移） +``` + +反查表 `TOKEN_NAMES: dict[int, str]` 把每个 ID 翻译成可读名（如 +`1000→""`、`2001→"u:compute"`、`7001→""`）。dump 打印时把 ID +序列翻译成可读字符串。 + +### 构造助手（薄封装，不引入 DSL） + +- `MsgTok`：给一个 message 分配固定的渲染 token 段，保证同一 message 在不同 + turn 渲染出相同 token（模拟干净 tokenizer）；漂移由测试显式注入。 +- `render_prompt(messages) -> list[int]`：拼成 token 序列（含 add_generation_prompt）。 +- `render_response(text) -> list[int]`：assistant 输出 token。 +- `turn(prompt_ids, response_ids, finish_reason, logprobs=None)`：构造 `TurnRecord`。 +- `drift(ids, at, sentinel=7001)`：在指定下标注入/替换哨兵 token，制造 + token-ID 漂移，返回新序列——让漂移点在测试代码里显式可见。 + +### 双态运行 + +- 每个 case 函数做严格 assertion。 +- `main()` 顺序跑所有 case；每个 case 跑完用 `_dump_helpers.dump_tree_txt` + 打印 tree，再打印线性化出的每个 Sample（token 翻译成可读名 + loss_mask + 对齐展示）。 +- 沿用现有文件的 `test_*` + `main()` 风格，可被 pytest 收集，也可直接 + `python -m` 跑供人眼审查。 + +## §2 Case 矩阵（按「层 × 分叉位置」组织） + +### 组 1 — 路由树层（断言 tree 形状） + +| # | Case | 分叉位置 | 预期树形 | +|---|------|---------|---------| +| 1.1 | 单 turn | 无 | system→user→assistant 一条链 | +| 1.2 | 干净多 turn（含 tool） | 无 | 一条链，每 message 一节点 | +| 1.3 | system 分叉 | system 不同 | root 下 2 子树 | +| 1.4 | user 分叉（共享 system） | user 不同 | system 共享，user 层 2 leaf | +| 1.5 | assistant message 分叉 | assistant 身份不同 | 共享 user，assistant 层 2 leaf | +| 1.6 | tool 分叉（同 assistant 不同 tool 结果） | tool 不同 | 共享 assistant，tool 层分叉 | +| 1.7 | token-only 漂移不分叉 | message 相同、prompt_ids 不同 | 树不分叉（DFS 忽略 token） | +| 1.8 | 多 tool message 逐节点挂载 | 一个 turn 多 tool | 每 tool 独立节点 | +| 1.9 | 跨 sid 隔离 | 不同 sid | 两棵独立树 | +| 1.10 | 空 response | 无 | assistant leaf messages=[]，turn_response_ids=[] | + +### 组 2 — 线性化层（断言 tokens/loss_mask/logprobs/reward） + +| # | Case | 触发 | 预期 | +|---|------|------|------| +| 2.1 | 单 turn 线性化 | — | tokens=p+r，loss 只覆盖 r | +| 2.2 | 干净多 turn 线性化 | LCP=cumulative | 1 Sample，prompt 尾 loss=0、resp loss=1 | +| 2.3 | drift case A（prompt 区漂移） | L 落在 prompt 区 | fork 成 2 Sample，不丢 token | +| 2.4 | drift case B1 短→replace | L 落在最近 resp 区、d0`（无全 mask 样本）。 +- token 期望值用构造助手拼出来，不手敲魔数。 + +### 打印格式（`main()` 时，每个 case 之后） + +``` +=== CASE 1.4 user 分叉（共享 system）=== +[tree] + +[samples] 2 个 + Sample#0 reward=1.0 resp_len=4 + tok : u:A 9001 + loss: 0 0 0 0 1 1 + Sample#1 ... +PASS 1.4 +``` + +token 与 loss_mask 上下对齐，漂移 token 显示成 ``，一眼看出分叉点 +和训练区。 + +### 结尾 + +`main()` 顺序跑全部 case，全 PASS 后打印 `ALL E2E CASES PASSED (N cases)`。 + +## 非目标（YAGNI） + +- 不引入 DSL / fluent builder。 +- 不改动 `trajectory_manager.py`。 +- 不依赖 sglang / 网络 / 真实 tokenizer。 +- 不修改现有 `test_trajectory_manager.py`。 From 4ea5d5ac6db1e0c4d64e18fbf11def6141d45f86 Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Mon, 8 Jun 2026 14:06:54 +0000 Subject: [PATCH 14/28] test(agent): end-to-end TrajectoryManager test matrix (append_turn/get_trajectory) 30 cases across 3 groups: routing-tree layer (message-identity forks), linearization layer (token-id drift A/B1/B2, dedup, reward split), and combined/stress (rewrite-merge, tree-fork+token-drift, deep multi-leaf, long mixed session). Semantic token vocab + reverse table for readable data; dual mode (strict assertions + human-readable tree/sample dumps). --- .../test_agent/test_trajectory_manager_e2e.py | 908 ++++++++++++++++++ 1 file changed, 908 insertions(+) create mode 100644 tests/test_agent/test_trajectory_manager_e2e.py diff --git a/tests/test_agent/test_trajectory_manager_e2e.py b/tests/test_agent/test_trajectory_manager_e2e.py new file mode 100644 index 0000000000..258bc5d4d0 --- /dev/null +++ b/tests/test_agent/test_trajectory_manager_e2e.py @@ -0,0 +1,908 @@ +"""End-to-end tests for TrajectoryManager via append_turn / get_trajectory. + +This script drives the two public interfaces of +``slime.agent.trajectory_manager.TrajectoryManager`` and exhaustively covers the +ways a trajectory can branch, organized as a two-axis matrix: + + * LAYER 1 — routing tree (append_turn). DFS merges on (role, node_match_key) + only, so MESSAGE IDENTITY决定 tree shape; token ids are irrelevant here. + * LAYER 2 — linearization (get_trajectory). TOKEN-ID prefix决定 how each leaf + chain becomes Samples (clean continuation / drift case A·B1·B2 / cross-leaf + dedup / reward split). + * COMBINED — both layers interacting (rewrite-merge, tree-fork + token-drift + stacked, deep multi-leaf dedup, long mixed session). + +Readability: + Token ids are SEMANTIC small integers (see TOKEN_NAMES). Each message renders + to ``[START, ...body, END]`` with a per-role band, so an id like 2001 reads as + ``u:compute`` and 7001 reads as ````. Expected token sequences are built + with the same render_* helpers used to feed append_turn, never hand-typed + magic numbers. + +Dual mode: + Every case is a ``test_*`` function doing strict assertions. ``main()`` runs + them all and, after each, prints the routing tree (token ids decoded to names) + and every linearized Sample with token / loss_mask aligned, so a human can read + exactly where each branch happened. Run with:: + + python -m tests.test_agent.test_trajectory_manager_e2e +""" + +from __future__ import annotations + +from tests.test_agent.test_claude_code_agent._dump_helpers import dump_tree_txt # noqa: E402 + +from slime.agent.adapters.common import TurnRecord # noqa: E402 +from slime.agent.trajectory_manager import TrajectoryManager, _lcp_len # noqa: E402 +from slime.utils.types import Sample # noqa: E402 + +# =========================================================================== +# §1 Semantic token vocabulary + reverse table +# =========================================================================== +# +# Per-role band. A message renders to [START, ...body, END]; the generation +# prompt appends the assistant START as the open-turn marker. + +_BANDS = { + "system": 1000, + "user": 2000, + "assistant": 9000, + "tool": 3000, +} +_GEN = _BANDS["assistant"] # add_generation_prompt marker +_DRIFT_BAND = 7000 + +# Reverse table: token id -> human-readable name. Filled lazily as messages are +# registered so dumps translate ids back to labels. +TOKEN_NAMES: dict[int, str] = {} +_ABBR = {"system": "sys", "user": "usr", "assistant": "ast", "tool": "tul"} +for _role, _base in _BANDS.items(): + TOKEN_NAMES[_base] = f"<{_ABBR[_role]}>" + TOKEN_NAMES[_base + 9] = f"" +TOKEN_NAMES[_GEN] = "" + + +def name_of(tok: int) -> str: + """Human-readable name for a token id (falls back to the raw int).""" + return TOKEN_NAMES.get(tok, str(tok)) + + +def _asst_body(label: str) -> int: + """Stable assistant body token for a response/message label. + + An assistant message replayed in a later prompt must render to the SAME + tokens the model generated for it, otherwise a clean continuation can never + hold (the cumulative prompt+response would not prefix the next prompt). So + both ``render_response`` and an assistant ``MsgTok`` derive their body token + from this one function, keyed on the label. + """ + body = _BANDS["assistant"] + 100 + (sum(ord(c) for c in label) % 800) + TOKEN_NAMES[body] = f"r:{label}" + return body + + +def render_ids(ids: list[int]) -> str: + """Decode an id list into a space-joined readable string.""" + return " ".join(name_of(t) for t in ids) + + +class MsgTok: + """A message bound to a fixed, deterministic token rendering. + + The same MsgTok always renders to the same token segment regardless of which + turn replays it (a clean tokenizer). Token-id drift is injected explicitly by + tests via ``drift`` — never by re-rendering. + """ + + _body_counter: dict[str, int] = {} + + def __init__(self, role: str, label: str) -> None: + self.role = role + self.label = label + base = _BANDS[role] + if role == "assistant": + # An assistant message must render to the same body token as the + # response it represents (label-keyed), so a replayed assistant in a + # later prompt token-matches the original generation -> clean + # continuation. See _asst_body. + self.body = _asst_body(label) + else: + # Allocate one stable body token per (role, label). Offset past the + # END marker (base+9): the counter is shared across cases, so bodies + # must never climb into base+9 (END) or they'd collide with it. + idx = MsgTok._body_counter.setdefault(role, 0) + 1 + MsgTok._body_counter[role] = idx + self.body = base + 10 + idx + TOKEN_NAMES[self.body] = f"{role}:{label}" + # message dict as the manager sees it (drives node_match_key). + self.message = {"role": role, "content": label} + + def render(self) -> list[int]: + """[START, body, END] for this message.""" + base = _BANDS[self.role] + return [base, self.body, base + 9] + + +def sys_msg(label: str) -> MsgTok: + return MsgTok("system", label) + + +def usr_msg(label: str) -> MsgTok: + return MsgTok("user", label) + + +def asst_msg(label: str) -> MsgTok: + return MsgTok("assistant", label) + + +def tool_msg(label: str) -> MsgTok: + return MsgTok("tool", label) + + +def render_prompt(msgs: list[MsgTok]) -> list[int]: + """Render a prompt message list, appending the generation-prompt marker.""" + out: list[int] = [] + for m in msgs: + out.extend(m.render()) + out.append(_GEN) + return out + + +def render_response(label: str) -> list[int]: + """Render an assistant response: [body, ]. + + The generation-prompt marker ```` equals the assistant START token, so + `` + render_response(x)`` == the assistant message ``[, body, + ]`` replayed in a later prompt. That identity is what makes a clean + continuation hold across turns. + """ + return [_asst_body(label), _BANDS["assistant"] + 9] + + +def messages(msgs: list[MsgTok]) -> list[dict]: + """The plain message dicts append_turn wants for prompt_messages.""" + return [m.message for m in msgs] + + +def drift(ids: list[int], at: int, sentinel: int = _DRIFT_BAND + 1) -> list[int]: + """Return a copy of ``ids`` with a sentinel spliced at index ``at``. + + The sentinel sits in the drift band (7000+), so a dump shows ```` at + the exact divergence point. Splicing (insert) makes ``len`` grow by one, + which is enough to make the lcp diverge at ``at``. + """ + TOKEN_NAMES[sentinel] = "" + return ids[:at] + [sentinel] + ids[at:] + + +def drift_replace(ids: list[int], at: int, sentinel: int = _DRIFT_BAND + 2) -> list[int]: + """Return a copy of ``ids`` with the token at ``at`` REPLACED by a sentinel. + + Unlike ``drift`` this keeps length constant — used when a test wants the + divergence inside a response span without changing the cumulative length. + """ + TOKEN_NAMES[sentinel] = "" + out = list(ids) + out[at] = sentinel + return out + + +def turn(prompt_ids, response_ids, *, finish_reason="stop", logprobs=None) -> TurnRecord: + return TurnRecord( + prompt_ids=list(prompt_ids), + output_ids=list(response_ids), + finish_reason=finish_reason, + output_log_probs=list(logprobs) if logprobs is not None else [], + ) + + +# A scratch space for the dual-mode printer: each case appends (title, mgr, sid, +# samples) so main() can render after the assertions pass. +_PRINT_LOG: list[tuple[str, object, str, list]] = [] + + +def _record(title: str, mgr, sid: str, samples: list) -> None: + _PRINT_LOG.append((title, mgr, sid, samples)) + + +# Convenience: append a turn with semantic messages, auto-rendering prompt unless +# an explicit prompt_ids is supplied (for drift injection). +def append( + mgr: TrajectoryManager, + sid: str, + prompt_msgs: list[MsgTok], + response_label: str | None, + *, + prompt_ids=None, + response_ids=None, + finish_reason="stop", + logprobs=None, + tools=None, + response_message=None, +): + p = list(prompt_ids) if prompt_ids is not None else render_prompt(prompt_msgs) + if response_ids is not None: + r = list(response_ids) + elif response_label is not None: + r = render_response(response_label) + else: + r = [] + rmsg = response_message + if rmsg is None and response_label is not None: + rmsg = {"role": "assistant", "content": response_label} + lp = logprobs + mgr.append_turn( + sid, + turn=turn(p, r, finish_reason=finish_reason, logprobs=lp), + prompt_messages=messages(prompt_msgs), + tools=tools, + response_message=rmsg, + ) + return p, r + + +def _leaves(mgr, sid): + return [leaf for leaf in mgr._trees[sid].leaves() if not leaf.is_root] + + +def _check_invariants(samples): + for s in samples: + assert len(s.loss_mask) == len(s.rollout_log_probs) == s.response_length, ( + "alignment broken", + len(s.loss_mask), + len(s.rollout_log_probs), + s.response_length, + ) + assert sum(s.loss_mask) > 0, "fully-masked sample emitted" + + +# =========================================================================== +# §2 Group 1 — routing tree layer (append_turn shapes the tree) +# =========================================================================== + + +def test_1_1_single_turn_chain(): + mgr = TrajectoryManager() + sid = "1.1" + s = sys_msg("S") + u = usr_msg("compute") + append(mgr, sid, [s, u], "ok") + chain = _leaves(mgr, sid)[0].path_from_root() + assert [n.role for n in chain] == ["system", "user", "assistant"] + assert chain[-1].turn_index == 1 + _record("1.1 single turn -> linear chain", mgr, sid, []) + print("PASS 1.1") + + +def test_1_2_clean_multiturn_with_tool(): + mgr = TrajectoryManager() + sid = "1.2" + s, u = sys_msg("S"), usr_msg("compute") + a1, t1 = asst_msg("call"), tool_msg("4") + append(mgr, sid, [s, u], "call", finish_reason="tool_calls") + append(mgr, sid, [s, u, a1, t1], "done") + chain = _leaves(mgr, sid)[0].path_from_root() + assert [n.role for n in chain] == ["system", "user", "assistant", "tool", "assistant"] + assert mgr.turn_count(sid) == 2 + _record("1.2 clean 2-turn with tool -> single chain", mgr, sid, []) + print("PASS 1.2") + + +def test_1_3_system_fork(): + mgr = TrajectoryManager() + sid = "1.3" + for sl in ["SA", "SB"]: + append(mgr, sid, [sys_msg(sl), usr_msg("u")], "a") + root = mgr._trees[sid] + assert len(root.children) == 2, "different system -> two subtrees at root" + assert len(_leaves(mgr, sid)) == 2 + _record("1.3 system fork -> two subtrees at root", mgr, sid, []) + print("PASS 1.3") + + +def test_1_4_user_fork_shared_system(): + mgr = TrajectoryManager() + sid = "1.4" + s = sys_msg("S") + for ul in ["A", "B"]: + append(mgr, sid, [s, usr_msg(ul)], ul.lower()) + root = mgr._trees[sid] + assert len(root.children) == 1, "system shared" + assert len(root.children[0].children) == 2, "user level forks" + assert len(_leaves(mgr, sid)) == 2 + _record("1.4 user fork (shared system)", mgr, sid, []) + print("PASS 1.4") + + +def test_1_5_assistant_message_fork(): + """Same (sys,user) prefix, two distinct assistant turns -> assistant fork.""" + mgr = TrajectoryManager() + sid = "1.5" + s, u = sys_msg("S"), usr_msg("u") + append(mgr, sid, [s, u], "a1") + append(mgr, sid, [s, u], "a2") + user_node = mgr._trees[sid].children[0].children[0] + assert len(user_node.children) == 2, "two assistant leaves hang off shared user" + _record("1.5 assistant fork under shared user", mgr, sid, []) + print("PASS 1.5") + + +def test_1_6_tool_fork_shared_assistant(): + """Same first assistant turn, two different tool results -> tool-level fork, + making the first assistant a shared snapshot node with 2 children.""" + mgr = TrajectoryManager() + sid = "1.6" + s, u, a1 = sys_msg("S"), usr_msg("u"), asst_msg("call") + append(mgr, sid, [s, u], "call", finish_reason="tool_calls") + append(mgr, sid, [s, u, a1, tool_msg("x")], "ax") + append(mgr, sid, [s, u, a1, tool_msg("y")], "ay") + asst1 = mgr._trees[sid].children[0].children[0].children[0] + assert asst1.role == "assistant" and asst1.turn_prompt_ids is not None + assert len(asst1.children) == 2, "shared assistant forks at the tool level" + assert len(_leaves(mgr, sid)) == 2 + _record("1.6 tool fork (shared assistant snapshot)", mgr, sid, []) + print("PASS 1.6") + + +def test_1_7_token_only_drift_no_fork(): + """Identical messages, tampered prompt_ids -> NO fork (DFS ignores tokens).""" + mgr = TrajectoryManager() + sid = "1.7" + s, u = sys_msg("S"), usr_msg("u") + p1, _ = append(mgr, sid, [s, u], "a") + tampered = drift(p1, 1) + append(mgr, sid, [s, u], "b", prompt_ids=tampered) + user_node = mgr._trees[sid].children[0].children[0] + assert len(user_node.children) == 2, "two assistant turns share the (sys,user) path" + # but the path above the assistant is single (not forked on tokens) + assert len(mgr._trees[sid].children) == 1 + _record("1.7 token-only drift -> no tree fork", mgr, sid, []) + print("PASS 1.7") + + +def test_1_8_multi_tool_per_turn(): + mgr = TrajectoryManager() + sid = "1.8" + s, u, a1 = sys_msg("S"), usr_msg("u"), asst_msg("call") + ta, tb = tool_msg("A"), tool_msg("B") + append(mgr, sid, [s, u], "call", finish_reason="tool_calls") + append(mgr, sid, [s, u, a1, ta, tb], "done") + chain = _leaves(mgr, sid)[0].path_from_root() + assert [n.role for n in chain] == ["system", "user", "assistant", "tool", "tool", "assistant"] + assert chain[3].messages == [ta.message] + assert chain[4].messages == [tb.message] + _record("1.8 multi-tool turn -> one node per tool", mgr, sid, []) + print("PASS 1.8") + + +def test_1_9_cross_sid_isolation(): + mgr = TrajectoryManager() + s = sys_msg("S") + for sid, ul in [("sid-a", "A"), ("sid-b", "B")]: + append(mgr, sid, [s, usr_msg(ul)], ul.lower()) + assert len(_leaves(mgr, "sid-a")) == 1 + assert len(_leaves(mgr, "sid-b")) == 1 + assert mgr._trees["sid-a"] is not mgr._trees["sid-b"] + print("PASS 1.9") + + +def test_1_10_empty_response(): + mgr = TrajectoryManager() + sid = "1.10" + s, u = sys_msg("S"), usr_msg("u") + append(mgr, sid, [s, u], None, response_ids=[], response_message=None, finish_reason="length") + asst = _leaves(mgr, sid)[0] + assert asst.role == "assistant" + assert asst.turn_response_ids == [] + assert asst.messages == [] + _record("1.10 empty response -> assistant leaf, no message", mgr, sid, []) + print("PASS 1.10") + + +# =========================================================================== +# §2 Group 2 — linearization layer (get_trajectory token routing) +# =========================================================================== + + +def test_2_1_single_turn_linearize(): + mgr = TrajectoryManager() + sid = "2.1" + s, u = sys_msg("S"), usr_msg("u") + p, r = append(mgr, sid, [s, u], "a", logprobs=None) + # attach explicit logprobs so we can check propagation + leaf = _leaves(mgr, sid)[0] + leaf.turn_response_logprobs = [-0.5] * len(r) + samples = mgr.get_trajectory(sid, base_sample=Sample(index=7, prompt="hi"), reward=1.0) + assert len(samples) == 1 + s0 = samples[0] + assert s0.tokens == p + r + assert s0.loss_mask == [1] * len(r) + assert s0.rollout_log_probs == [-0.5] * len(r) + assert s0.response_length == len(r) + assert s0.reward == 1.0 + _check_invariants(samples) + _record("2.1 single-turn linearize", mgr, sid, samples) + print("PASS 2.1") + + +def test_2_2_clean_multiturn_linearize(): + mgr = TrajectoryManager() + sid = "2.2" + s, u, a1, t1 = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("4") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) + p2, r2 = append(mgr, sid, [s, u, a1, t1], "done", logprobs=[-0.4] * 2) + samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + s0 = samples[0] + L = _lcp_len(p1 + r1, p2) + assert L == len(p1) + len(r1) + assert s0.tokens == p1 + r1 + p2[L:] + r2 + assert s0.loss_mask == [1] * len(r1) + [0] * (len(p2) - L) + [1] * len(r2) + assert s0.rollout_log_probs == [-0.5] * len(r1) + [0.0] * (len(p2) - L) + [-0.4] * len(r2) + _check_invariants(samples) + _record("2.2 clean 2-turn linearize", mgr, sid, samples) + print("PASS 2.2") + + +def test_2_3_drift_case_A_forks(): + """Drift inside a PROMPT region -> case A -> fork, no token dropped.""" + mgr = TrajectoryManager() + sid = "2.3" + s, u, a1, t = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("t") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) + p2_honest = render_prompt([s, u, a1, t]) + p2 = drift(p2_honest, len(p1) - 1) # inside p1's prompt region + p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2, logprobs=[-0.4] * 2) + samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + assert len(samples) == 2 + s1, s2 = samples + assert s1.tokens == p1 + r1 + assert s1.loss_mask == [1] * len(r1) + assert s2.tokens == p2 + r2 + assert s2.loss_mask == [1] * len(r2) + assert all(s.reward == 1.0 for s in samples) + _check_invariants(samples) + _record("2.3 drift case A (prompt region) -> fork", mgr, sid, samples) + print("PASS 2.3") + + +def test_2_4_drift_case_B1_short_replaces(): + """Small drift inside the most-recent response span -> replace.""" + mgr = TrajectoryManager() # default threshold 1024 + sid = "2.4" + s, u, a1, t = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("t") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) + p2_honest = render_prompt([s, u, a1, t]) + assert p2_honest[: len(p1) + len(r1)] == p1 + r1 + drift_idx = len(p1) + len(r1) - 1 # last token of r1's echo + p2 = drift_replace(p2_honest, drift_idx) + p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2, logprobs=[-0.4] * 2) + samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + s0 = samples[0] + L = _lcp_len(p1 + r1, p2) + assert L == drift_idx + assert s0.tokens == p2 + r2 + assert s0.loss_mask == [1] * (L - len(p1)) + [0] * (len(p2) - L) + [1] * len(r2) + assert s0.rollout_log_probs == [-0.5] * (L - len(p1)) + [0.0] * (len(p2) - L) + [-0.4] * len(r2) + _check_invariants(samples) + _record("2.4 drift case B1 (small) -> replace", mgr, sid, samples) + print("PASS 2.4") + + +def test_2_5_drift_case_B1_long_forks(): + mgr = TrajectoryManager(fork_merge_max_response_tokens=1) # d>=1 -> fork + sid = "2.5" + s, u, a1, t = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("t") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls") + p2_honest = render_prompt([s, u, a1, t]) + p2 = drift_replace(p2_honest, len(p1) + len(r1) - 1) + p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2) + samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + assert len(samples) == 2 + assert samples[0].tokens == p1 + r1 + assert samples[1].tokens == p2 + r2 + assert all(s.reward == 1.0 for s in samples) + _check_invariants(samples) + _record("2.5 drift case B1 (long) -> fork", mgr, sid, samples) + print("PASS 2.5") + + +def test_2_6_drift_case_B1_threshold_zero_forks(): + mgr = TrajectoryManager(fork_merge_max_response_tokens=0) + sid = "2.6" + s, u, a1, t = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("t") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls") + p2_honest = render_prompt([s, u, a1, t]) + p2 = drift_replace(p2_honest, len(p1) + len(r1) - 1) + append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2) + samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + assert len(samples) == 2 + _check_invariants(samples) + _record("2.6 drift case B1 threshold=0 -> fork", mgr, sid, samples) + print("PASS 2.6") + + +def test_2_7_drift_case_B2_earlier_turn_forks(): + """Drift inside an EARLIER turn's response span -> always fork.""" + mgr = TrajectoryManager() + sid = "2.7" + s, u = sys_msg("S"), usr_msg("u") + a1, t1 = asst_msg("a1"), tool_msg("t1") + a2, t2 = asst_msg("a2"), tool_msg("t2") + p1, r1 = append(mgr, sid, [s, u], "a1", finish_reason="tool_calls", logprobs=[-0.5] * 2) + p2, r2 = append(mgr, sid, [s, u, a1, t1], "a2", finish_reason="tool_calls", logprobs=[-0.4] * 2) + p3_honest = render_prompt([s, u, a1, t1, a2, t2]) + p3 = drift_replace(p3_honest, len(p1) + len(r1) - 1) # inside r1 (earlier span) + p3, r3 = append(mgr, sid, [s, u, a1, t1, a2, t2], "a3", prompt_ids=p3, logprobs=[-0.3] * 2) + samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + assert len(samples) == 2 + s1, s2 = samples + L12 = _lcp_len(p1 + r1, p2) + assert s1.tokens == p1 + r1 + p2[L12:] + r2 + assert s2.tokens == p3 + r3 + assert all(s.reward == 1.0 for s in samples) + _check_invariants(samples) + _record("2.7 drift case B2 (earlier turn) -> fork", mgr, sid, samples) + print("PASS 2.7") + + +def test_2_8_fork_reward_split(): + mgr = TrajectoryManager() + sid = "2.8" + s, u, a1, t = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("t") + p1, _ = append(mgr, sid, [s, u], "call", finish_reason="tool_calls") + p2_honest = render_prompt([s, u, a1, t]) + p2 = drift(p2_honest, len(p1) - 1) # prompt region -> case A fork + append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2) + samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=3.0) + assert len(samples) == 2 + assert all(abs(s.reward - 1.5) < 1e-9 for s in samples) + _record("2.8 fork reward split (3.0 / 2)", mgr, sid, samples) + print("PASS 2.8") + + +def test_2_9_two_leaves_reward_split(): + mgr = TrajectoryManager() + sid = "2.9" + s = sys_msg("S") + for ul in ["A", "B"]: + append(mgr, sid, [s, usr_msg(ul)], ul.lower()) + samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + assert len(samples) == 2 + assert all(s.reward == 1.0 for s in samples) + _check_invariants(samples) + _record("2.9 two leaves reward split (2.0 / 2)", mgr, sid, samples) + print("PASS 2.9") + + +def test_2_10_cross_leaf_dedup(): + """Shared assistant trained on first leaf only; second leaf re-emits it + as loss=0 context.""" + mgr = TrajectoryManager() + sid = "2.10" + s, u, a1 = sys_msg("S"), usr_msg("u"), asst_msg("call") + tx, ty = tool_msg("x"), tool_msg("y") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) + p2, r2 = append(mgr, sid, [s, u, a1, tx], "a2", logprobs=[-0.4] * 2) + p3, r3 = append(mgr, sid, [s, u, a1, ty], "a3", logprobs=[-0.3] * 2) + samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + assert len(samples) == 2 + s_first, s_second = samples + L2 = _lcp_len(p1 + r1, p2) + assert s_first.tokens == p2 + r2 + assert s_first.loss_mask == [1] * len(r1) + [0] * (len(p2) - L2) + [1] * len(r2) + assert s_second.tokens == p3 + r3 + assert s_second.loss_mask == [0] * (len(p3) - len(p1)) + [1] * len(r3) + assert s_second.rollout_log_probs == [0.0] * (len(p3) - len(p1)) + [-0.3] * len(r3) + _check_invariants(samples) + _record("2.10 cross-leaf dedup (shared assistant trained once)", mgr, sid, samples) + print("PASS 2.10") + + +def test_2_11_routing_only_assistant_filtered(): + """cc replays an assistant the manager never recorded -> mounts routing-only, + must be filtered out of the strict-prefix walk (no raise).""" + mgr = TrajectoryManager() + sid = "2.11" + s, u = sys_msg("S"), usr_msg("u") + a1, t1 = asst_msg("a1"), tool_msg("t1") + a2 = asst_msg("a2") + append(mgr, sid, [s, u], "a1", finish_reason="tool_calls") + append(mgr, sid, [s, u, a1, t1], "a2", finish_reason="tool_calls") + foreign = asst_msg("foreign") + t2 = tool_msg("t2") + append(mgr, sid, [s, u, a1, t1, a2, foreign, t2], "a3") + leaves = _leaves(mgr, sid) + assert len(leaves) == 1 + chain = leaves[0].path_from_root() + routing = [n for n in chain if n.role == "assistant" and n.turn_prompt_ids is None] + assert len(routing) == 1 and routing[0].messages[0]["content"] == "foreign" + samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt="")) + assert len(samples) == 1 + _record("2.11 routing-only assistant filtered (no raise)", mgr, sid, samples) + print("PASS 2.11") + + +def test_2_12_drop_clears_sid(): + mgr = TrajectoryManager() + sid = "2.12" + s, u = sys_msg("S"), usr_msg("u") + append(mgr, sid, [s, u], "a") + assert mgr.has_session(sid) + mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt="")) + assert not mgr.has_session(sid) + assert mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt="")) == [] + print("PASS 2.12") + + +# =========================================================================== +# §2 Group 3 — combined / stress (both layers interacting) +# =========================================================================== + + +def test_3_1_rewrite_merge_absorbs_short(): + mgr = TrajectoryManager() + sid = "3.1" + s, u = sys_msg("S"), usr_msg("u") + a1_rw = asst_msg("ok ") # cc-rewritten (different message identity) + t1 = tool_msg("t") + p1, r1 = append(mgr, sid, [s, u], "ok", finish_reason="tool_calls", logprobs=[-0.5] * 2) + p2, r2 = append(mgr, sid, [s, u, a1_rw, t1], "done", logprobs=[-0.4] * 2) + leaves = _leaves(mgr, sid) + assert len(leaves) == 1, "short rewrite absorbed, not forked" + chain = leaves[0].path_from_root() + merged = chain[2] + assert merged.turn_prompt_ids is None and merged.turn_index is None + assert merged.messages == [a1_rw.message] + assert merged.metadata["merged_rewrite"]["abandoned_turn_index"] == 1 + samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + assert samples[0].tokens == p2 + r2 + assert samples[0].loss_mask == [1] * len(r2) + _check_invariants(samples) + _record("3.1 rewrite-merge absorbs short assistant", mgr, sid, samples) + print("PASS 3.1") + + +def test_3_2_rewrite_merge_long_forks(): + mgr = TrajectoryManager(fork_merge_max_response_tokens=1) # r1 len 2 >= 1 + sid = "3.2" + s, u = sys_msg("S"), usr_msg("u") + a1_rw, t1 = asst_msg("ok2 "), tool_msg("t") + append(mgr, sid, [s, u], "ok", finish_reason="tool_calls") + append(mgr, sid, [s, u, a1_rw, t1], "done") + assert len(_leaves(mgr, sid)) == 2, "long rewrite forks" + samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + assert len(samples) == 2 + _check_invariants(samples) + _record("3.2 rewrite-merge long -> fork", mgr, sid, samples) + print("PASS 3.2") + + +def test_3_3_rewrite_merge_threshold_zero_forks(): + mgr = TrajectoryManager(fork_merge_max_response_tokens=0) + sid = "3.3" + s, u = sys_msg("S"), usr_msg("u") + a1_rw, t1 = asst_msg("ok3 "), tool_msg("t") + append(mgr, sid, [s, u], "ok", finish_reason="tool_calls") + append(mgr, sid, [s, u, a1_rw, t1], "done") + assert len(_leaves(mgr, sid)) == 2 + samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + _check_invariants(samples) + _record("3.3 rewrite-merge threshold=0 -> fork", mgr, sid, samples) + print("PASS 3.3") + + +def test_3_4_rewrite_merge_ambiguous_forks(): + mgr = TrajectoryManager() + sid = "3.4" + s, u = sys_msg("S"), usr_msg("u") + # two short assistant leaves under shared (sys,user) + append(mgr, sid, [s, u], "a") + append(mgr, sid, [s, u], "b") + a_c, t1 = asst_msg("c"), tool_msg("t") + append(mgr, sid, [s, u, a_c, t1], "d") + assert len(_leaves(mgr, sid)) == 3, "ambiguous candidates fork" + _record("3.4 rewrite-merge ambiguous -> fork", mgr, sid, []) + print("PASS 3.4") + + +def test_3_5_rewrite_merge_match_key_updated(): + """After merge, a later turn replaying the rewritten message must descend + through the merged node (match_key updated), not fork again.""" + mgr = TrajectoryManager() + sid = "3.5" + s, u = sys_msg("S"), usr_msg("u") + a1_rw = asst_msg("ok5 ") + t1, a2, t2 = tool_msg("t1"), asst_msg("second"), tool_msg("t2") + append(mgr, sid, [s, u], "ok", finish_reason="tool_calls") + p2, r2 = append(mgr, sid, [s, u, a1_rw, t1], "second", finish_reason="tool_calls", logprobs=[-0.4] * 2) + p3, r3 = append(mgr, sid, [s, u, a1_rw, t1, a2, t2], "third", logprobs=[-0.3] * 2) + leaves = _leaves(mgr, sid) + assert len(leaves) == 1, "match_key updated -> no spurious fork" + samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt="")) + assert len(samples) == 1 + L = _lcp_len(p2 + r2, p3) + assert samples[0].tokens == p2 + r2 + p3[L:] + r3 + _check_invariants(samples) + _record("3.5 rewrite-merge match_key updated", mgr, sid, samples) + print("PASS 3.5") + + +def test_3_6_tree_fork_plus_token_drift(): + """A tree fork (two leaves) where ONE leaf also drift-forks internally, + yielding 3 Samples total. Combines layer-1 (message fork) with layer-2 + (token drift fork).""" + mgr = TrajectoryManager() + sid = "3.6" + s, u = sys_msg("S"), usr_msg("u") + a1, tx, ty = asst_msg("call"), tool_msg("x"), tool_msg("y") + ax2 = asst_msg("ax2") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) + # Leaf X: clean continuation, then a third turn with a case-A prompt drift. + p2, r2 = append(mgr, sid, [s, u, a1, tx], "ax2", finish_reason="tool_calls", logprobs=[-0.4] * 2) + txx = tool_msg("xx") + p3_honest = render_prompt([s, u, a1, tx, ax2, txx]) + p3 = drift(p3_honest, len(p1) - 1) # case A drift -> fork inside leaf X + append(mgr, sid, [s, u, a1, tx, ax2, txx], "ax3", prompt_ids=p3, logprobs=[-0.2] * 2) + # Leaf Y: a separate tool result off the shared assistant. + append(mgr, sid, [s, u, a1, ty], "ay2", logprobs=[-0.1] * 2) + samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=3.0) + # Leaf X -> 2 samples (drift fork), Leaf Y -> 1 sample. Total 3. + assert len(samples) == 3, [s.tokens for s in samples] + assert all(abs(s.reward - 1.0) < 1e-9 for s in samples) + _check_invariants(samples) + _record("3.6 tree fork + token drift -> 3 samples", mgr, sid, samples) + print("PASS 3.6") + + +def test_3_7_deep_multi_leaf_dedup(): + """Three leaves sharing a 2-level assistant prefix; the shared turns are + trained exactly once across all leaves.""" + mgr = TrajectoryManager() + sid = "3.7" + s, u = sys_msg("S"), usr_msg("u") + a1, t1 = asst_msg("a1"), tool_msg("t1") + a2 = asst_msg("a2") + p1, r1 = append(mgr, sid, [s, u], "a1", finish_reason="tool_calls", logprobs=[-0.5] * 2) + p2, r2 = append(mgr, sid, [s, u, a1, t1], "a2", finish_reason="tool_calls", logprobs=[-0.4] * 2) + # three different tool results off a2 -> three leaves sharing a1+a2 + for lbl in ["p", "q", "r"]: + append(mgr, sid, [s, u, a1, t1, a2, tool_msg(lbl)], f"end-{lbl}", logprobs=[-0.3] * 2) + samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=3.0) + assert len(samples) == 3 + # Count trained tokens for r1 and r2 across all samples: each shared turn + # trained exactly once (first leaf), others loss=0. + # First leaf trains r1+r2+end; others train only their own end. + trained_first = sum(samples[0].loss_mask) + trained_rest = [sum(x.loss_mask) for x in samples[1:]] + assert trained_first > max(trained_rest), (trained_first, trained_rest) + _check_invariants(samples) + _record("3.7 deep multi-leaf dedup (3 leaves, shared trained once)", mgr, sid, samples) + print("PASS 3.7") + + +def test_3_8_long_mixed_session(): + """A ~7-turn session combining clean continuation, a mid-session B1 replace, + and a final case-A fork — verifying the mechanisms chain without interfering.""" + mgr = TrajectoryManager() + sid = "3.8" + s, u = sys_msg("S"), usr_msg("u") + a = [asst_msg(f"a{i}") for i in range(6)] + t = [tool_msg(f"t{i}") for i in range(6)] + lp = [-0.5, -0.5] + # turn 1 + p1, r1 = append(mgr, sid, [s, u], "a0", finish_reason="tool_calls", logprobs=lp) + # turns 2..4 clean + prefix = [s, u, a[0], t[0]] + append(mgr, sid, prefix, "a1", finish_reason="tool_calls", logprobs=lp) + prefix = prefix + [a[1], t[1]] + append(mgr, sid, prefix, "a2", finish_reason="tool_calls", logprobs=lp) + prefix = prefix + [a[2], t[2]] + # turn 5: B1 small replace — drift the last token of the previous response. + p5_honest = render_prompt(prefix) + p5 = drift_replace(p5_honest, len(p5_honest) - 2) # near tail, inside last resp echo region + append(mgr, sid, prefix, "a3", prompt_ids=p5, finish_reason="tool_calls", logprobs=lp) + prefix = prefix + [a[3], t[3]] + # turn 6: clean + append(mgr, sid, prefix, "a4", finish_reason="tool_calls", logprobs=lp) + prefix = prefix + [a[4], t[4]] + # turn 7: case-A fork (drift in early prompt region) + p7_honest = render_prompt(prefix) + p7 = drift(p7_honest, len(p1) - 1) + append(mgr, sid, prefix, "a5", prompt_ids=p7, logprobs=lp) + samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=4.0) + # The final case-A fork splits the single leaf chain into >=2 segments. + assert len(samples) >= 2, len(samples) + assert abs(sum(s.reward for s in samples) - 4.0) < 1e-9 + _check_invariants(samples) + _record(f"3.8 long mixed session -> {len(samples)} samples", mgr, sid, samples) + print("PASS 3.8") + + +# =========================================================================== +# §3 Dual-mode printer +# =========================================================================== + + +def _print_sample(idx: int, s: Sample) -> None: + toks = s.tokens + resp_start = len(toks) - s.response_length + # build aligned token/loss rows over the response region (the trained part); + # the leading prompt prefix has no loss_mask entry. + names = [name_of(t) for t in toks] + loss = ["-"] * resp_start + [str(x) for x in s.loss_mask] + widths = [max(len(names[i]), len(loss[i])) for i in range(len(toks))] + tok_row = " ".join(names[i].ljust(widths[i]) for i in range(len(toks))) + loss_row = " ".join(loss[i].ljust(widths[i]) for i in range(len(toks))) + print(f" Sample#{idx} reward={s.reward:.3f} resp_len={s.response_length}") + print(f" tok : {tok_row}") + print(f" loss: {loss_row}") + + +def _print_case(title: str, mgr, sid: str, samples: list) -> None: + print(f"\n=== CASE {title} ===") + txt = dump_tree_txt(mgr, sid) if mgr.has_session(sid) else "" + print("[tree]") + for line in txt.splitlines(): + print(" " + line) + if samples: + print(f"[samples] {len(samples)}") + for i, s in enumerate(samples): + _print_sample(i, s) + + +# =========================================================================== +# main +# =========================================================================== + + +_CASES = [ + test_1_1_single_turn_chain, + test_1_2_clean_multiturn_with_tool, + test_1_3_system_fork, + test_1_4_user_fork_shared_system, + test_1_5_assistant_message_fork, + test_1_6_tool_fork_shared_assistant, + test_1_7_token_only_drift_no_fork, + test_1_8_multi_tool_per_turn, + test_1_9_cross_sid_isolation, + test_1_10_empty_response, + test_2_1_single_turn_linearize, + test_2_2_clean_multiturn_linearize, + test_2_3_drift_case_A_forks, + test_2_4_drift_case_B1_short_replaces, + test_2_5_drift_case_B1_long_forks, + test_2_6_drift_case_B1_threshold_zero_forks, + test_2_7_drift_case_B2_earlier_turn_forks, + test_2_8_fork_reward_split, + test_2_9_two_leaves_reward_split, + test_2_10_cross_leaf_dedup, + test_2_11_routing_only_assistant_filtered, + test_2_12_drop_clears_sid, + test_3_1_rewrite_merge_absorbs_short, + test_3_2_rewrite_merge_long_forks, + test_3_3_rewrite_merge_threshold_zero_forks, + test_3_4_rewrite_merge_ambiguous_forks, + test_3_5_rewrite_merge_match_key_updated, + test_3_6_tree_fork_plus_token_drift, + test_3_7_deep_multi_leaf_dedup, + test_3_8_long_mixed_session, +] + + +def main() -> None: + for case in _CASES: + case() + # Replay the captured tree / sample snapshots as human-readable dumps. + print("\n" + "=" * 70) + print("HUMAN-READABLE DUMPS") + print("=" * 70) + for title, mgr, sid, samples in _PRINT_LOG: + _print_case(title, mgr, sid, samples) + print(f"\nALL E2E CASES PASSED ({len(_CASES)} cases)") + + +if __name__ == "__main__": + main() From cbee0decb3f3b3e0220e0af3503304d2b2437b73 Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Mon, 8 Jun 2026 14:12:01 +0000 Subject: [PATCH 15/28] test(agent): dump raw append_turn inputs in e2e readable output Each case now prints [raw turns] (the source prompt_ids/response_ids decoded to names, finish_reason, logprobs presence) before [tree] and [samples], so the full data flow source->tree->samples is visible. --- .../test_agent/test_trajectory_manager_e2e.py | 35 +++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/tests/test_agent/test_trajectory_manager_e2e.py b/tests/test_agent/test_trajectory_manager_e2e.py index 258bc5d4d0..a1a5ed841e 100644 --- a/tests/test_agent/test_trajectory_manager_e2e.py +++ b/tests/test_agent/test_trajectory_manager_e2e.py @@ -200,6 +200,11 @@ def turn(prompt_ids, response_ids, *, finish_reason="stop", logprobs=None) -> Tu # samples) so main() can render after the assertions pass. _PRINT_LOG: list[tuple[str, object, str, list]] = [] +# Raw append_turn inputs, keyed by sid, captured at call time so the printer can +# show the SOURCE data (prompt_ids / response_ids / finish / logprobs) that fed +# the tree — before any tree-building or linearization happened. +_TURN_LOG: dict[str, list[dict]] = {} + def _record(title: str, mgr, sid: str, samples: list) -> None: _PRINT_LOG.append((title, mgr, sid, samples)) @@ -231,6 +236,17 @@ def append( if rmsg is None and response_label is not None: rmsg = {"role": "assistant", "content": response_label} lp = logprobs + # Capture the raw turn inputs for the human-readable dump before the manager + # consumes them. + _TURN_LOG.setdefault(sid, []).append( + { + "prompt_msgs": [f"{m.role}:{m.label}" for m in prompt_msgs], + "prompt_ids": p, + "response_ids": r, + "finish": finish_reason, + "has_lp": lp is not None, + } + ) mgr.append_turn( sid, turn=turn(p, r, finish_reason=finish_reason, logprobs=lp), @@ -841,8 +857,27 @@ def _print_sample(idx: int, s: Sample) -> None: print(f" loss: {loss_row}") +def _print_raw_turns(sid: str) -> None: + """Print the raw append_turn inputs (the SOURCE data) for a sid. + + Shows, per turn, the prompt message labels and the actual prompt_ids / + response_ids decoded to readable names, plus finish_reason and whether + logprobs were attached. This is what fed the tree, before any building or + linearization. + """ + turns = _TURN_LOG.get(sid, []) + print(f"[raw turns] {len(turns)}") + for k, t in enumerate(turns, start=1): + msgs = " , ".join(t["prompt_msgs"]) + print(f" turn#{k} finish={t['finish']} has_logprobs={t['has_lp']}") + print(f" msgs : {msgs}") + print(f" prompt : {render_ids(t['prompt_ids'])}") + print(f" output : {render_ids(t['response_ids']) or ''}") + + def _print_case(title: str, mgr, sid: str, samples: list) -> None: print(f"\n=== CASE {title} ===") + _print_raw_turns(sid) txt = dump_tree_txt(mgr, sid) if mgr.has_session(sid) else "" print("[tree]") for line in txt.splitlines(): From a11c9b4193fcd4a05fb1415caf75b7e794bef7f5 Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Mon, 8 Jun 2026 14:20:00 +0000 Subject: [PATCH 16/28] test(agent): 1.7 now shows token drift's effect on the linearized sample - 1.7 calls get_trajectory and asserts the token lands in leaf 2's stripped prompt region (loss=0), proving token drift never corrupts a trained response while still being carried in the sample tokens. - get_traj wrapper snapshots the tree before get_trajectory drains the sid, so every case (incl. group 2/3) shows [tree] and [samples] together instead of . --- .../test_agent/test_trajectory_manager_e2e.py | 88 ++++++++++++++----- 1 file changed, 64 insertions(+), 24 deletions(-) diff --git a/tests/test_agent/test_trajectory_manager_e2e.py b/tests/test_agent/test_trajectory_manager_e2e.py index a1a5ed841e..56f0a10c58 100644 --- a/tests/test_agent/test_trajectory_manager_e2e.py +++ b/tests/test_agent/test_trajectory_manager_e2e.py @@ -261,6 +261,24 @@ def _leaves(mgr, sid): return [leaf for leaf in mgr._trees[sid].leaves() if not leaf.is_root] +# Tree text snapshot captured the instant before get_trajectory drains the sid, +# so the human-readable dump can show [tree] AND [samples] side by side even +# though get_trajectory consumes the session. +_TREE_SNAP: dict[str, str] = {} + + +def get_traj(mgr, sid, *args, **kwargs): + """get_trajectory wrapper that snapshots the tree before draining. + + Linearization (get_trajectory) pops the sid, so a later dump would only see + ````. Capturing the tree text here keeps the routing tree visible + next to the Samples it produced. + """ + if mgr.has_session(sid): + _TREE_SNAP[sid] = dump_tree_txt(mgr, sid) + return mgr.get_trajectory(sid, *args, **kwargs) + + def _check_invariants(samples): for s in samples: assert len(s.loss_mask) == len(s.rollout_log_probs) == s.response_length, ( @@ -361,18 +379,35 @@ def test_1_6_tool_fork_shared_assistant(): def test_1_7_token_only_drift_no_fork(): - """Identical messages, tampered prompt_ids -> NO fork (DFS ignores tokens).""" + """Identical messages, tampered prompt_ids -> NO tree fork (DFS ignores + tokens), but the drift DOES surface in the linearized sample: it lands in + leaf 2's prompt region (stripped / loss=0), proving token drift cannot + corrupt a trained response yet is still carried in the sample tokens.""" mgr = TrajectoryManager() sid = "1.7" s, u = sys_msg("S"), usr_msg("u") - p1, _ = append(mgr, sid, [s, u], "a") - tampered = drift(p1, 1) + pa, ra = append(mgr, sid, [s, u], "a") + tampered = drift(pa, 1) # spliced into the prompt at index 1 + rb = render_response("b") append(mgr, sid, [s, u], "b", prompt_ids=tampered) + # Tree: (sys,user) shared, two assistant turns hang off it -> two leaves; the + # path above the assistant is single (NOT forked on tokens). user_node = mgr._trees[sid].children[0].children[0] assert len(user_node.children) == 2, "two assistant turns share the (sys,user) path" - # but the path above the assistant is single (not forked on tokens) assert len(mgr._trees[sid].children) == 1 - _record("1.7 token-only drift -> no tree fork", mgr, sid, []) + + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + assert len(samples) == 2 + s_a, s_b = samples + # Leaf 1: clean. Leaf 2: tokens carry the token, but it sits in the + # stripped prompt region — loss_mask covers only the response (rb). + assert s_a.tokens == pa + ra + assert s_b.tokens == tampered + rb + assert s_b.loss_mask == [1] * len(rb), "drift confined to prompt; response fully trained" + assert (_DRIFT_BAND + 1) in s_b.tokens, "drift token is still carried in the sample" + assert (_DRIFT_BAND + 1) not in s_b.tokens[len(tampered) :], "drift not in the response region" + _check_invariants(samples) + _record("1.7 token-only drift -> no tree fork, drift lands in stripped prompt", mgr, sid, samples) print("PASS 1.7") @@ -428,7 +463,7 @@ def test_2_1_single_turn_linearize(): # attach explicit logprobs so we can check propagation leaf = _leaves(mgr, sid)[0] leaf.turn_response_logprobs = [-0.5] * len(r) - samples = mgr.get_trajectory(sid, base_sample=Sample(index=7, prompt="hi"), reward=1.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=7, prompt="hi"), reward=1.0) assert len(samples) == 1 s0 = samples[0] assert s0.tokens == p + r @@ -447,7 +482,7 @@ def test_2_2_clean_multiturn_linearize(): s, u, a1, t1 = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("4") p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) p2, r2 = append(mgr, sid, [s, u, a1, t1], "done", logprobs=[-0.4] * 2) - samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 1 s0 = samples[0] L = _lcp_len(p1 + r1, p2) @@ -469,7 +504,7 @@ def test_2_3_drift_case_A_forks(): p2_honest = render_prompt([s, u, a1, t]) p2 = drift(p2_honest, len(p1) - 1) # inside p1's prompt region p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2, logprobs=[-0.4] * 2) - samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) assert len(samples) == 2 s1, s2 = samples assert s1.tokens == p1 + r1 @@ -493,7 +528,7 @@ def test_2_4_drift_case_B1_short_replaces(): drift_idx = len(p1) + len(r1) - 1 # last token of r1's echo p2 = drift_replace(p2_honest, drift_idx) p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2, logprobs=[-0.4] * 2) - samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 1 s0 = samples[0] L = _lcp_len(p1 + r1, p2) @@ -514,7 +549,7 @@ def test_2_5_drift_case_B1_long_forks(): p2_honest = render_prompt([s, u, a1, t]) p2 = drift_replace(p2_honest, len(p1) + len(r1) - 1) p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2) - samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) assert len(samples) == 2 assert samples[0].tokens == p1 + r1 assert samples[1].tokens == p2 + r2 @@ -532,7 +567,7 @@ def test_2_6_drift_case_B1_threshold_zero_forks(): p2_honest = render_prompt([s, u, a1, t]) p2 = drift_replace(p2_honest, len(p1) + len(r1) - 1) append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2) - samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) assert len(samples) == 2 _check_invariants(samples) _record("2.6 drift case B1 threshold=0 -> fork", mgr, sid, samples) @@ -551,7 +586,7 @@ def test_2_7_drift_case_B2_earlier_turn_forks(): p3_honest = render_prompt([s, u, a1, t1, a2, t2]) p3 = drift_replace(p3_honest, len(p1) + len(r1) - 1) # inside r1 (earlier span) p3, r3 = append(mgr, sid, [s, u, a1, t1, a2, t2], "a3", prompt_ids=p3, logprobs=[-0.3] * 2) - samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) assert len(samples) == 2 s1, s2 = samples L12 = _lcp_len(p1 + r1, p2) @@ -571,7 +606,7 @@ def test_2_8_fork_reward_split(): p2_honest = render_prompt([s, u, a1, t]) p2 = drift(p2_honest, len(p1) - 1) # prompt region -> case A fork append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2) - samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=3.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=3.0) assert len(samples) == 2 assert all(abs(s.reward - 1.5) < 1e-9 for s in samples) _record("2.8 fork reward split (3.0 / 2)", mgr, sid, samples) @@ -584,7 +619,7 @@ def test_2_9_two_leaves_reward_split(): s = sys_msg("S") for ul in ["A", "B"]: append(mgr, sid, [s, usr_msg(ul)], ul.lower()) - samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) assert len(samples) == 2 assert all(s.reward == 1.0 for s in samples) _check_invariants(samples) @@ -602,7 +637,7 @@ def test_2_10_cross_leaf_dedup(): p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) p2, r2 = append(mgr, sid, [s, u, a1, tx], "a2", logprobs=[-0.4] * 2) p3, r3 = append(mgr, sid, [s, u, a1, ty], "a3", logprobs=[-0.3] * 2) - samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) assert len(samples) == 2 s_first, s_second = samples L2 = _lcp_len(p1 + r1, p2) @@ -634,7 +669,7 @@ def test_2_11_routing_only_assistant_filtered(): chain = leaves[0].path_from_root() routing = [n for n in chain if n.role == "assistant" and n.turn_prompt_ids is None] assert len(routing) == 1 and routing[0].messages[0]["content"] == "foreign" - samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt="")) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt="")) assert len(samples) == 1 _record("2.11 routing-only assistant filtered (no raise)", mgr, sid, samples) print("PASS 2.11") @@ -672,7 +707,7 @@ def test_3_1_rewrite_merge_absorbs_short(): assert merged.turn_prompt_ids is None and merged.turn_index is None assert merged.messages == [a1_rw.message] assert merged.metadata["merged_rewrite"]["abandoned_turn_index"] == 1 - samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 1 assert samples[0].tokens == p2 + r2 assert samples[0].loss_mask == [1] * len(r2) @@ -689,7 +724,7 @@ def test_3_2_rewrite_merge_long_forks(): append(mgr, sid, [s, u], "ok", finish_reason="tool_calls") append(mgr, sid, [s, u, a1_rw, t1], "done") assert len(_leaves(mgr, sid)) == 2, "long rewrite forks" - samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) assert len(samples) == 2 _check_invariants(samples) _record("3.2 rewrite-merge long -> fork", mgr, sid, samples) @@ -704,7 +739,7 @@ def test_3_3_rewrite_merge_threshold_zero_forks(): append(mgr, sid, [s, u], "ok", finish_reason="tool_calls") append(mgr, sid, [s, u, a1_rw, t1], "done") assert len(_leaves(mgr, sid)) == 2 - samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) _check_invariants(samples) _record("3.3 rewrite-merge threshold=0 -> fork", mgr, sid, samples) print("PASS 3.3") @@ -737,7 +772,7 @@ def test_3_5_rewrite_merge_match_key_updated(): p3, r3 = append(mgr, sid, [s, u, a1_rw, t1, a2, t2], "third", logprobs=[-0.3] * 2) leaves = _leaves(mgr, sid) assert len(leaves) == 1, "match_key updated -> no spurious fork" - samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt="")) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt="")) assert len(samples) == 1 L = _lcp_len(p2 + r2, p3) assert samples[0].tokens == p2 + r2 + p3[L:] + r3 @@ -764,7 +799,7 @@ def test_3_6_tree_fork_plus_token_drift(): append(mgr, sid, [s, u, a1, tx, ax2, txx], "ax3", prompt_ids=p3, logprobs=[-0.2] * 2) # Leaf Y: a separate tool result off the shared assistant. append(mgr, sid, [s, u, a1, ty], "ay2", logprobs=[-0.1] * 2) - samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=3.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=3.0) # Leaf X -> 2 samples (drift fork), Leaf Y -> 1 sample. Total 3. assert len(samples) == 3, [s.tokens for s in samples] assert all(abs(s.reward - 1.0) < 1e-9 for s in samples) @@ -786,7 +821,7 @@ def test_3_7_deep_multi_leaf_dedup(): # three different tool results off a2 -> three leaves sharing a1+a2 for lbl in ["p", "q", "r"]: append(mgr, sid, [s, u, a1, t1, a2, tool_msg(lbl)], f"end-{lbl}", logprobs=[-0.3] * 2) - samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=3.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=3.0) assert len(samples) == 3 # Count trained tokens for r1 and r2 across all samples: each shared turn # trained exactly once (first leaf), others loss=0. @@ -828,7 +863,7 @@ def test_3_8_long_mixed_session(): p7_honest = render_prompt(prefix) p7 = drift(p7_honest, len(p1) - 1) append(mgr, sid, prefix, "a5", prompt_ids=p7, logprobs=lp) - samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=4.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=4.0) # The final case-A fork splits the single leaf chain into >=2 segments. assert len(samples) >= 2, len(samples) assert abs(sum(s.reward for s in samples) - 4.0) < 1e-9 @@ -878,7 +913,12 @@ def _print_raw_turns(sid: str) -> None: def _print_case(title: str, mgr, sid: str, samples: list) -> None: print(f"\n=== CASE {title} ===") _print_raw_turns(sid) - txt = dump_tree_txt(mgr, sid) if mgr.has_session(sid) else "" + if mgr.has_session(sid): + txt = dump_tree_txt(mgr, sid) + else: + # Session already drained by get_trajectory; fall back to the snapshot + # captured by get_traj just before draining. + txt = _TREE_SNAP.get(sid, "") print("[tree]") for line in txt.splitlines(): print(" " + line) From f1b1792a65765c444dcc277f085faa8fbd126be7 Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Mon, 8 Jun 2026 14:25:53 +0000 Subject: [PATCH 17/28] test(agent): every case prints [samples] + mask info - All group-1 cases (1.1-1.6, 1.8, 1.9, 1.10) and 3.4 now call get_trajectory and record their samples, so [samples] with token/loss alignment is shown for every case (1.10 empty-response shows 0 samples, 1.9 records both sids). - Printer always emits the [samples] header (incl. 0). - _asst_body: counter-based label->token assignment (was a hash) so distinct labels never collide and mislabel dump tokens. --- .../test_agent/test_trajectory_manager_e2e.py | 83 +++++++++++++++---- 1 file changed, 65 insertions(+), 18 deletions(-) diff --git a/tests/test_agent/test_trajectory_manager_e2e.py b/tests/test_agent/test_trajectory_manager_e2e.py index 56f0a10c58..51fb936e63 100644 --- a/tests/test_agent/test_trajectory_manager_e2e.py +++ b/tests/test_agent/test_trajectory_manager_e2e.py @@ -67,6 +67,9 @@ def name_of(tok: int) -> str: return TOKEN_NAMES.get(tok, str(tok)) +_ASST_BODY: dict[str, int] = {} + + def _asst_body(label: str) -> int: """Stable assistant body token for a response/message label. @@ -74,11 +77,16 @@ def _asst_body(label: str) -> int: tokens the model generated for it, otherwise a clean continuation can never hold (the cumulative prompt+response would not prefix the next prompt). So both ``render_response`` and an assistant ``MsgTok`` derive their body token - from this one function, keyed on the label. + from this one function, keyed on the label. Bodies are assigned by a stable + per-label counter (NOT a hash) so distinct labels never collide on one id — + a collision would mislabel tokens in the dump and could spuriously match + across turns. """ - body = _BANDS["assistant"] + 100 + (sum(ord(c) for c in label) % 800) - TOKEN_NAMES[body] = f"r:{label}" - return body + if label not in _ASST_BODY: + body = _BANDS["assistant"] + 100 + len(_ASST_BODY) + _ASST_BODY[label] = body + TOKEN_NAMES[body] = f"r:{label}" + return _ASST_BODY[label] def render_ids(ids: list[int]) -> str: @@ -304,7 +312,10 @@ def test_1_1_single_turn_chain(): chain = _leaves(mgr, sid)[0].path_from_root() assert [n.role for n in chain] == ["system", "user", "assistant"] assert chain[-1].turn_index == 1 - _record("1.1 single turn -> linear chain", mgr, sid, []) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + _check_invariants(samples) + _record("1.1 single turn -> linear chain", mgr, sid, samples) print("PASS 1.1") @@ -318,7 +329,10 @@ def test_1_2_clean_multiturn_with_tool(): chain = _leaves(mgr, sid)[0].path_from_root() assert [n.role for n in chain] == ["system", "user", "assistant", "tool", "assistant"] assert mgr.turn_count(sid) == 2 - _record("1.2 clean 2-turn with tool -> single chain", mgr, sid, []) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + _check_invariants(samples) + _record("1.2 clean 2-turn with tool -> single chain", mgr, sid, samples) print("PASS 1.2") @@ -330,7 +344,10 @@ def test_1_3_system_fork(): root = mgr._trees[sid] assert len(root.children) == 2, "different system -> two subtrees at root" assert len(_leaves(mgr, sid)) == 2 - _record("1.3 system fork -> two subtrees at root", mgr, sid, []) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + assert len(samples) == 2 + _check_invariants(samples) + _record("1.3 system fork -> two subtrees at root", mgr, sid, samples) print("PASS 1.3") @@ -344,7 +361,10 @@ def test_1_4_user_fork_shared_system(): assert len(root.children) == 1, "system shared" assert len(root.children[0].children) == 2, "user level forks" assert len(_leaves(mgr, sid)) == 2 - _record("1.4 user fork (shared system)", mgr, sid, []) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + assert len(samples) == 2 + _check_invariants(samples) + _record("1.4 user fork (shared system)", mgr, sid, samples) print("PASS 1.4") @@ -357,7 +377,10 @@ def test_1_5_assistant_message_fork(): append(mgr, sid, [s, u], "a2") user_node = mgr._trees[sid].children[0].children[0] assert len(user_node.children) == 2, "two assistant leaves hang off shared user" - _record("1.5 assistant fork under shared user", mgr, sid, []) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + assert len(samples) == 2 + _check_invariants(samples) + _record("1.5 assistant fork under shared user", mgr, sid, samples) print("PASS 1.5") @@ -374,7 +397,12 @@ def test_1_6_tool_fork_shared_assistant(): assert asst1.role == "assistant" and asst1.turn_prompt_ids is not None assert len(asst1.children) == 2, "shared assistant forks at the tool level" assert len(_leaves(mgr, sid)) == 2 - _record("1.6 tool fork (shared assistant snapshot)", mgr, sid, []) + # The shared assistant (turn 1) is trained on the first leaf only; the second + # re-emits it as loss=0 context (cross-leaf dedup). + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + assert len(samples) == 2 + _check_invariants(samples) + _record("1.6 tool fork (shared assistant snapshot)", mgr, sid, samples) print("PASS 1.6") @@ -422,7 +450,10 @@ def test_1_8_multi_tool_per_turn(): assert [n.role for n in chain] == ["system", "user", "assistant", "tool", "tool", "assistant"] assert chain[3].messages == [ta.message] assert chain[4].messages == [tb.message] - _record("1.8 multi-tool turn -> one node per tool", mgr, sid, []) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + _check_invariants(samples) + _record("1.8 multi-tool turn -> one node per tool", mgr, sid, samples) print("PASS 1.8") @@ -434,6 +465,13 @@ def test_1_9_cross_sid_isolation(): assert len(_leaves(mgr, "sid-a")) == 1 assert len(_leaves(mgr, "sid-b")) == 1 assert mgr._trees["sid-a"] is not mgr._trees["sid-b"] + sa = get_traj(mgr, "sid-a", base_sample=Sample(index=0, prompt=""), reward=1.0) + sb = get_traj(mgr, "sid-b", base_sample=Sample(index=1, prompt=""), reward=1.0) + assert len(sa) == 1 and len(sb) == 1 + _check_invariants(sa) + _check_invariants(sb) + _record("1.9 cross-sid isolation (sid-a)", mgr, "sid-a", sa) + _record("1.9 cross-sid isolation (sid-b)", mgr, "sid-b", sb) print("PASS 1.9") @@ -446,7 +484,11 @@ def test_1_10_empty_response(): assert asst.role == "assistant" assert asst.turn_response_ids == [] assert asst.messages == [] - _record("1.10 empty response -> assistant leaf, no message", mgr, sid, []) + # Empty response -> the only turn has no trainable token, so its segment is + # dropped at linearization (no fully-masked sample). Zero samples is correct. + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 0 + _record("1.10 empty response -> assistant leaf, no message (0 samples)", mgr, sid, samples) print("PASS 1.10") @@ -681,9 +723,12 @@ def test_2_12_drop_clears_sid(): s, u = sys_msg("S"), usr_msg("u") append(mgr, sid, [s, u], "a") assert mgr.has_session(sid) - mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt="")) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 assert not mgr.has_session(sid) assert mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt="")) == [] + _check_invariants(samples) + _record("2.12 drop clears sid (2nd get_trajectory -> [])", mgr, sid, samples) print("PASS 2.12") @@ -755,7 +800,10 @@ def test_3_4_rewrite_merge_ambiguous_forks(): a_c, t1 = asst_msg("c"), tool_msg("t") append(mgr, sid, [s, u, a_c, t1], "d") assert len(_leaves(mgr, sid)) == 3, "ambiguous candidates fork" - _record("3.4 rewrite-merge ambiguous -> fork", mgr, sid, []) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=3.0) + assert len(samples) == 3 + _check_invariants(samples) + _record("3.4 rewrite-merge ambiguous -> fork", mgr, sid, samples) print("PASS 3.4") @@ -922,10 +970,9 @@ def _print_case(title: str, mgr, sid: str, samples: list) -> None: print("[tree]") for line in txt.splitlines(): print(" " + line) - if samples: - print(f"[samples] {len(samples)}") - for i, s in enumerate(samples): - _print_sample(i, s) + print(f"[samples] {len(samples)}") + for i, s in enumerate(samples): + _print_sample(i, s) # =========================================================================== From 0be5966e5f8a7d75274b2ca980c854314f2a99f5 Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Mon, 8 Jun 2026 14:30:39 +0000 Subject: [PATCH 18/28] test(agent): make reward-split explicit in dump + conservation assert The dump previously printed only the already-divided per-sample reward, so the 'reward / n_samples' averaging wasn't visible. Now the [samples] header shows the split (input / n = per-sample) and get_traj asserts the per-sample shares sum back to the input reward (the averaging invariant). --- .../test_agent/test_trajectory_manager_e2e.py | 29 +++++++++++++++++-- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/tests/test_agent/test_trajectory_manager_e2e.py b/tests/test_agent/test_trajectory_manager_e2e.py index 51fb936e63..b51897304f 100644 --- a/tests/test_agent/test_trajectory_manager_e2e.py +++ b/tests/test_agent/test_trajectory_manager_e2e.py @@ -274,17 +274,34 @@ def _leaves(mgr, sid): # though get_trajectory consumes the session. _TREE_SNAP: dict[str, str] = {} +# Input reward passed to get_trajectory, keyed by sid, so the dump can show the +# split (input_reward / n_samples == per_sample_reward) explicitly. +_REWARD_IN: dict[str, float] = {} + def get_traj(mgr, sid, *args, **kwargs): """get_trajectory wrapper that snapshots the tree before draining. Linearization (get_trajectory) pops the sid, so a later dump would only see ````. Capturing the tree text here keeps the routing tree visible - next to the Samples it produced. + next to the Samples it produced. The input ``reward`` is captured too so the + dump can show how it splits across the emitted samples. """ if mgr.has_session(sid): _TREE_SNAP[sid] = dump_tree_txt(mgr, sid) - return mgr.get_trajectory(sid, *args, **kwargs) + _REWARD_IN[sid] = kwargs.get("reward", 0.0) + samples = mgr.get_trajectory(sid, *args, **kwargs) + # Reward conservation: get_trajectory splits the input reward evenly across + # every emitted sample, so the per-sample shares must sum back to the input + # (modulo float error). This is the "averaged over sample count" invariant. + if samples: + total = sum(s.reward for s in samples) + assert abs(total - _REWARD_IN[sid]) < 1e-9, ( + "reward not conserved across split", + total, + _REWARD_IN[sid], + ) + return samples def _check_invariants(samples): @@ -970,7 +987,13 @@ def _print_case(title: str, mgr, sid: str, samples: list) -> None: print("[tree]") for line in txt.splitlines(): print(" " + line) - print(f"[samples] {len(samples)}") + n = len(samples) + if n: + r_in = _REWARD_IN.get(sid, 0.0) + per = r_in / n + print(f"[samples] {n} (reward split: {r_in:.3f} / {n} = {per:.3f} per sample)") + else: + print(f"[samples] {n}") for i, s in enumerate(samples): _print_sample(i, s) From fe7692a9bbddc46a24d141cafeb3b221ef522319 Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Mon, 8 Jun 2026 14:48:21 +0000 Subject: [PATCH 19/28] test(agent): set every case's input reward to 1.0 Previously cases used arbitrary input rewards (2.0/3.0/4.0) with no semantic meaning, which was confusing. Now every get_trajectory call uses reward=1.0; per-sample split varies only by sample count (1.0/N), and assertions check the even split generically instead of magic numbers. --- .../test_agent/test_trajectory_manager_e2e.py | 60 ++++++++++--------- 1 file changed, 31 insertions(+), 29 deletions(-) diff --git a/tests/test_agent/test_trajectory_manager_e2e.py b/tests/test_agent/test_trajectory_manager_e2e.py index b51897304f..cb4cd2e9ac 100644 --- a/tests/test_agent/test_trajectory_manager_e2e.py +++ b/tests/test_agent/test_trajectory_manager_e2e.py @@ -361,7 +361,7 @@ def test_1_3_system_fork(): root = mgr._trees[sid] assert len(root.children) == 2, "different system -> two subtrees at root" assert len(_leaves(mgr, sid)) == 2 - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 2 _check_invariants(samples) _record("1.3 system fork -> two subtrees at root", mgr, sid, samples) @@ -378,7 +378,7 @@ def test_1_4_user_fork_shared_system(): assert len(root.children) == 1, "system shared" assert len(root.children[0].children) == 2, "user level forks" assert len(_leaves(mgr, sid)) == 2 - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 2 _check_invariants(samples) _record("1.4 user fork (shared system)", mgr, sid, samples) @@ -394,7 +394,7 @@ def test_1_5_assistant_message_fork(): append(mgr, sid, [s, u], "a2") user_node = mgr._trees[sid].children[0].children[0] assert len(user_node.children) == 2, "two assistant leaves hang off shared user" - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 2 _check_invariants(samples) _record("1.5 assistant fork under shared user", mgr, sid, samples) @@ -416,7 +416,7 @@ def test_1_6_tool_fork_shared_assistant(): assert len(_leaves(mgr, sid)) == 2 # The shared assistant (turn 1) is trained on the first leaf only; the second # re-emits it as loss=0 context (cross-leaf dedup). - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 2 _check_invariants(samples) _record("1.6 tool fork (shared assistant snapshot)", mgr, sid, samples) @@ -441,7 +441,7 @@ def test_1_7_token_only_drift_no_fork(): assert len(user_node.children) == 2, "two assistant turns share the (sys,user) path" assert len(mgr._trees[sid].children) == 1 - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 2 s_a, s_b = samples # Leaf 1: clean. Leaf 2: tokens carry the token, but it sits in the @@ -563,14 +563,14 @@ def test_2_3_drift_case_A_forks(): p2_honest = render_prompt([s, u, a1, t]) p2 = drift(p2_honest, len(p1) - 1) # inside p1's prompt region p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2, logprobs=[-0.4] * 2) - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 2 s1, s2 = samples assert s1.tokens == p1 + r1 assert s1.loss_mask == [1] * len(r1) assert s2.tokens == p2 + r2 assert s2.loss_mask == [1] * len(r2) - assert all(s.reward == 1.0 for s in samples) + assert all(abs(s.reward - 1.0 / len(samples)) < 1e-9 for s in samples) _check_invariants(samples) _record("2.3 drift case A (prompt region) -> fork", mgr, sid, samples) print("PASS 2.3") @@ -608,11 +608,11 @@ def test_2_5_drift_case_B1_long_forks(): p2_honest = render_prompt([s, u, a1, t]) p2 = drift_replace(p2_honest, len(p1) + len(r1) - 1) p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2) - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 2 assert samples[0].tokens == p1 + r1 assert samples[1].tokens == p2 + r2 - assert all(s.reward == 1.0 for s in samples) + assert all(abs(s.reward - 1.0 / len(samples)) < 1e-9 for s in samples) _check_invariants(samples) _record("2.5 drift case B1 (long) -> fork", mgr, sid, samples) print("PASS 2.5") @@ -626,7 +626,7 @@ def test_2_6_drift_case_B1_threshold_zero_forks(): p2_honest = render_prompt([s, u, a1, t]) p2 = drift_replace(p2_honest, len(p1) + len(r1) - 1) append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2) - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 2 _check_invariants(samples) _record("2.6 drift case B1 threshold=0 -> fork", mgr, sid, samples) @@ -645,13 +645,13 @@ def test_2_7_drift_case_B2_earlier_turn_forks(): p3_honest = render_prompt([s, u, a1, t1, a2, t2]) p3 = drift_replace(p3_honest, len(p1) + len(r1) - 1) # inside r1 (earlier span) p3, r3 = append(mgr, sid, [s, u, a1, t1, a2, t2], "a3", prompt_ids=p3, logprobs=[-0.3] * 2) - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 2 s1, s2 = samples L12 = _lcp_len(p1 + r1, p2) assert s1.tokens == p1 + r1 + p2[L12:] + r2 assert s2.tokens == p3 + r3 - assert all(s.reward == 1.0 for s in samples) + assert all(abs(s.reward - 1.0 / len(samples)) < 1e-9 for s in samples) _check_invariants(samples) _record("2.7 drift case B2 (earlier turn) -> fork", mgr, sid, samples) print("PASS 2.7") @@ -665,10 +665,11 @@ def test_2_8_fork_reward_split(): p2_honest = render_prompt([s, u, a1, t]) p2 = drift(p2_honest, len(p1) - 1) # prompt region -> case A fork append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2) - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=3.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 2 - assert all(abs(s.reward - 1.5) < 1e-9 for s in samples) - _record("2.8 fork reward split (3.0 / 2)", mgr, sid, samples) + # reward 1.0 split evenly across the 2 forked samples -> 0.5 each. + assert all(abs(s.reward - 0.5) < 1e-9 for s in samples) + _record("2.8 fork reward split (1.0 / 2 = 0.5 each)", mgr, sid, samples) print("PASS 2.8") @@ -678,11 +679,12 @@ def test_2_9_two_leaves_reward_split(): s = sys_msg("S") for ul in ["A", "B"]: append(mgr, sid, [s, usr_msg(ul)], ul.lower()) - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 2 - assert all(s.reward == 1.0 for s in samples) + # reward 1.0 split evenly across the 2 leaves -> 0.5 each. + assert all(abs(s.reward - 0.5) < 1e-9 for s in samples) _check_invariants(samples) - _record("2.9 two leaves reward split (2.0 / 2)", mgr, sid, samples) + _record("2.9 two leaves reward split (1.0 / 2 = 0.5 each)", mgr, sid, samples) print("PASS 2.9") @@ -696,7 +698,7 @@ def test_2_10_cross_leaf_dedup(): p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) p2, r2 = append(mgr, sid, [s, u, a1, tx], "a2", logprobs=[-0.4] * 2) p3, r3 = append(mgr, sid, [s, u, a1, ty], "a3", logprobs=[-0.3] * 2) - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 2 s_first, s_second = samples L2 = _lcp_len(p1 + r1, p2) @@ -728,7 +730,7 @@ def test_2_11_routing_only_assistant_filtered(): chain = leaves[0].path_from_root() routing = [n for n in chain if n.role == "assistant" and n.turn_prompt_ids is None] assert len(routing) == 1 and routing[0].messages[0]["content"] == "foreign" - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt="")) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 1 _record("2.11 routing-only assistant filtered (no raise)", mgr, sid, samples) print("PASS 2.11") @@ -786,7 +788,7 @@ def test_3_2_rewrite_merge_long_forks(): append(mgr, sid, [s, u], "ok", finish_reason="tool_calls") append(mgr, sid, [s, u, a1_rw, t1], "done") assert len(_leaves(mgr, sid)) == 2, "long rewrite forks" - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 2 _check_invariants(samples) _record("3.2 rewrite-merge long -> fork", mgr, sid, samples) @@ -801,7 +803,7 @@ def test_3_3_rewrite_merge_threshold_zero_forks(): append(mgr, sid, [s, u], "ok", finish_reason="tool_calls") append(mgr, sid, [s, u, a1_rw, t1], "done") assert len(_leaves(mgr, sid)) == 2 - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=2.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) _check_invariants(samples) _record("3.3 rewrite-merge threshold=0 -> fork", mgr, sid, samples) print("PASS 3.3") @@ -817,7 +819,7 @@ def test_3_4_rewrite_merge_ambiguous_forks(): a_c, t1 = asst_msg("c"), tool_msg("t") append(mgr, sid, [s, u, a_c, t1], "d") assert len(_leaves(mgr, sid)) == 3, "ambiguous candidates fork" - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=3.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 3 _check_invariants(samples) _record("3.4 rewrite-merge ambiguous -> fork", mgr, sid, samples) @@ -837,7 +839,7 @@ def test_3_5_rewrite_merge_match_key_updated(): p3, r3 = append(mgr, sid, [s, u, a1_rw, t1, a2, t2], "third", logprobs=[-0.3] * 2) leaves = _leaves(mgr, sid) assert len(leaves) == 1, "match_key updated -> no spurious fork" - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt="")) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 1 L = _lcp_len(p2 + r2, p3) assert samples[0].tokens == p2 + r2 + p3[L:] + r3 @@ -864,10 +866,10 @@ def test_3_6_tree_fork_plus_token_drift(): append(mgr, sid, [s, u, a1, tx, ax2, txx], "ax3", prompt_ids=p3, logprobs=[-0.2] * 2) # Leaf Y: a separate tool result off the shared assistant. append(mgr, sid, [s, u, a1, ty], "ay2", logprobs=[-0.1] * 2) - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=3.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) # Leaf X -> 2 samples (drift fork), Leaf Y -> 1 sample. Total 3. assert len(samples) == 3, [s.tokens for s in samples] - assert all(abs(s.reward - 1.0) < 1e-9 for s in samples) + assert all(abs(s.reward - 1.0 / 3) < 1e-9 for s in samples) _check_invariants(samples) _record("3.6 tree fork + token drift -> 3 samples", mgr, sid, samples) print("PASS 3.6") @@ -886,7 +888,7 @@ def test_3_7_deep_multi_leaf_dedup(): # three different tool results off a2 -> three leaves sharing a1+a2 for lbl in ["p", "q", "r"]: append(mgr, sid, [s, u, a1, t1, a2, tool_msg(lbl)], f"end-{lbl}", logprobs=[-0.3] * 2) - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=3.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 3 # Count trained tokens for r1 and r2 across all samples: each shared turn # trained exactly once (first leaf), others loss=0. @@ -928,10 +930,10 @@ def test_3_8_long_mixed_session(): p7_honest = render_prompt(prefix) p7 = drift(p7_honest, len(p1) - 1) append(mgr, sid, prefix, "a5", prompt_ids=p7, logprobs=lp) - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=4.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) # The final case-A fork splits the single leaf chain into >=2 segments. assert len(samples) >= 2, len(samples) - assert abs(sum(s.reward for s in samples) - 4.0) < 1e-9 + assert abs(sum(s.reward for s in samples) - 1.0) < 1e-9 _check_invariants(samples) _record(f"3.8 long mixed session -> {len(samples)} samples", mgr, sid, samples) print("PASS 3.8") From 2ee5ee8204b7b0f47d5fde6167fc4170325762b9 Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Mon, 8 Jun 2026 14:56:53 +0000 Subject: [PATCH 20/28] =?UTF-8?q?test(agent):=20render=20whitespace=20in?= =?UTF-8?q?=20token=20labels=20as=20visible=20=E2=90=A3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Whitespace-only rewrite drift (e.g. cc turning 'ok' into 'ok ') was invisible in the dump, making 3.1's rewrite-merge trigger impossible to see. _vis() now shows spaces as ␣ in [raw turns] and [samples] labels. --- tests/test_agent/test_trajectory_manager_e2e.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/tests/test_agent/test_trajectory_manager_e2e.py b/tests/test_agent/test_trajectory_manager_e2e.py index cb4cd2e9ac..eaedeff4dd 100644 --- a/tests/test_agent/test_trajectory_manager_e2e.py +++ b/tests/test_agent/test_trajectory_manager_e2e.py @@ -67,6 +67,16 @@ def name_of(tok: int) -> str: return TOKEN_NAMES.get(tok, str(tok)) +def _vis(label: str) -> str: + """Make whitespace visible in a token label for the dump. + + Whitespace-only drift (e.g. a trailing space from a cc rewrite) is invisible + in a terminal, which makes ``r:ok`` vs ``r:ok `` indistinguishable. Render + spaces as ``␣`` so the difference is obvious in the readable output. + """ + return label.replace(" ", "␣") + + _ASST_BODY: dict[str, int] = {} @@ -85,7 +95,7 @@ def _asst_body(label: str) -> int: if label not in _ASST_BODY: body = _BANDS["assistant"] + 100 + len(_ASST_BODY) _ASST_BODY[label] = body - TOKEN_NAMES[body] = f"r:{label}" + TOKEN_NAMES[body] = f"r:{_vis(label)}" return _ASST_BODY[label] @@ -121,7 +131,7 @@ def __init__(self, role: str, label: str) -> None: idx = MsgTok._body_counter.setdefault(role, 0) + 1 MsgTok._body_counter[role] = idx self.body = base + 10 + idx - TOKEN_NAMES[self.body] = f"{role}:{label}" + TOKEN_NAMES[self.body] = f"{role}:{_vis(label)}" # message dict as the manager sees it (drives node_match_key). self.message = {"role": role, "content": label} @@ -248,7 +258,7 @@ def append( # consumes them. _TURN_LOG.setdefault(sid, []).append( { - "prompt_msgs": [f"{m.role}:{m.label}" for m in prompt_msgs], + "prompt_msgs": [f"{m.role}:{_vis(m.label)}" for m in prompt_msgs], "prompt_ids": p, "response_ids": r, "finish": finish_reason, From 8eed4dd4d5b5fbdb3d5bff2471b6ca9bff112f19 Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Mon, 8 Jun 2026 15:11:40 +0000 Subject: [PATCH 21/28] test(agent): add Group 4 (boundary/defensive/feature) -> 98% coverage Coverage of trajectory_manager.py rose 94%->98%. New cases: - 4.1 tools metadata attaches to first system node only - 4.2 logprobs/ids length mismatch raises - 4.3 empty prompt_messages skipped (no-op) - 4.4 default base_sample (None) - 4.5 mixed logprobs across turns (turn2 padded 0.0) - 4.6 case-B1 drift threshold boundary (d==threshold forks, d ValueError at append_turn.""" + mgr = TrajectoryManager() + sid = "4.2" + s, u = sys_msg("S"), usr_msg("u") + bad = TurnRecord( + prompt_ids=render_prompt([s, u]), + output_ids=[9101, 9102, 9103], + finish_reason="stop", + output_log_probs=[-0.1, -0.2], # length 2 != 3 + ) + raised = False + try: + mgr.append_turn( + sid, + turn=bad, + prompt_messages=messages([s, u]), + tools=None, + response_message={"role": "assistant", "content": "x"}, + ) + except ValueError as e: + raised = True + assert "output_log_probs" in str(e) + assert raised, "expected ValueError on logprobs/ids length mismatch" + print("PASS 4.2") + + +def test_4_3_empty_prompt_messages_skipped(): + """Empty prompt_messages -> append_turn is a no-op (warns, no node, no turn).""" + mgr = TrajectoryManager() + sid = "4.3" + mgr.append_turn( + sid, + turn=turn([1], [2], finish_reason="stop"), + prompt_messages=[], + tools=None, + response_message=None, + ) + assert mgr.turn_count(sid) == 0 + # The tree may be created empty (root only) or absent; either way no leaf. + assert not mgr.has_session(sid) or list(_leaves(mgr, sid)) == [] + print("PASS 4.3") + + +def test_4_4_default_base_sample(): + """get_trajectory with base_sample=None uses a default Sample(index=0).""" + mgr = TrajectoryManager() + sid = "4.4" + s, u = sys_msg("S"), usr_msg("u") + append(mgr, sid, [s, u], "a") + # snapshot tree before drain so the dump still renders it + _TREE_SNAP[sid] = dump_tree_txt(mgr, sid) + _REWARD_IN[sid] = 1.0 + samples = mgr.get_trajectory(sid, reward=1.0) # no base_sample + assert len(samples) == 1 + assert samples[0].index == 0 + _check_invariants(samples) + _record("4.4 default base_sample (None)", mgr, sid, samples) + print("PASS 4.4") + + +def test_4_5_mixed_logprobs_across_turns(): + """A trajectory where turn 1 carries logprobs and turn 2 does NOT: the + sample's turn-1 response region has real logprobs, the turn-2 region is + padded with 0.0 (the response is still trained, loss=1).""" + mgr = TrajectoryManager() + sid = "4.5" + s, u = sys_msg("S"), usr_msg("u") + a1, t1 = asst_msg("call"), tool_msg("t") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) + p2, r2 = append(mgr, sid, [s, u, a1, t1], "done", logprobs=None) # no logprobs + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + s0 = samples[0] + L = _lcp_len(p1 + r1, p2) + # turn-1 region: real logprobs; prompt tail: 0.0; turn-2 region: 0.0 (padded). + assert s0.rollout_log_probs == [-0.5] * len(r1) + [0.0] * (len(p2) - L) + [0.0] * len(r2) + # both responses still trained. + assert s0.loss_mask == [1] * len(r1) + [0] * (len(p2) - L) + [1] * len(r2) + _check_invariants(samples) + _record("4.5 mixed logprobs across turns (turn2 padded 0.0)", mgr, sid, samples) + print("PASS 4.5") + + +def test_4_6_drift_B1_threshold_boundary(): + """case-B1 threshold is exclusive: a drift tail of length d == threshold + forks, d == threshold-1 replaces. Verify both sides of the boundary.""" + + def run(threshold, drift_tail_len): + mgr = TrajectoryManager(fork_merge_max_response_tokens=threshold) + sid = f"4.6-{threshold}-{drift_tail_len}" + s, u = sys_msg("S"), usr_msg("u") + # 4-token response so the divergence can sit d tokens before its end. + p1 = render_prompt([s, u]) + r1 = [9001, 9002, 9003, 9004] + mgr.append_turn( + sid, + turn=turn(p1, r1, finish_reason="tool_calls"), + prompt_messages=messages([s, u]), + tools=None, + response_message={"role": "assistant", "content": "a1"}, + ) + a1m = {"role": "assistant", "content": "a1"} + tm = tool_msg("t") + # honest turn-2 prompt echoes p1 + r1 then the tool block + gen marker. + p2_honest = p1 + r1 + tm.render() + [_GEN] + # divergence d tokens before the end of r1's echo (inside its response span). + drift_idx = len(p1) + len(r1) - drift_tail_len + p2 = drift_replace(p2_honest, drift_idx) + r2 = [9101, 9102] + mgr.append_turn( + sid, + turn=turn(p2, r2, finish_reason="stop"), + prompt_messages=[*messages([s, u]), a1m, tm.message], + tools=None, + response_message={"role": "assistant", "content": "done"}, + ) + return get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + + forked = run(threshold=2, drift_tail_len=2) # d == threshold -> fork + assert len(forked) == 2, f"d==threshold must fork, got {len(forked)}" + replaced = run(threshold=2, drift_tail_len=1) # d < threshold -> replace + assert len(replaced) == 1, f"d None: test_3_6_tree_fork_plus_token_drift, test_3_7_deep_multi_leaf_dedup, test_3_8_long_mixed_session, + test_4_1_tools_metadata_on_first_system_only, + test_4_2_logprobs_length_mismatch_raises, + test_4_3_empty_prompt_messages_skipped, + test_4_4_default_base_sample, + test_4_5_mixed_logprobs_across_turns, + test_4_6_drift_B1_threshold_boundary, ] From cfad29db8c752eefa2d699d1c7a1fe7c5845540d Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Mon, 8 Jun 2026 15:32:20 +0000 Subject: [PATCH 22/28] test(agent): assert full output via golden token+loss strings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace hand-derived loss_mask index arithmetic (error-prone — it was wrong twice during review) with golden string assertions. Each sample renders to a readable line where trained tokens (loss=1) are wrapped in [...] and context (loss=0 / stripped prompt) is bare, e.g. system:S user:u [r:call] [] ... Every case now pins its FULL linearized output as one human-reviewable literal, so any change to tokens, response boundary, or which tokens carry training signal is caught. Verified: stripping the [...] brackets (a loss_mask regression) fails the assertion. 36 cases pass, 98% coverage. --- .../test_agent/test_trajectory_manager_e2e.py | 323 ++++++++++++++---- 1 file changed, 249 insertions(+), 74 deletions(-) diff --git a/tests/test_agent/test_trajectory_manager_e2e.py b/tests/test_agent/test_trajectory_manager_e2e.py index 1b7f61811f..47d1c33296 100644 --- a/tests/test_agent/test_trajectory_manager_e2e.py +++ b/tests/test_agent/test_trajectory_manager_e2e.py @@ -343,6 +343,35 @@ def _check_invariants(samples): assert sum(s.loss_mask) > 0, "fully-masked sample emitted" +def golden(sample) -> str: + """Render one Sample as a human-reviewable golden string. + + Every token is decoded to its readable name (````, ``r:done``, + ```` ...). The leading prompt prefix (no loss_mask entry) is shown as + plain names; the response region is shown with each TRAINED token (loss=1) + wrapped in ``[...]`` and each context token (loss=0) left bare. This makes the + full linearized result — tokens, where the response region starts, and + exactly which tokens carry training signal — a single literal a human can + eyeball and assert against, instead of hand-derived index arithmetic. + + Example: `` system:S user:u [r:ok] []`` + """ + toks = sample.tokens + resp_start = len(toks) - sample.response_length + parts: list[str] = [] + for i, t in enumerate(toks): + nm = name_of(t) + if i >= resp_start and sample.loss_mask[i - resp_start] == 1: + parts.append(f"[{nm}]") + else: + parts.append(nm) + return " ".join(parts) + + +def goldens(samples) -> list[str]: + return [golden(s) for s in samples] + + # =========================================================================== # §2 Group 1 — routing tree layer (append_turn shapes the tree) # =========================================================================== @@ -353,12 +382,15 @@ def test_1_1_single_turn_chain(): sid = "1.1" s = sys_msg("S") u = usr_msg("compute") - append(mgr, sid, [s, u], "ok") + p, r = append(mgr, sid, [s, u], "ok") chain = _leaves(mgr, sid)[0].path_from_root() assert [n.role for n in chain] == ["system", "user", "assistant"] assert chain[-1].turn_index == 1 samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 1 + assert goldens(samples) == [ + " system:S user:compute [r:ok] []", + ] _check_invariants(samples) _record("1.1 single turn -> linear chain", mgr, sid, samples) print("PASS 1.1") @@ -369,13 +401,17 @@ def test_1_2_clean_multiturn_with_tool(): sid = "1.2" s, u = sys_msg("S"), usr_msg("compute") a1, t1 = asst_msg("call"), tool_msg("4") - append(mgr, sid, [s, u], "call", finish_reason="tool_calls") - append(mgr, sid, [s, u, a1, t1], "done") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls") + p2, r2 = append(mgr, sid, [s, u, a1, t1], "done") chain = _leaves(mgr, sid)[0].path_from_root() assert [n.role for n in chain] == ["system", "user", "assistant", "tool", "assistant"] assert mgr.turn_count(sid) == 2 samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 1 + assert goldens(samples) == [ + " system:S user:compute [r:call] [] " + " tool:4 [r:done] []", + ] _check_invariants(samples) _record("1.2 clean 2-turn with tool -> single chain", mgr, sid, samples) print("PASS 1.2") @@ -390,7 +426,10 @@ def test_1_3_system_fork(): assert len(root.children) == 2, "different system -> two subtrees at root" assert len(_leaves(mgr, sid)) == 2 samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 2 + assert goldens(samples) == [ + " system:SA user:u [r:a] []", + " system:SB user:u [r:a] []", + ] _check_invariants(samples) _record("1.3 system fork -> two subtrees at root", mgr, sid, samples) print("PASS 1.3") @@ -407,7 +446,10 @@ def test_1_4_user_fork_shared_system(): assert len(root.children[0].children) == 2, "user level forks" assert len(_leaves(mgr, sid)) == 2 samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 2 + assert goldens(samples) == [ + " system:S user:A [r:a] []", + " system:S user:B [r:b] []", + ] _check_invariants(samples) _record("1.4 user fork (shared system)", mgr, sid, samples) print("PASS 1.4") @@ -423,7 +465,11 @@ def test_1_5_assistant_message_fork(): user_node = mgr._trees[sid].children[0].children[0] assert len(user_node.children) == 2, "two assistant leaves hang off shared user" samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 2 + # Two independent single-turn leaves sharing only the (sys,user) prefix. + assert goldens(samples) == [ + " system:S user:u [r:a1] []", + " system:S user:u [r:a2] []", + ] _check_invariants(samples) _record("1.5 assistant fork under shared user", mgr, sid, samples) print("PASS 1.5") @@ -442,10 +488,14 @@ def test_1_6_tool_fork_shared_assistant(): assert asst1.role == "assistant" and asst1.turn_prompt_ids is not None assert len(asst1.children) == 2, "shared assistant forks at the tool level" assert len(_leaves(mgr, sid)) == 2 - # The shared assistant (turn 1) is trained on the first leaf only; the second - # re-emits it as loss=0 context (cross-leaf dedup). samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 2 + # Leaf X owns the shared turn 1 (r:call trained); leaf Y shares it -> r:call + # demoted to loss=0 context, only r:ay trains (cross-leaf dedup). + assert goldens(samples) == [ + " system:S user:u [r:call] [] " + " tool:x [r:ax] []", + " system:S user:u r:call " " tool:y [r:ay] []", + ] _check_invariants(samples) _record("1.6 tool fork (shared assistant snapshot)", mgr, sid, samples) print("PASS 1.6") @@ -459,9 +509,8 @@ def test_1_7_token_only_drift_no_fork(): mgr = TrajectoryManager() sid = "1.7" s, u = sys_msg("S"), usr_msg("u") - pa, ra = append(mgr, sid, [s, u], "a") + pa, _ = append(mgr, sid, [s, u], "a") tampered = drift(pa, 1) # spliced into the prompt at index 1 - rb = render_response("b") append(mgr, sid, [s, u], "b", prompt_ids=tampered) # Tree: (sys,user) shared, two assistant turns hang off it -> two leaves; the # path above the assistant is single (NOT forked on tokens). @@ -471,12 +520,15 @@ def test_1_7_token_only_drift_no_fork(): samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 2 - s_a, s_b = samples - # Leaf 1: clean. Leaf 2: tokens carry the token, but it sits in the - # stripped prompt region — loss_mask covers only the response (rb). - assert s_a.tokens == pa + ra - assert s_b.tokens == tampered + rb - assert s_b.loss_mask == [1] * len(rb), "drift confined to prompt; response fully trained" + # Leaf 1: clean. Leaf 2: the token sits in the stripped prompt region + # (bare, no brackets); the response r:b is fully trained ([...]). + assert goldens(samples) == [ + " system:S user:u [r:a] []", + " system:S user:u [r:b] []", + ] + # Belt-and-suspenders on the drift placement: token present, but never inside + # the response region. + s_b = samples[1] assert (_DRIFT_BAND + 1) in s_b.tokens, "drift token is still carried in the sample" assert (_DRIFT_BAND + 1) not in s_b.tokens[len(tampered) :], "drift not in the response region" _check_invariants(samples) @@ -497,6 +549,10 @@ def test_1_8_multi_tool_per_turn(): assert chain[4].messages == [tb.message] samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 1 + assert goldens(samples) == [ + " system:S user:u [r:call] [] " + " tool:A tool:B [r:done] []", + ] _check_invariants(samples) _record("1.8 multi-tool turn -> one node per tool", mgr, sid, samples) print("PASS 1.8") @@ -512,7 +568,8 @@ def test_1_9_cross_sid_isolation(): assert mgr._trees["sid-a"] is not mgr._trees["sid-b"] sa = get_traj(mgr, "sid-a", base_sample=Sample(index=0, prompt=""), reward=1.0) sb = get_traj(mgr, "sid-b", base_sample=Sample(index=1, prompt=""), reward=1.0) - assert len(sa) == 1 and len(sb) == 1 + assert goldens(sa) == [" system:S user:A [r:a] []"] + assert goldens(sb) == [" system:S user:B [r:b] []"] _check_invariants(sa) _check_invariants(sb) _record("1.9 cross-sid isolation (sid-a)", mgr, "sid-a", sa) @@ -553,10 +610,8 @@ def test_2_1_single_turn_linearize(): samples = get_traj(mgr, sid, base_sample=Sample(index=7, prompt="hi"), reward=1.0) assert len(samples) == 1 s0 = samples[0] - assert s0.tokens == p + r - assert s0.loss_mask == [1] * len(r) + assert goldens(samples) == [" system:S user:u [r:a] []"] assert s0.rollout_log_probs == [-0.5] * len(r) - assert s0.response_length == len(r) assert s0.reward == 1.0 _check_invariants(samples) _record("2.1 single-turn linearize", mgr, sid, samples) @@ -573,9 +628,10 @@ def test_2_2_clean_multiturn_linearize(): assert len(samples) == 1 s0 = samples[0] L = _lcp_len(p1 + r1, p2) - assert L == len(p1) + len(r1) - assert s0.tokens == p1 + r1 + p2[L:] + r2 - assert s0.loss_mask == [1] * len(r1) + [0] * (len(p2) - L) + [1] * len(r2) + assert goldens(samples) == [ + " system:S user:u [r:call] [] " + " tool:4 [r:done] []", + ] assert s0.rollout_log_probs == [-0.5] * len(r1) + [0.0] * (len(p2) - L) + [-0.4] * len(r2) _check_invariants(samples) _record("2.2 clean 2-turn linearize", mgr, sid, samples) @@ -593,11 +649,13 @@ def test_2_3_drift_case_A_forks(): p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2, logprobs=[-0.4] * 2) samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 2 - s1, s2 = samples - assert s1.tokens == p1 + r1 - assert s1.loss_mask == [1] * len(r1) - assert s2.tokens == p2 + r2 - assert s2.loss_mask == [1] * len(r2) + # case-A fork: two coherent single-turn segments; the token stays in + # segment 2's stripped prompt region (bare), no token dropped. + assert goldens(samples) == [ + " system:S user:u [r:call] []", + " system:S user:u r:call " + " tool:t [r:done] []", + ] assert all(abs(s.reward - 1.0 / len(samples)) < 1e-9 for s in samples) _check_invariants(samples) _record("2.3 drift case A (prompt region) -> fork", mgr, sid, samples) @@ -620,8 +678,13 @@ def test_2_4_drift_case_B1_short_replaces(): s0 = samples[0] L = _lcp_len(p1 + r1, p2) assert L == drift_idx - assert s0.tokens == p2 + r2 - assert s0.loss_mask == [1] * (L - len(p1)) + [0] * (len(p2) - L) + [1] * len(r2) + # replace: the drifted r:call tail is dropped and re-supplied as loss=0 prompt + # context (the token marks the divergence); only the surviving head of + # r:call and the new r:done train. + assert goldens(samples) == [ + " system:S user:u [r:call] " + " tool:t [r:done] []", + ] assert s0.rollout_log_probs == [-0.5] * (L - len(p1)) + [0.0] * (len(p2) - L) + [-0.4] * len(r2) _check_invariants(samples) _record("2.4 drift case B1 (small) -> replace", mgr, sid, samples) @@ -638,8 +701,13 @@ def test_2_5_drift_case_B1_long_forks(): p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2) samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 2 - assert samples[0].tokens == p1 + r1 - assert samples[1].tokens == p2 + r2 + # Both segments single-turn (the drift forked them apart): each trains its own + # response. The sits in segment 2's stripped prompt. + assert goldens(samples) == [ + " system:S user:u [r:call] []", + " system:S user:u r:call " + " tool:t [r:done] []", + ] assert all(abs(s.reward - 1.0 / len(samples)) < 1e-9 for s in samples) _check_invariants(samples) _record("2.5 drift case B1 (long) -> fork", mgr, sid, samples) @@ -653,9 +721,14 @@ def test_2_6_drift_case_B1_threshold_zero_forks(): p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls") p2_honest = render_prompt([s, u, a1, t]) p2 = drift_replace(p2_honest, len(p1) + len(r1) - 1) - append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2) + p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2) samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 2 + assert goldens(samples) == [ + " system:S user:u [r:call] []", + " system:S user:u r:call " + " tool:t [r:done] []", + ] _check_invariants(samples) _record("2.6 drift case B1 threshold=0 -> fork", mgr, sid, samples) print("PASS 2.6") @@ -675,10 +748,14 @@ def test_2_7_drift_case_B2_earlier_turn_forks(): p3, r3 = append(mgr, sid, [s, u, a1, t1, a2, t2], "a3", prompt_ids=p3, logprobs=[-0.3] * 2) samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 2 - s1, s2 = samples - L12 = _lcp_len(p1 + r1, p2) - assert s1.tokens == p1 + r1 + p2[L12:] + r2 - assert s2.tokens == p3 + r3 + # Segment 1 = clean turns 1+2; segment 2 = turn 3 alone (forked because the + # drift hit an EARLIER turn's response span, which replace can't drop). + assert goldens(samples) == [ + " system:S user:u [r:a1] [] " + " tool:t1 [r:a2] []", + " system:S user:u r:a1 " + " tool:t1 r:a2 tool:t2 [r:a3] []", + ] assert all(abs(s.reward - 1.0 / len(samples)) < 1e-9 for s in samples) _check_invariants(samples) _record("2.7 drift case B2 (earlier turn) -> fork", mgr, sid, samples) @@ -689,14 +766,21 @@ def test_2_8_fork_reward_split(): mgr = TrajectoryManager() sid = "2.8" s, u, a1, t = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("t") - p1, _ = append(mgr, sid, [s, u], "call", finish_reason="tool_calls") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls") p2_honest = render_prompt([s, u, a1, t]) p2 = drift(p2_honest, len(p1) - 1) # prompt region -> case A fork - append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2) + p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2) samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 2 + # case-A fork: two single-turn segments, each trains its own response. + assert goldens(samples) == [ + " system:S user:u [r:call] []", + " system:S user:u r:call " + " tool:t [r:done] []", + ] # reward 1.0 split evenly across the 2 forked samples -> 0.5 each. assert all(abs(s.reward - 0.5) < 1e-9 for s in samples) + _check_invariants(samples) _record("2.8 fork reward split (1.0 / 2 = 0.5 each)", mgr, sid, samples) print("PASS 2.8") @@ -709,6 +793,10 @@ def test_2_9_two_leaves_reward_split(): append(mgr, sid, [s, usr_msg(ul)], ul.lower()) samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 2 + assert goldens(samples) == [ + " system:S user:A [r:a] []", + " system:S user:B [r:b] []", + ] # reward 1.0 split evenly across the 2 leaves -> 0.5 each. assert all(abs(s.reward - 0.5) < 1e-9 for s in samples) _check_invariants(samples) @@ -728,12 +816,14 @@ def test_2_10_cross_leaf_dedup(): p3, r3 = append(mgr, sid, [s, u, a1, ty], "a3", logprobs=[-0.3] * 2) samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 2 - s_first, s_second = samples - L2 = _lcp_len(p1 + r1, p2) - assert s_first.tokens == p2 + r2 - assert s_first.loss_mask == [1] * len(r1) + [0] * (len(p2) - L2) + [1] * len(r2) - assert s_second.tokens == p3 + r3 - assert s_second.loss_mask == [0] * (len(p3) - len(p1)) + [1] * len(r3) + s_second = samples[1] + # First leaf trains the shared r:call + its own r:a2; second leaf shares + # r:call (demoted to loss=0) and trains only r:a3. + assert goldens(samples) == [ + " system:S user:u [r:call] [] " + " tool:x [r:a2] []", + " system:S user:u r:call " " tool:y [r:a3] []", + ] assert s_second.rollout_log_probs == [0.0] * (len(p3) - len(p1)) + [-0.3] * len(r3) _check_invariants(samples) _record("2.10 cross-leaf dedup (shared assistant trained once)", mgr, sid, samples) @@ -760,6 +850,13 @@ def test_2_11_routing_only_assistant_filtered(): assert len(routing) == 1 and routing[0].messages[0]["content"] == "foreign" samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 1 + # The foreign assistant (r:foreign) is routing-only -> appears as bare context + # (no brackets); the three real turns r:a1/r:a2/r:a3 train. + assert goldens(samples) == [ + " system:S user:u [r:a1] [] " + " tool:t1 [r:a2] [] r:foreign " + " tool:t2 [r:a3] []", + ] _record("2.11 routing-only assistant filtered (no raise)", mgr, sid, samples) print("PASS 2.11") @@ -771,7 +868,7 @@ def test_2_12_drop_clears_sid(): append(mgr, sid, [s, u], "a") assert mgr.has_session(sid) samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 1 + assert goldens(samples) == [" system:S user:u [r:a] []"] assert not mgr.has_session(sid) assert mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt="")) == [] _check_invariants(samples) @@ -790,8 +887,8 @@ def test_3_1_rewrite_merge_absorbs_short(): s, u = sys_msg("S"), usr_msg("u") a1_rw = asst_msg("ok ") # cc-rewritten (different message identity) t1 = tool_msg("t") - p1, r1 = append(mgr, sid, [s, u], "ok", finish_reason="tool_calls", logprobs=[-0.5] * 2) - p2, r2 = append(mgr, sid, [s, u, a1_rw, t1], "done", logprobs=[-0.4] * 2) + append(mgr, sid, [s, u], "ok", finish_reason="tool_calls", logprobs=[-0.5] * 2) + append(mgr, sid, [s, u, a1_rw, t1], "done", logprobs=[-0.4] * 2) leaves = _leaves(mgr, sid) assert len(leaves) == 1, "short rewrite absorbed, not forked" chain = leaves[0].path_from_root() @@ -801,8 +898,11 @@ def test_3_1_rewrite_merge_absorbs_short(): assert merged.metadata["merged_rewrite"]["abandoned_turn_index"] == 1 samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 1 - assert samples[0].tokens == p2 + r2 - assert samples[0].loss_mask == [1] * len(r2) + # The abandoned turn-1 response (r:ok␣) is demoted to routing-only -> appears + # bare; only the surviving turn-2 r:done trains. + assert goldens(samples) == [ + " system:S user:u r:ok␣ " " tool:t [r:done] []", + ] _check_invariants(samples) _record("3.1 rewrite-merge absorbs short assistant", mgr, sid, samples) print("PASS 3.1") @@ -818,6 +918,12 @@ def test_3_2_rewrite_merge_long_forks(): assert len(_leaves(mgr, sid)) == 2, "long rewrite forks" samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 2 + # Leaf 1: the abandoned turn-1 standalone (r:ok). Leaf 2: turn-2 only (the + # rewritten r:ok2␣ assistant mounts routing-only and is filtered out). + assert goldens(samples) == [ + " system:S user:u [r:ok] []", + " system:S user:u r:ok2␣ " " tool:t [r:done] []", + ] _check_invariants(samples) _record("3.2 rewrite-merge long -> fork", mgr, sid, samples) print("PASS 3.2") @@ -832,6 +938,11 @@ def test_3_3_rewrite_merge_threshold_zero_forks(): append(mgr, sid, [s, u, a1_rw, t1], "done") assert len(_leaves(mgr, sid)) == 2 samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 2 + assert goldens(samples) == [ + " system:S user:u [r:ok] []", + " system:S user:u r:ok3␣ " " tool:t [r:done] []", + ] _check_invariants(samples) _record("3.3 rewrite-merge threshold=0 -> fork", mgr, sid, samples) print("PASS 3.3") @@ -849,6 +960,13 @@ def test_3_4_rewrite_merge_ambiguous_forks(): assert len(_leaves(mgr, sid)) == 3, "ambiguous candidates fork" samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 3 + # Leaves "a" and "b" are standalone single turns; leaf "d" carries the + # ambiguous-rewrite assistant (r:c) as routing-only (bare) -> trains only r:d. + assert goldens(samples) == [ + " system:S user:u [r:a] []", + " system:S user:u [r:b] []", + " system:S user:u r:c " " tool:t [r:d] []", + ] _check_invariants(samples) _record("3.4 rewrite-merge ambiguous -> fork", mgr, sid, samples) print("PASS 3.4") @@ -863,14 +981,18 @@ def test_3_5_rewrite_merge_match_key_updated(): a1_rw = asst_msg("ok5 ") t1, a2, t2 = tool_msg("t1"), asst_msg("second"), tool_msg("t2") append(mgr, sid, [s, u], "ok", finish_reason="tool_calls") - p2, r2 = append(mgr, sid, [s, u, a1_rw, t1], "second", finish_reason="tool_calls", logprobs=[-0.4] * 2) - p3, r3 = append(mgr, sid, [s, u, a1_rw, t1, a2, t2], "third", logprobs=[-0.3] * 2) + append(mgr, sid, [s, u, a1_rw, t1], "second", finish_reason="tool_calls", logprobs=[-0.4] * 2) + append(mgr, sid, [s, u, a1_rw, t1, a2, t2], "third", logprobs=[-0.3] * 2) leaves = _leaves(mgr, sid) assert len(leaves) == 1, "match_key updated -> no spurious fork" samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 1 - L = _lcp_len(p2 + r2, p3) - assert samples[0].tokens == p2 + r2 + p3[L:] + r3 + # Turn 1 (r:ok) was absorbed as routing-only (rewrite merge), so it appears + # bare; turns 2 and 3 (r:second / r:third) train in one clean chain. + assert goldens(samples) == [ + " system:S user:u r:ok5␣ " + " tool:t1 [r:second] [] tool:t2 [r:third] []", + ] _check_invariants(samples) _record("3.5 rewrite-merge match_key updated", mgr, sid, samples) print("PASS 3.5") @@ -887,7 +1009,7 @@ def test_3_6_tree_fork_plus_token_drift(): ax2 = asst_msg("ax2") p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) # Leaf X: clean continuation, then a third turn with a case-A prompt drift. - p2, r2 = append(mgr, sid, [s, u, a1, tx], "ax2", finish_reason="tool_calls", logprobs=[-0.4] * 2) + append(mgr, sid, [s, u, a1, tx], "ax2", finish_reason="tool_calls", logprobs=[-0.4] * 2) txx = tool_msg("xx") p3_honest = render_prompt([s, u, a1, tx, ax2, txx]) p3 = drift(p3_honest, len(p1) - 1) # case A drift -> fork inside leaf X @@ -895,8 +1017,18 @@ def test_3_6_tree_fork_plus_token_drift(): # Leaf Y: a separate tool result off the shared assistant. append(mgr, sid, [s, u, a1, ty], "ay2", logprobs=[-0.1] * 2) samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - # Leaf X -> 2 samples (drift fork), Leaf Y -> 1 sample. Total 3. assert len(samples) == 3, [s.tokens for s in samples] + assert goldens(samples) == [ + # Sample 0: leaf X first segment, trains the shared r:call (first claim) + r:ax2. + " system:S user:u [r:call] [] " + " tool:x [r:ax2] []", + # Sample 1: leaf X second segment, a FRESH segment after the case-A fork -> + # whole prompt stripped, only r:ax3 trains (the sits in its prompt). + " system:S user:u r:call " + " tool:x r:ax2 tool:xx [r:ax3] []", + # Sample 2: leaf Y, shares r:call (claimed by sample 0 -> bare), trains r:ay2. + " system:S user:u r:call " " tool:y [r:ay2] []", + ] assert all(abs(s.reward - 1.0 / 3) < 1e-9 for s in samples) _check_invariants(samples) _record("3.6 tree fork + token drift -> 3 samples", mgr, sid, samples) @@ -911,19 +1043,24 @@ def test_3_7_deep_multi_leaf_dedup(): s, u = sys_msg("S"), usr_msg("u") a1, t1 = asst_msg("a1"), tool_msg("t1") a2 = asst_msg("a2") - p1, r1 = append(mgr, sid, [s, u], "a1", finish_reason="tool_calls", logprobs=[-0.5] * 2) - p2, r2 = append(mgr, sid, [s, u, a1, t1], "a2", finish_reason="tool_calls", logprobs=[-0.4] * 2) + append(mgr, sid, [s, u], "a1", finish_reason="tool_calls", logprobs=[-0.5] * 2) + append(mgr, sid, [s, u, a1, t1], "a2", finish_reason="tool_calls", logprobs=[-0.4] * 2) # three different tool results off a2 -> three leaves sharing a1+a2 for lbl in ["p", "q", "r"]: append(mgr, sid, [s, u, a1, t1, a2, tool_msg(lbl)], f"end-{lbl}", logprobs=[-0.3] * 2) samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) assert len(samples) == 3 - # Count trained tokens for r1 and r2 across all samples: each shared turn - # trained exactly once (first leaf), others loss=0. - # First leaf trains r1+r2+end; others train only their own end. - trained_first = sum(samples[0].loss_mask) - trained_rest = [sum(x.loss_mask) for x in samples[1:]] - assert trained_first > max(trained_rest), (trained_first, trained_rest) + # Leaf 0 OWNS the shared r:a1 + r:a2 (both trained) and its own end-p; leaves + # 1 and 2 SHARE r:a1 + r:a2 (bare, claimed by leaf 0) and train only their own + # end response. + assert goldens(samples) == [ + " system:S user:u [r:a1] [] " + " tool:t1 [r:a2] [] tool:p [r:end-p] []", + " system:S user:u r:a1 " + " tool:t1 r:a2 tool:q [r:end-q] []", + " system:S user:u r:a1 " + " tool:t1 r:a2 tool:r [r:end-r] []", + ] _check_invariants(samples) _record("3.7 deep multi-leaf dedup (3 leaves, shared trained once)", mgr, sid, samples) print("PASS 3.7") @@ -959,8 +1096,26 @@ def test_3_8_long_mixed_session(): p7 = drift(p7_honest, len(p1) - 1) append(mgr, sid, prefix, "a5", prompt_ids=p7, logprobs=lp) samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - # The final case-A fork splits the single leaf chain into >=2 segments. - assert len(samples) >= 2, len(samples) + # The final case-A fork splits the single leaf chain into 3 segments. + assert goldens(samples) == [ + # Segment 1: turns 1-4 in one clean chain. The turn-5 B1 replace dropped a + # drifted tail (the is gone here) and realigned, so r:a0..r:a3 all + # train. + " system:S user:u [r:a0] [] " + " tool:t0 [r:a1] [] tool:t1 [r:a2] [] " + " tool:t2 [r:a3] []", + # Segment 2: cross-leaf-style dedup within the chain — the prior turns are + # re-emitted as bare context and only r:a4 trains. + " system:S user:u r:a0 " + " tool:t0 r:a1 tool:t1 r:a2 " + " tool:t2 r:a3 tool:t3 [r:a4] []", + # Segment 3: turn 7 after the case-A fork (the in the early prompt + # region); whole prefix bare, only r:a5 trains. + " system:S user:u r:a0 " + " tool:t0 r:a1 tool:t1 r:a2 " + " tool:t2 r:a3 tool:t3 r:a4 " + " tool:t4 [r:a5] []", + ] assert abs(sum(s.reward for s in samples) - 1.0) < 1e-9 _check_invariants(samples) _record(f"3.8 long mixed session -> {len(samples)} samples", mgr, sid, samples) @@ -993,7 +1148,10 @@ def test_4_1_tools_metadata_on_first_system_only(): others = [n for n in _iter_all(mgr._trees[sid]) if n is not sys_node] assert all(n.metadata.get("tools") is None for n in others), "tools attached exactly once" samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 1 + assert goldens(samples) == [ + " system:S user:u [r:call] [] " + " tool:t [r:done] []", + ] _check_invariants(samples) _record("4.1 tools metadata on first system only", mgr, sid, samples) print("PASS 4.1") @@ -1055,6 +1213,7 @@ def test_4_4_default_base_sample(): samples = mgr.get_trajectory(sid, reward=1.0) # no base_sample assert len(samples) == 1 assert samples[0].index == 0 + assert goldens(samples) == [" system:S user:u [r:a] []"] _check_invariants(samples) _record("4.4 default base_sample (None)", mgr, sid, samples) print("PASS 4.4") @@ -1074,10 +1233,14 @@ def test_4_5_mixed_logprobs_across_turns(): assert len(samples) == 1 s0 = samples[0] L = _lcp_len(p1 + r1, p2) - # turn-1 region: real logprobs; prompt tail: 0.0; turn-2 region: 0.0 (padded). + # both responses still trained (golden shows the loss layout)... + assert goldens(samples) == [ + " system:S user:u [r:call] [] " + " tool:t [r:done] []", + ] + # ...but turn-2's region carries padded 0.0 logprobs (it had none), while + # turn-1's region keeps its real logprobs. assert s0.rollout_log_probs == [-0.5] * len(r1) + [0.0] * (len(p2) - L) + [0.0] * len(r2) - # both responses still trained. - assert s0.loss_mask == [1] * len(r1) + [0] * (len(p2) - L) + [1] * len(r2) _check_invariants(samples) _record("4.5 mixed logprobs across turns (turn2 padded 0.0)", mgr, sid, samples) print("PASS 4.5") @@ -1116,12 +1279,24 @@ def run(threshold, drift_tail_len): tools=None, response_message={"role": "assistant", "content": "done"}, ) - return get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + return samples, p1, r1, p2, r2 - forked = run(threshold=2, drift_tail_len=2) # d == threshold -> fork + # d == threshold -> fork: two single-turn segments, each trains its own resp. + forked, p1, r1, p2, r2 = run(threshold=2, drift_tail_len=2) assert len(forked) == 2, f"d==threshold must fork, got {len(forked)}" - replaced = run(threshold=2, drift_tail_len=1) # d < threshold -> replace + assert forked[0].tokens == p1 + r1 + assert forked[0].loss_mask == [1] * len(r1) + assert forked[1].tokens == p2 + r2 + assert forked[1].loss_mask == [1] * len(r2) + # d < threshold -> replace: one coherent segment realigned to p2. + replaced, p1b, r1b, p2b, r2b = run(threshold=2, drift_tail_len=1) assert len(replaced) == 1, f"d Date: Mon, 8 Jun 2026 16:11:04 +0000 Subject: [PATCH 23/28] refactor(agent): drift-tolerant trajectory linearization Rewrite TrajectoryManager.get_trajectory to tolerate TITO re-tokenization drift instead of raising. Divergence index L is classified by where it falls: prompt region -> fork; inside most-recent response span -> replace if drifted tail < threshold else fork; inside an earlier response span -> always fork. Add cross-leaf dedup so shared snapshot nodes train exactly once. Rename fork_merge_max_response_tokens -> fork_threshold_tokens across the adapters and example generate.py. --- examples/coding_agent_rl/generate.py | 2 +- slime/agent/adapters/anthropic.py | 6 +- slime/agent/adapters/openai.py | 6 +- slime/agent/trajectory_manager.py | 403 +++++++++++++++------------ 4 files changed, 239 insertions(+), 178 deletions(-) diff --git a/examples/coding_agent_rl/generate.py b/examples/coding_agent_rl/generate.py index a007bf40e9..7956f52600 100644 --- a/examples/coding_agent_rl/generate.py +++ b/examples/coding_agent_rl/generate.py @@ -115,7 +115,7 @@ def __init__(self, args) -> None: sglang_url=sglang_url, tool_parser=self.tool_parser, reasoning_parser=self.reasoning_parser, - fork_merge_max_response_tokens=fork_merge_threshold, + fork_threshold_tokens=fork_merge_threshold, ) # handler_cancellation=True so a client disconnect cancels the handler # coroutine, arming the fire-and-forget /abort_request inside the diff --git a/slime/agent/adapters/anthropic.py b/slime/agent/adapters/anthropic.py index bc563682d4..de6cbafafb 100644 --- a/slime/agent/adapters/anthropic.py +++ b/slime/agent/adapters/anthropic.py @@ -67,7 +67,7 @@ def __init__( tool_parser=None, reasoning_parser=None, max_turns_per_sid: int | None = None, - fork_merge_max_response_tokens: int | None = None, + fork_threshold_tokens: int | None = None, on_turn_appended: Callable[..., None] | None = None, ) -> None: super().__init__( @@ -80,8 +80,8 @@ def __init__( # ``None`` means "caller did not specify" -> let TrajectoryManager use # its own default for the assistant-rewrite merge threshold. mgr_kwargs: dict[str, int] = {} - if fork_merge_max_response_tokens is not None: - mgr_kwargs["fork_merge_max_response_tokens"] = fork_merge_max_response_tokens + if fork_threshold_tokens is not None: + mgr_kwargs["fork_threshold_tokens"] = fork_threshold_tokens self.manager = TrajectoryManager(**mgr_kwargs) # Optional debug hook invoked after each successful append_turn. # Signature: (sid, prompt_messages, tools, response_message, diff --git a/slime/agent/adapters/openai.py b/slime/agent/adapters/openai.py index f051908a8f..99f1c12f03 100644 --- a/slime/agent/adapters/openai.py +++ b/slime/agent/adapters/openai.py @@ -74,7 +74,7 @@ def __init__( tool_parser=None, reasoning_parser=None, max_turns_per_sid: int | None = None, - fork_merge_max_response_tokens: int | None = None, + fork_threshold_tokens: int | None = None, on_turn_appended: Callable[..., None] | None = None, ) -> None: super().__init__( @@ -87,8 +87,8 @@ def __init__( # ``None`` means "caller did not specify" -> let TrajectoryManager use # its own default for the assistant-rewrite merge threshold. mgr_kwargs: dict[str, int] = {} - if fork_merge_max_response_tokens is not None: - mgr_kwargs["fork_merge_max_response_tokens"] = fork_merge_max_response_tokens + if fork_threshold_tokens is not None: + mgr_kwargs["fork_threshold_tokens"] = fork_threshold_tokens self.manager = TrajectoryManager(**mgr_kwargs) # Optional debug hook invoked after each successful append_turn. # Signature mirrors AnthropicAdapter.on_turn_appended: diff --git a/slime/agent/trajectory_manager.py b/slime/agent/trajectory_manager.py index 4053883f67..b8ff2f34bb 100644 --- a/slime/agent/trajectory_manager.py +++ b/slime/agent/trajectory_manager.py @@ -14,38 +14,62 @@ / ``turn_finish_reason`` / ``turn_index``. Non-assistant nodes carry no token attribution at all. -* ``get_trajectory`` linearizes each leaf turn-by-turn with a STRICT - exact-prefix contract: walking the leaf's assistant chain root→leaf, the - cumulative ``(prompt + response)`` tokens emitted so far MUST be an exact - prefix of the next turn's ``turn_prompt_ids``. When it is, the new prompt - tail ``prompt[len(cumulative):]`` is appended as loss_mask=0, then the - turn's ``response`` is appended as loss_mask=1 with real logprobs. - -* When the prefix does NOT match, the upstream tokenization drifted (the - same history re-tokenized differently across turns). That is a bug to - surface, not to paper over: ``get_trajectory`` raises ``ValueError`` with - the sid, turn_index, common-prefix length, and drift size so the - offending turn is locatable. Note a drift introduced at an early turn can - surface several turns later — the prefix check catches it whenever the - re-rendered early region first diverges from the accumulated tokens. - - (History: an earlier design tolerated drift via LCP drop-and-replace plus - optional drift-fork / drop-accounting / fork-merge. That machinery is - removed here in favor of failing loudly; it can be re-added as an explicit - layer later if real drift turns out to be unavoidable.) - -* ONE tolerated exception to the strict contract: an assistant-rewrite merge. - cc sometimes re-renders a previously-recorded assistant message when feeding - it back as prompt (tool_call arg order, whitespace). The message no longer - matches, so DFS forks at that assistant — which does NOT raise (the rewrite - mounts as a routing-only node, skipped at linearization) but leaves the - original short turn as a standalone stub leaf -> its own Sample, diluting the - trajectory's evenly-split reward. ``_try_merge_assistant_rewrite`` absorbs - such a rewrite onto the existing leaf when its response is short enough - (``fork_merge_max_response_tokens``), demoting that node to routing-only so it - contributes 0 training tokens. Non-assistant mismatches, long responses, and - ambiguous cases are left to fork as usual; same-message-different-token drift - still raises. +* ``get_trajectory`` linearizes each leaf turn-by-turn. Walking the leaf's + assistant chain root→leaf, the cumulative ``(prompt + response)`` tokens + emitted so far are matched against the next turn's ``turn_prompt_ids``: + + - **Clean continuation** — the cumulative tokens are an exact prefix of the + turn's prompt. The new prompt tail ``prompt[len(cumulative):]`` is appended + as loss_mask=0, then the turn's ``response`` is appended as loss_mask=1 with + real logprobs. + + - **Drift** — the same history re-tokenized differently across turns (TITO + drift: tool_call arg order, whitespace, reasoning-block reordering). Rather + than raise, ``get_trajectory`` tolerates the drift by where the divergence + index ``L`` (the common-prefix length) falls, never letting logprobs + misalign with tokens: + + * **case A** — ``L`` lands in a prompt region (outside every recorded + response span): a genuine prompt-level re-render → **fork** (finalize + the current coherent segment as its own Sample, restart a fresh segment + at this turn). Fork discards nothing. + * **case B1** — ``L`` lands inside the most-recent response span (the + immediately-previous turn's response got re-rendered). Let + ``d = len(cumulative) - L`` be the drifted tail length: ``d < + fork_threshold`` → **replace** (truncate to ``L``, silently drop the + drifted tail, realign to this turn's prompt); ``d >= fork_threshold`` → + **fork**. + * **case B2** — ``L`` lands inside an *earlier* turn's response span. + Replacing would discard that turn's tail plus every later turn, so this + always **forks** regardless of drift size. + + A fork splits one leaf into >=2 Samples; reward is split evenly across all + emitted Samples (see ``get_trajectory``). + +* ONE tolerated exception at the ROUTING (tree) layer: an assistant-rewrite + merge. cc sometimes re-renders a previously-recorded assistant message when + feeding it back as prompt (tool_call arg order, whitespace). The message no + longer matches, so DFS forks at that assistant — leaving the original short + turn as a standalone stub leaf -> its own Sample, diluting the trajectory's + evenly-split reward. ``_try_merge_assistant_rewrite`` absorbs such a rewrite + onto the existing leaf when its response is short enough + (``fork_threshold_tokens``), demoting that node to routing-only so it + contributes 0 training tokens. This is the MESSAGE-level dual of case B1's + TOKEN-level replace: the rewrite-merge triggers when the message dict differs + (DFS would fork), case B1 triggers when the message is identical but its + tokens drift inside one chain. They live at different layers and handle + different causes, so both are kept. + +* Cross-leaf dedup at the LINEARIZATION layer: a snapshot assistant node can be + shared by >=2 sibling leaves (it is the SAME Node object on each chain, since + ``_find_mount_point`` reuses children). Linearizing every leaf from the root + would otherwise train that shared prefix once per leaf. Instead the first leaf + to reach a node (DFS / build order) trains its response (loss=1); later leaves + re-emit it as loss=0 context (``trained=False`` in ``_Segment.extend``), so the + shared prefix is trained exactly once. The tree is left intact (unlike the + rewrite-merge's node demotion); only the per-leaf loss signal is masked. A + leaf's terminal turn is its own freshly-created node, never pre-claimed, so + every leaf keeps >=1 trained turn. """ from __future__ import annotations @@ -81,22 +105,6 @@ class Node: turn_index: int 1-based, monotonic per session """ - __slots__ = ( - # routing - "role", - "messages", - "metadata", - "parent", - "children", - "match_key", - # per-turn snapshot (assistant leaves) - "turn_prompt_ids", - "turn_response_ids", - "turn_response_logprobs", - "turn_finish_reason", - "turn_index", - ) - def __init__( self, *, @@ -172,6 +180,111 @@ def _lcp_len(a: list[int], b: list[int]) -> int: return i +# =========================================================================== +# Segment — one coherent linearized run (a fork boundary closes it) +# =========================================================================== + + +class _Segment: + """Token accumulator for one coherent linearized run within a chain. + + Holds the running ``buffer`` with aligned ``loss`` / ``logprobs`` and per-turn + response ``spans`` (``[start, end)`` half-open). Invariant: ``absorb_turn`` + always ends by appending a response, so ``spans[-1]`` is the most-recent + response span and the buffer ends at its end. A fork closes the segment and a + fresh one opens at the diverging turn (see module docstring case A/B1/B2). + """ + + def __init__(self) -> None: + self.buffer: list[int] = [] + self.loss: list[int] = [] + self.logprobs: list[float] = [] + # Each span: (start, end, turn_index) over self.buffer, half-open. + self.spans: list[tuple[int, int, int | None]] = [] + self.first_prompt_len: int = 0 + + def measure_drift(self, prompt_ids: list[int]) -> int: + """Length of the segment's token tail this turn's ``prompt_ids`` failed to reproduce. + + Normally 0 (clean continuation). A positive value IS token-id drift — the + same history re-tokenized differently this turn (TITO drift) — not an error. + """ + return len(self.buffer) - _lcp_len(self.buffer, prompt_ids) + + def can_absorb_drift(self, drift: int, fork_threshold: int) -> bool: + """Whether a ``drift``-token re-tokenization can be realigned into this segment. + + Realignable only when the drift is confined to the most-recent response + span (case B1) and shorter than ``fork_threshold``; otherwise the caller + forks. See module docstring for case A/B1/B2. + """ + if drift == 0: + return True + realign_at = len(self.buffer) - drift + if not self.spans or realign_at < self.spans[-1][0]: + return False # case A (prompt region) or B2 (earlier response span) + return fork_threshold > 0 and drift < fork_threshold + + def absorb_turn( + self, + drift: int, + prompt_ids: list[int], + response_ids: list[int], + response_logprobs: list[float] | None, + turn_index: int | None, + *, + trained: bool = True, + ) -> None: + """Append one turn, dropping any re-tokenization drift first so logprobs stay aligned. + + Drop the last ``drift`` tokens so the buffer re-anchors on the prefix this + turn's ``prompt_ids`` reproduced (shrinking the prior span if cut), then + append this turn's prompt tail (loss=0) and response (loss=1). A truncated + prior span stays loss=1: that region is both the prior turn's response and + this turn's prompt context. + + ``trained=False`` appends the response as loss=0 / logprob=0.0 instead — the + node is already owned by an earlier sibling leaf, so re-training it would + double-count the shared prefix (see ``_chain_to_sample`` claim-on-first-visit). + """ + realign_at = len(self.buffer) - drift + del self.buffer[realign_at:] + del self.loss[realign_at:] + del self.logprobs[realign_at:] + if self.spans and realign_at < self.spans[-1][1]: + s, _e, j = self.spans[-1] + self.spans[-1] = (s, realign_at, j) # may collapse to empty (s == realign_at); harmless + + is_first_turn = not self.spans + tail = prompt_ids[realign_at:] + self.buffer.extend(tail) + self.loss.extend([0] * len(tail)) + self.logprobs.extend([0.0] * len(tail)) + + start = len(self.buffer) + self.buffer.extend(response_ids) + self.loss.extend([1 if trained else 0] * len(response_ids)) + self.logprobs.extend( + response_logprobs if (trained and response_logprobs is not None) else [0.0] * len(response_ids) + ) + self.spans.append((start, len(self.buffer), turn_index)) + + if is_first_turn: + self.first_prompt_len = len(prompt_ids) # stripped at build time + + def response_strip(self) -> int: + """Start index of the response region (the leading first-turn prompt prefix).""" + return min(self.first_prompt_len, len(self.loss)) + + def has_trained_response(self) -> bool: + """Whether the response region carries any loss=1 token. + + False only when every turn in the segment was claimed by an earlier + sibling leaf (cross-leaf dedup) -> no training signal to emit. + """ + return any(self.loss[self.response_strip() :]) + + # =========================================================================== # TrajectoryManager # =========================================================================== @@ -185,16 +298,11 @@ class TrajectoryManager: ancestor) + exactly 1 assistant leaf carrying that turn's sglang snapshot. """ - def __init__(self, *, fork_merge_max_response_tokens: int | None = None) -> None: - # Drift fork/replace threshold (see module docstring + spec). Only a - # case-B1 drift (divergence inside the immediately-previous turn's - # response region) compares its drift length against this value: - # drift < threshold -> replace (truncate + realign, dropped tail - # counted in tito_dropped_*); drift >= threshold -> fork. case A - # (prompt region) and case B2 (drift in an earlier turn's response) - # always fork regardless. <=0 forces B1 to fork too (max fidelity). - # ``None`` from the caller means "use the default". - self._fork_threshold: int = 1024 if fork_merge_max_response_tokens is None else fork_merge_max_response_tokens + def __init__(self, *, fork_threshold_tokens: int | None = None) -> None: + # Drift fork/replace threshold for case-B1 (see module docstring case + # A/B1/B2). <=0 forces every B1 to fork (max fidelity); ``None`` means + # "use the default". + self._fork_threshold: int = 1024 if fork_threshold_tokens is None else fork_threshold_tokens self._trees: dict[str, Node] = {} self._turn_count: dict[str, int] = {} @@ -256,11 +364,17 @@ def get_trajectory( return [] samples: list[Sample] = [] + # Cross-leaf dedup (see module docstring): a snapshot node shared by sibling + # leaves is trained only by the first leaf to reach it. ``claimed`` carries + # node identity (id()) across leaves to enforce that. + claimed: set[int] = set() for routing_leaf in root.leaves(): if routing_leaf.is_root: continue chain = routing_leaf.path_from_root() - samples.extend(self._chain_to_sample(sid, chain, base_sample=base_sample, extra_metadata=extra_metadata)) + samples.extend( + self._chain_to_sample(chain, base_sample=base_sample, extra_metadata=extra_metadata, claimed=claimed) + ) # Reward is split evenly across every emitted sample (one per leaf); the # token-weighted reducer downstream then gives each loss token the @@ -308,22 +422,14 @@ def _try_merge_assistant_rewrite( ) -> tuple[Node, int]: """Absorb a short assistant-rewrite onto its existing node instead of forking. - cc sometimes re-renders a previously-recorded assistant message when - feeding it back as prompt (tool_call arg order, whitespace). That breaks - DFS at the assistant and forks a fresh subtree, leaving the original - short turn as a standalone stub leaf -> its own Sample, diluting the - trajectory's evenly-split reward. Forking does NOT raise (the rewritten - message mounts as a routing-only node and is already skipped at - linearization); this merge is purely a reward-hygiene / de-fragmentation - optimization. + See module docstring (rewrite-merge bullet) for the why. Purely a + reward-hygiene / de-fragmentation optimization — forking is already safe + (the rewrite mounts as a routing-only node, skipped at linearization). When the diverging message is an assistant and exactly one eligible - *short-response leaf* sibling exists, adopt the rewritten message onto - that node and DEMOTE it to routing-only (clear its turn snapshot), so it - contributes 0 training tokens -- handled by the existing - ``turn_prompt_ids is not None`` filter in ``_chain_to_sample`` (no change - needed there). Any other mismatch (non-assistant message, long response, - non-leaf or ambiguous candidates) is left to fork as usual. + *short-response leaf* sibling exists, adopt the rewritten message onto that + node and DEMOTE it to routing-only (clear its turn snapshot). Any other + mismatch (non-assistant, long response, non-leaf or ambiguous) forks as usual. """ if self._fork_threshold <= 0: return cur, i # feature off @@ -355,15 +461,13 @@ def _try_merge_assistant_rewrite( return cur, i sib = candidates[0] - # Observability breadcrumb only (NOT read by linearization). - sib.metadata["merged_rewrite"] = { + sib.metadata["merged_rewrite"] = { # observability breadcrumb only "abandoned_turn_index": sib.turn_index, "abandoned_response_tokens": len(sib.turn_response_ids or []), } - # Demote to routing-only: snapshot cleared -> skipped by the existing - # ``turn_prompt_ids is not None`` filter at linearization. Clearing - # turn_prompt_ids also prevents this node from being re-selected as a - # merge candidate on a later turn. + # Demote to routing-only: snapshot cleared -> skipped by the + # ``turn_prompt_ids is not None`` filter at linearization, and never + # re-selected as a merge candidate on a later turn. sib.turn_prompt_ids = None sib.turn_response_ids = None sib.turn_response_logprobs = None @@ -384,13 +488,10 @@ def _mount_prompt_messages( ) -> Node: """Attach each remaining prompt message as a routing node under ``cur``. - One node per message (``node.messages`` is a singleton list). Token - attribution happens at get_trajectory time, not here. The tools - metadata is placed only on the FIRST system node on the path — - ``_first_system_already_set(cur)`` walks ``cur → root`` looking for a - system ancestor that already carries it, and ``cur`` here is the - deepest node from descent (+ optional merge), so the walk sees every - ancestor that's already mounted. + One node per message; token attribution happens at get_trajectory time, not + here. The tools metadata is placed only on the FIRST system node on the path + (``_first_system_already_set`` walks ``cur → root``; ``cur`` is the deepest + mounted node, so the walk sees every already-mounted ancestor). """ for m in remaining_messages: role = m.get("role") @@ -425,96 +526,62 @@ def _attach_assistant_leaf( def _chain_to_sample( self, - sid: str, chain: list[Node], *, base_sample: Sample, extra_metadata: dict[str, Any] | None, + claimed: set[int], ) -> list[Sample]: - """Linearize one root→leaf chain into a single Sample (strict exact-prefix). + """Linearize one root→leaf chain into >=1 Samples (see module docstring). - Walk the chain's assistant nodes root→leaf, accumulating tokens. Each - turn's cumulative ``(prompt + response)`` so far MUST be an exact prefix - of the next turn's ``turn_prompt_ids``; otherwise the upstream - tokenization drifted and we raise (see module docstring). Reward is left - at 0.0 here and assigned by the caller. + Each turn either grows the current segment or forks a new one (case + A/B1-too-long/B2). ``claimed`` deduplicates snapshot nodes shared across + sibling leaves; a leaf's terminal node is never pre-claimed, so every leaf + keeps >=1 trained turn. Reward is left at 0.0 and assigned by the caller. """ # Only assistant leaves carrying this turn's sglang snapshot participate. # Routing assistant nodes mounted from prior-turn replay (turn_prompt_ids # is None) carry no token signal and are skipped. asst_chain = [n for n in chain if n.role == "assistant" and n.turn_prompt_ids is not None] - tokens: list[int] = [] - loss_mask: list[int] = [] - logprobs: list[float] = [] + segments: list[_Segment] = [] + seg = _Segment() for asst in asst_chain: - prompt = asst.turn_prompt_ids or [] - response = asst.turn_response_ids or [] - response_logprobs = asst.turn_response_logprobs - - # Strict prefix check: tokens accumulated so far must be an exact - # prefix of this turn's prompt. For the first turn `tokens` is empty - # so this trivially holds. - n = len(tokens) - if prompt[:n] != tokens: - self._raise_prefix_drift(sid, asst_chain, asst, tokens, prompt) - - new_prompt = prompt[n:] - tokens.extend(new_prompt) - loss_mask.extend([0] * len(new_prompt)) - logprobs.extend([0.0] * len(new_prompt)) - - tokens.extend(response) - loss_mask.extend([1] * len(response)) - logprobs.extend(response_logprobs if response_logprobs is not None else [0.0] * len(response)) - - first_prompt_len = len(asst_chain[0].turn_prompt_ids or []) if asst_chain else 0 + prompt_ids = asst.turn_prompt_ids or [] + drift = seg.measure_drift(prompt_ids) + if not seg.can_absorb_drift(drift, self._fork_threshold): + segments.append(seg) # fork: close this segment, start fresh here + seg = _Segment() + drift = 0 + trained = id(asst) not in claimed + claimed.add(id(asst)) + seg.absorb_turn( + drift, + prompt_ids, + asst.turn_response_ids or [], + asst.turn_response_logprobs, + asst.turn_index, + trained=trained, + ) + segments.append(seg) + + # Drop empty / fully-masked segments: an in-chain fork can isolate a run of + # turns all claimed by an earlier sibling leaf (no loss=1 token), which would + # trip the downstream "not fully masked" assert. A leaf's terminal turn is + # never pre-claimed, so its final segment always survives. return [ self._build_leaf_sample( base_sample=base_sample, extra_metadata=extra_metadata, - tokens=tokens, - loss_mask=loss_mask, - logprobs=logprobs, - first_prompt_len=first_prompt_len, + tokens=seg.buffer, + loss_mask=seg.loss, + logprobs=seg.logprobs, + strip=seg.response_strip(), ) + for seg in segments + if seg.buffer and seg.has_trained_response() ] - @staticmethod - def _raise_prefix_drift( - sid: str, - asst_chain: list[Node], - asst: Node, - tokens: list[int], - prompt: list[int], - ) -> None: - """Raise on a TITO drift: accumulated tokens are not a prefix of prompt. - - Reports the common-prefix length and drift size, plus which earlier - assistant turn's prompt region the divergence falls in — a drift - introduced at an early turn can surface only when a later turn - re-renders that early region differently. - """ - L = _lcp_len(tokens, prompt) - drift = len(tokens) - L - # Locate which turn's prompt region L lands in, so an early-turn drift - # that surfaces several turns later is still attributable. - drift_in_turn = None - for prior in asst_chain: - if prior is asst: - break - if L < len(prior.turn_prompt_ids or []): - drift_in_turn = prior.turn_index - break - raise ValueError( - f"get_trajectory(sid={sid} turn={asst.turn_index}): TITO drift — " - f"accumulated tokens are not a prefix of this turn's prompt " - f"(common_prefix_len={L}, drift={drift} tokens; divergence falls in " - f"turn {drift_in_turn}'s prompt region). The same history " - f"re-tokenized differently across turns; refusing to silently " - f"drop/realign." - ) - def _build_leaf_sample( self, *, @@ -523,26 +590,20 @@ def _build_leaf_sample( tokens: list[int], loss_mask: list[int], logprobs: list[float], - first_prompt_len: int, + strip: int, ) -> Sample: """Build one Sample from a linearized token segment. - ``loss_mask`` / ``logprobs`` are clamped to the response region (the - leading first-turn prompt prefix is stripped) per the slime contract: - ``response_length == len(loss_mask)`` and loss_mask/logprobs cover only - the response region. ``reward`` is left at 0.0; the caller assigns the - per-sample share. + ``loss_mask`` / ``logprobs`` are clamped to the response region (``strip`` + drops the leading first-turn prompt prefix) per the slime contract: + ``response_length == len(loss_mask)``, covering only the response region + (see backends/megatron_utils/data.py:139, ray/rollout.py:695). ``reward`` + is left at 0.0; the caller assigns the per-sample share. - Sample metadata carries only ``extra_metadata`` (empty on the - production path): the per-row dataset metadata and per-turn tool / - finish_reason snapshot are intentionally NOT propagated onto the leaf - Sample. Dump/analysis tooling reads those off the tree nodes instead. + Per-row dataset metadata and the per-turn tool / finish_reason snapshot are + intentionally NOT propagated here (dump/analysis tooling reads them off the + tree nodes); only ``extra_metadata`` rides along. """ - # Clamp loss_mask/logprobs to the response region: strip the leading - # first-turn prompt prefix so response_length == len(loss_mask), per the - # slime contract (see backends/megatron_utils/data.py:139, - # ray/rollout.py:695). - strip = min(first_prompt_len, len(loss_mask)) loss_resp, lp_resp = loss_mask[strip:], logprobs[strip:] metadata = dict(extra_metadata or {}) return Sample( From 054f89a521c61585fe1bc3954afd88048bb64abe Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Mon, 8 Jun 2026 16:22:12 +0000 Subject: [PATCH 24/28] chore(agent): untrack e2e test design doc and trajectory_manager tests Remove from version control while keeping local copies (git rm --cached): - docs/superpowers/specs/2026-06-08-trajectory-manager-e2e-tests-design.md - tests/test_agent/test_trajectory_manager.py - tests/test_agent/test_trajectory_manager_e2e.py --- ...-08-trajectory-manager-e2e-tests-design.md | 164 -- tests/test_agent/test_trajectory_manager.py | 1013 ------------ .../test_agent/test_trajectory_manager_e2e.py | 1424 ----------------- 3 files changed, 2601 deletions(-) delete mode 100644 docs/superpowers/specs/2026-06-08-trajectory-manager-e2e-tests-design.md delete mode 100644 tests/test_agent/test_trajectory_manager.py delete mode 100644 tests/test_agent/test_trajectory_manager_e2e.py diff --git a/docs/superpowers/specs/2026-06-08-trajectory-manager-e2e-tests-design.md b/docs/superpowers/specs/2026-06-08-trajectory-manager-e2e-tests-design.md deleted file mode 100644 index 0321ba7ab6..0000000000 --- a/docs/superpowers/specs/2026-06-08-trajectory-manager-e2e-tests-design.md +++ /dev/null @@ -1,164 +0,0 @@ -# TrajectoryManager 端到端测试脚本 — 设计文档 - -日期：2026-06-08 -状态：已批准设计，待实现 - -## 目标 - -写一个独立的端到端测试脚本，通过 `TrajectoryManager` 的两个公共接口 -`append_turn` 和 `get_trajectory`，从**数据结构角度**全面覆盖各种分叉情形： -prompt 分叉、assistant 分叉、token-ID 分叉，以及它们的组合。测试数据要 -**方便人类阅读**（语义化 token ID + 反查表），运行时既做严格 assertion， -又能打印可读的 tree / 线性化结果供人眼审查。 - -## 背景 - -`TrajectoryManager`（`slime/agent/trajectory_manager.py`）维护一棵 per-sid 的 -逐 message 路由树，并在 `get_trajectory` 时把每个 leaf 链线性化成 slime -`Sample`。它有两个正交的层： - -- **路由树层**（`append_turn`）：按 `(role, node_match_key)` 匹配，**只看 - message 身份**，与 token ID 无关。相同 message 前缀总落在同一路径。 -- **线性化层**（`get_trajectory`）：按 **token-ID 前缀**匹配累积 token，按 - 漂移位置走 case A / B1 / B2 路由（fork / replace），并做 cross-leaf dedup - 和 reward 均分。 - -现有 `tests/test_agent/test_trajectory_manager.py` 已有 28 个测试覆盖这些机制， -但用 `ord(char)` 映射 token、断言不易一眼看懂分叉点，且「两层同时发生」的 -组合格子覆盖较薄。本脚本是**独立新增**，与现有文件并存、互不依赖。 - -## 文件 - -`tests/test_agent/test_trajectory_manager_e2e.py` - -不改动 `trajectory_manager.py`；不依赖 sglang / 网络 / 真实 tokenizer；不碰 -现有 `test_trajectory_manager.py`。 - -## §1 基础设施 - -### 语义化 token 词表 - -token ID 用带语义段的小整数，看 assertion 一眼知道分叉在哪。每个 message -渲染成 `[START, ...body, END]`： - -``` -system : START=1000, END=1009, body 1001..1008 -user : START=2000, END=2009, body 2001..2008 -assistant: START=9000, END=9009, body 9001..9008 -tool : START=3000, END=3009, body 3001..3008 -gen-prompt 起手符 (add_generation_prompt): 9000 -漂移哨兵段: 7000..7099（不属于任何 role，dump 里一眼认出是人为漂移） -``` - -反查表 `TOKEN_NAMES: dict[int, str]` 把每个 ID 翻译成可读名（如 -`1000→""`、`2001→"u:compute"`、`7001→""`）。dump 打印时把 ID -序列翻译成可读字符串。 - -### 构造助手（薄封装，不引入 DSL） - -- `MsgTok`：给一个 message 分配固定的渲染 token 段，保证同一 message 在不同 - turn 渲染出相同 token（模拟干净 tokenizer）；漂移由测试显式注入。 -- `render_prompt(messages) -> list[int]`：拼成 token 序列（含 add_generation_prompt）。 -- `render_response(text) -> list[int]`：assistant 输出 token。 -- `turn(prompt_ids, response_ids, finish_reason, logprobs=None)`：构造 `TurnRecord`。 -- `drift(ids, at, sentinel=7001)`：在指定下标注入/替换哨兵 token，制造 - token-ID 漂移，返回新序列——让漂移点在测试代码里显式可见。 - -### 双态运行 - -- 每个 case 函数做严格 assertion。 -- `main()` 顺序跑所有 case；每个 case 跑完用 `_dump_helpers.dump_tree_txt` - 打印 tree，再打印线性化出的每个 Sample（token 翻译成可读名 + loss_mask - 对齐展示）。 -- 沿用现有文件的 `test_*` + `main()` 风格，可被 pytest 收集，也可直接 - `python -m` 跑供人眼审查。 - -## §2 Case 矩阵（按「层 × 分叉位置」组织） - -### 组 1 — 路由树层（断言 tree 形状） - -| # | Case | 分叉位置 | 预期树形 | -|---|------|---------|---------| -| 1.1 | 单 turn | 无 | system→user→assistant 一条链 | -| 1.2 | 干净多 turn（含 tool） | 无 | 一条链，每 message 一节点 | -| 1.3 | system 分叉 | system 不同 | root 下 2 子树 | -| 1.4 | user 分叉（共享 system） | user 不同 | system 共享，user 层 2 leaf | -| 1.5 | assistant message 分叉 | assistant 身份不同 | 共享 user，assistant 层 2 leaf | -| 1.6 | tool 分叉（同 assistant 不同 tool 结果） | tool 不同 | 共享 assistant，tool 层分叉 | -| 1.7 | token-only 漂移不分叉 | message 相同、prompt_ids 不同 | 树不分叉（DFS 忽略 token） | -| 1.8 | 多 tool message 逐节点挂载 | 一个 turn 多 tool | 每 tool 独立节点 | -| 1.9 | 跨 sid 隔离 | 不同 sid | 两棵独立树 | -| 1.10 | 空 response | 无 | assistant leaf messages=[]，turn_response_ids=[] | - -### 组 2 — 线性化层（断言 tokens/loss_mask/logprobs/reward） - -| # | Case | 触发 | 预期 | -|---|------|------|------| -| 2.1 | 单 turn 线性化 | — | tokens=p+r，loss 只覆盖 r | -| 2.2 | 干净多 turn 线性化 | LCP=cumulative | 1 Sample，prompt 尾 loss=0、resp loss=1 | -| 2.3 | drift case A（prompt 区漂移） | L 落在 prompt 区 | fork 成 2 Sample，不丢 token | -| 2.4 | drift case B1 短→replace | L 落在最近 resp 区、d0`（无全 mask 样本）。 -- token 期望值用构造助手拼出来，不手敲魔数。 - -### 打印格式（`main()` 时，每个 case 之后） - -``` -=== CASE 1.4 user 分叉（共享 system）=== -[tree] - -[samples] 2 个 - Sample#0 reward=1.0 resp_len=4 - tok : u:A 9001 - loss: 0 0 0 0 1 1 - Sample#1 ... -PASS 1.4 -``` - -token 与 loss_mask 上下对齐，漂移 token 显示成 ``，一眼看出分叉点 -和训练区。 - -### 结尾 - -`main()` 顺序跑全部 case，全 PASS 后打印 `ALL E2E CASES PASSED (N cases)`。 - -## 非目标（YAGNI） - -- 不引入 DSL / fluent builder。 -- 不改动 `trajectory_manager.py`。 -- 不依赖 sglang / 网络 / 真实 tokenizer。 -- 不修改现有 `test_trajectory_manager.py`。 diff --git a/tests/test_agent/test_trajectory_manager.py b/tests/test_agent/test_trajectory_manager.py deleted file mode 100644 index 3948f179a5..0000000000 --- a/tests/test_agent/test_trajectory_manager.py +++ /dev/null @@ -1,1013 +0,0 @@ -"""Unit tests for src_v2.trajectory_manager (Plan C: token-faithful). - -What we test: - (1) DFS merge only on (role, node_match_key): same prefix in messages - space always lands on the same path regardless of prompt_ids drift. - (2) get_trajectory linearization: turn 1 = full prompt + response; - turn k>=2 = strict exact-prefix append; tokens / loss_mask / - logprobs all stay in sync. - (3) TITO drift handling: when turn k+1.prompt diverges mid-stream from - cumulative tokens, get_trajectory RAISES (no silent drop/realign). -""" - -from __future__ import annotations - -import json - -from slime.agent.adapters.common import TurnRecord # noqa: E402 -from slime.agent.trajectory_manager import TrajectoryManager, _lcp_len, node_match_key # noqa: E402 -from slime.utils.types import Sample # noqa: E402 - - -def _turn(prompt_ids, response_ids, *, finish_reason, logprobs=None): - """Helper: build the TurnRecord the way call_sglang_generate would. - - ``logprobs=None`` maps to an empty ``output_log_probs`` (the dataclass - default) so the manager treats this turn as carrying no logprob signal. - Pass an explicit list to attach per-token logprobs. - """ - return TurnRecord( - prompt_ids=list(prompt_ids), - output_ids=list(response_ids), - finish_reason=finish_reason, - output_log_probs=list(logprobs) if logprobs is not None else [], - ) - - -# --------------------------------------------------------------------------- -# Helper-level tests -# --------------------------------------------------------------------------- - - -def test_node_match_key_is_dict_internal_sort_only(): - a = [{"role": "u", "content": "x"}] - b = [{"content": "x", "role": "u"}] - assert node_match_key(a) == node_match_key(b) - - c = [{"role": "u", "content": "x"}, {"role": "u", "content": "y"}] - d = [{"role": "u", "content": "y"}, {"role": "u", "content": "x"}] - assert node_match_key(c) != node_match_key(d) - - e = [{"role": "assistant", "tool_calls": [{"id": "1", "type": "function"}]}] - f = [{"role": "assistant", "tool_calls": [{"type": "function", "id": "1"}]}] - assert node_match_key(e) == node_match_key(f) - print("PASS test_node_match_key_is_dict_internal_sort_only") - - -def test_lcp_len(): - assert _lcp_len([], []) == 0 - assert _lcp_len([1, 2, 3], []) == 0 - assert _lcp_len([], [1, 2, 3]) == 0 - assert _lcp_len([1, 2, 3], [1, 2, 3]) == 3 - assert _lcp_len([1, 2, 3], [1, 2, 4]) == 2 - assert _lcp_len([1, 2, 3, 4, 5], [1, 2, 3]) == 3 - print("PASS test_lcp_len") - - -def test_manager_accepts_fork_threshold(): - # default 1024 when unspecified - m_default = TrajectoryManager() - assert m_default._fork_threshold == 1024 - # explicit value honored - m_explicit = TrajectoryManager(fork_merge_max_response_tokens=256) - assert m_explicit._fork_threshold == 256 - # None -> default - m_none = TrajectoryManager(fork_merge_max_response_tokens=None) - assert m_none._fork_threshold == 1024 - print("PASS test_manager_accepts_fork_threshold") - - -# --------------------------------------------------------------------------- -# Fake tokenizer (kept only as a shape-matching prompt/response generator; -# trajectory_manager doesn't invoke it under plan C). -# --------------------------------------------------------------------------- - - -class FakeTokenizer: - ROLE_START = {"system": 9001, "user": 9002, "assistant": 9003, "tool": 9004} - ROLE_END = {"system": 9101, "user": 9102, "assistant": 9103, "tool": 9104} - - def apply_chat_template(self, messages, *, tools=None, add_generation_prompt=False, **kwargs): - out: list[int] = [] - for m in messages: - role = m["role"] - content = m.get("content") or "" - if not isinstance(content, str): - content = json.dumps(content, ensure_ascii=False) - out.append(self.ROLE_START[role]) - out.extend(ord(c) for c in content) - out.append(self.ROLE_END[role]) - if add_generation_prompt: - out.append(self.ROLE_START["assistant"]) - return out - - -def _render_prompt(messages, tools=None, tokenizer=None): - tok = tokenizer or FakeTokenizer() - return tok.apply_chat_template(messages, tools=tools, add_generation_prompt=True) - - -def _render_response(content_str, tokenizer=None): - tok = tokenizer or FakeTokenizer() - return [ord(c) for c in content_str] + [tok.ROLE_END["assistant"]] - - -# --------------------------------------------------------------------------- -# Plan-C semantics tests -# --------------------------------------------------------------------------- - - -SYSTEM_MSG = "You are a python coding agent." -TOOLS_OPENAI = [ - { - "type": "function", - "function": { - "name": "run_python", - "description": "Run python code.", - "parameters": {"type": "object", "properties": {"code": {"type": "string"}}}, - }, - }, -] - - -def _three_turn_session(tok): - """3-turn linear session via append_turn. Returns mgr, sid, per-turn (p,r).""" - mgr = TrajectoryManager() - sid = "three-turn" - - sys_msg = {"role": "system", "content": SYSTEM_MSG} - user1 = {"role": "user", "content": "Compute 2+2."} - asst1 = {"role": "assistant", "content": "Computing."} - tool1 = {"role": "tool", "content": "4"} - asst2 = {"role": "assistant", "content": "Answer is 4."} - - p1 = _render_prompt([sys_msg, user1], tokenizer=tok) - r1 = _render_response("Computing.", tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p1, r1, finish_reason="tool_calls", logprobs=[-0.5] * len(r1)), - prompt_messages=[sys_msg, user1], - tools=TOOLS_OPENAI, - response_message=asst1, - ) - - p2 = _render_prompt([sys_msg, user1, asst1, tool1], tokenizer=tok) - r2 = _render_response("Answer is 4.", tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p2, r2, finish_reason="stop", logprobs=[-0.4] * len(r2)), - prompt_messages=[sys_msg, user1, asst1, tool1], - tools=TOOLS_OPENAI, - response_message=asst2, - ) - return mgr, sid, [(p1, r1), (p2, r2)] - - -def test_append_single_turn_shapes_tree(): - tok = FakeTokenizer() - mgr = TrajectoryManager() - sid = "single" - sys_msg = {"role": "system", "content": "S"} - user1 = {"role": "user", "content": "u"} - p = _render_prompt([sys_msg, user1], tokenizer=tok) - r = _render_response("a", tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p, r, finish_reason="stop"), - prompt_messages=[sys_msg, user1], - tools=None, - response_message={"role": "assistant", "content": "a"}, - ) - chain = list(mgr._trees[sid].leaves())[0].path_from_root() - roles = [n.role for n in chain] - assert roles == ["system", "user", "assistant"], roles - asst = chain[-1] - assert asst.turn_index == 1 - assert asst.turn_prompt_ids == p - assert asst.turn_response_ids == r - assert asst.turn_finish_reason == "stop" - print("PASS test_append_single_turn_shapes_tree") - - -def test_append_three_turn_chain_no_fork(): - """3-turn session with consistent prompts -> exactly 1 leaf.""" - tok = FakeTokenizer() - mgr, sid, _ = _three_turn_session(tok) - leaves = [leaf for leaf in mgr._trees[sid].leaves() if not leaf.is_root] - assert len(leaves) == 1 - chain = leaves[0].path_from_root() - roles = [n.role for n in chain] - assert roles == ["system", "user", "assistant", "tool", "assistant"], roles - assert mgr.turn_count(sid) == 2 - print("PASS test_append_three_turn_chain_no_fork") - - -def test_fork_on_text_diff(): - """Different user content under shared sys -> 2 leaves, sys shared.""" - tok = FakeTokenizer() - mgr = TrajectoryManager() - sid = "fork-text" - sys_msg = {"role": "system", "content": "S"} - - for content in ["uA", "uB"]: - user = {"role": "user", "content": content} - p = _render_prompt([sys_msg, user], tokenizer=tok) - r = _render_response(content[-1], tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p, r, finish_reason="stop"), - prompt_messages=[sys_msg, user], - tools=None, - response_message={"role": "assistant", "content": content[-1]}, - ) - - root = mgr._trees[sid] - assert len(root.children) == 1, "sys node must be shared" - sys_node = root.children[0] - assert len(sys_node.children) == 2, "user level must fork" - leaves = [leaf for leaf in root.leaves() if not leaf.is_root] - assert len(leaves) == 2 - print("PASS test_fork_on_text_diff") - - -def test_no_fork_on_token_only_diff(): - """Plan C: same text but tampered prompt_ids -> NO fork (DFS ignores tokens). - - This is the load-bearing behavior change vs the old prefix-match design. - """ - tok = FakeTokenizer() - mgr = TrajectoryManager() - sid = "tokens-diff-only" - sys_msg = {"role": "system", "content": "S"} - user1 = {"role": "user", "content": "u"} - pa = _render_prompt([sys_msg, user1], tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(pa, _render_response("a", tokenizer=tok), finish_reason="stop"), - prompt_messages=[sys_msg, user1], - tools=None, - response_message={"role": "assistant", "content": "a"}, - ) - tampered = list(pa) - tampered[1] = tampered[1] ^ 1 - mgr.append_turn( - sid, - turn=_turn(tampered, _render_response("b", tokenizer=tok), finish_reason="stop"), - prompt_messages=[sys_msg, user1], - tools=None, - response_message={"role": "assistant", "content": "b"}, - ) - root = mgr._trees[sid] - # Same (sys, user) path -> shared, but two different assistant turns - # produce two assistant leaves under the same user node. - assert len(root.children) == 1 - sys_node = root.children[0] - assert len(sys_node.children) == 1 - user_node = sys_node.children[0] - assert len(user_node.children) == 2, "two distinct assistant turns hang off shared user" - leaves = [leaf for leaf in root.leaves() if not leaf.is_root] - assert len(leaves) == 2 - print("PASS test_no_fork_on_token_only_diff") - - -def test_cross_sid_isolation(): - tok = FakeTokenizer() - mgr = TrajectoryManager() - sys_msg = {"role": "system", "content": "S"} - for sid, content in [("sid-a", "uA"), ("sid-b", "uB")]: - user = {"role": "user", "content": content} - p = _render_prompt([sys_msg, user], tokenizer=tok) - r = _render_response(content[-1], tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p, r, finish_reason="stop"), - prompt_messages=[sys_msg, user], - tools=None, - response_message={"role": "assistant", "content": content[-1]}, - ) - assert len(list(mgr._trees["sid-a"].leaves())) == 1 - assert len(list(mgr._trees["sid-b"].leaves())) == 1 - print("PASS test_cross_sid_isolation") - - -def test_role_tool_in_chain(): - tok = FakeTokenizer() - mgr = TrajectoryManager() - sid = "tool-chain" - sys_msg = {"role": "system", "content": "S"} - user1 = {"role": "user", "content": "u"} - asst1 = {"role": "assistant", "content": "a1"} - tool_a = {"role": "tool", "content": "tA"} - tool_b = {"role": "tool", "content": "tB"} - asst2 = {"role": "assistant", "content": "a2"} - - p1 = _render_prompt([sys_msg, user1], tokenizer=tok) - r1 = _render_response("a1", tokenizer=tok) - p2 = _render_prompt([sys_msg, user1, asst1, tool_a, tool_b], tokenizer=tok) - r2 = _render_response("a2", tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p1, r1, finish_reason="stop"), - prompt_messages=[sys_msg, user1], - tools=None, - response_message=asst1, - ) - mgr.append_turn( - sid, - turn=_turn(p2, r2, finish_reason="stop"), - prompt_messages=[sys_msg, user1, asst1, tool_a, tool_b], - tools=None, - response_message=asst2, - ) - - chain = list(mgr._trees[sid].leaves())[0].path_from_root() - roles = [n.role for n in chain] - # Per-message routing: the two tool_results mount as two separate nodes. - assert roles == ["system", "user", "assistant", "tool", "tool", "assistant"], roles - assert chain[3].messages == [tool_a] - assert chain[4].messages == [tool_b] - print("PASS test_role_tool_in_chain") - - -def test_response_logprobs_length_mismatch_raises(): - tok = FakeTokenizer() - mgr = TrajectoryManager() - sys_msg = {"role": "system", "content": "S"} - user1 = {"role": "user", "content": "u"} - p = _render_prompt([sys_msg, user1], tokenizer=tok) - bad_turn = TurnRecord( - prompt_ids=p, - output_ids=[1, 2, 3], - finish_reason="stop", - output_log_probs=[-0.1, -0.2], - ) - try: - mgr.append_turn( - "x", - turn=bad_turn, - prompt_messages=[sys_msg, user1], - tools=None, - response_message={"role": "assistant", "content": ""}, - ) - except ValueError as e: - assert "output_log_probs" in str(e) - print("PASS test_response_logprobs_length_mismatch_raises") - return - raise AssertionError("expected ValueError") - - -def test_response_ids_empty_ok(): - tok = FakeTokenizer() - mgr = TrajectoryManager() - sys_msg = {"role": "system", "content": "S"} - user1 = {"role": "user", "content": "u"} - p = _render_prompt([sys_msg, user1], tokenizer=tok) - mgr.append_turn( - "x", - turn=_turn(p, [], finish_reason="stop"), - prompt_messages=[sys_msg, user1], - tools=None, - response_message=None, - ) - chain = list(mgr._trees["x"].leaves())[0].path_from_root() - asst = chain[-1] - assert asst.role == "assistant" - assert asst.turn_response_ids == [] - assert asst.turn_prompt_ids == p - assert asst.messages == [] - print("PASS test_response_ids_empty_ok") - - -# --------------------------------------------------------------------------- -# get_trajectory linearization (Plan C heart of the matter) -# --------------------------------------------------------------------------- - - -def test_get_trajectory_single_turn(): - tok = FakeTokenizer() - mgr = TrajectoryManager() - sid = "g1" - sys_msg = {"role": "system", "content": "S"} - user = {"role": "user", "content": "u"} - p = _render_prompt([sys_msg, user], tokenizer=tok) - r = _render_response("a", tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p, r, finish_reason="stop", logprobs=[-0.5] * len(r)), - prompt_messages=[sys_msg, user], - tools=TOOLS_OPENAI, - response_message={"role": "assistant", "content": "a"}, - ) - samples = mgr.get_trajectory(sid, base_sample=Sample(index=7, prompt="hi"), reward=1.0) - assert len(samples) == 1 - s = samples[0] - assert s.tokens == p + r - # slime contract: loss_mask / rollout_log_probs cover only the response - # region (tokens after the initial prompt) and len(loss_mask) == - # response_length. The first turn's prompt prefix is stripped. - assert s.loss_mask == [1] * len(r) - assert s.rollout_log_probs == [-0.5] * len(r) - assert s.response_length == len(r) - assert s.reward == 1.0 - # Leaf Sample metadata is intentionally empty: no dataset-row passthrough, - # no per-turn tools / finish_reason snapshot (those live on the tree nodes). - assert s.metadata == {} - print("PASS test_get_trajectory_single_turn") - - -def test_get_trajectory_clean_multiturn(): - """Clean 2-turn session (no drift) linearizes as turn1 prompt+resp then - turn2 (prompt - LCP) + resp, with full coherent loss_mask / logprobs.""" - tok = FakeTokenizer() - mgr, sid, turns = _three_turn_session(tok) - samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 1 - s = samples[0] - - (p1, r1), (p2, r2) = turns - # LCP(p1+r1, p2) should equal len(p1)+len(r1) for our clean fake tokenizer - # (p2 starts exactly with p1 contents + asst response + tool block + new gen prompt) - L = _lcp_len(p1 + r1, p2) - assert L == len(p1) + len(r1), f"clean session LCP should equal cumulative, got {L}" - expected_tokens = p1 + r1 + p2[L:] + r2 - # loss_mask / rollout_log_probs are response-only (slime contract): - # turn1 prompt is stripped; turn1 response keeps mask=1, then the - # extra prompt slice (p2[L:]) is mask=0, then turn2 response is mask=1. - expected_loss = [1] * len(r1) + [0] * (len(p2) - L) + [1] * len(r2) - expected_logp = [-0.5] * len(r1) + [0.0] * (len(p2) - L) + [-0.4] * len(r2) - assert s.tokens == expected_tokens - assert s.loss_mask == expected_loss - assert s.rollout_log_probs == expected_logp - assert s.response_length == len(r1) + (len(p2) - L) + len(r2) - assert s.metadata == {} - print("PASS test_get_trajectory_clean_multiturn") - - -def test_get_trajectory_tito_drift_raises(): - """Strict exact-prefix: turn 2 prompt diverges mid-stream from cumulative - tokens. The accumulated turn-1 (prompt + response) is no longer a prefix of - turn 2's prompt, so get_trajectory must RAISE instead of dropping/realigning. - """ - tok = FakeTokenizer() - mgr = TrajectoryManager() - sid = "tito" - sys_msg = {"role": "system", "content": "S"} - user = {"role": "user", "content": "u"} - asst1 = {"role": "assistant", "content": "a1"} - tool = {"role": "tool", "content": "t"} - asst2 = {"role": "assistant", "content": "a2"} - - p1 = _render_prompt([sys_msg, user], tokenizer=tok) - r1 = _render_response("a1", tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p1, r1, finish_reason="tool_calls", logprobs=[-0.5] * len(r1)), - prompt_messages=[sys_msg, user], - tools=None, - response_message=asst1, - ) - # Build turn 2 prompt the "honest" way, then INJECT a synthetic divergence - # inside the assistant response region — simulating chat-template drift. We - # splice 3 fake tokens past the LCP so cumulative (p1+r1) is no longer a - # prefix of p2. - p2_honest = _render_prompt([sys_msg, user, asst1, tool], tokenizer=tok) - drift_at = len(p1) + 1 # inside r1 - p2 = list(p2_honest) - p2 = p2[:drift_at] + [77777, 77778, 77779] + p2[drift_at:] - r2 = _render_response("a2", tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p2, r2, finish_reason="stop", logprobs=[-0.4] * len(r2)), - prompt_messages=[sys_msg, user, asst1, tool], - tools=None, - response_message=asst2, - ) - - try: - mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - except ValueError as e: - msg = str(e) - assert "TITO drift" in msg, msg - assert f"turn={2}" in msg, msg # the turn whose prompt failed the check - assert "common_prefix_len" in msg, msg - print("PASS test_get_trajectory_tito_drift_raises") - return - raise AssertionError("expected ValueError on TITO drift") - - -def test_get_trajectory_tito_drift_late_surfacing_attributes_early_turn(): - """A drift introduced in an early turn's region but only re-rendered at a - later turn must still raise, and the message should attribute the divergence - to the early turn's prompt region (not just the failing turn). - """ - tok = FakeTokenizer() - mgr = TrajectoryManager() - sid = "tito-late" - sys_msg = {"role": "system", "content": "S"} - user = {"role": "user", "content": "u"} - asst1 = {"role": "assistant", "content": "a1"} - tool1 = {"role": "tool", "content": "t1"} - asst2 = {"role": "assistant", "content": "a2"} - tool2 = {"role": "tool", "content": "t2"} - asst3 = {"role": "assistant", "content": "a3"} - - # turn 1 + turn 2: clean continuation. - p1 = _render_prompt([sys_msg, user], tokenizer=tok) - r1 = _render_response("a1", tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p1, r1, finish_reason="tool_calls"), - prompt_messages=[sys_msg, user], - tools=None, - response_message=asst1, - ) - p2 = _render_prompt([sys_msg, user, asst1, tool1], tokenizer=tok) - r2 = _render_response("a2", tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p2, r2, finish_reason="tool_calls"), - prompt_messages=[sys_msg, user, asst1, tool1], - tools=None, - response_message=asst2, - ) - # turn 3: re-render the turn-1 region differently (splice a phantom token - # inside p1's range). Cumulative now diverges from p3 deep inside turn 1. - p3_honest = _render_prompt([sys_msg, user, asst1, tool1, asst2, tool2], tokenizer=tok) - drift_at = len(p1) - 1 # inside turn 1's prompt region - p3 = list(p3_honest) - p3 = p3[:drift_at] + [99999] + p3[drift_at:] - r3 = _render_response("a3", tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p3, r3, finish_reason="stop"), - prompt_messages=[sys_msg, user, asst1, tool1, asst2, tool2], - tools=None, - response_message=asst3, - ) - - try: - mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt="")) - except ValueError as e: - msg = str(e) - assert "TITO drift" in msg, msg - # The failing turn is turn 3, but the divergence falls in turn 1's region. - assert "turn=3" in msg, msg - assert "turn 1's prompt region" in msg, msg - print("PASS test_get_trajectory_tito_drift_late_surfacing_attributes_early_turn") - return - raise AssertionError("expected ValueError on late-surfacing TITO drift") - - -def test_get_trajectory_two_leaves_share_reward(): - """Forked tree (2 leaves) -> reward split evenly.""" - tok = FakeTokenizer() - mgr = TrajectoryManager() - sid = "split" - sys_msg = {"role": "system", "content": "S"} - for content in ["uA", "uB"]: - user = {"role": "user", "content": content} - p = _render_prompt([sys_msg, user], tokenizer=tok) - r = _render_response(content[-1], tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p, r, finish_reason="stop"), - prompt_messages=[sys_msg, user], - tools=None, - response_message={"role": "assistant", "content": content[-1]}, - ) - samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=2.0) - assert len(samples) == 2 - assert all(s.reward == 1.0 for s in samples) - print("PASS test_get_trajectory_two_leaves_share_reward") - - -def test_drop_clears_sid(): - tok = FakeTokenizer() - mgr, sid, _ = _three_turn_session(tok) - assert mgr.has_session(sid) - mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt="")) - assert not mgr.has_session(sid) - assert mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt="")) == [] - print("PASS test_drop_clears_sid") - - -def test_debug_dump_shape(): - from tests.test_agent.test_claude_code_agent._dump_helpers import dump_tree_json, dump_tree_txt - - tok = FakeTokenizer() - mgr, sid, _ = _three_turn_session(tok) - txt = dump_tree_txt(mgr, sid) - assert isinstance(txt, str) and txt - for needle in ("session=", "[system]", "[user]", "[assistant]", "[tool]", "turns=2"): - assert needle in txt, f"missing {needle!r}" - # Plan C: assistant rows show turn= / prompt_ids= / response_ids= - assert "turn=1" in txt - assert "turn=2" in txt - assert "prompt_ids=" in txt - assert "response_ids=" in txt - - j = dump_tree_json(mgr, sid) - assert j["found"] is True and j["sid"] == sid and j["turns"] == 2 - assert j["nodes_total"] == 5 - - miss = dump_tree_txt(mgr, "no-such") - assert miss == "" - miss_j = dump_tree_json(mgr, "no-such") - assert miss_j == {"sid": "no-such", "found": False} - print("PASS test_debug_dump_shape") - - -def test_get_trajectory_skips_routing_assistant_in_drift_loop(): - """cc replays a foreign assistant the manager never recorded as a leaf; - per-message routing mounts it as a routing-only assistant that must be - filtered out of the strict-prefix walk. - - With per-message routing each prompt message mounts its own node, so the - previously-recorded ``asst2`` leaf (a single-message node) still matches - by ``(role, node_match_key)`` and the chain descends through it — turn 2 - stays on the main path and re-enters the strict-prefix check. Only the - extra ``foreign`` assistant, which the manager never saw via append_turn, - has no matching leaf and mounts as a routing-only assistant - (``turn_prompt_ids`` / ``turn_index`` both None). - - Such routing assistants must be filtered out of ``asst_chain`` (they - carry no per-turn snapshot). If one leaked into the strict prefix walk it - would look like a turn with an empty prompt and trip the exact-prefix - check — raising spuriously. This guards that the filter keeps the - otherwise-clean trajectory from raising, while turns 1/2/3 all stay in - the single linearized chain. - - Regression for the 20260604-120030 batch where 6 instances surfaced - routing assistants at depth 23-39, exact pattern: cc replays a prior - assistant message that the manager never recorded as its own leaf. - """ - tok = FakeTokenizer() - mgr = TrajectoryManager() - sid = "routing-asst" - sys_msg = {"role": "system", "content": "S"} - user1 = {"role": "user", "content": "u1"} - asst1 = {"role": "assistant", "content": "real-a1"} - tool1 = {"role": "tool", "content": "t1"} - asst2 = {"role": "assistant", "content": "real-a2"} - - # turn 1 - p1 = _render_prompt([sys_msg, user1], tokenizer=tok) - r1 = _render_response("real-a1", tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p1, r1, finish_reason="tool_calls"), - prompt_messages=[sys_msg, user1], - tools=None, - response_message=asst1, - ) - # turn 2 — clean continuation - p2 = _render_prompt([sys_msg, user1, asst1, tool1], tokenizer=tok) - r2 = _render_response("real-a2", tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p2, r2, finish_reason="tool_calls"), - prompt_messages=[sys_msg, user1, asst1, tool1], - tools=None, - response_message=asst2, - ) - # turn 3 — cc replays an extra prior asst that manager never saw via - # add_turn. Per-message routing matches asst2's own leaf and descends, so - # only the unmatched `foreign` message mounts as a routing-only assistant. - foreign = {"role": "assistant", "content": "foreign-msg"} - tool2 = {"role": "tool", "content": "t2"} - asst3 = {"role": "assistant", "content": "real-a3"} - p3 = _render_prompt([sys_msg, user1, asst1, tool1, asst2, foreign, tool2], tokenizer=tok) - r3 = _render_response("real-a3", tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p3, r3, finish_reason="stop"), - prompt_messages=[sys_msg, user1, asst1, tool1, asst2, foreign, tool2], - tools=None, - response_message=asst3, - ) - - # Per-message routing keeps everything on one path: asst2 matched, so no - # spurious fork. - leaves = [leaf for leaf in mgr._trees[sid].leaves() if not leaf.is_root] - assert len(leaves) == 1, [n.messages for n in leaves] - leaf3 = leaves[0] - assert leaf3.messages[0].get("content") == "real-a3" - chain = leaf3.path_from_root() - asst_nodes = [n for n in chain if n.role == "assistant"] - - # Exactly one routing assistant (the foreign replay); turns 1/2/3 keep - # their snapshots and stay in asst_chain. - routing_asst = [n for n in asst_nodes if n.turn_prompt_ids is None] - assert len(routing_asst) == 1, ( - f"expected exactly one routing assistant; got " - f"{[(n.turn_index, n.turn_prompt_ids is not None) for n in asst_nodes]}" - ) - assert routing_asst[0].turn_index is None - assert routing_asst[0].messages[0].get("content") == "foreign-msg" - snapshot_turns = [n.turn_index for n in asst_nodes if n.turn_prompt_ids is not None] - assert snapshot_turns == [1, 2, 3], snapshot_turns - - # The routing assistant is filtered out, so the strict prefix walk sees a - # clean chain (turns 1/2/3) and does NOT raise. - samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt="")) - assert len(samples) == 1 - print("PASS test_get_trajectory_skips_routing_assistant_in_drift_loop") - - -# --------------------------------------------------------------------------- -# Assistant-rewrite merge (single tolerated exception to strict exact-prefix) -# --------------------------------------------------------------------------- - - -def test_rewrite_merge_absorbs_short_assistant(): - """cc re-renders a short prior assistant; the manager absorbs the rewrite - onto the existing leaf (demoted to routing-only) instead of forking a - reward-diluting stub. One leaf, one Sample, original response not trained. - """ - tok = FakeTokenizer() - mgr = TrajectoryManager() # default threshold 1024 -> merge ON - sid = "rw-merge" - sys_msg = {"role": "system", "content": "S"} - user1 = {"role": "user", "content": "u"} - asst1 = {"role": "assistant", "content": "ok"} # short raw output - asst1_rw = {"role": "assistant", "content": "ok "} # cc-rewritten (whitespace) - tool1 = {"role": "tool", "content": "t"} - asst2 = {"role": "assistant", "content": "done"} - - p1 = _render_prompt([sys_msg, user1], tokenizer=tok) - r1 = _render_response("ok", tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p1, r1, finish_reason="tool_calls", logprobs=[-0.5] * len(r1)), - prompt_messages=[sys_msg, user1], - tools=None, - response_message=asst1, - ) - p2 = _render_prompt([sys_msg, user1, asst1_rw, tool1], tokenizer=tok) - r2 = _render_response("done", tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p2, r2, finish_reason="stop", logprobs=[-0.4] * len(r2)), - prompt_messages=[sys_msg, user1, asst1_rw, tool1], - tools=None, - response_message=asst2, - ) - - # Single chain (no fork): the rewrite was absorbed. - leaves = [leaf for leaf in mgr._trees[sid].leaves() if not leaf.is_root] - assert len(leaves) == 1, [n.messages for n in leaves] - chain = leaves[0].path_from_root() - assert [n.role for n in chain] == ["system", "user", "assistant", "tool", "assistant"] - - merged = chain[2] - # Demoted to routing-only: snapshot cleared, adopted the rewritten message. - assert merged.turn_prompt_ids is None - assert merged.turn_index is None - assert merged.messages == [asst1_rw] - assert merged.metadata["merged_rewrite"]["abandoned_turn_index"] == 1 - assert merged.metadata["merged_rewrite"]["abandoned_response_tokens"] == len(r1) - - # Linearization: only turn 2 participates; the abandoned turn-1 response is - # NOT trained. tokens == p2 + r2; loss covers only r2. - samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 1 - s = samples[0] - assert s.tokens == p2 + r2 - assert s.loss_mask == [1] * len(r2) - assert s.rollout_log_probs == [-0.4] * len(r2) - assert s.reward == 1.0 - print("PASS test_rewrite_merge_absorbs_short_assistant") - - -def test_rewrite_merge_skips_long_assistant(): - """A long abandoned response (>= threshold) is NOT absorbed: it forks into - its own Sample (carrying enough real signal to train standalone). Forking - does not raise. - """ - tok = FakeTokenizer() - mgr = TrajectoryManager(fork_merge_max_response_tokens=2) # r1 (3 tok) >= 2 - sid = "rw-long" - sys_msg = {"role": "system", "content": "S"} - user1 = {"role": "user", "content": "u"} - asst1 = {"role": "assistant", "content": "ok"} - asst1_rw = {"role": "assistant", "content": "ok "} - tool1 = {"role": "tool", "content": "t"} - asst2 = {"role": "assistant", "content": "done"} - - p1 = _render_prompt([sys_msg, user1], tokenizer=tok) - r1 = _render_response("ok", tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p1, r1, finish_reason="tool_calls"), - prompt_messages=[sys_msg, user1], - tools=None, - response_message=asst1, - ) - p2 = _render_prompt([sys_msg, user1, asst1_rw, tool1], tokenizer=tok) - r2 = _render_response("done", tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p2, r2, finish_reason="stop"), - prompt_messages=[sys_msg, user1, asst1_rw, tool1], - tools=None, - response_message=asst2, - ) - - leaves = [leaf for leaf in mgr._trees[sid].leaves() if not leaf.is_root] - assert len(leaves) == 2, "long rewrite must fork, not merge" - samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt=""), reward=2.0) - assert len(samples) == 2 # no raise - print("PASS test_rewrite_merge_skips_long_assistant") - - -def test_rewrite_merge_disabled_by_zero_threshold(): - """fork_merge_max_response_tokens=0 disables merge: every rewrite forks.""" - tok = FakeTokenizer() - mgr = TrajectoryManager(fork_merge_max_response_tokens=0) - sid = "rw-off" - sys_msg = {"role": "system", "content": "S"} - user1 = {"role": "user", "content": "u"} - asst1 = {"role": "assistant", "content": "ok"} - asst1_rw = {"role": "assistant", "content": "ok "} - tool1 = {"role": "tool", "content": "t"} - - p1 = _render_prompt([sys_msg, user1], tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p1, _render_response("ok", tokenizer=tok), finish_reason="tool_calls"), - prompt_messages=[sys_msg, user1], - tools=None, - response_message=asst1, - ) - p2 = _render_prompt([sys_msg, user1, asst1_rw, tool1], tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p2, _render_response("done", tokenizer=tok), finish_reason="stop"), - prompt_messages=[sys_msg, user1, asst1_rw, tool1], - tools=None, - response_message={"role": "assistant", "content": "done"}, - ) - leaves = [leaf for leaf in mgr._trees[sid].leaves() if not leaf.is_root] - assert len(leaves) == 2, "merge disabled -> rewrite forks" - print("PASS test_rewrite_merge_disabled_by_zero_threshold") - - -def test_rewrite_merge_ambiguous_candidates_fork(): - """Two eligible short-leaf assistant siblings -> ambiguous -> fork (no - arbitrary merge). - """ - tok = FakeTokenizer() - mgr = TrajectoryManager() - sid = "rw-ambig" - sys_msg = {"role": "system", "content": "S"} - user1 = {"role": "user", "content": "u"} - - # Two turns sharing the (sys, user) prefix produce two assistant leaves. - for content in ["a", "b"]: - p = _render_prompt([sys_msg, user1], tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p, _render_response(content, tokenizer=tok), finish_reason="stop"), - prompt_messages=[sys_msg, user1], - tools=None, - response_message={"role": "assistant", "content": content}, - ) - user_node = mgr._trees[sid].children[0].children[0] - assert len(user_node.children) == 2 - - # A third turn rewrites at the assistant slot -> two merge candidates. - asst_c = {"role": "assistant", "content": "c"} - tool1 = {"role": "tool", "content": "t"} - p3 = _render_prompt([sys_msg, user1, asst_c, tool1], tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p3, _render_response("d", tokenizer=tok), finish_reason="stop"), - prompt_messages=[sys_msg, user1, asst_c, tool1], - tools=None, - response_message={"role": "assistant", "content": "d"}, - ) - leaves = [leaf for leaf in mgr._trees[sid].leaves() if not leaf.is_root] - assert len(leaves) == 3, "ambiguous candidates must fork, not merge" - print("PASS test_rewrite_merge_ambiguous_candidates_fork") - - -def test_rewrite_merge_non_assistant_mismatch_forks(): - """A non-assistant divergence (different user) is left to fork; the merge - hook does not touch it even with merge enabled. - """ - tok = FakeTokenizer() - mgr = TrajectoryManager() - sid = "rw-nonasst" - sys_msg = {"role": "system", "content": "S"} - for content in ["uA", "uB"]: - user = {"role": "user", "content": content} - p = _render_prompt([sys_msg, user], tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p, _render_response(content[-1], tokenizer=tok), finish_reason="stop"), - prompt_messages=[sys_msg, user], - tools=None, - response_message={"role": "assistant", "content": content[-1]}, - ) - sys_node = mgr._trees[sid].children[0] - assert len(sys_node.children) == 2, "user-level divergence still forks" - print("PASS test_rewrite_merge_non_assistant_mismatch_forks") - - -def test_rewrite_merge_match_key_updated_so_next_turn_descends(): - """Regression: after merge, the node's cached match_key must follow the - adopted (rewritten) message so a LATER turn's DFS descends through it - instead of forking again. Also exercises clean strict-prefix continuation - across the merged node. - """ - tok = FakeTokenizer() - mgr = TrajectoryManager() - sid = "rw-matchkey" - sys_msg = {"role": "system", "content": "S"} - user1 = {"role": "user", "content": "u"} - asst1 = {"role": "assistant", "content": "ok"} - asst1_rw = {"role": "assistant", "content": "ok "} - tool1 = {"role": "tool", "content": "t1"} - asst2 = {"role": "assistant", "content": "second"} - tool2 = {"role": "tool", "content": "t2"} - asst3 = {"role": "assistant", "content": "third"} - - # turn 1: short assistant. - p1 = _render_prompt([sys_msg, user1], tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p1, _render_response("ok", tokenizer=tok), finish_reason="tool_calls"), - prompt_messages=[sys_msg, user1], - tools=None, - response_message=asst1, - ) - # turn 2: rewrite asst1 -> merge. - p2 = _render_prompt([sys_msg, user1, asst1_rw, tool1], tokenizer=tok) - r2 = _render_response("second", tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p2, r2, finish_reason="tool_calls", logprobs=[-0.4] * len(r2)), - prompt_messages=[sys_msg, user1, asst1_rw, tool1], - tools=None, - response_message=asst2, - ) - # turn 3: prompt carries the rewritten asst1_rw again; DFS must descend - # through the merged node (match_key updated) and not fork. - p3 = _render_prompt([sys_msg, user1, asst1_rw, tool1, asst2, tool2], tokenizer=tok) - r3 = _render_response("third", tokenizer=tok) - mgr.append_turn( - sid, - turn=_turn(p3, r3, finish_reason="stop", logprobs=[-0.3] * len(r3)), - prompt_messages=[sys_msg, user1, asst1_rw, tool1, asst2, tool2], - tools=None, - response_message=asst3, - ) - - leaves = [leaf for leaf in mgr._trees[sid].leaves() if not leaf.is_root] - assert len(leaves) == 1, ("match_key not updated -> spurious fork", [n.messages for n in leaves]) - - # Strict prefix holds across the merged node: turns 2 and 3 linearize into - # one clean Sample (no raise). The demoted turn-1 node is filtered out. - samples = mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt="")) - assert len(samples) == 1 - s = samples[0] - L = _lcp_len(p2 + r2, p3) - assert s.tokens == p2 + r2 + p3[L:] + r3 - print("PASS test_rewrite_merge_match_key_updated_so_next_turn_descends") - - -# --------------------------------------------------------------------------- -# main -# --------------------------------------------------------------------------- - - -def main() -> None: - test_node_match_key_is_dict_internal_sort_only() - test_lcp_len() - test_append_single_turn_shapes_tree() - test_append_three_turn_chain_no_fork() - test_fork_on_text_diff() - test_no_fork_on_token_only_diff() - test_cross_sid_isolation() - test_role_tool_in_chain() - test_response_logprobs_length_mismatch_raises() - test_response_ids_empty_ok() - test_get_trajectory_single_turn() - test_get_trajectory_clean_multiturn() - test_get_trajectory_tito_drift_raises() - test_get_trajectory_tito_drift_late_surfacing_attributes_early_turn() - test_get_trajectory_two_leaves_share_reward() - test_drop_clears_sid() - test_debug_dump_shape() - test_get_trajectory_skips_routing_assistant_in_drift_loop() - test_rewrite_merge_absorbs_short_assistant() - test_rewrite_merge_skips_long_assistant() - test_rewrite_merge_disabled_by_zero_threshold() - test_rewrite_merge_ambiguous_candidates_fork() - test_rewrite_merge_non_assistant_mismatch_forks() - test_rewrite_merge_match_key_updated_so_next_turn_descends() - print("\nALL PLAN-C TESTS PASSED.") - - -if __name__ == "__main__": - main() diff --git a/tests/test_agent/test_trajectory_manager_e2e.py b/tests/test_agent/test_trajectory_manager_e2e.py deleted file mode 100644 index 47d1c33296..0000000000 --- a/tests/test_agent/test_trajectory_manager_e2e.py +++ /dev/null @@ -1,1424 +0,0 @@ -"""End-to-end tests for TrajectoryManager via append_turn / get_trajectory. - -This script drives the two public interfaces of -``slime.agent.trajectory_manager.TrajectoryManager`` and exhaustively covers the -ways a trajectory can branch, organized as a two-axis matrix: - - * LAYER 1 — routing tree (append_turn). DFS merges on (role, node_match_key) - only, so MESSAGE IDENTITY决定 tree shape; token ids are irrelevant here. - * LAYER 2 — linearization (get_trajectory). TOKEN-ID prefix决定 how each leaf - chain becomes Samples (clean continuation / drift case A·B1·B2 / cross-leaf - dedup / reward split). - * COMBINED — both layers interacting (rewrite-merge, tree-fork + token-drift - stacked, deep multi-leaf dedup, long mixed session). - -Readability: - Token ids are SEMANTIC small integers (see TOKEN_NAMES). Each message renders - to ``[START, ...body, END]`` with a per-role band, so an id like 2001 reads as - ``u:compute`` and 7001 reads as ````. Expected token sequences are built - with the same render_* helpers used to feed append_turn, never hand-typed - magic numbers. - -Dual mode: - Every case is a ``test_*`` function doing strict assertions. ``main()`` runs - them all and, after each, prints the routing tree (token ids decoded to names) - and every linearized Sample with token / loss_mask aligned, so a human can read - exactly where each branch happened. Run with:: - - python -m tests.test_agent.test_trajectory_manager_e2e -""" - -from __future__ import annotations - -from tests.test_agent.test_claude_code_agent._dump_helpers import dump_tree_txt # noqa: E402 - -from slime.agent.adapters.common import TurnRecord # noqa: E402 -from slime.agent.trajectory_manager import TrajectoryManager, _lcp_len # noqa: E402 -from slime.utils.types import Sample # noqa: E402 - -# =========================================================================== -# §1 Semantic token vocabulary + reverse table -# =========================================================================== -# -# Per-role band. A message renders to [START, ...body, END]; the generation -# prompt appends the assistant START as the open-turn marker. - -_BANDS = { - "system": 1000, - "user": 2000, - "assistant": 9000, - "tool": 3000, -} -_GEN = _BANDS["assistant"] # add_generation_prompt marker -_DRIFT_BAND = 7000 - -# Reverse table: token id -> human-readable name. Filled lazily as messages are -# registered so dumps translate ids back to labels. -TOKEN_NAMES: dict[int, str] = {} -_ABBR = {"system": "sys", "user": "usr", "assistant": "ast", "tool": "tul"} -for _role, _base in _BANDS.items(): - TOKEN_NAMES[_base] = f"<{_ABBR[_role]}>" - TOKEN_NAMES[_base + 9] = f"" -TOKEN_NAMES[_GEN] = "" - - -def name_of(tok: int) -> str: - """Human-readable name for a token id (falls back to the raw int).""" - return TOKEN_NAMES.get(tok, str(tok)) - - -def _vis(label: str) -> str: - """Make whitespace visible in a token label for the dump. - - Whitespace-only drift (e.g. a trailing space from a cc rewrite) is invisible - in a terminal, which makes ``r:ok`` vs ``r:ok `` indistinguishable. Render - spaces as ``␣`` so the difference is obvious in the readable output. - """ - return label.replace(" ", "␣") - - -_ASST_BODY: dict[str, int] = {} - - -def _asst_body(label: str) -> int: - """Stable assistant body token for a response/message label. - - An assistant message replayed in a later prompt must render to the SAME - tokens the model generated for it, otherwise a clean continuation can never - hold (the cumulative prompt+response would not prefix the next prompt). So - both ``render_response`` and an assistant ``MsgTok`` derive their body token - from this one function, keyed on the label. Bodies are assigned by a stable - per-label counter (NOT a hash) so distinct labels never collide on one id — - a collision would mislabel tokens in the dump and could spuriously match - across turns. - """ - if label not in _ASST_BODY: - body = _BANDS["assistant"] + 100 + len(_ASST_BODY) - _ASST_BODY[label] = body - TOKEN_NAMES[body] = f"r:{_vis(label)}" - return _ASST_BODY[label] - - -def render_ids(ids: list[int]) -> str: - """Decode an id list into a space-joined readable string.""" - return " ".join(name_of(t) for t in ids) - - -class MsgTok: - """A message bound to a fixed, deterministic token rendering. - - The same MsgTok always renders to the same token segment regardless of which - turn replays it (a clean tokenizer). Token-id drift is injected explicitly by - tests via ``drift`` — never by re-rendering. - """ - - _body_counter: dict[str, int] = {} - - def __init__(self, role: str, label: str) -> None: - self.role = role - self.label = label - base = _BANDS[role] - if role == "assistant": - # An assistant message must render to the same body token as the - # response it represents (label-keyed), so a replayed assistant in a - # later prompt token-matches the original generation -> clean - # continuation. See _asst_body. - self.body = _asst_body(label) - else: - # Allocate one stable body token per (role, label). Offset past the - # END marker (base+9): the counter is shared across cases, so bodies - # must never climb into base+9 (END) or they'd collide with it. - idx = MsgTok._body_counter.setdefault(role, 0) + 1 - MsgTok._body_counter[role] = idx - self.body = base + 10 + idx - TOKEN_NAMES[self.body] = f"{role}:{_vis(label)}" - # message dict as the manager sees it (drives node_match_key). - self.message = {"role": role, "content": label} - - def render(self) -> list[int]: - """[START, body, END] for this message.""" - base = _BANDS[self.role] - return [base, self.body, base + 9] - - -def sys_msg(label: str) -> MsgTok: - return MsgTok("system", label) - - -def usr_msg(label: str) -> MsgTok: - return MsgTok("user", label) - - -def asst_msg(label: str) -> MsgTok: - return MsgTok("assistant", label) - - -def tool_msg(label: str) -> MsgTok: - return MsgTok("tool", label) - - -def render_prompt(msgs: list[MsgTok]) -> list[int]: - """Render a prompt message list, appending the generation-prompt marker.""" - out: list[int] = [] - for m in msgs: - out.extend(m.render()) - out.append(_GEN) - return out - - -def render_response(label: str) -> list[int]: - """Render an assistant response: [body, ]. - - The generation-prompt marker ```` equals the assistant START token, so - `` + render_response(x)`` == the assistant message ``[, body, - ]`` replayed in a later prompt. That identity is what makes a clean - continuation hold across turns. - """ - return [_asst_body(label), _BANDS["assistant"] + 9] - - -def messages(msgs: list[MsgTok]) -> list[dict]: - """The plain message dicts append_turn wants for prompt_messages.""" - return [m.message for m in msgs] - - -def drift(ids: list[int], at: int, sentinel: int = _DRIFT_BAND + 1) -> list[int]: - """Return a copy of ``ids`` with a sentinel spliced at index ``at``. - - The sentinel sits in the drift band (7000+), so a dump shows ```` at - the exact divergence point. Splicing (insert) makes ``len`` grow by one, - which is enough to make the lcp diverge at ``at``. - """ - TOKEN_NAMES[sentinel] = "" - return ids[:at] + [sentinel] + ids[at:] - - -def drift_replace(ids: list[int], at: int, sentinel: int = _DRIFT_BAND + 2) -> list[int]: - """Return a copy of ``ids`` with the token at ``at`` REPLACED by a sentinel. - - Unlike ``drift`` this keeps length constant — used when a test wants the - divergence inside a response span without changing the cumulative length. - """ - TOKEN_NAMES[sentinel] = "" - out = list(ids) - out[at] = sentinel - return out - - -def turn(prompt_ids, response_ids, *, finish_reason="stop", logprobs=None) -> TurnRecord: - return TurnRecord( - prompt_ids=list(prompt_ids), - output_ids=list(response_ids), - finish_reason=finish_reason, - output_log_probs=list(logprobs) if logprobs is not None else [], - ) - - -# A scratch space for the dual-mode printer: each case appends (title, mgr, sid, -# samples) so main() can render after the assertions pass. -_PRINT_LOG: list[tuple[str, object, str, list]] = [] - -# Raw append_turn inputs, keyed by sid, captured at call time so the printer can -# show the SOURCE data (prompt_ids / response_ids / finish / logprobs) that fed -# the tree — before any tree-building or linearization happened. -_TURN_LOG: dict[str, list[dict]] = {} - - -def _record(title: str, mgr, sid: str, samples: list) -> None: - _PRINT_LOG.append((title, mgr, sid, samples)) - - -# Convenience: append a turn with semantic messages, auto-rendering prompt unless -# an explicit prompt_ids is supplied (for drift injection). -def append( - mgr: TrajectoryManager, - sid: str, - prompt_msgs: list[MsgTok], - response_label: str | None, - *, - prompt_ids=None, - response_ids=None, - finish_reason="stop", - logprobs=None, - tools=None, - response_message=None, -): - p = list(prompt_ids) if prompt_ids is not None else render_prompt(prompt_msgs) - if response_ids is not None: - r = list(response_ids) - elif response_label is not None: - r = render_response(response_label) - else: - r = [] - rmsg = response_message - if rmsg is None and response_label is not None: - rmsg = {"role": "assistant", "content": response_label} - lp = logprobs - # Capture the raw turn inputs for the human-readable dump before the manager - # consumes them. - _TURN_LOG.setdefault(sid, []).append( - { - "prompt_msgs": [f"{m.role}:{_vis(m.label)}" for m in prompt_msgs], - "prompt_ids": p, - "response_ids": r, - "finish": finish_reason, - "has_lp": lp is not None, - } - ) - mgr.append_turn( - sid, - turn=turn(p, r, finish_reason=finish_reason, logprobs=lp), - prompt_messages=messages(prompt_msgs), - tools=tools, - response_message=rmsg, - ) - return p, r - - -def _leaves(mgr, sid): - return [leaf for leaf in mgr._trees[sid].leaves() if not leaf.is_root] - - -def _iter_all(root): - """Yield every non-root node in the tree (pre-order).""" - stack = list(root.children) - while stack: - n = stack.pop() - yield n - stack.extend(n.children) - - -# A minimal OpenAI-shape tool spec, used to exercise tools-metadata routing. -TOOLS = [ - { - "type": "function", - "function": {"name": "run", "description": "Run.", "parameters": {"type": "object"}}, - } -] - - -# Tree text snapshot captured the instant before get_trajectory drains the sid, -# so the human-readable dump can show [tree] AND [samples] side by side even -# though get_trajectory consumes the session. -_TREE_SNAP: dict[str, str] = {} - -# Input reward passed to get_trajectory, keyed by sid, so the dump can show the -# split (input_reward / n_samples == per_sample_reward) explicitly. -_REWARD_IN: dict[str, float] = {} - - -def get_traj(mgr, sid, *args, **kwargs): - """get_trajectory wrapper that snapshots the tree before draining. - - Linearization (get_trajectory) pops the sid, so a later dump would only see - ````. Capturing the tree text here keeps the routing tree visible - next to the Samples it produced. The input ``reward`` is captured too so the - dump can show how it splits across the emitted samples. - """ - if mgr.has_session(sid): - _TREE_SNAP[sid] = dump_tree_txt(mgr, sid) - _REWARD_IN[sid] = kwargs.get("reward", 0.0) - samples = mgr.get_trajectory(sid, *args, **kwargs) - # Reward conservation: get_trajectory splits the input reward evenly across - # every emitted sample, so the per-sample shares must sum back to the input - # (modulo float error). This is the "averaged over sample count" invariant. - if samples: - total = sum(s.reward for s in samples) - assert abs(total - _REWARD_IN[sid]) < 1e-9, ( - "reward not conserved across split", - total, - _REWARD_IN[sid], - ) - return samples - - -def _check_invariants(samples): - for s in samples: - assert len(s.loss_mask) == len(s.rollout_log_probs) == s.response_length, ( - "alignment broken", - len(s.loss_mask), - len(s.rollout_log_probs), - s.response_length, - ) - assert sum(s.loss_mask) > 0, "fully-masked sample emitted" - - -def golden(sample) -> str: - """Render one Sample as a human-reviewable golden string. - - Every token is decoded to its readable name (````, ``r:done``, - ```` ...). The leading prompt prefix (no loss_mask entry) is shown as - plain names; the response region is shown with each TRAINED token (loss=1) - wrapped in ``[...]`` and each context token (loss=0) left bare. This makes the - full linearized result — tokens, where the response region starts, and - exactly which tokens carry training signal — a single literal a human can - eyeball and assert against, instead of hand-derived index arithmetic. - - Example: `` system:S user:u [r:ok] []`` - """ - toks = sample.tokens - resp_start = len(toks) - sample.response_length - parts: list[str] = [] - for i, t in enumerate(toks): - nm = name_of(t) - if i >= resp_start and sample.loss_mask[i - resp_start] == 1: - parts.append(f"[{nm}]") - else: - parts.append(nm) - return " ".join(parts) - - -def goldens(samples) -> list[str]: - return [golden(s) for s in samples] - - -# =========================================================================== -# §2 Group 1 — routing tree layer (append_turn shapes the tree) -# =========================================================================== - - -def test_1_1_single_turn_chain(): - mgr = TrajectoryManager() - sid = "1.1" - s = sys_msg("S") - u = usr_msg("compute") - p, r = append(mgr, sid, [s, u], "ok") - chain = _leaves(mgr, sid)[0].path_from_root() - assert [n.role for n in chain] == ["system", "user", "assistant"] - assert chain[-1].turn_index == 1 - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 1 - assert goldens(samples) == [ - " system:S user:compute [r:ok] []", - ] - _check_invariants(samples) - _record("1.1 single turn -> linear chain", mgr, sid, samples) - print("PASS 1.1") - - -def test_1_2_clean_multiturn_with_tool(): - mgr = TrajectoryManager() - sid = "1.2" - s, u = sys_msg("S"), usr_msg("compute") - a1, t1 = asst_msg("call"), tool_msg("4") - p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls") - p2, r2 = append(mgr, sid, [s, u, a1, t1], "done") - chain = _leaves(mgr, sid)[0].path_from_root() - assert [n.role for n in chain] == ["system", "user", "assistant", "tool", "assistant"] - assert mgr.turn_count(sid) == 2 - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 1 - assert goldens(samples) == [ - " system:S user:compute [r:call] [] " - " tool:4 [r:done] []", - ] - _check_invariants(samples) - _record("1.2 clean 2-turn with tool -> single chain", mgr, sid, samples) - print("PASS 1.2") - - -def test_1_3_system_fork(): - mgr = TrajectoryManager() - sid = "1.3" - for sl in ["SA", "SB"]: - append(mgr, sid, [sys_msg(sl), usr_msg("u")], "a") - root = mgr._trees[sid] - assert len(root.children) == 2, "different system -> two subtrees at root" - assert len(_leaves(mgr, sid)) == 2 - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert goldens(samples) == [ - " system:SA user:u [r:a] []", - " system:SB user:u [r:a] []", - ] - _check_invariants(samples) - _record("1.3 system fork -> two subtrees at root", mgr, sid, samples) - print("PASS 1.3") - - -def test_1_4_user_fork_shared_system(): - mgr = TrajectoryManager() - sid = "1.4" - s = sys_msg("S") - for ul in ["A", "B"]: - append(mgr, sid, [s, usr_msg(ul)], ul.lower()) - root = mgr._trees[sid] - assert len(root.children) == 1, "system shared" - assert len(root.children[0].children) == 2, "user level forks" - assert len(_leaves(mgr, sid)) == 2 - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert goldens(samples) == [ - " system:S user:A [r:a] []", - " system:S user:B [r:b] []", - ] - _check_invariants(samples) - _record("1.4 user fork (shared system)", mgr, sid, samples) - print("PASS 1.4") - - -def test_1_5_assistant_message_fork(): - """Same (sys,user) prefix, two distinct assistant turns -> assistant fork.""" - mgr = TrajectoryManager() - sid = "1.5" - s, u = sys_msg("S"), usr_msg("u") - append(mgr, sid, [s, u], "a1") - append(mgr, sid, [s, u], "a2") - user_node = mgr._trees[sid].children[0].children[0] - assert len(user_node.children) == 2, "two assistant leaves hang off shared user" - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - # Two independent single-turn leaves sharing only the (sys,user) prefix. - assert goldens(samples) == [ - " system:S user:u [r:a1] []", - " system:S user:u [r:a2] []", - ] - _check_invariants(samples) - _record("1.5 assistant fork under shared user", mgr, sid, samples) - print("PASS 1.5") - - -def test_1_6_tool_fork_shared_assistant(): - """Same first assistant turn, two different tool results -> tool-level fork, - making the first assistant a shared snapshot node with 2 children.""" - mgr = TrajectoryManager() - sid = "1.6" - s, u, a1 = sys_msg("S"), usr_msg("u"), asst_msg("call") - append(mgr, sid, [s, u], "call", finish_reason="tool_calls") - append(mgr, sid, [s, u, a1, tool_msg("x")], "ax") - append(mgr, sid, [s, u, a1, tool_msg("y")], "ay") - asst1 = mgr._trees[sid].children[0].children[0].children[0] - assert asst1.role == "assistant" and asst1.turn_prompt_ids is not None - assert len(asst1.children) == 2, "shared assistant forks at the tool level" - assert len(_leaves(mgr, sid)) == 2 - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - # Leaf X owns the shared turn 1 (r:call trained); leaf Y shares it -> r:call - # demoted to loss=0 context, only r:ay trains (cross-leaf dedup). - assert goldens(samples) == [ - " system:S user:u [r:call] [] " - " tool:x [r:ax] []", - " system:S user:u r:call " " tool:y [r:ay] []", - ] - _check_invariants(samples) - _record("1.6 tool fork (shared assistant snapshot)", mgr, sid, samples) - print("PASS 1.6") - - -def test_1_7_token_only_drift_no_fork(): - """Identical messages, tampered prompt_ids -> NO tree fork (DFS ignores - tokens), but the drift DOES surface in the linearized sample: it lands in - leaf 2's prompt region (stripped / loss=0), proving token drift cannot - corrupt a trained response yet is still carried in the sample tokens.""" - mgr = TrajectoryManager() - sid = "1.7" - s, u = sys_msg("S"), usr_msg("u") - pa, _ = append(mgr, sid, [s, u], "a") - tampered = drift(pa, 1) # spliced into the prompt at index 1 - append(mgr, sid, [s, u], "b", prompt_ids=tampered) - # Tree: (sys,user) shared, two assistant turns hang off it -> two leaves; the - # path above the assistant is single (NOT forked on tokens). - user_node = mgr._trees[sid].children[0].children[0] - assert len(user_node.children) == 2, "two assistant turns share the (sys,user) path" - assert len(mgr._trees[sid].children) == 1 - - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 2 - # Leaf 1: clean. Leaf 2: the token sits in the stripped prompt region - # (bare, no brackets); the response r:b is fully trained ([...]). - assert goldens(samples) == [ - " system:S user:u [r:a] []", - " system:S user:u [r:b] []", - ] - # Belt-and-suspenders on the drift placement: token present, but never inside - # the response region. - s_b = samples[1] - assert (_DRIFT_BAND + 1) in s_b.tokens, "drift token is still carried in the sample" - assert (_DRIFT_BAND + 1) not in s_b.tokens[len(tampered) :], "drift not in the response region" - _check_invariants(samples) - _record("1.7 token-only drift -> no tree fork, drift lands in stripped prompt", mgr, sid, samples) - print("PASS 1.7") - - -def test_1_8_multi_tool_per_turn(): - mgr = TrajectoryManager() - sid = "1.8" - s, u, a1 = sys_msg("S"), usr_msg("u"), asst_msg("call") - ta, tb = tool_msg("A"), tool_msg("B") - append(mgr, sid, [s, u], "call", finish_reason="tool_calls") - append(mgr, sid, [s, u, a1, ta, tb], "done") - chain = _leaves(mgr, sid)[0].path_from_root() - assert [n.role for n in chain] == ["system", "user", "assistant", "tool", "tool", "assistant"] - assert chain[3].messages == [ta.message] - assert chain[4].messages == [tb.message] - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 1 - assert goldens(samples) == [ - " system:S user:u [r:call] [] " - " tool:A tool:B [r:done] []", - ] - _check_invariants(samples) - _record("1.8 multi-tool turn -> one node per tool", mgr, sid, samples) - print("PASS 1.8") - - -def test_1_9_cross_sid_isolation(): - mgr = TrajectoryManager() - s = sys_msg("S") - for sid, ul in [("sid-a", "A"), ("sid-b", "B")]: - append(mgr, sid, [s, usr_msg(ul)], ul.lower()) - assert len(_leaves(mgr, "sid-a")) == 1 - assert len(_leaves(mgr, "sid-b")) == 1 - assert mgr._trees["sid-a"] is not mgr._trees["sid-b"] - sa = get_traj(mgr, "sid-a", base_sample=Sample(index=0, prompt=""), reward=1.0) - sb = get_traj(mgr, "sid-b", base_sample=Sample(index=1, prompt=""), reward=1.0) - assert goldens(sa) == [" system:S user:A [r:a] []"] - assert goldens(sb) == [" system:S user:B [r:b] []"] - _check_invariants(sa) - _check_invariants(sb) - _record("1.9 cross-sid isolation (sid-a)", mgr, "sid-a", sa) - _record("1.9 cross-sid isolation (sid-b)", mgr, "sid-b", sb) - print("PASS 1.9") - - -def test_1_10_empty_response(): - mgr = TrajectoryManager() - sid = "1.10" - s, u = sys_msg("S"), usr_msg("u") - append(mgr, sid, [s, u], None, response_ids=[], response_message=None, finish_reason="length") - asst = _leaves(mgr, sid)[0] - assert asst.role == "assistant" - assert asst.turn_response_ids == [] - assert asst.messages == [] - # Empty response -> the only turn has no trainable token, so its segment is - # dropped at linearization (no fully-masked sample). Zero samples is correct. - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 0 - _record("1.10 empty response -> assistant leaf, no message (0 samples)", mgr, sid, samples) - print("PASS 1.10") - - -# =========================================================================== -# §2 Group 2 — linearization layer (get_trajectory token routing) -# =========================================================================== - - -def test_2_1_single_turn_linearize(): - mgr = TrajectoryManager() - sid = "2.1" - s, u = sys_msg("S"), usr_msg("u") - p, r = append(mgr, sid, [s, u], "a", logprobs=None) - # attach explicit logprobs so we can check propagation - leaf = _leaves(mgr, sid)[0] - leaf.turn_response_logprobs = [-0.5] * len(r) - samples = get_traj(mgr, sid, base_sample=Sample(index=7, prompt="hi"), reward=1.0) - assert len(samples) == 1 - s0 = samples[0] - assert goldens(samples) == [" system:S user:u [r:a] []"] - assert s0.rollout_log_probs == [-0.5] * len(r) - assert s0.reward == 1.0 - _check_invariants(samples) - _record("2.1 single-turn linearize", mgr, sid, samples) - print("PASS 2.1") - - -def test_2_2_clean_multiturn_linearize(): - mgr = TrajectoryManager() - sid = "2.2" - s, u, a1, t1 = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("4") - p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) - p2, r2 = append(mgr, sid, [s, u, a1, t1], "done", logprobs=[-0.4] * 2) - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 1 - s0 = samples[0] - L = _lcp_len(p1 + r1, p2) - assert goldens(samples) == [ - " system:S user:u [r:call] [] " - " tool:4 [r:done] []", - ] - assert s0.rollout_log_probs == [-0.5] * len(r1) + [0.0] * (len(p2) - L) + [-0.4] * len(r2) - _check_invariants(samples) - _record("2.2 clean 2-turn linearize", mgr, sid, samples) - print("PASS 2.2") - - -def test_2_3_drift_case_A_forks(): - """Drift inside a PROMPT region -> case A -> fork, no token dropped.""" - mgr = TrajectoryManager() - sid = "2.3" - s, u, a1, t = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("t") - p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) - p2_honest = render_prompt([s, u, a1, t]) - p2 = drift(p2_honest, len(p1) - 1) # inside p1's prompt region - p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2, logprobs=[-0.4] * 2) - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 2 - # case-A fork: two coherent single-turn segments; the token stays in - # segment 2's stripped prompt region (bare), no token dropped. - assert goldens(samples) == [ - " system:S user:u [r:call] []", - " system:S user:u r:call " - " tool:t [r:done] []", - ] - assert all(abs(s.reward - 1.0 / len(samples)) < 1e-9 for s in samples) - _check_invariants(samples) - _record("2.3 drift case A (prompt region) -> fork", mgr, sid, samples) - print("PASS 2.3") - - -def test_2_4_drift_case_B1_short_replaces(): - """Small drift inside the most-recent response span -> replace.""" - mgr = TrajectoryManager() # default threshold 1024 - sid = "2.4" - s, u, a1, t = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("t") - p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) - p2_honest = render_prompt([s, u, a1, t]) - assert p2_honest[: len(p1) + len(r1)] == p1 + r1 - drift_idx = len(p1) + len(r1) - 1 # last token of r1's echo - p2 = drift_replace(p2_honest, drift_idx) - p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2, logprobs=[-0.4] * 2) - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 1 - s0 = samples[0] - L = _lcp_len(p1 + r1, p2) - assert L == drift_idx - # replace: the drifted r:call tail is dropped and re-supplied as loss=0 prompt - # context (the token marks the divergence); only the surviving head of - # r:call and the new r:done train. - assert goldens(samples) == [ - " system:S user:u [r:call] " - " tool:t [r:done] []", - ] - assert s0.rollout_log_probs == [-0.5] * (L - len(p1)) + [0.0] * (len(p2) - L) + [-0.4] * len(r2) - _check_invariants(samples) - _record("2.4 drift case B1 (small) -> replace", mgr, sid, samples) - print("PASS 2.4") - - -def test_2_5_drift_case_B1_long_forks(): - mgr = TrajectoryManager(fork_merge_max_response_tokens=1) # d>=1 -> fork - sid = "2.5" - s, u, a1, t = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("t") - p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls") - p2_honest = render_prompt([s, u, a1, t]) - p2 = drift_replace(p2_honest, len(p1) + len(r1) - 1) - p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2) - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 2 - # Both segments single-turn (the drift forked them apart): each trains its own - # response. The sits in segment 2's stripped prompt. - assert goldens(samples) == [ - " system:S user:u [r:call] []", - " system:S user:u r:call " - " tool:t [r:done] []", - ] - assert all(abs(s.reward - 1.0 / len(samples)) < 1e-9 for s in samples) - _check_invariants(samples) - _record("2.5 drift case B1 (long) -> fork", mgr, sid, samples) - print("PASS 2.5") - - -def test_2_6_drift_case_B1_threshold_zero_forks(): - mgr = TrajectoryManager(fork_merge_max_response_tokens=0) - sid = "2.6" - s, u, a1, t = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("t") - p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls") - p2_honest = render_prompt([s, u, a1, t]) - p2 = drift_replace(p2_honest, len(p1) + len(r1) - 1) - p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2) - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 2 - assert goldens(samples) == [ - " system:S user:u [r:call] []", - " system:S user:u r:call " - " tool:t [r:done] []", - ] - _check_invariants(samples) - _record("2.6 drift case B1 threshold=0 -> fork", mgr, sid, samples) - print("PASS 2.6") - - -def test_2_7_drift_case_B2_earlier_turn_forks(): - """Drift inside an EARLIER turn's response span -> always fork.""" - mgr = TrajectoryManager() - sid = "2.7" - s, u = sys_msg("S"), usr_msg("u") - a1, t1 = asst_msg("a1"), tool_msg("t1") - a2, t2 = asst_msg("a2"), tool_msg("t2") - p1, r1 = append(mgr, sid, [s, u], "a1", finish_reason="tool_calls", logprobs=[-0.5] * 2) - p2, r2 = append(mgr, sid, [s, u, a1, t1], "a2", finish_reason="tool_calls", logprobs=[-0.4] * 2) - p3_honest = render_prompt([s, u, a1, t1, a2, t2]) - p3 = drift_replace(p3_honest, len(p1) + len(r1) - 1) # inside r1 (earlier span) - p3, r3 = append(mgr, sid, [s, u, a1, t1, a2, t2], "a3", prompt_ids=p3, logprobs=[-0.3] * 2) - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 2 - # Segment 1 = clean turns 1+2; segment 2 = turn 3 alone (forked because the - # drift hit an EARLIER turn's response span, which replace can't drop). - assert goldens(samples) == [ - " system:S user:u [r:a1] [] " - " tool:t1 [r:a2] []", - " system:S user:u r:a1 " - " tool:t1 r:a2 tool:t2 [r:a3] []", - ] - assert all(abs(s.reward - 1.0 / len(samples)) < 1e-9 for s in samples) - _check_invariants(samples) - _record("2.7 drift case B2 (earlier turn) -> fork", mgr, sid, samples) - print("PASS 2.7") - - -def test_2_8_fork_reward_split(): - mgr = TrajectoryManager() - sid = "2.8" - s, u, a1, t = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("t") - p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls") - p2_honest = render_prompt([s, u, a1, t]) - p2 = drift(p2_honest, len(p1) - 1) # prompt region -> case A fork - p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2) - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 2 - # case-A fork: two single-turn segments, each trains its own response. - assert goldens(samples) == [ - " system:S user:u [r:call] []", - " system:S user:u r:call " - " tool:t [r:done] []", - ] - # reward 1.0 split evenly across the 2 forked samples -> 0.5 each. - assert all(abs(s.reward - 0.5) < 1e-9 for s in samples) - _check_invariants(samples) - _record("2.8 fork reward split (1.0 / 2 = 0.5 each)", mgr, sid, samples) - print("PASS 2.8") - - -def test_2_9_two_leaves_reward_split(): - mgr = TrajectoryManager() - sid = "2.9" - s = sys_msg("S") - for ul in ["A", "B"]: - append(mgr, sid, [s, usr_msg(ul)], ul.lower()) - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 2 - assert goldens(samples) == [ - " system:S user:A [r:a] []", - " system:S user:B [r:b] []", - ] - # reward 1.0 split evenly across the 2 leaves -> 0.5 each. - assert all(abs(s.reward - 0.5) < 1e-9 for s in samples) - _check_invariants(samples) - _record("2.9 two leaves reward split (1.0 / 2 = 0.5 each)", mgr, sid, samples) - print("PASS 2.9") - - -def test_2_10_cross_leaf_dedup(): - """Shared assistant trained on first leaf only; second leaf re-emits it - as loss=0 context.""" - mgr = TrajectoryManager() - sid = "2.10" - s, u, a1 = sys_msg("S"), usr_msg("u"), asst_msg("call") - tx, ty = tool_msg("x"), tool_msg("y") - p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) - p2, r2 = append(mgr, sid, [s, u, a1, tx], "a2", logprobs=[-0.4] * 2) - p3, r3 = append(mgr, sid, [s, u, a1, ty], "a3", logprobs=[-0.3] * 2) - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 2 - s_second = samples[1] - # First leaf trains the shared r:call + its own r:a2; second leaf shares - # r:call (demoted to loss=0) and trains only r:a3. - assert goldens(samples) == [ - " system:S user:u [r:call] [] " - " tool:x [r:a2] []", - " system:S user:u r:call " " tool:y [r:a3] []", - ] - assert s_second.rollout_log_probs == [0.0] * (len(p3) - len(p1)) + [-0.3] * len(r3) - _check_invariants(samples) - _record("2.10 cross-leaf dedup (shared assistant trained once)", mgr, sid, samples) - print("PASS 2.10") - - -def test_2_11_routing_only_assistant_filtered(): - """cc replays an assistant the manager never recorded -> mounts routing-only, - must be filtered out of the strict-prefix walk (no raise).""" - mgr = TrajectoryManager() - sid = "2.11" - s, u = sys_msg("S"), usr_msg("u") - a1, t1 = asst_msg("a1"), tool_msg("t1") - a2 = asst_msg("a2") - append(mgr, sid, [s, u], "a1", finish_reason="tool_calls") - append(mgr, sid, [s, u, a1, t1], "a2", finish_reason="tool_calls") - foreign = asst_msg("foreign") - t2 = tool_msg("t2") - append(mgr, sid, [s, u, a1, t1, a2, foreign, t2], "a3") - leaves = _leaves(mgr, sid) - assert len(leaves) == 1 - chain = leaves[0].path_from_root() - routing = [n for n in chain if n.role == "assistant" and n.turn_prompt_ids is None] - assert len(routing) == 1 and routing[0].messages[0]["content"] == "foreign" - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 1 - # The foreign assistant (r:foreign) is routing-only -> appears as bare context - # (no brackets); the three real turns r:a1/r:a2/r:a3 train. - assert goldens(samples) == [ - " system:S user:u [r:a1] [] " - " tool:t1 [r:a2] [] r:foreign " - " tool:t2 [r:a3] []", - ] - _record("2.11 routing-only assistant filtered (no raise)", mgr, sid, samples) - print("PASS 2.11") - - -def test_2_12_drop_clears_sid(): - mgr = TrajectoryManager() - sid = "2.12" - s, u = sys_msg("S"), usr_msg("u") - append(mgr, sid, [s, u], "a") - assert mgr.has_session(sid) - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert goldens(samples) == [" system:S user:u [r:a] []"] - assert not mgr.has_session(sid) - assert mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt="")) == [] - _check_invariants(samples) - _record("2.12 drop clears sid (2nd get_trajectory -> [])", mgr, sid, samples) - print("PASS 2.12") - - -# =========================================================================== -# §2 Group 3 — combined / stress (both layers interacting) -# =========================================================================== - - -def test_3_1_rewrite_merge_absorbs_short(): - mgr = TrajectoryManager() - sid = "3.1" - s, u = sys_msg("S"), usr_msg("u") - a1_rw = asst_msg("ok ") # cc-rewritten (different message identity) - t1 = tool_msg("t") - append(mgr, sid, [s, u], "ok", finish_reason="tool_calls", logprobs=[-0.5] * 2) - append(mgr, sid, [s, u, a1_rw, t1], "done", logprobs=[-0.4] * 2) - leaves = _leaves(mgr, sid) - assert len(leaves) == 1, "short rewrite absorbed, not forked" - chain = leaves[0].path_from_root() - merged = chain[2] - assert merged.turn_prompt_ids is None and merged.turn_index is None - assert merged.messages == [a1_rw.message] - assert merged.metadata["merged_rewrite"]["abandoned_turn_index"] == 1 - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 1 - # The abandoned turn-1 response (r:ok␣) is demoted to routing-only -> appears - # bare; only the surviving turn-2 r:done trains. - assert goldens(samples) == [ - " system:S user:u r:ok␣ " " tool:t [r:done] []", - ] - _check_invariants(samples) - _record("3.1 rewrite-merge absorbs short assistant", mgr, sid, samples) - print("PASS 3.1") - - -def test_3_2_rewrite_merge_long_forks(): - mgr = TrajectoryManager(fork_merge_max_response_tokens=1) # r1 len 2 >= 1 - sid = "3.2" - s, u = sys_msg("S"), usr_msg("u") - a1_rw, t1 = asst_msg("ok2 "), tool_msg("t") - append(mgr, sid, [s, u], "ok", finish_reason="tool_calls") - append(mgr, sid, [s, u, a1_rw, t1], "done") - assert len(_leaves(mgr, sid)) == 2, "long rewrite forks" - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 2 - # Leaf 1: the abandoned turn-1 standalone (r:ok). Leaf 2: turn-2 only (the - # rewritten r:ok2␣ assistant mounts routing-only and is filtered out). - assert goldens(samples) == [ - " system:S user:u [r:ok] []", - " system:S user:u r:ok2␣ " " tool:t [r:done] []", - ] - _check_invariants(samples) - _record("3.2 rewrite-merge long -> fork", mgr, sid, samples) - print("PASS 3.2") - - -def test_3_3_rewrite_merge_threshold_zero_forks(): - mgr = TrajectoryManager(fork_merge_max_response_tokens=0) - sid = "3.3" - s, u = sys_msg("S"), usr_msg("u") - a1_rw, t1 = asst_msg("ok3 "), tool_msg("t") - append(mgr, sid, [s, u], "ok", finish_reason="tool_calls") - append(mgr, sid, [s, u, a1_rw, t1], "done") - assert len(_leaves(mgr, sid)) == 2 - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 2 - assert goldens(samples) == [ - " system:S user:u [r:ok] []", - " system:S user:u r:ok3␣ " " tool:t [r:done] []", - ] - _check_invariants(samples) - _record("3.3 rewrite-merge threshold=0 -> fork", mgr, sid, samples) - print("PASS 3.3") - - -def test_3_4_rewrite_merge_ambiguous_forks(): - mgr = TrajectoryManager() - sid = "3.4" - s, u = sys_msg("S"), usr_msg("u") - # two short assistant leaves under shared (sys,user) - append(mgr, sid, [s, u], "a") - append(mgr, sid, [s, u], "b") - a_c, t1 = asst_msg("c"), tool_msg("t") - append(mgr, sid, [s, u, a_c, t1], "d") - assert len(_leaves(mgr, sid)) == 3, "ambiguous candidates fork" - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 3 - # Leaves "a" and "b" are standalone single turns; leaf "d" carries the - # ambiguous-rewrite assistant (r:c) as routing-only (bare) -> trains only r:d. - assert goldens(samples) == [ - " system:S user:u [r:a] []", - " system:S user:u [r:b] []", - " system:S user:u r:c " " tool:t [r:d] []", - ] - _check_invariants(samples) - _record("3.4 rewrite-merge ambiguous -> fork", mgr, sid, samples) - print("PASS 3.4") - - -def test_3_5_rewrite_merge_match_key_updated(): - """After merge, a later turn replaying the rewritten message must descend - through the merged node (match_key updated), not fork again.""" - mgr = TrajectoryManager() - sid = "3.5" - s, u = sys_msg("S"), usr_msg("u") - a1_rw = asst_msg("ok5 ") - t1, a2, t2 = tool_msg("t1"), asst_msg("second"), tool_msg("t2") - append(mgr, sid, [s, u], "ok", finish_reason="tool_calls") - append(mgr, sid, [s, u, a1_rw, t1], "second", finish_reason="tool_calls", logprobs=[-0.4] * 2) - append(mgr, sid, [s, u, a1_rw, t1, a2, t2], "third", logprobs=[-0.3] * 2) - leaves = _leaves(mgr, sid) - assert len(leaves) == 1, "match_key updated -> no spurious fork" - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 1 - # Turn 1 (r:ok) was absorbed as routing-only (rewrite merge), so it appears - # bare; turns 2 and 3 (r:second / r:third) train in one clean chain. - assert goldens(samples) == [ - " system:S user:u r:ok5␣ " - " tool:t1 [r:second] [] tool:t2 [r:third] []", - ] - _check_invariants(samples) - _record("3.5 rewrite-merge match_key updated", mgr, sid, samples) - print("PASS 3.5") - - -def test_3_6_tree_fork_plus_token_drift(): - """A tree fork (two leaves) where ONE leaf also drift-forks internally, - yielding 3 Samples total. Combines layer-1 (message fork) with layer-2 - (token drift fork).""" - mgr = TrajectoryManager() - sid = "3.6" - s, u = sys_msg("S"), usr_msg("u") - a1, tx, ty = asst_msg("call"), tool_msg("x"), tool_msg("y") - ax2 = asst_msg("ax2") - p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) - # Leaf X: clean continuation, then a third turn with a case-A prompt drift. - append(mgr, sid, [s, u, a1, tx], "ax2", finish_reason="tool_calls", logprobs=[-0.4] * 2) - txx = tool_msg("xx") - p3_honest = render_prompt([s, u, a1, tx, ax2, txx]) - p3 = drift(p3_honest, len(p1) - 1) # case A drift -> fork inside leaf X - append(mgr, sid, [s, u, a1, tx, ax2, txx], "ax3", prompt_ids=p3, logprobs=[-0.2] * 2) - # Leaf Y: a separate tool result off the shared assistant. - append(mgr, sid, [s, u, a1, ty], "ay2", logprobs=[-0.1] * 2) - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 3, [s.tokens for s in samples] - assert goldens(samples) == [ - # Sample 0: leaf X first segment, trains the shared r:call (first claim) + r:ax2. - " system:S user:u [r:call] [] " - " tool:x [r:ax2] []", - # Sample 1: leaf X second segment, a FRESH segment after the case-A fork -> - # whole prompt stripped, only r:ax3 trains (the sits in its prompt). - " system:S user:u r:call " - " tool:x r:ax2 tool:xx [r:ax3] []", - # Sample 2: leaf Y, shares r:call (claimed by sample 0 -> bare), trains r:ay2. - " system:S user:u r:call " " tool:y [r:ay2] []", - ] - assert all(abs(s.reward - 1.0 / 3) < 1e-9 for s in samples) - _check_invariants(samples) - _record("3.6 tree fork + token drift -> 3 samples", mgr, sid, samples) - print("PASS 3.6") - - -def test_3_7_deep_multi_leaf_dedup(): - """Three leaves sharing a 2-level assistant prefix; the shared turns are - trained exactly once across all leaves.""" - mgr = TrajectoryManager() - sid = "3.7" - s, u = sys_msg("S"), usr_msg("u") - a1, t1 = asst_msg("a1"), tool_msg("t1") - a2 = asst_msg("a2") - append(mgr, sid, [s, u], "a1", finish_reason="tool_calls", logprobs=[-0.5] * 2) - append(mgr, sid, [s, u, a1, t1], "a2", finish_reason="tool_calls", logprobs=[-0.4] * 2) - # three different tool results off a2 -> three leaves sharing a1+a2 - for lbl in ["p", "q", "r"]: - append(mgr, sid, [s, u, a1, t1, a2, tool_msg(lbl)], f"end-{lbl}", logprobs=[-0.3] * 2) - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 3 - # Leaf 0 OWNS the shared r:a1 + r:a2 (both trained) and its own end-p; leaves - # 1 and 2 SHARE r:a1 + r:a2 (bare, claimed by leaf 0) and train only their own - # end response. - assert goldens(samples) == [ - " system:S user:u [r:a1] [] " - " tool:t1 [r:a2] [] tool:p [r:end-p] []", - " system:S user:u r:a1 " - " tool:t1 r:a2 tool:q [r:end-q] []", - " system:S user:u r:a1 " - " tool:t1 r:a2 tool:r [r:end-r] []", - ] - _check_invariants(samples) - _record("3.7 deep multi-leaf dedup (3 leaves, shared trained once)", mgr, sid, samples) - print("PASS 3.7") - - -def test_3_8_long_mixed_session(): - """A ~7-turn session combining clean continuation, a mid-session B1 replace, - and a final case-A fork — verifying the mechanisms chain without interfering.""" - mgr = TrajectoryManager() - sid = "3.8" - s, u = sys_msg("S"), usr_msg("u") - a = [asst_msg(f"a{i}") for i in range(6)] - t = [tool_msg(f"t{i}") for i in range(6)] - lp = [-0.5, -0.5] - # turn 1 - p1, r1 = append(mgr, sid, [s, u], "a0", finish_reason="tool_calls", logprobs=lp) - # turns 2..4 clean - prefix = [s, u, a[0], t[0]] - append(mgr, sid, prefix, "a1", finish_reason="tool_calls", logprobs=lp) - prefix = prefix + [a[1], t[1]] - append(mgr, sid, prefix, "a2", finish_reason="tool_calls", logprobs=lp) - prefix = prefix + [a[2], t[2]] - # turn 5: B1 small replace — drift the last token of the previous response. - p5_honest = render_prompt(prefix) - p5 = drift_replace(p5_honest, len(p5_honest) - 2) # near tail, inside last resp echo region - append(mgr, sid, prefix, "a3", prompt_ids=p5, finish_reason="tool_calls", logprobs=lp) - prefix = prefix + [a[3], t[3]] - # turn 6: clean - append(mgr, sid, prefix, "a4", finish_reason="tool_calls", logprobs=lp) - prefix = prefix + [a[4], t[4]] - # turn 7: case-A fork (drift in early prompt region) - p7_honest = render_prompt(prefix) - p7 = drift(p7_honest, len(p1) - 1) - append(mgr, sid, prefix, "a5", prompt_ids=p7, logprobs=lp) - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - # The final case-A fork splits the single leaf chain into 3 segments. - assert goldens(samples) == [ - # Segment 1: turns 1-4 in one clean chain. The turn-5 B1 replace dropped a - # drifted tail (the is gone here) and realigned, so r:a0..r:a3 all - # train. - " system:S user:u [r:a0] [] " - " tool:t0 [r:a1] [] tool:t1 [r:a2] [] " - " tool:t2 [r:a3] []", - # Segment 2: cross-leaf-style dedup within the chain — the prior turns are - # re-emitted as bare context and only r:a4 trains. - " system:S user:u r:a0 " - " tool:t0 r:a1 tool:t1 r:a2 " - " tool:t2 r:a3 tool:t3 [r:a4] []", - # Segment 3: turn 7 after the case-A fork (the in the early prompt - # region); whole prefix bare, only r:a5 trains. - " system:S user:u r:a0 " - " tool:t0 r:a1 tool:t1 r:a2 " - " tool:t2 r:a3 tool:t3 r:a4 " - " tool:t4 [r:a5] []", - ] - assert abs(sum(s.reward for s in samples) - 1.0) < 1e-9 - _check_invariants(samples) - _record(f"3.8 long mixed session -> {len(samples)} samples", mgr, sid, samples) - print("PASS 3.8") - - -# =========================================================================== -# §2 Group 4 — boundary / defensive / feature-completion -# -# Fills coverage gaps the matrix above left open: tools-metadata routing, -# input-validation contracts, mixed-logprobs trajectories, the case-B1 drift -# threshold boundary, and the default-base_sample path. -# =========================================================================== - - -def test_4_1_tools_metadata_on_first_system_only(): - """tools passed to append_turn attach to the FIRST system node only; a later - turn carrying the same system must NOT re-attach (dedup via the - system-ancestor walk).""" - mgr = TrajectoryManager() - sid = "4.1" - s, u = sys_msg("S"), usr_msg("u") - a1, t1 = asst_msg("call"), tool_msg("t") - append(mgr, sid, [s, u], "call", finish_reason="tool_calls", tools=TOOLS) - append(mgr, sid, [s, u, a1, t1], "done", tools=TOOLS) - sys_node = mgr._trees[sid].children[0] - assert sys_node.role == "system" - assert sys_node.metadata.get("tools") == TOOLS, "tools land on the first system node" - # No other node carries tools. - others = [n for n in _iter_all(mgr._trees[sid]) if n is not sys_node] - assert all(n.metadata.get("tools") is None for n in others), "tools attached exactly once" - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert goldens(samples) == [ - " system:S user:u [r:call] [] " - " tool:t [r:done] []", - ] - _check_invariants(samples) - _record("4.1 tools metadata on first system only", mgr, sid, samples) - print("PASS 4.1") - - -def test_4_2_logprobs_length_mismatch_raises(): - """output_log_probs whose length != output_ids -> ValueError at append_turn.""" - mgr = TrajectoryManager() - sid = "4.2" - s, u = sys_msg("S"), usr_msg("u") - bad = TurnRecord( - prompt_ids=render_prompt([s, u]), - output_ids=[9101, 9102, 9103], - finish_reason="stop", - output_log_probs=[-0.1, -0.2], # length 2 != 3 - ) - raised = False - try: - mgr.append_turn( - sid, - turn=bad, - prompt_messages=messages([s, u]), - tools=None, - response_message={"role": "assistant", "content": "x"}, - ) - except ValueError as e: - raised = True - assert "output_log_probs" in str(e) - assert raised, "expected ValueError on logprobs/ids length mismatch" - print("PASS 4.2") - - -def test_4_3_empty_prompt_messages_skipped(): - """Empty prompt_messages -> append_turn is a no-op (warns, no node, no turn).""" - mgr = TrajectoryManager() - sid = "4.3" - mgr.append_turn( - sid, - turn=turn([1], [2], finish_reason="stop"), - prompt_messages=[], - tools=None, - response_message=None, - ) - assert mgr.turn_count(sid) == 0 - # The tree may be created empty (root only) or absent; either way no leaf. - assert not mgr.has_session(sid) or list(_leaves(mgr, sid)) == [] - print("PASS 4.3") - - -def test_4_4_default_base_sample(): - """get_trajectory with base_sample=None uses a default Sample(index=0).""" - mgr = TrajectoryManager() - sid = "4.4" - s, u = sys_msg("S"), usr_msg("u") - append(mgr, sid, [s, u], "a") - # snapshot tree before drain so the dump still renders it - _TREE_SNAP[sid] = dump_tree_txt(mgr, sid) - _REWARD_IN[sid] = 1.0 - samples = mgr.get_trajectory(sid, reward=1.0) # no base_sample - assert len(samples) == 1 - assert samples[0].index == 0 - assert goldens(samples) == [" system:S user:u [r:a] []"] - _check_invariants(samples) - _record("4.4 default base_sample (None)", mgr, sid, samples) - print("PASS 4.4") - - -def test_4_5_mixed_logprobs_across_turns(): - """A trajectory where turn 1 carries logprobs and turn 2 does NOT: the - sample's turn-1 response region has real logprobs, the turn-2 region is - padded with 0.0 (the response is still trained, loss=1).""" - mgr = TrajectoryManager() - sid = "4.5" - s, u = sys_msg("S"), usr_msg("u") - a1, t1 = asst_msg("call"), tool_msg("t") - p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) - p2, r2 = append(mgr, sid, [s, u, a1, t1], "done", logprobs=None) # no logprobs - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - assert len(samples) == 1 - s0 = samples[0] - L = _lcp_len(p1 + r1, p2) - # both responses still trained (golden shows the loss layout)... - assert goldens(samples) == [ - " system:S user:u [r:call] [] " - " tool:t [r:done] []", - ] - # ...but turn-2's region carries padded 0.0 logprobs (it had none), while - # turn-1's region keeps its real logprobs. - assert s0.rollout_log_probs == [-0.5] * len(r1) + [0.0] * (len(p2) - L) + [0.0] * len(r2) - _check_invariants(samples) - _record("4.5 mixed logprobs across turns (turn2 padded 0.0)", mgr, sid, samples) - print("PASS 4.5") - - -def test_4_6_drift_B1_threshold_boundary(): - """case-B1 threshold is exclusive: a drift tail of length d == threshold - forks, d == threshold-1 replaces. Verify both sides of the boundary.""" - - def run(threshold, drift_tail_len): - mgr = TrajectoryManager(fork_merge_max_response_tokens=threshold) - sid = f"4.6-{threshold}-{drift_tail_len}" - s, u = sys_msg("S"), usr_msg("u") - # 4-token response so the divergence can sit d tokens before its end. - p1 = render_prompt([s, u]) - r1 = [9001, 9002, 9003, 9004] - mgr.append_turn( - sid, - turn=turn(p1, r1, finish_reason="tool_calls"), - prompt_messages=messages([s, u]), - tools=None, - response_message={"role": "assistant", "content": "a1"}, - ) - a1m = {"role": "assistant", "content": "a1"} - tm = tool_msg("t") - # honest turn-2 prompt echoes p1 + r1 then the tool block + gen marker. - p2_honest = p1 + r1 + tm.render() + [_GEN] - # divergence d tokens before the end of r1's echo (inside its response span). - drift_idx = len(p1) + len(r1) - drift_tail_len - p2 = drift_replace(p2_honest, drift_idx) - r2 = [9101, 9102] - mgr.append_turn( - sid, - turn=turn(p2, r2, finish_reason="stop"), - prompt_messages=[*messages([s, u]), a1m, tm.message], - tools=None, - response_message={"role": "assistant", "content": "done"}, - ) - samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) - return samples, p1, r1, p2, r2 - - # d == threshold -> fork: two single-turn segments, each trains its own resp. - forked, p1, r1, p2, r2 = run(threshold=2, drift_tail_len=2) - assert len(forked) == 2, f"d==threshold must fork, got {len(forked)}" - assert forked[0].tokens == p1 + r1 - assert forked[0].loss_mask == [1] * len(r1) - assert forked[1].tokens == p2 + r2 - assert forked[1].loss_mask == [1] * len(r2) - # d < threshold -> replace: one coherent segment realigned to p2. - replaced, p1b, r1b, p2b, r2b = run(threshold=2, drift_tail_len=1) - assert len(replaced) == 1, f"d None: - toks = s.tokens - resp_start = len(toks) - s.response_length - # build aligned token/loss rows over the response region (the trained part); - # the leading prompt prefix has no loss_mask entry. - names = [name_of(t) for t in toks] - loss = ["-"] * resp_start + [str(x) for x in s.loss_mask] - widths = [max(len(names[i]), len(loss[i])) for i in range(len(toks))] - tok_row = " ".join(names[i].ljust(widths[i]) for i in range(len(toks))) - loss_row = " ".join(loss[i].ljust(widths[i]) for i in range(len(toks))) - print(f" Sample#{idx} reward={s.reward:.3f} resp_len={s.response_length}") - print(f" tok : {tok_row}") - print(f" loss: {loss_row}") - - -def _print_raw_turns(sid: str) -> None: - """Print the raw append_turn inputs (the SOURCE data) for a sid. - - Shows, per turn, the prompt message labels and the actual prompt_ids / - response_ids decoded to readable names, plus finish_reason and whether - logprobs were attached. This is what fed the tree, before any building or - linearization. - """ - turns = _TURN_LOG.get(sid, []) - print(f"[raw turns] {len(turns)}") - for k, t in enumerate(turns, start=1): - msgs = " , ".join(t["prompt_msgs"]) - print(f" turn#{k} finish={t['finish']} has_logprobs={t['has_lp']}") - print(f" msgs : {msgs}") - print(f" prompt : {render_ids(t['prompt_ids'])}") - print(f" output : {render_ids(t['response_ids']) or ''}") - - -def _print_case(title: str, mgr, sid: str, samples: list) -> None: - print(f"\n=== CASE {title} ===") - _print_raw_turns(sid) - if mgr.has_session(sid): - txt = dump_tree_txt(mgr, sid) - else: - # Session already drained by get_trajectory; fall back to the snapshot - # captured by get_traj just before draining. - txt = _TREE_SNAP.get(sid, "") - print("[tree]") - for line in txt.splitlines(): - print(" " + line) - n = len(samples) - if n: - r_in = _REWARD_IN.get(sid, 0.0) - per = r_in / n - print(f"[samples] {n} (reward split: {r_in:.3f} / {n} = {per:.3f} per sample)") - else: - print(f"[samples] {n}") - for i, s in enumerate(samples): - _print_sample(i, s) - - -# =========================================================================== -# main -# =========================================================================== - - -_CASES = [ - test_1_1_single_turn_chain, - test_1_2_clean_multiturn_with_tool, - test_1_3_system_fork, - test_1_4_user_fork_shared_system, - test_1_5_assistant_message_fork, - test_1_6_tool_fork_shared_assistant, - test_1_7_token_only_drift_no_fork, - test_1_8_multi_tool_per_turn, - test_1_9_cross_sid_isolation, - test_1_10_empty_response, - test_2_1_single_turn_linearize, - test_2_2_clean_multiturn_linearize, - test_2_3_drift_case_A_forks, - test_2_4_drift_case_B1_short_replaces, - test_2_5_drift_case_B1_long_forks, - test_2_6_drift_case_B1_threshold_zero_forks, - test_2_7_drift_case_B2_earlier_turn_forks, - test_2_8_fork_reward_split, - test_2_9_two_leaves_reward_split, - test_2_10_cross_leaf_dedup, - test_2_11_routing_only_assistant_filtered, - test_2_12_drop_clears_sid, - test_3_1_rewrite_merge_absorbs_short, - test_3_2_rewrite_merge_long_forks, - test_3_3_rewrite_merge_threshold_zero_forks, - test_3_4_rewrite_merge_ambiguous_forks, - test_3_5_rewrite_merge_match_key_updated, - test_3_6_tree_fork_plus_token_drift, - test_3_7_deep_multi_leaf_dedup, - test_3_8_long_mixed_session, - test_4_1_tools_metadata_on_first_system_only, - test_4_2_logprobs_length_mismatch_raises, - test_4_3_empty_prompt_messages_skipped, - test_4_4_default_base_sample, - test_4_5_mixed_logprobs_across_turns, - test_4_6_drift_B1_threshold_boundary, -] - - -def main() -> None: - for case in _CASES: - case() - # Replay the captured tree / sample snapshots as human-readable dumps. - print("\n" + "=" * 70) - print("HUMAN-READABLE DUMPS") - print("=" * 70) - for title, mgr, sid, samples in _PRINT_LOG: - _print_case(title, mgr, sid, samples) - print(f"\nALL E2E CASES PASSED ({len(_CASES)} cases)") - - -if __name__ == "__main__": - main() From bc3d304fb0db51aa1d97a39d48620d8baed00588 Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Mon, 8 Jun 2026 16:28:49 +0000 Subject: [PATCH 25/28] chore(agent): drop comments/docstrings in generate.py and TurnRecord Remove branch-added inline comments and docstrings in generate.py, drop the SLIME_DRIFT_FORK_MIN_LOSS_TOKENS warning block, and strip the TurnRecord docstring in adapters/common.py. --- examples/coding_agent_rl/generate.py | 25 +------------------------ slime/agent/adapters/common.py | 9 --------- 2 files changed, 1 insertion(+), 33 deletions(-) diff --git a/examples/coding_agent_rl/generate.py b/examples/coding_agent_rl/generate.py index 7956f52600..8123dc50e7 100644 --- a/examples/coding_agent_rl/generate.py +++ b/examples/coding_agent_rl/generate.py @@ -10,8 +10,7 @@ 2. ``sandbox.git_diff`` captures the model-produced patch. 3. ``sandbox.evaluate`` scores that patch in a second clean sandbox. 4. ``_merge_samples`` combines reward + the ``list[Sample]`` returned by - ``adapter.finish_session(sid)`` (which drains the per-sid trajectory - tree inside ``TrajectoryManager``). + ``adapter.finish_session(sid)``. All sandbox-side details live in ``sandbox.py``; the LLM plumbing (Anthropic <-> SGLang /generate, token capture, 3-kind segment split) uses @@ -97,19 +96,8 @@ def __init__(self, args) -> None: "Without it the sandbox cannot dial back and the rollout will " "silently abort." ) - # Assistant-rewrite merge threshold (see TrajectoryManager): when cc - # re-renders a short prior assistant, absorb it onto the existing leaf - # instead of forking a reward-diluting stub Sample. None -> manager - # default (1024); <=0 disables. fork_merge_threshold = os.environ.get("SLIME_FORK_MERGE_MAX_RESPONSE_TOKENS") fork_merge_threshold = int(fork_merge_threshold) if fork_merge_threshold else None - # drift-fork remains unimplemented in the strict core; warn if set. - if os.environ.get("SLIME_DRIFT_FORK_MIN_LOSS_TOKENS"): - logger.warning( - "[coding_agent_rl] SLIME_DRIFT_FORK_MIN_LOSS_TOKENS is set but " - "currently ignored: TrajectoryManager uses strict exact-prefix " - "linearization and raises on TITO drift." - ) self.adapter = AnthropicAdapter( tokenizer=self.tokenizer, sglang_url=sglang_url, @@ -145,9 +133,6 @@ def __init__(self, args) -> None: # --------------------------------------------------------------------------- # Trajectory -> Sample conversion -# adapter.finish_session(sid) drains the per-sid tree in TrajectoryManager and -# returns a list[Sample]. One trajectory yields >=1 samples because the agent -# may compact + reset mid-run, forking sub-trees that each become a sample. # --------------------------------------------------------------------------- @dataclass(frozen=True) class RewardResult: @@ -190,14 +175,6 @@ def _merge_samples( elapsed_sec: float, instance_id: str, ) -> list[Sample]: - """Decorate per-leaf Samples returned by TrajectoryManager.get_trajectory. - - The manager already filled tokens / loss_mask / rollout_log_probs / - response_length / reward (reward / N). We decode ``sample.response`` from - the response tokens slice -- slime's training logging path reads this - string. Per-trajectory metadata is intentionally NOT attached to the - samples (kept empty); revisit when dump/analysis needs it. - """ if not samples: return _abort_result(sample, "adapter_session_empty") diff --git a/slime/agent/adapters/common.py b/slime/agent/adapters/common.py index f34daf4bad..a302540683 100644 --- a/slime/agent/adapters/common.py +++ b/slime/agent/adapters/common.py @@ -22,15 +22,6 @@ @dataclasses.dataclass(frozen=True) class TurnRecord: - """Exact token snapshot for one assistant generation, returned by - :func:`call_sglang_generate`. - - ``prompt_ids`` is the full tokenized prompt sent to the generator for that - turn. ``output_ids`` is the raw generated output, and - ``output_log_probs`` is aligned with it when the rollout engine returns - per-token log probabilities. - """ - prompt_ids: list[int] output_ids: list[int] finish_reason: str From ece9007a25f3fbdba7ce3ca736d784e15a9939ed Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Mon, 8 Jun 2026 16:55:23 +0000 Subject: [PATCH 26/28] docs(agent): tighten trajectory_manager comments to why-not-what Trim the module/why docstrings to the repo's comment conventions: keep invariants, gotchas, and cross-layer contracts (cross-leaf dedup, truncated-span loss=1, sort_keys list-order, fully-masked-segment drop); drop comments that merely restate the code. Rewrite the module docstring around the append_turn / get_trajectory data flow. --- slime/agent/trajectory_manager.py | 252 ++++-------------------------- 1 file changed, 31 insertions(+), 221 deletions(-) diff --git a/slime/agent/trajectory_manager.py b/slime/agent/trajectory_manager.py index b8ff2f34bb..61e40c302c 100644 --- a/slime/agent/trajectory_manager.py +++ b/slime/agent/trajectory_manager.py @@ -1,75 +1,10 @@ -"""Per-message trajectory tree manager (C-plan: token-faithful). - -Design (Plan C, 2026-06-03; strict exact-prefix rewrite 2026-06-08; -per-message routing 2026-06-08): - -* The tree is a router only. DFS merge keys on ``(role, node_match_key)`` - alone, one tree node per message — no prompt_ids prefix check and no - same-role grouping. Same conversation prefix in ``messages`` space always - lands on the same path, regardless of any chat_template re-tokenization - drift across turns. - -* Each assistant leaf stores the THIS-TURN sglang snapshot: - ``turn_prompt_ids`` / ``turn_response_ids`` / ``turn_response_logprobs`` - / ``turn_finish_reason`` / ``turn_index``. Non-assistant nodes carry no - token attribution at all. - -* ``get_trajectory`` linearizes each leaf turn-by-turn. Walking the leaf's - assistant chain root→leaf, the cumulative ``(prompt + response)`` tokens - emitted so far are matched against the next turn's ``turn_prompt_ids``: - - - **Clean continuation** — the cumulative tokens are an exact prefix of the - turn's prompt. The new prompt tail ``prompt[len(cumulative):]`` is appended - as loss_mask=0, then the turn's ``response`` is appended as loss_mask=1 with - real logprobs. - - - **Drift** — the same history re-tokenized differently across turns (TITO - drift: tool_call arg order, whitespace, reasoning-block reordering). Rather - than raise, ``get_trajectory`` tolerates the drift by where the divergence - index ``L`` (the common-prefix length) falls, never letting logprobs - misalign with tokens: - - * **case A** — ``L`` lands in a prompt region (outside every recorded - response span): a genuine prompt-level re-render → **fork** (finalize - the current coherent segment as its own Sample, restart a fresh segment - at this turn). Fork discards nothing. - * **case B1** — ``L`` lands inside the most-recent response span (the - immediately-previous turn's response got re-rendered). Let - ``d = len(cumulative) - L`` be the drifted tail length: ``d < - fork_threshold`` → **replace** (truncate to ``L``, silently drop the - drifted tail, realign to this turn's prompt); ``d >= fork_threshold`` → - **fork**. - * **case B2** — ``L`` lands inside an *earlier* turn's response span. - Replacing would discard that turn's tail plus every later turn, so this - always **forks** regardless of drift size. - - A fork splits one leaf into >=2 Samples; reward is split evenly across all - emitted Samples (see ``get_trajectory``). - -* ONE tolerated exception at the ROUTING (tree) layer: an assistant-rewrite - merge. cc sometimes re-renders a previously-recorded assistant message when - feeding it back as prompt (tool_call arg order, whitespace). The message no - longer matches, so DFS forks at that assistant — leaving the original short - turn as a standalone stub leaf -> its own Sample, diluting the trajectory's - evenly-split reward. ``_try_merge_assistant_rewrite`` absorbs such a rewrite - onto the existing leaf when its response is short enough - (``fork_threshold_tokens``), demoting that node to routing-only so it - contributes 0 training tokens. This is the MESSAGE-level dual of case B1's - TOKEN-level replace: the rewrite-merge triggers when the message dict differs - (DFS would fork), case B1 triggers when the message is identical but its - tokens drift inside one chain. They live at different layers and handle - different causes, so both are kept. - -* Cross-leaf dedup at the LINEARIZATION layer: a snapshot assistant node can be - shared by >=2 sibling leaves (it is the SAME Node object on each chain, since - ``_find_mount_point`` reuses children). Linearizing every leaf from the root - would otherwise train that shared prefix once per leaf. Instead the first leaf - to reach a node (DFS / build order) trains its response (loss=1); later leaves - re-emit it as loss=0 context (``trained=False`` in ``_Segment.extend``), so the - shared prefix is trained exactly once. The tree is left intact (unlike the - rewrite-merge's node demotion); only the per-leaf loss signal is masked. A - leaf's terminal turn is its own freshly-created node, never pre-claimed, so - every leaf keeps >=1 trained turn. +"""Build a per-session training trajectory from multi-turn conversation data. + +The :class:`TrajectoryManager` builds one trajectory per session. ``append_turn`` +feeds in each turn (prompt messages + the served model's sglang snapshot), +routing it into a per-sid message tree; ``get_trajectory`` then linearizes that +tree into a ``list[Sample]`` of loss-masked training rows, tolerating TITO +re-tokenization drift via fork/replace. """ from __future__ import annotations @@ -91,20 +26,6 @@ class Node: - """One node in the trajectory tree. - - Routing fields (every node): - role, messages, parent, children, metadata - - Per-turn snapshot fields (assistant leaves only — None on non-assistant - and on internal assistant nodes that aren't a turn's own leaf): - turn_prompt_ids: list[int] sglang prompt as fed to /generate - turn_response_ids: list[int] sglang output ids - turn_response_logprobs: list[float] - turn_finish_reason: str | None - turn_index: int 1-based, monotonic per session - """ - def __init__( self, *, @@ -122,7 +43,6 @@ def __init__( # never the message list itself), so the routing key is computed once # here and reused on every descent instead of re-serializing per turn. self.match_key = node_match_key(self.messages) - # per-turn snapshot self.turn_prompt_ids: list[int] | None = None self.turn_response_ids: list[int] | None = None self.turn_response_logprobs: list[float] | None = None @@ -162,17 +82,12 @@ def leaves(self) -> Iterator[Node]: def node_match_key(messages: list[dict[str, Any]]) -> str: - """Identity key for a node's message list. - - json.dumps(sort_keys=True) sorts dict-internal keys recursively; list - element order is preserved (which is what we want: message order and - tool_calls order are both semantically significant). - """ + # sort_keys sorts dict-internal keys but PRESERVES list order -- message + # order and tool_calls order are both semantically significant. return json.dumps(messages, sort_keys=True, ensure_ascii=False) def _lcp_len(a: list[int], b: list[int]) -> int: - """Length of the longest common prefix between two int lists.""" n = min(len(a), len(b)) i = 0 while i < n and a[i] == b[i]: @@ -186,14 +101,7 @@ def _lcp_len(a: list[int], b: list[int]) -> int: class _Segment: - """Token accumulator for one coherent linearized run within a chain. - - Holds the running ``buffer`` with aligned ``loss`` / ``logprobs`` and per-turn - response ``spans`` (``[start, end)`` half-open). Invariant: ``absorb_turn`` - always ends by appending a response, so ``spans[-1]`` is the most-recent - response span and the buffer ends at its end. A fork closes the segment and a - fresh one opens at the diverging turn (see module docstring case A/B1/B2). - """ + """Token accumulator for one coherent linearized run within a chain.""" def __init__(self) -> None: self.buffer: list[int] = [] @@ -204,25 +112,14 @@ def __init__(self) -> None: self.first_prompt_len: int = 0 def measure_drift(self, prompt_ids: list[int]) -> int: - """Length of the segment's token tail this turn's ``prompt_ids`` failed to reproduce. - - Normally 0 (clean continuation). A positive value IS token-id drift — the - same history re-tokenized differently this turn (TITO drift) — not an error. - """ return len(self.buffer) - _lcp_len(self.buffer, prompt_ids) def can_absorb_drift(self, drift: int, fork_threshold: int) -> bool: - """Whether a ``drift``-token re-tokenization can be realigned into this segment. - - Realignable only when the drift is confined to the most-recent response - span (case B1) and shorter than ``fork_threshold``; otherwise the caller - forks. See module docstring for case A/B1/B2. - """ if drift == 0: return True realign_at = len(self.buffer) - drift if not self.spans or realign_at < self.spans[-1][0]: - return False # case A (prompt region) or B2 (earlier response span) + return False return fork_threshold > 0 and drift < fork_threshold def absorb_turn( @@ -235,24 +132,15 @@ def absorb_turn( *, trained: bool = True, ) -> None: - """Append one turn, dropping any re-tokenization drift first so logprobs stay aligned. - - Drop the last ``drift`` tokens so the buffer re-anchors on the prefix this - turn's ``prompt_ids`` reproduced (shrinking the prior span if cut), then - append this turn's prompt tail (loss=0) and response (loss=1). A truncated - prior span stays loss=1: that region is both the prior turn's response and - this turn's prompt context. - - ``trained=False`` appends the response as loss=0 / logprob=0.0 instead — the - node is already owned by an earlier sibling leaf, so re-training it would - double-count the shared prefix (see ``_chain_to_sample`` claim-on-first-visit). - """ + """Append one turn, dropping any re-tokenization drift first so logprobs stay aligned.""" realign_at = len(self.buffer) - drift del self.buffer[realign_at:] del self.loss[realign_at:] del self.logprobs[realign_at:] if self.spans and realign_at < self.spans[-1][1]: s, _e, j = self.spans[-1] + # NOTE: the truncated prior span stays loss=1 -- that region is both + # the prior turn's response and this turn's prompt context. self.spans[-1] = (s, realign_at, j) # may collapse to empty (s == realign_at); harmless is_first_turn = not self.spans @@ -273,15 +161,9 @@ def absorb_turn( self.first_prompt_len = len(prompt_ids) # stripped at build time def response_strip(self) -> int: - """Start index of the response region (the leading first-turn prompt prefix).""" return min(self.first_prompt_len, len(self.loss)) def has_trained_response(self) -> bool: - """Whether the response region carries any loss=1 token. - - False only when every turn in the segment was claimed by an earlier - sibling leaf (cross-leaf dedup) -> no training signal to emit. - """ return any(self.loss[self.response_strip() :]) @@ -291,17 +173,7 @@ def has_trained_response(self) -> bool: class TrajectoryManager: - """Per-sid trajectory tree manager. - - See module docstring for the C-plan invariants. Each ``append_turn`` - mounts >=0 prompt nodes (one per message, under the deepest matching - ancestor) + exactly 1 assistant leaf carrying that turn's sglang snapshot. - """ - def __init__(self, *, fork_threshold_tokens: int | None = None) -> None: - # Drift fork/replace threshold for case-B1 (see module docstring case - # A/B1/B2). <=0 forces every B1 to fork (max fidelity); ``None`` means - # "use the default". self._fork_threshold: int = 1024 if fork_threshold_tokens is None else fork_threshold_tokens self._trees: dict[str, Node] = {} self._turn_count: dict[str, int] = {} @@ -348,13 +220,12 @@ def get_trajectory( reward: float = 0.0, extra_metadata: dict[str, Any] | None = None, ) -> list: - """Drain a sid into slime ``Sample`` objects, then drop the session. + """Linearize this sid's routing tree into slime ``Sample`` objects and + consume the session. - ``get_trajectory`` is the lifecycle boundary where the message routing - tree is linearized into token-normalized ``Sample`` objects. Each - routing leaf yields exactly one Sample. ``reward`` is split evenly - across all emitted samples. The sid is consumed: a second call for the - same sid returns ``[]``. + Each routing leaf yields one or more Samples; ``reward`` is split evenly + across all of them. The sid is dropped afterwards, so a second call for + the same sid returns ``[]``. """ if base_sample is None: base_sample = Sample(index=0, prompt="") @@ -364,9 +235,6 @@ def get_trajectory( return [] samples: list[Sample] = [] - # Cross-leaf dedup (see module docstring): a snapshot node shared by sibling - # leaves is trained only by the first leaf to reach it. ``claimed`` carries - # node identity (id()) across leaves to enforce that. claimed: set[int] = set() for routing_leaf in root.leaves(): if routing_leaf.is_root: @@ -376,10 +244,7 @@ def get_trajectory( self._chain_to_sample(chain, base_sample=base_sample, extra_metadata=extra_metadata, claimed=claimed) ) - # Reward is split evenly across every emitted sample (one per leaf); the - # token-weighted reducer downstream then gives each loss token the - # trajectory's full R. Assigned after the fact so the per-leaf builder - # stays reward-agnostic. + # TODO custom reward func per_sample_reward = (reward / len(samples)) if samples else 0.0 for s in samples: s.reward = per_sample_reward @@ -391,13 +256,6 @@ def get_trajectory( # -------------------- internals ---------------------------------------- def _find_mount_point(self, root: Node, messages: list[dict[str, Any]]) -> tuple[Node, int]: - """DFS down the existing tree by ``(role, node_match_key)``, per message. - - Returns ``(cur, i)``: ``cur`` is the deepest node whose path matches - ``messages[:i]`` exactly; ``i`` is the index into ``messages`` of the - first message that diverges from anything mounted so far (i.e., where - this turn's new content begins). - """ cur = root i = 0 while i < len(messages): @@ -420,17 +278,8 @@ def _try_merge_assistant_rewrite( prompt_messages: list[dict[str, Any]], i: int, ) -> tuple[Node, int]: - """Absorb a short assistant-rewrite onto its existing node instead of forking. - - See module docstring (rewrite-merge bullet) for the why. Purely a - reward-hygiene / de-fragmentation optimization — forking is already safe - (the rewrite mounts as a routing-only node, skipped at linearization). - - When the diverging message is an assistant and exactly one eligible - *short-response leaf* sibling exists, adopt the rewritten message onto that - node and DEMOTE it to routing-only (clear its turn snapshot). Any other - mismatch (non-assistant, long response, non-leaf or ambiguous) forks as usual. - """ + """Absorb a short assistant-rewrite onto its node instead of forking -- reward + hygiene only; forking is already safe (a rewrite mounts as routing-only).""" if self._fork_threshold <= 0: return cur, i # feature off if i >= len(prompt_messages) or prompt_messages[i].get("role") != "assistant": @@ -440,17 +289,12 @@ def _try_merge_assistant_rewrite( c for c in cur.children if c.role == "assistant" - # Leaf == rewrite of the immediately-previous assistant: no later - # turn has extended it yet. A non-leaf assistant has already grown a - # subchain; merging onto it would tangle that history. and not c.children - # A real turn leaf carrying this turn's snapshot, not an already- - # demoted routing node (turn_prompt_ids cleared by a prior merge). - and c.turn_prompt_ids is not None and len(c.turn_response_ids or []) < self._fork_threshold + and c.turn_prompt_ids is not None + and len(c.turn_response_ids or []) < self._fork_threshold ] if len(candidates) != 1: if len(candidates) >= 2: - # Ambiguous: don't pick arbitrarily — fork as usual and hint. logger.warning( "append_turn(sid=%s turn=%s): %d eligible rewrite-merge " "candidates; forking instead (ambiguous mixed state).", @@ -461,21 +305,16 @@ def _try_merge_assistant_rewrite( return cur, i sib = candidates[0] - sib.metadata["merged_rewrite"] = { # observability breadcrumb only + sib.metadata["merged_rewrite"] = { "abandoned_turn_index": sib.turn_index, "abandoned_response_tokens": len(sib.turn_response_ids or []), } - # Demote to routing-only: snapshot cleared -> skipped by the - # ``turn_prompt_ids is not None`` filter at linearization, and never - # re-selected as a merge candidate on a later turn. + sib.turn_prompt_ids = None sib.turn_response_ids = None sib.turn_response_logprobs = None sib.turn_finish_reason = None sib.turn_index = None - # Adopt the rewritten message; the match_key cache MUST follow messages - # so a later turn's DFS descends through this (now rewritten) node - # instead of forking again. sib.messages = [prompt_messages[i]] sib.match_key = node_match_key(sib.messages) return sib, i + 1 @@ -486,13 +325,6 @@ def _mount_prompt_messages( remaining_messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None, ) -> Node: - """Attach each remaining prompt message as a routing node under ``cur``. - - One node per message; token attribution happens at get_trajectory time, not - here. The tools metadata is placed only on the FIRST system node on the path - (``_first_system_already_set`` walks ``cur → root``; ``cur`` is the deepest - mounted node, so the walk sees every already-mounted ancestor). - """ for m in remaining_messages: role = m.get("role") md: dict[str, Any] = {} @@ -510,7 +342,6 @@ def _attach_assistant_leaf( response_message: dict[str, Any] | None, metadata: dict[str, Any] | None, ) -> None: - """Attach this turn's assistant leaf carrying the sglang snapshot.""" asst = Node( role="assistant", messages=[response_message] if response_message is not None else [], @@ -532,16 +363,6 @@ def _chain_to_sample( extra_metadata: dict[str, Any] | None, claimed: set[int], ) -> list[Sample]: - """Linearize one root→leaf chain into >=1 Samples (see module docstring). - - Each turn either grows the current segment or forks a new one (case - A/B1-too-long/B2). ``claimed`` deduplicates snapshot nodes shared across - sibling leaves; a leaf's terminal node is never pre-claimed, so every leaf - keeps >=1 trained turn. Reward is left at 0.0 and assigned by the caller. - """ - # Only assistant leaves carrying this turn's sglang snapshot participate. - # Routing assistant nodes mounted from prior-turn replay (turn_prompt_ids - # is None) carry no token signal and are skipped. asst_chain = [n for n in chain if n.role == "assistant" and n.turn_prompt_ids is not None] segments: list[_Segment] = [] @@ -550,11 +371,14 @@ def _chain_to_sample( prompt_ids = asst.turn_prompt_ids or [] drift = seg.measure_drift(prompt_ids) if not seg.can_absorb_drift(drift, self._fork_threshold): - segments.append(seg) # fork: close this segment, start fresh here + segments.append(seg) seg = _Segment() drift = 0 trained = id(asst) not in claimed claimed.add(id(asst)) + # A snapshot node can be shared by sibling leaves (same Node object). + # Train it only on the first leaf to reach it; later leaves re-emit it + # as loss=0 context so the shared prefix isn't double-counted. seg.absorb_turn( drift, prompt_ids, @@ -565,8 +389,7 @@ def _chain_to_sample( ) segments.append(seg) - # Drop empty / fully-masked segments: an in-chain fork can isolate a run of - # turns all claimed by an earlier sibling leaf (no loss=1 token), which would + # Drop fully-masked segments (all turns claimed by an earlier leaf) -- they'd # trip the downstream "not fully masked" assert. A leaf's terminal turn is # never pre-claimed, so its final segment always survives. return [ @@ -592,18 +415,6 @@ def _build_leaf_sample( logprobs: list[float], strip: int, ) -> Sample: - """Build one Sample from a linearized token segment. - - ``loss_mask`` / ``logprobs`` are clamped to the response region (``strip`` - drops the leading first-turn prompt prefix) per the slime contract: - ``response_length == len(loss_mask)``, covering only the response region - (see backends/megatron_utils/data.py:139, ray/rollout.py:695). ``reward`` - is left at 0.0; the caller assigns the per-sample share. - - Per-row dataset metadata and the per-turn tool / finish_reason snapshot are - intentionally NOT propagated here (dump/analysis tooling reads them off the - tree nodes); only ``extra_metadata`` rides along. - """ loss_resp, lp_resp = loss_mask[strip:], logprobs[strip:] metadata = dict(extra_metadata or {}) return Sample( @@ -622,7 +433,6 @@ def _build_leaf_sample( @staticmethod def _first_system_already_set(start: Node) -> bool: - """Walk start->root looking for a system node already carrying tools.""" cur: Node | None = start while cur is not None and not cur.is_root: if cur.role == "system" and cur.metadata.get("tools") is not None: From 6565fe1fc88d9bd56251c529b3d390d85c43b252 Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Tue, 9 Jun 2026 03:05:03 +0000 Subject: [PATCH 27/28] refactor(agent): slim adapters and trajectory_manager, add e2e test Trim dead code from generate.py and the anthropic/openai/common adapters, streamline trajectory_manager linearization, and add an end-to-end trajectory manager test. --- examples/coding_agent_rl/generate.py | 78 +- slime/agent/adapters/anthropic.py | 27 - slime/agent/adapters/common.py | 38 +- slime/agent/adapters/openai.py | 25 - slime/agent/trajectory_manager.py | 45 +- .../test_agent/test_trajectory_manager_e2e.py | 1424 +++++++++++++++++ 6 files changed, 1503 insertions(+), 134 deletions(-) create mode 100644 tests/test_agent/test_trajectory_manager_e2e.py diff --git a/examples/coding_agent_rl/generate.py b/examples/coding_agent_rl/generate.py index 8123dc50e7..8624f8cd7b 100644 --- a/examples/coding_agent_rl/generate.py +++ b/examples/coding_agent_rl/generate.py @@ -9,8 +9,8 @@ 1. ``sandbox.run_claude_code`` prepares the agent sandbox and runs claude-code. 2. ``sandbox.git_diff`` captures the model-produced patch. 3. ``sandbox.evaluate`` scores that patch in a second clean sandbox. - 4. ``_merge_samples`` combines reward + the ``list[Sample]`` returned by - ``adapter.finish_session(sid)``. + 4. ``adapter.finish_session`` drains the session tree into reward-weighted + ``Sample`` objects with ``.response`` already decoded; ``generate`` logs. All sandbox-side details live in ``sandbox.py``; the LLM plumbing (Anthropic <-> SGLang /generate, token capture, 3-kind segment split) uses @@ -49,7 +49,6 @@ import secrets import time import traceback -from dataclasses import dataclass from typing import Any from slime.agent.adapters import AnthropicAdapter @@ -96,8 +95,7 @@ def __init__(self, args) -> None: "Without it the sandbox cannot dial back and the rollout will " "silently abort." ) - fork_merge_threshold = os.environ.get("SLIME_FORK_MERGE_MAX_RESPONSE_TOKENS") - fork_merge_threshold = int(fork_merge_threshold) if fork_merge_threshold else None + fork_merge_threshold = int(v) if (v := os.environ.get("SLIME_FORK_MERGE_MAX_RESPONSE_TOKENS")) else None self.adapter = AnthropicAdapter( tokenizer=self.tokenizer, sglang_url=sglang_url, @@ -132,15 +130,8 @@ def __init__(self, args) -> None: # --------------------------------------------------------------------------- -# Trajectory -> Sample conversion +# Session setup # --------------------------------------------------------------------------- -@dataclass(frozen=True) -class RewardResult: - reward: float - is_solved: bool - applied_cleanly: bool - - def _start_session( state: _State, sample: Sample, @@ -166,42 +157,11 @@ def _start_session( return session_id -def _merge_samples( - *, - sample: Sample, - state: _State, - samples: list[Sample], - reward_result: RewardResult, - elapsed_sec: float, - instance_id: str, -) -> list[Sample]: - if not samples: - return _abort_result(sample, "adapter_session_empty") - - for s in samples: - rlen = int(s.response_length or 0) - if rlen and s.tokens: - s.response = state.tokenizer.decode(s.tokens[-rlen:], skip_special_tokens=False) - else: - s.response = "" - - logger.info( - "[coding_agent_rl] %s: reward=%.2f solved=%s applied=%s elapsed=%.1fs segments=%d", - instance_id, - reward_result.reward, - reward_result.is_solved, - reward_result.applied_cleanly, - elapsed_sec, - len(samples), - ) - return samples - - # --------------------------------------------------------------------------- # Main per-sample agent function # # The four calls inside the timeout are the high-level rollout recipe: -# run_claude_code -> git_diff -> sandbox.evaluate -> merge_samples. +# run_claude_code -> git_diff -> sandbox.evaluate -> finish_session. # --------------------------------------------------------------------------- async def generate(args, sample: Sample, sampling_params: dict[str, Any]): """Per-sample agent function with wall-clock guard. See @@ -238,24 +198,26 @@ async def generate(args, sample: Sample, sampling_params: dict[str, Any]): pre_commands=md["pre_commands"], timeout_sec=SWE_EVAL_TIMEOUT_SEC, ) - reward_result = RewardResult( - reward=float(reward), - is_solved=bool(is_solved), - applied_cleanly=bool(applied_cleanly), - ) samples = await state.adapter.finish_session( session_id, base_sample=sample, - reward=float(reward_result.reward), + reward=float(reward), ) - return _merge_samples( - sample=sample, - state=state, - samples=samples, - reward_result=reward_result, - elapsed_sec=time.time() - t0, - instance_id=instance_id, + if not samples: + return _abort_result(sample, "adapter_session_empty") + + # finish_session already linearized, reward-weighted and decoded + # each segment's .response; here we only log a summary. + logger.info( + "[coding_agent_rl] %s: reward=%.2f solved=%s applied=%s elapsed=%.1fs segments=%d", + instance_id, + float(reward), + bool(is_solved), + bool(applied_cleanly), + time.time() - t0, + len(samples), ) + return samples except asyncio.TimeoutError: _log_timeout_diagnostic(t0) diff --git a/slime/agent/adapters/anthropic.py b/slime/agent/adapters/anthropic.py index de6cbafafb..e01f4de55b 100644 --- a/slime/agent/adapters/anthropic.py +++ b/slime/agent/adapters/anthropic.py @@ -36,7 +36,6 @@ ) from slime.agent.parsing import ParsedModelOutput, parse_model_output from slime.agent.trajectory_manager import TrajectoryManager -from slime.utils.types import Sample logger = logging.getLogger(__name__) @@ -98,32 +97,6 @@ def __init__( self.app.router.add_get("/healthz", _ok) self.app.router.add_get("/v1/models", _ok) - async def finish_session( - self, - sid: str, - *, - base_sample: Sample | None = None, - reward: float = 0.0, - extra_metadata: dict[str, Any] | None = None, - wait_timeout: float = 5.0, - ) -> list[Sample]: - """Drain a session's trajectory into Sample objects. - - Waits out in-flight requests for ``sid``, then linearises the - per-sid tree via ``TrajectoryManager.get_trajectory``. Idempotent -- - a second call for an already-popped sid returns ``[]``. - """ - await self.shutdown_session(sid, wait_timeout=wait_timeout) - # Drop the per-sid adapter Session; the trajectory itself is in - # manager._trees and will be popped by get_trajectory(drop=True). - self.store.pop(sid, None) - return self.manager.get_trajectory( - sid, - base_sample=base_sample, - reward=reward, - extra_metadata=extra_metadata, - ) - # ============================================================================= # Translation (Anthropic wire <-> chat-template messages) diff --git a/slime/agent/adapters/common.py b/slime/agent/adapters/common.py index a302540683..40abadfa32 100644 --- a/slime/agent/adapters/common.py +++ b/slime/agent/adapters/common.py @@ -32,8 +32,11 @@ class BaseAdapter: """Base HTTP adapter with per-instance session lifecycle state.""" session_cls: type + # Set by subclass __init__: the shared TrajectoryManager keyed by sid. + manager: Any def __init__(self, *, tokenizer, sglang_url, tool_parser=None, reasoning_parser=None) -> None: + self.tokenizer = tokenizer self.store: dict[str, Any] = {} self.inflight: dict[str, set[asyncio.Task]] = {} self.closed: set[str] = set() @@ -62,8 +65,39 @@ def open_session( async def shutdown_session(self, sid: str, *, wait_timeout: float = 5.0) -> None: await shutdown_session_tasks(sid, self.closed, self.inflight, wait_timeout=wait_timeout) - async def finish_session(self, sid: str, *, wait_timeout: float = 5.0) -> list: - raise NotImplementedError + async def finish_session( + self, + sid: str, + *, + base_sample=None, + reward: float = 0.0, + extra_metadata: dict | None = None, + wait_timeout: float = 5.0, + ) -> list: + """Drain a session's trajectory into fully-formed Sample objects. + + Waits out in-flight requests for ``sid``, linearises the per-sid tree + via ``TrajectoryManager.get_trajectory``, then decodes each sample's + trained tail into ``.response`` (the manager is tokenizer-free, so the + adapter that owns the tokenizer fills this in). Idempotent -- a second + call for an already-popped sid returns ``[]``. + """ + await self.shutdown_session(sid, wait_timeout=wait_timeout) + # Drop the per-sid adapter Session; the trajectory itself lives in the + # manager's per-sid tree and is popped by get_trajectory(drop=True). + self.store.pop(sid, None) + samples = self.manager.get_trajectory( + sid, + base_sample=base_sample, + reward=reward, + extra_metadata=extra_metadata, + ) + for s in samples: + rlen = int(s.response_length or 0) + s.response = ( + self.tokenizer.decode(s.tokens[-rlen:], skip_special_tokens=False) if rlen and s.tokens else "" + ) + return samples def request_session_id( diff --git a/slime/agent/adapters/openai.py b/slime/agent/adapters/openai.py index 99f1c12f03..6ff89d7004 100644 --- a/slime/agent/adapters/openai.py +++ b/slime/agent/adapters/openai.py @@ -43,7 +43,6 @@ ) from slime.agent.parsing import ParsedModelOutput, parse_model_output from slime.agent.trajectory_manager import TrajectoryManager -from slime.utils.types import Sample logger = logging.getLogger(__name__) @@ -104,30 +103,6 @@ def __init__( self.app.router.add_get("/healthz", _ok) self.app.router.add_get("/v1/models", _ok) - async def finish_session( - self, - sid: str, - *, - base_sample: Sample | None = None, - reward: float = 0.0, - extra_metadata: dict[str, Any] | None = None, - wait_timeout: float = 5.0, - ) -> list[Sample]: - """Drain a session's trajectory into Sample objects. - - Waits out in-flight requests for ``sid``, then linearises the - per-sid tree via ``TrajectoryManager.get_trajectory``. Idempotent -- - a second call for an already-popped sid returns ``[]``. - """ - await self.shutdown_session(sid, wait_timeout=wait_timeout) - self.store.pop(sid, None) - return self.manager.get_trajectory( - sid, - base_sample=base_sample, - reward=reward, - extra_metadata=extra_metadata, - ) - # ============================================================================= # Translation (OpenAI wire <-> chat-template messages) diff --git a/slime/agent/trajectory_manager.py b/slime/agent/trajectory_manager.py index 61e40c302c..5ca044a594 100644 --- a/slime/agent/trajectory_manager.py +++ b/slime/agent/trajectory_manager.py @@ -104,20 +104,20 @@ class _Segment: """Token accumulator for one coherent linearized run within a chain.""" def __init__(self) -> None: - self.buffer: list[int] = [] - self.loss: list[int] = [] + self.tokens: list[int] = [] + self.loss_mask: list[int] = [] self.logprobs: list[float] = [] - # Each span: (start, end, turn_index) over self.buffer, half-open. + # Each span: (start, end, turn_index) over self.tokens, half-open. self.spans: list[tuple[int, int, int | None]] = [] self.first_prompt_len: int = 0 - def measure_drift(self, prompt_ids: list[int]) -> int: - return len(self.buffer) - _lcp_len(self.buffer, prompt_ids) + def measure_token_drift(self, prompt_ids: list[int]) -> int: + return len(self.tokens) - _lcp_len(self.tokens, prompt_ids) def can_absorb_drift(self, drift: int, fork_threshold: int) -> bool: if drift == 0: return True - realign_at = len(self.buffer) - drift + realign_at = len(self.tokens) - drift if not self.spans or realign_at < self.spans[-1][0]: return False return fork_threshold > 0 and drift < fork_threshold @@ -133,9 +133,9 @@ def absorb_turn( trained: bool = True, ) -> None: """Append one turn, dropping any re-tokenization drift first so logprobs stay aligned.""" - realign_at = len(self.buffer) - drift - del self.buffer[realign_at:] - del self.loss[realign_at:] + realign_at = len(self.tokens) - drift + del self.tokens[realign_at:] + del self.loss_mask[realign_at:] del self.logprobs[realign_at:] if self.spans and realign_at < self.spans[-1][1]: s, _e, j = self.spans[-1] @@ -145,26 +145,26 @@ def absorb_turn( is_first_turn = not self.spans tail = prompt_ids[realign_at:] - self.buffer.extend(tail) - self.loss.extend([0] * len(tail)) + self.tokens.extend(tail) + self.loss_mask.extend([0] * len(tail)) self.logprobs.extend([0.0] * len(tail)) - start = len(self.buffer) - self.buffer.extend(response_ids) - self.loss.extend([1 if trained else 0] * len(response_ids)) + start = len(self.tokens) + self.tokens.extend(response_ids) + self.loss_mask.extend([1 if trained else 0] * len(response_ids)) self.logprobs.extend( response_logprobs if (trained and response_logprobs is not None) else [0.0] * len(response_ids) ) - self.spans.append((start, len(self.buffer), turn_index)) + self.spans.append((start, len(self.tokens), turn_index)) if is_first_turn: self.first_prompt_len = len(prompt_ids) # stripped at build time def response_strip(self) -> int: - return min(self.first_prompt_len, len(self.loss)) + return min(self.first_prompt_len, len(self.loss_mask)) def has_trained_response(self) -> bool: - return any(self.loss[self.response_strip() :]) + return any(self.loss_mask[self.response_strip() :]) # =========================================================================== @@ -369,7 +369,7 @@ def _chain_to_sample( seg = _Segment() for asst in asst_chain: prompt_ids = asst.turn_prompt_ids or [] - drift = seg.measure_drift(prompt_ids) + drift = seg.measure_token_drift(prompt_ids) if not seg.can_absorb_drift(drift, self._fork_threshold): segments.append(seg) seg = _Segment() @@ -396,13 +396,13 @@ def _chain_to_sample( self._build_leaf_sample( base_sample=base_sample, extra_metadata=extra_metadata, - tokens=seg.buffer, - loss_mask=seg.loss, + tokens=seg.tokens, + loss_mask=seg.loss_mask, logprobs=seg.logprobs, strip=seg.response_strip(), ) for seg in segments - if seg.buffer and seg.has_trained_response() + if seg.tokens and seg.has_trained_response() ] def _build_leaf_sample( @@ -419,7 +419,8 @@ def _build_leaf_sample( metadata = dict(extra_metadata or {}) return Sample( index=base_sample.index, - group_id=base_sample.group_id if base_sample.group_id is not None else base_sample.index, + group_index=base_sample.group_index, + rollout_id=base_sample.rollout_id if base_sample.rollout_id is not None else base_sample.index, prompt=base_sample.prompt, label=base_sample.label, tokens=list(tokens), diff --git a/tests/test_agent/test_trajectory_manager_e2e.py b/tests/test_agent/test_trajectory_manager_e2e.py new file mode 100644 index 0000000000..f5871eeb47 --- /dev/null +++ b/tests/test_agent/test_trajectory_manager_e2e.py @@ -0,0 +1,1424 @@ +"""End-to-end tests for TrajectoryManager via append_turn / get_trajectory. + +This script drives the two public interfaces of +``slime.agent.trajectory_manager.TrajectoryManager`` and exhaustively covers the +ways a trajectory can branch, organized as a two-axis matrix: + + * LAYER 1 — routing tree (append_turn). DFS merges on (role, node_match_key) + only, so MESSAGE IDENTITY决定 tree shape; token ids are irrelevant here. + * LAYER 2 — linearization (get_trajectory). TOKEN-ID prefix决定 how each leaf + chain becomes Samples (clean continuation / drift case A·B1·B2 / cross-leaf + dedup / reward split). + * COMBINED — both layers interacting (rewrite-merge, tree-fork + token-drift + stacked, deep multi-leaf dedup, long mixed session). + +Readability: + Token ids are SEMANTIC small integers (see TOKEN_NAMES). Each message renders + to ``[START, ...body, END]`` with a per-role band, so an id like 2001 reads as + ``u:compute`` and 7001 reads as ````. Expected token sequences are built + with the same render_* helpers used to feed append_turn, never hand-typed + magic numbers. + +Dual mode: + Every case is a ``test_*`` function doing strict assertions. ``main()`` runs + them all and, after each, prints the routing tree (token ids decoded to names) + and every linearized Sample with token / loss_mask aligned, so a human can read + exactly where each branch happened. Run with:: + + python -m tests.test_agent.test_trajectory_manager_e2e +""" + +from __future__ import annotations + +from tests.test_agent.test_claude_code_agent._dump_helpers import dump_tree_txt # noqa: E402 + +from slime.agent.adapters.common import TurnRecord # noqa: E402 +from slime.agent.trajectory_manager import TrajectoryManager, _lcp_len # noqa: E402 +from slime.utils.types import Sample # noqa: E402 + +# =========================================================================== +# §1 Semantic token vocabulary + reverse table +# =========================================================================== +# +# Per-role band. A message renders to [START, ...body, END]; the generation +# prompt appends the assistant START as the open-turn marker. + +_BANDS = { + "system": 1000, + "user": 2000, + "assistant": 9000, + "tool": 3000, +} +_GEN = _BANDS["assistant"] # add_generation_prompt marker +_DRIFT_BAND = 7000 + +# Reverse table: token id -> human-readable name. Filled lazily as messages are +# registered so dumps translate ids back to labels. +TOKEN_NAMES: dict[int, str] = {} +_ABBR = {"system": "sys", "user": "usr", "assistant": "ast", "tool": "tul"} +for _role, _base in _BANDS.items(): + TOKEN_NAMES[_base] = f"<{_ABBR[_role]}>" + TOKEN_NAMES[_base + 9] = f"" +TOKEN_NAMES[_GEN] = "" + + +def name_of(tok: int) -> str: + """Human-readable name for a token id (falls back to the raw int).""" + return TOKEN_NAMES.get(tok, str(tok)) + + +def _vis(label: str) -> str: + """Make whitespace visible in a token label for the dump. + + Whitespace-only drift (e.g. a trailing space from a cc rewrite) is invisible + in a terminal, which makes ``r:ok`` vs ``r:ok `` indistinguishable. Render + spaces as ``␣`` so the difference is obvious in the readable output. + """ + return label.replace(" ", "␣") + + +_ASST_BODY: dict[str, int] = {} + + +def _asst_body(label: str) -> int: + """Stable assistant body token for a response/message label. + + An assistant message replayed in a later prompt must render to the SAME + tokens the model generated for it, otherwise a clean continuation can never + hold (the cumulative prompt+response would not prefix the next prompt). So + both ``render_response`` and an assistant ``MsgTok`` derive their body token + from this one function, keyed on the label. Bodies are assigned by a stable + per-label counter (NOT a hash) so distinct labels never collide on one id — + a collision would mislabel tokens in the dump and could spuriously match + across turns. + """ + if label not in _ASST_BODY: + body = _BANDS["assistant"] + 100 + len(_ASST_BODY) + _ASST_BODY[label] = body + TOKEN_NAMES[body] = f"r:{_vis(label)}" + return _ASST_BODY[label] + + +def render_ids(ids: list[int]) -> str: + """Decode an id list into a space-joined readable string.""" + return " ".join(name_of(t) for t in ids) + + +class MsgTok: + """A message bound to a fixed, deterministic token rendering. + + The same MsgTok always renders to the same token segment regardless of which + turn replays it (a clean tokenizer). Token-id drift is injected explicitly by + tests via ``drift`` — never by re-rendering. + """ + + _body_counter: dict[str, int] = {} + + def __init__(self, role: str, label: str) -> None: + self.role = role + self.label = label + base = _BANDS[role] + if role == "assistant": + # An assistant message must render to the same body token as the + # response it represents (label-keyed), so a replayed assistant in a + # later prompt token-matches the original generation -> clean + # continuation. See _asst_body. + self.body = _asst_body(label) + else: + # Allocate one stable body token per (role, label). Offset past the + # END marker (base+9): the counter is shared across cases, so bodies + # must never climb into base+9 (END) or they'd collide with it. + idx = MsgTok._body_counter.setdefault(role, 0) + 1 + MsgTok._body_counter[role] = idx + self.body = base + 10 + idx + TOKEN_NAMES[self.body] = f"{role}:{_vis(label)}" + # message dict as the manager sees it (drives node_match_key). + self.message = {"role": role, "content": label} + + def render(self) -> list[int]: + """[START, body, END] for this message.""" + base = _BANDS[self.role] + return [base, self.body, base + 9] + + +def sys_msg(label: str) -> MsgTok: + return MsgTok("system", label) + + +def usr_msg(label: str) -> MsgTok: + return MsgTok("user", label) + + +def asst_msg(label: str) -> MsgTok: + return MsgTok("assistant", label) + + +def tool_msg(label: str) -> MsgTok: + return MsgTok("tool", label) + + +def render_prompt(msgs: list[MsgTok]) -> list[int]: + """Render a prompt message list, appending the generation-prompt marker.""" + out: list[int] = [] + for m in msgs: + out.extend(m.render()) + out.append(_GEN) + return out + + +def render_response(label: str) -> list[int]: + """Render an assistant response: [body, ]. + + The generation-prompt marker ```` equals the assistant START token, so + `` + render_response(x)`` == the assistant message ``[, body, + ]`` replayed in a later prompt. That identity is what makes a clean + continuation hold across turns. + """ + return [_asst_body(label), _BANDS["assistant"] + 9] + + +def messages(msgs: list[MsgTok]) -> list[dict]: + """The plain message dicts append_turn wants for prompt_messages.""" + return [m.message for m in msgs] + + +def drift(ids: list[int], at: int, sentinel: int = _DRIFT_BAND + 1) -> list[int]: + """Return a copy of ``ids`` with a sentinel spliced at index ``at``. + + The sentinel sits in the drift band (7000+), so a dump shows ```` at + the exact divergence point. Splicing (insert) makes ``len`` grow by one, + which is enough to make the lcp diverge at ``at``. + """ + TOKEN_NAMES[sentinel] = "" + return ids[:at] + [sentinel] + ids[at:] + + +def drift_replace(ids: list[int], at: int, sentinel: int = _DRIFT_BAND + 2) -> list[int]: + """Return a copy of ``ids`` with the token at ``at`` REPLACED by a sentinel. + + Unlike ``drift`` this keeps length constant — used when a test wants the + divergence inside a response span without changing the cumulative length. + """ + TOKEN_NAMES[sentinel] = "" + out = list(ids) + out[at] = sentinel + return out + + +def turn(prompt_ids, response_ids, *, finish_reason="stop", logprobs=None) -> TurnRecord: + return TurnRecord( + prompt_ids=list(prompt_ids), + output_ids=list(response_ids), + finish_reason=finish_reason, + output_log_probs=list(logprobs) if logprobs is not None else [], + ) + + +# A scratch space for the dual-mode printer: each case appends (title, mgr, sid, +# samples) so main() can render after the assertions pass. +_PRINT_LOG: list[tuple[str, object, str, list]] = [] + +# Raw append_turn inputs, keyed by sid, captured at call time so the printer can +# show the SOURCE data (prompt_ids / response_ids / finish / logprobs) that fed +# the tree — before any tree-building or linearization happened. +_TURN_LOG: dict[str, list[dict]] = {} + + +def _record(title: str, mgr, sid: str, samples: list) -> None: + _PRINT_LOG.append((title, mgr, sid, samples)) + + +# Convenience: append a turn with semantic messages, auto-rendering prompt unless +# an explicit prompt_ids is supplied (for drift injection). +def append( + mgr: TrajectoryManager, + sid: str, + prompt_msgs: list[MsgTok], + response_label: str | None, + *, + prompt_ids=None, + response_ids=None, + finish_reason="stop", + logprobs=None, + tools=None, + response_message=None, +): + p = list(prompt_ids) if prompt_ids is not None else render_prompt(prompt_msgs) + if response_ids is not None: + r = list(response_ids) + elif response_label is not None: + r = render_response(response_label) + else: + r = [] + rmsg = response_message + if rmsg is None and response_label is not None: + rmsg = {"role": "assistant", "content": response_label} + lp = logprobs + # Capture the raw turn inputs for the human-readable dump before the manager + # consumes them. + _TURN_LOG.setdefault(sid, []).append( + { + "prompt_msgs": [f"{m.role}:{_vis(m.label)}" for m in prompt_msgs], + "prompt_ids": p, + "response_ids": r, + "finish": finish_reason, + "has_lp": lp is not None, + } + ) + mgr.append_turn( + sid, + turn=turn(p, r, finish_reason=finish_reason, logprobs=lp), + prompt_messages=messages(prompt_msgs), + tools=tools, + response_message=rmsg, + ) + return p, r + + +def _leaves(mgr, sid): + return [leaf for leaf in mgr._trees[sid].leaves() if not leaf.is_root] + + +def _iter_all(root): + """Yield every non-root node in the tree (pre-order).""" + stack = list(root.children) + while stack: + n = stack.pop() + yield n + stack.extend(n.children) + + +# A minimal OpenAI-shape tool spec, used to exercise tools-metadata routing. +TOOLS = [ + { + "type": "function", + "function": {"name": "run", "description": "Run.", "parameters": {"type": "object"}}, + } +] + + +# Tree text snapshot captured the instant before get_trajectory drains the sid, +# so the human-readable dump can show [tree] AND [samples] side by side even +# though get_trajectory consumes the session. +_TREE_SNAP: dict[str, str] = {} + +# Input reward passed to get_trajectory, keyed by sid, so the dump can show the +# split (input_reward / n_samples == per_sample_reward) explicitly. +_REWARD_IN: dict[str, float] = {} + + +def get_traj(mgr, sid, *args, **kwargs): + """get_trajectory wrapper that snapshots the tree before draining. + + Linearization (get_trajectory) pops the sid, so a later dump would only see + ````. Capturing the tree text here keeps the routing tree visible + next to the Samples it produced. The input ``reward`` is captured too so the + dump can show how it splits across the emitted samples. + """ + if mgr.has_session(sid): + _TREE_SNAP[sid] = dump_tree_txt(mgr, sid) + _REWARD_IN[sid] = kwargs.get("reward", 0.0) + samples = mgr.get_trajectory(sid, *args, **kwargs) + # Reward conservation: get_trajectory splits the input reward evenly across + # every emitted sample, so the per-sample shares must sum back to the input + # (modulo float error). This is the "averaged over sample count" invariant. + if samples: + total = sum(s.reward for s in samples) + assert abs(total - _REWARD_IN[sid]) < 1e-9, ( + "reward not conserved across split", + total, + _REWARD_IN[sid], + ) + return samples + + +def _check_invariants(samples): + for s in samples: + assert len(s.loss_mask) == len(s.rollout_log_probs) == s.response_length, ( + "alignment broken", + len(s.loss_mask), + len(s.rollout_log_probs), + s.response_length, + ) + assert sum(s.loss_mask) > 0, "fully-masked sample emitted" + + +def golden(sample) -> str: + """Render one Sample as a human-reviewable golden string. + + Every token is decoded to its readable name (````, ``r:done``, + ```` ...). The leading prompt prefix (no loss_mask entry) is shown as + plain names; the response region is shown with each TRAINED token (loss=1) + wrapped in ``[...]`` and each context token (loss=0) left bare. This makes the + full linearized result — tokens, where the response region starts, and + exactly which tokens carry training signal — a single literal a human can + eyeball and assert against, instead of hand-derived index arithmetic. + + Example: `` system:S user:u [r:ok] []`` + """ + toks = sample.tokens + resp_start = len(toks) - sample.response_length + parts: list[str] = [] + for i, t in enumerate(toks): + nm = name_of(t) + if i >= resp_start and sample.loss_mask[i - resp_start] == 1: + parts.append(f"[{nm}]") + else: + parts.append(nm) + return " ".join(parts) + + +def goldens(samples) -> list[str]: + return [golden(s) for s in samples] + + +# =========================================================================== +# §2 Group 1 — routing tree layer (append_turn shapes the tree) +# =========================================================================== + + +def test_1_1_single_turn_chain(): + mgr = TrajectoryManager() + sid = "1.1" + s = sys_msg("S") + u = usr_msg("compute") + p, r = append(mgr, sid, [s, u], "ok") + chain = _leaves(mgr, sid)[0].path_from_root() + assert [n.role for n in chain] == ["system", "user", "assistant"] + assert chain[-1].turn_index == 1 + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + assert goldens(samples) == [ + " system:S user:compute [r:ok] []", + ] + _check_invariants(samples) + _record("1.1 single turn -> linear chain", mgr, sid, samples) + print("PASS 1.1") + + +def test_1_2_clean_multiturn_with_tool(): + mgr = TrajectoryManager() + sid = "1.2" + s, u = sys_msg("S"), usr_msg("compute") + a1, t1 = asst_msg("call"), tool_msg("4") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls") + p2, r2 = append(mgr, sid, [s, u, a1, t1], "done") + chain = _leaves(mgr, sid)[0].path_from_root() + assert [n.role for n in chain] == ["system", "user", "assistant", "tool", "assistant"] + assert mgr.turn_count(sid) == 2 + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + assert goldens(samples) == [ + " system:S user:compute [r:call] [] " + " tool:4 [r:done] []", + ] + _check_invariants(samples) + _record("1.2 clean 2-turn with tool -> single chain", mgr, sid, samples) + print("PASS 1.2") + + +def test_1_3_system_fork(): + mgr = TrajectoryManager() + sid = "1.3" + for sl in ["SA", "SB"]: + append(mgr, sid, [sys_msg(sl), usr_msg("u")], "a") + root = mgr._trees[sid] + assert len(root.children) == 2, "different system -> two subtrees at root" + assert len(_leaves(mgr, sid)) == 2 + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert goldens(samples) == [ + " system:SA user:u [r:a] []", + " system:SB user:u [r:a] []", + ] + _check_invariants(samples) + _record("1.3 system fork -> two subtrees at root", mgr, sid, samples) + print("PASS 1.3") + + +def test_1_4_user_fork_shared_system(): + mgr = TrajectoryManager() + sid = "1.4" + s = sys_msg("S") + for ul in ["A", "B"]: + append(mgr, sid, [s, usr_msg(ul)], ul.lower()) + root = mgr._trees[sid] + assert len(root.children) == 1, "system shared" + assert len(root.children[0].children) == 2, "user level forks" + assert len(_leaves(mgr, sid)) == 2 + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert goldens(samples) == [ + " system:S user:A [r:a] []", + " system:S user:B [r:b] []", + ] + _check_invariants(samples) + _record("1.4 user fork (shared system)", mgr, sid, samples) + print("PASS 1.4") + + +def test_1_5_assistant_message_fork(): + """Same (sys,user) prefix, two distinct assistant turns -> assistant fork.""" + mgr = TrajectoryManager() + sid = "1.5" + s, u = sys_msg("S"), usr_msg("u") + append(mgr, sid, [s, u], "a1") + append(mgr, sid, [s, u], "a2") + user_node = mgr._trees[sid].children[0].children[0] + assert len(user_node.children) == 2, "two assistant leaves hang off shared user" + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + # Two independent single-turn leaves sharing only the (sys,user) prefix. + assert goldens(samples) == [ + " system:S user:u [r:a1] []", + " system:S user:u [r:a2] []", + ] + _check_invariants(samples) + _record("1.5 assistant fork under shared user", mgr, sid, samples) + print("PASS 1.5") + + +def test_1_6_tool_fork_shared_assistant(): + """Same first assistant turn, two different tool results -> tool-level fork, + making the first assistant a shared snapshot node with 2 children.""" + mgr = TrajectoryManager() + sid = "1.6" + s, u, a1 = sys_msg("S"), usr_msg("u"), asst_msg("call") + append(mgr, sid, [s, u], "call", finish_reason="tool_calls") + append(mgr, sid, [s, u, a1, tool_msg("x")], "ax") + append(mgr, sid, [s, u, a1, tool_msg("y")], "ay") + asst1 = mgr._trees[sid].children[0].children[0].children[0] + assert asst1.role == "assistant" and asst1.turn_prompt_ids is not None + assert len(asst1.children) == 2, "shared assistant forks at the tool level" + assert len(_leaves(mgr, sid)) == 2 + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + # Leaf X owns the shared turn 1 (r:call trained); leaf Y shares it -> r:call + # demoted to loss=0 context, only r:ay trains (cross-leaf dedup). + assert goldens(samples) == [ + " system:S user:u [r:call] [] " + " tool:x [r:ax] []", + " system:S user:u r:call " " tool:y [r:ay] []", + ] + _check_invariants(samples) + _record("1.6 tool fork (shared assistant snapshot)", mgr, sid, samples) + print("PASS 1.6") + + +def test_1_7_token_only_drift_no_fork(): + """Identical messages, tampered prompt_ids -> NO tree fork (DFS ignores + tokens), but the drift DOES surface in the linearized sample: it lands in + leaf 2's prompt region (stripped / loss=0), proving token drift cannot + corrupt a trained response yet is still carried in the sample tokens.""" + mgr = TrajectoryManager() + sid = "1.7" + s, u = sys_msg("S"), usr_msg("u") + pa, _ = append(mgr, sid, [s, u], "a") + tampered = drift(pa, 1) # spliced into the prompt at index 1 + append(mgr, sid, [s, u], "b", prompt_ids=tampered) + # Tree: (sys,user) shared, two assistant turns hang off it -> two leaves; the + # path above the assistant is single (NOT forked on tokens). + user_node = mgr._trees[sid].children[0].children[0] + assert len(user_node.children) == 2, "two assistant turns share the (sys,user) path" + assert len(mgr._trees[sid].children) == 1 + + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 2 + # Leaf 1: clean. Leaf 2: the token sits in the stripped prompt region + # (bare, no brackets); the response r:b is fully trained ([...]). + assert goldens(samples) == [ + " system:S user:u [r:a] []", + " system:S user:u [r:b] []", + ] + # Belt-and-suspenders on the drift placement: token present, but never inside + # the response region. + s_b = samples[1] + assert (_DRIFT_BAND + 1) in s_b.tokens, "drift token is still carried in the sample" + assert (_DRIFT_BAND + 1) not in s_b.tokens[len(tampered) :], "drift not in the response region" + _check_invariants(samples) + _record("1.7 token-only drift -> no tree fork, drift lands in stripped prompt", mgr, sid, samples) + print("PASS 1.7") + + +def test_1_8_multi_tool_per_turn(): + mgr = TrajectoryManager() + sid = "1.8" + s, u, a1 = sys_msg("S"), usr_msg("u"), asst_msg("call") + ta, tb = tool_msg("A"), tool_msg("B") + append(mgr, sid, [s, u], "call", finish_reason="tool_calls") + append(mgr, sid, [s, u, a1, ta, tb], "done") + chain = _leaves(mgr, sid)[0].path_from_root() + assert [n.role for n in chain] == ["system", "user", "assistant", "tool", "tool", "assistant"] + assert chain[3].messages == [ta.message] + assert chain[4].messages == [tb.message] + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + assert goldens(samples) == [ + " system:S user:u [r:call] [] " + " tool:A tool:B [r:done] []", + ] + _check_invariants(samples) + _record("1.8 multi-tool turn -> one node per tool", mgr, sid, samples) + print("PASS 1.8") + + +def test_1_9_cross_sid_isolation(): + mgr = TrajectoryManager() + s = sys_msg("S") + for sid, ul in [("sid-a", "A"), ("sid-b", "B")]: + append(mgr, sid, [s, usr_msg(ul)], ul.lower()) + assert len(_leaves(mgr, "sid-a")) == 1 + assert len(_leaves(mgr, "sid-b")) == 1 + assert mgr._trees["sid-a"] is not mgr._trees["sid-b"] + sa = get_traj(mgr, "sid-a", base_sample=Sample(index=0, prompt=""), reward=1.0) + sb = get_traj(mgr, "sid-b", base_sample=Sample(index=1, prompt=""), reward=1.0) + assert goldens(sa) == [" system:S user:A [r:a] []"] + assert goldens(sb) == [" system:S user:B [r:b] []"] + _check_invariants(sa) + _check_invariants(sb) + _record("1.9 cross-sid isolation (sid-a)", mgr, "sid-a", sa) + _record("1.9 cross-sid isolation (sid-b)", mgr, "sid-b", sb) + print("PASS 1.9") + + +def test_1_10_empty_response(): + mgr = TrajectoryManager() + sid = "1.10" + s, u = sys_msg("S"), usr_msg("u") + append(mgr, sid, [s, u], None, response_ids=[], response_message=None, finish_reason="length") + asst = _leaves(mgr, sid)[0] + assert asst.role == "assistant" + assert asst.turn_response_ids == [] + assert asst.messages == [] + # Empty response -> the only turn has no trainable token, so its segment is + # dropped at linearization (no fully-masked sample). Zero samples is correct. + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 0 + _record("1.10 empty response -> assistant leaf, no message (0 samples)", mgr, sid, samples) + print("PASS 1.10") + + +# =========================================================================== +# §2 Group 2 — linearization layer (get_trajectory token routing) +# =========================================================================== + + +def test_2_1_single_turn_linearize(): + mgr = TrajectoryManager() + sid = "2.1" + s, u = sys_msg("S"), usr_msg("u") + p, r = append(mgr, sid, [s, u], "a", logprobs=None) + # attach explicit logprobs so we can check propagation + leaf = _leaves(mgr, sid)[0] + leaf.turn_response_logprobs = [-0.5] * len(r) + samples = get_traj(mgr, sid, base_sample=Sample(index=7, prompt="hi"), reward=1.0) + assert len(samples) == 1 + s0 = samples[0] + assert goldens(samples) == [" system:S user:u [r:a] []"] + assert s0.rollout_log_probs == [-0.5] * len(r) + assert s0.reward == 1.0 + _check_invariants(samples) + _record("2.1 single-turn linearize", mgr, sid, samples) + print("PASS 2.1") + + +def test_2_2_clean_multiturn_linearize(): + mgr = TrajectoryManager() + sid = "2.2" + s, u, a1, t1 = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("4") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) + p2, r2 = append(mgr, sid, [s, u, a1, t1], "done", logprobs=[-0.4] * 2) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + s0 = samples[0] + L = _lcp_len(p1 + r1, p2) + assert goldens(samples) == [ + " system:S user:u [r:call] [] " + " tool:4 [r:done] []", + ] + assert s0.rollout_log_probs == [-0.5] * len(r1) + [0.0] * (len(p2) - L) + [-0.4] * len(r2) + _check_invariants(samples) + _record("2.2 clean 2-turn linearize", mgr, sid, samples) + print("PASS 2.2") + + +def test_2_3_drift_case_A_forks(): + """Drift inside a PROMPT region -> case A -> fork, no token dropped.""" + mgr = TrajectoryManager() + sid = "2.3" + s, u, a1, t = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("t") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) + p2_honest = render_prompt([s, u, a1, t]) + p2 = drift(p2_honest, len(p1) - 1) # inside p1's prompt region + p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2, logprobs=[-0.4] * 2) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 2 + # case-A fork: two coherent single-turn segments; the token stays in + # segment 2's stripped prompt region (bare), no token dropped. + assert goldens(samples) == [ + " system:S user:u [r:call] []", + " system:S user:u r:call " + " tool:t [r:done] []", + ] + assert all(abs(s.reward - 1.0 / len(samples)) < 1e-9 for s in samples) + _check_invariants(samples) + _record("2.3 drift case A (prompt region) -> fork", mgr, sid, samples) + print("PASS 2.3") + + +def test_2_4_drift_case_B1_short_replaces(): + """Small drift inside the most-recent response span -> replace.""" + mgr = TrajectoryManager() # default threshold 1024 + sid = "2.4" + s, u, a1, t = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("t") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) + p2_honest = render_prompt([s, u, a1, t]) + assert p2_honest[: len(p1) + len(r1)] == p1 + r1 + drift_idx = len(p1) + len(r1) - 1 # last token of r1's echo + p2 = drift_replace(p2_honest, drift_idx) + p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2, logprobs=[-0.4] * 2) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + s0 = samples[0] + L = _lcp_len(p1 + r1, p2) + assert L == drift_idx + # replace: the drifted r:call tail is dropped and re-supplied as loss=0 prompt + # context (the token marks the divergence); only the surviving head of + # r:call and the new r:done train. + assert goldens(samples) == [ + " system:S user:u [r:call] " + " tool:t [r:done] []", + ] + assert s0.rollout_log_probs == [-0.5] * (L - len(p1)) + [0.0] * (len(p2) - L) + [-0.4] * len(r2) + _check_invariants(samples) + _record("2.4 drift case B1 (small) -> replace", mgr, sid, samples) + print("PASS 2.4") + + +def test_2_5_drift_case_B1_long_forks(): + mgr = TrajectoryManager(fork_threshold_tokens=1) # d>=1 -> fork + sid = "2.5" + s, u, a1, t = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("t") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls") + p2_honest = render_prompt([s, u, a1, t]) + p2 = drift_replace(p2_honest, len(p1) + len(r1) - 1) + p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 2 + # Both segments single-turn (the drift forked them apart): each trains its own + # response. The sits in segment 2's stripped prompt. + assert goldens(samples) == [ + " system:S user:u [r:call] []", + " system:S user:u r:call " + " tool:t [r:done] []", + ] + assert all(abs(s.reward - 1.0 / len(samples)) < 1e-9 for s in samples) + _check_invariants(samples) + _record("2.5 drift case B1 (long) -> fork", mgr, sid, samples) + print("PASS 2.5") + + +def test_2_6_drift_case_B1_threshold_zero_forks(): + mgr = TrajectoryManager(fork_threshold_tokens=0) + sid = "2.6" + s, u, a1, t = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("t") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls") + p2_honest = render_prompt([s, u, a1, t]) + p2 = drift_replace(p2_honest, len(p1) + len(r1) - 1) + p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 2 + assert goldens(samples) == [ + " system:S user:u [r:call] []", + " system:S user:u r:call " + " tool:t [r:done] []", + ] + _check_invariants(samples) + _record("2.6 drift case B1 threshold=0 -> fork", mgr, sid, samples) + print("PASS 2.6") + + +def test_2_7_drift_case_B2_earlier_turn_forks(): + """Drift inside an EARLIER turn's response span -> always fork.""" + mgr = TrajectoryManager() + sid = "2.7" + s, u = sys_msg("S"), usr_msg("u") + a1, t1 = asst_msg("a1"), tool_msg("t1") + a2, t2 = asst_msg("a2"), tool_msg("t2") + p1, r1 = append(mgr, sid, [s, u], "a1", finish_reason="tool_calls", logprobs=[-0.5] * 2) + p2, r2 = append(mgr, sid, [s, u, a1, t1], "a2", finish_reason="tool_calls", logprobs=[-0.4] * 2) + p3_honest = render_prompt([s, u, a1, t1, a2, t2]) + p3 = drift_replace(p3_honest, len(p1) + len(r1) - 1) # inside r1 (earlier span) + p3, r3 = append(mgr, sid, [s, u, a1, t1, a2, t2], "a3", prompt_ids=p3, logprobs=[-0.3] * 2) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 2 + # Segment 1 = clean turns 1+2; segment 2 = turn 3 alone (forked because the + # drift hit an EARLIER turn's response span, which replace can't drop). + assert goldens(samples) == [ + " system:S user:u [r:a1] [] " + " tool:t1 [r:a2] []", + " system:S user:u r:a1 " + " tool:t1 r:a2 tool:t2 [r:a3] []", + ] + assert all(abs(s.reward - 1.0 / len(samples)) < 1e-9 for s in samples) + _check_invariants(samples) + _record("2.7 drift case B2 (earlier turn) -> fork", mgr, sid, samples) + print("PASS 2.7") + + +def test_2_8_fork_reward_split(): + mgr = TrajectoryManager() + sid = "2.8" + s, u, a1, t = sys_msg("S"), usr_msg("u"), asst_msg("call"), tool_msg("t") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls") + p2_honest = render_prompt([s, u, a1, t]) + p2 = drift(p2_honest, len(p1) - 1) # prompt region -> case A fork + p2, r2 = append(mgr, sid, [s, u, a1, t], "done", prompt_ids=p2) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 2 + # case-A fork: two single-turn segments, each trains its own response. + assert goldens(samples) == [ + " system:S user:u [r:call] []", + " system:S user:u r:call " + " tool:t [r:done] []", + ] + # reward 1.0 split evenly across the 2 forked samples -> 0.5 each. + assert all(abs(s.reward - 0.5) < 1e-9 for s in samples) + _check_invariants(samples) + _record("2.8 fork reward split (1.0 / 2 = 0.5 each)", mgr, sid, samples) + print("PASS 2.8") + + +def test_2_9_two_leaves_reward_split(): + mgr = TrajectoryManager() + sid = "2.9" + s = sys_msg("S") + for ul in ["A", "B"]: + append(mgr, sid, [s, usr_msg(ul)], ul.lower()) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 2 + assert goldens(samples) == [ + " system:S user:A [r:a] []", + " system:S user:B [r:b] []", + ] + # reward 1.0 split evenly across the 2 leaves -> 0.5 each. + assert all(abs(s.reward - 0.5) < 1e-9 for s in samples) + _check_invariants(samples) + _record("2.9 two leaves reward split (1.0 / 2 = 0.5 each)", mgr, sid, samples) + print("PASS 2.9") + + +def test_2_10_cross_leaf_dedup(): + """Shared assistant trained on first leaf only; second leaf re-emits it + as loss=0 context.""" + mgr = TrajectoryManager() + sid = "2.10" + s, u, a1 = sys_msg("S"), usr_msg("u"), asst_msg("call") + tx, ty = tool_msg("x"), tool_msg("y") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) + p2, r2 = append(mgr, sid, [s, u, a1, tx], "a2", logprobs=[-0.4] * 2) + p3, r3 = append(mgr, sid, [s, u, a1, ty], "a3", logprobs=[-0.3] * 2) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 2 + s_second = samples[1] + # First leaf trains the shared r:call + its own r:a2; second leaf shares + # r:call (demoted to loss=0) and trains only r:a3. + assert goldens(samples) == [ + " system:S user:u [r:call] [] " + " tool:x [r:a2] []", + " system:S user:u r:call " " tool:y [r:a3] []", + ] + assert s_second.rollout_log_probs == [0.0] * (len(p3) - len(p1)) + [-0.3] * len(r3) + _check_invariants(samples) + _record("2.10 cross-leaf dedup (shared assistant trained once)", mgr, sid, samples) + print("PASS 2.10") + + +def test_2_11_routing_only_assistant_filtered(): + """cc replays an assistant the manager never recorded -> mounts routing-only, + must be filtered out of the strict-prefix walk (no raise).""" + mgr = TrajectoryManager() + sid = "2.11" + s, u = sys_msg("S"), usr_msg("u") + a1, t1 = asst_msg("a1"), tool_msg("t1") + a2 = asst_msg("a2") + append(mgr, sid, [s, u], "a1", finish_reason="tool_calls") + append(mgr, sid, [s, u, a1, t1], "a2", finish_reason="tool_calls") + foreign = asst_msg("foreign") + t2 = tool_msg("t2") + append(mgr, sid, [s, u, a1, t1, a2, foreign, t2], "a3") + leaves = _leaves(mgr, sid) + assert len(leaves) == 1 + chain = leaves[0].path_from_root() + routing = [n for n in chain if n.role == "assistant" and n.turn_prompt_ids is None] + assert len(routing) == 1 and routing[0].messages[0]["content"] == "foreign" + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + # The foreign assistant (r:foreign) is routing-only -> appears as bare context + # (no brackets); the three real turns r:a1/r:a2/r:a3 train. + assert goldens(samples) == [ + " system:S user:u [r:a1] [] " + " tool:t1 [r:a2] [] r:foreign " + " tool:t2 [r:a3] []", + ] + _record("2.11 routing-only assistant filtered (no raise)", mgr, sid, samples) + print("PASS 2.11") + + +def test_2_12_drop_clears_sid(): + mgr = TrajectoryManager() + sid = "2.12" + s, u = sys_msg("S"), usr_msg("u") + append(mgr, sid, [s, u], "a") + assert mgr.has_session(sid) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert goldens(samples) == [" system:S user:u [r:a] []"] + assert not mgr.has_session(sid) + assert mgr.get_trajectory(sid, base_sample=Sample(index=0, prompt="")) == [] + _check_invariants(samples) + _record("2.12 drop clears sid (2nd get_trajectory -> [])", mgr, sid, samples) + print("PASS 2.12") + + +# =========================================================================== +# §2 Group 3 — combined / stress (both layers interacting) +# =========================================================================== + + +def test_3_1_rewrite_merge_absorbs_short(): + mgr = TrajectoryManager() + sid = "3.1" + s, u = sys_msg("S"), usr_msg("u") + a1_rw = asst_msg("ok ") # cc-rewritten (different message identity) + t1 = tool_msg("t") + append(mgr, sid, [s, u], "ok", finish_reason="tool_calls", logprobs=[-0.5] * 2) + append(mgr, sid, [s, u, a1_rw, t1], "done", logprobs=[-0.4] * 2) + leaves = _leaves(mgr, sid) + assert len(leaves) == 1, "short rewrite absorbed, not forked" + chain = leaves[0].path_from_root() + merged = chain[2] + assert merged.turn_prompt_ids is None and merged.turn_index is None + assert merged.messages == [a1_rw.message] + assert merged.metadata["merged_rewrite"]["abandoned_turn_index"] == 1 + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + # The abandoned turn-1 response (r:ok␣) is demoted to routing-only -> appears + # bare; only the surviving turn-2 r:done trains. + assert goldens(samples) == [ + " system:S user:u r:ok␣ " " tool:t [r:done] []", + ] + _check_invariants(samples) + _record("3.1 rewrite-merge absorbs short assistant", mgr, sid, samples) + print("PASS 3.1") + + +def test_3_2_rewrite_merge_long_forks(): + mgr = TrajectoryManager(fork_threshold_tokens=1) # r1 len 2 >= 1 + sid = "3.2" + s, u = sys_msg("S"), usr_msg("u") + a1_rw, t1 = asst_msg("ok2 "), tool_msg("t") + append(mgr, sid, [s, u], "ok", finish_reason="tool_calls") + append(mgr, sid, [s, u, a1_rw, t1], "done") + assert len(_leaves(mgr, sid)) == 2, "long rewrite forks" + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 2 + # Leaf 1: the abandoned turn-1 standalone (r:ok). Leaf 2: turn-2 only (the + # rewritten r:ok2␣ assistant mounts routing-only and is filtered out). + assert goldens(samples) == [ + " system:S user:u [r:ok] []", + " system:S user:u r:ok2␣ " " tool:t [r:done] []", + ] + _check_invariants(samples) + _record("3.2 rewrite-merge long -> fork", mgr, sid, samples) + print("PASS 3.2") + + +def test_3_3_rewrite_merge_threshold_zero_forks(): + mgr = TrajectoryManager(fork_threshold_tokens=0) + sid = "3.3" + s, u = sys_msg("S"), usr_msg("u") + a1_rw, t1 = asst_msg("ok3 "), tool_msg("t") + append(mgr, sid, [s, u], "ok", finish_reason="tool_calls") + append(mgr, sid, [s, u, a1_rw, t1], "done") + assert len(_leaves(mgr, sid)) == 2 + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 2 + assert goldens(samples) == [ + " system:S user:u [r:ok] []", + " system:S user:u r:ok3␣ " " tool:t [r:done] []", + ] + _check_invariants(samples) + _record("3.3 rewrite-merge threshold=0 -> fork", mgr, sid, samples) + print("PASS 3.3") + + +def test_3_4_rewrite_merge_ambiguous_forks(): + mgr = TrajectoryManager() + sid = "3.4" + s, u = sys_msg("S"), usr_msg("u") + # two short assistant leaves under shared (sys,user) + append(mgr, sid, [s, u], "a") + append(mgr, sid, [s, u], "b") + a_c, t1 = asst_msg("c"), tool_msg("t") + append(mgr, sid, [s, u, a_c, t1], "d") + assert len(_leaves(mgr, sid)) == 3, "ambiguous candidates fork" + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 3 + # Leaves "a" and "b" are standalone single turns; leaf "d" carries the + # ambiguous-rewrite assistant (r:c) as routing-only (bare) -> trains only r:d. + assert goldens(samples) == [ + " system:S user:u [r:a] []", + " system:S user:u [r:b] []", + " system:S user:u r:c " " tool:t [r:d] []", + ] + _check_invariants(samples) + _record("3.4 rewrite-merge ambiguous -> fork", mgr, sid, samples) + print("PASS 3.4") + + +def test_3_5_rewrite_merge_match_key_updated(): + """After merge, a later turn replaying the rewritten message must descend + through the merged node (match_key updated), not fork again.""" + mgr = TrajectoryManager() + sid = "3.5" + s, u = sys_msg("S"), usr_msg("u") + a1_rw = asst_msg("ok5 ") + t1, a2, t2 = tool_msg("t1"), asst_msg("second"), tool_msg("t2") + append(mgr, sid, [s, u], "ok", finish_reason="tool_calls") + append(mgr, sid, [s, u, a1_rw, t1], "second", finish_reason="tool_calls", logprobs=[-0.4] * 2) + append(mgr, sid, [s, u, a1_rw, t1, a2, t2], "third", logprobs=[-0.3] * 2) + leaves = _leaves(mgr, sid) + assert len(leaves) == 1, "match_key updated -> no spurious fork" + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + # Turn 1 (r:ok) was absorbed as routing-only (rewrite merge), so it appears + # bare; turns 2 and 3 (r:second / r:third) train in one clean chain. + assert goldens(samples) == [ + " system:S user:u r:ok5␣ " + " tool:t1 [r:second] [] tool:t2 [r:third] []", + ] + _check_invariants(samples) + _record("3.5 rewrite-merge match_key updated", mgr, sid, samples) + print("PASS 3.5") + + +def test_3_6_tree_fork_plus_token_drift(): + """A tree fork (two leaves) where ONE leaf also drift-forks internally, + yielding 3 Samples total. Combines layer-1 (message fork) with layer-2 + (token drift fork).""" + mgr = TrajectoryManager() + sid = "3.6" + s, u = sys_msg("S"), usr_msg("u") + a1, tx, ty = asst_msg("call"), tool_msg("x"), tool_msg("y") + ax2 = asst_msg("ax2") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) + # Leaf X: clean continuation, then a third turn with a case-A prompt drift. + append(mgr, sid, [s, u, a1, tx], "ax2", finish_reason="tool_calls", logprobs=[-0.4] * 2) + txx = tool_msg("xx") + p3_honest = render_prompt([s, u, a1, tx, ax2, txx]) + p3 = drift(p3_honest, len(p1) - 1) # case A drift -> fork inside leaf X + append(mgr, sid, [s, u, a1, tx, ax2, txx], "ax3", prompt_ids=p3, logprobs=[-0.2] * 2) + # Leaf Y: a separate tool result off the shared assistant. + append(mgr, sid, [s, u, a1, ty], "ay2", logprobs=[-0.1] * 2) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 3, [s.tokens for s in samples] + assert goldens(samples) == [ + # Sample 0: leaf X first segment, trains the shared r:call (first claim) + r:ax2. + " system:S user:u [r:call] [] " + " tool:x [r:ax2] []", + # Sample 1: leaf X second segment, a FRESH segment after the case-A fork -> + # whole prompt stripped, only r:ax3 trains (the sits in its prompt). + " system:S user:u r:call " + " tool:x r:ax2 tool:xx [r:ax3] []", + # Sample 2: leaf Y, shares r:call (claimed by sample 0 -> bare), trains r:ay2. + " system:S user:u r:call " " tool:y [r:ay2] []", + ] + assert all(abs(s.reward - 1.0 / 3) < 1e-9 for s in samples) + _check_invariants(samples) + _record("3.6 tree fork + token drift -> 3 samples", mgr, sid, samples) + print("PASS 3.6") + + +def test_3_7_deep_multi_leaf_dedup(): + """Three leaves sharing a 2-level assistant prefix; the shared turns are + trained exactly once across all leaves.""" + mgr = TrajectoryManager() + sid = "3.7" + s, u = sys_msg("S"), usr_msg("u") + a1, t1 = asst_msg("a1"), tool_msg("t1") + a2 = asst_msg("a2") + append(mgr, sid, [s, u], "a1", finish_reason="tool_calls", logprobs=[-0.5] * 2) + append(mgr, sid, [s, u, a1, t1], "a2", finish_reason="tool_calls", logprobs=[-0.4] * 2) + # three different tool results off a2 -> three leaves sharing a1+a2 + for lbl in ["p", "q", "r"]: + append(mgr, sid, [s, u, a1, t1, a2, tool_msg(lbl)], f"end-{lbl}", logprobs=[-0.3] * 2) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 3 + # Leaf 0 OWNS the shared r:a1 + r:a2 (both trained) and its own end-p; leaves + # 1 and 2 SHARE r:a1 + r:a2 (bare, claimed by leaf 0) and train only their own + # end response. + assert goldens(samples) == [ + " system:S user:u [r:a1] [] " + " tool:t1 [r:a2] [] tool:p [r:end-p] []", + " system:S user:u r:a1 " + " tool:t1 r:a2 tool:q [r:end-q] []", + " system:S user:u r:a1 " + " tool:t1 r:a2 tool:r [r:end-r] []", + ] + _check_invariants(samples) + _record("3.7 deep multi-leaf dedup (3 leaves, shared trained once)", mgr, sid, samples) + print("PASS 3.7") + + +def test_3_8_long_mixed_session(): + """A ~7-turn session combining clean continuation, a mid-session B1 replace, + and a final case-A fork — verifying the mechanisms chain without interfering.""" + mgr = TrajectoryManager() + sid = "3.8" + s, u = sys_msg("S"), usr_msg("u") + a = [asst_msg(f"a{i}") for i in range(6)] + t = [tool_msg(f"t{i}") for i in range(6)] + lp = [-0.5, -0.5] + # turn 1 + p1, r1 = append(mgr, sid, [s, u], "a0", finish_reason="tool_calls", logprobs=lp) + # turns 2..4 clean + prefix = [s, u, a[0], t[0]] + append(mgr, sid, prefix, "a1", finish_reason="tool_calls", logprobs=lp) + prefix = prefix + [a[1], t[1]] + append(mgr, sid, prefix, "a2", finish_reason="tool_calls", logprobs=lp) + prefix = prefix + [a[2], t[2]] + # turn 5: B1 small replace — drift the last token of the previous response. + p5_honest = render_prompt(prefix) + p5 = drift_replace(p5_honest, len(p5_honest) - 2) # near tail, inside last resp echo region + append(mgr, sid, prefix, "a3", prompt_ids=p5, finish_reason="tool_calls", logprobs=lp) + prefix = prefix + [a[3], t[3]] + # turn 6: clean + append(mgr, sid, prefix, "a4", finish_reason="tool_calls", logprobs=lp) + prefix = prefix + [a[4], t[4]] + # turn 7: case-A fork (drift in early prompt region) + p7_honest = render_prompt(prefix) + p7 = drift(p7_honest, len(p1) - 1) + append(mgr, sid, prefix, "a5", prompt_ids=p7, logprobs=lp) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + # The final case-A fork splits the single leaf chain into 3 segments. + assert goldens(samples) == [ + # Segment 1: turns 1-4 in one clean chain. The turn-5 B1 replace dropped a + # drifted tail (the is gone here) and realigned, so r:a0..r:a3 all + # train. + " system:S user:u [r:a0] [] " + " tool:t0 [r:a1] [] tool:t1 [r:a2] [] " + " tool:t2 [r:a3] []", + # Segment 2: cross-leaf-style dedup within the chain — the prior turns are + # re-emitted as bare context and only r:a4 trains. + " system:S user:u r:a0 " + " tool:t0 r:a1 tool:t1 r:a2 " + " tool:t2 r:a3 tool:t3 [r:a4] []", + # Segment 3: turn 7 after the case-A fork (the in the early prompt + # region); whole prefix bare, only r:a5 trains. + " system:S user:u r:a0 " + " tool:t0 r:a1 tool:t1 r:a2 " + " tool:t2 r:a3 tool:t3 r:a4 " + " tool:t4 [r:a5] []", + ] + assert abs(sum(s.reward for s in samples) - 1.0) < 1e-9 + _check_invariants(samples) + _record(f"3.8 long mixed session -> {len(samples)} samples", mgr, sid, samples) + print("PASS 3.8") + + +# =========================================================================== +# §2 Group 4 — boundary / defensive / feature-completion +# +# Fills coverage gaps the matrix above left open: tools-metadata routing, +# input-validation contracts, mixed-logprobs trajectories, the case-B1 drift +# threshold boundary, and the default-base_sample path. +# =========================================================================== + + +def test_4_1_tools_metadata_on_first_system_only(): + """tools passed to append_turn attach to the FIRST system node only; a later + turn carrying the same system must NOT re-attach (dedup via the + system-ancestor walk).""" + mgr = TrajectoryManager() + sid = "4.1" + s, u = sys_msg("S"), usr_msg("u") + a1, t1 = asst_msg("call"), tool_msg("t") + append(mgr, sid, [s, u], "call", finish_reason="tool_calls", tools=TOOLS) + append(mgr, sid, [s, u, a1, t1], "done", tools=TOOLS) + sys_node = mgr._trees[sid].children[0] + assert sys_node.role == "system" + assert sys_node.metadata.get("tools") == TOOLS, "tools land on the first system node" + # No other node carries tools. + others = [n for n in _iter_all(mgr._trees[sid]) if n is not sys_node] + assert all(n.metadata.get("tools") is None for n in others), "tools attached exactly once" + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert goldens(samples) == [ + " system:S user:u [r:call] [] " + " tool:t [r:done] []", + ] + _check_invariants(samples) + _record("4.1 tools metadata on first system only", mgr, sid, samples) + print("PASS 4.1") + + +def test_4_2_logprobs_length_mismatch_raises(): + """output_log_probs whose length != output_ids -> ValueError at append_turn.""" + mgr = TrajectoryManager() + sid = "4.2" + s, u = sys_msg("S"), usr_msg("u") + bad = TurnRecord( + prompt_ids=render_prompt([s, u]), + output_ids=[9101, 9102, 9103], + finish_reason="stop", + output_log_probs=[-0.1, -0.2], # length 2 != 3 + ) + raised = False + try: + mgr.append_turn( + sid, + turn=bad, + prompt_messages=messages([s, u]), + tools=None, + response_message={"role": "assistant", "content": "x"}, + ) + except ValueError as e: + raised = True + assert "output_log_probs" in str(e) + assert raised, "expected ValueError on logprobs/ids length mismatch" + print("PASS 4.2") + + +def test_4_3_empty_prompt_messages_skipped(): + """Empty prompt_messages -> append_turn is a no-op (warns, no node, no turn).""" + mgr = TrajectoryManager() + sid = "4.3" + mgr.append_turn( + sid, + turn=turn([1], [2], finish_reason="stop"), + prompt_messages=[], + tools=None, + response_message=None, + ) + assert mgr.turn_count(sid) == 0 + # The tree may be created empty (root only) or absent; either way no leaf. + assert not mgr.has_session(sid) or list(_leaves(mgr, sid)) == [] + print("PASS 4.3") + + +def test_4_4_default_base_sample(): + """get_trajectory with base_sample=None uses a default Sample(index=0).""" + mgr = TrajectoryManager() + sid = "4.4" + s, u = sys_msg("S"), usr_msg("u") + append(mgr, sid, [s, u], "a") + # snapshot tree before drain so the dump still renders it + _TREE_SNAP[sid] = dump_tree_txt(mgr, sid) + _REWARD_IN[sid] = 1.0 + samples = mgr.get_trajectory(sid, reward=1.0) # no base_sample + assert len(samples) == 1 + assert samples[0].index == 0 + assert goldens(samples) == [" system:S user:u [r:a] []"] + _check_invariants(samples) + _record("4.4 default base_sample (None)", mgr, sid, samples) + print("PASS 4.4") + + +def test_4_5_mixed_logprobs_across_turns(): + """A trajectory where turn 1 carries logprobs and turn 2 does NOT: the + sample's turn-1 response region has real logprobs, the turn-2 region is + padded with 0.0 (the response is still trained, loss=1).""" + mgr = TrajectoryManager() + sid = "4.5" + s, u = sys_msg("S"), usr_msg("u") + a1, t1 = asst_msg("call"), tool_msg("t") + p1, r1 = append(mgr, sid, [s, u], "call", finish_reason="tool_calls", logprobs=[-0.5] * 2) + p2, r2 = append(mgr, sid, [s, u, a1, t1], "done", logprobs=None) # no logprobs + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + assert len(samples) == 1 + s0 = samples[0] + L = _lcp_len(p1 + r1, p2) + # both responses still trained (golden shows the loss layout)... + assert goldens(samples) == [ + " system:S user:u [r:call] [] " + " tool:t [r:done] []", + ] + # ...but turn-2's region carries padded 0.0 logprobs (it had none), while + # turn-1's region keeps its real logprobs. + assert s0.rollout_log_probs == [-0.5] * len(r1) + [0.0] * (len(p2) - L) + [0.0] * len(r2) + _check_invariants(samples) + _record("4.5 mixed logprobs across turns (turn2 padded 0.0)", mgr, sid, samples) + print("PASS 4.5") + + +def test_4_6_drift_B1_threshold_boundary(): + """case-B1 threshold is exclusive: a drift tail of length d == threshold + forks, d == threshold-1 replaces. Verify both sides of the boundary.""" + + def run(threshold, drift_tail_len): + mgr = TrajectoryManager(fork_threshold_tokens=threshold) + sid = f"4.6-{threshold}-{drift_tail_len}" + s, u = sys_msg("S"), usr_msg("u") + # 4-token response so the divergence can sit d tokens before its end. + p1 = render_prompt([s, u]) + r1 = [9001, 9002, 9003, 9004] + mgr.append_turn( + sid, + turn=turn(p1, r1, finish_reason="tool_calls"), + prompt_messages=messages([s, u]), + tools=None, + response_message={"role": "assistant", "content": "a1"}, + ) + a1m = {"role": "assistant", "content": "a1"} + tm = tool_msg("t") + # honest turn-2 prompt echoes p1 + r1 then the tool block + gen marker. + p2_honest = p1 + r1 + tm.render() + [_GEN] + # divergence d tokens before the end of r1's echo (inside its response span). + drift_idx = len(p1) + len(r1) - drift_tail_len + p2 = drift_replace(p2_honest, drift_idx) + r2 = [9101, 9102] + mgr.append_turn( + sid, + turn=turn(p2, r2, finish_reason="stop"), + prompt_messages=[*messages([s, u]), a1m, tm.message], + tools=None, + response_message={"role": "assistant", "content": "done"}, + ) + samples = get_traj(mgr, sid, base_sample=Sample(index=0, prompt=""), reward=1.0) + return samples, p1, r1, p2, r2 + + # d == threshold -> fork: two single-turn segments, each trains its own resp. + forked, p1, r1, p2, r2 = run(threshold=2, drift_tail_len=2) + assert len(forked) == 2, f"d==threshold must fork, got {len(forked)}" + assert forked[0].tokens == p1 + r1 + assert forked[0].loss_mask == [1] * len(r1) + assert forked[1].tokens == p2 + r2 + assert forked[1].loss_mask == [1] * len(r2) + # d < threshold -> replace: one coherent segment realigned to p2. + replaced, p1b, r1b, p2b, r2b = run(threshold=2, drift_tail_len=1) + assert len(replaced) == 1, f"d None: + toks = s.tokens + resp_start = len(toks) - s.response_length + # build aligned token/loss rows over the response region (the trained part); + # the leading prompt prefix has no loss_mask entry. + names = [name_of(t) for t in toks] + loss = ["-"] * resp_start + [str(x) for x in s.loss_mask] + widths = [max(len(names[i]), len(loss[i])) for i in range(len(toks))] + tok_row = " ".join(names[i].ljust(widths[i]) for i in range(len(toks))) + loss_row = " ".join(loss[i].ljust(widths[i]) for i in range(len(toks))) + print(f" Sample#{idx} reward={s.reward:.3f} resp_len={s.response_length}") + print(f" tok : {tok_row}") + print(f" loss: {loss_row}") + + +def _print_raw_turns(sid: str) -> None: + """Print the raw append_turn inputs (the SOURCE data) for a sid. + + Shows, per turn, the prompt message labels and the actual prompt_ids / + response_ids decoded to readable names, plus finish_reason and whether + logprobs were attached. This is what fed the tree, before any building or + linearization. + """ + turns = _TURN_LOG.get(sid, []) + print(f"[raw turns] {len(turns)}") + for k, t in enumerate(turns, start=1): + msgs = " , ".join(t["prompt_msgs"]) + print(f" turn#{k} finish={t['finish']} has_logprobs={t['has_lp']}") + print(f" msgs : {msgs}") + print(f" prompt : {render_ids(t['prompt_ids'])}") + print(f" output : {render_ids(t['response_ids']) or ''}") + + +def _print_case(title: str, mgr, sid: str, samples: list) -> None: + print(f"\n=== CASE {title} ===") + _print_raw_turns(sid) + if mgr.has_session(sid): + txt = dump_tree_txt(mgr, sid) + else: + # Session already drained by get_trajectory; fall back to the snapshot + # captured by get_traj just before draining. + txt = _TREE_SNAP.get(sid, "") + print("[tree]") + for line in txt.splitlines(): + print(" " + line) + n = len(samples) + if n: + r_in = _REWARD_IN.get(sid, 0.0) + per = r_in / n + print(f"[samples] {n} (reward split: {r_in:.3f} / {n} = {per:.3f} per sample)") + else: + print(f"[samples] {n}") + for i, s in enumerate(samples): + _print_sample(i, s) + + +# =========================================================================== +# main +# =========================================================================== + + +_CASES = [ + test_1_1_single_turn_chain, + test_1_2_clean_multiturn_with_tool, + test_1_3_system_fork, + test_1_4_user_fork_shared_system, + test_1_5_assistant_message_fork, + test_1_6_tool_fork_shared_assistant, + test_1_7_token_only_drift_no_fork, + test_1_8_multi_tool_per_turn, + test_1_9_cross_sid_isolation, + test_1_10_empty_response, + test_2_1_single_turn_linearize, + test_2_2_clean_multiturn_linearize, + test_2_3_drift_case_A_forks, + test_2_4_drift_case_B1_short_replaces, + test_2_5_drift_case_B1_long_forks, + test_2_6_drift_case_B1_threshold_zero_forks, + test_2_7_drift_case_B2_earlier_turn_forks, + test_2_8_fork_reward_split, + test_2_9_two_leaves_reward_split, + test_2_10_cross_leaf_dedup, + test_2_11_routing_only_assistant_filtered, + test_2_12_drop_clears_sid, + test_3_1_rewrite_merge_absorbs_short, + test_3_2_rewrite_merge_long_forks, + test_3_3_rewrite_merge_threshold_zero_forks, + test_3_4_rewrite_merge_ambiguous_forks, + test_3_5_rewrite_merge_match_key_updated, + test_3_6_tree_fork_plus_token_drift, + test_3_7_deep_multi_leaf_dedup, + test_3_8_long_mixed_session, + test_4_1_tools_metadata_on_first_system_only, + test_4_2_logprobs_length_mismatch_raises, + test_4_3_empty_prompt_messages_skipped, + test_4_4_default_base_sample, + test_4_5_mixed_logprobs_across_turns, + test_4_6_drift_B1_threshold_boundary, +] + + +def main() -> None: + for case in _CASES: + case() + # Replay the captured tree / sample snapshots as human-readable dumps. + print("\n" + "=" * 70) + print("HUMAN-READABLE DUMPS") + print("=" * 70) + for title, mgr, sid, samples in _PRINT_LOG: + _print_case(title, mgr, sid, samples) + print(f"\nALL E2E CASES PASSED ({len(_CASES)} cases)") + + +if __name__ == "__main__": + main() From ad122f417b92953c9dbd2795392e2b6b0970f2fd Mon Sep 17 00:00:00 2001 From: jingshenghang Date: Tue, 9 Jun 2026 03:24:54 +0000 Subject: [PATCH 28/28] refactor(agent): assert base_sample in get_trajectory instead of defaulting A None base_sample should never reach get_trajectory; replace the silent Sample(index=0) fallback with an assert so callers can't drop it. Update the e2e case to expect the assert and import pytest. --- slime/agent/trajectory_manager.py | 3 +-- tests/test_agent/test_trajectory_manager_e2e.py | 15 +++++---------- 2 files changed, 6 insertions(+), 12 deletions(-) diff --git a/slime/agent/trajectory_manager.py b/slime/agent/trajectory_manager.py index 5ca044a594..d48b7d1e09 100644 --- a/slime/agent/trajectory_manager.py +++ b/slime/agent/trajectory_manager.py @@ -227,8 +227,7 @@ def get_trajectory( across all of them. The sid is dropped afterwards, so a second call for the same sid returns ``[]``. """ - if base_sample is None: - base_sample = Sample(index=0, prompt="") + assert base_sample is not None, "get_trajectory requires a base_sample" root = self._trees.get(sid) if root is None: diff --git a/tests/test_agent/test_trajectory_manager_e2e.py b/tests/test_agent/test_trajectory_manager_e2e.py index f5871eeb47..d13bd2264f 100644 --- a/tests/test_agent/test_trajectory_manager_e2e.py +++ b/tests/test_agent/test_trajectory_manager_e2e.py @@ -30,6 +30,8 @@ from __future__ import annotations +import pytest # noqa: E402 + from tests.test_agent.test_claude_code_agent._dump_helpers import dump_tree_txt # noqa: E402 from slime.agent.adapters.common import TurnRecord # noqa: E402 @@ -1202,20 +1204,13 @@ def test_4_3_empty_prompt_messages_skipped(): def test_4_4_default_base_sample(): - """get_trajectory with base_sample=None uses a default Sample(index=0).""" + """get_trajectory without base_sample is rejected by an assert.""" mgr = TrajectoryManager() sid = "4.4" s, u = sys_msg("S"), usr_msg("u") append(mgr, sid, [s, u], "a") - # snapshot tree before drain so the dump still renders it - _TREE_SNAP[sid] = dump_tree_txt(mgr, sid) - _REWARD_IN[sid] = 1.0 - samples = mgr.get_trajectory(sid, reward=1.0) # no base_sample - assert len(samples) == 1 - assert samples[0].index == 0 - assert goldens(samples) == [" system:S user:u [r:a] []"] - _check_invariants(samples) - _record("4.4 default base_sample (None)", mgr, sid, samples) + with pytest.raises(AssertionError): + mgr.get_trajectory(sid, reward=1.0) # no base_sample print("PASS 4.4")