kaxil opened a new pull request, #76877:
URL: https://github.com/apache/airflow/pull/76877
`AgentOperator(durable=True)` and `@task.agent(durable=True)` refused to
load a Dag that passed a `FunctionToolset` without an `id`, or a function that
builds a toolset for each run:
```
ValueError: durable=True needs a unique id on every FunctionToolset, such as
FunctionToolset(..., id='orders'): ...
```
pydantic-ai names a durable step after the toolset's id, so the operator
asked the Dag author for one. That makes an ordinary agent fail as soon as
`durable=True` is added, and it blocks turning durable execution on by default
later in this stack. With this PR the operator fills in the ids itself, and
this agent needs none:
```python
AgentOperator(
task_id="durable_order_status",
prompt="Where is order 4411?",
llm_conn_id="pydanticai_default",
durable=True,
toolsets=[FunctionToolset(tools=[get_order]), carrier_tools],
)
```
- An id-less `FunctionToolset` is handed to the agent as a copy with an id
numbered in the order the toolsets are passed (`<airflow-function_toolset-0>`).
Its tool calls are recorded under the tool's name alone,
`durable_order_status__function_toolset.call_tool:get_order`, so adding a tool
to the toolset never renames another tool's step.
- A function that builds a toolset (and an id-less `DynamicToolset`) gets
its module and qualified name, `example_agent_durable.carrier_tools`. The
`unusual_prefix_<sha1 of the file path>_` that Airflow puts in a Dag file's
module name is left out, because the path can differ between workers and bundle
versions.
- A toolset whose id the Dag author chose keeps it in its step names,
exactly as before.
This covers toolsets passed in `toolsets=`, `agent_params["toolsets"]`, and
a `FunctionToolset` held directly by a `Toolset` capability.
`AirflowDurability` used in a plain `@task` is unchanged: there pydantic-ai
still raises for an id-less toolset.
Still refused when the Dag loads:
- a pydantic-ai `MCPToolset` without an id. No id is derived from the
server's URL or command, which can carry credentials. The old check tested
`leaf.id is None`, but pydantic-ai 2.62+ derives `.id` from an HTTP server's
hostname, so an id-less HTTP `MCPToolset` passed the load-time check. The check
now reads the id the author set (`_explicit_id`, falling back to `.id` before
2.62).
- a callable without a `__qualname__` (a `functools.partial`, a callable
object): pass it as `DynamicToolset(function, id=...)`.
- two functions with the same qualified name, such as two lambdas in one Dag
file, since their calls would be filed under one name.
## Design rationale
Renaming steps alone wasn't enough. pydantic-ai's durability capability
looks its durable toolsets up by id when the agent is built and swaps them in
by id on every run, so an id-less toolset still hits its `UserError`. The
operator sets `_id` on a copy instead (`_render_toolsets` already copies
toolsets to set rendered fields on them; the Dag author's object isn't
touched). The other options I weighed were overriding `_register_toolsets`,
`_wrap_and_register_leaf` and `get_wrapper_toolset` on
`BaseDurabilityCapability`, which means three private methods to keep in step
with pydantic-ai from 2.46 to 2.66, and journaling function tools in Airflow's
own wrapper, which renames the steps of every toolset that already has an id.
The old check refused derived ids on purpose, so that reordering the list
couldn't point a retry at another toolset's recorded calls. The number Airflow
assigns stays out of step names (an Airflow `JournalOperationNamer` drops it),
and pydantic-ai refuses two tools with the same name in one agent, so a step
named after the tool can't match the wrong toolset. Reordering or adding
toolsets changes the tools the model is offered, and the retry runs again from
the first model request, as it does today when someone edits a toolset that has
an id. My first version used a random id, which broke replay for a toolset with
instructions: pydantic-ai copies the toolset id into tool definitions and
instruction ids, and the request fingerprint hashes both. There's no colon in
the id because pydantic-ai rejects one on a toolset with instructions.
Old journals: the namer only strips ids Airflow assigned, and every toolset
journaled before this change had an id the author chose, because id-less ones
were refused at load. I ran the previous code against a JSON-file journal until
it recorded `orders_agent__function_toolset__orders.call_tool:lookup` and
failed on the next model request, then retried with this code over the same
file. The model request and `lookup` replayed (0 live `lookup` calls) and only
the final request ran live.
## End to end
A durable `AgentOperator` with a real model (`openai:gpt-5.4-mini` through a
gateway), `retries=1`, an id-less `FunctionToolset` with `get_order` and
`post_update`, and a `carrier_tools` function that builds a `FunctionToolset`
with `track_parcel`. `post_update` raises on try 1.
Try 1 runs all three tools live, then fails in `post_update`:
Try 2 replays `get_order` and `track_parcel` (returned in 0.00s, no `TOOL
RAN LIVE` line) and the recorded model requests, then runs `post_update` and
the final request:
```
try 1: Durable: replayed 0 steps (0 model, 0 tool, 0 other), recorded 8 new
steps (3 model, 2 tool, 3 other)
try 2: Durable: replayed 8 steps (3 model, 2 tool, 3 other), recorded 3 new
steps (1 model, 1 tool, 1 other)
```
## Known issues
A `FunctionToolset` nested deeper in the capabilities, inside
`CombinedCapability([...])`, `PrefixTools(wrapped=Toolset(...))` or a
`Toolset(build_fn)`, still doesn't get an id: the Dag loads and every durable
attempt fails with pydantic-ai's `UserError`, the same as before this PR. The
docs only promise the three places above; the operator would have to walk
`WrapperCapability.wrapped` and `CombinedCapability.capabilities` to reach the
rest.
--
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]
