planet/backend/scripts/system_restart_runner.py

from __future__ import annotations

import argparse
import os
import shlex
import subprocess
import sys
import time
from pathlib import Path
from urllib.error import URLError
from urllib.request import urlopen


ROOT_DIR = Path(__file__).resolve().parents[2]
BACKEND_DIR = ROOT_DIR / "backend"

if str(BACKEND_DIR) not in sys.path:
    sys.path.insert(0, str(BACKEND_DIR))

from app.services.system_control import (  # noqa: E402
    append_task_log,
    clear_active_task_id,
    get_allowed_command,
    get_action_recovery_mode,
    upsert_task_state,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--task-id", required=True)
    parser.add_argument("--action", required=True)
    return parser.parse_args()


def wait_for_http(url: str, timeout_seconds: int = 90, interval_seconds: float = 2.0) -> bool:
    deadline = time.time() + timeout_seconds
    success_streak = 0

    while time.time() < deadline:
        try:
            with urlopen(url, timeout=2) as response:
                if response.status == 200:
                    success_streak += 1
                    if success_streak >= 2:
                        return True
                else:
                    success_streak = 0
        except URLError:
            success_streak = 0
        except Exception:
            success_streak = 0

        time.sleep(interval_seconds)

    return False


def wait_for_recovery(action: str) -> tuple[bool, str]:
    recovery_mode = get_action_recovery_mode(action)
    if recovery_mode == "backend":
        return wait_for_http("http://localhost:8000/health"), "backend health recovery"
    if recovery_mode == "frontend":
        return wait_for_http("http://localhost:3000"), "frontend entrypoint recovery"
    if recovery_mode == "ai-provider":
        return wait_for_http("http://localhost:8010/health"), "ai provider health recovery"
    if recovery_mode == "database":
        return True, "database container restart completion"
    if recovery_mode == "system":
        backend_ok = wait_for_http("http://localhost:8000/health")
        frontend_ok = wait_for_http("http://localhost:3000")
        return backend_ok and frontend_ok, "system service recovery"
    return False, "unsupported recovery mode"


def main() -> int:
    args = parse_args()
    command = get_allowed_command(args.action)
    recovery_mode = get_action_recovery_mode(args.action)
    if command is None or recovery_mode is None:
        upsert_task_state(
            args.task_id,
            action=args.action,
            status="failed",
            stage="failed",
            message="Unsupported system action",
        )
        clear_active_task_id(args.task_id)
        return 1

    upsert_task_state(
        args.task_id,
        action=args.action,
        status="running",
        stage="spawning",
        message="Spawning restart command",
    )
    append_task_log(args.task_id, f"accepted {args.action} request")
    append_task_log(args.task_id, f"resolved command: {' '.join(command)}")

    env = os.environ.copy()
    env["PATH"] = f"{Path.home() / '.bun' / 'bin'}:{Path.home() / '.local' / 'bin'}:{env.get('PATH', '')}"

    try:
        upsert_task_state(
            args.task_id,
            action=args.action,
            status="running",
            stage="stopping",
            message="Restart command is running",
        )
        append_task_log(args.task_id, "restart command started")

        shell_command = shlex.join(command)
        completed = subprocess.run(
            ["zsh", "-ic", shell_command],
            cwd=str(ROOT_DIR),
            env=env,
            capture_output=True,
            text=True,
            check=False,
        )

        if completed.stdout.strip():
            for line in completed.stdout.strip().splitlines()[-20:]:
                append_task_log(args.task_id, line)
        if completed.stderr.strip():
            for line in completed.stderr.strip().splitlines()[-20:]:
                append_task_log(args.task_id, line)

        if completed.returncode != 0:
            upsert_task_state(
                args.task_id,
                action=args.action,
                status="failed",
                stage="failed",
                message=f"Restart command failed with exit code {completed.returncode}",
            )
            clear_active_task_id(args.task_id)
            return completed.returncode

        upsert_task_state(
            args.task_id,
            action=args.action,
            status="running",
            stage="waiting_for_health",
            message=f"Waiting for {recovery_mode} recovery",
        )
        append_task_log(args.task_id, f"waiting for {recovery_mode} recovery")

        recovered, recovery_label = wait_for_recovery(args.action)
        if recovered:
            upsert_task_state(
                args.task_id,
                action=args.action,
                status="succeeded",
                stage="healthy",
                message="Restart completed successfully",
            )
            append_task_log(args.task_id, f"{recovery_label} completed")
            clear_active_task_id(args.task_id)
            return 0

        upsert_task_state(
            args.task_id,
            action=args.action,
            status="timeout",
            stage="failed",
            message="Restart timed out waiting for recovery",
        )
        append_task_log(args.task_id, f"{recovery_label} timed out")
        clear_active_task_id(args.task_id)
        return 2
    except Exception as exc:
        append_task_log(args.task_id, f"runner exception: {exc}")
        upsert_task_state(
            args.task_id,
            action=args.action,
            status="failed",
            stage="failed",
            message=f"Restart runner failed: {exc}",
        )
        clear_active_task_id(args.task_id)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())