Add a POST /xcoms/{dag_id}/{run_id}/{task_id}/keys execution API endpoint

XComIterable, the result of an iterated task, stores one XCom per iteration
under distinct keys (return_value_0, return_value_1, ...) of the same task
instance. The existing slice endpoint for a mapped task's XComs ranges over
map_index for a single key, the inverse shape, so iterating or slicing an
XComIterable cost one GET per value.

The new endpoint takes a list of keys and returns their values in that order,
None for a key without an XCom, filtered by map_index (-1 by default), in one
database query. Around it:

* XComKeysRequest body model, and an AddXComKeysEndpoint execution API version
  change so older clients keep their contract.
* has_xcom_access reads the optional key from the request's path parameters,
  so the endpoint needs no separate router or dependency.
* GetXComByKeys supervisor message, XComOperations.get_by_keys client method,
  handle_get_xcom_by_keys handler registered with the task supervisor and the
  DAG processor, an AddGetXComByKeys supervisor schema version entry and the
  regenerated schema snapshot.
* BaseXCom.get_by_keys sends the message and deserializes the values, and
  XComIterable iterates and slices through it: one request instead of one per
  value. A single index stays one XCom read.

Stacked on the task iteration PR, which owns XComIterable; only this commit
belongs to the endpoint. Squashed from the earlier history of this branch and
re-based on that PR, with XComIterable's batched reads moved to its current
home in bases/xcom.py.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Author: dabla

airflow-core/src/airflow/api_fastapi/execution_api/datamodels/xcom.py

@@ -40,3 +40,9 @@ class XComSequenceSliceResponse(RootModel):
     """XCom schema with minimal structure for slice-based access."""
 
     root: list[JsonValue]
+
+
+class XComKeysRequest(BaseModel):
+    """Request body for fetching multiple XCom values by key list in a single query."""
+
+    keys: list[str]

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

