Rollout Server & API

Relevant source files

The following files were used as context for generating this wiki page:

• osmosis_ai/rollout/server/app.py
• osmosis_ai/rollout/server/serve.py
• osmosis_ai/rollout/core/schemas.py
• osmosis_ai/rollout/core/base.py
• osmosis_ai/rollout/config/settings.py
• osmosis_ai/rollout/server/api_key.py

The Rollout Server is a specialized FastAPI-based application component within the Osmosis SDK designed to host and execute RolloutAgentLoop implementations. Its primary purpose is to provide a standardized interface for TrainGate (the training orchestrator) to initiate agent rollouts, manage their lifecycle, and receive final results and metrics asynchronously.

The server handles protocol complexities such as concurrency control, idempotency for rollout requests, background task management, and optional platform registration. It acts as a bridge between the local agent logic and the remote Osmosis training infrastructure, ensuring that agent trajectories are captured and reported consistently.

Sources: osmosis_ai/rollout/server/app.py:1-40, osmosis_ai/rollout/server/serve.py:1-15, osmosis_ai/rollout/core/base.py:270-300

Architecture and Lifecycle

The Rollout Server operates on an asynchronous, event-driven architecture. When a rollout is initialized, the server immediately returns the available tools and then processes the agent logic in the background.

Server Lifecycle Flow

The following diagram illustrates the interaction between TrainGate and the Rollout Server during a typical execution.

The server uses a FastAPI lifespan context manager to manage startup (cleanup tasks, platform registration) and shutdown (canceling pending rollouts, releasing resources).

Sources: osmosis_ai/rollout/server/app.py:126-170, osmosis_ai/rollout/server/app.py:252-330

Core API Endpoints

The server exposes several endpoints to manage the rollout lifecycle and monitor health.

Endpoint Summary

Endpoint	Method	Description	Auth Required
/v1/rollout/init	POST	Initializes a new rollout session and starts background execution.	Yes (API Key)
/health	GET	Basic health check returning agent name and rollout counts.	No
/platform/health	GET	Authenticated health check for Osmosis Platform validation.	Yes (API Key)

Sources: osmosis_ai/rollout/server/app.py:217-250, osmosis_ai/rollout/server/app.py:252-260

Rollout Initialization (/v1/rollout/init)

This is the primary entry point for rollouts. It consumes a RolloutRequest and performs the following logic:

1. Authentication: Validates the Bearer token against the server's configured API key.
2. Idempotency: Checks if a rollout with the same idempotency_key or rollout_id is already active or recently completed.
3. Tool Discovery: Calls agent_loop.get_tools() to determine capabilities.
4. Task Dispatch: Launches the agent_loop.run() coroutine in the background using an asyncio.Task.
5. Response: Returns 202 Accepted with the tool definitions.

Sources: osmosis_ai/rollout/server/app.py:265-334, osmosis_ai/rollout/core/schemas.py:465-498

State Management and Concurrency

The server maintains internal state through the AppState class, which tracks active tasks and handles concurrency limits.

Component Diagram

The server uses an asyncio.Semaphore to enforce the max_concurrent_rollouts limit. If the limit is reached, background tasks wait until a slot becomes available. Completed records are maintained for a duration defined by record_ttl_seconds to support idempotency before being pruned by a background cleanup task.

Sources: osmosis_ai/rollout/server/app.py:100-115, osmosis_ai/rollout/config/settings.py:115-130

Security and Authentication

Authentication is handled via a simple API key mechanism. The server expects a Bearer token in the Authorization header.

• API Key Generation: The server can generate URL-safe API keys using secrets.token_urlsafe(32).
• Validation: Uses hmac.compare_digest for constant-time comparison to prevent timing attacks.
• Protocol: Requests from TrainGate to the Rollout Server MUST include the configured API key. Conversely, the RolloutRequest may contain an api_key that the Rollout Server uses to authenticate its own callbacks to TrainGate.

Sources: osmosis_ai/rollout/server/api_key.py:15-55, osmosis_ai/rollout/server/app.py:65-75

Configuration

Server behavior is controlled through RolloutServerSettings, which can be set via environment variables prefixed with OSMOSIS_ROLLOUT_SERVER_.

Setting	Default	Description
max_concurrent_rollouts	100	Limit on simultaneous active background tasks.
record_ttl_seconds	3600.0	Time to retain rollout IDs for idempotency checks.
cleanup_interval_seconds	60.0	Frequency of the internal pruning task.
request_timeout_seconds	600.0	Timeout for individual background operations.

Sources: osmosis_ai/rollout/config/settings.py:108-160

Execution Context

The RolloutContext is the primary interface provided to the agent during execution. It encapsulates:

• Metrics Tracking: Automatically records LLM latency, tool latency, and token usage.
• Debug Logging: If debug_dir is configured, it writes detailed event traces to {debug_dir}/{rollout_id}.jsonl.
• LLM Communication: Provides a chat() method that delegates to the specialized OsmosisLLMClient for callbacks to TrainGate.

Sources: osmosis_ai/rollout/core/base.py:95-150, osmosis_ai/rollout/core/base.py:210-260

Summary

The Rollout Server provides a robust, production-ready environment for executing LLM agents within the Osmosis ecosystem. By standardizing the API and lifecycle management, it allows developers to focus on agent logic while the SDK handles infrastructure concerns like concurrency, authentication, and asynchronous reporting of results to the training pipeline.