You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "kaxil (via GitHub)" <gi...@apache.org> on 2026/10/14 14:16:44 UTC
[PR] Allow durable agent execution with a CodeMode capability [airflow]
kaxil opened a new pull request, #74735:
URL: https://github.com/apache/airflow/pull/74735
`AgentOperator(durable=True)` refused a `CodeMode` capability with `ValueError("durable=True cannot be used with a CodeMode capability.")`. This PR removes that error, so an agent that uses code mode can retry from its last completed step like any other durable agent.
## How a retry replays code mode
`run_code` is not a journal step of its own. On retry each snippet runs again in Monty, and every tool call the snippet makes replays from the journal one by one, the same way pydantic-ai-harness handles Temporal. The REPL keeps state between snippets, so a recorded `run_code` result could not stand in for re-running it.
That alone did not work, because of the fingerprint. `CodeMode` puts a fresh `ToolReturnPart` for every nested call in the `run_code` return's `metadata["tool_returns"]`, and each one carries a new `timestamp`. Those timestamps were hashed into the next model request's fingerprint, so a retry diverged at the first model request after any `run_code` call and ran the rest of the run live.
## Design rationale
**The fingerprint drops only those timestamps.** `_strip_part` removes `timestamp` from each nested return, and only on a tool return whose metadata carries `code_mode: True`. That is the shape the harness writes and documents. The rest of `metadata` stays in the hash, because pydantic-ai keeps routing state there (`FallbackModel`'s continuation pin under `__pydantic_ai__`), and the existing test that metadata changes the fingerprint still passes. The nested calls' tool names, arguments, ids and results are all still hashed.
**A reordered `asyncio.gather` runs live; it never replays the wrong result.** Code mode numbers a snippet's nested call ids in the order the snippet makes the calls, and each call's journal fingerprint includes its id and arguments. When concurrent calls reach the journal in another order on retry, the first one out of place mismatches, and it and every later step run live.
**`eager` and `speculate` are not rejected.** They run tool calls while the model is still writing the snippet, which a positional journal cannot replay. The harness already turns both off whenever a `BaseDurabilityCapability` is in durable context, and `AirflowDurability` is one. That check exists in every harness release that has these options: `eager` from 0.31, `speculate` from 0.33. Refusing the options at parse time would reject a configuration that already runs safely. The new tests pin it: without durable, eager and speculation start calls early; with durable, they don't.
**The operator builds no `CodeMode` of its own** (`code_mode=True` went in #74312), so no stable capability id was needed. A user's `CodeMode()` stays anonymous, which #74314 and #74423 already made safe by taking `capability_id` and `deferred_capability_ids` out of the fingerprint.
**The `code-mode` extra floor goes from 0.24.0 to 0.35.0, and the extra joins the `dev` group.** Harness releases before 0.35.0 accept pydantic-monty 1.x but cannot run code with it. The lowest-dependencies job resolved harness 0.24.0 with monty 1.1.0, and every `run_code` failed. 0.35.0 is the first release that caps monty below 1.0. The lock moves harness from 0.34.0 to 0.35.0 and nothing else. The `dev` entry makes the new real-`CodeMode` tests run in CI instead of skipping, as the comment above it already does for `openai` and `anthropic`.
## Gotchas
The snippet's own code is not recorded, so anything it does outside a tool call happens again on retry: a write to a `mount` directory, a call to an `os_access` handler. `durable_execution.rst` now covers this in a new "Code mode" section.
## Evidence
The end-to-end run used a real `AgentOperator(durable=True, capabilities=[CodeMode()])` in breeze with `retries=1`, with a scripted `FunctionModel` standing in for the model. The run's first snippet gathers `fetch('a')` and `fetch('b')`, and its second calls `fetch('c')`. Try 1 fails on the third model request with a 503.
Try 2 without the fingerprint change replays one model request and two tool calls, then diverges at the second model request:
```text
Durable: the run took a different path from the previous attempt; this step and every step after it run again position=3 step=analyst__model.request reason=the request differs from the previous attempt's (prompt, model, settings, tools, arguments or conversation so far)
LIVE model request 2 on try 2
LIVE tool call: fetch(name='c') on try 2
LIVE model request 3 on try 2
Durable: replayed 3 steps (1 model, 2 tool, 0 other), recorded 3 new steps (2 model, 1 tool, 0 other)
```
Try 2 with it replays everything up to the failure and calls the model only for the request that failed:
```text
LIVE model request 3 on try 2
run_code returned [['a-data', 'b-data'], 'c-data']
Durable: replayed 5 steps (2 model, 3 tool, 0 other), recorded 1 new steps (1 model, 0 tool, 0 other)
```
With `fetch('a')` held back on retry until `fetch('b')` has started, the reordered calls run live and still return the right results:
```text
Durable: the run took a different path from the previous attempt; this step and every step after it run again position=1 step=analyst__function_toolset__warehouse.call_tool:fetch reason=the request differs from the previous attempt's (prompt, model, settings, tools, arguments or conversation so far)
LIVE tool call: fetch(name='b') on try 2
LIVE tool call: fetch(name='a') on try 2
...
run_code returned [['a-data', 'b-data'], 'c-data']
```
`test_a_retry_replays_the_steps_after_a_run_code_call` fails with the fingerprint change reverted (`{'model': 1, 'tool': 2}` replayed instead of `{'model': 2, 'tool': 3}`). The new tests pass on harness 0.34.0, 0.35.0 (the new floor, also under lowest-direct resolution) and 0.54.0.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow durable agent execution with a CodeMode capability [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #74735:
URL: https://github.com/apache/airflow/pull/74735#discussion_r4283246612
##########
providers/common/ai/tests/unit/common/ai/operators/test_agent.py:
##########
@@ -1474,19 +1450,20 @@ def test_code_mode_capability_turns_off_tool_approval(self, fake_harness, capabi
def test_harness_without_the_code_mode_extra_does_not_break_other_agents(
self, harness_without_code_mode_extra
):
- # Neither the durable check at construction nor the approval check at run time may
- # touch the harness root, whose CodeMode export would raise ImportError here.
- AgentOperator(task_id="t", prompt="p", llm_conn_id="llm", durable=True, capabilities=[Thinking()])
+ # The approval check at run time must not touch the harness root, whose CodeMode
+ # export would raise ImportError here.
op = AgentOperator(task_id="u", prompt="p", llm_conn_id="llm", capabilities=[Thinking()])
assert op._supports_tool_approval() is True
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="Per-tool approval needs Airflow >= 3.3")
def test_capability_function_is_left_for_the_run_to_resolve(self, fake_harness):
Review Comment:
Not deliberate, thanks. Split it in 9583d595a94: `test_capability_function_is_left_for_the_run_to_resolve` has `durable=True` again and runs on every Airflow version, and the `_supports_tool_approval()` assertion is now its own 3.3+ test.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow durable agent execution with a CodeMode capability [airflow]
Posted by "pankajastro (via GitHub)" <gi...@apache.org>.
pankajastro commented on code in PR #74735:
URL: https://github.com/apache/airflow/pull/74735#discussion_r4282523256
##########
providers/common/ai/tests/unit/common/ai/operators/test_agent.py:
##########
@@ -1474,19 +1450,20 @@ def test_code_mode_capability_turns_off_tool_approval(self, fake_harness, capabi
def test_harness_without_the_code_mode_extra_does_not_break_other_agents(
self, harness_without_code_mode_extra
):
- # Neither the durable check at construction nor the approval check at run time may
- # touch the harness root, whose CodeMode export would raise ImportError here.
- AgentOperator(task_id="t", prompt="p", llm_conn_id="llm", durable=True, capabilities=[Thinking()])
+ # The approval check at run time must not touch the harness root, whose CodeMode
+ # export would raise ImportError here.
op = AgentOperator(task_id="u", prompt="p", llm_conn_id="llm", capabilities=[Thinking()])
assert op._supports_tool_approval() is True
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="Per-tool approval needs Airflow >= 3.3")
def test_capability_function_is_left_for_the_run_to_resolve(self, fake_harness):
Review Comment:
This lost `durable=True` and is now skipped below Airflow 3.3, so it asserts nothing there. Was that deliberate? Could the `op.capabilities == [build]` check stay unskipped and only the `_supports_tool_approval()` assert be gated?
---
Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting
##########
providers/common/ai/tests/unit/common/ai/durable/test_code_mode.py:
##########
@@ -0,0 +1,230 @@
+# 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.
+"""
+AirflowDurability with pydantic-ai-harness's ``CodeMode``, through a real agent loop.
+
+``run_code`` is not a step of its own: a retry runs each snippet again, and the tool
+calls the snippet makes replay from the journal one by one.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import json
+from collections.abc import AsyncIterator
+from typing import Any
+
+import pytest
+from pydantic_ai import Agent, RunContext
+from pydantic_ai.capabilities import AbstractCapability
+from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart, ToolCallPart
+from pydantic_ai.models.function import AgentInfo, DeltaToolCall, DeltaToolCalls, FunctionModel
+from pydantic_ai.toolsets import FunctionToolset
+
+from airflow.providers.common.ai.durable import AirflowDurability
+from airflow.providers.common.ai.durable.journal import DurableJournal, journal_scope
+
+from unit.common.ai.durable.test_capability import Calls, attempt, responses_so_far, tool_returns
+
+code_mode = pytest.importorskip("pydantic_ai_harness.code_mode")
Review Comment:
`dev` now installs the harness, so should this be a plain import? A broken install would skip these tests and CI would stay green.
---
Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting
##########
providers/common/ai/tests/unit/common/ai/durable/test_code_mode.py:
##########
@@ -0,0 +1,230 @@
+# 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.
+"""
+AirflowDurability with pydantic-ai-harness's ``CodeMode``, through a real agent loop.
+
+``run_code`` is not a step of its own: a retry runs each snippet again, and the tool
+calls the snippet makes replay from the journal one by one.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import json
+from collections.abc import AsyncIterator
+from typing import Any
+
+import pytest
+from pydantic_ai import Agent, RunContext
+from pydantic_ai.capabilities import AbstractCapability
+from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart, ToolCallPart
+from pydantic_ai.models.function import AgentInfo, DeltaToolCall, DeltaToolCalls, FunctionModel
+from pydantic_ai.toolsets import FunctionToolset
+
+from airflow.providers.common.ai.durable import AirflowDurability
+from airflow.providers.common.ai.durable.journal import DurableJournal, journal_scope
+
+from unit.common.ai.durable.test_capability import Calls, attempt, responses_so_far, tool_returns
+
+code_mode = pytest.importorskip("pydantic_ai_harness.code_mode")
+CodeMode = code_mode.CodeMode
+
+GATHER = "import asyncio\na, b = await asyncio.gather(fetch(name='a'), fetch(name='b'))\n[a, b]"
+
+
+def fetch_toolset(calls: Calls, started: dict[str, asyncio.Event] | None = None) -> FunctionToolset[Any]:
+ toolset = FunctionToolset(id="db")
+
+ @toolset.tool_plain
+ async def fetch(name: str) -> str:
+ """Fetch a record."""
+ calls.bump(f"fetch:{name}")
+ if started is not None and name in started:
+ started[name].set()
+ return f"{name}-data"
+
+ return toolset
+
+
+def run_code_then_answer(calls: Calls, snippets: list[str], *, fail_final: list[bool]):
+ """A model that sends each snippet to ``run_code`` in turn, then answers with the last result."""
+
+ def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+ calls.bump("model")
+ if (sent := responses_so_far(messages)) < len(snippets):
+ return ModelResponse(parts=[ToolCallPart("run_code", {"code": snippets[sent]})])
+ if fail_final[0]:
+ raise RuntimeError("worker died")
+ return ModelResponse(parts=[TextPart(f"answer: {tool_returns(messages)[-1].content}")])
+
+ return model_fn
+
+
+def streamed(model_fn):
+ """The streaming form of ``model_fn``, for runs a capability streams (``eager``, ``speculate``)."""
+
+ async def stream_fn(messages: list[ModelMessage], info: AgentInfo) -> AsyncIterator[str | DeltaToolCalls]:
+ for part in model_fn(messages, info).parts:
+ if isinstance(part, ToolCallPart):
+ yield {0: DeltaToolCall(name=part.tool_name, json_args=part.args_as_json_str())}
+ elif isinstance(part, TextPart):
+ yield part.content
+
+ return stream_fn
+
+
+class _HoldBack(AbstractCapability[Any]):
+ """Holds one tool call back until another has started, so ``asyncio.gather`` starts them in another order."""
+
+ def __init__(self, held_back: str, until_started: asyncio.Event) -> None:
+ self.held_back = held_back
+ self.until_started = until_started
+
+ async def before_tool_execute(self, ctx: RunContext[Any], *, call, tool_def, args):
+ if args.get("name") == self.held_back:
+ await asyncio.wait_for(self.until_started.wait(), timeout=5)
+ return args
+
+
+class TestReplay:
+ @pytest.mark.asyncio
+ async def test_a_retry_replays_the_steps_after_a_run_code_call(self, memory_storage):
+ """
+ The snippet runs again and its nested calls replay, so the ``run_code`` return the
+ next model request sees differs only in the timestamps of those nested calls.
+ """
+ calls = Calls()
+ fail_final = [True]
+
+ def build() -> Agent[None, str]:
+ return Agent(
+ FunctionModel(
+ run_code_then_answer(calls, [GATHER, "await fetch(name='c')"], fail_final=fail_final)
+ ),
+ name="analyst",
+ toolsets=[fetch_toolset(calls)],
+ capabilities=[CodeMode(), AirflowDurability()],
+ )
+
+ with pytest.raises(RuntimeError, match="worker died"):
+ await attempt(memory_storage, build())
+ fail_final[0] = False
+ with journal_scope(journal := DurableJournal(memory_storage)):
+ result = await build().run("go")
+
+ assert result.output == "answer: c-data"
+ assert dict(journal.stats.replayed) == {"model": 2, "tool": 3}
+ assert dict(journal.stats.recorded) == {"model": 1}
+ assert (calls["fetch:a"], calls["fetch:b"], calls["fetch:c"]) == (1, 1, 1)
+
+ @pytest.mark.asyncio
+ async def test_nested_calls_started_in_another_order_run_live(self, memory_storage):
+ """
+ A nested call's fingerprint holds its arguments and its call id, which code mode
+ numbers in the order the snippet makes the calls. A call that reaches the journal
+ in another position never replays the result recorded for its neighbour.
+ """
+ calls = Calls()
+ fail_final = [True]
+
+ def build(*capabilities: AbstractCapability[Any], started=None) -> Agent[None, str]:
+ return Agent(
+ FunctionModel(run_code_then_answer(calls, [GATHER], fail_final=fail_final)),
+ name="analyst",
+ toolsets=[fetch_toolset(calls, started)],
+ capabilities=[CodeMode(), *capabilities, AirflowDurability()],
+ )
+
+ with pytest.raises(RuntimeError, match="worker died"):
+ await attempt(memory_storage, build())
+ fail_final[0] = False
+ b_started = asyncio.Event()
+ with journal_scope(journal := DurableJournal(memory_storage)):
+ result = await build(_HoldBack("a", until_started=b_started), started={"b": b_started}).run("go")
+
+ assert result.output == "answer: ['a-data', 'b-data']"
+ assert dict(journal.stats.replayed) == {"model": 1}
+ assert (calls["fetch:a"], calls["fetch:b"]) == (2, 2)
+
+
+class TestEarlyExecution:
+ """CodeMode's ``eager`` and ``speculate`` run tool calls ahead of the snippet; not while durable."""
+
+ @pytest.mark.parametrize("durable", [False, True])
+ @pytest.mark.asyncio
+ async def test_speculation_is_off_while_durable(self, memory_storage, durable):
+ calls = Calls()
+ speculating = CodeMode(speculate=["fetch"])
+ model_fn = run_code_then_answer(calls, ["await fetch(name='a')"], fail_final=[False])
+ agent = Agent(
+ FunctionModel(model_fn, stream_function=streamed(model_fn)),
+ name="analyst",
+ toolsets=[fetch_toolset(calls)],
+ capabilities=[speculating, AirflowDurability()],
+ )
+
+ if durable:
+ await attempt(memory_storage, agent)
+ else:
+ await agent.run("go")
+
+ assert (speculating.speculation_stats.launched > 0) is not durable
+ assert calls["fetch:a"] == 1
+
+ @pytest.mark.parametrize("durable", [False, True])
+ @pytest.mark.asyncio
+ async def test_eager_execution_is_off_while_durable(self, memory_storage, durable):
+ calls = Calls()
+ seen_mid_stream: list[int] = []
+ a_started = asyncio.Event()
+ code = "a = await fetch(name='a')\nb = await fetch(name='b')\n[a, b]"
+ args = json.dumps({"code": code})
+ # A statement counts as complete once the next one has streamed in full.
+ first_statement = args.index("\\n", args.index("\\n") + 2) + 2
+
+ async def stream_fn(
+ messages: list[ModelMessage], info: AgentInfo
+ ) -> AsyncIterator[str | DeltaToolCalls]:
+ if responses_so_far(messages):
+ yield "done"
+ return
+ yield {0: DeltaToolCall(name="run_code", json_args=args[:first_statement], tool_call_id="rc")}
+ # Give an eager pump time to run the first statement before the call completes.
+ with contextlib.suppress(TimeoutError):
+ await asyncio.wait_for(a_started.wait(), timeout=1)
Review Comment:
In the durable case nothing sets `a_started`, so the test always waits the full 1 s. Could it use a shorter timeout when `durable`?
---
Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting
##########
providers/common/ai/docs/durable_execution.rst:
##########
@@ -193,6 +193,33 @@ In a plain ``@task``, which has no such moment, the capability deletes a run's j
as soon as that run succeeds: after a failure, only the run that failed replays, and
the agents that ran before it run again.
+.. _durable-code-mode:
+
+Code mode
+---------
+
+With pydantic-ai-harness's ``CodeMode`` (:doc:`code_mode`), the model sends Python to one
+``run_code`` tool, and the snippet calls the agent's tools as functions. ``run_code`` is
+not a step of its own: the retry runs each snippet again, and the tool calls the snippet
+makes replay from the journal one by one, so a snippet that called three tools is three
+recorded steps. When the snippet returns what it returned before, the model request after
+it replays too.
+
+A call the snippet makes replays only when the call at that position in the journal has
+the same tool, arguments and call id. Code mode numbers a snippet's calls in the order the
+snippet makes them, so when calls under ``asyncio.gather`` reach the journal in another
+order on the retry, the first call out of place and every step after it run live. A call
+never gets the result recorded for another one.
+
+The snippet's own code is not recorded. Whatever it does outside a tool call happens again
Review Comment:
Replayed results are JSON-decoded, so `b"hi"` is `b'hi'` live but `'hi'` on replay. Tuple, datetime and Decimal were identical. Is that worth naming here?
---
Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting
##########
providers/common/ai/pyproject.toml:
##########
@@ -127,7 +127,9 @@ dependencies = [
# Provides the CodeMode capability for AgentOperator(capabilities=[CodeMode()]).
# pydantic-ai-harness is pre-1.0; an opt-in extra so its churn never breaks the
# base provider install.
-"code-mode" = ["pydantic-ai-harness[codemode]>=0.24.0"]
+# 0.35.0 is the first release that caps pydantic-monty below 1.0; earlier releases also
+# accept monty 1.x, which they cannot run code with.
+"code-mode" = ["pydantic-ai-harness[codemode]>=0.35.0"]
Review Comment:
Harness 0.35.0 requires `pydantic-ai-slim>=2.44.0`, but the base floor is `>=2.36.0`. Should the base floor move, or is the gap intentional?
---
Drafted-by: Claude Code (Sonnet 5.5); reviewed by @pankajastro before posting
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow durable agent execution with a CodeMode capability [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil merged PR #74735:
URL: https://github.com/apache/airflow/pull/74735
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Allow durable agent execution with a CodeMode capability [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #74735:
URL: https://github.com/apache/airflow/pull/74735#discussion_r4283248644
##########
providers/common/ai/tests/unit/common/ai/durable/test_code_mode.py:
##########
@@ -0,0 +1,230 @@
+# 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.
+"""
+AirflowDurability with pydantic-ai-harness's ``CodeMode``, through a real agent loop.
+
+``run_code`` is not a step of its own: a retry runs each snippet again, and the tool
+calls the snippet makes replay from the journal one by one.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import json
+from collections.abc import AsyncIterator
+from typing import Any
+
+import pytest
+from pydantic_ai import Agent, RunContext
+from pydantic_ai.capabilities import AbstractCapability
+from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart, ToolCallPart
+from pydantic_ai.models.function import AgentInfo, DeltaToolCall, DeltaToolCalls, FunctionModel
+from pydantic_ai.toolsets import FunctionToolset
+
+from airflow.providers.common.ai.durable import AirflowDurability
+from airflow.providers.common.ai.durable.journal import DurableJournal, journal_scope
+
+from unit.common.ai.durable.test_capability import Calls, attempt, responses_so_far, tool_returns
+
+code_mode = pytest.importorskip("pydantic_ai_harness.code_mode")
Review Comment:
Yes, switched to a plain import in 9583d595a94. Every CI job has the harness. The CI image syncs the provider `dev` group, which now includes `pydantic-ai-harness==0.35.0`. The compat jobs only uninstall editable packages and the installed providers before installing the target Airflow. The lowest-deps job installs `--all-extras`.
##########
providers/common/ai/pyproject.toml:
##########
@@ -127,7 +127,9 @@ dependencies = [
# Provides the CodeMode capability for AgentOperator(capabilities=[CodeMode()]).
# pydantic-ai-harness is pre-1.0; an opt-in extra so its churn never breaks the
# base provider install.
-"code-mode" = ["pydantic-ai-harness[codemode]>=0.24.0"]
+# 0.35.0 is the first release that caps pydantic-monty below 1.0; earlier releases also
+# accept monty 1.x, which they cannot run code with.
+"code-mode" = ["pydantic-ai-harness[codemode]>=0.35.0"]
Review Comment:
Intentional. The base floor stays at what the provider itself needs, and only the `code-mode` extra needs 2.44. The harness's own requirement raises it when the extra is installed. The lowest-deps job shows it working: it resolves harness 0.35.0 with pydantic-ai-slim 2.44.0, and the suite passes.
##########
providers/common/ai/tests/unit/common/ai/durable/test_code_mode.py:
##########
@@ -0,0 +1,230 @@
+# 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.
+"""
+AirflowDurability with pydantic-ai-harness's ``CodeMode``, through a real agent loop.
+
+``run_code`` is not a step of its own: a retry runs each snippet again, and the tool
+calls the snippet makes replay from the journal one by one.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import json
+from collections.abc import AsyncIterator
+from typing import Any
+
+import pytest
+from pydantic_ai import Agent, RunContext
+from pydantic_ai.capabilities import AbstractCapability
+from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart, ToolCallPart
+from pydantic_ai.models.function import AgentInfo, DeltaToolCall, DeltaToolCalls, FunctionModel
+from pydantic_ai.toolsets import FunctionToolset
+
+from airflow.providers.common.ai.durable import AirflowDurability
+from airflow.providers.common.ai.durable.journal import DurableJournal, journal_scope
+
+from unit.common.ai.durable.test_capability import Calls, attempt, responses_so_far, tool_returns
+
+code_mode = pytest.importorskip("pydantic_ai_harness.code_mode")
+CodeMode = code_mode.CodeMode
+
+GATHER = "import asyncio\na, b = await asyncio.gather(fetch(name='a'), fetch(name='b'))\n[a, b]"
+
+
+def fetch_toolset(calls: Calls, started: dict[str, asyncio.Event] | None = None) -> FunctionToolset[Any]:
+ toolset = FunctionToolset(id="db")
+
+ @toolset.tool_plain
+ async def fetch(name: str) -> str:
+ """Fetch a record."""
+ calls.bump(f"fetch:{name}")
+ if started is not None and name in started:
+ started[name].set()
+ return f"{name}-data"
+
+ return toolset
+
+
+def run_code_then_answer(calls: Calls, snippets: list[str], *, fail_final: list[bool]):
+ """A model that sends each snippet to ``run_code`` in turn, then answers with the last result."""
+
+ def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+ calls.bump("model")
+ if (sent := responses_so_far(messages)) < len(snippets):
+ return ModelResponse(parts=[ToolCallPart("run_code", {"code": snippets[sent]})])
+ if fail_final[0]:
+ raise RuntimeError("worker died")
+ return ModelResponse(parts=[TextPart(f"answer: {tool_returns(messages)[-1].content}")])
+
+ return model_fn
+
+
+def streamed(model_fn):
+ """The streaming form of ``model_fn``, for runs a capability streams (``eager``, ``speculate``)."""
+
+ async def stream_fn(messages: list[ModelMessage], info: AgentInfo) -> AsyncIterator[str | DeltaToolCalls]:
+ for part in model_fn(messages, info).parts:
+ if isinstance(part, ToolCallPart):
+ yield {0: DeltaToolCall(name=part.tool_name, json_args=part.args_as_json_str())}
+ elif isinstance(part, TextPart):
+ yield part.content
+
+ return stream_fn
+
+
+class _HoldBack(AbstractCapability[Any]):
+ """Holds one tool call back until another has started, so ``asyncio.gather`` starts them in another order."""
+
+ def __init__(self, held_back: str, until_started: asyncio.Event) -> None:
+ self.held_back = held_back
+ self.until_started = until_started
+
+ async def before_tool_execute(self, ctx: RunContext[Any], *, call, tool_def, args):
+ if args.get("name") == self.held_back:
+ await asyncio.wait_for(self.until_started.wait(), timeout=5)
+ return args
+
+
+class TestReplay:
+ @pytest.mark.asyncio
+ async def test_a_retry_replays_the_steps_after_a_run_code_call(self, memory_storage):
+ """
+ The snippet runs again and its nested calls replay, so the ``run_code`` return the
+ next model request sees differs only in the timestamps of those nested calls.
+ """
+ calls = Calls()
+ fail_final = [True]
+
+ def build() -> Agent[None, str]:
+ return Agent(
+ FunctionModel(
+ run_code_then_answer(calls, [GATHER, "await fetch(name='c')"], fail_final=fail_final)
+ ),
+ name="analyst",
+ toolsets=[fetch_toolset(calls)],
+ capabilities=[CodeMode(), AirflowDurability()],
+ )
+
+ with pytest.raises(RuntimeError, match="worker died"):
+ await attempt(memory_storage, build())
+ fail_final[0] = False
+ with journal_scope(journal := DurableJournal(memory_storage)):
+ result = await build().run("go")
+
+ assert result.output == "answer: c-data"
+ assert dict(journal.stats.replayed) == {"model": 2, "tool": 3}
+ assert dict(journal.stats.recorded) == {"model": 1}
+ assert (calls["fetch:a"], calls["fetch:b"], calls["fetch:c"]) == (1, 1, 1)
+
+ @pytest.mark.asyncio
+ async def test_nested_calls_started_in_another_order_run_live(self, memory_storage):
+ """
+ A nested call's fingerprint holds its arguments and its call id, which code mode
+ numbers in the order the snippet makes the calls. A call that reaches the journal
+ in another position never replays the result recorded for its neighbour.
+ """
+ calls = Calls()
+ fail_final = [True]
+
+ def build(*capabilities: AbstractCapability[Any], started=None) -> Agent[None, str]:
+ return Agent(
+ FunctionModel(run_code_then_answer(calls, [GATHER], fail_final=fail_final)),
+ name="analyst",
+ toolsets=[fetch_toolset(calls, started)],
+ capabilities=[CodeMode(), *capabilities, AirflowDurability()],
+ )
+
+ with pytest.raises(RuntimeError, match="worker died"):
+ await attempt(memory_storage, build())
+ fail_final[0] = False
+ b_started = asyncio.Event()
+ with journal_scope(journal := DurableJournal(memory_storage)):
+ result = await build(_HoldBack("a", until_started=b_started), started={"b": b_started}).run("go")
+
+ assert result.output == "answer: ['a-data', 'b-data']"
+ assert dict(journal.stats.replayed) == {"model": 1}
+ assert (calls["fetch:a"], calls["fetch:b"]) == (2, 2)
+
+
+class TestEarlyExecution:
+ """CodeMode's ``eager`` and ``speculate`` run tool calls ahead of the snippet; not while durable."""
+
+ @pytest.mark.parametrize("durable", [False, True])
+ @pytest.mark.asyncio
+ async def test_speculation_is_off_while_durable(self, memory_storage, durable):
+ calls = Calls()
+ speculating = CodeMode(speculate=["fetch"])
+ model_fn = run_code_then_answer(calls, ["await fetch(name='a')"], fail_final=[False])
+ agent = Agent(
+ FunctionModel(model_fn, stream_function=streamed(model_fn)),
+ name="analyst",
+ toolsets=[fetch_toolset(calls)],
+ capabilities=[speculating, AirflowDurability()],
+ )
+
+ if durable:
+ await attempt(memory_storage, agent)
+ else:
+ await agent.run("go")
+
+ assert (speculating.speculation_stats.launched > 0) is not durable
+ assert calls["fetch:a"] == 1
+
+ @pytest.mark.parametrize("durable", [False, True])
+ @pytest.mark.asyncio
+ async def test_eager_execution_is_off_while_durable(self, memory_storage, durable):
+ calls = Calls()
+ seen_mid_stream: list[int] = []
+ a_started = asyncio.Event()
+ code = "a = await fetch(name='a')\nb = await fetch(name='b')\n[a, b]"
+ args = json.dumps({"code": code})
+ # A statement counts as complete once the next one has streamed in full.
+ first_statement = args.index("\\n", args.index("\\n") + 2) + 2
+
+ async def stream_fn(
+ messages: list[ModelMessage], info: AgentInfo
+ ) -> AsyncIterator[str | DeltaToolCalls]:
+ if responses_so_far(messages):
+ yield "done"
+ return
+ yield {0: DeltaToolCall(name="run_code", json_args=args[:first_statement], tool_call_id="rc")}
+ # Give an eager pump time to run the first statement before the call completes.
+ with contextlib.suppress(TimeoutError):
+ await asyncio.wait_for(a_started.wait(), timeout=1)
Review Comment:
Done in 9583d595a94. It waits 0.5 s when durable and up to 5 s otherwise, and the non-durable case returns as soon as the eager call starts. The durable case now takes about 0.5 s.
##########
providers/common/ai/docs/durable_execution.rst:
##########
@@ -193,6 +193,33 @@ In a plain ``@task``, which has no such moment, the capability deletes a run's j
as soon as that run succeeds: after a failure, only the run that failed replays, and
the agents that ran before it run again.
+.. _durable-code-mode:
+
+Code mode
+---------
+
+With pydantic-ai-harness's ``CodeMode`` (:doc:`code_mode`), the model sends Python to one
+``run_code`` tool, and the snippet calls the agent's tools as functions. ``run_code`` is
+not a step of its own: the retry runs each snippet again, and the tool calls the snippet
+makes replay from the journal one by one, so a snippet that called three tools is three
+recorded steps. When the snippet returns what it returned before, the model request after
+it replays too.
+
+A call the snippet makes replays only when the call at that position in the journal has
+the same tool, arguments and call id. Code mode numbers a snippet's calls in the order the
+snippet makes them, so when calls under ``asyncio.gather`` reach the journal in another
+order on the retry, the first call out of place and every step after it run live. A call
+never gets the result recorded for another one.
+
+The snippet's own code is not recorded. Whatever it does outside a tool call happens again
Review Comment:
Not inside a snippet, at least on harness 0.35.0. `CodeMode` turns each nested result into JSON-compatible data before the snippet sees it, on the live run too. I tried a nested tool returning `b"hi"`, `(1, 2)`, a datetime and a `Decimal`. The snippet saw `str 'hi'`, `list [1, 2]`, `str '2026-01-01T00:00:00Z'` and `str '1.5'`, the same live and on replay. So I left the docs as they are. Where did you see `b'hi'` on the live run? If it was a tool called outside code mode, that's the general replay path rather than something this PR changes.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org