This is an automated email from the ASF dual-hosted git repository.
Lee-W 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 7811cbacf5a Reject unsupported output_type and bound HITL XCom pull
retries (#74450)
7811cbacf5a is described below
commit 7811cbacf5ab5227450c80bf2bd982c3c41eb371
Author: Wei Lee <[email protected]>
AuthorDate: Fri Oct 9 15:18:19 2026 +0900
Reject unsupported output_type and bound HITL XCom pull retries (#74450)
---
providers/common/ai/docs/hitl_review.rst | 9 ++++--
.../providers/common/ai/mixins/hitl_review.py | 21 ++++++++++--
.../airflow/providers/common/ai/operators/agent.py | 6 ++--
.../airflow/providers/common/ai/operators/llm.py | 9 ++++--
.../providers/common/ai/operators/llm_branch.py | 6 +++-
.../common/ai/operators/llm_schema_compare.py | 4 +++
.../providers/common/ai/operators/llm_sql.py | 5 ++-
.../ai/decorators/test_llm_schema_compare.py | 14 ++++++++
.../unit/common/ai/decorators/test_llm_sql.py | 12 +++++++
.../unit/common/ai/mixins/test_hitl_review.py | 37 ++++++++++++++++++++--
.../unit/common/ai/operators/test_llm_branch.py | 19 +++++------
.../common/ai/operators/test_llm_schema_compare.py | 6 ++++
.../tests/unit/common/ai/operators/test_llm_sql.py | 4 +++
13 files changed, 127 insertions(+), 25 deletions(-)
diff --git a/providers/common/ai/docs/hitl_review.rst
b/providers/common/ai/docs/hitl_review.rst
index 1b40a731dc8..40df880ccdb 100644
--- a/providers/common/ai/docs/hitl_review.rst
+++ b/providers/common/ai/docs/hitl_review.rst
@@ -23,7 +23,8 @@ Human-in-the-loop (HITL) review for agents
HITL Review adds an interactive feedback loop to agentic operators. After the
LLM Agent produces an initial output, a human reviewer can **approve**,
**reject**, or
**request changes** through a chat UI. The operator blocks until a
-terminal action, or until a timeout is reached or max_iterations reached.
+terminal action, or until a timeout is reached, max_iterations is reached, or
+polling the review XCom keeps failing.
This document describes the architecture, workflow, API, XCom schema, and
usage.
@@ -50,7 +51,8 @@ the API server and accesses the metadata database.
.. important::
**Worker slot usage**: Each HITL task **holds a worker slot for the entire
- review duration** (until approve, reject, or timeout or max_iterations).
The operator polls
+ review duration** (until approve, reject, timeout, max_iterations, or
+ repeated XCom polling failures). The operator polls
XCom with ``time.sleep``; it does not defer. With a 10-second poll interval
and review times of 30+ minutes, the worker is occupied for the duration.
@@ -133,7 +135,8 @@ Enable the review loop with ``enable_hitl_review=True``:
changes at iterations 1 to 4; the fifth output must be either approved or
rejected. Default ``5``.
- ``hitl_timeout``: Maximum wall-clock time to wait for all review rounds.
- ``None`` = no timeout (blocks until a terminal action).
+ ``None`` = no wall-clock timeout. The task still fails, re-raising the XCom
error, if polling
+ the human action XCom fails 10 times in a row.
- ``hitl_poll_interval``: Seconds between XCom polls while waiting for a
human response. Default ``10``.
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/mixins/hitl_review.py
b/providers/common/ai/src/airflow/providers/common/ai/mixins/hitl_review.py
index 612a902a223..8a1327a871e 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/mixins/hitl_review.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/mixins/hitl_review.py
@@ -35,6 +35,9 @@ from airflow.providers.common.ai.utils.hitl_review import (
log = logging.getLogger(__name__)
+# ``hitl_timeout`` may be ``None``, so a persistently failing poll would
otherwise loop forever.
+_MAX_CONSECUTIVE_XCOM_PULL_FAILURES = 10
+
if TYPE_CHECKING:
from airflow.sdk import Context
@@ -64,7 +67,9 @@ class HITLReviewMixin:
4. When a human sets action to ``approved``, returns the output.
5. When a human sets action to ``rejected``, raises a `HITLRejectException`
- The loop stops after ``hitl_timeout`` or ``max_hitl_iterations``.
+ The loop stops after ``hitl_timeout`` or ``max_hitl_iterations``. It also
stops, by
+ re-raising the underlying exception, when polling the human action XCom
fails
+ ``_MAX_CONSECUTIVE_XCOM_PULL_FAILURES`` times in a row, even if
``hitl_timeout`` is ``None``.
**Max iterations:** ``iteration`` counts outputs shown to the reviewer
(1 = initial, 2 = first regeneration, etc.). When the reviewer requests
@@ -107,6 +112,8 @@ class HITLReviewMixin:
:raises HITLMaxIterationsError: When max iterations reached without
approval.
:raises HITLRejectException: When the reviewer rejects the output.
:raises HITLTimeoutError: When hitl_timeout elapses with no response.
+ :raises Exception: The last XCom pull exception, when pulling the
human action
+ fails ``_MAX_CONSECUTIVE_XCOM_PULL_FAILURES`` times in a row.
"""
output_str = self._to_string(output) # type: ignore[attr-defined]
ti = context["task_instance"]
@@ -150,7 +157,8 @@ class HITLReviewMixin:
"""
Block until the session reaches a terminal state.
- This loops until the human approves, rejects, or the timeout/max
iterations is reached.
+ This loops until the human approves, rejects, the timeout/max
iterations is reached,
+ or pulling the human action XCom fails
``_MAX_CONSECUTIVE_XCOM_PULL_FAILURES`` times in a row.
"""
from airflow.providers.standard.exceptions import (
HITLRejectException,
@@ -159,6 +167,7 @@ class HITLReviewMixin:
last_seen_iteration = 0
first_poll = True
+ consecutive_pull_failures = 0
while True:
if deadline is not None and time.monotonic() > deadline:
@@ -181,8 +190,16 @@ class HITLReviewMixin:
key=XCOM_HUMAN_ACTION, task_ids=ti.task_id,
map_indexes=ti.map_index
)
except Exception:
+ consecutive_pull_failures += 1
+ if consecutive_pull_failures >=
_MAX_CONSECUTIVE_XCOM_PULL_FAILURES:
+ self.log.error(
+ "Giving up HITL review after %d consecutive XCom pull
failures",
+ consecutive_pull_failures,
+ )
+ raise
self.log.warning("Failed to pull XCom", exc_info=True)
continue
+ consecutive_pull_failures = 0
if action_raw is None:
# Human action may take some time to propagate; it must be
performed in the UI,
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 3c45bed198b..bc8ff30ac12 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
@@ -383,8 +383,10 @@ class AgentOperator(CancellableAgentRunMixin,
BaseOperator, HITLReviewMixin):
without calling the LLM. E.g. 5 allows changes at iterations 1–4.
Default ``5``.
:param hitl_timeout: Maximum wall-clock time to wait for
- all review rounds combined. ``None`` means no timeout (the
- operator blocks until a terminal action).
+ all review rounds combined. ``None`` means no wall-clock timeout;
+ the review still ends on a terminal action (approve or reject),
+ ``max_hitl_iterations``, or when polling the human action XCom fails
+ too many times in a row.
:param hitl_poll_interval: Seconds between XCom polls
while waiting for a human response. Default ``10``.
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
index 85b8465295b..c8837773caa 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm.py
@@ -60,6 +60,8 @@ except ImportError: # pragma: no cover - missing
``apache-airflow-task-sdk`` wa
_CORE_WALKER = False
if TYPE_CHECKING:
+ from types import UnionType
+
from pydantic_ai import Agent
from pydantic_ai.usage import UsageLimits
@@ -99,8 +101,9 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin):
:class:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook`
for how blank entries in the list are dropped.
:param system_prompt: System-level instructions for the LLM agent.
- :param output_type: Expected output type. Default ``str``. Set to a
Pydantic
- ``BaseModel`` subclass for structured output; the model instance is
+ :param output_type: Expected output type. Default ``str``. A single type
+ (``str``, ``int``, ``list[str]``, a Pydantic ``BaseModel`` subclass)
or a
+ union (``A | B``). For a ``BaseModel`` subclass, the model instance is
returned to XCom unchanged so downstream tasks can type-hint it
directly. The class must be defined at module scope -- nested classes
cannot be deserialized from XCom.
@@ -197,7 +200,7 @@ class LLMOperator(CancellableAgentRunMixin, BaseOperator,
LLMApprovalMixin):
model_id: str | None = None,
fallback_conn_ids: list[str] | None = None,
system_prompt: str = "",
- output_type: type = str,
+ output_type: type | UnionType = str,
agent_params: dict[str, Any] | None = None,
usage_limits: UsageLimits | dict[str, Any] | None = None,
require_approval: bool = False,
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py
index 2942b1601d4..21314fc2342 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py
@@ -160,7 +160,11 @@ class LLMBranchOperator(LLMOperator, BranchMixIn):
ignore_downstream_trigger_rules: bool = False,
**kwargs: Any,
) -> None:
- kwargs.pop("output_type", None)
+ if "output_type" in kwargs:
+ raise TypeError(
+ "LLMBranchOperator does not accept 'output_type'; it builds
the output type "
+ "itself from the downstream task IDs."
+ )
super().__init__(**kwargs)
self.branches = self._normalize_branches(branches)
self.allow_multiple_branches = allow_multiple_branches
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py
index cc2b07c853b..1865fcd6e5f 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_schema_compare.py
@@ -154,6 +154,10 @@ class LLMSchemaCompareOperator(LLMOperator):
system_prompt: str = DEFAULT_SYSTEM_PROMPT,
**kwargs: Any,
) -> None:
+ if "output_type" in kwargs:
+ raise TypeError(
+ "LLMSchemaCompareOperator does not accept 'output_type'; it
always returns SchemaCompareResult."
+ )
kwargs["output_type"] = SchemaCompareResult
# execute() always returns a dict, so the approval resume path must too
kwargs["serialize_output"] = True
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py
index 847e3d1ee26..879a8c440fc 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_sql.py
@@ -128,7 +128,10 @@ class LLMSQLQueryOperator(LLMOperator):
datasource_config: DataSourceConfig | None = None,
**kwargs: Any,
) -> None:
- kwargs.pop("output_type", None) # SQL operator always returns str
+ if "output_type" in kwargs:
+ raise TypeError(
+ "LLMSQLQueryOperator does not accept 'output_type'; it always
returns the generated SQL as str."
+ )
super().__init__(**kwargs)
self.db_conn_id = db_conn_id
self.table_names = table_names
diff --git
a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py
b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py
index b56e4eed118..7cfd2bece85 100644
---
a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py
+++
b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_schema_compare.py
@@ -26,6 +26,7 @@ from airflow.providers.common.ai.operators.llm_schema_compare
import (
LLMSchemaCompareOperator,
SchemaCompareResult,
)
+from airflow.providers.common.compat.sdk import DAG, task
from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS
@@ -48,6 +49,19 @@ class TestLLMSchemaCompareDecoratedOperator:
def test_custom_operator_name(self):
assert _LLMSchemaCompareDecoratedOperator.custom_operator_name ==
"@task.llm_schema_compare"
+ def test_output_type_rejected(self):
+ @task.llm_schema_compare(
+ llm_conn_id="my_llm", db_conn_ids=["conn_a", "conn_b"],
table_names=["t"], output_type=int
+ )
+ def my_prompt_fn():
+ return "Compare schemas"
+
+ with (
+ DAG(dag_id="dag"),
+ pytest.raises(TypeError, match="LLMSchemaCompareOperator does not
accept 'output_type'"),
+ ):
+ my_prompt_fn()
+
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
@patch.object(LLMSchemaCompareOperator, "_build_schema_context",
return_value="mocked schema")
def test_execute_calls_callable_and_uses_result_as_prompt(
diff --git
a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py
b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py
index e2d1f2c54b8..5aa2f730ae2 100644
--- a/providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py
+++ b/providers/common/ai/tests/unit/common/ai/decorators/test_llm_sql.py
@@ -22,6 +22,7 @@ import pytest
from pydantic_ai.messages import ImageUrl
from airflow.providers.common.ai.decorators.llm_sql import
_LLMSQLDecoratedOperator
+from airflow.providers.common.compat.sdk import DAG, task
from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS
@@ -30,6 +31,17 @@ class TestLLMSQLDecoratedOperator:
def test_custom_operator_name(self):
assert _LLMSQLDecoratedOperator.custom_operator_name == "@task.llm_sql"
+ def test_output_type_rejected(self):
+ @task.llm_sql(llm_conn_id="my_llm", output_type=int)
+ def my_prompt_fn():
+ return "Get all users"
+
+ with (
+ DAG(dag_id="dag"),
+ pytest.raises(TypeError, match="LLMSQLQueryOperator does not
accept 'output_type'"),
+ ):
+ my_prompt_fn()
+
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_execute_calls_callable_and_uses_result_as_prompt(self,
mock_hook_cls, make_mock_run_result):
"""The user's callable return value becomes the LLM prompt."""
diff --git
a/providers/common/ai/tests/unit/common/ai/mixins/test_hitl_review.py
b/providers/common/ai/tests/unit/common/ai/mixins/test_hitl_review.py
index 1bf319f5745..fba67629ce1 100644
--- a/providers/common/ai/tests/unit/common/ai/mixins/test_hitl_review.py
+++ b/providers/common/ai/tests/unit/common/ai/mixins/test_hitl_review.py
@@ -27,7 +27,10 @@ from datetime import timedelta
from unittest.mock import MagicMock, patch
from airflow.providers.common.ai.exceptions import HITLMaxIterationsError
-from airflow.providers.common.ai.mixins.hitl_review import HITLReviewMixin
+from airflow.providers.common.ai.mixins.hitl_review import (
+ _MAX_CONSECUTIVE_XCOM_PULL_FAILURES,
+ HITLReviewMixin,
+)
from airflow.providers.common.ai.utils.hitl_review import (
XCOM_AGENT_OUTPUT_PREFIX,
XCOM_AGENT_SESSION,
@@ -57,7 +60,7 @@ class FakeAgenticOperator(HITLReviewMixin):
self.max_hitl_iterations = 5
self.hitl_timeout = hitl_timeout or timedelta(seconds=30)
self.hitl_poll_interval = hitl_poll_interval
- self.log = MagicMock(spec=["info", "warning"])
+ self.log = MagicMock(spec=["info", "warning", "error"])
def regenerate_with_feedback(self, *, feedback: str, message_history):
return f"Revised: {feedback}", message_history
@@ -218,6 +221,36 @@ class TestHITLReviewMixin:
with pytest.raises(HITLTimeoutError, match="Task exceeded timeout"):
fake_op.run_hitl_review(context, "Output")
+ @patch("airflow.providers.common.ai.mixins.hitl_review.time.sleep",
autospec=True)
+ def test_xcom_pull_failures_raise_after_limit(self, mock_sleep, fake_op,
mock_ti, context):
+ fake_op.hitl_timeout = None
+
+ def failing_pull(*args, **kwargs):
+ if mock_ti.xcom_pull.call_count >
_MAX_CONSECUTIVE_XCOM_PULL_FAILURES + 2:
+ pytest.fail("XCom pull retried past the limit")
+ raise ConnectionError("api server down")
+
+ mock_ti.xcom_pull.side_effect = failing_pull
+
+ with pytest.raises(ConnectionError, match="api server down"):
+ fake_op.run_hitl_review(context, "Output")
+
+ assert mock_ti.xcom_pull.call_count ==
_MAX_CONSECUTIVE_XCOM_PULL_FAILURES
+ assert fake_op.log.warning.call_count ==
_MAX_CONSECUTIVE_XCOM_PULL_FAILURES - 1
+ fake_op.log.error.assert_called_once_with(
+ "Giving up HITL review after %d consecutive XCom pull failures",
+ _MAX_CONSECUTIVE_XCOM_PULL_FAILURES,
+ )
+
+ @patch("airflow.providers.common.ai.mixins.hitl_review.time.sleep",
autospec=True)
+ def test_xcom_pull_failure_counter_resets_on_success(self, mock_sleep,
fake_op, mock_ti, context):
+ fake_op.hitl_timeout = None
+ approve = HumanActionData(action="approve",
iteration=1).model_dump(mode="json")
+ failures = [ConnectionError("blip")] *
(_MAX_CONSECUTIVE_XCOM_PULL_FAILURES - 1)
+ mock_ti.xcom_pull.side_effect = [*failures, None, *failures, approve]
+
+ assert fake_op.run_hitl_review(context, "Output") == "Output"
+
@patch.object(FakeAgenticOperator, "regenerate_with_feedback")
@patch("airflow.providers.common.ai.mixins.hitl_review.time.sleep",
autospec=True)
def test_empty_feedback_with_changes_requested_treats_as_approve(
diff --git
a/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py
index 5ecb7d26abe..d57b4d45f40 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py
@@ -51,17 +51,14 @@ class TestLLMBranchOperator:
def test_template_fields(self):
assert set(LLMBranchOperator.template_fields) ==
{*LLMOperator.template_fields, "branches"}
- def test_output_type_ignored(self):
- """Passing output_type= doesn't break anything; it's silently
dropped."""
- op = LLMBranchOperator(
- task_id="test",
- prompt="pick a branch",
- llm_conn_id="my_llm",
- output_type=int,
- )
- # output_type is overridden to str (the LLMOperator default) since
- # the real output_type is built dynamically from downstream_task_ids
- assert op.output_type is str
+ def test_output_type_rejected(self):
+ with pytest.raises(TypeError, match="does not accept 'output_type'"):
+ LLMBranchOperator(
+ task_id="test",
+ prompt="pick a branch",
+ llm_conn_id="my_llm",
+ output_type=int,
+ )
@patch.object(LLMBranchOperator, "do_branch")
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
diff --git
a/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py
index 8be1f8c5e7d..3b0f4034938 100644
---
a/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py
+++
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_schema_compare.py
@@ -147,6 +147,12 @@ class TestLLMSchemaCompareOperator:
op = LLMSchemaCompareOperator(**_BASE_KWARGS, **kwargs)
assert op.context_strategy == "full"
+ def test_output_type_rejected(self):
+ with pytest.raises(TypeError, match="LLMSchemaCompareOperator does not
accept 'output_type'"):
+ LLMSchemaCompareOperator(
+ **_BASE_KWARGS, db_conn_ids=["conn_a", "conn_b"],
table_names=["t"], output_type=int
+ )
+
def test_init_succeeds_with_all_parameters(self):
ds = _make_ds_config(conn_id="ds")
op = LLMSchemaCompareOperator(
diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py
index 208f439319c..74161d00f53 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_sql.py
@@ -192,6 +192,10 @@ class TestLLMSQLQueryOperator:
assert issubclass(LLMSQLQueryOperator, LLMOperator)
+ def test_output_type_rejected(self):
+ with pytest.raises(TypeError, match="LLMSQLQueryOperator does not
accept 'output_type'"):
+ LLMSQLQueryOperator(task_id="test", prompt="get users",
llm_conn_id="my_llm", output_type=int)
+
def test_template_fields_include_parent_and_sql_specific(self):
expected = {
"prompt",
