bookshelf/src/log_thread.py

"""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