DemonGatanjieuandClaude Opus 5.5
6246df480d
Waiting on runs was the biggest share of the AI's calls in the user's real tasks (85 of 316 run_status calls, about ten per ten-minute run). A new Gates setting, "Waiting for runs", defaults to "Tell AI when done" (ai_gates.run_wait, stored beside the tool set in ai_tools.json). In that mode: - run_workflow still waits its 15 s, so a short run comes back done with its pictures. - Otherwise the answer's note, `next` and the your_run sentence tell the AI to say the run is going and end its turn, not to wait with run_status. run_status still honours wait_seconds, so a model that polls anyway is no worse off than before. When the run ends before the AI knows, the server marks the step `call_ai`. The workbench then shows lines to copy into the AI's chat (ui_ai_call.js): "It's done, take a look", "Again with another seed", "Go on", or "It failed, find out why". The mark goes as soon as the AI is told, or when a newer run starts, and the lines go with it. The AI's next call brings the result and the pictures. "AI waits" keeps the old behaviour. The setting takes effect at once, because the words are made for each answer. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
247 lines
15 KiB
Python
247 lines
15 KiB
Python
"""A run an AI started (api/ai_task_runs.py): where it stands, read from a stand-in ComfyUI queue
|
|
and progress registry, on its step and for the AI's run_status. Temporary folders only."""
|
|
import sys
|
|
import threading
|
|
import time
|
|
import unittest
|
|
from enum import Enum
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest import mock
|
|
|
|
PLUGIN = Path(__file__).resolve().parents[1]
|
|
sys.path[:0] = [str(PLUGIN), str(PLUGIN.parents[1]), str(PLUGIN / "tests")]
|
|
|
|
from api import ai_gates, ai_task_runs, ai_tasks # noqa: E402
|
|
from test_ai_tasks import isolate, ok # noqa: E402
|
|
|
|
|
|
class State(Enum):
|
|
Pending = "pending"
|
|
Running = "running"
|
|
Finished = "finished"
|
|
|
|
|
|
class Queue:
|
|
def __init__(self, running=(), pending=()):
|
|
self.running, self.pending = list(running), list(pending)
|
|
|
|
def get_current_queue(self):
|
|
return self.running, self.pending
|
|
|
|
|
|
def registry(prompt_id, nodes):
|
|
prompt = {"3": {"class_type": "KSampler", "_meta": {"title": "采样器"}}, "4": {"class_type": "CheckpointLoaderSimple"},
|
|
"8": {"class_type": "VAEDecode"}, "9": {"class_type": "SaveImage"}}
|
|
return SimpleNamespace(prompt_id=prompt_id, nodes=nodes,
|
|
dynprompt=SimpleNamespace(original_prompt=prompt, get_node=lambda node_id: prompt[node_id]))
|
|
|
|
|
|
class ProgressTests(unittest.TestCase):
|
|
def setUp(self):
|
|
isolate()
|
|
|
|
def test_queued_then_running_with_the_node_and_its_steps(self):
|
|
queue = Queue(running=[(5, "other")], pending=[(7, "mine"), (6, "first")])
|
|
with mock.patch.object(ai_task_runs, "_queue", return_value=queue):
|
|
self.assertEqual(ai_task_runs.progress_of(["mine"]), {"state": "queued", "ahead": 2})
|
|
queue.running, queue.pending = [(7, "mine")], []
|
|
nodes = {"4": {"state": State.Finished, "value": 1, "max": 1}, "3": {"state": State.Running, "value": 14, "max": 28}}
|
|
with mock.patch.object(ai_task_runs, "_registry", return_value=registry("mine", nodes)):
|
|
got = ai_task_runs.progress_of(["mine"])
|
|
self.assertEqual(got, {"state": "running", "prompt": 1, "prompts": 1, "nodes_done": 1, "nodes": 4,
|
|
"node": "采样器", "value": 14, "max": 28, "percent": 50})
|
|
with mock.patch.object(ai_task_runs, "_queue", side_effect=RuntimeError("no server")):
|
|
self.assertIsNone(ai_task_runs.progress_of(["mine"]))
|
|
|
|
def test_progress_reaches_the_step_and_run_status_waits_for_the_end(self):
|
|
ai_tasks.record("Claude", "run_workflow", {"batch": 1}, ok({"queued": 1}), time.time())
|
|
task_id = ai_tasks.list_tasks()[0]["id"]
|
|
step_id = ai_tasks.get_task(task_id)["steps"][0]["id"]
|
|
with mock.patch.object(ai_tasks, "_notify") as told:
|
|
ai_task_runs._show_progress(task_id, step_id, {"state": "running", "value": 3, "max": 20})
|
|
ai_task_runs._show_progress(task_id, step_id, {"state": "running", "value": 3, "max": 20})
|
|
self.assertEqual(told.call_count, 1, "the page hears only of changes")
|
|
self.assertEqual(told.call_args[0][1]["progress"]["value"], 3)
|
|
status = ai_task_runs.status("Claude")
|
|
self.assertEqual((status["state"], status["progress"]["max"]), ("running", 20))
|
|
self.assertIn("end your turn", status["next"], "by default the user tells the AI when it is done")
|
|
with mock.patch.object(ai_gates, "run_wait", return_value="wait"):
|
|
self.assertIn("run_status(wait_seconds=55)", ai_task_runs.status("Claude")["next"])
|
|
history = {"p1": {"status": {"completed": True}, "outputs": {
|
|
"9": {"images": [{"filename": "a.png", "subfolder": "x", "type": "output"}]},
|
|
"12": {"images": [{"filename": "pose_00001_.png", "type": "temp"}]}}}} # a PreviewImage: the AI's to look at
|
|
threading.Timer(0.3, ai_task_runs._finish_run, (task_id, step_id, ["p1"], history)).start()
|
|
status = ai_task_runs.status("Claude", wait_seconds=5)
|
|
self.assertEqual((status["state"], status["pictures"], status["previews"]), ("done", ["x/a.png"], ["pose_00001_.png"]))
|
|
self.assertIn("folder='temp'", status["previews_note"])
|
|
self.assertEqual([item["filename"] for item in ai_tasks.get_task(task_id)["steps"][0]["results"]], ["a.png"], "previews are no results")
|
|
self.assertNotIn("progress", ai_tasks.get_task(task_id)["steps"][0], "a finished run has no progress left")
|
|
|
|
def test_a_run_saves_only_when_something_changes(self):
|
|
ai_tasks.record("A", "run_workflow", {}, ok({}), time.time())
|
|
task_id = ai_tasks.list_tasks()[0]["id"]
|
|
step_id = ai_tasks.get_task(task_id)["steps"][0]["id"]
|
|
with mock.patch.object(ai_tasks, "_save") as saved:
|
|
ai_task_runs._finish_run(task_id, step_id, ["p1"], {})
|
|
ai_task_runs._finish_run(task_id, step_id, ["p1"], {})
|
|
self.assertEqual(saved.call_count, 1)
|
|
|
|
def test_no_run_yet(self):
|
|
ai_tasks.record("A", "describe_canvas", {}, ok({}), time.time(), True)
|
|
self.assertEqual(ai_task_runs.status("A")["state"], "none") # an answer: the task shows no failed step
|
|
|
|
def run_step(self):
|
|
step = ai_tasks.record("A", "run_workflow", {}, ok({"queued": 1}), time.time())
|
|
return ai_tasks.active_task_id("A"), step["id"]
|
|
|
|
def test_no_queued_prompt_is_not_success(self):
|
|
task, step = self.run_step()
|
|
ai_task_runs._finish_run(task, step, [], {})
|
|
self.assertEqual(ai_task_runs.status("A")["state"], "stopped")
|
|
self.assertFalse(ai_tasks.get_task(task)["steps"][-1]["ok"])
|
|
|
|
def test_failure_is_recorded_even_when_other_prompts_are_pending(self):
|
|
task, step = self.run_step()
|
|
ai_task_runs._finish_run(task, step, ["one", "two"], {})
|
|
ai_task_runs._finish_run(task, step, ["one", "two"], {"one": {"status": {"status_str": "error"}}})
|
|
self.assertFalse(ai_tasks.get_task(task)["steps"][-1]["ok"])
|
|
self.assertIn("error", ai_task_runs.status("A"))
|
|
|
|
def test_watcher_timeout_and_history_failure_stop_the_receipt(self):
|
|
# Run the actual watcher synchronously; replace only time and host queue boundaries.
|
|
def inline_thread(target, **kwargs):
|
|
return SimpleNamespace(start=target)
|
|
for fail_history in (False, True):
|
|
isolate()
|
|
task, step = self.run_step()
|
|
queue = mock.Mock()
|
|
queue.get_history.side_effect = OSError("unavailable")
|
|
with mock.patch.object(ai_task_runs.threading, "Thread", side_effect=inline_thread), \
|
|
mock.patch.object(ai_task_runs, "_queue", return_value=queue), \
|
|
mock.patch.object(ai_tasks, "RUN_WATCH_SECONDS", 60 if fail_history else 0):
|
|
ai_task_runs.watch(step, {"prompt_ids": ["mine"]})
|
|
status = ai_task_runs.status("A")
|
|
self.assertEqual(status["state"], "stopped")
|
|
self.assertIn("history" if fail_history else "timed out", status["error"])
|
|
|
|
def test_batch_progress_keeps_total_and_current_index(self):
|
|
task, step = self.run_step()
|
|
queue = Queue(running=[(2, "two")])
|
|
with mock.patch.object(ai_task_runs, "_queue", return_value=queue), mock.patch.object(ai_task_runs, "_registry", return_value=None):
|
|
self.assertEqual(ai_task_runs.progress_of(["one", "two"]), {"state": "running", "prompt": 2, "prompts": 2})
|
|
|
|
def test_two_tasks_collect_only_their_receipts_and_keep_batch_order(self):
|
|
def record(filename):
|
|
return {"status": {"completed": True}, "outputs": {"1": {"images": [{"filename": filename, "type": "output"}]}}}
|
|
history = {pid: record(pid + ".png") for pid in ("z-first", "a-second", "other-ai", "manual")}
|
|
queue = mock.Mock()
|
|
queue.get_history.side_effect = lambda prompt_id: {prompt_id: history[prompt_id]}
|
|
with mock.patch.object(ai_task_runs.threading, "Thread", side_effect=lambda target, **kw: SimpleNamespace(start=target)), \
|
|
mock.patch.object(ai_task_runs, "_queue", return_value=queue):
|
|
task, step = self.run_step()
|
|
ai_task_runs.watch(step, {"prompt_ids": ["z-first", "a-second"]})
|
|
second = ai_tasks.begin_task("B", "other")["id"]
|
|
other = ai_tasks.record("B", "run_workflow", {}, ok({}), time.time())
|
|
ai_task_runs.watch(other["id"], {"prompt_ids": ["other-ai"], "submission_error": "Only 1 of 2 prompts accepted"})
|
|
self.assertEqual(ai_task_runs.status("A")["pictures"], ["z-first.png", "a-second.png"])
|
|
self.assertEqual(ai_task_runs.status("B")["pictures"], ["other-ai.png"])
|
|
self.assertEqual(ai_task_runs.status("B")["state"], "error")
|
|
self.assertNotIn(mock.call(prompt_id="manual"), queue.get_history.call_args_list)
|
|
|
|
|
|
class NewsTests(unittest.TestCase):
|
|
"""The AI hears where its run stands in its other answers, and once that it ended."""
|
|
|
|
def setUp(self):
|
|
isolate()
|
|
self.step = ai_tasks.record("A", "run_workflow", {}, ok({"queued": 1}), time.time())
|
|
self.task = ai_tasks.active_task_id("A")
|
|
with ai_tasks._lock:
|
|
task = ai_tasks._load()[self.task]
|
|
task["steps"][-1]["prompt_ids"] = ["p1"]
|
|
ai_tasks._save(task)
|
|
|
|
def told(self):
|
|
return ai_task_runs.tell({"content": [{"type": "text", "text": "{}"}], "structuredContent": {"n": 1}}, self.task)
|
|
|
|
def test_while_it_goes_every_answer_says_where_it_stands(self):
|
|
with mock.patch.object(ai_tasks, "_notify"):
|
|
ai_task_runs._show_progress(self.task, self.step["id"], {"state": "running", "prompt": 1, "prompts": 1, "node": "KSampler",
|
|
"value": 14, "max": 28, "percent": 50, "nodes_done": 3, "nodes": 9})
|
|
for _ in range(2):
|
|
result = self.told()
|
|
run = result["structuredContent"]["your_run"]
|
|
self.assertEqual((run["state"], run["progress"]["percent"]), ("running", 50))
|
|
self.assertIn("KSampler step 14/28, 50%, 3 of 9 nodes done", run["summary"])
|
|
self.assertIn("your_run", result["content"][0]["text"], "the text says it too")
|
|
self.assertEqual(result["structuredContent"]["n"], 1, "the tool's own answer stays")
|
|
|
|
def test_the_end_is_told_once(self):
|
|
history = {"p1": {"status": {"completed": True}, "outputs": {"9": {"images": [{"filename": "a.png", "type": "output"}]}}}}
|
|
with mock.patch.object(ai_tasks, "_notify"):
|
|
ai_task_runs._finish_run(self.task, self.step["id"], ["p1"], history)
|
|
self.assertIn("finished", self.told()["structuredContent"]["your_run"]["summary"])
|
|
self.assertNotIn("your_run", self.told()["structuredContent"])
|
|
|
|
def test_run_status_telling_the_end_is_enough(self):
|
|
with mock.patch.object(ai_tasks, "_notify"):
|
|
ai_task_runs._finish_run(self.task, self.step["id"], ["p1"], {"p1": {"status": {"completed": True}}})
|
|
self.assertIn("finished", ai_task_runs.status("A")["summary"])
|
|
self.assertNotIn("your_run", self.told()["structuredContent"])
|
|
|
|
def test_the_pictures_come_once_with_the_answer_that_tells_the_end(self):
|
|
history = {"p1": {"status": {"completed": True}, "outputs": {"9": {"images": [
|
|
{"filename": f"{n}.png", "subfolder": "x", "type": "output"} for n in range(6)]}}}}
|
|
with mock.patch.object(ai_tasks, "_notify"):
|
|
ai_task_runs._finish_run(self.task, self.step["id"], ["p1"], history)
|
|
picture = {"type": "image", "data": "AAAA", "mimeType": "image/webp"}
|
|
with mock.patch.object(ai_task_runs, "_resolve_image", side_effect=lambda rel: (rel, "/out/" + rel)), mock.patch.object(ai_task_runs, "_image_content", return_value=picture):
|
|
before = ai_task_runs.show_pictures({"content": [{"type": "text", "text": "{}"}], "structuredContent": {}}, self.task)
|
|
self.assertEqual(len(before["content"]), 1, "nothing before the AI is told it ended")
|
|
result = ai_task_runs.show_pictures(self.told(), self.task)
|
|
self.assertEqual(result["structuredContent"]["pictures_attached"], ["x/0.png", "x/1.png", "x/2.png", "x/3.png"])
|
|
self.assertEqual([item["type"] for item in result["content"]], ["text"] + ["image"] * ai_task_runs.MAX_SHOWN)
|
|
self.assertIn("pictures_attached", result["content"][0]["text"])
|
|
again = ai_task_runs.show_pictures({"content": [{"type": "text", "text": "{}"}], "structuredContent": {}}, self.task)
|
|
self.assertEqual(len(again["content"]), 1, "once per run")
|
|
|
|
def test_a_line_to_paste_until_the_ai_knows(self):
|
|
with mock.patch.object(ai_tasks, "_notify"), mock.patch.object(ai_gates, "run_wait", return_value="call"):
|
|
ai_task_runs._finish_run(self.task, self.step["id"], ["p1"], {"p1": {"status": {"completed": True}}})
|
|
self.assertTrue(ai_tasks.get_task(self.task)["steps"][-1]["call_ai"], "the workbench offers the line")
|
|
self.told()
|
|
self.assertNotIn("call_ai", ai_tasks.get_task(self.task)["steps"][-1], "gone once the AI is told")
|
|
|
|
def test_no_line_when_the_ai_waits_or_a_new_run_starts(self):
|
|
with mock.patch.object(ai_tasks, "_notify"), mock.patch.object(ai_gates, "run_wait", return_value="wait"):
|
|
ai_task_runs._finish_run(self.task, self.step["id"], ["p1"], {"p1": {"status": {"completed": True}}})
|
|
self.assertNotIn("call_ai", ai_tasks.get_task(self.task)["steps"][-1])
|
|
with ai_tasks._lock:
|
|
task = ai_tasks._load()[self.task]
|
|
task["steps"][-1]["call_ai"] = True
|
|
ai_tasks._save(task)
|
|
newer = ai_tasks.record("A", "run_workflow", {}, ok({"queued": 1}), time.time())
|
|
with mock.patch("threading.Thread"):
|
|
ai_task_runs.watch(newer["id"], {"prompt_ids": ["p2"]})
|
|
self.assertNotIn("call_ai", ai_tasks.get_task(self.task)["steps"][0], "an older run offers no line once another starts")
|
|
|
|
def test_an_old_run_told_before_brings_no_pictures(self):
|
|
with mock.patch.object(ai_tasks, "_notify"):
|
|
ai_task_runs._finish_run(self.task, self.step["id"], ["p1"], {"p1": {"status": {"completed": True}, "outputs": {
|
|
"9": {"images": [{"filename": "a.png", "type": "output"}]}}}})
|
|
with ai_tasks._lock:
|
|
task = ai_tasks._load()[self.task]
|
|
task["steps"][-1]["told"] = True # told before pictures came with it: no "shown" mark
|
|
ai_tasks._save(task)
|
|
result = ai_task_runs.show_pictures({"content": [{"type": "text", "text": "{}"}], "structuredContent": {}}, self.task)
|
|
self.assertEqual(len(result["content"]), 1)
|
|
|
|
def test_no_run_no_news(self):
|
|
ai_tasks.begin_task("A", "other", new_task=True)
|
|
self.assertIsNone(ai_task_runs.news(ai_tasks.active_task_id("A")))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|