Skip to content

Commit eb05c72

Browse files
committed
refactor: replace MultiServerMCPClient with create_mcp_client and update MCP server validation
1 parent 0bb81c7 commit eb05c72

14 files changed

Lines changed: 609 additions & 24 deletions

File tree

apps/application/chat_pipeline/step/chat_step/impl/base_chat_step.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -445,7 +445,7 @@ def _handle_mcp_request(
445445
chat_model,
446446
system_prompt,
447447
message_list,
448-
json.dumps(mcp_servers_config),
448+
mcp_servers_config,
449449
mcp_output_enable,
450450
tool_init_params,
451451
source_id,

apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -446,7 +446,7 @@ def _handle_mcp_request(
446446
chat_model,
447447
system_prompt,
448448
message_list,
449-
json.dumps(mcp_servers_config),
449+
mcp_servers_config,
450450
mcp_output_enable,
451451
tool_init_params,
452452
source_id,

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
from typing import List
55

66
from django.db.models import QuerySet
7-
from langchain_mcp_adapters.client import MultiServerMCPClient
7+
from common.utils.mcp_client import create_mcp_client
88

99
from application.flow.i_step_node import NodeResult
1010
from application.flow.step_node.mcp_node.i_mcp_node import IMcpNode
@@ -46,7 +46,7 @@ def execute(self, mcp_servers, mcp_server, mcp_tool, mcp_tool_id, mcp_source, to
4646
params = self.handle_variables(params)
4747

4848
async def call_tool(t, a):
49-
client = MultiServerMCPClient(servers)
49+
client = create_mcp_client(servers)
5050
async with client.session(mcp_server) as s:
5151
return await s.call_tool(t, a)
5252

apps/application/flow/tools.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@
4545
from langchain_core.messages import AIMessageChunk, BaseMessage, BaseMessageChunk, ToolMessage
4646
from langchain_core.tools import StructuredTool
4747
from langchain_core.utils._merge import merge_lists as _original_merge_lists
48-
from langchain_mcp_adapters.client import MultiServerMCPClient
48+
from common.utils.mcp_client import create_mcp_client
4949
from langgraph.checkpoint.memory import MemorySaver
5050
from maxkb.const import CONFIG
5151
from pydantic import Field, create_model
@@ -396,7 +396,7 @@ def _extract_tool_id(raw_id):
396396

397397
async def _initialize_skills(mcp_servers, temp_dir):
398398
skills_dir = os.path.join(temp_dir, "skills")
399-
mcp_config = json.loads(mcp_servers)
399+
mcp_config = dict(mcp_servers) # Preserve server-generated InternalMCPConfig objects.
400400
if "skills" in mcp_config:
401401
skill_file_items = mcp_config.pop("skills")
402402
for skill_file in skill_file_items:
@@ -435,7 +435,7 @@ async def _initialize_skills(mcp_servers, temp_dir):
435435

436436
os.system("chmod -R g+rx " + temp_dir) # 确保技能目录可访问
437437

438-
client = MultiServerMCPClient(mcp_config)
438+
client = create_mcp_client(mcp_config)
439439

440440
return client
441441

apps/application/serializers/application.py

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@
5151
from knowledge.models import File, FileSourceType, Knowledge, KnowledgeScope
5252
from knowledge.serializers.common import BatchMoveSerializer, BatchSerializer
5353
from knowledge.serializers.knowledge import KnowledgeModelSerializer, KnowledgeSerializer
54-
from langchain_mcp_adapters.client import MultiServerMCPClient
54+
from common.utils.mcp_client import create_mcp_client
5555
from maxkb.conf import PROJECT_DIR
5656
from maxkb.const import CONFIG
5757
from models_provider.base_model_provider import ModelTypeConst
@@ -1154,7 +1154,7 @@ class PlayDemoTextRequest(serializers.Serializer):
11541154

11551155

11561156
async def get_mcp_tools(servers):
1157-
client = MultiServerMCPClient(servers)
1157+
client = create_mcp_client(servers)
11581158
return await client.get_tools()
11591159

11601160

@@ -1181,9 +1181,7 @@ def get_mcp_servers(self, instance, with_valid=True):
11811181
self.is_valid(raise_exception=True)
11821182
McpServersSerializer(data=instance).is_valid(raise_exception=True)
11831183
servers = json.loads(instance.get("mcp_servers"))
1184-
for server, config in servers.items():
1185-
if config.get("transport") not in ["sse", "streamable_http"]:
1186-
raise AppApiException(500, _("Only support transport=sse or transport=streamable_http"))
1184+
ToolExecutor().validate_mcp_transport(json.dumps(servers))
11871185
tools = []
11881186
for server in servers:
11891187
tools += [

apps/common/utils/mcp_client.py

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
"""MCP egress policy shared by validation, discovery and tool execution."""
2+
3+
import ipaddress
4+
from functools import partial
5+
6+
from langchain_mcp_adapters.client import MultiServerMCPClient
7+
8+
from common.utils.mcp_network import check_addresses, http_client_factory, parse_url
9+
10+
11+
class InternalMCPConfig(dict):
12+
"""In-memory provenance for configurations generated by ToolExecutor.
13+
14+
Never deserialize user input into this type. JSON round trips deliberately
15+
lose this privilege; keep runtime configurations in memory instead.
16+
"""
17+
18+
19+
def allowed_networks():
20+
from maxkb.const import CONFIG
21+
22+
return tuple(
23+
ipaddress.ip_network(value.strip())
24+
for value in CONFIG.get("MCP_ALLOWED_NETWORKS", "").split(",")
25+
if value.strip()
26+
)
27+
28+
29+
def validate_mcp_servers(servers):
30+
if not isinstance(servers, dict):
31+
raise ValueError("MCP servers must be an object")
32+
networks = allowed_networks()
33+
for config in servers.values():
34+
if not isinstance(config, dict) or config.get("transport") not in ("sse", "streamable_http"):
35+
raise ValueError("Only support transport=sse or transport=streamable_http")
36+
url = parse_url(config.get("url"))
37+
# Do not resolve user hostnames in the web process. The worker checks
38+
# every DNS result at connection time, including nonstandard IP notation.
39+
try:
40+
address = ipaddress.ip_address(url.host)
41+
except ValueError:
42+
continue
43+
check_addresses([str(address)], networks)
44+
45+
46+
def create_mcp_client(servers):
47+
if not isinstance(servers, dict):
48+
raise ValueError("MCP servers must be an object")
49+
connections = {}
50+
networks = allowed_networks()
51+
for name, config in servers.items():
52+
if not isinstance(config, dict):
53+
raise ValueError("MCP server configuration must be an object")
54+
internal = isinstance(config, InternalMCPConfig)
55+
if internal and config.get("transport") == "stdio":
56+
connections[name] = dict(config)
57+
continue
58+
if config.get("transport") not in ("sse", "streamable_http"):
59+
raise ValueError("Only support transport=sse or transport=streamable_http")
60+
url = parse_url(config.get("url"))
61+
if internal:
62+
connection = dict(config)
63+
connection["httpx_client_factory"] = partial(
64+
http_client_factory, url=str(url), internal=True,
65+
networks=(ipaddress.ip_network("127.0.0.1/32"),),
66+
)
67+
else:
68+
from common.utils.mcp_sandbox import sandbox_connection
69+
70+
connection = sandbox_connection(config, networks)
71+
connections[name] = connection
72+
return MultiServerMCPClient(connections)

apps/common/utils/mcp_network.py

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
"""HTTP destination checks usable without Django inside an MCP worker."""
2+
3+
import ipaddress
4+
import socket
5+
import ssl
6+
7+
import anyio
8+
import httpcore
9+
import httpx
10+
11+
12+
class MCPNetworkPolicyError(ValueError):
13+
"""Locally generated, credential-free policy failure safe to report."""
14+
15+
16+
def sandbox_failure_message(error):
17+
# SDK exception groups and chained HTTP errors can embed credentials. Only
18+
# report our own policy text, a numeric HTTP status, or a fixed description.
19+
errors, pending, seen = [], [error], set()
20+
while pending:
21+
current = pending.pop()
22+
if id(current) in seen:
23+
continue
24+
seen.add(id(current))
25+
errors.append(current)
26+
if isinstance(current, BaseExceptionGroup):
27+
pending.extend(current.exceptions)
28+
if current.__cause__ is not None:
29+
pending.append(current.__cause__)
30+
for current in errors:
31+
if isinstance(current, MCPNetworkPolicyError):
32+
return str(current)
33+
for current in errors:
34+
if isinstance(current, httpx.HTTPStatusError):
35+
status = current.response.status_code
36+
if 300 <= status < 400:
37+
return "MCP endpoint returned a redirect; configure its final URL"
38+
return f"MCP endpoint returned HTTP {status}; check endpoint and credentials"
39+
for exception_type, message in (
40+
(ssl.SSLCertVerificationError, "MCP TLS certificate verification failed"),
41+
(socket.gaierror, "MCP hostname resolution failed; check container DNS"),
42+
(PermissionError, "MCP access denied; check sandbox file and network policy"),
43+
((httpx.TimeoutException, TimeoutError), "MCP connection timed out"),
44+
(httpx.ConnectError, "MCP connection failed; check container connectivity and sandbox network policy"),
45+
):
46+
if any(isinstance(current, exception_type) for current in errors):
47+
return message
48+
return "MCP session failed; check endpoint, sandbox setup and network policy"
49+
50+
51+
def parse_url(value):
52+
if not isinstance(value, str) or not value or any(ord(c) <= 32 for c in value):
53+
raise ValueError("Invalid MCP server URL")
54+
try:
55+
url = httpx.URL(value)
56+
if (
57+
url.scheme not in ("http", "https") or not url.host or url.userinfo
58+
or url.fragment or "%" in url.host or "\\" in value
59+
or (url.port is not None and not 1 <= url.port <= 65535)
60+
):
61+
raise ValueError("Invalid MCP server URL")
62+
return url
63+
except (httpx.InvalidURL, ValueError) as exc:
64+
raise ValueError("Invalid MCP server URL") from exc
65+
66+
67+
def check_addresses(addresses, networks):
68+
if not addresses:
69+
raise MCPNetworkPolicyError("MCP server hostname has no addresses")
70+
for address in addresses:
71+
ip = ipaddress.ip_address(address)
72+
if isinstance(ip, ipaddress.IPv6Address) and ip.ipv4_mapped:
73+
ip = ip.ipv4_mapped
74+
# Do not let IPv6 transition mechanisms tunnel to restricted IPv4 hosts.
75+
transition = isinstance(ip, ipaddress.IPv6Address) and (
76+
ip.sixtofour is not None or ip.teredo is not None
77+
or ip in ipaddress.ip_network("64:ff9b::/96")
78+
or ip in ipaddress.ip_network("64:ff9b:1::/48")
79+
)
80+
public = ip.is_global and not ip.is_multicast and not transition
81+
if not public and not any(ip in network for network in networks):
82+
raise MCPNetworkPolicyError("MCP server address is not allowed by the network policy")
83+
84+
85+
class MCPNetworkBackend(httpcore.AnyIOBackend):
86+
def __init__(self, networks):
87+
self.networks = networks
88+
89+
async def connect_tcp(self, host, port, timeout=None, local_address=None, socket_options=None):
90+
try:
91+
with anyio.fail_after(timeout):
92+
# Resolve again at connection time, validate EVERY result, then
93+
# connect to the numeric address. HTTP Host and TLS SNI remain
94+
# the original hostname in httpcore, including certificate checks.
95+
results = await anyio.getaddrinfo(host, port, type=socket.SOCK_STREAM)
96+
addresses = list(dict.fromkeys(item[4][0] for item in results))
97+
check_addresses(addresses, self.networks)
98+
for index, address in enumerate(addresses):
99+
try:
100+
return await super().connect_tcp(
101+
address, port, timeout, local_address, socket_options
102+
)
103+
except (httpcore.ConnectError, httpcore.ConnectTimeout):
104+
if index == len(addresses) - 1:
105+
raise
106+
except TimeoutError as exc:
107+
raise httpcore.ConnectTimeout() from exc
108+
except OSError as exc:
109+
raise httpcore.ConnectError(str(exc)) from exc
110+
111+
112+
class MCPTransport(httpx.AsyncHTTPTransport):
113+
def __init__(self, url, networks, internal=False):
114+
super().__init__(trust_env=False)
115+
self.url = parse_url(url)
116+
self.internal = internal
117+
# HTTPX 0.28 has no public network_backend argument. Keep its standard
118+
# response/error handling and replace only the pool's connection backend.
119+
self._pool._network_backend = MCPNetworkBackend(networks)
120+
121+
async def handle_async_request(self, request):
122+
target = parse_url(str(request.url))
123+
if (target.scheme, target.host, target.port) != (self.url.scheme, self.url.host, self.url.port):
124+
raise ValueError("MCP requests must stay on the configured origin")
125+
if self.internal and target != self.url:
126+
raise ValueError("Internal MCP requests must use the generated endpoint")
127+
return await super().handle_async_request(request)
128+
129+
130+
def http_client_factory(headers=None, timeout=None, auth=None, *, url, networks, internal=False):
131+
return httpx.AsyncClient(
132+
headers=headers,
133+
timeout=timeout if timeout is not None else httpx.Timeout(30, read=300),
134+
auth=auth,
135+
follow_redirects=False,
136+
trust_env=False,
137+
transport=MCPTransport(url, networks, internal),
138+
)

apps/common/utils/mcp_sandbox.py

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
"""Build fixed stdio worker connections; user configuration never selects code."""
2+
3+
import json
4+
import pwd
5+
import sys
6+
from datetime import timedelta
7+
from pathlib import Path
8+
9+
from mcp.types import Implementation
10+
11+
12+
BOOTSTRAP_KEY = "maxkbSandbox"
13+
REMOTE_FIELDS = {"transport", "url", "headers", "timeout", "sse_read_timeout", "terminate_on_close"}
14+
15+
16+
def sandbox_settings():
17+
from maxkb.const import CONFIG
18+
19+
if not sys.platform.startswith("linux") or not bool(int(CONFIG.get("SANDBOX", 1))):
20+
raise ValueError("External MCP requires an enabled Linux sandbox")
21+
account = pwd.getpwnam("sandbox")
22+
sandbox_home = Path(CONFIG.get("SANDBOX_HOME", "/opt/maxkb-app/sandbox"))
23+
library = sandbox_home / "lib/sandbox.so"
24+
if not library.is_file() or not library.with_name(".sandbox.conf").is_file():
25+
raise ValueError("MCP sandbox library or configuration is missing")
26+
return {
27+
"uid": account.pw_uid,
28+
"gid": account.pw_gid,
29+
"library": str(library),
30+
"cwd": str(sandbox_home),
31+
"python_paths": CONFIG.get_sandbox_python_package_paths().split(","),
32+
"memory_mb": int(CONFIG.get("SANDBOX_PYTHON_PROCESS_LIMIT_MEM_MB", "256")),
33+
"cpu_cores": int(CONFIG.get("SANDBOX_PYTHON_PROCESS_LIMIT_CPU_CORES", "1")),
34+
"timeout": int(CONFIG.get("SANDBOX_PYTHON_PROCESS_LIMIT_TIMEOUT_SECONDS", "3600")),
35+
}
36+
37+
38+
def sandbox_connection(config, networks):
39+
settings = sandbox_settings()
40+
# Only transport data goes to the remote client. In particular, ignore user
41+
# command/env/factory/session_kwargs fields and never deserialize Python code.
42+
remote = {key: value for key, value in config.items() if key in REMOTE_FIELDS}
43+
bootstrap = {"connection": remote, "networks": [str(network) for network in networks]}
44+
# Check serializability before launching, and detach from mutable input.
45+
bootstrap = json.loads(json.dumps(bootstrap, allow_nan=False))
46+
return {
47+
"transport": "stdio",
48+
"command": sys.executable,
49+
"args": ["-I", str(Path(__file__).with_name("mcp_sandbox_worker.py"))],
50+
"cwd": settings["cwd"],
51+
"env": {
52+
"LD_PRELOAD": settings["library"],
53+
"MAXKB_MCP_WORKER_SETTINGS": json.dumps(settings),
54+
},
55+
"session_kwargs": {
56+
"read_timeout_seconds": timedelta(seconds=settings["timeout"]),
57+
# This field travels only over the child's stdio pipe. The worker
58+
# removes it before forwarding initialize to the remote server.
59+
"client_info": Implementation(name="maxkb-sandbox", version="1", **{BOOTSTRAP_KEY: bootstrap}),
60+
},
61+
}

0 commit comments

Comments
 (0)