[Feature] Add agent gateway

[zackcxb]

What does this PR do?

This PR adds uni_agent.gateway — an OpenAI-compatible session gateway runtime for multi-turn agent-style rollout in uni-agent, as a downstream integration of verl RFC #5790 and the upstream agent framework PR verl#6299.

Specifically:

• _GatewayActor (~200 lines) — thin FastAPI actor: /v1/chat/completions over sticky sessions; OpenAI-compatible error envelopes; capability gates (n>1 / response_format / tool_choice=required|dict → 400); typed I/O via ChatCompletionRequest / ChatCompletionResponse protocol types; lifecycle endpoints (/complete, finalize, abort).
• GatewaySession — per-session state + run_generation envelope (encode → generate → decode+commit). Commit-on-success isolation: failed backend/decode calls never pollute session state. Session fields are written in a single point under request_lock with no await / IO / parsing inside the commit segment.
• MessageCodec — model-scoped encode/decode/tool-parse/multimodal/sampling-param, stateless, shared across all sessions.
• GatewayManager — session-id → gateway-actor routing with least-active load balancing.
• GatewayServingRuntime — owns gateway actor lifecycle, injects the LLMServerClient as a duck-typed backend into the actor.
• GatewayActorConfig — frozen dataclass carrying model-/length-scoped knobs from the rollout yaml. Backend is separate so the codec/session boundary has no view of the LLM client lifecycle.

Response bodies follow the OpenAI Chat Completions spec (id / object / created / model / choices / usage). Request/response shapes are defined as TypedDicts in gateway/protocol.py — no openai Python SDK runtime dependency.

Token-truth (the RL correctness contract) is preserved inline: backend-sampled token IDs accumulate in TrajectoryBuffer with a response_mask that distinguishes generated tokens (mask=1, trainable) from chat-template interstitials (mask=0, masked). Response logprobs are copied directly from the backend TokenOutput — no decode→re-encode.

PR scope

Per maintainer request this PR has been split into three stacked PRs:

1. gateway (this PR) — uni_agent.gateway
2. framework — uni_agent.framework (follow-up, stacked on this PR)
3. deepeyes examples — examples/agent_train/deepeyes_gateway/ (follow-up, stacked on framework)

Only the gateway portion is reviewed here. Shared types (SessionHandle, Trajectory) moved from the old framework.types into gateway/types.py so the gateway is fully self-contained. There is zero uni_agent.framework import in the gateway package.

Checklist Before Starting

• Search for similar PRs: https://github.com/verl-project/uni-agent/pulls?q=is%3Apr+gateway+framework
• Format the PR title: [gateway] feat: add OpenAI-compatible session gateway runtime

Test

pytest tests/uni_agent/gateway/ -q

41 passed, 6 warnings (gateway actor / manager / session-runtime / typed-I/O response shape, ~170s wall-clock).

Critical regression gates included:

• test_gateway_actor_backend_failure_does_not_commit_partial_state (commit-on-success isolation)
• test_gateway_actor_context_change_splits_trajectory (branch-c materialized trajectory recovery)
• test_gateway_actor_continuation_budget_exhausted_materializes_length_stop (length-budget exhaust)
• test_gateway_serving_runtime_owns_gateway_lifecycle_and_session_runtime (runtime wiring)

All pass.

API and Usage

Public API:

• uni_agent.gateway — GatewayServingRuntime, GatewayManager, GatewayActor, GatewayActorConfig, SessionHandle, Trajectory

Minimum wiring (framework PR will provide a build_agent_framework() helper):

from uni_agent.gateway.config import GatewayActorConfig
from uni_agent.gateway.runtime import GatewayServingRuntime

gateway_actor_config = GatewayActorConfig(
    tokenizer=model_config.tokenizer,
    processor=model_config.processor,
    tool_parser_name=config.rollout.multi_turn.format,
    prompt_length=config.rollout.prompt_length,
    response_length=config.rollout.response_length,
)

