fleet-memory/hindsight-api/hindsight_api/api/__init__.py
DK09876 e798979733
Harden MCP server: fix routing, validation, and usage metering (#341)
* fix: move mental model usage metering into engine for MCP support

Mental model validation hooks (validate_mental_model_get, validate_mental_model_refresh)
were only called in REST HTTP handlers, not in the engine. MCP tools call engine methods
directly, so usage metering was skipped entirely for MCP mental model operations.

Moved pre-validation and post-completion hooks into memory_engine.py (matching the
retain/recall/reflect pattern) and removed the duplicate code from http.py.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: remove double validation from create_mental_model and add internal checks

- Remove pre-validation from create_mental_model since callers always call
  submit_async_refresh_mental_model next (which validates), preventing
  double credit checks
- Add is_internal checks to mental model metering validators (matching
  the existing pattern for recall/reflect) so background worker tasks
  skip billing

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: prevent 307 redirect on /mcp that breaks MCP tool discovery

Starlette's Mount class redirects /mcp to /mcp/ with a 307 Temporary
Redirect. Many MCP clients don't follow POST redirects, which causes
tool discovery to fail (0 tools discovered despite successful auth).

Add _MCPPathRewriteMiddleware that rewrites /mcp to /mcp/ at the ASGI
level before routing, preventing the redirect entirely. Both /mcp and
/mcp/ now work identically.

Add regression test test_mcp_no_trailing_slash_works to verify URLs
with and without trailing slashes discover tools correctly.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* harden MCP server for real-world usage

- Remove MCP_ENDPOINTS blocklist so banks named "sse"/"messages" route correctly
- Scope SSE body rewriting to text/event-stream responses only to prevent data corruption
- Add _validate_mental_model_inputs for name, source_query, max_tokens validation in MCP tools
- Improve "not found" error messages to include bank_id context
- Fix fragile tool count assertions (exact → minimum bounds)
- Add integration tests: tool execution, input validation, edge-case bank names
- Add unit tests for validation helper and tool-level validation

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* refactor: replace Mount + rewrite middleware with wrapping middleware

Starlette's Mount class redirects /mcp -> /mcp/ with 307, which MCP clients
don't follow. Previously we patched this with _MCPPathRewriteMiddleware.

Now MCPMiddleware wraps the FastAPI app directly via add_middleware, intercepting
/mcp* requests before they reach Starlette's router. No Mount means no redirect.

- Remove _MCPPathRewriteMiddleware (no longer needed)
- Remove app.mount() call
- Add prefix parameter to MCPMiddleware
- Use app.add_middleware() for proper Starlette integration
- Simplify path stripping (just remove prefix, no mount/root_path handling)
- Update routing test to match current behavior (no MCP_ENDPOINTS blocklist)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: update stale docstring referencing removed _MCPPathRewriteMiddleware

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 10:41:20 +01:00

136 lines
4.4 KiB
Python

"""
Unified API module for Hindsight.
Provides both HTTP REST API and MCP (Model Context Protocol) server.
"""
import logging
from contextlib import asynccontextmanager
from fastapi import FastAPI
from hindsight_api import MemoryEngine
logger = logging.getLogger(__name__)
def create_app(
memory: MemoryEngine,
http_api_enabled: bool = True,
mcp_api_enabled: bool = False,
mcp_mount_path: str = "/mcp",
initialize_memory: bool = True,
) -> FastAPI:
"""
Create and configure the unified Hindsight API application.
Args:
memory: MemoryEngine instance (already initialized with required parameters).
Migrations are controlled by the MemoryEngine's run_migrations parameter.
http_api_enabled: Whether to enable HTTP REST API endpoints (default: True)
mcp_api_enabled: Whether to enable MCP server (default: False)
mcp_mount_path: Path to mount MCP server (default: /mcp)
initialize_memory: Whether to initialize memory system on startup (default: True)
Returns:
Configured FastAPI application with enabled APIs
Example:
# HTTP only
app = create_app(memory)
# MCP only
app = create_app(memory, http_api_enabled=False, mcp_api_enabled=True)
# Both HTTP and MCP
app = create_app(memory, mcp_api_enabled=True)
"""
mcp_servers = None
# Create MCP servers first if enabled (we need their lifespans for chaining)
if mcp_api_enabled:
try:
from .mcp import MCPMiddleware, create_mcp_servers
mcp_servers = create_mcp_servers(memory=memory)
except ImportError as e:
logger.error(f"MCP server requested but dependencies not available: {e}")
logger.error("Install with: pip install hindsight-api[mcp]")
raise
# Import and create HTTP API if enabled
if http_api_enabled:
from .http import create_app as create_http_app
app = create_http_app(memory=memory, initialize_memory=initialize_memory)
logger.info("HTTP REST API enabled")
else:
# Create minimal FastAPI app
app = FastAPI(title="Hindsight API", version="0.0.7")
logger.info("HTTP REST API disabled")
# Add MCP middleware and chain its lifespan if enabled
if mcp_servers is not None:
multi_bank_server, single_bank_server, multi_bank_starlette_app, single_bank_starlette_app = mcp_servers
# Store the original lifespan
original_lifespan = app.router.lifespan_context
@asynccontextmanager
async def chained_lifespan(app_instance: FastAPI):
"""Chain both MCP lifespans with the main app lifespan."""
# Start both MCP lifespans (multi-bank and single-bank)
async with multi_bank_starlette_app.router.lifespan_context(multi_bank_starlette_app):
async with single_bank_starlette_app.router.lifespan_context(single_bank_starlette_app):
logger.info("MCP lifespans started (multi-bank and single-bank)")
# Then start the original app lifespan
async with original_lifespan(app_instance):
yield
logger.info("MCP lifespans stopped")
# Replace the app's lifespan with the chained version
app.router.lifespan_context = chained_lifespan
# Add MCP as a wrapping middleware — intercepts /mcp* requests directly,
# passes everything else through to the FastAPI app. No Starlette Mount
# means no 307 redirect for /mcp (no trailing slash).
app.add_middleware(
MCPMiddleware,
memory=memory,
prefix=mcp_mount_path,
multi_bank_app=multi_bank_starlette_app,
single_bank_app=single_bank_starlette_app,
multi_bank_server=multi_bank_server,
single_bank_server=single_bank_server,
)
logger.info(f"MCP server enabled at {mcp_mount_path}/")
return app
# Re-export commonly used items for backwards compatibility
from .http import (
CreateBankRequest,
DispositionTraits,
MemoryItem,
RecallRequest,
RecallResponse,
RecallResult,
ReflectRequest,
ReflectResponse,
RetainRequest,
)
__all__ = [
"create_app",
"RecallRequest",
"RecallResult",
"RecallResponse",
"MemoryItem",
"RetainRequest",
"ReflectRequest",
"ReflectResponse",
"CreateBankRequest",
"DispositionTraits",
]