ModelRouter/cmd/gui/main.js

// ModelRouter Desktop GUI — main process
// Clash-Verge-style: bundles and manages an embedded llmsproxy core instance
// in a per-user profile dir (userData), auto-starts it, keeps a tray icon,
// supports launch-at-login. Server-side users keep using the pure Go binary.
const {
  app,
  BrowserWindow,
  ipcMain,
  Menu,
  Tray,
  shell,
  nativeImage,
  nativeTheme,
} = require("electron");
const path = require("path");
const fs = require("fs");
const http = require("http");
const os = require("os");
const crypto = require("crypto");
const { spawn } = require("child_process");

const LOG_FILE = path.join(app.getPath("userData"), "gui.log");
function log(msg) {
  try {
    fs.appendFileSync(LOG_FILE, new Date().toISOString() + " " + msg + "\n");
  } catch (e) {}
}
const _consoleLog = console.log,
  _consoleErr = console.error;
console.log = function () {
  log(Array.prototype.slice.call(arguments).join(" "));
  _consoleLog.apply(null, arguments);
};
console.error = function () {
  log("ERR " + Array.prototype.slice.call(arguments).join(" "));
  _consoleErr.apply(null, arguments);
};

// ---------- embedded core profile ----------
// Per-user profile dir: <userData>/profile
const PROFILE_DIR = path.join(app.getPath("userData"), "profile");
const CONFIG_FILE = path.join(PROFILE_DIR, "config.yaml");
// Packaged: extraResource at process.resourcesPath/bin (outside asar, spawnable).
// Dev:   cmd/gui/bin, populated by prepare-core.js.
const CORE_EXE = path.join(
  app.isPackaged ? process.resourcesPath : __dirname,
  "bin",
  process.platform === "win32" ? "llmsproxy.exe" : "llmsproxy",
);
// default port for the embedded core (configurable via settings)
const DEFAULT_PORT = 8787;

let mainWindow = null;
let tray = null;
let coreProc = null;
let coreStartedAt = 0;
let coreReady = false; // true only after the embedded core actually answers HTTP
let settings = {}; // { port, autoStart, minimizeToTray }
let authRules = [];

const isDev = process.argv.includes("--dev");
const silentFlag = process.argv.includes("--silent");

function loadSettings() {
  try {
    const p = path.join(app.getPath("userData"), "settings.json");
    if (fs.existsSync(p)) settings = JSON.parse(fs.readFileSync(p, "utf-8"));
  } catch (e) {}
  if (typeof settings.port !== "number") settings.port = DEFAULT_PORT;
  if (typeof settings.silentStart !== "boolean") settings.silentStart = false;
  if (typeof settings.minimizeToTray !== "boolean")
    settings.minimizeToTray = true;
  return settings;
}
function saveSettings() {
  try {
    fs.writeFileSync(
      path.join(app.getPath("userData"), "settings.json"),
      JSON.stringify(settings, null, 2),
      "utf-8",
    );
  } catch (e) {}
}

// read admin key from embedded config.yaml (generated by core first run)
function readAdminKey() {
  // The embedded profile stores its admin key in `keys` (NOT gateway_keys):
  // gateway_keys entries get marked seed:true by the core (-> "replace the
  // initial key" warning in the web UI), while a plain keys entry authenticates
  // without any seed warning.
  try {
    if (!fs.existsSync(CONFIG_FILE)) return "";
    const raw = fs.readFileSync(CONFIG_FILE, "utf-8");
    // keys:
    //     - key: sk-gw-...
    const k = /^\s*keys:\s*\n(?:\s*-\s*key:\s*"?([^\s"#]+)"?[^\n]*\n)+/m.exec(
      raw,
    );
    if (k && k[1]) return k[1];
    // legacy: gateway_keys list (pre-migration configs)
    const g = /^\s*gateway_keys:\s*\n(?:\s*-\s*"?([^\s"#]+)"?\s*\n)+/m.exec(
      raw,
    );
    if (g && g[1]) return g[1];
    const g2 = /gateway_keys:[^[\n]*\[\s*"?([^\s"\]]+)"?/m.exec(raw);
    if (g2 && g2[1]) return g2[1];
  } catch (e) {}
  return "";
}

