diff --git a/tests/test_bm25.py b/tests/test_bm25.py index 625cf26a..2d7ed837 100644 --- a/tests/test_bm25.py +++ b/tests/test_bm25.py @@ -151,8 +151,8 @@ class TestBM25Reranking: # REORDER MODE (default rerank_filters=False, reactive tool/skill # search): an empty reranker result means the endpoint failed, so fall # back to BM25 order -- results are NEVER silently dropped. [bug-1] - # Guards the reorder-branch ``if not out: return pool[:k]`` -- removing - # it makes this return [] and fail. + # Guards the reorder-branch ``if not out: return ranked[:k]`` -- + # removing it makes this return [] and fail. index = BM25Index(self._DOCS, reranker=lambda q, d: []) assert index.search("alpha", k=5) == BM25Index(self._DOCS)._bm25_rank("alpha")[:5] @@ -217,7 +217,9 @@ class TestBM25Reranking: def test_recall_pool_capped_at_rerank_pool(self): # More matching docs than the recall cap: the reranker must receive - # exactly _RERANK_POOL docs, and outputs may only reference that set. + # exactly _RERANK_POOL docs. The cap bounds what the reranker SEES, + # not what the caller gets — matches past the pool trail in BM25 + # order, so a full-corpus k still returns every match. n = _RERANK_POOL + 10 docs = [f"alpha doc{i}" for i in range(n)] seen_len = {"n": -1} @@ -229,10 +231,58 @@ class TestBM25Reranking: index = BM25Index(docs, reranker=rec) result = index.search("alpha", k=n) assert seen_len["n"] == _RERANK_POOL # only the first 50 reached rerank - assert len(result) == _RERANK_POOL - # Every returned index is from the BM25 top-50 recall set. - recall = set(BM25Index(docs)._bm25_rank("alpha")[:_RERANK_POOL]) - assert set(result) == recall + # Identity rerank over the pool + BM25-ordered tail == full BM25 order. + assert result == BM25Index(docs)._bm25_rank("alpha") + + def test_matches_past_the_pool_trail_in_bm25_order(self): + # REORDER MODE with k > _RERANK_POOL: the reranked head comes first, + # then every match past the pool in BM25 order. Guards the tail loop — + # without it the result caps at the pool and tool_search's "top N of M" + # count floors at the pool size. [#941] + n = _RERANK_POOL + 10 + docs = [f"alpha doc{i}" for i in range(n)] + bm25_all = BM25Index(docs)._bm25_rank("alpha") + index = BM25Index(docs, reranker=lambda q, d: list(range(len(d)))[::-1]) + result = index.search("alpha", k=n) + assert result == bm25_all[:_RERANK_POOL][::-1] + bm25_all[_RERANK_POOL:] + + def test_reorder_fallbacks_not_capped_at_pool(self): + # Both reorder-mode fallbacks (reranker exception, empty result) must + # return the FULL BM25 ranking for k > _RERANK_POOL, not the pool + # slice — the fallback path must be as honest as the happy path. [#941] + def boom(q, d): + raise RuntimeError("rerank endpoint down") + + n = _RERANK_POOL + 10 + docs = [f"alpha doc{i}" for i in range(n)] + bm25_all = BM25Index(docs)._bm25_rank("alpha") + assert BM25Index(docs, reranker=boom).search("alpha", k=n) == bm25_all + assert BM25Index(docs, reranker=lambda q, d: []).search("alpha", k=n) == bm25_all + + def test_past_pool_tail_respects_k(self): + # k between the pool size and the match count: the tail must stop at + # k, not run to the end of the ranking — callers slice rows by the + # returned indices, so over-returning corrupts their result sets. + n = _RERANK_POOL + 10 + k = _RERANK_POOL + 5 + docs = [f"alpha doc{i}" for i in range(n)] + bm25_all = BM25Index(docs)._bm25_rank("alpha") + index = BM25Index(docs, reranker=lambda q, d: list(range(len(d)))[::-1]) + result = index.search("alpha", k=k) + assert len(result) == k + assert result == bm25_all[:_RERANK_POOL][::-1] + bm25_all[_RERANK_POOL:k] + + def test_filter_mode_error_fallback_keeps_pool_bound(self): + # FILTER MODE's error fallback stays pool-bounded: its happy path can + # never exceed the pool, and an endpoint failure must not return a + # longer list than a working endpoint ever could. + def boom(q, d): + raise RuntimeError("rerank endpoint down") + + n = _RERANK_POOL + 10 + docs = [f"alpha doc{i}" for i in range(n)] + index = BM25Index(docs, reranker=boom, rerank_filters=True) + assert index.search("alpha", k=n) == BM25Index(docs)._bm25_rank("alpha")[:_RERANK_POOL] def test_empty_query_skips_reranker(self): calls = {"n": 0} diff --git a/tests/test_mcp_hot_reload.py b/tests/test_mcp_hot_reload.py index feee0c79..8a654900 100644 --- a/tests/test_mcp_hot_reload.py +++ b/tests/test_mcp_hot_reload.py @@ -265,13 +265,17 @@ class TestErrorTracking: assert status["error"] == "" def test_error_cleared_on_remove(self) -> None: - """remove_server_sync cleans up _last_error entry.""" + """remove_server_sync cleans up _last_error entry and sweeps EVERY + user's discovery record for the removed name.""" mgr = MCPClientManager({"test": {"command": "echo"}}) mgr._last_error["test"] = "Connection refused" - mgr._pool_discovery_error["test"] = "stale discovery failure" + mgr._pool_discovery_error[("u1", "test")] = "stale discovery failure" + mgr._pool_discovery_error[("u2", "test")] = "stale discovery failure" + mgr._pool_discovery_error[("u1", "other")] = "unrelated server" mgr.remove_server_sync("test") assert "test" not in mgr._last_error - assert "test" not in mgr._pool_discovery_error + assert not any(sname == "test" for _uid, sname in mgr._pool_discovery_error) + assert mgr._pool_discovery_error[("u1", "other")] == "unrelated server" def test_all_server_status_includes_errors(self) -> None: """get_all_server_status propagates per-server errors.""" @@ -435,13 +439,15 @@ class TestReconcileSync: 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.""" + must clear discovery state — every user's record — 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._pool_discovery_error[("u1", "pool-srv")] = "old endpoint failed" + mgr._pool_discovery_error[("u2", "pool-srv")] = "old endpoint failed" mgr.reconcile_sync(_FakeStorage([])) - assert "pool-srv" not in mgr._pool_discovery_error + assert mgr._pool_discovery_error == {} row = _db_row( "pool-srv", @@ -451,7 +457,7 @@ class TestReconcileSync: ) row["auth_type"] = "oauth_user" mgr.reconcile_sync(_FakeStorage([row])) - assert mgr.get_server_status("pool-srv")["discovery_error"] == "" + assert mgr.get_server_status("pool-srv", "u1")["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 @@ -459,7 +465,7 @@ class TestReconcileSync: miss the flip.""" mgr = MCPClientManager({}) mgr._oauth_user_server_names = {"srv"} # previously oauth_user - mgr._pool_discovery_error["srv"] = "failure from old auth model" + mgr._pool_discovery_error[("u1", "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") @@ -469,7 +475,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 + assert 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 8928a29e..b69cbb0e 100644 --- a/tests/test_mcp_user_pool.py +++ b/tests/test_mcp_user_pool.py @@ -20,7 +20,7 @@ import threading import time from datetime import UTC, datetime, timedelta from types import SimpleNamespace -from typing import Any +from typing import TYPE_CHECKING, Any from unittest.mock import AsyncMock, MagicMock, patch import mcp.types as mcp_types @@ -36,6 +36,9 @@ from turnstone.core.mcp_client import MCPClientManager, PoolEntryState from turnstone.core.mcp_crypto import MCPTokenStore from turnstone.core.storage._sqlite import SQLiteBackend +if TYPE_CHECKING: + from collections.abc import Iterator + # --------------------------------------------------------------------------- # Fixtures and helpers # --------------------------------------------------------------------------- @@ -170,6 +173,44 @@ def _fake_pool_tools(server_name: str, tool_name: str) -> list[dict[str, Any]]: ] +def _run_failing_prime( + mgr: MCPClientManager, + loop: asyncio.AbstractEventLoop, + storage: SQLiteBackend, + exc: Exception, +) -> None: + """Register ``pool-srv`` (oauth_user) and run one failing prime for + ``user-1``: the token lookup succeeds, ``_prime_user_server`` raises + *exc*. + + Single copy of the prime-failure arrange (storage row, app_state, + registry, patches) — tests vary only the exception and their asserts, + so a prime-path wiring change lands in one place instead of drifting + across per-test copies. + """ + cipher = make_mcp_token_cipher() + mgr.set_storage(storage) + mgr.set_app_state(_make_app_state(storage, cipher=cipher)) + 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 exc + + 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")) + + # --------------------------------------------------------------------------- # Pool data structures # --------------------------------------------------------------------------- @@ -223,55 +264,103 @@ class _AsyncCM: return False +def _fake_connect_session() -> MagicMock: + """Session double for a successful ``_connect_one_pool``. + + ``_connect_one_pool`` discovers tools, resources, and prompts after + ``initialize()`` returns (resources/prompts capability-gated). The + capability stub advertises tools only so connect-path tests keep their + narrow focus; resources/prompts paths are exercised by the real-transport + tests in ``tests/test_mcp_user_catalog.py``. Single copy — divergent + per-test doubles would silently stop matching the connect path's + discovery sequence as it evolves. + """ + fake_session = MagicMock() + fake_session.initialize = AsyncMock(return_value=None) + fake_caps = MagicMock() + fake_caps.resources = None + fake_caps.prompts = None + fake_session.get_server_capabilities = MagicMock(return_value=fake_caps) + fake_session.list_tools = AsyncMock(return_value=MagicMock(tools=[])) + return fake_session + + +@contextlib.contextmanager +def _patched_pool_transport( + mgr: MCPClientManager, + fake_session: MagicMock, + observed_kwargs: dict[str, Any] | None = None, +) -> Iterator[None]: + """Patch the streamable-http transport under ``_connect_one_pool`` so a + connect succeeds against *fake_session* without a network. Records the + stream factory's url/headers into *observed_kwargs* when given. + """ + + def _stream_factory(*, url: str, headers: dict[str, str]) -> _AsyncCM: + if observed_kwargs is not None: + observed_kwargs["url"] = url + observed_kwargs["headers"] = dict(headers) + return _AsyncCM((AsyncMock(), AsyncMock(), lambda: None)) + + async def _probe(*_args: Any, **_kwargs: Any) -> None: + return None + + with ( + patch("turnstone.core.mcp_client.streamablehttp_client", side_effect=_stream_factory), + patch.object(mgr, "_tcp_probe", side_effect=_probe), + patch("turnstone.core.mcp_client.ClientSession", return_value=_AsyncCM(fake_session)), + ): + yield + + +_POOL_CONNECT_CFG = { + "type": "streamable-http", + "url": "https://mcp.example.com/sse", + "headers": {}, +} + + class TestLazyConnect: def test_connect_pool_injects_authorization_header(self, running_loop_mgr) -> None: - from unittest.mock import patch - mgr, loop, _ = running_loop_mgr observed_kwargs: dict[str, Any] = {} - - async def _probe(*_args: Any, **_kwargs: Any) -> None: - return None - - fake_session = MagicMock() - fake_session.initialize = AsyncMock(return_value=None) - # Phase 7b: ``_connect_one_pool`` discovers tools, resources, - # and prompts after ``initialize()`` returns (resources/prompts - # capability-gated). The capability stub returns a tools-only - # advertisement so the test can keep its narrow focus on the - # bearer-injection contract; resources/prompts paths are - # exercised by the real-transport tests in - # ``tests/test_mcp_user_catalog.py``. - fake_caps = MagicMock() - fake_caps.resources = None - fake_caps.prompts = None - fake_session.get_server_capabilities = MagicMock(return_value=fake_caps) - fake_session.list_tools = AsyncMock(return_value=MagicMock(tools=[])) - - def _stream_factory(*, url: str, headers: dict[str, str]) -> _AsyncCM: - observed_kwargs["url"] = url - observed_kwargs["headers"] = dict(headers) - return _AsyncCM((AsyncMock(), AsyncMock(), lambda: None)) - - with ( - patch("turnstone.core.mcp_client.streamablehttp_client", side_effect=_stream_factory), - patch.object(mgr, "_tcp_probe", side_effect=_probe), - patch("turnstone.core.mcp_client.ClientSession", return_value=_AsyncCM(fake_session)), - ): - cfg = { - "type": "streamable-http", - "url": "https://mcp.example.com/sse", - "headers": {}, - } + fake_session = _fake_connect_session() + with _patched_pool_transport(mgr, fake_session, observed_kwargs): entry = _run_on_loop( loop, - mgr._connect_one_pool(("user-1", "pool-srv"), cfg, "access-aaa"), + mgr._connect_one_pool( + ("user-1", "pool-srv"), dict(_POOL_CONNECT_CFG), "access-aaa" + ), ) assert entry.session is fake_session assert observed_kwargs["headers"]["Authorization"] == "Bearer access-aaa" + def test_successful_connect_clears_recorded_discovery_error(self, running_loop_mgr) -> None: + """Heal path, directly: a recorded discovery failure is cleared by + THAT user's next successful connect to THAT server. Another user's + record for the same server survives — their failure is still true — + as does the same user's record for another server.""" + mgr, loop, _ = running_loop_mgr + mgr._pool_discovery_error[("user-1", "pool-srv")] = "TimeoutError: prime failed" + mgr._pool_discovery_error[("user-2", "pool-srv")] = "500 from user-2's prime" + mgr._pool_discovery_error[("user-1", "other-srv")] = "unrelated failure" + + fake_session = _fake_connect_session() + with _patched_pool_transport(mgr, fake_session): + entry = _run_on_loop( + loop, + mgr._connect_one_pool( + ("user-1", "pool-srv"), dict(_POOL_CONNECT_CFG), "access-aaa" + ), + ) + + assert entry.session is fake_session + assert ("user-1", "pool-srv") not in mgr._pool_discovery_error + assert mgr._pool_discovery_error[("user-2", "pool-srv")] == "500 from user-2's prime" + assert mgr._pool_discovery_error[("user-1", "other-srv")] == "unrelated failure" + def test_connect_pool_rejects_non_http_transport(self, running_loop_mgr) -> None: mgr, loop, _ = running_loop_mgr cfg = {"type": "stdio", "command": "echo"} @@ -309,6 +398,16 @@ class TestEviction: def test_idle_eviction_closes_stale_entries(self, running_loop_mgr) -> None: mgr, loop, _ = running_loop_mgr mgr._user_pool_idle_ttl_s = 0.0 # everything is stale + # Records are reaped by the tick's orphan sweep alone, never by the + # entry drop itself: u0's record (no listener) survives the tick + # that evicts its entry and goes on the NEXT tick. u9's record has + # no entry but a LIVE listener, so no tick may touch it — and a + # name-wide reap is _drop_pool_discovery_errors' job alone. Full + # lifecycle in + # test_stub_eviction_keeps_discovery_record_while_session_lives. + mgr._pool_discovery_error[("u0", "pool-srv")] = "stale failure" + mgr._pool_discovery_error[("u9", "pool-srv")] = "no entry behind this" + mgr.add_listener(lambda: None, user_id="u9") async def _seed() -> list[PoolEntryState]: entries = [] @@ -325,6 +424,10 @@ class TestEviction: _run_on_loop(loop, _evict()) assert mgr._user_pool_entries == {} + assert mgr._pool_discovery_error[("u0", "pool-srv")] == "stale failure" + _run_on_loop(loop, _evict()) + assert ("u0", "pool-srv") not in mgr._pool_discovery_error + assert mgr._pool_discovery_error[("u9", "pool-srv")] == "no entry behind this" def test_eviction_skips_locked_entries(self, running_loop_mgr) -> None: mgr, loop, _ = running_loop_mgr @@ -1054,6 +1157,80 @@ class TestEviction: _run_on_loop(loop, mgr._evict_idle_pool_entries()) assert ("u-stub", "pool-srv") not in mgr._user_pool_entries + def test_stub_eviction_keeps_discovery_record_while_session_lives( + self, running_loop_mgr + ) -> None: + """Records are reaped ONLY by the tick's orphan sweep — an + entry-less key whose user has no live listener. A failed prime's + stub (``last_used=0.0``) full-drops on the FIRST tick, but a live + user's record must outlive it: the outage advisory has no + mid-session re-record path. A departed user's record goes one tick + after their entry does; removing the last listener retires a kept + record on the next tick — even with the entry map empty.""" + mgr, loop, _ = running_loop_mgr + mgr._user_pool_idle_ttl_s = 0.0 + mgr._oauth_user_server_names = {"pool-srv"} + mgr._pool_discovery_error[("u-live", "pool-srv")] = "TimeoutError: prime failed" + mgr._pool_discovery_error[("u-gone", "pool-srv")] = "stale failure" + + async def _seed() -> None: + await mgr._ensure_pool_entry(("u-live", "pool-srv")) + await mgr._ensure_pool_entry(("u-gone", "pool-srv")) + + _run_on_loop(loop, _seed()) + + def _cb() -> None: + return None + + mgr.add_listener(_cb, user_id="u-live") + + # Tick 1: both stubs full-drop; both records survive (u-gone's key + # still had its entry when the head-of-tick sweep ran). + _run_on_loop(loop, mgr._evict_idle_pool_entries()) + assert mgr._user_pool_entries == {} + assert mgr._pool_discovery_error[("u-live", "pool-srv")] == "TimeoutError: prime failed" + assert mgr._pool_discovery_error[("u-gone", "pool-srv")] == "stale failure" + + # Tick 2: the sweep reaps the departed user's entry-less record; + # the live user's survives. + _run_on_loop(loop, mgr._evict_idle_pool_entries()) + assert ("u-gone", "pool-srv") not in mgr._pool_discovery_error + assert mgr._pool_discovery_error[("u-live", "pool-srv")] == "TimeoutError: prime failed" + + # The user departs: the next tick reaps the kept record — and must + # do so despite the empty entry map. + mgr.remove_listener(_cb, user_id="u-live") + _run_on_loop(loop, mgr._evict_idle_pool_entries()) + assert mgr._pool_discovery_error == {} + + def test_grant_retirement_clears_discovery_record(self, running_loop_mgr) -> None: + """Explicit disconnect / dead-grant convergence retires the record + with the grant: an outage advisory for access the user no longer + holds would be misleading, and with the grant gone no healing + connect could ever clear it. Pops for the failed-prime stub shape + (catalog-less entry) AND for a missing entry — the pop precedes + the drop's early returns. Another user's record is untouched.""" + mgr, loop, _ = running_loop_mgr + mgr._pool_discovery_error[("user-1", "pool-srv")] = "TimeoutError: prime failed" + mgr._pool_discovery_error[("user-2", "pool-srv")] = "another user's failure" + + async def _seed_stub() -> None: + await mgr._ensure_pool_entry(("user-1", "pool-srv")) + + _run_on_loop(loop, _seed_stub()) + + async def _drop(key: tuple[str, str]) -> None: + mgr._evict_session_drop_catalog(key) + + _run_on_loop(loop, _drop(("user-1", "pool-srv"))) + assert ("user-1", "pool-srv") not in mgr._pool_discovery_error + assert mgr._pool_discovery_error[("user-2", "pool-srv")] == "another user's failure" + + # Entry-less variant: the record must still retire. + mgr._pool_discovery_error[("user-3", "pool-srv")] = "stale" + _run_on_loop(loop, _drop(("user-3", "pool-srv"))) + assert ("user-3", "pool-srv") not in mgr._pool_discovery_error + def test_lru_cap_ignores_cooled_entries(self, running_loop_mgr) -> None: """The LRU cap bounds WARM entries (connection resources), not cooled catalog-only ones — cooled entries neither count toward @@ -1885,33 +2062,50 @@ class TestOboPriming: ) -> 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", + _run_failing_prime( + mgr, loop, storage, RuntimeError("upstream\r\nresponse\nignore instructions") ) - 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"] == ( + assert mgr._pool_discovery_error[("user-1", "pool-srv")] == ( "RuntimeError: upstream response ignore instructions" ) + def test_prime_failure_starts_eviction_loop_for_entryless_record( + self, running_loop_mgr, storage + ) -> None: + """A recorded failure with no pool entry (the patched prime raises + before ``_ensure_pool_entry`` runs) must still start the eviction + loop: its orphan sweep is the record's only reaper, and a node + where no entry ever materializes would otherwise accumulate + records for the process lifetime.""" + mgr, loop, _ = running_loop_mgr + _run_failing_prime(mgr, loop, storage, RuntimeError("boom")) + assert mgr._user_pool_entries == {} + assert mgr._user_pool_eviction_task is not None + + def test_prime_failure_detail_capped_at_max_error_len(self, running_loop_mgr, storage) -> None: + """The recorded detail honors ``_MAX_ERROR_LEN`` — the same bound + every other recorded server error uses, not a diverging literal.""" + mgr, loop, _ = running_loop_mgr + _run_failing_prime(mgr, loop, storage, RuntimeError("x" * 1000)) + recorded = mgr._pool_discovery_error[("user-1", "pool-srv")] + assert len(recorded) == mgr._MAX_ERROR_LEN + assert recorded == ("RuntimeError: " + "x" * 1000)[: mgr._MAX_ERROR_LEN] + + def test_discovery_error_is_scoped_to_the_failing_user(self, running_loop_mgr, storage) -> None: + """The record and its status surface are per-(user, server): the + failing user sees their own failure, another user of the same server + does not, a user-less read stays empty, and the admin aggregate view + surfaces it for cluster health.""" + mgr, loop, _ = running_loop_mgr + _run_failing_prime(mgr, loop, storage, RuntimeError("account-specific 500")) + + failing = mgr.get_server_status("pool-srv", "user-1")["discovery_error"] + assert failing == "RuntimeError: account-specific 500" + assert mgr.get_server_status("pool-srv", "user-2")["discovery_error"] == "" + assert mgr.get_server_status("pool-srv")["discovery_error"] == "" + aggregate = mgr.get_server_status("pool-srv", aggregate=True)["discovery_error"] + assert aggregate == "RuntimeError: account-specific 500" + 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_sessions.py b/tests/test_sessions.py index 4daca094..defe07e3 100644 --- a/tests/test_sessions.py +++ b/tests/test_sessions.py @@ -1287,6 +1287,19 @@ class TestMCPActingUserBinding: # Merged tool list rebuilt under the new identity. mcp_client.get_tools.assert_any_call(user_id="bob") + def test_status_snapshot_follows_acting_user(self, tmp_db, mock_openai_client): + """The tool_search status snapshot scopes to the acting user, like + the get_tools call that builds the search corpus — owner-scoping + would render the owner's per-user pool state (including their + recorded discovery-failure text) into a non-owner participant's + search results.""" + session, mcp_client = self._make(mock_openai_client) + mcp_client.get_all_server_status.return_value = {} + + session.bind_acting_user("bob") + session._mcp_status_snapshot() + assert mcp_client.get_all_server_status.call_args.args == ("bob",) + def test_prepared_item_pins_identity_across_rebind(self, tmp_db, mock_openai_client): session, mcp_client = self._make(mock_openai_client) session.bind_acting_user("bob") diff --git a/tests/test_tool_search.py b/tests/test_tool_search.py index fca7ce38..2e67f12a 100644 --- a/tests/test_tool_search.py +++ b/tests/test_tool_search.py @@ -4,6 +4,7 @@ from __future__ import annotations import pytest +from turnstone.core.bm25 import _RERANK_POOL from turnstone.core.tool_search import ( BM25Index, ToolSearchManager, @@ -298,6 +299,23 @@ class TestToolSearchTruncation: manager.search("mcp") assert manager._last_total_matched == 5 + def test_total_matched_not_floored_by_rerank_pool(self): + # More matches than the rerank recall pool: the pool caps what the + # reranker reorders, not the recorded match count — "top N of M" must + # report the true M, not the pool size. [#941] + n = _RERANK_POOL + 10 + tools = [_make_tool(f"mcp__srv__tool{i}", f"alpha capability {i}") for i in range(n)] + mgr = ToolSearchManager( + tools, + always_on_names=set(), + max_results=3, + reranker=lambda q, d: list(range(len(d)))[::-1], + ) + results = mgr.search("alpha") + assert len(results) == 3 + assert mgr._last_total_matched == n + assert f"Showing the top 3 of {n}" in mgr.format_search_results(results) + class TestToolSearchUnavailableAdvisory: """A down/unauthorized server is surfaced, never silently treated as @@ -358,6 +376,17 @@ class TestToolSearchUnavailableAdvisory: text = mgr.format_search_results(mgr.search("github")) assert "unavailable" not in text.lower() + def test_discovery_error_not_flagged_for_user_with_warm_pool(self): + # A discovery record alongside connected=True in the per-user status + # snapshot is stale (the user's own successful connect clears their + # record), so the outage advisory must stay silent. [#941] + mgr = self._mgr( + [_make_tool("mcp__github__create_issue", "Create a github issue")], + {"github": {"connected": True, "discovery_error": "500 app failed"}}, + ) + 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")], @@ -395,6 +424,30 @@ class TestStatusReason: reason = _status_reason({"discovery_error": "TimeoutError: pool discovery"}) assert "discovery" in reason.lower() + def test_discovery_error_suppressed_when_connected(self): + # Belt-and-braces (#941): a successful connect clears the user's + # per-(user, server) record, so a record alongside a warm transport is + # stale — flagging a server the user's pool is serving would be false. + assert _status_reason({"connected": True, "discovery_error": "500 boom"}) == "" + + def test_discovery_error_fires_despite_retained_catalog(self): + # The record is per-(user, server), so a failure in this user's status + # IS their own — a retained idle catalog (#836, connected=False with a + # non-zero tools count) must not mask it: the transport behind those + # tools is failing and calls will too. + reason = _status_reason({"connected": False, "tools": 3, "discovery_error": "500 boom"}) + assert reason == "tool discovery failed: 500 boom" + + def test_discovery_error_fires_for_cold_pool(self): + reason = _status_reason({"connected": False, "tools": 0, "discovery_error": "500 boom"}) + assert reason == "tool discovery failed: 500 boom" + + def test_recorded_error_not_gated_by_connected(self): + # Deliberate asymmetry: only the discovery branch is per-user gated. A + # recorded error stays a hard signal even alongside connected=True + # (e.g. a flapping static server). + assert "boom" in _status_reason({"connected": True, "error": "boom"}) + @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"}) diff --git a/turnstone/core/bm25.py b/turnstone/core/bm25.py index 7b8250b5..fbc18060 100644 --- a/turnstone/core/bm25.py +++ b/turnstone/core/bm25.py @@ -71,7 +71,9 @@ class BM25Index: """Return indices of top-k documents by relevance. Stage 1 is BM25. When a reranker is attached, a recall pool of the top - ``_RERANK_POOL`` BM25 hits is reranked and the top-k spliced back. + ``_RERANK_POOL`` BM25 hits is reranked and the top-k spliced back. In + reorder mode, matches ranked past the pool trail in BM25 order rather + than being dropped, so a ``k`` larger than the pool sees every match. """ ranked = self._bm25_rank(query) if self._reranker is None: @@ -83,7 +85,12 @@ class BM25Index: try: order = list(self._reranker(query, docs)) except Exception: - return pool[:k] # reranker ERROR -> BM25 order (both modes) + # Reranker ERROR -> BM25 order (both modes). Filter mode keeps its + # pool bound — its happy path can never exceed the pool, and an + # endpoint failure must not return a longer list than a working + # endpoint ever could; reorder mode gets the full ranking, like + # its past-pool tail below. + return pool[:k] if self._rerank_filters else ranked[:k] seen: set[int] = set() out: list[int] = [] for pos in order: @@ -110,12 +117,21 @@ class BM25Index: # any pool items the reranker omitted (e.g. a top_n subset) are backfilled # in BM25 order so results are never silently lost. if not out: - return pool[:k] + return ranked[:k] for pos in range(len(pool)): if len(out) >= k: break if pos not in seen: out.append(pool[pos]) + # Matches ranked past the recall pool. A no-op when k <= _RERANK_POOL + # (the skills-find path passes k = min(len(rows), 50)); tool_search + # passes k = corpus size and DEPENDS on this tail — it counts + # matches from the result, and a pool-capped result would floor + # that count at the pool size. + for idx in ranked[len(pool) :]: + if len(out) >= k: + break + out.append(idx) return out def _score(self, q_tokens: list[str], doc_tokens: list[str], dl: int) -> float: diff --git a/turnstone/core/mcp_client.py b/turnstone/core/mcp_client.py index ae90b386..52bec425 100644 --- a/turnstone/core/mcp_client.py +++ b/turnstone/core/mcp_client.py @@ -677,13 +677,18 @@ 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] = {} + # Last tool-DISCOVERY error for pool (oauth_user / oauth_obo) servers, + # keyed (user_id, server_name) like the pool entries it describes: set + # when a user's prime/connect raises during discovery (transport, 5xx, + # timeout), cleared when THAT user's pool connects. Surfaced per-user + # 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 — and visible + # only to the user it happened to: the recorded exception text is + # per-user data (another account's failure must never render in this + # user's results), and another user's success must not erase a signal + # that is still true for this one. + self._pool_discovery_error: dict[tuple[str, str], str] = {} self._MAX_ERROR_LEN = 256 # Listener infrastructure (tool-change callbacks for ChatSession). @@ -976,13 +981,23 @@ class MCPClientManager: self._user_pool_locks[key] = lock entry = PoolEntryState(key=key, open_lock=lock) self._user_pool_entries[key] = entry - # First pool entry — start the idle-eviction loop. Deferred - # until the first pool key materializes so static-only - # deployments never spawn the task. - if self._user_pool_eviction_task is None: - self._user_pool_eviction_task = asyncio.create_task(self._user_pool_eviction_loop()) + self._ensure_eviction_loop() return entry + def _ensure_eviction_loop(self) -> None: + """Start the idle-eviction loop if it isn't running. Loop-only. + + Deferred until the first pool entry OR discovery record + materializes, so static-only deployments never spawn the task. + The record trigger matters: the eviction tick's orphan sweep is + the only reaper for entry-less discovery records, and a prime + can fail (and record) before ``_ensure_pool_entry`` ever runs — + on a node where no pool entry materializes, records would + otherwise accumulate for the process lifetime with no reaper. + """ + if self._user_pool_eviction_task is None: + self._user_pool_eviction_task = asyncio.create_task(self._user_pool_eviction_loop()) + def set_app_state(self, app_state: Any) -> None: """Wire the OAuth ``app.state`` into the manager. @@ -2923,9 +2938,11 @@ 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) + # Discovery just succeeded for THIS user's pool — clear their recorded + # failure so a healed outage doesn't linger on status. Other users' + # records for the server stay: this connect proves nothing about a + # failure another account experienced. + self._pool_discovery_error.pop(key, None) return entry # -- pool priming --------------------------------------------------------- @@ -3159,18 +3176,17 @@ class MCPClientManager: ) 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. + # THIS user's server status (and their tool_search + # 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[key] = self._sanitize_error_detail( + f"{type(exc).__name__}: {exc}" ) - self._pool_discovery_error[server_name] = detail[:200] + self._ensure_eviction_loop() log.debug( "mcp pool auto-prime failed user=%s server=%s", user_id, @@ -3658,12 +3674,29 @@ class MCPClientManager: those; they are bounded by live-session users × pool servers and reaped by the TTL pass within a tick of their user's last listener going away. + + Also the SINGLE reaper for discovery records: the full-drop path + deliberately leaves them in place (a live session must keep its + outage advisory past its stub's eviction), so this sweep pops any + record whose key has no backing entry once the user's last listener + is gone — within a tick of departure. It must run BEFORE the + empty-map early return: the orphan case has no entries by + definition, and a prime can fail (and record) without ever + allocating an entry. """ + live_uids = self._live_listener_uids() + # list(): sync-thread clears (_drop_pool_discovery_errors) mutate + # the record map concurrently with this loop-side iteration. + for key in [ + k + for k in list(self._pool_discovery_error) + if k not in self._user_pool_entries and k[0] not in live_uids + ]: + self._pool_discovery_error.pop(key, None) if not self._user_pool_entries: return now = time.monotonic() ttl = self._user_pool_idle_ttl_s - live_uids = self._live_listener_uids() # First pass: TTL-based eviction. Run closes in parallel so a tick # that needs to evict many entries doesn't block on serial teardowns. @@ -3776,6 +3809,14 @@ class MCPClientManager: had_catalog = self._entry_has_catalog(entry) self._user_pool_entries.pop(key, None) self._user_pool_last_used.pop(key, None) + # Deliberately NOT popping the discovery record here: it must + # outlive the entry exactly as long as a live session could + # still consult it (a failed prime leaves a catalog-less stub + # with last_used=0.0, so the FIRST tick full-drops it here, and + # there is no mid-session re-record path). The eviction tick's + # orphan sweep is the single reaper: once the key is entry-less + # AND the user's last listener is gone, the record goes — one + # policy site instead of a second liveness predicate here. if had_catalog: # Dropping the entry without rebuilding the per-user # catalogs would leave ``is_mcp_tool`` / per-user @@ -5687,7 +5728,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._drop_pool_discovery_errors(name) self._clear_static_push_state(name, markers=True) self._cb_clear(name) # Clear health-loop backoff/ping state so a later @@ -5730,7 +5771,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._drop_pool_discovery_errors(name) self._clear_static_push_state(name, markers=True) self._cb_clear(name) self._rebuild_tools() @@ -5746,10 +5787,53 @@ class MCPClientManager: log.info("Removed MCP server '%s'", name) return was_connected + def _sanitize_error_detail(self, msg: str) -> str: + """Single-line, bounded rendering for recorded server-error text. + + ONE pipeline for every recorded surface (``_last_error``, the pool + discovery record) so a future sanitizer hardening cannot land on + one and silently miss the other. + """ + return msg.replace("\n", " ").replace("\r", "").strip()[: self._MAX_ERROR_LEN] + def _set_error(self, name: str, msg: str) -> None: """Store a sanitized error string for a server.""" - clean = msg.replace("\n", " ").replace("\r", "") - self._last_error[name] = clean[: self._MAX_ERROR_LEN] + self._last_error[name] = self._sanitize_error_detail(msg) + + def _drop_pool_discovery_errors(self, server_name: str) -> None: + """Drop every user's recorded discovery failure for *server_name*. + + Registration-lifecycle clear (removal, reconcile-observed pool removal, + or a pool auth-type flip): the records describe a registration that no + longer exists, so a later same-name registration must not inherit them + — for ANY user. ``list()`` first: the keys copy is a single C-level + call, so a sync-thread caller (remove / reconcile) can't trip over the + mcp-loop's concurrent writes mid-iteration; the comprehension then + filters a private snapshot. + """ + for key in [k for k in list(self._pool_discovery_error) if k[1] == server_name]: + self._pool_discovery_error.pop(key, None) + + def _pool_discovery_error_for( + self, server_name: str, user_id: str | None, *, aggregate: bool = False + ) -> str: + """Recorded discovery failure, scoped like pool status itself: the + requesting user's own record, or any user's under the admin + ``aggregate`` view. No user context → ``""`` — a failure another + account experienced is per-user data, the same rule ``connected`` + follows in :meth:`_oauth_user_server_status`. + """ + if aggregate: + # list(): sync-thread status reads race mcp-loop writes; same + # snapshot discipline as the pool-entries scan in + # _oauth_user_server_status. + for (_uid, sname), detail in list(self._pool_discovery_error.items()): + if sname == server_name: + return detail + return "" + if not user_id: + return "" + return self._pool_discovery_error.get((user_id, server_name), "") def get_server_status( self, name: str, user_id: str | None = None, *, aggregate: bool = False @@ -5786,7 +5870,9 @@ 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, ""), + # Shape parity with the pool branch; only pool primes record + # discovery failures, so a static server has none by construction. + "discovery_error": "", "transport": transport, "command": cfg.get("command", "") if transport == "stdio" else "", "url": cfg.get("url", "") if transport != "stdio" else "", @@ -5852,7 +5938,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, ""), + "discovery_error": self._pool_discovery_error_for(name, user_id, aggregate=aggregate), "transport": "streamable-http", "command": "", "url": "", @@ -5933,7 +6019,7 @@ class MCPClientManager: # 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) + self._drop_pool_discovery_errors(name) # Pool servers newly registered OR migrated between pool auth types # since active sessions last primed. newly_added_pool = { @@ -7898,6 +7984,14 @@ class MCPClientManager: time. """ self._evict_session(key) + # The grant for this (user, server) is GONE — an outage advisory + # for access the user no longer holds would be misleading, and with + # the grant gone no healing connect can ever clear the record (prime + # returns at the token lookup; zero catalog tools means no lazy + # dispatch). The not-consented state has its own rails; the record + # retires with the grant. Before the early returns below: the + # record's stub is typically already catalog-less or dropped. + self._pool_discovery_error.pop(key, None) evict = self._user_pool_entries.get(key) if evict is None: return diff --git a/turnstone/core/session.py b/turnstone/core/session.py index 6d192788..6be5bb2a 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -3021,6 +3021,12 @@ class ChatSession: """Per-server MCP status for this session's user, consumed by ``ToolSearchManager`` to flag unavailable servers in search results. + Scoped to the EFFECTIVE user — the acting participant on a shared + workstream — matching the ``get_tools`` call that builds the search + corpus. Owner-scoping here would render the owner's per-user pool + state (including their recorded discovery-failure text) into a + non-owner participant's search results. + Returns ``{}`` when no MCP client is bound or the lookup fails — the advisory then simply stays silent rather than breaking tool search. """ @@ -3028,7 +3034,7 @@ class ChatSession: if client is None: return {} try: - return client.get_all_server_status(self._mcp_user_id) + return client.get_all_server_status(self._mcp_effective_user_id) except Exception: return {} diff --git a/turnstone/core/tool_search.py b/turnstone/core/tool_search.py index 8121af30..53ae56ca 100644 --- a/turnstone/core/tool_search.py +++ b/turnstone/core/tool_search.py @@ -66,6 +66,16 @@ def _status_reason(status: dict[str, Any]) -> str: 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. + + A recorded discovery failure is per-user state: the client records it + under the (user, server) pool key and the status snapshot is scoped to + the requesting user (the session wires ``get_all_server_status(user_id)``), + so a non-empty ``discovery_error`` here means THIS user's own discovery + failed — never another account's. ``connected`` still gates the branch as + belt-and-braces: a successful connect clears the user's record, so a warm + transport alongside a record is a stale signal, and flagging a server the + user's pool is actively serving would be false. A recorded ``error`` + stays ungated — it is a hard signal even while connected. """ if status.get("circuit_open"): return "circuit breaker open" @@ -75,7 +85,7 @@ def _status_reason(status: dict[str, Any]) -> str: if err: return f"error: {err[:120]}" disc = str(status.get("discovery_error") or "").replace("\n", " ").replace("\r", "").strip() - if disc: + if disc and not status.get("connected"): return f"tool discovery failed: {disc[:120]}" return ""