agentcache-core 0.9.9__py3-none-any.whl

agentcache/db.py

@@ -0,0 +1,359 @@
+import atexit
+import json
+import os
+import sqlite3
+import threading
+import time
+from typing import Any, Dict, List, Optional, TypeVar
+
+T = TypeVar("T")
+
+DB_PATH = os.path.join(os.path.expanduser("~"), ".agentcache", "agentcache.db")
+
+
+class StateKV:
+    def __init__(self, db_path: Optional[str] = None):
+        self.db_path = db_path or os.getenv("AGENTCACHE_DB_PATH") or DB_PATH
+        self._lock = threading.Lock()
+        # Per-thread persistent connection pool (A3.1)
+        self._local = threading.local()
+        self._init_db()
+        # Register WAL checkpoint on graceful shutdown (A3.2)
+        atexit.register(self._wal_checkpoint)
+
+    def _get_conn(self) -> sqlite3.Connection:
+        """Return a per-thread persistent connection, creating one if needed (A3.1)."""
+        conn = getattr(self._local, "conn", None)
+        if conn is None:
+            conn = sqlite3.connect(self.db_path, check_same_thread=False, timeout=30.0)
+            conn.row_factory = sqlite3.Row
+            conn.execute("PRAGMA journal_mode=WAL")
+            conn.execute("PRAGMA synchronous=NORMAL")
+            conn.execute("PRAGMA foreign_keys=ON")
+            conn.execute("PRAGMA journal_size_limit = 67108864")  # 64MB WAL limit
+            conn.execute("PRAGMA mmap_size = 268435456")  # 256MB mmap
+            self._local.conn = conn
+        return conn
+
+    def _wal_checkpoint(self) -> None:
+        """Flush WAL frames to the main database file on shutdown (A3.2)."""
+        try:
+            conn = getattr(self._local, "conn", None)
+            if conn:
+                conn.execute("PRAGMA wal_checkpoint(FULL)")
+                conn.commit()
+            else:
+                # Open a temporary connection just for the checkpoint
+                tmp = sqlite3.connect(self.db_path, check_same_thread=False, timeout=10)
+                try:
+                    tmp.execute("PRAGMA wal_checkpoint(FULL)")
+                    tmp.commit()
+                finally:
+                    tmp.close()
+            print("[db] WAL checkpoint completed on shutdown.")
+        except Exception as e:
+            print(f"[db] WAL checkpoint failed: {e}")
+
+    def teardown(self) -> None:
+        """Close the per-thread connection and flush WAL (for explicit cleanup)."""
+        self._wal_checkpoint()
+        conn = getattr(self._local, "conn", None)
+        if conn:
+            try:
+                conn.close()
+            except Exception:
+                pass
+            self._local.conn = None
+
+    def stats(self) -> Dict[str, Any]:
+        """Return DB statistics for the /health endpoint (A3.3).
+
+        Returns:
+            {
+                "db_size_bytes": int,
+                "kv_row_count": int,
+                "audit_row_count": int,
+                "wal_size_bytes": int,
+            }
+        """
+        result: Dict[str, Any] = {
+            "db_size_bytes": 0,
+            "kv_row_count": 0,
+            "audit_row_count": 0,
+            "wal_size_bytes": 0,
+        }
+        try:
+            if os.path.exists(self.db_path):
+                result["db_size_bytes"] = os.path.getsize(self.db_path)
+            wal_path = self.db_path + "-wal"
+            if os.path.exists(wal_path):
+                result["wal_size_bytes"] = os.path.getsize(wal_path)
+        except Exception:
+            pass
+        try:
+            conn = self._get_conn()
+            result["kv_row_count"] = conn.execute(
+                "SELECT COUNT(*) FROM kv_store"
+            ).fetchone()[0]
+            result["audit_row_count"] = conn.execute(
+                "SELECT COUNT(*) FROM audit_log"
+            ).fetchone()[0]
+        except Exception:
+            pass
+        return result
+
+    def _init_db(self):
+        try:
+            os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
+            # Use a temporary direct connection for initialization (before _local is set)
+            conn = sqlite3.connect(self.db_path, check_same_thread=False, timeout=30.0)
+            conn.row_factory = sqlite3.Row
+            conn.execute("PRAGMA journal_mode=WAL")
+            conn.execute("PRAGMA synchronous=NORMAL")
+            conn.execute("PRAGMA foreign_keys=ON")
+            conn.execute("PRAGMA journal_size_limit = 67108864")  # 64MB WAL limit
+            conn.execute("PRAGMA mmap_size = 268435456")  # 256MB mmap
+            try:
+                conn.execute("""
+                    CREATE TABLE IF NOT EXISTS kv_store (
+                        scope TEXT NOT NULL,
+                        key   TEXT NOT NULL,
+                        value TEXT NOT NULL,
+                        PRIMARY KEY (scope, key)
+                    )
+                """)
+                conn.execute(
+                    "CREATE INDEX IF NOT EXISTS idx_kv_scope ON kv_store(scope)"
+                )
+                conn.execute("""
+                    CREATE TABLE IF NOT EXISTS audit_log (
+                        id        INTEGER PRIMARY KEY AUTOINCREMENT,
+                        ts        INTEGER NOT NULL,
+                        agent_id  TEXT NOT NULL,
+                        message   TEXT NOT NULL
+                    )
+                """)
+                conn.execute("""
+                    CREATE TABLE IF NOT EXISTS sync_state_metadata (
+                        key   TEXT PRIMARY KEY,
+                        value TEXT NOT NULL
+                    )
+                """)
+                conn.commit()
+            finally:
+                conn.close()
+            print(f"[db] SQLite database initialized at {self.db_path}")
+        except Exception as e:
+            print(f"[db] WARNING initializing SQLite database: {e}")
+
+    def get(self, scope: str, key: str) -> Optional[Any]:
+        try:
+            conn = self._get_conn()
+            row = conn.execute(
+                "SELECT value FROM kv_store WHERE scope = ? AND key = ?", (scope, key)
+            ).fetchone()
+            if row:
+                val = json.loads(row["value"])
+                if isinstance(val, dict) and "id" not in val:
+                    val["id"] = key
+                return val
+            return None
+        except Exception as e:
+            print(f"[db] get failed (scope={scope}, key={key}): {e}")
+            return None
+
+    def _execute_write_with_retry(self, action_func):
+        """Helper to run a write transaction with exponential backoff on lock/busy errors."""
+        max_retries = 5
+        delay = 0.05
+        with self._lock:
+            for attempt in range(max_retries):
+                try:
+                    return action_func()
+                except sqlite3.OperationalError as e:
+                    err_msg = str(e).lower()
+                    if (
+                        "locked" in err_msg or "busy" in err_msg
+                    ) and attempt < max_retries - 1:
+                        time.sleep(delay)
+                        delay *= 2
+                        continue
+                    raise
+
+    def set(self, scope: str, key: str, value: Any) -> Any:
+        def action():
+            conn = self._get_conn()
+            conn.execute(
+                "INSERT OR REPLACE INTO kv_store (scope, key, value) VALUES (?, ?, ?)",
+                (scope, key, json.dumps(value)),
+            )
+            conn.commit()
+            return value
+
+        try:
+            return self._execute_write_with_retry(action)
+        except Exception as e:
+            print(f"[db] set failed (scope={scope}, key={key}): {e}")
+            return value
+
+    def delete(self, scope: str, key: str) -> bool:
+        def action():
+            conn = self._get_conn()
+            cur = conn.execute(
+                "DELETE FROM kv_store WHERE scope = ? AND key = ?", (scope, key)
+            )
+            conn.commit()
+            return cur.rowcount > 0
+
+        try:
+            return self._execute_write_with_retry(action)
+        except Exception as e:
+            print(f"[db] delete failed (scope={scope}, key={key}): {e}")
+            return False
+
+    def list(self, scope: str) -> List[Any]:
+        try:
+            conn = self._get_conn()
+            rows = conn.execute(
+                "SELECT key, value FROM kv_store WHERE scope = ?", (scope,)
+            ).fetchall()
+            results = []
+            for r in rows:
+                val = json.loads(r["value"])
+                if isinstance(val, dict) and "id" not in val:
+                    val["id"] = r["key"]
+                results.append(val)
+            return results
+        except Exception as e:
+            print(f"[db] list failed (scope={scope}): {e}")
+            return []
+
+    def update(self, scope: str, key: str, ops: List[Dict[str, Any]]) -> Optional[Any]:
+        def action():
+            conn = self._get_conn()
+            row = conn.execute(
+                "SELECT value FROM kv_store WHERE scope = ? AND key = ?",
+                (scope, key),
+            ).fetchone()
+            obj = json.loads(row["value"]) if row else {}
+            if not isinstance(obj, dict):
+                obj = {}
+            if "id" not in obj:
+                obj["id"] = key
+
+            for op in ops:
+                op_type = op.get("type")
+                path = op.get("path")
+                val = op.get("value")
+                if not path:
+                    continue
+                if op_type == "set":
+                    if "." in path:
+                        parts = path.split(".")
+                        curr = obj
+                        for part in parts[:-1]:
+                            if part not in curr or not isinstance(curr[part], dict):
+                                curr[part] = {}
+                            curr = curr[part]
+                        curr[parts[-1]] = val
+                    else:
+                        obj[path] = val
+                elif op_type == "delete":
+                    if "." in path:
+                        parts = path.split(".")
+                        curr = obj
+                        for part in parts[:-1]:
+                            if part not in curr or not isinstance(curr[part], dict):
+                                break
+                            curr = curr[part]
+                        else:
+                            curr.pop(parts[-1], None)
+                    else:
+                        obj.pop(path, None)
+
+            conn.execute(
+                "INSERT OR REPLACE INTO kv_store (scope, key, value) VALUES (?, ?, ?)",
+                (scope, key, json.dumps(obj)),
+            )
+            conn.commit()
+            return obj
+
+        try:
+            return self._execute_write_with_retry(action)
+        except Exception as e:
+            print(f"[db] update failed (scope={scope}, key={key}): {e}")
+            return None
+
+    def commit_version(self, message: str, agent_id: str) -> Optional[str]:
+        """Write an audit log entry instead of a Dolt commit."""
+        author = agent_id or "unknown-agent"
+
+        def action():
+            conn = self._get_conn()
+            cur = conn.execute(
+                "INSERT INTO audit_log (ts, agent_id, message) VALUES (?, ?, ?)",
+                (int(time.time() * 1000), author, message),
+            )
+            conn.commit()
+            row_id = str(cur.lastrowid)
+            print(f"[audit] {author}: {message} (id={row_id})")
+            return row_id
+
+        try:
+            return self._execute_write_with_retry(action)
+        except Exception as e:
+            print(f"[audit] commit_version failed: {e}")
+            return None
+
+    def get_audit_log(self, limit: int = 50) -> List[Dict[str, Any]]:
+        try:
+            conn = self._get_conn()
+            rows = conn.execute(
+                "SELECT id, ts, agent_id, message FROM audit_log ORDER BY ts DESC LIMIT ?",
+                (limit,),
+            ).fetchall()
+            return [dict(r) for r in rows]
+        except Exception as e:
+            print(f"[db] get_audit_log failed: {e}")
+            return []
+
+    def acquire_lock(self, lock_name: str, lease_seconds: int = 300) -> bool:
+        """Try to acquire a distributed lock in the database, atomic across processes."""
+        now = int(time.time())
+        lock_key = f"lock:{lock_name}"
+
+        def action():
+            conn = self._get_conn()
+            row = conn.execute(
+                "SELECT value FROM kv_store WHERE scope = ? AND key = ?",
+                ("mem:locks", lock_key),
+            ).fetchone()
+            if row:
+                val = json.loads(row["value"])
+                expires = val.get("expires", 0)
+                if now < expires:
+                    return False
+
+            # Lock is expired or doesn't exist. Write new lock.
+            lock_data = {
+                "expires": now + lease_seconds,
+                "acquired_at": now,
+                "owner": f"pid-{os.getpid()}",
+            }
+            conn.execute(
+                "INSERT OR REPLACE INTO kv_store (scope, key, value) VALUES (?, ?, ?)",
+                ("mem:locks", lock_key, json.dumps(lock_data)),
+            )
+            conn.commit()
+            return True
+
+        try:
+            return self._execute_write_with_retry(action)
+        except Exception as e:
+            print(f"[db] acquire_lock failed ({lock_name}): {e}")
+            return False
+
+    def release_lock(self, lock_name: str) -> None:
+        """Release a distributed lock in the database."""
+        lock_key = f"lock:{lock_name}"
+        self.delete("mem:locks", lock_key)

