cleveragents-core/tools/_block_store.py

"""SQLite-backed cross-process content block store.

Built for the pre-R3 chain of agents
(dispatcher → tier-dispatcher → task-implementor / pr-review-worker)
where each ``task`` tool hop summarised long context and pre-fetched
data sections (PR diff, comments, CI failure logs, reviews) bore the
brunt — a 36 KB section could shrink to 12 KB by depth-3, losing whole
sub-sections.

After the R3 wrapper-chain retirement (2026-05-17) the implementer
pool dispatches directly to ``task-implementor-tier-<slot>`` variants
with no intervening ``task`` hops, so the prefetched body now
generally survives intact. The block store is still load-bearing for
(a) ``pr-review-worker``, which is dispatched directly but whose
embedded prompt can still exceed practical size budgets, and (b) any
future workflow that re-introduces summarisation between dispatch
and worker.

This store lets the dispatcher externalise large sections into
durable rows and reference them in the prompt by short block keys
(e.g. ``pr-30-diff-deadbeefcafe``). Intermediate agents preserve the
keys; the worker re-fetches original content via the ``block_store``
MCP when the inline copy is no longer enough.

Storage: SQLite at ``/tmp/cleveragents-block-store/blocks.sqlite3``
(``BLOCK_STORE_DIR`` overrides). WAL mode so the dispatcher writer
and MCP-server reader coexist without lock contention. Content is
plain UTF-8 — callers JSON-encode their own payload.

Sizing: per-row cap 1 MB (covers every section observed; over-cap
inserts are a producer bug). No file-level cap; the dispatcher
startup janitor (:func:`janitor`) purges expired rows so the file
stays bounded by recent-PR working set.

TTL: default 1 hour (covers a full PR cycle with margin); pass
``ttl_s`` to :func:`register` for longer-lived rows.

Errors: operational errors raise :class:`BlockStoreError`. SQLite
I/O errors propagate as ``sqlite3.DatabaseError`` so callers can
distinguish "this row is bad" from "the store itself is broken".
"""

from __future__ import annotations

import datetime as _dt
import logging
import os
import re
import sqlite3
import threading
import time
from pathlib import Path
from typing import Any

_logger = logging.getLogger("block_store")

SCHEMA_VERSION = 1

# Module-level monotonic timestamp of the last janitor sweep, guarded
# by ``_LAST_JANITOR_LOCK`` so concurrent threads inside one process
# can't both clear the gate and double-sweep. Across processes the
# state is per-process — the two will race and both sweep, which is
# safe (SQLite DELETE is idempotent and the second pass is a no-op
# once the first commits).
_LAST_JANITOR_AT: float | None = None
_LAST_JANITOR_LOCK = threading.Lock()
_PERIODIC_JANITOR_INTERVAL_S = 3600  # 1 hour
_PERIODIC_JANITOR_INTERVAL_ENV = "BLOCK_STORE_PERIODIC_JANITOR_INTERVAL_S"

# ─── Defaults + env overrides ──────────────────────────────────────

_DEFAULT_STORE_DIR = Path("/tmp/cleveragents-block-store")
_STORE_DIR_ENV = "BLOCK_STORE_DIR"
_DISABLE_ENV = "BLOCK_STORE_DISABLE"

# Per-row content cap. 1 MB is large enough for any single PR's
# diff / comments / CI log feed we've observed, and well under any
# practical SQLite per-row ceiling. Operators can dial this UP for
# unusually large PRs but doing so should be a deliberate decision.
_DEFAULT_MAX_CONTENT_BYTES = 1_000_000
_MAX_CONTENT_BYTES_ENV = "BLOCK_STORE_MAX_CONTENT_BYTES"

# Default TTL — 1 hour covers a full review/implementer cycle with
# margin. Callers pass an explicit ``ttl_s`` to override (e.g. a
# slow tier-2 escalation that may run for several hours).
_DEFAULT_TTL_S = 3600
_DEFAULT_TTL_ENV = "BLOCK_STORE_DEFAULT_TTL_S"

# Key character whitelist. Keys flow through MCP -> agent prompt ->
# back to MCP, so they must be safe in shell args, JSON, and
# Markdown. Restrict to [A-Za-z0-9_:.-]. The pattern includes ``:``
# so callers can adopt namespaced keys like ``pr:30:diff:deadbeef``.
_KEY_RE = re.compile(r"^[A-Za-z0-9_:.\-]{1,256}$")


class BlockStoreError(RuntimeError):
    """Operational error from the block store. Distinct from
    ``sqlite3.*`` errors so callers can catch "expected, recover" cases
    (missing key, bad input) without swallowing genuine I/O failures."""


