LangBot/tests/unit_tests/api/service/test_mcp_service.py

"""
Unit tests for MCPService.

Tests MCP server CRUD operations including:
- MCP server listing with runtime info
- MCP server creation with limitations
- MCP server update with enable/disable
- MCP server deletion
- MCP server connection testing

Source: src/langbot/pkg/api/http/service/mcp.py
"""

from __future__ import annotations

import asyncio
import copy
import pytest
from unittest.mock import AsyncMock, Mock, MagicMock
from types import SimpleNamespace
import uuid

from langbot.pkg.api.http.authz import Permission
from langbot.pkg.api.http.context import (
    ExecutionContext,
    PrincipalContext,
    PrincipalType,
    RequestContext,
    WorkspaceContext,
)
from langbot.pkg.api.http.service.mcp import MCPService, redact_mcp_secrets, restore_mcp_secret_placeholders
from langbot.pkg.entity.persistence.mcp import MCPServer
from langbot.pkg.provider.tools.loaders.mcp_policy import MCPStdioDisabledError
from langbot.pkg.workspace.errors import WorkspaceNotFoundError


pytestmark = pytest.mark.asyncio

_CONTEXT = ExecutionContext(
    instance_uuid='instance-a',
    workspace_uuid='workspace-a',
    placement_generation=1,
)

_VIEWER_CONTEXT = RequestContext(
    instance_uuid='instance-a',
    placement_generation=1,
    request_id='request-a',
    auth_type='user_token',
    principal=PrincipalContext(
        principal_type=PrincipalType.ACCOUNT,
        account_uuid='account-a',
    ),
    workspace=WorkspaceContext(
        workspace_uuid='workspace-a',
        membership_uuid='membership-a',
        role='viewer',
        permissions=frozenset({Permission.RESOURCE_VIEW.value}),
    ),
)


def _service(ap: SimpleNamespace) -> MCPService:
    ap.workspace_service = SimpleNamespace(
        get_execution_binding=AsyncMock(return_value=SimpleNamespace(instance_uuid=_CONTEXT.instance_uuid))
    )
    if not hasattr(ap, 'logger'):
        ap.logger = Mock()
    return MCPService(ap)


def _create_mock_mcp_server(
    server_uuid: str = None,
    name: str = 'Test MCP Server',
    enable: bool = True,
    mode: str = 'stdio',
    extra_args: dict = None,
) -> Mock:
    """Helper to create mock MCPServer entity."""
    server = Mock(spec=MCPServer)
    server.uuid = server_uuid or str(uuid.uuid4())
    server.name = name
    server.enable = enable
    server.mode = mode
    server.extra_args = extra_args or {}
    return server


def _create_mock_result(items: list = None, first_item=None, *, scalar_value=0, rowcount=1):
    """Create mock result object for persistence queries."""
    result = Mock()
    result.all = Mock(return_value=items or [])
    result.first = Mock(return_value=first_item)
    result.scalar = Mock(return_value=scalar_value)
    result.rowcount = rowcount
    return result


class TestMCPServiceGetRuntimeInfo:
    """Tests for get_runtime_info method."""

    async def test_get_runtime_info_session_exists(self):
        """Returns runtime info when session exists."""
        # Setup
        ap = SimpleNamespace()
        ap.tool_mgr = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader = SimpleNamespace()

        mock_session = SimpleNamespace()
        mock_session.get_runtime_info_dict = Mock(return_value={'status': 'running', 'tools': 5})
        ap.tool_mgr.mcp_tool_loader.get_session = Mock(return_value=mock_session)

        service = _service(ap)

        # Execute
        result = await service.get_runtime_info(_CONTEXT, 'test-server')

        # Verify
        assert result is not None
        assert result['status'] == 'running'

    async def test_get_runtime_info_session_not_exists(self):
        """Returns None when session not exists."""
        # Setup
        ap = SimpleNamespace()
        ap.tool_mgr = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader.get_session = Mock(return_value=None)

        service = _service(ap)

        # Execute
        result = await service.get_runtime_info(_CONTEXT, 'nonexistent-server')

        # Verify
        assert result is None