agentcache/import_data.py

@@ -0,0 +1,86 @@
+import json
+import os
+import sys
+import urllib.parse
+
+from .db import StateKV
+
+
+def import_old_data(old_db_path: str, kv: StateKV) -> bool:
+    if not os.path.exists(old_db_path):
+        print(f"[import] Error: Path {old_db_path} does not exist.")
+        return False
+
+    print(f"[import] Starting migration from: {old_db_path}")
+    count = 0
+
+    for filename in os.listdir(old_db_path):
+        if not filename.endswith(".bin"):
+            continue
+        if not filename.startswith("mem%3A"):
+            continue
+
+        # URL decode the scope name (e.g. mem%3Asessions.bin -> mem:sessions)
+        decoded_name = urllib.parse.unquote(filename)
+        scope = decoded_name[:-4] if decoded_name.endswith(".bin") else decoded_name
+
+        filepath = os.path.join(old_db_path, filename)
+        try:
+            with open(filepath, "rb") as f:
+                data = f.read()
+
+            if not data:
+                continue
+
+            # Find the last matching JSON curly brace to strip binary padding/footers
+            try:
+                last_brace = data.rindex(b"}")
+                json_str = data[: last_brace + 1].decode("utf-8")
+                records = json.loads(json_str)
+            except Exception as e:
+                print(f"[import] Skipping {filename}: could not parse JSON ({e})")
+                continue
+
+            print(f"[import] Scope '{scope}': found {len(records)} items. Migrating...")
+
+            for key, val in records.items():
+                kv.set(scope, key, val)
+                count += 1
+
+        except Exception as e:
+            print(f"[import] Failed to process {filename}: {e}")
+
+    if count > 0:
+        print("[import] Creating Dolt version commit...")
+        kv.commit_version(
+            "Import legacy AgentCache database from Hugging Face", "system"
+        )
+
+    print(f"[import] Finished! Migrated {count} total records into Dolt SQL.")
+    return True
+
+
+if __name__ == "__main__":
+    # Prioritize the restored HF data in the user's home folder, falling back to local workspace data
+    home_path = os.path.expanduser(os.path.join("~", ".agentcache", "state_store.db"))
+    workspace_path = r"d:\Downloads\Projects\Other Projects\agentcache\agentcache\data\state_store.db"
+
+    if os.path.exists(home_path):
+        default_path = home_path
+        print(f"[import] Found restored HF data at: {default_path}")
+    else:
+        default_path = workspace_path
+        print(
+            f"[import] HF data not found at home path, falling back to workspace: {default_path}"
+        )
+
+    # Allow custom path via command line argument
+    if len(sys.argv) > 1:
+        default_path = sys.argv[1]
+        print(f"[import] Overriding path with CLI argument: {default_path}")
+
+    # Initialize connection
+    print("[import] Initializing Dolt StateKV...")
+    kv = StateKV()
+
+    import_old_data(default_path, kv)

