diff --git a/__init__.py b/__init__.py index c1e1dd5..2663b6a 100644 --- a/__init__.py +++ b/__init__.py @@ -158,6 +158,34 @@ except Exception as e: f"Failed to register usage hook: {e}" ) + +def _history_db_factory() -> "PromptDatabase": + """Initialize history on its worker even if monitoring initialization failed.""" + if _global_db is not None: + return _global_db + from .database.operations import PromptDatabase + + return PromptDatabase() + + +_history_registration = None +try: + from .py.config import server_instance + from .utils.execution_history import register_history_hook + + _history_registration = register_history_hook(server_instance, _history_db_factory) +except Exception: + try: + from .utils.logging_config import get_logger + + get_logger("prompt_manager.init").error("Failed to register execution history") + except Exception: + import logging + + logging.getLogger("prompt_manager.init").error( + "Failed to register execution history" + ) + __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] # Print startup message with loaded tools diff --git a/tests/test_execution_history_hooks.py b/tests/test_execution_history_hooks.py new file mode 100644 index 0000000..fa9e627 --- /dev/null +++ b/tests/test_execution_history_hooks.py @@ -0,0 +1,451 @@ +"""Actual dequeue and completion capture through the owned SQLite worker.""" + +from collections import namedtuple +import logging +import gc +import contextlib +import io +import os +import tempfile +import threading +import time +import sys +import types +import weakref +from types import SimpleNamespace +import unittest +from unittest.mock import patch + +from database import PromptDatabase +from utils import usage_tracking +from utils.execution_history_writer import HistoryWriter + +try: + from utils.execution_history import register_history_hook +except ImportError: + register_history_hook = None + +Status = namedtuple("ExecutionStatus", "status_str completed messages") +JOB = "a0c26da3-eb26-405c-b143-c662a84c3112" + + +class HostQueue: + def __init__(self): + self.results = [] + self.done_result = object() + self.done_calls = [] + + def get(self, *args, **kwargs): + return self.results.pop(0) if self.results else None + + def task_done(self, *args, **kwargs): + self.done_calls.append((args, kwargs)) + return self.done_result + + +def item(item_id, job=JOB, graph=None, extra=None, sensitive=None): + return ((0, job, graph or {}, extra or {}, [], sensitive or {}), item_id) + + +def status(kind="execution_success", job=JOB): + return Status("success", True, [(kind, {"prompt_id": job, "timestamp": 20})]) + + +class HistoryHooksTests(unittest.TestCase): + def setUp(self): + self.directory = tempfile.TemporaryDirectory() + self.addCleanup(self.directory.cleanup) + self.db = PromptDatabase(os.path.join(self.directory.name, "history.db")) + self.addCleanup(self.db.model.reset_schema_cache) + self.addCleanup(self.db.model.close_all) + self.queue = HostQueue() + self.server = SimpleNamespace(prompt_queue=self.queue) + self.registrations, self.releases = [], [] + self.addCleanup(self.cleanup_registrations) + for name in ("execution_history_writer", "execution_history"): + logger = logging.getLogger("prompt_manager." + name) + previous = logger.handlers, logger.propagate, logger.level + logger.handlers, logger.propagate, logger.level = ( + [logging.NullHandler()], + False, + logging.WARNING, + ) + self.addCleanup(self.restore_logger, logger, previous) + + def restore_logger(self, logger, previous): + logger.handlers, logger.propagate, logger.level = previous + + def cleanup_registrations(self): + for release in self.releases: + release.set() + for registration in self.registrations: + self.assertTrue(registration.dispose(3), "owned worker leaked") + + def wait_for(self, predicate): + deadline = time.monotonic() + 3 + while time.monotonic() < deadline: + if predicate(): + return + threading.Event().wait(0.005) + self.fail("finite deadline expired waiting for recorded history") + + def register(self, factory=None, ready=True): + self.assertTrue( + callable(register_history_hook), "history registration is absent" + ) + registration = register_history_hook(self.server, factory or (lambda: self.db)) + self.assertIsNotNone(registration) + self.registrations.append(registration) + if ready: + self.wait_for(lambda: registration.writer.status().accepting) + return registration + + def dequeue(self, item_id=0, **kwargs): + result = item(item_id, **kwargs) + self.queue.results.append(result) + self.assertIs(self.queue.get(), result) + return result + + def rows(self): + return [ + dict(row) + for row in self.db.model.get_connection() + .execute("SELECT * FROM execution_runs ORDER BY queue_item_id") + .fetchall() + ] + + def drain(self, registration): + self.wait_for(lambda: registration.writer.status().pending == 0) + + def test_package_history_initializes_when_monitor_and_usage_setup_fail(self): + try: + from tests.test_init_package import ( + close_every_database, + fake_server_module, + load_package, + ) + except ImportError: + from test_init_package import ( + close_every_database, + fake_server_module, + load_package, + ) + name = "prompt_manager_pkg_history_independent" + server = fake_server_module() + server.PromptServer.instance.prompt_queue = self.queue + server.PromptServer.instance.app = SimpleNamespace(middlewares=[]) + + def reject_usage(_): + raise RuntimeError("usage registration unavailable") + + server.PromptServer.instance.add_on_prompt_handler = reject_usage + output = os.path.join(self.directory.name, "output") + os.makedirs(output) + monitor_module = types.ModuleType(name + ".utils.image_monitor") + # Missing monitor import fails before shared database initialization. + with ( + patch.dict(sys.modules, {monitor_module.__name__: monitor_module}), + contextlib.redirect_stdout(io.StringIO()), + self.assertLogs("prompt_manager.init", level="ERROR"), + ): + package, _ = load_package(name, self.db.model.db_path, output, server) + self.addCleanup(sys.modules.pop, name, None) + self.addCleanup(close_every_database) + registration = getattr(package, "_history_registration", None) + self.assertIsNotNone(registration, "history setup depended on another hook") + self.registrations.append(registration) + self.addCleanup(registration.dispose) + self.assertIsNone(package._global_db) + self.wait_for(lambda: registration.writer.status().accepting) + self.dequeue(0) + self.queue.task_done(0, {}, status()) + self.drain(registration) + self.assertEqual(self.rows()[0]["outcome"], "succeeded") + + def test_distinct_actual_attempts_for_reused_job_and_cached_rerun(self): + registration = self.register() + self.dequeue(0) + self.queue.task_done(0, {"outputs": {}}, status("execution_error")) + self.dequeue(1) + self.queue.task_done(item_id=1, history_result={}, status=status()) + self.drain(registration) + rows = self.rows() + self.assertEqual(len(rows), 2) + self.assertEqual([r["job_id"] for r in rows], [JOB, JOB]) + self.assertEqual([r["outcome"] for r in rows], ["failed", "succeeded"]) + self.assertNotEqual(rows[0]["run_id"], rows[1]["run_id"]) + + def test_later_conflicting_completion_resolves_completed_stored_identity(self): + registration = self.register() + self.dequeue(0) + self.queue.task_done(0, {}, status()) + self.drain(registration) + self.queue.task_done(0, {}, status("execution_error")) + self.drain(registration) + self.assertEqual(self.rows()[0]["outcome"], "unknown") + self.assertEqual(self.rows()[0]["reason"], "conflicting_terminal") + + def test_repeated_registration_cleanup_and_reregistration_keep_session_and_wrappers( + self, + ): + before = self.queue.get + calls = [] + + def unrelated_before(*a, **k): + calls.append("before") + return before(*a, **k) + + self.queue.get = unrelated_before + first = self.register() + self.assertIs(register_history_hook(self.server, lambda: self.db), first) + observed = self.queue.get + + def unrelated_after(*a, **k): + calls.append("after") + return observed(*a, **k) + + self.queue.get = unrelated_after + self.dequeue(0) + self.queue.task_done(0, {}, status()) + self.assertTrue(first.dispose()) + self.assertTrue(first.dispose()) + self.assertIs(self.queue.get, unrelated_after) + second = self.register() + self.dequeue(1) + self.queue.task_done(1, {}, status()) + self.drain(second) + self.assertEqual(len(self.rows()), 2) + self.assertEqual(self.rows()[0]["session_id"], self.rows()[1]["session_id"]) + self.assertEqual(calls, ["after", "before"] * 2) + + def test_concurrent_handle_capacity_never_reassigns_active_attempts(self): + registration = self.register() + self.queue.results.extend(item(index) for index in range(129)) + ready = threading.Barrier(130) + errors = [] + + def dequeue(): + try: + ready.wait(3) + self.queue.get() + except Exception as error: + errors.append(error) + + threads = [threading.Thread(target=dequeue) for _ in range(129)] + for thread in threads: + thread.start() + ready.wait(3) + for thread in threads: + thread.join(3) + self.assertFalse(thread.is_alive()) + self.assertEqual(errors, []) + self.drain(registration) + rows = self.rows() + admitted = {row["queue_item_id"] for row in rows} + self.assertEqual(len(rows), 128) + missing = (set(range(129)) - admitted).pop() + oldest = min(admitted) + self.queue.task_done(oldest, {}, status()) + self.queue.task_done(missing, {}, status()) + self.drain(registration) + by_item = {row["queue_item_id"]: row for row in self.rows()} + self.assertEqual(by_item[oldest]["outcome"], "succeeded") + self.assertNotIn(missing, by_item) + self.assertGreater(registration.writer.status().recording_gaps, 0) + self.dequeue(129) + self.queue.task_done(129, {}, status()) + self.drain(registration) + self.assertEqual(self.rows()[-1]["queue_item_id"], 129) + + def test_active_handles_do_not_retain_graph_extra_or_sensitive_envelopes(self): + class Envelope(dict): + pass + + registration = self.register() + graph, extra, sensitive = ( + Envelope({"g": object()}), + Envelope({"e": object()}), + Envelope({"s": object()}), + ) + references = [weakref.ref(value) for value in (graph, extra, sensitive)] + self.dequeue(0, graph=graph, extra=extra, sensitive=sensitive) + self.drain(registration) + del graph, extra, sensitive + gc.collect() + self.assertEqual([reference() for reference in references], [None, None, None]) + self.queue.task_done(0, {}, status()) + self.drain(registration) + self.assertEqual(self.rows()[0]["outcome"], "succeeded") + + def test_capture_uses_no_database_methods_on_the_host_thread(self): + registration = self.register() + host_thread = threading.get_ident() + original = self.db.model.get_connection + + def connection(): + if threading.get_ident() == host_thread: + self.fail("history queried SQLite on the host thread") + return original() + + with patch.object(self.db.model, "get_connection", connection): + self.dequeue(0) + self.queue.task_done(0, {}, status()) + self.drain(registration) + self.assertEqual(self.rows()[0]["outcome"], "succeeded") + + def test_capture_never_copies_or_serializes_sensitive_objects(self): + class Sensitive: + def __deepcopy__(self, memo): + raise AssertionError("sensitive copy") + + def __repr__(self): + raise AssertionError("sensitive serialization") + + registration = self.register() + secret = Sensitive() + graph, extra, sensitive = {"graph": secret}, {"extra": secret}, {"auth": secret} + result = self.dequeue(0, graph=graph, extra=extra, sensitive=sensitive) + payload = {"outputs": secret, "meta": secret} + completion = status() + completion.messages[0][1].update( + {"traceback": secret, "current_inputs": secret} + ) + self.assertIs( + self.queue.task_done(0, payload, completion), self.queue.done_result + ) + self.drain(registration) + self.assertIs(result[0][2], graph) + self.assertIs(result[0][3], extra) + self.assertIs(result[0][5], sensitive) + self.assertEqual(len(self.rows()), 1) + self.assertEqual(self.rows()[0]["outcome"], "succeeded") + self.assertEqual( + set(completion.messages[0][1]), + {"prompt_id", "timestamp", "traceback", "current_inputs"}, + ) + + def test_callbacks_return_while_writer_blocked_and_disposal_is_bounded(self): + 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(0) + self.assertTrue(entered.wait(1)) + began = time.monotonic() + self.queue.task_done(0, {}, status()) + self.dequeue(1) + self.assertLess(time.monotonic() - began, 0.2) + began = time.monotonic() + self.assertFalse(registration.dispose()) + self.assertGreaterEqual(time.monotonic() - began, 1.8) + self.assertLess(time.monotonic() - began, 2.3) + release.set() + self.assertTrue(registration.dispose(2)) + + def test_restore_disables_capture_and_clears_handles_before_fresh_epoch(self): + try: + from tests.test_execution_history_writer import ManualClock + except ImportError: + from test_execution_history_writer import ManualClock + clock = ManualClock() + with patch( + "utils.execution_history.HistoryWriter", + lambda factory: HistoryWriter(factory, clock=clock), + ): + registration = self.register() + for index in range(128): + self.dequeue(index) + self.drain(registration) + old_epoch = registration.writer.status().epoch + backup = os.path.join(self.directory.name, "backup.db") + self.assertTrue(self.db.model.backup_database(backup)) + self.db.model.restore_from_file(backup) + self.assertFalse(registration.writer.status().accepting) + self.assertIsNone(registration.writer.status().epoch) + self.queue.task_done(0, {}, status()) + self.dequeue(128) + self.assertGreater(registration.writer.status().recording_gaps, 0) + clock.advance(5) + self.wait_for(lambda: registration.writer.status().accepting) + self.assertNotEqual(registration.writer.status().epoch, old_epoch) + # A late final descriptor cannot complete rows from the restored epoch. + self.queue.task_done(0, {}, status()) + self.dequeue(129) + self.queue.task_done(129, {}, status()) + self.drain(registration) + rows = self.rows() + self.assertEqual(len(rows), 129) + self.assertEqual(rows[0]["outcome"], "observed") + self.assertEqual(rows[-1]["queue_item_id"], 129) + self.assertEqual(rows[-1]["outcome"], "succeeded") + + def test_unavailable_database_and_invalid_identity_never_escape_host(self): + entered = threading.Event() + + def broken(): + entered.set() + raise RuntimeError("secret unavailable database") + + registration = self.register(broken, ready=False) + self.assertTrue(entered.wait(1)) + result = self.dequeue(0) + self.assertEqual(result[1], 0) + self.queue.task_done(0, {}, status()) + self.assertEqual(self.rows(), []) + self.assertGreater(registration.writer.status().recording_gaps, 0) + self.assertTrue(registration.dispose()) + registration = self.register() + self.dequeue(0, job="invalid id") + self.queue.task_done(True, {}, object()) + self.drain(registration) + self.assertEqual(self.rows(), []) + self.assertGreater(registration.writer.status().recording_gaps, 0) + + def test_invalid_evidence_is_recorded_unknown_and_raising_status_is_isolated(self): + registration = self.register() + self.dequeue(0) + self.queue.task_done(0, {}, object()) + self.drain(registration) + self.assertEqual(self.rows()[0]["outcome"], "unknown") + self.assertEqual(self.rows()[0]["reason"], "invalid_evidence") + + class BrokenStatus: + @property + def messages(self): + raise RuntimeError("secret failure") + + self.dequeue(1) + self.queue.task_done(1, {}, BrokenStatus()) + self.drain(registration) + self.assertEqual(self.rows()[1]["outcome"], "observed") + self.assertGreater(registration.writer.status().recording_gaps, 0) + + def test_unsupported_queue_does_not_create_writer_or_submit_fallback(self): + self.assertTrue( + callable(register_history_hook), "history registration is absent" + ) + for host in ( + object(), + SimpleNamespace(prompt_queue=SimpleNamespace(get=lambda: None)), + ): + self.assertIsNone( + register_history_hook(host, lambda: self.fail("DB touched")) + ) + + def test_existing_usage_callback_counts_once_alongside_history(self): + registration = self.register() + seen = [] + self.assertTrue(usage_tracking._wrap_queue_get(self.queue, seen.append)) + self.assertFalse(usage_tracking._wrap_queue_get(self.queue, seen.append)) + self.dequeue(0) + self.queue.task_done(0, {}, status()) + self.drain(registration) + self.assertEqual(seen, [JOB]) + self.assertEqual(len(self.rows()), 1) diff --git a/tests/test_execution_history_writer.py b/tests/test_execution_history_writer.py index 7d5c775..3ad2150 100644 --- a/tests/test_execution_history_writer.py +++ b/tests/test_execution_history_writer.py @@ -467,6 +467,20 @@ class HistoryWriterTests(unittest.TestCase): self.assertFalse(writer.status().alive) self.assertFalse(writer.submit(self.start_command())) + def test_external_capture_gaps_are_nonblocking_and_saturate_all_rejections(self): + writer = HistoryWriter(lambda: self.fail("producer accessed database")) + self.assertTrue(callable(getattr(writer, "record_gap", None))) + writer.record_gap() + self.assertEqual(writer.status().recording_gaps, 1) + with writer._lock: + writer._gaps = 2**63 - 2 + writer.record_gap() + writer.record_gap() + writer.submit(self.start_command()) + writer.submit(object()) + self.assertEqual(writer.status().recording_gaps, 2**63 - 1) + self.assertTrue(writer.close()) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_queue_observers.py b/tests/test_queue_observers.py new file mode 100644 index 0000000..e42b520 --- /dev/null +++ b/tests/test_queue_observers.py @@ -0,0 +1,132 @@ +"""Cooperative queue subscriptions preserve the host and other wrappers.""" + +import queue +import unittest + +try: + from utils.queue_observers import subscribe_dequeue, subscribe_task_done +except ImportError: + subscribe_dequeue = subscribe_task_done = None + + +class FakeQueue: + def __init__(self): + self.result = ((0, "job", {}, {}, [], {}), 0) + self.calls = [] + self.failure = None + self.done_result = object() + + def get(self, *args, **kwargs): + self.calls.append(("get", args, kwargs)) + if self.failure: + raise self.failure + return self.result + + def task_done(self, *args, **kwargs): + self.calls.append(("done", args, kwargs)) + if self.failure: + raise self.failure + return self.done_result + + +class QueueObserversTests(unittest.TestCase): + def require_api(self): + self.assertTrue(callable(subscribe_dequeue), "queue subscriptions are absent") + self.assertTrue(callable(subscribe_task_done), "task subscriptions are absent") + + def test_dequeue_preserves_identity_arguments_empty_and_exception(self): + self.require_api() + host, seen, sentinel = FakeQueue(), [], object() + subscribe_dequeue(host, "capture", seen.append) + original = host.result + self.assertIs(host.get(sentinel, timeout=sentinel), original) + self.assertEqual(host.calls, [("get", (sentinel,), {"timeout": sentinel})]) + self.assertEqual(seen, [original]) + host.result = None + self.assertIsNone(host.get()) + self.assertEqual(seen, [original]) + for error in (queue.Empty(), RuntimeError("host failure")): + host.failure = error + with self.assertRaises(type(error)) as raised: + host.get(timeout=sentinel) + self.assertIs(raised.exception, error) + self.assertEqual(seen, [original]) + + def test_task_done_observes_success_only_and_preserves_host_values(self): + self.require_api() + host, seen, sentinel = FakeQueue(), [], object() + subscribe_task_done(host, "capture", lambda a, k: seen.append((a, k))) + self.assertIs(host.task_done(0, sentinel, status=sentinel), host.done_result) + self.assertEqual(seen, [((0, sentinel), {"status": sentinel})]) + error = ValueError("host failure") + host.failure = error + with self.assertRaises(ValueError) as raised: + host.task_done(1) + self.assertIs(raised.exception, error) + self.assertEqual(len(seen), 1) + + def test_keys_are_idempotent_and_independently_removable(self): + self.require_api() + host, first, second = FakeQueue(), [], [] + unsubscribe = subscribe_dequeue(host, "a", first.append) + subscribe_dequeue(host, "a", lambda value: self.fail("duplicate callback")) + remove_second = subscribe_dequeue(host, "b", second.append) + host.get() + unsubscribe() + unsubscribe() + host.get() + self.assertEqual(len(first), 1) + self.assertEqual(len(second), 2) + remove_second() + subscribe_dequeue(host, "a", first.append) + unsubscribe() # old cleanup must not remove a new subscriber with this key + host.get() + self.assertEqual(len(first), 2) + + def test_other_wrappers_before_after_cleanup_and_resubscription_survive(self): + self.require_api() + host, calls, seen = FakeQueue(), [], [] + original = host.get + + def before(*a, **k): + calls.append("before") + return original(*a, **k) + + host.get = before + unsubscribe = subscribe_dequeue(host, "capture", seen.append) + observed = host.get + + def after(*a, **k): + calls.append("after") + return observed(*a, **k) + + host.get = after + host.get() + unsubscribe() + self.assertIs(host.get, after) + host.get() + subscribe_dequeue(host, "capture", seen.append) + host.get() + self.assertIs(host.get, after) + self.assertEqual(calls, ["after", "before"] * 3) + self.assertEqual(len(seen), 2) + + def test_failing_subscriber_cannot_disable_other_subscribers_or_host(self): + self.require_api() + host, seen = FakeQueue(), [] + + def fail(_): + raise ValueError("secret error") + + subscribe_dequeue(host, "bad", fail) + subscribe_dequeue(host, "good", seen.append) + with self.assertLogs("prompt_manager.queue_observers", level="WARNING") as logs: + self.assertIs(host.get(), host.result) + self.assertEqual(seen, [host.result]) + self.assertNotIn("secret error", "\n".join(logs.output)) + + def test_missing_methods_are_rejected_without_installing_a_fallback(self): + self.require_api() + for subscribe in (subscribe_dequeue, subscribe_task_done): + with self.assertRaises(ValueError): + subscribe(object(), "key", lambda *a: None) diff --git a/tests/test_usage_tracking.py b/tests/test_usage_tracking.py index f7070fb..eeabad5 100644 --- a/tests/test_usage_tracking.py +++ b/tests/test_usage_tracking.py @@ -373,6 +373,22 @@ class TestCountAtDequeue(DbTestCase): self.assertEqual(self._runs(positive), 1) self.assertEqual(calls, []) + def test_usage_subscription_remains_once_after_an_unrelated_wrapper(self): + positive = self._save("a cat") + prompt_queue = self.server.prompt_queue + original = prompt_queue.get + + def unrelated(*args, **kwargs): + return original(*args, **kwargs) + + prompt_queue.get = unrelated + self.assertFalse( + _wrap_queue_get(prompt_queue, lambda _: self.fail("duplicate")) + ) + self.assertIs(prompt_queue.get, unrelated) + self._dequeue(self._submit(workflow())) + self.assertEqual(self._runs(positive), 1) + def test_wrap_rejects_queues_without_a_callable_get(self): self.assertFalse(_wrap_queue_get(object(), lambda prompt_id: None)) self.assertFalse(_wrap_queue_get(None, lambda prompt_id: None)) diff --git a/utils/execution_history.py b/utils/execution_history.py new file mode 100644 index 0000000..c05c07b --- /dev/null +++ b/utils/execution_history.py @@ -0,0 +1,145 @@ +"""Own bounded queue-attempt capture without retaining host workflow payloads.""" + +import atexit +import threading +import time +from typing import Callable +from uuid import uuid4 + +try: + from ..database import PromptDatabase +except ImportError: + from database import PromptDatabase + +from .execution_history_events import make_finish, make_start +from .execution_history_writer import HistoryWriter +from .logging_config import get_logger +from .queue_observers import subscribe_dequeue, subscribe_task_done + +_logger = get_logger("prompt_manager.execution_history") +_STATE_ATTRIBUTE = "_prompt_manager_execution_history" +_registration_lock = threading.Lock() +_MAX_ACTIVE = 128 + + +class HistoryRegistration: + """Own subscriptions, one asynchronous writer and bounded attempt identities.""" + + def __init__( + self, queue: object, session_id: str, db_factory: Callable[[], PromptDatabase] + ) -> None: + self.writer = HistoryWriter(db_factory) + self._session_id = session_id + self._lock = threading.Lock() + self._active: dict[int, str] = {} + self._disposed = False + self._unsubscribers: list[Callable[[], None]] = [] + try: + self._unsubscribers.append( + subscribe_dequeue(queue, "execution_history", self._dequeue) + ) + self._unsubscribers.append( + subscribe_task_done(queue, "execution_history", self._task_done) + ) + self.writer.start() + atexit.register(self.dispose) + except Exception: + self.dispose() + raise + + 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()): + self._active.clear() + return epoch + + def _dequeue(self, result: object) -> None: + try: + with self._lock: + if self._disposed: + return + epoch = self._ready_epoch() + if epoch is None: + 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 + if self.writer.submit(command): + self._active[item_id] = epoch + except Exception: + self.writer.record_gap() + + def _task_done(self, args: tuple, kwargs: dict) -> None: + try: + with self._lock: + if self._disposed: + return + epoch = self._ready_epoch() + if epoch is None: + self.writer.record_gap() + return + item_id = args[0] if args else kwargs.get("item_id") + status = args[2] if len(args) > 2 else kwargs.get("status") + command = make_finish( + self._session_id, + item_id, + epoch, + status, + time.time_ns() // 1_000_000, + ) + if command is None: + self.writer.record_gap() + return + self._active.pop(command.queue_item_id, None) + self.writer.submit(command) + except Exception: + self.writer.record_gap() + + def dispose(self, timeout: float = 2.0) -> bool: + """Stop capture, release subscriptions and drain the owned writer boundedly.""" + with self._lock: + if not self._disposed: + self._disposed = True + self._active.clear() + for unsubscribe in self._unsubscribers: + unsubscribe() + self._unsubscribers.clear() + atexit.unregister(self.dispose) + return self.writer.close(timeout) + + +def register_history_hook( + server: object, db_factory: Callable[[], PromptDatabase] +) -> HistoryRegistration | None: + """Register only on queues supporting actual dequeue and final callbacks.""" + try: + queue = getattr(server, "prompt_queue", None) + if not all( + callable(getattr(queue, method, None)) for method in ("get", "task_done") + ): + return None + with _registration_lock: + state = getattr(queue, _STATE_ATTRIBUTE, None) + if state is None: + state = {"session_id": str(uuid4()), "registration": None} + setattr(queue, _STATE_ATTRIBUTE, state) + previous = state["registration"] + if previous is not None and not previous._disposed: + return previous + registration = HistoryRegistration(queue, state["session_id"], db_factory) + state["registration"] = registration + return registration + except Exception: + _logger.warning("Execution history queue capability unavailable") + return None diff --git a/utils/execution_history_writer.py b/utils/execution_history_writer.py index 7c66549..3cba705 100644 --- a/utils/execution_history_writer.py +++ b/utils/execution_history_writer.py @@ -26,6 +26,7 @@ from .logging_config import get_logger _logger = get_logger("prompt_manager.execution_history_writer") _MAX_PENDING = 1024 _PROBE_INTERVAL = 5.0 +_MAX_GAPS = 2**63 - 1 @dataclass(frozen=True) @@ -75,7 +76,7 @@ class HistoryWriter: encode_command(command) except ValueError: with self._lock: - self._gaps += 1 + self._record_gap_locked() self._wake.set() return False with self._lock: @@ -84,7 +85,7 @@ class HistoryWriter: or command.storage_epoch != self._epoch or self._pending >= _MAX_PENDING ): - self._gaps += 1 + self._record_gap_locked() accepted = False else: self._queue.put_nowait(command) @@ -93,6 +94,15 @@ class HistoryWriter: self._wake.set() return accepted + def _record_gap_locked(self) -> None: + self._gaps = min(_MAX_GAPS, self._gaps + 1) + + def record_gap(self) -> None: + """Account for unavailable capture without constructing a fake command.""" + with self._lock: + self._record_gap_locked() + self._wake.set() + def invalidate(self) -> None: with self._lock: self._accepting = False @@ -134,7 +144,7 @@ class HistoryWriter: except queue.Empty: return self._pending -= 1 - self._gaps += 1 + self._record_gap_locked() def _run(self) -> None: db: PromptDatabase | None = None @@ -187,7 +197,7 @@ class HistoryWriter: if db is not None: db.model.close() with self._lock: - self._gaps += 1 + self._record_gap_locked() if command is not None: failed = False try: @@ -206,7 +216,7 @@ class HistoryWriter: with self._lock: self._pending -= 1 if failed: - self._gaps += 1 + self._record_gap_locked() self._accepting = False self._epoch = None self._generation += 1 diff --git a/utils/queue_observers.py b/utils/queue_observers.py new file mode 100644 index 0000000..e07af2a --- /dev/null +++ b/utils/queue_observers.py @@ -0,0 +1,110 @@ +"""Keyed, independently removable observers for successful queue operations.""" + +import functools +import threading +import time +from typing import Callable + +from .logging_config import get_logger + +_logger = get_logger("prompt_manager.queue_observers") +_STATE_ATTRIBUTE = "_prompt_manager_queue_observers" +_registration_lock = threading.Lock() + + +class _Observation: + def __init__(self, queue: object, method: str) -> None: + self.queue = queue + self.method = method + self.lock = threading.Lock() + self.callbacks: dict[str, tuple[object, Callable]] = {} + self.original = None + self.wrapper = None + self.installed = False + self.next_warning = 0.0 + + def install(self) -> None: + original = getattr(self.queue, self.method, None) + if not callable(original): + raise ValueError("Queue capability unavailable") + + @functools.wraps(original) + def observe(*args: object, **kwargs: object) -> object: + result = original(*args, **kwargs) + if self.method == "task_done" or result is not None: + with self.lock: + callbacks = tuple(entry[1] for entry in self.callbacks.values()) + for callback in callbacks: + try: + if self.method == "get": + callback(result) + else: + callback(args, kwargs) + except Exception: + with self.lock: + now = time.monotonic() + warn = now >= self.next_warning + if warn: + self.next_warning = now + 5.0 + if warn: + _logger.warning("Queue observation callback failed") + return result + + setattr(self.queue, self.method, observe) + self.original, self.wrapper, self.installed = original, observe, True + + def subscribe(self, key: str, callback: Callable) -> Callable[[], None]: + with self.lock: + if not self.installed: + self.install() + 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 getattr(self.queue, self.method) is self.wrapper + ): + setattr(self.queue, self.method, self.original) + self.installed = False + + return unsubscribe + + +def _subscribe( + queue: object, method: str, key: str, callback: Callable +) -> Callable[[], None]: + with _registration_lock: + states = getattr(queue, _STATE_ATTRIBUTE, None) + if states is None: + states = {} + try: + setattr(queue, _STATE_ATTRIBUTE, states) + except (AttributeError, TypeError) as error: + raise ValueError("Queue does not support observation") from error + state = states.get(method) + if state is None: + state = _Observation(queue, method) + states[method] = state + return state.subscribe(key, callback) + + +def subscribe_dequeue( + queue: object, key: str, callback: Callable[[object], None] +) -> Callable[[], None]: + """Observe each actual get result once per key, after the host returns.""" + return _subscribe(queue, "get", key, callback) + + +def subscribe_task_done( + queue: object, key: str, callback: Callable[[tuple, dict], None] +) -> Callable[[], None]: + """Observe successful task_done calls with their original argument objects.""" + return _subscribe(queue, "task_done", key, callback) diff --git a/utils/usage_tracking.py b/utils/usage_tracking.py index 2acbdf2..4f0b094 100644 --- a/utils/usage_tracking.py +++ b/utils/usage_tracking.py @@ -19,7 +19,6 @@ Text from known pure string nodes is resolved at queue time too. Anything else (PromptSearchList batches, unknown nodes) is counted by the node as it executes. """ -import functools import threading import time import uuid @@ -32,6 +31,7 @@ try: from .prompt_graph import intrinsic_pair_nodes, resolve_text, run_prompt_components from .prompt_pair_resolver import resolve_text_input from .pair_execution import clear_pair_execution + from .queue_observers import subscribe_dequeue except ImportError: from utils.hashing import generate_prompt_hash from utils.logging_config import get_logger @@ -42,6 +42,7 @@ except ImportError: ) from utils.prompt_pair_resolver import resolve_text_input from utils.pair_execution import clear_pair_execution + from utils.queue_observers import subscribe_dequeue logger = get_logger("prompt_manager.usage_tracking") @@ -294,21 +295,20 @@ def _wrap_queue_get(prompt_queue: Any, on_dequeued: Callable[[str], None]) -> bo original = getattr(prompt_queue, "get", None) if not callable(original): return False - if getattr(original, _WRAPPED_MARKER, False): + if getattr(prompt_queue, _WRAPPED_MARKER, False) or getattr( + original, _WRAPPED_MARKER, False + ): return False - @functools.wraps(original) - def get(*args, **kwargs): - result = original(*args, **kwargs) - if result is not None: - try: - on_dequeued(str(result[0][1])) - except Exception as e: - logger.warning(f"Dequeue hook failed: {e}") - return result + def report(result: object) -> None: + try: + on_dequeued(str(result[0][1])) + except Exception as e: + logger.warning(f"Dequeue hook failed: {e}") - setattr(get, _WRAPPED_MARKER, True) - prompt_queue.get = get + subscribe_dequeue(prompt_queue, "usage_tracking", report) + setattr(prompt_queue.get, _WRAPPED_MARKER, True) + setattr(prompt_queue, _WRAPPED_MARKER, True) return True