Clear language SDK task XComs in the coordinator when a try starts

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

@@ -339,10 +339,12 @@ def ti_run(
                 },
             )
 
-        # Send the keys to the SDK so that the client requests to clear those XComs from the server.
-        # The reason we cannot do this here in the server is because we need to issue a purge on custom XCom backends
-        # too. With the current assumption, the workers ONLY have access to the custom XCom backends directly and they
-        # can issue the purge.
+        # Send the keys to the worker, which deletes those XComs. The server cannot delete them itself,
+        # because for a Python task the task runner also purges them from a custom XCom backend, and only
+        # workers can reach that backend. For a language SDK task (Java, Go, TypeScript), the supervisor
+        # deletes these XComs before the runtime starts the task and does not purge them, because those
+        # runtimes push their XComs straight to the API server, so the XCom values never pass through the
+        # backend.
 
         # However, do not clear it for deferral
         xcom_keys = []

task-sdk/src/airflow/sdk/coordinators/_subprocess.py

@@ -37,6 +37,7 @@
 import stat
 import subprocess
 import time
+from datetime import UTC, datetime
 from pathlib import PurePosixPath
 from typing import TYPE_CHECKING, NoReturn, TypeVar, cast
 
@@ -46,10 +47,11 @@
 
 from airflow.dag_processing.bundles.base import BundleVersionLock, unpack_bundle_version  # noqa: SDK002
 from airflow.sdk._shared.module_loading import import_string
-from airflow.sdk.api.datamodels._generated import BundleInfo
+from airflow.sdk.api.datamodels._generated import BundleInfo, TaskInstanceState
 from airflow.sdk.configuration import conf
 from airflow.sdk.coordinators._dag_importer import find_claiming_importer
 from airflow.sdk.execution_time.bundles import initialize_ti_bundle
+from airflow.sdk.execution_time.comms import RetryTask, TaskState
 from airflow.sdk.execution_time.coordinator import BaseCoordinator, TaskLaunchError
 from airflow.sdk.execution_time.schema import get_schema_version_migrator
 from airflow.sdk.execution_time.supervisor import ActivitySubprocess, NeverRaised, ProcessTracker
@@ -63,7 +65,7 @@
 
     from airflow.dag_processing.bundles.base import BaseDagBundle  # noqa: SDK002
     from airflow.sdk.api.client import Client
-    from airflow.sdk.api.datamodels._generated import TaskInstance
+    from airflow.sdk.api.datamodels._generated import TaskInstance, TIRunContext
 
     Tracked = TypeVar("Tracked", socket.socket, subprocess.Popen)
 
@@ -403,6 +405,41 @@ def start(  # type: ignore[override]
 
         return self
 
+    def _prepare_ti_context(self, ti: TaskInstance, ti_context: TIRunContext) -> TIRunContext | None:
+        """
+        Delete the XComs in ``ti_context.xcom_keys_to_clear`` so the launched runtime does not have to.
+
+        Unlike ``XCom.delete``, this method does not purge the XCom backend, for two reasons. XComs that a
+        runtime pushes go straight to the API server and never pass through the backend. Also, a purge may
+        need a connection that a hook running in the supervisor process cannot look up, such as a
+        connection that exists only in Airflow's metadata database. So if Python code stored an XCom for
+        this task instance through the backend, deleting that XCom here leaves the stored object in the
+        backend's storage.
+
+        If an XCom cannot be deleted, this try fails and the task does not run.
+        """
+        for key in ti_context.xcom_keys_to_clear or ():
+            self.process_log.debug("Clearing XCom with key", key=key)
+            try:
+                self.client.xcoms.delete(ti.dag_id, ti.run_id, ti.task_id, key, ti.map_index)
+            except Exception as e:
+                self.process_log.exception("Failed to clear XCom; not running the task", key=key)
+                end_date = datetime.now(tz=UTC)
+                reason = f"Failed to clear XCom {key!r}: {e}"
+                # wait() sends this outcome to the server after the runtime process exits, the same
+                # way it sends an outcome that the runtime itself reported.
+                self._pending_terminal_state_msg = (
+                    RetryTask(end_date=end_date, retry_reason=reason)
+                    if ti_context.should_retry
+                    else TaskState(state=TaskInstanceState.FAILED, end_date=end_date, retry_reason=reason)
+                )
+                self.kill(signal.SIGKILL)
+                return None
+        # The supervisor has already deleted the listed XComs, so it sends the runtime an empty list.
+        # If a runtime also deletes these XComs itself, the empty list keeps the runtime from deleting
+        # them a second time.
+        return ti_context.model_copy(update={"xcom_keys_to_clear": []})
+
     def wait(self) -> int:
         code = super().wait()
         self._close_unused_sockets(self._comm_server, self._logs_server)

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

@@ -191,6 +191,7 @@
     from structlog.typing import FilteringBoundLogger, WrappedLogger
 
     from airflow.executors.workloads import BundleInfo
+    from airflow.sdk.api.datamodels._generated import TIRunContext
     from airflow.sdk.bases.secrets_backend import BaseSecretsBackend
     from airflow.sdk.definitions.connection import Connection
     from airflow.sdk.types import RuntimeTaskInstanceProtocol as RuntimeTI
@@ -1735,6 +1736,10 @@ def _on_child_started(
                 return
             raise
 
+        if (prepared := self._prepare_ti_context(ti, ti_context)) is None:
+            return
+        ti_context = prepared
+
         # ti_context.start_date is only populated by the server when resuming from a deferral (to preserve the
         # original start_date rather than using the resume time). We fall back to now() otherwise. This ensures
         # that `context["ti"].start_date` always reflects the *first* start time. See TIRunContext.start_date
@@ -1759,6 +1764,17 @@ def _on_child_started(
             # Debug is fine, the process will have shown _something_ in it's last_chance exception handler
             log.debug("Couldn't send startup message to Subprocess - it died very early", pid=self.pid)
 
+    def _prepare_ti_context(self, ti: TaskInstance, ti_context: TIRunContext) -> TIRunContext | None:
+        """
+        Return the context the supervisor sends to the task process, or None when the task must not run.
+
+        An implementation that returns None must first stop the task process and record the outcome that
+        will be reported. This base version returns the context unchanged, because the Python task runner
+        itself deletes the XComs in ``ti_context.xcom_keys_to_clear`` with ``XCom.delete``, which also
+        purges them from the XCom backend.
+        """
+        return ti_context
+
     def wait(self) -> int:
         if self._wait_completed and self._exit_code is not None:
             return self._exit_code

task-sdk/tests/task_sdk/coordinators/test_subprocess.py

@@ -19,6 +19,7 @@
 
 import contextlib
 import json
+import logging
 import os
 import pathlib
 import signal
@@ -27,16 +28,20 @@
 import sys
 import threading
 import time
+from datetime import UTC, datetime
 from unittest.mock import ANY, MagicMock, call, patch
 
 import attrs
+import httpx2
+import msgspec
 import psutil
 import pytest
+from task_sdk import FAKE_BUNDLE, make_client
 from uuid6 import uuid7
 
 from airflow.dag_processing.bundles.base import BaseDagBundle, BundleVersion
 from airflow.sdk.api.client import Client, TaskInstanceOperations
-from airflow.sdk.api.datamodels._generated import BundleInfo, TaskInstance
+from airflow.sdk.api.datamodels._generated import BundleInfo, TaskInstance, TIRunContext
 from airflow.sdk.coordinators._dag_importer import CoordinatorDagImporter
 from airflow.sdk.coordinators._subprocess import (
     SubprocessCoordinator,
@@ -61,15 +66,15 @@
     pytest.skip("Coordinator is only compatible with Airflow >= 3.3.0", allow_module_level=True)
 
 
-def _make_ti(dag_id: str = "tutorial_dag", queue: str = "socket") -> TaskInstance:
+def _make_ti(dag_id: str = "tutorial_dag", queue: str = "socket", map_index: int = -1) -> TaskInstance:
     return TaskInstance(
         id=uuid7(),
         dag_version_id=uuid7(),
         task_id="task_1",
         dag_id=dag_id,
         run_id="run_1",
         try_number=1,
-        map_index=-1,
+        map_index=map_index,
         queue=queue,
     )
 
@@ -933,6 +938,174 @@ def test_register_pipe_readers_called_with_four_sockets(self, mock_client):
         assert mock_register.mock_calls == [call(ANY, ANY, ANY, ANY, data=ANY)]
 
 
+# This script connects back to the supervisor the way a language SDK runtime does, then waits for
+# StartupDetails. It exits with code 0 as soon as the supervisor sends anything or closes the
+# connection. It exits with code 3 if 10 seconds pass and neither has happened.
+_RUNTIME_AWAITING_STARTUP = """
+import socket
+import sys
+
+def connect(flag):
+    host, _, port = next(arg for arg in sys.argv if arg.startswith(flag)).split("=", 1)[1].rpartition(":")
+    return socket.create_connection((host, int(port)))
+
+comm = connect("--comm=")
+logs = connect("--logs=")
+comm.settimeout(10)
+try:
+    comm.recv(1)
+except TimeoutError:
+    sys.exit(3)
+"""
+
+
+class TestPopenActivitySubprocessClearXComs:
+    @pytest.mark.parametrize(
+        ("map_index", "query"),
+        [
+            pytest.param(-1, "", id="unmapped"),
+            pytest.param(0, "map_index=0", id="map-index-0"),
+            pytest.param(2, "map_index=2", id="map-index-2"),
+        ],
+    )
+    @pytest.mark.log_level(logging.DEBUG)
+    @patch.object(ActivitySubprocess, "_register_pipe_readers", autospec=True)
+    @patch("airflow.sdk.coordinators._subprocess._accept_connections", autospec=True)
+    @patch("airflow.sdk.coordinators._subprocess.subprocess.Popen", autospec=True)
+    def test_deletes_listed_xcoms_before_the_runtime_gets_the_task(
+        self, mock_popen, mock_accept, mock_register, make_ti_context_dict, captured_logs, map_index, query
+    ):
+        ti = _make_ti(map_index=map_index)
+        ti_context = {
+            **make_ti_context_dict(),
+            "xcom_keys_to_clear": ["return_value", "skipmixin_key"],
+            "start_date": "2024-12-01T00:30:00Z",
+            "arg_bindings": [{"name": "country", "kind": "literal", "value": "uk"}],
+        }
+        events: list = []
+
+        def logged_keys() -> list[tuple[str, str, str]]:
+            return [
+                (log["logger"], log["level"], log["key"])
+                for log in captured_logs
+                if log.get("event") == "Clearing XCom with key"
+            ]
+
+        def handle_request(request: httpx2.Request) -> httpx2.Response:
+            events.append((request.method, request.url.path, request.url.query.decode(), logged_keys()))
+            if request.url.path.endswith("/run"):
+                return httpx2.Response(200, json=ti_context)
+            return httpx2.Response(200, json={})
+
+        def record_frame(frame: bytearray) -> None:
+            body = msgspec.msgpack.decode(frame[4:])[1]
+            events.append((body["type"], TIRunContext.model_validate(body["ti_context"])))
+
+        comm_sock = MagicMock(spec=socket.socket)
+        comm_sock.sendall.side_effect = record_frame
+        mock_popen.return_value.pid = 12345
+        mock_accept.side_effect = lambda servers, drains, proc, **kw: (
+            {servers["comm"]: comm_sock, servers["logs"]: MagicMock(spec=socket.socket)},
+            {soc: b"" for soc in drains.values()},
+        )
+
+        _PopenActivitySubprocess.start(
+            what=ti,
+            dag_rel_path="bundle",
+            bundle_info=FAKE_BUNDLE,
+            client=make_client(transport=httpx2.MockTransport(handle_request)),
+            command=["/bin/true"],
+            subprocess_logs_to_stdout=False,
+        )
+
+        first = ("task", "debug", "return_value")
+        second = ("task", "debug", "skipmixin_key")
+        assert events == [
+            ("PATCH", f"/task-instances/{ti.id}/run", "", []),
+            ("DELETE", "/xcoms/tutorial_dag/run_1/task_1/return_value", query, [first]),
+            ("DELETE", "/xcoms/tutorial_dag/run_1/task_1/skipmixin_key", query, [first, second]),
+            ("StartupDetails", TIRunContext.model_validate({**ti_context, "xcom_keys_to_clear": []})),
+        ]
+
+    @pytest.mark.parametrize(
+        ("delete_error", "should_retry", "state", "error_text"),
+        [
+            pytest.param("forbidden", True, "up_for_retry", "Forbidden", id="forbidden-retry"),
+            pytest.param("forbidden", False, "failed", "Forbidden", id="forbidden-no-retry"),
+            pytest.param("unreachable", True, "up_for_retry", "Connection refused", id="unreachable-retry"),
+        ],
+    )
+    # The patch turns off the API client's retry policy, so the case where the API server is
+    # unreachable fails at once instead of after the client's retries and their waits.
+    @patch.object(Client, "_request_with_retry", Client._request_without_retry)
+    @patch.object(ActivitySubprocess, "send_msg", autospec=True)
+    def test_fails_the_try_without_starting_the_task_when_an_xcom_cannot_be_deleted(
+        self,
+        mock_send_msg,
+        make_ti_context_dict,
+        captured_logs,
+        time_machine,
+        delete_error,
+        should_retry,
+        state,
+        error_text,
+    ):
+        time_machine.move_to(datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC), tick=False)
+        ti = _make_ti()
+        ti_context = {
+            **make_ti_context_dict(),
+            "xcom_keys_to_clear": ["return_value", "skipmixin_key", "summary"],
+            "should_retry": should_retry,
+        }
+        requests: list[tuple[str, str]] = []
+        reports: list[tuple[str, str, str]] = []
+
+        def handle_request(request: httpx2.Request) -> httpx2.Response:
+            requests.append((request.method, request.url.path))
+            if request.url.path.endswith("/run"):
+                return httpx2.Response(200, json=ti_context)
+            if request.method == "DELETE":
+                if not request.url.path.endswith("/skipmixin_key"):
+                    return httpx2.Response(200, json={})
+                if delete_error == "unreachable":
+                    raise httpx2.ConnectError("Connection refused", request=request)
+                return httpx2.Response(403, json={"detail": "Forbidden"})
+            report = json.loads(request.content)
+            reports.append((report["state"], report["end_date"], report["retry_reason"]))
+            return httpx2.Response(204)
+
+        proc = _PopenActivitySubprocess.start(
+            what=ti,
+            dag_rel_path="bundle",
+            bundle_info=FAKE_BUNDLE,
+            client=make_client(transport=httpx2.MockTransport(handle_request)),
+            command=[sys.executable, "-c", _RUNTIME_AWAITING_STARTUP],
+            subprocess_logs_to_stdout=False,
+        )
+
+        assert proc.wait() == -signal.SIGKILL
+        mock_send_msg.assert_not_called()
+        assert requests == [
+            ("PATCH", f"/task-instances/{ti.id}/run"),
+            ("DELETE", "/xcoms/tutorial_dag/run_1/task_1/return_value"),
+            ("DELETE", "/xcoms/tutorial_dag/run_1/task_1/skipmixin_key"),
+            ("PATCH", f"/task-instances/{ti.id}/state"),
+        ]
+        assert reports == [
+            (state, "2026-01-02T03:04:05Z", f"Failed to clear XCom 'skipmixin_key': {error_text}"),
+        ]
+        assert proc.final_state == state
+        assert {
+            "event": "Failed to clear XCom; not running the task",
+            "key": "skipmixin_key",
+            "level": "error",
+            "logger": "task",
+            "timestamp": ANY,
+            "exc_info": ANY,
+            "loc": ANY,
+        } in captured_logs
+
+
 class TestResolveArtifactBundle:
     """Execute-time resolution of the Dag bundle holding the artifacts."""