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]

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".

[ashb]

not hasattr or is None feels a bit odd. Pick one to assert.

[ashb]

Given how tightly integrated into the exec API this change is, I'm not sure putting it in standard provider is right -- my first thought is that this should live in/with the dag parsing code, not with the stub operator code?

[jason810496]

I will look into whether it would be better to change the dag parsing side or to keep it as is.

[ashb]

"Mapped stubs would need per-map-index arg specs"

I don't think this is true -- it's the same function for each mapped index (by design) so each index would receive the same type of arguments.

[jason810496]

I will double check this part.

[jason810496]

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

[jason810496]

Actually, all the "data_type" here only comes from the annotation at the Stub Operator level. So what exactly the upstream task return doesn't really matter IMO.

There comes up another case that I needs to resolve. The case you give, user might annotate the argument as dict | bool, but I haven't deal with the union type annotation yet.

[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".

[jason810496]

I will look into whether it would be better to change the dag parsing side or to keep it as is.

[jason810496]

I will double check this part.

[Copilot]

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

[Copilot]

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

[jason810496]

Self-review comments.

[jason810496]

Self-review.

[jason810496]

Hi @ashb,

Here are the key updates since your last review:

• Replaced the data_type enum with pydantic-generated JSON-schema fragments, carried per argument as value_schema.
• Kept the argument materialization in the execution API:
  • In the existing Python world, the worker re-parses the Dag file and deserializes the operator, so it gets the call arguments for free.
  • A lang-SDK runtime can't parse Python, so it has to receive materialized bindings.
  • Doing this at Dag-processing time alone isn't enough -- resolving per-map-index values requires joining the TaskInstance at task runtime so ti_run in the execution API is the right place (it's also where API version negotiation strips the field for older clients).
• Deferred mapped-operator support to #70570 and #70571 to keep this one concise enough to review.

[github-actions]

uv.lock on main just moved via #71039 ("Record the relaxed types-paramiko bound in uv.lock"), commit 381339d and this PR currently conflicts.

Quickest fix:

git fetch upstream main && git rebase upstream/main
rm uv.lock && uv lock
git add uv.lock && git rebase --continue
git push --force-with-lease

Automated nudge — ignore if you're not ready to rebase. This comment is updated in place on future uv.lock bumps.

[uranusjr]

I can’t comment on the implementation, but the proposed interface looks very reasonable to me. It seems to me the described programming interface has not been fully implemented in this PR, and I don’t have enough knowledge to say whether this is on the right track or not tbh. Maybe it would be a good idea to edit the PR description to ground what exactly should be expected here.