basicmachines-co-basic-memory/src/basic_memory/mcp/server.py

"""
Basic Memory FastMCP server.
"""

import asyncio
import time
from contextlib import asynccontextmanager

from fastmcp import FastMCP
from loguru import logger
from sqlalchemy import text
from sqlalchemy.ext.asyncio import async_sessionmaker, AsyncSession

from basic_memory import db
from basic_memory.cli.auth import CLIAuth
from basic_memory.config import BasicMemoryConfig
from basic_memory.db import (
    scoped_session,
    _needs_semantic_embedding_backfill,
    _run_semantic_embedding_backfill,
)
from basic_memory.mcp.container import McpContainer, set_container
from basic_memory.services.initialization import initialize_app


async def _log_embedding_status(session_maker: async_sessionmaker[AsyncSession]) -> None:
    """Log a clear summary of semantic embedding status at startup."""
    try:
        async with scoped_session(session_maker) as session:
            entity_count = (
                await session.execute(text("SELECT COUNT(*) FROM entity"))
            ).scalar() or 0
            chunk_count = (
                await session.execute(text("SELECT COUNT(*) FROM search_vector_chunks"))
            ).scalar() or 0
            embedding_count = (
                await session.execute(text("SELECT COUNT(*) FROM search_vector_embeddings_rowids"))
            ).scalar() or 0

        if entity_count == 0:
            logger.info("Semantic embeddings: no entities yet")
        elif embedding_count == 0:
            logger.warning(
                f"Semantic embeddings: EMPTY — {entity_count} entities have no embeddings. "
                "Backfill running in background..."
            )
        else:
            logger.info(
                f"Semantic embeddings: {embedding_count} embeddings "
                f"across {chunk_count} chunks for {entity_count} entities"
            )
    except Exception as exc:
        logger.debug(f"Could not check embedding status at startup: {exc}")


async def _background_embedding_backfill(
    config: BasicMemoryConfig,
    session_maker: async_sessionmaker[AsyncSession],
) -> None:
    """Run semantic embedding backfill in the background without blocking startup."""
    try:
        if await _needs_semantic_embedding_backfill(config, session_maker):
            logger.info("Background embedding backfill starting...")
            await _run_semantic_embedding_backfill(config, session_maker)
            await _log_embedding_status(session_maker)
    except Exception as exc:
        logger.error(f"Background embedding backfill failed: {exc}")


@asynccontextmanager
async def lifespan(app: FastMCP):
    """Lifecycle manager for the MCP server.

    Handles:
    - Database initialization and migrations
    - File sync via SyncCoordinator (if enabled and not in cloud mode)
    - Proper cleanup on shutdown
    """
    # --- Composition Root ---
    # Create container and read config (single point of config access)
    container = McpContainer.create()
    set_container(container)

    config = container.config
    logger.info(f"Starting Basic Memory MCP server (mode={container.mode.name})")
    logger.info(
        f"Config: database_backend={config.database_backend.value}, "
        f"semantic_search_enabled={config.semantic_search_enabled}, "
        f"default_project={config.default_project}"
    )
    if config.semantic_search_enabled:
        logger.info(
            f"Semantic search: provider={config.semantic_embedding_provider}, "
            f"model={config.semantic_embedding_model}, "
            f"dimensions={config.semantic_embedding_dimensions or 'auto'}, "
            f"batch_size={config.semantic_embedding_batch_size}"
        )

    # Log configured projects with their routing mode
    for name, entry in config.projects.items():
        default = " (default)" if name == config.default_project else ""
        logger.info(f"Project: {name} -> {entry.path} [mode={entry.mode.value}]{default}")

    # Check cloud auth status (local file check, no network call)
    auth = CLIAuth(client_id=config.cloud_client_id, authkit_domain=config.cloud_domain)
    tokens = auth.load_tokens()
    if tokens is not None:
        if not auth.is_token_valid(tokens):
            expires_at = tokens.get("expires_at", 0)
            expired_ago = int(time.time() - expires_at)
            logger.warning(f"Cloud token expired {expired_ago}s ago - may need 'bm cloud login'")
        else:
            logger.info("Cloud: authenticated (OAuth token valid)")

    if config.cloud_api_key:
        logger.info("Cloud: API key configured")

    # Track if we created the engine (vs test fixtures providing it)
    # This prevents disposing an engine provided by test fixtures when
    # multiple Client connections are made in the same test
    engine_was_none = db._engine is None

    # Initialize app (runs migrations, reconciles projects)
    await initialize_app(container.config)

    # Log embedding status so it's easy to spot in the logs
    backfill_task: asyncio.Task | None = None  # type: ignore[type-arg]
    if config.semantic_search_enabled and db._session_maker is not None:
        await _log_embedding_status(db._session_maker)
        # Launch backfill in background so MCP server is ready immediately
        backfill_task = asyncio.create_task(
            _background_embedding_backfill(config, db._session_maker),
            name="embedding-backfill",
        )

    # Create and start sync coordinator (lifecycle centralized in coordinator)
    sync_coordinator = container.create_sync_coordinator()
    await sync_coordinator.start()

    try:
        yield
    finally:
        # Shutdown - coordinator handles clean task cancellation
        logger.debug("Shutting down Basic Memory MCP server")

        # Cancel embedding backfill if still running
        if backfill_task is not None and not backfill_task.done():
            backfill_task.cancel()
            try:
                await backfill_task
            except asyncio.CancelledError:
                logger.info("Background embedding backfill cancelled during shutdown")

        await sync_coordinator.stop()

        # Only shutdown DB if we created it (not if test fixture provided it)
        if engine_was_none:
            await db.shutdown_db()
            logger.debug("Database connections closed")
        else:  # pragma: no cover
            logger.debug("Skipping DB shutdown - engine provided externally")


mcp = FastMCP(
    name="Basic Memory",
    lifespan=lifespan,
)