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';
// v0.4.1: streamable HTTP 传输（MCP 当前主流远程传输方式）
import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.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 导致初始化崩溃
/** @visibleForTesting 纯函数，供安全表测直接断言 */
export 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 元字符
 */
/** @visibleForTesting 纯函数，供安全表测直接断言 */
export 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 配置文件传递，而非环境变量。
 */
/** @visibleForTesting 纯函数，供安全表测直接断言 */
export 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;
}

/**
 * v0.7.2 P2-8: 安全解析 MCP headers 列（JSON 对象，键值均为字符串）。
 *
 * 数据库中存储为 JSON 字符串（mcp_servers.headers 列自建表起即存在，
 * 此前从未被读写）。解析规则：
 * - 空/null → undefined（匿名连接）
 * - 非对象 / 数组 / 含非字符串键值 → 丢弃该条目并 WARN（保持可用性优先，
 *   不因单列损坏阻断整个 server 连接）
 */
/** @visibleForTesting 纯函数，供安全表测直接断言 */
export function safeParseHeaders(
  raw: string | null | undefined,
): Record<string, string> | undefined {
  if (!raw) return undefined;
  try {
    const parsed: unknown = JSON.parse(raw);
    if (!parsed || typeof parsed !== 'object' || Array.isArray(parsed)) return undefined;
    const out: Record<string, string> = {};
    for (const [k, v] of Object.entries(parsed as Record<string, unknown>)) {
      if (typeof v === 'string') {
        out[k] = v;
      } else {
        log.warn(`MCP headers entry "${k}" dropped: value is not a string`);
      }
    }
    return Object.keys(out).length > 0 ? out : undefined;
  } catch {
    log.warn(`MCP headers parse failed, ignoring: ${raw.slice(0, 100)}`);
    return undefined;
  }
}

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

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

export interface MCPServerConfig {
  id: string;
  name: string;
  /** v0.4.1: 新增 'streamable-http'（MCP 当前主流远程传输）；'sse' 保留向后兼容 */
  transport: 'stdio' | 'sse' | 'streamable-http';
  command?: string;
  args?: string[];
  url?: string;
  /**
   * v0.7.2 P2-8: 远程传输（sse / streamable-http）的自定义请求头。
   * 用于 Bearer/Basic 鉴权或网关路由头；仅 HTTP 传输生效（stdio 走 buildSafeEnv）。
   * 通过 requestInit 注入 —— MCP SDK 的 SSE/StreamableHTTP 客户端在
   * SSE GET 流与 JSON-RPC POST 中统一合并 _commonHeaders（含 requestInit.headers）。
   */
  headers?: Record<string, 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>();
  /**
   * 工具集合变更回调（v0.5.3）
   *
   * connectServer / disconnectServer 完成后触发，供调用方同步引擎工具列表。
   * 背景：setToolsAll 只作用于已存在的引擎；MCP 工具在运行中增删时，
   * 已打开会话的引擎不会自动感知 — 不回调同步则已有引擎持有失效工具定义
   * （调用报 Unknown tool）或缺失新工具（与 README"无需重启"的宣称不符）。
   * 懒创建引擎由 createEngine 从 registry 实时拉取（v0.5.2），无需此回调。
   */
  private toolsChangedCallback: (() => void) | null = null;

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

  /** 注册工具集合变更回调（main.ts 在 AgentEngineManager 创建后注入） */
  setOnToolsChanged(callback: () => void): void {
    this.toolsChangedCallback = callback;
  }

  /** 工具集合变更后的统一通知（连接注册完成 / 断开注销完成） */
  private notifyToolsChanged(): void {
    try {
      this.toolsChangedCallback?.();
    } catch (err) {
      // 回调失败不影响 MCP 主流程（仅引擎工具列表滞后，下轮 setToolsAll 兜底）
      log.warn('[MCPManager] toolsChanged callback failed:', err);
    }
  }

