[ColtenOuO]

Summary

The task state store is a per-task-instance key/value store that survives worker
crashes and retries within a Dag run — the place operators keep external job IDs,
checkpoints and progress metadata. Its REST API is addressed by four values that
together identify one task instance: dag_id + dag_run_id + task_id + map_index.

The two DELETE endpoints never checked that the task instance those four values
name actually exists, so they reported success while touching nothing. All six
endpoints declared 404 in create_openapi_http_exception_doc(...) — that doc is
maintained by hand and had drifted from the handlers in both directions, promising a
404 the deletes could not raise and one the list endpoint has no way to raise
either.

_require_ti() was written as a plain helper that each handler had to remember to
call, and it was called from PUT and PATCH only. The delete handlers went
straight to the query, and because those queries filter the four columns on
task_state_store directly — they never join task_instance — a request naming a
task instance that doesn't exist simply matches zero rows. At the SQL level that is
indistinguishable from "this task instance exists but has no state stored", so the
delete succeeded.

What that looked like

A typo in task_id made a delete report success while touching nothing:

$ curl -X DELETE .../taskInstances/proces/state-store    # note the missing 's'
204 No Content

Clearing state store is typically a pre-rerun cleanup step — clearing a task
instance in the UI deliberately does not wipe its state store, so a stale
checkpoint survives into the next attempt. An operator wiping a bad external job ID
before re-running gets 204, believes the state is gone, re-runs, and the task
resumes from exactly the state they thought they had removed. Nothing in the
response hints that the task_id was wrong.

The same resource also answered the same question two different ways depending on
the verb — a write refused, the matching delete claimed success:

$ curl -X PUT    .../taskInstances/proces/state-store/job_id -d '{"value":"x"}'
404 {"detail": "Task instance not found for dag_id='etl', ..., task_id='proces', map_index=-1"}

$ curl -X DELETE .../taskInstances/proces/state-store/job_id
204 No Content

Before / after

Endpoint	Before	After
DELETE .../state-store/{key}	204	404
DELETE .../state-store	204	404
PUT .../state-store/{key}	404	unchanged
PATCH .../state-store/{key}	404	unchanged
GET .../state-store	200 + empty collection	unchanged (spec drops its 404)
GET .../state-store/{key}	404, message names the key	unchanged

A missing key on an existing task instance still reports the key, unchanged:
Task state store key 'job_id' not found.

The reads are deliberately left alone (per review): an empty collection and a
"key not found" are already honest answers, and they would otherwise pay for an
extra task_instance lookup on every read. Since the list endpoint therefore has
no 404 path at all — requires_access_dag only raises 401/403, and
paginated_select raises nothing here — its create_openapi_http_exception_doc
entry is dropped, matching the other list endpoints (get_pools, get_variables,
…) that declare no error responses.

The fix

asset_state_store.py — the sibling module from the same feature — already solves
this correctly: _get_asset_or_404 is wired up as AssetIdDep = Annotated[int, Depends(...)] and every one of its endpoints takes asset_id: AssetIdDep, so the
check rides on the parameter and cannot be forgotten when an endpoint is added.

This applies the same pattern on the task side, split in two so the reads do not
pay for a check they don't need:

• TaskScopeDep maps the path and query parameters onto the TaskScope the
  handlers already needed, absorbing the old _get_scope() helper. Every endpoint
  takes it, so each handler drops from four path/query parameters to a single
  scope.
• ValidatedTaskScopeDep builds on it and adds the existence check. Only the
  writes and deletes take it.

Why clear needs its own dependency

DELETE .../state-store?all_map_indices=true means "clear this task at every map
index", and documents that map_index is ignored. ValidatedTaskScopeDep validates
one exact map index, defaulting to -1 — and an expanded mapped task has no
map_index=-1 instance: TaskMap.expand_mapped_task either converts that row into
map_index=0 or deletes it. Validating with ValidatedTaskScopeDep would therefore
reject a perfectly valid request:

# 'process' expanded to map_index 0, 1, 2 — there is no map_index=-1 row
$ curl -X DELETE .../taskInstances/process/state-store?all_map_indices=true
404 Task instance not found ... map_index=-1     # wrong

A dependency only receives the parameters it declares, so ValidatedTaskScopeDep
cannot see all_map_indices and can only ever take the strict path.
ValidatedClearTaskScopeDep declares the flag, and when it is set drops the
map_index predicate so any instance of that task in the run satisfies the check.
The 404 it raises names the flag rather than a map_index that was never applied:

$ curl -X DELETE .../taskInstances/proces/state-store?all_map_indices=true
404 {"detail": "Task instance not found for dag_id='etl', run_id='...', task_id='proces', all_map_indices=True"}

The strict path is untouched: ?map_index=7 on a task expanded to 0..2 still
returns 404.

Generated files

v2-rest-api-generated.yaml and ui/openapi-gen/ are regenerated. Two kinds of
change: the 404 removed from GET .../state-store, and parameter ordering —
moving path/query parameters behind a dependency changes the order FastAPI reports
them in, not the parameters themselves.

Tests

New TestUnknownTaskInstance covers every request shape against both a bad
task_id and a bad run_id, a non-existent map_index, the all_map_indices
error detail, and regression tests that the reads still answer as they did — the
list returns an empty collection, and a genuinely missing key still reports the key
rather than the task instance.

76 → 96 tests in this file; test_asset_state_store.py passes unchanged.

---

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

• Yes — Claude Code (Opus 5)

[jason810496]

Thanks for the fix.

It seems only those route calling the _require_ti helper need those new deps. Perhaps we can introduce different dependency injection helper (e.g. one for mapping the path parameters to the TaskScope another one leverage the first one but also validate the TI).

[ColtenOuO]

Thanks for reviewing.

I've addressed all comments.

[jason810496]

Thanks! LGTM overall after these comments.

We will need to update the PR description as well, since we don't validate those task_state_store routes now (list_task_state_store, get_task_state_store ).

[ColtenOuO]

Thanks for reviewing.

I've changed the code to show the all_map_indices=True in the error.

By the way, I noticed that list endpoint can never return 404 now that it doesn't validate the task instance, so I dropped its create_openapi_http_exception_doc entry — same as get_pools and get_variables. PR description is updated as well.