kaxil opened a new pull request, #76864:
URL: https://github.com/apache/airflow/pull/76864
A durable agent task replays every step the failed attempt completed, tool
calls included. That is wrong for some tools: one whose answer has to be
current (today's date, the state of a job), and one whose effect does not
outlive the attempt (a file in a temporary directory, a session on the worker).
Until now the only switch was `AirflowToolset.replayable`, which covers a whole
toolset and only Airflow's own.
This PR lets a single tool opt out with pydantic-ai's existing per-tool
mechanism, the tool metadata key of the durability engine:
`metadata={"airflow": False}` on the tool, or `SetToolMetadata(tools=[...],
airflow=False)` for tools of a toolset you don't build yourself, such as
`SQLToolset`. On a retry that tool runs live, and the model requests and other
tool calls still replay. pydantic-ai's Temporal integration documents the same
thing as `metadata={"temporal": False}`.
```python
AgentOperator(
task_id="durable_freshness_check",
prompt="Which tables in load_log have not been loaded today?",
llm_conn_id="pydanticai_default",
durable=True,
toolsets=[
FunctionToolset(tools=[Tool(today, metadata={"airflow": False})],
id="calendar"),
SQLToolset(db_conn_id="postgres_default",
allowed_tables=["load_log"]),
],
capabilities=[SetToolMetadata(tools=["query"], airflow=False)],
)
```
`True` or no key keeps the default. Any other value (`"false"`, `0`) fails
with a `UserError` that names the tool, so a typo doesn't silently keep a tool
replaying.
## Design rationale
**Both toolset paths read the same key.** For pydantic-ai's function and
dynamic toolsets, `AirflowDurability`'s operation backend now passes a
`RoleBasedOperationConfig(resolve_tool=...)` (it passed one with no resolver
before), which pydantic-ai calls with the tool definition handed down to the
call. For Airflow's own toolsets (`SQLToolset`, `HookToolset`, the provider's
`MCPToolset`, and any other toolset pydantic-ai doesn't run as a durable unit),
the journaling wrapper reads the same key. `AirflowToolset.replayable = False`
still opts a whole toolset out.
**An opted-out call keeps its journal position but is never recorded or
replayed.** pydantic-ai's own way to opt out is to have the resolver return
`False`, which takes the call out of the backend entirely. That would shift the
position of every step after the tool whenever someone turned the opt-out on or
off between attempts, and a failure in the tool would no longer belong to a
step, so the rule from the PR below this one (a tool call that fails twice in a
row with the same arguments makes the next retry ask the model again) would
stop applying and the retry would replay the same choice into the same error
until retries ran out. Instead the resolver returns a "live" config and the
backend claims the step with `replayable=False`, the path
`AirflowToolset.replayable` already uses. A test pins this: with the `False`
bypass, attempt 3 repeats the failure.
**The key is not part of the model request fingerprint.** Tool metadata is
hashed with the tool definitions, so adding `SetToolMetadata(airflow=False)`
after a stale replay would have made every model request miss and re-run the
whole agent, side effects included. The `airflow` key is now dropped before
hashing, and empty metadata hashes like none. Tools without the key hash as
before, so existing journals still match. One exception: a tool declared with
an explicit `metadata={}` hashes differently from before, so a retry that
crosses the upgrade runs that tool's following requests live once.
**MCP:** pydantic-ai refuses the opt-out on its own
`pydantic_ai.mcp.MCPToolset` ("MCP tools perform I/O and cannot run outside a
durable step"), keyed on the engine name, so it applies to Airflow too. I don't
think the reason holds for Airflow, where nothing runs outside the task
process, but lifting it would mean overriding a private pydantic-ai method, so
this PR leaves it and documents it. The provider's `MCPToolset` is an
`AirflowToolset` and accepts the opt-out.
## pydantic-ai-slim floor: 2.36.0 → 2.46.0
Before 2.46, `CombinedToolset.call_tool` passes the tool definition it
cached at `get_tools()` time, not the one it was handed, so metadata added by a
wrapper above it (`SetToolMetadata`'s `PreparedToolset`) never reaches the
leaf. On 2.36 to 2.45 the `SetToolMetadata` form would be silently ignored and
the tool would keep replaying. Metadata set on the tool itself works on every
version, but `SQLToolset`'s tools can only be reached through
`SetToolMetadata`. The lock already resolves 2.63, so `uv.lock` changes only in
its specifier lines. The whole common.ai suite passes at the new floor
(pydantic-ai-slim 2.46.0, pydantic-ai-harness 0.35.0: 4,670 passed).
## End to end
A durable `AgentOperator` with a real model (`openai:gpt-5.4-mini` through a
gateway), `retries=1`, a `FunctionToolset` with `today` (`metadata={"airflow":
False}`) and `publish_report`, and a `SQLToolset` over SQLite with
`SetToolMetadata(tools=["query"], airflow=False)`. `publish_report` raises on
try 1.
Try 1 runs `today`, `get_schema`, `query` and `publish_report`, which raises:
Try 2 runs `today` and `query` again (their own log line and the hook's
`Running statement:`), replays the four model requests and `get_schema`
(returned in 0.00s), then runs `publish_report` and the final request:
```
Durable: replayed 5 steps (4 model, 1 tool, 0 other), recorded 2 new steps
(1 model, 1 tool, 0 other)
```
## Known issues
A step without a fingerprint, such as another capability's
`@durable_operation`, replays only right after a step that replayed. After a
live call, a compaction summary therefore runs again even when the live result
didn't change, and the summary it writes can make the steps after it run again.
The docs say so. `ManagedAgentToolset` (`replayable=False`) has had the same
behaviour; this PR makes it reachable from any tool. Fingerprinting the
arguments of capability operations would fix it and is left for a follow-up.
--
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]
