feat(session): session server v2 — opt-in trajectory-tree serving; v1 preserved byte-exact

docs/user-guide/agentic-chat-template.md

@@ -1,9 +1,9 @@
 ---
-title: Agentic Chat Templates (TITO)
+title: Agentic Rollout (TITO)
 description: How to turn on and verify Token-In-Token-Out (TITO) for multi-turn agentic rollout.
 ---
 
-# Agentic Chat Templates (TITO)
+# Agentic Rollout (TITO)
 
 Multi-turn agentic rollout in Miles runs on **TITO** (Token-In-Token-Out): each turn's token sequence is a bit-perfect prefix of the next, so the trainer sees exactly the tokens the engine produced — no re-tokenization, no drift. The *why* is in the blog ([No Token Left Behind](https://lmsys.org/blog/2026-05-13-no-token-left-behind/)); this page is *how*.
 

docs/user-guide/index.md

@@ -12,7 +12,7 @@ description: Concepts, launch script walkthrough, customization hooks, and a com
 | [Customization](/user-guide/customization) | The 21 `--*-path` plug-points for custom Python — rollout, reward, filters, loss, hooks. |
 | [Rollout Endpoints](/user-guide/rollout-endpoints) | The `/generate` endpoint and the OpenAI chat endpoint for agentic sessions. |
 | [Fully Async Rollout](/user-guide/fully-async) | Queue-backed rollout production, tuning knobs, and when to use `train_async.py`. |
-| [Agentic Chat Templates](/user-guide/agentic-chat-template) | Turning on and verifying TITO so multi-turn agentic rollout stays append-only. |
+| [Agentic Rollout (TITO)](/user-guide/agentic-chat-template) | Turning on and verifying TITO so multi-turn agentic rollout stays append-only. |
 | [CLI Reference](/user-guide/cli-reference) | Every flag Miles accepts, grouped by subsystem. |
 | [Environments](/user-guide/environments) | Supplying an environment: dataset + reward, your own env via the plug points, or an external ecosystem. |
 

docs/user-guide/rollout-endpoints.md

@@ -181,6 +181,14 @@ remain a `messages` list. SGLang handles templating server-side.
 
 </Warning>
 
+<Warning>
+
+**Session server v2 output is a `list[Sample]`.** With `--use-session-server v2`, `agentic_tool_call.generate` returns one sample for each selected tree leaf. The v1 session server returns one scalar `Sample`.
+
+A custom reward model (`--custom-rm-path`) receives the v2 samples in batch form. `--group-rm`, `--partial-rollout`, and `--recompute-logprobs-via-prefill` are not supported with this v2 agentic output and are rejected explicitly.
+
+</Warning>
+
 ### Optional teardown: the `abort` hook
 
 The module named by `--custom-agent-function-path` may expose an optional `abort`
@@ -256,6 +264,6 @@ inherited across turns. Each request is tokenized independently.
 ## Next
 
 - [Customization](/user-guide/customization): the full catalog of `--*-path` hooks.
-- [Agentic Chat Templates](/user-guide/agentic-chat-template): verifying that a template is
+- [Agentic Rollout (TITO)](/user-guide/agentic-chat-template): verifying that a template is
   append-only across turns.
 - [Multi-agent example](/examples/multi-agent): full agentic walkthrough.

miles/rollout/generate_hub/agentic_tool_call.py

@@ -46,6 +46,7 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput:
         "agentic_tool_call.generate requires session_server_ip/session_server_ports. "
         "Pass --use-session-server to start the session server."
     )
+    use_v2 = getattr(input.args, "use_session_server", None) == "v2"
     tracer = await OpenAIEndpointTracer.create(input.args)
 
     custom_agent_function: Callable = load_function(input.args.custom_agent_function_path)
@@ -84,11 +85,11 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput:
     finally:
         # Collect even if the agent failed.
         logger.debug(f"{log_prefix} Calling collect_samples...")
