diff --git a/studio/backend/core/inference/mcp_client.py b/studio/backend/core/inference/mcp_client.py index a0d79bbdf4..5a36d90c5d 100644 --- a/studio/backend/core/inference/mcp_client.py +++ b/studio/backend/core/inference/mcp_client.py @@ -8,6 +8,7 @@ import json import os import shlex import sys +import time from typing import Any, Optional from loggers import get_logger @@ -16,6 +17,14 @@ logger = get_logger(__name__) MCP_TOOL_PREFIX = "mcp__" +# A failed probe isn't cached (a recovered server must come back), but it's +# recorded so a down server isn't re-probed -- and the chat send re-hung for +# the full timeout -- on every message. Cool off for this long after a failure; +# much longer for OAuth, whose probe can hang up to _OAUTH_PROBE_TIMEOUT, +# so that hang doesn't recur every minute. +FAILED_PROBE_COOLOFF_SECONDS = 60.0 +OAUTH_FAILED_PROBE_COOLOFF_SECONDS = 300.0 + _oauth_token_store = None @@ -227,6 +236,53 @@ async def list_tools_async( return await asyncio.wait_for(_fetch(), timeout = timeout) +# Discovered-tool cache, keyed by MCP server id. get_enabled_mcp_tools() +# probes a server only on a cache miss, keeping MCP discovery off the chat +# send's critical path -- tool schemas are stable within a session. The +# /refresh route warms it; a URL/header/OAuth change or a delete evicts it. +# Successful probes are cached indefinitely. +_tool_cache: dict[str, list[dict]] = {} + +# server_id -> monotonic time before which a failed server must not be +# re-probed (see record_probe_failure). Cleared on a successful probe or +# eviction. +_probe_cooloff_until: dict[str, float] = {} + +# MCP server fields whose change invalidates a server's discovered tools: the +# endpoint/auth used to probe it (url, headers, oauth) or whether it's used at +# all (is_enabled). A rename does not. The update route's eviction and +# get_enabled_mcp_tools' mid-probe guard both key off this so they can't drift. +TOOL_CACHE_INVALIDATING_FIELDS = frozenset({"url", "headers_json", "use_oauth", "is_enabled"}) + + +def get_cached_tools(server_id: str) -> Optional[list[dict]]: + return _tool_cache.get(server_id) + + +def cache_tools(server_id: str, tools: list[dict]) -> None: + _tool_cache[server_id] = tools + _probe_cooloff_until.pop(server_id, None) + + +def record_probe_failure(server_id: str, use_oauth: bool = False) -> None: + cooloff = OAUTH_FAILED_PROBE_COOLOFF_SECONDS if use_oauth else FAILED_PROBE_COOLOFF_SECONDS + _probe_cooloff_until[server_id] = time.monotonic() + cooloff + + +def in_failure_cooloff(server_id: str) -> bool: + return _probe_cooloff_until.get(server_id, 0.0) > time.monotonic() + + +def invalidate_tool_cache(server_id: Optional[str] = None) -> None: + """Evict one server's cached tools, or every entry when server_id is None.""" + if server_id is None: + _tool_cache.clear() + _probe_cooloff_until.clear() + else: + _tool_cache.pop(server_id, None) + _probe_cooloff_until.pop(server_id, None) + + def _flatten_result(result: Any) -> str: parts = [] for block in getattr(result, "content", None) or []: diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index b27fa6ff73..8c45706395 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -24,11 +24,16 @@ import urllib.request from core.inference.mcp_client import ( MCP_TOOL_PREFIX, + TOOL_CACHE_INVALIDATING_FIELDS, + cache_tools, call_tool_sync, + get_cached_tools, + in_failure_cooloff, is_stdio, list_tools_async, parse_server_headers, probe_timeout, + record_probe_failure, stdio_mcp_enabled, ) from storage import mcp_servers_db @@ -634,28 +639,56 @@ async def get_enabled_mcp_tools() -> list[dict]: if not servers: return [] - results = await asyncio.gather( - *( - list_tools_async( - url = s["url"], - headers = parse_server_headers(s), - timeout = probe_timeout(s["url"], bool(s.get("use_oauth"))), - use_oauth = bool(s.get("use_oauth")), - ) - for s in servers - ), - return_exceptions = True, - ) + # Skip servers still in their post-failure cool-off, otherwise a down + # server gets re-probed -- and blocks the send for the full timeout -- on + # every message. + uncached = [ + s for s in servers if get_cached_tools(s["id"]) is None and not in_failure_cooloff(s["id"]) + ] + if uncached: + results = await asyncio.gather( + *( + list_tools_async( + url = s["url"], + headers = parse_server_headers(s), + timeout = probe_timeout(s["url"], bool(s.get("use_oauth"))), + use_oauth = bool(s.get("use_oauth")), + ) + for s in uncached + ), + return_exceptions = True, + ) + # An edit/delete can land while we await a probe (up to 305 s for + # OAuth); its cache eviction is a no-op against an entry we haven't + # written yet. Re-read and drop a result whose server changed or + # was removed mid-probe, else a stale tool list caches indefinitely. + current = {s["id"]: s for s in mcp_servers_db.list_servers()} + for server, payload in zip(uncached, results): + # Guard the failure branch too: a stale failure must not park a + # cool-off on the fresh config, or the server the user just fixed + # is skipped for the whole window. + fresh = current.get(server["id"]) + if fresh is None or any( + fresh.get(k) != server.get(k) for k in TOOL_CACHE_INVALIDATING_FIELDS + ): + continue + if isinstance(payload, BaseException): + logger.warning( + "MCP server '%s' (%s) discovery failed: %s", + server.get("display_name") or server["id"], + server.get("url"), + payload, + ) + # Failures aren't cached, but record one so a down server + # isn't re-probed every send during the cool-off. + record_probe_failure(server["id"], bool(fresh.get("use_oauth"))) + continue + cache_tools(server["id"], payload) specs: list[dict] = [] - for server, payload in zip(servers, results): - if isinstance(payload, BaseException): - logger.warning( - "MCP server '%s' (%s) discovery failed: %s", - server.get("display_name") or server["id"], - server.get("url"), - payload, - ) + for server in servers: + payload = get_cached_tools(server["id"]) + if payload is None: continue specs.extend(_mcp_specs_for_server(server, payload)) return specs diff --git a/studio/backend/routes/mcp_servers.py b/studio/backend/routes/mcp_servers.py index 3001c6b7c9..37d99a222e 100644 --- a/studio/backend/routes/mcp_servers.py +++ b/studio/backend/routes/mcp_servers.py @@ -10,12 +10,16 @@ from fastapi import APIRouter, Depends, HTTPException from auth.authentication import get_current_subject from core.inference.mcp_client import ( + TOOL_CACHE_INVALIDATING_FIELDS, + cache_tools, clear_oauth_tokens_async, + invalidate_tool_cache, is_stdio, list_tools_async, parse_server_headers, parse_stdio_command, probe_timeout, + record_probe_failure, stdio_mcp_enabled, ) from core.inference.mcp_config_import import parse_mcp_config @@ -198,6 +202,11 @@ async def update_mcp_server( ): await clear_oauth_tokens_async(old["url"]) mcp_servers_db.update_server(server_id, changes) + # A new endpoint/auth makes cached tools wrong and disabling makes them + # unreachable, so drop them and let the next send re-probe; a rename + # leaves them valid. + if changes.keys() & TOOL_CACHE_INVALIDATING_FIELDS: + invalidate_tool_cache(server_id) return _row_to_response(mcp_servers_db.get_server(server_id)) @@ -209,6 +218,7 @@ async def delete_mcp_server(server_id: str, current_subject: str = Depends(get_c if old.get("use_oauth"): await clear_oauth_tokens_async(old["url"]) mcp_servers_db.delete_server(server_id) + invalidate_tool_cache(server_id) @router.post("/{server_id}/refresh", response_model = McpServerProbeResult) @@ -238,8 +248,23 @@ async def refresh_mcp_server_tools( error = str(exc), exc_info = True, ) + current = mcp_servers_db.get_server(server_id) + if current is not None and not any( + current.get(k) != server.get(k) for k in TOOL_CACHE_INVALIDATING_FIELDS + ): + # Start the cool-off so the next chat send doesn't immediately re-hang + # on this server's timeout. If the row changed while the probe was + # awaiting, the failure belongs to the old config and must not park + # the newly edited server. + record_probe_failure(server_id, use_oauth) return McpServerProbeResult(ok = False, error = safe_curated_detail(exc)) + # Warm the chat-path cache so the next send skips re-probing. + current = mcp_servers_db.get_server(server_id) + if current is not None and not any( + current.get(k) != server.get(k) for k in TOOL_CACHE_INVALIDATING_FIELDS + ): + cache_tools(server_id, tools) return McpServerProbeResult(ok = True, tool_count = len(tools)) diff --git a/studio/backend/tests/test_mcp_servers.py b/studio/backend/tests/test_mcp_servers.py index ede3cf15d4..90b1ade03c 100644 --- a/studio/backend/tests/test_mcp_servers.py +++ b/studio/backend/tests/test_mcp_servers.py @@ -598,3 +598,614 @@ def test_safetensors_agentic_empty_allowlist_still_means_allow_all(): ) # Empty allow-list = run anything (preserved contract). assert calls == [("python", {"code": "1"})] or len(calls) >= 1 + + +# ── discovery cache ───────────────────────────────────────────────── + + +def _one_tool(name = "echo"): + return [{"name": name, "inputSchema": {"type": "object", "properties": {}}}] + + +def test_get_enabled_mcp_tools_caches_discovery(tmp_path, monkeypatch): + """A second send must serve tools from cache instead of re-probing.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + calls: list[str] = [] + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + calls.append(url) + return _one_tool() + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + first = asyncio.run(tools_mod.get_enabled_mcp_tools()) + second = asyncio.run(tools_mod.get_enabled_mcp_tools()) + + assert len(calls) == 1 # probed once, cache hit on the second send + assert [t["function"]["name"] for t in first] == ["mcp__s1__echo"] + assert first == second + + +def test_get_enabled_mcp_tools_does_not_cache_failures(tmp_path, monkeypatch): + """A failed probe isn't cached: once the cool-off elapses, it's retried.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + attempts = {"n": 0} + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + attempts["n"] += 1 + if attempts["n"] == 1: + raise RuntimeError("server down") + return _one_tool() + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] # failure -> empty + # Expire the cool-off (an until-time in the past) so the server is retried. + mcp_client._probe_cooloff_until["s1"] = 0.0 + second = asyncio.run(tools_mod.get_enabled_mcp_tools()) + assert attempts["n"] == 2 # retried after the cool-off, not cached + assert [t["function"]["name"] for t in second] == ["mcp__s1__echo"] + + +def test_refresh_warms_tool_cache(tmp_path, monkeypatch): + """Clicking Refresh must populate the cache the chat path reads.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + async def fake_refresh( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + return _one_tool() + + monkeypatch.setattr(routes_mcp, "list_tools_async", fake_refresh) + res = asyncio.run(routes_mcp.refresh_mcp_server_tools("s1", current_subject = "u")) + assert res.ok and res.tool_count == 1 + + def boom(*a, **k): + raise AssertionError("chat path re-probed despite a warm cache") + + monkeypatch.setattr(tools_mod, "list_tools_async", boom) + specs = asyncio.run(tools_mod.get_enabled_mcp_tools()) + assert [t["function"]["name"] for t in specs] == ["mcp__s1__echo"] + + +def test_update_url_evicts_tool_cache(tmp_path, monkeypatch): + """Re-pointing the URL must drop the old endpoint's cached tools.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from models.mcp_servers import McpServerUpdate + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": _one_tool("stale")}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True) + + asyncio.run( + routes_mcp.update_mcp_server( + "s1", McpServerUpdate(url = "https://new/mcp"), current_subject = "u" + ) + ) + assert mcp_client.get_cached_tools("s1") is None + + +def test_update_display_name_keeps_tool_cache(tmp_path, monkeypatch): + """A rename touches no endpoint, so the cache must survive it.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from models.mcp_servers import McpServerUpdate + import routes.mcp_servers as routes_mcp + + cached = _one_tool() + monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": cached}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + asyncio.run( + routes_mcp.update_mcp_server("s1", McpServerUpdate(display_name = "B"), current_subject = "u") + ) + assert mcp_client.get_cached_tools("s1") == cached + + +def test_update_disable_evicts_tool_cache(tmp_path, monkeypatch): + """Disabling a server must drop its cached tools, not leave them unread.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from models.mcp_servers import McpServerUpdate + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": _one_tool()}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + asyncio.run( + routes_mcp.update_mcp_server("s1", McpServerUpdate(is_enabled = False), current_subject = "u") + ) + assert mcp_client.get_cached_tools("s1") is None + + +def test_delete_evicts_tool_cache(tmp_path, monkeypatch): + """Deleting a server must not leave its tools cached.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": _one_tool()}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + asyncio.run(routes_mcp.delete_mcp_server("s1", current_subject = "u")) + assert mcp_client.get_cached_tools("s1") is None + + +def test_invalidate_tool_cache_clears_all(monkeypatch): + from core.inference import mcp_client + + monkeypatch.setattr(mcp_client, "_tool_cache", {"a": _one_tool(), "b": _one_tool()}) + mcp_client.invalidate_tool_cache() + assert mcp_client.get_cached_tools("a") is None + assert mcp_client.get_cached_tools("b") is None + + +def test_get_enabled_mcp_tools_probes_only_uncached(tmp_path, monkeypatch): + """An already-cached server must not be re-probed alongside a cold one.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": _one_tool("cached")}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://a/mcp", is_enabled = True) + mcp_servers_db.create_server(id = "s2", display_name = "B", url = "https://b/mcp", is_enabled = True) + + probed: list[str] = [] + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + probed.append(url) + return _one_tool("fresh") + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + specs = asyncio.run(tools_mod.get_enabled_mcp_tools()) + assert probed == ["https://b/mcp"] # only the uncached server is probed + assert sorted(t["function"]["name"] for t in specs) == ["mcp__s1__cached", "mcp__s2__fresh"] + + +def test_get_enabled_mcp_tools_partial_failure_caches_healthy(tmp_path, monkeypatch): + """One server failing must not stop the others from being cached/served.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://bad/mcp", is_enabled = True) + mcp_servers_db.create_server(id = "s2", display_name = "B", url = "https://good/mcp", is_enabled = True) + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + if "bad" in url: + raise RuntimeError("down") + return _one_tool("ok") + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + specs = asyncio.run(tools_mod.get_enabled_mcp_tools()) + assert [t["function"]["name"] for t in specs] == ["mcp__s2__ok"] + assert mcp_client.get_cached_tools("s1") is None # failure not cached + assert mcp_client.get_cached_tools("s2") == _one_tool("ok") # healthy cached + + +def test_get_enabled_mcp_tools_caches_empty_tool_list(tmp_path, monkeypatch): + """A server exposing zero tools is cached as [] (a hit), not re-probed.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + calls: list[str] = [] + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + calls.append(url) + return [] + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] + assert len(calls) == 1 # [] is a cache hit, not re-probed every send + assert mcp_client.get_cached_tools("s1") == [] + + +def test_update_headers_evicts_tool_cache(tmp_path, monkeypatch): + """Changing auth headers must drop tools discovered under the old headers.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from models.mcp_servers import McpServerUpdate + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_tool_cache", {"s1": _one_tool()}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + asyncio.run( + routes_mcp.update_mcp_server( + "s1", + McpServerUpdate(headers = {"Authorization": "Bearer new"}), + current_subject = "u", + ) + ) + assert mcp_client.get_cached_tools("s1") is None + + +def test_get_enabled_mcp_tools_skips_cache_when_config_changes_mid_probe(tmp_path, monkeypatch): + """A config edit landing during an in-flight probe must not be clobbered + by the now-stale probe result (TOCTOU on the cache write).""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True) + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + # Simulate a PUT landing while we are awaiting the probe. + mcp_servers_db.update_server("s1", {"url": "https://new/mcp"}) + return _one_tool() + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + specs = asyncio.run(tools_mod.get_enabled_mcp_tools()) + assert specs == [] # stale result is neither served... + assert mcp_client.get_cached_tools("s1") is None # ...nor cached + + +def test_get_enabled_mcp_tools_no_cooloff_when_config_changes_mid_failed_probe( + tmp_path, monkeypatch +): + """An edit landing while a probe of the OLD config is failing must not park + a cool-off on the now-fresh config -- else the re-pointed server the user + just fixed is needlessly skipped for the whole cool-off window.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True) + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + # The user re-points the server while the old endpoint's probe fails. + mcp_servers_db.update_server("s1", {"url": "https://new/mcp"}) + raise RuntimeError("old endpoint down") + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] + # The failure was for the OLD config, so the new one must stay re-probable. + assert not mcp_client.in_failure_cooloff("s1") + + +def test_get_enabled_mcp_tools_no_cooloff_when_server_deleted_mid_failed_probe( + tmp_path, monkeypatch +): + """A delete landing while a probe fails must not leave an orphan cool-off + entry keyed by the since-removed server id.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + mcp_servers_db.delete_server("s1") + raise RuntimeError("down") + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] + assert "s1" not in mcp_client._probe_cooloff_until # no orphan cool-off + + +def test_get_enabled_mcp_tools_skips_failed_server_during_cooloff(tmp_path, monkeypatch): + """A down server is probed once, then skipped during the cool-off instead + of being re-probed (and re-hung) on every send.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + attempts = {"n": 0} + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + attempts["n"] += 1 + raise RuntimeError("down") + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] # probes, fails + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] # within cool-off + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] # still skipped + assert attempts["n"] == 1 # only the first send probed + + +def test_cache_tools_clears_failure_cooloff(monkeypatch): + """A successful probe lifts a server's failure cool-off.""" + from core.inference import mcp_client + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_client.record_probe_failure("s1") + assert mcp_client.in_failure_cooloff("s1") + mcp_client.cache_tools("s1", _one_tool()) + assert not mcp_client.in_failure_cooloff("s1") + + +def test_oauth_failure_cools_off_longer_than_plain(monkeypatch): + """An OAuth server's failure cools off longer than a plain server's, so its + multi-minute probe hang doesn't recur every minute.""" + from core.inference import mcp_client + + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_client.record_probe_failure("plain", use_oauth = False) + mcp_client.record_probe_failure("oauth", use_oauth = True) + assert mcp_client._probe_cooloff_until["oauth"] > mcp_client._probe_cooloff_until["plain"] + + +def test_invalidate_clears_failure_cooloff(monkeypatch): + """Eviction drops the failure cool-off so an edited server re-probes at once.""" + from core.inference import mcp_client + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {"s1": 1.0, "s2": 2.0}) + mcp_client.invalidate_tool_cache("s1") + assert "s1" not in mcp_client._probe_cooloff_until + assert "s2" in mcp_client._probe_cooloff_until + mcp_client.invalidate_tool_cache() + assert mcp_client._probe_cooloff_until == {} + + +def test_refresh_failure_records_cooloff(tmp_path, monkeypatch): + """A failed manual refresh starts the cool-off so the next chat send does + not immediately hang on the down server.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + async def boom( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + raise RuntimeError("down") + + monkeypatch.setattr(routes_mcp, "list_tools_async", boom) + res = asyncio.run(routes_mcp.refresh_mcp_server_tools("s1", current_subject = "u")) + assert res.ok is False + assert mcp_client.in_failure_cooloff("s1") + + +def test_refresh_drops_result_when_config_changes_mid_probe(tmp_path, monkeypatch): + """A manual refresh must not warm the chat cache with tools discovered + under an old config if the server is edited while the probe is in flight.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True) + + async def fake_refresh( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + mcp_servers_db.update_server("s1", {"url": "https://new/mcp"}) + return _one_tool("stale") + + monkeypatch.setattr(routes_mcp, "list_tools_async", fake_refresh) + res = asyncio.run(routes_mcp.refresh_mcp_server_tools("s1", current_subject = "u")) + assert res.ok and res.tool_count == 1 + assert mcp_client.get_cached_tools("s1") is None + + +def test_refresh_failure_no_cooloff_when_config_changes_mid_probe(tmp_path, monkeypatch): + """A manual refresh failure for an old config must not cool off the freshly + edited server.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + import routes.mcp_servers as routes_mcp + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://old/mcp", is_enabled = True) + + async def boom( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + mcp_servers_db.update_server("s1", {"url": "https://new/mcp"}) + raise RuntimeError("old endpoint down") + + monkeypatch.setattr(routes_mcp, "list_tools_async", boom) + res = asyncio.run(routes_mcp.refresh_mcp_server_tools("s1", current_subject = "u")) + assert res.ok is False + assert not mcp_client.in_failure_cooloff("s1") + + +def test_get_enabled_mcp_tools_drops_result_when_server_deleted_mid_probe(tmp_path, monkeypatch): + """A delete landing while a probe is in flight must drop the now-orphan + result -- the `fresh is None` arm of the mid-probe TOCTOU guard. The + result is neither served nor cached under the since-removed id.""" + import asyncio + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://x/mcp", is_enabled = True) + + async def fake( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + # Simulate a DELETE landing while we await the probe. + mcp_servers_db.delete_server("s1") + return _one_tool() + + monkeypatch.setattr(tools_mod, "list_tools_async", fake) + + specs = asyncio.run(tools_mod.get_enabled_mcp_tools()) + assert specs == [] # orphan result not served + assert mcp_client.get_cached_tools("s1") is None # nor cached under a gone id + + +def test_oauth_probe_failure_in_chat_path_uses_long_cooloff(tmp_path, monkeypatch): + """When an OAuth server fails discovery during a send, the chat path must + record the OAuth (long) cool-off, not the plain one -- otherwise its + multi-minute browser hang recurs every minute.""" + import asyncio + import time + + _reset_db(tmp_path, monkeypatch) + from core.inference import mcp_client + from core.inference import tools as tools_mod + + monkeypatch.setattr(mcp_client, "_tool_cache", {}) + monkeypatch.setattr(mcp_client, "_probe_cooloff_until", {}) + mcp_servers_db.create_server( + id = "s1", + display_name = "A", + url = "https://x/mcp", + is_enabled = True, + use_oauth = True, + ) + + async def boom( + url, + headers = None, + timeout = None, + use_oauth = False, + ): + raise RuntimeError("oauth down") + + monkeypatch.setattr(tools_mod, "list_tools_async", boom) + + assert asyncio.run(tools_mod.get_enabled_mcp_tools()) == [] + assert mcp_client.in_failure_cooloff("s1") + # The recorded window must exceed the plain cool-off, proving the OAuth + # branch (use_oauth=True) fired -- not the 60 s default. + remaining = mcp_client._probe_cooloff_until["s1"] - time.monotonic() + assert remaining > mcp_client.FAILED_PROBE_COOLOFF_SECONDS diff --git a/studio/backend/tests/test_mcp_stdio_pr5863.py b/studio/backend/tests/test_mcp_stdio_pr5863.py index c6a4898d5d..9a3e8d6882 100644 --- a/studio/backend/tests/test_mcp_stdio_pr5863.py +++ b/studio/backend/tests/test_mcp_stdio_pr5863.py @@ -19,6 +19,10 @@ from storage import mcp_servers_db def _reset_db(tmp_path, monkeypatch): monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) monkeypatch.setattr(mcp_servers_db, "_schema_ready", False) + # The discovered-tool cache is process-global and keyed by server id; tests + # reuse "stdio1", so clear it (and the failure cool-off) for isolation — + # otherwise a prior test's warm cache makes discovery skip its probe. + mcp_client.invalidate_tool_cache() def _enable(monkeypatch):