tatiana opened a new issue, #75406:
URL: https://github.com/apache/airflow/issues/75406
### Apache Airflow version
3.3.2 (apache-airflow-task-sdk 1.3.2)
### What happened?
A task that writes to `context["task_state_store"]` from `on_retry_callback`
gets `ServerResponseError`. The execution API answers `404
{"detail":{"reason":"not_found","message":"Task instance <id> not found"}}`: by
the time the callback runs, the task instance has been moved to `up_for_retry`
and the next try gets a new TI id, so the store accessor the callback holds
points at an id the API no longer accepts.
The same write from the final `on_failure_callback` (no retries left)
succeeds, and writes from inside `execute` succeed, so the retry callback is
the only place a task cannot save state.
### What you think should happen instead?
AIP-103 describes the task state store as state that survives retries, so
the retry callback is the natural place to flush what a task collected before
it failed. Either the write should be accepted (the store is keyed per task,
not per try), or the task runner should run the retry callback before the state
transition, or the docs should say the callback cannot write the store.
Our use case: the Cosmos `ExecutionMode.WATCHER` producer backs up its XCom
pushes so a retry can restore them instead of re-running finished dbt models.
On Airflow 3.3 it keeps that backup in the task state store, written after
every push; when the last write fails, the entries since the previous
successful write exist only in memory, and the retry callback cannot flush
them. We fall back to an Airflow Variable for that tail, which the callback can
still write: https://github.com/astronomer/astronomer-cosmos/pull/3076
(discussion in
https://github.com/astronomer/astronomer-cosmos/pull/3076#discussion_r4339873944
and
https://github.com/astronomer/astronomer-cosmos/pull/3076#discussion_r4339859212).
### How to reproduce
Save as `dags/task_store_retry_callback.py` and run `airflow dags test
task_store_retry_callback`.
```python
from datetime import datetime, timedelta
from airflow.sdk import dag, task
def on_retry(context):
store = context["task_state_store"]
try:
store.set("flushed_from_retry_callback", "value")
print("CALLBACK write ok")
except Exception as error:
response = getattr(error, "response", None)
print("CALLBACK write failed:", type(error).__name__, error)
if response is not None:
print("CALLBACK response:", response.status_code, response.text)
def on_failure(context):
try:
context["task_state_store"].set("flushed_from_failure_callback",
"value")
print("FAILURE CALLBACK write ok")
except Exception as error:
print("FAILURE CALLBACK write failed:", type(error).__name__, error)
@dag(dag_id="task_store_retry_callback", start_date=datetime(2026, 1, 1),
schedule=None, catchup=False)
def task_store_retry_callback():
@task(retries=1, retry_delay=timedelta(seconds=1),
on_retry_callback=on_retry, on_failure_callback=on_failure)
def producer(**context):
context["task_state_store"].set("written_in_task", "value")
raise RuntimeError("boom")
producer()
task_store_retry_callback()
```
Output (SQLite, LocalExecutor defaults):
```
[info] Task instance state updated ... new_state=up_for_retry
rows_affected=1 ti_id=01a163b2-749f-74d6-be88-83692c4f883c
CALLBACK write failed: ServerResponseError Server returned error
CALLBACK response: 404 {"detail":{"reason":"not_found","message":"Task
instance 01a163b2-749f-74d6-be88-83692c4f883c not found"}}
...
FAILURE CALLBACK write ok
```
The state update to `up_for_retry` is logged before the callback runs, and
try 2 runs under a different TI id.
### Operating System
macOS 15 (Darwin 25.6.0)
### Versions of Apache Airflow Providers
Not relevant.
### Deployment
Other (local `airflow dags test`)
### Anything else?
Same result with `dag.test()`.
### Are you willing to submit PR?
- [ ] Yes I am willing to submit a PR!
### Code of Conduct
- [x] I agree to follow this project's Code of Conduct
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
