From da4d6cdc84190ec7a14cc0ac019c9f6348cdbf04 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Sat, 17 Oct 2026 18:28:23 -0700 Subject: [PATCH] feat(history): bind exact runtime context to execution attempts --- prompt_manager.py | 11 +- prompt_manager_pair.py | 19 +- prompt_manager_text.py | 11 +- tests/test_execution_context_capture.py | 475 ++++++++++++++++++++++++ tests/test_execution_context_graph.py | 14 + tests/test_execution_history_hooks.py | 179 ++++++++- tests/test_execution_history_writer.py | 23 ++ tests/test_pair_execution.py | 20 + tests/test_prompt_pair_node.py | 42 +++ utils/execution_context_capture.py | 401 ++++++++++++++++++++ utils/execution_context_graph.py | 81 ++-- utils/execution_history.py | 106 +++++- utils/execution_history_writer.py | 14 + 13 files changed, 1348 insertions(+), 48 deletions(-) create mode 100644 tests/test_execution_context_capture.py create mode 100644 utils/execution_context_capture.py diff --git a/prompt_manager.py b/prompt_manager.py index b070e02..3f653e9 100644 --- a/prompt_manager.py +++ b/prompt_manager.py @@ -6,6 +6,11 @@ with persistent prompt storage and search capabilities. import time from typing import Any, Dict, Optional, Tuple +if __package__: + from .utils.execution_context_capture import record_runtime_context +else: + from utils.execution_context_capture import record_runtime_context + try: from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict except ImportError: @@ -256,7 +261,11 @@ class PromptManager(PromptManagerBase, ComfyNodeABC): ) self.logger.info(f"CLIP encoding completed, text: {repr(encoding_text)[:80]}") - return (conditioning, encoding_text) + result = (conditioning, encoding_text) + record_runtime_context( + type(self).__name__, text, prepend_text, append_text, encoding_text + ) + return result @classmethod def IS_CHANGED( diff --git a/prompt_manager_pair.py b/prompt_manager_pair.py index 75ae89b..dcef20d 100644 --- a/prompt_manager_pair.py +++ b/prompt_manager_pair.py @@ -5,6 +5,11 @@ import json import time from typing import Any, Dict, Optional, Tuple +if __package__: + from .utils.execution_context_capture import record_runtime_context +else: + from utils.execution_context_capture import record_runtime_context + try: from .prompt_manager import IO, InputTypeDict, PromptManagerFocused from .utils.hashing import generate_prompt_hash @@ -204,7 +209,19 @@ class PromptManagerPair(PromptManagerFocused): ) except Exception as error: self.logger.warning(f"Could not register paired metadata: {error}") - return positive_conditioning, negative_conditioning, positive, negative_text + result = positive_conditioning, negative_conditioning, positive, negative_text + record_runtime_context( + type(self).__name__, + text, + prepend_text, + append_text, + positive, + role="positive", + ) + record_runtime_context( + type(self).__name__, negative_text, "", "", negative_text, role="negative" + ) + return result @classmethod def IS_CHANGED( diff --git a/prompt_manager_text.py b/prompt_manager_text.py index 7ba3123..92f8133 100644 --- a/prompt_manager_text.py +++ b/prompt_manager_text.py @@ -6,6 +6,11 @@ without CLIP encoding, while maintaining all database and search features. import time from typing import Any, Dict, Optional, Tuple +if __package__: + from .utils.execution_context_capture import record_runtime_context +else: + from utils.execution_context_capture import record_runtime_context + try: from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict except ImportError: @@ -227,7 +232,11 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC): ) self.logger.debug(f"Text processing completed: {final_text[:100]}...") - return (final_text,) + result = (final_text,) + record_runtime_context( + type(self).__name__, text, prepend_text, append_text, final_text + ) + return result @classmethod def IS_CHANGED( diff --git a/tests/test_execution_context_capture.py b/tests/test_execution_context_capture.py new file mode 100644 index 0000000..d0962e4 --- /dev/null +++ b/tests/test_execution_context_capture.py @@ -0,0 +1,475 @@ +"""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 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_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_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_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) diff --git a/tests/test_execution_context_graph.py b/tests/test_execution_context_graph.py index e9dbdf4..56f3267 100644 --- a/tests/test_execution_context_graph.py +++ b/tests/test_execution_context_graph.py @@ -230,3 +230,17 @@ class ProjectionTests(unittest.TestCase): {"s": node("KSampler", positive=HostValue())}, ): self.m.project_contexts(g, None) + + def test_composed_affixes_and_retained_projection_are_bounded(self): + g = graph() + g["p"]["inputs"].update( + text="x" * 65536, prepend_text="y" * 65536, append_text="z" * 65536 + ) + result = self.m.project_contexts(g, None) + self.assertIsNone(result.recipes[0].positive.inferred) + g = graph() + g["p"]["inputs"]["text"] = "x" * 60000 + for i in range(15): + g[str(i)] = node("KSampler", positive=["p", 0], negative=["p", 0]) + result = self.m.project_contexts(g, None) + self.assertEqual(result.recipes, ()) diff --git a/tests/test_execution_history_hooks.py b/tests/test_execution_history_hooks.py index f54b56f..6d518c3 100644 --- a/tests/test_execution_history_hooks.py +++ b/tests/test_execution_history_hooks.py @@ -356,7 +356,7 @@ class HistoryHooksTests(unittest.TestCase): clock = ManualClock() with patch( "utils.execution_history.HistoryWriter", - lambda factory: HistoryWriter(factory, clock=clock), + lambda factory, **kwargs: HistoryWriter(factory, clock=clock, **kwargs), ): registration = self.register() for index in range(128): @@ -448,3 +448,180 @@ class HistoryHooksTests(unittest.TestCase): self.drain(registration) self.assertEqual(seen, [JOB]) self.assertEqual(len(self.rows()), 1) + + def test_rejected_new_start_revokes_previous_worker_capability(self): + from utils.execution_context_capture import active_capture + from utils.execution_history_events import make_start + + entered, release = threading.Event(), threading.Event() + self.releases.append(release) + original = self.db.write_execution_start + + def blocked(command): + entered.set() + release.wait(5) + return original(command) + + with patch.object(self.db, "write_execution_start", blocked): + registration = self.register() + self.dequeue(1) + self.assertTrue(entered.wait(1)) + previous = active_capture() + self.assertIsNotNone(previous) + epoch = registration.writer.status().epoch + for i in range(1023): + command = make_start(registration._session_id, item(1000 + i), epoch, 0) + self.assertTrue(registration.writer.submit(command)) + self.assertEqual(registration.writer.status().pending, 1024) + self.dequeue(2) + self.assertIsNone(active_capture()) + self.assertFalse(previous[0].valid(previous[1])) + release.set() + + def test_late_finish_does_not_clear_new_attempt_binding(self): + from utils.execution_context_capture import active_capture + + registration = self.register() + self.dequeue(1) + self.dequeue(2) + current = active_capture() + self.assertIsNotNone(current) + self.queue.task_done(1, {}, status()) + self.assertEqual(active_capture(), current) + # An identical callback replay keeps the already accepted binding. + registration._dequeue(item(2)) + self.assertEqual(active_capture(), current) + self.queue.task_done(2, {}, status()) + self.assertIsNone(active_capture()) + + def test_invalidation_releases_capture_immediately_outside_writer_locks(self): + from utils.execution_context_capture import active_capture + + registration = self.register() + self.dequeue(1) + current = active_capture() + pool = registration.writer.context_budget + self.assertGreater(pool.charged_bytes, 0) + registration.writer.invalidate() + self.assertEqual(pool.charged_bytes, 0) + self.assertFalse(current[0].valid(current[1])) + + def test_replayed_older_accepted_dequeue_cannot_revoke_current_binding(self): + from utils.execution_context_capture import active_capture + + registration = self.register() + old = self.dequeue(1) + self.dequeue(2) + current = active_capture() + registration._dequeue(old) + self.assertEqual(active_capture(), current) + + def test_broken_finish_releases_active_capture(self): + from utils.execution_context_capture import active_capture + + registration = self.register() + self.dequeue(1) + + class Broken: + @property + def messages(self): + raise RuntimeError("broken status") + + self.queue.task_done(1, {}, Broken()) + self.assertEqual(registration.writer.context_budget.charged_bytes, 0) + self.assertIsNone(active_capture()) + + def test_finished_accepted_replay_preserves_new_same_uuid_attempt(self): + from utils.execution_context_capture import active_capture + + registration = self.register() + old = self.dequeue(1) + self.queue.task_done(1, {}, status()) + self.dequeue(2) + current = active_capture() + registration._dequeue(old) + self.assertEqual(active_capture(), current) + self.queue.task_done(2, {}, status()) + self.drain(registration) + self.assertEqual(len(self.rows()), 2) + + def test_observed_upstream_survives_downstream_failure_in_database(self): + from utils.execution_context_capture import record_runtime_context + + g = { + "p": { + "class_type": "PromptManagerPair", + "inputs": {"text": "raw", "negative_text": ""}, + }, + "s": { + "class_type": "KSampler", + "inputs": {"positive": ["p", 0], "negative": ["p", 1]}, + }, + } + registration = self.register() + self.queue.results.append(((0, JOB, g, {}, ["s"]), 1)) + self.queue.get() + host = types.ModuleType("comfy_execution.utils") + host.get_executing_context = lambda: SimpleNamespace( + prompt_id=JOB, node_id="p", list_index=0 + ) + with patch.dict(sys.modules, {"comfy_execution.utils": host}): + record_runtime_context( + "PromptManagerPair", + " raw ", + " pre ", + "", + "pre raw TRIGGER", + role="positive", + ) + record_runtime_context("PromptManagerPair", "", "", "", "", role="negative") + self.queue.task_done(1, {}, status("execution_error")) + self.drain(registration) + row = self.rows()[0] + self.assertEqual(row["outcome"], "failed") + result = self.db.get_execution_contexts(row["run_id"]) + self.assertEqual(result["state"], "complete") + self.assertEqual( + result["contexts"][0]["positive"], + { + "raw": " raw ", + "effective": "pre raw TRIGGER", + "prepend": " pre ", + "append": "", + }, + ) + self.assertEqual(result["contexts"][0]["negative"]["effective"], "") + + def test_close_releases_capture_while_writer_is_blocked(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(1) + self.assertTrue(entered.wait(1)) + self.assertGreater(registration.writer.context_budget.charged_bytes, 0) + self.assertFalse(registration.writer.close(0)) + self.assertEqual(registration.writer.context_budget.charged_bytes, 0) + release.set() + + def test_reregistration_cannot_rebind_already_observed_queue_item(self): + from utils.execution_context_capture import active_capture + + first = self.register() + self.dequeue(10) + self.queue.task_done(10, {}, status()) + self.drain(first) + self.assertTrue(first.dispose()) + second = self.register() + self.dequeue(10) + self.assertIsNone(active_capture()) + self.queue.task_done(10, {}, status()) + self.drain(second) + self.assertEqual(len(self.rows()), 1) diff --git a/tests/test_execution_history_writer.py b/tests/test_execution_history_writer.py index 289818c..5c1be41 100644 --- a/tests/test_execution_history_writer.py +++ b/tests/test_execution_history_writer.py @@ -905,3 +905,26 @@ class HistoryWriterTests(unittest.TestCase): if __name__ == "__main__": unittest.main() + + def test_owner_invalidation_callback_runs_outside_both_locks_and_is_contained(self): + observations = [] + writer = None + + def callback(): + acquired = [ + lock.acquire(blocking=False) + for lock in (writer._lock, writer._context_commit_gate) + ] + observations.append(tuple(acquired)) + for lock, owned in zip( + (writer._lock, writer._context_commit_gate), acquired + ): + if owned: + lock.release() + raise RuntimeError("callback failure cannot escape") + + writer = HistoryWriter(lambda: self.db, on_invalidation=callback) + self.writers.append(writer) + writer.invalidate() + writer.close(0) + self.assertEqual(observations, [(True, True), (True, True)]) diff --git a/tests/test_pair_execution.py b/tests/test_pair_execution.py index 5f6bea2..a81399e 100644 --- a/tests/test_pair_execution.py +++ b/tests/test_pair_execution.py @@ -326,3 +326,23 @@ class TestPairExecution(unittest.TestCase): record_pair_execution(extra, "1", **kwargs) with self.assertRaises(ExecutionEvidenceError): read_pair_execution(extra, "1") + + def test_png_evidence_cannot_populate_unobserved_attempt(self): + from utils.execution_context_capture import ContextCapture + from utils.execution_context_budget import ContextBudget + from utils.execution_history_events import RunStart + from utils.execution_history_writer import HistoryWriterStatus + + identity = "11111111-1111-4111-8111-111111111111" + self.execute() + state = HistoryWriterStatus(identity, True, 0, 0, True, 1) + manager = ContextCapture(1, lambda: state, ContextBudget()) + self.addCleanup(manager.revoke_all) + cap = manager.begin( + RunStart(identity, identity, 1, identity, identity, 0), 1, graph(), None + ) + sealed = manager.seal(cap) + if sealed.lease: + self.addCleanup(sealed.lease.release) + self.assertIsNone(sealed.bundle.contexts[0].positive.effective.value) + self.assertIsNone(sealed.bundle.contexts[0].negative.effective.value) diff --git a/tests/test_prompt_pair_node.py b/tests/test_prompt_pair_node.py index 693a6c1..671d4cd 100644 --- a/tests/test_prompt_pair_node.py +++ b/tests/test_prompt_pair_node.py @@ -158,3 +158,45 @@ class TestPromptPairNode(unittest.TestCase): ) result = self.node.encode_prompt(self.clip, "cat", negative_text="blur") self.assertEqual(result[2:], ("cat LORA", "blur")) + + def test_failed_negative_encoding_cannot_publish_positive_runtime_proof(self): + from contextlib import ExitStack + import sys + import types + from utils.execution_context_capture import ContextCapture + from utils.execution_context_budget import ContextBudget + from utils.execution_history_events import RunStart + from utils.execution_history_writer import HistoryWriterStatus + from unittest.mock import patch + + identity = "11111111-1111-4111-8111-111111111111" + state = HistoryWriterStatus(identity, True, 0, 0, True, 1) + budget = ContextBudget() + manager = ContextCapture(1, lambda: state, budget) + self.addCleanup(manager.revoke_all) + cap = manager.begin( + RunStart(identity, identity, 1, identity, identity, 0), + 1, + pair_graph(), + None, + ) + host = types.ModuleType("comfy_execution.utils") + host.get_executing_context = lambda: types.SimpleNamespace( + prompt_id=identity, node_id="1", list_index=0 + ) + self.clip.encode_from_tokens_scheduled.side_effect = [ + ("conditioning", "cat LORA"), + RuntimeError("negative failed"), + ] + with ExitStack() as stack: + stack.enter_context( + patch.dict(sys.modules, {"comfy_execution.utils": host}) + ) + with self.assertRaisesRegex(RuntimeError, "negative failed"): + self.node.encode_prompt(self.clip, "cat", negative_text="blur") + sealed = manager.seal(cap) + if sealed.lease: + self.addCleanup(sealed.lease.release) + self.assertIsNone(sealed.bundle.contexts[0].positive.effective.value) + self.assertIsNone(sealed.bundle.contexts[0].negative.effective.value) + self.assertEqual(self.db.list_prompt_pairs()["total"], 0) diff --git a/utils/execution_context_capture.py b/utils/execution_context_capture.py new file mode 100644 index 0000000..a9d1ad0 --- /dev/null +++ b/utils/execution_context_capture.py @@ -0,0 +1,401 @@ +"""Revocable, bounded runtime evidence owned by an actual dequeue attempt.""" + +from __future__ import annotations + +from contextvars import ContextVar +from dataclasses import dataclass, replace +import threading +from typing import Callable, TYPE_CHECKING +from uuid import uuid4 +import weakref + +from .execution_context_budget import ContextBudget, PayloadLease, SealedContext +from .execution_context_graph import ( + GraphProjection, + GraphRole, + SourceRef, + project_contexts, +) +from .execution_contexts import ( + ContextBundle, + PromptSnapshot, + RecipeContext, + TextEvidence, + bundle_payload_size, + text_size, + MAX_ATTEMPT_BYTES, +) +from .execution_history_events import RunStart, encode_command + +if TYPE_CHECKING: + from .execution_history_writer import HistoryWriterStatus + + +@dataclass(frozen=True) +class AttemptCapability: + nonce: str + run_id: str + session_id: str + queue_item_id: int + job_id: str + storage_epoch: str + capture_generation: int + sources: tuple[SourceRef, ...] + + +@dataclass +class _Attempt: + capability: AttemptCapability + projection: GraphProjection + lease: PayloadLease | None + samples: dict + retained_bytes: int + reason: str = "" + + +_binding: ContextVar[tuple | None] = ContextVar("prompt_manager_attempt", default=None) + + +def active_capture() -> tuple[ContextCapture, AttemptCapability] | None: + binding = _binding.get() + if binding is not None: + manager = binding[0]() + if manager is not None and manager.valid(binding[1]): + return manager, binding[1] + return None + + +def unbind_current() -> None: + binding = _binding.get() + _binding.set(None) + if binding is not None: + manager = binding[0]() + if manager is not None: + manager.revoke(binding[1]) + + +def _unknown(reason: str = "missing") -> PromptSnapshot: + field = TextEvidence(None, "unknown", "unknown", reason) + return PromptSnapshot(field, field, field, field) + + +def _snapshot_bytes(snapshot: PromptSnapshot) -> int: + # Reuse the scalar validator, including provenance and per-field bounds. + bundle = ContextBundle( + (RecipeContext("x", snapshot, None, "inferred"),), "partial", "partial" + ) + return bundle_payload_size(bundle) + + +def _bounded_snapshot(snapshot: PromptSnapshot) -> tuple[PromptSnapshot, bool]: + if type(snapshot) is not PromptSnapshot: + raise ValueError("invalid_snapshot") + fields, limited = [], False + for field in (snapshot.raw, snapshot.effective, snapshot.prepend, snapshot.append): + if type(field) is not TextEvidence: + raise ValueError("invalid_field") + if field.value is not None: + try: + text_size(field.value) + except ValueError: + field = _unknown("limit_reached").raw + limited = True + fields.append(field) + if fields[1].value is not None and fields[1].provenance != "recorded": + fields[1] = _unknown("missing").effective + bounded = PromptSnapshot(*fields) + _snapshot_bytes(bounded) + return bounded, limited + + +class ContextCapture: + def __init__( + self, + queue_identity: int, + status_reader: Callable[[], HistoryWriterStatus], + budget: ContextBudget, + ) -> None: + self.queue_identity = queue_identity + self._status_reader = status_reader + self._budget = budget + self._lock = threading.RLock() + self._attempts: dict[str, _Attempt] = {} + + def _current(self, capability: AttemptCapability) -> _Attempt | None: + attempt = self._attempts.get(capability.nonce) + if attempt is None or attempt.capability is not capability: + return None + status = self._status_reader() + if ( + not status.accepting + or status.epoch != capability.storage_epoch + or status.capture_generation != capability.capture_generation + ): + self._drop(capability) + return None + return attempt + + def valid(self, capability: AttemptCapability) -> bool: + with self._lock: + return self._current(capability) is not None + + def begin( + self, start: RunStart, generation: int, graph: object, outputs: object + ) -> AttemptCapability: + unbind_current() + encode_command(start) + with self._lock: + cap = AttemptCapability( + str(uuid4()), + start.run_id, + start.session_id, + start.queue_item_id, + start.job_id, + start.storage_epoch, + generation, + (), + ) + if len(self._attempts) >= 128: + return cap + lease = self._budget.reserve(MAX_ATTEMPT_BYTES) + projection = GraphProjection((), (), "") + reason = "budget_exhausted" if lease is None else "" + retained = 0 + if lease is not None: + try: + projection = project_contexts(graph, outputs) + retained = sum( + _snapshot_bytes(role.inferred) if role.inferred else 0 + for recipe in projection.recipes + for role in (recipe.positive, recipe.negative) + ) + sum( + text_size(s.node_id, 512) + text_size(s.node_type, 128) + 32 + for s in projection.sources + ) + if retained > MAX_ATTEMPT_BYTES: + raise ValueError("projection_limit") + except ValueError: + projection = GraphProjection((), (), "") + reason = "graph_limit" + retained = 0 + cap = replace(cap, sources=projection.sources) + self._attempts[cap.nonce] = _Attempt( + cap, projection, lease, {}, retained, reason + ) + _binding.set((weakref.ref(self), cap)) + self._current(cap) + return cap + + def observe( + self, + capability: AttemptCapability, + source: SourceRef, + invocation_index: int, + snapshot: PromptSnapshot, + ) -> None: + try: + with self._lock: + attempt = self._current(capability) + if ( + attempt is None + or attempt.lease is None + or source not in capability.sources + or type(source) is not SourceRef + or type(invocation_index) is not int + or not 0 <= invocation_index <= 2**63 - 1 + ): + return + snapshot, limited = _bounded_snapshot(snapshot) + size = _snapshot_bytes(snapshot) + if ( + limited + or any( + field.reason == "limit_reached" + for field in ( + snapshot.raw, + snapshot.effective, + snapshot.prepend, + snapshot.append, + ) + ) + ) and attempt.reason not in ("multiple_samples", "limit_reached"): + attempt.reason = "field_limit" + key = source, invocation_index + if key in attempt.samples: + if attempt.samples[key] != snapshot: + attempt.reason = "multiple_samples" + return + if len(attempt.samples) >= 64: + attempt.reason = "multiple_samples" + return + if attempt.retained_bytes + size > MAX_ATTEMPT_BYTES: + attempt.reason = "limit_reached" + return + attempt.samples[key] = snapshot + attempt.retained_bytes += size + except (ValueError, AttributeError, TypeError): + return + + @staticmethod + def _discard_payload(attempt: _Attempt) -> None: + attempt.projection = GraphProjection((), (), "") + attempt.samples.clear() + attempt.retained_bytes = 0 + + def _drop(self, capability: AttemptCapability) -> _Attempt | None: + attempt = self._attempts.get(capability.nonce) + if attempt is not None and attempt.capability is capability: + self._attempts.pop(capability.nonce) + self._discard_payload(attempt) + if attempt.lease is not None: + attempt.lease.release() + return attempt + return None + + def revoke(self, capability: AttemptCapability) -> None: + with self._lock: + self._drop(capability) + + def revoke_all(self) -> None: + with self._lock: + for attempt in self._attempts.values(): + self._discard_payload(attempt) + if attempt.lease is not None: + attempt.lease.release() + self._attempts.clear() + + def seal(self, capability: AttemptCapability) -> SealedContext: + with self._lock: + attempt = self._current(capability) + if attempt is None: + return SealedContext(None, None, "unavailable", "unavailable") + self._attempts.pop(capability.nonce) + lease = attempt.lease + if lease is None: + return SealedContext(None, None, "unavailable", attempt.reason) + try: + ambiguous = attempt.reason == "multiple_samples" + + def role_snapshot(role: GraphRole) -> PromptSnapshot: + nonlocal ambiguous + if role.reason == "source-transformed": + return _unknown("unsupported") + if not role.sources: + return role.inferred or _unknown("unsupported") + values = [ + snapshot + for (source, _), snapshot in attempt.samples.items() + if source in role.sources + ] + if len(values) > 1 or attempt.reason == "multiple_samples": + ambiguous = True + return _unknown("ambiguous") + if len(role.sources) == 1 and len(values) == 1: + return values[0] + if role.inferred is not None: + s = role.inferred + return PromptSnapshot( + s.raw, _unknown().effective, s.prepend, s.append + ) + return _unknown() + + contexts = tuple( + RecipeContext( + r.anchor_id, + role_snapshot(r.positive), + role_snapshot(r.negative), + "inferred", + ) + for r in attempt.projection.recipes + ) + supported = any( + role.sources or role.inferred is not None + for r in attempt.projection.recipes + for role in (r.positive, r.negative) + ) + complete = bool(contexts) and all( + field.value is not None + for r in contexts + for s in (r.positive, r.negative) + for field in (s.raw, s.effective, s.prepend, s.append) + ) + state = ( + "ambiguous" if ambiguous else "complete" if complete else "partial" + ) + reason = ( + "multiple_samples" + if ambiguous + else "captured" if complete else "partial" + ) + if attempt.reason == "field_limit" and not ambiguous: + state, reason = "partial", "field_limit" + if not supported: + contexts, state, reason = (), "unavailable", "no_supported_context" + if attempt.reason in ("limit_reached", "graph_limit"): + contexts, state, reason = (), "unavailable", attempt.reason + bundle = ContextBundle(contexts, state, reason) + size = bundle_payload_size(bundle) + if size + attempt.retained_bytes > MAX_ATTEMPT_BYTES: + raise ValueError("assembly_limit") + self._discard_payload(attempt) + lease.shrink_to(size) + return SealedContext( + bundle, lease, state, reason, capability.capture_generation + ) + except ValueError: + self._discard_payload(attempt) + lease.release() + return SealedContext(None, None, "unavailable", "limit_reached") + + +def record_runtime_context( + node_type: str, + raw: str, + prepend: str, + append: str, + effective: str, + *, + role: str | None = None, +) -> None: + """Record only a successful node return in a trusted, bound host invocation.""" + try: + binding = active_capture() + if binding is None: + return + # Optional host import is independent of our package namespace. + from comfy_execution.utils import get_executing_context + + host = get_executing_context() + manager, cap = binding + if ( + host is None + or type(host.prompt_id) is not str + or host.prompt_id != cap.job_id + ): + return + if type(host.node_id) is not str or type(host.list_index) is not int: + return + if node_type == "PromptManagerPair": + slots = ( + (0, 2) if role == "positive" else (1, 3) if role == "negative" else () + ) + else: + slots = ( + (0,) + if node_type == "PromptManagerText" + else (0, 1) if role is None else () + ) + for slot in slots: + source = SourceRef(host.node_id, node_type, slot) + if source not in cap.sources: + continue + snapshot = PromptSnapshot( + *( + TextEvidence(value, "recorded", "runtime_return", None) + for value in (raw, effective, prepend, append) + ) + ) + manager.observe(cap, source, host.list_index, snapshot) + except Exception: + # Capture must never change the host result or expose input/exception text. + return diff --git a/utils/execution_context_graph.py b/utils/execution_context_graph.py index 4cc561f..82ff540 100644 --- a/utils/execution_context_graph.py +++ b/utils/execution_context_graph.py @@ -48,6 +48,10 @@ class _Unknown(ValueError): pass +class _Limit(_Unknown): + pass + + def _identifier(value: object) -> str: if type(value) is int: if abs(value) > 10**127: @@ -79,6 +83,7 @@ class _Projector: def __init__(self, graph: dict) -> None: self.graph = graph self.visits = 0 + self.retained = 0 self.transformed = False self.selected_seen = set() @@ -179,24 +184,29 @@ class _Projector: ) raw = child.effective.value prepend, append = before.effective.value, after.effective.value - effective = " ".join( - part for part in (prepend.strip(), raw, append.strip()) if part - ) + parts = [part for part in (prepend.strip(), raw, append.strip()) if part] + if sum(text_size(part) for part in parts) + max(0, len(parts) - 1) > 65536: + raise _Unknown("text_field_limit_reached") + effective = " ".join(parts) text_size(effective) - snapshot = PromptSnapshot( - child.effective, - TextEvidence(effective, "inferred", "graph_semantics", None), - before.effective, - after.effective, + snapshot = self.retain( + PromptSnapshot( + child.effective, + TextEvidence(effective, "inferred", "graph_semantics", None), + before.effective, + after.effective, + ) ) if name not in _PM: sources += before_sources + after_sources elif name in _PM: - snapshot = PromptSnapshot( - child.effective, - child.effective, - TextEvidence("", "inferred", "graph_literal", None), - TextEvidence("", "inferred", "graph_literal", None), + snapshot = self.retain( + PromptSnapshot( + child.effective, + child.effective, + TextEvidence("", "inferred", "graph_literal", None), + TextEvidence("", "inferred", "graph_literal", None), + ) ) else: snapshot = child @@ -254,8 +264,23 @@ class _Projector: return tuple(sources) return () - @staticmethod + def retain(self, snapshot: PromptSnapshot) -> PromptSnapshot: + size = sum( + text_size(field.value) + 64 + for field in ( + snapshot.raw, + snapshot.effective, + snapshot.prepend, + snapshot.append, + ) + ) + if self.retained + size > 1048576: + raise _Limit("projection_payload_limit_reached") + self.retained += size + return snapshot + def snapshot( + self, raw: str, effective: str, prepend: str, @@ -263,19 +288,21 @@ class _Projector: composed: bool = False, raw_composed: bool = False, ) -> PromptSnapshot: - return PromptSnapshot( - *( - TextEvidence( - value, - "inferred", - ( - "graph_semantics" - if (composed and index == 1 or raw_composed and index == 0) - else "graph_literal" - ), - None, + return self.retain( + PromptSnapshot( + *( + TextEvidence( + value, + "inferred", + ( + "graph_semantics" + if (composed and index == 1 or raw_composed and index == 0) + else "graph_literal" + ), + None, + ) + for index, value in enumerate((raw, effective, prepend, append)) ) - for index, value in enumerate((raw, effective, prepend, append)) ) ) @@ -288,6 +315,8 @@ class _Projector: return GraphRole( sources, snapshot, "source-transformed" if self.transformed else "" ) + except _Limit: + raise except ValueError as error: reason = str(error) if isinstance(error, _Unknown) else "invalid_text_value" return GraphRole( diff --git a/utils/execution_history.py b/utils/execution_history.py index c05c07b..6a4e952 100644 --- a/utils/execution_history.py +++ b/utils/execution_history.py @@ -1,5 +1,6 @@ """Own bounded queue-attempt capture without retaining host workflow payloads.""" +from collections import OrderedDict import atexit import threading import time @@ -12,6 +13,7 @@ except ImportError: from database import PromptDatabase from .execution_history_events import make_finish, make_start +from .execution_context_capture import ContextCapture, unbind_current from .execution_history_writer import HistoryWriter from .logging_config import get_logger from .queue_observers import subscribe_dequeue, subscribe_task_done @@ -26,13 +28,27 @@ class HistoryRegistration: """Own subscriptions, one asynchronous writer and bounded attempt identities.""" def __init__( - self, queue: object, session_id: str, db_factory: Callable[[], PromptDatabase] + self, + queue: object, + session_id: str, + db_factory: Callable[[], PromptDatabase], + *, + high_water: int = -1, ) -> None: - self.writer = HistoryWriter(db_factory) + self.writer = HistoryWriter( + db_factory, on_invalidation=self._invalidate_capture + ) + self.capture = ContextCapture( + id(queue), self.writer.status, self.writer.context_budget + ) + self._capabilities = {} + self._accepted = OrderedDict() + self._high_water = high_water self._session_id = session_id self._lock = threading.Lock() self._active: dict[int, str] = {} self._disposed = False + self._capture_generation = None self._unsubscribers: list[Callable[[], None]] = [] try: self._unsubscribers.append( @@ -47,40 +63,76 @@ class HistoryRegistration: self.dispose() raise + def _invalidate_capture(self) -> None: + self.capture.revoke_all() + def _ready_epoch(self) -> str | None: status = self.writer.status() epoch = status.epoch if status.accepting else None - if any(value != epoch for value in self._active.values()): + if status.capture_generation != self._capture_generation or any( + value != epoch for value in self._active.values() + ): self._active.clear() + self._capabilities.clear() + self.capture.revoke_all() + self._capture_generation = status.capture_generation return epoch def _dequeue(self, result: object) -> None: try: with self._lock: + epoch = self._ready_epoch() + command = ( + make_start( + self._session_id, result, epoch, time.time_ns() // 1_000_000 + ) + if epoch + else None + ) + accepted = ( + self._accepted.get(command.queue_item_id) + if command is not None + else None + ) + if command is not None and accepted == ( + command.job_id, + command.storage_epoch, + ): + return + unbind_current() if self._disposed: return - epoch = self._ready_epoch() - if epoch is None: + if command is None or command.queue_item_id in self._active: self.writer.record_gap() return - command = make_start( - self._session_id, result, epoch, time.time_ns() // 1_000_000 - ) - if command is None: - self.writer.record_gap() - return - item_id = command.queue_item_id - if item_id in self._active: - return if len(self._active) >= _MAX_ACTIVE: self.writer.record_gap() return + fresh = command.queue_item_id > self._high_water + self._high_water = max(self._high_water, command.queue_item_id) if self.writer.submit(command): - self._active[item_id] = epoch + self._accepted[command.queue_item_id] = ( + command.job_id, + command.storage_epoch, + ) + if len(self._accepted) > 128: + self._accepted.popitem(last=False) + self._active[command.queue_item_id] = epoch + if not fresh: + return + queue_item = result[0] + graph = queue_item[2] if len(queue_item) > 2 else None + outputs = queue_item[4] if len(queue_item) > 4 else None + cap = self.capture.begin( + command, self.writer.status().capture_generation, graph, outputs + ) + self._capabilities[command.queue_item_id] = cap except Exception: + unbind_current() self.writer.record_gap() def _task_done(self, args: tuple, kwargs: dict) -> None: + sealed = None try: with self._lock: if self._disposed: @@ -90,6 +142,11 @@ class HistoryRegistration: self.writer.record_gap() return item_id = args[0] if args else kwargs.get("item_id") + if type(item_id) is int: + cap = self._capabilities.pop(item_id, None) + self._active.pop(item_id, None) + if cap is not None: + sealed = self.capture.seal(cap) status = args[2] if len(args) > 2 else kwargs.get("status") command = make_finish( self._session_id, @@ -101,10 +158,16 @@ class HistoryRegistration: if command is None: self.writer.record_gap() return - self._active.pop(command.queue_item_id, None) - self.writer.submit(command) + if sealed is not None: + self.writer.submit_finalization(command, sealed) + sealed = None + else: + self.writer.submit(command) except Exception: self.writer.record_gap() + finally: + if sealed is not None and sealed.lease is not None: + sealed.lease.release() def dispose(self, timeout: float = 2.0) -> bool: """Stop capture, release subscriptions and drain the owned writer boundedly.""" @@ -112,6 +175,8 @@ class HistoryRegistration: if not self._disposed: self._disposed = True self._active.clear() + self._capabilities.clear() + self.capture.revoke_all() for unsubscribe in self._unsubscribers: unsubscribe() self._unsubscribers.clear() @@ -137,7 +202,12 @@ def register_history_hook( previous = state["registration"] if previous is not None and not previous._disposed: return previous - registration = HistoryRegistration(queue, state["session_id"], db_factory) + registration = HistoryRegistration( + queue, + state["session_id"], + db_factory, + high_water=previous._high_water if previous is not None else -1, + ) state["registration"] = registration return registration except Exception: diff --git a/utils/execution_history_writer.py b/utils/execution_history_writer.py index 3037068..2e4de1c 100644 --- a/utils/execution_history_writer.py +++ b/utils/execution_history_writer.py @@ -63,10 +63,12 @@ class HistoryWriter: *, clock: Callable[[], float] = time.monotonic, context_budget: ContextBudget | None = None, + on_invalidation: Callable[[], None] | None = None, ) -> None: self.context_budget = ( SHARED_CONTEXT_BUDGET if context_budget is None else context_budget ) + self._on_invalidation = on_invalidation self._db_factory = db_factory self._clock = clock self._context_commit_gate = threading.Lock() @@ -202,6 +204,13 @@ class HistoryWriter: ) yield current + def _notify_invalidation(self) -> None: + if self._on_invalidation is not None: + try: + self._on_invalidation() + except Exception: + pass + def invalidate(self) -> None: with self._context_commit_gate: with self._lock: @@ -210,6 +219,7 @@ class HistoryWriter: self._generation += 1 self._next_probe = self._clock() + _PROBE_INTERVAL self._discard_locked() + self._notify_invalidation() self._wake.set() def status(self) -> HistoryWriterStatus: @@ -233,6 +243,7 @@ class HistoryWriter: self._close_deadline = deadline self._accepting = False thread = self._thread + self._notify_invalidation() self._wake.set() if thread is not None: thread.join(max(0.0, deadline - time.monotonic())) @@ -369,6 +380,8 @@ class HistoryWriter: self._generation += 1 self._next_probe = self._clock() + _PROBE_INTERVAL self._discard_locked() + if failed: + self._notify_invalidation() with self._lock: gaps = self._gaps has_work = self._epoch is not None and not self._queue.empty() @@ -386,6 +399,7 @@ class HistoryWriter: self._accepting = False self._epoch = None self._discard_locked() + self._notify_invalidation() if unsubscribe is not None: unsubscribe() if db is not None: