AgentWorkflow

AgentWorkflow is the SDK contract for rollout behavior. You subclass it, implement one async run() method, and create one sample either by returning its message history or by calling the current policy through an Osmosis-supported agent integration that registers the sample source. The workflow should answer one question: given this dataset prompt, what should the agent do before the grader scores the result?

Base Class

from osmosis_ai.rollout import AgentWorkflow, AgentWorkflowContext


class MyWorkflow(AgentWorkflow):
    async def run(self, ctx: AgentWorkflowContext) -> None:
        # Build and run your agent here.
        pass

The SDK shape is:

class AgentWorkflow[TConfig: AgentWorkflowConfig](ABC):
    def __init__(self, config: TConfig | None = None):
        self.config = config

    @abstractmethod
    async def run(
        self, ctx: AgentWorkflowContext[TConfig]
    ) -> AgentWorkflowOutput | list[dict[str, Any]] | None:
        raise NotImplementedError

run() is called once for each workflow execution. Construct any per-execution agent or session objects inside the method and run the agent. Then either return the sample directly or let the integration register the resulting conversation with the active RolloutContext.

Workflow Return Value

run() accepts three shapes, all normalized to a single sample:

Return value	Effect
`AgentWorkflowOutput`	The `messages` field becomes the sample; optional `metrics` are attached
`list[dict[str, Any]]`	Wrapped as `AgentWorkflowOutput(messages=...)`
`None`	The SDK collects the sample from the source registered on the active `RolloutContext`

Supported integrations register a sample source for you, so most workflows return None: OsmosisStrandsAgent registers its Strands history, while OsmosisMemorySession registers an OpenAI Agents SDK session. Return AgentWorkflowOutput when you build the message history yourself:

from osmosis_ai.rollout import (
    AgentWorkflow,
    AgentWorkflowContext,
    AgentWorkflowOutput,
)


class ExplicitWorkflow(AgentWorkflow):
    async def run(self, ctx: AgentWorkflowContext) -> AgentWorkflowOutput:
        messages = [*ctx.prompt, {"role": "assistant", "content": "done"}]
        return AgentWorkflowOutput(
            messages=messages,
            metrics={"steps": 1.0},
        )

AgentWorkflowOutput.info is reserved and is not currently passed to graders. Use metrics for finite numeric sample measurements and ctx.artifacts_dir for files.

AgentWorkflowOutput rejects unknown top-level fields and non-finite metric values (NaN, inf, -inf). The same validation runs across Local and Harbor/container execution.

AgentWorkflowContext

The ctx object gives the workflow its input and config:

Field	Type	Description
`ctx.prompt`	`list[dict[str, Any]]`	Input messages for the current dataset row
`ctx.config`	`TConfig \| None`	Custom workflow config object, if one was provided
`ctx.metadata`	`dict[str, Any] \| None`	Per-row metadata from the dataset’s optional `metadata` column. `None` when the row has no metadata.
`ctx.artifacts_dir`	`pathlib.Path \| None`	Per-rollout directory where you can write log, trace, and other output files. `None` when the execution environment can’t provision a writable directory, so check for `None` before writing files.

If your dataset row contains system_prompt, user_prompt, and ground_truth, the prompt fields are assembled into ctx.prompt. The reference answer is not passed to the workflow; it is exposed to your grader as ctx.label. The same metadata object is available on both AgentWorkflowContext and GraderContext, so workflows and graders can read the same per-row context.

Keep task answers out of AgentWorkflow.run(). The workflow should produce behavior; the Grader should decide whether that behavior deserves reward.

Writing Artifacts

Use ctx.artifacts_dir to write files that shouldn’t be embedded in the sample payload, such as logs, traces, screenshots, or other large or binary outputs. Each rollout gets its own directory, but it’s None when the environment can’t provision one, so guard with if ctx.artifacts_dir: before writing to it. An unguarded write raises and fails the workflow.

import asyncio

from osmosis_ai.rollout import AgentWorkflow, AgentWorkflowContext


class TracingWorkflow(AgentWorkflow):
    async def run(self, ctx: AgentWorkflowContext) -> None:
        if ctx.artifacts_dir:
            log_file = ctx.artifacts_dir / "run.log"
            await asyncio.to_thread(log_file.write_text, "started\n")
        ...

