Support dynamic task mapping on stub tasks

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

@@ -0,0 +1,104 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Positional-argument binding spec for stub (foreign-runtime) tasks.
+
+Captured at parse time from the ``@task.stub`` TaskFlow call, stored in the serialized
+Dag, and delivered to the lang-SDK runtime via ``TIRunContext.arg_bindings``.
+"""
+
+from __future__ import annotations
+
+from functools import cache
+from typing import Annotated, Literal
+
+from pydantic import Field, JsonValue, TypeAdapter
+from typing_extensions import TypeAliasType
+
+from airflow.api_fastapi.core_api.base import BaseModel
+
+# A named, titled alias (like TaskArgBinding below) kept as free-form JSON rather than a
+# typed model, so unknown JSON-schema keywords survive re-serialization along the way.
+ArgValueSchema = TypeAliasType(
+    "ArgValueSchema", Annotated[dict[str, JsonValue], Field(title="ArgValueSchema")]
+)
+"""JSON-schema fragment constraining the value a stub-task argument binds to; generated
+by pydantic from the stub annotation, carried verbatim, unknown keywords ignored."""
+
+
+class XComArgBinding(BaseModel):
+    """One positional stub-task argument pulled from an upstream task's XCom."""
+
+    # No default: it would drop ``kind`` from ``required``, and the generated task-sdk
+    # client then types it ``Literal | None``, invalid as a tagged-union discriminator.
+    kind: Literal["xcom"]
+
+    name: str
+    """The stub function's parameter name this binding fills, in declaration order."""
+
+    value_schema: ArgValueSchema | None = None
+    """Schema fragment from the stub function's annotation; omitted when unconstrained."""
+
+    task_id: str
+    """Upstream task id whose ``return_value`` XCom is pulled."""
+
+    map_index: int = -1
+    """Map index of the upstream XCom row to pull; -1 is the unmapped row."""
+
+    element_index: int | None = None
+    """When set, the pulled value is a sequence and this binding takes the element at
+    this index (the stub was expanded over an unmapped upstream's output). The lang-SDK
+    side will GET the single row ``(task_id, map_index=-1)``, decode it, and take
+    ``value[element_index]``."""
+
+
+class LiteralArgBinding(BaseModel):
+    """One positional stub-task argument carrying an inline literal from the Dag file."""
+
+    kind: Literal["literal"]
+    """No default, for the same generated-client reason as ``XComArgBinding.kind``."""
+
+    name: str
+    """The stub function's parameter name this binding fills, in declaration order."""
+
+    value_schema: ArgValueSchema | None = None
+    """Schema fragment from the stub function's annotation; omitted when unconstrained."""
+
+    value: JsonValue | None = None
+    """The literal value from the Dag file."""
+
+    from_default: bool = False
+    """True when the value was filled from the stub signature's default rather than passed in the call."""
+
+
+# A named alias with an explicit title so the union lands in every schema as its own
+# named definition, which the supervisor-schema dump dedups with its task-sdk twin by title.
+TaskArgBinding = TypeAliasType(
+    "TaskArgBinding",
+    Annotated[XComArgBinding | LiteralArgBinding, Field(discriminator="kind", title="TaskArgBinding")],
+)
+"""One positional argument of a stub (foreign-runtime) task, in declaration order."""
+
+
+@cache
+def get_arg_bindings_adapter() -> TypeAdapter[list[TaskArgBinding]]:
+    """
+    Build (lazily, then cache) the adapter validating serialized dicts into ``TaskArgBinding``.
+
+    Only the stub-task path in the execution API needs it, so regular runs never pay for it.
+    """
+    return TypeAdapter(list[TaskArgBinding])

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

@@ -36,6 +36,7 @@
 from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel
 from airflow.api_fastapi.execution_api.datamodels.asset import AssetProfile
 from airflow.api_fastapi.execution_api.datamodels.connection import ConnectionResponse