# ─── Helpers ────────────────────────────────────────────────────────


def store_dir() -> Path:
    return Path(os.environ.get(_STORE_DIR_ENV) or str(_DEFAULT_STORE_DIR))


def store_path() -> Path:
    return store_dir() / "blocks.sqlite3"


def is_disabled() -> bool:
    raw = os.environ.get(_DISABLE_ENV, "").strip().lower()
    return raw in {"1", "true", "yes", "on"}


def _max_content_bytes() -> int:
    raw = os.environ.get(_MAX_CONTENT_BYTES_ENV)
    if raw:
        try:
            return max(1024, int(raw))
        except ValueError:
            pass
    return _DEFAULT_MAX_CONTENT_BYTES


def _default_ttl_s() -> int:
    raw = os.environ.get(_DEFAULT_TTL_ENV)
    if raw:
        try:
            return max(1, int(raw))
        except ValueError:
            pass
    return _DEFAULT_TTL_S


def _periodic_janitor_interval_s() -> int:
    raw = os.environ.get(_PERIODIC_JANITOR_INTERVAL_ENV)
    if raw:
        try:
            return max(60, int(raw))
        except ValueError:
            pass
    return _PERIODIC_JANITOR_INTERVAL_S


def _maybe_sweep_periodically() -> None:
    """Run :func:`janitor` if the last sweep was longer than
    ``_PERIODIC_JANITOR_INTERVAL_S`` ago. Per-process state — restart
    or first call always triggers a sweep on the next register.

    The gate check + stamp are guarded by ``_LAST_JANITOR_LOCK`` so
    concurrent threads inside one process see at most one sweep per
    interval. Cross-process races stay possible and are intentionally
    safe (idempotent SQLite DELETE).

    Self-contained so a long-running dispatcher doesn't depend on
    external scheduling to keep the SQLite file bounded between
    startup janitor passes."""
    global _LAST_JANITOR_AT
    now = time.monotonic()
    interval = _periodic_janitor_interval_s()
    with _LAST_JANITOR_LOCK:
        if _LAST_JANITOR_AT is not None and (now - _LAST_JANITOR_AT) < interval:
            return
        _LAST_JANITOR_AT = now
    try:
        janitor()
    except (OSError, sqlite3.DatabaseError):
        # Best-effort; never break a register() because the sweep
        # had a hiccup. janitor() already logs.
        pass


def _now() -> _dt.datetime:
    return _dt.datetime.now(_dt.timezone.utc)


def _validate_key(key: str) -> str:
    if not isinstance(key, str):
        raise BlockStoreError(f"key must be a string, got {type(key).__name__}")
    if not _KEY_RE.match(key):
        raise BlockStoreError(
            f"key {key!r} is invalid: must match [A-Za-z0-9_:.-]{{1,256}}"
        )
    return key


