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