from __future__ import annotations import time import traceback import uuid import quart import asyncio import sqlalchemy from pkg.entity.persistence.mcp import MCPServer from .....core import taskmgr from .. import group from sqlalchemy import insert @group.group_class('mcp', '/api/v1/mcp') class MCPRouterGroup(group.RouterGroup): async def initialize(self) -> None: @self.route('/servers', methods=['GET', 'POST'], auth_type=group.AuthType.USER_TOKEN) async def _() -> str: """获取MCP服务器列表""" if quart.request.method == 'GET': result = await self.ap.persistence_mgr.execute_async( sqlalchemy.select(MCPServer).order_by(MCPServer.created_at.desc()) ) raw_results = result.all() servers = [self.ap.persistence_mgr.serialize_model(MCPServer, row) for row in raw_results] servers_with_status = [] # 获取MCP工具加载器 mcp_loader = None for loader in self.ap.tool_mgr.loaders: if loader.__class__.__name__ == 'MCPLoader': mcp_loader = loader break for server in servers: # 设置状态 if server['enable']: status = 'enabled' else: status = 'disabled' # 构建 config 对象 (前端期望的格式) extra_args = server.get('extra_args', {}) config = { 'name': server['name'], 'mode': server['mode'], 'enable': server['enable'], } # 根据模式添加相应的配置 if server['mode'] == 'sse': config['url'] = extra_args.get('url', '') config['headers'] = extra_args.get('headers', {}) config['timeout'] = extra_args.get('timeout', 60) elif server['mode'] == 'stdio': config['command'] = extra_args.get('command', '') config['args'] = extra_args.get('args', []) config['env'] = extra_args.get('env', {}) # 从运行中的会话获取工具数量 tools_count = 0 if mcp_loader and hasattr(mcp_loader, 'sessions') and server['name'] in mcp_loader.sessions: session = mcp_loader.sessions[server['name']] tools_count = len(session.functions) server_info = { 'name': server['name'], 'mode': server['mode'], 'enable': server['enable'], 'status': status, 'tools': tools_count, # 从运行中的会话获取工具数量 'config': config, } servers_with_status.append(server_info) return self.success(data={'servers': servers_with_status}) elif quart.request.method == 'POST': data = await quart.request.json data = data['source'] try: # 检查服务器名称是否重复 result = await self.ap.persistence_mgr.execute_async( sqlalchemy.select(MCPServer).where(MCPServer.name == data['name']) ) if result.first() is not None: return self.http_status(400, -1, 'Server name already exists') # 创建新服务器配置 new_server = { 'uuid': str(uuid.uuid4()), 'name': data['name'], 'mode': 'sse', 'enable': data.get('enable', False), 'extra_args': { 'url':data.get('url',''), 'headers':data.get('headers',{}), 'timeout':data.get('timeout',60), 'ssereadtimeout':data.get('ssereadtimeout',300), }, } await self.ap.persistence_mgr.execute_async( sqlalchemy.insert(MCPServer).values(new_server) ) return self.success() except Exception: print(traceback.format_exc()) @self.route('/servers/', methods=['GET', 'PUT', 'DELETE'], auth_type=group.AuthType.USER_TOKEN) async def _(server_name: str) -> str: """获取、更新或删除MCP服务器配置""" result = await self.ap.persistence_mgr.execute_async( sqlalchemy.select(MCPServer).where(MCPServer.name == server_name) ) server = result.first() if server is None: return self.http_status(404, -1, 'Server not found') if quart.request.method == 'GET': server_data = self.ap.persistence_mgr.serialize_model(MCPServer, server) return self.success(data={'server': server_data}) elif quart.request.method == 'PUT': data = await quart.request.json update_data = { 'enable': data.get('enable', server.enable), } extra_args = server.extra_args or {} if server.mode == 'sse': extra_args.update({ 'url': data.get('url', extra_args.get('url','')), 'headers': data.get('headers', extra_args.get('headers',{})), 'timeout': data.get('timeout', extra_args.get('timeout',60)), 'ssereadtimeout': data.get('ssereadtimeout', extra_args.get('ssereadtimeout',300)), }) update_data['extra_args'] = extra_args await self.ap.persistence_mgr.execute_async( sqlalchemy.update(MCPServer).where(MCPServer.name == server_name).values(update_data) ) return self.success() elif quart.request.method == 'DELETE': await self.ap.persistence_mgr.execute_async( sqlalchemy.delete(MCPServer).where(MCPServer.name == server_name) ) return self.success() @self.route('/servers//test', methods=['POST'], auth_type=group.AuthType.USER_TOKEN) async def _(server_name: str) -> str: """测试MCP服务器连接""" result = await self.ap.persistence_mgr.execute_async( sqlalchemy.select(MCPServer).where(MCPServer.name == server_name) ) server = result.first() if server is None: return self.http_status(404, -1, 'Server not found') # 创建测试任务 ctx = taskmgr.TaskContext.new() wrapper = self.ap.task_mgr.create_user_task( self._test_mcp_server(server, ctx), kind='mcp-operation', name=f'mcp-test-{server_name}', label=f'Testing MCP server {server_name}', context=ctx, ) return self.success(data={'task_id': wrapper.id}) async def _test_mcp_server(self, server: MCPServer, ctx: taskmgr.TaskContext): """测试MCP服务器连接""" try: from .....provider.tools.loaders.mcp import RuntimeMCPSession ctx.current_action = f'Testing connection to {server.name}' print(server) # 创建临时会话进行测试 session = RuntimeMCPSession(server.name, { 'name': server.name, 'mode': server.mode, 'enable': server.enable, 'url': server.extra_args.get('url',''), 'headers': server.extra_args.get('headers',{}), 'timeout': server.extra_args.get('timeout',60), },enable=True, ap=self.ap) await session.start() # 获取工具列表作为测试 tools_count = len(session.functions) tool_name_list = [] for function in session.functions: tool_name_list.append(function.name) ctx.current_action = f'Successfully connected. Found {tools_count} tools.' # 关闭测试会话 await session.shutdown() return {'status': 'success', 'tools_count': tools_count,'tools_names_lists':tool_name_list} except Exception as e: print(traceback.format_exc()) ctx.current_action = f'Connection test failed: {str(e)}' raise e