You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ka...@apache.org on 2026/10/14 07:23:44 UTC
(airflow) branch main updated: Add AirflowDurability capability for durable agent replay across task retries (#74385)
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 68a4c3177f9 Add AirflowDurability capability for durable agent replay across task retries (#74385)
68a4c3177f9 is described below
commit 68a4c3177f94738508aade194256753ed5157e87
Author: Kaxil Naik <ka...@apache.org>
AuthorDate: Wed Oct 14 08:23:35 2026 +0100
Add AirflowDurability capability for durable agent replay across task retries (#74385)
AgentOperator(durable=True) recorded steps through CachingModel and
CachingToolset, which only saw what the operator wrapped itself: tools from
MCP servers, per-run toolsets and most capabilities, tool discovery, other
capabilities' @durable_operation methods (pydantic-ai-harness SpendLimits,
SummarizingCompaction) and streamed requests ran again on every retry.
Replace them with AirflowDurability, built on pydantic-ai's durable backend
API, over a framework-neutral journal kept in the task state store (or on
ObjectStorage below Airflow 3.3). durable=True attaches it, and a
pydantic-ai agent in a plain @task can attach it directly.
Under durable=True the agent keeps the first attempt's run_id across
retries, every FunctionToolset, MCPToolset and DynamicToolset needs an id,
and a tool result that is not JSON-serializable fails the task without
retrying. The pydantic-ai-slim floor goes to 2.36.0.
* Clean up nested durable runs and replay unloadable records live
Cleanup now deletes the steps of every run any attempt started, through
a registry kept with the journal: the task state store cannot list keys,
and a run an agent's tool started never starts again once that tool call
replays. A recorded model response or tool result that no longer loads
(written by another version) now runs live with a warning instead of
failing every retry. Also renames DurableJournal.run_id to get_run_id,
fixes two mypy errors in the tests, and drops the CachingToolset entry
from the provider yaml check.
* Keep durable replay matching after masked tool results and stale records
The journal records tool results with registered secrets masked, but a
live run hands the model the unmasked result. A retry replaying such a
result put ``***`` in the history, so the next model request's
fingerprint no longer matched and every later step ran live, including
tools that had already run. Model-request fingerprints now mask the
history the same way.
A recorded payload that no longer decodes is now checked with the result
codec pydantic-ai decodes it with, so a stale tool result, capability
operation or model response runs live instead of reaching the model as a
validation error or failing the retry. Also updates uv.lock for the
pydantic-ai-slim floor and the comment on AirflowToolset.replayable.
---
providers/common/ai/README.rst | 2 +-
providers/common/ai/docs/capabilities.rst | 18 +-
providers/common/ai/docs/durable_execution.rst | 332 +++++---
providers/common/ai/docs/index.rst | 12 +-
providers/common/ai/docs/operators/agent.rst | 24 +-
providers/common/ai/docs/toolsets/index.rst | 12 +-
providers/common/ai/docs/troubleshooting.rst | 26 +-
providers/common/ai/pyproject.toml | 17 +-
.../providers/common/ai/durable/__init__.py | 19 +
.../airflow/providers/common/ai/durable/base.py | 59 +-
.../providers/common/ai/durable/caching_model.py | 226 ------
.../providers/common/ai/durable/caching_toolset.py | 141 ----
.../providers/common/ai/durable/capability.py | 532 +++++++++++++
.../providers/common/ai/durable/fingerprint.py | 9 +-
.../airflow/providers/common/ai/durable/journal.py | 551 +++++++++++++
.../providers/common/ai/durable/replay_usage.py | 149 +++-
.../providers/common/ai/durable/step_counter.py | 52 --
.../airflow/providers/common/ai/durable/storage.py | 140 +---
.../common/ai/durable/task_state_store.py | 155 +---
.../ai/example_dags/example_agent_durable.py | 35 +-
.../src/airflow/providers/common/ai/exceptions.py | 10 +
.../airflow/providers/common/ai/operators/agent.py | 325 +++-----
.../airflow/providers/common/ai/toolsets/hook.py | 4 +-
.../providers/common/ai/toolsets/managed_agent.py | 6 +-
.../providers/common/ai/utils/toolset_base.py | 5 +
.../durable/{test_step_counter.py => conftest.py} | 22 +-
.../tests/unit/common/ai/durable/memory_storage.py | 58 ++
.../ai/tests/unit/common/ai/durable/test_base.py | 15 +-
.../unit/common/ai/durable/test_caching_model.py | 432 -----------
.../unit/common/ai/durable/test_caching_toolset.py | 268 -------
.../unit/common/ai/durable/test_capability.py | 855 +++++++++++++++++++++
.../unit/common/ai/durable/test_fingerprint.py | 12 +
.../tests/unit/common/ai/durable/test_journal.py | 427 ++++++++++
.../unit/common/ai/durable/test_replay_cost.py | 231 ++----
.../common/ai/durable/test_replay_verification.py | 138 ----
.../tests/unit/common/ai/durable/test_storage.py | 227 ++----
.../common/ai/durable/test_task_state_store.py | 212 +----
.../tests/unit/common/ai/operators/test_agent.py | 603 ++++++---------
.../unit/common/ai/utils/test_tool_metrics.py | 61 +-
.../in_container/run_provider_yaml_files_check.py | 3 -
uv.lock | 14 +-
41 files changed, 3595 insertions(+), 2844 deletions(-)
diff --git a/providers/common/ai/README.rst b/providers/common/ai/README.rst
index b00acddbbd2..a7f249e9a81 100644
--- a/providers/common/ai/README.rst
+++ b/providers/common/ai/README.rst
@@ -56,7 +56,7 @@ PIP package Version required
``apache-airflow`` ``>=2.11.0``
``apache-airflow-providers-common-compat`` ``>=1.21.0``
``apache-airflow-providers-standard`` ``>=1.20.0``
-``pydantic-ai-slim`` ``>=2.33.0``
+``pydantic-ai-slim`` ``>=2.36.0``
``structlog`` ``>=24.2.0``
========================================== ==================
diff --git a/providers/common/ai/docs/capabilities.rst b/providers/common/ai/docs/capabilities.rst
index 1a2b76e9959..12bf47565b6 100644
--- a/providers/common/ai/docs/capabilities.rst
+++ b/providers/common/ai/docs/capabilities.rst
@@ -96,7 +96,7 @@ against another capability, such as a guardrail.
With durable execution
----------------------
-With ``durable=True``, a retry replays completed steps from the cache instead of running them
+With ``durable=True``, a retry replays completed steps from the journal instead of running them
again. Whether a capability's work is replayed depends on where it runs:
.. list-table::
@@ -106,21 +106,19 @@ again. Whether a capability's work is replayed depends on where it runs:
- On retry
* - ``Thinking``, and ``WebSearch``, ``WebFetch`` or ``ImageGeneration`` when the model's
provider runs the tool natively
- - Replayed with the cached model response.
+ - Replayed with the recorded model response.
* - ``WebSearch``, ``WebFetch`` or ``ImageGeneration`` falling back to a local tool, for a
- provider without the native one
- - The local tool runs again.
- * - ``Toolset`` holding a toolset
- - Tool results are replayed.
- * - ``MCP``, ``PrefixTools``, ``CombinedCapability``, a ``Toolset`` built from a function,
- and capabilities from an agent spec file
- - Tools run again. Pass tools you need replayed in ``toolsets=`` instead.
+ provider without the native one; ``Toolset``, ``MCP``, ``PrefixTools``,
+ ``CombinedCapability``, and capabilities from an agent spec file
+ - Tool results are replayed. Each toolset they contribute needs an ``id``; pydantic-ai
+ names one without it when the agent is built. An ``id`` on a ``Toolset`` capability
+ does not reach the toolset inside it, so set it on that toolset.
* - pydantic-ai-harness ``CodeMode``
- Not allowed: the operator raises ``ValueError``. This includes a ``CodeMode`` inside a
``CombinedCapability`` or a wrapper such as ``PrefixTools``, but not one a capability
function builds when the run starts.
-See :doc:`durable_execution` for how the cache works.
+See :doc:`durable_execution` for how replay works.
Serialization
-------------
diff --git a/providers/common/ai/docs/durable_execution.rst b/providers/common/ai/docs/durable_execution.rst
index 0de990ce9c8..10555ab20e1 100644
--- a/providers/common/ai/docs/durable_execution.rst
+++ b/providers/common/ai/docs/durable_execution.rst
@@ -25,141 +25,247 @@ Durable execution
Experimental: this can change or be removed in a minor release of this provider.
See :ref:`howto/stability`.
-Agent tasks can involve multiple LLM calls and tool invocations. If a task
-fails mid-run (network error, timeout, transient API failure), a plain retry
-re-executes every LLM call and tool call from scratch -- repeating work that
-already succeeded and incurring additional cost.
+An agent task makes several model calls and tool calls. When it fails partway (a
+network error, a timeout, a provider outage), a plain Airflow retry runs every one of
+them again: the model calls are paid for twice and the tools' side effects happen
+twice. With durable execution, each step the agent completes is recorded in a journal
+kept for the task instance, and the retry replays the recorded steps instead of running
+them again. Only the steps the failed attempt never finished run live.
-Setting ``durable=True`` caches each LLM response and tool result as it
-completes. On retry, completed steps are replayed from the cache and only the
-remaining steps run against the live model and tools. The cache is deleted
-after successful completion.
+Durable execution only helps a task that has retries. To decide *whether* a failure is
+worth retrying, see :doc:`retry_policies`; a retried ``LLMBatchOperator`` re-attaches to
+its running batch instead (:ref:`llm-batch-reattach`). To resume a Strands Agents agent on a
+retry, see :ref:`howto/frameworks:strands-durable`.
-Durable execution only helps when the task has retries configured. Without
-retries there is nothing to replay.
+Turn it on
+----------
-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 resume a
-Strands Agents agent on a retry, see :ref:`howto/frameworks:strands-durable`.
+Set ``durable=True`` on ``AgentOperator``:
-**Configuration**
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent_durable.py
+ :language: python
+ :start-after: [START howto_operator_agent_durable]
+ :end-before: [END howto_operator_agent_durable]
+
+The same flag works on ``@task.agent``:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent_durable.py
+ :language: python
+ :start-after: [START howto_decorator_agent_durable]
+ :end-before: [END howto_decorator_agent_durable]
+
+``durable=True`` attaches the
+:class:`~airflow.providers.common.ai.durable.AirflowDurability` capability to the
+agent. A pydantic-ai agent you build yourself in a plain ``@task`` takes the capability
+directly:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent_durable.py
+ :language: python
+ :start-after: [START howto_capability_airflow_durability]
+ :end-before: [END howto_capability_airflow_durability]
+
+The capability needs a name for the agent's steps: the agent's ``name``, or
+``AirflowDurability(name=...)``. ``AgentOperator`` uses the agent's name when
+``agent_params`` sets one, and the task id otherwise. Outside a running Airflow task
+the capability does nothing, so the same agent runs normally in a test or a notebook.
+
+Give every toolset an id
+------------------------
+
+pydantic-ai records each call of a ``FunctionToolset``, ``MCPToolset`` or
+``DynamicToolset`` under the toolset's ``id``, and refuses to run one without it. Set
+the id on the toolset itself:
+
+.. code-block:: python
-On **Airflow >= 3.3** the cache is stored in the
-:doc:`task state store <apache-airflow:core-concepts/task-state-store>`,
-scoped to the task instance. No configuration is required; the store handles
-persistence across retries.
+ orders = FunctionToolset(tools=[get_order], id="orders")
-By default each cached step is written to the Airflow metadata database. Model
-responses and large tool results can be sizable, so for agents with large
-payloads configure ``[workers] state_store_backend`` to offload step values to
-external storage (e.g. object storage) instead of the metadata database; the
-provider then stores only a reference in the database.
+``AgentOperator(durable=True)`` checks this when the Dag is parsed and fails with:
-On **Airflow < 3.3** the cache is persisted to ObjectStorage and the location
-must be set in ``airflow.cfg``. The task raises ``ValueError`` at runtime if
-``durable=True`` and the option is missing.
+.. code-block:: text
+
+ ValueError: durable=True needs a unique id on every FunctionToolset, such as FunctionToolset(..., id='orders'): durable replay records each tool call under it. Set it on the toolset itself, and pass a function that builds a toolset as DynamicToolset(function, id=...); the id of a Toolset capability around a toolset does not reach it.
+
+Keep an id stable once tasks run with it: a retry finds its recorded tool calls by it.
+The provider's own toolsets (``SQLToolset``, ``HookToolset`` and the rest) and tools
+passed as ``agent_params={"tools": [...]}`` need nothing. A capability that contributes
+``@durable_operation`` methods needs an ``id`` for the same reason; pydantic-ai raises
+``UserError`` naming it when it has none. pydantic-ai-harness ``SpendLimits`` has one by
+default.
+
+Where the journal is kept
+-------------------------
+
+On **Airflow 3.3 and later** the journal is kept in the
+:doc:`task state store <apache-airflow:core-concepts/task-state-store>`, scoped to the
+task instance. No configuration is needed. Each step is written to the Airflow metadata
+database; for agents with large model responses or tool results, set
+``[workers] state_store_backend`` to offload the values to external storage instead.
+
+On **Airflow before 3.3** the journal is a JSON file on ObjectStorage, and the location
+must be set:
.. code-block:: ini
[common.ai]
- # Local filesystem -- suitable for development
+ # Local filesystem, suitable for development
durable_cache_path = file:///tmp/airflow_durable_cache
-The value is an ObjectStorage URI, so any supported backend works. For
-production, use a shared store so retries on a different worker can read the
-cache:
+Any ObjectStorage URI works. In production use a shared store, so a retry on another
+worker can read the journal:
.. code-block:: ini
[common.ai]
durable_cache_path = s3://my-bucket/airflow/durable-cache
-**Operator example**
+Without the option, the task fails with:
-.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent_durable.py
- :language: python
- :start-after: [START howto_operator_agent_durable]
- :end-before: [END howto_operator_agent_durable]
+.. code-block:: text
+
+ ValueError: durable=True requires [common.ai] durable_cache_path to be set. Example: durable_cache_path = file:///tmp/airflow_durable_cache
-**Decorator example**
-
-.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent_durable.py
- :language: python
- :start-after: [START howto_decorator_agent_durable]
- :end-before: [END howto_decorator_agent_durable]
-
-**How it works**
-
-1. On first execution, each LLM response and tool result is saved as the agent
- progresses, together with a fingerprint of the request that produced it
- (model, message history, settings, and tools for LLM steps; tool name,
- arguments, and call id for tool steps).
-2. If the task fails and Airflow retries it, completed steps are loaded from
- the cache and returned without calling the model or tool. Steps not yet in
- the cache proceed normally.
-3. Before a step is replayed, its stored fingerprint is compared against the
- current request. If anything changed between attempts -- the system
- prompt, the model, the toolset, model settings, or the conversation so
- far -- the stale entry is discarded, a warning is logged, and the step
- re-runs live. A divergence also invalidates the steps after it: re-running
- an LLM step produces fresh tool call ids, so tool results recorded under
- the old conversation no longer match. A changed agent costs a re-run; it
- never replays responses that belong to a different conversation.
-4. After successful completion, the cached steps are deleted.
-
-Replay verification compares the **requests** sent to models and tools, not
-the code behind them. Editing a tool's implementation between attempts does
-not invalidate an already-cached result for an identical call, and pointing
-``llm_conn_id`` at a different endpoint serving the same model name does not
-invalidate cached responses -- clear the cache to force a fully fresh run.
-
-When the run ends, successfully or not, an INFO line reports how many steps
-were replayed from the cache and how many new steps were cached. Per-step
-detail is available at DEBUG level.
-
-The cache is scoped to a single task instance (Dag id, run id, task id, and
-map index), so each run replays only its own steps. On Airflow >= 3.3 the cache
-lives in the task state store and is removed when the Dag run is cleaned up; on
-Airflow < 3.3 it is a JSON file named ``{dag_id}_{task_id}_{run_id}.json`` (with
-``_{map_index}`` appended for mapped tasks) under the configured
-``durable_cache_path``.
-
-.. note::
+How a retry replays
+-------------------
- Runs that fail permanently (exhaust all retries) leave their cached steps
- behind. These do not affect future Dag runs (each run is scoped separately).
- On Airflow >= 3.3 they are reclaimed when the Dag run is removed; on Airflow
- < 3.3 the orphaned JSON files consume storage until cleaned up, so add a
- lifecycle policy to the storage backend or remove them periodically.
-
-**Side effects and idempotency**
-
-Durable execution caches **return values**, not side effects. When a step is
-replayed, the tool's code does not run -- only the stored return value is
-returned. Two things follow from this:
-
-- If a tool completed successfully and its result was cached, the tool will
- **not** run again on retry. Any side effect it produced (writing a file,
- sending a message) already happened during the original run and is not
- repeated.
-- If a tool fails *before* its result is cached, it **will** run again on
- retry. A tool that partially completed (e.g. sent an email then raised an
- exception) may produce the side effect a second time.
-
-All built-in toolsets (``SQLToolset`` with ``allow_writes=False``,
-``HookToolset`` in read-only mode) are read-only and replay safely. For custom
-tools with non-idempotent side effects, design the tool to be idempotent. For
-example, check whether the operation already completed before acting, or
-use database constraints to prevent duplicate writes.
-
-Tool results must be JSON-serializable to be cached. If a tool returns a
-non-serializable value (e.g. ``BinaryContent`` from MCP tools), or a write to
-the task state store fails, the step is not cached and runs again on retry
-instead of replaying. The step itself still succeeds. Each such step logs a
-WARNING naming the tool, and the end-of-run summary lists every tool that was
-not cached. A step that re-runs can change what the model sees next, so the
-steps after it may re-run too.
+Each step the agent completes is recorded under its position in the run, with its
+name (which model or tool it was) and a fingerprint of what it was asked: the model,
+the message history, the settings and the tool definitions for a model request; the
+tool name, its arguments and the model-issued call id for a tool call.
+
+The retry runs the agent from the start. At each position, a step with the same name
+and the same fingerprint as the recorded one is replayed: the recorded result is
+returned and the model or tool is not called. The first step that differs means the run
+has taken another path, because the prompt, the model, a setting, a tool or the
+conversation changed, so that step and every step after it run live and are recorded
+again, and the task log says so:
+
+.. code-block:: text
+
+ Durable: the run took a different path from the previous attempt; this step and every step after it run again
+
+A changed agent costs a re-run; it never replays results that belong to another
+conversation. Settings that only affect transport, ``timeout`` and the prompt cache
+settings, are left out of the fingerprint, so changing them does not cause a re-run.
+Replay compares requests only: editing a tool's implementation does not invalidate a
+recorded result for an identical call.
+
+When a run fails, the journal notes where. The retry runs the step that failed again,
+and runs live everything the failed attempt did after it, such as the error handling of
+other capabilities. Tool calls already running alongside the failed one still replay. A
+step that raised and that the run recovered from, such as a tool error the model was
+told about, does not stop the steps after it from replaying.
+
+Some steps have no request to fingerprint: another capability's
+``@durable_operation``, or tool discovery. Such a step replays on its name, and only when
+the step before it also replayed.
+
+The journal records:
+
+- model requests, including streamed ones;
+- calls to tools from any toolset, including the provider's own toolsets, MCP servers,
+ toolsets built per run, and toolsets supplied by capabilities;
+- tool discovery for MCP and per-run toolsets;
+- ``@durable_operation`` methods of other capabilities, such as the spend accounting of
+ pydantic-ai-harness ``SpendLimits`` and the summarization of ``SummarizingCompaction``,
+ so a retry neither charges the spend twice nor pays for a summary again.
+
+A tool that raises ``ModelRetry``, ``ToolFailed``, ``ApprovalRequired`` or
+``CallDeferred`` replays as the same signal. ``AgentOperator`` does not support tool
+approval or deferred tools with ``durable=True``, so there such a call fails the task
+with ``UnsupportedToolDeferralError``, as it would without ``durable``. A tool that raises
+any other exception runs again on retry.
+
+A task can run more than one agent, and an agent's tool can run another agent; each
+agent run is replayed on its own, and one a tool started is replayed under that tool
+call. The journal is deleted when the task succeeds, including any steps an earlier
+attempt recorded beyond where the last one went, so a task cleared later starts fresh.
+In a plain ``@task``, which has no such moment, the capability deletes a run's journal
+as soon as that run succeeds: after a failure, only the run that failed replays, and
+the agents that ran before it run again.
+
+What the task log shows
+-----------------------
+
+``AgentOperator`` logs a summary at the end of every attempt, including the one that
+fails. Here an agent asked for two tools in one step: ``charge_card`` succeeded and
+``flaky_lookup`` raised, which failed the attempt. The failing attempt logged:
+
+.. code-block:: text
+
+ LLM run failed: requests=1, tool_calls=1, input_tokens=58, output_tokens=9, total_tokens=67
+ Durable: replayed 0 steps (0 model, 0 tool, 0 other), recorded 2 new steps (1 model, 1 tool, 0 other)
+
+The retry replayed the model response and ``charge_card``, and ran only ``flaky_lookup``
+and the final model request:
+
+.. code-block:: text
+
+ Durable: replayed 2 steps (1 model, 1 tool, 0 other), recorded 2 new steps (1 model, 1 tool, 0 other)
+ LLM run complete: model=test, requests=1, tool_calls=1, input_tokens=65, output_tokens=20, total_tokens=85
+
+These lines come from a test agent with two tools, not from the examples above.
+``charge_card`` ran once across both attempts, and the retry's usage counts only its own
+two live calls. A step that ran live but could not be recorded is named in a WARNING
+when it happens; the summary names such tool steps again and counts such model steps.
+Per-step detail is logged at DEBUG.
+
+The agent's ``run_id`` and ``conversation_id``
+----------------------------------------------
+
+Without ``durable``, each attempt's agent run gets a key for the attempt as its
+pydantic-ai ``run_id``: the task instance id on Airflow 3, which changes on every retry. With ``durable=True`` every attempt keeps the **first**
+attempt's id, and so does the run's ``conversation_id`` when there is no
+``message_history``. Capabilities that key their own state on these ids, such as
+pydantic-ai-harness ``SpendLimits`` and ``StepPersistence``, then find the records the
+replayed steps refer to: with a per-attempt id, a retry that replays ``SpendLimits``'
+accounting for a ``run`` window fails with a ``KeyError``. The ``run_id`` XCom, from a
+failed attempt or a successful one, and the ``gen_ai.agent.call.id`` span attribute carry
+the same id; the spans of each attempt still carry that attempt's own
+``airflow.task_instance.id`` and ``try_number``. The id is stored with the journal and
+deleted with it, so a task cleared after it succeeded starts a new one.
+
+Usage and limits
+----------------
+
+A replayed step adds nothing to the usage the run counts against ``usage_limits`` or
+reports in the ``usage`` XCom: not its request, tokens, cost or tool calls, nor the
+usage a replayed ``@durable_operation`` brings back. Each attempt counts only the calls
+it actually makes. This holds after clearing a failed task instance too: the clear
+starts a fresh budget but keeps the journal, and what the rerun replays from it is free.
+
+Side effects and idempotency
+----------------------------
+
+Durable execution records what a tool returned. A replayed tool's code does not run, so a
+side effect it had (a file written, a message sent) is not repeated.
+A tool that fails *before* its result is recorded runs again on retry, so a tool that
+sent an email and then raised sends it again. The same holds when the worker dies
+between a tool's side effect and the write of its result. Make tools with side effects
+idempotent: check whether the work already happened, or let a database constraint
+reject the duplicate.
+
+Tool results are masked by Airflow's secret masker before they are recorded, so a
+connection password in a result reaches the journal as ``***``.
+
+A toolset can opt out of replay. ``ManagedAgentToolset`` does by default
+(``replayable=False``), because a managed agent acts on a system Airflow cannot observe:
+its calls run again on every attempt, while the steps around them still replay.
+
+A result has to survive a JSON round trip to be recorded. One that cannot be encoded
+fails the task with ``DurableJournalError`` and no retry, since it would fail the same
+way every time; the model could not have read it either. A write the task state store
+refuses does not fail the task: the task log names
+the step, and a retry runs it again.
+
+Limitations
+-----------
+
+- ``durable=True`` cannot be combined with ``enable_hitl_review``, a pydantic-ai-harness
+ ``CodeMode`` capability, a ``SandboxToolset`` or tool approval. ``AgentOperator`` raises
+ ``ValueError`` for the first three when the Dag is parsed.
+- A task instance that is retried across an upgrade of this provider from a version
+ without ``AirflowDurability`` runs its steps live once: the journal keys changed.
See also
--------
diff --git a/providers/common/ai/docs/index.rst b/providers/common/ai/docs/index.rst
index 3962b0bd468..1c5c26e0f33 100644
--- a/providers/common/ai/docs/index.rst
+++ b/providers/common/ai/docs/index.rst
@@ -181,7 +181,7 @@ PIP package Version required
``apache-airflow`` ``>=2.11.0``
``apache-airflow-providers-common-compat`` ``>=1.21.0``
``apache-airflow-providers-standard`` ``>=1.20.0``
-``pydantic-ai-slim`` ``>=2.33.0``
+``pydantic-ai-slim`` ``>=2.36.0``
``structlog`` ``>=24.2.0``
========================================== ==================
@@ -220,12 +220,12 @@ Install them when installing from PyPI. For example:
=============== =======================================================================================================================================
Extra Dependencies
=============== =======================================================================================================================================
-``anthropic`` ``pydantic-ai-slim[anthropic]>=2.33.0``, ``anthropic>=1.0.0``
-``bedrock`` ``pydantic-ai-slim[bedrock]>=2.33.0``
-``google`` ``pydantic-ai-slim[google]>=2.33.0``
-``openai`` ``pydantic-ai-slim[openai]>=2.33.0``, ``openai>=2.47.0``
+``anthropic`` ``pydantic-ai-slim[anthropic]>=2.36.0``, ``anthropic>=1.0.0``
+``bedrock`` ``pydantic-ai-slim[bedrock]>=2.36.0``
+``google`` ``pydantic-ai-slim[google]>=2.36.0``
+``openai`` ``pydantic-ai-slim[openai]>=2.36.0``, ``openai>=2.47.0``
``typesafe`` ``typesafe-sdk>=0.6.0``
-``mcp`` ``pydantic-ai-slim[mcp]>=2.33.0``
+``mcp`` ``pydantic-ai-slim[mcp]>=2.36.0``
``modal`` ``apache-airflow-providers-modal``, ``modal>=1.5.2``
``opensandbox`` ``opensandbox>=1.1.0``
``boat`` ``boat-sdk>=1.9.3``
diff --git a/providers/common/ai/docs/operators/agent.rst b/providers/common/ai/docs/operators/agent.rst
index 08f01e4f475..7e9b41d463f 100644
--- a/providers/common/ai/docs/operators/agent.rst
+++ b/providers/common/ai/docs/operators/agent.rst
@@ -238,8 +238,10 @@ file, and ``system_prompt`` is added to the file's ``instructions``:
Build the path from ``__file__``: a relative path resolves against the worker's
working directory, not the Dag file.
-With ``durable=True``, tools from capabilities declared in the spec file are not
-replayed on retry; they run again. Pass tools you need replayed in ``toolsets=``.
+With ``durable=True``, the toolsets that spec-file capabilities contribute are replayed
+like any other, and need an id the same way (see :doc:`../durable_execution`). The
+operator cannot see inside a spec file when the Dag is parsed, so a missing id there
+fails the task when the agent is built rather than the Dag when it loads.
Agent features
--------------
@@ -424,20 +426,22 @@ Parameters
regeneration shares the count of the run before it whenever ``usage_limits``
is set, on every Airflow version; with ``usage_limits=None`` each
regeneration starts a fresh count.
-- ``durable``: When ``True``, enables step-level caching of model responses and
- tool results. On retry, cached steps are replayed instead of re-executing
- expensive LLM calls. On Airflow >= 3.3 the cache uses the task state store (no
- configuration needed); on older Airflow versions it requires the ``[common.ai]
- durable_cache_path`` config option to be set. Default ``False``. A replayed
+- ``durable``: When ``True``, records each step the agent completes (model
+ requests, tool calls, other capabilities' ``@durable_operation`` methods) in a
+ journal kept for the task instance, and replays the recorded steps on retry
+ instead of running them again. Every ``FunctionToolset``, ``MCPToolset`` and
+ ``DynamicToolset`` needs an ``id``. On Airflow >= 3.3 the journal uses the task
+ state store (no configuration needed); on older Airflow versions it requires the
+ ``[common.ai] durable_cache_path`` config option to be set. Default ``False``. A replayed
step adds nothing to the usage counted against ``usage_limits`` or reported
in the ``usage`` XCom -- not its request, tokens, cost, or tool calls -- so
each attempt counts only the model and tool calls it actually makes, on
every Airflow version. A step that re-runs live because the conversation
changed since the previous attempt is counted like any other live call, and
a retry whose cross-attempt total already sits at a limit can still start
- when the steps it needs are cached. Clearing a failed task instance starts
- a fresh budget but keeps the durable cache its attempts left behind, so what
- the rerun replays from that cache is free there too.
+ when the steps it needs are recorded. Clearing a failed task instance starts
+ a fresh budget but keeps the journal its attempts left behind, so what the
+ rerun replays from it is free there too. See :doc:`../durable_execution`.
- ``cache_prompt``: Ask the provider to cache the tool definitions, system prompt and
conversation so later requests read them back at a discount. Default ``True``; a no-op for
providers that cache on their own. See :ref:`agent-prompt-caching`.
diff --git a/providers/common/ai/docs/toolsets/index.rst b/providers/common/ai/docs/toolsets/index.rst
index 34f5d07fb8d..964539dc2fd 100644
--- a/providers/common/ai/docs/toolsets/index.rst
+++ b/providers/common/ai/docs/toolsets/index.rst
@@ -301,10 +301,9 @@ there is no agent budget, so each tool gets one correction unless its toolset se
Layering
--------
-:class:`~airflow.providers.common.ai.toolsets.logging.LoggingToolset` and the
-durable-execution ``CachingToolset`` are not alternatives to anything above. Both
-are wrappers: they take a toolset and return a toolset, adding per-call logging
-or replay from a durable cache. Choose a route first, then decide whether to wrap
+:class:`~airflow.providers.common.ai.toolsets.logging.LoggingToolset` is not an
+alternative to anything above. It is a wrapper: it takes a toolset and returns a
+toolset that logs each call. Choose a route first, then decide whether to wrap
it. The same applies to
:func:`~airflow.providers.common.ai.toolsets.langchain_bridge.airflow_toolset_to_langchain_tools`,
which converts a chosen toolset for a different agent framework rather than
@@ -332,8 +331,9 @@ context.
This approach gives you full control over the agent lifecycle -- you can call
``agent.run_sync()`` multiple times, swap models at runtime, or combine
results from several agents in a single task. The tradeoff is that you lose
-the durable execution (step-level caching with retry replay), HITL review
-integration, and automatic tool call logging that ``AgentOperator`` provides.
+the HITL review integration and automatic tool call logging that
+``AgentOperator`` provides. Durable replay is still available: attach
+``AirflowDurability`` to the agent (see :doc:`../durable_execution`).
Running the agent outside the data system
-----------------------------------------
diff --git a/providers/common/ai/docs/troubleshooting.rst b/providers/common/ai/docs/troubleshooting.rst
index a8c1c3978f1..5e0413c8f47 100644
--- a/providers/common/ai/docs/troubleshooting.rst
+++ b/providers/common/ai/docs/troubleshooting.rst
@@ -104,8 +104,19 @@ as a task failure.
on an older Airflow version. See :doc:`approval_gates` and :doc:`hitl_review`.
``durable=True and enable_hitl_review=True cannot be used together`` / ``durable=True cannot be used with a CodeMode capability``
- Durable replay assumes a stable step order across attempts, which neither a human
- review loop nor code mode provides. Pick one. See :doc:`durable_execution`.
+ Durable replay matches each step to the one at the same position in the previous
+ attempt, which neither a human review loop's regenerations nor the tool calls code
+ mode makes from inside ``run_code`` keep. Pick one. See :doc:`durable_execution`.
+
+``durable=True needs a unique id on every FunctionToolset, such as FunctionToolset(..., id='orders')``
+ Durable replay records each tool call under its toolset's id, so a
+ ``FunctionToolset``, ``MCPToolset`` or ``DynamicToolset`` without one is refused when
+ the Dag loads. Set ``id=`` on the toolset itself; an id on a ``Toolset`` capability
+ around it does not reach it. See :doc:`durable_execution`.
+
+``durable=True attaches AirflowDurability itself; pass durable=True or the capability, not both.``
+ ``AgentOperator(durable=True)`` adds the capability to the agent, so a second one in
+ ``capabilities`` would record every step twice. Drop one of them.
``message_history and enable_hitl_review=True cannot be used together``
The post-review transcript is not recoverable today, so the operator refuses rather
@@ -125,11 +136,16 @@ as a task failure.
Run-time errors
---------------
-``Agent model must be set when durable=True``
+``An agent needs to have a model in order to be used with Airflow, it cannot be set at agent run time.``
The agent was built without a model, usually because the connection has no
``model`` and no ``model_id`` was passed. Durable execution needs the model resolved
- up front so that replayed steps can be matched against their fingerprints. Fix the connection as described
- above.
+ up front so that replayed steps can be matched against their fingerprints. Fix the
+ connection as described above.
+
+``DurableJournalError: durable execution could not record a step's result, because it is not JSON-serializable``
+ A tool returned a value that cannot be encoded as JSON. The model could not have read
+ it either, so the task fails without retrying. Return text, numbers, lists, dicts or a
+ Pydantic model. See :doc:`durable_execution`.
``durable=True`` on Airflow below 3.3 fails with a ``ValueError`` about ``durable_cache_path``
On Airflow versions older than 3.3 the step cache lives in object storage and
diff --git a/providers/common/ai/pyproject.toml b/providers/common/ai/pyproject.toml
index 7078737fbf7..471fe221384 100644
--- a/providers/common/ai/pyproject.toml
+++ b/providers/common/ai/pyproject.toml
@@ -71,8 +71,9 @@ dependencies = [
# 2.33.0 is the first release that works with anthropic>=1: it moved to httpx2 alongside
# the SDK and stopped passing temperature/top_p/top_k as messages.create() kwargs, both
# of which raise TypeError on 2.31 and earlier. The cost API this provider relies on
- # (RunUsage.cost, UsageLimits.cost_limit) landed earlier, in 2.23.0.
- "pydantic-ai-slim>=2.33.0",
+ # (RunUsage.cost, UsageLimits.cost_limit) landed earlier, in 2.23.0. Durable execution
+ # (AirflowDurability) builds on the durable backend API pydantic-ai-slim published in 2.36.0.
+ "pydantic-ai-slim>=2.36.0",
# Airflow 3 brings structlog in through the Task SDK; Airflow 2 does not. 24.2.0 is the first
# release whose render_to_log_kwargs hands ``stacklevel`` to stdlib logging under that name.
"structlog>=24.2.0",
@@ -87,10 +88,10 @@ dependencies = [
# (2.45 and 2.46 raise TypeError on it); pydantic-ai-slim[openai]>=2.33.0 itself asks for
# openai>=3, so this floor only matters where that requirement is overridden, as in this
# repository's lock.
-"anthropic" = ["pydantic-ai-slim[anthropic]>=2.33.0", "anthropic>=1.0.0"]
-"bedrock" = ["pydantic-ai-slim[bedrock]>=2.33.0"]
-"google" = ["pydantic-ai-slim[google]>=2.33.0"]
-"openai" = ["pydantic-ai-slim[openai]>=2.33.0", "openai>=2.47.0"]
+"anthropic" = ["pydantic-ai-slim[anthropic]>=2.36.0", "anthropic>=1.0.0"]
+"bedrock" = ["pydantic-ai-slim[bedrock]>=2.36.0"]
+"google" = ["pydantic-ai-slim[google]>=2.36.0"]
+"openai" = ["pydantic-ai-slim[openai]>=2.36.0", "openai>=2.47.0"]
# TypeSafe Jev: answers typed questions (pick one of N, a bounded score, a yes/no
# probability) and cannot emit text, so it fits a branch or a classification but not a
# summary. It arrives through pydantic-ai's own ``typesafe:`` provider, so the hook needs
@@ -99,7 +100,7 @@ dependencies = [
# this extra only adds the pre-1.0 typesafe-sdk it talks to, which is why it is opt-in. It
# deliberately does not raise the pydantic-ai-slim floor to 2.45; the guide says 2.45+.
"typesafe" = ["typesafe-sdk>=0.6.0"]
-"mcp" = ["pydantic-ai-slim[mcp]>=2.33.0"]
+"mcp" = ["pydantic-ai-slim[mcp]>=2.36.0"]
# Modal-hosted sandboxes for SandboxToolset, the backend to use where the worker host
# cannot run microVMs -- Kubernetes, most notably. Needs no host-level installation,
# only Modal credentials, which the Modal provider's `modal` connection supplies.
@@ -173,7 +174,7 @@ dev = [
"apache-airflow-providers-standard",
# Additional devel dependencies (do not remove this line and add extra development dependencies)
"sqlglot>=30.0.0",
- "pydantic-ai-slim[mcp]>=2.33.0",
+ "pydantic-ai-slim[mcp]>=2.36.0",
"pydantic-ai-skills>=1.2.0",
"apache-airflow-providers-common-sql[datafusion]",
"langchain>=1.0.0",
diff --git a/providers/common/ai/src/airflow/providers/common/ai/durable/__init__.py b/providers/common/ai/src/airflow/providers/common/ai/durable/__init__.py
index 13a83393a91..7d824347e46 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/durable/__init__.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/durable/__init__.py
@@ -14,3 +14,22 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
+"""Durable execution: replay an agent's completed steps when Airflow retries its task."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+ from airflow.providers.common.ai.durable.capability import AirflowDurability
+
+__all__ = ["AirflowDurability"]
+
+
+def __getattr__(name: str) -> object:
+ # Lazy, so the framework-neutral journal can be imported without pydantic-ai's durable API.
+ if name == "AirflowDurability":
+ from airflow.providers.common.ai.durable.capability import AirflowDurability
+
+ return AirflowDurability
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
diff --git a/providers/common/ai/src/airflow/providers/common/ai/durable/base.py b/providers/common/ai/src/airflow/providers/common/ai/durable/base.py
index 99688f84f2a..888293f4ede 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/durable/base.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/durable/base.py
@@ -18,32 +18,32 @@
from __future__ import annotations
-from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
+from collections.abc import Iterable
+from typing import Any, Protocol, runtime_checkable
-if TYPE_CHECKING:
- from pydantic_ai.messages import ModelResponse
-
-# Marks a stored entry as a cached tool result; lets ``load_tool_result``
-# tell a cached ``None`` apart from a missing entry. Single source of truth so
-# the two backends cannot drift on the envelope shape.
-TOOL_RESULT_SENTINEL = "__durable_cached__"
-
-# Prefix for durable cache keys. On the task state store backend (>= 3.3) the
-# cache shares the task instance's key namespace with anything user code writes
+# Prefix for durable journal keys. On the task state store backend (>= 3.3) the
+# journal shares the task instance's key namespace with anything user code writes
# via ``context["task_state_store"]``; the reserved prefix keeps durable steps
# from colliding with user keys. No ``/`` -- task state store keys are a single,
# un-encoded URL path segment.
DURABLE_KEY_PREFIX = "__commonai_durable__"
-def build_model_step_key(step: int) -> str:
- """Build the durable cache key for the model response produced at ``step``."""
- return f"{DURABLE_KEY_PREFIX}model_step_{step}"
+def build_step_key(run: str | int, position: int) -> str:
+ """Build the journal key for the step at ``position`` of the agent run ``run``."""
+ return f"{DURABLE_KEY_PREFIX}run_{run}_step_{position}"
+
+
+def build_run_meta_key(run: str | int) -> str:
+ """Build the journal key for what the agent run ``run`` keeps about itself across attempts."""
+ return f"{DURABLE_KEY_PREFIX}run_{run}_meta"
-def build_tool_step_key(step: int) -> str:
- """Build the durable cache key for the tool result produced at ``step``."""
- return f"{DURABLE_KEY_PREFIX}tool_step_{step}"
+# Holds the id the task's agent run keeps on every attempt; see ``DurableJournal.get_run_id``.
+RUN_ID_KEY = f"{DURABLE_KEY_PREFIX}run_id"
+
+# Lists every agent run any attempt started, so cleanup reaches the ones this attempt did not.
+RUNS_KEY = f"{DURABLE_KEY_PREFIX}runs"
@runtime_checkable
@@ -51,24 +51,21 @@ class DurableStorageProtocol(Protocol):
"""
Persistence contract shared by the durable execution storage backends.
- Implemented by both :class:`~airflow.providers.common.ai.durable.storage.DurableStorage`
+ Implemented by :class:`~airflow.providers.common.ai.durable.storage.DurableStorage`
(ObjectStorage, Airflow < 3.3) and
:class:`~airflow.providers.common.ai.durable.task_state_store.TaskStateStoreDurableStorage`
- (AIP-103 task state store, Airflow >= 3.3). ``CachingModel`` and
- ``CachingToolset`` depend on this interface, not a concrete backend.
+ (AIP-103 task state store, Airflow >= 3.3). The journal depends on this interface,
+ not a concrete backend.
- Both ``save_*`` methods return whether the entry was written. A backend may
- skip a write (a tool result that is not JSON-serializable, a store write that
- fails) without failing the step; the step then re-runs live on retry, and the
- caller counts it as skipped rather than cached.
+ An entry is a JSON-compatible dict. ``save_step`` returns whether the entry was
+ written: a backend may skip a write (a value it cannot store, a store write that
+ fails) without failing the step, which then runs again on retry.
"""
- def save_model_response(self, key: str, response: ModelResponse, *, fingerprint: str | None) -> bool: ...
-
- def load_model_response(self, key: str) -> tuple[ModelResponse | None, str | None]: ...
-
- def save_tool_result(self, key: str, result: Any, *, fingerprint: str | None) -> bool: ...
+ def save_step(self, key: str, entry: dict[str, Any]) -> bool: ...
- def load_tool_result(self, key: str) -> tuple[bool, Any, str | None]: ...
+ def load_step(self, key: str) -> dict[str, Any] | None: ...
- def cleanup(self) -> None: ...
+ def delete_steps(self, keys: Iterable[str]) -> None:
+ """Delete these entries. Best-effort: runs after the work succeeded, so it must not raise."""
+ ...
diff --git a/providers/common/ai/src/airflow/providers/common/ai/durable/caching_model.py b/providers/common/ai/src/airflow/providers/common/ai/durable/caching_model.py
deleted file mode 100644
index 07084d5fdbd..00000000000
--- a/providers/common/ai/src/airflow/providers/common/ai/durable/caching_model.py
+++ /dev/null
@@ -1,226 +0,0 @@
-# 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.
-"""Caching model wrapper for durable execution."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any
-
-from pydantic_ai.messages import ModelResponse, ToolCallPart
-from pydantic_ai.models.wrapper import WrapperModel
-
-from airflow.providers.common.ai.durable.base import build_model_step_key, build_tool_step_key
-from airflow.providers.common.ai.durable.fingerprint import fingerprint_model_request
-from airflow.providers.common.ai.utils.task_logger import get_task_logger
-
-log = get_task_logger()
-
-if TYPE_CHECKING:
- from pydantic_ai.messages import ModelMessage
- from pydantic_ai.models import ModelRequestParameters
- from pydantic_ai.settings import ModelSettings
-
- from airflow.providers.common.ai.durable.base import DurableStorageProtocol
- from airflow.providers.common.ai.durable.replay_usage import ReplayUsageLedger
- from airflow.providers.common.ai.durable.step_counter import DurableStepCounter
-
-
-@dataclass(init=False)
-class CachingModel(WrapperModel):
- """
- Wraps a model to cache responses in ObjectStorage for durable execution.
-
- On each ``request()`` call, checks if a cached response exists for the
- current step index and was produced by an equivalent request (same model,
- message history, settings, and tools -- compared via fingerprint). If so,
- returns the cached response without calling the underlying model.
- Otherwise, calls the model and caches the response. A fingerprint
- mismatch means the agent changed between attempts; the stale entry is
- discarded and the step re-runs live.
-
- With a ``replay_usage`` ledger, a replay hit does not count toward the
- run's usage (see :class:`~airflow.providers.common.ai.durable.replay_usage.ReplayUsageLedger`).
- The cached response is returned unchanged -- it is part of what later
- steps' fingerprints hash, so zeroing its usage would make every step after
- this one diverge and re-run live -- and the ledger subtracts what the graph
- is about to add instead.
- """
-
- storage: DurableStorageProtocol = field(repr=False)
- counter: DurableStepCounter = field(repr=False)
- replay_usage: ReplayUsageLedger | None = field(default=None, repr=False)
-
- def __init__(
- self,
- wrapped: Any,
- *,
- storage: DurableStorageProtocol,
- counter: DurableStepCounter,
- replay_usage: ReplayUsageLedger | None = None,
- ) -> None:
- super().__init__(wrapped)
- self.storage = storage
- self.counter = counter
- self.replay_usage = replay_usage
- self._previous_suspended = False
- self._chain_function_calls = 0
- # Set by _credit_cached_successors when it loads the next model step's response
- # to decide whether to credit it; request() reuses it instead of loading again.
- self._peeked_model: tuple[str, ModelResponse, str | None] | None = None
-
- def credit_first_replay(self) -> None:
- """
- Credit the run's first model request if the cache holds a response for it.
-
- Called before the run starts, because pydantic-ai checks ``request_limit``
- before the first request reaches this model; without the credit, a retry whose
- seeded ``requests`` already equals ``request_limit`` could not start even when
- every step would replay for free.
- """
- if self.replay_usage is None:
- return
- cached, _ = self.storage.load_model_response(build_model_step_key(self.counter.total_steps))
- if cached is not None:
- self.replay_usage.credit_request()
-
- def _is_continuation(self, messages: list[ModelMessage]) -> bool:
- # pydantic-ai re-issues a suspended response (Anthropic ``pause_turn``, OpenAI
- # background mode) by sending it back as the last message; the graph counts
- # the whole chain as one request.
- return (
- self._previous_suspended
- and bool(messages)
- and isinstance(messages[-1], ModelResponse)
- and messages[-1].state == "suspended"
- )
-
- def _track_chain(self, response: ModelResponse, model_request_parameters: ModelRequestParameters) -> None:
- function_tools = {tool.name for tool in model_request_parameters.function_tools}
- self._chain_function_calls += sum(
- 1
- for part in response.parts
- if isinstance(part, ToolCallPart) and part.tool_name in function_tools
- )
- self._previous_suspended = response.state == "suspended"
-
- def _credit_cached_successors(self, ledger: ReplayUsageLedger, step: int) -> None:
- """
- Credit what the replayed, complete response at ``step`` is followed by in the cache.
-
- The step counter is shared with ``CachingToolset``, so this response's
- function-tool calls took the next step indices when they first ran, and the
- next model request took the index after them. Tool indices with a cached result
- are credited; a missing one is a call that raised. The scan stops at the next
- model step, and looks no further than one index per function-tool call in this
- response (and in the earlier segments of its continuation chain) plus one.
-
- The next model step's response, if any, is kept on ``self`` so ``request()`` does
- not load and deserialize it from storage a second time when it reaches that step.
- """
- calls = self._chain_function_calls
- tool_steps: list[int] = []
- for index in range(step + 1, step + calls + 2):
- if index <= step + calls:
- found, _, _ = self.storage.load_tool_result(build_tool_step_key(index))
- if found:
- tool_steps.append(index)
- continue
- next_key = build_model_step_key(index)
- next_model, next_fingerprint = self.storage.load_model_response(next_key)
- if next_model is not None:
- self._peeked_model = (next_key, next_model, next_fingerprint)
- ledger.credit_request()
- break
- ledger.credit_tool_steps(tool_steps, batch_calls=calls)
-
- async def request(
- self,
- messages: list[ModelMessage],
- model_settings: ModelSettings | None,
- model_request_parameters: ModelRequestParameters,
- ) -> ModelResponse:
- step = self.counter.next_step()
- key = build_model_step_key(step)
- continuation = self._is_continuation(messages)
- if not continuation:
- self._chain_function_calls = 0
- had_request_credit = self.replay_usage.settle() if self.replay_usage is not None else False
- # Fingerprint the *prepared* request, not the raw arguments. Concrete
- # models call ``prepare_request()`` at the start of ``request()`` to merge
- # their model-level ``settings`` and apply profile-specific transforms
- # (thinking resolution, native-tool handling, output-mode defaults) before
- # the provider sees the request. Fingerprinting the raw arguments would
- # miss a change that lives only at the model level -- e.g. a different
- # temperature or thinking setting on the connection -- and replay a stale
- # response. The raw arguments are still passed to ``wrapped.request()``,
- # which re-runs ``prepare_request()`` itself (it is pure and idempotent).
- prepared_settings, prepared_parameters = self.wrapped.prepare_request(
- model_settings, model_request_parameters
- )
- fingerprint = fingerprint_model_request(
- f"{self.wrapped.system}:{self.wrapped.model_name}",
- messages,
- prepared_settings,
- prepared_parameters,
- )
-
- cached: ModelResponse | None
- if self._peeked_model is not None and self._peeked_model[0] == key:
- _, cached, cached_fingerprint = self._peeked_model
- self._peeked_model = None
- else:
- self._peeked_model = None
- cached, cached_fingerprint = self.storage.load_model_response(key)
- if cached is not None:
- if cached_fingerprint == fingerprint:
- self.counter.replayed_model += 1
- log.debug("Durable: replayed cached model response", step=step)
- self._track_chain(cached, model_request_parameters)
- if self.replay_usage is not None:
- self.replay_usage.record_model_replay(cached, continuation=continuation)
- if cached.state != "suspended":
- self._credit_cached_successors(self.replay_usage, step)
- return cached
- log.warning(
- "Durable: cached model response does not match the current request; "
- "re-running this step instead of replaying",
- step=step,
- reason=(
- "entry predates fingerprinting or the request could not be fingerprinted"
- if fingerprint is None or cached_fingerprint is None
- else "model, prompt, message history, settings, or tools changed since "
- "the previous attempt"
- ),
- )
-
- if self.replay_usage is not None:
- self.replay_usage.record_live_model_request(had_request_credit=had_request_credit)
- response = await self.wrapped.request(messages, model_settings, model_request_parameters)
- if self.storage.save_model_response(key, response, fingerprint=fingerprint):
- self.counter.cached_model += 1
- log.debug("Durable: cached model response", step=step)
- else:
- self.counter.skipped_model += 1
- # A re-run model step returns fresh tool call ids, so every later step's
- # fingerprint changes and re-runs too.
- log.warning(
- "Durable: model response not cached; a retry re-runs this step and every step after it",
- step=step,
- )
- self._track_chain(response, model_request_parameters)
- return response
diff --git a/providers/common/ai/src/airflow/providers/common/ai/durable/caching_toolset.py b/providers/common/ai/src/airflow/providers/common/ai/durable/caching_toolset.py
deleted file mode 100644
index 1eb30127c01..00000000000
--- a/providers/common/ai/src/airflow/providers/common/ai/durable/caching_toolset.py
+++ /dev/null
@@ -1,141 +0,0 @@
-# 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.
-"""Caching toolset wrapper for durable execution."""
-
-from __future__ import annotations
-
-from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any
-
-from pydantic_ai.toolsets.wrapper import WrapperToolset
-
-from airflow.providers.common.ai.durable.base import build_tool_step_key
-from airflow.providers.common.ai.durable.fingerprint import fingerprint_tool_call
-from airflow.providers.common.ai.utils.task_logger import get_task_logger
-from airflow.providers.common.ai.utils.tool_metrics import record_tool_call
-from airflow.providers.common.ai.utils.toolset_base import AirflowToolset
-
-if TYPE_CHECKING:
- from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool
-
- from airflow.providers.common.ai.durable.base import DurableStorageProtocol
- from airflow.providers.common.ai.durable.replay_usage import ReplayUsageLedger
- from airflow.providers.common.ai.durable.step_counter import DurableStepCounter
-
-log = get_task_logger()
-
-
-@dataclass
-class CachingToolset(WrapperToolset[Any]):
- """
- Wraps a toolset to cache tool call results in ObjectStorage for durable execution.
-
- On each ``call_tool()`` invocation, checks if a cached result exists for
- the current step index and was produced by the same call (same tool name,
- arguments, and model-issued ``tool_call_id`` -- compared via fingerprint).
- If so, returns the cached result without executing the tool. Otherwise,
- executes the tool and caches the result. A fingerprint mismatch means the
- conversation diverged from the previous attempt; the stale entry is
- discarded and the tool runs live.
-
- The step index is grabbed before the first ``await``, so parallel tool
- calls via ``asyncio.gather`` get deterministic indices (tasks start
- executing their synchronous preamble in creation order).
-
- With a ``replay_usage`` ledger, a replayed call does not count toward the
- run's ``tool_calls`` (see
- :class:`~airflow.providers.common.ai.durable.replay_usage.ReplayUsageLedger`).
- """
-
- storage: DurableStorageProtocol = field(repr=False)
- counter: DurableStepCounter = field(repr=False)
- replay_usage: ReplayUsageLedger | None = field(default=None, repr=False)
-
- async def call_tool(
- self,
- name: str,
- tool_args: dict[str, Any],
- ctx: Any,
- tool: ToolsetTool[Any],
- ) -> Any:
- # Grab step index BEFORE any await -- ensures deterministic ordering
- # even when multiple tool calls run concurrently via asyncio.gather.
- step = self.counter.next_step()
-
- # The toolset a tool came from may declare that a completed call must not be served from
- # cache, because the call acted on a system Airflow cannot observe (a managed agent, for
- # instance). ``tool.toolset`` survives every pydantic-ai wrapper, so the check is per tool
- # and one such toolset inside a combined one does not stop its siblings from replaying.
- # The step still counts so later steps keep their keys.
- if not getattr(_innermost(tool.toolset), "replayable", True):
- log.debug("Durable: toolset is not replayable; running the tool", step=step, tool=name)
- if self.replay_usage is not None:
- self.replay_usage.record_live_tool_call(step)
- return await self.wrapped.call_tool(name, tool_args, ctx, tool)
-
- key = build_tool_step_key(step)
- fingerprint = fingerprint_tool_call(name, tool_args, ctx.tool_call_id)
- found, cached, cached_fingerprint = self.storage.load_tool_result(key)
- if found:
- if cached_fingerprint == fingerprint:
- self.counter.replayed_tool += 1
- log.debug("Durable: replayed cached tool result", step=step, tool=name)
- if self.replay_usage is not None:
- self.replay_usage.record_tool_replay(step)
- leaf = _innermost(self.wrapped)
- if not isinstance(leaf, AirflowToolset):
- # Inside a combined or dynamic toolset, the tool knows which one it came from.
- leaf = _innermost(tool.toolset)
- if isinstance(leaf, AirflowToolset):
- record_tool_call(type(leaf).__name__, "replayed")
- return cached
- log.warning(
- "Durable: cached tool result does not match the current tool call; "
- "re-running the tool instead of replaying",
- step=step,
- tool=name,
- reason=(
- "entry predates fingerprinting or the call could not be fingerprinted"
- if fingerprint is None or cached_fingerprint is None
- else "the conversation diverged from the previous attempt"
- ),
- )
-
- if self.replay_usage is not None:
- self.replay_usage.record_live_tool_call(step)
- result = await self.wrapped.call_tool(name, tool_args, ctx, tool)
- if self.storage.save_tool_result(key, result, fingerprint=fingerprint):
- self.counter.cached_tool += 1
- log.debug("Durable: cached tool result", step=step, tool=name)
- else:
- self.counter.skipped_tools.append(name)
- # Named here rather than only in the end-of-run summary: this warning is
- # logged on every path, including the failed attempt that Airflow retries.
- log.warning(
- "Durable: tool result not cached; a retry runs this tool again, "
- "and may re-run the steps after it",
- step=step,
- tool=name,
- )
- return result
-
-
-def _innermost(toolset: AbstractToolset[Any]) -> AbstractToolset[Any]:
- """Return the toolset under any wrappers, such as the masking wrapper AgentOperator adds."""
- while isinstance(toolset, WrapperToolset):
- toolset = toolset.wrapped
- return toolset
diff --git a/providers/common/ai/src/airflow/providers/common/ai/durable/capability.py b/providers/common/ai/src/airflow/providers/common/ai/durable/capability.py
new file mode 100644
index 00000000000..3a420b17adf
--- /dev/null
+++ b/providers/common/ai/src/airflow/providers/common/ai/durable/capability.py
@@ -0,0 +1,532 @@
+# 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.
+"""Durable execution for pydantic-ai agents in Airflow tasks, on pydantic-ai's durable backend API."""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable, Iterator, Mapping
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass, field
+from functools import partial
+from typing import TYPE_CHECKING, Any, ClassVar, TypeAlias, cast
+
+from pydantic_ai.durable_exec import (
+ JSON_CODEC,
+ BaseDurabilityCapability,
+ CapabilityOperationId,
+ DurabilityEngineSpec,
+ DurableOperationId,
+ JournalCallableOperationBackend,
+ ModelRequestId,
+ RoleBasedOperationConfig,
+ ToolsetCallToolId,
+ ToolsetValidateToolArgumentsId,
+)
+from pydantic_ai.exceptions import ApprovalRequired, CallDeferred, ModelRetry, ToolFailed
+from pydantic_ai.messages import ModelResponse, ToolReturn
+from pydantic_ai.tools import AgentDepsT
+from pydantic_ai.toolsets import DynamicToolset, FunctionToolset
+from pydantic_ai.toolsets.wrapper import WrapperToolset
+from pydantic_core import PydanticSerializationError
+
+from airflow.providers.common.ai.durable.fingerprint import fingerprint_model_request, fingerprint_tool_call
+from airflow.providers.common.ai.durable.journal import current_journal
+from airflow.providers.common.ai.durable.replay_usage import ReplayUsageLedger, is_successful_tool_payload
+from airflow.providers.common.ai.exceptions import DurableJournalError
+from airflow.providers.common.ai.utils.masking import mask_secrets
+from airflow.providers.common.ai.utils.tool_metrics import record_tool_call
+from airflow.providers.common.ai.utils.toolset_base import AirflowToolset
+
+if TYPE_CHECKING:
+ from pydantic_ai.agent import EventStreamHandler
+ from pydantic_ai.capabilities import WrapRunHandler
+
+ # The operation a backend binds is typed only in these modules; ``bind`` itself is public.
+ from pydantic_ai.durable_exec._operation import DurableOperation, ParamsT, ResultT, WireT
+ from pydantic_ai.durable_exec._operation_backend import BoundDurableOperation
+ from pydantic_ai.messages import ModelMessage
+ from pydantic_ai.models import Model, ModelRequestParameters
+ from pydantic_ai.run import AgentRunResult
+ from pydantic_ai.settings import ModelSettings
+ from pydantic_ai.tools import RunContext
+ from pydantic_ai.toolsets import AbstractToolset, ToolsetTool
+
+ from airflow.providers.common.ai.durable.journal import DurableRun, JournalStep, StepKind
+
+__all__ = ["AirflowDurability"]
+
+# What pydantic-ai's durable backend passes as ``cache_key``: a projection of the
+# operation's parameters that it types as an opaque tuple. These spell out the shapes
+# of the two projections fingerprinted here.
+_ModelRequestKey: TypeAlias = (
+ "tuple[str | None, list[ModelMessage], ModelSettings | None, ModelRequestParameters, RunContext[Any]]"
+)
+_ToolCallKey: TypeAlias = "tuple[str, dict[str, Any], RunContext[Any]]"
+
+try:
+ from pydantic_ai.mcp import MCPToolset
+except ImportError: # the ``mcp`` extra is not installed
+ DURABLE_UNIT_TOOLSETS: tuple[type, ...] = (FunctionToolset, DynamicToolset)
+else:
+ DURABLE_UNIT_TOOLSETS = (FunctionToolset, DynamicToolset, MCPToolset)
+"""Leaf toolsets pydantic-ai's durable backend runs as durable units; each needs a unique ``id``."""
+
+# Why a recorded step that matched is run live anyway: it was written by a version whose
+# types have changed since, so replaying it would fail on every retry.
+_UNLOADABLE = "the recorded result no longer loads, so it was written by another version"
+
+_NO_CONFIG: RoleBasedOperationConfig[None] = RoleBasedOperationConfig(
+ model=None, event=None, capability=None, tool=None
+)
+
+
+@dataclass(frozen=True)
+class _ActiveRun:
+ """The journal run an agent's steps go to, and the ledger that keeps their replays out of its usage."""
+
+ durable_run: DurableRun
+ ledger: ReplayUsageLedger
+
+
+_ACTIVE_RUN: ContextVar[_ActiveRun | None] = ContextVar("commonai_durability_run", default=None)
+
+
+@contextmanager
+def _active_run(active: _ActiveRun) -> Iterator[None]:
+ token = _ACTIVE_RUN.set(active)
+ try:
+ yield
+ finally:
+ _ACTIVE_RUN.reset(token)
+
+
+def _not_journalable(error: Exception) -> BaseException:
+ # The same value fails the same way on every retry, so retrying the task cannot help.
+ return DurableJournalError(
+ f"durable execution could not record a step's result, because it is not JSON-serializable: {error}"
+ )
+
+
+class AirflowDurability(BaseDurabilityCapability[AgentDepsT]):
+ """
+ Replay an agent's completed steps when Airflow retries the task it runs in.
+
+ Attach it to a pydantic-ai ``Agent`` running inside an Airflow task. Each model
+ request, tool call (function, MCP and dynamic toolsets, and Airflow's own toolsets),
+ tool discovery and ``@durable_operation`` of another capability
+ is recorded in the task instance's durable journal as it completes. When the task
+ fails and Airflow retries it, the agent runs again from the start and every step the
+ previous attempt completed is replayed from the journal instead of being run again:
+ no second model call, no second tool side effect, no second charge against a spend
+ limit. See :ref:`durable-execution` for how replay is verified.
+
+ ``AgentOperator(durable=True)`` and ``@task.agent(durable=True)`` attach it for you.
+ In a plain ``@task``, attach it yourself:
+
+ .. code-block:: python
+
+ @task(retries=3)
+ def research(question: str) -> str:
+ agent = Agent(
+ "anthropic:claude-sonnet-4-5",
+ name="researcher",
+ toolsets=[SQLToolset("warehouse")],
+ capabilities=[AirflowDurability()],
+ )
+ return agent.run_sync(question).output
+
+ Outside a running Airflow task the capability does nothing, so the same agent runs
+ normally in a test or a notebook. In a plain ``@task``, a run's steps are deleted as
+ soon as the run succeeds, so that clearing the task later starts it fresh: a task that
+ runs several agents replays only the run that failed, and runs the ones before it again.
+
+ Toolsets the agent calls through pydantic-ai's durable backend need a unique ``id``
+ (``FunctionToolset(id=...)``, ``MCPToolset(..., id=...)``), as do capabilities that
+ contribute ``@durable_operation`` methods; the ids name the steps in the journal.
+
+ :param models: Extra models the run may switch to, keyed by id; see pydantic-ai's
+ ``BaseDurabilityCapability``.
+ :param event_stream_handler: Optional handler for the run's events.
+ :param name: Prefix for the names of the agent's steps. Defaults to the agent's
+ ``name``; one of the two is required.
+ """
+
+ engine_spec: ClassVar[DurabilityEngineSpec] = DurabilityEngineSpec(
+ engine_name="Airflow",
+ durable_unit_noun="step",
+ durable_container_noun="task",
+ codec=JSON_CODEC,
+ serialization_failure=_not_journalable,
+ )
+
+ def __init__(
+ self,
+ *,
+ models: Mapping[str, Model] | None = None,
+ event_stream_handler: EventStreamHandler[AgentDepsT] | None = None,
+ name: str | None = None,
+ ) -> None:
+ super().__init__(models=models, event_stream_handler=event_stream_handler, name=name)
+ # Keyed by agent too: pydantic-ai binds a shallow copy of this capability to each
+ # agent it is attached to, and the copies share this dict.
+ self._fingerprint_models: dict[tuple[int, str | None], Model] = {}
+
+ @property
+ def in_durable_context(self) -> bool:
+ return current_journal() is not None
+
+ def get_durable_operation_backend(self) -> _AirflowOperationBackend:
+ return _AirflowOperationBackend(self)
+
+ def get_wrapper_toolset(self, toolset: AbstractToolset[AgentDepsT]) -> AbstractToolset[AgentDepsT] | None:
+ """Journal the leaf toolsets pydantic-ai's backend does not, then let the base wrap the rest."""
+ journaled = toolset.visit_and_replace(self._journal_other_leaf)
+ return super().get_wrapper_toolset(journaled) or journaled
+
+ def _journal_other_leaf(self, toolset: AbstractToolset[AgentDepsT]) -> AbstractToolset[AgentDepsT]:
+ # Function, dynamic and MCP toolsets are durable units of pydantic-ai's backend.
+ # Anything else, such as Airflow's own SQL, hook and managed-agent toolsets, is
+ # journaled here instead, or its calls would run again on every retry.
+ if isinstance(toolset, DURABLE_UNIT_TOOLSETS):
+ return toolset
+ return _JournaledToolset(wrapped=toolset, durability=self)
+
+ async def before_run(self, ctx: RunContext[AgentDepsT]) -> None:
+ # The base refuses a run with a ``cancellation_token`` inside the durable container,
+ # because for its engines the container runs elsewhere and cancelling it out of band
+ # would break replay. Airflow's container is the task process itself: AgentOperator's
+ # on_kill cancels the run there, the attempt fails, and the retry replays what the
+ # journal recorded. That check is all the base hook does, so it is not called.
+ return
+
+ async def wrap_run(self, ctx: RunContext[AgentDepsT], *, handler: WrapRunHandler) -> AgentRunResult[Any]:
+ journal = current_journal()
+ if journal is None:
+ return await super().wrap_run(ctx, handler=handler)
+ durable_run = journal.start_run()
+ # Keeps replays out of the usage the run counts and is limited by, which is the
+ # cross-attempt total when AgentOperator passes one as ``usage=``.
+ ledger = ReplayUsageLedger(run_usage=ctx.usage, usage_limits=ctx.usage_limits)
+ ledger.credit_first_replay(durable_run)
+ try:
+ with _active_run(_ActiveRun(durable_run, ledger)):
+ result = await super().wrap_run(ctx, handler=handler)
+ except BaseException as error:
+ durable_run.fail(error)
+ raise
+ finally:
+ ledger.settle()
+ if journal.clean_up_after_run:
+ durable_run.cleanup()
+ return result
+
+ async def _model_for_fingerprint(
+ self, model_id: str | None, run_context: RunContext[AgentDepsT]
+ ) -> Model:
+ """Return the model a request with ``model_id`` goes to, resolved once per agent and id."""
+ key = (id(self.agent), model_id)
+ if (model := self._fingerprint_models.get(key)) is None:
+ model = self._fingerprint_models[key] = await self._resolve_model_for_request(
+ model_id, run_context
+ )
+ return model
+
+
+class _AirflowOperationBackend(JournalCallableOperationBackend[None]):
+ """Runs each durable operation through the task's durable journal."""
+
+ def __init__(self, durability: AirflowDurability[Any]) -> None:
+ super().__init__(
+ agent_name=durability.name, default_model_id=durability.default_model_id, config=_NO_CONFIG
+ )
+ self._durability = durability
+ # How pydantic-ai decodes each operation's result, by operation.
+ self._loaders: dict[DurableOperationId, Callable[[object], object]] = {}
+
+ def bind(
+ self, operation: DurableOperation[ParamsT, WireT, ResultT]
+ ) -> BoundDurableOperation[ParamsT, WireT, ResultT]:
+ # pydantic-ai decodes the payload ``execute`` returns only after ``execute`` returns, when
+ # it is too late to run the step live instead. Keeping the decoder lets ``execute`` reject
+ # a recorded payload that no longer decodes, such as one another version wrote.
+ self._loaders[operation.operation_id] = operation.result_codec.load
+ return super().bind(operation)
+
+ def _loads(self, operation_id: DurableOperationId, payload: object) -> bool:
+ try:
+ self._loaders[operation_id](payload)
+ except (TypeError, ValueError):
+ return False
+ return True
+
+ def _claim(
+ self,
+ active: _ActiveRun,
+ operation_id: DurableOperationId,
+ name: str,
+ *,
+ kind: StepKind,
+ fingerprint: str | None,
+ ) -> JournalStep:
+ step = active.durable_run.claim(name, kind=kind, fingerprint=fingerprint)
+ if step.replayed and not self._loads(operation_id, step.payload):
+ active.durable_run.reject(step, reason=_UNLOADABLE)
+ return step
+
+ async def execute(
+ self,
+ *,
+ operation_id: DurableOperationId,
+ name: str,
+ body: Callable[[], Awaitable[object]],
+ cache_key: tuple[object, ...],
+ config: None,
+ ) -> object:
+ active = _ACTIVE_RUN.get()
+ if active is None:
+ return await body()
+ match operation_id:
+ case ModelRequestId(streaming=streaming):
+ return await self._model_request(
+ active, operation_id, name, body, cache_key, streaming=streaming
+ )
+ case ToolsetCallToolId():
+ step = _claim_tool_call(
+ active,
+ name,
+ _tool_fingerprint(cache_key),
+ valid_payload=partial(self._loads, operation_id),
+ )
+ if step.replayed:
+ return step.payload
+ # The masking wrapper AgentOperator adds sits outside this durable unit, so mask
+ # before recording: secrets must not reach the journal any more than the model.
+ return await step.run(body, to_record=mask_secrets)
+ case ToolsetValidateToolArgumentsId():
+ # Local validation of the arguments the recorded call was made with: nothing to
+ # save by replaying it, and it is the same on every attempt, so it takes no position.
+ return await body()
+ case CapabilityOperationId():
+ step = self._claim(active, operation_id, name, kind="other", fingerprint=None)
+ if step.replayed:
+ active.ledger.record_capability_replay(step.payload)
+ return step.payload
+ return await step.run(body)
+ case _:
+ # Tool discovery, compaction, event handling, and the operations later
+ # pydantic-ai versions add: replayed by name and position.
+ step = self._claim(active, operation_id, name, kind="other", fingerprint=None)
+ if step.replayed:
+ return step.payload
+ return await step.run(body)
+
+ async def _model_request(
+ self,
+ active: _ActiveRun,
+ operation_id: ModelRequestId,
+ name: str,
+ body: Callable[[], Awaitable[object]],
+ cache_key: tuple[object, ...],
+ *,
+ streaming: bool,
+ ) -> object:
+ model_id, messages, model_settings, parameters, run_context = cast("_ModelRequestKey", cache_key)
+ fingerprint = await self._fingerprint_model_request(
+ model_id, messages, model_settings, parameters, run_context
+ )
+ ledger = active.ledger
+ continuation = ledger.begin_model_request(messages)
+ had_request_credit = ledger.settle()
+ step = self._claim(active, operation_id, name, kind="model", fingerprint=fingerprint)
+ if step.replayed:
+ response = _load_model_response(step.payload, streaming=streaming)
+ ledger.record_model_replay(response, continuation=continuation)
+ ledger.track_chain(response, parameters)
+ if response.state != "suspended":
+ ledger.credit_successors(active.durable_run, step.position)
+ return step.payload
+ ledger.record_live_model_request(had_request_credit=had_request_credit)
+ payload = await step.run(body)
+ ledger.track_chain(_load_model_response(payload, streaming=streaming), parameters)
+ return payload
+
+ async def _fingerprint_model_request(
+ self,
+ model_id: str | None,
+ messages: list[ModelMessage],
+ model_settings: ModelSettings | None,
+ parameters: ModelRequestParameters,
+ run_context: RunContext[Any],
+ ) -> str | None:
+ # Fingerprint the request as the model will prepare it, not the raw arguments.
+ # ``prepare_request`` merges the model's own settings and applies profile
+ # transforms (thinking, native tools, output mode) before the provider sees the
+ # request, so a change that lives only on the connection, such as a different
+ # temperature, still invalidates the recorded response. It is pure, so calling it
+ # here as well as in the request itself is safe.
+ model = await self._durability._model_for_fingerprint(model_id, run_context)
+ prepared_settings, prepared_parameters = model.prepare_request(model_settings, parameters)
+ return fingerprint_model_request(
+ f"{model.system}:{model.model_name}", messages, prepared_settings, prepared_parameters
+ )
+
+
+@dataclass
+class _JournaledToolset(WrapperToolset[Any]):
+ """
+ Journals the calls of a leaf toolset that pydantic-ai's durable backend does not wrap.
+
+ Results, and the control-flow exceptions a tool raises for the model (``ModelRetry``,
+ ``ToolFailed``, ``ApprovalRequired``, ``CallDeferred``), are recorded as values, so a
+ retry replays them exactly. Any other exception is recorded as a failure and the call
+ runs again on retry. A result that is not JSON-serializable fails the task, as it does
+ for the toolsets pydantic-ai runs as durable steps.
+ """
+
+ durability: AirflowDurability[Any] = field(repr=False, kw_only=True)
+
+ def visit_and_replace(
+ self, visitor: Callable[[AbstractToolset[Any]], AbstractToolset[Any]]
+ ) -> AbstractToolset[Any]:
+ # A durable unit, like pydantic-ai's own durable toolsets: a later visit, such as the
+ # journaling pass of the next run, must not reach the leaf and wrap it again.
+ return self
+
+ async def call_tool(
+ self, name: str, tool_args: dict[str, Any], ctx: RunContext[Any], tool: ToolsetTool[Any]
+ ) -> Any:
+ active = _ACTIVE_RUN.get()
+ if active is None:
+ return await self.wrapped.call_tool(name, tool_args, ctx, tool)
+ leaf = self.wrapped
+ step = _claim_tool_call(
+ active,
+ f"{self.durability.name}__airflow_toolset__{leaf.id or type(leaf).__name__}.call_tool:{name}",
+ fingerprint_tool_call(name, tool_args, ctx.tool_call_id),
+ valid_payload=_is_journaled_tool_payload,
+ # A toolset whose calls act on a system Airflow cannot observe, such as a managed
+ # agent, runs them again on every attempt.
+ replayable=leaf.replayable if isinstance(leaf, AirflowToolset) else True,
+ )
+ if step.replayed:
+ if isinstance(leaf, AirflowToolset):
+ record_tool_call(type(leaf).__name__, "replayed")
+ return _decode_tool_payload(step.payload)
+ return await _run_and_record_tool(step, lambda: leaf.call_tool(name, tool_args, ctx, tool))
+
+
+def _claim_tool_call(
+ active: _ActiveRun,
+ name: str,
+ fingerprint: str | None,
+ *,
+ valid_payload: Callable[[Any], bool],
+ replayable: bool = True,
+) -> JournalStep:
+ """Claim a tool call's step and keep the ledger's count of tool calls in step with it."""
+ step = active.durable_run.claim(name, kind="tool", fingerprint=fingerprint, replayable=replayable)
+ if step.replayed and not valid_payload(step.payload):
+ active.durable_run.reject(step, reason=_UNLOADABLE)
+ if not step.replayed:
+ active.ledger.record_live_tool_call(step.position)
+ elif is_successful_tool_payload(step.payload):
+ active.ledger.record_tool_replay(step.position)
+ return step
+
+
+async def _run_and_record_tool(step: JournalStep, call: Callable[[], Awaitable[Any]]) -> Any:
+ try:
+ with step.executing():
+ result = await call()
+ except ModelRetry as e:
+ _record_tool_payload(step, {"kind": "model_retry", "message": e.message})
+ raise
+ except ToolFailed as e:
+ _record_tool_payload(step, {"kind": "tool_failed", "message": e.message})
+ raise
+ except ApprovalRequired as e:
+ _record_tool_payload(step, {"kind": "approval_required", "metadata": e.metadata})
+ raise
+ except CallDeferred as e:
+ _record_tool_payload(step, {"kind": "call_deferred", "metadata": e.metadata})
+ raise
+ except Exception as error:
+ step.fail(error)
+ raise
+ key = "tool_return" if isinstance(result, ToolReturn) else "result"
+ _record_tool_payload(step, {"kind": "tool_return", key: result})
+ return result
+
+
+def _record_tool_payload(step: JournalStep, payload: dict[str, Any]) -> None:
+ try:
+ encoded = JSON_CODEC.dump(Any, payload)
+ except (PydanticSerializationError, TypeError, ValueError) as error:
+ # As for the toolsets pydantic-ai runs as durable steps: the same value fails the same
+ # way on every attempt, and the model could not have read it either.
+ raise _not_journalable(error) from error
+ step.record(mask_secrets(encoded))
+
+
+def _decode_tool_payload(payload: Any) -> Any:
+ kind = payload.get("kind") if isinstance(payload, dict) else None
+ match kind:
+ case "tool_return" if "tool_return" in payload:
+ return JSON_CODEC.load(ToolReturn, payload["tool_return"])
+ case "tool_return":
+ return payload["result"]
+ case "model_retry":
+ raise ModelRetry(payload["message"])
+ case "tool_failed":
+ raise ToolFailed(payload["message"])
+ case "approval_required":
+ raise ApprovalRequired(metadata=payload.get("metadata"))
+ case "call_deferred":
+ raise CallDeferred(metadata=payload.get("metadata"))
+ raise DurableJournalError(f"durable execution found a tool result it cannot replay: {kind!r}")
+
+
+def _load_model_response(payload: object, *, streaming: bool) -> ModelResponse:
+ """Decode a model response payload; a streamed request records it with the events it streamed."""
+ raw = payload.get("response") if streaming and isinstance(payload, dict) else payload
+ return JSON_CODEC.load(ModelResponse, raw)
+
+
+def _is_journaled_tool_payload(payload: object) -> bool:
+ if not isinstance(payload, dict):
+ return False
+ match payload.get("kind"):
+ case "tool_return" if "tool_return" in payload:
+ try:
+ JSON_CODEC.load(ToolReturn, payload["tool_return"])
+ except (TypeError, ValueError):
+ return False
+ return True
+ case "tool_return":
+ return "result" in payload
+ case "model_retry" | "tool_failed":
+ return isinstance(payload.get("message"), str)
+ case "approval_required" | "call_deferred":
+ return True
+ return False
+
+
+def _tool_fingerprint(cache_key: tuple[object, ...]) -> str | None:
+ # Function and MCP calls project (name, args, ctx, tool); dynamic ones (name, args, ctx, tool_def).
+ tool_name, tool_args, ctx = cast("_ToolCallKey", cache_key[:3])
+ return fingerprint_tool_call(tool_name, tool_args, ctx.tool_call_id)
diff --git a/providers/common/ai/src/airflow/providers/common/ai/durable/fingerprint.py b/providers/common/ai/src/airflow/providers/common/ai/durable/fingerprint.py
index ad985e2e4dc..a7462216d99 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/durable/fingerprint.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/durable/fingerprint.py
@@ -35,7 +35,8 @@ Fields that pydantic-ai regenerates on every attempt (message-level
``timestamp``/``run_id``/``conversation_id``, part-level ``timestamp``, the
``capability_id`` of a tool from a capability without an ``id``) are excluded
from the fingerprint, and set-valued request
-parameters are sorted or dropped. Requests that cannot be serialized to JSON
+parameters are sorted or dropped. Registered secrets in the message history are
+masked before hashing, as they are in the tool results the journal records. Requests that cannot be serialized to JSON
fingerprint as ``None``, which degrades that step to unverified positional
replay (the pre-fingerprint behavior) rather than disabling caching.
"""
@@ -50,6 +51,7 @@ from pydantic import TypeAdapter
from pydantic_ai.messages import ModelMessagesTypeAdapter
from pydantic_ai.models import ModelRequestParameters
+from airflow.providers.common.ai.utils.masking import mask_secrets
from airflow.providers.common.ai.utils.prompt_cache import PROMPT_CACHE_SETTING_NAMES
from airflow.providers.common.ai.utils.task_logger import get_task_logger
@@ -154,7 +156,10 @@ def fingerprint_model_request(
return _digest(
{
"model": model_identifier,
- "messages": _strip_volatile(dumped),
+ # The journal records tool results with secrets masked, while a live run hands
+ # the model the unmasked result, so a replayed history holds ``***`` where the
+ # live one held the secret. Masking the history here as well makes both match.
+ "messages": _strip_volatile(mask_secrets(dumped)),
"settings": _content_settings(model_settings),
"params": _normalize_params(params),
}
diff --git a/providers/common/ai/src/airflow/providers/common/ai/durable/journal.py b/providers/common/ai/src/airflow/providers/common/ai/durable/journal.py
new file mode 100644
index 00000000000..4683634c1cd
--- /dev/null
+++ b/providers/common/ai/src/airflow/providers/common/ai/durable/journal.py
@@ -0,0 +1,551 @@
+# 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.
+"""
+The durable journal: replay an agent's completed steps when Airflow retries its task.
+
+An Airflow task retry is a fresh process that runs the agent again from the top. The
+journal records each step an agent run completes (a model response, a tool result, a
+capability operation) under the step's position in that run, and on the retry hands
+back the recorded result instead of running the step again. Steps the previous
+attempt never reached run live and are recorded in turn.
+
+A step is replayed only when the step at the same position in the previous attempt
+had the same name and the same fingerprint. The name says what kind of step it was
+(``agent__model.request``, ``agent__function_toolset__x.call_tool:query``); the
+fingerprint, computed by the framework adapter, says what it was asked (the message
+history, the tool arguments). The first mismatch means the run took a different path
+from the previous attempt, so that step and every step after it run live. So does
+everything the previous attempt did after a step that raised.
+
+A :class:`DurableJournal` belongs to one task attempt. Each agent run in it is a
+:class:`DurableRun`, numbered in the order the runs start, with positions of its own,
+so a task that runs several agents replays each of them independently.
+
+Nothing here depends on an agent framework. Payloads are JSON-compatible values and
+fingerprints are opaque strings.
+:class:`~airflow.providers.common.ai.durable.capability.AirflowDurability` drives the
+journal for pydantic-ai. Another framework's adapter starts a run with
+:meth:`DurableJournal.start_run` and calls one of two shapes of API on it:
+
+* :meth:`DurableRun.run` wraps a step:
+ ``await run.run(name, kind=..., fingerprint=..., body=...)``.
+* :meth:`DurableRun.claim` and :meth:`JournalStep.record` split it in two, for a
+ framework whose hooks see a step's input and its output in separate callbacks.
+"""
+
+from __future__ import annotations
+
+import collections
+from collections.abc import Awaitable, Callable, Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Literal
+
+from airflow.providers.common.ai.durable.base import RUN_ID_KEY, RUNS_KEY, build_run_meta_key, build_step_key
+from airflow.providers.common.ai.durable.storage import DurableStorage
+from airflow.providers.common.ai.observability import make_task_instance_run_key
+from airflow.providers.common.ai.utils.task_logger import get_task_logger
+from airflow.providers.common.compat.sdk import AirflowException, get_current_context
+from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_3_PLUS
+
+if TYPE_CHECKING:
+ import logging
+
+ from airflow.providers.common.ai.durable.base import DurableStorageProtocol
+ from airflow.sdk import Context
+
+log = get_task_logger()
+
+StepKind = Literal["model", "tool", "other"]
+"""What a step is, for the end-of-run summary and for adapters that treat model steps specially."""
+
+
+@dataclass
+class JournalStep:
+ """
+ One position in a run, claimed by :meth:`DurableRun.claim`.
+
+ When :attr:`replayed` is true, :attr:`payload` holds what the previous attempt
+ recorded and the caller returns it instead of running the step. Otherwise the
+ caller runs the step inside :meth:`executing` and passes its result to
+ :meth:`record`, or the exception it raised to :meth:`fail`; :meth:`run` does all three.
+ """
+
+ durable_run: DurableRun = field(repr=False)
+ position: int
+ name: str
+ kind: StepKind
+ fingerprint: str | None
+ replayable: bool = True
+ replayed: bool = False
+ payload: Any = None
+ # Agent runs started while this step executes; they are numbered under it.
+ nested_runs: int = field(default=0, repr=False)
+
+ async def run(
+ self, body: Callable[[], Awaitable[Any]], *, to_record: Callable[[Any], Any] | None = None
+ ) -> Any:
+ """
+ Run a step that was not replayed and record what ``body`` returns, or that it raised.
+
+ :param to_record: Transforms the live result before it is recorded, such as
+ masking secrets in a tool result. The live result is returned unchanged.
+ """
+ try:
+ with self.executing():
+ payload = await body()
+ except Exception as error:
+ self.fail(error)
+ raise
+ self.record(payload if to_record is None else to_record(payload))
+ return payload
+
+ @contextmanager
+ def executing(self) -> Iterator[None]:
+ """
+ Mark the code in the ``with`` block as this step's work.
+
+ An agent run that starts inside the block, such as an agent a tool calls, is
+ numbered under this step, so the runs of later steps keep their numbers on a
+ retry that replays this step instead of running it.
+ """
+ token = _EXECUTING_STEP.set(self)
+ try:
+ yield
+ finally:
+ _EXECUTING_STEP.reset(token)
+
+ def record(self, payload: Any) -> bool:
+ """
+ Record the step's result so a retry can replay it.
+
+ :param payload: A JSON-compatible value.
+ :return: Whether it was recorded. A step that was not recorded runs again on retry,
+ and one claimed with ``replayable=False`` is never recorded.
+ """
+ if not self.replayable:
+ return False
+ entry = {"name": self.name, "kind": self.kind, "fingerprint": self.fingerprint, "payload": payload}
+ if self.durable_run.save(self.position, entry):
+ self.durable_run.journal.stats.recorded[self.kind] += 1
+ log.debug("Durable: recorded step", position=self.position, step=self.name)
+ return True
+ self.durable_run.journal.stats.not_recorded.append((self.kind, self.name))
+ # Named here, not only in the summary: this line is logged on every path,
+ # including the failed attempt that Airflow retries.
+ log.warning(
+ "Durable: step not recorded; a retry runs it again, and may run the steps after it again",
+ position=self.position,
+ step=self.name,
+ )
+ return False
+
+ def fail(self, error: BaseException) -> None:
+ """
+ Note that the step raised ``error``.
+
+ Nothing is recorded for the step, so it runs again on retry. If ``error`` goes on to
+ fail the run, :meth:`DurableRun.fail` uses this to tell which steps the run had
+ already started when it failed; a run that recovers from the error is unaffected.
+ """
+ self.durable_run._raised.append((error, self.durable_run.position))
+
+
+@dataclass
+class JournalStats:
+ """What one task attempt replayed and recorded, by step kind."""
+
+ replayed: collections.Counter[StepKind] = field(default_factory=collections.Counter)
+ recorded: collections.Counter[StepKind] = field(default_factory=collections.Counter)
+ # Steps that ran live and could not be recorded, so a retry runs them again.
+ not_recorded: list[tuple[StepKind, str]] = field(default_factory=list)
+
+
+class DurableRun:
+ """
+ One agent run's steps in the journal.
+
+ Positions are handed out in the order steps are claimed. A step must be claimed
+ before its first ``await``, so concurrent steps (parallel tool calls) take their
+ positions in the order they were started rather than the order they finish.
+
+ Besides its steps, a run keeps one entry about itself: where the last attempt that
+ failed stopped being trustworthy, and the furthest position any attempt reached, so
+ a successful run deletes every step an earlier attempt left behind.
+ """
+
+ def __init__(self, journal: DurableJournal, key: str) -> None:
+ self.journal = journal
+ self.key = key
+ self.diverged = False
+ self._position = 0
+ # Read on first use: where the last failed attempt's entries stop being trusted,
+ # and how far any attempt got.
+ self._diverge_from: int | None = None
+ self._high_water = 0
+ self._meta_loaded = False
+ # Whether the step claimed last replayed; a step without a fingerprint replays only
+ # if it did, so it is never matched against a different preceding conversation.
+ self._previous_replayed = True
+ # Entries read ahead by ``peek`` and not yet claimed.
+ self._peeked: dict[int, dict[str, Any] | None] = {}
+ # Exceptions steps raised this attempt, with the next position when each was raised.
+ self._raised: list[tuple[BaseException, int]] = []
+ # The furthest position this attempt recorded a step at.
+ self._last_saved = -1
+
+ @property
+ def position(self) -> int:
+ """The position the next claimed step will take."""
+ return self._position
+
+ def _load_meta(self) -> None:
+ if self._meta_loaded:
+ return
+ self._meta_loaded = True
+ meta = self.journal.storage.load_step(build_run_meta_key(self.key)) or {}
+ diverge_from, high_water = meta.get("diverge_from"), meta.get("high_water")
+ self._diverge_from = diverge_from if isinstance(diverge_from, int) else None
+ self._high_water = high_water if isinstance(high_water, int) else 0
+
+ def peek(self, position: int) -> dict[str, Any] | None:
+ """
+ Return the entry the previous attempt recorded at ``position``, without claiming it.
+
+ ``None`` when nothing usable was recorded there, or the run has already diverged.
+ The entry is a dict with ``name``, ``kind``, ``fingerprint`` and ``payload``.
+ """
+ self._load_meta()
+ if self.diverged or (self._diverge_from is not None and position >= self._diverge_from):
+ return None
+ if position not in self._peeked:
+ self._peeked[position] = self.journal.storage.load_step(build_step_key(self.key, position))
+ return self._peeked[position]
+
+ def claim(
+ self, name: str, *, kind: StepKind, fingerprint: str | None, replayable: bool = True
+ ) -> JournalStep:
+ """
+ Take the next position and look up what the previous attempt recorded there.
+
+ :param name: What kind of step this is. Must be the same on every attempt for the
+ same step, and should differ between steps that are not interchangeable.
+ :param kind: ``"model"``, ``"tool"`` or ``"other"``.
+ :param fingerprint: What the step was asked, or ``None`` when it cannot be
+ fingerprinted. A step without one replays on a matching name, and only right
+ after a step that replayed.
+ :param replayable: ``False`` for a step that must always run live, such as a tool
+ whose effect Airflow cannot observe. It still takes a position, so the steps
+ after it keep theirs, but nothing is looked up or recorded for it.
+ """
+ position = self._position
+ self._position += 1
+ step = JournalStep(self, position, name, kind, fingerprint, replayable=replayable)
+ self._load_meta()
+ if not self.diverged and self._diverge_from is not None and position >= self._diverge_from:
+ self._diverge(position, name, reason="the previous attempt failed before reaching this step")
+ entry = self.peek(position) if replayable else None
+ self._peeked.pop(position, None)
+ previous_replayed, self._previous_replayed = self._previous_replayed, False
+ if entry is None:
+ return step
+ if entry.get("name") == name and entry.get("fingerprint") == fingerprint:
+ if fingerprint is None and not previous_replayed:
+ return step
+ step.replayed = self._previous_replayed = True
+ step.payload = entry.get("payload")
+ self.journal.stats.replayed[kind] += 1
+ log.debug("Durable: replayed step", position=position, step=name)
+ return step
+ self._diverge(
+ position,
+ name,
+ reason=(
+ f"the previous attempt ran {entry.get('name')!r} here"
+ if entry.get("name") != name
+ else "the request differs from the previous attempt's (prompt, model, settings, "
+ "tools, arguments or conversation so far)"
+ ),
+ )
+ return step
+
+ async def run(
+ self,
+ name: str,
+ *,
+ kind: StepKind,
+ fingerprint: str | None,
+ body: Callable[[], Awaitable[Any]],
+ replayable: bool = True,
+ to_record: Callable[[Any], Any] | None = None,
+ ) -> Any:
+ """Replay the step at the next position, or run ``body`` and record what it returns."""
+ step = self.claim(name, kind=kind, fingerprint=fingerprint, replayable=replayable)
+ if step.replayed:
+ return step.payload
+ return await step.run(body, to_record=to_record)
+
+ def save(self, position: int, entry: dict[str, Any]) -> bool:
+ """Store ``entry`` at ``position``; ``False`` when the storage skipped it."""
+ self._last_saved = max(self._last_saved, position)
+ return self.journal.storage.save_step(build_step_key(self.key, position), entry)
+
+ def fail(self, error: BaseException) -> None:
+ """
+ Record that the run failed with ``error``, so a retry does not replay its error path.
+
+ Steps the run started before the step that raised ``error`` (tool calls running
+ alongside it, say) still replay on retry; everything from then on, including what
+ error handling recorded, runs again.
+ """
+ failed_at = self._position
+ cause: BaseException | None = error
+ while cause is not None:
+ failed_at = min([failed_at, *(position for raised, position in self._raised if raised is cause)])
+ cause = cause.__cause__ or cause.__context__
+ self._load_meta()
+ # An attempt that was still replaying an earlier one when it failed, and recorded
+ # nothing from the failure on, leaves that earlier attempt's steps as they were.
+ diverge_from = failed_at if self.diverged or self._last_saved >= failed_at else self._diverge_from
+ meta = {"diverge_from": diverge_from, "high_water": max(self._high_water, self._position)}
+ self.journal.storage.save_step(build_run_meta_key(self.key), meta)
+
+ def cleanup(self) -> None:
+ """Delete this run's steps. Call only once the run, and anything that consumes it, has succeeded."""
+ self._load_meta()
+ end = max(self._high_water, self._position)
+ # An attempt killed before it could note how far it got may have gone further.
+ while self.journal.storage.load_step(build_step_key(self.key, end)) is not None:
+ end += 1
+ keys = [build_step_key(self.key, position) for position in range(end)]
+ self.journal.storage.delete_steps([*keys, build_run_meta_key(self.key)])
+ self.journal._forget_run(self.key)
+
+ def reject(self, step: JournalStep, *, reason: str) -> None:
+ """
+ Run a step that was going to replay live instead, and every step after it.
+
+ For an adapter that finds the recorded payload unusable after :meth:`claim`
+ matched it, such as one written by a version whose types have changed since.
+ """
+ if not step.replayed:
+ return
+ step.replayed = False
+ step.payload = None
+ self._previous_replayed = False
+ self.journal.stats.replayed[step.kind] -= 1
+ self._diverge(step.position, step.name, reason=reason)
+
+ def _diverge(self, position: int, name: str, *, reason: str) -> None:
+ self.diverged = True
+ self._peeked.clear()
+ log.warning(
+ "Durable: the run took a different path from the previous attempt; this step and "
+ "every step after it run again",
+ position=position,
+ step=name,
+ reason=reason,
+ )
+
+
+class DurableJournal:
+ """
+ One task attempt's view of the durable journal.
+
+ :param storage: Where steps are stored; see :func:`build_task_storage`.
+ :param clean_up_after_run: Delete a run's steps as soon as it succeeds. A journal the
+ operator manages leaves this off and calls :meth:`cleanup` itself after the
+ task's own post-run work has succeeded.
+ """
+
+ def __init__(self, storage: DurableStorageProtocol, *, clean_up_after_run: bool = False) -> None:
+ self.storage = storage
+ self.clean_up_after_run = clean_up_after_run
+ self.stats = JournalStats()
+ self._runs: list[DurableRun] = []
+ self._top_level_runs = 0
+ self._run_id: str | None = None
+ # Keys of every run any attempt started, read on first use. Cleanup walks it, so a run
+ # an earlier attempt started that this one never reaches (an agent called from a tool
+ # that now replays) is still deleted.
+ self._known_runs: list[str] | None = None
+
+ def get_run_id(self, *, default: str) -> str:
+ """
+ Return the id that names the task's agent run on every attempt.
+
+ The first attempt stores ``default`` beside the journal; a retry reads it back. It
+ is deleted with the journal when the task succeeds, so a later clear of the task
+ starts over with a new id. Capabilities that key their own state on the agent's
+ ``run_id``, such as pydantic-ai-harness ``SpendLimits`` and ``StepPersistence``,
+ need it to stay the same across a replay.
+
+ :param default: The id to use when no earlier attempt stored one, such as the
+ first attempt's task instance id.
+ """
+ if self._run_id is None:
+ entry = self.storage.load_step(RUN_ID_KEY)
+ stored = entry.get("run_id") if entry is not None else None
+ if isinstance(stored, str):
+ self._run_id = stored
+ else:
+ self._run_id = default
+ self.storage.save_step(RUN_ID_KEY, {"run_id": default})
+ return self._run_id
+
+ def start_run(self) -> DurableRun:
+ """
+ Start an agent run.
+
+ A run started while a step of this journal executes (an agent a tool calls) is
+ numbered under that step; any other run takes the next top-level number.
+ """
+ parent = _EXECUTING_STEP.get()
+ if parent is not None and parent.durable_run.journal is self:
+ parent.nested_runs += 1
+ key = f"{parent.durable_run.key}.{parent.position}.{parent.nested_runs}"
+ else:
+ key = str(self._top_level_runs)
+ self._top_level_runs += 1
+ durable_run = DurableRun(self, key)
+ self._runs.append(durable_run)
+ known = self._load_known_runs()
+ if key not in known:
+ known.append(key)
+ self.storage.save_step(RUNS_KEY, {"runs": known})
+ return durable_run
+
+ def _load_known_runs(self) -> list[str]:
+ if self._known_runs is None:
+ entry = self.storage.load_step(RUNS_KEY) or {}
+ runs = entry.get("runs")
+ self._known_runs = [key for key in runs if isinstance(key, str)] if isinstance(runs, list) else []
+ return self._known_runs
+
+ def cleanup(self) -> None:
+ """Delete the steps of every run in this attempt. Call only once the task's work has succeeded."""
+ started = {durable_run.key: durable_run for durable_run in self._runs}
+ for key in dict.fromkeys([*self._load_known_runs(), *started]):
+ (started.get(key) or DurableRun(self, key)).cleanup()
+ self.storage.delete_steps([RUN_ID_KEY])
+ log.debug("Durable journal cleaned up")
+
+ def _forget_run(self, key: str) -> None:
+ known = self._load_known_runs()
+ if key not in known:
+ return
+ known.remove(key)
+ if known:
+ self.storage.save_step(RUNS_KEY, {"runs": known})
+ else:
+ self.storage.delete_steps([RUNS_KEY])
+
+ def log_summary(self, logger: logging.Logger | Any) -> None:
+ """Log what this attempt replayed and recorded, and which steps a retry would run again."""
+ replayed, recorded = self.stats.replayed, self.stats.recorded
+ logger.info(
+ "Durable: replayed %d steps (%d model, %d tool, %d other), recorded %d new steps "
+ "(%d model, %d tool, %d other)",
+ replayed.total(),
+ replayed["model"],
+ replayed["tool"],
+ replayed["other"],
+ recorded.total(),
+ recorded["model"],
+ recorded["tool"],
+ recorded["other"],
+ )
+ if tools := [name for kind, name in self.stats.not_recorded if kind == "tool"]:
+ counts = collections.Counter(tools)
+ logger.warning(
+ "Durable: %d tool results were not recorded, and a retry runs them again: %s",
+ len(tools),
+ ", ".join(name if n == 1 else f"{name} (x{n})" for name, n in counts.items()),
+ )
+ if models := sum(1 for kind, _ in self.stats.not_recorded if kind == "model"):
+ logger.warning(
+ "Durable: %d model responses were not recorded, and a retry re-runs them "
+ "and every step after the first of them",
+ models,
+ )
+
+
+def build_task_storage(context: Context) -> DurableStorageProtocol:
+ """
+ Return the journal storage for the task instance ``context`` belongs to.
+
+ On Airflow >= 3.3 the journal lives in the AIP-103 task state store, which handles
+ persistence and large-value offload, so ``[common.ai] durable_cache_path`` is not
+ needed. On older versions it is a JSON file under that ObjectStorage path.
+ """
+ if AIRFLOW_V_3_3_PLUS:
+ # Imported lazily: NEVER_EXPIRE and the task state store accessor do not
+ # exist on Airflow versions before 3.3.
+ from airflow.providers.common.ai.durable.task_state_store import TaskStateStoreDurableStorage
+
+ return TaskStateStoreDurableStorage(context["task_state_store"])
+
+ ti = context["task_instance"]
+ return DurableStorage(
+ 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,
+ )
+
+
+_ACTIVE_JOURNAL: ContextVar[DurableJournal | None] = ContextVar("commonai_durable_journal", default=None)
+_EXECUTING_STEP: ContextVar[JournalStep | None] = ContextVar("commonai_durable_step", default=None)
+# The journal for the task attempt this process is running, when no caller set one
+# explicitly. Keyed by the attempt, so a second attempt in the same process (a test,
+# a long-lived worker) never reuses the first attempt's runs.
+_task_journal: tuple[str, DurableJournal] | None = None
+
+
+@contextmanager
+def journal_scope(journal: DurableJournal) -> Iterator[DurableJournal]:
+ """Make ``journal`` the active journal for the code in the ``with`` block."""
+ token = _ACTIVE_JOURNAL.set(journal)
+ try:
+ yield journal
+ finally:
+ _ACTIVE_JOURNAL.reset(token)
+
+
+def current_journal() -> DurableJournal | None:
+ """
+ Return the active journal, or ``None`` outside an Airflow task.
+
+ A journal set with :func:`journal_scope` wins. Otherwise, inside a running task,
+ this returns the task attempt's own journal, created on first use, which deletes the
+ steps of each run as soon as it succeeds. Outside a task there is nothing to replay
+ into, so durable execution is off and the agent runs as it would without it.
+ """
+ global _task_journal
+
+ if (journal := _ACTIVE_JOURNAL.get()) is not None:
+ return journal
+ try:
+ context = get_current_context()
+ except (RuntimeError, AirflowException):
+ # Airflow 3 raises RuntimeError outside a task, Airflow 2 AirflowException.
+ return None
+ attempt = make_task_instance_run_key(context["task_instance"])
+ if _task_journal is None or _task_journal[0] != attempt:
+ _task_journal = (attempt, DurableJournal(build_task_storage(context), clean_up_after_run=True))
+ return _task_journal[1]
diff --git a/providers/common/ai/src/airflow/providers/common/ai/durable/replay_usage.py b/providers/common/ai/src/airflow/providers/common/ai/durable/replay_usage.py
index 35bf1a2125f..688fa6ba065 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/durable/replay_usage.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/durable/replay_usage.py
@@ -20,20 +20,32 @@ from __future__ import annotations
import dataclasses
from copy import deepcopy
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
-from pydantic_ai.usage import RequestUsage, UsageLimits
+from pydantic import TypeAdapter
+from pydantic_ai.messages import ModelResponse, ToolCallPart
+from pydantic_ai.usage import RequestUsage, RunUsage, UsageLimits
if TYPE_CHECKING:
- from pydantic_ai.messages import ModelResponse
- from pydantic_ai.usage import RunUsage
+ from pydantic_ai.messages import ModelMessage
+ from pydantic_ai.models import ModelRequestParameters
-_REQUEST_USAGE_FIELDS = dataclasses.fields(RequestUsage)
+ from airflow.providers.common.ai.durable.journal import DurableRun
+_RUN_USAGE_ADAPTER = TypeAdapter(RunUsage)
-def subtract_request_usage(run_usage: RunUsage, usage: RequestUsage) -> None:
- """Subtract one response's token usage and cost from a run's usage, in place."""
- for field in _REQUEST_USAGE_FIELDS:
+# Tool-call payload kinds pydantic-ai counts as a successful call (``tool_calls += 1``);
+# a recorded ModelRetry, approval request or failure is replayed without being counted.
+SUCCESSFUL_TOOL_RESULT_KINDS = frozenset({"tool_return", "tool_content_result"})
+
+# How far past a replayed model response to look for the steps that followed it, beyond
+# one position per tool call: room for argument validation and capability operations.
+_SUCCESSOR_SCAN_SLACK = 16
+
+
+def subtract_request_usage(run_usage: RunUsage, usage: RequestUsage | RunUsage) -> None:
+ """Subtract one response's (or one durable operation's) usage and cost from a run's usage, in place."""
+ for field in dataclasses.fields(usage):
if field.name == "details":
for name, value in usage.details.items():
run_usage.details[name] = run_usage.details.get(name, 0) - value
@@ -48,10 +60,10 @@ def fill_replayed_cost(response: ModelResponse) -> None:
"""
Price a replayed response in place, the way pydantic-ai's graph is about to.
- ``CachingModel`` stores a live response before the graph prices it, so a cached
+ The journal records a live response before the graph prices it, so a replayed
response comes back with ``usage.cost=None`` even for a priced model. The graph fills
- the cost right after ``CachingModel.request`` returns and only when it is still unset,
- so pricing it here first means the cost this module subtracts is exactly the cost the
+ the cost right after the durable model request returns, and only when it is still
+ unset, by the same lookup, so the cost this module subtracts is exactly the cost the
graph then adds. Uses the public ``ModelResponse.cost()``, which runs the same price
lookup as the graph's own best-effort pricing; a response it cannot price stays
unpriced, as it would in the graph.
@@ -77,25 +89,25 @@ class ReplayUsageLedger:
"""
Nets durable replays out of the ``RunUsage`` a run is counted and limited against.
- pydantic-ai counts a replayed step exactly like a live one: after
- ``CachingModel.request`` returns, the graph does ``requests += 1`` and adds the
- response's usage, and after ``CachingToolset.call_tool`` returns, the tool manager
- does ``tool_calls += 1``. The ledger cancels each replay when it happens, so the run's
+ pydantic-ai counts a replayed step exactly like a live one: after a durable model
+ request returns, the graph does ``requests += 1`` and adds the response's usage, and
+ after a durable tool call returns successfully, the tool manager does
+ ``tool_calls += 1``. The ledger cancels each replay when it happens, so the run's
counts only move for live work. Nothing here is persisted: the result is the same
- whether the cache entry was written by the previous attempt, an attempt that failed
+ whether the journal entry was written by the previous attempt, an attempt that failed
halfway through its own replay, or an attempt before a clear (which resets the budget
- but keeps the durable cache).
+ but keeps the durable journal).
Two of pydantic-ai's limit checks run before the durable layer sees the step it
guards: ``check_before_request`` before each model request, and the up-front
``check_before_tool_call`` projection over a whole step's function-tool calls. For
those, the ledger takes a credit ahead of time -- when a replayed model response is
- followed in the cache by the next model step or by cached tool results -- and resolves
+ followed in the journal by the next model request or by recorded tool results -- and resolves
each credit when the step runs: a replay keeps it, a live call gives it back and
re-runs the check pydantic-ai made against the credited count. Credits that are never
resolved (the run took a different path, or stopped) are given back by
:meth:`settle`, so they never reach the persisted total; until then, a credit for a
- stale cache entry the run never reaches has only loosened that one up-front check.
+ stale journal entry the run never reaches has only loosened that one up-front check.
:param run_usage: The ``RunUsage`` passed to the run as ``usage=``.
:param usage_limits: The limits passed to the run; ``None`` means pydantic-ai's
@@ -115,16 +127,104 @@ class ReplayUsageLedger:
# segment was replayed; see record_model_replay.
self._chain_segment_usage: RequestUsage | None = None
self._chain_response_id: str | None = None
+ # Function-tool calls issued so far by the current continuation chain, and whether
+ # the last response was suspended (so the next request continues it).
+ self._chain_function_calls = 0
+ self._previous_suspended = False
+
+ def begin_model_request(self, messages: list[ModelMessage]) -> bool:
+ """
+ Note the start of a model request and return whether it continues a suspended response.
+
+ pydantic-ai re-issues a suspended response (Anthropic ``pause_turn``, OpenAI
+ background mode) by sending it back as the last message; the graph counts the whole
+ chain as one request.
+ """
+ continuation = (
+ self._previous_suspended
+ and bool(messages)
+ and isinstance(messages[-1], ModelResponse)
+ and messages[-1].state == "suspended"
+ )
+ if not continuation:
+ self._chain_function_calls = 0
+ return continuation
+
+ def track_chain(self, response: ModelResponse, model_request_parameters: ModelRequestParameters) -> None:
+ """Count the function-tool calls ``response`` issues, for crediting the steps after it."""
+ function_tools = {tool.name for tool in model_request_parameters.function_tools}
+ self._chain_function_calls += sum(
+ 1
+ for part in response.parts
+ if isinstance(part, ToolCallPart) and part.tool_name in function_tools
+ )
+ self._previous_suspended = response.state == "suspended"
+
+ def credit_first_replay(self, durable_run: DurableRun) -> None:
+ """
+ Credit the run's first model request if the journal holds a response for it.
+
+ Called before the run starts, because pydantic-ai checks ``request_limit`` before
+ the first request reaches the durable layer; without the credit, a retry whose
+ seeded ``requests`` already equals ``request_limit`` could not start even when
+ every step would replay for free. Steps recorded before the first model request
+ (tool discovery, capability operations) are skipped over.
+ """
+ position = durable_run.position
+ while (entry := durable_run.peek(position)) is not None:
+ if entry.get("kind") == "model":
+ self.credit_request()
+ return
+ position += 1
+
+ def credit_successors(self, durable_run: DurableRun, position: int) -> None:
+ """
+ Credit what the replayed, complete response at ``position`` is followed by in the journal.
+
+ The response's function-tool calls took positions after it when they first ran, and
+ the next model request a position after them. Tool calls with a recorded successful
+ result are credited; the scan stops at the next model request, which is credited too.
+ A position with nothing recorded is a call that raised or was not recorded, and the
+ scan gives up once there are more of those than calls in the chain.
+ """
+ calls = self._chain_function_calls
+ tool_positions: list[int] = []
+ gaps = 0
+ for index in range(position + 1, position + 2 * calls + _SUCCESSOR_SCAN_SLACK + 2):
+ entry = durable_run.peek(index)
+ if entry is None:
+ gaps += 1
+ if gaps > calls:
+ break
+ continue
+ kind = entry.get("kind")
+ if kind == "model":
+ self.credit_request()
+ break
+ if kind == "tool" and is_successful_tool_payload(entry.get("payload")):
+ tool_positions.append(index)
+ self.credit_tool_steps(tool_positions, batch_calls=calls)
+
+ def record_capability_replay(self, payload: Any) -> None:
+ """
+ Cancel the usage a replayed ``@durable_operation`` brings back with its result.
+
+ pydantic-ai adds the usage an operation accumulated (a summarization's model call,
+ say) to the run when the operation returns, including when it returns from the
+ journal; that usage was already counted by the attempt that ran it.
+ """
+ if isinstance(payload, dict) and isinstance(payload.get("usage_delta"), dict):
+ subtract_request_usage(self.run_usage, _RUN_USAGE_ADAPTER.validate_python(payload["usage_delta"]))
def credit_request(self) -> None:
- """Credit the next model request, which the cache says will be a replay."""
+ """Credit the next model request, which the journal says will be a replay."""
if not self._request_credit:
self._request_credit = True
self.run_usage.requests -= 1
def credit_tool_steps(self, steps: list[int], *, batch_calls: int) -> None:
"""
- Credit the tool calls at these step indices, which have cached results.
+ Credit the tool calls at these journal positions, which have recorded results.
:param batch_calls: The number of function-tool calls this replayed response
issued, i.e. the count pydantic-ai's up-front ``check_before_tool_call``
@@ -202,7 +302,7 @@ class ReplayUsageLedger:
def record_live_tool_call(self, step: int) -> None:
"""
- Give back a credit whose cached result did not replay, and re-check the limit.
+ Give back a credit whose recorded result did not replay, and re-check the limit.
pydantic-ai's up-front check projected the whole batch this credited call
belongs to as if every credited call in it would replay for free. With this one
@@ -238,3 +338,8 @@ class ReplayUsageLedger:
projected = deepcopy(self.run_usage)
projected.tool_calls = self._tool_batch_projection
self._limits.check_before_tool_call(projected)
+
+
+def is_successful_tool_payload(payload: Any) -> bool:
+ """Whether a recorded tool-call payload is a result pydantic-ai counts as a successful call."""
+ return isinstance(payload, dict) and payload.get("kind") in SUCCESSFUL_TOOL_RESULT_KINDS
diff --git a/providers/common/ai/src/airflow/providers/common/ai/durable/step_counter.py b/providers/common/ai/src/airflow/providers/common/ai/durable/step_counter.py
deleted file mode 100644
index 314195cbefb..00000000000
--- a/providers/common/ai/src/airflow/providers/common/ai/durable/step_counter.py
+++ /dev/null
@@ -1,52 +0,0 @@
-# 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.
-"""Shared step counter for durable execution caching."""
-
-from __future__ import annotations
-
-
-class DurableStepCounter:
- """
- Monotonically increasing counter shared between CachingModel and CachingToolset.
-
- Each model call and tool call increments the counter. The step index
- is used as the cache key; replay correctness is verified separately by
- comparing the request fingerprint stored with each cache entry (see
- ``airflow.providers.common.ai.durable.fingerprint``).
- """
-
- def __init__(self) -> None:
- self._step: int = 0
- self.replayed_model: int = 0
- self.replayed_tool: int = 0
- self.cached_model: int = 0
- self.cached_tool: int = 0
- # Steps that ran live but whose cache write the backend refused. They
- # re-run on retry, so they must not count as cached.
- self.skipped_model: int = 0
- self.skipped_tools: list[str] = []
-
- def next_step(self) -> int:
- """Return the current step and advance the counter."""
- step = self._step
- self._step += 1
- return step
-
- @property
- def total_steps(self) -> int:
- """Total number of steps executed so far."""
- return self._step
diff --git a/providers/common/ai/src/airflow/providers/common/ai/durable/storage.py b/providers/common/ai/src/airflow/providers/common/ai/durable/storage.py
index 4e98fe40eb4..3950b6f88b0 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/durable/storage.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/durable/storage.py
@@ -14,21 +14,16 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""ObjectStorage-backed durable storage for pydantic-ai agent step caching."""
+"""ObjectStorage-backed storage for the durable execution journal (Airflow < 3.3)."""
from __future__ import annotations
-import contextlib
import hashlib
import json
+from collections.abc import Iterable
from functools import lru_cache
from typing import Any
-from pydantic_ai.messages import ModelMessagesTypeAdapter, ModelResponse
-
-# Sentinel to distinguish "cached None" from "no cache entry" for tool results.
-# Shared with the task state store backend so the envelope shape cannot drift.
-from airflow.providers.common.ai.durable.base import TOOL_RESULT_SENTINEL as _SENTINEL
from airflow.providers.common.ai.utils.task_logger import get_task_logger
log = get_task_logger()
@@ -51,15 +46,16 @@ def _get_base_path():
class DurableStorage:
"""
- Stores step-level caches in a single JSON file on ObjectStorage.
+ Stores the durable journal in a single JSON file on ObjectStorage.
- All step caches (model responses and tool results) are stored as entries
+ All journal steps are stored as entries
in a single JSON blob, written to ``{base_path}/{cache_id}.json`` where
``cache_id`` is a hash of the task instance's identity (dag, task, run,
map index) so distinct task instances never share a file.
The file survives Airflow task retries since it lives outside the
- XCom system. It is deleted on successful task completion.
+ XCom system. :meth:`delete_steps` removes the steps of a run that succeeded, and
+ the file once nothing is left in it.
:param dag_id: DAG ID of the running task.
:param task_id: Task ID of the running task.
@@ -101,109 +97,49 @@ class DurableStorage:
return self._cache
- def _save_cache(self) -> None:
- """Persist the in-memory cache blob to storage."""
+ def _write(self, text: str) -> None:
+ """Persist the encoded journal to storage."""
path = self._get_path()
path.parent.mkdir(parents=True, exist_ok=True)
- path.write_text(json.dumps(self._cache))
-
- def save_model_response(self, key: str, response: ModelResponse, *, fingerprint: str | None) -> bool:
- """
- Serialize and store a ModelResponse with the request fingerprint that produced it.
-
- :return: Always ``True``. Unlike the task state store, this backend never skips a model response.
- """
- cache = self._load_cache()
- # Store the dumped messages as native JSON-compatible objects, not a
- # pre-encoded string: the whole cache is JSON-encoded once in
- # ``_save_cache``, so embedding a string here would double-encode the
- # (large) response payload.
- cache[key] = {
- "fingerprint": fingerprint,
- "data": ModelMessagesTypeAdapter.dump_python([response], mode="json"),
- }
- self._save_cache()
- return True
-
- def load_model_response(self, key: str) -> tuple[ModelResponse | None, str | None]:
- """
- Load a cached ModelResponse and its stored request fingerprint.
+ path.write_text(text)
- Returns ``(None, None)`` if not cached. Entries written before
- fingerprints existed load with a ``None`` fingerprint.
+ def save_step(self, key: str, entry: dict[str, Any]) -> bool:
"""
- cache = self._load_cache()
- raw = cache.get(key)
- if raw is None:
- return None, None
- try:
- if isinstance(raw, dict):
- messages = ModelMessagesTypeAdapter.validate_python(raw["data"])
- fingerprint = raw.get("fingerprint")
- else:
- # Legacy entry: the adapter JSON (a list) was stored directly as a string.
- messages = ModelMessagesTypeAdapter.validate_json(raw)
- fingerprint = None
- except (KeyError, IndexError, ValueError):
- # A torn or malformed entry degrades to a miss (the step re-runs),
- # never a task crash -- the cache is best-effort.
- log.warning("Durable: ignoring malformed cached model response", key=key)
- return None, None
- if not messages:
- return None, None
- return messages[0], fingerprint # type: ignore[return-value]
-
- def save_tool_result(self, key: str, result: Any, *, fingerprint: str | None) -> bool:
- """
- Store a tool call result with the call fingerprint that produced it.
-
- Non-serializable results (e.g. BinaryContent from MCP tools) are
- skipped with a warning -- the tool call still succeeds, but won't
- be replayed on retry.
+ Store one journal step and rewrite the journal file.
- :return: ``True`` if the entry was written, ``False`` if it was skipped.
+ :return: ``True`` if the entry was written, ``False`` if it could not be encoded.
"""
cache = self._load_cache()
+ cache[key] = entry
try:
- # Probe serializability before mutating the shared cache: a
- # non-serializable result must skip only this entry, not break the
- # whole-file ``_save_cache``. TypeError covers unsupported types;
- # ValueError covers circular references.
- json.dumps(result)
+ text = json.dumps(cache)
except (TypeError, ValueError):
- log.warning(
- "Durable: skipping cache for non-serializable tool result",
- key=key,
- type=type(result).__name__,
- )
+ # Only this entry is skipped; the rest of the journal is still written.
+ del cache[key]
+ log.warning("Durable: could not store step", key=key, exc_info=True)
return False
- cache[key] = {_SENTINEL: True, "value": result, "fingerprint": fingerprint}
- self._save_cache()
+ self._write(text)
return True
- def load_tool_result(self, key: str) -> tuple[bool, Any, str | None]:
- """
- Load a cached tool result and its stored call fingerprint.
+ def load_step(self, key: str) -> dict[str, Any] | None:
+ """Load one journal step, or ``None`` when there is none or it is not a journal entry."""
+ raw = self._load_cache().get(key)
+ return raw if isinstance(raw, dict) else None
- Returns a (found, value, fingerprint) tuple since the cached value
- itself could be None. Entries written before fingerprints existed
- load with a ``None`` fingerprint.
- """
+ def delete_steps(self, keys: Iterable[str]) -> None:
+ """Remove these steps, and the journal file once nothing is left in it."""
cache = self._load_cache()
- raw = cache.get(key)
- if raw is None:
- return False, None, None
- # Legacy entries were stored as a JSON string; new entries are native dicts.
- if isinstance(raw, str):
- raw = json.loads(raw)
- if not isinstance(raw, dict) or _SENTINEL not in raw:
- return False, None, None
- return True, raw["value"], raw.get("fingerprint")
-
- def cleanup(self) -> None:
- """Delete the cache file after successful execution."""
- # Best-effort cleanup
- with contextlib.suppress(FileNotFoundError, OSError):
- self._get_path().unlink()
- self._cache = None
- log.debug("Durable cache cleaned up")
+ for key in keys:
+ cache.pop(key, None)
+ # Runs after the run has already succeeded, so it must never raise; a file left
+ # behind holds only steps a retry of this run would replay.
+ try:
+ if cache:
+ self._write(json.dumps(cache))
+ else:
+ self._get_path().unlink()
+ self._cache = None
+ except FileNotFoundError:
+ self._cache = None
+ except Exception:
+ log.warning("Durable: could not clean up the journal file", exc_info=True)
diff --git a/providers/common/ai/src/airflow/providers/common/ai/durable/task_state_store.py b/providers/common/ai/src/airflow/providers/common/ai/durable/task_state_store.py
index c02a33881c1..0603b6ca679 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/durable/task_state_store.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/durable/task_state_store.py
@@ -15,36 +15,31 @@
# specific language governing permissions and limitations
# under the License.
"""
-Task-state-store-backed durable storage for pydantic-ai agent step caching.
+Task-state-store-backed storage for the durable execution journal.
Available on Airflow >= 3.3, where the AIP-103 task state store provides a
per-task-instance key/value store that survives retries within a run and is
-cleared when the run is removed. Each cached step is written under its own key
-(``model_step_{N}`` / ``tool_step_{N}``, each prefixed with the reserved
-``DURABLE_KEY_PREFIX`` so it cannot collide with user keys in the shared
-key namespace); the store handles persistence and,
-when ``[workers] state_store_backend`` is configured, transparently offloads
-large values to external storage. No ``[common.ai] durable_cache_path`` is
-needed.
+cleared when the run is removed. Each journal step is written under its own key
+(``step_{N}``, prefixed with the reserved ``DURABLE_KEY_PREFIX`` so it cannot
+collide with user keys in the shared key namespace); the store handles
+persistence and, when ``[workers] state_store_backend`` is configured,
+transparently offloads large values to external storage. No
+``[common.ai] durable_cache_path`` is needed.
This module is imported only on Airflow >= 3.3 (see
-``AgentOperator._build_durable_storage``); ``NEVER_EXPIRE`` does not exist on
-older airflow versions.
+:func:`~airflow.providers.common.ai.durable.journal.build_task_storage`);
+``NEVER_EXPIRE`` does not exist on older Airflow versions.
"""
from __future__ import annotations
-import json
from typing import TYPE_CHECKING, Any
-from pydantic_ai.messages import ModelMessagesTypeAdapter
-
-from airflow.providers.common.ai.durable.base import TOOL_RESULT_SENTINEL
from airflow.providers.common.ai.utils.task_logger import get_task_logger
from airflow.sdk.execution_time.context import NEVER_EXPIRE
if TYPE_CHECKING:
- from pydantic_ai.messages import ModelResponse
+ from collections.abc import Iterable
from airflow.sdk.execution_time.context import TaskStateStoreAccessor
@@ -53,15 +48,15 @@ log = get_task_logger()
class TaskStateStoreDurableStorage:
"""
- Stores step-level durable caches in the AIP-103 task state store.
+ Stores durable journal steps in the AIP-103 task state store.
- Each model response and tool result is written under its own key, scoped to
- the current task instance. Entries are written with ``NEVER_EXPIRE`` so a
- retry can replay them regardless of ``retry_delay`` or the global retention
- config, and the keys this run touched are deleted on successful completion.
+ Each step is written under its own key, scoped to the current task instance.
+ Entries are written with ``NEVER_EXPIRE`` so a retry can replay them regardless
+ of ``retry_delay`` or the global retention config, and the keys this attempt
+ touched are deleted by :meth:`delete_steps` once the run succeeds.
A run that fails permanently leaves its keys behind (``NEVER_EXPIRE`` skips
- garbage collection); they are removed when the DAG run is cleaned up, since
+ garbage collection); they are removed when the Dag run is cleaned up, since
task state store rows cascade with the run.
:param accessor: The task state store accessor for the current task
@@ -70,122 +65,36 @@ class TaskStateStoreDurableStorage:
def __init__(self, accessor: TaskStateStoreAccessor) -> None:
self._store = accessor
- # Keys written or replayed this run, deleted on cleanup. A divergent
- # retry that takes fewer steps may orphan keys from a longer earlier
- # attempt; those are reclaimed by the DAG-run cascade, not here.
- self._keys: set[str] = set()
- def save_model_response(self, key: str, response: ModelResponse, *, fingerprint: str | None) -> bool:
+ def save_step(self, key: str, entry: dict[str, Any]) -> bool:
"""
- Serialize and store a ModelResponse with the request fingerprint that produced it.
+ Store one journal step.
- Best-effort: the save runs *after* the live model call already succeeded, so a
- failed write (e.g. a value over the backend's size limit) must not fail the step.
- It is skipped with a warning and simply re-runs live on the next retry.
+ Best-effort: the save runs *after* the step already succeeded, so a failed
+ write (e.g. a value over the backend's size limit) must not fail the step.
+ It is skipped with a warning, and the step runs again on the next retry.
:return: ``True`` if the entry was written, ``False`` if it was skipped.
"""
try:
- self._store.set(
- key,
- {
- "fingerprint": fingerprint,
- "data": ModelMessagesTypeAdapter.dump_python([response], mode="json"),
- },
- retention=NEVER_EXPIRE,
- )
+ self._store.set(key, entry, retention=NEVER_EXPIRE)
except Exception:
- log.warning("Durable: skipping cache for model response", key=key, exc_info=True)
+ log.warning("Durable: could not store step", key=key, exc_info=True)
return False
- self._keys.add(key)
return True
- def load_model_response(self, key: str) -> tuple[ModelResponse | None, str | None]:
- """
- Load a cached ModelResponse and its stored request fingerprint.
-
- Returns ``(None, None)`` on a miss or a torn entry, so the step re-runs
- live rather than crashing the task.
- """
+ def load_step(self, key: str) -> dict[str, Any] | None:
+ """Load one journal step, or ``None`` when there is none or it is not a journal entry."""
raw = self._store.get(key)
- if not isinstance(raw, dict):
- return None, None
- try:
- messages = ModelMessagesTypeAdapter.validate_python(raw["data"])
- except (KeyError, IndexError, TypeError, ValueError):
- log.warning("Durable: ignoring malformed cached model response", key=key)
- return None, None
- # A foreign/torn entry can validate as a ModelRequest; only a response is replayable.
- if not messages or messages[0].kind != "response":
- return None, None
- self._keys.add(key)
- fingerprint = raw.get("fingerprint")
- return messages[0], fingerprint if isinstance(fingerprint, str) else None # type: ignore[return-value]
-
- def save_tool_result(self, key: str, result: Any, *, fingerprint: str | None) -> bool:
- """
- Store a tool call result with the call fingerprint that produced it.
-
- Non-serializable results (e.g. BinaryContent from MCP tools) are skipped
- with a warning -- the tool call still succeeds, but won't be replayed on
- retry.
-
- :return: ``True`` if the entry was written, ``False`` if it was skipped.
- """
- try:
- # The store validates against pydantic ``JsonValue``, which is stricter than
- # ``json.dumps``: it rejects tuples and non-string dict keys. Round-trip through
- # JSON to coerce those (tuple -> list, non-str keys -> str) -- matching the < 3.3
- # ObjectStorage backend -- so a common ``return (result, meta)`` is cached rather
- # than crashing the step. Non-serializable results (e.g. BinaryContent from MCP
- # tools) are skipped with a warning.
- normalized = json.loads(json.dumps(result))
- except (TypeError, ValueError):
- log.warning(
- "Durable: skipping cache for non-serializable tool result",
- key=key,
- type=type(result).__name__,
- )
- return False
- try:
- # Best-effort like the model-response save: a write that fails after the tool
- # already ran (e.g. an oversize value) must not fail the step -- skip and re-run
- # live on retry rather than surface an opaque comms error.
- self._store.set(
- key,
- {TOOL_RESULT_SENTINEL: True, "value": normalized, "fingerprint": fingerprint},
- retention=NEVER_EXPIRE,
- )
- except Exception:
- log.warning("Durable: skipping cache for tool result", key=key, exc_info=True)
- return False
- self._keys.add(key)
- return True
-
- def load_tool_result(self, key: str) -> tuple[bool, Any, str | None]:
- """
- Load a cached tool result and its stored call fingerprint.
+ return raw if isinstance(raw, dict) else None
- Returns a ``(found, value, fingerprint)`` tuple since the cached value
- itself may be ``None``.
- """
- raw = self._store.get(key)
- if not isinstance(raw, dict) or TOOL_RESULT_SENTINEL not in raw:
- return False, None, None
- self._keys.add(key)
- fingerprint = raw.get("fingerprint")
- return True, raw.get("value"), fingerprint if isinstance(fingerprint, str) else None
-
- def cleanup(self) -> None:
- """Delete the keys this run wrote or replayed after successful execution."""
- for key in self._keys:
- # Runs only after the task has already succeeded, so it must never raise
- # (that would fail a succeeded task). A key left behind by a failed delete
- # is reclaimed by the DAG-run cascade -- hence the deliberately broad catch.
+ def delete_steps(self, keys: Iterable[str]) -> None:
+ """Delete these keys; a key left behind by a failed delete is reclaimed by the Dag-run cascade."""
+ for key in keys:
+ # Runs only after the run has already succeeded, so it must never raise
+ # (that would fail a succeeded task) -- hence the deliberately broad catch.
# Log it so an offloaded value orphaned in external storage is at least visible.
try:
self._store.delete(key)
except Exception:
- log.warning("Durable: failed to delete cache key on cleanup", key=key, exc_info=True)
- self._keys.clear()
- log.debug("Durable cache cleaned up")
+ log.warning("Durable: failed to delete journal key", key=key, exc_info=True)
diff --git a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_durable.py b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_durable.py
index b7d2449fc1b..87ddf9e6367 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_durable.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_agent_durable.py
@@ -14,12 +14,14 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Example DAGs demonstrating durable execution with AgentOperator and @task.agent."""
+"""Example Dags demonstrating durable execution with AgentOperator, @task.agent and a plain @task."""
from __future__ import annotations
from datetime import timedelta
+from airflow.providers.common.ai.durable import AirflowDurability
+from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
from airflow.providers.common.ai.operators.agent import AgentOperator
from airflow.providers.common.compat.sdk import dag, task
@@ -29,7 +31,7 @@ except Exception:
SQLToolset = None # type: ignore[assignment,misc]
# ---------------------------------------------------------------------------
-# 1. Durable AgentOperator: resumes from last model call on retry
+# 1. Durable AgentOperator: a retry replays the steps the failed attempt completed
# ---------------------------------------------------------------------------
@@ -38,7 +40,7 @@ if SQLToolset is not None:
@dag(default_args={"retries": 3, "retry_delay": timedelta(seconds=30)}, tags=["example"])
def example_agent_durable_operator():
- """Agent with durable execution -- resumes from the last model call on retry."""
+ """Agent with durable execution: a retry replays the steps the failed attempt completed."""
AgentOperator(
task_id="durable_analyst",
prompt="What are the top 5 customers by order count?",
@@ -91,3 +93,30 @@ if SQLToolset is not None:
# [END howto_decorator_agent_durable]
example_agent_durable_decorator()
+
+
+# ---------------------------------------------------------------------------
+# 3. A pydantic-ai agent in a plain @task
+# ---------------------------------------------------------------------------
+
+
+# [START howto_capability_airflow_durability]
+if SQLToolset is not None:
+
+ @dag(default_args={"retries": 3, "retry_delay": timedelta(seconds=30)}, tags=["example"])
+ def example_airflow_durability_capability():
+ @task
+ def analyze(question: str) -> str:
+ agent = PydanticAIHook(llm_conn_id="pydanticai_default").create_agent(
+ name="orders_analyst",
+ instructions="You are a data analyst. Use tools to answer questions.",
+ toolsets=[SQLToolset(db_conn_id="postgres_default", allowed_tables=["orders"])],
+ capabilities=[AirflowDurability()],
+ )
+ return agent.run_sync(question).output
+
+ analyze("What was our total revenue last month?")
+
+ # [END howto_capability_airflow_durability]
+
+ example_airflow_durability_capability()
diff --git a/providers/common/ai/src/airflow/providers/common/ai/exceptions.py b/providers/common/ai/src/airflow/providers/common/ai/exceptions.py
index 2f9c61f9f4b..fff7281b5c4 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/exceptions.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/exceptions.py
@@ -53,6 +53,16 @@ class UnsupportedToolDeferralError(AirflowFailException):
"""
+class DurableJournalError(AirflowFailException):
+ """
+ Raised when durable execution cannot record or replay a step.
+
+ A step's result is not JSON-serializable, or the journal holds a result this version
+ cannot replay. The same step fails the same way on every attempt, so the task fails
+ without retrying.
+ """
+
+
class LLMFileAnalysisError(ValueError):
"""Base class for file-analysis validation errors."""
diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
index 52038818025..0da729eedcd 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py
@@ -18,7 +18,6 @@
from __future__ import annotations
-import collections
import copy
import hashlib
import json
@@ -36,6 +35,8 @@ from pydantic_ai.messages import ModelMessagesTypeAdapter
from pydantic_ai.toolsets.abstract import AbstractToolset
from pydantic_ai.usage import RunUsage
+from airflow.providers.common.ai.durable.capability import DURABLE_UNIT_TOOLSETS, AirflowDurability
+from airflow.providers.common.ai.durable.journal import DurableJournal, build_task_storage, journal_scope
from airflow.providers.common.ai.exceptions import (
ToolApprovalAlreadyRequestedError,
ToolApprovalError,
@@ -103,9 +104,6 @@ if TYPE_CHECKING:
from pydantic_ai.usage import UsageLimits
from airflow.providers.common.ai.durable.base import DurableStorageProtocol
- from airflow.providers.common.ai.durable.caching_model import CachingModel
- from airflow.providers.common.ai.durable.replay_usage import ReplayUsageLedger
- from airflow.providers.common.ai.durable.step_counter import DurableStepCounter
from airflow.providers.common.compat.sdk import TaskInstanceKey
from airflow.sdk import Context
from airflow.sdk.execution_time.context import TaskStateStoreAccessor
@@ -316,31 +314,30 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa
:ref:`the cross-attempt usage budget <agent-usage-budget>` for how
it is persisted, reset, and how ``durable`` replay and HITL
regeneration interact with it.
- :param durable: Experimental. When ``True``, enables step-level caching of model
- responses and tool results for durable execution. On retry, cached
- steps are replayed instead of re-executing. Each cached step is
- verified against the current request before replay: if the prompt,
- model, settings, tools, or message history changed since the failed
- attempt, the affected steps re-run live (with a warning) instead of
- replaying stale results. Default ``False``. A replayed step adds
- nothing to the usage counted against ``usage_limits`` or reported in
- the ``usage`` XCom -- not its request, tokens, cost, or tool calls --
- so every attempt counts only the model and tool calls it actually
- makes. This holds the same way after clearing a failed task
- instance: it starts a fresh budget but keeps the durable cache its
- attempts left behind, and whatever the rerun replays from that cache
- is free.
- On Airflow >= 3.3 the cache is kept in the AIP-103 task state store, so
+ :param durable: Experimental. When ``True``, attaches
+ :class:`~airflow.providers.common.ai.durable.AirflowDurability`, which
+ records each step the agent completes -- model responses, tool calls from
+ any toolset or capability, tool discovery, and other capabilities'
+ ``@durable_operation`` methods -- in a journal kept for the task instance.
+ On retry the agent runs again from the start and the steps the previous
+ attempt completed are replayed from the journal instead of running again.
+ Each step is replayed only if it matches what the previous attempt did at
+ the same point: if the prompt, model, settings, tools, tool arguments or
+ message history changed, that step and every step after it run live (with
+ a warning). Default ``False``. A replayed step adds nothing to the usage
+ counted against ``usage_limits`` or reported in the ``usage`` XCom -- not
+ its request, tokens, cost, or tool calls -- so every attempt counts only
+ the model and tool calls it actually makes. This holds the same way after
+ clearing a failed task instance: it starts a fresh budget but keeps the
+ journal its attempts left behind, and whatever the rerun replays from it
+ is free. The agent's ``run_id`` (and its ``conversation_id``, when there is
+ no message history) stays the first attempt's on every retry.
+ On Airflow >= 3.3 the journal is kept in the AIP-103 task state store, so
no extra configuration is needed. On older Airflow versions it is persisted to
ObjectStorage and requires ``[common.ai] durable_cache_path`` to be set.
- Tools are durably cached when provided via ``toolsets=`` or via a
- concrete pydantic-ai ``Toolset`` capability. Tools reaching the agent
- through any *other* capability -- ``MCP``, ``PrefixTools``,
- ``CombinedCapability``, a ``Toolset`` backed by a callable factory, or
- capabilities loaded from a ``spec_file`` -- are not cached and re-run on
- retry; put tools you need replayed in ``toolsets=``. Provider-native
- capabilities such as ``WebSearch`` and ``Thinking`` execute inside the
- model call and are covered by model-response caching.
+ Toolsets pydantic-ai runs as durable steps (``FunctionToolset``,
+ ``MCPToolset``, ``DynamicToolset``) need a unique ``id``, and so do
+ capabilities with ``@durable_operation`` methods.
Cannot be combined with a ``SandboxToolset`` (raises), attached or
not: a replayed tool result describes a workspace state the replay did
not reproduce, and the first call that misses the cache runs against
@@ -509,12 +506,8 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa
self.durable = durable
self.cache_prompt = cache_prompt
- # Populated per run in ``execute`` when durable=True. Declared here so
- # ``_build_agent`` -- also reached via ``regenerate_with_feedback``
- # outside ``execute`` -- can read them unconditionally.
- self._durable_storage: DurableStorageProtocol | None = None
- self._durable_counter: DurableStepCounter | None = None
- self._replay_usage: ReplayUsageLedger | None = None
+ # The task instance's durable journal, set per run in ``execute`` when durable=True.
+ self._durable_journal: DurableJournal | None = None
# Populated in ``execute``; also read (and, if unset, lazily initialized) by
# ``regenerate_with_feedback`` outside ``execute``, which is why they need a
@@ -538,14 +531,18 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa
raise ValueError("durable=True and enable_hitl_review=True cannot be used together.")
if durable and _contains_code_mode(self._declared_capabilities):
- # Durable replay caches individual model/tool steps via CachingModel /
- # CachingToolset and a shared step counter that assumes a stable call
- # order across runs. Code mode collapses tools into one ``run_code``
- # tool and lets the model emit arbitrary Python, so step counts and
- # ordering can differ between the original run and a retry, breaking
- # replay. Reject the combination rather than silently mis-replaying.
+ # Code mode collapses tools into one ``run_code`` tool and lets the model emit
+ # arbitrary Python, whose nested tool calls the journal does not see as steps of
+ # their own. Reject the combination rather than replay a partial ``run_code``.
raise ValueError("durable=True cannot be used with a CodeMode capability.")
+ if durable and any(
+ isinstance(capability, AirflowDurability) for capability in self._declared_capabilities
+ ):
+ raise ValueError(
+ "durable=True attaches AirflowDurability itself; pass durable=True or the capability, not both."
+ )
+
if message_history is not None and enable_hitl_review:
# The post-review transcript is not recoverable today (run_hitl_review
# returns only the final string), so emitting the pre-review transcript
@@ -556,6 +553,9 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa
if durable or enable_hitl_review:
self._reject_sandbox_without_continuity(durable=durable, enable_hitl_review=enable_hitl_review)
+ if durable:
+ self._require_durable_toolset_ids()
+
self.enable_hitl_review = enable_hitl_review
self.max_hitl_iterations = max_hitl_iterations
self.hitl_timeout = hitl_timeout
@@ -575,6 +575,38 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa
tool_approval_assigned_users, param="tool_approval_assigned_users"
)
+ def _require_durable_toolset_ids(self) -> None:
+ """
+ Refuse a toolset whose steps durable replay would have no name for.
+
+ pydantic-ai records each call of a ``FunctionToolset``, ``MCPToolset`` or
+ ``DynamicToolset`` under the toolset's ``id``, and refuses to run one without it.
+ The id is what a retry matches recorded calls by, so it has to come from the Dag
+ author: one derived from the toolset's position would point a retry at another
+ toolset's results as soon as the list is reordered. Checked here so the Dag fails
+ to load rather than at the first run.
+ """
+ missing = {
+ type(leaf).__name__
+ for toolset in self._declared_toolsets()
+ for leaf in iter_toolsets(toolset)
+ if isinstance(leaf, DURABLE_UNIT_TOOLSETS) and leaf.id is None
+ }
+ # A function that builds a toolset per run becomes a DynamicToolset with no id.
+ if any(
+ not isinstance(toolset, AbstractToolset)
+ for toolset in (*(self.toolsets or []), *(self.agent_params.get("toolsets") or []))
+ ):
+ missing.add("DynamicToolset")
+ if kinds := sorted(missing):
+ raise ValueError(
+ f"durable=True needs a unique id on every {' and '.join(kinds)}, such as "
+ f"{kinds[0]}(..., id='orders'): durable replay records each tool call under it. "
+ "Set it on the toolset itself, and pass a function that builds a toolset as "
+ "DynamicToolset(function, id=...); the id of a Toolset capability around a toolset "
+ "does not reach it."
+ )
+
def _reject_sandbox_without_continuity(self, *, durable: bool, enable_hitl_review: bool) -> None:
"""
Refuse a ``SandboxToolset`` under a feature that assumes the sandbox outlives the run.
@@ -720,13 +752,8 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa
# pydantic-ai wraps capability hooks in list order, so merging the two lists
# would pick an order the Dag author never wrote down.
raise ValueError("Pass capabilities either as capabilities=... or in agent_params, not both.")
- storage = self._durable_storage
- counter = self._durable_counter
if self.toolsets:
- # Innermost, so the durable cache only ever stores masked results.
toolsets: list[AbstractToolset] = [ensure_masked(ts) for ts in self.toolsets]
- if self.durable and storage is not None and counter is not None:
- toolsets = self._build_durable_toolsets(toolsets, storage, counter)
extra_kwargs["toolsets"] = toolsets
elif extra_kwargs.get("toolsets"):
extra_kwargs["toolsets"] = [ensure_masked(ts) for ts in extra_kwargs["toolsets"]]
@@ -734,11 +761,9 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa
_replace_capability_toolset(capability, ensure_masked)
for capability in self.capabilities or passed_through or []
]
- if self.durable and storage is not None and counter is not None:
- # Tools supplied through a ``Toolset`` capability bypass the
- # ``toolsets=`` wrapping above, so their results would re-execute on
- # every retry instead of replaying; wrap their inner toolset too.
- capabilities = self._build_durable_capabilities(capabilities, storage, counter)
+ if self.durable:
+ # Names the agent's steps in the journal; the task id is stable across retries.
+ capabilities.append(AirflowDurability(name=extra_kwargs.get("name") or self.task_id))
if self.cache_prompt:
capabilities.append(PromptCaching())
if self.enable_tool_logging and not _contains_capability(capabilities, ToolLoggingCapability):
@@ -812,119 +837,16 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa
if leaf.id is not None
]
- def _build_durable_toolsets(
- self, toolsets: list[AbstractToolset], storage: DurableStorageProtocol, counter: DurableStepCounter
- ) -> list[AbstractToolset]:
- """Wrap each toolset with CachingToolset for durable execution."""
- from airflow.providers.common.ai.durable.caching_toolset import CachingToolset
-
- return [
- CachingToolset(wrapped=ts, storage=storage, counter=counter, replay_usage=self._replay_usage)
- for ts in toolsets
- ]
-
- def _build_durable_capabilities(
- self, capabilities: list[Any], storage: DurableStorageProtocol, counter: DurableStepCounter
- ) -> list[Any]:
- """
- Wrap toolsets provided via a pydantic-ai ``Toolset`` capability for durable replay.
-
- Tools reaching the agent through ``capabilities=[Toolset(ts)]`` bypass the
- operator's ``toolsets=`` list, so the ``CachingToolset`` applied in
- :meth:`_build_durable_toolsets` never sees them and their results
- re-execute on every retry instead of replaying. Wrap each ``Toolset``
- capability's inner toolset with the same ``CachingToolset``, preserving
- the capability's other fields. Non-``Toolset`` capabilities pass through
- unchanged, as does a ``Toolset`` holding a callable factory rather than a
- concrete toolset (only a concrete toolset can be wrapped here).
- """
- from airflow.providers.common.ai.durable.caching_toolset import CachingToolset
-
- rewrapped: list[Any] = []
- for capability in capabilities:
- # ``Toolset.toolset`` can be a concrete toolset or a callable factory
- # resolved per run; only a concrete toolset can be wrapped here.
- toolset = _resolve_capability_toolset(capability)
- if toolset is not None:
- cached = CachingToolset(
- wrapped=toolset,
- storage=storage,
- counter=counter,
- replay_usage=self._replay_usage,
- )
- rewrapped.append(replace(capability, toolset=cached))
- continue
- if isinstance(capability, Toolset):
- # The toolset is a callable factory resolved per run, so there is
- # no concrete toolset to wrap; its results won't be cached for
- # replay. Warn so durable users aren't silently surprised on retry.
- self.log.warning(
- "durable=True: tools from a Toolset capability backed by a callable "
- "factory are not cached for replay; pass the toolset via `toolsets=` "
- "for durability."
- )
- rewrapped.append(capability)
- return rewrapped
-
- def _log_durable_summary(self, counter: DurableStepCounter) -> None:
- """
- Log what this attempt replayed and cached, and which steps it could not cache.
-
- A step whose cache write was skipped (a tool result that is not
- JSON-serializable, a store write that fails) ran live but is not cached,
- so a retry runs it again. For a tool with side effects the side effect
- repeats, so those tools are named rather than counted as cached.
- """
- self.log.info(
- "Durable: replayed %d cached steps (%d model, %d tool), cached %d new steps (%d model, %d tool)",
- counter.replayed_model + counter.replayed_tool,
- counter.replayed_model,
- counter.replayed_tool,
- counter.cached_model + counter.cached_tool,
- counter.cached_model,
- counter.cached_tool,
- )
- if counter.skipped_tools:
- calls = collections.Counter(counter.skipped_tools)
- self.log.warning(
- "Durable: %d tool results were not cached, and a retry runs them again: %s",
- len(counter.skipped_tools),
- ", ".join(name if n == 1 else f"{name} (x{n})" for name, n in calls.items()),
- )
- if counter.skipped_model:
- self.log.warning(
- "Durable: %d model responses were not cached, and a retry re-runs them "
- "and every step after the first of them",
- counter.skipped_model,
- )
-
def _build_durable_storage(self, context: Context) -> DurableStorageProtocol:
"""
- Return the durable storage backend for the current task instance.
+ Return where the durable journal of this task instance is stored.
- On Airflow >= 3.3 durable steps are cached in the AIP-103 task state
- store, which handles persistence and large-value offload natively, so no
- ``[common.ai] durable_cache_path`` is required. On older Airflow versions, fall back
- to the ObjectStorage backend configured via ``durable_cache_path``.
+ The task state store on Airflow >= 3.3; ObjectStorage under ``[common.ai]
+ durable_cache_path`` on older versions.
"""
- if AIRFLOW_V_3_3_PLUS:
- # Imported lazily: NEVER_EXPIRE and the task state store accessor do
- # not exist on Airflow versions before 3.3.
- from airflow.providers.common.ai.durable.task_state_store import TaskStateStoreDurableStorage
-
- return TaskStateStoreDurableStorage(context["task_state_store"])
-
- from airflow.providers.common.ai.durable.storage import DurableStorage
+ return build_task_storage(context)
- ti = context["task_instance"]
- return DurableStorage(
- 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,
- )
-
- def _report_failed_run(self, context: Context, run_usage: RunUsage) -> None:
+ def _report_failed_run(self, context: Context, run_usage: RunUsage, *, run_id: str | None = None) -> None:
"""
Log and XCom-push the usage a failed attempt incurred before it raised.
@@ -956,7 +878,8 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa
return
ti = context["task_instance"]
try:
- ti.xcom_push(key="run_id", value=make_task_instance_run_key(ti))
+ # The id the agent ran with, which is the durable run's on a durable retry.
+ ti.xcom_push(key="run_id", value=run_id or make_task_instance_run_key(ti))
except Exception:
self.log.warning("Failed to push run_id XCom for the failed run", exc_info=True)
if attempt_usage is not None:
@@ -971,7 +894,6 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa
prompt: Any,
*,
run_usage: RunUsage,
- caching_model: CachingModel | None = None,
**run_kwargs: Any,
) -> tuple[Any, RunUsage]:
"""
@@ -983,21 +905,11 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa
None`` because it is set lazily, but every caller of this method has already
ensured it is a real ``RunUsage`` by the time it gets here.
- With ``durable=True``, the replay ledger's unused credits are given back before
- the total is persisted (see ``ReplayUsageLedger.settle``).
+ With ``durable=True``, ``AirflowDurability`` keeps replayed steps out of
+ ``run_usage`` and settles its credits before the run returns, so neither reaches
+ the persisted total.
"""
- return self._run_tracked(
- agent,
- prompt,
- run_usage=run_usage,
- before_run=caching_model.credit_first_replay if caching_model is not None else None,
- **run_kwargs,
- )
-
- def _settle_tracked_usage(self) -> None:
- if self._replay_usage is not None:
- # A replay credit the run never used must not reach the persisted total.
- self._replay_usage.settle()
+ return self._run_tracked(agent, prompt, run_usage=run_usage, **run_kwargs)
def _run_and_report_on_failure(
self,
@@ -1005,18 +917,15 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa
agent: Agent[Any, Any],
run_usage: RunUsage,
run_kwargs: dict[str, Any],
- caching_model: CachingModel | None = None,
) -> tuple[Any, RunUsage]:
"""Run ``self.prompt`` via ``_run_agent_tracked``, reporting usage-at-failure on any raise."""
try:
- if caching_model is not None:
- with agent.override(model=caching_model):
- return self._run_agent_tracked(
- agent, self.prompt, run_usage=run_usage, caching_model=caching_model, **run_kwargs
- )
+ if self._durable_journal is not None:
+ with journal_scope(self._durable_journal):
+ return self._run_agent_tracked(agent, self.prompt, run_usage=run_usage, **run_kwargs)
return self._run_agent_tracked(agent, self.prompt, run_usage=run_usage, **run_kwargs)
except BaseException:
- self._report_failed_run(context, run_usage)
+ self._report_failed_run(context, run_usage, run_id=run_kwargs["run_id"])
raise
def execute(self, context: Context) -> Any:
@@ -1038,9 +947,7 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa
self._delete_approval_transcript(store)
ti = context["task_instance"]
- self._durable_storage = None
- self._durable_counter = None
- self._replay_usage = None
+ self._durable_journal = None
# Reads the state store before the expensive setup below (_build_agent, durable
# storage). None on < 3.3 or usage_limits=None -- see _build_usage_budget.
self._usage_budget = self._build_usage_budget(context, usage_limits, ti=ti)
@@ -1052,13 +959,7 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa
self._run_usage_base = copy_run_usage(run_usage)
if self.durable:
- from airflow.providers.common.ai.durable.replay_usage import ReplayUsageLedger
- from airflow.providers.common.ai.durable.step_counter import DurableStepCounter
-
- self._durable_storage = self._build_durable_storage(context)
- self._durable_counter = DurableStepCounter()
- # Built before _build_agent so the CachingToolset wrappers it creates share it.
- self._replay_usage = ReplayUsageLedger(run_usage=run_usage, usage_limits=usage_limits)
+ self._durable_journal = DurableJournal(self._build_durable_storage(context))
agent = self._build_agent()
@@ -1073,32 +974,26 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa
history = self._resolve_message_history()
if history is not None:
run_kwargs["message_history"] = history
+ if self._durable_journal is not None:
+ # With durable=True the run keeps the first attempt's id on every retry, because
+ # replayed steps carry it: capabilities that key their own state on the run id
+ # (pydantic-ai-harness SpendLimits, StepPersistence) would otherwise look up the
+ # retry's id in the previous attempt's records. Attempts still differ in their
+ # spans' task instance id and try number.
+ run_kwargs["run_id"] = self._durable_journal.get_run_id(default=run_kwargs["run_id"])
+ if not history:
+ # Without a history to inherit one from, pydantic-ai starts a new conversation
+ # on every attempt; keep one for the durable run, for the same reason.
+ run_kwargs["conversation_id"] = run_kwargs["run_id"]
- storage = self._durable_storage
- counter = self._durable_counter
- caching_model: CachingModel | None = None
# A killed run raises RunCancelled (see run_agent_sync), which propagates to fail the
- # task. The durable cache cleanup below is skipped on the raise, preserving it for retry.
- if self.durable and storage is not None and counter is not None:
- from pydantic_ai.models import infer_model
-
- from airflow.providers.common.ai.durable.caching_model import CachingModel
-
- if agent.model is None:
- raise ValueError("Agent model must be set when durable=True")
- resolved_model = infer_model(agent.model)
- caching_model = CachingModel(
- resolved_model, storage=storage, counter=counter, replay_usage=self._replay_usage
- )
-
+ # task. The journal cleanup in _complete_run is skipped on the raise, keeping it for the retry.
try:
- result, attempt_usage = self._run_and_report_on_failure(
- context, agent, run_usage, run_kwargs, caching_model
- )
+ result, attempt_usage = self._run_and_report_on_failure(context, agent, run_usage, run_kwargs)
finally:
# Also on a raise: the failed attempt is the one Airflow retries.
- if counter is not None:
- self._log_durable_summary(counter)
+ if self._durable_journal is not None:
+ self._durable_journal.log_summary(self.log)
return self._complete_run(context, result, attempt_usage=attempt_usage)
def _complete_run(self, context: Context, result: Any, *, attempt_usage: RunUsage) -> Any:
@@ -1142,13 +1037,13 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa
if self._serialize_model_output and isinstance(output, BaseModel):
output = output.model_dump()
- # Clean up the durable cache only after the run and every post-run step
+ # Clean up the durable journal only after the run and every post-run step
# that can still fail (the run-metadata and message-history XCom pushes
# above and output serialization) has succeeded. Cleaning up earlier and
- # then raising would leave the Airflow retry with an empty cache,
+ # then raising would leave the Airflow retry with an empty journal,
# re-executing every already-completed model and tool step.
- if self._durable_storage is not None:
- self._durable_storage.cleanup()
+ if self._durable_journal is not None:
+ self._durable_journal.cleanup()
self._clear_usage_budget(context)
return output
diff --git a/providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py b/providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py
index ecbff99d2af..f2713c0f090 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/hook.py
@@ -264,7 +264,7 @@ class HookToolset(AirflowToolset):
def _validate_args(self, bytes_params: frozenset[str], ctx: RunContext[Any], /, **tool_args: Any) -> None:
self._refuse_pinned(ctx, **tool_args)
# Only checked here: the decoded bytes are not kept, so what the toolset chain passes
- # on, and CachingToolset fingerprints, stays JSON.
+ # on, and the durable journal fingerprints, stays JSON.
_decode_bytes_args(tool_args, bytes_params)
def _refuse_pinned(self, ctx: RunContext[Any], /, **tool_args: Any) -> None:
@@ -288,7 +288,7 @@ class HookToolset(AirflowToolset):
# The framework bridges validate arguments without args_validator_func, so check again here.
self._refuse_pinned(ctx, **tool_args)
# Decoded here rather than in the args validator: validated arguments travel the whole
- # toolset chain, and CachingToolset fingerprints them with a plain json.dumps, so bytes
+ # toolset chain, and the durable journal fingerprints them with a plain json.dumps, so bytes
# upstream of this point would make every binary call unverifiable on durable replay.
tool_args = _decode_bytes_args(tool_args, self._bytes_params[method_name])
# A copy per call, so a method that modifies an argument it is given cannot change the pin.
diff --git a/providers/common/ai/src/airflow/providers/common/ai/toolsets/managed_agent.py b/providers/common/ai/src/airflow/providers/common/ai/toolsets/managed_agent.py
index dc11bf6007c..1a811aec23e 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/managed_agent.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/managed_agent.py
@@ -96,9 +96,9 @@ class BaseManagedAgentToolset(AirflowToolset):
rejects a request. ``0`` turns the first rejection into a hard error.
"""
- #: Whether ``durable=True`` may replay a completed invocation from its cache instead of
+ #: Whether ``durable=True`` may replay a completed invocation from its journal instead of
#: re-invoking. Off by default because a managed agent may act on systems Airflow cannot
- #: observe, so replaying a cached answer could skip a side effect. Read-only agents may
+ #: observe, so replaying a recorded answer could skip a side effect. Read-only agents may
#: opt in.
replayable: bool = False
@@ -288,7 +288,7 @@ class ManagedAgentToolset(BaseManagedAgentToolset):
``None`` means whatever the vendor client defaults to, which for Agent Engine is no
deadline at all.
:param max_retries: See :class:`BaseManagedAgentToolset`.
- :param replayable: Whether the durable cache may replay a completed call. Only set it
+ :param replayable: Whether durable execution may replay a completed call. Only set it
for an agent that is read-only.
:param vendor_options: Sent with every request as ``ManagedAgentRequest.vendor_options``,
for per-agent settings the vendor hook accepts there (Agent Engine's ``class_method``,
diff --git a/providers/common/ai/src/airflow/providers/common/ai/utils/toolset_base.py b/providers/common/ai/src/airflow/providers/common/ai/utils/toolset_base.py
index 0b2da926221..10caa772cc4 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/utils/toolset_base.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/utils/toolset_base.py
@@ -191,6 +191,11 @@ class AirflowToolset(AbstractToolset[Any]):
# A subclass that takes ``max_retries`` stores it here; ``None`` follows the run.
_max_retries: int | None = None
+ # Whether durable execution may replay a completed call on retry. A toolset whose
+ # calls act on a system Airflow cannot observe sets it to ``False``, so its calls
+ # run again on every attempt; see :ref:`durable-execution`.
+ replayable: bool = True
+
def _get_tool_max_retries(self, ctx: RunContext[Any]) -> int:
"""
Return how many times the model may correct a failed call to this toolset's tools.
diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_step_counter.py b/providers/common/ai/tests/unit/common/ai/durable/conftest.py
similarity index 55%
rename from providers/common/ai/tests/unit/common/ai/durable/test_step_counter.py
rename to providers/common/ai/tests/unit/common/ai/durable/conftest.py
index cd074989425..95a77a7abab 100644
--- a/providers/common/ai/tests/unit/common/ai/durable/test_step_counter.py
+++ b/providers/common/ai/tests/unit/common/ai/durable/conftest.py
@@ -16,23 +16,11 @@
# under the License.
from __future__ import annotations
-from airflow.providers.common.ai.durable.step_counter import DurableStepCounter
+import pytest
+from unit.common.ai.durable.memory_storage import MemoryStorage
-class TestDurableStepCounter:
- def test_starts_at_zero(self):
- counter = DurableStepCounter()
- assert counter.next_step() == 0
- def test_increments_monotonically(self):
- counter = DurableStepCounter()
- assert counter.next_step() == 0
- assert counter.next_step() == 1
- assert counter.next_step() == 2
-
- def test_total_steps_tracks_count(self):
- counter = DurableStepCounter()
- assert counter.total_steps == 0
- counter.next_step()
- counter.next_step()
- assert counter.total_steps == 2
+@pytest.fixture
+def memory_storage() -> MemoryStorage:
+ return MemoryStorage()
diff --git a/providers/common/ai/tests/unit/common/ai/durable/memory_storage.py b/providers/common/ai/tests/unit/common/ai/durable/memory_storage.py
new file mode 100644
index 00000000000..072223e9deb
--- /dev/null
+++ b/providers/common/ai/tests/unit/common/ai/durable/memory_storage.py
@@ -0,0 +1,58 @@
+# 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.
+"""An in-memory durable journal storage, for tests that simulate task attempts."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterable
+from typing import Any
+
+from airflow.providers.common.ai.durable.base import build_step_key
+
+
+class MemoryStorage:
+ """
+ A ``DurableStorageProtocol`` kept in a dict, shared by the attempts of one test.
+
+ Entries go through a JSON round-trip on the way in, as they do in both real backends,
+ so a test sees what a retry would read back rather than the live objects.
+ """
+
+ def __init__(self) -> None:
+ self.entries: dict[str, dict[str, Any]] = {}
+ self.refuse: set[str] = set()
+
+ def save_step(self, key: str, entry: dict[str, Any]) -> bool:
+ if key in self.refuse:
+ return False
+ self.entries[key] = json.loads(json.dumps(entry))
+ return True
+
+ def load_step(self, key: str) -> dict[str, Any] | None:
+ entry = self.entries.get(key)
+ return json.loads(json.dumps(entry)) if entry is not None else None
+
+ def delete_steps(self, keys: Iterable[str]) -> None:
+ for key in keys:
+ self.entries.pop(key, None)
+
+ def steps(self, run: int = 0) -> list[dict[str, Any]]:
+ """Return the entries of one run in position order."""
+ prefix = build_step_key(run, 0).removesuffix("0")
+ positions = sorted(int(key.removeprefix(prefix)) for key in self.entries if key.startswith(prefix))
+ return [self.entries[build_step_key(run, position)] for position in positions]
diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_base.py b/providers/common/ai/tests/unit/common/ai/durable/test_base.py
index 20826fd277c..ed164bad34a 100644
--- a/providers/common/ai/tests/unit/common/ai/durable/test_base.py
+++ b/providers/common/ai/tests/unit/common/ai/durable/test_base.py
@@ -18,7 +18,13 @@ from __future__ import annotations
import pytest
-from airflow.providers.common.ai.durable.base import DURABLE_KEY_PREFIX, DurableStorageProtocol
+from airflow.providers.common.ai.durable.base import (
+ DURABLE_KEY_PREFIX,
+ RUN_ID_KEY,
+ RUNS_KEY,
+ DurableStorageProtocol,
+ build_step_key,
+)
from airflow.providers.common.ai.durable.storage import DurableStorage
from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS
@@ -30,6 +36,13 @@ def test_key_prefix_is_a_single_path_segment():
assert "/" not in DURABLE_KEY_PREFIX
+def test_keys_are_reserved_and_distinct_per_run_and_position():
+ keys = {build_step_key(0, 1), build_step_key(1, 0), build_step_key(0, 10), RUN_ID_KEY, RUNS_KEY}
+
+ assert len(keys) == 5
+ assert all(key.startswith(DURABLE_KEY_PREFIX) and "/" not in key for key in keys)
+
+
class TestRealBackendsSatisfyProtocol:
"""The two shipped backends must stay in step with the protocol -- ``issubclass``
on a methods-only ``runtime_checkable`` protocol catches a renamed or dropped method."""
diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_caching_model.py b/providers/common/ai/tests/unit/common/ai/durable/test_caching_model.py
deleted file mode 100644
index 5091a59c457..00000000000
--- a/providers/common/ai/tests/unit/common/ai/durable/test_caching_model.py
+++ /dev/null
@@ -1,432 +0,0 @@
-# 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
-
-from decimal import Decimal
-from unittest.mock import AsyncMock, MagicMock, patch
-
-import pytest
-from pydantic_ai.exceptions import UsageLimitExceeded
-from pydantic_ai.messages import ModelResponse, TextPart, ToolCallPart
-from pydantic_ai.models import ModelRequestParameters
-from pydantic_ai.tools import ToolDefinition
-from pydantic_ai.usage import RequestUsage, RunUsage, UsageLimits
-
-from airflow.providers.common.ai.durable.base import DURABLE_KEY_PREFIX as P, DurableStorageProtocol
-from airflow.providers.common.ai.durable.caching_model import CachingModel
-from airflow.providers.common.ai.durable.fingerprint import fingerprint_model_request
-from airflow.providers.common.ai.durable.replay_usage import ReplayUsageLedger
-from airflow.providers.common.ai.durable.step_counter import DurableStepCounter
-
-from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS
-
-
-@pytest.fixture
-def mock_storage():
- storage = MagicMock(spec=DurableStorageProtocol)
- storage.load_model_response.return_value = (None, None)
- storage.save_model_response.return_value = True
- return storage
-
-
-@pytest.fixture
-def counter():
- return DurableStepCounter()
-
-
-@pytest.fixture
-def mock_model():
- model = MagicMock()
- model.model_name = "test-model"
- model.system = "test"
- model.profile = MagicMock()
- model.settings = None
- # CachingModel fingerprints the prepared request; identity keeps prepared == raw.
- model.prepare_request = lambda settings, params: (settings, params)
- return model
-
-
-@pytest.fixture(autouse=True)
-def _patch_infer_model():
- """Prevent WrapperModel.__init__ from resolving the mock as a real model."""
- with patch("pydantic_ai.models.wrapper.infer_model", side_effect=lambda m: m):
- yield
-
-
-@pytest.fixture
-def sample_response():
- return ModelResponse(parts=[TextPart(content="Hello!")])
-
-
-def request_fingerprint(messages=(), settings=None, params=None):
- """Fingerprint matching what CachingModel computes for the mock model."""
- return fingerprint_model_request(
- "test:test-model", list(messages), settings, params or ModelRequestParameters()
- )
-
-
-class TestCachingModelCacheHit:
- @pytest.mark.asyncio
- async def test_returns_cached_response_without_calling_model(
- self, mock_model, mock_storage, counter, sample_response
- ):
- mock_storage.load_model_response.return_value = (sample_response, request_fingerprint())
- caching = CachingModel(mock_model, storage=mock_storage, counter=counter)
-
- result = await caching.request([], None, ModelRequestParameters())
-
- assert result is sample_response
- mock_model.request.assert_not_called()
- mock_storage.load_model_response.assert_called_once_with(f"{P}model_step_0")
-
- @pytest.mark.asyncio
- async def test_advances_counter_on_cache_hit(self, mock_model, mock_storage, counter, sample_response):
- mock_storage.load_model_response.return_value = (sample_response, request_fingerprint())
- caching = CachingModel(mock_model, storage=mock_storage, counter=counter)
-
- await caching.request([], None, ModelRequestParameters())
-
- assert counter.total_steps == 1
-
- @pytest.mark.asyncio
- async def test_replay_returns_cached_response_unchanged_and_nets_out_its_usage(
- self, mock_model, mock_storage, counter
- ):
- """The cached response is returned as-is -- it flows into the message history the
- next step's fingerprint hashes, so changing its usage would break every later
- replay -- and the ledger subtracts exactly what the graph adds back after this
- call returns (``requests += 1`` and ``incr(response.usage)``)."""
- cached_original = ModelResponse(
- parts=[TextPart(content="Hello!")],
- usage=RequestUsage(input_tokens=100, output_tokens=50, cost=Decimal("0.10")),
- )
- mock_storage.load_model_response.side_effect = lambda key: (
- (cached_original, request_fingerprint()) if key == f"{P}model_step_0" else (None, None)
- )
- seed = RunUsage(requests=1, input_tokens=100, output_tokens=50, cost=Decimal("0.10"))
- run_usage = RunUsage(requests=1, input_tokens=100, output_tokens=50, cost=Decimal("0.10"))
- ledger = ReplayUsageLedger(run_usage=run_usage, usage_limits=None)
- caching = CachingModel(mock_model, storage=mock_storage, counter=counter, replay_usage=ledger)
-
- result = await caching.request([], None, ModelRequestParameters())
- run_usage.requests += 1
- run_usage.incr(result.usage)
-
- assert result is cached_original
- assert result.usage == RequestUsage(input_tokens=100, output_tokens=50, cost=Decimal("0.10"))
- assert run_usage == seed
-
- @pytest.mark.asyncio
- async def test_replay_prices_an_unpriced_cached_response_before_subtracting(
- self, mock_model, mock_storage, counter
- ):
- """A live response is cached before pydantic-ai prices it, so the cached copy has
- ``cost=None``. The replay must price it the way the graph is about to, or the
- graph adds a cost the ledger never subtracted."""
- cached = ModelResponse(
- parts=[TextPart(content="Hello!")],
- usage=RequestUsage(input_tokens=100_000, output_tokens=10_000),
- model_name="gpt-4o",
- provider_name="openai",
- )
- expected_cost = cached.cost().total_price
- assert expected_cost > 0
- mock_storage.load_model_response.side_effect = lambda key: (
- (cached, request_fingerprint()) if key == f"{P}model_step_0" else (None, None)
- )
- run_usage = RunUsage(requests=1, cost=expected_cost)
- caching = CachingModel(
- mock_model,
- storage=mock_storage,
- counter=counter,
- replay_usage=ReplayUsageLedger(run_usage=run_usage, usage_limits=None),
- )
-
- result = await caching.request([], None, ModelRequestParameters())
-
- assert result.usage.cost == expected_cost
- assert run_usage.cost == 0
-
- @pytest.mark.asyncio
- async def test_replay_without_ledger_returns_response_unchanged(self, mock_model, mock_storage, counter):
- """Without a ledger nothing is compensated, and the cached response's real usage
- is still returned unchanged (not zeroed) -- see the sibling test above for why."""
- cached = ModelResponse(parts=[TextPart(content="Hello!")], usage=RequestUsage(input_tokens=10))
- mock_storage.load_model_response.return_value = (cached, request_fingerprint())
- caching = CachingModel(mock_model, storage=mock_storage, counter=counter)
-
- result = await caching.request([], None, ModelRequestParameters())
-
- assert result is cached
- assert result.usage == RequestUsage(input_tokens=10)
-
-
-class TestCachingModelReplayCredits:
- @pytest.mark.asyncio
- async def test_credit_first_replay_lets_a_fully_replayable_retry_start_at_request_limit(
- self, mock_model, mock_storage, counter, sample_response
- ):
- """pydantic-ai checks ``requests >= request_limit`` before the first request reaches
- ``CachingModel``. A retry seeded with ``requests == request_limit`` whose first step
- is cached must pass that check, since the replay is free."""
- mock_storage.load_model_response.side_effect = lambda key: (
- (sample_response, request_fingerprint()) if key == f"{P}model_step_0" else (None, None)
- )
- limits = UsageLimits(request_limit=2)
- run_usage = RunUsage(requests=2)
- caching = CachingModel(
- mock_model,
- storage=mock_storage,
- counter=counter,
- replay_usage=ReplayUsageLedger(run_usage=run_usage, usage_limits=limits),
- )
-
- caching.credit_first_replay()
- limits.check_before_request(run_usage)
- await caching.request([], None, ModelRequestParameters())
- run_usage.requests += 1
-
- assert run_usage.requests == 2
-
- @pytest.mark.asyncio
- async def test_credited_request_that_runs_live_is_checked_again(
- self, mock_model, mock_storage, counter, sample_response
- ):
- """The credit assumed a free replay. When the fingerprint no longer matches, the
- request would be a real one past ``request_limit``, so it must not reach the model."""
- mock_storage.load_model_response.side_effect = lambda key: (
- (sample_response, "stale-fingerprint") if key == f"{P}model_step_0" else (None, None)
- )
- mock_model.request = AsyncMock(return_value=sample_response)
- limits = UsageLimits(request_limit=2)
- run_usage = RunUsage(requests=2)
- caching = CachingModel(
- mock_model,
- storage=mock_storage,
- counter=counter,
- replay_usage=ReplayUsageLedger(run_usage=run_usage, usage_limits=limits),
- )
-
- caching.credit_first_replay()
- limits.check_before_request(run_usage)
- with pytest.raises(UsageLimitExceeded, match="request_limit"):
- await caching.request([], None, ModelRequestParameters())
-
- mock_model.request.assert_not_called()
- assert run_usage.requests == 2
-
- @pytest.mark.asyncio
- async def test_replay_credits_cached_tool_steps_and_the_next_request(
- self, mock_model, mock_storage, counter
- ):
- """A replayed response with two function-tool calls credits the cached tool
- results at the next step indices (a missing one -- a call that raised -- is
- skipped) and the next model request; a call to a tool that is not a function
- tool takes no index and is not credited."""
- response = ModelResponse(
- parts=[
- ToolCallPart(tool_name="tool_a", args={}, tool_call_id="a"),
- ToolCallPart(tool_name="tool_b", args={}, tool_call_id="b"),
- ToolCallPart(tool_name="final_result", args={}, tool_call_id="o"),
- ]
- )
- params = ModelRequestParameters(
- function_tools=[ToolDefinition(name="tool_a"), ToolDefinition(name="tool_b")],
- output_tools=[ToolDefinition(name="final_result")],
- )
- models = {
- f"{P}model_step_0": (response, request_fingerprint(params=params)),
- f"{P}model_step_3": (response, "x"),
- }
- mock_storage.load_model_response.side_effect = lambda key: models.get(key, (None, None))
- mock_storage.load_tool_result.side_effect = lambda key: (
- (True, "A", "fp") if key == f"{P}tool_step_2" else (False, None, None)
- )
- run_usage = RunUsage(requests=5, tool_calls=5)
- ledger = ReplayUsageLedger(run_usage=run_usage, usage_limits=None)
- caching = CachingModel(mock_model, storage=mock_storage, counter=counter, replay_usage=ledger)
-
- await caching.request([], None, params)
-
- # -1 for the replayed request itself, -1 for the credited next request.
- assert run_usage.requests == 3
- assert run_usage.tool_calls == 4
- ledger.settle()
- assert (run_usage.requests, run_usage.tool_calls) == (4, 5)
-
- @pytest.mark.asyncio
- async def test_credited_next_request_reuses_the_peeked_response_without_reloading(
- self, mock_model, mock_storage, counter
- ):
- """``_credit_cached_successors`` already loads the next model step's response from
- storage to decide whether to credit it. Before this fix, ``request()`` loaded and
- deserialized that same entry again when it reached that step: two replayed steps
- cost 4 ``load_model_response`` calls (each step's own load, plus each step's scan
- re-loading the next one). Reusing the scan's object drops that to 3 -- the second
- step's own load is skipped because the first step's scan already produced it; only
- its own successor-scan (querying step 2, which is not cached) still hits storage."""
- response_0 = ModelResponse(parts=[TextPart(content="Hello!")])
- response_1 = ModelResponse(parts=[TextPart(content="World!")])
- models = {
- f"{P}model_step_0": (response_0, request_fingerprint()),
- f"{P}model_step_1": (response_1, request_fingerprint()),
- }
- mock_storage.load_model_response.side_effect = lambda key: models.get(key, (None, None))
- run_usage = RunUsage(requests=5)
- caching = CachingModel(
- mock_model,
- storage=mock_storage,
- counter=counter,
- replay_usage=ReplayUsageLedger(run_usage=run_usage, usage_limits=None),
- )
-
- result_0 = await caching.request([], None, ModelRequestParameters())
- calls_after_first_replay = mock_storage.load_model_response.call_count
- result_1 = await caching.request([], None, ModelRequestParameters())
-
- assert result_0 is response_0
- assert result_1 is response_1
- # Step 0's own load, plus its scan finding step 1 -- not yet the double-load this
- # test guards against, which would only show up once step 1's request() runs.
- assert calls_after_first_replay == 2
- # Without reuse this would be 4 (step 1 reloaded, then its scan queries step 2).
- assert mock_storage.load_model_response.call_count == 3
-
-
-class TestCachingModelCacheMiss:
- @pytest.mark.asyncio
- async def test_calls_model_and_caches_on_miss(self, mock_model, mock_storage, counter, sample_response):
- mock_model.request = AsyncMock(return_value=sample_response)
- caching = CachingModel(mock_model, storage=mock_storage, counter=counter)
-
- result = await caching.request([], None, ModelRequestParameters())
-
- assert result is sample_response
- mock_model.request.assert_called_once()
- mock_storage.save_model_response.assert_called_once_with(
- f"{P}model_step_0", sample_response, fingerprint=request_fingerprint()
- )
-
- @pytest.mark.asyncio
- async def test_sequential_calls_use_incrementing_keys(self, mock_model, mock_storage, counter):
- response_1 = ModelResponse(parts=[TextPart(content="First")])
- response_2 = ModelResponse(parts=[TextPart(content="Second")])
- mock_model.request = AsyncMock(side_effect=[response_1, response_2])
- caching = CachingModel(mock_model, storage=mock_storage, counter=counter)
-
- await caching.request([], None, ModelRequestParameters())
- await caching.request([], None, ModelRequestParameters())
-
- keys = [call[0][0] for call in mock_storage.save_model_response.call_args_list]
- assert keys == [f"{P}model_step_0", f"{P}model_step_1"]
-
- @pytest.mark.asyncio
- @pytest.mark.parametrize(
- ("written", "expected"),
- [pytest.param(True, (1, 0), id="written"), pytest.param(False, (0, 1), id="refused")],
- )
- async def test_counts_cached_only_when_storage_wrote_it(
- self, mock_model, mock_storage, counter, sample_response, written, expected
- ):
- """A response the backend did not store re-runs on retry, so it counts as skipped."""
- mock_model.request = AsyncMock(return_value=sample_response)
- mock_storage.save_model_response.return_value = written
- caching = CachingModel(mock_model, storage=mock_storage, counter=counter)
-
- result = await caching.request([], None, ModelRequestParameters())
-
- assert result is sample_response
- assert (counter.cached_model, counter.skipped_model) == expected
-
- @pytest.mark.asyncio
- @pytest.mark.skipif(
- not AIRFLOW_V_3_1_PLUS, reason="cap_structlog needs airflow._shared, which lands in Airflow 3.1"
- )
- @pytest.mark.parametrize("written", [True, False])
- async def test_warns_only_when_storage_skipped_the_write(
- self, mock_model, mock_storage, counter, sample_response, written, cap_structlog
- ):
- mock_model.request = AsyncMock(return_value=sample_response)
- mock_storage.save_model_response.return_value = written
- caching = CachingModel(mock_model, storage=mock_storage, counter=counter)
-
- await caching.request([], None, ModelRequestParameters())
-
- assert ({"step": 0, "log_level": "warning"} in cap_structlog) is not written
-
-
-class TestCachingModelReplayVerification:
- @pytest.mark.asyncio
- async def test_fingerprint_mismatch_treated_as_miss(
- self, mock_model, mock_storage, counter, sample_response
- ):
- """A cached entry recorded for a different request must not be replayed."""
- stale = ModelResponse(parts=[TextPart(content="stale")])
- mock_storage.load_model_response.return_value = (stale, "fp_of_old_conversation")
- mock_model.request = AsyncMock(return_value=sample_response)
- caching = CachingModel(mock_model, storage=mock_storage, counter=counter)
-
- result = await caching.request([], None, ModelRequestParameters())
-
- assert result is sample_response
- mock_model.request.assert_called_once()
- assert counter.replayed_model == 0
- mock_storage.save_model_response.assert_called_once_with(
- f"{P}model_step_0", sample_response, fingerprint=request_fingerprint()
- )
-
- @pytest.mark.asyncio
- async def test_legacy_entry_without_fingerprint_treated_as_miss(
- self, mock_model, mock_storage, counter, sample_response
- ):
- """Pre-fingerprint cache entries cannot be verified, so they re-run."""
- stale = ModelResponse(parts=[TextPart(content="stale")])
- mock_storage.load_model_response.return_value = (stale, None)
- mock_model.request = AsyncMock(return_value=sample_response)
- caching = CachingModel(mock_model, storage=mock_storage, counter=counter)
-
- result = await caching.request([], None, ModelRequestParameters())
-
- assert result is sample_response
- mock_model.request.assert_called_once()
-
- @pytest.mark.asyncio
- async def test_fingerprint_uses_prepared_request_not_raw_arguments(
- self, mock_storage, counter, sample_response
- ):
- """Concrete models merge model-level settings in ``prepare_request`` before the
- provider sees the request. The fingerprint must reflect the prepared settings,
- so a model-level change (e.g. a different temperature on the connection) is not
- invisible behind identical raw ``request()`` arguments."""
- model = MagicMock()
- model.model_name = "test-model"
- model.system = "test"
- model.profile = MagicMock()
- model.settings = None
- model.request = AsyncMock(return_value=sample_response)
- # Simulate prepare_request merging a model-level temperature into settings.
- model.prepare_request = lambda settings, params: ({"temperature": 0.9}, params)
- caching = CachingModel(model, storage=mock_storage, counter=counter)
-
- await caching.request([], None, ModelRequestParameters())
-
- stored_fingerprint = mock_storage.save_model_response.call_args.kwargs["fingerprint"]
- # Reflects the prepared settings, not the raw ``None`` the agent passed in.
- assert stored_fingerprint == fingerprint_model_request(
- "test:test-model", [], {"temperature": 0.9}, ModelRequestParameters()
- )
- assert stored_fingerprint != request_fingerprint()
diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_caching_toolset.py b/providers/common/ai/tests/unit/common/ai/durable/test_caching_toolset.py
deleted file mode 100644
index 307b7285df9..00000000000
--- a/providers/common/ai/tests/unit/common/ai/durable/test_caching_toolset.py
+++ /dev/null
@@ -1,268 +0,0 @@
-# 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
-
-from types import SimpleNamespace
-from unittest.mock import AsyncMock, MagicMock, patch
-
-import pytest
-from pydantic_ai.messages import ModelResponse, TextPart
-from pydantic_ai.models import ModelRequestParameters
-from pydantic_ai.toolsets.abstract import ToolsetTool
-from pydantic_ai.toolsets.combined import CombinedToolset
-from pydantic_ai.toolsets.function import FunctionToolset
-
-from airflow.providers.common.ai.durable.base import DURABLE_KEY_PREFIX as P, DurableStorageProtocol
-from airflow.providers.common.ai.durable.caching_model import CachingModel
-from airflow.providers.common.ai.durable.caching_toolset import CachingToolset
-from airflow.providers.common.ai.durable.fingerprint import fingerprint_tool_call
-from airflow.providers.common.ai.durable.step_counter import DurableStepCounter
-
-from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS
-
-
-@pytest.fixture
-def mock_storage():
- storage = MagicMock(spec=DurableStorageProtocol)
- storage.load_tool_result.return_value = (False, None, None)
- storage.load_model_response.return_value = (None, None)
- storage.save_tool_result.return_value = True
- storage.save_model_response.return_value = True
- return storage
-
-
-@pytest.fixture
-def counter():
- return DurableStepCounter()
-
-
-@pytest.fixture
-def mock_toolset():
- toolset = MagicMock()
- toolset.call_tool = AsyncMock(return_value="fresh result")
- toolset.get_tools = AsyncMock(return_value={})
- toolset.__aenter__ = AsyncMock(return_value=toolset)
- toolset.__aexit__ = AsyncMock(return_value=None)
- return toolset
-
-
-def ctx_for(tool_call_id: str | None = "call_1") -> SimpleNamespace:
- return SimpleNamespace(tool_call_id=tool_call_id)
-
-
-class TestCachingToolsetCacheHit:
- @pytest.mark.asyncio
- async def test_returns_cached_result_without_calling_tool(self, mock_toolset, mock_storage, counter):
- fingerprint = fingerprint_tool_call("search", {"q": "foo"}, "call_1")
- mock_storage.load_tool_result.return_value = (True, "cached result", fingerprint)
- caching = CachingToolset(wrapped=mock_toolset, storage=mock_storage, counter=counter)
-
- result = await caching.call_tool("search", {"q": "foo"}, ctx_for("call_1"), MagicMock())
-
- assert result == "cached result"
- mock_toolset.call_tool.assert_not_called()
- mock_storage.load_tool_result.assert_called_once_with(f"{P}tool_step_0")
-
- @pytest.mark.asyncio
- async def test_advances_counter_on_cache_hit(self, mock_toolset, mock_storage, counter):
- fingerprint = fingerprint_tool_call("search", {}, "call_1")
- mock_storage.load_tool_result.return_value = (True, "cached", fingerprint)
- caching = CachingToolset(wrapped=mock_toolset, storage=mock_storage, counter=counter)
-
- await caching.call_tool("search", {}, ctx_for("call_1"), MagicMock())
-
- assert counter.total_steps == 1
-
-
-class TestCachingToolsetCacheMiss:
- @pytest.mark.asyncio
- async def test_calls_tool_and_caches_on_miss(self, mock_toolset, mock_storage, counter):
- caching = CachingToolset(wrapped=mock_toolset, storage=mock_storage, counter=counter)
-
- result = await caching.call_tool("search", {"q": "foo"}, ctx_for("call_1"), MagicMock())
-
- assert result == "fresh result"
- mock_toolset.call_tool.assert_called_once()
- mock_storage.save_tool_result.assert_called_once_with(
- f"{P}tool_step_0",
- "fresh result",
- fingerprint=fingerprint_tool_call("search", {"q": "foo"}, "call_1"),
- )
-
- @pytest.mark.asyncio
- async def test_sequential_calls_use_incrementing_keys(self, mock_toolset, mock_storage, counter):
- mock_toolset.call_tool = AsyncMock(side_effect=["result_a", "result_b"])
- caching = CachingToolset(wrapped=mock_toolset, storage=mock_storage, counter=counter)
-
- await caching.call_tool("tool_a", {}, ctx_for(), MagicMock())
- await caching.call_tool("tool_b", {}, ctx_for(), MagicMock())
-
- keys = [call[0][0] for call in mock_storage.save_tool_result.call_args_list]
- assert keys == [f"{P}tool_step_0", f"{P}tool_step_1"]
-
- @pytest.mark.asyncio
- async def test_skipped_write_is_recorded_by_tool_name(self, mock_toolset, mock_storage, counter):
- """A result the backend did not store re-runs on retry, so it is not counted as cached."""
- mock_storage.save_tool_result.side_effect = [True, False]
- caching = CachingToolset(wrapped=mock_toolset, storage=mock_storage, counter=counter)
-
- await caching.call_tool("get_schema", {}, ctx_for("c1"), MagicMock())
- result = await caching.call_tool("run_query", {}, ctx_for("c2"), MagicMock())
-
- assert result == "fresh result"
- assert counter.cached_tool == 1
- assert counter.skipped_tools == ["run_query"]
-
- @pytest.mark.asyncio
- @pytest.mark.skipif(
- not AIRFLOW_V_3_1_PLUS, reason="cap_structlog needs airflow._shared, which lands in Airflow 3.1"
- )
- async def test_skipped_write_warns_by_tool_name(self, mock_toolset, mock_storage, counter, cap_structlog):
- """The warning names the tool on every path, not only in a successful run's summary."""
- mock_storage.save_tool_result.side_effect = [True, False]
- caching = CachingToolset(wrapped=mock_toolset, storage=mock_storage, counter=counter)
-
- await caching.call_tool("get_schema", {}, ctx_for("c1"), MagicMock())
- await caching.call_tool("run_query", {}, ctx_for("c2"), MagicMock())
-
- assert {"tool": "run_query", "step": 1, "log_level": "warning"} in cap_structlog
- assert {"tool": "get_schema", "log_level": "warning"} not in cap_structlog
-
-
-class TestCachingToolsetReplayVerification:
- @pytest.mark.asyncio
- async def test_different_tool_call_treated_as_miss(self, mock_toolset, mock_storage, counter):
- """A cached result recorded for a different tool call must not be replayed."""
- stale_fingerprint = fingerprint_tool_call("lookup_order", {"id": "A1"}, "old_call")
- mock_storage.load_tool_result.return_value = (True, "stale result", stale_fingerprint)
- caching = CachingToolset(wrapped=mock_toolset, storage=mock_storage, counter=counter)
-
- result = await caching.call_tool("charge_card", {"amount": 5}, ctx_for("new_call"), MagicMock())
-
- assert result == "fresh result"
- mock_toolset.call_tool.assert_called_once()
- assert counter.replayed_tool == 0
-
- @pytest.mark.asyncio
- async def test_changed_tool_call_id_treated_as_miss(self, mock_toolset, mock_storage, counter):
- """Same name/args but a new model-issued call id means the conversation diverged."""
- stale_fingerprint = fingerprint_tool_call("search", {"q": "foo"}, "old_call")
- mock_storage.load_tool_result.return_value = (True, "stale result", stale_fingerprint)
- caching = CachingToolset(wrapped=mock_toolset, storage=mock_storage, counter=counter)
-
- result = await caching.call_tool("search", {"q": "foo"}, ctx_for("new_call"), MagicMock())
-
- assert result == "fresh result"
- mock_toolset.call_tool.assert_called_once()
-
- @pytest.mark.asyncio
- async def test_legacy_entry_without_fingerprint_treated_as_miss(
- self, mock_toolset, mock_storage, counter
- ):
- """Pre-fingerprint cache entries cannot be verified, so the tool re-runs."""
- mock_storage.load_tool_result.return_value = (True, "stale result", None)
- caching = CachingToolset(wrapped=mock_toolset, storage=mock_storage, counter=counter)
-
- result = await caching.call_tool("search", {"q": "foo"}, ctx_for("call_1"), MagicMock())
-
- assert result == "fresh result"
- mock_toolset.call_tool.assert_called_once()
-
-
-class TestSharedCounter:
- @pytest.mark.asyncio
- async def test_model_and_toolset_share_counter(self, mock_toolset, mock_storage):
- """When CachingModel and CachingToolset share a counter, steps interleave correctly."""
- counter = DurableStepCounter()
-
- mock_model = MagicMock()
- mock_model.model_name = "test"
- mock_model.system = "test"
- mock_model.profile = MagicMock()
- mock_model.settings = None
- mock_model.prepare_request = lambda settings, params: (settings, params)
-
- response = ModelResponse(parts=[TextPart(content="response")])
- mock_model.request = AsyncMock(return_value=response)
-
- with patch("pydantic_ai.models.wrapper.infer_model", side_effect=lambda m: m):
- caching_model = CachingModel(mock_model, storage=mock_storage, counter=counter)
- caching_toolset = CachingToolset(wrapped=mock_toolset, storage=mock_storage, counter=counter)
-
- # Simulate: model call -> tool call -> model call
- await caching_model.request([], None, ModelRequestParameters())
- await caching_toolset.call_tool("search", {}, ctx_for(), MagicMock())
- await caching_model.request([], None, ModelRequestParameters())
-
- model_keys = [call[0][0] for call in mock_storage.save_model_response.call_args_list]
- tool_keys = [call[0][0] for call in mock_storage.save_tool_result.call_args_list]
-
- assert model_keys == [f"{P}model_step_0", f"{P}model_step_2"]
- assert tool_keys == [f"{P}tool_step_1"]
- assert counter.total_steps == 3
-
-
-class TestCachingToolsetReplayable:
- @pytest.mark.asyncio
- async def test_a_non_replayable_toolset_is_never_served_from_cache(
- self, mock_toolset, mock_storage, counter
- ):
- # A managed agent may have acted on a system Airflow cannot observe, so its toolset
- # declares replayable=False and a cached answer must not stand in for a fresh call.
- mock_toolset.replayable = False
- mock_storage.load_tool_result.return_value = (
- True,
- "stale cached result",
- fingerprint_tool_call("t", {}, "call_1"),
- )
- caching = CachingToolset(wrapped=mock_toolset, storage=mock_storage, counter=counter)
-
- tool = MagicMock(spec=ToolsetTool)
- tool.toolset = mock_toolset
-
- result = await caching.call_tool("t", {}, ctx_for(), tool=tool)
-
- assert result == "fresh result"
- mock_toolset.call_tool.assert_awaited_once()
- mock_storage.load_tool_result.assert_not_called()
- mock_storage.save_tool_result.assert_not_called()
- assert counter.replayed_tool == 0
- # The step is still consumed so later steps keep their keys.
- assert counter.next_step() == 1
-
- @pytest.mark.asyncio
- async def test_the_flag_is_read_per_tool_through_prefixed_and_combined_wrappers(
- self, mock_storage, counter
- ):
- # ``.prefixed()`` and ``CombinedToolset`` do not carry the attribute; the cache reads it off the
- # toolset each tool came from, so one non-replayable member does not stop its siblings replaying.
- class NonReplayable(FunctionToolset):
- replayable = False
-
- managed, plain = NonReplayable(), FunctionToolset()
- wrapped = CombinedToolset([plain, managed.prefixed("claims")])
- mock_storage.load_tool_result.return_value = (True, "stale", fingerprint_tool_call("t", {}, "call_1"))
- caching = CachingToolset(wrapped=wrapped, storage=mock_storage, counter=counter)
- from_managed, from_plain = MagicMock(spec=ToolsetTool), MagicMock(spec=ToolsetTool)
- from_managed.toolset, from_plain.toolset = managed, plain
-
- with patch.object(CombinedToolset, "call_tool", autospec=True, return_value="fresh") as call_tool:
- assert await caching.call_tool("t", {}, ctx_for(), tool=from_managed) == "fresh"
- assert await caching.call_tool("t", {}, ctx_for(), tool=from_plain) == "stale"
-
- call_tool.assert_awaited_once()
- mock_storage.load_tool_result.assert_called_once()
diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_capability.py b/providers/common/ai/tests/unit/common/ai/durable/test_capability.py
new file mode 100644
index 00000000000..02fb114ee6b
--- /dev/null
+++ b/providers/common/ai/tests/unit/common/ai/durable/test_capability.py
@@ -0,0 +1,855 @@
+# 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.
+"""
+AirflowDurability through a real pydantic-ai agent loop.
+
+Each test runs the scenario durable execution exists for: an attempt fails partway, and
+the retry runs the agent again from the top with a fresh journal over the same storage.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import dataclasses
+from typing import TYPE_CHECKING, Any
+
+import pytest
+from pydantic_ai import Agent, CancellationToken, RunContext
+from pydantic_ai.capabilities import AbstractCapability, durable_operation
+from pydantic_ai.exceptions import ModelRetry
+from pydantic_ai.messages import (
+ ModelMessage,
+ ModelResponse,
+ RetryPromptPart,
+ TextPart,
+ ToolCallPart,
+ ToolReturnPart,
+)
+from pydantic_ai.models.function import AgentInfo, FunctionModel
+from pydantic_ai.toolsets import FunctionToolset
+from pydantic_ai.usage import RequestUsage, RunUsage, UsageLimits
+
+from airflow.providers.common.ai.durable import AirflowDurability
+from airflow.providers.common.ai.durable.journal import DurableJournal, journal_scope
+from airflow.providers.common.ai.exceptions import DurableJournalError
+from airflow.providers.common.ai.utils.toolset_base import AirflowToolset, ensure_masked
+
+if TYPE_CHECKING:
+ from pydantic_ai.toolsets.abstract import ToolsetTool
+
+
+class Calls:
+ """Counts the live calls one test's model and tools make, across attempts."""
+
+ def __init__(self) -> None:
+ self.counts: dict[str, int] = {}
+
+ def bump(self, name: str) -> None:
+ self.counts[name] = self.counts.get(name, 0) + 1
+
+ def __getitem__(self, name: str) -> int:
+ return self.counts.get(name, 0)
+
+
+def responses_so_far(messages: list[ModelMessage]) -> int:
+ return sum(isinstance(message, ModelResponse) for message in messages)
+
+
+def tool_returns(messages: list[ModelMessage]) -> list[ToolReturnPart]:
+ return [part for message in messages for part in message.parts if isinstance(part, ToolReturnPart)]
+
+
+async def attempt(storage, agent: Agent[Any, Any], prompt: str = "go", **run_kwargs: Any) -> Any:
+ """Run one task attempt: a fresh journal over the storage the attempts share."""
+ with journal_scope(DurableJournal(storage)):
+ return await agent.run(prompt, **run_kwargs)
+
+
+class _WarehouseToolset(AirflowToolset):
+ """An Airflow toolset with one ``query`` tool, standing in for SQLToolset and friends."""
+
+ def __init__(self, calls: Calls, result: Any = "3 rows", *, replayable: bool = True) -> None:
+ self._calls = calls
+ self._result = result
+ self.replayable = replayable
+
+ def query(sql: str) -> str:
+ """Run a query."""
+ raise AssertionError("served by execute_tool")
+
+ self._inner = FunctionToolset(tools=[query])
+
+ @property
+ def id(self) -> str:
+ return "warehouse"
+
+ async def get_tools(self, ctx: RunContext[Any]) -> dict[str, ToolsetTool[Any]]:
+ tools = await self._inner.get_tools(ctx)
+ return {name: dataclasses.replace(tool, toolset=self) for name, tool in tools.items()}
+
+ async def execute_tool(self, name, tool_args, *, ctx, tool) -> Any:
+ self._calls.bump("query")
+ if isinstance(self._result, BaseException):
+ raise self._result
+ return self._result
+
+
+def tool_then_answer(calls: Calls, tool_name: str, *, fail_final: list[bool]):
+ """A model that calls ``tool_name`` once, then answers with what the tool returned."""
+
+ def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+ calls.bump("model")
+ if responses_so_far(messages) == 0:
+ return ModelResponse(parts=[ToolCallPart(tool_name, {"sql": "select 1"})])
+ if fail_final[0]:
+ raise RuntimeError("worker died")
+ return ModelResponse(parts=[TextPart(f"answer: {tool_returns(messages)[-1].content}")])
+
+ return model_fn
+
+
+class TestReplay:
+ @pytest.mark.asyncio
+ async def test_retry_replays_completed_steps_and_runs_only_the_failed_one(self, memory_storage):
+ calls = Calls()
+ fail_final = [True]
+
+ def build() -> Agent[None, str]:
+ toolset = FunctionToolset(id="db")
+
+ @toolset.tool_plain
+ def query(sql: str) -> str:
+ calls.bump("query")
+ return "3 rows"
+
+ return Agent(
+ FunctionModel(tool_then_answer(calls, "query", fail_final=fail_final)),
+ name="analyst",
+ toolsets=[toolset],
+ capabilities=[AirflowDurability()],
+ )
+
+ with pytest.raises(RuntimeError, match="worker died"):
+ await attempt(memory_storage, build())
+ fail_final[0] = False
+ result = await attempt(memory_storage, build())
+
+ assert result.output == "answer: 3 rows"
+ assert calls["query"] == 1
+ # The first model step replayed; only the one that failed runs again.
+ assert calls["model"] == 3
+
+ @pytest.mark.asyncio
+ async def test_changed_prompt_runs_everything_again(self, memory_storage):
+ calls = Calls()
+ fail_final = [True]
+
+ def build(instructions: str) -> Agent[None, str]:
+ toolset = FunctionToolset(id="db")
+
+ @toolset.tool_plain
+ def query(sql: str) -> str:
+ calls.bump("query")
+ return "3 rows"
+
+ return Agent(
+ FunctionModel(tool_then_answer(calls, "query", fail_final=fail_final)),
+ name="analyst",
+ instructions=instructions,
+ toolsets=[toolset],
+ capabilities=[AirflowDurability()],
+ )
+
+ with pytest.raises(RuntimeError):
+ await attempt(memory_storage, build("be terse"))
+ fail_final[0] = False
+ await attempt(memory_storage, build("be thorough"))
+
+ assert calls["query"] == 2
+ assert calls["model"] == 4
+
+ @pytest.mark.asyncio
+ async def test_parallel_tool_calls_replay_by_the_order_they_started(self, memory_storage):
+ calls = Calls()
+ delays: dict[int, float] = {}
+ fail_final = [True]
+
+ def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+ if responses_so_far(messages) == 0:
+ return ModelResponse(parts=[ToolCallPart("charge", {"amount": n}) for n in (1, 2, 3)])
+ if fail_final[0]:
+ raise RuntimeError("worker died")
+ return ModelResponse(parts=[TextPart(",".join(str(p.content) for p in tool_returns(messages)))])
+
+ def build() -> Agent[None, str]:
+ toolset = FunctionToolset(id="billing")
+
+ @toolset.tool_plain
+ async def charge(amount: int) -> str:
+ await asyncio.sleep(delays.get(amount, 0))
+ calls.bump(f"charge_{amount}")
+ return f"charged {amount}"
+
+ return Agent(
+ FunctionModel(model_fn), name="biller", toolsets=[toolset], capabilities=[AirflowDurability()]
+ )
+
+ # The first attempt finishes the calls in the reverse of the order it started them.
+ delays.update({1: 0.03, 2: 0.02, 3: 0.0})
+ with pytest.raises(RuntimeError):
+ await attempt(memory_storage, build())
+ fail_final[0] = False
+ delays.clear()
+ result = await attempt(memory_storage, build())
+
+ assert result.output == "charged 1,charged 2,charged 3"
+ assert calls.counts == {"charge_1": 1, "charge_2": 1, "charge_3": 1}
+
+ @pytest.mark.asyncio
+ async def test_outside_a_journal_the_agent_runs_normally(self, memory_storage):
+ calls = Calls()
+ toolset = FunctionToolset(id="db")
+
+ @toolset.tool_plain
+ def query(sql: str) -> str:
+ calls.bump("query")
+ return "3 rows"
+
+ agent = Agent(
+ FunctionModel(tool_then_answer(calls, "query", fail_final=[False])),
+ name="analyst",
+ toolsets=[toolset],
+ capabilities=[AirflowDurability()],
+ )
+
+ result = await agent.run("go")
+
+ assert result.output == "answer: 3 rows"
+ assert memory_storage.entries == {}
+
+
+class TestResultsThatCannotBeRecorded:
+ @pytest.mark.asyncio
+ async def test_a_function_tool_result_that_is_not_json_fails_without_retrying(self, memory_storage):
+ toolset = FunctionToolset(id="db")
+
+ @toolset.tool_plain
+ def query(sql: str) -> object:
+ return object()
+
+ agent = Agent(
+ FunctionModel(tool_then_answer(Calls(), "query", fail_final=[False])),
+ name="analyst",
+ toolsets=[toolset],
+ capabilities=[AirflowDurability()],
+ )
+
+ with pytest.raises(DurableJournalError, match="not JSON-serializable"):
+ await attempt(memory_storage, agent)
+
+ @pytest.mark.asyncio
+ async def test_an_airflow_toolset_result_that_is_not_json_fails_without_retrying(self, memory_storage):
+ agent = Agent(
+ FunctionModel(tool_then_answer(Calls(), "query", fail_final=[False])),
+ name="analyst",
+ toolsets=[_WarehouseToolset(Calls(), object())],
+ capabilities=[AirflowDurability()],
+ )
+
+ with pytest.raises(DurableJournalError, match="not JSON-serializable"):
+ await attempt(memory_storage, agent)
+
+
+class TestStreaming:
+ @pytest.mark.asyncio
+ async def test_a_streamed_model_request_replays(self, memory_storage):
+ calls = Calls()
+
+ async def stream_fn(messages: list[ModelMessage], info: AgentInfo):
+ calls.bump("model")
+ yield "streamed answer"
+
+ def build() -> Agent[None, str]:
+ return Agent(
+ FunctionModel(stream_function=stream_fn), name="streamer", capabilities=[AirflowDurability()]
+ )
+
+ outputs = []
+ for _ in range(2):
+ with journal_scope(DurableJournal(memory_storage)):
+ async with build().run_stream("go") as result:
+ outputs.append(await result.get_output())
+
+ assert outputs == ["streamed answer", "streamed answer"]
+ assert calls["model"] == 1
+
+
+class TestCancellation:
+ @pytest.mark.asyncio
+ async def test_a_run_inside_the_task_accepts_a_cancellation_token(self, memory_storage):
+ """AgentOperator.on_kill cancels the run in the task process, which is the durable container."""
+ agent = Agent(
+ FunctionModel(tool_then_answer(Calls(), "query", fail_final=[False])),
+ name="analyst",
+ toolsets=[_WarehouseToolset(Calls())],
+ capabilities=[AirflowDurability()],
+ )
+
+ result = await attempt(memory_storage, agent, cancellation_token=CancellationToken())
+
+ assert result.output == "answer: 3 rows"
+
+
+class TestUnloadableRecords:
+ """A recorded result written by another version runs live instead of failing every retry."""
+
+ @staticmethod
+ def build(calls: Calls, toolset) -> Agent[None, str]:
+ return Agent(
+ FunctionModel(tool_then_answer(calls, "query", fail_final=[False])),
+ name="analyst",
+ toolsets=[toolset],
+ capabilities=[AirflowDurability()],
+ )
+
+ @staticmethod
+ def function_toolset(calls: Calls) -> FunctionToolset:
+ toolset = FunctionToolset(id="db")
+
+ @toolset.tool_plain
+ def query(sql: str) -> str:
+ calls.bump("query")
+ return "3 rows"
+
+ return toolset
+
+ @pytest.mark.parametrize(
+ ("kind", "expected_calls"),
+ [
+ # Every step from the first model request on runs again.
+ pytest.param("model", {"model": 4, "query": 2}, id="model response"),
+ # The first model request replays; the tool call and the final request run again.
+ pytest.param("tool", {"model": 3, "query": 2}, id="tool result"),
+ ],
+ )
+ @pytest.mark.asyncio
+ async def test_a_recorded_step_that_no_longer_loads_runs_live(self, memory_storage, kind, expected_calls):
+ calls = Calls()
+ await attempt(memory_storage, self.build(calls, self.function_toolset(calls)))
+ for entry in memory_storage.steps():
+ if entry["kind"] == kind:
+ entry["payload"] = "written by another version"
+
+ journal = DurableJournal(memory_storage)
+ with journal_scope(journal):
+ result = await self.build(calls, self.function_toolset(calls)).run("go")
+
+ assert result.output == "answer: 3 rows"
+ assert calls.counts == expected_calls
+ # Counted as run live in the attempt's summary, not as replayed.
+ assert journal.stats.replayed[kind] == 0
+
+ @pytest.mark.parametrize(
+ "payload",
+ [
+ pytest.param({"kind": "tool_return"}, id="no result"),
+ pytest.param({"kind": "renamed_in_a_later_version"}, id="unknown kind"),
+ ],
+ )
+ @pytest.mark.asyncio
+ async def test_a_function_tool_record_that_no_longer_loads_runs_live(self, memory_storage, payload):
+ calls = Calls()
+ await attempt(memory_storage, self.build(calls, self.function_toolset(calls)))
+ for entry in memory_storage.steps():
+ if entry["kind"] == "tool":
+ entry["payload"] = payload
+
+ result = await attempt(memory_storage, self.build(calls, self.function_toolset(calls)))
+
+ assert result.output == "answer: 3 rows"
+ assert calls["query"] == 2
+
+ @pytest.mark.asyncio
+ async def test_a_capability_operation_record_that_no_longer_loads_runs_live(self, memory_storage):
+ calls = Calls()
+
+ def build() -> Agent[None, str]:
+ return Agent(
+ FunctionModel(lambda messages, info: ModelResponse(parts=[TextPart("done")])),
+ name="analyst",
+ capabilities=[AirflowDurability(), _Ledger(calls)],
+ )
+
+ await attempt(memory_storage, build())
+ for entry in memory_storage.steps():
+ if "accrue" in entry["name"]:
+ entry["payload"] = "not an amount"
+
+ await attempt(memory_storage, build())
+
+ assert calls["accrue"] == 2
+
+ @pytest.mark.asyncio
+ async def test_an_airflow_toolset_record_that_no_longer_loads_runs_live(self, memory_storage):
+ calls = Calls()
+ await attempt(memory_storage, self.build(calls, _WarehouseToolset(calls)))
+ for entry in memory_storage.steps():
+ if entry["kind"] == "tool":
+ entry["payload"] = {"kind": "renamed_in_a_later_version"}
+
+ result = await attempt(memory_storage, self.build(calls, _WarehouseToolset(calls)))
+
+ assert result.output == "answer: 3 rows"
+ assert calls["query"] == 2
+
+
+class TestAirflowToolsets:
+ """Airflow's own toolsets are not durable units of pydantic-ai's backend; the capability journals them."""
+
+ @pytest.mark.asyncio
+ async def test_airflow_toolset_calls_replay(self, memory_storage):
+ calls = Calls()
+ fail_final = [True]
+
+ def build() -> Agent[None, str]:
+ return Agent(
+ FunctionModel(tool_then_answer(calls, "query", fail_final=fail_final)),
+ name="analyst",
+ toolsets=[_WarehouseToolset(calls)],
+ capabilities=[AirflowDurability()],
+ )
+
+ with pytest.raises(RuntimeError):
+ await attempt(memory_storage, build())
+ fail_final[0] = False
+ result = await attempt(memory_storage, build())
+
+ assert result.output == "answer: 3 rows"
+ assert calls["query"] == 1
+
+ @pytest.mark.asyncio
+ async def test_a_model_retry_replays_without_running_the_tool(self, memory_storage):
+ calls = Calls()
+ fail_final = [True]
+
+ def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+ if responses_so_far(messages) == 0:
+ return ModelResponse(parts=[ToolCallPart("query", {"sql": "select nope"})])
+ if fail_final[0]:
+ raise RuntimeError("worker died")
+ retry = messages[-1].parts[0]
+ assert isinstance(retry, RetryPromptPart)
+ return ModelResponse(parts=[TextPart(f"retry said: {retry.content}")])
+
+ def build() -> Agent[None, str]:
+ return Agent(
+ FunctionModel(model_fn),
+ name="analyst",
+ toolsets=[_WarehouseToolset(calls, ModelRetry("no column nope"))],
+ capabilities=[AirflowDurability()],
+ )
+
+ with pytest.raises(RuntimeError):
+ await attempt(memory_storage, build())
+ fail_final[0] = False
+ result = await attempt(memory_storage, build())
+
+ assert result.output == "retry said: no column nope"
+ assert calls["query"] == 1
+
+ @pytest.mark.asyncio
+ async def test_a_toolset_that_is_not_replayable_runs_again(self, memory_storage):
+ calls = Calls()
+ fail_final = [True]
+
+ def build() -> Agent[None, str]:
+ return Agent(
+ FunctionModel(tool_then_answer(calls, "query", fail_final=fail_final)),
+ name="analyst",
+ toolsets=[_WarehouseToolset(calls, replayable=False)],
+ capabilities=[AirflowDurability()],
+ )
+
+ with pytest.raises(RuntimeError):
+ await attempt(memory_storage, build())
+ fail_final[0] = False
+ await attempt(memory_storage, build())
+
+ assert calls["query"] == 2
+ # The model step before it still replayed.
+ assert calls["model"] == 3
+
+
+@pytest.mark.enable_redact
+class TestMasking:
+ @pytest.mark.asyncio
+ async def test_journal_holds_function_tool_results_masked(self, memory_storage, registered_secret):
+ toolset = FunctionToolset(id="db")
+
+ @toolset.tool_plain
+ def query(sql: str) -> str:
+ return f"password={registered_secret}"
+
+ agent = Agent(
+ FunctionModel(tool_then_answer(Calls(), "query", fail_final=[False])),
+ name="analyst",
+ # AgentOperator puts the masking wrapper outside the durable unit.
+ toolsets=[ensure_masked(toolset)],
+ capabilities=[AirflowDurability()],
+ )
+
+ result = await attempt(memory_storage, agent)
+
+ assert result.output == "answer: password=***"
+ assert registered_secret not in str(memory_storage.entries)
+
+ @pytest.mark.asyncio
+ async def test_a_masked_tool_result_does_not_stop_the_retry_replaying(
+ self, memory_storage, registered_secret
+ ):
+ """The model saw the secret live but sees ``***`` on replay; the steps after it still replay."""
+ calls = Calls()
+ fail_final = [True]
+
+ def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+ calls.bump("model")
+ match responses_so_far(messages):
+ case 0:
+ return ModelResponse(parts=[ToolCallPart("query", {"sql": "select 1"})])
+ case 1:
+ return ModelResponse(parts=[ToolCallPart("notify", {})])
+ if fail_final[0]:
+ raise RuntimeError("worker died")
+ return ModelResponse(parts=[TextPart("done")])
+
+ def build() -> Agent[None, str]:
+ # No masking wrapper around it, as in a plain @task.
+ toolset = FunctionToolset(id="db")
+
+ @toolset.tool_plain
+ def query(sql: str) -> str:
+ return f"password={registered_secret}"
+
+ @toolset.tool_plain
+ def notify() -> str:
+ calls.bump("notify")
+ return "sent"
+
+ return Agent(
+ FunctionModel(model_fn),
+ name="analyst",
+ toolsets=[toolset],
+ capabilities=[AirflowDurability()],
+ )
+
+ with pytest.raises(RuntimeError):
+ await attempt(memory_storage, build())
+ fail_final[0] = False
+ await attempt(memory_storage, build())
+
+ assert calls["notify"] == 1
+ # Three model steps on the first attempt, and only the one that failed on the retry.
+ assert calls["model"] == 4
+
+ @pytest.mark.asyncio
+ async def test_journal_holds_airflow_toolset_results_masked(self, memory_storage, registered_secret):
+ agent = Agent(
+ FunctionModel(tool_then_answer(Calls(), "query", fail_final=[False])),
+ name="analyst",
+ toolsets=[_WarehouseToolset(Calls(), f"password={registered_secret}")],
+ capabilities=[AirflowDurability()],
+ )
+
+ await attempt(memory_storage, agent)
+
+ assert registered_secret not in str(memory_storage.entries)
+
+
+class _Ledger(AbstractCapability[Any]):
+ """Stands in for pydantic-ai-harness SpendLimits: it accrues through a durable operation."""
+
+ def __init__(self, calls: Calls, id: str | None = "ledger") -> None:
+ self._calls = calls
+ self.id = id
+
+ async def after_model_request(
+ self, ctx: RunContext[Any], *, request_context: Any, response: ModelResponse
+ ):
+ await self.accrue(ctx, 1)
+ return response
+
+ @durable_operation(name="accrue")
+ async def accrue(self, ctx: RunContext[Any], amount: int) -> int:
+ self._calls.bump("accrue")
+ return amount
+
+
+class TestSharedCapability:
+ @pytest.mark.asyncio
+ async def test_one_instance_on_two_agents_fingerprints_each_agents_own_model(self, memory_storage):
+ """Changing the second agent's model between attempts must not replay its old response."""
+
+ # FunctionModel is named after its function, so each is a different model.
+ def answer_a(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+ return ModelResponse(parts=[TextPart("a")])
+
+ def answer_b1(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+ return ModelResponse(parts=[TextPart("b1")])
+
+ def answer_b2(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+ return ModelResponse(parts=[TextPart("b2")])
+
+ async def task_attempt(second_model) -> str:
+ durability = AirflowDurability()
+ first = Agent(FunctionModel(answer_a), name="first", capabilities=[durability])
+ second = Agent(FunctionModel(second_model), name="second", capabilities=[durability])
+ with journal_scope(DurableJournal(memory_storage)):
+ await first.run("go")
+ return (await second.run("go")).output
+
+ await task_attempt(answer_b1)
+ output = await task_attempt(answer_b2)
+
+ assert output == "b2"
+
+
+class TestCapabilityOperations:
+ @pytest.mark.asyncio
+ async def test_durable_operations_of_other_capabilities_replay(self, memory_storage):
+ calls = Calls()
+ fail_final = [True]
+
+ def build() -> Agent[None, str]:
+ toolset = FunctionToolset(id="db")
+
+ @toolset.tool_plain
+ def query(sql: str) -> str:
+ return "3 rows"
+
+ return Agent(
+ FunctionModel(tool_then_answer(calls, "query", fail_final=fail_final)),
+ name="analyst",
+ toolsets=[toolset],
+ capabilities=[AirflowDurability(), _Ledger(calls)],
+ )
+
+ with pytest.raises(RuntimeError):
+ await attempt(memory_storage, build())
+ fail_final[0] = False
+ await attempt(memory_storage, build())
+
+ # One accrual per completed model request, none repeated by the retry.
+ assert calls["accrue"] == 2
+
+ @pytest.mark.asyncio
+ async def test_durable_operations_run_directly_outside_a_task(self):
+ calls = Calls()
+ agent = Agent(
+ FunctionModel(lambda messages, info: ModelResponse(parts=[TextPart("done")])),
+ name="analyst",
+ capabilities=[AirflowDurability(), _Ledger(calls)],
+ )
+
+ await agent.run("go")
+
+ assert calls["accrue"] == 1
+
+
+class TestRuns:
+ @pytest.mark.asyncio
+ async def test_each_run_of_an_attempt_replays_on_its_own(self, memory_storage):
+ """A task that runs two agents replays both when the second one failed."""
+ calls = Calls()
+ fail_final = [True]
+
+ def build(name: str, fail: list[bool]) -> Agent[None, str]:
+ toolset = FunctionToolset(id="db")
+
+ @toolset.tool_plain
+ def query(sql: str) -> str:
+ calls.bump(f"query_{name}")
+ return "3 rows"
+
+ return Agent(
+ FunctionModel(tool_then_answer(calls, "query", fail_final=fail)),
+ name=name,
+ toolsets=[toolset],
+ capabilities=[AirflowDurability()],
+ )
+
+ async def task_attempt() -> None:
+ with journal_scope(DurableJournal(memory_storage)):
+ await build("first", [False]).run("go")
+ await build("second", fail_final).run("go")
+
+ with pytest.raises(RuntimeError):
+ await task_attempt()
+ fail_final[0] = False
+ await task_attempt()
+
+ assert calls["query_first"] == 1
+ assert calls["query_second"] == 1
+
+ @pytest.mark.asyncio
+ async def test_agents_called_from_tools_replay_their_own_steps(self, memory_storage):
+ """
+ The outer agent's first tool runs an agent that succeeds, its second one an agent that fails.
+
+ On retry the first tool replays, so its agent never starts; the second tool's agent must
+ still find its own recorded tool call rather than the first agent's.
+ """
+ calls = Calls()
+ fail_second = [True]
+
+ def sub_agent(label: str, fail: list[bool]) -> Agent[None, str]:
+ toolset = FunctionToolset(id=f"{label}_tools")
+
+ @toolset.tool_plain
+ def fetch(sql: str) -> str:
+ calls.bump(f"fetch_{label}")
+ return f"{label} rows"
+
+ return Agent(
+ FunctionModel(tool_then_answer(Calls(), "fetch", fail_final=fail)),
+ name=f"sub_{label}",
+ toolsets=[toolset],
+ capabilities=[AirflowDurability()],
+ )
+
+ def outer_model(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+ done = [part.tool_name for part in tool_returns(messages)]
+ if not done:
+ return ModelResponse(parts=[ToolCallPart("first", {})])
+ if done == ["first"]:
+ return ModelResponse(parts=[ToolCallPart("second", {})])
+ return ModelResponse(parts=[TextPart("all done")])
+
+ def build() -> Agent[None, str]:
+ toolset = FunctionToolset(id="delegates")
+
+ @toolset.tool_plain
+ async def first() -> str:
+ return (await sub_agent("first", [False]).run("go")).output
+
+ @toolset.tool_plain
+ async def second() -> str:
+ return (await sub_agent("second", fail_second).run("go")).output
+
+ return Agent(
+ FunctionModel(outer_model),
+ name="outer",
+ toolsets=[toolset],
+ capabilities=[AirflowDurability()],
+ )
+
+ with pytest.raises(RuntimeError):
+ await attempt(memory_storage, build())
+ fail_second[0] = False
+ result = await attempt(memory_storage, build())
+
+ assert result.output == "all done"
+ assert calls.counts == {"fetch_first": 1, "fetch_second": 1}
+
+ @pytest.mark.asyncio
+ async def test_a_journal_that_cleans_up_after_each_run_deletes_a_run_that_succeeded(self, memory_storage):
+ toolset = FunctionToolset(id="db")
+
+ @toolset.tool_plain
+ def query(sql: str) -> str:
+ return "3 rows"
+
+ agent = Agent(
+ FunctionModel(tool_then_answer(Calls(), "query", fail_final=[False])),
+ name="analyst",
+ toolsets=[toolset],
+ capabilities=[AirflowDurability()],
+ )
+
+ with journal_scope(DurableJournal(memory_storage, clean_up_after_run=True)):
+ await agent.run("go")
+
+ assert memory_storage.entries == {}
+
+
+class TestReplayUsage:
+ """A replayed step adds nothing to the usage the run counts and limits."""
+
+ @staticmethod
+ def priced_model(calls: Calls, fail_final: list[bool]):
+ def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+ calls.bump("model")
+ if responses_so_far(messages) == 0:
+ return ModelResponse(
+ parts=[ToolCallPart("query", {"sql": "select 1"})],
+ usage=RequestUsage(input_tokens=100, output_tokens=10),
+ )
+ if fail_final[0]:
+ raise RuntimeError("worker died")
+ return ModelResponse(
+ parts=[TextPart("done")], usage=RequestUsage(input_tokens=200, output_tokens=20)
+ )
+
+ return model_fn
+
+ def build(self, calls: Calls, fail_final: list[bool]) -> Agent[None, str]:
+ toolset = FunctionToolset(id="db")
+
+ @toolset.tool_plain
+ def query(sql: str) -> str:
+ return "3 rows"
+
+ return Agent(
+ FunctionModel(self.priced_model(calls, fail_final)),
+ name="analyst",
+ toolsets=[toolset],
+ capabilities=[AirflowDurability()],
+ )
+
+ async def ledgered_attempt(self, storage, agent, usage: RunUsage, limits: UsageLimits | None = None):
+ """An attempt that, like AgentOperator, carries one usage object across attempts."""
+ return await attempt(storage, agent, usage=usage, usage_limits=limits)
+
+ @pytest.mark.asyncio
+ async def test_the_retry_counts_only_its_live_work(self, memory_storage):
+ calls = Calls()
+ fail_final = [True]
+ with pytest.raises(RuntimeError):
+ await self.ledgered_attempt(memory_storage, self.build(calls, fail_final), RunUsage())
+ fail_final[0] = False
+ usage = RunUsage()
+
+ await self.ledgered_attempt(memory_storage, self.build(calls, fail_final), usage)
+
+ assert (usage.requests, usage.tool_calls, usage.input_tokens, usage.output_tokens) == (1, 0, 200, 20)
+
+ @pytest.mark.asyncio
+ async def test_a_retry_whose_budget_is_spent_can_still_replay(self, memory_storage):
+ """The first request replays for free, so a request limit already reached must not stop the run."""
+ calls = Calls()
+ fail_final = [True]
+ with pytest.raises(RuntimeError):
+ await self.ledgered_attempt(memory_storage, self.build(calls, fail_final), RunUsage())
+ fail_final[0] = False
+ # The cross-attempt budget already holds the first attempt's two requests.
+ usage = RunUsage(requests=2)
+
+ result = await self.ledgered_attempt(
+ memory_storage, self.build(calls, fail_final), usage, UsageLimits(request_limit=3)
+ )
+
+ assert result.output == "done"
+ assert usage.requests == 3
diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_fingerprint.py b/providers/common/ai/tests/unit/common/ai/durable/test_fingerprint.py
index 38c5555955c..43d27486790 100644
--- a/providers/common/ai/tests/unit/common/ai/durable/test_fingerprint.py
+++ b/providers/common/ai/tests/unit/common/ai/durable/test_fingerprint.py
@@ -239,6 +239,18 @@ class TestModelRequestFingerprint:
assert fp1 != fp2
+ @pytest.mark.enable_redact
+ def test_stable_when_a_registered_secret_in_the_history_is_masked(self, registered_secret):
+ """The journal records tool results masked, so a replayed history holds ``***``."""
+ live = fingerprint_model_request(
+ "m", make_messages(user=f"password={registered_secret}"), None, ModelRequestParameters()
+ )
+ replayed = fingerprint_model_request(
+ "m", make_messages(user="password=***"), None, ModelRequestParameters()
+ )
+
+ assert live == replayed
+
def test_unserializable_request_returns_none(self):
fp = fingerprint_model_request("m", [object()], None, ModelRequestParameters()) # type: ignore[list-item]
diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_journal.py b/providers/common/ai/tests/unit/common/ai/durable/test_journal.py
new file mode 100644
index 00000000000..fe962087f08
--- /dev/null
+++ b/providers/common/ai/tests/unit/common/ai/durable/test_journal.py
@@ -0,0 +1,427 @@
+# 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.
+"""The framework-neutral journal, driven directly the way an adapter for any agent framework would."""
+
+from __future__ import annotations
+
+import logging
+from types import SimpleNamespace
+from unittest import mock
+
+import pytest
+
+from airflow.providers.common.ai.durable import journal as journal_module
+from airflow.providers.common.ai.durable.base import build_step_key
+from airflow.providers.common.ai.durable.journal import (
+ DurableJournal,
+ current_journal,
+ journal_scope,
+)
+
+
+async def ok(value):
+ return value
+
+
+async def boom():
+ raise RuntimeError("tool failed")
+
+
+async def raise_(error: BaseException):
+ raise error
+
+
+class TestReplay:
+ @pytest.mark.asyncio
+ async def test_a_matching_step_replays_and_its_body_does_not_run(self, memory_storage):
+ first = DurableJournal(memory_storage).start_run()
+ await first.run(
+ "agent__model.request", kind="model", fingerprint="fp", body=lambda: ok({"text": "hi"})
+ )
+ body = mock.AsyncMock(spec=ok)
+
+ retry = DurableJournal(memory_storage).start_run()
+ payload = await retry.run("agent__model.request", kind="model", fingerprint="fp", body=body)
+
+ assert payload == {"text": "hi"}
+ body.assert_not_called()
+ assert retry.journal.stats.replayed["model"] == 1
+
+ @pytest.mark.parametrize(
+ ("name", "fingerprint"),
+ [
+ pytest.param("agent__other.call", "fp", id="different step at the position"),
+ pytest.param("agent__model.request", "changed", id="same step asked something else"),
+ ],
+ )
+ @pytest.mark.asyncio
+ async def test_a_mismatch_runs_that_step_and_every_later_one_live(
+ self, memory_storage, name, fingerprint
+ ):
+ first = DurableJournal(memory_storage).start_run()
+ await first.run("agent__model.request", kind="model", fingerprint="fp", body=lambda: ok(1))
+ await first.run("agent__tool.call", kind="tool", fingerprint="t", body=lambda: ok(2))
+
+ retry = DurableJournal(memory_storage).start_run()
+ mismatched = await retry.run(name, kind="model", fingerprint=fingerprint, body=lambda: ok("live"))
+ # Would match position 1 of the first attempt, but the run has already diverged.
+ after = await retry.run("agent__tool.call", kind="tool", fingerprint="t", body=lambda: ok("live too"))
+
+ assert (mismatched, after) == ("live", "live too")
+ assert retry.diverged
+ assert memory_storage.steps()[1]["payload"] == "live too"
+
+ @pytest.mark.asyncio
+ async def test_a_step_without_a_fingerprint_replays_on_its_name(self, memory_storage):
+ first = DurableJournal(memory_storage).start_run()
+ await first.run(
+ "agent__capability__ledger.accrue", kind="other", fingerprint=None, body=lambda: ok(5)
+ )
+
+ retry = DurableJournal(memory_storage).start_run()
+ payload = await retry.run(
+ "agent__capability__ledger.accrue", kind="other", fingerprint=None, body=lambda: ok(6)
+ )
+
+ assert payload == 5
+
+ @pytest.mark.asyncio
+ async def test_a_step_that_was_not_recorded_runs_again_and_the_rest_still_replay(self, memory_storage):
+ memory_storage.refuse.add("__commonai_durable__run_0_step_0")
+ first = DurableJournal(memory_storage).start_run()
+ await first.run("a", kind="tool", fingerprint="1", body=lambda: ok("first"))
+ await first.run("b", kind="tool", fingerprint="2", body=lambda: ok("second"))
+ memory_storage.refuse.clear()
+
+ retry = DurableJournal(memory_storage).start_run()
+ a = await retry.run("a", kind="tool", fingerprint="1", body=lambda: ok("first again"))
+ b = await retry.run("b", kind="tool", fingerprint="2", body=lambda: ok("not replayed"))
+
+ assert (a, b) == ("first again", "second")
+ assert first.journal.stats.not_recorded == [("tool", "a")]
+
+ @pytest.mark.asyncio
+ async def test_a_step_that_is_not_replayable_keeps_its_position_but_is_never_recorded(
+ self, memory_storage
+ ):
+ first = DurableJournal(memory_storage).start_run()
+ await first.run("managed", kind="tool", fingerprint=None, body=lambda: ok("acted"), replayable=False)
+ await first.run("after", kind="tool", fingerprint="a", body=lambda: ok("recorded"))
+
+ retry = DurableJournal(memory_storage).start_run()
+ managed = await retry.run(
+ "managed", kind="tool", fingerprint=None, body=lambda: ok("acted again"), replayable=False
+ )
+ after = await retry.run("after", kind="tool", fingerprint="a", body=lambda: ok("live"))
+
+ assert (managed, after) == ("acted again", "recorded")
+ assert [step["name"] for step in memory_storage.steps()] == ["after"]
+
+ @pytest.mark.asyncio
+ async def test_to_record_transforms_only_what_is_recorded(self, memory_storage):
+ durable_run = DurableJournal(memory_storage).start_run()
+
+ payload = await durable_run.run(
+ "t", kind="tool", fingerprint=None, body=lambda: ok("secret"), to_record=lambda value: "***"
+ )
+
+ assert payload == "secret"
+ assert memory_storage.steps()[0]["payload"] == "***"
+
+
+class TestFailures:
+ @pytest.mark.asyncio
+ async def test_error_handling_of_a_failed_run_is_not_replayed(self, memory_storage):
+ """What error handling recorded after the step that failed the run must not replay into the retry."""
+ first = DurableJournal(memory_storage).start_run()
+ error = RuntimeError("tool failed")
+ with pytest.raises(RuntimeError):
+ await first.run("flaky", kind="tool", fingerprint="f", body=lambda: raise_(error))
+ await first.run("record_effect", kind="other", fingerprint="e", body=lambda: ok("failed"))
+ first.fail(error)
+
+ retry = DurableJournal(memory_storage).start_run()
+ flaky = await retry.run("flaky", kind="tool", fingerprint="f", body=lambda: ok("ok"))
+ effect = await retry.run("record_effect", kind="other", fingerprint="e", body=lambda: ok("completed"))
+
+ assert (flaky, effect) == ("ok", "completed")
+
+ @pytest.mark.asyncio
+ async def test_steps_started_before_the_failure_still_replay(self, memory_storage):
+ """A tool call running alongside the one that failed had already started, so it replays."""
+ first = DurableJournal(memory_storage).start_run()
+ error = RuntimeError("tool failed")
+ failing = first.claim("flaky", kind="tool", fingerprint="f")
+ sibling = first.claim("sibling", kind="tool", fingerprint="s")
+ failing.fail(error)
+ sibling.record("sibling result")
+ await first.run("on_error", kind="other", fingerprint="e", body=lambda: ok("cleanup"))
+ first.fail(error)
+
+ retry = DurableJournal(memory_storage).start_run()
+ await retry.run("flaky", kind="tool", fingerprint="f", body=lambda: ok("ok"))
+ replayed = await retry.run("sibling", kind="tool", fingerprint="s", body=lambda: ok("ran again"))
+ after = await retry.run("on_error", kind="other", fingerprint="e", body=lambda: ok("live"))
+
+ assert (replayed, after) == ("sibling result", "live")
+
+ @pytest.mark.asyncio
+ async def test_a_failure_the_run_recovered_from_does_not_stop_replay(self, memory_storage):
+ first = DurableJournal(memory_storage).start_run()
+ with pytest.raises(RuntimeError):
+ await first.run("flaky", kind="tool", fingerprint="f", body=boom)
+ await first.run("next", kind="tool", fingerprint="n", body=lambda: ok("recorded"))
+
+ retry = DurableJournal(memory_storage).start_run()
+ flaky = await retry.run("flaky", kind="tool", fingerprint="f", body=lambda: ok("ok"))
+ after = await retry.run("next", kind="tool", fingerprint="n", body=lambda: ok("ran again"))
+
+ assert (flaky, after) == ("ok", "recorded")
+
+ @pytest.mark.asyncio
+ async def test_a_run_failure_no_step_raised_keeps_everything_recorded(self, memory_storage):
+ first = DurableJournal(memory_storage).start_run()
+ await first.run("a", kind="tool", fingerprint="a", body=lambda: ok("recorded"))
+ first.fail(RuntimeError("usage limit"))
+
+ retry = DurableJournal(memory_storage).start_run()
+
+ assert await retry.run("a", kind="tool", fingerprint="a", body=lambda: ok("live")) == "recorded"
+
+ @pytest.mark.asyncio
+ async def test_a_step_without_a_fingerprint_does_not_replay_after_a_live_step(self, memory_storage):
+ """The live step may have produced something else, which the unfingerprinted step would depend on."""
+ memory_storage.refuse.add(build_step_key(0, 0))
+ first = DurableJournal(memory_storage).start_run()
+ await first.run("tool", kind="tool", fingerprint="t", body=lambda: ok("rows v1"))
+ await first.run("summarize", kind="other", fingerprint=None, body=lambda: ok("summary of v1"))
+ memory_storage.refuse.clear()
+
+ retry = DurableJournal(memory_storage).start_run()
+ await retry.run("tool", kind="tool", fingerprint="t", body=lambda: ok("rows v2"))
+ summary = await retry.run(
+ "summarize", kind="other", fingerprint=None, body=lambda: ok("summary of v2")
+ )
+
+ assert summary == "summary of v2"
+
+
+class TestReject:
+ @pytest.mark.asyncio
+ async def test_a_rejected_step_and_every_step_after_it_run_live(self, memory_storage):
+ first = DurableJournal(memory_storage).start_run()
+ await first.run("a", kind="model", fingerprint="a", body=lambda: ok("old a"))
+ await first.run("b", kind="tool", fingerprint="b", body=lambda: ok("old b"))
+
+ retry = DurableJournal(memory_storage).start_run()
+ step = retry.claim("a", kind="model", fingerprint="a")
+ retry.reject(step, reason="the recorded result no longer loads")
+ b = await retry.run("b", kind="tool", fingerprint="b", body=lambda: ok("live b"))
+
+ assert (step.replayed, step.payload, b) == (False, None, "live b")
+ assert retry.journal.stats.replayed.total() == 0
+
+
+class TestFailuresDuringReplay:
+ @pytest.mark.asyncio
+ async def test_an_attempt_killed_while_replaying_leaves_the_earlier_steps_replayable(
+ self, memory_storage
+ ):
+ first = DurableJournal(memory_storage).start_run()
+ for name in ("a", "b", "c"):
+ await first.run(
+ name, kind="tool", fingerprint=name, body=lambda name=name: ok(f"{name} recorded")
+ )
+ second = DurableJournal(memory_storage).start_run()
+ await second.run("a", kind="tool", fingerprint="a", body=lambda: ok("live"))
+ second.fail(TimeoutError("killed while replaying"))
+
+ third = DurableJournal(memory_storage).start_run()
+ results = [await third.run(n, kind="tool", fingerprint=n, body=lambda: ok("live")) for n in "abc"]
+
+ assert results == ["a recorded", "b recorded", "c recorded"]
+
+
+class TestCleanup:
+ @pytest.mark.asyncio
+ async def test_cleanup_reaches_a_run_started_by_a_step_that_now_replays(self, memory_storage):
+ """The nested run never starts again on the retry, but its steps are still deleted."""
+ first_journal = DurableJournal(memory_storage)
+ first = first_journal.start_run()
+
+ async def tool_that_runs_an_agent():
+ await first_journal.start_run().run("inner", kind="tool", fingerprint="i", body=lambda: ok(1))
+ return "done"
+
+ await first.run("tool", kind="tool", fingerprint="t", body=tool_that_runs_an_agent)
+ first.fail(RuntimeError("worker died"))
+ retry_journal = DurableJournal(memory_storage)
+ replayed = await retry_journal.start_run().run(
+ "tool", kind="tool", fingerprint="t", body=tool_that_runs_an_agent
+ )
+
+ retry_journal.cleanup()
+
+ assert replayed == "done"
+ assert memory_storage.entries == {}
+
+ @pytest.mark.asyncio
+ async def test_cleanup_deletes_what_an_earlier_attempt_recorded_beyond_this_one(self, memory_storage):
+ """Left behind, those steps would replay into a run started by clearing the task."""
+ first = DurableJournal(memory_storage).start_run()
+ for name in ("a", "b", "c"):
+ await first.run(name, kind="tool", fingerprint=name, body=lambda: ok(1))
+ first.fail(RuntimeError("worker died"))
+ retry_journal = DurableJournal(memory_storage)
+ await retry_journal.start_run().run("other", kind="tool", fingerprint="o", body=lambda: ok(2))
+
+ retry_journal.cleanup()
+
+ assert memory_storage.entries == {}
+
+ @pytest.mark.asyncio
+ async def test_cleanup_finds_steps_of_an_attempt_that_was_killed(self, memory_storage):
+ """A killed attempt never records how far it got."""
+ first = DurableJournal(memory_storage).start_run()
+ for name in ("a", "b", "c"):
+ await first.run(name, kind="tool", fingerprint=name, body=lambda: ok(1))
+ retry_journal = DurableJournal(memory_storage)
+ await retry_journal.start_run().run("other", kind="tool", fingerprint="o", body=lambda: ok(2))
+
+ retry_journal.cleanup()
+
+ assert memory_storage.entries == {}
+
+
+class TestRuns:
+ @pytest.mark.asyncio
+ async def test_runs_number_their_steps_separately(self, memory_storage):
+ journal = DurableJournal(memory_storage)
+ await journal.start_run().run("a", kind="other", fingerprint=None, body=lambda: ok(1))
+ await journal.start_run().run("b", kind="other", fingerprint=None, body=lambda: ok(2))
+
+ assert memory_storage.steps(0)[0]["name"] == "a"
+ assert memory_storage.steps(1)[0]["name"] == "b"
+
+ @pytest.mark.asyncio
+ async def test_a_run_started_by_a_step_is_numbered_under_it(self, memory_storage):
+ """So replaying a step that started a run never shifts the runs later steps start."""
+ journal = DurableJournal(memory_storage)
+ outer = journal.start_run()
+ nested_keys = []
+
+ async def tool_that_runs_an_agent():
+ nested_keys.append(journal.start_run().key)
+ return "done"
+
+ await outer.run("tool", kind="tool", fingerprint="t", body=tool_that_runs_an_agent)
+
+ assert nested_keys == ["0.0.1"]
+ assert journal.start_run().key == "1"
+
+ def test_the_run_id_is_the_first_attempts_on_every_retry(self, memory_storage):
+ first = DurableJournal(memory_storage).get_run_id(default="ti-try-1")
+ retry = DurableJournal(memory_storage).get_run_id(default="ti-try-2")
+
+ assert (first, retry) == ("ti-try-1", "ti-try-1")
+
+ @pytest.mark.asyncio
+ async def test_cleanup_deletes_the_steps_and_the_run_id(self, memory_storage):
+ journal = DurableJournal(memory_storage)
+ journal.get_run_id(default="ti-try-1")
+ await journal.start_run().run("a", kind="other", fingerprint=None, body=lambda: ok(1))
+
+ journal.cleanup()
+
+ assert memory_storage.entries == {}
+ assert DurableJournal(memory_storage).get_run_id(default="ti-try-2") == "ti-try-2"
+
+ @pytest.mark.asyncio
+ async def test_cleanup_deletes_replayed_steps_too(self, memory_storage):
+ await (
+ DurableJournal(memory_storage)
+ .start_run()
+ .run("a", kind="other", fingerprint=None, body=lambda: ok(1))
+ )
+ retry = DurableJournal(memory_storage)
+ await retry.start_run().run("a", kind="other", fingerprint=None, body=lambda: ok(2))
+
+ retry.cleanup()
+
+ assert memory_storage.entries == {}
+
+
+class TestSummary:
+ @pytest.mark.asyncio
+ async def test_summary_names_the_tools_a_retry_runs_again(self, memory_storage, caplog):
+ memory_storage.refuse.update({"__commonai_durable__run_0_step_1", "__commonai_durable__run_0_step_2"})
+ journal = DurableJournal(memory_storage)
+ durable_run = journal.start_run()
+ await durable_run.run("m", kind="model", fingerprint=None, body=lambda: ok(1))
+ await durable_run.run("send_email", kind="tool", fingerprint=None, body=lambda: ok(1))
+ await durable_run.run("send_email", kind="tool", fingerprint=None, body=lambda: ok(1))
+ logger = logging.getLogger("test_summary")
+
+ with caplog.at_level(logging.INFO, logger="test_summary"):
+ journal.log_summary(logger)
+
+ assert "recorded 1 new steps (1 model, 0 tool, 0 other)" in caplog.text
+ assert "2 tool results were not recorded, and a retry runs them again: send_email (x2)" in caplog.text
+
+
+class TestCurrentJournal:
+ def test_a_scoped_journal_wins(self, memory_storage):
+ journal = DurableJournal(memory_storage)
+
+ with journal_scope(journal):
+ assert current_journal() is journal
+
+ def test_outside_a_task_there_is_no_journal(self):
+ assert current_journal() is None
+
+ def test_inside_a_task_the_attempt_gets_one_journal_that_cleans_up_after_each_run(
+ self, memory_storage, monkeypatch
+ ):
+ monkeypatch.setattr(journal_module, "_task_journal", None)
+ context = {"task_instance": SimpleNamespace(id="ti-1")}
+ with (
+ mock.patch.object(journal_module, "get_current_context", autospec=True, return_value=context),
+ mock.patch.object(
+ journal_module, "build_task_storage", autospec=True, return_value=memory_storage
+ ),
+ ):
+ journal = current_journal()
+
+ assert journal is not None
+ assert journal.clean_up_after_run
+ assert current_journal() is journal
+
+ def test_a_new_attempt_in_the_same_process_gets_a_new_journal(self, memory_storage, monkeypatch):
+ monkeypatch.setattr(journal_module, "_task_journal", None)
+ ti = SimpleNamespace(id="ti-1")
+ with (
+ mock.patch.object(
+ journal_module, "get_current_context", autospec=True, return_value={"task_instance": ti}
+ ),
+ mock.patch.object(
+ journal_module, "build_task_storage", autospec=True, return_value=memory_storage
+ ),
+ ):
+ first = current_journal()
+ ti.id = "ti-2"
+
+ assert current_journal() is not first
diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_replay_cost.py b/providers/common/ai/tests/unit/common/ai/durable/test_replay_cost.py
index 4905d44c22a..da3420de9f6 100644
--- a/providers/common/ai/tests/unit/common/ai/durable/test_replay_cost.py
+++ b/providers/common/ai/tests/unit/common/ai/durable/test_replay_cost.py
@@ -15,26 +15,13 @@
# specific language governing permissions and limitations
# under the License.
"""
-Empirical check of durable replay cost accounting against ``cost_limit``.
-
-pydantic-ai's graph appends *every* model response's usage to the run's
-``RunUsage`` in ``_agent_graph.py::_append_response`` -- it cannot distinguish a
-response that came from a live model call from one ``CachingModel`` replayed
-from the durable cache; compensating for a replay is entirely on the caller.
-``CachingModel`` only does that compensation when it is given a ``replay_usage=``
-ledger (see ``durable/replay_usage.py``) -- ``AgentOperator`` always passes one
-with ``durable=True`` (around a real cross-attempt total on Airflow >= 3.3, or a
-fresh per-attempt ``RunUsage()`` otherwise), so its users never see the
-double-count described below. ``TestDurableReplayCostDuplication`` exercises the
-lower-level ``CachingModel`` API directly, with **no** ledger supplied, to pin down
-what happens without that compensation: each Airflow task attempt starts a
-fresh ``RunUsage`` (a new ``agent.run`` call), so a step that was already paid
-for in a prior, crashed attempt gets its cost added again to that fresh total.
-``TestDurableReplayCostWithSharedRunUsage`` is the contrasting case -- the same
-real ``CachingModel`` + ``DurableStorage`` + pydantic-ai ``Agent`` stack (no
-mocked cost arithmetic), but with a ``run_usage`` shared across both simulated
-attempts the way ``AgentOperator`` shares one, confirming the replayed cost is
-*not* recounted there.
+Durable replay cost accounting against ``cost_limit``, through the real stack.
+
+pydantic-ai's graph adds *every* model response's usage to the run's ``RunUsage``;
+it cannot tell a response the journal replayed from a live one. ``AirflowDurability``
+cancels each replay through a ``ReplayUsageLedger`` over the run's usage, which is the
+cross-attempt total when ``AgentOperator`` passes one as ``usage=``. These tests drive a real file-backed ``DurableStorage``, the journal and a
+pydantic-ai ``Agent`` -- no mocked cost arithmetic -- across simulated attempts.
"""
from __future__ import annotations
@@ -46,15 +33,13 @@ from unittest.mock import patch
import pytest
from pydantic_ai import Agent
-from pydantic_ai.exceptions import UsageLimitExceeded
from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart
from pydantic_ai.models.function import AgentInfo, FunctionModel
from pydantic_ai.models.wrapper import WrapperModel
from pydantic_ai.usage import RequestUsage, RunUsage, UsageLimits
-from airflow.providers.common.ai.durable.caching_model import CachingModel
-from airflow.providers.common.ai.durable.replay_usage import ReplayUsageLedger
-from airflow.providers.common.ai.durable.step_counter import DurableStepCounter
+from airflow.providers.common.ai.durable import AirflowDurability
+from airflow.providers.common.ai.durable.journal import DurableJournal, journal_scope
from airflow.providers.common.ai.durable.storage import DurableStorage
from airflow.providers.common.compat.sdk import ObjectStoragePath
@@ -85,139 +70,89 @@ class _Priced(WrapperModel):
)
-@pytest.fixture
-def durable_storage(tmp_path):
- """A real, file-backed DurableStorage -- exercises the actual JSON round-trip."""
- with patch("airflow.providers.common.ai.durable.storage._get_base_path") as mock_base:
- mock_base.return_value = ObjectStoragePath(f"file://{tmp_path.as_posix()}")
- yield DurableStorage(dag_id="dag", task_id="task", run_id="run_1", map_index=-1)
+@pytest.fixture(autouse=True)
+def journal_dir(tmp_path):
+ """Every ``DurableStorage`` built in a test reads and writes the same file-backed journal."""
+ with patch("airflow.providers.common.ai.durable.storage._get_base_path", autospec=True) as base:
+ base.return_value = ObjectStoragePath(f"file://{tmp_path.as_posix()}")
+ yield tmp_path
+
+
+def _storage() -> DurableStorage:
+ """A fresh ``DurableStorage`` for the same task instance -- what a new attempt (process) builds."""
+ return DurableStorage(dag_id="dag", task_id="task", run_id="run_1", map_index=-1)
async def _run_one_attempt(
- storage: DurableStorage,
- *,
- cost_limit: Decimal | None = None,
- counter: DurableStepCounter | None = None,
- run_usage: RunUsage | None = None,
- model: Any = None,
-):
- """Simulate one Airflow task attempt with a fresh Agent and shared cache.
-
- DurableStepCounter is fresh unless one is supplied. ``run_usage``, when given, is
- both what the ``ReplayUsageLedger`` compensates and the ``agent.run(usage=...)``
- seed -- mirrors how ``AgentOperator`` shares one ``RunUsage`` object between the two
- across attempts. ``None`` (the default) skips compensation entirely, matching a
- bare ``CachingModel`` used outside ``AgentOperator``."""
- counter = counter or DurableStepCounter()
+ *, model: Any = None, cost_limit: Decimal | None = None, run_usage: RunUsage | None = None
+) -> tuple[Any, DurableJournal]:
+ """
+ Simulate one Airflow task attempt with a fresh agent and journal over the shared file.
+
+ ``run_usage``, when given, is shared across attempts the way ``AgentOperator`` shares its
+ cross-attempt usage; otherwise each attempt starts a fresh ``RunUsage``. Either way the
+ capability keeps replays out of it.
+ """
limits = UsageLimits(cost_limit=cost_limit)
- ledger = ReplayUsageLedger(run_usage=run_usage, usage_limits=limits) if run_usage is not None else None
- caching = CachingModel(
- model or FunctionModel(_build_priced_response), storage=storage, counter=counter, replay_usage=ledger
+ journal = DurableJournal(_storage())
+ agent = Agent(
+ model or FunctionModel(_build_priced_response), name="a", capabilities=[AirflowDurability()]
)
- agent = Agent(model=caching)
- try:
- result = await agent.run("What is the answer?", usage=run_usage, usage_limits=limits)
- finally:
- if ledger is not None:
- ledger.settle()
- return result, counter
-
-
-def _reopen_storage() -> DurableStorage:
- """Build a fresh ``DurableStorage`` for the same dag/task/run -- simulates a new Airflow
- task attempt (new process) reloading the durable cache from disk via the public
- constructor, rather than reaching into the private ``_cache`` attribute."""
- return DurableStorage(dag_id="dag", task_id="task", run_id="run_1", map_index=-1)
-
+ with journal_scope(journal):
+ return await agent.run("What is the answer?", usage=run_usage, usage_limits=limits), journal
-class TestDurableReplayCostDuplication:
- """``CachingModel`` used directly, with no replay ledger -- not how ``AgentOperator``
- uses it (see the module docstring); ``TestDurableReplayCostWithSharedRunUsage`` below
- is the contrasting, ``AgentOperator``-shaped case."""
+class TestReplayedCostIsNotCountedAgain:
@pytest.mark.asyncio
- async def test_replayed_step_cost_is_recounted_on_retry_without_shared_run_usage(self, durable_storage):
- """A second attempt that only replays cached steps still reports the replayed cost
- as its own usage: pydantic-ai's graph can't tell a replay from a live call, and
- with no ledger supplied, ``CachingModel`` does not compensate (contrast with the
- shared-``run_usage`` test below)."""
- result1, counter1 = await _run_one_attempt(durable_storage)
- assert counter1.cached_model == 1
- assert counter1.replayed_model == 0
+ async def test_an_attempt_that_only_replays_reports_no_usage(self):
+ """Each attempt starts a fresh ``RunUsage``; the one that only replays counts nothing."""
+ result1, journal1 = await _run_one_attempt()
+ assert journal1.stats.recorded["model"] == 1
assert result1.usage.cost == PRICED_COST
- # New attempt: fresh process, so the cache is reloaded from disk via a new
- # DurableStorage -- this is what actually happens on an Airflow task retry.
- result2, counter2 = await _run_one_attempt(_reopen_storage())
+ result2, journal2 = await _run_one_attempt()
- # Zero new model calls this attempt ...
- assert counter2.cached_model == 0
- assert counter2.replayed_model == 1
- # ... yet the replayed step's cost is counted again, identically to attempt 1.
- assert result2.usage.cost == PRICED_COST
+ assert journal2.stats.replayed["model"] == 1
+ assert journal2.stats.recorded["model"] == 0
+ assert (result2.usage.requests, result2.usage.input_tokens, result2.usage.cost or 0) == (0, 0, 0)
@pytest.mark.asyncio
- async def test_retry_with_zero_new_spend_still_raises_usage_limit_exceeded(self, durable_storage):
- """Without a shared ``run_usage``, a retry that makes no new model calls can
- still raise UsageLimitExceeded, purely from replayed cost -- because
- check_cost() sees the run's cumulative usage, not "money spent in this
- attempt"."""
- # Attempt 1 stays comfortably under budget so it completes normally.
- await _run_one_attempt(durable_storage, cost_limit=PRICED_COST * 2)
-
- # Attempt 2 sets a limit below the already-paid-for replayed cost: zero new
- # spend, yet the replayed step alone pushes the cumulative usage over it.
- cost_limit = PRICED_COST / 2
- counter2 = DurableStepCounter()
- with pytest.raises(UsageLimitExceeded):
- await _run_one_attempt(_reopen_storage(), cost_limit=cost_limit, counter=counter2)
-
- # Zero new model calls this attempt -- the raise came purely from replayed cost.
- assert counter2.cached_model == 0
- assert counter2.replayed_model == 1
-
-
-class TestDurableReplayCostWithSharedRunUsage:
- """The ``AgentOperator``-shaped case: a single ``RunUsage`` object is both what the
- ``ReplayUsageLedger`` compensates and the ``agent.run(usage=...)`` seed across
- simulated attempts, as ``AgentOperator`` does. Contrast with
- ``TestDurableReplayCostDuplication`` above, which uses neither."""
+ async def test_a_retry_with_no_new_spend_stays_under_its_cost_limit(self):
+ await _run_one_attempt(cost_limit=PRICED_COST * 2)
+
+ # Below the cost the first attempt already paid, which the retry only replays.
+ result, journal = await _run_one_attempt(cost_limit=PRICED_COST / 2)
+
+ assert result.output == "the answer"
+ assert journal.stats.replayed["model"] == 1
@pytest.mark.asyncio
- async def test_replayed_step_cost_is_not_recounted_with_shared_run_usage(self, durable_storage):
- """The same scenario as the no-``run_usage`` test above, but with a ``RunUsage``
- shared across both simulated attempts: the replayed step's cost must not be
- added a second time to the cumulative total.
-
- The model leaves ``cost`` unset, as real providers do, so pydantic-ai prices
- each response itself after ``CachingModel`` has already written it to disk --
- the cached copy therefore comes back with ``cost=None``."""
+ async def test_replayed_cost_is_not_recounted_in_shared_run_usage(self):
+ """
+ The ``AgentOperator`` shape: one ``RunUsage`` carried across attempts.
+
+ The model leaves ``cost`` unset, as real providers do, so pydantic-ai prices each
+ response after the journal recorded it -- the recorded copy comes back with
+ ``cost=None``, and the ledger prices it the same way the graph does.
+ """
seed = RunUsage()
- model = FunctionModel(_build_unpriced_response)
- result1, counter1 = await _run_one_attempt(durable_storage, run_usage=seed, model=_Priced(model))
- assert counter1.cached_model == 1
- assert counter1.replayed_model == 0
+ model = _Priced(FunctionModel(_build_unpriced_response))
+ await _run_one_attempt(run_usage=seed, model=model)
live_cost = seed.cost
assert live_cost is not None
assert live_cost > 0
- # New attempt: fresh process (new DurableStorage), but the same cross-attempt
- # RunUsage carried forward -- what TaskStateStoreUsageBudget.load()/save() do
- # via the task state store on Airflow >= 3.3.
- result2, counter2 = await _run_one_attempt(_reopen_storage(), run_usage=seed, model=_Priced(model))
+ _, journal = await _run_one_attempt(run_usage=seed, model=model)
- # Zero new model calls this attempt ...
- assert counter2.cached_model == 0
- assert counter2.replayed_model == 1
- # ... and the cumulative cost is still just the one real charge, not double.
+ assert journal.stats.replayed["model"] == 1
assert seed.cost == live_cost
assert seed.requests == 1
assert seed.input_tokens == 100_000
-class TestDurableReplayOfContinuationChain:
+class TestReplayOfContinuationChain:
"""A continuation chain (Anthropic ``pause_turn``, OpenAI background mode) is one
- request to pydantic-ai but one ``CachingModel`` step per segment."""
+ request to pydantic-ai but one journal step per segment."""
@pytest.mark.parametrize(
("same_job", "segment_input_tokens"),
@@ -227,7 +162,7 @@ class TestDurableReplayOfContinuationChain:
],
)
@pytest.mark.asyncio
- async def test_replayed_chain_counts_nothing(self, durable_storage, same_job, segment_input_tokens):
+ async def test_replayed_chain_counts_nothing(self, same_job, segment_input_tokens):
"""Two suspended segments and a final one. Replaying the chain must add neither a
request per segment nor the segments' tokens: segments of a new response are
summed, while re-polls of the same provider response id replace each other."""
@@ -244,33 +179,31 @@ class TestDurableReplayOfContinuationChain:
)
seed = RunUsage()
- await _run_one_attempt(durable_storage, run_usage=seed, model=FunctionModel(build_segment))
+ await _run_one_attempt(run_usage=seed, model=FunctionModel(build_segment))
assert live_segments == [0, 1, 2]
live_total = (seed.requests, seed.input_tokens, seed.output_tokens)
assert live_total[0] == 1
live_segments.clear()
- _, counter = await _run_one_attempt(
- _reopen_storage(), run_usage=seed, model=FunctionModel(build_segment)
- )
+ _, journal = await _run_one_attempt(run_usage=seed, model=FunctionModel(build_segment))
assert live_segments == []
- assert counter.replayed_model == 3
+ assert journal.stats.replayed["model"] == 3
assert (seed.requests, seed.input_tokens, seed.output_tokens) == live_total
-class TestDurableStorageCostRoundTrip:
- def test_decimal_cost_survives_json_round_trip(self, durable_storage):
- """DurableStorage serializes the whole cache blob as JSON; confirm a Decimal
- ``usage.cost`` is not silently lost or coerced to float/None by that round-trip."""
- response = ModelResponse(
- parts=[TextPart(content="hi")],
- usage=RequestUsage(input_tokens=1, output_tokens=1, cost=Decimal("0.0123456789")),
- )
- durable_storage.save_model_response("model_step_0", response, fingerprint="fp")
+class TestCostRoundTrip:
+ @pytest.mark.asyncio
+ async def test_decimal_cost_survives_the_journal(self):
+ """The journal file is JSON; a Decimal ``usage.cost`` must come back a Decimal, not a float or None."""
+ cost = Decimal("0.0123456789")
+
+ def respond(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+ return ModelResponse(parts=[TextPart("hi")], usage=RequestUsage(input_tokens=1, cost=cost))
- loaded, _fingerprint = _reopen_storage().load_model_response("model_step_0")
+ await _run_one_attempt(model=FunctionModel(respond))
+ result, _ = await _run_one_attempt(model=FunctionModel(respond))
- assert loaded is not None
- assert loaded.usage.cost == Decimal("0.0123456789")
- assert isinstance(loaded.usage.cost, Decimal)
+ replayed = result.all_messages()[-1]
+ assert replayed.usage.cost == cost
+ assert isinstance(replayed.usage.cost, Decimal)
diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_replay_verification.py b/providers/common/ai/tests/unit/common/ai/durable/test_replay_verification.py
deleted file mode 100644
index 1e57ad45eaa..00000000000
--- a/providers/common/ai/tests/unit/common/ai/durable/test_replay_verification.py
+++ /dev/null
@@ -1,138 +0,0 @@
-# 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.
-"""
-End-to-end replay verification through a real pydantic-ai agent loop.
-
-Simulates the retry scenario durable execution exists for: attempt 1 fails
-partway with steps cached, attempt 2 starts a fresh counter against the same
-cache file. Replay must happen when the agent is unchanged and must NOT happen
-when the agent changed between attempts (the positional-keying staleness bug).
-"""
-
-from __future__ import annotations
-
-from unittest.mock import patch
-
-import pytest
-from pydantic_ai import Agent
-from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart, ToolCallPart
-from pydantic_ai.models.function import AgentInfo, FunctionModel
-from pydantic_ai.toolsets import FunctionToolset
-
-from airflow.providers.common.ai.durable.caching_model import CachingModel
-from airflow.providers.common.ai.durable.caching_toolset import CachingToolset
-from airflow.providers.common.ai.durable.step_counter import DurableStepCounter
-from airflow.providers.common.ai.durable.storage import DurableStorage
-from airflow.providers.common.compat.sdk import ObjectStoragePath
-
-
-@pytest.fixture
-def storage(tmp_path):
- with patch("airflow.providers.common.ai.durable.storage._get_base_path") as mock_base:
- mock_base.return_value = ObjectStoragePath(f"file://{tmp_path.as_posix()}")
- yield DurableStorage(dag_id="d", task_id="t", run_id="r")
-
-
-class AgentHarness:
- """A scripted two-step agent: one tool call, then a final answer."""
-
- def __init__(self, storage: DurableStorage, *, system_prompt: str, rate: float, fail: bool):
- self.live_model_calls = 0
- self.live_tool_calls = 0
- self.counter = DurableStepCounter()
-
- def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
- self.live_model_calls += 1
- if not any(isinstance(m, ModelResponse) for m in messages):
- return ModelResponse(parts=[ToolCallPart(tool_name="get_fx_rate", args={})])
- if fail:
- raise RuntimeError("simulated transient failure")
- returned = next(p.content for m in messages for p in m.parts if p.part_kind == "tool-return")
- return ModelResponse(parts=[TextPart(content=f"rate={returned}")])
-
- def get_fx_rate() -> float:
- """Return the USD to EUR exchange rate."""
- self.live_tool_calls += 1
- return rate
-
- # Constructor form: the @toolset.tool decorator is typed for
- # RunContext-taking functions only.
- toolset = FunctionToolset(tools=[get_fx_rate])
-
- self.agent = Agent(
- model=CachingModel(FunctionModel(model_fn), storage=storage, counter=self.counter),
- system_prompt=system_prompt,
- toolsets=[CachingToolset(wrapped=toolset, storage=storage, counter=self.counter)],
- )
-
- async def run(self) -> str:
- result = await self.agent.run("Convert 100 USD to EUR")
- return result.output
-
-
-async def failed_first_attempt(storage: DurableStorage, system_prompt: str = "currency bot") -> None:
- """Attempt 1: tool result 0.42 gets cached, then the final model call fails."""
- harness = AgentHarness(storage, system_prompt=system_prompt, rate=0.42, fail=True)
- with pytest.raises(RuntimeError, match="simulated transient failure"):
- await harness.run()
- storage._cache = None # retry runs in a new process: in-memory cache is cold
-
-
-class TestUnchangedRetryReplays:
- @pytest.mark.asyncio
- async def test_completed_steps_replay_only_failed_step_reruns(self, storage):
- await failed_first_attempt(storage)
-
- retry = AgentHarness(storage, system_prompt="currency bot", rate=0.42, fail=False)
- output = await retry.run()
-
- assert output == "rate=0.42"
- assert retry.counter.replayed_model == 1
- assert retry.counter.replayed_tool == 1
- assert retry.live_tool_calls == 0
- assert retry.live_model_calls == 1 # only the step that failed on attempt 1
-
-
-class TestChangedAgentDoesNotReplayStaleSteps:
- @pytest.mark.asyncio
- async def test_changed_system_prompt_invalidates_replay(self, storage):
- """Regression test for positional-keying staleness: a prompt tweak between
- attempts must re-run the conversation, not replay the old one."""
- await failed_first_attempt(storage)
-
- retry = AgentHarness(storage, system_prompt="careful currency bot", rate=0.99, fail=False)
- output = await retry.run()
-
- # The fixed tool ran in the new conversation; nothing stale was replayed.
- assert output == "rate=0.99"
- assert retry.counter.replayed_model == 0
- assert retry.counter.replayed_tool == 0
- assert retry.live_tool_calls == 1
-
- @pytest.mark.asyncio
- async def test_divergence_chains_to_downstream_tool_steps(self, storage):
- """A live model call mints fresh tool_call_ids, so cached tool results
- recorded under the old conversation cannot be cross-wired into the new one."""
- await failed_first_attempt(storage)
-
- retry = AgentHarness(storage, system_prompt="careful currency bot", rate=0.42, fail=False)
- await retry.run()
-
- # Same tool name and args as attempt 1, but the conversation diverged
- # at step 0 -- the tool must run live, not replay.
- assert retry.live_tool_calls == 1
- assert retry.counter.replayed_tool == 0
diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_storage.py b/providers/common/ai/tests/unit/common/ai/durable/test_storage.py
index c2670655195..aeae08f9e1a 100644
--- a/providers/common/ai/tests/unit/common/ai/durable/test_storage.py
+++ b/providers/common/ai/tests/unit/common/ai/durable/test_storage.py
@@ -20,33 +20,28 @@ import json
from unittest.mock import patch
import pytest
-from pydantic_ai.messages import (
- ModelMessagesTypeAdapter,
- ModelResponse,
- TextPart,
-)
-from pydantic_ai.usage import RequestUsage
+from airflow.providers.common.ai.durable import storage as storage_module
from airflow.providers.common.ai.durable.storage import DurableStorage
from airflow.providers.common.compat.sdk import ObjectStoragePath
+from tests_common.test_utils.config import conf_vars
+
@pytest.fixture
def tmp_cache_path(tmp_path):
- """Return a file:// path to a temporary directory for cache files."""
+ """Return a file:// path to a temporary directory for journal files."""
return f"file://{tmp_path.as_posix()}"
@pytest.fixture
def storage(tmp_cache_path):
- with patch("airflow.providers.common.ai.durable.storage._get_base_path") as mock_base:
+ with patch("airflow.providers.common.ai.durable.storage._get_base_path", autospec=True) as mock_base:
mock_base.return_value = ObjectStoragePath(tmp_cache_path)
yield DurableStorage(dag_id="test_dag", task_id="my_task", run_id="run_1", map_index=-1)
-@pytest.fixture
-def sample_response():
- return ModelResponse(parts=[TextPart(content="Hello!")])
+ENTRY = {"name": "agent__model.request", "kind": "model", "fingerprint": "fp", "payload": {"parts": []}}
class TestDurableStorageInit:
@@ -71,190 +66,72 @@ class TestDurableStorageInit:
assert a._cache_id != b._cache_id
-class TestSaveLoadModelResponse:
- def test_save_and_load_roundtrips(self, storage, sample_response):
- storage.save_model_response("model_step_0", sample_response, fingerprint="fp_abc")
-
- # Reset in-memory cache to force read from file
- storage._cache = None
- loaded, fingerprint = storage.load_model_response("model_step_0")
-
- assert loaded is not None
- assert loaded.parts[0].content == "Hello!"
- assert fingerprint == "fp_abc"
-
- def test_load_returns_none_when_no_cache(self, storage):
- assert storage.load_model_response("model_step_0") == (None, None)
-
- def test_metadata_carrying_response_roundtrips_byte_identical(self, storage):
- """Multi-step replay relies on cached responses round-tripping byte-identically:
- a later step's fingerprint includes earlier responses in history, metadata
- (usage, provider_response_id, finish_reason) and all. If a store/load cycle
- altered any of it, every multi-step replay would mismatch and re-run. Pin it."""
- resp = ModelResponse(
- parts=[TextPart(content="answer")],
- usage=RequestUsage(input_tokens=11, output_tokens=22),
- model_name="gpt-x",
- provider_response_id="resp_xyz",
- finish_reason="stop",
- )
- before = ModelMessagesTypeAdapter.dump_python([resp], mode="json")
-
- storage.save_model_response("model_step_0", resp, fingerprint="fp")
- storage._cache = None
- loaded, _ = storage.load_model_response("model_step_0")
-
- after = ModelMessagesTypeAdapter.dump_python([loaded], mode="json")
- assert after == before
-
- def test_stored_entry_is_single_encoded(self, storage, sample_response):
- """The response payload is stored as native JSON objects, not a nested
- JSON string -- the whole cache is encoded exactly once by ``_save_cache``."""
- storage.save_model_response("model_step_0", sample_response, fingerprint="fp")
-
- on_disk = json.loads(storage._get_path().read_text())
- entry = on_disk["model_step_0"]
-
- assert isinstance(entry, dict) # not a re-encoded JSON string
- assert isinstance(entry["data"], list) # not a nested JSON string
- assert entry["fingerprint"] == "fp"
-
- def test_legacy_entry_without_fingerprint_loads(self, storage, sample_response):
- """Entries written before fingerprinting (raw adapter JSON) load with a None fingerprint."""
- cache = storage._load_cache()
- cache["model_step_0"] = ModelMessagesTypeAdapter.dump_json([sample_response]).decode()
- storage._save_cache()
- storage._cache = None
-
- loaded, fingerprint = storage.load_model_response("model_step_0")
-
- assert loaded is not None
- assert loaded.parts[0].content == "Hello!"
- assert fingerprint is None
+class TestSaveLoad:
+ def test_entry_roundtrips(self, storage):
+ assert storage.save_step("k", ENTRY)
+ assert storage.load_step("k") == ENTRY
-class TestSaveLoadToolResult:
- def test_save_and_load_roundtrips(self, storage):
- storage.save_tool_result("tool_step_0", {"rows": [1, 2, 3]}, fingerprint="fp")
+ def test_a_retry_reads_what_the_previous_attempt_wrote(self, storage):
+ storage.save_step("k", ENTRY)
- storage._cache = None
- found, value, fingerprint = storage.load_tool_result("tool_step_0")
+ retry = DurableStorage(dag_id="test_dag", task_id="my_task", run_id="run_1", map_index=-1)
- assert found is True
- assert value == {"rows": [1, 2, 3]}
+ assert retry.load_step("k") == ENTRY
- def test_fingerprint_roundtrips(self, storage):
- storage.save_tool_result("tool_step_0", "result", fingerprint="fp_tool")
+ def test_missing_key_is_none(self, storage):
+ assert storage.load_step("missing") is None
- storage._cache = None
- found, value, fingerprint = storage.load_tool_result("tool_step_0")
-
- assert found is True
- assert fingerprint == "fp_tool"
-
- def test_legacy_entry_without_fingerprint_loads(self, storage):
- """Entries written before fingerprinting load with a None fingerprint."""
- cache = storage._load_cache()
- cache["tool_step_0"] = json.dumps({"__durable_cached__": True, "value": "old"})
- storage._save_cache()
- storage._cache = None
-
- found, value, fingerprint = storage.load_tool_result("tool_step_0")
-
- assert found is True
- assert value == "old"
- assert fingerprint is None
-
- def test_load_returns_false_when_no_cache(self, storage):
- found, value, fingerprint = storage.load_tool_result("tool_step_0")
- assert found is False
- assert value is None
- assert fingerprint is None
-
- def test_none_result_roundtrips(self, storage):
- storage.save_tool_result("tool_step_0", None, fingerprint="fp")
-
- storage._cache = None
- found, value, fingerprint = storage.load_tool_result("tool_step_0")
-
- assert found is True
- assert value is None
-
- def test_circular_reference_result_is_skipped_not_raised(self, storage):
- """A circular reference raises ValueError in json.dumps; it must skip the
- entry with a warning, not crash the (otherwise successful) tool step."""
+ def test_an_entry_that_cannot_be_encoded_is_skipped_and_the_rest_kept(self, storage):
circular: dict = {}
circular["self"] = circular
+ storage.save_step("good", ENTRY)
- storage.save_tool_result("tool_step_0", circular, fingerprint="fp") # must not raise
-
- found, _, _ = storage.load_tool_result("tool_step_0")
- assert found is False
-
-
-class TestSaveReturnsWhetherWritten:
- def test_written_entries_return_true(self, storage, sample_response):
- assert storage.save_model_response("model_step_0", sample_response, fingerprint="fp") is True
- assert storage.save_tool_result("tool_step_1", {"rows": [1]}, fingerprint="fp") is True
-
- def test_non_serializable_tool_result_returns_false(self, storage):
- assert storage.save_tool_result("tool_step_0", object(), fingerprint="fp") is False
-
-
-class TestMalformedEntries:
- def test_empty_data_list_degrades_to_miss(self, storage):
- """A torn entry whose data list is empty loads as a miss, not an IndexError."""
- cache = storage._load_cache()
- cache["model_step_0"] = {"fingerprint": "fp", "data": []}
- storage._save_cache()
- storage._cache = None
+ assert not storage.save_step("bad", {"payload": circular})
+ assert storage.load_step("good") == ENTRY
+ assert storage.load_step("bad") is None
- assert storage.load_model_response("model_step_0") == (None, None)
+ def test_a_non_dict_value_is_a_miss(self, storage):
+ storage._load_cache()["k"] = "legacy string entry"
- def test_entry_missing_data_key_degrades_to_miss(self, storage):
- cache = storage._load_cache()
- cache["model_step_0"] = {"fingerprint": "fp"}
- storage._save_cache()
- storage._cache = None
+ assert storage.load_step("k") is None
- assert storage.load_model_response("model_step_0") == (None, None)
+ def test_every_save_rewrites_one_file(self, storage, tmp_path):
+ storage.save_step("a", ENTRY)
+ storage.save_step("b", ENTRY)
+ files = list(tmp_path.iterdir())
+ assert len(files) == 1
+ assert set(json.loads(files[0].read_text())) == {"a", "b"}
-class TestCleanup:
- def test_cleanup_deletes_file(self, storage, sample_response):
- storage.save_model_response("model_step_0", sample_response, fingerprint="fp")
- path = storage._get_path()
- assert path.exists()
- storage.cleanup()
- assert not path.exists()
+class TestDeleteSteps:
+ def test_deleting_the_last_steps_removes_the_file(self, storage, tmp_path):
+ storage.save_step("a", ENTRY)
- def test_cleanup_on_nonexistent_file(self, storage):
- storage.cleanup() # Should not raise
+ storage.delete_steps(["a"])
+ assert list(tmp_path.iterdir()) == []
-class TestInMemoryCaching:
- def test_multiple_saves_write_single_file(self, storage, sample_response):
- storage.save_model_response("model_step_0", sample_response, fingerprint="fp")
- storage.save_tool_result("tool_step_1", "result", fingerprint="fp")
- storage.save_model_response("model_step_2", sample_response, fingerprint="fp")
+ def test_other_steps_are_kept(self, storage):
+ storage.save_step("a", ENTRY)
+ storage.save_step("b", ENTRY)
- assert "model_step_0" in storage._cache
- assert "tool_step_1" in storage._cache
- assert "model_step_2" in storage._cache
+ storage.delete_steps(["a"])
- def test_cache_survives_reload(self, storage, sample_response):
- """Simulate retry: save cache, reset in-memory, reload from file."""
- storage.save_model_response("model_step_0", sample_response, fingerprint="fp")
- storage.save_tool_result("tool_step_1", "tool result", fingerprint="fp")
+ assert storage.load_step("a") is None
+ assert storage.load_step("b") == ENTRY
- # Simulate new DurableStorage instance (as on retry)
- storage._cache = None
+ def test_deleting_when_there_is_no_file_does_not_raise(self, storage):
+ storage.delete_steps(["a"])
- loaded_response, _ = storage.load_model_response("model_step_0")
- assert loaded_response is not None
- assert loaded_response.parts[0].content == "Hello!"
- found, value, _ = storage.load_tool_result("tool_step_1")
- assert found is True
- assert value == "tool result"
+class TestMissingConfig:
+ def test_a_missing_cache_path_names_the_option(self):
+ storage_module._get_base_path.cache_clear()
+ try:
+ with conf_vars({("common.ai", "durable_cache_path"): ""}):
+ with pytest.raises(ValueError, match=r"\[common.ai\] durable_cache_path"):
+ DurableStorage(dag_id="d", task_id="t", run_id="r").load_step("k")
+ finally:
+ storage_module._get_base_path.cache_clear()
diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_task_state_store.py b/providers/common/ai/tests/unit/common/ai/durable/test_task_state_store.py
index e6990be2ffd..6454317fb9b 100644
--- a/providers/common/ai/tests/unit/common/ai/durable/test_task_state_store.py
+++ b/providers/common/ai/tests/unit/common/ai/durable/test_task_state_store.py
@@ -29,14 +29,7 @@ if not AIRFLOW_V_3_3_PLUS:
pytest.skip("task state store needs Airflow >= 3.3", allow_module_level=True)
from pydantic import JsonValue, TypeAdapter
-from pydantic_ai.messages import (
- ModelMessagesTypeAdapter,
- ModelResponse,
- TextPart,
-)
-from pydantic_ai.usage import RequestUsage
-
-from airflow.providers.common.ai.durable.base import TOOL_RESULT_SENTINEL
+
from airflow.providers.common.ai.durable.task_state_store import TaskStateStoreDurableStorage
from airflow.sdk.execution_time.context import NEVER_EXPIRE
@@ -80,193 +73,50 @@ def storage(accessor):
return TaskStateStoreDurableStorage(accessor)
-@pytest.fixture
-def sample_response():
- return ModelResponse(parts=[TextPart(content="Hello!")])
-
-
-class TestSaveLoadModelResponse:
- def test_save_and_load_roundtrips(self, storage, sample_response):
- storage.save_model_response("model_step_0", sample_response, fingerprint="fp_abc")
-
- loaded, fingerprint = storage.load_model_response("model_step_0")
-
- assert loaded is not None
- assert loaded.parts[0].content == "Hello!"
- assert fingerprint == "fp_abc"
-
- def test_stored_with_never_expire(self, storage, accessor, sample_response):
- """Cache entries must survive every retry regardless of retry_delay or retention config."""
- storage.save_model_response("model_step_0", sample_response, fingerprint="fp")
-
- assert accessor.set_retentions["model_step_0"] is NEVER_EXPIRE
-
- def test_stored_entry_is_native_json_not_a_string(self, storage, accessor, sample_response):
- storage.save_model_response("model_step_0", sample_response, fingerprint="fp")
-
- entry = accessor.store["model_step_0"]
- assert isinstance(entry, dict)
- assert isinstance(entry["data"], list)
- assert entry["fingerprint"] == "fp"
-
- def test_metadata_carrying_response_roundtrips_byte_identical(self, storage):
- """A later step fingerprints earlier responses in history, metadata and all;
- a store/load cycle that altered any of it would mismatch and re-run."""
- resp = ModelResponse(
- parts=[TextPart(content="answer")],
- usage=RequestUsage(input_tokens=11, output_tokens=22),
- model_name="gpt-x",
- provider_response_id="resp_xyz",
- finish_reason="stop",
- )
- before = ModelMessagesTypeAdapter.dump_python([resp], mode="json")
+ENTRY = {"name": "agent__model.request", "kind": "model", "fingerprint": "fp", "payload": {"parts": []}}
- storage.save_model_response("model_step_0", resp, fingerprint="fp")
- loaded, _ = storage.load_model_response("model_step_0")
- after = ModelMessagesTypeAdapter.dump_python([loaded], mode="json")
- assert after == before
+class TestSaveLoad:
+ def test_entry_roundtrips(self, storage):
+ assert storage.save_step("k", ENTRY)
- def test_load_returns_none_when_missing(self, storage):
- assert storage.load_model_response("model_step_0") == (None, None)
-
- def test_empty_data_list_degrades_to_miss(self, storage, accessor):
- accessor.store["model_step_0"] = {"fingerprint": "fp", "data": []}
- assert storage.load_model_response("model_step_0") == (None, None)
-
- def test_entry_missing_data_key_degrades_to_miss(self, storage, accessor):
- accessor.store["model_step_0"] = {"fingerprint": "fp"}
- assert storage.load_model_response("model_step_0") == (None, None)
-
-
-class TestSaveLoadToolResult:
- def test_save_and_load_roundtrips(self, storage):
- storage.save_tool_result("tool_step_0", {"rows": [1, 2, 3]}, fingerprint="fp_tool")
-
- found, value, fingerprint = storage.load_tool_result("tool_step_0")
-
- assert found is True
- assert value == {"rows": [1, 2, 3]}
- assert fingerprint == "fp_tool"
+ assert storage.load_step("k") == ENTRY
def test_stored_with_never_expire(self, storage, accessor):
- storage.save_tool_result("tool_step_0", "result", fingerprint="fp")
-
- assert accessor.set_retentions["tool_step_0"] is NEVER_EXPIRE
-
- def test_none_result_roundtrips(self, storage):
- """A cached ``None`` is a hit, distinguished from a missing entry by the sentinel."""
- storage.save_tool_result("tool_step_0", None, fingerprint="fp")
-
- found, value, _ = storage.load_tool_result("tool_step_0")
- assert found is True
- assert value is None
-
- def test_load_returns_false_when_missing(self, storage):
- assert storage.load_tool_result("tool_step_0") == (False, None, None)
+ """A retry can replay the step however long the retry delay or the retention config."""
+ storage.save_step("k", ENTRY)
- def test_non_dict_entry_is_a_miss(self, storage, accessor):
- accessor.store["tool_step_0"] = "not a dict"
- assert storage.load_tool_result("tool_step_0") == (False, None, None)
+ assert accessor.set_retentions["k"] is NEVER_EXPIRE
- def test_entry_without_sentinel_is_a_miss(self, storage, accessor):
- accessor.store["tool_step_0"] = {"value": "x"}
- assert storage.load_tool_result("tool_step_0") == (False, None, None)
- assert TOOL_RESULT_SENTINEL not in accessor.store["tool_step_0"]
+ def test_missing_key_is_none(self, storage):
+ assert storage.load_step("missing") is None
- def test_non_serializable_result_is_skipped_not_raised(self, storage, accessor):
- """A non-serializable tool result skips caching with a warning; the tool step still succeeds."""
- assert storage.save_tool_result("tool_step_0", object(), fingerprint="fp") is False # must not raise
+ def test_a_value_that_is_not_a_journal_entry_is_a_miss(self, storage, accessor):
+ accessor.store["k"] = "written by user code"
- assert "tool_step_0" not in accessor.store
- assert storage.load_tool_result("tool_step_0") == (False, None, None)
+ assert storage.load_step("k") is None
- def test_circular_reference_result_is_skipped_not_raised(self, storage, accessor):
- circular: dict = {}
- circular["self"] = circular
+ def test_a_write_the_store_rejects_is_skipped_not_raised(self, storage):
+ """The step already succeeded; a rejected write only loses its replay."""
+ assert not storage.save_step("k", {"payload": (1, 2)})
+ assert storage.load_step("k") is None
- storage.save_tool_result("tool_step_0", circular, fingerprint="fp") # must not raise
- assert "tool_step_0" not in accessor.store
-
- def test_tuple_result_is_normalized_and_cached(self, storage):
- """A tuple result (a common ``return (value, meta)``) is coerced to a list and cached.
-
- The store validates against ``JsonValue`` and rejects tuples, so an un-normalized
- write would fail a tool step that already succeeded -- normalization keeps it durable.
- """
- storage.save_tool_result("tool_step_0", (1, 2, 3), fingerprint="fp") # must not raise
-
- found, value, _ = storage.load_tool_result("tool_step_0")
- assert found is True
- assert value == [1, 2, 3]
-
- def test_non_string_keyed_dict_result_is_normalized_and_cached(self, storage):
- """A dict with non-string keys (e.g. ``DataFrame.to_dict()``) is coerced to string keys."""
- storage.save_tool_result("tool_step_0", {1: "a", 2: "b"}, fingerprint="fp") # must not raise
-
- found, value, _ = storage.load_tool_result("tool_step_0")
- assert found is True
- assert value == {"1": "a", "2": "b"}
-
-
-class TestSaveReturnsWhetherWritten:
- def test_written_entries_return_true(self, storage, sample_response):
- assert storage.save_model_response("model_step_0", sample_response, fingerprint="fp") is True
- assert storage.save_tool_result("tool_step_1", {"rows": [1]}, fingerprint="fp") is True
-
- @pytest.mark.parametrize("method", ["save_model_response", "save_tool_result"])
- def test_write_rejected_by_the_store_returns_false(self, storage, accessor, sample_response, method):
- """A store write that fails skips the entry and says so."""
-
- def reject(key, value, *, retention=None):
- raise ValueError("value exceeds the maximum size")
-
- accessor.set = reject
- value = sample_response if method == "save_model_response" else "result"
-
- assert getattr(storage, method)("step_0", value, fingerprint="fp") is False
- # Not tracked for cleanup either: a skipped key was never written.
- assert "step_0" not in storage._keys
-
-
-class TestCleanup:
- def test_cleanup_deletes_keys_written_this_run(self, storage, accessor, sample_response):
- storage.save_model_response("model_step_0", sample_response, fingerprint="fp")
- storage.save_tool_result("tool_step_1", "result", fingerprint="fp")
-
- storage.cleanup()
-
- assert "model_step_0" not in accessor.store
- assert "tool_step_1" not in accessor.store
-
- def test_cleanup_deletes_keys_only_replayed_this_run(self, accessor, sample_response):
- """A retry that replays an earlier attempt's keys (cache hits) must still clean them up."""
- TaskStateStoreDurableStorage(accessor).save_model_response(
- "model_step_0", sample_response, fingerprint="fp"
- )
-
- retry = TaskStateStoreDurableStorage(accessor)
- loaded, _ = retry.load_model_response("model_step_0") # cache hit, no re-write
- assert loaded is not None
-
- retry.cleanup()
- assert "model_step_0" not in accessor.store
-
- def test_cleanup_leaves_untracked_keys_untouched(self, storage, accessor, sample_response):
- """Cleanup deletes only the durable keys it touched, never the whole task instance namespace."""
+class TestDeleteSteps:
+ def test_deletes_only_the_given_keys(self, storage, accessor):
+ storage.save_step("a", ENTRY)
accessor.store["user_key"] = "kept"
- storage.save_model_response("model_step_0", sample_response, fingerprint="fp")
- storage.cleanup()
+ storage.delete_steps(["a"])
+
+ assert accessor.store == {"user_key": "kept"}
- assert "model_step_0" not in accessor.store
- assert accessor.store["user_key"] == "kept"
+ def test_a_failed_delete_does_not_raise(self, accessor):
+ class FailingDelete(FakeTaskStateStore):
+ def delete(self, key):
+ raise RuntimeError("api down")
- def test_cleanup_is_best_effort_on_delete_failure(self, accessor, sample_response):
- """A failing delete must not propagate out of cleanup."""
- storage = TaskStateStoreDurableStorage(accessor)
- storage.save_model_response("model_step_0", sample_response, fingerprint="fp")
- accessor.delete = lambda key: (_ for _ in ()).throw(RuntimeError("boom"))
+ storage = TaskStateStoreDurableStorage(FailingDelete())
+ storage.save_step("a", ENTRY)
- storage.cleanup() # must not raise
+ storage.delete_steps(["a"])
diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
index fb324664428..e45611283a7 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
@@ -31,7 +31,7 @@ from uuid import UUID
import pytest
from pydantic import BaseModel
-from pydantic_ai import Agent, DeferredToolRequests, Tool, ToolOutput
+from pydantic_ai import Agent, DeferredToolRequests, RunContext, Tool, ToolOutput
from pydantic_ai.capabilities import (
AbstractCapability,
CombinedCapability,
@@ -60,12 +60,9 @@ from pydantic_ai.toolsets.wrapper import WrapperToolset
from pydantic_ai.usage import RequestUsage, RunUsage, UsageLimits
from typing_extensions import TypeAliasType
-from airflow.providers.common.ai.durable.base import (
- DurableStorageProtocol,
- build_tool_step_key,
-)
-from airflow.providers.common.ai.durable.caching_toolset import CachingToolset
-from airflow.providers.common.ai.durable.step_counter import DurableStepCounter
+from airflow.providers.common.ai.durable import AirflowDurability
+from airflow.providers.common.ai.durable.base import RUN_ID_KEY, build_step_key
+from airflow.providers.common.ai.durable.journal import DurableJournal, DurableRun, journal_scope
from airflow.providers.common.ai.durable.storage import DurableStorage
from airflow.providers.common.ai.operators.agent import AgentOperator, HITLReviewLink
from airflow.providers.common.ai.sandbox.base import (
@@ -100,6 +97,7 @@ from airflow.providers.common.compat.sdk import (
from tests_common.test_utils.compat import OperatorSerialization
from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_1_PLUS, AIRFLOW_V_3_3_PLUS
+from unit.common.ai.durable.memory_storage import MemoryStorage
from unit.common.ai.sandbox.fake_tags import TaggedBackend
try:
@@ -198,42 +196,6 @@ def _build_priced_response(messages: list[ModelMessage], info: AgentInfo) -> Mod
)
-class _InMemoryDurableStorage:
- """In-memory DurableStorageProtocol backend for exercising real replay in tests.
-
- ``refuse_tool_writes`` stands in for a backend that skips a tool result
- (a store write that fails), so the step is not cached.
- """
-
- def __init__(self, *, refuse_tool_writes: bool = False):
- self.models: dict = {}
- self.tools: dict = {}
- self.refuse_tool_writes = refuse_tool_writes
-
- def save_model_response(self, key, response, *, fingerprint):
- self.models[key] = (response, fingerprint)
- return True
-
- def load_model_response(self, key):
- return self.models.get(key, (None, None))
-
- def save_tool_result(self, key, result, *, fingerprint):
- if self.refuse_tool_writes:
- return False
- self.tools[key] = (result, fingerprint)
- return True
-
- def load_tool_result(self, key):
- if key in self.tools:
- value, fingerprint = self.tools[key]
- return True, value, fingerprint
- return False, None, None
-
- def cleanup(self):
- self.models.clear()
- self.tools.clear()
-
-
class TestAgentOperatorValidation:
def test_requires_llm_conn_id(self):
# Airflow 2's BaseOperator reports a missing required argument as AirflowException.
@@ -1627,6 +1589,21 @@ class TestAgentOperatorRegenerateWithFeedback:
assert output == '{"text":"Revised","score":0.0}'
+def _tool_then_answer(fail_after_tool: list[bool], live: dict[str, int]):
+ """A model that calls ``my_tool`` once, then answers; it fails once after the tool if asked to."""
+
+ def model_fn(messages, info):
+ live["model"] += 1
+ if any(isinstance(p, ToolReturnPart) for m in messages for p in m.parts):
+ if fail_after_tool[0]:
+ fail_after_tool[0] = False
+ raise RuntimeError("transient model failure")
+ return ModelResponse(parts=[TextPart(content="done")])
+ return ModelResponse(parts=[ToolCallPart(tool_name="my_tool", args={})])
+
+ return model_fn
+
+
class TestAgentOperatorDurable:
def test_durable_param_stored(self):
op = AgentOperator(task_id="test", prompt="test", llm_conn_id="my_llm", durable=True)
@@ -1638,7 +1615,7 @@ class TestAgentOperatorDurable:
@pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="task state store backend requires Airflow >= 3.3")
def test_build_durable_storage_uses_task_state_store_on_3_3(self):
- """On Airflow >= 3.3 the cache lives in the task state store -- no durable_cache_path needed."""
+ """On Airflow >= 3.3 the journal lives in the task state store -- no durable_cache_path needed."""
# Imported inside the test: this module runs on all supported Airflow versions, but both symbols
# (and ``NEVER_EXPIRE``, pulled in by ``task_state_store``) only exist on 3.3+.
from airflow.providers.common.ai.durable.task_state_store import TaskStateStoreDurableStorage
@@ -1652,9 +1629,9 @@ class TestAgentOperatorDurable:
assert isinstance(storage, TaskStateStoreDurableStorage)
assert storage._store is accessor
- @patch("airflow.providers.common.ai.operators.agent.AIRFLOW_V_3_3_PLUS", False)
+ @patch("airflow.providers.common.ai.durable.journal.AIRFLOW_V_3_3_PLUS", False)
def test_build_durable_storage_falls_back_to_object_storage_below_3_3(self):
- """On Airflow < 3.3 the cache falls back to the ObjectStorage backend."""
+ """On Airflow < 3.3 the journal falls back to the ObjectStorage backend."""
ti = MagicMock(spec=["dag_id", "task_id", "run_id", "map_index"])
ti.configure_mock(dag_id="d", task_id="t", run_id="r", map_index=-1)
op = AgentOperator(task_id="t", prompt="p", llm_conn_id="c", durable=True)
@@ -1667,167 +1644,108 @@ class TestAgentOperatorDurable:
storage._cache_id == DurableStorage(dag_id="d", task_id="t", run_id="r", map_index=-1)._cache_id
)
- @patch("pydantic_ai.models.wrapper.infer_model", side_effect=lambda m: m)
- @patch("pydantic_ai.models.infer_model", autospec=True)
- @patch("airflow.providers.common.ai.operators.agent.AgentOperator._build_durable_storage")
- @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True)
- def test_execute_durable_wraps_model_and_cleans_up(
- self, mock_hook_cls, mock_build_storage, mock_infer, _, make_mock_run_result
- ):
- """durable=True wraps the model with CachingModel and cleans up the cache on success."""
- from airflow.providers.common.ai.durable.base import DurableStorageProtocol
-
- storage = MagicMock(spec=DurableStorageProtocol)
- # Empty cache: execute() looks up the first step before the run starts.
- storage.load_model_response.return_value = (None, None)
- mock_build_storage.return_value = storage
-
- mock_agent = MagicMock()
- mock_agent.run_sync.return_value = make_mock_run_result("ok")
- mock_agent.model = "test-model"
- mock_agent.override = MagicMock()
- mock_agent.override.return_value.__enter__ = MagicMock(return_value=None)
- mock_agent.override.return_value.__exit__ = MagicMock(return_value=False)
- mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent
-
- mock_infer.return_value = MagicMock()
+ def test_durable_attaches_airflow_durability_named_after_the_task(self):
+ op = AgentOperator(task_id="summarize", prompt="p", llm_conn_id="c", durable=True)
+ op.llm_hook = MagicMock(spec=["create_agent"])
- op = AgentOperator(task_id="test", prompt="test", llm_conn_id="my_llm", durable=True)
- result = op.execute(context=MagicMock())
-
- assert result == "ok"
- mock_agent.override.assert_called_once()
- assert "model" in mock_agent.override.call_args[1]
- storage.cleanup.assert_called_once()
-
- @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True)
- def test_execute_non_durable_does_not_wrap(self, mock_hook_cls, make_mock_run_result):
- """Default (durable=False) does not use override."""
- mock_agent = _make_mock_agent("ok", make_mock_run_result)
- mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent
-
- op = AgentOperator(task_id="test", prompt="test", llm_conn_id="my_llm")
- op.execute(context=_make_context())
-
- # run_sync called directly, no override
- mock_agent.run_sync.assert_called_once_with(
- "test", usage_limits=None, run_id="ti-1", cancellation_token=ANY, usage=ANY
- )
+ op._build_agent()
- def test_build_durable_capabilities_wraps_toolset_capability(self):
- """A ``Toolset`` capability's inner toolset is wrapped with CachingToolset;
- capabilities that are not ``Toolset`` pass through unchanged."""
- inner = FunctionToolset()
- passthrough = object()
- op = AgentOperator(task_id="t", prompt="p", llm_conn_id="c", durable=True)
+ capabilities = op.llm_hook.create_agent.call_args.kwargs["capabilities"]
+ durability = [c for c in capabilities if isinstance(c, AirflowDurability)]
+ assert [d.name for d in durability] == ["summarize"]
- result = op._build_durable_capabilities(
- [Toolset(inner), passthrough], MagicMock(spec=DurableStorageProtocol), DurableStepCounter()
+ def test_durable_uses_the_agents_own_name_when_it_has_one(self):
+ op = AgentOperator(
+ task_id="summarize", prompt="p", llm_conn_id="c", durable=True, agent_params={"name": "writer"}
)
+ op.llm_hook = MagicMock(spec=["create_agent"])
- assert isinstance(result[0], Toolset)
- assert isinstance(result[0].toolset, CachingToolset)
- assert result[0].toolset.wrapped is inner
- assert result[1] is passthrough
-
- def test_build_durable_capabilities_skips_callable_toolset_factory(self):
- """A ``Toolset`` holding a callable factory (resolved per run with
- RunContext) cannot be wrapped with CachingToolset, so it passes through."""
+ op._build_agent()
- def factory(ctx):
- return FunctionToolset()
+ capabilities = op.llm_hook.create_agent.call_args.kwargs["capabilities"]
+ assert [c.name for c in capabilities if isinstance(c, AirflowDurability)] == ["writer"]
- cap = Toolset(factory)
- op = AgentOperator(task_id="t", prompt="p", llm_conn_id="c", durable=True)
+ def test_non_durable_attaches_no_durability(self):
+ op = AgentOperator(task_id="t", prompt="p", llm_conn_id="c", cache_prompt=False)
+ op.llm_hook = MagicMock(spec=["create_agent"])
- result = op._build_durable_capabilities(
- [cap], MagicMock(spec=DurableStorageProtocol), DurableStepCounter()
- )
+ op._build_agent()
- assert result[0] is cap
+ capabilities = op.llm_hook.create_agent.call_args.kwargs["capabilities"]
+ assert not [c for c in capabilities if isinstance(c, AirflowDurability)]
@pytest.mark.parametrize("passed_as", ["capabilities", "agent_params"])
- def test_toolset_capability_logging_wraps_durable_cache(self, passed_as):
- """Concrete ``Toolset`` capabilities receive durable caching before agent creation."""
- inner = FunctionToolset()
+ def test_durable_keeps_masking_and_tool_logging_on_a_toolset_capability(self, passed_as):
+ inner = FunctionToolset(id="tools")
capability = Toolset(inner)
kwargs = (
{"capabilities": [capability]}
if passed_as == "capabilities"
else {"agent_params": {"capabilities": [capability]}}
)
- op = AgentOperator(
- task_id="t",
- prompt="p",
- llm_conn_id="c",
- durable=True,
- **kwargs,
- )
- op._durable_storage = MagicMock(spec=DurableStorageProtocol)
- op._durable_counter = DurableStepCounter()
- hook = MagicMock(spec=["create_agent"])
- op.llm_hook = hook
+ op = AgentOperator(task_id="t", prompt="p", llm_conn_id="c", durable=True, **kwargs)
+ op.llm_hook = MagicMock(spec=["create_agent"])
op._build_agent()
- capabilities = hook.create_agent.call_args.kwargs["capabilities"]
- assert isinstance(capabilities[0].toolset, CachingToolset)
- assert isinstance(capabilities[0].toolset.wrapped, MaskingToolset)
- assert capabilities[0].toolset.wrapped.wrapped is inner
- assert isinstance(capabilities[-1], ToolLoggingCapability)
-
- def test_toolset_capability_tool_replayed_on_retry(self):
- """A tool supplied via a ``Toolset`` capability is cached and replayed on a
- retry instead of re-executing. Regression: such tools bypassed the
- ``CachingToolset`` because they did not arrive via the ``toolsets=`` list."""
- calls = {"n": 0}
-
- def my_tool() -> str:
- calls["n"] += 1
- return "tool-result"
+ capabilities = op.llm_hook.create_agent.call_args.kwargs["capabilities"]
+ assert isinstance(capabilities[0].toolset, MaskingToolset)
+ assert capabilities[0].toolset.wrapped is inner
+ assert [type(c) for c in capabilities[1:]] == [
+ AirflowDurability,
+ PromptCaching,
+ ToolLoggingCapability,
+ ]
+
+ def test_durable_refuses_its_capability_passed_as_well(self):
+ with pytest.raises(ValueError, match="pass durable=True or the capability, not both"):
+ AgentOperator(
+ task_id="t", prompt="p", llm_conn_id="c", durable=True, capabilities=[AirflowDurability()]
+ )
- def model_fn(messages, info):
- saw_return = any(isinstance(p, ToolReturnPart) for m in messages for p in getattr(m, "parts", []))
- if saw_return:
- return ModelResponse(parts=[TextPart(content="done")])
- return ModelResponse(parts=[ToolCallPart(tool_name="my_tool", args={}, tool_call_id="c1")])
+ @pytest.mark.parametrize(
+ "kwargs",
+ [
+ pytest.param({"toolsets": [FunctionToolset()]}, id="toolsets"),
+ pytest.param({"agent_params": {"toolsets": [FunctionToolset()]}}, id="agent-params"),
+ pytest.param(
+ {"capabilities": [Toolset(FunctionToolset(), id="lookup")]}, id="toolset-capability"
+ ),
+ ],
+ )
+ def test_durable_refuses_a_toolset_without_an_id(self, kwargs):
+ with pytest.raises(ValueError, match="durable=True needs a unique id on every FunctionToolset"):
+ AgentOperator(task_id="t", prompt="p", llm_conn_id="c", durable=True, **kwargs)
- # Shared storage across two attempts; the second (a retry) must replay the
- # cached tool result rather than executing the tool a second time.
- storage = _InMemoryDurableStorage()
- for _ in range(2):
- op = AgentOperator(
+ def test_durable_refuses_a_function_that_builds_a_toolset(self):
+ with pytest.raises(ValueError, match="needs a unique id on every DynamicToolset"):
+ AgentOperator(
task_id="t",
- prompt="hi",
+ prompt="p",
llm_conn_id="c",
durable=True,
- enable_tool_logging=False,
- agent_params={"capabilities": [Toolset(FunctionToolset(tools=[my_tool]))]},
+ toolsets=[lambda ctx: FunctionToolset()],
)
- op._durable_storage = storage
- op._durable_counter = DurableStepCounter()
- hook = MagicMock(spec=["create_agent"])
- hook.create_agent.side_effect = lambda **kw: Agent(FunctionModel(model_fn), **kw)
- op.llm_hook = hook
- op._build_agent().run_sync("hi")
- assert calls["n"] == 1
+ def test_non_durable_accepts_a_toolset_without_an_id(self):
+ AgentOperator(task_id="t", prompt="p", llm_conn_id="c", toolsets=[FunctionToolset()])
@pytest.mark.parametrize(
- "capability",
+ "kwargs",
[
- pytest.param(lambda tool: Toolset(FunctionToolset([tool])), id="anonymous"),
- pytest.param(lambda tool: Toolset(FunctionToolset([tool]), id="lookup"), id="with-id"),
+ pytest.param(
+ lambda tool: {"toolsets": [FunctionToolset(tools=[tool], id="tools")]}, id="toolsets"
+ ),
+ pytest.param(
+ lambda tool: {"capabilities": [Toolset(FunctionToolset(tools=[tool], id="tools"))]},
+ id="toolset-capability",
+ ),
+ pytest.param(lambda tool: {"agent_params": {"tools": [tool]}}, id="agent-tools"),
],
)
- def test_retry_replays_steps_of_a_toolset_capability(self, capability):
- """
- pydantic-ai gives a capability without an ``id`` a random one per run and stamps it
- on its tools. A retry is a new run, so the model request differs only in that id;
- it must still replay; ``with-id`` is the control. The model issues fresh tool call ids,
- as a real provider does.
- """
- storage = _InMemoryDurableStorage()
+ def test_retry_replays_completed_steps(self, kwargs):
+ """However the tool reaches the agent, the retry replays it and the first model step."""
+ storage = MemoryStorage()
live = {"model": 0, "tool": 0}
fail_after_tool = [True]
@@ -1835,15 +1753,6 @@ class TestAgentOperatorDurable:
live["tool"] += 1
return "tool-result"
- def model_fn(messages, info):
- live["model"] += 1
- if any(isinstance(p, ToolReturnPart) for m in messages for p in m.parts):
- if fail_after_tool[0]:
- fail_after_tool[0] = False
- raise RuntimeError("transient model failure")
- return ModelResponse(parts=[TextPart(content="done")])
- return ModelResponse(parts=[ToolCallPart(tool_name="my_tool", args={})])
-
for try_number in (1, 2):
live.update(model=0, tool=0)
op = AgentOperator(
@@ -1852,106 +1761,103 @@ class TestAgentOperatorDurable:
llm_conn_id="c",
durable=True,
enable_tool_logging=False,
- capabilities=[capability(my_tool)],
+ **kwargs(my_tool),
)
op.llm_hook = MagicMock(spec=["create_agent"])
- op.llm_hook.create_agent.side_effect = lambda **kw: Agent(FunctionModel(model_fn), **kw)
+ op.llm_hook.create_agent.side_effect = lambda **kw: Agent(
+ FunctionModel(_tool_then_answer(fail_after_tool, live)), **kw
+ )
context = _make_context(ti=_make_ti(id=f"ti-{try_number}", try_number=try_number))
with (
patch.object(AgentOperator, "_build_durable_storage", autospec=True, return_value=storage),
pytest.raises(RuntimeError, match="transient") if try_number == 1 else nullcontext(),
):
op.execute(context=context)
- if try_number == 1:
- # Verified replay, not positional replay of unfingerprintable (None) steps.
- assert storage.models
- assert all(fingerprint is not None for _, fingerprint in storage.models.values())
# Attempt 2 replays model step 0 and the tool call; only the step that failed runs live.
assert live == {"model": 1, "tool": 0}
+ # And the journal is gone once the task succeeded.
+ assert storage.entries == {}
- def test_tool_result_refused_by_storage_is_counted_skipped_and_reruns(self):
- """A tool result the backend refuses to store is not counted as cached, and a
- retry runs the tool again instead of replaying it."""
- calls = {"n": 0}
+ def test_retry_keeps_the_first_attempts_run_id_and_conversation_id(self):
+ """Capabilities that key their state on the run id must see the same one on a replay."""
+ storage = MemoryStorage()
+ seen: list[tuple[str | None, str | None]] = []
+ fail_after_tool = [True]
+ live = {"model": 0}
- def my_tool() -> str:
- calls["n"] += 1
+ def my_tool(ctx: RunContext) -> str:
+ seen.append((ctx.run_id, ctx.conversation_id))
return "tool-result"
- def model_fn(messages, info):
- saw_return = any(isinstance(p, ToolReturnPart) for m in messages for p in getattr(m, "parts", []))
- if saw_return:
- return ModelResponse(parts=[TextPart(content="done")])
- return ModelResponse(parts=[ToolCallPart(tool_name="my_tool", args={}, tool_call_id="c1")])
-
- storage = _InMemoryDurableStorage(refuse_tool_writes=True)
- counters = []
- for _ in range(2):
+ for try_number in (1, 2):
op = AgentOperator(
task_id="t",
prompt="hi",
llm_conn_id="c",
durable=True,
enable_tool_logging=False,
- toolsets=[FunctionToolset(tools=[my_tool])],
+ toolsets=[FunctionToolset(tools=[my_tool], id="tools")],
)
- op._durable_storage = storage
- op._durable_counter = DurableStepCounter()
- hook = MagicMock(spec=["create_agent"])
- hook.create_agent.side_effect = lambda **kw: Agent(FunctionModel(model_fn), **kw)
- op.llm_hook = hook
- op._build_agent().run_sync("hi")
- counters.append(op._durable_counter)
-
- assert calls["n"] == 2
- first, retry = counters
- assert (first.cached_tool, first.skipped_tools) == (0, ["my_tool"])
- assert (retry.replayed_tool, retry.skipped_tools) == (0, ["my_tool"])
-
- def test_durable_summary_names_tools_that_were_not_cached(self, caplog):
- counter = DurableStepCounter()
- counter.cached_model = 2
- counter.cached_tool = 1
- counter.skipped_tools = ["run_query", "get_schema", "run_query"]
- counter.skipped_model = 1
- op = AgentOperator(task_id="t", prompt="p", llm_conn_id="c", durable=True)
+ op.llm_hook = MagicMock(spec=["create_agent"])
+ op.llm_hook.create_agent.side_effect = lambda **kw: Agent(
+ FunctionModel(_tool_then_answer(fail_after_tool, live)), **kw
+ )
+ context = _make_context(ti=_make_ti(id=f"ti-{try_number}", try_number=try_number))
+ with (
+ patch.object(AgentOperator, "_build_durable_storage", autospec=True, return_value=storage),
+ pytest.raises(RuntimeError) if try_number == 1 else nullcontext(),
+ ):
+ op.execute(context=context)
+ if try_number == 2:
+ pushes = {
+ c.kwargs["key"]: c.kwargs["value"]
+ for c in context["task_instance"].xcom_push.call_args_list
+ }
- with caplog.at_level("INFO"):
- op._log_durable_summary(counter)
+ assert seen == [("ti-1", "ti-1")]
+ assert pushes["run_id"] == "ti-1"
- assert (
- "Durable: replayed 0 cached steps (0 model, 0 tool), cached 3 new steps (2 model, 1 tool)"
- in caplog.messages
- )
- assert (
- "Durable: 3 tool results were not cached, and a retry runs them again: run_query (x2), get_schema"
- in caplog.messages
- )
- assert (
- "Durable: 1 model responses were not cached, and a retry re-runs them and every step after "
- "the first of them"
- ) in caplog.messages
+ def test_a_failed_retry_reports_the_durable_run_id(self):
+ """The run_id XCom of a failed attempt joins to the spans of the run it was part of."""
+ storage = MemoryStorage()
+ storage.save_step(RUN_ID_KEY, {"run_id": "ti-1"})
- def test_durable_summary_has_no_warning_when_everything_was_cached(self, caplog):
- counter = DurableStepCounter()
- counter.cached_model = 1
- op = AgentOperator(task_id="t", prompt="p", llm_conn_id="c", durable=True)
+ def model_fn(messages, info):
+ raise RuntimeError("transient model failure")
- with caplog.at_level("INFO"):
- op._log_durable_summary(counter)
+ op = AgentOperator(task_id="t", prompt="hi", llm_conn_id="c", durable=True)
+ op.llm_hook = MagicMock(spec=["create_agent"])
+ op.llm_hook.create_agent.side_effect = lambda **kw: Agent(FunctionModel(model_fn), **kw)
+ context = _make_context(ti=_make_ti(id="ti-2", try_number=2))
- assert (
- "Durable: replayed 0 cached steps (0 model, 0 tool), cached 1 new steps (1 model, 0 tool)"
- in caplog.messages
- )
- assert not [r for r in caplog.records if r.levelname == "WARNING"]
+ with (
+ patch.object(AgentOperator, "_build_durable_storage", autospec=True, return_value=storage),
+ pytest.raises(RuntimeError, match="transient"),
+ ):
+ op.execute(context=context)
+
+ pushes = {
+ c.kwargs["key"]: c.kwargs["value"] for c in context["task_instance"].xcom_push.call_args_list
+ }
+ assert pushes["run_id"] == "ti-1"
+
+ def test_non_durable_run_id_is_the_attempts_own(self, make_mock_run_result):
+ op = AgentOperator(task_id="t", prompt="hi", llm_conn_id="c")
+ agent = _make_mock_agent("ok", make_mock_run_result)
+ op.llm_hook = MagicMock(spec=["create_agent"])
+ op.llm_hook.create_agent.return_value = agent
+
+ op.execute(context=_make_context(ti=_make_ti(id="ti-2", try_number=2)))
- @patch("airflow.providers.common.ai.operators.agent.AgentOperator._build_durable_storage")
- def test_failed_run_logs_summary_naming_uncached_tools(self, mock_build_storage, caplog):
+ assert agent.run_sync.call_args.kwargs["run_id"] == "ti-2"
+ assert "conversation_id" not in agent.run_sync.call_args.kwargs
+
+ def test_failed_run_logs_summary_naming_tools_that_were_not_recorded(self, caplog):
"""The attempt that fails is the one Airflow retries, so its summary must name
- the tools that were not cached and will run again."""
- mock_build_storage.return_value = _InMemoryDurableStorage(refuse_tool_writes=True)
+ the tools that were not recorded and will run again."""
+ storage = MemoryStorage()
+ storage.refuse.add(build_step_key(0, 1))
def send_email() -> str:
return "sent"
@@ -1970,49 +1876,42 @@ class TestAgentOperatorDurable:
llm_conn_id="c",
durable=True,
enable_tool_logging=False,
- toolsets=[FunctionToolset(tools=[send_email, explode])],
+ toolsets=[FunctionToolset(tools=[send_email, explode], id="tools")],
)
- hook = MagicMock(spec=["create_agent"])
- hook.create_agent.side_effect = lambda **kw: Agent(FunctionModel(model_fn), **kw)
- op.llm_hook = hook
+ op.llm_hook = MagicMock(spec=["create_agent"])
+ op.llm_hook.create_agent.side_effect = lambda **kw: Agent(FunctionModel(model_fn), **kw)
- with caplog.at_level("INFO"), pytest.raises(RuntimeError, match="downstream failure"):
+ with (
+ patch.object(AgentOperator, "_build_durable_storage", autospec=True, return_value=storage),
+ caplog.at_level("INFO"),
+ pytest.raises(RuntimeError, match="downstream failure"),
+ ):
op.execute(context=_make_context())
+ assert "recorded 2 new steps (2 model, 0 tool, 0 other)" in caplog.text
assert (
- "Durable: replayed 0 cached steps (0 model, 0 tool), cached 2 new steps (2 model, 0 tool)"
- in caplog.messages
- )
- assert (
- "Durable: 1 tool results were not cached, and a retry runs them again: send_email"
- in caplog.messages
+ "1 tool results were not recorded, and a retry runs them again: "
+ "t__function_toolset__tools.call_tool:send_email" in caplog.text
)
- @patch("pydantic_ai.models.wrapper.infer_model", side_effect=lambda m: m)
- @patch("pydantic_ai.models.infer_model", autospec=True)
- @patch("airflow.providers.common.ai.operators.agent.AgentOperator._build_durable_storage")
- @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True)
- def test_cleanup_skipped_when_post_run_step_fails(
- self, mock_hook_cls, mock_build_storage, mock_infer, _, make_mock_run_result
- ):
- """Durable cleanup must not run if a post-run step (the message-history XCom
- push) fails, so the Airflow retry can still replay the cached steps."""
- storage = MagicMock(spec=DurableStorageProtocol)
- # Empty cache: execute() looks up the first step before the run starts.
- storage.load_model_response.return_value = (None, None)
- mock_build_storage.return_value = storage
-
- mock_agent = MagicMock(spec=["run_sync", "model", "override", "instrument"])
- mock_agent.run_sync.return_value = make_mock_run_result("ok")
- mock_agent.model = "test-model"
- mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent
-
+ def test_cleanup_skipped_when_post_run_step_fails(self):
+ """The journal must survive a post-run step (the message-history XCom push) that
+ fails, so the Airflow retry can still replay the recorded steps."""
+ storage = MemoryStorage()
op = AgentOperator(task_id="t", prompt="p", llm_conn_id="c", durable=True, message_history="[]")
- with patch.object(op, "_emit_message_history", side_effect=RuntimeError("xcom down")):
- with pytest.raises(RuntimeError, match="xcom down"):
- op.execute(context=_make_context())
+ op.llm_hook = MagicMock(spec=["create_agent"])
+ op.llm_hook.create_agent.side_effect = lambda **kw: Agent(
+ FunctionModel(lambda messages, info: ModelResponse(parts=[TextPart("ok")])), **kw
+ )
+
+ with (
+ patch.object(AgentOperator, "_build_durable_storage", autospec=True, return_value=storage),
+ patch.object(op, "_emit_message_history", autospec=True, side_effect=RuntimeError("xcom down")),
+ pytest.raises(RuntimeError, match="xcom down"),
+ ):
+ op.execute(context=_make_context())
- storage.cleanup.assert_not_called()
+ assert [step["name"] for step in storage.steps()] == ["t__model.request"]
def test_supports_durable_execution_marker(self):
assert AgentOperator._AgentOperator__supports_durable_execution is True
@@ -2150,37 +2049,32 @@ class TestAgentOperatorMessageHistory:
enable_hitl_review=True,
)
- @patch("pydantic_ai.models.wrapper.infer_model", side_effect=lambda m: m)
- @patch("pydantic_ai.models.infer_model", autospec=True)
- @patch("airflow.providers.common.ai.operators.agent.AgentOperator._build_durable_storage")
- @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True)
- def test_durable_path_also_seeds_message_history(
- self, mock_hook_cls, mock_build_storage, mock_infer, _, make_mock_run_result
- ):
- """The durable branch forwards message_history into the cached run too."""
- from airflow.providers.common.ai.durable.base import DurableStorageProtocol
-
- storage = MagicMock(spec=DurableStorageProtocol)
- # Empty cache: execute() looks up the first step before the run starts.
- storage.load_model_response.return_value = (None, None)
- mock_build_storage.return_value = storage
-
- mock_agent = MagicMock(spec=["run_sync", "model", "override", "instrument"])
- mock_agent.run_sync.return_value = make_mock_run_result("ok")
- mock_agent.model = "test-model"
- mock_agent.override.return_value.__enter__ = MagicMock(return_value=None)
- mock_agent.override.return_value.__exit__ = MagicMock(return_value=False)
- mock_hook_cls.get_hook.return_value.create_agent.return_value = mock_agent
- mock_infer.return_value = MagicMock()
-
- history_json = ModelMessagesTypeAdapter.dump_json(_sample_history()).decode()
+ def test_durable_path_also_seeds_message_history(self):
+ """The durable branch forwards message_history into the run, and keeps its conversation id."""
+ seen: list[int] = []
+ history = _sample_history()
op = AgentOperator(
- task_id="test", prompt="test", llm_conn_id="my_llm", durable=True, message_history=history_json
+ task_id="test",
+ prompt="test",
+ llm_conn_id="my_llm",
+ durable=True,
+ message_history=ModelMessagesTypeAdapter.dump_json(history).decode(),
)
- op.execute(context=MagicMock())
+ op.llm_hook = MagicMock(spec=["create_agent"])
- passed = mock_agent.run_sync.call_args.kwargs["message_history"]
- assert len(passed) == 2
+ def model_fn(messages, info):
+ seen.append(len(messages))
+ return ModelResponse(parts=[TextPart("ok")])
+
+ op.llm_hook.create_agent.side_effect = lambda **kw: Agent(FunctionModel(model_fn), **kw)
+
+ with patch.object(
+ AgentOperator, "_build_durable_storage", autospec=True, return_value=MemoryStorage()
+ ):
+ op.execute(context=_make_context())
+
+ # The two history messages, plus this run's request.
+ assert seen == [len(history) + 1]
class TestAgentOperatorCancellation:
@@ -2376,7 +2270,10 @@ class TestAgentOperatorSandboxContinuityGuards:
task_id="t",
prompt="p",
llm_conn_id="c",
- toolsets=[FunctionToolset().prefixed("fn"), CombinedToolset([FunctionToolset()])],
+ toolsets=[
+ FunctionToolset(id="a").prefixed("fn"),
+ CombinedToolset([FunctionToolset(id="b")]),
+ ],
**flag,
)
assert op.toolsets is not None
@@ -2982,7 +2879,7 @@ class TestAgentOperatorUsageBudget:
attempt 2 would seed requests=1 from attempt 1, replay step 1 (+1), reach
requests=2 >= request_limit=2, and raise before the live step 2 request is
attempted."""
- shared_storage = _InMemoryDurableStorage()
+ shared_storage = MemoryStorage()
mock_build_storage.return_value = shared_storage
call_count = {"tool": 0}
@@ -3016,7 +2913,7 @@ class TestAgentOperatorUsageBudget:
llm_conn_id="c",
durable=True,
enable_tool_logging=False,
- toolsets=[FunctionToolset(tools=[flaky_tool])],
+ toolsets=[FunctionToolset(tools=[flaky_tool], id="tools")],
usage_limits=limits,
)
with pytest.raises(RuntimeError, match="boom"):
@@ -3028,7 +2925,7 @@ class TestAgentOperatorUsageBudget:
llm_conn_id="c",
durable=True,
enable_tool_logging=False,
- toolsets=[FunctionToolset(tools=[flaky_tool])],
+ toolsets=[FunctionToolset(tools=[flaky_tool], id="tools")],
usage_limits=limits,
)
context2 = _make_context(task_state_store=task_state_store_accessor)
@@ -3063,7 +2960,7 @@ class TestAgentOperatorUsageBudget:
and ids, not ``usage``), so altering a replayed response's usage would make
every later step re-run live. Two tool calls -> two cached model steps must
both replay on the retry."""
- shared_storage = _InMemoryDurableStorage()
+ shared_storage = MemoryStorage()
mock_build_storage.return_value = shared_storage
call_count = {"tool_b": 0}
@@ -3094,7 +2991,7 @@ class TestAgentOperatorUsageBudget:
llm_conn_id="c",
durable=True,
enable_tool_logging=False,
- toolsets=[FunctionToolset(tools=[tool_a, tool_b])],
+ toolsets=[FunctionToolset(tools=[tool_a, tool_b], id="tools")],
)
with pytest.raises(RuntimeError, match="boom"):
op1.execute(context=_make_context(task_state_store=task_state_store_accessor))
@@ -3105,19 +3002,19 @@ class TestAgentOperatorUsageBudget:
llm_conn_id="c",
durable=True,
enable_tool_logging=False,
- toolsets=[FunctionToolset(tools=[tool_a, tool_b])],
+ toolsets=[FunctionToolset(tools=[tool_a, tool_b], id="tools")],
)
result = op2.execute(context=_make_context(task_state_store=task_state_store_accessor))
assert result == "done"
assert call_count["tool_b"] == 2
- counter = op2._durable_counter
- # Both step 0 (tool_a's model response) and step 1 (tool_b's) replayed from
- # cache; only the final "done" step is a genuinely new model call. If the
- # replayed response's usage were altered, step 1's fingerprint would mismatch
- # and this would instead be replayed_model=1, cached_model=2.
- assert counter.replayed_model == 2
- assert counter.cached_model == 1
+ stats = op2._durable_journal.stats
+ # Both of the model responses that asked for a tool replayed; only the final
+ # "done" step is a genuinely new model call. If the replayed response's usage
+ # were altered, the second request's fingerprint would mismatch and this would
+ # instead be one replayed and two recorded.
+ assert stats.replayed["model"] == 2
+ assert stats.recorded["model"] == 1
@pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="usage budget needs Airflow >= 3.3")
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True)
@@ -3412,7 +3309,7 @@ class _DurableBudgetScenario:
# Cache tool_a before tool_b fails; sibling cancellation in a parallel batch
# can otherwise prevent the successful result from being persisted.
- return FunctionToolset(tools=[Tool(tool_a, sequential=True), tool_b])
+ return FunctionToolset(tools=[Tool(tool_a, sequential=True), tool_b], id="tools")
def run_attempt(
self, *, max_tries=0, prompt="run", limits=None, fail_after_run=False, interrupt_at_step=None
@@ -3420,13 +3317,12 @@ class _DurableBudgetScenario:
"""Run one attempt; return ``(raised, usage XCom)``."""
ti = _make_ti(max_tries=max_tries)
context = _make_context(ti=ti, task_state_store=self.accessor)
- real_next_step = DurableStepCounter.next_step
+ real_claim = DurableRun.claim
- def advance_or_interrupt(counter):
- step = real_next_step(counter)
- if step == interrupt_at_step:
+ def claim_or_interrupt(durable_run, *args, **kwargs):
+ if durable_run.position == interrupt_at_step:
raise _make_airflow_task_timeout()
- return step
+ return real_claim(durable_run, *args, **kwargs)
op = AgentOperator(
task_id="t",
@@ -3452,7 +3348,7 @@ class _DurableBudgetScenario:
patch(
"airflow.providers.common.ai.operators.agent.PydanticAIHook", autospec=True
) as mock_hook_cls,
- patch.object(DurableStepCounter, "next_step", autospec=True, side_effect=advance_or_interrupt),
+ patch.object(DurableRun, "claim", autospec=True, side_effect=claim_or_interrupt),
post_run_failure,
):
mock_hook_cls.get_hook.return_value.create_agent.side_effect = lambda **kw: Agent(
@@ -3569,12 +3465,12 @@ class TestAgentOperatorDurableUsageBudgetEndToEnd:
scenario.assert_budget_matches_live_spend()
def test_retry_whose_replayed_tool_result_no_longer_matches(self):
- """The model step replays, so its tool calls were credited up front, but the cached
- ``tool_a`` entry no longer matches (``tool_step_1``'s fingerprint is overwritten
- here): ``tool_a`` runs live and must be counted, and checked against the limit."""
+ """The model step replays, so its tool calls were credited up front, but the recorded
+ ``tool_a`` entry no longer matches (its fingerprint is overwritten here): ``tool_a``,
+ and every step after it, runs live and must be counted, and checked against the limit."""
scenario = _DurableBudgetScenario()
scenario.run_attempt(fail_after_run=True)
- key = build_tool_step_key(1)
+ key = build_step_key(0, 1)
entry = json.loads(scenario.store[key])
entry["fingerprint"] = "stale"
scenario.store[key] = json.dumps(entry)
@@ -3584,10 +3480,11 @@ class TestAgentOperatorDurableUsageBudgetEndToEnd:
assert "tool_calls_limit" in str(raised)
assert scenario.live_tool_calls == 2
- raised, _ = scenario.run_attempt(limits=UsageLimits(request_limit=2, tool_calls_limit=3))
+ raised, _ = scenario.run_attempt(limits=UsageLimits(request_limit=3, tool_calls_limit=4))
assert raised is None
- assert (scenario.live_model_calls, scenario.live_tool_calls) == (2, 3)
+ # The model step before tool_a replayed; tool_a, tool_b and the final model step ran again.
+ assert (scenario.live_model_calls, scenario.live_tool_calls) == (3, 4)
scenario.assert_budget_matches_live_spend()
def test_retry_whose_replayed_tool_results_all_no_longer_match(self):
@@ -3598,7 +3495,7 @@ class TestAgentOperatorDurableUsageBudgetEndToEnd:
scenario = _DurableBudgetScenario()
scenario.run_attempt(fail_after_run=True)
for step in (1, 2):
- key = build_tool_step_key(step)
+ key = build_step_key(0, step)
entry = json.loads(scenario.store[key])
entry["fingerprint"] = "stale"
scenario.store[key] = json.dumps(entry)
@@ -3627,20 +3524,18 @@ class TestAgentOperatorMasksToolOutput:
@staticmethod
def _run(op: AgentOperator, storage=None) -> str:
- if storage is not None:
- op._durable_storage = storage
- op._durable_counter = DurableStepCounter()
hook = MagicMock(spec=["create_agent"])
hook.create_agent.side_effect = lambda **kw: Agent(FunctionModel(_echo_tool_result), **kw)
op.llm_hook = hook
- return op._build_agent().run_sync("hi").output
+ with journal_scope(DurableJournal(storage)) if storage is not None else nullcontext():
+ return op._build_agent().run_sync("hi").output
@staticmethod
def _dag_authors_toolset(secret: str) -> FunctionToolset:
def read_setting() -> str:
return f"api key: {secret}"
- return FunctionToolset(tools=[read_setting])
+ return FunctionToolset(tools=[read_setting], id="settings")
def test_a_toolset_passed_as_toolsets(self, registered_secret):
op = AgentOperator(
@@ -3670,8 +3565,8 @@ class TestAgentOperatorMasksToolOutput:
assert self._run(op) == "api key: ***"
- def test_a_toolset_capability_is_masked_before_the_durable_cache_stores_it(self, registered_secret):
- storage = _InMemoryDurableStorage()
+ def test_a_toolset_capability_is_masked_before_the_durable_journal_stores_it(self, registered_secret):
+ storage = MemoryStorage()
op = AgentOperator(
task_id="t",
prompt="hi",
@@ -3681,8 +3576,8 @@ class TestAgentOperatorMasksToolOutput:
)
assert self._run(op, storage) == "api key: ***"
- cached = [value for value, _ in storage.tools.values()]
- assert cached == ["api key: ***"]
+ recorded = [step["payload"] for step in storage.steps() if step["kind"] == "tool"]
+ assert recorded == [{"kind": "tool_return", "result": "api key: ***"}]
def test_a_toolset_that_masks_its_own_output_is_not_wrapped_again(self):
sql = SQLToolset("pg_default")
diff --git a/providers/common/ai/tests/unit/common/ai/utils/test_tool_metrics.py b/providers/common/ai/tests/unit/common/ai/utils/test_tool_metrics.py
index 3b1f26fd966..b22a7036ec9 100644
--- a/providers/common/ai/tests/unit/common/ai/utils/test_tool_metrics.py
+++ b/providers/common/ai/tests/unit/common/ai/utils/test_tool_metrics.py
@@ -20,15 +20,15 @@ import asyncio
from unittest.mock import MagicMock, call, patch
import pytest
-from pydantic_ai import RunContext
+from pydantic_ai import Agent, RunContext
from pydantic_ai.exceptions import ApprovalRequired, CallDeferred, ModelRetry
from pydantic_ai.models.test import TestModel
from pydantic_ai.toolsets.combined import CombinedToolset
from pydantic_ai.toolsets.function import FunctionToolset
from pydantic_ai.usage import RunUsage
-from airflow.providers.common.ai.durable.caching_toolset import CachingToolset
-from airflow.providers.common.ai.durable.step_counter import DurableStepCounter
+from airflow.providers.common.ai.durable import AirflowDurability
+from airflow.providers.common.ai.durable.journal import DurableJournal, journal_scope
from airflow.providers.common.ai.toolsets.sql import SQLToolset
from airflow.providers.common.ai.utils.tool_metrics import (
calling_framework,
@@ -36,7 +36,7 @@ from airflow.providers.common.ai.utils.tool_metrics import (
)
from airflow.providers.common.ai.utils.toolset_base import MaskingToolset, ensure_masked
-from unit.common.ai.operators.test_agent import _InMemoryDurableStorage
+from unit.common.ai.durable.memory_storage import MemoryStorage
from unit.common.ai.toolsets.test_sql import _make_mock_db_hook
from unit.common.ai.utils.test_toolset_base import _call as _call_scripted, _ScriptedToolset
@@ -57,6 +57,20 @@ def _sql_toolset(**hook_kwargs) -> SQLToolset:
return ts
+def _run_two_attempts(build_toolset, tool_name: str) -> None:
+ """Run a durable agent that calls ``tool_name`` twice over one journal: live, then replayed."""
+ storage = MemoryStorage()
+ for _ in range(2):
+ agent = Agent(
+ TestModel(call_tools=[tool_name]),
+ name="analyst",
+ toolsets=[build_toolset()],
+ capabilities=[AirflowDurability()],
+ )
+ with journal_scope(DurableJournal(storage)):
+ agent.run_sync("go")
+
+
def _call(toolset, name: str, args: dict):
async def run():
ctx = RunContext(deps=None, model=TestModel(), usage=RunUsage())
@@ -116,18 +130,7 @@ class TestToolsetsCountTheirCalls:
stats.incr.assert_not_called()
def test_a_durable_replay_is_counted_as_replayed_not_executed(self, stats):
- storage = _InMemoryDurableStorage()
- for _ in range(2):
- cached = CachingToolset(
- wrapped=ensure_masked(_sql_toolset()), storage=storage, counter=DurableStepCounter()
- )
- ctx = RunContext(deps=None, model=TestModel(), usage=RunUsage(), tool_call_id="c1")
-
- async def run(toolset=cached, ctx=ctx):
- tools = await toolset.get_tools(ctx)
- return await toolset.call_tool("list_tables", {}, ctx, tools["list_tables"])
-
- asyncio.run(run())
+ _run_two_attempts(lambda: ensure_masked(_sql_toolset()), "list_tables")
assert [c.kwargs["tags"]["outcome"] for c in stats.incr.call_args_list] == ["executed", "replayed"]
@@ -149,31 +152,13 @@ class TestOutcomes:
assert [c.kwargs["tags"]["outcome"] for c in stats.incr.call_args_list] == outcomes
def test_a_replay_through_a_wrapper_counts_the_toolset_underneath(self, stats):
- storage = _InMemoryDurableStorage()
- for _ in range(2):
- wrapped = ensure_masked(_sql_toolset().prefixed("wh"))
- cached = CachingToolset(wrapped=wrapped, storage=storage, counter=DurableStepCounter())
- ctx = RunContext(deps=None, model=TestModel(), usage=RunUsage(), tool_call_id="c1")
-
- async def run(toolset=cached, ctx=ctx):
- tools = await toolset.get_tools(ctx)
- return await toolset.call_tool("wh_list_tables", {}, ctx, tools["wh_list_tables"])
-
- asyncio.run(run())
+ _run_two_attempts(lambda: ensure_masked(_sql_toolset().prefixed("wh")), "wh_list_tables")
assert [c.kwargs["tags"]["outcome"] for c in stats.incr.call_args_list] == ["executed", "replayed"]
def test_a_replay_inside_a_combined_toolset_counts_the_toolset_it_came_from(self, stats):
- storage = _InMemoryDurableStorage()
- for _ in range(2):
- combined = CombinedToolset([_sql_toolset(), FunctionToolset([])])
- cached = CachingToolset(wrapped=combined, storage=storage, counter=DurableStepCounter())
- ctx = RunContext(deps=None, model=TestModel(), usage=RunUsage(), tool_call_id="c1")
-
- async def run(toolset=cached, ctx=ctx):
- tools = await toolset.get_tools(ctx)
- return await toolset.call_tool("list_tables", {}, ctx, tools["list_tables"])
-
- asyncio.run(run())
+ _run_two_attempts(
+ lambda: CombinedToolset([_sql_toolset(), FunctionToolset([], id="none")]), "list_tables"
+ )
assert [c.kwargs["tags"]["outcome"] for c in stats.incr.call_args_list] == ["executed", "replayed"]
diff --git a/scripts/in_container/run_provider_yaml_files_check.py b/scripts/in_container/run_provider_yaml_files_check.py
index 5ed46fe888c..2363ffbaeae 100755
--- a/scripts/in_container/run_provider_yaml_files_check.py
+++ b/scripts/in_container/run_provider_yaml_files_check.py
@@ -84,9 +84,6 @@ KNOWN_DEPRECATED_CLASSES = [
# check_all_provider_classes_are_registered's registration requirement. Contrast with
# e.g. LoggingToolset, which is documented in the toolsets how-to guide and is registered.
INTERNAL_UNREGISTERED_TOOLSET_CLASSES = {
- # Wraps a toolset with per-step result caching for durable execution; applied
- # automatically by AgentOperator, not part of the public toolsets how-to guide.
- "airflow.providers.common.ai.durable.caching_toolset.CachingToolset",
# Masks what a Dag author's toolset returns; applied automatically by AgentOperator.
"airflow.providers.common.ai.utils.toolset_base.MaskingToolset",
}
diff --git a/uv.lock b/uv.lock
index 3a470996ebd..b1bb74a45aa 100644
--- a/uv.lock
+++ b/uv.lock
@@ -4697,12 +4697,12 @@ requires-dist = [
{ name = "pydantic-ai-harness", extras = ["codemode"], marker = "extra == 'code-mode'", specifier = ">=0.24.0" },
{ name = "pydantic-ai-shields", marker = "extra == 'shields'", specifier = ">=0.3.4" },
{ name = "pydantic-ai-skills", marker = "extra == 'skills'", specifier = ">=1.2.0" },
- { name = "pydantic-ai-slim", specifier = ">=2.33.0" },
- { name = "pydantic-ai-slim", extras = ["anthropic"], marker = "extra == 'anthropic'", specifier = ">=2.33.0" },
- { name = "pydantic-ai-slim", extras = ["bedrock"], marker = "extra == 'bedrock'", specifier = ">=2.33.0" },
- { name = "pydantic-ai-slim", extras = ["google"], marker = "extra == 'google'", specifier = ">=2.33.0" },
- { name = "pydantic-ai-slim", extras = ["mcp"], marker = "extra == 'mcp'", specifier = ">=2.33.0" },
- { name = "pydantic-ai-slim", extras = ["openai"], marker = "extra == 'openai'", specifier = ">=2.33.0" },
+ { name = "pydantic-ai-slim", specifier = ">=2.36.0" },
+ { name = "pydantic-ai-slim", extras = ["anthropic"], marker = "extra == 'anthropic'", specifier = ">=2.36.0" },
+ { name = "pydantic-ai-slim", extras = ["bedrock"], marker = "extra == 'bedrock'", specifier = ">=2.36.0" },
+ { name = "pydantic-ai-slim", extras = ["google"], marker = "extra == 'google'", specifier = ">=2.36.0" },
+ { name = "pydantic-ai-slim", extras = ["mcp"], marker = "extra == 'mcp'", specifier = ">=2.36.0" },
+ { name = "pydantic-ai-slim", extras = ["openai"], marker = "extra == 'openai'", specifier = ">=2.36.0" },
{ name = "pypdf", marker = "extra == 'pdf'", specifier = ">=4.0.0" },
{ name = "python-docx", marker = "extra == 'docx'", specifier = ">=1.0.0" },
{ name = "sqlglot", marker = "extra == 'sql'", specifier = ">=30.0.0" },
@@ -4732,7 +4732,7 @@ dev = [
{ name = "openai", specifier = ">=2.47.0" },
{ name = "opensandbox", specifier = ">=1.1.0" },
{ name = "pydantic-ai-skills", specifier = ">=1.2.0" },
- { name = "pydantic-ai-slim", extras = ["mcp"], specifier = ">=2.33.0" },
+ { name = "pydantic-ai-slim", extras = ["mcp"], specifier = ">=2.36.0" },
{ name = "sqlglot", specifier = ">=30.0.0" },
]
docs = [{ name = "apache-airflow-devel-common", extras = ["docs"], editable = "devel-common" }]