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 3c254dbc58b Fix durable agent tasks failing when a tool returns bytes
(#76861)
3c254dbc58b is described below
commit 3c254dbc58b629b5065748a9a2c13b0a694e6706
Author: Kaxil Naik <[email protected]>
AuthorDate: Wed Nov 11 15:19:07 2026 +0000
Fix durable agent tasks failing when a tool returns bytes (#76861)
A tool that returned bytes that are not UTF-8 text, such as a PNG, failed a
durable agent
task: the journal encoded tool results with pydantic-ai's JSON codec, which
writes bytes as
UTF-8 text. Bytes that are UTF-8 text were recorded as text and replayed as
str, so the
retry showed the model something other than the base64 it first saw.
The journal now records bytes a tool returns, on their own or as values
inside dicts and
lists, as base64 and replays the same bytes. Secrets inside those bytes are
masked before
they are recorded, as they are for text. Entries recorded by earlier
versions load as before.
A result the journal still cannot encode fails the task as before.
DurableJournal takes
fail_on_unrecordable=False to log such a step and run it again on retry
instead, for a
journal whose durable execution was not requested explicitly.
Only ``bytes`` is converted to base64; ``bytearray``, ``memoryview`` and
sets of bytes are not, and the docs now say so.
From pydantic-ai 2.66, pydantic-ai records a function toolset's bytes in
a dataclass field itself, as base64, so such a result no longer fails to
record there. The docs no longer claim a fixed outcome for bytes in a
dataclass, and the test of a journal that leaves an unrecordable result
behind uses an Airflow toolset, whose results go through the journal's
own codec on every pydantic-ai version.
Also rename ``JournalStep._not_recorded`` to ``_mark_not_recorded``.
* Say what durable execution does with tool result bytes it does not convert
``bytearray``, ``memoryview``, sets of bytes and bytes in a dataclass or
Pydantic model field either cannot be encoded, which fails the task, or
come back changed, depending on the pydantic-ai version and the toolset.
The docs now say so in one place.
The test of a journal that leaves an unrecordable result behind covers
function toolsets again, with the journal's codec refusing the result,
since pydantic-ai 2.66 records that result itself.
* Scope the durable execution docs promise of a failure without retry
Only a result the journal itself cannot encode fails the task with
``DurableJournalError`` and no retry. Bytes in a Pydantic model field
returned by an Airflow toolset fail earlier, in secret masking, with an
error that is retried, so the docs no longer promise the no-retry
failure for every unconverted bytes type.
Name the ``set`` and ``frozenset`` types as code, so the docs spelling
check passes.
---
providers/common/ai/docs/durable_execution.rst | 16 +-
providers/common/ai/docs/toolsets/custom.rst | 5 +-
providers/common/ai/docs/troubleshooting.rst | 8 +-
.../providers/common/ai/durable/capability.py | 137 ++++++--
.../airflow/providers/common/ai/durable/journal.py | 54 +++-
.../unit/common/ai/durable/test_capability.py | 351 ++++++++++++++++++++-
6 files changed, 533 insertions(+), 38 deletions(-)
diff --git a/providers/common/ai/docs/durable_execution.rst
b/providers/common/ai/docs/durable_execution.rst
index 8b4cded05c5..1beaf3ef032 100644
--- a/providers/common/ai/docs/durable_execution.rst
+++ b/providers/common/ai/docs/durable_execution.rst
@@ -366,11 +366,17 @@ A toolset can opt out of replay. ``ManagedAgentToolset``
does by default
(``replayable=False``), because a managed agent acts on a system Airflow
cannot observe:
its calls run again on every attempt, while the steps around them still replay.
-A result has to survive a JSON round trip to be recorded. One that cannot be
encoded
-fails the task with ``DurableJournalError`` and no retry, since it would fail
the same
-way every time; the model could not have read it either. A write the task
state store
-refuses does not fail the task: the task log names
-the step, and a retry runs it again.
+A result has to survive a JSON round trip to be recorded. Bytes a tool
returns, on their
+own or as values anywhere inside dicts and lists, are recorded as base64, and
a retry
+replays the same bytes. Only ``bytes`` is converted. ``bytearray``,
``memoryview``, a
+``set`` or ``frozenset`` of bytes, and bytes in a field of a dataclass or
Pydantic model
+are not:
+depending on the pydantic-ai version and the toolset, such a result fails the
task, or
+it comes back changed (as text, a list or a dict), which can make the steps
after that
+tool call run again. Put bytes in a dict instead. A result the journal cannot
encode,
+such as an object with no JSON form, fails the task with
``DurableJournalError`` and no
+retry, since it would fail the same way every time. A write the task state
store
+refuses does not fail the task: the task log names the step, and a retry runs
it again.
Limitations
-----------
diff --git a/providers/common/ai/docs/toolsets/custom.rst
b/providers/common/ai/docs/toolsets/custom.rst
index 37090e292d4..ce92eb12feb 100644
--- a/providers/common/ai/docs/toolsets/custom.rst
+++ b/providers/common/ai/docs/toolsets/custom.rst
@@ -102,8 +102,9 @@ What a subclass implements
class's ``call_tool`` masks it on the way out. Implement ``execute_tool``
rather than
overriding ``call_tool``, since an override skips that masking and
``AgentOperator``
then wraps the toolset in a masking layer of its own. With
``durable=True`` a result
- has to be JSON-serializable, or the task fails with
``DurableJournalError`` and is
- not retried.
+ has to be JSON-serializable or bytes (on their own or as values inside
dicts and
+ lists), or the task fails with ``DurableJournalError`` and is not retried.
Put bytes
+ in a dict, not in a field of a dataclass or Pydantic model.
Two class attributes are optional:
diff --git a/providers/common/ai/docs/troubleshooting.rst
b/providers/common/ai/docs/troubleshooting.rst
index 545318a9392..e5867ec9532 100644
--- a/providers/common/ai/docs/troubleshooting.rst
+++ b/providers/common/ai/docs/troubleshooting.rst
@@ -160,9 +160,11 @@ Run-time errors
connection as described above.
``DurableJournalError: durable execution could not record a step's result,
because it is not JSON-serializable``
- A tool returned a value that cannot be encoded as JSON. The model could
not have read
- it either, so the task fails without retrying. Return text, numbers,
lists, dicts or a
- Pydantic model. See :doc:`durable_execution`.
+ A tool returned a value that cannot be encoded as JSON, such as an object
with no JSON
+ form or a dataclass holding bytes that are not UTF-8 text. It would fail
the same way
+ on every attempt, so the task fails without retrying. Return text,
numbers, bytes,
+ lists, dicts or a Pydantic model without bytes fields; put bytes in a dict
rather than
+ a field of your own class. See :doc:`durable_execution`.
Durable execution on Airflow below 3.3 fails with a ``ValueError`` about
``durable_cache_path``
On Airflow versions older than 3.3 the step cache lives in object storage
and
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/durable/capability.py
b/providers/common/ai/src/airflow/providers/common/ai/durable/capability.py
index 84b6ee31274..12d782d9010 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/durable/capability.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/durable/capability.py
@@ -18,6 +18,8 @@
from __future__ import annotations
+import base64
+import dataclasses
from collections.abc import Awaitable, Callable, Iterator, Mapping
from contextlib import contextmanager
from contextvars import ContextVar
@@ -29,6 +31,7 @@ from pydantic_ai.durable_exec import (
JSON_CODEC,
BaseDurabilityCapability,
CapabilityOperationId,
+ DurabilityCodec,
DurabilityEngineSpec,
DurableOperationId,
JournalCallableOperationBackend,
@@ -51,7 +54,7 @@ from pydantic_ai.toolsets.wrapper import WrapperToolset
from pydantic_core import PydanticSerializationError
from airflow.providers.common.ai.durable.fingerprint import
fingerprint_model_request, fingerprint_tool_call
-from airflow.providers.common.ai.durable.journal import current_journal
+from airflow.providers.common.ai.durable.journal import Unrecordable,
current_journal
from airflow.providers.common.ai.durable.replay_usage import
ReplayUsageLedger, is_successful_tool_payload
from airflow.providers.common.ai.exceptions import DurableJournalError
from airflow.providers.common.ai.utils.masking import mask_secrets
@@ -121,13 +124,102 @@ def _active_run(active: _ActiveRun) -> Iterator[None]:
_ACTIVE_RUN.reset(token)
-def _not_journalable(error: Exception) -> BaseException:
- # The same value fails the same way on every retry, so retrying the task
cannot help.
- return DurableJournalError(
- f"durable execution could not record a step's result, because it is
not JSON-serializable: {error}"
+# A recorded ``bytes`` value: a dict whose only key is this, holding the bytes
as base64.
+_BYTES_KEY = "__airflow_bytes__"
+
+# The kinds of pydantic-ai's recorded tool call result that hold the tool's
value in ``result``.
+_TOOL_RESULT_KINDS = frozenset({"tool_return", "tool_content_result"})
+
+
+def _tag_bytes(value: Any) -> Any:
+ """
+ Replace the ``bytes`` in a tool result with base64 the journal can hold.
+
+ The model sees bytes as base64 (pydantic-ai's ``tool_return_ta``), but
``JSON_CODEC``
+ encodes them as UTF-8 text: it fails on bytes that are not UTF-8 and turns
the rest into
+ ``str``, so the retry would show the model something else. Tagged, they
load back as the
+ same bytes. Only the parts of a result that a tool fills freely are
walked; a typed value
+ such as ``BinaryContent`` encodes its own bytes.
+ """
+ if isinstance(value, bytes):
+ return {_BYTES_KEY: base64.b64encode(value).decode("ascii")}
+ if isinstance(value, dict):
+ return {key: _tag_bytes(item) for key, item in value.items()}
+ if isinstance(value, list | tuple):
+ return [_tag_bytes(item) for item in value]
+ if isinstance(value, ToolReturn):
+ return dataclasses.replace(
+ value, return_value=_tag_bytes(value.return_value),
metadata=_tag_bytes(value.metadata)
+ )
+ if (
+ dataclasses.is_dataclass(value)
+ and not isinstance(value, type)
+ and getattr(value, "kind", None) in _TOOL_RESULT_KINDS
+ and any(attribute.name == "result" and attribute.init for attribute in
dataclasses.fields(value))
+ ):
+ # pydantic-ai's recorded tool call result, a shape its journals keep
loading across versions.
+ return dataclasses.replace(value, result=_tag_bytes(getattr(value,
"result")))
+ return value
+
+
+def _untag_bytes(payload: Any) -> Any:
+ """
+ Turn what :func:`_tag_bytes` recorded back into bytes.
+
+ A payload without tags loads as it is, and so does a dict a tool built
that only looks
+ like one, with a value that is not base64.
+ """
+ if isinstance(payload, dict):
+ if len(payload) == 1 and isinstance(encoded :=
payload.get(_BYTES_KEY), str):
+ try:
+ return base64.b64decode(encoded, validate=True)
+ except ValueError: # binascii.Error, or a value that is not ASCII
+ return payload
+ return {key: _untag_bytes(item) for key, item in payload.items()}
+ if isinstance(payload, list):
+ return [_untag_bytes(item) for item in payload]
+ return payload
+
+
+def _mask_payload(payload: Any) -> Any:
+ """Mask the secrets in an encoded payload, including inside its bytes,
which base64 hides from the masker."""
+ return _tag_bytes(mask_secrets(_untag_bytes(payload)))
+
+
+def _is_model_response(payload: Any) -> bool:
+ # Bytes are never tagged in a model response, so text the model wrote,
such as tool call
+ # arguments, is never read as a tag. The second shape is a streamed
response from
+ # pydantic-ai before 2.60, recorded with its events.
+ return isinstance(payload, dict) and (
+ payload.get("kind") == "response" or ("response" in payload and
"events" in payload)
)
+class _JournalCodec(DurabilityCodec):
+ """
+ pydantic-ai's ``JSON_CODEC``, recording bytes as base64 and returning
failures instead of raising them.
+
+ A result that does not encode comes back as :class:`Unrecordable`, so the
journal decides
+ whether that fails the task (:meth:`JournalStep.cannot_record`), and
``load`` hands the
+ live value on to the run either way.
+ """
+
+ def dump(self, tp: Any, value: Any) -> Any:
+ try:
+ return JSON_CODEC.dump(tp, _tag_bytes(value))
+ except (PydanticSerializationError, TypeError, ValueError,
RecursionError) as error:
+ # RecursionError: a result that contains itself, or is nested too
deep to walk.
+ return Unrecordable(value, error)
+
+ def load(self, tp: Any, payload: Any) -> Any:
+ if isinstance(payload, Unrecordable):
+ return payload.value
+ return JSON_CODEC.load(tp, payload if _is_model_response(payload) else
_untag_bytes(payload))
+
+
+_JOURNAL_CODEC = _JournalCodec()
+
+
class AirflowDurability(BaseDurabilityCapability[AgentDepsT]):
"""
Replay an agent's completed steps when Airflow retries the task it runs in.
@@ -176,8 +268,7 @@ class
AirflowDurability(BaseDurabilityCapability[AgentDepsT]):
engine_name="Airflow",
durable_unit_noun="step",
durable_container_noun="task",
- codec=JSON_CODEC,
- serialization_failure=_not_journalable,
+ codec=_JOURNAL_CODEC,
)
def __init__(
@@ -325,7 +416,7 @@ class
_AirflowOperationBackend(JournalCallableOperationBackend[None]):
return step.payload
# The masking wrapper AgentOperator adds sits outside this
durable unit, so mask
# before recording: secrets must not reach the journal any
more than the model.
- return await step.run(body, to_record=mask_secrets)
+ return await step.run(body, to_record=_mask_payload)
case ToolsetValidateToolArgumentsId():
# Local validation of the arguments the recorded call was made
with: nothing to
# save by replaying it, and it is the same on every attempt,
so it takes no position.
@@ -371,7 +462,12 @@ class
_AirflowOperationBackend(JournalCallableOperationBackend[None]):
return step.payload
ledger.record_live_model_request(had_request_credit=had_request_credit)
payload = await step.run(body)
- ledger.track_chain(_load_model_response(payload, streaming=streaming),
parameters)
+ if isinstance(payload, Unrecordable):
+ # Left unrecorded by a journal that does not fail on it: the live
response is the value.
+ response = payload.value if isinstance(payload.value,
ModelResponse) else payload.value.response
+ else:
+ response = _load_model_response(payload, streaming=streaming)
+ ledger.track_chain(response, parameters)
return payload
async def _fingerprint_model_request(
@@ -403,8 +499,8 @@ class _JournaledToolset(WrapperToolset[Any]):
Results, and the control-flow exceptions a tool raises for the model
(``ModelRetry``,
``ToolFailed``, ``ApprovalRequired``, ``CallDeferred``), are recorded as
values, so a
retry replays them exactly. Any other exception is recorded as a failure
and the call
- runs again on retry. A result that is not JSON-serializable fails the
task, as it does
- for the toolsets pydantic-ai runs as durable steps.
+ runs again on retry. Results are encoded as for the toolsets pydantic-ai
runs as durable
+ steps, bytes included, and one that cannot be encoded is handled the same
way.
"""
durability: AirflowDurability[Any] = field(repr=False, kw_only=True)
@@ -483,16 +579,15 @@ async def _run_and_record_tool(step: JournalStep, call:
Callable[[], Awaitable[A
def _record_tool_payload(step: JournalStep, payload: dict[str, Any]) -> None:
- try:
- encoded = JSON_CODEC.dump(Any, payload)
- except (PydanticSerializationError, TypeError, ValueError) as error:
- # As for the toolsets pydantic-ai runs as durable steps: the same
value fails the same
- # way on every attempt, and the model could not have read it either.
- raise _not_journalable(error) from error
- step.record(mask_secrets(encoded))
+ encoded = _JOURNAL_CODEC.dump(Any, payload)
+ if isinstance(encoded, Unrecordable):
+ step.cannot_record(encoded.error)
+ else:
+ step.record(_mask_payload(encoded))
def _decode_tool_payload(payload: Any) -> Any:
+ payload = _untag_bytes(payload)
kind = payload.get("kind") if isinstance(payload, dict) else None
match kind:
case "tool_return" if "tool_return" in payload:
@@ -525,14 +620,12 @@ def _is_journaled_tool_payload(payload: object) -> bool:
if not isinstance(payload, dict):
return False
match payload.get("kind"):
- case "tool_return" if "tool_return" in payload:
+ case "tool_return" if "tool_return" in payload or "result" in payload:
try:
- JSON_CODEC.load(ToolReturn, payload["tool_return"])
+ _decode_tool_payload(payload)
except (TypeError, ValueError):
return False
return True
- case "tool_return":
- return "result" in payload
case "model_retry" | "tool_failed":
return isinstance(payload.get("message"), str)
case "approval_required" | "call_deferred":
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 2d437834e4d..c4fad4b9422 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
@@ -59,6 +59,7 @@ from typing import TYPE_CHECKING, Any, Literal
from airflow.providers.common.ai.durable.base import RUN_ID_KEY, RUNS_KEY,
build_run_meta_key, build_step_key
from airflow.providers.common.ai.durable.storage import DurableStorage
+from airflow.providers.common.ai.exceptions import DurableJournalError
from airflow.providers.common.ai.observability import
make_task_instance_run_key
from airflow.providers.common.ai.utils.task_logger import get_task_logger
from airflow.providers.common.compat.sdk import AirflowException,
get_current_context
@@ -77,6 +78,19 @@ StepKind = Literal["model", "tool", "other"]
"""What a step is, for the end-of-run summary and for adapters that treat
model steps specially."""
+@dataclass(frozen=True)
+class Unrecordable:
+ """
+ A step result an adapter could not encode as JSON: the live value, and why.
+
+ A step body returns it in place of the payload; :meth:`JournalStep.run`
passes it to
+ :meth:`JournalStep.cannot_record` instead of recording it, and returns it
to the adapter.
+ """
+
+ value: Any
+ error: Exception
+
+
@dataclass
class JournalStep:
"""
@@ -127,7 +141,10 @@ class JournalStep:
except Exception as error:
self.fail(error)
raise
- self.record(payload if to_record is None else to_record(payload))
+ if isinstance(payload, Unrecordable):
+ self.cannot_record(payload.error)
+ else:
+ self.record(payload if to_record is None else to_record(payload))
return payload
@contextmanager
@@ -160,6 +177,28 @@ class JournalStep:
self.durable_run.journal.stats.recorded[self.kind] += 1
log.debug("Durable: recorded step", position=self.position,
step=self.name)
return True
+ self._mark_not_recorded()
+ return False
+
+ def cannot_record(self, error: Exception) -> None:
+ """
+ Note that the step's result has no form the journal can hold, such as
a value with no JSON form.
+
+ :raises DurableJournalError: When the journal fails on such results,
as it does by default.
+ Otherwise the step is left unrecorded, as when storage skips a
write, and runs again on retry.
+ """
+ if not self.replayable:
+ # Never recorded anyway, so there is nothing to fail or report.
+ return
+ if self.durable_run.journal.fail_on_unrecordable:
+ # The same value fails the same way on every retry, so retrying
the task cannot help.
+ raise DurableJournalError(
+ f"durable execution could not record a step's result, because
it is not JSON-serializable: {error}"
+ ) from error
+ log.warning("Durable: step result is not JSON-serializable",
step=self.name, error=str(error))
+ self._mark_not_recorded()
+
+ def _mark_not_recorded(self) -> None:
self.durable_run.journal.stats.not_recorded.append((self.kind,
self.name))
# Named here, not only in the summary: this line is logged on every
path,
# including the failed attempt that Airflow retries.
@@ -168,7 +207,6 @@ class JournalStep:
position=self.position,
step=self.name,
)
- return False
def fail(self, error: BaseException) -> None:
"""
@@ -482,11 +520,21 @@ class DurableJournal:
:param clean_up_after_run: Delete a run's steps as soon as it succeeds. A
journal the
operator manages leaves this off and calls :meth:`cleanup` itself
after the
task's own post-run work has succeeded.
+ :param fail_on_unrecordable: Fail the task on a step result the journal
cannot hold
+ (see :meth:`JournalStep.cannot_record`). With ``False`` the step is
logged and runs
+ again on retry instead.
"""
- def __init__(self, storage: DurableStorageProtocol, *, clean_up_after_run:
bool = False) -> None:
+ def __init__(
+ self,
+ storage: DurableStorageProtocol,
+ *,
+ clean_up_after_run: bool = False,
+ fail_on_unrecordable: bool = True,
+ ) -> None:
self.storage = storage
self.clean_up_after_run = clean_up_after_run
+ self.fail_on_unrecordable = fail_on_unrecordable
self.stats = JournalStats()
self._runs: list[DurableRun] = []
self._top_level_runs = 0
diff --git
a/providers/common/ai/tests/unit/common/ai/durable/test_capability.py
b/providers/common/ai/tests/unit/common/ai/durable/test_capability.py
index 9aba084a323..7393f0f9a77 100644
--- a/providers/common/ai/tests/unit/common/ai/durable/test_capability.py
+++ b/providers/common/ai/tests/unit/common/ai/durable/test_capability.py
@@ -24,6 +24,7 @@ the retry runs the agent again from the top with a fresh
journal over the same s
from __future__ import annotations
import asyncio
+import base64
import dataclasses
from collections.abc import Iterable
from typing import TYPE_CHECKING, Any
@@ -38,6 +39,7 @@ from pydantic_ai.messages import (
ModelResponse,
TextPart,
ToolCallPart,
+ ToolReturn,
ToolReturnPart,
)
from pydantic_ai.models.function import AgentInfo, FunctionModel
@@ -45,7 +47,8 @@ from pydantic_ai.toolsets import FunctionToolset
from pydantic_ai.usage import RequestUsage, RunUsage, UsageLimits
from airflow.providers.common.ai.durable import AirflowDurability
-from airflow.providers.common.ai.durable.journal import DurableJournal,
journal_scope
+from airflow.providers.common.ai.durable.capability import _JournalCodec
+from airflow.providers.common.ai.durable.journal import DurableJournal,
Unrecordable, journal_scope
from airflow.providers.common.ai.exceptions import DurableJournalError
from airflow.providers.common.ai.utils.toolset_base import AirflowToolset,
ensure_masked
@@ -54,6 +57,8 @@ from unit.common.ai.tool_parts import is_tool_result,
is_tool_retry
if TYPE_CHECKING:
from pydantic_ai.toolsets.abstract import ToolsetTool
+PNG = b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR"
+
class Calls:
"""Counts the live calls one test's model and tools make, across
attempts."""
@@ -333,6 +338,303 @@ class TestResultsThatCannotBeRecorded:
with pytest.raises(DurableJournalError, match="not JSON-serializable"):
await attempt(memory_storage, agent)
+ @pytest.mark.parametrize(
+ ("toolset", "codec_refuses"),
+ [
+ # From pydantic-ai 2.66 on, pydantic-ai records a function
toolset's bytes in a
+ # dataclass itself, as base64, so the codec refuses them to take
the same path on
+ # every version. An Airflow toolset's result goes through this
codec as it is.
+ pytest.param(lambda calls: _function_toolset(calls, _Image(PNG)),
True, id="function-toolset"),
+ pytest.param(lambda calls: _WarehouseToolset(calls, _Image(PNG)),
False, id="airflow-toolset"),
+ ],
+ )
+ @pytest.mark.asyncio
+ async def test_a_journal_that_does_not_fail_leaves_the_result_unrecorded(
+ self, memory_storage, monkeypatch, toolset, codec_refuses
+ ):
+ if codec_refuses:
+ dump = _JournalCodec.dump
+
+ def refusing_dump(codec: _JournalCodec, tp: Any, value: Any) ->
Any:
+ if isinstance(getattr(value, "result", None), _Image):
+ return Unrecordable(value, TypeError("no JSON form"))
+ return dump(codec, tp, value)
+
+ monkeypatch.setattr(_JournalCodec, "dump", refusing_dump)
+ calls = Calls()
+ fail_final = [True]
+
+ def build() -> Agent[None, str]:
+ return Agent(
+ FunctionModel(tool_then_answer(calls, "query",
fail_final=fail_final)),
+ name="analyst",
+ toolsets=[toolset(calls)],
+ capabilities=[AirflowDurability()],
+ )
+
+ journals = [DurableJournal(memory_storage, fail_on_unrecordable=False)
for _ in range(2)]
+ with journal_scope(journals[0]), pytest.raises(RuntimeError,
match="worker died"):
+ await build().run("go")
+ fail_final[0] = False
+ with journal_scope(journals[1]):
+ result = await build().run("go")
+
+ # The model reads the bytes inside the dataclass, so the run goes on;
the journal cannot.
+ assert result.output == f"answer: {_Image(PNG)}"
+ assert calls["query"] == 2
+ assert [kind for kind, _ in journals[0].stats.not_recorded] == ["tool"]
+
+ @pytest.mark.asyncio
+ async def test_a_result_that_contains_itself_fails_without_retrying(self,
memory_storage):
+ looped: dict[str, Any] = {}
+ looped["self"] = looped
+ agent = Agent(
+ FunctionModel(tool_then_answer(Calls(), "query",
fail_final=[False])),
+ name="analyst",
+ toolsets=[_function_toolset(Calls(), looped)],
+ capabilities=[AirflowDurability()],
+ )
+
+ with pytest.raises(DurableJournalError, match="not JSON-serializable"):
+ await attempt(memory_storage, agent)
+
+ @pytest.mark.parametrize("fail_on_unrecordable", [True, False])
+ @pytest.mark.asyncio
+ async def test_a_model_response_that_is_not_json(self, memory_storage,
fail_on_unrecordable):
+ def model_fn(messages: list[ModelMessage], info: AgentInfo) ->
ModelResponse:
+ return ModelResponse(parts=[TextPart("done")],
provider_details={"handle": object()})
+
+ agent = Agent(FunctionModel(model_fn), name="analyst",
capabilities=[AirflowDurability()])
+ journal = DurableJournal(memory_storage,
fail_on_unrecordable=fail_on_unrecordable)
+
+ with journal_scope(journal):
+ if fail_on_unrecordable:
+ with pytest.raises(DurableJournalError, match="not
JSON-serializable"):
+ await agent.run("go")
+ else:
+ assert (await agent.run("go")).output == "done"
+
+ assert [kind for kind, _ in journal.stats.not_recorded] == ([] if
fail_on_unrecordable else ["model"])
+
+ @pytest.mark.asyncio
+ async def
test_a_result_of_a_toolset_that_is_not_replayable_is_not_checked(self,
memory_storage):
+ calls = Calls()
+ agent = Agent(
+ FunctionModel(tool_then_answer(calls, "query",
fail_final=[False])),
+ name="analyst",
+ toolsets=[_WarehouseToolset(calls, _Image(PNG), replayable=False)],
+ capabilities=[AirflowDurability()],
+ )
+ journal = DurableJournal(memory_storage)
+
+ with journal_scope(journal):
+ result = await agent.run("go")
+
+ assert result.output == f"answer: {_Image(PNG)}"
+ assert journal.stats.not_recorded == []
+
+
[email protected]
+class _Image:
+ """A tool result whose bytes the model can read but the journal cannot
record."""
+
+ data: bytes
+
+
+def _function_toolset(calls: Calls, result: Any) -> FunctionToolset[None]:
+ """A function toolset with one ``query`` tool that returns ``result``."""
+ toolset: FunctionToolset[None] = FunctionToolset(id="db")
+
+ @toolset.tool_plain
+ def query(sql: str) -> Any:
+ calls.bump("query")
+ return result
+
+ return toolset
+
+
[email protected]
+class _Event:
+ """A tool result of the user's own whose ``kind`` matches pydantic-ai's
recorded tool return."""
+
+ name: str
+ kind: str = "tool_return"
+
+ @property
+ def result(self) -> str:
+ """Not a field, so it is not the ``result`` field of pydantic-ai's
recorded tool return."""
+ return self.name
+
+
+class TestBytesResults:
+ @pytest.mark.parametrize(
+ ("result", "content", "metadata"),
+ [
+ pytest.param(PNG, PNG, None, id="png"),
+ pytest.param(b"plain text", b"plain text", None, id="utf-8"),
+ pytest.param(
+ {"blob": b"\xff", "rows": [b"\xfe", 1]},
+ {"blob": b"\xff", "rows": [b"\xfe", 1]},
+ None,
+ id="nested",
+ ),
+ # A dict a tool built that only looks like a tag is kept as it is.
+ pytest.param(
+ {"__airflow_bytes__": "not base64!"},
+ {"__airflow_bytes__": "not base64!"},
+ None,
+ id="tag-shaped",
+ ),
+ pytest.param(
+ {"__airflow_bytes__": "café"}, {"__airflow_bytes__": "café"},
None, id="tag-shaped-non-ascii"
+ ),
+ # Recorded as a list, which is what the model sees either way.
+ pytest.param((PNG, 1), [PNG, 1], None, id="tuple"),
+ # pydantic-ai records a dict whose kind collides with its own tool
return differently.
+ pytest.param(
+ {"kind": "tool-return", "blob": PNG}, {"kind": "tool-return",
"blob": PNG}, None, id="kind"
+ ),
+ pytest.param(
+ ToolReturn(return_value={"png": PNG}, metadata={"thumbnail":
PNG}),
+ {"png": PNG},
+ {"thumbnail": PNG},
+ id="tool-return",
+ ),
+ ],
+ )
+ @pytest.mark.asyncio
+ async def test_a_function_tool_result_with_bytes_replays_the_same_bytes(
+ self, memory_storage, result, content, metadata
+ ):
+ calls = Calls()
+ fail_final = [True]
+ seen: list[tuple[Any, Any]] = []
+
+ def model_fn(messages: list[ModelMessage], info: AgentInfo) ->
ModelResponse:
+ calls.bump("model")
+ match responses_so_far(messages):
+ case 0:
+ return ModelResponse(parts=[ToolCallPart("query", {"sql":
"select 1"})])
+ case 1:
+ return ModelResponse(parts=[ToolCallPart("notify", {})])
+ if fail_final[0]:
+ raise RuntimeError("worker died")
+ part = tool_returns(messages)[0]
+ seen.append((part.content, part.metadata))
+ return ModelResponse(parts=[TextPart("done")])
+
+ def build() -> Agent[None, str]:
+ toolset = _function_toolset(calls, result)
+
+ @toolset.tool_plain
+ def notify() -> str:
+ calls.bump("notify")
+ return "sent"
+
+ return Agent(
+ FunctionModel(model_fn),
+ name="analyst",
+ toolsets=[toolset],
+ capabilities=[AirflowDurability()],
+ )
+
+ with pytest.raises(RuntimeError, match="worker died"):
+ await attempt(memory_storage, build())
+ fail_final[0] = False
+ await attempt(memory_storage, build())
+
+ assert seen == [(content, metadata)]
+ assert calls["query"] == 1
+ # The model request after the replayed bytes matches the recorded one,
so it replays too.
+ assert calls["notify"] == 1
+ assert calls["model"] == 4
+
+ @pytest.mark.parametrize(
+ "result",
+ [
+ pytest.param({"blob": PNG}, id="dict"),
+ pytest.param(ToolReturn(return_value={"blob": PNG}),
id="tool-return"),
+ ],
+ )
+ @pytest.mark.asyncio
+ async def
test_an_airflow_toolset_result_with_bytes_replays_the_same_bytes(self,
memory_storage, result):
+ calls = Calls()
+ fail_final = [True]
+ expected = {"blob": PNG}
+
+ def build() -> Agent[None, str]:
+ return Agent(
+ FunctionModel(tool_then_answer(calls, "query",
fail_final=fail_final)),
+ name="analyst",
+ toolsets=[_WarehouseToolset(calls, result)],
+ capabilities=[AirflowDurability()],
+ )
+
+ with pytest.raises(RuntimeError, match="worker died"):
+ await attempt(memory_storage, build())
+ fail_final[0] = False
+ run = await attempt(memory_storage, build())
+
+ assert run.output == f"answer: {expected}"
+ assert calls["query"] == 1
+
+ @pytest.mark.asyncio
+ async def test_a_result_with_its_own_kind_field_replays(self,
memory_storage):
+ calls = Calls()
+ fail_final = [True]
+
+ def build() -> Agent[None, str]:
+ return Agent(
+ FunctionModel(tool_then_answer(calls, "query",
fail_final=fail_final)),
+ name="analyst",
+ toolsets=[_function_toolset(calls, _Event("deploy"))],
+ capabilities=[AirflowDurability()],
+ )
+
+ with pytest.raises(RuntimeError, match="worker died"):
+ await attempt(memory_storage, build())
+ fail_final[0] = False
+ run = await attempt(memory_storage, build())
+
+ assert run.output == "answer: {'name': 'deploy', 'kind':
'tool_return'}"
+ assert calls["query"] == 1
+
+ @pytest.mark.asyncio
+ async def
test_a_replayed_model_response_keeps_arguments_that_look_like_bytes(self,
memory_storage):
+ """Only tool results are tagged, so what the model wrote replays as it
was written."""
+ fail_final = [True]
+ seen: list[Any] = []
+ args = {"payload": {"__airflow_bytes__": "aGk="}}
+
+ def model_fn(messages: list[ModelMessage], info: AgentInfo) ->
ModelResponse:
+ if responses_so_far(messages) == 0:
+ return ModelResponse(parts=[ToolCallPart("store", args)])
+ if fail_final[0]:
+ raise RuntimeError("worker died")
+ seen.append(next(part.args for part in messages[1].parts if
isinstance(part, ToolCallPart)))
+ return ModelResponse(parts=[TextPart("done")])
+
+ def build() -> Agent[None, str]:
+ toolset = FunctionToolset(id="db")
+
+ @toolset.tool_plain
+ def store(payload: dict[str, Any]) -> str:
+ return "stored"
+
+ return Agent(
+ FunctionModel(model_fn),
+ name="analyst",
+ toolsets=[toolset],
+ capabilities=[AirflowDurability()],
+ )
+
+ with pytest.raises(RuntimeError, match="worker died"):
+ await attempt(memory_storage, build())
+ fail_final[0] = False
+ await attempt(memory_storage, build())
+
+ assert seen == [args]
+
class TestStreaming:
@pytest.mark.asyncio
@@ -463,13 +765,20 @@ class TestUnloadableRecords:
assert calls["accrue"] == 2
+ @pytest.mark.parametrize(
+ "payload",
+ [
+ pytest.param({"kind": "renamed_in_a_later_version"},
id="unknown-kind"),
+ pytest.param({"kind": "tool_return", "tool_return": "not a tool
return"}, id="tool-return"),
+ ],
+ )
@pytest.mark.asyncio
- async def
test_an_airflow_toolset_record_that_no_longer_loads_runs_live(self,
memory_storage):
+ async def
test_an_airflow_toolset_record_that_no_longer_loads_runs_live(self,
memory_storage, payload):
calls = Calls()
await attempt(memory_storage, self.build(calls,
_WarehouseToolset(calls)))
for entry in memory_storage.steps():
if entry["kind"] == "tool":
- entry["payload"] = {"kind": "renamed_in_a_later_version"}
+ entry["payload"] = payload
result = await attempt(memory_storage, self.build(calls,
_WarehouseToolset(calls)))
@@ -741,6 +1050,42 @@ class TestMasking:
assert registered_secret not in str(memory_storage.entries)
+ @pytest.mark.parametrize(
+ "toolset",
+ [
+ pytest.param(lambda calls, result: _function_toolset(calls,
result), id="function-toolset"),
+ pytest.param(
+ lambda calls, result: _UnmaskedWarehouseToolset(calls,
result), id="airflow-toolset"
+ ),
+ ],
+ )
+ @pytest.mark.asyncio
+ async def test_journal_holds_secrets_inside_bytes_masked(
+ self, memory_storage, registered_secret, toolset
+ ):
+ result = f"password={registered_secret}".encode()
+ agent = Agent(
+ FunctionModel(tool_then_answer(Calls(), "query",
fail_final=[False])),
+ name="analyst",
+ # AgentOperator puts the masking wrapper outside the durable unit.
+ toolsets=[ensure_masked(toolset(Calls(), result))],
+ capabilities=[AirflowDurability()],
+ )
+
+ await attempt(memory_storage, agent)
+
+ journal = str(memory_storage.entries)
+ assert registered_secret not in journal
+ assert base64.b64encode(result).decode() not in journal
+ assert base64.b64encode(b"password=***").decode() in journal
+
+
+class _UnmaskedWarehouseToolset(_WarehouseToolset):
+ """An Airflow toolset whose ``call_tool`` override skips the masking the
base class does."""
+
+ async def call_tool(self, name, tool_args, ctx, tool) -> Any:
+ return await self.execute_tool(name, tool_args, ctx=ctx, tool=tool)
+
class _Ledger(AbstractCapability[Any]):
"""Stands in for pydantic-ai-harness SpendLimits: it accrues through a
durable operation."""
