From f8ada9a110c4dbef8c3c2636c78847ecd621ece7 Mon Sep 17 00:00:00 2001 From: liqiang-fit2cloud Date: Thu, 4 Dec 2025 15:40:50 +0800 Subject: [PATCH] refactor: remove init parameter from class ToolExecutor. --- .../step/chat_step/impl/base_chat_step.py | 2 +- .../ai_chat_step_node/impl/base_chat_node.py | 2 +- .../tool_lib_node/impl/base_tool_lib_node.py | 2 +- .../tool_node/impl/base_tool_node.py | 2 +- apps/common/utils/tool_code.py | 58 +++++++++---------- .../serializers/knowledge_workflow.py | 2 +- apps/tools/serializers/tool.py | 2 +- 7 files changed, 33 insertions(+), 37 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 ed953e13c..2d5e9b5a8 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 @@ -253,7 +253,7 @@ class BaseChatStep(IChatStep): tool = QuerySet(Tool).filter(id=tool_id).first() if tool is None or tool.is_active is False: continue - executor = ToolExecutor(CONFIG.get('SANDBOX')) + executor = ToolExecutor() if tool.init_params is not None: params = json.loads(rsa_long_decrypt(tool.init_params)) else: 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 630b2e9b6..85d046b2e 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 @@ -237,7 +237,7 @@ class BaseChatNode(IChatNode): tool = QuerySet(Tool).filter(id=tool_id).first() if not tool.is_active: continue - executor = ToolExecutor(CONFIG.get('SANDBOX')) + executor = ToolExecutor() if tool.init_params is not None: params = json.loads(rsa_long_decrypt(tool.init_params)) else: diff --git a/apps/application/flow/step_node/tool_lib_node/impl/base_tool_lib_node.py b/apps/application/flow/step_node/tool_lib_node/impl/base_tool_lib_node.py index 66dc921b8..7a0690e93 100644 --- a/apps/application/flow/step_node/tool_lib_node/impl/base_tool_lib_node.py +++ b/apps/application/flow/step_node/tool_lib_node/impl/base_tool_lib_node.py @@ -22,7 +22,7 @@ from common.utils.tool_code import ToolExecutor from maxkb.const import CONFIG from tools.models import Tool -function_executor = ToolExecutor(CONFIG.get('SANDBOX')) +function_executor = ToolExecutor() def write_context(step_variable: Dict, global_variable: Dict, node, workflow): diff --git a/apps/application/flow/step_node/tool_node/impl/base_tool_node.py b/apps/application/flow/step_node/tool_node/impl/base_tool_node.py index 6090b6dd1..5cc9e1711 100644 --- a/apps/application/flow/step_node/tool_node/impl/base_tool_node.py +++ b/apps/application/flow/step_node/tool_node/impl/base_tool_node.py @@ -17,7 +17,7 @@ from application.flow.step_node.tool_node.i_tool_node import IToolNode from common.utils.tool_code import ToolExecutor from maxkb.const import CONFIG -function_executor = ToolExecutor(CONFIG.get('SANDBOX')) +function_executor = ToolExecutor() def write_context(step_variable: Dict, global_variable: Dict, node, workflow): diff --git a/apps/common/utils/tool_code.py b/apps/common/utils/tool_code.py index 5eb68d4dc..a124eae6c 100644 --- a/apps/common/utils/tool_code.py +++ b/apps/common/utils/tool_code.py @@ -19,24 +19,21 @@ from maxkb.const import BASE_DIR, CONFIG from maxkb.const import PROJECT_DIR from textwrap import dedent +_enable_sandbox = bool(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') +_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: - enable_sandbox = bool(CONFIG.get('SANDBOX', 0)) - sandbox_path = CONFIG.get("SANDBOX_HOME", '/opt/maxkb-app/sandbox') if enable_sandbox else os.path.join(PROJECT_DIR, 'data', 'sandbox') - process_timeout_seconds = int(CONFIG.get("SANDBOX_PYTHON_PROCESS_TIMEOUT_SECONDS", '3600')) - process_limit_mem_mb = int(CONFIG.get("SANDBOX_PYTHON_PROCESS_LIMIT_MEM_MB", '256')) - 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不支持 - - def __init__(self, sandbox=False): - self.sandbox = sandbox - if sandbox: - self.user = 'sandbox' - else: - self.user = getpass.getuser() + def __init__(self): + pass @staticmethod def init_sandbox_dir(): - if not ToolExecutor.enable_sandbox: + if not _enable_sandbox: # 不启用sandbox就不初始化目录 return try: @@ -57,7 +54,7 @@ class ToolExecutor: if CONFIG.get("SANDBOX_TMP_DIR_ENABLED", '0') == "1": os.system("chmod g+rwx /tmp") # 初始化sandbox配置文件 - sandbox_lib_path = os.path.dirname(f'{ToolExecutor.sandbox_path}/lib/sandbox.so') + 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) @@ -70,8 +67,12 @@ class ToolExecutor: with open(sandbox_conf_file_path, "w") as f: f.write(f"SANDBOX_PYTHON_BANNED_HOSTS={banned_hosts}\n") f.write(f"SANDBOX_PYTHON_ALLOW_SUBPROCESS={allow_subprocess}\n") - os.system(f"chmod -R 550 {ToolExecutor.sandbox_path}") + os.system(f"chmod -R 550 {_sandbox_path}") + try: + init_sandbox_dir() + except Exception as e: + maxkb_logger.error(f'Exception: {e}', exc_info=True) def exec_code(self, code_str, keywords, function_name=None): _id = str(uuid.uuid7()) @@ -79,7 +80,7 @@ class ToolExecutor: err = '{"code":500,"msg":str(e),"data":None}' action_function = f'({function_name !a}, locals_v.get({function_name !a}))' if function_name else 'locals_v.popitem()' python_paths = CONFIG.get_sandbox_python_package_paths().split(',') - set_run_user = f'os.setgid({pwd.getpwnam(self.user).pw_gid});os.setuid({pwd.getpwnam(self.user).pw_uid});' if self.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 '' _exec_code = f""" try: import os, sys, json, base64, builtins @@ -98,7 +99,7 @@ try: exec_result=f(**keywords) builtins.print("\\n{_id}:"+base64.b64encode(json.dumps({success}, default=str).encode()).decode(), flush=True) except Exception as e: - if isinstance(e, MemoryError): e = Exception("Cannot allocate more memory: exceeded the limit of {ToolExecutor.process_limit_mem_mb} MB.") + if isinstance(e, MemoryError): e = Exception("Cannot allocate more memory: exceeded the limit of {_process_limit_mem_mb} MB.") builtins.print("\\n{_id}:"+base64.b64encode(json.dumps({err}, default=str).encode()).decode(), flush=True) """ maxkb_logger.debug(f"Sandbox execute code: {_exec_code}") @@ -184,7 +185,7 @@ except Exception as e: def generate_mcp_server_code(self, code_str, params): python_paths = CONFIG.get_sandbox_python_package_paths().split(',') code = self._generate_mcp_server_code(code_str, params) - set_run_user = f'os.setgid({pwd.getpwnam(self.user).pw_gid});os.setuid({pwd.getpwnam(self.user).pw_uid});' if self.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) @@ -208,9 +209,9 @@ exec({dedent(code)!a}) '-c', f'import base64,gzip; exec(gzip.decompress(base64.b64decode(\'{compressed_and_base64_encoded_code_str}\')).decode())', ], - 'cwd': ToolExecutor.sandbox_path, + 'cwd': _sandbox_path, 'env': { - 'LD_PRELOAD': f'{ToolExecutor.sandbox_path}/lib/sandbox.so', + 'LD_PRELOAD': f'{_sandbox_path}/lib/sandbox.so', }, 'transport': 'stdio', } @@ -218,31 +219,26 @@ exec({dedent(code)!a}) def _exec(self, execute_file): kwargs = {'cwd': BASE_DIR, 'env': { - 'LD_PRELOAD': f'{ToolExecutor.sandbox_path}/lib/sandbox.so', + 'LD_PRELOAD': f'{_sandbox_path}/lib/sandbox.so', }} try: subprocess_result = subprocess.run( [sys.executable, execute_file], - timeout=ToolExecutor.process_timeout_seconds, + timeout=_process_limit_timeout_seconds, text=True, capture_output=True, **kwargs, - preexec_fn=(lambda: None if (not self.sandbox or not sys.platform.startswith("linux")) else ( - resource.setrlimit(resource.RLIMIT_AS, (ToolExecutor.process_limit_mem_mb * 1024 * 1024,) * 2), - os.sched_setaffinity(0, set(random.sample(list(os.sched_getaffinity(0)), ToolExecutor.process_limit_cpu_cores))) + preexec_fn=(lambda: None if (not _enable_sandbox or not sys.platform.startswith("linux")) else ( + resource.setrlimit(resource.RLIMIT_AS, (_process_limit_mem_mb * 1024 * 1024,) * 2), + os.sched_setaffinity(0, set(random.sample(list(os.sched_getaffinity(0)), _process_limit_cpu_cores))) )) ) return subprocess_result except subprocess.TimeoutExpired: - raise Exception(_(f"Process execution timed out after {ToolExecutor.process_timeout_seconds} seconds.")) + raise Exception(_(f"Process execution timed out after {_process_limit_timeout_seconds} seconds.")) 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')) - -try: - ToolExecutor.init_sandbox_dir() -except Exception as e: - maxkb_logger.error(f'Exception: {e}', exc_info=True) \ No newline at end of file diff --git a/apps/knowledge/serializers/knowledge_workflow.py b/apps/knowledge/serializers/knowledge_workflow.py index 74b3b4154..f5d4fdabc 100644 --- a/apps/knowledge/serializers/knowledge_workflow.py +++ b/apps/knowledge/serializers/knowledge_workflow.py @@ -28,7 +28,7 @@ from system_manage.serializers.user_resource_permission import UserResourcePermi from tools.models import Tool from users.models import User -tool_executor = ToolExecutor(CONFIG.get('SANDBOX')) +tool_executor = ToolExecutor() class KnowledgeWorkflowModelSerializer(serializers.ModelSerializer): diff --git a/apps/tools/serializers/tool.py b/apps/tools/serializers/tool.py index b7be45512..b76153492 100644 --- a/apps/tools/serializers/tool.py +++ b/apps/tools/serializers/tool.py @@ -38,7 +38,7 @@ from system_manage.serializers.user_resource_permission import UserResourcePermi from tools.models import Tool, ToolScope, ToolFolder, ToolType from users.serializers.user import is_workspace_manage -tool_executor = ToolExecutor(CONFIG.get('SANDBOX')) +tool_executor = ToolExecutor() class ToolInstance: