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 1fe65f5f13a Keep parallel tool results replayable across retries in
OpenAI Agents durable runs (#76015)
1fe65f5f13a is described below
commit 1fe65f5f13ae060ac799a598e9c446561c0318fb
Author: Wei Lee <[email protected]>
AuthorDate: Fri Nov 6 08:38:35 2026 +0900
Keep parallel tool results replayable across retries in OpenAI Agents
durable runs (#76015)
Tool calls that run in parallel took their journal positions in the order
they
reached the tool wrapper, which follows how long run hooks, agent hooks and
tool
input guardrails took, not the order the model made the calls. When that
timing
differed between attempts, the positions no longer matched, the run
diverged and
tools that had completed ran again on retry.
The model wrapper now claims one position per function tool call in the
model's
own order and each tool picks up its position by tool call id. Also test the
parallel, handoff, streaming and approval cases.
* Release a durable run's reserved tool steps when it ends and reserve only
for wrapped tools
* Split the _SleepingTools test helper's docstring into a summary and a
description
The single-line docstring ran past the line width; it now opens with a
summary
sentence and keeps the timing, failure and call-count details below it.
* Keep tool calls that share a call id from taking each other's reserved
journal step
A response whose function tool calls repeat a call id, or leave it empty,
used to
have the later reservation overwrite the earlier one, so a tool could pick
up
another call's step and record its result there. A tool now takes a reserved
step only when the step's name and fingerprint are its own, and leaves it
for
the call it belongs to otherwise; a repeated id is reserved once. The SDK
rejects such responses before any tool runs, so this guards the reservation
made before that check.
* Count a reserved tool step as replayed only once its tool takes it
A journal position reserved for a tool call counted as replayed as soon as
it
was reserved, so a call whose tool never ran on the retry (stopped by a tool
input guardrail, for instance) was still counted. The journal now offers
``DurableRun.reserve``, which takes and matches a position like ``claim``
but
leaves the count to ``JournalStep.take``, called when a tool picks its step
up.
``claim`` is unchanged for every other caller, and ``reject`` only
subtracts a
replay it counted.
* Test that only wrapped tools get a reserved journal step
Nothing pinned the marker check in _reserve_tool_steps: widening it to every
FunctionTool left the whole file green. The new test puts a wrapped tool
and a
plain one in the same response and checks that only the wrapped call takes a
position.
* Keep a durable run's reserved tool steps on the run instead of in a
module global
- DurableRun.reserve now takes the response's tool calls, replaces what the
previous response left unclaimed, and keeps the steps on the run.
- DurableRun.claim_reserved hands a tool its reserved step by call id (or a
new position when none matches) and counts the replay.
- A reservation no tool took can no longer be picked up by a later call that
reuses its id.
- Remove _RESERVED_TOOL_STEPS and the cleanup in run_durably.
* Make the reserved-step replacement test fail without the run-level
reservation table
The earlier test also passed when a later response's reservation was skipped
because its call id was already reserved, since the two calls differed in
arguments. Use the same call id and arguments in both responses and check
the executed step lands at the new position.
Drop the test that rejected a reserved step nothing had taken: production
has
no such path, and it read the private reservation table.
* Inline the durable-tool check and rename JournalStep.take to
count_as_replayed
The check had a single caller, and count_as_replayed says what the method
does
rather than reading like a step taking its position.
* Restore the plain form of the durable-tool check
The inlined check wrapped main's condition in a double negation. Put back
main's own line so the diff no longer touches it.
---------
Co-authored-by: Kaxil Naik <[email protected]>
---
.../airflow/providers/common/ai/durable/journal.py | 70 ++-
.../providers/common/ai/durable/openai_agents.py | 64 ++-
.../tests/unit/common/ai/durable/test_journal.py | 76 ++++
.../unit/common/ai/durable/test_openai_agents.py | 486 ++++++++++++++++++++-
4 files changed, 678 insertions(+), 18 deletions(-)
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/durable/journal.py
b/providers/common/ai/src/airflow/providers/common/ai/durable/journal.py
index ffaf955098e..2d437834e4d 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/durable/journal.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/durable/journal.py
@@ -66,6 +66,7 @@ from airflow.providers.common.compat.version_compat import
AIRFLOW_V_3_3_PLUS
if TYPE_CHECKING:
import logging
+ from collections.abc import Sequence
from airflow.providers.common.ai.durable.base import DurableStorageProtocol
from airflow.sdk import Context
@@ -95,9 +96,22 @@ class JournalStep:
replayable: bool = True
replayed: bool = False
payload: Any = None
+ # Whether this step's replay is in ``stats.replayed``; see
:meth:`DurableRun.reserve`.
+ _replay_counted: bool = field(default=False, init=False, repr=False)
# Agent runs started while this step executes; they are numbered under it.
nested_runs: int = field(default=0, repr=False)
+ def count_as_replayed(self) -> None:
+ """
+ Count this step as replayed, once.
+
+ For a step from :meth:`DurableRun.reserve`, which
:meth:`DurableRun.claim_reserved` does;
+ a step from :meth:`DurableRun.claim` is already counted.
+ """
+ if self.replayed and not self._replay_counted:
+ self._replay_counted = True
+ self.durable_run.journal.stats.replayed[self.kind] += 1
+
async def run(
self, body: Callable[[], Awaitable[Any]], *, to_record:
Callable[[Any], Any] | None = None
) -> Any:
@@ -182,8 +196,9 @@ class DurableRun:
One agent run's steps in the journal.
Positions are handed out in the order steps are claimed. A step must be
claimed
- before its first ``await``, so concurrent steps (parallel tool calls) take
their
- positions in the order they were started rather than the order they finish.
+ before its first ``await``, so concurrent steps take their positions in
the order they
+ were started rather than the order they finish. An adapter whose
concurrent steps can
+ be delayed before their claim must claim them itself in a deterministic
order.
Besides its steps, a run keeps one entry about itself: where the last
attempt that
failed stopped being trustworthy, and the furthest position any attempt
reached, so
@@ -215,6 +230,8 @@ class DurableRun:
# Read with the rest of the run's entry: whether the previous attempt
failed because
# of the model's answers, which changes why the retry runs live from
``_diverge_from``.
self._model_at_fault = False
+ # Steps ``reserve`` numbered for the latest model response, by tool
call id, until taken.
+ self._reserved: dict[str, JournalStep] = {}
@property
def position(self) -> int:
@@ -245,6 +262,42 @@ class DurableRun:
self._peeked[position] =
self.journal.storage.load_step(build_step_key(self.key, position))
return self._peeked[position]
+ def reserve(self, calls: Sequence[tuple[str, str, str | None]], *, kind:
StepKind) -> None:
+ """
+ Take a position for each of ``calls`` in order, and keep them for
:meth:`claim_reserved`.
+
+ For an adapter that has to number steps before it knows whether they
will run. Each call
+ is ``(call_id, name, fingerprint)``. A replay counts in the stats only
once the step is
+ taken. Replaces what an earlier ``reserve`` left unclaimed, and a call
whose id is
+ already in ``calls`` takes no position. Like :meth:`claim`, call it
before the first ``await``.
+ """
+ reserved: dict[str, JournalStep] = {}
+ for call_id, name, fingerprint in calls:
+ if call_id not in reserved:
+ reserved[call_id] = self._claim(
+ name, kind=kind, fingerprint=fingerprint, replayable=True,
count_replay=False
+ )
+ self._reserved = reserved
+
+ def claim_reserved(
+ self, call_id: str, name: str, *, kind: StepKind, fingerprint: str |
None
+ ) -> JournalStep:
+ """
+ Take the step :meth:`reserve` kept for ``call_id``, counting its
replay.
+
+ A call whose id was not reserved, or whose reservation is already
taken or belongs to
+ another call sharing the id (some non-OpenAI adapters), takes a new
position as
+ :meth:`claim` does. Call it before the first ``await``.
+ """
+ step = self._reserved.pop(call_id, None)
+ if step is not None and (step.name, step.fingerprint) != (name,
fingerprint):
+ self._reserved[call_id] = step
+ step = None
+ if step is None:
+ step = self.claim(name, kind=kind, fingerprint=fingerprint)
+ step.count_as_replayed()
+ return step
+
def claim(
self, name: str, *, kind: StepKind, fingerprint: str | None,
replayable: bool = True
) -> JournalStep:
@@ -261,6 +314,11 @@ class DurableRun:
whose effect Airflow cannot observe. It still takes a position, so
the steps
after it keep theirs, but nothing is looked up or recorded for it.
"""
+ return self._claim(name, kind=kind, fingerprint=fingerprint,
replayable=replayable, count_replay=True)
+
+ def _claim(
+ self, name: str, *, kind: StepKind, fingerprint: str | None,
replayable: bool, count_replay: bool
+ ) -> JournalStep:
position = self._position
self._position += 1
step = JournalStep(self, position, name, kind, fingerprint,
replayable=replayable)
@@ -287,7 +345,9 @@ class DurableRun:
return step
step.replayed = self._previous_replayed = True
step.payload = entry.get("payload")
- self.journal.stats.replayed[kind] += 1
+ if count_replay:
+ step._replay_counted = True
+ self.journal.stats.replayed[kind] += 1
log.debug("Durable: replayed step", position=position, step=name)
return step
self._diverge(
@@ -397,7 +457,9 @@ class DurableRun:
step.replayed = False
step.payload = None
self._previous_replayed = False
- self.journal.stats.replayed[step.kind] -= 1
+ if step._replay_counted:
+ step._replay_counted = False
+ self.journal.stats.replayed[step.kind] -= 1
self._diverge(step.position, step.name, reason=reason)
def _diverge(self, position: int, name: str, *, reason: str) -> None:
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/durable/openai_agents.py
b/providers/common/ai/src/airflow/providers/common/ai/durable/openai_agents.py
index 1224565c184..9f2388c3dd5 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/durable/openai_agents.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/durable/openai_agents.py
@@ -58,6 +58,7 @@ try:
resolve_function_tool_failure_error_function,
set_function_tool_failure_error_function,
)
+ from openai.types.responses import ResponseFunctionToolCall
except ImportError as e:
raise AirflowOptionalProviderFeatureException(e)
@@ -103,6 +104,10 @@ _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")
+# Set on the copy of a function tool that reports through the journal, so that
reserving a position
+# skips tools nothing wraps, such as those of MCP servers or of a handoff
agent with a named model.
+_DURABLE_MARKER = "_airflow_durable_tool"
+
async def run_durably(
starting_agent: Agent[Any],
@@ -219,6 +224,7 @@ def _get_durable_tool(tool: Tool, durable_run: DurableRun)
-> Tool:
return tool
# A shallow copy, not dataclasses.replace: the tool's schema is already
prepared.
durable_tool = copy.copy(tool)
+ setattr(durable_tool, _DURABLE_MARKER, True)
invoke = tool.on_invoke_tool
# Exceptions a custom failure_error_function turned into text, by tool
call.
handled: dict[str, Exception] = {}
@@ -236,17 +242,9 @@ def _get_durable_tool(tool: Tool, durable_run: DurableRun)
-> 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)
+ name = f"tool:{tool.name}"
+ fingerprint = _build_tool_fingerprint(tool.name, arguments,
ctx.tool_call_id)
+ step = durable_run.claim_reserved(ctx.tool_call_id, name, kind="tool",
fingerprint=fingerprint)
if step.replayed:
return step.payload
try:
@@ -268,6 +266,19 @@ def _get_durable_tool(tool: Tool, durable_run: DurableRun)
-> Tool:
return durable_tool
+def _build_tool_fingerprint(name: str, arguments: str, call_id: str) -> str |
None:
+ try:
+ parsed = json.loads(arguments) if arguments else {}
+ except json.JSONDecodeError:
+ parsed = arguments
+ return fingerprint_tool_call(
+ 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}),
+ call_id,
+ )
+
+
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.
@@ -345,13 +356,16 @@ class _DurableModel(Model):
)
if step.replayed:
try:
- return ModelResponse(
+ replayed = 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}")
+ else:
+ self._reserve_tool_steps(replayed, tools)
+ return replayed
async def call() -> ModelResponse:
return await self.wrapped.get_response(
@@ -359,10 +373,34 @@ class _DurableModel(Model):
)
try:
- return await step.run(call, to_record=_record_response)
+ response = await step.run(call, to_record=_record_response)
except Exception:
self.failed_step = step
raise
+ self._reserve_tool_steps(response, tools)
+ return response
+
+ def _reserve_tool_steps(self, response: ModelResponse, tools: list[Tool])
-> None:
+ """
+ Claim a position for each of the response's tool calls, in the order
the model made them.
+
+ A call whose tool never runs (held for approval, stopped by a
guardrail) leaves its position
+ empty, which a retry runs live without disturbing the positions after
it. The steps the
+ previous response reserved and no tool took are dropped.
+ """
+ durable_names = {tool.name for tool in tools if getattr(tool,
_DURABLE_MARKER, False)}
+ self.durable_run.reserve(
+ [
+ (
+ item.call_id,
+ f"tool:{item.name}",
+ _build_tool_fingerprint(item.name, item.arguments,
item.call_id),
+ )
+ for item in response.output
+ if isinstance(item, ResponseFunctionToolCall) and item.name in
durable_names
+ ],
+ kind="tool",
+ )
def _retried(self, fingerprint: str | None) -> JournalStep | None:
"""
diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_journal.py
b/providers/common/ai/tests/unit/common/ai/durable/test_journal.py
index 41f70605409..f95c983156c 100644
--- a/providers/common/ai/tests/unit/common/ai/durable/test_journal.py
+++ b/providers/common/ai/tests/unit/common/ai/durable/test_journal.py
@@ -324,6 +324,82 @@ class TestReject:
assert retry.journal.stats.replayed.total() == 0
+class TestReserve:
+ @staticmethod
+ async def record_tool(memory_storage, *names):
+ first = DurableJournal(memory_storage).start_run()
+ for name in names:
+ await first.run(name, kind="tool", fingerprint=name, body=lambda:
ok("old"))
+
+ @pytest.mark.asyncio
+ async def test_a_reserved_step_counts_as_replayed_only_once_taken(self,
memory_storage):
+ await self.record_tool(memory_storage, "t")
+
+ retry = DurableJournal(memory_storage).start_run()
+ retry.reserve([("c", "t", "t")], kind="tool")
+ assert retry.journal.stats.replayed.total() == 0
+
+ step = retry.claim_reserved("c", "t", kind="tool", fingerprint="t")
+ step.count_as_replayed()
+
+ assert (step.position, step.replayed, step.payload) == (0, True, "old")
+ assert dict(retry.journal.stats.replayed) == {"tool": 1}
+
+ @pytest.mark.asyncio
+ async def test_a_reserved_id_is_claimed_once(self, memory_storage):
+ await self.record_tool(memory_storage, "t")
+
+ retry = DurableJournal(memory_storage).start_run()
+ retry.reserve([("c", "t", "t")], kind="tool")
+ retry.claim_reserved("c", "t", kind="tool", fingerprint="t")
+ again = retry.claim_reserved("c", "t", kind="tool", fingerprint="t")
+
+ assert (again.position, again.replayed) == (1, False)
+ assert dict(retry.journal.stats.replayed) == {"tool": 1}
+
+ @pytest.mark.asyncio
+ async def test_a_later_reserve_drops_the_steps_nothing_took(self,
memory_storage):
+ await self.record_tool(memory_storage, "t")
+
+ retry = DurableJournal(memory_storage).start_run()
+ retry.reserve([("old", "t", "t")], kind="tool")
+ retry.reserve([("new", "u", "u")], kind="tool")
+ stale = retry.claim_reserved("old", "t", kind="tool", fingerprint="t")
+
+ assert (stale.position, stale.replayed) == (2, False)
+
+ @pytest.mark.asyncio
+ async def test_a_reserve_without_calls_drops_the_steps_nothing_took(self,
memory_storage):
+ await self.record_tool(memory_storage, "t")
+
+ retry = DurableJournal(memory_storage).start_run()
+ retry.reserve([("old", "t", "t")], kind="tool")
+ retry.reserve([], kind="tool")
+ stale = retry.claim_reserved("old", "t", kind="tool", fingerprint="t")
+
+ assert (stale.position, stale.replayed) == (1, False)
+
+ @pytest.mark.asyncio
+ async def
test_a_call_that_shares_an_id_leaves_the_reservation_to_the_matching_call(self,
memory_storage):
+ await self.record_tool(memory_storage, "t")
+
+ retry = DurableJournal(memory_storage).start_run()
+ retry.reserve([("c", "t", "t")], kind="tool")
+ other = retry.claim_reserved("c", "t", kind="tool",
fingerprint="other")
+ own = retry.claim_reserved("c", "t", kind="tool", fingerprint="t")
+
+ assert (other.position, other.replayed) == (1, False)
+ assert (own.position, own.replayed) == (0, True)
+
+ def test_duplicate_ids_in_one_reserve_take_one_position(self,
memory_storage):
+ run = DurableJournal(memory_storage).start_run()
+ run.reserve([("c", "t", "a"), ("c", "t", "b"), ("d", "t", "d")],
kind="tool")
+
+ assert run.position == 2
+ assert run.claim_reserved("c", "t", kind="tool",
fingerprint="a").position == 0
+ assert run.claim_reserved("d", "t", kind="tool",
fingerprint="d").position == 1
+
+
class TestFailuresDuringReplay:
@pytest.mark.asyncio
async def
test_an_attempt_killed_while_replaying_leaves_the_earlier_steps_replayable(
diff --git
a/providers/common/ai/tests/unit/common/ai/durable/test_openai_agents.py
b/providers/common/ai/tests/unit/common/ai/durable/test_openai_agents.py
index 30baff77db8..da4adcaf7fb 100644
--- a/providers/common/ai/tests/unit/common/ai/durable/test_openai_agents.py
+++ b/providers/common/ai/tests/unit/common/ai/durable/test_openai_agents.py
@@ -16,10 +16,13 @@
# under the License.
from __future__ import annotations
+import asyncio
+import collections
import contextlib
import datetime
import json
from typing import Any
+from unittest import mock
import pytest
@@ -27,24 +30,36 @@ pytest.importorskip("agents")
from agents import (
Agent,
+ AgentHooks,
Model,
ModelProvider,
ModelResponse,
ModelSettings,
+ ModelTracing,
RunConfig,
+ RunHooks,
SQLiteSession,
+ ToolGuardrailFunctionOutput,
ToolOutputText,
Usage,
function_tool,
+ handoff,
+ tool_input_guardrail,
)
from agents.exceptions import MaxTurnsExceeded, ModelBehaviorError
from agents.retry import ModelRetryBackoffSettings, ModelRetrySettings
+from agents.tool_context import ToolContext
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.durable.openai_agents import (
+ _build_tool_fingerprint,
+ _DurableModel,
+ _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
@@ -139,6 +154,101 @@ class _CountingTool:
)
+class _ParallelModel(Model):
+ """
+ Asks for one ``lookup`` per key in a single response, then answers with
their results.
+
+ The calls get the ids ``call-<key>``. ``fail_on_answer`` makes the request
that would
+ answer raise, as a crashed worker would leave it.
+ """
+
+ def __init__(self, keys: tuple[str, ...] = ("a", "b"), *, fail_on_answer:
bool = False) -> None:
+ self.keys = keys
+ self.fail_on_answer = fail_on_answer
+ self.calls = 0
+
+ async def get_response(self, system_instructions, input, *args: Any,
**kwargs: Any) -> ModelResponse:
+ self.calls += 1
+ outputs = [
+ item["output"]
+ for item in (input if isinstance(input, list) else [])
+ if isinstance(item, dict) and item.get("type") ==
"function_call_output"
+ ]
+ output: list[ResponseFunctionToolCall | ResponseOutputMessage]
+ if not outputs:
+ output = [
+ ResponseFunctionToolCall(
+ id=f"fc-{key}",
+ call_id=f"call-{key}",
+ name="lookup",
+ arguments=json.dumps({"key": key}),
+ type="function_call",
+ )
+ for key in self.keys
+ ]
+ else:
+ if self.fail_on_answer:
+ raise RuntimeError("worker lost")
+ text = ResponseOutputText(text=f"answer: {sorted(outputs)}",
type="output_text", annotations=[])
+ output = [
+ ResponseOutputMessage(
+ id="msg", role="assistant", status="completed",
type="message", content=[text]
+ )
+ ]
+ return ModelResponse(output=list(output), usage=Usage(),
response_id=f"resp-{self.calls}")
+
+ def stream_response(self, *args: Any, **kwargs: Any):
+ raise NotImplementedError
+
+
+class _SleepingTools:
+ """
+ A ``lookup`` tool with controllable timing and failures.
+
+ It sleeps ``body_delays[key]`` seconds, fails for the keys in ``failing``,
and counts the calls per key.
+ """
+
+ def __init__(self, body_delays: dict[str, float] | None = None, failing:
set[str] | None = None) -> None:
+ self.body_delays = body_delays or {}
+ self.failing = failing or set()
+ self.calls: collections.Counter[str] = collections.Counter()
+
+ def tools(self) -> list[Any]:
+ async def function(arguments: dict[str, Any]) -> ToolResult:
+ key = arguments["key"]
+ self.calls[key] += 1
+ await asyncio.sleep(self.body_delays.get(key, 0))
+ if key in self.failing:
+ raise RuntimeError(f"lookup {key} failed")
+ return ToolResult(f"{key}=42")
+
+ return AirflowTools(
+ AirflowTool(
+ name="lookup",
+ description="Look a key up.",
+ parameters={"type": "object", "properties": {"key": {"type":
"string"}}},
+ function=function,
+ )
+ )
+
+
+class _TransferModel(Model):
+ """Hands the run over to the ``specialist`` agent on its first request."""
+
+ async def get_response(self, system_instructions, input, *args: Any,
**kwargs: Any) -> ModelResponse:
+ call = ResponseFunctionToolCall(
+ id="fc-transfer",
+ call_id="call-transfer",
+ name="transfer_to_specialist",
+ arguments="{}",
+ type="function_call",
+ )
+ return ModelResponse(output=[call], usage=Usage(),
response_id="resp-transfer")
+
+ def stream_response(self, *args: Any, **kwargs: Any):
+ raise NotImplementedError
+
+
class _Provider(ModelProvider):
"""Resolves every model name to one model."""
@@ -519,3 +629,377 @@ class TestRunDurably:
assert result.final_output == "answer: k0=42"
assert model.calls == 2
+
+
+class _SlowRunHooks(RunHooks):
+ """Makes every tool call wait ``delays[call id]`` seconds before it is
invoked."""
+
+ def __init__(self, delays: dict[str, float]) -> None:
+ self.delays = delays
+
+ async def on_tool_start(self, context, agent, tool) -> None:
+ await asyncio.sleep(self.delays[context.tool_call_id])
+
+
+class _SlowAgentHooks(AgentHooks):
+ """Makes every tool call of the agent wait ``delays[call id]`` seconds
before it is invoked."""
+
+ def __init__(self, delays: dict[str, float]) -> None:
+ self.delays = delays
+
+ async def on_tool_start(self, context, agent, tool) -> None:
+ await asyncio.sleep(self.delays[context.tool_call_id])
+
+
+def _attempt_parallel(
+ storage: MemoryStorage,
+ model: Model,
+ tools: _SleepingTools,
+ *,
+ agent_hooks: AgentHooks | None = None,
+ guardrail_delays: dict[str, float] | None = None,
+ rejected_calls: set[str] | None = None,
+ **kwargs: Any,
+):
+ """Run one attempt of an agent whose model asks for several tools at
once."""
+ function_tools = tools.tools()
+ if guardrail_delays is not None or rejected_calls is not None:
+
+ @tool_input_guardrail
+ async def wait(data):
+ await asyncio.sleep((guardrail_delays or
{}).get(data.context.tool_call_id, 0))
+ if rejected_calls and data.context.tool_call_id in rejected_calls:
+ return ToolGuardrailFunctionOutput.reject_content("stopped by
a guardrail")
+ return ToolGuardrailFunctionOutput.allow()
+
+ for function_tool_ in function_tools:
+ function_tool_.tool_input_guardrails = [wait]
+ agent = Agent(name="analyst", model=model, tools=function_tools,
hooks=agent_hooks)
+ journal = DurableJournal(storage, clean_up_after_run=True)
+ with journal_scope(journal):
+ return run_durably_sync(agent, "go", run_config=NO_TRACING, **kwargs),
journal
+
+
+class TestParallelToolCalls:
+ def
test_parallel_tool_calls_replay_when_their_bodies_finish_in_a_different_order(self):
+ """
+ Characterizes the documented claim that parallel calls keep the order
they started in.
+
+ The model's response reserves a position for each tool call, so how
long a tool body takes does not matter.
+ """
+ storage = MemoryStorage()
+ tools = _SleepingTools({"a": 0.05, "b": 0.01})
+ with pytest.raises(RuntimeError, match="worker lost"):
+ _attempt_parallel(storage, _ParallelModel(fail_on_answer=True),
tools)
+ assert tools.calls == {"a": 1, "b": 1}
+
+ tools.body_delays = {"a": 0.01, "b": 0.05}
+ retry_model = _ParallelModel()
+ _, journal = _attempt_parallel(storage, retry_model, tools)
+
+ assert journal.stats.replayed == {"model": 1, "tool": 2}
+ assert tools.calls == {"a": 1, "b": 1}
+ assert retry_model.calls == 1
+
+ def test_only_the_parallel_tool_call_that_failed_runs_again_on_retry(self):
+ storage = MemoryStorage()
+ tools = _SleepingTools(failing={"b"})
+ with pytest.raises(Exception, match="lookup b failed"):
+ _attempt_parallel(storage, _ParallelModel(("a", "b", "c")), tools)
+ assert tools.calls == {"a": 1, "b": 1, "c": 1}
+
+ tools.failing = set()
+ retry_model = _ParallelModel(("a", "b", "c"))
+ _, journal = _attempt_parallel(storage, retry_model, tools)
+
+ assert tools.calls == {"a": 1, "b": 2, "c": 1}
+ assert journal.stats.replayed == {"model": 1, "tool": 2}
+ assert retry_model.calls == 1
+
+ def
test_a_replayed_tool_call_a_guardrail_stops_is_not_counted_as_replayed(self):
+ storage = MemoryStorage()
+ tools = _SleepingTools()
+ rejected: set[str] = set()
+ with pytest.raises(RuntimeError, match="worker lost"):
+ _attempt_parallel(storage, _ParallelModel(fail_on_answer=True),
tools, rejected_calls=rejected)
+ assert tools.calls == {"a": 1, "b": 1}
+
+ rejected.add("call-b")
+ retry_model = _ParallelModel()
+ _, journal = _attempt_parallel(storage, retry_model, tools,
rejected_calls=rejected)
+
+ assert journal.stats.replayed == {"model": 1, "tool": 1}
+ assert tools.calls == {"a": 1, "b": 1}
+ assert retry_model.calls == 1
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("call_id", ["call-shared", ""],
ids=["duplicate-id", "empty-id"])
+ async def
test_tool_calls_sharing_a_call_id_each_replay_their_own_result(self, call_id):
+ """The SDK rejects a repeated id itself, but a model that emitted one
must not make the calls trade steps."""
+ storage = MemoryStorage()
+ tools = _SleepingTools()
+ response = ModelResponse(
+ output=[
+ ResponseFunctionToolCall(
+ id=f"fc-{key}",
+ call_id=call_id,
+ name="lookup",
+ arguments=json.dumps({"key": key}),
+ type="function_call",
+ )
+ for key in ("a", "b")
+ ],
+ usage=Usage(),
+ response_id="resp",
+ )
+
+ async def attempt(order: tuple[str, ...]) -> tuple[dict[str, Any],
DurableJournal]:
+ journal = DurableJournal(storage)
+ durable_run = journal.start_run()
+ durable_tools = [_get_durable_tool(tool, durable_run) for tool in
tools.tools()]
+ _DurableModel(mock.create_autospec(Model, instance=True),
durable_run)._reserve_tool_steps(
+ response, durable_tools
+ )
+ outputs = {}
+ for key in order:
+ context = ToolContext(
+ context=None,
+ usage=Usage(),
+ tool_name="lookup",
+ tool_call_id=call_id,
+ tool_arguments=json.dumps({"key": key}),
+ )
+ outputs[key] = await durable_tools[0].on_invoke_tool(context,
json.dumps({"key": key}))
+ return outputs, journal
+
+ await attempt(("a", "b"))
+ assert tools.calls == {"a": 1, "b": 1}
+
+ outputs, journal = await attempt(("b", "a"))
+
+ assert tools.calls == {"a": 1, "b": 1}
+ assert outputs == {"a": "a=42", "b": "b=42"}
+ assert journal.stats.replayed["tool"] == 2
+
+ def test_only_wrapped_tools_get_a_reserved_step(self):
+ """
+ A tool the model is handed but the durable run does not wrap never
takes a reserved step.
+
+ MCP tools and a handoff agent's tools reach the model this way, so
reserving for them would leave
+ positions nobody takes.
+ """
+
+ @function_tool
+ def plain_lookup(key: str) -> str:
+ """Look a key up."""
+ return key
+
+ journal = DurableJournal(MemoryStorage())
+ durable_run = journal.start_run()
+ durable_tool = _get_durable_tool(_SleepingTools().tools()[0],
durable_run)
+ fingerprint = _build_tool_fingerprint("lookup", json.dumps({"key":
"a"}), "call-durable")
+ response = ModelResponse(
+ output=[
+ ResponseFunctionToolCall(
+ id=f"fc-{call_id}",
+ call_id=call_id,
+ name=name,
+ arguments=json.dumps({"key": "a"}),
+ type="function_call",
+ )
+ for call_id, name in (("call-durable", "lookup"),
("call-plain", "plain_lookup"))
+ ],
+ usage=Usage(),
+ response_id="resp",
+ )
+
+ _DurableModel(mock.create_autospec(Model, instance=True),
durable_run)._reserve_tool_steps(
+ response, [durable_tool, plain_lookup]
+ )
+
+ assert durable_run.position == 1
+ durable = durable_run.claim_reserved(
+ "call-durable", "tool:lookup", kind="tool", fingerprint=fingerprint
+ )
+ plain = durable_run.claim_reserved("call-plain", "tool:plain_lookup",
kind="tool", fingerprint=None)
+ assert (durable.position, plain.position) == (0, 1)
+
+ @pytest.mark.asyncio
+ async def
test_a_later_response_reserves_a_new_step_for_a_call_with_the_same_id_and_arguments(self):
+ storage = MemoryStorage()
+ tools = _SleepingTools()
+ response = ModelResponse(
+ output=[
+ ResponseFunctionToolCall(
+ id="fc-x",
+ call_id="call-x",
+ name="lookup",
+ arguments=json.dumps({"key": "a"}),
+ type="function_call",
+ )
+ ],
+ usage=Usage(),
+ response_id="resp",
+ )
+ durable_run = DurableJournal(storage).start_run()
+ durable_tools = [_get_durable_tool(tool, durable_run) for tool in
tools.tools()]
+ model = _DurableModel(mock.create_autospec(Model, instance=True),
durable_run)
+ # No tool takes the first response's call; only the second response's
is invoked.
+ model._reserve_tool_steps(response, durable_tools)
+ model._reserve_tool_steps(response, durable_tools)
+ context = ToolContext(
+ context=None,
+ usage=Usage(),
+ tool_name="lookup",
+ tool_call_id="call-x",
+ tool_arguments=json.dumps({"key": "a"}),
+ )
+
+ await durable_tools[0].on_invoke_tool(context, json.dumps({"key":
"a"}))
+
+ assert build_step_key(durable_run.key, 1) in storage.entries
+ assert build_step_key(durable_run.key, 0) not in storage.entries
+
+ @pytest.mark.parametrize("source", ["run-hook", "agent-hook",
"tool-input-guardrail"])
+ @pytest.mark.parametrize(
+ "first_attempt_delays",
+ [{"call-a": 0.01, "call-b": 0.1}, {"call-a": 0.1, "call-b": 0.01}],
+ ids=["a-first", "b-first"],
+ )
+ def
test_parallel_tool_calls_replay_when_what_runs_before_them_takes_a_different_time(
+ self, source, first_attempt_delays
+ ):
+ """A hook or guardrail the SDK awaits before invoking a tool can let a
later call claim first."""
+ storage = MemoryStorage()
+ tools = _SleepingTools()
+
+ def attempt(model: Model, delays: dict[str, float]):
+ kwargs: dict[str, Any] = {}
+ if source == "run-hook":
+ kwargs["hooks"] = _SlowRunHooks(delays)
+ elif source == "agent-hook":
+ kwargs["agent_hooks"] = _SlowAgentHooks(delays)
+ else:
+ kwargs["guardrail_delays"] = delays
+ return _attempt_parallel(storage, model, tools, **kwargs)
+
+ with pytest.raises(RuntimeError, match="worker lost"):
+ attempt(_ParallelModel(fail_on_answer=True), first_attempt_delays)
+ assert tools.calls == {"a": 1, "b": 1}
+
+ _, journal = attempt(_ParallelModel(), {"call-a": 0.1, "call-b": 0.01})
+
+ assert journal.stats.replayed == {"model": 1, "tool": 2}
+ assert tools.calls == {"a": 1, "b": 1}
+
+
+class TestDurableExceptions:
+ def test_a_static_handoff_agent_replays_its_calls(self):
+ storage, tool = MemoryStorage(), _CountingTool()
+
+ def attempt(specialist_model: _ScriptedModel) -> DurableJournal:
+ specialist = Agent(name="specialist", model=specialist_model,
tools=AirflowTools(tool.tool()))
+ lead = Agent(name="lead", model=_TransferModel(),
handoffs=[specialist])
+ journal = DurableJournal(storage, clean_up_after_run=True)
+ with journal_scope(journal):
+ run_durably_sync(lead, "go", run_config=NO_TRACING)
+ return journal
+
+ # The transfer's own result is the first function output the
specialist sees.
+ with pytest.raises(RuntimeError, match="worker lost"):
+ attempt(_ScriptedModel(lookups=2, fail_on_answer=True))
+ assert tool.calls == 1
+
+ specialist_model = _ScriptedModel(lookups=2)
+ journal = attempt(specialist_model)
+
+ assert journal.stats.replayed == {"model": 2, "tool": 1}
+ assert specialist_model.calls == 1
+ assert tool.calls == 1
+
+ @pytest.mark.parametrize(
+ ("by_name", "replayed_models"),
+ [
+ pytest.param(False, 1, id="model-instance"),
+ pytest.param(True, 2, id="model-name"),
+ ],
+ )
+ def test_a_handoff_object_agent_runs_its_own_tools_live(self, by_name,
replayed_models):
+ storage, tool = MemoryStorage(), _CountingTool()
+
+ def attempt(specialist_model: _ScriptedModel) -> DurableJournal:
+ specialist = Agent(
+ name="specialist",
+ model="specialist-model" if by_name else specialist_model,
+ tools=AirflowTools(tool.tool()),
+ )
+ lead = Agent(name="lead", model=_TransferModel(),
handoffs=[handoff(specialist)])
+ config = RunConfig(tracing_disabled=True,
model_provider=_Provider(specialist_model))
+ journal = DurableJournal(storage, clean_up_after_run=True)
+ with journal_scope(journal):
+ run_durably_sync(lead, "go", run_config=config)
+ return journal
+
+ with pytest.raises(RuntimeError, match="worker lost"):
+ attempt(_ScriptedModel(lookups=2, fail_on_answer=True))
+ assert tool.calls == 1
+
+ journal = attempt(_ScriptedModel(lookups=2))
+
+ assert journal.stats.replayed == {"model": replayed_models}
+ assert tool.calls == 2
+
+ @pytest.mark.asyncio
+ async def test_a_streamed_model_call_is_not_recorded(self):
+ events = [object(), object()]
+
+ async def stream(*args: Any, **kwargs: Any):
+ for event in events:
+ yield event
+
+ wrapped = mock.create_autospec(Model, instance=True)
+ wrapped.stream_response.side_effect = stream
+ storage = MemoryStorage()
+ durable_run = DurableJournal(storage).start_run()
+ model = _DurableModel(wrapped, durable_run)
+
+ streamed = [
+ event
+ async for event in model.stream_response(
+ None,
+ "go",
+ ModelSettings(),
+ [],
+ None,
+ [],
+ ModelTracing.DISABLED,
+ previous_response_id=None,
+ conversation_id=None,
+ prompt=None,
+ )
+ ]
+
+ assert streamed == events
+ wrapped.get_response.assert_not_called()
+ assert durable_run.position == 0
+ assert storage.steps() == []
+
+ def test_a_tool_that_needs_approval_is_not_invoked_or_recorded(self):
+ calls = []
+
+ @function_tool(needs_approval=True)
+ def lookup(key: str) -> str:
+ """Look a key up."""
+ calls.append(key)
+ return f"{key}=42"
+
+ journal = DurableJournal(MemoryStorage(), clean_up_after_run=True)
+ with journal_scope(journal):
+ result = run_durably_sync(
+ Agent(name="analyst", model=_ScriptedModel(), tools=[lookup]),
"go", run_config=NO_TRACING
+ )
+
+ assert len(result.interruptions) == 1
+ assert calls == []
+ assert journal.stats.recorded == {"model": 1}
+ assert _get_durable_tool(lookup, journal.start_run()) is not lookup