@@ -30,6 +30,7 @@
 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.xcom import (
+    XComKeysRequest,
     XComResponse,
     XComSequenceIndexResponse,
     XComSequenceSliceResponse,
@@ -45,7 +46,6 @@ def has_xcom_access(
     dag_id: str,
     run_id: str,
     task_id: str,
-    xcom_key: Annotated[str, Path(alias="key", min_length=1)],
     request: Request,
     session: SessionDep,
     token=CurrentTIToken,
@@ -71,6 +71,7 @@ def has_xcom_access(
     from airflow.configuration import conf
 
     write = request.method not in {"GET", "HEAD", "OPTIONS"}
+    xcom_key = request.path_params.get("key")
 
     log.debug(
         "Checking %s XCom access for task instance '%s' to XCom '%s' on dag '%s'",
@@ -116,6 +117,7 @@ def has_xcom_access(
     dependencies=[Depends(has_xcom_access)],
 )
 
+
 log = logging.getLogger(__name__)
 
 
@@ -405,6 +407,38 @@ def get_xcom(
     return XComResponse(key=key, value=(result[0] if isinstance(result, tuple) else result).value)
 
 
+@router.post(
+    "/{dag_id}/{run_id}/{task_id}/keys",
+    summary="Get multiple XCom values by keys",
+    description=(
+        "Fetch multiple XCom values by key list in a single database query. "
+        "Optimised for XComIterable iteration, reducing N round-trips to one."
+    ),
+)
+def get_xcom_by_keys(
+    dag_id: str,
+    run_id: str,
+    task_id: str,
+    request_body: XComKeysRequest,
+    session: SessionDep,
+    map_index: Annotated[int, Query()] = -1,
+) -> XComSequenceSliceResponse:
+    """Fetch multiple XCom values by different keys in a single database query."""
+    key_list = request_body.keys
+    if not key_list:
+        return XComSequenceSliceResponse([])
+
+    query = select(XComModel.key, XComModel.value).where(
+        XComModel.dag_id == dag_id,
+        XComModel.run_id == run_id,
+        XComModel.task_id == task_id,
+        XComModel.map_index == map_index,
+        XComModel.key.in_(key_list),
+    )
+    rows = {row.key: row.value for row in session.execute(query)}
+    return XComSequenceSliceResponse([rows.get(key) for key in key_list])
+
+
 # TODO: once we have JWT tokens, then remove dag_id/run_id/task_id from the URL and just use the info in
 # the token
 @router.post(

airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py

@@ -59,6 +59,7 @@
     AddMultiTeamToTIRunContext,
     AddStoppedTaskReport,
     AddTerminalStateRetryReasonField,
+    AddXComKeysEndpoint,
     IdentifyArchivedTaskStateUpdates,
 )
 
@@ -73,6 +74,7 @@
         AddMultiTeamToTIRunContext,
         AddStoppedTaskReport,
         IdentifyArchivedTaskStateUpdates,
+        AddXComKeysEndpoint,
     ),
     Version(
         "2026-06-30",

airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py

@@ -108,3 +108,13 @@ class AddMultiTeamToTIRunContext(VersionChange):
     def remove_multi_team_field(response: ResponseInfo) -> None:  # type: ignore[misc]
         """Strip ``multi_team`` from the run context for older clients."""
         response.body.pop("multi_team", None)
+
+
+class AddXComKeysEndpoint(VersionChange):
+    """Add the ``xcoms/{dag_id}/{run_id}/{task_id}/keys`` endpoint that fetches several XCom values by key in one query."""
+
+    description = __doc__
+
+    instructions_to_migrate_to_previous_version = (
+        endpoint("/xcoms/{dag_id}/{run_id}/{task_id}/keys", ["POST"]).didnt_exist,
+    )

airflow-core/src/airflow/dag_processing/processor.py

@@ -55,6 +55,7 @@
     GetVariable,
     GetVariableKeys,
     GetXCom,
+    GetXComByKeys,
     GetXComCount,
     GetXComSequenceItem,
     GetXComSequenceSlice,
@@ -151,6 +152,7 @@ class DagFileParsingResult(BaseModel):
     | GetXCom
     | GetXComCount
     | GetXComSequenceItem
+    | GetXComByKeys
     | GetXComSequenceSlice
     | MaskSecret,
     Field(discriminator="type"),
@@ -644,6 +646,7 @@ def _handle_parsing_result(
         GetVariable,
         GetVariableKeys,
         GetXCom,
+        GetXComByKeys,
         GetXComCount,
         GetXComSequenceItem,
         GetXComSequenceSlice,

airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_xcoms.py

@@ -79,11 +79,10 @@ def _(
         dag_id: str = Path(),
         run_id: str = Path(),
         task_id: str = Path(),
-        xcom_key: str = Path(alias="key"),
         token=CurrentTIToken,
     ):
         with create_session() as session:
-            has_xcom_access(dag_id, run_id, task_id, xcom_key, request, session, token)
+            has_xcom_access(dag_id, run_id, task_id, request, session, token)
         raise HTTPException(
             status_code=status.HTTP_403_FORBIDDEN,
             detail={
@@ -875,3 +874,138 @@ def test_teamless_requester_scoping(self, client, session, dag_maker):
 
         assert forbidden.status_code == 403, forbidden.json()
         assert allowed.status_code == 200, allowed.json()
+
+
+class TestGetXComByKeys:
+    """Tests for the POST /{dag_id}/{run_id}/{task_id}/keys batched-fetch endpoint."""
+
+    def _insert_xcom(self, session, ti, key, value):
+        session.add(
+            XComModel(
+                key=key,
+                value=value,
+                dag_run_id=ti.dag_run.id,
+                run_id=ti.run_id,
+                task_id=ti.task_id,
+                dag_id=ti.dag_id,
+            )
+        )
+
+    def test_returns_values_in_request_key_order(self, client, create_task_instance, session):
+        """Values are returned in the same order as the requested keys."""
+        ti = create_task_instance()
+        self._insert_xcom(session, ti, "return_value_0", "alpha")
+        self._insert_xcom(session, ti, "return_value_1", "beta")
+        self._insert_xcom(session, ti, "return_value_2", "gamma")
+        session.commit()
+
+        response = client.post(
+            f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/keys",
+            json={"keys": ["return_value_2", "return_value_0", "return_value_1"]},
+        )
+
+        assert response.status_code == 200
+        assert response.json() == ["gamma", "alpha", "beta"]
+
+    def test_missing_key_returned_as_none(self, client, create_task_instance, session):
+        """Keys not found in the database are returned as null (None)."""
+        ti = create_task_instance()
+        self._insert_xcom(session, ti, "return_value_0", "exists")
+        session.commit()
+
+        response = client.post(
+            f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/keys",
+            json={"keys": ["return_value_0", "return_value_1"]},
+        )
+
+        assert response.status_code == 200
+        assert response.json() == ["exists", None]
+
+    def test_empty_key_list_returns_empty_list(self, client, create_task_instance, session):
+        ti = create_task_instance()
+        session.commit()
+
+        response = client.post(
+            f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/keys",
+            json={"keys": []},
+        )
+
+        assert response.status_code == 200
+        assert response.json() == []
+
+    def test_map_index_filter(self, client, dag_maker, session):
+        """Only XCom rows matching the requested map_index are returned."""
+
+        class MyOperator(EmptyOperator):
+            def __init__(self, *, x, **kwargs):
+                super().__init__(**kwargs)
+                self.x = x
+
+        with dag_maker(dag_id="dag"):
+            MyOperator.partial(task_id="task").expand(x=["a", "b"])
+        dag_run = dag_maker.create_dagrun(run_id="run")
+        tis = {ti.map_index: ti for ti in dag_run.task_instances}
+
+        # Insert XComs for map_index=0 and map_index=1 (both have real TI rows).
+        # map_index=-1 has no TI row for a mapped task, so it is not inserted.
+        session.add(
+            XComModel(
+                key="return_value_0",
+                value="for_index_0",
+                dag_run_id=dag_run.id,
+                run_id=dag_run.run_id,
+                task_id=tis[0].task_id,
+                dag_id=tis[0].dag_id,
+                map_index=0,
+            )
+        )
+        session.add(
+            XComModel(
+                key="return_value_0",
+                value="for_index_1",
+                dag_run_id=dag_run.id,
+                run_id=dag_run.run_id,
+                task_id=tis[1].task_id,
+                dag_id=tis[1].dag_id,
+                map_index=1,
+            )
+        )
+        session.commit()
+
+        dag_id = tis[0].dag_id
+        task_id = tis[0].task_id
+
+        # Default map_index=-1: no rows stored there, so null is returned.
+        response_default = client.post(
+            f"/execution/xcoms/{dag_id}/{dag_run.run_id}/{task_id}/keys",
+            json={"keys": ["return_value_0"]},
+        )
+        response_index_0 = client.post(
+            f"/execution/xcoms/{dag_id}/{dag_run.run_id}/{task_id}/keys",
+            params={"map_index": 0},
+            json={"keys": ["return_value_0"]},
+        )
+        response_index_1 = client.post(
+            f"/execution/xcoms/{dag_id}/{dag_run.run_id}/{task_id}/keys",
+            params={"map_index": 1},
+            json={"keys": ["return_value_0"]},
+        )
+
+        assert response_default.json() == [None]
+        assert response_index_0.json() == ["for_index_0"]
+        assert response_index_1.json() == ["for_index_1"]
+
+    def test_does_not_return_other_tasks_xcoms(self, client, create_task_instance, session):
+        """Keys from a different task_id are not returned."""
+        ti = create_task_instance()
+        other_ti = create_task_instance(dag_id="other_dag", task_id="other_task")
+        self._insert_xcom(session, other_ti, "return_value_0", "not_mine")
+        session.commit()
+
+        response = client.post(
+            f"/execution/xcoms/{ti.dag_id}/{ti.run_id}/{ti.task_id}/keys",
+            json={"keys": ["return_value_0"]},
+        )
+
+        assert response.status_code == 200
+        assert response.json() == [None]

airflow-core/tests/unit/jobs/test_triggerer_job.py

@@ -3054,6 +3054,7 @@ def get_type_names(union_type):
             "GetPreviousDagRun",
             "GetTaskBreadcrumbs",
             "GetTaskRescheduleStartDate",
+            "GetXComByKeys",
             "GetXComCount",
             "GetXComSequenceItem",
             "GetXComSequenceSlice",

task-sdk/src/airflow/sdk/api/client.py

@@ -746,6 +746,22 @@ def get_sequence_slice(
         resp = self.client.get(f"xcoms/{dag_id}/{run_id}/{task_id}/{key}/slice", params=params)
         return XComSequenceSliceResponse.model_validate_json(resp.read())
 
+    def get_by_keys(
+        self,
+        dag_id: str,
+        run_id: str,
+        task_id: str,
+        keys: list[str],
+        map_index: int = -1,
+    ) -> XComSequenceSliceResponse:
+        """Fetch multiple XCom values by key list in a single round-trip."""
+        resp = self.client.post(
+            f"xcoms/{dag_id}/{run_id}/{task_id}/keys",
+            params={"map_index": map_index} if map_index >= 0 else {},
+            json={"keys": keys},
+        )
+        return XComSequenceSliceResponse.model_validate_json(resp.read())
+
 
 class TaskStateStoreOperations:
     __slots__ = ("client",)

task-sdk/src/airflow/sdk/api/datamodels/_generated.py

@@ -529,6 +529,14 @@ class VariableResponse(BaseModel):
     value: Annotated[str | None, Field(title="Value")]
 
 
+class XComKeysRequest(BaseModel):
+    """
+    Request body for fetching multiple XCom values by key list in a single query.
+    """
+
+    keys: Annotated[list[str], Field(title="Keys")]
+
+
 class XComResponse(BaseModel):
     """
     XCom schema for responses with fields that are needed for Runtime.