Fix common.ai durable execution skipping Toolset-capability tools#69881
Open
kaxil wants to merge 1 commit into
Open
Fix common.ai durable execution skipping Toolset-capability tools#69881kaxil wants to merge 1 commit into
common.ai durable execution skipping Toolset-capability tools#69881kaxil wants to merge 1 commit into
Conversation
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.
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
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.
Summary
With
durable=True,AgentOperatorcaches 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'stoolsets=list, or a pydantic-aiToolsetcapability (agent_params={"capabilities": [Toolset(ts)]}). Only the first was wrapped withCachingToolset, so tools supplied via aToolsetcapability 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
Toolsetcapability with the sameCachingToolsetthe operator already applies totoolsets=, 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'stoolsets=path: no new keying scheme, no dependence on hook ordering.dataclasses.replace(cap, toolset=...)preserves the capability's other fields, and theCachingToolsetsurvives pydantic-ai's per-runfor_runcloning.Two latent durable-execution bugs fixed alongside
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.){dag}_{task}_{run}, so dagetl+ taskload_dataand dagetl_load+ taskdatamapped 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 concreteToolsetcapability. Tools provided through any other capability (MCP,PrefixTools,CombinedCapability, aToolsetbacked by a callable factory, or capabilities loaded from aspec_file) are not cached and re-run on retry; put tools you need replayed intoolsets=. A factory-backedToolsetadditionally logs a warning. Provider-native capabilities likeWebSearchandThinkingexecute inside the model call and are already covered by model-response caching.