kaxil commented on code in PR #76015:
URL: https://github.com/apache/airflow/pull/76015#discussion_r4431702509
##########
providers/common/ai/src/airflow/providers/common/ai/durable/openai_agents.py:
##########
@@ -345,24 +365,48 @@ async def get_response(
)
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(
system_instructions, input, model_settings, tools,
output_schema, handoffs, tracing, **kwargs
)
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. A position
+ that replays is counted in ``stats.replayed`` when reserved, so one
whose tool never runs
+ is over-counted.
+ """
+ durable_names = {tool.name for tool in tools if _is_durable_tool(tool)}
Review Comment:
`tools` here comes from `Agent.get_all_tools()`, which includes the agent's
MCP tools. The SDK builds those as plain `FunctionTool`s with a bool
`needs_approval` under the default policy, so they pass `_is_durable_tool`, but
`_get_durable_agent` only wraps `agent.tools` and nothing ever takes their
reservation. Same for the tools of an agent reached through `handoff()` whose
model is given by name. Each of those calls now takes a journal position it
never fills (and adds to the leak above), where before it took none. Could this
filter on the copies `_get_durable_tool` made, say with a marker on the copy,
rather than on the tool's type?
##########
providers/common/ai/tests/unit/common/ai/durable/test_openai_agents.py:
##########
@@ -519,3 +625,232 @@ def
test_outside_a_task_it_runs_the_agent_as_runner_run_does(self):
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,
+ **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:
+
+ @tool_input_guardrail
+ async def wait(data):
+ await asyncio.sleep(guardrail_delays[data.context.tool_call_id])
+ 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_finished_run_does_not_keep_its_reserved_tool_steps(self):
+ gc.collect()
+ kept_before = len(_RESERVED_TOOL_STEPS)
+
+ _attempt_parallel(MemoryStorage(), _ParallelModel(), _SleepingTools())
+ gc.collect()
+
+ assert len(_RESERVED_TOOL_STEPS) == kept_before
+
+ @pytest.mark.parametrize("source", ["run-hook", "agent-hook",
"tool-input-guardrail"])
+ def
test_parallel_tool_calls_replay_when_what_runs_before_them_takes_a_different_time(self,
source):
+ """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), {"call-a": 0.01,
"call-b": 0.1})
+ 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:
Review Comment:
These five pass against the `openai_agents.py` on `main` and cover paths
this PR doesn't change (streaming pass-through, the two handoff shapes,
approval). Removing the new `durable_names` filter still leaves the handoff
tests green. Are they meant to be part of this PR? If so, a line in the
description saying why would help; otherwise a separate PR keeps this one about
the reservation.
##########
providers/common/ai/src/airflow/providers/common/ai/durable/openai_agents.py:
##########
@@ -103,6 +105,13 @@
# 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")
+# Tool steps a model response claimed, by tool call id, until the tool takes
its own. The tools of
+# one response start in model order but reach their wrapper after guardrails
and hooks of varying
+# duration, so claiming there would number them by completion order.
+_RESERVED_TOOL_STEPS: weakref.WeakKeyDictionary[DurableRun, dict[str,
JournalStep]] = (
Review Comment:
`WeakKeyDictionary` only holds the key weakly. The value holds
`JournalStep`s, and `JournalStep.durable_run` is that same key, so while any
reservation is left unused this entry keeps the run, its journal and its
storage alive for the rest of the process. That happens whenever a reserved
call never reaches the wrapper, e.g. a `needs_approval=True` tool
(`test_a_tool_that_needs_approval_is_not_invoked_or_recorded` does exactly
this) or a tool input guardrail that rejects the call.
`test_a_finished_run_does_not_keep_its_reserved_tool_steps` only passes because
both calls get used.
Could `run_durably` own a plain dict and pass it down to the model provider
and the tool wrappers, or at least pop its entry in a `finally`? Then the
weakref and the gc test can go.
##########
providers/common/ai/tests/unit/common/ai/durable/test_openai_agents.py:
##########
@@ -519,3 +625,232 @@ def
test_outside_a_task_it_runs_the_agent_as_runner_run_does(self):
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,
+ **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:
+
+ @tool_input_guardrail
+ async def wait(data):
+ await asyncio.sleep(guardrail_delays[data.context.tool_call_id])
+ 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_finished_run_does_not_keep_its_reserved_tool_steps(self):
+ gc.collect()
+ kept_before = len(_RESERVED_TOOL_STEPS)
+
+ _attempt_parallel(MemoryStorage(), _ParallelModel(), _SleepingTools())
+ gc.collect()
+
+ assert len(_RESERVED_TOOL_STEPS) == kept_before
+
+ @pytest.mark.parametrize("source", ["run-hook", "agent-hook",
"tool-input-guardrail"])
+ def
test_parallel_tool_calls_replay_when_what_runs_before_them_takes_a_different_time(self,
source):
+ """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), {"call-a": 0.01,
"call-b": 0.1})
Review Comment:
Call a is the faster one on this first attempt, so claiming at invocation
already numbers the calls in model order. Deleting the `_reserve_tool_steps`
call after a live response (`openai_agents.py` line 389) leaves all 44 tests
green; only the replay-side reservation is pinned. Could this also run with the
first attempt's delays reversed (b fast first, then a fast on the retry)? That
case fails without line 389 and passes with it. I'd keep the current order too,
since it's the one that catches losing the replay-side call.
##########
providers/common/ai/tests/unit/common/ai/durable/test_openai_agents.py:
##########
@@ -519,3 +625,232 @@ def
test_outside_a_task_it_runs_the_agent_as_runner_run_does(self):
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,
+ **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:
+
+ @tool_input_guardrail
+ async def wait(data):
+ await asyncio.sleep(guardrail_delays[data.context.tool_call_id])
+ 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:
Review Comment:
I don't think any of these run in CI. `openai-agents` isn't a dependency of
the provider, so the regular provider jobs skip the file at the `importorskip`,
and the openai-agents agent framework job only runs what's in `TEST_PATHS` in
`scripts/in_container/run_agent_framework_tests.sh`, which doesn't list
`durable/test_openai_agents.py`. That job on this PR collected 187 items and
none from this file. The gap predates this PR, but could you add the file to
`TEST_PATHS` here so the fix's tests actually run?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
