Edouard Durand
3f1ddfed7f
Chat is an append-only chat.jsonl paged by sequence number: the welcome carries the latest 50 messages and GET /rigshare/api/chat pages further back. The panel loads older pages as you scroll up, keeps your place, and adds date dividers. ComfyUI reuses one slot element between custom sidebar tabs, so each tab now replaces the slot's content instead of appending to it.
327 lines
19 KiB
Python
327 lines
19 KiB
Python
"""End-to-end: HTTP API, login gate, websocket hub, folders over ComfyUI's API."""
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
from urllib.parse import quote
|
|
|
|
from helpers import PASSWORD, enc, rig_server
|
|
|
|
BASE_DOC = {"last_node_id": 2, "last_link_id": 0, "nodes": [
|
|
{"id": 1, "type": "A", "inputs": [], "outputs": [{"name": "o", "links": []}]},
|
|
{"id": 2, "type": "B", "inputs": [{"name": "i", "link": None}, {"name": "j", "link": None}], "outputs": []}],
|
|
"links": [], "groups": [], "extra": {}, "version": 0.4}
|
|
|
|
|
|
def test_login_gate_and_setup(run):
|
|
async def scenario():
|
|
async with rig_server() as rig:
|
|
h = rig.http
|
|
r = await h.get(rig.url("/"), headers={"Accept": "text/html"}, allow_redirects=False)
|
|
assert r.status == 302 and "/rigshare/login" in r.headers["Location"]
|
|
assert (await h.get(rig.url("/api/prompt"))).status == 401
|
|
status = await (await h.get(rig.url("/rigshare/api/status"))).json()
|
|
assert status["needs_setup"]
|
|
page = await (await h.get(rig.url("/rigshare/login"))).text()
|
|
assert "const SETUP = true" in page
|
|
await rig.setup_admin()
|
|
r = await h.post(rig.url("/rigshare/api/setup"), json={"username": "evil", "password": PASSWORD})
|
|
assert r.status == 409, "setup works only once"
|
|
r = await h.post(rig.url("/rigshare/api/login"), json={"username": "boss", "password": "nope-nope"})
|
|
assert r.status == 401
|
|
run(scenario())
|
|
|
|
|
|
def test_cookie_session_and_logout(run):
|
|
async def scenario():
|
|
async with rig_server() as rig:
|
|
await rig.setup_admin()
|
|
r = await rig.http.post(rig.url("/rigshare/api/login"), json={"username": "boss", "password": PASSWORD})
|
|
token = (await r.json())["token"]
|
|
assert "rigshare_session" in r.cookies
|
|
assert (await rig.http.get(rig.url("/api/prompt"))).status == 200, "cookie authenticates"
|
|
await rig.http.post(rig.url("/rigshare/api/logout"), headers={"Authorization": f"Bearer {token}"})
|
|
assert (await rig.http.get(rig.url("/api/prompt"), headers={"Authorization": f"Bearer {token}"})).status == 401
|
|
run(scenario())
|
|
|
|
|
|
def test_api_keys_over_http(run):
|
|
async def scenario():
|
|
async with rig_server() as rig:
|
|
admin = await rig.setup_admin()
|
|
r = await rig.http.post(rig.url("/rigshare/api/me/keys"), json={"name": "script"}, headers=admin)
|
|
body = await r.json()
|
|
key = {"Authorization": f"Bearer {body['key']}"}
|
|
assert (await rig.http.get(rig.url("/api/prompt"), headers=key)).status == 200
|
|
await rig.http.delete(rig.url(f"/rigshare/api/me/keys/{body['id']}"), headers=admin)
|
|
assert (await rig.http.get(rig.url("/api/prompt"), headers=key)).status == 401
|
|
run(scenario())
|
|
|
|
|
|
def test_permissions(run):
|
|
async def scenario():
|
|
async with rig_server() as rig:
|
|
admin = await rig.setup_admin()
|
|
viewer = await rig.add_user("viewer")
|
|
queuer = await rig.add_user("queuer", queue=True)
|
|
h = rig.http
|
|
assert (await h.post(rig.url("/api/prompt"), headers=viewer)).status == 403
|
|
assert (await h.post(rig.url("/api/prompt"), headers=queuer)).status == 200
|
|
assert (await h.post(rig.url("/api/manager/queue/install"), headers=queuer)).status == 403
|
|
assert (await h.get(rig.url("/customnode/getmappings"), headers=viewer)).status == 200, "read-only lookup"
|
|
assert (await h.post(rig.url("/api/manager/queue/install"), headers=admin)).status == 200
|
|
await h.patch(rig.url("/rigshare/api/users/queuer"), json={"perms": {"manager": True}}, headers=admin)
|
|
assert (await h.post(rig.url("/api/manager/queue/install"), headers=queuer)).status == 200
|
|
assert (await h.get(rig.url("/rigshare/api/users"), headers=viewer)).status == 403
|
|
run(scenario())
|
|
|
|
|
|
def test_template_filter_and_live_options(run, tmp_path):
|
|
templates = tmp_path / "templates"
|
|
templates.mkdir()
|
|
(templates / "index.json").write_text(json.dumps([
|
|
{"title": "Image", "templates": [{"name": "sdxl", "tags": ["Image"], "openSource": True},
|
|
{"name": "api_flux", "tags": ["API"], "openSource": False}]},
|
|
{"title": "Ads", "templates": [{"name": "x", "tags": ["API"]}]}]))
|
|
|
|
async def scenario():
|
|
async with rig_server(setup=lambda app: app.router.add_static("/templates", str(templates))) as rig:
|
|
admin = await rig.setup_admin()
|
|
alice = await rig.add_user("alice", edit=True)
|
|
sock = await rig.ws("alice")
|
|
full = await (await rig.http.get(rig.url("/templates/index.json"), headers=alice)).json()
|
|
assert sum(len(c["templates"]) for c in full) == 3
|
|
await rig.http.patch(rig.url("/rigshare/api/config"), json={"hide_api_templates": True, "hide_comfy_account": True}, headers=admin)
|
|
assert (await sock.until("server"))["server"]["hide_comfy_account"] is True
|
|
slim = await (await rig.http.get(rig.url("/templates/index.json"), headers=alice)).json()
|
|
assert [t["name"] for c in slim for t in c["templates"]] == ["sdxl"] and len(slim) == 1
|
|
await sock.close()
|
|
run(scenario())
|
|
|
|
|
|
def test_rooms_conflicts_and_folder_access(run):
|
|
async def scenario():
|
|
async with rig_server() as rig:
|
|
admin = await rig.setup_admin()
|
|
for name in ("alice", "bob", "carol"):
|
|
await rig.add_user(name, edit=True)
|
|
H = rig.headers
|
|
r = await rig.http.post(rig.url("/rigshare/api/workspace/mkdir"), json={"parent": "@shared", "name": "Team"}, headers=H["alice"])
|
|
assert r.status == 200
|
|
alice, bob, carol = [await rig.ws(n) for n in ("alice", "bob", "carol")]
|
|
room = "file:workflows/shared/Team/w.json"
|
|
await alice.send({"type": "join", "room": room, "name": "w", "doc": BASE_DOC})
|
|
joined = await alice.until("room")
|
|
assert joined["created"] and joined["manage"] and joined["folder"] == "shared/Team"
|
|
await bob.send({"type": "join", "room": room, "name": "w", "doc": BASE_DOC})
|
|
assert not (await bob.until("room"))["manage"], "only the folder owner or admins manage access"
|
|
await alice.drain(); await bob.drain()
|
|
|
|
def add(target_slot):
|
|
return {"nodes": [{"id": 3, "type": "N", "inputs": [], "outputs": [{"name": "o", "links": [1]}]}], "added": [3],
|
|
"links": {"upsert": [[1, 3, 0, 2, target_slot, "X"]], "added": [1]},
|
|
"set": {"last_node_id": 3, "last_link_id": 1}}
|
|
await alice.send({"type": "graph", "room": room, "patch": add(0)})
|
|
await bob.send({"type": "graph", "room": room, "patch": add(1)})
|
|
await asyncio.sleep(0.4)
|
|
doc = rig.hub.rooms[room].doc
|
|
assert sorted(n["id"] for n in doc["nodes"]) == [1, 2, 3, 4], "both nodes kept"
|
|
assert [s["link"] for s in doc["nodes"][1]["inputs"]] == [1, 2]
|
|
assert any(m["type"] == "doc" and m.get("reason") is None for m in await bob.drain()), "quiet resync"
|
|
|
|
# Access follows the folder: restricting it takes out whoever lost access.
|
|
r = await rig.http.patch(rig.url("/rigshare/api/workspace/folders"), json={"path": "shared/Team", "restricted": True, "members": {"carol": "view"}}, headers=H["bob"])
|
|
assert r.status == 403, "only the folder owner or an admin"
|
|
r = await rig.http.patch(rig.url("/rigshare/api/workspace/folders"), json={"path": "shared/Team", "restricted": True, "members": {"carol": "view"}}, headers=H["alice"])
|
|
assert r.status == 200
|
|
msgs = await bob.drain()
|
|
assert any(m["type"] == "room_denied" for m in msgs)
|
|
presence = [m for m in msgs if m["type"] == "presence"][-1]
|
|
assert all(r["key"] != room for r in presence["rooms"])
|
|
assert (await alice.until("room_info"))["restricted"] is True
|
|
|
|
await carol.send({"type": "join", "room": room, "name": "w", "doc": BASE_DOC})
|
|
assert (await carol.until("room"))["role"] == "view"
|
|
await carol.drain()
|
|
await carol.send({"type": "graph", "room": room, "patch": {"nodes": [dict(BASE_DOC["nodes"][0], pos=[9, 9])]}})
|
|
assert "view" in ((await carol.until("doc")).get("reason") or "")
|
|
await rig.http.patch(rig.url("/rigshare/api/workspace/folders"), json={"path": "shared/Team", "restricted": True, "members": {"carol": "edit"}}, headers=admin)
|
|
assert (await carol.until("room_info"))["role"] == "edit"
|
|
assert (await rig.http.put(rig.url("/rigshare/api/room/acl"), json={"key": room}, headers=admin)).status in (404, 405), "no per-workflow lists"
|
|
for s in (alice, bob, carol):
|
|
await s.close()
|
|
run(scenario())
|
|
|
|
|
|
def test_folders_over_comfy_api(run):
|
|
async def scenario():
|
|
async with rig_server() as rig:
|
|
await rig.setup_admin()
|
|
for name, perms in (("alice", {"edit": True}), ("bob", {"edit": True}), ("carol", {"edit": True}), ("dave", {})):
|
|
await rig.add_user(name, **perms)
|
|
H, h = rig.headers, rig.http
|
|
|
|
async def save(user, path):
|
|
return (await h.post(rig.url(f"/api/userdata/{enc(path)}"), data=b"{}", headers=H[user])).status
|
|
|
|
async def listing(user):
|
|
r = await h.get(rig.url("/api/userdata?dir=workflows&recurse=true&full_info=true"), headers=H[user])
|
|
return {i["path"] for i in await r.json()}
|
|
|
|
assert await save("alice", "workflows/users/alice/secret.json") == 200
|
|
assert await save("dave", "workflows/users/dave/mine.json") == 200
|
|
assert await save("dave", "workflows/common.json") == 403
|
|
assert await save("bob", "workflows/users/alice/evil.json") == 403
|
|
for raw in (enc("workflows/users/alice/secret.json"),
|
|
quote(enc("workflows/users/alice/secret.json"), safe=""),
|
|
enc("workflows/shared/../users/alice/secret.json"),
|
|
"workflows/users/alice/secret.json"):
|
|
assert (await h.get(rig.url(f"/api/userdata/{raw}"), headers=H["bob"])).status in (403, 404), raw
|
|
assert "users/alice/secret.json" not in await listing("bob")
|
|
assert "users/alice/secret.json" in await listing("alice")
|
|
|
|
r = await h.post(rig.url("/rigshare/api/workspace/mkdir"), json={"parent": "@shared", "name": "Team"}, headers=H["bob"])
|
|
assert r.status == 200
|
|
assert await save("alice", "workflows/shared/Team/plan.json") == 200
|
|
await h.patch(rig.url("/rigshare/api/workspace/folders"), json={"path": "shared/Team", "restricted": True, "members": {"carol": "view"}}, headers=H["bob"])
|
|
assert not any(p.startswith("shared/Team/") for p in await listing("alice"))
|
|
assert await save("carol", "workflows/shared/Team/plan.json") == 403
|
|
assert (await h.get(rig.url(f"/api/userdata/{enc('workflows/shared/Team/plan.json')}"), headers=H["carol"])).status == 200
|
|
|
|
r = await h.post(rig.url(f"/api/userdata/{enc('workflows/shared/Team/plan.json')}/move/{enc('workflows/users/bob/plan.json')}"), headers=H["bob"])
|
|
assert r.status == 200
|
|
assert (await h.delete(rig.url(f"/api/userdata/{enc('workflows/users/alice/secret.json')}"), headers=H["bob"])).status == 403
|
|
assert (await h.delete(rig.url(f"/api/userdata/{enc('workflows/users/alice/secret.json')}"), headers=H["alice"])).status == 204
|
|
|
|
sock = await rig.ws("bob")
|
|
await sock.send({"type": "join", "room": "file:workflows/users/alice/other.json", "name": "o", "doc": {"nodes": []}})
|
|
assert (await sock.until("room_denied"))["room"].endswith("alice/other.json")
|
|
await sock.close()
|
|
|
|
await save("alice", "workflows/users/alice/keep.json")
|
|
await h.patch(rig.url("/rigshare/api/users/alice"), json={"username": "alicia"}, headers=H["boss"])
|
|
assert os.path.exists(os.path.join(rig.user_root, "workflows", "users", "alicia", "keep.json"))
|
|
run(scenario())
|
|
|
|
|
|
def test_live_follows_folders(run):
|
|
async def scenario():
|
|
async with rig_server() as rig:
|
|
await rig.setup_admin()
|
|
for name in ("alice", "bob"):
|
|
await rig.add_user(name, edit=True)
|
|
alice, bob = await rig.ws("alice"), await rig.ws("bob")
|
|
|
|
# Your own private folder and unsaved tabs never become live.
|
|
for key in ("file:workflows/users/alice/mine.json", "tmp:abc"):
|
|
await alice.send({"type": "join", "room": key, "name": "x", "doc": BASE_DOC})
|
|
assert (await alice.until("room_denied"))["room"] == key
|
|
assert not rig.hub.rooms
|
|
|
|
# Anything else is live while open, and only open ones are listed.
|
|
key = "file:workflows/shared/team.json"
|
|
await alice.send({"type": "join", "room": key, "name": "team", "doc": BASE_DOC})
|
|
await alice.until("room")
|
|
presence = [m for m in await bob.drain() if m["type"] == "presence"][-1]
|
|
assert [r["key"] for r in presence["rooms"]] == [key]
|
|
await alice.send({"type": "join", "room": None})
|
|
presence = [m for m in await bob.drain() if m["type"] == "presence"][-1]
|
|
assert presence["rooms"] == [], "closed workflows leave the live list"
|
|
assert key in rig.hub.rooms, "its state is kept for the next person who opens it"
|
|
|
|
# Moving it into a private folder takes it out of live sharing.
|
|
path = "workflows/shared/team.json"
|
|
await rig.http.post(rig.url(f"/api/userdata/{enc(path)}"), data=b"{}", headers=rig.headers["alice"])
|
|
await bob.send({"type": "join", "room": key, "name": "team"})
|
|
await bob.until("room")
|
|
r = await rig.http.post(rig.url(f"/api/userdata/{enc(path)}/move/{enc('workflows/users/alice/team.json')}"), headers=rig.headers["boss"])
|
|
assert r.status == 200
|
|
assert not rig.hub.rooms
|
|
|
|
# Deleting a live file closes it for whoever has it open.
|
|
key2, path2 = "file:workflows/gone.json", "workflows/gone.json"
|
|
await rig.http.post(rig.url(f"/api/userdata/{enc(path2)}"), data=b"{}", headers=rig.headers["alice"])
|
|
await bob.send({"type": "join", "room": key2, "name": "gone", "doc": BASE_DOC})
|
|
await bob.until("room")
|
|
assert (await rig.http.delete(rig.url(f"/api/userdata/{enc(path2)}"), headers=rig.headers["boss"])).status == 204
|
|
assert (await bob.until("room_closed"))["room"] == key2
|
|
assert key2 not in rig.hub.rooms
|
|
await alice.close(); await bob.close()
|
|
run(scenario())
|
|
|
|
|
|
def test_moves_follow_open_tabs(run):
|
|
async def scenario():
|
|
async with rig_server() as rig:
|
|
await rig.setup_admin()
|
|
for name in ("alice", "bob", "carol", "dave"):
|
|
await rig.add_user(name, edit=True)
|
|
H, h = rig.headers, rig.http
|
|
socks = {n: await rig.ws(n) for n in ("alice", "bob", "carol", "dave")}
|
|
|
|
async def moves(name):
|
|
return [m for m in await socks[name].drain() if m["type"] in ("path_moved", "path_closed")]
|
|
|
|
# A live file renamed: everyone who could open it is told, and its room follows.
|
|
await h.post(rig.url(f"/api/userdata/{enc('workflows/a.json')}"), data=b"{}", headers=H["alice"])
|
|
await socks["bob"].send({"type": "join", "room": "file:workflows/a.json", "name": "a", "doc": BASE_DOC})
|
|
await socks["bob"].until("room")
|
|
for sock in socks.values():
|
|
await sock.drain()
|
|
r = await h.post(rig.url(f"/api/userdata/{enc('workflows/a.json')}/move/{enc('workflows/b.json')}"), headers=H["alice"])
|
|
assert r.status == 200
|
|
assert await moves("bob") == [{"type": "path_moved", "from": "workflows/a.json", "to": "workflows/b.json"}]
|
|
assert "file:workflows/b.json" in rig.hub.rooms and "file:workflows/a.json" not in rig.hub.rooms
|
|
|
|
# Into a private folder: the owner follows it, the others keep a private copy.
|
|
r = await h.post(rig.url(f"/api/userdata/{enc('workflows/b.json')}/move/{enc('workflows/users/alice/b.json')}"), headers=H["boss"])
|
|
assert r.status == 200
|
|
assert (await moves("alice"))[0]["type"] == "path_moved"
|
|
assert await moves("bob") == [{"type": "path_closed", "from": "workflows/b.json"}]
|
|
assert not rig.hub.rooms
|
|
|
|
# A restricted shared folder renamed: people without access learn nothing.
|
|
await h.post(rig.url("/rigshare/api/workspace/mkdir"), json={"parent": "@shared", "name": "Team"}, headers=H["alice"])
|
|
await h.patch(rig.url("/rigshare/api/workspace/folders"), json={"path": "shared/Team", "restricted": True, "members": {"carol": "view"}}, headers=H["alice"])
|
|
for sock in socks.values():
|
|
await sock.drain()
|
|
r = await h.patch(rig.url("/rigshare/api/workspace/folders"), json={"path": "shared/Team", "name": "Crew"}, headers=H["alice"])
|
|
assert r.status == 200
|
|
assert await moves("carol") == [{"type": "path_moved", "from": "workflows/shared/Team", "to": "workflows/shared/Crew"}]
|
|
assert await moves("dave") == []
|
|
for sock in socks.values():
|
|
await sock.close()
|
|
run(scenario())
|
|
|
|
|
|
def test_chat_history_pages(run):
|
|
async def scenario():
|
|
async with rig_server() as rig:
|
|
await rig.setup_admin()
|
|
sock = await rig.ws("boss")
|
|
for i in range(70):
|
|
await sock.send({"type": "chat", "text": f"hello {i}"})
|
|
await asyncio.sleep(0.3)
|
|
await sock.close()
|
|
sock = await rig.ws("boss")
|
|
await sock.close()
|
|
# The welcome carries only the newest page; the rest comes over HTTP.
|
|
h = rig.headers["boss"]
|
|
session = rig.http
|
|
ws = await session.ws_connect(rig.url("/rigshare/ws"), headers=h)
|
|
await ws.send_json({"type": "hello"})
|
|
while True:
|
|
welcome = json.loads((await ws.receive()).data)
|
|
if welcome["type"] == "welcome":
|
|
break
|
|
await ws.close()
|
|
assert len(welcome["chat"]) == 50 and welcome["chat_more"]
|
|
oldest = welcome["chat"][0]["seq"]
|
|
older = await (await session.get(rig.url(f"/rigshare/api/chat?before={oldest}&limit=50"), headers=h)).json()
|
|
texts = [m["text"] for m in older["messages"]]
|
|
assert not older["more"] and older["messages"][0]["seq"] == 0
|
|
seqs = [m["seq"] for m in older["messages"] + welcome["chat"]]
|
|
assert "hello 0" in texts and seqs == list(range(len(seqs))), "no gaps, no repeats"
|
|
assert (await session.get(rig.url("/rigshare/api/chat?before=x"))).status in (400, 401)
|
|
run(scenario())
|