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

2026-10-17 15:40:12 -07:00
parent 246496d175
commit eb9800a898
9 changed files with 924 additions and 18 deletions
@@ -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)