fixup! Run loop gates and advance iterations atomically

Author: ashb

airflow-core/docs/authoring-and-scheduling/loops.rst

@@ -58,7 +58,7 @@ count is reached.
 This example improves an estimate of the square root of two until the error
 is small enough:
 
-.. exampleinclude:: /authoring-and-scheduling/examples/example_task_loops.py
+.. exampleinclude:: /../src/airflow/example_dags/example_task_loops.py
    :start-after: [START refine_estimate]
    :end-before: [END refine_estimate]
 
@@ -81,15 +81,15 @@ For a fixed-count loop, omit ``until``. The definition ``refine.loop(max_iterati
 three iterations, carrying results between them. Reaching the cap completes
 a fixed-count loop successfully. Its gate is named ``__loop_gate`` within the group.
 
-.. exampleinclude:: /authoring-and-scheduling/examples/example_task_loops.py
+.. exampleinclude:: /../src/airflow/example_dags/example_task_loops.py
    :start-after: [START fixed_loop]
    :end-before: [END fixed_loop]
 
 For a task-group function with arguments, supply them with ``.partial()`` before
 calling ``.loop()``. Use ``.override()`` to configure the group, for example to
 give another loop a different ``group_id``:
 
-.. exampleinclude:: /authoring-and-scheduling/examples/example_task_loops.py
+.. exampleinclude:: /../src/airflow/example_dags/example_task_loops.py
    :start-after: [START partial_override_loop]
    :end-before: [END partial_override_loop]
 
@@ -187,7 +187,7 @@ single value or structure. A zero-length expansion skips the mapped task and,
 under the gate's ``all_success`` rule, skips the gate; no next iteration is
 created.
 
-.. exampleinclude:: /authoring-and-scheduling/examples/example_task_loops.py
+.. exampleinclude:: /../src/airflow/example_dags/example_task_loops.py
    :start-after: [START mapped_loop]
    :end-before: [END mapped_loop]
 

airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py

@@ -99,7 +99,11 @@
 from airflow.models.dynamic_region import SENTINEL_REGION_ID, AmbiguousProducerError
 from airflow.models.hitl import HITLDetail
 from airflow.models.log import Log
-from airflow.models.task_coordinates import TaskCoordinateResolver, public_map_index_expression
+from airflow.models.task_coordinates import (
+    LOOP_GATE_OPERATOR,
+    TaskCoordinateResolver,
+    public_map_index_expression,
+)
 from airflow.models.taskinstance import TaskInstance as TI, _stop_remaining_tasks
 from airflow.models.taskreschedule import TaskReschedule
 from airflow.models.trigger import Trigger, handle_event_submit
@@ -172,6 +176,7 @@ def ti_run(
             TI.dag_id,
             TI.run_id,
             TI.task_id,
+            TI.region_id,
             TI.region_index,
             TI.try_number,
             TI.max_tries,
@@ -450,7 +455,7 @@ def ti_update_state(
             select(TI.dag_id, TI.run_id).where(
                 TI.id == task_instance_id,
                 TI.working_set.is_(True),
-                TI.operator == "LoopGateOperator",
+                TI.operator == LOOP_GATE_OPERATOR,
             )
         ).one_or_none()
         if gate_run is not None:
@@ -465,9 +470,12 @@ def ti_update_state(
             )
             if loop_gate is not None:
                 loop_context = TaskCoordinateResolver(dag_bag, session).loop_context(loop_gate)
-                if loop_context is None or loop_context[0].gate_task_id != loop_gate.task_id:
+                if loop_context is None:
+                    loop_gate = None
+                elif loop_context[0].gate_task_id != loop_gate.task_id:
                     raise HTTPException(status_code=409, detail={"reason": "invalid_loop_gate"})
-                loop_group = loop_context[0]
+                else:
+                    loop_group = loop_context[0]
 
     old = (
         select(

airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py

@@ -64,6 +64,7 @@
     IdentifyArchivedTaskStateUpdates,
 )
 from airflow.api_fastapi.execution_api.versions.v2026_10_30_xcom_params import (
+    AddPreviousIterationToXComFilterParams,
     AddRegionSelectorsToXComFilterParams,
 )
 
@@ -77,6 +78,7 @@
         AddLoopContext,
         AddTerminalStateRetryReasonField,
         AddMultiTeamToTIRunContext,
+        AddPreviousIterationToXComFilterParams,
         AddStoppedTaskReport,
         AddTaskInstanceRegionCoordinates,
         AddRegionSelectorsToXComFilterParams,

airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30_xcom_params.py

@@ -35,3 +35,14 @@ class AddRegionSelectorsToXComFilterParams(VersionChange):
         schema(GetXcomFilterParams).field("region_id").didnt_exist,
         schema(GetXcomFilterParams).field("region_index").didnt_exist,
     )
+
+
+class AddPreviousIterationToXComFilterParams(VersionChange):
+    """Add the `previous_iteration` field to GetXComSliceFilterParams and GetXcomFilterParams."""
+
+    description = __doc__
+
+    instructions_to_migrate_to_previous_version = (
+        schema(GetXComSliceFilterParams).field("previous_iteration").didnt_exist,
+        schema(GetXcomFilterParams).field("previous_iteration").didnt_exist,
+    )

…cheduling/examples/example_task_loops.py …rflow/example_dags/example_task_loops.pyairflow-core/docs/authoring-and-scheduling/examples/example_task_loops.py renamed to airflow-core/src/airflow/example_dags/example_task_loops.py airflow-core/docs/authoring-and-scheduling/examples/example_task_loops.py renamed to airflow-core/src/airflow/example_dags/example_task_loops.py

@@ -114,7 +114,7 @@
 from airflow.models.log import resolve_team_name
 from airflow.models.pool import normalize_pool_name_for_stats
 from airflow.models.serialized_dag import SerializedDagModel
-from airflow.models.task_coordinates import TaskCoordinateResolver
+from airflow.models.task_coordinates import LOOP_GATE_OPERATOR, TaskCoordinateResolver
 from airflow.models.taskinstance import TaskInstance
 from airflow.models.team import Team
 from airflow.models.trigger import TRIGGER_FAIL_REPR, Trigger, TriggerFailureReason, handle_event_submit
@@ -1679,7 +1679,8 @@ def process_executor_events(
 
                     task = dag.get_task(ti.task_id)
                 except Exception:
-                    if state == TaskInstanceState.SUCCESS and ti.operator == "LoopGateOperator":
+                    if state == TaskInstanceState.SUCCESS and ti.operator == LOOP_GATE_OPERATOR:
+                        cls.logger().exception("Failing loop gate %s: %s", ti, msg)
                         ti.task = None
                         ti.handle_failure(error=msg, session=session)
                         continue

airflow-core/src/airflow/jobs/scheduler_job_runner.py

@@ -16,7 +16,7 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Protocol
+from typing import TYPE_CHECKING, Protocol
 from uuid import UUID
 
 import attrs
@@ -49,6 +49,22 @@
     from airflow.serialization.definitions.dag import SerializedDAG, SerializedOperator
 
 
+LOOP_GATE_OPERATOR = "LoopGateOperator"
+"""Operator class name the Task SDK gives the gate task of a loop."""
+
+
+@attrs.frozen(kw_only=True)
+class _ProducerRequest:
+    dag_id: str
+    run_id: str
+    task_id: str
+    is_mapped: bool
+    context: ProducerContext | None
+    map_indexes: int | Collection[int] | None
+    region_id: UUID | None
+    region_index: int | None
+
+
 class TaskCoordinate(Protocol):
     """Stored task identity shared by live and retained task data."""
 
@@ -249,7 +265,12 @@ def has_regions(self, dag_id: str, run_id: str | None, task_id: str) -> bool:
         return bool(self.session.scalar(select(query.exists())))
 
     def resolve_dependency(self, caller: TaskInstance, task_id: str) -> tuple[TaskInstance, ...]:
-        producer = self.get_task(caller.dag_id, caller.run_id, task_id, dag_version_id=caller.dag_version_id)
+        try:
+            producer = self.get_task(
+                caller.dag_id, caller.run_id, task_id, dag_version_id=caller.dag_version_id
+            )
+        except TaskNotFound:
+            return self.resolve(dag_id=caller.dag_id, run_id=caller.run_id, task_id=task_id, caller=caller)
         loop = enclosing_loop(producer)
         caller_task = self.get_task(
             caller.dag_id, caller.run_id, caller.task_id, dag_version_id=caller.dag_version_id
@@ -278,6 +299,18 @@ def resolve_dependency(self, caller: TaskInstance, task_id: str) -> tuple[TaskIn
             return tuple(gates[:1])
         return self.resolve(dag_id=caller.dag_id, run_id=caller.run_id, task_id=task_id, caller=caller)
 
+    @staticmethod
+    def _filter_region_indexes(
+        query: Select, region_index: int | None, map_indexes: int | Collection[int] | None
+    ) -> Select:
+        if region_index is not None:
+            query = query.where(TaskInstance.region_index == region_index)
+        if isinstance(map_indexes, int):
+            query = query.where(TaskInstance.region_index == map_indexes)
+        elif map_indexes is not None:
+            query = query.where(TaskInstance.region_index.in_(map_indexes))
+        return query
+
     def _filter_legacy_producers(
         self,
         query: Select,
@@ -295,13 +328,7 @@ def _filter_legacy_producers(
             TaskInstance.task_id == task_id,
             TaskInstance.region_id == SENTINEL_REGION_ID,
         )
-        if region_index is not None:
-            query = query.where(TaskInstance.region_index == region_index)
-        if isinstance(map_indexes, int):
-            query = query.where(TaskInstance.region_index == map_indexes)
-        elif map_indexes is not None:
-            query = query.where(TaskInstance.region_index.in_(map_indexes))
-        return query
+        return self._filter_region_indexes(query, region_index, map_indexes)
 
     def _filter_removed_task_producers(
         self,
@@ -324,13 +351,7 @@ def _filter_removed_task_producers(
         )
         if region_id is not None:
             query = query.where(TaskInstance.region_id == region_id)
-        if region_index is not None:
-            query = query.where(TaskInstance.region_index == region_index)
-        if isinstance(map_indexes, int):
-            query = query.where(TaskInstance.region_index == map_indexes)
-        elif map_indexes is not None:
-            query = query.where(TaskInstance.region_index.in_(map_indexes))
-        return query
+        return self._filter_region_indexes(query, region_index, map_indexes)
 
     def _is_legacy_lookup(
         self, dag_id: str, run_id: str, task_id: str, region_id: UUID | None, previous_iteration: bool
@@ -351,7 +372,7 @@ def _build_producer_request(
         region_index: int | None,
         map_indexes: int | Collection[int] | None,
         previous_iteration: bool,
-    ) -> dict[str, Any] | None:
+    ) -> _ProducerRequest | None:
         shared_run = caller is not None and (caller.dag_id, caller.run_id) == (dag_id, run_id)
         try:
             task = (
@@ -378,16 +399,16 @@ def _build_producer_request(
             )
         if previous_iteration and context is None:
             raise ValueError("Previous-iteration lookup requires a shared loop scope")
-        return {
-            "dag_id": dag_id,
-            "run_id": run_id,
-            "task_id": task_id,
-            "is_mapped": task.get_needs_expansion(),
-            "context": context,
-            "map_indexes": map_indexes,
-            "region_id": region_id,
-            "region_index": region_index,
-        }
+        return _ProducerRequest(
+            dag_id=dag_id,
+            run_id=run_id,
+            task_id=task_id,
+            is_mapped=task.get_needs_expansion(),
+            context=context,
+            map_indexes=map_indexes,
+            region_id=region_id,
+            region_index=region_index,
+        )
 
     def resolve(
         self,
@@ -437,7 +458,7 @@ def resolve(
                     ).order_by(TaskInstance.region_index)
                 )
             )
-        return resolve_current_producers(**request, session=self.session)
+        return resolve_current_producers(**attrs.asdict(request, recurse=False), session=self.session)
 
     def select_skip_target_ids(
         self, *, caller: TaskInstance, task_id: str, map_indexes: int | None = None
@@ -520,4 +541,4 @@ def select_producer_ids(
                 region_index=region_index,
                 map_indexes=map_indexes,
             )
-        return select_current_producer_ids(**request, session=self.session)
+        return select_current_producer_ids(**attrs.asdict(request, recurse=False), session=self.session)

airflow-core/src/airflow/models/task_coordinates.py

@@ -51,6 +51,7 @@
 from airflow.api_fastapi.execution_api.app import lifespan
 from airflow.api_fastapi.execution_api.datamodels.taskinstance import TISuccessStatePayload
 from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken
+from airflow.api_fastapi.execution_api.routes import task_instances as task_instances_route
 from airflow.api_fastapi.execution_api.routes.task_instances import _emit_task_span, ti_update_state
 from airflow.api_fastapi.execution_api.routes.xcoms import set_xcom
 from airflow.api_fastapi.execution_api.security import require_auth
@@ -62,7 +63,7 @@
 from airflow.models.dagbag import DBDagBag
 from airflow.models.dynamic_region import DynamicRegion
 from airflow.models.log import Log
-from airflow.models.task_coordinates import TaskCoordinateResolver
+from airflow.models.task_coordinates import LOOP_GATE_OPERATOR, TaskCoordinateResolver
 from airflow.models.task_state_store import TaskStateStoreModel
 from airflow.models.taskinstance import TaskInstance, clear_task_instances
 from airflow.models.xcom import XComModel, XComModelV2
@@ -1798,6 +1799,7 @@ def complete():
         ) == [0, 1]
         assert session.scalar(select(XComModel).where(XComModel.task_id == gate.task_id)) is None
 
+    @pytest.mark.backend("mysql", "postgres")
     def test_loop_decision_rewrite_and_completion_use_same_lock_order(
         self, session, running_loop_gate, mocker
     ):
@@ -1905,6 +1907,45 @@ def test_late_loop_decision_cannot_recreate_consumed_signal(self, client, sessio
         assert response.status_code == 409
         assert session.scalar(select(XComModel).where(XComModel.task_id == gate.task_id)) is None
 
+    def test_loop_gate_state_update_error_rolls_back_next_pass_and_keeps_decision(
+        self, client, session, running_loop_gate, mocker
+    ):
+        gate = running_loop_gate
+        mocker.patch.object(
+            task_instances_route,
+            "_create_ti_state_update_query_and_update_state",
+            autospec=True,
+            side_effect=StaleDataError("injected state update failure"),
+        )
+
+        response = client.patch(
+            f"/execution/task-instances/{gate.id}/state",
+            json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()},
+        )
+
+        assert response.status_code == 500
+        session.expire_all()
+        assert gate.state == State.RUNNING
+        assert len(gate.dag_run.get_task_instances(session=session)) == 2
+        assert session.scalar(select(XComModel.value).where(XComModel.task_id == gate.task_id)) == "continue"
+
+    @pytest.mark.parametrize("state", ["success", "failed"])
+    def test_non_loop_operator_named_loop_gate_completes_normally(
+        self, client, session, create_task_instance, state
+    ):
+        ti = create_task_instance(task_id="user_gate", start_date=DEFAULT_START_DATE, state=State.RUNNING)
+        ti.operator = LOOP_GATE_OPERATOR
+        session.commit()
+
+        response = client.patch(
+            f"/execution/task-instances/{ti.id}/state",
+            json={"state": state, "end_date": DEFAULT_END_DATE.isoformat()},
+        )
+
+        assert response.status_code == 204
+        session.expire_all()
+        assert ti.state == state
+
     def test_loop_gate_inner_insert_error_is_not_swallowed(self, client, session, running_loop_gate, mocker):
         gate = running_loop_gate
         bulk_insert = mocker.patch.object(

airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py

@@ -1174,7 +1174,11 @@ def body():
         assert (gate.id != original_id) is bool(retries)
         callback.assert_not_called()
         if retries:
-            archived = session.scalar(select(TaskInstance).where(TaskInstance.id == original_id))
+            archived = session.scalar(
+                select(TaskInstance)
+                .where(TaskInstance.id == original_id)
+                .execution_options(include_all_attempts=True)
+            )
             assert archived.working_set is None
         assert len(dr.get_task_instances(session=session)) == 2