attocode-code-intel 0.1.0__py3-none-any.whl

attocode_intel/tools/overlay_tools.py

@@ -0,0 +1,540 @@
+"""Named-overlay MCP tools — directory-based working-state swap.
+
+An overlay is a named copy of every local store file, stored under
+``.attocode/overlays/<name>/``. Switching between branches / working sets
+means calling ``overlay_activate(name)`` — an atomic per-file replacement
+that preserves the current state as a fresh overlay if you ask it to.
+
+Phase 2c-3 ships the minimal viable form: directory-based, no stacking,
+no tombstones, no dedup across overlays. It's the simplest thing that
+delivers the UX goal ("one command to swap code-intel state") and
+forward-compatible with a later stacked-layer design.
+
+Format of one overlay directory::
+
+    .attocode/overlays/<name>/
+        overlay.json                       # {created_at, description, source_hash}
+        stores/symbols.db                  # live-consistent SQLite backup
+        stores/embeddings.db
+        stores/memory.db
+        stores/adrs.db
+        stores/frecency.db
+        stores/query_history.db
+        stores/kw_index.db                 # copy
+        stores/trigrams.{lookup,postings,db}
+        cache_manifest.json                # project manifest at capture time
+
+Tools:
+
+  - ``overlay_create(name, description)`` — capture current state
+  - ``overlay_list``                      — enumerate all overlays
+  - ``overlay_activate(name, save_current_as)`` — swap over live state
+  - ``overlay_delete(name, confirm)``     — remove an overlay directory
+  - ``overlay_status``                    — currently-active overlay tag
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import shutil
+import sqlite3
+import time
+from typing import Any
+
+from attocode_intel._shared import _get_project_dir, mcp
+
+logger = logging.getLogger(__name__)
+
+OVERLAY_SCHEMA_VERSION = 1
+OVERLAY_DIR_NAME = "overlays"
+_OVERLAY_STATE_FILE = "_overlay_state.json"
+
+# Same shape as snapshot_tools' _SQLITE_STORES / _BINARY_FILES so the two
+# tool modules stay in lockstep. When a new store lands, update both lists.
+_SQLITE_STORES: tuple[tuple[str, str], ...] = (
+    ("index/symbols.db", "symbols.db"),
+    ("vectors/embeddings.db", "embeddings.db"),
+    ("index/kw_index.db", "kw_index.db"),
+    ("cache/memory.db", "memory.db"),
+    ("adrs.db", "adrs.db"),
+    ("frecency/frecency.db", "frecency.db"),
+    ("query_history/query_history.db", "query_history.db"),
+)
+_BINARY_FILES: tuple[tuple[str, str], ...] = (
+    ("index/trigrams.lookup", "trigrams.lookup"),
+    ("index/trigrams.postings", "trigrams.postings"),
+    ("index/trigrams.db", "trigrams.db"),
+)
+
+
+# ---------------------------------------------------------------------------
+# Paths + state helpers
+# ---------------------------------------------------------------------------
+
+
+def _attocode_dir(project_dir: str) -> str:
+    return os.path.join(project_dir, ".attocode")
+
+
+def _overlays_root(project_dir: str) -> str:
+    return os.path.join(_attocode_dir(project_dir), OVERLAY_DIR_NAME)
+
+
+def _overlay_dir(project_dir: str, name: str) -> str:
+    return os.path.join(_overlays_root(project_dir), name)
+
+
+def _state_path(project_dir: str) -> str:
+    return os.path.join(_overlays_root(project_dir), _OVERLAY_STATE_FILE)
+
+
+def _read_state(project_dir: str) -> dict[str, Any]:
+    path = _state_path(project_dir)
+    if not os.path.exists(path):
+        return {"active": None, "history": []}
+    try:
+        with open(path, encoding="utf-8") as f:
+            return json.loads(f.read())
+    except (OSError, json.JSONDecodeError):
+        return {"active": None, "history": []}
+
+
+def _write_state(project_dir: str, state: dict[str, Any]) -> None:
+    path = _state_path(project_dir)
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    tmp = path + ".tmp"
+    with open(tmp, "w", encoding="utf-8") as f:
+        json.dump(state, f, indent=2, sort_keys=True)
+    os.replace(tmp, path)
+
+
+def _is_valid_name(name: str) -> bool:
+    """Overlay names must be short filesystem-friendly slugs.
+
+    Rejected: empty strings, path separators, leading underscore
+    (reserved for internal state like ``_overlay_state.json``), and
+    anything with shell meta characters.
+    """
+    if not name or name.startswith("_") or name in (".", ".."):
+        return False
+    return all(c.isalnum() or c in ("-", "_", ".") for c in name)
+
+
+# ---------------------------------------------------------------------------
+# Store copy helpers (shared with snapshot_tools logic, inlined so the
+# two modules can evolve independently)
+# ---------------------------------------------------------------------------
+
+
+def _backup_sqlite(source: str, dest: str) -> None:
+    """Live-consistent SQLite backup via the online backup API."""
+    src_conn = sqlite3.connect(f"file:{source}?mode=ro", uri=True)
+    try:
+        dest_conn = sqlite3.connect(dest)
+        try:
+            src_conn.backup(dest_conn)
+        finally:
+            dest_conn.close()
+    finally:
+        src_conn.close()
+
+
+def _copy_store_to(dest_stores_dir: str, attocode: str) -> list[str]:
+    """Copy every tracked store into ``dest_stores_dir``.
+
+    Returns the list of copied store names (logical, not filenames).
+    """
+    os.makedirs(dest_stores_dir, exist_ok=True)
+    copied: list[str] = []
+
+    for src_rel, out_name in _SQLITE_STORES:
+        src = os.path.join(attocode, src_rel)
+        if not os.path.exists(src):
+            continue
+        dest = os.path.join(dest_stores_dir, out_name)
+        try:
+            _backup_sqlite(src, dest)
+        except sqlite3.Error as exc:
+            logger.warning(
+                "overlay: sqlite backup of %s failed (%s); falling back to shutil.copy",
+                src, exc,
+            )
+            shutil.copy2(src, dest)
+        copied.append(out_name)
+
+    for src_rel, out_name in _BINARY_FILES:
+        src = os.path.join(attocode, src_rel)
+        if not os.path.exists(src):
+            continue
+        shutil.copy2(src, os.path.join(dest_stores_dir, out_name))
+        copied.append(out_name)
+
+    return copied
+
+
+def _copy_store_from(src_stores_dir: str, attocode: str) -> list[str]:
+    """Copy every overlay store file back into the live ``.attocode/`` tree.
+
+    Each destination is atomically replaced via rename-tmp → rename so a
+    failure mid-copy can't produce a torn file. The overall restore is
+    atomic per-file, not per-collection — which is good enough for
+    Phase 2c-3 (single-user local).
+    """
+    restored: list[str] = []
+    mapping: list[tuple[str, str]] = list(_SQLITE_STORES) + list(_BINARY_FILES)
+    for src_rel, out_name in mapping:
+        src = os.path.join(src_stores_dir, out_name)
+        if not os.path.exists(src):
+            continue
+        dest = os.path.join(attocode, src_rel)
+        os.makedirs(os.path.dirname(dest), exist_ok=True)
+        tmp = dest + ".tmp"
+        shutil.copy2(src, tmp)
+        os.replace(tmp, dest)
+        restored.append(out_name)
+    return restored
+
+
+def _count_rows(db_path: str) -> int:
+    if not os.path.exists(db_path):
+        return 0
+    try:
+        conn = sqlite3.connect(f"file:{db_path}?mode=ro", uri=True)
+    except sqlite3.Error:
+        return 0
+    try:
+        # Sum across every user table so a newly-added store shows up too.
+        tables = [
+            r[0] for r in conn.execute(
+                "SELECT name FROM sqlite_master "
+                "WHERE type='table' AND name NOT LIKE 'sqlite_%'"
+            ).fetchall()
+        ]
+        total = 0
+        for t in tables:
+            try:
+                row = conn.execute(f"SELECT COUNT(*) FROM {t}").fetchone()
+                if row:
+                    total += int(row[0])
+            except sqlite3.OperationalError:
+                pass
+        return total
+    finally:
+        conn.close()
+
+
+def _safe_size(path: str) -> int:
+    try:
+        return os.path.getsize(path)
+    except OSError:
+        return 0
+
+
+def _fmt_bytes(n: int) -> str:
+    for unit in ("B", "KB", "MB", "GB"):
+        if n < 1024:
+            return f"{n}{unit}" if unit == "B" else f"{n:.1f}{unit}"
+        n //= 1024
+    return f"{n}TB"
+
+
+# ---------------------------------------------------------------------------
+# MCP tools
+# ---------------------------------------------------------------------------
+
+
+@mcp.tool()
+def overlay_create(name: str, description: str = "") -> str:
+    """Capture the current code-intel state as a named overlay.
+
+    Args:
+        name: Short filesystem-safe overlay name (e.g. ``"main"``,
+            ``"feature-auth"``). Must be alphanumeric plus ``- _ .``.
+        description: Optional human-readable description stored alongside
+            the overlay for reference.
+
+    Returns a summary of what was captured.
+    """
+    project_dir = _get_project_dir()
+    if not _is_valid_name(name):
+        return (
+            f"overlay_create: invalid overlay name {name!r}. "
+            f"Use alphanumeric plus '-', '_', '.'. Leading underscores reserved."
+        )
+    odir = _overlay_dir(project_dir, name)
+    if os.path.exists(odir):
+        return (
+            f"overlay_create: overlay {name!r} already exists at {odir}. "
+            f"Delete it first with overlay_delete, or pick a different name."
+        )
+
+    attocode = _attocode_dir(project_dir)
+    if not os.path.exists(attocode):
+        return "overlay_create: no .attocode/ directory to capture."
+
+    stores_dir = os.path.join(odir, "stores")
+    os.makedirs(stores_dir, exist_ok=True)
+    try:
+        copied = _copy_store_to(stores_dir, attocode)
+    except Exception as exc:  # noqa: BLE001
+        shutil.rmtree(odir, ignore_errors=True)
+        return f"overlay_create: copy failed, rolled back: {exc}"
+
+    if not copied:
+        shutil.rmtree(odir, ignore_errors=True)
+        return "overlay_create: nothing to capture (no store files present)."
+
+    # Copy cache_manifest.json if present (top-level — not in stores/).
+    manifest_src = os.path.join(attocode, "cache_manifest.json")
+    if os.path.exists(manifest_src):
+        shutil.copy2(manifest_src, os.path.join(odir, "cache_manifest.json"))
+
+    total_size = sum(
+        _safe_size(os.path.join(stores_dir, n)) for n in copied
+    )
+    metadata = {
+        "schema_version": OVERLAY_SCHEMA_VERSION,
+        "name": name,
+        "description": description,
+        "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+        "stores": copied,
+        "total_bytes": total_size,
+    }
+    with open(os.path.join(odir, "overlay.json"), "w", encoding="utf-8") as f:
+        json.dump(metadata, f, indent=2, sort_keys=True)
+
+    return (
+        f"overlay_create: captured {len(copied)} store(s) "
+        f"({_fmt_bytes(total_size)}) as {name!r} at {odir}"
+    )
+
+
+@mcp.tool()
+def overlay_list() -> str:
+    """List all named overlays under ``.attocode/overlays/``."""
+    project_dir = _get_project_dir()
+    root = _overlays_root(project_dir)
+    if not os.path.isdir(root):
+        return "overlay_list: no overlays/ directory."
+
+    state = _read_state(project_dir)
+    active = state.get("active")
+
+    entries: list[tuple[str, dict[str, Any]]] = []
+    for name in sorted(os.listdir(root)):
+        if name.startswith("_") or not _is_valid_name(name):
+            continue
+        odir = os.path.join(root, name)
+        if not os.path.isdir(odir):
+            continue
+        meta_path = os.path.join(odir, "overlay.json")
+        if os.path.exists(meta_path):
+            try:
+                with open(meta_path, encoding="utf-8") as f:
+                    meta = json.loads(f.read())
+            except (OSError, json.JSONDecodeError):
+                meta = {}
+        else:
+            meta = {}
+        entries.append((name, meta))
+
+    if not entries:
+        return "overlay_list: no named overlays found."
+
+    lines = [f"overlay_list: {len(entries)} overlay(s) (active={active or 'none'}):"]
+    for name, meta in entries:
+        marker = "*" if name == active else " "
+        created = meta.get("created_at", "?")
+        size = meta.get("total_bytes", 0)
+        desc = meta.get("description", "")
+        suffix = f"  {desc!r}" if desc else ""
+        lines.append(
+            f"  {marker} {name:20s}  "
+            f"{_fmt_bytes(size):>10s}  {created}{suffix}"
+        )
+    return "\n".join(lines)
+
+
+@mcp.tool()
+def overlay_status() -> str:
+    """Report which overlay (if any) is currently marked active.
+
+    The live code-intel state is always whatever is in ``.attocode/``;
+    the "active" marker just records the *name* of the overlay that
+    was last activated. Creating an overlay does not change the active
+    marker — activation does.
+    """
+    project_dir = _get_project_dir()
+    state = _read_state(project_dir)
+    active = state.get("active")
+    history = state.get("history", [])
+
+    if not active:
+        return (
+            "overlay_status: no overlay activated yet. "
+            "Create one with overlay_create(name)."
+        )
+
+    lines = [f"overlay_status: active={active!r}"]
+    if history:
+        lines.append(f"  recent: {', '.join(history[-5:])}")
+    return "\n".join(lines)
+
+
+@mcp.tool()
+def overlay_activate(
+    name: str,
+    save_current_as: str = "",
+    confirm: bool = False,
+) -> str:
+    """Swap the live ``.attocode/`` state with an overlay.
+
+    This copies every store file from ``.attocode/overlays/<name>/``
+    into the live ``.attocode/`` tree, overwriting anything in place.
+    Per-file atomic via tempfile + rename. If ``save_current_as`` is
+    provided, the current live state is captured as a new overlay
+    BEFORE the swap so you can return to it later.
+
+    Args:
+        name: Existing overlay name to restore.
+        save_current_as: Optional name for a pre-swap backup overlay.
+            Empty means "don't back up — just swap". Skipping this
+            silently loses the pre-swap state if no prior overlay has it.
+        confirm: Must be True to apply. Default False returns a dry run.
+    """
+    project_dir = _get_project_dir()
+    if not _is_valid_name(name):
+        return f"overlay_activate: invalid overlay name {name!r}"
+
+    odir = _overlay_dir(project_dir, name)
+    if not os.path.isdir(odir):
+        return f"overlay_activate: no overlay named {name!r}"
+
+    stores_src = os.path.join(odir, "stores")
+    if not os.path.isdir(stores_src):
+        return f"overlay_activate: overlay {name!r} missing stores/ directory"
+
+    attocode = _attocode_dir(project_dir)
+    state = _read_state(project_dir)
+
+    if save_current_as and not _is_valid_name(save_current_as):
+        return (
+            f"overlay_activate: invalid save_current_as name "
+            f"{save_current_as!r}"
+        )
+
+    # Preview the component list so the dry-run summary is useful.
+    components = sorted(os.listdir(stores_src)) if os.path.isdir(stores_src) else []
+
+    if not confirm:
+        lines = [
+            "overlay_activate: DRY RUN — nothing swapped.",
+            f"  source: {odir}",
+            f"  components: {len(components)}",
+        ]
+        if save_current_as:
+            lines.append(f"  will save current live state as: {save_current_as!r}")
+        else:
+            lines.append("  will NOT save current live state (pass save_current_as to keep it)")
+        lines.append("  Re-run with confirm=True to apply.")
+        return "\n".join(lines)
+
+    # 1. Save current state if requested.
+    if save_current_as:
+        save_dir = _overlay_dir(project_dir, save_current_as)
+        if os.path.exists(save_dir):
+            return (
+                f"overlay_activate: save_current_as={save_current_as!r} already "
+                f"exists. Delete it first or pick another name."
+            )
+        try:
+            os.makedirs(os.path.join(save_dir, "stores"), exist_ok=True)
+            saved = _copy_store_to(os.path.join(save_dir, "stores"), attocode)
+            manifest_src = os.path.join(attocode, "cache_manifest.json")
+            if os.path.exists(manifest_src):
+                shutil.copy2(manifest_src, os.path.join(save_dir, "cache_manifest.json"))
+            metadata = {
+                "schema_version": OVERLAY_SCHEMA_VERSION,
+                "name": save_current_as,
+                "description": f"auto-saved before activating {name!r}",
+                "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+                "stores": saved,
+                "total_bytes": sum(
+                    _safe_size(os.path.join(save_dir, "stores", n)) for n in saved
+                ),
+            }
+            with open(os.path.join(save_dir, "overlay.json"), "w", encoding="utf-8") as f:
+                json.dump(metadata, f, indent=2, sort_keys=True)
+        except Exception as exc:  # noqa: BLE001
+            shutil.rmtree(save_dir, ignore_errors=True)
+            return f"overlay_activate: failed to save current state: {exc}"
+
+    # 2. Restore from target overlay.
+    try:
+        restored = _copy_store_from(stores_src, attocode)
+    except Exception as exc:  # noqa: BLE001
+        return f"overlay_activate: restore failed: {exc}"
+
+    # 3. Copy cache_manifest.json if the overlay has one.
+    manifest_src = os.path.join(odir, "cache_manifest.json")
+    if os.path.exists(manifest_src):
+        shutil.copy2(manifest_src, os.path.join(attocode, "cache_manifest.json"))
+
+    # 4. Update state.
+    history = state.get("history", [])
+    prior_active = state.get("active")
+    if prior_active and prior_active != name:
+        history = [*history, prior_active][-10:]
+    state["active"] = name
+    state["history"] = history
+    _write_state(project_dir, state)
+
+    lines = [
+        f"overlay_activate: activated {name!r}",
+        f"  restored: {len(restored)} component(s)",
+    ]
+    if save_current_as:
+        lines.append(f"  saved prior state as: {save_current_as!r}")
+    return "\n".join(lines)
+
+
+@mcp.tool()
+def overlay_delete(name: str, confirm: bool = False) -> str:
+    """Remove a named overlay directory.
+
+    Args:
+        name: Overlay to delete.
+        confirm: Must be True to apply.
+    """
+    project_dir = _get_project_dir()
+    if not _is_valid_name(name):
+        return f"overlay_delete: invalid overlay name {name!r}"
+    odir = _overlay_dir(project_dir, name)
+    if not os.path.isdir(odir):
+        return f"overlay_delete: no overlay named {name!r}"
+
+    # Size preview.
+    total = 0
+    for root, _, files in os.walk(odir):
+        for f in files:
+            total += _safe_size(os.path.join(root, f))
+
+    if not confirm:
+        return (
+            f"overlay_delete: DRY RUN — would delete {odir} "
+            f"({_fmt_bytes(total)}). Re-run with confirm=True."
+        )
+
+    shutil.rmtree(odir, ignore_errors=False)
+
+    # If this was the active overlay, clear the state marker (the live
+    # .attocode/ tree is unaffected, but we no longer have a named
+    # reference to the state it matches).
+    state = _read_state(project_dir)
+    if state.get("active") == name:
+        state["active"] = None
+        _write_state(project_dir, state)
+
+    return f"overlay_delete: removed {name!r} ({_fmt_bytes(total)} freed)"