base_agent/tools/code_executor.py

"""
tools/code_executor.py
代码执行工具 —— 在沙箱中执行 Python 代码片段
配置通过 settings.tools['code_executor'] 读取
"""

import io
import sys
import textwrap
import time
import traceback
from contextlib import redirect_stderr, redirect_stdout

from config.settings import settings
from tools.base_tool import BaseTool
from utils.logger import get_logger

logger = get_logger("TOOL.CodeExecutor")


def _cfg(key: str, fallback=None):
    return settings.tools['code_executor'].get(key, fallback)


class Tool(BaseTool):
    name        = "code_executor"
    description = (
        "在安全沙箱中执行 Python 代码片段，返回标准输出和执行结果。"
        "适用于数据处理、计算、格式转换等任务。"
        "注意：沙箱模式下禁止文件系统写入、网络访问和系统调用。"
    )
    parameters  = {
        "type": "object",
        "properties": {
            "code": {
                "type":        "string",
                "description": "要执行的 Python 代码字符串",
            },
            "timeout": {
                "type":        "integer",
                "description": "执行超时秒数（默认来自 config.yaml code_executor.timeout）",
            },
        },
        "required": ["code"],
    }

    # 沙箱模式下禁止的模块和函数
    _FORBIDDEN_SANDBOX = [
        "import os", "import sys", "import subprocess",
        "import socket", "import requests", "import httpx",
        "import shutil", "open(", "__import__",
        "exec(", "eval(", "compile(",
    ]

    def execute(self, code: str = "", timeout: int | None = None, **_) -> str:
        if not code or not code.strip():
            return "❌ 参数错误: code 不能为空"

        sandbox = _cfg('sandbox', True)
        t       = timeout or _cfg('timeout', 5)
        code    = textwrap.dedent(code)

        logger.info(
            f"🐍 执行代码  sandbox={sandbox}  timeout={t}s  "
            f"[config timeout={_cfg('timeout')}s sandbox={_cfg('sandbox')}]\n"
            f"   代码预览: {code[:100]}"
        )

        # 沙箱安全检查
        if sandbox:
            err = self._sandbox_check(code)
            if err:
                return err

        # 使用线程超时执行
        return self._run_with_timeout(code, t)

    def _run_with_timeout(self, code: str, timeout: int) -> str:
        """在独立线程中执行代码，超时则终止"""
        import threading

        result_box: list[str] = []
        error_box:  list[str] = []

        def _run():
            stdout_buf = io.StringIO()
            stderr_buf = io.StringIO()
            local_ns: dict = {}
            start = time.time()
            try:
                with redirect_stdout(stdout_buf), redirect_stderr(stderr_buf):
                    exec(code, {"__builtins__": __builtins__}, local_ns)  # noqa: S102
                elapsed = time.time() - start
                stdout  = stdout_buf.getvalue()
                stderr  = stderr_buf.getvalue()
                # 尝试获取最后一个表达式的值
                last_val = ""
                lines    = [l.strip() for l in code.strip().splitlines() if l.strip()]
                if lines:
                    last_line = lines[-1]
                    if not last_line.startswith(("#", "print", "import", "from",
                                                 "def ", "class ", "if ", "for ",
                                                 "while ", "try:", "with ")):
                        try:
                            val = eval(last_line, {"__builtins__": __builtins__}, local_ns)  # noqa: S307
                            if val is not None:
                                last_val = f"\n返回值: {repr(val)}"
                        except Exception:
                            pass

                output = stdout + (f"\n[stderr]\n{stderr}" if stderr else "") + last_val
                result_box.append(
                    f"✅ 执行成功  耗时={elapsed:.3f}s\n"
                    f"{'─' * 40}\n"
                    f"{output.strip() or '(无输出)'}"
                )
            except Exception:
                elapsed = time.time() - start
                tb      = traceback.format_exc()
                error_box.append(
                    f"❌ 执行错误  耗时={elapsed:.3f}s\n"
                    f"{'─' * 40}\n{tb}"
                )

        thread = threading.Thread(target=_run, daemon=True)
        thread.start()
        thread.join(timeout=timeout)

        if thread.is_alive():
            return (
                f"⏰ 执行超时（>{timeout}s）\n"
                f"   请增大 config.yaml → tools.code_executor.timeout\n"
                f"   或优化代码逻辑"
            )

        if error_box:
            return error_box[0]
        return result_box[0] if result_box else "❌ 执行失败（未知错误）"

    def _sandbox_check(self, code: str) -> str | None:
        """沙箱模式下的静态安全检查"""
        for forbidden in self._FORBIDDEN_SANDBOX:
            if forbidden in code:
                return (
                    f"❌ 沙箱限制: 代码包含禁止操作 '{forbidden}'\n"
                    f"   如需完整权限请在 config.yaml → "
                    f"tools.code_executor.sandbox 设置为 false"
                )
        return None