mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-27 22:34:51 -06:00
Patrick Buckley
2bfc0f2c5d
* fix: harden MCP client against misbehaving servers Misbehaving/failed/misconfigured MCP servers could peg CPU at 100% due to anyio cancel-scope busy-loops (SDK #2147), uncancelled orphaned futures, and missing application-layer resilience. Five fixes: 1. Cancel orphaned futures on timeout — future.cancel() in all sync bridge methods prevents coroutine accumulation on the event loop 2. Per-server circuit breaker — 3-failure threshold with exponential cooldown (30s–5min), per-server jitter, auto-reconnect on half-open probe, McpError excluded (protocol errors from healthy servers) 3. Safe transport stream pre-close — store stream refs and close them before stack teardown in all error/shutdown paths, preventing the anyio zero-buffer CPU busy-loop 4. Notification debounce — 5s per-server rate limit on list_changed refresh storms from buggy servers 5. Periodic refresh backoff with auto-reconnect — disconnected servers get reconnection attempts with exponential backoff (60s–1hr) instead of being silently skipped forever * docs: add MCP resilience section to architecture docs and diagram Document the circuit breaker, future cancellation, stream pre-close, notification debounce, and periodic refresh backoff in the architecture guide and the MCP architecture PlantUML diagram. * fix: address review — stack leak on transport error, half-open comment - Widen _connect_one guard to check _per_server_stacks too, not just _sessions. Transport errors in sync dispatch methods evict the session but left the stack behind, leaking anyio tasks on reconnect. - Clarify half-open design: multiple callers are intentionally allowed through (reconnects serialize on the event loop, first failure re-trips).
1962 lines
73 KiB
Python
1962 lines
73 KiB
Python
"""Tests for turnstone.core.mcp_client — MCP client manager and config loading."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import concurrent.futures
|
|
import json
|
|
import time
|
|
from contextlib import AsyncExitStack
|
|
from typing import Any
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from turnstone.core.mcp_client import (
|
|
MCPClientManager,
|
|
_mcp_to_openai,
|
|
load_mcp_config,
|
|
)
|
|
from turnstone.core.tools import TOOLS, merge_mcp_tools
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _fake_mcp_tool(name: str = "search", description: str = "Search stuff") -> MagicMock:
|
|
"""Create a mock MCP tool object matching the SDK's Tool type."""
|
|
tool = MagicMock()
|
|
tool.name = name
|
|
tool.description = description
|
|
tool.inputSchema = {
|
|
"type": "object",
|
|
"properties": {"query": {"type": "string"}},
|
|
"required": ["query"],
|
|
}
|
|
return tool
|
|
|
|
|
|
def _fake_openai_tool(name: str = "mcp__test__search") -> dict[str, Any]:
|
|
"""Create a fake OpenAI-format tool dict."""
|
|
return {
|
|
"type": "function",
|
|
"function": {
|
|
"name": name,
|
|
"description": "[MCP: test] Search stuff",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {"query": {"type": "string"}},
|
|
"required": ["query"],
|
|
},
|
|
},
|
|
}
|
|
|
|
|
|
def _fake_mcp_resource(
|
|
uri: str = "file:///README.md",
|
|
name: str = "readme",
|
|
description: str = "Project readme",
|
|
mime_type: str = "text/plain",
|
|
) -> MagicMock:
|
|
"""Create a mock MCP Resource object matching the SDK's Resource type."""
|
|
res = MagicMock()
|
|
res.uri = uri
|
|
res.name = name
|
|
res.description = description
|
|
res.mimeType = mime_type
|
|
return res
|
|
|
|
|
|
def _fake_resource_dict(
|
|
uri: str = "file:///README.md",
|
|
name: str = "readme",
|
|
description: str = "Project readme",
|
|
mime_type: str = "text/plain",
|
|
server: str = "test",
|
|
) -> dict[str, Any]:
|
|
"""Create a fake resource dict as stored in per-server state."""
|
|
return {
|
|
"uri": uri,
|
|
"name": name,
|
|
"description": description,
|
|
"mimeType": mime_type,
|
|
"server": server,
|
|
}
|
|
|
|
|
|
def _fake_mcp_prompt(
|
|
name: str = "code_review",
|
|
description: str = "Generate a code review",
|
|
arguments: list[dict[str, Any]] | None = None,
|
|
) -> MagicMock:
|
|
"""Create a mock MCP Prompt object matching the SDK's Prompt type."""
|
|
prompt = MagicMock()
|
|
prompt.name = name
|
|
prompt.description = description
|
|
if arguments is None:
|
|
arg = MagicMock()
|
|
arg.name = "language"
|
|
arg.description = "Programming language"
|
|
arg.required = True
|
|
prompt.arguments = [arg]
|
|
else:
|
|
mock_args = []
|
|
for a in arguments:
|
|
arg = MagicMock()
|
|
arg.name = a["name"]
|
|
arg.description = a.get("description", "")
|
|
arg.required = a.get("required", False)
|
|
mock_args.append(arg)
|
|
prompt.arguments = mock_args
|
|
return prompt
|
|
|
|
|
|
def _fake_prompt_dict(
|
|
name: str = "mcp__test__code_review",
|
|
original_name: str = "code_review",
|
|
server: str = "test",
|
|
description: str = "Generate a code review",
|
|
) -> dict[str, Any]:
|
|
"""Create a fake prompt dict as stored in per-server state."""
|
|
return {
|
|
"name": name,
|
|
"original_name": original_name,
|
|
"server": server,
|
|
"description": description,
|
|
"arguments": [
|
|
{"name": "language", "description": "Programming language", "required": True}
|
|
],
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Schema conversion
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestMcpToOpenai:
|
|
def test_basic_conversion(self):
|
|
tool = _fake_mcp_tool("search_repos", "Search GitHub repos")
|
|
result = _mcp_to_openai("github", tool)
|
|
|
|
assert result["type"] == "function"
|
|
func = result["function"]
|
|
assert func["name"] == "mcp__github__search_repos"
|
|
assert func["description"] == "Search GitHub repos"
|
|
assert func["parameters"]["type"] == "object"
|
|
assert "query" in func["parameters"]["properties"]
|
|
|
|
def test_name_prefixing(self):
|
|
tool = _fake_mcp_tool("list_files")
|
|
result = _mcp_to_openai("fs", tool)
|
|
assert result["function"]["name"] == "mcp__fs__list_files"
|
|
|
|
def test_missing_input_schema(self):
|
|
tool = MagicMock()
|
|
tool.name = "ping"
|
|
tool.description = "Ping the server"
|
|
tool.inputSchema = None
|
|
result = _mcp_to_openai("test", tool)
|
|
assert result["function"]["parameters"] == {"type": "object", "properties": {}}
|
|
|
|
def test_empty_description(self):
|
|
tool = MagicMock()
|
|
tool.name = "noop"
|
|
tool.description = ""
|
|
tool.inputSchema = {"type": "object", "properties": {}}
|
|
result = _mcp_to_openai("test", tool)
|
|
assert result["function"]["description"] == ""
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Config loading
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestLoadMcpConfig:
|
|
def test_load_from_json_file(self, tmp_path):
|
|
config_file = tmp_path / "mcp.json"
|
|
config_file.write_text(
|
|
json.dumps(
|
|
{
|
|
"mcpServers": {
|
|
"github": {
|
|
"command": "npx",
|
|
"args": ["-y", "@modelcontextprotocol/server-github"],
|
|
"env": {"GITHUB_TOKEN": "test"},
|
|
}
|
|
}
|
|
}
|
|
)
|
|
)
|
|
result = load_mcp_config(str(config_file))
|
|
assert "github" in result
|
|
assert result["github"]["command"] == "npx"
|
|
assert result["github"]["env"]["GITHUB_TOKEN"] == "test"
|
|
|
|
def test_load_from_toml(self):
|
|
mock_config = {
|
|
"servers": {
|
|
"postgres": {
|
|
"type": "http",
|
|
"url": "https://mcp.example.com/mcp",
|
|
}
|
|
}
|
|
}
|
|
with patch("turnstone.core.mcp_client.load_config", return_value=mock_config):
|
|
result = load_mcp_config(None)
|
|
assert "postgres" in result
|
|
assert result["postgres"]["url"] == "https://mcp.example.com/mcp"
|
|
|
|
def test_empty_when_no_config(self):
|
|
with patch("turnstone.core.mcp_client.load_config", return_value={}):
|
|
result = load_mcp_config(None)
|
|
assert result == {}
|
|
|
|
def test_json_file_not_found(self, tmp_path):
|
|
with patch("turnstone.core.mcp_client.load_config", return_value={}):
|
|
result = load_mcp_config(str(tmp_path / "nonexistent.json"))
|
|
assert result == {}
|
|
|
|
def test_toml_config_path_redirect(self):
|
|
"""TOML [mcp] config_path redirects to JSON file."""
|
|
# load_config returns a section with config_path pointing to a nonexistent file
|
|
mock_config = {"config_path": "/tmp/nonexistent_mcp.json"}
|
|
with patch("turnstone.core.mcp_client.load_config", return_value=mock_config):
|
|
result = load_mcp_config(None)
|
|
assert result == {}
|
|
|
|
def test_invalid_json(self, tmp_path):
|
|
config_file = tmp_path / "bad.json"
|
|
config_file.write_text("not json")
|
|
with patch("turnstone.core.mcp_client.load_config", return_value={}):
|
|
result = load_mcp_config(str(config_file))
|
|
assert result == {}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# merge_mcp_tools
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestMergeTools:
|
|
def test_merge_preserves_builtin(self):
|
|
mcp_tools = [_fake_openai_tool()]
|
|
merged = merge_mcp_tools(TOOLS, mcp_tools)
|
|
# First N should be built-in
|
|
for i, t in enumerate(TOOLS):
|
|
assert merged[i] is t
|
|
|
|
def test_merge_appends_mcp(self):
|
|
mcp_tools = [_fake_openai_tool("mcp__a__x"), _fake_openai_tool("mcp__b__y")]
|
|
merged = merge_mcp_tools(TOOLS, mcp_tools)
|
|
assert len(merged) == len(TOOLS) + 2
|
|
assert merged[-2]["function"]["name"] == "mcp__a__x"
|
|
assert merged[-1]["function"]["name"] == "mcp__b__y"
|
|
|
|
def test_merge_empty_mcp(self):
|
|
merged = merge_mcp_tools(TOOLS, [])
|
|
assert merged == TOOLS
|
|
|
|
def test_merge_does_not_mutate_input(self):
|
|
mcp_tools = [_fake_openai_tool()]
|
|
original_len = len(TOOLS)
|
|
merge_mcp_tools(TOOLS, mcp_tools)
|
|
assert len(TOOLS) == original_len
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# MCPClientManager unit tests (no real MCP servers)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestMCPClientManager:
|
|
def test_init_state(self):
|
|
mgr = MCPClientManager({"test": {"command": "echo"}})
|
|
assert mgr.get_tools() == []
|
|
assert mgr.is_mcp_tool("anything") is False
|
|
assert mgr.server_count == 0
|
|
|
|
def test_get_tools_returns_copy(self):
|
|
mgr = MCPClientManager({})
|
|
mgr._tools = [_fake_openai_tool()]
|
|
tools = mgr.get_tools()
|
|
assert len(tools) == 1
|
|
tools.clear() # mutate the copy
|
|
assert len(mgr.get_tools()) == 1 # original unchanged
|
|
|
|
def test_is_mcp_tool(self):
|
|
mgr = MCPClientManager({})
|
|
mgr._tool_map["mcp__gh__search"] = ("gh", "search")
|
|
assert mgr.is_mcp_tool("mcp__gh__search") is True
|
|
assert mgr.is_mcp_tool("bash") is False
|
|
|
|
def test_server_count(self):
|
|
mgr = MCPClientManager({})
|
|
mgr._sessions["a"] = MagicMock()
|
|
mgr._sessions["b"] = MagicMock()
|
|
assert mgr.server_count == 2
|
|
|
|
def test_call_tool_sync_unknown_tool(self):
|
|
mgr = MCPClientManager({})
|
|
with pytest.raises(ValueError, match="Unknown MCP tool"):
|
|
mgr.call_tool_sync("mcp__no__such", {})
|
|
|
|
def test_call_tool_sync_disconnected_server(self):
|
|
mgr = MCPClientManager({})
|
|
mgr._tool_map["mcp__dead__ping"] = ("dead", "ping")
|
|
# No session registered for "dead", no config/loop → reconnect fails
|
|
with pytest.raises(RuntimeError, match="not connected"):
|
|
mgr.call_tool_sync("mcp__dead__ping", {})
|
|
|
|
def test_shutdown_on_unstarted_manager(self):
|
|
"""shutdown() should not raise when called on a manager that was never started."""
|
|
mgr = MCPClientManager({})
|
|
mgr.shutdown() # should be a no-op
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Session integration (mock MCP client)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestSessionIntegration:
|
|
@pytest.fixture()
|
|
def tmp_db(self, tmp_path):
|
|
from turnstone.core.storage import init_storage, reset_storage
|
|
|
|
reset_storage()
|
|
init_storage("sqlite", path=str(tmp_path / "test.db"), run_migrations=False)
|
|
yield
|
|
reset_storage()
|
|
|
|
def _make_session(self, mcp_client=None, **kwargs):
|
|
from turnstone.core.session import ChatSession
|
|
|
|
defaults: dict[str, Any] = dict(
|
|
client=MagicMock(),
|
|
model="test-model",
|
|
ui=MagicMock(),
|
|
instructions=None,
|
|
temperature=0.5,
|
|
max_tokens=4096,
|
|
tool_timeout=30,
|
|
mcp_client=mcp_client,
|
|
)
|
|
defaults.update(kwargs)
|
|
return ChatSession(**defaults)
|
|
|
|
def test_session_without_mcp(self, tmp_db):
|
|
session = self._make_session(mcp_client=None)
|
|
assert session._tools is TOOLS
|
|
assert session._mcp_client is None
|
|
|
|
def test_session_with_mcp(self, tmp_db):
|
|
mock_mcp = MagicMock()
|
|
mock_mcp.get_tools.return_value = [_fake_openai_tool()]
|
|
session = self._make_session(mcp_client=mock_mcp)
|
|
assert len(session._tools) == len(TOOLS) + 1
|
|
assert session._tools[-1]["function"]["name"] == "mcp__test__search"
|
|
|
|
def test_task_tools_include_mcp(self, tmp_db):
|
|
from turnstone.core.tools import TASK_AGENT_TOOLS
|
|
|
|
mock_mcp = MagicMock()
|
|
mock_mcp.get_tools.return_value = [_fake_openai_tool()]
|
|
session = self._make_session(mcp_client=mock_mcp)
|
|
assert len(session._task_tools) == len(TASK_AGENT_TOOLS) + 1
|
|
|
|
def test_agent_tools_include_mcp(self, tmp_db):
|
|
from turnstone.core.tools import AGENT_TOOLS
|
|
|
|
mock_mcp = MagicMock()
|
|
mock_mcp.get_tools.return_value = [_fake_openai_tool()]
|
|
session = self._make_session(mcp_client=mock_mcp)
|
|
assert len(session._agent_tools) == len(AGENT_TOOLS) + 1
|
|
|
|
def test_prepare_mcp_tool(self, tmp_db):
|
|
mock_mcp = MagicMock()
|
|
mock_mcp.get_tools.return_value = [_fake_openai_tool()]
|
|
mock_mcp.is_mcp_tool.return_value = True
|
|
session = self._make_session(mcp_client=mock_mcp)
|
|
|
|
tc = {
|
|
"id": "call_123",
|
|
"function": {
|
|
"name": "mcp__test__search",
|
|
"arguments": '{"query": "hello"}',
|
|
},
|
|
}
|
|
prepared = session._prepare_tool(tc)
|
|
assert prepared["func_name"] == "mcp__test__search"
|
|
assert prepared["needs_approval"] is True
|
|
assert "mcp:test/search" in prepared["header"]
|
|
assert callable(prepared["execute"])
|
|
|
|
def test_unknown_tool_without_mcp(self, tmp_db):
|
|
session = self._make_session(mcp_client=None)
|
|
tc = {
|
|
"id": "call_456",
|
|
"function": {"name": "nonexistent", "arguments": "{}"},
|
|
}
|
|
prepared = session._prepare_tool(tc)
|
|
assert "error" in prepared
|
|
assert "Unknown tool" in prepared["error"]
|
|
# Error lists available tools so the model can self-correct
|
|
assert "bash" in prepared["error"]
|
|
# Surfaces warning to user
|
|
session.ui.on_error.assert_called_once()
|
|
assert "nonexistent" in session.ui.on_error.call_args[0][0]
|
|
|
|
def test_prepare_tool_strips_whitespace_from_name(self, tmp_db):
|
|
"""Local models may produce tool names with leading/trailing whitespace."""
|
|
session = self._make_session(mcp_client=None)
|
|
tc = {
|
|
"id": "call_strip",
|
|
"function": {"name": " bash\n", "arguments": '{"command": "echo hi"}'},
|
|
}
|
|
prepared = session._prepare_tool(tc)
|
|
assert prepared["func_name"] == "bash"
|
|
assert "error" not in prepared
|
|
|
|
def test_prepare_tool_malformed_json_surfaces_error(self, tmp_db):
|
|
"""Malformed JSON args should surface a warning to the user and
|
|
give the model a hint about expected format."""
|
|
session = self._make_session(mcp_client=None)
|
|
tc = {
|
|
"id": "call_bad",
|
|
"function": {"name": "bash", "arguments": "{command: echo hi}"},
|
|
}
|
|
prepared = session._prepare_tool(tc)
|
|
assert "error" in prepared
|
|
assert "JSON parse error" in prepared["error"]
|
|
assert "command" in prepared["error"] # hint about expected key
|
|
assert "Please retry" in prepared["error"]
|
|
# User-facing warning
|
|
session.ui.on_error.assert_called_once()
|
|
assert "Malformed tool call" in session.ui.on_error.call_args[0][0]
|
|
|
|
def test_ensure_tool_call_ids_dict(self, tmp_db):
|
|
"""_ensure_tool_call_ids fills empty IDs on streaming-style dict."""
|
|
from turnstone.core.session import ChatSession
|
|
|
|
tool_calls_acc = {
|
|
0: {"id": "", "function": {"name": "bash", "arguments": "{}"}},
|
|
1: {"id": "", "function": {"name": "read_file", "arguments": "{}"}},
|
|
}
|
|
ChatSession._ensure_tool_call_ids(tool_calls_acc)
|
|
ids = [tc["id"] for tc in tool_calls_acc.values()]
|
|
assert all(id_.startswith("call_") for id_ in ids)
|
|
assert len(set(ids)) == 2 # unique
|
|
|
|
def test_ensure_tool_call_ids_list(self, tmp_db):
|
|
"""_ensure_tool_call_ids fills empty IDs on list (agent path)."""
|
|
from turnstone.core.session import ChatSession
|
|
|
|
tool_calls = [
|
|
{"id": None, "function": {"name": "bash", "arguments": "{}"}},
|
|
{"id": "call_existing", "function": {"name": "bash", "arguments": "{}"}},
|
|
]
|
|
ChatSession._ensure_tool_call_ids(tool_calls)
|
|
assert tool_calls[0]["id"].startswith("call_")
|
|
assert tool_calls[1]["id"] == "call_existing" # preserved
|
|
|
|
def test_mcp_command_no_client(self, tmp_db):
|
|
session = self._make_session(mcp_client=None)
|
|
session.handle_command("/mcp")
|
|
session.ui.on_info.assert_called_once()
|
|
assert "No MCP servers" in session.ui.on_info.call_args[0][0]
|
|
|
|
def test_mcp_command_with_tools(self, tmp_db):
|
|
mock_mcp = MagicMock()
|
|
mock_mcp.get_tools.return_value = [_fake_openai_tool()]
|
|
session = self._make_session(mcp_client=mock_mcp)
|
|
session.handle_command("/mcp")
|
|
session.ui.on_info.assert_called_once()
|
|
output = session.ui.on_info.call_args[0][0]
|
|
assert "MCP tools (1)" in output
|
|
assert "mcp__test__search" in output
|
|
|
|
def test_exec_mcp_tool(self, tmp_db):
|
|
mock_mcp = MagicMock()
|
|
mock_mcp.get_tools.return_value = [_fake_openai_tool()]
|
|
mock_mcp.is_mcp_tool.return_value = True
|
|
mock_mcp.call_tool_sync.return_value = "result text"
|
|
session = self._make_session(mcp_client=mock_mcp)
|
|
|
|
item = {
|
|
"call_id": "call_789",
|
|
"mcp_func_name": "mcp__test__search",
|
|
"mcp_args": {"query": "hello"},
|
|
}
|
|
call_id, output = session._exec_mcp_tool(item)
|
|
assert call_id == "call_789"
|
|
assert output == "result text"
|
|
mock_mcp.call_tool_sync.assert_called_once_with(
|
|
"mcp__test__search", {"query": "hello"}, timeout=30
|
|
)
|
|
|
|
def test_exec_mcp_tool_error(self, tmp_db):
|
|
mock_mcp = MagicMock()
|
|
mock_mcp.get_tools.return_value = [_fake_openai_tool()]
|
|
mock_mcp.is_mcp_tool.return_value = True
|
|
mock_mcp.call_tool_sync.side_effect = RuntimeError("server crashed")
|
|
session = self._make_session(mcp_client=mock_mcp)
|
|
|
|
item = {
|
|
"call_id": "call_err",
|
|
"mcp_func_name": "mcp__test__search",
|
|
"mcp_args": {"query": "hello"},
|
|
}
|
|
call_id, output = session._exec_mcp_tool(item)
|
|
assert call_id == "call_err"
|
|
assert "MCP tool error" in output
|
|
assert "server crashed" in output
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Server name validation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestServerNameValidation:
|
|
def test_double_underscore_in_name(self):
|
|
"""Server names with __ should be rejected during _connect_one."""
|
|
import asyncio
|
|
|
|
async def _run() -> None:
|
|
mgr = MCPClientManager({"my__bad": {"command": "echo"}})
|
|
async with AsyncExitStack() as stack:
|
|
mgr._exit_stack = stack
|
|
await mgr._connect_one("my__bad", {"command": "echo"})
|
|
# Should not have connected
|
|
assert "my__bad" not in mgr._sessions
|
|
assert mgr.get_tools() == []
|
|
|
|
asyncio.run(_run())
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# create_mcp_client guard
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestCreateMcpClient:
|
|
def test_returns_none_when_no_config(self):
|
|
with patch("turnstone.core.mcp_client.load_mcp_config", return_value={}):
|
|
from turnstone.core.mcp_client import create_mcp_client
|
|
|
|
result = create_mcp_client()
|
|
assert result is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tool refresh — _rebuild_tools, _refresh_server, listeners
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestRebuildTools:
|
|
def test_rebuild_from_per_server(self):
|
|
mgr = MCPClientManager({})
|
|
mgr._per_server_tools = {
|
|
"github": [_fake_openai_tool("mcp__github__search")],
|
|
"slack": [_fake_openai_tool("mcp__slack__send")],
|
|
}
|
|
mgr._rebuild_tools()
|
|
assert len(mgr._tools) == 2
|
|
names = {t["function"]["name"] for t in mgr._tools}
|
|
assert names == {"mcp__github__search", "mcp__slack__send"}
|
|
assert mgr._tool_map["mcp__github__search"] == ("github", "search")
|
|
assert mgr._tool_map["mcp__slack__send"] == ("slack", "send")
|
|
|
|
def test_rebuild_copy_on_write(self):
|
|
mgr = MCPClientManager({})
|
|
mgr._per_server_tools = {"a": [_fake_openai_tool("mcp__a__x")]}
|
|
mgr._rebuild_tools()
|
|
old_tools = mgr._tools
|
|
old_map = mgr._tool_map
|
|
mgr._per_server_tools["b"] = [_fake_openai_tool("mcp__b__y")]
|
|
mgr._rebuild_tools()
|
|
assert mgr._tools is not old_tools
|
|
assert mgr._tool_map is not old_map
|
|
|
|
def test_rebuild_empty(self):
|
|
mgr = MCPClientManager({})
|
|
mgr._per_server_tools = {}
|
|
mgr._rebuild_tools()
|
|
assert mgr._tools == []
|
|
assert mgr._tool_map == {}
|
|
|
|
|
|
class TestRefreshServer:
|
|
@staticmethod
|
|
def _add_empty_resource_prompt_mocks(
|
|
mgr: MCPClientManager, server_name: str, mock_session: MagicMock
|
|
) -> None:
|
|
"""Add empty list_resources/list_prompts mocks so _refresh_server works."""
|
|
mgr._supports_resources[server_name] = True
|
|
mgr._supports_prompts[server_name] = True
|
|
empty_res = MagicMock()
|
|
empty_res.resources = []
|
|
mock_session.list_resources = AsyncMock(return_value=empty_res)
|
|
empty_tmpl = MagicMock()
|
|
empty_tmpl.resourceTemplates = []
|
|
mock_session.list_resource_templates = AsyncMock(return_value=empty_tmpl)
|
|
empty_prompts = MagicMock()
|
|
empty_prompts.prompts = []
|
|
mock_session.list_prompts = AsyncMock(return_value=empty_prompts)
|
|
|
|
def test_refresh_detects_added_tools(self):
|
|
async def _run() -> None:
|
|
mgr = MCPClientManager({})
|
|
mock_session = MagicMock()
|
|
mock_result = MagicMock()
|
|
mock_result.tools = [
|
|
_fake_mcp_tool("search"),
|
|
_fake_mcp_tool("create"), # new tool
|
|
]
|
|
mock_session.list_tools = AsyncMock(return_value=mock_result)
|
|
self._add_empty_resource_prompt_mocks(mgr, "github", mock_session)
|
|
mgr._sessions["github"] = mock_session
|
|
mgr._per_server_tools["github"] = [_fake_openai_tool("mcp__github__search")]
|
|
mgr._rebuild_tools()
|
|
|
|
added, removed = await mgr._refresh_server("github")
|
|
assert "mcp__github__create" in added
|
|
assert removed == []
|
|
assert len(mgr._tools) == 2
|
|
|
|
asyncio.run(_run())
|
|
|
|
def test_refresh_detects_removed_tools(self):
|
|
async def _run() -> None:
|
|
mgr = MCPClientManager({})
|
|
mock_session = MagicMock()
|
|
mock_result = MagicMock()
|
|
mock_result.tools = [] # all tools removed
|
|
mock_session.list_tools = AsyncMock(return_value=mock_result)
|
|
self._add_empty_resource_prompt_mocks(mgr, "github", mock_session)
|
|
mgr._sessions["github"] = mock_session
|
|
mgr._per_server_tools["github"] = [_fake_openai_tool("mcp__github__search")]
|
|
mgr._rebuild_tools()
|
|
|
|
added, removed = await mgr._refresh_server("github")
|
|
assert added == []
|
|
assert "mcp__github__search" in removed
|
|
assert mgr._tools == []
|
|
|
|
asyncio.run(_run())
|
|
|
|
def test_refresh_no_changes(self):
|
|
async def _run() -> None:
|
|
mgr = MCPClientManager({})
|
|
mock_session = MagicMock()
|
|
mock_result = MagicMock()
|
|
mock_result.tools = [_fake_mcp_tool("search")]
|
|
mock_session.list_tools = AsyncMock(return_value=mock_result)
|
|
self._add_empty_resource_prompt_mocks(mgr, "github", mock_session)
|
|
mgr._sessions["github"] = mock_session
|
|
mgr._per_server_tools["github"] = [_fake_openai_tool("mcp__github__search")]
|
|
mgr._rebuild_tools()
|
|
|
|
added, removed = await mgr._refresh_server("github")
|
|
assert added == []
|
|
assert removed == []
|
|
|
|
asyncio.run(_run())
|
|
|
|
def test_refresh_disconnected_raises(self):
|
|
async def _run() -> None:
|
|
mgr = MCPClientManager({})
|
|
with pytest.raises(RuntimeError, match="not connected"):
|
|
await mgr._refresh_server_tools("ghost")
|
|
|
|
asyncio.run(_run())
|
|
|
|
|
|
class TestListeners:
|
|
def test_add_and_notify(self):
|
|
mgr = MCPClientManager({})
|
|
calls: list[int] = []
|
|
mgr.add_listener(lambda: calls.append(1))
|
|
mgr._per_server_tools = {"a": [_fake_openai_tool("mcp__a__x")]}
|
|
mgr._rebuild_tools()
|
|
assert len(calls) == 1
|
|
|
|
def test_remove_listener(self):
|
|
mgr = MCPClientManager({})
|
|
calls: list[int] = []
|
|
cb = lambda: calls.append(1) # noqa: E731
|
|
mgr.add_listener(cb)
|
|
mgr.remove_listener(cb)
|
|
mgr._rebuild_tools()
|
|
assert calls == []
|
|
|
|
def test_remove_nonexistent_listener(self):
|
|
mgr = MCPClientManager({})
|
|
mgr.remove_listener(lambda: None) # should not raise
|
|
|
|
def test_listener_error_does_not_propagate(self):
|
|
mgr = MCPClientManager({})
|
|
mgr.add_listener(lambda: 1 / 0) # will raise ZeroDivisionError
|
|
mgr._rebuild_tools() # should not raise
|
|
|
|
|
|
class TestServerNames:
|
|
def test_server_names_property(self):
|
|
mgr = MCPClientManager({"github": {}, "slack": {}})
|
|
assert sorted(mgr.server_names) == ["github", "slack"]
|
|
|
|
def test_server_names_empty(self):
|
|
mgr = MCPClientManager({})
|
|
assert mgr.server_names == []
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Session integration — tool refresh propagation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestSessionRefresh:
|
|
@pytest.fixture()
|
|
def tmp_db(self, tmp_path):
|
|
from turnstone.core.storage import init_storage, reset_storage
|
|
|
|
reset_storage()
|
|
init_storage("sqlite", path=str(tmp_path / "test.db"), run_migrations=False)
|
|
yield
|
|
reset_storage()
|
|
|
|
def _make_session(self, mcp_client=None, **kwargs):
|
|
from turnstone.core.session import ChatSession
|
|
|
|
defaults: dict[str, Any] = dict(
|
|
client=MagicMock(),
|
|
model="test-model",
|
|
ui=MagicMock(),
|
|
instructions=None,
|
|
temperature=0.5,
|
|
max_tokens=4096,
|
|
tool_timeout=30,
|
|
mcp_client=mcp_client,
|
|
)
|
|
defaults.update(kwargs)
|
|
return ChatSession(**defaults)
|
|
|
|
def test_listener_registered_on_init(self, tmp_db):
|
|
mock_mcp = MagicMock()
|
|
mock_mcp.get_tools.return_value = []
|
|
session = self._make_session(mcp_client=mock_mcp)
|
|
mock_mcp.add_listener.assert_called_once()
|
|
assert session._mcp_refresh_cb is not None
|
|
|
|
def test_no_listener_without_mcp(self, tmp_db):
|
|
session = self._make_session(mcp_client=None)
|
|
assert session._mcp_refresh_cb is None
|
|
|
|
def test_close_removes_listener(self, tmp_db):
|
|
mock_mcp = MagicMock()
|
|
mock_mcp.get_tools.return_value = []
|
|
session = self._make_session(mcp_client=mock_mcp)
|
|
session.close()
|
|
mock_mcp.remove_listener.assert_called_once()
|
|
assert session._mcp_refresh_cb is None
|
|
|
|
def test_close_idempotent(self, tmp_db):
|
|
mock_mcp = MagicMock()
|
|
mock_mcp.get_tools.return_value = []
|
|
session = self._make_session(mcp_client=mock_mcp)
|
|
session.close()
|
|
session.close() # should not raise
|
|
assert mock_mcp.remove_listener.call_count == 1
|
|
|
|
def test_on_mcp_tools_changed_rebuilds_tools(self, tmp_db):
|
|
mock_mcp = MagicMock()
|
|
mock_mcp.get_tools.return_value = [_fake_openai_tool("mcp__test__a")]
|
|
session = self._make_session(mcp_client=mock_mcp)
|
|
initial_count = len(session._tools)
|
|
|
|
# Simulate a tool refresh — MCP now has 2 tools
|
|
mock_mcp.get_tools.return_value = [
|
|
_fake_openai_tool("mcp__test__a"),
|
|
_fake_openai_tool("mcp__test__b"),
|
|
]
|
|
session._on_mcp_tools_changed()
|
|
assert len(session._tools) == initial_count + 1
|
|
|
|
def test_tool_search_preserved_across_refresh(self, tmp_db):
|
|
# Create enough MCP tools to trigger tool search
|
|
mcp_tools = [_fake_openai_tool(f"mcp__srv__tool{i}") for i in range(25)]
|
|
mock_mcp = MagicMock()
|
|
mock_mcp.get_tools.return_value = mcp_tools
|
|
session = self._make_session(
|
|
mcp_client=mock_mcp,
|
|
tool_search="auto",
|
|
tool_search_threshold=20,
|
|
)
|
|
assert session._tool_search is not None
|
|
|
|
# Expand a tool
|
|
session._tool_search.expand_visible(["mcp__srv__tool0"])
|
|
assert "mcp__srv__tool0" in session._tool_search.get_expanded_names()
|
|
|
|
# Refresh with same tools
|
|
session._on_mcp_tools_changed()
|
|
assert session._tool_search is not None
|
|
assert "mcp__srv__tool0" in session._tool_search.get_expanded_names()
|
|
|
|
def test_tool_search_prunes_removed_from_expanded(self, tmp_db):
|
|
mcp_tools = [_fake_openai_tool(f"mcp__srv__tool{i}") for i in range(25)]
|
|
mock_mcp = MagicMock()
|
|
mock_mcp.get_tools.return_value = mcp_tools
|
|
session = self._make_session(
|
|
mcp_client=mock_mcp,
|
|
tool_search="auto",
|
|
tool_search_threshold=20,
|
|
)
|
|
session._tool_search.expand_visible(["mcp__srv__tool0"])
|
|
|
|
# Refresh with tool0 removed
|
|
new_tools = [_fake_openai_tool(f"mcp__srv__tool{i}") for i in range(1, 25)]
|
|
mock_mcp.get_tools.return_value = new_tools
|
|
session._on_mcp_tools_changed()
|
|
# tool0 was removed, so it should no longer be in expanded
|
|
expanded = session._tool_search.get_expanded_names()
|
|
assert "mcp__srv__tool0" not in expanded
|
|
|
|
def test_mcp_refresh_command(self, tmp_db):
|
|
mock_mcp = MagicMock()
|
|
mock_mcp.get_tools.return_value = [_fake_openai_tool()]
|
|
mock_mcp.server_names = ["test"]
|
|
mock_mcp.refresh_sync.return_value = {"test": (["mcp__test__new"], [])}
|
|
session = self._make_session(mcp_client=mock_mcp)
|
|
|
|
session.handle_command("/mcp refresh")
|
|
mock_mcp.refresh_sync.assert_called_once_with(None)
|
|
session.ui.on_info.assert_called()
|
|
|
|
def test_mcp_refresh_specific_server(self, tmp_db):
|
|
mock_mcp = MagicMock()
|
|
mock_mcp.get_tools.return_value = [_fake_openai_tool()]
|
|
mock_mcp.server_names = ["github", "slack"]
|
|
mock_mcp.refresh_sync.return_value = {"github": ([], [])}
|
|
session = self._make_session(mcp_client=mock_mcp)
|
|
|
|
session.handle_command("/mcp refresh github")
|
|
mock_mcp.refresh_sync.assert_called_once_with("github")
|
|
|
|
def test_mcp_refresh_unknown_server(self, tmp_db):
|
|
mock_mcp = MagicMock()
|
|
mock_mcp.get_tools.return_value = [_fake_openai_tool()]
|
|
mock_mcp.server_names = ["github"]
|
|
session = self._make_session(mcp_client=mock_mcp)
|
|
|
|
session.handle_command("/mcp refresh nonexistent")
|
|
session.ui.on_error.assert_called_once()
|
|
assert "Unknown MCP server" in session.ui.on_error.call_args[0][0]
|
|
|
|
def test_mcp_refresh_error_handling(self, tmp_db):
|
|
mock_mcp = MagicMock()
|
|
mock_mcp.get_tools.return_value = [_fake_openai_tool()]
|
|
mock_mcp.server_names = ["test"]
|
|
mock_mcp.refresh_sync.side_effect = TimeoutError("timed out")
|
|
session = self._make_session(mcp_client=mock_mcp)
|
|
|
|
session.handle_command("/mcp refresh")
|
|
session.ui.on_error.assert_called_once()
|
|
assert "MCP refresh failed" in session.ui.on_error.call_args[0][0]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# MCP Resources
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestMCPResources:
|
|
def test_resource_discovery(self):
|
|
"""Mock list_resources() returning 2 resources, verify get_resources()."""
|
|
mgr = MCPClientManager({})
|
|
mgr._per_server_resources = {
|
|
"fs": [
|
|
_fake_resource_dict("file:///a.txt", "a", "File A", "text/plain", "fs"),
|
|
_fake_resource_dict("file:///b.txt", "b", "File B", "text/plain", "fs"),
|
|
],
|
|
}
|
|
mgr._rebuild_resources()
|
|
resources = mgr.get_resources()
|
|
assert len(resources) == 2
|
|
uris = {r["uri"] for r in resources}
|
|
assert uris == {"file:///a.txt", "file:///b.txt"}
|
|
assert all(r["server"] == "fs" for r in resources)
|
|
|
|
def test_rebuild_resources_copy_on_write(self):
|
|
"""Verify mutation safety — get_resources() returns independent copy."""
|
|
mgr = MCPClientManager({})
|
|
mgr._per_server_resources = {
|
|
"a": [_fake_resource_dict("file:///x", "x", "", "", "a")],
|
|
}
|
|
mgr._rebuild_resources()
|
|
old_resources = mgr._resources
|
|
old_map = mgr._resource_map
|
|
mgr._per_server_resources["b"] = [_fake_resource_dict("file:///y", "y", "", "", "b")]
|
|
mgr._rebuild_resources()
|
|
assert mgr._resources is not old_resources
|
|
assert mgr._resource_map is not old_map
|
|
|
|
def test_get_resources_returns_copy(self):
|
|
mgr = MCPClientManager({})
|
|
mgr._per_server_resources = {
|
|
"a": [_fake_resource_dict("file:///x", "x", "", "", "a")],
|
|
}
|
|
mgr._rebuild_resources()
|
|
resources = mgr.get_resources()
|
|
assert len(resources) == 1
|
|
resources.clear()
|
|
assert len(mgr.get_resources()) == 1
|
|
|
|
def test_read_resource_sync(self):
|
|
"""Mock session.read_resource(), verify text extraction."""
|
|
mgr = MCPClientManager({})
|
|
mgr._resource_map = {"file:///readme": ("fs", "file:///readme")}
|
|
mock_session = MagicMock()
|
|
mgr._sessions["fs"] = mock_session
|
|
mgr._loop = asyncio.new_event_loop()
|
|
|
|
# Mock the read_resource result
|
|
text_content = MagicMock(spec=["text"])
|
|
text_content.text = "Hello, world!"
|
|
mock_result = MagicMock()
|
|
mock_result.contents = [text_content]
|
|
mock_session.read_resource = AsyncMock(return_value=mock_result)
|
|
|
|
thread = None
|
|
try:
|
|
thread = __import__("threading").Thread(target=mgr._loop.run_forever, daemon=True)
|
|
thread.start()
|
|
output = mgr.read_resource_sync("file:///readme", timeout=5)
|
|
assert output == "Hello, world!"
|
|
mock_session.read_resource.assert_awaited_once_with("file:///readme")
|
|
finally:
|
|
mgr._loop.call_soon_threadsafe(mgr._loop.stop)
|
|
if thread:
|
|
thread.join(timeout=5)
|
|
mgr._loop.close()
|
|
|
|
def test_read_resource_sync_blob(self):
|
|
"""Verify base64 blob extraction."""
|
|
mgr = MCPClientManager({})
|
|
mgr._resource_map = {"file:///img.png": ("fs", "file:///img.png")}
|
|
mock_session = MagicMock()
|
|
mgr._sessions["fs"] = mock_session
|
|
mgr._loop = asyncio.new_event_loop()
|
|
|
|
blob_content = MagicMock(spec=["blob"])
|
|
blob_content.blob = "aGVsbG8="
|
|
mock_result = MagicMock()
|
|
mock_result.contents = [blob_content]
|
|
mock_session.read_resource = AsyncMock(return_value=mock_result)
|
|
|
|
thread = None
|
|
try:
|
|
thread = __import__("threading").Thread(target=mgr._loop.run_forever, daemon=True)
|
|
thread.start()
|
|
output = mgr.read_resource_sync("file:///img.png", timeout=5)
|
|
assert output == "aGVsbG8="
|
|
finally:
|
|
mgr._loop.call_soon_threadsafe(mgr._loop.stop)
|
|
if thread:
|
|
thread.join(timeout=5)
|
|
mgr._loop.close()
|
|
|
|
def test_read_resource_sync_unknown_uri(self):
|
|
mgr = MCPClientManager({})
|
|
with pytest.raises(ValueError, match="Unknown MCP resource"):
|
|
mgr.read_resource_sync("file:///nonexistent")
|
|
|
|
def test_read_resource_sync_disconnected(self):
|
|
mgr = MCPClientManager({})
|
|
mgr._resource_map = {"file:///x": ("dead", "file:///x")}
|
|
with pytest.raises(RuntimeError, match="not connected"):
|
|
mgr.read_resource_sync("file:///x")
|
|
|
|
def test_read_resource_sync_timeout(self):
|
|
"""Verify timeout handling."""
|
|
mgr = MCPClientManager({})
|
|
mgr._resource_map = {"file:///x": ("fs", "file:///x")}
|
|
mock_session = MagicMock()
|
|
mgr._sessions["fs"] = mock_session
|
|
mgr._loop = asyncio.new_event_loop()
|
|
|
|
async def _slow_read(_uri: str) -> None:
|
|
await asyncio.sleep(10)
|
|
|
|
mock_session.read_resource = _slow_read
|
|
|
|
thread = None
|
|
try:
|
|
thread = __import__("threading").Thread(target=mgr._loop.run_forever, daemon=True)
|
|
thread.start()
|
|
with pytest.raises(TimeoutError):
|
|
mgr.read_resource_sync("file:///x", timeout=1)
|
|
finally:
|
|
mgr._loop.call_soon_threadsafe(mgr._loop.stop)
|
|
if thread:
|
|
thread.join(timeout=5)
|
|
mgr._loop.close()
|
|
|
|
def test_resource_listener_notification(self):
|
|
"""Verify callback fires on rebuild."""
|
|
mgr = MCPClientManager({})
|
|
calls: list[int] = []
|
|
mgr.add_resource_listener(lambda: calls.append(1))
|
|
mgr._per_server_resources = {"a": [_fake_resource_dict()]}
|
|
mgr._rebuild_resources()
|
|
assert len(calls) == 1
|
|
|
|
def test_resource_listener_remove(self):
|
|
mgr = MCPClientManager({})
|
|
calls: list[int] = []
|
|
cb = lambda: calls.append(1) # noqa: E731
|
|
mgr.add_resource_listener(cb)
|
|
mgr.remove_resource_listener(cb)
|
|
mgr._rebuild_resources()
|
|
assert calls == []
|
|
|
|
def test_resource_listener_error_does_not_propagate(self):
|
|
mgr = MCPClientManager({})
|
|
mgr.add_resource_listener(lambda: 1 / 0)
|
|
mgr._rebuild_resources() # should not raise
|
|
|
|
def test_resource_refresh_on_notification(self):
|
|
"""Mock notification, verify re-fetch of resources."""
|
|
|
|
async def _run() -> None:
|
|
mgr = MCPClientManager({})
|
|
mock_session = MagicMock()
|
|
mgr._sessions["fs"] = mock_session
|
|
mgr._supports_resources["fs"] = True
|
|
|
|
# Initial state
|
|
mgr._per_server_resources["fs"] = [
|
|
_fake_resource_dict("file:///old", server="fs"),
|
|
]
|
|
mgr._rebuild_resources()
|
|
assert len(mgr.get_resources()) == 1
|
|
|
|
# Mock the re-fetch returning a new resource
|
|
new_res = _fake_mcp_resource("file:///new", "new")
|
|
mock_res_result = MagicMock()
|
|
mock_res_result.resources = [new_res]
|
|
mock_session.list_resources = AsyncMock(return_value=mock_res_result)
|
|
mock_tmpl_result = MagicMock()
|
|
mock_tmpl_result.resourceTemplates = []
|
|
mock_session.list_resource_templates = AsyncMock(return_value=mock_tmpl_result)
|
|
|
|
await mgr._refresh_server_resources("fs")
|
|
resources = mgr.get_resources()
|
|
assert len(resources) == 1
|
|
assert resources[0]["uri"] == "file:///new"
|
|
|
|
asyncio.run(_run())
|
|
|
|
def test_rebuild_resources_empty(self):
|
|
mgr = MCPClientManager({})
|
|
mgr._per_server_resources = {}
|
|
mgr._rebuild_resources()
|
|
assert mgr._resources == []
|
|
assert mgr._resource_map == {}
|
|
|
|
def test_rebuild_resources_multi_server(self):
|
|
mgr = MCPClientManager({})
|
|
mgr._per_server_resources = {
|
|
"fs": [_fake_resource_dict("file:///a", server="fs")],
|
|
"db": [_fake_resource_dict("db://table", name="table", server="db")],
|
|
}
|
|
mgr._rebuild_resources()
|
|
assert len(mgr._resources) == 2
|
|
assert mgr._resource_map["file:///a"] == ("fs", "file:///a")
|
|
assert mgr._resource_map["db://table"] == ("db", "db://table")
|
|
|
|
def test_template_prefix_matching(self):
|
|
"""Expanded URI matches template by prefix."""
|
|
mgr = MCPClientManager({})
|
|
mgr._per_server_resources = {
|
|
"db": [
|
|
{
|
|
"uri": "db://tables/{table}/rows/{id}",
|
|
"name": "row",
|
|
"description": "A row",
|
|
"mimeType": "application/json",
|
|
"server": "db",
|
|
"template": True,
|
|
},
|
|
],
|
|
}
|
|
mgr._rebuild_resources()
|
|
# Template should not be in resource_map
|
|
assert "db://tables/{table}/rows/{id}" not in mgr._resource_map
|
|
# But prefix matching should find it
|
|
result = mgr._match_template("db://tables/users/rows/1")
|
|
assert result is not None
|
|
server, template_uri = result
|
|
assert server == "db"
|
|
assert template_uri == "db://tables/{table}/rows/{id}"
|
|
|
|
def test_template_longest_prefix_wins(self):
|
|
"""When two templates have overlapping prefixes, the longer one wins."""
|
|
mgr = MCPClientManager({})
|
|
# Use templates with genuinely different prefix lengths:
|
|
# "db://data/" (6 chars after scheme) vs "db://data/tables/" (13 chars after scheme)
|
|
mgr._per_server_resources = {
|
|
"short": [
|
|
{
|
|
"uri": "db://data/{collection}",
|
|
"name": "collection",
|
|
"description": "",
|
|
"mimeType": "",
|
|
"server": "short",
|
|
"template": True,
|
|
},
|
|
],
|
|
"long": [
|
|
{
|
|
"uri": "db://data/tables/{table}",
|
|
"name": "table",
|
|
"description": "",
|
|
"mimeType": "",
|
|
"server": "long",
|
|
"template": True,
|
|
},
|
|
],
|
|
}
|
|
mgr._rebuild_resources()
|
|
# "db://data/tables/users" matches both prefixes ("db://data/" and
|
|
# "db://data/tables/") — the longer one should win
|
|
result = mgr._match_template("db://data/tables/users")
|
|
assert result is not None
|
|
server, template_uri = result
|
|
assert server == "long"
|
|
assert template_uri == "db://data/tables/{table}"
|
|
# URI that only matches the short prefix
|
|
result2 = mgr._match_template("db://data/views/active")
|
|
assert result2 is not None
|
|
assert result2[0] == "short"
|
|
|
|
def test_template_no_match_raises(self):
|
|
"""Completely unrelated URI still raises ValueError."""
|
|
mgr = MCPClientManager({})
|
|
mgr._per_server_resources = {
|
|
"db": [
|
|
{
|
|
"uri": "db://tables/{table}",
|
|
"name": "table",
|
|
"description": "",
|
|
"mimeType": "",
|
|
"server": "db",
|
|
"template": True,
|
|
},
|
|
],
|
|
}
|
|
mgr._rebuild_resources()
|
|
assert mgr._match_template("file:///something") is None
|
|
with pytest.raises(ValueError, match="Unknown MCP resource"):
|
|
mgr.read_resource_sync("file:///something")
|
|
|
|
def test_read_resource_sync_with_template_uri(self):
|
|
"""End-to-end: template discovered, expanded URI dispatched to correct server."""
|
|
mgr = MCPClientManager({})
|
|
mgr._per_server_resources = {
|
|
"db": [
|
|
{
|
|
"uri": "db://tables/{table}/rows/{id}",
|
|
"name": "row",
|
|
"description": "A row",
|
|
"mimeType": "application/json",
|
|
"server": "db",
|
|
"template": True,
|
|
},
|
|
],
|
|
}
|
|
mgr._rebuild_resources()
|
|
|
|
mock_session = MagicMock()
|
|
mgr._sessions["db"] = mock_session
|
|
mgr._loop = asyncio.new_event_loop()
|
|
|
|
text_content = MagicMock(spec=["text"])
|
|
text_content.text = '{"name": "Alice"}'
|
|
mock_result = MagicMock()
|
|
mock_result.contents = [text_content]
|
|
mock_session.read_resource = AsyncMock(return_value=mock_result)
|
|
|
|
thread = None
|
|
try:
|
|
thread = __import__("threading").Thread(target=mgr._loop.run_forever, daemon=True)
|
|
thread.start()
|
|
output = mgr.read_resource_sync("db://tables/users/rows/1", timeout=5)
|
|
assert output == '{"name": "Alice"}'
|
|
mock_session.read_resource.assert_awaited_once_with("db://tables/users/rows/1")
|
|
finally:
|
|
mgr._loop.call_soon_threadsafe(mgr._loop.stop)
|
|
if thread:
|
|
thread.join(timeout=5)
|
|
mgr._loop.close()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# MCP Prompts
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestMCPPrompts:
|
|
def test_prompt_discovery(self):
|
|
"""Mock list_prompts(), verify get_prompts() with correct prefixed names."""
|
|
mgr = MCPClientManager({})
|
|
mgr._per_server_prompts = {
|
|
"tmpl": [
|
|
_fake_prompt_dict("mcp__tmpl__code_review", "code_review", "tmpl"),
|
|
_fake_prompt_dict("mcp__tmpl__summarize", "summarize", "tmpl"),
|
|
],
|
|
}
|
|
mgr._rebuild_prompts()
|
|
prompts = mgr.get_prompts()
|
|
assert len(prompts) == 2
|
|
names = {p["name"] for p in prompts}
|
|
assert names == {"mcp__tmpl__code_review", "mcp__tmpl__summarize"}
|
|
# Verify map entries
|
|
assert mgr._prompt_map["mcp__tmpl__code_review"] == ("tmpl", "code_review")
|
|
assert mgr._prompt_map["mcp__tmpl__summarize"] == ("tmpl", "summarize")
|
|
|
|
def test_rebuild_prompts_copy_on_write(self):
|
|
"""Verify mutation safety."""
|
|
mgr = MCPClientManager({})
|
|
mgr._per_server_prompts = {
|
|
"a": [_fake_prompt_dict("mcp__a__p1", "p1", "a")],
|
|
}
|
|
mgr._rebuild_prompts()
|
|
old_prompts = mgr._prompts
|
|
old_map = mgr._prompt_map
|
|
mgr._per_server_prompts["b"] = [_fake_prompt_dict("mcp__b__p2", "p2", "b")]
|
|
mgr._rebuild_prompts()
|
|
assert mgr._prompts is not old_prompts
|
|
assert mgr._prompt_map is not old_map
|
|
|
|
def test_get_prompts_returns_copy(self):
|
|
mgr = MCPClientManager({})
|
|
mgr._per_server_prompts = {
|
|
"a": [_fake_prompt_dict("mcp__a__p1", "p1", "a")],
|
|
}
|
|
mgr._rebuild_prompts()
|
|
prompts = mgr.get_prompts()
|
|
assert len(prompts) == 1
|
|
prompts.clear()
|
|
assert len(mgr.get_prompts()) == 1
|
|
|
|
def test_get_prompt_sync(self):
|
|
"""Mock session.get_prompt(), verify message conversion."""
|
|
mgr = MCPClientManager({})
|
|
mgr._prompt_map = {"mcp__tmpl__review": ("tmpl", "review")}
|
|
mock_session = MagicMock()
|
|
mgr._sessions["tmpl"] = mock_session
|
|
mgr._loop = asyncio.new_event_loop()
|
|
|
|
# Build mock PromptMessage
|
|
msg1 = MagicMock()
|
|
msg1.role = "user"
|
|
msg1.content = MagicMock()
|
|
msg1.content.text = "Review this code"
|
|
msg2 = MagicMock()
|
|
msg2.role = "assistant"
|
|
msg2.content = MagicMock()
|
|
msg2.content.text = "Looks good!"
|
|
mock_result = MagicMock()
|
|
mock_result.messages = [msg1, msg2]
|
|
mock_session.get_prompt = AsyncMock(return_value=mock_result)
|
|
|
|
thread = None
|
|
try:
|
|
thread = __import__("threading").Thread(target=mgr._loop.run_forever, daemon=True)
|
|
thread.start()
|
|
messages = mgr.get_prompt_sync(
|
|
"mcp__tmpl__review", arguments={"language": "python"}, timeout=5
|
|
)
|
|
assert len(messages) == 2
|
|
assert messages[0] == {"role": "user", "content": "Review this code"}
|
|
assert messages[1] == {"role": "assistant", "content": "Looks good!"}
|
|
mock_session.get_prompt.assert_awaited_once_with(
|
|
"review", arguments={"language": "python"}
|
|
)
|
|
finally:
|
|
mgr._loop.call_soon_threadsafe(mgr._loop.stop)
|
|
if thread:
|
|
thread.join(timeout=5)
|
|
mgr._loop.close()
|
|
|
|
def test_get_prompt_sync_unknown(self):
|
|
mgr = MCPClientManager({})
|
|
with pytest.raises(ValueError, match="Unknown MCP prompt"):
|
|
mgr.get_prompt_sync("mcp__no__such")
|
|
|
|
def test_get_prompt_sync_disconnected(self):
|
|
mgr = MCPClientManager({})
|
|
mgr._prompt_map = {"mcp__dead__p": ("dead", "p")}
|
|
with pytest.raises(RuntimeError, match="not connected"):
|
|
mgr.get_prompt_sync("mcp__dead__p")
|
|
|
|
def test_get_prompt_sync_timeout(self):
|
|
"""Verify timeout handling."""
|
|
mgr = MCPClientManager({})
|
|
mgr._prompt_map = {"mcp__tmpl__slow": ("tmpl", "slow")}
|
|
mock_session = MagicMock()
|
|
mgr._sessions["tmpl"] = mock_session
|
|
mgr._loop = asyncio.new_event_loop()
|
|
|
|
async def _slow_prompt(_name: str, *, arguments: dict[str, str] | None = None) -> None:
|
|
await asyncio.sleep(10)
|
|
|
|
mock_session.get_prompt = _slow_prompt
|
|
|
|
thread = None
|
|
try:
|
|
thread = __import__("threading").Thread(target=mgr._loop.run_forever, daemon=True)
|
|
thread.start()
|
|
with pytest.raises(TimeoutError):
|
|
mgr.get_prompt_sync("mcp__tmpl__slow", timeout=1)
|
|
finally:
|
|
mgr._loop.call_soon_threadsafe(mgr._loop.stop)
|
|
if thread:
|
|
thread.join(timeout=5)
|
|
mgr._loop.close()
|
|
|
|
def test_prompt_listener_notification(self):
|
|
"""Verify callback fires on rebuild."""
|
|
mgr = MCPClientManager({})
|
|
calls: list[int] = []
|
|
mgr.add_prompt_listener(lambda: calls.append(1))
|
|
mgr._per_server_prompts = {"a": [_fake_prompt_dict()]}
|
|
mgr._rebuild_prompts()
|
|
assert len(calls) == 1
|
|
|
|
def test_prompt_listener_remove(self):
|
|
mgr = MCPClientManager({})
|
|
calls: list[int] = []
|
|
cb = lambda: calls.append(1) # noqa: E731
|
|
mgr.add_prompt_listener(cb)
|
|
mgr.remove_prompt_listener(cb)
|
|
mgr._rebuild_prompts()
|
|
assert calls == []
|
|
|
|
def test_prompt_listener_error_does_not_propagate(self):
|
|
mgr = MCPClientManager({})
|
|
mgr.add_prompt_listener(lambda: 1 / 0)
|
|
mgr._rebuild_prompts() # should not raise
|
|
|
|
def test_is_mcp_prompt(self):
|
|
"""Verify name lookup."""
|
|
mgr = MCPClientManager({})
|
|
mgr._prompt_map["mcp__tmpl__review"] = ("tmpl", "review")
|
|
assert mgr.is_mcp_prompt("mcp__tmpl__review") is True
|
|
assert mgr.is_mcp_prompt("nonexistent") is False
|
|
|
|
def test_prompt_refresh_on_notification(self):
|
|
"""Mock notification, verify re-fetch of prompts."""
|
|
|
|
async def _run() -> None:
|
|
mgr = MCPClientManager({})
|
|
mock_session = MagicMock()
|
|
mgr._sessions["tmpl"] = mock_session
|
|
mgr._supports_prompts["tmpl"] = True
|
|
|
|
# Initial state
|
|
mgr._per_server_prompts["tmpl"] = [
|
|
_fake_prompt_dict("mcp__tmpl__old", "old", "tmpl"),
|
|
]
|
|
mgr._rebuild_prompts()
|
|
assert len(mgr.get_prompts()) == 1
|
|
|
|
# Mock re-fetch returning a new prompt
|
|
new_prompt = _fake_mcp_prompt("new_prompt", "A new prompt")
|
|
mock_prompt_result = MagicMock()
|
|
mock_prompt_result.prompts = [new_prompt]
|
|
mock_session.list_prompts = AsyncMock(return_value=mock_prompt_result)
|
|
|
|
await mgr._refresh_server_prompts("tmpl")
|
|
prompts = mgr.get_prompts()
|
|
assert len(prompts) == 1
|
|
assert prompts[0]["name"] == "mcp__tmpl__new_prompt"
|
|
assert prompts[0]["original_name"] == "new_prompt"
|
|
|
|
asyncio.run(_run())
|
|
|
|
def test_rebuild_prompts_empty(self):
|
|
mgr = MCPClientManager({})
|
|
mgr._per_server_prompts = {}
|
|
mgr._rebuild_prompts()
|
|
assert mgr._prompts == []
|
|
assert mgr._prompt_map == {}
|
|
|
|
def test_rebuild_prompts_multi_server(self):
|
|
mgr = MCPClientManager({})
|
|
mgr._per_server_prompts = {
|
|
"a": [_fake_prompt_dict("mcp__a__p1", "p1", "a")],
|
|
"b": [_fake_prompt_dict("mcp__b__p2", "p2", "b")],
|
|
}
|
|
mgr._rebuild_prompts()
|
|
assert len(mgr._prompts) == 2
|
|
assert mgr._prompt_map["mcp__a__p1"] == ("a", "p1")
|
|
assert mgr._prompt_map["mcp__b__p2"] == ("b", "p2")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Shutdown cleans up new state
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestShutdownCleanup:
|
|
def test_shutdown_clears_resources_and_prompts(self):
|
|
mgr = MCPClientManager({})
|
|
mgr._per_server_resources = {"a": [_fake_resource_dict()]}
|
|
mgr._rebuild_resources()
|
|
mgr._per_server_prompts = {"a": [_fake_prompt_dict()]}
|
|
mgr._rebuild_prompts()
|
|
assert mgr.get_resources() != []
|
|
assert mgr.get_prompts() != []
|
|
|
|
mgr.shutdown()
|
|
assert mgr.get_resources() == []
|
|
assert mgr.get_prompts() == []
|
|
assert mgr._resource_map == {}
|
|
assert mgr._prompt_map == {}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TCP probe and unreachable server handling
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestTCPProbe:
|
|
"""MCPClientManager._tcp_probe should fail fast on unreachable servers."""
|
|
|
|
def test_tcp_probe_unreachable_raises_connection_error(self):
|
|
"""Unreachable host raises ConnectionError, not TimeoutError."""
|
|
mgr = MCPClientManager({})
|
|
|
|
async def _run():
|
|
with pytest.raises(ConnectionError, match="unreachable"):
|
|
await mgr._tcp_probe("test-server", "http://127.0.0.1:1")
|
|
|
|
asyncio.run(_run())
|
|
|
|
def test_tcp_probe_parses_url_correctly(self):
|
|
"""Port and host are extracted from the URL."""
|
|
mgr = MCPClientManager({})
|
|
|
|
async def _run():
|
|
# Non-routable port — should fail with ConnectionError
|
|
with pytest.raises(ConnectionError):
|
|
await mgr._tcp_probe("srv", "https://127.0.0.1:1/mcp")
|
|
|
|
asyncio.run(_run())
|
|
|
|
def test_tcp_probe_default_port_http(self):
|
|
"""Default port 80 used for http:// URLs without explicit port."""
|
|
mgr = MCPClientManager({})
|
|
|
|
async def _run():
|
|
# Will fail (nothing on port 80), but should not crash on parsing
|
|
with pytest.raises(ConnectionError):
|
|
await mgr._tcp_probe("srv", "http://127.0.0.1")
|
|
|
|
asyncio.run(_run())
|
|
|
|
def test_tcp_probe_dns_failure(self):
|
|
"""Unresolvable hostname raises ConnectionError."""
|
|
mgr = MCPClientManager({})
|
|
|
|
async def _run():
|
|
with pytest.raises(ConnectionError):
|
|
await mgr._tcp_probe("srv", "http://this.host.does.not.exist.invalid:8080/mcp")
|
|
|
|
asyncio.run(_run())
|
|
|
|
|
|
class TestConnectOneUnreachable:
|
|
"""_connect_one should handle unreachable HTTP servers gracefully."""
|
|
|
|
def test_unreachable_http_server_raises_connection_error(self):
|
|
"""Unreachable HTTP MCP server raises ConnectionError without spinning."""
|
|
mgr = MCPClientManager({})
|
|
mgr._loop = asyncio.new_event_loop()
|
|
|
|
async def _run():
|
|
with pytest.raises(ConnectionError, match="unreachable"):
|
|
await mgr._connect_one(
|
|
"bad-server",
|
|
{
|
|
"type": "http",
|
|
"url": "http://127.0.0.1:1/mcp",
|
|
},
|
|
)
|
|
|
|
mgr._loop.run_until_complete(_run())
|
|
mgr._loop.close()
|
|
|
|
# Server should NOT be in sessions (connection failed)
|
|
assert "bad-server" not in mgr._sessions
|
|
|
|
def test_connect_all_continues_after_unreachable_server(self):
|
|
"""_connect_all logs error and continues to next server."""
|
|
mgr = MCPClientManager(
|
|
{
|
|
"bad": {"type": "http", "url": "http://127.0.0.1:1/mcp"},
|
|
}
|
|
)
|
|
|
|
loop = asyncio.new_event_loop()
|
|
loop.run_until_complete(mgr._connect_all())
|
|
loop.close()
|
|
|
|
assert "bad" not in mgr._sessions
|
|
assert "bad" in mgr._last_error
|
|
|
|
|
|
class TestSafeCloseStack:
|
|
"""_safe_close_stack should suppress errors from broken anyio scopes."""
|
|
|
|
def test_suppresses_runtime_error(self):
|
|
"""RuntimeError from broken cancel scope is suppressed."""
|
|
|
|
async def _run():
|
|
stack = AsyncExitStack()
|
|
await stack.__aenter__()
|
|
|
|
# Simulate a broken close that raises RuntimeError
|
|
async def _broken_close():
|
|
raise RuntimeError("Attempted to exit cancel scope in a different task")
|
|
|
|
stack.aclose = _broken_close
|
|
# Should not raise
|
|
await MCPClientManager._safe_close_stack(stack)
|
|
|
|
asyncio.run(_run())
|
|
|
|
def test_suppresses_cancelled_error(self):
|
|
"""CancelledError during close is suppressed."""
|
|
|
|
async def _run():
|
|
stack = AsyncExitStack()
|
|
await stack.__aenter__()
|
|
|
|
async def _cancel_close():
|
|
raise asyncio.CancelledError()
|
|
|
|
stack.aclose = _cancel_close
|
|
await MCPClientManager._safe_close_stack(stack)
|
|
|
|
asyncio.run(_run())
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fix 1: Cancel orphaned futures on timeout
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestFutureCancellation:
|
|
"""Verify future.cancel() is called when sync bridge methods time out."""
|
|
|
|
def _make_manager_with_session(self) -> MCPClientManager:
|
|
mgr = MCPClientManager({"test": {"type": "stdio", "command": "echo"}})
|
|
mock_session = MagicMock()
|
|
# Prevent auto-spec from creating async coroutines that trigger warnings
|
|
mock_session.call_tool = MagicMock(return_value="sentinel")
|
|
mock_session.read_resource = MagicMock(return_value="sentinel")
|
|
mock_session.get_prompt = MagicMock(return_value="sentinel")
|
|
mgr._sessions["test"] = mock_session
|
|
mgr._loop = MagicMock()
|
|
mgr._tool_map["mcp__test__search"] = ("test", "search")
|
|
mgr._resource_map["file:///a.txt"] = ("test", "file:///a.txt")
|
|
mgr._prompt_map["mcp__test__review"] = ("test", "review")
|
|
return mgr
|
|
|
|
def test_call_tool_sync_cancels_future_on_timeout(self):
|
|
mgr = self._make_manager_with_session()
|
|
mock_future = MagicMock()
|
|
mock_future.result.side_effect = concurrent.futures.TimeoutError()
|
|
with (
|
|
patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
|
|
pytest.raises(TimeoutError, match="timed out"),
|
|
):
|
|
mgr.call_tool_sync("mcp__test__search", {"query": "x"}, timeout=1)
|
|
mock_future.cancel.assert_called_once()
|
|
|
|
def test_read_resource_sync_cancels_future_on_timeout(self):
|
|
mgr = self._make_manager_with_session()
|
|
mock_future = MagicMock()
|
|
mock_future.result.side_effect = concurrent.futures.TimeoutError()
|
|
with (
|
|
patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
|
|
pytest.raises(TimeoutError, match="timed out"),
|
|
):
|
|
mgr.read_resource_sync("file:///a.txt", timeout=1)
|
|
mock_future.cancel.assert_called_once()
|
|
|
|
def test_get_prompt_sync_cancels_future_on_timeout(self):
|
|
mgr = self._make_manager_with_session()
|
|
mock_future = MagicMock()
|
|
mock_future.result.side_effect = concurrent.futures.TimeoutError()
|
|
with (
|
|
patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
|
|
pytest.raises(TimeoutError, match="timed out"),
|
|
):
|
|
mgr.get_prompt_sync("mcp__test__review", timeout=1)
|
|
mock_future.cancel.assert_called_once()
|
|
|
|
def test_refresh_sync_cancels_future_on_timeout(self):
|
|
mgr = MCPClientManager({})
|
|
mgr._loop = MagicMock()
|
|
mock_future = MagicMock()
|
|
mock_future.result.side_effect = concurrent.futures.TimeoutError()
|
|
with (
|
|
patch.object(mgr, "_refresh_all", return_value=MagicMock()),
|
|
patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
|
|
pytest.raises(TimeoutError, match="timed out"),
|
|
):
|
|
mgr.refresh_sync(timeout=1)
|
|
mock_future.cancel.assert_called_once()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fix 2: Per-server circuit breaker
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestCircuitBreaker:
|
|
"""Verify per-server circuit breaker behavior."""
|
|
|
|
def test_circuit_stays_closed_below_threshold(self):
|
|
mgr = MCPClientManager({})
|
|
mgr._cb_record_failure("srv")
|
|
mgr._cb_record_failure("srv")
|
|
is_open, _ = mgr._cb_check("srv")
|
|
assert not is_open
|
|
|
|
def test_circuit_opens_at_threshold(self):
|
|
mgr = MCPClientManager({})
|
|
for _ in range(3):
|
|
mgr._cb_record_failure("srv")
|
|
is_open, cooldown_expired = mgr._cb_check("srv")
|
|
assert is_open
|
|
assert not cooldown_expired # just opened, cooldown not expired
|
|
|
|
def test_circuit_half_open_after_cooldown(self):
|
|
mgr = MCPClientManager({})
|
|
for _ in range(3):
|
|
mgr._cb_record_failure("srv")
|
|
# Simulate cooldown expiry
|
|
mgr._circuit_open_until["srv"] = time.monotonic() - 1
|
|
is_open, cooldown_expired = mgr._cb_check("srv")
|
|
assert is_open
|
|
assert cooldown_expired
|
|
|
|
def test_circuit_resets_on_success(self):
|
|
mgr = MCPClientManager({})
|
|
for _ in range(3):
|
|
mgr._cb_record_failure("srv")
|
|
assert "srv" in mgr._circuit_open_until
|
|
mgr._cb_record_success("srv")
|
|
is_open, _ = mgr._cb_check("srv")
|
|
assert not is_open
|
|
assert mgr._consecutive_failures.get("srv") is None
|
|
|
|
def test_success_decays_trip_count(self):
|
|
"""Success decays trip_count by 1 so flapping servers escalate backoff."""
|
|
mgr = MCPClientManager({})
|
|
mgr._circuit_trip_count["srv"] = 3
|
|
mgr._cb_record_success("srv")
|
|
assert mgr._circuit_trip_count["srv"] == 2
|
|
mgr._cb_record_success("srv")
|
|
assert mgr._circuit_trip_count["srv"] == 1
|
|
mgr._cb_record_success("srv")
|
|
assert "srv" not in mgr._circuit_trip_count
|
|
|
|
def test_cooldown_is_exponential(self):
|
|
mgr = MCPClientManager({})
|
|
# First trip (trip_count starts at 0)
|
|
for _ in range(3):
|
|
mgr._cb_record_failure("srv")
|
|
deadline1 = mgr._circuit_open_until["srv"]
|
|
base1 = deadline1 - time.monotonic()
|
|
# Reset circuit but keep trip_count at 1 (set by first trip)
|
|
mgr._cb_record_success("srv")
|
|
# trip_count decayed from 1 to 0 — manually set to 1 for test
|
|
mgr._circuit_trip_count["srv"] = 1
|
|
for _ in range(3):
|
|
mgr._cb_record_failure("srv")
|
|
deadline2 = mgr._circuit_open_until["srv"]
|
|
base2 = deadline2 - time.monotonic()
|
|
# Second trip should have longer cooldown (roughly 2x, within jitter)
|
|
assert base2 > base1 * 1.5
|
|
|
|
def test_cooldown_capped_at_max(self):
|
|
mgr = MCPClientManager({})
|
|
mgr._circuit_trip_count["srv"] = 100 # very high trip count
|
|
for _ in range(3):
|
|
mgr._cb_record_failure("srv")
|
|
deadline = mgr._circuit_open_until["srv"]
|
|
cooldown = deadline - time.monotonic()
|
|
# Should not exceed max (300s) + 10% jitter = 330s
|
|
assert cooldown <= mgr._CB_MAX_COOLDOWN * 1.11
|
|
|
|
def test_cb_gate_rejects_when_open(self):
|
|
mgr = MCPClientManager({})
|
|
for _ in range(3):
|
|
mgr._cb_record_failure("srv")
|
|
with pytest.raises(RuntimeError, match="circuit open"):
|
|
mgr._cb_gate("srv")
|
|
|
|
def test_cb_gate_allows_after_cooldown(self):
|
|
mgr = MCPClientManager({})
|
|
for _ in range(3):
|
|
mgr._cb_record_failure("srv")
|
|
mgr._circuit_open_until["srv"] = time.monotonic() - 1
|
|
# Should not raise
|
|
mgr._cb_gate("srv")
|
|
# Deadline should be removed (half-open probe allowed)
|
|
assert "srv" not in mgr._circuit_open_until
|
|
|
|
def test_cb_clear_removes_all_state(self):
|
|
mgr = MCPClientManager({})
|
|
for _ in range(3):
|
|
mgr._cb_record_failure("srv")
|
|
mgr._cb_clear("srv")
|
|
assert "srv" not in mgr._consecutive_failures
|
|
assert "srv" not in mgr._circuit_open_until
|
|
assert "srv" not in mgr._circuit_trip_count
|
|
|
|
@pytest.mark.filterwarnings("ignore::pytest.PytestUnraisableExceptionWarning")
|
|
@pytest.mark.filterwarnings("ignore:coroutine.*was never awaited:RuntimeWarning")
|
|
def test_call_tool_sync_records_failure_on_timeout(self):
|
|
mgr = MCPClientManager({"test": {"type": "stdio", "command": "echo"}})
|
|
mock_session = MagicMock()
|
|
mock_session.call_tool = MagicMock(return_value="sentinel")
|
|
mgr._sessions["test"] = mock_session
|
|
mgr._loop = MagicMock()
|
|
mgr._tool_map["mcp__test__ping"] = ("test", "ping")
|
|
mock_future = MagicMock()
|
|
mock_future.result.side_effect = concurrent.futures.TimeoutError()
|
|
with (
|
|
patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
|
|
pytest.raises(TimeoutError),
|
|
):
|
|
mgr.call_tool_sync("mcp__test__ping", {}, timeout=1)
|
|
assert mgr._consecutive_failures.get("test", 0) == 1
|
|
|
|
def test_call_tool_sync_records_success(self):
|
|
mgr = MCPClientManager({"test": {"type": "stdio", "command": "echo"}})
|
|
mock_session = MagicMock()
|
|
mock_session.call_tool = MagicMock(return_value="sentinel")
|
|
mgr._sessions["test"] = mock_session
|
|
mgr._loop = MagicMock()
|
|
mgr._tool_map["mcp__test__ping"] = ("test", "ping")
|
|
# Pre-set a failure
|
|
mgr._consecutive_failures["test"] = 2
|
|
mock_result = MagicMock()
|
|
mock_result.content = []
|
|
mock_result.isError = False
|
|
mock_future = MagicMock()
|
|
mock_future.result.return_value = mock_result
|
|
with patch("asyncio.run_coroutine_threadsafe", return_value=mock_future):
|
|
mgr.call_tool_sync("mcp__test__ping", {}, timeout=5)
|
|
assert mgr._consecutive_failures.get("test") is None
|
|
|
|
def test_connection_error_evicts_session(self):
|
|
mgr = MCPClientManager({"test": {"type": "stdio", "command": "echo"}})
|
|
mock_session = MagicMock()
|
|
mock_session.call_tool = MagicMock(return_value="sentinel")
|
|
mgr._sessions["test"] = mock_session
|
|
mgr._loop = MagicMock()
|
|
mgr._tool_map["mcp__test__ping"] = ("test", "ping")
|
|
mock_future = MagicMock()
|
|
mock_future.result.side_effect = BrokenPipeError("dead")
|
|
with (
|
|
patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
|
|
pytest.raises(BrokenPipeError),
|
|
):
|
|
mgr.call_tool_sync("mcp__test__ping", {}, timeout=5)
|
|
assert "test" not in mgr._sessions
|
|
|
|
def test_independent_circuits_per_server(self):
|
|
mgr = MCPClientManager({})
|
|
for _ in range(3):
|
|
mgr._cb_record_failure("a")
|
|
is_open_a, _ = mgr._cb_check("a")
|
|
is_open_b, _ = mgr._cb_check("b")
|
|
assert is_open_a
|
|
assert not is_open_b
|
|
|
|
def test_mcp_error_does_not_trip_circuit(self):
|
|
"""Protocol errors (McpError) should not count as transport failures."""
|
|
from mcp import McpError
|
|
from mcp.types import ErrorData
|
|
|
|
mgr = MCPClientManager({"test": {"type": "stdio", "command": "echo"}})
|
|
mock_session = MagicMock()
|
|
mock_session.call_tool = MagicMock(return_value="sentinel")
|
|
mgr._sessions["test"] = mock_session
|
|
mgr._loop = MagicMock()
|
|
mgr._tool_map["mcp__test__ping"] = ("test", "ping")
|
|
mock_future = MagicMock()
|
|
mock_future.result.side_effect = McpError(ErrorData(code=-32601, message="tool not found"))
|
|
with (
|
|
patch("asyncio.run_coroutine_threadsafe", return_value=mock_future),
|
|
pytest.raises(McpError),
|
|
):
|
|
mgr.call_tool_sync("mcp__test__ping", {}, timeout=5)
|
|
# Circuit should NOT have recorded a failure
|
|
assert mgr._consecutive_failures.get("test", 0) == 0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fix 3: Safe transport stream pre-close
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestSafeTransportStreams:
|
|
"""Verify stream references are stored and pre-closed."""
|
|
|
|
def test_pre_close_streams_closes_both(self):
|
|
mgr = MCPClientManager({})
|
|
stream_a = MagicMock()
|
|
stream_b = MagicMock()
|
|
mgr._server_streams["srv"] = (stream_a, stream_b)
|
|
|
|
async def _run():
|
|
await mgr._pre_close_streams("srv")
|
|
|
|
asyncio.run(_run())
|
|
stream_a.aclose.assert_called_once()
|
|
stream_b.aclose.assert_called_once()
|
|
assert "srv" not in mgr._server_streams
|
|
|
|
def test_pre_close_streams_ignores_missing(self):
|
|
mgr = MCPClientManager({})
|
|
|
|
async def _run():
|
|
await mgr._pre_close_streams("nonexistent")
|
|
|
|
asyncio.run(_run()) # should not raise
|
|
|
|
def test_pre_close_streams_suppresses_errors(self):
|
|
mgr = MCPClientManager({})
|
|
stream_a = MagicMock()
|
|
stream_a.aclose.side_effect = RuntimeError("boom")
|
|
stream_b = MagicMock()
|
|
mgr._server_streams["srv"] = (stream_a, stream_b)
|
|
|
|
async def _run():
|
|
await mgr._pre_close_streams("srv")
|
|
|
|
asyncio.run(_run()) # should not raise despite stream_a error
|
|
stream_b.aclose.assert_called_once()
|
|
|
|
def test_shutdown_clears_stream_refs(self):
|
|
mgr = MCPClientManager({})
|
|
mgr._server_streams["srv"] = (MagicMock(), MagicMock())
|
|
mgr.shutdown()
|
|
assert len(mgr._server_streams) == 0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fix 4: Notification debounce
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestNotificationDebounce:
|
|
"""Verify notification-triggered refreshes are debounced."""
|
|
|
|
def test_debounce_within_window(self):
|
|
mgr = MCPClientManager({})
|
|
mgr._last_notification_refresh["srv"] = time.monotonic()
|
|
# We can't easily call _on_notification (it's a closure), so test
|
|
# the debounce logic directly via the timestamp check
|
|
now = time.monotonic()
|
|
last = mgr._last_notification_refresh.get("srv", 0.0)
|
|
assert now - last < mgr._NOTIFICATION_DEBOUNCE
|
|
|
|
def test_debounce_passes_after_window(self):
|
|
mgr = MCPClientManager({})
|
|
# Set timestamp well in the past
|
|
mgr._last_notification_refresh["srv"] = time.monotonic() - 10
|
|
now = time.monotonic()
|
|
last = mgr._last_notification_refresh.get("srv", 0.0)
|
|
assert now - last >= mgr._NOTIFICATION_DEBOUNCE
|
|
|
|
def test_debounce_is_per_server(self):
|
|
mgr = MCPClientManager({})
|
|
mgr._last_notification_refresh["srv_a"] = time.monotonic()
|
|
# srv_b has no timestamp — should pass debounce
|
|
now = time.monotonic()
|
|
last_b = mgr._last_notification_refresh.get("srv_b", 0.0)
|
|
assert now - last_b >= mgr._NOTIFICATION_DEBOUNCE
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fix 5: Periodic refresh backoff
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestPeriodicRefreshBackoff:
|
|
"""Verify periodic refresh backoff and auto-reconnect."""
|
|
|
|
def test_backoff_set_on_failure(self):
|
|
mgr = MCPClientManager({})
|
|
mgr._refresh_failures["srv"] = 1
|
|
# Simulate what _periodic_refresh does on failure
|
|
failures = mgr._refresh_failures.get("srv", 0) + 1
|
|
mgr._refresh_failures["srv"] = failures
|
|
backoff = min(mgr._REFRESH_BACKOFF_BASE * (2 ** (failures - 1)), mgr._REFRESH_BACKOFF_MAX)
|
|
mgr._refresh_backoff_until["srv"] = time.monotonic() + backoff
|
|
assert mgr._refresh_backoff_until["srv"] > time.monotonic()
|
|
assert failures == 2
|
|
|
|
def test_backoff_doubles(self):
|
|
mgr = MCPClientManager({})
|
|
b1 = min(mgr._REFRESH_BACKOFF_BASE * (2**0), mgr._REFRESH_BACKOFF_MAX)
|
|
b2 = min(mgr._REFRESH_BACKOFF_BASE * (2**1), mgr._REFRESH_BACKOFF_MAX)
|
|
b3 = min(mgr._REFRESH_BACKOFF_BASE * (2**2), mgr._REFRESH_BACKOFF_MAX)
|
|
assert b1 == 60
|
|
assert b2 == 120
|
|
assert b3 == 240
|
|
|
|
def test_backoff_capped(self):
|
|
mgr = MCPClientManager({})
|
|
b = min(mgr._REFRESH_BACKOFF_BASE * (2**20), mgr._REFRESH_BACKOFF_MAX)
|
|
assert b == mgr._REFRESH_BACKOFF_MAX
|
|
|
|
def test_backoff_clears_on_success(self):
|
|
mgr = MCPClientManager({})
|
|
mgr._refresh_failures["srv"] = 3
|
|
mgr._refresh_backoff_until["srv"] = time.monotonic() + 1000
|
|
# Simulate success
|
|
mgr._refresh_failures.pop("srv", None)
|
|
mgr._refresh_backoff_until.pop("srv", None)
|
|
assert "srv" not in mgr._refresh_failures
|
|
assert "srv" not in mgr._refresh_backoff_until
|
|
|
|
def test_server_status_includes_circuit_info(self):
|
|
mgr = MCPClientManager({"srv": {"type": "stdio", "command": "echo"}})
|
|
status = mgr.get_server_status("srv")
|
|
assert "circuit_open" in status
|
|
assert "consecutive_failures" in status
|
|
assert status["circuit_open"] is False
|
|
assert status["consecutive_failures"] == 0
|
|
|
|
def test_server_status_shows_open_circuit(self):
|
|
mgr = MCPClientManager({"srv": {"type": "stdio", "command": "echo"}})
|
|
for _ in range(3):
|
|
mgr._cb_record_failure("srv")
|
|
status = mgr.get_server_status("srv")
|
|
assert status["circuit_open"] is True
|
|
assert status["consecutive_failures"] == 3
|