d8050387e4
* fix(mcp): handle Claude Code GET probe and make stateless_http configurable (#751) Claude Code v2.1.84+ sends a GET to /mcp/ before POST initialize, which fails with 405 (stateless) or 400 (stateful). Intercept sessionless GET requests in MCPMiddleware and return 200 OK so the client proceeds to POST initialize. Also make stateless_http configurable via HINDSIGHT_API_MCP_STATELESS (default: false/stateful) instead of hardcoding true. Closes #751 * docs: add HINDSIGHT_API_MCP_STATELESS to configuration reference
473 lines
20 KiB
Python
473 lines
20 KiB
Python
"""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 _ALL_TOOLS, MCPToolsConfig, register_mcp_tools
|
|
from hindsight_api.models import RequestContext
|
|
|
|
# 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)
|
|
|
|
# Context variable for MCP pre-authentication flag (set when MCP_AUTH_TOKEN validates)
|
|
_current_mcp_authenticated: ContextVar[bool] = ContextVar("current_mcp_authenticated", default=False)
|
|
|
|
|
|
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 get_current_mcp_authenticated() -> bool:
|
|
"""Get whether the request was pre-authenticated by MCP transport auth."""
|
|
return _current_mcp_authenticated.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
|
|
mcp_authenticated_resolver=get_current_mcp_authenticated, # Propagate MCP pre-auth flag
|
|
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)
|
|
global_config = _get_raw_config()
|
|
stateless = global_config.mcp_stateless
|
|
self.multi_bank_server = create_mcp_server(memory, multi_bank=True)
|
|
self.multi_bank_app = self.multi_bank_server.http_app(path="/", stateless_http=stateless)
|
|
self.single_bank_server = create_mcp_server(memory, multi_bank=False)
|
|
self.single_bank_app = self.single_bank_server.http_app(path="/", stateless_http=stateless)
|
|
|
|
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
|
|
|
|
# Handle GET-before-POST gracefully (Claude Code v2.1.84+ sends GET probe before POST initialize).
|
|
# Without a valid Mcp-Session-Id, GET has no meaningful response — return 200 OK so
|
|
# the client proceeds to POST initialize instead of marking the server as failed.
|
|
method = scope.get("method", "")
|
|
if method == "GET":
|
|
session_id = self._get_header(scope, "Mcp-Session-Id")
|
|
if not session_id:
|
|
logger.debug("MCP GET without session ID (client probe) — returning 200 OK")
|
|
await self._send_ok(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
|
|
mcp_pre_authenticated = False
|
|
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: mark as pre-authenticated so tenant extension won't re-validate
|
|
tenant_context = None
|
|
mcp_pre_authenticated = True
|
|
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, api_key_id, and mcp_authenticated 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
|
|
# Store MCP pre-authentication flag to skip tenant re-validation
|
|
mcp_auth_token = _current_mcp_authenticated.set(mcp_pre_authenticated)
|
|
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"] = ""
|
|
|
|
# Ensure Accept header includes required MIME types for MCP SDK.
|
|
# Some clients (e.g., Claude Code) don't send Accept, causing
|
|
# the SDK to reject with 406 Not Acceptable.
|
|
accept_header = self._get_header(new_scope, "accept")
|
|
if not accept_header or "text/event-stream" not in accept_header:
|
|
headers = [(k, v) for k, v in new_scope.get("headers", []) if k.lower() != b"accept"]
|
|
headers.append((b"accept", b"application/json, text/event-stream"))
|
|
new_scope["headers"] = headers
|
|
|
|
# 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)
|
|
_current_mcp_authenticated.reset(mcp_auth_token)
|
|
if schema_token is not None:
|
|
_current_schema.reset(schema_token)
|
|
|
|
async def _send_ok(self, send):
|
|
"""Send a 200 OK response with empty body (used for GET probes without session)."""
|
|
await send(
|
|
{
|
|
"type": "http.response.start",
|
|
"status": 200,
|
|
"headers": [(b"content-type", b"application/json")],
|
|
}
|
|
)
|
|
await send(
|
|
{
|
|
"type": "http.response.body",
|
|
"body": b"{}",
|
|
}
|
|
)
|
|
|
|
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)
|
|
"""
|
|
global_config = _get_raw_config()
|
|
stateless = global_config.mcp_stateless
|
|
|
|
multi_bank_server = create_mcp_server(memory, multi_bank=True)
|
|
multi_bank_app = multi_bank_server.http_app(path="/", stateless_http=stateless)
|
|
|
|
single_bank_server = create_mcp_server(memory, multi_bank=False)
|
|
single_bank_app = single_bank_server.http_app(path="/", stateless_http=stateless)
|
|
|
|
logger.info(f"MCP servers created (stateless_http={stateless})")
|
|
return multi_bank_server, single_bank_server, multi_bank_app, single_bank_app
|