LangBot/src/langbot/pkg/provider/tools/loaders/mcp.py

from __future__ import annotations

import enum
import os
import typing
from contextlib import AsyncExitStack
import traceback
from langbot_plugin.api.entities.events import pipeline_query
import sqlalchemy
import asyncio
import httpx

import pydantic
import uuid as uuid_module
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from mcp.client.sse import sse_client
from mcp.client.streamable_http import streamable_http_client
from mcp.client.websocket import websocket_client

from .. import loader
from ....core import app
import langbot_plugin.api.entities.builtin.resource.tool as resource_tool
import langbot_plugin.api.entities.builtin.provider.message as provider_message
from ....entity.persistence import mcp as persistence_mcp


class MCPSessionStatus(enum.Enum):
    CONNECTING = 'connecting'
    CONNECTED = 'connected'
    ERROR = 'error'


class MCPSessionErrorPhase(enum.Enum):
    """Which phase of the MCP lifecycle failed."""
    SESSION_CREATE = 'session_create'
    DEP_INSTALL = 'dep_install'
    PROCESS_START = 'process_start'
    RELAY_CONNECT = 'relay_connect'
    MCP_INIT = 'mcp_init'
    RUNTIME = 'runtime'
    TOOL_CALL = 'tool_call'


_VENV_DIRS = frozenset({'.venv', 'venv', 'env', '.env'})
_VENV_BIN_DIRS = frozenset({'bin', 'Scripts'})


class MCPServerBoxConfig(pydantic.BaseModel):
    """Structured configuration for running an MCP server inside a Box container."""

    image: str | None = None
    network: str = 'on'  # MCP servers need network for dependency installation
    host_path: str | None = None
    host_path_mode: str = 'ro'  # MCP servers default to read-only mount
    env: dict[str, str] = pydantic.Field(default_factory=dict)
    startup_timeout_sec: int = 120  # Longer default to allow pip install
    cpus: float | None = None
    memory_mb: int | None = None
    pids_limit: int | None = None
    read_only_rootfs: bool | None = None

    model_config = pydantic.ConfigDict(extra='ignore')