const embeddedBaseUrl = () =>
  "http://127.0.0.1:" + (settings.port || DEFAULT_PORT);

// ---------- embedded core lifecycle ----------
function isCoreRunning() {
  return coreProc !== null && coreProc.exitCode === null;
}

function coreStarted() {
  return coreStartedAt > 0 && isCoreRunning();
}

function writeDefaultEmbeddedConfig() {
  // Generate a self-contained embedded profile with a fresh random admin key.
  // No upstream sources are preconfigured — the desktop user adds them via
  // the WebUI's sources page (like Clash Verge's empty-by-default profile).
  // If a config already exists we do NOT overwrite it (keep user's edits).
  if (fs.existsSync(CONFIG_FILE)) return;
  const key = "sk-gw-" + crypto.randomBytes(16).toString("hex");
  const yaml = [
    "# ModelRouter embedded profile (auto-generated)",
    "listen: 127.0.0.1:" + (settings.port || DEFAULT_PORT),
    "default_model: AUTO",
    "adapter_dir: " + JSON.stringify(path.join(PROFILE_DIR, "adapters")),
    "runtime_file: " + JSON.stringify(path.join(PROFILE_DIR, "runtime.json")),
    "",
    // Admin key as a plain `keys` entry (no seed flag) so the web UI never
    // shows the replace-the-initial-key warning. gateway_keys entries are
    // marked seed:true by the core.
    "keys:",
    "  - key: " + key,
    "    role: admin",
    "    name: embedded",
  ].join("\n");
  try {
    fs.mkdirSync(PROFILE_DIR, { recursive: true });
    fs.writeFileSync(CONFIG_FILE, yaml, "utf-8");
  } catch (e) {
    console.error("write embedded config failed:", e);
  }
}

// does the embedded core answer already? returns its admin key if so
function probeEmbedded() {
  // Use the admin key from the embedded profile so an up core that requires
  // auth (401 without key) is still detected as healthy. Any HTTP response
  // (even 401/200) proves the listener is up; network errors mean it is not.
  return new Promise((resolve) => {
    const u = embeddedBaseUrl();
    const key = readAdminKey();
    const headers = key ? { Authorization: "Bearer " + key } : {};
    const req = http.get(u + "/v1/models", { headers }, (res) => {
      let body = "";
      res.on("data", (c) => (body += c));
      res.on("end", () => resolve({ ok: true, status: res.statusCode, body }));
    });
    req.on("error", () => resolve({ ok: false }));
    req.setTimeout(1200, () => {
      req.destroy();
      resolve({ ok: false });
    });
  });
}

function startCore() {
  if (coreProc && coreProc.exitCode === null) {
    console.log("core already running");
    return;
  }
  if (!fs.existsSync(CORE_EXE)) {
    console.error("embedded core binary missing:", CORE_EXE);
    return;
  }
  // ensure profile dir + config + local adapters exist
  fs.mkdirSync(PROFILE_DIR, { recursive: true });
  if (!fs.existsSync(CONFIG_FILE)) writeDefaultEmbeddedConfig();

  coreProc = spawn(CORE_EXE, ["-config", CONFIG_FILE], {
    cwd: PROFILE_DIR,
    env: Object.assign({}, process.env, {
      LLMS_PROXY_GUI: "1",
    }),
    stdio: ["ignore", "pipe", "pipe"],
  });
  coreStartedAt = Date.now();
  // capture core logs -> gui log
  coreProc.stdout.on("data", (d) => log("[core] " + String(d).trim()));
  coreProc.stderr.on("data", (d) => log("[core-err] " + String(d).trim()));
  coreProc.on("error", (err) => {
    console.error("core spawn error:", err.message);
    coreProc = null;
  });
  coreProc.on("exit", (code) => {
    const wasManaged = coreStartedAt > 0;
    coreStartedAt = 0;
    coreReady = false;
    coreProc = null;
    console.log("core exited:", code);
    // auto-restart if still running GUI (crash recovery, max every ~5s)
    if (wasManaged && !app.isQuitting) {
      const since = Date.now() - (coreLastExit || 0);
      if (since < 2000) return; // avoid tight crash loop
      coreLastExit = Date.now();
      setTimeout(() => {
        if (!app.isQuitting) startCore();
      }, 1500);
    }
  });
  // Broadcast starting (NOT ready) immediately. The renderer must not
  // attempt to load the iframe until the core actually answers HTTP —
  // otherwise the frame load races the listener and fails permanently.
  // (Fix for blank-screen: running=true was previously broadcast while the
  // core was still booting, so the iframe ERR_CONNECTION_REFUSED and the
  // renderer never retried.)
  notifyCoreState();
  // pump until the core answers, then inject the admin key and tell the
  // renderer the web UI is ready (the iframe must load only after auth is set)
  waitForCoreReady()
    .then((ok) => {
      if (ok) {
        setEmbeddedAuth();
        coreReady = true;
        notifyCoreState();
      }
    })
    .catch(() => {});
}

