Java SDK: Honor TaskFlow arg bindings sent by the supervisor

airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst

@@ -93,10 +93,7 @@ Java implementation
       }
 
       @Builder.Task(id = "transform")
-      public long transform(
-        Client client,
-        @Builder.XCom(task = "extract") long recordCount
-      ) {
+      public long transform(Client client, long recordCount) {
         var threshold = (String) client.getVariable("transform_threshold");
         // ... process data ...
         return transformedCount;
@@ -105,8 +102,11 @@ Java implementation
 
 .. note::
 
-  See how both ``transform`` in Python and Java need to have an argument to accept upstream XCom. The
-  Python one is needed to declare dependency, and the Java one is needed to actually retrieve the value.
+  The graph is declared once, in the Python Dag file: ``transform(extract())`` feeds the upstream's
+  return value into the downstream's parameter by calling tasks like functions. The supervisor sends
+  the resulting *argument bindings* to the Java runtime, and each Java data parameter receives
+  whatever the Python call site bound at its position — an upstream task's XCom or an inline
+  literal. See :ref:`java-sdk/arg-binding`.
 
 Java entry point
 ~~~~~~~~~~~~~~~~
@@ -115,7 +115,7 @@ Java entry point
 
     public class Main implements BundleBuilder {
       @Override
-      public Iterable<Dag> getDags() {
+      public Iterable<DagDef> getDags() {
         return List.of(SalesPipelineBuilder.build());  // SalesPipelineBuilder generated at compile time
       }
 
@@ -164,12 +164,17 @@ Annotate a plain Java class and let the SDK generate the boilerplate at compile
    * - ``@Builder.Task(id = "...")``
      - Marks a method as a task implementation.  The ``id`` must match the ``@task.stub`` function
        name in the Python Dag.  If ``id`` is omitted the method name is used.
-   * - ``@Builder.XCom(task = "...")``
-     - Injects the ``return_value`` XCom from the named upstream task as a method parameter.
-       The parameter type must be compatible with the stored value (see :ref:`java-sdk/types`).
+   * - ``TaskInput`` / ``@ArgName("...")``
+     - Marks a class as a task's input bundle, so keyword arguments bind by name instead of by
+       position: each public field receives the binding whose name matches it (the ``@ArgName``
+       value, or the verbatim field name).  See :ref:`java-sdk/arg-binding`.
+
+Besides the annotations, a task method may declare a ``Client`` and a ``Context`` parameter in any
+position; the SDK injects both.  Every other parameter is a *data parameter* and receives an
+argument bound by the Python ``@task.stub`` call site.
 
 The annotation processor generates a ``<ClassName>Builder`` class that wires up the task
-registry and handles XCom injection automatically.
+registry and resolves data parameters and XCom pushes automatically.
 
 .. code-block:: java
 
@@ -184,10 +189,7 @@ registry and handles XCom injection automatically.
       }
 
       @Builder.Task(id = "process")
-      public long process(
-        Client client,
-        @Builder.XCom(task = "fetch") String fetched
-      ) {
+      public long process(Client client, String fetched) {
         var threshold = (String) client.getVariable("process_threshold");
         // implement task logic
         return count;
@@ -203,7 +205,7 @@ Interface-based API
 ~~~~~~~~~~~~~~~~~~~
 
 Implement the ``Task`` interface directly for full control over how tasks are registered and how XComs are
-read.
+read.  Each task is registered as a ``TaskDef`` on a ``DagDef``.
 
 .. code-block:: java
 
@@ -224,16 +226,115 @@ Register tasks manually in a ``BundleBuilder``:
 
     public class MyBundle implements BundleBuilder {
       @Override
-      public Iterable<Dag> getDags() {
-        var dag = new Dag("my_dag");
-        dag.addTask("fetch", FetchTask.class);
-        dag.addTask("process", ProcessTask.class);
+      public Iterable<DagDef> getDags() {
+        var dag = new DagDef("my_dag")
+            .addTask(new TaskDef("fetch", FetchTask.class))
+            .addTask(new TaskDef("process", ProcessTask.class));
         return List.of(dag);
       }
     }
 
 See the `Java SDK API Reference <https://airflow.apache.org/docs/java-sdk/stable/>`__ for more details.
 
+.. _java-sdk/arg-binding:
+
+Binding stub arguments
+~~~~~~~~~~~~~~~~~~~~~~
+
+Calling a ``@task.stub`` TaskFlow-style in the Python Dag is what declares the graph, and the
+supervisor delivers the resulting argument bindings to the Java runtime with every task run.  A
+binding carries either an upstream task's ``return_value`` XCom or an inline literal written at the
+call site.
+
+Positional binding
+^^^^^^^^^^^^^^^^^^
+
+A task method's data parameters bind **by position**, in declaration order — the injected ``Client``
+and ``Context`` parameters do not take up a position.  Java parameter names are not part of the API,
+so renaming one in an IDE never rebinds an input.
+
+.. code-block:: python
+
+    @task.stub(queue="java")
+    def score(rows, threshold): ...
+
+
+    score(load_rows(), 0.75)
+
+.. code-block:: java
+
+    @Builder.Task(id = "score")
+    public long score(Client client, long rows, double threshold) {
+      // rows      <- the load_rows XCom (position 0)
+      // threshold <- the literal 0.75   (position 1)
+    }
+
+A primitive parameter cannot hold ``null``, so the task fails with ``MissingXComException`` when its
+binding resolves to nothing; declare a boxed type (``Long``, ``Double``, …) to receive ``null``
+instead.  Declaring more data parameters than the call site bound also fails the task, rather than
+running it with missing inputs.
+
+Named binding with a ``TaskInput`` bundle
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+To bind keyword arguments by name, declare a single parameter whose class implements ``TaskInput``.
+Each public non-final field receives the binding named by its ``@ArgName`` value, or by the verbatim
+field name — the deliberate, tagged boundary where the stub's ``snake_case`` argument names cross
+into ``camelCase`` Java fields.  The class needs a public no-argument constructor.
+
+.. code-block:: python
+
+    @task.stub(queue="java")
+    def score(region_code, threshold): ...
+
+
+    score(region_code="emea", threshold=load_threshold())
+
+.. code-block:: java
+
+    public static class ScoreInput implements TaskInput {
+      @ArgName("region_code")
+      public String region;
+
+      public double threshold;
+    }
+
+    @Builder.Task(id = "score")
+    public long score(Client client, ScoreInput input) { ... }
+
+A task declares flat data parameters **or** one ``TaskInput`` bundle, never both, so field names and
+flat positions cannot shift each other.  Mixing them, or declaring two bundles, fails the build.
+
+Reading bindings from the interface API
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Tasks written against the ``Task`` interface read the same bindings imperatively:
+
+.. list-table::
+   :header-rows: 1
+   :widths: 40 60
+
+   * - ``Client`` method
+     - Returns
+   * - ``hasArgs()``
+     - Whether the Python Dag called this stub with any TaskFlow arguments at all.
+   * - ``hasArg(int position)`` / ``hasArg(String name)``
+     - Whether an argument was bound at that position, or with that name.
+   * - ``getArg(int position)`` / ``getArg(String name)``
+     - The bound value — the inline literal, or the bound upstream's XCom.  Throws
+       ``IllegalArgumentException`` when nothing was bound there; probe with ``hasArg`` first.
+
+.. code-block:: java
+
+    public class ScoreTask implements Task {
+      @Override
+      public void execute(Context context, Client client) throws Exception {
+        var rows = client.getArg(0);
+        var threshold = client.hasArg("threshold") ? client.getArg("threshold") : 0.5;
+        // implement task logic
+      }
+    }
+
 .. _java-sdk/logging:
 
 Logging
@@ -428,7 +529,7 @@ represented as Java objects when read back via ``getXCom``.
 
 .. note::
 
-   An ``@Builder.XCom`` parameter that reads a value which was never pushed resolves to
+   A data parameter whose binding resolves to a value that was never pushed receives
    ``null``.  A boxed parameter (``Integer``, ``Long``, ``Boolean``, …) receives ``null``
    safely, but a primitive parameter (``int``, ``long``, ``boolean``, …) cannot represent
    ``null`` and the task fails with ``MissingXComException``.  Declare the parameter with a

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