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 [];
  }
}

// #6 修复: MCP stdio 命令安全校验
/**
 * 允许的 MCP Server 启动命令白名单。
 * 仅允许常见的 MCP Server 运行时，防止任意命令执行。
 */
const ALLOWED_MCP_COMMANDS = new Set([
  'npx', 'node', 'npm',
  'python', 'python3', 'uv', 'uvx',
  'bun', 'deno',
]);

/**
 * #6 修复: 校验 MCP Server 的 command 和 args，防止命令注入
 *
 * 1. 命令白名单：只允许已知的运行时命令
 * 2. 参数注入检测：拒绝包含 shell 元字符的参数
 *
 * @param command MCP Server 启动命令
 * @param args MCP Server 启动参数
 * @throws 如果命令不在白名单或参数包含 shell 元字符
 */
function validateMcpCommand(command: string, args: string[]): void {
  // 提取命令 basename（处理 /usr/bin/node、C:\node\node.exe 等路径）
  const baseCmd = command.split(/[\\/]/).pop()?.replace(/\.exe$/i, '') ?? command;

  if (!ALLOWED_MCP_COMMANDS.has(baseCmd)) {
    throw new Error(
      `MCP command "${baseCmd}" is not in the allowed list: ${[...ALLOWED_MCP_COMMANDS].join(', ')}. ` +
      `For security reasons, only standard MCP runtimes are permitted.`,
    );
  }

  // 审查修复: 移除 {}() 字符 — StdioClientTransport 用 spawn（不经 shell），
  // 这些字符无注入风险，但 MCP Server 的 args 常含 JSON 配置（如 --config {"port":3000}）会被误拒。
  // 保留 ; & | ` $ < > 换行 等高危字符。
  const shellMetacharPattern = /[;&|`$<>\n\r]/;
  for (const arg of args) {
    if (shellMetacharPattern.test(arg)) {
      throw new Error(
        `MCP command argument contains shell metacharacters and was rejected: ${arg.slice(0, 100)}`,
      );
    }
  }
}

/**
 * #6 修复 + 审查修复: 构建安全的子进程环境变量
 *
 * 审查修复: 原白名单方案过于激进，剥离了 MCP Server 运行所需的 npm_config_*、代理变量等，
 * 导致 MCP Server 无法启动。改为黑名单方案：剔除包含敏感后缀的变量，保留其余。
 *
 * 注意: GITHUB_TOKEN / SLACK_BOT_TOKEN 等含 _TOKEN 后缀的变量也会被过滤。
 * 如果 MCP Server 需要这些凭证，应通过 MCP Server 配置文件传递，而非环境变量。
 */
function buildSafeEnv(): Record<string, string> {
  // 敏感变量后缀黑名单 — 匹配这些后缀的变量不会被传递给子进程
  const SENSITIVE_SUFFIXES = [
    '_API_KEY', '_TOKEN', '_SECRET', '_PASSWORD', '_PASSWD',
    '_CREDENTIAL', '_CREDENTIALS', '_PRIVATE_KEY',
  ];
  // 敏感变量名黑名单（精确匹配）
  const SENSITIVE_KEYS = new Set([
    'DEEPSEEK_API_KEY', 'AGNES_API_KEY', 'MIMO_API_KEY',
    'GITEA_PASSWORD', 'DATABASE_PASSWORD',
  ]);

  const env: Record<string, string> = {};
  for (const [key, val] of Object.entries(process.env)) {
    if (!val) continue;
    // 跳过敏感变量名
    if (SENSITIVE_KEYS.has(key)) continue;
    // 跳过敏感后缀变量
    if (SENSITIVE_SUFFIXES.some((suffix) => key.toUpperCase().endsWith(suffix))) continue;
    env[key] = val;
  }
  return env;
}

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

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 并连接
   *
   * P1-11: 等待所有连接完成（或 30s 超时）后才返回，保证调用方广播 tools:ready
   * 时 MCP 工具已实际注册完成。原实现 connectServer 不 await，工具就绪广播
   * 早于工具注册，存在"广播后一段时间内 MCP 工具仍不可用"的窗口。
   *
   * 单个 server 连接超时（30s）用于兜底：stdio server 挂起时不阻塞整体就绪，
   * 超时的 server 会显示 error 状态，用户可在设置面板查看原因。
   */
  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;
    }>;

    const connectWithTimeout = (config: MCPServerConfig): Promise<unknown> =>
      Promise.race([
        this.connectServer(config),
        new Promise((resolve) => setTimeout(() => resolve('timeout'), 30_000)),
      ]);

    const settled = await Promise.allSettled(
      rows.map((row) => {
        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,
        };
        return connectWithTimeout(config).catch((err) => {
          log.warn(`MCP server "${config.name}" auto-connect failed: ${err}`);
        });
      }),
    );

    const failed = settled.filter((s) => s.status === 'rejected').length;
    log.info(`MCP Manager initialized: ${rows.length} server(s) configured (${failed} failed)`);
  }

  /**
   * 连接 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 ?? [];
        // #6 修复: 命令白名单 + 参数元字符检测，防止命令注入
        validateMcpCommand(config.command, args);
        transport = new StdioClientTransport({
          command: config.command,
          args,
          // #6 修复: 不透传完整 process.env，仅保留 MCP Server 运行所需的最小环境变量
          env: buildSafeEnv(),
        });
      } 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');
  }
}