LangBot/tests/unit_tests/provider/test_tool_manager.py

"""Unit tests for ToolManager.

Tests cover:
- Tool schema generation for OpenAI/LiteLLM
- Tool execution dispatch
"""

from __future__ import annotations

from importlib import import_module
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock

import langbot_plugin.api.entities.builtin.resource.tool as resource_tool
import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
import pytest

from langbot.pkg.provider.tools.errors import ToolExecutionDeniedError


def get_toolmgr_module():
    """Lazy import to avoid circular import issues."""
    return import_module('langbot.pkg.provider.tools.toolmgr')


class TestToolManagerInit:
    """Tests for ToolManager initialization."""

    def test_init_stores_app_reference(self):
        """Test that __init__ stores the Application reference."""
        toolmgr = get_toolmgr_module()

        mock_app = Mock()
        manager = toolmgr.ToolManager(mock_app)
        assert manager.ap is mock_app

    def test_init_no_tool_loaders(self):
        """Test that tool loaders are not initialized before initialize()."""
        toolmgr = get_toolmgr_module()

        mock_app = Mock()
        manager = toolmgr.ToolManager(mock_app)
        assert hasattr(manager, 'plugin_tool_loader') is False or manager.plugin_tool_loader is None


class TestToolManagerSchemaGeneration:
    """Tests for tool schema generation methods."""

    @pytest.fixture
    def mock_app(self):
        """Create mock app."""
        mock_app = Mock()
        mock_app.logger = Mock()
        return mock_app

    @pytest.fixture
    def sample_tools(self):
        """Create sample LLMTool list for testing."""

        def dummy_weather_func(**kwargs):
            return 'weather result'

        def dummy_calc_func(**kwargs):
            return 'calc result'

        tools = [
            resource_tool.LLMTool(
                name='get_weather',
                human_desc='Get current weather for a location',
                description='Get current weather for a location',
                parameters={
                    'type': 'object',
                    'properties': {'location': {'type': 'string', 'description': 'City name'}},
                    'required': ['location'],
                },
                func=dummy_weather_func,
            ),
            resource_tool.LLMTool(
                name='calculate',
                human_desc='Perform a calculation',
                description='Perform a calculation',
                parameters={
                    'type': 'object',
                    'properties': {'expression': {'type': 'string', 'description': 'Math expression'}},
                    'required': ['expression'],
                },
                func=dummy_calc_func,
            ),
        ]
        return tools

    @pytest.mark.asyncio
    async def test_generate_tools_for_openai(self, mock_app, sample_tools):
        """Test that generate_tools_for_openai produces correct schema."""
        toolmgr = get_toolmgr_module()

        manager = toolmgr.ToolManager(mock_app)
        result = await manager.generate_tools_for_openai(sample_tools)

        assert len(result) == 2

        # Verify first tool schema
        tool1 = result[0]
        assert tool1['type'] == 'function'
        assert tool1['function']['name'] == 'get_weather'
        assert tool1['function']['description'] == 'Get current weather for a location'
        assert 'parameters' in tool1['function']
        assert tool1['function']['parameters']['type'] == 'object'

        # Verify second tool schema
        tool2 = result[1]
        assert tool2['type'] == 'function'
        assert tool2['function']['name'] == 'calculate'

    @pytest.mark.asyncio
    async def test_generate_tools_empty_list(self, mock_app):
        """Test that generating tools from empty list returns empty list."""
        toolmgr = get_toolmgr_module()

        manager = toolmgr.ToolManager(mock_app)

        openai_result = await manager.generate_tools_for_openai([])
        assert openai_result == []

    @pytest.mark.asyncio
    async def test_openai_schema_fields_complete(self, mock_app, sample_tools):
        """Test that OpenAI schema includes all required fields."""
        toolmgr = get_toolmgr_module()

        manager = toolmgr.ToolManager(mock_app)
        result = await manager.generate_tools_for_openai(sample_tools)

        for tool_schema in result:
            assert 'type' in tool_schema
            assert tool_schema['type'] == 'function'
            assert 'function' in tool_schema
            func = tool_schema['function']
            assert 'name' in func
            assert 'description' in func
            assert 'parameters' in func