session_runtime = GatewayServingRuntime(
    llm_client=llm_client,
    gateway_count=int(af_cfg["gateway_count"]),
    gateway_actor_config=gateway_actor_config,
)

The runtime exposes a SessionRuntime-shaped surface (create_session / wait_for_completion / finalize_session / abort_session) consumed by the framework. Sessions communicate with the gateway via standard /v1/chat/completions requests:

handle = await session_runtime.create_session("session-id")
response = requests.post(
    f"{handle.base_url}/chat/completions",
    json={"model": "default", "messages": [{"role": "user", "content": "hello"}]},
)
await session_runtime.complete_session("session-id")
trajectories = await session_runtime.finalize_session("session-id")

Design & Code Changes

High-level structure:

• _GatewayActor — FastAPI routes, OpenAI error envelopes, capability gates, chat-completion JSON serialization (~200 lines).
• MessageCodec (codec.py) — model-scoped encode/decode/normalize/multimodal/sampling-param, stateless across sessions.
• GatewaySession (session.py) — per-session state, run_generation envelope, lifecycle methods. Returns GenerationOutcome business objects; never constructs HTTP responses.
• GatewayManager — session-to-actor routing via least-active count.
• GatewayServingRuntime — Ray actor lifecycle, backend injection, session-lifecycle delegation.

Request flow: HTTP request → _GatewayActor (capability gate) → GatewaySession.run_generation (encode → backend.generate → decode+commit) → GenerationOutcome → actor serializes ChatCompletionResponse JSON.

Key invariants:

• MessageCodec never touches session state.
• GatewaySession never constructs HTTP JSON.
• GenerationOutcome / TrajectoryBuffer carry token-truth; HTTP protocol types carry wire format.
• generation_lock serializes in-flight generation (implementation detail of the single-active trajectory model, not part of the public contract).

WIP / Follow-up

