rig-cli 1.0.0__py3-none-any.whl

rig/net/ports.py

@@ -0,0 +1,141 @@
+"""Loopback socket allocation, port reservation, and sticky lease computation."""
+
+from __future__ import annotations
+
+import contextlib
+import socket
+import time
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+from rig.core.constants import PORT_MAX, PORT_MIN, PORT_RELEASE_TIMEOUT_SECS
+from rig.net.registry import (
+    get_allocated_ports_for_others,
+    get_or_allocate_port,
+    port_is_free,
+)
+
+DEFAULT_PORT_WINDOW = 50
+
+_BASE_PORT_RULES = (
+    (("front", "web", "ui", "client", "next", "vite"), 3000),
+    (("back", "api", "server", "app", "worker"), 8000),
+    (("doc", "storybook", "admin"), 4000),
+)
+
+
+def _bind_candidate_port(port: int) -> tuple[socket.socket, int] | None:
+    if port < PORT_MIN or port > PORT_MAX:
+        return None
+    listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    try:
+        listener.bind(("127.0.0.1", port))
+        listener.listen(socket.SOMAXCONN)
+    except OSError:
+        listener.close()
+        return None
+    else:
+        return listener, port
+
+
+def _find_candidate_listener(
+    candidate_ports: Sequence[int],
+) -> tuple[socket.socket, int] | None:
+    return next((r for p in candidate_ports if (r := _bind_candidate_port(p)) is not None), None)
+
+
+def allocate_listener(candidate_ports: Sequence[int] | None = None) -> tuple[socket.socket, int]:
+    """Bind and listen on a loopback port, keeping ownership of it."""
+    if candidate_ports and (candidate := _find_candidate_listener(candidate_ports)) is not None:
+        return candidate
+    listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    try:
+        listener.bind(("127.0.0.1", 0))
+        listener.listen(socket.SOMAXCONN)
+    except BaseException:
+        listener.close()
+        raise
+    return listener, listener.getsockname()[1]
+
+
+def _safe_allocate_listener(
+    candidate_ports: Sequence[int] | None = None,
+) -> tuple[socket.socket, int]:
+    try:
+        return allocate_listener(candidate_ports) if candidate_ports else allocate_listener()
+    except TypeError:
+        return allocate_listener()
+
+
+def reserve_port(candidate_ports: Sequence[int] | None = None) -> int:
+    """Return a currently free loopback port for a service that cannot inherit a socket."""
+    listener, port = _safe_allocate_listener(candidate_ports)
+    listener.close()
+    return port
+
+
+def _safe_reserve_port(candidate_ports: Sequence[int] | None = None) -> int:
+    try:
+        return reserve_port(candidate_ports) if candidate_ports else reserve_port()
+    except TypeError:
+        return reserve_port()
+
+
+def wait_for_port_release(port: int, timeout: float = PORT_RELEASE_TIMEOUT_SECS) -> bool:
+    deadline = time.monotonic() + timeout
+    while True:
+        if port_is_free(port):
+            return True
+        if time.monotonic() >= deadline:
+            return False
+        time.sleep(0.05)
+
+
+def default_base_port_for_service(name: str) -> int:
+    """Return a human-friendly default base port based on service name conventions."""
+    slug = name.lower()
+    return next((port for keys, port in _BASE_PORT_RULES if any(k in slug for k in keys)), 5000)
+
+
+def _collect_excluded_ports(
+    state: Mapping[str, Any],
+    avoid: set[int] | None = None,
+    target: tuple[str, str] | None = None,
+) -> set[int]:
+    excluded = set(avoid or ())
+    for srec in state.get("services", {}).values():
+        if isinstance(srec, dict) and srec.get("port"):
+            with contextlib.suppress(ValueError, TypeError):
+                excluded.add(int(srec["port"]))
+    if target is not None:
+        excluded.update(get_allocated_ports_for_others(target))
+    return excluded
+
+
+def _resolve_base_port(service: Any, state: Mapping[str, Any]) -> int:
+    pref = getattr(service, "preferred_port", None)
+    if pref is not None and PORT_MIN <= pref <= PORT_MAX:
+        return pref
+    project = state.get("project") or getattr(service, "project", None)
+    if project:
+        base = default_base_port_for_service(service.name)
+        return get_or_allocate_port((str(project), service.name), base)
+    leased = state.get("ports", {}).get(service.name)
+    if leased is not None and isinstance(leased, int) and PORT_MIN <= leased <= PORT_MAX:
+        return leased
+    return default_base_port_for_service(service.name)
+
+
+def compute_candidate_ports(
+    service: Any,
+    state: Mapping[str, Any],
+    avoid: set[int] | None = None,
+) -> list[int]:
+    """Compute an ordered list of candidate ports for a service."""
+    project = state.get("project") or getattr(service, "project", None)
+    sname = getattr(service, "name", None)
+    target = (str(project), sname) if project and sname else None
+    excluded = _collect_excluded_ports(state, avoid, target)
+    base_port = _resolve_base_port(service, state)
+    limit = min(base_port + DEFAULT_PORT_WINDOW, PORT_MAX + 1)
+    return [port for port in range(base_port, limit) if port not in excluded]

