AIP-103: Add Execution API endpoints for task and asset states

[amoghrajesh]

---

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

• Yes - claude sonnet 4.6

Only the last commit is relevant, this has been built on top of #65759

closes: #66069

This PR is part of AIP-103 (Task State Management) and adds the execution API endpoints for task and asset state.

Summary

• New endpoints under /execution/state/:
  • GET/PUT/DELETE /state/ti/{task_instance_id}/{key} and
    DELETE /state/ti/{task_instance_id} to clear all keys for a task
  • GET/PUT/DELETE /state/asset/{name}/{key} and
    DELETE /state/asset/{name} to clear all keys for an asset
• Routes call the configured [state_store] backend (defaults to
  MetastoreStateBackend from PR 1). Custom worker-side backends
  bypass the API entirely (although I think I should hardcode
  MetastoreStateBackend in API server since we do not want state
  backend credentials there?)
• Task state is JWT-scoped via Security(require_auth, scopes=["ti:self"])
  on the router — a task can only access its own state.
• Single backend instance shared across both routers via a cached
  get_state_backend() in airflow.state.

A few design choices worth flagging

• API first model for custom state backends or not?

Two shapes:

A: API resolves backend, ie: worker → API → could be DB, S3, Redis, or even metastore.

B: API = always DB, ie: worker → backend and API → metastore only. Custom backends live worker-side only and bypass the API — anything stored on a custom state backend goes via the worker itself; if not configured, it falls through to metastore via the API.

Why Shape B?

1. XCom does this. One axis instead of two. Default is through API, custom backend is direct.
2. API stays minimal. No S3/Redis deps on the API server.
3. Lower latency. Custom backend = direct write, no API trip.

But I would love to hear thoughts on A. If not, I can hardcode MetastoreStateBackend in the API endpoints.

• Asset state routes use name (not integer asset_id) in the URL. Asset names are unique, directly available on the task's Asset object at runtime (no extra lookup), and consistent with how /assets/by-name already works. The integer asset_id is an internal DB detail that shouldn't leak into the API surface.

• Why we don't pass the route's session into the backend. First attempt threaded session=session from SessionDep into backend.set(...). mypy rightly complained — BaseStateBackend.set doesn't declare a session kwarg, only MetastoreStateBackend.set does (via @provide_session). I went back and forth on adding *, session: Any = None to the abstract methods, but it leaks a SQLAlchemy concept into an interface that S3/Redis/GCS backends have no use for.

  So: backend manages its own session via @provide_session → create_session(). Routes keep SessionDep for non-backend lookups (TaskInstance / AssetModel existence) — those are reads, no commit needed. Cost is one extra session per state op, negligible. Worth it to keep the abstraction clean.

• No Cadwyn migration. These are net-new endpoints. Old clients that don't call them are unaffected.

• Mapped task semantics. PUT, GET, single-key DELETE, and default DELETE-all all stay scoped to one mapped instance — task_instance_id resolves to a unique (dag_id, run_id, task_id, map_index). Fleet-wide wipe (across every map_index) is opt-in via ?all_map_indices=true on DELETE-all. Destructive bulk is never the default; SDK passes the flag only when the caller asks for it.

What's deferred / todo so far

• Per-task asset registration check. AIP-103 defined that the JWT security model "should be updated to only allow tasks to modify those Assets they are already registered with". Same gap exists in /assets and /asset-events today — the proper fix is a unified check across all asset routes, which I will probably try before merge here. TODO in asset_state.py.

Testing Manually

Setup

1. Breeze running: breeze start-airflow --backend postgres --load-example-dags
2. A DAG triggered so at least one TaskInstance exists
3. Use this script and run in breeze container and use that token:

