feat: replace global _mcp_servers with toolsets

Author: yamanahlawat

mcp-run-python/src/main.ts

@@ -96,13 +96,19 @@ The tools are injected into the global namespace automatically - no discovery fu
         .array(z.string())
         .optional()
         .describe('List of available tools for injection (enables tool injection when provided)'),
+      tool_name_mapping: z
+        .record(z.string())
+        .optional()
+        .describe('Mapping of python_name -> original_mcp_name for tool name conversion'),
     },
     async ({
       python_code,
       tools = [],
+      tool_name_mapping = {},
     }: {
       python_code: string
       tools?: string[]
+      tool_name_mapping?: Record<string, string>
     }) => {
       const logPromises: Promise<void>[] = []
 
@@ -169,6 +175,7 @@ The tools are injected into the global namespace automatically - no discovery fu
           {
             enableToolInjection: true,
             availableTools: tools,
+            toolNameMapping: tool_name_mapping,
             timeoutSeconds: 30,
             elicitationCallback,
           } as ToolInjectionConfig,

mcp-run-python/src/runCode.ts

@@ -13,6 +13,7 @@ export interface CodeFile {
 export interface ToolInjectionConfig {
   enableToolInjection: boolean
   availableTools: string[]
+  toolNameMapping?: Record<string, string> // python_name -> original_mcp_name
   timeoutSeconds: number
   // deno-lint-ignore no-explicit-any
   elicitationCallback?: (request: any) => Promise<any>
@@ -157,6 +158,7 @@ function injectToolFunctions(
     globals,
     config.availableTools,
     tool_callback,
+    config.toolNameMapping,
   )
 
   log('info', `Tool injection complete. Available tools: ${config.availableTools.join(', ')}`)

mcp-run-python/src/tool_injection.py

@@ -58,38 +58,52 @@ def _handle_tool_callback_result(result: Any, tool_name: str) -> Any:
         return result
 
 
-def _create_tool_function(tool_name: str, tool_callback: Callable[[Any], Any]) -> Callable[..., Any]:
+def _create_tool_function(
+    tool_name: str, tool_callback: Callable[[Any], Any], globals_dict: dict[str, Any]
+) -> Callable[..., Any]:
     """Create a tool function that can be called from Python."""
 
     def tool_function(*args: Any, **kwargs: Any) -> Any:
         """Synchronous tool function that handles the async callback properly."""
-        # Note: tool_callback is guaranteed to be not None due to check in inject_tool_functions
 
-        elicitation_request = _create_elicitation_request(tool_name, args, kwargs)
+        # Get the actual MCP tool name from the stored mapping
+        tool_mapping = globals_dict.get('__tool_name_mapping__', {})
+        actual_tool_name = tool_mapping.get(tool_name, tool_name)
+
+        elicitation_request = _create_elicitation_request(actual_tool_name, args, kwargs)
 
         try:
             result = tool_callback(elicitation_request)
-            return _handle_tool_callback_result(result, tool_name)
+            return _handle_tool_callback_result(result, actual_tool_name)
         except Exception as e:
-            raise Exception(f'Tool {tool_name} failed: {str(e)}')
+            raise Exception(f'Tool {actual_tool_name} failed: {str(e)}')
 
     return tool_function
 
 
 def inject_tool_functions(
-    globals_dict: dict[str, Any], available_tools: list[str], tool_callback: Callable[[Any], Any] | None = None
+    globals_dict: dict[str, Any],
+    available_tools: list[str],
+    tool_callback: Callable[[Any], Any] | None = None,
+    tool_name_mapping: dict[str, str] | None = None,
 ) -> None:
     """Inject tool functions into the global namespace.
 
     Args:
         globals_dict: Global namespace to inject tools into
-        available_tools: List of available tool names
+        available_tools: List of available tool names (should be Python-valid identifiers)
         tool_callback: Optional callback for tool execution
+        tool_name_mapping: Optional mapping of python_name -> original_mcp_name
     """
     if not available_tools:
         return
 
-    # Inject tool functions into globals
+    # Store the tool name mapping globally for elicitation callback to use
+    if tool_name_mapping:
+        globals_dict['__tool_name_mapping__'] = tool_name_mapping
+
+    # Inject tool functions into globals using Python-valid names
     for tool_name in available_tools:
         if tool_callback is not None:
-            globals_dict[tool_name] = _create_tool_function(tool_name, tool_callback)
+            # tool_name should already be a valid Python identifier from agent.py
+            globals_dict[tool_name] = _create_tool_function(tool_name, tool_callback, globals_dict)

pydantic_ai_slim/pydantic_ai/agent.py

@@ -12,6 +12,7 @@
 from types import FrameType
 from typing import TYPE_CHECKING, Any, Callable, ClassVar, Generic, cast, final, overload
 
+from mcp import types as mcp_types
 from opentelemetry.trace import NoOpTracer, use_span
 from pydantic.json_schema import GenerateJsonSchema
 from typing_extensions import Literal, Never, Self, TypeIs, TypeVar, deprecated
@@ -70,7 +71,6 @@
     from fasta2a.broker import Broker
     from fasta2a.schema import AgentProvider, Skill
     from fasta2a.storage import Storage
-    from mcp import types as mcp_types
     from starlette.middleware import Middleware
     from starlette.routing import BaseRoute, Route
     from starlette.types import ExceptionHandler, Lifespan
@@ -1814,6 +1814,20 @@ async def __aenter__(self) -> Self:
         async with self._enter_lock:
             if self._entered_count == 0:
                 self._exit_stack = AsyncExitStack()
+
+                for toolset in self._user_toolsets:
+                    if isinstance(toolset, MCPServer):
+                        if (
+                            hasattr(toolset, 'allow_elicitation')
+                            and toolset.allow_elicitation
+                            and toolset.elicitation_callback is None
+                        ):
+                            toolset.elicitation_callback = self._create_elicitation_callback()
+
+                            # Also setup auto-tool-injection for run_python_code if not already set
+                            if toolset.process_tool_call is None:
+                                toolset.process_tool_call = self._create_auto_tool_injection_callback()
+
                 toolset = self._get_toolset()
                 await self._exit_stack.enter_async_context(toolset)
             self._entered_count += 1
@@ -1866,23 +1880,36 @@ async def elicitation_callback(context: Any, params: Any) -> Any:
 
                     return mcp_types.ElicitResult(action='accept', content={'result': str(result)})
 
-                # Try MCP tools with name mapping
-                actual_tool_name = tool_name.replace('_', '-')
-
+                # Find the MCP server that has this tool
+                target_server = None
                 for toolset in self._user_toolsets:
                     if not isinstance(toolset, MCPServer):
                         continue
-                    mcp_server = toolset
-                    if 'mcp-run-python' in str(mcp_server):
+                    if 'mcp-run-python' in str(toolset):
                         continue
 
+                    # Check if this server has the tool
                     try:
-                        result = await mcp_server.direct_call_tool(actual_tool_name, tool_arguments)
-                        return mcp_types.ElicitResult(action='accept', content={'result': str(result)})
+                        server_tools = await toolset.list_tools()
+                        for tool_def in server_tools:
+                            if tool_def.name == tool_name:
+                                target_server = toolset
+                                break
+                        if target_server:
+                            break
                     except Exception:
                         continue
 
-                return mcp_types.ErrorData(code=mcp_types.INVALID_PARAMS, message=f'Tool {tool_name} not found')
+                if target_server:
+                    try:
+                        result = await target_server.direct_call_tool(tool_name, tool_arguments)
+                        return mcp_types.ElicitResult(action='accept', content={'result': str(result)})
+                    except Exception as e:
+                        return mcp_types.ErrorData(
+                            code=mcp_types.INTERNAL_ERROR, message=f'Tool execution failed: {str(e)}'
+                        )
+                else:
+                    return mcp_types.ErrorData(code=mcp_types.INVALID_PARAMS, message=f'Tool {tool_name} not found')
 
             except Exception as e:
                 return mcp_types.ErrorData(code=mcp_types.INTERNAL_ERROR, message=f'Tool execution failed: {str(e)}')
@@ -1900,30 +1927,37 @@ async def auto_inject_tools_callback(
         ) -> Any:
             """Auto-inject available tools into run_python_code calls."""
             if tool_name == 'run_python_code':
-                # Auto-inject available tools if not already provided
-                if 'tools' not in arguments or not arguments['tools']:
-                    available_tools: list[str] = []
-
-                    # Add function tools
-                    available_tools.extend(list(self._function_toolset.tools.keys()))
-
-                    for toolset in self._user_toolsets:
-                        if not isinstance(toolset, MCPServer):
-                            continue
-                        mcp_server = toolset
-                        if 'mcp-run-python' in str(mcp_server):
-                            continue
-
-                        try:
-                            server_tools = await mcp_server.list_tools()
-                            for tool_def in server_tools:
-                                python_name = tool_def.name.replace('-', '_')
-                                available_tools.append(python_name)
-                        except Exception:
-                            # Silently continue if we can't get tools from a server
-                            pass
-
-                    arguments['tools'] = available_tools
+                # Always auto-inject all available tools for Python code execution
+                available_tools: list[str] = []
+                tool_name_mapping: dict[str, str] = {}
+
+                # Add function tools
+                function_tools = list(self._function_toolset.tools.keys())
+                available_tools.extend(function_tools)
+                for func_tool_name in function_tools:
+                    tool_name_mapping[func_tool_name] = func_tool_name
+
+                # Add MCP server tools with proper name conversion
+                for toolset in self._user_toolsets:
+                    if not isinstance(toolset, MCPServer):
+                        continue
+                    if 'mcp-run-python' in str(toolset):
+                        continue
+
+                    try:
+                        server_tools = await toolset.list_tools()
+                        for tool_def in server_tools:
+                            original_name = tool_def.name
+                            python_name = original_name.replace('-', '_')
+                            available_tools.append(python_name)
+                            tool_name_mapping[python_name] = original_name
+                    except Exception:
+                        # Silently continue if we can't get tools from a server
+                        pass
+
+                # Always provide all available tools and mapping
+                arguments['tools'] = available_tools
+                arguments['tool_name_mapping'] = tool_name_mapping
 
             # Continue with normal processing
             return await call_tool_func(tool_name, arguments, None)
@@ -1954,17 +1988,16 @@ async def run_mcp_servers(
 
         for toolset in self._user_toolsets:
             if isinstance(toolset, MCPServer):
-                mcp_server = toolset
                 if (
-                    hasattr(mcp_server, 'allow_elicitation')
-                    and mcp_server.allow_elicitation
-                    and mcp_server.elicitation_callback is None
+                    hasattr(toolset, 'allow_elicitation')
+                    and toolset.allow_elicitation
+                    and toolset.elicitation_callback is None
                 ):
-                    mcp_server.elicitation_callback = self._create_elicitation_callback()
+                    toolset.elicitation_callback = self._create_elicitation_callback()
 
                     # Also setup auto-tool-injection for run_python_code if not already set
-                    if mcp_server.process_tool_call is None:
-                        mcp_server.process_tool_call = self._create_auto_tool_injection_callback()
+                    if toolset.process_tool_call is None:
+                        toolset.process_tool_call = self._create_auto_tool_injection_callback()
 
         async with self:
             yield

tests/test_mcp_elicitation.py

@@ -82,11 +82,15 @@ async def mock_elicitation(
         )
 
         model = TestModel(custom_output_text='Test response')
-        agent = Agent(model, mcp_servers=[server])
+        agent = Agent(model, toolsets=[server])
 
         # Verify the server is properly configured
-        assert len(agent._mcp_servers) == 1  # type: ignore
-        assert agent._mcp_servers[0].elicitation_callback is mock_elicitation  # type: ignore
+        toolsets = getattr(agent, '_user_toolsets', [])
+        mcp_servers = [ts for ts in toolsets if hasattr(ts, '__class__') and 'MCPServer' in ts.__class__.__name__]
+        assert len(mcp_servers) == 1
+        # Use getattr to safely check elicitation_callback
+        callback = getattr(mcp_servers[0], 'elicitation_callback', None)
+        assert callback is mock_elicitation
 
     async def test_elicitation_callback_error_handling(self):
         """Test error handling in elicitation callback."""
@@ -483,14 +487,19 @@ async def agent_tool_callback(
 
         # Create agent with the MCP server
         model = TestModel(custom_output_text='Tool injection test completed')
-        agent = Agent(model, mcp_servers=[mcp_server])
+        agent = Agent(model, toolsets=[mcp_server])
 
         # Verify the agent has the MCP server with elicitation callback
-        assert len(agent._mcp_servers) == 1  # type: ignore
-        assert agent._mcp_servers[0].elicitation_callback is agent_tool_callback  # type: ignore
+        # Note: Using getattr to safely access toolsets for testing
+        toolsets = getattr(agent, '_user_toolsets', [])
+        mcp_servers = [ts for ts in toolsets if hasattr(ts, '__class__') and 'MCPServer' in ts.__class__.__name__]
+        assert len(mcp_servers) == 1
+        # Use getattr to safely check elicitation_callback
+        callback = getattr(mcp_servers[0], 'elicitation_callback', None)
+        assert callback is agent_tool_callback
 
         # Test running agent with MCP servers
-        async with agent.run_mcp_servers():
+        async with agent:
             # Verify the MCP server is properly integrated
             tools = await mcp_server.list_tools()
             assert len(tools) == 1
@@ -780,7 +789,7 @@ async def test_mcp_run_python_code_execution(self):
 
         async with server:
             # Test basic Python execution
-            result = await server.call_tool(
+            result = await server.direct_call_tool(
                 'run_python_code', {'python_code': 'print("Hello, World!")\n"Hello from Python"'}
             )
 
@@ -827,7 +836,7 @@ async def python_code_callback(
         async with server:
             # Test Python code execution with tool injection
             # This should trigger the elicitation callback when tools are called
-            result = await server.call_tool(
+            result = await server.direct_call_tool(
                 'run_python_code',
                 {'python_code': 'print("Testing tool injection")', 'tools': ['web_search', 'calculate']},
             )
@@ -854,7 +863,7 @@ async def test_mcp_run_python_error_handling(self):
 
         async with server:
             # Test Python code with syntax error
-            result = await server.call_tool('run_python_code', {'python_code': 'print("Missing closing quote)'})
+            result = await server.direct_call_tool('run_python_code', {'python_code': 'print("Missing closing quote)'})
 
             # Should return error status
             assert isinstance(result, str)
@@ -1009,7 +1018,7 @@ async def test_mcp_run_python_with_dependencies(self):
 
         async with server:
             # Test code with dependencies
-            result = await server.call_tool(
+            result = await server.direct_call_tool(
                 'run_python_code',
                 {
                     'python_code': """
@@ -1049,7 +1058,8 @@ async def test_mcp_run_python_with_tool_prefix(self):
         async with server:
             tools = await server.list_tools()
             assert len(tools) == 1
-            assert tools[0].name == 'python_run_python_code'
+            # list_tools() returns original tool names without prefix
+            assert tools[0].name == 'run_python_code'
 
     async def test_mcp_run_python_timeout_setting(self):
         """Test mcp-run-python server with timeout setting."""
@@ -1071,7 +1081,7 @@ async def test_mcp_run_python_timeout_setting(self):
 
         async with server:
             # Test basic execution still works
-            result = await server.call_tool('run_python_code', {'python_code': 'print("Timeout test")'})
+            result = await server.direct_call_tool('run_python_code', {'python_code': 'print("Timeout test")'})
 
             assert isinstance(result, str)
             assert '<status>success</status>' in result