Support TaskFlow call syntax on stub tasks for the Go SDK

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

@@ -0,0 +1,58 @@
+# 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 a stub task's TaskFlow call (``@task.stub``), stored in the
+serialized Dag, and delivered to the lang-SDK runtime through ``TIRunContext.arg_bindings``
+so it can bind the values onto the native task function's parameters.
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+
+from pydantic import JsonValue
+
+from airflow.api_fastapi.core_api.base import BaseModel
+
+ArgBindingDataType = Literal["string", "integer", "number", "boolean", "object", "array", "any"]
+"""Language-neutral value type a stub-task argument binds to in the foreign runtime."""
+
+
+class TaskArgBinding(BaseModel):
+    """
+    One positional argument of a stub (foreign-runtime) task, in declaration order.
+
+    A deliberately flat shape (``kind`` discriminates instead of a union) so the JSON schema
+    generates a plain struct in the foreign-language SDKs consuming the supervisor schema.
+    """
+
+    kind: Literal["xcom", "literal"]
+    """Whether the value comes from an upstream task's XCom or is a literal from the Dag file."""
+
+    data_type: ArgBindingDataType = "any"
+    """Declared type from the stub function's annotation; runtimes type-check against it."""
+
+    task_id: str | None = None
+    """Upstream task id to pull the XCom from. Only set when ``kind`` is ``xcom``."""
+
+    key: str = "return_value"
+    """XCom key to pull. Only meaningful when ``kind`` is ``xcom``."""
+
+    value: JsonValue | None = None
+    """The literal value from the Dag file. Only set when ``kind`` is ``literal``."""

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

