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
from basic_memory import telemetry


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
    with telemetry.operation(
        "mcp.lifecycle.startup",
        entrypoint="mcp",
        mode=container.mode.name.lower(),
        default_project=config.default_project,
    ):
        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
        with telemetry.operation(
            "mcp.lifecycle.shutdown",
            entrypoint="mcp",
            mode=container.mode.name.lower(),
        ):
            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,
)