chopratejas--headroom/headroom/cli/install.py

"""Persistent install / deployment CLI commands."""

from __future__ import annotations

import shutil
import subprocess
from copy import deepcopy

import click

from headroom.install.health import probe_json, probe_ready
from headroom.install.models import (
    ConfigScope,
    DeploymentManifest,
    InstallPreset,
    ProviderSelectionMode,
    RuntimeKind,
    SupervisorKind,
)
from headroom.install.planner import build_manifest
from headroom.install.providers import apply_mutations, revert_mutations
from headroom.install.runtime import (
    acquire_runtime_start_lock,
    run_foreground,
    runtime_status,
    start_detached_agent,
    start_persistent_docker,
    stop_runtime,
    wait_ready,
)
from headroom.install.state import (
    ManifestError,
    delete_manifest,
    load_manifest,
    save_manifest,
)
from headroom.install.supervisors import (
    install_supervisor,
    remove_supervisor,
    start_supervisor,
    stop_supervisor,
)

from .main import main


@main.group()
def install() -> None:
    """Install and manage persistent Headroom deployments."""


def _require_manifest(profile: str) -> DeploymentManifest:
    try:
        manifest = load_manifest(profile)
    except ManifestError as e:
        raise click.ClickException(str(e)) from None
    if manifest is None:
        raise click.ClickException(f"No deployment profile named '{profile}' is installed.")
    return manifest


def _start_deployment(manifest: DeploymentManifest, *, assume_start_lock: bool = False) -> None:
    if not assume_start_lock:
        with acquire_runtime_start_lock(manifest.profile) as acquired:
            if not acquired:
                click.echo(f"Deployment '{manifest.profile}' start is already in progress.")
                return
            _start_deployment(manifest, assume_start_lock=True)
            return

    if probe_ready(manifest.health_url):
        return
    if manifest.preset == InstallPreset.PERSISTENT_DOCKER.value and shutil.which("docker") is None:
        raise click.ClickException(
            "Docker is required for this deployment but 'docker' was not found on PATH."
        )
    if runtime_status(manifest) == "running":
        if wait_ready(manifest, timeout_seconds=_STARTUP_READY_TIMEOUT_SECONDS):
            return
        stop_runtime(manifest)

    try:
        if manifest.preset == InstallPreset.PERSISTENT_DOCKER.value:
            start_persistent_docker(manifest)
        elif manifest.supervisor_kind == SupervisorKind.SERVICE.value:
            start_supervisor(manifest)
        else:
            start_detached_agent(manifest.profile)
    except FileNotFoundError as e:
        # A required external binary (docker, launchctl, systemctl) is missing.
        raise click.ClickException(f"Cannot start deployment '{manifest.profile}': {e}") from None
    except subprocess.CalledProcessError as e:
        raise click.ClickException(
            f"Cannot start deployment '{manifest.profile}': command failed "
            f"({' '.join(map(str, e.cmd)) if isinstance(e.cmd, list | tuple) else e.cmd})"
        ) from None

    if not wait_ready(manifest, timeout_seconds=45):
        raise click.ClickException(
            f"Deployment '{manifest.profile}' did not become ready after start."
        )


def _stop_deployment(manifest: DeploymentManifest) -> None:
    if manifest.supervisor_kind == SupervisorKind.SERVICE.value:
        stop_supervisor(manifest)
    stop_runtime(manifest)


def _remove_deployment(manifest: DeploymentManifest) -> None:
    try:
        _stop_deployment(manifest)
    except Exception:
        pass
    try:
        remove_supervisor(manifest)
    except Exception:
        pass
    try:
        revert_mutations(manifest)
    except Exception:
        pass
    delete_manifest(manifest.profile)


def _restore_deployment(manifest: DeploymentManifest) -> None:
    restored = deepcopy(manifest)
    restored.mutations = apply_mutations(restored)
    restored.artifacts = install_supervisor(restored)
    save_manifest(restored)
    _start_deployment(restored)


def _reject_task_lifecycle(manifest: DeploymentManifest, action: str) -> None:
    if manifest.supervisor_kind == SupervisorKind.TASK.value:
        raise click.ClickException(
            f"Deployment '{manifest.profile}' uses persistent-task scheduling; "
            f"`headroom install {action}` is not supported for task deployments."
        )


