""" Copyright 2024-2026 ChatterMate Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. """ """ ChatterMate - MCP Manager Tests """ import pytest import asyncio from unittest.mock import patch, AsyncMock, MagicMock from time import monotonic from uuid import uuid4 from app.models.mcp_tool import MCPTransportType from app.tools.mcp_connection import CONNECTION_FAILED, HeldConnection from app.tools.mcp_manager import ( CHAT_CONNECT_BUDGET_SECONDS, CONNECTOR_GUIDANCE_TEMPLATE, MAX_GUIDANCE_FUNCTIONS, CONNECT_TIMEOUT_MARGIN_SECONDS, DEFAULT_TIMEOUT_SECONDS, HANDSHAKE_REQUESTS, MCPToolsManager, TEST_CONNECT_BUDGET_SECONDS, _connect_error_message, _connect_timeout, _session_timeout, cleanup_mcp_tools, initialize_mcp_tools, ) @pytest.mark.asyncio async def test_initialize_mcp_tools_no_ids(): manager = MCPToolsManager() tools = await manager.initialize_mcp_tools(agent_id="", org_id="") assert tools == [] @pytest.mark.asyncio async def test_initialize_mcp_tools_stdio_success(): manager = MCPToolsManager() agent_id = str(uuid4()) org_id = str(uuid4()) # Mock repo to return a single STDIO tool config mock_tool_config = MagicMock() mock_tool_config.name = "FS Tool" mock_tool_config.transport_type = type("T", (), {"__eq__": lambda s, o: False})() # placeholder from app.models.mcp_tool import MCPTransportType mock_tool_config.transport_type = MCPTransportType.STDIO mock_tool_config.timeout = None mock_tool_config.command = "npx" mock_tool_config.args = ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"] mock_tool_config.env_vars = {"FOO": "bar"} with patch("app.tools.mcp_manager.SessionLocal") as mock_sess_local, \ patch("app.tools.mcp_manager.MCPToolRepository") as mock_repo_cls, \ patch("app.tools.mcp_manager.MCPTools") as mock_mcp_tools_cls: mock_db = MagicMock() mock_sess_local.return_value.__enter__.return_value = mock_db mock_repo = MagicMock() mock_repo.get_agent_mcp_tools.return_value = [mock_tool_config] mock_repo_cls.return_value = mock_repo # Mock MCPTools async context manager mock_mcp_instance = AsyncMock() # Simulate that functions are available after connect mock_mcp_instance.functions = {"list": MagicMock()} mock_mcp_tools_cls.return_value = mock_mcp_instance tools = await manager.initialize_mcp_tools(agent_id, org_id) assert len(tools) == 1 mock_repo.get_agent_mcp_tools.assert_called_once_with(agent_id) # __aenter__ called to connect mock_mcp_instance.__aenter__.assert_awaited() @pytest.mark.asyncio async def test_initialize_mcp_tools_stdio_missing_dirs_skips(): manager = MCPToolsManager() agent_id = str(uuid4()) org_id = str(uuid4()) # Tool without directories and without env var should be skipped mock_tool_config = MagicMock() from app.models.mcp_tool import MCPTransportType mock_tool_config.transport_type = MCPTransportType.STDIO mock_tool_config.timeout = None mock_tool_config.name = "FS Tool" mock_tool_config.command = "npx" mock_tool_config.args = ["-y", "@modelcontextprotocol/server-filesystem"] # no dirs after package mock_tool_config.env_vars = {} # no ALLOWED_DIRECTORIES with patch("app.tools.mcp_manager.SessionLocal") as mock_sess_local, \ patch("app.tools.mcp_manager.MCPToolRepository") as mock_repo_cls, \ patch("app.tools.mcp_manager.MCPTools") as mock_mcp_tools_cls: mock_db = MagicMock() mock_sess_local.return_value.__enter__.return_value = mock_db mock_repo = MagicMock() mock_repo.get_agent_mcp_tools.return_value = [mock_tool_config] mock_repo_cls.return_value = mock_repo tools = await manager.initialize_mcp_tools(agent_id, org_id) assert tools == [] # skipped mock_mcp_tools_cls.assert_not_called() @pytest.mark.asyncio async def test_cleanup_mcp_tools_closes_every_connected_tool(): manager = MCPToolsManager() tool1 = _connected_instance({"a": MagicMock()}) tool2 = _connected_instance({"b": MagicMock()}) with patch("app.tools.mcp_manager.MCPTools", side_effect=[tool1, tool2]): await manager._initialize_from_configs([_stdio_config(name="One"), _stdio_config(name="Two")]) await manager.cleanup_mcp_tools() tool1.__aexit__.assert_awaited() tool2.__aexit__.assert_awaited() assert manager.mcp_tools == [] @pytest.mark.asyncio async def test_initialize_and_cleanup_helpers(): # High level helper functions with patch("app.tools.mcp_manager.MCPToolsManager") as mock_mgr_cls: mock_mgr = AsyncMock() mock_mgr.initialize_mcp_tools.return_value = [] mock_mgr.cleanup_mcp_tools.return_value = None mock_mgr_cls.return_value = mock_mgr tools = await initialize_mcp_tools(agent_id="a", org_id="o") assert tools == [] await cleanup_mcp_tools([]) mock_mgr.cleanup_mcp_tools.assert_awaited() # ==================== Additional tests for better coverage ==================== from app.tools.mcp_manager import ChatAgentMCPMixin from app.models.mcp_tool import MCPTransportType @pytest.mark.asyncio async def test_initialize_mcp_tools_none_ids(): """Test with None values for agent_id and org_id""" manager = MCPToolsManager() tools = await manager.initialize_mcp_tools(agent_id=None, org_id=None) assert tools == [] @pytest.mark.asyncio async def test_initialize_mcp_tools_filesystem_with_env_dirs(): """Test filesystem tool with directories from env_vars""" manager = MCPToolsManager() agent_id = str(uuid4()) org_id = str(uuid4()) mock_tool_config = MagicMock() mock_tool_config.name = "FS Tool" mock_tool_config.transport_type = MCPTransportType.STDIO mock_tool_config.timeout = None mock_tool_config.command = "npx" mock_tool_config.args = ["-y", "@modelcontextprotocol/server-filesystem"] # No dirs in args mock_tool_config.env_vars = {"ALLOWED_DIRECTORIES": "/tmp, /home"} # Dirs in env with patch("app.tools.mcp_manager.SessionLocal") as mock_sess_local, \ patch("app.tools.mcp_manager.MCPToolRepository") as mock_repo_cls, \ patch("app.tools.mcp_manager.MCPTools") as mock_mcp_tools_cls: mock_db = MagicMock() mock_sess_local.return_value.__enter__.return_value = mock_db mock_repo = MagicMock() mock_repo.get_agent_mcp_tools.return_value = [mock_tool_config] mock_repo_cls.return_value = mock_repo mock_mcp_instance = AsyncMock() mock_mcp_instance.functions = {"list_files": MagicMock()} mock_mcp_tools_cls.return_value = mock_mcp_instance tools = await manager.initialize_mcp_tools(agent_id, org_id) assert len(tools) == 1 @pytest.mark.asyncio async def test_initialize_mcp_tools_filesystem_empty_env_dirs(): """Test filesystem tool with empty ALLOWED_DIRECTORIES env var""" manager = MCPToolsManager() agent_id = str(uuid4()) org_id = str(uuid4()) mock_tool_config = MagicMock() mock_tool_config.name = "FS Tool" mock_tool_config.transport_type = MCPTransportType.STDIO mock_tool_config.timeout = None mock_tool_config.command = "npx" mock_tool_config.args = ["-y", "@modelcontextprotocol/server-filesystem"] mock_tool_config.env_vars = {"ALLOWED_DIRECTORIES": ""} # Empty dirs with patch("app.tools.mcp_manager.SessionLocal") as mock_sess_local, \ patch("app.tools.mcp_manager.MCPToolRepository") as mock_repo_cls: mock_db = MagicMock() mock_sess_local.return_value.__enter__.return_value = mock_db mock_repo = MagicMock() mock_repo.get_agent_mcp_tools.return_value = [mock_tool_config] mock_repo_cls.return_value = mock_repo tools = await manager.initialize_mcp_tools(agent_id, org_id) assert tools == [] # Should be skipped @pytest.mark.asyncio async def test_initialize_mcp_tools_uvx_command(): """Test tool with uvx command""" manager = MCPToolsManager() agent_id = str(uuid4()) org_id = str(uuid4()) mock_tool_config = MagicMock() mock_tool_config.name = "Python Tool" mock_tool_config.transport_type = MCPTransportType.STDIO mock_tool_config.timeout = None mock_tool_config.command = "uvx" mock_tool_config.args = ["python-mcp-server"] mock_tool_config.env_vars = {} with patch("app.tools.mcp_manager.SessionLocal") as mock_sess_local, \ patch("app.tools.mcp_manager.MCPToolRepository") as mock_repo_cls, \ patch("app.tools.mcp_manager.MCPTools") as mock_mcp_tools_cls: mock_db = MagicMock() mock_sess_local.return_value.__enter__.return_value = mock_db mock_repo = MagicMock() mock_repo.get_agent_mcp_tools.return_value = [mock_tool_config] mock_repo_cls.return_value = mock_repo mock_mcp_instance = AsyncMock() mock_mcp_instance.functions = {"run": MagicMock()} mock_mcp_tools_cls.return_value = mock_mcp_instance tools = await manager.initialize_mcp_tools(agent_id, org_id) assert len(tools) == 1 @pytest.mark.asyncio async def test_initialize_mcp_tools_connection_timeout(): """Test tool connection timeout""" manager = MCPToolsManager() agent_id = str(uuid4()) org_id = str(uuid4()) mock_tool_config = MagicMock() mock_tool_config.name = "Slow Tool" mock_tool_config.transport_type = MCPTransportType.STDIO mock_tool_config.timeout = None mock_tool_config.command = "npx" mock_tool_config.args = ["-y", "slow-package"] mock_tool_config.env_vars = {} with patch("app.tools.mcp_manager.SessionLocal") as mock_sess_local, \ patch("app.tools.mcp_manager.MCPToolRepository") as mock_repo_cls, \ patch("app.tools.mcp_manager.MCPTools") as mock_mcp_tools_cls: mock_db = MagicMock() mock_sess_local.return_value.__enter__.return_value = mock_db mock_repo = MagicMock() mock_repo.get_agent_mcp_tools.return_value = [mock_tool_config] mock_repo_cls.return_value = mock_repo mock_mcp_instance = MagicMock() mock_mcp_instance.__aenter__ = AsyncMock(side_effect=asyncio.TimeoutError()) mock_mcp_instance.__aexit__ = AsyncMock() mock_mcp_tools_cls.return_value = mock_mcp_instance tools = await manager.initialize_mcp_tools(agent_id, org_id) assert tools == [] @pytest.mark.asyncio async def test_initialize_mcp_tools_package_not_found(): """Test tool with 404 not found error""" manager = MCPToolsManager() agent_id = str(uuid4()) org_id = str(uuid4()) mock_tool_config = MagicMock() mock_tool_config.name = "Missing Tool" mock_tool_config.transport_type = MCPTransportType.STDIO mock_tool_config.timeout = None mock_tool_config.command = "npx" mock_tool_config.args = ["-y", "nonexistent"] mock_tool_config.env_vars = {} with patch("app.tools.mcp_manager.SessionLocal") as mock_sess_local, \ patch("app.tools.mcp_manager.MCPToolRepository") as mock_repo_cls, \ patch("app.tools.mcp_manager.MCPTools") as mock_mcp_tools_cls: mock_db = MagicMock() mock_sess_local.return_value.__enter__.return_value = mock_db mock_repo = MagicMock() mock_repo.get_agent_mcp_tools.return_value = [mock_tool_config] mock_repo_cls.return_value = mock_repo mock_mcp_instance = MagicMock() mock_mcp_instance.__aenter__ = AsyncMock(side_effect=Exception("404 package not found")) mock_mcp_instance.__aexit__ = AsyncMock() mock_mcp_tools_cls.return_value = mock_mcp_instance tools = await manager.initialize_mcp_tools(agent_id, org_id) assert tools == [] @pytest.mark.asyncio async def test_initialize_mcp_tools_no_functions_loaded(): """Test tool that connects but has no functions""" manager = MCPToolsManager() agent_id = str(uuid4()) org_id = str(uuid4()) mock_tool_config = MagicMock() mock_tool_config.name = "Empty Tool" mock_tool_config.transport_type = MCPTransportType.STDIO mock_tool_config.timeout = None mock_tool_config.command = "npx" mock_tool_config.args = ["-y", "empty-package"] mock_tool_config.env_vars = {} with patch("app.tools.mcp_manager.SessionLocal") as mock_sess_local, \ patch("app.tools.mcp_manager.MCPToolRepository") as mock_repo_cls, \ patch("app.tools.mcp_manager.MCPTools") as mock_mcp_tools_cls: mock_db = MagicMock() mock_sess_local.return_value.__enter__.return_value = mock_db mock_repo = MagicMock() mock_repo.get_agent_mcp_tools.return_value = [mock_tool_config] mock_repo_cls.return_value = mock_repo mock_mcp_instance = MagicMock() mock_mcp_instance.__aenter__ = AsyncMock() mock_mcp_instance.__aexit__ = AsyncMock() mock_mcp_instance.functions = {} # Empty functions mock_mcp_tools_cls.return_value = mock_mcp_instance tools = await manager.initialize_mcp_tools(agent_id, org_id) assert tools == [] @pytest.mark.asyncio async def test_initialize_mcp_tools_sse_transport(): """A remote tool that can't be reached is excluded and recorded.""" manager = MCPToolsManager() agent_id = str(uuid4()) org_id = str(uuid4()) mock_tool_config = _remote_config(MCPTransportType.SSE, name="SSE Tool") with patch("app.tools.mcp_manager.SessionLocal") as mock_sess_local, \ patch("app.tools.mcp_manager.MCPToolRepository") as mock_repo_cls, \ patch("app.tools.mcp_manager.MCPTools") as mock_mcp_tools_cls: mock_db = MagicMock() mock_sess_local.return_value.__enter__.return_value = mock_db mock_repo = MagicMock() mock_repo.get_agent_mcp_tools.return_value = [mock_tool_config] mock_repo_cls.return_value = mock_repo mock_mcp_instance = MagicMock() mock_mcp_instance.__aenter__ = AsyncMock(side_effect=Exception("connection refused")) mock_mcp_instance.__aexit__ = AsyncMock() mock_mcp_tools_cls.return_value = mock_mcp_instance tools = await manager.initialize_mcp_tools(agent_id, org_id) assert tools == [] assert manager.failed_tools[-1]["name"] == "SSE Tool" @pytest.mark.asyncio async def test_initialize_mcp_tools_http_transport(): """A remote tool that can't be reached is excluded and recorded.""" manager = MCPToolsManager() agent_id = str(uuid4()) org_id = str(uuid4()) mock_tool_config = _remote_config(MCPTransportType.HTTP, name="HTTP Tool") with patch("app.tools.mcp_manager.SessionLocal") as mock_sess_local, \ patch("app.tools.mcp_manager.MCPToolRepository") as mock_repo_cls, \ patch("app.tools.mcp_manager.MCPTools") as mock_mcp_tools_cls: mock_db = MagicMock() mock_sess_local.return_value.__enter__.return_value = mock_db mock_repo = MagicMock() mock_repo.get_agent_mcp_tools.return_value = [mock_tool_config] mock_repo_cls.return_value = mock_repo mock_mcp_instance = MagicMock() mock_mcp_instance.__aenter__ = AsyncMock(side_effect=Exception("connection refused")) mock_mcp_instance.__aexit__ = AsyncMock() mock_mcp_tools_cls.return_value = mock_mcp_instance tools = await manager.initialize_mcp_tools(agent_id, org_id) assert tools == [] assert manager.failed_tools[-1]["name"] == "HTTP Tool" @pytest.mark.asyncio async def test_initialize_mcp_tools_missing_command(): """Test STDIO tool with missing command""" manager = MCPToolsManager() agent_id = str(uuid4()) org_id = str(uuid4()) mock_tool_config = MagicMock() mock_tool_config.name = "Bad Tool" mock_tool_config.transport_type = MCPTransportType.STDIO mock_tool_config.timeout = None mock_tool_config.command = None mock_tool_config.args = None with patch("app.tools.mcp_manager.SessionLocal") as mock_sess_local, \ patch("app.tools.mcp_manager.MCPToolRepository") as mock_repo_cls: mock_db = MagicMock() mock_sess_local.return_value.__enter__.return_value = mock_db mock_repo = MagicMock() mock_repo.get_agent_mcp_tools.return_value = [mock_tool_config] mock_repo_cls.return_value = mock_repo tools = await manager.initialize_mcp_tools(agent_id, org_id) assert tools == [] @pytest.mark.asyncio async def test_initialize_mcp_tools_db_error(): """Test error getting tools from database""" manager = MCPToolsManager() agent_id = str(uuid4()) org_id = str(uuid4()) with patch("app.tools.mcp_manager.SessionLocal") as mock_sess_local: mock_sess_local.return_value.__enter__.side_effect = Exception("DB error") tools = await manager.initialize_mcp_tools(agent_id, org_id) assert tools == [] @pytest.mark.asyncio async def test_cleanup_mcp_tools_empty(): """Test cleanup with no tools""" manager = MCPToolsManager() manager.mcp_tools = [] await manager.cleanup_mcp_tools() assert manager.mcp_tools == [] @pytest.mark.asyncio async def test_cleanup_mcp_tools_timeout(): """Test cleanup with timeout""" manager = MCPToolsManager() mock_tool = AsyncMock() mock_tool.__aexit__ = AsyncMock(side_effect=asyncio.TimeoutError()) manager.mcp_tools = [mock_tool] await manager.cleanup_mcp_tools() assert manager.mcp_tools == [] @pytest.mark.asyncio async def test_cleanup_mcp_tools_cancel_scope_warning(): """Test cleanup with cancel scope warning""" manager = MCPToolsManager() mock_tool = AsyncMock() mock_tool.__aexit__ = AsyncMock(side_effect=Exception("cancel scope different task")) manager.mcp_tools = [mock_tool] await manager.cleanup_mcp_tools() assert manager.mcp_tools == [] @pytest.mark.asyncio async def test_cleanup_mcp_tools_unexpected_error(): """Test cleanup with unexpected error""" manager = MCPToolsManager() mock_tool = AsyncMock() mock_tool.__aexit__ = AsyncMock(side_effect=Exception("Unexpected error")) manager.mcp_tools = [mock_tool] await manager.cleanup_mcp_tools() assert manager.mcp_tools == [] # ==================== Failure tracking + connection test ==================== def _stdio_config(name="Elasticsearch", command="npx", args=None, env_vars=None, timeout=None): config = MagicMock() config.name = name config.transport_type = MCPTransportType.STDIO config.command = command config.args = args if args is not None else ["-y", "@elastic/mcp-server-elasticsearch"] config.env_vars = env_vars or {} config.timeout = timeout return config @pytest.mark.asyncio async def test_failed_tools_records_connect_error(): """A tool that can't start (e.g. npx missing) lands in failed_tools with the real error, and connected_tool_names stays empty.""" manager = MCPToolsManager() with patch("app.tools.mcp_manager.SessionLocal") as mock_sess_local, \ patch("app.tools.mcp_manager.MCPToolRepository") as mock_repo_cls, \ patch("app.tools.mcp_manager.MCPTools") as mock_mcp_tools_cls: mock_sess_local.return_value.__enter__.return_value = MagicMock() mock_repo = MagicMock() config = _stdio_config() config.id = 7 mock_repo.get_by_ids.return_value = [config] mock_repo_cls.return_value = mock_repo mock_mcp_instance = MagicMock() mock_mcp_instance.__aenter__ = AsyncMock( side_effect=FileNotFoundError("[Errno 2] No such file or directory: 'npx'") ) mock_mcp_instance.__aexit__ = AsyncMock() mock_mcp_tools_cls.return_value = mock_mcp_instance tools = await manager.initialize_mcp_tools_by_ids(str(uuid4()), [7]) assert tools == [] assert manager.connected_tool_names == [] assert len(manager.failed_tools) == 1 assert manager.failed_tools[0]["name"] == "Elasticsearch" assert "npx" in manager.failed_tools[0]["error"] @pytest.mark.asyncio async def test_failed_tools_records_missing_ids(): """Requested tool ids that no longer resolve to configs are reported.""" manager = MCPToolsManager() with patch("app.tools.mcp_manager.SessionLocal") as mock_sess_local, \ patch("app.tools.mcp_manager.MCPToolRepository") as mock_repo_cls: mock_sess_local.return_value.__enter__.return_value = MagicMock() mock_repo = MagicMock() mock_repo.get_by_ids.return_value = [] mock_repo_cls.return_value = mock_repo tools = await manager.initialize_mcp_tools_by_ids(str(uuid4()), [42]) assert tools == [] assert manager.failed_tools == [ {"name": "Tool #42", "error": "Configured tool no longer exists or is disabled"} ] @pytest.mark.asyncio async def test_connected_tool_names_recorded_on_success(): manager = MCPToolsManager() with patch("app.tools.mcp_manager.MCPTools") as mock_mcp_tools_cls: mock_mcp_instance = AsyncMock() mock_mcp_instance.functions = {"search": MagicMock()} mock_mcp_tools_cls.return_value = mock_mcp_instance tools = await manager._initialize_from_configs([_stdio_config()]) assert len(tools) == 1 assert manager.connected_tool_names == ["Elasticsearch"] assert manager.failed_tools == [] @pytest.mark.asyncio async def test_test_tool_config_success(): manager = MCPToolsManager() with patch("app.tools.mcp_manager.MCPTools") as mock_mcp_tools_cls: mock_mcp_instance = AsyncMock() mock_mcp_instance.functions = {"search": MagicMock(), "get_index": MagicMock()} mock_mcp_tools_cls.return_value = mock_mcp_instance result = await manager.test_tool_config(_stdio_config()) assert result["success"] is True assert sorted(result["functions"]) == ["get_index", "search"] assert result["error"] is None # The probe tool is torn down, never registered for later use. assert manager.mcp_tools == [] mock_mcp_instance.__aexit__.assert_awaited() @pytest.mark.asyncio async def test_test_tool_config_connect_failure(): manager = MCPToolsManager() with patch("app.tools.mcp_manager.MCPTools") as mock_mcp_tools_cls: mock_mcp_instance = MagicMock() mock_mcp_instance.__aenter__ = AsyncMock( side_effect=FileNotFoundError("[Errno 2] No such file or directory: 'npx'") ) mock_mcp_instance.__aexit__ = AsyncMock() mock_mcp_tools_cls.return_value = mock_mcp_instance result = await manager.test_tool_config(_stdio_config()) assert result["success"] is False assert "npx" in result["error"] @pytest.mark.asyncio async def test_test_tool_config_misconfigured(): manager = MCPToolsManager() result = await manager.test_tool_config(_stdio_config(command=None, args=None)) assert result["success"] is False assert "misconfigured" in result["error"] @pytest.mark.asyncio async def test_test_tool_config_timeout(): manager = MCPToolsManager() async def hang(): await asyncio.sleep(60) with patch("app.tools.mcp_manager.MCPTools") as mock_mcp_tools_cls, \ patch("app.tools.mcp_manager._connect_timeout", return_value=0.01): mock_mcp_instance = MagicMock() mock_mcp_instance.__aenter__ = MagicMock(side_effect=hang) mock_mcp_instance.__aexit__ = AsyncMock() mock_mcp_tools_cls.return_value = mock_mcp_instance result = await manager.test_tool_config(_stdio_config()) assert result["success"] is False assert "Timed out" in result["error"] @pytest.mark.asyncio async def test_mixed_success_and_failure_bookkeeping(): """With one tool connecting and one failing, both lists are accurate — this is what run.connector_status is built from.""" manager = MCPToolsManager() good = _stdio_config(name="Good") bad = _stdio_config(name="Bad") good_instance = AsyncMock() good_instance.functions = {"search": MagicMock()} bad_instance = MagicMock() bad_instance.__aenter__ = AsyncMock(side_effect=Exception("connection refused")) bad_instance.__aexit__ = AsyncMock() with patch("app.tools.mcp_manager.MCPTools", side_effect=[good_instance, bad_instance]): tools = await manager._initialize_from_configs([good, bad]) assert len(tools) == 1 assert manager.connected_tool_names == ["Good"] assert manager.failed_tools == [{"name": "Bad", "error": "connection refused"}] @pytest.mark.asyncio async def test_test_tool_config_no_functions(): manager = MCPToolsManager() with patch("app.tools.mcp_manager.MCPTools") as mock_mcp_tools_cls: mock_mcp_instance = AsyncMock() mock_mcp_instance.functions = {} mock_mcp_tools_cls.return_value = mock_mcp_instance result = await manager.test_tool_config(_stdio_config()) assert result["success"] is False assert "no tools" in result["error"] def _remote_config(transport, name="Remote", url="https://example.com/mcp", timeout=None): config = MagicMock() config.name = name config.transport_type = transport config.url = url config.headers = None config.timeout = timeout config.sse_read_timeout = None config.terminate_on_close = True return config def test_stdio_tool_uses_configured_timeout(): """The configured timeout reaches the MCP session on the STDIO path. Without it agno falls back to its 5s default, which npx-launched servers straddle — the same config then connects or times out at random.""" with patch("app.tools.mcp_manager.MCPTools") as mock_mcp_tools_cls: MCPToolsManager._build_tool(_stdio_config(timeout=120)) assert mock_mcp_tools_cls.call_args.kwargs["timeout_seconds"] == 120 def test_stdio_tool_defaults_timeout_when_unset(): with patch("app.tools.mcp_manager.MCPTools") as mock_mcp_tools_cls: MCPToolsManager._build_tool(_stdio_config(timeout=None)) assert mock_mcp_tools_cls.call_args.kwargs["timeout_seconds"] == DEFAULT_TIMEOUT_SECONDS @pytest.mark.parametrize("transport", [MCPTransportType.SSE, MCPTransportType.HTTP]) def test_remote_tool_passes_timeout_to_session(transport): """Remote transports need timeout_seconds too: agno takes the *min* of it and the transport timeout, so leaving it unset pins the session to 5s no matter what the server params say.""" with patch("app.tools.mcp_manager.MCPTools") as mock_mcp_tools_cls: MCPToolsManager._build_tool(_remote_config(transport, timeout=90)) assert mock_mcp_tools_cls.call_args.kwargs["timeout_seconds"] == 90 def test_connect_budget_covers_the_whole_handshake(): """The outer guard has to sit above what the handshake can legitimately spend — initialize and list_tools each get the full per-request timeout — otherwise raising the timeout changes nothing, the guard just fires first.""" config = _stdio_config(timeout=120) assert _connect_timeout(config) > _session_timeout(config) * HANDSHAKE_REQUESTS assert _connect_timeout(config) == 120 * HANDSHAKE_REQUESTS + CONNECT_TIMEOUT_MARGIN_SECONDS def test_connect_budget_caps_a_generous_per_tool_timeout(): """An interactive caller's ceiling wins over the tool's own budget, so a connector that spawns but never speaks can't stall a live reply.""" manager = MCPToolsManager(connect_budget=12.0) config = _stdio_config(timeout=300) assert manager._budgeted_connect_timeout(config, monotonic() + 12.0) <= 12.0 @pytest.mark.asyncio async def test_connect_budget_is_shared_across_tools(): """The ceiling is for the whole set: once it's gone the remaining tools are reported as skipped rather than silently retried at full budget.""" manager = MCPToolsManager(connect_budget=0.05) first, second = _stdio_config(name="First"), _stdio_config(name="Second") async def slow_connect(): await asyncio.sleep(0.2) with patch("app.tools.mcp_manager.MCPTools") as mock_mcp_tools_cls: mock_instance = MagicMock() mock_instance.__aenter__ = MagicMock(side_effect=slow_connect) mock_instance.__aexit__ = AsyncMock() mock_mcp_tools_cls.return_value = mock_instance tools = await manager._initialize_from_configs([first, second]) assert tools == [] assert [failure["name"] for failure in manager.failed_tools] == ["First", "Second"] assert "Skipped" in manager.failed_tools[-1]["error"] def test_chat_turn_gets_a_tighter_budget_than_the_test_button(): """create_async runs per chat turn with a visitor waiting, so it can't inherit the patience the Test button needs for a cold npx launch.""" assert CHAT_CONNECT_BUDGET_SECONDS < TEST_CONNECT_BUDGET_SECONDS @pytest.mark.asyncio async def test_close_does_not_block_on_a_hanging_teardown(): """A STDIO server that never answered can leave __aexit__ unwinding for as long as the child process lives. Closing gives up after its timeout.""" tool = _connected_instance({"search": MagicMock()}) started = asyncio.Event() async def never_finishes(*args): started.set() await asyncio.sleep(3600) tool.__aexit__ = never_finishes connection = HeldConnection(tool) await connection.open(1.0) await asyncio.wait_for(connection.close(timeout=0.05), timeout=1.0) assert started.is_set() @pytest.mark.asyncio async def test_a_transport_that_cancels_itself_is_a_failed_connect_not_a_cancelled_request(): """mcp's streamable-HTTP transport cancels its own scope when the server refuses the connection. That must not reach the caller as a CancelledError, or leave its task flagged so the next connector is cancelled too.""" manager = MCPToolsManager() refused = MagicMock() async def cancel_own_task(): asyncio.current_task().cancel() await asyncio.sleep(0) refused.__aenter__ = MagicMock(side_effect=cancel_own_task) refused.__aexit__ = AsyncMock() good = _connected_instance({"search": MagicMock()}) with patch("app.tools.mcp_manager.MCPTools", side_effect=[refused, good]): tools = await manager._initialize_from_configs([_stdio_config(name="Down"), _stdio_config(name="Up")]) assert asyncio.current_task().cancelling() == 0 assert manager.connected_tool_names == ["Up"] assert manager.failed_tools == [{"name": "Down", "error": CONNECTION_FAILED}] assert tools == [good] await manager.cleanup_mcp_tools() @pytest.mark.asyncio async def test_a_tool_is_closed_by_the_task_that_opened_it(): """anyio cancel scopes must be exited in the task that entered them.""" tool = _connected_instance({"search": MagicMock()}) tasks = [] tool.__aenter__ = AsyncMock(side_effect=lambda: tasks.append(asyncio.current_task())) tool.__aexit__ = AsyncMock(side_effect=lambda *args: tasks.append(asyncio.current_task())) manager = MCPToolsManager() with patch("app.tools.mcp_manager.MCPTools", return_value=tool): await manager._initialize_from_configs([_stdio_config()]) await manager.cleanup_mcp_tools() opened_in, closed_in = tasks assert opened_in is closed_in assert opened_in is not asyncio.current_task() @pytest.mark.parametrize("error, message", [ (ExceptionGroup("unhandled errors in a TaskGroup", [ConnectionRefusedError("All connection attempts failed")]), "All connection attempts failed"), (ValueError("bad url"), "bad url"), ]) def test_connect_errors_name_the_real_cause(error, message): assert _connect_error_message(error) == message def test_no_connect_budget_leaves_the_tool_timeout_intact(): """Background callers (the investigation worker) get the full budget — they're bounded by their own wall clock instead.""" manager = MCPToolsManager() config = _stdio_config(timeout=300) assert manager._budgeted_connect_timeout(config, None) == _connect_timeout(config) @pytest.mark.asyncio async def test_connect_timeout_error_reports_configured_budget(): manager = MCPToolsManager() with patch("app.tools.mcp_manager.MCPTools") as mock_mcp_tools_cls, \ patch("app.tools.mcp_manager._connect_timeout", return_value=130.0): mock_instance = MagicMock() mock_instance.__aenter__ = AsyncMock(side_effect=asyncio.TimeoutError()) mock_instance.__aexit__ = AsyncMock() mock_mcp_tools_cls.return_value = mock_instance result = await manager.test_tool_config(_stdio_config(timeout=120)) assert result["success"] is False assert "130s" in result["error"] def _connected_instance(functions): instance = MagicMock() instance.__aenter__ = AsyncMock() instance.__aexit__ = AsyncMock() instance.functions = functions return instance def _function(parameters): function = MagicMock() function.parameters = parameters return function @pytest.mark.asyncio async def test_connected_tools_are_opted_out_of_provider_strict_mode(): """agno marks every tool strict whenever the agent has a response_model. A third-party schema can't satisfy that subset, and the resulting 400 kills every request carrying the tool (#303) — so these must opt out.""" manager = MCPToolsManager() search = _function({ "type": "object", "properties": {"queryBody": {"type": "object", "additionalProperties": {}}}, }) with patch("app.tools.mcp_manager.MCPTools") as mock_mcp_tools_cls: mock_mcp_tools_cls.return_value = _connected_instance({"search": search}) result = await manager.test_tool_config(_stdio_config()) assert result["success"] is True # False, not None: agno only defaults strict when it is still unset. assert search.strict is False @pytest.mark.asyncio async def test_the_advertised_schema_is_left_exactly_as_the_server_sent_it(): """Rewriting a foreign schema would change what the model may send; with strict off the provider accepts it as published.""" manager = MCPToolsManager() schema = {"type": "object", "properties": {"queryBody": {"type": "object", "additionalProperties": {}}}} search = _function(schema) with patch("app.tools.mcp_manager.MCPTools") as mock_mcp_tools_cls: mock_mcp_tools_cls.return_value = _connected_instance({"search": search}) await manager.test_tool_config(_stdio_config()) assert search.parameters == schema class TestChatAgentMCPMixin: """Tests for ChatAgentMCPMixin""" @pytest.mark.asyncio async def test_cleanup_with_manager(self): """Test cleanup using _mcp_manager""" class TestAgent(ChatAgentMCPMixin): def __init__(self): self._mcp_manager = MagicMock() self._mcp_manager.cleanup_mcp_tools = AsyncMock() agent = TestAgent() await agent.cleanup_mcp_tools() agent._mcp_manager.cleanup_mcp_tools.assert_awaited() @pytest.mark.asyncio async def test_cleanup_without_manager(self): """Test cleanup without _mcp_manager""" class TestAgent(ChatAgentMCPMixin): def __init__(self): self.mcp_tools = [MagicMock()] agent = TestAgent() with patch("app.tools.mcp_manager.cleanup_mcp_tools", new_callable=AsyncMock) as mock_cleanup: await agent.cleanup_mcp_tools() mock_cleanup.assert_awaited_once_with(agent.mcp_tools) def test_del_with_tools(self): """Test destructor with mcp_tools""" class TestAgent(ChatAgentMCPMixin): def __init__(self): self.mcp_tools = [MagicMock()] agent = TestAgent() agent.__del__() # Should not raise def test_del_without_tools(self): """Test destructor without mcp_tools""" class TestAgent(ChatAgentMCPMixin): pass agent = TestAgent() agent.__del__() # Should not raise # --- Per-connector usage guidance (#318) ------------------------------------- # # The model never sees a connector's name or description; MCP function names # arrive in one flat namespace shared by every connected server. Guidance is # what tells it which source a function belongs to and which fields to query. def _guided_config(guidance, name="Elastic prod", functions=None): config = MagicMock() config.name = name config.transport_type = MCPTransportType.STDIO config.timeout = None config.command = "npx" config.args = ["-y", "@elastic/mcp-server-elasticsearch"] config.env_vars = {} config.usage_guidance = guidance return config async def _connect(manager, config, functions=("search", "esql")): """Run a config through the real build/connect path with MCPTools mocked, and hand back the toolkit instance the manager registered.""" with patch("app.tools.mcp_manager.SessionLocal") as mock_sess_local, \ patch("app.tools.mcp_manager.MCPToolRepository") as mock_repo_cls, \ patch("app.tools.mcp_manager.MCPTools") as mock_mcp_tools_cls: mock_sess_local.return_value.__enter__.return_value = MagicMock() mock_repo = MagicMock() mock_repo.get_by_ids.return_value = [config] mock_repo_cls.return_value = mock_repo instance = AsyncMock() instance.functions = {name: MagicMock() for name in functions} instance.instructions = None instance.add_instructions = False mock_mcp_tools_cls.return_value = instance await manager.initialize_mcp_tools_by_ids(str(uuid4()), [1]) return instance @pytest.mark.asyncio async def test_guidance_becomes_toolkit_instructions_naming_the_connector(): config = _guided_config("Order id is fields.order_ref, not order_id.") instance = await _connect(MCPToolsManager(), config) assert instance.add_instructions is True assert "Order id is fields.order_ref" in instance.instructions # The connector name and its functions are what bind the guidance to the # right tools in a namespace shared with every other connector. assert "Elastic prod" in instance.instructions assert "search" in instance.instructions and "esql" in instance.instructions @pytest.mark.asyncio async def test_a_connector_without_guidance_leaves_the_prompt_untouched(): """The regression fence for every org that never opts in.""" instance = await _connect(MCPToolsManager(), _guided_config(None)) assert instance.instructions is None assert instance.add_instructions is False @pytest.mark.asyncio async def test_whitespace_only_guidance_injects_nothing(): instance = await _connect(MCPToolsManager(), _guided_config(" \n ")) assert instance.instructions is None assert instance.add_instructions is False @pytest.mark.asyncio async def test_guidance_is_truncated_to_the_schema_limit(): from app.models.schemas.mcp_tool import MAX_USAGE_GUIDANCE_CHARS instance = await _connect(MCPToolsManager(), _guided_config("x" * 5000)) assert "x" * MAX_USAGE_GUIDANCE_CHARS in instance.instructions assert "x" * (MAX_USAGE_GUIDANCE_CHARS + 1) not in instance.instructions @pytest.mark.asyncio async def test_the_function_list_is_capped_for_a_large_server(): names = [f"tool_{i}" for i in range(MAX_GUIDANCE_FUNCTIONS + 10)] instance = await _connect(MCPToolsManager(), _guided_config("Guidance."), functions=names) assert "(10 more)" in instance.instructions assert f"tool_{MAX_GUIDANCE_FUNCTIONS + 5}" not in instance.instructions @pytest.mark.asyncio async def test_a_connector_that_fails_to_connect_carries_no_guidance(): """Guidance must never describe a source that isn't actually there.""" manager = MCPToolsManager() config = _guided_config("Indices: app-logs-*.") with patch("app.tools.mcp_manager.SessionLocal") as mock_sess_local, \ patch("app.tools.mcp_manager.MCPToolRepository") as mock_repo_cls, \ patch("app.tools.mcp_manager.MCPTools") as mock_mcp_tools_cls: mock_sess_local.return_value.__enter__.return_value = MagicMock() mock_repo = MagicMock() mock_repo.get_by_ids.return_value = [config] mock_repo_cls.return_value = mock_repo instance = AsyncMock() instance.functions = {} # connected, but exposed nothing instance.instructions = None instance.add_instructions = False mock_mcp_tools_cls.return_value = instance tools = await manager.initialize_mcp_tools_by_ids(str(uuid4()), [1]) assert tools == [] assert instance.instructions is None assert instance.add_instructions is False @pytest.mark.asyncio async def test_chat_agents_get_connector_guidance_too(): """Connectors are shared between investigations and chat agents, and the guidance is documentation about the source — it applies on both paths.""" manager = MCPToolsManager() config = _guided_config("Indices: app-logs-*.") with patch("app.tools.mcp_manager.SessionLocal") as mock_sess_local, \ patch("app.tools.mcp_manager.MCPToolRepository") as mock_repo_cls, \ patch("app.tools.mcp_manager.MCPTools") as mock_mcp_tools_cls: mock_sess_local.return_value.__enter__.return_value = MagicMock() mock_repo = MagicMock() mock_repo.get_agent_mcp_tools.return_value = [config] mock_repo_cls.return_value = mock_repo instance = AsyncMock() instance.functions = {"search": MagicMock()} instance.instructions = None instance.add_instructions = False mock_mcp_tools_cls.return_value = instance await manager.initialize_mcp_tools(str(uuid4()), str(uuid4())) assert instance.add_instructions is True assert "Indices: app-logs-*." in instance.instructions def test_agno_still_renders_toolkit_instructions_into_the_system_message(): """Contract test against the real agno library, not a mock. Every other test here mocks MCPTools, so they would all keep passing if agno renamed or dropped `Toolkit.add_instructions` — and the feature would silently do nothing. This pins the one assumption the whole thing rests on: guidance set on a toolkit reaches the SYSTEM prompt, after the agent's own instructions. If an agno upgrade breaks this, fail here and loudly. """ from agno.agent import Agent from agno.models.openai import OpenAIChat from agno.tools.toolkit import Toolkit class _Connector(Toolkit): def __init__(self): super().__init__(name="MCPTools") self.register(self.search_logs) def search_logs(self, q: str) -> str: """Search logs.""" return "ok" toolkit = _Connector() toolkit.instructions = CONNECTOR_GUIDANCE_TEMPLATE.format( name="Elastic prod", functions="search_logs", guidance="Order id is fields.order_ref." ) toolkit.add_instructions = True agent = Agent( name="t", tools=[toolkit], instructions="BASE RULES.", markdown=False, model=OpenAIChat(id="gpt-4o-mini", api_key="not-used-no-request-is-made"), ) agent.determine_tools_for_model(model=agent.model, session_id="s1") system_message = agent.get_system_message(session_id="s1").content assert "Order id is fields.order_ref." in system_message assert "Elastic prod" in system_message # After the agent's own instructions, so the base rules are read first. assert system_message.index("") < system_message.index("data_source")