Restart ComfyUI automatically after plugin updates

web/js/o1keyApiSettings.js

@@ -1,10 +1,8 @@
 import { app } from "../../../scripts/app.js";
+import { startComfyUIRestart } from "./o1keyRestart.js";
 
 const DIALOG_ID = "o1key-api-settings-dialog";
 const STYLE_ID = "o1key-api-settings-styles";
-const RESTART_STATUS_URL = "/o1key/restart/status";
-const RESTART_TIMEOUT_MS = 120_000;
-let restartInProgress = false;
 
 function ensureStyles() {
     if (document.querySelector(`#${STYLE_ID}`)) return;
@@ -98,106 +96,11 @@ async function readJson(response) {
     return data;
 }
 
-function delay(milliseconds) {
-    return new Promise((resolve) => window.setTimeout(resolve, milliseconds));
-}
-
-async function fetchWithTimeout(url, options = {}, timeout = 2_500) {
-    const controller = new AbortController();
-    const timer = window.setTimeout(() => controller.abort(), timeout);
-    try {
-        return await fetch(url, {
-            ...options,
-            cache: "no-store",
-            signal: controller.signal,
-        });
-    } finally {
-        window.clearTimeout(timer);
-    }
-}
-
-async function getRestartStatus() {
-    const separator = RESTART_STATUS_URL.includes("?") ? "&" : "?";
-    const response = await fetchWithTimeout(
-        `${RESTART_STATUS_URL}${separator}_=${Date.now()}`
-    );
-    return readJson(response);
-}
-
-function setRestartButtonState(button, busy) {
-    button.disabled = busy;
-    button.setAttribute("aria-busy", String(busy));
-    button.setAttribute("aria-label", busy ? "ComfyUI 正在重启" : "重启");
-    button.setAttribute("title", busy ? "ComfyUI 正在重启" : "重启");
-
-    const icon = button.querySelector(".side-bar-button-icon");
-    if (icon) {
-        icon.className = busy
-            ? "pi pi-spinner pi-spin side-bar-button-icon"
-            : "pi pi-refresh side-bar-button-icon";
-    }
-
-    const label = button.querySelector(".side-bar-button-label");
-    if (label) label.textContent = busy ? "重启中" : "重启";
-}
-
-async function waitForRestart(oldBootId) {
-    const deadline = Date.now() + RESTART_TIMEOUT_MS;
-
-    while (Date.now() < deadline) {
-        try {
-            const status = await getRestartStatus();
-            if (status.ready && status.boot_id && status.boot_id !== oldBootId) {
-                const systemStatus = await fetchWithTimeout(
-                    `/api/system_stats?_=${Date.now()}`
-                );
-                if (systemStatus.ok) return;
-            }
-        } catch {
-            // 旧进程退出到新进程监听端口之间，接口暂时不可用是正常状态。
-        }
-        await delay(1_000);
-    }
-
-    throw new Error("等待 ComfyUI 重启超时，请查看终端中的错误信息。");
-}
-
-function reloadAfterRestart() {
-    const url = new URL(window.location.href);
-    url.searchParams.set("o1key_restart", Date.now().toString());
-    window.location.replace(url.toString());
-}
-
-async function restartComfyUI(event) {
-    if (restartInProgress) return;
+async function restartComfyUI() {
     if (!window.confirm("确定要重启 ComfyUI 吗？当前正在执行的任务会被中断。")) return;
-
-    const button = event.currentTarget;
-    restartInProgress = true;
-    setRestartButtonState(button, true);
-
     try {
-        let currentStatus = null;
-        try {
-            currentStatus = await getRestartStatus();
-        } catch {
-            // POST 响应还会返回旧进程的 boot_id，因此预检失败不阻断重启。
-        }
-
-        const response = await fetchWithTimeout("/o1key/restart", {
-            method: "POST",
-            headers: { "Content-Type": "application/json" },
-            body: "{}",
-        }, 5_000);
-        const result = await readJson(response);
-        const oldBootId = result.boot_id || currentStatus?.boot_id;
-        if (!oldBootId) throw new Error("无法确认当前 ComfyUI 进程标识。");
-
-        await waitForRestart(oldBootId);
-        reloadAfterRestart();
+        await startComfyUIRestart();
     } catch (error) {
-        restartInProgress = false;
-        setRestartButtonState(button, false);
         window.alert(error?.message || "ComfyUI 重启失败，请查看终端中的错误信息。");
     }
 }

web/js/o1keyRestart.js

@@ -0,0 +1,115 @@
+const RESTART_STATUS_URL = "/o1key/restart/status";
+const RESTART_TIMEOUT_MS = 120_000;
+let restartInProgress = false;
+
+async function readJson(response) {
+    const data = await response.json().catch(() => ({}));
+    if (!response.ok) throw new Error(data.error || `请求失败 (${response.status})`);
+    return data;
+}
+
+function delay(milliseconds) {
+    return new Promise((resolve) => window.setTimeout(resolve, milliseconds));
+}
+
+async function fetchWithTimeout(url, options = {}, timeout = 2_500) {
+    const controller = new AbortController();
+    const timer = window.setTimeout(() => controller.abort(), timeout);
+    try {
+        return await fetch(url, {
+            ...options,
+            cache: "no-store",
+            signal: controller.signal,
+        });
+    } finally {
+        window.clearTimeout(timer);
+    }
+}
+
+async function getRestartStatus() {
+    const separator = RESTART_STATUS_URL.includes("?") ? "&" : "?";
+    const response = await fetchWithTimeout(
+        `${RESTART_STATUS_URL}${separator}_=${Date.now()}`
+    );
+    return readJson(response);
+}
+
+function setRestartButtonState(button, busy) {
+    if (!button) return;
+    button.disabled = busy;
+    button.setAttribute("aria-busy", String(busy));
+    button.setAttribute("aria-label", busy ? "ComfyUI 正在重启" : "重启");
+    button.setAttribute("title", busy ? "ComfyUI 正在重启" : "重启");
+
+    const icon = button.querySelector(".side-bar-button-icon");
+    if (icon) {
+        icon.className = busy
+            ? "pi pi-spinner pi-spin side-bar-button-icon"
+            : "pi pi-refresh side-bar-button-icon";
+    }
+
+    const label = button.querySelector(".side-bar-button-label");
+    if (label) label.textContent = busy ? "重启中" : "重启";
+}
+
+async function waitForRestart(oldBootId) {
+    const deadline = Date.now() + RESTART_TIMEOUT_MS;
+
+    while (Date.now() < deadline) {
+        try {
+            const status = await getRestartStatus();
+            if (status.ready && status.boot_id && status.boot_id !== oldBootId) {
+                const systemStatus = await fetchWithTimeout(
+                    `/api/system_stats?_=${Date.now()}`
+                );
+                if (systemStatus.ok) return;
+            }
+        } catch {
+            // 旧进程退出到新进程监听端口之间，接口暂时不可用是正常状态。
+        }
+        await delay(1_000);
+    }
+
+    throw new Error("等待 ComfyUI 重启超时，请查看终端中的错误信息。");
+}
+
+function reloadAfterRestart() {
+    const url = new URL(window.location.href);
+    url.searchParams.set("o1key_restart", Date.now().toString());
+    window.location.replace(url.toString());
+}
+
+export async function startComfyUIRestart() {
+    if (restartInProgress) throw new Error("ComfyUI 正在重启，请稍候。");
+    restartInProgress = true;
+    let reloadRequested = false;
+    const button = document.querySelector("#o1key-restart-button");
+    setRestartButtonState(button, true);
+
+    try {
+        let currentStatus = null;
+        try {
+            currentStatus = await getRestartStatus();
+        } catch {
+            // POST 响应还会返回旧进程的 boot_id，因此预检失败不阻断重启。
+        }
+
+        const response = await fetchWithTimeout("/o1key/restart", {
+            method: "POST",
+            headers: { "Content-Type": "application/json" },
+            body: "{}",
+        }, 5_000);
+        const result = await readJson(response);
+        const oldBootId = result.boot_id || currentStatus?.boot_id;
+        if (!oldBootId) throw new Error("无法确认当前 ComfyUI 进程标识。");
+
+        await waitForRestart(oldBootId);
+        reloadAfterRestart();
+        reloadRequested = true;
+    } catch (error) {
+        setRestartButtonState(button, false);
+        throw error;
+    } finally {
+        if (!reloadRequested) restartInProgress = false;
+    }
+}

web/js/o1keyUpdateButton.js

@@ -1,13 +1,11 @@
 import { app } from "../../../scripts/app.js";
 import { api } from "../../../scripts/api.js";
+import { startComfyUIRestart } from "./o1keyRestart.js";
 
 const BUTTON_ID = "o1key-update-button";
-const DETECT_BUTTON_ID = "o1key-detect-button";
-const DETECT_NOTICE_ID = "o1key-detect-notice";
 const DIALOG_ID = "o1key-update-dialog";
 const STYLE_ID = "o1key-update-styles";
 let mountObserver = null;
-let detectNoticeTimer = null;
 
 const UPDATE_ERRORS = {
     local_changes: ["当前安装包含自定义修改，暂时无法自动更新。", "请联系技术支持处理，以免丢失现有内容。"],
@@ -207,26 +205,6 @@ function ensureStyles() {
         }
         #${DIALOG_ID} button:disabled { cursor: default; opacity: .55; }
         #${DIALOG_ID} button:disabled:hover { transform: none; }
-        #${DETECT_NOTICE_ID} {
-            position: fixed;
-            right: 20px;
-            bottom: 20px;
-            z-index: 100000;
-            display: flex;
-            align-items: center;
-            gap: 10px;
-            padding: 12px 16px;
-            border: 1px solid rgba(96, 165, 250, .45);
-            border-radius: 9px;
-            color: var(--p-primary-color, #60a5fa);
-            background: var(--comfy-menu-bg, #202020);
-            box-shadow: 0 12px 32px rgba(0, 0, 0, .32);
-            font-family: inherit;
-            font-size: 13px;
-            font-weight: 600;
-            line-height: 1.4;
-        }
-        #${DETECT_NOTICE_ID}[hidden] { display: none; }
         @media (max-width: 420px) {
             #${DIALOG_ID} .o1key-update-header { padding: 18px; }
             #${DIALOG_ID} .o1key-update-body { padding: 18px; }
@@ -328,15 +306,17 @@ function createUpdateDialog() {
 
     function setState(state) {
         dialog.dataset.state = state;
-        const busy = state === "checking" || state === "updating";
+        const busy = state === "checking" || state === "updating" || state === "restarting";
         primary.disabled = busy;
-        close.disabled = state === "updating";
-        secondary.disabled = state === "updating";
+        close.disabled = state === "updating" || state === "restarting";
+        secondary.disabled = state === "updating" || state === "restarting";
         primary.textContent = {
             checking: "检查中…",
             available: "立即更新",
             updating: "更新中…",
+            restarting: "重启中…",
             success: "完成",
+            restart_failed: "关闭",
         }[state] || "重新检查";
         secondary.textContent = state === "available" ? "稍后再说" : "关闭";
     }
@@ -414,14 +394,19 @@ function createUpdateDialog() {
                 return;
             }
             if (result.updated) {
+                if (!result.requirements_changed) {
+                    setState("restarting");
+                    show("更新完成，正在重启 ComfyUI…", "重启完成后页面会自动刷新，请稍候。", "busy");
+                    try {
+                        await startComfyUIRestart();
+                    } catch {
+                        setState("restart_failed");
+                        show("更新完成，但自动重启失败。", "请点击左侧“重启”按钮；如果问题持续，请联系技术支持。", "error");
+                    }
+                    return;
+                }
                 setState("success");
-                show(
-                    "更新完成。",
-                    result.requirements_changed
-                        ? "新版本还需要维护人员完成配置，请联系技术支持后重启 ComfyUI。"
-                        : "请重启 ComfyUI，并刷新页面以使用新版本。",
-                    "success",
-                );
+                show("更新完成。", "新版本还需要维护人员完成配置，请联系技术支持后重启 ComfyUI。", "success");
             } else {
                 setState("current");
                 show("已是最新版本。", "无需更新。", "success");
@@ -433,11 +418,11 @@ function createUpdateDialog() {
     }
     primary.addEventListener("click", () => {
         if (dialog.dataset.state === "available") return startUpdate();
-        if (dialog.dataset.state === "success") return dialog.close();
+        if (dialog.dataset.state === "success" || dialog.dataset.state === "restart_failed") return dialog.close();
         return checkForUpdate();
     });
     dialog.addEventListener("cancel", (event) => {
-        if (dialog.dataset.state === "updating") event.preventDefault();
+        if (dialog.dataset.state === "updating" || dialog.dataset.state === "restarting") event.preventDefault();
     });
     dialog.checkForUpdate = checkForUpdate;
 
@@ -458,29 +443,6 @@ function openUpdateDialog() {
     }
 }
 
-function showDetectNotice() {
-    ensureStyles();
-    let notice = document.querySelector(`#${DETECT_NOTICE_ID}`);
-    if (!notice) {
-        notice = document.createElement("div");
-        notice.id = DETECT_NOTICE_ID;
-        notice.setAttribute("role", "status");
-        const icon = document.createElement("span");
-        icon.className = "pi pi-check-circle";
-        icon.setAttribute("aria-hidden", "true");
-        const label = document.createElement("span");
-        label.textContent = "检测按钮已加载";
-        notice.append(icon, label);
-        document.body.appendChild(notice);
-    }
-    notice.hidden = false;
-    if (detectNoticeTimer !== null) window.clearTimeout(detectNoticeTimer);
-    detectNoticeTimer = window.setTimeout(() => {
-        notice.hidden = true;
-        detectNoticeTimer = null;
-    }, 3000);
-}
-
 function mountUpdateButton() {
     const tokenButton = document.querySelector("#o1key-token-manager-button");
     if (!tokenButton?.parentElement) return false;
@@ -506,26 +468,6 @@ function mountUpdateButton() {
 
     if (tokenButton.nextElementSibling !== updateButton) tokenButton.after(updateButton);
 
-    let detectButton = document.querySelector(`#${DETECT_BUTTON_ID}`);
-    if (!detectButton) {
-        detectButton = tokenButton.cloneNode(true);
-        detectButton.id = DETECT_BUTTON_ID;
-        detectButton.type = "button";
-        detectButton.classList.remove("side-bar-button-selected", "o1key-token-manager-button", "o1key-restart-button");
-        detectButton.classList.add("o1key-detect-button");
-        detectButton.setAttribute("aria-label", "检测");
-        detectButton.setAttribute("title", "检测按钮");
-        detectButton.removeAttribute("aria-pressed");
-
-        const icon = detectButton.querySelector(".side-bar-button-icon");
-        if (icon) icon.className = "pi pi-check-circle side-bar-button-icon";
-        const label = detectButton.querySelector(".side-bar-button-label");
-        if (label) label.textContent = "检测";
-
-        detectButton.addEventListener("click", showDetectNotice);
-    }
-
-    if (updateButton.nextElementSibling !== detectButton) updateButton.after(detectButton);
     return true;
 }