metona-ai-desktop/electron/ipc/mcp.ts

/**
 * IPC MCP Handlers — MCP 服务管理域（P2-9 从 handlers.ts 拆分）
 */

import { ipcMain } from 'electron';
import type { IPCContext } from './context';
import log from 'electron-log';
import { assertSafeConfigTarget } from '../harness/tools/built-in/ssrf-guard';

export function registerMCPHandlers(ctx: IPCContext): void {
  const { mcpManager } = ctx;

  ipcMain.handle('mcp:listServers', async () => {
    return mcpManager.getServerStates();
  });

  ipcMain.handle(
    'mcp:addServer',
    async (
      _event,
      config: {
        name: string;
        transport: string;
        command?: string;
        args?: string[];
        url?: string;
        headers?: Record<string, string>;
      },
    ) => {
      // M-38 修复: 完整参数校验，防止字段缺失或类型不符导致异常行为
      if (!config || typeof config !== 'object') {
        return { success: false, error: 'Invalid config' };
      }
      if (typeof config.name !== 'string' || !config.name.trim()) {
        return { success: false, error: 'Server name is required' };
      }
      // M-5 修复: transport 运行时校验（替代 as 断言）
      // v0.4.1: 新增 'streamable-http' 传输方式
      if (
        config.transport !== 'stdio' &&
        config.transport !== 'sse' &&
        config.transport !== 'streamable-http'
      ) {
        return {
          success: false,
          error: `Invalid transport: ${config.transport}. Must be 'stdio', 'sse', or 'streamable-http'`,
        };
      }
      // stdio 类型必须有 command
      if (
        config.transport === 'stdio' &&
        (typeof config.command !== 'string' || !config.command.trim())
      ) {
        return { success: false, error: 'command is required for stdio transport' };
      }
      // sse / streamable-http 类型必须有合法 url
      if (config.transport === 'sse' || config.transport === 'streamable-http') {
        if (typeof config.url !== 'string' || !config.url.trim()) {
          return { success: false, error: `url is required for ${config.transport} transport` };
        }
        try {
          new URL(config.url);
        } catch {
          return { success: false, error: 'Invalid url format' };
        }
        // v0.7.4 P2-9: 配置类 URL 高危目标校验（云元数据/链路本地/组播/保留段拦截；
        // 本地回环与私网实例放行——本地 MCP server 是合法用例）
        try {
          assertSafeConfigTarget(config.url);
        } catch (err) {
          return { success: false, error: (err as Error).message };
        }
      }
      // v0.7.2 P2-8: headers 校验 —— 可选，必须是扁平的 string→string 对象。
      // 上限约束（≤20 项、键 ≤128 字符、值 ≤4096 字符）防止把 headers
      // 当作数据通道滥用；非法条目整体拒绝而非静默丢弃（配置期显式失败）。
      if (config.headers !== undefined) {
        if (
          !config.headers ||
          typeof config.headers !== 'object' ||
          Array.isArray(config.headers)
        ) {
          return {
            success: false,
            error: 'headers must be an object of string key-value pairs',
          };
        }
        const entries = Object.entries(config.headers as Record<string, unknown>);
        if (entries.length > 20) {
          return { success: false, error: 'headers supports at most 20 entries' };
        }
        for (const [k, v] of entries) {
          if (typeof k !== 'string' || !k.trim() || k.length > 128) {
            return { success: false, error: `Invalid header name: ${String(k).slice(0, 40)}` };
          }
          if (typeof v !== 'string' || v.length > 4096) {
            return { success: false, error: `Invalid header value for "${k.slice(0, 40)}"` };
          }
        }
      }
      try {
        await mcpManager.addServer({
          name: config.name,
          transport: config.transport, // 已校验，无需断言
          command: config.command,
          args: config.args,
          url: config.url,
          headers: config.headers as Record<string, string> | undefined, // 已逐项校验
          enabled: true,
        });
        log.info(`MCP server added: ${config.name}`);
        return { success: true };
      } catch (error) {
        return { success: false, error: error instanceof Error ? error.message : String(error) };
      }
    },
  );

  ipcMain.handle('mcp:removeServer', async (_event, name: string) => {
    // M-38 修复: name 校验
    if (typeof name !== 'string' || !name.trim()) {
      return { success: false, error: 'Invalid server name' };
    }
    try {
      await mcpManager.removeServer(name);
      return { success: true };
    } catch (error) {
      return { success: false, error: error instanceof Error ? error.message : String(error) };
    }
  });

  ipcMain.handle('mcp:toggleServer', async (_event, name: string, enabled: boolean) => {
    // M-38 修复: name 和 enabled 校验
    if (typeof name !== 'string' || !name.trim()) {
      return { success: false, error: 'Invalid server name' };
    }
    if (typeof enabled !== 'boolean') {
      return { success: false, error: 'Invalid enabled flag' };
    }
    try {
      await mcpManager.toggleServer(name, enabled);
      return { success: true };
    } catch (error) {
      return { success: false, error: error instanceof Error ? error.message : String(error) };
    }
  });

  // ===== v0.8.0 P2-5: MCP Resources / Prompts 发现 =====
  ipcMain.handle('mcp:listServerContents', async (_event, name: unknown) => {
    if (typeof name !== 'string' || !name.trim()) {
      return { success: false, error: 'Invalid server name' };
    }
    return { success: true, data: mcpManager.getServerContents(name) };
  });

  // ===== v0.8.1 P1-4: Prompt 渲染（prompts/get，ChatInput 斜杠菜单消费） =====
  ipcMain.handle(
    'mcp:getPrompt',
    async (_event, serverName: unknown, promptName: unknown, args: unknown) => {
      if (typeof serverName !== 'string' || !serverName.trim()) {
        return { success: false, error: 'Invalid server name' };
      }
      if (typeof promptName !== 'string' || !promptName.trim()) {
        return { success: false, error: 'Invalid prompt name' };
      }
      // args 可选，必须是 string→string 扁平对象（MCP prompts/get arguments 契约）
      let promptArgs: Record<string, string> | undefined;
      if (args !== undefined && args !== null) {
        if (!args || typeof args !== 'object' || Array.isArray(args)) {
          return { success: false, error: 'Invalid arguments' };
        }
        promptArgs = {};
        for (const [k, v] of Object.entries(args as Record<string, unknown>)) {
          if (typeof v !== 'string') {
            return { success: false, error: `Argument "${k.slice(0, 40)}" must be a string` };
          }
          promptArgs[k] = v;
        }
      }
      try {
        const result = await mcpManager.getPrompt(serverName, promptName, promptArgs);
        return { success: true, data: result };
      } catch (error) {
        return { success: false, error: error instanceof Error ? error.message : String(error) };
      }
    },
  );

  // ===== v0.8.1 P1-4: Resource 读取（resources/read，@mcp 提及消费；512KB 上限） =====
  const MAX_RESOURCE_BYTES = 512 * 1024;
  ipcMain.handle('mcp:readResource', async (_event, serverName: unknown, uri: unknown) => {
    if (typeof serverName !== 'string' || !serverName.trim()) {
      return { success: false, error: 'Invalid server name' };
    }
    if (typeof uri !== 'string' || !uri.trim()) {
      return { success: false, error: 'Invalid resource uri' };
    }
    try {
      const result = await mcpManager.readResource(serverName, uri);
      if (result && result.text.length > MAX_RESOURCE_BYTES) {
        return {
          success: true,
          data: { text: result.text.slice(0, MAX_RESOURCE_BYTES), truncated: true },
        };
      }
      return { success: true, data: { text: result?.text ?? '', truncated: false } };
    } catch (error) {
      return { success: false, error: error instanceof Error ? error.message : String(error) };
    }
  });
}