ide_exec.py

from __future__ import annotations

import queue
import shlex
import subprocess
import threading
import time
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, Optional

from ide_fs import REPO_ROOT, resolve_dir

SAFE_COMMANDS = {
    "python",
    "python3",
    "py",
    "pip",
    "pip3",
    "pytest",
    "ruff",
    "mypy",
    "node",
    "npm",
    "pnpm",
    "yarn",
    "bun",
    "deno",
    "tsc",
    "eslint",
    "go",
    "cargo",
    "dotnet",
    "java",
    "javac",
    "rg",
    "git",
    "cmake",
    "make",
    "ninja",
}

MAX_LOG_LINES = 2000


def _command_base(cmd: str) -> str:
    if not cmd:
        return ""
    try:
        parts = shlex.split(cmd, posix=False)
    except ValueError:
        parts = cmd.split()
    if not parts:
        return ""
    name = Path(parts[0]).name
    return Path(name).stem.lower()


def is_safe_command(cmd: str) -> bool:
    base = _command_base(cmd)
    return base in SAFE_COMMANDS


@dataclass
class ExecSession:
    session_id: str
    cmd: str
    cwd: str
    unsafe: bool
    process: subprocess.Popen
    queue: "queue.Queue[dict]" = field(default_factory=queue.Queue)
    logs: list[dict] = field(default_factory=list)
    status: str = "running"
    return_code: Optional[int] = None
    started_at: float = field(default_factory=time.time)

    def append_log(self, stream: str, text: str) -> None:
        payload = {
            "type": "log",
            "stream": stream,
            "text": text,
            "timestamp": time.time(),
        }
        self.logs.append(payload)
        if len(self.logs) > MAX_LOG_LINES:
            self.logs = self.logs[-MAX_LOG_LINES:]
        self.queue.put(payload)

    def append_status(self, status: str) -> None:
        payload = {
            "type": "status",
            "status": status,
            "return_code": self.return_code,
            "timestamp": time.time(),
        }
        self.queue.put(payload)

    def terminate(self) -> None:
        if self.process and self.process.poll() is None:
            self.process.terminate()
            self.status = "stopping"
            self.append_status("stopping")

    def is_done(self) -> bool:
        return self.status in {"done", "failed"}


class ExecManager:
    def __init__(self) -> None:
        self.sessions: Dict[str, ExecSession] = {}

    def start(self, cmd: str, cwd: str | None, unsafe: bool = False) -> ExecSession:
        if not cmd:
            raise ValueError("Missing command")
        if not unsafe and not is_safe_command(cmd):
            raise ValueError("Command blocked by safe mode")

        work_dir = str(resolve_dir(cwd)) if cwd else str(REPO_ROOT)
        session_id = uuid.uuid4().hex
        process = subprocess.Popen(
            cmd,
            cwd=work_dir,
            shell=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            bufsize=1,
        )
        session = ExecSession(
            session_id=session_id,
            cmd=cmd,
            cwd=work_dir,
            unsafe=unsafe,
            process=process,
        )
        self.sessions[session_id] = session
        self._start_reader(session, process.stdout, "stdout")
        self._start_reader(session, process.stderr, "stderr")
        self._start_waiter(session)
        return session

    def _start_reader(self, session: ExecSession, stream, label: str) -> None:
        def _reader() -> None:
            if stream is None:
                return
            for line in iter(stream.readline, ""):
                clean = line.rstrip("\n")
                if clean:
                    session.append_log(label, clean)
            try:
                stream.close()
            except Exception:
                pass

        thread = threading.Thread(target=_reader, daemon=True)
        thread.start()

    def _start_waiter(self, session: ExecSession) -> None:
        def _waiter() -> None:
            session.process.wait()
            session.return_code = session.process.returncode
            session.status = "done" if session.return_code == 0 else "failed"
            session.append_status("exit")

        thread = threading.Thread(target=_waiter, daemon=True)
        thread.start()

    def get(self, session_id: str) -> Optional[ExecSession]:
        return self.sessions.get(session_id)

    def stop(self, session_id: str) -> bool:
        session = self.sessions.get(session_id)
        if not session:
            return False
        session.terminate()
        return True

    def backlog(self, session_id: str) -> list[dict]:
        session = self.sessions.get(session_id)
        if not session:
            return []
        return list(session.logs)


EXEC_MANAGER = ExecManager()