feat: implement SandboxMCPBackend for isolated worker support

Author: liuruibin

apps/application/flow/backend/sandbox_mcp.py

@@ -0,0 +1,40 @@
+"""MCP backend using isolated workers for user-configured remote servers."""
+
+from langchain_mcp_adapters.client import MultiServerMCPClient
+from mcp.types import CallToolResult
+
+from common.mcp.config import InternalMCPConfig, validate_mcp_servers
+from common.mcp.sandbox import sandbox_connection
+
+
+class SandboxMCPBackend(MultiServerMCPClient):
+    """Provide MCP tools and sessions backed by sandboxed stdio connections.
+
+    Inherited get_tools() creates tools whose later invocations also open sandbox
+    workers. This backend supplies the agent's tools; SandboxShellBackend remains
+    the agent backend for skill files and shell commands.
+    """
+
+    def __init__(self, servers: dict):
+        super().__init__(connections=self._build_connections(servers))
+
+    @staticmethod
+    def _build_connections(servers: dict) -> dict:
+        if not isinstance(servers, dict):
+            raise ValueError("MCP servers must be an object")
+        connections = {}
+        for name, config in servers.items():
+            if not isinstance(config, dict):
+                raise ValueError("MCP server configuration must be an object")
+            internal = isinstance(config, InternalMCPConfig)
+            if internal and config.get("transport") == "stdio":
+                connections[name] = dict(config)
+                continue
+            validate_mcp_servers({name: config})
+            connections[name] = dict(config) if internal else sandbox_connection(config)
+        return connections
+
+    async def call_tool(self, server_name: str, tool_name: str, arguments: dict | None = None) -> CallToolResult:
+        """Call one tool and close its session/worker, preserving the MCP result."""
+        async with self.session(server_name) as session:
+            return await session.call_tool(tool_name, arguments)

apps/application/flow/step_node/mcp_node/impl/base_mcp_node.py

@@ -4,7 +4,7 @@
 from typing import List
 
 from django.db.models import QuerySet
-from common.utils.mcp_client import create_mcp_client
+from application.flow.backend.sandbox_mcp import SandboxMCPBackend
 
 from application.flow.i_step_node import NodeResult
 from application.flow.step_node.mcp_node.i_mcp_node import IMcpNode
