From 8bfce62ad8786084cf64f83ff0f30dbc4efb1acc Mon Sep 17 00:00:00 2001 From: liqiang-fit2cloud Date: Tue, 18 Nov 2025 14:27:48 +0800 Subject: [PATCH] refactor: avoid using temp files in sandbox. --- .../step/chat_step/impl/base_chat_step.py | 12 +--- .../ai_chat_step_node/impl/base_chat_node.py | 12 +--- apps/common/utils/tool_code.py | 72 +++++++++---------- installer/start-maxkb.sh | 2 +- 4 files changed, 35 insertions(+), 63 deletions(-) diff --git a/apps/application/chat_pipeline/step/chat_step/impl/base_chat_step.py b/apps/application/chat_pipeline/step/chat_step/impl/base_chat_step.py index cd17d4a10..ebbc4c131 100644 --- a/apps/application/chat_pipeline/step/chat_step/impl/base_chat_step.py +++ b/apps/application/chat_pipeline/step/chat_step/impl/base_chat_step.py @@ -200,14 +200,6 @@ class BaseChatStep(IChatStep): mcp_enable, mcp_tool_ids, mcp_servers, mcp_source, tool_enable, tool_ids, mcp_output_enable) def get_details(self, manage, **kwargs): - # 删除临时生成的MCP代码文件 - if self.context.get('execute_ids'): - executor = ToolExecutor(CONFIG.get('SANDBOX')) - # 清理工具代码文件，延时删除，避免文件被占用 - for tool_id in self.context.get('execute_ids'): - code_path = f'{executor.sandbox_path}/execute/{tool_id}.py' - if os.path.exists(code_path): - os.remove(code_path) return { 'step_type': 'chat_step', 'run_time': self.context['run_time'], @@ -254,7 +246,6 @@ class BaseChatStep(IChatStep): if tool_enable: if tool_ids and len(tool_ids) > 0: # 如果有工具ID，则将其转换为MCP self.context['tool_ids'] = tool_ids - self.context['execute_ids'] = [] for tool_id in tool_ids: tool = QuerySet(Tool).filter(id=tool_id).first() if tool is None or tool.is_active is False: @@ -264,9 +255,8 @@ class BaseChatStep(IChatStep): params = json.loads(rsa_long_decrypt(tool.init_params)) else: params = {} - _id, tool_config = executor.get_tool_mcp_config(tool.code, params) + tool_config = executor.get_tool_mcp_config(tool.code, params) - self.context['execute_ids'].append(_id) mcp_servers_config[str(tool.id)] = tool_config if len(mcp_servers_config) > 0: diff --git a/apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py b/apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py index 20dc22b25..01bf8ce00 100644 --- a/apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py +++ b/apps/application/flow/step_node/ai_chat_step_node/impl/base_chat_node.py @@ -233,7 +233,6 @@ class BaseChatNode(IChatNode): if tool_enable: if tool_ids and len(tool_ids) > 0: # 如果有工具ID，则将其转换为MCP self.context['tool_ids'] = tool_ids - self.context['execute_ids'] = [] for tool_id in tool_ids: tool = QuerySet(Tool).filter(id=tool_id).first() if not tool.is_active: @@ -243,9 +242,8 @@ class BaseChatNode(IChatNode): params = json.loads(rsa_long_decrypt(tool.init_params)) else: params = {} - _id, tool_config = executor.get_tool_mcp_config(tool.code, params) + tool_config = executor.get_tool_mcp_config(tool.code, params) - self.context['execute_ids'].append(_id) mcp_servers_config[str(tool.id)] = tool_config if len(mcp_servers_config) > 0: @@ -307,14 +305,6 @@ class BaseChatNode(IChatNode): return result def get_details(self, index: int, **kwargs): - # 删除临时生成的MCP代码文件 - if self.context.get('execute_ids'): - executor = ToolExecutor(CONFIG.get('SANDBOX')) - # 清理工具代码文件，延时删除，避免文件被占用 - for tool_id in self.context.get('execute_ids'): - code_path = f'{executor.sandbox_path}/execute/{tool_id}.py' - if os.path.exists(code_path): - os.remove(code_path) return { 'name': self.node.properties.get('stepName'), "index": index, diff --git a/apps/common/utils/tool_code.py b/apps/common/utils/tool_code.py index e75f70869..4c764965b 100644 --- a/apps/common/utils/tool_code.py +++ b/apps/common/utils/tool_code.py @@ -1,24 +1,25 @@ # coding=utf-8 import ast +import base64 +import gzip import json import os +import socket import subprocess import sys -from textwrap import dedent -import socket import uuid_utils.compat as uuid +from common.utils.logger import maxkb_logger from django.utils.translation import gettext_lazy as _ from maxkb.const import BASE_DIR, CONFIG from maxkb.const import PROJECT_DIR -from common.utils.logger import maxkb_logger -import threading +from textwrap import dedent python_directory = sys.executable -class ToolExecutor: +class ToolExecutor: _dir_initialized = False - _lock = threading.Lock() + def __init__(self, sandbox=False): self.sandbox = sandbox if sandbox: @@ -29,29 +30,22 @@ class ToolExecutor: self.user = None self.banned_keywords = CONFIG.get("SANDBOX_PYTHON_BANNED_KEYWORDS", 'nothing_is_banned').split(','); self.sandbox_so_path = f'{self.sandbox_path}/sandbox.so' - with ToolExecutor._lock: - self._init_dir() + self._init_dir() def _init_dir(self): if ToolExecutor._dir_initialized: # 只初始化一次 return - execute_file_path = os.path.join(self.sandbox_path, 'execute') - os.makedirs(execute_file_path, 0o500, exist_ok=True) - result_file_path = os.path.join(self.sandbox_path, 'result') - os.makedirs(result_file_path, 0o300, exist_ok=True) if self.sandbox: os.system(f"chown {self.user}:root {self.sandbox_path}") - os.system(f"chown -R {self.user}:root {execute_file_path}") - os.system(f"chown -R {self.user}:root {result_file_path}") os.chmod(self.sandbox_path, 0o550) if CONFIG.get("SANDBOX_TMP_DIR_ENABLED", '0') == "1": tmp_dir_path = os.path.join(self.sandbox_path, 'tmp') os.makedirs(tmp_dir_path, 0o700, exist_ok=True) os.system(f"chown -R {self.user}:root {tmp_dir_path}") + if os.path.exists(self.sandbox_so_path): + os.chmod(self.sandbox_so_path, 0o440) try: - if os.path.exists(self.sandbox_so_path): - os.chmod(self.sandbox_so_path, 0o440) # 初始化host黑名单 banned_hosts_file_path = f'{self.sandbox_path}/.SANDBOX_BANNED_HOSTS' if os.path.exists(banned_hosts_file_path): @@ -74,13 +68,13 @@ class ToolExecutor: _id = str(uuid.uuid7()) success = '{"code":200,"msg":"成功","data":exec_result}' err = '{"code":500,"msg":str(e),"data":None}' - result_path = f'{self.sandbox_path}/result/{_id}.result' python_paths = CONFIG.get_sandbox_python_package_paths().split(',') _exec_code = f""" try: import os import sys import json + import base64 path_to_exclude = ['/opt/py3/lib/python3.11/site-packages', '/opt/maxkb-app/apps'] sys.path = [p for p in sys.path if p not in path_to_exclude] sys.path += {python_paths} @@ -92,21 +86,21 @@ try: for local in locals_v: globals_v[local] = locals_v[local] exec_result=f(**keywords) - with open({result_path!a}, 'w') as file: - file.write(json.dumps({success}, default=str)) + print(f"{_id}:" + base64.b64encode(json.dumps({success}, default=str).encode()).decode()) except Exception as e: - with open({result_path!a}, 'w') as file: - file.write(json.dumps({err})) + print(f"{_id}:" + base64.b64encode(json.dumps({err}, default=str).encode()).decode()) """ if self.sandbox: - subprocess_result = self._exec_sandbox(_exec_code, _id) + subprocess_result = self._exec_sandbox(_exec_code) else: subprocess_result = self._exec(_exec_code) if subprocess_result.returncode == 1: raise Exception(subprocess_result.stderr) - with open(result_path, 'r') as file: - result = json.loads(file.read()) - os.remove(result_path) + lines = subprocess_result.stdout.splitlines() + result_line = [line for line in lines if line.startswith(_id)] + if not result_line: + raise Exception("No result found.") + result = json.loads(base64.b64decode(result_line[0].split(":", 1)[1]).decode()) if result.get('code') == 200: return result.get('data') raise Exception(result.get('msg')) @@ -194,18 +188,16 @@ exec({dedent(code)!a}) """ def get_tool_mcp_config(self, code, params): - code = self.generate_mcp_server_code(code, params) - - _id = uuid.uuid7() - code_path = f'{self.sandbox_path}/execute/{_id}.py' - with open(code_path, 'w') as f: - f.write(code) + _code = self.generate_mcp_server_code(code, params) + maxkb_logger.debug(f"Python code of mcp tool: {_code}") + compressed_and_base64_encoded_code_str = base64.b64encode(gzip.compress(_code.encode())).decode() if self.sandbox: tool_config = { 'command': 'su', 'args': [ '-s', sys.executable, - '-c', f"exec(open('{code_path}', 'r').read())", + '-c', + f'import base64,gzip; exec(gzip.decompress(base64.b64decode(\'{compressed_and_base64_encoded_code_str}\')).decode())', self.user, ], 'cwd': self.sandbox_path, @@ -217,24 +209,24 @@ exec({dedent(code)!a}) else: tool_config = { 'command': sys.executable, - 'args': [code_path], + 'args': f'import base64,gzip; exec(gzip.decompress(base64.b64decode(\'{compressed_and_base64_encoded_code_str}\')).decode())', 'transport': 'stdio', } - return _id, tool_config + return tool_config - def _exec_sandbox(self, _code, _id): - exec_python_file = f'{self.sandbox_path}/execute/{_id}.py' - with open(exec_python_file, 'w') as file: - file.write(_code) + def _exec_sandbox(self, _code): kwargs = {'cwd': BASE_DIR} kwargs['env'] = { 'LD_PRELOAD': self.sandbox_so_path, } + maxkb_logger.debug(f"Sandbox execute code: {_code}") + compressed_and_base64_encoded_code_str = base64.b64encode(gzip.compress(_code.encode())).decode() subprocess_result = subprocess.run( - ['su', '-s', python_directory, '-c', "exec(open('" + exec_python_file + "').read())", self.user], + ['su', '-s', python_directory, '-c', + f'import base64,gzip; exec(gzip.decompress(base64.b64decode(\'{compressed_and_base64_encoded_code_str}\')).decode())', + self.user], text=True, capture_output=True, **kwargs) - os.remove(exec_python_file) return subprocess_result def validate_banned_keywords(self, code_str): diff --git a/installer/start-maxkb.sh b/installer/start-maxkb.sh index c9beef809..5ffd5ca1e 100644 --- a/installer/start-maxkb.sh +++ b/installer/start-maxkb.sh @@ -2,8 +2,8 @@ if [ ! -d /opt/maxkb/logs ]; then mkdir -p /opt/maxkb/logs - chmod 700 /opt/maxkb/logs fi +chmod -R 700 /opt/maxkb/logs if [ ! -d /opt/maxkb/local ]; then mkdir -p /opt/maxkb/local chmod 700 /opt/maxkb/local