Use task state store for common.ai durable execution on Airflow 3.3+#68926
Merged
Conversation
kaxil
changed the title
common.ai durable execution on Airflow 3.3+
amoghrajesh
approved these changes
Jun 24, 2026
kaxil
marked this pull request as draft
June 24, 2026 13:04
durable=True caching is per-task-instance, per-run, and cleared on success -- exactly the scope of the AIP-103 task state store. On Airflow >= 3.3 the cache now lives there instead of an ObjectStorage JSON file, so durable execution no longer needs the [common.ai] durable_cache_path config, and large step payloads are offloaded automatically through the configured worker state store backend. The ObjectStorage backend is kept as the fallback for Airflow < 3.3, selected at runtime by the operator. Both backends share one storage interface, so the caching model/toolset wrappers are unchanged.
…pelling
Three CI failures on this branch, all in the durable execution code:
- Compat 3.0.6/3.1.8/3.2.2: ``test_task_state_store`` used
``pytest.importorskip("airflow.sdk.execution_time.context")`` to skip on
pre-3.3 cores, but that module exists on older cores -- only the
``NEVER_EXPIRE`` name it pulls in (via ``task_state_store``) is 3.3+. Collection
blew up before the skip could fire. Gate on ``AIRFLOW_V_3_3_PLUS`` instead,
matching how ``AgentOperator._build_durable_storage`` guards the import.
- Non-DB / Low-dep core: ``test_providers_modules_should_have_tests`` flagged the
new ``durable/base.py`` as having no test. Add ``test_base.py`` covering the
reserved-key invariants and the ``DurableStorageProtocol`` runtime contract.
- Docs spellcheck: "sizeable" -> "sizable" in the agent operator docs.
- task_state_store docstring: "older cores" -> "older airflow versions". - test_agent: skip the task-state-store build test on < 3.3 instead of patching AIRFLOW_V_3_3_PLUS -- the inline 3.3-only import would otherwise fail on the compat runners once collection is fixed. - test_agent: give the accessor a TaskStateStoreAccessor spec and the task instance an attribute spec; hoist the version-agnostic DurableStorage import. - test_base: assert the two shipped backends (ObjectStorage always, task state store on 3.3+) satisfy DurableStorageProtocol, to catch method drift.
- save_tool_result: normalize results through JSON before writing. The store validates against pydantic JsonValue (stricter than json.dumps -- rejects tuples and non-string dict keys), so a tool returning `(value, meta)` or an int-keyed dict previously crashed a step that already succeeded. Normalizing (tuple->list, non-str keys->str) matches the <3.3 ObjectStorage backend and keeps the step durable. - save_model_response / save_tool_result: make the store write best-effort -- a failed write (e.g. an oversize value on the metadata DB) is skipped with a warning instead of failing an already-succeeded step and self-poisoning retry. - cleanup: log delete failures instead of silently suppressing, so a value orphaned in external storage is at least visible. - test: make FakeTaskStateStore mirror the real contract (reject None, validate JsonValue) and add tuple / non-string-key regression tests it now catches. - test_base: trim to the real-backend protocol-conformance tests + key-prefix invariant; drop the tests that only exercised typing.Protocol machinery. - docs: note durable keys carry the reserved DURABLE_KEY_PREFIX.
kaxil
force-pushed
the
common-ai-durable-task-state-store
branch
from
July 7, 2026 05:27
c3f52b0 to
db57c02
Compare
kaxil
marked this pull request as ready for review
July 7, 2026 05:29
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
durable=TrueonAgentOperator/@task.agentcaches each model response and toolresult so a retry replays completed steps instead of re-running expensive LLM calls.
Until now that cache was always an ObjectStorage JSON file, which required
[common.ai] durable_cache_pathto be set.The durable cache is per-task-instance, per-run, and cleared on success, which is exactly
the scope of the AIP-103 task state store.
On Airflow >= 3.3 the cache now lives there:
durable_cache_pathconfiguration needed.[workers] state_store_backendis set.On Airflow < 3.3 the ObjectStorage backend is unchanged and selected automatically.
Design
DurableStorageProtocol, soCachingModelandCachingToolsetdepend on the interface, not a concrete backend.AgentOperator._build_durable_storagepicks the backend by Airflow version.__commonai_durable__...) so they cannot collidewith keys user code writes via
context["task_state_store"]in the same task instance.NEVER_EXPIREso a retry can replay them regardless ofretry_delayor the global retention config;cleanup()removes only the keys this run touched.Gotchas
state_store_backend), cached step values are stored inthe metadata database. For agents with large responses, configure
[workers] state_store_backendto offload them to external storage.agent_params={"tools": [...]}is not wrapped by the caching toolset(only
toolsets=are), so only its model steps replay. This is existing behavior, unchanged here.Verification
Verified end-to-end on a real Airflow 3.3 standalone: a durable
AgentOperatorwhose toolfails on the first attempt replays the cached model step from the task state store on the
retry, and a direct round-trip task confirms cache, replay, and cleanup against the live
Execution API and metadata DB.