mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-25 21:34:47 -06:00
Patrick Buckley
cfc8a6c8c0
Light up production reachability of pool dispatch (RFC §3, invariant 8) by widening the public catalog API to optionally take a ``user_id``: - ``MCPClientManager.get_tools(user_id=None)`` returns the merged static + per-user pool view when ``user_id`` is supplied; the default preserves the legacy global-only contract. - ``is_mcp_tool(name, *, user_id=None)`` extends the lookup to the per-user ``_user_tool_map``. Pool tools become reachable from ``ChatSession._prepare_tool`` only when the session-bound user_id flows through — flipping invariant 8 from "must hold" to "satisfied". - Listener identity becomes ``(user_id, callback)``. Static-path changes fire ALL listeners (admin + every user); pool-entry changes fire only matching-user + admin (``None``) listeners. RFC §3.3. - Pool sessions discover their tool list on first connect (``_connect_one_pool`` → ``await session.list_tools()``); the notification closure binds to ``(user_id, server_name)`` so push-driven ``list_changed`` updates target the correct user's catalog. R6 verified empirically: ``list_tools()`` 401 propagates through anyio TaskGroup unwinding, no hang — plain ``await`` is fine, no carrier-race shape needed for discovery. - ``_evict_session`` drops ``entry.tools`` and rebuilds the user's index so an evicted-then-reconnected session doesn't carry stale catalog state. - ``web_search.resolve_web_search_client`` refuses ``auth_type=oauth_user`` backends (per-node web search can't carry per-user tokens). Resources / prompts pool dispatch deferred to Phase 7b — invariant 8 is satisfied by the tool path alone, and the resource/prompt path needs sibling ``_dispatch_pool_resource_sync`` / ``_dispatch_pool_prompt_sync`` helpers each with their own carrier-race plumbing (~400 LOC). Phase 7b will follow the patterns established here. CLI sessions default ``user_id=""`` and so cannot use oauth_user MCP servers — documented limitation; users must use the web UI. Round-1 review fixes (4-finder review applied, no push yet): - bug-1: get_tools(user_id) was iterating _user_pool_entries from sync threads while the mcp-loop concurrently mutated it (RuntimeError: dictionary changed size during iteration). Now reads from a sibling _user_tools dict updated atomically by _rebuild_user_tool_map. - bug-2: _close_pool_entry_if_idle (LRU/TTL eviction) skipped the catalog cleanup that _evict_session does — stale tools persisted in _user_tool_map and ChatSession's tool list never rebuilt. Now mirrors _evict_session. - perf-1: _last_pool_notification_refresh debounce dict was never pruned in either eviction path. Now popped alongside the entry. - perf-3: web_search resolver was issuing a sync SQL query per LLM turn to gate oauth_user backends. Now reads from the cached in-memory config. - sec-1: bearer token could leak into exc_info-rendered tracebacks via Sentry/faulthandler. log.debug now uses structured fields, not exc_info. - sec-2: tools-per-server response now capped at 1000 (defensive, mirrors _MAX_ERROR_LEN / _MAX_INSUFFICIENT_SCOPE_REPORTED). - Test cleanup: dropped two listener fan-out tests duplicating test_mcp_client.py coverage; renamed test_pool_session_notification_handler to match its actual scope (_refresh_pool_server_tools); removed stale comments referencing /tmp/r6-spike*.py scratchpads and a misleading "copy-on-write" comment. Round-2 pre-push review fixes (focused single-pass review applied): - round2-1: bug-2's catalog-cleanup block in _close_pool_entry_if_idle had no integration test (exactly the failure mode flagged in feedback_tests_through_boundaries.md). Added test_close_pool_entry_if_idle_clears_catalog_and_fires_listener driving the LRU/TTL eviction path through real streamablehttp_client + MockTransport. Negative-test verified: reverting the _rebuild_user_tool_map / _notify_user_tool_listeners calls makes the new test fail. - round2-3: documented the _oauth_user_server_names cache invariant in add_server_sync / remove_server_sync docstrings. Cache is reconcile_sync's sole owner — direct callers leave it stale, but _db_servers_to_config strips oauth_user rows so production paths are unaffected. Static→oauth_user transitions correctly leave the name in the cache because remove_server_sync drops the static connection, not the cache identity. - round2-6: strengthened test_rebuild_user_tool_map_populates and test_rebuild_user_tool_map_drops_empty_user to assert on the _user_tools sibling cache (bug-1 fix). Without this, a future revert dropping the sibling write would still pass the unit tests because get_tools coverage lives in separate tests. Round-3 full-stack review fixes (multi-stage review on the final state caught what the layered apply passes missed): - q-1 REGRESSION: pool tool-discovery used asyncio.wait_for around session.list_tools(), the exact pattern thef6a3b66fix (and feedback_asyncio_timeout_vs_wait_for.md) put in place to avoid. Python 3.11's asyncio.wait_for wraps the inner coroutine in a fresh task → cross-task scope-exit when the SDK's anyio TaskGroup unwinds on a 401. Switched to `async with asyncio.timeout(...):` pattern used by _safe_close_stack. - sec-2: TOCTOU in _connect_one_pool — entry.tools was published (via _rebuild_user_tool_map + listener fan-out) BEFORE entry.session was assigned. A sync-thread reader could observe a tool whose backing entry has session=None. Defence-in-depth — dispatch re-fetches its own token and lazy-reconnects on session=None — but reordering catches the race at the source. entry.session now publishes BEFORE catalog visibility. - bug-1: _close_pool_entry_if_idle's _user_pool_locks.pop ran unconditionally after the try/finally, but the early-return branches (entry None on re-check, in_flight > 0 under lock) skip it via Python's return-through-finally semantics. The lock was never popped on those paths. Now gated behind an `evicted` flag set only on the success path; in_flight > 0 leaves the lock for the active dispatcher to reuse, entry-None races leave the lock for re-allocation by _ensure_pool_entry. Comment now describes the actual semantics, not the original promise. - bug-2: softened the _rebuild_user_tool_map docstring's atomicity claim. The two-dict write is technically non-atomic across Python statements; in practice the window is sub-microsecond on the mcp-loop with no awaits between writes, and the listener fan-out fires AFTER both writes complete. Docstring now says "back-to-back on the mcp-loop" instead of "atomically alongside". - q-3: dropped `hasattr(mcp_client, "server_auth_type")` defensive check in web_search.py. The method ships in this commit; the hasattr created a silent fallthrough that would let a future rename silently re-enable oauth_user backends. - q-4: surfaced the CLI / empty-user_id limitation in a docstring comment at ChatSession.__init__'s self._user_id assignment. The note previously lived only inside is_mcp_tool's docstring — a future maintainer wiring CLI features against MCP pool servers wouldn't think to read is_mcp_tool to find the constraint. - q-2 + q-5: deleted a tautological duplicate test in test_mcp_user_catalog.py whose docstring claimed to test ChatSession.close but never instantiated a ChatSession (the manager-level identity semantics are already covered by test_listener_identity_includes_user_id in the same file and by test_session_close_removes_listener_with_same_user_id in test_mcp_client.py which DOES drive a ChatSession). Reworded a misleading "fixture provides only 5s" comment to point at the actual `_run_on_loop(..., timeout=5)` site. - q-6: the `self._user_id or None` collapse repeated at 8 sites across session.py. Cached once at __init__ as ``self._mcp_user_id`` (since ``_user_id`` is set once and never mutated); 8 call sites now read the cached value. The empty- string-is-CLI-sentinel invariant is documented at the assignment site, not re-asserted at each consumer. Deferred to follow-up: - sec-1: a hostile MCP server bound to user-A could craft a tool.name containing `__` to synthesize a prefixed-name collision in user-A's own catalog. Bounded impact: cross-tenant dispatch is prevented by the per-tenant token gate in _dispatch_pool, and user-B's get_tools(user_id="B") never includes user-A's pool entries. The fix needs policy decisions (reject vs. sanitize) and touches _mcp_to_openai which is shared between static and pool paths; better discussed in its own follow-up where the policy applies uniformly to static-path servers too. The threat model already requires user-A to have consented to a malicious server, who has many more dangerous vectors than tool-name shenanigans. Test count delta: +31 tests (5435 → 5466, ``-m "not live"``; one test deleted in round-3 apply per q-2): - ``tests/test_mcp_client.py`` +20 (per-user catalog state, listener identity, session thread-through) - ``tests/test_mcp_user_catalog.py`` +9 NEW (integration tests driving real ``streamablehttp_client`` + ``httpx.MockTransport`` per invariant 14: discovery on connect, user isolation, eviction + reconnect, LRU/TTL eviction (round2-1), R6 401-propagation regression, static byte-identical canonical regression; review passes dropped duplicate listener fan-out tests from earlier drafts whose coverage lived in test_mcp_client.py) - ``tests/test_web_search.py`` +2 (oauth_user backend rejection + static backend acceptance regression; updated to use the new ``server_auth_type`` in-memory accessor) (cherry picked from commita8b34bfe54)
2601 lines
101 KiB
Python
2601 lines
101 KiB
Python
"""Tests for turnstone.core.mcp_client — MCP client manager and config loading."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import concurrent.futures
|
||
import json
|
||
import time
|
||
from contextlib import AsyncExitStack
|
||
from typing import Any
|
||
from unittest.mock import AsyncMock, MagicMock, patch
|
||
|
||
import pytest
|
||
|
||
from 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:
|
||
loop.call_soon_threadsafe(loop.stop)
|
||
thread.join(timeout=2)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 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_agent_tools_include_mcp(self, tmp_db):
|
||
from turnstone.core.tools import AGENT_TOOLS
|
||
|
||
mock_mcp = MagicMock()
|
||
mock_mcp.get_tools.return_value = [_fake_openai_tool()]
|
||
session = self._make_session(mcp_client=mock_mcp)
|
||
assert len(session._agent_tools) == len(AGENT_TOOLS) + 1
|
||
|
||
def test_prepare_mcp_tool(self, tmp_db):
|
||
mock_mcp = MagicMock()
|
||
mock_mcp.get_tools.return_value = [_fake_openai_tool()]
|
||
mock_mcp.is_mcp_tool.return_value = True
|
||
session = self._make_session(mcp_client=mock_mcp)
|
||
|
||
tc = {
|
||
"id": "call_123",
|
||
"function": {
|
||
"name": "mcp__test__search",
|
||
"arguments": '{"query": "hello"}',
|
||
},
|
||
}
|
||
prepared = session._prepare_tool(tc)
|
||
assert prepared["func_name"] == "mcp__test__search"
|
||
assert prepared["needs_approval"] is True
|
||
assert "mcp:test/search" in prepared["header"]
|
||
assert callable(prepared["execute"])
|
||
|
||
def test_unknown_tool_without_mcp(self, tmp_db):
|
||
session = self._make_session(mcp_client=None)
|
||
tc = {
|
||
"id": "call_456",
|
||
"function": {"name": "nonexistent", "arguments": "{}"},
|
||
}
|
||
prepared = session._prepare_tool(tc)
|
||
assert "error" in prepared
|
||
assert "Unknown tool" in prepared["error"]
|
||
# Error lists available tools so the model can self-correct
|
||
assert "bash" in prepared["error"]
|
||
# Surfaces warning to user
|
||
session.ui.on_error.assert_called_once()
|
||
assert "nonexistent" in session.ui.on_error.call_args[0][0]
|
||
|
||
def test_prepare_tool_strips_whitespace_from_name(self, tmp_db):
|
||
"""Local models may produce tool names with leading/trailing whitespace."""
|
||
session = self._make_session(mcp_client=None)
|
||
tc = {
|
||
"id": "call_strip",
|
||
"function": {"name": " bash\n", "arguments": '{"command": "echo hi"}'},
|
||
}
|
||
prepared = session._prepare_tool(tc)
|
||
assert prepared["func_name"] == "bash"
|
||
assert "error" not in prepared
|
||
|
||
def test_prepare_tool_malformed_json_surfaces_error(self, tmp_db):
|
||
"""Malformed JSON args should surface a warning to the user and
|
||
give the model a hint about expected format."""
|
||
session = self._make_session(mcp_client=None)
|
||
tc = {
|
||
"id": "call_bad",
|
||
"function": {"name": "bash", "arguments": "{command: echo hi}"},
|
||
}
|
||
prepared = session._prepare_tool(tc)
|
||
assert "error" in prepared
|
||
assert "JSON parse error" in prepared["error"]
|
||
assert "command" in prepared["error"] # hint about expected key
|
||
assert "Please retry" in prepared["error"]
|
||
# User-facing warning
|
||
session.ui.on_error.assert_called_once()
|
||
assert "Malformed tool call" in session.ui.on_error.call_args[0][0]
|
||
|
||
def test_ensure_tool_call_ids_dict(self, tmp_db):
|
||
"""_ensure_tool_call_ids fills empty IDs on streaming-style dict."""
|
||
from turnstone.core.session import ChatSession
|
||
|
||
tool_calls_acc = {
|
||
0: {"id": "", "function": {"name": "bash", "arguments": "{}"}},
|
||
1: {"id": "", "function": {"name": "read_file", "arguments": "{}"}},
|
||
}
|
||
ChatSession._ensure_tool_call_ids(tool_calls_acc)
|
||
ids = [tc["id"] for tc in tool_calls_acc.values()]
|
||
assert all(id_.startswith("call_") for id_ in ids)
|
||
assert len(set(ids)) == 2 # unique
|
||
|
||
def test_ensure_tool_call_ids_list(self, tmp_db):
|
||
"""_ensure_tool_call_ids fills empty IDs on list (agent path)."""
|
||
from turnstone.core.session import ChatSession
|
||
|
||
tool_calls = [
|
||
{"id": None, "function": {"name": "bash", "arguments": "{}"}},
|
||
{"id": "call_existing", "function": {"name": "bash", "arguments": "{}"}},
|
||
]
|
||
ChatSession._ensure_tool_call_ids(tool_calls)
|
||
assert tool_calls[0]["id"].startswith("call_")
|
||
assert tool_calls[1]["id"] == "call_existing" # preserved
|
||
|
||
def test_mcp_command_no_client(self, tmp_db):
|
||
session = self._make_session(mcp_client=None)
|
||
session.handle_command("/mcp")
|
||
session.ui.on_info.assert_called_once()
|
||
assert "No MCP servers" in session.ui.on_info.call_args[0][0]
|
||
|
||
def test_mcp_command_with_tools(self, tmp_db):
|
||
mock_mcp = MagicMock()
|
||
mock_mcp.get_tools.return_value = [_fake_openai_tool()]
|
||
session = self._make_session(mcp_client=mock_mcp)
|
||
session.handle_command("/mcp")
|
||
session.ui.on_info.assert_called_once()
|
||
output = session.ui.on_info.call_args[0][0]
|
||
assert "MCP tools (1)" in output
|
||
assert "mcp__test__search" in output
|
||
|
||
def test_exec_mcp_tool(self, tmp_db):
|
||
mock_mcp = MagicMock()
|
||
mock_mcp.get_tools.return_value = [_fake_openai_tool()]
|
||
mock_mcp.is_mcp_tool.return_value = True
|
||
mock_mcp.call_tool_sync.return_value = "result text"
|
||
session = self._make_session(mcp_client=mock_mcp)
|
||
|
||
item = {
|
||
"call_id": "call_789",
|
||
"mcp_func_name": "mcp__test__search",
|
||
"mcp_args": {"query": "hello"},
|
||
}
|
||
call_id, output = session._exec_mcp_tool(item)
|
||
assert call_id == "call_789"
|
||
assert output == "result text"
|
||
mock_mcp.call_tool_sync.assert_called_once_with(
|
||
"mcp__test__search", {"query": "hello"}, user_id=None, timeout=30
|
||
)
|
||
|
||
def test_exec_mcp_tool_error(self, tmp_db):
|
||
mock_mcp = MagicMock()
|
||
mock_mcp.get_tools.return_value = [_fake_openai_tool()]
|
||
mock_mcp.is_mcp_tool.return_value = True
|
||
mock_mcp.call_tool_sync.side_effect = RuntimeError("server crashed")
|
||
session = self._make_session(mcp_client=mock_mcp)
|
||
|
||
item = {
|
||
"call_id": "call_err",
|
||
"mcp_func_name": "mcp__test__search",
|
||
"mcp_args": {"query": "hello"},
|
||
}
|
||
call_id, output = session._exec_mcp_tool(item)
|
||
assert call_id == "call_err"
|
||
assert "MCP tool error" in output
|
||
assert "server crashed" in output
|
||
|
||
# -- 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 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 == {}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 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.
|
||
assert refresh_event.wait(timeout=5), "refresh task was not scheduled"
|
||
assert session is new_session
|
||
|
||
def test_auto_reconnect_swallows_refresh_failure(self, running_loop_mgr):
|
||
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),
|
||
):
|
||
# Must not raise — refresh failures are non-fatal.
|
||
session = mgr._cb_auto_reconnect("srv")
|
||
# Background refresh actually started and exception was swallowed
|
||
# by the task without affecting the synchronous caller.
|
||
assert refresh_started.wait(timeout=5), "refresh task was not scheduled"
|
||
assert session is new_session
|