f6f1a7d889
* fix: zeroentropy rerank URL missing /v1 prefix and MCP routing tests - Fix ZeroEntropy reranker URL: /models/rerank -> /v1/models/rerank (#453) - Fix test_mcp_routing tests: update assertions to use submit_async_retain instead of the non-existent async_processing=False/retain_batch_async pattern * fix(openclaw): pass retainEveryNTurns through getPluginConfig and set it to 1 in tests getPluginConfig was not forwarding retainEveryNTurns from the raw config, so pluginConfig.retainEveryNTurns was always undefined (defaulting to 10). The integration tests use retainEveryNTurns: 1 so retain fires every turn.
454 lines
16 KiB
Python
454 lines
16 KiB
Python
"""Test MCP server routing with dynamic bank_id."""
|
|
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
|
|
|
|
@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
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mcp_context_variable():
|
|
"""Test that context variable works correctly."""
|
|
from hindsight_api.api.mcp import _current_bank_id, get_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 _current_bank_id, create_mcp_server
|
|
|
|
mcp_server = create_mcp_server(mock_memory)
|
|
|
|
# Get the tools
|
|
tools = mcp_server._tool_manager._tools
|
|
assert "retain" in tools
|
|
assert "recall" in tools
|
|
|
|
token = _current_bank_id.set("context-bank-id")
|
|
try:
|
|
retain_tool = tools["retain"]
|
|
result = await retain_tool.fn(content="test content", context="test_context")
|
|
assert result["status"] == "accepted"
|
|
|
|
# Verify the memory was called with the context bank_id
|
|
mock_memory.submit_async_retain.assert_called_once()
|
|
call_kwargs = mock_memory.submit_async_retain.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 _current_api_key, get_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 _current_api_key, _current_bank_id, create_mcp_server
|
|
|
|
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")
|
|
assert result["status"] == "accepted"
|
|
|
|
# Verify the memory was called with request_context containing api_key
|
|
mock_memory.submit_async_retain.assert_called_once()
|
|
call_kwargs = mock_memory.submit_async_retain.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)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tenant_id_context_variable():
|
|
"""Test that tenant_id and api_key_id context variables work correctly."""
|
|
from hindsight_api.api.mcp import (
|
|
_current_api_key_id,
|
|
_current_tenant_id,
|
|
get_current_api_key_id,
|
|
get_current_tenant_id,
|
|
)
|
|
|
|
# Initially None
|
|
assert get_current_tenant_id() is None
|
|
assert get_current_api_key_id() is None
|
|
|
|
# Set and verify
|
|
tenant_token = _current_tenant_id.set("org-123")
|
|
key_id_token = _current_api_key_id.set("key-456")
|
|
try:
|
|
assert get_current_tenant_id() == "org-123"
|
|
assert get_current_api_key_id() == "key-456"
|
|
finally:
|
|
_current_tenant_id.reset(tenant_token)
|
|
_current_api_key_id.reset(key_id_token)
|
|
|
|
# Back to None after reset
|
|
assert get_current_tenant_id() is None
|
|
assert get_current_api_key_id() is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mcp_tools_propagate_tenant_id_and_api_key_id(mock_memory):
|
|
"""Test that MCP tools propagate tenant_id and api_key_id to RequestContext.
|
|
|
|
This is the critical test for usage metering: the UsageMeteringValidator reads
|
|
request_context.tenant_id to identify the org for billing. Without this,
|
|
MCP operations get tenant_id="unknown" and billing is skipped entirely.
|
|
"""
|
|
from hindsight_api.api.mcp import (
|
|
_current_api_key,
|
|
_current_api_key_id,
|
|
_current_bank_id,
|
|
_current_tenant_id,
|
|
create_mcp_server,
|
|
)
|
|
|
|
mcp_server = create_mcp_server(mock_memory)
|
|
tools = mcp_server._tool_manager._tools
|
|
|
|
# Set all context vars (simulating what MCPMiddleware does after authenticate_mcp)
|
|
bank_token = _current_bank_id.set("test-bank")
|
|
api_key_token = _current_api_key.set("hsk_test_key")
|
|
tenant_token = _current_tenant_id.set("org-billing-123")
|
|
key_id_token = _current_api_key_id.set("key-uuid-456")
|
|
try:
|
|
retain_tool = tools["retain"]
|
|
await retain_tool.fn(content="test content", context="test_context")
|
|
|
|
# Verify the RequestContext passed to memory engine has all auth fields
|
|
mock_memory.submit_async_retain.assert_called_once()
|
|
request_context = mock_memory.submit_async_retain.call_args.kwargs["request_context"]
|
|
assert request_context.api_key == "hsk_test_key"
|
|
assert request_context.tenant_id == "org-billing-123"
|
|
assert request_context.api_key_id == "key-uuid-456"
|
|
finally:
|
|
_current_bank_id.reset(bank_token)
|
|
_current_api_key.reset(api_key_token)
|
|
_current_tenant_id.reset(tenant_token)
|
|
_current_api_key_id.reset(key_id_token)
|
|
|
|
|
|
def test_multi_bank_mode_exposes_all_tools(mock_memory):
|
|
"""Test that multi-bank mode exposes all tools including bank management and mental models."""
|
|
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
|
|
|
|
# Core tools
|
|
assert "retain" in tools
|
|
assert "recall" in tools
|
|
assert "reflect" in tools
|
|
assert "list_banks" in tools
|
|
assert "create_bank" in tools
|
|
|
|
# Mental model tools
|
|
assert "list_mental_models" in tools
|
|
assert "get_mental_model" in tools
|
|
assert "create_mental_model" in tools
|
|
assert "update_mental_model" in tools
|
|
assert "delete_mental_model" in tools
|
|
assert "refresh_mental_model" in tools
|
|
|
|
|
|
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 have bank-scoped tools
|
|
assert "retain" in tools
|
|
assert "recall" in tools
|
|
assert "reflect" in tools
|
|
|
|
# Mental model tools should also be present (they're bank-scoped)
|
|
assert "list_mental_models" in tools
|
|
assert "get_mental_model" in tools
|
|
assert "create_mental_model" in tools
|
|
assert "update_mental_model" in tools
|
|
assert "delete_mental_model" in tools
|
|
assert "refresh_mental_model" in tools
|
|
|
|
# Should NOT have bank management tools
|
|
assert "list_banks" not in tools
|
|
assert "create_bank" not in tools
|
|
|
|
|
|
def test_multi_bank_mode_tools_have_bank_id_param(mock_memory):
|
|
"""Test that multi-bank mode tools include bank_id parameter."""
|
|
import inspect
|
|
|
|
from hindsight_api.api.mcp import create_mcp_server
|
|
|
|
mcp_server = create_mcp_server(mock_memory, multi_bank=True)
|
|
tools = mcp_server._tool_manager._tools
|
|
|
|
# All bank-scoped tools should have bank_id parameter in multi-bank mode
|
|
bank_scoped_tools = [
|
|
"retain",
|
|
"recall",
|
|
"reflect",
|
|
"list_mental_models",
|
|
"get_mental_model",
|
|
"create_mental_model",
|
|
"update_mental_model",
|
|
"delete_mental_model",
|
|
"refresh_mental_model",
|
|
]
|
|
for tool_name in bank_scoped_tools:
|
|
tool = tools[tool_name]
|
|
sig = inspect.signature(tool.fn)
|
|
assert "bank_id" in sig.parameters, f"{tool_name} should have bank_id param in multi-bank mode"
|
|
|
|
|
|
def test_single_bank_mode_tools_no_bank_id_param(mock_memory):
|
|
"""Test that single-bank mode tools do NOT include bank_id parameter."""
|
|
import inspect
|
|
|
|
from hindsight_api.api.mcp import create_mcp_server
|
|
|
|
mcp_server = create_mcp_server(mock_memory, multi_bank=False)
|
|
tools = mcp_server._tool_manager._tools
|
|
|
|
# No bank-scoped tool should have bank_id parameter in single-bank mode
|
|
bank_scoped_tools = [
|
|
"retain",
|
|
"recall",
|
|
"reflect",
|
|
"list_mental_models",
|
|
"get_mental_model",
|
|
"create_mental_model",
|
|
"update_mental_model",
|
|
"delete_mental_model",
|
|
"refresh_mental_model",
|
|
]
|
|
for tool_name in bank_scoped_tools:
|
|
tool = tools[tool_name]
|
|
sig = inspect.signature(tool.fn)
|
|
assert "bank_id" not in sig.parameters, f"{tool_name} should NOT have bank_id param in single-bank mode"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
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
|
|
|
|
# Create middleware (single instance)
|
|
middleware = MCPMiddleware(None, mock_memory)
|
|
|
|
# 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
|
|
assert "list_mental_models" in multi_bank_tools
|
|
assert "create_mental_model" 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_mental_models" in single_bank_tools
|
|
assert "create_mental_model" in single_bank_tools
|
|
assert "list_banks" not in single_bank_tools
|
|
assert "create_bank" not in single_bank_tools
|
|
|
|
|
|
def test_global_mcp_enabled_tools_filter_restricts_registered_tools(mock_memory):
|
|
"""Test that global mcp_enabled_tools env setting restricts which tools are registered."""
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from hindsight_api.api.mcp import create_mcp_server
|
|
|
|
mock_cfg = MagicMock()
|
|
mock_cfg.mcp_enabled_tools = ["retain", "recall"]
|
|
|
|
with patch("hindsight_api.api.mcp._get_raw_config", return_value=mock_cfg):
|
|
mcp_server = create_mcp_server(mock_memory, multi_bank=True)
|
|
|
|
tools = mcp_server._tool_manager._tools
|
|
assert "retain" in tools
|
|
assert "recall" in tools
|
|
assert "reflect" not in tools
|
|
assert "list_banks" not in tools
|
|
assert "create_bank" not in tools
|
|
assert "list_mental_models" not in tools
|
|
|
|
|
|
def test_global_mcp_enabled_tools_none_exposes_all_tools(mock_memory):
|
|
"""Test that mcp_enabled_tools=None (default) exposes all tools."""
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from hindsight_api.api.mcp import create_mcp_server
|
|
|
|
mock_cfg = MagicMock()
|
|
mock_cfg.mcp_enabled_tools = None
|
|
|
|
with patch("hindsight_api.api.mcp._get_raw_config", return_value=mock_cfg):
|
|
mcp_server = create_mcp_server(mock_memory, multi_bank=True)
|
|
|
|
tools = mcp_server._tool_manager._tools
|
|
assert "retain" in tools
|
|
assert "recall" in tools
|
|
assert "reflect" in tools
|
|
assert "list_banks" in tools
|
|
assert "create_bank" in tools
|
|
|
|
|
|
def test_global_mcp_enabled_tools_intersects_with_single_bank_mode(mock_memory):
|
|
"""Test that global filter intersects with single-bank mode tool set.
|
|
|
|
list_banks is in the global allowlist but NOT in single-bank mode, so it
|
|
should be absent from the final registered set.
|
|
"""
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from hindsight_api.api.mcp import create_mcp_server
|
|
|
|
mock_cfg = MagicMock()
|
|
mock_cfg.mcp_enabled_tools = ["retain", "recall", "list_banks"]
|
|
|
|
with patch("hindsight_api.api.mcp._get_raw_config", return_value=mock_cfg):
|
|
mcp_server = create_mcp_server(mock_memory, multi_bank=False)
|
|
|
|
tools = mcp_server._tool_manager._tools
|
|
assert "retain" in tools
|
|
assert "recall" in tools
|
|
assert "list_banks" not in tools # single-bank mode excludes it regardless
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_routing_logic_from_url_path():
|
|
"""Test that routing correctly selects server based on URL structure.
|
|
|
|
Simulates the path parsing logic from MCPMiddleware.__call__ after the
|
|
prefix has been stripped. Any first path segment is treated as a bank_id.
|
|
"""
|
|
from hindsight_api.api.mcp import MCPMiddleware
|
|
|
|
# Mock memory
|
|
mock_memory = MagicMock()
|
|
|
|
# Create middleware
|
|
middleware = MCPMiddleware(None, mock_memory)
|
|
|
|
# Simulate different URL patterns and verify routing
|
|
# Path is what remains after stripping the /mcp prefix
|
|
test_cases = [
|
|
# (path_after_prefix_strip, 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"),
|
|
("/sse/", True, "sse", "Bank named 'sse' routes to single-bank"),
|
|
("/messages/", True, "messages", "Bank named 'messages' routes to single-bank"),
|
|
("/", False, None, "Root path, no bank ID"),
|
|
]
|
|
|
|
for path, expected_bank_from_path, expected_bank_id, description in test_cases:
|
|
bank_id = None
|
|
bank_id_from_path = False
|
|
|
|
if path.startswith("/") and len(path) > 1:
|
|
parts = path[1:].split("/", 1)
|
|
if parts[0]:
|
|
bank_id = parts[0]
|
|
bank_id_from_path = True
|
|
|
|
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})"
|