From 8b4f9ab20ef8e693e8e4408b17fb40f26e3adaff Mon Sep 17 00:00:00 2001 From: Chingfeng Li Date: Fri, 21 Nov 2025 14:27:18 +0800 Subject: [PATCH] Add ProgressFnT LoggingFnT... --- .../providers/tools/server_mcp/mcp_client.py | 46 +++++++++++++++---- .../providers/tools/server_mcp/mcp_manager.py | 21 +++++++-- 2 files changed, 56 insertions(+), 11 deletions(-) diff --git a/main/xiaozhi-server/core/providers/tools/server_mcp/mcp_client.py b/main/xiaozhi-server/core/providers/tools/server_mcp/mcp_client.py index c1126a64..c8c42d79 100644 --- a/main/xiaozhi-server/core/providers/tools/server_mcp/mcp_client.py +++ b/main/xiaozhi-server/core/providers/tools/server_mcp/mcp_client.py @@ -10,10 +10,13 @@ import concurrent.futures from contextlib import AsyncExitStack from typing import Optional, List, Dict, Any -from mcp import ClientSession, StdioServerParameters +from mcp import ClientSession, StdioServerParameters, Implementation +from mcp.client.session import SamplingFnT, ElicitationFnT, ListRootsFnT, LoggingFnT, MessageHandlerFnT from mcp.client.stdio import stdio_client from mcp.client.sse import sse_client from mcp.client.streamable_http import streamablehttp_client +from mcp.shared.session import ProgressFnT + from config.logger import setup_logging from core.utils.util import sanitize_tool_name @@ -41,13 +44,25 @@ class ServerMCPClient: self.tools_dict: Dict[str, Any] = {} self.name_mapping: Dict[str, str] = {} - async def initialize(self): + async def initialize(self, read_timeout_seconds: timedelta | None = None, + sampling_callback: SamplingFnT | None = None, + elicitation_callback: ElicitationFnT | None = None, + list_roots_callback: ListRootsFnT | None = None, + logging_callback: LoggingFnT | None = None, + message_handler: MessageHandlerFnT | None = None, + client_info: Implementation | None = None): """初始化MCP客户端连接""" if self._worker_task: return self._worker_task = asyncio.create_task( - self._worker(), name="ServerMCPClientWorker" + self._worker(read_timeout_seconds=read_timeout_seconds, + sampling_callback=sampling_callback, + elicitation_callback=elicitation_callback, + list_roots_callback=list_roots_callback, + logging_callback=logging_callback, + message_handler=message_handler, + client_info=client_info), name="ServerMCPClientWorker" ) await self._ready_evt.wait() @@ -97,12 +112,15 @@ class ServerMCPClient: for name, tool in self.tools_dict.items() ] - async def call_tool(self, name: str, args: dict) -> Any: + async def call_tool(self, name: str, arguments: dict, read_timeout_seconds: timedelta | None = None, progress_callback: ProgressFnT | None = None, *, meta: dict[str, Any] | None = None) -> Any: """调用指定工具 Args: name: 工具名称 - args: 工具参数 + arguments: 工具参数 + read_timeout_seconds: + progress_callback: 进度回调函数 + meta: Returns: Any: 工具执行结果 @@ -115,7 +133,7 @@ class ServerMCPClient: real_name = self.name_mapping.get(name, name) loop = self._worker_task.get_loop() - coro = self.session.call_tool(real_name, args) + coro = self.session.call_tool(real_name, arguments=arguments, read_timeout_seconds=read_timeout_seconds, progress_callback=progress_callback, meta=meta) if loop is asyncio.get_running_loop(): return await coro @@ -144,7 +162,13 @@ class ServerMCPClient: # 所有检查都通过，连接正常 return True - async def _worker(self): + async def _worker(self, read_timeout_seconds: timedelta | None = None, + sampling_callback: SamplingFnT | None = None, + elicitation_callback: ElicitationFnT | None = None, + list_roots_callback: ListRootsFnT | None = None, + logging_callback: LoggingFnT | None = None, + message_handler: MessageHandlerFnT | None = None, + client_info: Implementation | None = None): """MCP客户端工作协程""" async with AsyncExitStack() as stack: try: @@ -208,7 +232,13 @@ class ServerMCPClient: ClientSession( read_stream=read_stream, write_stream=write_stream, - read_timeout_seconds=timedelta(seconds=15), + read_timeout_seconds=read_timeout_seconds, + sampling_callback=sampling_callback, + elicitation_callback=elicitation_callback, + list_roots_callback=list_roots_callback, + logging_callback=logging_callback, + message_handler=message_handler, + client_info=client_info ) ) await self.session.initialize() diff --git a/main/xiaozhi-server/core/providers/tools/server_mcp/mcp_manager.py b/main/xiaozhi-server/core/providers/tools/server_mcp/mcp_manager.py index cf3deb21..6edc44d1 100644 --- a/main/xiaozhi-server/core/providers/tools/server_mcp/mcp_manager.py +++ b/main/xiaozhi-server/core/providers/tools/server_mcp/mcp_manager.py @@ -3,7 +3,14 @@ import asyncio import os import json +from datetime import timedelta from typing import Dict, Any, List + +from mcp import Implementation +from mcp.client.session import SamplingFnT, ElicitationFnT, ListRootsFnT, LoggingFnT, MessageHandlerFnT +from mcp.shared.session import ProgressFnT +from mcp.types import LoggingMessageNotificationParams + from config.config_loader import get_project_dir from config.logger import setup_logging from .mcp_client import ServerMCPClient @@ -56,7 +63,7 @@ class ServerMCPManager: # 初始化服务端MCP客户端 logger.bind(tag=TAG).info(f"初始化服务端MCP客户端: {name}") client = ServerMCPClient(srv_config) - await client.initialize() + await client.initialize(logging_callback=self.logging_callback) self.clients[name] = client client_tools = client.get_available_tools() self.tools.extend(client_tools) @@ -109,7 +116,7 @@ class ServerMCPManager: # 带重试机制的工具调用 for attempt in range(max_retries): try: - return await target_client.call_tool(tool_name, arguments) + return await target_client.call_tool(tool_name, arguments, progress_callback=self.progress_callback) except Exception as e: # 最后一次尝试失败时直接抛出异常 if attempt == max_retries - 1: @@ -131,7 +138,7 @@ class ServerMCPManager: config = self.load_config() if client_name in config: client = ServerMCPClient(config[client_name]) - await client.initialize() + await client.initialize(logging_callback=self.logging_callback) self.clients[client_name] = client target_client = client logger.bind(tag=TAG).info( @@ -159,3 +166,11 @@ class ServerMCPManager: except (asyncio.TimeoutError, Exception) as e: logger.bind(tag=TAG).error(f"关闭服务端MCP客户端 {name} 时出错: {e}") self.clients.clear() + + # 可选回调方法 + + async def logging_callback(self, params: LoggingMessageNotificationParams): + logger.bind(tag=TAG).info(f"[Server Log - {params.level.upper()}] {params.data}") + + async def progress_callback(self, progress: float, total: float | None, message: str | None) -> None: + logger.bind(tag=TAG).info(f"[Progress {progress}/{total}]: {message}") \ No newline at end of file