To change a picture from the gallery the AI had to open the saved card and set back, one by one, the model, LoRA and prompt the picture had been made with (8 to 12 calls in the exam). open_workflow now takes a picture's name too when no saved workflow has it: the server checks the picture is in the output folder and carries a workflow, and the page reads it from the PNG and opens it in a new tab, without the marks of the tab it ran on (the card that tab came from, which run_workflow would mistake it for, and its group of saves). The tool list is 18 characters shorter than before. Retaken with Gemini 3.8 Flash, questions 3 and 4 only: question 3 on a changed picture went from 28 calls (7 failed) to 21 (4), opening the picture's own workflow; written into docs/architecture/tool-exam.md. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
268 lines
16 KiB
Python
268 lines
16 KiB
Python
"""The MCP tools that build a workflow on the canvas (api/mcp_canvas_edit.py): node types read
|
|
from a stand-in /object_info, arguments checked before the page is asked, and the page asked
|
|
with exactly what it needs. No ComfyUI, no page."""
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
PLUGIN = Path(__file__).resolve().parents[1]
|
|
sys.path[:0] = [str(PLUGIN), str(PLUGIN.parents[1])]
|
|
|
|
from api import action_catalog, mcp_canvas_edit, mcp_fill, mcp_server # noqa: E402
|
|
from api.mcp_tools import ToolError, call_tool # noqa: E402
|
|
|
|
INFO = {
|
|
"KSampler": {"display_name": "KSampler", "category": "sampling", "description": "Denoises a latent",
|
|
"input": {"required": {"model": ["MODEL"], "steps": ["INT", {"default": 20, "min": 1, "max": 10000}],
|
|
"sampler_name": [["euler", "dpmpp_2m"]]},
|
|
"optional": {"note": ["STRING", {"multiline": True, "forceInput": True}]}},
|
|
"output": ["LATENT"], "output_name": ["LATENT"]},
|
|
"CheckpointLoaderSimple": {"display_name": "Load Checkpoint", "category": "loaders",
|
|
"input": {"required": {"ckpt_name": ["COMBO", {"options": [f"m{i}.safetensors" for i in range(80)]}]}},
|
|
"output": ["MODEL", "CLIP", "VAE"], "output_name": ["MODEL", "CLIP", "VAE"]},
|
|
"OldSampler": {"display_name": "Old sampler", "category": "sampling", "deprecated": True, "input": {}, "output": []},
|
|
}
|
|
|
|
|
|
class FindNodesTests(unittest.TestCase):
|
|
def setUp(self):
|
|
patcher = mock.patch.object(mcp_canvas_edit, "object_info", return_value=INFO)
|
|
patcher.start()
|
|
self.addCleanup(patcher.stop)
|
|
|
|
def test_search_gives_compact_candidates_without_choice_lists(self):
|
|
found = mcp_canvas_edit.find_nodes({}, query="load checkpoint")
|
|
node = found["nodes"][0]
|
|
self.assertEqual((node["type"], found["total_matches"]), ("CheckpointLoaderSimple", 1))
|
|
self.assertEqual(node["inputs"], ["ckpt_name:choice"])
|
|
self.assertEqual(node["outputs"], ["MODEL:MODEL", "CLIP:CLIP", "VAE:VAE"])
|
|
self.assertNotIn("m0.safetensors", str(found))
|
|
|
|
def test_search_is_bounded_even_with_many_large_node_definitions(self):
|
|
import json
|
|
with mock.patch.object(mcp_canvas_edit, "object_info", return_value={
|
|
"Large" + str(i): {"input": {"required": {"input" + str(k): ["STRING"] for k in range(100)}},
|
|
"description": "x" * 10000} for i in range(100)
|
|
}):
|
|
result = mcp_canvas_edit.find_nodes({}, query="Large")
|
|
self.assertEqual(result["total_matches"], 100)
|
|
self.assertLessEqual(len(result["nodes"]), 5)
|
|
self.assertLess(len(json.dumps(result)), 3600)
|
|
|
|
def test_settings_keep_their_range_and_sockets_say_so(self):
|
|
node = mcp_canvas_edit.find_nodes({}, types=["KSampler"])["nodes"][0]
|
|
steps = next(item for item in node["inputs"] if item["name"] == "steps")
|
|
self.assertEqual((steps["type"], steps["default"], steps["min"], steps["max"]), ("INT", 20, 1, 10000))
|
|
note = next(item for item in node["inputs"] if item["name"] == "note")
|
|
self.assertEqual((note["required"], note.get("socket_only")), (False, True))
|
|
|
|
def test_named_types_give_all_choices_up_to_a_limit_and_say_what_is_missing(self):
|
|
found = mcp_canvas_edit.find_nodes({}, types=["CheckpointLoaderSimple", "Nope"])
|
|
self.assertEqual(len(found["nodes"][0]["inputs"][0]["choices"]), mcp_canvas_edit.MAX_CHOICES_NAMED)
|
|
self.assertEqual(found["not_found"], ["Nope"])
|
|
|
|
def test_deprecated_nodes_come_last(self):
|
|
names = [node["type"] for node in mcp_canvas_edit.find_nodes({}, query="sampl")["nodes"]]
|
|
self.assertEqual(names, ["KSampler", "OldSampler"])
|
|
|
|
def test_nothing_asked(self):
|
|
with self.assertRaises(ToolError):
|
|
mcp_canvas_edit.find_nodes({}, query=" ")
|
|
|
|
|
|
class BuildToolTests(unittest.TestCase):
|
|
def setUp(self):
|
|
patcher = mock.patch.object(mcp_canvas_edit, "_page", return_value={"ok": True})
|
|
self.page = patcher.start()
|
|
self.addCleanup(patcher.stop)
|
|
|
|
def test_the_page_gets_what_it_needs(self):
|
|
self.assertIn("run_workflow", mcp_canvas_edit.add_node({}, "KSampler", {"steps": 28}, near_node_id=4)["next"])
|
|
self.assertEqual(self.page.call_args[0], ("add_node", {"type": "KSampler", "values": {"steps": 28}, "title": "", "near_node_id": 4,
|
|
"wire_from": [], "wire_to": [], "replace": False}))
|
|
mcp_canvas_edit.connect_nodes({}, 4, 3, "model", output="0")
|
|
self.assertEqual(self.page.call_args[0][1]["output"], 0, "a slot number given as text")
|
|
mcp_canvas_edit.connect_nodes({}, 4, 3, "clip", output="CLIP")
|
|
self.assertEqual(self.page.call_args[0][1]["output"], "CLIP")
|
|
mcp_canvas_edit.remove_nodes({}, [3, 4])
|
|
self.assertEqual(self.page.call_args[0], ("remove_nodes", {"node_ids": [3, 4]}))
|
|
|
|
def test_automatic_wiring_and_old_explicit_input_form(self):
|
|
mcp_canvas_edit.connect_nodes({}, 12, 15)
|
|
self.assertEqual(self.page.call_args[0][1], {"from_node_id": 12, "to_node_id": 15,
|
|
"input": None, "output": None, "replace": None})
|
|
mcp_canvas_edit.connect_nodes({}, 12, 15, input="pixels")
|
|
self.assertEqual(self.page.call_args[0][1]["output"], 0)
|
|
mcp_canvas_edit.add_node({}, "VAEEncodeForInpaint", wire_from=[12, 4], wire_to=3, replace=True)
|
|
payload = self.page.call_args[0][1]
|
|
self.assertEqual((payload["wire_from"], payload["wire_to"], payload["replace"]), ([12, 4], [3], True))
|
|
mcp_canvas_edit.add_node({}, "ImageScale", insert_after=8, wire_to=9)
|
|
self.assertEqual(self.page.call_args[0][1]["insert_after"], 8)
|
|
|
|
def test_registry_accepts_auto_calls_without_socket_names(self):
|
|
for tool, args in (("connect_nodes", {"from_node_id": 1, "to_node_id": 2}),
|
|
("add_node", {"type": "ImageScale", "insert_after": 8, "wire_to": [9]})):
|
|
result = call_tool(mcp_server.TOOLS, tool, args, {})
|
|
self.assertFalse(result.get("isError"), result)
|
|
self.assertEqual(self.page.call_count, 2)
|
|
|
|
def test_a_group_of_nodes_with_canvas_links(self):
|
|
nodes = {"a": {"class_type": "VAEEncode", "inputs": {"pixels": ["#12", "0"], "vae": ["#4", "VAE"]}},
|
|
"b": {"class_type": "SetLatentNoiseMask", "inputs": {"samples": ["a", 0], "mask": ["#12", "MASK"]}, "wire_to": 15}}
|
|
mcp_canvas_edit.add_nodes({}, nodes, replace=True)
|
|
action, payload = self.page.call_args[0]
|
|
self.assertEqual((action, payload["replace"], payload["near_node_id"]), ("add_nodes", True, None))
|
|
self.assertEqual(payload["nodes"]["a"]["inputs"]["pixels"], ["#12", 0], "a slot number given as text")
|
|
self.assertEqual(payload["nodes"]["b"]["wire_to"], [15])
|
|
for bad in ({"#1": {"class_type": "VAEEncode", "inputs": {}}},
|
|
{"a": {"class_type": "VAEEncode", "inputs": {"pixels": ["#12"]}}},
|
|
{"a": {"class_type": "VAEEncode", "inputs": {"pixels": ["#12", True]}}},
|
|
{"a": {"class_type": "VAEEncode", "inputs": {}, "wire_to": list(range(9))}}):
|
|
with self.assertRaises(ToolError, msg=bad):
|
|
mcp_canvas_edit.add_nodes({}, bad)
|
|
self.assertEqual(self.page.call_count, 1)
|
|
|
|
def test_settings_by_name_or_for_one_node(self):
|
|
mcp_canvas_edit.set_values({}, values={"steps": 30, "12": {"denoise": 0.5}})
|
|
self.assertEqual(self.page.call_args[0], ("set_values", {"node_id": None, "values": {"steps": 30, "12": {"denoise": 0.5}}}))
|
|
mcp_canvas_edit.set_values({}, 4, {"cfg": 5})
|
|
self.assertEqual(self.page.call_args[0], ("set_values", {"node_id": 4, "values": {"cfg": 5}}), "the old form")
|
|
for values in ({"sampler": {"denoise": 0.5}}, {"12": {"denoise": [1]}}, {}, {"12": {f"s{i}": 1 for i in range(61)}}):
|
|
with self.assertRaises(ToolError, msg=values):
|
|
mcp_canvas_edit.set_values({}, values=values)
|
|
|
|
def test_receipt_does_not_claim_disconnected_branch_can_run(self):
|
|
self.page.return_value = {"reaches_output": False, "unwired": []}
|
|
self.assertIn("does not reach", mcp_canvas_edit.add_node({}, "ImageScale")["next"])
|
|
|
|
def test_a_workflow_in_api_form(self):
|
|
nodes = {"1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "a.safetensors"}},
|
|
"2": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "steps": 20}}}
|
|
mcp_canvas_edit.new_workflow({}, "Portrait", nodes)
|
|
self.assertEqual(self.page.call_args, mock.call("new_workflow", {"name": "Portrait", "nodes": nodes}, timeout=60))
|
|
mcp_canvas_edit.new_workflow({})
|
|
self.assertIsNone(self.page.call_args[0][1]["nodes"], "a blank tab")
|
|
|
|
def test_a_saved_workflow_or_one_of_its_sets_opens(self):
|
|
for workflow, found in (("海边", ("recipe_b.json", "海边", None)),
|
|
("海边 / 夕阳", ("recipe_b.json", "海边 / 夕阳", "params_1.json"))):
|
|
with mock.patch.object(mcp_fill, "saved_workflow", return_value=found):
|
|
mcp_canvas_edit.open_workflow({}, workflow)
|
|
self.assertEqual(self.page.call_args, mock.call("open_workflow", {"filename": found[0], "name": found[1], "plan": found[2]}, timeout=60))
|
|
|
|
def test_a_gallery_pictures_own_workflow_opens(self):
|
|
import json
|
|
import tempfile
|
|
from PIL import Image, PngImagePlugin
|
|
from api import mcp_tools
|
|
root = Path(tempfile.mkdtemp())
|
|
self.addCleanup(__import__("shutil").rmtree, root, True)
|
|
(root / "sub").mkdir()
|
|
info = PngImagePlugin.PngInfo()
|
|
info.add_text("workflow", json.dumps({"nodes": [], "links": []}))
|
|
Image.new("RGB", (8, 8)).save(root / "sub" / "made_00001_.png", pnginfo=info)
|
|
Image.new("RGB", (8, 8)).save(root / "bare_00001_.png")
|
|
missing = ToolError("No saved workflow or plan called that; list_workflows names them.")
|
|
with mock.patch.object(mcp_tools.folder_paths, "get_output_directory", return_value=str(root)), \
|
|
mock.patch.object(mcp_fill, "saved_workflow", side_effect=missing):
|
|
mcp_canvas_edit.open_workflow({}, "sub/made_00001_.png")
|
|
self.assertEqual(self.page.call_args, mock.call(
|
|
"open_workflow", {"picture": {"subfolder": "sub", "filename": "made_00001_.png"}, "name": "made_00001_"}, timeout=60))
|
|
self.page.reset_mock()
|
|
for name, said in (("bare_00001_.png", "carries no workflow"), ("../made_00001_.png", "output"),
|
|
("gone_00001_.png", "No output image"), ("海边", "list_workflows")):
|
|
with self.subTest(name=name), self.assertRaises(ToolError) as caught:
|
|
mcp_canvas_edit.open_workflow({}, name)
|
|
self.assertIn(said, str(caught.exception))
|
|
self.page.assert_not_called()
|
|
|
|
def test_bad_arguments_never_reach_the_page(self):
|
|
bad = [
|
|
lambda: mcp_canvas_edit.add_node({}, ""),
|
|
lambda: mcp_canvas_edit.add_node({}, "KSampler", {"steps": [1, 2]}),
|
|
lambda: mcp_canvas_edit.add_node({}, "KSampler", {"text": "x" * (mcp_fill.MAX_TEXT + 1)}),
|
|
lambda: mcp_canvas_edit.add_node({}, "KSampler", near_node_id=True),
|
|
lambda: mcp_canvas_edit.add_node({}, "ImageScale", wire_from=[]),
|
|
lambda: mcp_canvas_edit.add_node({}, "ImageScale", wire_to=list(range(9))),
|
|
lambda: mcp_canvas_edit.add_node({}, "ImageScale", insert_type="IMAGE"),
|
|
lambda: mcp_canvas_edit.add_node({}, "ImageScale", replace="yes"),
|
|
lambda: mcp_canvas_edit.connect_nodes({}, 1, 2, replace="yes"),
|
|
lambda: mcp_canvas_edit.connect_nodes({}, 1, 2, ""),
|
|
lambda: mcp_canvas_edit.connect_nodes({}, {"id": 1}, 2, "model"),
|
|
lambda: mcp_canvas_edit.set_values({}, 3, {}),
|
|
lambda: mcp_canvas_edit.remove_nodes({}, []),
|
|
lambda: mcp_canvas_edit.remove_nodes({}, list(range(mcp_canvas_edit.MAX_REMOVE + 1))),
|
|
lambda: mcp_canvas_edit.new_workflow({}, "x", {}),
|
|
lambda: mcp_canvas_edit.new_workflow({}, "x", {"1": {"inputs": {}}}),
|
|
lambda: mcp_canvas_edit.new_workflow({}, "x", {"1": {"class_type": "A", "inputs": {"t": {"deep": 1}}}}),
|
|
]
|
|
for call in bad:
|
|
with self.assertRaises(ToolError):
|
|
call()
|
|
self.page.assert_not_called()
|
|
|
|
|
|
class PictureTests(unittest.TestCase):
|
|
"""add_nodes(use="picture"): the model picked for the task, an empty canvas, the family's files beside it."""
|
|
|
|
INFO = {"UNETLoader": {"input": {"required": {"unet_name": [["Anima\harem.safetensors"]]}}},
|
|
"CLIPLoader": {"input": {"required": {"clip_name": [["Anima\qwen_3_06b_base.safetensors"]]}}},
|
|
"VAELoader": {"input": {"required": {"vae_name": [["Anima\qwenImageGGUF_vae.safetensors"]]}}}}
|
|
|
|
def setUp(self):
|
|
self.canvas = None # what the page's sampler reads: None, an empty canvas
|
|
|
|
def page(action, args=None, timeout=30):
|
|
if action == "block_model":
|
|
if self.canvas is None:
|
|
raise ToolError("No sampler (KSampler) on the canvas: open a workflow that makes the picture first.")
|
|
return {"model": self.canvas}
|
|
return {"ok": True}
|
|
self.page = mock.MagicMock(side_effect=page)
|
|
root = str(Path("/models/diffusion_models"))
|
|
for patcher in (mock.patch.object(mcp_canvas_edit, "_page", self.page),
|
|
mock.patch.object(mcp_canvas_edit, "object_info", return_value=self.INFO),
|
|
mock.patch.object(mcp_canvas_edit, "_resolve_model",
|
|
return_value=("diffusion_models", 0, root, str(Path(root, "Anima", "harem.safetensors")))),
|
|
mock.patch.object(mcp_canvas_edit.model_advice, "base_of", return_value="Anima"),
|
|
mock.patch.object(mcp_canvas_edit.ai_tasks, "get_task", return_value={}),
|
|
mock.patch.object(mcp_canvas_edit.ai_toolbox, "effective",
|
|
return_value={"models": ["loras:0:style.safetensors", "diffusion_models:0:Anima/harem.safetensors"]})):
|
|
patcher.start()
|
|
self.addCleanup(patcher.stop)
|
|
|
|
def test_the_picked_model_on_an_empty_canvas(self):
|
|
mcp_canvas_edit.add_nodes({"task": "t"}, use="picture")
|
|
nodes = next(call[0][1]["nodes"] for call in self.page.call_args_list if call[0][0] == "add_nodes")
|
|
self.assertEqual((nodes["model"]["inputs"]["unet_name"], nodes["text_encoder"]["inputs"]["clip_name"]),
|
|
("Anima\harem.safetensors", "Anima\qwen_3_06b_base.safetensors"), "as ComfyUI lists them")
|
|
|
|
def test_a_canvas_that_makes_a_picture_or_no_model_is_refused(self):
|
|
self.canvas = "Anima\harem.safetensors"
|
|
with self.assertRaisesRegex(ToolError, "already makes a picture"):
|
|
mcp_canvas_edit.add_nodes({"task": "t"}, use="picture")
|
|
self.canvas = None
|
|
with mock.patch.object(mcp_canvas_edit.ai_toolbox, "effective", return_value=None):
|
|
with self.assertRaisesRegex(ToolError, "needs the model to draw with"):
|
|
mcp_canvas_edit.add_nodes({"task": "t"}, use="picture")
|
|
self.assertFalse(any(call[0][0] == "add_nodes" for call in self.page.call_args_list), "nothing built")
|
|
|
|
|
|
class RegistryTests(unittest.TestCase):
|
|
def test_served_with_their_place_in_the_flow(self):
|
|
for name in mcp_canvas_edit.CANVAS_TOOLS:
|
|
self.assertIn(name, mcp_server.TOOLS)
|
|
self.assertIn(name, action_catalog.TOOLS)
|
|
self.assertEqual(action_catalog.TOOLS["add_node"].gate, "canvas")
|
|
self.assertTrue(mcp_server.TOOLS["find_nodes"]["spec"]["annotations"]["readOnlyHint"])
|
|
self.assertTrue(mcp_server.TOOLS["remove_nodes"]["spec"]["annotations"]["destructiveHint"])
|
|
|
|
def test_unknown_arguments_are_refused_by_the_schema(self):
|
|
result = call_tool(mcp_server.TOOLS, "add_node", {"type": "KSampler", "delete_files": True}, {})
|
|
self.assertTrue(result["isError"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|