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 602ec1005a7 Pause an OpenAI Agents SDK durable run for tool approval
in Required Actions (#76786)
602ec1005a7 is described below
commit 602ec1005a7af892b05721266e31ac7568e040df
Author: Kaxil Naik <[email protected]>
AuthorDate: Wed Nov 11 16:19:20 2026 +0000
Pause an OpenAI Agents SDK durable run for tool approval in Required
Actions (#76786)
A function tool with needs_approval makes Runner.run return normally with
final_output set to None and the pending call in result.interruptions. A
task that returned run_durably_sync(...).final_output succeeded and pushed
None to XCom, and its journal was deleted as if the run had finished.
run_durably now raises UnsupportedToolDeferralError, naming each tool and
the agent that asked for it, and the task fails without retrying, as
AgentOperator already does when it cannot resolve an approval. The failure
counts against the model's last answer, so clearing the task with its task
state kept asks the model again from that turn: earlier turns replay, and an
agent tool's answer recorded while its own run waited for approval does not.
Outside a task, run_durably is still Runner.run and returns the
interruptions.
---
providers/common/ai/docs/changelog.rst | 12 +-
.../common/ai/docs/frameworks/openai_agents.rst | 90 ++++-
.../airflow/providers/common/ai/durable/base.py | 7 +-
.../airflow/providers/common/ai/durable/journal.py | 81 +++-
.../providers/common/ai/durable/openai_agents.py | 265 +++++++++++--
.../example_dags/example_openai_agents_approval.py | 96 +++++
.../src/airflow/providers/common/ai/exceptions.py | 11 +-
.../tests/unit/common/ai/durable/test_journal.py | 73 ++++
.../unit/common/ai/durable/test_openai_agents.py | 429 ++++++++++++++++++++-
9 files changed, 994 insertions(+), 70 deletions(-)
diff --git a/providers/common/ai/docs/changelog.rst
b/providers/common/ai/docs/changelog.rst
index 40f194e50a7..17fb111bd10 100644
--- a/providers/common/ai/docs/changelog.rst
+++ b/providers/common/ai/docs/changelog.rst
@@ -35,11 +35,13 @@ Changelog
see the ``@task.llm`` documentation for how to include text from function
arguments in it.
.. note::
- ``run_durably()`` and ``run_durably_sync()`` for the OpenAI Agents SDK now
fail the task with
- ``UnsupportedToolDeferralError`` when the run stops for a tool call that
needs approval
- (``needs_approval``, or ``require_approval`` on an MCP server). Earlier
versions returned the
- stopped run with its ``result.interruptions``. Code that approved them with
- ``result.to_state()`` and ``state.approve()`` after ``run_durably_sync()``
should call
+ ``run_durably()`` and ``run_durably_sync()`` for the OpenAI Agents SDK no
longer return a run
+ that stopped for a tool call that needs approval (``needs_approval``, or
``require_approval`` on
+ an MCP server). On Airflow 3.4+ with a standard provider that has
``request_input()``, a
+ ``@task`` or ``PythonOperator`` pauses in Required Actions and the run
resumes with the
+ reviewer's answer; otherwise the task fails with
``UnsupportedToolDeferralError``. Earlier
+ versions returned the stopped run with its ``result.interruptions``. Code
that approved them
+ with ``result.to_state()`` and ``state.approve()`` after
``run_durably_sync()`` should call
``Runner.run`` and handle ``result.interruptions`` itself.
1.1.0
diff --git a/providers/common/ai/docs/frameworks/openai_agents.rst
b/providers/common/ai/docs/frameworks/openai_agents.rst
index 9c894b89818..9f706c74de0 100644
--- a/providers/common/ai/docs/frameworks/openai_agents.rst
+++ b/providers/common/ai/docs/frameworks/openai_agents.rst
@@ -264,7 +264,7 @@ Any other wrapper fails with a ``TypeError`` that names it,
such as for
TypeError: AirflowTools cannot carry ApprovalRequiredToolset
(.approval_required()) to another framework's agent: only a Pydantic AI run
pauses for the approval it asks for. Pass the toolset without it and use the
framework's own tool approval, such as needs_approval in the OpenAI Agents SDK,
whose guide says how it behaves in a task, or run the agent with AgentOperator.
-:ref:`howto/frameworks:openai-agents-limits` says what ``needs_approval`` does
in a task.
+:ref:`howto/frameworks:openai-agents-approval` says what ``needs_approval``
does in a task.
Each toolset's page covers the rest of its options: :ref:`howto/toolsets`.
Tools
@@ -582,13 +582,8 @@ stands for building the agent as the task example does:
publish(answer("Which customer has the highest ARR?"))
-``run_durably`` does not resume a run that stops for a tool call that needs
approval
-(``needs_approval``, or ``require_approval`` on an MCP server). It fails the
task with
-``UnsupportedToolDeferralError``, without retrying, since a retry would most
likely stop at
-the same call. After you remove the approval requirement, clearing the task
with its task
-state kept (``keep_task_state``, or clearing the whole Dag run) replays the
turns before that
-call and asks the model again from there, so a tool call made in that same
answer runs again
-if the model asks for it. A clear that discards the task state starts the
agent over.
+A tool call that needs approval pauses the task until a person answers; see
+:ref:`howto/frameworks:openai-agents-approval`.
What a retry replays
^^^^^^^^^^^^^^^^^^^^
@@ -647,6 +642,73 @@ A ``SandboxToolset`` without ``attach_to`` is destroyed
when its ``with`` block
which a failed try also ends, so the retry starts with a new, empty sandbox
that the
replayed conversation does not describe.
+.. _howto/frameworks:openai-agents-approval:
+
+Have a person approve tool calls
+--------------------------------
+
+A refund, a deployment or an email to a customer shouldn't happen only because
a model
+asked for it. Set ``needs_approval=True`` on the function tool and run the
agent with
+``run_durably_sync``: each call to that tool waits in Required Actions until a
person
+approves or rejects it. To have a person review the agent's answer rather than
its tool
+calls, use ``AgentOperator(enable_hitl_review=True)``, or add the standard
provider's
+``ApprovalOperator`` as the next task.
+
+.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_openai_agents_approval.py
+ :language: python
+ :start-after: [START example_openai_agents_approval]
+ :end-before: [END example_openai_agents_approval]
+
+When the model asks to call ``issue_refund``, the task stops in the
``awaiting_input``
+state without holding a worker slot, and Required Actions shows a request
titled
+``support wants to call issue_refund``. Its body names the agent and the tool,
with the
+tool's description, and lists the call's arguments after Airflow's secret
masker has
+masked them. The reviewer chooses **Approve** or **Reject**, and can fill in a
reason.
+
+The answer resumes the task, which runs from the top again.
``run_durably_sync`` replays
+the calls it recorded before the pause instead of making them again (the calls
in "What
+always runs live" above still run live), and then:
+
+- On **Approve**, the tool runs and the agent continues.
+- On **Reject**, the tool does not run, and the model gets the reviewer's
reason as the
+ call's result, or the SDK's default rejection message when the reason is
empty.
+
+Things to know:
+
+- Pausing needs Airflow 3.4 or later, a standard provider that has
+ :func:`~airflow.providers.standard.operators.hitl.request_input`, and a
``@task`` or
+ ``PythonOperator``. Elsewhere the task fails with
``UnsupportedToolDeferralError``
+ instead, without retrying.
+- Code in the task before ``run_durably_sync`` runs again every time the task
resumes. So
+ does an earlier agent run in the same task, since what it recorded is
deleted when it
+ finishes: run an agent whose tools need approval in a task of its own.
+- If a step before the paused call returns something different when the task
resumes (a
+ tool that always runs live, a prompt built from live data, a Dag changed
while the task
+ waited), the run never comes back to the call the reviewer answered. The
task then fails
+ with ``DurableJournalError``, without retrying, rather than ask a question
the reviewer
+ did not see. Clearing it then asks about the call the run stopped at
instead, or, when
+ the run finished without stopping, replays that result.
+- When one answer from the model asks for several calls that need approval,
they are asked
+ one at a time, in the order the model made them, and the agent continues
once every one
+ is answered.
+- A retry asks again about calls that had not run. A call whose result was
recorded ran
+ after it was approved, so the retry approves it without asking.
+- A function tool whose ``needs_approval`` is a function, and a tool of an MCP
server in
+ ``mcp_servers`` whose ``require_approval`` applies, are not recorded (see
"What always
+ runs live" above). Once approved, such a tool runs again every time the task
resumes or
+ retries later in the run, so make it safe to repeat, or use
``needs_approval=True`` on a
+ function tool of your own.
+- The task fails with ``UnsupportedToolDeferralError`` instead of pausing when
the run has
+ a ``session``, when the call is made by an agent used through ``as_tool`` or
reached
+ through ``handoff()``, or when the approval is for a hosted MCP tool
+ (``HostedMCPTool``). After you remove the approval requirement, clearing the
task with
+ its task state kept (``keep_task_state``, or clearing the whole Dag run)
replays the turns
+ before that call and asks the model again from there, so a tool call made in
that same
+ answer runs again if the model asks for it. A clear that discards the task
state starts
+ the agent over.
+- Every pause carries the answers the task has received so far, so keep the
approvals one
+ run asks for to a handful.
+
.. _howto/frameworks:openai-agents-limits:
Sessions, approvals and streaming
@@ -655,11 +717,9 @@ Sessions, approvals and streaming
These SDK features interact with Airflow in ways the SDK documentation does
not cover:
- A ``session`` works, but under ``run_durably`` it makes every retry run
live, as
- above.
-- A tool's ``needs_approval`` is not connected to Airflow's approval UI.
``Runner.run``
- stops at the first call that needs approval and returns with
``final_output`` set to
- ``None`` and the pending call in ``result.interruptions``, and the task
succeeds without
- an answer unless your code handles them. ``run_durably`` fails the task
instead, as
- :ref:`howto/frameworks:openai-agents-durable` describes. To have a person
approve the
- agent's answer, add the standard provider's ``ApprovalOperator`` as the next
task.
+ above, and a run with a ``session`` cannot pause for approval.
+- ``Runner.run`` stops at the first tool call that needs approval and returns
with
+ ``final_output`` set to ``None`` and the pending call in
``result.interruptions``, and
+ the task succeeds without an answer unless your code handles them.
``run_durably`` asks
+ a person instead; see :ref:`howto/frameworks:openai-agents-approval`.
- A streamed run (``Runner.run_streamed``) is not recorded for durable retries.
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/durable/base.py
b/providers/common/ai/src/airflow/providers/common/ai/durable/base.py
index 888293f4ede..e9f1ff8165a 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/durable/base.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/durable/base.py
@@ -67,5 +67,10 @@ class DurableStorageProtocol(Protocol):
def load_step(self, key: str) -> dict[str, Any] | None: ...
def delete_steps(self, keys: Iterable[str]) -> None:
- """Delete these entries. Best-effort: runs after the work succeeded,
so it must not raise."""
+ """
+ Delete these entries. Best-effort, so it must not raise.
+
+ It runs after the work succeeded, and before a task pauses for input
to drop entries
+ the attempt replaced; there the caller reads the keys back to find any
that survived.
+ """
...
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/durable/journal.py
b/providers/common/ai/src/airflow/providers/common/ai/durable/journal.py
index c4fad4b9422..574aaacedcf 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/durable/journal.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/durable/journal.py
@@ -46,6 +46,9 @@ journal for pydantic-ai. Another framework's adapter starts a
run with
``await run.run(name, kind=..., fingerprint=..., body=...)``.
* :meth:`DurableRun.claim` and :meth:`JournalStep.record` split it in two, for
a
framework whose hooks see a step's input and its output in separate
callbacks.
+
+An adapter that stops the task to wait for input (``request_input()``) calls
+:meth:`DurableRun.pause` first, so that the resumed task replays what this
attempt recorded.
"""
from __future__ import annotations
@@ -247,6 +250,8 @@ class DurableRun:
self.journal = journal
self.key = key
self.diverged = False
+ # Where this attempt stopped replaying, when it did.
+ self._diverged_at: int | None = None
self._position = 0
# Read on first use: where the last failed attempt's entries stop
being trusted,
# and how far any attempt got.
@@ -268,6 +273,9 @@ class DurableRun:
# Read with the rest of the run's entry: whether the previous attempt
failed because
# of the model's answers, which changes why the retry runs live from
``_diverge_from``.
self._model_at_fault = False
+ # Read with the rest of the run's entry: whether ``_diverge_from`` is
where the task paused
+ # rather than where an attempt failed.
+ self._paused = False
# Steps ``reserve`` numbered for the latest model response, by tool
call id, until taken.
self._reserved: dict[str, JournalStep] = {}
@@ -285,6 +293,7 @@ class DurableRun:
self._diverge_from = diverge_from if isinstance(diverge_from, int)
else None
self._high_water = high_water if isinstance(high_water, int) else 0
self._model_at_fault = meta.get("model_at_fault") is True
+ self._paused = meta.get("paused") is True
def peek(self, position: int) -> dict[str, Any] | None:
"""
@@ -317,6 +326,11 @@ class DurableRun:
)
self._reserved = reserved
+ def has_recorded(self, call_id: str) -> bool:
+ """Whether the step :meth:`reserve` kept for ``call_id`` holds a
result an earlier attempt recorded."""
+ step = self._reserved.get(call_id)
+ return step is not None and step.replayed
+
def claim_reserved(
self, call_id: str, name: str, *, kind: StepKind, fingerprint: str |
None
) -> JournalStep:
@@ -364,15 +378,25 @@ class DurableRun:
self._last_model_position = position
self._load_meta()
if not self.diverged and self._diverge_from is not None and position
>= self._diverge_from:
- self._diverge(
- position,
- name,
- reason=(
- "the model's answers from this step on failed the previous
attempt, so it is asked again"
- if self._model_at_fault
- else "the previous attempt failed before reaching this
step"
- ),
- )
+ if self._paused:
+ # Nothing went wrong: the task paused here and resumed with
the input it waited for.
+ self.diverged, self._diverged_at = True, position
+ self._peeked.clear()
+ log.info(
+ "Durable: the task resumed after a pause; this step and
the ones after it run live",
+ position=position,
+ step=name,
+ )
+ else:
+ self._diverge(
+ position,
+ name,
+ reason=(
+ "the model's answers from this step on failed the
previous attempt, so it is asked again"
+ if self._model_at_fault
+ else "the previous attempt failed before reaching this
step"
+ ),
+ )
entry = self.peek(position) if replayable else None
self._peeked.pop(position, None)
previous_replayed, self._previous_replayed = self._previous_replayed,
False
@@ -465,10 +489,44 @@ class DurableRun:
# The entry at the previous attempt's point is still that
attempt's, as this one
# recorded nothing over it (the step there raised, say), so a
retry must not replay it.
diverge_from, rerun_model = previous, self._model_at_fault
+ self._save_meta(diverge_from, model_at_fault=rerun_model)
+
+ def pause(self) -> None:
+ """
+ Record that the task is about to pause, so that resuming it replays
what this attempt recorded.
+
+ An adapter calls this before it stops the task to wait for input.
Where an earlier
+ attempt's failure, or a different path this attempt took, left entries
that must not
+ replay, this attempt has run live from there: what it recorded becomes
trusted, and
+ the earlier entries at positions it left empty (a call still waiting
for approval, a
+ tool that raised) are deleted. Positions from here on still run live
after the pause.
+ """
+ self._load_meta()
+ start = min((p for p in (self._diverge_from, self._diverged_at) if p
is not None), default=None)
+ if start is None or self._position <= start:
+ # This attempt has not run anything live past the untrusted
entries yet.
+ return
+ stale = {
+ position: build_step_key(self.key, position)
+ for position in range(start, self._position)
+ if position not in self._saved
+ }
+ if stale:
+ self.journal.storage.delete_steps(stale.values())
+ # Deleting is best-effort; an entry that survived must still not
replay.
+ survived = [position for position, key in stale.items() if
self.journal.storage.load_step(key)]
+ if survived:
+ log.warning("Durable: could not delete entries this attempt
replaced", positions=survived)
+ trusted_until = survived[0] if survived else self._position
+ self._diverge_from, self._model_at_fault = trusted_until, False
+ self._save_meta(trusted_until, model_at_fault=False, paused=True)
+
+ def _save_meta(self, diverge_from: int | None, *, model_at_fault: bool,
paused: bool = False) -> None:
meta = {
"diverge_from": diverge_from,
"high_water": max(self._high_water, self._position),
- "model_at_fault": rerun_model,
+ "model_at_fault": model_at_fault,
+ "paused": paused,
}
self.journal.storage.save_step(build_run_meta_key(self.key), meta)
@@ -501,7 +559,8 @@ class DurableRun:
self._diverge(step.position, step.name, reason=reason)
def _diverge(self, position: int, name: str, *, reason: str) -> None:
- self.diverged = True
+ if not self.diverged:
+ self.diverged, self._diverged_at = True, position
self._peeked.clear()
log.warning(
"Durable: the run took a different path from the previous attempt;
this step and "
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/durable/openai_agents.py
b/providers/common/ai/src/airflow/providers/common/ai/durable/openai_agents.py
index d603a1dae6a..16cce631e80 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/durable/openai_agents.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/durable/openai_agents.py
@@ -24,7 +24,7 @@ from __future__ import annotations
import copy
import json
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, NoReturn
from pydantic import TypeAdapter, ValidationError
from pydantic_core import to_jsonable_python
@@ -64,21 +64,35 @@ except ImportError as e:
from airflow.providers.common.ai.durable.fingerprint import _digest,
fingerprint_tool_call
from airflow.providers.common.ai.durable.journal import _EXECUTING_STEP,
current_journal
-from airflow.providers.common.ai.exceptions import UnsupportedToolDeferralError
+from airflow.providers.common.ai.exceptions import DurableJournalError,
UnsupportedToolDeferralError
+from airflow.providers.common.ai.mixins.approval import LLMApprovalMixin
from airflow.providers.common.ai.tools.openai_agents import _REPLAY_RETRY,
_RETRY_COUNT, _RetryLimitReached
from airflow.providers.common.ai.utils.coroutines import run_coroutine_sync
-from airflow.providers.common.ai.utils.masking import mask_secrets
+from airflow.providers.common.ai.utils.masking import dumps_masked,
mask_secrets
from airflow.providers.common.ai.utils.task_logger import get_task_logger
+from airflow.providers.common.compat.sdk import Param, get_current_context
+from airflow.providers.common.compat.version_compat import AIRFLOW_V_3_4_PLUS
+from airflow.providers.standard.operators.python import PythonOperator
if TYPE_CHECKING:
- from collections.abc import AsyncIterator
+ from collections.abc import AsyncIterator, Callable
- from agents import AgentOutputSchemaBase, Handoff, ModelSettings,
ModelTracing, RunResult, Tool
+ from agents import (
+ AgentOutputSchemaBase,
+ Handoff,
+ ModelSettings,
+ ModelTracing,
+ RunResult,
+ RunState,
+ Tool,
+ ToolApprovalItem,
+ )
from agents.items import TResponseInputItem, TResponseStreamEvent
from agents.retry import ModelRetryAdvice, ModelRetryAdviceRequest
from agents.tool_context import ToolContext
from airflow.providers.common.ai.durable.journal import DurableRun,
JournalStep
+ from airflow.providers.standard.utils.hitl import InputResponse
__all__ = ["run_durably", "run_durably_sync"]
@@ -108,8 +122,9 @@ _DEFAULT_FAILURE_HANDLED =
"_function_tool_default_failure_handled"
_SERVER_SIDE_STATE = ("previous_response_id", "conversation_id",
"auto_previous_response_id")
# Set on the copy of a function tool that reports through the journal, so that
reserving a position
-# skips tools nothing wraps, such as those of MCP servers or of a handoff
agent with a named model.
-_DURABLE_MARKER = "_airflow_durable_tool"
+# skips tools nothing wraps, such as those of MCP servers or of a handoff
agent with a named model,
+# and on the copy of an agent whose calls replay, so that only its approvals
pause the task.
+_DURABLE_MARKER = "_airflow_durable"
async def run_durably(
@@ -162,6 +177,12 @@ async def run_durably(
A ``session`` makes every retry run live: the SDK adds the failed try's
prompt to it
before the first model call, so that call no longer asks what it asked
then.
+ A tool call that needs approval (``needs_approval``) is asked with the
standard
+ provider's ``request_input()``, which pauses the task in Required Actions.
The answer
+ resumes the task, the run replays up to the call, and the call is
approved, or rejected
+ with the reviewer's reason as its result. Steps are deleted only once a
run ends without
+ a pending approval.
+
Outside a task there is no journal, and this is ``Runner.run``.
:param starting_agent: The agent to run.
@@ -172,7 +193,9 @@ async def run_durably(
``auto_previous_response_id`` is passed: a replayed call never reaches
the
conversation OpenAI keeps for them.
:raises UnsupportedToolDeferralError: When the run stops for a tool call
that needs
- approval, which ``run_durably`` cannot resume. The task fails without
retrying.
+ approval and cannot pause for it: before Airflow 3.4 or without a
standard provider
+ that has ``request_input()``, with a ``session``, or for a call by an
agent whose
+ model is not recorded. The task fails without retrying.
"""
if (journal := current_journal()) is None:
return await Runner.run(starting_agent, input, run_config=run_config,
**kwargs)
@@ -187,24 +210,25 @@ async def run_durably(
if isinstance(config.model, Model):
config.model = _DurableModel(config.model, durable_run)
agent = _get_durable_agent(starting_agent, durable_run, {})
- try:
- result = await Runner.run(agent, input, run_config=config, **kwargs)
- except BaseException as error:
- durable_run.fail(error, model_at_fault=isinstance(error,
_MODEL_AT_FAULT))
- raise
- if result.interruptions:
- tools = ", ".join(sorted({f"{item.tool_name} ({item.agent.name})" for
item in result.interruptions}))
- unanswered = UnsupportedToolDeferralError(
- f"The agent stopped to ask for approval to call {tools}, and
run_durably() cannot "
- "resume a run that waits for approval. Remove the approval
requirement (needs_approval, "
- "or require_approval on an MCP server), or call Runner.run and
handle "
- "result.interruptions yourself."
+
+ async def run(run_input: str | list[TResponseInputItem] | RunState[Any])
-> RunResult:
+ try:
+ return await Runner.run(agent, run_input, run_config=config,
**kwargs)
+ except BaseException as error:
+ durable_run.fail(error, model_at_fault=isinstance(error,
_MODEL_AT_FAULT))
+ raise
+
+ result = await run(input)
+ # Each answer comes from a person through Required Actions. Asking pauses
the task, which
+ # then runs again from the top: the run replays up to here and gets the
answer back.
+ answered: set[str] = set()
+ while result.interruptions:
+ result = await run(
+ _answer_approvals(result, durable_run, answered,
has_session=kwargs.get("session") is not None)
)
- # Counted against the model's last answer, so a clear asks the model
again from there:
- # what was recorded after it, such as an agent tool's answer given
while it waited for
- # approval, never replays, and the turns before it still do.
- durable_run.fail(unanswered, model_at_fault=True)
- raise unanswered
+ if _answer_left_unread(durable_run, answered):
+ _did_not_come_back(durable_run)
+ # Only now: a run paused for approval replays what it recorded when the
task resumes.
if journal.clean_up_after_run:
durable_run.cleanup()
return result
@@ -221,11 +245,191 @@ def run_durably_sync(
return run_coroutine_sync(run_durably(starting_agent, input,
run_config=run_config, **kwargs))
+def _answer_approvals(
+ result: RunResult, durable_run: DurableRun, answered: set[str], *,
has_session: bool
+) -> RunState[Any]:
+ """
+ Return the run's state with every pending tool call approved or rejected
by a person.
+
+ A call is asked through ``request_input()``, which pauses the task the
first time and
+ returns the answer once the task resumes; its key goes into ``answered``.
A call whose
+ result is already recorded was approved and ran in an earlier run of the
task, so it is
+ approved again without asking.
+
+ :raises UnsupportedToolDeferralError: When the run cannot pause for the
approval.
+ :raises DurableJournalError: When the task resumed with an answer but the
run did not come
+ back to the call it answers.
+ """
+ if (request_input := _get_request_input()) is None:
+ _cannot_pause(
+ result,
+ durable_run,
+ "pausing a task for input needs Airflow 3.4+ and a standard
provider that has request_input()",
+ )
+ if (reason := _why_approval_cannot_pause(result, has_session)) is not None:
+ _cannot_pause(result, durable_run, reason)
+ durable_run.pause()
+ state = result.to_state()
+ for item in result.interruptions:
+ call = item.raw_item
+ if durable_run.has_recorded(call.call_id):
+ log.info(
+ "Durable: %s ran after it was approved in an earlier run of
this task; approving it again",
+ item.tool_name,
+ )
+ state.approve(item)
+ continue
+ # The call comes back the same on every resume, as the answer that
made it replays.
+ # Its fingerprint keeps an answer from applying to a different call
that reuses its id.
+ key =
f"{durable_run.key}:{call.call_id}:{_build_tool_fingerprint(call.name,
call.arguments, call.call_id)}"
+ if _answer_left_unread(durable_run, answered, asking=key):
+ # The task resumed with an answer to a call this run never
reached; asking about
+ # this one instead would never end, and request_input() refuses to.
+ _did_not_come_back(durable_run, stopped_at_another_call=True)
+ answer = request_input(
+ subject=f"{item.agent.name} wants to call {item.tool_name}",
+ body=_describe_call(item),
+ options=[LLMApprovalMixin.APPROVE, LLMApprovalMixin.REJECT],
+ # An empty default, not None: some UI versions show a null one as
"[object Object]".
+ params={
+ "reason": Param(
+ "",
+ type=["string", "null"],
+ description="Sent to the agent when you reject the call
(optional).",
+ )
+ },
+ key=key,
+ )
+ answered.add(key)
+ responder = LLMApprovalMixin._describe_responder(dict(answer))
+ if answer["chosen_options"] == [LLMApprovalMixin.APPROVE]:
+ log.info("Durable: %s approved by %s", item.tool_name, responder)
+ state.approve(item)
+ else:
+ log.info("Durable: %s rejected by %s", item.tool_name, responder)
+ # A field left empty comes back as None or "", and some UI
versions send the param's spec.
+ reason = answer["params_input"].get("reason")
+ reason = reason.strip() if isinstance(reason, str) else ""
+ state.reject(item, rejection_message=reason or None)
+ return state
+
+
+def _answer_left_unread(durable_run: DurableRun, answered: set[str], *,
asking: str | None = None) -> bool:
+ """
+ Whether the task resumed with an answer to an approval of this run that
the run never asked again for.
+
+ ``asking`` is the key the run is about to ask for; an answer already given
to it is read
+ by that call rather than left unread.
+ """
+ try:
+ # Optional: the answer log comes with request_input(), newer than the
standard provider common.ai requires.
+ from airflow.providers.standard.utils.hitl import answer_log_var
+ except ImportError:
+ return False
+ answer_log = answer_log_var.get()
+ return (
+ answer_log is not None
+ and answer_log.latest.startswith(f"{durable_run.key}:")
+ and answer_log.latest not in answered
+ and (asking is None or asking not in answer_log.answers)
+ )
+
+
+def _did_not_come_back(durable_run: DurableRun, *, stopped_at_another_call:
bool = False) -> NoReturn:
+ """
+ Fail the task resumed with an answer whose call the run never reached
again.
+
+ ``stopped_at_another_call`` is set when the run stopped at another call
that needs approval
+ (a clear asks about that one), and not when it finished without one (a
clear replays the result).
+ """
+ then = (
+ "Clear the task to ask about the call it stopped at instead."
+ if stopped_at_another_call
+ else "It finished without that call; clearing the task replays that
result without asking."
+ )
+ error = DurableJournalError(
+ "The task resumed with an answer to a tool approval, but the agent run
did not come back "
+ "to the call it answers: a step before that call gave a different
result this time, such "
+ "as a tool that always runs live, a prompt built from live data, or a
change to the Dag "
+ f"while the task waited. {then}"
+ )
+ durable_run.fail(error)
+ raise error
+
+
+def _get_request_input() -> Callable[..., InputResponse] | None:
+ """Return the standard provider's ``request_input()``, or ``None`` where
it cannot pause a task."""
+ if not AIRFLOW_V_3_4_PLUS:
+ return None
+ try:
+ # Optional: request_input() is newer than the standard provider
release common.ai requires.
+ from airflow.providers.standard.operators.hitl import request_input
+ except ImportError:
+ return None
+ return request_input
+
+
+def _cannot_pause(result: RunResult, durable_run: DurableRun, reason: str) ->
NoReturn:
+ tools = ", ".join(sorted({f"{item.tool_name} ({item.agent.name})" for item
in result.interruptions}))
+ error = UnsupportedToolDeferralError(
+ f"The agent stopped to ask for approval to call {tools}, and
run_durably() cannot "
+ f"pause the task for it: {reason}. Remove the approval requirement
(needs_approval, "
+ "or require_approval on an MCP server), or call Runner.run and handle "
+ "result.interruptions yourself."
+ )
+ # Counted against the model's last answer, so a clear asks the model again
from there:
+ # what was recorded after it, such as an agent tool's answer given while
it waited for
+ # approval, never replays, and the turns before it still do.
+ durable_run.fail(error, model_at_fault=True)
+ raise error
+
+
+def _why_approval_cannot_pause(result: RunResult, has_session: bool) -> str |
None:
+ if has_session:
+ return (
+ "a run with a session asks the model again on every resume, so the
call it paused on "
+ "never comes back to receive the answer"
+ )
+ for item in result.interruptions:
+ if not isinstance(item.raw_item, ResponseFunctionToolCall) or not
getattr(
+ item.agent, _DURABLE_MARKER, False
+ ):
+ return (
+ f"{item.tool_name} is not a function tool of an agent whose
model run_durably() "
+ "records, such as an agent used through as_tool() or reached
through handoff(), "
+ "so its call would not come back the same when the task
resumes"
+ )
+ try:
+ task = get_current_context()["task"]
+ except RuntimeError:
+ # A journal set with journal_scope() outside a running task, such as
in a unit test.
+ return "there is no running task to pause"
+ if not isinstance(task, PythonOperator):
+ return f"request_input() can pause a @task or PythonOperator, and this
is a {type(task).__name__}"
+ return None
+
+
+def _describe_call(item: ToolApprovalItem) -> str:
+ """Describe the call for the reviewer in Markdown: the agent, the tool and
its masked arguments."""
+ arguments = dumps_masked(_parse_arguments(item.raw_item.arguments),
indent=2)
+ tool = next(
+ (tool for tool in item.agent.tools if isinstance(tool, FunctionTool)
and tool.name == item.tool_name),
+ None,
+ )
+ description = f" — {tool.description}" if tool is not None and
tool.description else ""
+ return (
+ f"**Agent:** {item.agent.name}\n\n**Tool:**
{item.tool_name}{description}\n\n"
+ f"**Arguments**\n\n```json\n{arguments}\n```"
+ )
+
+
def _get_durable_agent(agent: Agent[Any], durable_run: DurableRun, seen:
dict[int, Agent[Any]]) -> Agent[Any]:
"""Return a copy of ``agent`` whose model, function tools and handoff
agents go through the journal."""
if (done := seen.get(id(agent))) is not None:
return done
durable_agent = agent.clone()
+ # Its calls replay the same on resume, so an approval it asks for can
pause the task.
+ setattr(durable_agent, _DURABLE_MARKER, True)
# Recorded before the handoffs are copied, so agents that hand off to each
other copy once.
seen[id(agent)] = durable_agent
if isinstance(agent.model, Model):
@@ -302,11 +506,16 @@ def _get_durable_tool(tool: Tool, durable_run:
DurableRun) -> Tool:
return durable_tool
-def _build_tool_fingerprint(name: str, arguments: str, call_id: str) -> str |
None:
+def _parse_arguments(arguments: str) -> Any:
+ """Return a tool call's arguments as the model sent them: parsed JSON, or
the raw text when it is not JSON."""
try:
- parsed = json.loads(arguments) if arguments else {}
+ return json.loads(arguments) if arguments else {}
except json.JSONDecodeError:
- parsed = arguments
+ return arguments
+
+
+def _build_tool_fingerprint(name: str, arguments: str, call_id: str) -> str |
None:
+ parsed = _parse_arguments(arguments)
return fingerprint_tool_call(
name,
# As recorded model answers are masked, so a replayed call's arguments
can be too.
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_openai_agents_approval.py
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_openai_agents_approval.py
new file mode 100644
index 00000000000..a6c2776d9eb
--- /dev/null
+++
b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_openai_agents_approval.py
@@ -0,0 +1,96 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Have a person approve each refund an OpenAI Agents SDK support agent wants to
issue.
+
+The agent looks orders up on its own, but ``issue_refund`` sets
``needs_approval=True``, so
+:func:`~airflow.providers.common.ai.durable.openai_agents.run_durably_sync`
pauses the task
+in Required Actions for every refund and resumes it with the reviewer's answer.
+
+Before running:
+
+1. Install the SDK: ``pip install "openai-agents>=0.20.0"``.
+2. Create a connection (``LLM_CONN_ID``, default ``openai_default``) whose
password is
+ an OpenAI API key. Set its host to the base URL of any other
OpenAI-compatible
+ endpoint, such as a gateway or a self-hosted server.
+3. Use Airflow 3.4 or later, with a standard provider that has
+ :func:`~airflow.providers.standard.operators.hitl.request_input`.
+"""
+
+from __future__ import annotations
+
+import os
+
+from airflow.providers.common.compat.sdk import BaseHook, dag, task
+
+LLM_CONN_ID = os.environ.get("LLM_CONN_ID", "openai_default")
+LLM_MODEL = os.environ.get("LLM_MODEL", "gpt-5.6-luna")
+
+ORDERS = {
+ "A-1001": {"total": 42.00, "status": "delivered damaged"},
+ "A-1002": {"total": 18.50, "status": "never delivered"},
+}
+DEFAULT_TICKET = "Order A-1001 arrived broken and A-1002 never came. Please
refund both in full."
+
+
+# [START example_openai_agents_approval]
+@dag(tags=["example"])
+def example_openai_agents_approval():
+ """Resolve a support ticket with an agent whose refunds a person approves
in Required Actions."""
+
+ @task
+ def handle_ticket(ticket: str = DEFAULT_TICKET) -> str:
+ # The SDK is an optional dependency, so it is imported when the task
runs, not at parse time.
+ from agents import Agent, OpenAIChatCompletionsModel, RunConfig,
function_tool
+ from openai import AsyncOpenAI
+
+ from airflow.providers.common.ai.durable.openai_agents import
run_durably_sync
+
+ @function_tool
+ def get_order(order_id: str) -> dict:
+ """Look an order up by its id."""
+ return ORDERS.get(order_id, {"error": f"No order {order_id}"})
+
+ @function_tool(needs_approval=True)
+ def issue_refund(order_id: str, amount: float, reason: str) -> str:
+ """Refund part or all of an order."""
+ return f"Refunded {amount:.2f} on {order_id}"
+
+ llm = BaseHook.get_connection(LLM_CONN_ID)
+ agent = Agent(
+ name="support",
+ model=OpenAIChatCompletionsModel(
+ model=LLM_MODEL,
+ openai_client=AsyncOpenAI(api_key=llm.password,
base_url=llm.host or None),
+ ),
+ instructions=(
+ "You resolve support tickets. Check each order with get_order,
then refund "
+ "only what the order supports with issue_refund, one call per
order."
+ ),
+ tools=[get_order, issue_refund],
+ )
+ # The SDK sends traces to OpenAI by default; see the guide's tracing
section.
+ result = run_durably_sync(agent, ticket,
run_config=RunConfig(tracing_disabled=True))
+ return result.final_output
+
+ handle_ticket()
+
+
+# [END example_openai_agents_approval]
+
+
+example_openai_agents_approval()
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 e2d190b461d..67facfa980e 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/exceptions.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/exceptions.py
@@ -70,8 +70,9 @@ class UnsupportedToolDeferralError(AirflowFailException):
In ``AgentOperator``, either the tool hands its work to an external
system, or it needs
approval where approval is not available (before Airflow 3.3, or with
``durable``,
``enable_hitl_review``, code mode or a ``SandboxToolset``). In
``run_durably`` for the
- OpenAI Agents SDK, a tool call needs approval. A retry would repeat the
same call, so the
- task fails without retrying.
+ OpenAI Agents SDK, a tool call needs approval and the run cannot pause the
task for it,
+ such as before Airflow 3.4 or with a ``session``. A retry would repeat the
same call, so
+ the task fails without retrying.
"""
@@ -80,8 +81,10 @@ class DurableJournalError(AirflowFailException):
Raised when durable execution cannot record or replay a step.
A step's result is not JSON-serializable, or the journal holds a result
this version
- cannot replay. The same step fails the same way on every attempt, so the
task fails
- without retrying.
+ cannot replay. Or a task resumed with the answer to an approval, and the
run, which
+ took a different path this time, never came back to the call the answer is
for. The
+ same step fails the same way on every attempt, and an answer cannot be
matched to a
+ call nobody was asked about, so the task fails without retrying.
"""
diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_journal.py
b/providers/common/ai/tests/unit/common/ai/durable/test_journal.py
index f95c983156c..ebdbff76f00 100644
--- a/providers/common/ai/tests/unit/common/ai/durable/test_journal.py
+++ b/providers/common/ai/tests/unit/common/ai/durable/test_journal.py
@@ -420,6 +420,79 @@ class TestFailuresDuringReplay:
assert results == ["a recorded", "b recorded", "c recorded"]
+async def _failed_attempt_then_retry_that_diverges_at_once(memory_storage):
+ """
+ Return a retry about to pause at position 6 that ran 0..5 live and left 1
unrecorded.
+
+ The failed attempt recorded 0..5, with ``try1-1`` at 1 and its model's
answer at 4, and
+ failed because of that answer, so it stopped being trusted from 4. The
retry's first step
+ asks something else, so it stops replaying at 0, below that; its tool at 1
raises.
+ """
+ first = DurableJournal(memory_storage).start_run()
+ for position in range(6):
+ await first.run(
+ f"s{position}",
+ kind="model" if position == 4 else "tool",
+ fingerprint=f"s{position}",
+ body=lambda p=position: ok(f"try1-{p}"),
+ )
+ first.fail(RuntimeError("the model's answers failed the run"),
model_at_fault=True)
+ assert
memory_storage.entries[build_run_meta_key(first.key)]["diverge_from"] == 4
+
+ retry = DurableJournal(memory_storage).start_run()
+ await retry.run("s0", kind="tool", fingerprint="changed", body=lambda:
ok("retry-0"))
+ with pytest.raises(RuntimeError):
+ await retry.run("s1", kind="tool", fingerprint="s1", body=boom)
+ for position in range(2, 6):
+ await retry.run(
+ f"s{position}",
+ kind="model" if position == 4 else "tool",
+ fingerprint=f"s{position}",
+ body=lambda p=position: ok(f"retry-{p}"),
+ )
+ return retry
+
+
+class TestPause:
+ @pytest.mark.asyncio
+ async def
test_entries_the_attempt_left_unrecorded_below_the_stored_boundary_do_not_replay(
+ self, memory_storage
+ ):
+ retry = await
_failed_attempt_then_retry_that_diverges_at_once(memory_storage)
+
+ retry.pause()
+
+ resumed = DurableJournal(memory_storage).start_run()
+ replayed = await resumed.run("s0", kind="tool", fingerprint="changed",
body=lambda: ok("live"))
+ rerun = await resumed.run("s1", kind="tool", fingerprint="s1",
body=lambda: ok("live"))
+ assert (replayed, rerun) == ("retry-0", "live")
+ assert
memory_storage.entries[build_run_meta_key(retry.key)]["diverge_from"] == 6
+
+ @pytest.mark.asyncio
+ async def test_an_entry_that_could_not_be_deleted_is_not_trusted(self,
memory_storage):
+ retry = await
_failed_attempt_then_retry_that_diverges_at_once(memory_storage)
+
+ with mock.patch.object(memory_storage, "delete_steps", autospec=True):
+ retry.pause()
+
+ # The failed attempt's entry at 1 is still there, so the resume runs
everything from 1 live.
+ assert
memory_storage.entries[build_run_meta_key(retry.key)]["diverge_from"] == 1
+
+ @pytest.mark.asyncio
+ async def
test_a_resume_still_replaying_leaves_the_boundary_where_it_is(self,
memory_storage):
+ retry = await
_failed_attempt_then_retry_that_diverges_at_once(memory_storage)
+ retry.pause()
+ meta = dict(memory_storage.entries[build_run_meta_key(retry.key)])
+
+ resumed = DurableJournal(memory_storage).start_run()
+ await resumed.run("s0", kind="tool", fingerprint="changed",
body=lambda: ok("live"))
+ resumed.pause()
+
+ assert memory_storage.entries[build_run_meta_key(resumed.key)] == meta
+ assert await resumed.run("s1", kind="tool", fingerprint="s1",
body=lambda: ok("live")) == "live"
+ assert await resumed.run("s2", kind="tool", fingerprint="s2",
body=lambda: ok("live")) == "retry-2"
+
+
class TestCleanup:
@pytest.mark.asyncio
async def
test_cleanup_reaches_a_run_started_by_a_step_that_now_replays(self,
memory_storage):
diff --git
a/providers/common/ai/tests/unit/common/ai/durable/test_openai_agents.py
b/providers/common/ai/tests/unit/common/ai/durable/test_openai_agents.py
index 39ecd4d8f58..be738e68f58 100644
--- a/providers/common/ai/tests/unit/common/ai/durable/test_openai_agents.py
+++ b/providers/common/ai/tests/unit/common/ai/durable/test_openai_agents.py
@@ -28,6 +28,7 @@ from typing import Any
from unittest import mock
import pytest
+import structlog
pytest.importorskip("agents")
@@ -65,10 +66,16 @@ from airflow.providers.common.ai.durable.openai_agents
import (
_get_durable_tool,
run_durably_sync,
)
-from airflow.providers.common.ai.exceptions import UnsupportedToolDeferralError
+from airflow.providers.common.ai.exceptions import DurableJournalError,
UnsupportedToolDeferralError
from airflow.providers.common.ai.tools import AirflowTool, ToolCallError,
ToolResult
from airflow.providers.common.ai.tools._from_toolset import
airflow_tools_from_toolset
from airflow.providers.common.ai.tools.openai_agents import AirflowTools
+from airflow.providers.standard.operators.hitl import request_input
+from airflow.sdk import DAG, BaseOperator, task
+from airflow.sdk.execution_time.comms import AwaitInputTask,
CreateHITLDetailPayload, SetXCom
+from airflow.sdk.execution_time.task_runner import run
+from airflow.sdk.serde import deserialize, serialize
+from airflow.utils.state import TaskInstanceState
from unit.common.ai.durable.memory_storage import MemoryStorage
@@ -93,9 +100,13 @@ class _ScriptedModel(Model):
argument: str = "key",
fail_first_request: bool = False,
call_id_prefix: str = "call",
+ fresh_call_ids: bool = False,
+ value_prefix: str = "k",
) -> None:
self.lookups = lookups
+ self.value_prefix = value_prefix
self.call_id_prefix = call_id_prefix
+ self.fresh_call_ids = fresh_call_ids
self.argument = argument
self.fail_first_request = fail_first_request
self.fail_on_answer = fail_on_answer
@@ -119,9 +130,11 @@ class _ScriptedModel(Model):
output = [
ResponseFunctionToolCall(
id=f"fc-{len(outputs)}",
- call_id=f"{self.call_id_prefix}-{len(outputs)}",
+ # A real model gives every call it makes a new id.
+ call_id=f"{self.call_id_prefix}-{len(outputs)}"
+ + (f"-r{self.calls}" if self.fresh_call_ids else ""),
name="lookup",
- arguments=json.dumps({self.argument: f"k{len(outputs)}"}),
+ arguments=json.dumps({self.argument:
f"{self.value_prefix}{len(outputs)}"}),
type="function_call",
)
]
@@ -1322,7 +1335,7 @@ class TestDurableExceptions:
[True, lambda context, arguments, call_id: True],
ids=["always", "by-rule"],
)
- def test_a_run_that_stops_for_approval_fails_naming_the_tool(self,
needs_approval):
+ def
test_without_request_input_a_run_that_stops_for_approval_fails_naming_the_tool(self,
needs_approval):
calls = []
@function_tool(needs_approval=needs_approval)
@@ -1332,8 +1345,16 @@ class TestDurableExceptions:
return f"{key}=42"
with (
+ mock.patch(
+
"airflow.providers.common.ai.durable.openai_agents._get_request_input",
+ autospec=True,
+ return_value=None,
+ ),
journal_scope(DurableJournal(MemoryStorage(),
clean_up_after_run=True)),
- pytest.raises(UnsupportedToolDeferralError, match=r"approval to
call lookup \(analyst\)"),
+ pytest.raises(
+ UnsupportedToolDeferralError,
+ match=r"approval to call lookup \(analyst\).*needs Airflow
3\.4\+ and a standard provider",
+ ),
):
run_durably_sync(
Agent(name="analyst", model=_ScriptedModel(), tools=[lookup]),
"go", run_config=NO_TRACING
@@ -1341,6 +1362,25 @@ class TestDurableExceptions:
assert calls == []
+ def
test_outside_a_task_a_run_that_stops_for_approval_fails_naming_the_tool(self):
+ """A journal set in a unit test has no task to pause, so the run fails
as it does without request_input()."""
+ calls: list[str] = []
+
+ with (
+ journal_scope(DurableJournal(MemoryStorage(),
clean_up_after_run=True)),
+ pytest.raises(
+ UnsupportedToolDeferralError,
+ match=r"approval to call lookup \(analyst\).*there is no
running task to pause",
+ ),
+ ):
+ run_durably_sync(
+ Agent(name="analyst", model=_ScriptedModel(),
tools=[_lookup_tool(calls)]),
+ "go",
+ run_config=NO_TRACING,
+ )
+
+ assert calls == []
+
def
test_a_clear_after_an_agent_tool_stopped_for_approval_asks_again_from_that_turn(self):
"""Earlier turns replay; the agent tool's answer, given while its own
run waited, does not."""
storage, calls = MemoryStorage(), []
@@ -1367,10 +1407,387 @@ class TestDurableExceptions:
return journal
# The specialist answers the first time it is asked, and waits for
approval the second time.
- with pytest.raises(UnsupportedToolDeferralError, match=r"approval to
call lookup \(specialist\)"):
+ with pytest.raises(UnsupportedToolDeferralError, match=r"lookup
\(specialist\).*through as_tool\(\)"):
attempt(lambda context, arguments, call_id: bool(calls),
call_id_prefix="first")
journal = attempt(False, call_id_prefix="second")
# The first turn and the specialist's first answer replay; the second
turn runs again.
assert journal.stats.replayed == {"model": 1, "tool": 1}
assert calls == ["k0", "k0"]
+
+
+def _answer(option: str, reason: str = "") -> dict[str, Any]:
+ """The answer event the Core API puts in the paused task's next_kwargs."""
+ return {
+ "chosen_options": [option],
+ "params_input": {"reason": reason},
+ "responded_at": datetime.datetime(2026, 11, 5, tzinfo=datetime.UTC),
+ "responded_by_user": {"id": "1", "name": "reviewer"},
+ "timedout": False,
+ }
+
+
+def _sent(mock_supervisor_comms, msg_type: type) -> list[Any]:
+ messages = (
+ call.kwargs.get("msg", call.args[0] if call.args else None)
+ for call in mock_supervisor_comms.send.call_args_list
+ )
+ return [msg for msg in messages if isinstance(msg, msg_type)]
+
+
+class _Reviewer:
+ """
+ Runs a task through the task runner and answers each Required Actions
request it parks on.
+
+ A resumed run is handed the parked kwargs plus the answer, as the Core API
response
+ handler hands them back, so ``PythonOperator.resume_with_input`` runs for
real.
+ """
+
+ def __init__(self, create_runtime_ti, mock_supervisor_comms) -> None:
+ self.create_runtime_ti = create_runtime_ti
+ self.comms = mock_supervisor_comms
+ self.decide: Any = lambda request: _answer("Approve")
+ self.keys: list[str] = []
+
+ @property
+ def requests(self) -> list[CreateHITLDetailPayload]:
+ return _sent(self.comms, CreateHITLDetailPayload)
+
+ def run(self, op) -> tuple[Any, BaseException | None]:
+ ti = self.create_runtime_ti(task=op)
+ state, msg, error = _run_ti(ti)
+ while isinstance(msg, AwaitInputTask):
+ parked = deserialize(msg.next_kwargs)
+ assert isinstance(parked, dict)
+ self.keys.append(parked["key"])
+ resumed = self.create_runtime_ti(task=op, ti_id=ti.id)
+ resumed._ti_context_from_server.next_method = msg.next_method
+ resumed._ti_context_from_server.next_kwargs = serialize(
+ {**parked, "event": self.decide(self.requests[-1])}
+ )
+ state, msg, error = _run_ti(resumed)
+ return state, error
+
+ def returned(self) -> list[Any]:
+ return [msg.value for msg in _sent(self.comms, SetXCom) if msg.key ==
"return_value"]
+
+
+def _run_ti(ti) -> tuple[Any, Any, BaseException | None]:
+ return run(ti, context=ti.get_template_context(),
log=structlog.get_logger())
+
+
[email protected]
+def reviewer(create_runtime_ti, mock_supervisor_comms):
+ return _Reviewer(create_runtime_ti, mock_supervisor_comms)
+
+
+def _lookup_tool(calls: list[str], needs_approval: Any = True) -> Any:
+ @function_tool(needs_approval=needs_approval)
+ def lookup(key: str) -> str:
+ """Look a key up."""
+ calls.append(key)
+ return f"{key}=42"
+
+ return lookup
+
+
+def _agent_task(storage: MemoryStorage, agent: Any, journals:
list[DurableJournal], **kwargs: Any):
+ """
+ A ``@task`` that runs an agent durably; each run, first or resumed, gets a
fresh journal over ``storage``.
+
+ ``agent`` is the agent, or a function of the run's number that returns the
agent that run uses.
+ """
+ with DAG(dag_id="agent_with_approvals"):
+
+ @task
+ def handle_ticket() -> str:
+ run_agent = agent(len(journals)) if callable(agent) else agent
+ journal = DurableJournal(storage, clean_up_after_run=True)
+ journals.append(journal)
+ with journal_scope(journal):
+ return run_durably_sync(run_agent, "go",
run_config=NO_TRACING, **kwargs).final_output
+
+ return handle_ticket().operator
+
+
+class TestToolApproval:
+ @pytest.mark.parametrize(
+ "needs_approval",
+ [True, lambda context, arguments, call_id: True],
+ ids=["always", "by-rule"],
+ )
+ def test_an_approved_call_runs_once_and_no_model_call_is_made_twice(self,
reviewer, needs_approval):
+ storage, calls, journals, model = MemoryStorage(), [], [],
_ScriptedModel()
+ agent = Agent(name="analyst", model=model, tools=[_lookup_tool(calls,
needs_approval)])
+
+ state, error = reviewer.run(_agent_task(storage, agent, journals))
+
+ assert (state, error) == (TaskInstanceState.SUCCESS, None)
+ assert reviewer.returned() == ["answer: k0=42"]
+ assert calls == ["k0"]
+ # The request that asked for the call replayed on resume; only the one
that answers is new.
+ assert model.calls == 2
+ assert journals[-1].stats.replayed == {"model": 1}
+ [request] = reviewer.requests
+ assert request.subject == "analyst wants to call lookup"
+ assert request.options == ["Approve", "Reject"]
+ assert "**Tool:** lookup — Look a key up." in request.body
+ assert '"key": "k0"' in request.body
+ assert reviewer.keys[0].split(":")[1] == "call-0"
+ assert storage.steps() == []
+
+ @pytest.mark.parametrize(
+ ("reason", "told_to_model"),
+ [
+ ("Refunds over 40 need a manager.", "Refunds over 40 need a
manager."),
+ (" ", "Tool execution was not approved."),
+ ],
+ ids=["with-reason", "blank-reason"],
+ )
+ def test_a_rejected_call_tells_the_model_the_reviewers_reason(self,
reviewer, reason, told_to_model):
+ storage, calls, journals, model = MemoryStorage(), [], [],
_ScriptedModel()
+ reviewer.decide = lambda request: _answer("Reject", reason)
+ agent = Agent(name="analyst", model=model, tools=[_lookup_tool(calls)])
+
+ state, _ = reviewer.run(_agent_task(storage, agent, journals))
+
+ assert state == TaskInstanceState.SUCCESS
+ assert calls == []
+ assert reviewer.returned() == [f"answer: {told_to_model}"]
+ assert model.calls == 2
+
+ def test_a_later_approval_replays_the_calls_before_it(self, reviewer):
+ storage, calls, journals, model = MemoryStorage(), [], [],
_ScriptedModel(lookups=2)
+ agent = Agent(name="analyst", model=model, tools=[_lookup_tool(calls)])
+
+ state, _ = reviewer.run(_agent_task(storage, agent, journals))
+
+ assert state == TaskInstanceState.SUCCESS
+ assert reviewer.returned() == ["answer: k1=42"]
+ assert len(reviewer.requests) == 2
+ # Across three runs of the task, every model call and every tool call
happened once.
+ assert model.calls == 3
+ assert calls == ["k0", "k1"]
+ assert journals[-1].stats.replayed == {"model": 2, "tool": 1}
+
+ def test_two_calls_in_one_answer_are_asked_one_at_a_time(self, reviewer):
+ storage, calls, journals, model = MemoryStorage(), [], [],
_ParallelModel()
+ reviewer.decide = lambda request: (
+ _answer("Approve") if '"key": "a"' in request.body else
_answer("Reject", "Not b.")
+ )
+ agent = Agent(name="analyst", model=model, tools=[_lookup_tool(calls)])
+
+ state, _ = reviewer.run(_agent_task(storage, agent, journals))
+
+ assert state == TaskInstanceState.SUCCESS
+ assert [key.split(":")[1] for key in reviewer.keys] == ["call-a",
"call-b"]
+ assert calls == ["a"]
+ assert reviewer.returned() == ["answer: ['Not b.', 'a=42']"]
+ assert model.calls == 2
+
+ def test_a_tool_of_an_agent_handed_off_to_pauses_the_same_way(self,
reviewer):
+ storage, calls, journals = MemoryStorage(), [], []
+ # The transfer's own result is the first function output the
specialist sees.
+ specialist = Agent(name="specialist", model=_ScriptedModel(lookups=2),
tools=[_lookup_tool(calls)])
+ lead = Agent(name="lead", model=_TransferModel(),
handoffs=[specialist])
+
+ state, _ = reviewer.run(_agent_task(storage, lead, journals))
+
+ assert state == TaskInstanceState.SUCCESS
+ assert reviewer.returned() == ["answer: k1=42"]
+ assert [request.subject for request in reviewer.requests] ==
["specialist wants to call lookup"]
+ assert calls == ["k1"]
+
+ def test_a_retry_approves_without_asking_a_call_that_already_ran(self,
reviewer):
+ storage, calls, journals = MemoryStorage(), [], []
+ lookup = _lookup_tool(calls)
+ failing = Agent(name="analyst",
model=_ScriptedModel(fail_on_answer=True), tools=[lookup])
+ state, error = reviewer.run(_agent_task(storage, failing, journals))
+ assert state == TaskInstanceState.FAILED
+ assert isinstance(error, RuntimeError)
+ assert len(reviewer.requests) == 1
+
+ retry_model = _ScriptedModel()
+ retry = Agent(name="analyst", model=retry_model, tools=[lookup])
+ state, _ = reviewer.run(_agent_task(storage, retry, journals))
+
+ assert state == TaskInstanceState.SUCCESS
+ assert reviewer.returned() == ["answer: k0=42"]
+ assert len(reviewer.requests) == 1
+ assert calls == ["k0"]
+ assert retry_model.calls == 1
+
+ def test_a_run_with_a_session_fails_instead_of_pausing(self, reviewer):
+ agent = Agent(name="analyst", model=_ScriptedModel(),
tools=[_lookup_tool([])])
+
+ state, error = reviewer.run(
+ _agent_task(MemoryStorage(), agent, [],
session=SQLiteSession("approval"))
+ )
+
+ assert state == TaskInstanceState.FAILED
+ assert isinstance(error, UnsupportedToolDeferralError)
+ assert "a run with a session" in str(error)
+ assert reviewer.requests == []
+
+ @pytest.mark.parametrize("fresh_call_ids", [True, False],
ids=["fresh-ids", "reused-ids"])
+ @pytest.mark.parametrize(
+ "decisions",
+ [("Approve",), ("Approve", "Approve"), ("Reject", "Approve")],
+ ids=["one-approval", "two-approvals", "reject-then-approve"],
+ )
+ def
test_a_retry_after_the_model_failed_pauses_and_resumes_on_its_own_answers(
+ self, reviewer, fresh_call_ids, decisions
+ ):
+ """A failed try's entries must neither replay into the retry nor make
its resumes ask the model again."""
+ storage, calls, journals = MemoryStorage(), [], []
+ model = _ScriptedModel(lookups=len(decisions),
fresh_call_ids=fresh_call_ids)
+ # The first try runs out of turns after its first answer, which fails
it because of the model.
+ unapproved = Agent(name="analyst", model=model,
tools=[_lookup_tool(calls, needs_approval=False)])
+ state, error = reviewer.run(_agent_task(storage, unapproved, journals,
max_turns=1))
+ assert (state, type(error)) == (TaskInstanceState.FAILED,
MaxTurnsExceeded)
+ assert calls == ["k0"]
+
+ reviewer.decide = lambda request:
_answer(decisions[len(reviewer.requests) - 1], "No.")
+ approved = Agent(name="analyst", model=model,
tools=[_lookup_tool(calls)])
+ state, error = reviewer.run(_agent_task(storage, approved, journals))
+
+ assert (state, error) == (TaskInstanceState.SUCCESS, None)
+ assert len(reviewer.requests) == len(decisions)
+ # The retry ran each approved call once, its answer to the first try's
call included.
+ assert calls == ["k0"] + [f"k{i}" for i, decision in
enumerate(decisions) if decision == "Approve"]
+ # The retry asked the model once per call and once for the answer;
every resume replayed.
+ assert model.calls == 1 + len(decisions) + 1
+ replayed_tools = sum(decision == "Approve" for decision in
decisions[:-1])
+ assert journals[-1].stats.replayed == (
+ {"model": len(decisions), "tool": replayed_tools} if
replayed_tools else {"model": len(decisions)}
+ )
+
+ def test_a_resume_that_does_not_come_back_to_the_answered_call_fails(self,
reviewer):
+ """An answer must not apply to a call the reviewer never saw, even one
reusing its id."""
+ storage, calls, journals = MemoryStorage(), [], []
+
+ def agent_for_run(run: int) -> Agent[Any]:
+ # From the resume on, the request differs, so the model is asked
again and answers
+ # with a call that reuses the paused call's id but not its
arguments.
+ return Agent(
+ name="analyst",
+ model=_ScriptedModel(argument="key" if run == 0 else "order"),
+ instructions=None if run == 0 else "Changed while the task
waited.",
+ tools=[_lookup_tool(calls)],
+ )
+
+ state, error = reviewer.run(_agent_task(storage, agent_for_run,
journals))
+
+ assert state == TaskInstanceState.FAILED
+ assert isinstance(error, DurableJournalError)
+ assert len(reviewer.requests) == 1
+ assert calls == []
+
+ def test_a_resume_that_finishes_without_the_answered_call_fails(self,
reviewer):
+ """Finishing would leave the answer unread and the journal deleted,
and a retry would start over."""
+ storage, calls, journals = MemoryStorage(), [], []
+
+ def agent_for_run(run: int) -> Agent[Any]:
+ # From the resume on, the request differs and the model asks for a
lookup that needs
+ # no approval, so the run finishes without asking about the call
that was answered.
+ return Agent(
+ name="analyst",
+ model=_ScriptedModel(value_prefix="k" if run == 0 else "x"),
+ instructions=None if run == 0 else "Changed while the task
waited.",
+ tools=[
+ _lookup_tool(calls, lambda context, arguments, call_id:
arguments["key"].startswith("k"))
+ ],
+ )
+
+ state, error = reviewer.run(_agent_task(storage, agent_for_run,
journals))
+
+ assert state == TaskInstanceState.FAILED
+ assert isinstance(error, DurableJournalError)
+ assert len(reviewer.requests) == 1
+ assert calls == ["x0"]
+ # Kept for a clear, which replays the finished run instead of starting
it over.
+ assert storage.steps() != []
+
+ def test_an_error_asking_for_approval_fails_the_task_with_that_error(self,
reviewer):
+ """Only an answer the run did not come back to fails it for good; any
other error is retried."""
+ calls: list[str] = []
+ failing = mock.create_autospec(
+ request_input, side_effect=RuntimeError("the API server dropped
the connection")
+ )
+ agent = Agent(name="analyst", model=_ScriptedModel(),
tools=[_lookup_tool(calls)])
+
+ with mock.patch(
+
"airflow.providers.common.ai.durable.openai_agents._get_request_input",
+ autospec=True,
+ return_value=failing,
+ ):
+ state, error = reviewer.run(_agent_task(MemoryStorage(), agent,
[]))
+
+ assert state == TaskInstanceState.FAILED
+ assert type(error) is RuntimeError
+ assert str(error) == "the API server dropped the connection"
+ assert calls == []
+
+ def test_a_task_that_runs_another_agent_first_resumes_past_it(self,
reviewer):
+ """The answer belongs to the second agent's run, so the first one
finishing does not claim it."""
+ storage, calls = MemoryStorage(), []
+ researcher = Agent(name="researcher", model=_ScriptedModel(),
tools=[_lookup_tool(calls, False)])
+ analyst = Agent(name="analyst",
model=_ScriptedModel(value_prefix="x"), tools=[_lookup_tool(calls)])
+ with DAG(dag_id="two_agents"):
+
+ @task
+ def research_then_analyse() -> str:
+ with journal_scope(DurableJournal(storage,
clean_up_after_run=True)):
+ run_durably_sync(researcher, "go", run_config=NO_TRACING)
+ return run_durably_sync(analyst, "go",
run_config=NO_TRACING).final_output
+
+ op = research_then_analyse().operator
+
+ state, error = reviewer.run(op)
+
+ assert (state, error) == (TaskInstanceState.SUCCESS, None)
+ assert len(reviewer.requests) == 1
+ # The first agent's record is deleted when it finishes, so it runs
again on the resume.
+ assert calls == ["k0", "k0", "x0"]
+
+ def test_an_operator_that_cannot_resume_fails_instead_of_pausing(self,
reviewer):
+ class _TicketOperator(BaseOperator):
+ def execute(self, context):
+ agent = Agent(name="analyst", model=_ScriptedModel(),
tools=[_lookup_tool([])])
+ with journal_scope(DurableJournal(MemoryStorage(),
clean_up_after_run=True)):
+ return run_durably_sync(agent, "go",
run_config=NO_TRACING).final_output
+
+ with DAG(dag_id="agent_in_an_operator"):
+ op = _TicketOperator(task_id="handle_ticket")
+
+ state, error = reviewer.run(op)
+
+ assert state == TaskInstanceState.FAILED
+ assert isinstance(error, UnsupportedToolDeferralError)
+ assert "_TicketOperator" in str(error)
+ assert reviewer.requests == []
+
+ def
test_a_tool_of_an_agent_reached_through_handoff_fails_instead_of_pausing(self,
reviewer):
+ specialist = Agent(name="specialist", model=_ScriptedModel(lookups=2),
tools=[_lookup_tool([])])
+ lead = Agent(name="lead", model=_TransferModel(),
handoffs=[handoff(specialist)])
+
+ state, error = reviewer.run(_agent_task(MemoryStorage(), lead, []))
+
+ assert state == TaskInstanceState.FAILED
+ assert isinstance(error, UnsupportedToolDeferralError)
+ assert "reached through handoff()" in str(error)
+ assert reviewer.requests == []
+
+ @pytest.mark.enable_redact
+ def test_a_secret_in_the_arguments_is_masked_for_the_reviewer(self,
reviewer, registered_secret):
+ calls: list[str] = []
+ model = _ScriptedModel(value_prefix=registered_secret)
+ agent = Agent(name="analyst", model=model, tools=[_lookup_tool(calls)])
+
+ state, _ = reviewer.run(_agent_task(MemoryStorage(), agent, []))
+
+ assert state == TaskInstanceState.SUCCESS
+ [request] = reviewer.requests
+ assert registered_secret not in request.body
+ assert "***" in request.body
+ assert len(calls) == 1
