Support TaskFlow call syntax on stub tasks for the Lang SDK

[jason810496]

Scope

This PR ships the Python side only of the contract: parse-time capture of the TaskFlow call into a serialized arg-binding spec, the wire model, and its delivery to SDK runtimes through the Execution API and supervisor schema.
Nothing in this PR binds arguments inside a task runtime — the Go snippet below illustrates the consumer and lives in the stacked follow-up.

Stacked on top of this PR:

• Bind TaskFlow stub-task call arguments in the Go SDK runtime #70209 — Bind TaskFlow stub-task call arguments in the Go SDK runtime (consumes arg_bindings)
• Support dynamic task mapping on stub tasks #70570 — Support dynamic task mapping on stub tasks (lifts the .expand()/.partial() rejection below)

Why

@task.stub tasks can only be declared argless today, so a Go task that needs an upstream's output has to hand-write GetXCom calls (with the upstream task_id hard-coded in Go, duplicating the wiring the Dag file already expresses). This PR ships the Python side of making the natural TaskFlow call work across the language boundary:

@task.stub(queue="golang")
def transform(country: str, extracted: dict): ...


with DAG(...):
    transform("uk", extract())  # extract() is a normal Python @task

// The runtime (stacked #70209) binds "uk" onto country and pulls
// extract's XCom into extracted.
func transform(ctx sdk.TIRunContext, log *slog.Logger, country string, extracted map[string]any) error

Supported TaskFlow syntax