class TestMCPServiceResources:
    """Tests for MCP resource helpers."""

    async def test_get_resource_templates_delegates_to_loader(self):
        ap = SimpleNamespace()
        ap.tool_mgr = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader.get_resource_templates = AsyncMock(
            return_value=[{'uri_template': 'file:///{path}', 'name': 'files'}]
        )

        service = _service(ap)
        service._require_server = AsyncMock(return_value=(_CONTEXT, {'name': 'docs'}))

        result = await service.get_mcp_server_resource_templates(_CONTEXT, 'docs')

        assert result == [{'uri_template': 'file:///{path}', 'name': 'files'}]
        ap.tool_mgr.mcp_tool_loader.get_resource_templates.assert_awaited_once_with(_CONTEXT, 'docs')

    async def test_read_resource_envelope_uses_ui_preview_source(self):
        ap = SimpleNamespace()
        ap.tool_mgr = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader.read_resource_envelope = AsyncMock(
            return_value={
                'server_name': 'docs',
                'uri': 'file:///README.md',
                'contents': [],
                'source': 'ui_preview',
            }
        )

        service = _service(ap)
        service._require_server = AsyncMock(return_value=(_CONTEXT, {'name': 'docs'}))

        result = await service.read_mcp_server_resource_envelope(
            _CONTEXT,
            'docs',
            'file:///README.md',
            max_bytes=4096,
            include_blob=True,
        )

        assert result['source'] == 'ui_preview'
        ap.tool_mgr.mcp_tool_loader.read_resource_envelope.assert_awaited_once_with(
            _CONTEXT,
            'docs',
            'file:///README.md',
            include_blob=True,
            source='ui_preview',
            max_bytes=4096,
        )


class TestMCPServiceGetMCPServers:
    """Tests for get_mcp_servers method."""

    async def test_get_mcp_servers_empty_list(self):
        """Returns empty list when no MCP servers exist."""
        # Setup
        ap = SimpleNamespace()
        ap.persistence_mgr = SimpleNamespace()
        mock_result = _create_mock_result([])
        ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
        ap.persistence_mgr.serialize_model = Mock(
            side_effect=lambda model_cls, entity: {
                'uuid': entity.uuid,
                'name': entity.name,
            }
        )
        ap.tool_mgr = SimpleNamespace(mcp_tool_loader=SimpleNamespace(get_session=Mock(return_value=None)))

        service = _service(ap)

        # Execute
        result = await service.get_mcp_servers(_CONTEXT)

        # Verify
        assert result == []

    async def test_get_mcp_servers_returns_serialized_list(self):
        """Returns serialized list of MCP servers."""
        # Setup
        ap = SimpleNamespace()
        ap.persistence_mgr = SimpleNamespace()

        server1 = _create_mock_mcp_server(server_uuid='uuid-1', name='Server 1')
        server2 = _create_mock_mcp_server(server_uuid='uuid-2', name='Server 2')

        mock_result = _create_mock_result([server1, server2])
        ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
        ap.persistence_mgr.serialize_model = Mock(
            side_effect=lambda model_cls, entity: {
                'uuid': entity.uuid,
                'name': entity.name,
                'enable': entity.enable,
                'mode': entity.mode,
            }
        )
        ap.tool_mgr = SimpleNamespace(mcp_tool_loader=SimpleNamespace(get_session=Mock(return_value=None)))

        service = _service(ap)

        # Execute
        result = await service.get_mcp_servers(_CONTEXT)

        # Verify
        assert len(result) == 2
        assert result[0]['name'] == 'Server 1'
        assert result[1]['name'] == 'Server 2'

    async def test_get_mcp_servers_with_runtime_info(self):
        """Returns MCP servers with runtime info when requested."""
        # Setup
        ap = SimpleNamespace()
        ap.persistence_mgr = SimpleNamespace()

        server1 = _create_mock_mcp_server(server_uuid='uuid-1', name='Server 1')

        mock_result = _create_mock_result([server1])
        ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
        ap.persistence_mgr.serialize_model = Mock(
            side_effect=lambda model_cls, entity: {
                'uuid': entity.uuid,
                'name': entity.name,
            }
        )
        ap.tool_mgr = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
        runtime_session = SimpleNamespace(get_runtime_info_dict=Mock(return_value={'status': 'connected'}))
        ap.tool_mgr.mcp_tool_loader.get_session = Mock(return_value=runtime_session)

        service = _service(ap)

        # Execute
        result = await service.get_mcp_servers(_CONTEXT, contain_runtime_info=True)

        # Verify - runtime info included
        assert result[0]['runtime_info'] == {'status': 'connected'}

    async def test_resource_view_list_and_detail_redact_secrets_without_mutating_raw_data(self):
        ap = SimpleNamespace()
        ap.persistence_mgr = SimpleNamespace()
        server = _create_mock_mcp_server(name='Secret Server')
        serialized = {
            'uuid': 'secret-uuid',
            'name': 'Secret Server',
            'enable': True,
            'extra_args': {
                'url': (
                    'https://mcp-user:mcp-password@mcp.invalid/connect'
                    '?token=url-secret&transport=streamable&sig=signed-secret'
                ),
                'headers': {
                    'Authorization': 'Bearer top-secret',
                    'X-API-Key': 'api-secret',
                    'Accept': 'application/json',
                },
                'env': {
                    'ACCESS_TOKEN': 'access-secret',
                    'TOKENIZER': 'public-model-name',
                },
                'credentials': {
                    'username': 'service-user',
                    'password': 'password-secret',
                },
                'public_key': 'public-value',
            },
        }
        original = copy.deepcopy(serialized)
        ap.persistence_mgr.execute_async = AsyncMock(
            side_effect=[
                _create_mock_result([server]),
                _create_mock_result(first_item=server),
            ]
        )
        ap.persistence_mgr.serialize_model = Mock(return_value=serialized)
        ap.tool_mgr = SimpleNamespace(mcp_tool_loader=SimpleNamespace(get_session=Mock(return_value=None)))
        service = _service(ap)

        listed = await service.get_mcp_servers(_VIEWER_CONTEXT)
        detail = await service.get_mcp_server_by_name(_VIEWER_CONTEXT, 'Secret Server')

        for response in (listed[0], detail):
            assert response['extra_args']['url'] == (
                'https://***@mcp.invalid/connect?token=***&transport=streamable&sig=***'
            )
            assert response['extra_args']['headers'] == {
                'Authorization': '***',
                'X-API-Key': '***',
                'Accept': 'application/json',
            }
            assert response['extra_args']['env'] == {
                'ACCESS_TOKEN': '***',
                'TOKENIZER': 'public-model-name',
            }
            assert response['extra_args']['credentials'] == {
                'username': '***',
                'password': '***',
            }
            assert response['extra_args']['public_key'] == 'public-value'
        assert serialized == original

    async def test_redacted_url_roundtrip_restores_persisted_credentials(self):
        persisted = {
            'extra_args': {'url': 'https://mcp-user:mcp-password@mcp.invalid/connect?token=url-secret&transport=http'}
        }

        submitted = redact_mcp_secrets(persisted)

        assert submitted['extra_args']['url'] == 'https://***@mcp.invalid/connect?token=***&transport=http'
        assert restore_mcp_secret_placeholders(submitted, persisted) == persisted


