fleet-memory/hindsight-api-slim/hindsight_api/api/mcp.py

"""Hindsight MCP Server implementation using FastMCP (HTTP transport)."""

import json
import logging
import os
from contextvars import ContextVar

from fastmcp import FastMCP

from hindsight_api import MemoryEngine
from hindsight_api.config import _get_raw_config
from hindsight_api.engine.memory_engine import _current_schema
from hindsight_api.extensions import MCPExtension, load_extension
from hindsight_api.extensions.tenant import AuthenticationError
from hindsight_api.mcp_tools import MCPToolsConfig, register_mcp_tools
from hindsight_api.models import RequestContext

# All tools available in the system (explicit list — no wildcards)
_ALL_TOOLS: frozenset[str] = frozenset(
    {
        "retain",
        "recall",
        "reflect",
        "list_banks",
        "create_bank",
        "list_mental_models",
        "get_mental_model",
        "create_mental_model",
        "update_mental_model",
        "delete_mental_model",
        "refresh_mental_model",
        "list_directives",
        "create_directive",
        "delete_directive",
        "list_memories",
        "get_memory",
        "delete_memory",
        "list_documents",
        "get_document",
        "delete_document",
        "list_operations",
        "get_operation",
        "cancel_operation",
        "list_tags",
        "get_bank",
        "get_bank_stats",
        "update_bank",
        "delete_bank",
        "clear_memories",
    }
)

# Configure logging from HINDSIGHT_API_LOG_LEVEL environment variable
_log_level_str = os.environ.get("HINDSIGHT_API_LOG_LEVEL", "info").lower()
_log_level_map = {
    "critical": logging.CRITICAL,
    "error": logging.ERROR,
    "warning": logging.WARNING,
    "info": logging.INFO,
    "debug": logging.DEBUG,
    "trace": logging.DEBUG,
}
logging.basicConfig(
    level=_log_level_map.get(_log_level_str, logging.INFO),
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
)
logger = logging.getLogger(__name__)

# Default bank_id from environment variable
DEFAULT_BANK_ID = os.environ.get("HINDSIGHT_MCP_BANK_ID", "default")

# Legacy MCP authentication token (for backwards compatibility)
# If set, this token is checked first before TenantExtension auth
MCP_AUTH_TOKEN = os.environ.get("HINDSIGHT_API_MCP_AUTH_TOKEN")

# Context variable to hold the current bank_id
_current_bank_id: ContextVar[str | None] = ContextVar("current_bank_id", default=None)

# Context variable to hold the current API key (for tenant auth propagation)
_current_api_key: ContextVar[str | None] = ContextVar("current_api_key", default=None)

# Context variables for tenant_id and api_key_id (set by authenticate, used by usage metering)
_current_tenant_id: ContextVar[str | None] = ContextVar("current_tenant_id", default=None)
_current_api_key_id: ContextVar[str | None] = ContextVar("current_api_key_id", default=None)


def get_current_bank_id() -> str | None:
    """Get the current bank_id from context."""
    return _current_bank_id.get()


def get_current_api_key() -> str | None:
    """Get the current API key from context."""
    return _current_api_key.get()


def get_current_tenant_id() -> str | None:
    """Get the current tenant_id from context."""
    return _current_tenant_id.get()


def get_current_api_key_id() -> str | None:
    """Get the current api_key_id from context."""
    return _current_api_key_id.get()


