chopratejas--headroom/headroom/memory/mcp_server.py

"""Headroom Memory MCP Server.

A stdio MCP server that exposes headroom's memory backend as tools
that Codex (or any MCP-compatible client) can call natively.

Tools:
    memory_search  — semantic search across stored memories
    memory_save    — persist a new fact/decision/convention

Design:
    - Embedder is pre-loaded at startup (no cold-start on first query)
    - On startup, any memories missing vector embeddings are re-indexed
      (fixes interop gap when memories were saved via a different path)
    - Save always generates embeddings inline

Usage:
    # Standalone (for testing):
    python -m headroom.memory.mcp_server --db /path/to/.headroom/memory.db

    # Registered in Codex config.toml (done by `headroom wrap codex --memory`):
    [mcp_servers.headroom_memory]
    command = "python"
    args = ["-m", "headroom.memory.mcp_server", "--user", "alice"]
    # When --db is omitted, the server resolves .headroom/memory.db from cwd.
"""

from __future__ import annotations

import argparse
import asyncio
import logging
import os
import sys
from pathlib import Path
from typing import Any

from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import TextContent, Tool

from headroom.memory.backends.local import LocalBackend, LocalBackendConfig

logger = logging.getLogger("headroom.memory.mcp")

# ---------------------------------------------------------------------------
# Tool definitions
# ---------------------------------------------------------------------------

_TOOLS = [
    Tool(
        name="memory_search",
        description=(
            "Search persistent memory for relevant knowledge from prior sessions. "
            "Use this for questions about architecture, conventions, prior decisions, "
            "project context, user preferences, org info, codenames, debugging history, "
            "or anything that might have been discussed before."
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Natural-language search query.",
                },
                "top_k": {
                    "type": "integer",
                    "description": "Max results to return (default 10).",
                    "default": 10,
                },
            },
            "required": ["query"],
        },
    ),
    Tool(
        name="memory_save",
        description=(
            "Save information to persistent memory for future sessions. "
            "Use this for decisions, conventions, architecture context, "
            "user preferences, project facts, or anything worth remembering.\n\n"
            "IMPORTANT: Break information into atomic facts — one fact per "
            "entry in the 'facts' array. Each fact should be a single, "
            "self-contained statement that answers one question. "
            "Do NOT combine multiple facts into one string.\n\n"
            "Good:  facts: ['Repo owner is Tejas C.', 'User prefers dark mode']\n"
            "Bad:   facts: ['Repo owner is Tejas C. Prefers dark mode.']"
        ),
        inputSchema={
            "type": "object",
            "properties": {
                "facts": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": (
                        "Array of atomic facts to save. Each entry should be "
                        "one self-contained fact. The system stores and indexes "
                        "each fact separately for precise retrieval."
                    ),
                },
                "importance": {
                    "type": "number",
                    "description": "0.0 (low) to 1.0 (critical). Default 0.7.",
                    "default": 0.7,
                },
            },
            "required": [],
        },
    ),
]


# ---------------------------------------------------------------------------
# Startup: pre-load embedder + re-index unembedded memories
# ---------------------------------------------------------------------------


async def _warm_up_backend(backend: LocalBackend, user_id: str) -> None:
    """Pre-load the embedder and re-index memories that lack embeddings.

    Memories saved via other paths (e.g. Claude Code proxy direct SQL)
    may exist in the store but have no vector embeddings.  This scans
    for those and re-indexes them so vector search works across agents.
    """
    await backend._ensure_initialized()
    hm = backend._hierarchical_memory
    if hm is None:
        return

    # Force-load the embedder now (not lazily on first search)
    _dummy = await hm._embedder.embed("warmup")
    logger.info("Memory MCP: embedder pre-loaded")

    # Ensure ALL memories are in the vector index.
    # Memories saved via other agents (Claude Code proxy, direct SQL) may
    # exist in the store but not be indexed — re-embed and index them all.
    all_memories = await backend.get_user_memories(user_id, limit=500)
    if not all_memories:
        return

    memories_missing_embeddings = [mem for mem in all_memories if mem.embedding is None]
    if memories_missing_embeddings:
        embeddings = await hm._embedder.embed_batch(
            [mem.content for mem in memories_missing_embeddings]
        )
        for mem, embedding in zip(memories_missing_embeddings, embeddings):
            mem.embedding = embedding
        await hm._store.save_batch(memories_missing_embeddings)

    indexed = await hm._vector_index.index_batch(all_memories)
    logger.info(f"Memory MCP: indexed {indexed} memories into vector store")