// waitForCoreReady polls probeEmbedded until the listener is up (or timeout).
function waitForCoreReady(timeoutMs = 10000) {
  const start = Date.now();
  return new Promise((resolve) => {
    const poll = () => {
      probeEmbedded().then((p) => {
        if (p.ok) return resolve(true);
        if (Date.now() - start > timeoutMs) return resolve(false);
        setTimeout(poll, 250);
      });
    };
    poll();
  });
}

let coreLastExit = 0;

function stopCore() {
  app.isQuitting = true;
  coreReady = false;
  if (coreProc && coreProc.exitCode === null) {
    try {
      coreProc.kill();
    } catch (e) {}
  }
  coreProc = null;
  coreStartedAt = 0;
  notifyCoreState();
}

function restartCore() {
  app.isQuitting = true;
  if (coreProc && coreProc.exitCode === null) {
    try {
      coreProc.kill();
    } catch (e) {}
  }
  coreProc = null;
  coreStartedAt = 0;
  app.isQuitting = false;
  setTimeout(() => startCore(), 400);
}

// ---------- auth rule injection (embedded only) ----------
function installAuthRule() {
  // Read the admin key live on every request so injection never depends on
  // ordering: as soon as the embedded core has written its config (key exists)
  // the web UI gets authenticated, even if the renderer loaded early.
  const apply = (details, callback) => {
    const url = details.url || "";
    const rule = authRules.find((r) => r.url && url.indexOf(r.url) === 0);
    if (!rule) return callback({ requestHeaders: details.requestHeaders });
    const h = Object.assign({}, details.requestHeaders);
    const key = rule.apiKey || readAdminKey() || "";
    // ModelRouter auth() checks gw_key cookie first; Chromium lets us inject
    // Cookie on top-level nav AND cross-origin XHR/fetch, which the CORS-
    // restricted Authorization header cannot do inside a sandboxed iframe.
    if (key && !/(^|;)\s*gw_key=/.test(h.Cookie || ""))
      h["Cookie"] = (h["Cookie"] ? h["Cookie"] + "; " : "") + "gw_key=" + key;
    if (rule.cookie)
      h["Cookie"] = (h["Cookie"] ? h["Cookie"] + "; " : "") + rule.cookie;
    if (rule.headers)
      Object.keys(rule.headers).forEach((k) => (h[k] = rule.headers[k]));
    callback({ requestHeaders: h });
  };
  require("electron").session.defaultSession.webRequest.onBeforeSendHeaders(
    (details, callback) => {
      try {
        apply(details, callback);
      } catch (e) {
        callback({ requestHeaders: details.requestHeaders });
      }
    },
  );
  require("electron").session.defaultSession.webRequest.onHeadersReceived(
    (details, callback) => {
      const h = Object.assign({}, details.responseHeaders || {});
      const url = details.url || "";
      if (authRules.some((r) => r.url && url.indexOf(r.url) === 0)) {
        h["Access-Control-Allow-Origin"] = ["*"];
        h["Access-Control-Allow-Methods"] = ["GET,POST,PUT,DELETE,OPTIONS"];
        h["Access-Control-Allow-Headers"] = [
          "Content-Type, X-API-Key, Authorization, Cookie, Range",
        ];
      }
      callback({ responseHeaders: h });
    },
  );
}