class RuntimeMCPSession:
    """运行时 MCP 会话"""

    ap: app.Application

    server_name: str

    server_uuid: str

    server_config: dict

    session: ClientSession | None

    exit_stack: AsyncExitStack

    functions: list[resource_tool.LLMTool] = []

    enable: bool

    # connected: bool
    status: MCPSessionStatus

    _lifecycle_task: asyncio.Task | None

    _shutdown_event: asyncio.Event

    _ready_event: asyncio.Event

    error_message: str | None = None

    error_phase: MCPSessionErrorPhase | None = None

    retry_count: int = 0

    def __init__(self, server_name: str, server_config: dict, enable: bool, ap: app.Application):
        self.server_name = server_name
        self.server_uuid = server_config.get('uuid', '')
        self.server_config = server_config
        self.ap = ap
        self.enable = enable
        self.session = None

        self.exit_stack = AsyncExitStack()
        self.functions = []

        self.status = MCPSessionStatus.CONNECTING

        self._lifecycle_task = None
        self._shutdown_event = asyncio.Event()
        self._ready_event = asyncio.Event()

        # Parse box config once
        self.box_config = MCPServerBoxConfig.model_validate(
            server_config.get('box', {})
        )

    async def _init_stdio_python_server(self):
        if self._uses_box_stdio():
            await self._init_box_stdio_server()
            return

        server_params = StdioServerParameters(
            command=self.server_config['command'],
            args=self.server_config['args'],
            env=self.server_config['env'],
        )

        stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))

        stdio, write = stdio_transport

        self.session = await self.exit_stack.enter_async_context(ClientSession(stdio, write))

        await self.session.initialize()

    async def _init_box_stdio_server(self):
        box_service = self.ap.box_service
        session_id = self._build_box_session_id()
        host_path = self._resolve_host_path()
        session_payload = self._build_box_session_payload(session_id, host_path)

        # Phase: session creation
        try:
            await box_service.create_session(
                session_payload,
                skip_host_mount_validation=True,
            )
        except Exception as e:
            self.error_phase = MCPSessionErrorPhase.SESSION_CREATE
            raise

        # Phase: dependency installation
        if host_path:
            install_cmd = self._detect_install_command(host_path)
            if install_cmd:
                self.ap.logger.info(
                    f'MCP server {self.server_name}: installing dependencies in Box '
                    f'with: {install_cmd}'
                )
                exec_payload = dict(session_payload)
                exec_payload['cmd'] = install_cmd
                exec_payload['timeout_sec'] = self.box_config.startup_timeout_sec or 120
                try:
                    result = await box_service.client.execute(
                        box_service.build_spec(exec_payload, skip_host_mount_validation=True)
                    )
                except Exception as e:
                    self.error_phase = MCPSessionErrorPhase.DEP_INSTALL
                    raise
                if not result.ok:
                    self.error_phase = MCPSessionErrorPhase.DEP_INSTALL
                    stderr_preview = (result.stderr or '')[:500]
                    raise Exception(
                        f'Dependency install failed (exit code {result.exit_code}): '
                        f'{stderr_preview}'
                    )

        # Phase: managed process start
        try:
            await box_service.start_managed_process(
                session_id,
                self._build_box_process_payload(host_path),
            )
        except Exception as e:
            self.error_phase = MCPSessionErrorPhase.PROCESS_START
            raise

        # Phase: WebSocket relay connection
        try:
            websocket_url = box_service.get_managed_process_websocket_url(session_id)
            transport = await self.exit_stack.enter_async_context(websocket_client(websocket_url))
            read_stream, write_stream = transport
            self.session = await self.exit_stack.enter_async_context(ClientSession(read_stream, write_stream))
        except Exception as e:
            self.error_phase = MCPSessionErrorPhase.RELAY_CONNECT
            raise

        # Phase: MCP protocol initialization
        try:
            await self.session.initialize()
        except Exception as e:
            self.error_phase = MCPSessionErrorPhase.MCP_INIT
            raise

    async def _init_sse_server(self):
        sse_transport = await self.exit_stack.enter_async_context(
            sse_client(
                self.server_config['url'],
                headers=self.server_config.get('headers', {}),
                timeout=self.server_config.get('timeout', 10),
                sse_read_timeout=self.server_config.get('ssereadtimeout', 30),
            )
        )

        sseio, write = sse_transport

        self.session = await self.exit_stack.enter_async_context(ClientSession(sseio, write))

        await self.session.initialize()

    async def _init_streamable_http_server(self):
        transport = await self.exit_stack.enter_async_context(
            streamable_http_client(
                self.server_config['url'],
                http_client=httpx.AsyncClient(
                    headers=self.server_config.get('headers', {}),
                    timeout=self.server_config.get('timeout', 10),
                    follow_redirects=True,
                ),
            )
        )

        read, write, _ = transport

        self.session = await self.exit_stack.enter_async_context(ClientSession(read, write))

        await self.session.initialize()

    _MAX_RETRIES = 3
    _RETRY_DELAYS = [2, 4, 8]

    async def _lifecycle_loop(self):
        """Manage the full MCP session lifecycle in a background task."""
        try:
            if self.server_config['mode'] == 'stdio':
                await self._init_stdio_python_server()
            elif self.server_config['mode'] == 'sse':
                await self._init_sse_server()
            elif self.server_config['mode'] == 'http':
                await self._init_streamable_http_server()
            else:
                raise ValueError(f'Unknown MCP server mode: {self.server_name}: {self.server_config}')

            await self.refresh()

            self.status = MCPSessionStatus.CONNECTED

            # Notify start() that connection is established
            self._ready_event.set()

            # Wait for shutdown signal, with optional health monitoring for Box stdio
            if self._uses_box_stdio():
                monitor_task = asyncio.create_task(self._monitor_box_process_health())
                shutdown_task = asyncio.create_task(self._shutdown_event.wait())
                done, pending = await asyncio.wait(
                    [shutdown_task, monitor_task],
                    return_when=asyncio.FIRST_COMPLETED,
                )
                for task in pending:
                    task.cancel()
                for task in done:
                    if task is monitor_task and not self._shutdown_event.is_set():
                        self.error_phase = MCPSessionErrorPhase.RUNTIME
                        raise Exception('Box managed process exited unexpectedly')
            else:
                await self._shutdown_event.wait()

        except Exception as e:
            self.status = MCPSessionStatus.ERROR
            self.error_message = str(e)
            self.ap.logger.error(f'Error in MCP session lifecycle {self.server_name}: {e}\n{traceback.format_exc()}')
            # Do NOT set _ready_event here — let _lifecycle_loop_with_retry
            # handle retries first. It will set the event when all retries
            # are exhausted or on success.
            raise  # Re-raise so _lifecycle_loop_with_retry can catch it
        finally:
            # Clean up all resources in the same task
            try:
                if self.exit_stack:
                    await self.exit_stack.aclose()
                    self.exit_stack = AsyncExitStack()
                self.functions.clear()
                self.session = None
            except Exception as e:
                self.ap.logger.error(f'Error cleaning up MCP session {self.server_name}: {e}\n{traceback.format_exc()}')
            finally:
                await self._cleanup_box_stdio_session()

    async def _lifecycle_loop_with_retry(self):
        """Wrap _lifecycle_loop with retry and exponential backoff."""
        for attempt in range(self._MAX_RETRIES + 1):
            try:
                await self._lifecycle_loop()
                return  # Normal shutdown, don't retry
            except Exception as e:
                self.retry_count = attempt + 1
                if self._shutdown_event.is_set():
                    return  # Shutdown requested, don't retry
                if attempt >= self._MAX_RETRIES:
                    self.status = MCPSessionStatus.ERROR
                    self.error_message = f'Failed after {self._MAX_RETRIES + 1} attempts: {e}'
                    self._ready_event.set()
                    return
                delay = self._RETRY_DELAYS[attempt]
                self.ap.logger.warning(
                    f'MCP session {self.server_name} failed (attempt {attempt + 1}), '
                    f'retrying in {delay}s: {e}'
                )
                await self._cleanup_box_stdio_session()
                # Reset status for retry
                self.status = MCPSessionStatus.CONNECTING
                self.error_message = None
                self.error_phase = None
                await asyncio.sleep(delay)

    async def _monitor_box_process_health(self):
        """Poll managed process status; return when process exits."""
        from ...box.models import BoxManagedProcessStatus

        session_id = self._build_box_session_id()
        while not self._shutdown_event.is_set():
            try:
                info = await self.ap.box_service.client.get_managed_process(session_id)
                if isinstance(info, dict):
                    status = info.get('status', '')
                else:
                    status = getattr(info, 'status', '')
                if status == BoxManagedProcessStatus.EXITED.value or status == BoxManagedProcessStatus.EXITED:
                    return
            except Exception:
                return  # Process or session gone
            await asyncio.sleep(5)

    async def start(self):
        if not self.enable:
            return

        # Create background task for lifecycle management with retry
        self._lifecycle_task = asyncio.create_task(self._lifecycle_loop_with_retry())

        # Wait for connection or failure (with timeout)
        startup_timeout = self.box_config.startup_timeout_sec if self._uses_box_stdio() else 30.0
        try:
            await asyncio.wait_for(self._ready_event.wait(), timeout=startup_timeout)
        except asyncio.TimeoutError:
            self.status = MCPSessionStatus.ERROR
            raise Exception(f'Connection timeout after {startup_timeout} seconds')

        # Check for errors
        if self.status == MCPSessionStatus.ERROR:
            raise Exception('Connection failed, please check URL')

    async def refresh(self):
        if not self.session:
            return

        self.functions.clear()

        tools = await self.session.list_tools()

        self.ap.logger.debug(f'Refresh MCP tools: {tools}')

        for tool in tools.tools:

            async def func(*, _tool=tool, **kwargs):
                if not self.session:
                    raise Exception('MCP session is not connected')

                result = await self.session.call_tool(_tool.name, kwargs)
                if result.isError:
                    error_texts = []
                    for content in result.content:
                        if content.type == 'text':
                            error_texts.append(content.text)
                    raise Exception('\n'.join(error_texts) if error_texts else 'Unknown error from MCP tool')

                result_contents: list[provider_message.ContentElement] = []
                for content in result.content:
                    if content.type == 'text':
                        result_contents.append(provider_message.ContentElement.from_text(content.text))
                    elif content.type == 'image':
                        result_contents.append(provider_message.ContentElement.from_image_base64(content.image_base64))
                    elif content.type == 'resource':
                        # TODO: Handle resource content
                        pass

                return result_contents

            func.__name__ = tool.name

            self.functions.append(
                resource_tool.LLMTool(
                    name=tool.name,
                    human_desc=tool.description or '',
                    description=tool.description or '',
                    parameters=tool.inputSchema,
                    func=func,
                )
            )

    def get_tools(self) -> list[resource_tool.LLMTool]:
        return self.functions

    def get_runtime_info_dict(self) -> dict:
        info = {
            'status': self.status.value,
            'error_message': self.error_message,
            'error_phase': self.error_phase.value if self.error_phase else None,
            'retry_count': self.retry_count,
            'tool_count': len(self.get_tools()),
            'tools': [
                {
                    'name': tool.name,
                    'description': tool.description,
                }
                for tool in self.get_tools()
            ],
        }
        if self._uses_box_stdio():
            info['box_session_id'] = self._build_box_session_id()
            info['box_enabled'] = True
        return info

    async def shutdown(self):
        """关闭会话并清理资源"""
        try:
            # 设置shutdown事件，通知lifecycle任务退出
            self._shutdown_event.set()

            # 等待lifecycle任务完成（带超时）
            if self._lifecycle_task and not self._lifecycle_task.done():
                try:
                    await asyncio.wait_for(self._lifecycle_task, timeout=5.0)
                except asyncio.TimeoutError:
                    self.ap.logger.warning(f'MCP session {self.server_name} shutdown timeout, cancelling task')
                    self._lifecycle_task.cancel()
                    try:
                        await self._lifecycle_task
                    except asyncio.CancelledError:
                        pass

            self.ap.logger.info(f'MCP session {self.server_name} shutdown complete')
        except Exception as e:
            self.ap.logger.error(f'Error shutting down MCP session {self.server_name}: {e}\n{traceback.format_exc()}')

    def _uses_box_stdio(self) -> bool:
        """Check whether this stdio MCP server should run inside a Box container.

        Returns True when mode is stdio AND the Box runtime is available.
        An explicit ``box`` key in server_config is NOT required — if the
        runtime is reachable, stdio servers default to Box isolation.
        """
        if self.server_config.get('mode') != 'stdio':
            return False
        try:
            return getattr(self.ap.box_service, 'available', False)
        except Exception:
            return False

    def _build_box_session_id(self) -> str:
        return f'mcp-{self.server_uuid}'

    def _rewrite_path(self, path: str, host_path: str | None) -> str:
        """Rewrite host path prefix to container /workspace prefix."""
        if not host_path or not path:
            return path
        normalized_host = os.path.realpath(host_path)
        if path.startswith(normalized_host + '/'):
            return '/workspace' + path[len(normalized_host):]
        if path == normalized_host:
            return '/workspace'
        return path

    def _infer_host_path(self) -> str | None:
        """Try to infer host_path from command and args absolute paths.

        Detects virtualenv patterns (e.g. .venv/bin/python) and walks up
        to the project root rather than using the bin directory.
        """
        candidates = []
        parts = [self.server_config.get('command', '')] + self.server_config.get('args', [])
        for part in parts:
            if not os.path.isabs(part):
                continue
            # Use the raw path for venv detection (before resolving symlinks)
            # because .venv/bin/python is often a symlink to the system python.
            if os.path.exists(part):
                directory = os.path.dirname(part)
                directory = self._unwrap_venv_path(directory)
                candidates.append(os.path.realpath(directory))
        if not candidates:
            return None
        common = os.path.commonpath(candidates)
        return common if common != '/' else None

    @staticmethod
    def _unwrap_venv_path(directory: str) -> str:
        """If directory looks like a virtualenv bin dir, return the project root.

        Recognized patterns:
          /project/.venv/bin        -> /project
          /project/venv/bin         -> /project
          /project/.venv/Scripts    -> /project  (Windows)
          /project/env/bin          -> /project
        """
        parts = directory.replace('\\', '/').split('/')
        # Look for patterns like .../(.venv|venv|env)/(bin|Scripts)
        for i in range(len(parts) - 1, 0, -1):
            if parts[i] in _VENV_BIN_DIRS and i >= 1:
                venv_dir = parts[i - 1]
                if venv_dir in _VENV_DIRS:
                    # Return everything before the venv directory
                    project_root = '/'.join(parts[:i - 1])
                    return project_root if project_root else '/'
        return directory

    def _resolve_host_path(self) -> str | None:
        """Resolve the effective host_path: explicit config > inference."""
        return self.box_config.host_path or self._infer_host_path()

    @staticmethod
    def _detect_install_command(host_path: str) -> str | None:
        """Detect how to install dependencies from the mounted project.

        Copies the project to a writable temp directory before installing,
        because /workspace may be mounted read-only and pip needs to write
        build artifacts in the source tree.
        """
        _COPY_AND_INSTALL = (
            'cp -r /workspace /tmp/_mcp_src'
            ' && pip install --no-cache-dir /tmp/_mcp_src'
            ' && rm -rf /tmp/_mcp_src'
        )
        _INSTALL_REQUIREMENTS = 'pip install --no-cache-dir -r /workspace/requirements.txt'

        if os.path.isfile(os.path.join(host_path, 'pyproject.toml')):
            return _COPY_AND_INSTALL
        if os.path.isfile(os.path.join(host_path, 'setup.py')):
            return _COPY_AND_INSTALL
        if os.path.isfile(os.path.join(host_path, 'requirements.txt')):
            return _INSTALL_REQUIREMENTS
        return None

    def _build_box_session_payload(self, session_id: str, host_path: str | None = None) -> dict:
        bc = self.box_config
        if host_path is None:
            host_path = self._resolve_host_path()

        payload: dict[str, typing.Any] = {
            'session_id': session_id,
            'workdir': '/workspace',
            'env': bc.env,
            # MCP sessions need network for dependency install and writable rootfs
            'network': bc.network,
            'read_only_rootfs': bc.read_only_rootfs if bc.read_only_rootfs is not None else False,
        }
        if host_path:
            payload['host_path'] = host_path
            payload['host_path_mode'] = bc.host_path_mode
        for key in ('image', 'cpus', 'memory_mb', 'pids_limit'):
            val = getattr(bc, key)
            if val is not None:
                payload[key] = val if not isinstance(val, enum.Enum) else val.value
        return payload

    def _build_box_process_payload(self, host_path: str | None = None) -> dict:
        if host_path is None:
            host_path = self._resolve_host_path()

        command = self.server_config['command']
        args = self.server_config.get('args', [])
        cwd = '/workspace'

        if host_path:
            # When host_path is resolved, we install deps in-container rather
            # than relying on the host venv.  Rewrite paths so the container
            # sees /workspace/... but replace venv python with plain "python".
            command = self._rewrite_venv_command(command, host_path)
            args = [self._rewrite_path(a, host_path) for a in args]
            cwd = self._rewrite_path(cwd, host_path)

        return {
            'command': command,
            'args': args,
            'env': self.server_config.get('env', {}),
            'cwd': cwd,
        }

    def _rewrite_venv_command(self, command: str, host_path: str) -> str:
        """Rewrite command: if it points to a venv python, use plain 'python'."""
        if not host_path or not command:
            return command
        normalized_host = os.path.realpath(host_path)
        if not command.startswith(normalized_host + '/'):
            return command
        # Check if command is a venv python interpreter
        rel = command[len(normalized_host) + 1:]  # e.g. ".venv/bin/python"
        parts = rel.replace('\\', '/').split('/')
        # Match patterns like .venv/bin/python*, venv/bin/python*, etc.
        if (len(parts) >= 3
                and parts[0] in _VENV_DIRS
                and parts[1] in _VENV_BIN_DIRS
                and parts[2].startswith('python')):
            return 'python'
        # Not a venv python — do normal path rewrite
        return self._rewrite_path(command, host_path)

    async def _cleanup_box_stdio_session(self) -> None:
        if not self._uses_box_stdio():
            return

        try:
            await self.ap.box_service.client.delete_session(self._build_box_session_id())
        except Exception as e:
            self.ap.logger.warning(f'Failed to cleanup Box session for MCP server {self.server_name}: {e}')


