diff --git a/.gitignore b/.gitignore index 9c5a8bd..1501364 100644 --- a/.gitignore +++ b/.gitignore @@ -180,3 +180,4 @@ venv.bak/ wheels/ # you could uncomment the following to ignore the entire vscode folder /AGENTS.md +docs/superpowers 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) 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/config.json b/src/.mitocode/config.json new file mode 100644 index 0000000..882ac5e --- /dev/null +++ b/src/.mitocode/config.json @@ -0,0 +1 @@ +{"name": "src", "description": "", "version": 2} \ No newline at end of file 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, diff --git a/src/mitocode/backend/exception/mitocode_exception.py b/src/mitocode/backend/exception/mitocode_exception.py index 086fab9..6764637 100644 --- a/src/mitocode/backend/exception/mitocode_exception.py +++ b/src/mitocode/backend/exception/mitocode_exception.py @@ -74,3 +74,27 @@ 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): + """Reserved for when an MCP connection cannot be established or is lost. + + Nothing raises it yet, and that is deliberate rather than an oversight: + :class:`~mitocode.backend.mcp.mcp_client.McpClient` records a failed + connection in ``state``/``last_error`` and returns ``False``, because one + unreachable server must not abort the startup of the servers that work. + Turning that into an exception belongs to whichever layer turns connection + failure into a user-visible abort. + """ + + +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..0244795 --- /dev/null +++ b/src/mitocode/backend/mcp/__init__.py @@ -0,0 +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_client.py b/src/mitocode/backend/mcp/mcp_client.py new file mode 100644 index 0000000..79063c2 --- /dev/null +++ b/src/mitocode/backend/mcp/mcp_client.py @@ -0,0 +1,464 @@ +"""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 +import re +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 + +#: 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]*)" +) + +#: An endpoint quoted inside free-form error text, in either of the two shapes +#: that can carry a credential. The first branch is a full URL — what httpx, a +#: proxy or fastmcp's own transport ``repr`` renders — and runs to the first +#: character that cannot belong to one: whitespace, the quotes an error wraps a +#: URL in, or the brackets a repr or traceback puts around it. The second is +#: the scheme-less reference a JSON-RPC server echoes back in its own error text +#: (``/mcp?access_token=…``), which is precisely the shape the MCP SDK's own +#: URL scrubbing cannot see because it only ever parses URLs it is handed. That +#: branch requires a ``key=`` in the query so ordinary prose containing a +#: question mark is not mistaken for an endpoint. +_ENDPOINT_RUN = re.compile( + r"[A-Za-z][A-Za-z0-9+.\-]*://[^\s'\"<>)\]}`]+" + r"|(?:/[^\s'\"<>)\]}`]*|[A-Za-z0-9_.~-]+)\?[A-Za-z0-9_.~-]+=[^\s'\"<>)\]}`]*" +) + +#: The three parts of a URL that can carry a credential. Each matches only the +#: first occurrence, and each is applied to a run :data:`_ENDPOINT_RUN` already +#: bounded, so none of them can reach past the endpoint into the sentence. +_ENDPOINT_USERINFO = re.compile(r"://[^/?#]*@") +_ENDPOINT_QUERY = re.compile(r"\?[^#\s]*") +_ENDPOINT_FRAGMENT = re.compile(r"#[^\s]*") + + +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) + + +def _scrub_endpoint_run(match: "re.Match[str]") -> str: + """Rewrite one matched endpoint with its credential-bearing parts removed. + + Args: + match: One :data:`_ENDPOINT_RUN` match, either a full URL or a + scheme-less reference. + + Returns: + The endpoint with its userinfo, query and fragment replaced by + :data:`_REDACTED`. Scheme, host, port and path are kept, and so is + the fact that something was there at all. + """ + run = match.group(0) + if "://" not in run: + # A scheme-less reference: only a query can have been scrubbed, because + # userinfo needs an authority to live in and there is none. + return f"{run.split('?', 1)[0]}?{_REDACTED}" + + scrubbed = _ENDPOINT_USERINFO.sub(f"://{_REDACTED}@", run, count=1) + # The query stops at a '#' so that a fragment stays a fragment rather than + # being swallowed into the query's redaction. + scrubbed = _ENDPOINT_QUERY.sub(f"?{_REDACTED}", scrubbed, count=1) + return _ENDPOINT_FRAGMENT.sub(f"#{_REDACTED}", scrubbed, count=1) + + +def _scrub_endpoint(text: str) -> str: + """Cut credentials out of any endpoint quoted in *text*. + + A remote MCP endpoint may legitimately carry a secret in its userinfo, its + query string or its fragment — ``?access_token=…`` is common enough that + :class:`~mitocode.backend.mcp.mcp_config.McpServerConfig` has to accept it + — and a handshake failure routinely quotes the endpoint it was made + against, in text that reaches the TUI error pane and the log. The MCP SDK + strips the same parts from the URLs *it* builds, but nothing strips them + from text a remote server chose to write. + + The scrub is deliberately URL-shaped rather than blanket: host, path and + the exception's own type survive, so the message still says which endpoint + failed, and a message that merely mentions a host, or nothing at all, is + left byte-identical. + + Args: + text: Error text as reported, possibly quoting an endpoint. + + Returns: + The same text with every endpoint's userinfo, query and fragment + replaced by :data:`_REDACTED`. + """ + return _ENDPOINT_RUN.sub(_scrub_endpoint_run, text) + + +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. + + No deadline is set on the client itself. ``Client(timeout=...)`` becomes + the session's default for *every* request, and the one request that + carries no deadline of its own is the ``list_tools`` discovery inside + :meth:`connect` — so with ``timeout`` below ``connect_timeout`` the + discovery would be bounded by the wrong deadline and report the wrong + one. Both callers apply their own instead: :meth:`connect` wraps the + handshake and the discovery in ``asyncio.wait_for``, and + :meth:`call_tool` passes ``timeout`` to the request itself. + """ + return Client( + self._config.build_transport(), + name=self._config.name, + ) + + @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. + + 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 + + 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 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. The + # shield keeps a *repeat* cancellation — a second Ctrl-C, or an + # enclosing wait_for firing as well — from interrupting it too, for + # as long as the loop outlives it. See disconnect() for why that + # is the limit of shield(). + cleanup_errors = await asyncio.shield(self._force_close()) + self._state = McpConnectionState.FAILED + self._last_error = f"{type(exc).__name__}: {_scrub_endpoint(str(exc))}" + if cleanup_errors: + # Appended, never substituted. The first half is what ended the + # connect; the rest is the teardown saying it could not release + # the transport either, which on the timeout path is the whole + # warning — a subprocess that may still be running. Discarding it + # the way this used to leaves the one diagnostic an operator needs. + self._last_error += "; " + "; ".join(cleanup_errors) + if isinstance(exc, Exception): + return False + raise + + self._state = McpConnectionState.CONNECTED + self._last_error = "" + return True + + async def _force_close(self) -> List[str]: + """Tear the connection down, swallowing but reporting 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. + + The collected errors are *returned* rather than written to + :attr:`last_error` here, because the two callers owe that attribute + different things: :meth:`disconnect` records them, while + :meth:`connect` is about to overwrite it with the failure that triggered + the cleanup and can only append them to that. + + Returns: + One description per handle that could not be released, empty when + the teardown was clean. + """ + 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__}: {_scrub_endpoint(str(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__}: " + f"{_scrub_endpoint(str(exc))}" + ) + return 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 — *unless* + they are not an :class:`Exception`. A :class:`BaseException` that is + neither one nor a cancellation — a framework's shutdown signal — keeps + travelling, because recording it away would silently downgrade the + second Ctrl-C or the enclosing cancellation that caused it, and + :attr:`last_error` stays untouched rather than blaming a cleanup + problem that never happened. + + The teardown is shielded for the same reason :meth:`connect`'s is: a + *repeat* cancellation — a second Ctrl-C, or an enclosing ``wait_for`` + firing in the same tick — must not interrupt ``Client.close()`` midway + and leave stdio's detached connect task holding the child process alive. + That holds as long as the event loop outlives the cancellation, which is + as far as :func:`asyncio.shield` reaches: a loop torn down with the + teardown still in flight cancels it like any other await, and the second + Ctrl-C is often about to bring that about. Only the fact that ``close`` + ran to completion tells a shielded teardown from a plain await; both + leave the handles cleared. + + :attr:`state` is set before that await rather than after it, for the + same reason :meth:`_force_close` drops the handles before its own + awaits: the teardown counts as begun from the instant it starts, so a + caller cancelled part-way through it must not find a client still + reporting ``CONNECTED``. That is not cosmetic — :meth:`connect` returns + early for a connected client, so the state left behind would hand the + next reconnect a client with no transport under it. The state means + "closing", not "released": for the one iteration between the assignment + and ``_force_close``'s first statement the handles are still live, and + nothing observable runs inside that window. + """ + if self._stack is None and self._client is None: + return + self._state = McpConnectionState.DISCONNECTED + errors = await asyncio.shield(self._force_close()) + if errors: + self._last_error = "; ".join(errors) + + 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. + + Whatever the server said about a failure is carried over as it wrote + it, which is what makes the error diagnosable — minus two scrubs, both + applied because this text is rendered into the TUI error pane and the + log: the arguments it quoted back, and the endpoint it names when a + call is refused. Host, path and the exception's own type survive both. + + 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: + # The server's own text is the most likely place for either kind of + # secret to reappear: an argument-validation failure quotes the + # value it rejected, and a refused call routinely names the endpoint + # it was made against. Both are cut before anything is + # interpolated. + raise MitoCodeMcpToolCallError( + f"MCP tool '{tool_name}' on server '{server}' failed: " + f"{_scrub_endpoint(_redact_arguments(str(exc)))}" + ) from exc + except asyncio.TimeoutError as exc: + # Scrubbed like the two branches above even though the only text the + # server chose here is the tool name: the rule is applied where the + # message is built, not where the leak was noticed, so a later edit + # that quotes something else here cannot reintroduce one silently. + raise MitoCodeMcpToolCallError( + _scrub_endpoint( + 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 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 -- and a + # transport failure quotes the endpoint for the same reason a + # handshake does, so the endpoint scrub applies here too. + raise MitoCodeMcpToolCallError( + f"MCP tool '{tool_name}' on server '{server}' raised " + f"{type(exc).__name__}: " + f"{_scrub_endpoint(_redact_arguments(str(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 "" diff --git a/src/mitocode/backend/mcp/mcp_config.py b/src/mitocode/backend/mcp/mcp_config.py new file mode 100644 index 0000000..dc29920 --- /dev/null +++ b/src/mitocode/backend/mcp/mcp_config.py @@ -0,0 +1,394 @@ +"""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. + + ``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. + 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. + + 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. + + Returns: + The validated endpoint. + + Raises: + MitoCodeMcpConfigError: If ``value`` is not an http(s) endpoint. + """ + 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: + """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. + + 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. + 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. The + message names the field and the bad element's index, but no values. + """ + value = raw.get(key, []) + 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. + 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. The + message names the field and the bad entry's key, but no values. + """ + value = raw.get(key, {}) + if not isinstance(value, dict): + 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): + raise MitoCodeMcpConfigError( + f"MCP server '{name}': '{key}' entry {entry_key!r} must map a " + "string to a string" + ) + 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/src/mitocode/backend/mcp/mcp_manager.py b/src/mitocode/backend/mcp/mcp_manager.py new file mode 100644 index 0000000..6f84926 --- /dev/null +++ b/src/mitocode/backend/mcp/mcp_manager.py @@ -0,0 +1,551 @@ +"""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, + 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, _scrub_endpoint +from mitocode.backend.mcp.mcp_config import McpServerConfig + +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: BaseException) -> str: + """Render an exception as its type plus its scrubbed message. + + The message is scrubbed of endpoint credentials because this string is what + :meth:`McpManager.status` publishes and what the log line carries: an + operation that fails while a server is unreachable is exactly the operation + whose text quotes the endpoint it failed to reach, and an endpoint is + allowed to carry a token. + + Args: + exc: The exception to describe. + + Returns: + ``": "``, or just the type when the exception + carries no message — a bare ``CancelledError`` has none. + """ + message = _scrub_endpoint(str(exc)) + return f"{type(exc).__name__}: {message}" if message else type(exc).__name__ + + +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`, 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, with + every name distinct. + + Raises: + MitoCodeMcpConfigError: If *config* is neither shape, if a list element + 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): + return McpServerConfig.parse_config(config) + + 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 + + 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. + + 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. 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, 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 + 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 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 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 — 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. + + Both phases are isolated per server, and nothing raised by the first + skips the second: the disconnection phase releases the subprocesses, + and leaving every one of them running because an unregistration was + cancelled — or aborted by a shutdown signal — 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: + BaseException: The first cancellation or shutdown signal a phase + raised, after both phases have run. + """ + deferred: Optional[BaseException] = None + for phase in (self._unregister_one, self._disconnect_one): + try: + await self._run_each(list(self._clients), phase) + except BaseException as exc: # pylint: disable=broad-except + # Remember the first one and carry on to the next phase; the + # caller is told afterwards, not in place of the teardown. The + # whole family is deferred, not just a cancellation: the work + # below re-raises anything that is not an Exception, and + # catching only the cancellation let every other one out of + # this method with the connections still open. + deferred = deferred or exc + if deferred is not None: + raise deferred + + 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. + + 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. + + Raises: + 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. + """ + 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 its most recent + error, then the disabled ones — which have no client and no tools. + """ + snapshot = { + name: McpServerStatus( + state=client.state.value, + # 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() + } + for name in self._disabled: + snapshot[name] = McpServerStatus(state="disabled", error="", tool_count=0) + return snapshot + + async def _run_each( + self, names: List[str], operation: Callable[[str], Awaitable[Any]] + ) -> 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. + + The same holds for the rest of that family. ``gather`` catches whatever + a child raises, so a ``KeyboardInterrupt`` or a framework's own shutdown + signal — neither an ``Exception`` — arrives here as a *result*; a child + failing that way is recorded against its server and re-raised once every + name has been visited, so no server's failure is reported without a + reason and none is swallowed. A cancellation still short-circuits: it + travels one server's worth of work, not the whole batch's. + + 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. + BaseException: The first non-``Exception`` ``BaseException`` any + server's work raised, after every server has been visited. + """ + results = await asyncio.gather( + *(operation(name) for name in names), + return_exceptions=True, + ) + deferred: Optional[BaseException] = None + for name, result in zip(names, results): + if isinstance(result, asyncio.CancelledError): + self._record_failure(name, result, "cancelled before it finished") + raise result + if isinstance(result, BaseException): + self._record_failure( + name, result, "aborted with a signal that is not an Exception" + ) + deferred = deferred or result + if deferred is not None: + raise deferred + + 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, "failed to connect") + return + + if not connected: + self._errors.pop(name, None) + # The client records rather than raises, so this -- not an + # exception -- is what an unreachable server looks like to the + # manager. Without the client's own text the log line says only + # that the server was unavailable, which is the very question the + # operator opened it to answer. + logger.warning( + f"MCP server '{name}' is unavailable; its tools are not registered" + f"{': ' + client.last_error if client.last_error else ''}" + ) + 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. + 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 + + # ``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``. + """ + await self._unregister_one(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, "failed to disconnect") + + 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 + 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. + + 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, []): + 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 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. + + Args: + 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. 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]}") + + 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/src/mitocode/backend/tool/tool_manager.py b/src/mitocode/backend/tool/tool_manager.py index 53e65a7..3d39ff5 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,16 @@ Dict, List, Optional, + Set, + TYPE_CHECKING, Union, get_args, get_origin, Literal, ) +from loguru import logger as loguru_logger + try: import types @@ -28,6 +36,84 @@ 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 = "__" + +#: Longest function name OpenAI's function-calling API accepts. Not an MCP +#: constraint — the specification imposes nothing here — but the composed name is +#: published to the model verbatim, so a long server name plus a long tool name +#: can break a request at the provider after passing every local check. +MAX_TOOL_NAME_LENGTH = 64 + + +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 _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. @@ -219,12 +305,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 +323,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 +341,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 +349,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,10 +366,26 @@ 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. Annotation + inference is the last resort. Between the two sits a third source, a + ``parameters_schema`` key in the ``declare_tool`` metadata, which + :func:`declare_tool` publishes no argument for: only a caller writing + ``func.__mitocode_tool_meta__["parameters_schema"]`` directly can set + it. It is honoured when present and non-``None``, like the explicit + argument — an empty dict is a schema that declares no arguments, not an + absent one — but it is not part of the public surface. + Returns: The registered :class:`ToolSpec`. """ meta = getattr(func, "__mitocode_tool_meta__", {}) + if parameters_schema is None: + # An empty dict is a schema, not an absence of one, in this tier too. + parameters_schema = meta.get("parameters_schema") + if parameters_schema is None: + parameters_schema = _build_parameters_schema(func) tool_name = name or meta.get("name") or func.__name__ human_description = ( @@ -292,13 +400,120 @@ def register( func=func, human_doc=human_description, agent_doc=agent_description, - parameters_schema=_build_parameters_schema(func), + parameters_schema=parameters_schema, 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. + + Registration is all-or-nothing. Every tool is read, normalised and + composed first, and ``self._tools`` is not touched until the whole + catalogue has survived that pass, because a caller can only ever + reclaim what it was told it registered: a catalogue whose later entry + is malformed would otherwise leave its healthy predecessors published, + unreported, and therefore removable by nobody. + + Publishing afterwards cannot itself fail for these inputs -- the name + is an f-string result, so it is always hashable, and the docs and + schema were normalised in the first pass -- which is what makes the + single pass sufficient rather than a prefix of a wider one. If + :meth:`register` ever grows a step that can raise, this needs a + rollback as well. + + The composed name is emitted straight into :meth:`tool_schemas`, so it + has to survive whatever the model API does with a function name. The + usual constraint is ``^[a-zA-Z0-9_-]{1,64}$``: the namespace prefix and + separator are already within that character set, but the *length* is + the server author's and the tool author's to spend, and a long server + name plus a long tool name can exceed 64 characters. Nothing here + fails such a tool — the name is not truncated, because a truncated name + is a different tool — but a warning says what the provider will say + later, when the whole request is rejected for a name this harness + composed. + + Args: + client: A connected MCP client exposing ``server_name`` and ``tools``. + + Returns: + The specs that were successfully registered, in discovery order. + """ + prepared: List[Dict[str, Any]] = [] + claimed: Set[str] = set() + for tool in client.tools: + tool_name = self.build_mcp_tool_name(client.server_name, tool.name) + if tool_name in claimed or self.get_spec(tool_name) is not None: + logger.warning( + f"Skipping MCP tool '{tool_name}': name already registered" + ) + continue + # Claimed as well as taken, so two tools in one catalogue composing + # the same name resolve exactly as they did when each was published + # as it was read: the first wins and the second is skipped. + claimed.add(tool_name) + + if len(tool_name) > MAX_TOOL_NAME_LENGTH: + logger.warning( + f"MCP tool '{tool_name}' is {len(tool_name)} characters, past " + f"the {MAX_TOOL_NAME_LENGTH} an OpenAI-compatible function name " + "may be; the model API will reject the request that carries it" + ) + + description = (tool.description or "").strip() + if not description: + description = ( + f"Tool '{tool.name}' provided by MCP server " + f"'{client.server_name}'." + ) + + prepared.append( + { + "func": _make_invoke(client, tool.name), + "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}", + } + ) + + registered: List[ToolSpec] = [] + for pending in prepared: + registered.append(self.register(**pending)) + 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/src/mitocode/tui/ui_main.py b/src/mitocode/tui/ui_main.py index 5aaad3c..3a02142 100644 --- a/src/mitocode/tui/ui_main.py +++ b/src/mitocode/tui/ui_main.py @@ -1,9 +1,4 @@ -"""Textual-based TUI for MitoCode. - -Renders the agent tree on the left, one or more :class:`AgentView` chat -panels in a dynamic grid, and routes streaming events from the -:class:`AgentOrchestrator` into the focused view. -""" +"""buggy, deprecated, dont reference this file any longer.""" import json import time diff --git a/test/test_mcp_client.py b/test/test_mcp_client.py new file mode 100644 index 0000000..6b9c20a --- /dev/null +++ b/test/test_mcp_client.py @@ -0,0 +1,1650 @@ +import asyncio +import contextlib +import json +import os +import sys +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Optional + +import pytest +from fastmcp import Client, FastMCP +from fastmcp.exceptions import ToolError + +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, + _redact_arguments, + _scrub_endpoint, +) +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 + + +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 _SlowCloseClient: + """Fake client whose ``close()`` blocks long enough to be cancelled inside. + + The window has to be opened *inside* ``close()``: the teardown the shield + protects is exactly that call, so a fake whose ``close()`` returns at once + would finish the await before any cancellation could be delivered and every + test below would pass whether or not the teardown is shielded. The pause is + announced through ``close_started`` so a test cancels against the teardown + it means to interrupt rather than against a sleep it cannot see. + """ + + def __init__(self, error: Optional[Exception] = None) -> None: + """Arm the recorder, optionally to raise *error* after the wait.""" + self.error = error + self.close_started = asyncio.Event() + self.close_finished = asyncio.Event() + + async def close(self) -> None: + """Record that teardown began, yield to the loop, then finish or fail.""" + self.close_started.set() + await asyncio.sleep(0.05) + self.close_finished.set() + if self.error is not None: + raise self.error + + +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. + """ + + +@contextlib.contextmanager +def _jsonrpc_error_stub(): + """Run a local HTTP server whose error messages quote the request path. + + A real handshake against it, rather than a hand-raised exception: the whole + question is what the MCP SDK and the HTTP layer actually put in the message, + so the guarantee has to be pinned against the genuine article. Bound to the + loopback interface on an ephemeral port, so the suite stays offline. + + Yields: + The ``http://127.0.0.1:`` base URL to point a client at. + """ + + class _Handler(BaseHTTPRequestHandler): + """Answers every call with a JSON-RPC error naming the requested path.""" + + def do_POST(self): # noqa: N802 - name fixed by BaseHTTPRequestHandler + """Refuse the call, quoting the path that was requested.""" + payload = json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "error": { + "code": -32603, + "message": f"Internal error while proxying {self.path}", + }, + } + ).encode() + self.send_response(500) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, *args): + """Stay quiet; the stub's output is not this test's subject.""" + + server = ThreadingHTTPServer(("127.0.0.1", 0), _Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + host, port = server.server_address[:2] + yield f"http://{host}:{port}" + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + + +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) + + +async def _cancel_during_disconnect_teardown( + client: McpClient, built: _SlowCloseClient +) -> None: + """Cancel a running ``disconnect()`` while its own ``close()`` is in flight. + + The 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. No message is attached, because on the 3.10 interpreter this was + observed on, a task that ends cancelled is re-cancelled with a fresh, empty + ``CancelledError`` by the loop -- the text a caller passes to ``cancel()`` + does not survive to the awaiter, so a ``match=`` here could only ever fail. + 3.14 keeps the message, which is the interpreter CI runs, so the omission + is a concession to the local venv rather than a claim about every runtime. + + Args: + client: The client whose :meth:`disconnect` must be cancelled. + built: The fake transport installed as the client's client handle. + """ + task = asyncio.ensure_future(client.disconnect()) + await asyncio.wait_for(built.close_started.wait(), timeout=5) + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + +@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_failure_never_echoes_a_credential_in_the_endpoint(): + """The connect path is where a credential in a URL reaches the error pane. + + ``McpServerConfig`` deliberately permits ``?access_token=…``, and a server + or proxy that fails the handshake routinely quotes the endpoint it was + called on — including its query string. That text ends up in ``last_error``, + which :meth:`~mitocode.backend.mcp.mcp_manager.McpManager.status` publishes + verbatim and the TUI pane an operator reads. The scrubbing is therefore + pinned against a *real* transport, not a hand-written exception: the + assertion has to survive whatever the MCP SDK chooses to put in the message. + """ + secret = "sk-CONNECT-QUERY-DO-NOT-LEAK" + with _jsonrpc_error_stub() as base_url: + config = McpServerConfig.parse( + "remote", + {"url": f"{base_url}/mcp?access_token={secret}", "connect_timeout": 10.0}, + ) + client = McpClient(config) + + assert await client.connect() is False + + assert client.state is McpConnectionState.FAILED + assert secret not in client.last_error, f"the URL query leaked: {client.last_error}" + assert "" in client.last_error, "and the endpoint is still reported" + assert "/mcp" in client.last_error, "so which endpoint failed survives" + + +@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_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. + + 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_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.parametrize( + "text", + [ + "MCPError: Internal error while proxying /mcp?access_token=sk-QUERY", + "Failed to reach https://host.invalid/mcp?access_token=sk-QUERY", + "url 'http://h.invalid/mcp?a=1&b=sk-QUERY' is not allowed", + "endpoint mcp?token=sk-QUERY refused", + "MCPError: handshake with https://host.invalid/mcp?a=sk-QUERY#frag failed", + ], +) +def test_endpoint_scrubbing_cuts_query_strings_wherever_they_are_quoted(text): + """A credential in a query must not survive, in either endpoint shape.""" + scrubbed = _scrub_endpoint(text) + + assert "sk-QUERY" not in scrubbed, f"the query survived: {scrubbed}" + assert "" in scrubbed + + +@pytest.mark.parametrize( + ("url", "expected"), + [ + # userinfo + ( + "https://user:pw@host.invalid/mcp", + "https://@host.invalid/mcp", + ), + # query + ( + "https://host.invalid/mcp?access_token=sk-X", + "https://host.invalid/mcp?", + ), + # fragment + ( + "https://host.invalid/mcp#frag", + "https://host.invalid/mcp#", + ), + # all three at once + ( + "http://u:p@host.invalid:8443/mcp?token=sk-X#frag", + "http://@host.invalid:8443/mcp?#", + ), + # a bare host with no path is still an endpoint + ("https://host.invalid?token=sk-X", "https://host.invalid?"), + # nothing to scrub: left byte-identical + ("https://host.invalid/mcp", "https://host.invalid/mcp"), + ("https://host.invalid/mcp?", "https://host.invalid/mcp?"), + ], +) +def test_endpoint_scrubbing_keeps_the_host_and_the_path(url, expected): + """The scrub is surgical: what an operator needs to fix the endpoint stays.""" + assert _scrub_endpoint(f"failed on {url} (retry later)") == ( + f"failed on {expected} (retry later)" + ) + + +@pytest.mark.parametrize( + "text", + [ + # prose that merely mentions a host is not an endpoint + "RuntimeError: handshake exploded", + "bad transport arguments", + "All connection attempts failed for host.invalid", + # a question mark in prose is not a query string + "MCPError: server said no? really", + "why did it fail? no idea", + # scheme-less references with nothing to cut + "POST /mcp returned 404", + ], +) +def test_endpoint_scrubbing_leaves_text_without_a_credential_untouched(text): + """A message with no credential in it must survive byte-identical.""" + assert _scrub_endpoint(text) == text + + +@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_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 + # Only a client that was connected has a state worth watching here: a fresh + # one already reads DISCONNECTED, which would make the assertion below true + # whatever disconnect() did. + client._state = McpConnectionState.CONNECTED + + 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 + # The state was set before the await too, so a teardown abandoned part-way + # still reports the transport as gone rather than leaving CONNECTED behind + # for the next connect() to adopt. + assert ( + client.state is McpConnectionState.DISCONNECTED + ), "a signal during teardown must not leave a CONNECTED client behind" + + +@pytest.mark.asyncio +async def test_force_close_does_not_swallow_a_shutdown_signal_from_the_stack(): + """The stack handle gets the same narrow catch as the client handle. + + ``stack.aclose()`` is the second of the two teardown handles and runs the + exit stack's own callbacks. Its handler is keyed on ``Exception`` for the + same reason the client's is: a shutdown signal or a cancellation delivered + *while* the stack unwinds has to keep travelling, or ``disconnect()`` + swallows whatever an enclosing ``wait_for`` was already trying to do. + Widening this one handler to ``BaseException`` is therefore a regression, + and the ``client``-side twin of this test is + ``test_force_close_does_not_swallow_a_shutdown_signal_raised_while_closing``. + """ + client = McpClient(_in_memory_config()) + recorder = _CloseRecorder() + client._client = recorder + + class _ShutdownOnStackClose: + """Stack that signals a shutdown part-way through unwinding.""" + + def __init__(self) -> None: + """Start with the stack still open.""" + self.closed = False + + async def aclose(self) -> None: + """Attempt the unwind, then announce a shutdown.""" + self.closed = True + raise _ShutdownSignal("stopping during teardown") + + built = _ShutdownOnStackClose() + client._stack = built + # Same precondition as the client-side twin above: the state is the thing + # under test, and a never-connected client already reads DISCONNECTED. + client._state = McpConnectionState.CONNECTED + + with pytest.raises(_ShutdownSignal): + await client.disconnect() + + assert built.closed is True + assert client.last_error == "", "a signal is not a cleanup failure to record" + assert client._client is None, "the client was still released first" + assert client._stack is None, "and the stack handle is dropped too" + assert client.state is McpConnectionState.DISCONNECTED, ( + "the state was set before the await, so an abandoned teardown still " + "reports the transport as gone" + ) + + +@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_cancellation_still_completes_the_transport_teardown(): + """A cancellation must finish ``close()`` before it propagates out. + + ``_force_close`` drops both handles *before* its first await, so a + cancellation landing inside ``close()`` propagates with the transport + unreleased and nothing left to retry: stdio's detached connect task is still + holding the child process alive, which is the exact orphan + :meth:`McpClient.connect` is careful to avoid on its own failure path. Only + the fact that ``close`` *ran to completion* distinguishes a shielded teardown + from a plain await -- both leave the handles cleared -- so the wait below is + the assertion. + """ + client = McpClient(_in_memory_config()) + built = _SlowCloseClient() + stack = _StackRecorder() + client._client = built + client._stack = stack + client._state = McpConnectionState.CONNECTED + + await _cancel_during_disconnect_teardown(client, built) + + await asyncio.wait_for(built.close_finished.wait(), timeout=5) + assert stack.closed is True, "the exit stack was stranded behind the client" + assert client._client is None + assert client._stack is None + assert client.state is McpConnectionState.DISCONNECTED, ( + "a client whose transport is gone must not keep reporting connected -- " + "connect() returns True early for a CONNECTED client, so the next " + "reconnect would adopt a client with no transport behind it" + ) + + +@pytest.mark.asyncio +async def test_disconnect_cancellation_is_not_recorded_as_a_cleanup_failure(): + """A cancellation stays a signal; it is not written to ``last_error``. + + Shielding the teardown must not turn an in-flight cancellation into a + recorded cleanup problem. ``_force_close``'s handler is ``Exception``-only + and ``disconnect`` records what it returns, so a caller that was cancelled + finds :attr:`last_error` exactly as it left it -- the same rule the + shutdown-signal tests pin from the other side. + """ + client = McpClient(_in_memory_config()) + built = _SlowCloseClient() + client._client = built + client._stack = _StackRecorder() + client._state = McpConnectionState.CONNECTED + + await _cancel_during_disconnect_teardown(client, built) + + await asyncio.wait_for(built.close_finished.wait(), timeout=5) + assert client.last_error == "", "a cancellation is not a cleanup failure" + + +@pytest.mark.asyncio +async def test_disconnect_records_the_teardown_error_after_a_slow_close(): + """A slow teardown still reports what it could not release. + + Shielding changes who waits for the teardown, not what it reports: the + errors ``_force_close`` collects are still joined into :attr:`last_error` in + the same shape, and the stack is still unwound behind a client that failed + to close. + """ + client = McpClient(_in_memory_config()) + built = _SlowCloseClient(error=RuntimeError("transport wedged")) + stack = _StackRecorder(OSError("process group already gone")) + client._client = built + client._stack = stack + client._state = McpConnectionState.CONNECTED + + await client.disconnect() + + assert client.last_error == ( + "close failed: RuntimeError: transport wedged; " + "stack close failed: OSError: process group already gone" + ) + assert built.close_finished.is_set(), "the slow close never ran to completion" + assert stack.closed is True + assert client.state is McpConnectionState.DISCONNECTED + + +@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 == [] + + +@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_and_the_teardown_error(): + """A failing teardown must neither replace nor be lost behind the failure. + + ``connect`` records the handshake error first — that is what an operator + needs in order to fix the configuration — and *appends* whatever + ``_force_close`` could not release. On the timeout path that appended half + is the whole warning: it is the difference between "the handshake timed out" + and "the handshake timed out and the child may still be running". + """ + 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.startswith( + "RuntimeError: handshake exploded" + ), "the failure that ended the connect leads the report" + assert client.last_error == ( + "RuntimeError: handshake exploded; " + "close failed: OSError: process group wedged" + ), "and the teardown's own diagnostic follows it instead of being dropped" + 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. 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 + # KNOWN LIMITATION (mcp_client.py:42-44, rationale at :34-41): text after a + # payload-supplied decoy marker is currently NOT redacted -- + # `sk-PART2LEAK` survives here. + # This assertion deliberately does NOT enforce that, because it cannot: + # `startswith` passes just as happily against an implementation that closes + # the leak. The assertion that actually breaks when the leak is closed is + # the exact-output comparison in + # test_redaction_keeps_the_report_after_an_injected_marker() below. + assert redacted.startswith("input_value=") + + +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 + + +#: An endpoint carrying a credential in its query, the shape a remote server +#: quotes back when it refuses a call. Distinct from the connect-side secret +#: so a failing assertion says which surface leaked. +_CALL_SECRET = "sk-CALLSECRET-DO-NOT-LEAK" +_CALL_ENDPOINT = f"https://host.invalid/mcp?access_token={_CALL_SECRET}" + + +@pytest.mark.asyncio +async def test_call_tool_tool_error_scrubs_the_endpoint_the_server_quoted(): + """A server that quotes its own endpoint must not publish its token. + + The connect path is scrubbed because a failed handshake quotes the URL it + was made against; the call path is the same shape reached from the other + direction — a tool error naming the endpoint the request went to, or the + proxy in front of it refusing. It lands in ``MitoCodeMcpToolCallError``, + which the TUI error pane and the log render, and the argument redaction + cannot see a URL at all. + """ + client = McpClient(_in_memory_config()) + + class _EndpointQuotingCall: + """Client whose tool error quotes the endpoint it was reached through.""" + + async def call_tool(self, name, arguments, **kwargs): + raise ToolError(f"upstream refused: GET {_CALL_ENDPOINT}") + + client._client = _EndpointQuotingCall() + + with pytest.raises(MitoCodeMcpToolCallError) as excinfo: + await client.call_tool("echo", {"text": "hi"}) + + message = str(excinfo.value) + assert _CALL_SECRET not in message, f"the endpoint's token leaked: {message}" + assert "" in message, "the credential was cut, not the whole endpoint" + assert "https://host.invalid/mcp" in message, "so which endpoint failed survives" + assert "upstream refused" in message, "the server's own text survived" + + +@pytest.mark.asyncio +async def test_call_tool_unexpected_error_scrubs_the_endpoint_it_quoted(): + """The generic handler quotes the same text, one prefix further along.""" + client = McpClient(_in_memory_config()) + + class _EndpointQuotingCall: + """Client that fails locally while naming the endpoint it dialed.""" + + async def call_tool(self, name, arguments, **kwargs): + raise RuntimeError(f"upstream refused: GET {_CALL_ENDPOINT}") + + client._client = _EndpointQuotingCall() + + with pytest.raises(MitoCodeMcpToolCallError) as excinfo: + await client.call_tool("echo", {"text": "hi"}) + + message = str(excinfo.value) + assert _CALL_SECRET not in message, f"the endpoint's token leaked: {message}" + assert "https://host.invalid/mcp?" in message, "host and path survive" + assert "upstream refused" in message, "the server's own text survived" + assert "MCP tool 'echo' on server 'probe' raised RuntimeError" 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_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 + 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 take its transport and name from the config. + + Every connect test swaps this method out, so without this the transport a + 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( + "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 diff --git a/test/test_mcp_config.py b/test/test_mcp_config.py new file mode 100644 index 0000000..ed750a0 --- /dev/null +++ b/test/test_mcp_config.py @@ -0,0 +1,756 @@ +"""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. + +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. +* **The guarantee does not extend to text this package never authored.** + A *validated* entry is what the messages above are about. Once a server is + actually contacted, the text that reaches ``last_error`` — and from there the + TUI pane and the log — is written by the transport and by the remote server, + which routinely quote the endpoint they failed on. Two narrower guarantees + cover that surface and are pinned in ``test_mcp_client.py`` instead: an + endpoint's userinfo, query and fragment are cut by + ``mcp_client._scrub_endpoint``, and a rejected tool *argument* is cut by + ``mcp_client._redact_arguments``. Anything else a remote server chooses to + quote into its own error text is still echoed: the guarantee is + "credentials we know the shape of", not "no server text reaches a log". +""" + +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 + + 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"] + + +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 + + +# -------------------------------------------------------------------------- +# 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")}) + + +@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( + 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) + + +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} + ) diff --git a/test/test_mcp_manager.py b/test/test_mcp_manager.py new file mode 100644 index 0000000..f2c2b03 --- /dev/null +++ b/test/test_mcp_manager.py @@ -0,0 +1,2619 @@ +import asyncio +import contextlib +import dataclasses +import os +import subprocess +import sys +from types import SimpleNamespace + +import pytest +from loguru import logger as loguru_logger + +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, 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. + + ``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. + + 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=(), 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 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._cancel_next = False + 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 + ] + 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 + self.attempts = 0 + + 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. + ``attempts`` counts every entry, so a test can tell "never given a turn" + from "given one and failed". + + 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. + """ + self.attempts += 1 + if self._cancel_next: + self._cancel_next = False + raise asyncio.CancelledError() + self.connected = self._succeed + if self._succeed: + self.state = McpConnectionState.CONNECTED + self.last_error = "" + 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.""" + 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 + + def go_online(self): + """Make every later connect attempt succeed, as a server coming back.""" + self._succeed = True + + 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 + return f"{tool_name}:{arguments}" + + +class _ShutdownSignal(BaseException): + """A shutdown request that is neither an ``Exception`` nor a cancellation. + + ``KeyboardInterrupt`` and ``SystemExit`` cannot be used to drive this + because asyncio re-raises those out of the event loop itself rather than + delivering them to the awaiting coroutine, which is the whole point here. + """ + + +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, + succeed=True, + fail_on="", + cancel_on="", + signal_on="", + message="", + tools=(), + ): + """Arm one method to raise, independently of the scripted outcome. + + Args: + config: The ``McpServerConfig`` this stand-in represents. + succeed: Value :meth:`connect` reports back when it does not raise. + fail_on: ``"connect"``, ``"disconnect"`` or ``""``. + cancel_on: ``"connect"``, ``"disconnect"`` or ``""``, for a + cancellation rather than a failure. Checked first, because a + cancelled teardown has to reach the caller while a failed one + must be recorded and swallowed. + signal_on: ``"connect"``, ``"disconnect"`` or ``""``, for a + ``BaseException`` that is neither of those two. + message: Message the raised failure carries, so a test can quote an + endpoint the way a real handshake failure would. + tools: Names this stand-in advertises once connected. + """ + super().__init__(config, succeed=succeed, tools=tools) + self._fail_on = fail_on + self._cancel_on = cancel_on + self._signal_on = signal_on + self._message = message + + async def connect(self): + """Raise instead of connecting when this client is armed for it. + + ``attempts`` counts every entry, whichever way this call ends, so a + test can never read ``0`` from a client that was handed a turn. + """ + if self._cancel_on == "connect": + self.attempts += 1 + raise asyncio.CancelledError() + if self._signal_on == "connect": + self.attempts += 1 + raise _ShutdownSignal("application is shutting down") + if self._fail_on == "connect": + self.attempts += 1 + raise RuntimeError(self._message or "handshake exploded") + return await super().connect() + + async def disconnect(self): + """Raise instead of closing when this client is armed for it.""" + if self._cancel_on == "disconnect": + raise asyncio.CancelledError() + if self._signal_on == "disconnect": + raise _ShutdownSignal("application is shutting down") + if self._fail_on == "disconnect": + raise RuntimeError(self._message or "teardown exploded") + await super().disconnect() + + +class ExplodingCatalogueClient: + """Stand-in whose tool catalogue cannot be read at all. + + The real :class:`McpClient` exposes a plain list, so the only way to reach + the registry's ``client.tools`` access is to substitute a client whose + attribute raises. That is the seam between "the server said it was + connected but its catalogue is unreadable" and a failure the manager can + report against that server alone. + """ + + def __init__(self, config, succeed=True): + """Connect on demand, with a catalogue that raises whenever it is read. + + Args: + config: The ``McpServerConfig`` this stand-in represents. + succeed: Value :meth:`connect` reports back. + """ + self.server_name = config.name + self._succeed = succeed + self.state = McpConnectionState.DISCONNECTED + self.last_error = "" + self.connected = False + + @property + def tools(self): + """Raise instead of returning the catalogue. + + Returns: + Never returns. + + Raises: + RuntimeError: Every time the registry asks for the catalogue. + """ + raise RuntimeError("catalogue exploded on access") + + async def connect(self): + """Connect without ever holding a readable catalogue.""" + self.connected = self._succeed + if not self._succeed: + self.state = McpConnectionState.FAILED + self.last_error = "connection refused" + return False + self.state = McpConnectionState.CONNECTED + return True + + async def disconnect(self): + """Close the connection and go back to DISCONNECTED.""" + self.connected = False + self.state = McpConnectionState.DISCONNECTED + + +def _mcp_tool(name, description): + """Build one advertised tool of the shape the registry reads. + + Args: + name: The tool name as advertised, without the namespace prefix. + description: Whatever the server sent. The malformed case sends a + value that is not a string, which is what a server advertising a + field of the wrong JSON type looks like from here. + + Returns: + A ``SimpleNamespace`` carrying ``name``, ``description`` and + ``input_schema``. + """ + return SimpleNamespace( + name=name, + description=description, + input_schema={"type": "object", "properties": {}}, + ) + + +class MalformedCatalogueClient: + """Stand-in advertising a healthy tool followed by an undescribable one. + + ``ToolManager.register_mcp_client`` reads every tool's description, so a + catalogue whose *second* entry carries something that is not a string is + what makes a registration fail part way through. The first tool is + entirely well formed, which is what gives the case its teeth: a registry + that publishes as it reads leaves that tool behind, and a caller told + "the registration failed" has no way to reclaim it. + """ + + #: A description type a well behaved server never sends. + MALFORMED_DESCRIPTION = 123 + + def __init__(self, config, healthy_tool="healthy"): + """Advertise one readable tool and one that cannot be described. + + Args: + config: The ``McpServerConfig`` this stand-in represents. + healthy_tool: Name of the well formed tool. A test sets it to a + name another layer already owns to reach the collision path. + """ + self.server_name = config.name + self.tools = [ + _mcp_tool(healthy_tool, f"{healthy_tool} on {config.name}"), + _mcp_tool("malformed", self.MALFORMED_DESCRIPTION), + ] + self.state = McpConnectionState.DISCONNECTED + self.last_error = "" + self.connected = False + + async def connect(self): + """Connect; publishing is the registry's business, not the client's.""" + self.state = McpConnectionState.CONNECTED + self.connected = True + return True + + async def disconnect(self): + """Close the connection and go back to DISCONNECTED.""" + self.connected = False + self.state = McpConnectionState.DISCONNECTED + + +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``. + """ + made = {} + + def _make(config): + succeed, tools = script.get(config.name, (True, ())) + client = FakeClient( + config, + succeed=succeed, + tools=tools, + keep_tools_on_failure=keep_tools_on_failure, + ) + made[config.name] = client + return client + + _make.made = made + return _make + + +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, + succeed=outcomes.get(config.name, True), + fail_on=script.get(config.name, ""), + ) + made[config.name] = client + return client + + _make.made = made + return _make + + +def _hostile_factory(script): + """Build a client_factory in which each server misbehaves in its own way. + + Args: + script: Mapping from server name to its misbehaviour. ``rude`` names the + method that raises ``"connect"`` or ``"disconnect"``, ``cancel_on`` + names the one that raises :class:`asyncio.CancelledError` instead, + ``catalogue_raises`` swaps in a client whose tool list cannot be + read, ``malformed`` swaps in one whose second advertised tool + cannot be described, ``signal_on`` names the one that raises a + ``BaseException`` which is neither an ``Exception`` nor a + cancellation, and ``succeed``/``tools``/``healthy_tool``/``message`` + configure the rest. A server the script does not mention is well + behaved. + + Returns: + The factory, with every client it built exposed as ``made``. + """ + made = {} + + def _make(config): + spec = dict(script.get(config.name, {})) + if spec.pop("malformed", False): + client = MalformedCatalogueClient( + config, healthy_tool=spec.pop("healthy_tool", "healthy") + ) + elif spec.pop("catalogue_raises", False): + client = ExplodingCatalogueClient(config, succeed=spec.pop("succeed", True)) + else: + client = RudeClient( + config, + succeed=spec.pop("succeed", True), + fail_on=spec.pop("rude", ""), + cancel_on=spec.pop("cancel_on", ""), + signal_on=spec.pop("signal_on", ""), + message=spec.pop("message", ""), + tools=spec.pop("tools", ()), + ) + made[config.name] = client + return client + + _make.made = made + return _make + + +#: Values standing in for credentials, one per configuration field that can +#: carry one. Each is unique so a failing assertion names *which* field leaked +#: rather than only that something did. +_SECRET_USERINFO = "pw-2d4f-secret" +_SECRET_HEADER = "Bearer 7c21-secret" +_SECRET_ENV = "env-5b8e-secret" +_SECRET_ARG = "--arg-9f3a-secret" +_SECRET_PATH = "/srv-6a0b-secret" + +_SECRETS = ( + _SECRET_USERINFO, + _SECRET_HEADER, + _SECRET_ENV, + _SECRET_ARG, + _SECRET_PATH, +) + + +def _http_secret_spec(server_name): + """A remote server entry carrying a secret in every field that takes one.""" + return { + "url": f"https://user:{_SECRET_USERINFO}@{server_name}.invalid/mcp", + "headers": {"Authorization": _SECRET_HEADER}, + "env": {"MCP_TOKEN": _SECRET_ENV}, + "args": [_SECRET_ARG], + "cwd": _SECRET_PATH, + } + + +def _stdio_secret_spec(): + """A local server entry carrying a secret in every field that takes one.""" + return { + "command": _SECRET_PATH, + "args": [_SECRET_ARG], + "env": {"MCP_TOKEN": _SECRET_ENV}, + "cwd": _SECRET_PATH, + "headers": {"Authorization": _SECRET_HEADER}, + } + + +def _assert_carries_no_secret(text, where): + """Fail naming every credential value that reached *text*. + + Args: + text: The message a caller can see -- the TUI pane or the log. + where: What produced it, used only in the failure message. + """ + leaked = [secret for secret in _SECRETS if secret in text] + assert not leaked, f"{where} leaked {leaked}" + + +@contextlib.contextmanager +def _captured_logs(): + """Collect every loguru line emitted inside the block. + + The manager's own logging is half of what it promises an operator: a + failure it records against a server is also named in the log, with the + tool that could not be released. ``status()`` deliberately carries less, + so the two surfaces have to be checked separately. + + Yields: + The list every emitted line is appended to. + """ + lines = [] + sink_id = loguru_logger.add(lines.append, level="DEBUG") + try: + yield lines + finally: + loguru_logger.remove(sink_id) + + +def _config(*entries): + """Build a ``{"mcpServers": {...}}`` document from (name, spec) pairs.""" + 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, + # A wedged child must fail the test rather than hang the whole suite. + timeout=120, + ) + + +@pytest.fixture +def stubborn_registry(monkeypatch): + """Build a registry that can be told to refuse to release particular tools. + + ``unregister`` is the seam the manager's ownership record is written + against, and "the removal raised" is the one outcome no client-side + stand-in can produce, so every test that needs a tool to stay in the + registry reaches it through here. + + Args: + monkeypatch: Patched here and undone by pytest when the test ends. + + Returns: + A factory taking the tool names to refuse. It returns a + ``(tool_manager, armed)`` pair; edit the ``armed`` set in place to arm + or disarm the refusal at any point in the test. + """ + + def _build(*names): + tool_manager = ToolManager() + real_unregister = tool_manager.unregister + armed = set(names) + + def _unregister(name): + if name in armed: + raise RuntimeError("registry exploded") + return real_unregister(name) + + monkeypatch.setattr(tool_manager, "unregister", _unregister) + return tool_manager, armed + + return _build + + +@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_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() + 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_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( + _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_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": "", "good": ""}) + manager = McpManager( + _config( + ("bad", {"url": "https://bad.invalid/mcp"}), + ("good", {"url": "https://good.invalid/mcp"}), + ), + tool_manager=tool_manager, + client_factory=factory, + ) + factory.made["bad"].cancel_next_connect() + + 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_start_keeps_owning_a_tool_it_could_not_unregister(stubborn_registry): + """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, armed = stubborn_registry() + 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.add("mcp__weather__x") + + 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.clear() + 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(stubborn_registry): + """The same guarantee holds on ``refresh``, the other path through it.""" + tool_manager, armed = stubborn_registry() + 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.add("mcp__weather__a") + + 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.clear() + 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_stop_completes_the_teardown_before_re_raising_a_shutdown_signal( + monkeypatch, +): + """A signal out of ``unregister`` must not skip the disconnections either. + + The phase loop deferred a cancellation and nothing else, but the work it + drives re-raises the whole ``BaseException`` family once every server has + been visited — so a signal out of the removal phase walked straight out of + ``stop()``, leaving every connection open and ``status`` still reporting + ``connected``. Deferring the family as a whole is what ``stop``'s own + docstring promises: neither phase can cost the other its work. + """ + 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 _ShutdownSignal("application is shutting down") + return real_unregister(name) + + monkeypatch.setattr(tool_manager, "unregister", _unregister) + + with pytest.raises(_ShutdownSignal): + await manager.stop() + + assert ( + factory.made["bad"].state is McpConnectionState.DISCONNECTED + ), "the signal skipped the disconnection phase" + assert factory.made["good"].state is McpConnectionState.DISCONNECTED + assert "shutting down" in manager.status()["bad"].error, "and was not swallowed" + # The removal was cut short, so the tool is still in the registry -- and + # still claimed there, so status() and the registry agree. + assert set(tool_manager.tools_map()) == {"mcp__bad__x"} + assert manager.status()["bad"].tool_count == 1 + + +@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.""" + 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_stop_completes_teardown_when_an_unregister_raises(stubborn_registry): + """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, armed = stubborn_registry() + 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.add("mcp__bad__x") + + 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(stubborn_registry): + """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, armed = stubborn_registry() + 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.add("mcp__bad__x") + await manager.stop() + assert manager.status()["bad"].tool_count == 1, "still owned, so still counted" + + armed.clear() + 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(): + """``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. + + ``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" + assert ( + manager.status()["weather"].tool_count == 1 + ), "the ownership record must shrink with the catalogue, or status() promises a tool that left" + + +@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_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( + "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(): + """Every name in ``__all__`` resolves; a stale entry is a failed import.""" + 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 + 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(): + """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 +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. + + ``start`` is inside the ``try`` on purpose. A child process outlives the + call that spawned it unless something closes its transport, so a failure + *during* the startup is exactly when the teardown matters most. + """ + 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, + ) + + try: + await manager.start() + 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" + + +@pytest.mark.asyncio +async def test_one_servers_failure_never_lands_on_another_servers_status(): + """Four servers, three broken in different ways, one healthy. + + Each failure has to be visible against the server that caused it and + invisible against every other: a shared error slot would report one reason + four times and say nothing about which server is actually broken. + """ + tool_manager = ToolManager() + factory = _hostile_factory( + { + "raises": {"rude": "connect"}, + "offline": {"succeed": False}, + "unreadable": {"catalogue_raises": True}, + "healthy": {"tools": ("ping",)}, + } + ) + manager = McpManager( + _config( + ("raises", {"url": "https://raises.invalid/mcp"}), + ("offline", {"url": "https://offline.invalid/mcp"}), + ("unreadable", {"url": "https://unreadable.invalid/mcp"}), + ("healthy", {"url": "https://healthy.invalid/mcp"}), + ), + tool_manager=tool_manager, + client_factory=factory, + ) + + await manager.start() + + status = manager.status() + assert set(tool_manager.tools_map()) == { + "mcp__healthy__ping" + }, "only the healthy server publishes tools" + assert ( + status["healthy"].state, + status["healthy"].error, + status["healthy"].tool_count, + ) == ("connected", "", 1), "the healthy server is untouched by its neighbours" + assert ( + status["raises"].state, + status["raises"].tool_count, + ) == ("disconnected", 0), "a connect that raised never reached the registry" + assert status["raises"].error == "RuntimeError: handshake exploded" + assert ( + status["offline"].state, + status["offline"].error, + status["offline"].tool_count, + ) == ("failed", "connection refused", 0) + assert ( + status["unreadable"].state, + status["unreadable"].error, + status["unreadable"].tool_count, + ) == ( + "connected", + "RuntimeError: catalogue exploded on access", + 0, + ), "an unreadable catalogue fails that server's registration, not its connect" + + reasons = {status[name].error for name in ("raises", "offline", "unreadable")} + for name, entry in status.items(): + if name in ("raises", "offline", "unreadable"): + continue + assert ( + entry.error not in reasons + ), f"'{name}' was handed another server's failure: {entry.error!r}" + + +@pytest.mark.asyncio +async def test_stop_after_an_isolated_start_closes_the_connection_that_did_open(): + """A partial start must still be fully torn down. + + ``start`` returns once every server has had its turn, failed or not, so + ``stop`` is the only thing left that can reclaim what did come up. + """ + tool_manager = ToolManager() + factory = _hostile_factory( + {"offline": {"succeed": False}, "healthy": {"tools": ("ping",)}} + ) + manager = McpManager( + _config( + ("offline", {"url": "https://offline.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() + + await manager.stop() + + assert tool_manager.tools_map() == {}, "no tool outlives stop()" + assert factory.made["healthy"].connected is False + assert factory.made["healthy"].state is McpConnectionState.DISCONNECTED + + +@pytest.mark.asyncio +async def test_a_registry_failing_on_the_second_server_leaves_the_first_published( + monkeypatch, +): + """The seam between two servers is one call, not one shared outcome.""" + tool_manager = ToolManager() + real_register = tool_manager.register_mcp_client + + def _register(client): + if client.server_name == "second": + raise ValueError("catalogue exploded") + return real_register(client) + + monkeypatch.setattr(tool_manager, "register_mcp_client", _register) + factory = _hostile_factory( + { + "first": {"tools": ("a",)}, + "second": {"tools": ("b",)}, + "third": {"tools": ("c",)}, + } + ) + manager = McpManager( + _config( + ("first", {"url": "https://first.invalid/mcp"}), + ("second", {"url": "https://second.invalid/mcp"}), + ("third", {"url": "https://third.invalid/mcp"}), + ), + tool_manager=tool_manager, + client_factory=factory, + ) + + await manager.start() + + assert set(tool_manager.tools_map()) == { + "mcp__first__a", + "mcp__third__c", + }, "the servers either side of the failing one still publish" + assert manager.status()["first"].tool_count == 1 + assert manager.status()["third"].tool_count == 1 + assert manager.status()["second"].tool_count == 0 + assert manager.status()["second"].error == "ValueError: catalogue exploded" + + +@pytest.mark.asyncio +async def test_a_tool_the_registry_already_holds_is_neither_adopted_nor_removed(): + """A name that is already taken belongs to whoever put it there. + + ``register_mcp_client`` skips a name that is still registered, so that tool + is absent from the result it returns. Recording the catalogue rather than + the registration result would make the manager its owner: it would count a + tool it never registered, and its ``stop`` would delete another layer's. + """ + tool_manager = ToolManager() + + async def pre_existing(**_kwargs): + """Stand-in for a tool some other layer registered. + + Returns: + A fixed string, so a dispatch proves the original survived. + """ + return "not from MCP" + + tool_manager.register(pre_existing, name="mcp__weather__a", human_doc="mine") + factory = _hostile_factory({"weather": {"tools": ("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"} + assert ( + tool_manager.get_spec("mcp__weather__a").human_doc == "mine" + ), "the name that was already taken was not overwritten" + assert ( + manager.status()["weather"].tool_count == 1 + ), "only the tool the manager actually registered is counted" + + await manager.stop() + + assert set(tool_manager.tools_map()) == { + "mcp__weather__a" + }, "stop must not remove a tool the manager never owned" + assert await tool_manager.tools_map()["mcp__weather__a"]() == "not from MCP" + + +@pytest.mark.asyncio +async def test_a_connected_server_that_advertises_nothing_is_distinguishable_from_a_failed_one(): + """Both contribute no tools; only one of them is broken. + + ``start`` registers only after a successful connect precisely so that an + empty result means "this server offers nothing" and never "this server + never came up". ``status`` has to keep the two apart as well. + """ + tool_manager = ToolManager() + factory = _hostile_factory({"empty": {"succeed": True}, "down": {"succeed": False}}) + manager = McpManager( + _config( + ("empty", {"url": "https://empty.invalid/mcp"}), + ("down", {"url": "https://down.invalid/mcp"}), + ), + tool_manager=tool_manager, + client_factory=factory, + ) + + await manager.start() + + assert ( + manager.status()["empty"].state, + manager.status()["empty"].error, + manager.status()["empty"].tool_count, + ) == ("connected", "", 0), "an empty catalogue is not a failure" + assert ( + manager.status()["down"].state, + manager.status()["down"].error, + manager.status()["down"].tool_count, + ) == ("failed", "connection refused", 0) + assert tool_manager.tools_map() == {} + + +def test_status_lists_every_enabled_server_before_the_disabled_ones(): + """The documented order is the enabled servers first, then the disabled.""" + manager = McpManager( + _config( + ("first", {"url": "https://a.invalid/mcp"}), + ("off", {"url": "https://b.invalid/mcp", "enabled": False}), + ("second", {"url": "https://c.invalid/mcp"}), + ), + tool_manager=ToolManager(), + client_factory=_hostile_factory({"first": {}, "second": {}}), + ) + + assert list(manager.status()) == ["first", "second", "off"] + + +def test_each_status_entry_is_a_frozen_snapshot(): + """A reader cannot rewrite what the next reader of ``status()`` will see.""" + status = McpServerStatus(state="connected", error="", tool_count=2) + + assert (status.state, status.error, status.tool_count) == ("connected", "", 2) + + with pytest.raises(dataclasses.FrozenInstanceError): + status.tool_count = 99 + + assert status.tool_count == 2, "the rejected assignment changed nothing" + + +@pytest.mark.asyncio +async def test_status_counts_exactly_the_tools_surviving_a_half_finished_stop( + stubborn_registry, +): + """One tool out of two refuses to leave; only that one is still owned. + + The ownership record is updated per tool and only once the registry has + confirmed each removal, so it must name the tool that stayed and not the + one that left. Keeping both would make ``status`` promise a tool nobody + could find in the registry. + """ + tool_manager, armed = stubborn_registry() + factory = _hostile_factory({"weather": {"tools": ("a", "b")}}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + armed.add("mcp__weather__a") + + await manager.stop() + + assert set(tool_manager.tools_map()) == { + "mcp__weather__a" + }, "'b' was released, so it must not still be claimed" + assert manager.status()["weather"].tool_count == 1 + assert manager.status()["weather"].error == "RuntimeError: registry exploded" + + +@pytest.mark.asyncio +async def test_stop_removes_the_tools_while_their_client_is_still_connected( + monkeypatch, +): + """The removal phase runs before the disconnection phase. + + The documented reason is the reverse failure: a tool still registered while + its client is gone dispatches into a dead client at call time, in the + model's hands rather than here. The order is observable from the registry + side -- a removal that runs after the connection has closed is removing a + tool nobody could serve any more. + """ + tool_manager = ToolManager() + factory = _hostile_factory({"weather": {"tools": ("get_forecast",)}}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + live_at_removal = {} + real_unregister = tool_manager.unregister + + def _unregister(name): + live_at_removal[name] = factory.made["weather"].connected + return real_unregister(name) + + monkeypatch.setattr(tool_manager, "unregister", _unregister) + + await manager.stop() + + assert live_at_removal == { + "mcp__weather__get_forecast": True + }, "the tool was unregistered after its client had already been disconnected" + + +@pytest.mark.asyncio +async def test_stop_propagates_a_cancellation_raised_by_a_disconnect(): + """A cancelled teardown is the caller unwinding, not a server that failed. + + Both phases still run and every tool is still released. The connection + that refused to close stays open, and ``status`` says so instead of + claiming a teardown that did not happen. + """ + tool_manager = ToolManager() + factory = _hostile_factory( + { + "stuck": {"tools": ("x",), "cancel_on": "disconnect"}, + "good": {"tools": ("y",)}, + } + ) + manager = McpManager( + _config( + ("stuck", {"url": "https://stuck.invalid/mcp"}), + ("good", {"url": "https://good.invalid/mcp"}), + ), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + + with pytest.raises(asyncio.CancelledError): + await manager.stop() + + assert tool_manager.tools_map() == {}, "the removal phase ran before the teardown" + assert factory.made["good"].state is McpConnectionState.DISCONNECTED + assert ( + factory.made["stuck"].state is McpConnectionState.CONNECTED + ), "the connection whose close was cancelled is still open" + assert manager.status()["stuck"].error == "CancelledError" + assert manager.status()["stuck"].tool_count == 0 + + +@pytest.mark.asyncio +async def test_a_cancelled_refresh_leaves_that_server_with_no_tools_and_the_others_alone(): + """A refresh cancelled mid-teardown publishes nothing. + + The tools go before the reconnect is attempted, so a cancellation between + the two leaves the server with none rather than with tools pointing at a + client that is not coming back. + """ + tool_manager = ToolManager() + factory = _hostile_factory( + { + "bad": {"tools": ("x",), "cancel_on": "disconnect"}, + "good": {"tools": ("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() + + with pytest.raises(asyncio.CancelledError): + await manager.refresh("bad") + + assert set(tool_manager.tools_map()) == { + "mcp__good__y" + }, "the cancelled server published nothing and the untouched one kept its tools" + assert manager.status()["bad"].tool_count == 0 + assert manager.status()["bad"].error == "CancelledError" + assert ( + manager.status()["good"].state, + manager.status()["good"].error, + manager.status()["good"].tool_count, + ) == ("connected", "", 1), "refreshing one server left the other alone" + + +@pytest.mark.asyncio +async def test_every_cancelled_server_is_visited_before_the_cancellation_is_raised(): + """One cancelled server must not cost the next one its turn. + + ``start`` runs the servers concurrently and inspects the results only + afterwards, so both cancelled servers are connected before the + cancellation is re-raised. A sequential loop that re-raised on the spot + would never reach the second one, and the healthy server would be left + never tried at all. + """ + tool_manager = ToolManager() + factory = _hostile_factory( + { + "first": {"tools": ("a",)}, + "second": {"tools": ("b",)}, + "healthy": {"tools": ("c",)}, + } + ) + manager = McpManager( + _config( + ("first", {"url": "https://first.invalid/mcp"}), + ("second", {"url": "https://second.invalid/mcp"}), + ("healthy", {"url": "https://healthy.invalid/mcp"}), + ), + tool_manager=tool_manager, + client_factory=factory, + ) + factory.made["first"].cancel_next_connect() + factory.made["second"].cancel_next_connect() + + with pytest.raises(asyncio.CancelledError): + await manager.start() + + assert factory.made["first"].attempts == 1 + assert ( + factory.made["second"].attempts == 1 + ), "the second cancelled server never got its turn" + assert set(tool_manager.tools_map()) == {"mcp__healthy__c"} + assert ( + manager.status()["first"].state, + manager.status()["first"].tool_count, + ) == ("disconnected", 0) + assert ( + manager.status()["healthy"].state, + manager.status()["healthy"].tool_count, + ) == ("connected", 1) + + +@pytest.mark.asyncio +async def test_status_error_is_the_exception_text_verbatim(stubborn_registry): + """The recorded error is the exception and nothing the manager adds. + + ``status`` is rendered into the TUI and written to the log, so forwarding + the exception's own text is the whole contract: a server name, a tool name + or a configuration value interpolated here would be text the exception + never contained, arriving at the one surface a user reads. + """ + tool_manager, _armed = stubborn_registry("mcp__weather__a") + factory = _hostile_factory({"weather": {"tools": ("a",)}}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + + await manager.stop() + + assert ( + manager.status()["weather"].error == "RuntimeError: registry exploded" + ), "no server name, no tool name, no decoration" + + +@pytest.mark.asyncio +async def test_the_log_names_the_tool_a_failed_unregistration_could_not_release( + stubborn_registry, +): + """The tool that stayed is named in the log, and not in ``status``. + + ``status`` describes a server; the log is where an operator reads which of + its tools is still holding a registry slot. A tool name is arbitrary text + from the server, so the two surfaces deliberately carry different amounts + and each has to be checked on its own. + """ + tool_manager, _armed = stubborn_registry("mcp__weather__a") + factory = _hostile_factory({"weather": {"tools": ("a",)}}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + + with _captured_logs() as lines: + await manager.stop() + + reported = [line for line in lines if "failed to unregister" in line] + assert len(reported) == 1, f"expected one unregistration failure, got {reported}" + assert "mcp__weather__a" in reported[0], "the log names the tool that stayed" + assert "weather" in reported[0], "and the server it belongs to" + assert "RuntimeError: registry exploded" in reported[0], "and why it stayed" + assert ( + "mcp__weather__a" not in manager.status()["weather"].error + ), "status describes the server, not the tool" + + +@pytest.mark.asyncio +async def test_a_connect_failure_does_not_publish_a_credential_from_the_endpoint(): + """The error the manager records is scrubbed of endpoint credentials. + + An unreachable server is the failure most likely to quote the endpoint it + was called on, and a remote endpoint may legitimately carry its token in the + query string. ``status()`` is the pane an operator reads, so the scrub has + to happen on the way in rather than being left to the reader to notice. + """ + manager = McpManager( + _config( + ( + "remote", + _http_secret_spec("remote"), + ) + ), + tool_manager=ToolManager(), + client_factory=_hostile_factory( + { + "remote": { + "rude": "connect", + "message": ( + "Internal error while proxying " + f"https://user:{_SECRET_USERINFO}@remote.invalid/mcp" + f"?access_token={_SECRET_HEADER}" + ), + } + } + ), + ) + + await manager.start() + + error = manager.status()["remote"].error + assert error != "", "the failure was recorded at all" + assert "remote.invalid" in error, "and it still says which server failed" + assert "" in error, "with the credential-bearing parts marked" + _assert_carries_no_secret(error, "status() for a failed connect") + await manager.stop() + + +@pytest.mark.asyncio +async def test_the_log_says_why_an_unavailable_server_was_skipped(): + """A client that reports failure rather than raising it still explains itself. + + ``McpClient.connect`` returns ``False`` for a server that is merely down, + which is the common case, so this branch is what an operator reads for + almost every misconfigured entry. Naming the server and nothing else + leaves the one question the log was opened to answer unanswered. + """ + factory = _hostile_factory({"weather": {"succeed": False}}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=ToolManager(), + client_factory=factory, + ) + + with _captured_logs() as lines: + await manager.start() + + reported = [line for line in lines if "is unavailable" in line] + assert len(reported) == 1, f"expected one unavailable server, got {reported}" + assert "weather" in reported[0], "the server is named" + assert "connection refused" in reported[0], "and so is why it was unavailable" + assert manager.status()["weather"].state == "failed" + + +@pytest.mark.asyncio +async def test_the_unavailable_warning_survives_a_client_that_reports_no_reason(): + """A client with nothing to say must not produce a dangling colon. + + The manager accepts any duck-typed client, so an empty ``last_error`` is + reachable even though the real one always explains itself. + """ + factory = _hostile_factory({"weather": {"succeed": False}}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=ToolManager(), + client_factory=factory, + ) + factory.made["weather"].last_error = "" + + with _captured_logs() as lines: + await manager.start() + + reported = [line for line in lines if "is unavailable" in line] + assert len(reported) == 1, f"expected one unavailable server, got {reported}" + assert not reported[0].rstrip().endswith(":"), "no dangling separator is left" + assert "not registered" in reported[0], "and the warning still says what it means" + + +@pytest.mark.asyncio +async def test_a_shutdown_signal_from_a_server_is_recorded_and_re_raised(): + """A ``BaseException`` that is neither an ``Exception`` nor a cancel is neither + swallowed nor silently dropped. + + ``gather(..., return_exceptions=True)`` hands back *any* exception a child + raised as a plain result, so a ``KeyboardInterrupt`` or a framework's own + shutdown signal arrives at the manager as data. Inspecting only for + ``Exception`` and ``CancelledError`` left it on the floor: no record, no log, + and ``start()`` returning normally while a server was being torn down by the + whole process. + """ + tool_manager = ToolManager() + factory = _hostile_factory( + { + "first": {"signal_on": "connect"}, + "second": {"tools": ("b",)}, + } + ) + manager = McpManager( + _config( + ("first", {"url": "https://first.invalid/mcp"}), + ("second", {"url": "https://second.invalid/mcp"}), + ), + tool_manager=tool_manager, + client_factory=factory, + ) + + with pytest.raises(_ShutdownSignal): + await manager.start() + + assert ( + factory.made["first"].attempts == 1 + ), "the server that signalled was still given its turn" + assert ( + "shutting down" in manager.status()["first"].error + ), "and the signal was recorded against it" + assert set(tool_manager.tools_map()) == { + "mcp__second__b" + }, "a signal from one server must not cost another its tools" + + +@pytest.mark.asyncio +async def test_a_shutdown_signal_from_a_teardown_is_recorded_and_re_raised(): + """The same rule governs ``stop()``, which runs the same gather.""" + tool_manager = ToolManager() + factory = _hostile_factory({"weather": {"tools": ("a",)}}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + factory.made["weather"]._signal_on = "disconnect" + + with pytest.raises(_ShutdownSignal): + await manager.stop() + + assert ( + "shutting down" in manager.status()["weather"].error + ), "the teardown's shutdown signal was recorded, not dropped" + + +@pytest.mark.asyncio +async def test_the_log_says_a_tool_the_registry_no_longer_holds_was_dropped( + monkeypatch, +): + """An unregistration the registry reports as a no-op is logged, not silent. + + The registry claims the tool is gone while the manager stops owning it. The + two agree in every real case, so a quiet line here is exactly how a + registry and a manager that disagree about who owns what would go unnoticed + — and the tool would be left in the registry with nobody to remove it. + """ + tool_manager = ToolManager() + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=_hostile_factory({"weather": {"tools": ("a",)}}), + ) + await manager.start() + monkeypatch.setattr(tool_manager, "unregister", lambda name: False) + + with _captured_logs() as lines: + await manager.stop() + + reported = [line for line in lines if "no longer owns it" in line] + assert len(reported) == 1, f"expected one dropped tool, got {reported}" + assert "mcp__weather__a" in reported[0], "the tool that was dropped is named" + assert "weather" in reported[0], "and the server it belonged to" + + +def test_a_duplicate_name_error_does_not_echo_any_configuration_value(): + """The refusal names the name and the positions, and nothing else. + + Both entries carry a distinct secret in every field that accepts one, so + the only way to keep them out of the message is for the message not to + carry either entry at all. + """ + config = [ + McpServerConfig.parse("dup", _http_secret_spec("one")), + McpServerConfig.parse("dup", _stdio_secret_spec()), + ] + + with pytest.raises(MitoCodeMcpConfigError) as excinfo: + McpManager( + config, tool_manager=ToolManager(), client_factory=_hostile_factory({}) + ) + + message = str(excinfo.value) + assert "dup" in message, "the duplicated name is named" + assert "1" in message, "and so is the offending position" + _assert_carries_no_secret(message, "the duplicate-name error") + + +def test_a_mistyped_list_entry_error_does_not_echo_any_configuration_value(): + """A rejected element is described by its position and its type. + + The rejected element *is* the secret-bearing one here, and its neighbours + are real entries carrying secrets too, so the message has to name the + element that was wrong without ever rendering any of them. + """ + config = [ + McpServerConfig.parse("first", _http_secret_spec("one")), + _http_secret_spec("rogue"), + McpServerConfig.parse("third", _stdio_secret_spec()), + ] + + with pytest.raises(MitoCodeMcpConfigError) as excinfo: + McpManager( + config, tool_manager=ToolManager(), client_factory=_hostile_factory({}) + ) + + message = str(excinfo.value) + assert "1" in message, "the offending position is named" + assert "dict" in message, "and so is the type that was found there" + _assert_carries_no_secret(message, "the mistyped-entry error") + + +@pytest.mark.asyncio +async def test_an_unknown_refresh_name_error_does_not_echo_any_configuration_value(): + """The refusal names the server asked for, not how it is configured. + + The existing probe puts a secret in ``headers`` only, so ``url``, + ``env``, ``args`` and ``cwd`` were never at risk from this module's own + message. Here each carries one of its own. + """ + manager = McpManager( + _config( + ("remote", _http_secret_spec("remote")), + ("local", _stdio_secret_spec()), + ), + tool_manager=ToolManager(), + client_factory=_hostile_factory({"remote": {}, "local": {}}), + ) + + with pytest.raises(MitoCodeMcpConfigError) as excinfo: + await manager.refresh("nope") + + message = str(excinfo.value) + assert "nope" in message, "the requested name is named" + assert "remote" in message and "local" in message, "and the configured servers" + _assert_carries_no_secret(message, "the unknown-server error") + + +@pytest.mark.asyncio +async def test_a_registration_that_raises_partway_publishes_nothing_at_all(): + """A catalogue that cannot be read whole publishes no part of itself. + + Driven through the real ``ToolManager``, not a patched one: the second + advertised tool carries a description that is not a string, so the + registry raises while reading it, after the first tool would otherwise + have been published. ``start`` is told the registration failed and + records no ownership for the server, so anything left behind would be a + dispatchable tool that ``status`` does not count and ``stop`` can never + reclaim -- the hazard the ``stuck`` merge defends against, arriving from + the registration side instead. + """ + tool_manager = ToolManager() + factory = _hostile_factory( + {"weather": {"malformed": True}, "healthy": {"tools": ("g",)}} + ) + manager = McpManager( + _config( + ("weather", {"url": "https://w.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__g" + }, "the other server still publishes, and the broken one publishes nothing" + assert manager.status()["weather"].error != "", "the failure was recorded" + assert ( + manager.status()["weather"].tool_count == 0 + ), "a server whose registration failed owns nothing" + + await manager.stop() + + assert ( + tool_manager.tools_map() == {} + ), "a tool the manager cannot name is a tool nobody can remove" + + +@pytest.mark.asyncio +async def test_a_failed_registration_does_not_adopt_a_name_another_layer_owns(): + """The constraint that rules out composing ownership from the catalogue. + + The catalogue's first tool collides with a tool some other layer + registered, so the registry skips it, and its second tool cannot be + described, so the registration fails. A manager that adopted the + advertised catalogue on failure would put both names into its ownership + record -- including the one it skipped -- and ``stop`` would then delete a + tool it never published. Green today; the lock on what any fix here has + to preserve. + """ + tool_manager = ToolManager() + + async def pre_existing(**_kwargs): + """Stand-in for a tool some other layer registered. + + Returns: + A fixed string, so a dispatch proves the original survived. + """ + return "not from MCP" + + tool_manager.register(pre_existing, name="mcp__weather__healthy", human_doc="mine") + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=_hostile_factory( + {"weather": {"malformed": True, "healthy_tool": "healthy"}} + ), + ) + + await manager.start() + + assert ( + manager.status()["weather"].tool_count == 0 + ), "the registration failed, so the manager owns nothing at all" + + await manager.stop() + + assert set(tool_manager.tools_map()) == { + "mcp__weather__healthy" + }, "the tool another layer owns survived a start/stop cycle" + assert await tool_manager.tools_map()["mcp__weather__healthy"]() == "not from MCP" + + +@pytest.mark.asyncio +async def test_a_stdio_server_that_dies_costs_only_its_own_tools(tmp_path): + """The isolation claim end to end, against two real subprocesses. + + One child exits immediately without speaking MCP; the other serves a tool + over a real stdio transport. No ``client_factory`` is involved, so this is + the feature's central promise proven through the real transport rather than + through a stand-in. + """ + server_script = tmp_path / "srv.py" + server_script.write_text( + "from fastmcp import FastMCP\n" + "app = FastMCP('live')\n" + "@app.tool\n" + "def ping() -> str:\n" + " '''Answer a ping.'''\n" + " return 'pong'\n" + "app.run()\n" + ) + + tool_manager = ToolManager() + manager = McpManager( + { + "mcpServers": { + "dead": { + "command": sys.executable, + "args": ["-c", "import sys; sys.exit(3)"], + "connect_timeout": 20.0, + }, + "live": { + "command": sys.executable, + "args": [str(server_script)], + "connect_timeout": 60.0, + }, + } + }, + tool_manager=tool_manager, + ) + + try: + await manager.start() + + dead = manager.status()["dead"] + assert (dead.state, dead.tool_count) == ( + "failed", + 0, + ), f"the dead server reported {dead.error!r}" + assert dead.error != "", "and it says why" + + live = manager.status()["live"] + assert (live.state, live.error, live.tool_count) == ("connected", "", 1) + assert set(tool_manager.tools_map()) == {"mcp__live__ping"} + assert await tool_manager.tools_map()["mcp__live__ping"]() == "pong" + finally: + await manager.stop() + + assert tool_manager.tools_map() == {} + assert manager.status()["live"].state == "disconnected" + + +@pytest.mark.asyncio +async def test_refresh_after_a_failed_refresh_publishes_the_catalogue_again(): + """A server that failed to reconnect is not permanently marked broken. + + The failed refresh leaves no tools behind, so the ownership record is empty + when the next one runs. That is what makes recovery possible: the retry has + to start from a clean slate rather than from a record naming tools that + already left. + """ + tool_manager = ToolManager() + factory = _hostile_factory({"weather": {"tools": ("a",)}}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + client = factory.made["weather"] + client.go_offline() + + await manager.refresh("weather") + + assert tool_manager.tools_map() == {}, "the failed refresh published nothing" + assert manager.status()["weather"].state == "failed" + + client.go_online() + # The failed connect dropped the cached catalogue, exactly as the real + # client does, so the server "coming back" is a server advertising + # something new rather than the same thing twice. + for name in ("a", "b"): + client.tools.append( + SimpleNamespace( + name=name, + description=f"{name} on weather", + input_schema={"type": "object", "properties": {}}, + ) + ) + await manager.refresh("weather") + + assert set(tool_manager.tools_map()) == { + "mcp__weather__a", + "mcp__weather__b", + }, "the recovery publishes the catalogue the server advertises now" + assert manager.status()["weather"] == McpServerStatus( + "connected", "", 2 + ), "a failure the retry resolved is history, not a current error" + + +@pytest.mark.asyncio +async def test_refresh_without_a_name_never_touches_a_disabled_server(): + """``refresh()`` with no name means every *enabled* server. + + A disabled entry has no client, so reaching one is not a failure a + ``try``/``except`` around a reconnect could absorb: it would be an + unrecorded ``KeyError`` out of a method documented never to raise for a + server. + """ + factory = _hostile_factory({"on": {"tools": ("a",)}}) + manager = McpManager( + _config( + ("on", {"url": "https://on.invalid/mcp"}), + ("off", {"url": "https://off.invalid/mcp", "enabled": False}), + ), + tool_manager=ToolManager(), + client_factory=factory, + ) + + await manager.refresh() + + assert "off" not in factory.made, "a disabled server is never given a client" + assert factory.made["on"].state is McpConnectionState.CONNECTED + assert manager.status()["off"] == McpServerStatus("disabled", "", 0) diff --git a/test/test_tool_manager.py b/test/test_tool_manager.py index f5e1887..70cac36 100644 --- a/test/test_tool_manager.py +++ b/test/test_tool_manager.py @@ -1,18 +1,117 @@ +import contextlib +import inspect import os +import subprocess import sys import unittest.mock as mock +from types import SimpleNamespace +from typing import Any, Dict, List, Literal, Optional, Set, Tuple + +import pytest +from loguru import logger as loguru_logger sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) 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 ( + MAX_TOOL_NAME_LENGTH, + MCP_TOOL_PREFIX, + MCP_TOOL_SEPARATOR, ToolManager, declare_tool, _build_parameters_schema, ) +#: 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 _BlockMcpSdk(importlib.abc.MetaPathFinder): + \"\"\"Make both MCP SDK distributions unresolvable, however they are spelled.\"\"\" + + def find_spec(self, fullname, path=None, target=None): + if fullname.split(".")[0] in ("fastmcp", "mcp"): + raise ImportError(f"{fullname} is not installed") + return None + + +class _DuckTypedClient: + server_name = "files" + tools = [SimpleNamespace(name="read", description="d", input_schema=None)] + + async def call_tool(self, tool_name, arguments): + return "" + + +def leaked(): + return sorted( + m + for m in sys.modules + if m.split(".")[0] in ("fastmcp", "mcp") or m.startswith("mitocode.backend.mcp") + ) + + +sys.meta_path.insert(0, _BlockMcpSdk()) +import mitocode.backend.tool.tool_manager as tool_manager + +assert not leaked(), f"mcp package imported at import time: {leaked()}" + +# 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 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 + 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() @@ -102,3 +201,1204 @@ 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_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_MCP_SDK, tmp_path) + 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 _malformed_tool(name, description=123): + """An advertised tool whose description is not a string. + + The type a server sends when it puts a number in a field the MCP schema + says is a string. It is what makes a catalogue unreadable part-way: the + registry has to strip the description, and that step can only fail on a + value that is not a string. + + Args: + name: The tool's name as advertised. + description: The undescribable value to carry. + """ + return _fake_tool(name, description=description) + + +@contextlib.contextmanager +def _captured_logs(): + """Collect every loguru line emitted inside the block. + + Asserting against a patched ``logger`` proves the call was made; asserting + against a real sink proves it reaches an operator's log file, which is the + half that matters for a warning nobody is watching the console for. + + Yields: + The list every emitted line is appended to. + """ + lines = [] + sink_id = loguru_logger.add(lines.append, level="DEBUG") + try: + yield lines + finally: + loguru_logger.remove(sink_id) + + +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"}), + ] + + +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}" + + +@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. + + Both spellings have no element or value type to advertise, so neither + invents an ``items`` or ``additionalProperties`` key. + """ + + 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. The tier is + therefore honoured but not advertised as public API — see ``register``'s + docstring. + """ + 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_keeps_an_empty_parameters_schema_from_the_metadata_tier(): + """Tier two reads an empty schema like tier one does, not by truthiness. + + An explicit ``parameters_schema={}`` says "this tool takes no arguments", + and that is a schema, not the absence of one. Deriving one instead would + advertise arguments the caller never agreed to, from a source that only + looked like a default. + """ + manager = ToolManager() + + @declare_tool(name="no_args") + def no_args(**kwargs): + return kwargs + + no_args.__mitocode_tool_meta__["parameters_schema"] = {} + + spec = manager.register(no_args) + + assert spec.parameters_schema == {}, "an empty schema from the metadata is kept" + + +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 not fake_logger.warning.called, "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 not fake_logger.warning.called, "a free name is not a warning" + + +def test_register_mcp_client_reaches_a_real_log_sink_when_it_skips_a_collision(): + """The skip is a log line an operator can find, not just a mock call. + + The ``logger.warning`` a skipped collision emits is the only trace that a + server advertises a tool this harness could not register; asserting against + a patched logger would pin the call site without proving anything reaches a + sink. + """ + 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 _captured_logs() as lines: + specs = manager.register_mcp_client(client) + + assert specs == [], "a colliding tool is not registered" + warnings = [line for line in lines if "Skipping MCP tool" in line] + assert len(warnings) == 1, f"expected one skipped collision, got {warnings}" + assert "mcp__files__read" in warnings[0], "the skipped name is in the line" + + +def test_register_mcp_client_logs_nothing_when_no_name_collides(): + """A clean registration emits no log line at all.""" + manager = ToolManager() + client = _fake_mcp_client("files", [_fake_tool("read", description="remote")]) + + with _captured_logs() as lines: + specs = manager.register_mcp_client(client) + + assert len(specs) == 1, "the tool was registered" + assert lines == [], f"a registration with no collision is silent, got {lines}" + + +def test_register_mcp_client_publishes_nothing_when_the_catalogue_is_unreadable(): + """Registration is all-or-nothing, owned here rather than downstream. + + The catalogue's second tool carries a description that is not a string, so + the read fails after the first tool would otherwise already be published. A + caller told only that "the registration failed" cannot name what to reclaim, + so anything left behind would be a dispatchable tool nobody owns. Driven + here, in the registry's own suite: this property belongs to + ``register_mcp_client`` and has to stay green even if the manager is never + exercised against it. + """ + manager = ToolManager() + client = _fake_mcp_client( + "weather", + [ + _fake_tool("healthy", description="healthy on weather"), + _malformed_tool("malformed"), + ], + ) + + with pytest.raises(AttributeError): + manager.register_mcp_client(client) + + assert manager.tools_map() == {}, "not even the healthy predecessor survives" + assert manager.tool_schemas() == [], "so there is nothing to dispatch either" + assert manager.get_spec("mcp__weather__healthy") is None + + +@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 neither the callable nor 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_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" + + +def test_register_mcp_client_warns_when_a_composed_name_is_too_long_for_a_model(): + """The composed name goes to the model verbatim, so its length is ours to warn about. + + The MCP specification imposes no name limit; the model API does, at 64 + characters. A server name plus a tool name that together exceed it produce + a tool that registers perfectly here and takes down the whole request at the + provider. The tool is still registered — truncating would rename it, and a + renamed tool is a different tool — but the warning says why. + """ + manager = ToolManager() + long_server = "s" * 40 + long_tool = "t" * 40 + client = _fake_mcp_client(long_server, [_fake_tool(long_tool, description="d")]) + composed = ToolManager.build_mcp_tool_name(long_server, long_tool) + + with _captured_logs() as lines: + specs = manager.register_mcp_client(client) + + assert len(composed) > MAX_TOOL_NAME_LENGTH, "the case is the one being tested" + assert [spec.name for spec in specs] == [composed], "and the tool still registers" + warnings = [line for line in lines if "characters, past the 64" in line] + assert len(warnings) == 1, f"expected one length warning, got {warnings}" + assert composed in warnings[0], "naming the composed name that is too long" + + +def test_register_mcp_client_does_not_warn_about_the_length_of_a_usable_name(): + """An ordinary namespaced name is well inside the limit and stays quiet.""" + manager = ToolManager() + client = _fake_mcp_client("weather", [_fake_tool("get_forecast", description="d")]) + + with _captured_logs() as lines: + specs = manager.register_mcp_client(client) + + assert [spec.name for spec in specs] == ["mcp__weather__get_forecast"] + assert lines == [], f"a 26-character name is not a warning, got {lines}" + + +@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. + + The defect this pins 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. ``_make_invoke`` closed it; the callable's signature is + pinned separately by + ``test_registered_mcp_tool_exposes_no_dispatch_bindings_as_keywords``. + """ + calls = [] + manager = ToolManager() + client = _fake_mcp_client("files", [_fake_tool("read", description="r")], calls) + manager.register_mcp_client(client) + + await manager.tools_map()["mcp__files__read"](_tool_name="write", path="a.txt") + + assert calls == [ + ("read", {"_tool_name": "write", "path": "a.txt"}) + ], "the dispatch must stay on the registered tool, with _tool_name as a plain argument" + + +@pytest.mark.asyncio +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 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 future regression. + """ + calls = [] + manager = ToolManager() + client = _fake_mcp_client("files", [_fake_tool("read", description="r")], calls) + manager.register_mcp_client(client) + + await manager.tools_map()["mcp__files__read"](_client="not-a-client") + + assert calls == [ + ("read", {"_client": "not-a-client"}) + ], "the real client must still receive the call, with _client as a plain argument" diff --git a/todo.md b/todo.md index b3f2e8f..78a964e 100644 --- a/todo.md +++ b/todo.md @@ -11,7 +11,8 @@ - [x] 重构Agent对外接口 (Agent → AgentSession/AgentRuntime/AgentOrchestrator 三层分离) - [ ] 添加上下文管理功能 - [ ] 添加索引功能 -- [ ] 添加MCP功能 +- [x] 添加MCP核心类 (配置解析 / 客户端连接 / 服务端生命周期 / 工具注册) +- [ ] MCP 接入 __main__ 与 TUI (创建McpManager、加载配置、展示status、配置项) - [ ] 添加Agent Skills功能 - [x] 完成TUI (TUI直接调用backend，移除HTTP/SSE层) - [x] 添加工具管理功能:PR #6