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/__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..e75c8de 100644 --- a/src/mitocode/backend/exception/mitocode_exception.py +++ b/src/mitocode/backend/exception/mitocode_exception.py @@ -74,3 +74,19 @@ class MitoCodeChatProviderModelError(MitoCodeChatProviderError): class MitoCodeChatProviderInitError(MitoCodeChatProviderError): """Raised when chat provider initialization fails.""" + + +class MitoCodeMcpError(MitoCodeBaseError): + """Base class for MCP-related errors.""" + + +class MitoCodeMcpConfigError(MitoCodeMcpError): + """Raised when an MCP server configuration entry is invalid.""" + + +class MitoCodeMcpConnectionError(MitoCodeMcpError): + """Raised when an MCP server connection cannot be established or is lost.""" + + +class MitoCodeMcpToolCallError(MitoCodeMcpError): + """Raised when an MCP tool call fails or returns an error result.""" diff --git a/src/mitocode/backend/mcp/__init__.py b/src/mitocode/backend/mcp/__init__.py new file mode 100644 index 0000000..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..0c5fc93 --- /dev/null +++ b/src/mitocode/backend/mcp/mcp_client.py @@ -0,0 +1,313 @@ +"""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]*)" +) + + +def _redact_arguments(text: str) -> str: + """Cut argument payloads out of error text while keeping it diagnosable. + + Only the value is removed: the tool, the server, the offending field name + and the description of what was expected all survive, so the message still + says *what* was wrong even though it no longer says *with what*. + + Args: + text: Error text as reported, possibly quoting the call's arguments. + + Returns: + The same text with every ``input_value=`` payload replaced by + :data:`_REDACTED`. + """ + return _ARGUMENT_PAYLOAD.sub(rf"input_value={_REDACTED}\g", text) + + +class McpConnectionState(str, Enum): + """Connection state of a single :class:`McpClient`.""" + + DISCONNECTED = "disconnected" + CONNECTED = "connected" + FAILED = "failed" + + +class McpClient: + """One MCP server connection plus its cached tool catalogue. + + Args: + config: The validated server definition to connect to. + """ + + def __init__(self, config: McpServerConfig): + """Store the server definition; the connection opens on :meth:`connect`. + + Args: + config: The validated server definition to connect to. + """ + self._config = config + self._client: Optional[Client] = None + self._stack: Optional[AsyncExitStack] = None + self._tools: List[Any] = [] + self._state = McpConnectionState.DISCONNECTED + self._last_error: str = "" + + def _build_fastmcp_client(self) -> Client: + """Construct the FastMCP client for this server. + + Exists as a seam so tests can substitute an in-memory client. + """ + return Client( + self._config.build_transport(), + name=self._config.name, + timeout=self._config.timeout, + ) + + @property + def server_name(self) -> str: + """Name this server's tools are namespaced under.""" + return self._config.name + + @property + def state(self) -> McpConnectionState: + """Current connection state, as of the last :meth:`connect` or :meth:`disconnect`.""" + return self._state + + @property + def last_error(self) -> str: + """Description of the most recent failure, or ``""`` when there is none.""" + return self._last_error + + @property + def tools(self) -> List[Any]: + """Tool objects cached by :meth:`connect`; empty while not connected.""" + return self._tools + + async def connect(self) -> bool: + """Open the connection and cache the server's tool catalogue. + + Failures are recorded rather than raised: an unreachable server must not + take down the servers that are reachable. Callers read :attr:`state` or + the return value to find out what happened. + + The timeout path deliberately closes the client rather than merely + abandoning the await. FastMCP's ``StdioTransport`` spawns its subprocess + inside a *detached* task and the calling coroutine only waits on that + task's ready event, so cancelling our await — which is all + ``asyncio.wait_for`` does — leaves the transport's connect task running + and the child process alive as an orphan. Only ``Client.close()`` sets + the transport's stop event and tears down the process group. + + 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. + await asyncio.shield(self._force_close()) + self._state = McpConnectionState.FAILED + self._last_error = f"{type(exc).__name__}: {exc}" + if isinstance(exc, Exception): + return False + raise + + self._state = McpConnectionState.CONNECTED + self._last_error = "" + return True + + async def _force_close(self) -> None: + """Tear the connection down, swallowing and recording cleanup errors. + + Shared by the failure path of :meth:`connect` and by + :meth:`disconnect` so there is exactly one way to release a connection. + The handles are dropped before the awaits so that a transport which + raises while closing still leaves the object in a clean state. + """ + client = self._client + stack = self._stack + self._client = None + self._stack = None + self._tools = [] + + errors = [] + if client is not None: + try: + await client.close() + except Exception as exc: # pylint: disable=broad-except + errors.append(f"close failed: {type(exc).__name__}: {exc}") + if stack is not None: + try: + await stack.aclose() + except Exception as exc: # pylint: disable=broad-except + errors.append(f"stack close failed: {type(exc).__name__}: {exc}") + if errors: + self._last_error = "; ".join(errors) + + async def disconnect(self) -> None: + """Close the connection if one is open. Idempotent. + + Errors raised while closing are swallowed and recorded in + :attr:`last_error` rather than propagated, because a failing teardown + must not prevent the remaining servers from being shut down. + """ + if self._stack is None and self._client is None: + return + await self._force_close() + self._state = McpConnectionState.DISCONNECTED + + async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> str: + """Invoke *tool_name* on the connected server and render the result as text. + + Args: + tool_name: The tool's name **as advertised by the server**, without + the ``mcp____`` namespace prefix. + arguments: Keyword arguments to forward to the tool. + + Returns: + The tool result as text. + + Raises: + MitoCodeMcpToolCallError: If the client is not connected, the call + times out, or the server reports a tool error. + """ + server = self._config.name + client = self._client + if client is None: + raise MitoCodeMcpToolCallError( + f"MCP server '{server}' is not connected, " + f"cannot call tool '{tool_name}'" + ) + + timeout = self._config.timeout + try: + result = await asyncio.wait_for( + client.call_tool(tool_name, arguments, timeout=timeout), + timeout=timeout, + ) + except ToolError as exc: + # The server's own text is the most likely place for the arguments to + # reappear: an argument-validation failure quotes the value it + # rejected, so the payload is cut before anything is interpolated. + raise MitoCodeMcpToolCallError( + f"MCP tool '{tool_name}' on server '{server}' failed: " + f"{_redact_arguments(str(exc))}" + ) from exc + except asyncio.TimeoutError as exc: + raise MitoCodeMcpToolCallError( + f"MCP tool '{tool_name}' on server '{server}' timed out " + f"after {timeout} seconds" + ) from exc + except MitoCodeMcpToolCallError: + raise + except Exception as exc: # pylint: disable=broad-except + # Arguments are never echoed here either: they routinely carry file + # contents and credentials, and this text reaches the TUI error pane + # and logs. A local validation failure can quote them just as a + # server-side one does, so the same redaction applies. + raise MitoCodeMcpToolCallError( + f"MCP tool '{tool_name}' on server '{server}' raised " + f"{type(exc).__name__}: {_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..6080360 --- /dev/null +++ b/src/mitocode/backend/mcp/mcp_manager.py @@ -0,0 +1,517 @@ +"""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 +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 message. + + Args: + exc: The exception to describe. + + Returns: + ``": "``, or just the type when the exception + carries no message — a bare ``CancelledError`` has none. + """ + message = 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 a cancellation in the first + does not skip the second: the disconnection phase releases the + subprocesses, and leaving every one of them running because an + unregistration was cancelled is the failure mode this method exists to + prevent. A registry that refuses to give a tool back must not skip that + teardown either, and a connection that refuses to close must not strand + the others; either is recorded against its server in :meth:`status` + instead. + + Raises: + asyncio.CancelledError: If one server's unregistration or + disconnection is cancelled, after both phases have run. + """ + cancelled: Optional[asyncio.CancelledError] = None + for phase in (self._unregister_one, self._disconnect_one): + try: + await self._run_each(list(self._clients), phase) + except asyncio.CancelledError as exc: + # Remember the first one and carry on to the next phase; the + # caller is told afterwards, not in place of the teardown. + cancelled = cancelled or exc + if cancelled is not None: + raise cancelled + + async def refresh(self, server_name: Optional[str] = None) -> None: + """Reconnect servers so a changed tool catalogue is picked up. + + 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. + + 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. + """ + results = await asyncio.gather( + *(operation(name) for name in names), + return_exceptions=True, + ) + for name, result in zip(names, results): + if isinstance(result, asyncio.CancelledError): + self._record_failure(name, result, "cancelled before it finished") + raise result + + async def _connect_one(self, name: str) -> None: + """Connect one server and register the tools it advertises. + + 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) + logger.warning( + f"MCP server '{name}' is unavailable; its tools are not registered" + ) + return + + # Whatever this server registered last time is released first, so a + # repeated start or refresh cannot leave the previous names behind as + # tools nobody owns and stop() cannot unregister. + 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..7c83098 100644 --- a/src/mitocode/backend/tool/tool_manager.py +++ b/src/mitocode/backend/tool/tool_manager.py @@ -2,6 +2,10 @@ Converts Python type annotations into JSON Schema tool definitions and keeps separate human/agent-facing documentation per registered tool. + +A caller may also supply the schema outright, which is what tools +discovered on an MCP server need: their arguments are described by the +server, not by anything introspectable in this process. """ import inspect @@ -12,12 +16,15 @@ Dict, List, Optional, + TYPE_CHECKING, Union, get_args, get_origin, Literal, ) +from loguru import logger as loguru_logger + try: import types @@ -28,6 +35,78 @@ from mitocode.backend.exception.mitocode_exception import MitoCodeToolDefinitionError +if TYPE_CHECKING: + # Imported for type checking only. The runtime stays on duck typing so + # this module keeps importing in an environment without fastmcp installed. + from mitocode.backend.mcp.mcp_client import McpClient + +logger = loguru_logger.bind(name=__name__) + +#: Prefix every namespaced MCP tool name carries, so a remote tool is +#: recognisably remote and cannot be confused with a built-in one. +MCP_TOOL_PREFIX = "mcp__" + +#: Separator between the server name and the tool name inside that prefix. +MCP_TOOL_SEPARATOR = "__" + + +def _normalize_input_schema(raw: Any) -> Dict[str, Any]: + """Coerce an MCP tool's advertised input schema into a usable JSON Schema. + + Servers are not required to send a complete schema. This fills in the + pieces :meth:`ToolManager.tool_schemas` always needs so a malformed + schema cannot poison an entire LLM request. + + Args: + raw: The ``input_schema`` value advertised by the MCP server. + + Returns: + A dict with at least ``type`` and ``properties`` keys. + """ + if not isinstance(raw, dict): + return {"type": "object", "properties": {}} + normalized = dict(raw) + normalized.setdefault("type", "object") + if not isinstance(normalized.get("properties"), dict): + normalized["properties"] = {} + return normalized + + +def _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 +298,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 +316,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 +334,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 +342,14 @@ def _decorator(func: Callable): return _decorator - def register( + def register( # pylint: disable=too-many-arguments self, func: Callable, *, name: Optional[str] = None, human_doc: Optional[str] = None, agent_doc: Optional[str] = None, + parameters_schema: Optional[Dict[str, Any]] = None, display_action: Optional[str] = None, display_param: Optional[str] = None, ) -> ToolSpec: @@ -274,6 +359,12 @@ def register( :func:`declare_tool`, which in turn takes precedence over the function's own name and docstring. + ``parameters_schema`` is needed for arguments that cannot be read off + a signature — an MCP server advertises its own JSON Schema — so when it + is supplied it is used verbatim instead of being derived, with a + ``parameters_schema`` entry in the ``declare_tool`` metadata as the + fallback and annotation inference last. + Returns: The registered :class:`ToolSpec`. """ @@ -292,13 +383,83 @@ def register( func=func, human_doc=human_description, agent_doc=agent_description, - parameters_schema=_build_parameters_schema(func), + parameters_schema=( + parameters_schema + if parameters_schema is not None + else (meta.get("parameters_schema") or _build_parameters_schema(func)) + ), display_action=display, display_param=param, ) self._tools[tool_name] = spec return spec + def unregister(self, name: str) -> bool: + """Remove a tool from the registry. + + Args: + name: The tool name to remove. + + Returns: + True if a tool was registered under *name* and has been removed, + False if no such tool existed. + """ + return self._tools.pop(name, None) is not None + + @staticmethod + def build_mcp_tool_name(server_name: str, tool_name: str) -> str: + """Compose the namespaced MitoCode name for an MCP tool. + + Args: + server_name: The MCP server's configured name. + tool_name: The tool's name as advertised by that server. + + Returns: + ``mcp____``. + """ + return f"{MCP_TOOL_PREFIX}{server_name}{MCP_TOOL_SEPARATOR}{tool_name}" + + def register_mcp_client(self, client: "McpClient") -> List[ToolSpec]: + """Register every tool an MCP client has discovered. + + The client must already be connected so that its cached tool catalogue + is populated. A tool whose composed name is already taken is skipped + with a warning rather than overwriting the existing registration. + + Args: + client: A connected MCP client exposing ``server_name`` and ``tools``. + + Returns: + The specs that were successfully registered, in discovery order. + """ + registered: List[ToolSpec] = [] + for tool in client.tools: + tool_name = self.build_mcp_tool_name(client.server_name, tool.name) + if self.get_spec(tool_name) is not None: + logger.warning( + f"Skipping MCP tool '{tool_name}': name already registered" + ) + continue + + description = (tool.description or "").strip() + if not description: + description = ( + f"Tool '{tool.name}' provided by MCP server " + f"'{client.server_name}'." + ) + + registered.append( + self.register( + _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}", + ) + ) + return registered + def tools_map(self) -> Dict[str, Callable]: """Return a name-to-callable mapping of all registered tools.""" return {name: spec.func for name, spec in self._tools.items()} diff --git a/test/test_mcp_client.py b/test/test_mcp_client.py new file mode 100644 index 0000000..c19d4be --- /dev/null +++ b/test/test_mcp_client.py @@ -0,0 +1,1230 @@ +import asyncio +import os +import sys +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, +) +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 _ShutdownSignal(BaseException): + """A shutdown request that is neither an ``Exception`` nor a cancellation. + + Stand-in for the third ``BaseException`` an awaited operation can raise: a + framework's own "stop everything" signal. ``KeyboardInterrupt`` and + ``SystemExit`` cannot be used here because asyncio re-raises those out of + the event loop itself instead of delivering them to the awaiting coroutine. + """ + + +async def _cancel_twice_during_teardown( + client: McpClient, built: _HandshakeHangingClient +) -> None: + """Cancel a hanging ``connect`` again while its own teardown is in flight. + + The second cancellation is timed off ``close_started`` so it cannot race + the teardown: without it the test would pass whether or not the teardown is + shielded. + + Args: + client: The client whose :meth:`connect` must be cancelled twice. + built: The fake transport installed as the client's build seam. + """ + task = asyncio.ensure_future(client.connect()) + await asyncio.wait_for(built.entered.wait(), timeout=5) + task.cancel() + + async def _cancel_mid_teardown() -> None: + await asyncio.wait_for(built.close_started.wait(), timeout=5) + task.cancel() + + canceller = asyncio.ensure_future(_cancel_mid_teardown()) + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.wait_for(canceller, timeout=5) + + +@pytest.mark.asyncio +async def test_connect_discovers_tools_from_server(monkeypatch): + config = _in_memory_config() + client = McpClient(config) + monkeypatch.setattr(client, "_build_fastmcp_client", lambda: Client(_server())) + + assert await client.connect() is True + assert client.state is McpConnectionState.CONNECTED + assert client.last_error == "" + assert sorted(t.name for t in client.tools) == ["add", "boom", "echo"] + + +@pytest.mark.asyncio +async def test_call_tool_returns_text_result(monkeypatch): + config = _in_memory_config() + client = McpClient(config) + monkeypatch.setattr(client, "_build_fastmcp_client", lambda: Client(_server())) + await client.connect() + + assert await client.call_tool("echo", {"text": "hi"}) == "echo:hi" + + +@pytest.mark.asyncio +async def test_call_tool_serializes_structured_result_as_json(monkeypatch): + config = _in_memory_config() + client = McpClient(config) + monkeypatch.setattr(client, "_build_fastmcp_client", lambda: Client(_server())) + await client.connect() + + result = await client.call_tool("add", {"a": 2, "b": 3}) + assert result == '{"sum": 5}' + + +@pytest.mark.asyncio +async def test_call_tool_raises_on_tool_error(monkeypatch): + config = _in_memory_config() + client = McpClient(config) + monkeypatch.setattr(client, "_build_fastmcp_client", lambda: Client(_server())) + await client.connect() + + with pytest.raises(MitoCodeMcpToolCallError) as excinfo: + await client.call_tool("boom", {}) + assert "kaboom" in str(excinfo.value) + + +@pytest.mark.asyncio +async def test_connect_failure_is_recorded_not_raised(): + config = McpServerConfig.parse( + "broken", + {"url": "https://127.0.0.1:1/mcp", "connect_timeout": 3.0}, + ) + client = McpClient(config) + + assert await client.connect() is False + assert client.state is McpConnectionState.FAILED + assert client.last_error != "" + assert client.tools == [] + + +@pytest.mark.asyncio +async def test_connect_timeout_is_honoured(): + client = McpClient(_slow_config(1.0)) + + loop = asyncio.get_event_loop() + started = loop.time() + assert await client.connect() is False + assert loop.time() - started < 5.0 + + +@pytest.mark.asyncio +async def test_connect_timeout_force_closes_the_client(): + """A timed-out connect must not leave a subprocess running. + + Review Focus #6: FastMCP's StdioTransport spawns the process inside a + detached task, so cancelling our await alone orphans it. Only + Client.close() tears the process group down. + """ + client = McpClient(_slow_config(1.0)) + + class _HangingClient: + """Async context manager whose __aenter__ never completes.""" + + closed = False + + async def __aenter__(self): + await asyncio.Event().wait() + + async def __aexit__(self, *exc_info): + return False + + async def close(self): + type(self).closed = True + + built = _HangingClient() + client._build_fastmcp_client = lambda: built + + assert await client.connect() is False + assert built.closed is True + assert client._client is None + + +@pytest.mark.asyncio +async def test_connect_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.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 + + with pytest.raises(_ShutdownSignal): + await client.disconnect() + + assert built.closed is True + # Not recorded: the recording branch is Exception-only, so a signal leaves + # last_error untouched rather than looking like a cleanup problem. + assert client.last_error == "" + # The handles were dropped before the await, so abandoning the teardown + # part-way still leaks no reference to the transport. + assert client._client is None + assert client._stack is None + + +@pytest.mark.asyncio +async def test_force_close_records_a_failing_stack_close(): + """A stack that cannot unwind is recorded instead of aborting the teardown.""" + client = McpClient(_in_memory_config()) + recorder = _CloseRecorder() + client._client = recorder + stack = _StackRecorder(OSError("process group already gone")) + client._stack = stack + + await client.disconnect() + + assert ( + client.last_error == "stack close failed: OSError: process group already gone" + ) + assert recorder.closed is True + assert stack.closed is True + + +@pytest.mark.asyncio +async def test_force_close_keeps_both_errors_when_both_handles_fail(): + """Neither teardown error may be lost behind the other.""" + client = McpClient(_in_memory_config()) + client._client = _CloseRecorder(RuntimeError("client wedged")) + client._stack = _StackRecorder(OSError("stack wedged")) + + await client.disconnect() + + assert client.last_error == ( + "close failed: RuntimeError: client wedged; " + "stack close failed: OSError: stack wedged" + ) + + +@pytest.mark.asyncio +async def test_disconnect_stays_idempotent_after_a_failing_close(monkeypatch): + """A failing teardown still leaves the object clean enough to be retried.""" + client = McpClient(_in_memory_config()) + monkeypatch.setattr(client, "_build_fastmcp_client", lambda: Client(_server())) + assert await client.connect() is True + assert client.tools + + client._client = _CloseRecorder(RuntimeError("transport wedged")) + + await client.disconnect() + + assert client.state is McpConnectionState.DISCONNECTED + assert client._client is None + assert client._stack is None + assert client.tools == [] + assert "transport wedged" in client.last_error + + # The second teardown is a no-op, not a retry: the error stays recorded. + await client.disconnect() + + assert client._client is None + assert client._stack is None + assert client.state is McpConnectionState.DISCONNECTED + assert client.last_error == "close failed: RuntimeError: transport wedged" + + +@pytest.mark.asyncio +async def test_disconnect_on_a_never_connected_client_changes_nothing(): + """Tearing down a client that never opened must not invent an error.""" + client = McpClient(_in_memory_config()) + client._state = McpConnectionState.FAILED + client._last_error = "connect never got that far" + + await client.disconnect() + + assert client.state is McpConnectionState.FAILED + assert client.last_error == "connect never got that far" + assert client._client is None + assert client._stack is None + assert client.tools == [] + + +@pytest.mark.asyncio +async def test_connect_recovers_when_no_handle_was_ever_created(monkeypatch): + """A transport that cannot even be built leaves nothing to close, cleanly.""" + + def _explode(): + raise RuntimeError("bad transport arguments") + + client = McpClient(_in_memory_config()) + monkeypatch.setattr(client, "_build_fastmcp_client", _explode) + + assert await client.connect() is False + + assert client.state is McpConnectionState.FAILED + assert client.last_error == "RuntimeError: bad transport arguments" + assert client._client is None + assert client._stack is None + assert client.tools == [] + + +@pytest.mark.asyncio +async def test_connect_repeat_cancellation_still_completes_the_transport_teardown(): + """A second cancellation must not abort the teardown the first one started. + + ``connect`` runs its cleanup outside the timed-out await and shields it, so + a repeat cancellation -- a second Ctrl-C, or an enclosing ``wait_for`` + firing in the same tick -- cannot interrupt ``Client.close()`` midway and + leave stdio's detached connect task holding the child process alive. Only + the fact that ``close`` *ran to completion* distinguishes the shield from a + plain await; both leave the handles cleared. + """ + client = McpClient(_in_memory_config()) + built = _HandshakeHangingClient() + client._build_fastmcp_client = lambda: built + + await _cancel_twice_during_teardown(client, built) + + await asyncio.wait_for(built.close_finished.wait(), timeout=5) + assert client._client is None + assert client._stack is None + assert client.tools == [] + + +@pytest.mark.asyncio +async def test_connect_repeat_cancellation_skips_the_failed_bookkeeping(): + """A repeat cancellation propagates before ``connect`` records the failure. + + The cancellation is re-raised out of the shielded await, so the two + statements that would set ``FAILED`` and record the error never run. The + object is still consistent -- no client, no stack, no tools -- but a reader + of :attr:`last_error` must not expect to find the cancellation logged here. + """ + client = McpClient(_in_memory_config()) + built = _HandshakeHangingClient() + client._build_fastmcp_client = lambda: built + + await _cancel_twice_during_teardown(client, built) + + assert client.state is McpConnectionState.DISCONNECTED + assert client.last_error == "" + + +@pytest.mark.asyncio +async def test_connect_shutdown_signal_closes_the_transport_and_still_propagates(): + """A ``BaseException`` that is not an ``Exception`` must be cleaned up after. + + ``CancelledError`` is only the best-known member of that family. Catching + ``(Exception, CancelledError)`` would pass every cancellation test while + leaving an application shutdown signal to orphan the subprocess, so the + handler has to be keyed on ``BaseException`` itself. + """ + client = McpClient(_in_memory_config()) + + class _ShuttingDownClient: + """Client whose handshake fails with a shutdown signal.""" + + def __init__(self) -> None: + """Start with the transport still open.""" + self.closed = False + + async def __aenter__(self): + """Abort the handshake the way a framework shutdown would.""" + raise _ShutdownSignal("application is shutting down") + + async def __aexit__(self, *exc_info): + """Never claim to have suppressed anything.""" + return False + + async def close(self) -> None: + """Record that the transport was released.""" + self.closed = True + + built = _ShuttingDownClient() + client._build_fastmcp_client = lambda: built + + with pytest.raises(_ShutdownSignal): + await client.connect() + + assert built.closed is True, "shutdown signal orphaned the transport" + assert client.state is McpConnectionState.FAILED + assert client.last_error == "_ShutdownSignal: application is shutting down" + assert client._client is None + assert client._stack is None + assert client.tools == [] + + +@pytest.mark.asyncio +async def test_connect_failure_reports_the_handshake_error_not_the_teardown_error(): + """A failing teardown must not replace the error that explains the failure. + + ``_force_close`` records cleanup problems in ``last_error`` and ``connect`` + then overwrites it with the exception that caused the failure. The message a + user reads is therefore the handshake error alone; the fact that teardown + *also* failed is only visible in the teardown's own effects. + """ + client = McpClient(_in_memory_config()) + + class _ExplodingHandshake: + """Client that fails to connect and cannot be closed either.""" + + def __init__(self) -> None: + """Start with the transport still open.""" + self.closed = False + + async def __aenter__(self): + """Refuse the handshake.""" + raise RuntimeError("handshake exploded") + + async def __aexit__(self, *exc_info): + """Never claim to have suppressed anything.""" + return False + + async def close(self) -> None: + """Attempt the teardown, which fails.""" + self.closed = True + raise OSError("process group wedged") + + built = _ExplodingHandshake() + client._build_fastmcp_client = lambda: built + + assert await client.connect() is False + + assert built.closed is True + assert client.last_error == "RuntimeError: handshake exploded" + assert client.state is McpConnectionState.FAILED + + +@pytest.mark.asyncio +async def test_connect_after_a_failure_clears_the_stale_error(monkeypatch): + """A successful retry must not leave the previous failure in the report.""" + client = McpClient(_in_memory_config()) + + class _ExplodingHandshake: + """Client that never completes a handshake.""" + + async def __aenter__(self): + """Refuse the handshake.""" + raise RuntimeError("transient outage") + + async def __aexit__(self, *exc_info): + """Never claim to have suppressed anything.""" + return False + + async def close(self) -> None: + """Succeed at the teardown.""" + return None + + client._build_fastmcp_client = lambda: _ExplodingHandshake() + assert await client.connect() is False + assert client.last_error == "RuntimeError: transient outage" + + server = FastMCP("probe") + + @server.tool + def ping() -> str: + """Report that the server is alive.""" + return "pong" + + monkeypatch.setattr(client, "_build_fastmcp_client", lambda: Client(server)) + + assert await client.connect() is True + assert client.state is McpConnectionState.CONNECTED + assert client.last_error == "" + assert sorted(t.name for t in client.tools) == ["ping"] + await client.disconnect() + + +def test_redaction_cuts_the_payload_up_to_the_last_injected_type_marker(): + """A value carrying pydantic's own suffix text is cut up to that decoy. + + The pattern consumes a payload up to ``, input_type=``; when the argument + itself contains that text the greedy match runs to the *last* occurrence on + the line. 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:39-41): text after a payload-supplied + # decoy marker is currently NOT redacted -- `sk-PART2LEAK` survives here, + # and this assertion is what lets it. Fixing the leak requires DELETING + # this assertion, not editing it to expect the fix. + assert redacted.startswith("input_value=") + + +def test_redaction_keeps_the_report_after_an_injected_marker(): + """Surrounding report text must survive a payload that fakes the suffix. + + Redaction that eats the sentence, the field name and the expected type + would make the error undiagnosable -- and a test asserting only "the + secret is gone" would pass just as happily against such an implementation. + """ + redacted = _redact_arguments( + "tool 'login' failed on field 'token': " + 'input_value=sk-DECOY", input_type=str; expected a valid string]' + ) + + assert "sk-DECOY" not in redacted + assert redacted == ( + "tool 'login' failed on field 'token': " + "input_value=, input_type=str; expected a valid string]" + ) + + +def test_redaction_cuts_each_error_of_a_multiline_report_independently(): + """One payload must not eat the next error's report on a following line. + + ``.`` does not cross a newline, so a validation report naming two offending + fields loses neither of them -- and each keeps the type pydantic expected. + """ + report = ( + "2 validation errors for login\n" + "token\n" + " Input should be a valid string [type=string_type, " + "input_value='sk-FIRST', input_type=str]\n" + "count\n" + " Input should be a valid integer [type=string_type, " + "input_value='sk-SECOND', input_type=int]" + ) + + redacted = _redact_arguments(report) + + assert redacted == ( + "2 validation errors for login\n" + "token\n" + " Input should be a valid string [type=string_type, " + "input_value=, input_type=str]\n" + "count\n" + " Input should be a valid integer [type=string_type, " + "input_value=, input_type=int]" + ) + + +def test_redaction_handles_an_empty_payload(): + """An ``input_value=`` with nothing after it is still a payload.""" + assert _redact_arguments( + "Input should be a valid string [type=string_type, input_value=, input_type=str]" + ) == ( + "Input should be a valid string [type=string_type, " + "input_value=, input_type=str]" + ) + assert _redact_arguments("nothing here: input_value=") == ( + "nothing here: input_value=" + ) + + +def test_redaction_cuts_a_suffixless_payload_to_the_end_of_its_line(): + """With no ``, input_type=`` to stop at, the whole rest of the line is cut. + + Stopping at the next space instead would leave the tail of the value -- and + a token is often longer than its first word. + """ + assert _redact_arguments("input_value=sk-BARE extra words follow") == ( + "input_value=" + ) + assert _redact_arguments("input_value=sk-BARE extra words follow\nnext line") == ( + "input_value=\nnext line" + ) + + +@pytest.mark.asyncio +async def test_call_tool_tool_error_redacts_the_payload_and_keeps_the_report(): + """The common validation failure must lose the value and keep everything else. + + This is the ``ToolError`` branch, which the generic handler cannot see, so + it carries the full pydantic report: the tool, the server, the offending + field and the expected type all have to remain, while the rejected value + does not. + """ + client = McpClient(_in_memory_config()) + secret = "sk-ADVERSARIAL-DO-NOT-LEAK" + + class _ValidatingCall: + """Client that rejects its argument the way fastmcp reports it.""" + + async def call_tool(self, name, arguments, **kwargs): + raise ToolError( + "1 validation error for login\n" + "token\n" + " Input should be a valid string [type=string_type, " + f"input_value='{secret}\", input_type=str]" + ) + + client._client = _ValidatingCall() + + with pytest.raises(MitoCodeMcpToolCallError) as excinfo: + await client.call_tool("login", {"token": secret}) + + message = str(excinfo.value) + assert secret not in message, f"the rejected argument leaked: {message}" + assert "MCP tool 'login' on server 'probe' failed" in message + assert "token" in message + assert "should be a valid string" in message + assert "input_type=str" in message + + +@pytest.mark.asyncio +async def test_call_tool_unexpected_error_redacts_the_payload_and_keeps_the_report(): + """A local validation failure quotes the arguments just as a server one does. + + The generic handler renders the exception type before the text, so the + redaction has to survive an extra prefix and must not cost the surrounding + report anything. + """ + client = McpClient(_in_memory_config()) + secret = "sk-LOCAL-VALIDATION-DO-NOT-LEAK" + + class _LocallyValidatingCall: + """Client that fails locally while validating its arguments.""" + + async def call_tool(self, name, arguments, **kwargs): + raise ValueError( + "field 'token' rejected\n" + " Input should be a valid string [type=string_type, " + f"input_value='{secret}', input_type=str]" + ) + + client._client = _LocallyValidatingCall() + + with pytest.raises(MitoCodeMcpToolCallError) as excinfo: + await client.call_tool("login", {"token": secret}) + + message = str(excinfo.value) + assert secret not in message, f"the rejected argument leaked: {message}" + assert "MCP tool 'login' on server 'probe' raised ValueError" in message + assert "token" in message + assert "should be a valid string" in message + assert "input_type=str" in message + + +@pytest.mark.asyncio +async def test_call_tool_forwards_the_configured_timeout_to_the_server_call(): + """The per-call deadline must reach the server call, not just the local wait. + + The deadline is applied twice on purpose: the outer ``wait_for`` bounds the + local await and the ``timeout`` keyword bounds the request the transport has + actually sent. Dropping the keyword leaves a hanging server unbounded + server-side while every local test still passes. + """ + client = McpClient( + McpServerConfig.parse( + "slowtool", {"url": "https://x.invalid/mcp", "timeout": 12.5} + ) + ) + seen = {} + + class _RecordingCall: + """Client that records how it was called.""" + + async def call_tool(self, name, arguments, **kwargs): + seen["name"] = name + seen["arguments"] = arguments + seen["timeout"] = kwargs.get("timeout", "absent") + return None + + client._client = _RecordingCall() + + await client.call_tool("echo", {"text": "hi"}) + + assert seen == {"name": "echo", "arguments": {"text": "hi"}, "timeout": 12.5} + + +@pytest.mark.asyncio +async def test_call_tool_keeps_non_ascii_structured_data_readable(): + """Structured output must not be escaped into ``\\uXXXX`` noise. + + ``ensure_ascii`` defaults to true, which turns a Japanese or accented field + into escape sequences; a result rendered that way still parses, so nothing + downstream notices except the reader. + """ + client = McpClient(_in_memory_config()) + + class _Result: + data = {"memo": "café", "city": "東京"} + content = [] + + class _Call: + async def call_tool(self, name, arguments, **kwargs): + return _Result() + + client._client = _Call() + + assert await client.call_tool("echo", {}) == '{"memo": "café", "city": "東京"}' + + +@pytest.mark.asyncio +async def test_call_tool_renders_an_unserialisable_result_value_with_str(): + """A result carrying a non-JSON value must render, not raise. + + ``json.dumps`` refuses a set, a handle or a custom object, and the failure + would escape ``call_tool`` as a bare ``TypeError`` because rendering happens + outside the error handling. ``default=str`` is what keeps that reachable. + """ + client = McpClient(_in_memory_config()) + + class _Handle: + """A value the model cannot receive as JSON.""" + + def __str__(self) -> str: + return "opaque-handle" + + class _Result: + data = {"handle": _Handle(), "ids": 3} + content = [] + + class _Call: + async def call_tool(self, name, arguments, **kwargs): + return _Result() + + client._client = _Call() + + assert await client.call_tool("echo", {}) == ( + '{"handle": "opaque-handle", "ids": 3}' + ) + + +def test_redaction_cuts_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..251a3d4 --- /dev/null +++ b/test/test_mcp_config.py @@ -0,0 +1,745 @@ +"""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. +""" + +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..487b807 --- /dev/null +++ b/test/test_mcp_manager.py @@ -0,0 +1,1308 @@ +import asyncio +import os +import subprocess +import sys +from types import SimpleNamespace + +import pytest + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) + +from mitocode.backend.exception.mitocode_exception import MitoCodeMcpConfigError +from mitocode.backend.mcp.mcp_client import McpClient, McpConnectionState +from mitocode.backend.mcp.mcp_config import McpServerConfig, 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 + + async def connect(self): + """Move to CONNECTED or FAILED according to the scripted outcome. + + One-shot cancellation is honoured first, so a test can arm it *after* + a successful ``start`` rather than making every later connect fail. + + Returns: + The scripted outcome. + + Raises: + asyncio.CancelledError: If :meth:`cancel_next_connect` armed this + attempt, which is how ``McpClient.connect`` behaves when the + caller cancels it. + """ + if self._cancel_next: + self._cancel_next = False + raise asyncio.CancelledError() + self.connected = self._succeed + if self._succeed: + self.state = McpConnectionState.CONNECTED + 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 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 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=""): + """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 ``""``. + """ + super().__init__(config, succeed=succeed) + self._fail_on = fail_on + + async def connect(self): + """Raise instead of connecting when this client is armed for it.""" + if self._fail_on == "connect": + raise RuntimeError("handshake exploded") + return await super().connect() + + async def disconnect(self): + """Raise instead of closing when this client is armed for it.""" + if self._fail_on == "disconnect": + raise RuntimeError("teardown exploded") + await super().disconnect() + + +def _factory(script, 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 _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.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(monkeypatch): + """A second ``start`` must not forget a tool the registry refused to remove. + + ``register_mcp_client`` skips a name that is still registered, so its + result omits the stuck tool. Assigning that result to the ownership record + would forget the tool while it is still in the registry — dispatchable, + counted nowhere, and reclaimable by nobody. + """ + tool_manager = ToolManager() + real_unregister = tool_manager.unregister + armed = {"bad": False} + + def _unregister(name): + if armed["bad"] and name == "mcp__weather__x": + raise RuntimeError("registry exploded") + return real_unregister(name) + + monkeypatch.setattr(tool_manager, "unregister", _unregister) + factory = _factory({"weather": (True, ("x",))}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + armed["bad"] = True + + await manager.start() + + assert ( + manager.status()["weather"].tool_count == 1 + ), "the tool is still in the registry, so status() must still count it" + assert "registry exploded" in manager.status()["weather"].error + + armed["bad"] = False + await manager.stop() + + assert ( + tool_manager.tools_map() == {} + ), "a later stop() must reclaim the tool nobody forgot" + + +@pytest.mark.asyncio +async def test_start_owns_the_tools_a_refresh_could_not_unregister(monkeypatch): + """The same guarantee holds on ``refresh``, the other path through it.""" + tool_manager = ToolManager() + real_unregister = tool_manager.unregister + armed = {"bad": False} + + def _unregister(name): + if armed["bad"] and name == "mcp__weather__a": + raise RuntimeError("registry exploded") + return real_unregister(name) + + monkeypatch.setattr(tool_manager, "unregister", _unregister) + factory = _factory({"weather": (True, ("a", "b"))}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + armed["bad"] = True + + await manager.refresh("weather") + + assert ( + manager.status()["weather"].tool_count == 2 + ), "'a' is stuck in the registry and 'b' was registered, so both are owned" + + armed["bad"] = False + await manager.stop() + + assert tool_manager.tools_map() == {} + + +@pytest.mark.asyncio +async def test_stop_completes_both_phases_and_then_propagates_a_cancellation( + monkeypatch, +): + """A cancelled unregistration must not skip the disconnections. + + The two phases run one after the other, so re-raising as soon as the first + one is cancelled would leave every server connected — the outcome the + per-server isolation exists to prevent, reached through the propagation path. + """ + tool_manager = ToolManager() + factory = _factory({"bad": (True, ("x",)), "good": (True, ("y",))}) + manager = McpManager( + _config( + ("bad", {"url": "https://bad.invalid/mcp"}), + ("good", {"url": "https://good.invalid/mcp"}), + ), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + + real_unregister = tool_manager.unregister + + def _unregister(name): + if name == "mcp__bad__x": + raise asyncio.CancelledError() + return real_unregister(name) + + monkeypatch.setattr(tool_manager, "unregister", _unregister) + + with pytest.raises(asyncio.CancelledError): + await manager.stop() + + assert factory.made["good"].state is McpConnectionState.DISCONNECTED + assert ( + factory.made["bad"].state is McpConnectionState.DISCONNECTED + ), "every connection was closed before the cancellation was re-raised" + assert manager.status()["bad"].error == "CancelledError" + + +@pytest.mark.asyncio +async def test_refresh_records_a_cancellation_and_lets_it_out(): + """``refresh`` propagates a cancellation for the same reason ``start`` does.""" + tool_manager = ToolManager() + factory = _rude_factory({"bad": "", "good": ""}) + manager = McpManager( + _config( + ("bad", {"url": "https://bad.invalid/mcp"}), + ("good", {"url": "https://good.invalid/mcp"}), + ), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + factory.made["bad"].cancel_next_connect() + + with pytest.raises(asyncio.CancelledError): + await manager.refresh("bad") + + assert manager.status()["bad"].error == "CancelledError" + assert ( + factory.made["good"].state is McpConnectionState.CONNECTED + ), "refreshing one server left the others alone" + + +@pytest.mark.asyncio +async def test_stop_disconnects_every_server_when_one_teardown_raises(): + """One failing teardown must not strand the other connections open.""" + 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(monkeypatch): + """A registry that refuses to give a tool back must not strand a connection. + + Unregistration runs before the disconnects, so an unguarded failure here + would skip the teardown entirely and leave every server connected. + """ + tool_manager = ToolManager() + real_unregister = tool_manager.unregister + armed = {"bad": False} + + def _unregister(name): + if armed["bad"] and name == "mcp__bad__x": + raise RuntimeError("registry exploded") + return real_unregister(name) + + monkeypatch.setattr(tool_manager, "unregister", _unregister) + factory = _factory({"bad": (True, ("x",)), "good": (True, ("y",))}) + manager = McpManager( + _config( + ("bad", {"url": "https://bad.invalid/mcp"}), + ("good", {"url": "https://good.invalid/mcp"}), + ), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + armed["bad"] = True + + await manager.stop() + + assert ( + factory.made["good"].state is McpConnectionState.DISCONNECTED + ), "the healthy server was torn down" + assert factory.made["bad"].state is McpConnectionState.DISCONNECTED + assert set(tool_manager.tools_map()) == { + "mcp__bad__x" + }, "only the tool the registry refused to remove is left" + assert "registry exploded" in manager.status()["bad"].error + + +@pytest.mark.asyncio +async def test_stop_retries_an_unregister_that_previously_failed(monkeypatch): + """A failed removal stays owned, so a later stop() can finish the job. + + The manager pops its ownership record only once a removal is confirmed; + forgetting a tool it never removed would leave it dispatchable forever. + """ + tool_manager = ToolManager() + real_unregister = tool_manager.unregister + armed = {"bad": False} + + def _unregister(name): + if armed["bad"] and name == "mcp__bad__x": + raise RuntimeError("registry exploded") + return real_unregister(name) + + monkeypatch.setattr(tool_manager, "unregister", _unregister) + factory = _factory({"bad": (True, ("x",))}) + manager = McpManager( + _config(("bad", {"url": "https://bad.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + armed["bad"] = True + await manager.stop() + assert manager.status()["bad"].tool_count == 1, "still owned, so still counted" + + armed["bad"] = False + await manager.stop() + + assert tool_manager.tools_map() == {} + assert manager.status()["bad"].tool_count == 0 + assert ( + manager.status()["bad"].error == "" + ), "a failure the retry resolved is history, not a current error" + + +@pytest.mark.asyncio +async def test_stop_forgets_a_tool_that_was_already_removed(monkeypatch): + """``unregister`` reporting "nothing was there" ends the manager's claim. + + A tool someone else removed from the registry cannot be removed again, so + keeping it on the ownership record would make every later stop() retry a + removal that can never succeed. + """ + tool_manager = ToolManager() + factory = _factory({"weather": (True, ("a",))}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + assert ( + tool_manager.unregister("mcp__weather__a") is True + ), "removed behind the manager's back" + + await manager.stop() + + assert manager.status()["weather"].tool_count == 0 + assert tool_manager.tools_map() == {} + + +@pytest.mark.asyncio +async def test_status_prefers_a_recorded_failure_over_a_stale_client_message(): + """The newest reason wins when the manager and the client both have one. + + The client's message describes the connect that failed earlier; the + recorded failure describes the teardown that just failed, which is the one + an operator can act on. + """ + tool_manager = ToolManager() + factory = _rude_factory({"weather": "disconnect"}, succeed={"weather": False}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + assert manager.status()["weather"].error == "connection refused" + + await manager.stop() + + assert manager.status()["weather"].error == ( + "RuntimeError: teardown exploded" + ), "the teardown failure is the most recent one" + + +@pytest.mark.asyncio +async def test_start_twice_then_stop_leaves_no_registrations(): + """A repeated start must not orphan the tools the first one registered. + + ``ToolManager.register_mcp_client`` never overwrites an existing name, so + a second registration pass registers nothing and empties the manager's + record of what it owns. ``stop()`` would then have no name left to + unregister and the first pass's tools would outlive the session. + """ + tool_manager = ToolManager() + factory = _factory({"weather": (True, ("a",))}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + + await manager.start() + await manager.start() + assert list(tool_manager.tools_map()) == ["mcp__weather__a"] + + await manager.stop() + + assert ( + tool_manager.tools_map() == {} + ), "every tool a repeated start registered must still be unregistered" + assert manager.status()["weather"].tool_count == 0 + + +@pytest.mark.asyncio +async def test_refresh_drops_tools_the_server_no_longer_advertises(): + """Refreshing must pick up the whole catalogue, not only its additions.""" + tool_manager = ToolManager() + factory = _factory({"weather": (True, ("a", "b"))}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + assert set(tool_manager.tools_map()) == {"mcp__weather__a", "mcp__weather__b"} + + factory.made["weather"].tools.pop() + await manager.refresh("weather") + + assert set(tool_manager.tools_map()) == { + "mcp__weather__a" + }, "a tool the server dropped must not stay dispatchable to a dead client" + + +@pytest.mark.asyncio +async def test_refresh_drops_stale_tools_when_the_reconnect_fails(): + """A failed refresh leaves that server with no tools, not with dead ones. + + The tools must be unregistered *before* the reconnect is attempted: kept + across a failure they would advertise a capability no model can call, and + would dispatch into a client that is not connected. + """ + tool_manager = ToolManager() + factory = _factory({"weather": (True, ("a", "b"))}) + manager = McpManager( + _config(("weather", {"url": "https://w.invalid/mcp"})), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + assert set(tool_manager.tools_map()) == {"mcp__weather__a", "mcp__weather__b"} + + factory.made["weather"].go_offline() + await manager.refresh("weather") + + assert ( + tool_manager.tools_map() == {} + ), "a server that could not be reached must publish nothing" + assert manager.status()["weather"].state == "failed" + + +@pytest.mark.asyncio +async def test_refresh_without_a_name_reconnects_every_server(): + tool_manager = ToolManager() + factory = _factory({"one": (True, ("a",)), "two": (True, ("b",))}) + manager = McpManager( + _config( + ("one", {"url": "https://one.invalid/mcp"}), + ("two", {"url": "https://two.invalid/mcp"}), + ), + tool_manager=tool_manager, + client_factory=factory, + ) + await manager.start() + + await manager.refresh() + + assert factory.made["one"].state is McpConnectionState.CONNECTED + assert factory.made["two"].state is McpConnectionState.CONNECTED + assert set(tool_manager.tools_map()) == {"mcp__one__a", "mcp__two__b"} + + +@pytest.mark.asyncio +async def test_refresh_of_a_disabled_server_is_rejected_by_name(): + """A disabled server has no client, so the rejection must say why.""" + manager = McpManager( + _config(("off", {"url": "https://off.invalid/mcp", "enabled": False})), + tool_manager=ToolManager(), + client_factory=_factory({}), + ) + + with pytest.raises(MitoCodeMcpConfigError) as excinfo: + await manager.refresh("off") + + assert "off" in str(excinfo.value) + assert "disabled" in str(excinfo.value) + + +@pytest.mark.asyncio +async def test_refresh_error_names_the_configured_servers(): + manager = McpManager( + _config( + ("weather", {"url": "https://w.invalid/mcp"}), + ("files", {"url": "https://f.invalid/mcp"}), + ), + tool_manager=ToolManager(), + client_factory=_factory({}), + ) + + with pytest.raises(MitoCodeMcpConfigError) as excinfo: + await manager.refresh("nope") + + message = str(excinfo.value) + assert "nope" in message + assert "weather" in message and "files" in message + + +@pytest.mark.asyncio +async def test_refresh_error_does_not_echo_server_credentials(): + """The message names a server, never the configuration behind it.""" + manager = McpManager( + _config( + ( + "files", + { + "url": "https://files.invalid/mcp", + "headers": {"Authorization": "Bearer s3cr3t-token"}, + }, + ) + ), + tool_manager=ToolManager(), + client_factory=_factory({}), + ) + + with pytest.raises(MitoCodeMcpConfigError) as excinfo: + await manager.refresh("nope") + + assert "s3cr3t-token" not in str(excinfo.value) + + +def test_a_list_of_configs_is_accepted_as_the_configuration(): + """Configurations normalised elsewhere need no round trip through JSON.""" + tool_manager = ToolManager() + factory = _factory({"weather": (True, ("a",))}) + manager = McpManager( + [McpServerConfig.parse("weather", {"url": "https://w.invalid/mcp"})], + tool_manager=tool_manager, + client_factory=factory, + ) + + assert "weather" in factory.made + assert manager.status()["weather"].state == "disconnected" + assert manager.status()["weather"].tool_count == 0 + + +def test_a_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. + """ + server_script = tmp_path / "srv.py" + server_script.write_text( + "from fastmcp import FastMCP\n" + "app = FastMCP('local')\n" + "@app.tool\n" + "def shout(word: str) -> str:\n" + " '''Shout a word.'''\n" + " return word.upper()\n" + "app.run()\n" + ) + + tool_manager = ToolManager() + manager = McpManager( + { + "mcpServers": { + "local": { + "command": sys.executable, + "args": [str(server_script)], + "connect_timeout": 60.0, + } + } + }, + tool_manager=tool_manager, + ) + + await manager.start() + try: + status = manager.status()["local"] + assert status.state == "connected", status.error + assert status.tool_count == 1 + assert "mcp__local__shout" in tool_manager.tools_map() + + result = await tool_manager.tools_map()["mcp__local__shout"](word="hey") + assert result == "HEY" + finally: + await manager.stop() + + assert tool_manager.tools_map() == {} + assert manager.status()["local"].state == "disconnected" diff --git a/test/test_tool_manager.py b/test/test_tool_manager.py index f5e1887..1187a4d 100644 --- a/test/test_tool_manager.py +++ b/test/test_tool_manager.py @@ -1,18 +1,114 @@ +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 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 ( + 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 +198,1054 @@ 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 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. That makes + the tier advertised but unreachable through the public API. + """ + manager = ToolManager() + from_meta = {"type": "object", "properties": {"remote": {"type": "string"}}} + + @declare_tool(name="remote_call") + def remote_call(**kwargs): + return kwargs + + remote_call.__mitocode_tool_meta__["parameters_schema"] = from_meta + spec = manager.register(remote_call) + + assert ( + spec.parameters_schema == from_meta + ), "the metadata schema must outrank the derived empty one" + assert ( + "remote" in spec.parameters_schema["properties"] + ), "the server's shape survives" + + +def test_register_prefers_the_explicit_parameters_schema_over_declare_tool_metadata(): + """Tier one beats tier two: the explicit argument wins over the metadata.""" + + @declare_tool(name="tiered") + def tiered(**kwargs): + return kwargs + + tiered.__mitocode_tool_meta__["parameters_schema"] = { + "type": "object", + "properties": {"from_meta": {"type": "string"}}, + } + explicit = {"type": "object", "properties": {"explicit": {"type": "boolean"}}} + + manager = ToolManager() + spec = manager.register(tiered, parameters_schema=explicit) + + assert spec.parameters_schema == explicit, "the explicit argument is the winner" + assert ( + "from_meta" not in spec.parameters_schema["properties"] + ), "the metadata tier must not be consulted at all" + + +def test_register_keeps_an_explicitly_empty_parameters_schema_instead_of_deriving_one(): + """An empty dict is a supplied schema, not an absent one. + + ``parameters_schema={}`` is how a caller says "this tool takes no + arguments"; treating it as absent would advertise a derived schema instead. + """ + + def ping() -> str: + return "pong" + + manager = ToolManager() + spec = manager.register(ping, parameters_schema={}) + + assert spec.parameters_schema == {}, "an explicit {} must be kept verbatim" + + +def test_register_replaces_a_registration_that_already_uses_the_name(): + """Registering the same name again overwrites rather than duplicating.""" + + def first(): + return "first" + + def second(): + return "second" + + manager = ToolManager() + manager.register(first, name="shared", agent_doc="first doc") + spec = manager.register(second, name="shared", agent_doc="second doc") + + assert list(manager.tools_map()) == ["shared"], "one name, one entry" + assert manager.tools_map()["shared"] is second, "the newest registration wins" + assert spec.agent_doc == "second doc", "the description is replaced with it" + + +def test_register_defaults_both_display_hints_to_empty_strings(): + """Display hints are opt-in; without them the spec carries empty strings.""" + + def plain(text: str): + return text + + manager = ToolManager() + spec = manager.register(plain) + + assert spec.display_action == "", "no display action was requested" + assert spec.display_param == "", "no display parameter was requested" + + +def test_tool_decorator_returns_the_decorated_function_unchanged(): + """The decorator registers the function and hands the same object back.""" + + def ping(): + return None + + manager = ToolManager() + decorated = manager.tool(agent_doc="doc")(ping) + + assert decorated is ping, "registration must not wrap or replace the callable" + + +def test_tool_decorator_forwards_an_explicit_parameters_schema_to_register(): + """The decorator must pass ``parameters_schema`` through, not swallow it.""" + manager = ToolManager() + schema = {"type": "object", "properties": {"city": {"type": "string"}}} + + @manager.tool(name="get_weather", parameters_schema=schema) + def get_weather(**kwargs): + return kwargs + + assert ( + manager.get_spec("get_weather").parameters_schema == schema + ), "the decorator forwards the schema to register()" + + +def test_get_spec_returns_none_for_a_name_that_was_never_registered(): + """An unknown name has no spec rather than raising.""" + manager = ToolManager() + assert manager.get_spec("absent") is None, "a missing tool reports None" + + +def test_get_spec_returns_the_same_object_register_returned(): + """``get_spec`` exposes the registered object, not a copy of it.""" + + def ping(): + return None + + manager = ToolManager() + spec = manager.register(ping, name="ping") + + assert manager.get_spec("ping") is spec, "callers must see the stored spec" + + +def test_tool_schemas_defaults_to_the_agent_audience(): + """The default audience is the agent's.""" + manager = ToolManager() + + @manager.tool(human_doc="human only", agent_doc="agent only") + def ping(): + return None + + assert manager.tool_schemas() == manager.tool_schemas( + "agent" + ), "the default audience is the agent's" + assert ( + manager.tool_schemas()[0]["function"]["description"] == "agent only" + ), "the agent's description is served by default" + + +def test_tool_schemas_serves_human_docs_for_the_human_audience(): + """``audience="human"`` switches the description to the human one.""" + manager = ToolManager() + + @manager.tool(human_doc="human only", agent_doc="agent only") + def ping(): + return None + + assert manager.tool_schemas("human")[0]["function"]["description"] == ( + "human only" + ), "the human audience gets the human description" + + +def test_tool_schemas_treats_an_unknown_audience_as_the_agent_one(): + """Only the literal ``"human"`` switches audience.""" + manager = ToolManager() + + @manager.tool(human_doc="human only", agent_doc="agent only") + def ping(): + return None + + assert manager.tool_schemas("operator")[0]["function"]["description"] == ( + "agent only" + ), "an unrecognised audience is served the agent description" + + +def test_register_mcp_client_registers_nothing_from_an_empty_catalogue(): + """A client that discovered no tools registers nothing and warns about nothing.""" + manager = ToolManager() + client = _fake_mcp_client("files", []) + + with mock.patch("mitocode.backend.tool.tool_manager.logger") as fake_logger: + specs = manager.register_mcp_client(client) + + assert specs == [], "an empty catalogue yields no specs" + assert manager.tools_map() == {}, "and leaves the registry untouched" + assert 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" + + +@pytest.mark.asyncio +async def test_register_mcp_client_keeps_the_first_of_two_clients_with_one_server_name(): + """A second client on the same server name never receives the dispatch.""" + first_calls = [] + second_calls = [] + first = _fake_mcp_client( + "files", [_fake_tool("read", description="from first")], first_calls + ) + second = _fake_mcp_client( + "files", [_fake_tool("read", description="from second")], second_calls + ) + manager = ToolManager() + manager.register_mcp_client(first) + manager.register_mcp_client(second) + + await manager.tools_map()["mcp__files__read"](path="a.txt") + + assert first_calls == [ + ("read", {"path": "a.txt"}) + ], "the first client keeps ownership of the composed name" + assert second_calls == [], "the second client must never be dispatched to" + + +def test_register_mcp_client_synthesizes_an_exact_description_for_a_blank_description(): + """A whitespace-only description is replaced by the documented sentence.""" + + manager = ToolManager() + client = _fake_mcp_client("files", [_fake_tool("read", description=" \n\t ")]) + + manager.register_mcp_client(client) + + assert manager.human_docs()["mcp__files__read"] == ( + "[files] Tool 'read' provided by MCP server 'files'." + ), "the synthesised description names the tool and the server" + assert manager.agent_docs()["mcp__files__read"] == ( + "[files] Tool 'read' provided by MCP server 'files'." + ), "both audiences get the same text" + + +def test_register_mcp_client_strips_surrounding_whitespace_from_a_description(): + """A real description keeps its content but loses its padding.""" + + manager = ToolManager() + client = _fake_mcp_client( + "files", [_fake_tool("read", description=" Read a file.\n")] + ) + + manager.register_mcp_client(client) + + assert manager.agent_docs()["mcp__files__read"] == ( + "[files] Read a file." + ), "the server's own description is used, trimmed" + + +def test_register_mcp_client_leaves_the_display_param_empty(): + """Only a display action is synthesised; there is no parameter to highlight.""" + + manager = ToolManager() + client = _fake_mcp_client("files", [_fake_tool("read", description="d")]) + + spec = manager.register_mcp_client(client)[0] + + assert ( + spec.display_param == "" + ), "the module does not know which argument to highlight" + + +def test_register_mcp_client_preserves_a_complete_input_schema_verbatim(): + """A server that advertises a full schema gets it back unchanged.""" + manager = ToolManager() + advertised = { + "type": "object", + "properties": {"path": {"type": "string"}, "limit": {"type": "integer"}}, + "required": ["path"], + "additionalProperties": False, + } + client = _fake_mcp_client("files", [_fake_tool("read", input_schema=advertised)]) + + manager.register_mcp_client(client) + + schema = manager.tool_schemas("agent")[0]["function"]["parameters"] + assert schema == advertised, "nothing is added to or removed from a complete schema" + + +def test_register_mcp_client_replaces_non_dict_properties_in_an_input_schema(): + """A ``properties`` value that is not a dict is replaced by an empty dict.""" + manager = ToolManager() + client = _fake_mcp_client( + "files", + [_fake_tool("read", input_schema={"type": "object", "properties": "city"})], + ) + + manager.register_mcp_client(client) + + schema = manager.tool_schemas("agent")[0]["function"]["parameters"] + assert schema["properties"] == {}, "the unusable value is dropped" + assert schema["type"] == "object", "a usable type survives" + + +@pytest.mark.asyncio +async def test_registered_mcp_tool_returns_the_server_result_unchanged(): + """The generated callable hands the server's return value straight back.""" + + manager = ToolManager() + client = _fake_mcp_client("files", [_fake_tool("read", description="r")]) + manager.register_mcp_client(client) + + result = await manager.tools_map()["mcp__files__read"](path="a.txt", limit=3) + + assert ( + result == "read got ['limit', 'path']" + ), "the server's string is passed through" + + +def test_registered_mcp_tool_schema_hides_the_private_dispatch_bindings(): + """``_client``/``_tool_name`` exist on the callable but not on the wire.""" + manager = ToolManager() + client = _fake_mcp_client( + "files", + [ + _fake_tool( + "read", + input_schema={ + "type": "object", + "properties": {"path": {"type": "string"}}, + }, + ) + ], + ) + + spec = manager.register_mcp_client(client)[0] + + assert set(spec.parameters_schema["properties"]) == { + "path" + }, "the advertised arguments are exactly the server's" + assert ( + "_client" not in spec.parameters_schema["properties"] + ), "an internal binding must never be advertised to a model" + assert ( + "_tool_name" not in spec.parameters_schema["properties"] + ), "an internal binding must never be advertised to a model" + + +@pytest.mark.asyncio +async def test_registered_mcp_tool_exposes_no_dispatch_bindings_as_keywords(): + """The generated callable must accept nothing but the model's own arguments.""" + manager = ToolManager() + client = _fake_mcp_client("files", [_fake_tool("read", description="r")]) + manager.register_mcp_client(client) + + parameters = inspect.signature(manager.tools_map()["mcp__files__read"]).parameters + + assert list(parameters) == [ + "kwargs" + ], "the client and upstream tool name must be captured, never bound as keywords" + + +@pytest.mark.xfail( + reason="defect: _invoke's defaults shared a namespace with **kwargs", + strict=False, +) +@pytest.mark.asyncio +async def test_registered_mcp_tool_cannot_redirect_the_dispatch_with_a_tool_name_argument(): + """A ``_tool_name`` argument must be forwarded, never bind over the captured name. + + FIXME(tool_manager): the defect was ``async def _invoke(_client=client, + _tool_name=tool.name, **kwargs)`` — the captured values sat in the same + namespace as the model's keywords, so this argument silently redirected the + call to a *different* tool on that server, bypassing the name-collision guard + entirely. Fixed by ``_make_invoke``; the ``xfail`` stays non-strict so the + suite is green both ways and can be promoted to a plain assertion once the + closure factory has shipped everywhere. + """ + calls = [] + manager = ToolManager() + client = _fake_mcp_client("files", [_fake_tool("read", description="r")], calls) + manager.register_mcp_client(client) + + 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.xfail( + reason="defect: _invoke's defaults shared a namespace with **kwargs", + strict=False, +) +@pytest.mark.asyncio +async def test_registered_mcp_tool_cannot_redirect_the_dispatch_with_a_client_argument(): + """A ``_client`` argument must be forwarded, never replace the captured client. + + FIXME(tool_manager): same namespace collision as ``_tool_name``. The + defect raised ``AttributeError`` from the model's own argument before the + server was ever reached; this asserts the dispatch contract instead of that + incidental error message, so it also accepts a clearer error from a fix. + """ + calls = [] + manager = ToolManager() + client = _fake_mcp_client("files", [_fake_tool("read", description="r")], calls) + manager.register_mcp_client(client) + + 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..a828a61 100644 --- a/todo.md +++ b/todo.md @@ -11,7 +11,7 @@ - [x] 重构Agent对外接口 (Agent → AgentSession/AgentRuntime/AgentOrchestrator 三层分离) - [ ] 添加上下文管理功能 - [ ] 添加索引功能 -- [ ] 添加MCP功能 +- [x] 添加MCP功能 - [ ] 添加Agent Skills功能 - [x] 完成TUI (TUI直接调用backend，移除HTTP/SSE层) - [x] 添加工具管理功能:PR #6