rig-cli 1.0.0__py3-none-any.whl

rig/commands/up/rollback.py

@@ -0,0 +1,49 @@
+"""Rollback lifecycle operations for failed service launches."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any
+
+from rig.commands.common import _stop_record
+from rig.commands.up.context import RollbackContext
+from rig.core.state import write_state
+from rig.manifest.models import Manifest
+
+
+def _rollback_step(
+    name: str,
+    ctx: tuple[dict[str, Any], Path, Path, Manifest | None],
+    failed_services: set[str],
+) -> None:
+    state, state_path, root, manifest = ctx
+    rec = state["services"].get(name)
+    if rec is None:
+        return
+    if manifest and any(
+        d in failed_services or d in state["services"] for d in manifest.dependents(name)
+    ):
+        return
+    outcome = _stop_record(rec, root)
+    if outcome in ("terminated", "killed", "stale"):
+        state["services"].pop(name, None)
+    else:
+        failed_services.add(name)
+    write_state(state_path, state)
+
+
+def rollback_started(started: Sequence[str], ctx: RollbackContext) -> None:
+    failed_services: set[str] = set()
+    for name in reversed(list(started)):
+        _rollback_step(name, (ctx.state, ctx.state_path, ctx.root, ctx.manifest), failed_services)
+
+
+def _rollback(
+    state: dict[str, Any], state_path: Path, started: Sequence[str], *args: Any, **kwargs: Any
+) -> None:
+    """External test adapter delegating to rollback_started."""
+    argv = list(args)
+    root = kwargs.get("root") or (argv.pop(0) if argv else Path("."))
+    man = kwargs.get("manifest") or (argv.pop(0) if argv else None)
+    rollback_started(started, RollbackContext(state, state_path, root, man))

rig/commands/up/runner.py

@@ -0,0 +1,118 @@
+"""Retry loop, rollback, and process supervisor for rig up."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from pathlib import Path
+from typing import Any
+
+from rig.commands.common import _stop_record, _values_for
+from rig.commands.up.context import RetryContext
+from rig.commands.up.rollback import _rollback, rollback_started
+from rig.commands.up.service import _await_ready, _start_service
+from rig.core.constants import PORT_RETRY_ATTEMPTS
+from rig.core.errors import RigError
+from rig.core.state import write_state
+from rig.manifest.models import Service
+from rig.net.ports import compute_candidate_ports
+
+_LAUNCH_EXCEPTIONS = (
+    KeyboardInterrupt,
+    SystemExit,
+    RigError,
+    OSError,
+    RuntimeError,
+    ValueError,
+)
+
+
+def _spawn_candidate(
+    service: Service,
+    ctx: tuple[Path, Path, str, Mapping[str, Any]],
+    cands: Sequence[int],
+) -> dict[str, Any]:
+    root, runtime, instance, values = ctx
+    return _start_service(service, root, runtime, instance, values, cands)
+
+
+def _handle_launch_failure(
+    exc: Exception | KeyboardInterrupt | SystemExit,
+    service: Service,
+    ctx: tuple[dict[str, Any], Path],
+) -> None:
+    state, state_path = ctx
+    partial = (
+        exc.details.get("partial_record") if isinstance(exc, RigError) and exc.details else None
+    )
+    if isinstance(partial, dict):
+        state["services"][service.name] = partial
+        write_state(state_path, state)
+    raise exc
+
+
+def _handle_unready(
+    record: dict[str, Any],
+    ctx: tuple[Path, Service, dict[str, Any], Path],
+    collided: set[int],
+) -> bool:
+    root, service, state, state_path = ctx
+    outcome = _stop_record(record, root)
+    if record.get("port"):
+        collided.add(int(record["port"]))
+    if outcome in ("terminated", "killed", "stale"):
+        state["services"].pop(service.name, None)
+        write_state(state_path, state)
+        return True
+    return False
+
+
+def _start_attempt(
+    service: Service,
+    paths: tuple[Path, Path, str],
+    ctx: tuple[dict[str, Any], Path, set[int]],
+) -> dict[str, Any] | None:
+    root, runtime, instance = paths
+    state, state_path, collided = ctx
+    values = _values_for(state, root, instance)
+    cands = compute_candidate_ports(service, state, avoid=collided)
+    try:
+        record = _spawn_candidate(service, (root, runtime, instance, values), cands)
+    except _LAUNCH_EXCEPTIONS as exc:
+        _handle_launch_failure(exc, service, (state, state_path))
+    state["services"][service.name] = record
+    write_state(state_path, state)
+    if _await_ready(service, record, root):
+        if record.get("port"):
+            state.setdefault("ports", {})[service.name] = record["port"]
+            write_state(state_path, state)
+        return record
+    _handle_unready(record, (root, service, state, state_path), collided)
+    return None
+
+
+def start_with_retry(service: Service, retry: RetryContext) -> dict[str, Any] | None:
+    attempts = PORT_RETRY_ATTEMPTS if service.type == "port" else 1
+    collided: set[int] = set()
+    for _ in range(attempts):
+        record = _start_attempt(
+            service,
+            (retry.root, retry.runtime, retry.instance),
+            (retry.state, retry.state_path, collided),
+        )
+        if record is not None:
+            return record
+    return None
+
+
+def _start_with_retry(
+    service: Service, root: Path, runtime: Path, *args: Any, **kwargs: Any
+) -> dict[str, Any] | None:
+    """External test adapter delegating to start_with_retry."""
+    argv = list(args)
+    inst = kwargs.get("instance") or (argv.pop(0) if argv else "")
+    st = kwargs.get("state") or (argv.pop(0) if argv else {})
+    sp = kwargs.get("state_path") or (argv.pop(0) if argv else Path("."))
+    return start_with_retry(service, RetryContext(root, runtime, inst, st, sp))
+
+
+__all__ = ["_rollback", "_start_with_retry", "rollback_started", "start_with_retry"]

