Run loop gates and advance iterations atomically

The scheduler, the API and the UI must never see every task of a loop finished
while its next pass does not exist yet. If we recorded a continuing gate
as successful and created the next pass afterwards, a crash between the two
would leave recovery unable to tell an unfinished advancement from a completed
decision. Gate completion and next-pass creation therefore happen in one
transaction under the DagRun lock, which also stops a duplicate completion from
advancing the loop twice.

The continue-or-stop decision is a reserved XCom stored under the gate's attempt
UUID. A retried gate starts without a decision, and a retired attempt cannot
advance the loop. An invalid decision moves the gate attempt to retry or failed
within the same request; rejecting the request instead would leave the attempt
running with nobody to finish it.

Gates use normal task dependencies and an explicit loop context. A barrier
that waited for every body task would override trigger rules that allow early
progress. Workers and dag.test supply the same coordinates, so decisions and
previous-iteration reads reach the intended gate and data.

Skipping downstream of a task from outside a loop skips every live pass.
Skipping only one would let later passes run work the author meant to skip.

Author: ashb

airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py

@@ -421,12 +421,24 @@ def safe_extract_from_orm(cls, data: Any) -> Any:
         return values
 
 
+class LoopContext(BaseModel):
+    """Pinned loop definition and enclosing iteration for a task execution."""
+
+    node_id: str
+    index: int = Field(ge=0)
+    max_iterations: int = Field(gt=0)
+    terminal_task_id: str
+    terminal_is_mapped: bool
+
+
 class TIRunContext(BaseModel):
     """Response schema for TaskInstance run context."""
 
     dag_run: DagRun
     """DAG run information for the task instance."""
 
+    loop: LoopContext | None = None
+
     task_reschedule_count: int = 0
     """How many times the task has been rescheduled."""
 

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

@@ -35,7 +35,7 @@
 from opentelemetry.trace import StatusCode
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 from pydantic import JsonValue, ValidationError
-from sqlalchemy import and_, func, or_, tuple_, update
+from sqlalchemy import and_, func, or_, tuple_, union, update
 from sqlalchemy.engine import CursorResult
 from sqlalchemy.exc import DataError, NoResultFound, SQLAlchemyError
 from sqlalchemy.orm import contains_eager, joinedload
