This is an automated email from the ASF dual-hosted git repository. kaxil pushed a commit to branch commonai-durable-code-mode in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 4a44fdf6f17d67835c3ffc117bff6f24b7d8029f Author: Kaxil Naik <[email protected]> AuthorDate: Wed Oct 14 13:17:53 2026 +0100 Allow durable agent execution with a CodeMode capability A retry now runs each run_code snippet again and replays the tool calls it makes from the journal. The fingerprint drops the timestamps of the nested calls CodeMode records on a run_code return, which changed on every retry and made the next model request run live. --- providers/common/ai/README.rst | 2 +- providers/common/ai/docs/capabilities.rst | 5 +- providers/common/ai/docs/code_mode.rst | 5 +- providers/common/ai/docs/durable_execution.rst | 33 ++- providers/common/ai/docs/index.rst | 2 +- providers/common/ai/docs/recovery.rst | 3 +- providers/common/ai/docs/troubleshooting.rst | 6 +- providers/common/ai/pyproject.toml | 6 +- .../providers/common/ai/durable/fingerprint.py | 40 +++- .../airflow/providers/common/ai/operators/agent.py | 12 +- .../tests/unit/common/ai/durable/test_code_mode.py | 230 +++++++++++++++++++++ .../unit/common/ai/durable/test_fingerprint.py | 59 ++++++ .../tests/unit/common/ai/operators/test_agent.py | 43 +--- uv.lock | 10 +- 14 files changed, 387 insertions(+), 69 deletions(-) diff --git a/providers/common/ai/README.rst b/providers/common/ai/README.rst index c53022d38f0..1fdef9e34cb 100644 --- a/providers/common/ai/README.rst +++ b/providers/common/ai/README.rst @@ -97,7 +97,7 @@ Extra Dependencies ``opensandbox`` ``opensandbox>=1.1.0`` ``boat`` ``boat-sdk>=1.9.3`` ``openshell`` ``openshell>=0.1.2,<0.2``, ``grpcio>=1.78.0``, ``protobuf>=6.31.1`` -``code-mode`` ``pydantic-ai-harness[codemode]>=0.24.0`` +``code-mode`` ``pydantic-ai-harness[codemode]>=0.35.0`` ``shields`` ``pydantic-ai-shields>=0.3.4`` ``skills`` ``apache-airflow-providers-git>=0.4.0``, ``pydantic-ai-skills>=1.2.0`` ``avro`` ``fastavro>=1.10.0; python_version < "3.14"``, ``fastavro>=1.12.1; python_version >= "3.14"`` diff --git a/providers/common/ai/docs/capabilities.rst b/providers/common/ai/docs/capabilities.rst index 43106a7cf78..08fc2aaf135 100644 --- a/providers/common/ai/docs/capabilities.rst +++ b/providers/common/ai/docs/capabilities.rst @@ -114,9 +114,8 @@ again. Whether a capability's work is replayed depends on where it runs: names one without it when the agent is built. An ``id`` on a ``Toolset`` capability does not reach the toolset inside it, so set it on that toolset. * - pydantic-ai-harness ``CodeMode`` - - Not allowed: the operator raises ``ValueError``. This includes a ``CodeMode`` inside a - ``CombinedCapability`` or a wrapper such as ``PrefixTools``, but not one a capability - function builds when the run starts. + - Each ``run_code`` snippet runs again, and the tool calls it makes are replayed; see + :ref:`durable-code-mode`. See :doc:`durable_execution` for how replay works. diff --git a/providers/common/ai/docs/code_mode.rst b/providers/common/ai/docs/code_mode.rst index c3abf0a500b..b4a5bb7b3b8 100644 --- a/providers/common/ai/docs/code_mode.rst +++ b/providers/common/ai/docs/code_mode.rst @@ -92,8 +92,9 @@ already made are not undone, and the model has to split the work across several agent fans out over a list, for example one tool call per row or per file, set it above the largest list you expect. -Code mode cannot be combined with ``durable=True`` (see :doc:`capabilities` for which -capabilities durable replay covers), and it turns off :doc:`tool_approval`. A tool that needs +With ``durable=True``, a retry runs each ``run_code`` snippet again and replays the tool +calls it makes, and ``eager`` and ``speculate`` are off; see :ref:`durable-code-mode`. +Code mode turns off :doc:`tool_approval`. A tool that needs approval and is called from ``run_code`` does not run: the model gets an error back and may carry on without it, so the task can still succeed. To keep such a tool out of ``run_code``, list the other tools in ``CodeMode(tools=[...])``; called directly, it fails the task. diff --git a/providers/common/ai/docs/durable_execution.rst b/providers/common/ai/docs/durable_execution.rst index 5c670045998..2d96bba3da2 100644 --- a/providers/common/ai/docs/durable_execution.rst +++ b/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 +on the retry, such as writing to a directory shared through ``mount`` or calling an +``os_access`` handler, and if what it reads that way has changed, the tool calls it makes +next differ and run live. Keep side effects in tools, whose results are recorded. + +``CodeMode``'s ``eager`` and ``speculate`` options start tool calls while the model is +still writing the snippet. Both are off while the run is durable, so a retry runs the +finished snippet the way the previous attempt did. + What the task log shows ----------------------- @@ -271,9 +298,9 @@ the step, and a retry runs it again. Limitations ----------- -- ``durable=True`` cannot be combined with ``enable_hitl_review``, a pydantic-ai-harness - ``CodeMode`` capability, a ``SandboxToolset`` or tool approval. ``AgentOperator`` raises - ``ValueError`` for the first three when the Dag is parsed. +- ``durable=True`` cannot be combined with ``enable_hitl_review``, a ``SandboxToolset`` + or tool approval. ``AgentOperator`` raises ``ValueError`` for the first two when the Dag + is parsed. - A task instance that is retried across an upgrade of this provider from a version without ``AirflowDurability`` runs its steps live once: the journal keys changed. diff --git a/providers/common/ai/docs/index.rst b/providers/common/ai/docs/index.rst index 9abac2e5f8b..f5b4aa71480 100644 --- a/providers/common/ai/docs/index.rst +++ b/providers/common/ai/docs/index.rst @@ -237,7 +237,7 @@ Extra Dependencies ``opensandbox`` ``opensandbox>=1.1.0`` ``boat`` ``boat-sdk>=1.9.3`` ``openshell`` ``openshell>=0.1.2,<0.2``, ``grpcio>=1.78.0``, ``protobuf>=6.31.1`` -``code-mode`` ``pydantic-ai-harness[codemode]>=0.24.0`` +``code-mode`` ``pydantic-ai-harness[codemode]>=0.35.0`` ``shields`` ``pydantic-ai-shields>=0.3.4`` ``skills`` ``apache-airflow-providers-git>=0.4.0``, ``pydantic-ai-skills>=1.2.0`` ``islo`` ``islo>=0.3.19`` diff --git a/providers/common/ai/docs/recovery.rst b/providers/common/ai/docs/recovery.rst index 5ccea3db064..481d1f047bd 100644 --- a/providers/common/ai/docs/recovery.rst +++ b/providers/common/ai/docs/recovery.rst @@ -160,5 +160,4 @@ What is not covered ``LLMBatchOperator`` re-attaches to its running batch (:ref:`llm-batch-reattach`). - A ``ManagedAgentToolset`` runs again on every attempt unless it is built with ``replayable=True``. Durable execution cannot be combined with ``enable_hitl_review``, - a ``CodeMode`` capability, a ``SandboxToolset`` or tool approval; see the limitations - in :doc:`durable_execution`. + a ``SandboxToolset`` or tool approval; see the limitations in :doc:`durable_execution`. diff --git a/providers/common/ai/docs/troubleshooting.rst b/providers/common/ai/docs/troubleshooting.rst index f83a43cd9f3..7072e02c07b 100644 --- a/providers/common/ai/docs/troubleshooting.rst +++ b/providers/common/ai/docs/troubleshooting.rst @@ -111,10 +111,10 @@ as a task failure. or later. Upgrade Airflow, or use ``on_uncertain="fail"`` and drop the review flags on an older Airflow version. See :doc:`approval_gates` and :doc:`hitl_review`. -``durable=True and enable_hitl_review=True cannot be used together`` / ``durable=True cannot be used with a CodeMode capability`` +``durable=True and enable_hitl_review=True cannot be used together`` Durable replay matches each step to the one at the same position in the previous - attempt, which neither a human review loop's regenerations nor the tool calls code - mode makes from inside ``run_code`` keep. Pick one. See :doc:`durable_execution`. + attempt, which a human review loop's regenerations do not keep. Pick one. See + :doc:`durable_execution`. ``durable=True needs a unique id on every FunctionToolset, such as FunctionToolset(..., id='orders')`` Durable replay records each tool call under its toolset's id, so a diff --git a/providers/common/ai/pyproject.toml b/providers/common/ai/pyproject.toml index 9aebf776494..af2619cf0a3 100644 --- a/providers/common/ai/pyproject.toml +++ b/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"] # Shield capabilities such as InputGuard, OutputGuard, ToolGuard, and CostTracking. "shields" = ["pydantic-ai-shields>=0.3.4"] # Agent Skills (agentskills.io) support. pydantic-ai-skills provides the toolset; @@ -197,6 +199,8 @@ dev = [ # needs them so the adapter tests, which build real SDK models, run instead of skipping. "openai>=2.47.0", "anthropic>=1.1.0", + # The code-mode extra, so the tests that run a real CodeMode under durable execution run. + "pydantic-ai-harness[codemode]>=0.35.0", "boat-sdk>=1.9.3", "islo>=0.3.19", "opensandbox>=1.1.0", diff --git a/providers/common/ai/src/airflow/providers/common/ai/durable/fingerprint.py b/providers/common/ai/src/airflow/providers/common/ai/durable/fingerprint.py index 3082f4df64f..48e717501e3 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/durable/fingerprint.py +++ b/providers/common/ai/src/airflow/providers/common/ai/durable/fingerprint.py @@ -33,7 +33,8 @@ longer match. Fields that pydantic-ai regenerates on every attempt (message-level ``timestamp``/``run_id``/``conversation_id``, part-level ``timestamp``, the -``capability_id`` of a tool from a capability without an ``id``) are excluded +``capability_id`` of a tool from a capability without an ``id``, the timestamps of +the nested calls code mode records on a ``run_code`` return) are excluded from the fingerprint, and set-valued request parameters are sorted or dropped. Registered secrets in the message history are masked before hashing, as they are in the tool results the journal records. Requests that cannot be serialized to JSON @@ -84,23 +85,46 @@ def _content_settings(model_settings: ModelSettings | None) -> dict[str, Any] | return content or None +def _strip_part(part: Any) -> Any: + """ + Drop a dumped part's timestamp, and those of the nested calls on a code mode return. + + pydantic-ai-harness's ``CodeMode`` does not journal ``run_code`` itself: a retry runs the + snippet again, its nested tool calls replay, and the return it builds holds a fresh + ``ToolReturnPart`` per nested call under ``metadata["tool_returns"]``, each with a new + timestamp. The rest of ``metadata`` is kept, as pydantic-ai keeps routing state there. + """ + if not isinstance(part, dict): + return part + cleaned = {k: v for k, v in part.items() if k != "timestamp"} + metadata = cleaned.get("metadata") + if ( + cleaned.get("part_kind") == "tool-return" + and isinstance(metadata, dict) + and metadata.get("code_mode") is True + and isinstance(nested := metadata.get("tool_returns"), dict) + ): + cleaned["metadata"] = { + **metadata, + "tool_returns": {call_id: _strip_part(ret) for call_id, ret in nested.items()}, + } + return cleaned + + def _strip_volatile(messages_dump: list[dict[str, Any]]) -> list[dict[str, Any]]: """ Drop per-attempt fields from a dumped message list. Only the levels pydantic-ai regenerates are touched (message-level ids and - timestamps, part-level timestamps); user data such as tool arguments is - never recursed into, so an argument legitimately named ``run_id`` still - affects the fingerprint. + timestamps, part-level timestamps, code mode's nested-call timestamps); user + data such as tool arguments is never recursed into, so an argument legitimately + named ``run_id`` still affects the fingerprint. """ stripped = [] for message in messages_dump: cleaned = {k: v for k, v in message.items() if k not in _VOLATILE_MESSAGE_KEYS} if isinstance(cleaned.get("parts"), list): - cleaned["parts"] = [ - {k: v for k, v in part.items() if k != "timestamp"} if isinstance(part, dict) else part - for part in cleaned["parts"] - ] + cleaned["parts"] = [_strip_part(part) for part in cleaned["parts"]] stripped.append(cleaned) return stripped diff --git a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py index 8bc303e6494..5648da03aa2 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py +++ b/providers/common/ai/src/airflow/providers/common/ai/operators/agent.py @@ -368,8 +368,10 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa Cannot be combined with a ``SandboxToolset`` (raises), attached or not: a replayed tool result describes a workspace state the replay did not reproduce, and the first call that misses the cache runs against - whatever the sandbox holds now. Cannot be combined with a pydantic-ai-harness - ``CodeMode`` capability (raises). + whatever the sandbox holds now. With a pydantic-ai-harness ``CodeMode`` + capability, a retry runs each ``run_code`` snippet again and replays the + tool calls it makes; ``CodeMode``'s ``eager`` and ``speculate`` options are + off while the run is durable. :param cache_prompt: When ``True`` (default), asks the provider to cache the tool definitions, system prompt and conversation so far, so the next request in the run -- and a mapped task's other instances within the @@ -560,12 +562,6 @@ class AgentOperator(CancellableAgentRunMixin, BaseOperator, HITLReviewMixin, Usa if durable and enable_hitl_review: raise ValueError("durable=True and enable_hitl_review=True cannot be used together.") - if durable and _contains_code_mode(self._declared_capabilities): - # Code mode collapses tools into one ``run_code`` tool and lets the model emit - # arbitrary Python, whose nested tool calls the journal does not see as steps of - # their own. Reject the combination rather than replay a partial ``run_code``. - raise ValueError("durable=True cannot be used with a CodeMode capability.") - if durable and any( isinstance(capability, AirflowDurability) for capability in self._declared_capabilities ): diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_code_mode.py b/providers/common/ai/tests/unit/common/ai/durable/test_code_mode.py new file mode 100644 index 00000000000..a69e3b7b400 --- /dev/null +++ b/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) + seen_mid_stream.append(calls["fetch:a"]) + yield {0: DeltaToolCall(json_args=args[first_statement:])} + + async def handler(ctx, events) -> None: + async for _ in events: + pass + + agent = Agent( + FunctionModel(stream_function=stream_fn), + name="analyst", + toolsets=[fetch_toolset(calls, started={"a": a_started})], + capabilities=[CodeMode(eager=True), AirflowDurability()], + ) + + if durable: + await attempt(memory_storage, agent, event_stream_handler=handler) + else: + await agent.run("go", event_stream_handler=handler) + + assert seen_mid_stream == [0 if durable else 1] + assert (calls["fetch:a"], calls["fetch:b"]) == (1, 1) diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_fingerprint.py b/providers/common/ai/tests/unit/common/ai/durable/test_fingerprint.py index 43d27486790..5594fc3df7b 100644 --- a/providers/common/ai/tests/unit/common/ai/durable/test_fingerprint.py +++ b/providers/common/ai/tests/unit/common/ai/durable/test_fingerprint.py @@ -26,6 +26,7 @@ from pydantic_ai.messages import ( SystemPromptPart, TextPart, ToolCallPart, + ToolReturnPart, UserPromptPart, ) from pydantic_ai.models import ModelRequestParameters @@ -220,6 +221,64 @@ class TestModelRequestFingerprint: assert fp1 != fp2 + @staticmethod + def code_mode_history(nested_timestamp, *, nested_content="a-data", code_mode=True): + """A ``run_code`` return as pydantic-ai-harness's ``CodeMode`` builds it, with one nested call.""" + nested = ToolReturnPart( + tool_name="fetch", content=nested_content, tool_call_id="rc__1", timestamp=nested_timestamp + ) + return [ + ModelRequest(parts=[UserPromptPart(content="q")]), + ModelResponse( + parts=[ToolCallPart(tool_name="run_code", args={"code": "..."}, tool_call_id="rc")] + ), + ModelRequest( + parts=[ + ToolReturnPart( + tool_name="run_code", + content="done", + tool_call_id="rc", + metadata={ + "code_mode": code_mode, + "tool_calls": {"rc__1": ToolCallPart(tool_name="fetch", tool_call_id="rc__1")}, + "tool_returns": {"rc__1": nested}, + }, + ) + ] + ), + ] + + @pytest.mark.parametrize("code_mode", [True, False]) + def test_nested_call_timestamps_count_only_outside_code_mode(self, code_mode): + """ + A retry runs a code mode snippet again, and its replayed nested calls get new + timestamps. Without code mode's marker, the metadata is someone else's and is + hashed as it is. + """ + fp1, fp2 = ( + fingerprint_model_request( + "m", + self.code_mode_history( + datetime.datetime(2026, 1, day, tzinfo=datetime.UTC), code_mode=code_mode + ), + None, + ModelRequestParameters(), + ) + for day in (1, 2) + ) + + assert (fp1 == fp2) is code_mode + + def test_changes_with_code_mode_nested_call_content(self): + """Only the nested calls' timestamps are dropped; what they returned still counts.""" + t = datetime.datetime(2026, 1, 1, tzinfo=datetime.UTC) + fp1 = fingerprint_model_request("m", self.code_mode_history(t), None, ModelRequestParameters()) + fp2 = fingerprint_model_request( + "m", self.code_mode_history(t, nested_content="b-data"), None, ModelRequestParameters() + ) + + assert fp1 != fp2 + def test_volatile_keys_inside_user_data_are_not_stripped(self): """Only pydantic-ai's own message/part fields are volatile; a tool argument legitimately named run_id must still affect the fingerprint.""" diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py index 27f5cda1ba2..fb22b27fd9b 100644 --- a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py +++ b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py @@ -1426,36 +1426,12 @@ class TestAgentOperatorCapabilities: assert "Thinking" not in json.dumps(serialized, default=str) - @pytest.mark.parametrize( - "wrap", - [ - pytest.param(lambda code_mode: code_mode, id="direct"), - pytest.param(lambda code_mode: CombinedCapability([Thinking(), code_mode]), id="combined"), - pytest.param(lambda code_mode: code_mode.prefix_tools("sandboxed"), id="prefixed"), - pytest.param( - lambda code_mode: CombinedCapability([PrefixTools(wrapped=code_mode, prefix="sandboxed")]), - id="prefixed-in-combined", - ), - ], - ) - @pytest.mark.parametrize("passed_as", ["capabilities", "agent_params"]) - def test_durable_refuses_a_code_mode_capability(self, fake_harness, wrap, passed_as): - capabilities = [wrap(_FakeCodeMode())] - kwargs = ( - {"capabilities": capabilities} - if passed_as == "capabilities" - else {"agent_params": {"capabilities": capabilities}} - ) - - with pytest.raises(ValueError, match="CodeMode capability"): - AgentOperator(task_id="t", prompt="p", llm_conn_id="llm", durable=True, **kwargs) + def test_durable_accepts_a_code_mode_capability(self, fake_harness): + code_mode = _FakeCodeMode() - def test_durable_accepts_capabilities_without_code_mode(self, fake_harness): - thinking = Thinking() + op = AgentOperator(task_id="t", prompt="p", llm_conn_id="llm", durable=True, capabilities=[code_mode]) - op = AgentOperator(task_id="t", prompt="p", llm_conn_id="llm", durable=True, capabilities=[thinking]) - - assert op.capabilities == [thinking] + assert op.capabilities == [code_mode] @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="Per-tool approval needs Airflow >= 3.3") @pytest.mark.parametrize( @@ -1474,19 +1450,20 @@ class TestAgentOperatorCapabilities: 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): def build(ctx): - return Thinking() + raise AssertionError("resolved before the run started") - op = AgentOperator(task_id="t", prompt="p", llm_conn_id="llm", durable=True, capabilities=[build]) + op = AgentOperator(task_id="t", prompt="p", llm_conn_id="llm", capabilities=[build]) + assert op._supports_tool_approval() is True assert op.capabilities == [build] diff --git a/uv.lock b/uv.lock index 79c86d58b35..a9200cbb334 100644 --- a/uv.lock +++ b/uv.lock @@ -4671,6 +4671,7 @@ dev = [ { name = "opensandbox" }, { name = "openshell" }, { name = "protobuf" }, + { name = "pydantic-ai-harness", extra = ["codemode"] }, { name = "pydantic-ai-skills" }, { name = "pydantic-ai-slim", extra = ["mcp"] }, { name = "sqlglot" }, @@ -4707,7 +4708,7 @@ requires-dist = [ { name = "protobuf", marker = "extra == 'openshell'", specifier = ">=6.31.1" }, { name = "pyarrow", marker = "python_full_version >= '3.14' and extra == 'parquet'", specifier = ">=22.0.0" }, { name = "pyarrow", marker = "python_full_version < '3.14' and extra == 'parquet'", specifier = ">=18.0.0" }, - { name = "pydantic-ai-harness", extras = ["codemode"], marker = "extra == 'code-mode'", specifier = ">=0.24.0" }, + { name = "pydantic-ai-harness", extras = ["codemode"], marker = "extra == 'code-mode'", specifier = ">=0.35.0" }, { name = "pydantic-ai-shields", marker = "extra == 'shields'", specifier = ">=0.3.4" }, { name = "pydantic-ai-skills", marker = "extra == 'skills'", specifier = ">=1.2.0" }, { name = "pydantic-ai-slim", specifier = ">=2.36.0" }, @@ -4747,6 +4748,7 @@ dev = [ { name = "opensandbox", specifier = ">=1.1.0" }, { name = "openshell", specifier = ">=0.1.2,<0.2" }, { name = "protobuf", specifier = ">=6.31.1" }, + { name = "pydantic-ai-harness", extras = ["codemode"], specifier = ">=0.35.0" }, { name = "pydantic-ai-skills", specifier = ">=1.2.0" }, { name = "pydantic-ai-slim", extras = ["mcp"], specifier = ">=2.36.0" }, { name = "sqlglot", specifier = ">=30.0.0" }, @@ -20044,7 +20046,7 @@ email = [ [[package]] name = "pydantic-ai-harness" -version = "0.34.0" +version = "0.35.0" source = { registry = "https://pypi.org/simple"; } dependencies = [ { name = "genai-prices" }, @@ -20052,9 +20054,9 @@ dependencies = [ { name = "json-repair" }, { name = "pydantic-ai-slim" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/16/32/e0759986b8bce033e17dac91e616e096f5ab99c41a6b55588609b29c9560/pydantic_ai_harness-0.34.0.tar.gz";, hash = "sha256:049ce17739b8e69067e9cdc8120fbbcc5f3c6c6bfb9cf2457a5c0b3ca8f046bf", size = 2991435, upload-time = "2026-09-23T02:50:57.536Z" } +sdist = { url = "https://files.pythonhosted.org/packages/ad/b6/d30279b62971960c2f740d7b8902a1f93f6de98cc787d56185a501a6905f/pydantic_ai_harness-0.35.0.tar.gz";, hash = "sha256:4bd643d09c9be349132e4eb2cbe727c8e78a3658f2f6c7a48b6411444cdd2395", size = 3208936, upload-time = "2026-09-25T05:03:26.736Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/18/e6/227553f3cc9b2b2d1350fd85ceba1e17a2bae6cfff2c6407f4bd4ef6e97e/pydantic_ai_harness-0.34.0-py3-none-any.whl";, hash = "sha256:df6ca8dc6b115d75c6236ed2c9f8b708ede02efbe2de8b0b21a7ec57a47b9393", size = 928236, upload-time = "2026-09-23T02:50:55.879Z" }, + { url = "https://files.pythonhosted.org/packages/a9/3f/7af10428985bee133d68a60439636450456200b4c896c77f04c9a096529b/pydantic_ai_harness-0.35.0-py3-none-any.whl";, hash = "sha256:ed1691d1045480bd9192b0dc13df5ae8a48448f2987cebb7344c11e3e319b293", size = 966091, upload-time = "2026-09-25T05:03:24.873Z" }, ] [package.optional-dependencies]
