zhuchaowe/maxkb
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: replace diskcache with pickle for result handling and organize sandbox directories

Browse Source
This commit is contained in:
CaptainB 2025-07-09 14:14:23 +08:00
parent 4da75ac3e0
commit dc04d46518
1 changed files with 13 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
apps/common/utils/tool_code.py
View File

@ -1,12 +1,12 @@
# coding=utf-8 # coding=utf-8
import os import os
import pickle
import subprocess import subprocess
import sys import sys
from textwrap import dedent from textwrap import dedent
import uuid_utils.compat as uuid import uuid_utils.compat as uuid
from diskcache import Cache
from maxkb.const import BASE_DIR, CONFIG from maxkb.const import BASE_DIR, CONFIG
from maxkb.const import PROJECT_DIR from maxkb.const import PROJECT_DIR
@ -31,6 +31,8 @@ class ToolExecutor:
old_mask = os.umask(0o077) old_mask = os.umask(0o077)
try: try:
os.makedirs(self.sandbox_path, 0o700, exist_ok=True) os.makedirs(self.sandbox_path, 0o700, exist_ok=True)
os.makedirs(os.path.join(self.sandbox_path, 'execute'), 0o700, exist_ok=True)
os.makedirs(os.path.join(self.sandbox_path, 'result'), 0o700, exist_ok=True)
finally: finally:
os.umask(old_mask) os.umask(old_mask)
@ -38,7 +40,7 @@ class ToolExecutor:
_id = str(uuid.uuid7()) _id = str(uuid.uuid7())
success = '{"code":200,"msg":"成功","data":exec_result}' success = '{"code":200,"msg":"成功","data":exec_result}'
err = '{"code":500,"msg":str(e),"data":None}' err = '{"code":500,"msg":str(e),"data":None}'
path = r'' + self.sandbox_path + '' result_path = f'{self.sandbox_path}/result/{_id}.result'
python_paths = CONFIG.get_sandbox_python_package_paths().split(',') python_paths = CONFIG.get_sandbox_python_package_paths().split(',')
_exec_code = f""" _exec_code = f"""
try: try:
@ -59,13 +61,12 @@ try:
for local in locals_v: for local in locals_v:
globals_v[local] = locals_v[local] globals_v[local] = locals_v[local]
exec_result=f(**keywords) exec_result=f(**keywords)
from diskcache import Cache import pickle
cache = Cache({path!a}) with open({result_path!a}, 'wb') as file:
cache.set({_id!a},{success}) file.write(pickle.dumps({success}))
except Exception as e: except Exception as e:
from diskcache import Cache with open({result_path!a}, 'wb') as file:
cache = Cache({path!a}) file.write(pickle.dumps({err}))
cache.set({_id!a},{err})
""" """
if self.sandbox: if self.sandbox:
subprocess_result = self._exec_sandbox(_exec_code, _id) subprocess_result = self._exec_sandbox(_exec_code, _id)
@ -73,15 +74,15 @@ except Exception as e:
subprocess_result = self._exec(_exec_code) subprocess_result = self._exec(_exec_code)
if subprocess_result.returncode == 1: if subprocess_result.returncode == 1:
raise Exception(subprocess_result.stderr) raise Exception(subprocess_result.stderr)
cache = Cache(self.sandbox_path) with open(result_path, 'rb') as file:
result = cache.get(_id) result = pickle.loads(file.read())
cache.delete(_id) os.remove(result_path)
if result.get('code') == 200: if result.get('code') == 200:
return result.get('data') return result.get('data')
raise Exception(result.get('msg')) raise Exception(result.get('msg'))
def _exec_sandbox(self, _code, _id): def _exec_sandbox(self, _code, _id):
exec_python_file = f'{self.sandbox_path}/{_id}.py' exec_python_file = f'{self.sandbox_path}/execute/{_id}.py'
with open(exec_python_file, 'w') as file: with open(exec_python_file, 'w') as file:
file.write(_code) file.write(_code)
os.system(f"chown {self.user}:root {exec_python_file}") os.system(f"chown {self.user}:root {exec_python_file}")