@@ -57,6 +57,7 @@
 from airflow.api_fastapi.execution_api.datamodels.taskinstance import (
     DagRunNoteUpdatePayload,
     InactiveAssetsResponse,
+    LoopContext,
     PreviousTIResponse,
     PrevSuccessfulDagRunResponse,
     TaskBreadcrumbsResponse,
@@ -94,7 +95,7 @@
 from airflow.models.asset import AssetActive
 from airflow.models.base import ID_LEN
 from airflow.models.dag import DagModel
-from airflow.models.dagrun import DagRun as DR
+from airflow.models.dagrun import DagRun as DR, InvalidLoopDecision
 from airflow.models.dynamic_region import SENTINEL_REGION_ID, AmbiguousProducerError
 from airflow.models.hitl import HITLDetail
 from airflow.models.log import Log
@@ -107,7 +108,7 @@
 from airflow.state import get_state_backend
 from airflow.triggers.base import TriggerEvent
 from airflow.utils.sqlalchemy import get_dialect_name
-from airflow.utils.state import DagRunState, TaskInstanceState, TerminalTIState
+from airflow.utils.state import DagRunState, IntermediateTIState, TaskInstanceState, TerminalTIState
 
 router = VersionedAPIRouter()
 
@@ -331,6 +332,19 @@ def ti_run(
             should_retry=_is_eligible_to_retry(previous_state, ti.try_number, ti.max_tries),
             multi_team=conf.getboolean("core", "multi_team"),
         )
+        resolver = TaskCoordinateResolver(dag_bag, session)
+        if loop_context := resolver.loop_context(ti):
+            group, index = loop_context
+            terminal = resolver.get_task(
+                ti.dag_id, ti.run_id, group.terminal_task_id, dag_version_id=ti.dag_version_id
+            )
+            context.loop = LoopContext(
+                node_id=group.node_id,
+                index=index,
+                max_iterations=group.max_iterations,
+                terminal_task_id=group.terminal_task_id,
+                terminal_is_mapped=terminal.get_needs_expansion(),
+            )
 
         # Only set for lang-SDK (foreign-runtime) tasks with a captured TaskFlow arg
         # spec; the route excludes unset fields, keeping regular responses lean.
@@ -430,6 +444,32 @@ def ti_update_state(
             raise HTTPException(status_code=409, detail={"reason": "invalid_state"})
         return Response(status_code=status.HTTP_204_NO_CONTENT)
 
+    loop_group = None
+    loop_gate = None
+    if isinstance(ti_patch_payload, (TISuccessStatePayload, TITerminalStatePayload)):
+        gate_run = session.execute(
+            select(TI.dag_id, TI.run_id).where(
+                TI.id == task_instance_id,
+                TI.working_set.is_(True),
+                TI.operator == "LoopGateOperator",
+            )
+        ).one_or_none()
+        if gate_run is not None:
+            session.execute(
+                select(DR).where(DR.dag_id == gate_run.dag_id, DR.run_id == gate_run.run_id).with_for_update()
+            ).scalar_one()
+            loop_gate = session.scalar(
+                select(TI)
+                .where(TI.id == task_instance_id, TI.working_set.is_(True))
+                .with_for_update(of=TI)
+                .execution_options(populate_existing=True)
+            )
+            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:
+                    raise HTTPException(status_code=409, detail={"reason": "invalid_loop_gate"})
+                loop_group = loop_context[0]
+
     old = (
         select(
             TI.state,
@@ -517,6 +557,25 @@ def ti_update_state(
                 detail={"reason": "invalid_partition_key", "message": str(e)},
             ) from e
 
+    gate_completed = False
+    if (
+        loop_gate is not None
+        and loop_group is not None
+        and isinstance(ti_patch_payload, TISuccessStatePayload)
+    ):
+        try:
+            loop_gate.dag_run.complete_loop_gate(
+                loop_gate, loop_group, TaskInstanceState.SUCCESS, session=session
+            )
+            gate_completed = True
+        except InvalidLoopDecision as error:
+            log.warning("Loop gate success rejected", error=str(error))
+            ti_patch_payload = _build_rejected_gate_payload(
+                ti_patch_payload,
+                reason=f"Loop gate success rejected: {error}",
+                retry=_is_eligible_to_retry(previous_state, try_number, max_tries),
+            )
+
     # We exclude_unset to avoid updating fields that are not set in the payload
     data = ti_patch_payload.model_dump(
         exclude={"task_outlets", "outlet_events", "retry_delay_seconds", "retry_reason"},
@@ -540,6 +599,8 @@ def ti_update_state(
         # Let DataErrorHandler return a 422 instead of silently marking the TI FAILED below.
         raise
     except Exception:
+        if loop_gate is not None:
+            raise
         # Set a task to failed in case any unexpected exception happened during task state update
         log.exception(
             "Error updating Task Instance state. Setting the task to failed.",
@@ -590,6 +651,9 @@ def ti_update_state(
         # Defer to app-level SQLAlchemyError handler (returns HTTP 500).
         raise
 
+    if loop_gate is not None and loop_group is not None and not gate_completed:
+        loop_gate.dag_run.complete_loop_gate(loop_gate, loop_group, updated_state, session=session)
+
     if updated_state == TaskInstanceState.SUCCESS:
         if conf.getboolean("state_store", "clear_on_success"):
             scope = TaskScope(
@@ -623,6 +687,15 @@ def ti_update_state(
         callback()
 
 
+def _build_rejected_gate_payload(
+    payload: TISuccessStatePayload, *, reason: str, retry: bool
+) -> TIRetryStatePayload | TITerminalStatePayload:
+    carried = payload.model_dump(include={"end_date", "rendered_map_index"}, exclude_unset=True)
+    if retry:
+        return TIRetryStatePayload(state=IntermediateTIState.UP_FOR_RETRY, retry_reason=reason, **carried)
+    return TITerminalStatePayload(state=TerminalStateNonSuccess.FAILED, retry_reason=reason, **carried)
+
+
 def _emit_task_span(ti, state, *, resolver: TaskCoordinateResolver):
     # just to be safe
     if not ti.dag_run:
@@ -926,21 +999,21 @@ def ti_skip_downstream(
     try:
         # A bare task_id skips every TI of that task, so an already expanded mapped task
         # (e.g. one mapped over a literal list) is skipped too, not only map_index -1.
-        selected_ids = {
-            ti.id
-            for task in tasks
-            for ti in resolver.resolve(
-                dag_id=dag_id,
-                run_id=run_id,
-                task_id=task if isinstance(task, str) else task[0],
+        selects = [
+            resolver.select_skip_target_ids(
                 caller=caller,
+                task_id=task if isinstance(task, str) else task[0],
                 map_indexes=None if isinstance(task, str) else task[1],
             )
-        }
+            for task in tasks
+        ]
     except AmbiguousProducerError as error:
         raise HTTPException(status.HTTP_409_CONFLICT, str(error)) from error
     except ValueError as error:
         raise HTTPException(status.HTTP_400_BAD_REQUEST, str(error)) from error
+    if not selects:
+        return
+    selected_ids = set(session.scalars(union(*selects)))
 
     # Don't overwrite tasks that are already executing or finished.
     # See: https://github.com/apache/airflow/issues/59378
@@ -958,6 +1031,7 @@ def ti_skip_downstream(
     query = (
         update(TI)
         .where(
+            TI.working_set.is_(True),
             TI.id.in_(selected_ids),
             skippable_state_clause,
         )