rp/pr/tools/command.py

import os
import select
import subprocess
import time

from pr.multiplexer import close_multiplexer, create_multiplexer, get_multiplexer

_processes = {}


def _register_process(pid: int, process):
    _processes[pid] = process
    return _processes


def _get_process(pid: int):
    return _processes.get(pid)


def kill_process(pid: int):
    try:
        process = _get_process(pid)
        if process:
            process.kill()
            _processes.pop(pid)

            mux_name = f"cmd-{pid}"
            if get_multiplexer(mux_name):
                close_multiplexer(mux_name)

            return {"status": "success", "message": f"Process {pid} has been killed"}
        else:
            return {"status": "error", "error": f"Process {pid} not found"}
    except Exception as e:
        return {"status": "error", "error": str(e)}


def tail_process(pid: int, timeout: int = 30):
    process = _get_process(pid)
    if process:
        mux_name = f"cmd-{pid}"
        mux = get_multiplexer(mux_name)

        if not mux:
            mux_name, mux = create_multiplexer(mux_name, show_output=True)

        try:
            start_time = time.time()
            timeout_duration = timeout
            stdout_content = ""
            stderr_content = ""

            while True:
                if process.poll() is not None:
                    remaining_stdout, remaining_stderr = process.communicate()
                    if remaining_stdout:
                        mux.write_stdout(remaining_stdout)
                        stdout_content += remaining_stdout
                    if remaining_stderr:
                        mux.write_stderr(remaining_stderr)
                        stderr_content += remaining_stderr

                    if pid in _processes:
                        _processes.pop(pid)

                    close_multiplexer(mux_name)

                    return {
                        "status": "success",
                        "stdout": stdout_content,
                        "stderr": stderr_content,
                        "returncode": process.returncode,
                    }

                if time.time() - start_time > timeout_duration:
                    return {
                        "status": "running",
                        "message": "Process is still running. Call tail_process again to continue monitoring.",
                        "stdout_so_far": stdout_content,
                        "stderr_so_far": stderr_content,
                        "pid": pid,
                    }

                ready, _, _ = select.select([process.stdout, process.stderr], [], [], 0.1)
                for pipe in ready:
                    if pipe == process.stdout:
                        line = process.stdout.readline()
                        if line:
                            mux.write_stdout(line)
                            stdout_content += line
                    elif pipe == process.stderr:
                        line = process.stderr.readline()
                        if line:
                            mux.write_stderr(line)
                            stderr_content += line
        except Exception as e:
            return {"status": "error", "error": str(e)}
    else:
        return {"status": "error", "error": f"Process {pid} not found"}


def run_command(command, timeout=30, monitored=False):
    mux_name = None
    try:
        process = subprocess.Popen(
            command,
            shell=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )
        _register_process(process.pid, process)

        mux_name, mux = create_multiplexer(f"cmd-{process.pid}", show_output=True)

        start_time = time.time()
        timeout_duration = timeout
        stdout_content = ""
        stderr_content = ""

        while True:
            if process.poll() is not None:
                remaining_stdout, remaining_stderr = process.communicate()
                if remaining_stdout:
                    mux.write_stdout(remaining_stdout)
                    stdout_content += remaining_stdout
                if remaining_stderr:
                    mux.write_stderr(remaining_stderr)
                    stderr_content += remaining_stderr

                if process.pid in _processes:
                    _processes.pop(process.pid)

                close_multiplexer(mux_name)

                return {
                    "status": "success",
                    "stdout": stdout_content,
                    "stderr": stderr_content,
                    "returncode": process.returncode,
                }

            if time.time() - start_time > timeout_duration:
                return {
                    "status": "running",
                    "message": f"Process still running after {timeout}s timeout. Use tail_process({process.pid}) to monitor or kill_process({process.pid}) to terminate.",
                    "stdout_so_far": stdout_content,
                    "stderr_so_far": stderr_content,
                    "pid": process.pid,
                    "mux_name": mux_name,
                }

            ready, _, _ = select.select([process.stdout, process.stderr], [], [], 0.1)
            for pipe in ready:
                if pipe == process.stdout:
                    line = process.stdout.readline()
                    if line:
                        mux.write_stdout(line)
                        stdout_content += line
                elif pipe == process.stderr:
                    line = process.stderr.readline()
                    if line:
                        mux.write_stderr(line)
                        stderr_content += line
    except Exception as e:
        if mux_name:
            close_multiplexer(mux_name)
        return {"status": "error", "error": str(e)}


def run_command_interactive(command):
    try:
        return_code = os.system(command)
        return {"status": "success", "returncode": return_code}
    except Exception as e:
        return {"status": "error", "error": str(e)}