@install.command("apply")
@click.option(
    "--preset",
    type=click.Choice([preset.value for preset in InstallPreset]),
    default=InstallPreset.PERSISTENT_SERVICE.value,
    show_default=True,
    help="Persistent runtime preset to install.",
)
@click.option(
    "--runtime",
    type=click.Choice([runtime.value for runtime in RuntimeKind]),
    default=RuntimeKind.PYTHON.value,
    show_default=True,
    help="Runtime used to execute Headroom for service/task modes.",
)
@click.option(
    "--scope",
    type=click.Choice([scope.value for scope in ConfigScope]),
    default=ConfigScope.USER.value,
    show_default=True,
    help="Where to apply persistent configuration.",
)
@click.option(
    "--providers",
    "provider_mode",
    type=click.Choice([mode.value for mode in ProviderSelectionMode]),
    default=ProviderSelectionMode.AUTO.value,
    show_default=True,
    help="Target selection mode for direct tool configuration.",
)
@click.option(
    "--target",
    "targets",
    multiple=True,
    type=click.Choice(["claude", "copilot", "codex", "aider", "cursor", "openclaw", "opencode"]),
    help="Tool target to configure when --providers manual is used.",
)
@click.option("--profile", default="default", show_default=True, help="Deployment profile name.")
@click.option(
    "--port",
    "-p",
    default=8787,
    type=click.IntRange(1, 65535),
    show_default=True,
    help="Persistent proxy port.",
)
@click.option(
    "--backend",
    default="anthropic",
    show_default=True,
    help="Proxy backend for the persistent runtime.",
)
@click.option(
    "--anyllm-provider",
    default=None,
    help="Provider for any-llm backends when --backend anyllm is used.",
)
@click.option("--region", default=None, help="Cloud region for Bedrock / Vertex style backends.")
@click.option(
    "--mode", "proxy_mode", default="token", show_default=True, help="Proxy optimization mode."
)
@click.option("--memory", is_flag=True, help="Enable persistent memory in the proxy runtime.")
@click.option(
    "--telemetry",
    is_flag=True,
    help="Opt in to anonymous telemetry in the runtime (off by default).",
)
@click.option(
    "--no-telemetry",
    is_flag=True,
    help="Force anonymous telemetry off in the runtime (already the default).",
)
@click.option(
    "--image",
    default="ghcr.io/chopratejas/headroom:latest",
    show_default=True,
    help="Docker image to use when runtime=docker or preset=persistent-docker.",
)
@click.option(
    "--no-http2",
    is_flag=True,
    help="Disable HTTP/2 in the persistent runtime (enabled by default).",
)
def install_apply(
    preset: str,
    runtime: str,
    scope: str,
    provider_mode: str,
    targets: tuple[str, ...],
    profile: str,
    port: int,
    backend: str,
    anyllm_provider: str | None,
    region: str | None,
    proxy_mode: str,
    memory: bool,
    telemetry: bool,
    no_telemetry: bool,
    image: str,
    no_http2: bool,
) -> None:
    """Install a persistent Headroom deployment."""

    if anyllm_provider and backend != "anyllm":
        click.echo(
            f"Warning: --anyllm-provider is ignored unless --backend anyllm "
            f"(got --backend {backend})."
        )

    if preset == InstallPreset.PERSISTENT_DOCKER.value:
        runtime = RuntimeKind.DOCKER.value

    manifest = build_manifest(
        profile=profile,
        preset=preset,
        runtime_kind=runtime,
        scope=scope,
        provider_mode=provider_mode,
        targets=list(targets),
        port=port,
        backend=backend,
        anyllm_provider=anyllm_provider,
        region=region,
        proxy_mode=proxy_mode,
        memory_enabled=memory,
        telemetry_enabled=telemetry and not no_telemetry,
        image=image,
        no_http2=no_http2,
    )

    try:
        existing = load_manifest(profile)
    except ManifestError as e:
        # A corrupt existing manifest shouldn't block a fresh apply; overwrite it.
        click.echo(f"Warning: {e}; overwriting.")
        existing = None
    if existing is not None:
        click.echo(f"Updating existing deployment profile '{profile}'...")
        _remove_deployment(existing)

    try:
        manifest.mutations = apply_mutations(manifest)
        manifest.artifacts = install_supervisor(manifest)
        save_manifest(manifest)
        _start_deployment(manifest)
    except Exception as exc:
        _remove_deployment(manifest)
        if existing is not None:
            click.echo(f"Restoring previous deployment '{profile}'...")
            _restore_deployment(existing)
        # Surface non-Click errors (OSError, CalledProcessError, …) as a clean
        # message rather than a raw traceback; Click errors pass through as-is.
        if isinstance(exc, click.ClickException | click.Abort):
            raise
        raise click.ClickException(f"Failed to install deployment '{profile}': {exc}") from exc

    click.echo(
        f"Installed persistent deployment '{profile}' "
        f"({manifest.preset}, runtime={manifest.runtime_kind}, scope={manifest.scope})."
    )
    click.echo(f"Health: {manifest.health_url}")
    if manifest.targets:
        click.echo(f"Targets: {', '.join(manifest.targets)}")


