Files
ComfyAssets-ComfyUI_PromptM…/tests/test_execution_output_hooks.py
T
Vito 4ef236aaf7 feat(history): persist exact run output references (#214)
* feat(history): define bounded output reference contracts

* feat(database): migrate execution output evidence to schema v3

* feat(history): persist immutable run output references

* test(history): isolate mocked node gallery monitor

* feat(history): capture bounded output descriptors with shared leases

* fix(history): authenticate payload owners and release failed captures

* feat(history): attach output evidence to exact finalized attempts

* test(history): verify run output evidence on real ComfyUI

* fix(history): release idle worker payload references

* docs(history): clarify defensive lease cleanup

* test(history): align output fixtures with native path styles

* test(history): separate restore completion from writer drain

* test(history): bound batch drains by progress and total time
2026-10-17 23:26:37 -07:00

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