From 759b821c00618afc00d7ee64b161ceea7b325732 Mon Sep 17 00:00:00 2001 From: Ben Copeland Date: Fri, 28 Aug 2026 10:20:37 +0100 Subject: [PATCH] mcp: run tool calls in worker threads FastMCP awaits an async tool but calls a sync one inline on the event loop. Every tool here is sync and makes a blocking request, so calls were serialised across all sessions: measured against a running server, a trivial list_trees took 58.66s instead of 0.14s while one slow call was in flight. Register an async wrapper that hands the call to a worker thread. The dashboard tools funnel through one registration site and inherit this; the Maestro tools are wrapped individually, so a new one needs the decorator adding. Concurrency also makes the per-call stdout redirect unsafe, as it mutates a global that overlapping calls can restore out of order, so that becomes a single redirect for the life of the server, entered inside the stdio transport once it has taken stdout for the protocol writer. anyio is imported where it is used, since kcidev.main imports every subcommand at startup and anyio comes only with the mcp extra. Signed-off-by: Ben Copeland --- kcidev/mcp/errors.py | 27 ++++++++++------------ kcidev/mcp/offload.py | 23 +++++++++++++++++++ kcidev/mcp/tools_dashboard.py | 5 +++- kcidev/mcp/tools_maestro.py | 6 +++++ kcidev/subcommands/mcp.py | 22 +++++++++++++++++- tests/test_kcidev.py | 36 +++++++++++++++++++++++++++++ tests/test_mcp_errors.py | 43 +++++++++++++++++++++++++---------- tests/test_mcp_server.py | 37 ++++++++++++++++++++++++++++++ 8 files changed, 170 insertions(+), 29 deletions(-) create mode 100644 kcidev/mcp/offload.py diff --git a/kcidev/mcp/errors.py b/kcidev/mcp/errors.py index 86ba066..8fd2770 100644 --- a/kcidev/mcp/errors.py +++ b/kcidev/mcp/errors.py @@ -1,8 +1,6 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- -import contextlib -import sys from functools import wraps import click @@ -18,18 +16,17 @@ class ToolExecutionError(Exception): def tool_errors(func): @wraps(func) def wrapper(*args, **kwargs): - with contextlib.redirect_stdout(sys.stderr): - try: - return func(*args, **kwargs) - except KciDevError as e: - raise ToolExecutionError(str(e)) from e - except click.ClickException as e: - raise ToolExecutionError(e.format_message()) from e - except click.Abort as e: - raise ToolExecutionError("KernelCI API request failed") from e - except SystemExit as e: - raise ToolExecutionError("KernelCI API request failed") from e - except requests.exceptions.RequestException as e: - raise ToolExecutionError(str(e)) from e + try: + return func(*args, **kwargs) + except KciDevError as e: + raise ToolExecutionError(str(e)) from e + except click.ClickException as e: + raise ToolExecutionError(e.format_message()) from e + except click.Abort as e: + raise ToolExecutionError("KernelCI API request failed") from e + except SystemExit as e: + raise ToolExecutionError("KernelCI API request failed") from e + except requests.exceptions.RequestException as e: + raise ToolExecutionError(str(e)) from e return wrapper diff --git a/kcidev/mcp/offload.py b/kcidev/mcp/offload.py new file mode 100644 index 0000000..bb0a1f9 --- /dev/null +++ b/kcidev/mcp/offload.py @@ -0,0 +1,23 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +import functools + +import anyio.to_thread + + +def tool_offload(func): + """Run a blocking MCP tool in a worker thread. + + FastMCP awaits async tools but calls sync ones inline on the event + loop, so a tool doing blocking I/O stalls every other request on the + server until it returns. Registering an async wrapper instead keeps + the loop free. The worker inherits a copy of the caller's context, + so tools reading contextvars still see what the caller set. + """ + + @functools.wraps(func) + async def wrapper(*args, **kwargs): + return await anyio.to_thread.run_sync(functools.partial(func, *args, **kwargs)) + + return wrapper diff --git a/kcidev/mcp/tools_dashboard.py b/kcidev/mcp/tools_dashboard.py index 07a3259..db790cd 100644 --- a/kcidev/mcp/tools_dashboard.py +++ b/kcidev/mcp/tools_dashboard.py @@ -7,6 +7,7 @@ from kcidev.api import KciDevError, KernelCIClient from kcidev.libs.filters import StatusFilter from kcidev.mcp.errors import tool_errors +from kcidev.mcp.offload import tool_offload _active_client = ContextVar("dashboard_tool_client", default=None) @@ -392,4 +393,6 @@ def bound_tool(*args, __tool=tool, **kwargs): finally: _active_client.reset(token) - server.tool(annotations=ToolAnnotations(readOnlyHint=True))(bound_tool) + server.tool(annotations=ToolAnnotations(readOnlyHint=True))( + tool_offload(bound_tool) + ) diff --git a/kcidev/mcp/tools_maestro.py b/kcidev/mcp/tools_maestro.py index 132586b..09fbfb6 100644 --- a/kcidev/mcp/tools_maestro.py +++ b/kcidev/mcp/tools_maestro.py @@ -2,6 +2,7 @@ # -*- coding: utf-8 -*- from kcidev.mcp.errors import tool_errors +from kcidev.mcp.offload import tool_offload def register_tools(server, client, api_url, pipeline_url, token): @@ -13,6 +14,7 @@ def register_tools(server, client, api_url, pipeline_url, token): if api_url: @server.tool(annotations=read_only) + @tool_offload @tool_errors def get_node(node_id: str): """Get a Maestro node (job, build or test run) by node id. @@ -25,6 +27,7 @@ def get_node(node_id: str): return client.get_node(node_id) @server.tool(annotations=read_only) + @tool_offload @tool_errors def list_nodes( filters: list[str] | None = None, @@ -52,6 +55,7 @@ def list_nodes( if pipeline_url and token: @server.tool(annotations=action) + @tool_offload @tool_errors def retry_job(node_id: str): """Retry a failed or incomplete KernelCI test job. @@ -62,6 +66,7 @@ def retry_job(node_id: str): return client.retry_job(node_id) @server.tool(annotations=action) + @tool_offload @tool_errors def trigger_checkout( giturl: str, @@ -83,6 +88,7 @@ def trigger_checkout( ) @server.tool(annotations=action) + @tool_offload @tool_errors def trigger_patchset( nodeid: str, diff --git a/kcidev/subcommands/mcp.py b/kcidev/subcommands/mcp.py index f3c5828..bf856f9 100644 --- a/kcidev/subcommands/mcp.py +++ b/kcidev/subcommands/mcp.py @@ -1,7 +1,9 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- +import contextlib import logging +import sys import click @@ -57,4 +59,22 @@ def mcp(ctx, transport, host, port): "Starting MCP server %s", "via stdio" if transport == "stdio" else f"on {host}:{port}", ) - server.run(transport="stdio" if transport == "stdio" else "streamable-http") + import anyio + + if transport == "stdio": + anyio.run(_run_stdio, server) + else: + with contextlib.redirect_stdout(sys.stderr): + server.run(transport="streamable-http") + + +async def _run_stdio(server): + from mcp.server.stdio import stdio_server + + async with stdio_server() as (read_stream, write_stream): + with contextlib.redirect_stdout(sys.stderr): + await server._mcp_server.run( + read_stream, + write_stream, + server._mcp_server.create_initialization_options(), + ) diff --git a/tests/test_kcidev.py b/tests/test_kcidev.py index fe49bb2..72b3ee2 100644 --- a/tests/test_kcidev.py +++ b/tests/test_kcidev.py @@ -916,3 +916,39 @@ def test_kcidev_mcp_help(): print(result.stderr) assert result.returncode == 0 assert "MCP" in result.stdout + + +def test_cli_imports_without_the_optional_mcp_extra(): + """The CLI must load when kci-dev is installed without [mcp]. + + kcidev.main imports every subcommand at startup, so anything the mcp + subcommand imports at module level becomes a hard dependency of the + whole CLI. anyio arrives only with the mcp extra. + """ + import sys + import textwrap + + script = textwrap.dedent( + """ + import sys + + class Blocker: + def find_spec(self, name, path=None, target=None): + if name.split(".")[0] in ("anyio", "mcp"): + raise ImportError(name) + return None + + sys.meta_path.insert(0, Blocker()) + import kcidev.main + + assert "anyio" not in sys.modules, "anyio was imported despite the blocker" + assert "mcp" not in sys.modules, "mcp was imported despite the blocker" + """ + ) + result = run( + [sys.executable, "-c", script], + stdout=PIPE, + stderr=PIPE, + universal_newlines=True, + ) + assert result.returncode == 0, result.stderr diff --git a/tests/test_mcp_errors.py b/tests/test_mcp_errors.py index f435c26..9178e8a 100644 --- a/tests/test_mcp_errors.py +++ b/tests/test_mcp_errors.py @@ -52,18 +52,6 @@ def fail(): fail() -def test_tool_errors_redirects_stdout_to_stderr(capsys): - @tool_errors - def noisy(): - print("chatter") - return 1 - - assert noisy() == 1 - captured = capsys.readouterr() - assert captured.out == "" - assert "chatter" in captured.err - - def test_tool_errors_preserves_signature(): @tool_errors def f(x: int, y: str = "a"): @@ -83,3 +71,34 @@ def fail(): with pytest.raises(ToolExecutionError, match="Dashboard build request failed"): fail() + + +def test_stdio_run_redirects_stdout_only_after_the_transport_takes_it(monkeypatch): + import contextlib + import sys + from unittest.mock import Mock + + import anyio + + from kcidev.subcommands import mcp as mcp_cmd + + seen = {} + + @contextlib.asynccontextmanager + async def fake_stdio_server(): + seen["at_capture"] = sys.stdout + yield (None, None) + + monkeypatch.setattr("mcp.server.stdio.stdio_server", fake_stdio_server) + + async def fake_run(read_stream, write_stream, options): + seen["during_run"] = sys.stdout + + server = Mock() + server._mcp_server.run = fake_run + server._mcp_server.create_initialization_options = Mock(return_value={}) + + anyio.run(mcp_cmd._run_stdio, server) + + assert seen["at_capture"] is not sys.stderr + assert seen["during_run"] is sys.stderr diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 0ed785b..f74d6b5 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -5,6 +5,7 @@ pytest.importorskip("mcp") import anyio +import anyio.to_thread import requests from mcp.shared.memory import ( create_connected_server_and_client_session as client_session, @@ -235,3 +236,39 @@ def test_list_nodes_http_error_keeps_api_detail(monkeypatch): result = _call_tool(create_server(CFG, "test"), "list_nodes", {}) assert result.isError is True assert "422" in result.content[0].text + + +def test_a_slow_tool_does_not_block_a_concurrent_one(monkeypatch): + import threading + + started = threading.Event() + release = threading.Event() + + def fake_get(url, *args, **kwargs): + response = Mock(status_code=200) + if "build/" in url: + release.set() + response.json.return_value = {"id": "maestro:b1"} + else: + started.set() + response.json.return_value = [{"released": release.wait(timeout=10)}] + return response + + monkeypatch.setattr(dashboard.kcidev_session, "get", Mock(side_effect=fake_get)) + + server = create_server() + outcome = {} + + async def run(): + async with client_session(server._mcp_server) as session: + + async def slow(): + outcome["slow"] = await session.call_tool("list_trees", {"days": 1}) + + async with anyio.create_task_group() as tg: + tg.start_soon(slow) + await anyio.to_thread.run_sync(started.wait) + await session.call_tool("get_build", {"build_id": "maestro:b1"}) + + anyio.run(run) + assert '"released": true' in outcome["slow"].content[0].text.lower()