This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 0eab1db3d2b Replay an OpenAI Agents SDK agent's completed calls when
Airflow retries its task (#74770)
0eab1db3d2b is described below
commit 0eab1db3d2b999332d1a4e759e3ff401f18da188
Author: Kaxil Naik <[email protected]>
AuthorDate: Sat Oct 17 13:55:15 2026 +0100
Replay an OpenAI Agents SDK agent's completed calls when Airflow retries
its task (#74770)
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.
A function tool's exception that failure_error_function turned into text was
recorded as the tool's result, so a retry replayed the error and never ran
the
tool again. Such a call is no longer recorded, as a call that raised is not.
The guide now splits what a retry replays from what always runs live, names
the journal's storage, and the tests cover unreadable recordings, structured
results, sessions, and a tool-raised ModelBehaviorError.
---
providers/common/ai/docs/durable_execution.rst | 3 +-
providers/common/ai/docs/features.rst | 3 +-
.../common/ai/docs/frameworks/openai_agents.rst | 72 ++-
providers/common/ai/docs/stability.rst | 3 +-
.../providers/common/ai/durable/openai_agents.py | 479 +++++++++++++++++++
.../example_dags/example_openai_agents_durable.py | 90 ++++
.../unit/common/ai/durable/test_openai_agents.py | 521 +++++++++++++++++++++
7 files changed, 1166 insertions(+), 5 deletions(-)
diff --git a/providers/common/ai/docs/durable_execution.rst
b/providers/common/ai/docs/durable_execution.rst
index 53d09702d74..6c0d34645ea 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 7dc46e3f8cb..94fc60c2654 100644
--- a/providers/common/ai/docs/features.rst
+++ b/providers/common/ai/docs/features.rst
@@ -272,7 +272,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 103caf46a53..a704a4298d9 100644
--- a/providers/common/ai/docs/frameworks/openai_agents.rst
+++ b/providers/common/ai/docs/frameworks/openai_agents.rst
@@ -142,11 +142,79 @@ 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``, except
+``previous_response_id``, ``conversation_id`` and
``auto_previous_response_id``, which they
+refuse: with them OpenAI keeps the conversation, and a replayed call never
reaches it.
+
+.. 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 Airflow 3.3 and later the journal is kept in
the task
+state store; on earlier versions, set ``[common.ai] durable_cache_path`` to
where it is
+kept. Once the run succeeds, what was recorded is deleted, so clearing the
task later
+starts over. Outside a task there is nothing to replay into, and the agent
runs as
+``Runner.run`` runs it.
+
+What a retry replays
+^^^^^^^^^^^^^^^^^^^^
+
+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: 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. Recorded tool results pass through Airflow's secret masker before they
are saved,
+and a replayed model call reports no token usage.
+
+From the first call that differs, that call and every call after it run live:
+
+- 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.
+- 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.
+
+What always runs live
+^^^^^^^^^^^^^^^^^^^^^
+
+- A call that was running when the try failed, so tools that change data must
still be
+ safe to repeat.
+- A tool call that failed, including one whose exception
``failure_error_function`` turned
+ into text for the model. A ``ModelBehaviorError`` that a tool raises, such
as arguments
+ that fail validation with ``failure_error_function=None``, counts as the
tool's failure:
+ the tool call runs again, but the answer that asked for it replays.
+- Tools other than function tools: hosted tools such as web search, the tools
of
+ ``mcp_servers``, and 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
too unless
+ the agent's ``model`` is a model name or not set.
+- A streamed run.
+- Every call, when the run has a ``session``: 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.
+
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 cced251110c..5feb431be94 100644
--- a/providers/common/ai/docs/stability.rst
+++ b/providers/common/ai/docs/stability.rst
@@ -261,7 +261,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..1224565c184
--- /dev/null
+++
b/providers/common/ai/src/airflow/providers/common/ai/durable/openai_agents.py
@@ -0,0 +1,479 @@
+# 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
+ from agents.tool import (
+ default_tool_error_function,
+ resolve_function_tool_failure_error_function,
+ set_function_tool_failure_error_function,
+ )
+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)
+
+# Set on a tool call's context by the SDK when its default
failure_error_function turned the
+# tool's exception into the text the model sees.
+_DEFAULT_FAILURE_HANDLED = "_function_tool_default_failure_handled"
+
+# 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
too unless the
+ agent's ``model`` is a model name or not set, and streamed runs. A tool
call that fails,
+ including one whose exception ``failure_error_function`` turns into text
for the model,
+ is not recorded either, so a retry runs it again. 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.
+ """
+ if (journal := current_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 = _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 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 _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()
+ # 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):
+ durable_agent.model = _DurableModel(agent.model, durable_run)
+ durable_agent.tools = [_get_durable_tool(tool, durable_run) for tool in
agent.tools]
+ durable_agent.handoffs = [
+ _get_durable_agent(target, durable_run, seen) if isinstance(target,
Agent) else target
+ for target in agent.handoffs
+ ]
+ return durable_agent
+
+
+def _get_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
+ # A shallow copy, not dataclasses.replace: the tool's schema is already
prepared.
+ durable_tool = copy.copy(tool)
+ invoke = tool.on_invoke_tool
+ # Exceptions a custom failure_error_function turned into text, by tool
call.
+ handled: dict[str, Exception] = {}
+ formatter = resolve_function_tool_failure_error_function(tool)
+ bind = getattr(invoke, "__agents_bind_function_tool__", None)
+ if formatter is not None and formatter is not default_tool_error_function
and bind is not None:
+
+ def remember(ctx: Any, error: Exception) -> Any:
+ handled[ctx.tool_call_id] = error
+ return formatter(ctx, error)
+
+ set_function_tool_failure_error_function(durable_tool, remember)
+ invoke = bind(durable_tool)
+
+ async def on_invoke_tool(ctx: ToolContext[Any], arguments: str) -> Any:
+ if _is_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,
+ )
+ step = durable_run.claim(f"tool:{tool.name}", kind="tool",
fingerprint=fingerprint)
+ if step.replayed:
+ return step.payload
+ try:
+ with step.executing():
+ output = await invoke(ctx, arguments)
+ except Exception as error:
+ step.fail(error)
+ raise
+ # A failure the SDK turned into text for the model is not recorded, so
a retry runs the
+ # call again, as it does a call that raised.
+ if (handled_error := handled.pop(ctx.tool_call_id, None)) is not None:
+ step.fail(handled_error)
+ elif not getattr(ctx, _DEFAULT_FAILURE_HANDLED, False):
+ # Secrets must not reach the journal any more than the model.
+ step.record(mask_secrets(_as_sent(output)))
+ return output
+
+ durable_tool.on_invoke_tool = on_invoke_tool
+ return durable_tool
+
+
+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 _is_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.wrapped = 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_step: 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 _is_inside_a_call(self.durable_run):
+ return await self.wrapped.get_response(
+ system_instructions, input, model_settings, tools,
output_schema, handoffs, tracing, **kwargs
+ )
+ fingerprint = _fingerprint_request(
+ self.wrapped, 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.wrapped.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 = 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_step = self.failed_step, 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.wrapped.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.wrapped, "_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.wrapped.get_retry_advice(request)
+
+ async def _cleanup_on_run_end(self, owner: object) -> None:
+ await self.wrapped._cleanup_on_run_end(owner)
+
+ async def close(self) -> None:
+ await self.wrapped.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..30baff77db8
--- /dev/null
+++ b/providers/common/ai/tests/unit/common/ai/durable/test_openai_agents.py
@@ -0,0 +1,521 @@
+# 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,
+ SQLiteSession,
+ ToolOutputText,
+ Usage,
+ function_tool,
+)
+from agents.exceptions import MaxTurnsExceeded, ModelBehaviorError
+from agents.retry import ModelRetryBackoffSettings, ModelRetrySettings
+from openai.types.responses import ResponseFunctionToolCall,
ResponseOutputMessage, ResponseOutputText
+from pydantic import BaseModel
+
+from airflow.providers.common.ai.durable.base import build_step_key
+from airflow.providers.common.ai.durable.journal import DurableJournal,
journal_scope
+from airflow.providers.common.ai.durable.openai_agents import
_get_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"),
+ pytest.param(
+ [ToolOutputText(text="3 rows"), ToolOutputText(text="2
orders")], id="structured-list"
+ ),
+ pytest.param({"type": "text", "text": "3 rows"}, id="typed-dict"),
+ ],
+ )
+ 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 _get_durable_tool(pay,
DurableJournal(MemoryStorage()).start_run()) is pay
+
+ @pytest.mark.parametrize(
+ ("option", "value"),
+ [
+ ("previous_response_id", "resp_1"),
+ ("conversation_id", "conv_1"),
+ ("auto_previous_response_id", True),
+ ],
+ )
+ def test_refuses_a_conversation_kept_on_openais_servers(self, option,
value):
+ with pytest.raises(ValueError, match=option):
+ _attempt(MemoryStorage(), _ScriptedModel(), _CountingTool(),
**{option: value})
+
+ def
test_a_recorded_answer_that_cannot_be_read_runs_live_with_every_call_after_it(self):
+ storage, tool = MemoryStorage(), _CountingTool()
+ with pytest.raises(RuntimeError):
+ _attempt(storage, _ScriptedModel(fail_on_answer=True), tool)
+ storage.entries[build_step_key(0, 0)]["payload"]["output"] = [{"type":
"not an output item"}]
+
+ retry_model = _ScriptedModel()
+ _, journal = _attempt(storage, retry_model, tool)
+
+ assert retry_model.calls == 2
+ assert tool.calls == 2
+ assert journal.stats.replayed == {"model": 0}
+
+ @pytest.mark.parametrize(
+ "failure_error_function",
+ [
+ pytest.param({}, id="default-formatter"),
+ pytest.param(
+ {"failure_error_function": lambda ctx, error: f"lookup failed:
{error}"}, id="custom"
+ ),
+ ],
+ )
+ def test_a_tool_call_that_failed_runs_again_on_retry(self,
failure_error_function):
+ """The SDK hands the model the error as text; that text must not
replay as the tool's result."""
+ calls = []
+
+ @function_tool(**failure_error_function)
+ def lookup(key: str) -> str:
+ """Look a key up."""
+ calls.append(key)
+ if len(calls) == 1:
+ raise ConnectionError("warehouse unreachable")
+ return f"{key}=42"
+
+ storage = MemoryStorage()
+
+ def attempt(model: _ScriptedModel) -> DurableJournal:
+ journal = DurableJournal(storage, clean_up_after_run=True)
+ with journal_scope(journal):
+ run_durably_sync(
+ Agent(name="analyst", model=model, tools=[lookup]), "go",
run_config=NO_TRACING
+ )
+ return journal
+
+ with pytest.raises(RuntimeError, match="worker lost"):
+ attempt(_ScriptedModel(fail_on_answer=True))
+ retry_model = _ScriptedModel()
+ attempt(retry_model)
+
+ assert len(calls) == 2
+ assert retry_model.calls == 1
+
+ def
test_a_model_behavior_error_a_tool_raises_reruns_the_tool_and_replays_the_answer_before_it(self):
+ calls = []
+
+ @function_tool(failure_error_function=None)
+ def lookup(key: str) -> str:
+ """Look a key up."""
+ calls.append(key)
+ if len(calls) == 1:
+ raise ModelBehaviorError("lookup was called with a key that
does not exist")
+ return f"{key}=42"
+
+ storage = MemoryStorage()
+
+ def attempt(model: _ScriptedModel) -> None:
+ with journal_scope(DurableJournal(storage,
clean_up_after_run=True)):
+ run_durably_sync(
+ Agent(name="analyst", model=model, tools=[lookup]), "go",
run_config=NO_TRACING
+ )
+
+ with pytest.raises(ModelBehaviorError):
+ attempt(_ScriptedModel())
+ retry_model = _ScriptedModel()
+ attempt(retry_model)
+
+ assert len(calls) == 2
+ # The answer that asked for the lookup replays; only the final answer
is asked for.
+ assert retry_model.calls == 1
+
+ def test_a_run_with_a_session_runs_live_on_retry(self, tmp_path):
+ """The session already holds the failed try's prompt, so the first
request differs."""
+ storage, tool = MemoryStorage(), _CountingTool()
+ session = SQLiteSession("analyst", str(tmp_path / "session.db"))
+ with pytest.raises(RuntimeError):
+ _attempt(storage, _ScriptedModel(fail_on_answer=True), tool,
session=session)
+
+ _, journal = _attempt(storage, _ScriptedModel(), tool, session=session)
+
+ assert journal.stats.replayed == {}
+
+ 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
