diff --git a/tests/test_mcp_hot_reload.py b/tests/test_mcp_hot_reload.py index f0b264fd..feee0c79 100644 --- a/tests/test_mcp_hot_reload.py +++ b/tests/test_mcp_hot_reload.py @@ -268,8 +268,10 @@ class TestErrorTracking: """remove_server_sync cleans up _last_error entry.""" mgr = MCPClientManager({"test": {"command": "echo"}}) mgr._last_error["test"] = "Connection refused" + mgr._pool_discovery_error["test"] = "stale discovery failure" mgr.remove_server_sync("test") assert "test" not in mgr._last_error + assert "test" not in mgr._pool_discovery_error def test_all_server_status_includes_errors(self) -> None: """get_all_server_status propagates per-server errors.""" @@ -431,12 +433,33 @@ class TestReconcileSync: mgr.reconcile_sync(_FakeStorage([row])) assert primed == [] + def test_removed_pool_server_does_not_restore_stale_discovery_error(self) -> None: + """Pool rows bypass remove_server_sync, so reconcile's registry diff + must clear discovery state before a same-name server is re-added.""" + mgr = MCPClientManager({}) + mgr._oauth_user_server_names = {"pool-srv"} + mgr._pool_discovery_error["pool-srv"] = "old endpoint failed" + + mgr.reconcile_sync(_FakeStorage([])) + assert "pool-srv" not in mgr._pool_discovery_error + + row = _db_row( + "pool-srv", + transport="streamable-http", + command="", + url="https://new.example/mcp", + ) + row["auth_type"] = "oauth_user" + mgr.reconcile_sync(_FakeStorage([row])) + assert mgr.get_server_status("pool-srv")["discovery_error"] == "" + def test_reprimes_on_pool_auth_type_flip(self) -> None: """A server MIGRATED in place oauth_user -> oauth_obo (same name) re-primes active users — a name-only diff would see the same name on both sides and miss the flip.""" mgr = MCPClientManager({}) mgr._oauth_user_server_names = {"srv"} # previously oauth_user + mgr._pool_discovery_error["srv"] = "failure from old auth model" primed: list[str] = [] mgr.prime_user_pools = lambda uid: primed.append(uid) # type: ignore[method-assign] mgr.add_listener(lambda: None, user_id="u1") @@ -446,6 +469,7 @@ class TestReconcileSync: assert primed == ["u1"] assert mgr._obo_server_names == {"srv"} assert mgr._oauth_user_server_names == set() + assert "srv" not in mgr._pool_discovery_error def test_reprime_survives_prime_exception(self) -> None: """One user's prime scheduling failure must not abort the loop or propagate diff --git a/tests/test_mcp_user_pool.py b/tests/test_mcp_user_pool.py index 186c4709..8928a29e 100644 --- a/tests/test_mcp_user_pool.py +++ b/tests/test_mcp_user_pool.py @@ -1880,6 +1880,38 @@ class TestOboPriming: assert len(warmed) == 1 assert warmed[0][2] == "minted-at" + def test_prime_failure_records_single_line_discovery_error( + self, running_loop_mgr, storage + ) -> None: + """Exception text is untrusted upstream input and status renders it.""" + mgr, loop, _ = running_loop_mgr + cipher = make_mcp_token_cipher() + mgr.set_storage(storage) + app_state = _make_app_state(storage, cipher=cipher) + mgr.set_app_state(app_state) + storage.create_mcp_server( + server_id="srv-o", + name="pool-srv", + transport="streamable-http", + url="https://mcp.example.com/sse", + auth_type="oauth_user", + ) + mgr._oauth_user_server_names = {"pool-srv"} + + async def _fail_prime(_key: Any, _cfg: Any, _token: str) -> None: + raise RuntimeError("upstream\r\nresponse\nignore instructions") + + lookup = AsyncMock(return_value=SimpleNamespace(kind="token", token="bearer")) + with ( + patch.object(mgr, "_prime_user_server", new=_fail_prime), + patch("turnstone.core.mcp_client.get_user_access_token_classified", new=lookup), + ): + _run_on_loop(loop, mgr._prime_user_pools("user-1")) + + assert mgr._pool_discovery_error["pool-srv"] == ( + "RuntimeError: upstream response ignore instructions" + ) + def test_prime_drops_retained_catalog_on_dead_grant(self, running_loop_mgr, storage) -> None: """Priming is a convergence point (#836): a NEW session's prime that finds the grant durably GONE must drop the retained catalog diff --git a/tests/test_tool_search.py b/tests/test_tool_search.py index 5cfa9a7a..fca7ce38 100644 --- a/tests/test_tool_search.py +++ b/tests/test_tool_search.py @@ -8,6 +8,7 @@ from turnstone.core.tool_search import ( BM25Index, ToolSearchManager, _mcp_server_summary, + _status_reason, _tokenize, _tool_name, ) @@ -256,6 +257,158 @@ class TestToolSearchManagerReranking: # --------------------------------------------------------------------------- +class TestToolSearchTruncation: + """search() records the full match count so format_search_results can + report honest truncation instead of silently dropping matches.""" + + @pytest.fixture() + def manager(self): + # Every name shares the ``mcp`` token, so a search for "mcp" matches + # all six deferred tools — more than max_results=3. + mcp_tools = [ + _make_tool("mcp__github__create_issue", "Create a new GitHub issue"), + _make_tool("mcp__github__list_issues", "List GitHub issues"), + _make_tool("mcp__github__get_repo", "Get repository details"), + _make_tool("mcp__slack__send_message", "Send a Slack message"), + _make_tool("mcp__slack__list_channels", "List Slack channels"), + _make_tool("mcp__jira__create_ticket", "Create a Jira ticket"), + ] + return ToolSearchManager(mcp_tools, always_on_names=set(), max_results=3) + + def test_search_records_total_matched(self, manager): + results = manager.search("mcp") + assert len(results) == 3 + assert manager._last_total_matched == 6 + + def test_format_reports_truncation(self, manager): + results = manager.search("mcp") + text = manager.format_search_results(results) + assert "Showing the top 3 of 6" in text + + def test_no_truncation_note_when_all_returned(self, manager): + results = manager.search("jira") # matches only the jira tool + text = manager.format_search_results(results) + assert "Showing the top" not in text + assert "Found 1" in text + + def test_truncation_excludes_already_expanded_from_total(self, manager): + # Expanding one match shrinks the reported total (only genuinely-new + # matches count), so "of N" never over-promises tools already loaded. + manager.expand_visible(["mcp__github__create_issue"]) + manager.search("mcp") + assert manager._last_total_matched == 5 + + +class TestToolSearchUnavailableAdvisory: + """A down/unauthorized server is surfaced, never silently treated as + 'no such tool'. Driven by the injected status_provider.""" + + def _mgr(self, tools, status): + return ToolSearchManager(tools, always_on_names=set(), status_provider=lambda: status) + + def test_empty_results_with_outage_explains_outage(self): + mgr = self._mgr( + [_make_tool("mcp__dhcp__GetLease", "Get a dhcp lease")], + {"DHCP-MCP": {"error": "500 app failed to start"}}, + ) + text = mgr.format_search_results(mgr.search("nonexistent_capability_xyz")) + assert "unavailable" in text.lower() or "outage" in text.lower() + assert "DHCP-MCP" in text + assert "500 app failed to start" in text + # Must NOT give the misleading "try a different query" line alone. + assert text != "No matching tools found. Try a different search query." + + def test_advisory_appended_when_results_present(self): + mgr = self._mgr( + [ + _make_tool("mcp__github__create_issue", "Create a github issue"), + _make_tool("mcp__dhcp__GetLease", "Get a dhcp lease"), + ], + {"DHCP-MCP": {"discovery_error": "500 app failed"}}, + ) + text = mgr.format_search_results(mgr.search("github")) + assert "Found 1" in text + assert "currently unavailable" in text + assert "DHCP-MCP" in text + assert "tool discovery failed" in text + + def test_circuit_open_flagged(self): + mgr = self._mgr( + [_make_tool("mcp__x__t", "thing")], + {"X": {"circuit_open": True}}, + ) + text = mgr.format_search_results([]) + assert "circuit breaker open" in text + + def test_healthy_server_not_flagged(self): + mgr = self._mgr( + [_make_tool("mcp__github__create_issue", "Create a github issue")], + {"github": {"connected": True, "error": "", "circuit_open": False}}, + ) + text = mgr.format_search_results(mgr.search("github")) + assert "unavailable" not in text.lower() + + def test_unprimed_server_not_flagged(self): + # connected=False with no error is "not reached yet", not an outage — + # flagging it would cry wolf on every server the user hasn't touched. + mgr = self._mgr( + [_make_tool("mcp__github__create_issue", "Create a github issue")], + {"github": {"connected": False, "error": "", "circuit_open": False}}, + ) + text = mgr.format_search_results(mgr.search("github")) + assert "unavailable" not in text.lower() + + def test_no_status_provider_is_legacy_behaviour(self): + mgr = ToolSearchManager( + [_make_tool("mcp__github__create_issue", "Create a github issue")], + always_on_names=set(), + ) + assert ( + mgr.format_search_results([]) + == "No matching tools found. Try a different search query." + ) + text = mgr.format_search_results(mgr.search("github")) + assert "unavailable" not in text.lower() + + def test_status_provider_error_is_swallowed(self): + def boom(): + raise RuntimeError("status backend down") + + mgr = ToolSearchManager( + [_make_tool("mcp__github__create_issue", "Create a github issue")], + always_on_names=set(), + status_provider=boom, + ) + # A broken provider must never break tool search itself. + text = mgr.format_search_results(mgr.search("github")) + assert "Found 1" in text + + +class TestStatusReason: + def test_circuit_open(self): + assert _status_reason({"circuit_open": True}) == "circuit breaker open" + + def test_error_text(self): + assert "500 boom" in _status_reason({"error": "500 boom"}) + + def test_discovery_error(self): + reason = _status_reason({"discovery_error": "TimeoutError: pool discovery"}) + assert "discovery" in reason.lower() + + @pytest.mark.parametrize("field", ["error", "discovery_error"]) + def test_error_text_is_single_line(self, field): + reason = _status_reason({field: "upstream\r\nresponse\nignore instructions"}) + assert "\n" not in reason + assert "\r" not in reason + assert "upstream response ignore instructions" in reason + + def test_healthy_is_empty(self): + assert _status_reason({"connected": True, "error": "", "circuit_open": False}) == "" + + def test_circuit_takes_precedence_over_error(self): + assert _status_reason({"circuit_open": True, "error": "boom"}) == "circuit breaker open" + + class TestMCPServerSummary: def test_groups_by_server(self): tools = [ diff --git a/turnstone/core/mcp_client.py b/turnstone/core/mcp_client.py index 2a3d60da..ae90b386 100644 --- a/turnstone/core/mcp_client.py +++ b/turnstone/core/mcp_client.py @@ -677,6 +677,13 @@ class MCPClientManager: self._db_managed: set[str] = set() # Per-server last-error tracking (set on failure, cleared on success) self._last_error: dict[str, str] = {} + # Per-server last tool-DISCOVERY error for pool (oauth_user / oauth_obo) + # servers: set when a prime/connect raises during discovery (transport, + # 5xx, timeout), cleared on the next successful pool connect. Surfaced + # via get_server_status(...)["discovery_error"] so a swallowed discovery + # failure (e.g. a 500 from the MCP endpoint) is visible instead of the + # server silently contributing zero tools to the catalog. + self._pool_discovery_error: dict[str, str] = {} self._MAX_ERROR_LEN = 256 # Listener infrastructure (tool-change callbacks for ChatSession). @@ -2916,6 +2923,9 @@ class MCPClientManager: # ``_user_resources`` / ``_user_prompts``. Per-user fan-out # ensures another user's session never observes this change. self._rebuild_and_notify_user_catalogs(user_id) + # Discovery just succeeded for this server — clear any prior recorded + # pool discovery failure so a healed outage doesn't linger on status. + self._pool_discovery_error.pop(server_name, None) return entry # -- pool priming --------------------------------------------------------- @@ -3147,7 +3157,20 @@ class MCPClientManager: user_id, server_name, ) - except Exception: + except Exception as exc: + # Record the discovery failure so it is visible via + # server status (and tool_search's unavailable-server + # advisory) instead of being swallowed to a debug log + # with the server silently contributing zero tools. + # Token-level outcomes (missing / dead grant) return + # earlier and never reach here, so this is a genuine + # connect/discovery failure (transport, 5xx, timeout). + detail = ( + f"{type(exc).__name__}: {exc}".replace("\n", " ") + .replace("\r", "") + .strip() + ) + self._pool_discovery_error[server_name] = detail[:200] log.debug( "mcp pool auto-prime failed user=%s server=%s", user_id, @@ -5664,6 +5687,7 @@ class MCPClientManager: # nothing would cover the dropped change. self._static_servers.pop(name, None) self._last_error.pop(name, None) + self._pool_discovery_error.pop(name, None) self._clear_static_push_state(name, markers=True) self._cb_clear(name) # Clear health-loop backoff/ping state so a later @@ -5706,6 +5730,7 @@ class MCPClientManager: self._server_configs.pop(name, None) self._static_servers.pop(name, None) self._last_error.pop(name, None) + self._pool_discovery_error.pop(name, None) self._clear_static_push_state(name, markers=True) self._cb_clear(name) self._rebuild_tools() @@ -5761,6 +5786,7 @@ class MCPClientManager: ), "prompts": len(state.prompts) if state is not None and state.session is not None else 0, "error": self._last_error.get(name, ""), + "discovery_error": self._pool_discovery_error.get(name, ""), "transport": transport, "command": cfg.get("command", "") if transport == "stdio" else "", "url": cfg.get("url", "") if transport != "stdio" else "", @@ -5826,6 +5852,7 @@ class MCPClientManager: "resources": len(rep.resources) if rep is not None and rep.resources else 0, "prompts": len(rep.prompts) if rep is not None and rep.prompts else 0, "error": self._last_error.get(name, ""), + "discovery_error": self._pool_discovery_error.get(name, ""), "transport": "streamable-http", "command": "", "url": "", @@ -5900,6 +5927,13 @@ class MCPClientManager: self._obo_server_names = new_obo_names new_pool_auth = {n: "oauth_user" for n in self._oauth_user_server_names} new_pool_auth.update(dict.fromkeys(self._obo_server_names, "oauth_obo")) + # Pool-backed rows never enter ``_server_configs`` and therefore do + # not pass through remove_server_sync. Clear the old registration's + # discovery failure here when a pool server is removed or changes its + # auth model, so a later same-name registration cannot inherit it. + for name, auth_type in prev_pool_auth.items(): + if new_pool_auth.get(name) != auth_type: + self._pool_discovery_error.pop(name, None) # Pool servers newly registered OR migrated between pool auth types # since active sessions last primed. newly_added_pool = { diff --git a/turnstone/core/session.py b/turnstone/core/session.py index a3252d6b..6d192788 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -1957,6 +1957,7 @@ class ChatSession: always_on_names=builtin_in_session, max_results=tool_search_max_results, reranker=self._bm25_reranker(), + status_provider=self._mcp_status_snapshot, ) # Converge with any MCP catalog change that fired during # construction: a listener callback landing between the @@ -3008,6 +3009,7 @@ class ChatSession: always_on_names=set(BUILTIN_TOOL_NAMES), max_results=self._tool_search_max_results, reranker=self._bm25_reranker(), + status_provider=self._mcp_status_snapshot, ) # Restore previously expanded tools that still exist if old_expanded: @@ -3015,6 +3017,21 @@ class ChatSession: else: self._tool_search = None + def _mcp_status_snapshot(self) -> dict[str, dict[str, Any]]: + """Per-server MCP status for this session's user, consumed by + ``ToolSearchManager`` to flag unavailable servers in search results. + + Returns ``{}`` when no MCP client is bound or the lookup fails — the + advisory then simply stays silent rather than breaking tool search. + """ + client = self._mcp_client + if client is None: + return {} + try: + return client.get_all_server_status(self._mcp_user_id) + except Exception: + return {} + def set_watch_runner(self, runner: Any, wake_fn: Callable[[], object] | None = None) -> None: """Inject the server-level WatchRunner and register a dispatch fn that routes watch results onto this session's NudgeQueue. diff --git a/turnstone/core/tool_search.py b/turnstone/core/tool_search.py index fd4cb275..8121af30 100644 --- a/turnstone/core/tool_search.py +++ b/turnstone/core/tool_search.py @@ -15,6 +15,8 @@ from typing import TYPE_CHECKING, Any from turnstone.core.bm25 import BM25Index, _tokenize # noqa: F401 if TYPE_CHECKING: + from collections.abc import Callable + from turnstone.core.rerank import Reranker # --------------------------------------------------------------------------- @@ -54,6 +56,30 @@ def _mcp_server_summary(tools: list[dict[str, Any]]) -> str: return ", ".join(parts) +def _status_reason(status: dict[str, Any]) -> str: + """Human-readable reason a server is unavailable, or ``""`` when it looks + healthy. Reads only fields ``MCPClientManager.get_server_status`` already + exposes, so this stays decoupled from the client internals. + + Deliberately conservative: a merely un-primed server (``connected=False`` + with no error) is NOT treated as a failure — surfacing that would cry wolf + on every server the user simply hasn't reached yet. Only hard signals + (open circuit breaker, a recorded error, or a recorded discovery failure) + mark a server unavailable. + """ + if status.get("circuit_open"): + return "circuit breaker open" + # The provider is injectable, so enforce the single-line output contract + # at the rendering boundary even though MCPClientManager also sanitizes. + err = str(status.get("error") or "").replace("\n", " ").replace("\r", "").strip() + if err: + return f"error: {err[:120]}" + disc = str(status.get("discovery_error") or "").replace("\n", " ").replace("\r", "").strip() + if disc: + return f"tool discovery failed: {disc[:120]}" + return "" + + class ToolSearchManager: """Session-scoped tool visibility manager with BM25 search. @@ -69,12 +95,21 @@ class ToolSearchManager: *, max_results: int = 5, reranker: Reranker | None = None, + status_provider: Callable[[], dict[str, dict[str, Any]]] | None = None, ) -> None: self._always_on: list[dict[str, Any]] = [] self._deferred: list[dict[str, Any]] = [] self._deferred_by_name: dict[str, dict[str, Any]] = {} self._expanded: dict[str, None] = {} # ordered set (preserves discovery order) self._max_results = max_results + # Optional callback -> {server_name: status_dict} (the shape + # MCPClientManager.get_all_server_status returns). Lets search results + # flag servers that are down/unauthorized so a failed discovery is + # never mistaken for "no such tool". None keeps the legacy behaviour. + self._status_provider = status_provider + # Total matches from the most recent search(), before the max_results + # slice — so format_search_results can report honest truncation. + self._last_total_matched = 0 for tool in all_tools: name = _tool_name(tool) @@ -108,16 +143,22 @@ class ToolSearchManager: """Search deferred tools by query, return top-k matches. Already-expanded tools are excluded so every result is genuinely new. + Records the full (pre-slice) match count on ``_last_total_matched`` so + callers can report honest truncation instead of silently dropping + matches past ``max_results``. """ - # Request extra results to compensate for filtering out expanded tools - indices = self._index.search(query, k=self._max_results + len(self._expanded)) - results = [] - for i in indices: - if _tool_name(self._deferred[i]) not in self._expanded: - results.append(self._deferred[i]) - if len(results) >= self._max_results: - break - return results + # Rank the WHOLE deferred corpus (k = len(deferred)) so the count + # reflects every match, not just the max_results slice the model gets. + # Indices map 1:1 to self._deferred (the index was built over it in + # order). Cheap: the tool corpus is small (tens, not thousands). + indices = self._index.search(query, k=max(len(self._deferred), 1)) + matched = [ + self._deferred[i] + for i in indices + if _tool_name(self._deferred[i]) not in self._expanded + ] + self._last_total_matched = len(matched) + return matched[: self._max_results] def get_expanded_names(self) -> list[str]: """Return names of currently expanded (discovered) tools.""" @@ -172,9 +213,53 @@ class ToolSearchManager: }, } + def _unavailable_servers(self) -> list[tuple[str, str]]: + """``(server_name, reason)`` for every MCP server the status provider + reports as failing, sorted by name. + + Empty when no provider is wired or every server looks healthy. A + provider that raises is treated as "no info" (empty) so a status + glitch can never break tool search itself. + """ + if self._status_provider is None: + return [] + try: + statuses = self._status_provider() or {} + except Exception: + return [] + out: list[tuple[str, str]] = [] + for name, status in statuses.items(): + if not isinstance(status, dict): + continue + reason = _status_reason(status) + if reason: + out.append((str(name), reason)) + out.sort() + return out + + @staticmethod + def _format_unavailable(items: list[tuple[str, str]]) -> str: + return ", ".join(f"{name} ({reason})" for name, reason in items) + def format_search_results(self, tools: list[dict[str, Any]]) -> str: - """Format search results as text for the tool_search response.""" + """Format search results as text for the tool_search response. + + Beyond the matched tools, this surfaces two things the raw list hides: + an honest truncation note when more tools matched than were returned, + and a warning when known MCP servers are currently unavailable — so a + down server (which contributes zero searchable tools) is never + mistaken for a missing capability. + """ + unavailable = self._unavailable_servers() if not tools: + if unavailable: + return ( + "No matching tools found for that query. Note: " + + self._format_unavailable(unavailable) + + ". A server that is unavailable contributes no searchable " + "tools until it recovers, so this may be an outage rather " + "than a missing capability." + ) return "No matching tools found. Try a different search query." lines = [] for tool in tools: @@ -182,8 +267,16 @@ class ToolSearchManager: name = fn.get("name", "") desc = fn.get("description", "") lines.append(f"- **{name}**: {desc}") - return ( - f"Found {len(tools)} matching tool(s):\n" - + "\n".join(lines) - + "\n\nThese tools are now available for use." - ) + out = f"Found {len(tools)} matching tool(s):\n" + "\n".join(lines) + if self._last_total_matched > len(tools): + out += ( + f"\n\n(Showing the top {len(tools)} of {self._last_total_matched} " + "matches — narrow your query to surface the rest.)" + ) + out += "\n\nThese tools are now available for use." + if unavailable: + out += ( + "\n\n⚠ Some tool servers are currently unavailable, so their " + "tools are not searchable: " + self._format_unavailable(unavailable) + "." + ) + return out