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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: agent-brain-cli
-Version: 10.3.1
+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.1,<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.1 → 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.1"
+__version__ = "10.4.0"

{agent_brain_cli-10.3.1 → 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.1 → agent_brain_cli-10.4.0}/agent_brain_cli/client/transport.py

@@ -32,6 +32,7 @@ Three §3.5 design-doc misuse cases surface as ``click.UsageError``
 
 from __future__ import annotations
 
+import os
 import shutil
 from pathlib import Path
 from typing import cast
@@ -63,6 +64,41 @@ def _resolve_state_dir_for_discovery() -> Path | None:
         return None
 
 
+def _stdio_command(state_dir: Path | None) -> str | list[str]:
+    """Build the ``agent-brain-mcp`` stdio command, pinning the state dir.
+
+    When ``state_dir`` is resolved, returns
+    ``["agent-brain-mcp", "--state-dir", <path>]`` so the spawned MCP
+    server resolves the same backend the CLI's UDS path would (CLI-MCP-04
+    byte-identical guarantee). ``McpStdioBackend`` filters the subprocess
+    env through ``DEFAULT_ENV_ALLOWLIST`` (which intentionally drops
+    ``AGENT_BRAIN_STATE_DIR``), so an explicit arg is the only reliable
+    way to propagate the state dir across the env boundary.
+
+    ``AGENT_BRAIN_STATE_DIR`` takes precedence over the discovery value,
+    mirroring the UDS leg (``resolve_transport`` →
+    ``resolve_socket_path`` honors the same override). Without this, a
+    user who sets ``AGENT_BRAIN_STATE_DIR`` and runs from a non-project
+    cwd would get the project-root discovery dir on the MCP leg but the
+    override on the UDS leg — diverging the two and breaking CLI-MCP-04.
+
+    When neither the env override nor ``state_dir`` resolves, returns the
+    bare ``"agent-brain-mcp"`` string and lets the subprocess fall back
+    to its own cwd-based discovery — preserving pre-fix behavior.
+    """
+    env_override = os.environ.get("AGENT_BRAIN_STATE_DIR")
+    if env_override:
+        resolved: str | None = str(Path(env_override).expanduser())
+    elif state_dir is not None:
+        resolved = str(state_dir)
+    else:
+        resolved = None
+
+    if resolved is None:
+        return "agent-brain-mcp"
+    return ["agent-brain-mcp", "--state-dir", resolved]
+
+
 def open_backend(ctx: click.Context, *, timeout: float = 30.0) -> BackendClient:
     """Construct a ``BackendClient`` over the transport selected by ``ctx``.
 
@@ -132,15 +168,32 @@ def open_backend(ctx: click.Context, *, timeout: float = 30.0) -> BackendClient:
                     "agent-brain-mcp not found on PATH; install "
                     "agent-brain-mcp into the same Python environment"
                 )
+            # Pin the resolved state dir into the subprocess as an
+            # explicit --state-dir arg (CLI-MCP-04). McpStdioBackend
+            # filters env through DEFAULT_ENV_ALLOWLIST, which drops
+            # AGENT_BRAIN_STATE_DIR, so without this the spawned
+            # agent-brain-mcp falls back to cwd discovery and may resolve
+            # a DIFFERENT backend than --transport uds. Passing the dir
+            # explicitly makes the byte-identical guarantee hold from any
+            # cwd. When state_dir is None (resolution failed) we keep the
+            # bare command and let agent-brain-mcp do its own cwd discovery.
+            command = _stdio_command(state_dir)
             # cast(): agent_brain_mcp ships with ignore_missing_imports=true
             # so mypy treats McpStdioBackend as Any. The runtime
             # @runtime_checkable Protocol contract is pinned by the
             # Phase 56-03 isinstance test in agent-brain-mcp/tests/.
-            return cast(BackendClient, McpStdioBackend(command="agent-brain-mcp"))
+            return cast(BackendClient, McpStdioBackend(command=command))
         # 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(
@@ -244,11 +297,16 @@ def open_mcp_backend(ctx: click.Context, *, timeout: float = 30.0) -> McpBackend
                 "agent-brain-mcp not found on PATH; install "
                 "agent-brain-mcp into the same Python environment"
             )
+        # Pin the resolved state dir into the subprocess (CLI-MCP-04) —
+        # same rationale as open_backend: DEFAULT_ENV_ALLOWLIST drops
+        # AGENT_BRAIN_STATE_DIR, so prompts/resources would otherwise run
+        # against a cwd-discovered backend.
+        command = _stdio_command(state_dir)
         # cast(): mirrors open_backend's pattern. The runtime
         # @runtime_checkable Protocol contract is pinned by the
         # Plan 59-01 isinstance test in agent-brain-mcp/tests/
         # test_mcp_backend_protocol_skeleton.py.
-        return cast(McpBackend, McpStdioBackend(command="agent-brain-mcp"))
+        return cast(McpBackend, McpStdioBackend(command=command))
     # mcp_transport == "http" — resolve_mcp_transport guarantees
     # mcp_target is not None for the http branch.
     assert mcp_target is not None  # noqa: S101

{agent_brain_cli-10.3.1 → 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.1 → 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.1 → agent_brain_cli-10.4.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "agent-brain-cli"
-version = "10.3.1"
+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.1"
+agent-brain-rag = "^10.4.0"
 agent-brain-uds = "^10.3.0"
 psutil = "^6.0"