State Management & Schemas

Relevant source files

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

• osmosis_ai/rollout/core/schemas.py
• osmosis_ai/rollout/config/settings.py
• osmosis_ai/cli_services/config.py
• osmosis_ai/rollout/validator.py
• tests/test_rollout_server.py
• osmosis_ai/rollout/core/init.py

The Osmosis SDK utilizes a structured approach to state management and data validation through Pydantic schemas and dedicated state tracking components. This system ensures that interactions between the RolloutServer, the Osmosis Platform (TrainGate), and remote LLM providers are type-safe, idempotent, and measurable.

At the core of this architecture is the "single source of truth" for protocol messages, defining how rollouts are initialized, executed, and completed. State management components track the lifecycle of asynchronous agent loops, preventing duplicate executions and managing resource cleanup through Time-to-Live (TTL) mechanisms.

State Management Architecture

State management is primarily handled by the AppState class (as seen in server implementations and tests), which maintains the runtime status of active and completed rollouts.

Rollout Lifecycle Tracking

The system tracks rollouts through three primary states: STARTED, COMPLETED, and ERROR. The state manager uses thread-safe structures to store task references and completion timestamps.

AppState Components

Component	Type	Description
rollout_tasks	Dict[str, Task]	Maps rollout_id to active background tasks for cancellation and monitoring.
completed_rollouts	Dict[str, float]	Stores completion timestamps for idempotency checks.
record_ttl_seconds	float	Configurable duration before completed records are pruned.

Sources: tests/test_rollout_server.py:53-110, osmosis_ai/rollout/config/settings.py:114-118

Protocol Schemas

The SDK defines strict Pydantic models for the Osmosis remote rollout protocol. These schemas handle tool definitions, chat completions, and execution metrics.

Tool & Function Definitions

Osmosis follows the OpenAI function calling format for tool definitions. This allows agent loops to describe capabilities to the LLM in a standardized way.

Communication Schemas

Key request and response models govern the exchange of data between the server and the training platform.

Model	Direction	Description
RolloutRequest	TrainGate -> Server	Initializes a rollout with messages, params, and callback URLs.
InitResponse	Server -> TrainGate	Acknowledges initialization and provides available tools.
RolloutResponse	Server -> TrainGate	Reports final trajectory, rewards, and metrics upon completion.
CompletionsRequest	Server -> TrainGate	Requests LLM generation during the agent loop execution.

Sources: osmosis_ai/rollout/core/schemas.py:284-414

Configuration & Settings

The system supports granular configuration via environment variables and typed settings classes, utilizing pydantic-settings when available.

Global Settings

The RolloutSettings class aggregates client and server configurations.

class RolloutSettings(_BaseSettings):
    client: RolloutClientSettings = Field(default_factory=RolloutClientSettings)
    server: RolloutServerSettings = Field(default_factory=RolloutServerSettings)
    max_metadata_size_bytes: int = Field(default=1048576) # 1MB default

Sources: osmosis_ai/rollout/config/settings.py:165-195

Configurable Limits

Setting	Default	Environment Variable
Max Metadata Size	1 MB	OSMOSIS_ROLLOUT_MAX_METADATA_SIZE_BYTES
Max Concurrent Rollouts	100	OSMOSIS_ROLLOUT_SERVER_MAX_CONCURRENT_ROLLOUTS
Record TTL	3600s	OSMOSIS_ROLLOUT_SERVER_RECORD_TTL_SECONDS
Max Turns	10	N/A (Defined in RolloutRequest)

Sources: osmosis_ai/rollout/config/settings.py:114-124, osmosis_ai/rollout/core/schemas.py:255-265

Validation Logic

The SDK includes a validation subsystem to ensure RolloutAgentLoop implementations conform to the required schemas before execution.

Validation Criteria

1. Name Attribute: Must be a non-empty string identifying the agent.
2. Tool Schema: get_tools() must return valid OpenAIFunctionToolSchema objects or compatible dictionaries.
3. Run Method: Must be an asynchronous function (async def).
4. Metadata Size: Metadata attached to requests must not exceed the thread-safe global limit (default 1MB).

Metrics Tracking

The RolloutMetrics schema provides detailed visibility into the execution of an agent loop, tracking latency and resource consumption.

class RolloutMetrics(BaseModel):
    total_latency_ms: float = 0.0
    llm_latency_ms: float = 0.0
    tool_latency_ms: float = 0.0
    num_llm_calls: int = 0
    num_tool_calls: int = 0
    prompt_tokens: int = 0
    response_tokens: int = 0

Sources: osmosis_ai/rollout/core/schemas.py:228-245

Conclusion

State management in the Osmosis SDK is designed for reliability in distributed training environments. By combining strict Pydantic schemas for the communication protocol with thread-safe state tracking in the server, the SDK provides a robust foundation for implementing complex, multi-turn agentic workflows. The inclusion of comprehensive validation and metrics ensures that implementations are both correct and performant.