Split out dynamic toolset handling

Author: strawgate

pydantic_ai_slim/pydantic_ai/agent.py

@@ -426,10 +426,12 @@ def __init__(
             self._output_toolset.max_retries = self._max_result_retries
 
         self._function_toolset = FunctionToolset(tools, max_retries=retries)
-        self._user_toolsets = [
-            toolset if isinstance(toolset, AbstractToolset) else DynamicToolset[AgentDepsT](toolset_func=toolset)
+        self._dynamic_toolsets = [
+            DynamicToolset[AgentDepsT](toolset_func=toolset)
             for toolset in toolsets or []
+            if not isinstance(toolset, AbstractToolset)
         ]
+        self._user_toolsets = [toolset for toolset in toolsets or [] if isinstance(toolset, AbstractToolset)]
 
         self.history_processors = history_processors or []
 
@@ -1680,7 +1682,7 @@ async def simple_toolset(ctx: RunContext[str]) -> AbstractToolset[str]:
         """
 
         def toolset_decorator(func_: ToolsetFunc[AgentDepsT]) -> ToolsetFunc[AgentDepsT]:
-            self._user_toolsets.append(DynamicToolset(func_, per_run_step=per_run_step))
+            self._dynamic_toolsets.append(DynamicToolset(func_, per_run_step=per_run_step))
             return func_
 
         return toolset_decorator if func is None else toolset_decorator(func)
@@ -1747,10 +1749,9 @@ def _get_toolset(
         else:
             user_toolsets = self._user_toolsets
 
-        dynamic_toolsets = [toolset.copy() for toolset in user_toolsets if isinstance(toolset, DynamicToolset)]
-        static_toolsets = [toolset for toolset in user_toolsets if not isinstance(toolset, DynamicToolset)]
+        dynamic_toolsets = [toolset.copy() for toolset in self._dynamic_toolsets]
 
-        all_toolsets = [self._function_toolset, *static_toolsets, *dynamic_toolsets]
+        all_toolsets = [self._function_toolset, *user_toolsets, *dynamic_toolsets]
 
         if self._prepare_tools:
             all_toolsets = [PreparedToolset(CombinedToolset(all_toolsets), self._prepare_tools)]

pydantic_ai_slim/pydantic_ai/toolsets/_dynamic.py

@@ -73,4 +73,4 @@ def apply(self, visitor: Callable[[AbstractToolset[AgentDepsT]], None]) -> None:
             self._toolset.apply(visitor)
 
     def copy(self) -> _DynamicToolset[AgentDepsT]:
-        return _DynamicToolset(toolset_func=self.toolset_func, per_run_step=self.per_run_step)
+        return _DynamicToolset(toolset_func=self.toolset_func, per_run_step=self.per_run_step)

tests/test_toolsets.py

@@ -16,7 +16,7 @@
 from pydantic_ai.messages import ToolCallPart
 from pydantic_ai.models.test import TestModel
 from pydantic_ai.tools import ToolDefinition
-from pydantic_ai.toolsets._dynamic import _DynamicToolset as DynamicToolset
+from pydantic_ai.toolsets._dynamic import _DynamicToolset as DynamicToolset  # pyright: ignore[reportPrivateUsage]
 from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool
 from pydantic_ai.toolsets.combined import CombinedToolset
 from pydantic_ai.toolsets.filtered import FilteredToolset
@@ -500,6 +500,25 @@ async def test_context_manager_failed_initialization():
     assert server1.is_running is False
 
 
+
+async def test_tool_manager_reuse_self():
+    """Test the retry logic with failed_tools and for_run_step method."""
+
+    run_context = build_run_context(None, run_step=1)
+
+    tool_manager = ToolManager[None](run_context, FunctionToolset[None](), tools={})
+
+    same_tool_manager = await tool_manager.for_run_step(ctx=run_context)
+
+    assert tool_manager is same_tool_manager
+
+    step_2_context = build_run_context(None, run_step=2)
+
+    updated_tool_manager = await tool_manager.for_run_step(ctx=step_2_context)
+
+    assert tool_manager != updated_tool_manager
+
+
 async def test_tool_manager_retry_logic():
     """Test the retry logic with failed_tools and for_run_step method."""
 
@@ -654,7 +673,7 @@ async def get_tools(self, ctx: RunContext[None]) -> dict[str, ToolsetTool[None]]
         async def call_tool(
             self, name: str, tool_args: dict[str, Any], ctx: RunContext[None], tool: ToolsetTool[None]
         ) -> Any:
-            return None
+            return None  # pragma: no cover
 
     def toolset_factory(ctx: RunContext[None]) -> AbstractToolset[None]:
         return EnterableToolset()
@@ -686,8 +705,8 @@ def visitor(toolset: AbstractToolset[None]) -> None:
 
     assert tools == {}
 
-async def test_dynamic_toolset_empty():
 
+async def test_dynamic_toolset_empty():
     def no_toolset_func(ctx: RunContext[None]) -> None:
         return None
 
@@ -697,4 +716,4 @@ def no_toolset_func(ctx: RunContext[None]) -> None:
 
     tools = await toolset.get_tools(run_context)
 
-    assert tools == {}
+    assert tools == {}