LangBot/src/langbot/pkg/api/http/service/mcp.py

from __future__ import annotations

import copy
import re
import uuid

import sqlalchemy

from ....core import app, taskmgr
from ....core.task_boundary import create_detached_task
from ....entity.persistence import mcp as persistence_mcp
from ....entity.persistence import plugin as persistence_plugin
from ....provider.tools.loaders.mcp import MCPSessionStatus, RuntimeMCPSession
from ....provider.tools.loaders.mcp_policy import require_stdio_mcp_enabled
from ....workspace.errors import WorkspaceNotFoundError
from ..context import ExecutionContext
from .secrets import is_url_key, redact_url_secrets, restore_url_secret_placeholders
from .tenant import TenantContext, require_workspace_uuid, scope_statement


_SECRET_MASK = '***'
_MISSING_SECRET = object()
_SENSITIVE_CONFIG_NAMES = frozenset(
    {
        'api_key',
        'apikey',
        'auth',
        'authorization',
        'cookie',
        'credentials',
        'database_url',
        'dsn',
        'key',
        'proxy_authorization',
        'set_cookie',
    }
)
_SENSITIVE_CONFIG_TOKENS = frozenset(
    {
        'credential',
        'credentials',
        'passwd',
        'password',
        'secret',
        'token',
    }
)
_SENSITIVE_KEY_QUALIFIERS = frozenset(
    {
        'access',
        'api',
        'auth',
        'bearer',
        'client',
        'debug',
        'encryption',
        'private',
        'signing',
    }
)


def _normalize_config_key(key: object) -> str:
    value = re.sub(r'([a-z0-9])([A-Z])', r'\1_\2', str(key or ''))
    return re.sub(r'[^a-zA-Z0-9]+', '_', value).strip('_').lower()


def _is_sensitive_config_key(key: object) -> bool:
    normalized = _normalize_config_key(key)
    if normalized in _SENSITIVE_CONFIG_NAMES:
        return True
    tokens = frozenset(token for token in normalized.split('_') if token)
    if tokens & _SENSITIVE_CONFIG_TOKENS:
        return True
    return 'key' in tokens and bool(tokens & _SENSITIVE_KEY_QUALIFIERS)


def _mask_secret_structure(value):
    if isinstance(value, dict):
        return {key: _mask_secret_structure(item) for key, item in value.items()}
    if isinstance(value, list):
        return [_mask_secret_structure(item) for item in value]
    if isinstance(value, tuple):
        return tuple(_mask_secret_structure(item) for item in value)
    if value is None or value == '':
        return value
    return _SECRET_MASK


def redact_mcp_secrets(value):
    """Return a recursively redacted copy of MCP configuration data."""

    if isinstance(value, dict):
        return {
            key: (
                _mask_secret_structure(item)
                if _is_sensitive_config_key(key)
                else redact_url_secrets(item)
                if is_url_key(key)
                else redact_mcp_secrets(item)
            )
            for key, item in value.items()
        }
    if isinstance(value, list):
        return [redact_mcp_secrets(item) for item in value]
    if isinstance(value, tuple):
        return tuple(redact_mcp_secrets(item) for item in value)
    return value


def restore_mcp_secret_placeholders(value, current_value=_MISSING_SECRET, *, sensitive: bool = False):
    """Restore masked leaves from the current MCP config before a write."""

    if sensitive and value == _SECRET_MASK:
        if current_value is _MISSING_SECRET:
            raise ValueError('Masked MCP secret has no existing value')
        return copy.deepcopy(current_value)
    if isinstance(value, dict):
        current_mapping = current_value if isinstance(current_value, dict) else {}
        return {
            key: (
                restore_url_secret_placeholders(
                    item,
                    current_mapping.get(key, _MISSING_SECRET),
                )
                if not sensitive and not _is_sensitive_config_key(key) and is_url_key(key)
                else restore_mcp_secret_placeholders(
                    item,
                    current_mapping.get(key, _MISSING_SECRET),
                    sensitive=sensitive or _is_sensitive_config_key(key),
                )
            )
            for key, item in value.items()
        }
    if isinstance(value, list):
        current_items = current_value if isinstance(current_value, (list, tuple)) else ()
        return [
            restore_mcp_secret_placeholders(
                item,
                current_items[index] if index < len(current_items) else _MISSING_SECRET,
                sensitive=sensitive,
            )
            for index, item in enumerate(value)
        ]
    if isinstance(value, tuple):
        current_items = current_value if isinstance(current_value, (list, tuple)) else ()
        return tuple(
            restore_mcp_secret_placeholders(
                item,
                current_items[index] if index < len(current_items) else _MISSING_SECRET,
                sensitive=sensitive,
            )
            for index, item in enumerate(value)
        )
    return value


