mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-12 23:12:23 -06:00
Patrick Buckley
c823156af5
Phase 0 of the OAuth-MCP RFC: prepare MCPClientManager for the per-(user,
server) session pool that lands in Phase 5, without changing static-path
behavior.
Two changes:
1. Hardening helpers _pre_close_streams and _tcp_probe rename their first
parameter from `name` to `key`. Type stays `str` for now; widening to
`str | tuple[str, str]` happens in Phase 5 when callers actually pass
tuples. _safe_close_stack takes the stack directly and is unchanged.
2. The eleven parallel name-keyed dicts (_sessions, _per_server_stacks,
_per_server_tools, _per_server_resources, _per_server_prompts,
_supports_list_changed, _supports_resources, _supports_resource_list_changed,
_supports_prompts, _supports_prompt_list_changed, _server_streams) are
consolidated into _static_servers: dict[str, StaticServerState]. Server-
level state (circuit breaker, notification debounce, last-error,
db-managed, merged catalog maps, listener lists) stays on the manager,
unchanged.
PoolEntryState is defined for Phase 5 use but no code instantiates it. The
typed map declarations (dict[str, StaticServerState] vs dict[tuple[str, str],
PoolEntryState]) make accidental cross-keying lookups easier to catch.
PR #296 hardening preserved exactly:
- pre-close-streams atomic take-and-clear before stack teardown
- stale-session-and-stack guard at _connect_one top: both state.session and
state.stack checked, cleared independently, entry preserved (not popped)
- transport-error session-eviction in dispatch sets state.session=None only,
leaving stack/streams for the next connect-time guard sweep
- _safe_close_stack CancelledError suppression unchanged
- TCP probe before streamablehttp_client unchanged
- future.cancel() after TimeoutError in all sync bridges unchanged
- notification debounce stays manager-level (not migrated into the dataclass)
Refresh helpers (_refresh_server_tools/_resources/_prompts) snapshot
state.session into a local immediately after the None guard so concurrent
transport-error eviction during await cannot null the session reference
mid-call.
Tests: shared _seed_static_state helper in tests/conftest.py replaces eleven
direct dict mutations; new test_reconnect_preserves_static_state_identity
guards the entry-preservation invariant. Pass count rises 5266 → 5267.
(cherry picked from commit be0950bb98)
413 lines
15 KiB
Python
413 lines
15 KiB
Python
"""Integration tests for MCPClientManager data flow.
|
|
|
|
Uses real storage (SQLite) and real MCPClientManager state manipulation,
|
|
but mock MCP sessions instead of wire-protocol connections. This validates
|
|
the full data pipeline: per-server data -> rebuild -> merged state ->
|
|
query methods -> storage sync -> shutdown cleanup.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from typing import Any
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
|
|
from tests.conftest import _seed_static_state
|
|
from turnstone.core.mcp_client import MCPClientManager
|
|
from turnstone.core.storage._sqlite import SQLiteBackend
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_resource(
|
|
uri: str, name: str, server: str, description: str = "", mime: str = "text/plain"
|
|
) -> dict[str, Any]:
|
|
return {
|
|
"uri": uri,
|
|
"name": name,
|
|
"description": description,
|
|
"mimeType": mime,
|
|
"server": server,
|
|
}
|
|
|
|
|
|
def _make_prompt(
|
|
prefixed_name: str,
|
|
original_name: str,
|
|
server: str,
|
|
description: str = "",
|
|
arguments: list[dict[str, Any]] | None = None,
|
|
) -> dict[str, Any]:
|
|
return {
|
|
"name": prefixed_name,
|
|
"original_name": original_name,
|
|
"server": server,
|
|
"description": description,
|
|
"arguments": arguments or [],
|
|
}
|
|
|
|
|
|
def _make_mock_session(
|
|
read_resource_result: Any = None,
|
|
get_prompt_result: Any = None,
|
|
) -> AsyncMock:
|
|
"""Build a mock ClientSession with configurable async return values."""
|
|
session = AsyncMock()
|
|
|
|
if read_resource_result is not None:
|
|
session.read_resource.return_value = read_resource_result
|
|
else:
|
|
# Default: single text content
|
|
content_item = MagicMock()
|
|
content_item.text = "resource content"
|
|
result = MagicMock()
|
|
result.contents = [content_item]
|
|
session.read_resource.return_value = result
|
|
|
|
if get_prompt_result is not None:
|
|
session.get_prompt.return_value = get_prompt_result
|
|
else:
|
|
msg = MagicMock()
|
|
msg.role = "user"
|
|
msg.content = MagicMock()
|
|
msg.content.text = "Hello, World!"
|
|
result = MagicMock()
|
|
result.messages = [msg]
|
|
session.get_prompt.return_value = result
|
|
|
|
return session
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Integration test class
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestFullLifecycleResourcesPrompts:
|
|
"""Integration test exercising real code paths with real SQLite storage
|
|
but mock MCP sessions.
|
|
|
|
Validates the complete data flow: per-server data population, rebuild
|
|
merging, query methods, resource/prompt dispatch through asyncio, storage
|
|
sync, and shutdown cleanup.
|
|
"""
|
|
|
|
@pytest.fixture()
|
|
def mgr(self) -> MCPClientManager:
|
|
"""Create an MCPClientManager with no server configs (no start())."""
|
|
return MCPClientManager({})
|
|
|
|
@pytest.fixture()
|
|
def db(self, tmp_path) -> SQLiteBackend:
|
|
"""Create a fresh SQLite backend for each test."""
|
|
backend = SQLiteBackend(str(tmp_path / "test.db"))
|
|
yield backend
|
|
backend.close()
|
|
|
|
def test_rebuild_resources_produces_merged_state(self, mgr: MCPClientManager) -> None:
|
|
"""_rebuild_resources merges per-server resources into a unified list."""
|
|
_seed_static_state(
|
|
mgr,
|
|
"alpha",
|
|
resources=[
|
|
_make_resource("file:///a.txt", "a", "alpha"),
|
|
_make_resource("file:///b.txt", "b", "alpha"),
|
|
],
|
|
)
|
|
_seed_static_state(mgr, "beta", resources=[_make_resource("file:///c.txt", "c", "beta")])
|
|
|
|
mgr._rebuild_resources()
|
|
|
|
resources = mgr.get_resources()
|
|
assert len(resources) == 3
|
|
uris = {r["uri"] for r in resources}
|
|
assert uris == {"file:///a.txt", "file:///b.txt", "file:///c.txt"}
|
|
# resource_map should have entries for all non-template resources
|
|
assert "file:///a.txt" in mgr._resource_map
|
|
assert "file:///c.txt" in mgr._resource_map
|
|
assert mgr.resource_count == 3
|
|
|
|
def test_rebuild_prompts_produces_merged_state(self, mgr: MCPClientManager) -> None:
|
|
"""_rebuild_prompts merges per-server prompts into a unified list."""
|
|
_seed_static_state(
|
|
mgr,
|
|
"alpha",
|
|
prompts=[_make_prompt("mcp__alpha__greet", "greet", "alpha", "Say hello")],
|
|
)
|
|
_seed_static_state(
|
|
mgr,
|
|
"beta",
|
|
prompts=[
|
|
_make_prompt("mcp__beta__summarize", "summarize", "beta", "Summarize text"),
|
|
_make_prompt("mcp__beta__translate", "translate", "beta", "Translate text"),
|
|
],
|
|
)
|
|
|
|
mgr._rebuild_prompts()
|
|
|
|
prompts = mgr.get_prompts()
|
|
assert len(prompts) == 3
|
|
names = {p["name"] for p in prompts}
|
|
assert names == {"mcp__alpha__greet", "mcp__beta__summarize", "mcp__beta__translate"}
|
|
# prompt_map should map prefixed -> (server, original)
|
|
assert mgr._prompt_map["mcp__alpha__greet"] == ("alpha", "greet")
|
|
assert mgr._prompt_map["mcp__beta__summarize"] == ("beta", "summarize")
|
|
assert mgr.prompt_count == 3
|
|
assert mgr.is_mcp_prompt("mcp__alpha__greet") is True
|
|
assert mgr.is_mcp_prompt("nonexistent") is False
|
|
|
|
def test_read_resource_sync_dispatches_correctly(self, mgr: MCPClientManager) -> None:
|
|
"""read_resource_sync dispatches to the correct session via a real asyncio loop."""
|
|
# Set up a real event loop in a thread (simulating start())
|
|
loop = asyncio.new_event_loop()
|
|
import threading
|
|
|
|
thread = threading.Thread(target=loop.run_forever, daemon=True)
|
|
thread.start()
|
|
mgr._loop = loop
|
|
|
|
try:
|
|
# Populate session and resource map
|
|
session = _make_mock_session()
|
|
_seed_static_state(
|
|
mgr,
|
|
"alpha",
|
|
session=session,
|
|
resources=[_make_resource("file:///readme.md", "readme", "alpha")],
|
|
)
|
|
mgr._rebuild_resources()
|
|
|
|
result = mgr.read_resource_sync("file:///readme.md", timeout=5)
|
|
assert result == "resource content"
|
|
session.read_resource.assert_awaited_once_with("file:///readme.md")
|
|
finally:
|
|
loop.call_soon_threadsafe(loop.stop)
|
|
thread.join(timeout=5)
|
|
loop.close()
|
|
|
|
def test_read_resource_sync_unknown_uri_raises(self, mgr: MCPClientManager) -> None:
|
|
"""read_resource_sync raises ValueError for an unknown URI."""
|
|
with pytest.raises(ValueError, match="Unknown MCP resource"):
|
|
mgr.read_resource_sync("file:///nonexistent")
|
|
|
|
def test_read_resource_via_template(self, mgr: MCPClientManager) -> None:
|
|
"""Expanded template URI dispatched to correct server via real asyncio loop."""
|
|
loop = asyncio.new_event_loop()
|
|
import threading
|
|
|
|
thread = threading.Thread(target=loop.run_forever, daemon=True)
|
|
thread.start()
|
|
mgr._loop = loop
|
|
|
|
try:
|
|
session = _make_mock_session()
|
|
# Register a template resource (no concrete resources)
|
|
_seed_static_state(
|
|
mgr,
|
|
"alpha",
|
|
session=session,
|
|
resources=[
|
|
{
|
|
"uri": "db://tables/{table}/rows/{id}",
|
|
"name": "row",
|
|
"description": "Fetch a row",
|
|
"mimeType": "application/json",
|
|
"server": "alpha",
|
|
"template": True,
|
|
},
|
|
],
|
|
)
|
|
mgr._rebuild_resources()
|
|
|
|
# Template should not be in _resource_map
|
|
assert "db://tables/{table}/rows/{id}" not in mgr._resource_map
|
|
# But expanded URI should resolve via prefix matching
|
|
result = mgr.read_resource_sync("db://tables/users/rows/42", timeout=5)
|
|
assert result == "resource content"
|
|
session.read_resource.assert_awaited_once_with("db://tables/users/rows/42")
|
|
finally:
|
|
loop.call_soon_threadsafe(loop.stop)
|
|
thread.join(timeout=5)
|
|
loop.close()
|
|
|
|
def test_get_prompt_sync_dispatches_correctly(self, mgr: MCPClientManager) -> None:
|
|
"""get_prompt_sync dispatches to the correct session via a real asyncio loop."""
|
|
loop = asyncio.new_event_loop()
|
|
import threading
|
|
|
|
thread = threading.Thread(target=loop.run_forever, daemon=True)
|
|
thread.start()
|
|
mgr._loop = loop
|
|
|
|
try:
|
|
session = _make_mock_session()
|
|
_seed_static_state(
|
|
mgr,
|
|
"alpha",
|
|
session=session,
|
|
prompts=[_make_prompt("mcp__alpha__greet", "greet", "alpha", "Say hello")],
|
|
)
|
|
mgr._rebuild_prompts()
|
|
|
|
messages = mgr.get_prompt_sync(
|
|
"mcp__alpha__greet", arguments={"name": "World"}, timeout=5
|
|
)
|
|
assert len(messages) == 1
|
|
assert messages[0]["role"] == "user"
|
|
assert messages[0]["content"] == "Hello, World!"
|
|
session.get_prompt.assert_awaited_once_with("greet", arguments={"name": "World"})
|
|
finally:
|
|
loop.call_soon_threadsafe(loop.stop)
|
|
thread.join(timeout=5)
|
|
loop.close()
|
|
|
|
def test_get_prompt_sync_unknown_name_raises(self, mgr: MCPClientManager) -> None:
|
|
"""get_prompt_sync raises ValueError for an unknown prompt name."""
|
|
with pytest.raises(ValueError, match="Unknown MCP prompt"):
|
|
mgr.get_prompt_sync("mcp__nosrv__nope")
|
|
|
|
def test_sync_prompts_to_storage_creates_templates(
|
|
self, mgr: MCPClientManager, db: SQLiteBackend
|
|
) -> None:
|
|
"""sync_prompts_to_storage creates governance templates in real SQLite."""
|
|
mgr.set_storage(db)
|
|
mgr._prompts = [
|
|
_make_prompt(
|
|
"mcp__alpha__greet",
|
|
"greet",
|
|
"alpha",
|
|
"Say hello",
|
|
[{"name": "user", "description": "Who to greet", "required": True}],
|
|
),
|
|
_make_prompt(
|
|
"mcp__beta__summarize",
|
|
"summarize",
|
|
"beta",
|
|
"Summarize text",
|
|
),
|
|
]
|
|
# Mark connected so set_storage triggers sync
|
|
mgr._connected.set()
|
|
# Re-set storage to trigger auto-sync
|
|
mgr.set_storage(db)
|
|
|
|
templates = db.list_prompt_templates()
|
|
assert len(templates) == 2
|
|
names = {t["name"] for t in templates}
|
|
assert names == {"mcp__alpha__greet", "mcp__beta__summarize"}
|
|
|
|
# Verify details on first template
|
|
tpl = db.get_prompt_template_by_name("mcp__alpha__greet")
|
|
assert tpl is not None
|
|
assert tpl["origin"] == "mcp"
|
|
assert tpl["mcp_server"] == "alpha"
|
|
assert tpl["readonly"] is True
|
|
assert tpl["category"] == "mcp"
|
|
assert "user" in tpl["variables"]
|
|
|
|
def test_sync_prompts_removes_stale_templates(
|
|
self, mgr: MCPClientManager, db: SQLiteBackend
|
|
) -> None:
|
|
"""sync_prompts_to_storage removes templates whose MCP prompts are gone."""
|
|
mgr.set_storage(db)
|
|
|
|
# Create an initial template via sync
|
|
mgr._prompts = [
|
|
_make_prompt("mcp__alpha__old", "old", "alpha", "Old prompt"),
|
|
]
|
|
mgr.sync_prompts_to_storage()
|
|
assert len(db.list_prompt_templates()) == 1
|
|
|
|
# Now the prompt is gone
|
|
mgr._prompts = []
|
|
result = mgr.sync_prompts_to_storage()
|
|
assert result["removed"] == ["mcp__alpha__old"]
|
|
assert len(db.list_prompt_templates()) == 0
|
|
|
|
def test_shutdown_clears_all_state(self, mgr: MCPClientManager) -> None:
|
|
"""shutdown() clears sessions, tools, resources, prompts, and listeners."""
|
|
# Populate state
|
|
_seed_static_state(
|
|
mgr,
|
|
"alpha",
|
|
session=MagicMock(),
|
|
tools=[
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "mcp__alpha__search",
|
|
"description": "Search",
|
|
"parameters": {},
|
|
},
|
|
}
|
|
],
|
|
resources=[
|
|
_make_resource("file:///a.txt", "a", "alpha"),
|
|
{
|
|
"uri": "db://tables/{table}",
|
|
"name": "table",
|
|
"description": "",
|
|
"mimeType": "",
|
|
"server": "alpha",
|
|
"template": True,
|
|
},
|
|
],
|
|
prompts=[_make_prompt("mcp__alpha__greet", "greet", "alpha")],
|
|
)
|
|
mgr._rebuild_tools()
|
|
mgr._rebuild_resources()
|
|
mgr._rebuild_prompts()
|
|
mgr._listeners.append(lambda: None)
|
|
mgr._resource_listeners.append(lambda: None)
|
|
mgr._prompt_listeners.append(lambda: None)
|
|
|
|
# Verify populated
|
|
assert len(mgr._static_servers) == 1
|
|
assert len(mgr._tools) == 1
|
|
assert len(mgr._resources) == 2 # 1 concrete + 1 template
|
|
assert len(mgr._template_prefixes) == 1
|
|
assert len(mgr._prompts) == 1
|
|
|
|
mgr.shutdown()
|
|
|
|
assert len(mgr._static_servers) == 0
|
|
assert len(mgr._tools) == 0
|
|
assert len(mgr._tool_map) == 0
|
|
assert len(mgr._resources) == 0
|
|
assert len(mgr._resource_map) == 0
|
|
assert len(mgr._template_prefixes) == 0
|
|
assert len(mgr._prompts) == 0
|
|
assert len(mgr._prompt_map) == 0
|
|
assert len(mgr._listeners) == 0
|
|
assert len(mgr._resource_listeners) == 0
|
|
assert len(mgr._prompt_listeners) == 0
|
|
|
|
def test_listener_notifications_fire_on_rebuild(self, mgr: MCPClientManager) -> None:
|
|
"""Rebuild methods fire the appropriate listener callbacks."""
|
|
tool_fired = []
|
|
resource_fired = []
|
|
prompt_fired = []
|
|
mgr.add_listener(lambda: tool_fired.append(1))
|
|
mgr.add_resource_listener(lambda: resource_fired.append(1))
|
|
mgr.add_prompt_listener(lambda: prompt_fired.append(1))
|
|
|
|
_seed_static_state(mgr, "alpha", tools=[])
|
|
mgr._rebuild_tools()
|
|
assert len(tool_fired) == 1
|
|
|
|
_seed_static_state(mgr, "alpha", resources=[_make_resource("file:///x.txt", "x", "alpha")])
|
|
mgr._rebuild_resources()
|
|
assert len(resource_fired) == 1
|
|
|
|
_seed_static_state(mgr, "alpha", prompts=[_make_prompt("mcp__alpha__p1", "p1", "alpha")])
|
|
mgr._rebuild_prompts()
|
|
assert len(prompt_fired) == 1
|
|
|
|
# Tool and resource listeners should not have been fired again
|
|
assert len(tool_fired) == 1
|
|
assert len(resource_fired) == 1
|