rig/net/probe.py

@@ -0,0 +1,56 @@
+"""Port listener inspection and process group querying."""
+
+from __future__ import annotations
+
+import os
+import shutil
+import subprocess
+from collections.abc import Sequence
+
+
+def run_query_command(cmd: Sequence[str], timeout: float = 1.0) -> str | None:
+    """Run a query subprocess safely, returning stdout if successful."""
+    try:
+        proc = subprocess.run(
+            cmd,
+            capture_output=True,
+            text=True,
+            check=False,
+            timeout=timeout,
+        )
+    except (subprocess.SubprocessError, OSError):
+        return None
+    if proc.returncode != 0 or not proc.stdout.strip():
+        return None
+    return proc.stdout.strip()
+
+
+def _matches_pgid(p: int, pgid: int) -> bool:
+    try:
+        return os.getpgid(p) == pgid
+    except (ProcessLookupError, PermissionError, OSError):
+        return False
+
+
+def _is_pid_owned(p: int, pid: int | None, pgid: int | None) -> bool:
+    if pid is not None and p == pid:
+        return True
+    return bool(pgid is not None and _matches_pgid(p, pgid))
+
+
+def _get_listener_pids(port: int) -> list[int]:
+    lsof_bin = shutil.which("lsof")
+    if not lsof_bin:
+        return []
+    output = run_query_command([lsof_bin, "-nP", f"-iTCP:{port}", "-sTCP:LISTEN", "-t"])
+    if not output:
+        return []
+    return [int(line.strip()) for line in output.splitlines() if line.strip().isdigit()]
+
+
+def port_listener_matches(port: int, pgid: int | None = None, pid: int | None = None) -> bool:
+    """Return ``True`` only when a listener on ``port`` belongs to ``pid`` or ``pgid``."""
+    if pid is None and pgid is None:
+        return False
+    pids = _get_listener_pids(port)
+    return bool(pids) and all(_is_pid_owned(p, pid, pgid) for p in pids)

rig/net/registry.py

