You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by as...@apache.org on 2026/10/18 03:45:08 UTC
(airflow) branch v3-4-test updated: [v3-4-test] Allow loop gates to add human-readable decisions (#74806) (#74932)
This is an automated email from the ASF dual-hosted git repository.
ashb pushed a commit to branch v3-4-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-4-test by this push:
new d5d60800688 [v3-4-test] Allow loop gates to add human-readable decisions (#74806) (#74932)
d5d60800688 is described below
commit d5d608006888ac89a4b6a45435efea80cf9f99e2
Author: Ash Berlin-Taylor <as...@astronomer.io>
AuthorDate: Sun Oct 18 04:44:58 2026 +0100
[v3-4-test] Allow loop gates to add human-readable decisions (#74806) (#74932)
As it landed on main, a loop's gate only recorded a go/no-go decision for
the next iteration, meaning the UI could show that a loop stopped after N
iterations, but a user would have to go to the logs to find out why,
exactly, the loop ended when it did.
This adds a `LoopDecision` object that lets an `until` callable include
human-readable text to show alongside its decision.
As well as introducing this new type, it also changes how the decision
makes it to the server. Prior to this change, the decision was shipped to
the API server via a specially named XCom value, in a normal SetXCom
request, which was then read when the gate's TISuccessState request
arrived. That "worked", but was inherently fragile: a separate write
with its own locking, validation and cleanup on every non-success path.
The decision now travels on the gate's terminal state update. The server
acts on it in the transaction that marks the gate finished, so the
stop/continue signal never needs storing. A gate that fails because until
stayed false at max_iterations sends its reason with the failure, which is
when an explanation helps most.
Only the reason is persisted, as an XCom on the gate's attempt that the
server writes in that same transaction. Keying it to the attempt means
each try keeps its own reason, and a cleared or retried gate starts with
none, so no clear path has to remember to delete a stale one. Task state
was considered instead, but it is keyed by the iteration rather than the
attempt, so a reason outlived the clears and tries it described, and
clear_on_success and retention could remove it. The server writes the
XCom directly to the metadata DB (not to a configured backend, so the UI
can always show it), and the reserved `_airflow_loop_` prefix stops users
and tasks forging it. The XCom list endpoint leaves those keys out when
asked, which the UI does.
* Show loop decision reasons in the UI
A finished loop's page shows the reason its final iteration gave, and every
task instance in an iteration shows that iteration's reason. The iteration
pickers are left alone because their rows are one-line labels.
A reason can run to 512 bytes over several lines, which pushed the header
stats around, so the headers cut it to one line and the tooltip shows it
in full.
(cherry picked from commit 4c4bab01e0635933fde25f8be14613112d1e2b05)
---
.../docs/authoring-and-scheduling/loops.rst | 30 ++-
airflow-core/pyproject.toml | 2 +-
.../api_fastapi/core_api/datamodels/ui/grid.py | 14 +-
.../api_fastapi/core_api/openapi/_private_ui.yaml | 5 +
.../core_api/openapi/v2-rest-api-generated.yaml | 11 +
.../api_fastapi/core_api/routes/public/xcom.py | 18 ++
.../core_api/services/public/task_instances.py | 11 +-
.../api_fastapi/core_api/services/ui/grid.py | 18 +-
.../execution_api/datamodels/taskinstance.py | 21 ++
.../execution_api/routes/task_instances.py | 22 +-
.../api_fastapi/execution_api/routes/xcoms.py | 58 +----
.../api_fastapi/execution_api/versions/__init__.py | 2 +
.../execution_api/versions/v2026_10_30.py | 12 +
.../src/airflow/example_dags/example_task_loops.py | 6 +-
airflow-core/src/airflow/models/dagrun.py | 44 ++--
airflow-core/src/airflow/models/dynamic_region.py | 3 +-
.../src/airflow/ui/openapi-gen/queries/common.ts | 5 +-
.../ui/openapi-gen/queries/ensureQueryData.ts | 6 +-
.../src/airflow/ui/openapi-gen/queries/prefetch.ts | 6 +-
.../src/airflow/ui/openapi-gen/queries/queries.ts | 6 +-
.../src/airflow/ui/openapi-gen/queries/suspense.ts | 6 +-
.../airflow/ui/openapi-gen/requests/schemas.gen.ts | 5 +
.../ui/openapi-gen/requests/services.gen.ts | 2 +
.../airflow/ui/openapi-gen/requests/types.gen.ts | 8 +
.../airflow/ui/public/i18n/locales/en/common.json | 1 +
.../ui/src/components/LoopDecisionReason.tsx | 37 +++
.../GroupTaskInstance/useLoopOutcomeStats.test.tsx | 64 ++++++
.../GroupTaskInstance/useLoopOutcomeStats.tsx | 16 +-
.../ui/src/pages/TaskInstance/Header.test.tsx | 34 ++-
.../airflow/ui/src/pages/TaskInstance/Header.tsx | 12 +
.../src/airflow/ui/src/pages/XCom/XCom.test.tsx | 18 ++
.../src/airflow/ui/src/pages/XCom/XCom.tsx | 1 +
.../core_api/routes/public/test_task_instances.py | 10 +-
.../core_api/routes/public/test_xcom.py | 48 ++++
.../api_fastapi/core_api/routes/ui/test_grid.py | 40 +++-
.../versions/head/test_task_instances.py | 251 +++++++--------------
.../execution_api/versions/head/test_xcoms.py | 53 +----
airflow-core/tests/unit/models/conftest.py | 8 +-
airflow-core/tests/unit/models/test_dagrun.py | 36 ++-
go-sdk/pkg/execution/genmodels/models.gen.go | 22 +-
go-sdk/schema/supervisor-schema.json | 49 ++++
task-sdk/docs/api.rst | 2 +
task-sdk/docs/index.rst | 1 +
task-sdk/src/airflow/sdk/__init__.py | 3 +
task-sdk/src/airflow/sdk/__init__.pyi | 2 +
task-sdk/src/airflow/sdk/api/client.py | 16 +-
.../src/airflow/sdk/api/datamodels/_generated.py | 14 ++
.../src/airflow/sdk/definitions/_internal/loop.py | 10 +-
.../sdk/definitions/decorators/task_group.py | 6 +-
task-sdk/src/airflow/sdk/definitions/loop.py | 33 +++
task-sdk/src/airflow/sdk/execution_time/comms.py | 2 +
task-sdk/src/airflow/sdk/execution_time/loop.py | 70 ++++--
.../airflow/sdk/execution_time/schema/schema.json | 49 ++++
.../execution_time/schema/versions/v2026_10_30.py | 12 +-
.../src/airflow/sdk/execution_time/supervisor.py | 2 +
.../src/airflow/sdk/execution_time/task_runner.py | 11 +-
task-sdk/tests/task_sdk/api/test_client.py | 27 +++
.../tests/task_sdk/execution_time/test_loop.py | 214 +++++++++++-------
.../task_sdk/execution_time/test_supervisor.py | 24 +-
ts-sdk/schema/supervisor-schema.json | 49 ++++
ts-sdk/src/generated/supervisor.ts | 14 ++
uv.lock | 2 +-
62 files changed, 1121 insertions(+), 463 deletions(-)
diff --git a/airflow-core/docs/authoring-and-scheduling/loops.rst b/airflow-core/docs/authoring-and-scheduling/loops.rst
index 67cd71a2a4f..161567a6935 100644
--- a/airflow-core/docs/authoring-and-scheduling/loops.rst
+++ b/airflow-core/docs/authoring-and-scheduling/loops.rst
@@ -51,7 +51,7 @@ function. Supply a positive integer ``max_iterations`` to limit the number of
iterations. An optional ``until`` callable decides when to stop early.
Airflow evaluates ``until`` in a gate task downstream of the body's terminal task.
-It must return a ``bool``:
+It must return a ``bool``, or a ``LoopDecision`` as described in :ref:`loops-decision-reason`:
* ``True`` means stop: the condition has been met.
* ``False`` means continue, provided another iteration is allowed.
@@ -110,6 +110,27 @@ give another loop a different ``group_id``:
:start-after: [START partial_override_loop]
:end-before: [END partial_override_loop]
+.. _loops-decision-reason:
+
+Explain a decision
+------------------
+
+As well as a ``bool``, ``until`` can return a ``LoopDecision``, which lets you provide free text
+for display to users in the UI. The text is shown as plain text: HTML and Markdown are not rendered.
+
+Returning a ``LoopDecision(stop=..., reason=...)`` instead of a ``bool`` lets you include more
+detailed information about why the loop continued or finished. ``stop`` has the same meaning as the
+plain ``bool`` return value. The ``reason`` field is optional text, written by the Dag author, that
+the Airflow UI shows next to the iteration's outcome, for example ``error 0.17, needs < 1e-06``. It
+is limited to 512 bytes of UTF-8; anything longer will result in the gate task failing outright. This
+decision will then show up in the UI on Task Instance Details section of the gate task.
+
+The reason is also recorded when the loop fails because it reached ``max_iterations`` without
+stopping.
+
+The ``accurate_enough`` condition in the ``refine_estimate`` example above returns a
+``LoopDecision``, so each iteration of that loop shows its error next to its outcome.
+
Read the loop context
=====================
@@ -293,9 +314,10 @@ counts iterations.
An exception in ``until`` fails the gate task and appears in its logs; the
gate is retried according to ``retries``. An ``until`` that returns anything
-other than a ``bool``, and a loop that reaches its iteration limit without
-meeting ``until``, fail the gate without retrying, because a retry would see the
-same result. In the second case the final gate fails; downstream tasks with the
+other than a ``bool`` or a valid ``LoopDecision``, a reason longer than 512
+bytes, and a loop that reaches its iteration limit without meeting ``until``,
+fail the gate without retrying, because a retry would see the same result. When
+the loop reaches its limit, the final gate fails; downstream tasks with the
default ``all_success`` trigger rule will not run, as described under
non-convergence above.
diff --git a/airflow-core/pyproject.toml b/airflow-core/pyproject.toml
index c475d71da89..839bf9a7cdd 100644
--- a/airflow-core/pyproject.toml
+++ b/airflow-core/pyproject.toml
@@ -135,7 +135,7 @@ dependencies = [
'pendulum>=3.1.0',
"pluggy>=1.5.0",
"psutil>=5.8.0",
- "pydantic>=2.11.0",
+ "pydantic>=2.12.0",
# Pygments 2.19.0 improperly renders .ini files with dictionaries as values
# See https://github.com/pygments/pygments/issues/2834
"pygments>=2.0.1,!=2.19.0",
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/grid.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/grid.py
index 3e2c391df65..b584fd7c4ca 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/grid.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/ui/grid.py
@@ -18,9 +18,13 @@
from __future__ import annotations
from datetime import datetime
-from typing import Literal
+from typing import Annotated, Literal
from uuid import UUID
+from pydantic import Field
+from pydantic.json_schema import SkipJsonSchema
+
+from airflow.api_fastapi.common.region import _drop_null_default
from airflow.api_fastapi.core_api.base import BaseModel
from airflow.utils.state import TaskInstanceState
@@ -57,6 +61,14 @@ class LoopIterationSummary(BaseModel):
state: TaskInstanceState | None = None
start_date: datetime | None = None
end_date: datetime | None = None
+ decision_reason: Annotated[
+ SkipJsonSchema[None] | str,
+ Field(
+ description="Why the loop's gate continued or stopped the loop at this iteration, if its until callable said.",
+ exclude_if=lambda reason: reason is None,
+ json_schema_extra=_drop_null_default,
+ ),
+ ] = None
class LoopInvocationResponse(BaseModel):
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml
index f80dd4b4f4b..150e27d3beb 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml
+++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml
@@ -4969,6 +4969,11 @@ components:
format: date-time
- type: 'null'
title: End Date
+ decision_reason:
+ type: string
+ title: Decision Reason
+ description: Why the loop's gate continued or stopped the loop at this iteration,
+ if its until callable said.
type: object
required:
- index
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
index ed84dbdc200..87e197465b5 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
+++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
@@ -7391,6 +7391,17 @@ paths:
minimum: -1
- type: 'null'
title: Region Index
+ - name: exclude_internal_keys
+ in: query
+ required: false
+ schema:
+ type: boolean
+ description: Leave out XComs whose keys Airflow reserves for itself, such
+ as loop decision reasons.
+ default: false
+ title: Exclude Internal Keys
+ description: Leave out XComs whose keys Airflow reserves for itself, such
+ as loop decision reasons.
- name: limit
in: query
required: false
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/xcom.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/xcom.py
index 792988a4187..0adaefa2de4 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/xcom.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/xcom.py
@@ -63,6 +63,7 @@ from airflow.api_fastapi.logging.decorators import action_logging
from airflow.exceptions import TaskNotFound
from airflow.models import DagRun as DR
from airflow.models.dag import DagModel
+from airflow.models.dynamic_region import LOOP_XCOM_PREFIX
from airflow.models.task_coordinates import TaskCoordinateResolver, public_map_index_expression
from airflow.models.taskinstance import TaskInstance
from airflow.models.xcom import XComModel, build_xcom_read_query, select_producers, xcom_entity
@@ -231,6 +232,12 @@ def get_xcom_entries(
map_index: Annotated[int | None, Query(ge=-1)] = None,
region_id: Annotated[UUID | None, Query()] = None,
region_index: Annotated[int | None, Query(ge=-1)] = None,
+ exclude_internal_keys: Annotated[
+ bool,
+ Query(
+ description="Leave out XComs whose keys Airflow reserves for itself, such as loop decision reasons."
+ ),
+ ] = False,
) -> XComCollectionResponse:
"""
Get all XCom entries.
@@ -285,6 +292,8 @@ def get_xcom_entries(
query = query.where(public_map_index == map_index)
if xcom_key is not None:
query = query.where(entity.key == xcom_key)
+ if exclude_internal_keys:
+ query = query.where(~entity.key.startswith(LOOP_XCOM_PREFIX, autoescape=True))
query, total_entries = paginated_select(
statement=query,
@@ -316,6 +325,12 @@ def get_xcom_entries(
return XComCollectionResponse(xcom_entries=entries, total_entries=total_entries)
+def _reject_reserved_loop_key(key: str) -> None:
+ # Loop gates write these keys through the Execution API, which validates their values.
+ if key.startswith(LOOP_XCOM_PREFIX):
+ raise HTTPException(status.HTTP_400_BAD_REQUEST, "Reserved loop XCom key")
+
+
@xcom_router.post(
"",
status_code=status.HTTP_201_CREATED,
@@ -342,6 +357,8 @@ def create_xcom_entry(
"""Create an XCom entry."""
from airflow.models.dagrun import DagRun
+ _reject_reserved_loop_key(request_body.key)
+
dag_run = session.scalar(select(DagRun).where(DagRun.dag_id == dag_id, DagRun.run_id == dag_run_id))
# Validate Dag ID
dag = get_dag_for_run_or_latest_version(dag_bag, dag_run, dag_id, session)
@@ -445,6 +462,7 @@ def update_xcom_entry(
dag_bag: DagBagDep,
) -> XComResponseNative:
"""Update an existing XCom entry."""
+ _reject_reserved_loop_key(xcom_key)
scope = resolve_task_scope(
dag_id=dag_id,
run_id=dag_run_id,
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py b/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py
index ed2b3316a97..b348176f5f9 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py
@@ -25,7 +25,7 @@ import structlog
from fastapi import HTTPException, status
from fastapi.exceptions import RequestValidationError
from pydantic import ValidationError
-from sqlalchemy import delete, select, tuple_
+from sqlalchemy import select, tuple_
from sqlalchemy.orm import joinedload
from sqlalchemy.orm.session import Session
@@ -57,7 +57,6 @@ from airflow.models.dag import DagModel
from airflow.models.dagbag import DBDagBag
from airflow.models.dagrun import lock_dag_runs
from airflow.models.dynamic_region import (
- LOOP_DECISION_KEY,
SENTINEL_REGION_ID,
DynamicRegion,
load_region_ancestry,
@@ -65,13 +64,11 @@ from airflow.models.dynamic_region import (
)
from airflow.models.renderedtifields import load_legacy_rendered_fields
from airflow.models.task_coordinates import (
- LOOP_GATE_OPERATOR,
TaskCoordinateResolver,
enclosing_loop,
public_map_index_expression,
)
from airflow.models.taskinstance import TaskInstance as TI, clear_task_instances, select_loop_clear_scope
-from airflow.models.xcom import XComModelV2
from airflow.serialization.definitions.dag import SerializedDAG
from airflow.state.metastore import _get_db_backend
from airflow.utils.state import TaskInstanceState
@@ -518,12 +515,6 @@ def _patch_selected_task_state(
[ti for ti in tis if ti.id not in superseded], later_loop_iterations=False, session=session
)
for ti in changed:
- if ti.operator == LOOP_GATE_OPERATOR:
- session.execute(
- delete(XComModelV2).where(
- XComModelV2.task_instance_id == ti.id, XComModelV2.key == LOOP_DECISION_KEY
- )
- )
if ti.state == TaskInstanceState.RESTARTING and ti.id in superseded:
ti.complete_restart(session=session, terminal_outcome=data["new_state"])
else:
diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py
index 10becd3a2f1..12cf433d84e 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/services/ui/grid.py
@@ -35,8 +35,9 @@ from airflow.api_fastapi.core_api.datamodels.ui.grid import (
LoopSummaryResponse,
)
from airflow.api_fastapi.core_api.services.ui.task_group import get_task_group_children_getter
-from airflow.models.dynamic_region import DynamicRegion, loop_position
+from airflow.models.dynamic_region import LOOP_DECISION_REASON_KEY, DynamicRegion, loop_position
from airflow.models.taskinstance import TaskInstance
+from airflow.models.xcom import XComModelV2
from airflow.serialization.definitions.baseoperator import SerializedBaseOperator
from airflow.serialization.definitions.mappedoperator import SerializedMappedOperator
from airflow.serialization.definitions.taskgroup import SerializedLoopTaskGroup, SerializedTaskGroup
@@ -268,6 +269,7 @@ class LoopRunInputs:
group: SerializedLoopTaskGroup | None = None
loop_definitions: dict[UUID | None, SerializedLoopTaskGroup] = field(default_factory=dict)
regions: dict[UUID, DynamicRegion] = field(default_factory=dict)
+ gate_reasons: dict[UUID, str] = field(default_factory=dict)
error: HTTPException | None = None
@@ -332,6 +334,7 @@ def load_loop_run_inputs(
if wanted_task_ids:
member_rows = session.execute(
select(
+ TaskInstance.id,
TaskInstance.run_id,
TaskInstance.task_id,
TaskInstance.region_id,
@@ -432,6 +435,9 @@ def summarize_loop_run(
LoopIterationSummary(
index=index,
task_count=len(tasks),
+ decision_reason=next(
+ (inputs.gate_reasons[ti.id] for ti in tasks if ti.id in inputs.gate_reasons), None
+ ),
state=agg_state(states),
start_date=min(starts) if starts else None,
end_date=max(ends) if ends and all(s in State.finished for s in states) else None,
@@ -493,4 +499,14 @@ def loop_run_summaries(
) -> list[LoopSummaryResponse]:
"""Summarize one run's loop invocations, loading its inputs through the batch loader."""
inputs = load_loop_run_inputs([run], group_id, session=session, dag_bag=dag_bag)[run.run_id]
+ gate_task_ids = {definition.gate_task_id for definition in inputs.loop_definitions.values()}
+ if gate_ids := [ti.id for ti in inputs.members if ti.task_id in gate_task_ids]:
+ inputs.gate_reasons = {
+ gate_id: reason
+ for gate_id, reason in session.execute(
+ select(XComModelV2.task_instance_id, XComModelV2.value).where(
+ XComModelV2.task_instance_id.in_(gate_ids), XComModelV2.key == LOOP_DECISION_REASON_KEY
+ )
+ )
+ }
return summarize_loop_run(inputs, group_id, loop_region_id=loop_region_id)
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py
index 81000c67d30..cda3ecd9a14 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py
@@ -29,6 +29,7 @@ from pydantic import (
Tag,
TypeAdapter,
WithJsonSchema,
+ field_validator,
model_validator,
)
@@ -39,6 +40,7 @@ from airflow.api_fastapi.execution_api.datamodels.asset import AssetProfile
from airflow.api_fastapi.execution_api.datamodels.connection import ConnectionResponse
from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import TaskArgBinding
from airflow.api_fastapi.execution_api.datamodels.variable import VariableResponse
+from airflow.models.dynamic_region import LOOP_DECISION_REASON_MAX_BYTES
from airflow.utils.state import (
DagRunState,
IntermediateTIState,
@@ -88,6 +90,23 @@ class TerminalStateNonSuccess(str, Enum):
SERVER_TERMINATED = "server_terminated"
+class LoopDecisionPayload(StrictBaseModel):
+ """A loop gate's evaluation of its ``until`` callable."""
+
+ stop: bool
+ reason: str | None = None
+
+ @field_validator("reason")
+ @classmethod
+ def _check_reason(cls, reason: str | None) -> str | None:
+ if reason is not None:
+ if len(reason.encode()) > LOOP_DECISION_REASON_MAX_BYTES:
+ raise ValueError(f"must be at most {LOOP_DECISION_REASON_MAX_BYTES} bytes of UTF-8")
+ if "\x00" in reason:
+ raise ValueError("cannot contain NUL bytes")
+ return reason
+
+
class TITerminalStatePayload(StrictBaseModel):
"""Report a terminal outcome other than SUCCESS, or acknowledge server-requested termination."""
@@ -99,6 +118,7 @@ class TITerminalStatePayload(StrictBaseModel):
retry_reason: str | None = None
hostname: str | None = None
pid: int | None = None
+ loop_decision: LoopDecisionPayload | None = None
class TISuccessStatePayload(StrictBaseModel):
@@ -122,6 +142,7 @@ class TISuccessStatePayload(StrictBaseModel):
task_outlets: Annotated[list[AssetProfile], Field(default_factory=list)]
outlet_events: Annotated[list[dict[str, Any]], Field(default_factory=list)]
rendered_map_index: str | None = None
+ loop_decision: LoopDecisionPayload | None = None
class TITargetStatePayload(StrictBaseModel):
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
index 3f462882ed9..3e204488fd3 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
@@ -598,9 +598,15 @@ def ti_update_state(
and loop_group is not None
and isinstance(ti_patch_payload, TISuccessStatePayload)
):
+ decision = ti_patch_payload.loop_decision
try:
loop_gate.dag_run.complete_loop_gate(
- loop_gate, loop_group, TaskInstanceState.SUCCESS, session=session
+ loop_gate,
+ loop_group,
+ TaskInstanceState.SUCCESS,
+ stop=decision.stop if decision is not None else None,
+ reason=decision.reason if decision is not None else None,
+ session=session,
)
gate_completed = True
except InvalidLoopDecision as error:
@@ -613,7 +619,7 @@ def ti_update_state(
# 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"},
+ exclude={"task_outlets", "outlet_events", "retry_delay_seconds", "retry_reason", "loop_decision"},
exclude_unset=True,
)
if "rendered_map_index" in data:
@@ -687,7 +693,17 @@ def ti_update_state(
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)
+ decision = (
+ ti_patch_payload.loop_decision if isinstance(ti_patch_payload, TITerminalStatePayload) else None
+ )
+ loop_gate.dag_run.complete_loop_gate(
+ loop_gate,
+ loop_group,
+ updated_state,
+ stop=None,
+ reason=decision.reason if decision is not None else None,
+ session=session,
+ )
if updated_state == TaskInstanceState.SUCCESS:
if conf.getboolean("state_store", "clear_on_success"):
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py
index ad9d1d42939..22ba2bc5665 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py
@@ -44,9 +44,8 @@ from airflow.api_fastapi.execution_api.in_process import TRIGGER_TI_ID_HEADER
from airflow.api_fastapi.execution_api.security import _IN_PROCESS_TRIGGERER_ID, CurrentTIToken
from airflow.api_fastapi.execution_api.versions.v2026_10_30 import IdentifyArchivedTaskStateUpdates
from airflow.exceptions import TaskNotFound
-from airflow.models.dagrun import DagRun, InvalidLoopDecision, validate_loop_decision
+from airflow.models.dagrun import DagRun
from airflow.models.dynamic_region import (
- LOOP_DECISION_KEY,
LOOP_XCOM_PREFIX,
SENTINEL_REGION_ID,
AmbiguousProducerError,
@@ -659,11 +658,11 @@ def get_xcom(
[
(
status.HTTP_400_BAD_REQUEST,
- "The key is empty, the value is too large to map, or is unserializable",
+ "The key is empty or reserved, the value is too large to map, or is unserializable",
),
(
status.HTTP_409_CONFLICT,
- "Several live task instances match the producer, or the loop gate is no longer running",
+ "Several live task instances match the producer",
),
]
),
@@ -674,7 +673,6 @@ def set_xcom(
task_id: str,
key: Annotated[str, Path(min_length=1)],
session: SessionDep,
- dag_bag: DagBagDep,
value: Annotated[
JsonValue,
Body(
@@ -705,43 +703,7 @@ def set_xcom(
"""Set an Airflow XCom."""
from airflow.configuration import conf
- region_id: UUID | None = None
- region_index: int | None = None
- if key == LOOP_DECISION_KEY:
- caller = session.get(TaskInstance, token.id)
- if (
- value not in ("continue", "stop")
- or mapped_length is not None
- or dag_result
- or map_index != -1
- or caller is None
- or (dag_id, run_id, task_id) != (caller.dag_id, caller.run_id, caller.task_id)
- ):
- raise HTTPException(status.HTTP_400_BAD_REQUEST, "Invalid loop decision write")
- session.execute(
- select(DagRun)
- .where(DagRun.dag_id == caller.dag_id, DagRun.run_id == caller.run_id)
- .with_for_update()
- ).scalar_one()
- caller = session.scalar(
- select(TaskInstance)
- .where(TaskInstance.id == token.id)
- .with_for_update()
- .execution_options(populate_existing=True)
- )
- if caller is None:
- raise HTTPException(status.HTTP_409_CONFLICT, "Loop gate is no longer running")
- loop = TaskCoordinateResolver(dag_bag, session).loop_context(caller)
- if loop is None or loop[0].gate_task_id != caller.task_id:
- raise HTTPException(status.HTTP_400_BAD_REQUEST, "Loop decisions require the pinned loop gate")
- if caller.state != TaskInstanceState.RUNNING or caller.working_set is not True:
- raise HTTPException(status.HTTP_409_CONFLICT, "Loop gate is no longer running")
- try:
- validate_loop_decision(value, caller.region_index, loop[0])
- except InvalidLoopDecision as error:
- raise HTTPException(status.HTTP_400_BAD_REQUEST, str(error)) from error
- region_id, region_index = caller.region_id, caller.region_index
- elif key.startswith(LOOP_XCOM_PREFIX):
+ if key.startswith(LOOP_XCOM_PREFIX):
raise HTTPException(status.HTTP_400_BAD_REQUEST, "Reserved loop XCom key")
# Validate that the provided key is not empty
@@ -790,8 +752,6 @@ def set_xcom(
run_id=run_id,
task_id=task_id,
map_index=map_index,
- region_id=region_id,
- region_index=region_index,
session=session,
),
serialize=False,
@@ -857,8 +817,6 @@ def _find_writer_id(
run_id: str,
task_id: str,
map_index: int,
- region_id: UUID | None = None,
- region_index: int | None = None,
session: SessionDep,
) -> UUID | None:
"""Resolve the attempt a write targets: the caller's own, or the live attempt of another task."""
@@ -866,9 +824,7 @@ def _find_writer_id(
TaskInstance.dag_id == dag_id,
TaskInstance.run_id == run_id,
TaskInstance.task_id == task_id,
- *build_coordinate_filters(
- TaskInstance, map_index=map_index, region_id=region_id, region_index=region_index
- ),
+ *build_coordinate_filters(TaskInstance, map_index=map_index),
)
own = session.scalar(
select(TaskInstance.id)
@@ -894,8 +850,6 @@ def _get_writer_id(
run_id: str,
task_id: str,
map_index: int,
- region_id: UUID | None = None,
- region_index: int | None = None,
session: SessionDep,
) -> UUID:
owner = _find_writer_id(
@@ -904,8 +858,6 @@ def _get_writer_id(
run_id=run_id,
task_id=task_id,
map_index=map_index,
- region_id=region_id,
- region_index=region_index,
session=session,
)
if owner is None:
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
index aa4d1028d2f..09515de289a 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
@@ -58,6 +58,7 @@ from airflow.api_fastapi.execution_api.versions.v2026_10_30 import (
AddDagRunNoteUpdateEndpoint,
AddFirstTaskRescheduleStartDateField,
AddLoopContext,
+ AddLoopDecisionField,
AddMultiTeamToTIRunContext,
AddStoppedTaskReport,
AddTaskInstanceRegionCoordinates,
@@ -76,6 +77,7 @@ bundle = VersionBundle(
AddCallbackRunEndpoint,
AddDagRunNoteUpdateEndpoint,
AddLoopContext,
+ AddLoopDecisionField,
AddTerminalStateRetryReasonField,
AddMultiTeamToTIRunContext,
AddPreviousIterationToXComFilterParams,
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py
index ec4b6ddf125..d6ac7dffd13 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py
@@ -32,6 +32,7 @@ from airflow.api_fastapi.execution_api.datamodels.taskinstance import (
TaskInstance,
TerminalStateNonSuccess,
TIRunContext,
+ TISuccessStatePayload,
TITerminalStatePayload,
)
@@ -122,6 +123,17 @@ class AddTerminalStateRetryReasonField(VersionChange):
)
+class AddLoopDecisionField(VersionChange):
+ """Add the `loop_decision` field a loop gate reports with its terminal state."""
+
+ description = __doc__
+
+ instructions_to_migrate_to_previous_version = (
+ schema(TISuccessStatePayload).field("loop_decision").didnt_exist,
+ schema(TITerminalStatePayload).field("loop_decision").didnt_exist,
+ )
+
+
class AddMultiTeamToTIRunContext(VersionChange):
"""Add ``multi_team`` so a worker can determine multi-team (e.g. for plugin scoping) without needing to trust its own config."""
diff --git a/airflow-core/src/airflow/example_dags/example_task_loops.py b/airflow-core/src/airflow/example_dags/example_task_loops.py
index 80e41a41468..8b178c4f6a2 100644
--- a/airflow-core/src/airflow/example_dags/example_task_loops.py
+++ b/airflow-core/src/airflow/example_dags/example_task_loops.py
@@ -19,7 +19,7 @@
from __future__ import annotations
# [START refine_estimate]
-from airflow.sdk import dag, task, task_group
+from airflow.sdk import LoopDecision, dag, task, task_group
@dag(schedule=None, catchup=False, tags=["example"])
@@ -38,7 +38,9 @@ def refine_estimate():
evaluate(improve())
def accurate_enough(*, loop):
- return loop.result["error"] < 0.000001
+ error = loop.result["error"]
+ threshold = 0.000001
+ return LoopDecision(stop=error < threshold, reason=f"error {error:.2g}, needs < {threshold:g}")
@task
def finished():
diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py
index 7d41528df5c..e052006b69e 100644
--- a/airflow-core/src/airflow/models/dagrun.py
+++ b/airflow-core/src/airflow/models/dagrun.py
@@ -90,7 +90,7 @@ from airflow.models.backfill import Backfill
from airflow.models.base import Base, StringID
from airflow.models.deadline_alert import DeadlineAlert as DeadlineAlertModel
from airflow.models.dynamic_region import (
- LOOP_DECISION_KEY,
+ LOOP_DECISION_REASON_KEY,
SENTINEL_REGION_ID,
DynamicRegion,
public_region_filter,
@@ -152,14 +152,14 @@ class InvalidLoopDecision(ValueError):
"""A gate cannot complete successfully with the available decision."""
-def validate_loop_decision(decision: str | None, region_index: int, group: SerializedLoopTaskGroup) -> None:
- """Raise :class:`InvalidLoopDecision` unless ``decision`` is valid for the gate of pass ``region_index``."""
+def validate_loop_decision(stop: bool | None, region_index: int, group: SerializedLoopTaskGroup) -> None:
+ """Raise :class:`InvalidLoopDecision` unless ``stop`` is valid for the gate of pass ``region_index``."""
at_limit = region_index + 1 >= group.max_iterations
- if decision not in ("continue", "stop"):
+ if stop is None:
raise InvalidLoopDecision("Successful loop gate requires a decision")
- if decision == "continue" and at_limit:
+ if not stop and at_limit:
raise InvalidLoopDecision("Loop cannot continue beyond its iteration limit")
- if not group.has_until and (decision == "stop") != at_limit:
+ if not group.has_until and stop != at_limit:
raise InvalidLoopDecision("Fixed-count loop decision does not match its iteration limit")
@@ -2323,24 +2323,34 @@ class DagRun(Base, LoggingMixin):
group: SerializedLoopTaskGroup,
state: TaskInstanceState,
*,
+ stop: bool | None,
+ reason: str | None,
session: Session,
) -> None:
- """Consume a gate decision while the caller holds the DagRun and TI locks."""
- from airflow.models.xcom import XComModelV2
+ """
+ Apply a gate's decision and record its reason while the caller holds the DagRun and TI locks.
+
+ A gate whose loop already has a later pass changes nothing, and its reason is not recorded.
+ """
+ # A module-level import of airflow.models.xcom is circular when this module is imported first.
+ from airflow.models.xcom import XComModel
from airflow.settings import task_instance_mutation_hook
if gate.dag_version_id is None:
raise InvalidLoopDecision("Loop gate requires a pinned DAG version")
- signal = XComModelV2.get_for_attempt(gate.id, LOOP_DECISION_KEY, session=session)
- later_gate = gate.has_later_loop_pass(group, session=session)
- if state != TaskInstanceState.SUCCESS or later_gate:
- if signal is not None:
- session.delete(signal)
+ if gate.has_later_loop_pass(group, session=session):
return
- decision = signal.value if signal is not None else None
- validate_loop_decision(decision, gate.region_index, group)
- session.delete(signal)
- if decision == "stop":
+ if state == TaskInstanceState.SUCCESS:
+ validate_loop_decision(stop, gate.region_index, group)
+ if reason:
+ XComModel.set_for_attempt(
+ task_instance_id=gate.id,
+ key=LOOP_DECISION_REASON_KEY,
+ value=reason,
+ serialize=False,
+ session=session,
+ )
+ if state != TaskInstanceState.SUCCESS or stop:
return
created_counts: dict[str, int] = defaultdict(int)
hook_is_noop: Literal[True, False] = getattr(task_instance_mutation_hook, "is_noop", False)
diff --git a/airflow-core/src/airflow/models/dynamic_region.py b/airflow-core/src/airflow/models/dynamic_region.py
index 1a722b775be..65470ea6bc5 100644
--- a/airflow-core/src/airflow/models/dynamic_region.py
+++ b/airflow-core/src/airflow/models/dynamic_region.py
@@ -52,7 +52,8 @@ from airflow.utils.sqlalchemy import CompactUUID, UtcDateTime
SENTINEL_REGION_ID = UUID(int=0)
LOOP_XCOM_PREFIX = "_airflow_loop_"
-LOOP_DECISION_KEY = f"{LOOP_XCOM_PREFIX}decision"
+LOOP_DECISION_REASON_KEY = f"{LOOP_XCOM_PREFIX}decision_reason"
+LOOP_DECISION_REASON_MAX_BYTES = 512
if TYPE_CHECKING:
from sqlalchemy import Select
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts
index b4aed6744dd..02ed3b0f2cd 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts
@@ -979,11 +979,12 @@ export const UseXcomServiceGetXcomEntryKeyFn = ({ dagId, dagRunId, deserialize,
export type XcomServiceGetXcomEntriesDefaultResponse = Awaited<ReturnType<typeof XcomService.getXcomEntries>>;
export type XcomServiceGetXcomEntriesQueryResult<TData = XcomServiceGetXcomEntriesDefaultResponse, TError = unknown> = UseQueryResult<TData, TError>;
export const useXcomServiceGetXcomEntriesKey = "XcomServiceGetXcomEntries";
-export const UseXcomServiceGetXcomEntriesKeyFn = ({ dagDisplayNamePattern, dagDisplayNamePrefixPattern, dagId, dagRunId, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, mapIndex, mapIndexFilter, offset, orderBy, regionId, regionIndex, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, taskId, taskIdPattern, taskIdPrefixPattern, teams, xcomKey, xcomKeyPattern, xcomKeyPrefixPattern }: {
+export const UseXcomServiceGetXcomEntriesKeyFn = ({ dagDisplayNamePattern, dagDisplayNamePrefixPattern, dagId, dagRunId, excludeInternalKeys, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, mapIndex, mapIndexFilter, offset, orderBy, regionId, regionIndex, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, taskId, taskIdPattern, taskIdPrefixPattern, teams, xcomKey, xcomKeyPattern, xcomKeyPrefixPattern }: {
dagDisplayNamePattern?: string;
dagDisplayNamePrefixPattern?: string;
dagId: string;
dagRunId: string;
+ excludeInternalKeys?: boolean;
limit?: number;
logicalDateGt?: string;
logicalDateGte?: string;
@@ -1008,7 +1009,7 @@ export const UseXcomServiceGetXcomEntriesKeyFn = ({ dagDisplayNamePattern, dagDi
xcomKey?: string;
xcomKeyPattern?: string;
xcomKeyPrefixPattern?: string;
-}, queryKey?: Array<unknown>) => [useXcomServiceGetXcomEntriesKey, ...(queryKey ?? [{ dagDisplayNamePattern, dagDisplayNamePrefixPattern, dagId, dagRunId, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, mapIndex, mapIndexFilter, offset, orderBy, regionId, regionIndex, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, taskId, taskIdPattern, taskIdPrefixPattern, teams, xcomKey, xcomKeyPattern, xcomKeyPrefixPattern }])];
+}, queryKey?: Array<unknown>) => [useXcomServiceGetXcomEntriesKey, ...(queryKey ?? [{ dagDisplayNamePattern, dagDisplayNamePrefixPattern, dagId, dagRunId, excludeInternalKeys, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, mapIndex, mapIndexFilter, offset, orderBy, regionId, regionIndex, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, taskId, taskIdPattern, taskIdPrefixPattern, teams, xcomKey, xcomKeyPattern, xcomKeyPrefixPattern }])];
export type TaskServiceGetTasksDefaultResponse = Awaited<ReturnType<typeof TaskService.getTasks>>;
export type TaskServiceGetTasksQueryResult<TData = TaskServiceGetTasksDefaultResponse, TError = unknown> = UseQueryResult<TData, TError>;
export const useTaskServiceGetTasksKey = "TaskServiceGetTasks";
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
index 2fb80f04de7..7624712dac2 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
@@ -1952,6 +1952,7 @@ export const ensureUseXcomServiceGetXcomEntryData = (queryClient: QueryClient, {
* @param data.mapIndex
* @param data.regionId
* @param data.regionIndex
+* @param data.excludeInternalKeys Leave out XComs whose keys Airflow reserves for itself, such as loop decision reasons.
* @param data.limit
* @param data.offset
* @param data.xcomKeyPattern Case-insensitive substring match (SQL `ILIKE`). Slower than `xcom_key_prefix_pattern` on large tables — see "Filtering with pattern parameters".
@@ -1976,11 +1977,12 @@ export const ensureUseXcomServiceGetXcomEntryData = (queryClient: QueryClient, {
* @returns XComCollectionResponse Successful Response
* @throws ApiError
*/
-export const ensureUseXcomServiceGetXcomEntriesData = (queryClient: QueryClient, { dagDisplayNamePattern, dagDisplayNamePrefixPattern, dagId, dagRunId, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, mapIndex, mapIndexFilter, offset, orderBy, regionId, regionIndex, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, taskId, taskIdPattern, taskIdPrefixPattern, teams, xcomKey, xcomKeyPattern, xcomKeyPrefixPattern }: {
+export const ensureUseXcomServiceGetXcomEntriesData = (queryClient: QueryClient, { dagDisplayNamePattern, dagDisplayNamePrefixPattern, dagId, dagRunId, excludeInternalKeys, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, mapIndex, mapIndexFilter, offset, orderBy, regionId, regionIndex, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, taskId, taskIdPattern, taskIdPrefixPattern, teams, xcomKey, xcomKeyPattern, xcomKeyPrefixPattern }: {
dagDisplayNamePattern?: string;
dagDisplayNamePrefixPattern?: string;
dagId: string;
dagRunId: string;
+ excludeInternalKeys?: boolean;
limit?: number;
logicalDateGt?: string;
logicalDateGte?: string;
@@ -2005,7 +2007,7 @@ export const ensureUseXcomServiceGetXcomEntriesData = (queryClient: QueryClient,
xcomKey?: string;
xcomKeyPattern?: string;
xcomKeyPrefixPattern?: string;
-}) => queryClient.ensureQueryData({ queryKey: Common.UseXcomServiceGetXcomEntriesKeyFn({ dagDisplayNamePattern, dagDisplayNamePrefixPattern, dagId, dagRunId, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, mapIndex, mapIndexFilter, offset, orderBy, regionId, regionIndex, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, taskId, taskIdPattern, taskIdPrefixPattern, teams, xcomKey, xcomKeyPattern, xcomKeyPrefixPattern }), queryFn: () => [...]
+}) => queryClient.ensureQueryData({ queryKey: Common.UseXcomServiceGetXcomEntriesKeyFn({ dagDisplayNamePattern, dagDisplayNamePrefixPattern, dagId, dagRunId, excludeInternalKeys, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, mapIndex, mapIndexFilter, offset, orderBy, regionId, regionIndex, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, taskId, taskIdPattern, taskIdPrefixPattern, teams, xcomKey, xcomKeyPattern, xcomKeyPrefixPatte [...]
/**
* Get Tasks
* Get tasks for Dag.
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
index b54bb5e4416..a4741bc8602 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
@@ -1952,6 +1952,7 @@ export const prefetchUseXcomServiceGetXcomEntry = (queryClient: QueryClient, { d
* @param data.mapIndex
* @param data.regionId
* @param data.regionIndex
+* @param data.excludeInternalKeys Leave out XComs whose keys Airflow reserves for itself, such as loop decision reasons.
* @param data.limit
* @param data.offset
* @param data.xcomKeyPattern Case-insensitive substring match (SQL `ILIKE`). Slower than `xcom_key_prefix_pattern` on large tables — see "Filtering with pattern parameters".
@@ -1976,11 +1977,12 @@ export const prefetchUseXcomServiceGetXcomEntry = (queryClient: QueryClient, { d
* @returns XComCollectionResponse Successful Response
* @throws ApiError
*/
-export const prefetchUseXcomServiceGetXcomEntries = (queryClient: QueryClient, { dagDisplayNamePattern, dagDisplayNamePrefixPattern, dagId, dagRunId, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, mapIndex, mapIndexFilter, offset, orderBy, regionId, regionIndex, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, taskId, taskIdPattern, taskIdPrefixPattern, teams, xcomKey, xcomKeyPattern, xcomKeyPrefixPattern }: {
+export const prefetchUseXcomServiceGetXcomEntries = (queryClient: QueryClient, { dagDisplayNamePattern, dagDisplayNamePrefixPattern, dagId, dagRunId, excludeInternalKeys, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, mapIndex, mapIndexFilter, offset, orderBy, regionId, regionIndex, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, taskId, taskIdPattern, taskIdPrefixPattern, teams, xcomKey, xcomKeyPattern, xcomKeyPrefixPattern }: {
dagDisplayNamePattern?: string;
dagDisplayNamePrefixPattern?: string;
dagId: string;
dagRunId: string;
+ excludeInternalKeys?: boolean;
limit?: number;
logicalDateGt?: string;
logicalDateGte?: string;
@@ -2005,7 +2007,7 @@ export const prefetchUseXcomServiceGetXcomEntries = (queryClient: QueryClient, {
xcomKey?: string;
xcomKeyPattern?: string;
xcomKeyPrefixPattern?: string;
-}) => queryClient.prefetchQuery({ queryKey: Common.UseXcomServiceGetXcomEntriesKeyFn({ dagDisplayNamePattern, dagDisplayNamePrefixPattern, dagId, dagRunId, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, mapIndex, mapIndexFilter, offset, orderBy, regionId, regionIndex, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, taskId, taskIdPattern, taskIdPrefixPattern, teams, xcomKey, xcomKeyPattern, xcomKeyPrefixPattern }), queryFn: () => X [...]
+}) => queryClient.prefetchQuery({ queryKey: Common.UseXcomServiceGetXcomEntriesKeyFn({ dagDisplayNamePattern, dagDisplayNamePrefixPattern, dagId, dagRunId, excludeInternalKeys, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, mapIndex, mapIndexFilter, offset, orderBy, regionId, regionIndex, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, taskId, taskIdPattern, taskIdPrefixPattern, teams, xcomKey, xcomKeyPattern, xcomKeyPrefixPattern [...]
/**
* Get Tasks
* Get tasks for Dag.
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
index fb85ecf5053..d3039078814 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
@@ -1952,6 +1952,7 @@ export const useXcomServiceGetXcomEntry = <TData = Common.XcomServiceGetXcomEntr
* @param data.mapIndex
* @param data.regionId
* @param data.regionIndex
+* @param data.excludeInternalKeys Leave out XComs whose keys Airflow reserves for itself, such as loop decision reasons.
* @param data.limit
* @param data.offset
* @param data.xcomKeyPattern Case-insensitive substring match (SQL `ILIKE`). Slower than `xcom_key_prefix_pattern` on large tables — see "Filtering with pattern parameters".
@@ -1976,11 +1977,12 @@ export const useXcomServiceGetXcomEntry = <TData = Common.XcomServiceGetXcomEntr
* @returns XComCollectionResponse Successful Response
* @throws ApiError
*/
-export const useXcomServiceGetXcomEntries = <TData = Common.XcomServiceGetXcomEntriesDefaultResponse, TError = unknown, TQueryKey extends Array<unknown> = unknown[]>({ dagDisplayNamePattern, dagDisplayNamePrefixPattern, dagId, dagRunId, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, mapIndex, mapIndexFilter, offset, orderBy, regionId, regionIndex, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, taskId, taskIdPattern, taskIdPrefixP [...]
+export const useXcomServiceGetXcomEntries = <TData = Common.XcomServiceGetXcomEntriesDefaultResponse, TError = unknown, TQueryKey extends Array<unknown> = unknown[]>({ dagDisplayNamePattern, dagDisplayNamePrefixPattern, dagId, dagRunId, excludeInternalKeys, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, mapIndex, mapIndexFilter, offset, orderBy, regionId, regionIndex, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, taskId, taskIdP [...]
dagDisplayNamePattern?: string;
dagDisplayNamePrefixPattern?: string;
dagId: string;
dagRunId: string;
+ excludeInternalKeys?: boolean;
limit?: number;
logicalDateGt?: string;
logicalDateGte?: string;
@@ -2005,7 +2007,7 @@ export const useXcomServiceGetXcomEntries = <TData = Common.XcomServiceGetXcomEn
xcomKey?: string;
xcomKeyPattern?: string;
xcomKeyPrefixPattern?: string;
-}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>, "queryKey" | "queryFn">) => useQuery<TData, TError>({ queryKey: Common.UseXcomServiceGetXcomEntriesKeyFn({ dagDisplayNamePattern, dagDisplayNamePrefixPattern, dagId, dagRunId, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, mapIndex, mapIndexFilter, offset, orderBy, regionId, regionIndex, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, taskId, taskIdPattern, ta [...]
+}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>, "queryKey" | "queryFn">) => useQuery<TData, TError>({ queryKey: Common.UseXcomServiceGetXcomEntriesKeyFn({ dagDisplayNamePattern, dagDisplayNamePrefixPattern, dagId, dagRunId, excludeInternalKeys, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, mapIndex, mapIndexFilter, offset, orderBy, regionId, regionIndex, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, task [...]
/**
* Get Tasks
* Get tasks for Dag.
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
index 4665742bd3b..d9e4d58598b 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
@@ -1952,6 +1952,7 @@ export const useXcomServiceGetXcomEntrySuspense = <TData = Common.XcomServiceGet
* @param data.mapIndex
* @param data.regionId
* @param data.regionIndex
+* @param data.excludeInternalKeys Leave out XComs whose keys Airflow reserves for itself, such as loop decision reasons.
* @param data.limit
* @param data.offset
* @param data.xcomKeyPattern Case-insensitive substring match (SQL `ILIKE`). Slower than `xcom_key_prefix_pattern` on large tables — see "Filtering with pattern parameters".
@@ -1976,11 +1977,12 @@ export const useXcomServiceGetXcomEntrySuspense = <TData = Common.XcomServiceGet
* @returns XComCollectionResponse Successful Response
* @throws ApiError
*/
-export const useXcomServiceGetXcomEntriesSuspense = <TData = Common.XcomServiceGetXcomEntriesDefaultResponse, TError = unknown, TQueryKey extends Array<unknown> = unknown[]>({ dagDisplayNamePattern, dagDisplayNamePrefixPattern, dagId, dagRunId, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, mapIndex, mapIndexFilter, offset, orderBy, regionId, regionIndex, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, taskId, taskIdPattern, taskI [...]
+export const useXcomServiceGetXcomEntriesSuspense = <TData = Common.XcomServiceGetXcomEntriesDefaultResponse, TError = unknown, TQueryKey extends Array<unknown> = unknown[]>({ dagDisplayNamePattern, dagDisplayNamePrefixPattern, dagId, dagRunId, excludeInternalKeys, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, mapIndex, mapIndexFilter, offset, orderBy, regionId, regionIndex, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, taskId, [...]
dagDisplayNamePattern?: string;
dagDisplayNamePrefixPattern?: string;
dagId: string;
dagRunId: string;
+ excludeInternalKeys?: boolean;
limit?: number;
logicalDateGt?: string;
logicalDateGte?: string;
@@ -2005,7 +2007,7 @@ export const useXcomServiceGetXcomEntriesSuspense = <TData = Common.XcomServiceG
xcomKey?: string;
xcomKeyPattern?: string;
xcomKeyPrefixPattern?: string;
-}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>, "queryKey" | "queryFn">) => useSuspenseQuery<TData, TError>({ queryKey: Common.UseXcomServiceGetXcomEntriesKeyFn({ dagDisplayNamePattern, dagDisplayNamePrefixPattern, dagId, dagRunId, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, mapIndex, mapIndexFilter, offset, orderBy, regionId, regionIndex, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, taskId, taskIdPat [...]
+}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>, "queryKey" | "queryFn">) => useSuspenseQuery<TData, TError>({ queryKey: Common.UseXcomServiceGetXcomEntriesKeyFn({ dagDisplayNamePattern, dagDisplayNamePrefixPattern, dagId, dagRunId, excludeInternalKeys, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, mapIndex, mapIndexFilter, offset, orderBy, regionId, regionIndex, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPatte [...]
/**
* Get Tasks
* Get tasks for Dag.
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts
index 74f4f6144e9..eab37cda3aa 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts
@@ -12211,6 +12211,11 @@ export const $LoopIterationSummary = {
}
],
title: 'End Date'
+ },
+ decision_reason: {
+ type: 'string',
+ title: 'Decision Reason',
+ description: "Why the loop's gate continued or stopped the loop at this iteration, if its until callable said."
}
},
type: 'object',
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
index 1ea77170ad0..b9d60978b19 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
@@ -4761,6 +4761,7 @@ export class XcomService {
* @param data.mapIndex
* @param data.regionId
* @param data.regionIndex
+ * @param data.excludeInternalKeys Leave out XComs whose keys Airflow reserves for itself, such as loop decision reasons.
* @param data.limit
* @param data.offset
* @param data.xcomKeyPattern Case-insensitive substring match (SQL `ILIKE`). Slower than `xcom_key_prefix_pattern` on large tables — see "Filtering with pattern parameters".
@@ -4799,6 +4800,7 @@ export class XcomService {
map_index: data.mapIndex,
region_id: data.regionId,
region_index: data.regionIndex,
+ exclude_internal_keys: data.excludeInternalKeys,
limit: data.limit,
offset: data.offset,
xcom_key_pattern: data.xcomKeyPattern,
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
index 1b33796bb11..956bc599eba 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
@@ -3156,6 +3156,10 @@ export type LoopIterationSummary = {
state?: TaskInstanceState | null;
start_date?: string | null;
end_date?: string | null;
+ /**
+ * Why the loop's gate continued or stopped the loop at this iteration, if its until callable said.
+ */
+ decision_reason?: string;
};
/**
@@ -5228,6 +5232,10 @@ export type GetXcomEntriesData = {
dagDisplayNamePrefixPattern?: string | null;
dagId: string;
dagRunId: string;
+ /**
+ * Leave out XComs whose keys Airflow reserves for itself, such as loop decision reasons.
+ */
+ excludeInternalKeys?: boolean;
limit?: number;
logicalDateGt?: string | null;
logicalDateGte?: string | null;
diff --git a/airflow-core/src/airflow/ui/public/i18n/locales/en/common.json b/airflow-core/src/airflow/ui/public/i18n/locales/en/common.json
index f237461162a..e1e03529b32 100644
--- a/airflow-core/src/airflow/ui/public/i18n/locales/en/common.json
+++ b/airflow-core/src/airflow/ui/public/i18n/locales/en/common.json
@@ -475,6 +475,7 @@
"hostname": "Hostname",
"id": "ID",
"iteration": "Iteration",
+ "loopDecision": "Decision",
"loopIterations": "Iterations run",
"maxTries": "Max Tries",
"pid": "PID",
diff --git a/airflow-core/src/airflow/ui/src/components/LoopDecisionReason.tsx b/airflow-core/src/airflow/ui/src/components/LoopDecisionReason.tsx
new file mode 100644
index 00000000000..b5fdea85bf7
--- /dev/null
+++ b/airflow-core/src/airflow/ui/src/components/LoopDecisionReason.tsx
@@ -0,0 +1,37 @@
+/*!
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+import { Text, type TextProps } from "@chakra-ui/react";
+
+import { Tooltip } from "src/system-components";
+
+export const LoopDecisionReason = ({ reason, ...rest }: { readonly reason: string } & TextProps) => (
+ // Portalled because the header card clips its overflow.
+ <Tooltip
+ content={
+ <Text overflowWrap="anywhere" whiteSpace="pre-line">
+ {reason}
+ </Text>
+ }
+ portalled
+ >
+ <Text maxWidth="md" truncate {...rest}>
+ {reason}
+ </Text>
+ </Tooltip>
+);
diff --git a/airflow-core/src/airflow/ui/src/pages/GroupTaskInstance/useLoopOutcomeStats.test.tsx b/airflow-core/src/airflow/ui/src/pages/GroupTaskInstance/useLoopOutcomeStats.test.tsx
new file mode 100644
index 00000000000..0642f0ede7a
--- /dev/null
+++ b/airflow-core/src/airflow/ui/src/pages/GroupTaskInstance/useLoopOutcomeStats.test.tsx
@@ -0,0 +1,64 @@
+/*!
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+import "@testing-library/jest-dom/vitest";
+import { render, screen } from "@testing-library/react";
+import { describe, expect, it, vi } from "vitest";
+
+import type { LoopSummaryResponse } from "openapi/requests/types.gen";
+
+import { Wrapper } from "src/utils/Wrapper";
+
+import { useLoopOutcomeStats } from "./useLoopOutcomeStats";
+
+const mockLoopSummary: { data?: Partial<LoopSummaryResponse> } = {};
+
+vi.mock("src/queries/useLoopSummary", () => ({
+ useLoopSummary: () => ({ data: mockLoopSummary.data }),
+}));
+
+const OutcomeStats = () => (
+ <>
+ {useLoopOutcomeStats({ dagId: "refine_estimate", groupId: "refine", runId: "run" }).map((stat) => (
+ <div key={stat.label}>{stat.value}</div>
+ ))}
+ </>
+);
+
+describe("useLoopOutcomeStats", () => {
+ it.each([
+ { shown: true, status: "stopped_early" },
+ { shown: true, status: "failed" },
+ { shown: false, status: "running" },
+ ] as const)("shows the final decision reason: $shown for a $status loop", ({ shown, status }) => {
+ mockLoopSummary.data = {
+ iterations: [
+ { decision_reason: "error 0.25, needs < 1e-06", index: 0 },
+ { decision_reason: "error 4.5e-12, needs < 1e-06", index: 1 },
+ ],
+ iterations_ran: 2,
+ max_iterations: 10,
+ status,
+ };
+
+ render(<OutcomeStats />, { wrapper: Wrapper });
+
+ expect(screen.queryByText("error 4.5e-12, needs < 1e-06") !== null).toBe(shown);
+ expect(screen.queryByText("error 0.25, needs < 1e-06")).not.toBeInTheDocument();
+ });
+});
diff --git a/airflow-core/src/airflow/ui/src/pages/GroupTaskInstance/useLoopOutcomeStats.tsx b/airflow-core/src/airflow/ui/src/pages/GroupTaskInstance/useLoopOutcomeStats.tsx
index 4983c94d744..ce1065593af 100644
--- a/airflow-core/src/airflow/ui/src/pages/GroupTaskInstance/useLoopOutcomeStats.tsx
+++ b/airflow-core/src/airflow/ui/src/pages/GroupTaskInstance/useLoopOutcomeStats.tsx
@@ -22,6 +22,8 @@ import { HStack, Icon, Text } from "@chakra-ui/react";
import { useTranslation } from "react-i18next";
import { MdLoop } from "react-icons/md";
+import { LoopDecisionReason } from "src/components/LoopDecisionReason";
+
import { useLoopSummary } from "src/queries/useLoopSummary";
import { reasonSentence } from "./LoopIterations/loopUtils";
@@ -42,13 +44,21 @@ export const useLoopOutcomeStats = ({
return [];
}
+ // A running loop's newest pass has not decided yet, so an older pass's reason would misdescribe it.
+ const finalReason = summary.status === "running" ? undefined : summary.iterations.at(-1)?.decision_reason;
+
return [
{
label: translate("loop.outcome"),
value: (
- <Text color={summary.status === "failed" ? "fg.error" : undefined}>
- {reasonSentence(translate, summary)}
- </Text>
+ <>
+ <Text color={summary.status === "failed" ? "fg.error" : undefined}>
+ {reasonSentence(translate, summary)}
+ </Text>
+ {finalReason === undefined ? undefined : (
+ <LoopDecisionReason color="fg.muted" reason={finalReason} />
+ )}
+ </>
),
},
{
diff --git a/airflow-core/src/airflow/ui/src/pages/TaskInstance/Header.test.tsx b/airflow-core/src/airflow/ui/src/pages/TaskInstance/Header.test.tsx
index 0558ad94261..27302696b11 100644
--- a/airflow-core/src/airflow/ui/src/pages/TaskInstance/Header.test.tsx
+++ b/airflow-core/src/airflow/ui/src/pages/TaskInstance/Header.test.tsx
@@ -20,7 +20,7 @@ import "@testing-library/jest-dom/vitest";
import { render, screen } from "@testing-library/react";
import { beforeEach, describe, expect, it, vi } from "vitest";
-import type { TaskInstanceResponse } from "openapi/requests/types.gen";
+import type { LoopSummaryResponse, TaskInstanceResponse } from "openapi/requests/types.gen";
import i18n from "src/i18n/config";
import { Wrapper } from "src/utils/Wrapper";
@@ -43,6 +43,11 @@ vi.mock("src/components/NotePreview", () => ({
}));
const mockConfig: Record<string, unknown> = { multi_team: false };
+const mockLoopSummary: { data?: Partial<LoopSummaryResponse> } = {};
+
+vi.mock("src/queries/useLoopSummary", () => ({
+ useLoopSummary: () => ({ data: mockLoopSummary.data }),
+}));
vi.mock("src/queries/useConfig", () => ({
useConfig: (key: string) => mockConfig[key],
@@ -106,6 +111,33 @@ describe("Header", () => {
const renderHeader = (overrides: Partial<TaskInstanceResponse>) =>
render(<Header taskInstance={{ ...baseTaskInstance, ...overrides }} />, { wrapper: Wrapper });
+describe("Header loop decision", () => {
+ beforeEach(() => {
+ mockLoopSummary.data = {
+ iterations: [
+ { decision_reason: "error 0.25, needs < 1e-06", index: 0 },
+ { index: 1 },
+ { decision_reason: "error 6e-06, needs < 1e-06", index: 2 },
+ ],
+ iterations_ran: 3,
+ };
+ });
+
+ it("shows the reason its own iteration's gate gave", () => {
+ renderHeader({ loop_iteration: { iteration: 2, loop_id: "refine" } });
+
+ expect(screen.getByText(i18n.t("common:taskInstance.loopDecision"))).toBeInTheDocument();
+ expect(screen.getByText("error 6e-06, needs < 1e-06")).toBeInTheDocument();
+ expect(screen.queryByText("error 0.25, needs < 1e-06")).not.toBeInTheDocument();
+ });
+
+ it("omits the stat for an iteration without a reason", () => {
+ renderHeader({ loop_iteration: { iteration: 1, loop_id: "refine" } });
+
+ expect(screen.queryByText(i18n.t("common:taskInstance.loopDecision"))).not.toBeInTheDocument();
+ });
+});
+
describe("Header state reason banner", () => {
// Without the bundle i18n.t() echoes the key, so the titles below would assert nothing.
beforeEach(() => {
diff --git a/airflow-core/src/airflow/ui/src/pages/TaskInstance/Header.tsx b/airflow-core/src/airflow/ui/src/pages/TaskInstance/Header.tsx
index 436f6f7e99b..dbf7381a925 100644
--- a/airflow-core/src/airflow/ui/src/pages/TaskInstance/Header.tsx
+++ b/airflow-core/src/airflow/ui/src/pages/TaskInstance/Header.tsx
@@ -30,6 +30,7 @@ import { ClearTaskInstanceButton } from "src/components/Clear";
import ClearTaskInstanceDialog from "src/components/Clear/TaskInstance/ClearTaskInstanceDialog";
import { DagVersion } from "src/components/DagVersion";
import { HeaderCard } from "src/components/HeaderCard";
+import { LoopDecisionReason } from "src/components/LoopDecisionReason";
import { MarkTaskInstanceAsButton } from "src/components/MarkAs";
import { NotePreview } from "src/components/NotePreview";
import ReactMarkdown from "src/components/ReactMarkdown";
@@ -75,6 +76,9 @@ export const Header = ({ taskInstance }: { readonly taskInstance: TaskInstanceRe
groupId: loopIteration?.loop_id ?? "",
runId: taskInstance.dag_run_id,
});
+ const decisionReason = loopSummary?.iterations.find(
+ (iteration) => iteration.index === loopIteration?.iteration,
+ )?.decision_reason;
const stats = [
{ label: translate("task.operator"), value: taskInstance.operator_name },
@@ -92,6 +96,14 @@ export const Header = ({ taskInstance }: { readonly taskInstance: TaskInstanceRe
value: `${loopIteration.iteration}/${loopSummary?.iterations_ran ?? loopIteration.iteration}`,
},
]),
+ ...(decisionReason === undefined
+ ? []
+ : [
+ {
+ label: translate("taskInstance.loopDecision"),
+ value: <LoopDecisionReason reason={decisionReason} />,
+ },
+ ]),
{ label: translate("startDate"), value: <Time datetime={taskInstance.start_date} /> },
{ label: translate("endDate"), value: <Time datetime={taskInstance.end_date} /> },
...(Boolean(taskInstance.start_date)
diff --git a/airflow-core/src/airflow/ui/src/pages/XCom/XCom.test.tsx b/airflow-core/src/airflow/ui/src/pages/XCom/XCom.test.tsx
index 18aab48cd2a..82532ba9c46 100644
--- a/airflow-core/src/airflow/ui/src/pages/XCom/XCom.test.tsx
+++ b/airflow-core/src/airflow/ui/src/pages/XCom/XCom.test.tsx
@@ -78,3 +78,21 @@ describe("XCom list expand/collapse", () => {
await waitFor(() => expect(screen.queryByTestId("monaco-editor")).toBeNull());
});
});
+
+describe("XCom list", () => {
+ it("asks the API to leave out the XComs Airflow reserves for itself", async () => {
+ const requested: Array<string | null> = [];
+
+ server.use(
+ http.get("/api/v2/dags/~/dagRuns/~/taskInstances/~/xcomEntries", ({ request }) => {
+ requested.push(new URL(request.url).searchParams.get("exclude_internal_keys"));
+
+ return HttpResponse.json({ total_entries: 1, xcom_entries: [xcomEntry] });
+ }),
+ );
+ render(<AppWrapper initialEntries={["/xcoms"]} />);
+
+ expect(await screen.findByText("return_value")).toBeInTheDocument();
+ expect(requested).toContain("true");
+ });
+});
diff --git a/airflow-core/src/airflow/ui/src/pages/XCom/XCom.tsx b/airflow-core/src/airflow/ui/src/pages/XCom/XCom.tsx
index ecf0710cbb1..4a07ca10206 100644
--- a/airflow-core/src/airflow/ui/src/pages/XCom/XCom.tsx
+++ b/airflow-core/src/airflow/ui/src/pages/XCom/XCom.tsx
@@ -239,6 +239,7 @@ export const XCom = () => {
...dagDisplayNameArg,
dagId,
dagRunId: runId,
+ excludeInternalKeys: true,
limit: pagination.pageSize,
logicalDateGte: logicalDateGte ?? undefined,
logicalDateLte: logicalDateLte ?? undefined,
diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
index 616126f411f..9ea64b24da8 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
@@ -43,7 +43,7 @@ from airflow.jobs.triggerer_job_runner import TriggererJobRunner
from airflow.models import DagModel, DagRun, Log, TaskInstance
from airflow.models.dag_version import DagVersion
from airflow.models.dagbundle import DagBundleModel
-from airflow.models.dynamic_region import LOOP_DECISION_KEY, DynamicRegion
+from airflow.models.dynamic_region import LOOP_DECISION_REASON_KEY, DynamicRegion
from airflow.models.renderedtifields import RenderedTaskInstanceFields as RTIF
from airflow.models.task_state_store import TaskStateStoreModel
from airflow.models.taskinstance import clear_loop_task_instances, clear_task_instances, uuid7
@@ -5983,7 +5983,7 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
assert tis["body.first", 0].note is None
assert tis["body.first", 2].note is None
- def test_manual_gate_success_after_rewind_stops_without_consuming_continue(
+ def test_manual_gate_success_after_rewind_stops_and_keeps_the_reason(
self, test_client, session, loop_instances
):
dr, loop, root, tis = loop_instances
@@ -5997,7 +5997,7 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
)
).one()
gate_id = gate.id
- session.add(XComModelV2(task_instance_id=gate_id, key=LOOP_DECISION_KEY, value="continue"))
+ session.add(XComModelV2(task_instance_id=gate_id, key=LOOP_DECISION_REASON_KEY, value="stale"))
session.commit()
response = test_client.patch(
@@ -6009,7 +6009,9 @@ class TestRegionalTaskStateControls(TestTaskInstanceEndpoint):
session.expire_all()
assert session.get(TaskInstance, gate_id).state == State.SUCCESS
assert all(ti.region_index == 0 for ti in dr.get_task_instances(session=session))
- assert not session.scalar(select(XComModelV2).where(XComModelV2.task_instance_id == gate_id))
+ assert session.scalars(
+ select(XComModelV2.value).where(XComModelV2.task_instance_id == gate_id)
+ ).all() == ["stale"]
@pytest.mark.parametrize("outcome", ["success", "failed", "skipped"])
def test_manual_terminal_settles_superseded_running_execution(
diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_xcom.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_xcom.py
index bee30deaea1..3444ca0f015 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_xcom.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_xcom.py
@@ -538,6 +538,38 @@ class TestGetXComEntries(TestXComEndpoint):
for e in entries
] == [(dag_id, dag_id, task_id, task_id, run_id, key) for dag_id, task_id, key in expected]
+ @pytest.mark.parametrize(
+ ("exclude_internal_keys", "expected_keys"),
+ [
+ (
+ False,
+ [
+ f"{TEST_XCOM_KEY}-0",
+ f"{TEST_XCOM_KEY}-1",
+ "_airflow_loop_decision_reason",
+ "xairflowxloopx",
+ ],
+ ),
+ (True, [f"{TEST_XCOM_KEY}-0", f"{TEST_XCOM_KEY}-1", "xairflowxloopx"]),
+ ],
+ ids=["included-by-default", "excluded"],
+ )
+ def test_exclude_internal_keys(self, test_client, exclude_internal_keys, expected_keys):
+ self._create_xcom_entries(TEST_DAG_ID, run_id, logical_date_parsed, TEST_TASK_ID)
+ for key in ("_airflow_loop_decision_reason", "xairflowxloopx"):
+ XComModel.set(
+ key=key, value="v", run_id=run_id, task_id=TEST_TASK_ID, dag_id=TEST_DAG_ID, map_index=-1
+ )
+
+ response = test_client.get(
+ f"/dags/{TEST_DAG_ID}/dagRuns/{run_id}/taskInstances/{TEST_TASK_ID}/xcomEntries",
+ params={"exclude_internal_keys": exclude_internal_keys, "order_by": "key"},
+ )
+
+ assert response.status_code == 200
+ assert response.json()["total_entries"] == len(expected_keys)
+ assert sorted(entry["key"] for entry in response.json()["xcom_entries"]) == sorted(expected_keys)
+
def test_should_respond_200_with_tilde(self, test_client):
self._create_xcom_entries(TEST_DAG_ID, run_id, logical_date_parsed, TEST_TASK_ID)
self._create_xcom_entries(TEST_DAG_ID_2, run_id, logical_date_parsed, TEST_TASK_ID_2)
@@ -961,6 +993,15 @@ class TestCreateXComEntry(TestXComEndpoint):
f"The XCom with key: `{TEST_XCOM_KEY}` with mentioned task instance already exists.",
id="xcom-already-exists",
),
+ pytest.param(
+ TEST_DAG_ID,
+ TEST_TASK_ID,
+ run_id,
+ XComCreateBody(key="_airflow_loop_decision_reason", value="converged"),
+ 400,
+ "Reserved loop XCom key",
+ id="reserved-loop-key",
+ ),
],
)
def test_create_xcom_entry(
@@ -1199,6 +1240,13 @@ class TestPatchXComEntry(TestXComEndpoint):
f"The XCom with key: `{TEST_XCOM_KEY}` with mentioned task instance doesn't exist.",
id="xcom-not-found",
),
+ pytest.param(
+ "_airflow_loop_decision_reason",
+ {"value": "converged"},
+ 400,
+ "Reserved loop XCom key",
+ id="reserved-loop-key",
+ ),
],
)
def test_patch_xcom_entry(self, key, patch_body, expected_status, expected_detail, test_client, session):
diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py
index 80158ac5839..87be7d87815 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_grid.py
@@ -45,7 +45,7 @@ from airflow.models.taskinstance import (
)
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.providers.standard.operators.python import PythonOperator
-from airflow.sdk import task, task_group
+from airflow.sdk import LoopDecision, task, task_group
from airflow.sdk.definitions.taskgroup import TaskGroup
from airflow.utils.session import NEW_SESSION, provide_session
from airflow.utils.state import DagRunState, TaskInstanceState
@@ -1329,6 +1329,42 @@ class TestGetGridDataEndpoint:
class TestGetLoopSummaryEndpoint:
gate_task_id = "body.__loop_gate"
+ @pytest.mark.parametrize(
+ ("until", "expected"),
+ [
+ (
+ lambda loop: LoopDecision(stop=loop.index >= 1, reason=f"iteration {loop.index} decided"),
+ ["iteration 0 decided", "iteration 1 decided"],
+ ),
+ (lambda loop: loop.index >= 1, [None, None]),
+ (lambda loop: LoopDecision(stop=loop.index >= 1), [None, None]),
+ ],
+ ids=["with_reason", "plain_bool", "decision_without_reason"],
+ )
+ def test_loop_summary_reports_the_reason_of_each_iteration(self, test_client, dag_maker, until, expected):
+ @task
+ def work():
+ return 1
+
+ @task_group
+ def body():
+ work()
+
+ with dag_maker(serialized=False) as dag:
+ body.loop(max_iterations=3, until=until)
+ run = dag.test()
+
+ response = test_client.get(f"/grid/loop/{run.dag_id}/{run.run_id}/body")
+
+ assert response.status_code == 200, response.text
+ iterations = response.json()["iterations"]
+ assert [iteration.get("decision_reason") for iteration in iterations] == expected
+ assert all(
+ "decision_reason" not in iteration
+ for iteration in iterations
+ if iteration.get("decision_reason") is None
+ )
+
@pytest.fixture
def executed_fixed_loop(self, dag_maker):
@task
@@ -1605,11 +1641,13 @@ class TestGetLoopSummaryEndpoint:
with (
capture_orm_selects("task_instance") as all_runs_tis,
capture_orm_selects("dynamic_region") as all_runs_regions,
+ capture_orm_selects("xcom_v2") as all_runs_xcoms,
):
response = test_client.get(path, params={"limit": 3})
assert len(all_runs_tis) == len(one_run_tis)
assert len(all_runs_regions) == len(one_run_regions)
+ assert all_runs_xcoms == []
assert [(run["run_id"], run["status"]) for run in response.json()["runs"]] == [
("history_0", "ran_to_cap"),
("history_1", "failed"),
diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
index 40b7f029a50..a2bba34a910 100644
--- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
+++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
@@ -18,11 +18,10 @@
from __future__ import annotations
import re
-import time
from concurrent.futures import ThreadPoolExecutor
from contextlib import nullcontext
from datetime import datetime, timedelta
-from threading import Barrier, Event, local
+from threading import Barrier, Event
from types import SimpleNamespace
from typing import TYPE_CHECKING
from unittest import mock
@@ -38,7 +37,7 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanE
from opentelemetry.trace import StatusCode
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
from pydantic import ValidationError
-from sqlalchemy import delete, event, select, update
+from sqlalchemy import delete, select, update
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import Session
@@ -51,13 +50,13 @@ from airflow._shared.timezones import timezone
from airflow.api_fastapi.auth.tokens import JWTGenerator, JWTValidator
from airflow.api_fastapi.execution_api.app import lifespan
from airflow.api_fastapi.execution_api.datamodels.taskinstance import (
+ LoopDecisionPayload,
TISuccessStatePayload,
TITerminalStatePayload,
)
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
from airflow.exceptions import AirflowSkipException, TaskNotFound
from airflow.jobs.scheduler_job_runner import SchedulerJobRunner
@@ -66,7 +65,7 @@ from airflow.models.asset import AssetActive, AssetAliasModel, AssetEvent, Asset
from airflow.models.dag import DagModel
from airflow.models.dag_version import DagVersion
from airflow.models.dagbag import DBDagBag
-from airflow.models.dynamic_region import DynamicRegion
+from airflow.models.dynamic_region import LOOP_DECISION_REASON_KEY, DynamicRegion
from airflow.models.log import Log
from airflow.models.task_coordinates import LOOP_GATE_OPERATOR, TaskCoordinateResolver
from airflow.models.task_state_store import TaskStateStoreModel
@@ -104,6 +103,16 @@ DEFAULT_START_DATE = timezone.parse("2024-10-31T11:00:00Z")
DEFAULT_END_DATE = timezone.parse("2024-10-31T12:00:00Z")
DEFAULT_RENDERED_MAP_INDEX = "test rendered map index"
+CONTINUE = {"stop": False}
+
+
+def _stored_loop_decision_reasons(session, ti_id):
+ return session.scalars(
+ select(XComModelV2.value).where(
+ XComModelV2.task_instance_id == ti_id, XComModelV2.key == LOOP_DECISION_REASON_KEY
+ )
+ ).all()
+
def _where_column_keys(statement) -> set[str]:
whereclause = getattr(statement, "whereclause", None)
@@ -1924,19 +1933,10 @@ class TestTIUpdateState:
assert gate in dr.task_instance_scheduling_decisions(session=session).schedulable_tis
gate.state, gate.start_date = State.RUNNING, DEFAULT_START_DATE
session.commit()
- exec_app = client.app.routes[-1].app
- exec_app.dependency_overrides[require_auth] = lambda: TIToken(id=gate.id, claims=TIClaims())
- assert (
- client.post(
- f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/{gate.task_id}/_airflow_loop_decision",
- json="continue",
- ).status_code
- == 201
- )
response = client.patch(
f"/execution/task-instances/{gate.id}/state",
- json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()},
+ json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat(), "loop_decision": CONTINUE},
)
assert response.status_code == 204, response.text
@@ -1951,7 +1951,7 @@ class TestTIUpdateState:
dr = gate.dag_run
response = client.patch(
f"/execution/task-instances/{gate.id}/state",
- json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()},
+ json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat(), "loop_decision": CONTINUE},
)
assert response.status_code == 204
session.expire_all()
@@ -2163,28 +2163,11 @@ class TestTIUpdateState:
gate.state = State.RUNNING
gate.start_date = DEFAULT_START_DATE
session.commit()
- if decision and rejected:
- XComModel.set_for_attempt(
- task_instance_id=gate.id,
- key="_airflow_loop_decision",
- value=decision,
- serialize=False,
- session=session,
- )
- session.commit()
- elif decision:
- exec_app = client.app.routes[-1].app
- exec_app.dependency_overrides[require_auth] = lambda: TIToken(id=gate.id, claims=TIClaims())
- response = client.post(
- f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/{gate.task_id}/_airflow_loop_decision",
- json=decision,
- )
- assert response.status_code == 201
+ payload = {"state": state, "end_date": DEFAULT_END_DATE.isoformat()}
+ if decision:
+ payload["loop_decision"] = {"stop": decision == "stop"}
- response = client.patch(
- f"/execution/task-instances/{gate.id}/state",
- json={"state": state, "end_date": DEFAULT_END_DATE.isoformat()},
- )
+ response = client.patch(f"/execution/task-instances/{gate.id}/state", json=payload)
assert response.status_code == 204, response.text
assert backend.clear.call_count == (state == State.SUCCESS and not rejected)
@@ -2201,13 +2184,8 @@ class TestTIUpdateState:
assert all(ti.region_id in region_ids for ti in current)
else:
assert all(ti.region_id == gate.region_id for ti in current)
- signal = session.scalar(select(XComModel).where(XComModel.task_id == gate.task_id))
- assert signal is None
if not rejected:
- duplicate = client.patch(
- f"/execution/task-instances/{gate.id}/state",
- json={"state": state, "end_date": DEFAULT_END_DATE.isoformat()},
- )
+ duplicate = client.patch(f"/execution/task-instances/{gate.id}/state", json=payload)
assert duplicate.status_code == 200
assert (
sorted(
@@ -2241,17 +2219,10 @@ class TestTIUpdateState:
dr = dag_maker.create_dagrun()
gate = next(ti for ti in dr.task_instances if ti.task_id == loop.gate_task_id)
gate.state, gate.start_date = State.RUNNING, DEFAULT_START_DATE
- XComModel.set_for_attempt(
- task_instance_id=gate.id,
- key="_airflow_loop_decision",
- value="continue",
- serialize=False,
- session=session,
- )
session.commit()
return gate
- def test_loop_gate_materialization_error_rolls_back_state_signal_and_successor(
+ def test_loop_gate_materialization_error_rolls_back_state_and_successor(
self, client, session, running_loop_gate, mocker
):
gate = running_loop_gate
@@ -2275,7 +2246,7 @@ class TestTIUpdateState:
mocker.patch.object(DagRun, "_create_task_instances", autospec=True, side_effect=fail_after_insert)
response = client.patch(
f"/execution/task-instances/{gate.id}/state",
- json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()},
+ json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat(), "loop_decision": CONTINUE},
)
assert response.status_code == 500
@@ -2283,7 +2254,6 @@ class TestTIUpdateState:
assert observed == [State.RUNNING, [0]]
assert gate.state == State.RUNNING
assert len(dr.get_task_instances(session=session)) == 2
- assert session.scalar(select(XComModel.value).where(XComModel.task_id == gate.task_id)) == "continue"
@pytest.mark.parametrize(
("max_tries", "expected_state"),
@@ -2296,7 +2266,6 @@ class TestTIUpdateState:
dr = gate.dag_run
gate_id = gate.id
gate.max_tries = max_tries
- session.execute(delete(XComModelV2).where(XComModelV2.task_instance_id == gate.id))
session.commit()
response = client.patch(
@@ -2324,7 +2293,11 @@ class TestTIUpdateState:
barrier.wait(timeout=10)
result = ti_update_state(
task_instance_id=gate_id,
- ti_patch_payload=TISuccessStatePayload(state="success", end_date=DEFAULT_END_DATE),
+ ti_patch_payload=TISuccessStatePayload(
+ state="success",
+ end_date=DEFAULT_END_DATE,
+ loop_decision=LoopDecisionPayload(stop=False),
+ ),
session=request_session,
dag_bag=DBDagBag(),
)
@@ -2338,88 +2311,6 @@ class TestTIUpdateState:
assert sorted(
ti.region_index for ti in dr.get_task_instances(session=session) if ti.task_id == gate.task_id
) == [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, request
- ):
- gate = running_loop_gate
- gate_id, dag_id, run_id, task_id = gate.id, gate.dag_id, gate.run_id, gate.task_id
- XComModel.set_for_attempt(
- task_instance_id=gate_id,
- key="_airflow_loop_decision",
- value="stop",
- serialize=False,
- session=session,
- )
- session.commit()
- bind = session.get_bind()
- role = local()
- writer_locked = Event()
- completion_at_lock = Event()
-
- def note_completion_lock(conn, cursor, statement, parameters, context, executemany):
- if getattr(role, "name", None) == "completion" and "FOR UPDATE" in statement:
- completion_at_lock.set()
-
- def pause_writer_after_first_lock(conn, cursor, statement, parameters, context, executemany):
- if (
- getattr(role, "name", None) == "writer"
- and "FOR UPDATE" in statement
- and not writer_locked.is_set()
- ):
- writer_locked.set()
- assert completion_at_lock.wait(timeout=10)
- # Give completion time to take its first lock, so a writer that locks in a different
- # order deadlocks against it.
- time.sleep(0.5)
-
- event.listen(bind, "before_cursor_execute", note_completion_lock)
- event.listen(bind, "after_cursor_execute", pause_writer_after_first_lock)
- request.addfinalizer(lambda: event.remove(bind, "before_cursor_execute", note_completion_lock))
- request.addfinalizer(
- lambda: event.remove(bind, "after_cursor_execute", pause_writer_after_first_lock)
- )
-
- def rewrite():
- role.name = "writer"
- with Session(bind=bind) as request_session:
- set_xcom(
- dag_id=dag_id,
- run_id=run_id,
- task_id=task_id,
- key="_airflow_loop_decision",
- session=request_session,
- dag_bag=DBDagBag(),
- value="continue",
- token=TIToken(id=gate_id, claims=TIClaims()),
- )
- request_session.commit()
-
- def complete():
- role.name = "completion"
- assert writer_locked.wait(timeout=10)
- with Session(bind=bind) as request_session:
- ti_update_state(
- task_instance_id=gate_id,
- ti_patch_payload=TISuccessStatePayload(state="success", end_date=DEFAULT_END_DATE),
- session=request_session,
- dag_bag=DBDagBag(),
- )
-
- with ThreadPoolExecutor(max_workers=2) as pool:
- futures = [pool.submit(rewrite), pool.submit(complete)]
- for future in futures:
- future.result(timeout=20)
- session.expire_all()
- assert gate.state == State.SUCCESS
- assert session.scalar(select(XComModel).where(XComModel.task_id == gate.task_id)) is None
- assert sorted(
- ti.region_index
- for ti in gate.dag_run.get_task_instances(session=session)
- if ti.task_id == gate.task_id
- ) == [0, 1]
def test_loop_gate_completion_uses_its_pinned_limit(self, client, session, running_loop_gate, dag_maker):
gate = running_loop_gate
@@ -2434,7 +2325,7 @@ class TestTIUpdateState:
response = client.patch(
f"/execution/task-instances/{gate.id}/state",
- json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()},
+ json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat(), "loop_decision": CONTINUE},
)
assert response.status_code == 204, response.text
@@ -2471,25 +2362,7 @@ class TestTIUpdateState:
assert gate.state == State.FAILED
assert [ti.region_index for ti in gate.dag_run.get_task_instances(session=session)] == [0, 0]
- def test_late_loop_decision_cannot_recreate_consumed_signal(self, client, session, running_loop_gate):
- gate = running_loop_gate
- response = client.patch(
- f"/execution/task-instances/{gate.id}/state",
- json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()},
- )
- assert response.status_code == 204
- exec_app = client.app.routes[-1].app
- exec_app.dependency_overrides[require_auth] = lambda: TIToken(id=gate.id, claims=TIClaims())
-
- response = client.post(
- f"/execution/xcoms/{gate.dag_id}/{gate.run_id}/{gate.task_id}/_airflow_loop_decision",
- json="continue",
- )
-
- 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(
+ def test_loop_gate_state_update_error_rolls_back_next_pass(
self, client, session, running_loop_gate, mocker
):
gate = running_loop_gate
@@ -2502,14 +2375,13 @@ class TestTIUpdateState:
response = client.patch(
f"/execution/task-instances/{gate.id}/state",
- json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()},
+ json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat(), "loop_decision": CONTINUE},
)
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(
@@ -2539,7 +2411,7 @@ class TestTIUpdateState:
response = client.patch(
f"/execution/task-instances/{gate.id}/state",
- json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()},
+ json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat(), "loop_decision": CONTINUE},
)
assert response.status_code == 500
@@ -2547,14 +2419,19 @@ class TestTIUpdateState:
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", [State.SUCCESS, State.FAILED, State.SKIPPED])
- def test_rerun_gate_preserves_existing_later_pass(self, client, session, running_loop_gate, state):
+ def test_rerun_gate_preserves_existing_later_pass_and_reason(
+ self, client, session, running_loop_gate, state
+ ):
gate = running_loop_gate
response = client.patch(
f"/execution/task-instances/{gate.id}/state",
- json={"state": "success", "end_date": DEFAULT_END_DATE.isoformat()},
+ json={
+ "state": "success",
+ "end_date": DEFAULT_END_DATE.isoformat(),
+ "loop_decision": {"stop": False, "reason": "first"},
+ },
)
assert response.status_code == 204
session.refresh(gate)
@@ -2580,7 +2457,11 @@ class TestTIUpdateState:
response = client.patch(
f"/execution/task-instances/{gate.id}/state",
- json={"state": state, "end_date": DEFAULT_END_DATE.isoformat()},
+ json={
+ "state": state,
+ "end_date": DEFAULT_END_DATE.isoformat(),
+ "loop_decision": {"stop": True, "reason": "second"},
+ },
)
assert response.status_code == 204, response.text
@@ -2589,6 +2470,48 @@ class TestTIUpdateState:
assert {
ti.id for ti in gate.dag_run.get_task_instances(session=session) if ti.region_index == 1
} == successor_ids
+ assert _stored_loop_decision_reasons(session, archived_id) == ["first"]
+ assert _stored_loop_decision_reasons(session, gate.id) == []
+
+ @pytest.mark.parametrize(
+ ("payload", "expected"),
+ [
+ ({"state": "success", "loop_decision": {"stop": False, "reason": "score 0.7"}}, ["score 0.7"]),
+ ({"state": "failed", "loop_decision": {"stop": False, "reason": "score 0.8"}}, ["score 0.8"]),
+ ({"state": "success", "loop_decision": {"stop": False}}, []),
+ ({"state": "success", "loop_decision": {"stop": False, "reason": ""}}, []),
+ ],
+ ids=["success", "failed-with-decision", "no-reason", "empty-reason"],
+ )
+ def test_loop_gate_records_the_reason_of_its_decision(
+ self, client, session, running_loop_gate, payload, expected
+ ):
+ gate = running_loop_gate
+
+ response = client.patch(
+ f"/execution/task-instances/{gate.id}/state",
+ json={**payload, "end_date": DEFAULT_END_DATE.isoformat()},
+ )
+
+ assert response.status_code == 204, response.text
+ session.expire_all()
+ assert _stored_loop_decision_reasons(session, gate.id) == expected
+
+ def test_loop_gate_rejects_a_reason_over_the_byte_limit(self, client, session, running_loop_gate):
+ gate = running_loop_gate
+
+ response = client.patch(
+ f"/execution/task-instances/{gate.id}/state",
+ json={
+ "state": "success",
+ "end_date": DEFAULT_END_DATE.isoformat(),
+ "loop_decision": {"stop": False, "reason": "é" * 257},
+ },
+ )
+
+ assert response.status_code == 422
+ session.expire_all()
+ assert gate.state == State.RUNNING
@pytest.mark.parametrize(
("interruption", "expected_state", "expected_status"),
diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py
index 80f844eb3cf..c694aa1270a 100644
--- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py
+++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py
@@ -642,57 +642,12 @@ def test_trigger_ti_id_claim_check(
)
-@pytest.mark.parametrize("key", ["_airflow_loop_decision", "_airflow_loop_other"])
-def test_normal_xcom_push_rejects_loop_reserved_prefix(client, create_task_instance, key):
+def test_normal_xcom_push_rejects_loop_reserved_prefix(client, create_task_instance):
ti = create_task_instance()
- response = client.post(f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/{key}", json="continue")
- assert response.status_code == 400
-
-
-@pytest.mark.parametrize(
- ("value", "region_index", "expected_status"),
- [
- ("continue", 0, 201),
- ("stop", 2, 201),
- ("continue", 2, 400),
- ("stop", 0, 400),
- (True, 2, 400),
- ("invalid", 2, 400),
- ({"decision": "stop"}, 2, 400),
- ],
-)
-def test_internal_loop_decision_uses_token_gate_coordinates(
- client, loop_xcoms, session, authenticate_as, value, region_index, expected_status
-):
- dr, _, consumer, _ = loop_xcoms
- gate = session.scalar(
- select(TaskInstance).where(
- TaskInstance.dag_id == dr.dag_id, TaskInstance.task_id == "body.__loop_gate"
- )
- )
- gate.region_id, gate.region_index = consumer.region_id, region_index
- gate.state = "running"
- session.commit()
- authenticate_as(gate)
- url = f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/{gate.task_id}/_airflow_loop_decision"
-
- response = client.post(url, json=value)
-
- assert response.status_code == expected_status
- stored = session.scalar(select(XComModel).where(XComModel.task_id == gate.task_id))
- if expected_status == 400:
- assert stored is None
- return
- assert (stored.region_id, stored.region_index, stored.value) == (gate.region_id, region_index, value)
- authenticate_as(consumer)
- assert client.post(url, json="stop").status_code == 400
- assert (
- client.post(
- f"/execution/xcoms/{dr.dag_id}/{dr.run_id}/{consumer.task_id}/_airflow_loop_decision",
- json="stop",
- ).status_code
- == 400
+ response = client.post(
+ f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/_airflow_loop_decision_reason", json="forged"
)
+ assert response.status_code == 400
def test_loop_result_item_read_loads_no_producer_rows(dag_maker, session, client, authenticate_as):
diff --git a/airflow-core/tests/unit/models/conftest.py b/airflow-core/tests/unit/models/conftest.py
index fb2e5fb5571..bc18469dcac 100644
--- a/airflow-core/tests/unit/models/conftest.py
+++ b/airflow-core/tests/unit/models/conftest.py
@@ -22,10 +22,9 @@ from sqlalchemy import select
from airflow._shared.timezones import timezone
from airflow.models.dagbag import DBDagBag
from airflow.models.dagrun import DagRun
-from airflow.models.dynamic_region import LOOP_DECISION_KEY, DynamicRegion
+from airflow.models.dynamic_region import DynamicRegion
from airflow.models.task_coordinates import TaskCoordinateResolver
from airflow.models.taskinstance import TaskInstance
-from airflow.models.xcom import XComModelV2
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.providers.standard.operators.python import PythonOperator
from airflow.sdk import task_group
@@ -102,10 +101,11 @@ def completed_loop(dag_maker, session):
ti.state = State.SUCCESS
ti.start_date = ti.end_date = timezone.utcnow()
gate = next(ti for ti in members if ti.task_id == loop.gate_task_id)
- session.add(XComModelV2(task_instance_id=gate.id, key=LOOP_DECISION_KEY, value=decision))
session.flush()
group, _ = resolver.loop_context(gate)
- dr.complete_loop_gate(gate, group, State.SUCCESS, session=session)
+ dr.complete_loop_gate(
+ gate, group, State.SUCCESS, stop=decision == "stop", reason=None, session=session
+ )
session.flush()
for index in range(5):
diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py
index 25795e40d8e..4feeec84245 100644
--- a/airflow-core/tests/unit/models/test_dagrun.py
+++ b/airflow-core/tests/unit/models/test_dagrun.py
@@ -68,7 +68,7 @@ from airflow.models.dagrun import (
)
from airflow.models.deadline import Deadline
from airflow.models.deadline_alert import DeadlineAlert as DeadlineAlertModel
-from airflow.models.dynamic_region import LOOP_DECISION_KEY, SENTINEL_REGION_ID, DynamicRegion
+from airflow.models.dynamic_region import LOOP_DECISION_REASON_KEY, SENTINEL_REGION_ID, DynamicRegion
from airflow.models.serialized_dag import SerializedDagModel
from airflow.models.task_coordinates import TaskCoordinateResolver
from airflow.models.taskinstance import (
@@ -6329,11 +6329,10 @@ def test_rerun_gate_continues_while_superseded_later_passes_terminate(completed_
clear_loop_task_instances([gates[2]], downstream=False, session=session)
rerun = session.get(TaskInstance, gates[2].id)
rerun.state = State.SUCCESS
- session.add(XComModelV2(task_instance_id=rerun.id, key=LOOP_DECISION_KEY, value="continue"))
session.flush()
group, _ = TaskCoordinateResolver(DBDagBag(), session).loop_context(rerun)
- dr.complete_loop_gate(rerun, group, State.SUCCESS, session=session)
+ dr.complete_loop_gate(rerun, group, State.SUCCESS, stop=False, reason=None, session=session)
assert gates[4].state == State.RESTARTING
assert session.scalars(
@@ -6346,6 +6345,37 @@ def test_rerun_gate_continues_while_superseded_later_passes_terminate(completed_
).all() == [None]
+@pytest.mark.parametrize(
+ ("index", "state", "stored"),
+ [(4, State.SUCCESS, ["why"]), (4, State.FAILED, ["why"]), (2, State.SUCCESS, []), (2, State.FAILED, [])],
+ ids=["last-success", "last-failed", "earlier-success", "earlier-failed"],
+)
+def test_complete_loop_gate_records_the_reason_only_without_a_later_pass(
+ completed_loop, session, index, state, stored
+):
+ dr, loop, _ = completed_loop
+ gate = next(
+ ti
+ for ti in dr.get_task_instances(session=session)
+ if ti.task_id == loop.gate_task_id and ti.region_index == index
+ )
+ resolver = TaskCoordinateResolver(DBDagBag(), session)
+ group, _ = resolver.loop_context(gate)
+ before = len(dr.get_task_instances(session=session))
+
+ dr.complete_loop_gate(gate, group, state, stop=index == 4, reason="why", session=session)
+
+ assert len(dr.get_task_instances(session=session)) == before
+ assert (
+ session.scalars(
+ select(XComModelV2.value).where(
+ XComModelV2.task_instance_id == gate.id, XComModelV2.key == LOOP_DECISION_REASON_KEY
+ )
+ ).all()
+ == stored
+ )
+
+
@pytest.fixture
def locked_dag_run_keys(session):
locked: list[list[tuple[str, str]]] = []
diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go
index 6e53729efe9..0a5dba694e6 100644
--- a/go-sdk/pkg/execution/genmodels/models.gen.go
+++ b/go-sdk/pkg/execution/genmodels/models.gen.go
@@ -1336,6 +1336,15 @@ type LoopContext struct {
TerminalTaskID string `msgpack:"terminal_task_id"`
}
+// A loop gate's evaluation of its “until“ callable.
+type LoopDecisionPayload struct {
+ // Reason corresponds to the JSON schema field "reason".
+ Reason interface{} `msgpack:"reason,omitempty"`
+
+ // Stop corresponds to the JSON schema field "stop".
+ Stop bool `msgpack:"stop"`
+}
+
// Add a new value to be redacted in task logs.
type MaskSecret struct {
// Name corresponds to the JSON schema field "name".
@@ -1648,6 +1657,9 @@ type SucceedTask struct {
// EndDate corresponds to the JSON schema field "end_date".
EndDate time.Time `msgpack:"end_date"`
+ // LoopDecision corresponds to the JSON schema field "loop_decision".
+ LoopDecision *LoopDecisionPayload `msgpack:"loop_decision,omitempty"`
+
// OutletEvents corresponds to the JSON schema field "outlet_events".
OutletEvents *OutletEvents `msgpack:"outlet_events,omitempty"`
@@ -1720,11 +1732,13 @@ type TIRunContext struct {
XcomKeysToClear []string `msgpack:"xcom_keys_to_clear,omitempty"`
}
+type TaskArgBinding interface{}
+
type TriggerKwargs map[string]interface{}
type Warnings []interface{}
-type TaskArgBinding interface{}
+type TaskIds []string
// Variable schema for responses with fields that are needed for Runtime.
type VariableResponse struct {
@@ -1735,8 +1749,6 @@ type VariableResponse struct {
Value interface{} `msgpack:"value"`
}
-type TaskIds []string
-
// Schema for TaskInstance model with minimal required fields needed for Runtime.
type TaskInstance struct {
// ContextCarrier corresponds to the JSON schema field "context_carrier".
@@ -1787,6 +1799,7 @@ const TaskInstanceStateRestarting TaskInstanceState = "restarting"
const TaskInstanceStateSuccess TaskInstanceState = "success"
const TaskInstanceStateRunning TaskInstanceState = "running"
const TaskInstanceStateQueued TaskInstanceState = "queued"
+
const TaskInstanceStateScheduled TaskInstanceState = "scheduled"
const TaskInstanceStateRemoved TaskInstanceState = "removed"
@@ -1864,6 +1877,9 @@ type TaskState struct {
// EndDate corresponds to the JSON schema field "end_date".
EndDate interface{} `msgpack:"end_date,omitempty"`
+ // LoopDecision corresponds to the JSON schema field "loop_decision".
+ LoopDecision *LoopDecisionPayload `msgpack:"loop_decision,omitempty"`
+
// RenderedMapIndex corresponds to the JSON schema field "rendered_map_index".
RenderedMapIndex interface{} `msgpack:"rendered_map_index,omitempty"`
diff --git a/go-sdk/schema/supervisor-schema.json b/go-sdk/schema/supervisor-schema.json
index 77888c5adfd..41a6ab4c6d3 100644
--- a/go-sdk/schema/supervisor-schema.json
+++ b/go-sdk/schema/supervisor-schema.json
@@ -3278,6 +3278,33 @@
"title": "LoopContext",
"type": "object"
},
+ "LoopDecisionPayload": {
+ "additionalProperties": false,
+ "description": "A loop gate's evaluation of its ``until`` callable.",
+ "properties": {
+ "stop": {
+ "title": "Stop",
+ "type": "boolean"
+ },
+ "reason": {
+ "anyOf": [
+ {
+ "type": "string"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null,
+ "title": "Reason"
+ }
+ },
+ "required": [
+ "stop"
+ ],
+ "title": "LoopDecisionPayload",
+ "type": "object"
+ },
"MaskSecret": {
"description": "Add a new value to be redacted in task logs.",
"properties": {
@@ -4114,6 +4141,17 @@
"default": null,
"title": "Rendered Map Index"
},
+ "loop_decision": {
+ "anyOf": [
+ {
+ "$ref": "#/$defs/LoopDecisionPayload"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null
+ },
"type": {
"const": "SucceedTask",
"default": "SucceedTask",
@@ -4360,6 +4398,17 @@
],
"default": null,
"title": "Retry Reason"
+ },
+ "loop_decision": {
+ "anyOf": [
+ {
+ "$ref": "#/$defs/LoopDecisionPayload"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null
}
},
"required": [
diff --git a/task-sdk/docs/api.rst b/task-sdk/docs/api.rst
index 1de617447fd..e6750cdf4bc 100644
--- a/task-sdk/docs/api.rst
+++ b/task-sdk/docs/api.rst
@@ -130,6 +130,8 @@ Tasks & Operators
-----------------
.. autoapiclass:: airflow.sdk.TaskGroup
+.. autoapiclass:: airflow.sdk.LoopDecision
+
.. autoclass:: airflow.sdk.TaskInstance
.. autoapiclass:: airflow.sdk.XComArg
diff --git a/task-sdk/docs/index.rst b/task-sdk/docs/index.rst
index 9e15dae85e5..c0653f6ee3f 100644
--- a/task-sdk/docs/index.rst
+++ b/task-sdk/docs/index.rst
@@ -100,6 +100,7 @@ Why use ``airflow.sdk``?
- :class:`airflow.sdk.RetryPolicy`
- :class:`airflow.sdk.RetryRule`
- :class:`airflow.sdk.TaskGroup`
+- :class:`airflow.sdk.LoopDecision`
- :class:`airflow.sdk.TaskInstanceState`
- :class:`airflow.sdk.DagRunState`
- :class:`airflow.sdk.WeightRule`
diff --git a/task-sdk/src/airflow/sdk/__init__.py b/task-sdk/src/airflow/sdk/__init__.py
index c163df52577..3a2a2fb4a0a 100644
--- a/task-sdk/src/airflow/sdk/__init__.py
+++ b/task-sdk/src/airflow/sdk/__init__.py
@@ -63,6 +63,7 @@ __all__ = [
"HourWindow",
"IdentityMapper",
"Label",
+ "LoopDecision",
"Metadata",
"MinimumCount",
"MonthWindow",
@@ -165,6 +166,7 @@ if TYPE_CHECKING:
from airflow.sdk.definitions.decorators import result, setup, task, teardown
from airflow.sdk.definitions.decorators.task_group import task_group
from airflow.sdk.definitions.edges import EdgeModifier, Label
+ from airflow.sdk.definitions.loop import LoopDecision
from airflow.sdk.definitions.param import Param, ParamsDict
from airflow.sdk.definitions.partition_mappers.allowed_key import AllowedKeyMapper
from airflow.sdk.definitions.partition_mappers.base import (
@@ -279,6 +281,7 @@ __lazy_imports: dict[str, str] = {
"HourWindow": ".definitions.partition_mappers.window",
"IdentityMapper": ".definitions.partition_mappers.identity",
"Label": ".definitions.edges",
+ "LoopDecision": ".definitions.loop",
"Metadata": ".definitions.asset.metadata",
"MinimumCount": ".definitions.partition_mappers.wait_policy",
"MonthWindow": ".definitions.partition_mappers.window",
diff --git a/task-sdk/src/airflow/sdk/__init__.pyi b/task-sdk/src/airflow/sdk/__init__.pyi
index 1f9e10b5f5a..221d3eb48a6 100644
--- a/task-sdk/src/airflow/sdk/__init__.pyi
+++ b/task-sdk/src/airflow/sdk/__init__.pyi
@@ -70,6 +70,7 @@ from airflow.sdk.definitions.decorators import (
)
from airflow.sdk.definitions.decorators.task_group import task_group as task_group
from airflow.sdk.definitions.edges import EdgeModifier as EdgeModifier, Label as Label
+from airflow.sdk.definitions.loop import LoopDecision as LoopDecision
from airflow.sdk.definitions.param import Param as Param
from airflow.sdk.definitions.partition_mappers.allowed_key import AllowedKeyMapper
from airflow.sdk.definitions.partition_mappers.base import PartitionMapper, RollupMapper
@@ -182,6 +183,7 @@ __all__ = [
"HourWindow",
"IdentityMapper",
"Label",
+ "LoopDecision",
"Metadata",
"MinimumCount",
"MonthWindow",
diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py
index 42afaace779..cc67a2697ab 100644
--- a/task-sdk/src/airflow/sdk/api/client.py
+++ b/task-sdk/src/airflow/sdk/api/client.py
@@ -63,6 +63,7 @@ from airflow.sdk.api.datamodels._generated import (
HITLDetailResponse,
HITLUser,
InactiveAssetsResponse,
+ LoopDecisionPayload,
Note,
PrevSuccessfulDagRunResponse,
ResultMessage,
@@ -281,6 +282,7 @@ class TaskInstanceOperations:
retry_reason: str | None = None,
*,
pid: int | None = None,
+ loop_decision: LoopDecisionPayload | None = None,
):
"""Report a terminal outcome or acknowledge server-requested termination."""
if state == TaskInstanceState.SUCCESS:
@@ -291,6 +293,8 @@ class TaskInstanceOperations:
rendered_map_index=rendered_map_index,
retry_reason=retry_reason,
)
+ if loop_decision is not None:
+ body.loop_decision = loop_decision
if state == TerminalStateNonSuccess.SERVER_TERMINATED:
body.hostname = get_hostname()
body.pid = pid
@@ -313,13 +317,23 @@ class TaskInstanceOperations:
)
self.client.patch(f"task-instances/{id}/state", content=body.model_dump_json())
- def succeed(self, id: uuid.UUID, when: datetime, task_outlets, outlet_events, rendered_map_index):
+ def succeed(
+ self,
+ id: uuid.UUID,
+ when: datetime,
+ task_outlets,
+ outlet_events,
+ rendered_map_index,
+ *,
+ loop_decision: LoopDecisionPayload | None = None,
+ ):
"""Tell the API server that this TI has succeeded."""
body = TISuccessStatePayload(
end_date=when,
task_outlets=task_outlets,
outlet_events=outlet_events,
rendered_map_index=rendered_map_index,
+ loop_decision=loop_decision,
)
self.client.patch(f"task-instances/{id}/state", content=body.model_dump_json())
diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
index b5343588148..63231a2c144 100644
--- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
+++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
@@ -245,6 +245,18 @@ class LoopContext(BaseModel):
terminal_is_mapped: Annotated[bool, Field(title="Terminal Is Mapped")]
+class LoopDecisionPayload(BaseModel):
+ """
+ A loop gate's evaluation of its ``until`` callable.
+ """
+
+ model_config = ConfigDict(
+ extra="forbid",
+ )
+ stop: Annotated[bool, Field(title="Stop")]
+ reason: Annotated[str | None, Field(title="Reason")] = None
+
+
class PrevSuccessfulDagRunResponse(BaseModel):
"""
Schema for response with previous successful DagRun information for Task Template Context.
@@ -387,6 +399,7 @@ class TISuccessStatePayload(BaseModel):
task_outlets: Annotated[list[AssetProfile] | None, Field(title="Task Outlets")] = None
outlet_events: Annotated[list[dict[str, Any]] | None, Field(title="Outlet Events")] = None
rendered_map_index: Annotated[str | None, Field(title="Rendered Map Index")] = None
+ loop_decision: LoopDecisionPayload | None = None
class TITargetStatePayload(BaseModel):
@@ -759,6 +772,7 @@ class TITerminalStatePayload(BaseModel):
retry_reason: Annotated[str | None, Field(title="Retry Reason")] = None
hostname: Annotated[str | None, Field(title="Hostname")] = None
pid: Annotated[int | None, Field(title="Pid")] = None
+ loop_decision: LoopDecisionPayload | None = None
class XComArgBinding(BaseModel):
diff --git a/task-sdk/src/airflow/sdk/definitions/_internal/loop.py b/task-sdk/src/airflow/sdk/definitions/_internal/loop.py
index 67f2f4dc95c..e3c138ae5a4 100644
--- a/task-sdk/src/airflow/sdk/definitions/_internal/loop.py
+++ b/task-sdk/src/airflow/sdk/definitions/_internal/loop.py
@@ -28,9 +28,9 @@ from airflow.sdk.definitions.taskgroup import MappedTaskGroup, TaskGroup
if TYPE_CHECKING:
from airflow.sdk.definitions.context import Context
from airflow.sdk.definitions.decorators.task_group import _TaskGroupFactory
+ from airflow.sdk.definitions.loop import LoopDecision
-LOOP_XCOM_PREFIX = "_airflow_loop_"
-LOOP_DECISION_KEY = f"{LOOP_XCOM_PREFIX}decision"
+LOOP_DECISION_REASON_MAX_BYTES = 512
def _get_until_doc(until: Callable[..., Any] | None) -> str | None:
@@ -162,12 +162,12 @@ def check_expansions_over_loop_tasks(dag: Any) -> None:
class LoopGateOperator(BaseOperator):
"""Definition of a loop gate task."""
- def __init__(self, *, until: Callable[..., bool] | None = None, **kwargs: Any):
+ def __init__(self, *, until: Callable[..., bool | LoopDecision] | None = None, **kwargs: Any):
super().__init__(**kwargs)
self.until = until
def execute(self, context: Context) -> None:
- # execution_time.loop imports LOOP_DECISION_KEY from this module.
+ # execution_time.loop imports LOOP_DECISION_REASON_MAX_BYTES from this module.
from airflow.sdk.execution_time.loop import execute_loop_gate
execute_loop_gate(self, context)
@@ -178,7 +178,7 @@ def create_loop(
/,
*,
max_iterations: int,
- until: Callable[..., bool] | None = None,
+ until: Callable[..., bool | LoopDecision] | None = None,
) -> LoopTaskGroup:
"""Construct a loop group and its gate."""
from airflow.sdk.bases.decorator import get_unique_task_id
diff --git a/task-sdk/src/airflow/sdk/definitions/decorators/task_group.py b/task-sdk/src/airflow/sdk/definitions/decorators/task_group.py
index b0148278fb4..3d010be1cdd 100644
--- a/task-sdk/src/airflow/sdk/definitions/decorators/task_group.py
+++ b/task-sdk/src/airflow/sdk/definitions/decorators/task_group.py
@@ -51,6 +51,7 @@ if TYPE_CHECKING:
OperatorExpandKwargsArgument,
)
from airflow.sdk.definitions.dag import DAG
+ from airflow.sdk.definitions.loop import LoopDecision
FParams = ParamSpec("FParams")
FReturn = TypeVar("FReturn", None, DAGNode)
@@ -132,7 +133,9 @@ class _TaskGroupFactory(ExpandableFactory, Generic[FParams, FReturn]):
self._task_group_created = True
return evolved
- def loop(self, *, max_iterations: int, until: Callable[..., bool] | None = None) -> TaskGroup:
+ def loop(
+ self, *, max_iterations: int, until: Callable[..., bool | LoopDecision] | None = None
+ ) -> TaskGroup:
"""
Repeat this group's tasks, creating each iteration when the previous gate continues.
@@ -142,6 +145,7 @@ class _TaskGroupFactory(ExpandableFactory, Generic[FParams, FReturn]):
:param max_iterations: Positive upper bound on the number of iterations.
:param until: Synchronous callable receiving task context keyword arguments.
True stops the loop. False at the iteration limit fails the gate.
+ May instead return a :class:`~airflow.sdk.LoopDecision` to attach a reason.
Omit for a fixed-count loop that succeeds at its limit.
"""
return create_loop(self, max_iterations=max_iterations, until=until)
diff --git a/task-sdk/src/airflow/sdk/definitions/loop.py b/task-sdk/src/airflow/sdk/definitions/loop.py
new file mode 100644
index 00000000000..2ce051ea9dd
--- /dev/null
+++ b/task-sdk/src/airflow/sdk/definitions/loop.py
@@ -0,0 +1,33 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import attrs
+
+
+@attrs.define(frozen=True)
+class LoopDecision:
+ """
+ What a loop's ``until`` callable returns to say whether to stop, with an optional explanation.
+
+ :param stop: True to stop the loop, False to run another iteration. False at ``max_iterations`` fails
+ the gate.
+ :param reason: Shown next to the iteration's outcome in the UI. At most 512 bytes of UTF-8.
+ """
+
+ stop: bool
+ reason: str | None = None
diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py b/task-sdk/src/airflow/sdk/execution_time/comms.py
index 9c6993f2bdb..bd5781f5dfe 100644
--- a/task-sdk/src/airflow/sdk/execution_time/comms.py
+++ b/task-sdk/src/airflow/sdk/execution_time/comms.py
@@ -79,6 +79,7 @@ from airflow.sdk.api.datamodels._generated import (
DagRunStateResponse,
HITLDetailRequest,
InactiveAssetsResponse,
+ LoopDecisionPayload,
PreviousTIResponse,
PrevSuccessfulDagRunResponse,
TaskBreadcrumbsResponse,
@@ -907,6 +908,7 @@ class TaskState(BaseModel):
type: Literal["TaskState"] = "TaskState"
rendered_map_index: str | None = None
retry_reason: str | None = None
+ loop_decision: LoopDecisionPayload | None = None
class SucceedTask(TISuccessStatePayload):
diff --git a/task-sdk/src/airflow/sdk/execution_time/loop.py b/task-sdk/src/airflow/sdk/execution_time/loop.py
index a3b52282b1f..755d5d10362 100644
--- a/task-sdk/src/airflow/sdk/execution_time/loop.py
+++ b/task-sdk/src/airflow/sdk/execution_time/loop.py
@@ -16,15 +16,17 @@
# under the License.
from __future__ import annotations
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
import attrs
+from airflow.sdk._shared.secrets_masker import redact
+from airflow.sdk.api.datamodels._generated import LoopDecisionPayload
from airflow.sdk.bases.decorator import determine_kwargs
-from airflow.sdk.definitions._internal.loop import LOOP_DECISION_KEY
+from airflow.sdk.definitions._internal.loop import LOOP_DECISION_REASON_MAX_BYTES
+from airflow.sdk.definitions.loop import LoopDecision
from airflow.sdk.definitions.xcom_arg import PlainXComArg
from airflow.sdk.exceptions import AirflowFailException
-from airflow.sdk.execution_time.comms import SetXCom
from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence
from airflow.sdk.execution_time.xcom import XCom
@@ -76,29 +78,53 @@ class LoopContextAccessor:
def execute_loop_gate(gate: LoopGateOperator, context: Context) -> None:
- # task_runner imports LoopContextAccessor from this module.
- from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
-
loop = context["loop"]
+ ti = cast("RuntimeTaskInstance", context["ti"])
+ reason: str | None = None
if gate.until is None:
stop = loop.index + 1 >= loop.max_iterations
else:
- stop = gate.until(**determine_kwargs(gate.until, (), context))
- if not isinstance(stop, bool):
- raise AirflowFailException(f"until must return a bool, got {type(stop).__name__}")
- if not stop and loop.index + 1 >= loop.max_iterations:
- raise LoopMaxIterationsExceeded(
- f"Loop condition remained false at max_iterations={loop.max_iterations}"
+ match result := gate.until(**determine_kwargs(gate.until, (), context)):
+ case bool():
+ stop = result
+ case LoopDecision(stop=bool(), reason=None | str()):
+ stop, reason = result.stop, result.reason or None
+ case LoopDecision(stop=bool()):
+ raise AirflowFailException(
+ f"LoopDecision reason must be a str or None, got {type(result.reason).__name__}"
+ )
+ case LoopDecision():
+ raise AirflowFailException(
+ f"LoopDecision stop must be a bool, got {type(result.stop).__name__}"
+ )
+ case _:
+ raise AirflowFailException(
+ f"until must return a bool or LoopDecision, got {type(result).__name__}"
+ )
+ if reason is not None:
+ # Mask here, where we can know what the task has masked.
+ reason = cast("str", redact(reason))
+ try:
+ encoded = reason.encode()
+ except UnicodeEncodeError:
+ raise AirflowFailException("LoopDecision reason must be valid UTF-8")
+ size = len(encoded)
+ if size > LOOP_DECISION_REASON_MAX_BYTES:
+ raise AirflowFailException(
+ f"LoopDecision reason must be at most {LOOP_DECISION_REASON_MAX_BYTES} bytes of UTF-8"
)
- decision = "stop" if stop else "continue"
- gate.log.info("Loop iteration %s: %s (max_iterations=%s)", loop.index, decision, loop.max_iterations)
- ti = context["ti"]
- SUPERVISOR_COMMS.send(
- SetXCom(
- dag_id=ti.dag_id,
- run_id=ti.run_id,
- task_id=ti.task_id,
- key=LOOP_DECISION_KEY,
- value=decision,
+ if "\x00" in reason:
+ raise AirflowFailException("LoopDecision reason cannot contain NUL bytes")
+ ti._loop_decision = LoopDecisionPayload(stop=stop, reason=reason)
+ suffix = "" if reason is None else f": {reason}"
+ if not stop and loop.index + 1 >= loop.max_iterations:
+ raise LoopMaxIterationsExceeded(
+ f"Loop condition remained false at max_iterations={loop.max_iterations}{suffix}"
)
+ gate.log.info(
+ "Loop iteration %s: %s (max_iterations=%s)%s",
+ loop.index,
+ "stop" if stop else "continue",
+ loop.max_iterations,
+ suffix,
)
diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json
index 77888c5adfd..41a6ab4c6d3 100644
--- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json
+++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json
@@ -3278,6 +3278,33 @@
"title": "LoopContext",
"type": "object"
},
+ "LoopDecisionPayload": {
+ "additionalProperties": false,
+ "description": "A loop gate's evaluation of its ``until`` callable.",
+ "properties": {
+ "stop": {
+ "title": "Stop",
+ "type": "boolean"
+ },
+ "reason": {
+ "anyOf": [
+ {
+ "type": "string"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null,
+ "title": "Reason"
+ }
+ },
+ "required": [
+ "stop"
+ ],
+ "title": "LoopDecisionPayload",
+ "type": "object"
+ },
"MaskSecret": {
"description": "Add a new value to be redacted in task logs.",
"properties": {
@@ -4114,6 +4141,17 @@
"default": null,
"title": "Rendered Map Index"
},
+ "loop_decision": {
+ "anyOf": [
+ {
+ "$ref": "#/$defs/LoopDecisionPayload"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null
+ },
"type": {
"const": "SucceedTask",
"default": "SucceedTask",
@@ -4360,6 +4398,17 @@
],
"default": null,
"title": "Retry Reason"
+ },
+ "loop_decision": {
+ "anyOf": [
+ {
+ "$ref": "#/$defs/LoopDecisionPayload"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null
}
},
"required": [
diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py
index d1efcf0b8c3..c9a80625c3e 100644
--- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py
+++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py
@@ -29,7 +29,12 @@ from __future__ import annotations
from cadwyn import VersionChange, schema
from airflow.dag_processing.processor import DagFileParsingResult # noqa: SDK002
-from airflow.sdk.api.datamodels._generated import PreviousTIResponse, TaskInstance, TIRunContext
+from airflow.sdk.api.datamodels._generated import (
+ PreviousTIResponse,
+ TaskInstance,
+ TIRunContext,
+ TISuccessStatePayload,
+)
from airflow.sdk.execution_time.comms import (
GetXCom,
GetXComCount,
@@ -40,7 +45,7 @@ from airflow.sdk.execution_time.comms import (
class AddLoopRuntime(VersionChange):
- """Carry loop context and previous-iteration reads."""
+ """Carry loop context, previous-iteration reads and loop gate decisions."""
description = __doc__
instructions_to_migrate_to_previous_version = (
@@ -49,6 +54,9 @@ class AddLoopRuntime(VersionChange):
schema(GetXComCount).field("previous_iteration").didnt_exist,
schema(GetXComSequenceItem).field("previous_iteration").didnt_exist,
schema(GetXComSequenceSlice).field("previous_iteration").didnt_exist,
+ schema(TaskState).field("loop_decision").didnt_exist,
+ # SucceedTask inherits loop_decision, and Cadwyn only removes a field from the class that declares it.
+ schema(TISuccessStatePayload).field("loop_decision").didnt_exist,
)
diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
index d6ff0b03f40..5f870d7be78 100644
--- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
+++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
@@ -1824,6 +1824,7 @@ class ActivitySubprocess(WatchedSubprocess):
when=msg.end_date or datetime.now(tz=UTC),
rendered_map_index=self._rendered_map_index,
retry_reason=msg.retry_reason,
+ loop_decision=msg.loop_decision,
)
elif isinstance(msg, SucceedTask):
self.client.task_instances.succeed(
@@ -1832,6 +1833,7 @@ class ActivitySubprocess(WatchedSubprocess):
task_outlets=msg.task_outlets,
outlet_events=msg.outlet_events,
rendered_map_index=self._rendered_map_index,
+ loop_decision=msg.loop_decision,
)
elif isinstance(msg, RetryTask):
self.client.task_instances.retry(
diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
index b90ec01e21e..77dbe0b1bc8 100644
--- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
+++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
@@ -52,6 +52,7 @@ from airflow.sdk.api.client import get_hostname, getuser
from airflow.sdk.api.datamodels._generated import (
AssetProfile,
DagRun,
+ LoopDecisionPayload,
PreviousTIResponse,
TaskInstance,
TaskInstanceState,
@@ -63,7 +64,7 @@ from airflow.sdk.bases.xcom import BaseXCom
from airflow.sdk.configuration import conf
from airflow.sdk.coordinators._dag_importer import find_claiming_importer
from airflow.sdk.definitions._internal.dag_parsing_context import _airflow_parsing_context_manager
-from airflow.sdk.definitions._internal.loop import LOOP_XCOM_PREFIX, is_loop_task_read_from_outside
+from airflow.sdk.definitions._internal.loop import is_loop_task_read_from_outside
from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet, is_arg_set
from airflow.sdk.definitions.asset import (
Asset,
@@ -239,6 +240,9 @@ class RuntimeTaskInstance(TaskInstance):
_ti_context_from_server: Annotated[TIRunContext | None, Field(repr=False)] = None
"""The Task Instance context from the API server, if any."""
+ _loop_decision: LoopDecisionPayload | None = None
+ """A loop gate's evaluation of ``until``, reported with its terminal state."""
+
max_tries: int = 0
"""The maximum number of retries for the task."""
@@ -1625,8 +1629,7 @@ def _run_task_and_map_outcome(
if ti._ti_context_from_server and (keys_to_delete := ti._ti_context_from_server.xcom_keys_to_clear):
for x in keys_to_delete:
log.debug("Clearing XCom with key", key=x)
- backend = BaseXCom if x.startswith(LOOP_XCOM_PREFIX) else XCom
- backend.delete(
+ XCom.delete(
key=x,
dag_id=ti.dag_id,
task_id=ti.task_id,
@@ -1708,6 +1711,7 @@ def _run_task_and_map_outcome(
state=TaskInstanceState.FAILED,
end_date=ti.end_date,
rendered_map_index=ti.rendered_map_index,
+ loop_decision=ti._loop_decision,
)
state = TaskInstanceState.FAILED
error = e
@@ -1840,6 +1844,7 @@ def _handle_current_task_success(
task_outlets=task_outlets,
outlet_events=outlet_events,
rendered_map_index=ti.rendered_map_index,
+ loop_decision=ti._loop_decision,
)
return msg, TaskInstanceState.SUCCESS
diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py
index 62fde3373e0..c92653386cf 100644
--- a/task-sdk/tests/task_sdk/api/test_client.py
+++ b/task-sdk/tests/task_sdk/api/test_client.py
@@ -44,6 +44,7 @@ from airflow.sdk.api.datamodels._generated import (
HITLDetailRequest,
HITLDetailResponse,
HITLUser,
+ LoopDecisionPayload,
TaskStateStoreResponse,
TerminalTIState,
VariableResponse,
@@ -475,6 +476,32 @@ class TestTaskInstanceOperations:
retry_reason=retry_reason,
)
+ @pytest.mark.parametrize(
+ ("method", "kwargs"),
+ [
+ ("finish", {"state": TerminalTIState.FAILED}),
+ ("succeed", {"task_outlets": [], "outlet_events": []}),
+ ],
+ )
+ def test_task_instance_terminal_state_sends_loop_decision(self, method, kwargs):
+ ti_id = uuid6.uuid7()
+ bodies = []
+
+ def handle_request(request: httpx2.Request) -> httpx2.Response:
+ bodies.append(json.loads(request.read()))
+ return httpx2.Response(status_code=204)
+
+ client = make_client(transport=httpx2.MockTransport(handle_request))
+ getattr(client.task_instances, method)(
+ ti_id,
+ when="2024-10-31T12:00:00Z",
+ rendered_map_index=None,
+ loop_decision=LoopDecisionPayload(stop=False, reason="error 0.2"),
+ **kwargs,
+ )
+
+ assert bodies[0]["loop_decision"] == {"stop": False, "reason": "error 0.2"}
+
def test_task_instance_heartbeat(self):
# Simulate a successful response from the server that sends a heartbeat for a ti
ti_id = uuid6.uuid7()
diff --git a/task-sdk/tests/task_sdk/execution_time/test_loop.py b/task-sdk/tests/task_sdk/execution_time/test_loop.py
index a2485584670..b226a8bf77b 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_loop.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_loop.py
@@ -16,21 +16,23 @@
# under the License.
from __future__ import annotations
+import re
+
import pytest
-from airflow.sdk import DAG, BaseOperator, TaskInstanceState, task, task_group
-from airflow.sdk.api.datamodels._generated import LoopContext
-from airflow.sdk.bases.xcom import BaseXCom
+from airflow.sdk import DAG, BaseOperator, LoopDecision, TaskInstanceState, task, task_group
+from airflow.sdk._shared.secrets_masker import _secrets_masker
+from airflow.sdk.api.datamodels._generated import LoopContext, LoopDecisionPayload
from airflow.sdk.definitions._internal.loop import create_loop
from airflow.sdk.exceptions import AirflowFailException
from airflow.sdk.execution_time import task_runner
from airflow.sdk.execution_time.comms import (
- DeleteXCom,
GetXCom,
GetXComCount,
GetXComSequenceItem,
GetXComSequenceSlice,
- SetXCom,
+ SucceedTask,
+ TaskState,
XComCountResponse,
XComResult,
XComSequenceIndexResult,
@@ -119,60 +121,141 @@ def test_previous_and_current_result_preserve_falsey_values(loop_ti, mock_superv
@pytest.mark.parametrize(
- ("index", "condition", "decision"),
- [(0, None, "continue"), (2, None, "stop"), (0, False, "continue"), (0, True, "stop"), (2, True, "stop")],
+ ("index", "until", "expected"),
+ [
+ (0, None, LoopDecisionPayload(stop=False)),
+ (2, None, LoopDecisionPayload(stop=True)),
+ (0, lambda: False, LoopDecisionPayload(stop=False)),
+ (0, lambda: True, LoopDecisionPayload(stop=True)),
+ (2, lambda: True, LoopDecisionPayload(stop=True)),
+ (0, lambda: LoopDecision(stop=True), LoopDecisionPayload(stop=True)),
+ (
+ 0,
+ lambda: LoopDecision(stop=False, reason="score 0.7, needs > 0.9"),
+ LoopDecisionPayload(stop=False, reason="score 0.7, needs > 0.9"),
+ ),
+ (
+ 0,
+ lambda: LoopDecision(stop=True, reason="é" * 256),
+ LoopDecisionPayload(stop=True, reason="é" * 256),
+ ),
+ (0, lambda: LoopDecision(stop=True, reason=""), LoopDecisionPayload(stop=True)),
+ (
+ 0,
+ lambda: LoopDecision(stop=True, reason="token hunter2 expired"),
+ LoopDecisionPayload(stop=True, reason="token *** expired"),
+ ),
+ ],
+ ids=[
+ "fixed_continue",
+ "fixed_stop_at_cap",
+ "bool_continue",
+ "bool_stop",
+ "bool_stop_at_cap",
+ "decision_without_reason",
+ "decision_with_reason",
+ "reason_at_byte_limit",
+ "empty_reason_is_no_reason",
+ "secret_in_reason_is_masked",
+ ],
)
-def test_gate_publishes_successful_decision(loop_ti, mock_supervisor_comms, index, condition, decision):
- def until(*, loop):
- assert loop.index == index
- return condition
+@pytest.mark.enable_redact
+def test_gate_sends_its_decision_with_success(loop_ti, mock_supervisor_comms, index, until, expected):
+ _secrets_masker().add_mask("hunter2", None)
+ ti = loop_ti(index=index, until=until)
+
+ state, message, error = task_runner.run(ti, context=ti.get_template_context(), log=ti.task.log)
- ti = loop_ti(index=index, until=until if condition is not None else None)
- context = ti.get_template_context()
+ assert (state, error) == (TaskInstanceState.SUCCESS, None)
+ assert isinstance(message, SucceedTask)
+ assert message.loop_decision == expected
- ti.task.execute(context)
- message = mock_supervisor_comms.send.call_args.args[0]
- assert isinstance(message, SetXCom)
- assert message.key == "_airflow_loop_decision"
- assert message.value == decision
+def _raise_value_error():
+ raise ValueError("condition failed")
@pytest.mark.parametrize(
- ("raises", "error", "message"),
- [(False, LoopMaxIterationsExceeded, "max_iterations"), (True, ValueError, "condition failed")],
+ ("until", "error", "message"),
+ [
+ (_raise_value_error, ValueError, "condition failed"),
+ (lambda: None, AirflowFailException, "got NoneType"),
+ (lambda: 0, AirflowFailException, "got int"),
+ (lambda: 1, AirflowFailException, "got int"),
+ (lambda: [], AirflowFailException, "got list"),
+ (lambda: "yes", AirflowFailException, "got str"),
+ (lambda: LoopDecision(stop=1), AirflowFailException, "LoopDecision stop must be a bool, got int"), # type: ignore[arg-type]
+ (
+ lambda: LoopDecision(stop=None, reason="r"), # type: ignore[arg-type]
+ AirflowFailException,
+ "LoopDecision stop must be a bool, got NoneType",
+ ),
+ (
+ lambda: LoopDecision(stop=True, reason=0.2), # type: ignore[arg-type]
+ AirflowFailException,
+ "LoopDecision reason must be a str or None, got float",
+ ),
+ (
+ lambda: LoopDecision(stop=True, reason=b"r"), # type: ignore[arg-type]
+ AirflowFailException,
+ "LoopDecision reason must be a str or None, got bytes",
+ ),
+ (lambda: LoopDecision(stop=True, reason="x" * 513), AirflowFailException, "512 bytes"),
+ (lambda: LoopDecision(stop=True, reason="é" * 257), AirflowFailException, "512 bytes"),
+ (lambda: LoopDecision(stop=True, reason="\ud800"), AirflowFailException, "valid UTF-8"),
+ ],
+ ids=[
+ "until_raises",
+ "none",
+ "zero",
+ "one",
+ "list",
+ "str",
+ "int_stop",
+ "none_stop",
+ "float_reason",
+ "bytes_reason",
+ "ascii_reason_over_limit",
+ "utf8_reason_over_limit",
+ "unpaired_surrogate_reason",
+ ],
)
-def test_unsuccessful_gate_does_not_publish_decision(loop_ti, mock_supervisor_comms, raises, error, message):
- def until(*, loop):
- if raises:
- raise ValueError("condition failed")
- return False
-
- ti = loop_ti(index=2, until=until)
- with pytest.raises(error, match=message):
- ti.task.execute(ti.get_template_context())
-
- mock_supervisor_comms.send.assert_not_called()
+def test_gate_fails_without_a_decision(loop_ti, mock_supervisor_comms, until, error, message):
+ ti = loop_ti(index=0, until=until)
+ state, sent, raised = task_runner.run(ti, context=ti.get_template_context(), log=ti.task.log)
-def test_gate_fails_without_retrying_when_the_condition_is_false_at_the_cap(loop_ti, mock_supervisor_comms):
- ti = loop_ti(index=2, until=lambda: False)
+ assert state == TaskInstanceState.FAILED
+ assert isinstance(raised, error)
+ assert re.search(message, str(raised))
+ assert isinstance(sent, TaskState)
+ assert sent.loop_decision is None
- with pytest.raises(LoopMaxIterationsExceeded, match="max_iterations=3") as exc_info:
- ti.task.execute(ti.get_template_context())
-
- assert isinstance(exc_info.value, AirflowFailException)
- mock_supervisor_comms.send.assert_not_called()
+@pytest.mark.parametrize(
+ ("result", "message", "reason"),
+ [
+ (False, r"max_iterations=3$", None),
+ (
+ LoopDecision(stop=False, reason="error 0.2, needs < 1e-06"),
+ r"max_iterations=3: error 0.2, needs",
+ "error 0.2, needs < 1e-06",
+ ),
+ ],
+ ids=["bool", "decision_with_reason"],
+)
+def test_gate_fails_without_retrying_when_the_condition_is_false_at_the_cap(
+ loop_ti, mock_supervisor_comms, result, message, reason
+):
+ ti = loop_ti(index=2, until=lambda: result)
-@pytest.mark.parametrize("result", [None, 0, 1, [], "yes"])
-def test_gate_fails_when_until_does_not_return_a_bool(loop_ti, mock_supervisor_comms, result):
- ti = loop_ti(index=0, until=lambda: result)
+ state, sent, error = task_runner.run(ti, context=ti.get_template_context(), log=ti.task.log)
- with pytest.raises(AirflowFailException, match=f"got {type(result).__name__}"):
- ti.task.execute(ti.get_template_context())
-
- mock_supervisor_comms.send.assert_not_called()
+ assert state == TaskInstanceState.FAILED
+ assert isinstance(error, LoopMaxIterationsExceeded)
+ assert re.search(message, str(error))
+ assert isinstance(sent, TaskState)
+ assert sent.loop_decision == LoopDecisionPayload(stop=False, reason=reason)
@pytest.mark.parametrize("previous_iteration", [False, True])
@@ -202,44 +285,3 @@ def test_mapped_terminal_keeps_iteration_across_lazy_reads(
assert result[::-1] == values[::-1]
if values:
assert result[-1] == values[-1]
-
-
-def test_loop_decision_bypasses_custom_backend(loop_ti, mock_supervisor_comms, mocker):
- serialize = mocker.patch("airflow.sdk.execution_time.xcom.XCom.serialize_value", autospec=True)
- ti = loop_ti()
-
- ti.task.execute(ti.get_template_context())
-
- serialize.assert_not_called()
- assert mock_supervisor_comms.send.call_args.args[0].value == "continue"
-
-
-def test_gate_retry_clears_old_signal_without_custom_backend(loop_ti, mock_supervisor_comms, monkeypatch):
- class CustomXCom(BaseXCom):
- @classmethod
- def purge(cls, xcom, *args):
- raise AssertionError("Control metadata must not reach custom backend")
-
- ti = loop_ti()
- ti._ti_context_from_server.xcom_keys_to_clear = ["_airflow_loop_decision"]
- monkeypatch.setattr(task_runner, "XCom", CustomXCom)
- mock_supervisor_comms.send.side_effect = lambda message: (
- XComResult(key="_airflow_loop_decision", value="stop") if isinstance(message, GetXCom) else None
- )
-
- state, _, error = task_runner.run(ti, context=ti.get_template_context(), log=ti.task.log)
-
- assert error is None
- assert state == TaskInstanceState.SUCCESS
- decisions = [
- call.args[0]
- for call in mock_supervisor_comms.send.call_args_list
- if call.args and isinstance(call.args[0], SetXCom)
- ]
- assert [decision.value for decision in decisions] == ["continue"]
- deletes = [
- call.args[0]
- for call in mock_supervisor_comms.send.call_args_list
- if call.args and isinstance(call.args[0], DeleteXCom)
- ]
- assert [delete.key for delete in deletes] == ["_airflow_loop_decision"]
diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
index 8aabe2ff5d1..7c3fee25cce 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
@@ -68,6 +68,7 @@ from airflow.sdk.api.datamodels._generated import (
DagRun,
DagRunState,
DagRunType,
+ LoopDecisionPayload,
PreviousTIResponse,
TaskInstance,
TaskInstanceState,
@@ -2676,7 +2677,9 @@ REQUEST_TEST_CASES = [
),
RequestTestCase(
message=SucceedTask(
- end_date=timezone.parse("2024-10-31T12:00:00Z"), rendered_map_index="test success task"
+ end_date=timezone.parse("2024-10-31T12:00:00Z"),
+ rendered_map_index="test success task",
+ loop_decision=LoopDecisionPayload(stop=True, reason="converged"),
),
client_mock=ClientMock(
method_path="task_instances.succeed",
@@ -2686,6 +2689,7 @@ REQUEST_TEST_CASES = [
"task_outlets": None,
"when": timezone.parse("2024-10-31T12:00:00Z"),
"rendered_map_index": "test success task",
+ "loop_decision": LoopDecisionPayload(stop=True, reason="converged"),
},
),
test_id="succeed_task",
@@ -4741,15 +4745,21 @@ class TestHandleRequest:
proc.client.task_instances.finish.assert_called_once()
assert proc._pending_terminal_state_msg is None
- def test_task_state_retry_reason_forwarded_to_finish(self, watched_subprocess, mocker):
- """A TaskState message's retry_reason must reach the deferred finish() call."""
+ @pytest.mark.parametrize(
+ "extra",
+ [
+ {"retry_reason": "auth error, do not retry"},
+ {"loop_decision": LoopDecisionPayload(stop=False, reason="error 0.2")},
+ ],
+ ids=["retry_reason", "loop_decision"],
+ )
+ def test_task_state_extras_forwarded_to_finish(self, watched_subprocess, mocker, extra):
+ """Optional TaskState fields must reach the deferred finish() call."""
watched_subprocess, _ = watched_subprocess
watched_subprocess._exit_code = 0
msg = TaskState(
- state=TaskInstanceState.FAILED,
- end_date=timezone.parse("2024-10-31T12:00:00Z"),
- retry_reason="auth error, do not retry",
+ state=TaskInstanceState.FAILED, end_date=timezone.parse("2024-10-31T12:00:00Z"), **extra
)
watched_subprocess._handle_request(msg, mocker.Mock(), req_id=1)
@@ -4760,7 +4770,7 @@ class TestHandleRequest:
state=TaskInstanceState.FAILED,
when=mocker.ANY,
rendered_map_index=None,
- retry_reason="auth error, do not retry",
+ **{"retry_reason": None, "loop_decision": None, **extra},
)
diff --git a/ts-sdk/schema/supervisor-schema.json b/ts-sdk/schema/supervisor-schema.json
index 77888c5adfd..41a6ab4c6d3 100644
--- a/ts-sdk/schema/supervisor-schema.json
+++ b/ts-sdk/schema/supervisor-schema.json
@@ -3278,6 +3278,33 @@
"title": "LoopContext",
"type": "object"
},
+ "LoopDecisionPayload": {
+ "additionalProperties": false,
+ "description": "A loop gate's evaluation of its ``until`` callable.",
+ "properties": {
+ "stop": {
+ "title": "Stop",
+ "type": "boolean"
+ },
+ "reason": {
+ "anyOf": [
+ {
+ "type": "string"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null,
+ "title": "Reason"
+ }
+ },
+ "required": [
+ "stop"
+ ],
+ "title": "LoopDecisionPayload",
+ "type": "object"
+ },
"MaskSecret": {
"description": "Add a new value to be redacted in task logs.",
"properties": {
@@ -4114,6 +4141,17 @@
"default": null,
"title": "Rendered Map Index"
},
+ "loop_decision": {
+ "anyOf": [
+ {
+ "$ref": "#/$defs/LoopDecisionPayload"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null
+ },
"type": {
"const": "SucceedTask",
"default": "SucceedTask",
@@ -4360,6 +4398,17 @@
],
"default": null,
"title": "Retry Reason"
+ },
+ "loop_decision": {
+ "anyOf": [
+ {
+ "$ref": "#/$defs/LoopDecisionPayload"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null
}
},
"required": [
diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts
index 7faf8c48770..8b018e48514 100644
--- a/ts-sdk/src/generated/supervisor.ts
+++ b/ts-sdk/src/generated/supervisor.ts
@@ -526,6 +526,8 @@ export type AssignedUsers1 = HITLUser[] | null;
export type Type54 = "HITLDetailRequestResult";
export type InactiveAssets = AssetProfile[] | null;
export type Type55 = "InactiveAssetsResult";
+export type Stop1 = boolean;
+export type Reason = string | null;
export type Name14 = string | null;
export type Type56 = "MaskSecret";
export type Ok = boolean;
@@ -1618,6 +1620,16 @@ export interface InactiveAssetsResult {
inactive_assets?: InactiveAssets;
type?: Type55;
}
+/**
+ * A loop gate's evaluation of its ``until`` callable.
+ *
+ * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema
+ * via the `definition` "LoopDecisionPayload".
+ */
+export interface LoopDecisionPayload {
+ stop: Stop1;
+ reason?: Reason;
+}
/**
* Add a new value to be redacted in task logs.
*
@@ -1853,6 +1865,7 @@ export interface SucceedTask {
task_outlets?: TaskOutlets;
outlet_events?: OutletEvents;
rendered_map_index?: RenderedMapIndex4;
+ loop_decision?: LoopDecisionPayload | null;
type?: Type75;
}
/**
@@ -1899,6 +1912,7 @@ export interface TaskState {
type?: Type79;
rendered_map_index?: RenderedMapIndex5;
retry_reason?: RetryReason1;
+ loop_decision?: LoopDecisionPayload | null;
}
/**
* Response to GetTaskStateStore; wraps the generated API response for supervisor to worker comms.
diff --git a/uv.lock b/uv.lock
index 05d90471521..f2742ccf22c 100644
--- a/uv.lock
+++ b/uv.lock
@@ -2130,7 +2130,7 @@ requires-dist = [
{ name = "pendulum", specifier = ">=3.1.0" },
{ name = "pluggy", specifier = ">=1.5.0" },
{ name = "psutil", specifier = ">=5.8.0" },
- { name = "pydantic", specifier = ">=2.11.0" },
+ { name = "pydantic", specifier = ">=2.12.0" },
{ name = "pygments", specifier = ">=2.0.1,!=2.19.0" },
{ name = "pygtrie", specifier = ">=2.5.0" },
{ name = "pyjwt", specifier = ">=2.11.0" },