agentcache/legacy.py

@@ -0,0 +1,94 @@
+"""Re-export shim — behaviour lives in core/*. See ADR-0001."""
+
+import threading
+from typing import Any, Dict, Optional
+
+from .core.audit_log import *  # noqa: F401,F403
+from .core.config import *  # noqa: F401,F403
+from .core.context_builder import *  # noqa: F401,F403
+from .core.graph import *  # noqa: F401,F403
+from .core.image_store import *  # noqa: F401,F403
+from .core.infer import *  # noqa: F401,F403
+from .core.kv_scopes import KV  # noqa: F401  re-exported for backward compat
+from .core.lessons import *  # noqa: F401,F403
+from .core.llm import *  # noqa: F401,F403
+from .core.memory_store import *  # noqa: F401,F403
+from .core.observation_store import (  # noqa: F401
+    normalize_folder_path,
+    validate_agent_id,
+)
+from .core.privacy import *  # noqa: F401,F403
+from .core.project_profile import *  # noqa: F401,F403
+from .core.search_service import IndexPersistence  # noqa: F401
+from .core.session_store import *  # noqa: F401,F403
+from .core.slots import *  # noqa: F401,F403
+from .storage.paths import fingerprint_id, generate_id  # noqa: F401
+
+# ---------------------------------------------------------------------------
+# Module-level state — intentionally kept here (see ADR-0001 "Out of Scope").
+# core/* modules access these via lazy `from .. import legacy as _legacy`.
+# ---------------------------------------------------------------------------
+_search_service = None  # type: ignore[assignment]  SearchService | None
+_stream_broadcaster = None  # Callable: (payload) -> None
+_dedup_locks: Dict[str, threading.Lock] = {}
+_dedup_locks_meta = threading.Lock()
+
+
+def set_search_service(service) -> None:
+    """Register the SearchService instance created by app.init_services()."""
+    global _search_service
+    _search_service = service
+
+
+def set_stream_broadcaster(broadcaster) -> None:
+    global _stream_broadcaster
+    _stream_broadcaster = broadcaster
+
+
+def broadcast_stream(payload: Dict[str, Any]) -> None:
+    if _stream_broadcaster:
+        try:
+            _stream_broadcaster(payload)
+        except Exception as e:
+            print(f"[broadcaster] Failed: {e}")
+
+
+def backfill_obs_lookup_if_needed(kv) -> None:
+    """Ensure every folder observation has an entry in KV.obs_lookup."""
+    from .core.session_store import _get_observation_store
+
+    store = _get_observation_store(kv)
+    store.backfill_lookup()
+
+
+def verify_index_sync_on_boot(kv, search_service: Optional[Any] = None) -> bool:
+    """Check if the search index size matches the database counts."""
+    try:
+        svc = search_service or _search_service
+        if svc is None:
+            from . import app as app_module
+
+            svc = getattr(app_module, "search_service", None)
+
+        folders = kv.list(KV.folders)
+        folder_obs_count = sum(int(f.get("obsCount", 0)) for f in folders)
+
+        memories = kv.list(KV.memories)
+        latest_memories_count = len(
+            [m for m in memories if m.get("isLatest") is not False]
+        )
+
+        total_db_count = folder_obs_count + latest_memories_count
+        index_size = svc.bm25_size if svc else 0
+
+        if total_db_count != index_size:
+            print(
+                f"[persistence] Index out of sync with DB (DB={total_db_count}, Index={index_size}). Rebuild required."
+            )
+            return False
+
+        print(f"[persistence] Index is in sync with DB (size={index_size}).")
+        return True
+    except Exception as e:
+        print(f"[persistence] verify_index_sync_on_boot failed: {e}")
+        return False