@@ -0,0 +1,132 @@
+"""Centralized machine-wide port registry and reservation management."""
+
+from __future__ import annotations
+
+import json
+import os
+import socket
+import tempfile
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any
+
+from rig.core.constants import DIR_MODE_PRIVATE, FILE_MODE_PRIVATE, PORT_MAX, PORT_MIN
+from rig.core.identity import get_state_home
+from rig.core.locks import exclusive_lock
+
+PORT_REGISTRY_FILE = "ports.json"
+PORT_LOCK_FILE = "ports.lock"
+
+
+def port_is_free(port: int) -> bool:
+    """Return True when nothing is listening on port on loopback."""
+    probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    with probe:
+        probe.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        try:
+            probe.bind(("127.0.0.1", port))
+        except OSError:
+            return False
+    return True
+
+
+def _ensure_dir(filename: str) -> Path:
+    home = get_state_home()
+    home.mkdir(parents=True, mode=DIR_MODE_PRIVATE, exist_ok=True)
+    return home / filename
+
+
+def read_port_registry(path: Path | None = None) -> dict[str, Any]:
+    """Read central port registry, returning empty schema on error."""
+    target = path or _ensure_dir(PORT_REGISTRY_FILE)
+    try:
+        data = json.loads(target.read_text())
+        if isinstance(data, dict):
+            allocs = data.get("allocations")
+            data["allocations"] = allocs if isinstance(allocs, dict) else {}
+            return data
+    except (OSError, json.JSONDecodeError):
+        pass
+    return {"version": 1, "allocations": {}}
+
+
+def _write_port_registry(target: Path, data: Mapping[str, Any]) -> None:
+    target.parent.mkdir(parents=True, mode=DIR_MODE_PRIVATE, exist_ok=True)
+    with tempfile.NamedTemporaryFile(
+        mode="w", dir=target.parent, delete=False, prefix=".ports.tmp."
+    ) as handle:
+        json.dump(data, handle, indent=2, sort_keys=True)
+        handle.write("\n")
+        handle.flush()
+        os.fsync(handle.fileno())
+        tmp = handle.name
+    os.chmod(tmp, FILE_MODE_PRIVATE)
+    os.replace(tmp, target)
+
+
+def allocation_key(project: str, service_name: str) -> str:
+    return f"{project}:{service_name}"
+
+
+def _find_free_port(base_port: int, window: int, used_ports: set[int]) -> int:
+    limit = min(base_port + window, PORT_MAX + 1)
+    for cand in range(base_port, limit):
+        if cand not in used_ports and port_is_free(cand):
+            return cand
+    for cand in range(limit, PORT_MAX + 1):
+        if cand not in used_ports and port_is_free(cand):
+            return cand
+    return base_port
+
+
+def get_or_allocate_port(target: tuple[str, str], base_port: int, window: int = 50) -> int:
+    """Return assigned port from central registry, allocating next free if absent."""
+    project, service_name = target
+    key = allocation_key(project, service_name)
+    reg_path, lock_path = _ensure_dir(PORT_REGISTRY_FILE), _ensure_dir(PORT_LOCK_FILE)
+    with exclusive_lock(lock_path):
+        data = read_port_registry(reg_path)
+        allocs = data["allocations"]
+        if key in allocs:
+            assigned = allocs[key]
+            if isinstance(assigned, int) and PORT_MIN <= assigned <= PORT_MAX:
+                return assigned
+        used = {int(p) for p in allocs.values() if isinstance(p, int) and PORT_MIN <= p <= PORT_MAX}
+        port = _find_free_port(base_port, window, used)
+        allocs[key] = port
+        _write_port_registry(reg_path, data)
+        return port
+
+
+def get_allocated_ports_for_others(target: tuple[str, str]) -> set[int]:
+    """Return set of ports allocated to other projects or services."""
+    project, service_name = target
+    key = allocation_key(project, service_name)
+    data = read_port_registry(_ensure_dir(PORT_REGISTRY_FILE))
+    allocs = data.get("allocations", {})
+    return {
+        int(p)
+        for k, p in allocs.items()
+        if k != key and isinstance(p, int) and PORT_MIN <= p <= PORT_MAX
+    }
+
+
+def release_port_allocation(project: str, service_name: str | None = None) -> list[int]:
+    """Release port allocation for a project or specific service."""
+    reg_path, lock_path = _ensure_dir(PORT_REGISTRY_FILE), _ensure_dir(PORT_LOCK_FILE)
+    with exclusive_lock(lock_path):
+        data = read_port_registry(reg_path)
+        allocs = data.get("allocations", {})
+        prefix = f"{project}:" if service_name is None else f"{project}:{service_name}"
+        released = [
+            allocs.pop(k)
+            for k in list(allocs.keys())
+            if k == prefix or (service_name is None and k.startswith(prefix))
+        ]
+        if released:
+            _write_port_registry(reg_path, data)
+        return released
+
+
+def list_port_allocations() -> dict[str, int]:
+    return dict(read_port_registry(_ensure_dir(PORT_REGISTRY_FILE)).get("allocations", {}))

rig/parser.py

