This is an automated email from the ASF dual-hosted git repository. kaxil pushed a commit to branch openai-agents-tool-approval in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 30af5e954ca2366af27a9e4fb622d6c3f41beb1c Author: Kaxil Naik <[email protected]> AuthorDate: Thu Nov 5 23:36:31 2026 +0000 Pause an OpenAI Agents SDK durable run for tool approval in Required Actions A function tool with needs_approval=True now pauses the task in Required Actions instead of failing it. run_durably asks each pending call through the standard provider's request_input(), one at a time in the model's order, with the agent, the tool and its masked arguments in the request, and Approve or Reject with an optional reason. The answer resumes the task: the run replays the calls it recorded before the pause, and the call is approved, or rejected with the reviewer's reason as its result. Steps are deleted only once a run ends with no pending approval. A retry approves a call whose result is already recorded without asking again. A pause moves the replay boundary an earlier model failure left, so the resume replays this try's answers instead of asking the model again. The answer's key includes the call's fingerprint, and a resume that does not come back to the answered call fails with DurableJournalError instead of looping. The task still fails with UnsupportedToolDeferralError where it cannot pause: before Airflow 3.4 or without a standard provider that has request_input(), outside a @task or PythonOperator, with a session, or for a call by an agent whose model is not recorded (as_tool, handoff()). --- .../common/ai/docs/frameworks/openai_agents.rst | 93 ++++- .../airflow/providers/common/ai/durable/base.py | 7 +- .../airflow/providers/common/ai/durable/journal.py | 81 ++++- .../providers/common/ai/durable/openai_agents.py | 256 ++++++++++++-- .../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 | 389 ++++++++++++++++++++- 8 files changed, 940 insertions(+), 66 deletions(-) diff --git a/providers/common/ai/docs/frameworks/openai_agents.rst b/providers/common/ai/docs/frameworks/openai_agents.rst index 5a5ce08600a..ade155b4cab 100644 --- a/providers/common/ai/docs/frameworks/openai_agents.rst +++ b/providers/common/ai/docs/frameworks/openai_agents.rst @@ -253,8 +253,9 @@ read, and building it fails: AttributeError: 'FilteredToolset' object has no attribute 'airflow_tools' Choose the tools with ``allowed_tables``, ``allowed_methods`` and ``pinned_arguments`` -instead. There is no per-call approval on this path; see -:ref:`howto/frameworks:openai-agents-limits`. Each toolset's page covers the rest of its +instead. ``AirflowTools`` has no approval option of its own; to have a person approve +each call, use a function tool with ``needs_approval``, as +:ref:`howto/frameworks:openai-agents-approval` shows. Each toolset's page covers the rest of its options: :ref:`howto/toolsets`. Tools @@ -571,13 +572,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 ^^^^^^^^^^^^^^^^^^^^ @@ -623,6 +619,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 @@ -631,11 +694,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 2d437834e4d..c36abaad406 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 @@ -209,6 +212,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. @@ -230,6 +235,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] = {} @@ -247,6 +255,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: """ @@ -279,6 +288,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: @@ -326,15 +340,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 @@ -427,10 +451,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) @@ -463,7 +521,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 8edeec898ff..4725aa050d1 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,20 +64,34 @@ 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.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"] @@ -106,8 +120,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( @@ -157,6 +172,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. @@ -167,7 +188,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) @@ -182,24 +205,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})) - error = 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(error, model_at_fault=True) - raise error + 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 @@ -216,11 +240,182 @@ 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)}" + try: + 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, + ) + except RuntimeError as error: + # Raised by request_input() when the answer the task resumed with is for a call that + # this run never reached; asking again would never end. + _did_not_come_back(durable_run, error) + 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]) -> bool: + """Whether the task resumed with an answer to an approval of this run that the run never asked again for.""" + 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 + ) + + +def _did_not_come_back(durable_run: DurableRun, cause: BaseException | None = None) -> NoReturn: + """ + Fail the task resumed with an answer whose call the run never reached again. + + ``cause`` is set when the run stopped at another call that needs approval (a clear asks + about that one), and ``None`` when it finished without one (a clear replays the result). + """ + then = ( + "Clear the task to ask about the call it stopped at instead." + if cause is not None + 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 from cause + + +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" + ) + task = get_current_context()["task"] + 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): @@ -282,11 +477,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 02984e1144b..cbba0316b74 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 @@ -25,6 +25,7 @@ from typing import Any from unittest import mock import pytest +import structlog pytest.importorskip("agents") @@ -60,9 +61,14 @@ 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, ToolResult from airflow.providers.common.ai.tools.openai_agents import AirflowTools +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 @@ -87,9 +93,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 @@ -113,9 +123,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", ) ] @@ -1000,7 +1012,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) @@ -1010,8 +1022,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 @@ -1045,10 +1065,367 @@ 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_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
