You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ka...@apache.org on 2026/10/14 22:03:54 UTC
(airflow) branch openai-agents-durable created (now ccb390c12d9)
This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a change to branch openai-agents-durable
in repository https://gitbox.apache.org/repos/asf/airflow.git
at ccb390c12d9 Replay an OpenAI Agents SDK agent's completed calls when Airflow retries its task
This branch includes the following new commits:
new ccb390c12d9 Replay an OpenAI Agents SDK agent's completed calls when Airflow retries its task
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(airflow) 01/01: Replay an OpenAI Agents SDK agent's completed calls when Airflow retries its task
Posted by ka...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch openai-agents-durable
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit ccb390c12d9ea8b6096295fa4786bcd6fae48468
Author: Kaxil Naik <ka...@apache.org>
AuthorDate: Wed Oct 14 21:48:44 2026 +0100
Replay an OpenAI Agents SDK agent's completed calls when Airflow retries its task
run_durably() and run_durably_sync() run an agent like Runner.run, recording
each model call and function tool call in the task's durable journal, the one
AgentOperator(durable=True) uses. A retry replays every call that asks what the
failed try asked, in order, and runs live from the first that differs. A run
the model failed (max turns, malformed output, output guardrail) asks the
model again from its last request; an agent called through as_tool replays as
one call; server-side conversations are refused because a replay never
reaches them.
---
providers/common/ai/docs/durable_execution.rst | 3 +-
providers/common/ai/docs/features.rst | 3 +-
.../common/ai/docs/frameworks/openai_agents.rst | 58 ++-
providers/common/ai/docs/stability.rst | 3 +-
.../providers/common/ai/durable/openai_agents.py | 448 +++++++++++++++++++++
.../example_dags/example_openai_agents_durable.py | 90 +++++
.../unit/common/ai/durable/test_openai_agents.py | 417 +++++++++++++++++++
7 files changed, 1017 insertions(+), 5 deletions(-)
diff --git a/providers/common/ai/docs/durable_execution.rst b/providers/common/ai/docs/durable_execution.rst
index 5c670045998..bf5476649ad 100644
--- a/providers/common/ai/docs/durable_execution.rst
+++ b/providers/common/ai/docs/durable_execution.rst
@@ -30,7 +30,8 @@ them again. Only the steps the failed attempt never finished run live.
Durable execution only helps a task that has retries. To decide *whether* a failure is
worth retrying, see :doc:`retry_policies`; a retried ``LLMBatchOperator`` re-attaches to
its running batch instead (:ref:`llm-batch-reattach`). To resume a Strands Agents agent on a
-retry, see :ref:`howto/frameworks:strands-durable`.
+retry, see :ref:`howto/frameworks:strands-durable`; for an OpenAI Agents SDK agent, see
+:ref:`howto/frameworks:openai-agents-durable`.
Turn it on
----------
diff --git a/providers/common/ai/docs/features.rst b/providers/common/ai/docs/features.rst
index 50b3dd6e5eb..f2420d51381 100644
--- a/providers/common/ai/docs/features.rst
+++ b/providers/common/ai/docs/features.rst
@@ -268,7 +268,8 @@ toolset except
- Gives a loop on the Anthropic SDK's ``tool_runner`` the toolsets.
- ``anthropic``
* - :doc:`OpenAI Agents SDK <frameworks/openai_agents>`
- - Gives an OpenAI Agents SDK agent the toolsets.
+ - Gives an OpenAI Agents SDK agent the toolsets. Run through ``run_durably``, a retry
+ replays its completed calls.
- None; install ``openai-agents``
* - :doc:`LangChain <hooks/langchain>`
- Builds LangChain chat and embedding models from a connection, and
diff --git a/providers/common/ai/docs/frameworks/openai_agents.rst b/providers/common/ai/docs/frameworks/openai_agents.rst
index 86d87cabeb4..763e0df9289 100644
--- a/providers/common/ai/docs/frameworks/openai_agents.rst
+++ b/providers/common/ai/docs/frameworks/openai_agents.rst
@@ -145,11 +145,65 @@ To turn the SDK's tracing off altogether, call ``agents.set_tracing_disabled(Tru
``OPENAI_AGENTS_DISABLE_TRACING=1`` in the worker's environment. That stops
``AirflowTraceProcessor`` from receiving anything too.
+.. _howto/frameworks:openai-agents-durable:
+
+Durable execution
+-----------------
+
+When a task that calls ``Runner.run`` fails and Airflow retries it, the agent starts
+again from the prompt: every model call is paid for again, and every tool call happens
+again. On a task with ``retries``, run the agent with
+:func:`~airflow.providers.common.ai.durable.openai_agents.run_durably_sync`, or
+:func:`~airflow.providers.common.ai.durable.openai_agents.run_durably` from async code,
+instead. They take the same arguments as ``Runner.run_sync`` and ``Runner.run``:
+
+.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_openai_agents_durable.py
+ :language: python
+ :start-after: [START example_openai_agents_durable]
+ :end-before: [END example_openai_agents_durable]
+
+Each model call and function tool call is recorded in the task instance's durable
+journal as it completes, the journal ``AgentOperator(durable=True)`` uses (see
+:doc:`../durable_execution`). On a retry the agent runs again from the start, and each
+call that asks what the failed try asked, in the same order, returns what it returned
+then. From the first call that differs, that call and every call after it run live. Once
+the run succeeds, what was recorded is deleted, so clearing the task later starts over.
+
+A retry runs some work again:
+
+- A call that was running when the try failed runs again, so tools that change data
+ must still be safe to repeat.
+- When the run failed because of the model's answers (``MaxTurnsExceeded``,
+ ``ModelBehaviorError``, ``ModelRefusalError``, or an output or tool input guardrail),
+ replaying those answers would fail it the same way, so the retry asks the model again
+ from its last request. A ``ModelBehaviorError`` that a tool raises, such as arguments
+ that fail validation with ``failure_error_function=None``, counts as the tool's failure
+ instead: the tool call runs again, but the answer that asked for it replays.
+- A retry whose prompt, instructions, model settings or tools differ from the failed
+ try's runs everything live from the first call that sees the difference.
+- These are not recorded and run live on every try: tools other than function tools
+ (hosted tools such as web search, the tools of ``mcp_servers``, shell and computer
+ tools); a function tool whose ``needs_approval`` is a function, since the SDK checks the
+ arguments it approved only through its own invoker; the function tools of an agent
+ reached through ``handoff()``, and its model unless that is given by name or left to
+ the default; and a streamed run. An agent used through ``as_tool`` is recorded as one
+ call.
+- 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 ``SandboxToolset`` without ``attach_to`` is destroyed when the try ends, so the retry
+ starts with a new, empty sandbox that the replayed conversation does not describe.
+
+Recorded tool results pass through Airflow's secret masker before they are saved, and a
+replayed model call reports no token usage. ``run_durably`` refuses
+``previous_response_id``, ``conversation_id`` and ``auto_previous_response_id``: with them
+OpenAI keeps the conversation, and a replayed call never reaches it. Outside a task there
+is nothing to replay into, and it runs the agent as ``Runner.run`` does.
+
Differences from ``AgentOperator``
----------------------------------
-- A task retry runs the agent again from the start, so tools that change data must be
- safe to repeat.
+- A task retry runs the agent again from the start unless it runs through
+ ``run_durably``; see :ref:`howto/frameworks:openai-agents-durable`.
- There is no human review step like ``AgentOperator(enable_hitl_review=True)``, and the
SDK's own ``needs_approval`` interruptions are not connected to Airflow's approval UI.
diff --git a/providers/common/ai/docs/stability.rst b/providers/common/ai/docs/stability.rst
index beefd7aa1eb..9298e171835 100644
--- a/providers/common/ai/docs/stability.rst
+++ b/providers/common/ai/docs/stability.rst
@@ -260,7 +260,8 @@ Everything this provider ships that is not in the table above is experimental.
* - The Anthropic tool runner adapter (:doc:`frameworks/anthropic`)
- Built on ``tool_runner``, a beta helper of the Anthropic SDK that can change between
SDK releases.
- * - The OpenAI Agents SDK adapter and trace processor (:doc:`frameworks/openai_agents`)
+ * - The OpenAI Agents SDK adapter, trace processor and ``run_durably``
+ (:doc:`frameworks/openai_agents`)
- Written against OpenAI Agents SDK 0.20 and 0.23. The SDK is pre-1.0 and changes
its tool and runner APIs between minor releases.
* - The ``common_ai.tool_calls`` metric 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
new file mode 100644
index 00000000000..bc0af085627
--- /dev/null
+++ b/providers/common/ai/src/airflow/providers/common/ai/durable/openai_agents.py
@@ -0,0 +1,448 @@
+# 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.
+"""
+Replay an OpenAI Agents SDK agent's completed model and tool calls when Airflow retries its task.
+
+.. note:: Experimental; see :ref:`howto/stability`.
+"""
+
+from __future__ import annotations
+
+import copy
+import json
+from typing import TYPE_CHECKING, Any
+
+from pydantic import TypeAdapter, ValidationError
+from pydantic_core import to_jsonable_python
+
+from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+try:
+ from agents import (
+ Agent,
+ FunctionTool,
+ Model,
+ ModelProvider,
+ ModelResponse,
+ RunConfig,
+ Runner,
+ ToolOutputFileContent,
+ ToolOutputImage,
+ ToolOutputText,
+ Usage,
+ )
+ from agents.exceptions import (
+ MaxTurnsExceeded,
+ ModelBehaviorError,
+ ModelRefusalError,
+ OutputGuardrailTripwireTriggered,
+ ToolInputGuardrailTripwireTriggered,
+ )
+ from agents.items import TResponseOutputItem
+except ImportError as e:
+ raise AirflowOptionalProviderFeatureException(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.utils.coroutines import run_coroutine_sync
+from airflow.providers.common.ai.utils.masking import mask_secrets
+from airflow.providers.common.ai.utils.task_logger import get_task_logger
+
+if TYPE_CHECKING:
+ from collections.abc import AsyncIterator
+
+ from agents import AgentOutputSchemaBase, Handoff, ModelSettings, ModelTracing, RunResult, Tool
+ 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
+
+__all__ = ["run_durably", "run_durably_sync"]
+
+log = get_task_logger()
+
+_OUTPUT_ITEMS: TypeAdapter[list[TResponseOutputItem]] = TypeAdapter(list[TResponseOutputItem])
+
+# Errors that come from the model's answers rather than from a step. The answers were
+# recorded, so a retry must ask the model again instead of replaying them into the same error.
+_MODEL_AT_FAULT = (
+ ModelBehaviorError,
+ ModelRefusalError,
+ MaxTurnsExceeded,
+ OutputGuardrailTripwireTriggered,
+ ToolInputGuardrailTripwireTriggered,
+)
+
+# The tool results the SDK hands the model as structured content rather than as text.
+_STRUCTURED_OUTPUTS = (ToolOutputText, ToolOutputImage, ToolOutputFileContent)
+
+# Arguments that keep the conversation on OpenAI's servers, which a replayed model call never reaches.
+_SERVER_SIDE_STATE = ("previous_response_id", "conversation_id", "auto_previous_response_id")
+
+
+async def run_durably(
+ starting_agent: Agent[Any],
+ input: str | list[TResponseInputItem],
+ *,
+ run_config: RunConfig | None = None,
+ **kwargs: Any,
+) -> RunResult:
+ """
+ Run an agent like ``Runner.run``, so that a retry of the task replays what already completed.
+
+ Each model call and each function tool call the run makes is recorded in the task's
+ durable journal as it completes. When the task fails and Airflow retries it, the agent
+ runs again from the start, and every call that asks what the previous attempt asked,
+ in the same order, returns the recorded result instead of calling the model or the tool
+ again. From the first call that differs, such as a model answer that changed, that call
+ and every call after it run live. Once the run succeeds what was recorded is deleted, so
+ clearing the task later starts the agent over.
+
+ .. code-block:: python
+
+ from agents import Agent
+
+ from airflow.providers.common.ai.durable.openai_agents import run_durably_sync
+ from airflow.providers.common.ai.tools.openai_agents import AirflowTools
+ from airflow.sdk import task
+
+
+ @task(retries=3)
+ def research(question: str) -> str:
+ agent = Agent(name="analyst", model=model, tools=AirflowTools(warehouse))
+ return run_durably_sync(agent, question).final_output
+
+ Recorded are the calls of every model given by name or in ``run_config``, the model and
+ function tools of the starting agent and of the agents in its ``handoffs`` list, and an
+ agent used through ``as_tool``, as one call. These run live on every attempt: tools
+ other than function tools (hosted tools such as web search, the tools of
+ ``mcp_servers``, shell and computer tools), a function tool whose ``needs_approval`` is
+ a function, since the SDK checks its arguments only through its own invoker, the
+ function tools of an agent reached through ``handoff()`` and its model unless that is
+ given by name or left to the default, and streamed runs. A replayed model call reports
+ no usage, since it costs nothing.
+
+ 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.
+
+ Outside a task there is no journal, and this is ``Runner.run``.
+
+ :param starting_agent: The agent to run.
+ :param input: The prompt, as for ``Runner.run``.
+ :param run_config: As for ``Runner.run``.
+ :param kwargs: Passed to ``Runner.run``, such as ``max_turns`` or ``context``.
+ :raises ValueError: When ``previous_response_id``, ``conversation_id`` or
+ ``auto_previous_response_id`` is passed: a replayed call never reaches the
+ conversation OpenAI keeps for them.
+ """
+ journal = current_journal()
+ if journal is None:
+ return await Runner.run(starting_agent, input, run_config=run_config, **kwargs)
+ if stateful := [name for name in _SERVER_SIDE_STATE if kwargs.get(name)]:
+ raise ValueError(
+ f"run_durably() cannot replay a run that keeps its conversation on OpenAI's servers "
+ f"({', '.join(stateful)}): a replayed model call never reaches it."
+ )
+ durable_run = journal.start_run()
+ config = copy.copy(run_config) if run_config is not None else RunConfig()
+ config.model_provider = _DurableModelProvider(config.model_provider, durable_run)
+ if isinstance(config.model, Model):
+ config.model = _DurableModel(config.model, durable_run)
+ agent = _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 journal.clean_up_after_run:
+ durable_run.cleanup()
+ return result
+
+
+def run_durably_sync(
+ starting_agent: Agent[Any],
+ input: str | list[TResponseInputItem],
+ *,
+ run_config: RunConfig | None = None,
+ **kwargs: Any,
+) -> RunResult:
+ """Run :func:`run_durably` from synchronous code, such as a plain ``@task``."""
+ return run_coroutine_sync(run_durably(starting_agent, input, run_config=run_config, **kwargs))
+
+
+def _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.clone()
+ # Recorded before the handoffs are copied, so agents that hand off to each other copy once.
+ seen[id(agent)] = durable
+ if isinstance(agent.model, Model):
+ durable.model = _DurableModel(agent.model, durable_run)
+ durable.tools = [_durable_tool(tool, durable_run) for tool in agent.tools]
+ durable.handoffs = [
+ _durable_agent(target, durable_run, seen) if isinstance(target, Agent) else target
+ for target in agent.handoffs
+ ]
+ return durable
+
+
+def _durable_tool(tool: Tool, durable_run: DurableRun) -> Tool:
+ # A tool whose approval depends on its arguments is left as it is: the SDK checks that the
+ # arguments it approved are the ones the tool runs with only through its own invoker.
+ if not isinstance(tool, FunctionTool) or callable(tool.needs_approval):
+ return tool
+ invoke = tool.on_invoke_tool
+
+ async def on_invoke_tool(ctx: ToolContext[Any], arguments: str) -> Any:
+ if _inside_a_call(durable_run):
+ return await invoke(ctx, arguments)
+ try:
+ parsed = json.loads(arguments) if arguments else {}
+ except json.JSONDecodeError:
+ parsed = arguments
+ fingerprint = fingerprint_tool_call(
+ tool.name,
+ # As recorded model answers are masked, so a replayed call's arguments can be too.
+ mask_secrets(parsed if isinstance(parsed, dict) else {"arguments": parsed}),
+ ctx.tool_call_id,
+ )
+ return await durable_run.run(
+ f"tool:{tool.name}",
+ kind="tool",
+ fingerprint=fingerprint,
+ body=lambda: invoke(ctx, arguments),
+ # Secrets must not reach the journal any more than the model.
+ to_record=lambda output: mask_secrets(_as_sent(output)),
+ )
+
+ # A shallow copy, not dataclasses.replace: the tool's schema is already prepared.
+ durable = copy.copy(tool)
+ durable.on_invoke_tool = on_invoke_tool
+ return durable
+
+
+def _as_sent(output: Any) -> Any:
+ """
+ Return a tool's result in the form the SDK sends it to the model, as a JSON-compatible value.
+
+ The SDK sends a ``ToolOutputText``, ``ToolOutputImage`` or ``ToolOutputFileContent``, a
+ ``dict`` with a ``type``, or a list of them, as structured content, and anything else as
+ ``str(output)``. Recording anything else as it is would replay it as a different text
+ (a model's fields as a dict, a tuple as a list), so the next model call would no longer
+ match the one the failed try made.
+ """
+ if isinstance(output, str):
+ return output
+ items = output if isinstance(output, list | tuple) else [output]
+ if items and all(
+ isinstance(item, _STRUCTURED_OUTPUTS) or (isinstance(item, dict) and "type" in item) for item in items
+ ):
+ return to_jsonable_python(output)
+ return str(output)
+
+
+def _inside_a_call(durable_run: DurableRun) -> bool:
+ """
+ Whether a recorded call of ``durable_run`` is running, such as an agent called through ``as_tool``.
+
+ That call is recorded as one step, and a retry replays it whole, so what it does inside
+ runs live rather than taking positions of the run it belongs to.
+ """
+ step = _EXECUTING_STEP.get()
+ return step is not None and step.durable_run is durable_run
+
+
+class _DurableModelProvider(ModelProvider):
+ """Wrap every model the run resolves by name in a :class:`_DurableModel`."""
+
+ def __init__(self, provider: ModelProvider, durable_run: DurableRun) -> None:
+ self._provider = provider
+ self._durable_run = durable_run
+ self._models: dict[str | None, _DurableModel] = {}
+
+ def get_model(self, model_name: str | None) -> Model:
+ if model_name not in self._models:
+ self._models[model_name] = _DurableModel(self._provider.get_model(model_name), self._durable_run)
+ return self._models[model_name]
+
+
+class _DurableModel(Model):
+ """Replay a model call the previous attempt recorded, or make it and record the response."""
+
+ def __init__(self, model: Model, durable_run: DurableRun) -> None:
+ self._model = model
+ self._durable_run = durable_run
+ # The step of a call that raised, which the SDK's own model retry may make again.
+ self._failed: JournalStep | None = None
+
+ async def get_response(
+ self,
+ system_instructions: str | None,
+ input: str | list[TResponseInputItem],
+ model_settings: ModelSettings,
+ tools: list[Tool],
+ output_schema: AgentOutputSchemaBase | None,
+ handoffs: list[Handoff],
+ tracing: ModelTracing,
+ **kwargs: Any,
+ ) -> ModelResponse:
+ if _inside_a_call(self._durable_run):
+ return await self._model.get_response(
+ system_instructions, input, model_settings, tools, output_schema, handoffs, tracing, **kwargs
+ )
+ fingerprint = _fingerprint_request(
+ self._model, system_instructions, input, model_settings, tools, output_schema, handoffs, kwargs
+ )
+ step = self._retried(fingerprint) or self._durable_run.claim(
+ "model", kind="model", fingerprint=fingerprint
+ )
+ if step.replayed:
+ try:
+ return ModelResponse(
+ output=_OUTPUT_ITEMS.validate_python(step.payload["output"]),
+ usage=Usage(),
+ response_id=step.payload.get("response_id"),
+ )
+ except (ValidationError, KeyError, TypeError) as error:
+ self._durable_run.reject(step, reason=f"the recorded model response cannot be read: {error}")
+
+ async def call() -> ModelResponse:
+ return await self._model.get_response(
+ system_instructions, input, model_settings, tools, output_schema, handoffs, tracing, **kwargs
+ )
+
+ try:
+ return await step.run(call, to_record=_record_response)
+ except Exception:
+ self._failed = step
+ raise
+
+ def _retried(self, fingerprint: str | None) -> JournalStep | None:
+ """
+ Return the step of the call that just raised, when this call is the SDK retrying it.
+
+ The SDK's model retry asks again with the same request and nothing in between; it is
+ the same call, so it takes the same position, and the next try lines up with this one.
+ """
+ failed, self._failed = self._failed, None
+ if failed is None or failed.fingerprint != fingerprint:
+ return None
+ return failed if self._durable_run.position == failed.position + 1 else None
+
+ def stream_response(
+ self,
+ system_instructions: str | None,
+ input: str | list[TResponseInputItem],
+ model_settings: ModelSettings,
+ tools: list[Tool],
+ output_schema: AgentOutputSchemaBase | None,
+ handoffs: list[Handoff],
+ tracing: ModelTracing,
+ **kwargs: Any,
+ ) -> AsyncIterator[TResponseStreamEvent]:
+ # A streamed response is not recorded: it reaches the caller event by event.
+ return self._model.stream_response(
+ system_instructions, input, model_settings, tools, output_schema, handoffs, tracing, **kwargs
+ )
+
+ def _supports_default_prompt_cache_key(self) -> bool:
+ # The runner asks the model whether to add a prompt cache key; the answer is the wrapped model's.
+ supports = getattr(self._model, "_supports_default_prompt_cache_key", None)
+ return bool(supports()) if callable(supports) else False
+
+ def get_retry_advice(self, request: ModelRetryAdviceRequest) -> ModelRetryAdvice | None:
+ return self._model.get_retry_advice(request)
+
+ async def _cleanup_on_run_end(self, owner: object) -> None:
+ await self._model._cleanup_on_run_end(owner)
+
+ async def close(self) -> None:
+ await self._model.close()
+
+
+def _record_response(response: ModelResponse) -> dict[str, Any]:
+ # Only the fields the response set, so a replayed item turns into the same next request as
+ # the live one: the SDK carries every field of an item into the conversation it sends.
+ output = [item.model_dump(mode="json", exclude_unset=True) for item in response.output]
+ # Masked, as the journal's other steps are; a replay hands the model its own masked answer.
+ return {"output": mask_secrets(output), "response_id": response.response_id}
+
+
+def _fingerprint_request(
+ model: Model,
+ system_instructions: str | None,
+ input: str | list[TResponseInputItem],
+ model_settings: ModelSettings,
+ tools: list[Tool],
+ output_schema: AgentOutputSchemaBase | None,
+ handoffs: list[Handoff],
+ options: dict[str, Any],
+) -> str | None:
+ """
+ Fingerprint what a model call asks: the model, instructions, conversation, settings and tools.
+
+ ``None`` when the request cannot be serialized; the call then replays on position alone,
+ and only right after a call that replayed.
+ """
+ try:
+ settings = to_jsonable_python(model_settings)
+ # Generated once per run, so it differs on every attempt without changing the request.
+ if isinstance(extra := settings.get("extra_args"), dict):
+ extra.pop("prompt_cache_key", None)
+ payload = {
+ "model": f"{type(model).__module__}.{type(model).__qualname__}:{getattr(model, 'model', None)}",
+ "instructions": system_instructions,
+ # Masked as the recorded steps are, so a replayed conversation matches a live one.
+ "input": _without_none(mask_secrets(to_jsonable_python(input))),
+ "settings": settings,
+ "tools": [_describe_tool(tool) for tool in tools],
+ "output": output_schema.json_schema()
+ if output_schema and not output_schema.is_plain_text()
+ else None,
+ "handoffs": [[handoff.tool_name, handoff.input_json_schema] for handoff in handoffs],
+ "prompt": to_jsonable_python(options.get("prompt")),
+ }
+ return _digest(payload)
+ except (TypeError, ValueError):
+ log.warning(
+ "Durable: could not fingerprint the model request; its recorded response replays "
+ "without verification"
+ )
+ return None
+
+
+def _without_none(value: Any) -> Any:
+ """
+ Drop ``None`` fields, so an item that leaves a field unset matches one that sets it to ``None``.
+
+ Which of the two the SDK sends depends on how the item was built, not on what was asked.
+ """
+ if isinstance(value, dict):
+ return {key: _without_none(item) for key, item in value.items() if item is not None}
+ if isinstance(value, list):
+ return [_without_none(item) for item in value]
+ return value
+
+
+def _describe_tool(tool: Tool) -> Any:
+ """Return what the model is told about a tool, for the request fingerprint."""
+ if isinstance(tool, FunctionTool):
+ return [tool.name, tool.description, tool.params_json_schema, tool.strict_json_schema]
+ # A hosted tool's configuration, such as a file search's vector stores, without its callbacks.
+ return [type(tool).__name__, to_jsonable_python(tool, fallback=lambda value: type(value).__name__)]
diff --git a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_openai_agents_durable.py b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_openai_agents_durable.py
new file mode 100644
index 00000000000..ca27e8da8ab
--- /dev/null
+++ b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_openai_agents_durable.py
@@ -0,0 +1,90 @@
+# 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.
+"""
+Run an OpenAI Agents SDK agent whose completed model and tool calls a task retry replays.
+
+The agent is the one in ``example_openai_agents_agent``: a model whose API key comes from
+an Airflow connection, and a ``SQLToolset``'s tools through
+:class:`~airflow.providers.common.ai.tools.openai_agents.AirflowTools`.
+:func:`~airflow.providers.common.ai.durable.openai_agents.run_durably_sync` runs it, so
+when the task fails and Airflow retries it, the calls the failed try completed return
+what they returned then instead of running again.
+
+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. Create a database connection (``DB_CONN_ID``, default ``sql_default``) whose
+ hook is a ``DbApiHook`` (e.g. SQLite, Postgres, MySQL).
+4. On Airflow before 3.3, set ``[common.ai] durable_cache_path`` to where the journal
+ is kept; from 3.3 on it is kept in the task state store.
+"""
+
+from __future__ import annotations
+
+import os
+from datetime import timedelta
+
+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")
+DB_CONN_ID = os.environ.get("DB_CONN_ID", "sql_default")
+
+DEFAULT_QUESTION = "Which tables exist, and how many rows does each contain?"
+
+
+# [START example_openai_agents_durable]
+@dag(tags=["example"])
+def example_openai_agents_durable():
+ """Answer a question about a database with an OpenAI Agents SDK agent that survives task retries."""
+
+ @task(retries=2, retry_delay=timedelta(minutes=1))
+ def run_durable_openai_agent(question: str = DEFAULT_QUESTION) -> str:
+ from agents import Agent, OpenAIChatCompletionsModel, RunConfig
+ from openai import AsyncOpenAI
+
+ from airflow.providers.common.ai.durable.openai_agents import run_durably_sync
+ from airflow.providers.common.ai.tools.openai_agents import AirflowTools
+ from airflow.providers.common.ai.toolsets.sql import SQLToolset
+
+ llm = BaseHook.get_connection(LLM_CONN_ID)
+ agent = Agent(
+ name="analyst",
+ model=OpenAIChatCompletionsModel(
+ model=LLM_MODEL,
+ openai_client=AsyncOpenAI(api_key=llm.password, base_url=llm.host or None),
+ ),
+ instructions=(
+ "You are a SQL analyst. Use list_tables and get_schema to explore "
+ "the database, then run read-only queries to answer the question."
+ ),
+ tools=AirflowTools(SQLToolset(db_conn_id=DB_CONN_ID)),
+ )
+ # The SDK sends traces to OpenAI by default; see the guide's tracing section.
+ result = run_durably_sync(agent, question, run_config=RunConfig(tracing_disabled=True))
+ return result.final_output
+
+ run_durable_openai_agent()
+
+
+# [END example_openai_agents_durable]
+
+
+example_openai_agents_durable()
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
new file mode 100644
index 00000000000..f1bffe35775
--- /dev/null
+++ b/providers/common/ai/tests/unit/common/ai/durable/test_openai_agents.py
@@ -0,0 +1,417 @@
+# 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.
+from __future__ import annotations
+
+import contextlib
+import datetime
+import json
+from typing import Any
+
+import pytest
+
+pytest.importorskip("agents")
+
+from agents import (
+ Agent,
+ Model,
+ ModelProvider,
+ ModelResponse,
+ ModelSettings,
+ RunConfig,
+ ToolOutputText,
+ Usage,
+ function_tool,
+)
+from agents.exceptions import MaxTurnsExceeded
+from agents.retry import ModelRetryBackoffSettings, ModelRetrySettings
+from openai.types.responses import ResponseFunctionToolCall, ResponseOutputMessage, ResponseOutputText
+from pydantic import BaseModel
+
+from airflow.providers.common.ai.durable.journal import DurableJournal, journal_scope
+from airflow.providers.common.ai.durable.openai_agents import _durable_tool, run_durably_sync
+from airflow.providers.common.ai.tools import AirflowTool, ToolResult
+from airflow.providers.common.ai.tools.openai_agents import AirflowTools
+
+from unit.common.ai.durable.memory_storage import MemoryStorage
+
+# Tracing off, so no test exports spans to OpenAI.
+NO_TRACING = RunConfig(tracing_disabled=True)
+
+
+class _ScriptedModel(Model):
+ """
+ Calls ``lookup`` until it has ``lookups`` results back, then answers with the last one.
+
+ ``calls`` counts the requests that reached the model. ``fail_on_answer`` makes the
+ request that would answer raise instead, as a crashed worker would leave it.
+ """
+
+ def __init__(
+ self,
+ *,
+ lookups: int = 1,
+ fail_on_answer: bool = False,
+ usage: int = 7,
+ argument: str = "key",
+ fail_first_request: bool = False,
+ ) -> None:
+ self.lookups = lookups
+ self.argument = argument
+ self.fail_first_request = fail_first_request
+ self.fail_on_answer = fail_on_answer
+ self.usage = usage
+ self.calls = 0
+ # How many tool results each request that reached the model carried.
+ self.tool_results_seen: list[int] = []
+
+ async def get_response(self, system_instructions, input, *args: Any, **kwargs: Any) -> ModelResponse:
+ self.calls += 1
+ if self.fail_first_request and self.calls == 1:
+ raise ConnectionError("connection reset")
+ outputs = [
+ item["output"]
+ for item in (input if isinstance(input, list) else [])
+ if isinstance(item, dict) and item.get("type") == "function_call_output"
+ ]
+ self.tool_results_seen.append(len(outputs))
+ output: list[ResponseFunctionToolCall | ResponseOutputMessage]
+ if len(outputs) < self.lookups:
+ output = [
+ ResponseFunctionToolCall(
+ id=f"fc-{len(outputs)}",
+ call_id=f"call-{len(outputs)}",
+ name="lookup",
+ arguments=json.dumps({self.argument: f"k{len(outputs)}"}),
+ type="function_call",
+ )
+ ]
+ else:
+ if self.fail_on_answer:
+ raise RuntimeError("worker lost")
+ text = ResponseOutputText(text=f"answer: {outputs[-1]}", type="output_text", annotations=[])
+ output = [
+ ResponseOutputMessage(
+ id="msg", role="assistant", status="completed", type="message", content=[text]
+ )
+ ]
+ usage = Usage(
+ requests=1, input_tokens=self.usage, output_tokens=self.usage, total_tokens=2 * self.usage
+ )
+ return ModelResponse(output=list(output), usage=usage, response_id=f"resp-{self.calls}")
+
+ def stream_response(self, *args: Any, **kwargs: Any):
+ raise NotImplementedError
+
+
+class _CountingTool:
+ """A ``lookup`` tool that counts the calls that reached it."""
+
+ def __init__(self, result: str = "42") -> None:
+ self.result = result
+ self.calls = 0
+
+ def tool(self) -> AirflowTool:
+ async def function(arguments: dict[str, Any]) -> ToolResult:
+ self.calls += 1
+ return ToolResult(f"{arguments['key']}={self.result}")
+
+ return AirflowTool(
+ name="lookup",
+ description="Look a key up.",
+ parameters={"type": "object", "properties": {"key": {"type": "string"}}},
+ function=function,
+ )
+
+
+class _Provider(ModelProvider):
+ """Resolves every model name to one model."""
+
+ def __init__(self, model: Model) -> None:
+ self.model = model
+
+ def get_model(self, model_name: str | None) -> Model:
+ return self.model
+
+
+class _Report(BaseModel):
+ rows: int
+
+
+def _attempt(storage: MemoryStorage, model: Model, tool: _CountingTool, prompt: str = "go", **kwargs: Any):
+ """Run one task attempt: a fresh journal over the storage the attempts share."""
+ agent = Agent(name="analyst", model=model, tools=AirflowTools(tool.tool()))
+ journal = DurableJournal(storage, clean_up_after_run=True)
+ with journal_scope(journal):
+ return run_durably_sync(agent, prompt, run_config=NO_TRACING, **kwargs), journal
+
+
+class TestRunDurably:
+ def test_a_retry_replays_the_calls_the_failed_attempt_completed(self):
+ storage, tool = MemoryStorage(), _CountingTool()
+ with pytest.raises(RuntimeError, match="worker lost"):
+ _attempt(storage, _ScriptedModel(fail_on_answer=True), tool)
+ assert tool.calls == 1
+
+ retry_model = _ScriptedModel()
+ result, journal = _attempt(storage, retry_model, tool)
+
+ assert result.final_output == "answer: k0=42"
+ # The first request and the tool call replay; only the request that answers is new.
+ assert retry_model.calls == 1
+ assert tool.calls == 1
+ assert journal.stats.replayed == {"model": 1, "tool": 1}
+
+ def test_every_completed_call_replays_not_only_the_first(self):
+ """A replayed answer must turn into the same next request as the live one did."""
+ storage, tool = MemoryStorage(), _CountingTool()
+ with pytest.raises(RuntimeError, match="worker lost"):
+ _attempt(storage, _ScriptedModel(lookups=3, fail_on_answer=True), tool)
+
+ retry_model = _ScriptedModel(lookups=3)
+ _, journal = _attempt(storage, retry_model, tool)
+
+ assert retry_model.calls == 1
+ assert tool.calls == 3
+ assert journal.stats.replayed == {"model": 3, "tool": 3}
+
+ def test_the_sdks_generated_prompt_cache_key_does_not_stop_a_replay(self):
+ """The runner adds a key of its own to every request of a model that accepts one."""
+
+ class _CachingModel(_ScriptedModel):
+ def _supports_default_prompt_cache_key(self) -> bool:
+ return True
+
+ storage, tool = MemoryStorage(), _CountingTool()
+ with pytest.raises(RuntimeError):
+ _attempt(storage, _CachingModel(fail_on_answer=True), tool)
+
+ retry_model = _CachingModel()
+ _attempt(storage, retry_model, tool)
+
+ assert retry_model.calls == 1
+
+ def test_a_replayed_model_call_reports_no_usage(self):
+ storage, tool = MemoryStorage(), _CountingTool()
+ with pytest.raises(RuntimeError):
+ _attempt(storage, _ScriptedModel(fail_on_answer=True, usage=7), tool)
+
+ result, _ = _attempt(storage, _ScriptedModel(usage=7), tool)
+
+ assert result.context_wrapper.usage.input_tokens == 7
+
+ def test_a_successful_run_deletes_what_it_recorded(self):
+ storage, tool = MemoryStorage(), _CountingTool()
+
+ _attempt(storage, _ScriptedModel(), tool)
+
+ assert storage.entries == {}
+
+ def test_a_different_prompt_runs_every_call_again(self):
+ storage, tool = MemoryStorage(), _CountingTool()
+ with pytest.raises(RuntimeError):
+ _attempt(storage, _ScriptedModel(fail_on_answer=True), tool, prompt="first question")
+
+ retry_model = _ScriptedModel()
+ _attempt(storage, retry_model, tool, prompt="second question")
+
+ assert retry_model.calls == 2
+ assert tool.calls == 2
+
+ def test_a_run_the_model_failed_asks_the_model_again(self):
+ """Running out of turns comes from the model's answers, which would replay into the same failure."""
+ storage, tool = MemoryStorage(), _CountingTool()
+ with pytest.raises(MaxTurnsExceeded):
+ _attempt(storage, _ScriptedModel(lookups=3), tool, max_turns=2)
+ assert tool.calls == 2
+
+ retry_model = _ScriptedModel(lookups=1)
+ result, _ = _attempt(storage, retry_model, tool, max_turns=5)
+
+ # The calls before the model's last request replay; that request is asked again, so
+ # the retry answers after the first lookup instead of replaying the second.
+ assert result.final_output == "answer: k0=42"
+ assert retry_model.calls == 1
+ assert tool.calls == 2
+
+ def test_a_model_resolved_by_name_is_recorded(self):
+ storage, tool = MemoryStorage(), _CountingTool()
+
+ def attempt(model: _ScriptedModel) -> None:
+ agent = Agent(name="analyst", model="scripted", tools=AirflowTools(tool.tool()))
+ config = RunConfig(tracing_disabled=True, model_provider=_Provider(model))
+ with journal_scope(DurableJournal(storage, clean_up_after_run=True)):
+ run_durably_sync(agent, "go", run_config=config)
+
+ with pytest.raises(RuntimeError):
+ attempt(_ScriptedModel(fail_on_answer=True))
+ retry_model = _ScriptedModel()
+ attempt(retry_model)
+
+ assert retry_model.calls == 1
+
+ def test_an_agent_used_as_a_tool_replays_as_one_call(self):
+ """Its own model calls run inside the outer tool call and must not take the outer run's positions."""
+ storage, tool = MemoryStorage(), _CountingTool()
+
+ def attempt(outer: _ScriptedModel, inner: _ScriptedModel) -> None:
+ specialist = Agent(name="specialist", model="inner", tools=AirflowTools(tool.tool()))
+ lead = Agent(
+ name="lead",
+ model=outer,
+ tools=[specialist.as_tool("lookup", "Ask the specialist.")],
+ )
+ config = RunConfig(tracing_disabled=True, model_provider=_Provider(inner))
+ with journal_scope(DurableJournal(storage, clean_up_after_run=True)):
+ run_durably_sync(lead, "go", run_config=config)
+
+ # ``as_tool`` takes the question for the specialist as ``input``.
+ with pytest.raises(RuntimeError):
+ attempt(_ScriptedModel(lookups=2, fail_on_answer=True, argument="input"), _ScriptedModel())
+ assert tool.calls == 2
+ outer, inner = _ScriptedModel(lookups=2, argument="input"), _ScriptedModel()
+ attempt(outer, inner)
+
+ # Both outer requests for the specialist and both of its answers replay.
+ assert outer.calls == 1
+ assert inner.calls == 0
+
+ @pytest.mark.enable_redact
+ def test_a_secret_in_a_tool_result_is_not_recorded(self, registered_secret):
+ """Any function tool, not only the Airflow tools, whose results are masked already."""
+
+ @function_tool
+ def lookup(key: str) -> str:
+ """Look a key up."""
+ return f"{key}={registered_secret}"
+
+ agent = Agent(name="analyst", model=_ScriptedModel(fail_on_answer=True), tools=[lookup])
+ storage = MemoryStorage()
+
+ with journal_scope(DurableJournal(storage, clean_up_after_run=True)), pytest.raises(RuntimeError):
+ run_durably_sync(agent, "go", run_config=NO_TRACING)
+
+ (recorded,) = [step for step in storage.steps() if step["name"] == "tool:lookup"]
+ assert recorded["payload"] == "k0=***"
+ assert registered_secret not in json.dumps(storage.entries)
+
+ @pytest.mark.parametrize(
+ "result",
+ [
+ pytest.param(_Report(rows=3), id="pydantic-model"),
+ pytest.param(("orders", 3), id="tuple"),
+ pytest.param(datetime.datetime(2026, 10, 14, 9, 30), id="datetime"),
+ pytest.param({"rows": 3}, id="dict"),
+ pytest.param(ToolOutputText(text="3 rows"), id="structured"),
+ ],
+ )
+ def test_a_tool_result_replays_as_the_model_saw_it(self, result):
+ """The SDK sends most results as str(); a replay must turn into the same next request."""
+
+ @function_tool
+ def lookup(key: str) -> Any:
+ """Look a key up."""
+ return result
+
+ storage = MemoryStorage()
+
+ def attempt(model: _ScriptedModel) -> DurableJournal:
+ agent = Agent(name="analyst", model=model, tools=[lookup])
+ journal = DurableJournal(storage, clean_up_after_run=True)
+ with journal_scope(journal):
+ run_durably_sync(agent, "go", run_config=NO_TRACING)
+ return journal
+
+ with pytest.raises(RuntimeError):
+ attempt(_ScriptedModel(lookups=2, fail_on_answer=True))
+ journal = attempt(_ScriptedModel(lookups=2))
+
+ assert journal.stats.replayed == {"model": 2, "tool": 2}
+
+ @pytest.mark.parametrize(
+ "change",
+ [
+ pytest.param({"instructions": "Answer in French."}, id="instructions"),
+ pytest.param({"model_settings": ModelSettings(temperature=0.1)}, id="settings"),
+ pytest.param({"output_type": _Report}, id="output-type"),
+ pytest.param({"handoffs": [Agent(name="auditor")]}, id="handoffs"),
+ pytest.param({"tools": []}, id="tools"),
+ ],
+ )
+ def test_a_request_that_changed_between_attempts_runs_live(self, change):
+ storage, tool = MemoryStorage(), _CountingTool()
+ with pytest.raises(RuntimeError):
+ _attempt(storage, _ScriptedModel(fail_on_answer=True), tool)
+
+ retry_model = _ScriptedModel()
+ agent = Agent(name="analyst", model=retry_model, tools=AirflowTools(tool.tool())).clone(**change)
+ # Some changes make the scripted answers invalid for the agent; only the first request matters.
+ with journal_scope(DurableJournal(storage, clean_up_after_run=True)), contextlib.suppress(Exception):
+ run_durably_sync(agent, "go", run_config=NO_TRACING)
+
+ # The first request reached the model, rather than replaying the failed try's answer.
+ assert retry_model.tool_results_seen[0] == 0
+
+ def test_a_request_the_sdk_retried_keeps_its_position(self):
+ """The SDK's model retry makes the same call again; the next try must still line up."""
+ settings = ModelSettings(
+ retry=ModelRetrySettings(
+ max_retries=1,
+ backoff=ModelRetryBackoffSettings(initial_delay=0, max_delay=0, jitter=False),
+ policy=lambda context: True,
+ )
+ )
+ storage, tool = MemoryStorage(), _CountingTool()
+
+ def attempt(model: _ScriptedModel) -> DurableJournal:
+ agent = Agent(
+ name="analyst", model=model, tools=AirflowTools(tool.tool()), model_settings=settings
+ )
+ journal = DurableJournal(storage, clean_up_after_run=True)
+ with journal_scope(journal):
+ run_durably_sync(agent, "go", run_config=NO_TRACING)
+ return journal
+
+ with pytest.raises(RuntimeError, match="worker lost"):
+ attempt(_ScriptedModel(fail_first_request=True, fail_on_answer=True))
+ journal = attempt(_ScriptedModel())
+
+ assert journal.stats.replayed == {"model": 1, "tool": 1}
+
+ def test_a_tool_whose_approval_depends_on_its_arguments_is_left_to_the_sdk(self):
+ async def policy(context, arguments, call_id) -> bool:
+ return arguments.get("amount", 0) > 100
+
+ @function_tool(needs_approval=policy)
+ def pay(amount: int = 1000) -> str:
+ """Pay an amount."""
+ return "paid"
+
+ assert _durable_tool(pay, DurableJournal(MemoryStorage()).start_run()) is pay
+
+ @pytest.mark.parametrize("option", ["previous_response_id", "conversation_id"])
+ def test_refuses_a_conversation_kept_on_openais_servers(self, option):
+ with pytest.raises(ValueError, match=option):
+ _attempt(MemoryStorage(), _ScriptedModel(), _CountingTool(), **{option: "resp_1"})
+
+ def test_outside_a_task_it_runs_the_agent_as_runner_run_does(self):
+ model, tool = _ScriptedModel(), _CountingTool()
+ agent = Agent(name="analyst", model=model, tools=AirflowTools(tool.tool()))
+
+ result = run_durably_sync(agent, "go", run_config=NO_TRACING)
+
+ assert result.final_output == "answer: k0=42"
+ assert model.calls == 2