Resolve a trigger's XCom requests as its deferred task instance

airflow-core/src/airflow/api_fastapi/execution_api/in_process.py

@@ -35,6 +35,8 @@
 
 logger = structlog.get_logger(logger_name=__name__)
 
+TRIGGER_TI_ID_HEADER = "Airflow-Trigger-Ti-Id"
+
 
 # Finalizer arguments must not reference the transport, or it cannot be garbage collected.
 def _shutdown_loop(
@@ -51,12 +53,16 @@ def _shutdown_loop(
     thread.join(timeout=5)
 
 
-def _configure_in_process_auth(app: FastAPI) -> None:
+def _configure_in_process_auth(app: FastAPI, *, triggerer_id: int | None = None) -> None:
     from airflow.api_fastapi.execution_api.datamodels.token import TIClaims, TIToken
     from airflow.api_fastapi.execution_api.routes.connections import has_connection_access
     from airflow.api_fastapi.execution_api.routes.variables import has_variable_access
     from airflow.api_fastapi.execution_api.routes.xcoms import has_xcom_access
-    from airflow.api_fastapi.execution_api.security import _IN_PROCESS_NON_TI_CALLER, _jwt_bearer
+    from airflow.api_fastapi.execution_api.security import (
+        _IN_PROCESS_NON_TI_CALLER,
+        _IN_PROCESS_TRIGGERER_ID,
+        _jwt_bearer,
+    )
 
     async def always_allow(request: Request):
         from uuid import UUID
@@ -71,6 +77,8 @@ async def always_allow(request: Request):
             "task_instance_id" not in request.path_params
             and "X-Airflow-In-Process-Attempt-Id" not in request.headers
         )
+        if triggerer_id is not None:
+            request.scope[_IN_PROCESS_TRIGGERER_ID] = triggerer_id
         claims = TIClaims(scope="execution")
         return TIToken(id=ti_id, claims=claims)
 
@@ -90,6 +98,7 @@ class InProcessExecutionAPI:
     """
 
     _app: FastAPI | None = None
+    triggerer_id: int | None = None
 
     @cached_property
     def app(self):
@@ -105,7 +114,7 @@ def app(self):
             self._app = create_task_execution_api_app(lifespan=attrs.evolve(lifespan, registry=registry))
             # Auth dependency overrides bypass validation, so the stub skips validator setup in lifespan().
             registry.register_value(JWTValidator, None)
-            _configure_in_process_auth(self._app)
+            _configure_in_process_auth(self._app, triggerer_id=self.triggerer_id)
 
         return self._app
 

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

@@ -24,6 +24,7 @@
 from uuid import UUID
 
 import attrs
+from cadwyn import current_dependency_solver
 from fastapi import APIRouter, Body, Depends, HTTPException, Path, Query, Request, Response, status
 from pydantic import JsonValue
 from sqlalchemy import select
@@ -33,18 +34,20 @@
 from airflow.api_fastapi.common.db.common import SessionDep
 from airflow.api_fastapi.core_api.base import BaseModel
 from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
-from airflow.api_fastapi.execution_api.datamodels.token import DagParseToken, ExecutionToken
+from airflow.api_fastapi.execution_api.datamodels.token import DagParseToken, ExecutionToken, TIToken
 from airflow.api_fastapi.execution_api.datamodels.xcom import (
     XComResponse,
     XComSequenceIndexResponse,
     XComSequenceSliceResponse,
 )
+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,
     CurrentExecutionToken,
-    CurrentTIToken,
     DagInGrantedBundle,
     ExecutionAPIRoute,
     ExecutionOrDagParseToken,
+    require_task_token,
 )
 from airflow.api_fastapi.execution_api.versions.v2026_10_30 import IdentifyArchivedTaskStateUpdates
 from airflow.exceptions import TaskNotFound
@@ -62,6 +65,7 @@
     enclosing_loop,
 )
 from airflow.models.taskinstance import TaskInstance
+from airflow.models.trigger import Trigger
 from airflow.models.xcom import XCOM_RETURN_KEY, XComModel, xcom_entity
 from airflow.utils.db import get_query_count
 from airflow.utils.state import TaskInstanceState
@@ -252,29 +256,86 @@ def _build_xcom_read(
         raise HTTPException(status.HTTP_400_BAD_REQUEST, str(error)) from error
 
 
-def xcom_query(
-    dag_id: str,
-    run_id: str,
-    task_id: str,
-    key: str,
+def _resolve_trigger_token(
+    request: Request,
     session: SessionDep,
-    dag_bag: DagBagDep,
-    map_index: Annotated[int | None, Query()] = None,
-    previous_iteration: bool = False,
     token: ExecutionToken = CurrentExecutionToken,
-) -> Select:
-    return _build_xcom_read(
-        dag_id=dag_id,
-        run_id=run_id,
-        task_id=task_id,
-        key=key,
-        session=session,
-        dag_bag=dag_bag,
-        token=token,
-        map_index=map_index,
-        previous_iteration=previous_iteration,
-        sequence=True,
-    ).statement
+    dependency_solver=Depends(current_dependency_solver),
+) -> ExecutionToken:
+    """
+    Return the identity an XCom request acts as.
+
+    A request made by a trigger's code carries its task instance's id in ``TRIGGER_TI_ID_HEADER``. The
+    triggerer subprocess can forge the header, so it is only honoured on the triggerer's in-process API.
+
+    A write acts as the claimed task instance only while it is deferred on a trigger assigned to the calling
+    triggerer. A read acts as it while that try is live, in any state, because a trigger's ``on_kill`` runs
+    after the task instance leaves ``deferred`` and the trigger row is gone; reads are not a security
+    boundary on this API, which already serves every XCom. Anything else gets a 403.
+    """
+    header = request.headers.get(TRIGGER_TI_ID_HEADER)
+    # The endpoint runs with Cadwyn's re-solved dependencies, so checking there reuses the endpoint's session.
+    if header is None or dependency_solver != "cadwyn":
+        return token
+
+    triggerer_id = request.scope.get(_IN_PROCESS_TRIGGERER_ID)
+    if triggerer_id is None:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail={
+                "reason": "access_denied",
+                "message": f"{TRIGGER_TI_ID_HEADER} is only accepted from the triggerer",
+            },
+        )
+
+    try:
+        claimed = UUID(header)
+    except ValueError:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail={"reason": "invalid_request", "message": f"{TRIGGER_TI_ID_HEADER} must be a UUID"},
+        )
+
+    if request.method in ("GET", "HEAD"):
+        permitted = (
+            session.scalar(
+                select(TaskInstance.working_set)
+                .where(TaskInstance.id == claimed)
+                .execution_options(include_all_attempts=True)
+            )
+            is True
+        )
+        denial = "Task instance is not a live attempt"
+    else:
+        permitted = (
+            session.scalar(
+                select(TaskInstance.id)
+                .join(Trigger, Trigger.id == TaskInstance.trigger_id)
+                .where(
+                    TaskInstance.id == claimed,
+                    Trigger.triggerer_id == triggerer_id,
+                    TaskInstance.state == TaskInstanceState.DEFERRED,
+                )
+            )
+            is not None
+        )
+        denial = "Task instance is not deferred on a trigger assigned to this triggerer"
+    if not permitted:
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail={"reason": "access_denied", "message": denial},
+        )
+    return token.model_copy(update={"id": claimed})
+
+
+_TriggerAwareToken = Depends(_resolve_trigger_token)
+
+
+def _resolve_trigger_ti_token(token: ExecutionToken = _TriggerAwareToken) -> TIToken:
+    return require_task_token(token)
+
+
+_TriggerAwareTIToken = Depends(_resolve_trigger_ti_token)
 
 
 @router.get(
@@ -297,7 +358,7 @@ def get_mapped_xcom_by_index(
     session: SessionDep,
     dag_bag: DagBagDep,
     previous_iteration: bool = False,
-    token: ExecutionToken = CurrentExecutionToken,
+    token: ExecutionToken = _TriggerAwareToken,
 ) -> XComSequenceIndexResponse:
     xcom_read = _build_xcom_read(
         dag_id=dag_id,
@@ -364,7 +425,7 @@ def get_mapped_xcom_by_slice(
     params: Annotated[GetXComSliceFilterParams, Query()],
     session: SessionDep,
     dag_bag: DagBagDep,
-    token: ExecutionToken = CurrentExecutionToken,
+    token: ExecutionToken = _TriggerAwareToken,
 ) -> XComSequenceSliceResponse:
     xcom_read = _build_xcom_read(
         dag_id=dag_id,
@@ -448,7 +509,10 @@ def get_mapped_xcom_by_slice(
     dependencies=[ExecutionOrDagParseToken, DagInGrantedBundle],
     responses={
         **create_openapi_http_exception_doc(
-            [(status.HTTP_400_BAD_REQUEST, "map_index cannot be specified in a HEAD request")]
+            [
+                (status.HTTP_400_BAD_REQUEST, "map_index cannot be specified in a HEAD request"),
+                (status.HTTP_409_CONFLICT, "Several live task instances produce the requested value"),
+            ]
         ),
         status.HTTP_200_OK: {
             "description": "Metadata about the number of matching XCom values",
@@ -463,10 +527,16 @@ def get_mapped_xcom_by_slice(
     description="Returns the count of mapped XCom values found in the `Content-Range` response header",
 )
 def head_xcom(
+    dag_id: str,
+    run_id: str,
+    task_id: str,
+    key: str,
     response: Response,
     session: SessionDep,
-    xcom_query: Annotated[Select, Depends(xcom_query)],
+    dag_bag: DagBagDep,
     map_index: Annotated[int | None, Query()] = None,
+    previous_iteration: bool = False,
+    token: ExecutionToken = _TriggerAwareToken,
 ) -> None:
     """Get the count of XComs from database - not other XCom Backends."""
     if map_index is not None:
@@ -475,6 +545,17 @@ def head_xcom(
             detail={"reason": "invalid_request", "message": "Cannot specify map_index in a HEAD request"},
         )
 
+    xcom_query = _build_xcom_read(
+        dag_id=dag_id,
+        run_id=run_id,
+        task_id=task_id,
+        key=key,
+        session=session,
+        dag_bag=dag_bag,
+        token=token,
+        previous_iteration=previous_iteration,
+        sequence=True,
+    ).statement
     count = get_query_count(xcom_query, session=session)
     # Tell the caller how many items in this query. We define a custom range unit (HTTP spec only defines
     # "bytes" but we can add our own)
@@ -509,7 +590,7 @@ def get_xcom(
     session: SessionDep,
     params: Annotated[GetXcomFilterParams, Query()],
     dag_bag: DagBagDep,
-    token: ExecutionToken = CurrentExecutionToken,
+    token: ExecutionToken = _TriggerAwareToken,
 ) -> XComResponse:
     """Get an Airflow XCom from database - not other XCom Backends."""
     owner = None
@@ -643,7 +724,7 @@ def set_xcom(
     mapped_length: Annotated[
         int | None, Query(ge=0, description="Number of mapped tasks this value expands into")
     ] = None,
-    token=CurrentTIToken,
+    token=_TriggerAwareTIToken,
 ):
     """Set an Airflow XCom."""
     from airflow.configuration import conf
@@ -773,7 +854,7 @@ def delete_xcom(
     task_id: str,
     key: Annotated[str, Path(min_length=1)],
     map_index: Annotated[int, Query()] = -1,
-    token=CurrentTIToken,
+    token=_TriggerAwareTIToken,
 ):
     """Delete a single XCom Value."""
     owner = _find_writer_id(

airflow-core/src/airflow/api_fastapi/execution_api/security.py

@@ -141,6 +141,7 @@
 _REQUEST_SCOPE_LIVE_ATTEMPT_KEY = "live_attempt_checked"
 _REQUEST_SCOPE_JOB_KEY = "dag_processor_job_id"
 _IN_PROCESS_NON_TI_CALLER = "airflow_in_process_non_ti_caller"
+_IN_PROCESS_TRIGGERER_ID = "airflow_in_process_triggerer_id"
 _SKIP_AUTO_TI_ATTEMPT_LIVE = "skip_auto_ti_attempt_live"
 _P = ParamSpec("_P")
 _R = TypeVar("_R")