+from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import TaskArgBinding
 from airflow.api_fastapi.execution_api.datamodels.variable import VariableResponse
 from airflow.utils.state import (
     DagRunState,
@@ -435,6 +436,13 @@ class TIRunContext(BaseModel):
     always reflects when the task *first* started, not when it was rescheduled/resumed.
     """
 
+    arg_bindings: list[TaskArgBinding] | None = None
+    """
+    Ordered positional-argument binding spec for stub (foreign-runtime) tasks.
+
+    ``None`` for regular tasks and for stub tasks that declare no parameters.
+    """
+
 
 class PrevSuccessfulDagRunResponse(BaseModel):
     """Schema for response with previous successful DagRun information for Task Template Context."""

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

@@ -32,7 +32,7 @@
 from opentelemetry import trace
 from opentelemetry.trace import StatusCode
 from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
-from pydantic import JsonValue
+from pydantic import JsonValue, ValidationError
 from sqlalchemy import and_, func, or_, tuple_, update
 from sqlalchemy.engine import CursorResult
 from sqlalchemy.exc import DataError, NoResultFound, SQLAlchemyError
@@ -49,6 +49,7 @@
 from airflow.api_fastapi.common.types import UtcDateTime
 from airflow.api_fastapi.compat import HTTP_422_UNPROCESSABLE_CONTENT
 from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc
+from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import get_arg_bindings_adapter
 from airflow.api_fastapi.execution_api.datamodels.taskinstance import (
     InactiveAssetsResponse,
     PreviousTIResponse,
@@ -75,6 +76,8 @@
     get_team_name_for_ti,
     require_auth,
 )
+from airflow.api_fastapi.execution_api.services.task_instances import STUB_TASK_TYPE, get_arg_bindings
+from airflow.api_fastapi.execution_api.versions import bundle
 from airflow.configuration import conf
 from airflow.exceptions import InvalidPartitionKeyError, TaskNotFound
 from airflow.models.asset import AssetActive
@@ -110,6 +113,22 @@
 log = structlog.get_logger(__name__)
 tracer = trace.get_tracer(__name__)
 
+# The first execution API version whose TIRunContext carries ``arg_bindings``.
+ARG_BINDINGS_API_VERSION = "2026-10-30"
+
+
+def _client_supports_arg_bindings() -> bool:
+    """
+    Whether the request's negotiated API version can receive ``arg_bindings``.
+
+    Clients on older versions never see the field (the version migration strips it from
+    the response), so the derivation -- and the structured failures it raises for
+    undeliverable mapped-stub specs -- must not run for them: a stub Dag that ran before
+    arg bindings existed keeps running against those clients.
+    """
+    version = bundle.api_version_var.get(None)
+    return version is None or str(version) >= ARG_BINDINGS_API_VERSION
+
 
 @ti_id_router.patch(
     "/{task_instance_id}/run",
@@ -163,6 +182,8 @@ def ti_run(
             TI.hostname,
             TI.unixname,
             TI.pid,
+            TI.operator,
+            TI.dag_version_id,
             # This selects the raw JSON value, bypassing the deserialization -- we want that to happen on the
             # client
             column("next_kwargs", JSON),
@@ -310,6 +331,30 @@ def ti_run(
             should_retry=_is_eligible_to_retry(previous_state, ti.try_number, ti.max_tries),
         )
 
+        # Only set for stub (foreign-runtime) tasks with a captured TaskFlow arg
+        # spec; the route excludes unset fields, keeping regular responses lean.
+        if (
+            ti.operator == STUB_TASK_TYPE
+            and _client_supports_arg_bindings()
+            and (arg_bindings := get_arg_bindings(dag_bag, ti, session=session))
+        ):
+            try:
+                context.arg_bindings = get_arg_bindings_adapter().validate_python(arg_bindings)
+            except ValidationError:
+                log.exception(
+                    "Serialized arg_bindings spec failed validation",
+                    dag_id=ti.dag_id,
+                    task_id=ti.task_id,
+                    dag_version_id=ti.dag_version_id,
+                )
+                raise HTTPException(
+                    status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                    detail={
+                        "reason": "invalid_arg_bindings",
+                        "message": "The serialized TaskFlow arg spec for this stub task is not valid.",
+                    },
+                )
+
         # Only set if they are non-null
         if ti.next_method:
             context.next_method = ti.next_method

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

@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.