146 lines
5.2 KiB
Python
146 lines
5.2 KiB
Python
"""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
|