kanibako-cli 1.5.0.dev14__py3-none-any.whl

kanibako/helpers.py

@@ -0,0 +1,339 @@
+"""Helper spawning: B-ary tree numbering and spawn budget management."""
+
+from __future__ import annotations
+
+import importlib.resources
+from dataclasses import dataclass
+from pathlib import Path
+
+from kanibako.config_io import dump_doc, load_doc
+
+# When breadth is unlimited (-1), use 2^16 for numbering purposes.
+# Large enough to never collide; small enough for human-readable numbers.
+UNLIMITED_BREADTH = 2**16
+
+
+def effective_breadth(breadth: int) -> int:
+    """Return the breadth used for numbering.
+
+    Maps -1 (unlimited) to ``UNLIMITED_BREADTH``.  Positive values pass
+    through unchanged.
+    """
+    if breadth == -1:
+        return UNLIMITED_BREADTH
+    if breadth < 1:
+        msg = f"breadth must be positive or -1, got {breadth}"
+        raise ValueError(msg)
+    return breadth
+
+
+def children_of(agent: int, breadth: int) -> tuple[int, int]:
+    """Return the (first_child, last_child) global numbers for *agent*.
+
+    Both bounds are inclusive.  The range always contains exactly
+    ``effective_breadth(breadth)`` slots, regardless of how many children
+    are actually spawned.
+    """
+    b = effective_breadth(breadth)
+    first = agent * b + 1
+    last = agent * b + b
+    return first, last
+
+
+def parent_of(agent: int, breadth: int) -> int | None:
+    """Return the global number of *agent*'s parent.
+
+    Returns ``None`` if *agent* is the director (agent 0).
+    """
+    if agent == 0:
+        return None
+    b = effective_breadth(breadth)
+    return (agent - 1) // b
+
+
+def agent_depth(agent: int, breadth: int) -> int:
+    """Return the depth of *agent* in the tree (director = 0)."""
+    depth = 0
+    current = agent
+    while current != 0:
+        current = parent_of(current, breadth)  # type: ignore[assignment]
+        depth += 1
+    return depth
+
+
+def nth_child(agent: int, n: int, breadth: int) -> int:
+    """Return the global number of *agent*'s *n*-th child (0-indexed).
+
+    Raises ``ValueError`` if *n* is out of range for the given breadth.
+    """
+    b = effective_breadth(breadth)
+    if n < 0 or n >= b:
+        msg = f"child index {n} out of range for breadth {b}"
+        raise ValueError(msg)
+    return agent * b + 1 + n
+
+
+def sibling_index(agent: int, breadth: int) -> int:
+    """Return the 0-based index of *agent* among its parent's children.
+
+    The director (agent 0) has no siblings; returns 0 by convention.
+    """
+    if agent == 0:
+        return 0
+    b = effective_breadth(breadth)
+    return (agent - 1) % b
+
+
+# ---------------------------------------------------------------------------
+# Spawn budget
+# ---------------------------------------------------------------------------
+
+DEFAULT_DEPTH = 4
+DEFAULT_BREADTH = 4
+
+
+@dataclass(frozen=True)
+class SpawnBudget:
+    """Spawn limits for an agent.  Immutable."""
+
+    depth: int = DEFAULT_DEPTH
+    breadth: int = DEFAULT_BREADTH
+
+
+def check_spawn_allowed(budget: SpawnBudget, current_children: int) -> str | None:
+    """Return an error message if spawning is not allowed, else ``None``."""
+    if budget.depth == 0:
+        return "spawn depth exhausted (depth=0)"
+    if budget.breadth != -1 and current_children >= budget.breadth:
+        return f"breadth limit reached ({current_children}/{budget.breadth})"
+    return None
+
+
+def child_budget(parent: SpawnBudget) -> SpawnBudget:
+    """Compute the spawn budget for a child of *parent*.
+
+    Depth is decremented by 1 (unless unlimited).  Breadth is inherited.
+    """
+    new_depth = parent.depth if parent.depth == -1 else parent.depth - 1
+    return SpawnBudget(depth=new_depth, breadth=parent.breadth)
+
+
+def resolve_spawn_budget(
+    ro_config: SpawnBudget | None,
+    host_config: SpawnBudget | None,
+    cli_depth: int | None,
+    cli_breadth: int | None,
+) -> SpawnBudget:
+    """Resolve the effective spawn budget using config precedence.
+
+    Order: RO config > host config > CLI flags > built-in defaults.
+    CLI flags only apply when neither RO nor host config exist.
+    """
+    if ro_config is not None:
+        return ro_config
+    if host_config is not None:
+        return host_config
+    depth = cli_depth if cli_depth is not None else DEFAULT_DEPTH
+    breadth = cli_breadth if cli_breadth is not None else DEFAULT_BREADTH
+    return SpawnBudget(depth=depth, breadth=breadth)
+
+
+# ---------------------------------------------------------------------------
+# Spawn config I/O
+# ---------------------------------------------------------------------------
+
+
+def read_spawn_config(path: Path) -> SpawnBudget | None:
+    """Read spawn limits from a config file (kanibako.yaml or RO spawn config).
+
+    Looks for a ``spawn`` section with ``depth`` and ``breadth`` keys.
+    Returns ``None`` if the file or section is absent.
+    """
+    if not path.exists():
+        return None
+    data = load_doc(path)
+    spawn = data.get("spawn")
+    if spawn is None:
+        return None
+    return SpawnBudget(
+        depth=int(spawn.get("depth", DEFAULT_DEPTH)),
+        breadth=int(spawn.get("breadth", DEFAULT_BREADTH)),
+    )
+
+
+def write_spawn_config(path: Path, budget: SpawnBudget) -> None:
+    """Write spawn limits as a ``spawn`` section in a config file.
+
+    For RO spawn configs this creates a standalone file.
+    For kanibako.yaml this preserves other sections.
+    """
+    existing = load_doc(path)
+    existing["spawn"] = {"depth": budget.depth, "breadth": budget.breadth}
+    dump_doc(path, existing)
+
+
+# ---------------------------------------------------------------------------
+# Directory structure
+# ---------------------------------------------------------------------------
+
+
+def create_helper_dirs(helpers_dir: Path, helper_num: int) -> Path:
+    """Create the directory layout for a single helper.
+
+    Creates vault (with ro, rw), workspace, playbook/scripts,
+    and peers directories.  Returns the helper's root directory.
+    """
+    root = helpers_dir / str(helper_num)
+    root.mkdir(parents=True, exist_ok=True)
+
+    # Vault with communication channels
+    vault = root / "vault"
+    vault.mkdir(exist_ok=True)
+    (vault / "ro").mkdir(exist_ok=True)
+    (vault / "rw").mkdir(exist_ok=True)
+
+    # Standard layout
+    (root / "workspace").mkdir(exist_ok=True)
+    playbook = root / "playbook"
+    playbook.mkdir(exist_ok=True)
+    (playbook / "scripts").mkdir(exist_ok=True)
+
+    # Peers directory
+    (root / "peers").mkdir(exist_ok=True)
+
+    return root
+
+
+def create_broadcast_dirs(helpers_dir: Path) -> Path:
+    """Create the broadcast channel directories under ``helpers/``.
+
+    Creates ``all/rw`` and ``all/ro``.  Idempotent.
+    Returns the ``all/`` directory.
+    """
+    all_dir = helpers_dir / "all"
+    (all_dir / "rw").mkdir(parents=True, exist_ok=True)
+    (all_dir / "ro").mkdir(parents=True, exist_ok=True)
+    return all_dir
+
+
+def create_peer_channels(
+    helpers_dir: Path,
+    new_helper: int,
+    existing_helpers: list[int],
+) -> None:
+    """Create peer channels between *new_helper* and each existing sibling.
+
+    For each pair (A, B) where A < B, creates:
+    - ``A:B-ro`` directory (A writes, B reads)
+    - ``B:A-ro`` directory (B writes, A reads)
+    - ``A:B-rw`` directory (shared read-write, owned by lower number)
+
+    The directories are created under ``helpers_dir`` and symlinked into
+    each helper's ``peers/`` directory.
+    """
+    channels_dir = helpers_dir / "channels"
+    channels_dir.mkdir(exist_ok=True)
+
+    for existing in existing_helpers:
+        lower = min(new_helper, existing)
+        higher = max(new_helper, existing)
+
+        # Create the three channel directories
+        ro_low_high = channels_dir / f"{lower}:{higher}-ro"
+        ro_high_low = channels_dir / f"{higher}:{lower}-ro"
+        rw_shared = channels_dir / f"{lower}:{higher}-rw"
+
+        ro_low_high.mkdir(exist_ok=True)
+        ro_high_low.mkdir(exist_ok=True)
+        rw_shared.mkdir(exist_ok=True)
+
+        # Symlink into each helper's peers/
+        _link_peer(helpers_dir, lower, f"{lower}:{higher}-ro", ro_low_high)
+        _link_peer(helpers_dir, lower, f"{higher}:{lower}-ro", ro_high_low)
+        _link_peer(helpers_dir, lower, f"{lower}:{higher}-rw", rw_shared)
+
+        _link_peer(helpers_dir, higher, f"{lower}:{higher}-ro", ro_low_high)
+        _link_peer(helpers_dir, higher, f"{higher}:{lower}-ro", ro_high_low)
+        _link_peer(helpers_dir, higher, f"{lower}:{higher}-rw", rw_shared)
+
+
+def _link_peer(helpers_dir: Path, helper_num: int, name: str, target: Path) -> None:
+    """Create a symlink in helper's peers/ pointing to a channel directory."""
+    link = helpers_dir / str(helper_num) / "peers" / name
+    if not link.exists():
+        link.symlink_to(target.resolve())
+
+
+def link_broadcast(helpers_dir: Path, helper_num: int) -> None:
+    """Create an ``all`` symlink in a helper's filesystem pointing to broadcast dirs."""
+    all_dir = helpers_dir / "all"
+    link = helpers_dir / str(helper_num) / "all"
+    if not link.exists():
+        link.symlink_to(all_dir.resolve())
+
+
+def remove_helper_dirs(
+    helpers_dir: Path,
+    helper_num: int,
+    sibling_helpers: list[int],
+) -> None:
+    """Remove a helper's directory tree and clean up its peer channels.
+
+    Removes:
+    - The helper's root directory (``helpers/{N}/``)
+    - Channel directories involving this helper
+    - Peer symlinks in siblings that pointed to removed channels
+    """
+    import shutil
+
+    # Remove peer symlinks in siblings and channel dirs
+    channels_dir = helpers_dir / "channels"
+    for sibling in sibling_helpers:
+        lower = min(helper_num, sibling)
+        higher = max(helper_num, sibling)
+        channel_names = [
+            f"{lower}:{higher}-ro",
+            f"{higher}:{lower}-ro",
+            f"{lower}:{higher}-rw",
+        ]
+        # Remove symlinks from the sibling's peers/
+        for name in channel_names:
+            link = helpers_dir / str(sibling) / "peers" / name
+            if link.is_symlink():
+                link.unlink()
+        # Remove channel directories
+        for name in channel_names:
+            chan = channels_dir / name
+            if chan.exists():
+                shutil.rmtree(chan)
+
+    # Remove the helper's root directory
+    helper_root = helpers_dir / str(helper_num)
+    if helper_root.exists():
+        shutil.rmtree(helper_root)
+
+
+# ---------------------------------------------------------------------------
+# helper-init.sh template
+# ---------------------------------------------------------------------------
+
+_INIT_SCRIPT_NAME = "helper-init.sh"
+
+
+def bundled_init_script() -> Path:
+    """Return the path to the bundled default ``helper-init.sh``."""
+    resource = importlib.resources.files("kanibako.scripts").joinpath(_INIT_SCRIPT_NAME)
+    return Path(str(resource))
+
+
+def resolve_init_script(parent_scripts_dir: Path | None) -> Path:
+    """Return the init script to use for helpers.
+
+    Checks the parent's ``playbook/scripts/`` for a custom version first,
+    then falls back to the bundled default.
+    """
+    if parent_scripts_dir is not None:
+        custom = parent_scripts_dir / _INIT_SCRIPT_NAME
+        if custom.is_file():
+            return custom
+    return bundled_init_script()

