[kaxil]

A task that runs a Strands Agents agent starts the agent over from its prompt when Airflow retries it: every model call is paid for again and every tool call happens again. The Strands page of the provider docs says so ("There is no step-level replay"). This adds invoke_durably in airflow.providers.common.ai.durable.strands, which runs the agent so that a retry resumes it from the last cycle boundary it reached, and TaskStateStoreStorage, a Strands Storage backed by the task instance's task state store. #74385 is the Pydantic AI side of the same feature.

from functools import partial

from strands import Agent

from airflow.providers.common.ai.durable.strands import invoke_durably


@task(retries=2)
def research(question: str) -> str:
    agent = partial(Agent, model=model, plugins=[AirflowTools(warehouse)])
    return str(invoke_durably(agent, question))

Design rationale

It uses Strands' own durable-execution contract rather than a per-call journal. Strands pauses an agent built with checkpointing=True at each boundary of a cycle that calls tools (once the model has asked for tools, once they have run), returns stop_reason="checkpoint", and resumes from a {"checkpointResume": ...} prompt, leaving the conversation to a session manager. Intercepting each model and tool call instead would mean Strands' private middleware (strands._middleware) or subclassing Agent and Model. The cost is granularity: a retry resumes from a cycle boundary, so a task that fails while a cycle's tools run runs that whole tool batch again. The docs say so, and a test pins it.

The pause point is saved inside the snapshot, in one write. At each pause, invoke_durably puts the checkpoint in agent.state and calls SnapshotSessionManager.save_snapshot, with save_latest_on="trigger" so Strands saves nothing on its own beyond redacting or trimming the conversation it holds. Two separate writes could disagree after a crash between them. More importantly, Strands saves the snapshot when an invocation fails, too: a max_tokens cut-off answer, for example, would otherwise become part of the conversation every retry resumes. A test covers that case.

invoke_durably takes a function that builds the agent. Strands restores the snapshot when the Agent is constructed, so the session manager has to go in at construction, and a retry with a different prompt needs a fresh agent after the old session is deleted. functools.partial(Agent, ...) fits the AgentFactory protocol.

TaskStateStoreStorage keeps an index of its keys under one extra key. Strands' Storage protocol needs list(prefix) (delete_session uses it), and the task state store accessor only has get, set and delete. Keys go under the reserved __commonai_durable__ prefix that #74385 also uses, with NEVER_EXPIRE.

AirflowTools no longer resets a tool's retry budget on a resume. Every resume is a new Strands invocation, and the plugin started a fresh budget per invocation, so with checkpointing a tool that fails on every turn would never reach its retry limit. An invocation that brings no new messages (a checkpoint or interrupt resume, or a call with no prompt) now carries on the budget of the one before it.

The default session id is the provider's existing compute_identity_key over the Dag id, task id, run id and map index, so it is the same on every try and short enough for the 512-character key limit.

Gotchas

• Strands marks checkpointing experimental (strands.experimental.checkpoint). It needs strands-agents 1.51 or later, which the docs' existing >=1.56.0 install line covers. Strands is not in uv.lock, so the new tests run in the agent-framework CI job, which this adds them to.
• invoke_durably returns at the agent's first stop that is not a pause (its answer, an interrupt, a cancellation) and deletes the saved state then, not when the task succeeds. A failure later in the task, such as in a second agent, runs the finished agent again on the retry. Add AirflowDurability capability for durable agent replay across task retries #74385 behaves the same way in a plain @task.
• A SandboxToolset without attach_to is destroyed when the try ends, so a resumed agent's conversation describes files the retry's new sandbox does not have. AirflowTools keeps only the converted tools, so invoke_durably cannot detect this the way AgentOperator(durable=True) refuses a sandbox; the docs say it instead.
• A model that keeps the conversation server-side (model.stateful) raises ValueError, since Strands drops a restored conversation for it.
• A retry, or a clear of the whole Dag run, resumes with the system prompt the snapshot was saved with. Clearing the task itself without keeping its task state starts over.
• The snapshot holds the whole conversation, including tool results, in the metadata database unless [workers] state_store_backend is set, and grows with it.
• On Airflow before 3.3 there is no task state store; pass a Strands storage such as S3Storage as storage=.

Follow-ups

• A prefix listing on the task state store accessor would let TaskStateStoreStorage drop its key index.

End-to-end run

A test Dag ran a Strands agent with a scripted model in a @task(retries=1). Cycle 1 called a tool that counts its calls in a file; cycle 2's model call failed on try 1 only. Try 2 logged the resume and finished, the tool ran once in total, and the task state store held no rows for the task afterwards.