feat: enhance MCP tool execution with unique ID generation and improved file handling

Author: liuruibin

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

@@ -27,7 +27,9 @@
 from application.flow.step_node.ai_chat_step_node.i_chat_node import IChatNode
 from application.flow.tools import Reasoning
 from common.utils.logger import maxkb_logger
+from common.utils.rsa_util import rsa_long_decrypt
 from common.utils.tool_code import ToolExecutor
+from maxkb.const import CONFIG
 from models_provider.models import Model
 from models_provider.tools import get_model_credential, get_model_instance_by_model_workspace_id
 from tools.models import Tool
@@ -283,25 +285,14 @@ def _handle_mcp_request(self, mcp_enable, tool_enable, mcp_source, mcp_servers,
                 self.context['execute_ids'] = []
                 for tool_id in tool_ids:
                     tool = QuerySet(Tool).filter(id=tool_id).first()
-                    executor = ToolExecutor(sandbox=True)
-                    code = executor.generate_mcp_server_code(tool.code)
-                    _id = uuid.uuid7()
+                    executor = ToolExecutor(CONFIG.get('SANDBOX'))
+                    if tool.init_params is not None:
+                        params = json.loads(rsa_long_decrypt(tool.init_params))
+                    else:
+                        params = {}
+                    _id, tool_config = executor.get_tool_mcp_config(tool.code, params)
+
                     self.context['execute_ids'].append(_id)
-                    code_path = f'{executor.sandbox_path}/execute/{_id}.py'
-                    with open(code_path, 'w') as f:
-                        f.write(code)
-                        os.system(f"chown {executor.user}:root {code_path}")
-
-                    tool_config = {
-                        'command': 'su',
-                        'args': [
-                            '-s', sys.executable,
-                            '-c', f"exec(open('{code_path}', 'r').read())",
-                            executor.user,
-                        ],
-                        'cwd': executor.sandbox_path,
-                        'transport': 'stdio',
-                    }
                     mcp_servers_config[str(tool.id)] = tool_config
 
         if len(mcp_servers_config) > 0:
@@ -347,7 +338,7 @@ def reset_message_list(message_list: List[BaseMessage], answer_text):
     def get_details(self, index: int, **kwargs):
         # 删除临时生成的MCP代码文件
         if self.context.get('execute_ids'):
-            executor = ToolExecutor(sandbox=True)
+            executor = ToolExecutor(CONFIG.get('SANDBOX'))
             # 清理工具代码文件，延时删除，避免文件被占用
             for tool_id in self.context.get('execute_ids'):
                 code_path = f'{executor.sandbox_path}/execute/{tool_id}.py'

apps/common/utils/tool_code.py

@@ -83,7 +83,7 @@ def exec_code(self, code_str, keywords):
             return result.get('data')
         raise Exception(result.get('msg'))
 
-    def _generate_mcp_server_code(self, _code):
+    def _generate_mcp_server_code(self, _code, params):
         self.validate_banned_keywords(_code)
 
         # 解析代码，提取导入语句和函数定义
@@ -100,7 +100,31 @@ def _generate_mcp_server_code(self, _code):
             if isinstance(node, ast.Import) or isinstance(node, ast.ImportFrom):
                 imports.append(ast.unparse(node))
             elif isinstance(node, ast.FunctionDef):
-                # 为函数添加 @mcp.tool() 装饰器
+                # 修改函数参数以包含 params 中的默认值
+                func_name = node.name
+                if func_name in params:
+                    func_params = params[func_name]
+                    # 为函数参数设置默认值
+                    for i, arg in enumerate(node.args.args):
+                        arg_name = arg.arg
+                        if arg_name in func_params:
+                            # 创建默认值节点
+                            default_value = func_params[arg_name]
+                            if isinstance(default_value, str):
+                                default_node = ast.Constant(value=default_value)
+                            elif isinstance(default_value, (int, float, bool)):
+                                default_node = ast.Constant(value=default_value)
+                            else:
+                                default_node = ast.Constant(value=str(default_value))
+
+                            # 添加到defaults列表
+                            if not hasattr(node.args, 'defaults') or node.args.defaults is None:
+                                node.args.defaults = []
+                            # 确保defaults列表长度正确
+                            while len(node.args.defaults) < len(node.args.args):
+                                node.args.defaults.insert(0, None)
+                            node.args.defaults[i] = default_node
+
                 func_code = ast.unparse(node)
                 functions.append(f"@mcp.tool()\n{func_code}\n")
             else:
@@ -116,9 +140,9 @@ def _generate_mcp_server_code(self, _code):
 
         return "\n".join(code_parts)
 
-    def generate_mcp_server_code(self, code_str):
+    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)
+        code = self._generate_mcp_server_code(code_str, params)
         return f"""
 import os
 import sys
@@ -132,6 +156,34 @@ def generate_mcp_server_code(self, code_str):
 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)
+        if self.sandbox:
+            os.system(f"chown {self.user}:root {code_path}")
+
+            tool_config = {
+                'command': 'su',
+                'args': [
+                    '-s', sys.executable,
+                    '-c', f"exec(open('{code_path}', 'r').read())",
+                    self.user,
+                ],
+                'cwd': self.sandbox_path,
+                'transport': 'stdio',
+            }
+        else:
+            tool_config = {
+                'command': sys.executable,
+                'args': [code_path],
+                'transport': 'stdio',
+            }
+        return _id, 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: