This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new eee6e2a4a50 Resume Strands agents from their last checkpoint when a
task retries (#74404)
eee6e2a4a50 is described below
commit eee6e2a4a5005dad21f5422a9c26f83c33ab2a0a
Author: Kaxil Naik <[email protected]>
AuthorDate: Wed Oct 14 01:38:31 2026 +0100
Resume Strands agents from their last checkpoint when a task retries
(#74404)
A task running a Strands Agents agent started the agent over from its
prompt on every retry. invoke_durably runs the agent with Strands'
checkpointing and a SnapshotSessionManager, saving the conversation and
the pause point together at each cycle boundary, so a retry resumes from
the last boundary instead. TaskStateStoreStorage keeps that state in the
task instance's task state store.
AirflowTools now keeps a tool's retry budget across checkpoint resumes,
which Strands runs as separate invocations.
---
docs/spelling_wordlist.txt | 1 +
providers/common/ai/docs/durable_execution.rst | 3 +-
providers/common/ai/docs/frameworks/strands.rst | 129 ++++++++-
.../airflow/providers/common/ai/durable/strands.py | 209 ++++++++++++++
.../providers/common/ai/durable/strands_storage.py | 137 +++++++++
.../ai/example_dags/example_strands_durable.py | 93 ++++++
.../airflow/providers/common/ai/tools/strands.py | 10 +-
.../tests/unit/common/ai/durable/test_strands.py | 314 +++++++++++++++++++++
.../unit/common/ai/durable/test_strands_storage.py | 156 ++++++++++
.../ai/tests/unit/common/ai/tools/test_strands.py | 54 +++-
scripts/in_container/run_agent_framework_tests.sh | 14 +-
11 files changed, 1106 insertions(+), 14 deletions(-)
diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt
index 22ad6f3c217..8d0cd1dd1f1 100644
--- a/docs/spelling_wordlist.txt
+++ b/docs/spelling_wordlist.txt
@@ -29,6 +29,7 @@ afterall
agentcore
AgentDetails
agentic
+AgentResult
AgentVersionDetails
ai
aio
diff --git a/providers/common/ai/docs/durable_execution.rst
b/providers/common/ai/docs/durable_execution.rst
index ab6647f2492..0de990ce9c8 100644
--- a/providers/common/ai/docs/durable_execution.rst
+++ b/providers/common/ai/docs/durable_execution.rst
@@ -40,7 +40,8 @@ retries there is nothing to replay.
This page is about making an ``AgentOperator`` retry cheap. Deciding *whether*
a task
should retry at all is :doc:`retry_policies`; a retried ``LLMBatchOperator``
re-attaches
-to its running batch instead of resubmitting (:ref:`llm-batch-reattach`).
+to its running batch instead of resubmitting (:ref:`llm-batch-reattach`). To
resume a
+Strands Agents agent on a retry, see :ref:`howto/frameworks:strands-durable`.
**Configuration**
diff --git a/providers/common/ai/docs/frameworks/strands.rst
b/providers/common/ai/docs/frameworks/strands.rst
index 41fd9184512..75401a9a2e1 100644
--- a/providers/common/ai/docs/frameworks/strands.rst
+++ b/providers/common/ai/docs/frameworks/strands.rst
@@ -159,17 +159,138 @@ tracer, so an ``Agent`` created earlier in the process,
such as at module level,
content capture on. See
:doc:`../observability`.
+.. _howto/frameworks:strands-durable:
+
+Durable execution with Strands
+------------------------------
+
+When a task that calls ``agent(question)`` fails and Airflow retries it, the
agent
+starts again from the prompt: every model call is paid for again, and every
tool
+call happens again. On a task with ``retries``, run the agent with
+:func:`~airflow.providers.common.ai.durable.strands.invoke_durably` instead,
and a
+retry resumes it from the last cycle boundary it reached:
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_strands_durable.py
+ :language: python
+ :start-after: [START example_strands_durable]
+ :end-before: [END example_strands_durable]
+
+The example needs Airflow 3.3 or later, ``strands-agents``, a connection
holding an
+Anthropic API key, and a database connection; its module docstring lists the
details.
+
+``invoke_durably`` takes a function that builds the agent, here
+``functools.partial(Agent, ...)``, and calls it with two more arguments:
+``checkpointing=True`` and a ``SnapshotSessionManager``. Strands
+`pauses an agent built with checkpointing
+<https://github.com/strands-agents/harness-sdk/blob/main/strands-py/src/strands/experimental/checkpoint/checkpoint.py>`__
+at each boundary of a cycle that calls tools: once the model has asked for
tools,
+and once they have run. It leaves keeping the conversation to a session
manager.
+At each pause ``invoke_durably`` has the session manager save the conversation,
+with where the agent paused in ``agent.state["airflow_checkpoint"]``, as one
+snapshot in the task instance's
+:doc:`task state store <apache-airflow:core-concepts/task-state-store>`, and
the
+agent carries on. Strands saves the snapshot on its own only to redact or trim
the
+conversation it already holds, so a failed invocation leaves the last pause's
+snapshot as it was. On a retry, the agent is rebuilt from the snapshot and
resumed
+from the pause. The task log says where. This line comes from a test Dag whose
+model failed in its second cycle; ``cycle`` counts from 0:
+
+.. code-block:: text
+
+ [2026-10-07 11:03:38] INFO - Resuming the Strands agent from its last
checkpoint
session_id=f711ace605c22a218c598a0f9e9a19a3d392fbcf6551055d8af473bdff33dd72
position=after_tools cycle=0
+
+``invoke_durably`` returns at the agent's first stop that is not a pause: its
final
+answer, or an interrupt or a cancellation, which it does not resume. It
deletes the
+saved state then, so clearing the task later starts the agent over. If that
delete
+fails, the task still succeeds and logs ``Could not delete the Strands agent's
saved
+state``, and a later clear of the whole Dag run resumes from the last pause.
+
+A retry runs some work again:
+
+- It resumes from a cycle boundary. If the task fails while a cycle's tools
run,
+ the retry runs all of that cycle's tools again, the finished ones included,
so
+ tools that change data must still be safe to repeat.
+- A task that fails before the model's first request for tools starts over.
+- A retry whose prompt differs from the first try's starts over.
+- The state is deleted when ``invoke_durably`` returns, not when the task
succeeds.
+ If the task fails after that, for example in a second agent or while
returning
+ the answer, the retry runs the finished agent again from its prompt.
+- A ``SandboxToolset`` without ``attach_to`` is destroyed when the try ends,
so a
+ resumed agent works from a conversation that describes files the retry's new,
+ empty sandbox does not have. Attach the toolset to a sandbox another task
owns,
+ or keep sandbox work out of a durable agent.
+
+A retry resumes the saved conversation with the system prompt it was saved
with.
+So does clearing the whole Dag run, which keeps task state. To run a changed
system
+prompt on a failed task, clear the task itself without keeping its task state.
+
+An agent whose model keeps the conversation server-side (``model.stateful``)
raises
+``ValueError``: Strands' `session manager
+<https://github.com/strands-agents/harness-sdk/blob/main/strands-py/src/strands/session/snapshot_session_manager.py>`__
+drops a restored conversation for such a model, so it cannot be resumed.
+
+The snapshot holds the whole conversation (the prompt, the model's replies,
every
+tool call and result), the system prompt and ``agent.state``. It goes to the
+metadata database unless ``[workers] state_store_backend`` sends task state to
+external storage, and anyone who can read the task instance's task state can
read
+it. Each pause rewrites it whole, and once it is larger than
+``[state_store] max_value_storage_bytes`` (64 KB by default) every write logs a
+warning; the write still goes through. ``AirflowTools`` masks
+registered secrets in its tools' results; results of your own ``@tool``
functions
+are saved as returned. A task that never succeeds leaves its state in the task
+state store until the Dag run is deleted. A snapshot offloaded to a custom
+``[workers] state_store_backend`` stays in that storage after the rows are
gone,
+and clearing the task does not reach it.
+
+Within one try, a tool's retry limit in ``AirflowTools`` counts across the
agent's
+pauses; a task retry starts the count again. The ``AgentResult`` that
+``invoke_durably`` returns comes from the agent's last invocation. Its metrics
+cover every invocation of the current try, one per resume in
+``metrics.agent_invocations``, and none of the earlier tries.
+
+When a task runs more than one agent with ``invoke_durably``, give each a
+``session_id`` of its own that stays the same across tries.
+
+The task state store needs Airflow 3.3 or later. On an earlier version,
+``invoke_durably`` without ``storage`` raises
+``AirflowOptionalProviderFeatureException``. Pass a Strands storage instead,
such
+as ``S3Storage``, inside the task; ``AwsBaseHook`` comes from the Amazon
provider:
+
+.. code-block:: python
+
+ from strands.storage import S3Storage
+
+ from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+
+ storage = S3Storage(
+ "my-agent-state",
+ prefix="strands/",
+ boto_session=AwsBaseHook(aws_conn_id="aws_default").get_session(),
+ )
+ result = invoke_durably(agent, question, storage=storage)
+
+The default ``session_id`` is derived from the Dag id, run id, task id and map
+index, so it differs between Dag runs. If you pass your own ``session_id`` and
the
+storage is shared across Dag runs, as a bucket is, it must differ between runs
too. ``invoke_durably`` deletes the
+state when the agent finishes or a try starts over; a task that never succeeds
+leaves it in the bucket.
+
+Strands marks checkpointing experimental
(``strands.experimental.checkpoint``). It
+needs ``strands-agents`` 1.56 or later, as in
:ref:`howto/frameworks:strands-install`.
+For Pydantic AI agents, see :doc:`../durable_execution`.
+
Differences from ``AgentOperator``
----------------------------------
The agent is Strands', so features that ``AgentOperator`` implements on top of
-Pydantic AI do not apply:
+Pydantic AI work differently or do not apply:
-- A task retry runs the agent again from the start. There is no step-level
- replay as with ``AgentOperator(durable=True)``, so tools that change data
must
- be safe to repeat.
+- Durable execution resumes from a cycle boundary rather than replaying each
+ model call and tool call. See :ref:`howto/frameworks:strands-durable`.
- There is no human review step like
``AgentOperator(enable_hitl_review=True)``.
+.. _howto/frameworks:strands-install:
+
Installation
------------
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/durable/strands.py
b/providers/common/ai/src/airflow/providers/common/ai/durable/strands.py
new file mode 100644
index 00000000000..7050c4e6ba8
--- /dev/null
+++ b/providers/common/ai/src/airflow/providers/common/ai/durable/strands.py
@@ -0,0 +1,209 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Resume a `Strands Agents <https://strandsagents.com/>`__ agent where it
stopped when Airflow retries its task.
+
+.. note:: Experimental, as Strands' checkpointing is; see
:mod:`strands.experimental.checkpoint`.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from typing import TYPE_CHECKING, Any, Protocol
+
+from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException, get_current_context
+
+try:
+ from strands.session import SnapshotSessionManager
+ from strands.types.session import encode_bytes_values
+except ImportError as e:
+ raise AirflowOptionalProviderFeatureException(e)
+
+from airflow.providers.common.ai.batch.state import compute_identity_key
+from airflow.providers.common.ai.utils.coroutines import run_coroutine_sync
+from airflow.providers.common.ai.utils.task_logger import get_task_logger
+from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_3_PLUS
+
+if AIRFLOW_V_3_3_PLUS:
+ from airflow.providers.common.ai.durable.strands_storage import
TaskStateStoreStorage
+
+if TYPE_CHECKING:
+ from strands import Agent
+ from strands.agent import AgentResult
+ from strands.session import SessionManager
+ from strands.storage import Storage
+ from strands.types.agent import AgentInput
+
+__all__ = ["AgentFactory", "invoke_durably"]
+
+log = get_task_logger()
+
+# The key in ``agent.state`` that holds where the agent paused and for which
prompt.
+_CHECKPOINT_STATE_KEY = "airflow_checkpoint"
+
+
+class AgentFactory(Protocol):
+ """
+ Builds the agent :func:`invoke_durably` runs, wiring in the two arguments
it is called with.
+
+ ``functools.partial(Agent, model=model, tools=[...])`` is one.
+ """
+
+ def __call__(self, *, session_manager: SessionManager, checkpointing:
bool) -> Agent: ...
+
+
+def invoke_durably(
+ agent_factory: AgentFactory,
+ prompt: AgentInput,
+ *,
+ session_id: str | None = None,
+ storage: Storage | None = None,
+) -> AgentResult:
+ """
+ Run a Strands agent so that a retry of the task resumes it from its last
checkpoint.
+
+ Call it inside a task. The agent is built with Strands'
``checkpointing=True`` and a
+ ``SnapshotSessionManager``, so it pauses at every boundary of a cycle that
calls tools:
+ once the model has asked for tools, and once they have run. At each pause
the
+ conversation is saved, together with where the agent paused, before the
agent carries
+ on. When the task fails and Airflow retries it, the agent is rebuilt from
the last
+ pause, so the model calls and tool calls of the cycles before it do not
run again.
+ Once the agent stops for any other reason, such as its final answer, an
interrupt or
+ a cancellation, what was saved is deleted, so clearing the task later
starts the
+ agent over. :ref:`howto/frameworks:strands-durable` describes what a retry
runs again.
+
+ Pass a function that builds the agent, such as ``functools.partial(Agent,
...)``:
+
+ .. code-block:: python
+
+ from functools import partial
+
+ from strands import Agent
+
+ from airflow.providers.common.ai.durable.strands import invoke_durably
+ from airflow.providers.common.ai.tools.strands import AirflowTools
+ from airflow.sdk import task
+
+
+ @task(retries=3)
+ def research(question: str) -> str:
+ agent = partial(Agent, model=model,
plugins=[AirflowTools(warehouse)])
+ return str(invoke_durably(agent, question))
+
+ :param agent_factory: Builds the agent, given the ``session_manager`` and
+ ``checkpointing`` arguments to pass to ``Agent``. It is called once
per try, or
+ twice when the saved state is from a different prompt.
+ :param prompt: The prompt for the agent; ignored when resuming.
+ :param session_id: Names the agent's saved state. Defaults to an id
derived from the
+ task instance (Dag id, run id, task id and map index), the same on
every try. Pass
+ a distinct one for each agent when a task runs more than one; it must
stay the same
+ across tries, and when ``storage`` is shared across Dag runs it must
differ between
+ them.
+ :param storage: Where the agent's state is saved. Defaults to the task
instance's task
+ state store, on Airflow 3.3 and later; on older versions pass a
Strands storage such
+ as ``S3Storage``.
+ :return: The result of the agent's last invocation. Its metrics cover
every invocation
+ of this try, one per resume in ``metrics.agent_invocations``, and none
of the
+ earlier tries.
+ :raises ValueError: If the agent's model keeps the conversation server-side
+ (``model.stateful``), since Strands drops a restored conversation for
it.
+ """
+ context = get_current_context()
+ if storage is None:
+ if not AIRFLOW_V_3_3_PLUS:
+ raise AirflowOptionalProviderFeatureException(
+ "invoke_durably() saves to the task state store, which needs
Airflow 3.3 or later. "
+ "On older versions pass a Strands storage, such as
strands.storage.S3Storage, as storage=."
+ )
+ storage = TaskStateStoreStorage(context["task_state_store"])
+ if session_id is None:
+ ti = context["ti"]
+ # A hash: short (task state store keys are at most 512 characters) and
free of the
+ # characters Strands refuses in a session id, whatever the Dag, run
and task ids hold.
+ session_id = compute_identity_key(
+ dag_id=ti.dag_id,
+ task_id=ti.task_id,
+ run_id=ti.run_id,
+ map_index=ti.map_index if ti.map_index is not None else -1,
+ )
+ prompt_fingerprint = _build_fingerprint(prompt)
+
+ session_manager, agent = _build_agent(agent_factory, session_id, storage)
+ saved = agent.state.get(_CHECKPOINT_STATE_KEY)
+ request: AgentInput | dict[str, Any]
+ if isinstance(saved, dict) and saved.get("prompt") == prompt_fingerprint:
+ checkpoint = saved["checkpoint"]
+ log.info(
+ "Resuming the Strands agent from its last checkpoint",
+ session_id=session_id,
+ position=checkpoint["position"],
+ cycle=checkpoint["cycle_index"],
+ )
+ request = _build_checkpoint_resume(checkpoint)
+ else:
+ if saved is not None:
+ # The agent was restored from an earlier try with a different
prompt, so its
+ # conversation cannot be resumed. Start over from an empty one.
+ log.warning("The saved Strands agent state is for a different
prompt; starting over")
+ run_coroutine_sync(session_manager.delete_session())
+ session_manager, agent = _build_agent(agent_factory, session_id,
storage)
+ request = prompt
+
+ # Strands types the prompt as AgentInput, which leaves out its
checkpointResume block.
+ result = agent(request) # type: ignore[arg-type]
+ while result.checkpoint is not None:
+ checkpoint = result.checkpoint.to_dict()
+ # In the snapshot itself, so the conversation and the pause it belongs
to are saved
+ # in one write and can never disagree.
+ agent.state.set(_CHECKPOINT_STATE_KEY, {"checkpoint": checkpoint,
"prompt": prompt_fingerprint})
+ run_coroutine_sync(session_manager.save_snapshot(agent,
is_latest=True))
+ log.debug("Saved a Strands checkpoint", session_id=session_id,
**checkpoint)
+ result = agent(_build_checkpoint_resume(checkpoint)) # type:
ignore[arg-type]
+
+ try:
+ run_coroutine_sync(session_manager.delete_session())
+ except Exception:
+ # The agent has finished; a failed delete must not fail the task. What
is left is
+ # removed with the Dag run, and a retry resumes from the last pause.
+ log.warning("Could not delete the Strands agent's saved state",
session_id=session_id, exc_info=True)
+ return result
+
+
+def _build_agent(
+ agent_factory: AgentFactory, session_id: str, storage: Storage
+) -> tuple[SnapshotSessionManager, Agent]:
+ # "trigger" with no trigger: Strands saves on its own only to redact or
trim the
+ # conversation it holds, not when an invocation ends or fails, so the
snapshot stays
+ # at the last pause.
+ session_manager = SnapshotSessionManager(session_id, storage=storage,
save_latest_on="trigger")
+ agent = agent_factory(session_manager=session_manager, checkpointing=True)
+ if agent.model.stateful:
+ raise ValueError(
+ "invoke_durably() cannot resume an agent whose model keeps the
conversation "
+ "server-side (model.stateful): Strands drops a restored
conversation for it."
+ )
+ return session_manager, agent
+
+
+def _build_fingerprint(prompt: AgentInput) -> str:
+ # Bytes, such as an image in a content block, encoded the way Strands
saves them.
+ return hashlib.sha256(json.dumps(encode_bytes_values(prompt),
sort_keys=True).encode()).hexdigest()
+
+
+def _build_checkpoint_resume(checkpoint: dict[str, Any]) -> dict[str, Any]:
+ return {"checkpointResume": {"checkpoint": checkpoint}}
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/durable/strands_storage.py
b/providers/common/ai/src/airflow/providers/common/ai/durable/strands_storage.py
new file mode 100644
index 00000000000..3d6c9e9849a
--- /dev/null
+++
b/providers/common/ai/src/airflow/providers/common/ai/durable/strands_storage.py
@@ -0,0 +1,137 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+A `Strands Agents <https://strandsagents.com/>`__ ``Storage`` backed by the
task state store.
+
+Available on Airflow >= 3.3, where the task state store keeps key/value state
for a
+task instance across its tries within a Dag run; ``NEVER_EXPIRE`` does not
exist on older
+Airflow versions.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+import bisect
+import builtins
+import threading
+from typing import TYPE_CHECKING
+
+from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
+
+try:
+ from strands.storage import Storage
+except ImportError as e:
+ raise AirflowOptionalProviderFeatureException(e)
+
+from airflow.providers.common.ai.durable.base import DURABLE_KEY_PREFIX
+from airflow.sdk.execution_time.context import NEVER_EXPIRE
+
+if TYPE_CHECKING:
+ from pydantic import JsonValue
+
+ from airflow.sdk.execution_time.context import TaskStateStoreAccessor
+
+__all__ = ["TaskStateStoreStorage"]
+
+_KEY_PREFIX = f"{DURABLE_KEY_PREFIX}strands/"
+# The task state store cannot list keys by prefix, so the keys this storage
holds are kept
+# in a sorted list under one more key.
+_INDEX_KEY = f"{DURABLE_KEY_PREFIX}strands_index"
+# Every instance over the same task instance shares that list, so updates to
it are
+# serialized across the process, not per instance.
+_index_lock = threading.Lock()
+
+
+class TaskStateStoreStorage(Storage):
+ """
+ A Strands ``Storage`` that keeps its keys in the running task instance's
task state store.
+
+ The keys share the task instance's key namespace with anything user code
writes through
+ ``context["task_state_store"]``, under a reserved prefix that keeps them
apart. They are
+ written with ``NEVER_EXPIRE``, so a retry finds them whatever
``retry_delay`` or the
+ retention config is. Keys left behind by a task that never succeeds are
removed with
+ the Dag run.
+
+ A value that is valid UTF-8, such as a Strands session snapshot, is stored
as text, and
+ anything else as base64.
+
+ :param accessor: The task state store of the running task instance,
+ ``context["task_state_store"]``.
+ """
+
+ def __init__(self, accessor: TaskStateStoreAccessor) -> None:
+ self._store = accessor
+
+ async def write(self, key: str, data: bytes) -> None:
+ await asyncio.to_thread(self._write, key, data)
+
+ async def read(self, key: str) -> bytes | None:
+ return _decode(await asyncio.to_thread(self._store.get, _KEY_PREFIX +
key))
+
+ async def delete(self, key: str) -> None:
+ await asyncio.to_thread(self._delete, key)
+
+ async def list(self, query: str = "") -> builtins.list[str]:
+ keys = await asyncio.to_thread(self._load_index)
+ return [key for key in keys if key.startswith(query)]
+
+ def _write(self, key: str, data: bytes) -> None:
+ with _index_lock:
+ keys = self._load_index()
+ if key not in keys:
+ # Index first: a listed key with no value is harmless to
delete, while a value
+ # missing from the index would outlive a delete of everything
listed.
+ bisect.insort(keys, key)
+ self._save_index(keys)
+ self._store.set(_KEY_PREFIX + key, _encode(data),
retention=NEVER_EXPIRE)
+
+ def _delete(self, key: str) -> None:
+ # Value first, so a failure in between leaves the key listed for the
next delete.
+ self._store.delete(_KEY_PREFIX + key)
+ with _index_lock:
+ keys = self._load_index()
+ if key not in keys:
+ return
+ keys.remove(key)
+ self._save_index(keys)
+
+ def _save_index(self, keys: builtins.list[str]) -> None:
+ if not keys:
+ self._store.delete(_INDEX_KEY)
+ return
+ index: builtins.list[JsonValue] = [*keys]
+ self._store.set(_INDEX_KEY, index, retention=NEVER_EXPIRE)
+
+ def _load_index(self) -> builtins.list[str]:
+ keys = self._store.get(_INDEX_KEY)
+ return [key for key in keys if isinstance(key, str)] if
isinstance(keys, builtins.list) else []
+
+
+def _encode(data: bytes) -> JsonValue:
+ try:
+ return data.decode("utf-8")
+ except UnicodeDecodeError:
+ return {"base64": base64.b64encode(data).decode("ascii")}
+
+
+def _decode(value: JsonValue) -> bytes | None:
+ if isinstance(value, str):
+ return value.encode("utf-8")
+ if isinstance(value, dict) and isinstance(encoded := value.get("base64"),
str):
+ return base64.b64decode(encoded)
+ return None
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_strands_durable.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_strands_durable.py
new file mode 100644
index 00000000000..b1d1e6515bd
--- /dev/null
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_strands_durable.py
@@ -0,0 +1,93 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Run a Strands Agents agent that a task retry resumes from its last checkpoint.
+
+The agent is the one in ``example_strands_agent``: an Anthropic model whose
API key
+comes from an Airflow connection, and a ``SQLToolset``'s tools through the
+:class:`~airflow.providers.common.ai.tools.strands.AirflowTools` plugin.
+:func:`~airflow.providers.common.ai.durable.strands.invoke_durably` runs it,
so when
+the task fails and Airflow retries it, the agent carries on from the last
cycle it
+finished instead of starting over.
+
+Before running:
+
+1. Use Airflow 3.3 or later, where the agent's state is kept in the task state
store.
+2. Install Strands: ``pip install "strands-agents>=1.56.0"``. The Anthropic
model
+ below also needs ``anthropic``, which the ``anthropic`` extra of this
provider
+ installs.
+3. Create a connection (``LLM_CONN_ID``, default ``anthropic_default``) whose
+ password is an Anthropic API key. Set its host only to route through a
gateway.
+4. Create a database connection (``DB_CONN_ID``, default ``sql_default``) whose
+ hook is a ``DbApiHook`` (e.g. SQLite, Postgres, MySQL).
+"""
+
+from __future__ import annotations
+
+import os
+from datetime import timedelta
+from functools import partial
+
+from airflow.providers.common.compat.sdk import BaseHook, dag, task
+
+LLM_CONN_ID = os.environ.get("LLM_CONN_ID", "anthropic_default")
+LLM_MODEL = os.environ.get("LLM_MODEL", "claude-sonnet-5")
+DB_CONN_ID = os.environ.get("DB_CONN_ID", "sql_default")
+
+DEFAULT_QUESTION = "Which tables exist, and how many rows does each contain?"
+
+
+# [START example_strands_durable]
+@dag(tags=["example"])
+def example_strands_durable():
+ """Answer a question about a database with a Strands agent that survives
task retries."""
+
+ @task(retries=2, retry_delay=timedelta(minutes=1))
+ def run_durable_strands_agent(question: str = DEFAULT_QUESTION) -> str:
+ from strands import Agent
+ from strands.models.anthropic import AnthropicModel
+
+ from airflow.providers.common.ai.durable.strands import invoke_durably
+ from airflow.providers.common.ai.tools.strands import AirflowTools
+ from airflow.providers.common.ai.toolsets.sql import SQLToolset
+
+ llm = BaseHook.get_connection(LLM_CONN_ID)
+ model = AnthropicModel(
+ client_args={"api_key": llm.password, "base_url": llm.host or
None},
+ model_id=LLM_MODEL,
+ max_tokens=2048,
+ )
+ # invoke_durably builds the agent, adding the session manager and
checkpointing.
+ agent = partial(
+ Agent,
+ model=model,
+ plugins=[AirflowTools(SQLToolset(db_conn_id=DB_CONN_ID))],
+ system_prompt=(
+ "You are a SQL analyst. Use list_tables and get_schema to
explore "
+ "the database, then run read-only queries to answer the
question."
+ ),
+ callback_handler=None,
+ )
+ return str(invoke_durably(agent, question))
+
+ run_durable_strands_agent()
+
+
+# [END example_strands_durable]
+
+
+example_strands_durable()
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/tools/strands.py
b/providers/common/ai/src/airflow/providers/common/ai/tools/strands.py
index 3df6777269e..d9e80cf91c6 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/tools/strands.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/tools/strands.py
@@ -87,7 +87,8 @@ class AirflowTools(Plugin):
result passes through Airflow's secret masker first. A failure the model
can
correct, such as a query naming a missing column, reaches it as a Strands
result
with ``status="error"``, counted against the tool's retry limit once per
model turn
- and afresh on every run of the agent. A tool that is still failing once
the limit is
+ and afresh each time the agent is given a new prompt; resuming from a
checkpoint
+ keeps the count. A tool that is still failing once the limit is
used up, or a failure the model cannot fix, such as a hook raising, fails
the agent
run, so the task fails and Airflow retries it. Strands on its own would
hand that
failure to the model and carry on. Tools a toolset marks sequential, such
as the
@@ -116,8 +117,11 @@ class AirflowTools(Plugin):
@hook
def _start_a_run(self, event: BeforeInvocationEvent) -> None:
- # Each call on the agent is a run of its own, with a fresh retry
budget per tool.
- self._run = object()
+ # Each call on the agent with a new prompt is a run of its own, with a
fresh retry
+ # budget per tool. A call that brings no new messages carries on the
run before it:
+ # resuming from a checkpoint or an interrupt, or calling the agent
with no prompt.
+ if event.messages != []:
+ self._run = object()
@hook
def _fail_the_run_on_tool_call_error(self, event: AfterToolCallEvent) ->
None:
diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_strands.py
b/providers/common/ai/tests/unit/common/ai/durable/test_strands.py
new file mode 100644
index 00000000000..a5125bda811
--- /dev/null
+++ b/providers/common/ai/tests/unit/common/ai/durable/test_strands.py
@@ -0,0 +1,314 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import asyncio
+import json
+import uuid
+from collections import Counter
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+from unittest import mock
+
+import pytest
+
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS
+
+if not AIRFLOW_V_3_3_PLUS:
+ pytest.skip("task state store needs Airflow >= 3.3",
allow_module_level=True)
+
+pytest.importorskip("strands")
+
+from strands import Agent, tool
+from strands.storage import InMemoryStorage
+from strands.types.exceptions import MaxTokensReachedException
+
+from airflow.providers.common.ai.durable import strands as durable_strands
+from airflow.providers.common.ai.durable.strands import invoke_durably
+from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException
+
+from unit.common.ai.durable.test_task_state_store import FakeTaskStateStore
+from unit.common.ai.tools.test_strands import _OneToolCallModel
+
+if TYPE_CHECKING:
+ from strands.types.agent import AgentInput
+
+
+class _WorkerKilled(SystemExit):
+ """Stands in for the worker dying: Strands hands an exception from a tool
to the model."""
+
+
+class _ScriptedModel(_OneToolCallModel):
+ """
+ Turn 0 calls ``first``, turn 1 calls ``second`` and ``third`` together,
turn 2 answers.
+
+ The turn is the number of tool results in the conversation, so an agent
rebuilt from a
+ snapshot carries on from where the last one stopped.
+ """
+
+ def __init__(
+ self, turns: list[int], fail_at_turn: int | None, truncate_at_turn:
int | None = None
+ ) -> None:
+ super().__init__("first", {})
+ self._turns = turns
+ self._fail_at_turn = fail_at_turn
+ self._truncate_at_turn = truncate_at_turn
+
+ async def stream(self, messages, tool_specs=None, system_prompt=None,
**kwargs: Any):
+ turn = sum(1 for message in messages if any("toolResult" in block for
block in message["content"]))
+ self._turns.append(turn)
+ if turn == self._fail_at_turn:
+ raise RuntimeError(f"model unavailable at turn {turn}")
+ yield {"messageStart": {"role": "assistant"}}
+ if turn == self._truncate_at_turn:
+ # Strands keeps the cut-off answer in the conversation, then
raises.
+ yield {"contentBlockDelta": {"delta": {"text": "cut off"}}}
+ yield {"contentBlockStop": {}}
+ yield {"messageStop": {"stopReason": "max_tokens"}}
+ return
+ calls = {0: ["first"], 1: ["second", "third"]}.get(turn, [])
+ for name in calls:
+ start = {"toolUse": {"toolUseId": f"{name}-call", "name": name}}
+ yield {"contentBlockStart": {"start": start}}
+ yield {"contentBlockDelta": {"delta": {"toolUse": {"input":
json.dumps({})}}}}
+ yield {"contentBlockStop": {}}
+ if calls:
+ yield {"messageStop": {"stopReason": "tool_use"}}
+ return
+ yield {"contentBlockDelta": {"delta": {"text": "done"}}}
+ yield {"contentBlockStop": {}}
+ yield {"messageStop": {"stopReason": "end_turn"}}
+
+
+@dataclass
+class _TaskInstance:
+ dag_id: str = "research"
+ run_id: str = "manual__2026-10-07T00:00:00+00:00"
+ task_id: str = "answer"
+ map_index: int = -1
+ # Every try of a task instance gets a new id; the agent's state must not
be keyed on it.
+ id: uuid.UUID = field(default_factory=uuid.uuid4)
+
+
+class _Task:
+ """A task running a durable agent. Each ``run`` is one try; only stored
state carries over."""
+
+ def __init__(self) -> None:
+ self.store = FakeTaskStateStore()
+ self.tool_calls: Counter[str] = Counter()
+ self.model_turns: list[int] = []
+ self.agents: list[Agent] = []
+
+ def run(
+ self,
+ prompt: AgentInput = "go",
+ *,
+ fail_at_turn: int | None = None,
+ truncate_at_turn: int | None = None,
+ kill_in_tool: str | None = None,
+ map_index: int = -1,
+ **kwargs: Any,
+ ) -> str:
+ ti = _TaskInstance(map_index=map_index)
+ model = _ScriptedModel(self.model_turns, fail_at_turn,
truncate_at_turn)
+
+ def call(name: str) -> str:
+ self.tool_calls[name] += 1
+ if name == kill_in_tool:
+ raise _WorkerKilled
+ return f"{name} done"
+
+ @tool
+ def first() -> str:
+ """First tool."""
+ return call("first")
+
+ @tool
+ def second() -> str:
+ """Second tool."""
+ return call("second")
+
+ @tool
+ def third() -> str:
+ """Third tool."""
+ return call("third")
+
+ def factory(**durable_kwargs: Any) -> Agent:
+ agent = Agent(model=model, tools=[first, second, third],
callback_handler=None, **durable_kwargs)
+ self.agents.append(agent)
+ return agent
+
+ context = {"ti": ti, "task_state_store": self.store}
+ with mock.patch.object(durable_strands, "get_current_context",
autospec=True, return_value=context):
+ return str(invoke_durably(factory, prompt, **kwargs)).strip()
+
+
+def _get_texts(agent: Agent) -> list[str]:
+ """The text blocks of the agent's conversation: the prompts it was given
and its answers."""
+ return [block["text"] for message in agent.messages for block in
message["content"] if "text" in block]
+
+
[email protected]
+def task():
+ return _Task()
+
+
+class TestInvokeDurably:
+ def test_runs_the_agent_to_its_answer(self, task):
+ assert task.run() == "done"
+ assert task.tool_calls == {"first": 1, "second": 1, "third": 1}
+
+ def test_a_retry_resumes_after_the_last_finished_cycle(self, task):
+ with pytest.raises(RuntimeError, match="model unavailable at turn 1"):
+ task.run(fail_at_turn=1)
+
+ assert task.run() == "done"
+
+ # Cycle 1's model call and tool ran once; the retry started from cycle
2's model call.
+ assert task.tool_calls == {"first": 1, "second": 1, "third": 1}
+ assert task.model_turns == [0, 1, 1, 2]
+ assert _get_texts(task.agents[-1]) == ["go", "done"]
+
+ def test_a_retry_after_the_last_tool_cycle_only_asks_for_the_answer(self,
task):
+ with pytest.raises(RuntimeError):
+ task.run(fail_at_turn=2)
+
+ assert task.run() == "done"
+
+ assert task.tool_calls == {"first": 1, "second": 1, "third": 1}
+ assert task.model_turns == [0, 1, 2, 2]
+
+ def
test_a_failed_invocation_leaves_nothing_of_itself_in_the_saved_state(self,
task):
+ """Strands appends a cut-off answer before raising; the retry must not
inherit it."""
+ with pytest.raises(MaxTokensReachedException):
+ task.run(truncate_at_turn=1)
+
+ assert task.run() == "done"
+
+ assert _get_texts(task.agents[-1]) == ["go", "done"]
+ assert task.tool_calls == {"first": 1, "second": 1, "third": 1}
+
+ @pytest.mark.parametrize(
+ ("killed_in", "tool_calls"),
+ [
+ pytest.param("first", {"first": 2, "second": 1, "third": 1},
id="cycle 1"),
+ pytest.param("third", {"first": 1, "second": 2, "third": 2},
id="cycle 2"),
+ ],
+ )
+ def test_a_crash_inside_a_cycle_runs_all_of_its_tools_again(self, task,
killed_in, tool_calls):
+ """Strands resumes from a cycle boundary, so a cycle's tool batch is
all or nothing."""
+ with pytest.raises(_WorkerKilled):
+ task.run(kill_in_tool=killed_in)
+
+ assert task.run() == "done"
+
+ assert task.tool_calls == tool_calls
+ assert _get_texts(task.agents[-1]) == ["go", "done"]
+
+ def test_a_prompt_with_bytes_resumes(self, task):
+ """A document in the prompt is fingerprinted the way Strands saves
it."""
+ prompt = [
+ {"text": "go"},
+ {"document": {"format": "txt", "name": "orders", "source":
{"bytes": b"42"}}},
+ ]
+ with pytest.raises(RuntimeError):
+ task.run(prompt, fail_at_turn=1)
+
+ assert task.run(prompt) == "done"
+
+ assert task.tool_calls == {"first": 1, "second": 1, "third": 1}
+
+ def
test_a_failure_before_the_first_checkpoint_starts_over_with_one_prompt(self,
task):
+ with pytest.raises(RuntimeError):
+ task.run(fail_at_turn=0)
+
+ assert task.run() == "done"
+
+ assert _get_texts(task.agents[-1]) == ["go", "done"]
+
+ def test_a_retry_with_a_different_prompt_starts_over(self, task):
+ with pytest.raises(RuntimeError):
+ task.run(fail_at_turn=1)
+
+ assert task.run("go, but differently") == "done"
+
+ assert task.tool_calls == {"first": 2, "second": 1, "third": 1}
+ assert _get_texts(task.agents[-1]) == ["go, but differently", "done"]
+
+ @pytest.mark.parametrize("fail_at_turn", [None, 1], ids=["first try",
"after a retry"])
+ def test_the_saved_state_is_deleted_once_the_agent_finishes(self, task,
fail_at_turn):
+ if fail_at_turn is not None:
+ with pytest.raises(RuntimeError):
+ task.run(fail_at_turn=fail_at_turn)
+ assert task.store.store
+
+ task.run()
+
+ assert task.store.store == {}
+
+ def
test_a_failed_delete_after_the_agent_finishes_does_not_fail_the_task(self,
task):
+ class _DeleteFailsStore(FakeTaskStateStore):
+ def delete(self, key):
+ if key in self.store:
+ raise RuntimeError("api down")
+ super().delete(key)
+
+ task.store = _DeleteFailsStore()
+
+ assert task.run() == "done"
+
+ def test_a_strands_storage_can_replace_the_task_state_store(self, task):
+ storage = InMemoryStorage()
+ with pytest.raises(RuntimeError):
+ task.run(fail_at_turn=1, storage=storage)
+
+ assert task.run(storage=storage) == "done"
+
+ assert task.tool_calls == {"first": 1, "second": 1, "third": 1}
+ assert task.store.store == {}
+ assert asyncio.run(storage.list("")) == []
+
+ def
test_mapped_task_instances_sharing_a_storage_keep_their_own_state(self, task):
+ storage = InMemoryStorage()
+ with pytest.raises(RuntimeError):
+ task.run(fail_at_turn=1, map_index=0, storage=storage)
+
+ # Another mapped task instance starts its own agent rather than
resuming the first one's.
+ assert task.run(map_index=1, storage=storage) == "done"
+ assert task.tool_calls["first"] == 2
+
+ # And the first one still resumes where it stopped.
+ assert task.run(map_index=0, storage=storage) == "done"
+ assert task.tool_calls["first"] == 2
+
+ def test_refuses_a_stateful_model(self, task):
+ """Strands drops a restored conversation for a model that keeps it
server-side."""
+ with (
+ mock.patch.object(_ScriptedModel, "stateful",
new_callable=mock.PropertyMock, return_value=True),
+ pytest.raises(ValueError, match="stateful"),
+ ):
+ # The model fails at once if it is ever called, rather than loop
on lost history.
+ task.run(fail_at_turn=0)
+
+ assert task.model_turns == []
+
+ def test_needs_a_storage_below_airflow_3_3(self, task):
+ with (
+ mock.patch.object(durable_strands, "AIRFLOW_V_3_3_PLUS", False),
+ pytest.raises(AirflowOptionalProviderFeatureException,
match="Airflow 3.3"),
+ ):
+ task.run()
diff --git
a/providers/common/ai/tests/unit/common/ai/durable/test_strands_storage.py
b/providers/common/ai/tests/unit/common/ai/durable/test_strands_storage.py
new file mode 100644
index 00000000000..5fed3a69620
--- /dev/null
+++ b/providers/common/ai/tests/unit/common/ai/durable/test_strands_storage.py
@@ -0,0 +1,156 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import asyncio
+import time
+
+import pytest
+
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS
+
+if not AIRFLOW_V_3_3_PLUS:
+ pytest.skip("task state store needs Airflow >= 3.3",
allow_module_level=True)
+
+pytest.importorskip("strands")
+
+from strands import Agent
+from strands.session import SnapshotSessionManager
+
+from airflow.providers.common.ai.durable.base import DURABLE_KEY_PREFIX
+from airflow.providers.common.ai.durable.strands_storage import
TaskStateStoreStorage
+from airflow.sdk.execution_time.context import NEVER_EXPIRE
+
+from unit.common.ai.durable.test_task_state_store import FakeTaskStateStore
+from unit.common.ai.tools.test_strands import _OneToolCallModel
+
+
[email protected]
+def accessor():
+ return FakeTaskStateStore()
+
+
[email protected]
+def storage(accessor):
+ return TaskStateStoreStorage(accessor)
+
+
+class TestTaskStateStoreStorage:
+ @pytest.mark.parametrize(
+ "data",
+ [
+ pytest.param(b'{"messages": ["\xe2\x82\xac 42"]}', id="utf-8"),
+ pytest.param(b"\x00\xff\xfe binary", id="binary"),
+ pytest.param(b"", id="empty"),
+ ],
+ )
+ def test_read_returns_the_bytes_written(self, storage, data):
+ asyncio.run(storage.write("session/a/snapshot.json", data))
+
+ assert asyncio.run(storage.read("session/a/snapshot.json")) == data
+
+ def test_read_of_a_missing_key_returns_none(self, storage):
+ assert asyncio.run(storage.read("session/missing.json")) is None
+
+ def test_writes_under_the_reserved_prefix_without_expiry(self, storage,
accessor):
+ asyncio.run(storage.write("session/a/snapshot.json", b"{}"))
+
+ assert set(accessor.store) == {
+ f"{DURABLE_KEY_PREFIX}strands/session/a/snapshot.json",
+ f"{DURABLE_KEY_PREFIX}strands_index",
+ }
+ assert set(accessor.set_retentions.values()) == {NEVER_EXPIRE}
+
+ def test_leaves_user_keys_alone(self, storage, accessor):
+ accessor.set("job_id", "abc")
+ asyncio.run(storage.write("job_id", b"agent"))
+ asyncio.run(storage.delete("job_id"))
+
+ assert accessor.store == {"job_id": "abc"}
+
+ def test_list_returns_matching_keys_sorted(self, storage):
+ for key in ("session/b/x.json", "context/a/y.json",
"session/a/z.json"):
+ asyncio.run(storage.write(key, b"{}"))
+
+ assert asyncio.run(storage.list("session/")) == ["session/a/z.json",
"session/b/x.json"]
+ assert asyncio.run(storage.list("")) == ["context/a/y.json",
"session/a/z.json", "session/b/x.json"]
+
+ def test_overwriting_a_key_lists_it_once(self, storage):
+ asyncio.run(storage.write("session/a.json", b"1"))
+ asyncio.run(storage.write("session/a.json", b"2"))
+
+ assert asyncio.run(storage.list("")) == ["session/a.json"]
+ assert asyncio.run(storage.read("session/a.json")) == b"2"
+
+ def test_deleting_every_key_leaves_nothing_in_the_store(self, storage,
accessor):
+ asyncio.run(storage.write("session/a.json", b"1"))
+ asyncio.run(storage.write("session/b.json", b"2"))
+
+ asyncio.run(storage.delete("session/a.json"))
+ assert asyncio.run(storage.list("")) == ["session/b.json"]
+ asyncio.run(storage.delete("session/b.json"))
+
+ assert accessor.store == {}
+
+ def test_deleting_a_missing_key_is_a_no_op(self, storage, accessor):
+ asyncio.run(storage.delete("session/missing.json"))
+
+ assert accessor.store == {}
+
+ def test_concurrent_writes_keep_every_key_listed(self):
+ """Strands deletes a session's keys concurrently; index updates must
not lose each other."""
+
+ class _SlowStore(FakeTaskStateStore):
+ def get(self, key, default=None):
+ value = super().get(key, default)
+ # Between this read of the index and the write that follows
it, so that
+ # unserialized updates would overwrite each other.
+ time.sleep(0.01)
+ return value
+
+ storage = TaskStateStoreStorage(_SlowStore())
+ keys = [f"session/{n}.json" for n in range(8)]
+
+ async def write_all():
+ await asyncio.gather(*(storage.write(key, b"{}") for key in keys))
+
+ asyncio.run(write_all())
+
+ assert asyncio.run(storage.list("")) == keys
+
+ def test_instances_over_one_task_instance_share_their_keys(self, accessor):
+ asyncio.run(TaskStateStoreStorage(accessor).write("session/a.json",
b"1"))
+ asyncio.run(TaskStateStoreStorage(accessor).write("session/b.json",
b"2"))
+
+ assert asyncio.run(TaskStateStoreStorage(accessor).list("")) ==
["session/a.json", "session/b.json"]
+
+ def test_backs_a_snapshot_session_manager(self, storage, accessor):
+ session_manager = SnapshotSessionManager("run-1", storage=storage)
+ agent = Agent(
+ model=_OneToolCallModel("lookup", {}),
session_manager=session_manager, callback_handler=None
+ )
+ agent.messages.append({"role": "user", "content": [{"text": "hello"}]})
+ session_manager.sync_agent(agent)
+
+ restored_manager = SnapshotSessionManager("run-1", storage=storage)
+ restored = Agent(
+ model=_OneToolCallModel("lookup", {}),
session_manager=restored_manager, callback_handler=None
+ )
+ assert restored.messages == [{"role": "user", "content": [{"text":
"hello"}]}]
+
+ asyncio.run(restored_manager.delete_session())
+ assert accessor.store == {}
diff --git a/providers/common/ai/tests/unit/common/ai/tools/test_strands.py
b/providers/common/ai/tests/unit/common/ai/tools/test_strands.py
index bc08c8b442d..2da97fb74aa 100644
--- a/providers/common/ai/tests/unit/common/ai/tools/test_strands.py
+++ b/providers/common/ai/tests/unit/common/ai/tools/test_strands.py
@@ -19,7 +19,7 @@ from __future__ import annotations
import asyncio
import copy
import json
-from typing import Any
+from typing import TYPE_CHECKING, Any
from unittest.mock import MagicMock, patch
import pytest
@@ -40,6 +40,9 @@ from airflow.providers.common.ai.toolsets.sql import
SQLToolset
from unit.common.ai.toolsets.test_sql import _make_mock_db_hook
+if TYPE_CHECKING:
+ from strands.agent import AgentResult
+
class _OneToolCallModel(Model):
"""Asks for one call of ``tool_name``, then answers with the status and
text it got back."""
@@ -74,6 +77,37 @@ class _OneToolCallModel(Model):
yield {"messageStop": {"stopReason": "end_turn"}}
+class _RetryingModel(_OneToolCallModel):
+ """Calls ``tool_name`` again after every error result, up to ``attempts``
times, then answers."""
+
+ def __init__(self, tool_name: str, tool_input: dict[str, Any], *,
attempts: int) -> None:
+ super().__init__(tool_name, tool_input)
+ self._attempts = attempts
+
+ async def stream(self, messages, tool_specs=None, system_prompt=None,
**kwargs: Any):
+ calls = sum(1 for message in messages if any("toolResult" in block for
block in message["content"]))
+ yield {"messageStart": {"role": "assistant"}}
+ if calls < self._attempts:
+ start = {"toolUse": {"toolUseId": f"call-{calls}", "name":
self._tool_name}}
+ yield {"contentBlockStart": {"start": start}}
+ yield {"contentBlockDelta": {"delta": {"toolUse": {"input":
json.dumps(self._tool_input)}}}}
+ yield {"contentBlockStop": {}}
+ yield {"messageStop": {"stopReason": "tool_use"}}
+ return
+ yield {"contentBlockDelta": {"delta": {"text": "gave up"}}}
+ yield {"contentBlockStop": {}}
+ yield {"messageStop": {"stopReason": "end_turn"}}
+
+
+def _run_through_checkpoints(agent: Agent, prompt: str) -> AgentResult:
+ """Resume a checkpointing agent from each checkpoint it returns, until it
finishes."""
+ result = agent(prompt)
+ while result.checkpoint is not None:
+ # Strands types the prompt as AgentInput, which leaves out its
checkpointResume block.
+ result = agent({"checkpointResume": {"checkpoint":
result.checkpoint.to_dict()}}) # type: ignore[arg-type]
+ return result
+
+
def _tool(result: ToolResult | Exception) -> AirflowTool:
async def function(arguments: dict[str, Any]) -> ToolResult:
if isinstance(result, Exception):
@@ -188,6 +222,24 @@ class TestAgentRun:
assert str(agent("go")).startswith("error:")
assert str(agent("go again")).startswith("error:")
+ def test_resuming_from_a_checkpoint_keeps_the_retry_budget(self):
+ """With checkpointing every cycle is its own invocation; the budget
must still span the run."""
+
+ def lookup(key: str) -> str:
+ raise ModelRetry("no such key")
+
+ agent = Agent(
+ model=_RetryingModel("lookup", {"key": "a"}, attempts=3),
+
plugins=[AirflowTools(*airflow_tools_from_toolset(FunctionToolset([lookup],
max_retries=1)))],
+ callback_handler=None,
+ checkpointing=True,
+ )
+
+ with pytest.raises(EventLoopException) as caught:
+ _run_through_checkpoints(agent, "go")
+
+ assert isinstance(caught.value.original_exception, ToolCallError)
+
@pytest.mark.enable_redact
def test_a_secret_in_a_result_reaches_the_model_masked(self,
registered_secret):
answer = _run_agent(
diff --git a/scripts/in_container/run_agent_framework_tests.sh
b/scripts/in_container/run_agent_framework_tests.sh
index 2053c68fdc1..c2b584697f1 100755
--- a/scripts/in_container/run_agent_framework_tests.sh
+++ b/scripts/in_container/run_agent_framework_tests.sh
@@ -20,9 +20,9 @@
#
# Strands Agents caps mcp, and Google ADK caps opentelemetry and websockets,
below the versions
# uv.lock resolves, so their adapter tests are skipped everywhere else in CI.
This installs one
-# framework into the CI image and runs the tests of the framework-neutral
tools and their adapters;
-# the other framework's tests skip themselves. One framework per invocation,
so a bad release of
-# one cannot mask the other.
+# framework into the CI image and runs the tests of the framework-neutral
tools and their adapters,
+# and of durable execution for Strands; the other framework's tests skip
themselves. One framework
+# per invocation, so a bad release of one cannot mask the other.
#
# By default every package already in the image is held at its installed
version with uv's
# --override, so the framework is tested against the same dependencies as the
rest of Airflow and
@@ -32,7 +32,11 @@
# The newest framework release older than the repository's uv exclude-newer
window is installed.
set -euo pipefail
-TEST_PATH="providers/common/ai/tests/unit/common/ai/tools"
+TEST_PATHS=(
+ "providers/common/ai/tests/unit/common/ai/tools"
+ "providers/common/ai/tests/unit/common/ai/durable/test_strands.py"
+ "providers/common/ai/tests/unit/common/ai/durable/test_strands_storage.py"
+)
framework="${1:-}"
case "${framework}" in
@@ -87,4 +91,4 @@ uv pip freeze | grep -iE
'^(strands-agents|google-adk|mcp|opentelemetry-(api|sdk
python -c "${import_check}"
# --skip-db-tests: the job runs with backend "none", which has no database to
set up.
-pytest "${TEST_PATH}" --skip-db-tests -p no:cacheprovider --color=yes -ra
+pytest "${TEST_PATHS[@]}" --skip-db-tests -p no:cacheprovider --color=yes -ra
