[AgentGateway]feat: Prefix Trie Storage for Multi Trajectories

[gxlvera]

Summary

Implements the per-session prefix trie that replaces the single linear active_trajectory, per RFC #51. Today a session keeps one message_history + one active trajectory, so only the latest branch can be re-attached; this generalizes it into a trie where an incoming request longest-prefix matches any path and continues from the nearest checkpoint — enabling sub-agent / best-of-N / context-condensation / warm-start.

Fully gated behind GatewayActorConfig.gateway_trie_enabled (default off); flag-off reproduces the legacy single-active behavior exactly (covered by a parity test).

What's in this PR (full M1)

Data model — gateway/trie.py (new)

• MessageKey (+ canonicalize_content / make_message_key): hashable per-message identity; excludes reasoning_content (M1; TODO thinking-model replay) and keys multimodal content via a stable image/video digest (pixels live on the node, not the key).
• BranchCheckpoint, TrieNode, PrefixTrie with the prepare → commit lifecycle and a finalize traversal (iter_export_nodes → terminal checkpoints only by default).
• prepare materializes incoming prompt-side messages as pending/structural nodes (no checkpoint, never exported) and clones the nearest checkpoint into a request-local buffer.
• TrajectoryBuffer re-homed to gateway/types.py so the trie stays free of the verl-importing codec.

Session integration — gateway/session.py

• _prepare_generation_inputs_trie: routes through trie.prepare (longest-prefix match + nearest-checkpoint clone) and reuses the existing encode_full / encode_incremental paths — only the state model changes.
• commit attaches the assistant child via trie.commit; finalize traverses the trie → one branch-aware Trajectory per terminal checkpoint (reward stamped per the current one-session-one-reward contract).
• Per-branch multimodal stored on the committed node (new-this-turn images) and reconstructed along the path. Failed generations abandon the pending node. snapshot_state reports num_branches / num_inflight_generations in trie mode.
• Tools-change gate preserved (full re-encode), mirroring the legacy active_tool_schemas gate.

Wiring: config.py / gateway.py thread gateway_trie_enabled into sessions.

Tests (CPU-only)

• tests/uni_agent/gateway/test_trie_on_cpu.py — 16 trie-structural tests covering the RFC appendix fork types: A sequential extension, B system-keyed split (parallel sub-agents), C context condensation, D best-of-N + idempotent retry, E warm-start; plus pending-node/no-export, clone isolation, terminal-vs-all export, multimodal digest + per-branch reconstruction.
• tests/uni_agent/gateway/test_session_trie_on_cpu.py — 3 session-integration tests: flag parity (trie-on yields identical backend prompts + identical finalized trajectory as trie-off on a linear multi-turn conversation — the compatibility gate), best-of-N fan-out (N siblings → N trajectories), multi-turn reattach (single continuous branch).

Scope / deferrals

• Tokenization boundary: reuse current remove_system_prompt (text via tokenizer, multimodal via processor). Per-model seam handling (Qwen trailing \n, GLM 4.7 ambiguous bos/eos) + a token-sequence verifier → M2 (TITO-style merge).
• Concurrency: generation_lock kept (serial) → M3.
• GC/eviction: none — the trie dies with the session at finalize/abort.
• Reward: one-session-one-reward; reward contract unchanged.
• No new caller API — branching is purely message-driven.

Note that this PR is a more complete version of the closed #64 and this description is produced by @ChangyiYang.

[gemini-code-assist]

Code Review

This pull request introduces a prefix trie (PrefixTrie) to support multi-trajectory session storage in the gateway, enabling features like multi-branch reattachment, best-of-N, sub-agents, and warm-starts. The changes integrate the trie into GatewaySession under a new configuration flag and add comprehensive CPU-only unit and integration tests. The review feedback focuses on a key memory efficiency improvement: removing the redundant messages list from BranchCheckpoint to avoid O(N^2) space complexity, as the full conversation history can already be reconstructed on demand using rebuild_messages.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Memory Efficiency Improvement: Redundant messages in BranchCheckpoint

Storing messages in BranchCheckpoint introduces O(N^2) space complexity for message storage over a session of N turns, because each checkpoint stores a copy of the entire conversation history up to that point.

Since PrefixTrie already implements rebuild_messages(node) (which reconstructs the exact same message list by walking up the parent pointers in O(N) time), storing messages on the checkpoint is redundant.

By removing messages from BranchCheckpoint, we can reduce the space complexity to O(N) and avoid significant memory bloat in long conversations (especially those with large system prompts or tool outputs).

Suggested change

	trajectory_buffer: TrajectoryBuffer
	request_tools: list[dict[str, Any]] | None = None
	chat_template_kwargs_key: tuple | None = None
	messages: list[dict[str, Any]] | None = None
	image_data: list[Any] | None = None
	video_data: list[Any] | None = None
	extra_fields: dict[str, Any] | None = None
	trajectory_buffer: TrajectoryBuffer
	request_tools: list[dict[str, Any]] | None = None
	chat_template_kwargs_key: tuple | None = None
	image_data: list[Any] | None = None
	video_data: list[Any] | None = None
	extra_fields: dict[str, Any] | None = None

