Coverage for aipyapp/cli/command/cmd_mcp.py: 0%

1from rich import print 

2import argparse 

3 

4from ... import T 

5from .base_parser import ParserCommand 

6from .utils import print_records, print_table 

7 

8class MCPCommand(ParserCommand): 

9 name = 'mcp' 

10 description = T('MCP operations') 

11 

12 def add_subcommands(self, subparsers): 

13 subparsers.add_parser('status', help=T('Show MCP status')) 

14 

15 common_parser = argparse.ArgumentParser(add_help=False) 

16 common_parser.add_argument('--user', action='store_true', help=T('User MCP')) 

17 common_parser.add_argument('--sys', action='store_true', help=T('System MCP')) 

18 subparsers.add_parser('enable', help=T('Enable MCP'), parents=[common_parser]) 

19 subparsers.add_parser('disable', help=T('Disable MCP'), parents=[common_parser]) 

20 

21 # 服务器管理子命令 

22 server_parser = subparsers.add_parser('server', help=T('MCP server operations')) 

23 server_group = server_parser.add_mutually_exclusive_group() 

24 server_group.add_argument('--list', action='store_true', help=T('List MCP servers')) 

25 server_group.add_argument('--enable', nargs='?', const='*',help=T('Enable MCP server')) 

26 server_group.add_argument('--disable', nargs='?', const='*', help=T('Disable MCP server')) 

27 

28 def _enable_mcp(self, args, ctx, enable): 

29 mcp = ctx.tm.mcp 

30 if args.user: 

31 ret = mcp.enable_user_mcp(enable) 

32 elif args.sys: 

33 ret = mcp.enable_sys_mcp(enable) 

34 else: 

35 ctx.console.print(T('Please specify --user or --sys'), style="bold red") 

36 ret = None 

37 ctx.console.print(T('Success') if ret else T('Failed'), style="bold green" if ret else "bold red") 

38 return ret 

39 

40 def cmd_status(self, args, ctx): 

41 mcp = ctx.tm.mcp 

42 ret = mcp.get_status() 

43 rows = [ 

44 (T('System MCP'), T('Enabled') if ret['sys_mcp_enabled'] else T('Disabled')), 

45 (T('User MCP'), T('Enabled') if ret['user_mcp_enabled'] else T('Disabled')), 

46 (T('Total Servers'), ret['total_servers']), 

47 (T('Enabled Servers'), ret['enabled_servers']), 

48 (T('Total Tools'), ret['total_tools']), 

49 (T('Enabled Tools'), ret['enabled_tools']) 

50 ] 

51 print_table(rows, title=T('MCP status'), headers=["Name", "Value"]) 

52 return ret 

53 

54 def cmd_enable(self, args, ctx): 

55 ret = self._enable_mcp(args, ctx, True) 

56 return ret 

57 

58 def cmd_disable(self, args, ctx): 

59 ret = self._enable_mcp(args, ctx, False) 

60 return ret 

61 

62 def cmd_server(self, args, ctx): 

63 mcp = ctx.tm.mcp 

64 if args.enable: 

65 ret = mcp.enable_user_server(args.enable) 

66 elif args.disable: 

67 ret = mcp.enable_user_server(args.disable, False) 

68 else: 

69 servers = mcp.list_user_servers() 

70 print_records(servers, title=T('MCP servers')) 

71 ret = True 

72 ctx.console.print(T('Success') if ret else T('Failed'), style="bold green" if ret else "bold red") 

73 return ret 

74 

75 def cmd(self, args, ctx): 

76 return self.cmd_status(args, ctx)