tools/code_executor.py

import json
import subprocess
import tempfile
import os
from google.genai import types


def execute_python_code(code: str, timeout: int = 30) -> dict:
    """Execute Python code in a sandboxed subprocess and return output."""
    result = {"stdout": "", "stderr": "", "exit_code": -1, "success": False}

    # Security: block dangerous operations
    blocked = ["os.system", "subprocess", "eval(", "exec(", "import os", "__import__"]
    for b in blocked:
        if b in code:
            result["stderr"] = f"Blocked: '{b}' is not allowed for security reasons."
            return result

    try:
        with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
            f.write(code)
            tmp_path = f.name

        proc = subprocess.run(
            ["python3", tmp_path],
            capture_output=True,
            text=True,
            timeout=timeout,
            cwd=tempfile.gettempdir(),
            env={**os.environ, "HOME": tempfile.gettempdir()},
        )
        result["stdout"] = proc.stdout
        result["stderr"] = proc.stderr
        result["exit_code"] = proc.returncode
        result["success"] = proc.returncode == 0

    except subprocess.TimeoutExpired:
        result["stderr"] = f"Code execution timed out after {timeout}s"
    except Exception as e:
        result["stderr"] = str(e)
    finally:
        try:
            os.unlink(tmp_path)
        except Exception:
            pass

    return result


def execute_javascript_code(code: str, timeout: int = 10) -> dict:
    """Execute JavaScript code using Node.js."""
    result = {"stdout": "", "stderr": "", "exit_code": -1, "success": False}

    try:
        with tempfile.NamedTemporaryFile(mode="w", suffix=".js", delete=False) as f:
            f.write(code)
            tmp_path = f.name

        proc = subprocess.run(
            ["node", tmp_path],
            capture_output=True,
            text=True,
            timeout=timeout,
        )
        result["stdout"] = proc.stdout
        result["stderr"] = proc.stderr
        result["exit_code"] = proc.returncode
        result["success"] = proc.returncode == 0

    except FileNotFoundError:
        result["stderr"] = "Node.js is not available on this server."
    except subprocess.TimeoutExpired:
        result["stderr"] = f"JS execution timed out after {timeout}s"
    except Exception as e:
        result["stderr"] = str(e)
    finally:
        try:
            os.unlink(tmp_path)
        except Exception:
            pass

    return result


# Tool declarations for google-genai function calling
code_tool_declarations = [
    types.FunctionDeclaration(
        name="execute_code",
        description="Execute Python or JavaScript code and return output. Use to calculate, process data, or test algorithms. Never use for system access or network requests.",
        parameters=types.Schema(
            type="OBJECT",
            properties={
                "language": types.Schema(
                    type="STRING",
                    enum=["python", "javascript"],
                    description="Programming language",
                ),
                "code": types.Schema(type="STRING", description="The code to execute"),
            },
            required=["language", "code"],
        ),
    )
]


def handle_code_execution(args: dict) -> str:
    """Handle code execution tool call."""
    language = args.get("language", "python")
    code = args.get("code", "")

    if language == "python":
        result = execute_python_code(code)
    elif language == "javascript":
        result = execute_javascript_code(code)
    else:
        return f"Unsupported language: {language}"

    output = ""
    if result["success"]:
        if result["stdout"]:
            output += result["stdout"]
        if result["stderr"]:
            output += f"\n[stderr]: {result['stderr']}"
        if not output:
            output = "Code executed successfully (no output)."
    else:
        output = f"Error (exit {result['exit_code']}): {result['stderr']}"
        if result["stdout"]:
            output = f"{result['stdout']}\n{output}"

    return output.strip()