keepm 0.2.0__py3-none-any.whl

keepm/server.py

@@ -0,0 +1,450 @@
+from __future__ import annotations
+
+import argparse
+import asyncio
+import dataclasses
+import datetime as dt
+import hashlib
+import sqlite3
+from contextlib import asynccontextmanager
+from pathlib import Path
+from collections.abc import Sequence
+from typing import Any
+
+from keepm import __version__
+from keepm.config import ConfigStore, KeepMConfig, MemoryLanguage, default_data_root
+from keepm.database import SCHEMA_VERSION, MemoryDatabase
+from keepm.models import MemoryDocument
+from keepm.projects import resolve_workspace
+from keepm.service import MemoryService
+from keepm.storage import MemoryStorage
+from keepm.sync import SyncEngine
+from keepm.watcher import WatcherService
+
+try:
+    from fastmcp import Context, FastMCP
+except ImportError:
+    Context = Any  # type: ignore[misc,assignment]
+    FastMCP = None  # type: ignore[assignment]
+
+
+def _vault_id(config: KeepMConfig) -> str:
+    identity = f"{config.vault_path}\0{config.memory_root.resolve()}"
+    return hashlib.sha256(identity.encode("utf-8")).hexdigest()[:24]
+
+
+def _jsonable(value: Any) -> Any:
+    if dataclasses.is_dataclass(value):
+        return _jsonable(dataclasses.asdict(value))
+    if isinstance(value, Path):
+        return str(value)
+    if isinstance(value, dict):
+        return {str(key): _jsonable(item) for key, item in value.items()}
+    if isinstance(value, (tuple, list, set)):
+        return [_jsonable(item) for item in value]
+    return value
+
+
+def _document_result(
+    document: MemoryDocument, checksum: str, *, max_chars: int
+) -> dict[str, Any]:
+    if max_chars < 1:
+        raise ValueError("max_chars must be positive")
+    payload = _jsonable(document)
+    body = str(payload["body"])
+    truncated = len(body) > max_chars
+    payload["body"] = body[:max_chars]
+    return {
+        "memory": payload,
+        "checksum": checksum,
+        "output_chars": len(payload["body"]),
+        "truncated": truncated,
+    }
+
+
+class Runtime:
+    def __init__(
+        self,
+        *,
+        config_path: Path | None = None,
+        data_root: Path | None = None,
+    ) -> None:
+        self.config_store = ConfigStore(config_path)
+        self.data_root = (data_root or default_data_root()).expanduser().resolve()
+        self.config: KeepMConfig | None = None
+        self.database: MemoryDatabase | None = None
+        self.storage: MemoryStorage | None = None
+        self.sync_engine: SyncEngine | None = None
+        self.service: MemoryService | None = None
+        self.watcher: WatcherService | None = None
+        self.last_error: str | None = None
+        self.recovered_index: str | None = None
+        self._lock = asyncio.Lock()
+
+    async def start(self) -> None:
+        async with self._lock:
+            if self.service is not None:
+                return
+            try:
+                config = self.config_store.load()
+            except FileNotFoundError:
+                return
+            except Exception as error:
+                self.last_error = f"cannot load KeepM configuration: {error}"
+                return
+            try:
+                await self._build_locked(config)
+            except Exception as error:
+                self.last_error = f"cannot initialize KeepM: {error}"
+
+    async def configure(
+        self,
+        *,
+        vault_path: str,
+        memory_dir: str = "Memory",
+        stale_after_seconds: int = 30,
+        language: MemoryLanguage | None = None,
+    ) -> dict[str, Any]:
+        vault = Path(vault_path).expanduser().resolve()
+        if not vault.is_dir():
+            raise ValueError(f"vault_path must be an existing directory: {vault}")
+        existing = self.config
+        if existing is None:
+            existing = self.config_store.load_optional()
+        config = KeepMConfig(
+            vault_path=vault,
+            memory_dir=memory_dir,
+            stale_after_seconds=stale_after_seconds,
+            language=language or (existing.language if existing is not None else "en"),
+        )
+        async with self._lock:
+            await self._teardown_locked()
+            self.config_store.save(config)
+            await self._build_locked(config)
+            self.last_error = None
+            return self.status()
+
+    async def _build_locked(self, config: KeepMConfig) -> None:
+        self.config = config
+        self.storage = MemoryStorage(config)
+        self.database = self._open_database(config)
+        self.sync_engine = SyncEngine(self.storage, self.database)
+        self.service = MemoryService(
+            config, self.storage, self.database, self.sync_engine
+        )
+        report = self.sync_engine.quick_reconcile()
+        self.database.set_state("dirty", "1" if report.errors else "0")
+        self.watcher = WatcherService(config, self.sync_engine)
+        await self.watcher.start()
+
+    def _open_database(self, config: KeepMConfig) -> MemoryDatabase:
+        database_path = self.data_root / "indexes" / f"{_vault_id(config)}.sqlite3"
+        database: MemoryDatabase | None = None
+        try:
+            database = MemoryDatabase(database_path)
+            if not database.quick_check():
+                raise sqlite3.DatabaseError("PRAGMA quick_check failed")
+            if database.get_state("schema_version") != SCHEMA_VERSION:
+                raise sqlite3.DatabaseError("incompatible KeepM index schema")
+            canonical_root = database.get_state("canonical_root")
+            expected_root = str(config.memory_root.resolve())
+            if canonical_root not in (None, expected_root):
+                raise sqlite3.DatabaseError("index belongs to a different Memory root")
+        except (OSError, RuntimeError, sqlite3.DatabaseError) as error:
+            if database is not None:
+                database.close()
+            if database_path.exists():
+                suffix = dt.datetime.now().strftime("%Y%m%d-%H%M%S-%f")
+                preserved = database_path.with_name(
+                    f"{database_path.name}.corrupt-{suffix}"
+                )
+                database_path.replace(preserved)
+                for sidecar_suffix in ("-wal", "-shm"):
+                    sidecar = Path(str(database_path) + sidecar_suffix)
+                    if sidecar.exists():
+                        sidecar.replace(Path(str(preserved) + sidecar_suffix))
+                self.recovered_index = f"{preserved}: {error}"
+            database = MemoryDatabase(database_path)
+
+        database.set_state("canonical_root", str(config.memory_root.resolve()))
+        return database
+
+    async def _teardown_locked(self) -> None:
+        if self.watcher is not None:
+            await self.watcher.stop()
+        if self.database is not None:
+            self.database.close()
+        self.config = None
+        self.database = None
+        self.storage = None
+        self.sync_engine = None
+        self.service = None
+        self.watcher = None
+
+    async def close(self) -> None:
+        async with self._lock:
+            await self._teardown_locked()
+
+    def require_service(self) -> MemoryService:
+        if self.service is None:
+            raise RuntimeError(
+                "KeepM is not configured. Call memory_configure with your Obsidian Vault path."
+            )
+        return self.service
+
+    def status(self, workspace: Path | None = None) -> dict[str, Any]:
+        if self.service is None:
+            return {
+                "configured": False,
+                "config_path": str(self.config_store.path),
+                "error": self.last_error,
+            }
+        status = self.service.status(workspace=workspace)
+        status.update(
+            {
+                "watcher_running": bool(self.watcher and self.watcher.running),
+                "watcher_error": self.watcher.last_error if self.watcher else None,
+                "recovered_index": self.recovered_index,
+                "error": self.last_error,
+            }
+        )
+        return status
+
+
+async def _workspace_from_context(
+    context: Any, workspace_path: str | None
+) -> Path:
+    client_roots: list[str] = []
+    list_roots = getattr(context, "list_roots", None)
+    if list_roots is not None:
+        try:
+            roots = await list_roots()
+            for root in roots or ():
+                client_roots.append(str(getattr(root, "uri", root)))
+        except Exception:
+            client_roots = []
+    return resolve_workspace(client_roots, workspace_path, Path.cwd())
+
+
+def create_mcp(
+    *,
+    config_path: Path | None = None,
+    data_root: Path | None = None,
+) -> Any:
+    if FastMCP is None:
+        raise RuntimeError(
+            "FastMCP is not installed. Install KeepM with its declared dependencies."
+        )
+    runtime = Runtime(config_path=config_path, data_root=data_root)
+
+    @asynccontextmanager
+    async def lifespan(server: Any):
+        await runtime.start()
+        try:
+            yield {}
+        finally:
+            await runtime.close()
+
+    server = FastMCP("KeepM", lifespan=lifespan)
+    read_only = {"readOnlyHint": True, "openWorldHint": False}
+    local_write = {
+        "readOnlyHint": False,
+        "destructiveHint": False,
+        "openWorldHint": False,
+    }
+
+    @server.tool(annotations=local_write)
+    async def memory_configure(
+        vault_path: str,
+        memory_dir: str = "Memory",
+        stale_after_seconds: int = 30,
+        language: MemoryLanguage | None = None,
+    ) -> dict[str, Any]:
+        """Configure the one local Obsidian Vault used by KeepM."""
+        return await runtime.configure(
+            vault_path=vault_path,
+            memory_dir=memory_dir,
+            stale_after_seconds=stale_after_seconds,
+            language=language,
+        )
+
+    @server.tool(annotations=read_only)
+    async def memory_status(
+        ctx: Context, workspace_path: str | None = None
+    ) -> dict[str, Any]:
+        """Report configuration, project, watcher, database, and sync state."""
+        workspace = None
+        if runtime.service is not None:
+            workspace = await _workspace_from_context(ctx, workspace_path)
+        return runtime.status(workspace)
+
+    @server.tool(annotations=read_only)
+    async def memory_context(
+        ctx: Context,
+        query: str,
+        workspace_path: str | None = None,
+        limit: int = 10,
+        max_chars: int = 6_000,
+    ) -> dict[str, Any]:
+        """Return bounded pinned and query-relevant project/global memories."""
+        service = runtime.require_service()
+        workspace = await _workspace_from_context(ctx, workspace_path)
+        await service.ensure_fresh()
+        items = service.context(
+            workspace=workspace, query=query, limit=limit, max_chars=max_chars
+        )
+        output_chars = sum(len(item.snippet) for item in items)
+        return {
+            "items": _jsonable(items),
+            "count": len(items),
+            "output_chars": output_chars,
+            "truncated": len(items) >= limit or output_chars >= max_chars,
+            "language": service.config.language,
+        }
+
+    @server.tool(annotations=read_only)
+    async def memory_search(
+        ctx: Context,
+        query: str,
+        workspace_path: str | None = None,
+        scope: str | None = None,
+        limit: int = 10,
+    ) -> dict[str, Any]:
+        """Search the active project and global Markdown memory index."""
+        service = runtime.require_service()
+        workspace = await _workspace_from_context(ctx, workspace_path)
+        await service.ensure_fresh()
+        items = service.search(
+            workspace=workspace, query=query, scope=scope, limit=limit
+        )
+        return {"items": _jsonable(items), "count": len(items), "limit": limit}
+
+    @server.tool(annotations=read_only)
+    async def memory_read(
+        ctx: Context,
+        identifier: str,
+        workspace_path: str | None = None,
+        max_chars: int = 20_000,
+    ) -> dict[str, Any]:
+        """Read one authoritative Markdown memory by ID, name, or safe path."""
+        service = runtime.require_service()
+        workspace = await _workspace_from_context(ctx, workspace_path)
+        await service.ensure_fresh()
+        document, checksum = service.read_with_checksum(
+            workspace=workspace, identifier=identifier
+        )
+        return _document_result(document, checksum, max_chars=max_chars)
+
+    @server.tool(annotations=local_write)
+    async def memory_write(
+        ctx: Context,
+        name: str,
+        description: str,
+        memory_type: str,
+        body: str,
+        scope: str = "project",
+        tags: list[str] | None = None,
+        pinned: bool = False,
+        workspace_path: str | None = None,
+        expected_checksum: str | None = None,
+    ) -> dict[str, Any]:
+        """Create or update an atomic Markdown memory and its derived index."""
+        service = runtime.require_service()
+        workspace = await _workspace_from_context(ctx, workspace_path)
+        await service.ensure_fresh()
+        document = service.write(
+            workspace=workspace,
+            scope=scope,
+            name=name,
+            description=description,
+            memory_type=memory_type,
+            body=body,
+            tags=tuple(tags or ()),
+            pinned=pinned,
+            expected_checksum=expected_checksum,
+        )
+        record = service.database.get_by_id(document.id)
+        if record is None:
+            raise RuntimeError("memory was written but is absent from the index")
+        return _document_result(document, record.checksum, max_chars=2_000)
+
+    @server.tool(
+        annotations={
+            "readOnlyHint": False,
+            "destructiveHint": True,
+            "openWorldHint": False,
+        }
+    )
+    async def memory_delete(
+        ctx: Context,
+        identifier: str,
+        workspace_path: str | None = None,
+    ) -> dict[str, Any]:
+        """Soft-delete one memory into Memory/.trash."""
+        service = runtime.require_service()
+        workspace = await _workspace_from_context(ctx, workspace_path)
+        await service.ensure_fresh()
+        trashed = service.delete(workspace=workspace, identifier=identifier)
+        return {"deleted": True, "trash_path": str(trashed)}
+
+    @server.tool(annotations=read_only)
+    async def memory_links(
+        ctx: Context,
+        identifier: str,
+        direction: str = "both",
+        workspace_path: str | None = None,
+    ) -> dict[str, Any]:
+        """Return forward WikiLinks and/or backlinks for one memory."""
+        service = runtime.require_service()
+        workspace = await _workspace_from_context(ctx, workspace_path)
+        await service.ensure_fresh()
+        result = _jsonable(
+            service.links(
+                workspace=workspace, identifier=identifier, direction=direction
+            )
+        )
+        result["forward_count"] = len(result.get("forward", []))
+        result["backlink_count"] = len(result.get("backlinks", []))
+        result["count"] = result["forward_count"] + result["backlink_count"]
+        return result
+
+    @server.tool(annotations=local_write)
+    async def memory_sync(full: bool = False) -> dict[str, Any]:
+        """Run quick checksum reconciliation or a full derived-index rebuild."""
+        report = runtime.require_service().sync(full=full)
+        payload = _jsonable(report)
+        payload["error_count"] = len(payload["errors"])
+        return {"report": payload}
+
+    @server.tool(annotations=read_only)
+    async def memory_doctor() -> dict[str, Any]:
+        """Diagnose Markdown and SQLite consistency without modifying memories."""
+        findings = runtime.require_service().doctor()
+        return {"findings": _jsonable(findings), "count": len(findings)}
+
+    setattr(server, "_keepm_runtime", runtime)
+    return server
+
+
+try:
+    mcp = create_mcp()
+except RuntimeError:
+    mcp = None
+
+
+def main(argv: Sequence[str] | None = None) -> None:
+    parser = argparse.ArgumentParser(
+        prog="keepm",
+        description="Run the KeepM local Markdown memory MCP over stdio.",
+    )
+    parser.add_argument("--version", action="version", version=f"keepm {__version__}")
+    parser.parse_args(argv)
+    if mcp is None:
+        parser.error(
+            "FastMCP is not installed; install KeepM with its declared dependencies"
+        )
+    mcp.run(transport="stdio")
+
+
+if __name__ == "__main__":
+    main()