class TestMCPServiceCreateMCPServer:
    """Tests for create_mcp_server method."""

    async def test_create_stdio_rejected_by_independent_instance_gate(self):
        ap = SimpleNamespace(
            instance_config=SimpleNamespace(
                data={
                    'mcp': {'stdio': {'enabled': False}},
                    'system': {'limitation': {'max_extensions': -1}},
                }
            ),
            persistence_mgr=SimpleNamespace(execute_async=AsyncMock()),
            tool_mgr=None,
        )
        service = _service(ap)

        with pytest.raises(MCPStdioDisabledError, match='disabled by instance policy'):
            await service.create_mcp_server(
                _CONTEXT,
                {'name': 'local', 'mode': 'stdio', 'enable': True, 'extra_args': {}},
            )

        ap.persistence_mgr.execute_async.assert_not_awaited()

    async def test_create_mcp_server_max_extensions_reached_raises(self):
        """Raises ValueError when max_extensions limit reached."""
        # Setup
        ap = SimpleNamespace()
        ap.persistence_mgr = SimpleNamespace()
        ap.instance_config = SimpleNamespace()
        ap.instance_config.data = {'system': {'limitation': {'max_extensions': 2}}}
        ap.plugin_connector = SimpleNamespace()
        ap.plugin_connector.list_plugins = AsyncMock(return_value=[Mock(), Mock()])  # 2 plugins

        # Mock get_mcp_servers to return 0 servers (2 plugins already)
        ap.persistence_mgr.execute_async = AsyncMock(
            side_effect=[
                _create_mock_result(scalar_value=0),
                _create_mock_result(scalar_value=2),
            ]
        )
        ap.persistence_mgr.serialize_model = Mock(return_value={})
        ap.tool_mgr = SimpleNamespace(mcp_tool_loader=SimpleNamespace(get_session=Mock(return_value=None)))

        service = _service(ap)

        # Execute & Verify - 2 plugins + new server would exceed limit
        with pytest.raises(ValueError, match='Maximum number of extensions'):
            await service.create_mcp_server(_CONTEXT, {'name': 'New Server'})

    async def test_create_mcp_server_no_limit(self):
        """Creates MCP server without limit when max_extensions=-1."""
        # Setup
        ap = SimpleNamespace()
        ap.persistence_mgr = SimpleNamespace()
        ap.instance_config = SimpleNamespace()
        ap.instance_config.data = {
            'system': {
                'limitation': {
                    'max_extensions': -1  # No limit
                }
            }
        }
        ap.tool_mgr = None

        mock_result = _create_mock_result([])
        ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
        ap.persistence_mgr.serialize_model = Mock(return_value={'uuid': 'new-uuid'})

        service = _service(ap)

        # Execute
        server_uuid = await service.create_mcp_server(_CONTEXT, {'name': 'New Server'})

        # Verify
        assert server_uuid is not None
        assert len(server_uuid) == 36  # UUID format

    async def test_create_mcp_server_duplicate_name_raises(self):
        """Rejects duplicate MCP server names."""
        # Setup
        ap = SimpleNamespace()
        ap.persistence_mgr = SimpleNamespace()
        ap.instance_config = SimpleNamespace()
        ap.instance_config.data = {'system': {'limitation': {'max_extensions': -1}}}
        ap.tool_mgr = None

        existing_server = _create_mock_mcp_server(name='Existing Server')
        ap.persistence_mgr.execute_async = AsyncMock(return_value=_create_mock_result(first_item=existing_server))
        ap.persistence_mgr.serialize_model = Mock(return_value={})

        service = _service(ap)

        # Execute & Verify
        with pytest.raises(ValueError, match='MCP server already exists: Existing Server'):
            await service.create_mcp_server(_CONTEXT, {'name': 'Existing Server'})

    async def test_create_mcp_server_loads_server(self):
        """Loads server into tool_mgr when enabled."""
        # Setup
        ap = SimpleNamespace()
        ap.persistence_mgr = SimpleNamespace()
        ap.instance_config = SimpleNamespace()
        ap.instance_config.data = {'system': {'limitation': {'max_extensions': -1}}}
        ap.tool_mgr = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader.host_mcp_server = AsyncMock()
        ap.tool_mgr.mcp_tool_loader._hosted_mcp_tasks = []

        # Create mock server entity
        server_entity = _create_mock_mcp_server(server_uuid='new-uuid', enable=True)

        call_count = 0

        async def mock_execute(query):
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                return _create_mock_result([])  # Empty result for duplicate-name check
            elif call_count == 2:
                return Mock()  # Insert
            return _create_mock_result(first_item=server_entity)  # Select created

        ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
        ap.persistence_mgr.serialize_model = Mock(
            return_value={'uuid': 'new-uuid', 'name': 'New Server', 'enable': True}
        )

        service = _service(ap)

        # Execute
        await service.create_mcp_server(_CONTEXT, {'name': 'New Server', 'enable': True})

        # Verify - host_mcp_server was called
        ap.tool_mgr.mcp_tool_loader.host_mcp_server.assert_called_once()

    async def test_create_mcp_server_does_not_start_host_until_transaction_commits(self):
        """The Runtime must not observe a server row that can still roll back."""

        gate = asyncio.get_running_loop().create_future()

        class PersistenceManagerStub:
            def create_after_commit_gate(self):
                return gate

        ap = SimpleNamespace()
        ap.persistence_mgr = PersistenceManagerStub()
        ap.instance_config = SimpleNamespace(data={'system': {'limitation': {'max_extensions': -1}}})
        observed = []

        async def host_mcp_server(context, config):
            observed.append((context, config))

        ap.tool_mgr = SimpleNamespace(
            mcp_tool_loader=SimpleNamespace(
                host_mcp_server=host_mcp_server,
                _hosted_mcp_tasks=[],
            )
        )
        server_entity = _create_mock_mcp_server(server_uuid='new-uuid', enable=True)
        results = [
            _create_mock_result([]),
            Mock(),
            _create_mock_result(first_item=server_entity),
        ]
        ap.persistence_mgr.execute_async = AsyncMock(side_effect=results)
        ap.persistence_mgr.serialize_model = Mock(
            return_value={'uuid': 'new-uuid', 'name': 'New Server', 'enable': True}
        )
        service = _service(ap)

        await service.create_mcp_server(_CONTEXT, {'name': 'New Server', 'enable': True})
        await asyncio.sleep(0)
        assert observed == []

        gate.set_result(None)
        await ap.tool_mgr.mcp_tool_loader._hosted_mcp_tasks[0]
        assert observed == [
            (
                _CONTEXT,
                {'uuid': 'new-uuid', 'name': 'New Server', 'enable': True},
            )
        ]

    async def test_create_mcp_server_disabled_no_load(self):
        """Does not load server when disabled."""
        # Setup
        ap = SimpleNamespace()
        ap.persistence_mgr = SimpleNamespace()
        ap.instance_config = SimpleNamespace()
        ap.instance_config.data = {'system': {'limitation': {'max_extensions': -1}}}
        ap.tool_mgr = None

        mock_result = _create_mock_result([])
        ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
        ap.persistence_mgr.serialize_model = Mock(return_value={'uuid': 'new-uuid'})

        service = _service(ap)

        # Execute with enable=False
        server_uuid = await service.create_mcp_server(_CONTEXT, {'name': 'New Server', 'enable': False})

        # Verify - no tool_mgr load attempt
        assert server_uuid is not None


