feat(session): session server v2 — opt-in trajectory-tree serving; v1 preserved byte-exact#1771
Open
guapisolo wants to merge 16 commits into
Open
feat(session): session server v2 — opt-in trajectory-tree serving; v1 preserved byte-exact#1771guapisolo wants to merge 16 commits into
guapisolo wants to merge 16 commits into
Conversation
TITOTokenizer.render_messages was a one-line delegate to the module-level apply_chat_template that only binds self.tokenizer and **self.chat_template_kwargs; having two names for the same rendering operation split the vocabulary between the method and the function it wraps. Rename the method (body and signature unchanged) so the instance method matches both the module function and the HF-native tokenizer.apply_chat_template convention, and update all call sites. Inside tito_tokenizer.py the module function is now referenced as template.apply_chat_template so the same-named delegation does not read as recursion and cannot be confused with the method or the HF method. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Streaming clients hitting POST /sessions/{id}/v1/chat/completions got a
deterministic 400 today: the session server always injects
return_meta_info=True, and sglang rejects that combination with
stream=true. TITO also needs the complete non-streaming response
(choices[0].message + meta_info.output_token_logprobs), so real
streaming passthrough cannot work either.
Instead, pop stream/stream_options in Phase 1 so the backend call and
the SessionRecord stay non-streaming, then honor the client's intent at
response time: synthesize a single chat.completion.chunk carrying the
whole message as one delta (role/content/reasoning_content/tool_calls
with index; all tool_calls ride in one chunk because some clients
mis-assemble fragmented arguments), followed by data: [DONE]. Errors
keep their real status codes as JSON since they all occur before the
SSE body is built. Adapted from NVIDIA-NeMo/ProRL-Agent-Server and
THUDM/slime.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Black-box agent harnesses mostly consume chat completions as SSE, so the session-server e2e tests should exercise the fake-streaming path by default instead of leaving it fast-test-only. Add miles/utils/test_utils/openai_stream_client.py, a black-box SSE client that rebuilds the non-streaming response shape from OpenAI delta chunks (content/reasoning_content concatenation, tool_calls merged by index with the streaming-only index key stripped). session_verify_agent (test_session_server_multi_role, 7 models) and session_tool_agent now default to stream=true, with a request_kwargs stream=False opt-out. The rebuilt assistant message must survive the next turn's TITO prefix check (message_matches compares role/content/reasoning_content/ tool_calls); two new fast tests pin that down: a multi-turn streaming conversation extends the session without rollback, and rebuilt tool_calls dict-equal the record's stored assistant tool_calls. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Streaming clients rebuild the assistant message from the fake-stream chunk and, per OpenAI delta semantics, drop the streaming-only tool_call `index` key. SGLang's non-streaming wire always carries it (ToolCall.index, int or null), so the stored assistant message never matched its replay: every multi-role e2e tool_result request died with HTTP 400 "rollback failed: no assistant message found in the first 2 matched messages" and stage-c-4-gpu-h200 (qwen3 / deepseekv4 / nemotron3) timed out with zero progress (run 29881124917). message_matches' contract is template-relevant comparison, so normalize tool_calls before comparing: drop `index` and null-valued keys from each call on both sides. No chat template reads them, and Jinja renders an absent key and a null-valued key identically. MockSGLangServer now serializes `index` like real SGLang (the tool's position in the tools list, via ToolCallItem.tool_index) — its previous shape is what let the fast suite miss this. Fixtures and exact-shape expectations follow the wire shape, and the new streamed tool-call -> tool_result continuation test reproduces the CI failure without the template.py fix. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…keys Reviewers reasonably ask why the stateful server does not instead repair the replayed message by filling the missing tool_call index back from stored state. Record the rationale at the decision site: matching only answers "same history?" (stored checkpoints stay the token/record authority, nothing downstream reads replayed keys), enrichment would presuppose the per-call correspondence the comparison is establishing, and it cannot handle clients that renumber index instead of dropping it. Also spell out that index is an sglang extension absent from the OpenAI non-streaming message shape. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
compute_samples_from_openai_records, its per-record helper, and truncate_samples_by_total_tokens move from the client-side generate_utils module to miles/rollout/session/samples/merge.py, unchanged: the follow-up commit makes the session server assemble samples where the records live, and server-owned logic must not live in a driver-named module (core.py importing generate_utils.sample_utils.merge_samples set the dependency-direction precedent). Their tests move alongside to tests/fast/rollout/session/test_samples.py. Also deletes tests/e2e/sglang/utils/logprob_verify_generate.py: already unreferenced, and built on the records-collection path the follow-up commit retires. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ds never leave it
POST /sessions/{session_id}/samples (registered before the catch-all proxy, which would otherwise forward it to the inference backend): the owning instance runs compute -> truncate -> merge synchronously on its event loop (the lock-free get_session invariant) and replies with one safetensors payload — one named tensor per computed array field, with the scalar JSON riding as the rank-one uint8 tensor _samples_meta (safetensors.numpy.load exposes no header metadata and __metadata__ is reserved by the format), so malformed payloads fail inside safetensors' validated parser instead of hand-rolled framing checks. Only the ten COMPUTED_FIELDS cross the wire on blank templates; the driver overlays them onto deepcopies of its input sample, and an import-time guard forces every new Sample field to be classified computed-or-template. The codec lives in samples/codec.py as a black box (encode_samples_reply / decode_samples_reply / SamplesReply), depending only on Sample, NumPy, and safetensors (>=0.8.0, now an explicit dependency); deterministic assembly failures return 422 with the assertion text; empty replies carry a no_records/all_truncated discriminator preserving today's ABORTED semantics.
The old records path (GET the full dump, GiB-scale under R3 with per-turn full-prefix arrays, parsed on the driver's single interpreter) is retired from the training path. collect_samples posts once via the new http_utils.post_bytes_no_retry (post() force-decodes json/text and blind-retries; a 5xx here means the owning instance died with the records, a 422 is deterministic), raises on non-2xx with the body text, raises on timeout instead of silently ABORTing the sample (a measured data-loss bug in the records path), and attempts the session DELETE on every path so failed sessions cannot accumulate. agentic_tool_call.generate shrinks to a shell: agent call -> collect_samples -> empty_reason mapping -> metadata application in today's order.
Tests assert golden Sample field values derived from the records fixture (per-turn, merged, truncated), the wire-codec round-trip, drift-guard and malformed-payload contracts, the 422/404/empty_reason/route-order contracts, and the client's behavior deltas.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…from the samples codec The import-time guard demanded a classification edit for every new Sample field, though most new fields are driver-side and irrelevant to the wire (main's multi-lora adapter/reward_spec broke import of the whole session package the day it merged, killing pytest collection at conftest). The guard also could not force correct classification: dumping a name into TEMPLATE_FIELDS silenced it while shipping the very stale-value bug it existed to catch. The wire contract is now the allowlist alone: only COMPUTED_FIELDS cross the samples wire; every other Sample field never crosses and the driver overlay keeps its local deepcopy's value verbatim. The load-bearing checks stay: every computed field needs exactly one wire representation, and _assert_overlay_template_defaults still guards overlay equivalence. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ple semantics
The flag made every turn a separate full-context sample (skipping the TITO
merge), which is wrong: sample boundaries are a property of the trajectory,
not a CLI switch, and per-turn siblings share one group_index so they pollute
group baselines. The session server now always folds linear turns into one
merged sample (a new sample only starts at a genuine context discontinuity,
i.e. future compaction/subagent support), the samples-op body shrinks to
{"max_seq_len"}, and agentic_tool_call.generate always returns list[Sample] —
including the aborted path, so a batch never mixes scalar and list shapes.
Consequences: multi_turn.generate always TITO-accumulates one scalar sample;
custom RMs on the agentic path are now called in batch form with
list[Sample]; check_reward_nonzero_std flattens nested groups; agentic +
--group-rm/--partial-rollout/--recompute-logprobs-via-prefill is documented
as unsupported.
…ample]] A dynamic-filter group mixes scalar and list elements (agentic generate returns list[Sample] per call), which is exactly why _flatten_samples exists; the old list[Sample] annotation misdocumented that contract.
guapisolo
requested review from
Shi-Dong,
Zhichenzzz,
fzyzcjy,
jybsuper,
maocheng23 and
yueming-yuan
as code owners
July 23, 2026 10:08
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Additive files only, verified green against the current tree untouched: nine pins locking the default session server's retry semantics at the HTTP layer — pure-drop and divergent retries roll back one assistant and regenerate; deep and no-anchor rollbacks 400 with exact error texts and no state mutation; failed first turn keeps empty-session accept-anything; exact-history resend is a degenerate extension; the disallowed-role 400 keeps its rollback side effect; collect after retry yields one sample; and the few-shot prompt-assistant anchoring quirk is pinned as a characterization (such shapes are outside v1's append-only + one-step retry envelope — flows that need them belong to session server v2). These pins exist so the upcoming opt-in v2 tree server can never silently change the default path. The dir-level conftest re-exports test_sessions' router_env fixture so additive modules can share it without a shadowed import. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…plate surface
Two structural facts about mid-path assistant re-rendering, promoted to
the registry so every template surface agrees ("never cut think"):
- The preserve kwargs (clear_thinking / truncate_history_thinking /
drop_thinking / preserve_thinking) become family constants on all
{tool} rows, not just the plain-chat rows; spike-verified byte-identical
on pure tool loops. Prompt assistants carrying reasoning (few-shot,
compaction) now render with their think blocks preserved — that is the
fix, not a side effect.
- supports_midpath_assistant_rerender (True on the base class) marks the
one structural exception: DeepSeek-V3.2's sglang encoder hardcodes the
think gate with no escape kwarg, so re-rendering an assistant that
precedes a later user turn cannot be canonical there and consumers must
fail loud instead of silently cutting reasoning.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
First piece of session server v2 (opt-in tree serving; MULTI_LINEAGE_DESIGN v5), new package only — nothing imports it yet and v1 stays byte-identical. A session is a forest of TrajectoryNodes; a node is the end of ONE model generation (delta_messages = env msgs + client-carried foreign assistants + the sampled assistant; token_ids = full root->node snapshot so every leaf is its own assembled sequence; completion_span in snapshot coordinates; seq = per-session logical commit order, THE ordering key; response_id = the agent-branch <-> leaf join key). find_attach_point is a pure deepest-prefix search (full-delta consumption gate, twin ties to latest seq); create_node is append-only under the owning session lock with a MAX_NODES=1024 runaway backstop. Unit matrix copied verbatim from the v5 line (18 tests). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…server v2 The v5 multi-lineage design lands as an OPT-IN second implementation instead of replacing v1: --use-session-server grows an optional value (bare flag or "v1" = the untouched linear server, "v2" = tree serving), and the whole tree stack lives in miles/rollout/session/v2/. The v1 modules stay byte-identical to their pre-branch state. v2 package (carried from the v5 line, imports adjusted only): - session_state: always-branch serving over the forest — deepest attach point, suffix becomes the branch delta, non-extensions grow siblings or new roots; retry semantics moves entirely to the samples-op picker. Branch suffixes may carry client assistants (compaction); snapshot splice with zero-inheritance-root fallback on canonical-prefix divergence. Truncated-path extension is 409 (TruncatedGenerationError, defined here — v1 never raises it). --session-strict-append-only is the single-chain guard: branch shapes fail loud with attach diagnostics. - assembly: per-leaf fold (compute -> truncate -> fold), default pick (temporal-supersession retry trim; roots never trimmed; non-retry-shaped trees 422) and default merge (exactly-once completion masking over the surviving set, rewards keyed by response id, fold < agent < server metadata layering); both replaceable via --session-sample-picker-path / --session-merge-function-path (sync-only, loaded in-process). - core: v2 twin of SessionCore reusing v1's HTTP plumbing unchanged. Shared seams, each inert for v1 by construction: - sessions.py dispatches registry/core on the flag and forwards the collect body's "metadata" (agent semantic layer) only under v2. - codec: encode/decode parameterized by a fields tuple; v1 default keeps the wire byte-identical, v2 adds reward + per-sample metadata (COMPUTED_FIELDS_V2) with conditional overlay semantics on decode. - arguments: flag upgrade (nargs="?"), unknown values rejected, and the three v2 flags require --use-session-server v2. Tests: v2 HTTP matrix (tree pins incl. deep/root divergence branching, strict guard, truncation + compaction), session_state unit matrix, and the samples-op suite (golden, multi-leaf trim/masking/rewards, hook lanes) — all against a v2-flagged server; the restored v1 suites and the byte-exact A-list (git diff against the pre-branch base is empty) pin that the default path did not move. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…nal overlay
The driver half of session server v2, gated on
args.use_session_server == "v2" so the v1 path is byte-equal to before:
- OpenAIEndpointTracer captures the samples-wire allowlist at create()
(v1 default for direct constructions); collect_samples gains
agent_metadata and puts it on the body as "metadata" ONLY when set —
the v1 collect body stays exactly {"max_seq_len": ...} (pinned).
- agentic_tool_call: under v2 the agent's metadata (incl. per-branch
rewards keyed by response id) travels through collect_samples and comes
back applied by the server-side merge, so the driver-side overlays
(agent_metadata onto every sample, session_metadata onto the last) are
v1-only branches now.
- session_verify_runner re-serializes an explicit version string instead
of silently downgrading it to the bare flag.
Pins: the untouched v1 payload/behavior tests keep passing as-is; new v2
pins cover the metadata channel, reward/metadata overlay decode, and
create()'s version-to-fields selection.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…z + edge matrix The v2 legality suite, carried from the v5 line (v2 import paths + the extended wire tuple are the only changes): an independent oracle re-derives the picker rules and earliest-surviving-leaf ownership from the design and checks every assembled sample for token/lockstep legality, exactly-once training over owned completion spans, precise trim sets (guardrail 422s must name the violator), and reward/order placement. Three seeded fuzz banks (retry-shaped forests, max_seq_len truncation with ownership re-derived over the surviving material, shorter-replacement shapes driving the guardrail lane) plus targeted edges: chained retries, equal-length twins, fork-below-fork ownership nesting, mid-path length-turn fold stop. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
guapisolo
force-pushed
the
feat/session-rollback-mode
branch
from
July 23, 2026 19:42
96f91e7 to
58bc6d5
Compare
guapisolo
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
The session server needs multi-lineage trajectory support (retries, compaction with carried assistants, subagent forests). Earlier revisions of this branch replaced the linear implementation in place; that made a brand-new system look like an evolution of the old one and put the default production path at risk. This revision restructures the work so that:
git diffagainst the pre-branch base is empty for every v1 module and test file.--use-session-server v2), fully contained inmiles/rollout/session/v2/.Commit series (atop the stacked picks)
test(session): nine byte-exact HTTP pins for the v1 retry/rollback surface, verified against the untouched tree — the default path can never move silently. Includes a characterization pin for v1's few-shot prompt-assistant anchoring quirk (out of v1's append-only + one-step-retry envelope; such flows belong to v2).fix(tito): preserve-think kwargs become family constants on every template surface ("never cut think"), plussupports_midpath_assistant_rerendermarking DeepSeek-V3.2's structural exception. Spike-verified byte-identical on pure tool loops.feat(session): v2 trajectory tree — node = end of one model generation, full root→node token snapshots, pure deepest-prefix attach search, append-only commits under the session lock (MAX_NODES=1024backstop).feat(session): session server v2 — always-branch serving (mismatches grow siblings/new roots; retry semantics moves to the samples-op picker),--session-strict-append-onlysingle-chain guard, compaction splice with zero-inheritance-root fallback, 409 on extending truncated generations, per-leaf assembly with default pick (temporal-supersession trim, roots never trimmed, non-retry-shaped trees 422) and default merge (exactly-once completion masking, rewards keyed by response id), both replaceable via--session-sample-picker-path/--session-merge-function-path. Version dispatch is a ~10-line seam insessions.py; the codec is parameterized so the v1 wire stays byte-identical while v2 addsreward/metadata.feat(rollout): driver-side seam — the agent's semantic layer (per-branch rewards) travels throughcollect_samplesunder v2 and comes back applied by the server-side merge; the v1 driver path (metadata overlays,{"max_seq_len": ...}collect body) is byte-equal to before, pinned.test(session): tree-merge legality invariants — an independent oracle over randomized forests (seeded fuzz banks for retry shapes, truncation, guardrail-422 shapes) plus targeted edges (chained retries, equal-length twins, fork-below-fork ownership, mid-path length turns).Testing
git diffagainst base; new v1 pins green against the pristine tree.Follow-ups
session_verify_runnerbranching-harness variant) on 4×H200 CI.🤖 Generated with Claude Code
https://claude.ai/code/session_01NmYsBAGK6NDt1LzjA7VSqH