@install.command("status")
@click.option("--profile", default="default", show_default=True, help="Deployment profile name.")
def install_status(profile: str) -> None:
    """Show persistent deployment status."""

    manifest = _require_manifest(profile)
    payload = probe_json(manifest.health_url.replace("/readyz", "/health"))
    click.echo(f"Profile:    {manifest.profile}")
    click.echo(f"Preset:     {manifest.preset}")
    click.echo(f"Runtime:    {manifest.runtime_kind}")
    click.echo(f"Supervisor: {manifest.supervisor_kind}")
    click.echo(f"Scope:      {manifest.scope}")
    click.echo(f"Port:       {manifest.port}")
    click.echo(f"Status:     {runtime_status(manifest)}")
    click.echo(f"Healthy:    {'yes' if probe_ready(manifest.health_url) else 'no'}")
    if payload and isinstance(payload, dict):
        click.echo(f"Health URL: {manifest.health_url.replace('/readyz', '/health')}")
        click.echo(f"Backend:    {payload.get('config', {}).get('backend', manifest.backend)}")


@install.command("start")
@click.option("--profile", default="default", show_default=True, help="Deployment profile name.")
def install_start(profile: str) -> None:
    """Start a persistent deployment."""

    manifest = _require_manifest(profile)
    _reject_task_lifecycle(manifest, "start")
    _start_deployment(manifest)
    click.echo(f"Started deployment '{profile}'.")


@install.command("stop")
@click.option("--profile", default="default", show_default=True, help="Deployment profile name.")
def install_stop(profile: str) -> None:
    """Stop a persistent deployment."""

    manifest = _require_manifest(profile)
    _reject_task_lifecycle(manifest, "stop")
    _stop_deployment(manifest)
    click.echo(f"Stopped deployment '{profile}'.")


@install.command("restart")
@click.option("--profile", default="default", show_default=True, help="Deployment profile name.")
def install_restart(profile: str) -> None:
    """Restart a persistent deployment."""

    manifest = _require_manifest(profile)
    _reject_task_lifecycle(manifest, "restart")
    _stop_deployment(manifest)
    _start_deployment(manifest)
    click.echo(f"Restarted deployment '{profile}'.")


@install.command("remove")
@click.option("--profile", default="default", show_default=True, help="Deployment profile name.")
def install_remove(profile: str) -> None:
    """Remove a persistent deployment and undo managed config."""

    manifest = _require_manifest(profile)
    try:
        if manifest.supervisor_kind == SupervisorKind.SERVICE.value:
            stop_supervisor(manifest)
    except Exception:
        pass
    try:
        stop_runtime(manifest)
    except Exception:
        pass
    try:
        remove_supervisor(manifest)
    except Exception:
        pass
    revert_mutations(manifest)
    delete_manifest(profile)
    click.echo(f"Removed deployment '{profile}'.")


@install.group("agent", hidden=True)
def install_agent() -> None:
    """Hidden runtime helpers used by persistent supervisors."""


@install_agent.command("run")
@click.option("--profile", default="default", show_default=True, help="Deployment profile name.")
def install_agent_run(profile: str) -> None:
    """Run the persistent runtime in the foreground."""

    manifest = _require_manifest(profile)
    raise SystemExit(run_foreground(manifest))


_STARTUP_READY_TIMEOUT_SECONDS = 15


@install_agent.command("ensure")
@click.option("--profile", default="default", show_default=True, help="Deployment profile name.")
def install_agent_ensure(profile: str) -> None:
    """Ensure a persistent deployment is healthy, starting it when needed."""

    manifest = _require_manifest(profile)
    if probe_ready(manifest.health_url):
        click.echo(f"Deployment '{profile}' is already healthy.")
        return
    with acquire_runtime_start_lock(manifest.profile) as acquired:
        if not acquired:
            click.echo(f"Deployment '{profile}' start is already in progress.")
            return
        # Double-check after acquiring the lock — another ensure may have
        # started the runtime while we waited for the lock.
        if probe_ready(manifest.health_url):
            click.echo(f"Deployment '{profile}' is already healthy.")
            return
        if runtime_status(manifest) == "running":
            # Runtime exists but isn't ready yet — give it a grace period
            # before deciding it's wedged and restarting.
            if wait_ready(manifest, timeout_seconds=_STARTUP_READY_TIMEOUT_SECONDS):
                click.echo(f"Deployment '{profile}' is healthy.")
                return
            stop_runtime(manifest)
        _start_deployment(manifest, assume_start_lock=True)
    click.echo(f"Deployment '{profile}' is healthy.")