def _connect() -> sqlite3.Connection:
    """Open (and lazily initialise) the SQLite database with WAL on
    for concurrent reader/writer access from dispatcher + MCP server."""
    target = store_path()
    target.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(
        str(target),
        timeout=10.0,
        isolation_level=None,  # autocommit; explicit BEGIN/COMMIT in writes
    )
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode = WAL")
    conn.execute("PRAGMA synchronous = NORMAL")
    conn.execute("PRAGMA foreign_keys = ON")
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS blocks (
            key             TEXT PRIMARY KEY,
            content         TEXT NOT NULL,
            content_bytes   INTEGER NOT NULL,
            content_type    TEXT NOT NULL DEFAULT 'text/plain',
            source          TEXT NOT NULL DEFAULT 'dispatcher',
            created_at      TEXT NOT NULL,
            expires_at      TEXT NOT NULL,
            pr_number       INTEGER,
            block_type      TEXT
        )
        """
    )
    conn.execute(
        "CREATE INDEX IF NOT EXISTS idx_blocks_expires_at ON blocks(expires_at)"
    )
    conn.execute("CREATE INDEX IF NOT EXISTS idx_blocks_pr_number ON blocks(pr_number)")
    return conn


# ─── Public API ─────────────────────────────────────────────────────


def register(
    key: str,
    content: str,
    *,
    ttl_s: int | None = None,
    source: str = "dispatcher",
    content_type: str = "text/plain",
    pr_number: int | None = None,
    block_type: str | None = None,
) -> dict[str, Any]:
    """Insert or replace a block.

    Returns ``{"key", "bytes", "expires_at"}`` on success. Raises
    :class:`BlockStoreError` on validation failure, sqlite3 errors
    on I/O failure.

    ``ttl_s`` defaults to ``BLOCK_STORE_DEFAULT_TTL_S`` (1h).
    ``content`` is stored verbatim; callers serialise structured
    data themselves so the store stays format-agnostic. The
    optional ``pr_number`` and ``block_type`` columns are indexed
    so :func:`list_keys` can filter by PR cheaply for telemetry /
    debugging.
    """
    if is_disabled():
        raise BlockStoreError("block store is disabled via BLOCK_STORE_DISABLE")
    _maybe_sweep_periodically()
    _validate_key(key)
    if not isinstance(content, str):
        raise BlockStoreError(f"content must be a string, got {type(content).__name__}")
    encoded = content.encode("utf-8")
    cap = _max_content_bytes()
    if len(encoded) > cap:
        raise BlockStoreError(
            f"content for {key!r} is {len(encoded)} bytes, exceeds cap {cap}"
        )
    effective_ttl = ttl_s if ttl_s is not None else _default_ttl_s()
    if effective_ttl <= 0:
        raise BlockStoreError(f"ttl_s must be positive, got {effective_ttl}")
    now = _now()
    expires = now + _dt.timedelta(seconds=int(effective_ttl))
    conn = _connect()
    try:
        with conn:
            conn.execute(
                """
                INSERT INTO blocks (
                    key, content, content_bytes, content_type,
                    source, created_at, expires_at, pr_number, block_type
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                ON CONFLICT(key) DO UPDATE SET
                    content       = excluded.content,
                    content_bytes = excluded.content_bytes,
                    content_type  = excluded.content_type,
                    source        = excluded.source,
                    created_at    = excluded.created_at,
                    expires_at    = excluded.expires_at,
                    pr_number     = excluded.pr_number,
                    block_type    = excluded.block_type
                """,
                (
                    key,
                    content,
                    len(encoded),
                    content_type,
                    source,
                    now.isoformat(),
                    expires.isoformat(),
                    pr_number,
                    block_type,
                ),
            )
    finally:
        conn.close()
    return {
        "key": key,
        "bytes": len(encoded),
        "expires_at": expires.isoformat(),
    }


def fetch(key: str) -> dict[str, Any]:
    """Return a block by key, or raise :class:`BlockStoreError` if
    missing / expired.

    Returns ``{"key", "content", "bytes", "content_type", "source",
    "created_at", "expires_at", "pr_number", "block_type"}``.
    """
    if is_disabled():
        raise BlockStoreError("block store is disabled via BLOCK_STORE_DISABLE")
    _validate_key(key)
    conn = _connect()
    try:
        row = conn.execute("SELECT * FROM blocks WHERE key = ?", (key,)).fetchone()
    finally:
        conn.close()
    if row is None:
        raise BlockStoreError(f"no block with key {key!r}")
    expires_at = row["expires_at"]
    try:
        expires_dt = _dt.datetime.fromisoformat(expires_at)
    except (TypeError, ValueError) as exc:
        raise BlockStoreError(
            f"block {key!r} has malformed expires_at {expires_at!r}"
        ) from exc
    if expires_dt <= _now():
        # Expired rows are not returned; the next janitor pass cleans
        # them up. Don't delete inline here — keeps fetch a pure read.
        raise BlockStoreError(
            f"block {key!r} expired at {expires_at}; re-register from source"
        )
    return {
        "key": row["key"],
        "content": row["content"],
        "bytes": row["content_bytes"],
        "content_type": row["content_type"],
        "source": row["source"],
        "created_at": row["created_at"],
        "expires_at": row["expires_at"],
        "pr_number": row["pr_number"],
        "block_type": row["block_type"],
    }


def list_keys(
    *,
    prefix: str | None = None,
    pr_number: int | None = None,
    include_expired: bool = False,
    limit: int = 200,
) -> list[dict[str, Any]]:
    """List block summaries (without content) for debugging / probing.

    Returns ``[{"key", "bytes", "expires_at", "source", "block_type",
    "pr_number"}]`` ordered by ``created_at`` DESC. ``limit`` caps
    the result count so a debugger probing a many-PR store doesn't
    pull megabytes of rows.
    """
    if is_disabled():
        raise BlockStoreError("block store is disabled via BLOCK_STORE_DISABLE")
    if limit <= 0 or limit > 1000:
        raise BlockStoreError(f"limit must be 1..1000, got {limit}")
    where: list[str] = []
    args: list[Any] = []
    if prefix is not None:
        if not isinstance(prefix, str):
            raise BlockStoreError(
                f"prefix must be a string, got {type(prefix).__name__}"
            )
        where.append("key LIKE ? ESCAPE '\\'")
        # Escape LIKE wildcards in the operator-supplied prefix so a
        # literal "%" doesn't expand into a wildcard match.
        safe_prefix = (
            prefix.replace("\\", "\\\\").replace("%", r"\%").replace("_", r"\_")
        )
        args.append(safe_prefix + "%")
    if pr_number is not None:
        try:
            args.append(int(pr_number))
        except (TypeError, ValueError) as exc:
            raise BlockStoreError(
                f"pr_number must be an int, got {pr_number!r}"
            ) from exc
        where.append("pr_number = ?")
    if not include_expired:
        where.append("expires_at > ?")
        args.append(_now().isoformat())
    where_sql = (" WHERE " + " AND ".join(where)) if where else ""
    sql = (
        "SELECT key, content_bytes, expires_at, source, block_type, pr_number "
        f"FROM blocks{where_sql} ORDER BY created_at DESC LIMIT ?"
    )
    args.append(limit)
    conn = _connect()
    try:
        rows = conn.execute(sql, args).fetchall()
    finally:
        conn.close()
    return [
        {
            "key": r["key"],
            "bytes": r["content_bytes"],
            "expires_at": r["expires_at"],
            "source": r["source"],
            "block_type": r["block_type"],
            "pr_number": r["pr_number"],
        }
        for r in rows
    ]


def invalidate(key: str) -> bool:
    """Remove a block by key. Returns ``True`` if a row was deleted,
    ``False`` if no such key existed. Idempotent."""
    if is_disabled():
        raise BlockStoreError("block store is disabled via BLOCK_STORE_DISABLE")
    _validate_key(key)
    conn = _connect()
    try:
        with conn:
            cur = conn.execute("DELETE FROM blocks WHERE key = ?", (key,))
            return cur.rowcount > 0
    finally:
        conn.close()


def janitor() -> int:
    """Delete every expired row. Returns the number of rows removed.

    Idempotent and best-effort — any exception is logged and swallowed
    so a dispatcher startup never aborts because the block store had
    a transient I/O hiccup (unwritable dir, locked DB, etc.). The
    catch is intentionally broad: this is "cleanup if you can",
    not "fail loudly if you can't"."""
    if is_disabled():
        return 0
    conn: sqlite3.Connection | None = None
    try:
        conn = _connect()
        with conn:
            cur = conn.execute(
                "DELETE FROM blocks WHERE expires_at <= ?",
                (_now().isoformat(),),
            )
            return int(cur.rowcount or 0)
    except (OSError, sqlite3.DatabaseError) as exc:
        _logger.warning("block store janitor swallowed error: %s", exc)
        return 0
    finally:
        if conn is not None:
            try:
                conn.close()
            except sqlite3.DatabaseError:
                pass


# ─── Convenience helpers for callers building keys ──────────────────


def sanitise_sha_suffix(head_sha: str | None) -> str:
    """Strip ``head_sha`` to hex-lowercase, cap at 12 chars. Returns
    ``""`` for None / non-hex inputs. Single source of truth so the
    key-builders in this module and :mod:`_block_prompt` produce
    matching suffixes."""
    if not head_sha:
        return ""
    return re.sub(r"[^a-f0-9]", "", str(head_sha).lower())[:12]


def make_pr_block_key(
    pr_number: int,
    block_type: str,
    head_sha: str | None = None,
) -> str:
    """Canonical key shape: ``pr-{N}-{block_type}[-{head_sha[:12]}]``.

    The optional ``head_sha`` suffix lets callers scope a block to a
    specific commit (so a force-push invalidates the old key naturally
    when the new dispatcher cycle registers under the new SHA)."""
    if not isinstance(pr_number, int) or pr_number <= 0:
        raise BlockStoreError(f"pr_number must be positive int, got {pr_number!r}")
    if not block_type or not re.match(r"^[a-z0-9_]+$", block_type):
        raise BlockStoreError(f"block_type must match [a-z0-9_]+, got {block_type!r}")
    base = f"pr-{pr_number}-{block_type}"
    safe_sha = sanitise_sha_suffix(head_sha)
    if safe_sha:
        base = f"{base}-{safe_sha}"
    _validate_key(base)
    return base


__all__ = (
    "BlockStoreError",
    "SCHEMA_VERSION",
    "fetch",
    "invalidate",
    "is_disabled",
    "janitor",
    "list_keys",
    "sanitise_sha_suffix",
    "make_pr_block_key",
    "register",
    "store_dir",
    "store_path",
)