rig/commands/up/service.py

@@ -0,0 +1,136 @@
+"""Service execution and health checking logic for rig up."""
+
+from __future__ import annotations
+
+import shutil
+import sys
+import time
+from collections.abc import Mapping, Sequence
+from pathlib import Path
+from typing import Any
+
+from rig.commands.up.context import ServiceContext
+from rig.compose.starter import _start_compose_service
+from rig.compose.supervisor import compose_record_alive
+from rig.core.constants import EXIT_EXTERNAL_TOOL
+from rig.core.env import build_service_env, render
+from rig.core.errors import RigError
+from rig.core.state import redact
+from rig.manifest.models import Service
+from rig.net.health import wait_for_http
+from rig.net.ports import _safe_reserve_port, compute_candidate_ports
+from rig.net.probe import port_listener_matches
+from rig.proc.process import identity_matches, pid_alive
+from rig.proc.spawn import spawn_fd_service, spawn_port_service, uvicorn_argv
+
+
+def _validate_service_env(
+    service: Service, root: Path, svc_vals: Mapping[str, Any]
+) -> tuple[Path, dict[str, str]]:
+    cwd = (Path(root) / service.cwd).resolve()
+    if not cwd.is_dir():
+        raise RigError(f"service {service.name!r} working directory {cwd} does not exist")
+    if service.type in ("port", "fd") and not shutil.which("lsof"):
+        raise RigError(
+            "'lsof' required but not on PATH", code="E_EXTERNAL_TOOL", exit_code=EXIT_EXTERNAL_TOOL
+        )
+    env = build_service_env(service.env, service.inherit, Path(root), svc_vals, service.env_files)
+    return cwd, env
+
+
+def _spawn_local_service(
+    service: Service,
+    paths: tuple[Path, Path],
+    ctx: tuple[dict[str, str], Mapping[str, Any], Sequence[int]],
+) -> dict[str, Any]:
+    cwd, log_path = paths
+    env, svc_vals, cands = ctx
+    if service.type == "fd":
+        py_bin = str(render(service.python or sys.executable, svc_vals))
+        cmd = service.command or uvicorn_argv(py_bin, service.app or "", service.factory)
+        return spawn_fd_service(service.name, cmd, cwd, env, log_path, svc_vals, cands)
+    p = _safe_reserve_port(cands)
+    return spawn_port_service(service.name, service.command, cwd, env, log_path, p, svc_vals)
+
+
+def start_service(service: Service, ctx: ServiceContext) -> dict[str, Any]:
+    cwd = (Path(ctx.root) / service.cwd).resolve()
+    log_path = ctx.runtime / "logs" / f"{service.name}.log"
+    cwd, env = _validate_service_env(service, ctx.root, {**ctx.values, "cwd": str(cwd)})
+    if service.type == "compose":
+        return _start_compose_service(service, ctx.root, ctx.instance, env=env)
+    ports = (
+        ctx.candidate_ports
+        if ctx.candidate_ports is not None
+        else compute_candidate_ports(service, {})
+    )
+    rec = _spawn_local_service(
+        service, (cwd, log_path), (env, {**ctx.values, "cwd": str(cwd)}, ports)
+    )
+    rec.update(
+        {
+            "log": str(log_path),
+            "env": redact(env),
+            "depends_on": list(service.depends_on),
+            "health": service.healthcheck_path,
+            "healthcheck_path": service.healthcheck_path,
+        }
+    )
+    return rec
+
+
+def _start_service(
+    service: Service, root: Path, runtime: Path, *args: Any, **kwargs: Any
+) -> dict[str, Any]:
+    """External test adapter delegating to start_service."""
+    argv = list(args)
+    inst = kwargs.get("instance") or (argv.pop(0) if argv else "")
+    vals = kwargs.get("values") or (argv.pop(0) if argv else {})
+    cands = kwargs.get("candidate_ports") or (argv.pop(0) if argv else None)
+    return start_service(service, ServiceContext(root, runtime, inst, vals, cands))
+
+
+def _check_port_listener(port: int, ids: tuple[Any, Any], record: Mapping[str, Any]) -> bool:
+    pid, pgid = ids
+    p_id = pid if isinstance(pid, int) else None
+    pg_id = pgid if isinstance(pgid, int) else None
+    if not port_listener_matches(port, pgid=pg_id, pid=p_id):
+        return False
+    return pid_alive(p_id) and identity_matches(record) if isinstance(p_id, int) else True
+
+
+def _await_process_service(
+    service: Service, record: Mapping[str, Any], target: tuple[Any, Any, int | None]
+) -> bool:
+    pid, pgid, port = target
+    p_id = pid if isinstance(pid, int) else None
+    pg_id = pgid if isinstance(pgid, int) else None
+    if service.healthcheck_path and port is not None:
+        wait_target = (service.healthcheck_timeout, p_id, pg_id)
+        ok = wait_for_http(port, service.healthcheck_path, wait_target)
+        return ok and _check_port_listener(port, (pid, pgid), record)
+    if isinstance(pid, int):
+        time.sleep(0.3)
+        matched = port is None or port_listener_matches(port, pgid=pg_id, pid=p_id)
+        return matched and pid_alive(pid) and identity_matches(record)
+    return True
+
+
+def _await_ready(service: Service, record: Mapping[str, Any], root: Path = Path(".")) -> bool:
+    pid, pgid = record.get("pid"), record.get("pgid")
+    if isinstance(pid, int) and not pid_alive(pid):
+        return False
+    port = int(record["port"]) if record.get("port") and str(record["port"]).isdigit() else None
+    if service.type == "compose":
+        p_id = pid if isinstance(pid, int) else None
+        pg_id = pgid if isinstance(pgid, int) else None
+        if (
+            service.healthcheck_path
+            and port is not None
+            and not wait_for_http(
+                port, service.healthcheck_path, (service.healthcheck_timeout, p_id, pg_id)
+            )
+        ):
+            return False
+        return compose_record_alive(record, root)
+    return _await_process_service(service, record, (pid, pgid, port))

