Add asset and task store UI

[bbovenzi]

Add CRUDs action for Asset and Task Stores

Asset Store:

• Create a new tab navigation on an asset page to switch between events and asset state

Task Store:

• Move xcoms and task state into a "Storage" tab with an xcom and task state sub-tabs. Xcoms url is preserved.

---

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

• Yes Claude Sonnet 4.6

---

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

[amoghrajesh]

Two thins related to editing a task state:

1. When I try to edit a task state, this is how the popup looks like:

Is it possible to pretty print here too?

2. Delete persists correct value to DB but the edit dialog later shows then old value

[amoghrajesh]

One other thing I noticed about adding a task / asset state is, UI allows adding things like incomplete jsons, something like:

curl --location --request PUT 'http://localhost:28080/api/v2/dags/my_dag/dagRuns/manual__2026-05-22T07:59:31.188183+00:00/taskInstances/t1/states/job_id' \
--header 'Content-Type: application/json' \
--header 'Authorization: ••••••' \
--data '{ "value": "incomplete}'

The API responds with 422 error but UI is sending it as:

curl 'http://localhost:28080/api/v2/dags/my_dag/dagRuns/manual__2026-05-22T07:59:31.188183+00:00/taskInstances/t1/states/abcd?map_index=-1' \
-X 'PUT' \
--data-raw '{"value":"{\"abcd\": \"a}"}'

I think that the UI should validate the form field before constructing the JSON body to send

[amoghrajesh]

Thanks for the awesome work, @bbovenzi!

Looking really nice, I left some comments on the PRs for issues after testing these many things, you have my dags but the last one for mapped is here:

from __future__ importannotationsimportjsonimportrandomfromdatetimeimportdatetime, timezonefromairflow.sdkimportDAG, taskTABLES= ["orders", "customers", "products"]
withDAG(
dag_id="example_task_state_mapped",
schedule=None,
start_date=datetime(2026, 1, 1),
catchup=False,
tags=["example", "aip-103", "task-state", "mapped"],
doc_md=__doc__,
) asdag:
@taskdefget_tables() ->list[str]:
"""Return the list of tables to process."""returnTABLES@taskdefprocess_table(table: str, **context) ->dict:
"""Process one table — each mapped instance gets its own task state."""ts=context["task_state"]
map_index=context["task_instance"].map_indexrow_count=random.randint(100, 10000)
result= {
"table": table,
"map_index": map_index,
"row_count": row_count,
"processed_at": datetime.now(tz=timezone.utc).isoformat(timespec="seconds"),
}
ts.set("table", table)
ts.set("status", "complete")
ts.set("row_count", str(row_count))
ts.set("result", json.dumps(result))
print(f"[map_index={map_index}] Processed {table}: {row_count} rows")
returnresulttables=get_tables()
process_table.expand(table=tables)

Task State — Spark DAG

• All keys visible after a completed run (job_id, submitted_at, status, poll_result, completed_at)
• poll_result JSON is pretty-printed
• job_id shows Never in Expires At, other keys show a date
• After retry-reattach: same job_id persists, status updates to complete
• Delete a single key — row gone, others intact
• Edit a key — new value shows immediately
• Clear all — table goes empty

Asset State — Watermark DAG

• First run: watermark, total_runs=1, last_run_summary appear on asset detail page
• Subsequent runs: total_runs increments, watermark advances, prev_watermark matches previous run
• Consumer DAG fires automatically after each producer run
• Clear asset state then re-trigger: total_runs=1, prev_watermark=null

Mapped Tasks — Mapped DAG (example_task_state_mapped)

• Trigger DAG — 3 mapped instances run (map_index 0, 1, 2 for orders/customers/products)
• Each mapped TI shows its own table, row_count, result in Storage tab — no bleed between instances
• Switching between map_index 0/1/2 in the UI shows different state values
• Clear single instance (map_index=0) — only that instance's state is gone, others intact
• Clear all (all_map_indices=true) — state wiped across all 3 instances

[amoghrajesh]

I found another bug related to the core API where editing a task state field overwrote the expiry to NEVER and also we didn't provide an option for users to set expiry for a task state when creating a new one. For the task state Storage tab, here's how the UI should call the API:

Adding a new task state:

So the modal for adding a task state ^ will need a new field for expiry_date - I imagine a datetime picker along with something that can serve three expiry options:

1. "default" pre-selected
2. Maybe a radio button for "never expire"
3. Datetime picker for selecting a datetime for expiry

Once a value is picked, call the PUT /states/{key} with:

• {"value": "...", "expires_at": "default"} — apply server default retention
• {"value": "...", "expires_at": null} — never expire
  {"value": "...", "expires_at": "2026-06-01T00:00:00Z"} — specific datetime from the date picker

Editing an existing task state:
Edit existing key (value only) now can call the PATCH /states/{key} with {"value": "..."}. Expiry is always preserved, no expiry field needed.

This is being fixed in #67319

[amoghrajesh]

This one will also serve as a good enhancement here: #67530

[amoghrajesh]

Taking it for a spin again!

[amoghrajesh]

Thanks for the awesome work, @bbovenzi!

Looking really nice, I left some comments on the PRs for issues after testing these many things, you have my dags but the last one for mapped is here:

from __future__ importannotationsimportjsonimportrandomfromdatetimeimportdatetime, timezonefromairflow.sdkimportDAG, taskTABLES= ["orders", "customers", "products"]
withDAG(
dag_id="example_task_state_mapped",
schedule=None,
start_date=datetime(2026, 1, 1),
catchup=False,
tags=["example", "aip-103", "task-state", "mapped"],
doc_md=__doc__,
) asdag:
@taskdefget_tables() ->list[str]:
"""Return the list of tables to process."""returnTABLES@taskdefprocess_table(table: str, **context) ->dict:
"""Process one table — each mapped instance gets its own task state."""ts=context["task_state"]
map_index=context["task_instance"].map_indexrow_count=random.randint(100, 10000)
result= {
"table": table,
"map_index": map_index,
"row_count": row_count,
"processed_at": datetime.now(tz=timezone.utc).isoformat(timespec="seconds"),
}
ts.set("table", table)
ts.set("status", "complete")
ts.set("row_count", str(row_count))
ts.set("result", json.dumps(result))
print(f"[map_index={map_index}] Processed {table}: {row_count} rows")
returnresulttables=get_tables()
process_table.expand(table=tables)

Task State — Spark DAG

• All keys visible after a completed run (job_id, submitted_at, status, poll_result, completed_at)
• poll_result JSON is pretty-printed
• job_id shows Never in Expires At, other keys show a date
• After retry-reattach: same job_id persists, status updates to complete
• Delete a single key — row gone, others intact
• Edit a key — new value shows immediately
• Clear all — table goes empty

Asset State — Watermark DAG

• First run: watermark, total_runs=1, last_run_summary appear on asset detail page
• Subsequent runs: total_runs increments, watermark advances, prev_watermark matches previous run
• Consumer DAG fires automatically after each producer run
• Clear asset state then re-trigger: total_runs=1, prev_watermark=null

Mapped Tasks — Mapped DAG (example_task_state_mapped)

• Trigger DAG — 3 mapped instances run (map_index 0, 1, 2 for orders/customers/products)
• Each mapped TI shows its own table, row_count, result in Storage tab — no bleed between instances
• Switching between map_index 0/1/2 in the UI shows different state values
• Clear single instance (map_index=0) — only that instance's state is gone, others intact
• Clear all (all_map_indices=true) — state wiped across all 3 instances

I am retesting similar scenarios but with the example dags for task store and asset store committed to repo.

1. The section: ## Task State — Spark DAG still looks great! And love the validation here.

2. Same with: Asset State — Watermark DAG section, looks great.
   Validation on UI is great too.

3. Using this dag for mapped

