[train] feat: Multiple Linear Chains for Gateway Sessions by zhanghuiyao · Pull Request #66 · verl-project/uni-agent

zhanghuiyao · 2026-06-23T02:22:43Z

Summary

Implements per-session multiple linear trajectory chains for GatewaySession. Multiple chains are now the default and only session state model: live chains are stored in active_chains, closed chains in materialized_chains, and the previous single-active-chain state is removed.

This lets subagent/system splits, context compaction, repeated prompts, retries, and best-of-N-style sibling requests coexist without overwriting unrelated trajectory state. Same-session backend generation can overlap by default; request preparation, decoding, commits, and lifecycle transitions remain serialized by the session lock.

What's in this PR

Multiple-chain state and selection

Adds chain-local message history, cumulative SHA-256 prefix hashes, tool schemas, token buffers, multimodal state, logprob completeness, and ordering metadata.
Reuses the longest active chain only when the incoming normalized history is a proven prefix continuation and tool schemas match.
Uses (updated_seq, chain_id) as the deterministic tie-break for equally long compatible siblings.
Starts a new chain when no safe continuation exists and retains all chains until finalize.

Lifecycle and trajectory materialization

Clones working buffers during prepare and commits only after backend success.
Closes only the selected chain on length exhaustion without persisting unsent incremental tokens or media.
Rejects commits that return after finalize or abort.
Sorts early-closed and still-active chains by accepted outcome order before returning the existing Trajectory schema.
Keeps response logprobs all-or-none: aligned with response_ids or emitted as None.

Default same-session backend concurrency

Allows backend generation calls from one session to overlap by default.
Reserves a selected existing chain while its backend request is in flight, so concurrent requests cannot overwrite the same live chain.
Runs prepare, decode, commit, finalize, and abort under request_lock; reservations are released on success, failure, decode failure, cancellation, or late terminal-state rejection.
Uses session_id as the LLM client sticky-routing key. The LLM client still creates a unique server request ID for each generation call.
Accepts successful results in backend completion order. Concurrent siblings therefore share the session-level reward whose target can be timing-dependent.

Codec, sampling, and framework integration

Uses stable deterministic message-prefix hashing and canonicalizes tool-call arguments for comparison.
Preserves tool-call IDs and tool_call_id in prefix identity, so rewritten committed IDs split safely.
Uses configured data.apply_chat_template_kwargs consistently for full and incremental encoding; request-level chat_template_kwargs is not part of the public request contract.
Passes train, validation, and per-sample greedy sampling defaults from the framework into each gateway session while preserving allowed per-request overrides.
Keeps reward and TransferQueue schemas unchanged: the final finalized trajectory is scored and its score is broadcast to every trajectory from that session.

Tests

CPU coverage spans:

direct session behavior: linear parity, subagent return-to-main, compaction, sibling selection, prefix hashing, tools, length exhaustion, multimodal isolation, logprobs, backend failures, cancellation, and late finalize/abort results;
actor behavior: sampling and budget forwarding, default same-session backend overlap, OpenAI response envelopes, session state, multimodal flows, and error mapping;
manager/HTTP behavior: session routing, default multiple-chain finalization, and reward metadata;
framework behavior: gateway configuration, train/validation/greedy sampling, finalize ordering, reward target selection, broadcast scoring, and TransferQueue output.

Representative validation for the current PR head:

RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 .venv/bin/python -m pytest \
  tests/uni_agent/gateway/test_session_multiple_chains_on_cpu.py \
  tests/uni_agent/gateway/test_gateway_actor_on_cpu.py \
  tests/uni_agent/gateway/test_gateway_manager_on_cpu.py \
  tests/uni_agent/framework/test_generate_sequences_on_cpu.py -q
# 92 passed

git diff --check
# passed

Design

The core rule is: reuse a chain's token buffer only when the incoming normalized message history proves it is a safe continuation of that chain; otherwise create a sibling chain and preserve the existing state.

Backend generation is the only phase allowed to overlap within a session. Reserving selected chain IDs prevents two concurrent requests from committing over the same chain. Requests that cannot select a reserved chain prepare independently and materialize as siblings. Completion order defines order_seq, which also defines the final trajectory used by the existing session-level reward policy.

Scope / deferrals

