fleet-memory/hindsight-api/tests/test_mcp_routing.py

"""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():
    """Create a mock MemoryEngine."""
    memory = MagicMock()
    memory.retain_batch_async = AsyncMock()
    memory.submit_async_retain = AsyncMock(return_value={"operation_id": "test-op-123"})
    memory.recall_async = AsyncMock(return_value=MagicMock(results=[]))
    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."""
    from hindsight_api.api.mcp import get_current_bank_id, _current_bank_id

    # Initially None
    assert get_current_bank_id() is None

    # Set and verify
    token = _current_bank_id.set("test-bank-123")
    try:
        assert get_current_bank_id() == "test-bank-123"
    finally:
        _current_bank_id.reset(token)

    # Back to None after reset
    assert get_current_bank_id() is None


@pytest.mark.asyncio
async def test_mcp_tools_use_context_bank_id(mock_memory):
    """Test that MCP tools use bank_id from context."""
    from hindsight_api.api.mcp import create_mcp_server, _current_bank_id

    mcp_server = create_mcp_server(mock_memory)

    # Get the tools
    tools = mcp_server._tool_manager._tools
    assert "retain" in tools
    assert "recall" in tools

    # Test retain with bank_id from context (use async_processing=False for synchronous test)
    token = _current_bank_id.set("context-bank-id")
    try:
        retain_tool = tools["retain"]
        result = await retain_tool.fn(content="test content", context="test_context", async_processing=False)
        assert "successfully" in result.lower()

        # Verify the memory was called with the context bank_id
        mock_memory.retain_batch_async.assert_called_once()
        call_kwargs = mock_memory.retain_batch_async.call_args.kwargs
        assert call_kwargs["bank_id"] == "context-bank-id"
    finally:
        _current_bank_id.reset(token)


def test_path_parsing_logic():
    """Test the path parsing logic for bank_id extraction."""
    def parse_path(path):
        """Simulate the path parsing logic from MCPMiddleware."""
        if not path.startswith("/") or len(path) <= 1:
            return None, None  # Error case

        parts = path[1:].split("/", 1)
        if not parts[0]:
            return None, None  # Error case

        bank_id = parts[0]
        new_path = "/" + parts[1] if len(parts) > 1 else "/"
        return bank_id, new_path

    # Test bank-specific paths
    bank_id, remaining = parse_path("/my-bank/")
    assert bank_id == "my-bank"
    assert remaining == "/"

    bank_id, remaining = parse_path("/my-bank")
    assert bank_id == "my-bank"
    assert remaining == "/"

    # Test error case - no bank_id
    bank_id, remaining = parse_path("/")
    assert bank_id is None

    # Test with complex bank_id
    bank_id, remaining = parse_path("/user_12345/")
    assert bank_id == "user_12345"
    assert remaining == "/"

    # Test with additional path after bank_id
    bank_id, remaining = parse_path("/my-bank/some/path")
    assert bank_id == "my-bank"
    assert remaining == "/some/path"


@pytest.mark.asyncio
async def test_api_key_context_variable():
    """Test that API key context variable works correctly."""
    from hindsight_api.api.mcp import get_current_api_key, _current_api_key

    # Initially None
    assert get_current_api_key() is None

    # Set and verify
    token = _current_api_key.set("test-api-key-123")
    try:
        assert get_current_api_key() == "test-api-key-123"
    finally:
        _current_api_key.reset(token)

    # Back to None after reset
    assert get_current_api_key() is None


@pytest.mark.asyncio
async def test_mcp_tools_propagate_api_key(mock_memory):
    """Test that MCP tools propagate API key to RequestContext."""
    from hindsight_api.api.mcp import create_mcp_server, _current_bank_id, _current_api_key

    mcp_server = create_mcp_server(mock_memory)
    tools = mcp_server._tool_manager._tools

    # Set both bank_id and api_key context
    bank_token = _current_bank_id.set("test-bank")
    api_key_token = _current_api_key.set("test-bearer-token")
    try:
        retain_tool = tools["retain"]
        result = await retain_tool.fn(content="test content", context="test_context", async_processing=False)
        assert "successfully" in result.lower()

        # Verify the memory was called with request_context containing api_key
        mock_memory.retain_batch_async.assert_called_once()
        call_kwargs = mock_memory.retain_batch_async.call_args.kwargs
        assert call_kwargs["request_context"].api_key == "test-bearer-token"
    finally:
        _current_bank_id.reset(bank_token)
        _current_api_key.reset(api_key_token)


# --- Middleware authentication tests ---


@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


@pytest.fixture
def memory_with_default_auth():
    """Create a mock MemoryEngine with DefaultTenantExtension (no auth)."""
    memory = MagicMock()
    memory._tenant_extension = DefaultTenantExtension({})
    return memory


@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."""
    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)

    assert status == 401
    assert b"Authentication failed" in body


@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."""
    from hindsight_api.api.mcp import MCPMiddleware

    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)

    assert status == 401
    assert b"Authentication failed" in body


@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

    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(),
        )


@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(),
        )