diff --git a/README.md b/README.md index 09e2407..e4874fd 100644 --- a/README.md +++ b/README.md @@ -15,7 +15,8 @@ A ComfyUI custom node package that silently tracks which nodes, packages, and mo - **Smart aging** — items gradually move from "recently unused" to "safe to remove" over time - **Uninstall detection** — removed packages/models are flagged separately, historical data preserved - **Expandable detail** — click any package to see individual node-level stats -- **One-click disable** — disable unused packages straight from the dialog via ComfyUI Manager (per-package or in bulk), reversible at any time +- **One-click disable** — disable unused packages straight from the dialog (per-package or in bulk), reversible at any time. Uses ComfyUI Manager when it manages the pack, and **falls back to a native disable** (moving the pack into `custom_nodes/.disabled/`) for hand-cloned packs, loose single-file nodes, or when Manager isn't installed. Natively-disabled packs get an **Enable** button in the *Uninstalled* tier to move them back — no Manager needed (restart to apply) +- **Whitelist** — click the ☆ star on any package to protect it: whitelisted packs move into their own pinned group, never show a Disable button, and are skipped by the 7-day trial auto-disable - **Workflow tab** — on loading a workflow, splits unresolved nodes into *Missing* (install permanently or on a trial) and *Disabled* (enable permanently or on a trial), with a rolling **7-day trial** that auto-disables packages left unused - **Mirror search** — a standalone palette (⌕ button / `Ctrl/Cmd+Shift+D`) that searches nodes belonging to currently-disabled packages, draws an imitation node box (real inputs/widgets/outputs, parsed from source), and re-enables the pack on the spot - **Non-blocking** — DB writes happen in a background thread, no impact on workflow execution diff --git a/__init__.py b/__init__.py index 858edb1..968b203 100644 --- a/__init__.py +++ b/__init__.py @@ -7,6 +7,7 @@ from server import PromptServer from .mapper import NodePackageMapper, ModelMapper from .node_introspect import find_disabled_pack_path, get_node_schema +from .pack_fs import disable_pack_native, enable_pack_native, list_disabled_packs from .tracker import UsageTracker logger = logging.getLogger(__name__) @@ -144,6 +145,100 @@ async def get_node_schema_route(request): return web.json_response({"error": "internal error"}, status=500) +@routes.post("/nodes-stats/native-disable") +async def native_disable(request): + """Disable a pack by moving it into custom_nodes/.disabled/ (no Manager). + + Fallback for packs ComfyUI Manager doesn't manage. A restart is required for + ComfyUI to unload the pack. + """ + try: + data = await request.json() + package = data.get("package") + if not package: + return web.json_response({"error": "package required"}, status=400) + + ok, message = await asyncio.get_event_loop().run_in_executor( + None, disable_pack_native, package + ) + if not ok: + return web.json_response({"status": "error", "message": message}, status=409) + mapper.invalidate() + return web.json_response({"status": "ok", "message": message}) + except Exception: + logger.error("nodes-stats: error in native disable", exc_info=True) + return web.json_response({"error": "internal error"}, status=500) + + +@routes.post("/nodes-stats/native-enable") +async def native_enable(request): + """Re-enable a pack by moving it out of custom_nodes/.disabled/ (no Manager).""" + try: + data = await request.json() + package = data.get("package") + if not package: + return web.json_response({"error": "package required"}, status=400) + + ok, message = await asyncio.get_event_loop().run_in_executor( + None, enable_pack_native, package + ) + if not ok: + return web.json_response({"status": "error", "message": message}, status=409) + mapper.invalidate() + return web.json_response({"status": "ok", "message": message}) + except Exception: + logger.error("nodes-stats: error in native enable", exc_info=True) + return web.json_response({"error": "internal error"}, status=500) + + +@routes.get("/nodes-stats/disabled-packs") +async def get_disabled_packs(request): + """List packs present in custom_nodes/.disabled/ (re-enable candidates).""" + try: + names = await asyncio.get_event_loop().run_in_executor(None, list_disabled_packs) + return web.json_response(sorted(names)) + except Exception: + logger.error("nodes-stats: error listing disabled packs", exc_info=True) + return web.json_response({"error": "internal error"}, status=500) + + +@routes.get("/nodes-stats/whitelist") +async def get_whitelist(request): + try: + return web.json_response(sorted(tracker.get_whitelist())) + except Exception: + logger.error("nodes-stats: error getting whitelist", exc_info=True) + return web.json_response({"error": "internal error"}, status=500) + + +@routes.post("/nodes-stats/whitelist/add") +async def whitelist_add(request): + try: + data = await request.json() + package = data.get("package") + if not package: + return web.json_response({"error": "package required"}, status=400) + tracker.add_to_whitelist(package) + return web.json_response({"status": "ok"}) + except Exception: + logger.error("nodes-stats: error adding to whitelist", exc_info=True) + return web.json_response({"error": "internal error"}, status=500) + + +@routes.post("/nodes-stats/whitelist/remove") +async def whitelist_remove(request): + try: + data = await request.json() + package = data.get("package") + if not package: + return web.json_response({"error": "package required"}, status=400) + tracker.remove_from_whitelist(package) + return web.json_response({"status": "ok"}) + except Exception: + logger.error("nodes-stats: error removing from whitelist", exc_info=True) + return web.json_response({"error": "internal error"}, status=500) + + @routes.get("/nodes-stats/trials") async def get_trials(request): try: diff --git a/js/nodes_stats.js b/js/nodes_stats.js index 767221f..7b66d6d 100644 --- a/js/nodes_stats.js +++ b/js/nodes_stats.js @@ -119,6 +119,10 @@ function unresolvedNodeTypes() { // Latest workflow scan, shared so showStatsDialog can render the Workflow tab. let _lastWorkflowScan = { disabled: [], missing: [] }; +// Lowercased names of packs currently in custom_nodes/.disabled/ — the packs the +// extension can natively re-enable. Refreshed each time the dialog opens. +let _disabledPacksSet = new Set(); + async function onWorkflowLoaded() { const types = unresolvedNodeTypes(); _lastWorkflowScan = await classifyUnresolved(types); @@ -133,12 +137,14 @@ async function onWorkflowLoaded() { async function showStatsDialog(initialTab = "nodes") { let data, modelData, managerInfo, trials = []; try { - const [pkgResp, modelResp, mgr, trialsResp] = await Promise.all([ + const [pkgResp, modelResp, mgr, trialsResp, disabledPacks] = await Promise.all([ fetch("/nodes-stats/packages"), fetch("/nodes-stats/models"), fetchManagerInfo(), fetch("/nodes-stats/trials").catch(() => null), + fetchDisabledPacks(), ]); + _disabledPacksSet = disabledPacks; if (!pkgResp.ok) { alert("Failed to load node stats: HTTP " + pkgResp.status); return; } if (!modelResp.ok) { alert("Failed to load model stats: HTTP " + modelResp.status); return; } data = await pkgResp.json(); @@ -248,6 +254,8 @@ async function showStatsDialog(initialTab = "nodes") { wireDisableButtons(dialog, managerInfo); wireWorkflowButtons(dialog); + wireWhitelistButtons(dialog); + wireNativeEnableButtons(dialog); switchTab(TABS.includes(initialTab) ? initialTab : "nodes"); @@ -285,6 +293,7 @@ function dialogStyle() { #nodes-stats-dialog .ns-btn:hover:not(:disabled){background:#3a2020;border-color:#e44;color:#fff;} #nodes-stats-dialog .ns-btn:disabled{opacity:0.5;cursor:default;} #nodes-stats-dialog .ns-disable-all-btn{border-color:#a33;color:#e88;} + #nodes-stats-dialog .ns-enable-native-btn:hover:not(:disabled){background:#203a20;border-color:#4a4;color:#fff;} ${rows} `; } @@ -306,7 +315,11 @@ function summaryBar(items) { } function buildNodesTabContent(custom, managerInfo) { - const byStatus = (s) => custom.filter((p) => p.status === s); + // Whitelisted packs are pulled out of their usage tiers into a pinned group + // and never offered for disable. Everything else classifies as usual. + const whitelisted = custom.filter((p) => p.whitelisted); + const rest = custom.filter((p) => !p.whitelisted); + const byStatus = (s) => rest.filter((p) => p.status === s); const safeToRemove = byStatus("safe_to_remove"); const considerRemoving = byStatus("consider_removing"); const unusedNew = byStatus("unused_new"); @@ -320,6 +333,7 @@ function buildNodesTabContent(custom, managerInfo) { { count: used.length, status: "used", label: "used", id: "nodes-stats-used-badge" }, ]); + html += renderWhitelistSection(whitelisted, managerInfo); html += renderSection("Safe to Remove", "Unused for 2+ months", "safe_to_remove", safeToRemove, managerInfo); html += renderSection("Consider Removing", "Unused for 1-2 months", "consider_removing", considerRemoving, managerInfo); html += renderSection("Recently Unused", "Unused for less than 1 month", "unused_new", unusedNew, managerInfo); @@ -333,8 +347,12 @@ function renderSection(title, subtitle, status, packages, managerInfo) { if (packages.length === 0) return ""; const color = STATUS_META[status].color; - const withActions = !!managerInfo && DISABLEABLE_TIERS.has(status); - const eligible = withActions + // Disable actions no longer require Manager: an installed pack Manager doesn't + // cover is disabled natively (moved into custom_nodes/.disabled by us). + // The Uninstalled tier gets an Enable action for packs recoverable from + // .disabled/ (also native, so it works with or without Manager). + const withActions = DISABLEABLE_TIERS.has(status) || sectionHasEnable(status, packages); + const eligible = DISABLEABLE_TIERS.has(status) ? packages.filter((p) => isDisableEligible(p, managerInfo)).map((p) => p.package) : []; @@ -352,12 +370,48 @@ function renderSection(title, subtitle, status, packages, managerInfo) { return html; } -// A package can be disabled only if ComfyUI Manager knows it (by directory -// name) and it is currently active (any state other than already-disabled). +// How a package can be disabled: "manager" when ComfyUI Manager knows it (by +// directory name) and it's currently active; "native" when it's installed but +// Manager doesn't cover it (hand-cloned repo, loose file, or Manager absent) — +// the extension moves it into custom_nodes/.disabled itself; null when it can't +// be disabled (uninstalled, or whitelisted/protected). +function disableMode(pkg, managerInfo) { + if (pkg.whitelisted || !pkg.installed) return null; + const info = managerInfo && managerInfo[pkg.package]; + if (info) { + // Manager knows this pack: offer disable only while it's still active. + return info.state && info.state !== "disabled" ? "manager" : null; + } + // Manager doesn't cover it (hand-cloned, loose file, or Manager absent): + // the extension moves it into custom_nodes/.disabled itself. + return "native"; +} + function isDisableEligible(pkg, managerInfo) { - if (!managerInfo || !pkg.installed) return false; - const info = managerInfo[pkg.package]; - return !!(info && info.state && info.state !== "disabled"); + return disableMode(pkg, managerInfo) !== null; +} + +// An uninstalled pack can be natively re-enabled iff its source is still sitting +// in custom_nodes/.disabled/. (Whitelisted packs live in their own group.) +function nativeEnableEligible(pkg) { + return pkg.status === "uninstalled" && _disabledPacksSet.has(pkg.package.toLowerCase()); +} + +function sectionHasEnable(status, packages) { + return status === "uninstalled" && packages.some(nativeEnableEligible); +} + +const WHITELIST_COLOR = "#e5c04b"; + +// Pinned group of protected packs. They keep their own row background (mixed +// statuses) but never show a Disable action — the star toggles protection off. +function renderWhitelistSection(packages, managerInfo) { + if (packages.length === 0) return ""; + let html = `
+

★ Whitelist — protected, never disabled

+
`; + html += buildTable(packages, null, false, managerInfo); + return html; } function buildModelsTabContent(modelData) { @@ -462,6 +516,15 @@ function sectionHeader(title, subtitle, color) { return html; } +// Clickable whitelist toggle rendered before a package name. Filled gold star +// when protected, hollow grey when not. Wired by wireWhitelistButtons. +function whitelistStar(pkg) { + const on = !!pkg.whitelisted; + return `${on ? "★" : "☆"}`; +} + function buildTable(packages, status, withActions, managerInfo) { const colspan = withActions ? 7 : 6; @@ -480,19 +543,23 @@ function buildTable(packages, status, withActions, managerInfo) { const hasNodes = pkg.nodes && pkg.nodes.length > 0; const lastSeen = pkg.last_seen ? new Date(pkg.last_seen).toLocaleDateString() : "—"; - html += ` + html += ` ${hasNodes ? "▶" : " "} - ${escapeHtml(pkg.package)} + ${whitelistStar(pkg)}${escapeHtml(pkg.package)} ${pkg.total_nodes} ${pkg.used_nodes}/${pkg.total_nodes} ${pkg.total_executions} ${lastSeen}`; if (withActions) { - const eligible = isDisableEligible(pkg, managerInfo); - const cell = eligible - ? `` - : ``; + let cell; + if (nativeEnableEligible(pkg)) { + cell = ``; + } else if (isDisableEligible(pkg, managerInfo)) { + cell = ``; + } else { + cell = ``; + } html += `${cell}`; } html += ``; @@ -799,8 +866,8 @@ function disablePayload(dirName, info) { } function wireDisableButtons(dialog, managerInfo) { - if (!managerInfo) return; - + // No early return on a missing Manager: native-disable buttons are rendered + // for installed packs Manager doesn't cover, and must work with no Manager. dialog.querySelectorAll(".ns-disable-btn").forEach((btn) => { btn.addEventListener("click", (e) => { e.stopPropagation(); @@ -831,48 +898,156 @@ function wireWorkflowButtons(dialog) { b.addEventListener("click", (e) => { e.stopPropagation(); handleTrialInstall(b.dataset.pkg, dialog, false); })); } -async function handleDisable(pkgNames, dialog, managerInfo) { - // Only act on packages Manager still reports as active (guards against - // double-clicks and stale buttons after a partial batch). - pkgNames = pkgNames.filter((n) => managerInfo[n] && managerInfo[n].state !== "disabled"); - if (pkgNames.length === 0) return; +// Star toggles on package rows: add/remove the pack from the whitelist, then +// re-render the Nodes tab so the pack moves in/out of the pinned group. +function wireWhitelistButtons(dialog) { + dialog.querySelectorAll(".ns-wl-btn").forEach((btn) => { + btn.addEventListener("click", (e) => { + e.stopPropagation(); + toggleWhitelist(btn.dataset.pkg, btn.dataset.on !== "1"); + }); + }); +} - const what = pkgNames.length === 1 ? `"${pkgNames[0]}"` : `${pkgNames.length} packages`; +async function toggleWhitelist(pkg, add) { + try { + const r = await fetch(add ? "/nodes-stats/whitelist/add" : "/nodes-stats/whitelist/remove", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ package: pkg }), + }); + if (!r.ok) throw new Error("HTTP " + r.status); + } catch (e) { + notify("Whitelist update failed: " + e.message, "error"); + return; + } + showStatsDialog("nodes"); // re-fetch + rebuild so the pack regroups +} + +async function fetchWhitelist() { + try { + const r = await fetch("/nodes-stats/whitelist"); + if (r.ok) return new Set(await r.json()); + } catch { /* treat as empty */ } + return new Set(); +} + +// Enable buttons on recoverable uninstalled packs: move the pack out of +// custom_nodes/.disabled via the native route. Restart is left to the user. +function wireNativeEnableButtons(dialog) { + dialog.querySelectorAll(".ns-enable-native-btn").forEach((btn) => { + btn.addEventListener("click", (e) => { + e.stopPropagation(); + handleNativeEnable(btn.dataset.pkg, btn, dialog); + }); + }); +} + +async function handleNativeEnable(pkg, btn, dialog) { + if (!confirm( + `Enable "${pkg}"?\n\nIt will be moved out of custom_nodes/.disabled and ` + + `load on the next ComfyUI restart.` + )) return; + btn.disabled = true; + try { + const r = await fetch("/nodes-stats/native-enable", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ package: pkg }), + }); + if (!r.ok) { + let msg = "HTTP " + r.status; + try { const j = await r.json(); if (j && j.message) msg = j.message; } catch { /* keep status */ } + throw new Error(msg); + } + _disabledPacksSet.delete(pkg.toLowerCase()); + const cell = dialog.querySelector(`.ns-action-cell[data-pkg="${cssEscape(pkg)}"]`); + if (cell) cell.innerHTML = `✓ enabled · restart`; + showRestartBanner(dialog); + notify(`Enabled ${pkg}. Restart ComfyUI to load it.`, "success"); + } catch (e) { + btn.disabled = false; + notify("Failed to enable: " + e.message, "error"); + } +} + +// Lowercased set of pack names in custom_nodes/.disabled/ (native re-enable +// candidates). Lowercased so lookups match package names case-insensitively. +async function fetchDisabledPacks() { + try { + const r = await fetch("/nodes-stats/disabled-packs"); + if (r.ok) return new Set((await r.json()).map((n) => n.toLowerCase())); + } catch { /* treat as empty */ } + return new Set(); +} + +async function handleDisable(pkgNames, dialog, managerInfo) { + // Partition into packages Manager can disable (still reported active) vs. the + // rest, which we disable natively by moving them into custom_nodes/.disabled. + // The eligible lists already excluded whitelisted/uninstalled packs. + const managerNames = pkgNames.filter( + (n) => managerInfo && managerInfo[n] && managerInfo[n].state && managerInfo[n].state !== "disabled" + ); + const nativeNames = pkgNames.filter((n) => !managerNames.includes(n)); + if (managerNames.length === 0 && nativeNames.length === 0) return; + + const total = managerNames.length + nativeNames.length; + const what = total === 1 ? `"${managerNames[0] || nativeNames[0]}"` : `${total} packages`; + const via = managerNames.length && nativeNames.length + ? "ComfyUI Manager / directly" + : managerNames.length ? "ComfyUI Manager" : "the extension"; const confirmMsg = - `Disable ${what} via ComfyUI Manager?\n\n` + + `Disable ${what} via ${via}?\n\n` + `They will be moved to custom_nodes/.disabled and a ComfyUI restart is ` + - `required to take effect. You can re-enable them anytime from ComfyUI Manager.`; + `required to take effect. You can re-enable them anytime.`; if (!confirm(confirmMsg)) return; setDisableButtonsBusy(dialog, true); + const succeeded = []; + const failed = []; try { - const pre = await mgrFetch("/manager/queue/status").then((r) => (r && r.ok ? r.json() : null)).catch(() => null); - if (pre && pre.is_processing) { - notify("ComfyUI Manager is busy. Please try again in a moment.", "warn"); - setDisableButtonsBusy(dialog, false); - return; + if (managerNames.length) { + const pre = await mgrFetch("/manager/queue/status").then((r) => (r && r.ok ? r.json() : null)).catch(() => null); + if (pre && pre.is_processing) { + notify("ComfyUI Manager is busy. Please try again in a moment.", "warn"); + failed.push(...managerNames); + } else { + try { + await runManagerDisable(managerNames.map((n) => disablePayload(n, managerInfo[n]))); + // Reconcile against Manager's actual state: disabled only if no longer + // reported as active on disk. + const after = await fetchManagerInfo(); + const isStillActive = (n) => after && after[n] && after[n].state !== "disabled"; + for (const n of managerNames) { + if (after ? !isStillActive(n) : true) { + succeeded.push(n); + if (managerInfo[n]) managerInfo[n].state = "disabled"; + } else { + failed.push(n); + } + } + } catch (e) { + failed.push(...managerNames); + notify("ComfyUI Manager disable failed: " + e.message, "error"); + } + } } - const payloads = pkgNames.map((n) => disablePayload(n, managerInfo[n])); - await runManagerDisable(payloads); + if (nativeNames.length) { + const res = await runNativeDisable(nativeNames); + succeeded.push(...res.succeeded); + failed.push(...res.failed); + } - // Reconcile against Manager's actual state: a package is considered - // disabled only if it's no longer reported as active on disk. - const after = await fetchManagerInfo(); - const isStillActive = (n) => after && after[n] && after[n].state !== "disabled"; - const succeeded = after ? pkgNames.filter((n) => !isStillActive(n)) : pkgNames; - const failed = pkgNames.filter((n) => !succeeded.includes(n)); - - succeeded.forEach((n) => { if (managerInfo[n]) managerInfo[n].state = "disabled"; }); markPackagesDisabled(dialog, succeeded); - updateBulkButtons(dialog, managerInfo); + updateBulkButtons(dialog, new Set(succeeded)); if (succeeded.length > 0) { showRestartBanner(dialog); notify(`Disabled ${succeeded.length} package${succeeded.length !== 1 ? "s" : ""}. Restart ComfyUI to apply.`, "success"); } if (failed.length > 0) { - notify(`ComfyUI Manager could not disable: ${failed.join(", ")}`, "error"); + notify(`Could not disable: ${failed.join(", ")}`, "error"); } } catch (e) { notify("Failed to disable: " + e.message, "error"); @@ -881,6 +1056,26 @@ async function handleDisable(pkgNames, dialog, managerInfo) { } } +// Disable packs by moving each into custom_nodes/.disabled via the extension's +// own route (no Manager). Returns { succeeded, failed } name lists. +async function runNativeDisable(pkgNames) { + const succeeded = []; + const failed = []; + for (const n of pkgNames) { + try { + const r = await fetch("/nodes-stats/native-disable", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ package: n }), + }); + (r.ok ? succeeded : failed).push(n); + } catch { + failed.push(n); + } + } + return { succeeded, failed }; +} + // v4-legacy: queue a batch of operations ({install:[...]}, {disable:[...]}, ...) // through the single endpoint that also starts the worker, then wait for it to // finish. A reset first clears any stale queue (harmless if empty). @@ -1391,26 +1586,39 @@ async function processExpiredTrials() { if (r.ok) trials = await r.json(); } catch { return; } - const expired = trials.filter((t) => t.expired); + let expired = trials.filter((t) => t.expired); + if (!expired.length) return; + + // Whitelisted packs are protected: never auto-disable them. Clear their trial + // row so they stop being tracked as temporary (they're now keepers). + const whitelist = await fetchWhitelist(); + const protectedExpired = expired.filter((t) => whitelist.has(t.package)); + for (const t of protectedExpired) await stopTrial(t.package); + expired = expired.filter((t) => !whitelist.has(t.package)); if (!expired.length) return; const mgr = await fetchManagerInfo(); - if (!mgr) return; // Manager unavailable — leave rows for a later session // Don't reset Manager's queue out from under an in-progress operation // (e.g. startup install work); the expired rows persist and retry next session. - if (await managerIsBusy()) return; + if (mgr && await managerIsBusy()) return; const done = []; for (const t of expired) { - const info = mgr[t.package]; - if (!info || info.state === "disabled") { + const info = mgr && mgr[t.package]; + if (info && info.state === "disabled") { await stopTrial(t.package); done.push(t.package); continue; } try { - await runManagerDisable([disablePayload(t.package, info)]); + if (info) { + await runManagerDisable([disablePayload(t.package, info)]); + } else { + // Manager absent or doesn't know this pack — disable it natively. + const res = await runNativeDisable([t.package]); + if (!res.succeeded.includes(t.package)) throw new Error("native disable failed"); + } await stopTrial(t.package); done.push(t.package); } catch { /* keep the row; retry next session */ } @@ -1454,13 +1662,13 @@ function markPackagesDisabled(dialog, pkgNames) { } } -// Recompute "Disable all (N)" counts after a batch; hide buttons with nothing -// left to disable. -function updateBulkButtons(dialog, managerInfo) { +// Recompute "Disable all (N)" counts after a batch, dropping packages that were +// just disabled; hide buttons with nothing left. Mode-agnostic (Manager+native). +function updateBulkButtons(dialog, succeededSet) { dialog.querySelectorAll(".ns-disable-all-btn").forEach((btn) => { let names = []; try { names = JSON.parse(btn.dataset.pkgs); } catch { names = []; } - const remaining = names.filter((n) => managerInfo[n] && managerInfo[n].state !== "disabled"); + const remaining = names.filter((n) => !succeededSet.has(n)); if (remaining.length === 0) { btn.style.display = "none"; } else { @@ -1492,6 +1700,14 @@ function showRestartBanner(dialog) { } async function rebootComfy() { + // The reboot endpoint is provided by ComfyUI Manager. Without it we can't + // restart the server ourselves — tell the user to do it manually (this is the + // native-only disable path). + const api = await detectManagerApi(); + if (!api) { + notify("No ComfyUI Manager to trigger a restart — please restart ComfyUI manually to apply the changes.", "warn"); + return; + } if (!confirm("Restart ComfyUI now? The server will go down briefly and the page will reconnect.")) return; notify("Restarting ComfyUI…", "info"); try { diff --git a/pack_fs.py b/pack_fs.py new file mode 100644 index 0000000..d8a7212 --- /dev/null +++ b/pack_fs.py @@ -0,0 +1,138 @@ +"""Native enable/disable of custom-node packages, independent of ComfyUI Manager. + +Disabling a pack is just a filesystem convention that ComfyUI core honors: a +directory (or ``*.py`` file) moved into ``custom_nodes/.disabled/`` is skipped at +boot, because core ignores any entry named ``.disabled`` / starting with a dot. +Manager uses the same convention. We only *move* files here — never import, +delete, or re-clone — so the operation is reversible and safe. + +This is the fallback the extension uses when ComfyUI Manager is absent or does +not recognize a pack (a hand-cloned repo, an oddly-installed pack, or a loose +single-file node). A restart is required for ComfyUI to pick up the change, +exactly as with a Manager disable. +""" + +import logging +import os +import shutil + +logger = logging.getLogger(__name__) + +# find_disabled_pack_path already knows how to locate a pack under .disabled/, +# including Manager's @version suffix. Import it in a way that works both as a +# package (runtime) and as a top-level module (tests add the root to sys.path). +try: + from .node_introspect import find_disabled_pack_path +except ImportError: # pragma: no cover - exercised via top-level import in tests + from node_introspect import find_disabled_pack_path + + +def _custom_node_roots(): + import folder_paths + + return [os.path.normpath(r) for r in folder_paths.get_folder_paths("custom_nodes")] + + +def _valid_name(pack_name): + """Reject path-y input; pack names are plain directory/file names.""" + return bool(pack_name) and not any(c in pack_name for c in ("/", "\\")) and ".." not in pack_name + + +def list_disabled_packs(): + """Return clean names of packs sitting in any custom_nodes/.disabled/ dir. + + Strips a Manager ``@version`` suffix and a ``.py`` extension so the names + match package names as reported in the stats (and accepted by + ``enable_pack_native``). + """ + names = set() + for root in _custom_node_roots(): + ddir = os.path.join(root, ".disabled") + try: + entries = os.listdir(ddir) + except Exception: + continue + for e in entries: + if e.startswith("."): + continue + base = e.split("@", 1)[0] + stem = base[:-3] if base.endswith(".py") else base + if stem: + names.add(stem) + return names + + +def find_active_pack_path(pack_name): + """Locate an installed (active) pack directly under a custom_nodes root. + + Returns an absolute path to the pack directory or single ``.py`` file, or + None. Matches case-insensitively and tolerates a Manager ``@version`` suffix. + Only real packs (a directory, or a ``.py`` file) qualify. + """ + if not _valid_name(pack_name): + return None + target = pack_name.lower() + for root in _custom_node_roots(): + try: + entries = os.listdir(root) + except Exception: + continue + for e in entries: + if e.startswith("."): + continue + base = e.split("@", 1)[0] + stem = base[:-3] if base.endswith(".py") else base + if target not in (stem.lower(), base.lower(), e.lower()): + continue + full = os.path.join(root, e) + if os.path.isdir(full) or (os.path.isfile(full) and full.endswith(".py")): + return os.path.normpath(full) + return None + + +def disable_pack_native(pack_name): + """Move an active pack into ``custom_nodes/.disabled/``. + + Returns ``(ok, message)``. Never raises for expected conditions (pack not + found, collision, permission error); those come back as ``(False, reason)``. + """ + src = find_active_pack_path(pack_name) + if not src: + return False, "pack not found on disk" + root = os.path.dirname(src) + ddir = os.path.join(root, ".disabled") + dest = os.path.join(ddir, os.path.basename(src)) + if os.path.exists(dest): + return False, "a disabled copy already exists at " + dest + try: + os.makedirs(ddir, exist_ok=True) + shutil.move(src, dest) + except Exception as e: + logger.warning("nodes-stats: native disable failed for %s", pack_name, exc_info=True) + return False, str(e) + return True, "disabled" + + +def enable_pack_native(pack_name): + """Move a pack from ``custom_nodes/.disabled/`` back to its root. + + Returns ``(ok, message)``. Drops any Manager ``@version`` suffix from the + destination so the pack lands as a clean, importable directory name — the + same shape Manager itself restores on enable (``ComfyMath@nightly`` -> + ``ComfyMath``). A single ``.py`` file keeps its extension. + """ + src = find_disabled_pack_path(pack_name) + if not src: + return False, "disabled pack not found" + ddir = os.path.dirname(src) # .../custom_nodes/.disabled + root = os.path.dirname(ddir) # .../custom_nodes + dest_name = os.path.basename(src).split("@", 1)[0] + dest = os.path.join(root, dest_name) + if os.path.exists(dest): + return False, "an active copy already exists at " + dest + try: + shutil.move(src, dest) + except Exception as e: + logger.warning("nodes-stats: native enable failed for %s", pack_name, exc_info=True) + return False, str(e) + return True, "enabled" diff --git a/pyproject.toml b/pyproject.toml index b0673f4..55ad0b6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-nodes-stats" description = "Track usage statistics for all ComfyUI nodes and packages" -version = "1.7.2" +version = "1.8.0" license = "MIT" [project.urls] diff --git a/tests/test_pack_fs.py b/tests/test_pack_fs.py new file mode 100644 index 0000000..a77c900 --- /dev/null +++ b/tests/test_pack_fs.py @@ -0,0 +1,150 @@ +import os + +import pytest + +import pack_fs + + +@pytest.fixture +def custom_nodes(tmp_path): + """A custom_nodes root wired into the mocked folder_paths.""" + import folder_paths + + root = tmp_path / "custom_nodes" + root.mkdir() + folder_paths.get_folder_paths.return_value = [str(root)] + return root + + +def _make_pack(root, name): + pack = root / name + pack.mkdir() + (pack / "__init__.py").write_text("NODE_CLASS_MAPPINGS = {}\n") + return pack + + +# --- find_active_pack_path ------------------------------------------------ + +def test_find_active_dir(custom_nodes): + _make_pack(custom_nodes, "MyPack") + found = pack_fs.find_active_pack_path("MyPack") + assert found == os.path.normpath(str(custom_nodes / "MyPack")) + + +def test_find_active_case_insensitive(custom_nodes): + _make_pack(custom_nodes, "MyPack") + assert pack_fs.find_active_pack_path("mypack") is not None + + +def test_find_active_single_py_file(custom_nodes): + (custom_nodes / "loose_node.py").write_text("NODE_CLASS_MAPPINGS = {}\n") + found = pack_fs.find_active_pack_path("loose_node") + assert found == os.path.normpath(str(custom_nodes / "loose_node.py")) + + +def test_find_active_ignores_disabled_dir(custom_nodes): + (custom_nodes / ".disabled").mkdir() + _make_pack(custom_nodes / ".disabled", "MyPack") + assert pack_fs.find_active_pack_path("MyPack") is None + + +def test_find_active_rejects_path_traversal(custom_nodes): + assert pack_fs.find_active_pack_path("../evil") is None + assert pack_fs.find_active_pack_path("a/b") is None + assert pack_fs.find_active_pack_path("") is None + + +# --- disable_pack_native -------------------------------------------------- + +def test_disable_moves_dir_into_disabled(custom_nodes): + _make_pack(custom_nodes, "MyPack") + ok, msg = pack_fs.disable_pack_native("MyPack") + assert ok, msg + assert not (custom_nodes / "MyPack").exists() + assert (custom_nodes / ".disabled" / "MyPack" / "__init__.py").exists() + + +def test_disable_moves_single_py_file(custom_nodes): + (custom_nodes / "loose_node.py").write_text("x = 1\n") + ok, msg = pack_fs.disable_pack_native("loose_node") + assert ok, msg + assert not (custom_nodes / "loose_node.py").exists() + assert (custom_nodes / ".disabled" / "loose_node.py").exists() + + +def test_disable_missing_pack_fails(custom_nodes): + ok, msg = pack_fs.disable_pack_native("Ghost") + assert not ok + assert "not found" in msg + + +def test_disable_collision_fails(custom_nodes): + _make_pack(custom_nodes, "MyPack") + (custom_nodes / ".disabled").mkdir() + (custom_nodes / ".disabled" / "MyPack").mkdir() # pre-existing disabled copy + ok, msg = pack_fs.disable_pack_native("MyPack") + assert not ok + assert "already exists" in msg + # original left untouched + assert (custom_nodes / "MyPack").exists() + + +# --- enable_pack_native --------------------------------------------------- + +def test_enable_moves_back(custom_nodes): + _make_pack(custom_nodes, "MyPack") + assert pack_fs.disable_pack_native("MyPack")[0] + ok, msg = pack_fs.enable_pack_native("MyPack") + assert ok, msg + assert (custom_nodes / "MyPack" / "__init__.py").exists() + assert not (custom_nodes / ".disabled" / "MyPack").exists() + + +def test_enable_missing_fails(custom_nodes): + ok, msg = pack_fs.enable_pack_native("Ghost") + assert not ok + assert "not found" in msg + + +def test_enable_strips_version_suffix(custom_nodes): + # A Manager-disabled pack on disk carries an @version suffix; enabling should + # restore it as a clean, importable directory name. + ddir = custom_nodes / ".disabled" + ddir.mkdir() + pack = ddir / "ComfyMath@nightly" + pack.mkdir() + (pack / "__init__.py").write_text("NODE_CLASS_MAPPINGS = {}\n") + + ok, msg = pack_fs.enable_pack_native("ComfyMath") + assert ok, msg + assert (custom_nodes / "ComfyMath" / "__init__.py").exists() + assert not (custom_nodes / "ComfyMath@nightly").exists() + assert not (ddir / "ComfyMath@nightly").exists() + + +def test_disable_then_enable_roundtrip_py_file(custom_nodes): + (custom_nodes / "loose_node.py").write_text("x = 1\n") + assert pack_fs.disable_pack_native("loose_node")[0] + assert pack_fs.enable_pack_native("loose_node")[0] + assert (custom_nodes / "loose_node.py").exists() + + +# --- list_disabled_packs -------------------------------------------------- + +def test_list_disabled_empty(custom_nodes): + assert pack_fs.list_disabled_packs() == set() + + +def test_list_disabled_strips_version_and_ext(custom_nodes): + ddir = custom_nodes / ".disabled" + ddir.mkdir() + (ddir / "ComfyMath@nightly").mkdir() # Manager-style @version suffix + (ddir / "PlainPack").mkdir() + (ddir / "loose_node.py").write_text("x = 1\n") # single-file pack + assert pack_fs.list_disabled_packs() == {"ComfyMath", "PlainPack", "loose_node"} + + +def test_list_disabled_reflects_a_native_disable(custom_nodes): + _make_pack(custom_nodes, "MyPack") + assert pack_fs.disable_pack_native("MyPack")[0] + assert "MyPack" in pack_fs.list_disabled_packs() diff --git a/tests/test_whitelist.py b/tests/test_whitelist.py new file mode 100644 index 0000000..2aa1834 --- /dev/null +++ b/tests/test_whitelist.py @@ -0,0 +1,72 @@ +import pytest + +from tracker import UsageTracker + + +@pytest.fixture +def tracker(tmp_path): + return UsageTracker(db_path=str(tmp_path / "test.db")) + + +def test_whitelist_starts_empty(tracker): + assert tracker.get_whitelist() == set() + + +def test_add_and_get(tracker): + tracker.add_to_whitelist("My-Pack") + assert tracker.get_whitelist() == {"My-Pack"} + + +def test_add_is_idempotent(tracker): + tracker.add_to_whitelist("My-Pack") + tracker.add_to_whitelist("My-Pack") + assert tracker.get_whitelist() == {"My-Pack"} + + +def test_remove(tracker): + tracker.add_to_whitelist("My-Pack") + tracker.remove_from_whitelist("My-Pack") + assert tracker.get_whitelist() == set() + + +def test_remove_absent_is_noop(tracker): + tracker.remove_from_whitelist("Nope") # must not raise + assert tracker.get_whitelist() == set() + + +def test_reset_clears_whitelist(tracker): + tracker.add_to_whitelist("My-Pack") + tracker.reset() + assert tracker.get_whitelist() == set() + + +class _Mapper: + """Minimal stand-in for NodePackageMapper with a fixed mapping.""" + + def __init__(self, mapping): + self.mapping = mapping + + def get_package(self, ct): + return self.mapping.get(ct, "__unknown__") + + def get_all_packages(self): + return set(self.mapping.values()) - {"__builtin__"} + + +def test_package_stats_flags_whitelisted(tracker): + mapper = _Mapper({"NodeA": "Pack-A", "NodeB": "Pack-B"}) + tracker.record_usage(["NodeA", "NodeB"], mapper) + tracker.add_to_whitelist("Pack-A") + + stats = {p["package"]: p for p in tracker.get_package_stats(mapper)} + assert stats["Pack-A"]["whitelisted"] is True + assert stats["Pack-B"]["whitelisted"] is False + + +def test_package_stats_whitelist_is_case_insensitive(tracker): + mapper = _Mapper({"NodeA": "Pack-A"}) + tracker.record_usage(["NodeA"], mapper) + tracker.add_to_whitelist("pack-a") # different case than the package name + + stats = {p["package"]: p for p in tracker.get_package_stats(mapper)} + assert stats["Pack-A"]["whitelisted"] is True diff --git a/tracker.py b/tracker.py index ab09cc6..c09c04d 100644 --- a/tracker.py +++ b/tracker.py @@ -57,6 +57,11 @@ CREATE TABLE IF NOT EXISTS trial_packages ( budget INTEGER NOT NULL DEFAULT 7 ); +CREATE TABLE IF NOT EXISTS whitelist_packages ( + package TEXT PRIMARY KEY, + added_at TEXT NOT NULL +); + CREATE INDEX IF NOT EXISTS idx_node_usage_package ON node_usage(package); CREATE INDEX IF NOT EXISTS idx_prompt_log_timestamp ON prompt_log(timestamp); CREATE INDEX IF NOT EXISTS idx_model_usage_type ON model_usage(model_type); @@ -286,6 +291,13 @@ class UsageTracker: tracking_start, one_month_ago, two_months_ago, "unused_new" ) + # Flag whitelisted packages so the UI can pull them into their own group + # and suppress disable actions. Whitelist entries are kept verbatim; match + # case-insensitively since directory names vary by how users clone/symlink. + whitelist = {w.lower() for w in self.get_whitelist()} + for entry in packages.values(): + entry["whitelisted"] = entry["package"].lower() in whitelist + result = [p for p in packages.values() if p["package"].lower() not in EXCLUDED_PACKAGES] result.sort(key=lambda p: p["total_executions"]) return result @@ -480,6 +492,44 @@ class UsageTracker: finally: conn.close() + def add_to_whitelist(self, package): + """Protect a package: keep it out of disable actions and its own group.""" + now = datetime.now(timezone.utc).isoformat() + with self._lock: + self._ensure_db() + conn = self._connect() + try: + conn.execute( + """INSERT INTO whitelist_packages (package, added_at) VALUES (?, ?) + ON CONFLICT(package) DO NOTHING""", + (package, now), + ) + conn.commit() + finally: + conn.close() + + def remove_from_whitelist(self, package): + """Remove a package from the whitelist.""" + with self._lock: + self._ensure_db() + conn = self._connect() + try: + conn.execute("DELETE FROM whitelist_packages WHERE package = ?", (package,)) + conn.commit() + finally: + conn.close() + + def get_whitelist(self): + """Return the set of whitelisted package names.""" + with self._lock: + self._ensure_db() + conn = self._connect() + try: + rows = conn.execute("SELECT package FROM whitelist_packages").fetchall() + return {r[0] for r in rows} + finally: + conn.close() + def reset(self): """Clear all tracked data.""" with self._lock: @@ -490,6 +540,7 @@ class UsageTracker: conn.execute("DELETE FROM prompt_log") conn.execute("DELETE FROM model_usage") conn.execute("DELETE FROM trial_packages") + conn.execute("DELETE FROM whitelist_packages") conn.commit() finally: conn.close()