replypier/backend/tests/tools/test_mcp_manager.py

"""
Copyright 2024-2026 ReplyPier

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.
"""

"""
ReplyPier - 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("</instructions>") < system_message.index("data_source")