rig/compose/__init__.py

@@ -0,0 +1 @@
+"""Docker and Docker Compose integration for rig."""

rig/compose/client.py

@@ -0,0 +1,144 @@
+"""Docker and Docker Compose CLI client invocation."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+from collections.abc import Mapping, Sequence
+from pathlib import Path
+from typing import Any
+
+from rig.compose.context import (
+    INHERIT_DOCKER_HOST,
+    _pin_docker_endpoint,
+)
+from rig.core.constants import (
+    COMPOSE_DISCOVERY_TIMEOUT_SECS,
+    COMPOSE_UP_TIMEOUT_SECS,
+    DOCKER_CLIENT_ENV_PASSTHROUGH,
+    EXIT_EXTERNAL_TOOL,
+)
+from rig.core.errors import RigError
+
+
+def compose_argv(
+    instance: str,
+    root: Path,
+    *extra_args: Any,
+    **kwargs: Any,
+) -> list[str]:
+    """Build a Compose command scoped to this checkout."""
+    p = list(extra_args)
+    compose_file = p.pop(0) if p else kwargs["compose_file"]
+    args = p.pop(0) if p else kwargs.get("args", ())
+    context = p.pop(0) if p else kwargs.get("context")
+
+    return [
+        "docker",
+        *(["--context", str(context)] if context else []),
+        "compose",
+        "--project-directory",
+        str(root),
+        "-p",
+        instance,
+        "-f",
+        str(compose_file),
+        *args,
+    ]
+
+
+def parse_compose_port(output: str) -> int:
+    """Extract the host port from ``docker compose port`` output."""
+    line = output.strip().splitlines()[-1].strip() if output.strip() else ""
+    _, sep, port = line.rpartition(":")
+    if not sep or not port.isdigit():
+        raise RigError(f"no published host port in compose output {output!r}")
+    return int(port)
+
+
+def compose_file_present(record: Mapping[str, Any], root: Path) -> bool:
+    """Return ``True`` when this record's compose file is still on disk."""
+    compose_file = record.get("compose_file")
+    if not compose_file:
+        return False
+    return (Path(root) / str(compose_file)).is_file()
+
+
+def _exec_docker_cmd(
+    argv: list[str],
+    env: Mapping[str, str],
+    timeout_spec: tuple[float, str],
+) -> subprocess.CompletedProcess:
+    timeout, timeout_msg = timeout_spec
+    try:
+        return subprocess.run(argv, capture_output=True, text=True, timeout=timeout, env=env)
+    except FileNotFoundError:
+        msg = "docker is not installed or not on PATH"
+        raise RigError(msg, code="E_EXTERNAL_TOOL", exit_code=EXIT_EXTERNAL_TOOL) from None
+    except subprocess.TimeoutExpired:
+        raise RigError(timeout_msg) from None
+
+
+def _unpack_compose_call(
+    args: Sequence[Any], kwargs: Mapping[str, Any]
+) -> tuple[Path, Sequence[str], str | None, float, Mapping[str, str] | None, Any]:
+    p = list(args)
+    cfile = Path(p.pop(0) if p else kwargs["compose_file"])
+    cmd_args = p.pop(0) if p else kwargs.get("args", ())
+    context = p.pop(0) if p else kwargs.get("context")
+    timeout = float(p.pop(0) if p else kwargs.get("timeout", COMPOSE_UP_TIMEOUT_SECS))
+    env = p.pop(0) if p else kwargs.get("env")
+    host = p.pop(0) if p else kwargs.get("docker_host", INHERIT_DOCKER_HOST)
+    return cfile, cmd_args, context, timeout, env, host
+
+
+def run_compose(
+    instance: str,
+    root: Path,
+    *args: Any,
+    **kwargs: Any,
+) -> subprocess.CompletedProcess:
+    cfile, cmd_args, context, timeout, env, host = _unpack_compose_call(args, kwargs)
+    argv = compose_argv(instance, root, cfile, cmd_args, context)
+    cmd_env = dict(os.environ) if env is None else dict(env)
+    _pin_docker_endpoint(cmd_env, context, host)
+    msg = f"compose command timed out: {' '.join(cmd_args)}"
+    return _exec_docker_cmd(argv, cmd_env, (timeout, msg))
+
+
+def _apply_env_passthrough(cmd_env: dict[str, str], env: Mapping[str, str]) -> None:
+    for name in DOCKER_CLIENT_ENV_PASSTHROUGH:
+        value = env.get(name)
+        if value is None:
+            cmd_env.pop(name, None)
+        else:
+            cmd_env[name] = str(value)
+
+
+def _prepare_docker_env(
+    env: Mapping[str, str] | None, context: str | None, docker_host: Any
+) -> dict[str, str]:
+    cmd_env = dict(os.environ)
+    if env is not None:
+        _apply_env_passthrough(cmd_env, env)
+    _pin_docker_endpoint(cmd_env, context, docker_host)
+    return cmd_env
+
+
+def run_docker(
+    args: Sequence[str],
+    *extra: Any,
+    **kwargs: Any,
+) -> subprocess.CompletedProcess:
+    """Run one plain ``docker`` command against one explicit Docker endpoint."""
+    p = list(extra)
+    context = p.pop(0) if p else kwargs.get("context")
+    timeout = float(p.pop(0) if p else kwargs.get("timeout", COMPOSE_DISCOVERY_TIMEOUT_SECS))
+    docker_host = p.pop(0) if p else kwargs.get("docker_host", INHERIT_DOCKER_HOST)
+    env = p.pop(0) if p else kwargs.get("env")
+
+    argv = ["docker", "--context", str(context)] if context else ["docker"]
+    argv.extend(args)
+    cmd_env = _prepare_docker_env(env, context, docker_host)
+    msg = f"docker command timed out: {' '.join(args)}"
+    return _exec_docker_cmd(argv, cmd_env, (timeout, msg))