+        collect_kwargs = {"max_seq_len": max_seq_len}
+        if use_v2:
+            collect_kwargs["agent_metadata"] = agent_metadata
         try:
-            result = await tracer.collect_samples(
-                input.sample,
-                max_seq_len=max_seq_len,
-            )
+            result = await tracer.collect_samples(input.sample, **collect_kwargs)
         except asyncio.TimeoutError:
             collect_timed_out = True
             logger.warning(f"{log_prefix} Timed out collecting samples", exc_info=True)
@@ -101,7 +102,7 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput:
     if collect_timed_out:
         sample = deepcopy(input.sample)
         sample.status = Sample.Status.ABORTED
-        return GenerateFnOutput(samples=sample)
+        return GenerateFnOutput(samples=[sample] if use_v2 else sample)
 
     if not result.samples:
         if result.empty_reason == "all_truncated":
@@ -110,11 +111,15 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput:
             logger.warning("No model calls recorded for sample")
         sample = deepcopy(input.sample)
         sample.status = Sample.Status.ABORTED
-        return GenerateFnOutput(samples=sample)
+        return GenerateFnOutput(samples=[sample] if use_v2 else sample)
 
     samples = result.samples
-    for s in samples:
-        s.metadata.update(agent_metadata or {})
+    if not use_v2:
+        # v1: the agent's metadata is applied driver-side. Under v2 it traveled
+        # through collect_samples and came back applied by the server-side
+        # merge (per-sample metadata/reward on the wire) — no overlay here.
+        for s in samples:
+            s.metadata.update(agent_metadata or {})
 
     # If the agent function reports wall-clock time spent outside policy generation
     # (env/tool steps), surface it on Sample.non_generation_time so throughput
@@ -124,6 +129,9 @@ async def generate(input: GenerateFnInput) -> GenerateFnOutput:
         for s in samples:
             s.non_generation_time = ngt
 
+    if use_v2:
+        return GenerateFnOutput(samples=samples)
+
     (sample,) = samples
     sample.metadata.update(result.session_metadata)
     return GenerateFnOutput(samples=sample)

miles/rollout/generate_utils/openai_endpoint_utils.py

@@ -7,7 +7,12 @@
 import random
 from argparse import Namespace
 
-from miles.rollout.session.samples.codec import SamplesReply, decode_samples_and_merge_input_sample
+from miles.rollout.session.samples.codec import (
+    COMPUTED_FIELDS,
+    COMPUTED_FIELDS_V2,
+    SamplesReply,
+    decode_samples_and_merge_input_sample,
+)
 from miles.utils.http_utils import post, post_bytes_no_retry
 from miles.utils.types import Sample
 
@@ -17,11 +22,21 @@
 
 
 class OpenAIEndpointTracer:
-    def __init__(self, router_url: str, session_id: str, session_server_instance_id: str | None = None):
+    def __init__(
+        self,
+        router_url: str,
+        session_id: str,
+        session_server_instance_id: str | None = None,
+        samples_wire_fields: tuple[str, ...] = COMPUTED_FIELDS,
+    ):
         self.router_url = router_url
         self.session_id = session_id
         self.base_url = f"{router_url}/sessions/{session_id}"
         self.session_server_instance_id = session_server_instance_id
+        # The samples-wire allowlist must match the server's encode: v1 default,
+        # extended under --use-session-server v2 (create() selects from args;
+        # direct constructions keep v1).
+        self.samples_wire_fields = samples_wire_fields
 
     @property
     def session_server_id(self) -> str:
@@ -45,19 +60,26 @@ async def create(args: Namespace):
         session_server_instance_id = instance_ids.get(session_port)
         response = await post(f"{session_url}/sessions", {}, action="post")
         session_id = response["session_id"]
+        use_v2 = getattr(args, "use_session_server", None) == "v2"
         return OpenAIEndpointTracer(
             router_url=session_url,
             session_id=session_id,
             session_server_instance_id=session_server_instance_id,
+            samples_wire_fields=COMPUTED_FIELDS_V2 if use_v2 else COMPUTED_FIELDS,
         )
 
