Jaroslav Beneš
5e75948069
The last unbuilt capability, and built the way CLAUDE.md said it had to be: a
ToolDef reaching resolve_tools plus a permission and a capability flag, not a new
code path. The only genuinely new UI is one branch in the transcript.
services/images/ is three modules. comfy.py speaks HTTP -- submit, poll /history,
fetch the PNG, /free, and an /object_info discovery for the admin page only.
Polled and not socketed, because holding a connection open for the length of a
generation is the live-connection state the whole ssh.py design forbids, and the
thing being waited for takes tens of seconds anyway. The base URL is exempt from
the SSRF guard by construction, exactly as Connection.base_url and the audio
endpoints are -- said out loud in the docstring, because a default of
127.0.0.1:8188 is precisely the shape that guard exists to refuse and therefore
reads as a hole rather than a decision.
workflow.py fills a template, and the one thing that matters is that it walks the
parsed JSON rather than the text of it. A value that is exactly "{{steps}}"
becomes the number 20; ComfyUI validates types and refuses the string. A
placeholder inside a longer string is still text, which is what makes
"{{prompt}}, masterpiece" work -- and text substitution would additionally mean a
prompt containing a quotation mark produced a document that no longer parses, on
the one input guaranteed to hold arbitrary text. Which node holds the prompt is
the administrator's statement rather than a guess from node types: sniffing for
the first CLIPTextEncode works on the shipped workflow and on nothing else, and
swaps positive for negative the first time somebody reorders them. seed has no
fixed default, because one would make every unspecified generation identical and
make the retry loop redraw the same rejected picture four times.
tool.py is one call, one finished image. Returning every attempt to the
conversation would cost a round each, make the ceiling advisory rather than
enforced, and walk the reader past every reject -- so the reviewer lives inside
the tool and is asked about *bytes*: an attempt about to be discarded should not
leave an Attachment behind, so it sees a downscaled preview built in memory and
only the kept image is written. Anything that goes wrong in review is a keep;
losing a picture because a judging request timed out would be the check
destroying the thing it was checking. The last attempt is kept whatever the
verdict, so a request always produces something. Rejects are recorded, not
stored.
Preserve VRAM unloads the chat's own connection and nothing else, because the
memory being freed belongs to one machine: local llama-swap answers GET /unload,
and a box on the network has no reason to be unloaded when ComfyUI wants memory
here. The swap goes round the review rather than round the tool, which costs two
model loads per retry -- so the two settings are independent and the page warns
when both are on. Nothing loads the LLM back: the reply's next request does, and
that step exists in the description and not in the code, so the code says so.
Two rules elsewhere had to be drawn for the first time. message_payload sends
images only on user turns -- no assistant message had ever carried one, and the
moment one does the multimodal list form on an assistant turn is rejected by
OpenAI and most local runners, breaking every later turn in the chat. And
files.store gained keep_original, because _process_image turns anything without
alpha into JPEG q85 at 1400px: right for a phone photo, a visible loss on the one
output this feature exists to produce.
/image sends the ordinary message with force_tool, which becomes tool_choice for
the first round only -- left in place the reply would draw a picture, be asked
again, and draw another. FORCEABLE_TOOLS is an allow list because the name is
read off a form.
ToolContext gained chat_id, and that fixed a tool nobody had ever successfully
run: _run_scratch_write read context.chat_id on a dataclass with no such field,
so every call raised AttributeError, swallowed by run_tool's blanket except into
"the scratch_write tool failed" -- indistinguishable from a model calling it
wrongly. The test that existed asserted the family and the risk, which are
properties of the declaration rather than of the code.
Verified against the real ComfyUI 0.27.0 on this machine rather than against
documentation: every endpoint shape here was read off it, a generation ran end to
end through the client, the reviewer was shown a matching and a mismatched prompt
and answered KEEP and RETRY correctly, and the unload hook fired for the local
llama-swap and not for the remote box.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
273 lines
9.1 KiB
Python
273 lines
9.1 KiB
Python
"""The image generation admin page, and the workflows behind it."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
from sqlalchemy import select
|
|
|
|
import lembas
|
|
from lembas.db.models import ImageWorkflow
|
|
from lembas.services import settings_store
|
|
from lembas.services.crypto import UNCHANGED_SENTINEL, decrypt
|
|
|
|
BASE_TEXT = (Path(lembas.__file__).parent / "services/images/base_workflow.json").read_text(
|
|
encoding="utf-8"
|
|
)
|
|
BASE = json.loads(BASE_TEXT)
|
|
|
|
|
|
@pytest.fixture
|
|
def admin(client: TestClient, db, registered):
|
|
from lembas.db.models import User
|
|
|
|
db.query(User).update({"role": "admin"})
|
|
db.commit()
|
|
return client
|
|
|
|
|
|
def _settings(client, **overrides):
|
|
data = {
|
|
"enabled": "true",
|
|
"base_url": "http://comfy.test:8188",
|
|
"api_key": "",
|
|
"timeout": "600",
|
|
"checkpoints": "sd.safetensors\nother.safetensors",
|
|
"default_workflow_id": "",
|
|
"review_enabled": "true",
|
|
"review_model_id": "",
|
|
"max_tries": "4",
|
|
"preserve_vram": "",
|
|
"instructions": "",
|
|
}
|
|
data.update(overrides)
|
|
return client.post("/admin/images", data=data, follow_redirects=False)
|
|
|
|
|
|
# --- The page ------------------------------------------------------------------
|
|
def test_the_page_renders_and_is_in_the_nav(admin):
|
|
page = admin.get("/admin/images").text
|
|
assert 'href="/admin/images"' in page
|
|
assert "ComfyUI" in page
|
|
|
|
|
|
def test_settings_are_saved(admin, db):
|
|
assert _settings(admin).status_code == 303
|
|
|
|
values = settings_store.images(db)
|
|
assert values["enabled"] is True
|
|
assert values["base_url"] == "http://comfy.test:8188"
|
|
assert values["checkpoints"] == ["sd.safetensors", "other.safetensors"]
|
|
assert values["review_enabled"] is True
|
|
|
|
|
|
def test_a_trailing_slash_is_stripped_from_the_url(admin, db):
|
|
_settings(admin, base_url="http://comfy.test:8188/")
|
|
assert settings_store.images(db)["base_url"] == "http://comfy.test:8188"
|
|
|
|
|
|
def test_an_absent_checkbox_is_off(admin, db):
|
|
"""An unticked box is simply missing from a form post -- that absence *is*
|
|
the off signal."""
|
|
data = {"base_url": "http://x", "timeout": "600", "max_tries": "4"}
|
|
admin.post("/admin/images", data=data, follow_redirects=False)
|
|
values = settings_store.images(db)
|
|
assert values["enabled"] is False
|
|
assert values["review_enabled"] is False
|
|
assert values["preserve_vram"] is False
|
|
|
|
|
|
def test_the_numbers_are_clamped(admin, db):
|
|
_settings(admin, max_tries="500", timeout="1")
|
|
values = settings_store.images(db)
|
|
assert values["max_tries"] == 10
|
|
assert values["timeout"] == 10.0
|
|
|
|
|
|
def test_max_tries_never_reaches_zero(admin, db):
|
|
"""Zero would mean the tool generates nothing and reports success. Unlike
|
|
the zeroes elsewhere in the settings, there is no reading of it anybody
|
|
wants."""
|
|
_settings(admin, max_tries="0")
|
|
assert settings_store.images(db)["max_tries"] == 1
|
|
|
|
|
|
def test_the_key_survives_a_save_that_did_not_touch_it(admin, db):
|
|
_settings(admin, api_key="secret-key")
|
|
assert decrypt(settings_store.images(db)["api_key_encrypted"]) == "secret-key"
|
|
|
|
_settings(admin, api_key=UNCHANGED_SENTINEL, instructions="something else")
|
|
assert decrypt(settings_store.images(db)["api_key_encrypted"]) == "secret-key"
|
|
|
|
|
|
def test_an_emptied_key_is_removed(admin, db):
|
|
_settings(admin, api_key="secret-key")
|
|
_settings(admin, api_key="")
|
|
assert settings_store.images(db)["api_key_encrypted"] == ""
|
|
|
|
|
|
def test_the_key_is_never_rendered(admin, db):
|
|
_settings(admin, api_key="secret-key")
|
|
assert "secret-key" not in admin.get("/admin/images").text
|
|
|
|
|
|
def test_testing_without_a_url_says_so(admin, db):
|
|
_settings(admin, base_url="")
|
|
assert "Set a base URL first" in admin.post("/admin/images/test").text
|
|
|
|
|
|
# --- Workflows -----------------------------------------------------------------
|
|
def test_new_is_not_read_as_an_id(admin):
|
|
"""FastAPI matches in registration order, so `/workflows/new` has to be
|
|
registered before `/workflows/{id}` or "new" is an id and 404s. This has
|
|
been a bug twice in this codebase."""
|
|
response = admin.get("/admin/images/workflows/new")
|
|
assert response.status_code == 200
|
|
assert "Export (API)" in response.text
|
|
|
|
|
|
def test_the_default_workflow_is_offered_to_start_from(admin):
|
|
"""Somebody setting this up for the first time should not have to find a
|
|
working API-format document before they can try anything."""
|
|
page = admin.get("/admin/images/workflows/new").text
|
|
assert "CheckpointLoaderSimple" in page
|
|
assert "{{prompt}}" in page
|
|
|
|
|
|
def test_a_workflow_is_created(admin, db):
|
|
response = admin.post(
|
|
"/admin/images/workflows",
|
|
data={
|
|
"slug": "sdxl",
|
|
"name": "SDXL",
|
|
"description": "photo",
|
|
"workflow": BASE_TEXT,
|
|
"enabled": "true",
|
|
},
|
|
follow_redirects=False,
|
|
)
|
|
assert response.status_code == 303
|
|
|
|
row = db.scalar(select(ImageWorkflow).where(ImageWorkflow.slug == "sdxl"))
|
|
assert row is not None
|
|
assert row.workflow_json["4"]["inputs"]["ckpt_name"] == "{{model}}"
|
|
|
|
|
|
def test_broken_json_is_refused_with_a_sentence(admin, db):
|
|
"""A sentence and not a 422: losing forty lines of pasted JSON to a
|
|
validation error is not a thing to do to somebody, so the form comes back
|
|
with what was typed still in it."""
|
|
response = admin.post(
|
|
"/admin/images/workflows",
|
|
data={"slug": "x", "name": "X", "workflow": "{not json"},
|
|
)
|
|
assert response.status_code == 200
|
|
assert "not valid JSON" in response.text
|
|
assert "{not json" in response.text, "and what was typed is still there"
|
|
assert db.scalar(select(ImageWorkflow)) is None
|
|
|
|
|
|
def test_a_workflow_with_no_prompt_placeholder_is_refused(admin, db):
|
|
"""It would draw the same picture whatever anybody typed, and would look
|
|
like a broken model rather than an unparameterised template."""
|
|
without = json.loads(BASE_TEXT)
|
|
without["6"]["inputs"]["text"] = "a fixed prompt"
|
|
response = admin.post(
|
|
"/admin/images/workflows",
|
|
data={"slug": "x", "name": "X", "workflow": json.dumps(without)},
|
|
)
|
|
assert "{{prompt}}" in response.text
|
|
assert "every image would be the same" in response.text
|
|
assert db.scalar(select(ImageWorkflow)) is None
|
|
|
|
|
|
def test_an_unknown_placeholder_is_refused(admin, db):
|
|
response = admin.post(
|
|
"/admin/images/workflows",
|
|
data={
|
|
"slug": "x",
|
|
"name": "X",
|
|
"workflow": json.dumps({"1": {"a": "{{prompt}} {{lora}}"}}),
|
|
},
|
|
)
|
|
assert "lora" in response.text
|
|
assert db.scalar(select(ImageWorkflow)) is None
|
|
|
|
|
|
def test_a_ui_format_export_is_refused_helpfully(admin, db):
|
|
"""The two ComfyUI formats look alike enough that this is the mistake
|
|
everybody makes first, and "invalid" would not tell them which button to
|
|
press instead."""
|
|
response = admin.post(
|
|
"/admin/images/workflows",
|
|
data={"slug": "x", "name": "X", "workflow": "[]"},
|
|
)
|
|
assert "Export (API)" in response.text
|
|
|
|
|
|
def test_a_duplicate_slug_is_refused(admin, db):
|
|
for _ in range(2):
|
|
response = admin.post(
|
|
"/admin/images/workflows",
|
|
data={"slug": "dup", "name": "Dup", "workflow": BASE_TEXT},
|
|
)
|
|
assert "already a workflow" in response.text
|
|
assert len(list(db.scalars(select(ImageWorkflow)))) == 1
|
|
|
|
|
|
def test_a_rejected_edit_leaves_the_stored_row_alone(admin, db):
|
|
admin.post(
|
|
"/admin/images/workflows",
|
|
data={"slug": "keep", "name": "Keep", "workflow": BASE_TEXT},
|
|
)
|
|
row = db.scalar(select(ImageWorkflow))
|
|
|
|
admin.post(
|
|
f"/admin/images/workflows/{row.id}",
|
|
data={"slug": "keep", "name": "Renamed", "workflow": "{broken"},
|
|
)
|
|
|
|
db.expire_all()
|
|
row = db.scalar(select(ImageWorkflow))
|
|
assert row.name == "Keep", "validated against a draft, so nothing was written"
|
|
|
|
|
|
def test_a_workflow_is_deleted(admin, db):
|
|
admin.post(
|
|
"/admin/images/workflows",
|
|
data={"slug": "gone", "name": "Gone", "workflow": BASE_TEXT},
|
|
)
|
|
row = db.scalar(select(ImageWorkflow))
|
|
admin.post(f"/admin/images/workflows/{row.id}/delete", follow_redirects=False)
|
|
|
|
db.expire_all()
|
|
assert db.scalar(select(ImageWorkflow)) is None
|
|
|
|
|
|
def test_only_an_administrator_can_reach_it(client: TestClient, db, registered):
|
|
"""The first account registered is an administrator, so this needs a second
|
|
one -- the `test_admin_tools` fixture's reason for existing."""
|
|
from lembas.db.models import User
|
|
|
|
client.post("/auth/logout")
|
|
client.post(
|
|
"/auth/register",
|
|
data={"name": "Sam", "email": "sam@shire.test", "password": "correct horse battery"},
|
|
follow_redirects=False,
|
|
)
|
|
user = db.scalar(select(User).where(User.email == "sam@shire.test"))
|
|
user.role = "user"
|
|
user.active = True
|
|
db.commit()
|
|
client.post(
|
|
"/auth/login",
|
|
data={"email": "sam@shire.test", "password": "correct horse battery"},
|
|
follow_redirects=False,
|
|
)
|
|
|
|
assert client.get("/admin/images").status_code in (302, 303, 403, 404)
|
|
assert client.post("/admin/images/workflows", data={}).status_code in (302, 303, 403, 404)
|