Support TaskFlow call syntax on stub tasks for the Lang SDK

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

@@ -0,0 +1,93 @@
+# 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 _ArgBindingBase(BaseModel):
+    """Fields every :class:`TaskArgBinding` variant carries, regardless of ``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."""
+
+
+class XComArgBinding(_ArgBindingBase):
+    """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"]
+
+    task_id: str
+    """Upstream task id whose ``return_value`` XCom is pulled."""
+
+
+class LiteralArgBinding(_ArgBindingBase):
+    """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``."""
+
+    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,11 @@
     get_team_name_for_ti,
     require_auth,
 )
+from airflow.api_fastapi.execution_api.services.task_instances import (
+    LANG_SDK_OPERATORS,
+    client_supports_arg_bindings,
+    get_arg_bindings,
+)
 from airflow.configuration import conf
 from airflow.exceptions import InvalidPartitionKeyError, TaskNotFound
 from airflow.models.asset import AssetActive
@@ -163,6 +169,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 +318,30 @@ def ti_run(
             should_retry=_is_eligible_to_retry(previous_state, ti.try_number, ti.max_tries),
         )
 
+        # Only set for lang-SDK (foreign-runtime) tasks with a captured TaskFlow arg
+        # spec; the route excludes unset fields, keeping regular responses lean.
+        if (
+            ti.operator in LANG_SDK_OPERATORS
+            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.

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

@@ -0,0 +1,67 @@
+# 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.
+"""Business logic backing the task-instance execution routes."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from sqlalchemy.orm import Session
+
+    from airflow.models.dagbag import DBDagBag
+
+# Task types (``TaskInstance.operator``, the operator class name) whose tasks carry a
+# lang-SDK ``arg_bindings`` spec. Used to gate the serialized-Dag lookup so regular tasks
+# never pay for it. The gate matches exact class names; a new lang-SDK operator adds its
+# name here.
+LANG_SDK_OPERATORS = frozenset({"_StubOperator"})
+
+
+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 must not run for them.
+
+    Rather than comparing the negotiated version by date, we check the
+    ``VersionChangeWithSideEffects`` subclass's ``is_applied`` flag; see
+    https://docs.cadwyn.dev/concepts/version_changes/#version-changes-with-side-effects
+    """
+    # Imported locally: the versions package transitively imports the routes, which import
+    # this module, so a top-level import here would be circular.
+    from airflow.api_fastapi.execution_api.versions.v2026_10_30 import AddArgBindingsToTIRunContext
+
+    return AddArgBindingsToTIRunContext.is_applied
+
+
+def get_arg_bindings(dag_bag: DBDagBag, ti: Any, *, session: Session) -> list | None:
+    """
+    Extract the stub task's TaskFlow arg spec from its Dag version.
+
+    Mapped (``.expand()``) stubs never capture a parse-time spec, so they resolve to
+    ``None`` here and keep the legacy ignored-args behavior; per-map-index delivery
+    lands in a follow-up.
+    """
+    if ti.dag_version_id is None:
+        return None
+    if (dag := dag_bag.get_dag(ti.dag_version_id, session=session)) is None:
+        return None
+    if (task := dag.task_dict.get(ti.task_id)) is None:
+        return None
+    return getattr(task, "_arg_bindings", None)

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

@@ -51,9 +51,11 @@
     AddTeamNameField,
     AddVariableKeysEndpoint,
 )
+from airflow.api_fastapi.execution_api.versions.v2026_10_30 import AddArgBindingsToTIRunContext
 
 bundle = VersionBundle(
     HeadVersion(),
+    Version("2026-10-30", AddArgBindingsToTIRunContext),
     Version(
         "2026-06-30",
         AddVariableKeysEndpoint,

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

@@ -0,0 +1,42 @@
+# 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.
+
+from __future__ import annotations
+
+from cadwyn import (
+    ResponseInfo,
+    VersionChangeWithSideEffects,
+    convert_response_to_previous_version_for,
+    schema,
+)
+
+from airflow.api_fastapi.execution_api.datamodels.taskinstance import TIRunContext
+
+
+class AddArgBindingsToTIRunContext(VersionChangeWithSideEffects):
+    """Add the ``arg_bindings`` argument-binding spec for stub (foreign-runtime) tasks."""
+
+    description = __doc__
+
+    # A side-effect change, not just a schema one, so ti_run can gate the server-side spec
+    # derivation on ``is_applied``: clients older than this version never receive the field.
+    instructions_to_migrate_to_previous_version = (schema(TIRunContext).field("arg_bindings").didnt_exist,)
+
+    @convert_response_to_previous_version_for(TIRunContext)  # type: ignore[arg-type]
+    def remove_arg_bindings_field(response: ResponseInfo) -> None:  # type: ignore[misc]
+        """Strip ``arg_bindings`` from the run context for older clients."""
+        response.body.pop("arg_bindings", None)

airflow-core/src/airflow/serialization/schema.json

@@ -142,6 +142,48 @@
       "description": "A python dictionary containing values of any type",
       "type": "object"
     },
+    "typed_dict": {
+      "type": "object",
+      "properties": {
+        "__type": {
+          "type": "string",
+          "const": "dict"
+        },
+        "__var": { "$ref": "#/definitions/dict" }
+      },
+      "required": [
+        "__type",
+        "__var"
+      ],
+      "additionalProperties": false
+    },
+    "arg_binding": {
+      "$comment": "One captured TaskFlow call argument of a @task.stub task, in dict-encoded form. The inner object stays open so future binding fields keep validating on older cores",
+      "type": "object",
+      "properties": {
+        "__type": {
+          "type": "string",
+          "const": "dict"
+        },
+        "__var": {
+          "type": "object",
+          "properties": {
+            "name": { "type": "string" },
+            "kind": { "type": "string", "enum": [ "xcom", "literal" ] },
+            "value_schema": { "$ref": "#/definitions/typed_dict" },
+            "task_id": { "type": "string" },
+            "value": {},
+            "from_default": { "type": "boolean" }
+          },
+          "required": [ "name", "kind" ]
+        }
+      },
+      "required": [
+        "__type",
+        "__var"
+      ],
+      "additionalProperties": false
+    },
     "color": {
       "type": "string",
       "pattern": "^#[a-fA-F0-9]{3,6}$"
@@ -345,7 +387,12 @@
         "is_teardown": {"type": "boolean", "default": false},
         "on_failure_fail_dagrun": {"type": "boolean", "default": false},
         "max_active_tis_per_dag": {"type": "integer"},
-        "max_active_tis_per_dagrun": {"type": "integer"}
+        "max_active_tis_per_dagrun": {"type": "integer"},
+        "_arg_bindings": {
+          "$comment": "Only present on @task.stub tasks called with TaskFlow arguments",
+          "type": "array",
+          "items": { "$ref": "#/definitions/arg_binding" }
+        }
       },
       "dependencies": {
         "expand_input": ["partial_kwargs", "_is_mapped"],