Lee-W · Oct 9, 2026 · Oct 7, 2026 · Oct 8, 2026 · Oct 8, 2026 · Oct 8, 2026
@@ -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``.
 

@@ -35,6 +35,9 @@
 
 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 @@ def run_hitl_review(
         :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 @@ def _poll_loop(
         """
         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 @@ def _poll_loop(
 
         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 @@ def _poll_loop(
                     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,

@@ -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``.
 

@@ -60,6 +60,8 @@
     _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 @@ def __init__(
         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,

@@ -160,7 +160,11 @@ def __init__(
         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

@@ -154,6 +154,10 @@ def __init__(
         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

@@ -128,7 +128,10 @@ def __init__(
         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

@@ -26,6 +26,7 @@
     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(

@@ -22,6 +22,7 @@
 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."""

@@ -27,7 +27,10 @@
 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 @@ def __init__(
         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 @@ def test_timeout_raises_hitl_timeout_error(
         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(

@@ -51,17 +51,14 @@ def test_inherits_from_skipmixin_is_true(self):
     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)

@@ -147,6 +147,12 @@ def test_init_succeeds(self, kwargs):
         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(

@@ -192,6 +192,10 @@ def test_inherits_from_llm_operator(self):
 
         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",