Files
DemonGatanjieu-Anomalous_Mo…/tests/test_ai_run_map.py
T
DemonGatanjieuandClaude Opus 5.5 a74f885c53 Show what the task needs now over the message box; take dropped and pasted pictures
- ui_ai_now.js: over the message box, a run's progress with stop, the AI's step still going, a
  gate to allow or refuse there, and the AI's open question with its options as buttons (own
  words or skip too; several picks and pictures lead to the stage). progressWords is shared
  with the run card; the panel's and gate's answers are exported for it.
- ui_ai_drop.js: pictures dropped on the box or pasted are kept in output/ai_inbox
  (POST /anomalous/ai_pictures/upload, ai_pictures.save_upload: read as a picture first,
  25 MB cap, never over a file) and attached; ComfyUI output pictures dragged in attach as they are.
- Long AI notes are folded to their first lines with "Show all".

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-23 21:24:48 +08:00

180 lines
10 KiB
Python

"""A run's shape and what changed (ai_run_map.py), the task's own canvas tab (ai_task_canvas.py),
pictures handed to the AI (ai_pictures.py) and show_run_map."""
import io
import json
import os
import time
import unittest
from unittest import mock
from PIL import Image
from PIL.PngImagePlugin import PngInfo
from test_ai_tasks import OUTPUT, isolate, ok
from api import ai_panels, ai_pictures, ai_run_map, ai_task_canvas, ai_tasks, mcp_task_tools
def lora_workflow(steps=20, strength=0.65):
return {
"4": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "sdxl/Kirazuri_v4.safetensors"}, "_meta": {"title": "Kirazuri"}},
"10": {"class_type": "LoraLoader", "inputs": {"model": ["4", 0], "clip": ["4", 1], "lora_name": "agefx.safetensors",
"strength_model": strength, "strength_clip": 1}},
"11": {"class_type": "Reroute", "inputs": {"": ["10", 0]}},
"6": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["10", 1], "text": "irena, grey hair, witch hat"}},
"7": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["10", 1], "text": "lowres"}},
"5": {"class_type": "EmptyLatentImage", "inputs": {"width": 1024, "height": 1536, "batch_size": 2}},
"3": {"class_type": "KSampler", "inputs": {"model": ["11", 0], "positive": ["6", 0], "negative": ["7", 0], "latent_image": ["5", 0],
"seed": 123, "steps": steps, "cfg": 5, "sampler_name": "euler", "scheduler": "normal", "denoise": 1}},
"8": {"class_type": "VAEDecode", "inputs": {"samples": ["3", 0], "vae": ["4", 2]}},
"9": {"class_type": "SaveImage", "inputs": {"images": ["8", 0], "filename_prefix": "irena"}},
}
class MapTests(unittest.TestCase):
def test_boxes_roles_facts_and_links_through_plumbing(self):
shape = ai_run_map.map_of(lora_workflow())
boxes = {box["id"]: box for box in shape["boxes"]}
self.assertNotIn("11", boxes, "a reroute is wired through, not shown")
self.assertEqual(boxes["4"]["role"], "model")
self.assertEqual(boxes["4"]["facts"]["file"], "Kirazuri_v4")
self.assertEqual((boxes["10"]["role"], boxes["10"]["facts"]["strength"]), ("lora", 0.65))
self.assertEqual(boxes["6"]["facts"], {"text": "irena, grey hair, witch hat", "polarity": "positive"})
self.assertEqual(boxes["7"]["facts"]["polarity"], "negative")
self.assertEqual(boxes["5"]["facts"], {"width": 1024, "height": 1536, "batch": 2})
self.assertEqual(boxes["3"]["facts"], {"steps": 20, "cfg": 5, "sampler_name": "euler", "scheduler": "normal", "denoise": 1})
self.assertIn({"from": "10", "to": "3", "type": "MODEL", "input": "model"}, [
{**link, "type": "MODEL"} for link in shape["links"] if link["from"] == "10" and link["to"] == "3"])
self.assertNotIn("3.seed", shape["params"], "the seed changes every run; it is no change to report")
def test_changes_since_the_run_before(self):
before = ai_run_map.map_of(lora_workflow())
after_workflow = lora_workflow(steps=30, strength=0.8)
after_workflow["12"] = {"class_type": "LoraLoader", "inputs": {"model": ["10", 0], "lora_name": "light.safetensors"}}
del after_workflow["9"]
changes = ai_run_map.changes(before, ai_run_map.map_of(after_workflow))
self.assertIn({"node": "12", "title": "LoraLoader", "change": "added"}, changes)
self.assertIn({"node": "9", "title": "SaveImage", "change": "removed"}, changes)
self.assertIn({"node": "3", "title": "KSampler", "key": "steps", "from": "20", "to": "30"}, changes)
self.assertIn({"node": "10", "title": "LoraLoader", "key": "strength_model", "from": "0.65", "to": "0.8"}, changes)
self.assertEqual(ai_run_map.changes(None, before), [])
def test_a_huge_workflow_keeps_what_makes_the_picture(self):
workflow = lora_workflow()
for index in range(40):
workflow[f"x{index}"] = {"class_type": "SomeHelper", "inputs": {"value": index}}
shape = ai_run_map.map_of(workflow)
self.assertEqual(len(shape["boxes"]), ai_run_map.MAX_BOXES)
self.assertTrue({"3", "4", "10", "6", "5"} <= {box["id"] for box in shape["boxes"]})
class TaskCanvasTests(unittest.TestCase):
def setUp(self):
isolate()
self.task = ai_tasks.begin_task("A", "Irena portrait")["id"]
def test_the_task_gets_its_tab_and_the_ai_is_told_once(self):
self.assertEqual(ai_task_canvas.default_mode(), "new")
self.assertIsNone(ai_task_canvas.spec("set_values"), "outside a tool call there is no task")
holder = ai_task_canvas.begin_call(self.task)
self.assertEqual(ai_task_canvas.spec("set_values"), {"mode": "new", "tab": None, "title": "Irena portrait"})
self.assertIsNone(ai_task_canvas.spec("open_model"), "not a canvas action")
ai_task_canvas.remember({"path": "workflows/Irena portrait 101500.json", "name": "Irena portrait 101500", "made": True})
self.assertEqual(ai_tasks.get_task(self.task)["canvas"]["tab"]["path"], "workflows/Irena portrait 101500.json")
told = ai_task_canvas.tell(ok({"done": True}), holder)
self.assertIn("Irena portrait 101500", told["structuredContent"]["canvas_tab"]["note"])
again = ai_task_canvas.begin_call(self.task)
self.assertEqual(ai_task_canvas.spec("add_node")["tab"]["name"], "Irena portrait 101500")
self.assertNotIn("canvas_tab", ai_task_canvas.tell(ok({}), again)["structuredContent"])
def test_the_user_chooses_per_task_and_by_default(self):
ai_task_canvas.set_mode(self.task, "current")
ai_task_canvas.begin_call(self.task)
self.assertEqual(ai_task_canvas.spec("set_values")["mode"], "current")
other = ai_tasks.begin_task("B", "Other", new_task=True)["id"]
ai_task_canvas.set_default("current")
self.assertEqual(ai_task_canvas.mode_of(ai_tasks.get_task(other)), "current")
with self.assertRaises(ValueError):
ai_task_canvas.set_mode(self.task, "sideways")
with self.assertRaises(ValueError):
ai_task_canvas.set_default("sideways")
ai_task_canvas.set_default("new")
def picture(name, workflow=None):
info = PngInfo()
if workflow:
info.add_text("prompt", json.dumps(workflow))
Image.new("RGB", (64, 96), (40, 80, 120)).save(os.path.join(OUTPUT, name), pnginfo=info)
return name
class PictureTests(unittest.TestCase):
def setUp(self):
isolate()
self.task = ai_tasks.begin_task("A", "Look at these")["id"]
self.patch = mock.patch.object(ai_tasks.folder_paths, "get_output_directory", return_value=OUTPUT)
self.patch.start()
self.addCleanup(self.patch.stop)
def test_pictures_reach_the_ai_themselves_with_what_made_them(self):
ai_pictures.hand(self.task, [picture("one.png", lora_workflow())])
ai_pictures.hand(self.task, [picture("two.png")])
steps = ai_tasks.get_task(self.task)["steps"]
self.assertEqual(len(steps), 1, "pictures handed before the AI looked join one step")
self.assertEqual([item["filename"] for item in steps[0]["results"]], ["one.png", "two.png"])
result = ai_panels.attach_pending(ok({}), "A", task_id=self.task)
message = result["structuredContent"]["user_messages"][0]
self.assertEqual((message["kind"], message["images"]), ("pictures", ["one.png", "two.png"]))
made = message["made_with"][0]
self.assertEqual((made["models"], made["loras"], made["positive"], made["size"]),
(["Kirazuri_v4"], ["agefx (0.65)"], "irena, grey hair, witch hat", "1024x1536"))
self.assertEqual(message["made_with"][1], {"image": "two.png"})
self.assertEqual(sum(1 for item in result["content"] if item.get("type") == "image"), 2)
self.assertNotIn("_pictures", message)
def test_only_output_pictures_and_open_tasks(self):
for images in ([], ["../secret.png"], ["missing.png"], [picture("x.png")] * 9):
with self.assertRaises(ai_pictures.PicturesError):
ai_pictures.hand(self.task, images)
with self.assertRaises(ai_pictures.PicturesError):
ai_pictures.hand("nope", [picture("y.png")])
def test_a_dropped_picture_is_kept_in_the_inbox_never_over_a_file(self):
data = io.BytesIO()
Image.new("RGB", (8, 8), "red").save(data, "JPEG")
first = ai_pictures.save_upload(data.getvalue(), "../../草稿 参考.png", now=1000)
second = ai_pictures.save_upload(data.getvalue(), "../../草稿 参考.png", now=1000)
self.addCleanup(lambda: [os.remove(os.path.join(OUTPUT, rel)) for rel in (first, second)])
self.assertTrue(first.startswith("ai_inbox/草稿 参考_") and first.endswith(".jpg"), first) # its stem, its real type
self.assertNotEqual(first, second)
ai_pictures.hand(self.task, [first]) # an output picture like any other
for bad in (b"", b"not a picture", b"x" * (ai_pictures.MAX_UPLOAD + 1)):
with self.assertRaises(ai_pictures.PicturesError):
ai_pictures.save_upload(bad, "x.png")
class ShowRunMapTests(unittest.TestCase):
def setUp(self):
isolate()
self.task = ai_tasks.begin_task("A", "Run")["id"]
self.ctx = {"client": "A", "task": self.task}
def test_the_stage_shows_the_latest_run_that_has_its_shape(self):
from api.mcp_tools import ToolError
with self.assertRaises(ToolError):
mcp_task_tools.show_run_map(self.ctx, "结构")
step = ai_tasks.record("A", "run_workflow", {}, ok({}), time.time())
with ai_tasks._lock:
task = ai_tasks._load()[self.task]
task["steps"][-1].update(map=ai_run_map.map_of(lora_workflow()), changes=[])
ai_tasks._save(task)
answer = mcp_task_tools.show_run_map(self.ctx, "结构")
self.assertEqual(answer["run"], step["id"])
panel = ai_tasks.get_task(self.task)["steps"][-1]["panel"]
self.assertEqual(panel["blocks"][0]["type"], "run_map")
self.assertEqual(panel["blocks"][0]["step"], step["id"])
if __name__ == "__main__":
unittest.main()