holetron-lab/fleet-memory
1
0
Fork 0
Code Releases Packages Activity Actions

feat: improve mcp tools based on endpoint (#318)

Browse source 
* feat: improve mcp tools based on endpoint

* feat: improve mcp tools based on endpoint

* test: add integration test for MCP endpoint routing

- Add test_mcp_endpoint_routing.py to verify single-bank vs multi-bank tool exposure
- Verifies /mcp/ exposes all tools with bank_id parameters
- Verifies /mcp/{bank_id}/ only exposes scoped tools without bank_id parameters
- Regression test for issue #317

Related: #317, #318

* test: use StreamableHTTP client for MCP endpoint routing test

Replace httpx AsyncClient SSE parsing with proper MCP StreamableHTTP
client. This correctly tests the MCP server using the actual protocol
that clients will use.

Fixes #317
This commit is contained in:
Nicolò Boschi 2026-02-08 09:28:59 +01:00 committed by GitHub
parent d0f67c9f8b
commit d90588b3e1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 284 additions and 245 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
hindsight-api/hindsight_api/api/__init__.py
View file

@ -72,22 +72,24 @@ def create_app(
# Mount MCP server and chain its lifespan if enabled
if mcp_app is not None:
# Get the MCP app's underlying Starlette app for lifespan access
mcp_starlette_app = mcp_app.mcp_app
# Get both MCP apps' underlying Starlette apps for lifespan access
multi_bank_starlette_app = mcp_app.multi_bank_app
single_bank_starlette_app = mcp_app.single_bank_app
# Store the original lifespan
original_lifespan = app.router.lifespan_context
@asynccontextmanager
async def chained_lifespan(app_instance: FastAPI):
"""Chain the MCP lifespan with the main app lifespan."""
# Start MCP lifespan first
async with mcp_starlette_app.router.lifespan_context(mcp_starlette_app):
logger.info("MCP lifespan started")
# Then start the original app lifespan
async with original_lifespan(app_instance):
yield
logger.info("MCP lifespan stopped")
"""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

89
hindsight-api/hindsight_api/api/mcp.py
View file

@ -54,12 +54,14 @@ def get_current_api_key() -> str | None:
return _current_api_key.get()
def create_mcp_server(memory: MemoryEngine) -> FastMCP:
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 with stateless_http enabled
@ -71,8 +73,8 @@ def create_mcp_server(memory: MemoryEngine) -> FastMCP:
config = MCPToolsConfig(
bank_id_resolver=get_current_bank_id,
api_key_resolver=get_current_api_key, # Propagate API key for tenant auth
include_bank_id_param=True, # HTTP MCP supports multi-bank via parameter
tools=None, # All tools
include_bank_id_param=multi_bank,
tools=None if multi_bank else {"retain", "recall", "reflect"}, # Scoped tools for single-bank mode
retain_fire_and_forget=False, # HTTP MCP supports sync/async modes
)
@ -88,7 +90,7 @@ def create_mcp_server(memory: MemoryEngine) -> FastMCP:
class MCPMiddleware:
"""ASGI middleware that handles authentication and extracts bank_id from header or path.
"""ASGI middleware that handles authentication and routes to appropriate MCP server.
Authentication:
1. If HINDSIGHT_API_MCP_AUTH_TOKEN is set (legacy), validates against that token
@ -96,12 +98,25 @@ class MCPMiddleware:
- DefaultTenantExtension: no auth required (local dev)
- ApiKeyTenantExtension: validates against env var
Bank ID can be provided via:
1. X-Bank-Id header (recommended for Claude Code)
2. URL path: /mcp/{bank_id}/
3. Environment variable HINDSIGHT_MCP_BANK_ID (fallback default)
Two modes based on URL structure:
For Claude Code, configure with:
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
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>"
"""
@ -110,10 +125,23 @@ class MCPMiddleware:
self.app = app
self.memory = memory
self.tenant_extension = memory._tenant_extension
self.mcp_server = create_mcp_server(memory)
self.mcp_app = self.mcp_server.http_app(path="/")
# Expose the lifespan for the parent app to chain
self.lifespan = self.mcp_app.lifespan_handler if hasattr(self.mcp_app, "lifespan_handler") else None
# Create two server instances:
# 1. Multi-bank server (for /mcp/ root endpoint)
self.multi_bank_server = create_mcp_server(memory, multi_bank=True)
self.multi_bank_app = self.multi_bank_server.http_app(path="/")
# 2. Single-bank server (for /mcp/{bank_id}/ endpoints)
self.single_bank_server = create_mcp_server(memory, multi_bank=False)
self.single_bank_app = self.single_bank_server.http_app(path="/")
# Backward compatibility: expose multi_bank_app as mcp_app
self.mcp_app = self.multi_bank_app
# Expose the lifespan for the parent app to chain (use multi-bank as default)
self.lifespan = (
self.multi_bank_app.lifespan_handler if hasattr(self.multi_bank_app, "lifespan_handler") else None
)
def _get_header(self, scope: dict, name: str) -> str | None:
"""Extract a header value from ASGI scope."""
@ -125,7 +153,7 @@ class MCPMiddleware:
async def __call__(self, scope, receive, send):
if scope["type"] != "http":
await self.mcp_app(scope, receive, send)
await self.multi_bank_app(scope, receive, send)
return
# Extract auth token from header (for tenant auth propagation)
@ -173,8 +201,13 @@ class MCPMiddleware:
elif path == "/mcp":
path = "/"
# Ensure path has leading slash (needed after stripping mount path)
if path and not path.startswith("/"):
path = "/" + path
# Try to get bank_id from header first (for Claude Code compatibility)
bank_id = self._get_header(scope, "X-Bank-Id")
bank_id_from_path = False
# MCP endpoint paths that should not be treated as bank_ids
MCP_ENDPOINTS = {"sse", "messages"}
@ -187,6 +220,7 @@ class MCPMiddleware:
if parts[0] and parts[0] not in MCP_ENDPOINTS:
# First segment looks like a bank_id
bank_id = parts[0]
bank_id_from_path = True
new_path = "/" + parts[1] if len(parts) > 1 else "/"
# Fall back to default bank_id
@ -194,6 +228,11 @@ class MCPMiddleware:
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 and api_key context
bank_id_token = _current_bank_id.set(bank_id)
# Store the auth token for tenant extension to validate
@ -206,7 +245,7 @@ class MCPMiddleware:
# Wrap send to rewrite the SSE endpoint URL to include bank_id if using path-based routing
async def send_wrapper(message):
if message["type"] == "http.response.body":
if message["type"] == "http.response.body" and bank_id_from_path:
body = message.get("body", b"")
if body and b"/messages" in body:
# Rewrite /messages to /{bank_id}/messages in SSE endpoint event
@ -214,7 +253,7 @@ class MCPMiddleware:
message = {**message, "body": body}
await send(message)
await self.mcp_app(new_scope, receive, send_wrapper)
await target_app(new_scope, receive, send_wrapper)
finally:
_current_bank_id.reset(bank_id_token)
if api_key_token is not None:
@ -242,15 +281,23 @@ class MCPMiddleware:
def create_mcp_app(memory: MemoryEngine):
"""
Create an ASGI app that handles MCP requests.
Create an ASGI app that handles MCP requests with dynamic tool exposure.
Authentication:
Uses the TenantExtension from the MemoryEngine (same auth as REST API).
Bank ID can be provided via:
1. X-Bank-Id header: claude mcp add --transport http hindsight http://localhost:8888/mcp --header "X-Bank-Id: my-bank"
2. URL path: /mcp/{bank_id}/
3. Environment variable HINDSIGHT_MCP_BANK_ID (fallback, default: "default")
Two modes based on URL structure:
1. Single-bank mode (recommended for agent isolation):
- URL: /mcp/{bank_id}/
- Tools: retain, recall, reflect (no bank_id parameter)
- Example: claude mcp add --transport http my-agent http://localhost:8888/mcp/my-agent-bank/
2. Multi-bank mode (for cross-bank operations):
- URL: /mcp/
- Tools: retain, recall, reflect, list_banks, create_bank (all with bank_id parameter)
- Bank ID from: X-Bank-Id header or HINDSIGHT_MCP_BANK_ID env var (default: "default")
- Example: claude mcp add --transport http hindsight http://localhost:8888/mcp --header "X-Bank-Id: my-bank"
Args:
memory: MemoryEngine instance

78
hindsight-api/tests/test_mcp_endpoint_routing.py Normal file
View file

@ -0,0 +1,78 @@
"""Integration test for MCP endpoint routing.
This test verifies that /mcp/ and /mcp/{bank_id}/ expose different tool sets.
"""
import httpx
import pytest
from mcp.client.session import ClientSession
from mcp.client.streamable_http import streamable_http_client
@pytest.mark.asyncio
async def test_mcp_endpoint_routing_integration(memory):
"""Test that multi-bank and single-bank endpoints expose different tools using StreamableHTTP.
This is a regression test for issue #317 where /mcp/{bank_id}/ was incorrectly
exposing all tools (including list_banks) and bank_id parameters.
"""
from hindsight_api.api import create_app
# Create app with MCP enabled
app = create_app(memory, mcp_api_enabled=True, initialize_memory=False)
# Use the app's lifespan context to properly initialize MCP servers
async with app.router.lifespan_context(app):
# Create an HTTPX client that routes to our ASGI app
from httpx import ASGITransport
async with httpx.AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as http_client:
# Test 1: Multi-bank endpoint /mcp/
async with streamable_http_client("http://test/mcp/", http_client=http_client) as (
read_stream,
write_stream,
_,
):
async with ClientSession(read_stream, write_stream) as session:
await session.initialize()
multi_result = await session.list_tools()
multi_tools = {t.name for t in multi_result.tools}
# Multi-bank should have all tools including bank management
assert "retain" in multi_tools
assert "recall" in multi_tools
assert "reflect" in multi_tools
assert "list_banks" in multi_tools, "Multi-bank should expose list_banks"
assert "create_bank" in multi_tools, "Multi-bank should expose create_bank"
# Multi-bank retain should have bank_id parameter
retain_tool = next((t for t in multi_result.tools if t.name == "retain"), None)
assert retain_tool is not None
multi_params = set(retain_tool.inputSchema.get("properties", {}).keys())
assert "bank_id" in multi_params, "Multi-bank retain should have bank_id parameter"
# Test 2: Single-bank endpoint /mcp/test-bank/
async with streamable_http_client("http://test/mcp/test-bank/", http_client=http_client) as (
read_stream,
write_stream,
_,
):
async with ClientSession(read_stream, write_stream) as session:
await session.initialize()
single_result = await session.list_tools()
single_tools = {t.name for t in single_result.tools}
# Single-bank should only have scoped tools (no bank management)
assert "retain" in single_tools
assert "recall" in single_tools
assert "reflect" in single_tools
assert "list_banks" not in single_tools, "Single-bank should NOT expose list_banks"
assert "create_bank" not in single_tools, "Single-bank should NOT expose create_bank"
# Single-bank retain should NOT have bank_id parameter
retain_tool = next((t for t in single_result.tools if t.name == "retain"), None)
assert retain_tool is not None
single_params = set(retain_tool.inputSchema.get("properties", {}).keys())
assert "bank_id" not in single_params, "Single-bank retain should NOT have bank_id parameter"

340
hindsight-api/tests/test_mcp_routing.py
View file

@ -1,12 +1,8 @@
"""Test MCP server routing with dynamic bank_id."""
import json
import pytest
from unittest.mock import AsyncMock, MagicMock
from hindsight_api.extensions.builtin.tenant import ApiKeyTenantExtension, DefaultTenantExtension
@pytest.fixture
def mock_memory():
@ -18,41 +14,6 @@ def mock_memory():
return memory
def _make_scope(path="/mcp", headers=None):
"""Build a minimal ASGI HTTP scope."""
raw_headers = []
for k, v in (headers or {}).items():
raw_headers.append((k.lower().encode(), v.encode()))
# MCP requires Accept header
raw_headers.append((b"accept", b"application/json, text/event-stream"))
raw_headers.append((b"content-type", b"application/json"))
return {
"type": "http",
"path": path,
"root_path": "",
"headers": raw_headers,
}
async def _collect_response(middleware, scope, body=b""):
"""Send a request through the middleware and collect the response status and body."""
status = None
response_body = b""
async def receive():
return {"type": "http.request", "body": body}
async def send(message):
nonlocal status, response_body
if message["type"] == "http.response.start":
status = message["status"]
elif message["type"] == "http.response.body":
response_body += message.get("body", b"")
await middleware(scope, receive, send)
return status, response_body
@pytest.mark.asyncio
async def test_mcp_context_variable():
"""Test that context variable works correctly."""
@ -182,199 +143,150 @@ async def test_mcp_tools_propagate_api_key(mock_memory):
_current_api_key.reset(api_key_token)
# --- Middleware authentication tests ---
def test_multi_bank_mode_exposes_all_tools(mock_memory):
"""Test that multi-bank mode exposes all tools including bank management."""
from hindsight_api.api.mcp import create_mcp_server
# Create server in multi-bank mode (default)
mcp_server = create_mcp_server(mock_memory, multi_bank=True)
tools = mcp_server._tool_manager._tools
# Should have all tools
assert "retain" in tools
assert "recall" in tools
assert "reflect" in tools
assert "list_banks" in tools
assert "create_bank" in tools
@pytest.fixture
def memory_with_api_key_auth():
"""Create a mock MemoryEngine with ApiKeyTenantExtension."""
memory = MagicMock()
memory._tenant_extension = ApiKeyTenantExtension({"api_key": "test-secret-123"})
return memory
def test_single_bank_mode_excludes_bank_management_tools(mock_memory):
"""Test that single-bank mode only exposes bank-scoped tools."""
from hindsight_api.api.mcp import create_mcp_server
# Create server in single-bank mode
mcp_server = create_mcp_server(mock_memory, multi_bank=False)
tools = mcp_server._tool_manager._tools
# Should only have bank-scoped tools
assert "retain" in tools
assert "recall" in tools
assert "reflect" in tools
# Should NOT have bank management tools
assert "list_banks" not in tools
assert "create_bank" not in tools
@pytest.fixture
def memory_with_default_auth():
"""Create a mock MemoryEngine with DefaultTenantExtension (no auth)."""
memory = MagicMock()
memory._tenant_extension = DefaultTenantExtension({})
return memory
def test_multi_bank_mode_tools_have_bank_id_param(mock_memory):
"""Test that multi-bank mode tools include bank_id parameter."""
from hindsight_api.api.mcp import create_mcp_server
import inspect
mcp_server = create_mcp_server(mock_memory, multi_bank=True)
tools = mcp_server._tool_manager._tools
# Check that tools have bank_id parameter
retain_tool = tools["retain"]
retain_sig = inspect.signature(retain_tool.fn)
assert "bank_id" in retain_sig.parameters
recall_tool = tools["recall"]
recall_sig = inspect.signature(recall_tool.fn)
assert "bank_id" in recall_sig.parameters
reflect_tool = tools["reflect"]
reflect_sig = inspect.signature(reflect_tool.fn)
assert "bank_id" in reflect_sig.parameters
def test_single_bank_mode_tools_no_bank_id_param(mock_memory):
"""Test that single-bank mode tools do NOT include bank_id parameter."""
from hindsight_api.api.mcp import create_mcp_server
import inspect
mcp_server = create_mcp_server(mock_memory, multi_bank=False)
tools = mcp_server._tool_manager._tools
# Check that tools do NOT have bank_id parameter
retain_tool = tools["retain"]
retain_sig = inspect.signature(retain_tool.fn)
assert "bank_id" not in retain_sig.parameters
recall_tool = tools["recall"]
recall_sig = inspect.signature(recall_tool.fn)
assert "bank_id" not in recall_sig.parameters
reflect_tool = tools["reflect"]
reflect_sig = inspect.signature(reflect_tool.fn)
assert "bank_id" not in reflect_sig.parameters
@pytest.mark.asyncio
async def test_mcp_middleware_rejects_no_auth(memory_with_api_key_auth):
"""MCP middleware returns 401 when no Authorization header is provided."""
async def test_middleware_handles_both_endpoints(mock_memory):
"""Test that MCPMiddleware routes to correct server based on URL path."""
from hindsight_api.api.mcp import MCPMiddleware
middleware = MCPMiddleware(None, memory_with_api_key_auth)
scope = _make_scope(path="/mcp")
status, body = await _collect_response(middleware, scope)
# Create middleware (single instance)
middleware = MCPMiddleware(None, mock_memory)
assert status == 401
assert b"Authentication failed" in body
# Verify both server instances exist
assert middleware.multi_bank_app is not None
assert middleware.single_bank_app is not None
# Verify they expose different tools
multi_bank_tools = middleware.multi_bank_server._tool_manager._tools
single_bank_tools = middleware.single_bank_server._tool_manager._tools
# Multi-bank should have all tools
assert "retain" in multi_bank_tools
assert "recall" in multi_bank_tools
assert "list_banks" in multi_bank_tools
assert "create_bank" in multi_bank_tools
# Single-bank should only have scoped tools
assert "retain" in single_bank_tools
assert "recall" in single_bank_tools
assert "list_banks" not in single_bank_tools
assert "create_bank" not in single_bank_tools
@pytest.mark.asyncio
async def test_mcp_middleware_rejects_wrong_key(memory_with_api_key_auth):
"""MCP middleware returns 401 when an invalid API key is provided."""
async def test_routing_logic_from_url_path():
"""Test that routing correctly selects server based on URL structure."""
from hindsight_api.api.mcp import MCPMiddleware
from unittest.mock import AsyncMock
middleware = MCPMiddleware(None, memory_with_api_key_auth)
scope = _make_scope(path="/mcp", headers={"Authorization": "Bearer wrong-key"})
status, body = await _collect_response(middleware, scope)
# Mock memory
mock_memory = MagicMock()
assert status == 401
assert b"Authentication failed" in body
# Create middleware
middleware = MCPMiddleware(None, mock_memory)
# Simulate different URL patterns and verify routing
test_cases = [
# (path_after_stripping_mcp, expected_bank_id_from_path, expected_bank_id, description)
("/alice/messages", True, "alice", "Bank ID in path with endpoint"),
("/my-agent-123/", True, "my-agent-123", "Bank ID in path with trailing slash"),
("ciccio/messages", True, "ciccio", "Bank ID without leading slash (after mount strip)"),
("bob", True, "bob", "Bank ID only, no leading slash"),
("/messages", False, None, "MCP endpoint, no bank ID"),
("/", False, None, "Root path, no bank ID"),
]
@pytest.mark.asyncio
async def test_mcp_middleware_accepts_valid_key(memory_with_api_key_auth):
"""MCP middleware passes through when a valid API key is provided."""
from hindsight_api.api.mcp import MCPMiddleware
for path, expected_bank_from_path, expected_bank_id, description in test_cases:
# Simulate the path parsing logic with leading slash normalization
if path and not path.startswith("/"):
path = "/" + path
middleware = MCPMiddleware(None, memory_with_api_key_auth)
scope = _make_scope(
path="/mcp",
headers={"Authorization": "Bearer test-secret-123"},
)
# FastMCP raises RuntimeError because its lifespan isn't initialized in unit tests.
# If we get that error, auth passed — the request made it past the middleware.
with pytest.raises(RuntimeError, match="Task group is not initialized"):
await _collect_response(
middleware,
scope,
body=json.dumps({"jsonrpc": "2.0", "id": 1, "method": "tools/list"}).encode(),
)
bank_id = None
bank_id_from_path = False
MCP_ENDPOINTS = {"sse", "messages"}
if path.startswith("/") and len(path) > 1:
parts = path[1:].split("/", 1)
if parts[0] and parts[0] not in MCP_ENDPOINTS:
bank_id = parts[0]
bank_id_from_path = True
@pytest.mark.asyncio
async def test_mcp_middleware_default_tenant_no_auth_required(memory_with_default_auth):
"""MCP middleware passes through with no auth when DefaultTenantExtension is used."""
from hindsight_api.api.mcp import MCPMiddleware
middleware = MCPMiddleware(None, memory_with_default_auth)
scope = _make_scope(path="/mcp")
# Same as above — RuntimeError means auth passed and request reached FastMCP internals.
with pytest.raises(RuntimeError, match="Task group is not initialized"):
await _collect_response(
middleware,
scope,
body=json.dumps({"jsonrpc": "2.0", "id": 1, "method": "tools/list"}).encode(),
)
class MultiTenantTestExtension:
"""Test extension that maps API keys to tenant schemas."""
def __init__(self, key_to_schema: dict[str, str]):
self.key_to_schema = key_to_schema
async def authenticate(self, context):
from hindsight_api.extensions.tenant import AuthenticationError, TenantContext
if not context.api_key:
raise AuthenticationError("API key required")
schema = self.key_to_schema.get(context.api_key)
if not schema:
raise AuthenticationError("Invalid API key")
return TenantContext(schema_name=schema)
async def authenticate_mcp(self, context):
"""MCP auth delegates to authenticate by default."""
return await self.authenticate(context)
@pytest.mark.asyncio
async def test_mcp_middleware_sets_schema_from_tenant_context():
"""MCP middleware sets _current_schema from tenant context for multi-tenant isolation."""
from hindsight_api.api.mcp import MCPMiddleware
from hindsight_api.engine.memory_engine import _current_schema
# Create extension that maps keys to different schemas
tenant_ext = MultiTenantTestExtension({
"key-for-tenant-alpha": "tenant_alpha",
"key-for-tenant-beta": "tenant_beta",
})
memory = MagicMock()
memory._tenant_extension = tenant_ext
middleware = MCPMiddleware(None, memory)
# Track what schema was set during request processing
captured_schema = None
# Patch the mcp_app to capture the schema instead of actually processing
async def mock_mcp_app(scope, receive, send):
nonlocal captured_schema
captured_schema = _current_schema.get()
# Send a minimal response
await send({"type": "http.response.start", "status": 200, "headers": []})
await send({"type": "http.response.body", "body": b"{}"})
middleware.mcp_app = mock_mcp_app
# Test tenant alpha
scope = _make_scope(path="/mcp", headers={"Authorization": "Bearer key-for-tenant-alpha"})
await _collect_response(middleware, scope)
assert captured_schema == "tenant_alpha", f"Expected tenant_alpha, got {captured_schema}"
# Test tenant beta
scope = _make_scope(path="/mcp", headers={"Authorization": "Bearer key-for-tenant-beta"})
await _collect_response(middleware, scope)
assert captured_schema == "tenant_beta", f"Expected tenant_beta, got {captured_schema}"
@pytest.mark.asyncio
async def test_mcp_legacy_auth_token(monkeypatch):
"""MCP middleware supports legacy MCP_AUTH_TOKEN for backwards compatibility."""
import hindsight_api.api.mcp as mcp_module
from hindsight_api.api.mcp import MCPMiddleware
# Set legacy auth token
monkeypatch.setattr(mcp_module, "MCP_AUTH_TOKEN", "legacy-secret-token")
memory = MagicMock()
# Even with ApiKeyTenantExtension, legacy token should work
memory._tenant_extension = ApiKeyTenantExtension({"api_key": "different-key"})
middleware = MCPMiddleware(None, memory)
# Wrong token should fail
scope = _make_scope(path="/mcp", headers={"Authorization": "Bearer wrong-token"})
status, body = await _collect_response(middleware, scope)
assert status == 401
assert b"Invalid authentication token" in body
# Correct legacy token should pass
scope = _make_scope(path="/mcp", headers={"Authorization": "Bearer legacy-secret-token"})
with pytest.raises(RuntimeError, match="Task group is not initialized"):
await _collect_response(
middleware,
scope,
body=json.dumps({"jsonrpc": "2.0", "id": 1, "method": "tools/list"}).encode(),
)
@pytest.mark.asyncio
async def test_mcp_auth_disabled_flag():
"""ApiKeyTenantExtension with mcp_auth_disabled=true skips MCP auth."""
from hindsight_api.api.mcp import MCPMiddleware
memory = MagicMock()
# Create extension with MCP auth disabled
memory._tenant_extension = ApiKeyTenantExtension({
"api_key": "test-secret-123",
"mcp_auth_disabled": "true",
})
middleware = MCPMiddleware(None, memory)
# No auth should pass when mcp_auth_disabled=true
scope = _make_scope(path="/mcp")
with pytest.raises(RuntimeError, match="Task group is not initialized"):
await _collect_response(
middleware,
scope,
body=json.dumps({"jsonrpc": "2.0", "id": 1, "method": "tools/list"}).encode(),
)
assert bank_id_from_path == expected_bank_from_path, f"Failed for: {description} (path={path})"
assert bank_id == expected_bank_id, f"Failed bank_id for: {description} (path={path}, got={bank_id})"