didactyl/tests/harness/agent_process.py

from __future__ import annotations

import os
import re
import signal
import subprocess
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

from .didactyl_client import DidactylClient


@dataclass
class AgentProcess:
    binary_path: str
    config_path: str
    api_port: int = 8485
    api_bind: str = "127.0.0.1"
    debug_level: int = 5
    log_file: Optional[str] = None
    base_url: Optional[str] = None

    def __post_init__(self) -> None:
        self.binary_path = str(Path(self.binary_path).resolve())
        self.config_path = str(Path(self.config_path).resolve())
        self.log_file = self.log_file or "tests/results/agent_debug.log"
        scheme = "https"
        self.base_url = self.base_url or f"{scheme}://{self.api_bind}:{self.api_port}"
        self.process: Optional[subprocess.Popen[str]] = None

    def _command(self) -> list[str]:
        return [
            self.binary_path,
            "--config",
            self.config_path,
            "--debug",
            str(self.debug_level),
            "--api-port",
            str(self.api_port),
            "--api-bind",
            self.api_bind,
        ]

    def _pid_exists(self, pid: int) -> bool:
        try:
            os.kill(pid, 0)
            return True
        except OSError:
            return False

    def _collect_listener_pids(self) -> set[int]:
        pids: set[int] = set()
        commands = [
            ["lsof", "-t", f"-iTCP:{self.api_port}", "-sTCP:LISTEN"],
            ["fuser", f"{self.api_port}/tcp"],
            ["ss", "-ltnp"],
        ]

        for cmd in commands:
            try:
                proc = subprocess.run(cmd, capture_output=True, text=True, check=False)
            except FileNotFoundError:
                continue

            output = f"{proc.stdout}\n{proc.stderr}"

            if cmd and cmd[0] == "ss":
                port_token = f":{self.api_port}"
                for line in output.splitlines():
                    if port_token not in line:
                        continue
                    for m in re.finditer(r"pid=(\d+)", line):
                        pids.add(int(m.group(1)))
                continue

            for token in output.split():
                if token.isdigit():
                    pids.add(int(token))

        return pids

    def _kill_stale_port_processes(self) -> None:
        current_pid = os.getpid()
        tracked_pid = self.process.pid if self.process else None
        pids = self._collect_listener_pids()
        targets = [p for p in pids if p != current_pid and p != tracked_pid]

        for pid in targets:
            try:
                os.kill(pid, signal.SIGTERM)
            except OSError:
                pass

        deadline = time.time() + 2.0
        while time.time() < deadline:
            alive = [pid for pid in targets if self._pid_exists(pid)]
            if not alive:
                return
            time.sleep(0.1)

        for pid in targets:
            if self._pid_exists(pid):
                try:
                    os.kill(pid, signal.SIGKILL)
                except OSError:
                    pass

    def _read_log_since(self, offset: int) -> tuple[str, int]:
        path = Path(self.log_file)
        if not path.exists():
            return "", offset

        with path.open("rb") as f:
            f.seek(0, os.SEEK_END)
            end = f.tell()
            if offset > end:
                offset = 0
            f.seek(offset, os.SEEK_SET)
            data = f.read()
            return data.decode("utf-8", errors="replace"), end

    def start(self, timeout: float = 30.0) -> bool:
        if self.is_alive():
            return True

        self._kill_stale_port_processes()

        env = os.environ.copy()
        env["DIDACTYL_LOG_FILE"] = str(self.log_file)

        Path(self.log_file).parent.mkdir(parents=True, exist_ok=True)
        Path(self.log_file).write_text("", encoding="utf-8")

        self.process = subprocess.Popen(
            self._command(),
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            env=env,
        )

        client = DidactylClient(base_url=self.base_url, timeout=2.0, verify_tls=False)
        deadline = time.time() + timeout
        log_offset = 0
        saw_ready_log = False

        while time.time() < deadline:
            if self.process and self.process.poll() is not None:
                return False

            log_chunk, log_offset = self._read_log_since(log_offset)
            if log_chunk:
                if "HTTP API disabled due to bind/init failure" in log_chunk:
                    return False
                if "startup checklist [18] READY: ok" in log_chunk or "entering main poll loop" in log_chunk:
                    saw_ready_log = True

            try:
                data = client.status()
                if data.get("success") and saw_ready_log:
                    return True
            except Exception:
                pass

            time.sleep(0.5)

        return False

    def stop(self, timeout: float = 10.0) -> bool:
        if not self.process:
            return True
        if self.process.poll() is not None:
            return True

        try:
            self.process.send_signal(signal.SIGTERM)
            self.process.wait(timeout=timeout)
            return True
        except subprocess.TimeoutExpired:
            self.process.kill()
            self.process.wait(timeout=5)
            return False

    def restart(self, timeout: float = 30.0) -> bool:
        self.stop()
        return self.start(timeout=timeout)

    def is_alive(self) -> bool:
        return self.process is not None and self.process.poll() is None

    def pid(self) -> Optional[int]:
        return None if not self.process else self.process.pid

    def return_code(self) -> Optional[int]:
        return None if not self.process else self.process.poll()

    def read_pipes(self) -> tuple[str, str]:
        if not self.process:
            return "", ""
        out = ""
        err = ""
        if self.process.stdout:
            out = self.process.stdout.read() or ""
        if self.process.stderr:
            err = self.process.stderr.read() or ""
        return out, err