def create_mcp_server(memory: MemoryEngine, multi_bank: bool = True) -> FastMCP:
    """
    Create and configure the Hindsight MCP server.

    Args:
        memory: MemoryEngine instance (required)
        multi_bank: If True, expose all tools with bank_id parameters (default).
                   If False, only expose bank-scoped tools without bank_id parameters.

    Returns:
        Configured FastMCP server instance
    """
    mcp = FastMCP("hindsight-mcp-server")

    global_config = _get_raw_config()

    # Tools available for this mode (multi-bank exposes all tools; single-bank excludes bank-management tools)
    _SINGLE_BANK_TOOLS: frozenset[str] = frozenset(
        {
            "retain",
            "recall",
            "reflect",
            "list_mental_models",
            "get_mental_model",
            "create_mental_model",
            "update_mental_model",
            "delete_mental_model",
            "refresh_mental_model",
            "list_directives",
            "create_directive",
            "delete_directive",
            "list_memories",
            "get_memory",
            "delete_memory",
            "list_documents",
            "get_document",
            "delete_document",
            "list_operations",
            "get_operation",
            "cancel_operation",
            "list_tags",
            "get_bank",
            "update_bank",
            "delete_bank",
            "clear_memories",
        }
    )
    base_tools: frozenset[str] | None = None if multi_bank else _SINGLE_BANK_TOOLS

    # Apply global mcp_enabled_tools filter (env-level allowlist)
    if global_config.mcp_enabled_tools is not None:
        allowed = frozenset(global_config.mcp_enabled_tools)
        base_tools = (base_tools if base_tools is not None else _ALL_TOOLS) & allowed

    # Configure and register tools using shared module
    config = MCPToolsConfig(
        bank_id_resolver=get_current_bank_id,
        api_key_resolver=get_current_api_key,  # Propagate API key for tenant auth
        tenant_id_resolver=get_current_tenant_id,  # Propagate tenant_id for usage metering
        api_key_id_resolver=get_current_api_key_id,  # Propagate api_key_id for usage metering
        include_bank_id_param=multi_bank,
        tools=base_tools,
    )

    register_mcp_tools(mcp, memory, config)

    # Load and register additional tools from MCP extension if configured
    mcp_extension = load_extension("MCP", MCPExtension)
    if mcp_extension:
        logger.info(f"Loading MCP extension: {mcp_extension.__class__.__name__}")
        mcp_extension.register_tools(mcp, memory)

    # Make all tools tolerant of extra arguments from LLMs (e.g., "explanation")
    _make_tools_tolerant(mcp)

    return mcp


def _make_tools_tolerant(mcp: FastMCP) -> None:
    """Wrap all tool run methods to strip unknown arguments before validation.

    LLMs frequently add extra fields like "explanation" or "reasoning" to tool calls.
    FastMCP's Pydantic TypeAdapter rejects these with "Unexpected keyword argument".
    This wraps each tool's run() to filter arguments to only known parameters.
    """
    try:
        for name, tool in mcp._tool_manager._tools.items():
            if hasattr(tool, "parameters") and tool.parameters:
                allowed = set(tool.parameters.get("properties", {}).keys())
                original_run = tool.run

                async def _tolerant_run(arguments, _allowed=allowed, _orig=original_run):
                    extra_keys = set(arguments.keys()) - _allowed
                    if extra_keys:
                        logger.debug(f"Stripping unknown arguments from tool call: {extra_keys}")
                        arguments = {k: v for k, v in arguments.items() if k in _allowed}
                    return await _orig(arguments)

                # FunctionTool is a Pydantic model with extra='forbid', so use
                # object.__setattr__ to bypass Pydantic's setter validation.
                object.__setattr__(tool, "run", _tolerant_run)
    except (AttributeError, KeyError) as e:
        logger.warning(f"Could not make tools tolerant of extra arguments: {e}")