No prefix trie, warm-start structural nodes, or nearest-checkpoint traversal.
No enable_multiple_chains or enable_parallel_session_generation feature flag; multiple chains and backend overlap are the default session model.
No ignore_cch_for_prefix_hash compatibility mode; historical cch= changes remain ordinary semantic prefix changes.
No public branch IDs, parent-tip hints, request idempotency keys, or branch metadata in OpenAI responses or Trajectory.extra_fields.
No per-chain reward policy and no Trajectory or TransferQueue schema changes.
No request-level chat_template_kwargs; chat-template configuration remains gateway/model scoped.

…eway sessions - Introduced `enable_multiple_chains` configuration option in GatewayActorConfig to toggle multiple chain support. - Updated GatewayManager and GatewayActor to handle multiple chains. - Enhanced GatewaySession to manage active chains, including selection, materialization, and closure of chains. - Implemented chain state management with ChainState and MaterializedChain data classes. - Refactored message encoding and response handling to accommodate multiple chains. - Added methods for computing message prefix hashes and managing chain compatibility.

gemini-code-assist

Code Review

This pull request introduces support for managing multiple active linear trajectory chains per gateway session, allowing for branching conversations (e.g., subagents) within a single session. The changes span the framework configuration, gateway actor, and session management, backed by a comprehensive suite of unit tests. The feedback identifies a potential issue where the calculated remaining response budget could become negative if the generated response exceeds the budget, which may cause LLM backends to crash; clamping this value to a minimum of zero is recommended.

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.

…ro and update MessageCodec for dynamic system prompts

- Removed enable_multiple_chains parameter from session and gateway configurations. - Simplified GatewaySession to handle multiple chains without conditional checks. - Updated MessageCodec to derive system prompts from processors when available. - Consolidated trajectory handling to streamline the materialization of active chains. - Enhanced tests for multiple chains to ensure consistent behavior and coverage. - Removed legacy code related to single chain handling, improving clarity and maintainability.

- Introduced `enable_parallel_session_generation` and `ignore_cch_for_prefix_hash` flags in GatewayActorConfig to control session behavior. - Updated GatewaySession to handle parallel backend generation and manage inflight generations. - Enhanced tests to cover new functionality, including concurrent requests and validation of boolean flags. - Modified message prefix hash computation to optionally ignore CCH markers for prefix comparisons. - Ensured backward compatibility while implementing new features.

…n decode failures

…and_trajectory_assembly

…and related configurations

…d GatewaySession

zackcxb · 2026-07-10T06:24:56Z

+            f"got {type(value).__name__}"
+        )
+    return value
+


consider inline this helper function. For example:
enable_parallel_session_generation = af_cfg.get("enable_parallel_session_generation", False)
if not isinstance(enable_parallel_session_generation, bool):
raise ValueError(...)

zackcxb · 2026-07-10T08:36:08Z

+        # present: encode_incremental() slices processor-produced ids by this length on the
+        # default path, and a tokenizer-derived prefix would mis-slice the continuation delta
+        # whenever the processor adds BOS / multimodal-aware prefix tokens the bare tokenizer
+        # never emits.


This docstr is redundant imo. The code is self-explanatory enough

zackcxb · 2026-07-10T08:48:57Z

+
+    omitted = codec.encode_incremental([{"role": "user", "content": "delta"}])
+    assert _PlainTokenizer().decode(omitted) == "user:delta\nassistant:"
+


The first several test cases seem redundant, and some seem artificial (e.g., adding an extra bos in the processor, and monkeypatch). Consider removing them

zackcxb · 2026-07-10T08:49:53Z

Add docstrs for tests

zackcxb · 2026-07-13T01:39:03Z

+            return await self._run_generation_prepared(payload, request_context, backend)
+
        async with self.generation_lock:
+            return await self._run_generation_prepared(payload, request_context, backend)


Is generation_lock still necessary?

zackcxb · 2026-07-13T03:34:20Z

-        materialized_trajectory = None
-        image_data = None
-        video_data = None
+        incoming_message_prefix_hashes = self._compute_message_prefix_hashes(messages)


consider inline this helper _compute_message_prefix_hashes

…d remove unused function

zackcxb · 2026-07-13T06:27:53Z

+        candidates = [
+            chain
+            for chain in self.active_chains
+            if self._is_chain_request_compatible(


consider inlining

zackcxb · 2026-07-13T07:00:09Z

+        ("enable_parallel_session_generation", 1),
+    ],
+)
+def test_build_gateway_manager_rejects_non_bool_m2_flags(flag_name, bad_value):


seems redundant.

zackcxb · 2026-07-13T07:02:49Z



