iai-mcp-opencode/src/iai_mcp/provenance_queue.py

"""Plan 05-14 — async provenance write queue (OPS-10 / M-02).

Moves provenance writes off the recall critical path. A single daemon
thread drains a bounded queue.Queue of (record_id, entry) pairs and
flushes them via the existing ``MemoryStore.append_provenance_batch``
exactly as the sync path did.

Why this is the right shape:
- provenance writes are pure SIDE EFFECTS; pipeline_recall never reads
  their result. Textbook fire-and-forget candidate.
- The biological analogue: consolidation writes happen during rest, not
  during retrieval (CLS / sleep replay).
- The existing ``AsyncWriteQueue`` is for record inserts,
  which must be durable before their return (S4 viability check reads
  them back). Provenance has no such contract — a simpler, purpose-built
  queue avoids the coroutine/event-loop machinery that asyncio imposes.

Constitutional fences:
- Rule 1: worker swallows all exceptions (recall must never fail due
  to a provenance-write failure).
- entries are never dropped during normal operation; on shutdown
  the atexit hook drains the queue. W1/when the
  in-memory queue is full under overload, batches are spilled to
  ``~/.iai-mcp/.provenance-overflow/<unix_ms>-<n>.jsonl``. The worker
  drains the spill dir on idle and re-enqueues the batches. Zero drops
  on the happy path; the only path that can drop is disk-write failure
  (alarmed via the ``provenance_queue_spill_failed`` stderr event).
- C3 / C6: stdlib only. No extra dependencies.

Python 3.11+.
"""
from __future__ import annotations

import atexit
import json
import queue
import sys
import threading
import time
from pathlib import Path
from typing import TYPE_CHECKING
from uuid import UUID

if TYPE_CHECKING:
    from iai_mcp.store import MemoryStore


# Sentinel pushed on the queue to wake the worker for stop/flush.
_STOP = object()
_FLUSH = object()

# W1/D-01 — overflow spill-to-disk.
OVERFLOW_DIR_NAME = ".provenance-overflow"
# Worker idle poll: 5s upper bound on overflow-drain responsiveness.
# Bounded so under sustained overload the spill drain catches up
# within a small constant time after _q clears.
_WORKER_IDLE_POLL_S = 5.0


