This is an automated email from the ASF dual-hosted git repository. kaxil pushed a commit to branch commonai-tool-approval-repeat in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 474626c60ebdc67b61d7bc887bfa23e0140f6e87 Author: Kaxil Naik <[email protected]> AuthorDate: Wed Oct 14 02:49:13 2026 +0100 Let an agent ask for tool approval in more than one step On Airflow 3.4 and later, an AgentOperator run pauses for approval at every step that calls a marked tool, instead of failing the task on its second request. A real agent decides each gated call after it sees the previous result, so the once-per-Dag-run rule failed it on its second approval. The rule existed because Airflow 3.3 keeps the first request's subject and body on a repeat request. Airflow 3.4 replaces them and clears the earlier answer (#73235), so a second pause shows the reviewer only its own calls. On 3.3 the rule stays. Core still accepts an answer without knowing which request it was for, so each request now carries its number in a hidden form field and in its subject. An answer that names an earlier request, from a form left open on that request, is not applied: the task asks about the waiting calls again. Each resume gets its own run_id, since pydantic-ai refuses a run_id already in the history, and keeps the transcript when the resumed run pauses again. Per-tool approval moves to the stable table. --- providers/common/ai/docs/observability.rst | 4 +- providers/common/ai/docs/stability.rst | 22 +- providers/common/ai/docs/tool_approval.rst | 56 +++--- providers/common/ai/docs/toolsets/hook.rst | 4 +- providers/common/ai/docs/toolsets/index.rst | 4 +- providers/common/ai/docs/toolsets/mcp.rst | 4 +- providers/common/ai/docs/toolsets/sql.rst | 6 +- .../src/airflow/providers/common/ai/exceptions.py | 9 +- .../airflow/providers/common/ai/operators/agent.py | 159 +++++++++++---- .../ai/operators/test_agent_tool_approval.py | 221 +++++++++++++++++++-- 10 files changed, 390 insertions(+), 99 deletions(-) diff --git a/providers/common/ai/docs/observability.rst b/providers/common/ai/docs/observability.rst index e4a6fa02060..48f4c24a56e 100644 --- a/providers/common/ai/docs/observability.rst +++ b/providers/common/ai/docs/observability.rst @@ -79,8 +79,8 @@ How it works With ``enable_hitl_review`` the ``run_id`` and ``usage`` reflect the initial model run, not the human-feedback regenerations. A run that resumes after a tool approval (see :doc:`tool_approval`) continues as - ``<task-instance id>-resumed``, which is the ``run_id`` the operator pushes; - ``usage`` covers both parts. + ``<task-instance id>-resumed-<n>`` after pause ``n``, which is the ``run_id`` + the operator pushes; ``usage`` covers 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 identity keys without ``airflow.task_instance.id``. diff --git a/providers/common/ai/docs/stability.rst b/providers/common/ai/docs/stability.rst index 79ec4e8fe50..278377648fa 100644 --- a/providers/common/ai/docs/stability.rst +++ b/providers/common/ai/docs/stability.rst @@ -50,8 +50,21 @@ Pydantic AI toolsets, but the Pydantic AI class they inherit from can change. always dumped to a dict. ``decision_policy`` is experimental; see below. * - ``@task.agent`` and :class:`~airflow.providers.common.ai.operators.agent.AgentOperator` - Runs an agent with the model from ``llm_conn_id`` and the given toolsets, and - returns its output under the same rules as ``@task.llm``. ``durable``, - code mode and per-tool approval are experimental; see below. + returns its output under the same rules as ``@task.llm``. ``durable`` and + code mode are experimental; see below. + * - Per-tool approval on ``AgentOperator`` (:doc:`tool_approval`) + - On Airflow 3.3 and later, a call to a tool marked with Pydantic AI's approval API + (``approval_required(...)`` on a toolset, or ``requires_approval=True`` on a + function tool) pauses the task before the tool runs. **Approve** runs the call and + the agent carries on; **Reject** tells the agent the call was denied and it carries + on without it. ``tool_approval_timeout``, ``on_tool_approval_timeout`` (``"fail"`` + by default, or ``"deny"``) and ``tool_approval_assigned_users`` keep their meaning. + On Airflow 3.4 and later, every later step that calls a marked tool pauses again; + on Airflow 3.3 a task instance can ask once per Dag run. With ``durable=True``, + ``enable_hitl_review=True`` or a ``SandboxToolset`` that provisions its own sandbox, + and on Airflow versions before 3.3, a call to such a tool fails the task; under a + ``CodeMode`` capability, a call from ``run_code`` does not run and is reported back + to the model. * - ``message_history`` on ``AgentOperator`` - Seeds the run with the given conversation, as a list of messages or its JSON form, and publishes the finished conversation under the ``message_history`` XCom key so a @@ -182,11 +195,6 @@ Everything this provider ships that is not in the table above is experimental. * - ``durable=True`` on ``AgentOperator`` (:doc:`durable_execution`) - Which tool results are replayed on retry will change, so that a tool can declare whether its result may be reused. - * - Per-tool approval on ``AgentOperator`` (``tool_approval_timeout``, - ``on_tool_approval_timeout`` and ``tool_approval_assigned_users``; - :doc:`tool_approval`) - - A task instance can pause for approval only once per Dag run, so an agent that - needs a second approval later in the run fails. Lifting that limit is planned. * - Code mode (:doc:`code_mode`) and the ``shields`` extra (used in :doc:`capabilities`) - Thin integrations of ``pydantic-ai-harness`` and ``pydantic-ai-shields``, which are diff --git a/providers/common/ai/docs/tool_approval.rst b/providers/common/ai/docs/tool_approval.rst index 8c1c9ebdbfb..744e9f36b1d 100644 --- a/providers/common/ai/docs/tool_approval.rst +++ b/providers/common/ai/docs/tool_approval.rst @@ -20,11 +20,6 @@ Approve an agent's tool calls ============================= -.. note:: - - Experimental: this can change or be removed in a minor release of this provider. - See :ref:`howto/stability`. - .. seealso:: To approve, edit or reject an LLM operator's output instead, see :doc:`approval_gates`; to review an agent's final answer over several rounds, see :doc:`hitl_review`. @@ -86,20 +81,33 @@ approval are decided together and the others run straight away, once. ``usage_limits`` covers both sides of the pause, so a ``cost_limit`` is not reset by it. -One approval per task instance ------------------------------- - -A task instance asks for approval at most once per Dag run, and that includes its -retries and clears. Airflow keeps a single approval request per task instance, -and a second request would show the reviewer the first one's tool call while -asking about the new one. So when the agent asks again -- later in the same run, -or in a retry after the first request -- the task fails with -``ToolApprovalAlreadyRequestedError`` and does not retry. - -Have the agent request the gated calls in one step, or give each irreversible -action its own task. A task whose approved call ran and that then failed cannot -ask again in the same Dag run; trigger a new run, and keep gated tools -idempotent, since the approved call has already happened once. +An agent that needs approval more than once +------------------------------------------- + +An agent usually decides its next call after it sees the result of the last one, so +a refund agent may ask to refund one order, read the result, and then ask to refund +another. On Airflow 3.4 and later, each step that calls a marked tool pauses the task +again, and the reviewer sees only that step's calls. ``usage_limits`` counts across +every pause, so a run that pauses many times can reach ``request_limit`` (50 by +default) partway through, after earlier approved calls have run. + +Each request is numbered in its subject, and the review form sends the number back +with the answer in a hidden field. An answer from a form left open on an earlier +request, such as a second reviewer who never reloaded the page, is not applied to the +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. + +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 +instance asks again, so a second request would show the reviewer the earlier call +while asking about the new one. The task fails with +``ToolApprovalAlreadyRequestedError`` instead, and does not retry. On 3.3, have the +agent request the gated calls in one step, or give each irreversible action its own +task. Timeouts and who decides ------------------------ @@ -141,8 +149,8 @@ Requirements and limits 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 task resumes, whether the resumed run - succeeds or fails, and a later try deletes one left behind by a try that ended + 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 @@ -150,6 +158,6 @@ Requirements and limits call against a connection the reviewer did not see. The check compares toolset ids, which name the connection for ``SQLToolset``, ``MCPToolset`` and ``HookToolset``; it does not notice an edit to the connection itself. -- The resumed run gets its own pydantic-ai ``run_id``, ``<task-instance id>-resumed``. - The ``run_id`` XCom holds that id, and the ``usage`` XCom covers both sides of - the pause. +- 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. diff --git a/providers/common/ai/docs/toolsets/hook.rst b/providers/common/ai/docs/toolsets/hook.rst index 6ad8b8d2e26..b00e83c6493 100644 --- a/providers/common/ai/docs/toolsets/hook.rst +++ b/providers/common/ai/docs/toolsets/hook.rst @@ -190,8 +190,8 @@ reflection-based adapter, so the work is choosing the method list. :ref:`defense-layer table <toolset-defense-layers>` states this outright. Choose methods whose worst case you accept, not methods you intend to constrain later. To expose a method that changes something and have a person approve the call - first, wrap the toolset with ``.approval_required()``. A task instance can pause - for approval once per Dag run; see :doc:`../tool_approval`. + first, wrap the toolset with ``.approval_required()``. On Airflow 3.3 a task + instance can pause for approval only once per Dag run; see :doc:`../tool_approval`. - Its calls act as barriers. The tools are registered with ``sequential=True`` and each hook method runs in a worker thread, one blocking hook call at a time in the task process, so a slow call holds up every other tool the model emitted diff --git a/providers/common/ai/docs/toolsets/index.rst b/providers/common/ai/docs/toolsets/index.rst index 34f5d07fb8d..1e669affde4 100644 --- a/providers/common/ai/docs/toolsets/index.rst +++ b/providers/common/ai/docs/toolsets/index.rst @@ -69,8 +69,8 @@ with is resolved in the worker and is not part of any of that. Two controls work on every toolset. ``.approval_required()`` pauses the task before a matching call runs, until a person approves or rejects it on the **Required Actions** -page. It needs Airflow 3.3 or later, pauses a task instance at most once per Dag run, -and does not combine with ``durable=True`` or a ``SandboxToolset`` that provisions its +page. It needs Airflow 3.3 or later, pauses at every step that calls a matching tool on +Airflow 3.4 and later (once per Dag run on 3.3), and does not combine with ``durable=True`` or a ``SandboxToolset`` that provisions its own sandbox (:doc:`../tool_approval` lists every limit). ``.filtered()`` drops tools from what the model is offered, as the :ref:`MCP guide <howto/toolset:mcp-filtered>` shows. diff --git a/providers/common/ai/docs/toolsets/mcp.rst b/providers/common/ai/docs/toolsets/mcp.rst index 64b69d3723f..65ca204f769 100644 --- a/providers/common/ai/docs/toolsets/mcp.rst +++ b/providers/common/ai/docs/toolsets/mcp.rst @@ -196,8 +196,8 @@ agent is offered. It does not change what the server lets the connection's crede do, so give the connection the narrowest credential the server accepts. To keep a tool available but have a person approve a call to it before it runs, wrap -the toolset with ``.approval_required()`` instead. A task instance can pause for -approval once per Dag run; see :doc:`../tool_approval`. +the toolset with ``.approval_required()`` instead. On Airflow 3.3 a task instance can +pause for approval only once per Dag run; see :doc:`../tool_approval`. Direct pydantic-ai MCP toolsets ------------------------------- diff --git a/providers/common/ai/docs/toolsets/sql.rst b/providers/common/ai/docs/toolsets/sql.rst index df0a4caca78..1e44b4b87ad 100644 --- a/providers/common/ai/docs/toolsets/sql.rst +++ b/providers/common/ai/docs/toolsets/sql.rst @@ -318,9 +318,9 @@ Parameters (``DESCRIBE``/``SHOW``) statements are permitted. To have a person approve a ``query`` call before it runs, wrap the toolset with ``.approval_required()`` and return ``tool_def.name == "query"`` from its function. Reads and writes both go - through ``query``, and a task instance can pause for approval once per Dag run, so - this fits an agent that runs one query, such as a single write; see - :doc:`../tool_approval`. + through ``query``, so every query pauses for a decision. On Airflow 3.3 a task + instance can pause only once per Dag run, so there this fits an agent that runs one + query, such as a single write; see :doc:`../tool_approval`. - ``max_rows``: Maximum rows returned from the ``query`` tool. Default ``50``. Rows beyond it are not read out of a DBAPI cursor; what the driver has already transferred is its own call. See :ref:`bounded-query-results`. diff --git a/providers/common/ai/src/airflow/providers/common/ai/exceptions.py b/providers/common/ai/src/airflow/providers/common/ai/exceptions.py index 2f9c61f9f4b..6c538744417 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/exceptions.py +++ b/providers/common/ai/src/airflow/providers/common/ai/exceptions.py @@ -34,11 +34,12 @@ class ToolApprovalError(RuntimeError): class ToolApprovalAlreadyRequestedError(AirflowFailException): """ - Raised when a task instance asks for a second tool approval. + Raised on Airflow 3.3 when a task instance asks for a second tool approval. - Airflow keeps one approval request per task instance, across retries and clears, and a - second request would show the reviewer the first one's details. The task fails without - retrying, since a retry would ask again. + Airflow 3.3 keeps one approval request per task instance, across retries and clears, and + a second request would show the reviewer the first one's details. The task fails without + retrying, since a retry would ask again. Airflow 3.4 replaces a repeat request's details, + so there the run pauses again instead. """ diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py index 52038818025..afd0e9ea78a 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 @@ -32,7 +32,7 @@ from typing import TYPE_CHECKING, Any, ClassVar, Literal, NoReturn from pydantic import BaseModel, TypeAdapter from pydantic_ai import DeferredToolRequests, DeferredToolResults, ToolDenied from pydantic_ai.capabilities import AbstractCapability, Toolset, WrapperCapability -from pydantic_ai.messages import ModelMessagesTypeAdapter +from pydantic_ai.messages import ModelMessagesTypeAdapter, ModelResponse, ToolCallPart from pydantic_ai.toolsets.abstract import AbstractToolset from pydantic_ai.usage import RunUsage @@ -75,7 +75,11 @@ from airflow.providers.common.compat.sdk import ( conf, redact, ) -from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_1_PLUS, AIRFLOW_V_3_3_PLUS +from airflow.providers.common.compat.version_compat import ( + AIRFLOW_V_3_1_PLUS, + AIRFLOW_V_3_3_PLUS, + AIRFLOW_V_3_4_PLUS, +) from airflow.providers.standard.exceptions import HITLTimeoutError, HITLTriggerEventError if AIRFLOW_V_3_3_PLUS: @@ -111,11 +115,18 @@ if TYPE_CHECKING: from airflow.sdk.execution_time.context import TaskStateStoreAccessor from airflow.sdk.execution_time.hitl import HITLUser -# Task state store keys: the transcript of a run paused for tool approval, and a marker that -# this task instance has asked once. The store is keyed by Dag run, task and map index, so the -# marker survives retries and clears, as the task instance's single approval request does. +# Task state store keys: the transcript of a run paused for tool approval, and, on Airflow 3.3, +# a marker that this task instance has asked once. The store is keyed by Dag run, task and map +# 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" +# 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 +# A hidden form field holding the pause number. The review form sends back the values it loaded, +# so a form left open on an earlier pause answers with that pause's number, which the resume +# refuses to apply to the calls that are waiting now. +_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) @@ -204,6 +215,22 @@ def _contains_code_mode(capabilities: Iterable[AgentCapability[Any]]) -> bool: return code_mode_cls is not None and _contains_capability(capabilities, code_mode_cls) +def _answered_request(event: dict[str, Any]) -> int | None: + """Return the approval request a review answer was for, or ``None`` when it does not say.""" + answered = (event.get("params_input") or {}).get(_APPROVAL_REQUEST_PARAM) + if isinstance(answered, dict): + # Some UI versions send back the whole parameter spec instead of its value. + answered = answered.get("value") + if answered is None or answered == "": + # An API caller that sent only the chosen option. + return None + try: + return int(answered) + except (TypeError, ValueError): + # Unreadable, so it matches no request and the calls are asked about again. + return -1 + + 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( @@ -403,7 +430,7 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa :param hitl_poll_interval: Seconds between XCom polls while waiting for a human response. Default ``10``. - **Per-tool approval** (Airflow 3.3+, experimental): + **Per-tool approval** (Airflow 3.3+): Mark the tools a human must approve with pydantic-ai's own API -- ``toolset.approval_required(...)``, or ``requires_approval=True`` on a function @@ -411,9 +438,10 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa arguments, appear on the **Required Actions** page; the task waits in the ``awaiting_input`` state without holding a worker slot. On **Approve** the calls run and the agent carries on. On **Reject** the agent is told the call was denied - (with the reviewer's reason, when given) and carries on without it. A task - instance asks at most once per Dag run, across retries and clears; a second - request fails the task. ``usage_limits`` applies to both sides of the pause. + (with the reviewer's reason, when given) and carries on without it. On Airflow + 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`` capability, or a ``SandboxToolset`` that provisions its own sandbox; there, a tool that requires approval fails @@ -421,13 +449,13 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa which does not run and is reported back to the model. A ``SandboxToolset`` attached to a sandbox another task owns is fine: the sandbox outlives the pause. - :param tool_approval_timeout: Experimental. How long the pause waits for a decision. + :param tool_approval_timeout: How long the pause waits for a decision. ``None`` (default) waits indefinitely. Must be positive. - :param on_tool_approval_timeout: Experimental. What a timed-out pause does: ``"fail"`` + :param on_tool_approval_timeout: What a timed-out pause does: ``"fail"`` (default) fails the task, ``"deny"`` rejects the pending calls so the agent carries on without them, and needs a ``tool_approval_timeout``. There is no approve-on-timeout. - :param tool_approval_assigned_users: Experimental. Users allowed to decide. ``None`` (default) + :param tool_approval_assigned_users: Users allowed to decide. ``None`` (default) leaves it to anyone who can act on the task's Required Actions. :param serialize_output: If ``True`` and ``output_type`` is a Pydantic @@ -522,6 +550,8 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa self._usage_budget: TaskStateStoreUsageBudget | None = None self._run_usage: RunUsage | None = None 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 # The answer the HITL reviewer is looking at, as the agent returned it. self._output_under_review: Any = None @@ -1162,9 +1192,11 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa the task waited, and the rendered toolset ids, so it refuses to run approved calls against a connection id the reviewer did not see. - A task instance asks at most once. Airflow keeps one approval request per task - instance, across retries and clears, and a repeat request keeps the first one's - subject and body, so the reviewer would approve a new call while reading the old one. + Airflow keeps one approval request per task instance, across retries and clears. From + Airflow 3.4 a repeat request replaces its details and clears the earlier answer, so the + run can pause again in a later step or a later try. Before 3.4 a repeat request kept the + first one's subject and body, so the reviewer would approve a new call while reading the + old one; there a task instance asks at most once. """ requests: DeferredToolRequests = result.output if requests.calls: @@ -1182,13 +1214,13 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa "a CodeMode capability or a SandboxToolset." ) store = context["task_state_store"] - if store.get(_TOOL_APPROVAL_REQUESTED_KEY): + 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. Airflow keeps one approval request per " - "task instance and would show the reviewer the earlier request's details, so the task " - "fails instead of pausing. Have the agent ask for the gated calls in one step, give " - "each irreversible action its own task, or trigger a new Dag run." + "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." ) transcript = ModelMessagesTypeAdapter.dump_json(result.all_messages()).decode() retention = ( @@ -1197,17 +1229,35 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa else NEVER_EXPIRE ) store.set(_TOOL_APPROVAL_TRANSCRIPT_KEY, transcript, retention=retention) + self._ask_for_tool_approval( + context, + requests.approvals, + continuation={ + "tool_call_ids": [call.tool_call_id for call in requests.approvals], + "usage": _RUN_USAGE_ADAPTER.dump_python(result.usage, mode="json"), + "attempt_usage": _RUN_USAGE_ADAPTER.dump_python(attempt_usage, mode="json"), + "transcript_sha256": hashlib.sha256(transcript.encode()).hexdigest(), + "toolset_ids": self._toolset_ids(), + }, + ) + def _ask_for_tool_approval( + self, context: Context, calls: Sequence[ToolCallPart], *, continuation: dict[str, Any] + ) -> 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) # Tool arguments can carry credentials (an HTTP header, an MCP token); mask them first. pending = "\n\n".join( f"**{call.tool_name}**\n\n```json\n" f"{json.dumps(redact(call.args_as_dict()), indent=2, default=str)}\n```" - for call in requests.approvals + for call in calls ) + pause_number = self._approval_pause_number + 1 upsert_hitl_detail( ti_id=context["task_instance"].id, options=[LLMApprovalMixin.APPROVE, LLMApprovalMixin.REJECT], - subject=f"Approve tool call for task `{self.task_id}`", + # Numbered, so each pause of a run looks different on the Required Actions page. + subject=f"Approve {pending_names} for task `{self.task_id}` (request {pause_number})", body=f"The agent wants to run:\n\n{pending}", defaults=[LLMApprovalMixin.REJECT] if self.on_tool_approval_timeout == "deny" else None, multiple=False, @@ -1221,21 +1271,22 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa "description": "Sent to the agent when you reject the call (optional).", "schema": {"type": ["string", "null"]}, }, + # A "const" schema renders as a hidden field the reviewer cannot change. + _APPROVAL_REQUEST_PARAM: { + "value": pause_number, + "description": "Which approval request of this run the answer is for.", + "schema": {"type": "integer", "const": pause_number}, + }, }, assigned_users=self.tool_approval_assigned_users, ) - # Only once the request exists: a failed request must not block the retry from asking. - store.set(_TOOL_APPROVAL_REQUESTED_KEY, True, retention=NEVER_EXPIRE) + if not _REPEAT_APPROVAL_SUPPORTED: + # Only once the request exists: a failed request must not block the retry from asking. + 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", - kwargs={ - "tool_call_ids": [call.tool_call_id for call in requests.approvals], - "usage": _RUN_USAGE_ADAPTER.dump_python(result.usage, mode="json"), - "attempt_usage": _RUN_USAGE_ADAPTER.dump_python(attempt_usage, mode="json"), - "transcript_sha256": hashlib.sha256(transcript.encode()).hexdigest(), - "toolset_ids": self._toolset_ids(), - }, + kwargs={**continuation, "pause_number": pause_number}, timeout=self.tool_approval_timeout, ) @@ -1256,16 +1307,23 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa toolset_ids: list[str], event: dict[str, Any], attempt_usage: dict[str, Any] | None = None, + pause_number: int = 1, ) -> Any: """Continue a run paused by :meth:`_pause_for_tool_approval` with the reviewer's decision.""" store = context["task_state_store"] + self._approval_pause_number = pause_number try: - return self._resume_after_tool_approval( + output = self._resume_after_tool_approval( context, store, tool_call_ids, usage, transcript_sha256, toolset_ids, event, attempt_usage ) - finally: - # Whatever the outcome. A second pause fails closed before writing a new transcript. + except TaskAwaitingInput: + # The resumed run paused again: the transcript now belongs to that new pause. + raise + except BaseException: self._delete_approval_transcript(store) + raise + self._delete_approval_transcript(store) + return output def _resume_after_tool_approval( self, @@ -1295,6 +1353,32 @@ 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, + ) + calls = [ + part + for message in ModelMessagesTypeAdapter.validate_json(transcript) + if isinstance(message, ModelResponse) + for part in message.parts + if isinstance(part, ToolCallPart) and part.tool_call_id in tool_call_ids + ] + self._ask_for_tool_approval( + context, + calls, + continuation={ + "tool_call_ids": tool_call_ids, + "usage": usage, + "attempt_usage": attempt_usage, + "transcript_sha256": transcript_sha256, + "toolset_ids": toolset_ids, + }, + ) approval: bool | ToolDenied if event.get("timedout"): @@ -1350,9 +1434,10 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa approvals={tool_call_id: approval for tool_call_id in tool_call_ids} ), usage_limits=usage_limits, - # pydantic-ai refuses a run_id already in the history; the task-instance id stays - # the prefix, so the resumed run still joins back to the task. - run_id=f"{ti.id}-resumed", + # pydantic-ai refuses a run_id already in the history, so each resume of the try + # gets its own; the task-instance id stays the prefix, so the resumed run still + # joins back to the task. + run_id=f"{ti.id}-resumed-{self._approval_pause_number}", ) except BaseException: self._report_failed_run(context, run_usage) 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 7334b98fc2e..11e4cbb8850 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 @@ -46,7 +46,7 @@ from airflow.providers.common.ai.toolsets.sql import SQLToolset from airflow.providers.common.ai.utils.usage_budget import USAGE_BUDGET_KEY, dump_run_usage from airflow.providers.standard.exceptions import HITLTimeoutError -from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS +from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS, AIRFLOW_V_3_4_PLUS pytestmark = pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="Per-tool approval needs Airflow 3.3+.") @@ -54,6 +54,7 @@ if AIRFLOW_V_3_3_PLUS: from airflow.sdk.exceptions import TaskAwaitingInput UPSERT = "airflow.providers.common.ai.operators.agent.upsert_hitl_detail" +REPEAT_SUPPORTED = "airflow.providers.common.ai.operators.agent._REPEAT_APPROVAL_SUPPORTED" class _FakeTaskStateStore: @@ -108,13 +109,20 @@ def _lookup_and_refund(messages, info: AgentInfo) -> ModelResponse: return ModelResponse(parts=[TextPart(" | ".join(_returns(messages)))]) -def _two_refunds(messages, info: AgentInfo) -> ModelResponse: - """Refund order 1, then order 2, one call per step: two separate approvals.""" - done = _returns(messages) - if len(done) < 2: - n = len(done) + 1 - return ModelResponse(parts=[ToolCallPart("refund", {"order_id": n}, tool_call_id=f"c-{n}")]) - return ModelResponse(parts=[TextPart(" | ".join(done))]) +def _refunds_one_per_step(count: int): + """Refund orders 1..count, one call per step, so each refund needs its own approval.""" + + def model(messages, info: AgentInfo) -> ModelResponse: + done = _returns(messages) + if len(done) < count: + n = len(done) + 1 + return ModelResponse(parts=[ToolCallPart("refund", {"order_id": n}, tool_call_id=f"c-{n}")]) + return ModelResponse(parts=[TextPart(" | ".join(done))]) + + return model + + +_two_refunds = _refunds_one_per_step(2) def _operator(model_fn, toolsets, **kwargs) -> AgentOperator: @@ -175,13 +183,14 @@ class TestPause: assert "lookup" not in body assert upsert.call_args.kwargs["options"] == ["Approve", "Reject"] assert upsert.call_args.kwargs["defaults"] is None - assert upsert.call_args.kwargs["subject"] == "Approve tool call for task `t`" + assert upsert.call_args.kwargs["subject"] == "Approve refund for task `t` (request 1)" assert upsert.call_args.kwargs["multiple"] is False # Optional in the review form: without "null" the UI requires a reason even to approve. assert upsert.call_args.kwargs["params"]["reason"]["schema"] == {"type": ["string", "null"]} + # Hidden in the review form ("const"), and sent back with the answer. + assert upsert.call_args.kwargs["params"]["request"]["schema"] == {"type": "integer", "const": 1} assert upsert.call_args.kwargs["assigned_users"] == [{"id": "u1", "name": "alice"}] assert exc.value.timeout == timedelta(hours=1) - assert store.data[_TOOL_APPROVAL_REQUESTED_KEY] is True def test_a_function_tool_marked_requires_approval_pauses_too(self): refunds = [] @@ -218,11 +227,13 @@ class TestPause: assert "s3cr3t" not in upsert.call_args.kwargs["body"] - def test_a_later_try_fails_closed_once_an_approval_was_requested(self): - """Core keeps one approval request per task instance across retries and clears, with the - first request's subject and body, so a retry's request would show the earlier call.""" + @patch(REPEAT_SUPPORTED, False) + def test_before_airflow_3_4_a_later_try_fails_closed_once_an_approval_was_requested(self): + """Before 3.4, core keeps one approval request per task instance across retries and clears, + with the first request's subject and body, so a retry's request would show the earlier call.""" shop, store = _Shop(), _FakeTaskStateStore() _pause(_operator(_lookup_and_refund, [shop.toolset()]), _context(store)) + assert store.data[_TOOL_APPROVAL_REQUESTED_KEY] is True with patch(UPSERT, autospec=True) as upsert: with pytest.raises(ToolApprovalAlreadyRequestedError, match="already asked once"): @@ -230,6 +241,43 @@ class TestPause: upsert.assert_not_called() + @patch(REPEAT_SUPPORTED, True) + def test_from_airflow_3_4_no_once_per_run_marker_is_written(self): + store = _FakeTaskStateStore() + _pause(_operator(_lookup_and_refund, [_Shop().toolset()]), _context(store)) + assert _TOOL_APPROVAL_REQUESTED_KEY not in store.data + + @pytest.mark.skipif(AIRFLOW_V_3_4_PLUS, reason="The once-per-run rule applies before Airflow 3.4.") + def test_on_the_installed_airflow_before_3_4_a_second_request_fails(self): + """No patch: checks the real version gate on the Airflow this suite runs against.""" + shop, store = _Shop(), _FakeTaskStateStore() + _pause(_operator(_lookup_and_refund, [shop.toolset()]), _context(store)) + with patch(UPSERT, autospec=True): + with pytest.raises(ToolApprovalAlreadyRequestedError): + _operator(_lookup_and_refund, [shop.toolset()]).execute(_context(store, try_number=2)) + + @pytest.mark.skipif(not AIRFLOW_V_3_4_PLUS, reason="Repeat requests need Airflow 3.4.") + def test_on_the_installed_airflow_3_4_a_second_request_pauses(self): + """No patch: checks the real version gate on the Airflow this suite runs against.""" + shop, store = _Shop(), _FakeTaskStateStore() + _pause(_operator(_lookup_and_refund, [shop.toolset()]), _context(store)) + with patch(UPSERT, autospec=True): + with pytest.raises(TaskAwaitingInput): + _operator(_lookup_and_refund, [shop.toolset()]).execute(_context(store, try_number=2)) + + @patch(REPEAT_SUPPORTED, True) + def test_a_later_try_asks_again(self): + shop, store = _Shop(), _FakeTaskStateStore() + _pause(_operator(_lookup_and_refund, [shop.toolset()]), _context(store)) + + with patch(UPSERT, autospec=True) as upsert: + with pytest.raises(TaskAwaitingInput): + _operator(_lookup_and_refund, [shop.toolset()]).execute(_context(store, try_number=2)) + + upsert.assert_called_once() + assert '"order_id": 1' in upsert.call_args.kwargs["body"] + + @patch(REPEAT_SUPPORTED, False) def test_a_failed_request_does_not_block_the_retry(self): shop, store = _Shop(), _FakeTaskStateStore() with patch(UPSERT, autospec=True, side_effect=RuntimeError("api down")): @@ -382,9 +430,10 @@ class TestResume: assert "No reviewer answered within the approval timeout" in output assert shop.calls == [("lookup", 1)] - def test_a_second_approval_in_the_same_try_fails_closed(self): - """Airflow keeps one approval request per task instance and a second one would show the - first one's details, so the reviewer would approve refund 2 while reading refund 1.""" + @patch(REPEAT_SUPPORTED, False) + def test_before_airflow_3_4_a_second_approval_in_the_same_try_fails_closed(self): + """Before 3.4, a second request would show the first one's details, so the reviewer would + approve refund 2 while reading refund 1.""" shop, store = _Shop(), _FakeTaskStateStore() first = _pause(_operator(_two_refunds, [shop.toolset()]), _context(store)) @@ -398,6 +447,146 @@ class TestResume: upsert.assert_not_called() assert _TOOL_APPROVAL_TRANSCRIPT_KEY not in store.data + @patch(REPEAT_SUPPORTED, True) + def test_a_later_step_pauses_again_with_its_own_call(self): + shop, store = _Shop(), _FakeTaskStateStore() + first = _pause(_operator(_two_refunds, [shop.toolset()]), _context(store)) + + with patch(UPSERT, autospec=True) as upsert: + with pytest.raises(TaskAwaitingInput) as second: + _operator(_two_refunds, [shop.toolset()]).resume_after_tool_approval( + _context(store), **first.kwargs, event=APPROVE + ) + + assert shop.calls == [("refund", 1)] + assert first.kwargs["pause_number"] == 1 + assert second.value.kwargs["pause_number"] == 2 + assert second.value.kwargs["tool_call_ids"] == ["c-2"] + assert upsert.call_args.kwargs["subject"] == "Approve refund for task `t` (request 2)" + body = upsert.call_args.kwargs["body"] + assert '"order_id": 2' in body + assert '"order_id": 1' not in body + # The new pause's transcript, which its resume needs, survives the resume that paused again. + assert _TOOL_APPROVAL_TRANSCRIPT_KEY in store.data + + @patch(REPEAT_SUPPORTED, True) + def test_an_agent_that_needs_ten_approvals_pauses_ten_times_and_finishes(self): + shop, store = _Shop(), _FakeTaskStateStore() + model = _refunds_one_per_step(10) + paused = _pause(_operator(model, [shop.toolset()]), _context(store)) + pauses = 1 + output = None + while output is None: + ctx = _context(store) + with patch(UPSERT, autospec=True): + try: + output = _operator(model, [shop.toolset()]).resume_after_tool_approval( + ctx, **paused.kwargs, event=APPROVE + ) + except TaskAwaitingInput as again: + paused, pauses = again, pauses + 1 + + assert pauses == 10 + # Every resume gets its own run_id; the last one is what the run_id XCom holds. + ctx["task_instance"].xcom_push.assert_any_call(key="run_id", value="ti-1-resumed-10") + assert shop.calls == [("refund", n) for n in range(1, 11)] + assert output == " | ".join(f"refunded order {n}" for n in range(1, 11)) + assert _TOOL_APPROVAL_TRANSCRIPT_KEY not in store.data + + @patch(REPEAT_SUPPORTED, True) + def test_a_rejected_later_step_carries_on_without_that_call(self): + shop, store = _Shop(), _FakeTaskStateStore() + first = _pause(_operator(_two_refunds, [shop.toolset()]), _context(store)) + with patch(UPSERT, autospec=True): + with pytest.raises(TaskAwaitingInput) as second: + _operator(_two_refunds, [shop.toolset()]).resume_after_tool_approval( + _context(store), **first.kwargs, event=APPROVE + ) + + output = _operator(_two_refunds, [shop.toolset()]).resume_after_tool_approval( + _context(store), **second.value.kwargs, event=_reject("order 2 is outside policy") + ) + + assert shop.calls == [("refund", 1)] + assert "refunded order 1" in output + assert "order 2 is outside policy" in output + + @patch(REPEAT_SUPPORTED, True) + @pytest.mark.parametrize( + "stale_value", + [ + pytest.param(1, id="number"), + pytest.param("1", id="string"), + pytest.param({"value": 1, "schema": {"type": "integer", "const": 1}}, id="param-spec"), + pytest.param("not-a-number", id="unreadable"), + ], + ) + @patch(REPEAT_SUPPORTED, True) + def test_an_answer_from_a_form_left_open_on_an_earlier_pause_asks_again(self, stale_value): + """Reviewer B still has request 1's form open; their Approve must not run refund 2.""" + shop, store = _Shop(), _FakeTaskStateStore() + first = _pause(_operator(_two_refunds, [shop.toolset()]), _context(store)) + with patch(UPSERT, autospec=True): + with pytest.raises(TaskAwaitingInput) as second: + _operator(_two_refunds, [shop.toolset()]).resume_after_tool_approval( + _context(store), **first.kwargs, event=APPROVE + ) + stale = {**APPROVE, "params_input": {"request": stale_value}} + + with patch(UPSERT, autospec=True) as upsert: + with pytest.raises(TaskAwaitingInput) as asked_again: + _operator(_two_refunds, [shop.toolset()]).resume_after_tool_approval( + _context(store), **second.value.kwargs, event=stale + ) + + assert shop.calls == [("refund", 1)] + assert asked_again.value.kwargs["tool_call_ids"] == ["c-2"] + assert asked_again.value.kwargs["pause_number"] == 3 + assert asked_again.value.kwargs["transcript_sha256"] == second.value.kwargs["transcript_sha256"] + assert '"order_id": 2' in upsert.call_args.kwargs["body"] + assert _TOOL_APPROVAL_TRANSCRIPT_KEY in store.data + + output = _operator(_two_refunds, [shop.toolset()]).resume_after_tool_approval( + _context(store), **asked_again.value.kwargs, event={**APPROVE, "params_input": {"request": 3}} + ) + assert shop.calls == [("refund", 1), ("refund", 2)] + assert output == "refunded order 1 | refunded order 2" + + @patch(REPEAT_SUPPORTED, True) + def test_an_answer_for_the_waiting_request_is_applied(self): + shop, store = _Shop(), _FakeTaskStateStore() + first = _pause(_operator(_two_refunds, [shop.toolset()]), _context(store)) + with patch(UPSERT, autospec=True): + with pytest.raises(TaskAwaitingInput) as second: + _operator(_two_refunds, [shop.toolset()]).resume_after_tool_approval( + _context(store), **first.kwargs, event={**APPROVE, "params_input": {"request": 1}} + ) + + output = _operator(_two_refunds, [shop.toolset()]).resume_after_tool_approval( + _context(store), **second.value.kwargs, event={**APPROVE, "params_input": {"request": 2}} + ) + + assert output == "refunded order 1 | refunded order 2" + + def test_a_continuation_without_pause_number_resumes_as_the_first_pause(self): + """A task paused by an earlier provider version has no ``pause_number`` in its kwargs; it + resumes as pause 1, so a later pause is numbered 2 and its resume gets a fresh run_id.""" + shop, store = _Shop(), _FakeTaskStateStore() + paused = _pause(_operator(_two_refunds, [shop.toolset()]), _context(store)) + kwargs = {k: v for k, v in paused.kwargs.items() if k != "pause_number"} + + with patch(UPSERT, autospec=True): + with pytest.raises(TaskAwaitingInput) as second: + _operator(_two_refunds, [shop.toolset()]).resume_after_tool_approval( + _context(store), **kwargs, event=APPROVE + ) + output = _operator(_two_refunds, [shop.toolset()]).resume_after_tool_approval( + _context(store), **second.value.kwargs, event=APPROVE + ) + + assert second.value.kwargs["pause_number"] == 2 + assert output == "refunded order 1 | refunded order 2" + @pytest.mark.parametrize( ("kwargs_override", "event", "error"), [