function setEmbeddedAuth() {
  authRules = [];
  authRules.push({ url: embeddedBaseUrl(), apiKey: readAdminKey() || "" });
}

// ---------- tray ----------
function createTray() {
  // Tray icon: prefer the packaged asset (inside asar) or the dev build dir.
  // GNOME renders tray icons via AppIndicator: a fully transparent PNG with
  // only colored pixels draws as an odd box, so ship an opaque rounded-blue
  // M icon with a subtle white ring for dark/light trays alike.
  const candidates = [
    path.join(__dirname, "build", "icon.png"),
    path.join(process.resourcesPath || "", "build", "icon.png"),
  ];
  let img = null;
  for (const c of candidates) {
    try {
      const t = nativeImage.createFromPath(c);
      if (!t.isEmpty()) {
        img = t;
        break;
      }
    } catch (e) {}
  }
  if (!img) {
    // clear blue-M fallback (never the transparent fabric pixel)
    img = nativeImage.createFromDataURL(
      "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABYAAAAWCAYAAADEtGw7AAAACXBIWXMAAAsTAAALEwEAmpwYAAABkUlEQVQ4jb2Vz0sCQRTH/S+yP8Ydkg7RKRD6FwpzxiWoQyL+Awohab8g6RD04yBeIuhQt7pVkCeFLv4Hzdtf7vhiZnHTXVzUtR485u3j7Wdn38x8J5H4L0vpfFnLASOU1wiD02jnR4QCle9EQgkFnVBuEAY4m3PQGOQnQ9mswIBToKHfJ3PNNDxzkuVJH6x6Ghvqe+4XzHg9WJCvWli68HyjYPh5GQ/zsibcDl4bAUMjWPDWcXFo9abj52tN28+/d0UILFlTgcUA8bXt+vmXtqtyscHdnkDTRlzdBUzrgGANsNMT8cF3z301sqqJuUNTxbdP/fjgwrmlxssHBxv3jooPzqz44M2Sgdwc4OeXwI+uUK3IFI344EzRUAvmCkTH9RZSbrmFgE9aXgukHbecWcC8HgXerniLJm2rbEaCxw6IFI9gQfnaxqtHB9f3DVzRQcXS5ZZb2/OeKzd2CJyisDMuQoxDXJ3QKOdp9r00pnAag/xCBSigyVR+dZ6ZToT68CxPyiJ17UxzNcnaUQ3+a/sBfMAWejRz3/oAAAAASUVORK5CYII=",
    );
  }
  tray = new Tray(img);
  // Rebuild the context menu when the OS theme flips (dark/light) so menu
  // text colors follow the desktop theme; nativeTheme is the Electron API
  // that tracks GTK/appindicator dark mode on Linux.
  const themeRebuild = () => {
    try {
      if (global.__rebuildTray) global.__rebuildTray();
    } catch (e) {}
  };
  try {
    nativeTheme.on("updated", themeRebuild);
  } catch (e) {}
  // Status rows used to be enabled:false, which GTK renders in a fixed light
  // grey that never flips with the system dark/light theme — invisible on
  // light desktops. They are now enabled (normal theme-aware fg) with a
  // no-op click so they stay display-only but inherit the menu text color.
  const statusRow = (text) => ({ label: text, enabled: true, click: () => {} });
  const rebuild = () => {
    const running = coreStarted();
    const menu = Menu.buildFromTemplate([
      statusRow("ModelRouter Desktop"),
      statusRow(running ? "● 内核运行中" : "○ 内核已停止"),
      { type: "separator" },
      {
        label: "显示主窗口",
        click: () => {
          if (mainWindow) {
            mainWindow.show();
            mainWindow.focus();
          }
        },
      },
      { type: "separator" },
      {
        label: "开启内核 (核心网关启动)",
        enabled: !running,
        click: () => startCore(),
      },
      {
        label: "重启内核",
        enabled: running,
        click: () => restartCore(),
      },
      {
        label: "停止内核",
        enabled: running,
        click: () => stopCore(),
      },
      { type: "separator" },
      {
        label: "开机自启",
        type: "checkbox",
        checked: !!settings.autoStart,
        click: () => {
          settings.autoStart = !settings.autoStart;
          applyAutoStart();
          saveSettings();
          rebuild();
        },
      },
      {
        label: "静默启动",
        type: "checkbox",
        checked: !!settings.silentStart,
        click: () => {
          settings.silentStart = !settings.silentStart;
          saveSettings();
          rebuild();
        },
      },
      { type: "separator" },
      {
        label: "退出",
        click: () => app.quit(),
      },
    ]);
    tray.setContextMenu(menu);
  };
  tray.setToolTip("ModelRouter");
  tray.on("click", () => {
    if (mainWindow) {
      mainWindow.isVisible() ? mainWindow.focus() : mainWindow.show();
    }
  });
  rebuild();
  global.__rebuildTray = rebuild;
  // keep a weak ref; main process keeps tray alive
}

