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.
This commit is contained in:
@@ -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
|
||||
|
||||
+1
-1
@@ -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 |
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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())
|
||||
|
||||
@@ -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
|
||||
|
||||
+35
-4
@@ -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")
|
||||
|
||||
@@ -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())
|
||||
|
||||
@@ -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);
|
||||
|
||||
+46
-2
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user