FusionAGI/fusionagi/tools/connectors/code_runner.py

"""Code runner connector: execute code in a sandboxed subprocess."""

import subprocess
import tempfile
from pathlib import Path
from typing import Any

from fusionagi._logger import logger
from fusionagi.tools.connectors.base import BaseConnector

SUPPORTED_LANGUAGES = {
    "python": {"ext": ".py", "cmd": ["python3"]},
    "javascript": {"ext": ".js", "cmd": ["node"]},
    "bash": {"ext": ".sh", "cmd": ["bash"]},
    "ruby": {"ext": ".rb", "cmd": ["ruby"]},
}


class CodeRunnerConnector(BaseConnector):
    """Execute code snippets in sandboxed subprocesses.

    Supports Python, JavaScript (Node), Bash, and Ruby.
    Execution is timeout-bounded (default 30s) and captures stdout/stderr.
    """

    name = "code_runner"

    def __init__(self, timeout: float = 30.0, max_output: int = 10000) -> None:
        self._timeout = timeout
        self._max_output = max_output

    def invoke(self, action: str, params: dict[str, Any]) -> Any:
        if action == "run":
            return self._run(
                params.get("code", ""),
                params.get("language", "python"),
                params.get("timeout"),
            )
        if action == "languages":
            return {"languages": list(SUPPORTED_LANGUAGES.keys())}
        return {"error": f"Unknown action: {action}"}

    def _run(self, code: str, language: str, timeout: float | None = None) -> dict[str, Any]:
        if not code.strip():
            return {"stdout": "", "stderr": "", "exit_code": 0, "error": "Empty code"}

        lang = language.lower()
        if lang not in SUPPORTED_LANGUAGES:
            return {
                "stdout": "",
                "stderr": "",
                "exit_code": 1,
                "error": f"Unsupported language: {lang}. Supported: {list(SUPPORTED_LANGUAGES.keys())}",
            }

        spec = SUPPORTED_LANGUAGES[lang]
        effective_timeout = timeout or self._timeout

        try:
            with tempfile.NamedTemporaryFile(
                mode="w", suffix=spec["ext"], delete=False, dir="/tmp"
            ) as f:
                f.write(code)
                f.flush()
                script_path = f.name

            result = subprocess.run(
                [*spec["cmd"], script_path],
                capture_output=True,
                text=True,
                timeout=effective_timeout,
                cwd="/tmp",
            )

            Path(script_path).unlink(missing_ok=True)

            return {
                "stdout": result.stdout[: self._max_output],
                "stderr": result.stderr[: self._max_output],
                "exit_code": result.returncode,
                "error": None,
            }

        except subprocess.TimeoutExpired:
            logger.warning("CodeRunner timeout", extra={"language": lang, "timeout": effective_timeout})
            return {
                "stdout": "",
                "stderr": f"Execution timed out after {effective_timeout}s",
                "exit_code": -1,
                "error": "timeout",
            }
        except FileNotFoundError:
            return {
                "stdout": "",
                "stderr": f"Runtime not found for {lang}: {spec['cmd'][0]}",
                "exit_code": -1,
                "error": f"Runtime '{spec['cmd'][0]}' not installed",
            }
        except Exception as e:
            logger.warning("CodeRunner failed", extra={"error": str(e)})
            return {"stdout": "", "stderr": str(e), "exit_code": -1, "error": str(e)}

    def schema(self) -> dict[str, Any]:
        return {
            "name": self.name,
            "actions": ["run", "languages"],
            "parameters": {"code": "string", "language": "string", "timeout": "number"},
        }