Let an agent ask for tool approval in more than one step

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``.

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

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,15 +149,15 @@ 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
   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
   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.

providers/common/ai/docs/toolsets/hook.rst

@@ -211,8 +211,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

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.

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
 -------------------------------

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`.

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.
     """