JARVIS/backend/app/agents/tools/sandbox_executor.py

"""
Sandbox Executor - 沙盒执行器
在高风险任务在隔离的临时目录中执行
"""

import asyncio
import os
import shutil
import tempfile
from dataclasses import dataclass, field
from pathlib import Path
from typing import AsyncGenerator

from app.agents.tools.ai_adapter import AICLIAdapter, CodeExecutionResult


@dataclass
class SandboxEnvironment:
    """沙盒环境"""

    workspace_path: Path
    session_id: str

    @staticmethod
    async def create(prefix: str = "jarvis_code_") -> "SandboxEnvironment":
        """创建新的沙盒环境"""
        temp_dir = tempfile.mkdtemp(prefix=prefix)
        session_id = Path(temp_dir).name
        return SandboxEnvironment(
            workspace_path=Path(temp_dir),
            session_id=session_id,
        )

    async def cleanup(self) -> None:
        """清理沙盒环境"""
        if self.workspace_path.exists():
            shutil.rmtree(self.workspace_path)

    def list_created_files(self) -> list[str]:
        """列出沙盒中创建的文件"""
        if not self.workspace_path.exists():
            return []
        files = []
        for root, dirs, filenames in os.walk(self.workspace_path):
            for filename in filenames:
                full_path = os.path.join(root, filename)
                rel_path = os.path.relpath(full_path, self.workspace_path)
                files.append(rel_path)
        return files


@dataclass
class ExecutionResult:
    """执行结果"""

    success: bool
    exit_code: int
    stdout: str
    stderr: str
    files_created: list[str] = field(default_factory=list)


class SandboxExecutor:
    """沙盒执行器"""

    def __init__(self, adapter: AICLIAdapter, timeout: int = 300):
        self.adapter = adapter
        self.timeout = timeout
        self._sessions: dict[str, SandboxEnvironment] = {}

    async def execute(
        self,
        prompt: str,
        session_id: str | None = None,
    ) -> AsyncGenerator[str, None]:
        """
        执行代码任务，yield 实时输出

        Args:
            prompt: 任务描述
            session_id: 会话 ID（可选，用于复用沙盒）

        Yields:
            str: 实时输出
        """
        # 1. 创建或复用沙盒环境
        if session_id and session_id in self._sessions:
            env = self._sessions[session_id]
        else:
            env = await SandboxEnvironment.create()
            self._sessions[env.session_id] = env
            session_id = env.session_id

        # 2. 构建命令
        cmd = self.adapter.build_command(prompt, env.workspace_path)

        # 3. 异步执行，实时 yield 输出
        process = await asyncio.create_subprocess_exec(
            *cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd=str(env.workspace_path),
            env={**os.environ, "TERM": "xterm-256color"},
        )

        # 4. 实时读取输出
        stdout_lines = []
        stderr_lines = []

        while True:
            try:
                line_bytes = await asyncio.wait_for(
                    process.stdout.readline(),
                    timeout=self.timeout,
                )
                if not line_bytes:
                    break
                line = line_bytes.decode("utf-8", errors="replace")
                stdout_lines.append(line)
                yield line
            except asyncio.TimeoutError:
                process.kill()
                yield f"\n[ERROR] Execution timed out after {self.timeout}s\n"
                break

        # 5. 读取 stderr
        stderr_bytes = await process.communicate()
        if stderr_bytes[1]:
            stderr = stderr_bytes[1].decode("utf-8", errors="replace")
            stderr_lines.append(stderr)
            yield f"\n[STDERR]\n{stderr}\n"

        # 6. 返回结果（通过 yield 完成标记）
        result = ExecutionResult(
            success=process.returncode == 0,
            exit_code=process.returncode or 0,
            stdout="".join(stdout_lines),
            stderr="".join(stderr_lines),
            files_created=env.list_created_files(),
        )
        yield f"\n[EXIT_CODE] {result.exit_code}\n"
        yield f"\n[COMPLETE] success={result.success}\n"

    async def get_result(
        self,
        prompt: str,
        session_id: str | None = None,
    ) -> ExecutionResult:
        """同步执行并返回完整结果"""
        output_parts = []
        async for line in self.execute(prompt, session_id):
            output_parts.append(line)
        # 解析最后的结果
        # 实际使用中可能需要更复杂的结果收集
        return ExecutionResult(
            success="[COMPLETE] success=True" in "".join(output_parts),
            exit_code=0,
            stdout="".join(output_parts),
            stderr="",
            files_created=[],
        )

    async def cleanup_session(self, session_id: str) -> bool:
        """清理指定会话"""
        if session_id in self._sessions:
            await self._sessions[session_id].cleanup()
            del self._sessions[session_id]
            return True
        return False

    def get_session(self, session_id: str) -> SandboxEnvironment | None:
        """获取会话环境"""
        return self._sessions.get(session_id)