After the rollout finishes, collected files appear alongside its sample in the run’s Artifacts panel on the Osmosis Platform, mirroring the layout you write under ctx.artifacts_dir. Artifact collection never affects rewards or rollout status. Write artifacts off the event loop with asyncio.to_thread, as shown above. In managed runs, ctx.artifacts_dir can be backed by remote storage, so closing a file may wait for its upload. A synchronous write then blocks every rollout on the server and delays result polling.

Saved Trajectories

Rollout servers created with create_rollout_server() save finished samples as ATIF (Agent Trajectory Interchange Format) documents alongside run artifacts. Persistence is best-effort and never changes reward or rollout status. Calling LocalBackend.run_workflow() or another backend method directly in a custom harness does not install this server lifecycle; that harness must save any trajectory it needs. Files land next to the artifacts directory on the platform-managed host:

~/.osmosis/<rollout_id>/
├── trajectory.json          # ATIF document for the rollout's sample
└── artifacts/...            # files you wrote under ctx.artifacts_dir

RolloutSample.messages preserves the framework-native history that graders read. For ATIF persistence, built-in integrations make a separate, best-effort normalized copy in trajectory_messages; normalization failure leaves the native sample intact and skips trajectory persistence. Explicit workflow output uses its returned messages for both views when they can be copied.

If you build a custom sample source whose native history isn’t already OpenAI chat-completions-shaped, set RolloutSample.trajectory_messages on the returned sample to control what gets persisted (an explicit None skips trajectory saving for that sample).

ATIF includes usage, model, and timestamp data only when the source messages or server report provide those fields. Missing metadata stays absent; the SDK does not fabricate it.

Live Trajectory Previews

Live trajectory previews are disabled unless the hosting environment enables preview publication. When enabled, the rollout server samples supported conversations about every three seconds and publishes at most the 100 most recent ATIF steps; size limits can shorten the preview further. Previews are diagnostic only and do not replace the grading sample, affect rewards, or change the saved trajectory.json. Before publication, the SDK redacts the configured LLM API key and recognized credential patterns, including authorization headers and URL credentials. When the host enables previews, these sources provide snapshots automatically:

Source	Preview contents
`OsmosisStrandsAgent`	Recent Strands agent history
`OsmosisMemorySession` (OpenAI Agents SDK)	Recent session items, including tool and reasoning turns
Harbor Terminus-2, mini-swe-agent, OpenCode	The agent’s latest native trajectory or completed turns

Workflows that only return messages from run() at completion have no live source, so they produce no preview. Bundles built with an older SDK version also produce no preview and keep their final-result behavior. Mini-swe-agent versions that write their trajectory only on exit can’t provide running previews.

Opt Out for a Rollout

To keep a rollout out of live previews, set preview_enabled = False on the active rollout context. Do this at the start of run(), before you construct the agent or session and before the first await:

from osmosis_ai.rollout import AgentWorkflow, AgentWorkflowContext, get_rollout_context
from osmosis_ai.rollout.integrations.agents.strands import (
    OsmosisRolloutModel,
    OsmosisStrandsAgent,
)


class PrivateStrandsWorkflow(AgentWorkflow):
    async def run(self, ctx: AgentWorkflowContext) -> None:
        rollout_context = get_rollout_context()
        if rollout_context is not None:
            rollout_context.preview_enabled = False
        agent = OsmosisStrandsAgent(
            model=OsmosisRolloutModel(),
            messages=ctx.prompt,
            callback_handler=None,
        )
        await agent.invoke_async()

Setting the flag later stops further snapshots but does not retract a preview the SDK already captured.

Add Previews to a Custom Sample Source

Custom SampleSource implementations publish no preview by default. To opt in, override async get_preview(max_messages) and return a RolloutSample holding a copy of at most max_messages recent messages:

from copy import deepcopy

from osmosis_ai.rollout import RolloutSample, SampleSource


class MySampleSource(SampleSource):
    def __init__(self, messages: list[dict]):
        self._messages = messages

    async def get_sample(self) -> RolloutSample:
        return RolloutSample(messages=list(self._messages))

    async def get_preview(self, max_messages: int) -> RolloutSample | None:
        if max_messages <= 0:
            return None
        return RolloutSample(messages=deepcopy(self._messages[-max_messages:]))

