ewilliams28/turnstone
1
0
Fork 0
mirror of https://github.com/turnstonelabs/turnstone.git synced 2026-08-12 23:12:23 -06:00
Code Issues Packages Projects Releases Wiki Activity

fix(tool-search): discovery-failure records are per-user, rerank counts honest

Browse Source 
Follow-up to #938; closes #941. The unavailable-server advisory fired for
users whose own pool was warm: _pool_discovery_error was keyed by server
name while pool connections are per-(user, server), so one account's
failed prime rendered its exception text into every user's search results.

- mcp_client: re-key _pool_discovery_error to (user_id, server_name).
  Written by the failing user's prime (single sanitize-and-cap pipeline
  shared with _set_error), cleared by that user's successful connect,
  retired with the grant on explicit disconnect / dead-grant convergence,
  and swept name-wide on registration lifecycle (removal, reconcile
  auth-type flips) via a snapshot-safe helper. Departed users' records
  are reaped by the eviction tick's orphan sweep — the single tick-side
  reaper; a live user's record survives its stub's eviction because the
  advisory has no mid-session re-record path. The eviction loop also
  starts on record write, so records written before any pool entry
  exists cannot outlive their users. Status reads scope to the
  requesting user, with an any-user view under the admin aggregate flag.
- tool_search: _status_reason treats discovery_error as an outage only
  when the requesting user's own status is not connected — with per-user
  records this is belt-and-braces, since a successful connect clears the
  user's record.
- session: the tool-search status snapshot scopes to the EFFECTIVE user
  (the acting participant on shared workstreams), matching the get_tools
  call that builds the search corpus, so an owner's pool state never
  renders into a non-owner's results.
- bm25: with a reranker attached, matches ranked past the recall pool
  trail in BM25 order (reorder mode), so tool_search's "top N of M"
  count no longer floors at the pool size; the exception fallback is
  mode-aware (filter mode keeps its pool bound, byte-for-byte).
This commit is contained in:
Patrick Buckley
2026-08-02 01:43:20 -07:00
parent c4b2dd7135
commit e526df95d0
9 changed files with 555 additions and 113 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_bm25.py
+57 -7
View File
@@ -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}
tests/test_mcp_hot_reload.py
+15 -9
View File
@@ -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
tests/test_mcp_user_pool.py
+255 -61
View File
@@ -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
tests/test_sessions.py
+13
View File
@@ -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")
tests/test_tool_search.py
+53
View File
@@ -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"})
turnstone/core/bm25.py
+19 -3
View File
@@ -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:
turnstone/core/mcp_client.py
+125 -31
View File
@@ -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
turnstone/core/session.py
+7 -1
View File
@@ -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 {}
turnstone/core/tool_search.py
+11 -1
View File
@@ -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 ""