Batch XCom lookups when resolving a mapped task's expand() arguments

[ColtenOuO]

Summary

DictOfListsExpandInput.resolve() calls .resolve(context) on every XComArg kwarg of a mapped task's .expand() one at a time, and for the common case (a non-mapped upstream) each of those turns into its own ti.xcom_pull() call — one Execution API round trip per kwarg. A mapped task with N XComArg kwargs makes N sequential round trips today. Both call sites in expandinput.py already carried a TODO admitting this. This PR adds a batch XCom lookup endpoint and wires the common case through it, so N kwargs collapse into one round trip instead of N.

Change

New batch endpoint (POST /execution/xcoms/{dag_id}/{run_id}/batch) — looks up multiple (task_id, key, map_index) triples in a single request and a single DB query (tuple_(...).in_(...)), scoped to one dag run. A missing item is reported as found=False instead of failing the whole batch, since a partial miss (an upstream that hasn't pushed yet) is an expected batch outcome, not an error. Access control reuses the existing multi-team dag-level check. The request body is capped at MAX_XCOM_BATCH_ITEMS (1000) items — each item expands to 3 SQL bind parameters, so an unbounded batch from an authenticated task could otherwise push a query past Postgres/MySQL/SQLite bind-parameter limits or spend unbounded API-server work on a single request; an oversized request is rejected with 422 before touching the database.

Task SDK plumbing — XComOperations.get_batch() on the generated client, a GetXComBatch/XComBatchResult message pair on the supervisor comms protocol, and the matching supervisor-side dispatch handler, following the exact same shape as the existing single-item GetXCom path.

Wiring in DictOfListsExpandInput.resolve() — before resolving each kwarg individually, it now collects the eligible ones (plain XComArg, non-mapped upstream, not inside a mapped task group — the same condition that already made the per-item path pull with map_indexes=None) and resolves them with a single batched call. Everything else (mapped-upstream LazyXComSequence, MapXComArg/ZipXComArg/ConcatXComArg) is unchanged and still resolves per-item; batching those is left for a follow-up, since they need additional design work (lazy count/slice semantics, composite-tree walking) with smaller marginal payoff.

A custom XCom backend (core.xcom_backend) also disables batching entirely, falling back to the existing per-item path — a custom backend's get_one/get_all may bypass the batch endpoint's semantics, so this sidesteps that risk rather than trying to make batching custom-backend-aware. If the batch response is ever neither a valid result nor the "not supported" error (a protocol contract violation, not an expected outcome), it raises a clear TypeError instead of silently resolving to wrong values.

Backward compatibility

This PR touches both sides of the wire (API server and Task SDK client), so all four version-skew combinations that can show up during a rolling upgrade are handled explicitly:

• New server + old client — the batch endpoint is registered on a Cadwyn migration (AddXComBatchEndpoint, endpoint(...).didnt_exist for versions before it). An old Task SDK is pinned to an older API version, so the endpoint doesn't exist in the contract it negotiates against; it never calls it and keeps using the pre-existing per-item GetXCom path unchanged.
• Old server + new client — a new Task SDK's batch call gets a 404 from a server that doesn't have the endpoint yet. The client surfaces this as ErrorType.XCOM_BATCH_NOT_SUPPORTED, and DictOfListsExpandInput.resolve() catches exactly that and falls back to resolving the group per-item. The task still succeeds; it just doesn't get the speed-up until the server catches up.
• New server + new client — batching applies as designed.
• Old server + old client — unaffected; the batch endpoint doesn't exist on either side.

Everything added (the route, the GetXComBatch/XComBatchResult comms messages, XComOperations.get_batch()) is strictly additive — no existing endpoint, message shape, or client method changes shape — so the two sides can roll out in either order without a hard dependency.

Before / After

Before: N XComArg kwargs on one .expand() call = N sequential Execution API round trips (worker → supervisor → API server → DB), one per kwarg.

After: the eligible kwargs resolve in a single round trip regardless of N.

Measured locally (verification script, not included in this PR) against a real FastAPI execution-API app and a real breeze Postgres database (in-process ASGI transport — this excludes only the raw TCP hop a separate worker process would add on top, so these numbers are a conservative floor, not an inflated one):

N kwargs	current (s)	batched (s)	speedup
2	0.054	0.023	2.3x
5	0.086	0.025	3.4x
10	0.176	0.026	6.7x
20	0.337	0.026	12.8x
50	0.890	0.028	32.2x

The current path scales linearly with N (~17ms/call); the batched path stays flat (~25ms total) regardless of N.

Testing

• Full task-sdk suite: 2798 passed, 7 skipped (pre-existing/environment, unrelated), 0 failed.
• New execution-API route tests: multiple items found, partial miss (found=False), empty batch, over-the-cap request rejected with 422, a batch item with an explicit non-default map_index, and multi-team access-control (same-team allowed, cross-team forbidden) for the batch endpoint.
• New Task SDK client tests: successful batch call, and the 404 → ErrorType.XCOM_BATCH_NOT_SUPPORTED fallback signal.
• New supervisor dispatch test (added to the existing REQUEST_TEST_CASES table) covering GetXComBatch → XComBatchResult.
• New expandinput/xcom_arg tests: N kwargs collapse into exactly one GetXComBatch call; correct per-kwarg value distribution; old-server fallback (404 → per-item calls, correct values); custom-backend skip (batching never attempted); mixed eligibility (a mapped-task-group XComArg alongside plain ones — only the plain ones batch); the found=False path (a missing return_value XCom resolves to None, a missing non-return_value key raises XComNotFound, matching the per-item xcom_pull path exactly). Mocks in the new SDK-level tests are autospec'd against BaseOperator/RuntimeTaskInstance rather than bare MagicMock, so a signature or attribute drift on those classes would fail the test instead of passing silently.
• Fixed 3 pre-existing tests (test_map_cross_product, test_map_product_same, test_mapped_render_template_fields_validating_operator) whose hand-rolled comms mocks didn't know about the new GetXComBatch message — a real regression this PR would otherwise have introduced, caught by running the full suite rather than only the new tests.
• mypy clean on both task-sdk and the touched airflow-core files.
• Both Cadwyn version-check prek hooks (check-execution-api-versions, check-supervisor-schemas-versions) pass.

Scope / follow-ups

Batching for a mapped-upstream LazyXComSequence (the GetXComCount length lookup) and for MapXComArg/ZipXComArg/ConcatXComArg composite resolution is intentionally out of scope here — those need their own design work and have smaller marginal payoff for typical DAG shapes (most .expand() calls use a handful of plain kwargs, not many mapped-upstream ones). Tracked as a natural next step, not filed as a separate issue since no workaround is being shipped here.

---

Was generative AI tooling used to co-author this PR?

• Yes — Claude Code (Sonnet 5)

Generated-by: Claude Code (Sonnet 5) following the guidelines

[ColtenOuO]

self review and check all 22 files.

[ColtenOuO]

Updared the PR's description (added Backward compatibility)

[ColtenOuO]

Fixed DAG File Processor forward