Let an XCom read tell the value it loads which iteration it came from

loop.previous read an iterated terminal's XComIterable from the previous
iteration, then patched the pointer from outside so its items were read
from that iteration too: an isinstance check in loop.py, and any other read
with previous_iteration (a lazy sequence item) kept the bug.

The XCom read now passes its previous_iteration flag to _load, which holds
it in a context variable while the value is deserialized, and
XComIterable.deserialize builds the iterable with it. The flag cannot go
through deserialize_value, whose signature XCom backends implement.
loop.py no longer needs to know about XComIterable.

LazyXComSequence.aget loaded through deserialize_value directly and so
skipped a backend's load_value; it now goes through XCom._load.

Author: dabla

task-sdk/src/airflow/sdk/bases/xcom.py

@@ -143,6 +143,12 @@ class _OwnTaskInstance(Protocol):
     "airflow_current_trigger_ti", default=None
 )
 
+# Whether the XCom value being loaded was read from the iteration before the reader's, so a value that points
+# at other XComs reads them from that iteration too; see ``XComIterable.deserialize``.
+_loading_previous_iteration: contextvars.ContextVar[bool] = contextvars.ContextVar(
+    "airflow_loading_previous_iteration", default=False
+)
+
 
 class TIKeyProtocol(Protocol):
     dag_id: str
