feat(rollout_gateway): add inference api & backend-agnostic token-level trajectory capture

[luyuzhe111]

What

Adds agentcore_rl_toolkit.rollout_gateway — an in-repo layer that captures loss-maskable, token-level trajectories from agent rollouts for RL training. It's the successor to the external rllm-model-gateway dependency the current backends/{slime,verl} integrations rely on.

Why

RL training needs per-token ids, logprobs, and a loss mask for every model turn — not just final text. The current path (rllm-model-gateway) proxies /v1/chat/completions and scrapes token_ids out of the vLLM response. This gateway instead owns tokenization: an agent points its client at it, and the gateway renders messages → token_ids, calls a token-in/token-out backend, and records the trajectory as a side effect.

Improvements over rllm-model-gateway

This replaces rllm-model-gateway; each point below is a concrete limitation it has that this design removes.

Aspect	rllm-model-gateway	this gateway
Inference backends	vLLM only — token_ids scraped from vLLM-shaped responses	Backend-agnostic SamplingBackend seam: vLLM, SGLang (HTTP), Tinker (in-process SDK). New engine = one class.
Wire protocols	OpenAI /v1/chat/completions only	OpenAI Chat Completions and Anthropic Messages, normalized to one canonical form
Multi-turn / sub-agents	Linear, flat prefix-extension contract; a turn that breaks prefix-extension is added as a separate turn without cumulative mode	Per-session message tree with CLEAN/REALIGN/FORK drift healing — branching, parallel tool calls, and dynamically spawned sub-agents are captured, not dropped
Tokenization authority	Renders only in an opt-in cumulative_token_mode; otherwise vLLM server's chat completion endpoint	Always renders in-process (single authority) → loss-mask boundaries exact, no retokenization drift; also lets sample-only backends (Tinker) work
Session identity	In the URL path (/sessions/{sid}/v1/...), needs a create_session→get_session_url handshake	In the api-key/Bearer slot; fixed base_url, no per-session URLs, no handshake
Storage / deps	SQLite + HTTP-CRUD store; standalone service	In-memory, keyed by session id; runs as a thread in the trainer (in-process trace read). Torch-free, aiohttp-free core
Output boundary	gateway-specific TraceRecord shape	framework-neutral TraceRecord → converted to each backend's native sample type in that backend's own process

What's included

• Core (torch-free): TraceRecord, TrajectoryManager (per-session message tree: multi-turn, parallel tool calls, CLEAN/REALIGN/FORK drift healing).
• Renderer seam: HfTemplateRenderer (default; dependency-free regex + </think> parsing) and TinkerRenderer.
• Sampling backend seam: VllmHttpBackend, SglangHttpBackend,TinkerSdkBackend.
• Wire-protocol adapters: OpenAI (/v1/chat/completions), Anthropic (/v1/messages). An agent framework speaks its native protocol — Strands / OpenAI-style clients hit Chat Completions, Claude Code / Anthropic clients hit Messages — and drives the same gateway unmodified (change base_url, nothing else). This is what lets a single training run capture trajectories from heterogeneous agents across APIs; both protocols normalize to one canonical message form and fold into the same per-session trajectory tree.
• RolloutGateway: assembles the above on one aiohttp app, all adapters sharing one TrajectoryManager.
• Docs (AGENTS.md) + [gateway] extra + dev/CI wiring.

Dependencies

Base install unchanged. Gateway deps (aiohttp, transformers) live behind the [gateway] extra and are added to [dev] so CI runs the suite. The package core imports torch-free and aiohttp-free (RolloutGateway is lazily exposed). Tinker (tinker, tinker-cookbook, Python ≥3.11) is installed manually rather than as an extra, to keep the project's ≥3.10 floor resolvable.

Scope / not included

This PR is the capture layer only. Wiring it into a training backend (the per-backend rollout function, TraceRecord → native sample conversion, and ACR dispatch/reward-join glue) is a follow-up — a prototype dispatcher is parked on wip/online-rl-dispatch until a backend consumes it end-to-end.

Testing

• uv run pytest tests/rollout_gateway/ — 14 tests: trajectory drift/masking, OpenAI adapter pipeline (fake backend, over aiohttp test client), renderer round-trip, torch-free/aiohttp-free import assertion.
• VllmHttpBackend + full gateway verified end-to-end against a live vLLM server (multi-turn capture; captured prompt token_ids match a direct apply_chat_template, drift 0).
• ruff check clean; uv lock --locked in sync.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

[lyzustc]

Will we track the latest updates from slime for them?

[luyuzhe111]

good question! yea that's the plan — it's a manual re-sync, but the process is documented so it's not ad-hoc. AGENTS.md has a "Vendored from slime (upstream baselines)" section that records each vendored file, its slime source path, and the baseline commit (90c212b5) we lifted from. To check for upstream changes: git -C <slime> diff 90c212b5..HEAD -- slime/agent/<file>. Our copies are intentionally modified (torch-free, Sample→TraceRecord, injected backend/renderer seams), so a re-sync is a review-and-port, not a merge — and we bump the recorded baseline commit when we do one.

[lyzustc]

Is this a typo? "/healthz" -> "/health"

[luyuzhe111]

thanks for the careful review! it's actually not a typo — /healthz is the intentional Kubernetes-style liveness-probe convention (the trailing z is idiomatic; it came in with the lifted slime adapter)

[lyzustc]

Will RolloutGateway class be the place where gateway server is launched? If so I think when launching the server, we should provide just underlying vllm/sglang/tinker server url and model name, while adapter, backend, renderer and tokenizer should be auto created within the gateway.

[luyuzhe111]

Agreed! I was planning to add this ergonomic constructor in the integration PR rather than here so it can be directly tested.

The plan is a convenience classmethod like RolloutGateway.from_engine(engine_url=..., model=..., engine="vllm") that auto-builds backend/renderer/tokenizer for the HTTP-engine common case, while keeping the explicit __init__(backend=, renderer=, tokenizer=) for Tinker and custom setups. That lives naturally with the integration wiring (engine selection + config), which is the next PR — this one is the capture layer only : )