This is an automated email from the ASF dual-hosted git repository. kaxil pushed a commit to branch durable-tool-approval in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 243ba5a475fe67e44be1bf93618c8336493158fa Author: Kaxil Naik <[email protected]> AuthorDate: Wed Nov 11 01:20:45 2026 +0000 Let durable agent tasks pause for tool approval AgentOperator(durable=True) and @task.agent(durable=True) on Airflow 3.3+ now pause for a tool that needs approval instead of failing with UnsupportedToolDeferralError. The decision is a step of the durable journal: the resume runs the agent again from the prompt, every step before the pause replays, and the approval step gets the reviewer's answer. A retry after that replays the decision, so the reviewer is not asked again and an approved call that already ran is not run again. Answers are kept in the task state store beside the journal and belong to the approval step that asked, the calls the reviewer saw (tool and arguments) and the toolsets at the time. The approval step runs first, applies the answers, then lets the other capabilities' handle_deferred_tool_calls handlers settle the rest, and asks about what is left, so no handler can overrule a reviewer. An approval reached inside another tool's work, such as an agent a tool runs, fails the task, since a replay of that tool would not reach it again. The journal treats a pause as a pause rather than a failure: the resume trusts every step the paused attempt recorded. --- providers/common/ai/docs/durable_execution.rst | 17 +- providers/common/ai/docs/observability.rst | 4 +- providers/common/ai/docs/recovery.rst | 4 +- providers/common/ai/docs/tool_approval.rst | 67 +- .../providers/common/ai/durable/capability.py | 38 +- .../providers/common/ai/durable/fingerprint.py | 20 + .../airflow/providers/common/ai/durable/journal.py | 7 + .../src/airflow/providers/common/ai/exceptions.py | 5 +- .../airflow/providers/common/ai/operators/agent.py | 544 +++++++++++--- .../unit/common/ai/durable/test_fingerprint.py | 41 +- .../tests/unit/common/ai/durable/test_journal.py | 76 ++ .../tests/unit/common/ai/operators/test_agent.py | 5 +- .../ai/operators/test_agent_tool_approval.py | 829 ++++++++++++++++++++- 13 files changed, 1526 insertions(+), 131 deletions(-) diff --git a/providers/common/ai/docs/durable_execution.rst b/providers/common/ai/docs/durable_execution.rst index 01e6dc07408..ce5e9ecc2c0 100644 --- a/providers/common/ai/docs/durable_execution.rst +++ b/providers/common/ai/docs/durable_execution.rst @@ -240,7 +240,8 @@ instead. 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 step before it also replayed. A step that settles tool calls waiting for approval is +fingerprinted by those calls, so it replays only for them. The journal records: @@ -255,9 +256,12 @@ The journal records: summary follows a tool set to run on every attempt. 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 +``CallDeferred`` replays as the same signal. With ``AgentOperator`` on Airflow 3.3 and +later, a tool that needs approval pauses the task (except with code mode or +``enable_hitl_review``), and the reviewer's answer is recorded as a step, so a retry does +not ask again, and replays an approved call that returned (:ref:`tool-approval-durable` +has the exceptions). A deferred tool (``CallDeferred``) fails the task with +``UnsupportedToolDeferralError``, as it would without ``durable``. A tool that raises any other exception runs again on retry; once it has failed the run on two attempts in a row, the next retry asks the model again, as described above. @@ -514,9 +518,8 @@ offered the tool, and the task is not retried: Limitations ----------- -- ``durable=True`` cannot be combined with ``enable_hitl_review``, a ``SandboxToolset`` - or tool approval. ``AgentOperator`` raises ``ValueError`` for the first two when the Dag - is parsed. +- ``durable=True`` cannot be combined with ``enable_hitl_review`` or a ``SandboxToolset``. + ``AgentOperator`` raises ``ValueError`` for both 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. diff --git a/providers/common/ai/docs/observability.rst b/providers/common/ai/docs/observability.rst index 9673b5cb240..3021f6ca10e 100644 --- a/providers/common/ai/docs/observability.rst +++ b/providers/common/ai/docs/observability.rst @@ -81,8 +81,8 @@ How it works from the cache are not part of ``usage``; with ``enable_hitl_review`` the ``run_id`` and ``usage`` reflect the initial model run, not the human-feedback regenerations; and a run that resumes after a tool approval (see - :doc:`tool_approval`) continues as ``<task-instance id>-resumed-<n>`` after - pause ``n``, which is the ``run_id`` the operator pushes, with ``usage`` + :doc:`tool_approval`) without ``durable`` continues as ``<task-instance id>-resumed-<n>`` + after pause ``n``, which is the ``run_id`` the operator pushes, with ``usage`` covering every part. Airflow 2 has no task-instance id, so there the key is ``<dag_id>/<run_id>/<task_id>/<map_index>/<try_number>``, and spans carry the five diff --git a/providers/common/ai/docs/recovery.rst b/providers/common/ai/docs/recovery.rst index 481d1f047bd..d9fb479fc31 100644 --- a/providers/common/ai/docs/recovery.rst +++ b/providers/common/ai/docs/recovery.rst @@ -159,5 +159,5 @@ What is not covered Airflow 3.3 and later (:ref:`howto/frameworks:strands-durable`), and a retried ``LLMBatchOperator`` re-attaches to its running batch (:ref:`llm-batch-reattach`). - A ``ManagedAgentToolset`` runs again on every attempt unless it is built with - ``replayable=True``. Durable execution cannot be combined with ``enable_hitl_review``, - a ``SandboxToolset`` or tool approval; see the limitations in :doc:`durable_execution`. + ``replayable=True``. Durable execution cannot be combined with ``enable_hitl_review`` + or a ``SandboxToolset``; see the limitations in :doc:`durable_execution`. diff --git a/providers/common/ai/docs/tool_approval.rst b/providers/common/ai/docs/tool_approval.rst index 9baa0662483..de849125d59 100644 --- a/providers/common/ai/docs/tool_approval.rst +++ b/providers/common/ai/docs/tool_approval.rst @@ -97,9 +97,11 @@ request, such as a second reviewer who never reloaded the page, is not applied t calls waiting now: the task asks about them again under a new number. An answer sent through the REST API without that field is applied to the waiting request. -A retry or a clear asks again for the calls the new attempt makes, including ones a -reviewer approved in an earlier attempt that then failed. The approved call already -ran once, so keep gated tools idempotent. +Without ``durable=True``, a retry or a clear asks again for the calls the new attempt +makes, including ones a reviewer approved in an earlier attempt that then failed. The +approved call already ran once, so keep gated tools idempotent. With ``durable=True`` a +retry replays an approved call that returned instead of asking again +(:ref:`next section <tool-approval-durable>`); a call that failed partway still runs again. On Airflow 3.3 a task instance can ask only once per Dag run, retries and clears included. Airflow 3.3 keeps the first request's subject and body when the same task @@ -109,6 +111,42 @@ while asking about the new one. The task fails with agent request the gated calls in one step, or give each irreversible action its own task. +.. _tool-approval-durable: + +With durable execution +---------------------- + +With ``durable=True`` (see :doc:`durable_execution`), the reviewer's answer is a step in +the task's journal. When the reviewer answers, the task runs the agent again from the +prompt: the model requests and tool calls before the pause replay from the journal +instead of running again, and the approval step gets the answer. The resumed run's +``Durable:`` line in the task log counts what replayed. With ``@task.agent``, the function +that returns the prompt runs again on each resume, as it does on a retry. + +Every answer is kept in the task state store beside the journal, and deleted with it when +the task succeeds, so a retry uses them too. An answer belongs to the approval request it was +given for: the run's first, second or later approval step, counting steps that replay, and +the calls the reviewer saw there, matched by tool and arguments. When a retry reaches that +step with those calls, the reviewer is not asked again, an approved call that already ran +replays instead of running again, a rejected call stays rejected, and on Airflow 3.3 the retry +does not fail with ``ToolApprovalAlreadyRequestedError``. A later step that asks for the same +tool and arguments is asked about again. A clear that deletes the task's state +(:ref:`durable-clear`) starts the agent over, and it asks again. + +The run can reach a step with other calls, or a different number of approval steps before +it: after the Dag file changed during the pause, or when a tool that +:ref:`runs on every attempt <durable-tool-opt-out>` returned something new, so the model +requests after it ran again and the calls they chose run again, approved ones included. The +run then asks about the calls no answer covers (on Airflow 3.3 the task fails instead, as for +any second request). Answered calls the run never reaches do not run; when a resumed run +finishes without them, the task log names their tools. + +The approval step also runs the other capabilities' ``handle_deferred_tool_calls`` +handlers, such as the decision model below, on the calls no answer covers, so a handler +cannot overrule a reviewer. A handler runs again, model call included, whenever the step +runs live with calls no answer covers, such as on a resume after a pause that left some calls +to it. A retry after the step was recorded replays its decisions. + Timeouts and who decides ------------------------ @@ -205,7 +243,7 @@ Requirements and limits - Airflow 3.3 or later. On older versions a tool marked for approval fails the task, as it did before. -- Not together with ``durable=True``, ``enable_hitl_review=True``, a ``CodeMode`` +- Not together with ``enable_hitl_review=True``, a ``CodeMode`` capability, or a ``SandboxToolset`` that provisions its own sandbox. Each assumes the run finishes in one go; that sandbox, for one, is destroyed when the @@ -217,17 +255,24 @@ Requirements and limits (:ref:`sandbox-attach`). - Tools that hand work to an external system (pydantic-ai's ``CallDeferred``) are not supported; the task fails without retrying. -- The conversation so far, including tool results, is kept in the task's - :doc:`task state store <apache-airflow:core-concepts/task-state-store>` while - the task waits. It is deleted when the resumed run succeeds or fails, and replaced +- With ``durable=True``, a tool that needs approval has to be called by the task's agent + itself. One reached from inside another tool, such as the task's agent run again by a tool + (``ctx.agent.run``), fails the task with ``UnsupportedToolDeferralError`` unless the tool + catches it: once that tool is recorded, a retry replays it without reaching the approval + again. +- Without ``durable=True``, the conversation so far, including tool results, is kept + in the task's :doc:`task state store <apache-airflow:core-concepts/task-state-store>` + while the task waits. It is deleted when the resumed run succeeds or fails, and replaced when it pauses again; a later try deletes one left behind by a try that ended while waiting. With a ``tool_approval_timeout``, it also expires a day after the timeout. - If a templated connection id (see :ref:`sql-toolset-templated-connection`) renders differently when the task resumes, the task fails rather than run an approved - call against a connection the reviewer did not see. The check compares toolset + call against a connection the reviewer did not see. With ``durable=True``, a retry + that renders it differently does not apply the earlier answers and asks again. The check compares toolset ids, which name the connection for ``SQLToolset``, ``MCPToolset`` and ``HookToolset``; it does not notice an edit to the connection itself. -- Each resumed run gets its own pydantic-ai ``run_id``, - ``<task-instance id>-resumed-<n>`` when it resumes from pause ``n``. The ``run_id`` XCom - holds the last one, and the ``usage`` XCom covers every part of the run. +- Without ``durable=True``, each resumed run gets its own pydantic-ai ``run_id``, + ``<task-instance id>-resumed-<n>`` when it resumes from pause ``n``; with it, the resumed + run keeps the durable run's id. The ``run_id`` XCom holds the last one, and the ``usage`` + XCom covers every part of the run. 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 index 6dd5bc249e8..2b4f7c5fef6 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/durable/capability.py +++ b/providers/common/ai/src/airflow/providers/common/ai/durable/capability.py @@ -51,13 +51,17 @@ from pydantic_ai.exceptions import ( UnexpectedModelBehavior, ) from pydantic_ai.messages import ModelResponse, ToolReturn -from pydantic_ai.tools import AgentDepsT +from pydantic_ai.tools import AgentDepsT, DeferredToolRequests from pydantic_ai.toolsets import AbstractToolset, DynamicToolset, FunctionToolset, ToolsetTool from pydantic_ai.toolsets.wrapper import WrapperToolset from pydantic_core import PydanticSerializationError from airflow.providers.common.ai.durable.base import TOOL_CONFIG_KEY -from airflow.providers.common.ai.durable.fingerprint import fingerprint_model_request, fingerprint_tool_call +from airflow.providers.common.ai.durable.fingerprint import ( + fingerprint_deferred_tool_requests, + fingerprint_model_request, + fingerprint_tool_call, +) from airflow.providers.common.ai.durable.journal import Unrecordable, current_journal from airflow.providers.common.ai.durable.replay_usage import ReplayUsageLedger, is_successful_tool_payload from airflow.providers.common.ai.exceptions import DurableJournalError, DurableToolMetadataError @@ -65,6 +69,11 @@ 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 from airflow.providers.common.ai.utils.toolsets import iter_toolsets +from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_3_PLUS + +if AIRFLOW_V_3_3_PLUS: + # Waiting for input parks a task in the awaiting_input state, which older Airflow versions do not have. + from airflow.sdk.exceptions import TaskAwaitingInput if TYPE_CHECKING: from pydantic_ai.agent import EventStreamHandler @@ -86,11 +95,12 @@ __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. +# of the projections fingerprinted here. _ModelRequestKey: TypeAlias = ( "tuple[str | None, list[ModelMessage], ModelSettings | None, ModelRequestParameters, RunContext[Any]]" ) _ToolCallKey: TypeAlias = "tuple[str, dict[str, Any], RunContext[Any]]" +_CapabilityOperationKey: TypeAlias = "tuple[str | None, dict[str, Any], RunContext[Any]]" try: from pydantic_ai.mcp import MCPToolset @@ -475,6 +485,10 @@ class AirflowDurability(BaseDurabilityCapability[AgentDepsT]): with _active_run(_ActiveRun(durable_run, ledger)): result = await super().wrap_run(ctx, handler=handler) except BaseException as error: + if AIRFLOW_V_3_3_PLUS and isinstance(error, TaskAwaitingInput): + # A step is waiting for a person, such as a tool approval: the task resumes from here. + durable_run.pause() + raise # pydantic-ai raises UnexpectedModelBehavior when the model's answers fail its checks # too often (a tool's ModelRetry, output validation): the answers are recorded, so a # retry has to ask the model again rather than replay them into the same error. @@ -602,7 +616,9 @@ class _AirflowOperationBackend(CallableOperationBackend[_StepMode]): # 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) + step = self._claim( + active, operation_id, name, kind="other", fingerprint=_capability_fingerprint(cache_key) + ) if step.replayed: active.ledger.record_capability_replay(step.payload) return step.payload @@ -816,6 +832,20 @@ def _is_journaled_tool_payload(payload: object) -> bool: return False +def _capability_fingerprint(cache_key: tuple[object, ...]) -> str | None: + """ + Fingerprint a capability operation that settles deferred tool calls, such as an approval, by those calls. + + Its answer is about these calls, so it must not replay for others. Any other capability + operation has no fingerprint, and replays by name and position. + """ + _, arguments, _ = cast("_CapabilityOperationKey", cache_key) + for argument in arguments.values(): + if isinstance(argument, DeferredToolRequests): + return fingerprint_deferred_tool_requests(argument) + return None + + 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]) 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 b6a7297dd2f..14c71e03b6c 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 @@ -118,6 +118,7 @@ if TYPE_CHECKING: from pydantic_ai.messages import ModelMessage from pydantic_ai.settings import ModelSettings + from pydantic_ai.tools import DeferredToolRequests log = get_task_logger() @@ -641,3 +642,22 @@ def fingerprint_tool_call( step=step, ) return None + + +def fingerprint_deferred_tool_requests(requests: DeferredToolRequests) -> str | None: + """ + Fingerprint tool calls waiting to be settled, such as approved, from each call's :func:`fingerprint_tool_call`. + + The answer recorded for these calls then replays only for calls with the same names, + arguments and call ids, which a replayed model response keeps. ``None`` when any call + cannot be fingerprinted. + """ + calls = { + kind: [ + fingerprint_tool_call(call.tool_name, call.args_as_dict(), call.tool_call_id) for call in pending + ] + for kind, pending in (("approvals", requests.approvals), ("calls", requests.calls)) + } + if any(fingerprint is None for pending in calls.values() for fingerprint in pending): + return None + return _digest(calls) 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 index f77e25cd154..4b5850074c9 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/durable/journal.py +++ b/providers/common/ai/src/airflow/providers/common/ai/durable/journal.py @@ -928,6 +928,13 @@ def build_task_storage(context: Context) -> DurableStorageProtocol: _ACTIVE_JOURNAL: ContextVar[DurableJournal | None] = ContextVar("commonai_durable_journal", default=None) _EXECUTING_STEP: ContextVar[JournalStep | None] = ContextVar("commonai_durable_step", default=None) + + +def executing_step() -> JournalStep | None: + """Return the journal step whose work is running now, such as a tool call, or ``None`` outside one.""" + return _EXECUTING_STEP.get() + + # 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. 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 201445725db..ed5a2a0f725 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/exceptions.py +++ b/providers/common/ai/src/airflow/providers/common/ai/exceptions.py @@ -48,7 +48,8 @@ class ToolApprovalError(RuntimeError): Raised when a run paused for tool approval cannot resume safely. The saved transcript, or the agent's rendered toolset ids, no longer match what the - reviewer saw. A retry starts a fresh run. + reviewer saw. A retry runs the agent again and asks the reviewer again; on Airflow 3.3, + where a task instance asks once, the retry fails instead. """ @@ -68,7 +69,7 @@ class UnsupportedToolDeferralError(AirflowFailException): Raised when an agent defers a tool call that nothing in the task can resolve. In ``AgentOperator``, either the tool hands its work to an external system, or it needs - approval where approval is not available (before Airflow 3.3, or with ``durable``, + approval where approval is not available (before Airflow 3.3, or with ``enable_hitl_review``, code mode or a ``SandboxToolset``). In ``run_durably`` for the OpenAI Agents SDK, a tool call needs approval and the run cannot pause the task for it, such as before Airflow 3.4 or with a ``session``. A retry would repeat the same call, so 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 61e44e4fc1d..d78a7ebe4da 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 @@ -22,16 +22,25 @@ import copy import hashlib import json import sys -from collections.abc import Callable, Iterable, Sequence -from dataclasses import replace +from collections.abc import Awaitable, Callable, Iterable, Sequence +from dataclasses import dataclass, field, replace from datetime import timedelta from functools import cached_property -from typing import TYPE_CHECKING, Any, ClassVar, Literal, NoReturn +from typing import TYPE_CHECKING, Any, ClassVar, Literal, NoReturn, TypeGuard from pydantic import TypeAdapter from pydantic_ai import DeferredToolRequests, DeferredToolResults, ToolDenied -from pydantic_ai.capabilities import AbstractCapability, Toolset, WrapperCapability +from pydantic_ai.capabilities import ( + AbstractCapability, + CapabilityOrdering, + Toolset, + WrapperCapability, + durable_operation, +) from pydantic_ai.messages import ModelMessagesTypeAdapter, ModelResponse, ToolCallPart + +# Resolved at run time: pydantic-ai reads the durable operation's signature to build its schema. +from pydantic_ai.tools import RunContext # noqa: TC002 from pydantic_ai.toolsets import DynamicToolset from pydantic_ai.toolsets.abstract import AbstractToolset from pydantic_ai.usage import RunUsage @@ -42,7 +51,13 @@ from airflow.providers.common.ai.durable.capability import ( _get_chosen_toolset_id, _give_durable_ids, ) -from airflow.providers.common.ai.durable.journal import DurableJournal, build_task_storage, journal_scope +from airflow.providers.common.ai.durable.fingerprint import fingerprint_tool_call +from airflow.providers.common.ai.durable.journal import ( + DurableJournal, + build_task_storage, + executing_step, + journal_scope, +) from airflow.providers.common.ai.exceptions import ( ToolApprovalAlreadyRequestedError, ToolApprovalError, @@ -127,6 +142,9 @@ if TYPE_CHECKING: # index, so the marker survives retries and clears, as the task instance's approval request does. _TOOL_APPROVAL_TRANSCRIPT_KEY = "common_ai_tool_approval_transcript" _TOOL_APPROVAL_REQUESTED_KEY = "common_ai_tool_approval_requested" +# Under durable, every answer a reviewer gave this task instance's agent run, kept with the journal +# and deleted with it, so a resume, a later pause or a retry applies each to the calls it was for. +_TOOL_APPROVAL_ANSWERS_KEY = "common_ai_tool_approval_answers" # Airflow 3.4 replaces a repeat approval request's details and clears the earlier answer, so a # run can pause again; see ``_pause_for_tool_approval`` for why 3.3 cannot. _REPEAT_APPROVAL_SUPPORTED = AIRFLOW_V_3_4_PLUS @@ -137,6 +155,11 @@ _APPROVAL_REQUEST_PARAM = "request" # How long the transcript outlives a timed pause, so a resume that runs late still finds it. _TRANSCRIPT_RETENTION_MARGIN = timedelta(days=1) _RUN_USAGE_ADAPTER: TypeAdapter[RunUsage] = TypeAdapter(RunUsage) +_STORED_ANSWERS_ADAPTER: TypeAdapter[list[dict[str, Any]]] = TypeAdapter(list[dict[str, Any]]) +# Names the durable approval step in the journal, so it must never change. +_TOOL_APPROVAL_CAPABILITY_ID = "airflow_tool_approval" +# The method a durable run paused for tool approval resumes with. +_DURABLE_APPROVAL_RESUME = "resume_after_durable_tool_approval" class HITLReviewLink(BaseOperatorLink): @@ -238,11 +261,134 @@ def _answered_request(event: dict[str, Any]) -> int | None: return -1 -def _is_awaiting_input(error: BaseException) -> bool: +def _is_awaiting_input(error: BaseException) -> TypeGuard[TaskAwaitingInput]: """Whether ``error`` is the pause for tool approval, which only Airflow 3.3+ raises.""" return AIRFLOW_V_3_3_PLUS and isinstance(error, TaskAwaitingInput) +def _refuse_external_calls(requests: DeferredToolRequests) -> None: + """Fail on tool calls that hand their work to an external system, which AgentOperator cannot resolve.""" + if requests.calls: + raise UnsupportedToolDeferralError( + "The agent called tools that need external execution " + f"({', '.join(call.tool_name for call in requests.calls)}); AgentOperator only " + "supports tools that need approval." + ) + + +def _call_content(call: ToolCallPart) -> str | None: + """Fingerprint what a reviewer sees of a call, its tool and arguments, but not the id the model gave it.""" + return fingerprint_tool_call(call.tool_name, call.args_as_dict(), None) + + +@dataclass +class _ToolApprovalAnswer: + """ + A reviewer's answer, and what it was given for. + + :param step: Which approval step of the run asked, counting replayed ones, so the answer + never reaches another step, such as a later call with the same arguments. + :param calls: The calls the reviewer saw, as ``(tool name, content fingerprint)``. + :param toolset_ids: The agent's rendered toolset ids when the reviewer answered. + """ + + step: int + calls: list[tuple[str, str | None]] + toolset_ids: list[str] + approval: bool | ToolDenied + + def take(self, pending: Sequence[ToolCallPart]) -> dict[str, bool | ToolDenied]: + """ + Return the answer for each pending call it was given for, by the call's current id. + + A call matches on its tool and arguments, so a model request asked again on resume, whose + calls get new ids, still gets the answer, and a different call with a reused id does not. + Each answered call is used once: a step asked again can hold the same call twice where + the reviewer saw it once. + """ + settled: dict[str, bool | ToolDenied] = {} + for call in pending: + answered = (call.tool_name, _call_content(call)) + if answered[1] is not None and answered in self.calls: + self.calls.remove(answered) + settled[call.tool_call_id] = self.approval + return settled + + def dump(self) -> dict[str, Any]: + denied = self.approval.message if isinstance(self.approval, ToolDenied) else None + return { + "step": self.step, + "calls": [list(call) for call in self.calls], + "toolset_ids": self.toolset_ids, + "denied": denied, + } + + @classmethod + def load(cls, stored: dict[str, Any]) -> _ToolApprovalAnswer: + denied = stored["denied"] + return cls( + step=stored["step"], + calls=[(name, fingerprint) for name, fingerprint in stored["calls"]], + toolset_ids=stored["toolset_ids"], + approval=True if denied is None else ToolDenied(denied), + ) + + +@dataclass +class _DurableToolApproval(AbstractCapability[Any]): + """ + Settles the tool calls a durable run is waiting to have approved, as a step of its journal. + + It runs first (outermost) and settles every call or parks the task, so the whole decision is + one step, ``decide``: the reviewers' answers, then the other capabilities' + ``handle_deferred_tool_calls`` handlers, then a request for whatever is left. A parked step + is not recorded, so the resumed run, which replays every step before it, reaches it again + with the answer. Once recorded, a retry replays the decision instead of asking anyone again, + and replays the approved calls instead of running them twice. + + :param next_step: Number the approval step about to run for these calls, replayed or not, + or return ``None`` when ``decide`` itself is asking the other handlers. + :param settle: The body of ``decide``. + """ + + next_step: Callable[[DeferredToolRequests], int | None] = field(repr=False) + settle: Callable[[RunContext[Any], DeferredToolRequests, int], Awaitable[DeferredToolResults]] = field( + repr=False + ) + + @classmethod + def get_serialization_name(cls) -> str | None: + return None + + def get_ordering(self) -> CapabilityOrdering: + # Ahead of any handler the Dag author ordered first, so none decides a call before the reviewer. + return CapabilityOrdering(position="outermost") + + async def handle_deferred_tool_calls( + self, ctx: RunContext[Any], *, requests: DeferredToolRequests + ) -> DeferredToolResults | None: + if (step := self.next_step(requests)) is None: + return None + return await self.decide(ctx, requests=requests, step=step) + + @durable_operation(name="tool_approval") + async def decide( + self, ctx: RunContext[Any], *, requests: DeferredToolRequests, step: int + ) -> DeferredToolResults | None: + return await self.settle(ctx, requests, step) + + +def _only_for(requests: DeferredToolRequests, results: DeferredToolResults) -> DeferredToolResults: + """Keep the results for the calls in ``requests``, so a handler cannot settle a call it was not given.""" + approval_ids = {call.tool_call_id for call in requests.approvals} + call_ids = {call.tool_call_id for call in requests.calls} + return DeferredToolResults( + approvals={i: result for i, result in results.approvals.items() if i in approval_ids}, + calls={i: result for i, result in results.calls.items() if i in call_ids}, + metadata={i: meta for i, meta in results.metadata.items() if i in approval_ids | call_ids}, + ) + + def _declares_agent_template_fields(toolset: Any) -> bool: """Whether *toolset*, or a toolset it wraps or combines, has connection IDs to render.""" return isinstance(toolset, AbstractToolset) and any( @@ -474,7 +620,9 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa 3.4+, each later step that calls a marked tool pauses again; on Airflow 3.3 a task instance asks at most once per Dag run, across retries and clears, and a second request fails the task. ``usage_limits`` applies to every side of every pause. - Not available together with ``durable``, ``enable_hitl_review``, a ``CodeMode`` + With ``durable``, the answer is recorded in the journal and the resumed run replays + up to it, so a retry neither asks again nor runs an approved call twice. + Not available together with ``enable_hitl_review``, a ``CodeMode`` capability, or a ``SandboxToolset`` that provisions its own sandbox; there, a tool that requires approval fails the task as before, except one called from inside ``CodeMode``'s ``run_code``, @@ -584,6 +732,16 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa self._run_usage_base: RunUsage = RunUsage() # Which tool approval pause of this try the run is resuming from; 0 before the first. self._approval_pause_number = 0 + # Under durable: every answer reviewers gave this run, read from the task state store, and + # whether the store holds any; the one this resume brought, until a call takes it; the + # approval steps the run has reached; whether an approval step is asking the other + # handlers; and the run's usage when it paused, with this attempt's share of it. + self._tool_approval_answers: list[_ToolApprovalAnswer] = [] + self._tool_approval_answers_stored = False + self._new_tool_approval_answer: _ToolApprovalAnswer | None = None + self._approval_steps = 0 + self._asking_other_handlers = False + self._usage_at_pause: tuple[RunUsage, RunUsage] | None = None # The answer the HITL reviewer is looking at, as the agent returned it. self._output_under_review: Any = None @@ -824,8 +982,12 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa } return PydanticAIHook.get_hook(self.llm_conn_id, hook_params=hook_params) - def _build_agent(self) -> Agent[object, Any]: - """Build and return a pydantic-ai Agent from the operator's config.""" + def _build_agent(self, *, approval_context: Context | None = None) -> Agent[object, Any]: + """ + Build and return a pydantic-ai Agent from the operator's config. + + :param approval_context: The task's context, where a durable run can pause for tool approval. + """ extra_kwargs = dict(self.agent_params) passed_through = extra_kwargs.pop("capabilities", None) if passed_through is not None and self.capabilities is not None: @@ -853,6 +1015,18 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa 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 approval_context is not None and self._supports_tool_approval(): + # First, so the other handlers run inside its step (see _DurableToolApproval). + capabilities.insert( + 0, + _DurableToolApproval( + next_step=self._next_approval_step, + settle=lambda ctx, requests, step: self._settle_tool_approvals( + approval_context, ctx, requests, step + ), + id=_TOOL_APPROVAL_CAPABILITY_ID, + ), + ) if self.cache_prompt: capabilities.append(PromptCaching()) if self.enable_tool_logging and not _contains_capability(capabilities, ToolLoggingCapability): @@ -871,15 +1045,13 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa """ Whether a tool that requires approval pauses the task instead of failing it. - Each excluded feature assumes the run finishes in one go: durable replay counts - steps across a single run, HITL review and code mode wrap the run, and a - sandbox the toolset provisions itself is destroyed when the run ends, so its - files would be gone on resume. A sandbox another task owns (``attach_to``) - outlives the pause, and the resumed run attaches to it again. + Each excluded feature assumes the run finishes in one go: HITL review and code + mode wrap the run, and a sandbox the toolset provisions itself is destroyed when + the run ends, so its files would be gone on resume. A sandbox another task owns + (``attach_to``) outlives the pause, and the resumed run attaches to it again. """ if ( not AIRFLOW_V_3_3_PLUS - or self.durable or self.enable_hitl_review or _contains_code_mode(self._declared_capabilities) ): @@ -891,9 +1063,9 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa Return ``output_type``, plus ``DeferredToolRequests`` when tool approval is supported or ``durable``. pydantic-ai drops ``DeferredToolRequests`` from the output schema the model sees; - it only lets the run end on a tool call awaiting approval. A durable run cannot pause, - but ending it there lets ``_pause_for_tool_approval`` fail the task with - ``UnsupportedToolDeferralError`` instead of pydantic-ai's ``UserError``. + it only lets the run end on a tool call awaiting approval. Under ``durable`` where + approval is not available, ending it there lets ``_pause_for_tool_approval`` fail the + task with ``UnsupportedToolDeferralError`` instead of pydantic-ai's ``UserError``. ``self.output_type`` is left alone because it is part of the serialized Dag. """ if not (self._supports_tool_approval() or self.durable): @@ -1028,8 +1200,19 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa 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, run_id=run_kwargs["run_id"]) + except BaseException as error: + if not _is_awaiting_input(error): + self._report_failed_run(context, run_usage, run_id=run_kwargs["run_id"]) + elif error.method_name == _DURABLE_APPROVAL_RESUME: + # The durable run paused for approval. Its usage so far, known only now that the run + # has settled it, goes with the pause, so the resume reports both sides of it. + error.kwargs = { + **(error.kwargs or {}), + "usage": _RUN_USAGE_ADAPTER.dump_python(run_usage, mode="json"), + "attempt_usage": _RUN_USAGE_ADAPTER.dump_python( + subtract_run_usage(run_usage, self._run_usage_base), mode="json" + ), + } raise def execute(self, context: Context) -> Any: @@ -1039,25 +1222,39 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa # A try that paused and then ended some other way (marked failed while waiting, a failed # request) leaves its transcript behind; a fresh run never reads it. ``.get``: a context # built by hand in a unit test has no task state store, and nothing to clean up. - if self._supports_tool_approval() and (store := context.get("task_state_store")) is not None: - self._delete_approval_transcript(store) + if ( + not self.durable + and self._supports_tool_approval() + and (store := context.get("task_state_store")) is not None + ): + self._delete_approval_state(store, _TOOL_APPROVAL_TRANSCRIPT_KEY) ti = context["task_instance"] 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) - self._run_usage = self._usage_budget.load() if self._usage_budget else RunUsage() + if self._usage_at_pause is None: + self._run_usage = self._usage_budget.load() if self._usage_budget else RunUsage() + self._run_usage_base = copy_run_usage(self._run_usage) + else: + # A durable run resumed after tool approval. The budget, when there is one, saved the + # same total when the run paused; this attempt's usage started before the pause. + usage_at_pause, attempt_usage_at_pause = self._usage_at_pause + self._run_usage = usage_at_pause + self._run_usage_base = subtract_run_usage(usage_at_pause, attempt_usage_at_pause) # A local, non-Optional alias of `self._run_usage` for the rest of this method -- # see `_run_agent_tracked`'s docstring for why callers pass this explicitly instead # of letting callees read `self._run_usage` (which mypy can't narrow past None). run_usage: RunUsage = self._run_usage - self._run_usage_base = copy_run_usage(run_usage) if self.durable: self._durable_journal = DurableJournal(self._build_durable_storage(context)) + self._approval_steps = 0 + if self._supports_tool_approval() and (store := context.get("task_state_store")) is not None: + self._load_tool_approval_answers(store) - agent = self._build_agent() + agent = self._build_agent(approval_context=context) self._run_identity_attrs = build_run_identity_attributes(ti) stamp_identity_on_agent_spans(agent, self._run_identity_attrs) @@ -1085,12 +1282,22 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa # A killed run raises RunCancelled (see run_agent_sync), which propagates to fail the # 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) + result, _ = 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 self._durable_journal is not None: self._durable_journal.log_summary(self.log) - return self._complete_run(context, result, attempt_usage=attempt_usage) + if (answer := self._new_tool_approval_answer) is not None and answer.calls: + # The replay went another way, such as after a change to the Dag, and never reached them. + self.log.warning( + "The resumed run finished without reaching the tool calls the answer was for (%s), " + "so they did not run.", + ", ".join(name for name, _ in answer.calls), + ) + # From where this attempt started, which for a resumed run is before the pause. + return self._complete_run( + context, result, attempt_usage=subtract_run_usage(run_usage, self._run_usage_base) + ) def _complete_run(self, context: Context, result: Any, *, attempt_usage: RunUsage) -> Any: """Finish a run, or pause it when the agent is waiting on a tool call to be approved.""" @@ -1136,6 +1343,8 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa # re-executing every already-completed model and tool step. if self._durable_journal is not None: self._durable_journal.cleanup() + if self._tool_approval_answers_stored and (store := context.get("task_state_store")) is not None: + self._delete_approval_state(store, _TOOL_APPROVAL_ANSWERS_KEY) self._clear_usage_budget(context) return output @@ -1162,29 +1371,18 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa old one; there a task instance asks at most once. """ requests: DeferredToolRequests = result.output - if requests.calls: - raise UnsupportedToolDeferralError( - "The agent called tools that need external execution " - f"({', '.join(call.tool_name for call in requests.calls)}); AgentOperator only " - "supports tools that need approval." - ) - pending_names = ", ".join(call.tool_name for call in requests.approvals) + _refuse_external_calls(requests) if not self._supports_tool_approval(): - # Reached under durable=True, or with DeferredToolRequests in a user-set output_type. + # Reached where approval is unsupported (a durable run's approval step is not added + # then), or with DeferredToolRequests in a user-set output_type. raise UnsupportedToolDeferralError( - f"The agent called tools that need approval ({pending_names}), but tool approval " - "needs Airflow 3.3+ and is not available with durable, enable_hitl_review, " - "a CodeMode capability or a SandboxToolset." + "The agent called tools that need approval " + f"({', '.join(call.tool_name for call in requests.approvals)}), but tool approval " + "needs Airflow 3.3+ and is not available with enable_hitl_review, a CodeMode " + "capability or a SandboxToolset." ) store = context["task_state_store"] - if not _REPEAT_APPROVAL_SUPPORTED and store.get(_TOOL_APPROVAL_REQUESTED_KEY): - raise ToolApprovalAlreadyRequestedError( - f"The agent asked for a second tool approval ({pending_names}), but this task " - "instance already asked once in this Dag run. Before Airflow 3.4, a repeat approval " - "request shows the reviewer the earlier request's details, so the task fails instead " - "of pausing. Upgrade to Airflow 3.4, have the agent ask for the gated calls in one " - "step, or give each irreversible action its own task." - ) + self._refuse_repeat_approval_request(context, requests.approvals) transcript = ModelMessagesTypeAdapter.dump_json(result.all_messages()).decode() retention = ( self.tool_approval_timeout + _TRANSCRIPT_RETENTION_MARGIN @@ -1204,8 +1402,24 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa }, ) + def _refuse_repeat_approval_request(self, context: Context, calls: Sequence[ToolCallPart]) -> None: + """Before Airflow 3.4, fail rather than ask a task instance's reviewer a second time.""" + if not _REPEAT_APPROVAL_SUPPORTED and context["task_state_store"].get(_TOOL_APPROVAL_REQUESTED_KEY): + raise ToolApprovalAlreadyRequestedError( + f"The agent asked for a second tool approval ({', '.join(call.tool_name for call in calls)}), " + "but this task instance already asked once in this Dag run. Before Airflow 3.4, a repeat " + "approval request shows the reviewer the earlier request's details, so the task fails " + "instead of pausing. Upgrade to Airflow 3.4, have the agent ask for the gated calls in " + "one step, or give each irreversible action its own task." + ) + def _ask_for_tool_approval( - self, context: Context, calls: Sequence[ToolCallPart], *, continuation: dict[str, Any] + self, + context: Context, + calls: Sequence[ToolCallPart], + *, + continuation: dict[str, Any], + resume_with: str = "resume_after_tool_approval", ) -> NoReturn: """Write the approval request for ``calls`` and park the task until it is answered.""" pending_names = ", ".join(call.tool_name for call in calls) @@ -1248,18 +1462,18 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa context["task_state_store"].set(_TOOL_APPROVAL_REQUESTED_KEY, True, retention=NEVER_EXPIRE) self.log.info("Waiting for approval of %s", pending_names) raise TaskAwaitingInput( - method_name="resume_after_tool_approval", + method_name=resume_with, kwargs={**continuation, "pause_number": pause_number}, timeout=self.tool_approval_timeout, ) - def _delete_approval_transcript(self, store: TaskStateStoreAccessor) -> None: - # Best-effort: the transcript holds tool results, but failing the task over its cleanup - # would be worse, and the row goes with the Dag run anyway. + def _delete_approval_state(self, store: TaskStateStoreAccessor, key: str) -> None: + # Best-effort: the transcript and the answers say what the agent did, but failing the task + # over their cleanup would be worse, and the rows go with the Dag run anyway. try: - store.delete(_TOOL_APPROVAL_TRANSCRIPT_KEY) + store.delete(key) except Exception: - self.log.warning("Could not delete the tool approval transcript", exc_info=True) + self.log.warning("Could not delete %s from the task state store", key, exc_info=True) def resume_after_tool_approval( self, @@ -1283,9 +1497,9 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa # The resumed run paused again: the transcript now belongs to that new pause. raise except BaseException: - self._delete_approval_transcript(store) + self._delete_approval_state(store, _TOOL_APPROVAL_TRANSCRIPT_KEY) raise - self._delete_approval_transcript(store) + self._delete_approval_state(store, _TOOL_APPROVAL_TRANSCRIPT_KEY) return output def _resume_after_tool_approval( @@ -1299,16 +1513,7 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa event: dict[str, Any], attempt_usage: dict[str, Any] | None, ) -> Any: - if "error" in event: - if event.get("error_type") == "timeout": - raise HITLTimeoutError(f"Tool approval timed out: {event['error']}") - raise HITLTriggerEventError(event) - if (current := self._toolset_ids()) != toolset_ids: - raise ToolApprovalError( - f"The agent's toolsets changed while it waited for approval (paused with {toolset_ids}, " - f"resumed with {current}), so the approved calls would reach a connection the reviewer " - "did not see." - ) + self._check_review_answer(event, toolset_ids) transcript = store.get(_TOOL_APPROVAL_TRANSCRIPT_KEY) if not isinstance(transcript, str) or ( hashlib.sha256(transcript.encode()).hexdigest() != transcript_sha256 @@ -1316,14 +1521,7 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa raise ToolApprovalError( "The transcript saved when the task paused for tool approval is missing or was modified." ) - if (answered := _answered_request(event)) is not None and answered != self._approval_pause_number: - # A review form left open on an earlier pause: its answer was about other calls. - self.log.warning( - "An answer for approval request %s arrived while request %s was waiting; " - "asking again instead of applying it.", - answered, - self._approval_pause_number, - ) + if not self._answers_waiting_request(event): calls = [ part for message in ModelMessagesTypeAdapter.validate_json(transcript) @@ -1343,27 +1541,7 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa }, ) - approval: bool | ToolDenied - if event.get("timedout"): - # on_tool_approval_timeout="deny": nobody refused, so do not tell the agent a person did. - approval = ToolDenied( - "No reviewer answered within the approval timeout, so this call was not run." - ) - self.log.info( - "Tool calls denied: nobody answered within tool_approval_timeout=%s", - self.tool_approval_timeout, - ) - elif LLMApprovalMixin.APPROVE in event["chosen_options"]: - approval = True - self.log.info("Tool calls approved by %s", LLMApprovalMixin._describe_responder(event)) - else: - reason = (event.get("params_input") or {}).get("reason") - # Only a typed reason reaches the agent; an untouched field can come back as None, "", - # or, from some UI versions, the whole parameter spec. - if not isinstance(reason, str) or not reason.strip(): - reason = "A reviewer denied this tool call." - approval = ToolDenied(reason) - self.log.info("Tool calls rejected by %s", LLMApprovalMixin._describe_responder(event)) + approval = self._approval_from_answer(event) usage_limits = coerce_usage_limits(self.usage_limits) ti = context["task_instance"] @@ -1409,6 +1587,184 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa context, result, attempt_usage=subtract_run_usage(run_usage, self._run_usage_base) ) + def _check_review_answer(self, event: dict[str, Any], toolset_ids: list[str]) -> None: + """Refuse to resume when the review ended without an answer, or the toolsets changed while it waited.""" + if "error" in event: + if event.get("error_type") == "timeout": + raise HITLTimeoutError(f"Tool approval timed out: {event['error']}") + raise HITLTriggerEventError(event) + if (current := self._toolset_ids()) != toolset_ids: + raise ToolApprovalError( + f"The agent's toolsets changed while it waited for approval (paused with {toolset_ids}, " + f"resumed with {current}), so the approved calls would reach a connection the reviewer " + "did not see." + ) + + def _answers_waiting_request(self, event: dict[str, Any]) -> bool: + """Whether the answer is for the approval request waiting now, not one a form left open showed.""" + if (answered := _answered_request(event)) is None or answered == self._approval_pause_number: + return True + # A review form left open on an earlier pause: its answer was about other calls. + self.log.warning( + "An answer for approval request %s arrived while request %s was waiting; " + "asking again instead of applying it.", + answered, + self._approval_pause_number, + ) + return False + + def _approval_from_answer(self, event: dict[str, Any]) -> bool | ToolDenied: + """Turn the reviewer's answer into the result each paused call gets: run it, or tell the agent why not.""" + if event.get("timedout"): + # on_tool_approval_timeout="deny": nobody refused, so do not tell the agent a person did. + self.log.info( + "Tool calls denied: nobody answered within tool_approval_timeout=%s", + self.tool_approval_timeout, + ) + return ToolDenied("No reviewer answered within the approval timeout, so this call was not run.") + if LLMApprovalMixin.APPROVE in event["chosen_options"]: + self.log.info("Tool calls approved by %s", LLMApprovalMixin._describe_responder(event)) + return True + reason = (event.get("params_input") or {}).get("reason") + # Only a typed reason reaches the agent; an untouched field can come back as None, "", + # or, from some UI versions, the whole parameter spec. + if not isinstance(reason, str) or not reason.strip(): + reason = "A reviewer denied this tool call." + self.log.info("Tool calls rejected by %s", LLMApprovalMixin._describe_responder(event)) + return ToolDenied(reason) + + def _load_tool_approval_answers(self, store: TaskStateStoreAccessor) -> None: + """Read the answers reviewers gave this run, leaving out ones given for other toolsets.""" + stored = _STORED_ANSWERS_ADAPTER.validate_python(store.get(_TOOL_APPROVAL_ANSWERS_KEY) or []) + self._tool_approval_answers_stored = bool(stored) + toolset_ids = self._toolset_ids() + answers = [_ToolApprovalAnswer.load(raw) for raw in stored] + if self._new_tool_approval_answer is not None and answers: + # The resume stored it last; keep it as the object the warnings look at. + answers[-1] = self._new_tool_approval_answer + # Like a resume, an answer never reaches a connection the reviewer did not see. + self._tool_approval_answers = [answer for answer in answers if answer.toolset_ids == toolset_ids] + if skipped := len(answers) - len(self._tool_approval_answers): + self.log.warning( + "Not applying %s approval answer(s) given while the agent had other toolsets.", + skipped, + ) + + def _next_approval_step(self, requests: DeferredToolRequests) -> int | None: + if self._asking_other_handlers: + return None + if (step := executing_step()) is not None: + # A tool's own work (an agent it runs, or calls it makes inline) is not run again once + # the tool is recorded, so an approval reached there would not be reached on replay, and + # every later approval step would be numbered one lower and get another step's answer. + if not requests.approvals: + _refuse_external_calls(requests) + raise UnsupportedToolDeferralError( + "The agent called tools that need approval " + f"({', '.join(call.tool_name for call in requests.approvals)}) from inside the " + f"{step.name!r} step, such as an agent run by a tool. With durable, a tool that needs " + "approval has to be called by the task's agent itself." + ) + self._approval_steps += 1 + return self._approval_steps + + async def _settle_tool_approvals( + self, context: Context, ctx: RunContext[Any], requests: DeferredToolRequests, step: int + ) -> DeferredToolResults: + """ + Settle every call a durable run is waiting on, as its approval step number ``step``. + + The answers reviewers gave this step go first, each to the calls it was given for, so no + handler can overrule one. The other capabilities' ``handle_deferred_tool_calls`` + handlers settle what they can of the rest, through the run's own chain, as they would + without durable. Whatever is left is asked about, and the task parks: the resume runs + the agent again, which replays every step before this one and reaches it with the answer. + """ + settled = DeferredToolResults() + for answer in self._tool_approval_answers: + if answer.step != step: + continue + pending = [call for call in requests.approvals if call.tool_call_id not in settled.approvals] + if taken := answer.take(pending): + self.log.info( + "Applying the %s answer given at approval step %s to %s", + "approve" if answer.approval is True else "reject", + step, + ", ".join(call.tool_name for call in pending if call.tool_call_id in taken), + ) + settled.approvals.update(taken) + remaining: DeferredToolRequests | None = ( + requests.remaining(settled) if settled.approvals else requests + ) + if remaining is not None and (root := ctx.root_capability) is not None: + self._asking_other_handlers = True + try: + result = await root.handle_deferred_tool_calls(ctx, requests=remaining) + finally: + self._asking_other_handlers = False + if result: + result = _only_for(remaining, result) + settled.update(result) + remaining = remaining.remaining(result) + if remaining is None: + return settled + _refuse_external_calls(remaining) + if (new := self._new_tool_approval_answer) is not None and new.calls: + self.log.warning( + "The resumed run reached tool calls no answer covers (%s); asking about them.", + ", ".join(call.tool_name for call in remaining.approvals), + ) + self._refuse_repeat_approval_request(context, remaining.approvals) + self._ask_for_tool_approval( + context, + remaining.approvals, + continuation={ + "step": step, + "calls": [[call.tool_name, _call_content(call)] for call in remaining.approvals], + "toolset_ids": self._toolset_ids(), + }, + resume_with=_DURABLE_APPROVAL_RESUME, + ) + + def resume_after_durable_tool_approval( + self, + context: Context, + step: int, + calls: list[list[str | None]], + toolset_ids: list[str], + usage: dict[str, Any], + attempt_usage: dict[str, Any], + event: dict[str, Any], + pause_number: int, + ) -> Any: + """ + Continue a durable run that paused for tool approval, with the reviewer's decision. + + The answer joins the run's earlier ones in the task state store, for approval step + ``step`` and the calls in ``calls`` (tool names and argument fingerprints). The task + then runs again from ``execute``, under the task's journal: every step before the pause + replays, and the approval step, which the pause left unrecorded, gets the answer. + """ + self._approval_pause_number = pause_number + self._check_review_answer(event, toolset_ids) + if self._answers_waiting_request(event): + answer = _ToolApprovalAnswer( + step=step, + calls=[(str(name), fingerprint) for name, fingerprint in calls], + toolset_ids=toolset_ids, + approval=self._approval_from_answer(event), + ) + store = context["task_state_store"] + stored = _STORED_ANSWERS_ADAPTER.validate_python(store.get(_TOOL_APPROVAL_ANSWERS_KEY) or []) + store.set(_TOOL_APPROVAL_ANSWERS_KEY, [*stored, answer.dump()], retention=NEVER_EXPIRE) + self._new_tool_approval_answer = answer + # Without an answer for these calls, the run reaches them and asks again. + self._usage_at_pause = ( + _RUN_USAGE_ADAPTER.validate_python(usage), + _RUN_USAGE_ADAPTER.validate_python(attempt_usage), + ) + return self.execute(context) + def _resolve_message_history(self) -> list[ModelMessage] | None: """ Deserialize :attr:`message_history` into a list of pydantic-ai messages. 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 c3143fb3b4b..a2076fd5c32 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 @@ -55,7 +55,7 @@ from pydantic_ai.messages import ( from pydantic_ai.models import ModelRequestParameters from pydantic_ai.models.function import FunctionModel from pydantic_ai.settings import ToolOrOutput -from pydantic_ai.tools import ToolDefinition +from pydantic_ai.tools import DeferredToolRequests, ToolDefinition from pydantic_core import to_jsonable_python from airflow.providers.common.ai.durable import fingerprint as fingerprint_module @@ -69,6 +69,7 @@ from airflow.providers.common.ai.durable.fingerprint import ( _seed_free, _template, _unpaired, + fingerprint_deferred_tool_requests, fingerprint_model_request, fingerprint_tool_call, ) @@ -2043,3 +2044,41 @@ class TestDigestsAcrossHashSeeds: assert first is not None assert first == second + + +class TestDeferredToolRequests: + """The calls an approval answer is for: it must not replay for any other call.""" + + @staticmethod + def _approvals(*calls: tuple[str, dict[str, Any], str]) -> DeferredToolRequests: + return DeferredToolRequests( + approvals=[ToolCallPart(name, args, tool_call_id=i) for name, args, i in calls] + ) + + def test_the_same_calls_hash_the_same(self): + calls = ("refund", {"order_id": 1}, "c-1"), ("refund", {"order_id": 2}, "c-2") + assert fingerprint_deferred_tool_requests( + self._approvals(*calls) + ) == fingerprint_deferred_tool_requests(self._approvals(*calls)) + + @pytest.mark.parametrize( + "changed", + [ + pytest.param(("refund", {"order_id": 2}, "c-1"), id="arguments"), + pytest.param(("refund", {"order_id": 1}, "c-9"), id="call-id"), + pytest.param(("cancel", {"order_id": 1}, "c-1"), id="tool"), + ], + ) + def test_a_different_call_hashes_differently(self, changed): + before = fingerprint_deferred_tool_requests(self._approvals(("refund", {"order_id": 1}, "c-1"))) + assert fingerprint_deferred_tool_requests(self._approvals(changed)) != before + + def test_an_approval_and_an_external_call_hash_differently(self): + call = ToolCallPart("refund", {"order_id": 1}, tool_call_id="c-1") + assert fingerprint_deferred_tool_requests( + DeferredToolRequests(approvals=[call]) + ) != fingerprint_deferred_tool_requests(DeferredToolRequests(calls=[call])) + + def test_a_call_that_cannot_be_fingerprinted_leaves_none(self): + with patch.object(fingerprint_module, "fingerprint_tool_call", autospec=True, return_value=None): + assert fingerprint_deferred_tool_requests(self._approvals(("refund", {}, "c-1"))) is None 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 index 4dc10253820..8648db8ea73 100644 --- a/providers/common/ai/tests/unit/common/ai/durable/test_journal.py +++ b/providers/common/ai/tests/unit/common/ai/durable/test_journal.py @@ -666,6 +666,82 @@ async def _failed_attempt_then_retry_that_diverges_at_once(memory_storage): class TestPause: + """A run that waits for input resumes by replaying what the attempt that paused recorded.""" + + @pytest.mark.asyncio + async def test_the_resume_replays_what_the_paused_attempt_recorded_after_an_earlier_failure( + self, memory_storage + ): + """The earlier failure was at a step the paused attempt ran live and so did not record.""" + first = DurableJournal(memory_storage).start_run() + await first.run("model", kind="model", fingerprint="m1", body=lambda: ok("look the order up")) + with pytest.raises(RuntimeError): + await first.run("lookup", kind="tool", fingerprint="l", body=boom) + first.fail(RuntimeError("tool failed")) + + retry = DurableJournal(memory_storage).start_run() + await retry.run("model", kind="model", fingerprint="m1", body=lambda: ok("live")) + await retry.run("lookup", kind="tool", fingerprint="l", body=lambda: ok("$10"), replayable=False) + await retry.run("model", kind="model", fingerprint="m2", body=lambda: ok("refund it")) + retry.claim("approve", kind="other", fingerprint="calls") + retry.pause() + + resume = DurableJournal(memory_storage).start_run() + await resume.run("model", kind="model", fingerprint="m1", body=lambda: ok("live")) + await resume.run("lookup", kind="tool", fingerprint="l", body=lambda: ok("$10"), replayable=False) + chosen = await resume.run("model", kind="model", fingerprint="m2", body=lambda: ok("asked again")) + answer = await resume.run("approve", kind="other", fingerprint="calls", body=lambda: ok("approved")) + + assert (chosen, answer) == ("refund it", "approved") + assert dict(resume.journal.stats.replayed) == {"model": 2} + + @pytest.mark.asyncio + async def test_an_entry_an_earlier_attempt_left_where_the_paused_one_recorded_nothing_runs_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")) + first.fail(RuntimeError("usage limit")) + + retry = DurableJournal(memory_storage).start_run() + await retry.run("a", kind="model", fingerprint="a2", body=lambda: ok("new a")) + memory_storage.refuse.add(build_step_key(0, 1)) + await retry.run("b", kind="tool", fingerprint="b", body=lambda: ok("new b")) + retry.claim("approve", kind="other", fingerprint="calls") + retry.pause() + memory_storage.refuse.clear() + + resume = DurableJournal(memory_storage).start_run() + a = await resume.run("a", kind="model", fingerprint="a2", body=lambda: ok("live a")) + b = await resume.run("b", kind="tool", fingerprint="b", body=lambda: ok("live b")) + + assert (a, b) == ("new a", "live b") + + @pytest.mark.asyncio + async def test_a_pause_short_of_an_earlier_one_leaves_none_of_its_later_steps_to_replay( + self, memory_storage + ): + first = DurableJournal(memory_storage).start_run() + for name in "abcd": + await first.run(name, kind="tool", fingerprint=name, body=lambda name=name: ok(f"old {name}")) + first.claim("approve", kind="other", fingerprint="calls") + first.pause() + + resumed = DurableJournal(memory_storage).start_run() + await resumed.run("a", kind="tool", fingerprint="a", body=lambda: ok("live a")) + await resumed.run("b", kind="tool", fingerprint="b2", body=lambda: ok("new b")) + resumed.claim("approve", kind="other", fingerprint="other calls") + resumed.pause() + + again = DurableJournal(memory_storage).start_run() + for name, fingerprint in (("a", "a"), ("b", "b2")): + await again.run(name, kind="tool", fingerprint=fingerprint, body=lambda: ok("live")) + await again.run("approve", kind="other", fingerprint="other calls", body=lambda: ok("approved")) + d = await again.run("d", kind="tool", fingerprint="d", body=lambda: ok("live d")) + + assert d == "live d" + @pytest.mark.asyncio async def test_entries_the_attempt_left_unrecorded_below_the_stored_boundary_do_not_replay( self, memory_storage 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 3eb80fd5a07..c95d96c9355 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 @@ -2237,7 +2237,10 @@ class TestAgentOperatorDurable: assert listed.id is not None assert orders.id is None - def test_durable_fails_a_tool_that_raises_approval_required_with_unsupported_tool_deferral_error(self): + @patch("airflow.providers.common.ai.operators.agent.AIRFLOW_V_3_3_PLUS", False) + def test_durable_before_airflow_3_3_fails_a_tool_that_raises_approval_required(self): + """Approval parks the task in awaiting_input, which Airflow before 3.3 does not have.""" + def refund() -> str: raise ApprovalRequired diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_agent_tool_approval.py b/providers/common/ai/tests/unit/common/ai/operators/test_agent_tool_approval.py index 51ed73095de..87b3a41a7c9 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_agent_tool_approval.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_agent_tool_approval.py @@ -18,24 +18,37 @@ from __future__ import annotations +import json from datetime import timedelta from typing import Any from unittest.mock import MagicMock, patch import pytest -from pydantic_ai import Agent, CallDeferred, DeferredToolRequests, Tool +from pydantic_ai import ( + Agent, + ApprovalRequired, + CallDeferred, + DeferredToolRequests, + DeferredToolResults, + RunContext, + Tool, +) +from pydantic_ai.capabilities import CapabilityOrdering, HandleDeferredToolCalls from pydantic_ai.exceptions import UsageLimitExceeded from pydantic_ai.messages import ModelResponse, TextPart, ToolCallPart from pydantic_ai.models.function import AgentInfo, FunctionModel from pydantic_ai.toolsets.function import FunctionToolset from pydantic_ai.usage import RunUsage +from airflow.providers.common.ai.decorators.agent import _AgentDecoratedOperator +from airflow.providers.common.ai.durable.journal import DurableRun from airflow.providers.common.ai.exceptions import ( ToolApprovalAlreadyRequestedError, ToolApprovalError, UnsupportedToolDeferralError, ) from airflow.providers.common.ai.operators.agent import ( + _TOOL_APPROVAL_ANSWERS_KEY, _TOOL_APPROVAL_REQUESTED_KEY, _TOOL_APPROVAL_TRANSCRIPT_KEY, AgentOperator, @@ -68,7 +81,8 @@ class _FakeTaskStateStore: return self.data.get(key, default) def set(self, key, value, *, retention=None): - self.data[key] = value + # Stored as JSON, as the real store does, so a test sees what a resume would read back. + self.data[key] = json.loads(json.dumps(value)) def delete(self, key): self.data.pop(key, None) @@ -80,7 +94,9 @@ class _Shop: def __init__(self): self.calls: list[tuple[str, int]] = [] - def toolset(self): + def toolset(self, *, live_lookup: bool = False): + """``live_lookup`` opts ``lookup`` out of the durable journal, so it runs on every attempt.""" + def lookup(order_id: int) -> str: self.calls.append(("lookup", order_id)) return f"order {order_id} costs $10" @@ -89,7 +105,8 @@ class _Shop: self.calls.append(("refund", order_id)) return f"refunded order {order_id}" - return FunctionToolset(tools=[lookup, refund]).approval_required( + tools = [Tool(lookup, metadata={"airflow": False}) if live_lookup else lookup, refund] + return FunctionToolset(tools=tools).approval_required( lambda ctx, tool_def, args: tool_def.name == "refund" ) @@ -310,8 +327,8 @@ class TestPause: assert upsert.call_args.kwargs["defaults"] == ["Reject"] def test_a_user_set_deferred_output_type_fails_where_approval_is_unsupported(self): - """Adding DeferredToolRequests by hand must not open a pause that durable, HITL review, - code mode or a sandbox cannot survive.""" + """Adding DeferredToolRequests by hand must not open a pause that HITL review, code mode + or a sandbox cannot survive.""" op = _operator( _lookup_and_refund, [_Shop().toolset()], @@ -320,7 +337,7 @@ class TestPause: ) with patch(UPSERT, autospec=True) as upsert: - with pytest.raises(UnsupportedToolDeferralError, match="not available with durable"): + with pytest.raises(UnsupportedToolDeferralError, match="not available with enable_hitl_review"): op.execute(_context(_FakeTaskStateStore())) upsert.assert_not_called() @@ -785,3 +802,801 @@ class TestWhenApprovalApplies: durable=True, agent_params={"toolsets": [SandboxToolset(_NoopBackend())]}, ) + + +class _Counted: + """A FunctionModel function that counts the requests reaching the model, which replays skip.""" + + def __init__(self, model_fn): + self.model_fn = model_fn + self.calls = 0 + + def __call__(self, messages, info: AgentInfo) -> ModelResponse: + self.calls += 1 + return self.model_fn(messages, info) + + +def _fails_once_after_refund(model_fn): + """``model_fn``, except that the request after the refund ran fails the first time.""" + failed = [] + + def model(messages, info: AgentInfo) -> ModelResponse: + if not failed and any("refunded" in result for result in _returns(messages)): + failed.append(True) + raise RuntimeError("model unavailable") + return model_fn(messages, info) + + return model + + +def _durable(model_fn, toolsets, **kwargs) -> AgentOperator: + return _operator(model_fn, toolsets, durable=True, **kwargs) + + +def _durable_pause(op: AgentOperator, ctx: Any) -> TaskAwaitingInput: + paused = _pause(op, ctx) + assert paused.method_name == "resume_after_durable_tool_approval" + return paused + + +class TestDurable: + """With ``durable=True`` the approval is a step of the journal, and the resume replays up to it.""" + + def test_pauses_without_saving_a_transcript_or_reporting_a_failed_run(self): + shop, store = _Shop(), _FakeTaskStateStore() + ctx = _context(store) + + with patch(UPSERT, autospec=True) as upsert: + with pytest.raises(TaskAwaitingInput) as exc: + _durable(_lookup_and_refund, [shop.toolset()]).execute(ctx) + + ctx["task_instance"].xcom_push.assert_not_called() + + assert exc.value.method_name == "resume_after_durable_tool_approval" + assert [name for name, _ in exc.value.kwargs["calls"]] == ["refund"] + assert exc.value.kwargs["pause_number"] == 1 + assert RunUsage(**exc.value.kwargs["usage"]).requests == 1 + assert shop.calls == [("lookup", 1)] + assert '"order_id": 1' in upsert.call_args.kwargs["body"] + assert _TOOL_APPROVAL_TRANSCRIPT_KEY not in store.data + + def test_approve_replays_the_steps_before_the_pause_and_runs_the_call_once(self): + shop, store = _Shop(), _FakeTaskStateStore() + model = _Counted(_lookup_and_refund) + paused = _durable_pause(_durable(model, [shop.toolset()]), _context(store)) + + output = _durable(model, [shop.toolset()]).resume_after_durable_tool_approval( + _context(store), **paused.kwargs, event=APPROVE + ) + + assert output == "order 1 costs $10 | refunded order 1" + assert shop.calls == [("lookup", 1), ("refund", 1)] + # One request before the pause, and one after it: the first is replayed on resume. + assert model.calls == 2 + + def test_reject_tells_the_agent_why_and_skips_the_call(self): + shop, store = _Shop(), _FakeTaskStateStore() + paused = _durable_pause(_durable(_lookup_and_refund, [shop.toolset()]), _context(store)) + + output = _durable(_lookup_and_refund, [shop.toolset()]).resume_after_durable_tool_approval( + _context(store), **paused.kwargs, event=_reject("over the refund limit") + ) + + assert shop.calls == [("lookup", 1)] + assert "over the refund limit" in output + + @pytest.mark.parametrize("repeat_supported", [False, True], ids=["airflow-3.3", "airflow-3.4"]) + def test_a_retry_after_a_failed_resume_neither_asks_again_nor_runs_the_call_twice(self, repeat_supported): + """On Airflow 3.3 asking again would also fail the task, as it asks once per task instance.""" + shop, store = _Shop(), _FakeTaskStateStore() + model = _fails_once_after_refund(_lookup_and_refund) + with patch(REPEAT_SUPPORTED, repeat_supported): + paused = _durable_pause(_durable(model, [shop.toolset()]), _context(store)) + with pytest.raises(RuntimeError, match="model unavailable"): + _durable(model, [shop.toolset()]).resume_after_durable_tool_approval( + _context(store), **paused.kwargs, event=APPROVE + ) + + with patch(UPSERT, autospec=True) as upsert: + output = _durable(model, [shop.toolset()]).execute(_context(store, try_number=2)) + + upsert.assert_not_called() + assert shop.calls == [("lookup", 1), ("refund", 1)] + assert output == "order 1 costs $10 | refunded order 1" + + def test_a_tool_that_raises_approval_required_itself_is_not_asked_about_again_on_retry(self): + refunds = [] + + def refund(ctx: RunContext[None], order_id: int) -> str: + if not ctx.tool_call_approved: + raise ApprovalRequired + refunds.append(order_id) + return f"refunded order {order_id}" + + def model(messages, info: AgentInfo) -> ModelResponse: + if _returns(messages): + return ModelResponse(parts=[TextPart(" | ".join(_returns(messages)))]) + return ModelResponse(parts=[ToolCallPart("refund", {"order_id": 1}, tool_call_id="c-refund")]) + + flaky = _fails_once_after_refund(model) + store = _FakeTaskStateStore() + paused = _durable_pause(_durable(flaky, [FunctionToolset(tools=[refund])]), _context(store)) + with pytest.raises(RuntimeError, match="model unavailable"): + _durable(flaky, [FunctionToolset(tools=[refund])]).resume_after_durable_tool_approval( + _context(store), **paused.kwargs, event=APPROVE + ) + + with patch(UPSERT, autospec=True) as upsert: + output = _durable(flaky, [FunctionToolset(tools=[refund])]).execute(_context(store, try_number=2)) + + upsert.assert_not_called() + assert refunds == [1] + assert output == "refunded order 1" + + @patch(REPEAT_SUPPORTED, True) + def test_a_later_step_pauses_again_and_the_last_resume_replays_both_approvals(self): + shop, store = _Shop(), _FakeTaskStateStore() + model = _Counted(_two_refunds) + first = _durable_pause(_durable(model, [shop.toolset()]), _context(store)) + with patch(UPSERT, autospec=True) as upsert: + with pytest.raises(TaskAwaitingInput) as second: + _durable(model, [shop.toolset()]).resume_after_durable_tool_approval( + _context(store), **first.kwargs, event=APPROVE + ) + + output = _durable(model, [shop.toolset()]).resume_after_durable_tool_approval( + _context(store), **second.value.kwargs, event=APPROVE + ) + + assert second.value.kwargs["pause_number"] == 2 + assert [name for name, _ in second.value.kwargs["calls"]] == ["refund"] + assert upsert.call_args.kwargs["subject"] == "Approve refund for task `t` (request 2)" + assert shop.calls == [("refund", 1), ("refund", 2)] + assert output == "refunded order 1 | refunded order 2" + assert model.calls == 3 + + @patch(REPEAT_SUPPORTED, False) + def test_before_airflow_3_4_a_second_approval_in_the_same_try_fails_closed(self): + shop, store = _Shop(), _FakeTaskStateStore() + first = _durable_pause(_durable(_two_refunds, [shop.toolset()]), _context(store)) + + with pytest.raises(ToolApprovalAlreadyRequestedError, match="already asked once"): + _durable(_two_refunds, [shop.toolset()]).resume_after_durable_tool_approval( + _context(store), **first.kwargs, event=APPROVE + ) + assert shop.calls == [("refund", 1)] + + @patch(REPEAT_SUPPORTED, True) + def test_an_answer_from_a_form_left_open_on_an_earlier_pause_asks_again(self): + shop, store = _Shop(), _FakeTaskStateStore() + first = _durable_pause(_durable(_two_refunds, [shop.toolset()]), _context(store)) + with patch(UPSERT, autospec=True): + with pytest.raises(TaskAwaitingInput) as second: + _durable(_two_refunds, [shop.toolset()]).resume_after_durable_tool_approval( + _context(store), **first.kwargs, event=APPROVE + ) + + with patch(UPSERT, autospec=True) as upsert: + with pytest.raises(TaskAwaitingInput) as asked_again: + _durable(_two_refunds, [shop.toolset()]).resume_after_durable_tool_approval( + _context(store), **second.value.kwargs, event={**APPROVE, "params_input": {"request": 1}} + ) + + assert shop.calls == [("refund", 1)] + assert asked_again.value.kwargs["pause_number"] == 3 + assert [name for name, _ in asked_again.value.kwargs["calls"]] == ["refund"] + assert '"order_id": 2' in upsert.call_args.kwargs["body"] + + def test_changed_toolsets_refuse_to_run_the_approved_call(self): + shop, store = _Shop(), _FakeTaskStateStore() + paused = _durable_pause(_durable(_lookup_and_refund, [shop.toolset()]), _context(store)) + + with pytest.raises(ToolApprovalError, match="toolsets changed"): + _durable(_lookup_and_refund, [shop.toolset()]).resume_after_durable_tool_approval( + _context(store), **{**paused.kwargs, "toolset_ids": ["sql-tenant_acme"]}, event=APPROVE + ) + assert shop.calls == [("lookup", 1)] + + def test_usage_output_covers_both_sides_of_the_pause(self): + shop, store = _Shop(), _FakeTaskStateStore() + paused = _durable_pause(_durable(_lookup_and_refund, [shop.toolset()]), _context(store)) + ctx = _context(store) + + _durable(_lookup_and_refund, [shop.toolset()]).resume_after_durable_tool_approval( + ctx, **paused.kwargs, event=APPROVE + ) + + pushed = {c.kwargs["key"]: c.kwargs["value"] for c in ctx["task_instance"].xcom_push.call_args_list} + assert pushed["usage"]["requests"] == 2 + + def test_usage_limits_span_the_pause(self): + shop, store = _Shop(), _FakeTaskStateStore() + limits = {"request_limit": 1} + paused = _durable_pause( + _durable(_lookup_and_refund, [shop.toolset()], usage_limits=limits), _context(store) + ) + + with pytest.raises(UsageLimitExceeded): + _durable( + _lookup_and_refund, [shop.toolset()], usage_limits=limits + ).resume_after_durable_tool_approval(_context(store), **paused.kwargs, event=APPROVE) + + def test_the_journal_run_is_paused_not_failed(self): + """A failure can hand the next try back to an earlier try's boundary; a pause must not.""" + with ( + patch.object(DurableRun, "pause", autospec=True, side_effect=DurableRun.pause) as pause, + patch.object(DurableRun, "fail", autospec=True, side_effect=DurableRun.fail) as fail, + ): + _durable_pause(_durable(_lookup_and_refund, [_Shop().toolset()]), _context(_FakeTaskStateStore())) + + pause.assert_called_once() + fail.assert_not_called() + + def test_external_execution_calls_are_refused(self): + def slow_job() -> str: + raise CallDeferred + + def model(messages, info): + return ModelResponse(parts=[ToolCallPart("slow_job", {})]) + + with patch(UPSERT, autospec=True) as upsert: + with pytest.raises(UnsupportedToolDeferralError, match="need external execution"): + _durable(model, [FunctionToolset(tools=[slow_job])]).execute(_context(_FakeTaskStateStore())) + upsert.assert_not_called() + + def test_a_retry_replays_the_answer_after_a_tool_opted_out_of_the_journal(self): + """The live call is claimed just before the answer, so the answer replays by its calls alone.""" + shop, store = _Shop(), _FakeTaskStateStore() + + model = _fails_once_after_refund(_lookup_and_refund) + paused = _durable_pause(_durable(model, [shop.toolset(live_lookup=True)]), _context(store)) + with pytest.raises(RuntimeError, match="model unavailable"): + _durable(model, [shop.toolset(live_lookup=True)]).resume_after_durable_tool_approval( + _context(store), **paused.kwargs, event=APPROVE + ) + + with patch(UPSERT, autospec=True) as upsert: + retry = _durable(model, [shop.toolset(live_lookup=True)]) + output = retry.execute(_context(store, try_number=2)) + + upsert.assert_not_called() + # The recorded decision replays rather than being made again from the stored answer. + assert retry._durable_journal is not None + assert retry._durable_journal.stats.replayed["other"] == 1 + assert shop.calls.count(("refund", 1)) == 1 + assert output == "order 1 costs $10 | refunded order 1" + + @patch(REPEAT_SUPPORTED, True) + def test_an_answer_is_not_applied_to_another_call_with_the_same_id(self): + """The Dag changed during the pause, so the resumed run asks the model again and gets another call.""" + shop, store = _Shop(), _FakeTaskStateStore() + paused = _durable_pause(_durable(_lookup_and_refund, [shop.toolset()]), _context(store)) + + def refunds_order_2(messages, info: AgentInfo) -> ModelResponse: + # Some providers number calls per response, so the id can repeat for another call. + return ModelResponse(parts=[ToolCallPart("refund", {"order_id": 2}, tool_call_id="c-refund")]) + + changed = _durable(refunds_order_2, [shop.toolset()]) + changed.prompt = "Refund order 2" + with patch(UPSERT, autospec=True) as upsert: + with pytest.raises(TaskAwaitingInput) as asked: + changed.resume_after_durable_tool_approval(_context(store), **paused.kwargs, event=APPROVE) + + assert ("refund", 2) not in shop.calls + assert asked.value.method_name == "resume_after_durable_tool_approval" + assert [name for name, _ in asked.value.kwargs["calls"]] == ["refund"] + assert '"order_id": 2' in upsert.call_args.kwargs["body"] + + def test_an_answer_applies_to_the_same_call_under_a_new_id(self): + """ + A tool opted out of the journal returns something new on resume, so the model is asked again. + + It asks for the same refund under a new call id, as a real model does, and the answer still applies. + """ + refunds, store = [], _FakeTaskStateStore() + asked = [] + + def stock_level() -> int: + return len(asked) + + def refund(order_id: int) -> str: + refunds.append(order_id) + return f"refunded order {order_id}" + + def model(messages, info: AgentInfo) -> ModelResponse: + done = _returns(messages) + asked.append(True) + if not done: + return ModelResponse( + parts=[ToolCallPart("stock_level", {}, tool_call_id=f"c-stock-{len(asked)}")] + ) + if len(done) == 1: + return ModelResponse( + parts=[ToolCallPart("refund", {"order_id": 1}, tool_call_id=f"c-refund-{len(asked)}")] + ) + return ModelResponse(parts=[TextPart(" | ".join(done))]) + + def toolsets(): + return [ + FunctionToolset( + tools=[ + Tool(stock_level, metadata={"airflow": False}), + Tool(refund, requires_approval=True), + ] + ) + ] + + paused = _durable_pause(_durable(model, toolsets()), _context(store)) + with patch(UPSERT, autospec=True) as upsert: + output = _durable(model, toolsets()).resume_after_durable_tool_approval( + _context(store), **paused.kwargs, event=APPROVE + ) + + upsert.assert_not_called() + assert refunds == [1] + assert output.endswith("refunded order 1") + + def test_a_handler_of_the_dag_authors_settles_only_the_calls_the_reviewer_did_not(self): + """The handler approves order 2 and leaves order 1 to a person, who rejects it.""" + refunds, store = [], _FakeTaskStateStore() + + def refund(order_id: int) -> str: + refunds.append(order_id) + return f"refunded order {order_id}" + + def model(messages, info: AgentInfo) -> ModelResponse: + if done := _returns(messages): + return ModelResponse(parts=[TextPart(" | ".join(done))]) + return ModelResponse( + parts=[ + ToolCallPart("refund", {"order_id": 1}, tool_call_id="c-1"), + ToolCallPart("refund", {"order_id": 2}, tool_call_id="c-2"), + ] + ) + + verdicts = [] + + seen: list[str] = [] + + def judge(ctx, requests: DeferredToolRequests): + # Approves order 2 the first time, then every call it saw then, given to it or not. + verdicts.append(True) + if len(verdicts) > 1: + return DeferredToolResults(approvals=dict.fromkeys(seen, True)) + seen.extend(c.tool_call_id for c in requests.approvals) + return requests.build_results( + approvals={ + c.tool_call_id: True for c in requests.approvals if c.args_as_dict()["order_id"] == 2 + } + ) + + def operator(): + return _durable( + model, + [FunctionToolset(tools=[Tool(refund, requires_approval=True)])], + capabilities=[HandleDeferredToolCalls(handler=judge)], + ) + + paused = _durable_pause(operator(), _context(store)) + assert [name for name, _ in paused.kwargs["calls"]] == ["refund"] + operator().resume_after_durable_tool_approval( + _context(store), **paused.kwargs, event=_reject("order 1 needs a ticket") + ) + + assert refunds == [2] + + @patch(REPEAT_SUPPORTED, True) + def test_a_judge_that_changes_its_mind_on_resume_still_lets_the_run_finish(self): + """At each ask the judge settles one of two calls, a different one each time it is asked.""" + refunds, store = [], _FakeTaskStateStore() + asked = [] + + def refund(order_id: int) -> str: + refunds.append(order_id) + return f"refunded order {order_id}" + + def model(messages, info: AgentInfo) -> ModelResponse: + if done := _returns(messages): + return ModelResponse(parts=[TextPart(" | ".join(sorted(done)))]) + return ModelResponse( + parts=[ + ToolCallPart("refund", {"order_id": 1}, tool_call_id="c-1"), + ToolCallPart("refund", {"order_id": 2}, tool_call_id="c-2"), + ] + ) + + def judge(ctx, requests: DeferredToolRequests): + asked.append(True) + chosen = 1 + len(asked) % 2 + return requests.build_results( + approvals={ + c.tool_call_id: True for c in requests.approvals if c.args_as_dict()["order_id"] == chosen + } + ) + + def operator(): + return _durable( + model, + [FunctionToolset(tools=[Tool(refund, requires_approval=True)])], + capabilities=[HandleDeferredToolCalls(handler=judge)], + ) + + paused, output, pauses = _durable_pause(operator(), _context(store)), None, 1 + while output is None and pauses < 5: + with patch(UPSERT, autospec=True): + try: + output = operator().resume_after_durable_tool_approval( + _context(store), **paused.kwargs, event=APPROVE + ) + except TaskAwaitingInput as again: + paused, pauses = again, pauses + 1 + + assert output == "refunded order 1 | refunded order 2" + assert sorted(refunds) == [1, 2] + + @patch(REPEAT_SUPPORTED, True) + def test_the_same_call_asked_for_again_later_is_asked_about_again(self): + """An answer covers the call the reviewer saw once, not every later call with the same arguments.""" + refunds, store = [], _FakeTaskStateStore() + + def refund(order_id: int) -> str: + refunds.append(order_id) + return f"refunded order {order_id}" + + def model(messages, info: AgentInfo) -> ModelResponse: + done = _returns(messages) + if len(done) < 2: + return ModelResponse( + parts=[ToolCallPart("refund", {"order_id": 1}, tool_call_id=f"c-{len(done)}")] + ) + return ModelResponse(parts=[TextPart("done")]) + + def operator(): + return _durable(model, [FunctionToolset(tools=[Tool(refund, requires_approval=True)])]) + + first = _durable_pause(operator(), _context(store)) + with patch(UPSERT, autospec=True) as upsert: + with pytest.raises(TaskAwaitingInput): + operator().resume_after_durable_tool_approval(_context(store), **first.kwargs, event=APPROVE) + + upsert.assert_called_once() + assert refunds == [1] + + def test_an_answer_survives_a_handler_that_fails_after_it_was_applied(self): + """The handler leaves order 1 to a person, then fails once on resume deciding order 2.""" + refunds, store = [], _FakeTaskStateStore() + verdicts = [] + + def refund(order_id: int) -> str: + refunds.append(order_id) + return f"refunded order {order_id}" + + def model(messages, info: AgentInfo) -> ModelResponse: + if done := _returns(messages): + return ModelResponse(parts=[TextPart(" | ".join(sorted(done)))]) + return ModelResponse( + parts=[ + ToolCallPart("refund", {"order_id": 1}, tool_call_id="c-1"), + ToolCallPart("refund", {"order_id": 2}, tool_call_id="c-2"), + ] + ) + + def flaky(ctx, requests: DeferredToolRequests): + verdicts.append(True) + if len(verdicts) == 2: + raise RuntimeError("judge unavailable") + return requests.build_results( + approvals={ + c.tool_call_id: True for c in requests.approvals if c.args_as_dict()["order_id"] == 2 + } + ) + + def operator(): + return _durable( + model, + [FunctionToolset(tools=[Tool(refund, requires_approval=True)])], + capabilities=[HandleDeferredToolCalls(handler=flaky)], + ) + + paused = _durable_pause(operator(), _context(store)) + with pytest.raises(RuntimeError, match="judge unavailable"): + operator().resume_after_durable_tool_approval(_context(store), **paused.kwargs, event=APPROVE) + with patch(UPSERT, autospec=True) as upsert: + output = operator().execute(_context(store, try_number=2)) + + upsert.assert_not_called() + assert output == "refunded order 1 | refunded order 2" + assert sorted(refunds) == [1, 2] + assert _TOOL_APPROVAL_ANSWERS_KEY not in store.data + + def test_a_handler_ordered_outermost_cannot_overrule_the_reviewer(self): + shop, store = _Shop(), _FakeTaskStateStore() + asked = [] + + class _OutermostJudge(HandleDeferredToolCalls): + def get_ordering(self): + return CapabilityOrdering(position="outermost") + + def judge(ctx, requests: DeferredToolRequests): + # Leaves the call to a person, and approves it whenever asked again. + asked.append(True) + return requests.build_results(approve_all=True) if len(asked) > 1 else None + + def operator(): + return _durable( + _lookup_and_refund, [shop.toolset()], capabilities=[_OutermostJudge(handler=judge)] + ) + + paused = _durable_pause(operator(), _context(store)) + operator().resume_after_durable_tool_approval(_context(store), **paused.kwargs, event=_reject()) + + assert ("refund", 1) not in shop.calls + + def test_a_task_agent_resumes_with_the_prompt_its_function_returns(self): + shop, store = _Shop(), _FakeTaskStateStore() + + def operator(): + op = _AgentDecoratedOperator( + task_id="t", + python_callable=lambda: "Refund order 1", + llm_conn_id="llm", + toolsets=[shop.toolset()], + durable=True, + ) + op.llm_hook = MagicMock(spec=["create_agent"]) + op.llm_hook.create_agent.side_effect = lambda **kw: Agent(FunctionModel(_lookup_and_refund), **kw) + return op + + paused = _durable_pause(operator(), _context(store)) + output = operator().resume_after_durable_tool_approval( + _context(store), **paused.kwargs, event=APPROVE + ) + + assert output == "order 1 costs $10 | refunded order 1" + assert shop.calls == [("lookup", 1), ("refund", 1)] + + def test_a_handler_built_per_run_is_asked_as_built(self): + """The capability the Dag declares leaves the call to a person; the one it builds per run approves it.""" + shop = _Shop() + + class _PerRunJudge(HandleDeferredToolCalls): + async def for_run(self, ctx): + return _PerRunJudge(handler=lambda ctx, requests: requests.build_results(approve_all=True)) + + with patch(UPSERT, autospec=True) as upsert: + output = _durable( + _lookup_and_refund, + [shop.toolset()], + capabilities=[_PerRunJudge(handler=lambda ctx, requests: None)], + ).execute(_context(_FakeTaskStateStore())) + + upsert.assert_not_called() + assert output == "order 1 costs $10 | refunded order 1" + + +def _refunds_order_1_twice(fail_once_after_first: bool = False): + """Refund order 1, then refund order 1 again in a later step; optionally fail once between them.""" + failed: list[bool] = [] + + def model(messages, info: AgentInfo) -> ModelResponse: + done = _returns(messages) + if len(done) == 1 and fail_once_after_first and not failed: + failed.append(True) + raise RuntimeError("model unavailable") + if len(done) < 2: + return ModelResponse( + parts=[ToolCallPart("refund", {"order_id": 1}, tool_call_id=f"c-{len(done)}")] + ) + return ModelResponse(parts=[TextPart(" | ".join(done))]) + + return model + + +class TestDurableAnswersStayWithTheirRequest: + """An answer settles the approval step that asked for it, and no later step with the same call.""" + + @pytest.mark.parametrize( + ("answers", "refunds"), + [ + pytest.param([APPROVE, _reject()], [1], id="approve-then-reject"), + pytest.param([_reject(), APPROVE], [1], id="reject-then-approve"), + ], + ) + @patch(REPEAT_SUPPORTED, True) + def test_two_requests_for_the_same_call_each_get_their_own_answer(self, answers, refunds): + shop, store = _Shop(), _FakeTaskStateStore() + model = _refunds_order_1_twice() + paused = _durable_pause(_durable(model, [shop.toolset()]), _context(store)) + with patch(UPSERT, autospec=True): + with pytest.raises(TaskAwaitingInput) as second: + _durable(model, [shop.toolset()]).resume_after_durable_tool_approval( + _context(store), **paused.kwargs, event=answers[0] + ) + _durable(model, [shop.toolset()]).resume_after_durable_tool_approval( + _context(store), **second.value.kwargs, event=answers[1] + ) + + assert [order for name, order in shop.calls if name == "refund"] == refunds + + @pytest.mark.parametrize("repeat_supported", [False, True], ids=["airflow-3.3", "airflow-3.4"]) + def test_a_retry_asks_about_a_later_call_like_one_already_approved(self, repeat_supported): + shop, store = _Shop(), _FakeTaskStateStore() + model = _refunds_order_1_twice(fail_once_after_first=True) + with patch(REPEAT_SUPPORTED, repeat_supported): + paused = _durable_pause(_durable(model, [shop.toolset()]), _context(store)) + with pytest.raises(RuntimeError, match="model unavailable"): + _durable(model, [shop.toolset()]).resume_after_durable_tool_approval( + _context(store), **paused.kwargs, event=APPROVE + ) + with patch(UPSERT, autospec=True): + with pytest.raises((TaskAwaitingInput, ToolApprovalAlreadyRequestedError)): + _durable(model, [shop.toolset()]).execute(_context(store, try_number=2)) + + assert [order for name, order in shop.calls if name == "refund"] == [1] + + def test_a_retry_does_not_apply_an_answer_given_for_another_connection(self): + shop, store = _Shop(), _FakeTaskStateStore() + model = _fails_once_after_refund(_lookup_and_refund) + + def operator(customer: str) -> AgentOperator: + op = _durable(model, [shop.toolset(), SQLToolset(db_conn_id="tenant_{{ params.customer }}")]) + op.render_template_fields({"params": {"customer": customer}}) + return op + + paused = _durable_pause(operator("acme"), _context(store)) + with pytest.raises(RuntimeError, match="model unavailable"): + operator("acme").resume_after_durable_tool_approval( + _context(store), **paused.kwargs, event=APPROVE + ) + with patch(REPEAT_SUPPORTED, True), patch(UPSERT, autospec=True) as upsert: + with pytest.raises(TaskAwaitingInput): + operator("globex").execute(_context(store, try_number=2)) + + upsert.assert_called_once() + assert shop.calls.count(("refund", 1)) == 1 + + def test_a_handler_the_model_has_not_loaded_is_not_asked(self): + shop = _Shop() + judge = HandleDeferredToolCalls( + handler=lambda ctx, requests: requests.build_results(approve_all=True), + id="judge", + defer_loading=True, + ) + + with patch(UPSERT, autospec=True) as upsert: + with pytest.raises(TaskAwaitingInput): + _durable(_lookup_and_refund, [shop.toolset()], capabilities=[judge]).execute( + _context(_FakeTaskStateStore()) + ) + + upsert.assert_called_once() + assert ("refund", 1) not in shop.calls + + def test_a_tool_that_needs_approval_called_from_inside_another_tool_fails_the_task(self): + """A tool runs the agent again, and that run asks for approval: a replay of the tool would skip it.""" + refunds, store = [], _FakeTaskStateStore() + + async def delegate(ctx: RunContext[None]) -> str: + assert ctx.agent is not None + result = await ctx.agent.run("CHILD") + return f"delegate got {result.output!r}" + + def refund(order_id: int) -> str: + refunds.append(order_id) + return f"refunded order {order_id}" + + def model(messages, info: AgentInfo) -> ModelResponse: + done = _returns(messages) + if "CHILD" in str(messages[0].parts[-1].content) and not done: + return ModelResponse(parts=[ToolCallPart("refund", {"order_id": 1}, tool_call_id="c-child")]) + if not done: + return ModelResponse(parts=[ToolCallPart("delegate", {}, tool_call_id="c-delegate")]) + return ModelResponse(parts=[TextPart(" | ".join(done))]) + + toolsets = [ + FunctionToolset(tools=[Tool(delegate, takes_ctx=True), Tool(refund, requires_approval=True)]) + ] + with patch(UPSERT, autospec=True) as upsert: + with pytest.raises(UnsupportedToolDeferralError, match="from inside the"): + _durable(model, toolsets).execute(_context(store)) + + upsert.assert_not_called() + assert refunds == [] + + @patch(REPEAT_SUPPORTED, True) + def test_one_answer_covers_one_call_when_the_resumed_step_has_it_twice(self): + """A live tool changes on resume, and the model asked again wants refund(1) twice in that step.""" + refunds, store, asked, refund_asks = [], _FakeTaskStateStore(), [], [] + + def stock_level() -> int: + return len(asked) + + def refund(order_id: int) -> str: + refunds.append(order_id) + return f"refunded order {order_id}" + + def model(messages, info: AgentInfo) -> ModelResponse: + done = _returns(messages) + asked.append(True) + if not done: + return ModelResponse(parts=[ToolCallPart("stock_level", {}, tool_call_id=f"s-{len(asked)}")]) + if len(done) == 1: + refund_asks.append(True) + return ModelResponse( + parts=[ + ToolCallPart("refund", {"order_id": 1}, tool_call_id=f"r-{len(asked)}-{i}") + for i in range(len(refund_asks)) + ] + ) + return ModelResponse(parts=[TextPart(" | ".join(done))]) + + def toolsets(): + return [ + FunctionToolset( + tools=[ + Tool(stock_level, metadata={"airflow": False}), + Tool(refund, requires_approval=True), + ] + ) + ] + + paused = _durable_pause(_durable(model, toolsets()), _context(store)) + with patch(UPSERT, autospec=True) as upsert: + with pytest.raises(TaskAwaitingInput) as again: + _durable(model, toolsets()).resume_after_durable_tool_approval( + _context(store), **paused.kwargs, event=APPROVE + ) + + assert len(refund_asks) == 2 + assert refunds == [] + upsert.assert_called_once() + assert len(again.value.kwargs["calls"]) == 1 + + @patch(REPEAT_SUPPORTED, True) + def test_a_step_a_handler_helped_settle_keeps_its_number_for_the_next_one(self): + """Step 1: a judge approves order 2 and the reviewer order 1. Step 2: the reviewer approves order 3.""" + refunds, store = [], _FakeTaskStateStore() + + def refund(order_id: int) -> str: + refunds.append(order_id) + return f"refunded order {order_id}" + + def model(messages, info: AgentInfo) -> ModelResponse: + done = _returns(messages) + if not done: + return ModelResponse( + parts=[ + ToolCallPart("refund", {"order_id": 1}, tool_call_id="c-1"), + ToolCallPart("refund", {"order_id": 2}, tool_call_id="c-2"), + ] + ) + if len(done) == 2: + return ModelResponse(parts=[ToolCallPart("refund", {"order_id": 3}, tool_call_id="c-3")]) + return ModelResponse(parts=[TextPart(" | ".join(sorted(done)))]) + + def judge(ctx, requests: DeferredToolRequests): + return requests.build_results( + approvals={ + c.tool_call_id: True for c in requests.approvals if c.args_as_dict()["order_id"] == 2 + } + ) + + def operator(): + return _durable( + model, + [FunctionToolset(tools=[Tool(refund, requires_approval=True)])], + capabilities=[HandleDeferredToolCalls(handler=judge)], + ) + + first = _durable_pause(operator(), _context(store)) + with patch(UPSERT, autospec=True): + with pytest.raises(TaskAwaitingInput) as second: + operator().resume_after_durable_tool_approval(_context(store), **first.kwargs, event=APPROVE) + output = operator().resume_after_durable_tool_approval( + _context(store), **second.value.kwargs, event=APPROVE + ) + + assert second.value.method_name == "resume_after_durable_tool_approval" + assert (first.kwargs["step"], second.value.kwargs["step"]) == (1, 2) + assert sorted(refunds) == [1, 2, 3] + assert output == "refunded order 1 | refunded order 2 | refunded order 3"