# @loader.loader_class('mcp')
class MCPLoader(loader.ToolLoader):
    """MCP 工具加载器。

    在此加载器中管理所有与 MCP Server 的连接。
    """

    sessions: dict[str, RuntimeMCPSession]

    _last_listed_functions: list[resource_tool.LLMTool]

    _hosted_mcp_tasks: list[asyncio.Task]

    def __init__(self, ap: app.Application):
        super().__init__(ap)
        self.sessions = {}
        self._last_listed_functions = []
        self._hosted_mcp_tasks = []

    async def initialize(self):
        await self.load_mcp_servers_from_db()

    async def load_mcp_servers_from_db(self):
        self.ap.logger.info('Loading MCP servers from db...')

        self.sessions = {}

        result = await self.ap.persistence_mgr.execute_async(sqlalchemy.select(persistence_mcp.MCPServer))
        servers = result.all()

        for server in servers:
            config = self.ap.persistence_mgr.serialize_model(persistence_mcp.MCPServer, server)

            task = asyncio.create_task(self.host_mcp_server(config))
            self._hosted_mcp_tasks.append(task)

    async def host_mcp_server(self, server_config: dict):
        self.ap.logger.debug(f'Loading MCP server {server_config}')
        try:
            session = await self.load_mcp_server(server_config)
            self.sessions[server_config['name']] = session
        except Exception as e:
            self.ap.logger.error(
                f'Failed to load MCP server from db: {server_config["name"]}({server_config["uuid"]}): {e}\n{traceback.format_exc()}'
            )
            return

        self.ap.logger.debug(f'Starting MCP server {server_config["name"]}({server_config["uuid"]})')
        try:
            await session.start()
        except Exception as e:
            self.ap.logger.error(
                f'Failed to start MCP server {server_config["name"]}({server_config["uuid"]}): {e}\n{traceback.format_exc()}'
            )
            return

        self.ap.logger.debug(f'Started MCP server {server_config["name"]}({server_config["uuid"]})')

    async def load_mcp_server(self, server_config: dict) -> RuntimeMCPSession:
        """加载 MCP 服务器到运行时

        Args:
            server_config: 服务器配置字典，必须包含:
                - name: 服务器名称
                - mode: 连接模式 (stdio/sse/http)
                - enable: 是否启用
                - extra_args: 额外的配置参数 (可选)
        """
        uuid_ = server_config.get('uuid')
        if not uuid_:
            self.ap.logger.warning('Server UUID is None for MCP server, maybe testing in the config page.')
            uuid_ = str(uuid_module.uuid4())
            server_config['uuid'] = uuid_

        name = server_config['name']
        uuid = server_config['uuid']
        mode = server_config['mode']
        enable = server_config['enable']
        extra_args = server_config.get('extra_args', {})

        mixed_config = {
            'name': name,
            'uuid': uuid,
            'mode': mode,
            'enable': enable,
            **extra_args,
        }

        session = RuntimeMCPSession(name, mixed_config, enable, self.ap)

        return session

    async def get_tools(self, bound_mcp_servers: list[str] | None = None) -> list[resource_tool.LLMTool]:
        all_functions = []

        for session in self.sessions.values():
            # If bound_mcp_servers is specified, only include tools from those servers
            if bound_mcp_servers is not None:
                if session.server_uuid in bound_mcp_servers:
                    all_functions.extend(session.get_tools())
            else:
                # If no bound servers specified, include all tools
                all_functions.extend(session.get_tools())

        self._last_listed_functions = all_functions

        return all_functions

    async def has_tool(self, name: str) -> bool:
        """检查工具是否存在"""
        for session in self.sessions.values():
            for function in session.get_tools():
                if function.name == name:
                    return True
        return False

    async def invoke_tool(self, name: str, parameters: dict, query: pipeline_query.Query) -> typing.Any:
        """执行工具调用"""
        for session in self.sessions.values():
            for function in session.get_tools():
                if function.name == name:
                    self.ap.logger.debug(f'Invoking MCP tool: {name} with parameters: {parameters}')
                    try:
                        result = await function.func(**parameters)
                        self.ap.logger.debug(f'MCP tool {name} executed successfully')
                        return result
                    except Exception as e:
                        self.ap.logger.error(f'Error invoking MCP tool {name}: {e}\n{traceback.format_exc()}')
                        raise

        raise ValueError(f'Tool not found: {name}')

    async def remove_mcp_server(self, server_name: str):
        """移除 MCP 服务器"""
        if server_name not in self.sessions:
            self.ap.logger.warning(f'MCP server {server_name} not found in sessions, skipping removal')
            return

        session = self.sessions.pop(server_name)
        await session.shutdown()
        self.ap.logger.info(f'Removed MCP server: {server_name}')

    def get_session(self, server_name: str) -> RuntimeMCPSession | None:
        """获取指定名称的 MCP 会话"""
        return self.sessions.get(server_name)

    def has_session(self, server_name: str) -> bool:
        """检查是否存在指定名称的 MCP 会话"""
        return server_name in self.sessions

    def get_all_server_names(self) -> list[str]:
        """获取所有已加载的 MCP 服务器名称"""
        return list(self.sessions.keys())

    def get_server_tool_count(self, server_name: str) -> int:
        """获取指定服务器的工具数量"""
        session = self.get_session(server_name)
        return len(session.get_tools()) if session else 0

    def get_all_servers_info(self) -> dict[str, dict]:
        """获取所有服务器的信息"""
        info = {}
        for server_name, session in self.sessions.items():
            info[server_name] = {
                'name': server_name,
                'mode': session.server_config.get('mode'),
                'enable': session.enable,
                'tools_count': len(session.get_tools()),
                'tool_names': [f.name for f in session.get_tools()],
            }
        return info

    async def shutdown(self):
        """关闭所有工具"""
        self.ap.logger.info('Shutting down all MCP sessions...')
        for server_name, session in list(self.sessions.items()):
            try:
                await session.shutdown()
                self.ap.logger.debug(f'Shutdown MCP session: {server_name}')
            except Exception as e:
                self.ap.logger.error(f'Error shutting down MCP session {server_name}: {e}\n{traceback.format_exc()}')
        self.sessions.clear()
        self.ap.logger.info('All MCP sessions shutdown complete')