Register an instance from inside run() with set_sample_source() on the active rollout context; each rollout accepts only one source. get_preview() must return a bounded, read-only snapshot. Don’t call get_sample() from it or mutate the history that graders read. Previews from custom sources show unknown step IDs and turn counts.

Model Routing Requirement

LLM calls inside run() must route through the RolloutContext installed by the execution backend. The training cluster uses the rollout-scoped chat-completions URL from this context to serve the current policy, collect traces, and connect the reward to the sample.

Use one of the supported integrations:

Framework	Use	Integration objects
Strands Agents	Strands tools, Strands message history, migration from `strands.Agent`	`OsmosisStrandsAgent`, `OsmosisRolloutModel`
OpenAI Agents SDK	`Runner.run`, sessions, handoffs, OpenAI-style tools	`OsmosisAgent`, `OsmosisRolloutModel`, `OsmosisMemorySession`

Do not call litellm, the OpenAI SDK, or another provider SDK directly with a hard-coded policy model from run(). Direct calls bypass the rollout context and are not compatible with training.

Strands Pattern

For Strands, pass ctx.prompt directly as messages and call invoke_async():

from osmosis_ai.rollout import AgentWorkflow, AgentWorkflowContext
from osmosis_ai.rollout.integrations.agents.strands import (
    OsmosisRolloutModel,
    OsmosisStrandsAgent,
)


class SimpleStrandsWorkflow(AgentWorkflow):
    async def run(self, ctx: AgentWorkflowContext) -> None:
        agent = OsmosisStrandsAgent(
            name="simple-strands-agent",
            model=OsmosisRolloutModel(params={"temperature": 1.0}),
            messages=ctx.prompt,
            callback_handler=None,
        )
        await agent.invoke_async()

Constructing OsmosisStrandsAgent inside run() binds it to the active rollout context and registers the agent as a sample source. See Strands Integration for tool examples, migration steps, and details about OsmosisRolloutModel.

OpenAI Agents Pattern

For OpenAI Agents, construct an OsmosisAgent, attach OpenAI Agents ModelSettings, create one OsmosisMemorySession, and pass that session to Runner.run():

from agents import ModelSettings, Runner
from osmosis_ai.rollout import AgentWorkflow, AgentWorkflowContext
from osmosis_ai.rollout.integrations.agents.openai_agents import (
    OsmosisAgent,
    OsmosisMemorySession,
    OsmosisRolloutModel,
)


class SimpleOpenAIWorkflow(AgentWorkflow):
    async def run(self, ctx: AgentWorkflowContext) -> None:
        agent = OsmosisAgent(
            name="simple-openai-agent",
            instructions="Answer the user's request clearly.",
            model=OsmosisRolloutModel(),
            model_settings=ModelSettings(temperature=1.0, max_tokens=4096),
        )
        session = OsmosisMemorySession()
        await Runner.run(
            agent,
            ctx.prompt,
            session=session,
        )

The session is what records the OpenAI Agents SDK conversation for grading. Create it inside run() so it registers with the current RolloutContext. See OpenAI Agents Integration for session behavior, tracing notes, and migration steps.

Custom Configuration

Custom configs extend AgentWorkflowConfig. Define a module-level config instance in your rollout entrypoint and pass it to the backend:

from osmosis_ai.rollout import (
    AgentWorkflow,
    AgentWorkflowConfig,
    AgentWorkflowContext,
    ConcurrencyConfig,
)


class SearchWorkflowConfig(AgentWorkflowConfig):
    name: str = "search-workflow"
    max_iterations: int = 8
    temperature: float = 1.0
    concurrency: ConcurrencyConfig = ConcurrencyConfig(max_concurrent=4)


class SearchWorkflow(AgentWorkflow[SearchWorkflowConfig]):
    async def run(self, ctx: AgentWorkflowContext[SearchWorkflowConfig]) -> None:
        config = ctx.config or SearchWorkflowConfig()
        max_iterations = config.max_iterations
        temperature = config.temperature
        # Use these values when constructing your agent.


search_workflow_config = SearchWorkflowConfig()

Pass the config instance explicitly to the backend constructor, for example LocalBackend(workflow=SearchWorkflow, workflow_config=search_workflow_config). Eval and training TOML files do not currently set workflow config fields directly. BaseConfig allows extra fields, so simple rollout configs usually do not need additional Pydantic boilerplate.