function notifyCoreState() {
  if (mainWindow && !mainWindow.isDestroyed()) {
    mainWindow.webContents.send("core:state", {
      running: coreStarted() && coreReady,
      ready: coreReady,
      startedAt: coreStartedAt,
    });
  }
  try {
    if (global.__rebuildTray) global.__rebuildTray();
  } catch (e) {}
}

// ---------- auto start / launch at login ----------
function setupAutoStart() {
  if (settings.autoStart === undefined) settings.autoStart = true;
  applyAutoStart();
}
function applyAutoStart() {
  const on = !!settings.autoStart;
  if (process.platform === "linux") {
    // Linux: explicit XDG autostart desktop entry (app.setLoginItemSettings
    // is a no-op on Linux).
    const d = path.join(
      os.homedir(),
      ".config",
      "autostart",
      "modelrouter-gui.desktop",
    );
    try {
      if (on) {
        const exe = process.execPath;
        fs.mkdirSync(path.dirname(d), { recursive: true });
        fs.writeFileSync(
          d,
          '[Desktop Entry]\nType=Application\nName=ModelRouter\nComment=Unified LLM gateway\nExec="' +
            exe +
            '" --no-sandbox' +
            (settings.silentStart ? " --silent" : "") +
            "\nTerminal=false\nX-GNOME-Autostart-enabled=true\n",
          "utf-8",
        );
      } else if (fs.existsSync(d)) {
        fs.unlinkSync(d);
      }
    } catch (e) {
      console.warn("linux autostart failed:", e.message);
    }
    return;
  }
  try {
    app.setLoginItemSettings({ openAtLogin: on });
  } catch (e) {
    console.warn("auto-start not supported:", e.message);
  }
}

// ---------- window ----------
function createWindow() {
  const menu = Menu.buildFromTemplate([]);
  Menu.setApplicationMenu(menu);
  mainWindow = new BrowserWindow({
    // Window icon (taskbar/dock): same asset as the tray and package.
    icon: path.join(__dirname, "build", "icon.png"),
    // GNOME taskbar icons match the window's WM_CLASS (lowercase
    // "modelrouter-gui") against StartupWMClass in the .desktop file.
    backgroundColor: "#1f2937",
    width: 1280,
    height: 860,
    minWidth: 900,
    minHeight: 600,
    title: "ModelRouter",
    frame: false,
    show: false,
    webPreferences: {
      preload: path.join(__dirname, "preload.js"),
      contextIsolation: true,
      nodeIntegration: false,
    },
  });
  mainWindow.loadFile(path.join(__dirname, "renderer", "index.html"));
  mainWindow.webContents.setWindowOpenHandler(({ url }) => {
    if (url && url.startsWith("http")) shell.openExternal(url);
    return { action: "deny" };
  });
  // Silent start: `--silent` CLI flag or settings.silentStart keeps the app
  // in the tray without showing the window (Clash-Verge style). First run
  // always shows the window so the user sees the dashboard.
  mainWindow.once("ready-to-show", () => {
    const firstRun = !fs.existsSync(
      path.join(app.getPath("userData"), "settings.json"),
    );
    if (firstRun || !(silentFlag || settings.silentStart)) {
      mainWindow.show();
    }
  });
  mainWindow.on("close", (e) => {
    if (settings.minimizeToTray && !app.isQuitting) {
      e.preventDefault();
      mainWindow.hide();
      tray &&
        tray.displayBalloon &&
        tray.displayBalloon({
          title: "ModelRouter",
          content: "已最小化到托盘，内核继续运行",
        });
    }
  });
  mainWindow.on("closed", () => (mainWindow = null));
  if (isDev) mainWindow.webContents.openDevTools();
}

