"""Unit tests for custom MCP deferred load + user-level tool cache.""" from __future__ import annotations import asyncio from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import pytest from langchain_core.tools import BaseTool, StructuredTool from octop.infra.connectors.mcp_tool_cache import ( fingerprint_mcp_spec, wrap_tools_for_shared_use, ) def test_fingerprint_stable_and_ignores_enabled() -> None: a = fingerprint_mcp_spec( { "transport": "stdio", "command": "npx", "args": ["-y", "pkg"], "enabled": True, } ) b = fingerprint_mcp_spec( { "transport": "stdio", "command": "npx", "args": ["-y", "pkg"], "enabled": False, } ) assert a == b c = fingerprint_mcp_spec( { "transport": "stdio", "command": "npx", "args": ["-y", "other"], } ) assert a != c @pytest.mark.asyncio async def test_wrap_tools_serialize_ainvoke() -> None: order: list[str] = [] class _FakeTool: name = "demo_tool" description = "demo" args_schema = None async def ainvoke(self, *_a: Any, **_k: Any) -> str: order.append("start") await asyncio.sleep(0.05) order.append("end") return "ok" def invoke(self, *_a: Any, **_k: Any) -> str: return "ok" lock = asyncio.Lock() wrapped = wrap_tools_for_shared_use([_FakeTool()], lock) assert len(wrapped) == 1 assert isinstance(wrapped[0], BaseTool) assert isinstance(wrapped[0], StructuredTool) async def _call() -> Any: return await wrapped[0].ainvoke({}) await asyncio.gather(_call(), _call()) assert order == ["start", "end", "start", "end"] @pytest.mark.asyncio async def test_get_or_load_caches_and_reuses() -> None: from octop.infra.agents.manager import AgentManager mgr = object.__new__(AgentManager) mgr._mcp_tool_cache = {} mgr._mcp_tool_cache_locks = {} mgr._mcp_tool_cache_guard = asyncio.Lock() mgr._paths = MagicMock(root=MagicMock()) fake = MagicMock() fake.name = "s_tool1" fake.description = "t" fake.args_schema = None fake.metadata = None fake.ainvoke = AsyncMock(return_value="ok") aload = AsyncMock(return_value=[fake]) spec = {"transport": "stdio", "command": "npx", "args": []} with patch("octop_harness.mcp.aload_mcp_tools", aload): first = await mgr._get_or_load_mcp_tools(1, "s", spec) second = await mgr._get_or_load_mcp_tools(1, "s", spec) assert aload.await_count == 1 assert first is second assert isinstance(first[0], StructuredTool) assert first[0].name == "s_tool1" @pytest.mark.asyncio async def test_get_or_load_misses_on_fingerprint_change() -> None: from octop.infra.agents.manager import AgentManager mgr = object.__new__(AgentManager) mgr._mcp_tool_cache = {} mgr._mcp_tool_cache_locks = {} mgr._mcp_tool_cache_guard = asyncio.Lock() mgr._paths = MagicMock(root=MagicMock()) def _tool(name: str) -> MagicMock: t = MagicMock() t.name = name t.description = "t" t.args_schema = None t.metadata = None t.ainvoke = AsyncMock(return_value="ok") return t aload = AsyncMock(side_effect=[[_tool("s_a")], [_tool("s_b")]]) with patch("octop_harness.mcp.aload_mcp_tools", aload): a = await mgr._get_or_load_mcp_tools( 1, "s", {"transport": "stdio", "command": "npx", "args": ["a"]} ) b = await mgr._get_or_load_mcp_tools( 1, "s", {"transport": "stdio", "command": "npx", "args": ["b"]} ) assert aload.await_count == 2 assert a[0].name == "s_a" assert b[0].name == "s_b" assert len(mgr._mcp_tool_cache) == 1 @pytest.mark.asyncio async def test_prepare_chat_mcp_injects_custom_from_cache() -> None: from octop.infra.agents.manager import AgentManager mgr = object.__new__(AgentManager) mgr._mcp_tool_cache = {} mgr._mcp_tool_cache_locks = {} mgr._mcp_tool_cache_guard = asyncio.Lock() mgr._paths = MagicMock(root=MagicMock()) tool = MagicMock() tool.name = "deepwiki_ask" tool.description = "ask" tool.args_schema = None tool.metadata = None tool.ainvoke = AsyncMock(return_value="ok") aload = AsyncMock(return_value=[tool]) agent = MagicMock() agent._mcp_tool_name_set = frozenset() agent.config.mcp_server_configs = {"deepwiki": {}} agent.append_mcp_tools = MagicMock( side_effect=lambda tools: setattr( agent, "_mcp_tool_name_set", frozenset(getattr(t, "name", "") for t in tools) | frozenset(agent._mcp_tool_name_set), ) ) row = MagicMock() mgr.get_agent = MagicMock(return_value=agent) mgr.get_row = MagicMock(return_value=row) mgr._connector_uid_for = MagicMock(return_value=7) mgr._connector_svc = MagicMock() mgr._connector_svc.custom_harness_configs.return_value = { "deepwiki": { "transport": "streamable_http", "url": "https://mcp.deepwiki.com/mcp", "headers": {"Accept": "application/json, text/event-stream"}, } } mgr.reload_connectors = AsyncMock() with patch("octop_harness.mcp.aload_mcp_tools", aload): failed = await mgr.prepare_chat_mcp("A1", ["deepwiki"], connector_user_id=7) failed2 = await mgr.prepare_chat_mcp("A1", ["deepwiki"], connector_user_id=7) assert failed == [] assert failed2 == [] assert aload.await_count == 1 assert agent.append_mcp_tools.call_count == 1 mgr.reload_connectors.assert_not_awaited() assert agent.config.mcp_server_configs["deepwiki"]["transport"] == "streamable_http" injected = agent.append_mcp_tools.call_args.args[0] assert isinstance(injected[0], StructuredTool) @pytest.mark.asyncio async def test_prepare_chat_mcp_attaches_gateway_tools_without_reload() -> None: """Gateway connectors inject in-process; rebuilding would drop the runtime. A rebuild ends in the very same injection, but first closes the harness agent (and the checkpointer pool an in-flight turn still writes to) and holds ``AgentManager._lock`` across workspace I/O. """ from octop.infra.agents.manager import AgentManager from octop.infra.connectors.builder import mcp_server_name instance_id = "01TESTGATEWAYINSTANCE0001" server = mcp_server_name("qq-mail", instance_id) inst = MagicMock() inst.instance_id = instance_id inst.kind = "qq-mail" inst.status = "active" inst.mcp_server_name = server agent = MagicMock() agent._mcp_tools = [] agent._mcp_tool_name_set = frozenset() agent.config.mcp_server_configs = {} def _inject(tools: list[Any]) -> None: agent._mcp_tools = [*tools, *agent._mcp_tools] agent._mcp_tool_name_set = frozenset(str(getattr(t, "name", "")) for t in agent._mcp_tools) agent.inject_mcp_tools = MagicMock(side_effect=_inject) mgr = object.__new__(AgentManager) mgr._repos = MagicMock() mgr._repos.connector_repo.list_visible.return_value = [inst] mgr._connector_svc = MagicMock() mgr._connector_svc.custom_harness_configs.return_value = {} mgr._connector_svc.decrypt.return_value = {"email": "a@qq.com", "password": "code"} mgr._connector_svc.ensure_fresh_credentials = AsyncMock(return_value={}) mgr.get_agent = MagicMock(return_value=agent) mgr.get_row = MagicMock(return_value=MagicMock()) mgr._connector_uid_for = MagicMock(return_value=7) mgr.reload_connectors = AsyncMock() failed = await mgr.prepare_chat_mcp("A1", [server], connector_user_id=7) # Asserted before the tool checks so a reintroduced rebuild reports itself # rather than surfacing as "tools missing". mgr.reload_connectors.assert_not_awaited() assert failed == [] assert agent.inject_mcp_tools.call_count == 1 assert all(name.startswith(f"{server}_") for name in agent._mcp_tool_name_set) assert server in agent.config.mcp_server_configs mgr._connector_svc.ensure_fresh_credentials.assert_awaited_once_with(instance_id, "qq-mail") def test_wrap_tools_for_shared_use() -> None: lock = asyncio.Lock() inner = MagicMock() inner.name = "x" inner.description = "d" inner.args_schema = None inner.metadata = None wrapped = wrap_tools_for_shared_use([inner], lock) assert len(wrapped) == 1 assert isinstance(wrapped[0], StructuredTool)