Field	Type	Default	Description
`name`	`str`	required	Identifier for the workflow
`description`	`str \| None`	`None`	Optional description
`concurrency`	`ConcurrencyConfig`	unlimited	Maximum concurrent workflow executions

Tool-Using Workflows

Tool use belongs inside your agent framework, not in the backend. Define tools the way your framework expects, pass them into the Osmosis-wrapped agent, and let the integration record the resulting messages. For example, a Strands workflow can keep its tool list in config:

from typing import Any

from strands import tool
from osmosis_ai.rollout import (
    AgentWorkflow,
    AgentWorkflowConfig,
    AgentWorkflowContext,
)
from osmosis_ai.rollout.integrations.agents.strands import (
    OsmosisRolloutModel,
    OsmosisStrandsAgent,
)


@tool(name="search")
def search_tool(query: str) -> str:
    """Search for information."""
    return f"results for {query}"


class ToolWorkflowConfig(AgentWorkflowConfig):
    name: str = "tool-workflow"
    model: Any = OsmosisRolloutModel(params={"temperature": 1.0})
    tools: Any = [search_tool]
    max_iterations: int = 8


tool_workflow_config = ToolWorkflowConfig()


class ToolWorkflow(AgentWorkflow[ToolWorkflowConfig]):
    async def run(self, ctx: AgentWorkflowContext[ToolWorkflowConfig]) -> None:
        config = ctx.config or ToolWorkflowConfig()
        agent = OsmosisStrandsAgent(
            name="search-agent",
            model=config.model,
            tools=config.tools,
            messages=ctx.prompt,
            callback_handler=None,
        )

        for _ in range(config.max_iterations):
            result = await agent.invoke_async()
            content = result.message.get("content", [])
            if not any("toolUse" in block for block in content):
                break

Entry Point Wiring

Workflow classes and config objects require explicit wiring. Select them in the backend constructor and expose that backend through the rollout server:

from osmosis_ai.rollout import LocalBackend
from osmosis_ai.rollout.server import create_rollout_server

backend = LocalBackend(
    workflow=SearchWorkflow,
    workflow_config=search_workflow_config,
)
app = create_rollout_server(backend=backend)

Multiple concrete AgentWorkflow subclasses can exist in the same module; only the class passed as workflow runs. Submit preflight imports the entrypoint once to surface import-time errors and does not inspect the module namespace; see Files in a Rollout for when that import is skipped and what it executes locally.

Next Steps

Strands Integration

Build a Strands-based rollout with tools and OsmosisStrandsAgent.

OpenAI Agents Integration

Build an OpenAI Agents SDK rollout with OsmosisAgent and OsmosisMemorySession.

Grader

Define reward logic for the sample your workflow produces.

Evaluation

Submit an evaluation run to test your workflow and grader before a training run.

Get started

Core API

Integrations

Execution

Agent tools

Base Class

Workflow Return Value

AgentWorkflowContext

Writing Artifacts

Saved Trajectories

Live Trajectory Previews

Opt Out for a Rollout

Add Previews to a Custom Sample Source

Model Routing Requirement

Strands Pattern

OpenAI Agents Pattern

Custom Configuration

Tool-Using Workflows

Entry Point Wiring

Next Steps

Strands Integration

OpenAI Agents Integration

Grader

Evaluation

​Base Class

​Workflow Return Value

​AgentWorkflowContext

​Writing Artifacts

​Saved Trajectories

​Live Trajectory Previews

​Opt Out for a Rollout

​Add Previews to a Custom Sample Source

​Model Routing Requirement

​Strands Pattern

​OpenAI Agents Pattern

​Custom Configuration

​Tool-Using Workflows

​Entry Point Wiring

​Next Steps

Strands Integration

OpenAI Agents Integration

Grader

Evaluation

Base Class

Workflow Return Value

AgentWorkflowContext

Writing Artifacts

Saved Trajectories

Live Trajectory Previews

Opt Out for a Rollout

Add Previews to a Custom Sample Source

Model Routing Requirement

Strands Pattern

OpenAI Agents Pattern

Custom Configuration

Tool-Using Workflows

Entry Point Wiring

Next Steps