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.

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")

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())