# ---------------------------------------------------------------------------
# MCP Server
# ---------------------------------------------------------------------------


def create_memory_server(db_path: str, user_id: str = "default") -> Server:
    """Create an MCP server backed by headroom's local memory."""

    server = Server("headroom-memory")
    _backend: LocalBackend | None = None
    _init_task: asyncio.Task | None = None

    async def _init_backend() -> LocalBackend:
        """Initialize backend with ONNX embedder (fast, no PyTorch)."""
        nonlocal _backend
        config = LocalBackendConfig(db_path=db_path, embedder_backend="onnx")
        _backend = LocalBackend(config)
        await _warm_up_backend(_backend, user_id)
        logger.info(f"Memory MCP: ready (db={db_path}, user={user_id})")
        return _backend

    async def _get_backend() -> LocalBackend:
        nonlocal _backend, _init_task
        if _backend is not None:
            return _backend
        # Wait for background init if it's running
        if _init_task is not None:
            await _init_task
            return _backend  # type: ignore[return-value]
        # Fallback: init inline (shouldn't normally happen)
        return await _init_backend()

    @server.list_tools()
    async def list_tools() -> list[Tool]:
        # Kick off background init on first list_tools (called at MCP handshake)
        nonlocal _init_task
        if _backend is None and _init_task is None:
            _init_task = asyncio.create_task(_init_backend())
        return _TOOLS

    @server.call_tool()
    async def call_tool(name: str, arguments: dict) -> list[TextContent]:
        backend = await _get_backend()

        if name == "memory_search":
            return await _handle_search(backend, arguments, user_id)
        elif name == "memory_save":
            return await _handle_save(backend, arguments, user_id)

        return [TextContent(type="text", text=f"Unknown tool: {name}")]

    return server


async def _handle_search(
    backend: LocalBackend, arguments: dict[str, Any], user_id: str
) -> list[TextContent]:
    query = arguments.get("query", "")
    top_k = arguments.get("top_k", 10)

    if not query:
        return [TextContent(type="text", text="Error: query is required")]

    try:
        # Over-fetch to compensate for filtering out superseded memories
        results = await backend.search_memories(
            query=query,
            user_id=user_id,
            top_k=top_k * 3,
            include_related=True,
        )

        if not results:
            return [TextContent(type="text", text="No memories found.")]

        # Filter out superseded memories — only return current/active ones.
        # Re-check the store because in-memory HNSW metadata may be stale.
        active_results = []
        for r in results:
            if getattr(r.memory, "superseded_by", None):
                continue
            # Double-check against the store for recently superseded memories
            try:
                stored = await backend.get_memory(r.memory.id)
                if stored and getattr(stored, "superseded_by", None):
                    continue
            except Exception:
                pass
            active_results.append(r)

        if not active_results:
            return [TextContent(type="text", text="No memories found.")]

        # Trim to requested top_k
        active_results = active_results[:top_k]

        try:
            await backend.record_access([r.memory.id for r in active_results])
        except Exception as e:
            # Usage metadata must never make a successful retrieval fail.
            logger.warning(f"Memory MCP: failed to record access: {e}")

        lines = []
        for i, r in enumerate(active_results, 1):
            score = f"{r.score:.2f}" if hasattr(r, "score") else "?"
            lines.append(f"{i}. [relevance={score}] {r.memory.content}")
            if hasattr(r, "related_entities") and r.related_entities:
                lines.append(f"   Related: {', '.join(r.related_entities[:3])}")

        return [TextContent(type="text", text="\n".join(lines))]
    except Exception as e:
        logger.error(f"memory_search failed: {e}")
        return [TextContent(type="text", text=f"Search error: {e}")]


# Similarity threshold for auto-supersession: if a new memory is this
# similar to an existing one, it replaces (supersedes) the old one.
_SUPERSEDE_SIMILARITY = 0.70