class MCPMiddleware:
    """ASGI middleware that intercepts MCP requests and routes to appropriate MCP server.

    This middleware wraps the main FastAPI app and intercepts requests matching the
    configured prefix (default: /mcp). Non-MCP requests pass through to the inner app.

    Authentication:
        1. If HINDSIGHT_API_MCP_AUTH_TOKEN is set (legacy), validates against that token
        2. Otherwise, uses TenantExtension.authenticate_mcp() from the MemoryEngine
           - DefaultTenantExtension: no auth required (local dev)
           - ApiKeyTenantExtension: validates against env var

    Two modes based on URL structure:

    1. Multi-bank mode (for /mcp/ root endpoint):
       - Exposes all tools: retain, recall, reflect, list_banks, create_bank
       - All tools include optional bank_id parameter for cross-bank operations
       - Bank ID from: X-Bank-Id header or HINDSIGHT_MCP_BANK_ID env var

    2. Single-bank mode (for /mcp/{bank_id}/ endpoints):
       - Exposes bank-scoped tools only: retain, recall, reflect
       - No bank_id parameter (comes from URL)
       - No bank management tools (list_banks, create_bank)
       - Recommended for agent isolation

    Bank ID resolution priority:
        1. URL path (e.g., /mcp/{bank_id}/) → single-bank mode
        2. X-Bank-Id header → multi-bank mode
        3. HINDSIGHT_MCP_BANK_ID env var → multi-bank mode (default: "default")

    Examples:
        # Single-bank mode (recommended for agent isolation)
        claude mcp add --transport http my-agent http://localhost:8888/mcp/my-agent-bank/ \\
            --header "Authorization: Bearer <token>"

        # Multi-bank mode (for cross-bank operations)
        claude mcp add --transport http hindsight http://localhost:8888/mcp \\
            --header "X-Bank-Id: my-bank" --header "Authorization: Bearer <token>"
    """

    def __init__(
        self,
        app,
        memory: MemoryEngine,
        prefix: str = "/mcp",
        multi_bank_app=None,
        single_bank_app=None,
        multi_bank_server=None,
        single_bank_server=None,
    ):
        self.app = app
        self.prefix = prefix
        self.memory = memory
        self.tenant_extension = memory._tenant_extension

        if multi_bank_app and single_bank_app:
            # Pre-created servers (used when called via add_middleware from create_app)
            self.multi_bank_app = multi_bank_app
            self.single_bank_app = single_bank_app
            self.multi_bank_server = multi_bank_server
            self.single_bank_server = single_bank_server
        else:
            # Create servers internally (for direct construction / tests)
            self.multi_bank_server = create_mcp_server(memory, multi_bank=True)
            self.multi_bank_app = self.multi_bank_server.http_app(path="/", stateless_http=True)
            self.single_bank_server = create_mcp_server(memory, multi_bank=False)
            self.single_bank_app = self.single_bank_server.http_app(path="/", stateless_http=True)

    def _get_header(self, scope: dict, name: str) -> str | None:
        """Extract a header value from ASGI scope."""
        name_lower = name.lower().encode()
        for header_name, header_value in scope.get("headers", []):
            if header_name.lower() == name_lower:
                return header_value.decode()
        return None

    async def __call__(self, scope, receive, send):
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        path = scope.get("path", "")

        # Check if this is an MCP request (matches prefix)
        if not (path == self.prefix or path.startswith(self.prefix + "/")):
            # Not an MCP request — pass through to the inner app
            await self.app(scope, receive, send)
            return

        # Strip prefix from path
        path = path[len(self.prefix) :] or "/"

        # Extract auth token from header (for tenant auth propagation)
        auth_header = self._get_header(scope, "Authorization")
        auth_token: str | None = None
        if auth_header:
            # Support both "Bearer <token>" and direct token
            auth_token = auth_header[7:].strip() if auth_header.startswith("Bearer ") else auth_header.strip()

        # Authenticate: check legacy MCP_AUTH_TOKEN first, then TenantExtension
        tenant_context = None
        auth_tenant_id: str | None = None
        auth_api_key_id: str | None = None
        if MCP_AUTH_TOKEN:
            # Legacy authentication mode - validate against static token
            if not auth_token:
                await self._send_error(send, 401, "Authorization header required")
                return
            if auth_token != MCP_AUTH_TOKEN:
                await self._send_error(send, 401, "Invalid authentication token")
                return
            # Legacy mode doesn't use tenant schemas
            tenant_context = None
        else:
            # Use TenantExtension.authenticate_mcp() for auth
            try:
                auth_context = RequestContext(api_key=auth_token)
                tenant_context = await self.tenant_extension.authenticate_mcp(auth_context)
                # Capture tenant_id and api_key_id set by authenticate() for usage metering
                auth_tenant_id = auth_context.tenant_id
                auth_api_key_id = auth_context.api_key_id
            except AuthenticationError as e:
                await self._send_error(send, 401, str(e), extra_headers=e.headers)
                return

        # Set schema from tenant context so downstream DB queries use the correct schema
        schema_token = (
            _current_schema.set(tenant_context.schema_name) if tenant_context and tenant_context.schema_name else None
        )

        # Resolve bank_id: path takes priority over header.
        # Path = user's explicit connection endpoint (e.g., /mcp/my-bank/).
        # X-Bank-Id header = per-request override for multi-bank mode only.
        bank_id = None
        bank_id_from_path = False
        new_path = path

        # First, try to extract from path: /{bank_id}/...
        if path.startswith("/") and len(path) > 1:
            parts = path[1:].split("/", 1)
            if parts[0]:
                bank_id = parts[0]
                bank_id_from_path = True
                new_path = "/" + parts[1] if len(parts) > 1 else "/"

        # If no path-based bank_id, try X-Bank-Id header (multi-bank mode)
        if not bank_id:
            bank_id = self._get_header(scope, "X-Bank-Id")

        # Fall back to default bank_id
        if not bank_id:
            bank_id = DEFAULT_BANK_ID
            logger.debug(f"Using default bank_id: {bank_id}")

        # Select the appropriate MCP app based on how bank_id was provided:
        # - Path-based bank_id → single-bank app (no bank_id param, scoped tools)
        # - Header/env bank_id → multi-bank app (bank_id param, all tools)
        target_app = self.single_bank_app if bank_id_from_path else self.multi_bank_app

        # Set bank_id, api_key, tenant_id, and api_key_id context
        bank_id_token = _current_bank_id.set(bank_id)
        # Store the auth token for tenant extension to validate
        api_key_token = _current_api_key.set(auth_token) if auth_token else None
        # Store tenant_id and api_key_id from authentication for usage metering
        tenant_id_token = _current_tenant_id.set(auth_tenant_id) if auth_tenant_id else None
        api_key_id_token = _current_api_key_id.set(auth_api_key_id) if auth_api_key_id else None
        try:
            new_scope = scope.copy()
            new_scope["path"] = new_path
            # Clear root_path since we're passing directly to the app
            new_scope["root_path"] = ""

            # Wrap send to rewrite the SSE endpoint URL to include bank_id if using path-based routing.
            # Only rewrite SSE (text/event-stream) responses to avoid corrupting tool results
            # that might contain the literal string "data: /messages".
            is_sse_response = False

            async def send_wrapper(message):
                nonlocal is_sse_response
                if message["type"] == "http.response.start":
                    for header_name, header_value in message.get("headers", []):
                        if header_name == b"content-type" and b"text/event-stream" in header_value:
                            is_sse_response = True
                            break
                if message["type"] == "http.response.body" and bank_id_from_path and is_sse_response:
                    body = message.get("body", b"")
                    if body and b"/messages" in body:
                        # Rewrite /messages to /{bank_id}/messages in SSE endpoint event
                        body = body.replace(b"data: /messages", f"data: /{bank_id}/messages".encode())
                        message = {**message, "body": body}
                await send(message)

            await target_app(new_scope, receive, send_wrapper)
        finally:
            _current_bank_id.reset(bank_id_token)
            if api_key_token is not None:
                _current_api_key.reset(api_key_token)
            if tenant_id_token is not None:
                _current_tenant_id.reset(tenant_id_token)
            if api_key_id_token is not None:
                _current_api_key_id.reset(api_key_id_token)
            if schema_token is not None:
                _current_schema.reset(schema_token)

    async def _send_error(self, send, status: int, message: str, extra_headers: dict[str, str] | None = None):
        """Send an error response."""
        body = json.dumps({"error": message}).encode()
        headers = [(b"content-type", b"application/json")]
        for key, value in (extra_headers or {}).items():
            headers.append((key.encode(), value.encode()))
        await send(
            {
                "type": "http.response.start",
                "status": status,
                "headers": headers,
            }
        )
        await send(
            {
                "type": "http.response.body",
                "body": body,
            }
        )


def create_mcp_servers(memory: MemoryEngine):
    """Create multi-bank and single-bank MCP servers and their Starlette apps.

    Returns the servers and apps separately so lifespans can be chained before
    the middleware wraps the main app.

    Returns:
        Tuple of (multi_bank_server, single_bank_server, multi_bank_app, single_bank_app)
    """
    multi_bank_server = create_mcp_server(memory, multi_bank=True)
    multi_bank_app = multi_bank_server.http_app(path="/", stateless_http=True)

    single_bank_server = create_mcp_server(memory, multi_bank=False)
    single_bank_app = single_bank_server.http_app(path="/", stateless_http=True)

    return multi_bank_server, single_bank_server, multi_bank_app, single_bank_app