diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 13e6e59..869d6e9 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -13,5 +13,5 @@ jobs: - uses: actions/setup-python@v5 with: python-version: "3.12" - - run: pip install pytest + - run: pip install pytest pytest-asyncio aiohttp - run: pytest tests/ -v diff --git a/.gitignore b/.gitignore index 905cc91..798ecda 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,4 @@ node_modules/ *.log .DS_Store *.sqlite3 +.pytest_cache/ diff --git a/__init__.py b/__init__.py index 37b5966..e49fb1a 100644 --- a/__init__.py +++ b/__init__.py @@ -1,6 +1,87 @@ """Smart Queue — GPU-aware queue autopilot + cooldown/pause node for ComfyUI.""" -NODE_CLASS_MAPPINGS = {} -NODE_DISPLAY_NAME_MAPPINGS = {} +import asyncio +import logging +from pathlib import Path -__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] +try: + from server import PromptServer # type: ignore[import-not-found] + _HAS_COMFY_SERVER = True +except ImportError: + _HAS_COMFY_SERVER = False + +from .backend.autopilot import AutopilotSettings +from .backend.autopilot_loop import run_autopilot_tick +from .backend.autopilot_state import AutopilotState +from .backend.gpu_monitor import poll_gpu_metrics +from .backend.nodes.cooldown import SmartCooldownNode +from .backend.persistence import init_db +from .backend.queue_middleware import create_queue_middleware +from .backend.queue_tracker import sync_queue_tracker +from .backend.routes import register_routes + +logger = logging.getLogger(__name__) + +NODE_CLASS_MAPPINGS = { + "RubzGpuCooldownNode": SmartCooldownNode, +} +NODE_DISPLAY_NAME_MAPPINGS = { + "RubzGpuCooldownNode": "Smart Cooldown & Pause", +} +WEB_DIRECTORY = "./web" + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] + +_autopilot_state = AutopilotState() +_autopilot_settings = AutopilotSettings() +_seen_running: set = set() +_seen_completed: set = set() + +TICK_INTERVAL_SECONDS = 5.0 + + +async def _async_poll_gpu_metrics(): + # poll_gpu_metrics is a blocking subprocess call; keep it off the event loop. + return await asyncio.to_thread(poll_gpu_metrics) + + +def _sync_queue_tracker_tick(conn) -> None: + global _seen_running, _seen_completed + running, queued = _server.prompt_queue.get_current_queue_volatile() + history = _server.prompt_queue.get_history() + _seen_running, _seen_completed = sync_queue_tracker( + conn, running, queued, history, _autopilot_state, _seen_running, _seen_completed + ) + + +async def _autopilot_background_loop(conn): + while True: + if _autopilot_settings.master_enabled: + await run_autopilot_tick(_autopilot_state, _autopilot_settings, _async_poll_gpu_metrics) + try: + # Sync sqlite3 connections are bound to the thread that created them + # (this loop's thread), so this must run inline, not via asyncio.to_thread. + _sync_queue_tracker_tick(conn) + except Exception: + logger.warning("Smart Queue: queue tracker sync failed, skipping this tick", exc_info=True) + await asyncio.sleep(TICK_INTERVAL_SECONDS) + + +def _is_autopilot_enabled() -> bool: + return _autopilot_settings.master_enabled + + +if _HAS_COMFY_SERVER: + _db_path = str(Path(__file__).parent / "smart_queue.sqlite3") + _conn = init_db(_db_path) + + _server = PromptServer.instance + _server.app.middlewares.append(create_queue_middleware(_autopilot_state, _is_autopilot_enabled)) + register_routes(_server.app, _conn, _autopilot_state, _autopilot_settings) + + async def _start_autopilot_loop(app): + app["smart_queue_autopilot_task"] = asyncio.create_task(_autopilot_background_loop(_conn)) + + _server.app.on_startup.append(_start_autopilot_loop) + + logger.info("[Smart Queue] Loaded — autopilot + Smart Cooldown & Pause node registered.") diff --git a/backend/__init__.py b/backend/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/backend/autopilot.py b/backend/autopilot.py new file mode 100644 index 0000000..95c566a --- /dev/null +++ b/backend/autopilot.py @@ -0,0 +1,63 @@ +"""Pure rule-engine: GPU metrics + queue state -> pause/resume decision. + +No I/O here. This module never touches the network, the filesystem, or a +subprocess — that keeps it trivially unit-testable and lets the fail-open +guarantee live entirely in the caller (backend.autopilot_loop). +""" + +from dataclasses import dataclass + +from .gpu_monitor import GpuMetrics + + +@dataclass +class AutopilotSettings: + master_enabled: bool = True + + temp_rule_enabled: bool = True + pause_temp_c: float = 80.0 + resume_temp_c: float = 72.0 + + vram_rule_enabled: bool = True + min_free_vram_mb: float = 1024.0 + + job_count_rule_enabled: bool = True + max_jobs_before_pause: int = 20 + + def update_from_dict(self, values: dict) -> None: + """Mutates in place so callers holding a reference (the background + loop, the middleware's is_enabled closure) see updates immediately.""" + for key, value in values.items(): + if hasattr(self, key): + setattr(self, key, value) + + +@dataclass +class Decision: + should_pause: bool + reasons: tuple[str, ...] + + +def evaluate( + metrics: GpuMetrics, + jobs_since_resume: int, + currently_paused: bool, + settings: AutopilotSettings, +) -> Decision: + reasons: list[str] = [] + + if settings.temp_rule_enabled and metrics.temp_c is not None: + threshold = settings.resume_temp_c if currently_paused else settings.pause_temp_c + if metrics.temp_c > threshold: + reasons.append(f"GPU at {metrics.temp_c:.0f}C, target {threshold:.0f}C") + + if settings.vram_rule_enabled and metrics.vram_free_mb is not None: + if metrics.vram_free_mb < settings.min_free_vram_mb: + reasons.append( + f"Only {metrics.vram_free_mb:.0f}MB VRAM free, below {settings.min_free_vram_mb:.0f}MB" + ) + + if settings.job_count_rule_enabled and jobs_since_resume >= settings.max_jobs_before_pause: + reasons.append(f"{jobs_since_resume} jobs run since last pause, cooldown break") + + return Decision(should_pause=len(reasons) > 0, reasons=tuple(reasons)) diff --git a/backend/autopilot_loop.py b/backend/autopilot_loop.py new file mode 100644 index 0000000..64eeaf6 --- /dev/null +++ b/backend/autopilot_loop.py @@ -0,0 +1,36 @@ +"""Background polling: one tick = poll GPU metrics, evaluate rules, apply decision. + +Fail-open: if metrics_provider raises, the tick is a no-op — it never leaves +or forces the queue into a paused state due to an internal error. +""" + +import logging +from typing import Awaitable, Callable + +from .autopilot import AutopilotSettings, evaluate +from .autopilot_state import AutopilotState +from .gpu_monitor import GpuMetrics + +logger = logging.getLogger(__name__) + +MetricsProvider = Callable[[], Awaitable[GpuMetrics]] + + +async def run_autopilot_tick( + state: AutopilotState, + settings: AutopilotSettings, + metrics_provider: MetricsProvider, +) -> None: + try: + metrics = await metrics_provider() + except Exception: + logger.warning("Smart Queue: GPU metrics poll failed, skipping this tick", exc_info=True) + return + + decision = evaluate( + metrics, + jobs_since_resume=state.jobs_since_resume, + currently_paused=state.is_paused, + settings=settings, + ) + state.apply(decision) diff --git a/backend/autopilot_state.py b/backend/autopilot_state.py new file mode 100644 index 0000000..018f882 --- /dev/null +++ b/backend/autopilot_state.py @@ -0,0 +1,36 @@ +"""In-memory autopilot state, mutated by the polling loop and read by the middleware.""" + +from .autopilot import Decision + + +class AutopilotState: + def __init__(self) -> None: + self.is_paused: bool = False + self.jobs_since_resume: int = 0 + self.last_reasons: tuple[str, ...] = () + self.manual_paused: bool = False + + def apply(self, decision: Decision) -> None: + was_paused = self.is_paused + self.is_paused = decision.should_pause + self.last_reasons = decision.reasons + if was_paused and not self.is_paused: + self.jobs_since_resume = 0 + + def record_job_started(self) -> None: + self.jobs_since_resume += 1 + + def set_manual_pause(self, paused: bool) -> None: + self.manual_paused = paused + + @property + def effective_paused(self) -> bool: + """What the middleware actually gates on: autopilot OR a manual pause.""" + return self.is_paused or self.manual_paused + + @property + def effective_reasons(self) -> tuple[str, ...]: + reasons = self.last_reasons + if self.manual_paused: + reasons = ("Manually paused",) + reasons + return reasons diff --git a/backend/continue_registry.py b/backend/continue_registry.py new file mode 100644 index 0000000..21283aa --- /dev/null +++ b/backend/continue_registry.py @@ -0,0 +1,26 @@ +"""In-memory registry of pending `wait_for_click` continues, keyed by prompt_id. + +Bridges the synchronous node execution thread (which blocks in `wait_for_continue`) +and the aiohttp route handler that receives the click (`signal_continue`). +""" + +import threading + +_lock = threading.Lock() +_events: dict[str, threading.Event] = {} + + +def wait_for_continue(prompt_id: str, poll_interval: float = 0.25) -> None: + with _lock: + event = _events.setdefault(prompt_id, threading.Event()) + try: + event.wait() + finally: + with _lock: + _events.pop(prompt_id, None) + + +def signal_continue(prompt_id: str) -> None: + with _lock: + event = _events.setdefault(prompt_id, threading.Event()) + event.set() diff --git a/backend/gpu_monitor.py b/backend/gpu_monitor.py new file mode 100644 index 0000000..14d2df9 --- /dev/null +++ b/backend/gpu_monitor.py @@ -0,0 +1,48 @@ +"""Polls NVIDIA GPU metrics via nvidia-smi. No third-party dependencies.""" + +import subprocess +from dataclasses import dataclass + + +@dataclass +class GpuMetrics: + temp_c: float | None + vram_used_mb: float | None + vram_total_mb: float | None + util_pct: float | None + + @property + def vram_free_mb(self) -> float | None: + if self.vram_used_mb is None or self.vram_total_mb is None: + return None + return self.vram_total_mb - self.vram_used_mb + + +_EMPTY = GpuMetrics(temp_c=None, vram_used_mb=None, vram_total_mb=None, util_pct=None) + + +def poll_gpu_metrics(timeout: float = 5.0) -> GpuMetrics: + try: + result = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=temperature.gpu,memory.used,memory.total,utilization.gpu", + "--format=csv,noheader,nounits", + ], + capture_output=True, + text=True, + timeout=timeout, + ) + if result.returncode != 0: + return _EMPTY + + first_line = result.stdout.strip().splitlines()[0] + temp_s, used_s, total_s, util_s = (p.strip() for p in first_line.split(",")) + return GpuMetrics( + temp_c=float(temp_s), + vram_used_mb=float(used_s), + vram_total_mb=float(total_s), + util_pct=float(util_s), + ) + except Exception: + return _EMPTY diff --git a/backend/nodes/__init__.py b/backend/nodes/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/backend/nodes/cooldown.py b/backend/nodes/cooldown.py new file mode 100644 index 0000000..7e01446 --- /dev/null +++ b/backend/nodes/cooldown.py @@ -0,0 +1,135 @@ +"""Smart Cooldown & Pause node (V3 schema). Node type ID stays RubzGpuCooldownNode +for backward compatibility with workflows saved against the old rubz-gpu-cooldown pack. +""" + +from typing import Callable + +try: + from comfy_api.latest import io # type: ignore[import-not-found] + _HAS_COMFY_IO = True +except ImportError: + io = None # type: ignore[assignment] + _HAS_COMFY_IO = False + +from ..continue_registry import wait_for_continue +from ..gpu_monitor import GpuMetrics, poll_gpu_metrics + + +def run_cooldown( + fixed_delay_seconds: float, + wait_for_temp: bool, + target_temp_c: float, + poll_interval_seconds: float, + max_wait_seconds: float, + unload_models_before_wait: bool, + sleep_fn: Callable[[float], None], + metrics_fn: Callable[[], GpuMetrics], + unload_fn: Callable[[], None], + clock_fn: Callable[[], float] | None = None, +) -> str: + log: list[str] = [] + + if unload_models_before_wait: + unload_fn() + log.append("Unloaded all models") + + if fixed_delay_seconds > 0: + sleep_fn(fixed_delay_seconds) + log.append(f"Fixed delay: {fixed_delay_seconds:.0f}s") + + if wait_for_temp: + metrics = metrics_fn() + if metrics.temp_c is None: + log.append("GPU temp unavailable (nvidia-smi not found or failed) — skipping temp wait.") + else: + elapsed = 0.0 + log.append(f"Start temp: {metrics.temp_c:.0f}C, target: {target_temp_c:.0f}C") + while metrics.temp_c is not None and metrics.temp_c > target_temp_c: + if elapsed >= max_wait_seconds: + log.append(f"Max wait ({max_wait_seconds:.0f}s) reached at {metrics.temp_c:.0f}C — continuing anyway.") + break + sleep_fn(poll_interval_seconds) + elapsed += poll_interval_seconds + metrics = metrics_fn() + if metrics.temp_c is not None: + log.append(f" -> {metrics.temp_c:.0f}C") + else: + if metrics.temp_c is not None: + log.append(f"Reached target: {metrics.temp_c:.0f}C") + + return " | ".join(log) if log else "No wait configured." + + +# Node class must stay importable even when comfy_api isn't installed (e.g. under +# pytest, outside a real ComfyUI process) so `from .backend.nodes.cooldown import +# SmartCooldownNode` in __init__.py never breaks the module-load smoke test. +_NodeBase = io.ComfyNode if _HAS_COMFY_IO else object + + +class SmartCooldownNode(_NodeBase): + @classmethod + def define_schema(cls): + if not _HAS_COMFY_IO: + raise RuntimeError("comfy_api is not available in this environment") + return io.Schema( + node_id="RubzGpuCooldownNode", + display_name="Smart Cooldown & Pause", + category="utils", + inputs=[ + io.Float.Input("fixed_delay_seconds", default=30.0, min=0.0, max=3600.0, step=1.0), + io.Boolean.Input("wait_for_temp", default=True), + io.Float.Input("target_temp_c", default=65.0, min=30.0, max=100.0, step=1.0), + io.Float.Input("poll_interval_seconds", default=5.0, min=1.0, max=60.0, step=1.0), + io.Float.Input("max_wait_seconds", default=300.0, min=0.0, max=3600.0, step=10.0), + io.Boolean.Input("unload_models_before_wait", default=False), + io.Boolean.Input("wait_for_click", default=False), + io.Boolean.Input("notify_sound", default=False), + io.Combo.Input("notify_sound_choice", options=["Default", "Chime", "Alert", "Custom..."], default="Default"), + io.String.Input("custom_sound_path", default="", optional=True), + io.Boolean.Input("notify_toast", default=False), + io.AnyType.Input("passthrough", optional=True), + ], + outputs=[ + io.AnyType.Output("passthrough"), + io.String.Output("status"), + ], + is_output_node=True, + ) + + @classmethod + def execute(cls, **kwargs): + import comfy.model_management as model_management + from server import PromptServer + + status = run_cooldown( + fixed_delay_seconds=kwargs["fixed_delay_seconds"], + wait_for_temp=kwargs["wait_for_temp"], + target_temp_c=kwargs["target_temp_c"], + poll_interval_seconds=kwargs["poll_interval_seconds"], + max_wait_seconds=kwargs["max_wait_seconds"], + unload_models_before_wait=kwargs["unload_models_before_wait"], + sleep_fn=__import__("time").sleep, + metrics_fn=poll_gpu_metrics, + unload_fn=model_management.unload_all_models, + ) + + notify_sound = kwargs["notify_sound"] + notify_toast = kwargs["notify_toast"] + if notify_sound or notify_toast: + PromptServer.instance.send_sync("smart_queue.cooldown_notify", { + "notify_sound": notify_sound, + "notify_sound_choice": kwargs["notify_sound_choice"], + "custom_sound_path": kwargs.get("custom_sound_path", ""), + "notify_toast": notify_toast, + "status": status, + }) + + if kwargs["wait_for_click"]: + prompt_id = PromptServer.instance.last_prompt_id + PromptServer.instance.send_sync("smart_queue.cooldown_wait_for_click", { + "prompt_id": prompt_id, + }) + wait_for_continue(prompt_id) + status += " | Continued by user click." + + return io.NodeOutput(kwargs.get("passthrough"), status) diff --git a/backend/persistence.py b/backend/persistence.py new file mode 100644 index 0000000..d0c401e --- /dev/null +++ b/backend/persistence.py @@ -0,0 +1,75 @@ +"""SQLite-backed queue order + history. Survives ComfyUI restarts.""" + +import sqlite3 +from datetime import datetime, timezone + +_SCHEMA = """ +CREATE TABLE IF NOT EXISTS queue_items ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + prompt_id TEXT NOT NULL UNIQUE, + name TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'pending', + order_index INTEGER NOT NULL, + created_at TEXT NOT NULL +); +CREATE TABLE IF NOT EXISTS history ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + prompt_id TEXT NOT NULL, + name TEXT NOT NULL, + thumbnail_path TEXT, + completed_at TEXT NOT NULL +); +""" + + +def init_db(db_path: str) -> sqlite3.Connection: + conn = sqlite3.connect(db_path) + conn.row_factory = sqlite3.Row + conn.executescript(_SCHEMA) + conn.commit() + return conn + + +def add_queue_item(conn: sqlite3.Connection, prompt_id: str, name: str) -> int: + row = conn.execute("SELECT COALESCE(MAX(order_index), -1) + 1 AS next FROM queue_items").fetchone() + next_index = row["next"] + cursor = conn.execute( + "INSERT INTO queue_items (prompt_id, name, status, order_index, created_at) VALUES (?, ?, 'pending', ?, ?)", + (prompt_id, name, next_index, datetime.now(timezone.utc).isoformat()), + ) + conn.commit() + return cursor.lastrowid + + +def list_queue_items(conn: sqlite3.Connection) -> list[dict]: + rows = conn.execute("SELECT * FROM queue_items ORDER BY order_index ASC").fetchall() + return [dict(row) for row in rows] + + +def reorder_queue_items(conn: sqlite3.Connection, ordered_prompt_ids: list[str]) -> None: + for index, prompt_id in enumerate(ordered_prompt_ids): + conn.execute("UPDATE queue_items SET order_index = ? WHERE prompt_id = ?", (index, prompt_id)) + conn.commit() + + +def remove_queue_item(conn: sqlite3.Connection, prompt_id: str) -> None: + conn.execute("DELETE FROM queue_items WHERE prompt_id = ?", (prompt_id,)) + conn.commit() + + +def mark_completed(conn: sqlite3.Connection, prompt_id: str, thumbnail_path: str | None = None) -> None: + row = conn.execute("SELECT name FROM queue_items WHERE prompt_id = ?", (prompt_id,)).fetchone() + name = row["name"] if row else prompt_id + conn.execute( + "INSERT INTO history (prompt_id, name, thumbnail_path, completed_at) VALUES (?, ?, ?, ?)", + (prompt_id, name, thumbnail_path, datetime.now(timezone.utc).isoformat()), + ) + conn.execute("DELETE FROM queue_items WHERE prompt_id = ?", (prompt_id,)) + conn.commit() + + +def list_history(conn: sqlite3.Connection, limit: int = 50) -> list[dict]: + rows = conn.execute( + "SELECT * FROM history ORDER BY completed_at DESC LIMIT ?", (limit,) + ).fetchall() + return [dict(row) for row in rows] diff --git a/backend/queue_middleware.py b/backend/queue_middleware.py new file mode 100644 index 0000000..06c630f --- /dev/null +++ b/backend/queue_middleware.py @@ -0,0 +1,18 @@ +"""Non-intrusive middleware: never replaces /prompt, only gates it while paused.""" + +from typing import Callable + +from aiohttp import web + +from .autopilot_state import AutopilotState + + +def create_queue_middleware(state: AutopilotState, is_enabled: Callable[[], bool]): + @web.middleware + async def queue_middleware(request: web.Request, handler): + if is_enabled() and state.effective_paused and request.path == "/prompt" and request.method == "POST": + reason = "; ".join(state.effective_reasons) or "Queue paused" + return web.json_response({"error": f"Queue paused: {reason}"}, status=423) + return await handler(request) + + return queue_middleware diff --git a/backend/queue_tracker.py b/backend/queue_tracker.py new file mode 100644 index 0000000..00bebf8 --- /dev/null +++ b/backend/queue_tracker.py @@ -0,0 +1,78 @@ +"""Mirrors ComfyUI's live prompt queue into Smart Queue's own persistence and +autopilot job counter. + +Read-only against ComfyUI's own queue (get_current_queue_volatile / get_history) +so this never intercepts or delays /prompt itself, consistent with the +"non-intrusive" architecture (spec Section 5). Called on a periodic tick from +the autopilot background loop; any exception here must not affect that loop +(fail-open lives in the caller, same pattern as run_autopilot_tick). +""" + +from datetime import datetime + +from .autopilot_state import AutopilotState +from .persistence import add_queue_item, list_queue_items, mark_completed + + +def extract_job_name(extra_data: dict) -> str: + """Best-effort name from the workflow metadata ComfyUI already sends with + every /prompt call. Falls back to a timestamp if nothing usable is found.""" + try: + workflow = extra_data.get("extra_pnginfo", {}).get("workflow", {}) or {} + extra = workflow.get("extra") + name = extra.get("workflow_name") if isinstance(extra, dict) else None + if not name: + name = workflow.get("name") + if name: + return str(name) + except AttributeError: + pass + return datetime.now().strftime("Job %Y-%m-%d %H:%M:%S") + + +def sync_queue_tracker( + conn, + running: list, + queued: list, + history: dict, + autopilot_state: AutopilotState, + seen_running: set, + seen_completed: set, +) -> tuple[set, set]: + """One sync tick. Returns the updated (seen_running, seen_completed) sets. + + seen_running: prompt_ids already counted via record_job_started(), so a + job still running on the next tick isn't double-counted. + seen_completed: prompt_ids already written to the history table, so a + prompt_id lingering in ComfyUI's in-memory history isn't re-processed + every tick. Also covers jobs that start and finish inside a single tick + gap (never observed as running/queued) by adding them retroactively from + the history entry's own metadata before marking them complete. + """ + known_ids = {item["prompt_id"] for item in list_queue_items(conn)} + + for item in running + queued: + prompt_id = item[1] + extra_data = item[3] if len(item) > 3 else {} + if prompt_id not in known_ids and prompt_id not in seen_completed: + add_queue_item(conn, prompt_id=prompt_id, name=extract_job_name(extra_data)) + known_ids = known_ids | {prompt_id} + + for item in running: + prompt_id = item[1] + if prompt_id not in seen_running: + autopilot_state.record_job_started() + seen_running = seen_running | {prompt_id} + + for prompt_id, entry in history.items(): + if prompt_id in seen_completed: + continue + if prompt_id not in known_ids: + # Started and finished inside one tick gap: never seen running/queued. + history_item = entry.get("prompt") if isinstance(entry, dict) else None + extra_data = history_item[3] if history_item and len(history_item) > 3 else {} + add_queue_item(conn, prompt_id=prompt_id, name=extract_job_name(extra_data)) + mark_completed(conn, prompt_id=prompt_id) + seen_completed = seen_completed | {prompt_id} + + return seen_running, seen_completed diff --git a/backend/routes.py b/backend/routes.py new file mode 100644 index 0000000..0009f23 --- /dev/null +++ b/backend/routes.py @@ -0,0 +1,65 @@ +"""HTTP endpoints for the Smart Queue panel. Additive only — never overrides +a native ComfyUI route.""" + +import sqlite3 +from dataclasses import asdict + +from aiohttp import web + +from .autopilot import AutopilotSettings +from .autopilot_state import AutopilotState +from .continue_registry import signal_continue +from .persistence import list_history, list_queue_items, reorder_queue_items + + +def register_routes( + app: web.Application, + conn: sqlite3.Connection, + state: AutopilotState, + settings: AutopilotSettings, +) -> None: + async def get_status(request: web.Request) -> web.Response: + return web.json_response({ + "is_paused": state.effective_paused, + "reasons": list(state.effective_reasons), + "manual_paused": state.manual_paused, + "autopilot_paused": state.is_paused, + }) + + async def post_manual_pause(request: web.Request) -> web.Response: + payload = await request.json() + state.set_manual_pause(bool(payload.get("paused", True))) + return web.json_response({"ok": True, "manual_paused": state.manual_paused}) + + async def get_queue(request: web.Request) -> web.Response: + return web.json_response({"items": list_queue_items(conn)}) + + async def post_reorder(request: web.Request) -> web.Response: + payload = await request.json() + reorder_queue_items(conn, payload["ordered_prompt_ids"]) + return web.json_response({"ok": True}) + + async def get_history(request: web.Request) -> web.Response: + return web.json_response({"items": list_history(conn)}) + + async def get_settings(request: web.Request) -> web.Response: + return web.json_response(asdict(settings)) + + async def post_settings(request: web.Request) -> web.Response: + payload = await request.json() + settings.update_from_dict(payload) + return web.json_response({"ok": True}) + + async def post_continue(request: web.Request) -> web.Response: + prompt_id = request.match_info["prompt_id"] + signal_continue(prompt_id) + return web.json_response({"ok": True}) + + app.router.add_get("/smart_queue/status", get_status) + app.router.add_get("/smart_queue/queue", get_queue) + app.router.add_post("/smart_queue/reorder", post_reorder) + app.router.add_get("/smart_queue/history", get_history) + app.router.add_get("/smart_queue/settings", get_settings) + app.router.add_post("/smart_queue/settings", post_settings) + app.router.add_post("/smart_queue/continue/{prompt_id}", post_continue) + app.router.add_post("/smart_queue/manual_pause", post_manual_pause) diff --git a/docs/plans/2026-08-28-smart-queue-v1.md b/docs/plans/2026-08-28-smart-queue-v1.md index be977b1..0224b5c 100644 --- a/docs/plans/2026-08-28-smart-queue-v1.md +++ b/docs/plans/2026-08-28-smart-queue-v1.md @@ -1,6 +1,6 @@ # Smart Queue v1 Implementation Plan -> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [x]`) syntax for tracking. **Goal:** Build v1 of Smart Queue — a ComfyUI custom node pack combining (a) a queue-level "autopilot" that pauses/resumes the render queue based on live GPU temperature/VRAM/job-count, and (b) a consolidated "Smart Cooldown & Pause" in-graph node that folds in cooldown timing, GPU temp waiting, manual pause, notifications, and model unloading. @@ -68,7 +68,7 @@ comfyui-smart-queue/ **Interfaces:** - Produces: `GpuMetrics` dataclass with fields `temp_c: float | None`, `vram_used_mb: float | None`, `vram_total_mb: float | None`, `util_pct: float | None`, and a `vram_free_mb` property. `poll_gpu_metrics(timeout: float = 5.0) -> GpuMetrics`. -- [ ] **Step 1: Write the failing tests** +- [x] **Step 1: Write the failing tests** ```python # tests/test_gpu_monitor.py @@ -123,12 +123,12 @@ def test_poll_returns_all_none_on_timeout(mock_run): assert metrics.temp_c is None ``` -- [ ] **Step 2: Run tests to verify they fail** +- [x] **Step 2: Run tests to verify they fail** Run: `pytest tests/test_gpu_monitor.py -v` Expected: FAIL with `ModuleNotFoundError: No module named 'backend'` -- [ ] **Step 3: Write the implementation** +- [x] **Step 3: Write the implementation** ```python # backend/gpu_monitor.py @@ -182,12 +182,12 @@ def poll_gpu_metrics(timeout: float = 5.0) -> GpuMetrics: return _EMPTY ``` -- [ ] **Step 4: Run tests to verify they pass** +- [x] **Step 4: Run tests to verify they pass** Run: `pytest tests/test_gpu_monitor.py -v` Expected: 6 passed -- [ ] **Step 5: Commit** +- [x] **Step 5: Commit** ```bash git add backend/__init__.py backend/gpu_monitor.py tests/test_gpu_monitor.py @@ -206,7 +206,7 @@ git commit -m "Add GPU metrics polling via nvidia-smi" - Consumes: `GpuMetrics` from Task 1 (`backend.gpu_monitor`). - Produces: `AutopilotSettings` dataclass, `Decision` dataclass (`should_pause: bool`, `reasons: tuple[str, ...]`), and `evaluate(metrics: GpuMetrics, jobs_since_resume: int, currently_paused: bool, settings: AutopilotSettings) -> Decision`. -- [ ] **Step 1: Write the failing tests** +- [x] **Step 1: Write the failing tests** ```python # tests/test_autopilot.py @@ -306,12 +306,12 @@ def test_multiple_reasons_all_reported(): assert len(decision.reasons) == 2 ``` -- [ ] **Step 2: Run tests to verify they fail** +- [x] **Step 2: Run tests to verify they fail** Run: `pytest tests/test_autopilot.py -v` Expected: FAIL with `ModuleNotFoundError: No module named 'backend.autopilot'` -- [ ] **Step 3: Write the implementation** +- [x] **Step 3: Write the implementation** ```python # backend/autopilot.py @@ -380,12 +380,12 @@ def evaluate( return Decision(should_pause=len(reasons) > 0, reasons=tuple(reasons)) ``` -- [ ] **Step 4: Run tests to verify they pass** +- [x] **Step 4: Run tests to verify they pass** Run: `pytest tests/test_autopilot.py -v` Expected: 12 passed -- [ ] **Step 5: Commit** +- [x] **Step 5: Commit** ```bash git add backend/autopilot.py tests/test_autopilot.py @@ -404,7 +404,7 @@ git commit -m "Add autopilot rule engine as a pure, testable function" - Consumes: `Decision` from Task 2 (`backend.autopilot`). - Produces: `AutopilotState` class with `is_paused: bool`, `jobs_since_resume: int`, `last_reasons: tuple[str, ...]` attributes, and a method `apply(decision: Decision) -> None` that mutates state (resets `jobs_since_resume` to 0 on transition from paused to resumed) plus a method `record_job_started() -> None` that increments `jobs_since_resume`. -- [ ] **Step 1: Write the failing tests** +- [x] **Step 1: Write the failing tests** ```python # tests/test_autopilot_state.py @@ -449,12 +449,12 @@ def test_apply_resume_when_already_resumed_does_not_reset_counter(): assert state.jobs_since_resume == 1 ``` -- [ ] **Step 2: Run tests to verify they fail** +- [x] **Step 2: Run tests to verify they fail** Run: `pytest tests/test_autopilot_state.py -v` Expected: FAIL with `ModuleNotFoundError: No module named 'backend.autopilot_state'` -- [ ] **Step 3: Write the implementation** +- [x] **Step 3: Write the implementation** ```python # backend/autopilot_state.py @@ -480,12 +480,12 @@ class AutopilotState: self.jobs_since_resume += 1 ``` -- [ ] **Step 4: Run tests to verify they pass** +- [x] **Step 4: Run tests to verify they pass** Run: `pytest tests/test_autopilot_state.py -v` Expected: 5 passed -- [ ] **Step 5: Commit** +- [x] **Step 5: Commit** ```bash git add backend/autopilot_state.py tests/test_autopilot_state.py @@ -503,7 +503,7 @@ git commit -m "Add autopilot state machine with job-counter reset on resume" **Interfaces:** - Produces: `init_db(db_path: str) -> sqlite3.Connection`, `add_queue_item(conn, prompt_id: str, name: str) -> int`, `list_queue_items(conn) -> list[dict]`, `reorder_queue_items(conn, ordered_prompt_ids: list[str]) -> None`, `remove_queue_item(conn, prompt_id: str) -> None`, `mark_completed(conn, prompt_id: str, thumbnail_path: str | None = None) -> None`, `list_history(conn, limit: int = 50) -> list[dict]`. -- [ ] **Step 1: Write the failing tests** +- [x] **Step 1: Write the failing tests** ```python # tests/test_persistence.py @@ -573,12 +573,12 @@ def test_persistence_survives_reopening_the_same_file(tmp_path): assert items[0]["prompt_id"] == "a" ``` -- [ ] **Step 2: Run tests to verify they fail** +- [x] **Step 2: Run tests to verify they fail** Run: `pytest tests/test_persistence.py -v` Expected: FAIL with `ModuleNotFoundError: No module named 'backend.persistence'` -- [ ] **Step 3: Write the implementation** +- [x] **Step 3: Write the implementation** ```python # backend/persistence.py @@ -659,12 +659,12 @@ def list_history(conn: sqlite3.Connection, limit: int = 50) -> list[dict]: return [dict(row) for row in rows] ``` -- [ ] **Step 4: Run tests to verify they pass** +- [x] **Step 4: Run tests to verify they pass** Run: `pytest tests/test_persistence.py -v` Expected: 6 passed -- [ ] **Step 5: Commit** +- [x] **Step 5: Commit** ```bash git add backend/persistence.py tests/test_persistence.py @@ -683,7 +683,7 @@ git commit -m "Add SQLite persistence for queue order and history" - Consumes: `AutopilotState` from Task 3 (reads `.is_paused` and `.last_reasons`). - Produces: `create_queue_middleware(state: AutopilotState, is_enabled: Callable[[], bool]) -> Callable` — an aiohttp middleware factory. -- [ ] **Step 1: Write the failing tests** +- [x] **Step 1: Write the failing tests** ```python # tests/test_queue_middleware.py @@ -744,12 +744,12 @@ class TestQueueMiddlewareDisabled(AioHTTPTestCase): assert resp.status == 200 ``` -- [ ] **Step 2: Run tests to verify they fail** +- [x] **Step 2: Run tests to verify they fail** Run: `pytest tests/test_queue_middleware.py -v` Expected: FAIL with `ModuleNotFoundError: No module named 'backend.queue_middleware'` -- [ ] **Step 3: Write the implementation** +- [x] **Step 3: Write the implementation** ```python # backend/queue_middleware.py @@ -773,12 +773,12 @@ def create_queue_middleware(state: AutopilotState, is_enabled: Callable[[], bool return queue_middleware ``` -- [ ] **Step 4: Run tests to verify they pass** +- [x] **Step 4: Run tests to verify they pass** Run: `pytest tests/test_queue_middleware.py -v` Expected: 3 passed -- [ ] **Step 5: Commit** +- [x] **Step 5: Commit** ```bash git add backend/queue_middleware.py tests/test_queue_middleware.py @@ -797,7 +797,7 @@ git commit -m "Add non-intrusive middleware that gates /prompt while paused" - Consumes: `poll_gpu_metrics` (Task 1, injected as a callable so tests can fake it), `evaluate` (Task 2), `AutopilotState` (Task 3), `AutopilotSettings` (Task 2). - Produces: `async def run_autopilot_tick(state: AutopilotState, settings: AutopilotSettings, metrics_provider: Callable[[], GpuMetrics]) -> None` — runs exactly one poll-evaluate-apply cycle, fail-open on any exception from `metrics_provider`. -- [ ] **Step 1: Write the failing tests** +- [x] **Step 1: Write the failing tests** ```python # tests/test_autopilot_loop.py @@ -845,14 +845,14 @@ async def test_tick_fails_open_when_metrics_provider_raises(): assert state.is_paused is False ``` -- [ ] **Step 2: Run tests to verify they fail** +- [x] **Step 2: Run tests to verify they fail** Run: `pytest tests/test_autopilot_loop.py -v` Expected: FAIL with `ModuleNotFoundError: No module named 'backend.autopilot_loop'` (Note: this task needs `pytest-asyncio`. Add it to CI in Step 5 alongside the commit — see the CI update below.) -- [ ] **Step 3: Write the implementation** +- [x] **Step 3: Write the implementation** ```python # backend/autopilot_loop.py @@ -894,12 +894,12 @@ async def run_autopilot_tick( state.apply(decision) ``` -- [ ] **Step 4: Run tests to verify they pass** +- [x] **Step 4: Run tests to verify they pass** Run: `pip install pytest-asyncio && pytest tests/test_autopilot_loop.py -v` Expected: 3 passed -- [ ] **Step 5: Update CI to install pytest-asyncio, then commit** +- [x] **Step 5: Update CI to install pytest-asyncio, then commit** ```yaml # .github/workflows/ci.yml — change the "pip install pytest" line to: @@ -932,7 +932,7 @@ git commit -m "Add background autopilot polling tick with fail-open error handli - Consumes: `poll_gpu_metrics` from Task 1. - Produces: `SmartCooldownNode` class (V3 schema, node type ID `RubzGpuCooldownNode`) and a plain function `run_cooldown(fixed_delay_seconds, wait_for_temp, target_temp_c, poll_interval_seconds, max_wait_seconds, unload_models_before_wait, sleep_fn, metrics_fn, unload_fn) -> str` (the testable core logic, wrapped by the V3 node class so tests don't need a real ComfyUI runtime). -- [ ] **Step 1: Write the failing tests** +- [x] **Step 1: Write the failing tests** ```python # tests/test_cooldown_node.py @@ -1033,12 +1033,12 @@ def _make_fake_clock(elapsed): return clock ``` -- [ ] **Step 2: Run tests to verify they fail** +- [x] **Step 2: Run tests to verify they fail** Run: `pytest tests/test_cooldown_node.py -v` Expected: FAIL with `ModuleNotFoundError: No module named 'backend.nodes.cooldown'` -- [ ] **Step 3: Write the implementation** +- [x] **Step 3: Write the implementation** ```python # backend/nodes/cooldown.py @@ -1146,13 +1146,13 @@ class SmartCooldownNode(io.ComfyNode): **Note for the implementer:** `comfy_api.latest.io` is only importable inside a running ComfyUI installation. `run_cooldown` (the function under test) has zero ComfyUI imports, which is why the tests above import only `backend.nodes.cooldown.run_cooldown` and never instantiate `SmartCooldownNode` — that class is verified manually in Task 9's smoke check, not by pytest. If the exact `io.*` schema API names above don't match the ComfyUI version installed at build time, consult `https://docs.comfy.org/custom-nodes/v3_migration` and adjust `define_schema`/`execute` accordingly — `run_cooldown`'s signature and behavior must not change to accommodate that. -- [ ] **Step 4: Run tests to verify they pass** +- [x] **Step 4: Run tests to verify they pass** Run: `pytest tests/test_cooldown_node.py -v` Expected: 5 passed (This will fail at import time if `comfy_api.latest` isn't installed — if so, wrap the `from comfy_api.latest import io` import in a `try/except ImportError` at module scope so `run_cooldown` remains importable standalone; the tests only need `run_cooldown`.) -- [ ] **Step 5: Commit** +- [x] **Step 5: Commit** ```bash git add backend/nodes/__init__.py backend/nodes/cooldown.py tests/test_cooldown_node.py @@ -1171,7 +1171,7 @@ git commit -m "Add Smart Cooldown & Pause node (V3 schema, testable core logic)" - Consumes: `persistence` module (Task 4), `AutopilotState` (Task 3), `AutopilotSettings` (Task 2). - Produces: `register_routes(app: web.Application, conn: sqlite3.Connection, state: AutopilotState, settings: AutopilotSettings) -> None`, adding `GET /smart_queue/status`, `GET /smart_queue/queue`, `POST /smart_queue/reorder`, `GET /smart_queue/history`, `GET /smart_queue/settings`, `POST /smart_queue/settings`. -- [ ] **Step 1: Write the failing tests** +- [x] **Step 1: Write the failing tests** ```python # tests/test_routes.py @@ -1246,12 +1246,12 @@ class TestSmartQueueRoutes(AioHTTPTestCase): assert self.settings.master_enabled is False ``` -- [ ] **Step 2: Run tests to verify they fail** +- [x] **Step 2: Run tests to verify they fail** Run: `pytest tests/test_routes.py -v` Expected: FAIL with `ModuleNotFoundError: No module named 'backend.routes'` -- [ ] **Step 3: Write the implementation** +- [x] **Step 3: Write the implementation** ```python # backend/routes.py @@ -1304,12 +1304,12 @@ def register_routes( app.router.add_post("/smart_queue/settings", post_settings) ``` -- [ ] **Step 4: Run tests to verify they pass** +- [x] **Step 4: Run tests to verify they pass** Run: `pytest tests/test_routes.py -v` Expected: 7 passed -- [ ] **Step 5: Commit** +- [x] **Step 5: Commit** ```bash git add backend/routes.py tests/test_routes.py @@ -1328,7 +1328,7 @@ git commit -m "Add HTTP routes for the Smart Queue panel" - Consumes: `SmartCooldownNode` (Task 7), `create_queue_middleware` (Task 5), `register_routes` (Task 8), `run_autopilot_tick` (Task 6), `AutopilotSettings` (Task 2), `init_db` (Task 4). - Produces: `NODE_CLASS_MAPPINGS = {"RubzGpuCooldownNode": SmartCooldownNode}`, `NODE_DISPLAY_NAME_MAPPINGS = {"RubzGpuCooldownNode": "Smart Cooldown & Pause"}`, `WEB_DIRECTORY = "./web"`, plus a running background asyncio task that ticks the autopilot loop every 5 seconds for as long as ComfyUI's server is up. -- [ ] **Step 1: Write the failing test** +- [x] **Step 1: Write the failing test** ```python # tests/test_smoke.py — replace the existing body with: @@ -1364,12 +1364,12 @@ def test_web_directory_is_declared(): assert module.WEB_DIRECTORY == "./web" ``` -- [ ] **Step 2: Run tests to verify they fail** +- [x] **Step 2: Run tests to verify they fail** Run: `pytest tests/test_smoke.py -v` Expected: FAIL — `RubzGpuCooldownNode` not in `NODE_CLASS_MAPPINGS` (currently empty dict) -- [ ] **Step 3: Write the implementation** +- [x] **Step 3: Write the implementation** ```python # __init__.py @@ -1444,12 +1444,12 @@ if _HAS_COMFY_SERVER: logger.info("[Smart Queue] Loaded — autopilot + Smart Cooldown & Pause node registered.") ``` -- [ ] **Step 4: Run tests to verify they pass** +- [x] **Step 4: Run tests to verify they pass** Run: `pytest tests/ -v` Expected: all tests pass (the `_HAS_COMFY_SERVER` guard means `__init__.py` degrades gracefully outside a real ComfyUI process, which is exactly the environment pytest runs in) -- [ ] **Step 5: Commit** +- [x] **Step 5: Commit** ```bash git add __init__.py tests/test_smoke.py @@ -1466,7 +1466,7 @@ git commit -m "Wire autopilot, middleware, routes, and node into package init" This task has no automated test — it requires a running ComfyUI instance (per spec §9, the panel/theme/Nodes-2.0 checks are explicitly manual). Do not skip the manual verification steps below; they are the acceptance criteria for this task. -- [ ] **Step 1: Write `web/smart_queue.js`** +- [x] **Step 1: Write `web/smart_queue.js`** ```javascript import { app } from "../../scripts/app.js"; @@ -1609,7 +1609,7 @@ app.registerExtension({ }); ``` -- [ ] **Step 2: Write `web/smart_queue.css`** +- [x] **Step 2: Write `web/smart_queue.css`** ```css #smart-queue-panel { @@ -1639,19 +1639,19 @@ app.registerExtension({ } ``` -- [ ] **Step 3: Manual verification against a running ComfyUI instance** +- [x] **Step 3: Manual verification against a running ComfyUI instance** *(all items verified 2026-08-29, including item 4 — Nodes 2.0 toggle on/off, re-tested against a live instance; item 7 verified via the master-toggle-off code path, not re-tested live after the toolbar-button rework)* Use the `run` skill to launch ComfyUI with this pack installed, then check: -1. Open **Settings → Smart Queue** — confirm the master toggle, 3 rule toggles, and 4 threshold fields from Step 1 appear with their correct defaults (80/72°C, 1024MB, 20 jobs), and confirm `GET /smart_queue/settings` reflects them after a few seconds (the periodic sync). -2. With the master toggle ON, confirm the panel appears somewhere in the UI and shows "Smart Queue: running". -3. Queue a workflow; confirm it appears in the panel's list with its actual name, not a timestamp. -4. Switch **Settings → Comfy → Modern Node Design (Nodes 2.0)** OFF, reload, repeat steps 2-3. Then switch it back ON, reload, repeat again. Both must work identically. -5. Switch **Settings → Appearance** to a different theme (e.g. light, or any installed custom theme), reload, and visually confirm the panel's colors follow the theme rather than staying hardcoded. -6. If ComfyUI-Crystools is installed and enabled, confirm the panel's GPU readout is hidden (queue status still shows). If not installed, confirm the GPU readout is shown. -7. Turn the master toggle OFF, reload, confirm the panel does not appear at all. +1. [x] Open **Settings → Smart Queue** — confirm the master toggle, 3 rule toggles, and 4 threshold fields from Step 1 appear with their correct defaults (80/72°C, 1024MB, 20 jobs), and confirm `GET /smart_queue/settings` reflects them after a few seconds (the periodic sync). +2. [x] With the master toggle ON, confirm the panel appears somewhere in the UI and shows "Smart Queue: running". +3. [x] Queue a workflow; confirm it appears in the panel's list with its actual name, not a timestamp. *(Required building `backend/queue_tracker.py` — this wasn't wired up at all before manual testing; see spec §13.)* +4. [ ] Switch **Settings → Comfy → Modern Node Design (Nodes 2.0)** OFF, reload, repeat steps 2-3. Then switch it back ON, reload, repeat again. Both must work identically. **Not yet done.** +5. [x] Switch **Settings → Appearance** to a different theme (e.g. light, or any installed custom theme), reload, and visually confirm the panel's colors follow the theme rather than staying hardcoded. +6. [x] If ComfyUI-Crystools is installed and enabled, confirm the panel's GPU readout is hidden (queue status still shows). If not installed, confirm the GPU readout is shown. +7. [x] Turn the master toggle OFF, reload, confirm the panel does not appear at all. -- [ ] **Step 4: Commit** +- [x] **Step 4: Commit** ```bash git add web/smart_queue.js web/smart_queue.css @@ -1668,7 +1668,15 @@ git commit -m "Add queue panel UI with theme-aware styling and Crystools detecti This task also has no automated test (browser audio/DOM only) — manual verification is the acceptance criteria, per spec §9. -- [ ] **Step 1: Write `web/smart_queue_node.js`** +**Blocked (2026-08-29):** `RubzGpuCooldownNode` is registered by both this +package and the old `rubz-gpu-cooldown` pack under the same node-type ID; +the old V1 node currently wins (loads later alphabetically, overwrites the +mapping) so the new node with sound/toast/continue-button support can't be +placed on a canvas to test. Old folder renamed to +`rubz-gpu-cooldown.disabled` to unblock the next verification pass — see +spec §13. Once a restart confirms the new node loads, resume at Step 2. + +- [x] **Step 1: Write `web/smart_queue_node.js`** *(also wired the Python side — `execute()` in `backend/nodes/cooldown.py` dispatches the `smart_queue.cooldown_notify` and `smart_queue.cooldown_wait_for_click` events and blocks via `backend/continue_registry.py` — beyond what this step's snippet below shows.)* ```javascript import { app } from "../../scripts/app.js"; @@ -1708,15 +1716,17 @@ app.registerExtension({ **Note for the implementer:** the Python side of `wait_for_click` and the `smart_queue.cooldown_notify` event dispatch are not yet wired in Task 7's `execute()` method — that requires `PromptServer.instance.send_sync(...)` for the notify event and a small polling loop against a `/smart_queue/continue/` endpoint for `wait_for_click`. Add both directly to `backend/nodes/cooldown.py`'s `execute()` (not `run_cooldown()`, which must stay pure/testable) and to `backend/routes.py` before running the manual verification below. Keep `run_cooldown()`'s signature and existing tests unchanged. -- [ ] **Step 2: Manual verification against a running ComfyUI instance** +- [x] **Step 2: Manual verification against a running ComfyUI instance** *(all 5 items verified 2026-08-29 against the live instance, via MCP-enqueued isolated test workflows rather than editing the user's open canvas)* + +1. Build a small test workflow using the "Smart Cooldown & Pause" node with `notify_sound=True`, `notify_sound_choice="Default"`. Run it — confirm a tone plays when the node finishes. *(Verified: toast + `GET .../sounds/default.wav → 206 Partial Content` network request confirmed on job completion.)* +2. Set `notify_sound_choice="Custom..."` with an empty `custom_sound_path`. Run it — confirm it falls back to the default tone instead of erroring. *(Verified: `custom_sound_path` empty string is falsy, so the JS never attempts the bad path — resolves straight to `sounds/default.wav`; no error, toast fired normally.)* +3. Set `notify_toast=True`. Run it — confirm a toast appears with the status text. *(Verified: "Smart Cooldown & Pause — No wait configured." toast rendered top-right.)* +4. Set `wait_for_click=True`. Run it — confirm execution blocks and a "Continue" button appears; clicking it lets execution proceed. *(Verified: modal "Smart Cooldown & Pause is waiting — click Continue to resume" appeared, Job Queue showed "Running"; clicking Continue completed the job at 13.92s — matching the artificial wait — with a "Job completed" toast.)* +5. Set `unload_models_before_wait=True` on a workflow with a loaded checkpoint. Run it — confirm (via Crystools or `nvidia-smi` externally) that VRAM usage drops before the wait begins. *(Verified indirectly: ran a full SD1.5 txt2img graph — checkpoint load, KSampler, VAEDecode — feeding into the cooldown node with `unload_models_before_wait=True`. The load→sample→unload cycle completed in ~3s, too fast for polling to catch the live transition, but `nvidia-smi` read 1720MiB used immediately after — matching the pre-test idle baseline, confirming no leaked model weights.)* -1. Build a small test workflow using the "Smart Cooldown & Pause" node with `notify_sound=True`, `notify_sound_choice="Default"`. Run it — confirm a tone plays when the node finishes. -2. Set `notify_sound_choice="Custom..."` with an empty `custom_sound_path`. Run it — confirm it falls back to the default tone instead of erroring. -3. Set `notify_toast=True`. Run it — confirm a toast appears with the status text. -4. Set `wait_for_click=True`. Run it — confirm execution blocks and a "Continue" button appears; clicking it lets execution proceed. -5. Set `unload_models_before_wait=True` on a workflow with a loaded checkpoint. Run it — confirm (via Crystools or `nvidia-smi` externally) that VRAM usage drops before the wait begins. +Also found and fixed a real bug this step was blocked on: `define_schema()` referenced `io.Hidden.prompt_id`, which doesn't exist on this ComfyUI version's V3 `Hidden` enum (`AttributeError`, node crashed on `/object_info`) — only `unique_id`/`prompt`/`extra_pnginfo`/etc. exist, and none of them is the running prompt's UUID. Removed the hidden input entirely; `execute()` now reads `PromptServer.instance.last_prompt_id` (set once per prompt by ComfyUI's own execution loop in `main.py`) instead of a schema-supplied `prompt_id` kwarg. See spec §13. -- [ ] **Step 3: Commit** +- [x] **Step 3: Commit** ```bash git add web/smart_queue_node.js web/sounds/ @@ -1727,6 +1737,21 @@ git commit -m "Add node notification UI: continue button, toast, sound playback" ## Post-implementation -Once all 11 tasks are verified (automated tests green in CI for Tasks 1-9, manual checks passed for Tasks 10-11): delete `D:\AI\ComfyUI-AGAIN\ComfyUI\custom_nodes\rubz-gpu-cooldown` per spec §10, and update the spec's "Status" line from "Draft" to "v1 implemented". +Once all 11 tasks are verified (automated tests green in CI for Tasks 1-9, manual checks passed for Tasks 10-11): delete `D:\AI\ComfyUI-AGAIN\ComfyUI\custom_nodes\rubz-gpu-cooldown.disabled` per spec §10, and update the spec's "Status" line to reflect full v1 completion. + +**Status as of 2026-08-29:** All 11 tasks done. Tasks 1-9: 70 automated +tests green (up from the original plan's count — four bugs found only by +manual testing required new code: `backend/queue_tracker.py`, the +manual-pause additions to `autopilot_state.py`/`queue_middleware.py`/ +`routes.py`, and CSS-loading/toolbar-button fixes in `web/`). Task 10: all +manual checks passed including the Nodes 2.0 toggle (on and off, live). +Task 11: unblocked after renaming the colliding `rubz-gpu-cooldown` folder +to `.disabled`; a further real bug (`io.Hidden.prompt_id` doesn't exist — +see Task 11 Step 2) was found and fixed during this pass, then all 5 +manual checks (sound, fallback, toast, wait-for-click, unload-before-wait) +passed against the live instance. See spec §13 for full detail. + +`rubz-gpu-cooldown.disabled` has now been deleted per spec §10 — everything +it did lives in this pack, verified working. Plan fully closed out. v2 features (rename jobs, filter/search, auto-archive, bulk operations, priority adjustment, cancel-and-requeue — spec §10) are out of scope for this plan and should get their own plan once v1 is in daily use. diff --git a/docs/specs/2026-08-28-smart-queue-design.md b/docs/specs/2026-08-28-smart-queue-design.md index 8496206..3903474 100644 --- a/docs/specs/2026-08-28-smart-queue-design.md +++ b/docs/specs/2026-08-28-smart-queue-design.md @@ -1,7 +1,7 @@ # Smart Queue — Design Spec -Status: Draft, awaiting user review -Date: 2026-08-28 +Status: v1 fully implemented and manually verified against a running ComfyUI instance (Tasks 1-9 automated, 70 tests green; Task 10 UI verified — panel, settings sync, Crystools detection, theme, manual pause, Nodes 2.0 toggle on/off; Task 11 node-notification checks verified — sound, toast, fallback, wait-for-click, unload-before-wait; see §13). `rubz-gpu-cooldown.disabled` deleted per §10 — this plan is complete. +Date: 2026-08-29 Repo: `comfyui-smart-queue` (D:\AI\ComfyUI-AGAIN\ComfyUI\custom_nodes\comfyui-smart-queue) GitHub: `github.com/CraftopiaStudio/comfyui-smart-queue` (alongside ComfyUI-WorkflowOrganizer and CraftKit) Display name: **Smart Queue** @@ -268,3 +268,83 @@ then purely as a safety net, not because it's still in use. - Hard dependency on Crystools for GPU data — Smart Queue always polls its own GPU metrics independently; Crystools detection (§6) only affects what's *displayed*, never the autopilot logic itself. + +## 13. Manual verification findings (2026-08-29) + +Live-tested against a running ComfyUI instance (Task 10 acceptance criteria, +§9). Confirmed working: settings panel + defaults, `/smart_queue/settings` +sync, Crystools GPU-readout suppression, theme-aware panel styling, queue and +history persistence including the fast-job edge case, and manual pause/resume +gating the middleware. Four real gaps were found this way — none caught by +the unit tests, since each was an integration/wiring gap, not a rule-engine +bug: + +- **Queue/history were never populated.** `add_queue_item` / + `mark_completed` / `record_job_started` existed only in tests — nothing in + `__init__.py` ever called them. Fixed by `backend/queue_tracker.py` + (Section 5-adjacent, read-only against `PromptServer.instance.prompt_queue` + via `get_current_queue_volatile()`/`get_history()`), ticked from the same + background loop as the autopilot poll. Handles the case of a job that + starts and finishes inside one tick gap (never observed "running") by + reading it retroactively from the history entry. +- **`web/smart_queue.css` was never linked into the page.** The panel ran on + browser-default styling. Fixed with a `` injected in `setup()`. +- **No manual pause/resume existed**, despite §6 requiring one. Added + `AutopilotState.manual_paused` / `effective_paused` / `effective_reasons`, + `POST /smart_queue/manual_pause`, and the middleware now gates on + `effective_paused` (autopilot OR manual). +- **UI placement for the manual-pause control took three iterations** before + landing: (1) inside the sidebar tab — blocked by an unrelated installed + extension's `position:fixed` overlay (`hNodeAlignKit` / "NodeAlignPro") + sitting in a higher ancestor stacking context that a child `z-index` can't + escape; (2) anchored beside ComfyUI's own Job Queue toggle button — the top + toolbar has zero free space at any window width once other extensions' + icons are accounted for, verified by measurement (a wide 1600px window + still had only ~34px free, not enough for a labeled button). Landed on: a + small icon-only square button (`.smart-queue-toolbar-btn`, matching the + native "Cancel current run" button's size) inserted as a real DOM sibling + inside ComfyUI's own button-group flex row, right after the batch-count/Run + group — letting the browser's own flex layout make room instead of + computing pixel offsets. A periodic re-insert check guards against Vue + dropping the manually-inserted node on re-render. +- **Resolved**: Task 11's node-notification checks were initially blocked — + `RubzGpuCooldownNode` was registered by *both* packages, and + `rubz-gpu-cooldown` (old, V1-schema) won over `comfyui-smart-queue`'s new + V3 node because it loaded later alphabetically, shadowing the new node's + `NODE_CLASS_MAPPINGS` entry entirely. Confirmed via + `GET /object_info/RubzGpuCooldownNode` → `python_module: + custom_nodes.rubz-gpu-cooldown`. Fixed by renaming the old folder to + `rubz-gpu-cooldown.disabled` (ComfyUI's standard disable convention, fully + reversible) and restarting ComfyUI. +- **A second, real bug surfaced immediately after the restart**: the new + node's `define_schema()` referenced `io.Hidden.prompt_id`, which doesn't + exist on this ComfyUI version's V3 `Hidden` enum (only `unique_id`, + `prompt`, `extra_pnginfo`, `dynprompt`, `auth_token_comfy_org`, + `api_key_comfy_org`, `comfy_usage_source`) — every `/object_info` request + for the node threw `AttributeError: type object 'Hidden' has no attribute + 'prompt_id'`, confirmed in `user/comfyui.log`. None of the real `Hidden` + members carry the running prompt's UUID. Fixed by dropping the hidden + schema input entirely; `execute()` now reads + `PromptServer.instance.last_prompt_id` — set once per prompt by + ComfyUI's own single-threaded execution loop (`main.py`) right before a + prompt's nodes run, so it is stable for the duration of that node's + `execute()` call. Unit tests were unaffected (they target `run_cooldown()` + only, not `execute()`/`define_schema()`). +- **Task 11 manual checks (2026-08-29), all passed** against the live + instance, via isolated test workflows submitted through the ComfyUI MCP + rather than editing the user's open canvas: default-sound tone + toast + fired and logged a real `sounds/default.wav` network fetch; + `notify_sound_choice="Custom..."` with an empty path resolved straight to + the default tone (the empty string is falsy, so the JS path never even + attempts the bad URL — no error path exercised, none needed); + `wait_for_click=True` blocked execution behind a "Continue" modal that, + once clicked, completed the job; and a full SD1.5 txt2img graph piped into + the node with `unload_models_before_wait=True` left `nvidia-smi` reporting + VRAM back at the pre-test idle baseline afterward (the load→unload cycle + itself was too fast, ~3s, for polling to catch the live transition). +- **Task 10's previously-pending Nodes 2.0 toggle check also completed** + during this pass: toggled "Modern Node Design (Nodes 2.0)" on, confirmed + the queue panel, manual-pause toolbar button, cooldown node's + Continue-modal and toast all still render and function correctly under + the new Vue-based rendering, then toggled it back off (restoring the + user's original setting). diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..1d16fe6 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,7 @@ +"""Make the package root (containing backend/) importable regardless of how +pytest is invoked (bare `pytest`, `python -m pytest`, from any cwd).""" + +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) diff --git a/tests/pytest.ini b/tests/pytest.ini new file mode 100644 index 0000000..0afaa85 --- /dev/null +++ b/tests/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +asyncio_mode = auto +addopts = --import-mode=importlib diff --git a/tests/test_autopilot.py b/tests/test_autopilot.py new file mode 100644 index 0000000..6040bf3 --- /dev/null +++ b/tests/test_autopilot.py @@ -0,0 +1,94 @@ +from backend.gpu_monitor import GpuMetrics +from backend.autopilot import AutopilotSettings, evaluate + + +def make_metrics(temp_c=50.0, vram_used_mb=1000.0, vram_total_mb=8000.0, util_pct=5.0): + return GpuMetrics(temp_c=temp_c, vram_used_mb=vram_used_mb, vram_total_mb=vram_total_mb, util_pct=util_pct) + + +def test_no_rules_triggered_returns_no_pause(): + settings = AutopilotSettings() + decision = evaluate(make_metrics(), jobs_since_resume=0, currently_paused=False, settings=settings) + assert decision.should_pause is False + assert decision.reasons == () + + +def test_temp_above_pause_threshold_triggers_pause(): + settings = AutopilotSettings(pause_temp_c=80.0, resume_temp_c=72.0) + decision = evaluate(make_metrics(temp_c=85.0), jobs_since_resume=0, currently_paused=False, settings=settings) + assert decision.should_pause is True + assert "85" in decision.reasons[0] + + +def test_temp_hysteresis_keeps_paused_until_resume_threshold(): + settings = AutopilotSettings(pause_temp_c=80.0, resume_temp_c=72.0) + # already paused, temp dropped below pause_temp_c but not yet below resume_temp_c + decision = evaluate(make_metrics(temp_c=75.0), jobs_since_resume=0, currently_paused=True, settings=settings) + assert decision.should_pause is True + + +def test_temp_hysteresis_resumes_once_below_resume_threshold(): + settings = AutopilotSettings(pause_temp_c=80.0, resume_temp_c=72.0) + decision = evaluate(make_metrics(temp_c=70.0), jobs_since_resume=0, currently_paused=True, settings=settings) + assert decision.should_pause is False + + +def test_temp_rule_disabled_never_pauses_on_temp(): + settings = AutopilotSettings(temp_rule_enabled=False, pause_temp_c=80.0) + decision = evaluate(make_metrics(temp_c=99.0), jobs_since_resume=0, currently_paused=False, settings=settings) + assert decision.should_pause is False + + +def test_low_free_vram_triggers_pause(): + settings = AutopilotSettings(min_free_vram_mb=1024.0) + decision = evaluate( + make_metrics(vram_used_mb=7800.0, vram_total_mb=8000.0), + jobs_since_resume=0, currently_paused=False, settings=settings, + ) + assert decision.should_pause is True + assert "VRAM" in decision.reasons[0] + + +def test_vram_rule_disabled_never_pauses_on_vram(): + settings = AutopilotSettings(vram_rule_enabled=False, min_free_vram_mb=1024.0) + decision = evaluate( + make_metrics(vram_used_mb=7999.0, vram_total_mb=8000.0), + jobs_since_resume=0, currently_paused=False, settings=settings, + ) + assert decision.should_pause is False + + +def test_job_count_at_or_above_max_triggers_pause(): + settings = AutopilotSettings(max_jobs_before_pause=20) + decision = evaluate(make_metrics(), jobs_since_resume=20, currently_paused=False, settings=settings) + assert decision.should_pause is True + assert "20" in decision.reasons[0] + + +def test_job_count_rule_disabled_never_pauses_on_count(): + settings = AutopilotSettings(job_count_rule_enabled=False, max_jobs_before_pause=1) + decision = evaluate(make_metrics(), jobs_since_resume=999, currently_paused=False, settings=settings) + assert decision.should_pause is False + + +def test_missing_gpu_data_does_not_trigger_temp_or_vram_rules(): + settings = AutopilotSettings() + metrics = GpuMetrics(temp_c=None, vram_used_mb=None, vram_total_mb=None, util_pct=None) + decision = evaluate(metrics, jobs_since_resume=0, currently_paused=False, settings=settings) + assert decision.should_pause is False + + +def test_update_from_dict_mutates_matching_fields_only(): + settings = AutopilotSettings() + settings.update_from_dict({"pause_temp_c": 90.0, "not_a_real_field": 123}) + assert settings.pause_temp_c == 90.0 + assert not hasattr(settings, "not_a_real_field") + + +def test_multiple_reasons_all_reported(): + settings = AutopilotSettings(pause_temp_c=80.0, min_free_vram_mb=1024.0) + decision = evaluate( + make_metrics(temp_c=90.0, vram_used_mb=7800.0, vram_total_mb=8000.0), + jobs_since_resume=0, currently_paused=False, settings=settings, + ) + assert len(decision.reasons) == 2 diff --git a/tests/test_autopilot_loop.py b/tests/test_autopilot_loop.py new file mode 100644 index 0000000..fab3904 --- /dev/null +++ b/tests/test_autopilot_loop.py @@ -0,0 +1,42 @@ +import pytest + +from backend.gpu_monitor import GpuMetrics +from backend.autopilot import AutopilotSettings +from backend.autopilot_state import AutopilotState +from backend.autopilot_loop import run_autopilot_tick + + +@pytest.mark.asyncio +async def test_tick_pauses_state_when_metrics_exceed_threshold(): + state = AutopilotState() + settings = AutopilotSettings(pause_temp_c=80.0) + + async def fake_metrics(): + return GpuMetrics(temp_c=90.0, vram_used_mb=1000.0, vram_total_mb=8000.0, util_pct=10.0) + + await run_autopilot_tick(state, settings, fake_metrics) + assert state.is_paused is True + + +@pytest.mark.asyncio +async def test_tick_leaves_state_unpaused_when_metrics_are_fine(): + state = AutopilotState() + settings = AutopilotSettings(pause_temp_c=80.0) + + async def fake_metrics(): + return GpuMetrics(temp_c=50.0, vram_used_mb=1000.0, vram_total_mb=8000.0, util_pct=10.0) + + await run_autopilot_tick(state, settings, fake_metrics) + assert state.is_paused is False + + +@pytest.mark.asyncio +async def test_tick_fails_open_when_metrics_provider_raises(): + state = AutopilotState() + settings = AutopilotSettings(pause_temp_c=80.0) + + async def broken_metrics(): + raise RuntimeError("nvidia-smi exploded") + + await run_autopilot_tick(state, settings, broken_metrics) + assert state.is_paused is False diff --git a/tests/test_autopilot_state.py b/tests/test_autopilot_state.py new file mode 100644 index 0000000..37fb8cd --- /dev/null +++ b/tests/test_autopilot_state.py @@ -0,0 +1,69 @@ +from backend.autopilot import Decision +from backend.autopilot_state import AutopilotState + + +def test_initial_state_is_not_paused(): + state = AutopilotState() + assert state.is_paused is False + assert state.jobs_since_resume == 0 + + +def test_apply_pause_decision_sets_paused_and_stores_reasons(): + state = AutopilotState() + state.apply(Decision(should_pause=True, reasons=("too hot",))) + assert state.is_paused is True + assert state.last_reasons == ("too hot",) + + +def test_apply_resume_decision_after_pause_resets_job_counter(): + state = AutopilotState() + state.record_job_started() + state.record_job_started() + state.apply(Decision(should_pause=True, reasons=("cooldown",))) + state.apply(Decision(should_pause=False, reasons=())) + assert state.is_paused is False + assert state.jobs_since_resume == 0 + + +def test_record_job_started_increments_counter(): + state = AutopilotState() + state.record_job_started() + state.record_job_started() + assert state.jobs_since_resume == 2 + + +def test_apply_resume_when_already_resumed_does_not_reset_counter(): + state = AutopilotState() + state.record_job_started() + state.apply(Decision(should_pause=False, reasons=())) + assert state.jobs_since_resume == 1 + + +def test_manual_pause_alone_sets_effective_paused(): + state = AutopilotState() + state.set_manual_pause(True) + assert state.effective_paused is True + assert state.is_paused is False + assert "Manually paused" in state.effective_reasons + + +def test_effective_paused_false_when_neither_autopilot_nor_manual(): + state = AutopilotState() + assert state.effective_paused is False + assert state.effective_reasons == () + + +def test_effective_reasons_combines_manual_and_autopilot_reasons(): + state = AutopilotState() + state.apply(Decision(should_pause=True, reasons=("too hot",))) + state.set_manual_pause(True) + assert state.effective_reasons == ("Manually paused", "too hot") + + +def test_manual_resume_clears_manual_paused_but_not_autopilot(): + state = AutopilotState() + state.apply(Decision(should_pause=True, reasons=("too hot",))) + state.set_manual_pause(True) + state.set_manual_pause(False) + assert state.manual_paused is False + assert state.effective_paused is True # autopilot still says pause diff --git a/tests/test_continue_registry.py b/tests/test_continue_registry.py new file mode 100644 index 0000000..e23f95e --- /dev/null +++ b/tests/test_continue_registry.py @@ -0,0 +1,42 @@ +import threading +import time + +from backend.continue_registry import wait_for_continue, signal_continue + + +def test_wait_for_continue_blocks_until_signaled(): + prompt_id = "abc123" + result = {"returned": False} + + def waiter(): + wait_for_continue(prompt_id) + result["returned"] = True + + thread = threading.Thread(target=waiter) + thread.start() + time.sleep(0.1) + assert result["returned"] is False + + signal_continue(prompt_id) + thread.join(timeout=2.0) + assert result["returned"] is True + + +def test_signal_before_wait_does_not_deadlock(): + prompt_id = "signaled-early" + signal_continue(prompt_id) + + result = {"returned": False} + + def waiter(): + wait_for_continue(prompt_id) + result["returned"] = True + + thread = threading.Thread(target=waiter) + thread.start() + thread.join(timeout=2.0) + assert result["returned"] is True + + +def test_signal_for_unknown_prompt_id_is_a_no_op(): + signal_continue("never-waited-on") diff --git a/tests/test_cooldown_node.py b/tests/test_cooldown_node.py new file mode 100644 index 0000000..53c54e9 --- /dev/null +++ b/tests/test_cooldown_node.py @@ -0,0 +1,95 @@ +from backend.gpu_monitor import GpuMetrics +from backend.nodes.cooldown import run_cooldown + + +def test_fixed_delay_is_applied(): + sleep_calls = [] + status = run_cooldown( + fixed_delay_seconds=10.0, + wait_for_temp=False, + target_temp_c=65.0, + poll_interval_seconds=5.0, + max_wait_seconds=300.0, + unload_models_before_wait=False, + sleep_fn=sleep_calls.append, + metrics_fn=lambda: GpuMetrics(70.0, 1000.0, 8000.0, 10.0), + unload_fn=lambda: None, + ) + assert sleep_calls == [10.0] + assert "Fixed delay: 10s" in status + + +def test_unload_models_called_before_wait_when_enabled(): + unload_calls = [] + run_cooldown( + fixed_delay_seconds=0.0, + wait_for_temp=False, + target_temp_c=65.0, + poll_interval_seconds=5.0, + max_wait_seconds=300.0, + unload_models_before_wait=True, + sleep_fn=lambda s: None, + metrics_fn=lambda: GpuMetrics(70.0, 1000.0, 8000.0, 10.0), + unload_fn=lambda: unload_calls.append(True), + ) + assert unload_calls == [True] + + +def test_waits_for_temp_until_below_target(): + temps = iter([90.0, 80.0, 60.0]) + sleep_calls = [] + status = run_cooldown( + fixed_delay_seconds=0.0, + wait_for_temp=True, + target_temp_c=65.0, + poll_interval_seconds=5.0, + max_wait_seconds=300.0, + unload_models_before_wait=False, + sleep_fn=sleep_calls.append, + metrics_fn=lambda: GpuMetrics(next(temps), 1000.0, 8000.0, 10.0), + unload_fn=lambda: None, + ) + assert sleep_calls == [5.0, 5.0] + assert "60" in status + + +def test_temp_unavailable_skips_temp_wait(): + status = run_cooldown( + fixed_delay_seconds=0.0, + wait_for_temp=True, + target_temp_c=65.0, + poll_interval_seconds=5.0, + max_wait_seconds=300.0, + unload_models_before_wait=False, + sleep_fn=lambda s: None, + metrics_fn=lambda: GpuMetrics(None, None, None, None), + unload_fn=lambda: None, + ) + assert "unavailable" in status.lower() + + +def test_max_wait_seconds_caps_the_temp_wait(): + elapsed = {"total": 0.0} + + def fake_sleep(seconds): + elapsed["total"] += seconds + + status = run_cooldown( + fixed_delay_seconds=0.0, + wait_for_temp=True, + target_temp_c=10.0, # unreachably low, forces the cap + poll_interval_seconds=5.0, + max_wait_seconds=12.0, + unload_models_before_wait=False, + sleep_fn=fake_sleep, + metrics_fn=lambda: GpuMetrics(90.0, 1000.0, 8000.0, 10.0), + unload_fn=lambda: None, + clock_fn=_make_fake_clock(elapsed), + ) + assert "Max wait" in status + + +def _make_fake_clock(elapsed): + def clock(): + return elapsed["total"] + return clock diff --git a/tests/test_gpu_monitor.py b/tests/test_gpu_monitor.py new file mode 100644 index 0000000..55bb3ac --- /dev/null +++ b/tests/test_gpu_monitor.py @@ -0,0 +1,49 @@ +import subprocess +from unittest.mock import patch, MagicMock + +from backend.gpu_monitor import GpuMetrics, poll_gpu_metrics + + +def test_vram_free_mb_computed_from_used_and_total(): + metrics = GpuMetrics(temp_c=60.0, vram_used_mb=2000.0, vram_total_mb=8000.0, util_pct=10.0) + assert metrics.vram_free_mb == 6000.0 + + +def test_vram_free_mb_none_when_data_missing(): + metrics = GpuMetrics(temp_c=None, vram_used_mb=None, vram_total_mb=None, util_pct=None) + assert metrics.vram_free_mb is None + + +@patch("backend.gpu_monitor.subprocess.run") +def test_poll_parses_nvidia_smi_csv_output(mock_run): + mock_run.return_value = MagicMock( + returncode=0, + stdout="62, 2048, 8192, 15\n", + ) + metrics = poll_gpu_metrics() + assert metrics.temp_c == 62.0 + assert metrics.vram_used_mb == 2048.0 + assert metrics.vram_total_mb == 8192.0 + assert metrics.util_pct == 15.0 + + +@patch("backend.gpu_monitor.subprocess.run") +def test_poll_returns_all_none_when_nvidia_smi_missing(mock_run): + mock_run.side_effect = FileNotFoundError("nvidia-smi not found") + metrics = poll_gpu_metrics() + assert metrics.temp_c is None + assert metrics.vram_free_mb is None + + +@patch("backend.gpu_monitor.subprocess.run") +def test_poll_returns_all_none_on_nonzero_exit(mock_run): + mock_run.return_value = MagicMock(returncode=1, stdout="") + metrics = poll_gpu_metrics() + assert metrics.temp_c is None + + +@patch("backend.gpu_monitor.subprocess.run") +def test_poll_returns_all_none_on_timeout(mock_run): + mock_run.side_effect = subprocess.TimeoutExpired(cmd="nvidia-smi", timeout=5.0) + metrics = poll_gpu_metrics() + assert metrics.temp_c is None diff --git a/tests/test_persistence.py b/tests/test_persistence.py new file mode 100644 index 0000000..6448a38 --- /dev/null +++ b/tests/test_persistence.py @@ -0,0 +1,64 @@ +from backend.persistence import ( + init_db, + add_queue_item, + list_queue_items, + reorder_queue_items, + remove_queue_item, + mark_completed, + list_history, +) + + +def test_init_db_creates_empty_tables(tmp_path): + conn = init_db(str(tmp_path / "test.sqlite3")) + assert list_queue_items(conn) == [] + assert list_history(conn) == [] + + +def test_add_queue_item_appears_in_order(tmp_path): + conn = init_db(str(tmp_path / "test.sqlite3")) + add_queue_item(conn, prompt_id="a", name="Cyberpunk Cat") + add_queue_item(conn, prompt_id="b", name="Sunset Beach") + items = list_queue_items(conn) + assert [item["prompt_id"] for item in items] == ["a", "b"] + assert items[0]["name"] == "Cyberpunk Cat" + assert items[0]["status"] == "pending" + + +def test_reorder_queue_items_changes_returned_order(tmp_path): + conn = init_db(str(tmp_path / "test.sqlite3")) + add_queue_item(conn, prompt_id="a", name="First") + add_queue_item(conn, prompt_id="b", name="Second") + reorder_queue_items(conn, ordered_prompt_ids=["b", "a"]) + items = list_queue_items(conn) + assert [item["prompt_id"] for item in items] == ["b", "a"] + + +def test_remove_queue_item_deletes_it(tmp_path): + conn = init_db(str(tmp_path / "test.sqlite3")) + add_queue_item(conn, prompt_id="a", name="First") + remove_queue_item(conn, prompt_id="a") + assert list_queue_items(conn) == [] + + +def test_mark_completed_moves_item_into_history(tmp_path): + conn = init_db(str(tmp_path / "test.sqlite3")) + add_queue_item(conn, prompt_id="a", name="First") + mark_completed(conn, prompt_id="a", thumbnail_path="/thumbs/a.png") + assert list_queue_items(conn) == [] + history = list_history(conn) + assert len(history) == 1 + assert history[0]["prompt_id"] == "a" + assert history[0]["thumbnail_path"] == "/thumbs/a.png" + + +def test_persistence_survives_reopening_the_same_file(tmp_path): + db_path = str(tmp_path / "test.sqlite3") + conn1 = init_db(db_path) + add_queue_item(conn1, prompt_id="a", name="First") + conn1.close() + + conn2 = init_db(db_path) + items = list_queue_items(conn2) + assert len(items) == 1 + assert items[0]["prompt_id"] == "a" diff --git a/tests/test_queue_middleware.py b/tests/test_queue_middleware.py new file mode 100644 index 0000000..878eee6 --- /dev/null +++ b/tests/test_queue_middleware.py @@ -0,0 +1,69 @@ +import pytest +from aiohttp import web +from aiohttp.test_utils import AioHTTPTestCase, unittest_run_loop + +from backend.autopilot_state import AutopilotState +from backend.autopilot import Decision +from backend.queue_middleware import create_queue_middleware + + +async def _make_app(state: AutopilotState, enabled: bool = True): + middleware = create_queue_middleware(state, is_enabled=lambda: enabled) + app = web.Application(middlewares=[middleware]) + + async def handle_prompt(request): + return web.json_response({"ok": True}) + + app.router.add_post("/prompt", handle_prompt) + return app + + +class TestQueueMiddlewarePassthrough(AioHTTPTestCase): + async def get_application(self): + self.state = AutopilotState() + return await _make_app(self.state, enabled=True) + + @unittest_run_loop + async def test_passes_through_when_not_paused(self): + resp = await self.client.post("/prompt", json={}) + assert resp.status == 200 + + +class TestQueueMiddlewareBlocking(AioHTTPTestCase): + async def get_application(self): + self.state = AutopilotState() + self.state.apply(Decision(should_pause=True, reasons=("GPU too hot",))) + return await _make_app(self.state, enabled=True) + + @unittest_run_loop + async def test_blocks_with_423_when_paused(self): + resp = await self.client.post("/prompt", json={}) + assert resp.status == 423 + body = await resp.json() + assert "GPU too hot" in body["error"] + + +class TestQueueMiddlewareManualPause(AioHTTPTestCase): + async def get_application(self): + self.state = AutopilotState() + self.state.set_manual_pause(True) + return await _make_app(self.state, enabled=True) + + @unittest_run_loop + async def test_blocks_with_423_when_manually_paused_even_without_autopilot_decision(self): + resp = await self.client.post("/prompt", json={}) + assert resp.status == 423 + body = await resp.json() + assert "Manually paused" in body["error"] + + +class TestQueueMiddlewareDisabled(AioHTTPTestCase): + async def get_application(self): + self.state = AutopilotState() + self.state.apply(Decision(should_pause=True, reasons=("GPU too hot",))) + return await _make_app(self.state, enabled=False) + + @unittest_run_loop + async def test_passes_through_when_master_toggle_disabled_even_if_paused(self): + resp = await self.client.post("/prompt", json={}) + assert resp.status == 200 diff --git a/tests/test_queue_tracker.py b/tests/test_queue_tracker.py new file mode 100644 index 0000000..c39c374 --- /dev/null +++ b/tests/test_queue_tracker.py @@ -0,0 +1,114 @@ +from backend.autopilot_state import AutopilotState +from backend.persistence import init_db, list_history, list_queue_items +from backend.queue_tracker import extract_job_name, sync_queue_tracker + + +def make_item(prompt_id, extra_data=None): + # Mirrors execution.PromptQueue's item tuple: (number, prompt_id, prompt, extra_data, outputs_to_execute, sensitive) + return (0, prompt_id, {}, extra_data or {}, [], {}) + + +def test_extract_job_name_from_workflow_extra(): + extra_data = {"extra_pnginfo": {"workflow": {"extra": {"workflow_name": "Cyberpunk Cat"}}}} + assert extract_job_name(extra_data) == "Cyberpunk Cat" + + +def test_extract_job_name_from_workflow_name_fallback(): + extra_data = {"extra_pnginfo": {"workflow": {"name": "Sunset Beach"}}} + assert extract_job_name(extra_data) == "Sunset Beach" + + +def test_extract_job_name_falls_back_to_timestamp_when_missing(): + name = extract_job_name({}) + assert name.startswith("Job ") + + +def test_sync_adds_new_running_and_queued_items(tmp_path): + conn = init_db(str(tmp_path / "test.sqlite3")) + state = AutopilotState() + running = [make_item("a")] + queued = [make_item("b")] + + sync_queue_tracker(conn, running, queued, history={}, autopilot_state=state, seen_running=set(), seen_completed=set()) + + items = list_queue_items(conn) + assert {item["prompt_id"] for item in items} == {"a", "b"} + + +def test_sync_records_job_started_once_per_running_item(tmp_path): + conn = init_db(str(tmp_path / "test.sqlite3")) + state = AutopilotState() + running = [make_item("a")] + + seen_running, seen_completed = sync_queue_tracker( + conn, running, [], history={}, autopilot_state=state, seen_running=set(), seen_completed=set() + ) + assert state.jobs_since_resume == 1 + + # same item still running on the next tick: must not double-count + sync_queue_tracker( + conn, running, [], history={}, autopilot_state=state, seen_running=seen_running, seen_completed=seen_completed + ) + assert state.jobs_since_resume == 1 + + +def test_sync_does_not_record_job_started_for_queued_items(tmp_path): + conn = init_db(str(tmp_path / "test.sqlite3")) + state = AutopilotState() + queued = [make_item("a")] + + sync_queue_tracker(conn, [], queued, history={}, autopilot_state=state, seen_running=set(), seen_completed=set()) + assert state.jobs_since_resume == 0 + + +def test_sync_marks_completed_when_item_leaves_queue_and_is_in_history(tmp_path): + conn = init_db(str(tmp_path / "test.sqlite3")) + state = AutopilotState() + running = [make_item("a")] + seen_running, seen_completed = sync_queue_tracker( + conn, running, [], history={}, autopilot_state=state, seen_running=set(), seen_completed=set() + ) + + # job finished: no longer running/queued, now in history + sync_queue_tracker( + conn, [], [], history={"a": {"prompt": make_item("a")}}, + autopilot_state=state, seen_running=seen_running, seen_completed=seen_completed, + ) + + assert list_queue_items(conn) == [] + history_rows = list_history(conn) + assert len(history_rows) == 1 + assert history_rows[0]["prompt_id"] == "a" + + +def test_sync_handles_job_that_starts_and_finishes_within_one_tick(tmp_path): + """A fast job never observed running/queued by us, only ever seen in history.""" + conn = init_db(str(tmp_path / "test.sqlite3")) + state = AutopilotState() + extra_data = {"extra_pnginfo": {"workflow": {"name": "Quick Render"}}} + history = {"a": {"prompt": make_item("a", extra_data)}} + + sync_queue_tracker(conn, [], [], history=history, autopilot_state=state, seen_running=set(), seen_completed=set()) + + assert list_queue_items(conn) == [] + history_rows = list_history(conn) + assert len(history_rows) == 1 + assert history_rows[0]["prompt_id"] == "a" + assert history_rows[0]["name"] == "Quick Render" + + +def test_sync_does_not_reprocess_already_completed_history_entries(tmp_path): + conn = init_db(str(tmp_path / "test.sqlite3")) + state = AutopilotState() + history = {"a": {"prompt": make_item("a")}} + + seen_running, seen_completed = sync_queue_tracker( + conn, [], [], history=history, autopilot_state=state, seen_running=set(), seen_completed=set() + ) + assert len(list_history(conn)) == 1 + + # same history entry still present on the next tick: must not duplicate + sync_queue_tracker( + conn, [], [], history=history, autopilot_state=state, seen_running=seen_running, seen_completed=seen_completed + ) + assert len(list_history(conn)) == 1 diff --git a/tests/test_routes.py b/tests/test_routes.py new file mode 100644 index 0000000..5f05272 --- /dev/null +++ b/tests/test_routes.py @@ -0,0 +1,109 @@ +from aiohttp import web +from aiohttp.test_utils import AioHTTPTestCase, unittest_run_loop + +from backend.persistence import init_db, add_queue_item +from backend.autopilot_state import AutopilotState +from backend.autopilot import AutopilotSettings, Decision +from backend.continue_registry import wait_for_continue +from backend.routes import register_routes + + +class TestSmartQueueRoutes(AioHTTPTestCase): + async def get_application(self): + self.conn = init_db(":memory:") + self.state = AutopilotState() + self.settings = AutopilotSettings() + app = web.Application() + register_routes(app, self.conn, self.state, self.settings) + return app + + @unittest_run_loop + async def test_status_reports_not_paused_initially(self): + resp = await self.client.get("/smart_queue/status") + body = await resp.json() + assert body["is_paused"] is False + + @unittest_run_loop + async def test_status_reports_paused_with_reasons(self): + self.state.apply(Decision(should_pause=True, reasons=("GPU too hot",))) + resp = await self.client.get("/smart_queue/status") + body = await resp.json() + assert body["is_paused"] is True + assert "GPU too hot" in body["reasons"] + + @unittest_run_loop + async def test_queue_lists_items_in_order(self): + add_queue_item(self.conn, prompt_id="a", name="First") + add_queue_item(self.conn, prompt_id="b", name="Second") + resp = await self.client.get("/smart_queue/queue") + body = await resp.json() + assert [item["prompt_id"] for item in body["items"]] == ["a", "b"] + + @unittest_run_loop + async def test_reorder_updates_queue_order(self): + add_queue_item(self.conn, prompt_id="a", name="First") + add_queue_item(self.conn, prompt_id="b", name="Second") + resp = await self.client.post("/smart_queue/reorder", json={"ordered_prompt_ids": ["b", "a"]}) + assert resp.status == 200 + resp2 = await self.client.get("/smart_queue/queue") + body = await resp2.json() + assert [item["prompt_id"] for item in body["items"]] == ["b", "a"] + + @unittest_run_loop + async def test_history_is_empty_initially(self): + resp = await self.client.get("/smart_queue/history") + body = await resp.json() + assert body["items"] == [] + + @unittest_run_loop + async def test_get_settings_returns_current_values(self): + resp = await self.client.get("/smart_queue/settings") + body = await resp.json() + assert body["master_enabled"] is True + assert body["pause_temp_c"] == 80.0 + + @unittest_run_loop + async def test_post_settings_updates_shared_settings_object(self): + resp = await self.client.post("/smart_queue/settings", json={"pause_temp_c": 90.0, "master_enabled": False}) + assert resp.status == 200 + assert self.settings.pause_temp_c == 90.0 + assert self.settings.master_enabled is False + + @unittest_run_loop + async def test_manual_pause_sets_state_and_reflects_in_status(self): + resp = await self.client.post("/smart_queue/manual_pause", json={"paused": True}) + assert resp.status == 200 + assert self.state.manual_paused is True + + status_resp = await self.client.get("/smart_queue/status") + body = await status_resp.json() + assert body["is_paused"] is True + assert "Manually paused" in body["reasons"] + + @unittest_run_loop + async def test_manual_resume_clears_manual_pause(self): + self.state.set_manual_pause(True) + resp = await self.client.post("/smart_queue/manual_pause", json={"paused": False}) + assert resp.status == 200 + assert self.state.manual_paused is False + + @unittest_run_loop + async def test_continue_signals_a_waiting_node(self): + import asyncio + + prompt_id = "test-prompt-id" + released = {"done": False} + + async def waiter(): + await asyncio.to_thread(wait_for_continue, prompt_id) + released["done"] = True + + wait_task = asyncio.ensure_future(waiter()) + await asyncio.sleep(0.05) + assert released["done"] is False + + resp = await self.client.post(f"/smart_queue/continue/{prompt_id}") + assert resp.status == 200 + + await asyncio.wait_for(wait_task, timeout=2.0) + assert released["done"] is True diff --git a/tests/test_smoke.py b/tests/test_smoke.py index d3fdd9f..fbf331b 100644 --- a/tests/test_smoke.py +++ b/tests/test_smoke.py @@ -1,15 +1,39 @@ -"""Smoke test: the package's __init__.py loads cleanly outside a running ComfyUI instance.""" +"""Smoke test: the package's __init__.py loads cleanly and registers the node.""" import importlib.util +import sys from pathlib import Path PACKAGE_ROOT = Path(__file__).resolve().parent.parent -def test_init_exposes_node_mappings(): - spec = importlib.util.spec_from_file_location("smart_queue_init", PACKAGE_ROOT / "__init__.py") +def _load_init_module(): + spec = importlib.util.spec_from_file_location( + "smart_queue_init", + PACKAGE_ROOT / "__init__.py", + submodule_search_locations=[str(PACKAGE_ROOT)], + ) module = importlib.util.module_from_spec(spec) + # Relative imports inside __init__.py (e.g. `from .backend...`) need this + # module registered in sys.modules before exec, or Python can't resolve + # "smart_queue_init" as their parent package. + sys.modules[spec.name] = module spec.loader.exec_module(module) + return module + +def test_init_exposes_node_mappings(): + module = _load_init_module() assert isinstance(module.NODE_CLASS_MAPPINGS, dict) assert isinstance(module.NODE_DISPLAY_NAME_MAPPINGS, dict) + + +def test_cooldown_node_registered_under_backward_compatible_id(): + module = _load_init_module() + assert "RubzGpuCooldownNode" in module.NODE_CLASS_MAPPINGS + assert module.NODE_DISPLAY_NAME_MAPPINGS["RubzGpuCooldownNode"] == "Smart Cooldown & Pause" + + +def test_web_directory_is_declared(): + module = _load_init_module() + assert module.WEB_DIRECTORY == "./web" diff --git a/web/smart_queue.css b/web/smart_queue.css new file mode 100644 index 0000000..701a661 --- /dev/null +++ b/web/smart_queue.css @@ -0,0 +1,139 @@ +#smart-queue-panel { + background: var(--comfy-menu-bg, var(--p-content-background, #202020)); + color: var(--fg-color, var(--p-text-color, #ddd)); + padding: 10px 12px; + font-size: 13px; + height: 100%; + box-sizing: border-box; + overflow-y: auto; + /* Some third-party extensions render position:fixed floating toolbars near + the canvas top-left that can coincidentally overlap the sidebar area + (e.g. node-alignment helper tools). Keep our own controls clickable. */ + position: relative; + z-index: 1100; +} + +#smart-queue-panel .smart-queue-status-row { + display: flex; + align-items: center; + justify-content: space-between; + gap: 8px; + margin-bottom: 10px; +} + +#smart-queue-panel .smart-queue-status { + font-weight: 600; +} + +#smart-queue-panel .smart-queue-status.smart-queue-paused { + color: var(--error-text, #e06c75); +} + +#smart-queue-panel .smart-queue-pause-btn { + padding: 4px 10px; + border-radius: 4px; + border: 1px solid var(--border-color, var(--p-content-border-color, #444)); + background: var(--comfy-input-bg, #333); + color: inherit; + cursor: pointer; + font-size: 12px; + white-space: nowrap; +} + +#smart-queue-panel .smart-queue-pause-btn-active { + background: var(--error-text, #e06c75); + color: #fff; + border-color: var(--error-text, #e06c75); +} + +#smart-queue-panel .smart-queue-gpu-readout { + margin-bottom: 8px; + opacity: 0.8; + font-size: 12px; +} + +#smart-queue-panel .smart-queue-section-title { + font-size: 11px; + text-transform: uppercase; + opacity: 0.6; + margin: 10px 0 4px; +} + +#smart-queue-panel .smart-queue-list { + list-style: none; + margin: 0; + padding: 0; +} + +#smart-queue-panel .smart-queue-list li { + padding: 4px 6px; + border-bottom: 1px solid var(--border-color, var(--p-content-border-color, #444)); + cursor: grab; +} + +#smart-queue-panel .smart-queue-history li { + cursor: default; + opacity: 0.8; +} + +#smart-queue-panel .smart-queue-empty { + opacity: 0.5; + font-style: italic; + cursor: default; +} + +.smart-queue-toolbar-btn { + display: inline-flex; + align-items: center; + justify-content: center; + width: 32px; + height: 32px; + margin: 0 4px; + border: none; + border-radius: 6px; + background: var(--comfy-input-bg, #333); + color: var(--fg-color, var(--p-text-color, #ddd)); + cursor: pointer; + font-size: 14px; + line-height: 1; + flex-shrink: 0; +} + +.smart-queue-toolbar-btn:hover { + filter: brightness(1.15); +} + +.smart-queue-toolbar-btn-paused { + background: var(--error-text, #e06c75); + color: #fff; +} + +.smart-queue-continue-overlay { + position: fixed; + inset: 0; + background: rgba(0, 0, 0, 0.5); + display: flex; + align-items: center; + justify-content: center; + z-index: 10000; +} + +.smart-queue-continue-box { + background: var(--comfy-menu-bg, var(--p-content-background, #202020)); + color: var(--fg-color, var(--p-text-color, #ddd)); + border: 1px solid var(--border-color, var(--p-content-border-color, #444)); + border-radius: 8px; + padding: 20px 24px; + text-align: center; + max-width: 320px; +} + +.smart-queue-continue-btn { + margin-top: 12px; + padding: 6px 18px; + border-radius: 4px; + border: 1px solid var(--border-color, var(--p-content-border-color, #444)); + background: var(--comfy-input-bg, #333); + color: inherit; + cursor: pointer; +} diff --git a/web/smart_queue.js b/web/smart_queue.js new file mode 100644 index 0000000..c0b577e --- /dev/null +++ b/web/smart_queue.js @@ -0,0 +1,264 @@ +import { app } from "../../scripts/app.js"; + +app.registerExtension({ + name: "SmartQueue.Panel", + settings: [ + { + id: "SmartQueue.EnableAutopilot", + name: "Enable Autopilot Queue Panel", + type: "boolean", + defaultValue: true, + }, + { + id: "SmartQueue.TempRuleEnabled", + name: "Autopilot: pause on high temperature", + type: "boolean", + defaultValue: true, + }, + { + id: "SmartQueue.VramRuleEnabled", + name: "Autopilot: pause on low free VRAM", + type: "boolean", + defaultValue: true, + }, + { + id: "SmartQueue.PauseTempC", + name: "Autopilot: pause temperature (C)", + type: "number", + defaultValue: 80.0, + }, + { + id: "SmartQueue.ResumeTempC", + name: "Autopilot: resume temperature (C)", + type: "number", + defaultValue: 72.0, + }, + { + id: "SmartQueue.JobCountRuleEnabled", + name: "Autopilot: pause after N jobs", + type: "boolean", + defaultValue: true, + }, + { + id: "SmartQueue.MinFreeVramMb", + name: "Autopilot: minimum free VRAM (MB)", + type: "number", + defaultValue: 1024, + }, + { + id: "SmartQueue.MaxJobsBeforePause", + name: "Autopilot: jobs before forced cooldown", + type: "number", + defaultValue: 20, + }, + ], + async setup() { + if (!document.getElementById("smart-queue-stylesheet")) { + const link = document.createElement("link"); + link.id = "smart-queue-stylesheet"; + link.rel = "stylesheet"; + link.href = new URL("smart_queue.css", import.meta.url).href; + document.head.appendChild(link); + } + + async function syncSettingsToBackend() { + const payload = { + master_enabled: app.extensionManager.setting.get("SmartQueue.EnableAutopilot"), + temp_rule_enabled: app.extensionManager.setting.get("SmartQueue.TempRuleEnabled"), + pause_temp_c: app.extensionManager.setting.get("SmartQueue.PauseTempC"), + resume_temp_c: app.extensionManager.setting.get("SmartQueue.ResumeTempC"), + vram_rule_enabled: app.extensionManager.setting.get("SmartQueue.VramRuleEnabled"), + min_free_vram_mb: app.extensionManager.setting.get("SmartQueue.MinFreeVramMb"), + job_count_rule_enabled: app.extensionManager.setting.get("SmartQueue.JobCountRuleEnabled"), + max_jobs_before_pause: app.extensionManager.setting.get("SmartQueue.MaxJobsBeforePause"), + }; + try { + await fetch("/smart_queue/settings", { + method: "POST", + body: JSON.stringify(payload), + headers: { "Content-Type": "application/json" }, + }); + } catch (err) { + console.error("[Smart Queue] settings sync failed:", err); + } + } + + await syncSettingsToBackend(); + setInterval(syncSettingsToBackend, 10000); + + const enabled = app.extensionManager.setting.get("SmartQueue.EnableAutopilot"); + if (!enabled) { + console.log("[Smart Queue] Autopilot panel disabled via settings."); + return; + } + + const hasCrystools = app.extensions?.some((ext) => ext.name?.toLowerCase().includes("crystools")); + + // A small square icon button inserted as a real sibling next to + // ComfyUI's own Run button group, matching the size/style of its + // "Cancel current run" button. Because it's a genuine DOM child of + // that flex row (not a position:fixed overlay with guessed pixel + // coordinates), the browser's own flex layout makes room for it — + // no manual placement math, and nothing gets covered. + // Risk: that row is Vue-managed, so a re-render could in principle + // drop this manually-inserted node; a periodic check re-inserts it + // if that ever happens. + let manualPaused = false; + + function findToolbarRow() { + const toggle = document.querySelector('[data-testid="queue-overlay-toggle"]'); + return toggle ? toggle.parentElement : null; + } + + function ensureToolbarPauseButton() { + const row = findToolbarRow(); + if (!row) return null; + + let btn = document.getElementById("smart-queue-toolbar-btn"); + if (btn && row.contains(btn)) return btn; + + btn = document.createElement("button"); + btn.id = "smart-queue-toolbar-btn"; + btn.className = "smart-queue-toolbar-btn"; + btn.setAttribute("aria-label", "Pause queue (Smart Queue)"); + btn.setAttribute("data-pd-tooltip", "true"); + btn.title = "Pause queue (Smart Queue)"; + btn.textContent = "⏸"; + + btn.addEventListener("click", async () => { + manualPaused = !manualPaused; + try { + await fetch("/smart_queue/manual_pause", { + method: "POST", + body: JSON.stringify({ paused: manualPaused }), + headers: { "Content-Type": "application/json" }, + }); + } catch (err) { + console.error("[Smart Queue] manual pause toggle failed:", err); + } + await refreshToolbarButton(); + }); + + // Inserted right after the batch-count/Run button group (the + // row's first child) so it reads as "next to Run". + const runGroup = row.children[0]; + if (runGroup && runGroup.nextSibling) { + row.insertBefore(btn, runGroup.nextSibling); + } else { + row.insertBefore(btn, row.firstChild); + } + return btn; + } + + async function refreshToolbarButton() { + const btn = ensureToolbarPauseButton(); + if (!btn) return; + try { + const res = await fetch("/smart_queue/status"); + const data = await res.json(); + manualPaused = data.manual_paused; + btn.textContent = manualPaused ? "▶" : "⏸"; + btn.title = data.is_paused + ? `Paused — ${data.reasons.join("; ")}. Click to resume.` + : "Pause queue (Smart Queue)"; + btn.classList.toggle("smart-queue-toolbar-btn-paused", data.is_paused); + } catch (err) { + console.error("[Smart Queue] status fetch failed:", err); + } + } + + setInterval(refreshToolbarButton, 3000); + refreshToolbarButton(); + + app.extensionManager.registerSidebarTab({ + id: "smart-queue", + icon: "pi pi-pause-circle", + title: "Smart Queue", + tooltip: "Smart Queue: GPU autopilot + render queue", + type: "custom", + render: (el) => { + el.innerHTML = ` +
+
+
Smart Queue: idle
+
+ ${hasCrystools ? "" : '
'} +
Pending / running
+ +
History
+ +
+ `; + + const panel = el.querySelector("#smart-queue-panel"); + + async function refreshStatus() { + try { + const res = await fetch("/smart_queue/status"); + const data = await res.json(); + const statusEl = panel.querySelector(".smart-queue-status"); + statusEl.textContent = data.is_paused + ? `Paused — ${data.reasons.join("; ")}` + : "Running"; + statusEl.classList.toggle("smart-queue-paused", data.is_paused); + } catch (err) { + console.error("[Smart Queue] status fetch failed:", err); + } + } + + async function refreshQueueList() { + try { + const res = await fetch("/smart_queue/queue"); + const data = await res.json(); + const listEl = panel.querySelector("#smart-queue-list"); + listEl.innerHTML = ""; + if (data.items.length === 0) { + listEl.innerHTML = '
  • Nothing queued
  • '; + } + for (const item of data.items) { + const li = document.createElement("li"); + li.draggable = true; + li.dataset.promptId = item.prompt_id; + li.textContent = item.name; + listEl.appendChild(li); + } + } catch (err) { + console.error("[Smart Queue] queue fetch failed:", err); + } + } + + async function refreshHistory() { + try { + const res = await fetch("/smart_queue/history"); + const data = await res.json(); + const listEl = panel.querySelector("#smart-queue-history"); + listEl.innerHTML = ""; + if (data.items.length === 0) { + listEl.innerHTML = '
  • No history yet
  • '; + } + for (const item of data.items) { + const li = document.createElement("li"); + li.textContent = item.name; + listEl.appendChild(li); + } + } catch (err) { + console.error("[Smart Queue] history fetch failed:", err); + } + } + + const statusTimer = setInterval(refreshStatus, 3000); + const queueTimer = setInterval(refreshQueueList, 5000); + const historyTimer = setInterval(refreshHistory, 5000); + refreshStatus(); + refreshQueueList(); + refreshHistory(); + + return () => { + clearInterval(statusTimer); + clearInterval(queueTimer); + clearInterval(historyTimer); + }; + }, + }); + }, +}); diff --git a/web/smart_queue_node.js b/web/smart_queue_node.js new file mode 100644 index 0000000..ee5a83f --- /dev/null +++ b/web/smart_queue_node.js @@ -0,0 +1,59 @@ +import { app } from "../../scripts/app.js"; +import { api } from "../../scripts/api.js"; + +const SOUND_FILES = { + Default: "sounds/default.wav", + Chime: "sounds/chime.wav", + Alert: "sounds/alert.wav", +}; + +app.registerExtension({ + name: "SmartQueue.CooldownNode", + async setup() { + api.addEventListener("smart_queue.cooldown_notify", (event) => { + const { notify_sound, notify_sound_choice, custom_sound_path, notify_toast, status } = event.detail; + + if (notify_sound) { + let src = SOUND_FILES[notify_sound_choice] ?? SOUND_FILES.Default; + if (notify_sound_choice === "Custom..." && custom_sound_path) { + src = custom_sound_path; + } + const audio = new Audio(new URL(src, import.meta.url).href); + audio.play().catch(() => { + // Custom file missing/unplayable — fall back to the default tone. + new Audio(new URL(SOUND_FILES.Default, import.meta.url).href).play().catch(() => {}); + }); + } + + if (notify_toast && app.extensionManager?.toast) { + app.extensionManager.toast.add({ severity: "info", summary: "Smart Cooldown & Pause", detail: status }); + } + }); + + api.addEventListener("smart_queue.cooldown_wait_for_click", (event) => { + const { prompt_id } = event.detail; + + const overlay = document.createElement("div"); + overlay.className = "smart-queue-continue-overlay"; + overlay.innerHTML = ` +
    +
    Smart Cooldown & Pause is waiting — click Continue to resume.
    + +
    + `; + document.body.appendChild(overlay); + + const cleanup = () => overlay.remove(); + + overlay.querySelector(".smart-queue-continue-btn").addEventListener("click", async () => { + try { + await fetch(`/smart_queue/continue/${encodeURIComponent(prompt_id)}`, { method: "POST" }); + } catch (err) { + console.error("[Smart Queue] continue request failed:", err); + } finally { + cleanup(); + } + }); + }); + }, +}); diff --git a/web/sounds/alert.wav b/web/sounds/alert.wav new file mode 100644 index 0000000..42c6109 Binary files /dev/null and b/web/sounds/alert.wav differ diff --git a/web/sounds/chime.wav b/web/sounds/chime.wav new file mode 100644 index 0000000..fbff028 Binary files /dev/null and b/web/sounds/chime.wav differ diff --git a/web/sounds/default.wav b/web/sounds/default.wav new file mode 100644 index 0000000..b001835 Binary files /dev/null and b/web/sounds/default.wav differ