diff --git a/CHANGELOG.md b/CHANGELOG.md index 17751ec..4c4f5a7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ ### Added +- **Queue** tab: what is running and waiting, with the workflow's name and who queued it. Cancel your own runs from there; admins can cancel anyone's. The server enforces it too: cancelling someone else's prompt, or clearing the whole queue, needs *Admin* +- Panel tabs in a new order: Live, Chat, People, Queue, Server, Admin - Admins can clear the whole chat history (*Admin → Chat*, or `DELETE /rigshare/api/chat`). Everyone's chat empties at once, followed by a notice saying who cleared it ### Fixed diff --git a/docs/accounts.md b/docs/accounts.md index c18e85c..196e544 100644 --- a/docs/accounts.md +++ b/docs/accounts.md @@ -6,7 +6,7 @@ |---|---| | *(none)*, a viewer | Watch shared workflows, chat, follow people, open snapshots as private copies | | **Edit** | Change shared workflows, save or rename workflow files, take and restore snapshots, upload files | -| **Queue** | Queue prompts, interrupt, free memory | +| **Queue** | Queue prompts, stop or remove your own runs, free memory. Stopping or removing someone else's, or clearing the queue, needs *Admin* | | **Manager** | Open ComfyUI Manager: install, update or remove custom nodes and models, snapshots, restart | | **Admin** | Everything, plus deleting workflow files, accounts and server settings | diff --git a/docs/sharing.md b/docs/sharing.md index b412d6f..c63daca 100644 --- a/docs/sharing.md +++ b/docs/sharing.md @@ -19,6 +19,8 @@ Only the tab on screen is live. Switch to another tab and you leave the room; sw **Live** shows the tab on screen: whether it is live, who else is in it, and buttons for **Snapshot** and **History**. For a private or unsaved tab it says why it is not live. Below that is the list of **live workflows**: those someone has open right now, with their folder and the people in each one. Open any of them with the arrow. The Files tab marks live files with a green dot. +**Queue** shows the run in progress and the ones waiting, each with the workflow's name, who queued it and when. The button on a row stops or removes it: your own runs, or anyone's for an admin. + **People** lists who is online and which workflow they are on. The eye button *follows* someone: it jumps to their tab and keeps your view on theirs until you click or scroll the canvas. ## What syncs, and how diff --git a/rigshare_core/hub.py b/rigshare_core/hub.py index 6470e96..f94e3da 100644 --- a/rigshare_core/hub.py +++ b/rigshare_core/hub.py @@ -15,6 +15,7 @@ import json import logging import secrets import time +from collections import OrderedDict from aiohttp import web, WSMsgType @@ -106,6 +107,7 @@ class Hub: self.stats = None self._task = None self.workspace = None # set by __init__ (folders, private spaces) + self.queue_owners = OrderedDict() # prompt_id -> {user, name, workflow, at} # ----- lifecycle ---------------------------------------------------- @@ -142,10 +144,39 @@ class Hub: continue try: self.stats = await loop.run_in_executor(None, collect_stats, self.server) + self.stats["queue_items"] = self.with_owners(self.stats.get("queue_items")) await self.broadcast({"type": "stats", "stats": self.stats}) except Exception as e: log.debug(f"[RigShare] stats error: {e}") + # ----- queue ---------------------------------------------------------- + + def record_prompt(self, prompt_id, user, workflow): + """Remember who queued a prompt (ComfyUI itself doesn't).""" + if not prompt_id: + return + self.queue_owners[prompt_id] = { + "user": user["username"] if user else None, + "name": (user.get("display_name") or user["username"]) if user else None, + "workflow": clean_text(workflow, 120) or None, "at": int(time.time())} + while len(self.queue_owners) > 2000: + self.queue_owners.popitem(last=False) + + def owner_of(self, prompt_id): + return (self.queue_owners.get(prompt_id) or {}).get("user") + + def with_owners(self, items): + if items is None: + return None + return [{**it, **(self.queue_owners.get(it["id"]) or {})} for it in items] + + def running_prompt_ids(self): + try: + running, _ = self.server.prompt_queue.get_current_queue_volatile() + return [entry[1] for entry in running] + except Exception: + return [] + def snapshot(self, room, author, label=None): if room.doc is None: return None diff --git a/rigshare_core/routes.py b/rigshare_core/routes.py index f1cc921..7f9aa46 100644 --- a/rigshare_core/routes.py +++ b/rigshare_core/routes.py @@ -8,6 +8,7 @@ import json import logging import os import time +from urllib.parse import unquote from collections import defaultdict from aiohttp import web @@ -616,6 +617,34 @@ def setup(server, store, hub): action = {"GET": "get", "HEAD": "get", "POST": "save", "DELETE": "delete"}.get(request.method) return (action, src, None) if action else None + def api_path(request): + return request.path[4:] if request.path.startswith("/api/") else request.path + + async def check_queue_ownership(request): + """Cancelling a queued prompt or stopping the running one: your own, unless admin.""" + if request.method != "POST" or api_path(request) not in ("/queue", "/interrupt"): + return None + if trusted(request.remote or "") or perms_for(request).get("admin"): + return None + user = current_user(request) + me = user["username"] if user else None + try: + body = await request.json() if request.body_exists else {} + except ValueError: + body = {} + body = body if isinstance(body, dict) else {} + if api_path(request) == "/queue": + if body.get("clear"): + return "only admins can clear the whole queue" + ids = body.get("delete") or [] + else: + ids = [body["prompt_id"]] if body.get("prompt_id") else hub.running_prompt_ids() + for prompt_id in ids: + owner = hub.owner_of(prompt_id) + if owner and owner != me: + return "you can only cancel your own prompts" + return None + def snapshot_saved_file(request, path, body): if not path.endswith(".json") or Workspace.workflows_rel(path) is None: return @@ -726,9 +755,19 @@ def setup(server, store, hub): "details": "", "extra_info": {}}, "node_errors": {}}, status=403) break + refusal = await check_queue_ownership(request) + if refusal: + return web.json_response({"error": f"RigShare: {refusal}"}, status=403) if userdata and userdata[0] == "save": await request.read() # keep the body for the snapshot below (aiohttp caches it) response = await handler(request) + # Remember who queued each prompt, for the Queue tab. + if request.method == "POST" and api_path(request) == "/prompt" and getattr(response, "status", 0) == 200: + try: + prompt_id = json.loads(response.body).get("prompt_id") + except (ValueError, TypeError, AttributeError): + prompt_id = None + hub.record_prompt(prompt_id, current_user(request), unquote(request.headers.get("X-RigShare-Workflow", ""))) # Every save of a workflow keeps a snapshot, named after when it was saved. if userdata and userdata[0] == "save" and getattr(response, "status", 0) == 200: snapshot_saved_file(request, userdata[1], await request.read()) diff --git a/rigshare_core/stats.py b/rigshare_core/stats.py index b53b21b..4a3720e 100644 --- a/rigshare_core/stats.py +++ b/rigshare_core/stats.py @@ -51,6 +51,23 @@ def _gpus(): return out +def queue_items(server): + """Running and pending prompts: [{id, number, status, nodes}], running first.""" + try: + running, pending = server.prompt_queue.get_current_queue_volatile() + except Exception: + return None + items = [] + for status, entries in (("running", running), ("pending", sorted(pending, key=lambda e: e[0]))): + for entry in entries: + try: + prompt = entry[2] if isinstance(entry[2], dict) else {} + items.append({"id": entry[1], "number": entry[0], "status": status, "nodes": len(prompt)}) + except (IndexError, TypeError): + continue + return items + + def _queue(server): try: running, pending = server.prompt_queue.get_current_queue_volatile() @@ -76,4 +93,5 @@ def collect_stats(server=None): stats["gpus"] = [] if server is not None: stats["queue"] = _queue(server) + stats["queue_items"] = queue_items(server) return stats diff --git a/tests/python/helpers.py b/tests/python/helpers.py index c046ab0..353d405 100644 --- a/tests/python/helpers.py +++ b/tests/python/helpers.py @@ -33,17 +33,28 @@ def enc(path): return quote(path, safe="") +class FakeQueue: + """ComfyUI's PromptQueue, as far as RigShare reads it.""" + + def __init__(self): + self.running, self.pending, self.number = [], [], 0 + + def get_current_queue_volatile(self): + return list(self.running), list(self.pending) + + class FakePromptServer: def __init__(self): self.app = web.Application() self.sent = [] + self.prompt_queue = FakeQueue() def send_sync(self, event, data, sid=None): self.sent.append((event, sid)) -def comfy_routes(user_root): - """Minimal ComfyUI /userdata routes for the single default user.""" +def comfy_routes(user_root, queue=None): + """Minimal ComfyUI /userdata routes for the single default user, plus the queue.""" def resolve(name): if "%" in name: @@ -83,6 +94,24 @@ def comfy_routes(user_root): async def ok(request): return web.json_response({"ok": True}) + async def prompt(request): + if request.method != "POST" or queue is None: + return web.json_response({"ok": True}) + body = await request.json() if request.body_exists else {} + queue.number += 1 + prompt_id = f"p{queue.number}" + queue.pending.append((queue.number, prompt_id, (body or {}).get("prompt") or {"1": {}}, {}, [])) + return web.json_response({"prompt_id": prompt_id, "number": queue.number, "node_errors": {}}) + + async def queue_route(request): + body = await request.json() if request.body_exists else {} + if queue is not None: + gone = set((body or {}).get("delete") or []) + queue.pending = [e for e in queue.pending if e[1] not in gone] + if (body or {}).get("clear"): + queue.pending = [] + return web.Response(status=200) + table = web.RouteTableDef() for prefix in ("", "/api"): table.get(prefix + "/userdata")(listing) @@ -90,7 +119,9 @@ def comfy_routes(user_root): table.post(prefix + "/userdata/{file}")(post) table.delete(prefix + "/userdata/{file}")(delete) table.post(prefix + "/userdata/{file}/move/{dest}")(move) - table.route("*", prefix + "/prompt")(ok) + table.route("*", prefix + "/prompt")(prompt) + table.post(prefix + "/queue")(queue_route) + table.post(prefix + "/interrupt")(ok) table.route("*", prefix + "/manager/{tail:.*}")(ok) table.route("*", prefix + "/customnode/{tail:.*}")(ok) return table @@ -184,7 +215,7 @@ async def rig_server(config=None, workspace=True, setup=None): if workspace: hub.workspace = Workspace(store, os.path.join(user_root, "workflows")) routes.setup(ps, store, hub) - ps.app.add_routes(comfy_routes(user_root)) + ps.app.add_routes(comfy_routes(user_root, ps.prompt_queue)) if setup: setup(ps.app) server = TestServer(ps.app, host="127.0.0.1") diff --git a/tests/python/test_server.py b/tests/python/test_server.py index 0259a52..80c4723 100644 --- a/tests/python/test_server.py +++ b/tests/python/test_server.py @@ -463,3 +463,35 @@ def test_admin_clears_chat_history(run): assert [m["text"] for m in history["messages"]] == [notice["text"]] and not history["more"] await sock.close() run(scenario()) + + +def test_queue_shows_who_launched_and_guards_cancelling(run): + async def scenario(): + async with rig_server() as rig: + admin = await rig.setup_admin() + alice = await rig.add_user("alice", queue=True) + bob = await rig.add_user("bob", queue=True) + h = rig.http + + async def queue(user, name): + r = await h.post(rig.url("/api/prompt"), json={"prompt": {"1": {}, "2": {}}}, + headers={**user, "X-RigShare-Workflow": quote(name)}) + return (await r.json())["prompt_id"] + a = await queue(alice, "Face swap é") + b = await queue(bob, "Upscale") + rig.ps.prompt_queue.running.append(rig.ps.prompt_queue.pending.pop(0)) # alice's starts + + from rigshare_core.stats import queue_items + items = rig.hub.with_owners(queue_items(rig.ps)) + assert [(i["status"], i["user"], i["workflow"], i["nodes"]) for i in items] == [ + ("running", "alice", "Face swap é", 2), ("pending", "bob", "Upscale", 2)] + + # Your own runs only, unless admin; clearing everyone's is admin-only. + assert (await h.post(rig.url("/api/queue"), json={"delete": [b]}, headers=alice)).status == 403 + assert (await h.post(rig.url("/api/interrupt"), json={}, headers=bob)).status == 403, "alice's is running" + assert (await h.post(rig.url("/api/interrupt"), json={"prompt_id": a}, headers=alice)).status == 200 + assert (await h.post(rig.url("/api/queue"), json={"clear": True}, headers=bob)).status == 403 + assert (await h.post(rig.url("/api/queue"), json={"delete": [b]}, headers=bob)).status == 200 + assert not rig.ps.prompt_queue.pending + assert (await h.post(rig.url("/api/queue"), json={"clear": True}, headers=admin)).status == 200 + run(scenario()) diff --git a/web/rigshare/main.js b/web/rigshare/main.js index 9cabf9f..71b27b1 100644 --- a/web/rigshare/main.js +++ b/web/rigshare/main.js @@ -51,12 +51,18 @@ function setting(id, fallback) { } } -function installAuthHeaders(client) { +function installAuthHeaders(client, sync) { const original = api.fetchApi.bind(api); api.fetchApi = (route, options = {}) => { let headers = options.headers || {}; if (headers instanceof Headers) headers = Object.fromEntries(headers.entries()); - return original(route, { ...options, headers: { ...headers, ...client.authHeaders() } }); + headers = { ...headers, ...client.authHeaders() }; + // Queueing: name the workflow, so the Queue tab can show what each run is. + if (route === "/prompt" && (options.method || "GET").toUpperCase() === "POST") { + const wf = sync.store?.activeWorkflow; + if (wf) headers["X-RigShare-Workflow"] = encodeURIComponent(sync.tabName(wf)); + } + return original(route, { ...options, headers }); }; } @@ -297,7 +303,7 @@ app.registerExtension({ panel.chatToasts = setting("RigShare.ChatToasts", false); client.on("toast", (t) => app.extensionManager?.toast?.add({ life: 4000, ...t })); - installAuthHeaders(client); + installAuthHeaders(client, sync); installQueueGuard(client); installManagerGuard(client); installComfyAccountGuard(client); diff --git a/web/rigshare/panel.js b/web/rigshare/panel.js index fc456e3..0887517 100644 --- a/web/rigshare/panel.js +++ b/web/rigshare/panel.js @@ -10,8 +10,9 @@ const ICON_URL = new URL("./assets/icon.svg", import.meta.url).href; const VIEWS = [ { id: "workflows", icon: "pi-share-alt", label: "Live" }, - { id: "people", icon: "pi-users", label: "People" }, { id: "chat", icon: "pi-comments", label: "Chat" }, + { id: "people", icon: "pi-users", label: "People" }, + { id: "queue", icon: "pi-list", label: "Queue" }, { id: "server", icon: "pi-server", label: "Server" }, { id: "admin", icon: "pi-shield", label: "Admin", admin: true }, ]; @@ -60,7 +61,10 @@ export class RigSharePanel { client.on("status", () => this.render()); client.on("self", () => this.render()); client.on("presence", () => this.renderSoft()); - client.on("stats", () => { if (this.view === "server") this.renderBody(); }); + client.on("stats", () => { + if (this.view === "server" || this.view === "queue") this.renderBody(); + this.renderNav(); // queue count + }); client.on("chat-history", () => this.renderChatLog()); client.on("chat", (m) => this.onChat(m)); sync.addEventListener("state", () => { if (this.view === "workflows") this.renderBody(); }); @@ -171,6 +175,7 @@ export class RigSharePanel { people: c.users.length || null, chat: this.unread || null, workflows: c.rooms?.length || null, + queue: c.stats?.queue_items?.length || null, }; this.nav.replaceChildren(...VIEWS.filter((v) => !v.admin || c.perms.admin).map((v) => h("button", { @@ -199,6 +204,7 @@ export class RigSharePanel { people: () => this.peopleView(), chat: () => [this.chatView], server: () => this.serverView(), + queue: () => this.queueView(), admin: () => this.adminView(), account: () => this.accountView(), }; @@ -566,6 +572,44 @@ export class RigSharePanel { // ----- server --------------------------------------------------------- + // ----- queue ---------------------------------------------------------- + + queueView() { + const c = this.client; + const items = c.stats?.queue_items; + if (items == null) return [this.section("pi-list", "Queue", h("p", { class: "rs-muted" }, "Waiting for the server…"))]; + const running = items.filter((it) => it.status === "running"); + const pending = items.filter((it) => it.status === "pending"); + return [ + this.section("pi-play", "Running", ...(running.length ? running.map((it) => this.queueRow(it)) : [h("p", { class: "rs-muted" }, "Nothing running.")])), + this.section("pi-list", `Pending${pending.length ? ` (${pending.length})` : ""}`, + ...(pending.length ? pending.map((it, i) => this.queueRow(it, i + 1)) : [h("p", { class: "rs-muted" }, "Nothing waiting.")])), + ]; + } + + queueRow(it, position) { + const c = this.client; + const online = it.user ? c.users.find((u) => u.key === it.user) : null; + const who = it.name || (it.user ? `@${it.user}` : "Unknown"); + const mine = it.user && it.user === c.self?.key; + const canCancel = c.perms.queue && (mine || c.perms.admin); + const running = it.status === "running"; + const meta = [position ? `#${position}` : null, `${it.nodes} nodes`, it.at ? `queued ${timeAgo(it.at * 1000)}` : null].filter(Boolean).join(" · "); + return h("div", { class: `rs-row-item rs-queue-row ${running ? "running" : ""}` }, + h("span", { class: "rs-avatar", style: `background:${online?.color ?? "var(--rs-border)"}`, title: who }, initials(who)), + h("div", { class: "rs-grow rs-min0" }, + h("div", { class: "rs-name rs-ellipsis" }, it.workflow || "Workflow", + running ? h("i", { class: "pi pi-spin pi-spinner rs-inline-icon rs-muted", title: "Running" }) : null), + h("div", { class: "rs-muted rs-small rs-ellipsis" }, `${who}${mine ? " (you)" : ""}`), + h("div", { class: "rs-muted rs-small" }, meta)), + canCancel ? h("button", { + class: "rs-btn rs-btn-icon rs-danger", title: running ? "Stop this run" : "Remove from the queue", + onclick: () => this.run(() => running + ? c.request("POST", "/api/interrupt", { prompt_id: it.id }) + : c.request("POST", "/api/queue", { delete: [it.id] }), running ? "Stopping…" : "Removed from the queue"), + }, h("i", { class: `pi ${running ? "pi-stop" : "pi-times"}` })) : null); + } + serverView() { const st = this.client.stats; const floating = this.widget.enabled;