Files
DemonGatanjieu-Anomalous_Mo…/tests/test_mcp_canvas_edit.py
T
DemonGatanjieuandClaude Opus 5.5 05e2d269b4 Let open_workflow open the workflow a gallery picture carries
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>
2026-11-06 19:57:13 +08:00

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