diff --git a/navi/api/deps.py b/navi/api/deps.py index 3700207..d3e6950 100644 --- a/navi/api/deps.py +++ b/navi/api/deps.py @@ -103,25 +103,17 @@ async def register_mcp_tools(registry: ToolRegistry, manager: McpManager) -> None: - """Register all currently connected MCP tools into the tool registry.""" - from navi.mcp.tools import McpTool + """Register the tools of every connected server, replacing what this layer registered. - for name, client in manager.clients.items(): - if not client.connected: - continue - try: - tools = await client.list_tools() - for tool in tools: - mcp_tool = McpTool( - server_name=name, - tool_name=tool.name, - description=tool.description or "", - parameters=tool.inputSchema, - manager=manager, - ) - registry.register_external(mcp_tool) - except Exception as exc: - log.warning("Failed to register MCP tools for %r: %s", name, exc) + The rebuild itself lives in :func:`navi.mcp.tools.build_mcp_tools`; what + matters here is that it is *read-only until it succeeds*, then committed in + one synchronous swap (``ToolRegistry.replace_external_tools``). Adding to + the registry as we went is what used to leave the model holding MCP tool + names that resolve to nothing. + """ + from navi.mcp.tools import build_mcp_tools + + registry.replace_external_tools(await build_mcp_tools(manager)) # Backward-compat re-exports so nothing outside breaks diff --git a/navi/api/routes/admin.py b/navi/api/routes/admin.py index fd22494..f5181d9 100644 --- a/navi/api/routes/admin.py +++ b/navi/api/routes/admin.py @@ -528,8 +528,8 @@ """Hot-reload user tools, context providers and MCP servers. Same work as the ``reload_tools`` built-in, for admins whose profile does - not carry that tool. ``register_mcp_tools`` is shared, so a reload from here - and one from the tool cannot leave different toolsets behind. + not carry that tool — one ``reload_all``, so a reload from here and one from + the tool cannot leave different toolsets behind. """ from navi.api.deps import get_cp_registry, get_mcp_manager, get_tool_registry from navi.core.reload import reload_all diff --git a/navi/core/registry.py b/navi/core/registry.py index a9f5164..0ef2db5 100644 --- a/navi/core/registry.py +++ b/navi/core/registry.py @@ -70,16 +70,27 @@ self._external_names.discard(name) self._tools.pop(name, None) - def clear_external_tools(self) -> list[str]: - """Drop every external tool and return the names removed. + def replace_external_tools(self, tools: list[Tool]) -> list[str]: + """Swap the whole external set for *tools* in one synchronous step. - External registrations are rebuilt from scratch on reload: an MCP - server removed from the config (or a tool it stopped exposing) has no - other way to leave the registry. + External registrations are rebuilt from scratch on reload: an MCP server + removed from the config (or a tool it stopped exposing) has no other way + to leave the registry. The caller assembles the new set *first* — that is + where a slow or silent server can fail — and commits it here, so a failed + rebuild leaves the previous registrations in place rather than an empty + registry, with the model holding names that resolve to nothing. + + Synchronous by design and with no ``await`` anywhere in it: the swap + cannot be interrupted halfway by a deadline. Returns the names dropped. """ - removed = sorted(self._external_names) + incoming = {tool.name for tool in tools} + removed = sorted(self._external_names - incoming) for name in removed: - self.unregister_external(name) + self._external_names.discard(name) + self._tools.pop(name, None) + for tool in tools: + self._tools[tool.name] = tool + self._external_names.add(tool.name) return removed def add_middleware(self, middleware) -> None: diff --git a/navi/core/reload.py b/navi/core/reload.py index ddea2b9..2950f03 100644 --- a/navi/core/reload.py +++ b/navi/core/reload.py @@ -2,11 +2,19 @@ One implementation serves both surfaces — the ``reload_tools`` built-in and the admin endpoint behind the settings button — so the two can never drift apart. -The MCP half especially: ``register_mcp_tools`` only ever *adds* registrations, -so reconnecting without dropping the old ones left ghost tools for every server -removed from the config or every tool the server stopped exposing. + +The MCP half is the delicate one, in two ways. MCP registrations are rebuilt from +what the servers answer *right now*, so a server removed from the config (or a +tool it stopped exposing) has no other way to leave the registry — which is why +they are replaced wholesale, and not added to. And because that rebuild talks to +every server, it is done in two steps: collect what each server offers, then +commit the whole set in one synchronous swap. A rebuild that fails or runs out of +budget on the way therefore leaves the previous registrations exactly as they +were, instead of an empty registry with the model holding names that resolve to +nothing. """ +import asyncio from dataclasses import dataclass, field import structlog @@ -15,6 +23,15 @@ log = structlog.get_logger() +# How long the whole MCP half may take: reconnect every server, list every +# server's tools. A healthy fleet answers in seconds; the budget exists so a +# fleet with a server that has stopped answering still produces a report — with +# the previous registrations left untouched — instead of hanging the turn that +# called reload_tools. It is the backstop, not the first line of defence: every +# individual MCP operation carries its own bound (see navi/mcp/client.py), so a +# single silent server costs its own bound here, not the whole budget. +MCP_RELOAD_BUDGET_S = 120.0 + @dataclass class ReloadReport: @@ -73,7 +90,10 @@ lines.append(f"Context provider errors ({len(self.provider_errors)}):") lines.extend(f" {f}: {e}" for f, e in self.provider_errors.items()) if self.mcp_error is not None: - lines.append(f"MCP reload error: {self.mcp_error}") + lines.append( + f"MCP reload error: {self.mcp_error} — the MCP tools registered " + "before it were left untouched." + ) else: lines.append( f"MCP tools ({len(self.mcp_tools)}): {', '.join(self.mcp_tools)}." @@ -99,18 +119,21 @@ async def _reload_mcp(registry, mcp_manager) -> tuple[list[str], str | None]: """Reconnect every MCP server and rebuild its tool registrations. - Every external registration is dropped first — see the module docstring. - Servers that are down right now lose their tools until the health-check - reconnects them, which re-registers them through the manager's - on_server_connected callback. + Budgeted as a whole, and committed in one step: whatever happens — a server + that never answers, a teardown that will not finish — the registry either + receives the complete new set or keeps the one it had. Servers that are down + right now lose their tools until the health-check reconnects them, which + re-registers them through the manager's on_server_connected callback. """ - from navi.api.deps import register_mcp_tools - from navi.mcp.tools import is_mcp_tool + from navi.mcp.tools import build_mcp_tools - registry.clear_external_tools() - await mcp_manager.reload_all() - await register_mcp_tools(registry, mcp_manager) - return sorted(t.name for t in registry.all() if is_mcp_tool(t.name)), None + async with asyncio.timeout(MCP_RELOAD_BUDGET_S): + await mcp_manager.reload_all() + tools = await build_mcp_tools(mcp_manager) + + # No await from here to the swap: the registry is never seen half-built. + registry.replace_external_tools(tools) + return sorted(t.name for t in tools), None async def reload_all(registry, cp_registry=None, mcp_manager=None) -> ReloadReport: @@ -132,6 +155,11 @@ if mcp_manager is not None and registry is not None: try: report.mcp_tools, report.mcp_error = await _reload_mcp(registry, mcp_manager) + except TimeoutError: + # The phase ran out of budget and was cancelled: nothing was + # committed, so the registry still holds the tools it had. + log.warning("reload.mcp_timeout", budget_s=MCP_RELOAD_BUDGET_S) + report.mcp_error = f"did not finish within {MCP_RELOAD_BUDGET_S:.0f}s" except Exception as exc: # noqa: BLE001 — a failed MCP pass must not lose the rest log.warning("reload.mcp_failed", error=str(exc)) report.mcp_error = f"{type(exc).__name__}: {exc}" diff --git a/navi/mcp/client.py b/navi/mcp/client.py index a39b435..4c746e6 100644 --- a/navi/mcp/client.py +++ b/navi/mcp/client.py @@ -6,6 +6,7 @@ import time from contextlib import AsyncExitStack from dataclasses import dataclass, field +from datetime import timedelta from typing import Any import httpx @@ -20,6 +21,61 @@ logger = logging.getLogger(__name__) +# ── Budgets ────────────────────────────────────────────────────────────────── +# Every wait on this path is bounded, because a server that accepted a request +# and went silent is indistinguishable from a slow one — and nothing above this +# layer could tell them apart. Before these existed the wait was *infinite*: +# with ``read_timeout_seconds=None`` the SDK waits inside +# ``anyio.fail_after(None)``, the caller's future is never resolved, and the +# runner task stays wedged for the life of the process. Every later reload then +# blocked on that same client (``disconnect_all`` is sequential), so one silent +# server made ``reload_tools`` look permanently broken. +# +# Two layers, deliberately ordered: the SDK-side read timeout fires first and +# cleans up *inside* the runner, so the client stays usable; the caller-side +# bound is the backstop, covering also a caller queued behind a stuck operation +# (and a server that stops reading its requests, where the SDK never gets to its +# read wait). It aborts the runner, which is rebuilt on the next command. +_SESSION_READ_TIMEOUT_S = 30.0 +"""Per-request read timeout handed to ``ClientSession``: initialize + tools/list.""" +_CALL_TOOL_READ_TIMEOUT_S = 240.0 +"""Per-request override for tools/call — a remote tool may legitimately take minutes.""" +_CONNECT_TIMEOUT_S = 60.0 +"""Caller-side. Twice the SDK bound, so a silent handshake fails SDK-side first.""" +_LIST_TOOLS_TIMEOUT_S = 45.0 +"""Caller-side, likewise above the SDK bound (see the retry in the runner).""" +_CALL_TOOL_TIMEOUT_S = 300.0 +_QUICK_OP_TIMEOUT_S = 15.0 +"""Ops that do no I/O — only the wait for their turn behind another op.""" +_DISCONNECT_TIMEOUT_S = 15.0 +"""Transport teardown. One that will not finish is abandoned (see disconnect()).""" + +# httpx bounds for the HTTP transports. These are the SDK's own recommended +# values (``mcp.shared._httpx_utils``), which we were not getting: we hand the +# transport our own AsyncClient for its headers, and a bare ``httpx.AsyncClient`` +# defaults to 5s for *everything* — so the streamable-http GET stream died every +# time its server went quiet for five seconds. ``read`` must stay generous: a +# streamable-http server holds the response stream open. +_HTTP_TIMEOUT_S = 30.0 +_HTTP_STREAM_READ_TIMEOUT_S = 300.0 + + +class McpOperationTimeout(RuntimeError): + """An MCP operation did not answer within its budget. + + Carries the server and op so the caller can name the culprit — on prod the + log line ``mcp.op_timeout`` is the only thing that says *which* server was + silent, because a server that never answers writes nothing itself. + """ + + def __init__(self, server: str, op: str, timeout: float) -> None: + super().__init__( + f"MCP server {server!r} did not answer '{op}' within {timeout:.0f}s" + ) + self.server = server + self.op = op + self.timeout = timeout + @dataclass class _Cmd: @@ -82,10 +138,18 @@ async def connect(self) -> None: """Open transport, initialise session. Idempotent when already connected.""" - await self._send("connect") + await self._send("connect", _CONNECT_TIMEOUT_S) async def disconnect(self) -> None: - """Close transport, stop the runner task, and reset state.""" + """Close transport, stop the runner task, and reset state. + + Bounded, because ``McpManager.disconnect_all`` walks the whole pool and + this is where a hang used to become *every* reload's problem: the + teardown runs the MCP SDK's exit stack, which can wait on a server that + stopped answering. A teardown that will not finish in + :data:`_DISCONNECT_TIMEOUT_S` is abandoned — the runner is cancelled and + its transport dies with the task. + """ runner = self._runner_task queue = self._queue # Ask the runner to stop. It closes the transport *inside its own @@ -95,10 +159,24 @@ queue.put_nowait(_Cmd("disconnect", None, {})) if runner is not None and not runner.done(): # Shield so a lifespan cancellation does not interrupt the - # runner's own transport teardown; the runner finishes in its - # own task regardless. + # runner's own transport teardown; the wait around it is bounded. try: - await asyncio.shield(runner) + await asyncio.wait_for(asyncio.shield(runner), _DISCONNECT_TIMEOUT_S) + except TimeoutError: + logger.warning( + "mcp.disconnect_timeout server=%s after=%.0fs — abandoning the runner", + self.name, + _DISCONNECT_TIMEOUT_S, + ) + self._abort_runner() + except asyncio.CancelledError: + # The caller is being cancelled (shutdown, or an outer deadline + # such as the budget on the MCP reload phase). Re-raising is + # what lets that deadline actually fire — and the runner, still + # shielded and possibly stuck in this same teardown, is + # cancelled on the way out. + self._abort_runner() + raise except BaseException: pass self._runner_task = None @@ -108,7 +186,9 @@ async def list_tools(self) -> list[Tool]: """Return the tools exposed by the remote MCP server.""" - return await self._send("list_tools") # type: ignore[return-value] + return await self._send( # type: ignore[return-value] + "list_tools", _LIST_TOOLS_TIMEOUT_S + ) async def call_tool(self, tool_name: str, arguments: dict[str, Any] | None = None) -> tuple[str, bool]: """Execute a remote tool and return its output as a string. @@ -119,7 +199,10 @@ Text content is concatenated; images are reported as a placeholder. """ return await self._send( # type: ignore[return-value] - "call_tool", tool_name=tool_name, arguments=arguments + "call_tool", + _CALL_TOOL_TIMEOUT_S, + tool_name=tool_name, + arguments=arguments, ) async def mark_disconnected(self) -> None: @@ -128,7 +211,7 @@ Used by the health-check loop when a server drops silently: the next operation will reconnect (with backoff). """ - await self._send("mark_disconnected") + await self._send("mark_disconnected", _QUICK_OP_TIMEOUT_S) # ── Runner plumbing ─────────────────────────────────────────────────── @@ -140,14 +223,52 @@ self._queue = asyncio.Queue() self._runner_task = loop.create_task(self._runner()) - async def _send(self, op: str, **payload: Any) -> Any: - """Enqueue a command and await its result.""" + def _abort_runner(self) -> None: + """Cancel the runner task and drop our handles to it. + + The runner closes the transport in its own ``finally``, so cancelling it + is what actually frees the stream. Clearing the handles here — rather + than letting the dying task do it — is what keeps the next command from + queueing behind a runner that is already unwinding, and the mirrored + state has to be cleared here too: the dying runner's own teardown is + ignored (it is no longer the current one), so nothing else would. + """ + runner = self._runner_task + self._runner_task = None + self._queue = None + self._connected = False + self._instructions = None + if runner is not None and not runner.done(): + runner.cancel() + + async def _send(self, op: str, timeout: float, **payload: Any) -> Any: + """Enqueue a command and await its result, for at most *timeout* seconds. + + *timeout* is required so no future op can quietly reintroduce an + unbounded wait; the module constants above say what each op is worth. + """ self._ensure_runner() loop = asyncio.get_running_loop() fut: asyncio.Future = loop.create_future() assert self._queue is not None await self._queue.put(_Cmd(op, fut, payload)) - return await fut + try: + return await asyncio.wait_for(fut, timeout) + except TimeoutError: + # The runner is somewhere inside this op and will not come back: + # abort it, or every later command queues behind the same silence. + logger.warning( + "mcp.op_timeout server=%s op=%s after=%.0fs", self.name, op, timeout + ) + self._abort_runner() + raise McpOperationTimeout(self.name, op, timeout) from None + except asyncio.CancelledError: + # The caller gave up (outer deadline, stop, shutdown). Nothing will + # consume this result, and a runner still waiting on a silent server + # would stay wedged for the life of the process. + fut.cancel() + self._abort_runner() + raise async def _runner(self) -> None: """Owns the AsyncExitStack; processes commands until ``disconnect``. @@ -158,6 +279,7 @@ """ queue = self._queue assert queue is not None + me = asyncio.current_task() exit_stack = AsyncExitStack() session: ClientSession | None = None @@ -166,6 +288,17 @@ backoff: float = 1.0 last_attempt: float | None = None + def publish_state() -> None: + """Mirror this runner's state — unless this runner was replaced. + + A runner that a timed-out operation aborted keeps unwinding its own + teardown for a moment, and its dying state must not be written over + the fresh runner a later command has already started. + """ + if self._runner_task is me: + self._connected = connected + self._instructions = instructions + async def close_all() -> None: nonlocal session, connected, instructions try: @@ -180,8 +313,7 @@ session = None connected = False instructions = None - self._connected = False - self._instructions = None + publish_state() async def open_transport() -> None: nonlocal session, instructions @@ -206,13 +338,25 @@ if not cfg.url: raise ValueError("sse transport requires 'url'") transport = await exit_stack.enter_async_context( - sse_client(cfg.url, headers=cfg.headers) + sse_client( + cfg.url, + headers=cfg.headers, + timeout=_HTTP_TIMEOUT_S, + sse_read_timeout=_HTTP_STREAM_READ_TIMEOUT_S, + ) ) elif cfg.is_streamable_http: if not cfg.url: raise ValueError("streamable_http transport requires 'url'") http_client = await exit_stack.enter_async_context( - httpx.AsyncClient(headers=cfg.headers or {}) + httpx.AsyncClient( + headers=cfg.headers or {}, + # Never the httpx default: 5s on the GET stream read + # tears the transport down on an idle server. + timeout=httpx.Timeout( + _HTTP_TIMEOUT_S, read=_HTTP_STREAM_READ_TIMEOUT_S + ), + ) ) transport = await exit_stack.enter_async_context( streamable_http_client(cfg.url, http_client=http_client) @@ -224,7 +368,12 @@ read_stream, write_stream = transport session = await exit_stack.enter_async_context( - ClientSession(read_stream, write_stream) + ClientSession( + read_stream, + write_stream, + # Without this the SDK waits forever on every request. + read_timeout_seconds=timedelta(seconds=_SESSION_READ_TIMEOUT_S), + ) ) init_result = await session.initialize() instructions = getattr(init_result, "instructions", None) @@ -262,8 +411,7 @@ connected = True backoff = 1.0 last_attempt = None - self._connected = True - self._instructions = instructions + publish_state() logger.info("MCP server %r connected (%s)", self.name, self.config.transport) try: @@ -277,7 +425,7 @@ break # finally -> close_all elif cmd.op == "mark_disconnected": connected = False - self._connected = False + publish_state() self._resolve(cmd.fut, None) elif cmd.op == "list_tools": await ensure_connected() @@ -294,13 +442,21 @@ await ensure_connected() try: result = await session.call_tool( # type: ignore[union-attr] - tool_name, arguments or {} + tool_name, + arguments or {}, + read_timeout_seconds=timedelta( + seconds=_CALL_TOOL_READ_TIMEOUT_S + ), ) except Exception: await close_all() await ensure_connected() result = await session.call_tool( # type: ignore[union-attr] - tool_name, arguments or {} + tool_name, + arguments or {}, + read_timeout_seconds=timedelta( + seconds=_CALL_TOOL_READ_TIMEOUT_S + ), ) parts: list[str] = [] for item in result.content: @@ -328,16 +484,14 @@ # dead endpoint). Treat it as an ordinary connect failure: # one unreachable MCP server must not take down the whole # startup (McpManager.load_all only catches Exception). - self._connected = connected - self._instructions = instructions + publish_state() self._fail(cmd.fut, RuntimeError( f"MCP server {self.name!r} connect failed (anyio CancelledError)" )) except Exception as exc: # On a connect/ensure failure, mirror the disconnected # state before surfacing the error to the caller. - self._connected = connected - self._instructions = instructions + publish_state() self._fail(cmd.fut, exc) finally: await close_all() diff --git a/navi/mcp/manager.py b/navi/mcp/manager.py index d91eea3..d8c56fb 100644 --- a/navi/mcp/manager.py +++ b/navi/mcp/manager.py @@ -148,7 +148,16 @@ await self.load_all(configs) async def disconnect_all(self) -> None: - """Close every open connection and clear the client pool.""" + """Close every open connection and clear the client pool. + + Each ``disconnect`` carries its own bound (see + :meth:`McpClient.disconnect`), so one server stuck in teardown delays + this by its bound and no longer blocks every reload that follows it. + + Cancellation is *not* swallowed here: this runs under the MCP reload's + phase budget, and a swallowed cancel between two clients is exactly how + such a deadline stops being enforceable. + """ if not self._clients and not self._user_clients: return self._visible_cache.clear() @@ -156,7 +165,7 @@ try: await client.disconnect() except asyncio.CancelledError: - pass + raise except Exception as exc: logger.warning("MCP server %r disconnect error: %s", name, exc) self._clients.clear() @@ -165,7 +174,7 @@ try: await client.disconnect() except asyncio.CancelledError: - pass + raise except Exception as exc: logger.warning( "MCP user client (server=%s user=%s) disconnect error: %s", @@ -463,7 +472,15 @@ await asyncio.sleep(5.0) async def _run_health_check(self) -> None: - """Probe all servers: reconnect dead ones, verify live ones are still alive.""" + """Probe all servers: reconnect dead ones, verify live ones are still alive. + + Every probe goes through :class:`McpClient`, where each operation carries + its own bound — so a server that has stopped answering delays this sweep + by its bound and can never freeze the loop. That matters beyond the sweep + itself: this loop is what reconnects a dropped server *and* re-registers + its tools through the ``on_server_connected`` callback, so a frozen loop + used to leave those tools missing until the process was restarted. + """ from navi.core.event_bus import get_event_bus from navi.core.events import McpStatusUpdate diff --git a/navi/mcp/tools.py b/navi/mcp/tools.py index 38bb5e5..9fb560a 100644 --- a/navi/mcp/tools.py +++ b/navi/mcp/tools.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +import logging from typing import Any from pathlib import Path @@ -17,6 +18,8 @@ from .manager import McpManager +logger = logging.getLogger(__name__) + # Single source of truth for MCP tool naming. _MCP_SEP = "__" @@ -143,3 +146,34 @@ return ToolResult(success=True, output=output, metadata=metadata) except Exception as exc: return ToolResult(success=False, output="", error=str(exc)) + + +async def build_mcp_tools(manager: McpManager) -> list[McpTool]: + """Every tool every *connected* server exposes right now, as proxies. + + Read-only on purpose. The reload assembles the complete set here and only + then swaps it into the registry, so a server that fails or stops answering + halfway cannot leave the registry with a half-built toolset. A server that + will not list its tools is skipped — the health-check reconnects it, and the + manager's ``on_server_connected`` callback registers it then. + """ + tools: list[McpTool] = [] + for name, client in manager.clients.items(): + if not client.connected: + continue + try: + listed = await client.list_tools() + except Exception as exc: + logger.warning("MCP server %r list_tools failed: %s", name, exc) + continue + tools.extend( + McpTool( + server_name=name, + tool_name=tool.name, + description=tool.description or "", + parameters=tool.inputSchema, + manager=manager, + ) + for tool in listed + ) + return tools diff --git a/tests/integration/test_mcp_integration.py b/tests/integration/test_mcp_integration.py index 3c0217b..e6d8554 100644 --- a/tests/integration/test_mcp_integration.py +++ b/tests/integration/test_mcp_integration.py @@ -2,6 +2,7 @@ import asyncio import sys +import time import pytest @@ -112,3 +113,32 @@ await client.connect() assert "NAVI_TEST_DEFINITELY_UNSET" in str(excinfo.value) assert not client.connected + + +@pytest.mark.anyio +async def test_a_real_child_that_never_answers_is_bounded(monkeypatch): + """The prod shape, against a real subprocess: a child that reads its stdin + and never replies. Without a read timeout the handshake waits forever and + the runner task stays wedged for the life of the process — which is what + made every later reload block on this one client.""" + import navi.mcp.client as client_mod + + monkeypatch.setattr(client_mod, "_SESSION_READ_TIMEOUT_S", 0.3) + monkeypatch.setattr(client_mod, "_CONNECT_TIMEOUT_S", 3.0) + + client = McpClient( + "mute", + McpServerConfig( + transport="stdio", + command=sys.executable, + args=["-c", "import sys; sys.stdin.read()"], + ), + ) + + started = time.monotonic() + with pytest.raises(Exception): + await asyncio.wait_for(client.connect(), timeout=10) + + assert time.monotonic() - started < 5 + assert not client.connected + await client.disconnect() diff --git a/tests/unit/core/test_registry.py b/tests/unit/core/test_registry.py index 4801f25..e8558d2 100644 --- a/tests/unit/core/test_registry.py +++ b/tests/unit/core/test_registry.py @@ -41,6 +41,34 @@ assert "builtin" in reg._builtin_names assert "user" not in reg._builtin_names + def test_replace_external_tools_swaps_the_whole_set(self): + """The rebuild commits in one step, so a caller that failed halfway + never has to leave the registry missing what it had.""" + reg = ToolRegistry() + reg.register(FakeTool("builtin"), builtin=True) + reg.register_external(FakeTool("mcp__a__one")) + reg.register_external(FakeTool("mcp__a__two")) + + removed = reg.replace_external_tools( + [FakeTool("mcp__a__two"), FakeTool("mcp__b__one")] + ) + + assert removed == ["mcp__a__one"] + assert sorted(t.name for t in reg.all()) == [ + "builtin", "mcp__a__two", "mcp__b__one", + ] + + def test_replace_external_tools_leaves_builtins_and_user_tools_alone(self): + reg = ToolRegistry() + reg.register(FakeTool("builtin"), builtin=True) + reg.register(FakeTool("user")) + reg.register_external(FakeTool("mcp__a__one")) + + reg.replace_external_tools([]) + + assert sorted(t.name for t in reg.all()) == ["builtin", "user"] + assert reg._external_names == set() + class TestProfileRegistry: def test_register_and_get(self): diff --git a/tests/unit/core/test_reload.py b/tests/unit/core/test_reload.py index 91c1875..037fa4d 100644 --- a/tests/unit/core/test_reload.py +++ b/tests/unit/core/test_reload.py @@ -1,11 +1,14 @@ """reload_all() — the single implementation behind reload_tools and the admin button. -The MCP half is the delicate one: register_mcp_tools only adds registrations, so -a reload that reconnects without dropping the old ones leaves ghost tools — for -every server taken out of the config, and for every tool a server stopped -exposing. They stay visible to the model and fail only when called. +The MCP half is the delicate one. MCP registrations are rebuilt from what the +servers answer *right now* — that is the only way a server removed from the +config, or a tool a server stopped exposing, ever leaves the registry. So the +rebuild happens in two steps: collect everything, then commit it in one swap. A +rebuild that fails, or a server that never answers, must leave the previous +registrations exactly as they were rather than an empty registry. """ +import asyncio import textwrap import pytest @@ -34,12 +37,15 @@ class FakeClient: - def __init__(self, server, tools): + def __init__(self, server, tools, silent=False): self.server = server self.connected = True self._tools = tools + self.silent = silent async def list_tools(self): + if self.silent: + raise RuntimeError(f"no answer from {self.server}") return [ type("T", (), {"name": n, "description": "d", "inputSchema": {}}) for n in self._tools @@ -53,10 +59,16 @@ self.servers = servers self.reloads = 0 self.fail = False + # Servers that stop answering: ``reload_all`` succeeds, listing them + # does not — the shape of a server that is connected but silent. + self.silent: set[str] = set() @property def clients(self): - return {s: FakeClient(s, t) for s, t in self.servers.items()} + return { + s: FakeClient(s, t, silent=s in self.silent) + for s, t in self.servers.items() + } async def reload_all(self): self.reloads += 1 @@ -64,6 +76,15 @@ raise RuntimeError("mcp is down") +class HangingMcpManager: + """A manager whose reload never returns — the shape of the prod hang.""" + + clients: dict = {} + + async def reload_all(self): + await asyncio.sleep(3600) + + @pytest.fixture def isolated_settings(tmp_path, monkeypatch): """Point tools_dir at a scratch directory and control enabled.json.""" @@ -200,6 +221,60 @@ assert registry.get("alpha") is not None +async def test_a_failed_mcp_pass_keeps_the_tools_it_had(isolated_settings): + """The registry is committed in one step at the end of the rebuild, so a + rebuild that dies halfway leaves what it had — not an empty registry with + the model holding MCP names that resolve to nothing.""" + registry = ToolRegistry() + manager = FakeMcpManager({"alpha": ["ping", "pong"]}) + await reload_all(registry, None, manager) + + manager.fail = True + report = await reload_all(registry, None, manager) + + assert report.mcp_error is not None + assert registry.get("mcp__alpha__ping") is not None + assert registry.get("mcp__alpha__pong") is not None + assert any("left untouched" in line for line in report.lines()) + + +async def test_a_silent_server_cannot_hang_the_reload(isolated_settings, monkeypatch): + """A server that never answers must cost the turn nothing: the phase is + budgeted, and the report says so instead of the tool call hanging.""" + import navi.core.reload as reload_mod + + monkeypatch.setattr(reload_mod, "MCP_RELOAD_BUDGET_S", 0.2) + registry = ToolRegistry() + manager = FakeMcpManager({"alpha": ["ping"]}) + await reload_all(registry, None, manager) + + report = await asyncio.wait_for( + reload_all(registry, None, HangingMcpManager()), timeout=5 + ) + + assert report.ok is False + assert "did not finish within" in report.mcp_error + assert registry.get("mcp__alpha__ping") is not None + assert any("left untouched" in line for line in report.lines()) + + +async def test_the_tools_of_a_silent_server_do_not_take_the_others_down(isolated_settings): + """One server that will not list its tools loses its own registrations for + this round — the health check re-registers them when it answers again — + while every other server keeps its own.""" + registry = ToolRegistry() + manager = FakeMcpManager({"alpha": ["ping"], "beta": ["pong"]}) + await reload_all(registry, None, manager) + + manager.silent = {"alpha"} + report = await reload_all(registry, None, manager) + + assert report.mcp_error is None + assert report.mcp_tools == ["mcp__beta__pong"] + assert registry.get("mcp__beta__pong") is not None + assert [t.name for t in registry.all()] == ["mcp__beta__pong"] + + def test_report_serialises_for_the_api(): report = ReloadReport(tools=["a"], total_tools=2, tool_errors={"x.py": "boom"}) data = report.as_dict() diff --git a/tests/unit/mcp/test_client_bounds.py b/tests/unit/mcp/test_client_bounds.py new file mode 100644 index 0000000..1ca9a67 --- /dev/null +++ b/tests/unit/mcp/test_client_bounds.py @@ -0,0 +1,228 @@ +"""Every wait on the MCP path is bounded. + +A server that accepts a request and never answers is the shape of the prod +failure. With no read timeout the SDK waits forever (``anyio.fail_after(None)``), +the caller's future is never resolved, and the runner task stays wedged for the +life of the process — so every later reload blocked on that same client and +``disconnect_all`` walked the pool one stuck client at a time. + +These tests talk to a fake transport through a *real* ``ClientSession``: it reads +requests and replies to nothing, which is what a silent server looks like from +here. The bounds are the code under test, so each test patches them down to +fractions of a second. +""" + +import asyncio +import contextlib +import time + +import anyio +import httpx +import pytest +from mcp.shared.message import SessionMessage +from mcp.types import ( + LATEST_PROTOCOL_VERSION, + Implementation, + InitializeResult, + JSONRPCMessage, + JSONRPCResponse, + ServerCapabilities, +) + +import navi.mcp.client as client_mod +from navi.mcp.client import McpClient, McpOperationTimeout +from navi.mcp.config import McpServerConfig +from navi.mcp.manager import McpManager + + +class SilentServer: + """A stand-in stdio transport: reads every request, answers almost none. + + ``answer_handshake=False`` makes it silent from the very first byte — the + server that accepts the connection and says nothing at all. Otherwise it + completes the ``initialize`` handshake (so the session is real and usable) + and then goes quiet, which is how a server that died mid-session behaves. + """ + + def __init__(self, *, answer_handshake: bool = True) -> None: + self.answer_handshake = answer_handshake + self.requests: list[str] = [] + self._tasks: list[asyncio.Task] = [] + + def __call__(self, params): # stands in for stdio_client(params) + return self + + async def __aenter__(self): + # (read_stream, write_stream) from the client's side: it reads what we + # send, and we read what it writes. Streams are buffered, so a client + # write never blocks on us reading it. + to_client_w, to_client = anyio.create_memory_object_stream[SessionMessage](8) + from_client, from_client_r = anyio.create_memory_object_stream[SessionMessage](8) + self._to_client = to_client_w + self._from_client = from_client + self._tasks.append(asyncio.create_task(self._serve(from_client_r))) + return to_client, from_client + + async def __aexit__(self, *exc_info) -> bool: + for task in self._tasks: + task.cancel() + for task in self._tasks: + with contextlib.suppress(BaseException): + await task + self._tasks.clear() + return False + + async def _serve(self, reader) -> None: + async for session_message in reader: + request = session_message.message.root + method = getattr(request, "method", None) + if method is None: + continue + self.requests.append(method) + if method != "initialize" or not self.answer_handshake: + continue # read it, never answer it + result = InitializeResult( + protocolVersion=LATEST_PROTOCOL_VERSION, + capabilities=ServerCapabilities(), + serverInfo=Implementation(name="silent", version="1"), + ) + await self._to_client.send( + SessionMessage( + message=JSONRPCMessage( + JSONRPCResponse( + jsonrpc="2.0", + id=request.id, + result=result.model_dump(by_alias=True, mode="json"), + ) + ) + ) + ) + + +def _cfg() -> McpServerConfig: + return McpServerConfig(transport="stdio", command="python", args=["-c", "pass"]) + + +@pytest.fixture +def quiet_after_handshake(monkeypatch): + """A server that completes the handshake, then never answers again.""" + server = SilentServer(answer_handshake=True) + monkeypatch.setattr(client_mod, "stdio_client", server) + return server + + +@pytest.fixture +def silent_from_the_start(monkeypatch): + """A server that accepts the connection and answers nothing at all.""" + server = SilentServer(answer_handshake=False) + monkeypatch.setattr(client_mod, "stdio_client", server) + return server + + +async def test_a_silent_server_cannot_wedge_list_tools(quiet_after_handshake, monkeypatch): + """The SDK-level read timeout is what ends a request nobody answers — and it + ends it *inside* the runner, so the client is still usable afterwards.""" + monkeypatch.setattr(client_mod, "_SESSION_READ_TIMEOUT_S", 0.2) + monkeypatch.setattr(client_mod, "_LIST_TOOLS_TIMEOUT_S", 10.0) + client = McpClient("quiet", _cfg()) + await asyncio.wait_for(client.connect(), timeout=5) + + started = time.monotonic() + with pytest.raises(Exception, match="Timed out"): + await asyncio.wait_for(client.list_tools(), timeout=5) + + assert time.monotonic() - started < 3 + assert quiet_after_handshake.requests.count("tools/list") >= 1 + await client.disconnect() + + +async def test_the_caller_bound_covers_an_op_the_sdk_never_times_out( + quiet_after_handshake, monkeypatch +): + """The layer below the SDK's: an op that never reaches the SDK's read wait — + a server that stops reading its requests, a caller queued behind one — is + caught here, and the wedged runner is aborted rather than left for the next + command to queue behind.""" + monkeypatch.setattr(client_mod, "_SESSION_READ_TIMEOUT_S", 60.0) + monkeypatch.setattr(client_mod, "_LIST_TOOLS_TIMEOUT_S", 0.2) + monkeypatch.setattr(client_mod, "_CONNECT_TIMEOUT_S", 5.0) + client = McpClient("wedged", _cfg()) + await client.connect() + + with pytest.raises(McpOperationTimeout) as excinfo: + await asyncio.wait_for(client.list_tools(), timeout=5) + + assert excinfo.value.server == "wedged" + assert excinfo.value.op == "list_tools" + # Aborted, not merely abandoned: the next command starts a fresh runner. + assert client._runner_task is None + assert client.connected is False + + +async def test_a_server_that_never_answers_cannot_hang_load_all( + silent_from_the_start, monkeypatch +): + """Startup (and every reload) connects to every configured server. One that + accepts the connection and says nothing used to hang the whole loop.""" + monkeypatch.setattr(client_mod, "_SESSION_READ_TIMEOUT_S", 0.2) + manager = McpManager() + + started = time.monotonic() + await asyncio.wait_for( + manager.load_all({"quiet": _cfg()}), timeout=5 + ) + + assert time.monotonic() - started < 3 + assert "quiet" in manager.clients # kept in the pool for the health check + assert manager._connected_status["quiet"] is False + await manager.disconnect_all() + + +async def test_a_wedged_client_does_not_block_load_all(silent_from_the_start, monkeypatch): + """The prod symptom, at the manager level: ``load_all`` disconnects the whole + pool first, so one client stuck in teardown blocked every reload after it.""" + monkeypatch.setattr(client_mod, "_SESSION_READ_TIMEOUT_S", 60.0) + monkeypatch.setattr(client_mod, "_CONNECT_TIMEOUT_S", 60.0) + monkeypatch.setattr(client_mod, "_DISCONNECT_TIMEOUT_S", 0.2) + + manager = McpManager() + wedged = McpClient("wedged", _cfg()) + manager._clients["wedged"] = wedged + pending = asyncio.create_task(wedged.connect()) # never answers + await asyncio.sleep(0.1) + + started = time.monotonic() + await asyncio.wait_for(manager.load_all({}), timeout=5) + + assert time.monotonic() - started < 3 + assert manager.clients == {} + + # The caller that was waiting on the abandoned runner is released, not left + # hanging: a reload's budget has to be able to reach the end of its work. + with pytest.raises(BaseException): + await asyncio.wait_for(pending, timeout=2) + + +async def test_the_streamable_http_transport_is_not_left_on_httpx_defaults(monkeypatch): + """A bare ``httpx.AsyncClient`` defaults to 5s for *everything*, so the GET + stream of an idle streamable-http server — and the POST that carries a slow + tool call — died on a timeout nobody chose.""" + monkeypatch.setattr(client_mod, "_SESSION_READ_TIMEOUT_S", 0.2) + captured: dict = {} + + class RecordingClient(httpx.AsyncClient): + def __init__(self, **kwargs): + captured.update(kwargs) + super().__init__(**kwargs) + + monkeypatch.setattr(client_mod.httpx, "AsyncClient", RecordingClient) + client = McpClient( + "streamy", + McpServerConfig(transport="streamable_http", url="http://127.0.0.1:9/mcp"), + ) + with pytest.raises(Exception): + await client.connect() + + assert captured["timeout"].read == client_mod._HTTP_STREAM_READ_TIMEOUT_S + assert captured["timeout"].connect == client_mod._HTTP_TIMEOUT_S + await client.disconnect()