156 lines
5.9 KiB
Python
156 lines
5.9 KiB
Python
"""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
|