Default Python tool to direct provider (#1467)

Signed-off-by: tamohannes <hovhannes.tamoyan@gmail.com>

Author: tamohannes

nemo_skills/mcp/servers/python_tool.py

@@ -26,7 +26,6 @@
 
 from nemo_skills.code_execution.sandbox import get_sandbox
 from nemo_skills.mcp.tool_manager import Tool
-from nemo_skills.mcp.tool_providers import MCPClientTool
 from nemo_skills.mcp.utils import add_config_args, load_mcp_config
 
 logger = logging.getLogger(__name__)
@@ -110,74 +109,24 @@ def main():
 # ==============================
 
 
-class PythonTool(MCPClientTool):
-    def __init__(self, exec_timeout_s: float = DEFAULT_EXEC_TIMEOUT_S) -> None:
-        super().__init__()
-        # Defaults for stdio Python MCP using explicit client class
-        self.apply_config_updates(
-            {
-                "client": "nemo_skills.mcp.clients.MCPStdioClient",
-                "client_params": {
-                    "command": "python",
-                    "args": ["-m", "nemo_skills.mcp.servers.python_tool"],
-                },
-                # hide args from schemas and sanitize at runtime
-                "hide_args": {"stateful_python_code_exec": ["session_id", "timeout"]},
-                # use explicit Hydra connector built from full context by default
-                "init_hook": "hydra",
-                # execution-specific default
-                "exec_timeout_s": exec_timeout_s,
-            }
-        )
-        self.requests_to_sessions = defaultdict(lambda: None)
-
-    def _description(self) -> str:
-        return get_python_tool_description(self._config.get("exec_timeout_s", DEFAULT_EXEC_TIMEOUT_S))
-
-    async def list_tools(self) -> List[Dict[str, Any]]:
-        tools = await super().list_tools()
-        return [
-            {**tool, "description": self._description()} if tool.get("name") == "stateful_python_code_exec" else tool
-            for tool in tools
-        ]
-
-    async def execute(self, tool_name: str, arguments: Dict[str, Any], extra_args: Dict[str, Any] | None = None):
-        # Ensure timeout is sent via extra_args (post-sanitize), not in main arguments
-        arguments = dict(arguments)
-        # TODO: error handling?
-        request_id = extra_args.pop("request_id")
-        merged_extra = dict(extra_args or {})
-        merged_extra.setdefault("timeout", self._config.get("exec_timeout_s", DEFAULT_EXEC_TIMEOUT_S))
-        merged_extra["session_id"] = self.requests_to_sessions[request_id]
-        result = await self._client.call_tool(tool=tool_name, args=arguments, extra_args=merged_extra)
-        self.requests_to_sessions[request_id] = result["session_id"]
-        output = f"{result['output_dict']['stdout']}{result['output_dict']['stderr']}"
-        if output.endswith("\n"):  # there is always a trailing newline, removing it
-            output = output[:-1]
-        return output
-
-    async def shutdown(self) -> None:
-        return None
-
-
 class DirectPythonTool(Tool):
     """Python code execution tool that calls the sandbox directly, bypassing MCP.
 
-    This is a drop-in replacement for PythonTool that eliminates the MCP protocol
-    overhead (subprocess spawning, MCP session initialization, JSON-RPC serialization)
-    by calling sandbox.execute_code() directly via HTTP.
+    This is the in-process implementation used by PythonTool. It eliminates the
+    MCP protocol overhead (subprocess spawning, MCP session initialization,
+    JSON-RPC serialization) by calling sandbox.execute_code() directly via HTTP.
 
-    Shared config keys with PythonTool (so switching is just changing the module spec):
+    Config keys:
         - hide_args: controls which args are stripped from schemas and sanitized at runtime
         - exec_timeout_s: default execution timeout
 
     Usage:
-        tool_modules=["nemo_skills.mcp.servers.python_tool::DirectPythonTool"]
+        tool_modules=["nemo_skills.mcp.servers.python_tool::PythonTool"]
     """
 
     def __init__(self, exec_timeout_s: float = DEFAULT_EXEC_TIMEOUT_S) -> None:
         self._config: Dict[str, Any] = {
-            # Same keys/defaults as PythonTool (minus MCP-specific: client, client_params, init_hook)
+            # MCP-specific keys (client, client_params, init_hook) are intentionally absent here
             "hide_args": {"stateful_python_code_exec": ["session_id", "timeout"]},
             "exec_timeout_s": exec_timeout_s,
             "sandbox": {},
@@ -299,5 +248,17 @@ async def cleanup_request(self, request_id: str) -> None:
         self.requests_to_sessions.pop(request_id, None)
 
 
+class PythonTool(DirectPythonTool):
+    """Default Python tool implementation.
+
+    Uses the direct in-process provider (DirectPythonTool) instead of stdio MCP
+    transport so generation jobs do not depend on a subprocess teardown path.
+    Kept as a subclass for backward compatibility with existing
+    ``tool_modules=[...python_tool::PythonTool]`` references.
+    """
+
+    pass
+
+
 if __name__ == "__main__":
     main()

tests/test_mcp_clients.py

@@ -637,7 +637,7 @@ async def get_items(count: int):
 
 
 # ==============================
-# Comparison tests: MCP PythonTool vs DirectPythonTool
+# DirectPythonTool functional tests
 # ==============================
 
 
@@ -815,74 +815,50 @@ async def test_direct_python_tool_cleanup_request_deletes_session():
 
 
 @pytest.mark.asyncio
-async def test_mcp_vs_direct_python_tool_parity():
-    """MCP-based PythonTool and DirectPythonTool produce identical results for the same tool calls."""
-    from nemo_skills.mcp.servers.python_tool import DirectPythonTool, PythonTool
+async def test_direct_python_tool_session_sequence():
+    """DirectPythonTool runs a multi-step stateful session and computes correct results."""
+    from nemo_skills.mcp.servers.python_tool import DirectPythonTool
 
     sandbox_context = {"sandbox": {"sandbox_type": "local"}}
 
-    # Set up DirectPythonTool
     direct = DirectPythonTool()
     direct.configure(context=sandbox_context)
 
-    # Set up MCP PythonTool
-    mcp_tool = PythonTool()
-    mcp_tool.configure(context=sandbox_context)
-
-    # Verify both expose the same tool name
-    direct_tools = await direct.list_tools()
-    mcp_tools = await mcp_tool.list_tools()
-    assert direct_tools[0]["name"] == mcp_tools[0]["name"] == "stateful_python_code_exec"
+    assert (await direct.list_tools())[0]["name"] == "stateful_python_code_exec"
 
-    # Define a sequence of tool calls that exercises session persistence
+    # A sequence of tool calls that exercises session persistence
     tool_calls = [
-        {"code": "import math", "request_id": "parity"},
-        {"code": "result = math.factorial(10)", "request_id": "parity"},
-        {"code": "print(result)", "request_id": "parity"},
-        {"code": "x = [i**2 for i in range(5)]", "request_id": "parity"},
-        {"code": "print(sum(x))", "request_id": "parity"},
+        {"code": "import math", "request_id": "seq"},
+        {"code": "result = math.factorial(10)", "request_id": "seq"},
+        {"code": "print(result)", "request_id": "seq"},
+        {"code": "x = [i**2 for i in range(5)]", "request_id": "seq"},
+        {"code": "print(sum(x))", "request_id": "seq"},
     ]
 
-    direct_results = await _run_tool_sequence(direct, tool_calls)
-    mcp_results = await _run_tool_sequence(mcp_tool, tool_calls)
-
-    for i, (d, m) in enumerate(zip(direct_results, mcp_results)):
-        assert d == m, f"Mismatch at step {i}: direct={d!r}, mcp={m!r}"
+    results = await _run_tool_sequence(direct, tool_calls)
 
     # Verify the actual computed values are correct
-    assert direct_results[2] == "3628800"  # 10!
-    assert direct_results[4] == "30"  # 0 + 1 + 4 + 9 + 16
+    assert results[2] == "3628800"  # 10!
+    assert results[4] == "30"  # 0 + 1 + 4 + 9 + 16
 
     await direct.shutdown()
-    await mcp_tool.shutdown()
 
 
 @pytest.mark.asyncio
-async def test_mcp_vs_direct_error_parity():
-    """Both implementations handle errors the same way."""
-    from nemo_skills.mcp.servers.python_tool import DirectPythonTool, PythonTool
+async def test_direct_python_tool_surfaces_exceptions():
+    """DirectPythonTool surfaces Python exceptions in the returned output."""
+    from nemo_skills.mcp.servers.python_tool import DirectPythonTool
 
     sandbox_context = {"sandbox": {"sandbox_type": "local"}}
 
     direct = DirectPythonTool()
     direct.configure(context=sandbox_context)
 
-    mcp_tool = PythonTool()
-    mcp_tool.configure(context=sandbox_context)
-
-    tool_calls = [
-        {"code": "1 / 0", "request_id": "err"},
-    ]
-
-    direct_results = await _run_tool_sequence(direct, tool_calls)
-    mcp_results = await _run_tool_sequence(mcp_tool, tool_calls)
+    results = await _run_tool_sequence(direct, [{"code": "1 / 0", "request_id": "err"}])
 
-    # Both should contain ZeroDivisionError
-    assert "ZeroDivisionError" in direct_results[0]
-    assert "ZeroDivisionError" in mcp_results[0]
+    assert "ZeroDivisionError" in results[0]
 
     await direct.shutdown()
-    await mcp_tool.shutdown()
 
 
 # ==============================
@@ -938,24 +914,21 @@ def test_python_tools_accept_exec_timeout_argument():
     assert DirectPythonTool(exec_timeout_s=42).default_config()["exec_timeout_s"] == 42
 
 
+def test_python_tool_defaults_to_direct_provider():
+    from nemo_skills.mcp.servers.python_tool import DirectPythonTool, PythonTool
+
+    # PythonTool is now the in-process direct provider (no stdio MCP transport).
+    assert isinstance(PythonTool(), DirectPythonTool)
+
+
 @pytest.mark.asyncio
 async def test_python_tool_description_uses_exec_timeout_override():
     from nemo_skills.mcp.servers.python_tool import PythonTool
 
-    class FakeClient:
-        async def list_tools(self):
-            return [
-                {
-                    "name": "stateful_python_code_exec",
-                    "description": "stale server description",
-                    "input_schema": {"type": "object", "properties": {"code": {"type": "string"}}},
-                }
-            ]
-
     tool = PythonTool()
     tool.configure(
-        overrides={"client": FakeClient, "client_params": {}, "exec_timeout_s": 37, "init_hook": None},
-        context={},
+        overrides={"exec_timeout_s": 37},
+        context={"sandbox": {"sandbox_type": "local"}},
     )
 
     tools = await tool.list_tools()
@@ -964,45 +937,6 @@ async def list_tools(self):
     assert "10.0 seconds" not in tools[0]["description"]
 
 
-@pytest.mark.asyncio
-async def test_python_tool_uses_exec_timeout_argument():
-    from nemo_skills.mcp.servers.python_tool import PythonTool
-
-    class FakeClient:
-        def __init__(self):
-            self.extra_args = None
-
-        async def list_tools(self):
-            return [
-                {
-                    "name": "stateful_python_code_exec",
-                    "description": "stale server description",
-                    "input_schema": {"type": "object", "properties": {"code": {"type": "string"}}},
-                }
-            ]
-
-        async def call_tool(self, tool, args, extra_args=None):
-            self.extra_args = extra_args
-            return {"session_id": "session", "output_dict": {"stdout": "done\n", "stderr": ""}}
-
-    fake_client = FakeClient()
-    tool = PythonTool(exec_timeout_s=42)
-    tool._client = fake_client
-
-    result = await tool.execute(
-        "stateful_python_code_exec",
-        {"code": "print('done')"},
-        extra_args={"request_id": "timeout-arg"},
-    )
-
-    assert result == "done"
-    assert fake_client.extra_args["timeout"] == 42
-
-    tools = await tool.list_tools()
-    assert "42.0 seconds" in tools[0]["description"]
-    assert "10.0 seconds" not in tools[0]["description"]
-
-
 @pytest.mark.asyncio
 async def test_direct_python_tool_uses_exec_timeout_argument():
     from nemo_skills.mcp.servers.python_tool import DirectPythonTool