From 6e522aa5e113491b7d63651a93331f08906cbf55 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 27 Aug 2026 15:09:13 +0000 Subject: [PATCH 01/20] Add workspace-backed startup and provider config Co-authored-by: alanwang233233 <107164653+alanwang233233@users.noreply.github.com> --- src/mitocode/__main__.py | 219 +++++++++++++----- .../backend/agent/agent_orchestrator.py | 12 +- src/mitocode/tui/ui_main.py | 65 +++++- test/test_main_entrypoint.py | 40 ++++ test/test_orchestrator_and_bridge.py | 41 ++++ 5 files changed, 322 insertions(+), 55 deletions(-) create mode 100644 test/test_main_entrypoint.py diff --git a/src/mitocode/__main__.py b/src/mitocode/__main__.py index 0a0fa30..12936ed 100644 --- a/src/mitocode/__main__.py +++ b/src/mitocode/__main__.py @@ -1,96 +1,209 @@ -"""WIP""" +"""CLI entrypoint for running the MitoCode TUI.""" import argparse import asyncio -from pyexpat.errors import messages - -from httpx import stream +import json +import os +from pathlib import Path +from typing import Any, Dict +from mitocode import APP_VERSION, MITOCODE_CONFIG_PATH from mitocode.backend.agent.agent_orchestrator import AgentOrchestrator -from mitocode.backend.chat_provider.scenario_provider import ScenarioProvider from mitocode.backend.chat_provider.chat_provider_openai import ChatProviderOpenAI +from mitocode.backend.chat_provider.scenario_provider import ScenarioProvider +from mitocode.backend.chat_provider.sample_chat_provider import SampleChatProvider from mitocode.backend.tool.tool_manager import ToolManager +from mitocode.backend.workspace.workspace import Workspace from mitocode.tui.ui_main import VibeFrontend -def _register_tools(tool_manager: ToolManager): - """Register built-in demo tools (bash, read_file, write_file) on the manager.""" +def _register_tools(tool_manager: ToolManager, workspace: Workspace) -> None: + """Register workspace-backed tools (bash, read_file, write_file).""" @tool_manager.tool( - human_doc="Run a shell command", - agent_doc="Execute a bash command and return its output", + human_doc="Run a shell command inside workspace", + agent_doc="Execute a shell command in the workspace directory and return output", display_action="bash", ) async def bash(command: str) -> str: - return f"$ {command}\n[bash] command executed successfully" + proc = await asyncio.create_subprocess_shell( + command, + cwd=str(workspace.path), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + stdout_bytes, stderr_bytes = await proc.communicate() + stdout = stdout_bytes.decode("utf-8", errors="replace") + stderr = stderr_bytes.decode("utf-8", errors="replace") + output = stdout if stdout else "(no stdout)" + if stderr: + output += f"\n[stderr]\n{stderr}" + return ( + f"$ {command}\n[cwd] {workspace.path}\n[exit_code] {proc.returncode}\n{output}" + ) @tool_manager.tool( human_doc="Read file content", - agent_doc="Read the full content of a file at the given path", + agent_doc="Read UTF-8 text file content from a workspace-relative path", display_action="file.read", display_param="file_path", ) async def read_file(file_path: str) -> str: - return ( - "def bubble_sort(arr):\n" - " n = len(arr)\n" - " for i in range(n):\n" - " swapped = False\n" - " for j in range(0, n - i - 1):\n" - " if arr[j] > arr[j + 1]:\n" - " arr[j], arr[j + 1] = arr[j + 1], arr[j]\n" - " swapped = True\n" - " if not swapped:\n" - " break\n" - " return arr\n" - ) + safe_path = workspace._validate_within_workspace(Path(file_path)) + if not safe_path.exists(): + raise FileNotFoundError(f"File not found: {safe_path}") + if not safe_path.is_file(): + raise IsADirectoryError(f"Path is not a file: {safe_path}") + return await asyncio.to_thread(safe_path.read_text, encoding="utf-8") @tool_manager.tool( human_doc="Write content to a file", - agent_doc="Create or overwrite a file with the given content", + agent_doc="Write UTF-8 content into a workspace-relative file path", display_action="file.write", display_param="file_path", ) async def write_file(file_path: str, content: str) -> str: - return f"Successfully wrote {len(content)} bytes to {file_path}" + safe_path = workspace._validate_within_workspace(Path(file_path)) + await asyncio.to_thread(safe_path.parent.mkdir, parents=True, exist_ok=True) + await asyncio.to_thread(safe_path.write_text, content, encoding="utf-8") + return f"Successfully wrote {len(content)} bytes to {safe_path}" + + +def _parse_provider_settings(raw: str) -> Dict[str, Any]: + """Parse user-provided provider settings JSON string.""" + if not raw: + return {} + parsed = json.loads(raw) + if not isinstance(parsed, dict): + raise ValueError("--provider-settings must be a JSON object") + return parsed + + +def _default_provider_kind() -> str: + return ( + "openai" + if os.getenv("OPENAI_API_KEY") or os.getenv("OPENROUTER_API_KEY") + else "scenario" + ) -def main(): +def _build_provider(args: argparse.Namespace) -> SampleChatProvider: + """Build and initialize the provider from CLI arguments.""" + if args.provider == "scenario": + settings = {"delay": args.speed} + settings.update(_parse_provider_settings(args.provider_settings)) + return asyncio.run( + ScenarioProvider.factory( + default_model=args.model or "scenario-model", + default_model_list=["scenario-model"], + default_messages=None, + default_tools=None, + stream=True, + settings=settings, + delay=settings.get("delay", args.speed), + ) + ) + + api_key = args.api_key or os.getenv("OPENAI_API_KEY") or os.getenv("OPENROUTER_API_KEY") + if not api_key: + raise ValueError("OpenAI provider requires --api-key or OPENAI_API_KEY/OPENROUTER_API_KEY") + + settings: Dict[str, Any] = { + "api_key": api_key, + "timeout": args.timeout, + } + if args.base_url: + settings["base_url"] = args.base_url + settings.update(_parse_provider_settings(args.provider_settings)) + + default_model = args.model or "gpt-4o-mini" + return asyncio.run( + ChatProviderOpenAI.factory( + default_model=default_model, + default_model_list=[default_model], + default_messages=None, + default_tools=None, + stream=True, + settings=settings, + ) + ) + + +def _init_workspace(path_arg: str, name: str) -> Workspace: + """Initialize or create workspace at the given path.""" + workspace_root = Path(path_arg).expanduser().resolve() + workspace_root.mkdir(parents=True, exist_ok=True) + + config_file = workspace_root / MITOCODE_CONFIG_PATH / "config.json" + if config_file.exists(): + return Workspace.init_from_path(workspace_root) + + return Workspace.new_workspace( + workspace_path=workspace_root, + name=name, + description=f"MitoCode workspace (app v{APP_VERSION})", + override=False, + ) + + +def main() -> None: """Parse CLI arguments and launch the MitoCode TUI application.""" parser = argparse.ArgumentParser(description="MitoCode") parser.add_argument( "--speed", type=float, default=0.03, - help="Delay between chunks in seconds (default: 0.03)", + help="Delay between scenario provider chunks in seconds (default: 0.03)", + ) + parser.add_argument( + "--workspace", + type=str, + default=".", + help="Workspace root path (default: current directory)", + ) + parser.add_argument( + "--workspace-name", + type=str, + default="mitocode-workspace", + help="Workspace name for first-time initialization", + ) + parser.add_argument( + "--provider", + choices=["scenario", "openai"], + default=_default_provider_kind(), + help="LLM provider backend", + ) + parser.add_argument("--model", type=str, default="", help="Provider model name") + parser.add_argument( + "--api-key", + type=str, + default="", + help="OpenAI/OpenRouter API key (or use OPENAI_API_KEY / OPENROUTER_API_KEY)", + ) + parser.add_argument( + "--base-url", + type=str, + default=os.getenv("OPENAI_BASE_URL", ""), + help="OpenAI-compatible API base URL override", + ) + parser.add_argument( + "--timeout", + type=float, + default=60, + help="Provider request timeout in seconds", + ) + parser.add_argument( + "--provider-settings", + type=str, + default="", + help="Extra provider settings as JSON object", ) args = parser.parse_args() + workspace = _init_workspace(args.workspace, args.workspace_name) tool_manager = ToolManager() - _register_tools(tool_manager) - - provider = asyncio.run(ChatProviderOpenAI.factory( - default_model= "nvidia/nemotron-3.5-lightning:free", - default_model_list= None, - default_messages= None, - default_tools= None, - stream= True, - settings={ - "api_key": "sk-别看了这次我不会再把key提交到仓库了！", - "base_url": "https://openrouter.ai/api/v1", - "timeout": 60, - }) - ) - - # print(asyncio.run(provider.chat( - # model ="nvidia/nemotron-3.5-lightning:free", - # tools =None, - # messages = [ - # {"role": "system", "content": "You are a helpful assistant."}, - # {"role": "user", "content": "Hello!"} - # ], - # stream = False - # ))) + _register_tools(tool_manager, workspace) + provider = _build_provider(args) orchestrator = AgentOrchestrator( tool_manager=tool_manager, default_provider=provider @@ -98,7 +211,7 @@ def main(): asyncio.run( orchestrator.create_agent( name="Main", - description="Chief orchestrator", + description=f"Workspace: {workspace.path}", ) ) diff --git a/src/mitocode/backend/agent/agent_orchestrator.py b/src/mitocode/backend/agent/agent_orchestrator.py index 0f39318..70a46f1 100644 --- a/src/mitocode/backend/agent/agent_orchestrator.py +++ b/src/mitocode/backend/agent/agent_orchestrator.py @@ -264,8 +264,18 @@ async def _run_session(self, agent_id: int) -> None: except asyncio.CancelledError: session.state = AgentState.CANCELLED raise - except Exception: # pylint: disable=broad-exception-caught + except Exception as e: # pylint: disable=broad-exception-caught session.state = AgentState.ERROR + await self.event_queue.put( + ( + agent_id, + { + "type": "error", + "agent_id": agent_id, + "message": str(e) or "Unknown runtime error", + }, + ) + ) finally: if self._tasks.get(agent_id) is asyncio.current_task(): self._tasks[agent_id] = None diff --git a/src/mitocode/tui/ui_main.py b/src/mitocode/tui/ui_main.py index 5aaad3c..27e8296 100644 --- a/src/mitocode/tui/ui_main.py +++ b/src/mitocode/tui/ui_main.py @@ -360,6 +360,47 @@ def render(self): ) +class WelcomePanel(Static): + """Startup panel shown before first message is sent.""" + + panel_type = "welcome" + + def render(self): + return Panel( + RichMarkdown( + "## 欢迎使用 MitoCode TUI\n\n" + "- 在左侧选择 Agent\n" + "- 在下方输入框发送消息\n" + "- 使用 `Quick create` 快速创建子 Agent\n" + "- `F12` 打开调试日志" + ), + title="[bold cyan]Welcome[/bold cyan]", + border_style="", + box=box.SIMPLE, + padding=(1, 1), + ) + + +class ErrorPanel(Static): + """Error panel used for runtime/system errors.""" + + panel_type = "error" + + def __init__(self, text: str): + super().__init__() + self.error_text = text + + def render(self): + return Panel( + self.error_text, + title="[bold red]Error[/bold red]", + border_style="red", + box=box.SIMPLE, + style="on #3a1f1f", + padding=(0, 1), + ) + + class AgentView(VerticalScroll): """Scrollable per-agent chat view assembling message and tool panels. @@ -374,11 +415,19 @@ def __init__(self, agent_id: int, agent_name: str = ""): self.border_title = f"Agent {agent_id}" self.styles.border = ("solid", "cyan") self._streaming_panel: Optional[Static] = None + self._welcome_panel: Optional[WelcomePanel] = None - def on_mount(self): + async def on_mount(self): """Enable auto vertical scrolling with a stable scrollbar gutter.""" self.styles.overflow_y = "auto" self.styles.scrollbar_gutter = "stable" + self._welcome_panel = WelcomePanel() + await self.mount(self._welcome_panel) + + async def _hide_welcome(self): + if self._welcome_panel is not None: + await self._welcome_panel.remove() + self._welcome_panel = None def set_active(self, active: bool): """Highlight (or reset) the border indicating the focused agent.""" @@ -405,6 +454,7 @@ def _scroll_if_was_at_bottom(self, was_at_bottom: bool): async def append_user_message(self, text: str) -> None: """Mount a new user-message panel and scroll to it.""" + await self._hide_welcome() self._streaming_panel = None panel = UserMessagePanel(text) await self.mount(panel) @@ -417,6 +467,7 @@ async def append_plain_text(self, delta: str) -> None: plain text, otherwise appends in place; keeps the view pinned to the bottom only if the user was already there. """ + await self._hide_welcome() if self._last_panel_type != "plain_text": self._streaming_panel = None panel = PlainTextPanel() @@ -430,6 +481,7 @@ async def append_plain_text(self, delta: str) -> None: async def append_reasoning(self, delta: str) -> None: """Stream reasoning deltas into a shared reasoning-text panel.""" + await self._hide_welcome() if self._last_panel_type != "reasoning_text": self._streaming_panel = None panel = ReasoningTextPanel() @@ -450,6 +502,7 @@ async def start_tool_call( guessed from the target file extension; other actions get a compact one-line preview panel. """ + await self._hide_welcome() self._streaming_panel = None if action in ("file.write",): @@ -498,6 +551,14 @@ async def finish_tool_call(self, result: str, success: bool = True) -> None: self._streaming_panel = None self._scroll_to_bottom() + async def append_error(self, message: str) -> None: + """Mount a runtime error panel.""" + await self._hide_welcome() + self._streaming_panel = None + panel = ErrorPanel(message) + await self.mount(panel) + self._scroll_to_bottom() + class DynamicGrid(Grid): """Grid container arranging agent views: one large + up to three small. @@ -874,6 +935,8 @@ async def _handle_agent_event(self, agent_id: int, event: dict) -> None: result = event.get("result", "") success = event.get("success", False) await agent_view.finish_tool_call(result, success) + elif event_type == "error": + await agent_view.append_error(event.get("message", "Unknown error")) # ------------------------------------------------------------------ # Quick create & chat input diff --git a/test/test_main_entrypoint.py b/test/test_main_entrypoint.py new file mode 100644 index 0000000..3484e45 --- /dev/null +++ b/test/test_main_entrypoint.py @@ -0,0 +1,40 @@ +import os +import sys + +import pytest + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) + +from mitocode.__main__ import ( # noqa: E402 + _default_provider_kind, + _init_workspace, + _parse_provider_settings, +) + + +def test_parse_provider_settings_accepts_empty(): + assert _parse_provider_settings("") == {} + + +def test_parse_provider_settings_rejects_non_object(): + with pytest.raises(ValueError): + _parse_provider_settings('["bad"]') + + +def test_init_workspace_creates_then_loads(tmp_path): + root = tmp_path / "ws" + ws1 = _init_workspace(str(root), "demo") + assert ws1.config["name"] == "demo" + assert ws1.config_file.exists() + + ws2 = _init_workspace(str(root), "ignored") + assert ws2.config["name"] == "demo" + + +def test_default_provider_kind_uses_env(monkeypatch): + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + monkeypatch.delenv("OPENROUTER_API_KEY", raising=False) + assert _default_provider_kind() == "scenario" + + monkeypatch.setenv("OPENAI_API_KEY", "dummy") + assert _default_provider_kind() == "openai" diff --git a/test/test_orchestrator_and_bridge.py b/test/test_orchestrator_and_bridge.py index 00b4e6d..c0a3520 100644 --- a/test/test_orchestrator_and_bridge.py +++ b/test/test_orchestrator_and_bridge.py @@ -54,6 +54,29 @@ async def _stream(): return _stream() +class FailingProvider(_SampleChatProvider): + def __init__(self): + self.stream = True + + @staticmethod + def get_model_list(): + return ["failing-model"] + + @classmethod + def get_settings_schema(cls): + return {"type": "object", "properties": {}, "additionalProperties": True} + + @staticmethod + def check_settings_schema(settings): + return True + + def init(self): + pass + + async def chat(self, tools=None, messages=None, **kwargs): + raise RuntimeError("provider exploded") + + @pytest.mark.asyncio async def test_orchestrator_create_and_send_message(): orchestrator = AgentOrchestrator(default_provider=EchoProvider()) @@ -167,6 +190,24 @@ async def test_create_agent_alias_matches_create_session_id(): assert session.name == "Alias" +@pytest.mark.asyncio +async def test_orchestrator_emits_error_event_on_runtime_failure(): + orchestrator = AgentOrchestrator(default_provider=FailingProvider()) + agent_id = await orchestrator.create_agent(name="Broken") + task = await orchestrator.send_message(agent_id, "hello") + assert task is not None + await task + + got_error_event = False + while not orchestrator.event_queue.empty(): + aid, event = await orchestrator.event_queue.get() + if aid == agent_id and event.get("type") == "error": + got_error_event = True + assert "provider exploded" in event.get("message", "") + break + assert got_error_event is True + + def test_event_bridge_token(): bridge = EventBridge() result = bridge.adapt({"event": "agent.text", "delta": "hello"}, agent_id=1) From b28bff249f8e0ae8a6129f15bed9aa2194d75757 Mon Sep 17 00:00:00 2001 From: Wang Allen Date: Sat, 10 Oct 2026 10:41:37 +0800 Subject: [PATCH 02/20] feat(mcp): add MCP server management --- .../backend/agent/agent_orchestrator.py | 12 +--- src/mitocode/tui/ui_main.py | 65 +------------------ test/test_main_entrypoint.py | 40 ------------ test/test_orchestrator_and_bridge.py | 41 ------------ 4 files changed, 2 insertions(+), 156 deletions(-) delete mode 100644 test/test_main_entrypoint.py diff --git a/src/mitocode/backend/agent/agent_orchestrator.py b/src/mitocode/backend/agent/agent_orchestrator.py index 70a46f1..0f39318 100644 --- a/src/mitocode/backend/agent/agent_orchestrator.py +++ b/src/mitocode/backend/agent/agent_orchestrator.py @@ -264,18 +264,8 @@ async def _run_session(self, agent_id: int) -> None: except asyncio.CancelledError: session.state = AgentState.CANCELLED raise - except Exception as e: # pylint: disable=broad-exception-caught + except Exception: # pylint: disable=broad-exception-caught session.state = AgentState.ERROR - await self.event_queue.put( - ( - agent_id, - { - "type": "error", - "agent_id": agent_id, - "message": str(e) or "Unknown runtime error", - }, - ) - ) finally: if self._tasks.get(agent_id) is asyncio.current_task(): self._tasks[agent_id] = None diff --git a/src/mitocode/tui/ui_main.py b/src/mitocode/tui/ui_main.py index 27e8296..5aaad3c 100644 --- a/src/mitocode/tui/ui_main.py +++ b/src/mitocode/tui/ui_main.py @@ -360,47 +360,6 @@ def render(self): ) -class WelcomePanel(Static): - """Startup panel shown before first message is sent.""" - - panel_type = "welcome" - - def render(self): - return Panel( - RichMarkdown( - "## 欢迎使用 MitoCode TUI\n\n" - "- 在左侧选择 Agent\n" - "- 在下方输入框发送消息\n" - "- 使用 `Quick create` 快速创建子 Agent\n" - "- `F12` 打开调试日志" - ), - title="[bold cyan]Welcome[/bold cyan]", - border_style="", - box=box.SIMPLE, - padding=(1, 1), - ) - - -class ErrorPanel(Static): - """Error panel used for runtime/system errors.""" - - panel_type = "error" - - def __init__(self, text: str): - super().__init__() - self.error_text = text - - def render(self): - return Panel( - self.error_text, - title="[bold red]Error[/bold red]", - border_style="red", - box=box.SIMPLE, - style="on #3a1f1f", - padding=(0, 1), - ) - - class AgentView(VerticalScroll): """Scrollable per-agent chat view assembling message and tool panels. @@ -415,19 +374,11 @@ def __init__(self, agent_id: int, agent_name: str = ""): self.border_title = f"Agent {agent_id}" self.styles.border = ("solid", "cyan") self._streaming_panel: Optional[Static] = None - self._welcome_panel: Optional[WelcomePanel] = None - async def on_mount(self): + def on_mount(self): """Enable auto vertical scrolling with a stable scrollbar gutter.""" self.styles.overflow_y = "auto" self.styles.scrollbar_gutter = "stable" - self._welcome_panel = WelcomePanel() - await self.mount(self._welcome_panel) - - async def _hide_welcome(self): - if self._welcome_panel is not None: - await self._welcome_panel.remove() - self._welcome_panel = None def set_active(self, active: bool): """Highlight (or reset) the border indicating the focused agent.""" @@ -454,7 +405,6 @@ def _scroll_if_was_at_bottom(self, was_at_bottom: bool): async def append_user_message(self, text: str) -> None: """Mount a new user-message panel and scroll to it.""" - await self._hide_welcome() self._streaming_panel = None panel = UserMessagePanel(text) await self.mount(panel) @@ -467,7 +417,6 @@ async def append_plain_text(self, delta: str) -> None: plain text, otherwise appends in place; keeps the view pinned to the bottom only if the user was already there. """ - await self._hide_welcome() if self._last_panel_type != "plain_text": self._streaming_panel = None panel = PlainTextPanel() @@ -481,7 +430,6 @@ async def append_plain_text(self, delta: str) -> None: async def append_reasoning(self, delta: str) -> None: """Stream reasoning deltas into a shared reasoning-text panel.""" - await self._hide_welcome() if self._last_panel_type != "reasoning_text": self._streaming_panel = None panel = ReasoningTextPanel() @@ -502,7 +450,6 @@ async def start_tool_call( guessed from the target file extension; other actions get a compact one-line preview panel. """ - await self._hide_welcome() self._streaming_panel = None if action in ("file.write",): @@ -551,14 +498,6 @@ async def finish_tool_call(self, result: str, success: bool = True) -> None: self._streaming_panel = None self._scroll_to_bottom() - async def append_error(self, message: str) -> None: - """Mount a runtime error panel.""" - await self._hide_welcome() - self._streaming_panel = None - panel = ErrorPanel(message) - await self.mount(panel) - self._scroll_to_bottom() - class DynamicGrid(Grid): """Grid container arranging agent views: one large + up to three small. @@ -935,8 +874,6 @@ async def _handle_agent_event(self, agent_id: int, event: dict) -> None: result = event.get("result", "") success = event.get("success", False) await agent_view.finish_tool_call(result, success) - elif event_type == "error": - await agent_view.append_error(event.get("message", "Unknown error")) # ------------------------------------------------------------------ # Quick create & chat input diff --git a/test/test_main_entrypoint.py b/test/test_main_entrypoint.py deleted file mode 100644 index 3484e45..0000000 --- a/test/test_main_entrypoint.py +++ /dev/null @@ -1,40 +0,0 @@ -import os -import sys - -import pytest - -sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) - -from mitocode.__main__ import ( # noqa: E402 - _default_provider_kind, - _init_workspace, - _parse_provider_settings, -) - - -def test_parse_provider_settings_accepts_empty(): - assert _parse_provider_settings("") == {} - - -def test_parse_provider_settings_rejects_non_object(): - with pytest.raises(ValueError): - _parse_provider_settings('["bad"]') - - -def test_init_workspace_creates_then_loads(tmp_path): - root = tmp_path / "ws" - ws1 = _init_workspace(str(root), "demo") - assert ws1.config["name"] == "demo" - assert ws1.config_file.exists() - - ws2 = _init_workspace(str(root), "ignored") - assert ws2.config["name"] == "demo" - - -def test_default_provider_kind_uses_env(monkeypatch): - monkeypatch.delenv("OPENAI_API_KEY", raising=False) - monkeypatch.delenv("OPENROUTER_API_KEY", raising=False) - assert _default_provider_kind() == "scenario" - - monkeypatch.setenv("OPENAI_API_KEY", "dummy") - assert _default_provider_kind() == "openai" diff --git a/test/test_orchestrator_and_bridge.py b/test/test_orchestrator_and_bridge.py index c0a3520..00b4e6d 100644 --- a/test/test_orchestrator_and_bridge.py +++ b/test/test_orchestrator_and_bridge.py @@ -54,29 +54,6 @@ async def _stream(): return _stream() -class FailingProvider(_SampleChatProvider): - def __init__(self): - self.stream = True - - @staticmethod - def get_model_list(): - return ["failing-model"] - - @classmethod - def get_settings_schema(cls): - return {"type": "object", "properties": {}, "additionalProperties": True} - - @staticmethod - def check_settings_schema(settings): - return True - - def init(self): - pass - - async def chat(self, tools=None, messages=None, **kwargs): - raise RuntimeError("provider exploded") - - @pytest.mark.asyncio async def test_orchestrator_create_and_send_message(): orchestrator = AgentOrchestrator(default_provider=EchoProvider()) @@ -190,24 +167,6 @@ async def test_create_agent_alias_matches_create_session_id(): assert session.name == "Alias" -@pytest.mark.asyncio -async def test_orchestrator_emits_error_event_on_runtime_failure(): - orchestrator = AgentOrchestrator(default_provider=FailingProvider()) - agent_id = await orchestrator.create_agent(name="Broken") - task = await orchestrator.send_message(agent_id, "hello") - assert task is not None - await task - - got_error_event = False - while not orchestrator.event_queue.empty(): - aid, event = await orchestrator.event_queue.get() - if aid == agent_id and event.get("type") == "error": - got_error_event = True - assert "provider exploded" in event.get("message", "") - break - assert got_error_event is True - - def test_event_bridge_token(): bridge = EventBridge() result = bridge.adapt({"event": "agent.text", "delta": "hello"}, agent_id=1) From 5cbed57d5f2f86562dc4c32f58aab04f8422fce0 Mon Sep 17 00:00:00 2001 From: Wang Allen Date: Wed, 14 Oct 2026 10:11:20 +0800 Subject: [PATCH 03/20] refactor: simplify __main__.py --- src/mitocode/__main__.py | 219 ++++++++++----------------------------- 1 file changed, 53 insertions(+), 166 deletions(-) diff --git a/src/mitocode/__main__.py b/src/mitocode/__main__.py index 12936ed..0a0fa30 100644 --- a/src/mitocode/__main__.py +++ b/src/mitocode/__main__.py @@ -1,209 +1,96 @@ -"""CLI entrypoint for running the MitoCode TUI.""" +"""WIP""" import argparse import asyncio -import json -import os -from pathlib import Path -from typing import Any, Dict +from pyexpat.errors import messages + +from httpx import stream -from mitocode import APP_VERSION, MITOCODE_CONFIG_PATH from mitocode.backend.agent.agent_orchestrator import AgentOrchestrator -from mitocode.backend.chat_provider.chat_provider_openai import ChatProviderOpenAI from mitocode.backend.chat_provider.scenario_provider import ScenarioProvider -from mitocode.backend.chat_provider.sample_chat_provider import SampleChatProvider +from mitocode.backend.chat_provider.chat_provider_openai import ChatProviderOpenAI from mitocode.backend.tool.tool_manager import ToolManager -from mitocode.backend.workspace.workspace import Workspace from mitocode.tui.ui_main import VibeFrontend -def _register_tools(tool_manager: ToolManager, workspace: Workspace) -> None: - """Register workspace-backed tools (bash, read_file, write_file).""" +def _register_tools(tool_manager: ToolManager): + """Register built-in demo tools (bash, read_file, write_file) on the manager.""" @tool_manager.tool( - human_doc="Run a shell command inside workspace", - agent_doc="Execute a shell command in the workspace directory and return output", + human_doc="Run a shell command", + agent_doc="Execute a bash command and return its output", display_action="bash", ) async def bash(command: str) -> str: - proc = await asyncio.create_subprocess_shell( - command, - cwd=str(workspace.path), - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - stdout_bytes, stderr_bytes = await proc.communicate() - stdout = stdout_bytes.decode("utf-8", errors="replace") - stderr = stderr_bytes.decode("utf-8", errors="replace") - output = stdout if stdout else "(no stdout)" - if stderr: - output += f"\n[stderr]\n{stderr}" - return ( - f"$ {command}\n[cwd] {workspace.path}\n[exit_code] {proc.returncode}\n{output}" - ) + return f"$ {command}\n[bash] command executed successfully" @tool_manager.tool( human_doc="Read file content", - agent_doc="Read UTF-8 text file content from a workspace-relative path", + agent_doc="Read the full content of a file at the given path", display_action="file.read", display_param="file_path", ) async def read_file(file_path: str) -> str: - safe_path = workspace._validate_within_workspace(Path(file_path)) - if not safe_path.exists(): - raise FileNotFoundError(f"File not found: {safe_path}") - if not safe_path.is_file(): - raise IsADirectoryError(f"Path is not a file: {safe_path}") - return await asyncio.to_thread(safe_path.read_text, encoding="utf-8") + return ( + "def bubble_sort(arr):\n" + " n = len(arr)\n" + " for i in range(n):\n" + " swapped = False\n" + " for j in range(0, n - i - 1):\n" + " if arr[j] > arr[j + 1]:\n" + " arr[j], arr[j + 1] = arr[j + 1], arr[j]\n" + " swapped = True\n" + " if not swapped:\n" + " break\n" + " return arr\n" + ) @tool_manager.tool( human_doc="Write content to a file", - agent_doc="Write UTF-8 content into a workspace-relative file path", + agent_doc="Create or overwrite a file with the given content", display_action="file.write", display_param="file_path", ) async def write_file(file_path: str, content: str) -> str: - safe_path = workspace._validate_within_workspace(Path(file_path)) - await asyncio.to_thread(safe_path.parent.mkdir, parents=True, exist_ok=True) - await asyncio.to_thread(safe_path.write_text, content, encoding="utf-8") - return f"Successfully wrote {len(content)} bytes to {safe_path}" - - -def _parse_provider_settings(raw: str) -> Dict[str, Any]: - """Parse user-provided provider settings JSON string.""" - if not raw: - return {} - parsed = json.loads(raw) - if not isinstance(parsed, dict): - raise ValueError("--provider-settings must be a JSON object") - return parsed - - -def _default_provider_kind() -> str: - return ( - "openai" - if os.getenv("OPENAI_API_KEY") or os.getenv("OPENROUTER_API_KEY") - else "scenario" - ) + return f"Successfully wrote {len(content)} bytes to {file_path}" -def _build_provider(args: argparse.Namespace) -> SampleChatProvider: - """Build and initialize the provider from CLI arguments.""" - if args.provider == "scenario": - settings = {"delay": args.speed} - settings.update(_parse_provider_settings(args.provider_settings)) - return asyncio.run( - ScenarioProvider.factory( - default_model=args.model or "scenario-model", - default_model_list=["scenario-model"], - default_messages=None, - default_tools=None, - stream=True, - settings=settings, - delay=settings.get("delay", args.speed), - ) - ) - - api_key = args.api_key or os.getenv("OPENAI_API_KEY") or os.getenv("OPENROUTER_API_KEY") - if not api_key: - raise ValueError("OpenAI provider requires --api-key or OPENAI_API_KEY/OPENROUTER_API_KEY") - - settings: Dict[str, Any] = { - "api_key": api_key, - "timeout": args.timeout, - } - if args.base_url: - settings["base_url"] = args.base_url - settings.update(_parse_provider_settings(args.provider_settings)) - - default_model = args.model or "gpt-4o-mini" - return asyncio.run( - ChatProviderOpenAI.factory( - default_model=default_model, - default_model_list=[default_model], - default_messages=None, - default_tools=None, - stream=True, - settings=settings, - ) - ) - - -def _init_workspace(path_arg: str, name: str) -> Workspace: - """Initialize or create workspace at the given path.""" - workspace_root = Path(path_arg).expanduser().resolve() - workspace_root.mkdir(parents=True, exist_ok=True) - - config_file = workspace_root / MITOCODE_CONFIG_PATH / "config.json" - if config_file.exists(): - return Workspace.init_from_path(workspace_root) - - return Workspace.new_workspace( - workspace_path=workspace_root, - name=name, - description=f"MitoCode workspace (app v{APP_VERSION})", - override=False, - ) - - -def main() -> None: +def main(): """Parse CLI arguments and launch the MitoCode TUI application.""" parser = argparse.ArgumentParser(description="MitoCode") parser.add_argument( "--speed", type=float, default=0.03, - help="Delay between scenario provider chunks in seconds (default: 0.03)", - ) - parser.add_argument( - "--workspace", - type=str, - default=".", - help="Workspace root path (default: current directory)", - ) - parser.add_argument( - "--workspace-name", - type=str, - default="mitocode-workspace", - help="Workspace name for first-time initialization", - ) - parser.add_argument( - "--provider", - choices=["scenario", "openai"], - default=_default_provider_kind(), - help="LLM provider backend", - ) - parser.add_argument("--model", type=str, default="", help="Provider model name") - parser.add_argument( - "--api-key", - type=str, - default="", - help="OpenAI/OpenRouter API key (or use OPENAI_API_KEY / OPENROUTER_API_KEY)", - ) - parser.add_argument( - "--base-url", - type=str, - default=os.getenv("OPENAI_BASE_URL", ""), - help="OpenAI-compatible API base URL override", - ) - parser.add_argument( - "--timeout", - type=float, - default=60, - help="Provider request timeout in seconds", - ) - parser.add_argument( - "--provider-settings", - type=str, - default="", - help="Extra provider settings as JSON object", + help="Delay between chunks in seconds (default: 0.03)", ) args = parser.parse_args() - workspace = _init_workspace(args.workspace, args.workspace_name) tool_manager = ToolManager() - _register_tools(tool_manager, workspace) - provider = _build_provider(args) + _register_tools(tool_manager) + + provider = asyncio.run(ChatProviderOpenAI.factory( + default_model= "nvidia/nemotron-3.5-lightning:free", + default_model_list= None, + default_messages= None, + default_tools= None, + stream= True, + settings={ + "api_key": "sk-别看了这次我不会再把key提交到仓库了！", + "base_url": "https://openrouter.ai/api/v1", + "timeout": 60, + }) + ) + + # print(asyncio.run(provider.chat( + # model ="nvidia/nemotron-3.5-lightning:free", + # tools =None, + # messages = [ + # {"role": "system", "content": "You are a helpful assistant."}, + # {"role": "user", "content": "Hello!"} + # ], + # stream = False + # ))) orchestrator = AgentOrchestrator( tool_manager=tool_manager, default_provider=provider @@ -211,7 +98,7 @@ def main() -> None: asyncio.run( orchestrator.create_agent( name="Main", - description=f"Workspace: {workspace.path}", + description="Chief orchestrator", ) ) From be1d90937046906376418e07884e84a00430d5c5 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 14 Oct 2026 02:44:04 +0000 Subject: [PATCH 04/20] ci: auto-format code --- src/mitocode/__main__.py | 24 ++++++++++--------- .../chat_provider/scenario_provider.py | 1 - 2 files changed, 13 insertions(+), 12 deletions(-) diff --git a/src/mitocode/__main__.py b/src/mitocode/__main__.py index 0a0fa30..f511246 100644 --- a/src/mitocode/__main__.py +++ b/src/mitocode/__main__.py @@ -69,17 +69,19 @@ def main(): tool_manager = ToolManager() _register_tools(tool_manager) - provider = asyncio.run(ChatProviderOpenAI.factory( - default_model= "nvidia/nemotron-3.5-lightning:free", - default_model_list= None, - default_messages= None, - default_tools= None, - stream= True, - settings={ - "api_key": "sk-别看了这次我不会再把key提交到仓库了！", - "base_url": "https://openrouter.ai/api/v1", - "timeout": 60, - }) + provider = asyncio.run( + ChatProviderOpenAI.factory( + default_model="nvidia/nemotron-3.5-lightning:free", + default_model_list=None, + default_messages=None, + default_tools=None, + stream=True, + settings={ + "api_key": "sk-别看了这次我不会再把key提交到仓库了！", + "base_url": "https://openrouter.ai/api/v1", + "timeout": 60, + }, + ) ) # print(asyncio.run(provider.chat( diff --git a/src/mitocode/backend/chat_provider/scenario_provider.py b/src/mitocode/backend/chat_provider/scenario_provider.py index e1fd5c0..fdb0804 100644 --- a/src/mitocode/backend/chat_provider/scenario_provider.py +++ b/src/mitocode/backend/chat_provider/scenario_provider.py @@ -12,7 +12,6 @@ from mitocode.backend.chat_provider.sample_chat_provider import _SampleChatProvider - def _make_chunk( content: Optional[str] = None, tool_calls: Optional[List[Any]] = None, From 43e6382901ecaed86e15e4ee02e8a6efdaf776f0 Mon Sep 17 00:00:00 2001 From: Allen Wang Date: Fri, 16 Oct 2026 23:30:16 +0800 Subject: [PATCH 05/20] feat(mcp): add MCP server config parsing and exceptions Normalises Claude Desktop style MCP server entries into a typed McpServerConfig, validating that url and command are mutually exclusive and building the matching StreamableHttp or Stdio transport object. Constraint: Config arrives as a parameter; this layer never touches the filesystem Rejected: fastmcp-slim[client] vs fastmcp split | fastmcp-slim already ships the full server code in 4.1, the split is only extras Confidence: high Scope-risk: narrow --- pyproject.toml | 1 + requirements.txt | 3 +- .../backend/exception/mitocode_exception.py | 16 + src/mitocode/backend/mcp/__init__.py | 1 + src/mitocode/backend/mcp/mcp_config.py | 354 ++++++++++++++++++ test/test_mcp_config.py | 119 ++++++ 6 files changed, 493 insertions(+), 1 deletion(-) create mode 100644 src/mitocode/backend/mcp/__init__.py create mode 100644 src/mitocode/backend/mcp/mcp_config.py create mode 100644 test/test_mcp_config.py diff --git a/pyproject.toml b/pyproject.toml index 813d958..c10bfb4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,7 @@ dependencies = [ "httpx-sse", "textual", "fastapi", + "fastmcp>=4.1,<5", ] [project.optional-dependencies] diff --git a/requirements.txt b/requirements.txt index ed3cb8f..d8458a9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -12,4 +12,5 @@ httpx httpx-sse textual fastapi -jsonschema \ No newline at end of file +jsonschema +fastmcp>=4.1,<5 \ No newline at end of file diff --git a/src/mitocode/backend/exception/mitocode_exception.py b/src/mitocode/backend/exception/mitocode_exception.py index 086fab9..e75c8de 100644 --- a/src/mitocode/backend/exception/mitocode_exception.py +++ b/src/mitocode/backend/exception/mitocode_exception.py @@ -74,3 +74,19 @@ class MitoCodeChatProviderModelError(MitoCodeChatProviderError): class MitoCodeChatProviderInitError(MitoCodeChatProviderError): """Raised when chat provider initialization fails.""" + + +class MitoCodeMcpError(MitoCodeBaseError): + """Base class for MCP-related errors.""" + + +class MitoCodeMcpConfigError(MitoCodeMcpError): + """Raised when an MCP server configuration entry is invalid.""" + + +class MitoCodeMcpConnectionError(MitoCodeMcpError): + """Raised when an MCP server connection cannot be established or is lost.""" + + +class MitoCodeMcpToolCallError(MitoCodeMcpError): + """Raised when an MCP tool call fails or returns an error result.""" diff --git a/src/mitocode/backend/mcp/__init__.py b/src/mitocode/backend/mcp/__init__.py new file mode 100644 index 0000000..48bdae8 --- /dev/null +++ b/src/mitocode/backend/mcp/__init__.py @@ -0,0 +1 @@ +"""MCP client support: configuration, connections, and lifecycle management.""" diff --git a/src/mitocode/backend/mcp/mcp_config.py b/src/mitocode/backend/mcp/mcp_config.py new file mode 100644 index 0000000..8faf27b --- /dev/null +++ b/src/mitocode/backend/mcp/mcp_config.py @@ -0,0 +1,354 @@ +"""Normalisation and validation of MCP server configuration entries. + +The Claude Desktop configuration shape — ``{"mcpServers": {"": {...}}}`` +— is the accepted input format. An entry with a ``url`` describes a remote +server reached over Streamable HTTP; an entry with a ``command`` describes a +local server launched as a subprocess and spoken to over stdio. + +This layer only normalises configuration. It never touches the filesystem and +never opens a connection: :meth:`McpServerConfig.build_transport` merely +assembles the transport object that a later layer hands to a live client. +""" + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Dict, List, NoReturn + +from mitocode.backend.exception.mitocode_exception import MitoCodeMcpConfigError + +try: + from fastmcp.client.transports import ( + ClientTransport, + StdioTransport, + StreamableHttpTransport, + ) +except ImportError as exc: # pragma: no cover - fastmcp is a hard dependency + raise ImportError( + "MCP support requires the 'fastmcp' package; " + "install it with 'pip install fastmcp>=4.1,<5'" + ) from exc + + +class McpServerKind(str, Enum): + """Transport used to reach an MCP server.""" + + STDIO = "stdio" + HTTP = "http" + + +@dataclass(frozen=True) +class McpServerConfig: # pylint: disable=too-many-instance-attributes + """A validated MCP server definition. + + Attributes: + name: Unique server name used to namespace its tools. + kind: Which transport this server is reached over. + url: Streamable HTTP endpoint; empty for stdio servers. + command: Executable to launch; empty for HTTP servers. + args: Command-line arguments passed to ``command``. + env: Extra environment variables for the subprocess. + cwd: Working directory for the subprocess. + headers: Extra HTTP headers, typically carrying a static token. + timeout: Per-call timeout in seconds. + connect_timeout: Seconds allowed for the connection handshake. + Generous by default because a cold ``npx -y`` start can take + tens of seconds. + enabled: Whether the manager should attempt to connect. + """ + + name: str + kind: McpServerKind + url: str = "" + command: str = "" + args: List[str] = field(default_factory=list) + env: Dict[str, str] = field(default_factory=dict) + cwd: str = "" + headers: Dict[str, str] = field(default_factory=dict) + timeout: float = 60.0 + connect_timeout: float = 30.0 + enabled: bool = True + + @classmethod + def parse(cls, name: str, raw: dict) -> "McpServerConfig": + """Normalise one raw server entry into a validated configuration. + + Every field is type-checked but only the ones that belong to the + resolved transport are meaningful; entries that carry a stray field for + the other transport are accepted and simply carried along, because + being lenient about shape keeps hand-written configuration usable. + + Args: + name: Unique server name, normally the key of the ``mcpServers`` + mapping. + raw: One server entry. It must carry exactly one of ``url`` (a + remote server) or ``command`` (a local stdio server). + + Returns: + A frozen configuration with every field populated. + + Raises: + MitoCodeMcpConfigError: If ``name`` is blank or not a string, if + ``raw`` is not a mapping, if the transport selector is missing, + ambiguous or malformed, or if any optional field has the wrong + type. + """ + if not isinstance(name, str) or not name.strip(): + raise MitoCodeMcpConfigError( + f"MCP server name must be a non-empty string, got {name!r}" + ) + if not isinstance(raw, dict): + raise MitoCodeMcpConfigError( + f"MCP server '{name}': entry must be an object, " + f"got {type(raw).__name__}" + ) + + has_url = "url" in raw + has_command = "command" in raw + if has_url and has_command: + raise MitoCodeMcpConfigError( + f"MCP server '{name}': 'url' and 'command' are mutually " + "exclusive, provide exactly one" + ) + if not has_url and not has_command: + raise MitoCodeMcpConfigError( + f"MCP server '{name}': either 'url' for a remote server or " + f"'command' for a stdio server is required, got keys {sorted(map(str, raw))}" + ) + + url = "" + command = "" + if has_url: + kind = McpServerKind.HTTP + url = _read_http_url(name, raw["url"]) + else: + kind = McpServerKind.STDIO + command = _read_stdio_command(name, raw["command"]) + + enabled = raw.get("enabled", True) + if not isinstance(enabled, bool): + _invalid(name, "enabled", "a boolean", enabled) + + return cls( + name=name, + kind=kind, + url=url, + command=command, + args=_read_str_list(name, raw, "args"), + env=_read_str_mapping(name, raw, "env"), + cwd=_read_str(name, raw, "cwd"), + headers=_read_str_mapping(name, raw, "headers"), + timeout=_read_positive_number(name, raw, "timeout", 60.0), + connect_timeout=_read_positive_number(name, raw, "connect_timeout", 30.0), + enabled=enabled, + ) + + @classmethod + def parse_config(cls, config: dict) -> List["McpServerConfig"]: + """Normalise a whole ``{"mcpServers": {...}}`` document. + + Args: + config: The full configuration mapping. + + Returns: + One validated configuration per entry, in document order. + + Raises: + MitoCodeMcpConfigError: If ``config`` is not a mapping, has no + ``"mcpServers"`` key, or if that key does not hold a mapping. + Individual entries are reported by :meth:`parse`. + """ + if not isinstance(config, dict): + raise MitoCodeMcpConfigError( + "MCP configuration must be an object holding an 'mcpServers' " + f"mapping, got {type(config).__name__}" + ) + if "mcpServers" not in config: + raise MitoCodeMcpConfigError( + "MCP configuration is missing the 'mcpServers' key, " + f"got keys {sorted(map(str, config))}" + ) + servers = config["mcpServers"] + if not isinstance(servers, dict): + raise MitoCodeMcpConfigError( + "MCP configuration 'mcpServers' must be an object, " + f"got {type(servers).__name__}" + ) + return [cls.parse(name, raw) for name, raw in servers.items()] + + def build_transport(self) -> ClientTransport: + """Build the fastmcp transport object matching this server's kind. + + No I/O happens here — the transport only carries connection parameters + until a client actually uses it. + + Returns: + A ``StreamableHttpTransport`` for remote servers or a + ``StdioTransport`` for local ones. Empty optional mappings are + passed as ``None`` because that is how fastmcp spells "unset". + + Raises: + MitoCodeMcpConfigError: If ``kind`` is not a known transport kind. + """ + if self.kind is McpServerKind.HTTP: + return StreamableHttpTransport(url=self.url, headers=self.headers or None) + if self.kind is McpServerKind.STDIO: + return StdioTransport( + command=self.command, + args=self.args, + env=self.env or None, + cwd=self.cwd or None, + ) + raise MitoCodeMcpConfigError( + f"MCP server '{self.name}': unknown transport kind {self.kind!r}" + ) + + +def _invalid(name: str, key: str, expected: str, value: Any) -> NoReturn: + """Report one configuration field that failed validation. + + Args: + name: Name of the offending MCP server. + key: The configuration key that failed validation. + expected: Human-readable description of the accepted shape. + value: The offending value. + + Raises: + MitoCodeMcpConfigError: Always. + """ + raise MitoCodeMcpConfigError( + f"MCP server '{name}': '{key}' must be {expected}, got {value!r}" + ) + + +def _read_http_url(name: str, value: Any) -> str: + """Validate the ``url`` field of a remote server entry. + + Args: + name: Name of the MCP server, used in the error message. + value: The raw ``url`` value. + + Returns: + The validated endpoint. + + Raises: + MitoCodeMcpConfigError: If ``value`` is not an http(s) endpoint. + """ + if not isinstance(value, str) or not value.startswith(("http://", "https://")): + raise MitoCodeMcpConfigError( + f"MCP server '{name}': 'url' must be an http:// or https:// " + f"endpoint, got {value!r}" + ) + return value + + +def _read_stdio_command(name: str, value: Any) -> str: + """Validate the ``command`` field of a local server entry. + + Args: + name: Name of the MCP server, used in the error message. + value: The raw ``command`` value. + + Returns: + The validated executable name. + + Raises: + MitoCodeMcpConfigError: If ``value`` is not a non-empty string. + """ + if not isinstance(value, str) or not value.strip(): + raise MitoCodeMcpConfigError( + f"MCP server '{name}': 'command' must be a non-empty string, " + f"got {value!r}" + ) + return value + + +def _read_str(name: str, raw: dict, key: str) -> str: + """Read an optional string field, defaulting to the empty string. + + Args: + name: Name of the MCP server, used in the error message. + raw: The raw server entry. + key: The configuration key to read. + + Returns: + The field value, or ``""`` when the key is absent. + + Raises: + MitoCodeMcpConfigError: If the value is present but not a string. + """ + value = raw.get(key, "") + if not isinstance(value, str): + _invalid(name, key, "a string", value) + return value + + +def _read_str_list(name: str, raw: dict, key: str) -> List[str]: + """Read an optional list-of-strings field, defaulting to an empty list. + + Args: + name: Name of the MCP server, used in the error message. + raw: The raw server entry. + key: The configuration key to read. + + Returns: + A copy of the field value, so later mutation of the caller's list + cannot reach the frozen configuration. + + Raises: + MitoCodeMcpConfigError: If the value is not a list of strings. + """ + value = raw.get(key, []) + if not isinstance(value, list) or not all(isinstance(item, str) for item in value): + _invalid(name, key, "a list of strings", value) + return list(value) + + +def _read_str_mapping(name: str, raw: dict, key: str) -> Dict[str, str]: + """Read an optional string-to-string mapping, defaulting to an empty dict. + + Args: + name: Name of the MCP server, used in the error message. + raw: The raw server entry. + key: The configuration key to read. + + Returns: + A copy of the field value, so later mutation of the caller's mapping + cannot reach the frozen configuration. + + Raises: + MitoCodeMcpConfigError: If the value is not a mapping of strings. + """ + value = raw.get(key, {}) + if not isinstance(value, dict): + _invalid(name, key, "an object mapping strings to strings", value) + for entry_key, entry_value in value.items(): + if not isinstance(entry_key, str) or not isinstance(entry_value, str): + _invalid(name, key, "an object mapping strings to strings", value) + return dict(value) + + +def _read_positive_number(name: str, raw: dict, key: str, default: float) -> float: + """Read an optional strictly positive number field. + + ``bool`` is rejected explicitly: it is a subclass of ``int`` in Python, and + ``timeout: true`` is a configuration mistake rather than a timeout of 1. + The comparison is spelled ``not value > 0`` rather than ``value <= 0`` so a + NaN timeout is rejected too — it is not greater than zero. + + Args: + name: Name of the MCP server, used in the error message. + raw: The raw server entry. + key: The configuration key to read. + default: Value used when the key is absent. + + Returns: + The field value as a float. + + Raises: + MitoCodeMcpConfigError: If the value is a boolean, is not a number, or + is not greater than zero. + """ + value = raw.get(key, default) + if isinstance(value, bool) or not isinstance(value, (int, float)) or not value > 0: + _invalid(name, key, "a positive number", value) + return float(value) diff --git a/test/test_mcp_config.py b/test/test_mcp_config.py new file mode 100644 index 0000000..0def97e --- /dev/null +++ b/test/test_mcp_config.py @@ -0,0 +1,119 @@ +import os +import sys + +import pytest + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) + +from mitocode.backend.exception.mitocode_exception import MitoCodeMcpConfigError +from mitocode.backend.mcp.mcp_config import McpServerConfig, McpServerKind + + +def test_mcp_config_error_is_a_mitocode_base_error(): + from mitocode.backend.exception.mitocode_exception import MitoCodeMcpError + + assert issubclass(MitoCodeMcpConfigError, MitoCodeMcpError) + + +def test_parse_http_server_config(): + config = McpServerConfig.parse( + "weather", {"url": "https://weather.example.com/mcp"} + ) + assert config.name == "weather" + assert config.kind is McpServerKind.HTTP + assert config.url == "https://weather.example.com/mcp" + assert config.timeout == 60.0 + assert config.enabled is True + + +def test_parse_stdio_server_config_with_env_and_cwd(): + config = McpServerConfig.parse( + "files", + { + "command": "npx", + "args": ["-y", "@mcp/files"], + "env": {"TOKEN": "x"}, + "cwd": "/tmp", + }, + ) + assert config.kind is McpServerKind.STDIO + assert config.command == "npx" + assert config.args == ["-y", "@mcp/files"] + assert config.env == {"TOKEN": "x"} + assert config.cwd == "/tmp" + + +def test_parse_rejects_both_url_and_command(): + with pytest.raises(MitoCodeMcpConfigError): + McpServerConfig.parse("bad", {"url": "https://x/mcp", "command": "npx"}) + + +def test_parse_rejects_neither_url_nor_command(): + with pytest.raises(MitoCodeMcpConfigError): + McpServerConfig.parse("bad", {"env": {}}) + + +def test_parse_rejects_non_http_url(): + with pytest.raises(MitoCodeMcpConfigError): + McpServerConfig.parse("bad", {"url": "ftp://x/mcp"}) + + +def test_parse_rejects_non_string_argv(): + with pytest.raises(MitoCodeMcpConfigError): + McpServerConfig.parse("bad", {"command": "npx", "args": ["ok", 3]}) + + +def test_parse_rejects_non_positive_timeout(): + with pytest.raises(MitoCodeMcpConfigError): + McpServerConfig.parse("bad", {"url": "https://x/mcp", "timeout": 0}) + + +def test_parse_defaults_timeout_and_connect_timeout(): + config = McpServerConfig.parse("a", {"url": "https://x/mcp"}) + assert config.timeout == 60.0 + assert config.connect_timeout == 30.0 + assert config.enabled is True + + +def test_parse_rejects_non_positive_connect_timeout(): + with pytest.raises(MitoCodeMcpConfigError): + McpServerConfig.parse("a", {"url": "https://x/mcp", "connect_timeout": -1}) + + +def test_parse_config_reads_mcp_servers_wrapper(): + configs = McpServerConfig.parse_config( + { + "mcpServers": { + "weather": {"url": "https://weather.example.com/mcp"}, + "files": {"command": "npx", "args": ["-y", "files-mcp"]}, + } + } + ) + assert [c.name for c in configs] == ["weather", "files"] + assert [c.kind for c in configs] == [McpServerKind.HTTP, McpServerKind.STDIO] + + +def test_parse_config_rejects_missing_mcp_servers_key(): + with pytest.raises(MitoCodeMcpConfigError): + McpServerConfig.parse_config({"weather": {"url": "https://x/mcp"}}) + + +def test_build_transport_returns_http_transport(): + from fastmcp.client.transports import StreamableHttpTransport + + config = McpServerConfig.parse( + "weather", {"url": "https://weather.example.com/mcp", "headers": {"A": "b"}} + ) + transport = config.build_transport() + assert isinstance(transport, StreamableHttpTransport) + assert transport.url == "https://weather.example.com/mcp" + + +def test_build_transport_returns_stdio_transport(): + from fastmcp.client.transports import StdioTransport + + config = McpServerConfig.parse("files", {"command": "npx", "args": ["-y", "f"]}) + transport = config.build_transport() + assert isinstance(transport, StdioTransport) + assert transport.command == "npx" + assert transport.args == ["-y", "f"] From 51afe9d70f3c2c7adc036d957fbe2b7bc1034d50 Mon Sep 17 00:00:00 2001 From: Allen Wang Date: Fri, 16 Oct 2026 23:35:09 +0800 Subject: [PATCH 06/20] fix(mcp): never echo credential values into config errors env and headers routinely carry static tokens, and a single malformed entry caused the whole mapping to be rendered into the exception message, which then lands in the TUI error pane and in logs. Report only the offending entry's key. Auditing the remaining helpers found the same pattern in url (a credential can sit in userinfo or the query string) and args (command-line flags such as --token ), so both now report a scheme, an index or a type name instead of the value. That is also the more actionable message: 'args[3]' beats a dump of the whole argv. Constraint: headers/env must never appear by value in any exception message Rejected: type-name-only messages for timeout/cwd/command too | those fields are not credential-bearing by definition, and reporting 'got 0' or 'got -1' is what makes a numeric mistake fixable Confidence: high Scope-risk: narrow --- src/mitocode/backend/mcp/mcp_config.py | 64 +++++++++++++++++++++----- test/test_mcp_config.py | 58 +++++++++++++++++++++++ 2 files changed, 110 insertions(+), 12 deletions(-) diff --git a/src/mitocode/backend/mcp/mcp_config.py b/src/mitocode/backend/mcp/mcp_config.py index 8faf27b..dc29920 100644 --- a/src/mitocode/backend/mcp/mcp_config.py +++ b/src/mitocode/backend/mcp/mcp_config.py @@ -206,6 +206,10 @@ def build_transport(self) -> ClientTransport: def _invalid(name: str, key: str, expected: str, value: Any) -> NoReturn: """Report one configuration field that failed validation. + ``value`` is rendered into the message, so this helper is only for fields + that cannot carry credentials. Secret-bearing fields must raise their own + error naming just the offending entry's key — see :func:`_read_str_mapping`. + Args: name: Name of the offending MCP server. key: The configuration key that failed validation. @@ -223,6 +227,10 @@ def _invalid(name: str, key: str, expected: str, value: Any) -> NoReturn: def _read_http_url(name: str, value: Any) -> str: """Validate the ``url`` field of a remote server entry. + Only the scheme is reported on failure. A URL may legitimately carry a + credential in its userinfo or query string, so the rest of the value must + not reach an exception message either. + Args: name: Name of the MCP server, used in the error message. value: The raw ``url`` value. @@ -233,12 +241,18 @@ def _read_http_url(name: str, value: Any) -> str: Raises: MitoCodeMcpConfigError: If ``value`` is not an http(s) endpoint. """ - if not isinstance(value, str) or not value.startswith(("http://", "https://")): - raise MitoCodeMcpConfigError( - f"MCP server '{name}': 'url' must be an http:// or https:// " - f"endpoint, got {value!r}" - ) - return value + if isinstance(value, str): + if value.startswith(("http://", "https://")): + return value + if "://" in value: + detail = f"got scheme {value.split('://', 1)[0]!r}" + else: + detail = "got no http:// or https:// scheme" + else: + detail = f"got {type(value).__name__}" + raise MitoCodeMcpConfigError( + f"MCP server '{name}': 'url' must be an http:// or https:// endpoint, {detail}" + ) def _read_stdio_command(name: str, value: Any) -> str: @@ -285,6 +299,10 @@ def _read_str(name: str, raw: dict, key: str) -> str: def _read_str_list(name: str, raw: dict, key: str) -> List[str]: """Read an optional list-of-strings field, defaulting to an empty list. + A malformed element is reported by index and type, never by value. ``args`` + routinely carries command-line credentials such as ``--token ``, and + a message naming the offending position is what makes the mistake fixable. + Args: name: Name of the MCP server, used in the error message. raw: The raw server entry. @@ -295,17 +313,32 @@ def _read_str_list(name: str, raw: dict, key: str) -> List[str]: cannot reach the frozen configuration. Raises: - MitoCodeMcpConfigError: If the value is not a list of strings. + MitoCodeMcpConfigError: If the value is not a list of strings. The + message names the field and the bad element's index, but no values. """ value = raw.get(key, []) - if not isinstance(value, list) or not all(isinstance(item, str) for item in value): - _invalid(name, key, "a list of strings", value) + if not isinstance(value, list): + raise MitoCodeMcpConfigError( + f"MCP server '{name}': '{key}' must be a list of strings, " + f"got {type(value).__name__}" + ) + for index, item in enumerate(value): + if not isinstance(item, str): + raise MitoCodeMcpConfigError( + f"MCP server '{name}': '{key}[{index}]' must be a string, " + f"got {type(item).__name__}" + ) return list(value) def _read_str_mapping(name: str, raw: dict, key: str) -> Dict[str, str]: """Read an optional string-to-string mapping, defaulting to an empty dict. + Only the offending entry's *key* is ever reported, never the mapping and + never any of its values. ``env`` and ``headers`` routinely carry + credentials, and exception messages end up in the TUI error pane and in + logs, so a rejected entry must not echo the token sitting next to the typo. + Args: name: Name of the MCP server, used in the error message. raw: The raw server entry. @@ -316,14 +349,21 @@ def _read_str_mapping(name: str, raw: dict, key: str) -> Dict[str, str]: cannot reach the frozen configuration. Raises: - MitoCodeMcpConfigError: If the value is not a mapping of strings. + MitoCodeMcpConfigError: If the value is not a mapping of strings. The + message names the field and the bad entry's key, but no values. """ value = raw.get(key, {}) if not isinstance(value, dict): - _invalid(name, key, "an object mapping strings to strings", value) + raise MitoCodeMcpConfigError( + f"MCP server '{name}': '{key}' must be an object mapping strings to " + f"strings, got {type(value).__name__}" + ) for entry_key, entry_value in value.items(): if not isinstance(entry_key, str) or not isinstance(entry_value, str): - _invalid(name, key, "an object mapping strings to strings", value) + raise MitoCodeMcpConfigError( + f"MCP server '{name}': '{key}' entry {entry_key!r} must map a " + "string to a string" + ) return dict(value) diff --git a/test/test_mcp_config.py b/test/test_mcp_config.py index 0def97e..107be2c 100644 --- a/test/test_mcp_config.py +++ b/test/test_mcp_config.py @@ -117,3 +117,61 @@ def test_build_transport_returns_stdio_transport(): assert isinstance(transport, StdioTransport) assert transport.command == "npx" assert transport.args == ["-y", "f"] + + +def test_parse_error_message_never_echoes_header_values(): + secret = "Bearer sk-live-DO-NOT-LEAK" + with pytest.raises(MitoCodeMcpConfigError) as excinfo: + McpServerConfig.parse( + "weather", + { + "url": "https://weather.example.com/mcp", + "headers": {"Authorization": secret, "X-Retry": 5}, + }, + ) + message = str(excinfo.value) + assert "X-Retry" in message + assert secret not in message + + +def test_parse_error_message_never_echoes_env_values(): + secret = "env-token-DO-NOT-LEAK" + with pytest.raises(MitoCodeMcpConfigError) as excinfo: + McpServerConfig.parse( + "files", {"command": "npx", "env": {"MCP_TOKEN": secret, "PORT": 8080}} + ) + message = str(excinfo.value) + assert "PORT" in message + assert secret not in message + + +def test_parse_error_message_never_echoes_a_bare_header_string(): + secret = "sk-live-BARE-DO-NOT-LEAK" + with pytest.raises(MitoCodeMcpConfigError) as excinfo: + McpServerConfig.parse( + "weather", + {"url": "https://weather.example.com/mcp", "headers": secret}, + ) + message = str(excinfo.value) + assert "headers" in message + assert secret not in message + + +def test_parse_error_message_never_echoes_a_url_credential(): + secret = "url-password-DO-NOT-LEAK" + with pytest.raises(MitoCodeMcpConfigError) as excinfo: + McpServerConfig.parse("weather", {"url": f"ftp://user:{secret}@x/mcp"}) + message = str(excinfo.value) + assert "ftp" in message + assert secret not in message + + +def test_parse_error_message_never_echoes_argv_credentials(): + secret = "argv-token-DO-NOT-LEAK" + with pytest.raises(MitoCodeMcpConfigError) as excinfo: + McpServerConfig.parse( + "files", {"command": "npx", "args": ["-y", "--token", secret, 3]} + ) + message = str(excinfo.value) + assert "args[3]" in message + assert secret not in message From eb028723789968cc63d157fd7b85d5b297e673c5 Mon Sep 17 00:00:00 2001 From: Allen Wang Date: Fri, 16 Oct 2026 23:52:12 +0800 Subject: [PATCH 07/20] test(mcp): harden McpServerConfig suite to full branch coverage Enumerate every validation branch of mcp_config.py and pin each one: the non-dict/blank/whitespace entry shapes, non-bool enabled, NaN and bool timeouts, and the unreachable-from-parse unknown-kind guard. Adds sentinel credentials and asserts they never appear in exception messages, and pins that parse copies the caller's lists and mappings rather than aliasing them. 89% to 100% statement and branch coverage. A 37-mutant sweep against a scratch copy killed 36; the survivor is an equivalent reflow of one message literal, which is the expected result. Constraint: src/ is read-only to the test writer; no source file changed here Rejected: hypothesis property tests | no hypothesis in the dependency set, and the sweep already covers the two invariants Confidence: high Scope-risk: narrow --- test/test_mcp_config.py | 509 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 509 insertions(+) diff --git a/test/test_mcp_config.py b/test/test_mcp_config.py index 107be2c..92bf4ab 100644 --- a/test/test_mcp_config.py +++ b/test/test_mcp_config.py @@ -1,13 +1,44 @@ +"""Unit tests for :mod:`mitocode.backend.mcp.mcp_config`. + +The module is pure: it normalises a ``{"mcpServers": {...}}`` document into +frozen :class:`McpServerConfig` values and assembles a fastmcp transport +object. Nothing here touches the filesystem or a network, so the whole suite +runs offline and fast. + +Two properties are load-bearing beyond "it raises on bad input": + +* every optional field is *copied* out of the caller's raw mapping, so later + mutation of the document cannot reach a config that has already been + validated; and +* no exception message ever echoes a value from a secret-bearing field + (``url``, ``args``, ``env``, ``headers``), because those messages end up in + the TUI error pane and in logs. +""" + import os +import re import sys +from dataclasses import FrozenInstanceError +from unittest.mock import patch import pytest sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) from mitocode.backend.exception.mitocode_exception import MitoCodeMcpConfigError +from mitocode.backend.mcp import mcp_config from mitocode.backend.mcp.mcp_config import McpServerConfig, McpServerKind +# A valid https endpoint used wherever the transport selector is incidental. +HTTP_ENDPOINT = "https://weather.example.com/mcp" + +# Distinct sentinel credentials. None of these strings may ever appear inside +# an exception message raised while validating a secret-bearing field. +SECRET_HEADER = "Bearer sk-live-HEADER-DO-NOT-LEAK" +SECRET_ENV = "env-token-ENV-DO-NOT-LEAK" +SECRET_ARGV = "argv-token-ARGV-DO-NOT-LEAK" +SECRET_URL = "url-password-URL-DO-NOT-LEAK" + def test_mcp_config_error_is_a_mitocode_base_error(): from mitocode.backend.exception.mitocode_exception import MitoCodeMcpError @@ -175,3 +206,481 @@ def test_parse_error_message_never_echoes_argv_credentials(): message = str(excinfo.value) assert "args[3]" in message assert secret not in message + + +# -------------------------------------------------------------------------- +# Enum surface +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("kind", "expected_token"), + [(McpServerKind.STDIO, "stdio"), (McpServerKind.HTTP, "http")], + ids=["stdio", "http"], +) +def test_mcp_server_kind_is_a_string_enum_token_used_in_config_documents( + kind, expected_token +): + """A kind doubles as the lowercase token written into a config document.""" + assert ( + kind.value == expected_token + ), f"{kind!r} must serialise as {expected_token!r}" + assert isinstance( + kind, str + ), "McpServerKind must subclass str so config values can be compared to it" + + +# -------------------------------------------------------------------------- +# parse() — accepted shapes and defaulting +# -------------------------------------------------------------------------- + + +def test_parse_leaves_the_mutable_optional_fields_empty_when_they_are_absent(): + """Absent optional fields must default to empty containers, never None.""" + config = McpServerConfig.parse("weather", {"url": HTTP_ENDPOINT}) + + assert config.args == [], "args must default to an empty list" + assert config.env == {}, "env must default to an empty dict" + assert config.headers == {}, "headers must default to an empty dict" + assert config.cwd == "", "cwd must default to the empty string" + assert config.command == "", "an HTTP entry must not gain a command" + + +def test_parse_keeps_stray_fields_that_belong_to_the_other_transport(): + """Leniency contract: a stdio entry may carry headers/cwd and keeps them.""" + config = McpServerConfig.parse( + "files", + {"command": "npx", "cwd": "/tmp", "headers": {"X-Key": "value"}}, + ) + + assert config.kind is McpServerKind.STDIO, "the entry is still a stdio server" + assert config.headers == { + "X-Key": "value" + }, "a stray header mapping must be carried along, not dropped" + assert config.cwd == "/tmp", "a stray cwd must be carried along, not dropped" + + +def test_parse_honours_enabled_false(): + """enabled=False must survive parsing, not be coerced to the True default.""" + config = McpServerConfig.parse("weather", {"url": HTTP_ENDPOINT, "enabled": False}) + + assert ( + config.enabled is False + ), "an explicitly disabled server must stay disabled after parsing" + + +def test_parse_copies_the_callers_args_list(): + """Mutation of the caller's argv after parse must not reach the config.""" + args = ["-y", "files-mcp"] + config = McpServerConfig.parse("files", {"command": "npx", "args": args}) + + args.append("--token") + args[0] = "mutated" + + assert config.args == [ + "-y", + "files-mcp", + ], "the config holds its own copy of argv, so later mutation cannot reach it" + assert config.args is not args, "parse must not alias the caller's list" + + +def test_parse_copies_the_callers_env_and_headers_mappings(): + """Mutation of the caller's mappings after parse must not reach the config.""" + env = {"MCP_TOKEN": "seed"} + headers = {"Authorization": SECRET_HEADER} + config = McpServerConfig.parse( + "files", {"command": "npx", "env": env, "headers": headers} + ) + + env["INJECTED"] = "later" + headers["X-Added"] = "later" + + assert config.env == { + "MCP_TOKEN": "seed" + }, "the config holds its own copy of env, so later mutation cannot reach it" + assert config.headers == { + "Authorization": SECRET_HEADER + }, "the config holds its own copy of headers, so later mutation cannot reach it" + assert config.env is not env, "parse must not alias the caller's env mapping" + assert ( + config.headers is not headers + ), "parse must not alias the caller's headers mapping" + + +def test_mcp_server_config_is_frozen(): + """A validated config is shared, so its fields must not be assignable.""" + config = McpServerConfig.parse("weather", {"url": HTTP_ENDPOINT}) + + with pytest.raises(FrozenInstanceError, match="cannot assign to field 'timeout'"): + config.timeout = 1.0 + + +# -------------------------------------------------------------------------- +# parse() — the entry itself +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "bad_name", + [None, 7, "", " ", "\t\n"], + ids=["none", "int", "empty", "spaces", "whitespace"], +) +def test_parse_rejects_a_blank_or_non_string_name(bad_name): + """A blank name would namespace tools as '', so it must be refused.""" + with pytest.raises(MitoCodeMcpConfigError, match="name must be a non-empty string"): + McpServerConfig.parse(bad_name, {"url": HTTP_ENDPOINT}) + + +@pytest.mark.parametrize( + ("bad_raw", "expected_type"), + [ + (["https://x/mcp"], "list"), + ("https://x/mcp", "str"), + (None, "NoneType"), + (7, "int"), + ], + ids=["list", "str", "none", "int"], +) +def test_parse_rejects_an_entry_that_is_not_a_mapping(bad_raw, expected_type): + """A non-mapping entry must be refused, naming the type that arrived.""" + with pytest.raises( + MitoCodeMcpConfigError, match=f"entry must be an object, got {expected_type}$" + ): + McpServerConfig.parse("bad", bad_raw) + + +def test_parse_error_message_lists_entry_keys_without_echoing_values(): + """'got keys [...]' may name keys but must never name their values.""" + with pytest.raises( + MitoCodeMcpConfigError, match=r"got keys \['MCP_TOKEN'\]" + ) as excinfo: + McpServerConfig.parse("files", {"MCP_TOKEN": SECRET_ENV}) + + message = str(excinfo.value) + assert ( + SECRET_ENV not in message + ), f"the missing-selector message leaked a credential: {message!r}" + + +# -------------------------------------------------------------------------- +# parse() — the transport selector +# -------------------------------------------------------------------------- + + +def test_parse_rejects_a_url_without_a_scheme(): + """A bare host has no scheme, so the detail line says so without the value.""" + with pytest.raises( + MitoCodeMcpConfigError, match="got no http:// or https:// scheme" + ): + McpServerConfig.parse("bad", {"url": "weather.example.com/mcp"}) + + +@pytest.mark.parametrize( + ("bad_url", "expected_type"), + [ + (42, "int"), + (None, "NoneType"), + ([HTTP_ENDPOINT], "list"), + ({"url": HTTP_ENDPOINT}, "dict"), + ], + ids=["int", "none", "list", "dict"], +) +def test_parse_rejects_a_non_string_url(bad_url, expected_type): + """A non-string url must be refused, naming the type that arrived.""" + with pytest.raises( + MitoCodeMcpConfigError, + match=( + "'url' must be an http:// or https:// endpoint, " f"got {expected_type}$" + ), + ): + McpServerConfig.parse("bad", {"url": bad_url}) + + +@pytest.mark.parametrize( + "bad_command", + ["", " ", None, 7, ["npx"]], + ids=["empty", "blank", "none", "int", "list"], +) +def test_parse_rejects_a_blank_or_non_string_command(bad_command): + """An empty command would build a transport that can never spawn.""" + with pytest.raises( + MitoCodeMcpConfigError, match="'command' must be a non-empty string" + ): + McpServerConfig.parse("bad", {"command": bad_command}) + + +# -------------------------------------------------------------------------- +# parse() — the optional fields +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("bad_enabled", "expected_repr"), + [("yes", "'yes'"), (1, "1"), (0, "0"), (None, "None")], + ids=["str", "int-one", "int-zero", "none"], +) +def test_parse_rejects_a_non_boolean_enabled(bad_enabled, expected_repr): + """'enabled' is a flag; ints and strings are typos, not truthy values.""" + with pytest.raises( + MitoCodeMcpConfigError, + match=f"'enabled' must be a boolean, got {expected_repr}$", + ): + McpServerConfig.parse("weather", {"url": HTTP_ENDPOINT, "enabled": bad_enabled}) + + +@pytest.mark.parametrize( + ("bad_cwd", "expected_repr"), + [ + (5, "5"), + (None, "None"), + (["/tmp"], "['/tmp']"), + ({"dir": "/tmp"}, "{'dir': '/tmp'}"), + ], + ids=["int", "none", "list", "dict"], +) +def test_parse_rejects_a_non_string_cwd(bad_cwd, expected_repr): + """A non-string cwd must be refused, echoing the value that arrived. + + ``cwd`` is not secret-bearing, so unlike ``env``/``headers`` its value is + rendered into the message on purpose: a wrong path is what makes the typo + findable. + """ + with pytest.raises( + MitoCodeMcpConfigError, + match=f"'cwd' must be a string, got {re.escape(expected_repr)}$", + ): + McpServerConfig.parse("files", {"command": "npx", "cwd": bad_cwd}) + + +@pytest.mark.parametrize("field", ["timeout", "connect_timeout"]) +def test_parse_rejects_a_boolean_timeout(field): + """bool subclasses int, so 'timeout: true' must not become a 1s timeout.""" + with pytest.raises( + MitoCodeMcpConfigError, match=rf"'{field}' must be a positive number, got True" + ): + McpServerConfig.parse("files", {"command": "npx", field: True}) + + +@pytest.mark.parametrize("field", ["timeout", "connect_timeout"]) +def test_parse_rejects_a_nan_timeout(field): + """NaN is not greater than zero, so it must be refused rather than stored.""" + with pytest.raises( + MitoCodeMcpConfigError, match=rf"'{field}' must be a positive number, got nan" + ): + McpServerConfig.parse("files", {"command": "npx", field: float("nan")}) + + +def test_parse_rejects_a_non_numeric_timeout(): + """A quoted timeout is a typo, not a string that coerces later.""" + with pytest.raises( + MitoCodeMcpConfigError, match="'timeout' must be a positive number, got '30'" + ): + McpServerConfig.parse("files", {"command": "npx", "timeout": "30"}) + + +def test_parse_coerces_an_integer_timeout_to_a_float(): + """An int timeout is accepted, but the field is documented to hold a float.""" + config = McpServerConfig.parse("files", {"command": "npx", "timeout": 90}) + + assert config.timeout == 90.0, "the integer 90 must read back as 90.0" + assert isinstance( + config.timeout, float + ), "an int timeout must be normalised to float so callers need no int check" + + +# -------------------------------------------------------------------------- +# Secret-bearing fields — nothing may reach the exception message +# -------------------------------------------------------------------------- + + +def test_parse_error_message_never_echoes_argv_supplied_as_a_bare_string(): + """'args: "--token "' is refused by type alone, value untouched.""" + with pytest.raises( + MitoCodeMcpConfigError, match="'args' must be a list of strings, got str" + ) as excinfo: + McpServerConfig.parse( + "files", {"command": "npx", "args": f"--token {SECRET_ARGV}"} + ) + + message = str(excinfo.value) + assert ( + SECRET_ARGV not in message + ), f"the not-a-list argv message leaked a credential: {message!r}" + + +def test_parse_error_message_never_echoes_env_supplied_as_a_bare_string(): + """'env: "TOKEN="' is refused by type alone, value untouched.""" + with pytest.raises( + MitoCodeMcpConfigError, + match="'env' must be an object mapping strings to strings, got str", + ) as excinfo: + McpServerConfig.parse( + "files", {"command": "npx", "env": f"MCP_TOKEN={SECRET_ENV}"} + ) + + message = str(excinfo.value) + assert ( + SECRET_ENV not in message + ), f"the not-a-mapping env message leaked a credential: {message!r}" + + +@pytest.mark.parametrize("field", ["env", "headers"]) +def test_parse_error_message_names_a_non_string_entry_key_but_not_its_value(field): + """The offending key is named so the typo is fixable; the value is not.""" + with pytest.raises( + MitoCodeMcpConfigError, + match=rf"'{field}' entry 5 must map a string to a string", + ) as excinfo: + McpServerConfig.parse("files", {"command": "npx", field: {5: SECRET_ENV}}) + + message = str(excinfo.value) + assert ( + SECRET_ENV not in message + ), f"the bad-entry-key message for '{field}' leaked a credential: {message!r}" + + +def test_parse_error_message_never_echoes_a_scheme_less_url(): + """The no-scheme detail is a fixed sentence, so the url cannot leak.""" + with pytest.raises( + MitoCodeMcpConfigError, match="got no http:// or https:// scheme" + ) as excinfo: + McpServerConfig.parse("bad", {"url": f"//user:{SECRET_URL}@x/mcp"}) + + message = str(excinfo.value) + assert ( + SECRET_URL not in message + ), f"the no-scheme url message leaked a credential: {message!r}" + + +def test_parse_config_error_message_lists_top_level_keys_without_echoing_values(): + """The missing-key detail names keys only, never the document's values.""" + with pytest.raises( + MitoCodeMcpConfigError, match=r"got keys \['mcpToken'\]" + ) as excinfo: + McpServerConfig.parse_config({"mcpToken": SECRET_HEADER}) + + message = str(excinfo.value) + assert ( + SECRET_HEADER not in message + ), f"the missing-'mcpServers' message leaked a credential: {message!r}" + + +# -------------------------------------------------------------------------- +# parse_config() +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("bad_config", "expected_type"), + [([], "list"), ("mcpServers", "str"), (None, "NoneType"), (7, "int")], + ids=["list", "str", "none", "int"], +) +def test_parse_config_rejects_a_document_that_is_not_a_mapping( + bad_config, expected_type +): + """The whole document must be a mapping before its keys mean anything.""" + with pytest.raises( + MitoCodeMcpConfigError, + match=( + "must be an object holding an 'mcpServers' mapping, " + f"got {expected_type}$" + ), + ): + McpServerConfig.parse_config(bad_config) + + +@pytest.mark.parametrize( + ("bad_servers", "expected_type"), + [([], "list"), ("weather", "str"), (None, "NoneType")], + ids=["list", "str", "none"], +) +def test_parse_config_rejects_mcp_servers_that_is_not_a_mapping( + bad_servers, expected_type +): + """'mcpServers' must itself be a mapping of entries.""" + with pytest.raises( + MitoCodeMcpConfigError, + match=f"'mcpServers' must be an object, got {expected_type}$", + ): + McpServerConfig.parse_config({"mcpServers": bad_servers}) + + +def test_parse_config_returns_an_empty_list_for_an_empty_servers_mapping(): + """A document with no servers is valid, not an error.""" + assert McpServerConfig.parse_config({"mcpServers": {}}) == [] + + +def test_parse_config_surfaces_the_failing_entry_name(): + """An entry that parse() rejects must fail the whole document, by name.""" + with pytest.raises( + MitoCodeMcpConfigError, + match="MCP server 'broken': 'url' and 'command' are mutually exclusive", + ): + McpServerConfig.parse_config( + { + "mcpServers": { + "weather": {"url": HTTP_ENDPOINT}, + "broken": {"url": HTTP_ENDPOINT, "command": "npx"}, + } + } + ) + + +# -------------------------------------------------------------------------- +# build_transport() +# -------------------------------------------------------------------------- + + +def test_build_transport_rejects_an_unknown_transport_kind(): + """Only direct construction can reach the guard, since parse sets the kind.""" + config = McpServerConfig(name="legacy", kind="sse", url=HTTP_ENDPOINT) + + with pytest.raises( + MitoCodeMcpConfigError, + match="MCP server 'legacy': unknown transport kind 'sse'", + ): + config.build_transport() + + +def test_build_transport_passes_none_for_an_unset_env_and_cwd(): + """Unset stdio options must reach fastmcp as None, not as {} or ''.""" + config = McpServerConfig.parse("files", {"command": "npx", "args": []}) + + transport = config.build_transport() + + assert ( + transport.env is None + ), "an empty env must be handed over as 'unset' (None), not as {}" + assert ( + transport.cwd is None + ), "an empty cwd must be handed over as 'unset' (None), not as ''" + + +def test_build_transport_carries_populated_env_and_cwd_into_the_stdio_transport(): + """Configured stdio options must survive the hand-off to fastmcp.""" + config = McpServerConfig.parse( + "files", + { + "command": "npx", + "args": ["-y", "f"], + "env": {"MCP_TOKEN": "seed"}, + "cwd": "/tmp", + }, + ) + + transport = config.build_transport() + + assert transport.command == "npx", "the command must reach the transport" + assert transport.args == ["-y", "f"], "the argv must reach the transport" + assert transport.env == {"MCP_TOKEN": "seed"}, "env must reach the transport" + assert transport.cwd == "/tmp", "cwd must reach the transport" + + +def test_build_transport_passes_none_for_unset_headers(): + """fastmcp stores headers as {}, so only the call site shows 'unset'.""" + config = McpServerConfig.parse("weather", {"url": HTTP_ENDPOINT}) + + with patch.object(mcp_config, "StreamableHttpTransport") as transport_cls: + config.build_transport() + + transport_cls.assert_called_once_with(url=HTTP_ENDPOINT, headers=None) From 9a0322222895b5f433cb07d2e815aadb463d5032 Mon Sep 17 00:00:00 2001 From: Allen Wang Date: Fri, 16 Oct 2026 23:54:18 +0800 Subject: [PATCH 08/20] test(mcp): pin zero-timeout boundary and populated HTTP headers Two gaps found in review of the McpServerConfig suite: - The strict "greater than zero" bound was only pinned incidentally, by timeout: 0 in the pre-existing rejection test and connect_timeout: -1. Add an explicit boundary test over int 0, float 0.0 and -1 across both timeout fields so the strictness is pinned directly rather than inferred from a neighbouring case. - build_transport passed headers=self.headers or None, but only the empty case was observed. Replacing that expression with a bare headers=None survived the whole suite, so a server silently losing its Authorization header had no signal. Pin the populated case at the call site, where fastmcp's own headers or {} normalisation hides the distinction. Also record the deliberately unpinned limitations of the module in the test module docstring: shallow immutability of the frozen dataclass, the auto-generated repr printing url verbatim, float("inf") passing the timeout check, "http://" with no host being accepted, and the credential-in-key position that is echoed by design. Co-Authored-By: Claude Opus 4.8 (1M context) --- test/test_mcp_config.py | 59 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/test/test_mcp_config.py b/test/test_mcp_config.py index 92bf4ab..251a3d4 100644 --- a/test/test_mcp_config.py +++ b/test/test_mcp_config.py @@ -13,6 +13,33 @@ * no exception message ever echoes a value from a secret-bearing field (``url``, ``args``, ``env``, ``headers``), because those messages end up in the TUI error pane and in logs. + +Known limitations, deliberately unpinned +---------------------------------------- + +The credential guarantee above is scoped to *exception messages raised while +validating*. These neighbouring behaviours are accepted by the module and are +deliberately not pinned here, because pinning them would lock in a decision +that belongs to the module's owner, not to this suite: + +* **Shallow immutability.** ``frozen=True`` blocks field *assignment* only. + ``config.args``, ``config.env`` and ``config.headers`` remain mutable in + place, and ``build_transport`` hands those live containers to fastmcp by + reference — mutating them afterwards changes what a later connection uses. +* **The auto-generated ``repr`` is not covered by the guarantee.** The + dataclass ``repr`` prints every field, ``url`` included, so ``log(config)`` + or a traceback that formats the config can still echo an embedded URL + password. +* **``float("inf")`` is accepted as a timeout.** ``inf > 0`` holds, so an + infinite timeout parses and then hangs the caller forever. +* **``"http://"`` with no host is accepted.** The scheme check is a + ``startswith`` prefix test, not a URL parse, so a hostless URL reaches + fastmcp and fails later at connect time rather than at config time. +* **A credential placed in an entry *key* is echoed.** + ``_read_str_mapping`` names the offending key on purpose, and + ``parse``/``parse_config`` list entry keys via ``sorted(map(str, ...))``, so + ``env: {"sk-live-XXXX": 8080}`` puts the token in the message. Only the + *values* are guaranteed silent; that is the documented contract. """ import os @@ -470,6 +497,23 @@ def test_parse_rejects_a_nan_timeout(field): McpServerConfig.parse("files", {"command": "npx", field: float("nan")}) +@pytest.mark.parametrize("field", ["timeout", "connect_timeout"]) +@pytest.mark.parametrize( + ("bad_timeout", "expected_repr"), + [(0, "0"), (0.0, "0.0"), (-1, "-1")], + ids=["int-zero", "float-zero", "negative"], +) +def test_parse_rejects_a_timeout_at_the_zero_boundary( + field, bad_timeout, expected_repr +): + """The bound is strict: zero is refused, not accepted as a 0-second timeout.""" + with pytest.raises( + MitoCodeMcpConfigError, + match=f"'{field}' must be a positive number, got {re.escape(expected_repr)}$", + ): + McpServerConfig.parse("files", {"command": "npx", field: bad_timeout}) + + def test_parse_rejects_a_non_numeric_timeout(): """A quoted timeout is a typo, not a string that coerces later.""" with pytest.raises( @@ -684,3 +728,18 @@ def test_build_transport_passes_none_for_unset_headers(): config.build_transport() transport_cls.assert_called_once_with(url=HTTP_ENDPOINT, headers=None) + + +def test_build_transport_carries_populated_headers_into_the_http_transport(): + """Configured auth headers must survive the hand-off to fastmcp.""" + config = McpServerConfig.parse( + "weather", + {"url": HTTP_ENDPOINT, "headers": {"Authorization": SECRET_HEADER}}, + ) + + with patch.object(mcp_config, "StreamableHttpTransport") as transport_cls: + config.build_transport() + + transport_cls.assert_called_once_with( + url=HTTP_ENDPOINT, headers={"Authorization": SECRET_HEADER} + ) From 1fde6ccd7e35eab017b9ae0711a23edc65d140d8 Mon Sep 17 00:00:00 2001 From: Allen Wang Date: Sat, 17 Oct 2026 12:27:54 +0800 Subject: [PATCH 09/20] test(mcp): cover McpClient teardown error paths Force close records rather than raises a failed client.close() or stack.aclose(), so those branches needed direct coverage. Adds fakes for both handles, an idempotent second disconnect after a wedged teardown, a no-op disconnect on a never-connected client, and a connect that fails before any handle exists. --- test/test_mcp_client.py | 536 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 536 insertions(+) create mode 100644 test/test_mcp_client.py diff --git a/test/test_mcp_client.py b/test/test_mcp_client.py new file mode 100644 index 0000000..6e38685 --- /dev/null +++ b/test/test_mcp_client.py @@ -0,0 +1,536 @@ +import asyncio +import os +import sys +from typing import Optional + +import pytest +from fastmcp import Client, FastMCP + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) + +from mitocode.backend.exception.mitocode_exception import MitoCodeMcpToolCallError +from mitocode.backend.mcp.mcp_client import McpClient, McpConnectionState +from mitocode.backend.mcp.mcp_config import McpServerConfig + + +def _server() -> FastMCP: + """Build an in-memory FastMCP server exposing three tools.""" + server = FastMCP("probe") + + @server.tool + def echo(text: str) -> str: + """Echo the given text back.""" + return f"echo:{text}" + + @server.tool + def add(a: int, b: int) -> dict: + """Add two integers.""" + return {"sum": a + b} + + @server.tool + def boom() -> str: + """Always fail.""" + raise ValueError("kaboom") + + return server + + +def _in_memory_config() -> McpServerConfig: + """Config placeholder reused by connect() monkeypatching.""" + return McpServerConfig.parse("probe", {"url": "https://probe.invalid/mcp"}) + + +def _slow_config(connect_timeout: float) -> McpServerConfig: + """Build a config pointing at an address that never answers.""" + return McpServerConfig.parse( + "slow", + {"url": "https://10.255.255.1/mcp", "connect_timeout": connect_timeout}, + ) + + +class _CloseRecorder: + """Stand-in for a fastmcp ``Client`` that records whether close ran.""" + + def __init__(self, error: Optional[Exception] = None): + """Arm the recorder, optionally to raise *error* instead of returning.""" + self.error = error + self.closed = False + + async def close(self) -> None: + """Mark the client closed, then fail if this recorder was armed.""" + self.closed = True + if self.error is not None: + raise self.error + + +class _StackRecorder: + """Stand-in for an ``AsyncExitStack`` that records whether aclose ran.""" + + def __init__(self, error: Optional[Exception] = None): + """Arm the recorder, optionally to raise *error* instead of returning.""" + self.error = error + self.closed = False + + async def aclose(self) -> None: + """Mark the stack closed, then fail if this recorder was armed.""" + self.closed = True + if self.error is not None: + raise self.error + + +@pytest.mark.asyncio +async def test_connect_discovers_tools_from_server(monkeypatch): + config = _in_memory_config() + client = McpClient(config) + monkeypatch.setattr(client, "_build_fastmcp_client", lambda: Client(_server())) + + assert await client.connect() is True + assert client.state is McpConnectionState.CONNECTED + assert client.last_error == "" + assert sorted(t.name for t in client.tools) == ["add", "boom", "echo"] + + +@pytest.mark.asyncio +async def test_call_tool_returns_text_result(monkeypatch): + config = _in_memory_config() + client = McpClient(config) + monkeypatch.setattr(client, "_build_fastmcp_client", lambda: Client(_server())) + await client.connect() + + assert await client.call_tool("echo", {"text": "hi"}) == "echo:hi" + + +@pytest.mark.asyncio +async def test_call_tool_serializes_structured_result_as_json(monkeypatch): + config = _in_memory_config() + client = McpClient(config) + monkeypatch.setattr(client, "_build_fastmcp_client", lambda: Client(_server())) + await client.connect() + + result = await client.call_tool("add", {"a": 2, "b": 3}) + assert result == '{"sum": 5}' + + +@pytest.mark.asyncio +async def test_call_tool_raises_on_tool_error(monkeypatch): + config = _in_memory_config() + client = McpClient(config) + monkeypatch.setattr(client, "_build_fastmcp_client", lambda: Client(_server())) + await client.connect() + + with pytest.raises(MitoCodeMcpToolCallError) as excinfo: + await client.call_tool("boom", {}) + assert "kaboom" in str(excinfo.value) + + +@pytest.mark.asyncio +async def test_connect_failure_is_recorded_not_raised(): + config = McpServerConfig.parse( + "broken", + {"url": "https://127.0.0.1:1/mcp", "connect_timeout": 3.0}, + ) + client = McpClient(config) + + assert await client.connect() is False + assert client.state is McpConnectionState.FAILED + assert client.last_error != "" + assert client.tools == [] + + +@pytest.mark.asyncio +async def test_connect_timeout_is_honoured(): + client = McpClient(_slow_config(1.0)) + + loop = asyncio.get_event_loop() + started = loop.time() + assert await client.connect() is False + assert loop.time() - started < 5.0 + + +@pytest.mark.asyncio +async def test_connect_timeout_force_closes_the_client(): + """A timed-out connect must not leave a subprocess running. + + Review Focus #6: FastMCP's StdioTransport spawns the process inside a + detached task, so cancelling our await alone orphans it. Only + Client.close() tears the process group down. + """ + client = McpClient(_slow_config(1.0)) + + class _HangingClient: + """Async context manager whose __aenter__ never completes.""" + + closed = False + + async def __aenter__(self): + await asyncio.Event().wait() + + async def __aexit__(self, *exc_info): + return False + + async def close(self): + type(self).closed = True + + built = _HangingClient() + client._build_fastmcp_client = lambda: built + + assert await client.connect() is False + assert built.closed is True + assert client._client is None + + +@pytest.mark.asyncio +async def test_connect_timeout_during_tool_discovery_is_honoured(): + """Discovery is a second round trip and needs the connect deadline too. + + A server that completes the handshake but then never answers list_tools + would hang forever if only __aenter__ were wrapped in wait_for. + """ + client = McpClient(_slow_config(1.0)) + + class _StallingDiscoveryClient: + """Client that connects instantly but never finishes listing tools.""" + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc_info): + return False + + async def list_tools(self): + await asyncio.Event().wait() + + async def close(self): + return None + + client._build_fastmcp_client = lambda: _StallingDiscoveryClient() + + loop = asyncio.get_event_loop() + started = loop.time() + # The outer deadline is a backstop: without wait_for around list_tools the + # connect hangs forever, and this test should fail rather than hang CI. + assert await asyncio.wait_for(client.connect(), timeout=10.0) is False + assert loop.time() - started < 5.0 + assert client.state is McpConnectionState.FAILED + assert client.tools == [] + + +@pytest.mark.asyncio +async def test_call_tool_raises_when_not_connected(): + config = _in_memory_config() + client = McpClient(config) + + with pytest.raises(MitoCodeMcpToolCallError) as excinfo: + await client.call_tool("echo", {"text": "hi"}) + # Not merely "some error": the message must say the server is not connected, + # otherwise an unreachable server is indistinguishable from a bad tool call. + assert "not connected" in str(excinfo.value) + + +@pytest.mark.asyncio +async def test_call_tool_notes_non_text_content(): + """Review Focus #4: a non-text-only result must not become an empty string.""" + client = McpClient(_in_memory_config()) + + class _Block: + type = "image" + + class _Result: + data = None + content = [_Block()] + is_error = False + + async def _fake_call_tool(name, arguments, **kwargs): + return _Result() + + fake_client = type("C", (), {})() + fake_client.call_tool = _fake_call_tool + client._client = fake_client + + result = await client.call_tool("anything", {}) + assert "non-text" in result + + +@pytest.mark.asyncio +async def test_call_tool_renders_empty_result_as_empty_string(monkeypatch): + client = McpClient(_in_memory_config()) + monkeypatch.setattr(client, "_build_fastmcp_client", lambda: Client(_server())) + await client.connect() + + class _EmptyResult: + data = None + content = [] + + class _NullClient: + async def call_tool(self, name, arguments, **kwargs): + return _EmptyResult() + + client._client = _NullClient() + + assert await client.call_tool("echo", {}) == "" + + +@pytest.mark.asyncio +async def test_disconnect_clears_state_and_is_idempotent(monkeypatch): + client = McpClient(_in_memory_config()) + monkeypatch.setattr(client, "_build_fastmcp_client", lambda: Client(_server())) + await client.connect() + assert client.tools + + await client.disconnect() + + assert client.state is McpConnectionState.DISCONNECTED + assert client.tools == [] + assert client._client is None + await client.disconnect() + assert client.state is McpConnectionState.DISCONNECTED + + +@pytest.mark.asyncio +async def test_disconnect_records_cleanup_failure_without_raising(monkeypatch): + """A teardown error must not stop the remaining servers from closing.""" + client = McpClient(_in_memory_config()) + monkeypatch.setattr(client, "_build_fastmcp_client", lambda: Client(_server())) + await client.connect() + assert client.last_error == "" + + class _BadClose: + async def close(self): + raise RuntimeError("transport wedged") + + client._client = _BadClose() + + await client.disconnect() + + assert client.state is McpConnectionState.DISCONNECTED + assert "transport wedged" in client.last_error + assert client.tools == [] + + +@pytest.mark.asyncio +async def test_connect_returns_early_when_already_connected(monkeypatch): + client = McpClient(_in_memory_config()) + builds = [] + monkeypatch.setattr( + client, "_build_fastmcp_client", lambda: builds.append(1) or Client(_server()) + ) + await client.connect() + + assert await client.connect() is True + + assert len(builds) == 1 + + +def test_fresh_client_reports_its_server_name_and_disconnected_state(): + client = McpClient(_in_memory_config()) + + assert client.server_name == "probe" + assert client.state is McpConnectionState.DISCONNECTED + assert client.tools == [] + assert client.last_error == "" + + +@pytest.mark.asyncio +async def test_call_tool_times_out_and_reports_the_deadline(): + client = McpClient( + McpServerConfig.parse( + "slowtool", {"url": "https://x.invalid/mcp", "timeout": 0.05} + ) + ) + + class _HangingCall: + async def call_tool(self, name, arguments, **kwargs): + await asyncio.Event().wait() + + client._client = _HangingCall() + + with pytest.raises(MitoCodeMcpToolCallError) as excinfo: + await client.call_tool("echo", {"text": "hi"}) + assert "timed out" in str(excinfo.value) + assert "0.05" in str(excinfo.value) + + +@pytest.mark.asyncio +async def test_call_tool_wraps_unexpected_error_without_leaking_arguments(): + """Arguments routinely carry file contents and credentials, so they must + never reach an exception message that ends up in the TUI and in logs.""" + client = McpClient(_in_memory_config()) + + class _ExplodingCall: + async def call_tool(self, name, arguments, **kwargs): + raise RuntimeError("transport gave up") + + client._client = _ExplodingCall() + secret = {"token": "sk-do-not-log-me"} + + with pytest.raises(MitoCodeMcpToolCallError) as excinfo: + await client.call_tool("echo", secret) + + message = str(excinfo.value) + assert "transport gave up" in message + assert "RuntimeError" in message + assert "sk-do-not-log-me" not in message + + +@pytest.mark.asyncio +async def test_call_tool_does_not_re_wrap_its_own_error(): + client = McpClient(_in_memory_config()) + original = MitoCodeMcpToolCallError("precise original message") + + class _RaisingCall: + async def call_tool(self, name, arguments, **kwargs): + raise original + + client._client = _RaisingCall() + + with pytest.raises(MitoCodeMcpToolCallError) as excinfo: + await client.call_tool("echo", {}) + + assert excinfo.value is original + + +@pytest.mark.asyncio +async def test_call_tool_joins_text_content_blocks(): + client = McpClient(_in_memory_config()) + + class _TextBlock: + def __init__(self, text): + self.text = text + + class _Result: + data = None + content = [_TextBlock("first "), _TextBlock("second")] + + class _Call: + async def call_tool(self, name, arguments, **kwargs): + return _Result() + + client._client = _Call() + + assert await client.call_tool("echo", {}) == "first second" + + +@pytest.mark.asyncio +async def test_force_close_records_the_type_of_a_failing_client_close(): + """A wedged transport must become a recorded error, never a raised one.""" + client = McpClient(_in_memory_config()) + recorder = _CloseRecorder(RuntimeError("transport wedged")) + client._client = recorder + client._stack = _StackRecorder() + + await client.disconnect() + + assert client.last_error == "close failed: RuntimeError: transport wedged" + assert recorder.closed is True + assert client._client is None + assert client._stack is None + assert client.tools == [] + + +@pytest.mark.asyncio +async def test_force_close_still_closes_the_stack_when_client_close_raises(): + """A client that cannot be closed must not strand the stack holding it.""" + client = McpClient(_in_memory_config()) + client._client = _CloseRecorder(RuntimeError("client wedged")) + stack = _StackRecorder() + client._stack = stack + + await client.disconnect() + + assert stack.closed is True + assert client.last_error == "close failed: RuntimeError: client wedged" + + +@pytest.mark.asyncio +async def test_force_close_records_a_failing_stack_close(): + """A stack that cannot unwind is recorded instead of aborting the teardown.""" + client = McpClient(_in_memory_config()) + recorder = _CloseRecorder() + client._client = recorder + stack = _StackRecorder(OSError("process group already gone")) + client._stack = stack + + await client.disconnect() + + assert ( + client.last_error == "stack close failed: OSError: process group already gone" + ) + assert recorder.closed is True + assert stack.closed is True + + +@pytest.mark.asyncio +async def test_force_close_keeps_both_errors_when_both_handles_fail(): + """Neither teardown error may be lost behind the other.""" + client = McpClient(_in_memory_config()) + client._client = _CloseRecorder(RuntimeError("client wedged")) + client._stack = _StackRecorder(OSError("stack wedged")) + + await client.disconnect() + + assert client.last_error == ( + "close failed: RuntimeError: client wedged; " + "stack close failed: OSError: stack wedged" + ) + + +@pytest.mark.asyncio +async def test_disconnect_stays_idempotent_after_a_failing_close(monkeypatch): + """A failing teardown still leaves the object clean enough to be retried.""" + client = McpClient(_in_memory_config()) + monkeypatch.setattr(client, "_build_fastmcp_client", lambda: Client(_server())) + assert await client.connect() is True + assert client.tools + + client._client = _CloseRecorder(RuntimeError("transport wedged")) + + await client.disconnect() + + assert client.state is McpConnectionState.DISCONNECTED + assert client._client is None + assert client._stack is None + assert client.tools == [] + assert "transport wedged" in client.last_error + + # The second teardown is a no-op, not a retry: the error stays recorded. + await client.disconnect() + + assert client._client is None + assert client._stack is None + assert client.state is McpConnectionState.DISCONNECTED + assert client.last_error == "close failed: RuntimeError: transport wedged" + + +@pytest.mark.asyncio +async def test_disconnect_on_a_never_connected_client_changes_nothing(): + """Tearing down a client that never opened must not invent an error.""" + client = McpClient(_in_memory_config()) + client._state = McpConnectionState.FAILED + client._last_error = "connect never got that far" + + await client.disconnect() + + assert client.state is McpConnectionState.FAILED + assert client.last_error == "connect never got that far" + assert client._client is None + assert client._stack is None + assert client.tools == [] + + +@pytest.mark.asyncio +async def test_connect_recovers_when_no_handle_was_ever_created(monkeypatch): + """A transport that cannot even be built leaves nothing to close, cleanly.""" + + def _explode(): + raise RuntimeError("bad transport arguments") + + client = McpClient(_in_memory_config()) + monkeypatch.setattr(client, "_build_fastmcp_client", _explode) + + assert await client.connect() is False + + assert client.state is McpConnectionState.FAILED + assert client.last_error == "RuntimeError: bad transport arguments" + assert client._client is None + assert client._stack is None + assert client.tools == [] From 588c577205f43a2a91f4037a3dccfdabb423acff Mon Sep 17 00:00:00 2001 From: Allen Wang Date: Sat, 17 Oct 2026 12:28:31 +0800 Subject: [PATCH 10/20] feat(mcp): add McpClient with cached tool discovery and safe failure One MCP server connection per client. connect() opens the transport and caches the server's tool catalogue in one step, and records failures into state and last_error instead of raising, so one unreachable server cannot abort a batch connect. call_tool() renders a result to text because AgentRuntime calls str() on whatever a tool returns. The timeout path calls Client.close() rather than merely abandoning the await. StdioTransport spawns its subprocess in a detached task and the caller only waits on that task's ready event, so asyncio.wait_for's cancellation would leave the connect task running and the child process alive as an orphan. Handles are assigned before the context is entered, and cleanup runs outside the timed-out await so the teardown itself is never cancelled. Constraint: Result must render to a string, because AgentRuntime calls str() on tool return values Rejected: ClientGroup.from_config | its __aenter__ connects every client, so one bad server kills all tools Confidence: high Scope-risk: narrow --- src/mitocode/backend/mcp/mcp_client.py | 259 +++++++++++++++++++++++++ 1 file changed, 259 insertions(+) create mode 100644 src/mitocode/backend/mcp/mcp_client.py diff --git a/src/mitocode/backend/mcp/mcp_client.py b/src/mitocode/backend/mcp/mcp_client.py new file mode 100644 index 0000000..7c80e70 --- /dev/null +++ b/src/mitocode/backend/mcp/mcp_client.py @@ -0,0 +1,259 @@ +"""Connection lifecycle for a single MCP server. + +:class:`McpClient` owns one server connection: it opens the transport, caches +the tool catalogue the server advertises, and invokes tools by name. It +deliberately owns *only* the connection — assembling a set of servers and +deciding which ones to start is a later layer's job. + +The asymmetry that shapes this module is between :meth:`McpClient.connect` and +:meth:`McpClient.call_tool`. Connecting is a *best-effort* operation: a server +that is down, misconfigured or simply slow must not abort the startup of the +servers that do work, so :meth:`connect` records the failure into +:attr:`McpClient.state` and :attr:`McpClient.last_error` and returns ``False`` +instead of raising. Invoking a tool, by contrast, is a direct answer to a +direct request from the model, so :meth:`call_tool` raises +:class:`MitoCodeMcpToolCallError` and lets the caller decide what to do. +""" + +import asyncio +import json +from contextlib import AsyncExitStack +from enum import Enum +from typing import Any, Dict, List, Optional + +from fastmcp import Client +from fastmcp.exceptions import ToolError + +from mitocode.backend.exception.mitocode_exception import MitoCodeMcpToolCallError +from mitocode.backend.mcp.mcp_config import McpServerConfig + + +class McpConnectionState(str, Enum): + """Connection state of a single :class:`McpClient`.""" + + DISCONNECTED = "disconnected" + CONNECTED = "connected" + FAILED = "failed" + + +class McpClient: + """One MCP server connection plus its cached tool catalogue. + + Args: + config: The validated server definition to connect to. + """ + + def __init__(self, config: McpServerConfig): + """Store the server definition; the connection opens on :meth:`connect`. + + Args: + config: The validated server definition to connect to. + """ + self._config = config + self._client: Optional[Client] = None + self._stack: Optional[AsyncExitStack] = None + self._tools: List[Any] = [] + self._state = McpConnectionState.DISCONNECTED + self._last_error: str = "" + + def _build_fastmcp_client(self) -> Client: + """Construct the FastMCP client for this server. + + Exists as a seam so tests can substitute an in-memory client. + """ + return Client( + self._config.build_transport(), + name=self._config.name, + timeout=self._config.timeout, + ) + + @property + def server_name(self) -> str: + """Name this server's tools are namespaced under.""" + return self._config.name + + @property + def state(self) -> McpConnectionState: + """Current connection state, as of the last :meth:`connect` or :meth:`disconnect`.""" + return self._state + + @property + def last_error(self) -> str: + """Description of the most recent failure, or ``""`` when there is none.""" + return self._last_error + + @property + def tools(self) -> List[Any]: + """Tool objects cached by :meth:`connect`; empty while not connected.""" + return self._tools + + async def connect(self) -> bool: + """Open the connection and cache the server's tool catalogue. + + Failures are recorded rather than raised: an unreachable server must not + take down the servers that are reachable. Callers read :attr:`state` or + the return value to find out what happened. + + The timeout path deliberately closes the client rather than merely + abandoning the await. FastMCP's ``StdioTransport`` spawns its subprocess + inside a *detached* task and the calling coroutine only waits on that + task's ready event, so cancelling our await — which is all + ``asyncio.wait_for`` does — leaves the transport's connect task running + and the child process alive as an orphan. Only ``Client.close()`` sets + the transport's stop event and tears down the process group. + + Returns: + ``True`` when the client is connected and its tools are cached, + ``False`` otherwise. + """ + if self._state is McpConnectionState.CONNECTED: + return True + + timeout = self._config.connect_timeout + try: + # Assigned before entering the context so that a handshake which + # never completes still leaves close() reachable from _force_close. + client = self._build_fastmcp_client() + self._client = client + stack = AsyncExitStack() + self._stack = stack + await asyncio.wait_for(stack.enter_async_context(client), timeout=timeout) + # Discovery is a second round trip and needs the same deadline; + # without it a server that accepts connections but never answers + # would hang forever. + self._tools = list( + await asyncio.wait_for(client.list_tools(), timeout=timeout) + ) + except Exception as exc: # pylint: disable=broad-except + # Cleanup runs outside the timed-out await, or the teardown itself + # would be cancelled and the subprocess orphaned after all. + await self._force_close() + self._state = McpConnectionState.FAILED + self._last_error = f"{type(exc).__name__}: {exc}" + return False + + self._state = McpConnectionState.CONNECTED + self._last_error = "" + return True + + async def _force_close(self) -> None: + """Tear the connection down, swallowing and recording cleanup errors. + + Shared by the failure path of :meth:`connect` and by + :meth:`disconnect` so there is exactly one way to release a connection. + The handles are dropped before the awaits so that a transport which + raises while closing still leaves the object in a clean state. + """ + client = self._client + stack = self._stack + self._client = None + self._stack = None + self._tools = [] + + errors = [] + if client is not None: + try: + await client.close() + except Exception as exc: # pylint: disable=broad-except + errors.append(f"close failed: {type(exc).__name__}: {exc}") + if stack is not None: + try: + await stack.aclose() + except Exception as exc: # pylint: disable=broad-except + errors.append(f"stack close failed: {type(exc).__name__}: {exc}") + if errors: + self._last_error = "; ".join(errors) + + async def disconnect(self) -> None: + """Close the connection if one is open. Idempotent. + + Errors raised while closing are swallowed and recorded in + :attr:`last_error` rather than propagated, because a failing teardown + must not prevent the remaining servers from being shut down. + """ + if self._stack is None and self._client is None: + return + await self._force_close() + self._state = McpConnectionState.DISCONNECTED + + async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> str: + """Invoke *tool_name* on the connected server and render the result as text. + + Args: + tool_name: The tool's name **as advertised by the server**, without + the ``mcp____`` namespace prefix. + arguments: Keyword arguments to forward to the tool. + + Returns: + The tool result as text. + + Raises: + MitoCodeMcpToolCallError: If the client is not connected, the call + times out, or the server reports a tool error. + """ + server = self._config.name + client = self._client + if client is None: + raise MitoCodeMcpToolCallError( + f"MCP server '{server}' is not connected, " + f"cannot call tool '{tool_name}'" + ) + + timeout = self._config.timeout + try: + result = await asyncio.wait_for( + client.call_tool(tool_name, arguments, timeout=timeout), + timeout=timeout, + ) + except ToolError as exc: + raise MitoCodeMcpToolCallError( + f"MCP tool '{tool_name}' on server '{server}' failed: {exc}" + ) from exc + except asyncio.TimeoutError as exc: + raise MitoCodeMcpToolCallError( + f"MCP tool '{tool_name}' on server '{server}' timed out " + f"after {timeout} seconds" + ) from exc + except MitoCodeMcpToolCallError: + raise + except Exception as exc: # pylint: disable=broad-except + # Arguments are never echoed: they routinely carry file contents and + # credentials, and this text reaches the TUI error pane and logs. + raise MitoCodeMcpToolCallError( + f"MCP tool '{tool_name}' on server '{server}' raised " + f"{type(exc).__name__}: {exc}" + ) from exc + return self._render_result(result) + + @staticmethod + def _render_result(result: Any) -> str: + """Render a fastmcp ``CallToolResult`` as plain text. + + Structured results are JSON-encoded so the model sees data rather than a + Python repr. A result carrying only non-text blocks — images, audio, + embedded resources — is reported as such instead of collapsing to an + empty string, which would read to the model as a successful empty answer. + + Args: + result: The ``CallToolResult`` returned by fastmcp. + + Returns: + The rendered text, possibly the empty string. + """ + data = getattr(result, "data", None) + if data is not None: + if isinstance(data, str): + return data + return json.dumps(data, ensure_ascii=False, default=str) + + content = getattr(result, "content", None) or [] + texts = [block.text for block in content if hasattr(block, "text")] + joined = "".join(texts) + if joined: + return joined + if content: + return ( + f"[MCP server returned {len(content)} non-text content block(s), " + "which are not supported here]" + ) + return "" From c18605fce7e606835da170421728c4205227e845 Mon Sep 17 00:00:00 2001 From: Allen Wang Date: Sat, 17 Oct 2026 12:35:28 +0800 Subject: [PATCH 11/20] fix(mcp): keep arguments out of tool errors and orphaned subprocesses Two review findings on McpClient, both about the same invariant: never orphan a subprocess, never leak a credential. The ToolError path interpolated fastmcp's text verbatim, and for an argument-validation failure that text is pydantic's, whose input_value= segment is the argument itself. A rejected token reached the exception message and from there the TUI error pane and the logs, contradicting the no-leak constraint the module's own comment asserts. Payloads are now cut out by _redact_arguments; tool, server, field name, expected type and input_type all survive, so the message still says what was wrong, just not with what. Cancellation derives from BaseException, so it bypassed the except Exception handler entirely and left _client/_stack populated with stdio's detached connect task still holding the child process -- exactly the orphan the rest of connect() is built to prevent. connect() now catches BaseException, cleans up under asyncio.shield so a repeat cancellation cannot interrupt the teardown, records the failure, and re-raises for anything that is not an ordinary Exception. Constraint: connect() keeps returning False for ordinary failures so one unreachable server cannot abort a batch connect; only cancellation re-raises. Constraint: no credential value may reach an exception message. Rejected: never raising on cancellation | swallows Ctrl-C and TUI shutdown. Rejected: re-raising with a bare 'raise Exception' | loses the cancellation identity the caller needs to unwind cleanly. Confidence: high Scope-risk: narrow --- src/mitocode/backend/mcp/mcp_client.py | 70 +++++++++++++++-- test/test_mcp_client.py | 100 ++++++++++++++++++++++++- 2 files changed, 161 insertions(+), 9 deletions(-) diff --git a/src/mitocode/backend/mcp/mcp_client.py b/src/mitocode/backend/mcp/mcp_client.py index 7c80e70..0c5fc93 100644 --- a/src/mitocode/backend/mcp/mcp_client.py +++ b/src/mitocode/backend/mcp/mcp_client.py @@ -17,6 +17,7 @@ import asyncio import json +import re from contextlib import AsyncExitStack from enum import Enum from typing import Any, Dict, List, Optional @@ -27,6 +28,38 @@ from mitocode.backend.exception.mitocode_exception import MitoCodeMcpToolCallError from mitocode.backend.mcp.mcp_config import McpServerConfig +#: Stands in for an argument payload that has been cut out of error text. +_REDACTED = "" + +#: The argument payload pydantic quotes back after rejecting a value. A server +#: that validates its input hands the offending value straight to the client in +#: its error text, and that value is whatever the model sent — routinely a token +#: or a file's contents. The payload runs up to ``input_type=`` when pydantic +#: supplies that trailing context and to the end of the line when it does not; +#: ``.*`` is greedy so a value that itself contains the text ``", input_type="`` +#: is consumed up to the real suffix instead of leaking its tail. ``.`` does not +#: cross a newline, so one error's payload never eats the next error's report. +_ARGUMENT_PAYLOAD = re.compile( + r"input_value=(?:(?P.*)(?P, input_type=[^\]\n]*)|[^\n]*)" +) + + +def _redact_arguments(text: str) -> str: + """Cut argument payloads out of error text while keeping it diagnosable. + + Only the value is removed: the tool, the server, the offending field name + and the description of what was expected all survive, so the message still + says *what* was wrong even though it no longer says *with what*. + + Args: + text: Error text as reported, possibly quoting the call's arguments. + + Returns: + The same text with every ``input_value=`` payload replaced by + :data:`_REDACTED`. + """ + return _ARGUMENT_PAYLOAD.sub(rf"input_value={_REDACTED}\g", text) + class McpConnectionState(str, Enum): """Connection state of a single :class:`McpClient`.""" @@ -102,9 +135,20 @@ async def connect(self) -> bool: and the child process alive as an orphan. Only ``Client.close()`` sets the transport's stop event and tears down the process group. + Cancellation is the same story with a different trigger: ``Ctrl-C``, a + TUI shutdown or an enclosing ``wait_for`` delivers + :class:`asyncio.CancelledError`, which derives from ``BaseException`` + and so bypasses an ``except Exception`` handler entirely. It is caught, + cleaned up after and then re-raised — swallowing a cancellation is not + this method's call to make, but orphaning a subprocess is not an option + either. + Returns: ``True`` when the client is connected and its tools are cached, ``False`` otherwise. + + Raises: + asyncio.CancelledError: If the caller cancels the connection. """ if self._state is McpConnectionState.CONNECTED: return True @@ -124,13 +168,17 @@ async def connect(self) -> bool: self._tools = list( await asyncio.wait_for(client.list_tools(), timeout=timeout) ) - except Exception as exc: # pylint: disable=broad-except + except BaseException as exc: # pylint: disable=broad-except # Cleanup runs outside the timed-out await, or the teardown itself - # would be cancelled and the subprocess orphaned after all. - await self._force_close() + # would be cancelled and the subprocess orphaned after all. The + # shield keeps a *repeat* cancellation — a second Ctrl-C, or an + # enclosing wait_for firing as well — from interrupting it too. + await asyncio.shield(self._force_close()) self._state = McpConnectionState.FAILED self._last_error = f"{type(exc).__name__}: {exc}" - return False + if isinstance(exc, Exception): + return False + raise self._state = McpConnectionState.CONNECTED self._last_error = "" @@ -206,8 +254,12 @@ async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> str: timeout=timeout, ) except ToolError as exc: + # The server's own text is the most likely place for the arguments to + # reappear: an argument-validation failure quotes the value it + # rejected, so the payload is cut before anything is interpolated. raise MitoCodeMcpToolCallError( - f"MCP tool '{tool_name}' on server '{server}' failed: {exc}" + f"MCP tool '{tool_name}' on server '{server}' failed: " + f"{_redact_arguments(str(exc))}" ) from exc except asyncio.TimeoutError as exc: raise MitoCodeMcpToolCallError( @@ -217,11 +269,13 @@ async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> str: except MitoCodeMcpToolCallError: raise except Exception as exc: # pylint: disable=broad-except - # Arguments are never echoed: they routinely carry file contents and - # credentials, and this text reaches the TUI error pane and logs. + # Arguments are never echoed here either: they routinely carry file + # contents and credentials, and this text reaches the TUI error pane + # and logs. A local validation failure can quote them just as a + # server-side one does, so the same redaction applies. raise MitoCodeMcpToolCallError( f"MCP tool '{tool_name}' on server '{server}' raised " - f"{type(exc).__name__}: {exc}" + f"{type(exc).__name__}: {_redact_arguments(str(exc))}" ) from exc return self._render_result(result) diff --git a/test/test_mcp_client.py b/test/test_mcp_client.py index 6e38685..99a0b8a 100644 --- a/test/test_mcp_client.py +++ b/test/test_mcp_client.py @@ -9,7 +9,11 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) from mitocode.backend.exception.mitocode_exception import MitoCodeMcpToolCallError -from mitocode.backend.mcp.mcp_client import McpClient, McpConnectionState +from mitocode.backend.mcp.mcp_client import ( + McpClient, + McpConnectionState, + _redact_arguments, +) from mitocode.backend.mcp.mcp_config import McpServerConfig @@ -179,6 +183,51 @@ async def close(self): assert client._client is None +@pytest.mark.asyncio +async def test_connect_cancellation_closes_the_transport_and_still_propagates(): + """Cancelling a handshake must tear the transport down, then keep cancelling. + + ``CancelledError`` derives from ``BaseException``, so it walks straight past + an ``except Exception`` handler: Ctrl-C, a TUI shutdown or an enclosing + ``wait_for`` would leave ``_client``/``_stack`` populated and stdio's + detached connect task holding the child process alive. Swallowing the + cancellation would be wrong too — cleanup, then propagate. + """ + client = McpClient(_slow_config(1.0)) + entered = asyncio.Event() + + class _HangingClient: + """Async context manager whose __aenter__ never completes.""" + + closed = False + + async def __aenter__(self): + entered.set() + await asyncio.Event().wait() + + async def __aexit__(self, *exc_info): + return False + + async def close(self): + type(self).closed = True + + built = _HangingClient() + client._build_fastmcp_client = lambda: built + + task = asyncio.ensure_future(client.connect()) + await asyncio.wait_for(entered.wait(), timeout=5) + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + assert built.closed is True, "cancellation orphaned the transport" + assert client._client is None + assert client._stack is None + assert client.tools == [] + assert client.state is McpConnectionState.FAILED + + @pytest.mark.asyncio async def test_connect_timeout_during_tool_discovery_is_honoured(): """Discovery is a second round trip and needs the connect deadline too. @@ -389,6 +438,55 @@ async def call_tool(self, name, arguments, **kwargs): assert excinfo.value is original +@pytest.mark.asyncio +async def test_tool_error_message_never_echoes_the_arguments_it_validated( + monkeypatch, +): + """A schema mistake must not turn the rejected argument into an error message. + + fastmcp re-raises pydantic's validation text as a ``ToolError``, and every + ``input_value=`` segment in it is the value the model just sent. This is + the *common* failure — the model calls a tool with the wrong shape — and the + generic ``except Exception`` path cannot help because it never sees it. + """ + secret = "sk-TOOL-ERROR-DO-NOT-LEAK" + server = FastMCP("probe") + + @server.tool + def login(token: str) -> str: + """Accept a token.""" + return "ok" + + client = McpClient(_in_memory_config()) + monkeypatch.setattr(client, "_build_fastmcp_client", lambda: Client(server)) + assert await client.connect() is True + + with pytest.raises(MitoCodeMcpToolCallError) as excinfo: + await client.call_tool("login", {"token": [secret]}) + + message = str(excinfo.value) + assert secret not in message, f"the rejected argument leaked: {message}" + # Redaction must cost diagnosis nothing: which tool, on which server, which + # field, and what was expected of it all have to survive. + assert "login" in message + assert "probe" in message + assert "token" in message + assert "should be a valid string" in message + await client.disconnect() + + +def test_redaction_covers_an_input_value_with_no_type_suffix(): + """Text that quotes an argument without pydantic's type suffix is still cut. + + ``input_type=`` is the common tail, but the rule must not depend on it: an + argument payload is never safe to echo, whatever surrounds it. + """ + redacted = _redact_arguments("server said no: input_value=sk-BARE-DO-NOT-LEAK") + + assert "sk-BARE-DO-NOT-LEAK" not in redacted + assert "server said no" in redacted + + @pytest.mark.asyncio async def test_call_tool_joins_text_content_blocks(): client = McpClient(_in_memory_config()) From 79b08098bc706c959c770554bf764f4914cb8c57 Mon Sep 17 00:00:00 2001 From: Allen Wang Date: Sat, 17 Oct 2026 12:51:17 +0800 Subject: [PATCH 12/20] test(mcp): pin argument redaction and subprocess teardown against mutation The suite already reported 100% statement and branch coverage of mcp_client.py, so this adds no lines to the coverage report. It pins what that report cannot see: behaviours the existing tests let through. Redaction (the credential leak shipped in c18605f): - adversarial payloads carrying a decoy ", input_type=" marker, both with and without trailing report text after the decoy - two payloads on one line, which pins the greedy match to the *last* marker -- a non-greedy pattern silently leaves the second argument in the log - a multi-error, multi-line pydantic report, so one payload cannot eat the next error's field name and expected type - empty and suffixless payloads - end-to-end through both call_tool error branches, asserting the tool name, server name, field name and real input_type= suffix all survive, so redaction cannot be "achieved" by deleting the report Teardown (the orphaned subprocess shipped in c18605f): - repeat cancellation while Client.close() is in flight: only a shielded teardown runs to completion, so this is what pins the shield itself - a repeat cancellation skips the FAILED bookkeeping entirely - a BaseException that is neither Exception nor CancelledError still tears the transport down and propagates, which is what pins the handler on BaseException rather than (Exception, CancelledError) Also pins behaviour that no test touched: the timeout keyword reaching the server call, ensure_ascii=False and default=str in result rendering, last_error being cleared by a successful retry, a teardown error being overwritten by the handshake error, the real _build_fastmcp_client wiring, and the not-connected report naming both server and tool. Verified against 24 mutants of a scratch copy of the module under /tmp; all 24 are killed. Coverage stays 100% statement and branch. --- test/test_mcp_client.py | 548 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 548 insertions(+) diff --git a/test/test_mcp_client.py b/test/test_mcp_client.py index 99a0b8a..5af24aa 100644 --- a/test/test_mcp_client.py +++ b/test/test_mcp_client.py @@ -5,6 +5,7 @@ import pytest from fastmcp import Client, FastMCP +from fastmcp.exceptions import ToolError sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) @@ -82,6 +83,74 @@ async def aclose(self) -> None: raise self.error +class _HandshakeHangingClient: + """Fake client whose handshake never completes and whose close is slow. + + The slow ``close`` is the whole point: it opens a window in which a second + cancellation can land while the teardown started by the first one is still + awaiting, which is the only way to tell a shielded teardown from an + unshielded one. + """ + + def __init__(self) -> None: + """Arm the events a test needs to synchronise on.""" + self.entered = asyncio.Event() + self.close_started = asyncio.Event() + self.close_finished = asyncio.Event() + + async def __aenter__(self): + """Announce that the handshake started, then wait forever.""" + self.entered.set() + await asyncio.Event().wait() + + async def __aexit__(self, *exc_info): + """Never claim to have suppressed anything.""" + return False + + async def close(self) -> None: + """Record that teardown began, yield to the loop, then record it ended.""" + self.close_started.set() + await asyncio.sleep(0.05) + self.close_finished.set() + + +class _ShutdownSignal(BaseException): + """A shutdown request that is neither an ``Exception`` nor a cancellation. + + Stand-in for the third ``BaseException`` an awaited operation can raise: a + framework's own "stop everything" signal. ``KeyboardInterrupt`` and + ``SystemExit`` cannot be used here because asyncio re-raises those out of + the event loop itself instead of delivering them to the awaiting coroutine. + """ + + +async def _cancel_twice_during_teardown( + client: McpClient, built: _HandshakeHangingClient +) -> None: + """Cancel a hanging ``connect`` again while its own teardown is in flight. + + The second cancellation is timed off ``close_started`` so it cannot race + the teardown: without it the test would pass whether or not the teardown is + shielded. + + Args: + client: The client whose :meth:`connect` must be cancelled twice. + built: The fake transport installed as the client's build seam. + """ + task = asyncio.ensure_future(client.connect()) + await asyncio.wait_for(built.entered.wait(), timeout=5) + task.cancel() + + async def _cancel_mid_teardown() -> None: + await asyncio.wait_for(built.close_started.wait(), timeout=5) + task.cancel() + + canceller = asyncio.ensure_future(_cancel_mid_teardown()) + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.wait_for(canceller, timeout=5) + + @pytest.mark.asyncio async def test_connect_discovers_tools_from_server(monkeypatch): config = _in_memory_config() @@ -632,3 +701,482 @@ def _explode(): assert client._client is None assert client._stack is None assert client.tools == [] + + +@pytest.mark.asyncio +async def test_connect_repeat_cancellation_still_completes_the_transport_teardown(): + """A second cancellation must not abort the teardown the first one started. + + ``connect`` runs its cleanup outside the timed-out await and shields it, so + a repeat cancellation -- a second Ctrl-C, or an enclosing ``wait_for`` + firing in the same tick -- cannot interrupt ``Client.close()`` midway and + leave stdio's detached connect task holding the child process alive. Only + the fact that ``close`` *ran to completion* distinguishes the shield from a + plain await; both leave the handles cleared. + """ + client = McpClient(_in_memory_config()) + built = _HandshakeHangingClient() + client._build_fastmcp_client = lambda: built + + await _cancel_twice_during_teardown(client, built) + + await asyncio.wait_for(built.close_finished.wait(), timeout=5) + assert client._client is None + assert client._stack is None + assert client.tools == [] + + +@pytest.mark.asyncio +async def test_connect_repeat_cancellation_skips_the_failed_bookkeeping(): + """A repeat cancellation propagates before ``connect`` records the failure. + + The cancellation is re-raised out of the shielded await, so the two + statements that would set ``FAILED`` and record the error never run. The + object is still consistent -- no client, no stack, no tools -- but a reader + of :attr:`last_error` must not expect to find the cancellation logged here. + """ + client = McpClient(_in_memory_config()) + built = _HandshakeHangingClient() + client._build_fastmcp_client = lambda: built + + await _cancel_twice_during_teardown(client, built) + + assert client.state is McpConnectionState.DISCONNECTED + assert client.last_error == "" + + +@pytest.mark.asyncio +async def test_connect_shutdown_signal_closes_the_transport_and_still_propagates(): + """A ``BaseException`` that is not an ``Exception`` must be cleaned up after. + + ``CancelledError`` is only the best-known member of that family. Catching + ``(Exception, CancelledError)`` would pass every cancellation test while + leaving an application shutdown signal to orphan the subprocess, so the + handler has to be keyed on ``BaseException`` itself. + """ + client = McpClient(_in_memory_config()) + + class _ShuttingDownClient: + """Client whose handshake fails with a shutdown signal.""" + + def __init__(self) -> None: + """Start with the transport still open.""" + self.closed = False + + async def __aenter__(self): + """Abort the handshake the way a framework shutdown would.""" + raise _ShutdownSignal("application is shutting down") + + async def __aexit__(self, *exc_info): + """Never claim to have suppressed anything.""" + return False + + async def close(self) -> None: + """Record that the transport was released.""" + self.closed = True + + built = _ShuttingDownClient() + client._build_fastmcp_client = lambda: built + + with pytest.raises(_ShutdownSignal): + await client.connect() + + assert built.closed is True, "shutdown signal orphaned the transport" + assert client.state is McpConnectionState.FAILED + assert client.last_error == "_ShutdownSignal: application is shutting down" + assert client._client is None + assert client._stack is None + assert client.tools == [] + + +@pytest.mark.asyncio +async def test_connect_failure_reports_the_handshake_error_not_the_teardown_error(): + """A failing teardown must not replace the error that explains the failure. + + ``_force_close`` records cleanup problems in ``last_error`` and ``connect`` + then overwrites it with the exception that caused the failure. The message a + user reads is therefore the handshake error alone; the fact that teardown + *also* failed is only visible in the teardown's own effects. + """ + client = McpClient(_in_memory_config()) + + class _ExplodingHandshake: + """Client that fails to connect and cannot be closed either.""" + + def __init__(self) -> None: + """Start with the transport still open.""" + self.closed = False + + async def __aenter__(self): + """Refuse the handshake.""" + raise RuntimeError("handshake exploded") + + async def __aexit__(self, *exc_info): + """Never claim to have suppressed anything.""" + return False + + async def close(self) -> None: + """Attempt the teardown, which fails.""" + self.closed = True + raise OSError("process group wedged") + + built = _ExplodingHandshake() + client._build_fastmcp_client = lambda: built + + assert await client.connect() is False + + assert built.closed is True + assert client.last_error == "RuntimeError: handshake exploded" + assert client.state is McpConnectionState.FAILED + + +@pytest.mark.asyncio +async def test_connect_after_a_failure_clears_the_stale_error(monkeypatch): + """A successful retry must not leave the previous failure in the report.""" + client = McpClient(_in_memory_config()) + + class _ExplodingHandshake: + """Client that never completes a handshake.""" + + async def __aenter__(self): + """Refuse the handshake.""" + raise RuntimeError("transient outage") + + async def __aexit__(self, *exc_info): + """Never claim to have suppressed anything.""" + return False + + async def close(self) -> None: + """Succeed at the teardown.""" + return None + + client._build_fastmcp_client = lambda: _ExplodingHandshake() + assert await client.connect() is False + assert client.last_error == "RuntimeError: transient outage" + + server = FastMCP("probe") + + @server.tool + def ping() -> str: + """Report that the server is alive.""" + return "pong" + + monkeypatch.setattr(client, "_build_fastmcp_client", lambda: Client(server)) + + assert await client.connect() is True + assert client.state is McpConnectionState.CONNECTED + assert client.last_error == "" + assert sorted(t.name for t in client.tools) == ["ping"] + await client.disconnect() + + +def test_redaction_cuts_the_payload_up_to_the_last_injected_type_marker(): + """A value carrying pydantic's own suffix text is cut up to that decoy. + + The pattern consumes a payload up to ``, input_type=``; when the argument + itself contains that text the greedy match runs to the *last* occurrence on + the line. Everything before the decoy is cut, which is the guarantee the + module documents: text after the marker is indistinguishable from the + framework's suffix and is deliberately kept so the report stays readable. + """ + redacted = _redact_arguments("input_value='sk-PART1\", input_type=sk-PART2LEAK'") + + assert "sk-PART1" not in redacted + assert redacted == "input_value=, input_type=sk-PART2LEAK'" + + +def test_redaction_keeps_the_report_after_an_injected_marker(): + """Surrounding report text must survive a payload that fakes the suffix. + + Redaction that eats the sentence, the field name and the expected type + would make the error undiagnosable -- and a test asserting only "the + secret is gone" would pass just as happily against such an implementation. + """ + redacted = _redact_arguments( + "tool 'login' failed on field 'token': " + 'input_value=sk-DECOY", input_type=str; expected a valid string]' + ) + + assert "sk-DECOY" not in redacted + assert redacted == ( + "tool 'login' failed on field 'token': " + "input_value=, input_type=str; expected a valid string]" + ) + + +def test_redaction_cuts_each_error_of_a_multiline_report_independently(): + """One payload must not eat the next error's report on a following line. + + ``.`` does not cross a newline, so a validation report naming two offending + fields loses neither of them -- and each keeps the type pydantic expected. + """ + report = ( + "2 validation errors for login\n" + "token\n" + " Input should be a valid string [type=string_type, " + "input_value='sk-FIRST', input_type=str]\n" + "count\n" + " Input should be a valid integer [type=string_type, " + "input_value='sk-SECOND', input_type=int]" + ) + + redacted = _redact_arguments(report) + + assert redacted == ( + "2 validation errors for login\n" + "token\n" + " Input should be a valid string [type=string_type, " + "input_value=, input_type=str]\n" + "count\n" + " Input should be a valid integer [type=string_type, " + "input_value=, input_type=int]" + ) + + +def test_redaction_handles_an_empty_payload(): + """An ``input_value=`` with nothing after it is still a payload.""" + assert _redact_arguments( + "Input should be a valid string [type=string_type, input_value=, input_type=str]" + ) == ( + "Input should be a valid string [type=string_type, " + "input_value=, input_type=str]" + ) + assert _redact_arguments("nothing here: input_value=") == ( + "nothing here: input_value=" + ) + + +def test_redaction_cuts_a_suffixless_payload_to_the_end_of_its_line(): + """With no ``, input_type=`` to stop at, the whole rest of the line is cut. + + Stopping at the next space instead would leave the tail of the value -- and + a token is often longer than its first word. + """ + assert _redact_arguments("input_value=sk-BARE extra words follow") == ( + "input_value=" + ) + assert _redact_arguments("input_value=sk-BARE extra words follow\nnext line") == ( + "input_value=\nnext line" + ) + + +@pytest.mark.asyncio +async def test_call_tool_tool_error_redacts_the_payload_and_keeps_the_report(): + """The common validation failure must lose the value and keep everything else. + + This is the ``ToolError`` branch, which the generic handler cannot see, so + it carries the full pydantic report: the tool, the server, the offending + field and the expected type all have to remain, while the rejected value + does not. + """ + client = McpClient(_in_memory_config()) + secret = "sk-ADVERSARIAL-DO-NOT-LEAK" + + class _ValidatingCall: + """Client that rejects its argument the way fastmcp reports it.""" + + async def call_tool(self, name, arguments, **kwargs): + raise ToolError( + "1 validation error for login\n" + "token\n" + " Input should be a valid string [type=string_type, " + f"input_value='{secret}\", input_type=str]" + ) + + client._client = _ValidatingCall() + + with pytest.raises(MitoCodeMcpToolCallError) as excinfo: + await client.call_tool("login", {"token": secret}) + + message = str(excinfo.value) + assert secret not in message, f"the rejected argument leaked: {message}" + assert "MCP tool 'login' on server 'probe' failed" in message + assert "token" in message + assert "should be a valid string" in message + assert "input_type=str" in message + + +@pytest.mark.asyncio +async def test_call_tool_unexpected_error_redacts_the_payload_and_keeps_the_report(): + """A local validation failure quotes the arguments just as a server one does. + + The generic handler renders the exception type before the text, so the + redaction has to survive an extra prefix and must not cost the surrounding + report anything. + """ + client = McpClient(_in_memory_config()) + secret = "sk-LOCAL-VALIDATION-DO-NOT-LEAK" + + class _LocallyValidatingCall: + """Client that fails locally while validating its arguments.""" + + async def call_tool(self, name, arguments, **kwargs): + raise ValueError( + "field 'token' rejected\n" + " Input should be a valid string [type=string_type, " + f"input_value='{secret}', input_type=str]" + ) + + client._client = _LocallyValidatingCall() + + with pytest.raises(MitoCodeMcpToolCallError) as excinfo: + await client.call_tool("login", {"token": secret}) + + message = str(excinfo.value) + assert secret not in message, f"the rejected argument leaked: {message}" + assert "MCP tool 'login' on server 'probe' raised ValueError" in message + assert "token" in message + assert "should be a valid string" in message + assert "input_type=str" in message + + +@pytest.mark.asyncio +async def test_call_tool_forwards_the_configured_timeout_to_the_server_call(): + """The per-call deadline must reach the server call, not just the local wait. + + The deadline is applied twice on purpose: the outer ``wait_for`` bounds the + local await and the ``timeout`` keyword bounds the request the transport has + actually sent. Dropping the keyword leaves a hanging server unbounded + server-side while every local test still passes. + """ + client = McpClient( + McpServerConfig.parse( + "slowtool", {"url": "https://x.invalid/mcp", "timeout": 12.5} + ) + ) + seen = {} + + class _RecordingCall: + """Client that records how it was called.""" + + async def call_tool(self, name, arguments, **kwargs): + seen["name"] = name + seen["arguments"] = arguments + seen["timeout"] = kwargs.get("timeout", "absent") + return None + + client._client = _RecordingCall() + + await client.call_tool("echo", {"text": "hi"}) + + assert seen == {"name": "echo", "arguments": {"text": "hi"}, "timeout": 12.5} + + +@pytest.mark.asyncio +async def test_call_tool_keeps_non_ascii_structured_data_readable(): + """Structured output must not be escaped into ``\\uXXXX`` noise. + + ``ensure_ascii`` defaults to true, which turns a Japanese or accented field + into escape sequences; a result rendered that way still parses, so nothing + downstream notices except the reader. + """ + client = McpClient(_in_memory_config()) + + class _Result: + data = {"memo": "café", "city": "東京"} + content = [] + + class _Call: + async def call_tool(self, name, arguments, **kwargs): + return _Result() + + client._client = _Call() + + assert await client.call_tool("echo", {}) == '{"memo": "café", "city": "東京"}' + + +@pytest.mark.asyncio +async def test_call_tool_renders_an_unserialisable_result_value_with_str(): + """A result carrying a non-JSON value must render, not raise. + + ``json.dumps`` refuses a set, a handle or a custom object, and the failure + would escape ``call_tool`` as a bare ``TypeError`` because rendering happens + outside the error handling. ``default=str`` is what keeps that reachable. + """ + client = McpClient(_in_memory_config()) + + class _Handle: + """A value the model cannot receive as JSON.""" + + def __str__(self) -> str: + return "opaque-handle" + + class _Result: + data = {"handle": _Handle(), "ids": 3} + content = [] + + class _Call: + async def call_tool(self, name, arguments, **kwargs): + return _Result() + + client._client = _Call() + + assert await client.call_tool("echo", {}) == ( + '{"handle": "opaque-handle", "ids": 3}' + ) + + +def test_redaction_cuts_the_last_payload_when_a_line_reports_several(): + """Every payload on a line must be cut, not just the first one. + + Matching greedily to the *last* ``, input_type=`` is what makes a payload + containing a decoy marker safe: stopping at the first marker instead would + leave the remainder of the line -- including the next payload -- untouched, + so a hostile value that merely looks like the framework's suffix would hand + the following argument straight back to the log. + """ + redacted = _redact_arguments( + "probe failed: input_value=sk-FIRST, input_type=str; " + "input_value=sk-SECOND, input_type=int" + ) + + assert "sk-FIRST" not in redacted + assert "sk-SECOND" not in redacted + assert redacted == "probe failed: input_value=, input_type=int" + + +def test_build_fastmcp_client_wires_the_configured_transport_and_name(): + """The real builder must use the config, not a hardcoded transport. + + Every connect test swaps this method out, so without this the transport a + server is reached through -- HTTP versus stdio, and its command line -- is + never exercised at all. + """ + remote = McpClient( + McpServerConfig.parse( + "remote", {"url": "https://remote.invalid/mcp", "timeout": 7.5} + ) + )._build_fastmcp_client() + local = McpClient( + McpServerConfig.parse( + "local", {"command": "my-server", "args": ["--flag", "value"]} + ) + )._build_fastmcp_client() + + assert isinstance(remote, Client) + assert remote.name == "remote" + assert type(remote.transport).__name__ == "StreamableHttpTransport" + assert remote.transport.url == "https://remote.invalid/mcp" + assert isinstance(local, Client) + assert local.name == "local" + assert type(local.transport).__name__ == "StdioTransport" + assert local.transport.command == "my-server" + assert local.transport.args == ["--flag", "value"] + + +@pytest.mark.asyncio +async def test_call_tool_names_the_server_and_the_tool_when_not_connected(): + """The not-connected report must say which server and which tool. + + With several servers configured, "not connected" on its own does not say + which one to fix; the caller needs both halves of the address before it + built one. + """ + client = McpClient(_in_memory_config()) + + with pytest.raises(MitoCodeMcpToolCallError) as excinfo: + await client.call_tool("echo", {"text": "hi"}) + + message = str(excinfo.value) + assert "MCP server 'probe' is not connected" in message + assert "cannot call tool 'echo'" in message From 2e51ccb1ed3282c24465c1ad38510206d90f98fd Mon Sep 17 00:00:00 2001 From: Allen Wang Date: Sat, 17 Oct 2026 13:22:02 +0800 Subject: [PATCH 13/20] test(mcp): stop asserting a redaction leak, and pin _force_close's narrow catch Three review fixes, all confined to test/test_mcp_client.py. An assertion was enforcing a credential leak. test_redaction_cuts_the_payload_up_to_the_last_injected_type_marker asserted the redacted output character-for-character, including that `sk-PART2LEAK` survives a payload that supplies its own `", input_type="` decoy. That green-lights a leak in a suite whose binding constraint is never to leak argument values, and worse, it makes a plausible fix -- requiring the suffix to authenticate as a real type token -- fail as if it were a regression. The test now asserts the security property and the documented prefix, and names the residual leak in a comment saying that closing it requires DELETING the assertion. Coverage of the decoy shape is not lost: test_redaction_keeps_the_report_after_an_injected_marker still pins the exact output and now kills the suffix-dropping and suffix-eating mutants that the weakened assertion used to catch. _force_close's handlers were untested one call frame below connect's. _mcp_client.py:205 and :210 catch Exception while connect's handler at :171 catches BaseException, and widening the teardown handlers to BaseException survived the whole suite. test_force_close_does_not_swallow_a_shutdown_ signal_raised_while_closing pins that a signal from close() escapes disconnect() and leaves last_error untouched -- which is what stops disconnect() from swallowing an in-flight cancellation. Verified: that mutation now fails the suite. test_redaction_cuts_the_last_payload_when_a_line_reports_several was misnamed; it cuts everything up to the last marker, not the last payload. Renamed to test_redaction_cuts_up_to_the_last_marker_when_a_line_reports_ several, docstring included. Also softened the builder test's docstring, which overclaimed full config wiring: fastmcp 4.1 accepts Client(timeout=) but keeps it off the client, so that keyword stays unpinned and the docstring now says so. 48 in the file, 279 repo-wide, mcp_client.py still 100% statement and branch. 14 mutants re-checked against a scratch copy under /tmp, all killed. --- test/test_mcp_client.py | 68 +++++++++++++++++++++++++++++++++++------ 1 file changed, 58 insertions(+), 10 deletions(-) diff --git a/test/test_mcp_client.py b/test/test_mcp_client.py index 5af24aa..c19d4be 100644 --- a/test/test_mcp_client.py +++ b/test/test_mcp_client.py @@ -608,6 +608,47 @@ async def test_force_close_still_closes_the_stack_when_client_close_raises(): assert client.last_error == "close failed: RuntimeError: client wedged" +@pytest.mark.asyncio +async def test_force_close_does_not_swallow_a_shutdown_signal_raised_while_closing(): + """A ``BaseException`` from ``close()`` must escape, not be recorded away. + + ``_force_close`` deliberately catches ``Exception``, not ``BaseException``, + and the asymmetry with :meth:`connect`'s handler is the point: a shutdown + signal or a cancellation delivered *while* the teardown is running has to + keep travelling, or ``disconnect()`` swallows an in-flight cancellation. + Widening the handler to ``BaseException`` is therefore a regression, and + this is the test that says so. + """ + client = McpClient(_in_memory_config()) + + class _ShutdownOnClose: + """Client that signals a shutdown part-way through closing.""" + + def __init__(self) -> None: + """Start with the transport still open.""" + self.closed = False + + async def close(self) -> None: + """Attempt the teardown, then announce a shutdown.""" + self.closed = True + raise _ShutdownSignal("stopping during teardown") + + built = _ShutdownOnClose() + client._client = built + + with pytest.raises(_ShutdownSignal): + await client.disconnect() + + assert built.closed is True + # Not recorded: the recording branch is Exception-only, so a signal leaves + # last_error untouched rather than looking like a cleanup problem. + assert client.last_error == "" + # The handles were dropped before the await, so abandoning the teardown + # part-way still leaks no reference to the transport. + assert client._client is None + assert client._stack is None + + @pytest.mark.asyncio async def test_force_close_records_a_failing_stack_close(): """A stack that cannot unwind is recorded instead of aborting the teardown.""" @@ -875,14 +916,18 @@ def test_redaction_cuts_the_payload_up_to_the_last_injected_type_marker(): The pattern consumes a payload up to ``, input_type=``; when the argument itself contains that text the greedy match runs to the *last* occurrence on - the line. Everything before the decoy is cut, which is the guarantee the - module documents: text after the marker is indistinguishable from the - framework's suffix and is deliberately kept so the report stays readable. + the line. What precedes the decoy is therefore cut, which is the guarantee + this test pins -- deliberately as a prefix, not as the whole output, so that + closing the leak below shows up as a deletion rather than as a regression. """ redacted = _redact_arguments("input_value='sk-PART1\", input_type=sk-PART2LEAK'") assert "sk-PART1" not in redacted - assert redacted == "input_value=, input_type=sk-PART2LEAK'" + # KNOWN LIMITATION (mcp_client.py:39-41): text after a payload-supplied + # decoy marker is currently NOT redacted -- `sk-PART2LEAK` survives here, + # and this assertion is what lets it. Fixing the leak requires DELETING + # this assertion, not editing it to expect the fix. + assert redacted.startswith("input_value=") def test_redaction_keeps_the_report_after_an_injected_marker(): @@ -1116,10 +1161,10 @@ async def call_tool(self, name, arguments, **kwargs): ) -def test_redaction_cuts_the_last_payload_when_a_line_reports_several(): - """Every payload on a line must be cut, not just the first one. +def test_redaction_cuts_up_to_the_last_marker_when_a_line_reports_several(): + """Everything up to the *last* marker on a line is cut, not just the first. - Matching greedily to the *last* ``, input_type=`` is what makes a payload + Matching greedily to the last ``, input_type=`` is what makes a payload containing a decoy marker safe: stopping at the first marker instead would leave the remainder of the line -- including the next payload -- untouched, so a hostile value that merely looks like the framework's suffix would hand @@ -1136,11 +1181,14 @@ def test_redaction_cuts_the_last_payload_when_a_line_reports_several(): def test_build_fastmcp_client_wires_the_configured_transport_and_name(): - """The real builder must use the config, not a hardcoded transport. + """The real builder must take its transport and name from the config. Every connect test swaps this method out, so without this the transport a - server is reached through -- HTTP versus stdio, and its command line -- is - never exercised at all. + server is actually reached through -- HTTP versus stdio, the endpoint, and + the command line -- is never exercised. The ``timeout`` keyword is + deliberately not asserted: fastmcp 4.1 accepts it but keeps it off the + client, so pinning it would mean reaching into private attributes, and the + deadline is already pinned where it is observable, in ``call_tool``. """ remote = McpClient( McpServerConfig.parse( From 6d736493c54ce50c9b9384b41e65013283a85cc4 Mon Sep 17 00:00:00 2001 From: Allen Wang Date: Sat, 17 Oct 2026 13:37:05 +0800 Subject: [PATCH 14/20] feat(tools): let ToolManager accept MCP tool catalogues register() now honours a caller-supplied parameters_schema because MCP servers ship their own JSON Schema that annotation inference cannot derive. register_mcp_client() namespaces tools as mcp____ and refuses to overwrite an existing registration; unregister() lets a later layer drop a server's tools again. Constraint: McpClient referenced only under TYPE_CHECKING so tool_manager imports without fastmcp installed Rejected: Registering via AgentSession's tools_map/tool_schemas kwargs | those parameters are deprecated Confidence: high Scope-risk: moderate --- src/mitocode/backend/tool/tool_manager.py | 140 +++++++++++++- test/test_tool_manager.py | 223 ++++++++++++++++++++++ 2 files changed, 360 insertions(+), 3 deletions(-) diff --git a/src/mitocode/backend/tool/tool_manager.py b/src/mitocode/backend/tool/tool_manager.py index 53e65a7..58279e2 100644 --- a/src/mitocode/backend/tool/tool_manager.py +++ b/src/mitocode/backend/tool/tool_manager.py @@ -2,6 +2,10 @@ Converts Python type annotations into JSON Schema tool definitions and keeps separate human/agent-facing documentation per registered tool. + +A caller may also supply the schema outright, which is what tools +discovered on an MCP server need: their arguments are described by the +server, not by anything introspectable in this process. """ import inspect @@ -12,12 +16,15 @@ Dict, List, Optional, + TYPE_CHECKING, Union, get_args, get_origin, Literal, ) +from loguru import logger as loguru_logger + try: import types @@ -28,6 +35,42 @@ from mitocode.backend.exception.mitocode_exception import MitoCodeToolDefinitionError +if TYPE_CHECKING: + # Imported for type checking only. The runtime stays on duck typing so + # this module keeps importing in an environment without fastmcp installed. + from mitocode.backend.mcp.mcp_client import McpClient + +logger = loguru_logger.bind(name=__name__) + +#: Prefix every namespaced MCP tool name carries, so a remote tool is +#: recognisably remote and cannot be confused with a built-in one. +MCP_TOOL_PREFIX = "mcp__" + +#: Separator between the server name and the tool name inside that prefix. +MCP_TOOL_SEPARATOR = "__" + + +def _normalize_input_schema(raw: Any) -> Dict[str, Any]: + """Coerce an MCP tool's advertised input schema into a usable JSON Schema. + + Servers are not required to send a complete schema. This fills in the + pieces :meth:`ToolManager.tool_schemas` always needs so a malformed + schema cannot poison an entire LLM request. + + Args: + raw: The ``input_schema`` value advertised by the MCP server. + + Returns: + A dict with at least ``type`` and ``properties`` keys. + """ + if not isinstance(raw, dict): + return {"type": "object", "properties": {}} + normalized = dict(raw) + normalized.setdefault("type", "object") + if not isinstance(normalized.get("properties"), dict): + normalized["properties"] = {} + return normalized + def _annotation_to_schema(annotation: Any) -> Dict[str, Any]: """Convert a Python type annotation to a JSON Schema fragment. @@ -219,12 +262,13 @@ def __init__(self): """ self._tools: Dict[str, ToolSpec] = {} - def tool( + def tool( # pylint: disable=too-many-arguments self, name: Optional[str] = None, *, human_doc: Optional[str] = None, agent_doc: Optional[str] = None, + parameters_schema: Optional[Dict[str, Any]] = None, display_action: Optional[str] = None, display_param: Optional[str] = None, ): @@ -236,6 +280,9 @@ def tool( human_doc: Human-readable description for UI purposes. agent_doc: Description sent to the LLM in tool schemas. Defaults to *human_doc* if not provided. + parameters_schema: JSON Schema of accepted arguments. + Defaults to one derived from the function's type + annotations. display_action: Optional action key used by the TUI for styling the tool invocation. display_param: Parameter whose value should be highlighted @@ -251,6 +298,7 @@ def _decorator(func: Callable): name=name, human_doc=human_doc, agent_doc=agent_doc, + parameters_schema=parameters_schema, display_action=display_action, display_param=display_param, ) @@ -258,13 +306,14 @@ def _decorator(func: Callable): return _decorator - def register( + def register( # pylint: disable=too-many-arguments self, func: Callable, *, name: Optional[str] = None, human_doc: Optional[str] = None, agent_doc: Optional[str] = None, + parameters_schema: Optional[Dict[str, Any]] = None, display_action: Optional[str] = None, display_param: Optional[str] = None, ) -> ToolSpec: @@ -274,6 +323,12 @@ def register( :func:`declare_tool`, which in turn takes precedence over the function's own name and docstring. + ``parameters_schema`` is needed for arguments that cannot be read off + a signature — an MCP server advertises its own JSON Schema — so when it + is supplied it is used verbatim instead of being derived, with a + ``parameters_schema`` entry in the ``declare_tool`` metadata as the + fallback and annotation inference last. + Returns: The registered :class:`ToolSpec`. """ @@ -292,13 +347,92 @@ def register( func=func, human_doc=human_description, agent_doc=agent_description, - parameters_schema=_build_parameters_schema(func), + parameters_schema=( + parameters_schema + if parameters_schema is not None + else (meta.get("parameters_schema") or _build_parameters_schema(func)) + ), display_action=display, display_param=param, ) self._tools[tool_name] = spec return spec + def unregister(self, name: str) -> bool: + """Remove a tool from the registry. + + Args: + name: The tool name to remove. + + Returns: + True if a tool was registered under *name* and has been removed, + False if no such tool existed. + """ + return self._tools.pop(name, None) is not None + + @staticmethod + def build_mcp_tool_name(server_name: str, tool_name: str) -> str: + """Compose the namespaced MitoCode name for an MCP tool. + + Args: + server_name: The MCP server's configured name. + tool_name: The tool's name as advertised by that server. + + Returns: + ``mcp____``. + """ + return f"{MCP_TOOL_PREFIX}{server_name}{MCP_TOOL_SEPARATOR}{tool_name}" + + def register_mcp_client(self, client: "McpClient") -> List[ToolSpec]: + """Register every tool an MCP client has discovered. + + The client must already be connected so that its cached tool catalogue + is populated. A tool whose composed name is already taken is skipped + with a warning rather than overwriting the existing registration. + + Args: + client: A connected MCP client exposing ``server_name`` and ``tools``. + + Returns: + The specs that were successfully registered, in discovery order. + """ + registered: List[ToolSpec] = [] + for tool in client.tools: + tool_name = self.build_mcp_tool_name(client.server_name, tool.name) + if self.get_spec(tool_name) is not None: + logger.warning( + f"Skipping MCP tool '{tool_name}': name already registered" + ) + continue + + async def _invoke(_client=client, _tool_name=tool.name, **kwargs): + """Forward the model's keyword arguments to the server tool. + + The default arguments capture *this* iteration's client and + tool name. Closing over the loop variables instead would + route every registered tool to whichever tool came last. + """ + return await _client.call_tool(_tool_name, kwargs) + + description = (tool.description or "").strip() + if not description: + description = ( + f"Tool '{tool.name}' provided by MCP server " + f"'{client.server_name}'." + ) + + registered.append( + self.register( + _invoke, + name=tool_name, + human_doc=f"[{client.server_name}] {description}", + agent_doc=f"[{client.server_name}] {description}", + parameters_schema=_normalize_input_schema(tool.input_schema), + display_action=f"mcp:{client.server_name}", + ) + ) + return registered + def tools_map(self) -> Dict[str, Callable]: """Return a name-to-callable mapping of all registered tools.""" return {name: spec.func for name, spec in self._tools.items()} diff --git a/test/test_tool_manager.py b/test/test_tool_manager.py index f5e1887..96dd373 100644 --- a/test/test_tool_manager.py +++ b/test/test_tool_manager.py @@ -1,6 +1,10 @@ import os +import subprocess import sys import unittest.mock as mock +from types import SimpleNamespace + +import pytest sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) @@ -13,6 +17,30 @@ _build_parameters_schema, ) +#: Imports tool_manager in a process where fastmcp cannot be resolved at all, +#: so the one-way tool_manager -> mcp -> fastmcp dependency can be checked for +#: real rather than by reading the source. Run out-of-process because the test +#: session has already imported fastmcp. +_IMPORT_WITHOUT_FASTMCP = """ +import importlib.abc +import sys + + +class _BlockFastMcp(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname == "fastmcp" or fullname.startswith("fastmcp."): + raise ImportError("fastmcp is not installed") + return None + + +sys.meta_path.insert(0, _BlockFastMcp()) +import mitocode.backend.tool.tool_manager + +assert "fastmcp" not in sys.modules, "fastmcp was imported" +leaked = [m for m in sys.modules if m.startswith("mitocode.backend.mcp")] +assert not leaked, f"mcp package imported at runtime: {leaked}" +""" + def test_tool_manager_decorator_registration_and_docs(): manager = ToolManager() @@ -102,3 +130,198 @@ def test_tool(city: str = "Beijing", country: str = "China"): assert props["city"]["default"] == "Beijing" assert props["country"]["default"] == "China" assert props["city"] is not props["country"] + + +def test_tool_manager_imports_without_fastmcp_installed(): + """tool_manager must stay importable where fastmcp is absent. + + Dependency runs one way — tool_manager -> mcp -> fastmcp — so the MCP + reference in ``register_mcp_client``'s annotation must stay under + ``TYPE_CHECKING`` and the module must duck-type the client at runtime. + """ + env = dict(os.environ) + env["PYTHONPATH"] = os.pathsep.join( + [ + os.path.join(os.path.dirname(__file__), "..", "src"), + env.get("PYTHONPATH", ""), + ] + ) + result = subprocess.run( + [sys.executable, "-c", _IMPORT_WITHOUT_FASTMCP], + capture_output=True, + text=True, + env=env, + check=False, + ) + assert result.returncode == 0, result.stderr + + +def test_register_accepts_explicit_parameters_schema(): + manager = ToolManager() + schema = {"type": "object", "properties": {"city": {"type": "string"}}} + + def get_weather(**kwargs): + return kwargs + + spec = manager.register(get_weather, name="get_weather", parameters_schema=schema) + assert spec.parameters_schema == schema + assert manager.tool_schemas("agent")[0]["function"]["parameters"] == schema + + +def test_register_without_parameters_schema_still_derives_from_annotations(): + manager = ToolManager() + + def add(a: int, b: int): + return a + b + + spec = manager.register(add) + assert spec.parameters_schema["properties"]["a"]["type"] == "integer" + + +def test_unregister_removes_tool_and_reports_presence(): + manager = ToolManager() + + @manager.tool(agent_doc="noop") + def noop(): + return None + + assert manager.unregister("noop") is True + assert manager.unregister("noop") is False + assert "noop" not in manager.tools_map() + assert manager.tool_schemas("agent") == [] + + +def _fake_mcp_client(server_name, tools, calls=None): + """Duck-typed stand-in for McpClient, avoiding a real fastmcp dependency. + + Args: + server_name: Value for the stand-in's ``server_name``. + tools: Value for the stand-in's ``tools`` catalogue. + calls: Optional list collecting one ``(tool_name, arguments)`` pair per + invocation, so a test can observe what a registered tool dispatches. + """ + + async def call_tool(tool_name, arguments): + if calls is not None: + calls.append((tool_name, arguments)) + return f"{tool_name} got {sorted(arguments)}" + + return SimpleNamespace(server_name=server_name, tools=tools, call_tool=call_tool) + + +def _fake_tool(name, description="", input_schema=None): + """Duck-typed stand-in for one MCP server tool.""" + return SimpleNamespace( + name=name, description=description, input_schema=input_schema + ) + + +def test_build_mcp_tool_name_uses_double_underscore_namespace(): + assert ( + ToolManager.build_mcp_tool_name("weather", "get_forecast") + == "mcp__weather__get_forecast" + ) + + +def test_register_mcp_client_registers_namespaced_tools(): + manager = ToolManager() + client = _fake_mcp_client( + "weather", + [ + _fake_tool( + "get_forecast", + description="Get the forecast.", + input_schema={ + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + ) + ], + ) + + specs = manager.register_mcp_client(client) + assert [s.name for s in specs] == ["mcp__weather__get_forecast"] + assert "mcp__weather__get_forecast" in manager.tools_map() + assert manager.agent_docs()["mcp__weather__get_forecast"] == ( + "[weather] Get the forecast." + ) + schema = manager.tool_schemas("agent")[0]["function"]["parameters"] + assert schema["properties"]["city"]["type"] == "string" + + +def test_register_mcp_client_synthesizes_description_when_missing(): + manager = ToolManager() + client = _fake_mcp_client("files", [_fake_tool("read", description=None)]) + + manager.register_mcp_client(client) + doc = manager.agent_docs()["mcp__files__read"] + assert "files" in doc + assert "read" in doc + + +def test_register_mcp_client_normalizes_missing_input_schema(): + manager = ToolManager() + client = _fake_mcp_client("files", [_fake_tool("read", input_schema=None)]) + + manager.register_mcp_client(client) + schema = manager.tool_schemas("agent")[0]["function"]["parameters"] + assert schema["type"] == "object" + assert schema["properties"] == {} + + +def test_register_mcp_client_completes_a_partial_input_schema_without_mutating_it(): + manager = ToolManager() + advertised = {"properties": {"path": {"type": "string"}}} + client = _fake_mcp_client("files", [_fake_tool("read", input_schema=advertised)]) + + manager.register_mcp_client(client) + schema = manager.tool_schemas("agent")[0]["function"]["parameters"] + assert schema["type"] == "object" + assert schema["properties"] == {"path": {"type": "string"}} + # The server's own catalogue object must come back untouched. + assert advertised == {"properties": {"path": {"type": "string"}}} + + +def test_register_mcp_client_skips_name_collision(): + manager = ToolManager() + + @manager.tool(agent_doc="built-in") + def mcp__files__read(): + return "builtin" + + client = _fake_mcp_client("files", [_fake_tool("read", description="remote")]) + specs = manager.register_mcp_client(client) + + assert specs == [] + assert manager.agent_docs()["mcp__files__read"] == "built-in" + + +def test_register_mcp_client_sets_display_action_for_tui(): + manager = ToolManager() + client = _fake_mcp_client("files", [_fake_tool("read", description="d")]) + + spec = manager.register_mcp_client(client)[0] + assert spec.display_action == "mcp:files" + assert spec.human_doc == "[files] d" + + +@pytest.mark.asyncio +async def test_registered_mcp_tools_dispatch_to_their_own_server_tool(): + calls = [] + manager = ToolManager() + client = _fake_mcp_client( + "files", + [_fake_tool("read", description="r"), _fake_tool("write", description="w")], + calls, + ) + manager.register_mcp_client(client) + tools = manager.tools_map() + + assert await tools["mcp__files__read"](path="a.txt") == "read got ['path']" + await tools["mcp__files__write"](path="b.txt", body="hi") + + assert calls == [ + ("read", {"path": "a.txt"}), + ("write", {"path": "b.txt", "body": "hi"}), + ] From 797176cce41b3321ca0c2fe55343011c60bb8634 Mon Sep 17 00:00:00 2001 From: Allen Wang Date: Sat, 17 Oct 2026 13:55:54 +0800 Subject: [PATCH 15/20] test(tools): pin the collision warning, the dependency guard, and _invoke's dispatch namespace The 18 existing tests pinned the MCP registration paths but left four things unverified, plus 15 uncovered branches in the schema derivation. - register_mcp_client's logger.warning on a name collision was never asserted. It is now pinned with the exact message, and a clean registration is pinned as quiet. - The fastmcp dependency guard was weaker than advertised in two ways: PYTHONPATH was joined with a trailing ':' and built from a relative path, so the child's cwd joined sys.path and a stray mitocode/ package would shadow the module under test; and an import executed inside a function body sailed past it. The probe now uses an absolute PYTHONPATH with no empty element and an empty tmp_path cwd, and a second probe registers a client before re-checking sys.modules. - _invoke binds its captured client and tool name as default parameters, so a server tool argument named _client or _tool_name binds over them. Both are pinned with their current (defective) behaviour and reported, not fixed. - register()'s parameters_schema meta tier is documented but unreachable: declare_tool writes only name/human_doc/agent_doc. Pinned by writing the private __mitocode_tool_meta__ attribute directly, with a comment saying so. - Closes 15 uncovered branches: _annotation_to_schema's unannotated/Union/PEP 604/ Literal/sequence/mapping paths and its unsupported-annotation raise, _build_parameters_schema's positional-only raise and *args/**kwargs skip, the omitted required key, _normalize_input_schema's non-dict properties, and register()'s precedence tiers, docstring fallbacks and display-hint defaults. Module line+branch coverage 75% -> 97%; the two remaining statements are unreachable on Python 3.10 (the types.UnionType ImportError fallback and the single-member Union unwrap). --- test/test_tool_manager.py | 951 +++++++++++++++++++++++++++++++++++++- 1 file changed, 936 insertions(+), 15 deletions(-) diff --git a/test/test_tool_manager.py b/test/test_tool_manager.py index 96dd373..56e8414 100644 --- a/test/test_tool_manager.py +++ b/test/test_tool_manager.py @@ -3,6 +3,7 @@ import sys import unittest.mock as mock from types import SimpleNamespace +from typing import Any, Dict, List, Literal, Optional, Set, Tuple import pytest @@ -11,7 +12,10 @@ from mitocode.backend.agent.agent_session import AgentSession from mitocode.backend.agent.session_factory import SessionFactory from mitocode.backend.chat_provider.sample_chat_provider import _SampleChatProvider +from mitocode.backend.exception.mitocode_exception import MitoCodeToolDefinitionError from mitocode.backend.tool.tool_manager import ( + MCP_TOOL_PREFIX, + MCP_TOOL_SEPARATOR, ToolManager, declare_tool, _build_parameters_schema, @@ -41,6 +45,76 @@ def find_spec(self, fullname, path=None, target=None): assert not leaked, f"mcp package imported at runtime: {leaked}" """ +#: The same blocked-fastmcp interpreter, but this one also *calls* the MCP entry +#: point. ``_IMPORT_WITHOUT_FASTMCP`` only observes imports that happen while +#: ``tool_manager`` is being imported, so a lazy ``import`` inside a function +#: body would sail past it; this probe closes that hole by exercising +#: ``register_mcp_client`` and then re-checking ``sys.modules``. +_IMPORT_WITHOUT_FASTMCP_WHEN_CALLED = """ +import importlib.abc +import sys + + +class _BlockFastMcp(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path=None, target=None): + if fullname == "fastmcp" or fullname.startswith("fastmcp."): + raise ImportError("fastmcp is not installed") + return None + + +class _DuckTypedClient: + server_name = "files" + tools = [] + + async def call_tool(self, tool_name, arguments): + return "" + + +sys.meta_path.insert(0, _BlockFastMcp()) +from mitocode.backend.tool.tool_manager import ToolManager + +leaked_at_import = [m for m in sys.modules if m.startswith("mitocode.backend.mcp")] +assert not leaked_at_import, f"mcp package imported at import time: {leaked_at_import}" + +registered = ToolManager().register_mcp_client(_DuckTypedClient()) +assert registered == [], f"an empty catalogue registers nothing, got {registered!r}" + +leaked_after_call = [m for m in sys.modules if m.startswith("mitocode.backend.mcp")] +assert not leaked_after_call, f"register_mcp_client imported the mcp package: {leaked_after_call}" +assert "fastmcp" not in sys.modules, "fastmcp was imported" +""" + + +def _run_import_probe(script, cwd): + """Run *script* in a child interpreter that cannot resolve ``fastmcp``. + + ``PYTHONPATH`` is absolute and carries no empty trailing entry, and the child + runs in an empty *cwd*: ``python -c`` puts the working directory at the front + of ``sys.path`` and an empty ``PYTHONPATH`` element adds it again, so a + ``cwd`` holding a ``mitocode/`` package would shadow the module under test + and quietly turn the probe into a test of the decoy. + + Args: + script: Source for the child interpreter. + cwd: Directory the child runs in; must contain no importable package. + + Returns: + The completed ``subprocess.CompletedProcess``. + """ + src_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "src")) + env = dict(os.environ) + env["PYTHONPATH"] = os.pathsep.join([src_dir, env.get("PYTHONPATH", "")]).rstrip( + os.pathsep + ) + return subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + env=env, + cwd=str(cwd), + check=False, + ) + def test_tool_manager_decorator_registration_and_docs(): manager = ToolManager() @@ -132,27 +206,28 @@ def test_tool(city: str = "Beijing", country: str = "China"): assert props["city"] is not props["country"] -def test_tool_manager_imports_without_fastmcp_installed(): +def test_tool_manager_imports_without_fastmcp_installed(tmp_path): """tool_manager must stay importable where fastmcp is absent. Dependency runs one way — tool_manager -> mcp -> fastmcp — so the MCP reference in ``register_mcp_client``'s annotation must stay under ``TYPE_CHECKING`` and the module must duck-type the client at runtime. + + *tmp_path* is an empty directory so the child interpreter's ``sys.path`` + cannot pick up a stray ``mitocode/`` package and shadow the module under + test. """ - env = dict(os.environ) - env["PYTHONPATH"] = os.pathsep.join( - [ - os.path.join(os.path.dirname(__file__), "..", "src"), - env.get("PYTHONPATH", ""), - ] - ) - result = subprocess.run( - [sys.executable, "-c", _IMPORT_WITHOUT_FASTMCP], - capture_output=True, - text=True, - env=env, - check=False, - ) + result = _run_import_probe(_IMPORT_WITHOUT_FASTMCP, tmp_path) + assert result.returncode == 0, result.stderr + + +def test_tool_manager_does_not_import_the_mcp_package_when_called(tmp_path): + """Calling the MCP entry point must not pull the mcp package in either. + + The import-time probe above cannot see a function-body import, so this one + registers a duck-typed client and re-checks ``sys.modules`` afterwards. + """ + result = _run_import_probe(_IMPORT_WITHOUT_FASTMCP_WHEN_CALLED, tmp_path) assert result.returncode == 0, result.stderr @@ -325,3 +400,849 @@ async def test_registered_mcp_tools_dispatch_to_their_own_server_tool(): ("read", {"path": "a.txt"}), ("write", {"path": "b.txt", "body": "hi"}), ] + + +def test_mcp_tool_namespace_constants_are_the_documented_strings(): + """The prefix and separator are module-public; the composed name depends on them.""" + assert MCP_TOOL_PREFIX == "mcp__", "every remote tool must be recognisably remote" + assert MCP_TOOL_SEPARATOR == "__", "server and tool names are joined by __" + + +def test_build_mcp_tool_name_joins_names_verbatim_without_escaping(): + """A server name holding the separator is not escaped, so composition is ambiguous.""" + composed = ToolManager.build_mcp_tool_name("my__files", "read") + assert composed == "mcp__my__files__read", "names are concatenated as given" + assert ( + ToolManager.build_mcp_tool_name("my", "files__read") == "mcp__my__files__read" + ), "a different server/tool pair composes to the same string" + + +@pytest.mark.parametrize( + "annotation, expected_type", + [ + (str, "string"), + (int, "integer"), + (float, "number"), + (bool, "boolean"), + (dict, "object"), + (list, "array"), + (type(None), "null"), + ], +) +def test_build_parameters_schema_types_each_supported_scalar(annotation, expected_type): + """Every supported base annotation maps to its documented JSON Schema type.""" + + def probe(value: annotation): + return value + + schema = _build_parameters_schema(probe) + assert list(schema["properties"]) == ["value"], "one entry per keyword argument" + assert schema["properties"]["value"] == { + "type": expected_type + }, f"{annotation!r} must render as {expected_type!r}" + + +def test_two_tools_with_the_same_annotation_get_independent_schema_dicts(): + """One tool's parameter schema must never be the same object as another's.""" + + def first(text: str): + return text + + def second(text: str): + return text + + first_property = _build_parameters_schema(first)["properties"]["text"] + second_property = _build_parameters_schema(second)["properties"]["text"] + first_property["type"] = "integer" + + assert second_property == { + "type": "string" + }, "a schema fragment shared with the module's type map would be corrupted here" + + +def test_build_parameters_schema_types_an_unannotated_parameter_as_a_required_string(): + """An absent annotation means string, and the parameter stays required.""" + + def unannotated(value): + return value + + schema = _build_parameters_schema(unannotated) + assert schema["properties"]["value"] == { + "type": "string" + }, "an un-annotated parameter is described as a string" + assert schema["required"] == ["value"], "no default was given, so it is required" + + +def test_build_parameters_schema_renders_an_optional_parameter_as_a_two_branch_any_of(): + """``Optional[str]`` is a two-member union, so it is not a bare string.""" + + def maybe(label: Optional[str]): + return label + + schema = _build_parameters_schema(maybe) + assert schema["properties"]["label"] == { + "anyOf": [{"type": "string"}, {"type": "null"}] + }, "Optional[str] must offer both string and null" + + +def test_build_parameters_schema_renders_a_pep604_union_as_any_of(): + """``str | None`` must be treated exactly like ``Optional[str]``.""" + + def maybe(label: str | None): + return label + + schema = _build_parameters_schema(maybe) + assert schema["properties"]["label"] == { + "anyOf": [{"type": "string"}, {"type": "null"}] + }, "the PEP 604 union type is a union too" + + +def test_build_parameters_schema_renders_a_literal_as_an_enum_of_its_first_members_type(): + """The enum's JSON type is taken from its first member.""" + + def choose(mode: Literal["read", "write"]): + return mode + + schema = _build_parameters_schema(choose) + assert schema["properties"]["mode"] == { + "type": "string", + "enum": ["read", "write"], + }, "a Literal of strings is a string constrained to those values" + + +def test_build_parameters_schema_renders_an_empty_literal_as_a_string(): + """A Literal with no members has no member whose type could be sampled.""" + + def choose(mode: Literal[()]): + return mode + + schema = _build_parameters_schema(choose) + assert schema["properties"]["mode"] == { + "type": "string" + }, "an empty enum falls back to the un-annotated default" + + +@pytest.mark.parametrize( + "annotation, expected", + [ + (List[int], {"type": "array", "items": {"type": "integer"}}), + (Tuple[int, str], {"type": "array", "items": {"type": "integer"}}), + (Set[float], {"type": "array", "items": {"type": "number"}}), + ], +) +def test_build_parameters_schema_types_a_parameterised_sequence_as_an_array( + annotation, expected +): + """A list, tuple or set annotation is an array described by its first item.""" + + def collect(values: annotation): + return values + + schema = _build_parameters_schema(collect) + assert ( + schema["properties"]["values"] == expected + ), f"{annotation!r} must render as {expected!r}" + + +def test_build_parameters_schema_types_a_bare_list_annotation_as_an_array_without_items(): + """An un-parameterised ``list`` has no item type to advertise.""" + + def collect(values: list): + return values + + schema = _build_parameters_schema(collect) + assert schema["properties"]["values"] == { + "type": "array" + }, "no items key is invented for a bare list" + + +@pytest.mark.parametrize( + "annotation, expected", + [ + (list, {"type": "array"}), + (List, {"type": "array"}), + (dict, {"type": "object"}), + (Dict, {"type": "object"}), + ], +) +def test_build_parameters_schema_types_an_unparameterised_container_without_element_type( + annotation, expected +): + """A bare builtin and its bare ``typing`` alias must describe the same thing.""" + + def collect(values: annotation): + return values + + schema = _build_parameters_schema(collect) + assert ( + schema["properties"]["values"] == expected + ), f"{annotation!r} has no element type to advertise, so none is invented" + + +@pytest.mark.parametrize("annotation", [tuple, set]) +def test_register_rejects_a_bare_tuple_or_set_annotation(annotation): + """A bare ``list`` is supported but a bare ``tuple``/``set`` is not. + + The supported list in the docstring names str/int/float/bool/dict/list, so + ``tuple`` and ``set`` are rejected even though their parameterised forms are + accepted. + """ + + def collect(values: annotation): + return values + + manager = ToolManager() + with pytest.raises(MitoCodeToolDefinitionError, match="Unsupported type"): + manager.register(collect, name="collect") + + assert ( + manager.tools_map() == {} + ), "a rejected annotation must leave the registry untouched" + + +def test_build_parameters_schema_types_a_mapping_annotation_by_its_value_type(): + """The key type of a mapping annotation is dropped; the value type is kept.""" + + def index(table: Dict[str, int]): + return table + + schema = _build_parameters_schema(index) + assert schema["properties"]["table"] == { + "type": "object", + "additionalProperties": {"type": "integer"}, + }, "Dict[str, int] describes an object whose values are integers" + + +@pytest.mark.parametrize("annotation", [complex, Any]) +def test_register_rejects_an_unsupported_annotation_without_registering(annotation): + """An annotation outside the supported set fails loudly, leaving no half-built tool.""" + + def broken(value: annotation): + return value + + manager = ToolManager() + with pytest.raises( + MitoCodeToolDefinitionError, match="Unsupported type annotation" + ): + manager.register(broken, name="broken") + + assert manager.tools_map() == {}, "a failed derivation must not register a tool" + + +def test_register_rejects_a_positional_only_parameter_without_registering(): + """LLMs invoke tools by keyword, so a positional-only parameter is unrepresentable.""" + + def broken(value: str, /): + return value + + manager = ToolManager() + with pytest.raises( + MitoCodeToolDefinitionError, match="positional-only parameter 'value'" + ): + manager.register(broken, name="broken") + + assert ( + manager.tools_map() == {} + ), "a rejected signature must leave the registry untouched" + + +def test_register_names_a_callable_object_by_its_class_in_the_rejection_message(): + """A tool without ``__name__`` is still named in the error, by its class.""" + + class BrokenTool: + def __call__(self, value: str, /): + return value + + manager = ToolManager() + with pytest.raises(MitoCodeToolDefinitionError, match="Tool function 'BrokenTool'"): + manager.register(BrokenTool(), name="broken") + + assert ( + manager.tools_map() == {} + ), "a rejected signature must leave the registry untouched" + + +def test_build_parameters_schema_omits_star_args_and_star_kwargs(): + """``*args``/``**kwargs`` are not keyword arguments a model could name.""" + + def collect(required: int, *args, **kwargs): + return required, args, kwargs + + schema = _build_parameters_schema(collect) + assert list(schema["properties"]) == [ + "required" + ], "only named parameters are described" + assert schema["required"] == [ + "required" + ], "the named parameter keeps its requiredness" + + +def test_build_parameters_schema_omits_the_required_key_when_nothing_is_required(): + """An all-defaults signature advertises defaults and no required list.""" + + def defaulted(count: int = 0): + return count + + schema = _build_parameters_schema(defaulted) + assert "required" not in schema, "no required parameter means no required list" + assert schema["properties"]["count"] == { + "type": "integer", + "default": 0, + }, "the default is advertised so the model may omit the argument" + + +def test_build_parameters_schema_always_forbids_undeclared_arguments(): + """The top-level schema is always a closed object schema.""" + + def simple(text: str): + return text + + schema = _build_parameters_schema(simple) + assert schema["type"] == "object", "a parameter list describes an object" + assert ( + schema["additionalProperties"] is False + ), "an undeclared argument cannot be passed by a model, so it is forbidden" + + +def test_declare_tool_without_a_name_falls_back_to_the_function_name(): + """An unnamed declaration keeps the decorated function's own name.""" + + @declare_tool(human_doc="doc", agent_doc="agent") + def fallback_name(): + return None + + manager = ToolManager() + manager.register(fallback_name) + + assert list(manager.tools_map()) == [ + "fallback_name" + ], "no name means the function's __name__" + + +def test_register_prefers_an_explicit_name_over_declare_tool_metadata(): + """The explicit ``name`` argument outranks the declared metadata name.""" + + @declare_tool(name="meta_name") + def conflict(text: str): + return text + + manager = ToolManager() + manager.register(conflict, name="explicit_name") + + assert list(manager.tools_map()) == [ + "explicit_name" + ], "the explicit name is the one the model sees" + + +def test_register_prefers_explicit_docs_over_declare_tool_metadata(): + """The explicit doc arguments outrank the declared metadata docs.""" + + @declare_tool(human_doc="meta human", agent_doc="meta agent") + def conflict(text: str): + return text + + manager = ToolManager() + manager.register(conflict, human_doc="explicit human", agent_doc="explicit agent") + + assert manager.human_docs()["conflict"] == "explicit human", "explicit wins" + assert manager.agent_docs()["conflict"] == "explicit agent", "explicit wins" + + +def test_register_falls_back_to_the_function_docstring_for_both_audiences(): + """With no metadata and no arguments, the function's docstring is the description.""" + + def documented(): + """Search the workspace.""" + + manager = ToolManager() + manager.register(documented) + + assert manager.human_docs()["documented"] == "Search the workspace.", "docstring" + assert manager.agent_docs()["documented"] == ( + "Search the workspace." + ), "the agent description falls back to the human one" + + +def test_register_gives_an_undocumented_tool_empty_docs_for_both_audiences(): + """No docstring, no metadata and no arguments means no description at all.""" + + def bare(): + return None + + manager = ToolManager() + manager.register(bare) + + assert ( + manager.human_docs()["bare"] == "" + ), "an absent docstring is not the string 'None'" + assert manager.agent_docs()["bare"] == "", "and the agent gets nothing either" + + +def test_register_falls_back_to_the_human_doc_when_only_that_is_given(): + """An agent-only description may be omitted; the human one then serves both.""" + + @declare_tool(human_doc="human only") + def only_human(): + return None + + manager = ToolManager() + manager.register(only_human) + + assert manager.human_docs()["only_human"] == "human only", "the declared human doc" + assert manager.agent_docs()["only_human"] == ( + "human only" + ), "the agent description defaults to the human description" + + +def test_register_reads_parameters_schema_from_declare_tool_metadata(): + """The documented metadata tier for ``parameters_schema`` does work. + + ``declare_tool`` writes only ``name``, ``human_doc`` and ``agent_doc``, so + there is no public way to place ``parameters_schema`` into the metadata this + tier reads; the private attribute has to be written by hand. That makes + the tier advertised but unreachable through the public API. + """ + manager = ToolManager() + from_meta = {"type": "object", "properties": {"remote": {"type": "string"}}} + + @declare_tool(name="remote_call") + def remote_call(**kwargs): + return kwargs + + remote_call.__mitocode_tool_meta__["parameters_schema"] = from_meta + spec = manager.register(remote_call) + + assert ( + spec.parameters_schema == from_meta + ), "the metadata schema must outrank the derived empty one" + assert ( + "remote" in spec.parameters_schema["properties"] + ), "the server's shape survives" + + +def test_register_prefers_the_explicit_parameters_schema_over_declare_tool_metadata(): + """Tier one beats tier two: the explicit argument wins over the metadata.""" + + @declare_tool(name="tiered") + def tiered(**kwargs): + return kwargs + + tiered.__mitocode_tool_meta__["parameters_schema"] = { + "type": "object", + "properties": {"from_meta": {"type": "string"}}, + } + explicit = {"type": "object", "properties": {"explicit": {"type": "boolean"}}} + + manager = ToolManager() + spec = manager.register(tiered, parameters_schema=explicit) + + assert spec.parameters_schema == explicit, "the explicit argument is the winner" + assert ( + "from_meta" not in spec.parameters_schema["properties"] + ), "the metadata tier must not be consulted at all" + + +def test_register_keeps_an_explicitly_empty_parameters_schema_instead_of_deriving_one(): + """An empty dict is a supplied schema, not an absent one. + + ``parameters_schema={}`` is how a caller says "this tool takes no + arguments"; treating it as absent would advertise a derived schema instead. + """ + + def ping() -> str: + return "pong" + + manager = ToolManager() + spec = manager.register(ping, parameters_schema={}) + + assert spec.parameters_schema == {}, "an explicit {} must be kept verbatim" + + +def test_register_replaces_a_registration_that_already_uses_the_name(): + """Registering the same name again overwrites rather than duplicating.""" + + def first(): + return "first" + + def second(): + return "second" + + manager = ToolManager() + manager.register(first, name="shared", agent_doc="first doc") + spec = manager.register(second, name="shared", agent_doc="second doc") + + assert list(manager.tools_map()) == ["shared"], "one name, one entry" + assert manager.tools_map()["shared"] is second, "the newest registration wins" + assert spec.agent_doc == "second doc", "the description is replaced with it" + + +def test_register_defaults_both_display_hints_to_empty_strings(): + """Display hints are opt-in; without them the spec carries empty strings.""" + + def plain(text: str): + return text + + manager = ToolManager() + spec = manager.register(plain) + + assert spec.display_action == "", "no display action was requested" + assert spec.display_param == "", "no display parameter was requested" + + +def test_tool_decorator_returns_the_decorated_function_unchanged(): + """The decorator registers the function and hands the same object back.""" + + def ping(): + return None + + manager = ToolManager() + decorated = manager.tool(agent_doc="doc")(ping) + + assert decorated is ping, "registration must not wrap or replace the callable" + + +def test_tool_decorator_forwards_an_explicit_parameters_schema_to_register(): + """The decorator must pass ``parameters_schema`` through, not swallow it.""" + manager = ToolManager() + schema = {"type": "object", "properties": {"city": {"type": "string"}}} + + @manager.tool(name="get_weather", parameters_schema=schema) + def get_weather(**kwargs): + return kwargs + + assert ( + manager.get_spec("get_weather").parameters_schema == schema + ), "the decorator forwards the schema to register()" + + +def test_get_spec_returns_none_for_a_name_that_was_never_registered(): + """An unknown name has no spec rather than raising.""" + manager = ToolManager() + assert manager.get_spec("absent") is None, "a missing tool reports None" + + +def test_get_spec_returns_the_same_object_register_returned(): + """``get_spec`` exposes the registered object, not a copy of it.""" + + def ping(): + return None + + manager = ToolManager() + spec = manager.register(ping, name="ping") + + assert manager.get_spec("ping") is spec, "callers must see the stored spec" + + +def test_tool_schemas_defaults_to_the_agent_audience(): + """The default audience is the agent's.""" + manager = ToolManager() + + @manager.tool(human_doc="human only", agent_doc="agent only") + def ping(): + return None + + assert manager.tool_schemas() == manager.tool_schemas( + "agent" + ), "the default audience is the agent's" + assert ( + manager.tool_schemas()[0]["function"]["description"] == "agent only" + ), "the agent's description is served by default" + + +def test_tool_schemas_serves_human_docs_for_the_human_audience(): + """``audience="human"`` switches the description to the human one.""" + manager = ToolManager() + + @manager.tool(human_doc="human only", agent_doc="agent only") + def ping(): + return None + + assert manager.tool_schemas("human")[0]["function"]["description"] == ( + "human only" + ), "the human audience gets the human description" + + +def test_tool_schemas_treats_an_unknown_audience_as_the_agent_one(): + """Only the literal ``"human"`` switches audience.""" + manager = ToolManager() + + @manager.tool(human_doc="human only", agent_doc="agent only") + def ping(): + return None + + assert manager.tool_schemas("operator")[0]["function"]["description"] == ( + "agent only" + ), "an unrecognised audience is served the agent description" + + +def test_register_mcp_client_registers_nothing_from_an_empty_catalogue(): + """A client that discovered no tools registers nothing and warns about nothing.""" + manager = ToolManager() + client = _fake_mcp_client("files", []) + + with mock.patch("mitocode.backend.tool.tool_manager.logger") as fake_logger: + specs = manager.register_mcp_client(client) + + assert specs == [], "an empty catalogue yields no specs" + assert manager.tools_map() == {}, "and leaves the registry untouched" + assert ( + fake_logger.warning.called is False + ), "nothing collided, so nothing to warn about" + + +def test_register_mcp_client_returns_specs_in_catalogue_order(): + """The returned specs follow discovery order, not alphabetical order.""" + manager = ToolManager() + client = _fake_mcp_client( + "files", + [ + _fake_tool("write", description="w"), + _fake_tool("delete", description="d"), + _fake_tool("read", description="r"), + ], + ) + + specs = manager.register_mcp_client(client) + + assert [spec.name for spec in specs] == [ + "mcp__files__write", + "mcp__files__delete", + "mcp__files__read", + ], "discovery order is the contract" + assert list(manager.tools_map()) == [ + spec.name for spec in specs + ], "the registry keeps the same order" + + +def test_register_mcp_client_namespaces_the_same_tool_name_per_server(): + """Two servers offering one tool name both register under their own namespace.""" + manager = ToolManager() + alpha = _fake_mcp_client("alpha", [_fake_tool("search", description="from alpha")]) + beta = _fake_mcp_client("beta", [_fake_tool("search", description="from beta")]) + + specs = manager.register_mcp_client(alpha) + manager.register_mcp_client(beta) + + assert [spec.name for spec in specs] == [ + "mcp__alpha__search", + "mcp__beta__search", + ], "the server name keeps two identical tool names apart" + assert manager.agent_docs()["mcp__alpha__search"] == ( + "[alpha] from alpha" + ), "each keeps its own description" + + +def test_register_mcp_client_warns_and_keeps_the_existing_tool_on_collision(): + """A skipped collision is logged, not silent, and the incumbent survives.""" + manager = ToolManager() + + @manager.tool(agent_doc="built-in") + def mcp__files__read(): + return "builtin" + + client = _fake_mcp_client("files", [_fake_tool("read", description="remote")]) + with mock.patch("mitocode.backend.tool.tool_manager.logger") as fake_logger: + specs = manager.register_mcp_client(client) + + assert specs == [], "a colliding tool is not registered" + assert manager.agent_docs()["mcp__files__read"] == "built-in", "the incumbent wins" + fake_logger.warning.assert_called_once_with( + "Skipping MCP tool 'mcp__files__read': name already registered" + ), "the skip must be visible in the log" + + +def test_register_mcp_client_does_not_warn_when_every_name_is_free(): + """Only collisions warn; a clean registration is quiet.""" + manager = ToolManager() + client = _fake_mcp_client("files", [_fake_tool("read", description="remote")]) + + with mock.patch("mitocode.backend.tool.tool_manager.logger") as fake_logger: + specs = manager.register_mcp_client(client) + + assert len(specs) == 1, "the tool was registered" + assert fake_logger.warning.called is False, "a free name is not a warning" + + +@pytest.mark.asyncio +async def test_register_mcp_client_keeps_the_first_of_two_clients_with_one_server_name(): + """A second client on the same server name never receives the dispatch.""" + first_calls = [] + second_calls = [] + first = _fake_mcp_client( + "files", [_fake_tool("read", description="from first")], first_calls + ) + second = _fake_mcp_client( + "files", [_fake_tool("read", description="from second")], second_calls + ) + manager = ToolManager() + manager.register_mcp_client(first) + manager.register_mcp_client(second) + + await manager.tools_map()["mcp__files__read"](path="a.txt") + + assert first_calls == [ + ("read", {"path": "a.txt"}) + ], "the first client keeps ownership of the composed name" + assert second_calls == [], "the second client must never be dispatched to" + + +def test_register_mcp_client_synthesizes_an_exact_description_for_a_blank_description(): + """A whitespace-only description is replaced by the documented sentence.""" + + manager = ToolManager() + client = _fake_mcp_client("files", [_fake_tool("read", description=" \n\t ")]) + + manager.register_mcp_client(client) + + assert manager.human_docs()["mcp__files__read"] == ( + "[files] Tool 'read' provided by MCP server 'files'." + ), "the synthesised description names the tool and the server" + assert manager.agent_docs()["mcp__files__read"] == ( + "[files] Tool 'read' provided by MCP server 'files'." + ), "both audiences get the same text" + + +def test_register_mcp_client_strips_surrounding_whitespace_from_a_description(): + """A real description keeps its content but loses its padding.""" + + manager = ToolManager() + client = _fake_mcp_client( + "files", [_fake_tool("read", description=" Read a file.\n")] + ) + + manager.register_mcp_client(client) + + assert manager.agent_docs()["mcp__files__read"] == ( + "[files] Read a file." + ), "the server's own description is used, trimmed" + + +def test_register_mcp_client_leaves_the_display_param_empty(): + """Only a display action is synthesised; there is no parameter to highlight.""" + + manager = ToolManager() + client = _fake_mcp_client("files", [_fake_tool("read", description="d")]) + + spec = manager.register_mcp_client(client)[0] + + assert ( + spec.display_param == "" + ), "the module does not know which argument to highlight" + + +def test_register_mcp_client_preserves_a_complete_input_schema_verbatim(): + """A server that advertises a full schema gets it back unchanged.""" + manager = ToolManager() + advertised = { + "type": "object", + "properties": {"path": {"type": "string"}, "limit": {"type": "integer"}}, + "required": ["path"], + "additionalProperties": False, + } + client = _fake_mcp_client("files", [_fake_tool("read", input_schema=advertised)]) + + manager.register_mcp_client(client) + + schema = manager.tool_schemas("agent")[0]["function"]["parameters"] + assert schema == advertised, "nothing is added to or removed from a complete schema" + + +def test_register_mcp_client_replaces_non_dict_properties_in_an_input_schema(): + """A ``properties`` value that is not a dict is replaced by an empty dict.""" + manager = ToolManager() + client = _fake_mcp_client( + "files", + [_fake_tool("read", input_schema={"type": "object", "properties": "city"})], + ) + + manager.register_mcp_client(client) + + schema = manager.tool_schemas("agent")[0]["function"]["parameters"] + assert schema["properties"] == {}, "the unusable value is dropped" + assert schema["type"] == "object", "a usable type survives" + + +@pytest.mark.asyncio +async def test_registered_mcp_tool_returns_the_server_result_unchanged(): + """The generated callable hands the server's return value straight back.""" + + manager = ToolManager() + client = _fake_mcp_client("files", [_fake_tool("read", description="r")]) + manager.register_mcp_client(client) + + result = await manager.tools_map()["mcp__files__read"](path="a.txt", limit=3) + + assert ( + result == "read got ['limit', 'path']" + ), "the server's string is passed through" + + +def test_registered_mcp_tool_schema_hides_the_private_dispatch_bindings(): + """``_client``/``_tool_name`` exist on the callable but not on the wire.""" + manager = ToolManager() + client = _fake_mcp_client( + "files", + [ + _fake_tool( + "read", + input_schema={ + "type": "object", + "properties": {"path": {"type": "string"}}, + }, + ) + ], + ) + + spec = manager.register_mcp_client(client)[0] + + assert set(spec.parameters_schema["properties"]) == { + "path" + }, "the advertised arguments are exactly the server's" + assert ( + "_client" not in spec.parameters_schema["properties"] + ), "an internal binding must never be advertised to a model" + assert ( + "_tool_name" not in spec.parameters_schema["properties"] + ), "an internal binding must never be advertised to a model" + + +@pytest.mark.asyncio +async def test_registered_mcp_tool_lets_a_model_argument_override_the_captured_tool_name(): + """Documents a defect: a ``_tool_name`` argument silently redirects the dispatch. + + ``_invoke`` binds its captured client and tool name as ordinary default + parameters, so they share one namespace with the model's ``**kwargs``. A + server tool advertising an argument called ``_tool_name`` therefore binds + over the captured name and the call lands on a different server tool. The + docstring promises only the model's keyword arguments are forwarded. + """ + calls = [] + manager = ToolManager() + client = _fake_mcp_client("files", [_fake_tool("read", description="r")], calls) + manager.register_mcp_client(client) + + result = await manager.tools_map()["mcp__files__read"]( + _tool_name="write", path="a.txt" + ) + + assert calls == [ + ("write", {"path": "a.txt"}) + ], "the captured tool name was overridden by the model's argument" + assert result == "write got ['path']", "the dispatch went to 'write', not 'read'" + + +@pytest.mark.asyncio +async def test_registered_mcp_tool_lets_a_model_argument_override_the_captured_client(): + """Documents a defect: a ``_client`` argument replaces the captured client. + + Same namespace collision as ``_tool_name``: the value binds over the + captured client, so the call fails instead of reaching the server at all. + """ + calls = [] + manager = ToolManager() + client = _fake_mcp_client("files", [_fake_tool("read", description="r")], calls) + manager.register_mcp_client(client) + + with pytest.raises(AttributeError, match="'call_tool'"): + await manager.tools_map()["mcp__files__read"](_client="not-a-client") + + assert calls == [], "the real client was never reached" From a9d94c5d5c0590a882e7f935b430fa5797ea63a3 Mon Sep 17 00:00:00 2001 From: Allen Wang Date: Sat, 17 Oct 2026 14:12:59 +0800 Subject: [PATCH 16/20] fix(tools): close the MCP dispatch argument-injection hole The generated callable bound its client and upstream tool name as default arguments, putting both in the same namespace as the model's keyword arguments. A server tool advertising an argument literally named _tool_name redirected the call to a different tool on that server -- including one that was never registered, bypassing the name-collision guard register_mcp_client relies on. Replaced with a closure factory that exposes neither value as a keyword. The two tests that pinned the defective behaviour are now xfail(strict=False) asserting the secure contract, so they XPASS on the fix instead of blocking it, and are renamed to read as a known defect rather than a capability. This is the same pattern the branch already ruled against once, on a redaction test that asserted a secret survived. Also: advertise a tool in the dependency-direction probe so the loop body actually executes -- with an empty catalogue a lazy in-loop import passed both probes silently; merge the two subprocess probes, which were 80% of the file's runtime; and block the standalone mcp SDK alongside fastmcp. Constraint: McpClient stays referenced only under TYPE_CHECKING; tool_manager must import without fastmcp installed Confidence: high Scope-risk: narrow --- src/mitocode/backend/tool/tool_manager.py | 47 ++++-- test/test_tool_manager.py | 191 +++++++++++----------- 2 files changed, 134 insertions(+), 104 deletions(-) diff --git a/src/mitocode/backend/tool/tool_manager.py b/src/mitocode/backend/tool/tool_manager.py index 58279e2..7c83098 100644 --- a/src/mitocode/backend/tool/tool_manager.py +++ b/src/mitocode/backend/tool/tool_manager.py @@ -72,6 +72,42 @@ def _normalize_input_schema(raw: Any) -> Dict[str, Any]: return normalized +def _make_invoke(client: "McpClient", upstream_name: str) -> Callable: + """Build the async callable that dispatches one MCP tool. + + The client and the upstream tool name are captured in the closure, so the + generated callable takes ``**kwargs`` and nothing else. Two earlier shapes + were wrong and must not come back: + + * Binding them as default arguments — ``async def _invoke(_client=client, + _tool_name=tool.name, **kwargs)`` — puts them in the *same* namespace as + the model's keyword arguments. A server tool advertising an argument + literally named ``_client`` or ``_tool_name`` then binds over the captured + values, and ``_tool_name`` in particular silently redirects the call to a + different tool on that server, bypassing the name-collision guard that + ``register_mcp_client`` relies on. + * Closing over the loop variables directly would route every tool registered + from one catalogue to whichever tool happened to come last. + + A closure factory gives the isolation the default arguments were reaching + for without exposing either value as a keyword. + + Args: + client: The connected MCP client to dispatch through. + upstream_name: The tool's name *on the server*, which is not the + namespaced name MitoCode registers it under. + + Returns: + An async callable forwarding its keyword arguments to that one tool. + """ + + async def _invoke(**kwargs): + """Forward the model's keyword arguments to the server tool.""" + return await client.call_tool(upstream_name, kwargs) + + return _invoke + + def _annotation_to_schema(annotation: Any) -> Dict[str, Any]: """Convert a Python type annotation to a JSON Schema fragment. @@ -405,15 +441,6 @@ def register_mcp_client(self, client: "McpClient") -> List[ToolSpec]: ) continue - async def _invoke(_client=client, _tool_name=tool.name, **kwargs): - """Forward the model's keyword arguments to the server tool. - - The default arguments capture *this* iteration's client and - tool name. Closing over the loop variables instead would - route every registered tool to whichever tool came last. - """ - return await _client.call_tool(_tool_name, kwargs) - description = (tool.description or "").strip() if not description: description = ( @@ -423,7 +450,7 @@ async def _invoke(_client=client, _tool_name=tool.name, **kwargs): registered.append( self.register( - _invoke, + _make_invoke(client, tool.name), name=tool_name, human_doc=f"[{client.server_name}] {description}", agent_doc=f"[{client.server_name}] {description}", diff --git a/test/test_tool_manager.py b/test/test_tool_manager.py index 56e8414..1187a4d 100644 --- a/test/test_tool_manager.py +++ b/test/test_tool_manager.py @@ -1,3 +1,4 @@ +import inspect import os import subprocess import sys @@ -21,72 +22,65 @@ _build_parameters_schema, ) -#: Imports tool_manager in a process where fastmcp cannot be resolved at all, -#: so the one-way tool_manager -> mcp -> fastmcp dependency can be checked for -#: real rather than by reading the source. Run out-of-process because the test -#: session has already imported fastmcp. -_IMPORT_WITHOUT_FASTMCP = """ +#: Imports tool_manager in a process where the MCP SDKs cannot be resolved at +#: all, then drives the MCP entry point, so the one-way +#: tool_manager -> mcp -> fastmcp dependency can be checked for real rather than by +#: reading the source. Run out-of-process because the test session has already +#: imported fastmcp. +#: +#: Both halves matter. The import-time assertion only sees imports executed while +#: ``tool_manager`` is being imported, so a lazy ``import`` inside a function body +#: sails past it; the registration below is what closes that hole, and the client +#: must advertise a tool for the registration loop to actually run. +_IMPORT_WITHOUT_MCP_SDK = """ import importlib.abc import sys +from types import SimpleNamespace -class _BlockFastMcp(importlib.abc.MetaPathFinder): - def find_spec(self, fullname, path=None, target=None): - if fullname == "fastmcp" or fullname.startswith("fastmcp."): - raise ImportError("fastmcp is not installed") - return None - - -sys.meta_path.insert(0, _BlockFastMcp()) -import mitocode.backend.tool.tool_manager - -assert "fastmcp" not in sys.modules, "fastmcp was imported" -leaked = [m for m in sys.modules if m.startswith("mitocode.backend.mcp")] -assert not leaked, f"mcp package imported at runtime: {leaked}" -""" - -#: The same blocked-fastmcp interpreter, but this one also *calls* the MCP entry -#: point. ``_IMPORT_WITHOUT_FASTMCP`` only observes imports that happen while -#: ``tool_manager`` is being imported, so a lazy ``import`` inside a function -#: body would sail past it; this probe closes that hole by exercising -#: ``register_mcp_client`` and then re-checking ``sys.modules``. -_IMPORT_WITHOUT_FASTMCP_WHEN_CALLED = """ -import importlib.abc -import sys - +class _BlockMcpSdk(importlib.abc.MetaPathFinder): + \"\"\"Make both MCP SDK distributions unresolvable, however they are spelled.\"\"\" -class _BlockFastMcp(importlib.abc.MetaPathFinder): def find_spec(self, fullname, path=None, target=None): - if fullname == "fastmcp" or fullname.startswith("fastmcp."): - raise ImportError("fastmcp is not installed") + if fullname.split(".")[0] in ("fastmcp", "mcp"): + raise ImportError(f"{fullname} is not installed") return None class _DuckTypedClient: server_name = "files" - tools = [] + tools = [SimpleNamespace(name="read", description="d", input_schema=None)] async def call_tool(self, tool_name, arguments): return "" -sys.meta_path.insert(0, _BlockFastMcp()) -from mitocode.backend.tool.tool_manager import ToolManager +def leaked(): + return sorted( + m + for m in sys.modules + if m.split(".")[0] in ("fastmcp", "mcp") or m.startswith("mitocode.backend.mcp") + ) + -leaked_at_import = [m for m in sys.modules if m.startswith("mitocode.backend.mcp")] -assert not leaked_at_import, f"mcp package imported at import time: {leaked_at_import}" +sys.meta_path.insert(0, _BlockMcpSdk()) +import mitocode.backend.tool.tool_manager as tool_manager -registered = ToolManager().register_mcp_client(_DuckTypedClient()) -assert registered == [], f"an empty catalogue registers nothing, got {registered!r}" +assert not leaked(), f"mcp package imported at import time: {leaked()}" -leaked_after_call = [m for m in sys.modules if m.startswith("mitocode.backend.mcp")] -assert not leaked_after_call, f"register_mcp_client imported the mcp package: {leaked_after_call}" -assert "fastmcp" not in sys.modules, "fastmcp was imported" +# The advertised tool is what makes register_mcp_client's loop body run, which is +# exactly where a helpful-looking lazy import would hide. +registered = tool_manager.ToolManager().register_mcp_client(_DuckTypedClient()) +assert [spec.name for spec in registered] == ["mcp__files__read"], ( + f"the loop body did not run, so this probe would prove nothing: {registered!r}" +) + +assert not leaked(), f"register_mcp_client imported the mcp package: {leaked()}" """ def _run_import_probe(script, cwd): - """Run *script* in a child interpreter that cannot resolve ``fastmcp``. + """Run *script* in a child interpreter that cannot resolve the MCP SDKs. ``PYTHONPATH`` is absolute and carries no empty trailing entry, and the child runs in an empty *cwd*: ``python -c`` puts the working directory at the front @@ -206,28 +200,23 @@ def test_tool(city: str = "Beijing", country: str = "China"): assert props["city"] is not props["country"] -def test_tool_manager_imports_without_fastmcp_installed(tmp_path): - """tool_manager must stay importable where fastmcp is absent. +def test_tool_manager_imports_and_registers_without_the_mcp_sdk(tmp_path): + """tool_manager must stay importable and usable without either MCP SDK. Dependency runs one way — tool_manager -> mcp -> fastmcp — so the MCP reference in ``register_mcp_client``'s annotation must stay under ``TYPE_CHECKING`` and the module must duck-type the client at runtime. + The probe checks ``sys.modules`` twice: once after importing ``tool_manager`` + and once after registering a tool from a client. A lazy ``import`` in a + function body is invisible to the first check, which is why the second one + drives the registration loop rather than merely calling the entry point. + *tmp_path* is an empty directory so the child interpreter's ``sys.path`` cannot pick up a stray ``mitocode/`` package and shadow the module under test. """ - result = _run_import_probe(_IMPORT_WITHOUT_FASTMCP, tmp_path) - assert result.returncode == 0, result.stderr - - -def test_tool_manager_does_not_import_the_mcp_package_when_called(tmp_path): - """Calling the MCP entry point must not pull the mcp package in either. - - The import-time probe above cannot see a function-body import, so this one - registers a duck-typed client and re-checks ``sys.modules`` afterwards. - """ - result = _run_import_probe(_IMPORT_WITHOUT_FASTMCP_WHEN_CALLED, tmp_path) + result = _run_import_probe(_IMPORT_WITHOUT_MCP_SDK, tmp_path) assert result.returncode == 0, result.stderr @@ -544,18 +533,6 @@ def collect(values: annotation): ), f"{annotation!r} must render as {expected!r}" -def test_build_parameters_schema_types_a_bare_list_annotation_as_an_array_without_items(): - """An un-parameterised ``list`` has no item type to advertise.""" - - def collect(values: list): - return values - - schema = _build_parameters_schema(collect) - assert schema["properties"]["values"] == { - "type": "array" - }, "no items key is invented for a bare list" - - @pytest.mark.parametrize( "annotation, expected", [ @@ -568,7 +545,11 @@ def collect(values: list): def test_build_parameters_schema_types_an_unparameterised_container_without_element_type( annotation, expected ): - """A bare builtin and its bare ``typing`` alias must describe the same thing.""" + """A bare builtin and its bare ``typing`` alias must describe the same thing. + + Both spellings have no element or value type to advertise, so neither + invents an ``items`` or ``additionalProperties`` key. + """ def collect(values: annotation): return values @@ -985,9 +966,7 @@ def test_register_mcp_client_registers_nothing_from_an_empty_catalogue(): assert specs == [], "an empty catalogue yields no specs" assert manager.tools_map() == {}, "and leaves the registry untouched" - assert ( - fake_logger.warning.called is False - ), "nothing collided, so nothing to warn about" + assert not fake_logger.warning.called, "nothing collided, so nothing to warn about" def test_register_mcp_client_returns_specs_in_catalogue_order(): @@ -1059,7 +1038,7 @@ def test_register_mcp_client_does_not_warn_when_every_name_is_free(): specs = manager.register_mcp_client(client) assert len(specs) == 1, "the tool was registered" - assert fake_logger.warning.called is False, "a free name is not a warning" + assert not fake_logger.warning.called, "a free name is not a warning" @pytest.mark.asyncio @@ -1206,43 +1185,67 @@ def test_registered_mcp_tool_schema_hides_the_private_dispatch_bindings(): @pytest.mark.asyncio -async def test_registered_mcp_tool_lets_a_model_argument_override_the_captured_tool_name(): - """Documents a defect: a ``_tool_name`` argument silently redirects the dispatch. - - ``_invoke`` binds its captured client and tool name as ordinary default - parameters, so they share one namespace with the model's ``**kwargs``. A - server tool advertising an argument called ``_tool_name`` therefore binds - over the captured name and the call lands on a different server tool. The - docstring promises only the model's keyword arguments are forwarded. +async def test_registered_mcp_tool_exposes_no_dispatch_bindings_as_keywords(): + """The generated callable must accept nothing but the model's own arguments.""" + manager = ToolManager() + client = _fake_mcp_client("files", [_fake_tool("read", description="r")]) + manager.register_mcp_client(client) + + parameters = inspect.signature(manager.tools_map()["mcp__files__read"]).parameters + + assert list(parameters) == [ + "kwargs" + ], "the client and upstream tool name must be captured, never bound as keywords" + + +@pytest.mark.xfail( + reason="defect: _invoke's defaults shared a namespace with **kwargs", + strict=False, +) +@pytest.mark.asyncio +async def test_registered_mcp_tool_cannot_redirect_the_dispatch_with_a_tool_name_argument(): + """A ``_tool_name`` argument must be forwarded, never bind over the captured name. + + FIXME(tool_manager): the defect was ``async def _invoke(_client=client, + _tool_name=tool.name, **kwargs)`` — the captured values sat in the same + namespace as the model's keywords, so this argument silently redirected the + call to a *different* tool on that server, bypassing the name-collision guard + entirely. Fixed by ``_make_invoke``; the ``xfail`` stays non-strict so the + suite is green both ways and can be promoted to a plain assertion once the + closure factory has shipped everywhere. """ calls = [] manager = ToolManager() client = _fake_mcp_client("files", [_fake_tool("read", description="r")], calls) manager.register_mcp_client(client) - result = await manager.tools_map()["mcp__files__read"]( - _tool_name="write", path="a.txt" - ) + await manager.tools_map()["mcp__files__read"](_tool_name="write", path="a.txt") assert calls == [ - ("write", {"path": "a.txt"}) - ], "the captured tool name was overridden by the model's argument" - assert result == "write got ['path']", "the dispatch went to 'write', not 'read'" + ("read", {"_tool_name": "write", "path": "a.txt"}) + ], "the dispatch must stay on the registered tool, with _tool_name as a plain argument" +@pytest.mark.xfail( + reason="defect: _invoke's defaults shared a namespace with **kwargs", + strict=False, +) @pytest.mark.asyncio -async def test_registered_mcp_tool_lets_a_model_argument_override_the_captured_client(): - """Documents a defect: a ``_client`` argument replaces the captured client. +async def test_registered_mcp_tool_cannot_redirect_the_dispatch_with_a_client_argument(): + """A ``_client`` argument must be forwarded, never replace the captured client. - Same namespace collision as ``_tool_name``: the value binds over the - captured client, so the call fails instead of reaching the server at all. + FIXME(tool_manager): same namespace collision as ``_tool_name``. The + defect raised ``AttributeError`` from the model's own argument before the + server was ever reached; this asserts the dispatch contract instead of that + incidental error message, so it also accepts a clearer error from a fix. """ calls = [] manager = ToolManager() client = _fake_mcp_client("files", [_fake_tool("read", description="r")], calls) manager.register_mcp_client(client) - with pytest.raises(AttributeError, match="'call_tool'"): - await manager.tools_map()["mcp__files__read"](_client="not-a-client") + await manager.tools_map()["mcp__files__read"](_client="not-a-client") - assert calls == [], "the real client was never reached" + assert calls == [ + ("read", {"_client": "not-a-client"}) + ], "the real client must still receive the call, with _client as a plain argument" From 61851dd50f1dc9c05dd8a006e92cefd7dcada702 Mon Sep 17 00:00:00 2001 From: Allen Wang Date: Sat, 17 Oct 2026 14:26:32 +0800 Subject: [PATCH 17/20] feat(mcp): add McpManager with per-server failure isolation start() connects every configured server concurrently and registers only the ones that came up, so an unreachable server costs you its own tools and nothing else. client_factory is the seam that lets the tests drive both fake and real stdio servers through the same lifecycle. Each server's tools are registered only after its connect() reports success, and released before every re-registration: register_mcp_client never overwrites a name, so a repeated start or a refresh that failed would otherwise leave tools the manager no longer owns and stop() cannot remove. A failure the client did not report itself is recorded against that server and surfaced through status(), so "connected but no tools" always has a reason next to it. Constraint: start() must run before any AgentSession is created, since AgentSession snapshots the tool registry Rejected: ClientGroup.from_config | a single failing server fails the whole group's connect Confidence: high Scope-risk: moderate Co-Authored-By: Claude Opus 4.8 (1M context) --- src/mitocode/backend/mcp/__init__.py | 13 + src/mitocode/backend/mcp/mcp_manager.py | 350 ++++++++++++ test/test_mcp_manager.py | 701 ++++++++++++++++++++++++ todo.md | 2 +- 4 files changed, 1065 insertions(+), 1 deletion(-) create mode 100644 src/mitocode/backend/mcp/mcp_manager.py create mode 100644 test/test_mcp_manager.py diff --git a/src/mitocode/backend/mcp/__init__.py b/src/mitocode/backend/mcp/__init__.py index 48bdae8..0244795 100644 --- a/src/mitocode/backend/mcp/__init__.py +++ b/src/mitocode/backend/mcp/__init__.py @@ -1 +1,14 @@ """MCP client support: configuration, connections, and lifecycle management.""" + +from .mcp_client import McpClient, McpConnectionState +from .mcp_config import McpServerConfig, McpServerKind +from .mcp_manager import McpManager, McpServerStatus + +__all__ = [ + "McpClient", + "McpConnectionState", + "McpManager", + "McpServerConfig", + "McpServerKind", + "McpServerStatus", +] diff --git a/src/mitocode/backend/mcp/mcp_manager.py b/src/mitocode/backend/mcp/mcp_manager.py new file mode 100644 index 0000000..b0a56a8 --- /dev/null +++ b/src/mitocode/backend/mcp/mcp_manager.py @@ -0,0 +1,350 @@ +"""Lifecycle management for the set of configured MCP servers. + +:class:`McpManager` owns one :class:`~mitocode.backend.mcp.mcp_client.McpClient` +per configured server and decides when each of them starts, refreshes and +stops. It takes no interest in what those tools look like: registration is +delegated to :meth:`ToolManager.register_mcp_client`, which namespaces each +tool under the server that advertises it. + +Two decisions shape the module. Connections are driven one server at a time +rather than through a fastmcp ``ClientGroup``, because a group connects its +members as a unit and one unreachable server would take down every tool of +every other server. And no operation here is allowed to raise on account of +a single server: :meth:`McpManager.start` runs once during startup, and one +broken entry in a shared configuration file must not stop the harness from +starting with the tools that do work. A server that fails costs you that +server's tools and — because :meth:`McpManager.status` reports the reason — +nothing else. +""" + +import asyncio +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional, Set, Union + +from loguru import logger as loguru_logger + +from mitocode.backend.exception.mitocode_exception import MitoCodeMcpConfigError +from mitocode.backend.mcp.mcp_client import McpClient +from mitocode.backend.mcp.mcp_config import McpServerConfig +from mitocode.backend.tool.tool_manager import ToolManager + +logger = loguru_logger.bind(name=__name__) + + +def _describe(exc: Exception) -> str: + """Render an exception as its type plus its message. + + Args: + exc: The exception to describe. + + Returns: + ``": "``. + """ + return f"{type(exc).__name__}: {exc}" + + +def _normalise_servers( + config: Union[Dict[str, Any], List[McpServerConfig]], +) -> List[McpServerConfig]: + """Turn whatever the caller passed as configuration into validated entries. + + A raw ``{"mcpServers": {...}}`` document is handed to + :meth:`McpServerConfig.parse_config`; a list of already-validated entries + — from another layer that loaded them — is taken as it is. + + Args: + config: The MCP configuration, in either accepted shape. + + Returns: + One validated configuration per entry, in configuration order. + + Raises: + MitoCodeMcpConfigError: If *config* is neither shape, if a list element + is not an :class:`McpServerConfig`, or if a raw document holds an + invalid entry. Only an element's position and type are ever + reported, never a value, because ``env``, ``headers`` and ``args`` + carry credentials. + """ + if isinstance(config, dict): + return McpServerConfig.parse_config(config) + + if isinstance(config, list): + servers: List[McpServerConfig] = [] + for index, entry in enumerate(config): + if not isinstance(entry, McpServerConfig): + raise MitoCodeMcpConfigError( + f"MCP server entry {index} must be an McpServerConfig, " + f"got {type(entry).__name__}" + ) + servers.append(entry) + return servers + + raise MitoCodeMcpConfigError( + "MCP configuration must be an object holding an 'mcpServers' mapping or " + f"a list of McpServerConfig, got {type(config).__name__}" + ) + + +@dataclass(frozen=True) +class McpServerStatus: + """A snapshot of one MCP server's state. + + Attributes: + state: ``"connected"``, ``"disconnected"``, ``"failed"`` or ``"disabled"``. + error: The last connection error, empty when there is none. + tool_count: How many tools this server contributed to the registry. + """ + + state: str + error: str + tool_count: int + + +class McpManager: + """Owns the lifecycle of every configured MCP server. + + One client is built per enabled server up front, so a configuration + mistake is reported by the constructor rather than by the first connect + attempt. Registration goes through :meth:`ToolManager.register_mcp_client` + and is undone by name, which is what lets one server be refreshed or + stopped without disturbing another's tools. + + Args: + config: The MCP configuration: a ``{"mcpServers": {...}}`` document, + or a list of :class:`McpServerConfig` entries. + tool_manager: Registry the servers' tools are published into. + client_factory: Builds the client for one server. Exists as a seam so + tests can drive the lifecycle without a subprocess; the default is + the real :class:`McpClient`. + + Note: + ``AgentSession`` snapshots the tool registry at construction time, so + :meth:`start` must complete before any agent session is created, + otherwise existing sessions will not see the MCP tools. + + Note: + The connection deadline is not a manager setting: it comes from each + server's own ``connect_timeout`` (30 seconds by default), because a + cold ``npx -y`` start can take tens of seconds and a single short + global timeout would report a healthy server as failed. + """ + + def __init__( + self, + config: Union[Dict[str, Any], List[McpServerConfig]], + *, + tool_manager: ToolManager, + client_factory: Callable[[McpServerConfig], McpClient] = McpClient, + ) -> None: + """Validate the configuration and build a client for each enabled server. + + Args: + config: The MCP configuration, in either accepted shape. + tool_manager: Registry the servers' tools are published into. + client_factory: Builds the client for one server. + + Raises: + MitoCodeMcpConfigError: If the configuration is invalid in either + shape. Raised before any client is built, so a bad document + cannot leave a half-built manager behind. + """ + servers = _normalise_servers(config) + self._tool_manager = tool_manager + self._client_factory = client_factory + self._clients: Dict[str, McpClient] = {} + self._registered: Dict[str, List[str]] = {} + self._disabled: Set[str] = set() + self._errors: Dict[str, str] = {} + + for server in servers: + if server.enabled: + self._clients[server.name] = client_factory(server) + else: + self._disabled.add(server.name) + + async def start(self) -> None: + """Connect every enabled server and register the tools they advertise. + + Servers are connected concurrently and independently: one that is + unreachable, slow or broken costs its own tools and is described by + :meth:`status`, while every other server still comes up. This method + never raises. + + A server's tools are registered only after its ``connect`` has + reported success, because a disconnected client has no catalogue — + registering then would yield an empty list that reads exactly like a + server that genuinely advertises nothing. + """ + await asyncio.gather( + *(self._connect_one(name) for name in list(self._clients)), + return_exceptions=True, + ) + + async def stop(self) -> None: + """Unregister every MCP tool and close every connection. Idempotent. + + Tools are removed before the connections are closed, so no tool can be + dispatched to a client that is going away. One connection that + refuses to close is recorded against its server in :meth:`status` + instead of stranding the others. + """ + for name in list(self._clients): + await self._unregister_server(name) + + await asyncio.gather( + *(self._disconnect_one(name) for name in list(self._clients)), + return_exceptions=True, + ) + + async def refresh(self, server_name: Optional[str] = None) -> None: + """Reconnect servers so a changed tool catalogue is picked up. + + Every target is unregistered *before* it is reconnected, so the + registry ends up holding the catalogue the server advertises now — + including the tools it has dropped. A reconnect that fails therefore + leaves that server with no tools rather than with tools that dispatch + to a client that is not connected. + + Args: + server_name: The server to refresh; ``None`` refreshes every + enabled server. + + Raises: + MitoCodeMcpConfigError: If *server_name* is neither an enabled nor + a disabled configured server. Never raised for a server that + merely fails to reconnect. + """ + targets = self._refresh_targets(server_name) + await asyncio.gather( + *(self._reconnect_one(name) for name in targets), + return_exceptions=True, + ) + + def status(self) -> Dict[str, McpServerStatus]: + """Describe every configured server, connected or not. + + Returns: + One snapshot per server, keyed by server name. Enabled servers are + reported first, each with its client's own state and error text, + then the disabled ones — which have no client and no tools. + """ + snapshot = { + name: McpServerStatus( + state=client.state.value, + error=client.last_error or self._errors.get(name, ""), + tool_count=len(self._registered.get(name, [])), + ) + for name, client in self._clients.items() + } + for name in self._disabled: + snapshot[name] = McpServerStatus(state="disabled", error="", tool_count=0) + return snapshot + + async def _connect_one(self, name: str) -> None: + """Connect one server and register the tools it advertises. + + Never raises. A failure the client reported itself is already in + :attr:`McpClient.last_error`; one it raised instead is recorded here, + because a server that failed with no stated reason is exactly the case + that is hard to diagnose from the outside. + + Args: + name: Name of the server to connect, as keyed in ``_clients``. + """ + client = self._clients[name] + try: + connected = await client.connect() + except Exception as exc: # pylint: disable=broad-except + self._record_failure(name, exc, "connect") + return + + if not connected: + self._errors.pop(name, None) + logger.warning( + f"MCP server '{name}' is unavailable; its tools are not registered" + ) + return + + # Whatever this server registered last time is released first, so a + # repeated start or refresh cannot leave the previous names behind as + # tools nobody owns and stop() cannot unregister. + await self._unregister_server(name) + try: + specs = self._tool_manager.register_mcp_client(client) + except Exception as exc: # pylint: disable=broad-except + self._record_failure(name, exc, "register its tools") + return + + self._registered[name] = [spec.name for spec in specs] + self._errors.pop(name, None) + + async def _reconnect_one(self, name: str) -> None: + """Unregister, close and reconnect one server. + + Args: + name: Name of the server to reconnect, as keyed in ``_clients``. + """ + await self._unregister_server(name) + await self._disconnect_one(name) + await self._connect_one(name) + + async def _disconnect_one(self, name: str) -> None: + """Close one server's connection, recording a failure instead of raising. + + A teardown that raises is recorded against the server, so a + connection left open is visible in :meth:`status` rather than silently + swallowed. + + Args: + name: Name of the server to close, as keyed in ``_clients``. + """ + try: + await self._clients[name].disconnect() + except Exception as exc: # pylint: disable=broad-except + self._record_failure(name, exc, "disconnect") + + async def _unregister_server(self, name: str) -> None: + """Remove every tool one server registered from the registry. + + Args: + name: Name of the server whose tools are removed. + """ + for tool_name in self._registered.pop(name, []): + self._tool_manager.unregister(tool_name) + + def _record_failure(self, name: str, exc: Exception, action: str) -> None: + """Remember a failure the client did not report itself, and log it. + + Args: + name: Name of the MCP server that failed. + exc: The exception that ended the operation. + action: What was being attempted, phrased to read as + ``failed to ``. + """ + self._errors[name] = _describe(exc) + logger.error(f"MCP server '{name}': failed to {action}: {self._errors[name]}") + + def _refresh_targets(self, server_name: Optional[str]) -> List[str]: + """Resolve the ``refresh`` argument into the servers to reconnect. + + Args: + server_name: The requested server, or ``None`` for all of them. + + Returns: + The names to reconnect. + + Raises: + MitoCodeMcpConfigError: If *server_name* is not configured. + """ + if server_name is None: + return list(self._clients) + if server_name in self._clients: + return [server_name] + if server_name in self._disabled: + raise MitoCodeMcpConfigError( + f"MCP server '{server_name}' is disabled and cannot be refreshed" + ) + raise MitoCodeMcpConfigError( + f"MCP server '{server_name}' is not configured; the configured " + f"servers are {sorted(self._clients)}" + ) diff --git a/test/test_mcp_manager.py b/test/test_mcp_manager.py new file mode 100644 index 0000000..ca0b17e --- /dev/null +++ b/test/test_mcp_manager.py @@ -0,0 +1,701 @@ +import os +import sys +from types import SimpleNamespace + +import pytest + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) + +from mitocode.backend.exception.mitocode_exception import MitoCodeMcpConfigError +from mitocode.backend.mcp.mcp_client import McpClient, McpConnectionState +from mitocode.backend.mcp.mcp_config import McpServerConfig +from mitocode.backend.mcp.mcp_manager import McpManager, McpServerStatus +from mitocode.backend.tool.tool_manager import ToolManager + + +class FakeClient: + """Stand-in for McpClient whose connect outcome is scripted. + + ``state`` tracks the real client's three states rather than inferring one + from ``connected``, so a client that has just been stopped stays + distinguishable from one that never came up. ``status()`` reports that + difference and the manager's contract rests on it. + """ + + def __init__(self, config, succeed=True, tools=()): + """Script one connect outcome and the tool catalogue it advertises. + + Args: + config: The ``McpServerConfig`` this stand-in represents. + succeed: Value :meth:`connect` reports back. + tools: Names of the tools this stand-in advertises. + """ + self._config = config + self._succeed = succeed + self.server_name = config.name + self.tools = [ + SimpleNamespace( + name=name, + description=f"{name} on {config.name}", + input_schema={"type": "object", "properties": {}}, + ) + for name in tools + ] + self.state = McpConnectionState.DISCONNECTED + self.last_error = "" + self.connected = False + self.calls = 0 + + async def connect(self): + """Move to CONNECTED or FAILED according to the scripted outcome.""" + self.connected = self._succeed + if self._succeed: + self.state = McpConnectionState.CONNECTED + self.last_error = "" + else: + self.state = McpConnectionState.FAILED + self.last_error = "connection refused" + return self._succeed + + async def disconnect(self): + """Close the connection and go back to DISCONNECTED.""" + self.connected = False + self.state = McpConnectionState.DISCONNECTED + + def go_offline(self): + """Make every later connect attempt fail, as a server going away would.""" + self._succeed = False + + async def call_tool(self, tool_name, arguments): + """Record the call and render it, as the real client would.""" + self.calls += 1 + return f"{tool_name}:{arguments}" + + +class RudeClient(FakeClient): + """FakeClient that raises from one chosen method. + + The real :class:`McpClient` swallows its own failures and reports them, so + this stand-in is how a test reaches the manager's own guard against a + client that breaks its contract instead. + """ + + def __init__(self, config, fail_on=""): + """Arm one method to raise. + + Args: + config: The ``McpServerConfig`` this stand-in represents. + fail_on: ``"connect"``, ``"disconnect"`` or ``""`` for neither. + """ + super().__init__(config, succeed=fail_on != "connect") + self._fail_on = fail_on + + async def connect(self): + """Raise instead of connecting when this client is armed for it.""" + if self._fail_on == "connect": + raise RuntimeError("handshake exploded") + return await super().connect() + + async def disconnect(self): + """Raise instead of closing when this client is armed for it.""" + if self._fail_on == "disconnect": + raise RuntimeError("teardown exploded") + await super().disconnect() + + +def _factory(script): + """Build a client_factory driven by a {server_name: (succeed, tools)} map. + + Args: + script: Mapping from server name to its scripted outcome. A server + the script does not mention connects successfully with no tools. + + Returns: + The factory, with every client it built exposed as ``made``. + """ + made = {} + + def _make(config): + succeed, tools = script.get(config.name, (True, ())) + client = FakeClient(config, succeed=succeed, tools=tools) + made[config.name] = client + return client + + _make.made = made + return _make + + +def _rude_factory(script): + """Build a client_factory whose clients raise from a scripted method. + + Args: + script: Mapping from server name to the method it raises from, or + ``""`` for a well-behaved client. + + Returns: + The factory, with every client it built exposed as ``made``. + """ + made = {} + + def _make(config): + client = RudeClient(config, fail_on=script.get(config.name, "")) + made[config.name] = client + return client + + _make.made = made + return _make + + +def _config(*entries): + """Build a ``{"mcpServers": {...}}`` document from (name, spec) pairs.""" + return {"mcpServers": {name: spec for name, spec in entries}} + + +@pytest.mark.asyncio +async def test_start_registers_tools_from_every_server(): + tool_manager = ToolManager() + factory = _factory( + { + "weather": (True, ("get_forecast",)), + "files": (True, ("read",)), + } + ) + manager = McpManager( + _config( + ("weather", {"url": "https://w.invalid/mcp"}), + ("files", {"command": "npx", "args": ["files-mcp"]}), + ), + tool_manager=tool_manager, + client_factory=factory, + ) + + await manager.start() + + assert "mcp__weather__get_forecast" in tool_manager.tools_map() + assert "mcp__files__read" in tool_manager.tools_map() + + +@pytest.mark.asyncio +async def test_start_isolates_failures_across_servers(): + tool_manager = ToolManager() + factory = _factory( + { + "broken": (False, ()), + "healthy": (True, ("ping",)), + } + ) + manager = McpManager( + _config( + ("broken", {"url": "https://broken.invalid/mcp"}), + ("healthy", {"url": "https://healthy.invalid/mcp"}), + ), + tool_manager=tool_manager, + client_factory=factory, + ) + + await manager.start() + + assert "mcp__healthy__ping" in tool_manager.tools_map() + assert tool_manager.tools_map().get("mcp__broken__ping") is None + assert manager.status()["broken"].error != "" + assert manager.status()["broken"].state == "failed" + assert manager.status()["healthy"].state == "connected" + + +@pytest.mark.asyncio +async def test_status_reports_state_error_and_tool_count(): + tool_manager = ToolManager() + factory = _factory({"weather": (True, ("a", "b"))}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + + await manager.start() + + status = manager.status()["weather"] + assert isinstance(status, McpServerStatus) + assert status.state == "connected" + assert status.error == "" + assert status.tool_count == 2 + + +@pytest.mark.asyncio +async def test_stop_unregisters_tools_and_disconnects(): + tool_manager = ToolManager() + factory = _factory({"weather": (True, ("get_forecast",))}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + + await manager.start() + await manager.stop() + + assert "mcp__weather__get_forecast" not in tool_manager.tools_map() + assert factory.made["weather"].connected is False + assert manager.status()["weather"].state == "disconnected" + assert manager.status()["weather"].tool_count == 0 + + +@pytest.mark.asyncio +async def test_refresh_picks_up_new_tools(): + tool_manager = ToolManager() + factory = _factory({"weather": (True, ("a",))}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + + await manager.start() + assert "mcp__weather__a" in tool_manager.tools_map() + + factory.made["weather"].tools.append( + SimpleNamespace( + name="b", + description="b", + input_schema={"type": "object", "properties": {}}, + ) + ) + await manager.refresh("weather") + + assert "mcp__weather__a" in tool_manager.tools_map() + assert "mcp__weather__b" in tool_manager.tools_map() + + +@pytest.mark.asyncio +async def test_refresh_rejects_unknown_server_name(): + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=ToolManager(), + client_factory=_factory({}), + ) + + with pytest.raises(MitoCodeMcpConfigError): + await manager.refresh("nope") + + +@pytest.mark.asyncio +async def test_disabled_server_is_not_connected(): + tool_manager = ToolManager() + factory = _factory({"off": (True, ("x",))}) + manager = McpManager( + _config(("off", {"url": "https://off.invalid/mcp", "enabled": False})), + tool_manager=tool_manager, + client_factory=factory, + ) + + await manager.start() + + assert "off" not in factory.made + assert tool_manager.tools_map() == {} + assert manager.status()["off"].state == "disabled" + assert manager.status()["off"].tool_count == 0 + assert manager.status()["off"].error == "" + + +@pytest.mark.asyncio +async def test_start_does_not_raise_when_a_server_connect_raises(): + """A client that breaks its own contract must not stop the other servers. + + The real ``McpClient`` swallows its own failures and returns ``False``, so + this reaches the manager's own guard against a client that raises instead + of reporting. The caller still gets a return, and the healthy server still + publishes its tools. + """ + tool_manager = ToolManager() + factory = _rude_factory({"bad": "connect", "good": ""}) + manager = McpManager( + _config( + ("bad", {"url": "https://bad.invalid/mcp"}), + ("good", {"url": "https://good.invalid/mcp"}), + ), + tool_manager=tool_manager, + client_factory=factory, + ) + + await manager.start() + + assert manager.status()["good"].state == "connected" + assert "RuntimeError" in manager.status()["bad"].error + assert "handshake exploded" in manager.status()["bad"].error + assert manager.status()["bad"].state != "connected" + assert ( + manager.status()["bad"].tool_count == 0 + ), "a server that raised must not look like one that advertised no tools" + + +@pytest.mark.asyncio +async def test_start_does_not_raise_when_a_registration_raises(monkeypatch): + """A catalogue the registry cannot digest must cost only its own server. + + This is the seam the manager cannot otherwise reach: the real client + reports its own failures, and a registration is the manager's own code + raising underneath it. + """ + tool_manager = ToolManager() + real_register = tool_manager.register_mcp_client + + def _register(client): + if client.server_name == "bad": + raise ValueError("catalogue exploded") + return real_register(client) + + monkeypatch.setattr(tool_manager, "register_mcp_client", _register) + factory = _factory({"bad": (True, ("a",)), "good": (True, ("g",))}) + manager = McpManager( + _config( + ("bad", {"url": "https://bad.invalid/mcp"}), + ("good", {"url": "https://good.invalid/mcp"}), + ), + tool_manager=tool_manager, + client_factory=factory, + ) + + await manager.start() + + assert set(tool_manager.tools_map()) == { + "mcp__good__g" + }, "the healthy server still publishes its tools" + assert "catalogue exploded" in manager.status()["bad"].error + assert manager.status()["bad"].tool_count == 0 + + +@pytest.mark.asyncio +async def test_stop_disconnects_every_server_when_one_teardown_raises(): + """One failing teardown must not strand the other connections open.""" + tool_manager = ToolManager() + factory = _rude_factory({"bad": "disconnect", "good": ""}) + manager = McpManager( + _config( + ("bad", {"url": "https://bad.invalid/mcp"}), + ("good", {"url": "https://good.invalid/mcp"}), + ), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + + await manager.stop() + + assert factory.made["good"].state is McpConnectionState.DISCONNECTED + assert tool_manager.tools_map() == {} + assert "RuntimeError" in manager.status()["bad"].error + + +@pytest.mark.asyncio +async def test_start_twice_then_stop_leaves_no_registrations(): + """A repeated start must not orphan the tools the first one registered. + + ``ToolManager.register_mcp_client`` never overwrites an existing name, so + a second registration pass registers nothing and empties the manager's + record of what it owns. ``stop()`` would then have no name left to + unregister and the first pass's tools would outlive the session. + """ + tool_manager = ToolManager() + factory = _factory({"weather": (True, ("a",))}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + + await manager.start() + await manager.start() + assert list(tool_manager.tools_map()) == ["mcp__weather__a"] + + await manager.stop() + + assert ( + tool_manager.tools_map() == {} + ), "every tool a repeated start registered must still be unregistered" + assert manager.status()["weather"].tool_count == 0 + + +@pytest.mark.asyncio +async def test_refresh_drops_tools_the_server_no_longer_advertises(): + """Refreshing must pick up the whole catalogue, not only its additions.""" + tool_manager = ToolManager() + factory = _factory({"weather": (True, ("a", "b"))}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + assert set(tool_manager.tools_map()) == {"mcp__weather__a", "mcp__weather__b"} + + factory.made["weather"].tools.pop() + await manager.refresh("weather") + + assert set(tool_manager.tools_map()) == { + "mcp__weather__a" + }, "a tool the server dropped must not stay dispatchable to a dead client" + + +@pytest.mark.asyncio +async def test_refresh_drops_stale_tools_when_the_reconnect_fails(): + """A failed refresh leaves that server with no tools, not with dead ones. + + The tools must be unregistered *before* the reconnect is attempted: kept + across a failure they would advertise a capability no model can call, and + would dispatch into a client that is not connected. + """ + tool_manager = ToolManager() + factory = _factory({"weather": (True, ("a", "b"))}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + assert set(tool_manager.tools_map()) == {"mcp__weather__a", "mcp__weather__b"} + + factory.made["weather"].go_offline() + await manager.refresh("weather") + + assert ( + tool_manager.tools_map() == {} + ), "a server that could not be reached must publish nothing" + assert manager.status()["weather"].state == "failed" + + +@pytest.mark.asyncio +async def test_refresh_without_a_name_reconnects_every_server(): + tool_manager = ToolManager() + factory = _factory({"one": (True, ("a",)), "two": (True, ("b",))}) + manager = McpManager( + _config( + ("one", {"url": "https://one.invalid/mcp"}), + ("two", {"url": "https://two.invalid/mcp"}), + ), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + + await manager.refresh() + + assert factory.made["one"].state is McpConnectionState.CONNECTED + assert factory.made["two"].state is McpConnectionState.CONNECTED + assert set(tool_manager.tools_map()) == {"mcp__one__a", "mcp__two__b"} + + +@pytest.mark.asyncio +async def test_refresh_of_a_disabled_server_is_rejected_by_name(): + """A disabled server has no client, so the rejection must say why.""" + manager = McpManager( + _config(("off", {"url": "https://off.invalid/mcp", "enabled": False})), + tool_manager=ToolManager(), + client_factory=_factory({}), + ) + + with pytest.raises(MitoCodeMcpConfigError) as excinfo: + await manager.refresh("off") + + assert "off" in str(excinfo.value) + assert "disabled" in str(excinfo.value) + + +@pytest.mark.asyncio +async def test_refresh_error_names_the_configured_servers(): + manager = McpManager( + _config( + ("weather", {"url": "https://w.invalid/mcp"}), + ("files", {"url": "https://f.invalid/mcp"}), + ), + tool_manager=ToolManager(), + client_factory=_factory({}), + ) + + with pytest.raises(MitoCodeMcpConfigError) as excinfo: + await manager.refresh("nope") + + message = str(excinfo.value) + assert "nope" in message + assert "weather" in message and "files" in message + + +@pytest.mark.asyncio +async def test_refresh_error_does_not_echo_server_credentials(): + """The message names a server, never the configuration behind it.""" + manager = McpManager( + _config( + ( + "files", + { + "url": "https://files.invalid/mcp", + "headers": {"Authorization": "Bearer s3cr3t-token"}, + }, + ) + ), + tool_manager=ToolManager(), + client_factory=_factory({}), + ) + + with pytest.raises(MitoCodeMcpConfigError) as excinfo: + await manager.refresh("nope") + + assert "s3cr3t-token" not in str(excinfo.value) + + +def test_a_list_of_configs_is_accepted_as_the_configuration(): + """Configurations normalised elsewhere need no round trip through JSON.""" + tool_manager = ToolManager() + factory = _factory({"weather": (True, ("a",))}) + manager = McpManager( + [McpServerConfig.parse("weather", {"url": "https://w.invalid/mcp"})], + tool_manager=tool_manager, + client_factory=factory, + ) + + assert "weather" in factory.made + assert manager.status()["weather"].state == "disconnected" + assert manager.status()["weather"].tool_count == 0 + + +def test_a_configuration_that_is_neither_a_document_nor_a_list_is_rejected(): + with pytest.raises(MitoCodeMcpConfigError): + McpManager( + "mcpServers", tool_manager=ToolManager(), client_factory=_factory({}) + ) + + +def test_a_list_entry_that_is_not_a_config_is_rejected(): + with pytest.raises(MitoCodeMcpConfigError) as excinfo: + McpManager( + [{"url": "https://w.invalid/mcp"}], + tool_manager=ToolManager(), + client_factory=_factory({}), + ) + + assert "0" in str(excinfo.value), "the offending position is named" + + +def test_an_invalid_server_entry_is_reported_before_any_client_is_built(): + """Configuration is validated as a whole, so no half-built manager escapes.""" + factory = _factory({"good": (True, ("a",))}) + config = { + "mcpServers": { + "good": {"url": "https://good.invalid/mcp"}, + "bad": {"url": "not-a-url"}, + } + } + + with pytest.raises(MitoCodeMcpConfigError): + McpManager(config, tool_manager=ToolManager(), client_factory=factory) + + assert factory.made == {}, "a rejected configuration connects to nothing" + + +@pytest.mark.asyncio +async def test_start_after_stop_brings_the_tools_back(): + """Stopping is a pause, not a one-way door: the same manager restarts.""" + tool_manager = ToolManager() + factory = _factory({"weather": (True, ("a",))}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + await manager.stop() + assert tool_manager.tools_map() == {} + + await manager.start() + + assert set(tool_manager.tools_map()) == {"mcp__weather__a"} + assert manager.status()["weather"].state == "connected" + assert manager.status()["weather"].tool_count == 1 + + +@pytest.mark.asyncio +async def test_start_with_no_servers_at_all_does_nothing(): + tool_manager = ToolManager() + manager = McpManager( + {"mcpServers": {}}, tool_manager=tool_manager, client_factory=_factory({}) + ) + + await manager.start() + + assert manager.status() == {} + assert tool_manager.tools_map() == {} + + +@pytest.mark.asyncio +async def test_status_before_start_reports_a_disconnected_server(): + """A manager that has not started yet says so rather than looking idle.""" + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=ToolManager(), + client_factory=_factory({"weather": (True, ("a",))}), + ) + + status = manager.status()["weather"] + + assert status.state == "disconnected" + assert status.tool_count == 0 + assert status.error == "" + + +def test_mcp_package_exports_the_manager_and_its_neighbours(): + from mitocode.backend import mcp + + assert mcp.McpManager is McpManager + assert mcp.McpServerStatus is McpServerStatus + assert mcp.McpClient is McpClient + assert mcp.McpConnectionState is McpConnectionState + assert mcp.McpServerConfig is McpServerConfig + + +@pytest.mark.asyncio +async def test_local_stdio_server_end_to_end(tmp_path): + """The local-server path, proven against a real subprocess. + + No ``client_factory`` here: this is the only test in the feature that + exercises the whole chain — a real ``McpClient``, a real stdio transport + and a real ``sys.executable`` child serving over MCP — so it is what + justifies the transport plumbing underneath it. + """ + server_script = tmp_path / "srv.py" + server_script.write_text( + "from fastmcp import FastMCP\n" + "app = FastMCP('local')\n" + "@app.tool\n" + "def shout(word: str) -> str:\n" + " '''Shout a word.'''\n" + " return word.upper()\n" + "app.run()\n" + ) + + tool_manager = ToolManager() + manager = McpManager( + { + "mcpServers": { + "local": { + "command": sys.executable, + "args": [str(server_script)], + "connect_timeout": 60.0, + } + } + }, + tool_manager=tool_manager, + ) + + await manager.start() + try: + status = manager.status()["local"] + assert status.state == "connected", status.error + assert status.tool_count == 1 + assert "mcp__local__shout" in tool_manager.tools_map() + + result = await tool_manager.tools_map()["mcp__local__shout"](word="hey") + assert result == "HEY" + finally: + await manager.stop() + + assert tool_manager.tools_map() == {} + assert manager.status()["local"].state == "disconnected" diff --git a/todo.md b/todo.md index b3f2e8f..a828a61 100644 --- a/todo.md +++ b/todo.md @@ -11,7 +11,7 @@ - [x] 重构Agent对外接口 (Agent → AgentSession/AgentRuntime/AgentOrchestrator 三层分离) - [ ] 添加上下文管理功能 - [ ] 添加索引功能 -- [ ] 添加MCP功能 +- [x] 添加MCP功能 - [ ] 添加Agent Skills功能 - [x] 完成TUI (TUI直接调用backend，移除HTTP/SSE层) - [x] 添加工具管理功能:PR #6 From 3289605a19b8937a9f3a8c50d19afe094f630b4e Mon Sep 17 00:00:00 2001 From: Allen Wang Date: Sat, 17 Oct 2026 14:44:25 +0800 Subject: [PATCH 18/20] fix(mcp): hold the one-way dependency, one status per server, and a guarded teardown MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three defects the review found, all in the manager's own file. The mcp package imported tool_manager at runtime, breaking the plan's one-way dependency and making mitocode.backend.mcp unimportable without the tool layer. The import moves under TYPE_CHECKING with the annotation quoted, matching tool_manager's own guard for McpClient; the manager only ever needed the two methods it calls. A repeated server name left a tool registered, dispatchable and reported as the other entry's — "disabled", tool_count 0. Names are refused in _normalise_servers instead, naming the position and the name but never a value. Two different servers offering the same tool name are unaffected: that collision guard keys on mcp____. stop()'s unregister loop was the module's only unguarded await, and it ran before the disconnects, so a registry that refused to give a tool back skipped teardown entirely and left every connection open. Both phases now run through _run_each, which isolates each server and inspects the gather results: return_exceptions=True also swallows a child's CancelledError, which except Exception cannot catch, so a cancelled connect was dropped with no trace. It is recorded and re-raised. unregister()'s bool now decides ownership instead of being discarded: a tool that was removed, or that was already gone, is dropped from the record, while one whose removal raised stays owned so a later stop() can finish the job. status() prefers a recorded failure over the client's own older message, and stops reporting one a successful retry resolved. Constraint: start() must still never raise for a server that fails; only a cancellation propagates Rejected: setdefault in status() | hides the collision instead of refusing a configuration that cannot describe itself consistently Rejected: disconnect before unregister | leaves a registered tool pointing at a dead client, failing in the model's hands instead of here Confidence: high Scope-risk: moderate Co-Authored-By: Claude Opus 4.8 (1M context) --- src/mitocode/backend/mcp/mcp_manager.py | 233 +++++++++--- test/test_mcp_manager.py | 481 +++++++++++++++++++++++- 2 files changed, 647 insertions(+), 67 deletions(-) diff --git a/src/mitocode/backend/mcp/mcp_manager.py b/src/mitocode/backend/mcp/mcp_manager.py index b0a56a8..1712a32 100644 --- a/src/mitocode/backend/mcp/mcp_manager.py +++ b/src/mitocode/backend/mcp/mcp_manager.py @@ -19,28 +19,48 @@ import asyncio from dataclasses import dataclass -from typing import Any, Callable, Dict, List, Optional, Set, Union +from typing import ( + Any, + Awaitable, + Callable, + Dict, + List, + Optional, + Set, + TYPE_CHECKING, + Union, +) from loguru import logger as loguru_logger from mitocode.backend.exception.mitocode_exception import MitoCodeMcpConfigError from mitocode.backend.mcp.mcp_client import McpClient from mitocode.backend.mcp.mcp_config import McpServerConfig -from mitocode.backend.tool.tool_manager import ToolManager + +if TYPE_CHECKING: + # The tool layer is reached only through the object handed in, never by + # importing it: the dependency runs one way, tool_manager -> mcp. A runtime + # import here would make mitocode.backend.mcp unimportable on its own and + # would break the symmetry with tool_manager's own TYPE_CHECKING guard for + # McpClient. The annotation is quoted because this module has no + # `from __future__ import annotations`, so annotations are evaluated. + from mitocode.backend.tool.tool_manager import ToolManager logger = loguru_logger.bind(name=__name__) -def _describe(exc: Exception) -> str: +def _describe(exc: BaseException) -> str: """Render an exception as its type plus its message. Args: exc: The exception to describe. Returns: - ``": "``. + ``": "``, or just the type when the exception + carries no message — a bare ``CancelledError`` has none. """ - return f"{type(exc).__name__}: {exc}" + message = str(exc) + return f"{type(exc).__name__}: {message}" if message else type(exc).__name__ def _normalise_servers( @@ -49,20 +69,31 @@ def _normalise_servers( """Turn whatever the caller passed as configuration into validated entries. A raw ``{"mcpServers": {...}}`` document is handed to - :meth:`McpServerConfig.parse_config`; a list of already-validated entries - — from another layer that loaded them — is taken as it is. + :meth:`McpServerConfig.parse_config`, whose keys are unique by + construction; a list of already-validated entries — from another layer that + loaded them — is taken as it is, subject to the same uniqueness rule. + + A repeated name is refused rather than resolved. The name is both the key + a client is stored under and the ``mcp____`` namespace its tools are + published in, so letting two entries share one would leave a tool that is + registered and dispatchable while :meth:`McpManager.status` describes it as + the *other* entry's — ``disabled`` with ``tool_count=0``, say. Whichever + entry happened to come last would silently win, and the tools of the one + that lost would be unreachable with no way to tell that is what happened. Args: config: The MCP configuration, in either accepted shape. Returns: - One validated configuration per entry, in configuration order. + One validated configuration per entry, in configuration order, with + every name distinct. Raises: MitoCodeMcpConfigError: If *config* is neither shape, if a list element - is not an :class:`McpServerConfig`, or if a raw document holds an - invalid entry. Only an element's position and type are ever - reported, never a value, because ``env``, ``headers`` and ``args`` + is not an :class:`McpServerConfig`, if two entries share a name, or + if a raw document holds an invalid entry. Only an element's + position, its type and the duplicated name are ever reported, never + a configuration value, because ``env``, ``headers`` and ``args`` carry credentials. """ if isinstance(config, dict): @@ -70,12 +101,20 @@ def _normalise_servers( if isinstance(config, list): servers: List[McpServerConfig] = [] + seen: Dict[str, int] = {} for index, entry in enumerate(config): if not isinstance(entry, McpServerConfig): raise MitoCodeMcpConfigError( f"MCP server entry {index} must be an McpServerConfig, " f"got {type(entry).__name__}" ) + if entry.name in seen: + raise MitoCodeMcpConfigError( + f"MCP server entry {index} reuses the name {entry.name!r} " + f"already taken by entry {seen[entry.name]}; server names " + "must be unique because they namespace that server's tools" + ) + seen[entry.name] = index servers.append(entry) return servers @@ -133,20 +172,28 @@ def __init__( self, config: Union[Dict[str, Any], List[McpServerConfig]], *, - tool_manager: ToolManager, + tool_manager: "ToolManager", client_factory: Callable[[McpServerConfig], McpClient] = McpClient, ) -> None: """Validate the configuration and build a client for each enabled server. + The tool registry is held as a duck-typed object: the module never + imports the tool layer, so any object exposing + :meth:`~mitocode.backend.tool.tool_manager.ToolManager.register_mcp_client` + and :meth:`~mitocode.backend.tool.tool_manager.ToolManager.unregister` + will do. + Args: - config: The MCP configuration, in either accepted shape. + config: The MCP configuration, in either accepted shape. Every + server name must be distinct. tool_manager: Registry the servers' tools are published into. client_factory: Builds the client for one server. Raises: MitoCodeMcpConfigError: If the configuration is invalid in either - shape. Raised before any client is built, so a bad document - cannot leave a half-built manager behind. + shape, or repeats a server name. Raised before any client is + built, so a bad document cannot leave a half-built manager + behind. """ servers = _normalise_servers(config) self._tool_manager = tool_manager @@ -168,33 +215,41 @@ async def start(self) -> None: Servers are connected concurrently and independently: one that is unreachable, slow or broken costs its own tools and is described by :meth:`status`, while every other server still comes up. This method - never raises. + never raises for a server that fails. A server's tools are registered only after its ``connect`` has reported success, because a disconnected client has no catalogue — registering then would yield an empty list that reads exactly like a server that genuinely advertises nothing. + + Raises: + asyncio.CancelledError: If a server's connect is cancelled. + ``McpClient.connect`` re-raises cancellation by contract, so it + is recorded against that server in :meth:`status` and then + propagated: a cancellation is the caller's decision to unwind, + not a server that failed. """ - await asyncio.gather( - *(self._connect_one(name) for name in list(self._clients)), - return_exceptions=True, - ) + await self._run_each(list(self._clients), self._connect_one) async def stop(self) -> None: """Unregister every MCP tool and close every connection. Idempotent. Tools are removed before the connections are closed, so no tool can be - dispatched to a client that is going away. One connection that - refuses to close is recorded against its server in :meth:`status` - instead of stranding the others. - """ - for name in list(self._clients): - await self._unregister_server(name) + dispatched to a client that is going away — the opposite order would + leave a registered tool pointing at a dead client, which fails at + call time in the model's hands rather than here. - await asyncio.gather( - *(self._disconnect_one(name) for name in list(self._clients)), - return_exceptions=True, - ) + Both phases are isolated per server. A registry that refuses to give a + tool back must not skip the teardown that releases the subprocesses, + and a connection that refuses to close must not strand the others; + either is recorded against its server in :meth:`status` instead. + + Raises: + asyncio.CancelledError: If one server's unregistration or + disconnection is cancelled. + """ + await self._run_each(list(self._clients), self._unregister_one) + await self._run_each(list(self._clients), self._disconnect_one) async def refresh(self, server_name: Optional[str] = None) -> None: """Reconnect servers so a changed tool catalogue is picked up. @@ -205,6 +260,10 @@ async def refresh(self, server_name: Optional[str] = None) -> None: leaves that server with no tools rather than with tools that dispatch to a client that is not connected. + A server that has not been started yet is connected by this call. + Refusing that would make a manager being recovered after a crash reject + the one call meant to recover it. + Args: server_name: The server to refresh; ``None`` refreshes every enabled server. @@ -213,25 +272,29 @@ async def refresh(self, server_name: Optional[str] = None) -> None: MitoCodeMcpConfigError: If *server_name* is neither an enabled nor a disabled configured server. Never raised for a server that merely fails to reconnect. + asyncio.CancelledError: If one target's work is cancelled. """ - targets = self._refresh_targets(server_name) - await asyncio.gather( - *(self._reconnect_one(name) for name in targets), - return_exceptions=True, - ) + await self._run_each(self._refresh_targets(server_name), self._reconnect_one) def status(self) -> Dict[str, McpServerStatus]: """Describe every configured server, connected or not. + A server name resolves to exactly one entry: :func:`_normalise_servers` + refuses a repeated name, so a server is never both enabled and + disabled. + Returns: One snapshot per server, keyed by server name. Enabled servers are - reported first, each with its client's own state and error text, - then the disabled ones — which have no client and no tools. + reported first, each with its client's own state and its most recent + error, then the disabled ones — which have no client and no tools. """ snapshot = { name: McpServerStatus( state=client.state.value, - error=client.last_error or self._errors.get(name, ""), + # A failure the manager recorded is the more recent of the two: + # the client's own text describes an earlier attempt, and the + # recorded one describes the step that just failed. + error=self._errors.get(name, "") or client.last_error, tool_count=len(self._registered.get(name, [])), ) for name, client in self._clients.items() @@ -240,6 +303,39 @@ def status(self) -> Dict[str, McpServerStatus]: snapshot[name] = McpServerStatus(state="disabled", error="", tool_count=0) return snapshot + async def _run_each( + self, names: List[str], operation: Callable[[str], Awaitable[None]] + ) -> None: + """Run *operation* once per name, concurrently, isolating failures. + + Every result is awaited and every name is visited, so one server that + raises cannot cost any other server its turn — which is the whole point + of the gather and the reason it is not a plain loop. + + The results are then inspected rather than discarded. + ``return_exceptions=True`` hands back a child's ``CancelledError`` as a + *result* instead of raising it, and ``except Exception`` cannot catch + one, so a cancelled connect would otherwise be dropped with no trace + and the server would report a failure with no reason at all. It is + recorded against that server and re-raised, because a cancellation is + an instruction to unwind, not a server that failed. + + Args: + names: Servers to run *operation* for, each visited exactly once. + operation: One server's worth of work, e.g. :meth:`_connect_one`. + + Raises: + asyncio.CancelledError: If *operation* was cancelled for one server. + """ + results = await asyncio.gather( + *(operation(name) for name in names), + return_exceptions=True, + ) + for name, result in zip(names, results): + if isinstance(result, asyncio.CancelledError): + self._record_failure(name, result, "cancelled before it finished") + raise result + async def _connect_one(self, name: str) -> None: """Connect one server and register the tools it advertises. @@ -255,7 +351,7 @@ async def _connect_one(self, name: str) -> None: try: connected = await client.connect() except Exception as exc: # pylint: disable=broad-except - self._record_failure(name, exc, "connect") + self._record_failure(name, exc, "failed to connect") return if not connected: @@ -268,11 +364,11 @@ async def _connect_one(self, name: str) -> None: # Whatever this server registered last time is released first, so a # repeated start or refresh cannot leave the previous names behind as # tools nobody owns and stop() cannot unregister. - await self._unregister_server(name) + await self._unregister_one(name) try: specs = self._tool_manager.register_mcp_client(client) except Exception as exc: # pylint: disable=broad-except - self._record_failure(name, exc, "register its tools") + self._record_failure(name, exc, "failed to register its tools") return self._registered[name] = [spec.name for spec in specs] @@ -284,7 +380,7 @@ async def _reconnect_one(self, name: str) -> None: Args: name: Name of the server to reconnect, as keyed in ``_clients``. """ - await self._unregister_server(name) + await self._unregister_one(name) await self._disconnect_one(name) await self._connect_one(name) @@ -301,28 +397,65 @@ async def _disconnect_one(self, name: str) -> None: try: await self._clients[name].disconnect() except Exception as exc: # pylint: disable=broad-except - self._record_failure(name, exc, "disconnect") + self._record_failure(name, exc, "failed to disconnect") - async def _unregister_server(self, name: str) -> None: + async def _unregister_one(self, name: str) -> None: """Remove every tool one server registered from the registry. + The ownership record is updated per tool and only once the registry has + confirmed the outcome, which is the reason ``unregister``'s ``bool`` + return is not discarded. Three outcomes are possible and only two let + go of the name: + + * removed — the tool is gone, so drop it from the record; + * already absent — somebody removed it behind the manager's back, so + there is nothing left to retry and it is dropped too; + * raised — the outcome is unknown, so the name stays owned and a later + :meth:`stop` can finish the job. Forgetting it here would leave a + tool in the registry that no one will ever remove. + + A failure here is recorded and never raised, so it cannot skip the + teardown that releases this and every other server's resources. + Args: name: Name of the server whose tools are removed. """ - for tool_name in self._registered.pop(name, []): - self._tool_manager.unregister(tool_name) + remaining: List[str] = [] + for tool_name in self._registered.get(name, []): + try: + removed = self._tool_manager.unregister(tool_name) + except Exception as exc: # pylint: disable=broad-except + remaining.append(tool_name) + self._record_failure( + name, exc, f"failed to unregister the tool {tool_name!r}" + ) + continue + if not removed: + logger.warning( + f"MCP tool '{tool_name}' was not in the registry; " + f"server '{name}' no longer owns it" + ) + if remaining: + self._registered[name] = remaining + return + self._registered.pop(name, None) + # Nothing failed and nothing is left, so whatever this server last + # complained about is history. Leaving it in place would have status() + # report a failure that no longer describes anything. + self._errors.pop(name, None) - def _record_failure(self, name: str, exc: Exception, action: str) -> None: + def _record_failure(self, name: str, exc: BaseException, problem: str) -> None: """Remember a failure the client did not report itself, and log it. Args: name: Name of the MCP server that failed. exc: The exception that ended the operation. - action: What was being attempted, phrased to read as - ``failed to ``. + problem: A complete phrase saying what went wrong, e.g. + ``"failed to connect"`` — it is logged verbatim and may name a + tool, which is never a credential. """ self._errors[name] = _describe(exc) - logger.error(f"MCP server '{name}': failed to {action}: {self._errors[name]}") + logger.error(f"MCP server '{name}': {problem}: {self._errors[name]}") def _refresh_targets(self, server_name: Optional[str]) -> List[str]: """Resolve the ``refresh`` argument into the servers to reconnect. diff --git a/test/test_mcp_manager.py b/test/test_mcp_manager.py index ca0b17e..cb79065 100644 --- a/test/test_mcp_manager.py +++ b/test/test_mcp_manager.py @@ -1,4 +1,6 @@ +import asyncio import os +import subprocess import sys from types import SimpleNamespace @@ -8,10 +10,86 @@ from mitocode.backend.exception.mitocode_exception import MitoCodeMcpConfigError from mitocode.backend.mcp.mcp_client import McpClient, McpConnectionState -from mitocode.backend.mcp.mcp_config import McpServerConfig +from mitocode.backend.mcp.mcp_config import McpServerConfig, McpServerKind from mitocode.backend.mcp.mcp_manager import McpManager, McpServerStatus from mitocode.backend.tool.tool_manager import ToolManager +#: Imports the ``mcp`` package in a process where the tool layer cannot be +#: resolved at all, so the one-way ``mcp`` -> (never) ``tool_manager`` +#: dependency is checked for real rather than by reading the source. +#: +#: Both ``mitocode`` and ``mitocode.backend`` are stubbed as bare packages: +#: their real ``__init__`` files re-export the agent layer, which imports +#: ``tool_manager`` itself, so a real import of either parent would prove +#: nothing about the module under test. The stubs carry only ``__path__``, +#: which is what lets ``mitocode.backend.mcp`` load from disk as usual. +_IMPORT_WITHOUT_TOOL_LAYER = """ +import importlib.abc +import os +import sys +import types + +SRC_DIR = os.environ["MITOCODE_SRC"] + + +class _BlockToolLayer(importlib.abc.MetaPathFinder): + \"\"\"Make the whole tool layer unimportable, however it is reached.\"\"\" + + def find_spec(self, fullname, path=None, target=None): + if fullname.startswith("mitocode.backend.tool"): + raise ImportError(f"{fullname} is not importable here") + return None + + +def _stub_package(name, path): + \"\"\"Register a package that never runs its real __init__. + + Args: + name: Fully qualified package name. + path: Directory its submodules live in. + \"\"\" + module = types.ModuleType(name) + module.__path__ = [path] + sys.modules[name] = module + + +_stub_package("mitocode", SRC_DIR) +_stub_package("mitocode.backend", os.path.join(SRC_DIR, "mitocode", "backend")) +sys.meta_path.insert(0, _BlockToolLayer()) + +import mitocode.backend.mcp as mcp + + +def leaked(): + \"\"\"List the tool-layer modules this import dragged in.\"\"\" + return sorted(m for m in sys.modules if m.startswith("mitocode.backend.tool")) + + +assert not leaked(), f"importing the mcp package pulled in the tool layer: {leaked()}" + + +class _StubRegistry: + \"\"\"Duck-typed stand-in for ToolManager, to drive the manager offline.\"\"\" + + def register_mcp_client(self, client): + \"\"\"Register nothing and report nothing registered.\"\"\" + return [] + + def unregister(self, name): + \"\"\"Report a removal that never happened.\"\"\" + return False + + +# The manager must work with anything exposing those two methods: that is what +# lets the mcp package stay unaware of the tool layer. +manager = mcp.McpManager( + {"mcpServers": {"w": {"url": "https://w.invalid/mcp"}}}, + tool_manager=_StubRegistry(), +) +assert manager.status()["w"].state == "disconnected" +print("ok") +""" + class FakeClient: """Stand-in for McpClient whose connect outcome is scripted. @@ -20,18 +98,27 @@ class FakeClient: from ``connected``, so a client that has just been stopped stays distinguishable from one that never came up. ``status()`` reports that difference and the manager's contract rests on it. + + The catalogue follows the real client's rule that a client which is not + connected has no tools: it is emptied on construction when the scripted + outcome is a failure, and again by every failed :meth:`connect`. """ - def __init__(self, config, succeed=True, tools=()): + def __init__(self, config, succeed=True, tools=(), keep_tools_on_failure=False): """Script one connect outcome and the tool catalogue it advertises. Args: config: The ``McpServerConfig`` this stand-in represents. succeed: Value :meth:`connect` reports back. - tools: Names of the tools this stand-in advertises. + tools: Names of the tools this stand-in advertises once connected. + keep_tools_on_failure: Leave the catalogue populated even while + not connected. Deliberately unfaithful: it is how a test + reaches an implementation that registers before connecting, + which the real client would not let it get away with. """ self._config = config self._succeed = succeed + self.keep_tools_on_failure = keep_tools_on_failure self.server_name = config.name self.tools = [ SimpleNamespace( @@ -41,21 +128,30 @@ def __init__(self, config, succeed=True, tools=()): ) for name in tools ] + if not succeed and not keep_tools_on_failure: + self.tools.clear() self.state = McpConnectionState.DISCONNECTED self.last_error = "" self.connected = False self.calls = 0 async def connect(self): - """Move to CONNECTED or FAILED according to the scripted outcome.""" + """Move to CONNECTED or FAILED according to the scripted outcome. + + Returns: + The scripted outcome. + """ self.connected = self._succeed if self._succeed: self.state = McpConnectionState.CONNECTED self.last_error = "" - else: - self.state = McpConnectionState.FAILED - self.last_error = "connection refused" - return self._succeed + return True + # The real client drops its cached catalogue when a connect fails. + if not self.keep_tools_on_failure: + self.tools.clear() + self.state = McpConnectionState.FAILED + self.last_error = "connection refused" + return False async def disconnect(self): """Close the connection and go back to DISCONNECTED.""" @@ -80,20 +176,23 @@ class RudeClient(FakeClient): client that breaks its contract instead. """ - def __init__(self, config, fail_on=""): - """Arm one method to raise. + def __init__(self, config, succeed=True, fail_on=""): + """Arm one method to raise, independently of the scripted outcome. Args: config: The ``McpServerConfig`` this stand-in represents. - fail_on: ``"connect"``, ``"disconnect"`` or ``""`` for neither. + succeed: Value :meth:`connect` reports back when it does not raise. + fail_on: ``"connect"``, ``"disconnect"``, ``"cancel"`` or ``""``. """ - super().__init__(config, succeed=fail_on != "connect") + super().__init__(config, succeed=succeed) self._fail_on = fail_on async def connect(self): - """Raise instead of connecting when this client is armed for it.""" + """Raise, or cancel, instead of connecting when this client is armed.""" if self._fail_on == "connect": raise RuntimeError("handshake exploded") + if self._fail_on == "cancel": + raise asyncio.CancelledError() return await super().connect() async def disconnect(self): @@ -103,12 +202,14 @@ async def disconnect(self): await super().disconnect() -def _factory(script): +def _factory(script, keep_tools_on_failure=False): """Build a client_factory driven by a {server_name: (succeed, tools)} map. Args: script: Mapping from server name to its scripted outcome. A server the script does not mention connects successfully with no tools. + keep_tools_on_failure: Give every client a catalogue that survives a + failed connect, for the registration-ordering probe. Returns: The factory, with every client it built exposed as ``made``. @@ -117,7 +218,12 @@ def _factory(script): def _make(config): succeed, tools = script.get(config.name, (True, ())) - client = FakeClient(config, succeed=succeed, tools=tools) + client = FakeClient( + config, + succeed=succeed, + tools=tools, + keep_tools_on_failure=keep_tools_on_failure, + ) made[config.name] = client return client @@ -125,20 +231,27 @@ def _make(config): return _make -def _rude_factory(script): +def _rude_factory(script, succeed=None): """Build a client_factory whose clients raise from a scripted method. Args: script: Mapping from server name to the method it raises from, or ``""`` for a well-behaved client. + succeed: Optional mapping from server name to the connect outcome used + when that method does not raise, overriding the default of ``True``. Returns: The factory, with every client it built exposed as ``made``. """ made = {} + outcomes = succeed or {} def _make(config): - client = RudeClient(config, fail_on=script.get(config.name, "")) + client = RudeClient( + config, + succeed=outcomes.get(config.name, True), + fail_on=script.get(config.name, ""), + ) made[config.name] = client return client @@ -151,6 +264,33 @@ def _config(*entries): return {"mcpServers": {name: spec for name, spec in entries}} +def _run_import_probe(script): + """Run *script* in a child interpreter that cannot resolve the tool layer. + + ``MITOCODE_SRC`` is absolute and the child runs in an empty *cwd*, so + ``python -c`` cannot pick up a stray ``mitocode/`` package that would + shadow the module under test. + + Args: + script: Source for the child interpreter. + + Returns: + The completed ``subprocess.CompletedProcess``. + """ + env = dict(os.environ) + env["MITOCODE_SRC"] = os.path.abspath( + os.path.join(os.path.dirname(__file__), "..", "src") + ) + return subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + env=env, + cwd=env["MITOCODE_SRC"], + check=False, + ) + + @pytest.mark.asyncio async def test_start_registers_tools_from_every_server(): tool_manager = ToolManager() @@ -221,6 +361,39 @@ async def test_status_reports_state_error_and_tool_count(): assert status.tool_count == 2 +@pytest.mark.asyncio +async def test_tools_are_registered_only_after_a_successful_connect(): + """A server that failed to connect must publish nothing. + + The client is armed with a catalogue it never connected well enough to + discover, which the real client does not do — that is precisely the shape + an implementation that registers *before* connecting would publish. + """ + tool_manager = ToolManager() + factory = _factory( + {"broken": (False, ("ping",)), "healthy": (True, ("pong",))}, + keep_tools_on_failure=True, + ) + manager = McpManager( + _config( + ("broken", {"url": "https://broken.invalid/mcp"}), + ("healthy", {"url": "https://healthy.invalid/mcp"}), + ), + tool_manager=tool_manager, + client_factory=factory, + ) + + await manager.start() + + assert set(tool_manager.tools_map()) == { + "mcp__healthy__pong" + }, "only the connected server's tools may be registered" + assert manager.status()["broken"].tool_count == 0 + assert ( + factory.made["broken"].tools != [] + ), "the probe catalogue survived, so it proves the point" + + @pytest.mark.asyncio async def test_stop_unregisters_tools_and_disconnects(): tool_manager = ToolManager() @@ -266,6 +439,29 @@ async def test_refresh_picks_up_new_tools(): assert "mcp__weather__b" in tool_manager.tools_map() +@pytest.mark.asyncio +async def test_refresh_before_start_connects_the_server_it_names(): + """``refresh`` is not gated on ``start``: it brings up what it names. + + Documented rather than rejected, because refusing it would make a manager + that is being reconnected after a crash refuse the one call meant to + recover it. + """ + tool_manager = ToolManager() + factory = _factory({"weather": (True, ("a",))}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + + await manager.refresh("weather") + + assert factory.made["weather"].state is McpConnectionState.CONNECTED + assert set(tool_manager.tools_map()) == {"mcp__weather__a"} + assert manager.status()["weather"].tool_count == 1 + + @pytest.mark.asyncio async def test_refresh_rejects_unknown_server_name(): manager = McpManager( @@ -364,6 +560,35 @@ def _register(client): assert manager.status()["bad"].tool_count == 0 +@pytest.mark.asyncio +async def test_start_records_a_cancelled_server_and_lets_the_cancellation_out(): + """A cancellation is recorded *and* propagated, never swallowed. + + ``return_exceptions=True`` hands a child's ``CancelledError`` back as a + result rather than raising it, and ``except Exception`` cannot catch one, + so without this the server would report a failure with no reason at all. + """ + tool_manager = ToolManager() + factory = _rude_factory({"bad": "cancel", "good": ""}) + manager = McpManager( + _config( + ("bad", {"url": "https://bad.invalid/mcp"}), + ("good", {"url": "https://good.invalid/mcp"}), + ), + tool_manager=tool_manager, + client_factory=factory, + ) + + with pytest.raises(asyncio.CancelledError): + await manager.start() + + assert manager.status()["bad"].error == "CancelledError" + assert manager.status()["bad"].tool_count == 0 + assert ( + factory.made["good"].state is McpConnectionState.CONNECTED + ), "the server that did connect still published its tools" + + @pytest.mark.asyncio async def test_stop_disconnects_every_server_when_one_teardown_raises(): """One failing teardown must not strand the other connections open.""" @@ -386,6 +611,136 @@ async def test_stop_disconnects_every_server_when_one_teardown_raises(): assert "RuntimeError" in manager.status()["bad"].error +@pytest.mark.asyncio +async def test_stop_completes_teardown_when_an_unregister_raises(monkeypatch): + """A registry that refuses to give a tool back must not strand a connection. + + Unregistration runs before the disconnects, so an unguarded failure here + would skip the teardown entirely and leave every server connected. + """ + tool_manager = ToolManager() + real_unregister = tool_manager.unregister + armed = {"bad": False} + + def _unregister(name): + if armed["bad"] and name == "mcp__bad__x": + raise RuntimeError("registry exploded") + return real_unregister(name) + + monkeypatch.setattr(tool_manager, "unregister", _unregister) + factory = _factory({"bad": (True, ("x",)), "good": (True, ("y",))}) + manager = McpManager( + _config( + ("bad", {"url": "https://bad.invalid/mcp"}), + ("good", {"url": "https://good.invalid/mcp"}), + ), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + armed["bad"] = True + + await manager.stop() + + assert ( + factory.made["good"].state is McpConnectionState.DISCONNECTED + ), "the healthy server was torn down" + assert factory.made["bad"].state is McpConnectionState.DISCONNECTED + assert set(tool_manager.tools_map()) == { + "mcp__bad__x" + }, "only the tool the registry refused to remove is left" + assert "registry exploded" in manager.status()["bad"].error + + +@pytest.mark.asyncio +async def test_stop_retries_an_unregister_that_previously_failed(monkeypatch): + """A failed removal stays owned, so a later stop() can finish the job. + + The manager pops its ownership record only once a removal is confirmed; + forgetting a tool it never removed would leave it dispatchable forever. + """ + tool_manager = ToolManager() + real_unregister = tool_manager.unregister + armed = {"bad": False} + + def _unregister(name): + if armed["bad"] and name == "mcp__bad__x": + raise RuntimeError("registry exploded") + return real_unregister(name) + + monkeypatch.setattr(tool_manager, "unregister", _unregister) + factory = _factory({"bad": (True, ("x",))}) + manager = McpManager( + _config(("bad", {"url": "https://bad.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + armed["bad"] = True + await manager.stop() + assert manager.status()["bad"].tool_count == 1, "still owned, so still counted" + + armed["bad"] = False + await manager.stop() + + assert tool_manager.tools_map() == {} + assert manager.status()["bad"].tool_count == 0 + assert ( + manager.status()["bad"].error == "" + ), "a failure the retry resolved is history, not a current error" + + +@pytest.mark.asyncio +async def test_stop_forgets_a_tool_that_was_already_removed(monkeypatch): + """``unregister`` reporting "nothing was there" ends the manager's claim. + + A tool someone else removed from the registry cannot be removed again, so + keeping it on the ownership record would make every later stop() retry a + removal that can never succeed. + """ + tool_manager = ToolManager() + factory = _factory({"weather": (True, ("a",))}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + assert ( + tool_manager.unregister("mcp__weather__a") is True + ), "removed behind the manager's back" + + await manager.stop() + + assert manager.status()["weather"].tool_count == 0 + assert tool_manager.tools_map() == {} + + +@pytest.mark.asyncio +async def test_status_prefers_a_recorded_failure_over_a_stale_client_message(): + """The newest reason wins when the manager and the client both have one. + + The client's message describes the connect that failed earlier; the + recorded failure describes the teardown that just failed, which is the one + an operator can act on. + """ + tool_manager = ToolManager() + factory = _rude_factory({"weather": "disconnect"}, succeed={"weather": False}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + assert manager.status()["weather"].error == "connection refused" + + await manager.stop() + + assert manager.status()["weather"].error == ( + "RuntimeError: teardown exploded" + ), "the teardown failure is the most recent one" + + @pytest.mark.asyncio async def test_start_twice_then_stop_leaves_no_registrations(): """A repeated start must not orphan the tools the first one registered. @@ -557,6 +912,74 @@ def test_a_list_of_configs_is_accepted_as_the_configuration(): assert manager.status()["weather"].tool_count == 0 +def test_a_repeated_server_name_in_a_list_is_rejected(): + """Two entries under one name would make status() and the registry disagree. + + The name is both the client key and the ``mcp____`` tool namespace, + so a duplicate leaves a tool registered, dispatchable and reported as + ``disabled`` with ``tool_count=0``. Refusing the configuration is the only + way to keep those three facts in agreement. + """ + factory = _factory({"dup": (True, ("a",))}) + config = [ + McpServerConfig.parse("dup", {"url": "https://one.invalid/mcp"}), + McpServerConfig.parse("dup", {"url": "https://two.invalid/mcp"}), + ] + + with pytest.raises(MitoCodeMcpConfigError) as excinfo: + McpManager(config, tool_manager=ToolManager(), client_factory=factory) + + assert "dup" in str(excinfo.value), "the duplicated name is named" + assert "1" in str(excinfo.value), "and so is the offending position" + assert factory.made == {}, "a rejected configuration connects to nothing" + + +def test_a_repeated_server_name_is_rejected_even_when_only_one_is_enabled(): + """Enabled-then-disabled is the pair that produced the worst mismatch.""" + factory = _factory({"dup": (True, ("a",))}) + config = [ + McpServerConfig.parse("dup", {"url": "https://one.invalid/mcp"}), + McpServerConfig.parse( + "dup", {"url": "https://two.invalid/mcp", "enabled": False} + ), + ] + + with pytest.raises(MitoCodeMcpConfigError) as excinfo: + McpManager(config, tool_manager=ToolManager(), client_factory=factory) + + assert "dup" in str(excinfo.value) + assert factory.made == {} + + +@pytest.mark.asyncio +async def test_two_servers_offering_the_same_tool_name_both_register(): + """A duplicate *tool* name across servers is legal and must keep working. + + Only duplicate *server* names collide; the ``mcp____`` prefix keeps + two servers' identically-named tools apart. + """ + tool_manager = ToolManager() + factory = _factory({"alpha": (True, ("search",)), "beta": (True, ("search",))}) + manager = McpManager( + _config( + ("alpha", {"url": "https://alpha.invalid/mcp"}), + ("beta", {"url": "https://beta.invalid/mcp"}), + ), + tool_manager=tool_manager, + client_factory=factory, + ) + + await manager.start() + + assert set(tool_manager.tools_map()) == { + "mcp__alpha__search", + "mcp__beta__search", + } + assert manager.status()["alpha"].tool_count == 1 + assert manager.status()["beta"].tool_count == 1 + assert tool_manager.agent_docs()["mcp__alpha__search"] == "[alpha] search on alpha" + + def test_a_configuration_that_is_neither_a_document_nor_a_list_is_rejected(): with pytest.raises(MitoCodeMcpConfigError): McpManager( @@ -642,6 +1065,7 @@ async def test_status_before_start_reports_a_disconnected_server(): def test_mcp_package_exports_the_manager_and_its_neighbours(): + """Every name in ``__all__`` resolves; a stale entry is a failed import.""" from mitocode.backend import mcp assert mcp.McpManager is McpManager @@ -649,6 +1073,29 @@ def test_mcp_package_exports_the_manager_and_its_neighbours(): assert mcp.McpClient is McpClient assert mcp.McpConnectionState is McpConnectionState assert mcp.McpServerConfig is McpServerConfig + assert mcp.McpServerKind is McpServerKind + assert set(mcp.__all__) == { + "McpClient", + "McpConnectionState", + "McpManager", + "McpServerConfig", + "McpServerKind", + "McpServerStatus", + }, "the package advertises exactly the six documented names" + + +def test_mcp_package_imports_without_the_tool_layer(tmp_path): + """The ``mcp`` package must be importable with no tool layer at all. + + The plan fixes the dependency direction: ``mcp`` never imports + ``tool_manager``. A reference kept under ``TYPE_CHECKING`` is invisible to + this probe, which is why the block covers the whole ``mitocode.backend.tool`` + package rather than one name. + """ + result = _run_import_probe(_IMPORT_WITHOUT_TOOL_LAYER) + + assert result.returncode == 0, result.stderr + assert result.stdout.strip().splitlines()[-1] == "ok" @pytest.mark.asyncio From c26e0b5304953b6935ce660ce42bcfa59eb49211 Mon Sep 17 00:00:00 2001 From: Allen Wang Date: Sat, 17 Oct 2026 14:56:30 +0800 Subject: [PATCH 19/20] fix(mcp): keep owning a tool a re-registration could not release MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _connect_one assigned register_mcp_client's result straight over the ownership record. That list omits any tool whose removal had just failed, because the tool is still in the registry and register_mcp_client skips a taken name — so a second start, or a refresh, silently forgot a dispatchable tool: counted nowhere, and reclaimable by nobody. A third stop() could not get it back either. _unregister_one now returns the names it could not release and _connect_one merges them, keeping the recorded error too, so status() and the registry agree again. stop() also propagated a child's CancelledError as soon as its first phase reported one, skipping the disconnections that release the subprocesses — the strand-everything outcome the per-server isolation exists to prevent, reached through the propagation path. Both phases now run and the cancellation is re-raised afterwards. _test hygiene: the import probe no longer takes an unused fixture, and its child has a timeout so a wedged interpreter fails the test instead of the suite. Constraint: start() must still never raise for a server that fails; only a cancellation propagates Confidence: high Scope-risk: moderate Co-Authored-By: Claude Opus 4.8 (1M context) --- src/mitocode/backend/mcp/mcp_manager.py | 64 ++++++--- test/test_mcp_manager.py | 172 +++++++++++++++++++++++- 2 files changed, 215 insertions(+), 21 deletions(-) diff --git a/src/mitocode/backend/mcp/mcp_manager.py b/src/mitocode/backend/mcp/mcp_manager.py index 1712a32..6080360 100644 --- a/src/mitocode/backend/mcp/mcp_manager.py +++ b/src/mitocode/backend/mcp/mcp_manager.py @@ -239,17 +239,29 @@ async def stop(self) -> None: leave a registered tool pointing at a dead client, which fails at call time in the model's hands rather than here. - Both phases are isolated per server. A registry that refuses to give a - tool back must not skip the teardown that releases the subprocesses, - and a connection that refuses to close must not strand the others; - either is recorded against its server in :meth:`status` instead. + Both phases are isolated per server, and a cancellation in the first + does not skip the second: the disconnection phase releases the + subprocesses, and leaving every one of them running because an + unregistration was cancelled is the failure mode this method exists to + prevent. A registry that refuses to give a tool back must not skip that + teardown either, and a connection that refuses to close must not strand + the others; either is recorded against its server in :meth:`status` + instead. Raises: asyncio.CancelledError: If one server's unregistration or - disconnection is cancelled. + disconnection is cancelled, after both phases have run. """ - await self._run_each(list(self._clients), self._unregister_one) - await self._run_each(list(self._clients), self._disconnect_one) + cancelled: Optional[asyncio.CancelledError] = None + for phase in (self._unregister_one, self._disconnect_one): + try: + await self._run_each(list(self._clients), phase) + except asyncio.CancelledError as exc: + # Remember the first one and carry on to the next phase; the + # caller is told afterwards, not in place of the teardown. + cancelled = cancelled or exc + if cancelled is not None: + raise cancelled async def refresh(self, server_name: Optional[str] = None) -> None: """Reconnect servers so a changed tool catalogue is picked up. @@ -304,7 +316,7 @@ def status(self) -> Dict[str, McpServerStatus]: return snapshot async def _run_each( - self, names: List[str], operation: Callable[[str], Awaitable[None]] + self, names: List[str], operation: Callable[[str], Awaitable[Any]] ) -> None: """Run *operation* once per name, concurrently, isolating failures. @@ -323,6 +335,7 @@ async def _run_each( Args: names: Servers to run *operation* for, each visited exactly once. operation: One server's worth of work, e.g. :meth:`_connect_one`. + Whatever it returns is discarded; it manages its own state. Raises: asyncio.CancelledError: If *operation* was cancelled for one server. @@ -364,19 +377,31 @@ async def _connect_one(self, name: str) -> None: # Whatever this server registered last time is released first, so a # repeated start or refresh cannot leave the previous names behind as # tools nobody owns and stop() cannot unregister. - await self._unregister_one(name) + stuck = await self._unregister_one(name) try: specs = self._tool_manager.register_mcp_client(client) except Exception as exc: # pylint: disable=broad-except self._record_failure(name, exc, "failed to register its tools") return - self._registered[name] = [spec.name for spec in specs] - self._errors.pop(name, None) + # ``stuck`` is non-empty only when a removal above failed. Those tools + # are still in the registry, so register_mcp_client skipped their names + # and they are absent from ``specs``: recording ``specs`` alone would + # forget the one record a later stop() needs to reclaim them, leaving a + # dispatchable tool nobody owns. The error stays recorded for the same + # reason. + self._registered[name] = [spec.name for spec in specs] + stuck + if not stuck: + self._errors.pop(name, None) async def _reconnect_one(self, name: str) -> None: """Unregister, close and reconnect one server. + The unregistration here is what makes a *failed* reconnect publish + nothing, since :meth:`_connect_one` only registers after a successful + connect. Its leftover names need no handling: they stay in the + ownership record, and :meth:`_connect_one` retries them. + Args: name: Name of the server to reconnect, as keyed in ``_clients``. """ @@ -399,7 +424,7 @@ async def _disconnect_one(self, name: str) -> None: except Exception as exc: # pylint: disable=broad-except self._record_failure(name, exc, "failed to disconnect") - async def _unregister_one(self, name: str) -> None: + async def _unregister_one(self, name: str) -> List[str]: """Remove every tool one server registered from the registry. The ownership record is updated per tool and only once the registry has @@ -419,6 +444,11 @@ async def _unregister_one(self, name: str) -> None: Args: name: Name of the server whose tools are removed. + + Returns: + The tools that are still in the registry because their removal + failed, so the caller can keep owning them. Empty when everything + this server held was released. """ remaining: List[str] = [] for tool_name in self._registered.get(name, []): @@ -437,12 +467,13 @@ async def _unregister_one(self, name: str) -> None: ) if remaining: self._registered[name] = remaining - return + return remaining self._registered.pop(name, None) # Nothing failed and nothing is left, so whatever this server last # complained about is history. Leaving it in place would have status() # report a failure that no longer describes anything. self._errors.pop(name, None) + return [] def _record_failure(self, name: str, exc: BaseException, problem: str) -> None: """Remember a failure the client did not report itself, and log it. @@ -451,8 +482,11 @@ def _record_failure(self, name: str, exc: BaseException, problem: str) -> None: name: Name of the MCP server that failed. exc: The exception that ended the operation. problem: A complete phrase saying what went wrong, e.g. - ``"failed to connect"`` — it is logged verbatim and may name a - tool, which is never a credential. + ``"failed to connect"``. It is logged verbatim and may name a + tool. Such a name is arbitrary text from the server, so it is + not guaranteed to look innocuous — but it is already published + through :meth:`status` and the tool schemas, so naming it here + adds no exposure of its own. """ self._errors[name] = _describe(exc) logger.error(f"MCP server '{name}': {problem}: {self._errors[name]}") diff --git a/test/test_mcp_manager.py b/test/test_mcp_manager.py index cb79065..487b807 100644 --- a/test/test_mcp_manager.py +++ b/test/test_mcp_manager.py @@ -119,6 +119,7 @@ def __init__(self, config, succeed=True, tools=(), keep_tools_on_failure=False): self._config = config self._succeed = succeed self.keep_tools_on_failure = keep_tools_on_failure + self._cancel_next = False self.server_name = config.name self.tools = [ SimpleNamespace( @@ -138,9 +139,20 @@ def __init__(self, config, succeed=True, tools=(), keep_tools_on_failure=False): async def connect(self): """Move to CONNECTED or FAILED according to the scripted outcome. + One-shot cancellation is honoured first, so a test can arm it *after* + a successful ``start`` rather than making every later connect fail. + Returns: The scripted outcome. + + Raises: + asyncio.CancelledError: If :meth:`cancel_next_connect` armed this + attempt, which is how ``McpClient.connect`` behaves when the + caller cancels it. """ + if self._cancel_next: + self._cancel_next = False + raise asyncio.CancelledError() self.connected = self._succeed if self._succeed: self.state = McpConnectionState.CONNECTED @@ -162,6 +174,10 @@ def go_offline(self): """Make every later connect attempt fail, as a server going away would.""" self._succeed = False + def cancel_next_connect(self): + """Cancel only the next connect, leaving any later one to succeed.""" + self._cancel_next = True + async def call_tool(self, tool_name, arguments): """Record the call and render it, as the real client would.""" self.calls += 1 @@ -182,17 +198,15 @@ def __init__(self, config, succeed=True, fail_on=""): Args: config: The ``McpServerConfig`` this stand-in represents. succeed: Value :meth:`connect` reports back when it does not raise. - fail_on: ``"connect"``, ``"disconnect"``, ``"cancel"`` or ``""``. + fail_on: ``"connect"``, ``"disconnect"`` or ``""``. """ super().__init__(config, succeed=succeed) self._fail_on = fail_on async def connect(self): - """Raise, or cancel, instead of connecting when this client is armed.""" + """Raise instead of connecting when this client is armed for it.""" if self._fail_on == "connect": raise RuntimeError("handshake exploded") - if self._fail_on == "cancel": - raise asyncio.CancelledError() return await super().connect() async def disconnect(self): @@ -288,6 +302,8 @@ def _run_import_probe(script): env=env, cwd=env["MITOCODE_SRC"], check=False, + # A wedged child must fail the test rather than hang the whole suite. + timeout=120, ) @@ -569,7 +585,7 @@ async def test_start_records_a_cancelled_server_and_lets_the_cancellation_out(): so without this the server would report a failure with no reason at all. """ tool_manager = ToolManager() - factory = _rude_factory({"bad": "cancel", "good": ""}) + factory = _rude_factory({"bad": "", "good": ""}) manager = McpManager( _config( ("bad", {"url": "https://bad.invalid/mcp"}), @@ -578,6 +594,7 @@ async def test_start_records_a_cancelled_server_and_lets_the_cancellation_out(): tool_manager=tool_manager, client_factory=factory, ) + factory.made["bad"].cancel_next_connect() with pytest.raises(asyncio.CancelledError): await manager.start() @@ -589,6 +606,149 @@ async def test_start_records_a_cancelled_server_and_lets_the_cancellation_out(): ), "the server that did connect still published its tools" +@pytest.mark.asyncio +async def test_start_keeps_owning_a_tool_it_could_not_unregister(monkeypatch): + """A second ``start`` must not forget a tool the registry refused to remove. + + ``register_mcp_client`` skips a name that is still registered, so its + result omits the stuck tool. Assigning that result to the ownership record + would forget the tool while it is still in the registry — dispatchable, + counted nowhere, and reclaimable by nobody. + """ + tool_manager = ToolManager() + real_unregister = tool_manager.unregister + armed = {"bad": False} + + def _unregister(name): + if armed["bad"] and name == "mcp__weather__x": + raise RuntimeError("registry exploded") + return real_unregister(name) + + monkeypatch.setattr(tool_manager, "unregister", _unregister) + factory = _factory({"weather": (True, ("x",))}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + armed["bad"] = True + + await manager.start() + + assert ( + manager.status()["weather"].tool_count == 1 + ), "the tool is still in the registry, so status() must still count it" + assert "registry exploded" in manager.status()["weather"].error + + armed["bad"] = False + await manager.stop() + + assert ( + tool_manager.tools_map() == {} + ), "a later stop() must reclaim the tool nobody forgot" + + +@pytest.mark.asyncio +async def test_start_owns_the_tools_a_refresh_could_not_unregister(monkeypatch): + """The same guarantee holds on ``refresh``, the other path through it.""" + tool_manager = ToolManager() + real_unregister = tool_manager.unregister + armed = {"bad": False} + + def _unregister(name): + if armed["bad"] and name == "mcp__weather__a": + raise RuntimeError("registry exploded") + return real_unregister(name) + + monkeypatch.setattr(tool_manager, "unregister", _unregister) + factory = _factory({"weather": (True, ("a", "b"))}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + armed["bad"] = True + + await manager.refresh("weather") + + assert ( + manager.status()["weather"].tool_count == 2 + ), "'a' is stuck in the registry and 'b' was registered, so both are owned" + + armed["bad"] = False + await manager.stop() + + assert tool_manager.tools_map() == {} + + +@pytest.mark.asyncio +async def test_stop_completes_both_phases_and_then_propagates_a_cancellation( + monkeypatch, +): + """A cancelled unregistration must not skip the disconnections. + + The two phases run one after the other, so re-raising as soon as the first + one is cancelled would leave every server connected — the outcome the + per-server isolation exists to prevent, reached through the propagation path. + """ + tool_manager = ToolManager() + factory = _factory({"bad": (True, ("x",)), "good": (True, ("y",))}) + manager = McpManager( + _config( + ("bad", {"url": "https://bad.invalid/mcp"}), + ("good", {"url": "https://good.invalid/mcp"}), + ), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + + real_unregister = tool_manager.unregister + + def _unregister(name): + if name == "mcp__bad__x": + raise asyncio.CancelledError() + return real_unregister(name) + + monkeypatch.setattr(tool_manager, "unregister", _unregister) + + with pytest.raises(asyncio.CancelledError): + await manager.stop() + + assert factory.made["good"].state is McpConnectionState.DISCONNECTED + assert ( + factory.made["bad"].state is McpConnectionState.DISCONNECTED + ), "every connection was closed before the cancellation was re-raised" + assert manager.status()["bad"].error == "CancelledError" + + +@pytest.mark.asyncio +async def test_refresh_records_a_cancellation_and_lets_it_out(): + """``refresh`` propagates a cancellation for the same reason ``start`` does.""" + tool_manager = ToolManager() + factory = _rude_factory({"bad": "", "good": ""}) + manager = McpManager( + _config( + ("bad", {"url": "https://bad.invalid/mcp"}), + ("good", {"url": "https://good.invalid/mcp"}), + ), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + factory.made["bad"].cancel_next_connect() + + with pytest.raises(asyncio.CancelledError): + await manager.refresh("bad") + + assert manager.status()["bad"].error == "CancelledError" + assert ( + factory.made["good"].state is McpConnectionState.CONNECTED + ), "refreshing one server left the others alone" + + @pytest.mark.asyncio async def test_stop_disconnects_every_server_when_one_teardown_raises(): """One failing teardown must not strand the other connections open.""" @@ -1084,7 +1244,7 @@ def test_mcp_package_exports_the_manager_and_its_neighbours(): }, "the package advertises exactly the six documented names" -def test_mcp_package_imports_without_the_tool_layer(tmp_path): +def test_mcp_package_imports_without_the_tool_layer(): """The ``mcp`` package must be importable with no tool layer at all. The plan fixes the dependency direction: ``mcp`` never imports From b4dbaf41efbc316d471689bcce3ea3b8e6e784da Mon Sep 17 00:00:00 2001 From: Allen Wang Date: Sat, 17 Oct 2026 18:52:13 +0800 Subject: [PATCH 20/20] chore(mcp): remove superseded experiment scripts --- experiments/mcp_agent.py | 263 ------------------------ experiments/mcp_multi_thread.py | 354 -------------------------------- experiments/multiagent.py | 27 --- 3 files changed, 644 deletions(-) delete mode 100644 experiments/mcp_agent.py delete mode 100644 experiments/mcp_multi_thread.py delete mode 100644 experiments/multiagent.py diff --git a/experiments/mcp_agent.py b/experiments/mcp_agent.py deleted file mode 100644 index c93ec96..0000000 --- a/experiments/mcp_agent.py +++ /dev/null @@ -1,263 +0,0 @@ -import subprocess -import sys -import os -import asyncio -import json -from contextlib import AsyncExitStack -from mcp.client.session import ClientSession -from mcp.client.stdio import stdio_client, StdioServerParameters -from openai import AsyncOpenAI - -# Configuration -OPENAI_API_KEY = "" - -# 增加超时时间（从30秒改为120秒） -COMMAND_TIMEOUT = 120 - - -async def ensure_dependencies(): - print("Checking and installing dependencies...") - try: - # Check Node.js and npm - subprocess.run( - ["npm", "--version"], check=True, capture_output=True, timeout=10 - ) - except ( - subprocess.CalledProcessError, - FileNotFoundError, - subprocess.TimeoutExpired, - ): - pass - - try: - # Check uv - subprocess.run(["uv", "--version"], check=True, capture_output=True, timeout=10) - except ( - subprocess.CalledProcessError, - FileNotFoundError, - subprocess.TimeoutExpired, - ): - print("Installing dependencies...") - subprocess.run( - [sys.executable, "-m", "pip", "install", "uv", "mcp", "openai"], - check=True, - timeout=300, - ) - - -async def start_mcp_servers(): - # 设置环境变量，解决编码和 Git 分页器问题 - env = os.environ.copy() - env["PYTHONIOENCODING"] = "utf-8" - env["GIT_PAGER"] = "cat" # 禁用 Git 分页器，防止等待用户输入 - env["PAGER"] = "cat" - - server_params_list = [ - # Git server parameters - StdioServerParameters( - command=sys.executable, - args=["-m", "uv", "tool", "run", "-q", "mcp-server-git", "-r", os.getcwd()], - env=env, - ), - # Terminal controller MCP server parameters - StdioServerParameters( - command=sys.executable, - args=[ - "-m", - "uv", - "tool", - "run", - "-q", - "--from", - "terminal-controller", - "terminal_controller", - ], - env=env, - ), - ] - - return server_params_list - - -async def main(): - if not OPENAI_API_KEY: - print( - "WARNING: OPENAI_API_KEY environment variable is not set. The OpenAI API call will fail unless you provide one." - ) - - await ensure_dependencies() - - print("\nStarting MCP servers...") - server_params_list = await start_mcp_servers() - - sessions = [] - transport_managers = [] - - # We use AsyncExitStack to manage multiple server connections easily - async with AsyncExitStack() as stack: - for params in server_params_list: - print(f"Connecting to server: {params.command} {' '.join(params.args)}") - transport_manager = stdio_client(params) - read_stream, write_stream = await stack.enter_async_context( - transport_manager - ) - session = await stack.enter_async_context( - ClientSession(read_stream, write_stream) - ) - await session.initialize() - sessions.append(session) - transport_managers.append(transport_manager) - print( - f"Connected to {params.command} {params.args[1] if len(params.args) > 1 else ''} MCP Server." - ) - - # Aggregate all tools from all sessions - all_tools = [] - tool_to_session_map = {} - for session in sessions: - response = await session.list_tools() - for tool in response.tools: - all_tools.append(tool) - tool_to_session_map[tool.name] = session - print(f"Loaded tool: {tool.name}") - - # Prepare tools for OpenAI mapping - openai_tools = [] - for t in all_tools: - openai_tools.append( - { - "type": "function", - "function": { - "name": t.name, - "description": t.description, - "parameters": t.inputSchema, - }, - } - ) - - print("\nAll MCP servers initialized and tools loaded.") - print(f"Total tools available to AI: {len(openai_tools)}") # 修复 f-string - - if not OPENAI_API_KEY: - return - - # Example interactions with AI - client = AsyncOpenAI( - api_key=OPENAI_API_KEY, - base_url="https://app.nextchat.club/api/siliconflow/v1", # 修复 URL 多余空格 - ) - messages = [{"role": "user", "content": "run winver"}] - - print("\nSending prompt to OpenAI:") - print(messages[0]["content"]) - - try: - while True: - print("\nWaiting for OpenAI model response...") - response = await client.chat.completions.create( - model="Qwen/Qwen3-Coder-30B-A3B-Instruct", - messages=messages, - tools=openai_tools, - tool_choice="auto", - timeout=COMMAND_TIMEOUT, # 添加 API 请求超时 - ) - - response_message = response.choices[0].message - messages.append(response_message) - - if response_message.tool_calls: - for tool_call in response_message.tool_calls: - tool_name = tool_call.function.name - tool_args = json.loads(tool_call.function.arguments) - print( - f"\nModel wants to call tool '{tool_name}' with args {tool_args}" - ) - - if tool_name in tool_to_session_map: - session = tool_to_session_map[tool_name] - try: - # 添加超时控制 - result = await asyncio.wait_for( - session.call_tool(tool_name, arguments=tool_args), - timeout=COMMAND_TIMEOUT, - ) - tool_result_content = "" - if hasattr(result, "content"): - for content_item in result.content: - if hasattr(content_item, "text"): - tool_result_content += ( - content_item.text + "\n" - ) - else: - tool_result_content = str(result) - - print( - f"Tool execution result: {tool_result_content[:500]}..." - ) # 增加显示长度 - messages.append( - { - "role": "tool", - "tool_call_id": tool_call.id, - "name": tool_name, - "content": tool_result_content, - } - ) - except asyncio.TimeoutError: - print( - f"Error: Tool '{tool_name}' execution timed out after {COMMAND_TIMEOUT} seconds" - ) - messages.append( - { - "role": "tool", - "tool_call_id": tool_call.id, - "name": tool_name, - "content": f"Error: Command timed out after {COMMAND_TIMEOUT} seconds. Try breaking the command into smaller steps.", - } - ) - except Exception as e: - print(f"Error calling tool {tool_name}: {e}") - messages.append( - { - "role": "tool", - "tool_call_id": tool_call.id, - "name": tool_name, - "content": f"Error: {str(e)}", - } - ) - else: - print(f"Unknown tool requested: {tool_name}") - messages.append( - { - "role": "tool", - "tool_call_id": tool_call.id, - "name": tool_name, - "content": f"Error: unknown tool '{tool_name}'", - } - ) - else: - # No more tool calls, AI is providing final text response - print("\nFinal response from AI:") - print(response_message.content) - break - finally: - # 确保所有会话正确关闭 - print("\nCleaning up sessions...") - for session in sessions: - try: - await session.close() - except Exception as e: - print(f"Warning: Error closing session: {e}") - - -if __name__ == "__main__": - # Windows 上设置事件循环策略，避免子进程问题 - if sys.platform == "win32": - asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy()) - - try: - asyncio.run(main()) - except KeyboardInterrupt: - print("\nInterrupted by user") - except Exception as e: - print(f"\nFatal error: {e}") - raise diff --git a/experiments/mcp_multi_thread.py b/experiments/mcp_multi_thread.py deleted file mode 100644 index 213299f..0000000 --- a/experiments/mcp_multi_thread.py +++ /dev/null @@ -1,354 +0,0 @@ -import subprocess -import sys -import os -import asyncio -import json -import threading -from concurrent.futures import Future, ThreadPoolExecutor -from contextlib import AsyncExitStack -from mcp.client.session import ClientSession -from mcp.client.stdio import stdio_client, StdioServerParameters -from openai import OpenAI - -# Configuration -OPENAI_API_KEY = "" - -# 增加超时时间（从30秒改为120秒） -COMMAND_TIMEOUT = 120 - - -def ensure_dependencies(): - print("Checking and installing dependencies...") - try: - # Check Node.js and npm - subprocess.run( - ["npm", "--version"], check=True, capture_output=True, timeout=10 - ) - except ( - subprocess.CalledProcessError, - FileNotFoundError, - subprocess.TimeoutExpired, - ): - pass - - try: - # Check uv - subprocess.run(["uv", "--version"], check=True, capture_output=True, timeout=10) - except ( - subprocess.CalledProcessError, - FileNotFoundError, - subprocess.TimeoutExpired, - ): - print("Installing dependencies...") - subprocess.run( - [sys.executable, "-m", "pip", "install", "uv", "mcp", "openai"], - check=True, - timeout=300, - ) - - -def build_server_params(): - # 设置环境变量，解决编码和 Git 分页器问题 - env = os.environ.copy() - env["PYTHONIOENCODING"] = "utf-8" - env["GIT_PAGER"] = "cat" # 禁用 Git 分页器，防止等待用户输入 - env["PAGER"] = "cat" - - server_params_list = [ - # Git server parameters - StdioServerParameters( - command=sys.executable, - args=["-m", "uv", "tool", "run", "-q", "mcp-server-git", "-r", os.getcwd()], - env=env, - ), - # Terminal controller MCP server parameters - StdioServerParameters( - command=sys.executable, - args=[ - "-m", - "uv", - "tool", - "run", - "-q", - "--from", - "terminal-controller", - "terminal_controller", - ], - env=env, - ), - ] - - return server_params_list - - -class MCPSessionManager: - """Manages MCP server connections in a dedicated background thread. - - Since the MCP library uses async I/O internally, a private asyncio event - loop is kept alive on a daemon thread. All async operations are submitted - to that loop via asyncio.run_coroutine_threadsafe, which returns a - concurrent.futures.Future that the calling thread can wait on. This lets - the rest of the program stay fully synchronous / thread-based. - """ - - def __init__(self, server_params_list: list[StdioServerParameters]): - self.server_params_list = server_params_list - self._loop = asyncio.new_event_loop() - self._loop_thread = threading.Thread( - target=self._run_loop, daemon=True, name="mcp-event-loop" - ) - self._stack = AsyncExitStack() - self.sessions: list[ClientSession] = [] - self.all_tools: list = [] - self.tool_to_session_map: dict = {} - - # ------------------------------------------------------------------ - # Internal helpers - # ------------------------------------------------------------------ - - def _run_loop(self): - asyncio.set_event_loop(self._loop) - self._loop.run_forever() - - def _run_coro(self, coro, timeout=None): - """Submit a coroutine to the background loop and block until done.""" - future: Future = asyncio.run_coroutine_threadsafe(coro, self._loop) - return future.result(timeout=timeout) - - async def _connect_all(self): - for params in self.server_params_list: - print(f"Connecting to server: {params.command} {' '.join(params.args)}") - transport_manager = stdio_client(params) - read_stream, write_stream = await self._stack.enter_async_context( - transport_manager - ) - session = await self._stack.enter_async_context( - ClientSession(read_stream, write_stream) - ) - await session.initialize() - self.sessions.append(session) - print( - f"Connected to {params.command} {params.args[1] if len(params.args) > 1 else ''} MCP Server." - ) - - for session in self.sessions: - response = await session.list_tools() - for tool in response.tools: - self.all_tools.append(tool) - self.tool_to_session_map[tool.name] = session - print(f"Loaded tool: {tool.name}") - - async def _call_tool_async(self, name: str, args: dict): - session = self.tool_to_session_map[name] - return await session.call_tool(name, arguments=args) - - async def _cleanup_async(self): - await self._stack.aclose() - - # ------------------------------------------------------------------ - # Public interface (synchronous) - # ------------------------------------------------------------------ - - def start(self): - """Start the background event loop and connect to all MCP servers.""" - self._loop_thread.start() - self._run_coro(self._connect_all()) - - def call_tool(self, name: str, args: dict, timeout=COMMAND_TIMEOUT): - """Call an MCP tool synchronously, with an optional timeout (seconds).""" - return self._run_coro(self._call_tool_async(name, args), timeout=timeout) - - def cleanup(self): - """Disconnect from all servers and stop the background event loop.""" - if self._loop.is_running(): - self._run_coro(self._cleanup_async()) - self._loop.call_soon_threadsafe(self._loop.stop) - self._loop_thread.join() - - -def call_tool_in_thread( - manager: MCPSessionManager, - tool_name: str, - tool_args: dict, - result_holder: list, - error_holder: list, -): - """Worker function executed in a ThreadPoolExecutor thread to call an MCP tool.""" - try: - result = manager.call_tool(tool_name, tool_args, timeout=COMMAND_TIMEOUT) - result_holder.append(result) - except TimeoutError as exc: - error_holder.append(exc) - except Exception as exc: - error_holder.append(exc) - - -def main(): - if not OPENAI_API_KEY: - print( - "WARNING: OPENAI_API_KEY environment variable is not set. The OpenAI API call will fail unless you provide one." - ) - - ensure_dependencies() - - print("\nStarting MCP servers...") - server_params_list = build_server_params() - - manager = MCPSessionManager(server_params_list) - manager.start() - - try: - # Prepare tools for OpenAI - openai_tools = [] - for t in manager.all_tools: - openai_tools.append( - { - "type": "function", - "function": { - "name": t.name, - "description": t.description, - "parameters": t.inputSchema, - }, - } - ) - - print("\nAll MCP servers initialized and tools loaded.") - print(f"Total tools available to AI: {len(openai_tools)}") - - if not OPENAI_API_KEY: - return - - # Synchronous OpenAI client — no async needed - client = OpenAI( - api_key=OPENAI_API_KEY, - base_url="https://app.nextchat.club/api/siliconflow/v1", - ) - messages = [{"role": "user", "content": "run winver"}] - - print("\nSending prompt to OpenAI:") - print(messages[0]["content"]) - - # Use a thread pool for parallel tool calls when the model requests multiple at once - with ThreadPoolExecutor( - max_workers=4, thread_name_prefix="tool-worker" - ) as executor: - while True: - print("\nWaiting for OpenAI model response...") - response = client.chat.completions.create( - model="Qwen/Qwen3-Coder-30B-A3B-Instruct", - messages=messages, - tools=openai_tools, - tool_choice="auto", - timeout=COMMAND_TIMEOUT, # 添加 API 请求超时 - ) - - response_message = response.choices[0].message - messages.append(response_message) - - if response_message.tool_calls: - # Dispatch all requested tool calls concurrently via the thread pool - futures = {} - result_holders = {} - error_holders = {} - for tool_call in response_message.tool_calls: - tool_name = tool_call.function.name - tool_args = json.loads(tool_call.function.arguments) - print( - f"\nModel wants to call tool '{tool_name}' with args {tool_args}" - ) - - if tool_name in manager.tool_to_session_map: - result_holders[tool_call.id] = [] - error_holders[tool_call.id] = [] - futures[tool_call.id] = executor.submit( - call_tool_in_thread, - manager, - tool_name, - tool_args, - result_holders[tool_call.id], - error_holders[tool_call.id], - ) - else: - print(f"Unknown tool requested: {tool_name}") - messages.append( - { - "role": "tool", - "tool_call_id": tool_call.id, - "name": tool_name, - "content": f"Error: unknown tool '{tool_name}'", - } - ) - - # Collect results in original order - for tool_call in response_message.tool_calls: - tool_name = tool_call.function.name - if tool_call.id not in futures: - continue - # Wait for the worker to finish. Add a small grace period beyond the - # internal timeout so we never block here indefinitely. - futures[tool_call.id].result(timeout=COMMAND_TIMEOUT + 5) - - if error_holders[tool_call.id]: - exc = error_holders[tool_call.id][0] - if isinstance(exc, TimeoutError): - print( - f"Error: Tool '{tool_name}' execution timed out after {COMMAND_TIMEOUT} seconds" - ) - messages.append( - { - "role": "tool", - "tool_call_id": tool_call.id, - "name": tool_name, - "content": f"Error: Command timed out after {COMMAND_TIMEOUT} seconds. Try breaking the command into smaller steps.", - } - ) - else: - print(f"Error calling tool {tool_name}: {exc}") - messages.append( - { - "role": "tool", - "tool_call_id": tool_call.id, - "name": tool_name, - "content": f"Error: {str(exc)}", - } - ) - else: - result = result_holders[tool_call.id][0] - tool_result_content = "" - if hasattr(result, "content"): - for content_item in result.content: - if hasattr(content_item, "text"): - tool_result_content += content_item.text + "\n" - else: - tool_result_content = str(result) - - print( - f"Tool execution result: {tool_result_content[:500]}..." - ) # show first 500 chars - messages.append( - { - "role": "tool", - "tool_call_id": tool_call.id, - "name": tool_name, - "content": tool_result_content, - } - ) - else: - # No more tool calls, AI is providing final text response - print("\nFinal response from AI:") - print(response_message.content) - break - finally: - # Ensure all sessions are closed properly - print("\nCleaning up sessions...") - manager.cleanup() - - -if __name__ == "__main__": - try: - main() - except KeyboardInterrupt: - print("\nInterrupted by user") - except Exception as e: - print(f"\nFatal error: {e}") - raise diff --git a/experiments/multiagent.py b/experiments/multiagent.py deleted file mode 100644 index f5067aa..0000000 --- a/experiments/multiagent.py +++ /dev/null @@ -1,27 +0,0 @@ -import threading -from concurrent.futures import ThreadPoolExecutor -from backend.agent.agent import Agent -from backend.agent import mcp_manager - -executor = ThreadPoolExecutor(max_workers=4) -agents: dict[int, Agent] = {} -agents_cnt = 0 -lock = threading.Lock() - - -def submit_exist(agent_id: int, message: str): - with lock: - agent = agents[agent_id] - agent.append_str_message(message) - return executor.submit(agent.run) - - -def submit_new( - prompt: list[dict], globalmcpmanager: mcp_manager.MCPManager, **settings -): - with lock: - agent = Agent(prompt, globalmcpmanager, settings=settings) - global agents_cnt - agents[agents_cnt] = agent - agents_cnt += 1 - return executor.submit(agent.run)