Files
ComfyAssets-ComfyUI_PromptM…/tests/test_execution_output_database.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

429 lines
16 KiB
Python

"""Frozen output occurrences belong to an exact stored execution attempt."""
import sqlite3
import unittest
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from dataclasses import replace
from threading import Barrier
from unittest.mock import patch
from database.execution_history import HistoryEpochChanged
from utils.execution_outputs import OutputBundle, OutputReference
try:
from tests import test_execution_history_database as history_tests
except ImportError:
import test_execution_history_database as history_tests
def bundle(
filename="first.png", *, root="a" * 64, node="9", ordinal=0, path_style="posix"
):
return OutputBundle(
(OutputReference(node, ordinal, "output", "", filename, path_style, root),),
"complete" if root else "partial",
"reported" if root else "root_unavailable",
)
class OutputDatabaseTests(unittest.TestCase):
setUp = history_tests.HistoryDatabaseTests.setUp
close_databases = history_tests.HistoryDatabaseTests.close_databases
database = history_tests.HistoryDatabaseTests.database
start = history_tests.HistoryDatabaseTests.start
finish = history_tests.HistoryDatabaseTests.finish
def capture(self, db, value=None, item=0):
start = self.start(db, item)
db.write_execution_start(start)
self.assertTrue(
db.write_execution_outputs_for_finish(self.finish(start), value or bundle())
)
return start
def rows(self, db):
return [
tuple(row)
for row in db.model.get_connection().execute(
"SELECT * FROM execution_run_outputs ORDER BY run_id,"
"output_node_id,image_ordinal"
)
]
def test_exact_attempt_parent_not_recomputed_uuid(self):
db = self.database()
stored_id = "55555555-5555-4555-8555-555555555555"
start = replace(self.start(db), run_id=stored_id)
db.write_execution_start(start)
self.assertTrue(
db.write_execution_outputs_for_finish(self.finish(start), bundle())
)
self.assertEqual(self.rows(db)[0][0], stored_id)
self.assertEqual(db.get_execution_outputs(stored_id)["state"], "complete")
other = self.capture(db, item=1)
self.assertEqual(len(self.rows(db)), 2)
self.assertEqual(
db.get_execution_outputs(other.run_id)["references"][0]["filename"],
"first.png",
)
self.assertEqual(db.get_execution_run(stored_id)["outcome"], "observed")
def test_exact_replay_noop_and_conflict_sticky(self):
db = self.database()
first = bundle()
refs = (replace(first.references[0], output_node_id="2"),) + (
replace(first.references[0], image_ordinal=0),
replace(first.references[0], image_ordinal=1),
)
first = replace(first, references=refs)
start = self.capture(db, first)
finish = self.finish(start)
original = self.rows(db)
# Construct identical DTOs from dictionaries with different insertion orders.
reordered = tuple(
OutputReference(**dict(reversed(list(vars(ref).items())))) for ref in refs
)
self.assertTrue(
db.write_execution_outputs_for_finish(
finish, replace(first, references=reordered)
)
)
self.assertEqual(db.get_execution_outputs(start.run_id)["state"], "complete")
self.assertEqual(self.rows(db), original)
self.assertTrue(
db.write_execution_outputs_for_finish(finish, bundle("other.png"))
)
capture = db.get_execution_outputs(start.run_id)
self.assertEqual(
(capture["state"], capture["reason"]), ("conflicting", "conflicting")
)
self.assertEqual(
[r["filename"] for r in capture["references"]], ["first.png"] * 3
)
self.assertEqual(capture["references"][0]["context_membership"], "unknown")
for replay in (first, bundle("third.png")):
db.write_execution_outputs_for_finish(finish, replay)
self.assertEqual(db.get_execution_outputs(start.run_id), capture)
self.assertEqual(self.rows(db), original)
def test_header_and_every_reference_scalar_participate_in_replay(self):
variants = (
replace(bundle(), state="partial", reason="unsupported"),
bundle(root="b" * 64),
bundle(node="8"),
bundle(ordinal=1),
replace(
bundle(),
references=(replace(bundle().references[0], folder_type="temp"),),
),
replace(
bundle(),
references=(replace(bundle().references[0], subfolder="nested"),),
),
replace(
bundle(),
references=(replace(bundle().references[0], path_style="windows"),),
),
OutputBundle((), "complete", "none_reported"),
)
db = self.database()
for item, variant in enumerate(variants):
with self.subTest(item=item):
start = self.capture(db, item=item)
db.write_execution_outputs_for_finish(self.finish(start), variant)
capture = db.get_execution_outputs(start.run_id)
self.assertEqual(capture["state"], "conflicting")
self.assertEqual(capture["references"][0]["root_key"], "a" * 64)
self.assertEqual(capture["references"][0]["output_node_id"], "9")
def test_concurrent_replay_keeps_first_rows(self):
db = self.database()
start = self.start(db)
db.write_execution_start(start)
finish = self.finish(start)
peers = [self.database() for _ in range(4)]
barrier = Barrier(len(peers))
def write(index):
barrier.wait(timeout=5)
return peers[index].write_execution_outputs_for_finish(
finish, bundle("first.png" if index % 2 else "second.png")
)
with ThreadPoolExecutor(max_workers=4) as executor:
self.assertEqual(list(executor.map(write, range(4))), [True] * 4)
capture = db.get_execution_outputs(start.run_id)
self.assertEqual(
(capture["state"], capture["reason"]), ("conflicting", "conflicting")
)
self.assertEqual(len(capture["references"]), 1)
frozen = self.rows(db)
for value in (bundle(), bundle("second.png")):
db.write_execution_outputs_for_finish(finish, value)
self.assertEqual(self.rows(db), frozen)
def test_epoch_and_generation_guard_cover_commit(self):
db = self.database()
start = self.start(db)
db.write_execution_start(start)
observations = []
conn = db.model.get_connection()
active = False
@contextmanager
def guard():
nonlocal active
self.assertTrue(conn.in_transaction)
active = True
try:
yield True
finally:
observations.append(("exit", conn.in_transaction))
active = False
def trace(sql):
if sql == "COMMIT":
observations.append(("commit", active))
conn.set_trace_callback(trace)
self.addCleanup(conn.set_trace_callback, None)
self.assertTrue(
db.write_execution_outputs_for_finish(
self.finish(start), bundle(), generation_guard=guard
)
)
self.assertEqual(observations, [("commit", True), ("exit", False)])
self.assertFalse(active)
@contextmanager
def stale():
yield False
other = self.start(db, item=1)
db.write_execution_start(other)
self.assertFalse(
db.write_execution_outputs_for_finish(
self.finish(other), bundle(), generation_guard=stale
)
)
self.assertEqual(
db.get_execution_outputs(other.run_id)["state"], "not_captured"
)
old_epoch = start.storage_epoch
with conn:
conn.execute(
"UPDATE execution_history_metadata SET storage_epoch=?",
("66666666-6666-4666-8666-666666666666",),
)
with self.assertRaises(HistoryEpochChanged):
db.write_execution_outputs_for_finish(
self.finish(start), bundle("late.png")
)
self.assertEqual(db.get_execution_run(start.run_id)["storage_epoch"], old_epoch)
self.assertEqual(
db.get_execution_outputs(start.run_id)["references"][0]["filename"],
"first.png",
)
def test_failure_rolls_back_header_and_keeps_guard_until_rollback(self):
db = self.database()
start = self.start(db)
db.write_execution_start(start)
conn = db.model.get_connection()
with conn:
conn.execute(
"CREATE TRIGGER reject_output BEFORE INSERT ON execution_run_outputs "
"BEGIN SELECT RAISE(ABORT,'injected'); END"
)
states = []
@contextmanager
def guard():
states.append(conn.in_transaction)
try:
yield True
finally:
states.append(conn.in_transaction)
with self.assertRaises(sqlite3.IntegrityError):
db.write_execution_outputs_for_finish(
self.finish(start), bundle(), generation_guard=guard
)
self.assertEqual(states, [True, False])
self.assertEqual(
db.get_execution_outputs(start.run_id)["state"], "not_captured"
)
self.assertEqual(self.rows(db), [])
def test_missing_header_or_parent_rejected(self):
db = self.database()
start = self.start(db)
self.assertFalse(
db.write_execution_outputs_for_finish(self.finish(start), bundle())
)
self.assertEqual(
db.get_execution_outputs(start.run_id),
{
"state": "not_captured",
"reason": None,
"format_version": None,
"source": None,
"references": [],
},
)
db.write_execution_start(start)
conn = db.model.get_connection()
conn.execute("PRAGMA foreign_keys=OFF")
self.assertEqual(conn.execute("PRAGMA foreign_keys").fetchone()[0], 0)
with conn:
conn.execute(
"INSERT INTO execution_run_outputs "
"VALUES(?, '9',0,'output','','orphan.png','posix',?)",
(start.run_id, "a" * 64),
)
with self.assertRaises(sqlite3.DatabaseError):
db.get_execution_outputs(start.run_id)
with self.assertRaises(sqlite3.DatabaseError):
db.write_execution_outputs_for_finish(self.finish(start), bundle())
self.assertEqual(
conn.execute("SELECT count(*) FROM execution_output_capture").fetchone()[0],
0,
)
with conn:
conn.execute("DELETE FROM execution_run_outputs")
db.write_execution_outputs_for_finish(self.finish(start), bundle())
with conn:
conn.execute("DELETE FROM execution_runs")
with self.assertRaises(sqlite3.DatabaseError):
db.get_execution_outputs(start.run_id)
def test_reads_never_access_files_or_rebind(self):
db = self.database()
value = bundle(root=None)
value = replace(
value,
references=(
replace(
value.references[0],
folder_type="temp",
subfolder=" Exact Ω ",
filename="Keep Case.png",
),
),
)
start = self.capture(db, value)
conn = db.model.get_connection()
before = list(conn.iterdump())
with (
patch("builtins.open", side_effect=AssertionError("file read")),
patch("os.stat", side_effect=AssertionError("stat")),
patch("os.path.realpath", side_effect=AssertionError("resolve")),
):
capture = db.get_execution_outputs(start.run_id)
self.assertEqual(
capture,
{
"state": "partial",
"reason": "root_unavailable",
"format_version": 1,
"source": "host_final_history",
"references": [
{
"output_node_id": "9",
"image_ordinal": 0,
"folder_type": "temp",
"subfolder": " Exact Ω ",
"filename": "Keep Case.png",
"path_style": "posix",
"root_key": None,
"file_identity": "unverified",
"newly_generated": "unknown",
"context_membership": "unknown",
"transient": True,
}
],
},
)
self.assertEqual(list(conn.iterdump()), before)
for table in (
"prompts",
"prompt_pairs",
"generated_images",
"execution_run_contexts",
):
self.assertEqual(
conn.execute(f"SELECT count(*) FROM {table}").fetchone()[0], 0
)
def test_read_uses_consistent_snapshot_and_preserves_callers_transaction(self):
db = self.database()
start = self.capture(db)
peer = self.database()
conn = db.model.get_connection()
changed = False
def trace(sql):
nonlocal changed
if "FROM execution_run_outputs" in sql and not changed:
changed = True
with peer.model.get_connection() as writer:
writer.execute(
"UPDATE execution_run_outputs SET filename='later.png'"
)
conn.set_trace_callback(trace)
self.addCleanup(conn.set_trace_callback, None)
self.assertEqual(
db.get_execution_outputs(start.run_id)["references"][0]["filename"],
"first.png",
)
conn.set_trace_callback(None)
self.assertTrue(changed)
self.assertFalse(conn.in_transaction)
self.assertEqual(
db.get_execution_outputs(start.run_id)["references"][0]["filename"],
"later.png",
)
conn.execute("BEGIN")
db.get_execution_outputs(start.run_id)
self.assertTrue(conn.in_transaction)
conn.rollback()
def test_empty_unavailable_and_partial_headers_remain_distinct(self):
db = self.database()
for item, (state, reason) in enumerate(
(
("complete", "none_reported"),
("unavailable", "unavailable"),
("partial", "unsupported"),
)
):
start = self.capture(db, OutputBundle((), state, reason), item)
capture = db.get_execution_outputs(start.run_id)
self.assertEqual(
(capture["state"], capture["reason"], capture["references"]),
(state, reason, []),
)
def test_invalid_scalars_rejected_before_transaction(self):
db = self.database()
start = self.start(db)
db.write_execution_start(start)
trace = []
conn = db.model.get_connection()
conn.set_trace_callback(trace.append)
self.addCleanup(conn.set_trace_callback, None)
for invalid in (
replace(
bundle(),
references=(replace(bundle().references[0], filename="../unsafe.png"),),
),
replace(bundle(), source="filesystem_scan"),
):
with self.assertRaises(ValueError):
db.write_execution_outputs_for_finish(self.finish(start), invalid)
with self.assertRaises(ValueError):
db.write_execution_outputs_for_finish(start, bundle())
self.assertEqual(trace, [])