class TestMCPServiceGetMCPServerByName:
    """Tests for get_mcp_server_by_name method."""

    async def test_get_mcp_server_by_name_found(self):
        """Returns MCP server when found by name."""
        # Setup
        ap = SimpleNamespace()
        ap.persistence_mgr = SimpleNamespace()

        server = _create_mock_mcp_server(name='Found Server')
        mock_result = _create_mock_result(first_item=server)
        ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)
        ap.persistence_mgr.serialize_model = Mock(
            return_value={
                'uuid': 'test-uuid',
                'name': 'Found Server',
                'runtime_info': None,
            }
        )
        ap.tool_mgr = SimpleNamespace(mcp_tool_loader=SimpleNamespace(get_session=Mock(return_value=None)))

        service = _service(ap)
        # Execute
        result = await service.get_mcp_server_by_name(_CONTEXT, 'Found Server')

        # Verify
        assert result is not None
        assert result['name'] == 'Found Server'

    async def test_get_mcp_server_by_name_not_found(self):
        """Returns None when MCP server not found."""
        # Setup
        ap = SimpleNamespace()
        ap.persistence_mgr = SimpleNamespace()

        mock_result = _create_mock_result(first_item=None)
        ap.persistence_mgr.execute_async = AsyncMock(return_value=mock_result)

        service = _service(ap)

        # Execute
        result = await service.get_mcp_server_by_name(_CONTEXT, 'Nonexistent Server')

        # Verify
        assert result is None