• GatewayServingRuntime + GatewayManager merge (runtime's thin session-method delegates can fold into the runtime)
• Anthropic Messages format support (/v1/messages) as a sibling route + protocol type — not yet needed
• Subagent support
• Turn-wise or trie storage

Checklist Before Submitting

• Read the Contribute Guide.
• Add unit tests to cover all new code — 41 CPU tests included, following the *_on_cpu.py naming convention.
• Public classes / methods / fields carry docstrings.
• Apply pre-commit checks: pre-commit install && pre-commit run --all-files
• Add / update documentation — inline docstrings ship with this PR; module-level documentation deferred to a follow-up.

[gemini-code-assist]

Code Review

This pull request introduces a comprehensive agent framework and gateway system designed to facilitate agentic workflows within a training environment. Key components include a factory for constructing frameworks, an OpenAI-compatible framework implementation that manages sequence generation and trajectory logging, and a gateway system that provides an OpenAI-compatible API for agent interactions. The gateway handles session lifecycle, trajectory buffering, and multimodal data processing. Feedback identifies a critical issue with the incremental token encoding logic in the gateway, which may produce malformed sequences due to assumptions about tokenizer stability and turn separators. Further recommendations include parallelizing reward calculations to improve performance and replacing blocking ray.get calls with asynchronous operations to avoid event loop starvation.

[gemini-code-assist]

The incremental encoding logic is fragile and likely to produce malformed token sequences. Slicing tokens based on the length of a pre-encoded system prompt assumes that the tokenizer is prefix-stable and that the chat template doesn't insert turn separators or special tokens between the system prompt and the first message. Furthermore, concatenating these incremental IDs to the previous turn's response IDs (at line 542) will miss the necessary turn separators (e.g., <|im_end|> and <|im_start|>user) required by most chat templates. It is safer to re-encode the full message history and identify the delta, or simply rely on the backend's prefix caching by sending the full prompt.

[gemini-code-assist]

Using ray.get inside an async context (called via build_agent_framework) will block the event loop, preventing other concurrent tasks from making progress. Since a helper _await_ray_ref is already defined in this file, you should consider moving the gateway startup logic to an async initialization method that can be awaited, rather than performing blocking calls in the constructor.

[wangtiance]

为什么放在trainer目录下？我觉得这是黑盒调用训推通用的流程。我偏向往上提一级，直接放uni_agent/framework和uni_agent/gateway.

[yyDing1]

The current entry point binds a single runner via agent_runner_fqn + agent_runner_kwargs. This works for a single-task recipe like DeepEyes, but it doesn't scale to multi-task rollout.

We may introduce an AgentRunner abstract base with a minimal run() contract:

# uni_agent/trainer/framework/runner.py
class AgentRunner(ABC):
    name: str = ""
    @abstractmethod
    async def run(
        self,
        *,
        raw_prompt: list[dict],
        session: SessionHandle,
        session_runtime: SessionRuntime,
        sample_index: int,
        tools_kwargs: dict[str, Any] | None = None,
    ) -> None:
        ...

Each sample carries the runner name; config mounts a name → runner map.

The config could be in the following format:

# agent_runner.yaml
- name: deepeyes
  _target_: examples.agent_train.deepeyes_gateway.runner.DeepEyesAgentRunner
  max_turns: 5
  tools:
    - name: image_zoom_in_tool
      config_path: examples/agent_train/deepeyes_gateway/configs/image_zoom_in_tool_config.yaml

- name: swe
  _target_: examples.agent_train.swe_gateway.runner.SweAgentRunner
  max_turns: 50
  env:
    deployment:
      type: vefaas
      command: ...
  tools:
    - name: str_replace_editor
    - name: execute_bash

Then the framework resolves the runner per-session by sample["agent_runner_name"], like:

# framework.py:_run_session
runner = self._runners_by_name[sample_fields["agent_runner_name"]]
await runner.run(
    raw_prompt=raw_prompt,
    session=session,
    session_runtime=self.session_runtime,
    sample_index=sample_index,
    tools_kwargs=sample_fields.get("tools_kwargs"),
)

This could be similar to verl's existing agent_loop_config pattern, and we can adopt the same shape here.

[gxlvera]

Hi, I would like to propose using Prefix Trie for multi-trajectory storage for Agentgateway. My RFC is here:#51
This approach could address the following limitations of current implementation:

• Single active branch only: A session keeps one message_history and one active trajectory. When switching sub-agents, picking a resample path, or returning to an older branch, new requests cannot reattach to historical branches. A trie keeps every branch; incoming messages longest-prefix-match against any path and continue from there.
• Repeated encoding of shared prefixes: Message/token prefixes shared across trajectories are re-materialized and re-tokenized on every branch switch. A trie stores checkpoints on shared nodes; later calls clone from the matched node and tokenize incrementally.
• No concurrent inference: 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.

For detailed explanation, please also refer to this comment: verl-project/verl#6299 (comment)

[sl-1314]

Hi, I noticed that, in the original verl AgentLoopManager/AgentLoopManagerTQ, it spawns num_workers independent AgentLoopWorker(Ray actors), distributing the total train_batch×rollout.n agent loops across these actors for parallelism. However, OpenAICompatibleAgentFramework.generate_sequences currently runs all train_batch×rollout.n sessions in a single asyncio event loop inside the PPOTrainer process. This means CPU-bound operations in one agent loop (e.g. tool execution) will block all other concurrent sessions.

Suggestion: Refer the original verl AgentLoopManager pattern — introduce multiple AgentLoopWorker Ray actors, partition the batch tasks across them.

# refer AgentLoopManager._init_agent_loop_workers()
for i in range(num_workers):
    worker = AgentLoopWorker.options(
        scheduling_strategy=NodeAffinitySchedulingStrategy(node_id=..., soft=True)
    ).remote(config, ...)
    self.workers.append(worker)

async def generate_sequences(self, prompts):
    chunks = prompts.chunk(len(self.workers))
    await asyncio.gather(*[
        w.generate_sequences.remote(chunk)
        for w, chunk in zip(self.workers, chunks)
    ])

This separates two independent concurrency axes: gateway_count for LLM serving throughput, num_workers for agent execution parallelism — consistent with the original verl design.

[wuxibin89]

1. P0: Add doc string for all public class, method, fields and functions.
2. P1: Separate this PR into 3 PRs: gateway, framework, deepeyes examples. We can only review the gateway part before sepration.

[wuxibin89]

Why not just decorate it with @ray.remote?

@ray.remote
class GatewayActor:
    ...

[zackcxb]

Hi, I noticed that, in the original verl AgentLoopManager/AgentLoopManagerTQ, it spawns num_workers independent AgentLoopWorker(Ray actors), distributing the total train_batch×rollout.n agent loops across these actors for parallelism. However, OpenAICompatibleAgentFramework.generate_sequences currently runs all train_batch×rollout.n sessions in a single asyncio event loop inside the PPOTrainer process. This means CPU-bound operations in one agent loop (e.g. tool execution) will block all other concurrent sessions.

Suggestion: Refer the original verl AgentLoopManager pattern — introduce multiple AgentLoopWorker Ray actors, partition the batch tasks across them.

# refer AgentLoopManager._init_agent_loop_workers()
for i in range(num_workers):
    worker = AgentLoopWorker.options(
        scheduling_strategy=NodeAffinitySchedulingStrategy(node_id=..., soft=True)
    ).remote(config, ...)
    self.workers.append(worker)

async def generate_sequences(self, prompts):
    chunks = prompts.chunk(len(self.workers))
    await asyncio.gather(*[
        w.generate_sequences.remote(chunk)
        for w, chunk in zip(self.workers, chunks)
    ])

This separates two independent concurrency axes: gateway_count for LLM serving throughput, num_workers for agent execution parallelism — consistent with the original verl design.

Hi, I noticed that, in the original verl AgentLoopManager/AgentLoopManagerTQ, it spawns num_workers independent AgentLoopWorker(Ray actors), distributing the total train_batch×rollout.n agent loops across these actors for parallelism. However, OpenAICompatibleAgentFramework.generate_sequences currently runs all train_batch×rollout.n sessions in a single asyncio event loop inside the PPOTrainer process. This means CPU-bound operations in one agent loop (e.g. tool execution) will block all other concurrent sessions.

Suggestion: Refer the original verl AgentLoopManager pattern — introduce multiple AgentLoopWorker Ray actors, partition the batch tasks across them.

# refer AgentLoopManager._init_agent_loop_workers()
for i in range(num_workers):
    worker = AgentLoopWorker.options(
        scheduling_strategy=NodeAffinitySchedulingStrategy(node_id=..., soft=True)
    ).remote(config, ...)
    self.workers.append(worker)

async def generate_sequences(self, prompts):
    chunks = prompts.chunk(len(self.workers))
    await asyncio.gather(*[
        w.generate_sequences.remote(chunk)
        for w, chunk in zip(self.workers, chunks)
    ])

This separates two independent concurrency axes: gateway_count for LLM serving throughput, num_workers for agent execution parallelism — consistent with the original verl design.

Thanks for the comment, this is a valid point. The current framework path fans out all
batch_size * rollout.n sessions with asyncio.gather inside the trainer process, and the
existing semaphore only caps coroutine concurrency; it does not isolate CPU-bound tool or
sandbox work. We have also seen this show up with SWE-style runners where sandbox work
can block unrelated sessions.

Since I have split the current PR into three separate PRs (gateway, framework, and examples), I will treat this as a framework-layer follow-up and address it in the next PR instead of this one.
For the framework PR, I plan to keep the reference OpenAICompatibleAgentFramework the lightweight inline implementation but documenting the dispatching mode. Heavy runners such as SWE can provide a framework subclass that overrides the session execution boundary and runs the agent runner
in process/Ray workers.