zcbot/tools/run_python.py

"""run_python: 在 subprocess 里执行 Python 代码 (Hybrid 范式的关键)。

JSON tool call 处理离散操作 (read/edit/shell),run_python 处理连续逻辑
(数据计算、批处理、生成 .pptx/.docx)。让模型自己挑工具。

阶段 1 (本地): subprocess + 工作目录限制 + 敏感环境变量过滤
阶段 2 (后期): Docker / E2B 替换执行后端
"""
from __future__ import annotations

import os
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Optional

from core import procs

from .base import Tool, compact_tool_output

_SENSITIVE_PATTERNS = ("API_KEY", "TOKEN", "SECRET", "PASSWORD", "PRIVATE_KEY")
# 刻意放行的例外:research skill 在 sandbox 里直连 paper_server,这把只读文献库 key
# 的消费方就是沙盒代码本身;低价值、可随时在 paper_server admin 撤销。
_ENV_ALLOWLIST = frozenset({"PAPER_SERVER_API_KEY"})


class RunPythonTool(Tool):
    name = "run_python"
    description = (
        "Execute Python code in a subprocess. Returns stdout/stderr/exit_code.\n"
        "Use script_path for non-trivial code: write the .py under scripts/ (e.g. scripts/analyze.py) first, "
        "then execute it so the full source stays in files instead of conversation history.\n"
        "Use inline code only for short throwaway snippets (a quick calc / one-line probe) — these run from a temp file and leave no trace.\n"
        "Good for: data analysis, batch file ops, document generation (.pptx/.docx), "
        "matplotlib charts, or any task where Python is more natural than chaining tools.\n"
        "Working directory is the agent's base dir (task_dir); relative paths resolve against it. "
        "Keep process scripts in scripts/; write deliverables to task_dir root or the SKILL-specified path.\n"
        "Available libs (install with shell pip if missing): "
        "pandas, numpy, matplotlib, python-pptx, python-docx, requests, pypdf."
    )
    parameters = {
        "type": "object",
        "properties": {
            "code": {
                "type": "string",
                "description": "Short Python source. For longer code, write a .py file and pass script_path.",
            },
            "script_path": {
                "type": "string",
                "description": "Path to an existing .py file to execute (relative to task_dir). Prefer this for non-trivial code; keep such scripts under scripts/.",
            },
            "timeout": {"type": "integer", "default": 120, "description": "Seconds before kill. background=true 时含义变为最长运行时长,默认 7200s。"},
            "background": {
                "type": "boolean",
                "description": (
                    "true = 后台执行:立即返回 proc_id,进程 detach 独立运行(不受本轮对话、"
                    "服务重启影响),用 check_process 查进度/结果。适用:预计运行超过约 1 分钟的"
                    "计算/训练/批处理。快速脚本保持默认 false(同步等结果)。前台超时被杀的"
                    "长任务应改用 background=true 重新发起。"
                ),
            },
        },
        "required": [],
    }

    def __init__(
        self,
        base_dir: Optional[Path] = None,
        user_root: Optional[Path] = None,
        task_id: Optional[str] = None,
    ) -> None:
        super().__init__(base_dir, user_root=user_root)
        self.task_id = str(task_id) if task_id else "default"

    def _filtered_env(self) -> dict:
        env = os.environ.copy()
        for k in list(env):
            u = k.upper()
            if u not in _ENV_ALLOWLIST and any(p in u for p in _SENSITIVE_PATTERNS):
                del env[k]
        env["PYTHONIOENCODING"] = "utf-8"
        env["PYTHONPATH"] = str(self.base_dir) + os.pathsep + env.get("PYTHONPATH", "")
        return env

    def _execute_background(
        self, code: str | None, script_path: str | None, timeout: int | None
    ) -> str:
        anchor = self.user_root or self.base_dir
        if procs.count_running(anchor) >= procs.MAX_RUNNING_PER_USER:
            return (
                f"[Error] 已有 {procs.MAX_RUNNING_PER_USER} 个后台进程在跑(上限)。"
                f"用 check_process 查看,等待完成或 kill 掉不需要的再启动。"
            )
        timeout_s = procs.clamp_timeout(timeout if timeout and timeout > 120 else None)
        if script_path:
            script = self._resolve(script_path)
            if not script.is_file():
                return f"[Error] script_path not found: {self._display(script)}"
            proc_id, _ = procs.launch_host(
                anchor, self.task_id,
                kind="python",
                command_display=f"python {self._display(script)}",
                argv=[sys.executable, str(script)],
                cwd=self.base_dir, timeout_s=timeout_s, env=self._filtered_env(),
            )
            shown = self._display(script)
        elif isinstance(code, str) and code.strip():
            proc_id, _ = procs.launch_host(
                anchor, self.task_id,
                kind="python",
                command_display=f"python <inline {len(code)} chars>",
                cwd=self.base_dir, timeout_s=timeout_s, env=self._filtered_env(),
                inline_code=code,
            )
            shown = "<inline code>"
        else:
            return "[Error] run_python requires code or script_path"
        return (
            f"[Background] 已启动后台进程 proc_id={proc_id}({shown}),最长运行 {timeout_s}s。\n"
            f"用 check_process(proc_id=\"{proc_id}\") 查进度和日志。进程独立于本轮对话运行,"
            f"服务重启也不中断。现在可以继续其他工作;若无事可做,结束回合并告知用户稍后询问进度。"
        )

    def execute(
        self,
        code: str | None = None,
        script_path: str | None = None,
        timeout: int = 120,
        background: bool = False,
    ) -> str:
        if background:
            return self._execute_background(code, script_path, timeout)
        cleanup_script = False
        if script_path:
            script = self._resolve(script_path)
            if not script.is_file():
                return f"[Error] script_path not found: {self._display(script)}"
        elif isinstance(code, str):
            # 写到临时文件,避免 -c 转义问题
            with tempfile.NamedTemporaryFile(
                suffix=".py", mode="w", delete=False, encoding="utf-8"
            ) as f:
                f.write(code)
                script = Path(f.name)
            cleanup_script = True
        else:
            return "[Error] run_python requires code or script_path"

        try:
            env = self._filtered_env()
            result = subprocess.run(
                [sys.executable, str(script)],
                cwd=str(self.base_dir),
                capture_output=True,
                timeout=timeout,
                text=True,
                encoding="utf-8",
                errors="replace",
                env=env,
            )
        except subprocess.TimeoutExpired:
            return (
                f"[Error] python script timed out after {timeout}s. "
                f"若任务本身需要长时间运行,用 background=true 重新发起(后台执行,check_process 查进度)。"
            )
        finally:
            if cleanup_script:
                try:
                    script.unlink()
                except OSError:
                    pass

        parts = []
        if result.stdout:
            parts.append(f"[stdout]\n{result.stdout.rstrip()}")
        if result.stderr:
            parts.append(f"[stderr]\n{result.stderr.rstrip()}")
        parts.append(f"[exit {result.returncode}]")
        return compact_tool_output("\n".join(parts))