-    async def collect_samples(self, input_sample: Sample, *, max_seq_len: int | None) -> SamplesReply:
+    async def collect_samples(
+        self, input_sample: Sample, *, max_seq_len: int | None, agent_metadata: dict | None = None
+    ) -> SamplesReply:
         """Fetch server-assembled training samples for this session."""
+        body: dict = {"max_seq_len": max_seq_len}
+        if agent_metadata is not None:
+            body["metadata"] = agent_metadata
         try:
             # `asyncio.TimeoutError` propagates after cleanup is attempted for `agentic_tool_call.generate` to handle.
             payload = await post_bytes_no_retry(
                 f"{self.base_url}/samples",
-                {"max_seq_len": max_seq_len},
+                body,
                 timeout=_SESSION_REQUEST_TIMEOUT,
             )
         finally:
@@ -69,4 +91,4 @@ async def collect_samples(self, input_sample: Sample, *, max_seq_len: int | None
             except Exception as e:
                 logger.warning(f"Failed to delete session {self.session_id} after collecting samples: {e}")
 
-        return decode_samples_and_merge_input_sample(payload, input_sample)
+        return decode_samples_and_merge_input_sample(payload, input_sample, fields=self.samples_wire_fields)

miles/rollout/session/core.py

@@ -142,6 +142,92 @@ def proxy_result_to_response(result: dict) -> Response:
     return Response(content=_render_json(data), status_code=status_code, headers=headers, media_type=JSON_MEDIA_TYPE)
 
 
+def prepare_chat_request(body: bytes, args, tito_tokenizer) -> tuple:
+    """Parse and normalize a chat request body — the session-independent half
+    of chat dispatch, shared verbatim by the v1 and v2 cores. Returns
+    ``(request_body, client_stream, tito_tokenizer)``; the tokenizer may be a
+    request-scoped clone.
+    """
+    try:
+        request_body = json.loads(body) if body else {}
+    except json.JSONDecodeError as e:
+        raise MessageValidationError(f"invalid JSON body: {e}") from e
+
+    # Fake streaming: the backend must stay non-streaming (TITO needs the
+    # complete message + meta_info, and sglang rejects return_meta_info
+    # with stream=true), so pop the client's intent here and honor it
+    # when rendering the client response.
+    client_stream = bool(request_body.pop("stream", False))
+    request_body.pop("stream_options", None)
+
+    # TITO token tracking needs Miles-owned input_ids plus SGLang output
+    # metadata: logprobs=True populates meta_info.output_token_logprobs and
+    # return_meta_info wraps it in choice.meta_info. Hardcoded (not
+    # setdefault) so agent-side overrides cannot break token accumulation.
+    request_body["logprobs"] = True
+    request_body["return_meta_info"] = True
+    if getattr(args, "use_rollout_routing_replay", False):
+        request_body["return_routed_experts"] = True
+    if getattr(args, "use_rollout_indexer_replay", False):
+        request_body["return_indexer_topk"] = True
+    # Must be False so stop-token text is trimmed from assistant content;
+    # token IDs still come from logprobs below.
+    request_body["no_stop_trim"] = False
+    # Serve the adapter being trained instead of the base weights.
+    if is_lora_enabled(args):
+        request_body["lora_path"] = LORA_ADAPTER_NAME
+    # FIXME(session): Only nested `chat_template_kwargs` reach the local renderer;
+    # top-level `reasoning` and `reasoning_effort` are not mapped to template kwargs.
+    request_ctk = request_body.get("chat_template_kwargs")
+    if request_ctk is not None and not isinstance(request_ctk, dict):
+        raise MessageValidationError("chat_template_kwargs must be an object")
+    if request_ctk:
+        try:
+            tito_tokenizer = tito_tokenizer.clone_with_chat_template_kwargs(request_ctk)
+        except ValueError as e:
+            raise MessageValidationError(str(e)) from e
+    if tito_tokenizer.chat_template_kwargs:
+        request_body["chat_template_kwargs"] = dict(tito_tokenizer.chat_template_kwargs)
+    else:
+        request_body.pop("chat_template_kwargs", None)
+    return request_body, client_stream, tito_tokenizer
+
+
+def extract_completion(result: dict) -> tuple:
+    """Decode and validate the backend chat response — shared verbatim by the
+    v1 and v2 cores. Returns ``(response, choice, assistant_message,
+    completion_token_ids)``; malformed upstream payloads raise
+    ``UpstreamResponseError``.
+    """
+    response = json.loads(result["response_body"])
+    choice = response.get("choices", [{}])[0]
+
+    meta_info = choice.get("meta_info")
+    if not isinstance(meta_info, dict) or "output_token_logprobs" not in meta_info:
+        raise UpstreamResponseError("meta_info and output_token_logprobs must be in choice (requires logprobs=True)")
+    assistant_message = choice.get("message") or {}
+    if assistant_message.get("content") is None:
+        raise UpstreamResponseError(
+            "assistant message content is None, when tool call parser failed SGLang should still return "
+            "an empty content rather than None. Please check your modified SGLang version."
+        )
+
+    output_token_logprobs = meta_info["output_token_logprobs"]
+    completion_tokens = meta_info["completion_tokens"]
+
+    actual_output_logprobs_len = len(output_token_logprobs)
+    if actual_output_logprobs_len != completion_tokens:
+        raise UpstreamResponseError(
+            "invalid chat completion response: "
+            f"len(output_token_logprobs)={actual_output_logprobs_len} "
+            f"!= completion_tokens={completion_tokens}. "
+            f"Please check whether you use the correct SGLang branch which has fix the tokenizer batch decode issue."
+        )
+
+    completion_token_ids = [t[1] for t in output_token_logprobs]
+    return response, choice, assistant_message, completion_token_ids
+
+
 class SessionCore:
     """HTTP session operations over one ``SessionRegistry``."""
 
@@ -232,8 +318,8 @@ async def chat_completions(
 
         Flow: prepare pretokenized input_ids (lock held briefly) → proxy to
         backend (NO lock) → validate response → update trajectory checkpoint and
-        append record (lock held briefly). The lock is NOT held during the slow
-        proxy call so DELETE/other ops are not blocked if the agent disconnects.
+        append record (lock held briefly). The lock is NOT held during the long
+        inference call so DELETE/other ops are not blocked if the agent disconnects.
         """
         request_timestamp = time.time()
         session = self.registry.get_session(session_id)
@@ -245,50 +331,9 @@ async def chat_completions(
             if session.closing:
                 raise SessionNotFoundError(f"session not found: session_id={session_id}")
 
-            try:
-                request_body = json.loads(body) if body else {}
-            except json.JSONDecodeError as e:
-                raise MessageValidationError(f"invalid JSON body: {e}") from e
-
-            # Fake streaming: the backend must stay non-streaming (TITO needs the
-            # complete message + meta_info, and sglang rejects return_meta_info
-            # with stream=true), so pop the client's intent here and honor it
-            # when rendering the client response.
-            client_stream = bool(request_body.pop("stream", False))
-            request_body.pop("stream_options", None)
-
-            # TITO token tracking needs Miles-owned input_ids plus SGLang output
-            # metadata: logprobs=True populates meta_info.output_token_logprobs and
-            # return_meta_info wraps it in choice.meta_info. Hardcoded (not
-            # setdefault) so agent-side overrides cannot break token accumulation.
-            request_body["logprobs"] = True
-            request_body["return_meta_info"] = True
-            if getattr(self.args, "use_rollout_routing_replay", False):
-                request_body["return_routed_experts"] = True
-            if getattr(self.args, "use_rollout_indexer_replay", False):
-                request_body["return_indexer_topk"] = True
-            # Must be False so stop-token text is trimmed from assistant content;
-            # token IDs still come from logprobs below.
-            request_body["no_stop_trim"] = False
-            # Without this the engine serves the base weights, so the adapter being
-            # trained would never shape the trajectories it is scored on.
-            if is_lora_enabled(self.args):
-                request_body["lora_path"] = LORA_ADAPTER_NAME
-            # FIXME(session): Only nested `chat_template_kwargs` reach the local renderer;
-            # top-level `reasoning` and `reasoning_effort` are not mapped to template kwargs.
-            request_ctk = request_body.get("chat_template_kwargs")
-            if request_ctk is not None and not isinstance(request_ctk, dict):
-                raise MessageValidationError("chat_template_kwargs must be an object")
-            tito_tokenizer = self.registry.tito_tokenizer
-            if request_ctk:
-                try:
-                    tito_tokenizer = tito_tokenizer.clone_with_chat_template_kwargs(request_ctk)
-                except ValueError as e:
-                    raise MessageValidationError(str(e)) from e
-            if tito_tokenizer.chat_template_kwargs:
-                request_body["chat_template_kwargs"] = dict(tito_tokenizer.chat_template_kwargs)
-            else:
-                request_body.pop("chat_template_kwargs", None)
+            request_body, client_stream, tito_tokenizer = prepare_chat_request(
+                body, self.args, self.registry.tito_tokenizer
+            )
 
             request_messages = request_body.get("messages", [])
             prompt_token_ids = session.prepare_pretokenized(
@@ -314,34 +359,7 @@ async def chat_completions(
         if result["status_code"] != 200:
             return proxy_result_to_response(result)
 
-        response = json.loads(result["response_body"])
-        choice = response.get("choices", [{}])[0]
-
-        meta_info = choice.get("meta_info")
-        if not isinstance(meta_info, dict) or "output_token_logprobs" not in meta_info:
-            raise UpstreamResponseError(
-                "meta_info and output_token_logprobs must be in choice (requires logprobs=True)"
-            )
-        assistant_message = choice.get("message") or {}
-        if assistant_message.get("content") is None:
-            raise UpstreamResponseError(
-                "assistant message content is None, when tool call parser failed SGLang should still return "
-                "an empty content rather than None. Please check your modified SGLang version."
-            )
-
-        output_token_logprobs = meta_info["output_token_logprobs"]
-        completion_tokens = meta_info["completion_tokens"]
-
-        actual_output_logprobs_len = len(output_token_logprobs)
-        if actual_output_logprobs_len != completion_tokens:
-            raise UpstreamResponseError(
-                "invalid chat completion response: "
-                f"len(output_token_logprobs)={actual_output_logprobs_len} "
-                f"!= completion_tokens={completion_tokens}. "
-                f"Please check whether you use the correct SGLang branch which has fix the tokenizer batch decode issue."
-            )
-
-        completion_token_ids = [t[1] for t in output_token_logprobs]
+        response, _, assistant_message, completion_token_ids = extract_completion(result)
 
         # --- Phase 3: update state (lock held briefly) ---
         async with session.lock:

miles/rollout/session/errors.py

@@ -5,6 +5,7 @@
 SessionError (base)
 ├── SessionNotFoundError       → 404  session does not exist
 ├── MessageValidationError     → 400  messages structure/content invalid
+├── TruncatedGenerationError   → 409  extending a length-truncated generation (v2)
 ├── TokenizationError          → 500  TITO tokenizer / prefix mismatch
 └── UpstreamResponseError      → 502  SGLang response invalid or unexpected
 """
@@ -32,6 +33,14 @@ class MessageValidationError(SessionError):
     status_code: int = 400
 
 
+class TruncatedGenerationError(SessionError):
+    """Raised when a request extends a generation that ended with
+    finish_reason='length'; only the v2 tree server raises it (v1 never
+    branches)."""
+
+    status_code: int = 409
+
+
 class TokenizationError(SessionError):
     """Raised when TITO tokenization invariants are violated.