class TestMCPServiceUpdateMCPServer:
    """Tests for update_mcp_server method."""

    async def test_update_mcp_server_disable_enabled_server(self):
        """Removes server when disabling previously enabled server."""
        # Setup
        ap = SimpleNamespace()
        ap.persistence_mgr = SimpleNamespace()
        ap.tool_mgr = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader.sessions = {'Old Server': Mock()}
        ap.tool_mgr.mcp_tool_loader.remove_mcp_server = AsyncMock()
        ap.tool_mgr.mcp_tool_loader.has_session = Mock(return_value=True)

        old_server = _create_mock_mcp_server(name='Old Server', enable=True)
        updated_server = _create_mock_mcp_server(name='Old Server', enable=False)

        call_count = 0

        async def mock_execute(query):
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                return _create_mock_result(first_item=old_server)
            if call_count == 2:
                return _create_mock_result()
            return _create_mock_result(first_item=updated_server)

        ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
        ap.persistence_mgr.serialize_model = Mock(
            side_effect=lambda _model, entity: {
                'uuid': 'test-uuid',
                'name': entity.name,
                'enable': entity.enable,
            }
        )

        service = _service(ap)

        # Execute - disable server
        await service.update_mcp_server(_CONTEXT, 'test-uuid', {'enable': False})

        # Verify - server was removed
        ap.tool_mgr.mcp_tool_loader.remove_mcp_server.assert_called_once()

    async def test_update_mcp_server_enable_disabled_server(self):
        """Loads server when enabling previously disabled server."""
        # Setup
        ap = SimpleNamespace()
        ap.persistence_mgr = SimpleNamespace()
        ap.tool_mgr = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader.sessions = {}
        ap.tool_mgr.mcp_tool_loader.host_mcp_server = AsyncMock()
        ap.tool_mgr.mcp_tool_loader._hosted_mcp_tasks = []
        ap.tool_mgr.mcp_tool_loader.has_session = Mock(return_value=False)

        old_server = _create_mock_mcp_server(name='Old Server', enable=False)

        updated_server = _create_mock_mcp_server(name='Old Server', enable=True)

        call_count = 0

        async def mock_execute(query):
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                return _create_mock_result(first_item=old_server)
            elif call_count == 2:
                return Mock()  # Update
            return _create_mock_result(first_item=updated_server)  # Select updated

        ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
        ap.persistence_mgr.serialize_model = Mock(
            return_value={'uuid': 'test-uuid', 'name': 'Old Server', 'enable': True}
        )

        service = _service(ap)

        # Execute - enable server
        await service.update_mcp_server(_CONTEXT, 'test-uuid', {'enable': True})

        # Verify - server was loaded
        ap.tool_mgr.mcp_tool_loader.host_mcp_server.assert_called_once()

    async def test_update_mcp_server_update_enabled_server(self):
        """Removes and reloads server when updating enabled server."""
        # Setup
        ap = SimpleNamespace()
        ap.persistence_mgr = SimpleNamespace()
        ap.tool_mgr = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader.sessions = {'Old Server': Mock()}
        ap.tool_mgr.mcp_tool_loader.remove_mcp_server = AsyncMock()
        ap.tool_mgr.mcp_tool_loader.host_mcp_server = AsyncMock()
        ap.tool_mgr.mcp_tool_loader._hosted_mcp_tasks = []
        ap.tool_mgr.mcp_tool_loader.has_session = Mock(return_value=True)

        old_server = _create_mock_mcp_server(name='Old Server', enable=True)

        # Mock for: first select -> update -> second select (for updated server)
        call_count = 0

        async def mock_execute(query):
            nonlocal call_count
            call_count += 1
            # All selects return the server
            return _create_mock_result(first_item=old_server)

        ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
        ap.persistence_mgr.serialize_model = Mock(
            return_value={'uuid': 'test-uuid', 'name': 'Old Server', 'enable': True}
        )

        service = _service(ap)

        # Execute - update enabled server (keep enabled, update extra_args)
        await service.update_mcp_server(_CONTEXT, 'test-uuid', {'enable': True, 'extra_args': {'new': 'args'}})

        # Verify - remove and reload
        ap.tool_mgr.mcp_tool_loader.remove_mcp_server.assert_called_once_with(_CONTEXT, 'Old Server')
        ap.tool_mgr.mcp_tool_loader.host_mcp_server.assert_called_once()

    async def test_update_mcp_server_no_tool_mgr(self):
        """Updates persistence without tool_mgr operations."""
        # Setup
        ap = SimpleNamespace()
        ap.persistence_mgr = SimpleNamespace()
        # Set mcp_tool_loader to None, not tool_mgr itself
        ap.tool_mgr = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader = None

        old_server = _create_mock_mcp_server(name='Server', enable=True)

        # Mock execute for select and update
        call_count = 0

        async def mock_execute(query):
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                return _create_mock_result(first_item=old_server)
            return Mock()  # Update

        ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
        ap.persistence_mgr.serialize_model = Mock(
            return_value={
                'uuid': 'test-uuid',
                'name': 'Server',
                'enable': True,
            }
        )

        service = _service(ap)

        # Execute - should not raise
        await service.update_mcp_server(_CONTEXT, 'test-uuid', {'enable': False})

        # Verify - persistence was called
        assert ap.persistence_mgr.execute_async.call_count >= 2

    async def test_update_restores_existing_masked_secrets_and_preserves_explicit_changes(self):
        ap = SimpleNamespace()
        ap.persistence_mgr = SimpleNamespace()
        ap.tool_mgr = SimpleNamespace(mcp_tool_loader=None)
        old_server = _create_mock_mcp_server(name='Server', enable=True)
        old_data = {
            'uuid': 'test-uuid',
            'name': 'Server',
            'enable': True,
            'mode': 'streamable_http',
            'extra_args': {
                'headers': {
                    'Authorization': 'Bearer original-secret',
                    'X-API-Key': 'original-api-key',
                    'Cookie': 'original-cookie',
                }
            },
        }
        captured_updates = []

        async def mock_execute(statement):
            if not captured_updates:
                captured_updates.append(None)
                return _create_mock_result(first_item=old_server)
            captured_updates[0] = statement
            return _create_mock_result()

        ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
        ap.persistence_mgr.serialize_model = Mock(return_value=old_data)
        service = _service(ap)

        await service.update_mcp_server(
            _CONTEXT,
            'test-uuid',
            {
                'extra_args': {
                    'headers': {
                        'Authorization': '***',
                        'X-API-Key': 'replacement-api-key',
                        'Cookie': '',
                    }
                }
            },
        )

        persisted = captured_updates[0].compile().params['extra_args']
        assert persisted['headers'] == {
            'Authorization': 'Bearer original-secret',
            'X-API-Key': 'replacement-api-key',
            'Cookie': '',
        }

    async def test_update_rejects_masked_secret_without_existing_value(self):
        ap = SimpleNamespace()
        ap.persistence_mgr = SimpleNamespace()
        ap.tool_mgr = SimpleNamespace(mcp_tool_loader=None)
        old_server = _create_mock_mcp_server(name='Server', enable=True)
        ap.persistence_mgr.execute_async = AsyncMock(return_value=_create_mock_result(first_item=old_server))
        ap.persistence_mgr.serialize_model = Mock(
            return_value={
                'uuid': 'test-uuid',
                'name': 'Server',
                'enable': True,
                'extra_args': {'headers': {'Accept': 'application/json'}},
            }
        )
        service = _service(ap)

        with pytest.raises(ValueError, match='Masked MCP secret has no existing value'):
            await service.update_mcp_server(
                _CONTEXT,
                'test-uuid',
                {'extra_args': {'headers': {'Authorization': '***'}}},
            )

        assert ap.persistence_mgr.execute_async.await_count == 1


