Draw a picture, on a ComfyUI you are running

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>
This commit is contained in:
Jaroslav Beneš
2026-08-05 14:13:19 +02:00
co-authored by Claude Opus 5
parent 9f5ff72e32
commit 47d1ddbc3c
38 changed files with 3958 additions and 18 deletions
+73 -5
View File
@@ -113,6 +113,12 @@ FAMILY_SCRATCH = "scratch"
# services/agent/session.py:resolve, which answers all three at once.
FAMILY_AGENT = "agent"
# Drawing a picture on a ComfyUI an administrator configured. A family of its
# own for the reason `fetch` is one: an instance may reasonably want a model
# that can look things up but not spend a minute of GPU on every request, and
# the whole cost of this one is somewhere else.
FAMILY_IMAGE = "image"
# The built-in families, in the order they are offered.
FAMILIES = (
FAMILY_SEARCH,
@@ -123,6 +129,7 @@ FAMILIES = (
FAMILY_SKILLS,
FAMILY_SCRATCH,
FAMILY_ASK,
FAMILY_IMAGE,
FAMILY_AGENT,
)
@@ -164,6 +171,11 @@ class ToolContext:
"""
owner_id: str
# Which conversation this call belongs to. Needed by anything that writes
# something the chat owns rather than something the *reader* owns -- the
# scratch document, a generated image -- and empty for a call with no chat
# behind it, which is what those runners check first.
chat_id: str = ""
search_config: dict[str, Any] = field(default_factory=dict)
# Which knowledge bases this chat is scoped to. Empty means "everything the
# owner can see", which is what a chat with none attached should do.
@@ -186,6 +198,17 @@ class ToolContext:
# a model can name a skill it was never shown and the runner would fetch it
# anyway. Same rule as "what may be run is what was offered".
skills_off: frozenset[str] = field(default_factory=frozenset)
# Image generation, snapshotted like everything else here. `image_config` is
# the instance settings group; the two below are this chat's preferences,
# used when the model names neither. `model_id` and `connection_id` are what
# the reviewer and the Preserve VRAM unload need to find the chat's own
# endpoint -- its own, and no other, because the VRAM being freed belongs to
# one machine.
image_config: dict[str, Any] = field(default_factory=dict)
image_workflow_id: str = ""
image_checkpoint: str = ""
model_id: str = ""
connection_id: str = ""
@dataclass
@@ -1097,7 +1120,7 @@ REGISTRY: dict[str, ToolDef] = {
def _family_allowed(
family: str, *, config: dict, capabilities: dict, allowed: dict
family: str, *, config: dict, capabilities: dict, allowed: dict, images: bool = False
) -> bool:
"""Whether one family is on for this chat.
@@ -1122,6 +1145,14 @@ def _family_allowed(
# attach path keeps working, because that one is a person's instruction
# rather than a model's choice.
return bool(allowed.get("tools.fetch") and config.get("fetch_enabled"))
if gate == FAMILY_IMAGE:
# Its own branch rather than a name in the tuple below, and the second
# half is why: an instance with no ComfyUI, or one with no checkpoints
# listed, must not offer this at all. A model that calls it there spends
# a round to be told the thing it was offered does not work, which is
# the shape `resolve_tools` already refuses for `skill_get` with an
# empty library. `settings_store.images_ready` answers all three.
return bool(allowed.get("tools.image") and images)
if gate in (FAMILY_CUSTOM, FAMILY_MCP, FAMILY_ASK, FAMILY_AGENT, FAMILY_SCRATCH):
# Deliberately without `library.use`: an HTTP endpoint an administrator
# wrote has nothing to do with this person's own documents and notes,
@@ -1167,6 +1198,20 @@ def _agent_defs(db: DBSession, chat: Chat | None, user: User | None) -> list[Too
return agent_tools.tool_defs(context)
def _image_defs(db: DBSession, values: dict | None = None) -> list[ToolDef]:
"""The image tool, whose schema carries this instance's own choices.
Built per request rather than at import, because the templates a model may
name and the checkpoints it may draw with are rows and settings. That is the
same reason a custom tool cannot live in `REGISTRY`, and it is why this has
to be listed in `registry(db)` below as well -- a name that resolves to no
family is a tool whose guidance never reaches the model.
"""
from lembas.services.images import tool as image_tool
return [image_tool.tool_def(db, values if values is not None else settings_store.images(db))]
def _book(defs: list[ToolDef]) -> dict[str, ToolDef]:
"""Keyed by name, first claim winning.
@@ -1196,7 +1241,9 @@ def registry(db: DBSession) -> dict[str, ToolDef]:
"""
from lembas.services.agent import tools as agent_tools
return _book([*_row_defs(db, None, everything=True), *agent_tools.tool_defs()])
return _book(
[*_row_defs(db, None, everything=True), *agent_tools.tool_defs(), *_image_defs(db)]
)
def families(db: DBSession) -> tuple[str, ...]:
@@ -1220,10 +1267,21 @@ def resolve_tools(db: DBSession, chat: Chat, user: User | None) -> ToolSet:
allowed = permissions.resolve(db, user)
config = settings_store.search(db)
image_values = settings_store.images(db)
images_ready = settings_store.images_ready(db)
# Resolved against what this reader may see, not against everything that
# exists: a tool restricted to a group is not offered outside it.
book = _book([*_row_defs(db, user), *_agent_defs(db, chat, user)])
# exists: a tool restricted to a group is not offered outside it. The image
# tool is built only when it could be offered, because building its schema
# reads the workflow table and there is no sense doing that for an instance
# with no ComfyUI.
book = _book(
[
*_row_defs(db, user),
*_agent_defs(db, chat, user),
*(_image_defs(db, image_values) if images_ready else []),
]
)
# What this chat has switched off, applied AFTER the gates and never
# instead of them. A chat can only ever *narrow* what the model's
@@ -1240,7 +1298,11 @@ def resolve_tools(db: DBSession, chat: Chat, user: User | None) -> ToolSet:
tool
for tool in book.values()
if _family_allowed(
tool.family, config=config, capabilities=capabilities, allowed=allowed
tool.family,
config=config,
capabilities=capabilities,
allowed=allowed,
images=images_ready,
)
and gate_of(tool.family) not in off
# Nothing to read and nothing to improve. Offering `skill_get` with
@@ -1327,7 +1389,13 @@ def context_for(
return ToolContext(
agent=agent_session.resolve(db, chat, user) if chat is not None else None,
owner_id=user.id if user else "",
chat_id=chat.id if chat is not None else "",
search_config=settings_store.search(db),
image_config=settings_store.images(db),
image_workflow_id=(chat.image_workflow_id or "") if chat is not None else "",
image_checkpoint=(chat.image_checkpoint or "") if chat is not None else "",
model_id=(chat.model_id or "") if chat is not None else "",
connection_id=(chat.connection_id or "") if chat is not None else "",
base_ids=[base.id for base in chat.knowledge_bases] if chat is not None else [],
skills_off=scoped_skills_off(chat),
tools=tools.by_name if tools is not None else None,