class MCPService:
    """Workspace-scoped MCP configuration and runtime facade."""

    ap: app.Application

    def __init__(self, ap: app.Application) -> None:
        self.ap = ap

    async def _execution_context(self, context: TenantContext) -> ExecutionContext:
        workspace_uuid = require_workspace_uuid(context)
        instance_uuid = str(getattr(context, 'instance_uuid', '') or '').strip()
        generation = getattr(context, 'placement_generation', None)
        if not instance_uuid or not isinstance(generation, int) or isinstance(generation, bool) or generation <= 0:
            raise ValueError('MCP operations require an explicit fenced execution context')
        binding = await self.ap.workspace_service.get_execution_binding(
            workspace_uuid,
            expected_generation=generation,
        )
        if binding.instance_uuid != instance_uuid:
            raise ValueError('MCP execution context belongs to another LangBot instance')
        return ExecutionContext(
            instance_uuid=instance_uuid,
            workspace_uuid=workspace_uuid,
            placement_generation=generation,
            bot_uuid=getattr(context, 'bot_uuid', None),
            pipeline_uuid=getattr(context, 'pipeline_uuid', None),
            query_uuid=getattr(context, 'query_uuid', None),
        )

    async def get_runtime_info(self, context: TenantContext, server_name: str) -> dict | None:
        execution_context = await self._execution_context(context)
        session = self.ap.tool_mgr.mcp_tool_loader.get_session(execution_context, server_name)
        return session.get_runtime_info_dict() if session else None

    async def get_mcp_servers(self, context: TenantContext, contain_runtime_info: bool = False) -> list[dict]:
        execution_context = await self._execution_context(context)
        result = await self.ap.persistence_mgr.execute_async(
            scope_statement(sqlalchemy.select(persistence_mcp.MCPServer), persistence_mcp.MCPServer, context)
        )
        serialized_servers = [
            redact_mcp_secrets(self.ap.persistence_mgr.serialize_model(persistence_mcp.MCPServer, server))
            for server in result.all()
        ]
        if contain_runtime_info:
            for server in serialized_servers:
                session = self.ap.tool_mgr.mcp_tool_loader.get_session(execution_context, server['name'])
                server['runtime_info'] = session.get_runtime_info_dict() if session else None
        return serialized_servers

    async def create_mcp_server(self, context: TenantContext, server_data: dict) -> str:
        execution_context = await self._execution_context(context)
        workspace_uuid = execution_context.workspace_uuid

        # This gate is independent of Box availability.  Cloud v2 disables
        # stdio MCP even though Box Runtime itself remains available.
        require_stdio_mcp_enabled(self.ap, server_data)

        limitation = self.ap.instance_config.data.get('system', {}).get('limitation', {})
        max_extensions = limitation.get('max_extensions', -1)
        if max_extensions >= 0:
            mcp_count_result = await self.ap.persistence_mgr.execute_async(
                sqlalchemy.select(sqlalchemy.func.count(persistence_mcp.MCPServer.uuid)).where(
                    persistence_mcp.MCPServer.workspace_uuid == workspace_uuid
                )
            )
            plugin_count_result = await self.ap.persistence_mgr.execute_async(
                sqlalchemy.select(sqlalchemy.func.count())
                .select_from(persistence_plugin.PluginSetting)
                .where(persistence_plugin.PluginSetting.workspace_uuid == workspace_uuid)
            )
            if (mcp_count_result.scalar() or 0) + (plugin_count_result.scalar() or 0) >= max_extensions:
                raise ValueError(f'Maximum number of extensions ({max_extensions}) reached')

        payload = dict(server_data)
        payload.pop('workspace_uuid', None)
        server_name = str(payload.get('name') or '').strip()
        if not server_name:
            raise ValueError('MCP server name is required')
        payload['name'] = server_name
        payload['workspace_uuid'] = workspace_uuid
        payload['uuid'] = str(uuid.uuid4())

        existing_result = await self.ap.persistence_mgr.execute_async(
            sqlalchemy.select(persistence_mcp.MCPServer).where(
                persistence_mcp.MCPServer.workspace_uuid == workspace_uuid,
                persistence_mcp.MCPServer.name == server_name,
            )
        )
        if existing_result.first() is not None:
            raise ValueError(f'MCP server already exists: {server_name}')

        await self.ap.persistence_mgr.execute_async(sqlalchemy.insert(persistence_mcp.MCPServer).values(payload))
        created = await self._get_mcp_server_by_uuid_raw(execution_context, payload['uuid'])
        if created and self.ap.tool_mgr.mcp_tool_loader:
            task = create_detached_task(
                self.ap.tool_mgr.mcp_tool_loader.host_mcp_server(execution_context, created),
                after_commit_manager=self.ap.persistence_mgr,
                workspace_uuid=execution_context.workspace_uuid,
            )
            tracker = getattr(
                self.ap.tool_mgr.mcp_tool_loader,
                'track_hosted_task',
                None,
            )
            if callable(tracker):
                tracker(task, execution_context)
            else:
                self.ap.tool_mgr.mcp_tool_loader._hosted_mcp_tasks.append(task)
        return payload['uuid']

    async def get_mcp_server_by_uuid(self, context: TenantContext, server_uuid: str) -> dict | None:
        execution_context = await self._execution_context(context)
        server_data = await self._get_mcp_server_by_uuid_raw(execution_context, server_uuid)
        return redact_mcp_secrets(server_data) if server_data is not None else None

    async def _get_mcp_server_by_uuid_raw(
        self,
        execution_context: ExecutionContext,
        server_uuid: str,
    ) -> dict | None:
        result = await self.ap.persistence_mgr.execute_async(
            scope_statement(
                sqlalchemy.select(persistence_mcp.MCPServer).where(persistence_mcp.MCPServer.uuid == server_uuid),
                persistence_mcp.MCPServer,
                execution_context,
            )
        )
        server = result.first()
        return self.ap.persistence_mgr.serialize_model(persistence_mcp.MCPServer, server) if server else None

    async def get_mcp_server_by_name(self, context: TenantContext, server_name: str) -> dict | None:
        execution_context = await self._execution_context(context)
        server_data = await self._get_mcp_server_by_name_raw(execution_context, server_name)
        if server_data is None:
            return None
        session = self.ap.tool_mgr.mcp_tool_loader.get_session(execution_context, server_name)
        response_data = {
            **server_data,
            'runtime_info': session.get_runtime_info_dict() if session else None,
        }
        return redact_mcp_secrets(response_data)

    async def _get_mcp_server_by_name_raw(
        self,
        execution_context: ExecutionContext,
        server_name: str,
    ) -> dict | None:
        result = await self.ap.persistence_mgr.execute_async(
            scope_statement(
                sqlalchemy.select(persistence_mcp.MCPServer).where(persistence_mcp.MCPServer.name == server_name),
                persistence_mcp.MCPServer,
                execution_context,
            )
        )
        server = result.first()
        if server is None:
            return None
        return self.ap.persistence_mgr.serialize_model(persistence_mcp.MCPServer, server)

    async def update_mcp_server(self, context: TenantContext, server_uuid: str, server_data: dict) -> None:
        execution_context = await self._execution_context(context)
        old_server = await self._get_mcp_server_by_uuid_raw(execution_context, server_uuid)
        if old_server is None:
            raise WorkspaceNotFoundError('MCP server not found')

        payload = dict(server_data)
        payload.pop('uuid', None)
        payload.pop('workspace_uuid', None)
        payload = restore_mcp_secret_placeholders(payload, old_server)
        if 'name' in payload:
            payload['name'] = str(payload['name'] or '').strip()
            if not payload['name']:
                raise ValueError('MCP server name is required')
            duplicate = await self._get_mcp_server_by_name_raw(execution_context, payload['name'])
            if duplicate is not None and duplicate['uuid'] != server_uuid:
                raise ValueError(f'MCP server already exists: {payload["name"]}')

        effective_server = {**old_server, **payload}
        # Existing disabled rows remain readable/deletable.  Switching away
        # from stdio or explicitly disabling one is also allowed, but an
        # update may never leave a disabled stdio server enabled.
        if bool(effective_server.get('enable', True)):
            require_stdio_mcp_enabled(self.ap, effective_server)

        result = await self.ap.persistence_mgr.execute_async(
            scope_statement(
                sqlalchemy.update(persistence_mcp.MCPServer)
                .where(persistence_mcp.MCPServer.uuid == server_uuid)
                .values(payload),
                persistence_mcp.MCPServer,
                execution_context,
            )
        )
        if getattr(result, 'rowcount', None) == 0:
            raise WorkspaceNotFoundError('MCP server not found')

        loader = self.ap.tool_mgr.mcp_tool_loader
        if loader is None:
            return
        old_name = old_server['name']
        old_enable = bool(old_server['enable'])
        updated = await self._get_mcp_server_by_uuid_raw(execution_context, server_uuid)
        if updated is None:
            raise WorkspaceNotFoundError('MCP server not found')
        new_enable = bool(updated['enable'])
        if old_enable and loader.has_session(execution_context, old_name):
            await loader.remove_mcp_server(execution_context, old_name)
        if new_enable:
            task = create_detached_task(
                loader.host_mcp_server(execution_context, updated),
                after_commit_manager=self.ap.persistence_mgr,
                workspace_uuid=execution_context.workspace_uuid,
            )
            tracker = getattr(loader, 'track_hosted_task', None)
            if callable(tracker):
                tracker(task, execution_context)
            else:
                loader._hosted_mcp_tasks.append(task)

    async def delete_mcp_server(self, context: TenantContext, server_uuid: str) -> None:
        execution_context = await self._execution_context(context)
        server = await self._get_mcp_server_by_uuid_raw(execution_context, server_uuid)
        if server is None:
            raise WorkspaceNotFoundError('MCP server not found')
        result = await self.ap.persistence_mgr.execute_async(
            scope_statement(
                sqlalchemy.delete(persistence_mcp.MCPServer).where(persistence_mcp.MCPServer.uuid == server_uuid),
                persistence_mcp.MCPServer,
                execution_context,
            )
        )
        if getattr(result, 'rowcount', None) == 0:
            raise WorkspaceNotFoundError('MCP server not found')
        loader = self.ap.tool_mgr.mcp_tool_loader
        if loader and loader.has_session(execution_context, server['name']):
            await loader.remove_mcp_server(execution_context, server['name'])

    async def _require_server(self, context: TenantContext, server_name: str) -> tuple[ExecutionContext, dict]:
        execution_context = await self._execution_context(context)
        server = await self._get_mcp_server_by_name_raw(execution_context, server_name)
        if server is None:
            raise WorkspaceNotFoundError('MCP server not found')
        return execution_context, server

    async def get_mcp_server_resources(self, context: TenantContext, server_name: str) -> list[dict]:
        execution_context, _ = await self._require_server(context, server_name)
        return await self.ap.tool_mgr.mcp_tool_loader.get_resources(execution_context, server_name)

    async def get_mcp_server_resource_templates(self, context: TenantContext, server_name: str) -> list[dict]:
        execution_context, _ = await self._require_server(context, server_name)
        return await self.ap.tool_mgr.mcp_tool_loader.get_resource_templates(execution_context, server_name)

    async def read_mcp_server_resource_envelope(
        self,
        context: TenantContext,
        server_name: str,
        uri: str,
        *,
        max_bytes: int | None = None,
        include_blob: bool = False,
    ) -> dict:
        execution_context, _ = await self._require_server(context, server_name)
        kwargs = {'include_blob': include_blob, 'source': 'ui_preview'}
        if max_bytes is not None:
            kwargs['max_bytes'] = max_bytes
        return await self.ap.tool_mgr.mcp_tool_loader.read_resource_envelope(
            execution_context,
            server_name,
            uri,
            **kwargs,
        )

    async def read_mcp_server_resource(self, context: TenantContext, server_name: str, uri: str) -> list[dict]:
        execution_context, _ = await self._require_server(context, server_name)
        return await self.ap.tool_mgr.mcp_tool_loader.read_resource(execution_context, server_name, uri)

    async def test_mcp_server(self, context: TenantContext, server_name: str, server_data: dict) -> int:
        execution_context = await self._execution_context(context)
        runtime_mcp_session: RuntimeMCPSession | None = None
        test_session: RuntimeMCPSession | None = None
        ctx = taskmgr.TaskContext.new()

        if server_name != '_':
            _, persisted_server = await self._require_server(execution_context, server_name)
            require_stdio_mcp_enabled(self.ap, persisted_server)
            runtime_mcp_session = self.ap.tool_mgr.mcp_tool_loader.get_session(execution_context, server_name)
            if runtime_mcp_session is None:
                raise WorkspaceNotFoundError('MCP server not found')
            persisted_session = runtime_mcp_session

            async def _refresh_and_report() -> None:
                needs_start = persisted_session.status == MCPSessionStatus.ERROR or persisted_session.session is None
                if needs_start:
                    await persisted_session.start()
                else:
                    try:
                        await persisted_session.refresh()
                    except Exception:
                        await persisted_session.start()
                ctx.metadata['runtime_info'] = persisted_session.get_runtime_info_dict()

            coroutine = _refresh_and_report()
        else:
            payload = dict(server_data)
            payload.pop('workspace_uuid', None)
            payload['workspace_uuid'] = execution_context.workspace_uuid
            require_stdio_mcp_enabled(self.ap, payload)
            runtime_mcp_session = await self.ap.tool_mgr.mcp_tool_loader.load_mcp_server(
                execution_context,
                payload,
            )
            test_session = runtime_mcp_session

            async def _run_and_cleanup() -> None:
                try:
                    await test_session.start()
                    ctx.metadata['runtime_info'] = test_session.get_runtime_info_dict()
                finally:
                    try:
                        await test_session.shutdown()
                    except Exception as exc:
                        self.ap.logger.warning(
                            f'Failed to tear down transient MCP test session '
                            f'{test_session.server_name}: {type(exc).__name__}: {exc}'
                        )

            coroutine = _run_and_cleanup()

        try:
            wrapper = self.ap.task_mgr.create_user_task(
                coroutine,
                kind='mcp-operation',
                name=f'mcp-test-{execution_context.workspace_uuid}-{server_name}',
                label=f'Testing MCP server {server_name}',
                context=ctx,
                instance_uuid=execution_context.instance_uuid,
                workspace_uuid=execution_context.workspace_uuid,
                placement_generation=execution_context.placement_generation,
            )
        except taskmgr.TaskCapacityError:
            if test_session is not None:
                try:
                    await test_session.shutdown()
                except Exception as exc:
                    self.ap.logger.warning(
                        f'Failed to tear down rejected transient MCP test session '
                        f'{test_session.server_name}: {type(exc).__name__}: {exc}'
                    )
            raise
        return wrapper.id

    async def get_mcp_server_logs(
        self,
        context: TenantContext,
        server_name: str,
        limit: int = 200,
        level: str | None = None,
    ) -> list[dict]:
        execution_context, _ = await self._require_server(context, server_name)
        session = self.ap.tool_mgr.mcp_tool_loader.get_session(execution_context, server_name)
        if not session:
            return []
        logs = list(session._log_buffer)
        if level:
            logs = [log for log in logs if log.get('level') == level]
        return logs[-limit:]