class TestMCPServiceDeleteMCPServer:
    """Tests for delete_mcp_server method."""

    async def test_delete_mcp_server_calls_remove_and_delete(self):
        """Calls both persistence delete and tool_mgr remove."""
        # Setup
        ap = SimpleNamespace()
        ap.persistence_mgr = SimpleNamespace()
        ap.tool_mgr = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader.sessions = {'Server to Delete': Mock()}
        ap.tool_mgr.mcp_tool_loader.remove_mcp_server = AsyncMock()
        ap.tool_mgr.mcp_tool_loader.has_session = Mock(return_value=True)

        server = _create_mock_mcp_server(name='Server to Delete')

        call_count = 0

        async def mock_execute(query):
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                return _create_mock_result(first_item=server)
            return Mock()  # Delete

        ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
        ap.persistence_mgr.serialize_model = Mock(
            return_value={
                'uuid': 'test-uuid',
                'name': 'Server to Delete',
                'enable': True,
            }
        )

        service = _service(ap)

        # Execute
        await service.delete_mcp_server(_CONTEXT, 'test-uuid')

        # Verify
        ap.tool_mgr.mcp_tool_loader.remove_mcp_server.assert_called_once_with(_CONTEXT, 'Server to Delete')
        ap.persistence_mgr.execute_async.assert_called()

    async def test_delete_mcp_server_not_in_sessions(self):
        """Does not attempt remove if server not in sessions."""
        # Setup
        ap = SimpleNamespace()
        ap.persistence_mgr = SimpleNamespace()
        ap.tool_mgr = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader.sessions = {}  # Server not in sessions
        ap.tool_mgr.mcp_tool_loader.remove_mcp_server = AsyncMock()
        ap.tool_mgr.mcp_tool_loader.has_session = Mock(return_value=False)

        server = _create_mock_mcp_server(name='Not in Sessions')

        call_count = 0

        async def mock_execute(query):
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                return _create_mock_result(first_item=server)
            return Mock()

        ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)
        ap.persistence_mgr.serialize_model = Mock(
            return_value={
                'uuid': 'test-uuid',
                'name': 'Not in Sessions',
                'enable': True,
            }
        )

        service = _service(ap)

        # Execute
        await service.delete_mcp_server(_CONTEXT, 'test-uuid')

        # Verify - remove not called (server not in sessions)
        ap.tool_mgr.mcp_tool_loader.remove_mcp_server.assert_not_called()

    async def test_delete_mcp_server_nonexistent_uuid(self):
        """Delete operation completes even for nonexistent UUID."""
        # Setup
        ap = SimpleNamespace()
        ap.persistence_mgr = SimpleNamespace()
        ap.tool_mgr = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader.sessions = {}
        ap.tool_mgr.mcp_tool_loader.remove_mcp_server = AsyncMock()
        ap.tool_mgr.mcp_tool_loader.has_session = Mock(return_value=False)

        # No server found
        call_count = 0

        async def mock_execute(query):
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                return _create_mock_result(first_item=None)
            return Mock()

        ap.persistence_mgr.execute_async = AsyncMock(side_effect=mock_execute)

        service = _service(ap)

        with pytest.raises(WorkspaceNotFoundError, match='MCP server not found'):
            await service.delete_mcp_server(_CONTEXT, 'nonexistent-uuid')

        assert ap.persistence_mgr.execute_async.await_count == 1


