core/engine/src/runtime/sandbox.py - main - alexefan136/flowstack

core/engine/src/runtime/sandbox.py
172 строки6 KB
Alexander Efanov
upd fix
31 июл 2026, 19:17
31 июл 2026, 19:17d146d86
"""Sandbox — изолированное выполнение кода (subprocess + timeout)."""

from __future__ import annotations

import asyncio
import json
import sys
from dataclasses import dataclass
from typing import Any

import structlog

from src.primitives import elapsed_ms, monotonic

logger = structlog.get_logger()

# Шаблон обёртки: выполняет код и сериализует переменную `result` (если есть)
_WRAPPER_TEMPLATE = """
import json

_inputs = {inputs_json}
result = None

{code}

try:
    _out = result
except NameError:
    _out = None

print("\\n__SANDBOX_RESULT__" + json.dumps({{"result": _out}}, default=str))
"""

_RESULT_MARKER = "__SANDBOX_RESULT__"


@dataclass
class SandboxResult:
    """Результат выполнения кода в песочнице."""

    success: bool
    stdout: str = ""
    stderr: str = ""
    result: Any = None
    exit_code: int = 0
    duration_ms: float = 0.0
    error: str | None = None

    def to_dict(self) -> dict[str, Any]:
        """Преобразовать в словарь."""
        return {
            "success": self.success,
            "stdout": self.stdout,
            "stderr": self.stderr,
            "result": self.result,
            "exit_code": self.exit_code,
            "duration_ms": self.duration_ms,
            "error": self.error,
        }


class Sandbox:
    """
    Песочница для выполнения Python-кода.

    Выполняет код в отдельном subprocess с timeout. Код имеет доступ к
    переменной ``inputs`` (dict) и может установить ``result`` для возврата.

    Note: для production рекомендуется Docker/gVisor-изоляция. Данная
    реализация — subprocess + timeout (базовый уровень).
    """

    def __init__(
        self,
        timeout_seconds: float = 30.0,
        python_executable: str | None = None,
    ) -> None:
        self.timeout = timeout_seconds
        self.python = python_executable or sys.executable
        logger.info("sandbox.initialized", timeout=timeout_seconds)

    def _build_wrapper(self, code: str, inputs: dict[str, Any] | None) -> str:
        """Собрать код-обёртку с сериализацией результата."""
        inputs_json = json.dumps(inputs or {}, ensure_ascii=False, default=str)
        return _WRAPPER_TEMPLATE.format(inputs_json=inputs_json, code=code)

    def _parse_output(self, stdout: str) -> tuple[str, Any]:
        """Извлечь чистый stdout и result из вывода."""
        if _RESULT_MARKER in stdout:
            before, _, after = stdout.partition(_RESULT_MARKER)
            try:
                parsed = json.loads(after.strip())
                return before.strip(), parsed.get("result")
            except json.JSONDecodeError:
                return stdout.strip(), None
        return stdout.strip(), None

    async def execute(
        self,
        code: str,
        inputs: dict[str, Any] | None = None,
        timeout_seconds: float | None = None,
    ) -> SandboxResult:
        """
        Выполнить Python-код в изолированном subprocess.

        Args:
            code: код для выполнения (может использовать ``inputs`` и ``result``)
            inputs: входные данные (доступны как переменная ``inputs``)
            timeout_seconds: таймаут (по умолчанию из настроек)

        Returns:
            SandboxResult с stdout/stderr/result.
        """
        start = monotonic()
        timeout = timeout_seconds or self.timeout
        wrapper = self._build_wrapper(code, inputs)

        try:
            proc = await asyncio.create_subprocess_exec(
                self.python,
                "-I",  # isolated mode (без user site-packages)
                "-c",
                wrapper,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
            )

            stdout_bytes, stderr_bytes = await asyncio.wait_for(proc.communicate(), timeout=timeout)

            stdout = stdout_bytes.decode("utf-8", errors="replace")
            stderr = stderr_bytes.decode("utf-8", errors="replace")
            clean_stdout, result = self._parse_output(stdout)
            duration = elapsed_ms(start)

            success = proc.returncode == 0
            return SandboxResult(
                success=success,
                stdout=clean_stdout,
                stderr=stderr,
                result=result,
                exit_code=proc.returncode or 0,
                duration_ms=duration,
                error=stderr if not success and stderr else None,
            )

        except TimeoutError:
            logger.warning("sandbox.timeout", timeout=timeout)
            try:
                proc.kill()  # type: ignore[possibly-undefined]
            except Exception:
                pass
            return SandboxResult(
                success=False,
                error=f"Execution timeout ({timeout}s)",
                duration_ms=elapsed_ms(start),
            )

        except Exception as e:
            logger.error("sandbox.execution_failed", error=str(e))
            return SandboxResult(
                success=False,
                error=str(e),
                duration_ms=elapsed_ms(start),
            )

    async def execute_expression(self, expression: str) -> SandboxResult:
        """Вычислить выражение и вернуть результат."""
        return await self.execute(f"result = ({expression})")


__all__ = ["Sandbox", "SandboxResult"]