kanibako/hygiene.py

@@ -0,0 +1,296 @@
+"""Shell directory cleanup: remove waste files, compress old conversation logs."""
+
+from __future__ import annotations
+
+import gzip
+import os
+import shutil
+import time
+from pathlib import Path
+
+from kanibako.log import get_logger
+
+# Directories whose *contents* are always safe to delete.
+_WASTE_DIRS = (
+    ".claude/telemetry",
+    ".claude/debug",
+)
+
+# Subdirectories under .cache/ that are safe to purge.
+# We intentionally keep pip, uv, npm, etc. — only remove known-waste dirs.
+_CACHE_WASTE_DIRS = (
+    ".cache/claude",
+    ".cache/sentry",
+    ".cache/@anthropic",
+)
+
+# Files older than this many days get compressed (conversation logs).
+_COMPRESS_AGE_DAYS = 7
+
+
+def cleanup_shell_dir(
+    shell_dir: Path,
+    dry_run: bool = False,
+) -> list[str]:
+    """Remove stale/waste files from a persistent shell directory.
+
+    Returns a list of human-readable action strings describing what was
+    (or would be, in dry_run mode) cleaned up.  The list is empty when
+    there is nothing to do.
+    """
+    logger = get_logger("hygiene")
+    actions: list[str] = []
+
+    if not shell_dir.is_dir():
+        return actions
+
+    # 1. Delete known waste directories.
+    actions.extend(_clean_waste_dirs(shell_dir, dry_run, logger))
+
+    # 2. Delete .cache waste subdirectories.
+    actions.extend(_clean_cache_waste(shell_dir, dry_run, logger))
+
+    # 3. Remove duplicate claude binaries outside .local/.
+    actions.extend(_clean_duplicate_binaries(shell_dir, dry_run, logger))
+
+    # 4. Compress old conversation logs.
+    actions.extend(_compress_old_logs(shell_dir, dry_run, logger))
+
+    if actions:
+        total = len(actions)
+        prefix = "[dry-run] " if dry_run else ""
+        logger.info("%sHygiene: %d action(s) taken", prefix, total)
+
+    return actions
+
+
+def _clean_waste_dirs(
+    shell_dir: Path,
+    dry_run: bool,
+    logger: object,
+) -> list[str]:
+    """Delete contents of known waste directories."""
+    actions: list[str] = []
+    for rel in _WASTE_DIRS:
+        target = shell_dir / rel
+        if not target.is_dir():
+            continue
+        freed = _dir_size(target)
+        if freed == 0:
+            continue
+        desc = (
+            f"{'[dry-run] ' if dry_run else ''}"
+            f"Removed {rel}/ contents ({_fmt_size(freed)})"
+        )
+        if not dry_run:
+            _remove_dir_contents(target)
+        actions.append(desc)
+    return actions
+
+
+def _clean_cache_waste(
+    shell_dir: Path,
+    dry_run: bool,
+    logger: object,
+) -> list[str]:
+    """Delete waste subdirectories under .cache/."""
+    actions: list[str] = []
+    for rel in _CACHE_WASTE_DIRS:
+        target = shell_dir / rel
+        if not target.is_dir():
+            continue
+        freed = _dir_size(target)
+        if freed == 0:
+            continue
+        desc = (
+            f"{'[dry-run] ' if dry_run else ''}"
+            f"Removed {rel}/ ({_fmt_size(freed)})"
+        )
+        if not dry_run:
+            shutil.rmtree(target, ignore_errors=True)
+        actions.append(desc)
+    return actions
+
+
+def _clean_duplicate_binaries(
+    shell_dir: Path,
+    dry_run: bool,
+    logger: object,
+) -> list[str]:
+    """Remove claude binary copies outside .local/.
+
+    The legitimate binary lives at .local/bin/claude (bind-mounted from
+    host).  Anything else that looks like a large claude binary is waste
+    — typically 200+ MB copies left by install scripts or updates.
+    """
+    actions: list[str] = []
+    # Minimum size to consider: 100 MB (real binary is ~227 MB).
+    min_size = 100 * 1024 * 1024
+
+    for candidate in _find_claude_binaries(shell_dir):
+        # Skip the legitimate location.
+        try:
+            rel = candidate.relative_to(shell_dir / ".local")
+            # It's under .local — leave it alone.
+            _ = rel
+            continue
+        except ValueError:
+            pass
+
+        try:
+            size = candidate.stat().st_size
+        except OSError:
+            continue
+
+        if size < min_size:
+            continue
+
+        rel_path = candidate.relative_to(shell_dir)
+        desc = (
+            f"{'[dry-run] ' if dry_run else ''}"
+            f"Removed duplicate binary {rel_path} ({_fmt_size(size)})"
+        )
+        if not dry_run:
+            try:
+                candidate.unlink()
+            except OSError:
+                continue
+        actions.append(desc)
+
+    return actions
+
+
+def _find_claude_binaries(shell_dir: Path) -> list[Path]:
+    """Find files named 'claude' that look like binaries in shell_dir.
+
+    Only walks directories that are likely to contain stray copies:
+    the top level and a few common subdirectories.  Does NOT recurse
+    the entire tree (that would be too slow).
+    """
+    candidates: list[Path] = []
+
+    # Check top-level and common locations for stray binaries.
+    search_dirs = [
+        shell_dir,
+        shell_dir / ".claude" / "bin",
+        shell_dir / "bin",
+        shell_dir / ".bin",
+        shell_dir / ".npm" / "_npx",
+    ]
+    for d in search_dirs:
+        if not d.is_dir():
+            continue
+        claude_file = d / "claude"
+        if claude_file.is_file() and not claude_file.is_symlink():
+            candidates.append(claude_file)
+
+    return candidates
+
+
+def _compress_old_logs(
+    shell_dir: Path,
+    dry_run: bool,
+    logger: object,
+) -> list[str]:
+    """Gzip conversation logs older than _COMPRESS_AGE_DAYS.
+
+    Looks for .jsonl files under .claude/projects/*/conversation_logs/.
+    """
+    actions: list[str] = []
+    cutoff = time.time() - (_COMPRESS_AGE_DAYS * 86400)
+
+    projects_dir = shell_dir / ".claude" / "projects"
+    if not projects_dir.is_dir():
+        return actions
+
+    # Glob for conversation log files.
+    for log_file in projects_dir.glob("*/conversation_logs/*.jsonl"):
+        if not log_file.is_file():
+            continue
+        # Already compressed?
+        if log_file.suffix == ".gz":
+            continue
+
+        try:
+            mtime = log_file.stat().st_mtime
+        except OSError:
+            continue
+
+        if mtime >= cutoff:
+            continue
+
+        original_size = log_file.stat().st_size
+        if original_size == 0:
+            continue
+
+        rel_path = log_file.relative_to(shell_dir)
+        gz_path = log_file.with_suffix(log_file.suffix + ".gz")
+
+        if not dry_run:
+            try:
+                _gzip_file(log_file, gz_path)
+                compressed_size = gz_path.stat().st_size
+            except OSError:
+                continue
+        else:
+            compressed_size = original_size  # estimate unavailable in dry-run
+
+        desc = (
+            f"{'[dry-run] ' if dry_run else ''}"
+            f"Compressed {rel_path} ({_fmt_size(original_size)}"
+        )
+        if not dry_run:
+            desc += f" -> {_fmt_size(compressed_size)}"
+        desc += ")"
+        actions.append(desc)
+
+    return actions
+
+
+def _gzip_file(src: Path, dst: Path) -> None:
+    """Compress *src* to *dst* with gzip and remove the original."""
+    with open(src, "rb") as f_in, gzip.open(dst, "wb") as f_out:
+        shutil.copyfileobj(f_in, f_out)
+    # Preserve modification time on the compressed file.
+    stat = src.stat()
+    os.utime(dst, (stat.st_atime, stat.st_mtime))
+    src.unlink()
+
+
+def _remove_dir_contents(d: Path) -> None:
+    """Remove all entries inside *d* without removing *d* itself."""
+    for entry in d.iterdir():
+        if entry.is_dir() and not entry.is_symlink():
+            shutil.rmtree(entry, ignore_errors=True)
+        else:
+            try:
+                entry.unlink()
+            except OSError:
+                pass
+
+
+def _dir_size(d: Path) -> int:
+    """Return total size of all files under *d* (non-recursive symlink-safe)."""
+    total = 0
+    try:
+        for entry in d.rglob("*"):
+            if entry.is_file() and not entry.is_symlink():
+                try:
+                    total += entry.stat().st_size
+                except OSError:
+                    pass
+    except OSError:
+        pass
+    return total
+
+
+def _fmt_size(nbytes: int) -> str:
+    """Format byte count as a human-readable string."""
+    if nbytes < 1024:
+        return f"{nbytes} B"
+    elif nbytes < 1024 * 1024:
+        return f"{nbytes / 1024:.1f} KB"
+    elif nbytes < 1024 * 1024 * 1024:
+        return f"{nbytes / (1024 * 1024):.1f} MB"
+    else:
+        return f"{nbytes / (1024 * 1024 * 1024):.1f} GB"