+@pytest.mark.asyncio
+async def test_generate_sequences_preserves_sorted_trajectory_order_and_rewards_final_target(fake_tq):


Is this really needed? We already have test cases for session order, TQ order, and reward target

zackcxb · 2026-07-13T07:03:41Z

+
+
+@pytest.mark.asyncio
+async def test_generate_sequences_preserves_normal_multiple_chain_order_and_rewards_final_main(fake_tq):


Is this really needed? focus on functionalities that are yet covered by other tests

zackcxb · 2026-07-13T07:04:13Z

+
+
+@pytest.mark.asyncio
+async def test_score_trajectories_uses_last_finalized_trajectory_as_reward_target():


Seems redundant

zackcxb · 2026-07-13T07:15:36Z

+        enable_parallel_session_generation: bool = False,
    ):
        """Create an active session bound to a handle and model codec."""
+        if type(enable_parallel_session_generation) is not bool:


this has already been validated in config.py

- Removed the `enable_parallel_session_generation` flag from session creation and related tests. - Updated session handling to support concurrent backend generation without the need for the flag. - Refactored tests to align with the new session management approach, ensuring proper handling of concurrent requests and chain reservations. - Enhanced the `GatewaySession` class to manage reserved chains and streamline the generation process. - Adjusted the codec and framework to accommodate the removal of the parallel generation flag while maintaining functionality.

zackcxb · 2026-07-14T07:56:55Z

+        # Same-session requests overlap backend generation and commit in backend
+        # completion order. The framework currently scores session_trajectories[-1]
+        # and broadcasts that reward, so concurrent siblings share one reward target.
+        return await self._run_generation_prepared(payload, request_context, backend)


remove this wrapper

zackcxb · 2026-07-14T08:19:54Z

+            chain
+            for chain in self.active_chains
+            if chain.chain_id not in self.reserved_chain_ids
+            and self._is_chain_request_compatible(


this helper can be replaced by chain.active_tool_schemas == tools

zackcxb · 2026-07-14T08:28:10Z

+            incoming_message_prefix_hashes=list(incoming_message_prefix_hashes),
+            logprobs_complete=logprobs_complete,
        )



move these public api to the beginning of this class, they are currently interrupting the code reading flow. For example:
init
run_generation
set_reward_info
finalize
abort
snapshot_state

_prepare_generation_inputs
_select_chain
...

… GatewaySession

zackcxb · 2026-07-14T09:50:54Z

+        else:
+            buffer = self._copy_trajectory_buffer(selected_chain.buffer)
+            image_data, video_data = self._copy_chain_media(selected_chain)
+            logprobs_complete = selected_chain.logprobs_complete


Could we avoid carrying logprobs_complete through both ChainState and EncodedData? logprobs is already a session-level sampling invariant (rollout.calculate_log_probs is forwarded through the session sampling params), and it is not request-overridable by default.
When the effective sampling params request logprobs, a missing or length-mismatched output.log_probs seems like a backend contract violation and should fail fast. When logprobs are disabled, the trajectory can simply expose response_logprobs=None. The buffer can then keep the simpler invariant that response_logprobs is either empty or aligned with response_ids, without duplicating completeness state across prepare, commit, and materialization. Incremental masked tokens should use the same effective logprobs flag when appending 0.0.

zackcxb · 2026-07-14T12:20:30Z

+                    logprobs_complete=logprobs_complete,
+                )
+
+            if incremental_messages:


This if-check can be removed. Relevant variables are initialized as empty:
incremental_ids = []
new_image_data = None
new_video_data = None

…tests

zackcxb · 2026-07-15T03:57:56Z

-        sampling_params = self._codec.build_sampling_params(payload)
        remaining_response_budget = (
-            self._response_length - len(buffer.response_mask) if self._response_length is not None else None
+            max(0, self._response_length - len(buffer.response_mask)) if self._response_length is not None else None


this is not needed, since self._response_length < len(buffer.response_mask) is already gated by
if self._response_length is not None
and len(buffer.response_mask) + len(incremental_ids) >= self._response_length

zackcxb · 2026-07-15T07:04:25Z

+                        stop_reason=output.stop_reason,
+                    )
+                except Exception as e:
+                    raise HTTPException(status_code=500, detail=f"{e.__class__.__name__}: {e}") from e


This exception handler is not quite needed. I will also add a global handler for internal errors in the Anthropic PR

zackcxb · 2026-07-15T12:46:26Z

                    sampling_params=encoded.sampling_params,
