executors/python_executor.py

"""
executors/python_executor.py
-----------------------------
Python code executor for AutoDevAgent.

Runs generated Python code in a sandboxed subprocess with a strict
timeout. Captures stdout, stderr, and execution time. Returns a typed
ExecutionResult — never raises on code failures, only on executor
failures (e.g. subprocess could not be spawned).

Design decisions:
  - Uses subprocess.run() with timeout, not exec() or eval().
    This gives genuine process isolation — a crash in generated code
    cannot affect the parent process.
  - Code is written to a temp file rather than passed via -c flag.
    -c has shell escaping issues with multiline code and special chars.
  - stderr is always captured and cleaned before surfacing to the LLM.
    Raw tracebacks contain temp file paths that confuse the debug agent.
  - Timeout kills the process and returns ErrorType.TIMEOUT cleanly.

Known limitation:
  - subprocess isolation is not Docker-grade. A malicious script could
    access the filesystem or make network calls. Acceptable for a
    free-tier HuggingFace Spaces deployment; document in README.

Usage:
    from executors.python_executor import PythonExecutor
    from pipeline.state import PipelineState, Language

    executor = PythonExecutor()
    state = PipelineState(task="...", language=Language.PYTHON)
    state.generated_code = "print('hello')"
    result = executor.run(state)
    print(result["execution_result"].stdout)   # "hello"
    print(result["execution_result"].success)  # True
"""

import logging
import os
import subprocess
import sys
import tempfile
import time
from typing import Any

from config import settings
from pipeline.state import (
    ExecutionResult,
    PipelineState,
    PipelineStatus,
)

logger = logging.getLogger(__name__)


class PythonExecutor:
    """
    Executes Python code in a sandboxed subprocess.

    Writes generated code to a temporary file, runs it with the
    current Python interpreter under a configurable timeout, captures
    all output, and returns a typed ExecutionResult.

    Attributes:
        timeout: Maximum seconds allowed per execution run.
    """

    def __init__(self) -> None:
        """Initialise with timeout from central config."""
        self.timeout: int = settings.python_execution_timeout

    def run(self, state: PipelineState) -> dict[str, Any]:
        """
        Execute the generated code in state.generated_code.

        Writes code to a temp file, spawns a subprocess, enforces
        the timeout, captures stdout/stderr, and cleans error messages
        before returning.

        Args:
            state: Current PipelineState. Reads: generated_code.

        Returns:
            Partial state dict with keys:
              - "execution_result": ExecutionResult
              - "status": PipelineStatus.EXECUTING
        """
        code = state.generated_code

        if not code or not code.strip():
            logger.warning("PythonExecutor received empty code.")
            return {
                "execution_result": ExecutionResult(
                    success=False,
                    error_msg="No code was provided to execute.",
                ),
                "status": PipelineStatus.EXECUTING,
            }

        logger.info("PythonExecutor running code (%d lines)", len(code.splitlines()))

        # ── Dependency resolution ─────────────────────────────────── #
        # Detect imports and attempt to install missing packages before
        # execution so the code doesn't fail on ImportError needlessly.
        try:
            from utils.dependency_manager import DependencyManager
            dep_result = DependencyManager().resolve(code)
            if dep_result.failed:
                logger.warning(
                    "PythonExecutor: could not install: %s", dep_result.failed
                )
            elif dep_result.installed:
                logger.info(
                    "PythonExecutor: auto-installed: %s", dep_result.installed
                )
        except Exception as dep_err:
            # Dependency resolution failure is non-fatal — attempt execution anyway
            logger.warning("DependencyManager error (non-fatal): %s", dep_err)


        # Write code to a named temp file so subprocess can run it cleanly
        tmp_path = None
        try:
            tmp_file = tempfile.NamedTemporaryFile(
                mode="w",
                suffix=".py",
                delete=False,
                encoding="utf-8",
            )
            tmp_file.write(code)
            tmp_file.close()
            tmp_path = tmp_file.name

            result = self._execute(tmp_path)

        finally:
            # Always clean up the temp file
            if tmp_path and os.path.exists(tmp_path):
                os.unlink(tmp_path)

        return {
            "execution_result": result,
            "status": PipelineStatus.EXECUTING,
        }

    # ---------------------------------------------------------------- #
    #  Internal execution                                               #
    # ---------------------------------------------------------------- #

    def _execute(self, file_path: str) -> ExecutionResult:
        """
        Spawn a subprocess to run the code file and capture output.

        Args:
            file_path: Absolute path to the temporary Python file.

        Returns:
            ExecutionResult with success flag, stdout, stderr,
            cleaned error message, and execution time.
        """
        start = time.perf_counter()

        try:
            proc = subprocess.run(
                [sys.executable, file_path],
                capture_output=True,
                text=True,
                timeout=self.timeout,
            )
            elapsed = round(time.perf_counter() - start, 3)

            if proc.returncode == 0:
                logger.info("PythonExecutor success in %.3fs", elapsed)
                return ExecutionResult(
                    success=True,
                    stdout=proc.stdout.strip(),
                    stderr=proc.stderr.strip(),
                    exec_time=elapsed,
                )
            else:
                error_msg = _clean_error(proc.stderr, file_path)
                logger.info(
                    "PythonExecutor failed in %.3fs: %s",
                    elapsed,
                    error_msg[:120],
                )
                return ExecutionResult(
                    success=False,
                    stdout=proc.stdout.strip(),
                    stderr=proc.stderr.strip(),
                    error_msg=error_msg,
                    exec_time=elapsed,
                )

        except subprocess.TimeoutExpired:
            elapsed = round(time.perf_counter() - start, 3)
            msg = (
                f"Execution timed out after {self.timeout} seconds. "
                "The code may contain an infinite loop or expensive computation."
            )
            logger.warning("PythonExecutor timeout after %.3fs", elapsed)
            return ExecutionResult(
                success=False,
                error_msg=msg,
                exec_time=elapsed,
            )

        except Exception as e:
            # Executor itself failed — subprocess could not be spawned
            elapsed = round(time.perf_counter() - start, 3)
            msg = f"Executor internal error: {e}"
            logger.error(msg)
            return ExecutionResult(
                success=False,
                error_msg=msg,
                exec_time=elapsed,
            )


# ------------------------------------------------------------------ #
#  Helpers                                                             #
# ------------------------------------------------------------------ #

def _clean_error(stderr: str, file_path: str) -> str:
    """
    Clean a raw Python traceback for the debug agent.

    Raw tracebacks reference the temp file path (e.g.
    '/tmp/tmpXYZ123.py'), which is meaningless to the LLM. This
    function replaces those paths with 'generated_code.py' and
    strips redundant 'During handling...' noise.

    Also extracts just the final error line if the full traceback is
    very long, since the error type + message is what the debug agent
    needs most.

    Args:
        stderr:    Raw stderr string from the subprocess.
        file_path: Path to the temp file to scrub from the output.

    Returns:
        Cleaned error string suitable for the debug agent prompt.
    """
    if not stderr:
        return "Unknown error — no stderr output captured."

    # Replace temp file path with a clean placeholder
    cleaned = stderr.replace(file_path, "generated_code.py")

    # Remove 'During handling of the above exception...' noise
    noise_marker = "During handling of the above exception"
    if noise_marker in cleaned:
        cleaned = cleaned[: cleaned.index(noise_marker)].strip()

    # If the traceback is very long, keep the last 10 lines
    # (the traceback header + the final error line are most useful)
    lines = cleaned.strip().splitlines()
    if len(lines) > 10:
        cleaned = "\n".join(lines[-10:])

    return cleaned.strip()