Skip to content

Fix common.ai durable execution skipping Toolset-capability tools#69881

Open
kaxil wants to merge 1 commit into
apache:mainfrom
astronomer:durable-caching-toolset-capabilities
Open

Fix common.ai durable execution skipping Toolset-capability tools#69881
kaxil wants to merge 1 commit into
apache:mainfrom
astronomer:durable-caching-toolset-capabilities

Conversation

@kaxil

@kaxil kaxil commented Jul 14, 2026

Copy link
Copy Markdown
Member

Summary

With durable=True, AgentOperator caches each model response and tool result so a retried task replays completed steps instead of re-running them. Tools reach the agent two ways: the operator's toolsets= list, or a pydantic-ai Toolset capability (agent_params={"capabilities": [Toolset(ts)]}). Only the first was wrapped with CachingToolset, so tools supplied via a Toolset capability were re-executed on every retry instead of replaying, silently defeating durability for that path (a side-effecting tool would run again).

This wraps the inner toolset of a concrete Toolset capability with the same CachingToolset the operator already applies to toolsets=, so both paths get identical replay semantics.

Why this approach

The wrapping happens at the toolset layer (CachingToolset.call_tool), reusing the existing caching + step-counter + fingerprint machinery, rather than re-expressing durability as a pydantic-ai capability hook (wrap_tool_execute). Caching stays below the capability middleware, so it inherits the exact replay semantics and edge-case behavior of today's toolsets= path: no new keying scheme, no dependence on hook ordering. dataclasses.replace(cap, toolset=...) preserves the capability's other fields, and the CachingToolset survives pydantic-ai's per-run for_run cloning.

Two latent durable-execution bugs fixed alongside

  • Cleanup ran too early. cleanup() deleted the cache right after the run, before the message-history XCom push and output serialization. If either failed, the retry started with an empty cache and re-ran every completed step. Cleanup now runs after those steps. (If a post-run step fails on every attempt the cache lingers until the retention sweep, an accepted tradeoff, and the same orphan class already existed for run failures.)
  • Pre-3.3 cache filename aliased distinct tasks. The ObjectStorage backend named its file {dag}_{task}_{run}, so dag etl + task load_data and dag etl_load + task data mapped to the same file, letting one task read, overwrite, or delete another's cache. The filename is now a hash of the identity components.

Coverage / limitations

Durable caching covers tools from toolsets= and from a concrete Toolset capability. Tools provided through any other capability (MCP, PrefixTools, CombinedCapability, a Toolset backed by a callable factory, or capabilities loaded from a spec_file) are not cached and re-run on retry; put tools you need replayed in toolsets=. A factory-backed Toolset additionally logs a warning. Provider-native capabilities like WebSearch and Thinking execute inside the model call and are already covered by model-response caching.

When durable=True, tools supplied via a pydantic-ai Toolset capability
(capabilities=[Toolset(ts)]) bypassed the CachingToolset applied to
toolsets=, so their results re-executed on every retry instead of
replaying. Wrap the inner toolset of concrete Toolset capabilities with
the same CachingToolset. A Toolset capability backed by a callable
factory can't be wrapped (it resolves per run) and now logs a warning.

Also fixes two latent durable-execution bugs found in review:

- cleanup() ran before the message-history XCom push and output
  serialization. A failure in either wiped the cache, so the retry
  re-ran every already-completed step. Move cleanup to after them.

- the pre-3.3 ObjectStorage cache filename joined dag/task/run with "_",
  aliasing distinct tasks (dag "etl"/task "load_data" and dag
  "etl_load"/task "data" both mapped to one file, so one task could
  read, overwrite, or delete another's cache). Hash the identity
  components instead.
@kaxil kaxil marked this pull request as ready for review July 15, 2026 13:57
@kaxil kaxil requested a review from gopidesupavan as a code owner July 15, 2026 13:57
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant