metona-ai-desktop/electron/services/mcp-manager.service.ts

/**
 * MCP Manager Service — MCP Server 生命周期管理
 *
 * 负责：
 * 1. MCP Server 的连接/断开/重连
 * 2. 工具发现与动态注册到 ToolRegistry
 * 3. Server 状态管理与健康检查
 *
 * 使用 @modelcontextprotocol/sdk 官方库。
 *
 * @see docs/生产级通用 AI Agent 智能体桌面应用：完整设计与构建指南.html — 第七章
 * @see standard/开发规范.md — 优先使用第三方成熟库
 */

import { Client } from '@modelcontextprotocol/sdk/client/index.js';
import { StdioClientTransport } from '@modelcontextprotocol/sdk/client/stdio.js';
import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js';
import type { Tool } from '@modelcontextprotocol/sdk/types.js';
import { nanoid } from 'nanoid';
import type Database from 'better-sqlite3';
import log from 'electron-log';
import type { ToolRegistry } from '../harness/tools/registry';
import type { IMetonaTool, ToolExecutionContext } from '../harness/types/metona-tool';
import type { MetonaToolDef } from '../harness/types';
import { MetonaToolCategory, MetonaRiskLevel } from '../harness/types';

// v0.3.0 修复: 安全解析 JSON args，防止数据库中存储了非法 JSON 导致初始化崩溃
function safeParseArgs(raw: string): string[] {
  try {
    const parsed = JSON.parse(raw);
    return Array.isArray(parsed) ? parsed : [];
  } catch {
    log.warn(`MCP args parse failed, using empty array: ${raw.slice(0, 100)}`);
    return [];
  }
}

// ===== 类型定义 =====

export type MCPServerStatus = 'connecting' | 'connected' | 'disconnected' | 'error';

export interface MCPServerConfig {
  id: string;
  name: string;
  transport: 'stdio' | 'sse';
  command?: string;
  args?: string[];
  url?: string;
  enabled: boolean;
}

export interface MCPServerState {
  config: MCPServerConfig;
  status: MCPServerStatus;
  client: Client | null;
  tools: Tool[];
  error?: string;
  connectedAt?: number;
}

// ===== MCP Tool Adapter =====

/**
 * 将 MCP Tool 适配为 IMetonaTool 接口
 */
class MCPToolAdapter implements IMetonaTool {
  readonly definition: MetonaToolDef;

  constructor(
    private mcpTool: Tool,
    private client: Client,
    private serverName: string,
  ) {
    this.definition = {
      name: `mcp_${serverName}_${mcpTool.name}`,
      description: mcpTool.description ?? `MCP tool from ${serverName}`,
      parameters: this.convertSchema(mcpTool.inputSchema),
      category: MetonaToolCategory.MCP,
      riskLevel: MetonaRiskLevel.MEDIUM,
      requiresPermission: false,
      timeoutMs: 30_000,
    };
  }

  async execute(args: Record<string, unknown>, _context: ToolExecutionContext): Promise<unknown> {
    const result = await this.client.callTool({
      name: this.mcpTool.name,
      arguments: args,
    });
    return result.content;
  }

  /**
   * 将 MCP JSON Schema 转换为 MetonaToolParams
   */
  private convertSchema(schema: Record<string, unknown>): MetonaToolDef['parameters'] {
    const properties: Record<string, { type: 'string' | 'number' | 'boolean' | 'object' | 'array'; description: string }> = {};
    const schemaProps = (schema.properties ?? {}) as Record<string, Record<string, unknown>>;

    for (const [key, prop] of Object.entries(schemaProps)) {
      properties[key] = {
        type: (prop.type as 'string' | 'number' | 'boolean' | 'object' | 'array') ?? 'string',
        description: (prop.description as string) ?? '',
      };
    }

    return {
      type: 'object',
      properties,
      required: schema.required as string[] | undefined,
    };
  }
}

// ===== MCP Manager =====

export class MCPManager {
  private servers = new Map<string, MCPServerState>();

  constructor(
    private getDB: () => Database.Database,
    private toolRegistry: ToolRegistry,
  ) {}

  /**
   * 初始化：从数据库加载已启用的 MCP Server 并连接
   */
  async initialize(): Promise<void> {
    const db = this.getDB();
    const rows = db.prepare(`
      SELECT * FROM mcp_servers WHERE enabled = 1
    `).all() as Array<{
      id: string; name: string; transport: string;
      command: string | null; args: string | null; url: string | null;
    }>;

    for (const row of rows) {
      const config: MCPServerConfig = {
        id: row.id,
        name: row.name,
        transport: row.transport as 'stdio' | 'sse',
        command: row.command ?? undefined,
        args: row.args ? safeParseArgs(row.args) : undefined,
        url: row.url ?? undefined,
        enabled: true,
      };

      // 异步连接，不阻塞启动
      this.connectServer(config).catch((err) => {
        log.warn(`MCP server "${config.name}" auto-connect failed: ${err}`);
      });
    }

    log.info(`MCP Manager initialized: ${rows.length} server(s) configured`);
  }

