agent-brain-cli 10.3.2__tar.gz → 10.4.0__tar.gz

{agent_brain_cli-10.3.2 → agent_brain_cli-10.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: agent-brain-cli
-Version: 10.3.2
+Version: 10.4.0
 Summary: Agent Brain CLI - Command-line interface for managing AI agent memory and knowledge retrieval
 Home-page: https://github.com/SpillwaveSolutions/agent-brain
 License: MIT
@@ -15,7 +15,7 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Requires-Dist: agent-brain-rag (>=10.3.2,<11.0.0)
+Requires-Dist: agent-brain-rag (>=10.4.0,<11.0.0)
 Requires-Dist: agent-brain-uds (>=10.3.0,<11.0.0)
 Requires-Dist: click (>=8.1.0,<9.0.0)
 Requires-Dist: httpx (>=0.28.0,<0.29.0)

{agent_brain_cli-10.3.2 → agent_brain_cli-10.4.0}/agent_brain_cli/__init__.py

@@ -1,3 +1,3 @@
 """Doc-Serve CLI - Command-line interface for managing Doc-Serve server."""
 
-__version__ = "10.3.2"
+__version__ = "10.4.0"

{agent_brain_cli-10.3.2 → agent_brain_cli-10.4.0}/agent_brain_cli/cli.py

@@ -12,6 +12,7 @@ from .commands import (
     config_group,
     doctor_command,
     folders_group,
+    graph_group,
     index_command,
     init_command,
     inject_command,
@@ -164,6 +165,9 @@ cli.add_command(mcp_group, name="mcp")
 cli.add_command(prompt_command, name="prompt")
 cli.add_command(resources_group, name="resources")
 
+# Register graph management commands
+cli.add_command(graph_group, name="graph")
+
 # Register server interaction commands
 cli.add_command(status_command, name="status")
 cli.add_command(query_command, name="query")

{agent_brain_cli-10.3.2 → agent_brain_cli-10.4.0}/agent_brain_cli/client/transport.py

@@ -186,7 +186,14 @@ def open_backend(ctx: click.Context, *, timeout: float = 30.0) -> BackendClient:
         # mcp_transport == "http" — resolve_mcp_transport guarantees
         # mcp_target is not None for the http branch.
         assert mcp_target is not None  # noqa: S101
-        return cast(BackendClient, McpHttpBackend(url=mcp_target, timeout=timeout))
+        # Thread state_dir so FileTokenStorage is keyed correctly when
+        # AGENT_BRAIN_MCP_AUTH=oauth is set (Phase 69 Plan 03).
+        # The opt-in flows via the AGENT_BRAIN_MCP_AUTH env var — no extra
+        # param needed here.
+        return cast(
+            BackendClient,
+            McpHttpBackend(url=mcp_target, timeout=timeout, state_dir=state_dir),
+        )
 
     # --- HTTP / UDS branch (existing v1/v2 path) -----------------
     transport, target = resolve_transport(

{agent_brain_cli-10.3.2 → agent_brain_cli-10.4.0}/agent_brain_cli/commands/__init__.py

@@ -4,6 +4,7 @@ from .cache import cache_group
 from .config import config_group
 from .doctor import doctor_command
 from .folders import folders_group
+from .graph import graph_group
 from .index import index_command
 from .init import init_command
 from .inject import inject_command
@@ -26,6 +27,7 @@ __all__ = [
     "config_group",
     "doctor_command",
     "folders_group",
+    "graph_group",
     "index_command",
     "inject_command",
     "init_command",

agent_brain_cli-10.4.0/agent_brain_cli/commands/graph.py

@@ -0,0 +1,154 @@
+"""`agent-brain graph` Click sub-group for graph management commands (Phase 64).
+
+Subcommands:
+    restore-from-snapshot — replay the latest (or a specific) kuzu snapshot
+        back into the live graph. Confirm-by-default; --yes skips the prompt
+        for non-interactive/CI use; --dry-run reports the plan without mutating.
+
+Closes #184 bug 1: when kuzu opens cleanly but the live graph is STALE after
+an AGENT_BRAIN_JOB_TIMEOUT rollback, operators can replay the latest valid
+snapshot to bring the graph back in sync.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import click
+from rich.console import Console
+
+from agent_brain_cli.config import STATE_DIR_NAME, resolve_project_root
+
+# Re-use diagnostics helpers for state-dir + store-type resolution and the
+# server-running guard. These are local imports inside the CLI so the module
+# loads without requiring the server package to be installed.
+from agent_brain_cli.diagnostics import (
+    _graph_index_dir,
+    _read_graphrag_block,
+    _server_is_running,
+)
+
+# agent-brain-rag is a declared dependency of agent-brain-cli (^10.x), so this
+# import is safe at module level. The lazy try/except below is a belt-and-braces
+# guard for edge cases where the server wheel is not installed in the env.
+try:
+    from agent_brain_server.storage.graph_store import GraphStoreManager
+except ImportError:  # pragma: no cover
+    GraphStoreManager = None
+
+console = Console()
+
+
+@click.group(name="graph")
+def graph_group() -> None:
+    """Manage the GraphRAG graph index (kuzu backend).
+
+    \b
+    Commands:
+      restore-from-snapshot  Replay a snapshot back into kuzu.
+    """
+
+
+@graph_group.command("restore-from-snapshot")
+@click.option(
+    "--snapshot",
+    "snapshot_path",
+    type=click.Path(exists=True, path_type=Path),
+    default=None,
+    help=(
+        "Snapshot file to restore (default: latest valid snapshot on disk). "
+        "Must be a .json snapshot written by the agent-brain-server."
+    ),
+)
+@click.option(
+    "--dry-run",
+    is_flag=True,
+    help="Report what would be restored and exit without mutating kuzu.",
+)
+@click.option(
+    "--yes",
+    "assume_yes",
+    is_flag=True,
+    help="Skip the confirmation prompt (non-interactive/CI use).",
+)
+def restore_from_snapshot(
+    snapshot_path: Path | None,
+    dry_run: bool,
+    assume_yes: bool,
+) -> None:
+    """Replay a kuzu snapshot back into the live graph.
+
+    Resolves the project's graph_index directory, reads the store type from
+    config, then replays the latest valid (or a specific) snapshot back into
+    the kuzu database.
+
+    Default (no flags): prints a summary of what WILL be restored and prompts
+    for confirmation.  Use --yes for CI/scripted invocations; --dry-run to
+    preview without touching the database.
+
+    The server MUST be stopped before running this command because kuzu does
+    not allow concurrent writers.
+    """
+    # --- Resolve state dir + store type ----------------------------------- #
+    project_root = resolve_project_root()
+    state_dir = project_root / STATE_DIR_NAME
+
+    block = _read_graphrag_block(state_dir)
+    store_type = str((block or {}).get("store_type") or "simple").lower()
+
+    if store_type != "kuzu":
+        console.print(
+            "[yellow]Graph restore only applies to the kuzu backend "
+            f"(configured store_type={store_type!r}). Nothing to do.[/yellow]"
+        )
+        return
+
+    # --- Guard: server must be stopped ------------------------------------ #
+    if _server_is_running(state_dir):
+        console.print(
+            "[red]Stop the server first (`agent-brain stop`) — restore "
+            "mutates the kuzu database the server holds open.[/red]"
+        )
+        raise SystemExit(1)
+
+    # --- Build the graph store manager ------------------------------------ #
+    graph_dir = _graph_index_dir(state_dir)
+    if GraphStoreManager is None:  # pragma: no cover
+        console.print(
+            "[red]Cannot import agent-brain-server. Install it with "
+            "`uv pip install agent-brain-rag`.[/red]"
+        )
+        raise SystemExit(1)
+
+    mgr = GraphStoreManager(graph_dir, "kuzu")
+
+    # --- Preview (plan) --------------------------------------------------- #
+    plan = mgr.plan_restore(snapshot_path)
+    if plan is None:
+        console.print(
+            "[yellow]No snapshot available to restore at "
+            f"{graph_dir / 'snapshots'}.[/yellow]"
+        )
+        raise SystemExit(1)
+
+    snap_file, triplet_count = plan
+    kuzu_db_path = graph_dir / "kuzu_db"
+    console.print(
+        f"Will restore [bold]{triplet_count}[/bold] triplets from "
+        f"[cyan]{snap_file.name}[/cyan] into [cyan]{kuzu_db_path}[/cyan]."
+    )
+
+    if dry_run:
+        console.print("(dry run — nothing changed)")
+        return
+
+    # --- Interactive confirmation ------------------------------------------ #
+    if not assume_yes:
+        confirmed = click.confirm("Proceed with restore?", default=False)
+        if not confirmed:
+            console.print("Aborted.")
+            return
+
+    # --- Execute restore --------------------------------------------------- #
+    restored = mgr.restore_from_snapshot(snapshot_path)
+    console.print(f"Restored [bold]{restored}[/bold] triplets.")

{agent_brain_cli-10.3.2 → agent_brain_cli-10.4.0}/agent_brain_cli/diagnostics.py

@@ -29,6 +29,15 @@ from agent_brain_cli.config import (
     resolve_project_root_with_strategy,
 )
 
+# agent-brain-rag is a declared dependency of agent-brain-cli. Import at module
+# level so apply_safe_fixes can be patched cleanly in tests.
+try:
+    from agent_brain_server.storage.graph_store import (
+        GraphStoreManager as GraphStoreManager,
+    )
+except ImportError:  # pragma: no cover
+    GraphStoreManager = None
+
 #: Severity returned by every diagnostic check.
 SEVERITY_OK = "ok"
 SEVERITY_WARN = "warn"
@@ -463,6 +472,130 @@ def _check_graph_store_health(state_dir: Path) -> CheckResult | None:
     )
 
 
+def _get_live_kuzu_relationship_count(graph_dir: Path) -> int | None:
+    """Open the kuzu DB read-only and return the live relationship count.
+
+    Returns ``None`` on any failure (not-installed, corrupt, etc.) — the
+    caller treats a None as "can't determine" and skips the staleness check.
+
+    Args:
+        graph_dir: The graph_index directory containing ``kuzu_db``.
+    """
+    try:
+        import kuzu  # noqa: PLC0415
+    except ImportError:
+        return None
+    kuzu_db_path = graph_dir / "kuzu_db"
+    try:
+        db = kuzu.Database(str(kuzu_db_path), read_only=True)
+        conn = kuzu.Connection(db)
+        result = conn.execute("MATCH ()-[r]->() RETURN COUNT(r)")
+        row = result.get_next()
+        return int(row[0]) if row else 0
+    except (IndexError, RuntimeError, OSError, TypeError, Exception):  # noqa: BLE001
+        return None
+
+
+def _check_graph_staleness(state_dir: Path) -> CheckResult | None:
+    """Warn when the live kuzu graph appears STALE (snapshot > live contents).
+
+    This check is SEPARATE from ``_check_graph_store_health`` — that one
+    verifies the DB opens without corruption; this one verifies the live
+    relationship count is not significantly below the latest snapshot count,
+    which signals an ``AGENT_BRAIN_JOB_TIMEOUT`` rollback left the graph
+    partially empty.
+
+    Returns ``None`` when:
+    - GraphRAG is not configured or store_type != "kuzu"
+    - The kuzu DB does not exist on disk yet
+    - No snapshot exists (nothing to compare against)
+    - The live COUNT query fails (could be a format change; don't FAIL here)
+
+    Args:
+        state_dir: Project state directory.
+    """
+    block = _read_graphrag_block(state_dir)
+    if not block or not block.get("enabled"):
+        return None
+    store_type = str(block.get("store_type") or "simple").lower()
+    if store_type != "kuzu":
+        return None
+
+    graph_dir = _graph_index_dir(state_dir)
+    kuzu_db_path = graph_dir / "kuzu_db"
+    if not kuzu_db_path.exists():
+        return None
+
+    # Check that at least one snapshot exists.
+    try:
+        from agent_brain_server.storage.graph_snapshot import (  # noqa: PLC0415
+            GraphSnapshotManager,
+        )
+    except ImportError:
+        # Server package not installed — fall back to manual JSON scan.
+        snap_dir = graph_dir / "snapshots"
+        if not snap_dir.is_dir() or not any(snap_dir.iterdir()):
+            return None
+        # Best-effort: pick the newest file and count triplets.
+        candidates = sorted(
+            (p for p in snap_dir.iterdir() if p.is_file() and p.suffix == ".json"),
+            key=lambda p: (p.stat().st_mtime, p.name),
+            reverse=True,
+        )
+        snapshot_count = 0
+        for snap in candidates:
+            try:
+                payload = json.loads(snap.read_text())
+                snapshot_count = len(payload.get("triplets") or [])
+                break
+            except (OSError, json.JSONDecodeError):
+                continue
+    else:
+        snap_mgr = GraphSnapshotManager(graph_dir)
+        if snap_mgr.latest() is None:
+            return None
+        loaded = snap_mgr.load_latest_valid()
+        if loaded is None:
+            return None
+        _snap_path, triplets = loaded
+        snapshot_count = len(triplets)
+
+    if snapshot_count == 0:
+        return None
+
+    live_count = _get_live_kuzu_relationship_count(graph_dir)
+    if live_count is None:
+        # Can't determine live count — skip rather than raise a false alarm.
+        return None
+
+    if snapshot_count > live_count:
+        return CheckResult(
+            "graph_staleness",
+            SEVERITY_WARN,
+            f"Live kuzu graph has {live_count} relationships but the latest "
+            f"snapshot has {snapshot_count} — the graph looks stale (likely an "
+            "AGENT_BRAIN_JOB_TIMEOUT rollback).",
+            fix=(
+                "Run `agent-brain graph restore-from-snapshot` (or "
+                "`agent-brain doctor --fix`) to replay the snapshot."
+            ),
+            details={
+                "live_relationships": live_count,
+                "snapshot_triplets": snapshot_count,
+            },
+        )
+
+    return CheckResult(
+        "graph_staleness",
+        SEVERITY_OK,
+        "Live kuzu graph is consistent with the latest snapshot.",
+        details={
+            "live_relationships": live_count,
+            "snapshot_triplets": snapshot_count,
+        },
+    )
+
+
 def _server_is_running(state_dir: Path) -> bool:
     """Best-effort check for an active server lock under state_dir."""
     return (state_dir / "server.lock").exists() or (state_dir / "lock").exists()
@@ -538,6 +671,11 @@ def run_doctor(server_url_override: str | None = None) -> DoctorReport:
     if graph_check is not None:
         checks.append(graph_check)
 
+    # Phase 64 GSTAB-02 — warn when live kuzu is stale (snapshot > live count).
+    stale_check = _check_graph_staleness(state_dir)
+    if stale_check is not None:
+        checks.append(stale_check)
+
     checks.append(_check_gitignore(project_root))
 
     checks.append(_check_server(server_url, runtime_file))
@@ -627,6 +765,34 @@ def apply_safe_fixes(report: DoctorReport) -> list[str]:
                     "start empty. Re-run `agent-brain index <folder>` to "
                     "rebuild."
                 )
+        elif check.name == "graph_staleness" and check.status == SEVERITY_WARN:
+            # Phase 64 GSTAB-02 — restore a stale (but non-corrupt) kuzu graph.
+            if _server_is_running(state_dir):
+                actions.append(
+                    "Skipped graph restore: server appears to be running. "
+                    "Stop it with `agent-brain stop` first, then re-run "
+                    "`agent-brain doctor --fix`."
+                )
+                continue
+            graph_dir = _graph_index_dir(state_dir)
+            if GraphStoreManager is None:  # pragma: no cover
+                actions.append(
+                    "Could not import agent-brain-server to perform restore. "
+                    "Run `agent-brain graph restore-from-snapshot` manually."
+                )
+                continue
+            try:
+                mgr = GraphStoreManager(graph_dir, "kuzu")
+                restored_count = mgr.restore_from_snapshot(None)
+                actions.append(
+                    f"Restored {restored_count} triplets from the latest "
+                    f"snapshot into {graph_dir}."
+                )
+            except Exception as exc:  # noqa: BLE001
+                actions.append(
+                    f"Graph restore failed: {exc}. "
+                    "Run `agent-brain graph restore-from-snapshot` manually."
+                )
     return actions
 
 

{agent_brain_cli-10.3.2 → agent_brain_cli-10.4.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "agent-brain-cli"
-version = "10.3.2"
+version = "10.4.0"
 description = "Agent Brain CLI - Command-line interface for managing AI agent memory and knowledge retrieval"
 authors = ["Spillwave Solutions"]
 readme = "README.md"
@@ -27,7 +27,7 @@ httpx = "^0.28.0"
 rich = "^13.9.0"
 pyyaml = "^6.0.0"
 pydantic = "^2.10.0"
-agent-brain-rag = "^10.3.2"
+agent-brain-rag = "^10.4.0"
 agent-brain-uds = "^10.3.0"
 psutil = "^6.0"