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()