"""Attempt fencing and bounded runtime evidence, without a running host.""" from contextvars import copy_context from dataclasses import replace import gc import sys import threading import types import unittest import weakref from unittest.mock import patch from test_execution_context_graph import transformed_input_graph from utils.execution_context_budget import ContextBudget from utils.execution_context_graph import SourceRef from utils.execution_contexts import PromptSnapshot, TextEvidence from utils.execution_history_events import RunStart from utils.execution_history_writer import HistoryWriterStatus try: from utils import execution_context_capture as capture except ImportError: capture = None SESSION = "11111111-1111-4111-8111-111111111111" EPOCH = "22222222-2222-4222-8222-222222222222" JOB = "33333333-3333-4333-8333-333333333333" def graph(): return { "p": { "class_type": "PromptManagerPair", "inputs": {"text": "raw", "negative_text": "negative"}, }, "s": { "class_type": "KSampler", "inputs": {"positive": ["p", 0], "negative": ["p", 1]}, }, } def snapshot(raw=" raw ", effective="injected", prepend=" pre ", append=" post "): return PromptSnapshot( *( TextEvidence(v, "recorded", "runtime_return", None) for v in (raw, effective, prepend, append) ) ) class CaptureTests(unittest.TestCase): def setUp(self): self.assertIsNotNone(capture, "attempt capture API is missing") self.pool = ContextBudget() self.status = HistoryWriterStatus(EPOCH, True, 0, 0, True, 3) self.manager = capture.ContextCapture(10, lambda: self.status, self.pool) self.addCleanup(self.manager.revoke_all) self.start = RunStart(SESSION, SESSION, 1, JOB, EPOCH, 0) def begin(self, item=1, g=None): return self.manager.begin( replace(self.start, queue_item_id=item), 3, graph() if g is None else g, None, ) def seal(self, cap): result = self.manager.seal(cap) if result.lease: self.addCleanup(result.lease.release) return result def test_unexpected_projection_error_releases_payload_but_keeps_capability(self): with patch.object( capture, "project_contexts", side_effect=RuntimeError("private") ): cap = self.begin() self.assertTrue(self.manager.valid(cap)) self.assertEqual(self.pool.charged_bytes, 0) result = self.seal(cap) self.assertEqual( (result.state, result.reason), ("unavailable", "invalid_evidence") ) self.assertEqual(self.pool.charged_bytes, 0) def test_unexpected_seal_error_discards_popped_attempt_and_lease(self): cap = self.begin() with patch.object( capture, "bundle_payload_size", side_effect=RuntimeError("private") ): result = self.seal(cap) self.assertEqual( (result.state, result.reason), ("unavailable", "invalid_evidence") ) self.assertFalse(self.manager.valid(cap)) self.assertEqual(self.pool.charged_bytes, 0) def test_runtime_keeps_raw_affix_whitespace_and_injected_effective_text(self): cap = self.begin() self.manager.observe(cap, SourceRef("p", "PromptManagerPair", 0), 0, snapshot()) result = self.seal(cap) side = result.bundle.contexts[0].positive self.assertEqual( [ side.raw.value, side.effective.value, side.prepend.value, side.append.value, ], [" raw ", "injected", " pre ", " post "], ) self.assertIsNone(result.bundle.contexts[0].negative.effective.value) self.assertEqual(result.state, "partial") self.assertLess(self.pool.charged_bytes, 1048576) def test_reused_uuid_late_async_and_same_epoch_generation_are_fenced(self): a = self.begin() inherited = copy_context() b = self.begin(2) inherited.run( self.manager.observe, a, SourceRef("p", "PromptManagerPair", 0), 0, snapshot(), ) self.assertIsNone(self.seal(a).bundle) self.status = replace(self.status, capture_generation=4) self.manager.observe(b, SourceRef("p", "PromptManagerPair", 0), 0, snapshot()) self.assertIsNone(self.seal(b).bundle) self.assertEqual(self.pool.charged_bytes, 0) def test_unbound_thread_and_forged_metadata_cannot_record(self): cap = self.begin() host = types.ModuleType("comfy_execution.utils") host.get_executing_context = lambda: types.SimpleNamespace( prompt_id=JOB, node_id="p", list_index=0 ) with patch.dict(sys.modules, {"comfy_execution.utils": host}): thread = threading.Thread( target=capture.record_runtime_context, args=("PromptManagerPair", "bad", "", "", "bad"), kwargs={"role": "positive"}, ) thread.start() thread.join(2) self.manager.observe( replace(cap, nonce="forged"), SourceRef("p", "PromptManagerPair", 0), 0, snapshot(), ) host.get_executing_context = lambda: types.SimpleNamespace( prompt_id="forged", node_id="p", list_index=0 ) capture.record_runtime_context( "PromptManagerPair", "bad", "", "", "bad", role="positive" ) self.assertIsNone(self.seal(cap).bundle.contexts[0].positive.effective.value) def test_new_dequeue_revokes_unfinished_attempt(self): a = self.begin() b = self.begin(2) self.assertEqual(capture.active_capture(), (self.manager, b)) self.manager.revoke(a) self.assertEqual(capture.active_capture(), (self.manager, b)) self.manager.revoke_all() self.assertIsNone(capture.active_capture()) self.assertEqual(self.pool.charged_bytes, 0) def test_equal_distinct_invocations_are_ambiguous_and_replay_is_idempotent(self): for count, state in ((1, "partial"), (64, "ambiguous"), (65, "ambiguous")): cap = self.begin(count) for i in range(count): self.manager.observe( cap, SourceRef("p", "PromptManagerPair", 0), i, snapshot() ) self.manager.observe( cap, SourceRef("p", "PromptManagerPair", 0), i, snapshot() ) result = self.seal(cap) self.assertEqual(result.state, state) if count > 1: self.assertIsNone(result.bundle.contexts[0].positive.effective.value) def test_projection_never_retains_original_containers_and_budget_denial(self): class Envelope(dict): pass g = Envelope(graph()) ref = weakref.ref(g) cap = self.begin(g=g) del g gc.collect() self.assertIsNone(ref()) self.manager.revoke(cap) small = capture.ContextCapture(10, lambda: self.status, ContextBudget(1)) cap = small.begin(self.start, 3, graph(), None) self.assertEqual(small.seal(cap).reason, "budget_exhausted") def test_crossed_roles_and_multiple_anchors_use_exact_slots(self): g = graph() g["t"] = { "class_type": "KSampler", "inputs": {"positive": ["p", 1], "negative": ["p", 0]}, } cap = self.begin(g=g) self.manager.observe( cap, SourceRef("p", "PromptManagerPair", 0), 0, snapshot(effective="POS") ) self.manager.observe( cap, SourceRef("p", "PromptManagerPair", 1), 0, snapshot(effective="NEG") ) result = self.seal(cap) self.assertEqual(result.state, "complete") self.assertEqual( [ (r.positive.effective.value, r.negative.effective.value) for r in result.bundle.contexts ], [("POS", "NEG"), ("NEG", "POS")], ) def test_transformed_sources_cannot_replace_composed_text(self): g = graph() g["text"] = { "class_type": "StringConcatenate", "inputs": {"string_a": ["p", 2], "string_b": " suffix"}, } g["clip"] = {"class_type": "CLIPTextEncode", "inputs": {"text": ["text", 0]}} g["s"]["inputs"]["positive"] = ["clip", 0] cap = self.begin(g=g) self.manager.observe( cap, SourceRef("p", "PromptManagerPair", 2), 0, snapshot(effective="TRIGGER"), ) result = self.seal(cap) self.assertIsNone(result.bundle.contexts[0].positive.effective.value) self.assertNotEqual(result.state, "complete") def test_upstream_transform_preserves_final_pm_runtime_snapshot(self): for field in ("text", "prepend_text", "append_text"): for unresolved in (False, True): with self.subTest(field=field, unresolved=unresolved): cap = self.begin(g=transformed_input_graph(field, unresolved)) self.host_context() capture.record_runtime_context( "PromptManagerPair", " actual raw ", " actual pre ", " actual post ", " actual injected effective ", role="positive", ) side = self.seal(cap).bundle.contexts[0].positive for name, expected in ( ("raw", " actual raw "), ("prepend", " actual pre "), ("append", " actual post "), ("effective", " actual injected effective "), ): evidence = getattr(side, name) self.assertEqual(evidence.value, expected) self.assertEqual(evidence.provenance, "recorded") self.assertEqual(evidence.source, "runtime_return") def test_unsupported_graph_has_no_supported_context(self): result = self.seal(self.begin(g={})) self.assertEqual(result.reason, "no_supported_context") def host_context(self, kind="PromptManagerPair", index=0, job=JOB, node_id="p"): host = types.ModuleType("comfy_execution.utils") host.get_executing_context = lambda: types.SimpleNamespace( prompt_id=job, node_id=node_id, list_index=index ) stack = self.enter_stack() stack.enter_context(patch.dict(sys.modules, {"comfy_execution.utils": host})) return host def enter_stack(self): from contextlib import ExitStack stack = ExitStack() self.addCleanup(stack.close) return stack def test_nodes_record_only_successful_returns_and_save_failure_preserves_text(self): import logging from unittest.mock import Mock from prompt_manager import ( PromptManager, PromptManagerFocused, PromptManagerSplit, ) from prompt_manager_text import PromptManagerText from prompt_manager_pair import PromptManagerPair for cls in ( PromptManager, PromptManagerFocused, PromptManagerSplit, PromptManagerText, PromptManagerPair, ): kind = cls.__name__ g = graph() g["p"]["class_type"] = kind if kind != "PromptManagerPair": g["s"]["inputs"]["negative"] = ["p", 0] if kind == "PromptManagerText": g["c"] = {"class_type": "CLIPTextEncode", "inputs": {"text": ["p", 0]}} g["s"]["inputs"] = {"positive": ["c", 0], "negative": ["c", 0]} cap = self.begin(g=g) self.host_context(kind) node = cls.__new__(cls) node.db = Mock() node.db.get_prompt_by_hash.return_value = None node.db.save_prompt.side_effect = RuntimeError("offline") node.db.save_prompt_pair.side_effect = RuntimeError("offline") node.logger = logging.getLogger("test.capture.nodes") node.prompt_tracker = Mock() node.comfyui_integration = Mock() node._inject_lora_trigger_words = lambda text: text + " TRIGGER" clip = Mock() clip.tokenize.side_effect = lambda text: text clip.encode_from_tokens_scheduled.side_effect = lambda text: ( "encoded", text, ) kwargs = dict(text=" raw ", prepend_text=" pre ", append_text=" post ") if kind != "PromptManagerText": kwargs["clip"] = clip if kind == "PromptManagerPair": kwargs["negative_text"] = " neg " output = ( node.process_text(**kwargs) if kind == "PromptManagerText" else node.encode_prompt(**kwargs) ) result = self.seal(cap) side = result.bundle.contexts[0].positive self.assertEqual(side.effective.provenance, "recorded", kind) self.assertEqual(side.raw.value, " raw ") self.assertEqual(side.prepend.value, " pre ") self.assertEqual(side.append.value, " post ") self.assertEqual( side.effective.value, output[-2] if kind == "PromptManagerPair" else output[-1], ) if kind == "PromptManagerPair": self.assertEqual( result.bundle.contexts[0].negative.effective.value, " neg " ) if kind != "PromptManagerText": cap = self.begin(g=g) clip.encode_from_tokens_scheduled.side_effect = RuntimeError( "encode failed" ) with self.assertRaises(RuntimeError): node.encode_prompt(**kwargs) self.assertIsNone( self.seal(cap).bundle.contexts[0].positive.effective.value ) def test_oversized_projection_and_observation_release_all_reservations(self): g = graph() g["p"]["inputs"].update( text="x" * 65536, prepend_text=" " * 65536, append_text=" " * 65536 ) cap = self.begin(g=g) self.manager.revoke(cap) self.assertEqual(self.pool.charged_bytes, 0) cap = self.begin() self.manager.observe( cap, SourceRef("p", "PromptManagerPair", 0), 0, snapshot(effective="x" * 65537), ) self.assertIsNone(self.seal(cap).bundle.contexts[0].positive.effective.value) def test_recorder_rejects_wrong_class_slot_and_noninteger_invocation(self): for node_type, role, index, node_id in ( ("PromptManager", None, 0, "p"), ("PromptManagerPair", "invalid", 0, "p"), ("PromptManagerPair", "positive", True, "p"), ("PromptManagerPair", "positive", -1, "p"), ("PromptManagerPair", "positive", 0, "other"), ): cap = self.begin() self.host_context(index=index, node_id=node_id) capture.record_runtime_context(node_type, "bad", "", "", "bad", role=role) self.assertIsNone( self.seal(cap).bundle.contexts[0].positive.effective.value ) def test_async_inherited_capability_after_await_cannot_record_new_attempt(self): import asyncio async def scenario(): old = self.begin() entered, release = asyncio.Event(), asyncio.Event() async def late(): entered.set() await release.wait() capture.record_runtime_context( "PromptManagerPair", "late", "", "", "late", role="positive" ) self.host_context() task = asyncio.create_task(late()) await entered.wait() new = self.begin(2) release.set() await task self.assertIsNone(self.seal(old).bundle) self.assertIsNone( self.seal(new).bundle.contexts[0].positive.effective.value ) asyncio.run(scenario()) def test_writer_restore_or_close_revokes_callback_paused_during_encoding(self): import asyncio from utils.execution_history_writer import HistoryWriter for action in ("invalidate", "close"): with self.subTest(action=action): pool = ContextBudget() current = [HistoryWriterStatus(EPOCH, True, 0, 0, True, 3)] manager = capture.ContextCapture(10, lambda: current[0], pool) self.addCleanup(manager.revoke_all) writer = HistoryWriter( lambda: self.fail("This callback test must not open storage"), context_budget=pool, on_invalidation=manager.revoke_all, ) self.addCleanup(writer.close) self.host_context() async def scenario(): old = manager.begin(self.start, 3, graph(), None) manager.observe( old, SourceRef("p", "PromptManagerPair", 0), 0, snapshot() ) entered, release = asyncio.Event(), asyncio.Event() async def encoding_return(): entered.set() await release.wait() capture.record_runtime_context( "PromptManagerPair", "late", "", "", "late", role="positive" ) task = asyncio.create_task(encoding_return()) await entered.wait() self.assertGreater(pool.charged_bytes, 0) if action == "invalidate": writer.invalidate() else: self.assertTrue(writer.close()) self.assertEqual(pool.charged_bytes, 0) self.assertIsNone(capture.active_capture()) new_epoch = "44444444-4444-4444-8444-444444444444" current[0] = HistoryWriterStatus(new_epoch, True, 0, 0, True, 4) new = manager.begin( replace( self.start, run_id=JOB, queue_item_id=2, storage_epoch=new_epoch, ), 4, graph(), None, ) release.set() await task self.assertIsNone(manager.seal(old).bundle) sealed = manager.seal(new) try: self.assertIsNone( sealed.bundle.contexts[0].positive.effective.value ) finally: if sealed.lease: sealed.lease.release() self.assertEqual(pool.charged_bytes, 0) asyncio.run(scenario()) def test_oversized_runtime_field_is_unknown_without_erasing_other_fields(self): cap = self.begin() self.host_context() capture.record_runtime_context( "PromptManagerPair", " actual ", " pre ", " post ", "x" * 65537, role="positive", ) result = self.seal(cap) side = result.bundle.contexts[0].positive self.assertEqual(side.raw.value, " actual ") self.assertEqual(side.prepend.value, " pre ") self.assertIsNone(side.effective.value) self.assertEqual(side.effective.reason, "limit_reached") self.assertEqual(result.reason, "field_limit") def test_mixed_cache_provenance_cannot_promote_inferred_pm_effective(self): cap = self.begin() field = TextEvidence("literal", "inferred", "graph_literal", None) cached = TextEvidence("cached TRIGGER", "recorded", "host_cache_output", None) self.manager.observe( cap, SourceRef("p", "PromptManagerPair", 0), 0, PromptSnapshot(field, cached, field, field), ) side = self.seal(cap).bundle.contexts[0].positive self.assertEqual(side.raw.provenance, "inferred") self.assertEqual(side.effective.value, "cached TRIGGER") cap = self.begin(2) self.manager.observe( cap, SourceRef("p", "PromptManagerPair", 0), 0, PromptSnapshot(field, field, field, field), ) self.assertIsNone(self.seal(cap).bundle.contexts[0].positive.effective.value) def test_restore_and_close_release_weakly_bound_payload_immediately(self): cap = self.begin() binding = copy_context() self.manager.revoke_all() self.assertEqual(self.pool.charged_bytes, 0) self.assertIsNone(binding.run(capture.active_capture)) self.manager.observe(cap, SourceRef("p", "PromptManagerPair", 0), 0, snapshot()) self.assertIsNone(self.seal(cap).bundle) owner = capture.ContextCapture(11, lambda: self.status, self.pool) owner.begin(self.start, 3, graph(), None) binding = copy_context() owner.revoke_all() ref = weakref.ref(owner) del owner gc.collect() self.assertIsNone(ref()) self.assertIsNone(binding.run(capture.active_capture)) def test_conflicting_replay_stays_ambiguous_after_limited_other_source(self): cap = self.begin() source = SourceRef("p", "PromptManagerPair", 0) self.manager.observe(cap, source, 0, snapshot(effective="first")) self.manager.observe(cap, source, 0, snapshot(effective="different")) self.manager.observe( cap, SourceRef("p", "PromptManagerPair", 1), 0, snapshot(effective="x" * 65537), ) sealed = self.seal(cap) self.assertEqual(sealed.state, "ambiguous") self.assertIsNone(sealed.bundle.contexts[0].positive.effective.value) def test_capture_owner_bounds_unavailable_attempts_to_128(self): owner = capture.ContextCapture(12, lambda: self.status, ContextBudget(0)) self.addCleanup(owner.revoke_all) capabilities = [ copy_context().run( owner.begin, replace(self.start, queue_item_id=i), 3, {}, None ) for i in range(129) ] self.assertEqual(sum(owner.valid(cap) for cap in capabilities), 128) def test_revocation_discards_owned_samples_before_releasing_budget(self): cap = self.begin() value = snapshot() ref = weakref.ref(value.raw) self.manager.observe(cap, SourceRef("p", "PromptManagerPair", 0), 0, value) del value lease = self.manager._attempts[cap.nonce].lease original = lease.release def release(): self.assertIsNone(ref(), "sample was uncharged while still owned") original() with patch.object(lease, "release", release): self.manager.revoke(cap)