Skip downstream tasks on LLMBranchOperator reject

providers/common/ai/docs/operators/llm_branch.rst

@@ -92,9 +92,10 @@ against the downstream task IDs before branching:
     :start-after: [START howto_operator_llm_branch_approval]
     :end-before: [END howto_operator_llm_branch_approval]
 
-Rejecting the review, or letting ``approval_timeout`` expire, **fails** the
-task (``HITLRejectException`` / ``HITLTimeoutError``), so downstream tasks
-end up ``upstream_failed`` rather than skipped.
+Rejecting the review **skips every downstream task**, matching
+:class:`~airflow.providers.standard.operators.hitl.ApprovalOperator`. Set
+``fail_on_reject=True`` to fail the task instead (generally discouraged).
+Letting ``approval_timeout`` expire fails the task (``HITLTimeoutError``).
 
 ``require_approval=True`` requires a string prompt: a decorated callable
 returning a ``Sequence[UserContent]`` raises ``TypeError`` before the LLM
@@ -133,6 +134,8 @@ Parameters
   means wait indefinitely.  Default ``None``.
 - ``allow_modifications``: If ``True``, the reviewer can change the chosen
   branch(es) before approving.  Default ``False``.
+- ``fail_on_reject``: If ``True``, a rejected review fails the task instead of
+  skipping every downstream task.  Generally discouraged.  Default ``False``.
 
 Logging
 -------

providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py

@@ -25,6 +25,7 @@
 
 from airflow.providers.common.ai.operators.llm import LLMOperator
 from airflow.providers.common.ai.utils.logging import log_run_summary
+from airflow.providers.standard.exceptions import HITLRejectException
 from airflow.providers.standard.operators.branch import BranchMixIn
 
 if TYPE_CHECKING:
@@ -46,14 +47,21 @@ class LLMBranchOperator(LLMOperator, BranchMixIn):
     :param system_prompt: System-level instructions for the LLM agent.
     :param allow_multiple_branches: When ``False`` (default) the LLM returns a
         single task ID. When ``True`` the LLM may return one or more task IDs.
+    :param fail_on_reject: If ``True``, a rejected review fails the task
+        instead of skipping every downstream task. Generally discouraged,
+        as for :class:`~airflow.providers.standard.operators.hitl.ApprovalOperator`.
+        Default ``False``.
     :param agent_params: Additional keyword arguments passed to the pydantic-ai
         ``Agent`` constructor (e.g. ``retries``, ``model_settings``, ``tools``).
 
     Human-in-the-Loop approval parameters are inherited from
     :class:`~airflow.providers.common.ai.operators.llm.LLMOperator`
     (``require_approval``, ``approval_timeout``, ``allow_modifications``).
     The task pauses after the LLM chooses the branch(es) and only skips the
-    unselected downstream tasks once a reviewer approves. The review form
+    unselected downstream tasks once a reviewer approves. Rejecting the
+    review skips every downstream task, matching
+    :class:`~airflow.providers.standard.operators.hitl.ApprovalOperator`;
+    set ``fail_on_reject=True`` to fail the task instead. The review form
     lists the valid downstream task IDs; with ``allow_modifications=True``
     the editable choice is rendered as a dropdown of those IDs (single-branch
     mode) or a multi-select of them (``allow_multiple_branches=True``), and
@@ -69,11 +77,13 @@ def __init__(
         self,
         *,
         allow_multiple_branches: bool = False,
+        fail_on_reject: bool = False,
         **kwargs: Any,
     ) -> None:
         kwargs.pop("output_type", None)
         super().__init__(**kwargs)
         self.allow_multiple_branches = allow_multiple_branches
+        self.fail_on_reject = fail_on_reject
 
     def execute(self, context: Context) -> str | Iterable[str] | None:
         if self.require_approval:
@@ -133,7 +143,13 @@ def execute(self, context: Context) -> str | Iterable[str] | None:
 
     def execute_complete(self, context: Context, generated_output: str, event: dict[str, Any]) -> Any:
         """Resume after human review, validating the reviewed choice before branching."""
-        output = super().execute_complete(context, generated_output, event)
+        try:
+            output = super().execute_complete(context, generated_output, event)
+        except HITLRejectException:
+            if self.fail_on_reject:
+                raise
+            self.log.info("Rejected. Skipping all downstream tasks...")
+            return self.do_branch(context, None)
         branches = self._parse_reviewed_branches(output)
         selected = {branches} if isinstance(branches, str) else set(branches)
         invalid = selected - self.downstream_task_ids

providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py

@@ -26,6 +26,7 @@
 from airflow.providers.common.ai.operators.llm import LLMOperator
 from airflow.providers.common.ai.operators.llm_branch import LLMBranchOperator
 from airflow.providers.common.compat.sdk import Param, ParamValidationError, TaskDeferred
+from airflow.providers.standard.exceptions import HITLRejectException
 
 from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS, AIRFLOW_V_3_3_PLUS
 
@@ -397,6 +398,30 @@ def test_execute_complete_approved_multiple_branches(self, mock_do_branch):
         assert result == ["task_a", "task_c"]
         mock_do_branch.assert_called_once_with(ctx, ["task_a", "task_c"])
 
+    @patch.object(LLMBranchOperator, "do_branch")
+    def test_execute_complete_reject_skips_all_downstream(self, mock_do_branch):
+        mock_do_branch.return_value = None
+        op = LLMBranchOperator(task_id="t", prompt="p", llm_conn_id="c")
+        op.downstream_task_ids = {"task_a", "task_b"}
+        event = {"chosen_options": ["Reject"], "responded_by_user": "admin"}
+        ctx = _make_context()
+
+        result = op.execute_complete(ctx, generated_output="task_a", event=event)
+
+        assert result is None
+        mock_do_branch.assert_called_once_with(ctx, None)
+
+    @patch.object(LLMBranchOperator, "do_branch")
+    def test_execute_complete_reject_fails_with_fail_on_reject(self, mock_do_branch):
+        op = LLMBranchOperator(task_id="t", prompt="p", llm_conn_id="c", fail_on_reject=True)
+        op.downstream_task_ids = {"task_a", "task_b"}
+        event = {"chosen_options": ["Reject"], "responded_by_user": "admin"}
+
+        with pytest.raises(HITLRejectException, match="rejected"):
+            op.execute_complete(_make_context(), generated_output="task_a", event=event)
+
+        mock_do_branch.assert_not_called()
+
     @patch.object(LLMBranchOperator, "do_branch")
     def test_execute_complete_with_modified_branch(self, mock_do_branch):
         """A reviewer-modified branch is used when it is a valid downstream task."""