mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-27 22:34:51 -06:00
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)
223 lines
8.8 KiB
Python
223 lines
8.8 KiB
Python
"""Tests for turnstone.core.web_search — pluggable web search backends."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from turnstone.core.web_search import (
|
|
DuckDuckGoClient,
|
|
MCPSearchClient,
|
|
TavilyClient,
|
|
_format_ddg,
|
|
_format_tavily,
|
|
resolve_web_search_client,
|
|
)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Formatters
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestFormatTavily:
|
|
def test_formats_answer_and_results(self):
|
|
data = {
|
|
"answer": "Python is great",
|
|
"results": [
|
|
{"title": "Python.org", "url": "https://python.org", "content": "Official site"},
|
|
{"title": "PyPI", "url": "https://pypi.org", "content": "Package index"},
|
|
],
|
|
}
|
|
out = _format_tavily(data, "python")
|
|
assert "Answer: Python is great" in out
|
|
assert "[Python.org](https://python.org)" in out
|
|
assert "[PyPI](https://pypi.org)" in out
|
|
|
|
def test_no_results(self):
|
|
out = _format_tavily({"results": []}, "nothing")
|
|
assert "No results for 'nothing'" in out
|
|
|
|
def test_no_answer(self):
|
|
data = {
|
|
"results": [{"title": "T", "url": "http://t", "content": "C"}],
|
|
}
|
|
out = _format_tavily(data, "q")
|
|
assert "Answer:" not in out
|
|
assert "[T](http://t)" in out
|
|
|
|
|
|
class TestFormatDDG:
|
|
def test_formats_results(self):
|
|
results = [
|
|
{"title": "DDG Result", "href": "https://ddg.example.com", "body": "Search body"},
|
|
]
|
|
out = _format_ddg(results, "test")
|
|
assert "[DDG Result](https://ddg.example.com)" in out
|
|
assert "Search body" in out
|
|
|
|
def test_no_results(self):
|
|
out = _format_ddg([], "nothing")
|
|
assert "No results for 'nothing'" in out
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Client tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestTavilyClient:
|
|
def test_search_calls_api(self):
|
|
mock_resp = MagicMock()
|
|
mock_resp.json.return_value = {
|
|
"answer": "42",
|
|
"results": [{"title": "T", "url": "http://t", "content": "C"}],
|
|
}
|
|
with patch("turnstone.core.web_search.httpx.post", return_value=mock_resp) as mock_post:
|
|
client = TavilyClient("test-key", timeout=10)
|
|
result = client.search("meaning of life", max_results=3)
|
|
|
|
mock_post.assert_called_once()
|
|
call_kwargs = mock_post.call_args
|
|
assert call_kwargs.kwargs["json"]["query"] == "meaning of life"
|
|
assert call_kwargs.kwargs["json"]["max_results"] == 3
|
|
assert "Answer: 42" in result
|
|
|
|
|
|
class TestDuckDuckGoClient:
|
|
def test_integration_via_mock_ddgs(self):
|
|
"""Patch the ddgs import inside DuckDuckGoClient.search."""
|
|
mock_ddgs = MagicMock()
|
|
mock_ddgs.__enter__ = MagicMock(return_value=mock_ddgs)
|
|
mock_ddgs.__exit__ = MagicMock(return_value=False)
|
|
mock_ddgs.text.return_value = [
|
|
{"title": "DDG Result", "href": "https://ddg.co", "body": "Found it"},
|
|
]
|
|
mock_module = MagicMock()
|
|
mock_module.DDGS.return_value = mock_ddgs
|
|
with patch.dict("sys.modules", {"ddgs": mock_module}):
|
|
client = DuckDuckGoClient(timeout=10)
|
|
result = client.search("test query", max_results=3)
|
|
mock_ddgs.text.assert_called_once_with("test query", max_results=3)
|
|
assert "[DDG Result](https://ddg.co)" in result
|
|
assert "Found it" in result
|
|
|
|
|
|
class TestMCPSearchClient:
|
|
def test_delegates_to_mcp(self):
|
|
mcp = MagicMock()
|
|
mcp.call_tool_sync.return_value = "MCP search results"
|
|
client = MCPSearchClient(mcp, "mcp__ddg__search", timeout=30)
|
|
result = client.search("test", max_results=3, topic="news")
|
|
mcp.call_tool_sync.assert_called_once_with(
|
|
"mcp__ddg__search",
|
|
{"query": "test", "max_results": 3, "topic": "news"},
|
|
timeout=30,
|
|
)
|
|
assert result == "MCP search results"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Resolver
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestResolveClient:
|
|
def test_auto_tavily_when_key_present(self):
|
|
client = resolve_web_search_client("", tavily_key="key")
|
|
assert isinstance(client, TavilyClient)
|
|
|
|
def test_auto_ddg_when_no_tavily(self):
|
|
with patch("turnstone.core.web_search._ddg_available", return_value=True):
|
|
client = resolve_web_search_client("", tavily_key=None)
|
|
assert isinstance(client, DuckDuckGoClient)
|
|
|
|
def test_auto_none_when_nothing_available(self):
|
|
with patch("turnstone.core.web_search._ddg_available", return_value=False):
|
|
client = resolve_web_search_client("", tavily_key=None)
|
|
assert client is None
|
|
|
|
def test_explicit_tavily(self):
|
|
client = resolve_web_search_client("tavily", tavily_key="key")
|
|
assert isinstance(client, TavilyClient)
|
|
|
|
def test_explicit_tavily_no_key(self):
|
|
client = resolve_web_search_client("tavily", tavily_key=None)
|
|
assert client is None
|
|
|
|
def test_explicit_ddg(self):
|
|
with patch("turnstone.core.web_search._ddg_available", return_value=True):
|
|
client = resolve_web_search_client("ddg", tavily_key=None)
|
|
assert isinstance(client, DuckDuckGoClient)
|
|
|
|
def test_explicit_ddg_not_installed(self):
|
|
with patch("turnstone.core.web_search._ddg_available", return_value=False):
|
|
client = resolve_web_search_client("ddg", tavily_key=None)
|
|
assert client is None
|
|
|
|
def test_mcp_backend(self):
|
|
mcp = MagicMock()
|
|
mcp.is_mcp_tool.return_value = True
|
|
client = resolve_web_search_client("mcp:ddg:search", tavily_key=None, mcp_client=mcp)
|
|
assert isinstance(client, MCPSearchClient)
|
|
mcp.is_mcp_tool.assert_called_with("mcp__ddg__search")
|
|
|
|
def test_mcp_backend_not_connected(self):
|
|
mcp = MagicMock()
|
|
mcp.is_mcp_tool.return_value = False
|
|
client = resolve_web_search_client("mcp:ddg:search", tavily_key=None, mcp_client=mcp)
|
|
assert client is None
|
|
|
|
def test_mcp_backend_no_client(self):
|
|
client = resolve_web_search_client("mcp:ddg:search", tavily_key=None, mcp_client=None)
|
|
assert client is None
|
|
|
|
def test_unknown_backend_returns_none(self):
|
|
client = resolve_web_search_client("typo_backend", tavily_key="key")
|
|
assert client is None
|
|
|
|
def test_resolve_web_search_client_rejects_oauth_user_backend(self):
|
|
"""A web_search backend pointing at an ``auth_type=oauth_user``
|
|
MCP server MUST be rejected at boot — per-node web_search
|
|
cannot carry per-user tokens, so resolving the backend would
|
|
guarantee a 401-on-call instead of a clean disablement.
|
|
|
|
Phase 7 invariant 8 corollary: pool tools are user-scoped;
|
|
every entry point that lacks per-user identity (web_search
|
|
boot resolver, eval harness, CLI default) MUST refuse them
|
|
rather than silently produce a broken client.
|
|
|
|
Verified by reverting the ``server_auth_type(...) == 'oauth_user'``
|
|
guard in ``resolve_web_search_client``: the resolver returns
|
|
an ``MCPSearchClient`` whose ``call_tool_sync`` would surface
|
|
a 401 / consent_required structured error on every search.
|
|
"""
|
|
mcp = MagicMock()
|
|
mcp.is_mcp_tool.return_value = True # name resolves
|
|
mcp.server_auth_type.return_value = "oauth_user"
|
|
client = resolve_web_search_client(
|
|
"mcp:oauth-search:search", tavily_key=None, mcp_client=mcp
|
|
)
|
|
assert client is None, (
|
|
"oauth_user-backed web_search backend resolved to a non-None client; "
|
|
"boot-time guard missing or regressed."
|
|
)
|
|
# Per-turn callers must read from the in-memory cache, never
|
|
# the SQL helper — perf regression guard.
|
|
mcp.server_auth_type.assert_called_with("oauth-search")
|
|
assert not mcp._lookup_server_row.called, (
|
|
"resolver issued a SQL roundtrip via _lookup_server_row; "
|
|
"per-turn web_search backend resolution must use the "
|
|
"in-memory server_auth_type accessor."
|
|
)
|
|
|
|
def test_resolve_web_search_client_accepts_static_backend(self):
|
|
"""Static-path (``auth_type=none`` or ``static``) MCP backends
|
|
still resolve cleanly — the new guard ONLY rejects oauth_user.
|
|
"""
|
|
mcp = MagicMock()
|
|
mcp.is_mcp_tool.return_value = True
|
|
mcp.server_auth_type.return_value = None
|
|
client = resolve_web_search_client(
|
|
"mcp:static-search:search", tavily_key=None, mcp_client=mcp
|
|
)
|
|
assert isinstance(client, MCPSearchClient)
|