"""Thread-safe AI logging helpers for use from thread pool workers. Provides start_entry() / finish_entry() that schedule log operations on the event loop via call_soon_threadsafe, making them safe to call from executor threads. Also provides a ContextVar so plugin/entity context flows through asyncio.run_in_executor() calls automatically. Initialized by logic/ai_log.py at app startup via set_app_loop(). Importable by both logic/ and plugins/ without circular dependencies. """ import concurrent.futures import time from collections.abc import Callable from contextvars import ContextVar from dataclasses import dataclass from typing import Any try: import asyncio as _asyncio _AbstractEventLoop = _asyncio.AbstractEventLoop except ImportError: # pragma: no cover _AbstractEventLoop = Any # type: ignore[assignment,misc] import asyncio @dataclass class _LogCtx: plugin_id: str entity_type: str entity_id: str # ContextVar propagated automatically into asyncio executor threads. _ctx: ContextVar[_LogCtx | None] = ContextVar("_log_ctx", default=None) # Initialized at startup by set_app_loop(). _loop: asyncio.AbstractEventLoop | None = None _log_start_fn: Callable[..., str] | None = None _log_finish_fn: Callable[..., None] | None = None def set_app_loop( loop: asyncio.AbstractEventLoop, log_start: Callable[..., str], log_finish: Callable[..., None], ) -> None: """Store the running event loop and logging callables. Must be called once at app startup from the async context. Args: loop: The running asyncio event loop. log_start: Synchronous log_start function from logic.ai_log. log_finish: Synchronous log_finish function from logic.ai_log. """ global _loop, _log_start_fn, _log_finish_fn _loop = loop _log_start_fn = log_start _log_finish_fn = log_finish def set_log_ctx(plugin_id: str, entity_type: str, entity_id: str) -> None: """Set the current log context for this thread/task. Call before run_in_executor() to propagate context into executor threads. Or call directly inside a thread to set context for subsequent calls in the same thread. Args: plugin_id: Plugin ID to attribute log entries to. entity_type: Entity type (e.g. ``"books"``). entity_id: Entity ID. """ _ctx.set(_LogCtx(plugin_id=plugin_id, entity_type=entity_type, entity_id=entity_id)) def start_entry(model: str, request_summary: str) -> str: """Start a log entry from a thread pool worker. Reads context from the ContextVar set by set_log_ctx(). Schedules log_start on the event loop and blocks briefly to obtain the entry ID. Returns empty string if context or loop is unavailable. Args: model: Model name used for the request. request_summary: Short human-readable description. Returns: Log entry ID string, or ``""`` if logging is unavailable. """ ctx = _ctx.get() if ctx is None or _loop is None or _log_start_fn is None: return "" fut: concurrent.futures.Future[str] = concurrent.futures.Future() fn = _log_start_fn pid, et, eid = ctx.plugin_id, ctx.entity_type, ctx.entity_id def _call() -> None: try: entry_id = fn(pid, et, eid, model, request_summary) fut.set_result(entry_id) except Exception as exc: # noqa: BLE001 fut.set_exception(exc) _loop.call_soon_threadsafe(_call) try: return fut.result(timeout=5) except Exception: return "" def finish_entry(entry_id: str, status: str, response: str, started_at: float) -> None: """Finish a log entry from a thread pool worker (fire-and-forget). Schedules log_finish on the event loop. Does nothing if entry_id is empty or the loop is unavailable. Args: entry_id: ID returned by start_entry(). status: ``"ok"`` or ``"error"``. response: Short summary of response or error message. started_at: ``time.time()`` value recorded before the request. """ if not entry_id or _loop is None or _log_finish_fn is None: return fn = _log_finish_fn _loop.call_soon_threadsafe(fn, entry_id, status, response, started_at) def timed_start(model: str, request_summary: str) -> tuple[str, float]: """Convenience wrapper: start an entry and record the start time. Returns: Tuple of (entry_id, started_at) for passing to finish_entry(). """ started_at = time.time() entry_id = start_entry(model, request_summary) return entry_id, started_at