[train] feat: Multiple Linear Chains for Gateway Sessions

docs/source/start/agent_train.md

@@ -152,6 +152,8 @@ Start with the script defaults, then tune these first:
 - `max_prompt_length`, `max_response_length`: context budget for the agent trajectory.
 - `staleness_threshold`, `trigger_parameter_sync_step`, `require_batches`, `partial_rollout`: fully async scheduling and weight synchronization behavior.
 
+Gateway sessions keep multiple active linear chains by default for subagent, compaction, and retry-style flows. Backend generation calls within one session can overlap by default, while prepare, decode, and commit remain serialized under the session `request_lock`. Successful outcomes are accepted in completion order, so the final accepted trajectory used as the scoring target can be timing-dependent. Its score is broadcast to all trajectories from that session.
+
 For MoE or large models, also check tensor, pipeline, context, and expert parallelism settings such as `GEN_TP`, `TP`, `PP`, `CP`, and `EP` in `train_qwen3p5_moe.sh`.
 
 ---

examples/gateway/debug_launcher.py

@@ -468,15 +468,29 @@ def capture_debug_snapshot(actor: _GatewayActor, session_id: str, urls: dict[str
             "session_state": None,
             "message_history": [],
             "active_tool_schemas": None,
+            "active_chains": [],
             "provider_urls": dict(urls),
             "note": "session was not available when debug snapshot was captured",
         }
-    return {
+    active_chains = [
+        {
+            "chain_id": chain.chain_id,
+            "message_history": list(chain.message_history),
+            "active_tool_schemas": chain.active_tool_schemas,
+        }
+        for chain in session.active_chains
+    ]
+    snapshot = {
         "session_state": session.snapshot_state(),
-        "message_history": list(session.message_history),
-        "active_tool_schemas": session.active_tool_schemas,
+        # Keep the original single-chain fields for existing debug consumers.
+        "message_history": active_chains[0]["message_history"] if len(active_chains) == 1 else [],
+        "active_tool_schemas": active_chains[0]["active_tool_schemas"] if len(active_chains) == 1 else None,
+        "active_chains": active_chains,
         "provider_urls": dict(urls),
     }
+    if len(active_chains) > 1:
+        snapshot["note"] = "message_history and active_tool_schemas are omitted when multiple chains are active"
+    return snapshot
 
 
 async def run_claude_once(
@@ -594,6 +608,7 @@ async def run_debug_session_once(
         "session_state": None,
         "message_history": [],
         "active_tool_schemas": None,
+        "active_chains": [],
         "provider_urls": {},
         "note": "session was not created before debug snapshot was captured",
     }

tests/uni_agent/framework/test_generate_sequences_on_cpu.py

@@ -83,7 +83,16 @@ async def _build_framework_with_agent_runners(
             "actor_rollout_ref": {
                 "rollout": {
                     "n": n,
-                    "val_kwargs": {"n": val_n},
+                    "temperature": 0.7,
+                    "top_p": 0.8,
+                    "top_k": 20,
+                    "calculate_log_probs": True,
+                    "val_kwargs": {
+                        "n": val_n,
+                        "temperature": 0,
+                        "top_p": 0.95,
+                        "top_k": -1,
+                    },
                     "custom": {"agent_framework": agent_framework_cfg},
                 }
             }
@@ -96,6 +105,65 @@ async def _build_framework_with_agent_runners(
     )
 
 
+@pytest.mark.parametrize(
+    ("data_config", "expected_chat_template_kwargs"),
+    [
+        ({}, {}),
+        ({"apply_chat_template_kwargs": {"thinking": True}}, {"thinking": True}),
+    ],
+)
+def test_build_gateway_manager_wires_gateway_config_defaults(
+    monkeypatch,
+    data_config,
+    expected_chat_template_kwargs,
+):
+    from omegaconf import OmegaConf
+
+    from uni_agent.framework import entry as entry_module
+
+    class _ModelConfig:
+        tokenizer = object()
+        processor = None
+
+    captured = {}
+
+    class _FakeGatewayManager:
+        def __init__(self, *, llm_client, gateway_count, gateway_actor_config):
+            captured["llm_client"] = llm_client
+            captured["gateway_count"] = gateway_count
+            captured["gateway_actor_config"] = gateway_actor_config
+
+    monkeypatch.setattr(entry_module, "omega_conf_to_dataclass", lambda _config: _ModelConfig())
+    monkeypatch.setattr(entry_module, "GatewayManager", _FakeGatewayManager)
+
+    llm_client = object()
+    config = OmegaConf.create(
+        {
+            "data": data_config,
+            "actor_rollout_ref": {
+                "model": {},
+                "rollout": {
+                    "prompt_length": 128,
+                    "response_length": 64,
+                    "multi_turn": {"format": "hermes"},
+                    "custom": {"agent_framework": {"gateway_count": 2}},
+                },
+            },
+        }
+    )
+
+    manager = entry_module.build_gateway_manager(config=config, llm_client=llm_client)
+
+    assert isinstance(manager, _FakeGatewayManager)
+    assert captured["llm_client"] is llm_client
+    assert captured["gateway_count"] == 2
+    assert captured["gateway_actor_config"].prompt_length == 128
+    assert captured["gateway_actor_config"].response_length == 64
+    assert captured["gateway_actor_config"].tool_parser_name == "hermes"
+    assert isinstance(captured["gateway_actor_config"].apply_chat_template_kwargs, dict)
+    assert captured["gateway_actor_config"].apply_chat_template_kwargs == expected_chat_template_kwargs
+
+
 class _FakeTransferQueue:
     def __init__(self):
         self.puts = []
@@ -131,6 +199,7 @@ class _FakeGatewayManager:
     def __init__(self, finalized_by_session_prefix: dict[str, list[Trajectory]]):
         self._finalized_by_prefix = finalized_by_session_prefix
         self.created_sessions = []
+        self.created_session_kwargs = []
         self.finalized_sessions = []
         self.aborted_sessions = []
 
@@ -142,6 +211,7 @@ def _lookup(self, session_id: str) -> list[Trajectory]:
 
     async def create_session(self, session_id: str, **kwargs):
         self.created_sessions.append(session_id)
+        self.created_session_kwargs.append(dict(kwargs))
         return SessionHandle(
             session_id=session_id,
             base_url=f"http://fake/{session_id}/v1",
@@ -156,20 +226,29 @@ async def abort_session(self, session_id: str) -> None:
         self.aborted_sessions.append(session_id)
 
 
-def _build_prompts(count: int = 2, *, global_steps: int = 7, validate: bool = False):
+def _build_prompts(
+    count: int = 2,
+    *,
+    global_steps: int = 7,
+    validate: bool = False,
+    do_sample: bool | None = None,
+):
     non_tensor_dict = {"global_steps": global_steps}
     if validate:
         non_tensor_dict["validate"] = True
+    tensor_dict = {
+        "raw_prompt": [[{"role": "user", "content": f"sample {i}"}] for i in range(count)],
+        "uid": [f"uid-{i}" for i in range(count)],
+        "data_source": ["deepeyes"] * count,
+        "reward_model": [{"ground_truth": f"answer-{i}"} for i in range(count)],
+        "extra_info": [{"index": i} for i in range(count)],
+        "tools_kwargs": [{"tool": i} for i in range(count)],
+        "agent_name": ["deepeyes"] * count,
+    }
+    if do_sample is not None:
+        tensor_dict["__do_sample__"] = [do_sample] * count
     return tu.get_tensordict(
-        tensor_dict={
-            "raw_prompt": [[{"role": "user", "content": f"sample {i}"}] for i in range(count)],
-            "uid": [f"uid-{i}" for i in range(count)],
-            "data_source": ["deepeyes"] * count,
-            "reward_model": [{"ground_truth": f"answer-{i}"} for i in range(count)],
-            "extra_info": [{"index": i} for i in range(count)],
-            "tools_kwargs": [{"tool": i} for i in range(count)],
-            "agent_name": ["deepeyes"] * count,
-        },
+        tensor_dict=tensor_dict,
         non_tensor_dict=non_tensor_dict,
     )
 
@@ -179,6 +258,7 @@ def _trajectory(
     prompt_ids: list[int] | None = None,
     response_ids: list[int] | None = None,
     response_logprobs: list[float] | None = None,
+    reward_info: dict[str, object] | None = None,
     num_turns: int = 2,
     extra_fields: dict[str, object] | None = None,
 ):
@@ -189,6 +269,7 @@ def _trajectory(
         response_ids=response_ids,
         response_mask=[1] * len(response_ids),
         response_logprobs=response_logprobs,
+        reward_info=dict(reward_info or {}),
         reward_score=None,
         num_turns=num_turns,
         multi_modal_data={"images": ["raw-image-should-not-be-written"]},
@@ -258,6 +339,68 @@ async def test_agent_runners_registry_materializes_runners_and_selects_by_agent_
     assert all("gateway_manager" not in call["kwargs"] for call in calls)
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("validate", "do_sample", "expected_sampling_params"),
+    [
+        (
+            False,
+            None,
+            {
+                "temperature": 0.7,
+                "top_p": 0.8,
+                "top_k": 20,
+                "repetition_penalty": 1.0,
+                "logprobs": True,
+            },
+        ),
+        (
+            True,
+            None,
+            {
+                "temperature": 0,
+                "top_p": 0.95,
+                "top_k": -1,
+                "repetition_penalty": 1.0,
+                "logprobs": True,
+            },
+        ),
+        (
+            False,
+            False,
+            {
+                "temperature": 0,
+                "top_p": 1.0,
+                "top_k": -1,
+                "repetition_penalty": 1.0,
+                "logprobs": True,
+            },
+        ),
+    ],
+)
+async def test_framework_binds_sampling_defaults_to_gateway_sessions(
+    fake_tq,
+    validate,
+    do_sample,
+    expected_sampling_params,
+):
+    runtime = _FakeGatewayManager({"session-0-0": [_trajectory()]})
+    framework = await _build_framework_with_agent_runners(
+        agent_runners={"runner": _inline_runner_config(_async_noop_runner)},
+        gateway_manager=runtime,
+    )
+
+    await framework.generate_sequences(
+        _build_prompts(
+            count=1,
+            validate=validate,
+            do_sample=do_sample,
+        )
+    )
+
+    assert runtime.created_session_kwargs == [{"sampling_params": expected_sampling_params}]
+
+
 @pytest.mark.asyncio
 async def test_generate_sequences_writes_tq_schema_for_each_session(monkeypatch, fake_tq):
     """Full ``generate_sequences`` path writes one TQ batch per successful

tests/uni_agent/gateway/test_debug_launcher.py

@@ -3,6 +3,7 @@
 import os
 import subprocess
 from pathlib import Path
+from types import SimpleNamespace
 
 import httpx
 import pytest
@@ -137,6 +138,31 @@ def test_write_debug_snapshot_json_writes_message_history_and_session_state(tmp_
     assert record["message_history"] == [{"role": "user", "content": "hi"}]
 
 
+def test_capture_debug_snapshot_preserves_multiple_active_chains():
+    chains = [
+        SimpleNamespace(
+            chain_id=1,
+            message_history=[{"role": "user", "content": "main"}],
+            active_tool_schemas=None,
+        ),
+        SimpleNamespace(
+            chain_id=2,
+            message_history=[{"role": "user", "content": "subagent"}],
+            active_tool_schemas=[{"type": "function"}],
+        ),
+    ]
+    session = SimpleNamespace(active_chains=chains, snapshot_state=lambda: {"num_active_chains": 2})
+    actor = SimpleNamespace(_sessions={"s-debug": session})
+
+    snapshot = debug_launcher.capture_debug_snapshot(actor, "s-debug", {"openai_base_url": "http://local"})
+
+    assert [chain["chain_id"] for chain in snapshot["active_chains"]] == [1, 2]
+    assert snapshot["active_chains"][1]["message_history"][0]["content"] == "subagent"
+    assert snapshot["message_history"] == []
+    assert snapshot["active_tool_schemas"] is None
+    assert "multiple chains" in snapshot["note"]
+
+
 def test_provider_urls_strips_v1_for_anthropic_and_keeps_openai_base_url():
     """The printed Anthropic base URL omits /v1 for Claude Code while the
     OpenAI-compatible URL keeps the session-scoped /v1 suffix."""
@@ -560,6 +586,7 @@ def fake_run(cmd, *, env, check, text, capture_output, timeout):
     debug_snapshot = json.loads(result.debug_snapshot_path.read_text())
     assert debug_snapshot["session_state"]["has_active_trajectory"] is True
     assert debug_snapshot["message_history"][0] == {"role": "user", "content": "Reply with OK only."}
+    assert debug_snapshot["active_chains"][0]["message_history"] == debug_snapshot["message_history"]
 
 
 @pytest.mark.asyncio