From bb8a771b9f67208738576d4fa935f5dfbb0830a9 Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Mon, 15 Jun 2026 21:36:17 +0000 Subject: [PATCH 1/7] Add XCom read access to callback supervisor comms channel Accept workload tokens on the get_xcom, connections, and variables execution API routes so a deadline callback subprocess can read XCom values (and connections/variables) via the supervisor comms channel, mirroring task access. --- .../api_fastapi/execution_api/routes/xcoms.py | 5 +- .../sdk/execution_time/callback_supervisor.py | 19 +++-- .../test_callback_supervisor.py | 75 ++++++++++++++++++- 3 files changed, 89 insertions(+), 10 deletions(-) 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 9ab8eb200712a..2368bffefffff 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 @@ -20,7 +20,7 @@ import logging from typing import Annotated -from fastapi import APIRouter, Body, Depends, HTTPException, Path, Query, Request, Response, status +from fastapi import APIRouter, Body, Depends, HTTPException, Path, Query, Request, Response, Security, status from pydantic import JsonValue from sqlalchemy import delete from sqlalchemy.sql.selectable import Select @@ -32,7 +32,7 @@ XComSequenceIndexResponse, XComSequenceSliceResponse, ) -from airflow.api_fastapi.execution_api.security import CurrentTIToken +from airflow.api_fastapi.execution_api.security import CurrentTIToken, require_auth from airflow.models.taskmap import TaskMap from airflow.models.xcom import XComModel from airflow.utils.db import get_query_count @@ -305,6 +305,7 @@ class GetXcomFilterParams(BaseModel): @router.get( "/{dag_id}/{run_id}/{task_id}/{key:path}", description="Get a single XCom Value", + dependencies=[Security(require_auth, scopes=["token:execution", "token:workload"])], ) def get_xcom( dag_id: str, 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 3e62a90f23539..ad819a673049e 100644 --- a/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py @@ -47,12 +47,14 @@ GetDagRun, 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 ( @@ -102,11 +104,11 @@ class CallbackContextFetchError(RuntimeError): # 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, plus GetDagRun for -# building context from DagRun identifiers. +# This is a minimal subset of ToSupervisor: read-only access to Connections, +# Variables, and XCom values, plus MaskSecret for the secrets masker, plus +# GetDagRun for building context from DagRun identifiers. CallbackToSupervisor = Annotated[ - GetConnection | GetDagRun | GetVariable | GetVariableKeys | MaskSecret, + GetConnection | GetDagRun | GetVariable | GetVariableKeys | GetXCom | MaskSecret, Field(discriminator="type"), ] @@ -220,9 +222,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 ``XCom.get()`` via the + supervisor's API client. """ client: Client # The HTTP client to use for communication with the API server. @@ -435,6 +438,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 647bda8ea87f8..a999a18626bc7 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 @@ -33,7 +33,7 @@ from airflow.sdk._shared.template_rendering import render_callback_kwargs from airflow.sdk._shared.timezones import timezone -from airflow.sdk.api.datamodels._generated import DagRun, DagRunState, DagRunType +from airflow.sdk.api.datamodels._generated import DagRun, DagRunState, DagRunType, XComResponse from airflow.sdk.execution_time.callback_supervisor import ( CALLBACK_CONTEXT_FETCH_EXIT_CODE, CallbackContextFetchError, @@ -43,8 +43,17 @@ supervise_callback, ) from airflow.sdk.execution_time.comms import ( + BundleInfo, + ConnectionResult, ErrorResponse, + GetConnection, GetDagRun, + GetVariable, + GetVariableKeys, + GetXCom, + MaskSecret, + VariableKeysResult, + VariableResult, _RequestFrame, ) @@ -172,6 +181,70 @@ class RequestCase: response=_MOCK_DAG_RUN, ), ), + RequestCase( + message=GetConnection(conn_id="test_conn"), + test_id="get_connection_with_password", + client_mock=ClientMock( + method_path="connections.get", + args=("test_conn",), + response=ConnectionResult(conn_id="test_conn", conn_type="mysql", password="secret"), + ), + mask_secret_args=("secret",), + ), + RequestCase( + message=GetVariable(key="test_key"), + test_id="get_variable", + client_mock=ClientMock( + method_path="variables.get", + args=("test_key",), + response=VariableResult(key="test_key", value="test_value"), + ), + ), + RequestCase( + message=GetVariableKeys(prefix="test_"), + test_id="get_variable_keys", + client_mock=ClientMock( + method_path="variables.keys", + kwargs={"prefix": "test_", "limit": 1000, "offset": 0}, + 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", + map_index=None, + ), + test_id="get_xcom", + client_mock=ClientMock( + method_path="xcoms.get", + args=("test_dag", "test_run_1", "upstream_task", "return_value", None, 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), + response=XComResponse(key="custom_key", value={"nested": "data"}), + ), + ), + RequestCase( + message=MaskSecret(value="super_secret", name="api_key"), + test_id="mask_secret", + mask_secret_args=("super_secret", "api_key"), + ), ] @pytest.fixture From b33c4959338beaad9d2b850f5347d81e82c016b6 Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Fri, 19 Jun 2026 21:41:48 +0000 Subject: [PATCH 2/7] Grant token:workload to XCom read routes via ExecutionAPIRoute Callback supervisors carry a workload-scoped JWT, but the XCom read routes only accepted execution-scoped tokens, so a callback's XCom read got 403 ("Token type 'workload' not allowed"). Mirror the variables/connections routes: set route_class=ExecutionAPIRoute on the router (so per-route token:* scopes are honored) and add token:workload to the GET item/slice and HEAD routes alongside the existing get_xcom route. POST/DELETE stay execution-only. --- .../src/airflow/api_fastapi/execution_api/routes/xcoms.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) 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 2368bffefffff..6c8b4ad6d7468 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 @@ -32,7 +32,7 @@ XComSequenceIndexResponse, XComSequenceSliceResponse, ) -from airflow.api_fastapi.execution_api.security import CurrentTIToken, require_auth +from airflow.api_fastapi.execution_api.security import CurrentTIToken, ExecutionAPIRoute, require_auth from airflow.models.taskmap import TaskMap from airflow.models.xcom import XComModel from airflow.utils.db import get_query_count @@ -105,6 +105,7 @@ def has_xcom_access( router = APIRouter( + route_class=ExecutionAPIRoute, responses={ status.HTTP_401_UNAUTHORIZED: {"description": "Unauthorized"}, status.HTTP_403_FORBIDDEN: {"description": "Task does not have access to the XCom"}, @@ -135,6 +136,7 @@ async def xcom_query( @router.get( "/{dag_id}/{run_id}/{task_id}/{key:path}/item/{offset}", + dependencies=[Security(require_auth, scopes=["token:execution", "token:workload"])], description="Get a single XCom value from a mapped task by sequence index", ) def get_mapped_xcom_by_index( @@ -180,6 +182,7 @@ class GetXComSliceFilterParams(BaseModel): @router.get( "/{dag_id}/{run_id}/{task_id}/{key:path}/slice", + dependencies=[Security(require_auth, scopes=["token:execution", "token:workload"])], description="Get XCom values from a mapped task by sequence slice", ) def get_mapped_xcom_by_slice( @@ -273,6 +276,7 @@ def get_mapped_xcom_by_slice( }, }, }, + dependencies=[Security(require_auth, scopes=["token:execution", "token:workload"])], description="Returns the count of mapped XCom values found in the `Content-Range` response header", ) def head_xcom( From ef8576e292a23db4272103a1eac39d2a2f629b63 Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Mon, 22 Jun 2026 17:48:52 +0000 Subject: [PATCH 3/7] Re-trigger CI (flaky subprocess/triggerer test timeouts; mypy-providers is pre-existing main breakage in amazon/login.py, unrelated) From 52e67b7a82304ae00020b4ad97fea2843eef60b6 Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Mon, 22 Jun 2026 18:55:32 +0000 Subject: [PATCH 4/7] Re-trigger CI From 94fc5550d33b7d65e88ab997c604f7a9164af8d9 Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Mon, 22 Jun 2026 19:22:12 +0000 Subject: [PATCH 5/7] Re-trigger CI From 906c102194d115fa0e3c1b21c4a6db97ce78b8db Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Tue, 23 Jun 2026 02:36:55 +0000 Subject: [PATCH 6/7] Update connections token-scope test: workload tokens accepted (read) (route_class=ExecutionAPIRoute), so deadline callback subprocesses can read connections. Update test_workload_scope_rejected_on_connections_endpoint -> test_workload_scope_accepted_on_connections_endpoint: a workload token now passes the token-type check and reaches the route (404 for a missing connection) rather than being rejected with 403. Verified locally: 1 passed. --- .../execution_api/versions/head/test_task_instances.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) 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 6912d02d1be0f..ca597b17ebf09 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 @@ -4050,7 +4050,14 @@ def test_workload_scope_rejected_on_state_endpoint(self, client, session, create assert "Token type 'workload' not allowed" in resp.json()["detail"] def test_workload_scope_accepted_on_connections_endpoint(self, client, session, create_task_instance): - """Workload scoped tokens are accepted on GET /connections for deadline callback subprocesses.""" + """Workload scoped tokens are accepted on GET /connections (read access for callbacks). + + The connections router declares ``token:workload`` on its read route and sets + ``route_class=ExecutionAPIRoute`` so the scope is enforced; deadline callback + subprocesses (which carry workload tokens) must be able to read connections. + A missing connection therefore returns 404 (the request reached the route), + not a 403 token-type rejection. + """ ti = create_task_instance(task_id="test_workload_conn", state=State.RUNNING) session.commit() @@ -4059,6 +4066,7 @@ def test_workload_scope_accepted_on_connections_endpoint(self, client, session, resp = client.get("/execution/connections/test_conn") # Workload tokens are now accepted; 404 because the connection doesn't exist in the test DB. assert resp.status_code == 404 + assert "Token type 'workload' not allowed" not in str(resp.json()) def test_execution_scope_accepted_on_all_endpoints(self, client, session, create_task_instance): """Execution scoped tokens should be accepted on all endpoints.""" From d8bb9ae821c988e3692d16620e323d8d192e3018 Mon Sep 17 00:00:00 2001 From: Sean Ghaeli Date: Tue, 23 Jun 2026 08:16:40 +0000 Subject: [PATCH 7/7] Remove unused BundleInfo import (ruff F401) from rebase union Conflict resolution when rebasing onto #66608 unioned two import lists and left BundleInfo imported but unused. Drop it to satisfy ruff. --- .../tests/task_sdk/execution_time/test_callback_supervisor.py | 1 - 1 file changed, 1 deletion(-) 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 a999a18626bc7..14ee6854d42f4 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 @@ -43,7 +43,6 @@ supervise_callback, ) from airflow.sdk.execution_time.comms import ( - BundleInfo, ConnectionResult, ErrorResponse, GetConnection,