From ed4b78dc37bd31a59d0831f2054527a1845e2e9d Mon Sep 17 00:00:00 2001
From: 836304831 <836304831@qq.com>
Date: Wed, 12 Mar 2025 23:33:37 +0800
Subject: [PATCH] update python_execute safe

---
 app/tool/python_execute.py | 73 +++++++++++++++++++++-----------------
 1 file changed, 41 insertions(+), 32 deletions(-)

diff --git a/app/tool/python_execute.py b/app/tool/python_execute.py
index 88e1aab..e9c8140 100644
--- a/app/tool/python_execute.py
+++ b/app/tool/python_execute.py
@@ -1,4 +1,6 @@
-import threading
+import sys
+from io import StringIO
+import multiprocessing
 from typing import Dict
 
 from app.tool.base import BaseTool
@@ -20,6 +22,20 @@ class PythonExecute(BaseTool):
         "required": ["code"],
     }
 
+    def _run_code(self, code: str, result_dict: dict, safe_globals: dict) -> None:
+        original_stdout = sys.stdout
+        try:
+            output_buffer = StringIO()
+            sys.stdout = output_buffer
+            exec(code, safe_globals, safe_globals)
+            result_dict["observation"] = output_buffer.getvalue()
+            result_dict["success"] = True
+        except Exception as e:
+            result_dict["observation"] = str(e)
+            result_dict["success"] = False
+        finally:
+            sys.stdout = original_stdout
+
     async def execute(
         self,
         code: str,
@@ -35,36 +51,29 @@ class PythonExecute(BaseTool):
         Returns:
             Dict: Contains 'output' with execution output or error message and 'success' status.
         """
-        result = {"observation": ""}
 
-        def run_code():
-            try:
-                safe_globals = {"__builtins__": dict(__builtins__)}
+        with multiprocessing.Manager() as manager:
+            result = manager.dict({
+                "observation": "",
+                "success": False
+            })
+            if isinstance(__builtins__, dict):
+                safe_globals = {"__builtins__": __builtins__}
+            else:
+                safe_globals = {"__builtins__": __builtins__.__dict__.copy()}
+            proc = multiprocessing.Process(
+                target=self._run_code,
+                args=(code, result, safe_globals)
+            )
+            proc.start()
+            proc.join(timeout)
 
-                import sys
-                from io import StringIO
-
-                output_buffer = StringIO()
-                sys.stdout = output_buffer
-
-                exec(code, safe_globals, {})
-
-                sys.stdout = sys.__stdout__
-
-                result["observation"] = output_buffer.getvalue()
-
-            except Exception as e:
-                result["observation"] = str(e)
-                result["success"] = False
-
-        thread = threading.Thread(target=run_code)
-        thread.start()
-        thread.join(timeout)
-
-        if thread.is_alive():
-            return {
-                "observation": f"Execution timeout after {timeout} seconds",
-                "success": False,
-            }
-
-        return result
+            # timeout process
+            if proc.is_alive():
+                proc.terminate()
+                proc.join(1)
+                return {
+                    "observation": f"Execution timeout after {timeout} seconds",
+                    "success": False,
+                }
+            return dict(result)