class TestToolManagerExecuteFuncCall:
    """Tests for execute_func_call method."""

    @pytest.fixture
    def mock_app_with_loaders(self):
        """Create mock app with mock tool loaders.

        Returns (app, plugin_loader, mcp_loader). The native and skill loaders
        are attached directly to the app for tests that don't need to assert
        against them — they all default to ``has_tool == False`` so the
        execute_func_call probe falls through to the plugin/mcp pair.
        """
        mock_app = Mock()
        mock_app.logger = Mock()

        def _make_inert_loader():
            loader = Mock()
            loader.has_tool = AsyncMock(return_value=False)
            loader.invoke_tool = AsyncMock(return_value=None)
            loader.initialize = AsyncMock()
            loader.shutdown = AsyncMock()
            return loader

        # Create mock plugin loader
        mock_plugin_loader = _make_inert_loader()
        mock_plugin_loader.invoke_tool = AsyncMock(return_value='plugin_result')

        # Create mock MCP loader
        mock_mcp_loader = _make_inert_loader()
        mock_mcp_loader.invoke_tool = AsyncMock(return_value='mcp_result')

        # Stash inert native/skill loaders so the ToolManager probe order
        # (native → plugin → mcp → skill) doesn't AttributeError. Tests that
        # need to override these can replace the attributes on the manager.
        mock_app._inert_native_loader = _make_inert_loader()
        mock_app._inert_skill_loader = _make_inert_loader()

        return mock_app, mock_plugin_loader, mock_mcp_loader

    @staticmethod
    def _wire_loaders(manager, mock_app, plugin_loader, mcp_loader):
        """Attach all four loaders (native + plugin + mcp + skill) to manager."""
        manager.native_tool_loader = mock_app._inert_native_loader
        manager.plugin_tool_loader = plugin_loader
        manager.mcp_tool_loader = mcp_loader
        manager.skill_tool_loader = mock_app._inert_skill_loader

    @pytest.fixture
    def sample_query(self):
        """Create sample query for testing."""
        query = Mock(spec=pipeline_query.Query)
        return query

    @pytest.mark.asyncio
    async def test_execute_calls_plugin_loader_when_has_tool(self, mock_app_with_loaders, sample_query):
        """Test that execute_func_call uses plugin loader when tool exists there."""
        toolmgr = get_toolmgr_module()

        mock_app, mock_plugin_loader, mock_mcp_loader = mock_app_with_loaders
        mock_plugin_loader.has_tool = AsyncMock(return_value=True)

        manager = toolmgr.ToolManager(mock_app)
        self._wire_loaders(manager, mock_app, mock_plugin_loader, mock_mcp_loader)

        result = await manager.execute_func_call('test_tool', {'param': 'value'}, sample_query)

        assert result == 'plugin_result'
        mock_plugin_loader.invoke_tool.assert_called_once_with('test_tool', {'param': 'value'}, sample_query)
        # MCP loader should not be called
        mock_mcp_loader.invoke_tool.assert_not_called()

    @pytest.mark.asyncio
    async def test_execute_calls_mcp_loader_when_plugin_not_found(self, mock_app_with_loaders, sample_query):
        """Test that execute_func_call uses MCP loader when plugin doesn't have tool."""
        toolmgr = get_toolmgr_module()

        mock_app, mock_plugin_loader, mock_mcp_loader = mock_app_with_loaders
        mock_plugin_loader.has_tool = AsyncMock(return_value=False)
        mock_mcp_loader.has_tool = AsyncMock(return_value=True)

        manager = toolmgr.ToolManager(mock_app)
        self._wire_loaders(manager, mock_app, mock_plugin_loader, mock_mcp_loader)

        result = await manager.execute_func_call('test_tool', {'param': 'value'}, sample_query)

        assert result == 'mcp_result'
        mock_mcp_loader.invoke_tool.assert_called_once_with('test_tool', {'param': 'value'}, sample_query)

    @pytest.mark.asyncio
    async def test_execute_raises_when_tool_not_found(self, mock_app_with_loaders, sample_query):
        """Test that execute_func_call raises ToolNotFoundError when tool not found."""
        toolmgr = get_toolmgr_module()

        mock_app, mock_plugin_loader, mock_mcp_loader = mock_app_with_loaders
        mock_plugin_loader.has_tool = AsyncMock(return_value=False)
        mock_mcp_loader.has_tool = AsyncMock(return_value=False)

        manager = toolmgr.ToolManager(mock_app)
        self._wire_loaders(manager, mock_app, mock_plugin_loader, mock_mcp_loader)

        with pytest.raises(toolmgr.ToolNotFoundError, match='Tool not found: unknown_tool'):
            await manager.execute_func_call('unknown_tool', {}, sample_query)

    @pytest.mark.asyncio
    async def test_plugin_loader_checked_first(self, mock_app_with_loaders, sample_query):
        """Test that plugin loader is checked before MCP loader."""
        toolmgr = get_toolmgr_module()

        mock_app, mock_plugin_loader, mock_mcp_loader = mock_app_with_loaders
        # Both loaders have the tool, but plugin should be used
        mock_plugin_loader.has_tool = AsyncMock(return_value=True)
        mock_mcp_loader.has_tool = AsyncMock(return_value=True)

        manager = toolmgr.ToolManager(mock_app)
        self._wire_loaders(manager, mock_app, mock_plugin_loader, mock_mcp_loader)

        await manager.execute_func_call('test_tool', {}, sample_query)

        # Plugin loader should be invoked, MCP should not
        mock_plugin_loader.invoke_tool.assert_called_once()
        mock_mcp_loader.invoke_tool.assert_not_called()

    @pytest.mark.asyncio
    async def test_bound_mcp_source_wins_over_unbound_plugin_name_collision(self, mock_app_with_loaders):
        toolmgr = get_toolmgr_module()
        mock_app, mock_plugin_loader, mock_mcp_loader = mock_app_with_loaders
        mock_plugin_loader.has_tool = AsyncMock(return_value=True)
        mock_mcp_loader.has_tool = AsyncMock(return_value=True)
        manager = toolmgr.ToolManager(mock_app)
        self._wire_loaders(manager, mock_app, mock_plugin_loader, mock_mcp_loader)
        query = Mock(variables={})

        result = await manager.execute_func_call(
            'shared_tool',
            {'value': 1},
            query,
            source_ref={'source': 'mcp', 'source_id': 'bound-mcp'},
        )

        assert result == 'mcp_result'
        mock_plugin_loader.has_tool.assert_not_awaited()
        mock_plugin_loader.invoke_tool.assert_not_awaited()
        mock_mcp_loader.has_tool.assert_awaited_once_with('shared_tool', source_id='bound-mcp')
        mock_mcp_loader.invoke_tool.assert_awaited_once_with(
            'shared_tool',
            {'value': 1},
            query,
            source_id='bound-mcp',
        )

    @pytest.mark.asyncio
    async def test_bound_plugin_source_selects_one_of_two_same_name_plugins(self, mock_app_with_loaders):
        toolmgr = get_toolmgr_module()
        mock_app, mock_plugin_loader, mock_mcp_loader = mock_app_with_loaders
        mock_plugin_loader.has_tool = AsyncMock(return_value=True)
        mock_mcp_loader.has_tool = AsyncMock(return_value=True)
        manager = toolmgr.ToolManager(mock_app)
        self._wire_loaders(manager, mock_app, mock_plugin_loader, mock_mcp_loader)
        query = Mock(variables={})

        result = await manager.execute_func_call(
            'shared_tool',
            {},
            query,
            source_ref={'source': 'plugin', 'source_id': 'authorized/plugin'},
        )

        assert result == 'plugin_result'
        mock_plugin_loader.has_tool.assert_awaited_once_with('shared_tool', source_id='authorized/plugin')
        mock_plugin_loader.invoke_tool.assert_awaited_once_with(
            'shared_tool',
            {},
            query,
            source_id='authorized/plugin',
        )
        mock_mcp_loader.has_tool.assert_not_awaited()

    @pytest.mark.asyncio
    async def test_disabled_synthetic_mcp_tool_is_monitored_as_error(
        self,
        mock_app_with_loaders,
        sample_query,
    ):
        toolmgr = get_toolmgr_module()
        mock_app, mock_plugin_loader, mock_mcp_loader = mock_app_with_loaders
        mock_app.monitoring_service = SimpleNamespace(record_tool_call=AsyncMock())
        sample_query.variables = {}
        sample_query.launcher_type = None
        sample_query.launcher_id = None
        sample_query.bot_uuid = 'bot-1'
        mock_mcp_loader.has_tool = AsyncMock(return_value=True)
        mock_mcp_loader.invoke_tool = AsyncMock(
            side_effect=ToolExecutionDeniedError(
                'langbot_mcp_read_resource',
                'MCP resource agent reads are disabled',
            )
        )
        manager = toolmgr.ToolManager(mock_app)
        self._wire_loaders(manager, mock_app, mock_plugin_loader, mock_mcp_loader)

        with pytest.raises(ToolExecutionDeniedError):
            await manager.execute_func_call(
                'langbot_mcp_read_resource',
                {'server_name': 'docs', 'uri': 'file:///README.md'},
                sample_query,
                source_ref={'source': 'mcp', 'source_id': None},
            )

        record = mock_app.monitoring_service.record_tool_call
        record.assert_awaited_once()
        assert record.await_args.kwargs['status'] == 'error'
        assert record.await_args.kwargs['tool_source'] == 'mcp'
        assert 'MCP resource agent reads are disabled' in record.await_args.kwargs['error_message']


