Decorate custom state refs with an envelope for UI clarity

[amoghrajesh]

---

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

• Yes - claude sonnet 4.6

What problem are we solving?

When a custom worker backend (e.g. S3, GCS) stores a state value externally and writes a reference string back to the DB, the UI has no way to tell whether a value like s3://bucket/ti_123/job_id is:

• The user's actual state value (a plain string they stored), or
• An opaque reference to externally-stored data

Without this distinction, the UI would show the raw path as if it were the value, which can be confusing and misleading.

Current behaviour

Custom backends return a reference string from serialize_task_state_to_ref(), which is stored verbatim in the DB. The DB value column contains either a plain JSON value or a reference string with no structural difference between the two — the UI cannot differentiate them.

Proposed change

When a custom worker backend is configured, the framework now automatically wraps the reference returned by serialize_task_state_to_ref() in a typed envelope before storing:

{"__airflow_state_ref__": "s3://bucket/ti_123/job_id"}

On read, the framework detects the envelope, extracts the ref, and passes it to deserialize_task_state_from_ref() and the backend never sees the envelope. If a stored value does not carry the marker (e.g. a corrupt row), the raw value is returned and a warning is logged.

The default path (no custom backend) is unaffected, plain JSON values are stored and returned as before.

UI Impact

UI PR for reference #67292

The UI reads state values directly from the DB and displays them as-is. With this change, when a custom backend is in use, the UI will show {"__airflow_state_ref__": "..."} instead of a raw reference string, making it visually clear that the value is a pointer to externally-stored data rather than the actual state value.

Testing

Created a custom worker side backend based on file system:

from __future__ importannotationsimportjsonfrompathlibimportPathfromtypingimportTYPE_CHECKINGfromairflow.sdk.stateimportBaseStateBackendifTYPE_CHECKING:
fromdatetimeimportdatetimefrompydanticimportJsonValuefromsqlalchemy.ext.asyncioimportAsyncSessionfromsqlalchemy.ormimportSessionfromairflow_shared.stateimportStateScopeBASE_DIR=Path("/tmp/airflow_state")
classFileStateBackend(BaseStateBackend):
"""Stores task/asset state values as local JSON files; returns the path as the ref."""defserialize_task_state_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_state_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: StateScope, key: str, *, session: Session|None=None) ->str|None:
raiseNotImplementedError(
"FileStateBackend is a worker-side backend; server uses MetastoreStateBackend"
)
defset(
self,
scope: StateScope,
key: str,
value: str,
*,
expires_at: datetime|None=None,
session: Session|None=None,
) ->None:
raiseNotImplementedErrordefdelete(self, scope: StateScope, key: str, *, session: Session|None=None) ->None:
raiseNotImplementedErrordefclear(
self, scope: StateScope, *, all_map_indices: bool=False, session: Session|None=None
) ->None:
raiseNotImplementedErrorasyncdefaget(self, scope: StateScope, key: str, *, session: AsyncSession|None=None) ->str|None:
raiseNotImplementedErrorasyncdefaset(
self,
scope: StateScope,
key: str,
value: str,
*,
expires_at: datetime|None=None,
session: AsyncSession|None=None,
) ->None:
raiseNotImplementedErrorasyncdefadelete(self, scope: StateScope, key: str, *, session: AsyncSession|None=None) ->None:
raiseNotImplementedErrorasyncdefaclear(
self, scope: StateScope, *, all_map_indices: bool=False, session: AsyncSession|None=None
) ->None:
raiseNotImplementedError

Ran breeze with: export AIRFLOW__WORKERS__STATE_BACKEND=file_state_backend.FileStateBackend

DAG:

@dag(schedule=None, start_date=datetime(2026, 4, 23), catchup=True)defsimple_task_state():
@taskdefmy_task(**context: Context):
task_state=context["task_state"]
task_state.set("job_id", "12345")
task_state.set("secret-dict", {"key": "value"})
task_state.set("int_value", 42)
print("Fetching task states I stored earlier")
print("job_id:", task_state.get("job_id"), type(task_state.get("job_id")))
print("secret-dict:", task_state.get("secret-dict"), type(task_state.get("secret-dict")))
print("int_value:", task_state.get("int_value"), type(task_state.get("int_value")))
my_task()

The task is agnostic to the envelope.

File system is updated with the custom backend generated files:

[Breeze:3.10.20] root@78d3c819fefd:/tmp/airflow_state/ti_019e6821-0661-7bdc-ad37-8dad5ac1fa14$ pwd
/tmp/airflow_state/ti_019e6821-0661-7bdc-ad37-8dad5ac1fa14
[Breeze:3.10.20] root@78d3c819fefd:/tmp/airflow_state/ti_019e6821-0661-7bdc-ad37-8dad5ac1fa14$ ll
bash: ll: command not found
[Breeze:3.10.20] root@78d3c819fefd:/tmp/airflow_state/ti_019e6821-0661-7bdc-ad37-8dad5ac1fa14$ ls -l
total 12
-rw-r--r-- 1 root root 2 May 27 06:30 int_value.json
-rw-r--r-- 1 root root 7 May 27 06:30 job_id.json
-rw-r--r-- 1 root root 16 May 27 06:30 secret-dict.json
[Breeze:3.10.20] root@78d3c819fefd:/tmp/airflow_state/ti_019e6821-0661-7bdc-ad37-8dad5ac1fa14$ cat int_value.json
42[Breeze:3.10.20] root@78d3c819fefd:/tmp/airflow_state/ti_019e6821-0661-7bdc-ad37-8dad5ac1fa14$ cat job_id.json
"12345"[Breeze:3.10.20] root@78d3c819fefd:/tmp/airflow_state/ti_019e6821-0661-7bdc-ad37-8dad5ac1fa14$ cat secret-dict.json
{"key": "value"}[Breeze:3.10.20] root@78d3c819fefd:/tmp/airflow_state/ti_019e6821-0661-7bdc-ad37-8dad5ac1fa14$

Core API will return the external reference for UI to build upon, the extra encoding like {\"__airflow_state_ref__\": \"/tmp/airflow_state/ti_019e6821-0661-7bdc-ad37-8dad5ac1fa14/int_value.json\"} will be fixed by: #67547

---

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

[kaxil]

Thanks for the thorough follow-ups -- all the points from the last pass are addressed. The move off the __type/__var serialization namespace to __airflow_state_ref__ removes the collision, the _wrap/_unwrap helpers keep the wire format in one place, and the warning-log on the mismatch path is a good safety net. LGTM.

[kaxil]

Rebased your PR on main since it had unrelated spelling failures

[amoghrajesh]

Thanks for rebasing, @kaxil. I am merging this one in now