@greatlhd/ailo-desktop 1.0.0

package/src/mcp/session.ts

@@ -0,0 +1,182 @@
+import { spawn, type ChildProcess } from "child_process";
+import type { ToolCapability } from "@greatlhd/ailo-endpoint-sdk";
+
+export interface PendingRPC {
+  resolve: (v: unknown) => void;
+  reject: (e: Error) => void;
+}
+
+export interface StdioSession {
+  kind: "stdio";
+  proc: ChildProcess;
+  tools: ToolCapability[];
+  buffer: string;
+  nextId: number;
+  pendingRequests: Map<number, PendingRPC>;
+  config?: MCPServerConfig;
+}
+
+export interface SSESession {
+  kind: "sse";
+  tools: ToolCapability[];
+  nextId: number;
+  pendingRequests: Map<number, PendingRPC>;
+  messageEndpoint: string;
+  abortController: AbortController;
+  config?: MCPServerConfig;
+}
+
+export type MCPSession = StdioSession | SSESession;
+
+export interface MCPServerConfig {
+  transport?: "stdio" | "sse";
+  command?: string;
+  args?: string[];
+  env?: Record<string, string>;
+  url?: string;
+  enabled?: boolean;
+}
+
+function shellEscape(arg: string): string {
+  if (process.platform === "win32") {
+    if (!/[ "&|<>^%!]/.test(arg)) return arg;
+    return `"${arg.replace(/"/g, '""')}"`;
+  }
+  if (!/[^a-zA-Z0-9_./:=@-]/.test(arg)) return arg;
+  return `'${arg.replace(/'/g, "'\\''")}'`;
+}
+
+export async function startStdioServer(
+  config: MCPServerConfig,
+  onExit: (code: number | null) => void,
+): Promise<StdioSession> {
+  const env = { ...process.env, PYTHONIOENCODING: "utf-8", PYTHONUTF8: "1", ...(config.env ?? {}) };
+  const args = config.args ?? [];
+  const shellCmd = [config.command!, ...args].map(shellEscape).join(" ");
+  const proc = spawn(shellCmd, [], {
+    stdio: ["pipe", "pipe", "pipe"],
+    env,
+    shell: true,
+  });
+
+  const session: StdioSession = {
+    kind: "stdio",
+    proc,
+    tools: [],
+    buffer: "",
+    nextId: 1,
+    pendingRequests: new Map(),
+  };
+
+  proc.stderr?.on("data", (chunk: Buffer) => {
+    const text = chunk.toString("utf-8").trim();
+    if (text) console.error(`[mcp:stderr]`, text);
+  });
+
+  proc.on("exit", (code) => {
+    onExit(code);
+    for (const [, pending] of session.pendingRequests) {
+      pending.reject(new Error(`MCP process exited`));
+    }
+    session.pendingRequests.clear();
+  });
+
+  return session;
+}
+
+export async function startSSEServer(
+  config: MCPServerConfig,
+  onMessage: (session: SSESession, data: string) => void,
+  onEnd: () => void,
+): Promise<SSESession> {
+  const baseUrl = config.url!.replace(/\/+$/, "");
+  const sseUrl = `${baseUrl}/sse`;
+  const abortController = new AbortController();
+
+  const session: SSESession = {
+    kind: "sse",
+    tools: [],
+    nextId: 1,
+    pendingRequests: new Map(),
+    messageEndpoint: "",
+    abortController,
+  };
+
+  const response = await fetch(sseUrl, {
+    signal: abortController.signal,
+    headers: { Accept: "text/event-stream" },
+  });
+  if (!response.ok) throw new Error(`SSE connect failed: HTTP ${response.status}`);
+  if (!response.body) throw new Error("SSE response has no body");
+
+  const reader = response.body.getReader();
+  const decoder = new TextDecoder();
+  let sseBuffer = "";
+  let endpointResolved = false;
+
+  const readLoop = async () => {
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) {
+        if (!abortController.signal.aborted) {
+          onEnd();
+          for (const [, pending] of session.pendingRequests) {
+            pending.reject(new Error(`SSE stream ended`));
+          }
+        }
+        break;
+      }
+      sseBuffer += decoder.decode(value, { stream: true });
+      const lines = sseBuffer.split("\n");
+      sseBuffer = lines.pop() ?? "";
+      let currentEvent = "";
+      let currentData = "";
+      for (const line of lines) {
+        if (line.startsWith("event:")) {
+          currentEvent = line.slice(6).trim();
+        } else if (line.startsWith("data:")) {
+          currentData = line.slice(5).trim();
+        } else if (line.trim() === "" && currentEvent) {
+          if (currentEvent === "endpoint" && !endpointResolved) {
+            endpointResolved = true;
+            let ep = currentData;
+            if (ep.startsWith("/")) ep = `${baseUrl}${ep}`;
+            session.messageEndpoint = ep;
+          } else if (currentEvent === "message") {
+            onMessage(session, currentData);
+          }
+          currentEvent = "";
+          currentData = "";
+        }
+      }
+    }
+  };
+
+  readLoop().catch((err) => {
+    if (!abortController.signal.aborted) {
+      console.error(`[mcp:sse] read error:`, err.message);
+      onEnd();
+    }
+  });
+
+  return session;
+}
+
+export function stopSession(session: MCPSession): Promise<void> {
+  if (session.kind === "stdio") {
+    try { session.proc.kill("SIGTERM"); } catch {}
+    return new Promise((resolve) => {
+      const timer = setTimeout(() => {
+        try { session.proc.kill("SIGKILL"); } catch {}
+        resolve();
+      }, 3000);
+      session.proc.once("exit", () => { clearTimeout(timer); resolve(); });
+    });
+  } else {
+    session.abortController.abort();
+    for (const [, pending] of session.pendingRequests) {
+      pending.reject(new Error(`SSE session stopped`));
+    }
+    return Promise.resolve();
+  }
+}