  /**
   * 连接 MCP Server
   */
  async connectServer(config: MCPServerConfig): Promise<void> {
    const { name } = config;

    // 断开已有连接
    if (this.servers.has(name)) {
      await this.disconnectServer(name);
    }

    this.servers.set(name, {
      config,
      status: 'connecting',
      client: null,
      tools: [],
    });

    try {
      let transport;

      if (config.transport === 'stdio' && config.command) {
        // stdio 模式
        const args = config.args ?? [];
        transport = new StdioClientTransport({
          command: config.command,
          args,
        });
      } else if (config.transport === 'sse' && config.url) {
        // SSE 模式（远程 HTTP）
        transport = new SSEClientTransport(new URL(config.url));
      } else {
        throw new Error(
          `Unsupported transport "${config.transport}". ` +
          `'stdio' requires 'command', 'sse' requires 'url'.`,
        );
      }

      const client = new Client(
        { name: 'metona-ai-desktop', version: '1.0.0' },
        { capabilities: {} },
      );

      await client.connect(transport);

      // 发现工具
      const toolsResult = await client.listTools();
      const tools = toolsResult.tools ?? [];

      // 注册到 ToolRegistry
      for (const tool of tools) {
        const adapter = new MCPToolAdapter(tool, client, name);
        this.toolRegistry.registerMCP(name, adapter);
      }

      // 更新状态
      const state = this.servers.get(name)!;
      state.status = 'connected';
      state.client = client;
      state.tools = tools;
      state.connectedAt = Date.now();
      state.error = undefined;

      // 更新数据库
      const db = this.getDB();
      db.prepare(`
        UPDATE mcp_servers SET last_connected = ?, error_message = NULL WHERE name = ?
      `).run(Date.now(), name);

      log.info(`MCP server "${name}" connected: ${tools.length} tool(s)`);
    } catch (error) {
      const state = this.servers.get(name);
      if (state) {
        state.status = 'error';
        state.error = (error as Error).message;
      }

      // 更新数据库
      const db = this.getDB();
      db.prepare(`
        UPDATE mcp_servers SET error_message = ? WHERE name = ?
      `).run((error as Error).message, name);

      log.error(`MCP server "${name}" connection failed:`, error);
      throw error;
    }
  }

  /**
   * 断开 MCP Server
   */
  async disconnectServer(name: string): Promise<void> {
    const state = this.servers.get(name);
    if (!state) return;

    // 从 ToolRegistry 注销
    this.toolRegistry.unregisterMCPTools(name);

    // 关闭客户端
    if (state.client) {
      try {
        await state.client.close();
      } catch {
        // 忽略关闭错误
      }
    }

    state.status = 'disconnected';
    state.client = null;
    state.tools = [];

    log.info(`MCP server "${name}" disconnected`);
  }

  /**
   * 切换 Server 启用/禁用
   */
  async toggleServer(name: string, enabled: boolean): Promise<void> {
    const db = this.getDB();
    db.prepare(`
      UPDATE mcp_servers SET enabled = ?, updated_at = ? WHERE name = ?
    `).run(enabled ? 1 : 0, Date.now(), name);

    if (enabled) {
      const row = db.prepare('SELECT * FROM mcp_servers WHERE name = ?').get(name) as {
        id: string; name: string; transport: string;
        command: string | null; args: string | null; url: string | null;
      } | undefined;
      if (row) {
        await this.connectServer({
          id: row.id, name: row.name,
          transport: row.transport as 'stdio' | 'sse',
          command: row.command ?? undefined,
          args: row.args ? safeParseArgs(row.args) : undefined,
          url: row.url ?? undefined,
          enabled: true,
        });
      }
    } else {
      await this.disconnectServer(name);
    }
  }

  /**
   * 添加新的 MCP Server
   */
  async addServer(config: Omit<MCPServerConfig, 'id'>): Promise<void> {
    const db = this.getDB();
    const id = `mcp_${nanoid(8)}`;

    db.prepare(`
      INSERT INTO mcp_servers (id, name, transport, command, args, url, enabled)
      VALUES (?, ?, ?, ?, ?, ?, 1)
    `).run(
      id, config.name, config.transport,
      config.command ?? null,
      config.args ? JSON.stringify(config.args) : null,
      config.url ?? null,
    );

    if (config.enabled !== false) {
      await this.connectServer({ ...config, id, enabled: true });
    }
  }

  /**
   * 移除 MCP Server
   */
  async removeServer(name: string): Promise<void> {
    await this.disconnectServer(name);
    const db = this.getDB();
    db.prepare('DELETE FROM mcp_servers WHERE name = ?').run(name);
    this.servers.delete(name);
    log.info(`MCP server "${name}" removed`);
  }

  /**
   * 获取所有 Server 状态
   */
  getServerStates(): Array<{
    name: string;
    status: MCPServerStatus;
    toolCount: number;
    error?: string;
  }> {
    return Array.from(this.servers.values()).map((s) => ({
      name: s.config.name,
      status: s.status,
      toolCount: s.tools.length,
      error: s.error,
    }));
  }

  /**
   * 获取单个 Server 状态
   */
  getServerState(name: string): {
    name: string;
    status: MCPServerStatus;
    toolCount: number;
    error?: string;
  } | null {
    const state = this.servers.get(name);
    if (!state) return null;
    return {
      name: state.config.name,
      status: state.status,
      toolCount: state.tools.length,
      error: state.error,
    };
  }

  /**
   * 关闭所有连接
   */
  async shutdown(): Promise<void> {
    const names = Array.from(this.servers.keys());
    await Promise.allSettled(names.map((n) => this.disconnectServer(n)));
    log.info('MCP Manager shut down');
  }
}