This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 2bf8c2e6053 Let an agent ask for tool approval in more than one step
(#74686)
2bf8c2e6053 is described below
commit 2bf8c2e6053db6addb2b08bc12b95d94e5349c41
Author: Kaxil Naik <[email protected]>
AuthorDate: Wed Oct 14 09:38:49 2026 +0100
Let an agent ask for tool approval in more than one step (#74686)
* 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.
* Pin the repeat-approval gate in the continuation-resume test
The test resumes two pauses but relied on the real gate, so it failed on
the Airflow 3.3 compat job, where a second request is refused. Patch the
gate open as the other multi-pause tests do; the decorator had been left
on the stale-answer test above it.
---
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 27c9dd84118..089d8949630 100644
--- a/providers/common/ai/docs/toolsets/hook.rst
+++ b/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
diff --git a/providers/common/ai/docs/toolsets/index.rst
b/providers/common/ai/docs/toolsets/index.rst
index 964539dc2fd..be80e9313c0 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 fff7281b5c4..a6ecd9a1b5b 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 0da729eedcd..6215733d224 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
@@ -31,7 +31,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
@@ -76,7 +76,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:
@@ -109,11 +113,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)
@@ -202,6 +213,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(
@@ -400,7 +427,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
@@ -408,9 +435,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
@@ -418,13 +446,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
@@ -515,6 +543,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
@@ -1057,9 +1087,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:
@@ -1077,13 +1109,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 = (
@@ -1092,17 +1124,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,
@@ -1116,21 +1166,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,
)
@@ -1151,16 +1202,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,
@@ -1190,6 +1248,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"):
@@ -1245,9 +1329,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..2777e238d9d 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
+
+ @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"
+
+ @patch(REPEAT_SUPPORTED, True)
+ 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"),
[
