Require WrapperToolset subclasses to implement their own prepare_for_run

Author: DouweM

pydantic_ai_slim/pydantic_ai/toolset.py

@@ -441,6 +441,10 @@ async def __aexit__(
     ) -> bool | None:
         return await self.wrapped.__aexit__(exc_type, exc_value, traceback)
 
+    @abstractmethod
+    async def prepare_for_run(self, ctx: RunContext[AgentDepsT]) -> RunToolset[AgentDepsT]:
+        raise NotImplementedError()
+
     @property
     def tool_defs(self) -> list[ToolDefinition]:
         return self.wrapped.tool_defs

tests/test_mcp.py

@@ -232,7 +232,7 @@ def get_none() -> None:  # pragma: no cover
         with pytest.raises(
             UserError,
             match=re.escape(
-                "MCPServerStdio(command='python', args=['-m', 'tests.mcp_server'], tool_prefix=None) defines a tool whose name conflicts with existing tool from FunctionToolset: 'get_none'. Consider setting `tool_prefix` to avoid name conflicts."
+                "MCPServerStdio(command='python', args=['-m', 'tests.mcp_server'], tool_prefix=None) defines a tool whose name conflicts with existing tool from Function toolset: 'get_none'. Consider setting `tool_prefix` to avoid name conflicts."
             ),
         ):
             await agent.run('Get me a conflict')