diff --git a/apps/common/utils/tool_code.py b/apps/common/utils/tool_code.py index 3ff8ef04e2c..901ff235335 100644 --- a/apps/common/utils/tool_code.py +++ b/apps/common/utils/tool_code.py @@ -5,41 +5,40 @@ import gzip import json import os -import pwd import random -import resource import socket import subprocess import sys import tempfile import time -from contextlib import contextmanager, suppress +from contextlib import contextmanager +from contextlib import suppress from textwrap import dedent import uuid_utils.compat as uuid from django.utils.translation import gettext_lazy as _ -from maxkb.const import BASE_DIR, CONFIG, PROJECT_DIR from common.utils.logger import maxkb_logger +from maxkb.const import BASE_DIR, CONFIG +from maxkb.const import PROJECT_DIR -_enable_sandbox = bool(int(CONFIG.get("SANDBOX", 0))) -_run_user = "sandbox" if _enable_sandbox else getpass.getuser() -_sandbox_path = ( - CONFIG.get("SANDBOX_HOME", "/opt/maxkb-app/sandbox") - if _enable_sandbox - else os.path.join(PROJECT_DIR, "data", "sandbox") -) -_sandbox_python_sys_path = CONFIG.get_sandbox_python_package_paths().split(",") -_process_limit_timeout_seconds = int(CONFIG.get("SANDBOX_PYTHON_PROCESS_LIMIT_TIMEOUT_SECONDS", "3600")) -_process_limit_cpu_cores = ( - min(max(int(CONFIG.get("SANDBOX_PYTHON_PROCESS_LIMIT_CPU_CORES", "1")), 1), len(os.sched_getaffinity(0))) - if sys.platform.startswith("linux") - else os.cpu_count() -) # 只支持linux，window和mac不支持 -_process_limit_mem_mb = int(CONFIG.get("SANDBOX_PYTHON_PROCESS_LIMIT_MEM_MB", "256")) +pwd = None +resource = None +_enable_sandbox = bool(int(CONFIG.get('SANDBOX', 0))) +_run_user = 'sandbox' if _enable_sandbox else getpass.getuser() +_sandbox_path = CONFIG.get("SANDBOX_HOME", '/opt/maxkb-app/sandbox') if _enable_sandbox else os.path.join(PROJECT_DIR, + 'data', + 'sandbox') +_sandbox_python_sys_path = CONFIG.get_sandbox_python_package_paths().split(',') +_process_limit_timeout_seconds = int(CONFIG.get("SANDBOX_PYTHON_PROCESS_LIMIT_TIMEOUT_SECONDS", '3600')) +_process_limit_cpu_cores = min(max(int(CONFIG.get("SANDBOX_PYTHON_PROCESS_LIMIT_CPU_CORES", '1')), 1), + len(os.sched_getaffinity(0))) if sys.platform.startswith( + "linux") else os.cpu_count() # 只支持linux，window和mac不支持 +_process_limit_mem_mb = int(CONFIG.get("SANDBOX_PYTHON_PROCESS_LIMIT_MEM_MB", '256')) class ToolExecutor: + def __init__(self): pass @@ -50,9 +49,8 @@ def init_sandbox_dir(): return try: # 只初始化一次 - fd = os.open( - os.path.join(PROJECT_DIR, "tmp", "tool_executor_init_dir.lock"), os.O_CREAT | os.O_EXCL | os.O_WRONLY - ) + fd = os.open(os.path.join(PROJECT_DIR, 'tmp', 'tool_executor_init_dir.lock'), + os.O_CREAT | os.O_EXCL | os.O_WRONLY) os.close(fd) except FileExistsError: # 文件已存在 → 已初始化过 @@ -62,32 +60,30 @@ def init_sandbox_dir(): os.system("chmod -R g-rwx /dev/shm /dev/mqueue") os.system("chmod o-rwx /run/postgresql") except Exception as e: - maxkb_logger.warning(f"Exception: {e}", exc_info=True) + maxkb_logger.warning(f'Exception: {e}', exc_info=True) pass - if CONFIG.get("SANDBOX_TMP_DIR_ENABLED", "0") == "1": + if CONFIG.get("SANDBOX_TMP_DIR_ENABLED", '0') == "1": os.system("chmod g+rwx /tmp") # 初始化sandbox配置文件 - sandbox_lib_path = os.path.dirname(f"{_sandbox_path}/lib/sandbox.so") - sandbox_conf_file_path = f"{sandbox_lib_path}/.sandbox.conf" + sandbox_lib_path = os.path.dirname(f'{_sandbox_path}/lib/sandbox.so') + sandbox_conf_file_path = f'{sandbox_lib_path}/.sandbox.conf' if os.path.exists(sandbox_conf_file_path): os.remove(sandbox_conf_file_path) - banned_hosts = CONFIG.get("SANDBOX_PYTHON_BANNED_HOSTS", "").strip() - allow_dl_paths = CONFIG.get("SANDBOX_PYTHON_ALLOW_DL_PATHS", "").strip() - allow_dl_open = CONFIG.get("SANDBOX_PYTHON_ALLOW_DL_OPEN", "0") - allow_subprocess = CONFIG.get("SANDBOX_PYTHON_ALLOW_SUBPROCESS", "0") - allow_syscall = CONFIG.get("SANDBOX_PYTHON_ALLOW_SYSCALL", "0") + banned_hosts = CONFIG.get("SANDBOX_PYTHON_BANNED_HOSTS", '').strip() + allow_dl_paths = CONFIG.get("SANDBOX_PYTHON_ALLOW_DL_PATHS", '').strip() + allow_dl_open = CONFIG.get("SANDBOX_PYTHON_ALLOW_DL_OPEN", '0') + allow_subprocess = CONFIG.get("SANDBOX_PYTHON_ALLOW_SUBPROCESS", '0') + allow_syscall = CONFIG.get("SANDBOX_PYTHON_ALLOW_SYSCALL", '0') if banned_hosts: hostname = socket.gethostname() local_ip = socket.gethostbyname(hostname) banned_hosts = f"{banned_hosts},{local_ip}" banned_hosts = ",".join( - s.strip() for s in banned_hosts.split(",") if s.strip() and s.strip().lower() != hostname.lower() - ) - with open(sandbox_conf_file_path, "w", encoding="utf-8") as f: + s.strip() for s in banned_hosts.split(",") if s.strip() and s.strip().lower() != hostname.lower()) + with open(sandbox_conf_file_path, "w", encoding='utf-8') as f: f.write(f"SANDBOX_PYTHON_BANNED_HOSTS={banned_hosts}\n") f.write( - f"SANDBOX_PYTHON_ALLOW_DL_PATHS={','.join(sorted(set(filter(None, sys.path + _sandbox_python_sys_path + allow_dl_paths.split(',')))))}\n" - ) + f"SANDBOX_PYTHON_ALLOW_DL_PATHS={','.join(sorted(set(filter(None, sys.path + _sandbox_python_sys_path + allow_dl_paths.split(',')))))}\n") f.write(f"SANDBOX_PYTHON_ALLOW_DL_OPEN={allow_dl_open}\n") f.write(f"SANDBOX_PYTHON_ALLOW_SUBPROCESS={allow_subprocess}\n") f.write(f"SANDBOX_PYTHON_ALLOW_SYSCALL={allow_syscall}\n") @@ -96,18 +92,12 @@ def init_sandbox_dir(): try: init_sandbox_dir() except Exception as e: - maxkb_logger.error(f"Exception: {e}", exc_info=True) + maxkb_logger.error(f'Exception: {e}', exc_info=True) def exec_code(self, code_str, keywords, function_name=None): _id = str(uuid.uuid7()) - action_function = ( - f"({function_name!a}, locals_v.get({function_name!a}))" if function_name else "locals_v.popitem()" - ) - set_run_user = ( - f"os.setgid({pwd.getpwnam(_run_user).pw_gid});os.setuid({pwd.getpwnam(_run_user).pw_uid});" - if _enable_sandbox - else "" - ) + action_function = f'({function_name !a}, locals_v.get({function_name !a}))' if function_name else 'locals_v.popitem()' + set_run_user = f'os.setgid({pwd.getpwnam(_run_user).pw_gid});os.setuid({pwd.getpwnam(_run_user).pw_uid});' if _enable_sandbox else '' _exec_code = f""" try: import os, sys, json @@ -136,7 +126,7 @@ def exec_code(self, code_str, keywords, function_name=None): sys.stdout.flush() """ maxkb_logger.debug(f"Tool execution({_id}) execute code: {_exec_code}") - with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=True) as f: + with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=True) as f: f.write(_exec_code) f.flush() with execution_timer(_id): @@ -149,10 +139,10 @@ def exec_code(self, code_str, keywords, function_name=None): last_line = lines[-2] if not last_line.startswith(_id): raise Exception("No result found.") - result = json.loads(last_line[len(_id) :]) - if result.get("code") == 200: - return result.get("data") - raise Exception(result.get("msg") + (f"\n{subprocess_result.stderr}" if subprocess_result.stderr else "")) + result = json.loads(last_line[len(_id):]) + if result.get('code') == 200: + return result.get('data') + raise Exception(result.get('msg') + (f'\n{subprocess_result.stderr}' if subprocess_result.stderr else '')) def _generate_mcp_server_code(self, _code, params, name=None, description=None, tool_id=None): # 解析代码,提取导入语句和函数定义 @@ -167,7 +157,7 @@ def _generate_mcp_server_code(self, _code, params, name=None, description=None, if isinstance(node, ast.Import) or isinstance(node, ast.ImportFrom): imports.append(ast.unparse(node)) elif isinstance(node, ast.FunctionDef): - if node.name.startswith("_"): + if node.name.startswith('_'): other_code.append(ast.unparse(node)) continue # 修改函数参数以包含 params 中的默认值 @@ -201,25 +191,9 @@ def _generate_mcp_server_code(self, _code, params, name=None, description=None, # 将不支持 JSON Schema 的参数类型注解替换为 Any， # 避免 FastMCP/Pydantic 生成 schema 时崩溃（如 requests.Response） _safe_annotation_names = { - "str", - "int", - "float", - "bool", - "dict", - "list", - "tuple", - "set", - "bytes", - "Any", - "Optional", - "Union", - "List", - "Dict", - "Tuple", - "Set", - "Sequence", - "None", - "NoneType", + 'str', 'int', 'float', 'bool', 'dict', 'list', 'tuple', + 'set', 'bytes', 'Any', 'Optional', 'Union', 'List', + 'Dict', 'Tuple', 'Set', 'Sequence', 'None', 'NoneType', } def _is_safe_annotation(node_ann): @@ -241,9 +215,9 @@ def _is_safe_annotation(node_ann): for arg in node.args.args: if not _is_safe_annotation(arg.annotation): - arg.annotation = ast.Name(id="Any", ctx=ast.Load()) + arg.annotation = ast.Name(id='Any', ctx=ast.Load()) # 修改返回类型注解为 Result - node.returns = ast.Name(id="Result", ctx=ast.Load()) + node.returns = ast.Name(id='Result', ctx=ast.Load()) # 修改 return 语句为 return Result(result=..., tool_id=...) class ReturnTransformer(ast.NodeTransformer): @@ -255,24 +229,24 @@ def visit_Return(self, node): # return 语句没有返回值 new_return = ast.Return( value=ast.Call( - func=ast.Name(id="Result", ctx=ast.Load()), + func=ast.Name(id='Result', ctx=ast.Load()), args=[], keywords=[ - ast.keyword(arg="result", value=ast.Constant(value=None)), - ast.keyword(arg="tool_id", value=ast.Constant(value=tool_id)), - ], + ast.keyword(arg='result', value=ast.Constant(value=None)), + ast.keyword(arg='tool_id', value=ast.Constant(value=tool_id)) + ] ) ) else: # return 语句有返回值 new_return = ast.Return( value=ast.Call( - func=ast.Name(id="Result", ctx=ast.Load()), + func=ast.Name(id='Result', ctx=ast.Load()), args=[], keywords=[ - ast.keyword(arg="result", value=node.value), - ast.keyword(arg="tool_id", value=ast.Constant(value=tool_id)), - ], + ast.keyword(arg='result', value=node.value), + ast.keyword(arg='tool_id', value=ast.Constant(value=tool_id)) + ] ) ) return ast.copy_location(new_return, node) @@ -282,7 +256,7 @@ def visit_Return(self, node): ast.fix_missing_locations(node) func_code = ast.unparse(node) # 有些模型不支持name是中文,例如: deepseek, 其他模型未知 - escaped_desc = (name + " " + description).replace("\n", " ").replace("'", " ") + escaped_desc = (name + ' ' + description).replace('\n', ' ').replace("'", " ") functions.append(f"@mcp.tool(description='{escaped_desc}')\n{func_code}\n") else: other_code.append(ast.unparse(node)) @@ -294,19 +268,15 @@ def visit_Return(self, node): code_parts.append(f"\nclass Result(BaseModel):") code_parts.append(f"\n\tresult: Any") code_parts.append(f"\n\ttool_id: str\n") - code_parts.append(f'\nmcp = FastMCP("{uuid.uuid7()}")\n') + code_parts.append(f"\nmcp = FastMCP(\"{uuid.uuid7()}\")\n") code_parts.extend(other_code) code_parts.extend(functions) - code_parts.append('\nmcp.run(transport="stdio")\n') + code_parts.append("\nmcp.run(transport=\"stdio\")\n") return "\n".join(code_parts) def generate_mcp_server_code(self, code_str, params, name, description, tool_id): code = self._generate_mcp_server_code(code_str, params, name, description, tool_id) - set_run_user = ( - f"os.setgid({pwd.getpwnam(_run_user).pw_gid});os.setuid({pwd.getpwnam(_run_user).pw_uid});" - if _enable_sandbox - else "" - ) + set_run_user = f'os.setgid({pwd.getpwnam(_run_user).pw_gid});os.setuid({pwd.getpwnam(_run_user).pw_uid});' if _enable_sandbox else '' return f""" import os, sys, logging logging.basicConfig(level=logging.WARNING) @@ -325,45 +295,40 @@ def get_tool_mcp_config(self, tool, params): maxkb_logger.debug(f"Python code of mcp tool: {_code}") compressed_and_base64_encoded_code_str = base64.b64encode(gzip.compress(_code.encode())).decode() tool_config = { - "command": sys.executable, - "args": [ - "-c", - f"import base64,gzip; exec(gzip.decompress(base64.b64decode('{compressed_and_base64_encoded_code_str}')).decode())", + 'command': sys.executable, + 'args': [ + '-c', + f'import base64,gzip; exec(gzip.decompress(base64.b64decode(\'{compressed_and_base64_encoded_code_str}\')).decode())', ], - "cwd": _sandbox_path, - "env": { - "LD_PRELOAD": f"{_sandbox_path}/lib/sandbox.so", + 'cwd': _sandbox_path, + 'env': { + 'LD_PRELOAD': f'{_sandbox_path}/lib/sandbox.so', }, - "transport": "stdio", + 'transport': 'stdio', } return tool_config def get_app_mcp_config(self, api_key): app_config = { - "url": f"http://127.0.0.1:8080{CONFIG.get_chat_path()}/api/mcp", - "transport": "streamable_http", - "headers": { - "Authorization": f"Bearer {api_key}", + 'url': f'http://127.0.0.1:8080{CONFIG.get_chat_path()}/api/mcp', + 'transport': 'streamable_http', + 'headers': { + 'Authorization': f'Bearer {api_key}', }, } return app_config def _exec(self, execute_file, _id): - kwargs = { - "cwd": BASE_DIR, - "env": { - "LD_PRELOAD": f"{_sandbox_path}/lib/sandbox.so", - "_ID": _id, - }, - } + kwargs = {'cwd': BASE_DIR, 'env': { + 'LD_PRELOAD': f'{_sandbox_path}/lib/sandbox.so', + '_ID': _id, + }} def _set_resource_limit(): - if not _enable_sandbox or not sys.platform.startswith("linux"): - return - with suppress(Exception): - resource.setrlimit(resource.RLIMIT_AS, (_process_limit_mem_mb * 1024 * 1024,) * 2) - with suppress(Exception): - os.sched_setaffinity(0, set(random.sample(list(os.sched_getaffinity(0)), _process_limit_cpu_cores))) + if not _enable_sandbox or not sys.platform.startswith("linux"): return + with suppress(Exception): resource.setrlimit(resource.RLIMIT_AS, (_process_limit_mem_mb * 1024 * 1024,) * 2) + with suppress(Exception): os.sched_setaffinity(0, set(random.sample(list(os.sched_getaffinity(0)), + _process_limit_cpu_cores))) try: subprocess_result = subprocess.run( @@ -372,7 +337,7 @@ def _set_resource_limit(): text=True, capture_output=True, **kwargs, - preexec_fn=_set_resource_limit, + preexec_fn=_set_resource_limit ) return subprocess_result except subprocess.TimeoutExpired: @@ -381,8 +346,8 @@ def _set_resource_limit(): def validate_mcp_transport(self, code_str): servers = json.loads(code_str) for server, config in servers.items(): - if config.get("transport") not in ["sse", "streamable_http"]: - raise Exception(_("Only support transport=sse or transport=streamable_http")) + if config.get('transport') not in ['sse', 'streamable_http']: + raise Exception(_('Only support transport=sse or transport=streamable_http')) @contextmanager