@@ -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 TaskArgBinding
 from airflow.api_fastapi.execution_api.datamodels.taskinstance import (
     InactiveAssetsResponse,
     PreviousTIResponse,
@@ -110,6 +111,36 @@
 log = structlog.get_logger(__name__)
 tracer = trace.get_tracer(__name__)
 
+# Task type recorded on the TI row (``TaskInstance.operator``) for
+# ``airflow.providers.standard.decorators.stub._StubOperator``. Used to gate the
+# serialized-dag lookup for ``arg_bindings`` so regular tasks never pay for it.
+_STUB_TASK_TYPE = "_StubOperator"
+
+
+def _get_arg_bindings(dag_version_id: UUID | None, task_id: str, *, session) -> list[dict] | None:
+    """Extract the stub task's serialized positional-arg spec from the serialized Dag blob."""
+    # Imported here on purpose: only the Multi-Lang stub-task path touches the
+    # serialized-dag machinery, so keep it off the module's top-level imports.
+    from airflow.models.dag_version import DagVersion
+    from airflow.serialization.enums import Encoding
+    from airflow.serialization.serialized_objects import BaseSerialization
+
+    if dag_version_id is None:
+        return None
+    dag_version = session.get(DagVersion, dag_version_id)
+    if dag_version is None or dag_version.serialized_dag is None:
+        return None
+    data = dag_version.serialized_dag.data
+    if not data:
+        return None
+    for task in data.get("dag", {}).get("tasks", []):
+        var = task.get(Encoding.VAR) or {}
+        if var.get("task_id") == task_id:
+            if encoded := var.get("_arg_bindings"):
+                return BaseSerialization.deserialize(encoded)
+            return None
+    return None
+
 
 @ti_id_router.patch(
     "/{task_instance_id}/run",
@@ -163,6 +194,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 +343,13 @@ 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 (
+            arg_bindings := _get_arg_bindings(ti.dag_version_id, ti.task_id, session=session)
+        ):
+            context.arg_bindings = [TaskArgBinding.model_validate(arg) for arg in arg_bindings]
+
         # 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/versions/__init__.py

@@ -41,6 +41,7 @@
     RemoveUpstreamMapIndexesField,
 )
 from airflow.api_fastapi.execution_api.versions.v2026_06_30 import (
+    AddArgBindingsToTIRunContext,
     AddAssetsByAliasEndpoint,
     AddAwaitingInputStatePayload,
     AddConnectionTestEndpoint,
@@ -65,6 +66,7 @@
         AddTaskAndAssetStateStoreEndpoints,
         AddAssetsByAliasEndpoint,
         AddPartitionDateField,
+        AddArgBindingsToTIRunContext,
     ),
     Version(
         "2026-04-06",

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

@@ -141,3 +141,16 @@ def remove_partition_date_from_dag_run(response: ResponseInfo) -> None:  # type:
         """Strip ``partition_date`` from the nested ``dag_run`` payload for older clients."""
         if "dag_run" in response.body and isinstance(response.body["dag_run"], dict):
             response.body["dag_run"].pop("partition_date", None)
+
+
+class AddArgBindingsToTIRunContext(VersionChange):
+    """Add the ``arg_bindings`` positional-argument binding spec for stub (foreign-runtime) tasks."""
+
+    description = __doc__
+
+    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/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py

@@ -370,6 +370,43 @@ async def workload_token(request: Request) -> TIToken:
         assert extras["scope"] == "execution"
         assert extras["sub"] == str(ti.id)
 
+    def test_ti_run_returns_arg_bindings_for_stub_task(self, client, dag_maker):
+        """A stub task's TaskFlow arg spec is extracted from the serialized dag and returned."""
+        from airflow.providers.standard.decorators.stub import stub
+
+        def extract(): ...
+
+        def transform(country: str, extracted: dict): ...
+
+        with dag_maker("test_arg_bindings_dag", serialized=True):
+            stub(transform)("uk", stub(extract)())
+
+        dr = dag_maker.create_dagrun()
+        tis = {ti.task_id: ti for ti in dr.get_task_instances()}
+        for ti in tis.values():
+            ti.set_state(State.QUEUED)
+        dag_maker.session.flush()
+
+        payload = {
+            "state": "running",
+            "hostname": "random-hostname",
+            "unixname": "random-unixname",
+            "pid": 100,
+            "start_date": "2024-09-30T12:00:00Z",
+        }
+
+        response = client.patch(f"/execution/task-instances/{tis['transform'].id}/run", json=payload)
+        assert response.status_code == 200
+        assert response.json()["arg_bindings"] == [
+            {"kind": "literal", "data_type": "string", "value": "uk"},
+            {"kind": "xcom", "data_type": "object", "task_id": "extract", "key": "return_value"},
+        ]
+
+        # An argless stub has no captured spec, so the field stays unset.
+        response = client.patch(f"/execution/task-instances/{tis['extract'].id}/run", json=payload)
+        assert response.status_code == 200
+        assert "arg_bindings" not in response.json()
+
     def test_dynamic_task_mapping_with_parse_time_value(self, client, dag_maker):
         """Test that dynamic task mapping works correctly with parse-time values."""
         with dag_maker("test_dynamic_task_mapping_with_parse_time_value", serialized=True):

airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_04_17/test_task_instances.py

@@ -84,3 +84,46 @@ def test_head_version_includes_team_name_field(self, client, session, create_tas
             response = client.patch(f"/execution/task-instances/{ti.id}/run", json=RUN_PATCH_BODY)
         assert response.status_code == 200
         assert response.json()["dag_run"]["team_name"] is None
+
+
+class TestArgBindingsFieldBackwardCompat:
+    @pytest.fixture(autouse=True)
+    def _freeze_time(self, time_machine):
+        time_machine.move_to(TIMESTAMP_STR, tick=False)
+
+    def setup_method(self):
+        clear_db_runs()
+
+    def teardown_method(self):
+        clear_db_runs()
+
+    @pytest.fixture
+    def stub_ti(self, dag_maker):
+        from airflow.providers.standard.decorators.stub import stub
+
+        def extract(): ...
+
+        def transform(country: str, extracted: dict): ...
+
+        with dag_maker("test_arg_bindings_compat_dag", serialized=True):
+            stub(transform)("uk", stub(extract)())
+
+        dr = dag_maker.create_dagrun()
+        tis = {ti.task_id: ti for ti in dr.get_task_instances()}
+        for ti in tis.values():
+            ti.set_state(State.QUEUED)
+        dag_maker.session.flush()
+        return tis["transform"]
+
+    def test_old_version_strips_arg_bindings_even_when_set(self, old_ver_client, stub_ti):
+        response = old_ver_client.patch(f"/execution/task-instances/{stub_ti.id}/run", json=RUN_PATCH_BODY)
+        assert response.status_code == 200
+        assert "arg_bindings" not in response.json()
+
+    def test_head_version_includes_arg_bindings(self, client, stub_ti):
+        response = client.patch(f"/execution/task-instances/{stub_ti.id}/run", json=RUN_PATCH_BODY)
+        assert response.status_code == 200
+        assert response.json()["arg_bindings"] == [
+            {"kind": "literal", "data_type": "string", "value": "uk"},
+            {"kind": "xcom", "data_type": "object", "task_id": "extract", "key": "return_value"},
+        ]

airflow-core/tests/unit/serialization/test_dag_serialization.py

@@ -77,7 +77,7 @@
 from airflow.serialization.definitions.param import SerializedParam
 from airflow.serialization.definitions.xcom_arg import SchedulerPlainXComArg
 from airflow.serialization.encoders import ensure_serialized_asset
-from airflow.serialization.enums import Encoding
+from airflow.serialization.enums import DagAttributeTypes as DAT, Encoding
 from airflow.serialization.json_schema import load_dag_schema_dict
 from airflow.serialization.serialized_objects import (
     BaseSerialization,
@@ -3405,6 +3405,46 @@ def inner():
     assert serialized3["python_callable_name"] == "empty_function"
 
 
+def test_stub_task_args_round_trip():
+    """The stub task's TaskFlow arg spec (``_arg_bindings``) survives Dag serialization."""
+    from airflow.providers.standard.decorators.stub import stub
+
+    def extract(): ...
+
+    def transform(country: str, extracted: dict): ...
+
+    with DAG(dag_id="arg_bindings_dag", schedule=None) as dag:
+        stub(transform)("uk", stub(extract)())
+
+    ser_dag = DagSerialization.to_dict(dag)
+    encoded_tasks = {t[Encoding.VAR]["task_id"]: t[Encoding.VAR] for t in ser_dag["dag"]["tasks"]}
+    assert "_arg_bindings" not in encoded_tasks["extract"], "argless stubs must not serialize a spec"
+    assert encoded_tasks["transform"]["_arg_bindings"] == [
+        {
+            Encoding.TYPE: DAT.DICT,
+            Encoding.VAR: {"kind": "literal", "data_type": "string", "value": "uk"},
+        },
+        {
+            Encoding.TYPE: DAT.DICT,
+            Encoding.VAR: {
+                "kind": "xcom",
+                "data_type": "object",
+                "task_id": "extract",
+                "key": "return_value",
+            },
+        },
+    ]
+
+    round_tripped = DagSerialization.from_dict(ser_dag)
+    assert round_tripped.task_dict["transform"]._arg_bindings == [
+        {"kind": "literal", "data_type": "string", "value": "uk"},
+        {"kind": "xcom", "data_type": "object", "task_id": "extract", "key": "return_value"},
+    ]
+    assert not hasattr(round_tripped.task_dict["extract"], "_arg_bindings") or (
+        round_tripped.task_dict["extract"]._arg_bindings is None
+    )
+
+
 def test_handle_v1_serdag():
     v1 = {
         "__version": 1,