  /**
   * 初始化：从数据库加载已启用的 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;
      headers: 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,
          // v0.4.1: 支持三种传输方式（stdio / sse / streamable-http）
          transport: row.transport as MCPServerConfig['transport'],
          command: row.command ?? undefined,
          args: row.args ? safeParseArgs(row.args) : undefined,
          url: row.url ?? undefined,
          // v0.7.2 P2-8: headers 列接线（此前建表即存在但从未读写）
          headers: safeParseHeaders(row.headers),
          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，旧式传输，保留向后兼容）
        // v0.7.2 P2-8: 注入自定义请求头（非空时）—— requestInit.headers 同时
        // 覆盖 SSE GET 流与 JSON-RPC POST（SDK _commonHeaders 统一合并点）
        const headerInit = {
          ...(config.headers && Object.keys(config.headers).length > 0
            ? { requestInit: { headers: config.headers } }
            : {}),
        };
        transport = new SSEClientTransport(new URL(config.url), headerInit);
      } else if (config.transport === 'streamable-http' && config.url) {
        // v0.4.1: streamable HTTP 模式（MCP 当前主流远程传输）
        const headerInit = {
          ...(config.headers && Object.keys(config.headers).length > 0
            ? { requestInit: { headers: config.headers } }
            : {}),
        };
        transport = new StreamableHTTPClientTransport(new URL(config.url), headerInit);
      } else {
        throw new Error(
          `Unsupported transport "${config.transport}". ` +
            `'stdio' requires 'command', 'sse'/'streamable-http' 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（v0.6.4: registerMCP 对重名冲突返回 false，此处聚合上报）
      let registeredCount = 0;
      let skippedCount = 0;
      for (const tool of tools) {
        const adapter = new MCPToolAdapter(tool, client, name);
        if (this.toolRegistry.registerMCP(name, adapter)) {
          registeredCount++;
        } else {
          skippedCount++;
        }
      }

      // v0.5.3: 工具集合已变化 — 通知调用方同步引擎工具列表（已存在引擎热更新）
      this.notifyToolsChanged();

      // 更新状态
      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: ${registeredCount} tool(s) registered` +
          (skippedCount > 0 ? `, ${skippedCount} skipped due to name conflicts` : ''),
      );
    } 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 = [];

    // v0.5.3: 工具集合已变化 — 通知调用方同步引擎（已有引擎需移除失效工具定义）
    this.notifyToolsChanged();

    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;
            headers: string | null;
          }
        | undefined;
      if (row) {
        await this.connectServer({
          id: row.id,
          name: row.name,
          transport: row.transport as MCPServerConfig['transport'],
          command: row.command ?? undefined,
          args: row.args ? safeParseArgs(row.args) : undefined,
          url: row.url ?? undefined,
          headers: safeParseHeaders(row.headers),
          enabled: true,
        });
      }
    } else {
      await this.disconnectServer(name);
    }
  }

  /**
   * 添加新的 MCP Server
   *
   * v0.4.1: 校验 transport 与对应字段匹配（stdio→command，sse/streamable-http→url）
   * v0.7.2 P2-8: headers 持久化（JSON 列；仅远程传输消费，stdio 忽略）
   */
  async addServer(config: Omit<MCPServerConfig, 'id'>): Promise<void> {
    const db = this.getDB();
    const id = `mcp_${nanoid(8)}`;

    // 校验传输方式与必填字段
    if (config.transport === 'stdio' && !config.command) {
      throw new Error('stdio transport requires "command"');
    }
    if ((config.transport === 'sse' || config.transport === 'streamable-http') && !config.url) {
      throw new Error(`${config.transport} transport requires "url"`);
    }

    db.prepare(
      `
      INSERT INTO mcp_servers (id, name, transport, command, args, url, headers, enabled)
      VALUES (?, ?, ?, ?, ?, ?, ?, 1)
    `,
    ).run(
      id,
      config.name,
      config.transport,
      config.command ?? null,
      config.args ? JSON.stringify(config.args) : null,
      config.url ?? null,
      config.headers && Object.keys(config.headers).length > 0
        ? JSON.stringify(config.headers)
        : 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');
  }
}