rig/compose/context.py

@@ -0,0 +1,56 @@
+"""Docker context and endpoint resolution."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+from collections.abc import Mapping
+from typing import Any
+
+INHERIT_DOCKER_HOST: Any = object()
+
+
+def record_docker_endpoint(record: Mapping[str, Any]) -> tuple[Any, Any]:
+    """Return the Docker context and host one record was started against."""
+    return record.get("docker_context"), record.get("docker_host", INHERIT_DOCKER_HOST)
+
+
+def resolve_current_docker_context(env: Mapping[str, str] | None = None) -> str | None:
+    """Return the name of the Docker context that is active right now."""
+    cmd_env = dict(os.environ) if env is None else dict(env)
+    named = cmd_env.get("DOCKER_CONTEXT")
+    if named:
+        return named
+    try:
+        probe = subprocess.run(
+            ["docker", "context", "show"],
+            capture_output=True,
+            text=True,
+            timeout=5.0,
+            env=cmd_env,
+        )
+    except (OSError, subprocess.SubprocessError):
+        return None
+    if probe.returncode != 0:
+        return None
+    return probe.stdout.strip() or None
+
+
+def _pin_docker_endpoint(cmd_env: dict[str, str], context: Any, docker_host: Any) -> dict[str, str]:
+    """Point one Docker command at the endpoint its record was started against."""
+    cmd_env.pop("DOCKER_CONTEXT", None)
+    if docker_host is INHERIT_DOCKER_HOST:
+        return cmd_env
+    if docker_host:
+        cmd_env["DOCKER_HOST"] = str(docker_host)
+    else:
+        cmd_env.pop("DOCKER_HOST", None)
+    return cmd_env
+
+
+def record_compose_env(record: Mapping[str, Any]) -> dict[str, str] | None:
+    """Return the environment one compose record was started with, if it was recorded."""
+    env = record.get("compose_env")
+    if not isinstance(env, Mapping):
+        return None
+    return {str(key): str(value) for key, value in env.items()}

rig/compose/discovery.py

@@ -0,0 +1,121 @@
+"""Compose container discovery, initialization, and port resolution."""
+
+from __future__ import annotations
+
+import subprocess
+import time
+from pathlib import Path
+from typing import Any
+
+from rig.compose.client import parse_compose_port
+from rig.core.constants import EXIT_OP_FAILED
+from rig.core.errors import RigError
+from rig.core.state import redact
+
+
+def init_compose_record(
+    service: Any,
+    instance: str,
+    target: tuple[Path, str | None, str | None, dict[str, str] | None],
+) -> dict[str, Any]:
+    """Build the initial state record for a starting compose service."""
+    cfile, ctx, host, env = target
+    rec: dict[str, Any] = {
+        "name": service.name,
+        "type": "compose",
+        "pid": None,
+        "pgid": None,
+        "instance": instance,
+        "compose_file": str(cfile),
+        "compose_service": service.compose_service,
+        "docker_context": ctx,
+        "docker_host": host,
+        "container": "",
+        "port": None,
+        "url": None,
+        "started_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
+        "depends_on": list(service.depends_on),
+        "health": service.healthcheck_path,
+        "healthcheck_path": service.healthcheck_path,
+    }
+    return {**rec, **({"compose_env": redact(env)} if env is not None else {})}
+
+
+INTERRUPTION_EXCEPTIONS = (
+    KeyboardInterrupt,
+    SystemExit,
+    OSError,
+    RigError,
+    subprocess.SubprocessError,
+    RuntimeError,
+    ValueError,
+)
+
+
+def stranded_error(msg: str, record: dict[str, Any]) -> RigError:
+    """Format RigError when a container failed to start and could not be cleaned up."""
+    return RigError(
+        f"{msg}; the partial container could not be removed and stays recorded",
+        code="E_COMPOSE_FAILED",
+        exit_code=EXIT_OP_FAILED,
+        hint="run 'rig down' or 'rig prune --force' to reclaim it",
+        details={"partial_record": record},
+    )
+
+
+def handle_compose_interruption(
+    exc: BaseException,
+    info: tuple[Any, dict[str, Any]],
+    cleanup_fn: Any,
+    *args: Any,
+) -> None:
+    """Handle interruption during compose start or discovery."""
+    action = args[0] if args else "the start of"
+    service, record = info
+    if not cleanup_fn():
+        name = type(exc).__name__
+        msg = f"{action} {service.name!r} was interrupted by {name}"
+        raise stranded_error(msg, record) from exc
+    raise exc
+
+
+def _discover_container_id(service: Any, compose_fn: Any, fail_fn: Any) -> str:
+    try:
+        ids = compose_fn(["ps", "-q", str(service.compose_service)], timeout=60.0)
+    except (RigError, OSError, subprocess.SubprocessError, RuntimeError) as exc:
+        raise fail_fn(f"compose failed to list containers for {service.name!r}: {exc}") from exc
+    if ids.returncode != 0:
+        err = ids.stderr.strip() or ids.stdout.strip()
+        raise fail_fn(f"compose failed to list containers for {service.name!r}: {err}")
+    container = ids.stdout.strip().splitlines()[0].strip() if ids.stdout.strip() else ""
+    if not container:
+        raise fail_fn(f"compose reported no container for {service.name!r}")
+    return container
+
+
+def _discover_service_port(service: Any, compose_fn: Any, fail_fn: Any) -> int:
+    try:
+        pub = compose_fn(
+            ["port", str(service.compose_service), str(service.compose_port)],
+            timeout=60.0,
+        )
+    except (RigError, OSError, subprocess.SubprocessError, RuntimeError) as exc:
+        raise fail_fn(f"compose failed to resolve port for {service.name!r}: {exc}") from exc
+    if pub.returncode != 0:
+        err = pub.stderr.strip() or pub.stdout.strip() or "no output"
+        raise fail_fn(f"compose failed to resolve port for {service.name!r}: {err}")
+    return parse_compose_port(pub.stdout)
+
+
+def discover_container_and_port(
+    service: Any,
+    record: dict[str, Any],
+    helpers: tuple[Any, Any],
+) -> None:
+    """Populate record with running container ID and resolved loopback port."""
+    compose_fn, fail_fn = helpers
+    record["container"] = _discover_container_id(service, compose_fn, fail_fn)
+    if service.compose_port:
+        port = _discover_service_port(service, compose_fn, fail_fn)
+        record["port"] = port
+        record["url"] = f"http://127.0.0.1:{port}"