Lee-W commented on code in PR #74404: URL: https://github.com/apache/airflow/pull/74404#discussion_r4259144959
########## 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 Review Comment: ```suggestion from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException try: from strands.session import SnapshotSessionManager from strands.types.session import encode_bytes_values except ImportError as e: raise AirflowOptionalProviderFeatureException(e) import hashlib import json from typing import TYPE_CHECKING, Any, Protocol from airflow.providers.common.ai.batch.state import compute_identity_key from airflow.providers.common.compat.sdk import get_current_context 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 ``` ########## providers/common/ai/docs/frameworks/strands.rst: ########## @@ -159,17 +159,139 @@ 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.51 or later, which the version in Review Comment: Should we make it 1.56.0 or change the install section to 1.51 instead? ########## 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 = _fingerprint(prompt) + + session_manager, agent = _build(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 = _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_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(_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_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 _fingerprint(prompt: AgentInput) -> str: Review Comment: ```suggestion def _build_fingerprint(prompt: AgentInput) -> str: ``` ########## 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 = _fingerprint(prompt) + + session_manager, agent = _build(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 = _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_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(_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_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 _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 _resume(checkpoint: dict[str, Any]) -> dict[str, Any]: Review Comment: ```suggestion def _build_checkpoint_resume(checkpoint: dict[str, Any]) -> dict[str, Any]: ``` ########## 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 _texts(agent: Agent) -> list[str]: Review Comment: ```suggestion def _build_texts(agent: Agent) -> list[str]: ``` ########## 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) + Review Comment: ```suggestion from __future__ import annotations from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException try: from strands.storage import Storage except ImportError as e: raise AirflowOptionalProviderFeatureException(e) import asyncio import base64 import bisect import builtins import threading from typing import TYPE_CHECKING ``` ########## 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") Review Comment: ```suggestion 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") import asyncio import time ``` ########## 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 Review Comment: ```suggestion 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") 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 from strands import Agent, tool ``` ########## 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) Review Comment: ```suggestion from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException try: from strands.session import SnapshotSessionManager from strands.types.session import encode_bytes_values except ImportError as e: raise AirflowOptionalProviderFeatureException(e) import hashlib import json from typing import TYPE_CHECKING, Any, Protocol from airflow.providers.common.compat.sdk import get_current_context ``` ########## 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 = _fingerprint(prompt) + + session_manager, agent = _build(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 = _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_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(_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_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 _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 _resume(checkpoint: dict[str, Any]) -> dict[str, Any]: Review Comment: ```suggestion def _build_checkpoint_resume(checkpoint: dict[str, Any]) -> dict[str, Any]: ``` ########## providers/common/ai/tests/unit/common/ai/tools/test_strands.py: ########## @@ -188,6 +222,25 @@ def lookup(key: str) -> str: 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: + """Look a key up.""" Review Comment: ```suggestion ``` ########## 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 = _fingerprint(prompt) + + session_manager, agent = _build(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 = _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_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(_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_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 _fingerprint(prompt: AgentInput) -> str: Review Comment: ```suggestion def _build_fingerprint(prompt: AgentInput) -> str: ``` ########## providers/common/ai/docs/frameworks/strands.rst: ########## @@ -159,17 +159,139 @@ 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.51 or later, which the version in Review Comment: The version we use the in the example Dag is 1.56.0. should we unify them to one verison -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
