olivv-cs-ComfyUI-FunPack/app/ui/shell/maintenance.js

// Keeping the install alive: update, switch branch, roll back, restart ComfyUI. Every one ends in a restart and a page
// reload, so edits are saved first and a card holds the screen until the server answers again.
import { composer as c } from "../composer/composer.js";

const BRANCH_NOTES = { main: "Most stable available version of FunUI", dev: "Test new experimental features before they appear on main", "v4-classic": "Old version of FunUI with non-modular backend" };
const wait = (ms) => new Promise((r) => setTimeout(r, ms));

export function createMaintenance({ api, flush }) {
  const tell = (text) => c.toast.warn({ text });

  // -> git status when it is usable, else says why and returns null
  async function status(needClean) {
    let gs;
    try { gs = await api.gitFull(); } catch (err) { return tell(err.message), null; }
    if (!gs || !gs.ok) return tell((gs && gs.detail) || "Git is not available for this install."), null;
    if (needClean && gs.dirty) return tell("Local changes in the FunUI folder: commit or stash them first."), null;
    return gs;
  }
  const ask = (title, message) => c.modal.dialogue({ title, message, confirmLabel: "Go ahead" }).result;

  async function reloadWhenBack(card) {
    const start = Date.now();
    await wait(3500);
    for (;;) {
      try { if ((await api.health()).ok) return location.reload(); } catch { /* still down */ }
      if (Date.now() - start > 90000) card.setText("Still waiting on ComfyUI… it may need a manual restart.");
      await wait(2000);
    }
  }
  async function run(message, action, said) {
    try { await flush(); } catch (err) { return tell(err.message); }
    const text = c.text.md({ text: message });
    const card = c.modal.generic({ title: "Working", size: "sm", closeOnOutside: false, closeOnEsc: false, body: c.region.stack({ gap: "md", children: [c.progress.indeterminate({ label: "Working" }), text] }) });
    card.setText = (t) => text.setText ? text.setText(t) : (text.node.textContent = t);
    try {
      const res = await action(card);
      if (res && res.restarting === false) {       // nothing restarted: say what is true, and stay on the page
        card.close();
        return c.modal.dialogue({ title: "Not restarted", message: `${res.blocked || (said ? said(res) : "Nothing changed.")} ${res.blocked ? "Restart ComfyUI yourself when the generation has finished." : ""}`.trim(), confirmLabel: "OK", cancelLabel: "Close" }).result;
      }
      const bad = res && res.requirements && res.requirements.ran && !res.requirements.ok;
      if (bad) { await c.modal.dialogue({ title: "Dependencies failed", message: `${res.requirements.detail || ""}\n\nFunUI was updated, but installing its dependencies failed: it may not load until they are installed by hand. ComfyUI will restart now.`, confirmLabel: "Restart", cancelLabel: "Restart" }).result; }
      const line = said ? said(res) : message;
      card.setText(/Restarting ComfyUI/.test(line) ? line : `${line}\nRestarting ComfyUI…`);
    } catch (err) {
      if (action.dropsConnection) { /* the server went down as it restarted: expected */ } else { card.close(); return tell(err.message); }
    }
    return reloadWhenBack(card);
  }
  // Said after every update, whether or not anything was missing: the check is part of the update, so its result is too.
  const deps = (res) => {
    const q = res && res.requirements;
    if (!q) return "";
    if (!q.ran) return q.ok === false ? ` Requirements could not be checked: ${String(q.detail || "").split("\n")[0]}` : " Requirements: all present.";
    return q.ok ? " Requirements installed." : ` Requirement install FAILED: ${String(q.detail || "").split("\n").slice(0, 2).join(" ")}`;
  };

  return {
    /** What each action would do right now: {update, branch, rollback, ok} for the cards' hints. */
    async hints() {
      let gs; try { gs = await api.gitFull(); } catch { gs = null; }
      if (!gs || !gs.ok) return { ok: false, update: "Git unavailable for this install", branch: "Git unavailable for this install", rollback: "Git unavailable for this install" };
      return { ok: true,
        update: gs.fetch_ok === false ? "Could not reach origin: update status unknown" : gs.dirty ? "Local changes — commit first" : gs.behind > 0 ? `${gs.behind} commit(s) behind origin/${gs.branch}` : `origin/${gs.branch} up to date`,
        branch: `On ${gs.branch}${gs.dirty ? " · local changes — commit first" : " · pick another"}`,
        rollbackOk: Boolean(gs.rollback_target),
        rollback: gs.rollback_target ? `Undo last update — back to ${String(gs.rollback_target.commit).slice(0, 8)}` : "Nothing to undo yet" };
    },
    async update() {
      const gs = await status(false); if (!gs) return;
      const branch = gs.branch || "dev";
      // Local changes stop the code from moving, not the packages: the requirements check still runs, on its own.
      if (gs.dirty) {
        if (!(await ask("Check requirements", "FunUI has local changes, so the code cannot be updated. Check this install's requirements and install any that are missing?"))) return;
        return run("Checking requirements…", () => api.requirements(), (r) => deps({ requirements: r }));
      }
      if (!(await ask("Update", `Pull the latest "${branch}" from origin, check requirements and install any that are missing, then restart ComfyUI. If a generation is running, the restart waits until you press Restart after it finishes.${gs.behind > 0 ? ` ${gs.behind} commit(s) available.` : ""}`))) return;
      // Requirements first and on their own, so a failed pull (no network, a refused merge) does not skip them.
      const action = async () => {
        let requirements; try { requirements = await api.requirements(); } catch (err) { requirements = { ran: false, ok: false, detail: err.message }; }
        return { ...(await api.git("update", { branch })), requirements };
      };
      return run(`Pulling origin/${branch}…`, action, (r) => `${r.updated ? `Updated ${r.before} → ${r.after}.` : "Already up to date."}${deps(r)}`);
    },
    async switchBranch() {
      const gs = await status(true); if (!gs) return;
      if (!(gs.branches || []).length) return tell("No git branches found.");
      const branch = await c.modal.choice({ title: "Switch FunUI branch", items: gs.branches.map((b) => ({ id: b, label: b, hint: [b === gs.branch ? "current" : "", BRANCH_NOTES[b]].filter(Boolean).join(" · ") })) }).result;
      if (!branch || branch === gs.branch) return;
      if (!(await ask("Switch branch", `Switch to "${branch}", pull from origin, and restart ComfyUI. If a generation is running, the restart waits until you press Restart after it finishes.`))) return;
      return run(`Switching to ${branch}…`, () => api.git("checkout", { branch }), (r) => `On ${branch}${r.updated ? ` (${r.before} → ${r.after})` : ", already up to date"}.${deps(r)}`);
    },
    async rollback() {
      const gs = await status(true); if (!gs) return;
      const target = gs.rollback_target;
      if (!target) return tell("Nothing to roll back to — no update or branch switch on record.");
      if (!(await ask("Roll back", `Roll back to ${String(target.commit).slice(0, 8)}${target.subject ? ` ("${target.subject}")` : ""}? If that update changed requirements.txt, dependencies are not reinstalled: run pip by hand afterwards if things do not load.`))) return;
      return run(`Rolling back to ${String(target.commit).slice(0, 8)}…`, () => api.git("rollback"), (r) => `Rolled back ${r.before} → ${r.after}.`);
    },
    /** Switch torch to its CUDA 13 build (same versions), then restart. `why` is the warning that offered it. */
    async cuda13(why) {
      if (!(await ask("Switch torch to CUDA 13", `${why || ""}\n\nDownload the CUDA 13 build of the same torch, torchvision and torchaudio versions (several GB), install it, check it imports, and restart ComfyUI. If it fails, the old build is put back. This can take 10+ minutes.`.trim()))) return;
      // Started, then polled: the download outlives one request through a rental's proxy. A server that comes back
      // with no job and no result has restarted on its own after a success -- run() then waits for it and reloads.
      const action = async (card) => {
        await api.torchSwap();
        for (;;) {
          await wait(1000);
          let s; try { s = await api.torchSwapState(); } catch { continue; }         // down while restarting: keep asking
          if (s.step) card.setText(s.step);
          if (s.result) return s.result;
          if (!s.running) return { restarting: true };
        }
      };
      return run("Switching torch to CUDA 13…", action, (r) => r.message || "");
    },
    async restart() {
      let waiting = []; try { waiting = (await api.addons()).filter((a) => a.state === "pending").map((a) => a.title); } catch { /* not knowing is not a reason to refuse */ }
      const installs = waiting.length ? `\n\nOn the way up it will also download and install: ${waiting.join(", ")}. That can take a few minutes more.` : "";
      if (!(await ask("Restart ComfyUI", `Restart now? The server is down for 10–40 s and any running generation is lost. This page reloads when it is back.${installs}`))) return;
      const action = () => api.git("restart"); action.dropsConnection = true;
      return run("Restarting ComfyUI…", action);
    },
  };
}