package/src/mcp_manager.ts

@@ -0,0 +1,273 @@
+import { MCPConfigManager, startStdioServer, startSSEServer, stopSession, initializeSession, createStdioRpc, createSSERpc, handleSSEMessage } from "./mcp/index.js";
+import type { MCPServerConfig, MCPSession, StdioSession, SSESession } from "./mcp/index.js";
+import type { ToolCapability } from "@greatlhd/ailo-endpoint-sdk";
+
+type Args = Record<string, unknown>;
+
+export class LocalMCPManager {
+  private configManager = new MCPConfigManager();
+  private sessions = new Map<string, MCPSession>();
+  private watchTimer: ReturnType<typeof setInterval> | null = null;
+  private onToolsChanged: (() => void) | null = null;
+
+  setOnToolsChanged(cb: () => void): void {
+    this.onToolsChanged = cb;
+  }
+
+  async init(): Promise<void> {
+    await this.configManager.load();
+    for (const [name, cfg] of this.configManager.getAll()) {
+      if (cfg.enabled !== false) {
+        await this.startServer(name, cfg).catch((e) => console.error(`[mcp] failed to start ${name}:`, e.message));
+      }
+    }
+  }
+
+  startWatching(intervalMs = 2000): void {
+    if (this.watchTimer) return;
+    this.watchTimer = setInterval(() => this.checkConfigChange(), intervalMs);
+  }
+
+  stopWatching(): void {
+    if (this.watchTimer) { clearInterval(this.watchTimer); this.watchTimer = null; }
+  }
+
+  async handle(args: Args): Promise<{ text: string; toolsChanged: boolean }> {
+    const action = args.action as string;
+    switch (action) {
+      case "list": return { text: this.doList(), toolsChanged: false };
+      case "create": return await this.doCreate(args);
+      case "delete": return await this.doDelete(args);
+      case "start": return await this.doStart(args);
+      case "stop": return await this.doStop(args);
+      default: throw new Error(`未知 action: ${action}`);
+    }
+  }
+
+  getAllPrivateTools(): ToolCapability[] {
+    const all: ToolCapability[] = [];
+    for (const session of this.sessions.values()) {
+      all.push(...session.tools);
+    }
+    return all;
+  }
+
+  async shutdown(): Promise<void> {
+    this.stopWatching();
+    for (const [name] of this.sessions) {
+      await this.stopServer(name);
+    }
+  }
+
+  getConfigs(): Map<string, MCPServerConfig> {
+    return this.configManager.getAll();
+  }
+
+  isRunning(name: string): boolean {
+    return this.sessions.has(name);
+  }
+
+  getToolsForServer(name: string): ToolCapability[] {
+    return this.sessions.get(name)?.tools ?? [];
+  }
+
+  private async checkConfigChange(): Promise<void> {
+    const changed = await this.configManager.checkChanged();
+    if (!changed) return;
+
+    console.log("[mcp] config file changed, reloading...");
+    await this.configManager.load();
+    const newConfigs = this.configManager.getAll();
+    let toolsChanged = false;
+
+    for (const [name] of this.sessions) {
+      if (!newConfigs.has(name) || newConfigs.get(name)?.enabled === false) {
+        await this.stopServer(name);
+        toolsChanged = true;
+      }
+    }
+
+    for (const [name, cfg] of newConfigs) {
+      if (cfg.enabled === false) continue;
+      const existing = this.sessions.get(name);
+      if (!existing) {
+        await this.startServer(name, cfg).catch((e) =>
+          console.error(`[mcp] reload failed for ${name}:`, e.message),
+        );
+        toolsChanged = true;
+      }
+    }
+
+    if (toolsChanged) this.onToolsChanged?.();
+  }
+
+  private async startServer(name: string, config: MCPServerConfig): Promise<void> {
+    if (this.sessions.has(name)) await this.stopServer(name);
+
+    const transport = config.transport ?? "stdio";
+    if (transport === "sse") {
+      await this.startSSEServer(name, config);
+    } else {
+      await this.startStdioServer(name, config);
+    }
+  }
+
+  private async startStdioServer(name: string, config: MCPServerConfig): Promise<void> {
+    const session = await startStdioServer(config, (code) => {
+      console.log(`[mcp] ${name} exited with code ${code}`);
+      this.sessions.delete(name);
+    });
+
+    session.config = config;
+    this.sessions.set(name, session);
+
+    const { request, notify, processBuffer } = createStdioRpc(session);
+    session.proc.stdout!.on("data", (chunk: Buffer) => {
+      session.buffer += chunk.toString("utf-8");
+      processBuffer();
+    });
+
+    await this.initMCPSession(name, session, request, notify);
+  }
+
+  private async startSSEServer(name: string, config: MCPServerConfig): Promise<void> {
+    const session = await startSSEServer(config, (sess, data) => {
+      handleSSEMessage(sess, data);
+    }, () => {
+      console.log(`[mcp:sse] ${name} SSE stream ended`);
+      this.sessions.delete(name);
+    });
+
+    session.config = config;
+    this.sessions.set(name, session);
+    console.log(`[mcp:sse] ${name} connected, message endpoint: ${session.messageEndpoint}`);
+
+    const { request, notify } = createSSERpc(session);
+    await this.initMCPSession(name, session, request, notify);
+  }
+
+  private async initMCPSession(
+    name: string,
+    session: MCPSession,
+    request: (method: string, params: unknown) => Promise<unknown>,
+    notify: (method: string, params: unknown) => void,
+  ): Promise<void> {
+    try {
+      const tools = await initializeSession(session, request, notify);
+      session.tools = tools.map((t) => ({
+        name: `${name}:${t.name}`,
+        description: t.description,
+        parameters: t.parameters,
+      }));
+      console.log(`[mcp] ${name} started (${session.kind}), discovered ${session.tools.length} tool(s)`);
+    } catch (e: unknown) {
+      const msg = e instanceof Error ? e.message : String(e);
+      console.error(`[mcp] ${name} init failed:`, msg);
+      await this.stopServer(name);
+      throw e;
+    }
+  }
+
+  private async stopServer(name: string): Promise<void> {
+    const session = this.sessions.get(name);
+    if (!session) return;
+    this.sessions.delete(name);
+    await stopSession(session);
+  }
+
+  async executeToolCall(serverName: string, toolName: string, args: Record<string, unknown>): Promise<unknown> {
+    const session = this.sessions.get(serverName);
+    if (!session) throw new Error(`MCP server not running: ${serverName}`);
+
+    if (session.kind === "stdio") {
+      const { request } = createStdioRpc(session);
+      return request("tools/call", { name: toolName, arguments: args });
+    } else {
+      const { request } = createSSERpc(session);
+      return request("tools/call", { name: toolName, arguments: args });
+    }
+  }
+
+  private doList(): string {
+    const configs = this.configManager.getAll();
+    if (configs.size === 0) return "没有已注册的 MCP 服务";
+    const lines: string[] = ["MCP 服务列表："];
+    for (const [name, cfg] of configs) {
+      const session = this.sessions.get(name);
+      const status = session ? "运行中" : (cfg.enabled !== false ? "已停止" : "已禁用");
+      const transport = cfg.transport ?? "stdio";
+      const toolCount = session?.tools.length ?? 0;
+      lines.push(`- ${name}: ${status} | ${transport} | ${toolCount} 工具`);
+      if (session && session.tools.length > 0) {
+        for (const t of session.tools) {
+          lines.push(`  - ${t.name}: ${t.description ?? ""}`);
+        }
+      }
+    }
+    return lines.join("\n");
+  }
+
+  private async doCreate(args: Args): Promise<{ text: string; toolsChanged: boolean }> {
+    const name = args.name as string;
+    if (!name) throw new Error("name 必填");
+    const transport = (args.transport as string) ?? "stdio";
+    let config: MCPServerConfig;
+    if (transport === "sse") {
+      const url = args.url as string;
+      if (!url) throw new Error("SSE 模式下 url 必填");
+      config = { transport: "sse", url, env: args.env as Record<string, string>, enabled: true };
+    } else {
+      const command = args.command as string;
+      if (!command) throw new Error("command 必填");
+      config = { transport: "stdio", command, args: args.args as string[], env: args.env as Record<string, string>, enabled: true };
+    }
+    this.configManager.set(name, config);
+    await this.configManager.save();
+    try {
+      await this.startServer(name, config);
+      const session = this.sessions.get(name);
+      const toolCount = session?.tools.length ?? 0;
+      return { text: `已创建并启动 MCP 服务 ${name}，发现 ${toolCount} 个工具`, toolsChanged: true };
+    } catch (e: unknown) {
+      const msg = e instanceof Error ? e.message : String(e);
+      return { text: `已创建 MCP 服务 ${name}，但启动失败: ${msg}`, toolsChanged: false };
+    }
+  }
+
+  private async doDelete(args: Args): Promise<{ text: string; toolsChanged: boolean }> {
+    const name = args.name as string;
+    if (!name) throw new Error("name 必填");
+    const hadSession = this.sessions.has(name);
+    if (hadSession) await this.stopServer(name);
+    this.configManager.delete(name);
+    await this.configManager.save();
+    return { text: `已删除 MCP 服务 ${name}`, toolsChanged: hadSession };
+  }
+
+  private async doStart(args: Args): Promise<{ text: string; toolsChanged: boolean }> {
+    const name = args.name as string;
+    if (!name) throw new Error("name 必填");
+    const cfg = this.configManager.get(name);
+    if (!cfg) throw new Error(`MCP 服务不存在: ${name}`);
+    cfg.enabled = true;
+    this.configManager.set(name, cfg);
+    await this.configManager.save();
+    await this.startServer(name, cfg);
+    const session = this.sessions.get(name);
+    const toolCount = session?.tools.length ?? 0;
+    return { text: `已启动 MCP 服务 ${name}，发现 ${toolCount} 个工具`, toolsChanged: true };
+  }
+
+  private async doStop(args: Args): Promise<{ text: string; toolsChanged: boolean }> {
+    const name = args.name as string;
+    if (!name) throw new Error("name 必填");
+    const cfg = this.configManager.get(name);
+    if (!cfg) throw new Error(`MCP 服务不存在: ${name}`);
+    cfg.enabled = false;
+    this.configManager.set(name, cfg);
+    await this.configManager.save();
+    const hadSession = this.sessions.has(name);
+    await this.stopServer(name);
+    return { text: `已停止 MCP 服务 ${name}`, toolsChanged: hadSession };
+  }
+}