agent-Specialization/modules/toolbox_container.py

# modules/toolbox_container.py - 专用工具容器管理器

import asyncio
import time
import uuid
import shlex
from pathlib import Path
from typing import Optional, Dict

from modules.persistent_terminal import PersistentTerminal
from config import (
    TERMINAL_SANDBOX_MODE,
    TERMINAL_SANDBOX_IMAGE,
    TERMINAL_SANDBOX_MOUNT_PATH,
    TERMINAL_SANDBOX_SHELL,
    TERMINAL_SANDBOX_NETWORK,
    TERMINAL_SANDBOX_CPUS,
    TERMINAL_SANDBOX_MEMORY,
    TERMINAL_SANDBOX_BINDS,
    TERMINAL_SANDBOX_BIN,
    TERMINAL_SANDBOX_NAME_PREFIX,
    TERMINAL_SANDBOX_ENV,
    TERMINAL_SANDBOX_REQUIRE,
    TOOLBOX_TERMINAL_IDLE_SECONDS,
)


def _build_sandbox_options() -> Dict:
    """构造与终端一致的沙箱配置."""
    name_prefix = TERMINAL_SANDBOX_NAME_PREFIX
    if name_prefix:
        name_prefix = f"{name_prefix}-toolbox"
    else:
        name_prefix = "toolbox-term"
    return {
        "image": TERMINAL_SANDBOX_IMAGE,
        "mount_path": TERMINAL_SANDBOX_MOUNT_PATH,
        "shell": TERMINAL_SANDBOX_SHELL,
        "network": TERMINAL_SANDBOX_NETWORK,
        "cpus": TERMINAL_SANDBOX_CPUS,
        "memory": TERMINAL_SANDBOX_MEMORY,
        "binds": list(TERMINAL_SANDBOX_BINDS),
        "bin": TERMINAL_SANDBOX_BIN,
        "name_prefix": name_prefix,
        "env": dict(TERMINAL_SANDBOX_ENV),
        "require": TERMINAL_SANDBOX_REQUIRE,
    }


class ToolboxContainer:
    """为 run_command/run_python 提供的专用容器."""

    def __init__(
        self,
        project_path: str,
        sandbox_mode: Optional[str] = None,
        sandbox_options: Optional[Dict] = None,
        idle_timeout: int = TOOLBOX_TERMINAL_IDLE_SECONDS,
    ):
        self.project_path = Path(project_path).resolve()
        self.sandbox_mode = (sandbox_mode or TERMINAL_SANDBOX_MODE or "host").lower()
        options = _build_sandbox_options()
        if sandbox_options:
            for key, value in sandbox_options.items():
                if key == "binds" and isinstance(value, list):
                    options[key] = list(value)
                elif key == "env" and isinstance(value, dict):
                    options[key] = dict(value)
                else:
                    options[key] = value
        self.sandbox_options = options
        self.idle_timeout = max(0, int(idle_timeout)) if idle_timeout is not None else 0

        self._terminal: Optional[PersistentTerminal] = None
        self._lock = asyncio.Lock()
        self._session_name = f"toolbox-{uuid.uuid4().hex[:10]}"
        self._last_used = 0.0

    async def _ensure_terminal(self) -> PersistentTerminal:
        """确保容器已启动。"""
        async with self._lock:
            if self._terminal and self._terminal.is_running:
                return self._terminal

            terminal = PersistentTerminal(
                session_name=self._session_name,
                working_dir=str(self.project_path),
                broadcast_callback=None,
                project_path=str(self.project_path),
                sandbox_mode=self.sandbox_mode,
                sandbox_options=self.sandbox_options,
            )

            if not terminal.start():
                raise RuntimeError("工具容器启动失败，请检查 Docker 或本地 shell 环境。")

            self._terminal = terminal
            self._last_used = time.time()
            print(f"[Toolbox] 工具容器已启动 (session={self._session_name}, mode={terminal.execution_mode})")
            return terminal

    def _wrap_command(self, command: str, work_path: Path, execution_mode: Optional[str]) -> str:
        """根据执行模式动态拼接 cd 指令。"""
        if work_path == self.project_path:
            return command

        try:
            relative = work_path.relative_to(self.project_path).as_posix()
        except ValueError:
            relative = ""

        if not relative:
            return command

        if execution_mode == "docker":
            mount_path = self.sandbox_options.get("mount_path", "/workspace").rstrip("/")
            target = f"{mount_path}/{relative}"
        else:
            target = str(work_path)

        return f"cd {shlex.quote(target)} && {command}"

    async def run(self, command: str, work_path: Path, timeout: Optional[float] = None) -> Dict:
        """在容器/主机终端中执行命令。"""
        terminal = await self._ensure_terminal()
        shell_command = self._wrap_command(command, work_path, terminal.execution_mode)
        result = await asyncio.to_thread(
            terminal.send_command,
            shell_command,
            True,
            timeout,
        )
        self._last_used = time.time()
        self._cleanup_if_idle()
        return result

    def _cleanup_if_idle(self):
        if self.idle_timeout <= 0 or not self._terminal:
            return
        if time.time() - self._last_used >= self.idle_timeout:
            self._terminal.close()
            print(f"[Toolbox] 工具容器空闲超时，已释放 (session={self._session_name})")
            self._terminal = None

    def shutdown(self):
        """立即关闭容器。"""
        if self._terminal:
            self._terminal.close()
            print(f"[Toolbox] 工具容器已关闭 (session={self._session_name})")
            self._terminal = None

    @property
    def execution_mode(self) -> Optional[str]:
        if self._terminal:
            return self._terminal.execution_mode
        return None