class TestToolManagerSourceResolution:
    @pytest.mark.asyncio
    @pytest.mark.parametrize('source', ['mcp', 'plugin'])
    async def test_same_name_implementations_in_one_scope_are_hidden(self, source):
        toolmgr = get_toolmgr_module()
        app = Mock(logger=Mock())
        manager = toolmgr.ToolManager(app)
        manager.get_tool_catalog = AsyncMock(
            return_value=[
                {'name': 'shared_tool', 'source': source, 'source_id': 'source-a'},
                {'name': 'shared_tool', 'source': source, 'source_id': 'source-b'},
                {'name': 'unique_tool', 'source': 'builtin'},
            ]
        )

        catalog = await manager.get_resolved_tool_catalog()

        assert [item['name'] for item in catalog] == ['unique_tool']
        app.logger.warning.assert_called_once()


class TestToolManagerShutdown:
    """Tests for shutdown method."""

    @pytest.mark.asyncio
    async def test_shutdown_calls_loader_shutdown(self):
        """Test that shutdown calls shutdown on every registered loader."""
        toolmgr = get_toolmgr_module()

        mock_app = Mock()

        def _make_loader():
            loader = Mock()
            loader.shutdown = AsyncMock()
            return loader

        mock_native_loader = _make_loader()
        mock_plugin_loader = _make_loader()
        mock_mcp_loader = _make_loader()
        mock_skill_loader = _make_loader()

        manager = toolmgr.ToolManager(mock_app)
        manager.native_tool_loader = mock_native_loader
        manager.plugin_tool_loader = mock_plugin_loader
        manager.mcp_tool_loader = mock_mcp_loader
        manager.skill_tool_loader = mock_skill_loader

        await manager.shutdown()

        mock_native_loader.shutdown.assert_called_once()
        mock_plugin_loader.shutdown.assert_called_once()
        mock_mcp_loader.shutdown.assert_called_once()
        mock_skill_loader.shutdown.assert_called_once()