kaxil opened a new pull request, #69881:
URL: https://github.com/apache/airflow/pull/69881
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 (which rebuilds via `replace(..., wrapped=...)`).
A `Toolset` capability backed by a callable factory (resolved per run) has
no concrete toolset to wrap at build time, so it is left unwrapped and logs a
warning. Pass such toolsets via `toolsets=` for durability.
## 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.
- **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.
## Gotchas / tradeoffs
- Factory-backed `Toolset` capabilities are not durably cached (warned at
build time; documented on the `durable` param).
- Changing the pre-3.3 cache filename means a task that cached under the old
name and retries after this upgrade gets a one-time cache miss (re-runs its
steps). It is bounded and one-time. The task-state-store backend (Airflow 3.3+)
is unaffected; its keys are unchanged.
## Follow-up (not in this PR)
If a post-run step fails deterministically on every attempt, the cache is
never cleaned up and lingers (the task-state-store `clean` sweep / TTL is the
backstop). This pre-dates the PR (cleanup already only ran on success); a
targeted "cleanup on final attempt" is a separate change.
--
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]
