feat(history): observe actual queue attempts without altering execution (#209)

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

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

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)

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