Let toolsets be built dynamically based on run context

[strawgate]

This PR adds a decorator that can be used to dynamically add toolsets to an Agent run similar to how the system prompt and instructions can be added to via decorated functions.

• Closes Dynamic toolsets based on run context #2331

[DouweM]

@strawgate Thanks! My main concern is that the function taking ctx suggests that a new toolset is going to be built when the context changes (on each run step) or at least when a new run is performed (and the ctx.deps would likely change), so that you can safely use run/user-specific details from ctx.deps like in your filesystem root directory example, without accidentally using a toolset built for another user/run.

But in reality it is tied to toolset enters/exits and there's no such guarantee... Those could align with agent runs, but wouldn't if you have your own async with agent: wrapping multiple agent runs (which you may want to do to keep MCP servers running during an entire request that could include multiple agent runs) because of the entered_count stuff we do in Agent.__aenter__ and CombinedToolset.__aenter__. Maybe we should just rip that out, so that we know for sure we enter/exit each toolset whenever a new agent run starts? Maybe we should also pass only deps instead of RunContext to make it clear you're not actually getting run-step specific context?

I want to make sure we don't accidentally introduce a footgun where the behavior will change completely when you add an async with agent: around a request or the entire FastAPI lifecycle or something, thinking you're just optimizing e.g. MCP server start/shutdowns.

Maybe we need to add a new hook to toolsets that's only called when a new run/step is actually started, so we don't rely on the enters/exits which could happen at any level?

This is all a bit complicated and my analysis here may very well be wrong. Either way I think we'll need tests with different patterns to make sure we get the desired behavior in each case.

We should probably chat about this directly, would you mind joining our public Slack so we can huddle?

[strawgate]

@DouweM i added the toolset decorator and added toolsetfunc to agent init, the tests will be passing shortly, I'll play around with this in my poc environment

[strawgate]

@DouweM the tests are passing and I've been using this on my fork for a little bit with good success, interested in your thoughts

[DouweM]

@strawgate Thank you, I like the decorator approach a lot. As I mentioned on our call, I want the default behavior to be that it gets rebuilt for every run step, as the presence of the RunContext suggests, so I implemented that and pushed into the PR.

We now have a new per_run_step argument on the decorator that you'll want to set to False, plus a DynamicToolset that tracks the toolset and handles enters/exits, so we're also sort of back to where this PR started :) Let me know what you think!

The only things missing now is a test to make sure the enter/exit behavior is correct, and docs.

[DouweM]

@strawgate Thanks William, I take it the new implementation has been working well for you?

[strawgate]

@strawgate Thanks William, I take it the new implementation has been working well for you?

Yes! I am using it here and here

I was mostly testing the per_run_step=true but just switched to per_run_step=false and am seeing good results.

[strawgate]

looks like unrelated test is failing

[DouweM]

@strawgate I made some tweaks to the docs, having this replace the WrapperToolset workaround for dynamic toolsets that encouraged some bad patterns. If this all looks good to you I think we can merge soon!

[strawgate]

Looks great

[HamzaFarhan]

@DouweM I feel like after this it makes even more sense for output_type to also accept toolsets
Right now if we have 10 output_types, and 5 toolsets with let's say 5-6 tools per toolset, we can add a toolset dynamically. In other words, add 5-6 tools with a single check. But we would need to add the output_types using if else in a single prepare function.
If they were toolsets, the same trigger could be used to load 5-6 tools and 2-3 corresponding output_types.

[DouweM]

@HamzaFarhan For proper type checking of the agent.result.output, the possible output types need to be statically known at agent definition time, which wouldn't work with toolsets. Type safety is pretty core to Pydantic (and AI's) philosophy, so that makes this a non-starter unfortunately.