@@ -37,12 +37,8 @@ def execute(self, mcp_servers, mcp_server, mcp_tool, mcp_tool_id, mcp_source, to
         params = json.loads(json.dumps(tool_params))
         params = self.handle_variables(params)
 
-        async def call_tool(t, a):
-            client = create_mcp_client(servers)
-            async with client.session(mcp_server) as s:
-                return await s.call_tool(t, a)
-
-        res = asyncio.run(call_tool(mcp_tool, params))
+        backend = SandboxMCPBackend(servers)
+        res = asyncio.run(backend.call_tool(mcp_server, mcp_tool, params))
         return NodeResult(
             {"result": [content.text for content in res.content], "tool_params": params, "mcp_tool": mcp_tool}, {}
         )

apps/application/flow/tools.py

@@ -45,12 +45,12 @@
 from langchain_core.messages import AIMessageChunk, BaseMessage, BaseMessageChunk, ToolMessage
 from langchain_core.tools import StructuredTool
 from langchain_core.utils._merge import merge_lists as _original_merge_lists
-from common.utils.mcp_client import create_mcp_client
 from langgraph.checkpoint.memory import MemorySaver
 from maxkb.const import CONFIG
 from pydantic import Field, create_model
 from tools.models import Tool, ToolRecord, ToolType, ToolWorkflowVersion
 
+from application.flow.backend.sandbox_mcp import SandboxMCPBackend
 from application.flow.backend.sandbox_shell import SandboxShellBackend
 from application.flow.common import Workflow, WorkflowMode
 from application.flow.i_step_node import ToolWorkflowPostHandler, WorkFlowPostHandler
@@ -394,7 +394,7 @@ def _extract_tool_id(raw_id):
     return tool_id or raw_id
 
 
-async def _initialize_skills(mcp_servers, temp_dir):
+async def _initialize_skills(mcp_servers, temp_dir) -> SandboxMCPBackend:
     skills_dir = os.path.join(temp_dir, "skills")
     mcp_config = dict(mcp_servers)  # Preserve server-generated InternalMCPConfig objects.
     if "skills" in mcp_config:
@@ -435,9 +435,7 @@ async def _initialize_skills(mcp_servers, temp_dir):
 
         os.system("chmod -R g+rx " + temp_dir)  # 确保技能目录可访问
 
-    client = create_mcp_client(mcp_config)
-
-    return client
+    return SandboxMCPBackend(mcp_config)
 
 
 async def _yield_mcp_response(
@@ -455,8 +453,8 @@ async def _yield_mcp_response(
 ):
     try:
         checkpointer = MemorySaver()
-        client = await _initialize_skills(mcp_servers, temp_dir)
-        tools = await client.get_tools()
+        mcp_backend = await _initialize_skills(mcp_servers, temp_dir)
+        tools = await mcp_backend.get_tools()
         for tool in tools:
             tool.handle_tool_error = True
         if extra_tools:

apps/application/serializers/application.py

@@ -51,7 +51,7 @@
 from knowledge.models import File, FileSourceType, Knowledge, KnowledgeScope
 from knowledge.serializers.common import BatchMoveSerializer, BatchSerializer
 from knowledge.serializers.knowledge import KnowledgeModelSerializer, KnowledgeSerializer
-from common.utils.mcp_client import create_mcp_client
+from application.flow.backend.sandbox_mcp import SandboxMCPBackend
 from maxkb.conf import PROJECT_DIR
 from maxkb.const import CONFIG
 from models_provider.models import Model
@@ -1095,8 +1095,8 @@ class PlayDemoTextRequest(serializers.Serializer):
 
 
 async def get_mcp_tools(servers):
-    client = create_mcp_client(servers)
-    return await client.get_tools()
+    backend = SandboxMCPBackend(servers)
+    return await backend.get_tools()
 
 
 class McpServersSerializer(serializers.Serializer):

apps/common/mcp/__init__.py

@@ -0,0 +1 @@
+"""Shared MCP configuration and sandbox workers."""

apps/common/mcp/client.py

@@ -0,0 +1,12 @@
+"""Compatibility exports and client factory for the dedicated MCP backend."""
+
+from application.flow.backend.sandbox_mcp import SandboxMCPBackend
+from common.mcp.config import InternalMCPConfig, validate_mcp_servers
+
+
+__all__ = ["InternalMCPConfig", "validate_mcp_servers", "create_mcp_client"]
+
+
+def create_mcp_client(servers):
+    """Keep existing callers compatible with the dedicated MCP backend."""
+    return SandboxMCPBackend(servers)

apps/common/mcp/config.py

@@ -0,0 +1,19 @@
+"""Shared MCP configuration types and validation."""
+
+
+class InternalMCPConfig(dict):
+    """In-memory provenance for configurations generated by ToolExecutor.
+
+    Never deserialize user input into this type. JSON round trips deliberately
+    lose this privilege; keep runtime configurations in memory instead.
+    """
+
+
+def validate_mcp_servers(servers):
+    if not isinstance(servers, dict):
+        raise ValueError("MCP servers must be an object")
+    for config in servers.values():
+        if not isinstance(config, dict) or config.get("transport") not in ("sse", "streamable_http"):
+            raise ValueError("Only support transport=sse or transport=streamable_http")
+        if not isinstance(config.get("url"), str) or not config["url"].strip():
+            raise ValueError("MCP server URL must be a non-empty string")

apps/common/utils/mcp_sandbox.py apps/common/mcp/sandbox.pyapps/common/utils/mcp_sandbox.py renamed to apps/common/mcp/sandbox.py

@@ -4,18 +4,19 @@
 import pwd
 import sys
 from datetime import timedelta
+from importlib.machinery import PathFinder
 from pathlib import Path
 
 from mcp.types import Implementation
 
+from maxkb.const import CONFIG
+
 
 BOOTSTRAP_KEY = "maxkbSandbox"
 REMOTE_FIELDS = {"transport", "url", "headers", "timeout", "sse_read_timeout", "terminate_on_close"}
 
 
 def sandbox_settings():
-    from maxkb.const import CONFIG
-
     if not sys.platform.startswith("linux") or not bool(int(CONFIG.get("SANDBOX", 1))):
         raise ValueError("External MCP requires an enabled Linux sandbox")
     account = pwd.getpwnam("sandbox")
@@ -37,6 +38,11 @@ def sandbox_settings():
 
 def sandbox_connection(config):
     settings = sandbox_settings()
+    # Release builds replace source files with adjacent, sourceless .pyc files.
+    # Search only our installed directory, never a user-controlled module path.
+    worker = PathFinder.find_spec("sandbox_worker", [str(Path(__file__).parent)])
+    if worker is None or worker.origin is None or Path(worker.origin).suffix not in (".py", ".pyc"):
+        raise RuntimeError("MCP sandbox worker is missing or has an unsupported format")
     # Only transport data goes to the remote client. In particular, ignore user
     # command/env/factory/session_kwargs fields and never deserialize Python code.
     remote = {key: value for key, value in config.items() if key in REMOTE_FIELDS}
@@ -46,7 +52,7 @@ def sandbox_connection(config):
     return {
         "transport": "stdio",
         "command": sys.executable,
-        "args": ["-I", str(Path(__file__).with_name("mcp_sandbox_worker.py"))],
+        "args": ["-I", worker.origin],
         "cwd": settings["cwd"],
         "env": {
             "LD_PRELOAD": settings["library"],

apps/common/utils/mcp_sandbox_proxy.py apps/common/mcp/sandbox_proxy.pyapps/common/utils/mcp_sandbox_proxy.py renamed to apps/common/mcp/sandbox_proxy.py

@@ -1,19 +1,51 @@
 """Forward MCP messages without converting tools, results or notifications."""
 
 from contextlib import asynccontextmanager
-from functools import partial
 import logging
 import os
+import socket
+import ssl
 import sys
 
 import anyio
 import httpx
 from mcp.client.sse import sse_client
 from mcp.client.streamable_http import streamable_http_client
 from mcp.server.stdio import stdio_server
+from mcp.shared._httpx_utils import create_mcp_http_client
 from mcp.types import JSONRPCRequest
 
 
+def sandbox_failure_message(error):
+    # SDK exception groups and chained HTTP errors can embed credentials. Only
+    # report numeric HTTP statuses or fixed descriptions, never exception text.
+    errors, pending, seen = [], [error], set()
+    while pending:
+        current = pending.pop()
+        if id(current) in seen:
+            continue
+        seen.add(id(current))
+        errors.append(current)
+        if isinstance(current, BaseExceptionGroup):
+            pending.extend(current.exceptions)
+        if current.__cause__ is not None:
+            pending.append(current.__cause__)
+    for current in errors:
+        if isinstance(current, httpx.HTTPStatusError):
+            return f"MCP endpoint returned HTTP {current.response.status_code}; check endpoint and credentials"
+    for exception_type, message in (
+        (ssl.SSLCertVerificationError, "MCP TLS certificate verification failed"),
+        (socket.gaierror, "MCP hostname resolution failed; check container DNS"),
+        (PermissionError, "MCP access denied; check sandbox file and network policy"),
+        ((httpx.TimeoutException, TimeoutError), "MCP connection timed out"),
+        (httpx.TooManyRedirects, "MCP endpoint returned too many redirects"),
+        (httpx.ConnectError, "MCP connection failed; check container connectivity and sandbox network policy"),
+    ):
+        if any(isinstance(current, exception_type) for current in errors):
+            return message
+    return "MCP session failed; check endpoint, sandbox setup and network policy"
+
+
 class PipeInput:
     """Cancellable pipe reads; a blocked readline thread would delay shutdown."""
 
@@ -75,11 +107,10 @@ def extract_bootstrap(message):
 
 
 @asynccontextmanager
-async def remote_transport(bootstrap, http_factory):
+async def remote_transport(bootstrap):
     config = bootstrap["connection"]
     if config.get("transport") not in ("sse", "streamable_http"):
         raise ValueError("Unsupported external MCP transport")
-    factory = partial(http_factory, url=config["url"])
     timeout = config.get("timeout", 5 if config["transport"] == "sse" else 30)
     read_timeout = config.get("sse_read_timeout", 300)
     if config["transport"] == "sse":
@@ -88,11 +119,13 @@ async def remote_transport(bootstrap, http_factory):
             headers=config.get("headers"),
             timeout=timeout,
             sse_read_timeout=read_timeout,
-            httpx_client_factory=factory,
         ) as streams:
             yield streams
     else:
-        async with factory(headers=config.get("headers"), timeout=httpx.Timeout(timeout, read=read_timeout)) as client:
+        async with create_mcp_http_client(
+            headers=config.get("headers"),
+            timeout=httpx.Timeout(timeout, read=read_timeout),
+        ) as client:
             async with streamable_http_client(
                 config["url"],
                 http_client=client,
@@ -111,19 +144,19 @@ async def forward(source, destination, cancel_scope):
         cancel_scope.cancel()
 
 
-async def proxy(http_factory):
+async def proxy():
     async with stdio_server(stdin=PipeInput(), stdout=PipeOutput()) as (local_read, local_write):
         with anyio.fail_after(30):
             first = await local_read.receive()
             bootstrap = extract_bootstrap(first)
-        async with remote_transport(bootstrap, http_factory) as (remote_read, remote_write):
+        async with remote_transport(bootstrap) as (remote_read, remote_write):
             async with anyio.create_task_group() as tasks:
                 tasks.start_soon(forward, remote_read, local_write, tasks.cancel_scope)
                 await remote_write.send(first)
                 tasks.start_soon(forward, local_read, remote_write, tasks.cancel_scope)
 
 
-def run(http_factory):
+def run():
     # Remote SDK exceptions may contain authorization headers or URL parameters.
     logging.disable(logging.CRITICAL)
-    anyio.run(proxy, http_factory)
+    anyio.run(proxy)

apps/common/utils/mcp_sandbox_worker.py apps/common/mcp/sandbox_worker.pyapps/common/utils/mcp_sandbox_worker.py renamed to apps/common/mcp/sandbox_worker.py

@@ -3,6 +3,7 @@
 import ctypes
 from contextlib import contextmanager
 import errno
+import importlib.machinery
 import importlib.util
 import ipaddress
 import json
@@ -17,7 +18,7 @@
 
 
 class MCPWorkerFailure(Exception):
-    """A failure whose message was sanitized by the fixed network module."""
+    """A failure whose message was sanitized by the protocol proxy."""
 
 
 class DlInfo(ctypes.Structure):
@@ -168,24 +169,23 @@ def enter_sandbox(settings):
 
 def main():
     settings = json.loads(os.environ.pop("MAXKB_MCP_WORKER_SETTINGS"))
-    # Load only fixed, installed modules before dropping access to the app tree.
-    # Neither module imports Django nor reads the application configuration.
-    modules = {}
-    for name in ("mcp_network", "mcp_sandbox_proxy"):
-        path = Path(__file__).with_name(name + ".py")
-        spec = importlib.util.spec_from_file_location(name, path)
-        module = importlib.util.module_from_spec(spec)
-        spec.loader.exec_module(module)
-        modules[name] = module
+    # Load the fixed proxy before dropping access to the app tree. It does not
+    # import Django or read application configuration. The loader supports both
+    # source and sourceless release layouts, only from our installed directory.
+    spec = importlib.machinery.PathFinder.find_spec("sandbox_proxy", [str(Path(__file__).parent)])
+    if spec is None or spec.loader is None:
+        raise RuntimeError("MCP sandbox dependency is missing")
+    proxy = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(proxy)
     # Remove the application directory while retaining approved package paths.
     app_path = str(Path(__file__).resolve().parents[2])
     sys.path = [p for p in sys.path if p != app_path]
     sys.path.extend(p for p in settings["python_paths"] if p and p not in sys.path)
     enter_sandbox(settings)
     try:
-        modules["mcp_sandbox_proxy"].run(modules["mcp_network"].http_client_factory)
+        proxy.run()
     except Exception as exc:
-        raise MCPWorkerFailure(modules["mcp_network"].sandbox_failure_message(exc)) from None
+        raise MCPWorkerFailure(proxy.sandbox_failure_message(exc)) from None
 
 
 if __name__ == "__main__":