class ProvenanceWriteQueue:
    """Single-daemon-thread coalescing queue for provenance batches.

    Usage:
        q = ProvenanceWriteQueue(store, coalesce_ms=50)
        q.start()                                # idempotent
        q.enqueue([(record_id, entry_dict), ...])  # non-blocking
        q.flush(timeout=2.0)                     # drain + wait
        q.stop()                                 # drain + join

    The worker loop:
        1. Blocking `.get()` on the queue (wakes on enqueue or sentinel).
        2. Opportunistic drain up to ``max_batch_pairs`` pairs OR until
           the queue has been empty for ``coalesce_ms``.
        3. Single call to ``store.append_provenance_batch(pairs,
           records_cache=None)``.
        4. Back to (1).

    All worker exceptions are logged to stderr as structured JSON events
    and swallowed.
    """

    def __init__(
        self,
        store: "MemoryStore",
        *,
        coalesce_ms: int = 50,
        max_queue_size: int = 4096,
        max_batch_pairs: int = 256,
    ) -> None:
        self._store = store
        self._coalesce_s = max(1, int(coalesce_ms)) / 1000.0
        self._max_batch = int(max_batch_pairs)
        # Queue items are either lists of (UUID, dict) pairs or the
        # _STOP / _FLUSH sentinels.
        self._q: queue.Queue = queue.Queue(maxsize=int(max_queue_size))
        self._thread: threading.Thread | None = None
        self._started = False
        self._stop_requested = False
        # flush synchronisation: drained_event is set by the worker when
        # it has processed everything up to a _FLUSH sentinel.
        self._flush_event = threading.Event()
        self._atexit_registered = False
        self._lock = threading.Lock()

    # ------------------------------------------------------------------ lifecycle

    def start(self) -> None:
        """Start the worker thread. Idempotent."""
        with self._lock:
            if self._started:
                return
            self._started = True
            self._stop_requested = False
            self._thread = threading.Thread(
                target=self._run,
                name="iai-mcp-provenance-queue",
                daemon=True,
            )
            self._thread.start()
            if not self._atexit_registered:
                atexit.register(self._atexit_flush)
                self._atexit_registered = True

    def stop(self) -> None:
        """Signal the worker, drain remaining items, join the thread.

        Idempotent. After stop the queue is no longer usable; call
        start() to revive (fresh worker, same queue instance).
        """
        with self._lock:
            if not self._started:
                return
            self._stop_requested = True
            try:
                self._q.put_nowait(_STOP)
            except queue.Full:
                # Drop one item to make room for the sentinel.
                try:
                    self._q.get_nowait()
                    self._q.put_nowait(_STOP)
                except queue.Empty:
                    pass
            t = self._thread
        if t is not None:
            t.join(timeout=5.0)
        with self._lock:
            self._started = False
            self._thread = None

    def flush(self, timeout: float = 2.0) -> None:
        """Wait until the worker has drained everything enqueued so far.

        Puts a _FLUSH sentinel; the worker signals _flush_event once it
        has processed all pairs that were in the queue at that point.
        Times out silently — the caller is responsible for deciding
        whether to retry; recall latency is never blocked by flush().
        """
        if not self._started:
            return
        self._flush_event.clear()
        try:
            self._q.put(_FLUSH, timeout=timeout)
        except queue.Full:
            return
        self._flush_event.wait(timeout=timeout)

    # ---------------------------------------------------------------- public write

    def enqueue(self, pairs: "list[tuple[UUID, dict]]") -> None:
        """Non-blocking enqueue.

        W1/when the in-memory queue is full, the batch
        spills to ``~/.iai-mcp/.provenance-overflow/<ts>-<n>.jsonl``.
        The worker thread drains the spill dir on idle and re-enqueues
        the batches. zero drops under overload (only path that
        can drop is disk-write failure, which is itself alarmed).
        """
        if not pairs:
            return
        try:
            self._q.put_nowait(list(pairs))
            return
        except queue.Full:
            pass
        # In-memory queue full — spill to disk. Worker will pick this
        # up on its next idle cycle. Recall hot path is unaffected
        # (this branch only fires on the WRITE side under overload).
        self._spill_to_disk(list(pairs))
        try:
            sys.stderr.write(
                '{"event":"provenance_queue_overflow_spill","n_pairs":'
                + str(len(pairs))
                + "}\n"
            )
        except Exception:
            pass

    # ---------------------------------------------------------------- spill / drain

    def _spill_to_disk(self, pairs: list) -> None:
        """Persist a rejected batch to ``~/.iai-mcp/.provenance-overflow/``.

        Per-batch JSONL file: one line per (uuid_str, entry_dict) pair.
        File-level atomicity — the worker re-enqueues the entire file's
        contents in one call, then unlinks. Format:

            {"id": "<uuid>", "entry": {...}}\n
            {"id": "<uuid>", "entry": {...}}\n

        Failure modes:
        - Disk full / permission denied: emits structured stderr event
          ``provenance_queue_spill_failed``. This is the ONLY drop path
          remaining post-07.9 W1; it's a system-level alarm condition,
          not a normal-operation outcome.
        """
        if not pairs:
            return
        try:
            overflow_dir = Path.home() / ".iai-mcp" / OVERFLOW_DIR_NAME
            overflow_dir.mkdir(parents=True, exist_ok=True)
            ts_ms = int(time.time() * 1000)
            # Tag with the batch length and a short pid suffix so two
            # spills inside the same millisecond never collide.
            fpath = overflow_dir / f"{ts_ms}-{len(pairs)}-{id(pairs) & 0xFFFF:04x}.jsonl"
            tmp_path = fpath.with_suffix(fpath.suffix + ".tmp")
            with tmp_path.open("w", encoding="utf-8") as fh:
                for rid, entry in pairs:
                    fh.write(json.dumps({"id": str(rid), "entry": entry}) + "\n")
            tmp_path.rename(fpath)  # atomic rename keeps drain from
            # ever reading a half-written file.
        except Exception as exc:
            try:
                sys.stderr.write(
                    '{"event":"provenance_queue_spill_failed","error":'
                    + _json_str(str(exc))
                    + ',"n_pairs":' + str(len(pairs)) + '}\n'
                )
            except Exception:
                pass

    def _drain_overflow_dir(self) -> int:
        """Re-enqueue any spilled batches into ``_q``.

        Called by the worker on idle (between blocking `_q.get()` cycles).
        Per-file atomicity: re-enqueue ALL pairs from a file via a single
        ``_q.put`` call, then unlink. If ``_q`` is still full, leave the
        file on disk for the next idle cycle.

        Returns the number of pairs successfully re-enqueued in this pass.
        """
        overflow_dir = Path.home() / ".iai-mcp" / OVERFLOW_DIR_NAME
        if not overflow_dir.exists():
            return 0
        n_re_enqueued = 0
        # sorted() so older spill files drain first (FIFO durability).
        for fpath in sorted(overflow_dir.glob("*.jsonl")):
            try:
                pairs: list = []
                with fpath.open(encoding="utf-8") as fh:
                    for line in fh:
                        line = line.strip()
                        if not line:
                            continue
                        obj = json.loads(line)
                        pairs.append((UUID(obj["id"]), obj["entry"]))
                if not pairs:
                    fpath.unlink()
                    continue
                # Short-timeout put: this is the worker thread, so
                # blocking briefly is fine, but a long block would
                # delay normal-path enqueues that arrive during drain.
                try:
                    self._q.put(pairs, timeout=0.5)
                except queue.Full:
                    # Queue still saturated — leave the file for the
                    # next idle cycle. Don't unlink.
                    return n_re_enqueued
                fpath.unlink()
                n_re_enqueued += len(pairs)
            except Exception as exc:
                # Malformed spill file: preserve evidence, do not lose data.
                try:
                    failed = fpath.with_suffix(f".failed-{int(time.time())}.jsonl")
                    fpath.rename(failed)
                    sys.stderr.write(
                        '{"event":"provenance_queue_spill_drain_failed","error":'
                        + _json_str(str(exc)) + '}\n'
                    )
                except Exception:
                    pass
        return n_re_enqueued

    # ------------------------------------------------------------------ internals

    def _run(self) -> None:
        """Worker loop.

        W1/between blocking `_q.get()` cycles the worker
        drains any spilled overflow files at ``~/.iai-mcp/.provenance-overflow/``.
        Bounded poll: idle-timeout = ``_WORKER_IDLE_POLL_S`` so the spill
        drain runs at most once per ``_WORKER_IDLE_POLL_S`` seconds when
        the queue is empty.
        """
        while True:
            try:
                item = self._q.get(timeout=_WORKER_IDLE_POLL_S)
            except queue.Empty:
                # Idle tick — try to drain the overflow dir back into _q.
                # Defensive: any error during drain is logged + swallowed.
                try:
                    self._drain_overflow_dir()
                except Exception:
                    pass
                continue
            except Exception:
                continue
            if item is _STOP:
                # Drain remaining real items before exit.
                self._drain_remaining()
                return
            if item is _FLUSH:
                # Drain everything enqueued before this sentinel.
                self._drain_remaining()
                self._flush_event.set()
                continue
            # Normal batch. Coalesce: pull more pending items until we
            # hit max_batch_pairs or a short idle window.
            pairs: list = list(item)
            while len(pairs) < self._max_batch:
                try:
                    nxt = self._q.get(timeout=self._coalesce_s)
                except queue.Empty:
                    break
                if nxt is _STOP:
                    # Flush what we have, then exit.
                    self._flush_batch(pairs)
                    self._drain_remaining()
                    return
                if nxt is _FLUSH:
                    self._flush_batch(pairs)
                    self._drain_remaining()
                    self._flush_event.set()
                    pairs = []
                    break
                pairs.extend(nxt)
            if pairs:
                self._flush_batch(pairs)

    def _drain_remaining(self) -> None:
        """Pull everything currently in the queue and flush as one batch."""
        pairs: list = []
        saw_flush = False
        while True:
            try:
                item = self._q.get_nowait()
            except queue.Empty:
                break
            if item is _STOP:
                continue
            if item is _FLUSH:
                saw_flush = True
                continue
            pairs.extend(item)
        if pairs:
            self._flush_batch(pairs)
        if saw_flush:
            self._flush_event.set()

    def _flush_batch(self, pairs: list) -> None:
        """Call store.append_provenance_batch, swallow all exceptions (Rule 1)."""
        if not pairs:
            return
        try:
            self._store.append_provenance_batch(pairs, records_cache=None)
        except Exception as exc:
            try:
                sys.stderr.write(
                    '{"event":"provenance_queue_flush_failed","n_pairs":'
                    + str(len(pairs))
                    + ',"error":'
                    + _json_str(str(exc))
                    + "}\n"
                )
            except Exception:
                pass

    def _atexit_flush(self) -> None:
        """atexit handler — drain and stop the worker. Never raises."""
        try:
            if self._started:
                self.flush(timeout=2.0)
                self.stop()
        except Exception:
            pass


def _json_str(s: str) -> str:
    """Minimal JSON string escape for stderr structured logs."""
    return '"' + s.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n") + '"'