class TestMCPServiceTestMCPServer:
    """Tests for test_mcp_server method."""

    async def test_transient_stdio_test_rejected_by_instance_gate(self):
        ap = SimpleNamespace(
            instance_config=SimpleNamespace(data={'mcp': {'stdio': {'enabled': False}}}),
            tool_mgr=SimpleNamespace(mcp_tool_loader=SimpleNamespace(load_mcp_server=AsyncMock())),
            task_mgr=SimpleNamespace(create_user_task=Mock()),
        )
        service = _service(ap)

        with pytest.raises(MCPStdioDisabledError, match='disabled by instance policy'):
            await service.test_mcp_server(
                _CONTEXT,
                '_',
                {'name': 'local', 'mode': 'stdio', 'enable': True, 'extra_args': {}},
            )

        ap.tool_mgr.mcp_tool_loader.load_mcp_server.assert_not_awaited()
        ap.task_mgr.create_user_task.assert_not_called()

    async def test_test_mcp_server_existing_server(self):
        """Tests existing MCP server connection."""
        # Setup
        ap = SimpleNamespace()
        ap.tool_mgr = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader = SimpleNamespace()

        from langbot.pkg.provider.tools.loaders.mcp import MCPSessionStatus

        mock_session = MagicMock()
        mock_session.status = MCPSessionStatus.ERROR
        mock_session.start = AsyncMock()
        mock_session.refresh = AsyncMock()
        ap.tool_mgr.mcp_tool_loader.get_session = Mock(return_value=mock_session)

        ap.task_mgr = SimpleNamespace()

        service = _service(ap)
        service._require_server = AsyncMock(return_value=(_CONTEXT, {'name': 'existing-server'}))

        def create_user_task(coroutine, **_kwargs):
            coroutine.close()
            return SimpleNamespace(id=123)

        ap.task_mgr.create_user_task = Mock(side_effect=create_user_task)

        # Execute
        task_id = await service.test_mcp_server(_CONTEXT, 'existing-server', {})

        # Verify - returns task ID
        assert task_id == 123

    async def test_test_mcp_server_not_found_raises(self):
        """Raises ValueError when server not found."""
        # Setup
        ap = SimpleNamespace()
        ap.tool_mgr = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader.get_session = Mock(return_value=None)

        service = _service(ap)
        service._require_server = AsyncMock(side_effect=WorkspaceNotFoundError('MCP server not found'))

        # Execute & Verify
        with pytest.raises(WorkspaceNotFoundError, match='MCP server not found'):
            await service.test_mcp_server(_CONTEXT, 'nonexistent-server', {})

    async def test_test_mcp_server_new_server(self):
        """Tests new MCP server with underscore name."""
        # Setup
        ap = SimpleNamespace()
        ap.tool_mgr = SimpleNamespace()
        ap.tool_mgr.mcp_tool_loader = SimpleNamespace()

        mock_session = MagicMock()
        mock_session.start = AsyncMock()
        ap.tool_mgr.mcp_tool_loader.load_mcp_server = AsyncMock(return_value=mock_session)

        ap.task_mgr = SimpleNamespace()

        service = _service(ap)

        def create_user_task(coroutine, **_kwargs):
            coroutine.close()
            return SimpleNamespace(id=456)

        ap.task_mgr.create_user_task = Mock(side_effect=create_user_task)

        # Execute with '_' name (new server)
        task_id = await service.test_mcp_server(_CONTEXT, '_', {'name': 'New Server'})

        # Verify - load_mcp_server called
        ap.tool_mgr.mcp_tool_loader.load_mcp_server.assert_called_once()
        assert task_id == 456