@@ -0,0 +1,72 @@
+"""CLI argument parser setup and error handling for rig."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from typing import Any
+
+from rig.core.constants import EXIT_USAGE, __version__
+from rig.core.errors import RigError, print_json_error
+
+
+class RigArgumentParser(argparse.ArgumentParser):
+    """ArgumentParser that outputs structured JSON errors when requested."""
+
+    def __init__(self, *args: Any, as_json: bool = False, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self.as_json = as_json
+
+    def error(self, message: str) -> None:
+        if self.as_json:
+            err = RigError(message, code="E_USAGE", exit_code=EXIT_USAGE)
+            print_json_error(err, command="cli")
+            sys.exit(EXIT_USAGE)
+        super().error(message)
+
+
+def _add_service_parsers(sub: Any) -> None:
+    p_up = sub.add_parser("up", help="start services")
+    p_up.add_argument("--scope", default="full")
+    p_up.add_argument("--mode", default=None)
+    p_up.add_argument("--switch", action="store_true")
+
+    p_down = sub.add_parser("down", help="stop services")
+    p_down.add_argument("target", nargs="?", default=None)
+    p_down.add_argument("--all", action="store_true", dest="all_instances")
+    p_down.add_argument("--scope", default="full")
+
+
+def _add_admin_parsers(sub: Any) -> None:
+    sub.add_parser("status", help="status")
+    p_ps = sub.add_parser("ps", aliases=["ls", "list"], help="ps")
+    p_ps.add_argument("--health", action="store_true")
+    p_ps.add_argument("-w", "--wide", action="store_true")
+    sub.add_parser("prune", help="prune").add_argument("--force", action="store_true")
+    sub.add_parser("check", help="check").add_argument("--mode", default=None)
+    p_init = sub.add_parser("init", help="init")
+    p_init.add_argument("--dry-run", action="store_true")
+    p_init.add_argument("--force", action="store_true")
+    p_init.add_argument("--up", action="store_true")
+    sub.add_parser("schema", help="schema")
+    p_logs = sub.add_parser("logs", help="logs")
+    p_logs.add_argument("service", nargs="?", default=None, help="service name")
+    p_logs.add_argument("-n", "--tail", type=int, default=50, help="number of lines")
+    p_logs.add_argument("--mode", default=None, help="mode overlay")
+
+
+def build_parser(as_json: bool = False) -> argparse.ArgumentParser:
+    """Construct the command line parser for rig."""
+    parser = RigArgumentParser(prog="rig", as_json=as_json)
+    parser.add_argument("-v", "--version", action="version", version=f"%(prog)s {__version__}")
+    parser.add_argument("--root", default=None)
+    parser.add_argument("--manifest", default=None)
+    parser.add_argument("--json", action="store_true")
+    sub = parser.add_subparsers(
+        dest="command",
+        required=True,
+        parser_class=lambda **kw: RigArgumentParser(as_json=as_json, **kw),
+    )
+    _add_service_parsers(sub)
+    _add_admin_parsers(sub)
+    return parser

rig/proc/__init__.py

@@ -0,0 +1 @@
+"""Process lifecycle, spawning, inspection, and teardown for rig."""

rig/proc/process.py

@@ -0,0 +1,133 @@
+"""Process identity inspection and liveness verification."""
+
+from __future__ import annotations
+
+import os
+import shutil
+import subprocess
+import time
+from collections.abc import Mapping
+from typing import Any
+
+_OWN_CHILDREN: dict[int, subprocess.Popen] = {}
+
+
+def _ps(pid: int, fields: str) -> str | None:
+    try:
+        completed = subprocess.run(
+            ["ps", "-ww", "-p", str(pid), "-o", fields],
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+    except (OSError, subprocess.SubprocessError):
+        return None
+    if completed.returncode != 0:
+        return None
+    line = completed.stdout.strip()
+    return line or None
+
+
+def process_start_time(pid: int) -> str | None:
+    """Return the kernel-reported start time of ``pid``."""
+    value = _ps(pid, "lstart=")
+    return " ".join(value.split()) if value else None
+
+
+def process_args(pid: int) -> str | None:
+    value = _ps(pid, "args=")
+    return " ".join(value.split()) if value else None
+
+
+def stable_process_args(pid: int, settle: float = 1.0, interval: float = 0.06) -> str | None:
+    """Return the command line once two consecutive reads agree."""
+    previous = process_args(pid)
+    deadline = time.monotonic() + max(settle, 0.0)
+    while time.monotonic() < deadline:
+        time.sleep(interval)
+        current = process_args(pid)
+        if current is None or current == previous:
+            return current if current is not None else previous
+        previous = current
+    return previous
+
+
+def _is_zombie(pid: int) -> bool:
+    state = _ps(pid, "state=")
+    return bool(state) and state.lstrip().upper().startswith("Z")
+
+
+def _collect(pid: int) -> bool:
+    """Reap a child this process started. Returns ``True`` when it has exited."""
+    child = _OWN_CHILDREN.get(pid)
+    if child is None or child.poll() is None:
+        return False
+    _OWN_CHILDREN.pop(pid, None)
+    return True
+
+
+def _is_live_signal(pid: int) -> bool:
+    try:
+        os.kill(pid, 0)
+    except ProcessLookupError:
+        return False
+    except PermissionError:
+        return True
+    else:
+        return True
+
+
+def pid_alive(pid: int) -> bool:
+    """Return ``True`` when ``pid`` names a live, non-zombie process."""
+    if _collect(pid) or not _is_live_signal(pid):
+        return False
+    if _is_zombie(pid):
+        _collect(pid)
+        return False
+    return True
+
+
+def resolve_binary(argv0: str) -> str:
+    """Return the canonical executable path for ``argv0``."""
+    if os.sep in argv0:
+        return os.path.realpath(argv0)
+    found = shutil.which(argv0)
+    return os.path.realpath(found) if found else argv0
+
+
+def _binary_matches(recorded: str, observed_argv0: str) -> bool:
+    if os.sep in observed_argv0:
+        return os.path.realpath(observed_argv0) == recorded
+    return os.path.basename(recorded).startswith(observed_argv0)
+
+
+def identity_baseline(record: Mapping[str, Any]) -> str:
+    """Return the command line that ``pid`` must still report to be considered ours."""
+    observed = record.get("identity")
+    if observed:
+        return str(observed)
+    argv = record.get("argv") or []
+    return " ".join(str(item) for item in argv)
+
+
+def _record_fields_valid(record: Mapping[str, Any], baseline: str) -> bool:
+    pid = record.get("pid")
+    return (
+        isinstance(pid, int)
+        and pid > 0
+        and bool(record.get("start_time") and baseline and record.get("binary"))
+    )
+
+
+def identity_matches(record: Mapping[str, Any]) -> bool:
+    """Return ``True`` only when ``pid`` still runs exactly the recorded program."""
+    baseline = identity_baseline(record)
+    if not _record_fields_valid(record, baseline):
+        return False
+    pid = int(record["pid"])
+    if not pid_alive(pid) or process_start_time(pid) != record.get("start_time"):
+        return False
+    observed = process_args(pid)
+    if not observed or observed != baseline:
+        return False
+    return _binary_matches(str(record.get("binary")), observed.split()[0])

rig/proc/record.py

@@ -0,0 +1,54 @@
+"""Process record construction and metadata formatting."""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import time
+from collections.abc import Mapping
+from typing import Any
+
+from rig.proc.process import process_start_time, resolve_binary, stable_process_args
+
+
+def build_record(
+    proc: subprocess.Popen,
+    name: str,
+    meta: tuple[str, list[str], int | None, Mapping[str, Any] | None],
+) -> dict[str, Any]:
+    """Construct structured process record dictionary."""
+    kind, resolved, port, extra = meta
+    observed = stable_process_args(proc.pid)
+    argv0 = observed.split()[0] if observed else str(resolved[0])
+    try:
+        pgid = os.getpgid(proc.pid)
+    except ProcessLookupError:
+        pgid = proc.pid
+    rec = {
+        "name": name,
+        "type": kind,
+        "pid": proc.pid,
+        "pgid": pgid,
+        "binary": resolve_binary(argv0),
+        "argv": list(resolved),
+        "identity": observed,
+        "start_time": process_start_time(proc.pid),
+        "port": port,
+        "url": f"http://127.0.0.1:{port}" if port else None,
+        "started_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
+    }
+    return {**rec, **(extra or {})}
+
+
+def _record(
+    name: str,
+    kind: str,
+    proc: subprocess.Popen,
+    *args: Any,
+    **kwargs: Any,
+) -> dict[str, Any]:
+    argv = list(args)
+    resolved = kwargs.get("argv") or (argv.pop(0) if argv else [])
+    port = kwargs.get("port") if "port" in kwargs else (argv.pop(0) if argv else None)
+    extra = kwargs.get("extra") if "extra" in kwargs else (argv.pop(0) if argv else None)
+    return build_record(proc, name, (kind, resolved, port, extra))