// ---------- IPC ----------
ipcMain.handle("window:minimize", (e) =>
  BrowserWindow.fromWebContents(e.sender)?.minimize(),
);
ipcMain.handle("window:toggleMaximize", (e) => {
  const w = BrowserWindow.fromWebContents(e.sender);
  if (!w) return;
  if (w.isMaximized()) w.unmaximize();
  else w.maximize();
});
ipcMain.handle("window:close", (e) =>
  BrowserWindow.fromWebContents(e.sender)?.close(),
);
ipcMain.handle(
  "window:isMaximized",
  (e) => !!BrowserWindow.fromWebContents(e.sender)?.isMaximized(),
);

ipcMain.handle("core:state", () => ({
  running: coreStarted() && coreReady,
  ready: coreReady,
  startedAt: coreStartedAt,
  baseUrl: embeddedBaseUrl(),
  port: settings.port || DEFAULT_PORT,
  profileDir: PROFILE_DIR,
  configFile: CONFIG_FILE,
  coreExe: CORE_EXE,
  coreExists: fs.existsSync(CORE_EXE),
}));
ipcMain.handle("core:start", () => {
  app.isQuitting = false;
  startCore();
  setTimeout(() => setEmbeddedAuth(), 400);
  return coreStarted();
});
ipcMain.handle("core:stop", () => {
  stopCore();
  return true;
});
ipcMain.handle("core:restart", () => {
  restartCore();
  setTimeout(() => setEmbeddedAuth(), 500);
  return true;
});
ipcMain.handle("core:probe", async () => {
  const p = await probeEmbedded();
  return p;
});
ipcMain.handle("core:key", () => readAdminKey());

ipcMain.handle("settings:get", () => loadSettings());
ipcMain.handle("settings:set", (_, patch) => {
  Object.assign(settings, patch);
  // port change -> rewrite ONLY the listen line, preserving sources/keys/etc.
  if (typeof patch.port === "number" && patch.port !== DEFAULT_PORT) {
    try {
      if (!fs.existsSync(CONFIG_FILE)) writeDefaultEmbeddedConfig();
      let raw = fs.readFileSync(CONFIG_FILE, "utf-8");
      raw = raw.replace(/^listen:.*$/m, "listen: 127.0.0.1:" + patch.port);
      fs.writeFileSync(CONFIG_FILE, raw, "utf-8");
    } catch (e) {
      console.error("port rewrite failed:", e);
    }
    if (coreStarted()) restartCore();
  }
  if (typeof patch.autoStart === "boolean") applyAutoStart();
  saveSettings();
  setTimeout(() => setEmbeddedAuth(), 300);
  return settings;
});
ipcMain.handle("settings:logFile", () => LOG_FILE);

// open paths in file manager
ipcMain.handle("shell:openPath", (_, p) => {
  shell.openPath(p);
  return true;
});

// window control from renderer uses preload -> ipcMain above

// ---------- bootstrap ----------
const gotLock = app.requestSingleInstanceLock();
if (gotLock) {
  app.on("second-instance", () => {
    if (mainWindow) {
      if (mainWindow.isMinimized()) mainWindow.restore();
      mainWindow.show();
      mainWindow.focus();
    }
  });

  app.whenReady().then(() => {
    installAuthRule();
    loadSettings();
    setupAutoStart();
    setEmbeddedAuth();
    createWindow();
    createTray();

    // start embedded core if not already running (another instance of llmsproxy)
    setTimeout(async () => {
      const p = await probeEmbedded();
      if (!p.ok) {
        startCore();
      }
    }, 600);
  });

  app.on("activate", () => {
    if (mainWindow === null) createWindow();
  });
  app.on("before-quit", () => {
    app.isQuitting = true;
    stopCore();
    try {
      if (tray) tray.destroy();
    } catch (e) {}
  });
} else {
  app.quit();
}