From 541f1f26c85d1cdb7c53f7acc8fa2819ecc3869b Mon Sep 17 00:00:00 2001
From: CaptainB <bin@fit2cloud.com>
Date: Tue, 12 Aug 2025 17:22:24 +0800
Subject: [PATCH] feat: refactor MCP code generation and update cleanup logic

---
 .../step_node/ai_chat_step_node/impl/base_chat_node.py |  6 +++---
 apps/common/utils/tool_code.py                         | 10 ++--------
 2 files changed, 5 insertions(+), 11 deletions(-)

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 5505a0277..6aee4ebd3 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
@@ -281,7 +281,7 @@ class BaseChatNode(IChatNode):
                 for tool_id in tool_ids:
                     tool = QuerySet(Tool).filter(id=tool_id).first()
                     executor = ToolExecutor()
-                    code = executor.get_exec_code(tool.code)
+                    code = executor.generate_mcp_server_code(tool.code)
                     code_path = f'{executor.sandbox_path}/execute/{tool_id}.py'
                     with open(code_path, 'w') as f:
                         f.write(code)
@@ -340,8 +340,8 @@ class BaseChatNode(IChatNode):
             # 清理工具代码文件，延时删除，避免文件被占用
             for tool_id in self.context.get('tool_ids'):
                 code_path = f'{executor.sandbox_path}/execute/{tool_id}.py'
-                if os.path.exists(code_path):
-                    os.remove(code_path)
+                # 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 cfe41a563..e127d98c4 100644
--- a/apps/common/utils/tool_code.py
+++ b/apps/common/utils/tool_code.py
@@ -116,14 +116,13 @@ except Exception as e:
 
         return "\n".join(code_parts)
 
-    def get_exec_code(self, code_str):
+    def generate_mcp_server_code(self, code_str):
         python_paths = CONFIG.get_sandbox_python_package_paths().split(',')
         code = self._generate_mcp_server_code(code_str)
         return f"""
 try:
     import os
     import sys
-    import pickle
     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}
@@ -131,12 +130,7 @@ try:
     for key in list(env.keys()):
         if key in os.environ and (key.startswith('MAXKB') or key.startswith('POSTGRES') or key.startswith('PG') or key.startswith('REDIS') or key == 'PATH'):
             del os.environ[key]
-    locals_v={'{}'}
-    globals_v=globals()
-    exec({dedent(code)!a}, globals_v, locals_v)
-    f_name, f = locals_v.popitem()
-    for local in locals_v:
-        globals_v[local] = locals_v[local]
+    exec({dedent(code)!a})
 except Exception as e:
     pass
 """