From feb3cdcb432a7786971cca5eb35a4dea6f624cff Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Wed, 29 Jul 2026 22:35:23 +0000 Subject: [PATCH 1/2] Add XCom read access to deadline callbacks run by executors Executor-run deadline callbacks could read Connections and Variables through the callback supervisor but not XComs, so an alert could not include an upstream task's output. Serve GetXCom through the shared handle_get_xcom handler, and add context["xcom_pull"] bound to the callback's Dag run so a callback only names the task to pull from. Only reads are exposed. The request uses the execution token from the callback token exchange, so no API routes or token scopes change. --- airflow-core/docs/howto/deadline-alerts.rst | 3 + .../sdk/execution_time/callback_supervisor.py | 49 ++++++++++++-- .../test_callback_supervisor.py | 65 +++++++++++++++++++ 3 files changed, 111 insertions(+), 6 deletions(-) diff --git a/airflow-core/docs/howto/deadline-alerts.rst b/airflow-core/docs/howto/deadline-alerts.rst index 538173c35055b..714397dafc0e0 100644 --- a/airflow-core/docs/howto/deadline-alerts.rst +++ b/airflow-core/docs/howto/deadline-alerts.rst @@ -387,6 +387,9 @@ as long as the variables being templated are included in the simplified context. ID and the calculated deadline time of the Deadline Alert as well as the data included in the ``GET`` REST API response for Dag Run. Support for more comprehensive context and templating will be added in future versions. +Synchronous callbacks can also read XComs from tasks in the same Dag run through ``context["xcom_pull"]``, +for example ``context["xcom_pull"]("extract")`` or ``context["xcom_pull"]("extract", key="row_count")``. + Deadline Calculation ^^^^^^^^^^^^^^^^^^^^ diff --git a/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py b/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py index bc38e2abb20e2..3c772a828591b 100644 --- a/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py @@ -39,12 +39,14 @@ GetConnection, GetVariable, GetVariableKeys, + GetXCom, MaskSecret, ) from airflow.sdk.execution_time.request_handlers import ( handle_get_connection, handle_get_variable, handle_get_variable_keys, + handle_get_xcom, handle_mask_secret, ) from airflow.sdk.execution_time.supervisor import ( @@ -78,14 +80,45 @@ class _BundleInfoLike(Protocol): # The set of messages that a callback subprocess can send to the supervisor. -# This is a minimal subset of ToSupervisor: read-only access to Connections -# and Variables, plus MaskSecret for the secrets masker. +# This is a minimal subset of ToSupervisor: read-only access to Connections, +# Variables, and XCom values, plus MaskSecret for the secrets masker. CallbackToSupervisor = Annotated[ - GetConnection | GetVariable | GetVariableKeys | MaskSecret, + GetConnection | GetVariable | GetVariableKeys | GetXCom | MaskSecret, Field(discriminator="type"), ] +def _add_xcom_pull_to_context(callback_kwargs: dict) -> None: + """ + Add an ``xcom_pull`` helper bound to the callback's Dag run to its context. + + Callbacks have no task instance, so the helper takes ``dag_id`` and ``dag_run_id`` from the + ``dag_run`` in the context and the callback only names the task to pull from. + """ + context = callback_kwargs.get("context") + if not isinstance(context, dict): + return + dag_run = context.get("dag_run") + if not isinstance(dag_run, dict) or not dag_run.get("dag_id") or not dag_run.get("dag_run_id"): + return + + def xcom_pull( + task_ids: str, *, key: str = "return_value", map_indexes: int | None = None, default: Any = None + ) -> Any: + from airflow.sdk.execution_time.xcom import XCom + + value = XCom.get_one( + key=key, + dag_id=dag_run["dag_id"], + run_id=dag_run["dag_run_id"], + task_id=task_ids, + map_index=map_indexes, + ) + return default if value is None else value + + context["xcom_pull"] = xcom_pull + + SERDE_KWARGS_ENCODING = "serde" @@ -209,9 +242,10 @@ class CallbackSubprocess(WatchedSubprocess): Uses the WatchedSubprocess infrastructure for fork/monitor/signal handling while keeping a simple lifecycle: start, run callback, exit. - Provides a limited set of comms channels (Connections and Variables) so - that callback code can access runtime services like - ``Connection.get()`` and ``Variable.get()`` via the supervisor's API client. + Provides a limited set of comms channels (Connections, Variables, and XCom) + so that callback code can access runtime services like + ``Connection.get()``, ``Variable.get()``, and ``context["xcom_pull"]()`` via + the supervisor's API client. """ client: Client # The HTTP client to use for communication with the API server. @@ -270,6 +304,7 @@ def _target(): exc_info=True, ) + _add_xcom_pull_to_context(callback_kwargs) success, error_msg = execute_callback( callback_path=callback_path, callback_kwargs=callback_kwargs, @@ -391,6 +426,8 @@ def _handle_request(self, msg: CallbackToSupervisor, log: FilteringBoundLogger, resp, dump_opts = handle_get_variable(self.client, msg) elif isinstance(msg, GetVariableKeys): resp, dump_opts = handle_get_variable_keys(self.client, msg) + elif isinstance(msg, GetXCom): + resp, dump_opts = handle_get_xcom(self.client, msg) elif isinstance(msg, MaskSecret): handle_mask_secret(msg) else: diff --git a/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py index 7eb5bc9c949ef..08c49d6cb2513 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py @@ -31,6 +31,7 @@ import pytest import structlog +from airflow.sdk.api.datamodels._generated import XComResponse from airflow.sdk.execution_time.callback_supervisor import ( CallbackSubprocess, Path, @@ -43,9 +44,11 @@ GetConnection, GetVariable, GetVariableKeys, + GetXCom, MaskSecret, VariableKeysResult, VariableResult, + XComResult, _RequestFrame, ) @@ -320,6 +323,38 @@ class RequestCase: response=VariableKeysResult(keys=["test_key"], total_entries=1), ), ), + RequestCase( + message=GetXCom( + key="return_value", + dag_id="test_dag", + run_id="test_run_1", + task_id="upstream_task", + ), + test_id="get_xcom", + client_mock=ClientMock( + method_path="xcoms.get", + args=("test_dag", "test_run_1", "upstream_task", "return_value", None, False), + kwargs={"previous_iteration": False}, + response=XComResponse(key="return_value", value="xcom_payload"), + ), + ), + RequestCase( + message=GetXCom( + key="custom_key", + dag_id="dag_a", + run_id="run_42", + task_id="task_b", + map_index=3, + include_prior_dates=True, + ), + test_id="get_xcom_with_map_index", + client_mock=ClientMock( + method_path="xcoms.get", + args=("dag_a", "run_42", "task_b", "custom_key", 3, True), + kwargs={"previous_iteration": False}, + response=XComResponse(key="custom_key", value={"nested": "data"}), + ), + ), RequestCase( message=MaskSecret(value="super_secret", name="api_key"), test_id="mask_secret", @@ -639,6 +674,36 @@ def test_callback_supervisor_should_exit_on_error(self, base_start_kwargs): assert exc_info.value.code == 1 + def _run_target_with_context(self, base_start_kwargs, context): + CallbackSubprocess.start(**{**base_start_kwargs, "callback_kwargs": {"context": context}}) + self.mock_super_start.call_args.kwargs["target"]() + return self.mock_execute_callback.call_args.kwargs["callback_kwargs"]["context"] + + @pytest.mark.parametrize( + ("xcom_value", "expected"), + [ + pytest.param("payload", "payload", id="value"), + pytest.param(None, "fallback", id="missing"), + ], + ) + def test_context_xcom_pull_uses_dag_run_ids(self, base_start_kwargs, xcom_value, expected): + comms = self.mock_comms_decoder.__getitem__.return_value.return_value + comms.send.return_value = XComResult(key="out", value=xcom_value) + + context = self._run_target_with_context( + base_start_kwargs, {"dag_run": {"dag_id": "dag_a", "dag_run_id": "run_1"}} + ) + + assert context["xcom_pull"]("upstream", key="out", map_indexes=2, default="fallback") == expected + comms.send.assert_called_once_with( + GetXCom(key="out", dag_id="dag_a", run_id="run_1", task_id="upstream", map_index=2) + ) + + def test_context_without_dag_run_has_no_xcom_pull(self, base_start_kwargs): + context = self._run_target_with_context(base_start_kwargs, {"deadline": {"id": "x"}}) + + assert "xcom_pull" not in context + @pytest.mark.parametrize("option_value", ["True", "False"]) def test_start_keeps_bare_fork_regardless_of_exec_option(self, base_start_kwargs, option_value): """ From 5f7e186d2755e3f74f886e680d35ba4ddad64f4b Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Thu, 22 Oct 2026 18:08:10 +0000 Subject: [PATCH 2/2] Use singular argument names in the callback xcom_pull and warn when it is unavailable --- .../src/airflow/sdk/execution_time/callback_supervisor.py | 7 ++++--- .../task_sdk/execution_time/test_callback_supervisor.py | 5 +++-- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py b/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py index 3c772a828591b..a2861de413cd7 100644 --- a/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py @@ -100,10 +100,11 @@ def _add_xcom_pull_to_context(callback_kwargs: dict) -> None: return dag_run = context.get("dag_run") if not isinstance(dag_run, dict) or not dag_run.get("dag_id") or not dag_run.get("dag_run_id"): + log.warning("Callback context has no Dag run, so xcom_pull is not available") return def xcom_pull( - task_ids: str, *, key: str = "return_value", map_indexes: int | None = None, default: Any = None + task_id: str, *, key: str = "return_value", map_index: int | None = None, default: Any = None ) -> Any: from airflow.sdk.execution_time.xcom import XCom @@ -111,8 +112,8 @@ def xcom_pull( key=key, dag_id=dag_run["dag_id"], run_id=dag_run["dag_run_id"], - task_id=task_ids, - map_index=map_indexes, + task_id=task_id, + map_index=map_index, ) return default if value is None else value diff --git a/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py index 08c49d6cb2513..36c0cf52ee8a3 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py @@ -694,15 +694,16 @@ def test_context_xcom_pull_uses_dag_run_ids(self, base_start_kwargs, xcom_value, base_start_kwargs, {"dag_run": {"dag_id": "dag_a", "dag_run_id": "run_1"}} ) - assert context["xcom_pull"]("upstream", key="out", map_indexes=2, default="fallback") == expected + assert context["xcom_pull"]("upstream", key="out", map_index=2, default="fallback") == expected comms.send.assert_called_once_with( GetXCom(key="out", dag_id="dag_a", run_id="run_1", task_id="upstream", map_index=2) ) - def test_context_without_dag_run_has_no_xcom_pull(self, base_start_kwargs): + def test_context_without_dag_run_has_no_xcom_pull(self, base_start_kwargs, caplog): context = self._run_target_with_context(base_start_kwargs, {"deadline": {"id": "x"}}) assert "xcom_pull" not in context + assert "Callback context has no Dag run, so xcom_pull is not available" in caplog @pytest.mark.parametrize("option_value", ["True", "False"]) def test_start_keeps_bare_fork_regardless_of_exec_option(self, base_start_kwargs, option_value):