def main() -> None:
    from airflow.configuration import conf

    secret = conf.get("api_auth", "jwt_secret", fallback=None)
    pkey_path = conf.get("api_auth", "jwt_private_key_path", fallback=None)
    print(
        f"[diag] jwt_secret first 8 chars: {(secret or '')[:8]!r} (empty means using private key file or auto-gen)"
    )
    print(f"[diag] jwt_private_key_path: {pkey_path!r}")
    print(f"[diag] jwt_audience: {conf.get('execution_api', 'jwt_audience', fallback=None)!r}")

    with create_session() as session:
        ti = session.scalar(select(TaskInstance).limit(1))
        if ti is None:
            print(
                "No task instances found. Trigger a DAG first (e.g. `airflow dags trigger example_bash_operator`)."
            )
            return
        gen = _jwt_generator()
        token = gen.generate(extras={"sub": str(ti.id), "scope": "execution"})
        print(f"ti_id:  {ti.id}")
        print(f"dag_id: {ti.dag_id}  run_id: {ti.run_id}  task_id: {ti.task_id}  map_index: {ti.map_index}")
        print(f"token:  {token}")


if __name__ == "__main__":
    main()

Example:

[Breeze:3.10.20] root@f43294940632:/opt/airflow$ python dev/gen_ti_token.py
2026-04-29T08:26:29.764581Z [info     ] setup plugin alembic.autogenerate.schemas [alembic.runtime.plugins] loc=plugins.py:37
2026-04-29T08:26:29.765573Z [info     ] setup plugin alembic.autogenerate.tables [alembic.runtime.plugins] loc=plugins.py:37
2026-04-29T08:26:29.766389Z [info     ] setup plugin alembic.autogenerate.types [alembic.runtime.plugins] loc=plugins.py:37
2026-04-29T08:26:29.766492Z [info     ] setup plugin alembic.autogenerate.constraints [alembic.runtime.plugins] loc=plugins.py:37
2026-04-29T08:26:29.766558Z [info     ] setup plugin alembic.autogenerate.defaults [alembic.runtime.plugins] loc=plugins.py:37
2026-04-29T08:26:29.766622Z [info     ] setup plugin alembic.autogenerate.comments [alembic.runtime.plugins] loc=plugins.py:37
[diag] jwt_secret first 8 chars: 'tIyIg3gY' (empty means using private key file or auto-gen)
[diag] jwt_private_key_path: None
[diag] jwt_audience: 'urn:airflow.apache.org:task'
2026-04-29T08:26:30.272426Z [warning  ] The HMAC key is 24 bytes long, which is below the minimum recommended length of 64 bytes for SHA512. See RFC 7518 Section 3.2. [py.warnings] category=InsecureKeyLengthWarning filename=/usr/python/lib/python3.10/site-packages/jwt/api_jwt.py lineno=147
ti_id:  019dd858-888f-7160-b16f-033390d4386a
dag_id: my_dag  run_id: manual__2026-04-29T08:26:14.011123+00:00  task_id: t1  map_index: -1
token:  eyJhbGciOiJIUzUxMiIsImtpZCI6Im5vdC11c2VkIiwidHlwIjoiSldUIn0.eyJzdWIiOiIwMTlkZDg1OC04ODhmLTcxNjAtYjE2Zi0wMzMzOTBkNDM4NmEiLCJzY29wZSI6ImV4ZWN1dGlvbiIsImp0aSI6Ijc2NmM2MDgyZDEzZTQ1MzU4Y2U5YjNkMmZmMjZhOWRiIiwiaXNzIjoiYWlyZmxvdyIsImF1ZCI6InVybjphaXJmbG93LmFwYWNoZS5vcmc6dGFzayIsIm5iZiI6MTc3NzQ1MTE5MCwiZXhwIjoxNzc3NDUxNzkwLCJpYXQiOjE3Nzc0NTExOTB9.Y-nAYstD5zcBiX7JcXqCml1-sIFlmjqjkAYSNbOOQNaRmdMaYPmxcTIAX1fKr0a_fgRyla3sGOMhjdACBD_sEQ

5. I created a postman collections with AIP 103 related API endpoints, but I will showcase running them below.

Testing

Regular Tasks

