From 0bdb33230fcaad45107242ca3d4f5afa19f9637a Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Sat, 17 Oct 2026 18:38:56 -0700 Subject: [PATCH] feat(history): capture cached prompt strings without mutating host caches --- tests/test_execution_context_graph.py | 19 ++ tests/test_execution_context_host.py | 397 ++++++++++++++++++++++++++ tests/test_execution_history_hooks.py | 79 +++++ utils/execution_context_capture.py | 68 ++++- utils/execution_context_graph.py | 30 +- utils/execution_context_host.py | 155 ++++++++++ utils/execution_history.py | 19 +- 7 files changed, 755 insertions(+), 12 deletions(-) create mode 100644 tests/test_execution_context_host.py create mode 100644 utils/execution_context_host.py diff --git a/tests/test_execution_context_graph.py b/tests/test_execution_context_graph.py index 56f3267..6759ac2 100644 --- a/tests/test_execution_context_graph.py +++ b/tests/test_execution_context_graph.py @@ -244,3 +244,22 @@ class ProjectionTests(unittest.TestCase): g[str(i)] = node("KSampler", positive=["p", 0], negative=["p", 0]) result = self.m.project_contexts(g, None) self.assertEqual(result.recipes, ()) + + def test_partial_source_snapshots_keep_literal_fields_without_stateful_substitution( + self, + ): + for upstream in ("RuntimeGenerator", "PromptManagerText"): + g = graph() + g["dynamic"] = node(upstream, text="looks literal") + g["p"]["inputs"]["text"] = ["dynamic", 0] + projection = self.m.project_contexts(g, None) + positive = next( + snapshot + for ref, snapshot in projection.source_snapshots + if ref.output_slot == 0 + ) + self.assertIsNone(positive.raw.value) + self.assertIsNone(positive.effective.value) + self.assertEqual(positive.prepend.value, "pre") + self.assertEqual(positive.append.value, "post") + self.assertEqual(positive.append.provenance, "inferred") diff --git a/tests/test_execution_context_host.py b/tests/test_execution_context_host.py new file mode 100644 index 0000000..e65aae1 --- /dev/null +++ b/tests/test_execution_context_host.py @@ -0,0 +1,397 @@ +"""Read-only cached STRING evidence and cooperative executor notifications.""" + +from collections import namedtuple +from contextvars import Context +from dataclasses import replace +import gc +import sys +import weakref +import threading +import types +import unittest +from unittest.mock import patch + +from utils.execution_context_budget import ContextBudget +from utils.execution_context_capture import ContextCapture +from utils.execution_history_events import RunStart +from utils.execution_history_writer import HistoryWriterStatus + +try: + from utils import execution_context_host as host +except ImportError: + host = None + +SESSION = "11111111-1111-4111-8111-111111111111" +EPOCH = "22222222-2222-4222-8222-222222222222" +JOB = "33333333-3333-4333-8333-333333333333" +Entry = namedtuple("CacheEntry", "ui outputs") + + +class HostTests(unittest.TestCase): + def setUp(self): + self.assertIsNotNone(host, "read-only host adapter is missing") + self.queue = object() + self.pool = ContextBudget() + self.status = HistoryWriterStatus(EPOCH, True, 0, 0, True, 3) + self.capture = ContextCapture(id(self.queue), lambda: self.status, self.pool) + self.addCleanup(self.capture.revoke_all) + self.calls = [] + calls = self.calls + + class BasicCache: + def __init__(self): + self.entries = {} + + def get_local(self, node_id): + calls.append(("local", node_id)) + return self.entries.get(node_id) + + async def get(self, node_id): + raise AssertionError("async get must never be called") + + class HierarchicalCache(BasicCache): + pass + + class LRUCache(BasicCache): + pass + + class RAMPressureCache(LRUCache): + pass + + class NullCache: + def get_local(self, node_id): + raise AssertionError("null cache is not usable") + + class PromptExecutor: + def add_message(self, event, data, broadcast): + calls.append(("original", event, data, broadcast)) + if event == "raise": + raise RuntimeError("original") + return "original-return" + + self.execution = types.ModuleType("execution") + self.execution.PromptExecutor = PromptExecutor + self.execution.CacheEntry = Entry + self.caching = types.ModuleType("comfy_execution.caching") + for cls in ( + BasicCache, + HierarchicalCache, + LRUCache, + RAMPressureCache, + NullCache, + ): + setattr(self.caching, cls.__name__, cls) + package = types.ModuleType("comfy_execution") + package.caching = self.caching + self.patch = patch.dict( + sys.modules, + { + "comfy_execution": package, + "execution": self.execution, + "comfy_execution.caching": self.caching, + }, + ) + self.patch.start() + self.addCleanup(self.patch.stop) + self.executor = PromptExecutor() + self.executor.server = types.SimpleNamespace(prompt_queue=self.queue) + self.executor.caches = types.SimpleNamespace(outputs=HierarchicalCache()) + + def begin( + self, kind="PromptManager", raw=" raw ", prepend=" pre ", append=" post " + ): + g = { + "p": { + "class_type": kind, + "inputs": { + "text": raw, + "negative_text": "negative", + "prepend_text": prepend, + "append_text": append, + }, + }, + "clip": {"class_type": "CLIPTextEncode", "inputs": {"text": ["p", 0]}}, + "s": { + "class_type": "KSampler", + "inputs": { + "positive": ( + ["clip", 0] if kind == "PromptManagerText" else ["p", 0] + ), + "negative": ( + ["p", 1] if kind == "PromptManagerPair" else ["clip", 0] + ), + }, + }, + } + if type(raw) is list or type(prepend) is list: + g["dynamic"] = {"class_type": "RuntimeGenerator", "inputs": {}} + g["stateful"] = { + "class_type": "PromptManagerText", + "inputs": {"text": "looks literal"}, + } + self.cap = self.capture.begin( + RunStart(SESSION, SESSION, 1, JOB, EPOCH, 0), 3, g, None + ) + return self.cap + + def observe(self, event, executor=None, **data): + host.observe_cached_context( + executor or self.executor, event, {"prompt_id": JOB, **data} + ) + + def seal(self): + sealed = self.capture.seal(self.cap) + if sealed.lease: + self.addCleanup(sealed.lease.release) + return sealed.bundle.contexts[0] + + def test_original_return_exception_and_wrapper_cleanup_preserved(self): + cls = self.execution.PromptExecutor + original = cls.add_message + + def earlier(executor, *args, **kwargs): + self.calls.append("earlier") + return original(executor, *args, **kwargs) + + cls.add_message = earlier + seen = [] + stop_a = host.subscribe_executor_messages("a", lambda *args: seen.append(args)) + stop_b = host.subscribe_executor_messages("b", lambda *args: seen.append(args)) + wrapped = cls.add_message + + def later(executor, *args, **kwargs): + self.calls.append("later") + return wrapped(executor, *args, **kwargs) + + cls.add_message = later + data = {"prompt_id": JOB} + self.assertEqual( + self.executor.add_message("ok", data, broadcast=False), "original-return" + ) + self.assertEqual(len(seen), 2) + self.assertIs(seen[0][0], self.executor) + self.assertIs(seen[0][1][1], data) + with self.assertRaisesRegex(RuntimeError, "original"): + self.executor.add_message("raise", data, False) + self.assertEqual(len(seen), 2) + stop_a() + self.executor.add_message("ok", data, False) + self.assertEqual(len(seen), 3) + stop_b() + self.assertIs(cls.add_message, later) + self.executor.add_message("ok", data, False) + self.assertEqual(len(seen), 3) + stop_c = host.subscribe_executor_messages("c", lambda *args: seen.append(args)) + self.executor.add_message("ok", data, False) + self.assertEqual(len(seen), 4) + stop_c() + + def test_cache_reads_only_local_string_slots(self): + for mode in ("HierarchicalCache", "LRUCache", "RAMPressureCache"): + for kind, values in ( + ("PromptManager", [object(), ["cached"]]), + ("PromptManagerFocused", [object(), ["cached"]]), + ("PromptManagerSplit", [object(), ["cached"]]), + ("PromptManagerText", [["cached"]]), + ("PromptManagerPair", [object(), object(), ["cached"], [""]]), + ): + with self.subTest(mode=mode, kind=kind): + self.begin(kind) + cache = getattr(self.caching, mode)() + entry = Entry({"secret": object()}, values) + cache.entries["p"] = entry + self.executor.caches.outputs = cache + self.observe("execution_start") + self.observe("execution_cached", nodes=["p"]) + result = self.seal() + self.assertEqual(result.positive.effective.value, "cached") + self.assertEqual( + result.positive.effective.source, "host_cache_output" + ) + if kind == "PromptManagerPair": + self.assertEqual(result.negative.effective.value, "") + self.assertIs(cache.entries["p"], entry) + self.assertEqual(set(cache.__dict__), {"entries"}) + self.assertTrue(all(call[0] == "local" for call in self.calls)) + + def test_changed_trigger_configuration_keeps_cached_effective_text(self): + self.begin() + self.executor.caches.outputs.entries["p"] = Entry( + None, [object(), ["old injected text"]] + ) + self.observe("execution_start") + self.observe("execution_cached", nodes=["p"]) + result = self.seal().positive + self.assertEqual(result.effective.value, "old injected text") + self.assertEqual(result.raw.value, " raw ") + self.assertEqual(result.prepend.value, " pre ") + self.assertEqual(result.append.value, " post ") + self.assertEqual(result.raw.provenance, "inferred") + + def test_dynamic_cached_raw_is_unknown(self): + for dynamic in (["dynamic", 0], ["stateful", 0]): + for field in ("raw", "prepend"): + with self.subTest(dynamic=dynamic, field=field): + self.begin(**{field: dynamic}) + self.executor.caches.outputs.entries["p"] = Entry( + None, [object(), ["cached"]] + ) + self.observe("execution_start") + self.observe("execution_cached", nodes=["p"]) + result = self.seal().positive + self.assertEqual(result.effective.value, "cached") + self.assertIsNone(getattr(result, field).value) + self.assertEqual(result.append.value, " post ") + if field == "raw": + self.assertEqual(result.prepend.value, " pre ") + else: + self.assertEqual(result.raw.value, " raw ") + + def test_provider_local_reset_eviction_and_unknown_modes_fail_closed(self): + for mode in ( + "NullCache", + "unknown", + "reset", + "evicted", + "malformed", + "provider-local", + ): + with self.subTest(mode=mode): + self.begin() + cache = self.caching.HierarchicalCache() + if mode == "NullCache": + cache = self.caching.NullCache() + elif mode == "unknown": + cache = types.SimpleNamespace( + get_local=lambda _: self.fail("unknown cache read") + ) + elif mode == "malformed": + cache.entries["p"] = Entry(None, [object(), [object()]]) + elif mode in ("reset", "provider-local"): + cache.entries["p"] = Entry( + None, [object(), ["resident provider output"]] + ) + self.executor.caches.outputs = cache + self.observe("execution_start") + if mode == "reset": + self.executor.caches.outputs = self.caching.HierarchicalCache() + self.observe("execution_cached", nodes=["p"]) + result = self.seal().positive + self.assertEqual( + result.effective.value, + "resident provider output" if mode == "provider-local" else None, + ) + + def test_worker_queue_job_executor_and_generation_fences(self): + for failure in ( + "unbound", + "thread", + "queue", + "job", + "executor", + "recursive", + "generation", + "no-start", + ): + with self.subTest(failure=failure): + self.begin() + self.executor.server.prompt_queue = self.queue + self.executor.caches.outputs.entries["p"] = Entry( + None, [object(), ["cached"]] + ) + if failure != "no-start": + self.observe("execution_start") + target = self.executor + if failure == "queue": + self.executor.server.prompt_queue = object() + if failure in ("executor", "recursive"): + target = self.execution.PromptExecutor() + target.server = self.executor.server + target.caches = self.executor.caches + if failure == "recursive": + self.observe("execution_start", target) + target = self.executor + data = { + "prompt_id": "wrong" if failure == "job" else JOB, + "nodes": ["p"], + } + if failure == "generation": + self.status = replace(self.status, capture_generation=4) + + def action(): + host.observe_cached_context(target, "execution_cached", data) + + if failure == "unbound": + Context().run(action) + elif failure == "thread": + thread = threading.Thread(target=action) + thread.start() + thread.join(2) + else: + action() + if failure == "generation": + self.assertIsNone(self.capture.seal(self.cap).bundle) + self.status = replace(self.status, capture_generation=3) + else: + self.assertIsNone(self.seal().positive.effective.value) + + def test_missing_cache_preserves_known_affixes_when_raw_is_dynamic(self): + self.begin(raw=["dynamic", 0]) + self.observe("execution_start") + self.observe("execution_cached", nodes=["p"]) + result = self.seal().positive + self.assertIsNone(result.raw.value) + self.assertIsNone(result.effective.value) + self.assertEqual(result.prepend.value, " pre ") + self.assertEqual(result.append.value, " post ") + + def test_cache_entries_tensors_and_executor_are_not_retained(self): + class Tensor: + pass + + tensor = Tensor() + tensor_ref = weakref.ref(tensor) + self.begin() + self.executor.caches.outputs.entries["p"] = Entry(tensor, [tensor, ["cached"]]) + self.observe("execution_start") + self.observe("execution_cached", nodes=["p"]) + self.executor.caches.outputs.entries.clear() + executor_ref = weakref.ref(self.executor) + del tensor + del self.executor + gc.collect() + self.assertIsNone(tensor_ref()) + self.assertIsNone(executor_ref()) + self.assertEqual(self.seal().positive.effective.value, "cached") + + def test_unknown_and_malformed_cache_payloads_are_bounded(self): + for values in (["x"] * 65, ["x" * 65537], [object()], ["x", "x"]): + with self.subTest(count=len(values)): + self.begin() + self.executor.caches.outputs.entries["p"] = Entry( + None, [object(), values] + ) + self.observe("execution_start") + self.observe("execution_cached", nodes=["p"]) + result = self.seal().positive + self.assertIsNone(result.effective.value) + self.assertEqual( + result.append.value, " post " if len(values) != 2 else None + ) + + def test_disposed_subscriptions_do_not_retain_dynamic_executor_classes(self): + def install_and_remove(): + class TemporaryExecutor: + def add_message(self, *args, **kwargs): + return None + + self.execution.PromptExecutor = TemporaryExecutor + ref = weakref.ref(TemporaryExecutor) + stop = host.subscribe_executor_messages("temporary", lambda *args: None) + stop() + del self.execution.PromptExecutor + return ref + + ref = install_and_remove() + gc.collect() + self.assertIsNone(ref()) diff --git a/tests/test_execution_history_hooks.py b/tests/test_execution_history_hooks.py index 6d518c3..78c745d 100644 --- a/tests/test_execution_history_hooks.py +++ b/tests/test_execution_history_hooks.py @@ -625,3 +625,82 @@ class HistoryHooksTests(unittest.TestCase): self.queue.task_done(10, {}, status()) self.drain(second) self.assertEqual(len(self.rows()), 1) + + def test_cached_executor_subscription_persists_strings_and_disposes_cooperatively( + self, + ): + from utils.execution_context_host import subscribe_executor_messages + + calls = [] + execution = types.ModuleType("execution") + caching = types.ModuleType("comfy_execution.caching") + package = types.ModuleType("comfy_execution") + entry_type = namedtuple("CacheEntry", "ui outputs") + + class PromptExecutor: + def add_message(self, event, data, broadcast): + calls.append(event) + return "host-result" + + class HierarchicalCache: + def get_local(self, node_id): + return entry_type(None, [object(), object(), ["old trigger"], [""]]) + + execution.PromptExecutor, execution.CacheEntry = PromptExecutor, entry_type + caching.HierarchicalCache = HierarchicalCache + package.caching = caching + with patch.dict( + sys.modules, + { + "execution": execution, + "comfy_execution": package, + "comfy_execution.caching": caching, + }, + ): + observer_calls = [] + independent = subscribe_executor_messages( + "independent", lambda *args: observer_calls.append(args) + ) + try: + registration = self.register() + g = { + "p": { + "class_type": "PromptManagerPair", + "inputs": {"text": "raw", "negative_text": ""}, + }, + "s": { + "class_type": "KSampler", + "inputs": {"positive": ["p", 0], "negative": ["p", 1]}, + }, + } + self.queue.results.append(((0, JOB, g, {}, ["s"]), 1)) + self.queue.get() + executor = PromptExecutor() + executor.server = self.server + executor.caches = SimpleNamespace(outputs=HierarchicalCache()) + self.assertEqual( + executor.add_message("execution_start", {"prompt_id": JOB}, False), + "host-result", + ) + executor.add_message( + event="execution_cached", + data={"prompt_id": JOB, "nodes": ["p"]}, + broadcast=False, + ) + self.queue.task_done(1, {}, status()) + self.drain(registration) + result = self.db.get_execution_contexts(self.rows()[0]["run_id"]) + self.assertEqual( + result["contexts"][0]["positive"]["effective"], "old trigger" + ) + self.assertEqual(result["contexts"][0]["negative"]["effective"], "") + self.assertTrue(registration.dispose()) + executor.add_message( + "execution_cached", {"prompt_id": JOB, "nodes": ["p"]}, False + ) + self.assertEqual(len(observer_calls), 3) + finally: + independent() + self.assertEqual( + calls, ["execution_start", "execution_cached", "execution_cached"] + ) diff --git a/utils/execution_context_capture.py b/utils/execution_context_capture.py index a9d1ad0..05cfaee 100644 --- a/utils/execution_context_capture.py +++ b/utils/execution_context_capture.py @@ -51,6 +51,8 @@ class _Attempt: samples: dict retained_bytes: int reason: str = "" + executor_identity: int | None = None + executor_rejected: bool = False _binding: ContextVar[tuple | None] = ContextVar("prompt_manager_attempt", default=None) @@ -164,13 +166,20 @@ class ContextCapture: 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 + retained = ( + sum( + _snapshot_bytes(role.inferred) if role.inferred else 0 + for recipe in projection.recipes + for role in (recipe.positive, recipe.negative) + ) + + sum( + _snapshot_bytes(snapshot) + for _, snapshot in projection.source_snapshots + ) + + 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") @@ -186,6 +195,36 @@ class ContextCapture: self._current(cap) return cap + def bind_executor(self, capability: AttemptCapability, identity: int) -> None: + """Bind one scalar host identity; recursive starts invalidate cache proof.""" + with self._lock: + attempt = self._current(capability) + if attempt is not None: + if attempt.executor_identity is not None: + attempt.executor_rejected = True + else: + attempt.executor_identity = identity + + def matches_executor(self, capability: AttemptCapability, identity: int) -> bool: + with self._lock: + attempt = self._current(capability) + return bool( + attempt is not None + and not attempt.executor_rejected + and attempt.executor_identity == identity + ) + + def source_snapshot( + self, capability: AttemptCapability, source: SourceRef + ) -> PromptSnapshot | None: + with self._lock: + attempt = self._current(capability) + if attempt is not None: + for ref, snapshot in attempt.projection.source_snapshots: + if ref == source: + return snapshot + return None + def observe( self, capability: AttemptCapability, @@ -292,8 +331,19 @@ class ContextCapture: return _unknown("ambiguous") if len(role.sources) == 1 and len(values) == 1: return values[0] - if role.inferred is not None: - s = role.inferred + inferred = role.inferred + if len(role.sources) == 1: + snapshots = attempt.projection.source_snapshots + inferred = next( + ( + snapshot + for source, snapshot in snapshots + if source == role.sources[0] + ), + inferred, + ) + if inferred is not None: + s = inferred return PromptSnapshot( s.raw, _unknown().effective, s.prepend, s.append ) diff --git a/utils/execution_context_graph.py b/utils/execution_context_graph.py index 82ff540..985a145 100644 --- a/utils/execution_context_graph.py +++ b/utils/execution_context_graph.py @@ -32,6 +32,7 @@ class GraphProjection: recipes: tuple[DeclaredRecipe, ...] sources: tuple[SourceRef, ...] reason: str + source_snapshots: tuple[tuple[SourceRef, PromptSnapshot], ...] = () _SAMPLERS = {"KSampler", "KSamplerAdvanced", "SamplerCustom"} @@ -214,6 +215,26 @@ class _Projector: raise _Unknown("source_reference_limit_reached") return snapshot, sources + def source_snapshot(self, source: SourceRef) -> PromptSnapshot: + name, inputs = self.node(source.node_id, 1) + adapter = _DEFAULT_ADAPTERS[(name, source.output_slot)] + unknown = TextEvidence(None, "unknown", "unknown", "missing") + + def literal(value: object) -> TextEvidence: + try: + snapshot, sources = self.text(value, "string", 2) + # Stateful PM outputs cannot prove their injected result from inputs. + return unknown if sources else snapshot.effective + except _Limit: + raise + except ValueError: + return unknown + + raw = literal(inputs.get(adapter.input_name)) + before = literal(inputs.get("prepend_text", "") if adapter.compose else "") + after = literal(inputs.get("append_text", "") if adapter.compose else "") + return self.retain(PromptSnapshot(raw, unknown, before, after)) + def discover_sources( self, value: object, kind: str, depth: int = 1, active: frozenset = frozenset() ) -> tuple[SourceRef, ...]: @@ -266,7 +287,7 @@ class _Projector: def retain(self, snapshot: PromptSnapshot) -> PromptSnapshot: size = sum( - text_size(field.value) + 64 + (text_size(field.value) if field.value is not None else 0) + 64 for field in ( snapshot.raw, snapshot.effective, @@ -388,7 +409,12 @@ def project_contexts(graph: object, outputs: object) -> GraphProjection: for role in (recipe.positive, recipe.negative) ): reason = "some_roles_require_inferred_or_unresolved_text" - return GraphProjection(tuple(recipes), sources, reason) + return GraphProjection( + tuple(recipes), + sources, + reason, + tuple((source, resolver.source_snapshot(source)) for source in sources), + ) except ValueError as error: reason = str(error) if isinstance(error, _Unknown) else "invalid_text_value" return GraphProjection((), (), reason) diff --git a/utils/execution_context_host.py b/utils/execution_context_host.py new file mode 100644 index 0000000..c3d8ba3 --- /dev/null +++ b/utils/execution_context_host.py @@ -0,0 +1,155 @@ +"""Pinned-host v1 adapter for local cached STRING evidence only.""" + +import functools +import threading +from typing import Callable + +from .execution_context_capture import active_capture +from .execution_contexts import PromptSnapshot, TextEvidence + +_STATE_ATTRIBUTE = "_prompt_manager_executor_messages_v1" +_registration_lock = threading.Lock() + + +class _Messages: + def __init__(self, executor_class: type) -> None: + self.executor_class = executor_class + self.lock = threading.Lock() + self.callbacks = {} + self.wrapper = None + self.original = None + self.installed = False + + def subscribe(self, key: str, callback: Callable) -> Callable[[], None]: + with self.lock: + if not self.installed: + original = self.executor_class.add_message + if not callable(original): + raise ValueError("Executor message capability unavailable") + + @functools.wraps(original) + def observe(executor, *args, **kwargs): + result = original(executor, *args, **kwargs) + with self.lock: + callbacks = tuple(entry[1] for entry in self.callbacks.values()) + for notify in callbacks: + try: + notify(executor, args, kwargs) + except Exception: + # Best-effort observers cannot change host results. + pass + return result + + self.executor_class.add_message = observe + self.original, self.wrapper, self.installed = original, observe, True + existing = self.callbacks.get(key) + token = existing[0] if existing is not None else object() + if existing is None: + self.callbacks[key] = (token, callback) + + def unsubscribe() -> None: + with self.lock: + current = self.callbacks.get(key) + if current is None or current[0] is not token: + return + del self.callbacks[key] + if ( + not self.callbacks + and self.executor_class.add_message is self.wrapper + ): + self.executor_class.add_message = self.original + self.installed = False + + return unsubscribe + + +def subscribe_executor_messages(key: str, callback: Callable) -> Callable[[], None]: + """Observe successful original calls; preserve earlier and later wrappers.""" + from execution import PromptExecutor + + with _registration_lock: + state = vars(PromptExecutor).get(_STATE_ATTRIBUTE) + if state is None: + state = _Messages(PromptExecutor) + setattr(PromptExecutor, _STATE_ATTRIBUTE, state) + return state.subscribe(key, callback) + + +def observe_cached_context( + executor: object, message_type: object, message_data: object +) -> None: + """Copy bounded local strings without cache/provider execution or retention.""" + try: + binding = active_capture() + if binding is None or type(message_type) is not str: + return + manager, cap = binding + if ( + type(message_data) is not dict + or type(message_data.get("prompt_id")) is not str + or message_data["prompt_id"] != cap.job_id + or id(executor.server.prompt_queue) != manager.queue_identity + ): + return + if message_type == "execution_start": + manager.bind_executor(cap, id(executor)) + return + if message_type != "execution_cached" or not manager.matches_executor( + cap, id(executor) + ): + return + nodes = message_data.get("nodes") + if type(nodes) not in (list, tuple) or len(nodes) > 500: + return + from comfy_execution import caching + from execution import CacheEntry + + cache = executor.caches.outputs + supported = tuple( + cls + for name in ("HierarchicalCache", "LRUCache", "RAMPressureCache") + if isinstance(cls := getattr(caching, name, None), type) + ) + if type(cache) not in supported: + return + local = getattr(cache, "get_local", None) + if not callable(local): + return + # Only source IDs retained by the bounded submitted graph are queried. + for source in cap.sources: + if not any(type(node) is str and node == source.node_id for node in nodes): + continue + slot = ( + (3 if source.output_slot in (1, 3) else 2) + if source.node_type == "PromptManagerPair" + else 0 if source.node_type == "PromptManagerText" else 1 + ) + entry = local(source.node_id) + if type(entry) is not CacheEntry: + continue + outputs = entry.outputs + if type(outputs) not in (list, tuple) or not slot < len(outputs) <= 4: + continue + values = outputs[slot] + if type(values) not in (list, tuple) or len(values) > 64: + continue + inferred = manager.source_snapshot(cap, source) + if inferred is None: + continue + for index, value in enumerate(values): + if type(value) is not str: + continue + manager.observe( + cap, + source, + index, + PromptSnapshot( + inferred.raw, + TextEvidence(value, "recorded", "host_cache_output", None), + inferred.prepend, + inferred.append, + ), + ) + except Exception: + # Unknown host versions and malformed entries leave evidence unavailable. + return diff --git a/utils/execution_history.py b/utils/execution_history.py index 6a4e952..f63475e 100644 --- a/utils/execution_history.py +++ b/utils/execution_history.py @@ -13,7 +13,8 @@ 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_context_capture import ContextCapture, active_capture, unbind_current +from .execution_context_host import observe_cached_context, subscribe_executor_messages from .execution_history_writer import HistoryWriter from .logging_config import get_logger from .queue_observers import subscribe_dequeue, subscribe_task_done @@ -57,12 +58,28 @@ class HistoryRegistration: self._unsubscribers.append( subscribe_task_done(queue, "execution_history", self._task_done) ) + try: + self._unsubscribers.append( + subscribe_executor_messages( + f"execution_history_{id(self)}", self._executor_message + ) + ) + except (ImportError, AttributeError, TypeError, ValueError): + pass self.writer.start() atexit.register(self.dispose) except Exception: self.dispose() raise + def _executor_message(self, executor: object, args: tuple, kwargs: dict) -> None: + binding = active_capture() + if binding is None or binding[0] is not self.capture: + return + event = args[0] if args else kwargs.get("event") + data = args[1] if len(args) > 1 else kwargs.get("data") + observe_cached_context(executor, event, data) + def _invalidate_capture(self) -> None: self.capture.revoke_all()