Durable Agents with Pydantic AI

This tutorial does not run on the current runtime yet: its generic action parameters lose their fields on the way to the worker (#771), and parts of its workflow loop don't compile. It will be updated once #771 is fixed. Until then, Building an Agent From Scratch shows the same pattern without the generics.

Pydantic AI handles model requests, tool schemas, and message history behind an Agent API. To resume an agent after a process stops, we need to save its conversation and dependencies between calls.

Building an Agent From Scratch introduces that loop without an agent framework: send messages, run tools, send the results back. Read it first if those steps are unfamiliar.

This draft explores how to put Waymark actions around Pydantic AI model calls and tool calls. The runtime blockers above prevent it from running today; the examples explain the intended state and execution boundaries.

Define the agent

Define the agent as a module-level object. Actions can access it in their worker process; we will persist the agent's data, not the live Agent object.

from pydantic import BaseModel
from pydantic_ai import Agent

class SupportDeps(BaseModel):
    customer_id: str
    tenant: str = "default"

support_agent = Agent(
    "openai:gpt-5.5",
    deps_type=SupportDeps,
    output_type=str,
    instructions=(
        "You are a concise support agent. Use tools before making "
        "claims about orders."
    ),
)

Pydantic AI tools are native Python functions. When you decorate one with agent.tool, Pydantic AI reads the function signature for the parameter names and types, reads the docstring for the tool description, and sends that metadata with the model request. This way the model doesn't need a handwritten schema that can drift away from the code.

Pydantic AI validates the model's requested arguments before calling the tool. For this example, the tool returns a fixed order status so we can focus on passing dependencies and history between calls.

from pydantic_ai import RunContext

@support_agent.tool
async def lookup_order(ctx: RunContext[SupportDeps], order_id: str) -> str:
    """Look up the current status for a customer order."""
    return (
        f"Customer {ctx.deps.customer_id} has order {order_id}: "
        "packed and waiting for carrier pickup."
    )

SupportDeps carries the customer and tenant across tool calls. Persisting it lets a later action recover that context without holding on to a live client or database session.

Type the state

Alongside dependencies, the state needs the conversation history and the final output, when available. Pydantic AI provides a message serializer; our wrapper keeps its JSON output together with those other values.

from typing import Generic, TypeVar

DepsT = TypeVar("DepsT", bound=BaseModel)
OutputT = TypeVar("OutputT")

class DurableAgentState(BaseModel, Generic[DepsT, OutputT]):
    deps: DepsT
    user_prompt: str | None = None
    messages_json: str | None = None
    is_done: bool = False
    final_output: OutputT | None = None

DepsT and OutputT describe each agent's dependencies and output. A support agent can store SupportDeps and return str, while another agent could use a different deps model and a structured Pydantic output.

We also need to identify which agent owns the history. An in-memory reference won't survive a restart. This draft stores the module and variable names so a worker can import the agent later.

def _agent_ref(agent: Agent[DepsT, OutputT]) -> tuple[str, str]:
    # Hard-coded module paths and variable names are invisible to linters, so
    # derive the durable reference from the actual global agent object instead.
    for name, value in globals().items():
        if value is agent:
            return __name__, name
    raise ValueError("Agent must be assigned to a module-level variable")

The helper finds the agent among this module's globals. Those names become persisted data: renaming or moving the agent can break older state, even if a type checker accepts the refactor. The state model stores the names and excludes the live object from serialization.

import importlib
from typing import Generic, Self, TypeVar, cast
 
from pydantic import Field, InstanceOf, model_validator
 
DepsT = TypeVar("DepsT", bound=BaseModel)
OutputT = TypeVar("OutputT")
 
class DurableAgentState(BaseModel, Generic[DepsT, OutputT]):
    agent: InstanceOf[Agent[DepsT, OutputT]] | None = Field(
        default=None,
        exclude=True,
        repr=False,
    )
    agent_module: str | None = None
    agent_variable: str | None = None
 
    deps: DepsT
    user_prompt: str | None = None
    messages_json: str | None = None
    is_done: bool = False
    final_output: OutputT | None = None
 
    @model_validator(mode="after")
    def resolve_agent_ref(self) -> Self:
        if self.agent is not None:
            if not isinstance(self.deps, self.agent.deps_type):
                raise TypeError("Agent deps_type does not match deps")
            self.agent_module, self.agent_variable = _agent_ref(self.agent)
            self.agent = None
 
        if self.agent_module is None or self.agent_variable is None:
            raise ValueError("Pass agent=... when creating agent state")
 
        return self
 
    def get_agent(self) -> Agent[DepsT, OutputT]:
        if self.agent_module is None or self.agent_variable is None:
            raise ValueError("Agent reference has not been resolved")
        module = importlib.import_module(self.agent_module)
        agent = getattr(module, self.agent_variable)
        return cast(Agent[DepsT, OutputT], agent)

Another worker can resolve the reference if the same module and variable are available in its environment. The cast supplies a static type hint; it does not validate the imported agent's types at runtime.

The message history needs its own translation layer. Pydantic AI messages are polymorphic: the list can contain model requests and model responses, and each message can contain different part types such as user prompts, tool calls, tool returns, and text output. A plain list[dict[str, object]] loses the Python types Pydantic AI expects on resume.

ModelMessagesTypeAdapter serializes that union to JSON and validates it back into Pydantic AI message objects when an action resumes the conversation.

from collections.abc import Sequence
from pydantic_ai.messages import ModelMessage, ModelMessagesTypeAdapter
 
class DurableAgentState(BaseModel, Generic[DepsT, OutputT]):
    agent: InstanceOf[Agent[DepsT, OutputT]] | None = Field(
        default=None,
        exclude=True,
        repr=False,
    )
    agent_module: str | None = None
    agent_variable: str | None = None
 
    deps: DepsT
    user_prompt: str | None = None
    messages_json: str | None = None
    is_done: bool = False
    final_output: OutputT | None = None
 
    @model_validator(mode="after")
    def resolve_agent_ref(self) -> Self:
        if self.agent is not None:
            if not isinstance(self.deps, self.agent.deps_type):
                raise TypeError("Agent deps_type does not match deps")
            self.agent_module, self.agent_variable = _agent_ref(self.agent)
            self.agent = None
 
        if self.agent_module is None or self.agent_variable is None:
            raise ValueError("Pass agent=... when creating agent state")
 
        return self
 
    def get_agent(self) -> Agent[DepsT, OutputT]:
        if self.agent_module is None or self.agent_variable is None:
            raise ValueError("Agent reference has not been resolved")
        module = importlib.import_module(self.agent_module)
        agent = getattr(module, self.agent_variable)
        return cast(Agent[DepsT, OutputT], agent)
 
    def get_messages(self) -> list[ModelMessage]:
        if not self.messages_json:
            return []
        return ModelMessagesTypeAdapter.validate_json(self.messages_json)
 
    def set_messages(self, messages: Sequence[ModelMessage]) -> None:
        self.messages_json = ModelMessagesTypeAdapter.dump_json(
            list(messages)
        ).decode("utf-8")

Own the loop

Before we implement the actions, sketch the workflow loop they need to serve. The workflow is going to keep handing one durable state payload from action to action until state.is_done flips to True.

roundtrip_agent will own one model turn. It either updates the state with final output or returns tool calls. handle_tool_call will run each requested tool and return results for the next roundtrip.

from waymark import Workflow


class AgentWorkflowBase(Workflow, Generic[DepsT, OutputT]):
    async def run_agent(
        self,
        state: DurableAgentState[DepsT, OutputT],
    ) -> DurableAgentState[DepsT, OutputT]:
        tool_results: list[AgentToolResult[DepsT, OutputT]] = []
        while True:
            roundtrip = await self.run_action(
                roundtrip_agent(state, tool_results=tool_results)
            )
            state = roundtrip.state
            tool_results = []

            if state.is_done:
                return state

            if not roundtrip.tool_calls:
                raise RuntimeError("Agent stopped without output or tool calls")

            tool_results = [
                await self.run_action(
                    handle_tool_call(
                        AgentToolCall(state=state, tool_call=tool_call)
                    )
                )
                for tool_call in roundtrip.tool_calls
            ]

The base workflow passes state between actions until is_done is true. The actions handle model messages and tool execution.

Initialize state

Create the initial state inside an action, where Python can run the model's validators and resolve the agent reference. Its return value supplies the customer context and first prompt to the workflow.

from waymark import action

@action
async def initialize_support_agent(
    customer_id: str,
    user_prompt: str,
) -> DurableAgentState[SupportDeps, str]:
    return DurableAgentState[SupportDeps, str](
        agent=support_agent,
        deps=SupportDeps(customer_id=customer_id),
        user_prompt=user_prompt,
    )

Defer tool calls

Calling support_agent.run(...) inside one action would also run its tools there. To save progress between tool calls, we need Pydantic AI to return those requests to the workflow so Waymark can execute them as separate actions.

Pydantic AI already has that escape hatch: a tool can raise CallDeferred. When the run's output_type includes DeferredToolRequests, Pydantic AI returns the requested calls instead of treating the exception as a failure.

Tools should defer during a model turn and execute when a Waymark action calls them. A context variable lets the wrapper distinguish those two cases.

from collections.abc import Iterator
from contextlib import contextmanager
from contextvars import ContextVar

_DEFER_TOOLS = ContextVar("_DEFER_TOOLS", default=False)

@contextmanager
def defer_tool_calls() -> Iterator[None]:
    token = _DEFER_TOOLS.set(True)
    try:
        yield
    finally:
        _DEFER_TOOLS.reset(token)

The tool wrapper checks that variable. Inside with defer_tool_calls(), it raises CallDeferred; outside the block, it calls the original function.

from collections.abc import Callable
from functools import wraps
from inspect import isawaitable
from typing import Any

from pydantic_ai import CallDeferred

ToolFuncT = TypeVar("ToolFuncT", bound=Callable[..., object])

def deferable_tool(
    agent: Agent[DepsT, OutputT],
) -> Callable[[ToolFuncT], ToolFuncT]:
    def decorator(fn: ToolFuncT) -> ToolFuncT:
        @wraps(fn)
        async def wrapper(*args: Any, **kwargs: Any) -> object:
            if _DEFER_TOOLS.get():
                raise CallDeferred()

            result = fn(*args, **kwargs)
            if isawaitable(result):
                return await result
            return result

        agent.tool(wrapper)
        return fn

    return decorator

wraps preserves the function's signature and docstring for Pydantic AI's schema generation. Register the order tool through this wrapper.

from pydantic_ai import RunContext
 
-@support_agent.tool
@deferable_tool(support_agent)
async def lookup_order(ctx: RunContext[SupportDeps], order_id: str) -> str:
    """Look up the current status for a customer order."""
    return (
        f"Customer {ctx.deps.customer_id} has order {order_id}: "
        "packed and waiting for carrier pickup."
    )

The model action runs the agent with tool deferral enabled. It saves the new history and returns either requested tools or the final output.

from pydantic import Field
from pydantic_ai import DeferredToolRequests, DeferredToolResults
from pydantic_ai.messages import ToolCallPart

class AgentRoundtrip(BaseModel, Generic[DepsT, OutputT]):
    state: DurableAgentState[DepsT, OutputT]
    tool_calls: list[ToolCallPart] = Field(default_factory=list)

class AgentToolResult(BaseModel, Generic[DepsT, OutputT]):
    tool_call: ToolCallPart
    content: object

@action
async def roundtrip_agent(
    state: DurableAgentState[DepsT, OutputT],
    tool_results: list[AgentToolResult[DepsT, OutputT]] | None = None,
) -> AgentRoundtrip[DepsT, OutputT]:
    next_state = state.model_copy(deep=True)
    history = next_state.get_messages()
    agent = next_state.get_agent()
    with defer_tool_calls():
        result = await agent.run(
            user_prompt=next_state.user_prompt if not history else None,
            deps=next_state.deps,
            message_history=history or None,
            deferred_tool_results=(
                DeferredToolResults(
                    calls={
                        result.tool_call.tool_call_id: result.content
                        for result in tool_results
                    }
                )
                if tool_results
                else None
            ),
            output_type=[agent.output_type, DeferredToolRequests],
        )
    next_state.set_messages(result.all_messages())

    if isinstance(result.output, DeferredToolRequests):
        return AgentRoundtrip(
            state=next_state,
            tool_calls=[*result.output.calls, *result.output.approvals],
        )

    next_state.is_done = True
    next_state.final_output = result.output
    return AgentRoundtrip(state=next_state)

After the first model call, every later roundtrip is driven by the stored Pydantic AI message history. When the previous step produced tool results, Pydantic AI's deferred_tool_results parameter turns them into the right tool-return messages before asking the model to continue.

Run tools durably

Each requested tool runs in its own action. We ask the agent's public toolsets for the registered tools, let Pydantic AI validate the arguments, and call the same function it would have called inside a normal agent run.

from pydantic_ai.models.test import TestModel
from pydantic_ai.usage import RunUsage

class AgentToolCall(BaseModel, Generic[DepsT, OutputT]):
    state: DurableAgentState[DepsT, OutputT]
    tool_call: ToolCallPart

@action
async def handle_tool_call(
    payload: AgentToolCall[DepsT, OutputT],
) -> AgentToolResult[DepsT, OutputT]:
    next_state = payload.state.model_copy(deep=True)
    agent = next_state.get_agent()

    ctx = RunContext(
        deps=next_state.deps,
        model=TestModel(),
        usage=RunUsage(),
        agent=agent,
        prompt=next_state.user_prompt,
        messages=next_state.get_messages(),
        tool_call_id=payload.tool_call.tool_call_id,
        tool_name=payload.tool_call.tool_name,
    )

    for toolset in agent.toolsets:
        tools = await toolset.get_tools(ctx)
        if payload.tool_call.tool_name not in tools:
            continue

        tool = tools[payload.tool_call.tool_name]
        content = await toolset.call_tool(
            payload.tool_call.tool_name,
            payload.tool_call.args_as_dict(),
            ctx,
            tool,
        )
        return AgentToolResult(
            tool_call=payload.tool_call,
            content=content,
        )

    raise ValueError(f"Unknown tool: {payload.tool_call.tool_name}")

For this tutorial, tools return plain JSONable values. Pydantic AI also supports richer return shapes for retry prompts and other tool results; those can go into the same DeferredToolResults object later.

TestModel() is only there because we are constructing RunContext ourselves. The real provider call already happened in roundtrip_agent. The next roundtrip will pair this content with the original tool_call_id through Pydantic AI's DeferredToolResults, so we do not need to hand-build ToolReturnPart.

Wire the workflow

The support workflow supplies the concrete dependency and output types. It initializes the state, runs the loop, and checks that a final answer exists.

from waymark import action, workflow

@action
async def require_support_output(
    state: DurableAgentState[SupportDeps, str],
) -> str:
    if state.final_output is None:
        raise ValueError("Agent finished without output")
    return state.final_output

@workflow
class DurableSupportAgentWorkflow(AgentWorkflowBase[SupportDeps, str]):
    async def run(self, customer_id: str, user_prompt: str) -> str:
        state = await self.run_action(
            initialize_support_agent(
                customer_id=customer_id,
                user_prompt=user_prompt,
            )
        )
        state = await self.run_agent(state)
        return await self.run_action(require_support_output(state))

Production

This draft needs the runtime fixes listed at the top before it can be used in production. The intended extension point is the agent's dependencies, tools, and output type; message serialization and deferred tool handling can be shared.

When adding real tools, keep live clients inside actions and make side effects idempotent so a repeated action does not send a duplicate email or issue a second refund.