Nicolò Boschi
6582e26ef9
The knowledge base was reachable only over HTTP, so an MCP client had to fall back to a second integration path to browse or maintain it. Register the seven agent-facing operations as native MCP tools, with the same bank scoping, tenant auth and operation-validator behaviour as the existing tools: get_knowledge_base_tree, search_knowledge_base, get_knowledge_page, create_knowledge_folder, create_knowledge_page, update_knowledge_node, delete_knowledge_node export_knowledge_base stays HTTP/CLI-only — it returns the whole bank as one markdown bundle, which does not belong in an agent's context window. Two places where the MCP surface cannot mirror the HTTP one, both commented at the call site: - MCP arguments cannot express an explicit null, so update_knowledge_node reads parent_id="root" as "move to the top level". Node ids are prefixed kf-/kp-, so the literal cannot collide with a real folder id. - The page refresh trigger is flattened to a single refresh_after_consolidation flag, matching how create/update_mental_model already expose it. It is sent as a patch, so an unstated flag leaves the knowledge-page defaults (delta mode, observation-only) intact — the regression #3506 fixed. get_knowledge_page returns the rendered markdown document once instead of the HTTP body+markdown pair, which would double the tokens for no new information. search_knowledge_base clamps limit instead of rejecting it: an agent that asked for 500 pages wants results, not a 422. Also adds a structural guard that the three hand-maintained tool allowlists (_ALL_TOOLS, register_mcp_tools()'s default set, and the single-bank set in create_mcp_server) agree with what is actually registered — a name added to one but not the others silently drops the tool from the endpoint.
632 lines
27 KiB
Python
632 lines
27 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 import __version__ as HINDSIGHT_VERSION
|
|
from hindsight_api.api.passthrough_headers import collect_passthrough_headers
|
|
from hindsight_api.config import DEFAULT_MCP_RECALL_DESCRIPTION, DEFAULT_MCP_RETAIN_DESCRIPTION, _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)
|
|
|
|
# Context variable for the headers an operator opted into forwarding to extensions
|
|
# (HINDSIGHT_API_EXTENSION_PASSTHROUGH_HEADERS). Defaults to None rather than {}
|
|
# so no single dict is shared as a default across requests.
|
|
_current_extra_headers: ContextVar[dict[str, str] | None] = ContextVar("current_extra_headers", 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 get_current_mcp_authenticated() -> bool:
|
|
"""Get whether the request was pre-authenticated by MCP transport auth."""
|
|
return _current_mcp_authenticated.get()
|
|
|
|
|
|
def get_current_extra_headers() -> dict[str, str]:
|
|
"""Get the allowlisted passthrough headers for the current request.
|
|
|
|
Returns a copy: every RequestContext built during the request owns its dict,
|
|
so extension code mutating one cannot alter what the next tool call sees.
|
|
"""
|
|
return dict(_current_extra_headers.get() or {})
|
|
|
|
|
|
def _build_mcp_tool_descriptions(extra_instructions: str | None) -> tuple[str | None, str | None]:
|
|
"""Return custom retain/recall descriptions when server-level MCP instructions are set."""
|
|
if not isinstance(extra_instructions, str):
|
|
return None, None
|
|
|
|
extra_instructions = extra_instructions.strip()
|
|
if not extra_instructions:
|
|
return None, None
|
|
|
|
suffix = f"\n\nAdditional instructions: {extra_instructions}"
|
|
return DEFAULT_MCP_RETAIN_DESCRIPTION + suffix, DEFAULT_MCP_RECALL_DESCRIPTION + suffix
|
|
|
|
|
|
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", version=HINDSIGHT_VERSION)
|
|
|
|
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",
|
|
"sync_retain",
|
|
"recall",
|
|
"reflect",
|
|
"list_mental_models",
|
|
"get_mental_model",
|
|
"create_mental_model",
|
|
"update_mental_model",
|
|
"delete_mental_model",
|
|
"refresh_mental_model",
|
|
"clear_mental_model",
|
|
"list_directives",
|
|
"create_directive",
|
|
"delete_directive",
|
|
"list_memories",
|
|
"get_memory",
|
|
"update_memory",
|
|
"invalidate_memory",
|
|
"list_documents",
|
|
"get_document",
|
|
"delete_document",
|
|
"list_operations",
|
|
"get_operation",
|
|
"cancel_operation",
|
|
"list_tags",
|
|
"get_bank",
|
|
"update_bank",
|
|
"delete_bank",
|
|
"clear_memories",
|
|
"get_knowledge_base_tree",
|
|
"search_knowledge_base",
|
|
"get_knowledge_page",
|
|
"create_knowledge_folder",
|
|
"create_knowledge_page",
|
|
"update_knowledge_node",
|
|
"delete_knowledge_node",
|
|
}
|
|
)
|
|
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
|
|
|
|
retain_description, recall_description = _build_mcp_tool_descriptions(
|
|
getattr(global_config, "mcp_instructions", None)
|
|
)
|
|
|
|
# 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
|
|
extra_headers_resolver=get_current_extra_headers, # Propagate allowlisted headers to extensions
|
|
include_bank_id_param=multi_bank,
|
|
tools=base_tools,
|
|
retain_description=retain_description,
|
|
recall_description=recall_description,
|
|
)
|
|
|
|
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 _get_mcp_tools(mcp: FastMCP) -> dict:
|
|
"""Get tool name→object mapping, compatible with FastMCP 2.x and 3.x."""
|
|
# FastMCP 2.x: _tool_manager._tools
|
|
if hasattr(mcp, "_tool_manager"):
|
|
return mcp._tool_manager._tools # type: ignore[union-attr]
|
|
# FastMCP 3.x: _local_provider._components with "tool:" prefix
|
|
if hasattr(mcp, "_local_provider"):
|
|
return {
|
|
k.split(":")[1].split("@")[0]: v
|
|
for k, v in mcp._local_provider._components.items() # type: ignore[union-attr]
|
|
if k.startswith("tool:")
|
|
}
|
|
msg = "Cannot locate tools on FastMCP instance"
|
|
raise AttributeError(msg)
|
|
|
|
|
|
def _make_tools_tolerant(mcp: FastMCP) -> None:
|
|
"""Wrap all tool run methods to strip unknown arguments and coerce string-encoded JSON.
|
|
|
|
LLMs frequently add extra fields like "explanation" or "reasoning" to tool calls.
|
|
FastMCP's Pydantic TypeAdapter rejects these with "Unexpected keyword argument".
|
|
|
|
LLMs also frequently serialize list/dict arguments as JSON strings instead of native
|
|
types (e.g., tags='["a","b"]' instead of tags=["a","b"]). This auto-coerces them.
|
|
|
|
This wraps each tool's run() to apply both fixes before validation.
|
|
"""
|
|
try:
|
|
tools = _get_mcp_tools(mcp)
|
|
for name, tool in tools.items():
|
|
if hasattr(tool, "parameters") and tool.parameters:
|
|
properties = tool.parameters.get("properties", {})
|
|
allowed = set(properties.keys())
|
|
|
|
# Build sets of parameter names that expect array or object types.
|
|
# Handles both direct types {"type": "array"} and anyOf/oneOf unions
|
|
# like {"anyOf": [{"type": "array", ...}, {"type": "null"}]}.
|
|
array_params: set[str] = set()
|
|
object_params: set[str] = set()
|
|
for param_name, param_schema in properties.items():
|
|
_collect_coercible_types(param_schema, param_name, array_params, object_params)
|
|
|
|
original_run = tool.run
|
|
|
|
async def _tolerant_run(
|
|
arguments,
|
|
_allowed=allowed,
|
|
_orig=original_run,
|
|
_array_params=array_params,
|
|
_object_params=object_params,
|
|
):
|
|
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}
|
|
|
|
# Coerce string-encoded JSON for list/dict parameters
|
|
arguments = _coerce_string_json(arguments, _array_params, _object_params)
|
|
|
|
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}")
|
|
|
|
|
|
def _collect_coercible_types(schema: dict, param_name: str, array_params: set[str], object_params: set[str]) -> None:
|
|
"""Check a JSON Schema property and add param_name to array_params/object_params if applicable."""
|
|
# Direct type
|
|
schema_type = schema.get("type")
|
|
if schema_type == "array":
|
|
array_params.add(param_name)
|
|
return
|
|
if schema_type == "object":
|
|
object_params.add(param_name)
|
|
return
|
|
|
|
# anyOf / oneOf unions (e.g., list[str] | None → {"anyOf": [{"type": "array"}, {"type": "null"}]})
|
|
for variant in schema.get("anyOf", []) + schema.get("oneOf", []):
|
|
variant_type = variant.get("type")
|
|
if variant_type == "array":
|
|
array_params.add(param_name)
|
|
return
|
|
if variant_type == "object":
|
|
object_params.add(param_name)
|
|
return
|
|
|
|
|
|
def _coerce_string_json(arguments: dict, array_params: set[str], object_params: set[str]) -> dict:
|
|
"""Auto-coerce string-encoded JSON arrays/objects to native types.
|
|
|
|
LLM agents frequently serialize list and dict tool arguments as JSON strings.
|
|
This is backward-compatible: native arrays/objects pass through unchanged.
|
|
"""
|
|
for param_name in array_params:
|
|
val = arguments.get(param_name)
|
|
if isinstance(val, str):
|
|
try:
|
|
parsed = json.loads(val)
|
|
if isinstance(parsed, list):
|
|
arguments = {**arguments, param_name: parsed}
|
|
logger.debug(f"Coerced string to list for parameter '{param_name}'")
|
|
except (json.JSONDecodeError, TypeError):
|
|
pass
|
|
|
|
for param_name in object_params:
|
|
val = arguments.get(param_name)
|
|
if isinstance(val, str):
|
|
try:
|
|
parsed = json.loads(val)
|
|
if isinstance(parsed, dict):
|
|
arguments = {**arguments, param_name: parsed}
|
|
logger.debug(f"Coerced string to dict for parameter '{param_name}'")
|
|
except (json.JSONDecodeError, TypeError):
|
|
pass
|
|
|
|
return arguments
|
|
|
|
|
|
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
|
|
|
|
def _get_extra_headers(self, scope: dict) -> dict[str, str]:
|
|
"""Collect the headers an operator opted into forwarding to extensions.
|
|
|
|
Shares ``collect_passthrough_headers`` with the HTTP transport, so both
|
|
agree on decoding and on what a duplicated header means. Empty unless
|
|
HINDSIGHT_API_EXTENSION_PASSTHROUGH_HEADERS names a header the request
|
|
actually carries.
|
|
"""
|
|
return collect_passthrough_headers(scope.get("headers", []), _get_raw_config().extension_passthrough_headers)
|
|
|
|
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()
|
|
|
|
# Resolved before authentication so authenticate_mcp() can read a
|
|
# passthrough header, not just the bearer token. Named for the request
|
|
# side: _send_error()'s `extra_headers` below is *response* headers.
|
|
passthrough_headers = self._get_extra_headers(scope)
|
|
|
|
# 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, extra_headers=dict(passthrough_headers))
|
|
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)
|
|
# Store the allowlisted passthrough headers so per-tool RequestContexts carry them
|
|
extra_headers_token = _current_extra_headers.set(passthrough_headers)
|
|
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)
|
|
_current_extra_headers.reset(extra_headers_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
|