@@ -527,7 +533,7 @@ def get_one(
             raise TypeError(f"Expected XComResult, received: {type(msg)} {msg}")
 
         if msg.value is not None:
-            return cls._load(msg)
+            return cls._load(msg, previous_iteration=_previous_iteration)
         log.debug(
             "No XCom value found; defaulting to None.",
             key=key,
@@ -591,7 +597,7 @@ async def aget_one(
             raise TypeError(f"Expected XComResult, received: {type(msg)} {msg}")
 
         if msg.value is not None:
-            return await asyncio.to_thread(cls._load, msg)
+            return await asyncio.to_thread(cls._load, msg, previous_iteration=_previous_iteration)
         log.debug(
             "No XCom value found; defaulting to None.",
             key=key,
@@ -743,7 +749,14 @@ def _load_many(cls, values: list[JsonValue]) -> list[Any]:
         return [cls._load(_XComValueWrapper(value)) for value in values]
 
     @classmethod
-    def _load(cls, result: Any) -> Any:
+    def _load(cls, result: Any, *, previous_iteration: bool = False) -> Any:
+        if previous_iteration:
+            # The flag cannot travel through deserialize_value, whose signature backends implement.
+            token = _loading_previous_iteration.set(True)
+            try:
+                return cls._load(result)
+            finally:
+                _loading_previous_iteration.reset(token)
         # Legacy backends keep receiving the result object deserialize_value has always been given.
         if _is_preferred_over(cls, "load_value", "deserialize_value"):
             return cls.load_value(result.value)
@@ -828,8 +841,8 @@ class XComIterable(Sequence):
     iterating and slicing: N elements are N requests. Reading several in one request needs an
     Execution API endpoint that takes a list of keys, which does not exist yet.
 
-    Inside a loop, the result an iterated terminal pushed in the previous iteration
-    (``loop.previous``) has ``previous_iteration`` set, so its elements are read from that
+    An iterable read from the previous iteration (``loop.previous`` inside a loop) is built with
+    ``previous_iteration`` set (see :meth:`deserialize`), so its elements are read from that
     iteration too and not from the current one. "Previous" is relative to whoever reads, so
     such an iterable cannot be pushed on: :meth:`serialize` refuses it (see there).
     """
@@ -948,8 +961,13 @@ def serialize(self) -> dict:
 
     @classmethod
     def deserialize(cls, data: dict, version: int):
-        """Ensure the object is JSON deserializable."""
-        return cls(**data)
+        """
+        Ensure the object is JSON deserializable.
+
+        The iterable is a pointer whose item reads resolve to the reader's iteration, so one loaded by a
+        read from the previous iteration reads its items from that iteration as well.
+        """
+        return cls(**data, previous_iteration=_loading_previous_iteration.get())
 
 
 class _AsyncXComIterator:

task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py

@@ -141,7 +141,7 @@ async def aget(self, index: int) -> T:
             raise IndexError(index)
         if not isinstance(msg, XComSequenceIndexResult):
             raise TypeError(f"Got unexpected response to GetXComSequenceItem: {msg!r}")
-        return XCom.deserialize_value(_XComWrapper(msg.root))
+        return XCom._load(_XComWrapper(msg.root), previous_iteration=self._previous_iteration)
 
     def __len__(self) -> int:
         if self._len is None:
@@ -224,7 +224,10 @@ def __getitem__(self, key: int | slice) -> T | Sequence[T]:
             )
             if not isinstance(msg, XComSequenceSliceResult):
                 raise TypeError(f"Got unexpected response to GetXComSequenceSlice: {msg!r}")
-            return [XCom._load(_XComWrapper(value)) for value in msg.root]
+            return [
+                XCom._load(_XComWrapper(value), previous_iteration=self._previous_iteration)
+                for value in msg.root
+            ]
 
         if not isinstance(key, int):
             if (index := getattr(key, "__index__", None)) is not None:
@@ -246,7 +249,7 @@ def __getitem__(self, key: int | slice) -> T | Sequence[T]:
             raise IndexError(key)
         if not isinstance(msg, XComSequenceIndexResult):
             raise TypeError(f"Got unexpected response to GetXComSequenceItem: {msg!r}")
-        return XCom._load(_XComWrapper(msg.root))
+        return XCom._load(_XComWrapper(msg.root), previous_iteration=self._previous_iteration)
 
 
 def _coerce_slice_index(value: Any) -> int | None:

task-sdk/src/airflow/sdk/execution_time/loop.py

@@ -23,7 +23,6 @@
 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.bases.xcom import XComIterable
 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
@@ -69,17 +68,13 @@ def _result(self, *, previous_iteration: bool) -> Any:
             return LazyXComSequence(
                 xcom_arg=PlainXComArg(task), ti=self._ti, previous_iteration=previous_iteration
             )
-        result = XCom.get_one(
+        return XCom.get_one(
             key=XCom.XCOM_RETURN_KEY,
             dag_id=self._ti.dag_id,
             run_id=self._ti.run_id,
             task_id=task.task_id,
             _previous_iteration=previous_iteration,
         )
-        if previous_iteration and isinstance(result, XComIterable):
-            # An iterated terminal pushes a pointer to its items: read them in the same iteration.
-            result.previous_iteration = True
-        return result
 
 
 def execute_loop_gate(gate: LoopGateOperator, context: Context) -> None:

task-sdk/tests/task_sdk/bases/test_xcom.py

@@ -337,6 +337,35 @@ def test_an_iterable_read_through_loop_previous_is_refused_when_pushed(self):
         with pytest.raises(TypeError, match="read through loop.previous cannot be pushed as is"):
             BaseXCom.serialize_value(iterable)
 
+    @pytest.mark.parametrize("previous_iteration", [False, True])
+    @pytest.mark.parametrize("use_async", [False, True])
+    @pytest.mark.asyncio
+    async def test_an_iterable_read_from_the_previous_iteration_reads_its_items_from_it(
+        self, previous_iteration, use_async, mock_supervisor_comms
+    ):
+        """The XCom read that loads the pointer tells it which iteration it was read from."""
+        response = XComResult(
+            key="return_value", value=BaseXCom.serialize_value(self.make_iterable(length=2))
+        )
+        coordinates = dict(key="return_value", dag_id="dag", task_id="task", run_id="run")
+
+        if use_async:
+            mock_supervisor_comms.asend = AsyncMock(return_value=response)
+            result = await XCom.aget_one(**coordinates, _previous_iteration=previous_iteration)
+        else:
+            mock_supervisor_comms.send.return_value = response
+            result = XCom.get_one(**coordinates, _previous_iteration=previous_iteration)
+
+        assert isinstance(result, XComIterable)
+        assert result.previous_iteration is previous_iteration
+        # The flag lives for the load only: the next read starts from the current iteration again.
+        assert (
+            XComIterable.deserialize(
+                {"task_id": "t", "dag_id": "d", "run_id": "r"}, version=1
+            ).previous_iteration
+            is False
+        )
+
     def test_deserialize_restores_fields(self):
         data = {"task_id": "task", "dag_id": "dag", "run_id": "run", "map_index": 2, "length": 5}
         iterable = XComIterable.deserialize(data, version=1)

task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py

@@ -22,7 +22,7 @@
 import pytest
 
 import airflow
-from airflow.sdk.bases.xcom import BaseXCom
+from airflow.sdk.bases.xcom import BaseXCom, XComIterable
 from airflow.sdk.exceptions import ErrorType
 from airflow.sdk.execution_time.comms import (
     ErrorResponse,
@@ -224,6 +224,38 @@ def test_getitem_loads_through_load_value(
     assert lazy_sequence[index] == expected
 
 
+@pytest.mark.asyncio
+async def test_aget_loads_through_load_value(monkeypatch, mock_supervisor_comms, lazy_sequence):
+    mock_supervisor_comms.asend = AsyncMock(return_value=XComSequenceIndexResult(root="ref-4"))
+    monkeypatch.setattr(airflow.sdk.execution_time.xcom, "XCom", LoadValueXCom)
+
+    assert await lazy_sequence.aget(4) == "loaded ref-4"
+
+
+@pytest.mark.parametrize("previous_iteration", [False, True])
+@pytest.mark.parametrize("read", ["index", "slice", "aget"])
+@pytest.mark.asyncio
+async def test_an_iterable_item_read_from_the_previous_iteration_reads_its_items_from_it(
+    read, previous_iteration, mock_supervisor_comms, mock_xcom_arg, mock_ti
+):
+    """A mapped terminal's item can be an iterated result: it reads its items from the iteration it came from."""
+    stored = BaseXCom.serialize_value(XComIterable(task_id="t", dag_id="dag", run_id="run", length=1))
+    sequence = LazyXComSequence(mock_xcom_arg, mock_ti, previous_iteration=previous_iteration)
+
+    if read == "aget":
+        mock_supervisor_comms.asend = AsyncMock(return_value=XComSequenceIndexResult(root=stored))
+        item = await sequence.aget(0)
+    elif read == "slice":
+        mock_supervisor_comms.send.return_value = XComSequenceSliceResult(root=[stored])
+        (item,) = sequence[0:1]
+    else:
+        mock_supervisor_comms.send.return_value = XComSequenceIndexResult(root=stored)
+        item = sequence[0]
+
+    assert isinstance(item, XComIterable)
+    assert item.previous_iteration is previous_iteration
+
+
 def test_getitem_indexerror(mock_supervisor_comms, lazy_sequence):
     mock_supervisor_comms.send.return_value = ErrorResponse(
         error=ErrorType.XCOM_NOT_FOUND,