from __future__ importannotationsimportrandomfromdatetimeimportdatetime, timezonefromairflow.sdkimportDAG, taskTABLES= ["orders", "customers", "products"]
withDAG(
dag_id="example_task_store_mapped",
schedule=None,
start_date=datetime(2026, 1, 1),
catchup=False,
tags=["example", "aip-103", "task-state", "mapped"],
doc_md=__doc__,
) asdag:
@taskdefget_tables() ->list[str]:
"""Return the list of tables to process."""returnTABLES@taskdefprocess_table(table: str, **context) ->dict:
"""Process one table — each mapped instance gets its own task state."""ts=context["task_store"]
map_index=context["task_instance"].map_indexrow_count=random.randint(100, 10000)
result= {
"table": table,
"map_index": map_index,
"row_count": row_count,
"processed_at": datetime.now(tz=timezone.utc).isoformat(timespec="seconds"),
}
ts.set("table", table)
ts.set("status", "complete")
ts.set("row_count", row_count)
ts.set("result", result)
print(f"[map_index={map_index}] Processed {table}: {row_count} rows")
returnresulttables=get_tables()
process_table.expand(table=tables)

Looks great too, tried, get, set, clear, clear all etc.

Now for custom backends, using this backend:

# 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."""File-based state backend for testing the ExternalState envelope.Workers write state values as JSON files under /tmp/airflow_state/ and storethe file path as the external reference. The DB therefore holds: {"__type": "ExternalState", "__var": "/tmp/airflow_state/ti_<id>/job_id.json"}instead of the raw value, which lets you verify the envelope behaviour end-to-end.Configure in airflow.cfg (or via env var) before starting a worker: [workers] state_backend = dev.file_state_backend.FileStateBackendThe server-side abstract methods (get/set/delete/clear and their async variants)raise NotImplementedError — this backend is purely a worker-side serialization hook."""from __future__ importannotationsimportjsonfrompathlibimportPathfromtypingimportTYPE_CHECKINGfromairflow.sdk.stateimportBaseStoreBackendifTYPE_CHECKING:
fromdatetimeimportdatetimefrompydanticimportJsonValuefromsqlalchemy.ext.asyncioimportAsyncSessionfromsqlalchemy.ormimportSessionfromairflow_shared.stateimportStoreScopeBASE_DIR=Path("/tmp/airflow_state")
classFileStateBackend(BaseStoreBackend):
"""Stores task/asset state values as local JSON files; returns the path as the ref."""defserialize_task_store_to_ref(self, *, value: JsonValue, key: str, ti_id: str) ->str:
path=BASE_DIR/f"ti_{ti_id}"/f"{key}.json"path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(json.dumps(value))
returnstr(path)
defdeserialize_task_store_from_ref(self, stored: str) ->JsonValue:
returnjson.loads(Path(stored).read_text())
defserialize_asset_state_to_ref(self, *, value: JsonValue, key: str, asset_ref: str) ->str:
safe=asset_ref.replace("/", "_").replace(":", "")
path=BASE_DIR/"assets"/safe/f"{key}.json"path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(json.dumps(value))
returnstr(path)
defdeserialize_asset_state_from_ref(self, stored: str) ->JsonValue:
returnjson.loads(Path(stored).read_text())
defget(self, scope: StoreScope, key: str, *, session: Session|None=None) ->str|None:
raiseNotImplementedError(
"FileStateBackend is a worker-side backend; server uses MetastoreStateBackend"
)
defset(
self,
scope: StoreScope,
key: str,
value: str,
*,
expires_at: datetime|None=None,
session: Session|None=None,
) ->None:
raiseNotImplementedErrordefdelete(self, scope: StoreScope, key: str, *, session: Session|None=None) ->None:
raiseNotImplementedErrordefclear(
self, scope: StoreScope, *, all_map_indices: bool=False, session: Session|None=None
) ->None:
raiseNotImplementedErrorasyncdefaget(self, scope: StoreScope, key: str, *, session: AsyncSession|None=None) ->str|None:
raiseNotImplementedErrorasyncdefaset(
self,
scope: StoreScope,
key: str,
value: str,
*,
expires_at: datetime|None=None,
session: AsyncSession|None=None,
) ->None:
raiseNotImplementedErrorasyncdefadelete(self, scope: StoreScope, key: str, *, session: AsyncSession|None=None) ->None:
raiseNotImplementedErrorasyncdefaclear(
self, scope: StoreScope, *, all_map_indices: bool=False, session: AsyncSession|None=None
) ->None:
raiseNotImplementedError

And tried out whether the custom ref envelope shows up, and it looks fine:

[pierrejeambrun]

Code looks good to me.

Just a few suggestions/nit but nothing blocking.

[amoghrajesh]

Amazing, thanks brent!