Files
ComfyAssets-ComfyUI_PromptM…/tests/test_execution_output_hooks.py
T
Vito b1e7abd380 fix(threads): stop and join the prompt cleanup thread, the image worker and the observer (#302, #370) (#372)
* fix(threads): stop and join the prompt cleanup thread, the image worker and the observer (#302, #370)

- PromptTracker: the cleanup thread waits on a stop event instead of
  time.sleep; new shutdown(timeout=2) joins it and returns whether it ended.
- ImageGenerationHandler.shutdown(timeout=5) now joins the worker and
  returns a bool; new close() is final: schedule() refuses and a sealed
  run's history re-check is dropped, so no worker restarts afterwards.
- ImageMonitor.stop_monitoring(timeout=5) closes the handler first, then
  stops and joins the watchdog observer within one deadline.
- Every test that builds a real tracker, handler or monitor ends with its
  final stop and an assertion, before its folder and databases go.
- #370: the wrapper/session history test drains before dispose(); the
  held-write bound test uses 30 s ceilings instead of tight wall clocks.

* test(history): drain the writer before asserting a bounded close it does not test (#370)

* fix(threads): a later stop waits again for a worker or observer an earlier stop left running (#302)

Review findings on #372:
- ImageGenerationHandler keeps the workers a stop signalled that had not
  ended; every later shutdown/close joins them again, so a second call no
  longer returns True while the first worker still runs.
- ImageMonitor.stop_monitoring keeps an observer or handler it left
  running and stops and joins it again on the next call.
- A closed handler takes no awaiting history path, and a re-check that
  raced close() is not kept due.
- Tests: second stops of a busy worker and a live observer, the closed
  handler's re-check paths, the handle test asserts the worker ended and
  stops before it asserts, and a no-new-thread check instead of set
  equality.

* test(browser): the System reset restores every settings value, safety_backup included

The frontend coverage shard of #372 failed in workspace-system-tagging:
the review read "No safety backup (the policy is off)". The backups spec
saves safety_backup=false, and /__test__/system/reset restored config.json
but not PromptManagerConfig.SAFETY_BACKUP (nor INFINITE_SCROLL), so a
later spec on the same worker server saw the policy off. The reset now
captures and restores every non-tagging value POST /settings can change.
2026-10-27 02:15:23 -07:00

406 lines
16 KiB
Python

"""Exact successful host callbacks persist independently bounded output evidence."""
from dataclasses import replace
import os
from pathlib import Path
import threading
import sys
from types import SimpleNamespace
from unittest.mock import patch
from utils.execution_history import HistoryRegistration
from utils.execution_context_budget import ContextBudget
from utils.execution_outputs import snapshot_output_roots
from utils.queue_observers import subscribe_dequeue, subscribe_task_done
try:
from tests import test_execution_history_hooks as fixtures
except ImportError:
import test_execution_history_hooks as fixtures
def history(filename="a.png"):
return {
"outputs": {
"save": {
"images": [{"filename": filename, "subfolder": "", "type": "output"}]
}
}
}
class OutputHooksTests(fixtures.unittest.TestCase):
setUp = fixtures.HistoryHooksTests.setUp
restore_logger = fixtures.HistoryHooksTests.restore_logger
cleanup_registrations = fixtures.HistoryHooksTests.cleanup_registrations
wait_for = fixtures.HistoryHooksTests.wait_for
dequeue = fixtures.HistoryHooksTests.dequeue
rows = fixtures.HistoryHooksTests.rows
drain = fixtures.HistoryHooksTests.drain
def register(self, roots_reader=None, pool=None):
if roots_reader is None:
def roots_reader():
return snapshot_output_roots("posix", "/output", "/temp")
if pool is None:
registration = HistoryRegistration(
self.queue, fixtures.JOB, lambda: self.db, roots_reader=roots_reader
)
else:
from utils.execution_history_writer import HistoryWriter
with patch(
"utils.execution_history.HistoryWriter",
lambda factory, **kw: HistoryWriter(factory, context_budget=pool, **kw),
):
registration = HistoryRegistration(
self.queue, fixtures.JOB, lambda: self.db, roots_reader=roots_reader
)
self.registrations.append(registration)
self.wait_for(lambda: registration.writer.status().accepting)
return registration
def outputs(self, item_id):
row = next(row for row in self.rows() if row["queue_item_id"] == item_id)
return self.db.get_execution_outputs(row["run_id"])
def test_positional_keyword_and_text_only_outputs_use_exact_fresh_attempt(self):
registration = self.register()
for i in range(3):
self.dequeue(i)
payload = history("a.png" if i < 2 else "b.png")
if i == 1:
self.queue.task_done(
item_id=i,
history_result=payload,
status=fixtures.status("execution_error"),
)
else:
self.assertIs(
self.queue.task_done(i, payload, fixtures.status()),
self.queue.done_result,
)
self.drain(registration)
self.assertEqual(
[self.outputs(i)["references"][0]["filename"] for i in range(3)],
["a.png", "a.png", "b.png"],
)
self.assertEqual(
[r["outcome"] for r in self.rows()], ["succeeded", "failed", "succeeded"]
)
self.assertEqual(registration.writer.context_budget.charged_bytes, 0)
def test_budget_denied_context_still_captures_independently_affordable_outputs(
self,
):
pool = ContextBudget(2048)
registration = self.register(pool=pool)
self.dequeue(0)
self.queue.task_done(0, history(), fixtures.status())
self.drain(registration)
self.assertEqual(self.outputs(0)["state"], "complete")
self.assertEqual(
self.db.get_execution_contexts(self.rows()[0]["run_id"])["reason"],
"budget_exhausted",
)
self.assertEqual(pool.charged_bytes, 0)
def test_missing_history_and_explicit_none_are_distinct_from_empty_evidence(self):
registration = self.register()
for i in range(4):
self.dequeue(i)
if i == 0:
self.queue.task_done(item_id=i, status=fixtures.status())
else:
self.queue.task_done(
i,
[None, {"outputs": {}}, {"outputs": {"text": {"text": ["ok"]}}}][
i - 1
],
fixtures.status(),
)
self.drain(registration)
self.assertEqual(
[(self.outputs(i)["state"], self.outputs(i)["reason"]) for i in range(4)],
[
("unavailable", "unavailable"),
("unavailable", "invalid_evidence"),
("complete", "none_reported"),
("complete", "none_reported"),
],
)
def test_replay_old_dequeue_reused_uuid_and_forged_capability_cannot_rebind(self):
registration = self.register()
old = self.dequeue(1)
self.queue.task_done(1, history("first.png"), fixtures.status())
self.dequeue(2)
registration._dequeue(old)
self.queue.task_done(1, history("replayed.png"), fixtures.status())
self.queue.task_done(2, history("second.png"), fixtures.status())
self.dequeue(3)
cap = registration._capabilities[3]
registration._capabilities[3] = replace(cap)
self.queue.task_done(3, history("forged.png"), fixtures.status())
registration.capture.revoke(cap)
self.drain(registration)
self.assertEqual(self.outputs(1)["references"][0]["filename"], "first.png")
self.assertEqual(self.outputs(2)["references"][0]["filename"], "second.png")
self.assertEqual(self.outputs(3)["state"], "not_captured")
self.assertEqual(registration.writer.context_budget.charged_bytes, 0)
def test_root_snapshot_mutation_and_reader_errors_preserve_outcome_and_reference(
self,
):
roots = [snapshot_output_roots("posix", "/before", "/temp")]
def reader():
if isinstance(roots[0], Exception):
raise roots[0]
return roots[0]
registration = self.register(reader)
self.dequeue(0)
roots[0] = snapshot_output_roots("posix", "/after", "/temp")
self.queue.task_done(0, history(), fixtures.status())
self.dequeue(1)
roots[0] = RuntimeError("private")
self.queue.task_done(1, history(), fixtures.status())
self.dequeue(2)
self.queue.task_done(2, history(), fixtures.status())
self.drain(registration)
for i in range(3):
self.assertEqual(self.outputs(i)["reason"], "root_unavailable")
self.assertIsNone(self.outputs(i)["references"][0]["root_key"])
self.assertEqual([r["outcome"] for r in self.rows()], ["succeeded"] * 3)
def test_unexpected_output_producer_error_preserves_sealed_context_and_outcome(
self,
):
registration = self.register()
graph = {
"p": {
"class_type": "PromptManagerPair",
"inputs": {"text": "raw", "negative_text": "neg"},
},
"s": {
"class_type": "KSampler",
"inputs": {"positive": ["p", 0], "negative": ["p", 1]},
},
}
self.queue.results.append(((0, fixtures.JOB, graph, {}, ["s"]), 0))
self.queue.get()
with patch(
"utils.execution_output_capture._extract",
side_effect=RuntimeError("private"),
):
self.queue.task_done(0, history(), fixtures.status())
self.drain(registration)
self.assertEqual(self.rows()[0]["outcome"], "succeeded")
self.assertEqual(self.outputs(0)["reason"], "invalid_evidence")
self.assertEqual(
self.db.get_execution_contexts(self.rows()[0]["run_id"])["reason"],
"partial",
)
self.assertEqual(
self.db.get_execution_contexts(self.rows()[0]["run_id"])["contexts"][0][
"positive"
]["raw"],
"raw",
)
self.assertEqual(registration.writer.context_budget.charged_bytes, 0)
def test_unexpected_context_seal_and_projection_errors_do_not_lose_output(self):
registration = self.register()
with patch(
"utils.execution_context_capture.project_contexts",
side_effect=RuntimeError("private"),
):
self.dequeue(0)
self.queue.task_done(0, history(), fixtures.status())
self.dequeue(1)
with patch(
"utils.execution_context_capture.bundle_payload_size",
side_effect=RuntimeError("private"),
):
self.queue.task_done(1, history(), fixtures.status())
self.drain(registration)
self.assertEqual([self.outputs(i)["state"] for i in range(2)], ["complete"] * 2)
self.assertEqual([r["outcome"] for r in self.rows()], ["succeeded"] * 2)
self.assertEqual(registration.writer.context_budget.charged_bytes, 0)
def test_original_failure_missing_callback_disposal_and_observer_coexistence(self):
calls = []
before_get, before_done = self.queue.get, self.queue.task_done
unsubscribe_get = subscribe_dequeue(
self.queue, "other", lambda result: calls.append(result[1])
)
unsubscribe_done = subscribe_task_done(
self.queue, "other", lambda args, kwargs: calls.append(args[0])
)
self.addCleanup(unsubscribe_get)
self.addCleanup(unsubscribe_done)
registration = self.register()
self.dequeue(0)
late = registration._task_done
# Drain first: the 2 s bound is not under test here (#370).
self.drain(registration)
self.assertTrue(registration.dispose())
thread = threading.Thread(
target=lambda: late((0, history(), fixtures.status()), {})
)
thread.start()
thread.join(2)
self.assertFalse(thread.is_alive())
self.queue.task_done(0, history(), fixtures.status())
self.assertEqual(calls, [0, 0])
self.assertEqual(self.outputs(0)["state"], "not_captured")
unsubscribe_get()
unsubscribe_done()
self.assertEqual(self.queue.get, before_get)
self.assertEqual(self.queue.task_done, before_done)
self.assertEqual(registration.writer.context_budget.charged_bytes, 0)
def test_rejected_invalid_revoked_or_unfresh_dequeue_never_reads_roots(self):
reads = []
def reader():
reads.append(1)
return snapshot_output_roots("posix", "/output", "/temp")
registration = self.register(reader)
self.dequeue(10)
self.dequeue(9)
self.queue.task_done(9, history(), fixtures.status())
self.queue.task_done(10, history(), fixtures.status())
self.dequeue(11, job="invalid")
self.queue.task_done(11, history(), fixtures.status())
self.drain(registration)
self.assertEqual(reads, [1]) # 10 was revoked by the next actual dequeue.
self.assertEqual(self.outputs(9)["state"], "not_captured")
self.assertEqual(self.outputs(10)["state"], "not_captured")
self.assertEqual(registration.writer.context_budget.charged_bytes, 0)
def test_original_failures_preserve_exception_and_leave_outputs_missing(
self,
):
error = RuntimeError("host error")
def fail(*args, **kwargs):
raise error
original_get = self.queue.get
self.queue.get = fail
registration = self.register()
with self.assertRaises(RuntimeError) as raised:
self.queue.get()
self.assertIs(raised.exception, error)
self.assertEqual(self.rows(), [])
self.assertTrue(registration.dispose())
self.queue.get = original_get
self.queue.task_done = fail
registration = self.register()
self.dequeue(1)
with self.assertRaises(RuntimeError) as raised:
self.queue.task_done(1, history(), fixtures.status())
self.assertIs(raised.exception, error)
self.drain(registration)
self.assertEqual(self.rows()[0]["outcome"], "observed")
self.assertEqual(self.outputs(1)["state"], "not_captured")
self.assertTrue(registration.dispose())
self.assertEqual(registration.writer.context_budget.charged_bytes, 0)
def test_default_adapter_uses_configured_output_and_temp_namespaces(self):
native_root = Path(self.directory.name)
cases = (
(
os.name,
str(native_root / "output"),
str(native_root / "temp"),
"windows" if os.name == "nt" else "posix",
),
("posix", "/configured", "/configured", "posix"),
("nt", r"D:\configured", r"D:\configured", "windows"),
)
for item, (platform, output_root, temp_root, style) in enumerate(cases):
folder_paths = SimpleNamespace(
get_output_directory=lambda: output_root,
get_temp_directory=lambda: temp_root,
)
with (
self.subTest(platform=platform, root=output_root),
patch.dict(sys.modules, {"folder_paths": folder_paths}),
patch("utils.execution_history.os", SimpleNamespace(name=platform)),
):
registration = HistoryRegistration(
self.queue, fixtures.JOB, lambda: self.db
)
self.registrations.append(registration)
self.wait_for(lambda: registration.writer.status().accepting)
self.dequeue(item)
payload = history()
payload["outputs"]["save"]["images"].append(
{"filename": "a.png", "subfolder": "", "type": "temp"}
)
self.queue.task_done(item, payload, fixtures.status())
self.drain(registration)
evidence = self.outputs(item)
self.assertEqual(evidence["state"], "complete")
references = evidence["references"]
self.assertEqual(len(references[0]["root_key"]), 64)
self.assertEqual(len(references[1]["root_key"]), 64)
self.assertEqual([r["path_style"] for r in references], [style] * 2)
self.assertEqual(
[r["folder_type"] for r in references], ["output", "temp"]
)
self.assertNotEqual(
references[0]["root_key"], references[1]["root_key"]
)
self.assertEqual(registration._roots, {})
self.assertTrue(registration.dispose())
def test_sealed_context_is_discarded_before_output_reservation(self):
pool = ContextBudget(1048576)
registration = self.register(pool=pool)
self.dequeue(0)
self.assertEqual(pool.charged_bytes, 1048576)
self.queue.task_done(0, history(), fixtures.status())
self.drain(registration)
self.assertEqual(self.outputs(0)["state"], "complete")
self.assertEqual(pool.charged_bytes, 0)
def test_saturated_start_cannot_capture_outputs_or_retain_previous_roots(self):
entered, release = threading.Event(), threading.Event()
self.releases.append(release)
original = self.db.write_execution_start
def blocked(command):
entered.set()
release.wait(3)
return original(command)
with patch.object(self.db, "write_execution_start", blocked):
registration = self.register()
self.dequeue(0)
self.assertTrue(entered.wait(1))
from utils.execution_history_events import make_start
epoch = registration.writer.status().epoch
for item in range(1023):
command = make_start(
registration._session_id, fixtures.item(1000 + item), epoch, 0
)
self.assertTrue(registration.writer.submit(command))
self.dequeue(1)
self.assertEqual(registration._roots, {})
self.queue.task_done(0, history("revoked.png"), fixtures.status())
self.queue.task_done(1, history("rejected.png"), fixtures.status())
registration.writer.invalidate()
release.set()
self.assertTrue(registration.dispose(2))
self.assertEqual(self.outputs(0)["state"], "not_captured")
self.assertEqual(len(self.rows()), 1)
self.assertEqual(registration.writer.context_budget.charged_bytes, 0)