ewilliams28/turnstone
1
0
Fork 0
mirror of https://github.com/turnstonelabs/turnstone.git synced 2026-08-12 23:12:23 -06:00
Code Issues Packages Projects Releases Wiki Activity
Files
3f5ee333fbcc4c024f3dcdab92139b91264059da
turnstone/tests/test_mcp_client.py
T
Patrick Buckley 3f5ee333fb fix(mcp): close the shutdown drain race + close the owned loop 
Review feedback: (1) gating the drain on a main-thread truthiness check
of _background_tasks could skip cancellation when a spawn queued via
call_soon_threadsafe had not reached the set yet — submit whenever the
loop is RUNNING and snapshot on the loop, where FIFO callback order
guarantees earlier-queued spawns have landed; (2) shutdown stopped the
loop thread but never closed the loop or cleared _loop/_thread, leaking
selector resources for embedders that cycle managers — close + clear
when we own the thread and it actually stopped (loud warning when it
does not); unowned loops (tests wiring _loop directly) stay untouched;
(3) the bare await-in-suppress drain loops become
asyncio.gather(return_exceptions=True) in both the shutdown drain and
the test fixture.
2026-06-11 18:15:58 -07:00

2818 lines
110 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""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, suppress
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from tests.conftest import _seed_static_state
from turnstone.core.mcp_client import (
MCPClientManager,
_db_servers_to_config,
_mcp_to_openai,
load_mcp_config,
)
from turnstone.core.tools import INTERACTIVE_TOOLS, 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}
],
}
@pytest.fixture
def running_loop_mgr():
"""Yield a (mgr, loop, thread) triple with a background loop already running.
Spawns an MCPClientManager with a default config of {"srv": stdio echo}
and a fresh asyncio loop driven by a daemon thread. Tests that need a
different config can mutate ``mgr._server_configs`` directly. The
fixture stops the loop and joins the thread on teardown so each test
leaves a clean slate.
"""
import threading as _threading
cfg = {"srv": {"type": "stdio", "command": "echo"}}
mgr = MCPClientManager(cfg)
loop = asyncio.new_event_loop()
thread = _threading.Thread(target=loop.run_forever, daemon=True)
thread.start()
mgr._loop = loop
try:
yield mgr, loop, thread
finally:
# Drain BEFORE stopping: a task left pending (or finished-but-
# unretrieved) on a stopped loop becomes cross-test global state —
# asyncio reports it at GC time, mid-suite, onto whatever stream
# pytest has attached THEN (the "I/O operation on closed file"
# spew), and a silently-abandoned loop thread keeps running
# manager code against torn-down mocks.
async def _cancel_pending() -> None:
tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
for t in tasks:
t.cancel()
await asyncio.gather(*tasks, return_exceptions=True)
with suppress(Exception):
asyncio.run_coroutine_threadsafe(_cancel_pending(), loop).result(timeout=5)
loop.call_soon_threadsafe(loop.stop)
thread.join(timeout=5)
assert not thread.is_alive(), "mcp test loop thread failed to stop within 5s"
loop.close()
# ---------------------------------------------------------------------------
# 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 == {}
class TestDBServersToConfig:
"""``_db_servers_to_config`` shapes DB rows for the MCP client."""
def test_static_streamable_http_row_passes_through(self) -> None:
rows = [
{
"name": "static-srv",
"transport": "streamable-http",
"url": "https://mcp.example.com",
"headers": '{"Authorization": "Bearer token"}',
"auth_type": "static",
}
]
result = _db_servers_to_config(rows)
assert "static-srv" in result
assert result["static-srv"]["url"] == "https://mcp.example.com"
assert result["static-srv"]["headers"] == {"Authorization": "Bearer token"}
def test_db_servers_to_config_skips_oauth_user_rows(self) -> None:
"""Rows with auth_type=oauth_user must be invisible to the static
auto-connect path.
Auto-connecting these with empty headers fails the AS check and
trips the circuit breaker on startup. Per-user OAuth servers
come online lazily once the user has consented.
"""
rows = [
{
"name": "static-srv",
"transport": "streamable-http",
"url": "https://static.example.com",
"headers": "{}",
"auth_type": "static",
},
{
"name": "oauth-srv",
"transport": "streamable-http",
"url": "https://oauth.example.com",
"headers": "{}",
"auth_type": "oauth_user",
},
{
"name": "stdio-srv",
"transport": "stdio",
"command": "echo",
"args": "[]",
"env": "{}",
"auth_type": "none",
},
]
result = _db_servers_to_config(rows)
assert set(result) == {"static-srv", "stdio-srv"}
assert "oauth-srv" not in 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({})
_seed_static_state(mgr, "a", session=MagicMock())
_seed_static_state(mgr, "b", session=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
# -- Phase 7: per-user catalog scoping ---------------------------------
def test_is_mcp_tool_user_id_default_none_unchanged(self):
"""Sanity: default ``user_id=None`` answers static-only.
The legacy single-arg call still works, and unknown names still
return False — Phase 7 adds an optional keyword without
rewriting the static-path semantics.
"""
mgr = MCPClientManager({})
mgr._tool_map["mcp__static__list"] = ("static", "list")
# Legacy single-arg call still works.
assert mgr.is_mcp_tool("mcp__static__list") is True
assert mgr.is_mcp_tool("mcp__static__list", user_id=None) is True
assert mgr.is_mcp_tool("nonexistent") is False
assert mgr.is_mcp_tool("nonexistent", user_id=None) is False
def test_is_mcp_tool_user_keyed_pool_tool(self):
"""A name visible only via ``_user_tool_map`` resolves only for
the matching ``user_id``.
Verifies the new branch: ``_tool_map`` miss + ``user_id`` hit.
"""
mgr = MCPClientManager({})
mgr._user_tool_map["user-1"] = {
"mcp__pool-srv__do": ("pool-srv", "do"),
}
# Visible to user-1.
assert mgr.is_mcp_tool("mcp__pool-srv__do", user_id="user-1") is True
# Invisible to None caller (admin / web-search backend resolution).
assert mgr.is_mcp_tool("mcp__pool-srv__do", user_id=None) is False
# Invisible to a different user.
assert mgr.is_mcp_tool("mcp__pool-srv__do", user_id="user-2") is False
def test_is_mcp_tool_static_wins_for_any_user(self):
"""Static-path tools are visible regardless of ``user_id`` —
the merged view is ``static ∪ user-pool``."""
mgr = MCPClientManager({})
mgr._tool_map["mcp__static__list"] = ("static", "list")
# Even for an unknown user, a static tool is still reachable —
# static-path is process-global.
assert mgr.is_mcp_tool("mcp__static__list", user_id="user-1") is True
assert mgr.is_mcp_tool("mcp__static__list", user_id="anybody") is True
def test_get_tools_user_id_none_returns_static_only(self):
"""``get_tools(user_id=None)`` returns the global static catalog.
Pool tools are NEVER included in the default-arg view — that's
the legacy contract every pre-Phase-7 caller relies on.
"""
mgr = MCPClientManager({})
mgr._tools = [_fake_openai_tool("mcp__static__list")]
# Seed a pool entry that should NOT appear in the default view.
from turnstone.core.mcp_client import PoolEntryState
entry = PoolEntryState(key=("user-1", "pool-srv"), open_lock=MagicMock())
entry.tools = [_fake_openai_tool("mcp__pool-srv__do")]
mgr._user_pool_entries[("user-1", "pool-srv")] = entry
tools = mgr.get_tools()
names = [t["function"]["name"] for t in tools]
assert names == ["mcp__static__list"]
def test_get_tools_user_id_merges_pool(self):
"""``get_tools(user_id='user-1')`` merges static + that user's pool tools.
Other users' pool entries MUST NOT leak into the result —
privacy / RBAC invariant.
"""
from turnstone.core.mcp_client import PoolEntryState
mgr = MCPClientManager({})
mgr._tools = [_fake_openai_tool("mcp__static__list")]
e1 = PoolEntryState(key=("user-1", "pool-srv"), open_lock=MagicMock())
e1.tools = [_fake_openai_tool("mcp__pool-srv__do")]
e2 = PoolEntryState(key=("user-2", "pool-srv"), open_lock=MagicMock())
e2.tools = [_fake_openai_tool("mcp__pool-srv__other")]
mgr._user_pool_entries[("user-1", "pool-srv")] = e1
mgr._user_pool_entries[("user-2", "pool-srv")] = e2
# Production invariant: ``_connect_one_pool`` /
# ``_refresh_pool_server_tools`` / ``_evict_session`` /
# ``_close_pool_entry_if_idle`` all call ``_rebuild_user_tool_map``
# immediately after mutating ``_user_pool_entries``. Tests that
# seed pool entries directly must mirror that invariant —
# ``get_tools(user_id=...)`` reads from the ``_user_tools``
# snapshot (built by ``_rebuild_user_tool_map``), never iterating
# ``_user_pool_entries`` directly.
mgr._rebuild_user_tool_map("user-1")
mgr._rebuild_user_tool_map("user-2")
u1_names = [t["function"]["name"] for t in mgr.get_tools(user_id="user-1")]
assert sorted(u1_names) == ["mcp__pool-srv__do", "mcp__static__list"]
u2_names = [t["function"]["name"] for t in mgr.get_tools(user_id="user-2")]
assert sorted(u2_names) == ["mcp__pool-srv__other", "mcp__static__list"]
# Default still global-only — unaffected by either user's entries.
default_names = [t["function"]["name"] for t in mgr.get_tools()]
assert default_names == ["mcp__static__list"]
def test_get_tools_user_id_returns_copies(self):
"""Mirror existing ``test_get_tools_returns_copy``: caller mutation
of the returned list MUST NOT affect the manager's catalog."""
from turnstone.core.mcp_client import PoolEntryState
mgr = MCPClientManager({})
mgr._tools = [_fake_openai_tool("mcp__static__a")]
entry = PoolEntryState(key=("user-1", "pool-srv"), open_lock=MagicMock())
entry.tools = [_fake_openai_tool("mcp__pool-srv__b")]
mgr._user_pool_entries[("user-1", "pool-srv")] = entry
# See note in ``test_get_tools_user_id_merges_pool``.
mgr._rebuild_user_tool_map("user-1")
tools = mgr.get_tools(user_id="user-1")
assert len(tools) == 2
tools.clear()
# Re-fetch — original catalog unchanged.
assert len(mgr.get_tools(user_id="user-1")) == 2
def test_get_tools_user_with_none_tools_skipped(self):
"""A pool entry that hasn't completed discovery (``entry.tools is None``)
contributes no tools — the merged view skips it cleanly."""
from turnstone.core.mcp_client import PoolEntryState
mgr = MCPClientManager({})
mgr._tools = [_fake_openai_tool("mcp__static__a")]
# Brand-new pool entry, discovery not yet run.
entry = PoolEntryState(key=("user-1", "pool-srv"), open_lock=MagicMock())
assert entry.tools is None
mgr._user_pool_entries[("user-1", "pool-srv")] = entry
# Rebuild observes ``entry.tools is None`` and skips this entry.
mgr._rebuild_user_tool_map("user-1")
names = [t["function"]["name"] for t in mgr.get_tools(user_id="user-1")]
assert names == ["mcp__static__a"]
def test_rebuild_user_tool_map_populates(self):
"""``_rebuild_user_tool_map`` materializes the per-user index from
pool entries owned by that user."""
from turnstone.core.mcp_client import PoolEntryState
mgr = MCPClientManager({})
entry = PoolEntryState(key=("user-1", "pool-srv"), open_lock=MagicMock())
entry.tools = [_fake_openai_tool("mcp__pool-srv__do")]
mgr._user_pool_entries[("user-1", "pool-srv")] = entry
mgr._rebuild_user_tool_map("user-1")
assert mgr._user_tool_map["user-1"] == {"mcp__pool-srv__do": ("pool-srv", "do")}
# Sibling _user_tools cache (bug-1 fix) MUST be populated alongside
# the map — otherwise get_tools(user_id="user-1") would silently
# return the static-only view despite is_mcp_tool returning True.
assert mgr._user_tools["user-1"] == [_fake_openai_tool("mcp__pool-srv__do")]
def test_rebuild_user_tool_map_drops_empty_user(self):
"""Rebuilding for a user with no pool entries removes the key
rather than retaining an empty-dict sentinel."""
from turnstone.core.mcp_client import PoolEntryState
mgr = MCPClientManager({})
entry = PoolEntryState(key=("user-1", "pool-srv"), open_lock=MagicMock())
entry.tools = [_fake_openai_tool("mcp__pool-srv__do")]
mgr._user_pool_entries[("user-1", "pool-srv")] = entry
mgr._rebuild_user_tool_map("user-1")
assert "user-1" in mgr._user_tool_map
assert "user-1" in mgr._user_tools
# Drop the entry, rebuild — user_id key should be removed from BOTH
# the map and the sibling tool list (bug-1 fix). A drop in only one
# would leave get_tools and is_mcp_tool out of sync.
mgr._user_pool_entries.clear()
mgr._rebuild_user_tool_map("user-1")
assert "user-1" not in mgr._user_tool_map
assert "user-1" not in mgr._user_tools
def test_rebuild_user_tool_map_isolates_users(self):
"""Rebuilding for ``user-1`` MUST NOT touch ``user-2``'s entry."""
from turnstone.core.mcp_client import PoolEntryState
mgr = MCPClientManager({})
e1 = PoolEntryState(key=("user-1", "pool-srv"), open_lock=MagicMock())
e1.tools = [_fake_openai_tool("mcp__pool-srv__one")]
e2 = PoolEntryState(key=("user-2", "pool-srv"), open_lock=MagicMock())
e2.tools = [_fake_openai_tool("mcp__pool-srv__two")]
mgr._user_pool_entries[("user-1", "pool-srv")] = e1
mgr._user_pool_entries[("user-2", "pool-srv")] = e2
mgr._rebuild_user_tool_map("user-1")
mgr._rebuild_user_tool_map("user-2")
assert mgr._user_tool_map["user-1"] == {"mcp__pool-srv__one": ("pool-srv", "one")}
assert mgr._user_tool_map["user-2"] == {"mcp__pool-srv__two": ("pool-srv", "two")}
# Clear user-1's entry only; rebuild user-1; user-2 must remain.
mgr._user_pool_entries.pop(("user-1", "pool-srv"))
mgr._rebuild_user_tool_map("user-1")
assert "user-1" not in mgr._user_tool_map
assert mgr._user_tool_map["user-2"] == {"mcp__pool-srv__two": ("pool-srv", "two")}
def test_rebuild_user_tool_map_does_not_touch_static(self):
"""Invariant 1: per-user rebuild must NOT mutate ``_tool_map``."""
from turnstone.core.mcp_client import PoolEntryState
mgr = MCPClientManager({})
mgr._tool_map["mcp__static__list"] = ("static", "list")
entry = PoolEntryState(key=("user-1", "pool-srv"), open_lock=MagicMock())
entry.tools = [_fake_openai_tool("mcp__pool-srv__do")]
mgr._user_pool_entries[("user-1", "pool-srv")] = entry
before = dict(mgr._tool_map)
mgr._rebuild_user_tool_map("user-1")
assert mgr._tool_map == before
# ---------------------------------------------------------------------------
# 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)
# Interactive session surface — coordinator tools excluded.
assert session._tools is INTERACTIVE_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(INTERACTIVE_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_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"},
user_id=None,
timeout=30,
is_interactive_for_consent=True,
)
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
# -- Phase 7: per-user catalog scoping ---------------------------------
def test_session_passes_user_id_to_get_tools(self, tmp_db):
"""ChatSession threads its ``user_id`` into ``get_tools`` so the
merged static + pool view is scoped to the session's user."""
mock_mcp = MagicMock()
mock_mcp.get_tools.return_value = [_fake_openai_tool()]
self._make_session(mcp_client=mock_mcp, user_id="user-7")
mock_mcp.get_tools.assert_called_with(user_id="user-7")
def test_session_get_tools_empty_user_id_passes_none(self, tmp_db):
"""Sentinel ``user_id=""`` (CLI / service / unknown) collapses to
``user_id=None`` so the static-only view is returned."""
mock_mcp = MagicMock()
mock_mcp.get_tools.return_value = [_fake_openai_tool()]
self._make_session(mcp_client=mock_mcp, user_id="")
mock_mcp.get_tools.assert_called_with(user_id=None)
def test_session_passes_user_id_to_add_listener(self, tmp_db):
"""ChatSession registers its tool-change listener under its own
``user_id`` so pool-only changes for OTHER users do not fire it."""
mock_mcp = MagicMock()
mock_mcp.get_tools.return_value = [_fake_openai_tool()]
self._make_session(mcp_client=mock_mcp, user_id="user-7")
# ``add_listener`` was called with ``user_id="user-7"``.
listener_calls = mock_mcp.add_listener.call_args_list
assert listener_calls, "ChatSession did not register a tool listener"
first_call = listener_calls[0]
assert first_call.kwargs.get("user_id") == "user-7"
def test_session_close_removes_listener_with_same_user_id(self, tmp_db):
"""R4 critical: register and remove MUST agree on ``user_id`` —
the listener identity is ``(user_id, callback)``."""
mock_mcp = MagicMock()
mock_mcp.get_tools.return_value = [_fake_openai_tool()]
session = self._make_session(mcp_client=mock_mcp, user_id="user-7")
session.close()
# ``remove_listener`` must be called with the same ``user_id``.
remove_calls = mock_mcp.remove_listener.call_args_list
assert remove_calls, "ChatSession.close did not unregister a tool listener"
first_remove = remove_calls[0]
assert first_remove.kwargs.get("user_id") == "user-7"
# And the callback identity must match what was registered.
registered_cb = mock_mcp.add_listener.call_args_list[0].args[0]
removed_cb = first_remove.args[0]
assert registered_cb is removed_cb
def test_session_unknown_tool_lists_user_scoped_catalog(self, tmp_db):
"""The "Unknown tool" error message lists tools the session can
actually invoke — drawn from the merged user-scoped catalog,
not the manager's private static-only ``_tool_map``."""
mock_mcp = MagicMock()
# Pretend the user's merged view contains a static + pool entry.
mock_mcp.get_tools.return_value = [
_fake_openai_tool("mcp__static__list"),
_fake_openai_tool("mcp__pool-srv__do"),
]
mock_mcp.is_mcp_tool.return_value = False
session = self._make_session(mcp_client=mock_mcp, user_id="user-7")
# Reset the call counter so we observe only the _prepare_tool call.
mock_mcp.get_tools.reset_mock()
tc = {
"id": "call_unknown",
"function": {"name": "no_such_tool", "arguments": "{}"},
}
prepared = session._prepare_tool(tc)
assert "error" in prepared
# The error mentions both static and pool tools — proves we're
# consulting the merged catalog rather than ``_tool_map``.
assert "mcp__static__list" in prepared["error"]
assert "mcp__pool-srv__do" in prepared["error"]
# And the catalog request was scoped to this session's user.
assert any(
call.kwargs.get("user_id") == "user-7" for call in mock_mcp.get_tools.call_args_list
)
# ---------------------------------------------------------------------------
# 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._static_servers
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({})
_seed_static_state(mgr, "github", tools=[_fake_openai_tool("mcp__github__search")])
_seed_static_state(mgr, "slack", tools=[_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({})
_seed_static_state(mgr, "a", tools=[_fake_openai_tool("mcp__a__x")])
mgr._rebuild_tools()
old_tools = mgr._tools
old_map = mgr._tool_map
_seed_static_state(mgr, "b", tools=[_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._static_servers = {}
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."""
_seed_static_state(mgr, server_name, supports_resources=True, supports_prompts=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)
_seed_static_state(
mgr,
"github",
session=mock_session,
tools=[_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)
_seed_static_state(
mgr,
"github",
session=mock_session,
tools=[_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)
_seed_static_state(
mgr,
"github",
session=mock_session,
tools=[_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 TestLastRefreshTracking:
"""Phase 9 admin status pill — ``_last_refresh`` is written on every
refresh path so the admin UI reflects manual-refresh AND auto-
reconnect outcomes uniformly. This test class pins the contract.
"""
@staticmethod
def _seed_minimal(mgr: MCPClientManager, name: str = "srv") -> MagicMock:
mock_session = MagicMock()
mock_session.list_tools = AsyncMock(return_value=MagicMock(tools=[]))
mock_session.list_resources = AsyncMock(return_value=MagicMock(resources=[]))
mock_session.list_resource_templates = AsyncMock(
return_value=MagicMock(resourceTemplates=[])
)
mock_session.list_prompts = AsyncMock(return_value=MagicMock(prompts=[]))
_seed_static_state(
mgr,
name,
session=mock_session,
tools=[],
supports_resources=True,
supports_prompts=True,
)
return mock_session
def test_last_refresh_written_on_success(self) -> None:
async def _run() -> None:
mgr = MCPClientManager({})
self._seed_minimal(mgr)
assert "srv" not in mgr._last_refresh
await mgr._refresh_server("srv")
entry = mgr._last_refresh.get("srv")
assert entry is not None
ts, outcome = entry
assert outcome == "ok"
assert isinstance(ts, float) and ts > 0
asyncio.run(_run())
def test_last_refresh_written_on_tool_refresh_failure(self) -> None:
"""When ``_refresh_server_tools`` raises, the outcome reflects
the exception class and the exception still propagates."""
async def _run() -> None:
mgr = MCPClientManager({})
mock_session = self._seed_minimal(mgr)
mock_session.list_tools = AsyncMock(side_effect=RuntimeError("upstream down"))
with pytest.raises(RuntimeError, match="upstream down"):
await mgr._refresh_server("srv")
entry = mgr._last_refresh.get("srv")
assert entry is not None
_, outcome = entry
assert outcome == "error:RuntimeError"
asyncio.run(_run())
def test_last_refresh_records_first_exception_when_multiple_fail(
self,
) -> None:
"""``return_exceptions=True`` lets sibling tasks complete; the
outcome reflects the FIRST exception encountered."""
async def _run() -> None:
mgr = MCPClientManager({})
mock_session = self._seed_minimal(mgr)
# Tools succeeds; resources raises first (gather preserves
# argument order in its results list, so resources is the
# first failure regardless of which awaitable finished first
# in wall-clock terms).
mock_session.list_resources = AsyncMock(side_effect=ValueError("res boom"))
mock_session.list_prompts = AsyncMock(side_effect=KeyError("prompts boom"))
with pytest.raises((ValueError, KeyError)):
await mgr._refresh_server("srv")
entry = mgr._last_refresh.get("srv")
assert entry is not None
_, outcome = entry
# Either of the two failing tasks could be "first" in
# gather's results list ordering — the order is positional
# so resources (arg #2) comes before prompts (arg #3).
assert outcome == "error:ValueError"
asyncio.run(_run())
def test_refresh_all_overwrites_stale_ok_on_reconnect_failure(
self,
) -> None:
"""The chokepoint bug-1 fix: a prior successful refresh's ``'ok'``
entry MUST be overwritten when a subsequent reconnect fails —
otherwise the admin pill shows misleading "ok" while the server
is in fact broken."""
async def _run() -> None:
mgr = MCPClientManager({})
# Server is configured but has no live session — _refresh_all
# routes to the reconnect branch.
mgr._server_configs["srv"] = {"type": "stdio", "command": "x"}
# Pre-seed a stale "ok" from an earlier successful refresh.
mgr._last_refresh["srv"] = (1000.0, "ok")
async def _raise(*_a: object, **_kw: object) -> None:
raise ConnectionError("reconnect failed")
mgr._connect_one = _raise # type: ignore[assignment]
await mgr._refresh_all("srv")
entry = mgr._last_refresh.get("srv")
assert entry is not None
ts, outcome = entry
# Outcome reflects the new failure, not the stale ok.
assert outcome == "error:ConnectionError"
assert ts > 1000.0
asyncio.run(_run())
def test_get_server_status_surfaces_last_refresh_fields(self) -> None:
"""``get_server_status`` surfaces ``last_refresh_at`` and
``last_refresh_outcome`` for the admin pill — null when no
refresh has occurred yet, populated after one."""
mgr = MCPClientManager({})
mgr._server_configs["srv"] = {"type": "stdio", "command": "x"}
# No refresh yet — fields must be present and null so the JS
# renderer can branch on absence cleanly.
status = mgr.get_server_status("srv")
assert status["last_refresh_at"] is None
assert status["last_refresh_outcome"] is None
# Populate the tuple directly and re-read.
mgr._last_refresh["srv"] = (12345.5, "ok")
status = mgr.get_server_status("srv")
assert status["last_refresh_at"] == 12345.5
assert status["last_refresh_outcome"] == "ok"
class TestListeners:
def test_add_and_notify(self):
mgr = MCPClientManager({})
calls: list[int] = []
mgr.add_listener(lambda: calls.append(1))
_seed_static_state(mgr, "a", tools=[_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
# -- Phase 7: user-keyed listener fan-out ------------------------------
def test_add_listener_records_user_id(self):
"""``add_listener`` stores ``(user_id, callback)`` tuples — the
listener identity carries the user_id."""
mgr = MCPClientManager({})
cb_admin = lambda: None # noqa: E731
cb_user = lambda: None # noqa: E731
mgr.add_listener(cb_admin) # default: user_id=None (admin)
mgr.add_listener(cb_user, user_id="user-1")
assert (None, cb_admin) in mgr._listeners
assert ("user-1", cb_user) in mgr._listeners
def test_remove_listener_requires_matching_user_id(self):
"""Removing with a different ``user_id`` must NOT remove the
original registration — listener identity is the pair."""
mgr = MCPClientManager({})
calls: list[int] = []
cb = lambda: calls.append(1) # noqa: E731
mgr.add_listener(cb, user_id="user-1")
# Try to remove with the wrong user_id — should be a no-op.
mgr.remove_listener(cb, user_id="user-2")
# The user-1 listener should still be live.
mgr._notify_user_tool_listeners("user-1")
assert calls == [1]
# Now remove with the right user_id.
mgr.remove_listener(cb, user_id="user-1")
mgr._notify_user_tool_listeners("user-1")
assert calls == [1] # not invoked again
def test_static_change_fires_all_listeners(self):
"""``_rebuild_tools`` (static-path change) fires ALL registered
listeners — admin + every user. RFC §3.3."""
mgr = MCPClientManager({})
admin_calls: list[int] = []
u1_calls: list[int] = []
u2_calls: list[int] = []
mgr.add_listener(lambda: admin_calls.append(1))
mgr.add_listener(lambda: u1_calls.append(1), user_id="user-1")
mgr.add_listener(lambda: u2_calls.append(1), user_id="user-2")
_seed_static_state(mgr, "a", tools=[_fake_openai_tool("mcp__a__x")])
mgr._rebuild_tools()
assert admin_calls == [1]
assert u1_calls == [1]
assert u2_calls == [1]
def test_user_tool_listeners_only_fire_for_matching_user(self):
"""``_notify_user_tool_listeners('user-1')`` fires admin (None)
and user-1 listeners; user-2's listener is silent."""
mgr = MCPClientManager({})
admin_calls: list[int] = []
u1_calls: list[int] = []
u2_calls: list[int] = []
mgr.add_listener(lambda: admin_calls.append(1))
mgr.add_listener(lambda: u1_calls.append(1), user_id="user-1")
mgr.add_listener(lambda: u2_calls.append(1), user_id="user-2")
mgr._notify_user_tool_listeners("user-1")
assert admin_calls == [1]
assert u1_calls == [1]
assert u2_calls == []
mgr._notify_user_tool_listeners("user-2")
assert admin_calls == [1, 1]
assert u1_calls == [1]
assert u2_calls == [1]
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({})
_seed_static_state(
mgr,
"fs",
resources=[
_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({})
_seed_static_state(mgr, "a", resources=[_fake_resource_dict("file:///x", "x", "", "", "a")])
mgr._rebuild_resources()
old_resources = mgr._resources
old_map = mgr._resource_map
_seed_static_state(mgr, "b", resources=[_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({})
_seed_static_state(mgr, "a", resources=[_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()
_seed_static_state(mgr, "fs", session=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()
_seed_static_state(mgr, "fs", session=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()
_seed_static_state(mgr, "fs", session=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))
_seed_static_state(mgr, "a", resources=[_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()
_seed_static_state(
mgr,
"fs",
session=mock_session,
supports_resources=True,
resources=[_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._static_servers = {}
mgr._rebuild_resources()
assert mgr._resources == []
assert mgr._resource_map == {}
def test_rebuild_resources_multi_server(self):
mgr = MCPClientManager({})
_seed_static_state(mgr, "fs", resources=[_fake_resource_dict("file:///a", server="fs")])
_seed_static_state(
mgr,
"db",
resources=[_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({})
_seed_static_state(
mgr,
"db",
resources=[
{
"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)
_seed_static_state(
mgr,
"short",
resources=[
{
"uri": "db://data/{collection}",
"name": "collection",
"description": "",
"mimeType": "",
"server": "short",
"template": True,
},
],
)
_seed_static_state(
mgr,
"long",
resources=[
{
"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({})
_seed_static_state(
mgr,
"db",
resources=[
{
"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({})
_seed_static_state(
mgr,
"db",
resources=[
{
"uri": "db://tables/{table}/rows/{id}",
"name": "row",
"description": "A row",
"mimeType": "application/json",
"server": "db",
"template": True,
},
],
)
mgr._rebuild_resources()
mock_session = MagicMock()
_seed_static_state(mgr, "db", session=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({})
_seed_static_state(
mgr,
"tmpl",
prompts=[
_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({})
_seed_static_state(mgr, "a", prompts=[_fake_prompt_dict("mcp__a__p1", "p1", "a")])
mgr._rebuild_prompts()
old_prompts = mgr._prompts
old_map = mgr._prompt_map
_seed_static_state(mgr, "b", prompts=[_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({})
_seed_static_state(mgr, "a", prompts=[_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()
_seed_static_state(mgr, "tmpl", session=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()
_seed_static_state(mgr, "tmpl", session=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))
_seed_static_state(mgr, "a", prompts=[_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()
_seed_static_state(
mgr,
"tmpl",
session=mock_session,
supports_prompts=True,
prompts=[_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._static_servers = {}
mgr._rebuild_prompts()
assert mgr._prompts == []
assert mgr._prompt_map == {}
def test_rebuild_prompts_multi_server(self):
mgr = MCPClientManager({})
_seed_static_state(mgr, "a", prompts=[_fake_prompt_dict("mcp__a__p1", "p1", "a")])
_seed_static_state(mgr, "b", prompts=[_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({})
_seed_static_state(mgr, "a", resources=[_fake_resource_dict()])
mgr._rebuild_resources()
_seed_static_state(mgr, "a", prompts=[_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 == {}
def test_shutdown_closes_owned_loop_and_clears_refs(self):
"""When the manager owns the loop thread, shutdown must close the loop
(selector resources leak otherwise) and drop both refs; a second
shutdown is then a clean no-op."""
import threading as _threading
mgr = MCPClientManager({})
loop = asyncio.new_event_loop()
thread = _threading.Thread(target=loop.run_forever, daemon=True)
thread.start()
mgr._loop = loop
mgr._thread = thread
mgr.shutdown()
assert loop.is_closed()
assert mgr._loop is None
assert mgr._thread is None
mgr.shutdown() # idempotent
def test_shutdown_leaves_unowned_loop_open(self):
"""Tests (and any embedder) that wire ``_loop`` directly without a
thread own the loop's lifecycle — shutdown must not close it."""
mgr = MCPClientManager({})
loop = asyncio.new_event_loop()
mgr._loop = loop
try:
mgr.shutdown()
assert not loop.is_closed()
finally:
loop.close()
# ---------------------------------------------------------------------------
# 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 have a live session (connection failed)
bad_state = mgr._static_servers.get("bad-server")
assert bad_state is None or bad_state.session is None
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()
bad_state = mgr._static_servers.get("bad")
assert bad_state is None or bad_state.session is None
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")
_seed_static_state(mgr, "test", session=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")
_seed_static_state(mgr, "test", session=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")
_seed_static_state(mgr, "test", session=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")
# Seed both session and stack so the test can verify stack survives.
old_stack = MagicMock()
old_streams = (MagicMock(), MagicMock())
_seed_static_state(mgr, "test", session=mock_session, stack=old_stack, streams=old_streams)
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)
# Session evicted, but stack/streams remain for the stale-and-stack
# guard in _connect_one to clean up on next reconnect attempt.
state = mgr._static_servers["test"]
assert state.session is None
assert state.stack is old_stack
assert state.streams is old_streams
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")
_seed_static_state(mgr, "test", session=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()
_seed_static_state(mgr, "srv", streams=(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()
# Streams cleared, but the state entry itself can remain.
assert mgr._static_servers["srv"].streams is None
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()
_seed_static_state(mgr, "srv", streams=(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({})
_seed_static_state(mgr, "srv", streams=(MagicMock(), MagicMock()))
mgr.shutdown()
assert len(mgr._static_servers) == 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
# ---------------------------------------------------------------------------
# reconnect_sync — operator-driven full reconnect
# ---------------------------------------------------------------------------
class TestReconnectSync:
"""Verify reconnect_sync tears down old session, clears CB, calls _connect_one."""
def test_reconnect_unknown_server_returns_error(self):
mgr = MCPClientManager({})
result = mgr.reconnect_sync("missing")
assert result == {
"connected": False,
"tools": 0,
"resources": 0,
"prompts": 0,
"error": "unknown server",
}
def test_reconnect_clears_circuit_breaker(self, running_loop_mgr):
mgr, _loop, _thread = running_loop_mgr
async def _fake_connect_one(name: str, _cfg: dict[str, Any]) -> None:
_seed_static_state(mgr, name, session=MagicMock())
# Pre-trip the breaker
for _ in range(3):
mgr._cb_record_failure("srv")
assert "srv" in mgr._circuit_open_until
with (
patch.object(mgr, "_connect_one", side_effect=_fake_connect_one),
patch.object(mgr, "_pre_close_streams", new=AsyncMock()),
):
result = mgr.reconnect_sync("srv")
assert result["connected"] is True
assert result["error"] == ""
assert "srv" not in mgr._circuit_open_until
assert "srv" not in mgr._consecutive_failures
def test_reconnect_closes_old_session_then_calls_connect_one(self, running_loop_mgr):
mgr, _loop, _thread = running_loop_mgr
order: list[str] = []
old_stack = MagicMock(spec=AsyncExitStack)
async def _pre_close(name: str) -> None:
order.append("pre_close")
async def _safe_close(stack: Any) -> None:
order.append("safe_close")
assert stack is old_stack
async def _connect_one(name: str, _cfg: dict[str, Any]) -> None:
order.append("connect_one")
_seed_static_state(mgr, name, session=MagicMock())
# Seed the old session/stack/streams that the guard should clear.
_seed_static_state(
mgr,
"srv",
session=MagicMock(),
stack=old_stack,
streams=(MagicMock(), MagicMock()),
)
with (
patch.object(mgr, "_pre_close_streams", side_effect=_pre_close),
patch.object(mgr, "_safe_close_stack", side_effect=_safe_close),
patch.object(mgr, "_connect_one", side_effect=_connect_one),
):
result = mgr.reconnect_sync("srv")
assert result["connected"] is True
assert order == ["pre_close", "safe_close", "connect_one"]
# The old stack reference should have been cleared from state.
assert mgr._static_servers["srv"].stack is not old_stack
def test_reconnect_failure_returns_error_dict(self, running_loop_mgr):
mgr, _loop, _thread = running_loop_mgr
async def _connect_one(name: str, _cfg: dict[str, Any]) -> None:
raise RuntimeError("handshake failed")
with (
patch.object(mgr, "_connect_one", side_effect=_connect_one),
patch.object(mgr, "_pre_close_streams", new=AsyncMock()),
):
result = mgr.reconnect_sync("srv")
assert result["connected"] is False
assert "handshake failed" in result["error"]
def test_reconnect_failure_clears_stale_catalog(self, running_loop_mgr):
# bug-2: when _connect_one fails mid-reconnect, the per-server
# catalog must be dropped so the merged tool/resource/prompt maps
# don't keep advertising entries with no live session.
mgr, _loop, _thread = running_loop_mgr
# Seed catalog state from a previous successful connect.
_seed_static_state(
mgr,
"srv",
tools=[_fake_openai_tool("mcp__srv__t")],
resources=[_fake_resource_dict(server="srv")],
prompts=[_fake_prompt_dict(server="srv")],
)
mgr._rebuild_tools()
mgr._rebuild_resources()
mgr._rebuild_prompts()
async def _connect_one(name: str, _cfg: dict[str, Any]) -> None:
raise RuntimeError("handshake failed")
with (
patch.object(mgr, "_connect_one", side_effect=_connect_one),
patch.object(mgr, "_pre_close_streams", new=AsyncMock()),
):
result = mgr.reconnect_sync("srv")
assert result["connected"] is False
# Per-server catalog should be cleared and merged maps drained.
srv_state = mgr._static_servers.get("srv")
assert srv_state is not None
assert srv_state.tools == []
assert srv_state.resources == []
assert srv_state.prompts == []
assert "mcp__srv__t" not in mgr._tool_map
def test_reconnect_preserves_static_state_identity(self, running_loop_mgr):
# q-3: PR #296 invariant 5 — _static_servers[name] must be the SAME
# object across a connect → transient-failure → reconnect cycle.
# Guards against future refactors that pop-and-repopulate the entry,
# which would invalidate any references held by concurrent readers.
mgr, _loop, _thread = running_loop_mgr
# First connect: seed an initial entry as if _connect_one succeeded.
async def _first_connect(name: str, _cfg: dict[str, Any]) -> None:
_seed_static_state(mgr, name, session=MagicMock())
with (
patch.object(mgr, "_connect_one", side_effect=_first_connect),
patch.object(mgr, "_pre_close_streams", new=AsyncMock()),
):
mgr.reconnect_sync("srv")
state_before = mgr._static_servers["srv"]
id_before = id(state_before)
# Simulate a transient transport failure: evict the session (as
# call_tool_sync would on BrokenPipeError) but keep the entry.
state_before.session = None
# Reconnect.
async def _reconnect(name: str, _cfg: dict[str, Any]) -> None:
_seed_static_state(mgr, name, session=MagicMock())
with (
patch.object(mgr, "_connect_one", side_effect=_reconnect),
patch.object(mgr, "_pre_close_streams", new=AsyncMock()),
):
result = mgr.reconnect_sync("srv")
assert result["connected"] is True
state_after = mgr._static_servers["srv"]
assert id(state_after) == id_before
assert state_after is state_before
# ---------------------------------------------------------------------------
# _cb_auto_reconnect — refresh-on-reconnect
# ---------------------------------------------------------------------------
class TestCBAutoReconnectRefresh:
"""Verify _cb_auto_reconnect schedules catalog refresh after a successful reconnect.
The refresh runs as a fire-and-forget background task on the loop so it
doesn't block the caller (perf-1). Tests wait briefly for the scheduled
task to run and observe its effect.
"""
def test_auto_reconnect_schedules_refresh_server_on_success(self, running_loop_mgr):
import threading as _threading
mgr, _loop, _thread = running_loop_mgr
new_session = MagicMock()
refresh_event = _threading.Event()
async def _connect_one(name: str, _cfg: dict[str, Any]) -> None:
_seed_static_state(mgr, name, session=new_session)
async def _refresh(name: str) -> tuple[list[str], list[str]]:
refresh_event.set()
return [], []
with (
patch.object(mgr, "_connect_one", side_effect=_connect_one),
patch.object(mgr, "_refresh_server", side_effect=_refresh),
):
session = mgr._cb_auto_reconnect("srv")
# Wait for the scheduled refresh task to actually run on the loop,
# then for the tracked task to DRAIN — exiting the patch context
# while the task is still in flight would hand the un-patched
# method to its tail.
assert refresh_event.wait(timeout=5), "refresh task was not scheduled"
deadline = time.time() + 5
while mgr._background_tasks and time.time() < deadline:
time.sleep(0.02)
assert not mgr._background_tasks, "background refresh task never drained"
assert session is new_session
def test_auto_reconnect_retrieves_and_logs_refresh_failure(self, running_loop_mgr):
"""A refresh failure must be RETRIEVED and logged by the task's
done-callback — not abandoned for asyncio to report as "Task exception
was never retrieved" at GC time (which lands on whatever stream pytest
has attached by then: the closed-file CI spew)."""
import threading as _threading
mgr, _loop, _thread = running_loop_mgr
new_session = MagicMock()
refresh_started = _threading.Event()
async def _connect_one(name: str, _cfg: dict[str, Any]) -> None:
_seed_static_state(mgr, name, session=new_session)
async def _refresh_failing(name: str) -> tuple[list[str], list[str]]:
refresh_started.set()
raise RuntimeError("catalog fetch broke")
with (
patch.object(mgr, "_connect_one", side_effect=_connect_one),
patch.object(mgr, "_refresh_server", side_effect=_refresh_failing),
patch("turnstone.core.mcp_client.log") as mock_log,
):
# Must not raise — refresh failures are non-fatal to the caller.
session = mgr._cb_auto_reconnect("srv")
assert refresh_started.wait(timeout=5), "refresh task was not scheduled"
# Poll for the WARNING while the patch is still active — gating on
# set-emptiness alone would race the un-patch (review-caught: the
# warning could land on the restored real logger).
deadline = time.time() + 5
warn_calls = []
while not warn_calls and time.time() < deadline:
warn_calls = [
c for c in mock_log.warning.call_args_list if "MCP background" in str(c.args[0])
]
time.sleep(0.02)
# The tracked task must also fully drain (emptiness now implies
# "done AND reported" — discard is the callback's LAST step).
deadline = time.time() + 5
while mgr._background_tasks and time.time() < deadline:
time.sleep(0.02)
assert not mgr._background_tasks, "background refresh task never drained"
assert session is new_session
assert warn_calls, (
"the refresh failure must be logged by the done-callback, not left "
"for GC-time reporting"
)
exc = warn_calls[0].kwargs.get("exc_info")
assert isinstance(exc, RuntimeError)
assert "catalog fetch broke" in str(exc)
Reference in New Issue View Git Blame Copy Permalink