refactor: simplify MCP elicitation by removing tool name mapping complexity

Author: yamanahlawat

mcp-run-python/src/tool_injection.py

@@ -64,19 +64,14 @@ def _create_tool_function(
     """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."""
-
-        # 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)
+        """Tool function that calls the MCP elicitation callback."""
+        elicitation_request = _create_elicitation_request(tool_name=tool_name, args=args, kwargs=kwargs)
 
         try:
             result = tool_callback(elicitation_request)
-            return _handle_tool_callback_result(result, actual_tool_name)
+            return _handle_tool_callback_result(result, tool_name)
         except Exception as e:
-            raise Exception(f'Tool {actual_tool_name} failed: {str(e)}')
+            raise Exception(f'Tool {tool_name} failed: {str(e)}')
 
     return tool_function
 
@@ -85,25 +80,18 @@ def inject_tool_functions(
     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 (should be Python-valid identifiers)
+        available_tools: List of available tool names
         tool_callback: Optional callback for tool execution
-        tool_name_mapping: Optional mapping of python_name -> original_mcp_name
     """
     if not available_tools:
         return
 
-    # 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:
-            # 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)
+            python_name = tool_name.replace('-', '_')
+            globals_dict[python_name] = _create_tool_function(tool_name, tool_callback, globals_dict)

pydantic_ai_slim/pydantic_ai/agent.py

@@ -12,7 +12,6 @@
 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
@@ -34,7 +33,7 @@
 from ._agent_graph import HistoryProcessor
 from ._output import OutputToolset
 from ._tool_manager import ToolManager
-from .mcp import MCPServer
+from .mcp import MCPServer, create_auto_tool_injection_callback, create_tool_elicitation_callback
 from .models.instrumented import InstrumentationSettings, InstrumentedModel, instrument_model
 from .output import OutputDataT, OutputSpec
 from .profiles import ModelProfile
@@ -1697,7 +1696,6 @@ def _get_toolset(
             if self._prepare_output_tools:
                 output_toolset = PreparedToolset(output_toolset, self._prepare_output_tools)
             all_toolsets = [output_toolset, *all_toolsets]
-
         return CombinedToolset(all_toolsets)
 
     def _infer_name(self, function_frame: FrameType | None) -> None:
@@ -1796,19 +1794,6 @@ async def __aenter__(self) -> Self:
             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
@@ -1837,113 +1822,50 @@ def _set_sampling_model(toolset: AbstractToolset[AgentDepsT]) -> None:
 
         self._get_toolset().apply(_set_sampling_model)
 
-    def _create_elicitation_callback(self):
-        """Create an elicitation callback that routes to this agent's tools."""
+    def set_mcp_elicitation_toolset(self, toolset_for_elicitation: AbstractToolset[Any] | None = None) -> None:
+        """Set the toolset to use for MCP elicitation callbacks.
 
-        async def elicitation_callback(context: Any, params: Any) -> Any:
-            """Handle elicitation requests by delegating to the agent's tools."""
-            try:
-                tool_execution_data = json.loads(params.message)
-                tool_name = tool_execution_data.get('tool_name')
-                tool_arguments = tool_execution_data.get('arguments', {})
-
-                # Try function tools first
-                function_tools = self._function_toolset.tools
-                if tool_name in function_tools:
-                    tool_func = function_tools[tool_name].function_schema.function
-
-                    # Handle both sync and async functions
-
-                    if inspect.iscoroutinefunction(tool_func):
-                        result = await tool_func(**tool_arguments)
-                    else:
-                        result = tool_func(**tool_arguments)
-
-                    return mcp_types.ElicitResult(action='accept', content={'result': str(result)})
-
-                # Find the MCP server that has this tool
-                target_server = None
-                for toolset in self._user_toolsets:
-                    if not isinstance(toolset, MCPServer):
-                        continue
-                    if 'mcp-run-python' in str(toolset):
-                        continue
-
-                    # Check if this server has the tool
-                    try:
-                        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
-
-                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)}')
-
-        return elicitation_callback
-
-    def _create_auto_tool_injection_callback(self):
-        """Create a callback that auto-injects available tools into run_python_code calls."""
-
-        async def auto_inject_tools_callback(
-            ctx: RunContext[Any],
-            call_tool_func: Callable[[str, dict[str, Any], dict[str, Any] | None], Awaitable[Any]],
-            tool_name: str,
-            arguments: dict[str, Any],
-        ) -> Any:
-            """Auto-inject available tools into run_python_code calls."""
-            if tool_name == 'run_python_code':
-                # 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)
-
-        return auto_inject_tools_callback
+        This method configures all MCP servers in the agent's toolsets to use the provided
+        toolset for handling elicitation requests (tool injection). This enables Python code
+        executed via mcp-run-python to call back to the agent's tools.
+
+        Args:
+            toolset_for_elicitation: Toolset to use for tool injection via elicitation.
+                                    If None, uses the agent's complete toolset.
+
+        Example:
+            ```python
+            agent = Agent('openai:gpt-4o')
+            agent.tool(web_search)
+            agent.tool(send_email)
+
+            mcp_server = MCPServerStdio(command='deno', args=[...], allow_elicitation=True)
+            agent.add_toolset(mcp_server)
+
+            # Enable tool injection with all agent tools
+            agent.set_mcp_elicitation_toolset()
+
+            # Or use specific toolset
+            custom_toolset = FunctionToolset(web_search)
+            agent.set_mcp_elicitation_toolset(custom_toolset)
+            ```
+        """
+        if toolset_for_elicitation is None:
+            # Use complete toolset for both elicitation and injection
+            toolset_for_elicitation = self._get_toolset()
+
+        # Set up callbacks for all MCP servers
+        def _set_elicitation_toolset(toolset: AbstractToolset[Any]) -> None:
+            if isinstance(toolset, MCPServer) and toolset.allow_elicitation:
+                # Set up elicitation callback
+                if toolset.elicitation_callback is None:
+                    toolset.elicitation_callback = create_tool_elicitation_callback(toolset=toolset_for_elicitation)
+
+                # Set up tool injection callback
+                if toolset.process_tool_call is None:
+                    toolset.process_tool_call = create_auto_tool_injection_callback(toolset=toolset_for_elicitation)
+
+        self._get_toolset().apply(_set_elicitation_toolset)
 
     @asynccontextmanager
     @deprecated(
@@ -1965,21 +1887,6 @@ async def run_mcp_servers(
             if model is not None:
                 raise
 
-        # Auto-setup elicitation callback if allow_elicitation is True and no callback is set
-
-        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()
-
         async with self:
             yield