async def _handle_save(
    backend: LocalBackend, arguments: dict[str, Any], user_id: str
) -> list[TextContent]:
    facts = arguments.get("facts", [])
    importance = arguments.get("importance", 0.7)

    # Backward compat: accept single "content" string too
    if not facts:
        content = arguments.get("content", "")
        if content:
            facts = [content]

    if not facts:
        return [TextContent(type="text", text="Error: facts array is required")]

    try:
        saved = 0
        superseded = 0
        results_lines: list[str] = []

        for fact in facts:
            fact = fact.strip()
            if not fact:
                continue

            # Check for semantically similar existing memory to auto-supersede
            superseded_id: str | None = None
            try:
                existing = await backend.search_memories(
                    query=fact,
                    user_id=user_id,
                    top_k=3,
                )
                for r in existing:
                    if getattr(r.memory, "superseded_by", None):
                        continue
                    if r.score >= _SUPERSEDE_SIMILARITY:
                        superseded_id = r.memory.id
                        logger.info(
                            f"Memory MCP: auto-superseding [{r.memory.id[:8]}] "
                            f"(similarity={r.score:.2f}): {r.memory.content[:60]}"
                        )
                        break
            except Exception:
                pass

            if superseded_id:
                memory = await backend.update_memory(
                    memory_id=superseded_id,
                    new_content=fact,
                )
                results_lines.append(
                    f"  updated [{superseded_id[:8]}→{memory.id[:8]}]: {fact[:60]}"
                )
                superseded += 1
            else:
                memory = await backend.save_memory(
                    content=fact,
                    user_id=user_id,
                    importance=importance,
                )
                results_lines.append(f"  saved [{memory.id[:8]}]: {fact[:60]}")
                saved += 1

        summary = f"Saved {saved} new, updated {superseded} existing ({saved + superseded} total)"
        return [TextContent(type="text", text=summary + "\n" + "\n".join(results_lines))]
    except Exception as e:
        logger.error(f"memory_save failed: {e}")
        return [TextContent(type="text", text=f"Save error: {e}")]


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------


async def _run(db_path: str, user_id: str) -> None:
    server = create_memory_server(db_path, user_id)
    async with stdio_server() as (read_stream, write_stream):
        await server.run(read_stream, write_stream, server.create_initialization_options())


def _memory_mcp_startup_context(
    configured_db: str, cwd: Path, db_flag_present: bool
) -> dict[str, str | bool]:
    """Describe the DB path the memory MCP server will try to open."""
    configured_path = Path(configured_db).expanduser()
    resolved_path = (
        configured_path if configured_path.is_absolute() else (cwd / configured_path)
    ).resolve(strict=False)
    active_project_db = (cwd / ".headroom" / "memory.db").resolve(strict=False)
    if not db_flag_present:
        config_source = "cwd-default"
        resolution = "dynamic-cwd"
    else:
        config_source = "cli-flag"
        resolution = "static-cli"
    if resolved_path == active_project_db:
        storage_scope = "active-project"
    elif resolved_path.name == "memory.db":
        storage_scope = "external-memory-db"
    else:
        storage_scope = "custom-db-path"
    path_exists = resolved_path.exists()
    path_readable = path_exists and os.access(resolved_path, os.R_OK)
    return {
        "configured_db": str(configured_path),
        "resolved_db": str(resolved_path),
        "config_source": config_source,
        "cwd": str(cwd),
        "project_root": str(cwd),
        "storage_scope": storage_scope,
        "path_exists": path_exists,
        "path_readable": path_readable,
        "resolution": resolution,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="Headroom Memory MCP Server")
    parser.add_argument(
        "--db",
        default=str(Path.cwd() / ".headroom" / "memory.db"),
        help="Path to memory SQLite database",
    )
    parser.add_argument(
        "--user",
        default=os.environ.get("USER", os.environ.get("USERNAME", "default")),
        help="User ID for memory scoping",
    )
    args = parser.parse_args()

    # Skip HuggingFace model freshness checks — use cached models only.
    # This eliminates 1-2s of HTTP HEAD requests on every startup.
    os.environ.setdefault("HF_HUB_OFFLINE", "1")
    os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")

    # Log to stderr (MCP uses stdout for protocol)
    logging.basicConfig(
        level=logging.INFO,
        stream=sys.stderr,
        format="%(name)s: %(message)s",
    )

    startup = _memory_mcp_startup_context(
        configured_db=args.db,
        cwd=Path.cwd(),
        db_flag_present=any(arg == "--db" or arg.startswith("--db=") for arg in sys.argv[1:]),
    )
    logger.info(
        "Memory MCP startup: configured_db=%s, resolved_db=%s, config_source=%s, "
        "cwd=%s, project_root=%s, storage_scope=%s, path_exists=%s, "
        "path_readable=%s, resolution=%s",
        startup["configured_db"],
        startup["resolved_db"],
        startup["config_source"],
        startup["cwd"],
        startup["project_root"],
        startup["storage_scope"],
        startup["path_exists"],
        startup["path_readable"],
        startup["resolution"],
    )

    asyncio.run(_run(args.db, args.user))


if __name__ == "__main__":
    main()