Every form below parses, serializes, and round-trips through the execution API (provider capture matrix in test_stub.py; the Dag shown is #70209's taskflow_binding_dag example):

@task.stub(queue="golang")
def via_flat_args(
    name: str,
    count: int,
    ratio: float,
    enabled: bool,
    tags: list,
    config: dict,
    numbers: list,
    note: str | None = None,
): ...


@dag(dag_id="taskflow_binding_dag")
def taskflow_binding_dag():
    via_flat_args(
        "summary",
        3,
        2.5,
        True,  # positional scalar literals (str/int/float/bool)
        ["metrics", "hourly"],  # array literal
        config=make_config(),  # keyword arg: XCom from another @task.stub
        numbers=make_numbers(),  # XCom binding onto a typed array parameter
    )  # `note` unpassed: its None default is captured as from_default
    region = make_region()
    via_struct_no_tags(RegionCode=region, Threshold=0.75)  # one XCom fanned into several calls
    via_struct_arg_tag(region_code=region, threshold=0.75)  # literal + XCom mixed as kwargs
    via_struct_unmatched_arg(region_code=region)  # defaulted param left unpassed

• Literals of every JSON shape, positional or keyword; the wire value_schema is the JSON-schema fragment pydantic generates from the parameter annotation (TypeAdapter(annotation).json_schema() with a GenerateJsonSchema subclass layering OpenAPI's int64/double numeric formats): str→{"type": "string"}, int→{"type": "integer", "format": "int64"}, dict[str, int]→{"type": "object", "additionalProperties": {...}}, list[int]→{"type": "array", "items": {...}}, Literal["a", "b"]→{"type": "string", "enum": [...]}, datetime/date/time/timedelta→string with the standard date-time/date/time/duration formats, unions→standard anyOf (str | None→{"anyOf": [{"type": "string"}, {"type": "null"}]}); annotations pydantic cannot schema (arbitrary classes, unresolvable names) and untyped/Any parameters omit value_schema entirely (decode-only binding).
• XComArg return values from any upstream — a normal @task or another @task.stub — which also wires the dependency edge (transform("uk", extract()) implies extract >> transform).
• Keyword arguments normalize to declaration order through signature binding, so the serialized spec is always positional.
• Defaults left unpassed are captured with from_default: true, letting keyword-style consumers (the Go sdk.TaskInput struct mode) leave them unclaimed.
• Argless calls behave exactly as before: no spec is serialized, and pre-existing signatures (**kwargs, context-key parameter names) keep parsing.

Rejected loudly at parse time (v1 scope): .expand()/.partial() on a stub, stubs called with arguments inside a mapped task group, map/zip/concat XComArgs, indexing an upstream by custom XCom key, *args/**kwargs or context-key parameter names (only when the call actually passes arguments), and non-JSON-serializable literals (including NaN/Infinity).

How

• Parse time (providers/standard): _StubOperator binds the TaskFlow call to the stub's signature and serializes an ordered positional-arg spec with the Dag (_arg_bindings); the cross-version imports it needs are routed through the common.compat sdk seam.
• Wire model: each spec entry is one variant of a kind-discriminated union — XComArgBinding (pull of an upstream's return-value XCom by task_id) or LiteralArgBinding (inline JSON value, from_default-flagged when captured from a signature default) — carrying the stub parameter's name and an annotation-derived value_schema. The fragment is deliberately free-form (dict[str, JsonValue], not a typed model) so every keyword pydantic generated survives the server → supervisor → runtime trip verbatim; consumers validate the keywords they understand and ignore the rest, per JSON-schema semantics. The exact fragment shape follows the pydantic version active at Dag-parse time.
• Server: ti_run returns the spec as a new optional TIRunContext.arg_bindings field, resolved through the shared DBDagBag only for _StubOperator tasks and stripped for older clients by a new execution-API version 2026-10-30 (Airflow 3.4 target); the supervisor schema gets a mirror version 2026-10-30 with a downgrade path, and the generated ts-sdk models plus the Go SDK's SupervisorSchemaVersion pin follow the schema bump in this PR. The serialized-Dag schema.json documents the optional per-task _arg_bindings property (no SERIALIZER_VERSION bump: optional field, no serialization-logic change).
• Anything outside the v1 contract fails loudly at parse time (the rejected list above).

---

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

• Yes, with help of Claude Code (Fable 5) following the guidelines

[jason810496]

I will wait until next review then address my own comments to avoid CI-rerun.

[ashb]

Overall I like the direction. Almost all of my comments I can be challenged on, don't just make the changes if you think the current way is better/more correct

[ashb]

Hmmmmm, I wonder if this should not allow none, and make it an empty list in that case. I don't think it functionally makes a difference but... 🤔

[ashb]

You resolved this but didn't change anything or say why it's fine as it is.

[ashb]

I wonder if introducing arg_bindings should result in a bump in the serialization version? @amoghrajesh WDYT?

[jason810496]

I will update the airflow-core/src/airflow/serialization/schema.json to show the new args_binding field (the native Dag TaskFlow will leverage the args_binding field as well). However, I don't think it's not necessary to bump the "serialization version".

[amoghrajesh]

I don't think this needs a bump.

Its purely additive and only appears on @task.stub tasks that were called with taskflow args. definitions.operator also has additionalProperties: true and validate_schema only runs on the write path, so nothing on the read side changes for a v3 blob. conversion_v3_to_v4 would probably be an empty function or very less.

[ashb]

I agree it's not "needed", no, but is it cleaner/a better forward pattern if we start making the version mean something?

In Airflow 2 the schema version was always 1, so it effectively never got used, and I'm not sure that's a good pattern.

[jason810496]

Thanks you Ash for the review. I will address the comments shortly.

[jason810496]

I will update the airflow-core/src/airflow/serialization/schema.json to show the new args_binding field (the native Dag TaskFlow will leverage the args_binding field as well). However, I don't think it's not necessary to bump the "serialization version".

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[ashb]

Elsewhere we do this:

    state: Annotated[
        Literal[TIState.RUNNING],
        # Specify a default in the schema, but not in code.
        WithJsonSchema({"type": "string", "enum": [TIState.RUNNING], "default": TIState.RUNNING}),
    ]

It might be worth seeing if we can do the same thing here to apply a default too?

Or this approach from ExecuteTask workload:

    type: Literal["ExecuteTask"] = Field(init=False, default="ExecuteTask")

[uranusjr]

Gating with operator type name is not a good idea. Add a flag on _StubOperator, serialize it, and read it here instead.

[ashb]

Yes, this is what we do for EmptyOperator -- we have a "is_empty" field or similar.

[uranusjr]

There’s also still a mention to sdk.TaskInput that does not exist. Otherwise I think this is good (aside from the comment above).

[ashb]

I'm still not sure what to make of the schema/data type on an xcom arg. How will it be used?

[ashb]

Why is this change in this PR?

[jason810496]

My bad, this is a bad rebase.

[ashb]

Stub operator should be made not-usable on Airflow2 if someone tries to instantiate one -- it wouldn't ever work anyway without the tasksdk machinery at execution time

[ashb]

I don't think this limitation should stay.

[ashb]

I can mayve see concat, and zip not being supported, but why no map?

Is this a "not for now" or a "it's really hard, possibly never"?

[jason810496]

All the mapped operator cases will be supported in #70570 as follow-up. It's a bit difficult but it's feasible.

[ashb]

I'm still not convinced about all this type handling logic being in the standard provider. A lot of this feels like it belongs in the dag parser, keeping the stub very simple.