1. Put on task_state

Validation:

2. Get task state

3. Delete task state

Validation:

4. Write and then task state for same task

6. Demonstrating clear_all

• First adding one more key

• Sending API call

All gone:

Mapped Tasks

• Everything related to PUT, GET, DELETE remains same because ti_id is unique for mapped tasks and it will populate those values right but DELETE-all is intentionally different. A worker calling DELETE /state/ti/{ti_id} wipes state across every map_index of the task, not just its own. The reasoning: clear-all is a "task is done with its state" operation, not a per-instance reset. Per-instance reset is already covered by DELETE. Showcasing that below.

Using this DAG:

from airflow import DAG
from airflow.sdk import task
from datetime import datetime


with DAG(
    dag_id="example_mapped_tasks",
    start_date=datetime(2024, 1, 1),
    schedule=None,
    catchup=False,
) as dag:

    @task
    def generate_numbers():
        return [1, 2, 3, 4, 5]

    @task
    def square(x: int):
        return x * x

    numbers = generate_numbers()

    squared = square.expand(x=numbers)

Ran it and this is TI:

Pushed in task_state for map_index: 0,1,4

1. Now trying DELETE ALL with ti_id: 00000000-0000-0000-0000-000000000001 (Negative case)

2. Trying to run DELETE ALL API but with TI of task_instance as 4

(It just deleted it for that task_state with map_index as 4)

3. Trying to run DELETE ALL API but with TI of task_instance as -1 (019ddd86-34b6-7eb5-9dcb-96803d1c49f5)

4. Passing in the destructive flag: all_map_indices to see if it deletes all

All have been deleted

Assets being tested

• Used this DAG and triggered once and created JWT token as earlier:

from __future__ import annotations

import pendulum

from airflow.sdk import DAG, Asset, task

aip103_test_asset = Asset(
    name="aip103_test_asset",
    uri="s3://aip103-test/watermarks/orders",
    group="asset",
)


with DAG(
    dag_id="aip103_asset_producer",
    start_date=pendulum.datetime(2026, 1, 1, tz="UTC"),
    schedule=None,
    catchup=False,
    tags=["aip-103", "asset-state-test"],
):

    @task(outlets=[aip103_test_asset])
    def produce():
        print(f"Producer running for {aip103_test_asset.uri!r}")

    produce()


with DAG(
    dag_id="aip103_asset_consumer",
    start_date=pendulum.datetime(2026, 1, 1, tz="UTC"),
    schedule=None,
    catchup=False,
    tags=["aip-103", "asset-state-test"],
):

    @task(inlets=[aip103_test_asset])
    def consume():
        print(f"Consumer running for {aip103_test_asset.uri!r}")

    consume()

5a. Put asset state:

5b. Get asset state:

5c. Delete asset state

5d. Delete all asset state

For that adding multiple keys, for same asset

Running delete all request

All gone:

---

• Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
• For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
• When adding dependency, check compliance with the ASF 3rd Party License Policy.
• For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

[jscheffl]

Looks good in general.

[amoghrajesh]

Rerunning with full tests

[amoghrajesh]

Merging this PR in right now to unblock further PRs, I have handled all the review comments on the PR, thanks for reviews @kaxil @jscheffl!

[github-actions]

Backport failed to create: v3-2-test. View the failure log Run details

Note: As of Merging PRs targeted for Airflow 3.X
the committer who merges the PR is responsible for backporting the PRs that are bug fixes (generally speaking) to the maintenance branches.

In matter of doubt please ask in #release-management Slack channel.

Status	Branch	Result
❌	v3-2-test

You can attempt to backport this manually by running:

cherry_picker ffb1b8a v3-2-test

This should apply the commit to the v3-2-test branch and leave the commit in conflict state marking
the files that need manual conflict resolution.

After you have resolved the conflicts, you can continue the backport process by running:

cherry_picker --continue

If you don't have cherry-picker installed, see the installation guide.

[amoghrajesh]

No need to rebase