airflow.providers.common.ai.durable.capability

Durable execution for pydantic-ai agents in Airflow tasks, on pydantic-ai’s durable backend API.

Classes

AirflowDurability

Replay an agent's completed steps when Airflow retries the task it runs in.

Module Contents

class airflow.providers.common.ai.durable.capability.AirflowDurability(*, models=None, event_stream_handler=None, name=None)[source]

Bases: pydantic_ai.durable_exec.BaseDurabilityCapability[pydantic_ai.tools.AgentDepsT]

Replay an agent’s completed steps when Airflow retries the task it runs in.

Attach it to a pydantic-ai Agent running inside an Airflow task. Each model request, tool call (function, MCP and dynamic toolsets, and Airflow’s own toolsets), tool discovery and @durable_operation of another capability is recorded in the task instance’s durable journal as it completes. When the task fails and Airflow retries it, the agent runs again from the start and every step the previous attempt completed is replayed from the journal instead of being run again: no second model call, no second tool side effect, no second charge against a spend limit. See Durable execution for how replay is verified.

AgentOperator(durable=True) and @task.agent(durable=True) attach it for you. In a plain @task, attach it yourself:

@task(retries=3)
def research(question: str) -> str:
    agent = Agent(
        "anthropic:claude-sonnet-4-5",
        name="researcher",
        toolsets=[SQLToolset("warehouse")],
        capabilities=[AirflowDurability()],
    )
    return agent.run_sync(question).output

Outside a running Airflow task the capability does nothing, so the same agent runs normally in a test or a notebook. In a plain @task, a run’s steps are deleted as soon as the run succeeds, so that clearing the task later starts it fresh: a task that runs several agents replays only the run that failed, and runs the ones before it again.

A tool that must run on every attempt, such as one that reads the current time or creates a temporary table, opts out with metadata={"airflow": False}; for tools of a toolset you do not define, such as SQLToolset, add the capability SetToolMetadata(tools=[...], airflow=False). pydantic-ai’s own MCPToolset refuses the opt-out; the provider’s MCPToolset accepts it.

Toolsets the agent calls through pydantic-ai’s durable backend need a unique id (FunctionToolset(id=...), MCPToolset(..., id=...)), as do capabilities that contribute @durable_operation methods; the ids name the steps in the journal. AgentOperator and @task.agent fill in the ids of a FunctionToolset or a function that builds a toolset passed in toolsets or agent_params["toolsets"], and of a FunctionToolset held directly by a Toolset capability (Toolset ids); in a plain @task you set them.

Parameters:

  • models (collections.abc.Mapping[str, pydantic_ai.models.Model] | None) – Extra models the run may switch to, keyed by id; see pydantic-ai’s BaseDurabilityCapability.

  • event_stream_handler (pydantic_ai.agent.EventStreamHandler[pydantic_ai.tools.AgentDepsT] | None) – Optional handler for the run’s events.

  • name (str | None) – Prefix for the names of the agent’s steps. Defaults to the agent’s name; one of the two is required.

engine_spec: ClassVar[pydantic_ai.durable_exec.DurabilityEngineSpec][source]

Declarative configuration for this durable execution engine.

property in_durable_context: bool[source]

Whether execution is currently inside this engine’s durable container (workflow or flow).

get_durable_operation_backend()[source]

Return the backend that dispatches this capability’s durable operations.

async before_model_request(ctx, request_context)[source]

Validate the final request and repair durability after an outer hook swaps the model.

get_wrapper_toolset(toolset)[source]

Journal the leaf toolsets pydantic-ai’s backend does not, then let the base wrap the rest.

async before_run(ctx)[source]
async wrap_run(ctx, *, handler)[source]

Force sequential tool execution when required by a sequence-keyed durable engine.

Was this entry helpful?