-                    image_data=encoded.image_data,
-                    video_data=encoded.video_data,
+                    image_data=self._copy_media_list(encoded.image_data),


Why do we need copying here?

…ains # Conflicts: # tests/uni_agent/gateway/test_gateway_actor_on_cpu.py # uni_agent/gateway/gateway.py # uni_agent/gateway/session/codec.py # uni_agent/gateway/session/protocol.py # uni_agent/gateway/session/session.py

zackcxb · 2026-07-16T11:41:07Z

        "messages": [_normalize_message(message) for message in messages],
        "tools": tools,
-        "chat_template_kwargs": dict(chat_template_kwargs) if chat_template_kwargs else {},
+        "chat_template_kwargs": {},


This can be removed altogether, along with the chat_template_kwargs fields in TypedDict

gemini-code-assist Bot reviewed Jun 23, 2026

View reviewed changes

Comment thread uni_agent/gateway/session/session.py

zhanghuiyao marked this pull request as draft June 23, 2026 02:26

zhanghuiyao changed the title ~~[AgentGateway] feat: add support for multiple active linear trajectory chains in gat…~~ [AgentGateway] feat: Multiple Linear Chains for Gateway Sessions Jun 23, 2026

zhanghuiyao added 3 commits June 23, 2026 12:10

feat: enhance GatewaySession to clamp remaining response budget to ze…

47bb69a

…ro and update MessageCodec for dynamic system prompts

zhanghuiyao marked this pull request as ready for review June 25, 2026 02:55

fix: Update error handling in GatewaySession to raise HTTPException o…

6634f6d

…n decode failures

zhaizhiqiangA requested a review from yyDing1 July 1, 2026 08:48

zhanghuiyao added 3 commits July 1, 2026 17:45

feat: add test_multiple_chains_claude_code_subagent_cch_prefix_match_…

589dde1

…and_trajectory_assembly

refactor: remove ignore_cch_for_prefix_hash flag from GatewaySession …

88c3001

…and related configurations

refactor: streamline chat template kwargs handling in MessageCodec an…

55f57ab

…d GatewaySession

zackcxb reviewed Jul 13, 2026

View reviewed changes

refactor: simplify boolean flag retrieval in build_gateway_manager an…

388718b

…d remove unused function

zackcxb reviewed Jul 13, 2026

View reviewed changes

zhanghuiyao added 3 commits July 13, 2026 19:26

refactor: removing unused tests

ffe3ea3

Reject invalid response budgets and close exhausted sessions earlier

28deb35

zackcxb reviewed Jul 14, 2026

View reviewed changes

refactor: improve request handling and validation in MessageCodec and…

a2ab4bf

… GatewaySession

zackcxb reviewed Jul 14, 2026

View reviewed changes

refactor: streamline logprobs handling in GatewaySession and related …

1ba3592

…tests

zackcxb reviewed Jul 15, 2026

View reviewed changes

zhanghuiyao added 2 commits July 15, 2026 23:33

Simplify gateway chain state tracking and media handling

6d435c4

Refine GatewaySession generation and update tests

c46fa29

zhanghuiyao changed the title ~~[AgentGateway] feat: Multiple Linear Chains for Gateway Sessions~~ [train] feat: Multiple Linear Chains for Gateway Sessions Jul 16, 2026

zackcxb reviewed Jul 16, 2026

View reviewed changes

zhanghuiyao added 3 commits July 16, 2026 20:07

Remove request-level chat template kwargs from gateway adapters

9a45e71

Update gateway tests for length materialization and tool-call prefixing

9f394fb

Normalize tool calls as dictionaries

4b0f4f7

yyDing1 merged commit 1a1a11e into verl-project:main Jul 17, 2026
4 checks passed


		omitted = codec.encode_incremental([{"role": "user", "content": "delta"}])
		assert _PlainTokenizer().decode(omitted) == "user:delta\nassistant:"



		@pytest.mark.asyncio
		async def test_generate_sequences_preserves_sorted_trajectory_order_and_rewards_final_target(fake_tq):



		@pytest.mark.asyncio
		async def test_score_trajectories_uses_last_finalized_trajectory_as_reward_target():

Uh oh!

Conversation

zhanghuiyao commented Jun 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

What's in this PR

Tests

Design

Scope / deferrals

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zackcxb Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

zhanghuiyao commented Jun 23, 2026 •

edited

Loading

zackcxb Jul 13, 2026 •

edited

Loading