diff --git a/tests/test_mcp_hot_reload.py b/tests/test_mcp_hot_reload.py index e151adc1..f0b264fd 100644 --- a/tests/test_mcp_hot_reload.py +++ b/tests/test_mcp_hot_reload.py @@ -430,3 +430,38 @@ class TestReconcileSync: row["auth_type"] = "oauth_obo" mgr.reconcile_sync(_FakeStorage([row])) assert primed == [] + + 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 + primed: list[str] = [] + mgr.prime_user_pools = lambda uid: primed.append(uid) # type: ignore[method-assign] + mgr.add_listener(lambda: None, user_id="u1") + row = _db_row("srv", transport="streamable-http", command="", url="https://srv:8443/") + row["auth_type"] = "oauth_obo" # flipped in place + mgr.reconcile_sync(_FakeStorage([row])) + assert primed == ["u1"] + assert mgr._obo_server_names == {"srv"} + assert mgr._oauth_user_server_names == set() + + def test_reprime_survives_prime_exception(self) -> None: + """One user's prime scheduling failure must not abort the loop or propagate + out of reconcile_sync (which would 500 the reload endpoint).""" + mgr = MCPClientManager({}) + primed: list[str] = [] + + def _prime(uid: str) -> None: + if uid == "boom-user": + raise RuntimeError("scheduling blew up") + primed.append(uid) + + mgr.prime_user_pools = _prime # type: ignore[method-assign] + mgr.add_listener(lambda: None, user_id="boom-user") + mgr.add_listener(lambda: None, user_id="ok-user") + row = _db_row("azobo", transport="streamable-http", command="", url="https://azobo:8443/") + row["auth_type"] = "oauth_obo" + mgr.reconcile_sync(_FakeStorage([row])) # must not raise + assert "ok-user" in primed # the other user was still primed diff --git a/turnstone/core/mcp_client.py b/turnstone/core/mcp_client.py index 64dec043..dbdac57f 100644 --- a/turnstone/core/mcp_client.py +++ b/turnstone/core/mcp_client.py @@ -4706,11 +4706,14 @@ class MCPClientManager: log.warning("reconcile_sync: failed to read mcp_servers table", exc_info=True) return {"added": [], "removed": [], "updated": []} - # Pool-backed (oauth_user / oauth_obo) names known BEFORE this - # reconcile — diffed against the refreshed sets below to spot servers - # that appeared since active sessions last primed (see the re-prime - # self-heal near the end of this method). - prev_pool_names = self._oauth_user_server_names | self._obo_server_names + # Prior (name -> pool auth_type) view, reconstructed from the tracked + # pool-name sets. Diffing auth_type — not just names — below catches a + # server MIGRATED in place between the two pool auth types + # (oauth_user <-> oauth_obo, the flip the OBO feature enables): a + # name-only diff sees the same name on both sides and misses it. + # Mirrors the explicit two-type split of the set rebuilds just below. + prev_pool_auth = {n: "oauth_user" for n in self._oauth_user_server_names} + prev_pool_auth.update(dict.fromkeys(self._obo_server_names, "oauth_obo")) # Refresh the in-memory oauth_user name cache from the rows we # just read — feeds :meth:`server_auth_type` so callers (e.g. @@ -4722,9 +4725,13 @@ class MCPClientManager: self._obo_server_names = { row["name"] for row in rows if row.get("auth_type") == "oauth_obo" } - newly_added_pool = ( - self._oauth_user_server_names | self._obo_server_names - ) - prev_pool_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 servers newly registered OR migrated between pool auth types + # since active sessions last primed. + newly_added_pool = { + name for name, at in new_pool_auth.items() if prev_pool_auth.get(name) != at + } desired = _db_servers_to_config(rows) desired_names = set(desired) @@ -4785,31 +4792,48 @@ class MCPClientManager: # active session's user so a mid-session registration surfaces its # tools automatically — no reconnect or fresh workstream needed. if newly_added_pool: - self._reprime_active_users(newly_added_pool) + self._reprime_active_users(len(newly_added_pool)) return {"added": added, "removed": removed, "updated": updated} - def _reprime_active_users(self, new_servers: set[str]) -> None: - """Re-warm active sessions' pools after new pool-backed servers appear. + def _reprime_active_users(self, changed_count: int) -> None: + """Schedule a pool re-warm for every active session's user after + *changed_count* pool-backed servers were newly registered or migrated + between pool auth types this reconcile (see :meth:`reconcile_sync`). - Called from :meth:`reconcile_sync` when a reconcile reveals an - oauth_user / oauth_obo server that active sessions never primed. Active - users come from the tool-listener registry (each open ChatSession + Active users come from the tool-listener registry (each open ChatSession registers ``(user_id, callback)`` via :meth:`add_listener`); the ``user_id=None`` global/admin listener is skipped. ``prime_user_pools`` - is idempotent and fire-and-forget, so re-priming an already-warm pool - is a cheap skip and re-priming a user without the stored token / - captured credential is a no-op. + is fire-and-forget — it SCHEDULES the mint/connect onto the mcp-loop and + returns, no-opping for an already-warm pool, a user without a captured + credential, or a down loop — so the log below reports what was + SCHEDULED, not what completed. The call is guarded per-user (mirroring + the session.py call sites) so one user's scheduling failure can't abort + the loop or 500 the reload endpoint. + + Fan-out note: only the infrequent, operator-driven reload path reaches + here and each user's mint concurrency is already capped + (``_PRIME_MAX_CONCURRENCY``); a shared cross-user mint cap is left as a + follow-up if IdP rate-limiting is ever observed under a large active-user + count. """ with self._listeners_lock: user_ids = {uid for uid, _cb in self._listeners if uid} for user_id in user_ids: - self.prime_user_pools(user_id) + try: + self.prime_user_pools(user_id) + except Exception: + log.debug( + "mcp reconcile re-prime scheduling failed user=%s", + user_id, + exc_info=True, + ) if user_ids: log.info( - "MCP reconcile: re-primed %d active user(s) after %d new pool server(s) appeared", + "MCP reconcile: scheduled pool re-prime for %d active session " + "user(s) after %d pool server change(s)", len(user_ids), - len(new_servers), + changed_count, ) # -- query methods -------------------------------------------------------