[gemini-code-assist]

Simplify checkpoint_messages retrieval by directly calling self.rebuild_messages(ckpt_node) now that messages is removed from BranchCheckpoint.

Suggested change

	checkpoint_messages = ckpt.messages or self.rebuild_messages(ckpt_node)
	checkpoint_messages = self.rebuild_messages(ckpt_node)

[gemini-code-assist]

Redundant messages in BranchCheckpoint (Commit Step)

As noted in the BranchCheckpoint definition, storing the full conversation history at each checkpoint is redundant because rebuild_messages can reconstruct it on demand. Removing covered_messages here avoids O(N^2) memory overhead.

        checkpoint = BranchCheckpoint(
            trajectory_buffer=trajectory_buffer,
            request_tools=request_tools,
            chat_template_kwargs_key=chat_template_kwargs_key,
            image_data=image_data,
            video_data=video_data,
            extra_fields=extra_fields,
        )

[gemini-code-assist]

Simplify messages retrieval by directly calling self._trie.rebuild_messages(node) now that messages is removed from BranchCheckpoint.

Suggested change

	messages = checkpoint.messages or self._trie.rebuild_messages(node)
	messages = self._trie.rebuild_messages(node)

[gxlvera]

With prefix trie, I think maybe we can get rid of generation lock but only use request lock? the reason is stated in details in: #51 as

One shared state requires a generation lock and serial LLM calls. With a trie, each call owns a cloned branch state; tokenize and commit can interleave—supporting sub-agents, best-of-n, etc.

[zhanghuiyao]

Two trajectory parity comments from local review.

[zhanghuiyao]

P1: when prepared.request_tools != tools, this forces a full re-encode but does not materialize/export the previous checkpoint. The new checkpoint is still committed as a descendant, while finalize exports only terminal checkpoints, so the parent trajectory's trainable response tokens are dropped. This should either create an explicit export boundary, or be documented/tested as an intentional flag-on/off parity exception.

[ChangyiYang]

Good catch — confirmed this is a real correctness gap, not just a parity nit. Here's the fix we landed for it.

Root cause. finalize exported only terminal checkpoints (a committed node with no committed descendant), which silently assumes a deeper checkpoint contains all of its ancestors' trainable tokens. That holds for an incremental continuation (the child buffer is clone(parent) + delta, so the parent's mask=1 tokens are carried forward), but not for a full re-encode: on a tools change the child re-tokenizes the whole history, so the parent's generated tokens become non-trainable prompt in the child and remain mask=1 only on the parent. Since the parent isn't terminal, it was dropped — exactly the trainable tokens the legacy flow preserves via materialized_trajectory.

Fix (absorb-by-incremental-edge). Instead of inferring "absorbed" from tree depth, track it from how the buffer was actually built:

• TrieNode.absorbed: bool
• commit(..., incremental: bool): when incremental (buffer cloned from the nearest ancestor checkpoint and extended), mark that ancestor absorbed = True. A full re-encode marks nothing.
• iter_export_nodes yields every checkpoint with absorbed == False (the old _collect_terminals post-order pass is removed).

use_incremental — the same prepared.request_tools == tools gate flagged here — is threaded from the session into commit.

Why it handles the case. On a mid-branch tools change the re-encoded child does not mark its prefix-ancestor absorbed, so both export as independent trajectories: the pre-change turn keeps its mask=1 tokens (== legacy's materialize-on-context-break), and the re-encoded turn carries the prior turns as mask=0 prompt. It's also robust to the mixed case (a node with both an incremental child and a full-re-encode child): the incremental child absorbs the ancestor, the re-encode child doesn't, so the ancestor's trainable tokens land in exactly one exported trajectory — no drop, no double-count.

Added a regression test (test_trie_exports_pre_reencode_branch_on_tools_change) that is red on the old deepest-terminal rule (1 trajectory, turn-1 generation lost) and green after (2 trajectories, turn-1 mask=1 preserved). Full gateway suite stays green.

Scoped this to the export-correctness issue (this comment); the retry-vs-best-of-N distinction in your other comment is a separate axis and we're treating it independently.

[gxlvera]

I think Changyi's explanation makes sense

[zhanghuiyao]

P1/P2: using only MessageKey here conflates idempotent retry with independent best-of-N samples. If two sampled completions are token-identical, the second commit reuses this child and overwrites the checkpoint, so trie mode exports 1 trajectory where legacy keeps 2 samples. Training semantics should distinguish retry from duplicate samples, or document/test this as an intentional contract.

[gxlvera]

I have a question, for best-of-N samples, is it really necessary to keep the two identical trajectory?

[gxlvera]

Also, what do you mean by "Training semantics should distinguish retry from duplicate samples"? Thank you~

[zhanghuiyao]

By "distinguish retry from duplicate samples", I mean:

• retry: the same logical generation attempt should be idempotent and may reuse/overwrite the same node;
• duplicate sample: a separate best-of-N attempt should preserve multiplicity, even if the generated tokens are identical, unless we explicitly define trie mode as deduplicating outputs.

[zhanghuiyao]

But according to our discussion, a better approach might be to only keep the selected one 😊