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 a7066062a13 Fix pulling `AgentOperator` output for list, marker and
function types (#74729)
a7066062a13 is described below
commit a7066062a1394829ff04e0b35c84c2c25eb81518
Author: PoAn Yang <[email protected]>
AuthorDate: Sat Oct 24 01:18:54 2026 +0800
Fix pulling `AgentOperator` output for list, marker and function types
(#74729)
---
providers/common/ai/docs/changelog.rst | 23 ++
providers/common/ai/docs/structured_output.rst | 48 +++-
.../airflow/providers/common/ai/operators/agent.py | 54 +++-
.../tests/unit/common/ai/operators/test_agent.py | 293 ++++++++++++++++++++-
4 files changed, 395 insertions(+), 23 deletions(-)
diff --git a/providers/common/ai/docs/changelog.rst
b/providers/common/ai/docs/changelog.rst
index 9d22ae90a69..4093a86db8f 100644
--- a/providers/common/ai/docs/changelog.rst
+++ b/providers/common/ai/docs/changelog.rst
@@ -25,6 +25,29 @@
Changelog
---------
+.. note::
+ The worker cannot register a model class for XCom deserialization when
``output_type`` is a
+ list of output types such as ``[Summary, int]``, an output marker such as
+ ``ToolOutput(Summary)``, an output function, a ``TypeAliasType`` or a
``NewType``. For these
+ output types, ``AgentOperator`` and ``@task.agent`` now push each model in
the output as a
+ JSON-compatible dict (``model_dump(mode="json")``), including models inside
a list, tuple or
+ dict. On Airflow < 3.3, the worker registers no model class at all, so a
model inside a list,
+ tuple or dict is now pushed as a JSON-compatible dict for every
``output_type``. A top-level
+ model is still dumped with plain ``model_dump()`` when
``serialize_output=True`` is set or on
+ Airflow < 3.3.
+
+ The new behaviour changes the type that a downstream task receives after an
upgrade. Versions
+ 1.0.0 and 1.1.0 pushed these models as model instances. They dumped a
top-level model with
+ plain ``model_dump()`` only in HITL review, with ``serialize_output=True``
or on
+ Airflow < 3.3. A downstream task loaded a model instance when the model's
class was listed in
+ ``[core] allowed_deserialization_classes``, or on Airflow >= 3.3 when
another task in the
+ same Dag declared the class as its ``output_type``. That downstream task now
gets a dict and
+ must read ``result["field"]`` instead of ``result.field``. In HITL review on
Airflow >= 3.3
+ without ``serialize_output=True``, values such as ``datetime``, ``date``,
``timedelta``,
+ ``UUID`` and ``Decimal`` in a top-level model now arrive as strings, for
example
+ ``"PT1M30S"`` for a ``timedelta``. Tuples and sets in that model now arrive
as lists. See
+ :doc:`structured_output`.
+
1.1.0
.....
diff --git a/providers/common/ai/docs/structured_output.rst
b/providers/common/ai/docs/structured_output.rst
index b5e325695ee..8be2063916e 100644
--- a/providers/common/ai/docs/structured_output.rst
+++ b/providers/common/ai/docs/structured_output.rst
@@ -26,7 +26,7 @@ so downstream tasks can type-hint the class directly
(``def downstream(result: MyModel)``) and use attribute access
(``result.field``).
The declared ``output_type`` (and any ``BaseModel`` reachable from
-``Union``/``Optional``/``list`` shapes) is registered for XCom deserialization
by
+``Union``/``Optional``/``list[...]`` types) is registered for XCom
deserialization by
the worker when it loads the Dag, before any task runs -- so no edit to
``[core] allowed_deserialization_classes`` is needed. The Pydantic class must
be
defined at **module scope** and bound to an attribute matching its
``__name__``;
@@ -50,6 +50,52 @@ loaded Dag's tasks at startup and registers each declared
class, so it also work
for mapped producers (``.expand(...)``) and for workers that load Dags from a
cache that bypasses operator construction.
+The worker finds no class to register when ``output_type`` is a list of output
+types such as ``[Summary, int]``, an output marker such as
+``ToolOutput(Summary)``, an output function, a ``TypeAliasType`` or a
+``NewType``. For these output types, ``AgentOperator`` and ``@task.agent`` push
+each model in the output as a JSON-compatible dict, including models inside a
+list, tuple or dict. A model is pushed as a dict even if its class is listed in
+``[core] allowed_deserialization_classes``, or if another task in the same Dag
+declares the class as its ``output_type``. A downstream task therefore reads
+``result["field"]``, not ``result.field``.
+
+In a JSON-compatible dict, values such as ``datetime``, ``date``, ``UUID`` and
+``Decimal`` arrive as strings, tuples and sets arrive as lists, and ``Enum``
+members arrive as their values.
+
+A model at the top level is still dumped with plain ``model_dump()`` when
+``serialize_output=True`` is set or on Airflow versions before 3.3. Either
+condition is enough. A plain ``Enum`` member (one that does not also subclass
+``str`` or ``int``), an ``HttpUrl`` or a ``Path`` in that model then stays a
+Python object. The default XCom backend cannot serialize these objects, so the
+XCom push fails with ``TypeError``. Before Airflow 3.3 the worker registers no
+model class. ``AgentOperator`` and ``@task.agent`` therefore push a model
+inside a list, tuple or dict as a JSON-compatible dict for every
+``output_type``, not only for the output types listed above.
+
+Upgrading from common.ai 1.0.0 or 1.1.0 can change the type that a downstream
+task receives:
+
+* On Airflow 3.3 and later, common.ai 1.0.0 and 1.1.0 pushed each model in the
+ output as a model instance when ``output_type`` was one of the output types
+ listed above. They dumped a top-level model with plain ``model_dump()`` only
+ when ``serialize_output=True`` was set or in HITL review. A downstream task
+ loaded the instance when the class was listed in
+ ``[core] allowed_deserialization_classes`` or when another task in the same
+ Dag declared the class as its ``output_type``. That downstream task now gets
+ a dict, and ``result.field`` raises ``AttributeError``.
+* In HITL review on Airflow 3.3 and later without ``serialize_output=True``,
+ common.ai 1.0.0 and 1.1.0 dumped a top-level model of these output types with
+ plain ``model_dump()`` after the reviewer approved the model. Values such as
+ ``datetime``, ``date``, ``timedelta``, ``UUID`` and ``Decimal`` in that model
+ now arrive as strings, for example ``"PT1M30S"`` for a ``timedelta``. Tuples
+ and sets in that model now arrive as lists.
+* Before Airflow 3.3, common.ai 1.0.0 and 1.1.0 pushed a model inside a list,
+ tuple or dict as a model instance, for every ``output_type``. A downstream
+ task that loaded the instance because the class was listed in
+ ``[core] allowed_deserialization_classes`` now gets a dict.
+
The Airflow UI's XCom viewer renders Pydantic instances via the
``stringify`` path, which produces a representation like
``my_module.MyModel@version=1(field=value,...)`` without consulting the
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 58f4c8b730c..407dad59148 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
@@ -241,6 +241,26 @@ def _declares_agent_template_fields(toolset: Any) -> bool:
)
+def _dump_unregistered_models(value: Any, registered: set[type]) -> Any:
+ """Return *value* with each model whose class is not in *registered*
dumped to a JSON-compatible dict."""
+ if isinstance(value, BaseModel):
+ if type(value) in registered:
+ return value
+ # A python-mode model_dump() leaves values such as plain Enum members,
HttpUrl and Path in the
+ # dict, and the default XCom serializer raises TypeError for those
values.
+ return value.model_dump(mode="json")
+ # The default XCom serializer already pushes a subclass of list, tuple or
dict, such as a NamedTuple
+ # or an OrderedDict, as the plain type. Building the plain type here does
not change what a
+ # downstream task receives.
+ if isinstance(value, list):
+ return [_dump_unregistered_models(item, registered) for item in value]
+ if isinstance(value, tuple):
+ return tuple(_dump_unregistered_models(item, registered) for item in
value)
+ if isinstance(value, dict):
+ return {key: _dump_unregistered_models(item, registered) for key, item
in value.items()}
+ return value
+
+
# CancellableAgentRunMixin must precede BaseOperator so its on_kill overrides
BaseOperator's
# no-op. The other mixins only add methods, so they can trail BaseOperator.
See the MRO guard
# test in tests/unit/common/ai/mixins/test_cancellable_run.py.
@@ -1105,22 +1125,11 @@ class AgentOperator(CancellableAgentRunMixin,
BaseOperator, HITLReviewMixin, Usa
# The reviewer can approve, reject, or ask for a new answer, but
cannot edit the answer text.
# So the approved answer is the last answer the agent returned,
and there is no need to parse
# it back from that text.
- output = self._output_under_review
- if isinstance(output, BaseModel) and (
- self._serialize_model_output or type(output) not in
iter_pydantic_models(self.output_type)
- ):
- # A downstream task can load a model from XCom only if the
worker registered the model's
- # class. The worker registers the classes that
iter_pydantic_models finds in output_type.
- # iter_pydantic_models finds no class in a list of output
types, an output marker, an output
- # function, a TypeAliasType or a NewType, so the operator
pushes the model as a dict. When
- # iter_pydantic_models could not be imported,
_serialize_model_output is always True, so the
- # condition above never reaches the iter_pydantic_models call.
- output = output.model_dump()
+ output = self._convert_output_for_xcom(self._output_under_review)
self._clear_usage_budget(context)
return output
- if self._serialize_model_output and isinstance(output, BaseModel):
- output = output.model_dump()
+ output = self._convert_output_for_xcom(output)
# Clean up the durable journal only after the run and every post-run
step
# that can still fail (the run-metadata and message-history XCom pushes
@@ -1132,6 +1141,25 @@ class AgentOperator(CancellableAgentRunMixin,
BaseOperator, HITLReviewMixin, Usa
self._clear_usage_budget(context)
return output
+ def _convert_output_for_xcom(self, output: Any) -> Any:
+ """Return the XCom value, with dicts in place of models that a
downstream task could not load."""
+ if not self.do_xcom_push:
+ return output
+ if self._serialize_model_output and isinstance(output, BaseModel):
+ return output.model_dump()
+ if not _CORE_WALKER:
+ # iter_pydantic_models could not be imported, so the worker
registers no class.
+ # _serialize_model_output is always True when that import fails,
so a model at the top level was
+ # already dumped above. The class of a model inside a list, tuple
or dict is not registered
+ # either, so this call passes an empty set of registered classes
to dump that model too.
+ return _dump_unregistered_models(output, set())
+ # A downstream task can load a model from XCom only if the worker
registered the model's class.
+ # The worker registers the classes that iter_pydantic_models finds in
output_type.
+ # iter_pydantic_models finds no class in a list of output types, an
output marker, an output
+ # function, a TypeAliasType or a NewType. For these output types the
operator pushes each model
+ # as a dict, including models inside a list, tuple or dict.
+ return _dump_unregistered_models(output,
set(iter_pydantic_models(self.output_type)))
+
def _pause_for_tool_approval(self, context: Context, result: Any, *,
attempt_usage: RunUsage) -> NoReturn:
"""
Park the task until a human approves or rejects the tool calls the
agent is waiting on.
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 178db398166..8ee85870a7b 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
@@ -21,16 +21,18 @@ import dataclasses
import json
import logging
import sys
+from collections import OrderedDict
from contextlib import nullcontext
from datetime import timedelta
from decimal import Decimal
+from enum import Enum
from types import ModuleType, SimpleNamespace
-from typing import Any, NewType
+from typing import Any, NamedTuple, NewType
from unittest.mock import ANY, MagicMock, PropertyMock, call, patch
from uuid import UUID
import pytest
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict, HttpUrl
from pydantic_ai import Agent, DeferredToolRequests, RunContext, Tool,
ToolOutput
from pydantic_ai.capabilities import (
AbstractCapability,
@@ -106,7 +108,11 @@ from unit.common.ai.durable.memory_storage import
MemoryStorage
from unit.common.ai.sandbox.fake_tags import TaggedBackend
try:
- from airflow.sdk.serde import SUPPORTS_OPERATOR_DESERIALIZATION_WALKER as
_CORE_WALKER
+ from airflow.sdk.serde import (
+ SUPPORTS_OPERATOR_DESERIALIZATION_WALKER as _CORE_WALKER,
+ deserialize,
+ serialize,
+ )
except ImportError:
_CORE_WALKER = False
@@ -129,6 +135,42 @@ SummaryAlias = TypeAliasType("SummaryAlias", Summary)
SummaryNewType = NewType("SummaryNewType", Summary)
+class SummaryKind(Enum):
+ SHORT = "short"
+
+
+class LinkedSummary(BaseModel):
+ kind: SummaryKind
+ url: HttpUrl
+
+
+class DetailedSummary(Summary):
+ detail: str = "more"
+
+
+class SummaryPair(NamedTuple):
+ summary: Summary
+ total: int
+
+
+class SummaryList(list):
+ pass
+
+
+class Unserializable:
+ pass
+
+
+class UnserializableSummary(BaseModel):
+ model_config = ConfigDict(arbitrary_types_allowed=True)
+
+ value: Unserializable
+
+
+class BytesSummary(BaseModel):
+ data: bytes
+
+
def _make_mock_agent(output, make_mock_run_result, *, cost=None):
"""Create a mock agent that returns the given output.
@@ -1034,6 +1076,206 @@ class TestAgentOperatorExecute:
assert result.text == "Great"
assert result.score == 0.95
+ @pytest.mark.parametrize(
+ "output_type",
+ [
+ pytest.param([Summary, int], id="list-of-output-types"),
+ pytest.param(ToolOutput(Summary), id="output-marker"),
+ pytest.param(_summarize, id="output-function"),
+ pytest.param(SummaryAlias, id="type-alias"),
+ pytest.param(SummaryNewType, id="new-type"),
+ ],
+ )
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ def test_execute_dumps_a_model_the_worker_does_not_register(
+ self, mock_hook_cls, make_mock_run_result, output_type
+ ):
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
_make_mock_agent(
+ Summary(text="ok"), make_mock_run_result
+ )
+ op = AgentOperator(
+ task_id="test", prompt="Analyze this", llm_conn_id="my_llm",
output_type=output_type
+ )
+
+ result = op.execute(context=MagicMock())
+
+ assert result == {"text": "ok", "score": 0.0}
+
+ @requires_typed_xcom
+ @pytest.mark.parametrize(
+ ("output_type", "output", "expected", "serialize_output"),
+ [
+ pytest.param(
+ ToolOutput(list[Summary]),
+ [Summary(text="ok")],
+ [{"text": "ok", "score": 0.0}],
+ False,
+ id="list-in-output-marker",
+ ),
+ pytest.param(
+ [list[Summary], int],
+ [Summary(text="ok")],
+ [{"text": "ok", "score": 0.0}],
+ False,
+ id="list-in-list-of-output-types",
+ ),
+ pytest.param(
+ ToolOutput(tuple[Summary, int]),
+ (Summary(text="ok"), 1),
+ ({"text": "ok", "score": 0.0}, 1),
+ False,
+ id="tuple-in-output-marker",
+ ),
+ pytest.param(
+ ToolOutput(dict[str, Summary]),
+ {"a": Summary(text="ok")},
+ {"a": {"text": "ok", "score": 0.0}},
+ False,
+ id="dict-in-output-marker",
+ ),
+ pytest.param(
+ ToolOutput(list[tuple[dict[str, list[Summary]], int]]),
+ [({"a": [Summary(text="ok")]}, 1)],
+ [({"a": [{"text": "ok", "score": 0.0}]}, 1)],
+ False,
+ id="nested-containers-in-output-marker",
+ ),
+ pytest.param(
+ ToolOutput(SummaryPair),
+ SummaryPair(Summary(text="ok"), 1),
+ ({"text": "ok", "score": 0.0}, 1),
+ False,
+ id="named-tuple-in-output-marker",
+ ),
+ pytest.param(
+ ToolOutput(list[Summary]),
+ SummaryList([Summary(text="ok")]),
+ [{"text": "ok", "score": 0.0}],
+ False,
+ id="list-subclass-in-output-marker",
+ ),
+ pytest.param(
+ ToolOutput(OrderedDict[str, Summary]),
+ OrderedDict(a=Summary(text="ok")),
+ {"a": {"text": "ok", "score": 0.0}},
+ False,
+ id="ordered-dict-in-output-marker",
+ ),
+ pytest.param(
+ list[Summary],
+ [DetailedSummary(text="ok")],
+ [{"text": "ok", "score": 0.0, "detail": "more"}],
+ False,
+ id="subclass-of-a-registered-model",
+ ),
+ pytest.param(
+ ToolOutput(list[Summary]),
+ [Summary(text="ok")],
+ [{"text": "ok", "score": 0.0}],
+ True,
+ id="list-in-output-marker-with-serialize-output",
+ ),
+ pytest.param(
+ ToolOutput(LinkedSummary),
+ LinkedSummary(kind=SummaryKind.SHORT,
url=HttpUrl("https://example.com/a";)),
+ {"kind": "short", "url": "https://example.com/a"},
+ False,
+ id="enum-and-url-fields",
+ ),
+ ],
+ )
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ def test_execute_pushes_each_model_the_worker_does_not_register_as_json(
+ self, mock_hook_cls, make_mock_run_result, output_type, output,
expected, serialize_output
+ ):
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
_make_mock_agent(
+ output, make_mock_run_result
+ )
+ op = AgentOperator(
+ task_id="test",
+ prompt="Analyze this",
+ llm_conn_id="my_llm",
+ output_type=output_type,
+ serialize_output=serialize_output,
+ )
+
+ result = op.execute(context=MagicMock())
+
+ assert result == expected
+ assert deserialize(serialize(result)) == expected
+
+ @requires_typed_xcom
+ @pytest.mark.parametrize(
+ ("output_type", "output", "key"),
+ [
+ pytest.param(list[Summary], [Summary(text="ok")], 0, id="list"),
+ pytest.param(tuple[Summary, int], (Summary(text="ok"), 1), 0,
id="tuple"),
+ pytest.param(dict[str, Summary], {"a": Summary(text="ok")}, "a",
id="dict"),
+ ],
+ )
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ def test_execute_keeps_each_model_the_worker_registers(
+ self, mock_hook_cls, make_mock_run_result, output_type, output, key
+ ):
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
_make_mock_agent(
+ output, make_mock_run_result
+ )
+ op = AgentOperator(
+ task_id="test", prompt="Analyze this", llm_conn_id="my_llm",
output_type=output_type
+ )
+
+ result = op.execute(context=MagicMock())
+
+ assert result == output
+ assert type(result[key]) is Summary
+
+ @pytest.mark.parametrize(
+ "output_type",
+ [
+ pytest.param(ToolOutput(list[Summary]),
id="list-in-output-marker"),
+ pytest.param(list[Summary], id="list"),
+ ],
+ )
+ @patch("airflow.providers.common.ai.operators.agent._CORE_WALKER", False)
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ def test_execute_dumps_nested_models_without_the_worker_walk(
+ self, mock_hook_cls, make_mock_run_result, output_type
+ ):
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
_make_mock_agent(
+ [Summary(text="ok")], make_mock_run_result
+ )
+ op = AgentOperator(
+ task_id="test", prompt="Analyze this", llm_conn_id="my_llm",
output_type=output_type
+ )
+
+ result = op.execute(context=MagicMock())
+
+ assert result == [{"text": "ok", "score": 0.0}]
+
+ @pytest.mark.parametrize(
+ "output",
+ [
+ pytest.param(UnserializableSummary(value=Unserializable()),
id="arbitrary-type-field"),
+ pytest.param(BytesSummary(data=b"\xff"), id="non-utf-8-bytes"),
+ ],
+ )
+ @patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
+ def test_execute_returns_the_output_unchanged_when_nothing_is_pushed(
+ self, mock_hook_cls, make_mock_run_result, output
+ ):
+ mock_hook_cls.get_hook.return_value.create_agent.return_value =
_make_mock_agent(
+ output, make_mock_run_result
+ )
+ op = AgentOperator(
+ task_id="test",
+ prompt="Analyze this",
+ llm_conn_id="my_llm",
+ output_type=[type(output), int],
+ do_xcom_push=False,
+ )
+
+ assert op.execute(context=MagicMock()) is output
+
def test_declares_output_type_for_deserialization(self):
"""Declares ``output_type`` so the worker-side DAG walk registers it
for deserialization."""
assert "output_type" in
AgentOperator.deserialization_allowed_class_fields
@@ -1103,14 +1345,14 @@ class TestAgentOperatorExecute:
def test_execute_with_enable_hitl_review_delegates_to_run_hitl_review(
self, mock_hook_cls, mock_run_hitl, make_mock_run_result
):
- """When enable_hitl_review=True, execute delegates to run_hitl_review
with output and message_history."""
+ """execute returns the agent's output instead of run_hitl_review's
text."""
msg_history = [MagicMock()]
mock_result = make_mock_run_result("Initial output")
mock_result.all_messages.return_value = msg_history
mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = mock_result
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
- mock_run_hitl.return_value = "Initial output"
+ mock_run_hitl.return_value = "Approved output"
op = AgentOperator(
task_id="test",
@@ -1131,16 +1373,15 @@ class TestAgentOperatorExecute:
)
@patch("airflow.providers.common.ai.operators.agent.AgentOperator.run_hitl_review",
autospec=True)
@patch("airflow.providers.common.ai.operators.agent.PydanticAIHook",
autospec=True)
- def test_execute_with_hitl_rehydrates_base_model(
+ def test_execute_with_hitl_returns_the_model_instance(
self, mock_hook_cls, mock_run_hitl, make_mock_run_result
):
- """When enable_hitl_review=True and output_type is BaseModel, execute
returns the model instance."""
+ """execute returns the agent's model instance, not a model parsed from
run_hitl_review's text."""
mock_result = make_mock_run_result(Summary(text="Approved summary",
score=0.9))
mock_agent = MagicMock(spec=["run_sync", "instrument"])
mock_agent.run_sync.return_value = mock_result
mock_hook_cls.get_hook.return_value.create_agent.return_value =
mock_agent
- # run_hitl_review returns JSON string (as stored in
session.current_output)
- mock_run_hitl.return_value = '{"text": "Approved summary", "score":
0.9}'
+ mock_run_hitl.return_value = '{"text": "Edited"}'
op = AgentOperator(
task_id="test",
@@ -1204,6 +1445,19 @@ class TestAgentOperatorExecute:
pytest.param(_summarize, Summary(text="ok"), {"text": "ok",
"score": 0.0}, id="output-function"),
pytest.param(SummaryAlias, Summary(text="ok"), {"text": "ok",
"score": 0.0}, id="type-alias"),
pytest.param(SummaryNewType, Summary(text="ok"), {"text": "ok",
"score": 0.0}, id="new-type"),
+ pytest.param(
+ ToolOutput(list[Summary]),
+ [Summary(text="ok")],
+ [{"text": "ok", "score": 0.0}],
+ id="list-in-output-marker",
+ ),
+ pytest.param(
+ ToolOutput(LinkedSummary),
+ LinkedSummary(kind=SummaryKind.SHORT,
url=HttpUrl("https://example.com/a";)),
+ {"kind": "short", "url": "https://example.com/a"},
+ id="enum-and-url-fields",
+ marks=requires_typed_xcom,
+ ),
],
)
@patch("airflow.providers.common.ai.operators.agent.AgentOperator.run_hitl_review",
autospec=True)
@@ -2040,6 +2294,27 @@ class TestAgentOperatorDurable:
assert [step["name"] for step in storage.steps()] ==
["t__model.request"]
+ def test_cleanup_skipped_when_output_conversion_fails(self):
+ """The journal must survive a failed conversion of the output for
XCom, so the Airflow retry
+ can still replay the recorded steps."""
+ storage = MemoryStorage()
+ op = AgentOperator(task_id="t", prompt="p", llm_conn_id="c",
durable=True)
+ op.llm_hook = MagicMock(spec=["create_agent"])
+ op.llm_hook.create_agent.side_effect = lambda **kw: Agent(
+ FunctionModel(lambda messages, info:
ModelResponse(parts=[TextPart("ok")])), **kw
+ )
+
+ with (
+ patch.object(AgentOperator, "_build_durable_storage",
autospec=True, return_value=storage),
+ patch.object(
+ op, "_convert_output_for_xcom", autospec=True,
side_effect=RuntimeError("cannot dump")
+ ),
+ pytest.raises(RuntimeError, match="cannot dump"),
+ ):
+ op.execute(context=_make_context())
+
+ assert [step["name"] for step in storage.steps()] ==
["t__model.request"]
+
def test_supports_durable_execution_marker(self):
assert AgentOperator._AgentOperator__supports_durable_execution is True