agentcache/mcp_stdio.py

@@ -0,0 +1,141 @@
+#!/usr/bin/env python3
+"""
+stdio MCP wrapper — bridges the MCP stdio protocol to the
+agentmemory Flask HTTP API running on localhost.
+
+Usage: python mcp_stdio.py
+
+Antigravity sync functions have been moved to examples/antigravity_sync.py (D1.3).
+All tool calls are proxied to the HTTP API.
+"""
+
+import json
+import os
+import sys
+
+import requests
+
+BASE = (
+    os.getenv("AGENTCACHE_URL")
+    or os.getenv("AGENTMEMORY_URL")
+    or "http://127.0.0.1:3111"
+).rstrip("/")
+if not BASE.endswith("/agentcache") and not BASE.endswith("/agentmemory"):
+    BASE = f"{BASE}/agentcache"
+
+_secret = os.getenv("AGENTCACHE_SECRET") or os.getenv("AGENTMEMORY_SECRET")
+
+
+def headers():
+    h = {"Content-Type": "application/json"}
+    if _secret:
+        h["Authorization"] = f"Bearer {_secret}"
+    return h
+
+
+def send(msg):
+    line = json.dumps(msg, separators=(",", ":"))
+    sys.stdout.write(line + "\n")
+    sys.stdout.flush()
+
+
+def handle(req):
+    method = req.get("method", "")
+    req_id = req.get("id")
+    params = req.get("params") or {}
+
+    if method == "initialize":
+        send(
+            {
+                "jsonrpc": "2.0",
+                "id": req_id,
+                "result": {
+                    "protocolVersion": "2024-11-05",
+                    "capabilities": {"tools": {}},
+                    "serverInfo": {"name": "agentcache-local", "version": "0.9.9"},
+                },
+            }
+        )
+
+    elif method == "initialized":
+        pass  # notification — no response
+
+    elif method == "ping":
+        send({"jsonrpc": "2.0", "id": req_id, "result": {}})
+
+    elif method == "tools/list":
+        try:
+            r = requests.get(f"{BASE}/mcp/tools", headers=headers(), timeout=5)
+            tools = r.json().get("tools", [])
+            send({"jsonrpc": "2.0", "id": req_id, "result": {"tools": tools}})
+        except Exception as e:
+            send(
+                {
+                    "jsonrpc": "2.0",
+                    "id": req_id,
+                    "error": {
+                        "code": -32000,
+                        "message": f"agentmemory unreachable: {e}",
+                    },
+                }
+            )
+
+    elif method == "tools/call":
+        name = params.get("name")
+        args = params.get("arguments") or {}
+        try:
+            r = requests.post(
+                f"{BASE}/mcp/tools",
+                headers=headers(),
+                json={"name": name, "arguments": args},
+                timeout=30,
+            )
+            result = r.json()
+            # MCP expects content array
+            if "content" not in result:
+                result = {"content": [{"type": "text", "text": json.dumps(result)}]}
+            send({"jsonrpc": "2.0", "id": req_id, "result": result})
+        except Exception as e:
+            send(
+                {
+                    "jsonrpc": "2.0",
+                    "id": req_id,
+                    "error": {"code": -32000, "message": str(e)},
+                }
+            )
+
+    elif req_id is not None:
+        send(
+            {
+                "jsonrpc": "2.0",
+                "id": req_id,
+                "error": {"code": -32601, "message": "Method not found"},
+            }
+        )
+
+
+def main():
+    for raw_line in sys.stdin:
+        raw_line = raw_line.strip()
+        if not raw_line:
+            continue
+        try:
+            req = json.loads(raw_line)
+        except json.JSONDecodeError:
+            continue
+        try:
+            handle(req)
+        except Exception as e:
+            req_id = req.get("id") if isinstance(req, dict) else None
+            if req_id is not None:
+                send(
+                    {
+                        "jsonrpc": "2.0",
+                        "id": req_id,
+                        "error": {"code": -32603, "message": f"Internal error: {e}"},
+                    }
+                )
+
+
+if __name__ == "__main__":
+    main()