From 531fec2ec9008bb5fa51ffd95b15a11abf3ff766 Mon Sep 17 00:00:00 2001 From: Edouard Durand Date: Sat, 26 Sep 2026 08:20:44 +0000 Subject: [PATCH] Queue tab: running and pending prompts with who queued them The server records who queued each prompt (and the tab's name, sent by the client) and publishes queue items with the stats. Cancelling someone else's prompt or clearing the queue needs Admin. Tabs: Live, Chat, People, Queue, Server, Admin. --- CHANGELOG.md | 2 ++ docs/accounts.md | 2 +- docs/sharing.md | 2 ++ rigshare_core/hub.py | 31 ++++++++++++++++++++++++ rigshare_core/routes.py | 39 ++++++++++++++++++++++++++++++ rigshare_core/stats.py | 18 ++++++++++++++ tests/python/helpers.py | 39 ++++++++++++++++++++++++++---- tests/python/test_server.py | 32 +++++++++++++++++++++++++ web/rigshare/main.js | 12 +++++++--- web/rigshare/panel.js | 48 +++++++++++++++++++++++++++++++++++-- 10 files changed, 215 insertions(+), 10 deletions(-) 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;