src/ai_agent/utils/shutdown.py

# utils/shutdown.py
"""Periodic cleanup and shutdown hooks.

Two mechanisms keep the app tidy at runtime and on exit:

**Background cleanup thread** (started by :func:`register`):
- Sweeps expired cache rows every ``CLEANUP_INTERVAL_SECONDS`` (default 7200).
- Purges old log files on the same interval (files older than
  ``LOG_RETENTION_DAYS`` days, default 7).

**atexit hook** (also registered by :func:`register`):
- Runs a final cache sweep, then VACUUM and closes the connection cleanly
  (triggers a WAL checkpoint).  Handles the case where the process exits
  before the next background interval fires.

Call :func:`register` once at startup (see ``cli.py``).
"""

from __future__ import annotations

import atexit
import logging
import os
import threading
import time
from pathlib import Path

log = logging.getLogger("ai_agent.shutdown")

LOG_RETENTION_DAYS: int = int(os.getenv("LOG_RETENTION_DAYS", "7"))
CLEANUP_INTERVAL_SECONDS: int = int(os.getenv("CLEANUP_INTERVAL_SECONDS", "7200"))


# ---------------------------------------------------------------------------
# Cache DB helpers
# ---------------------------------------------------------------------------

def _sweep_cache_db() -> None:
    """Delete expired rows from the cache DB (lightweight, runs periodically)."""
    from ai_agent.utils.cache_db import get_cache_db_or_none  # noqa: PLC0415

    db = get_cache_db_or_none()
    if db is None:
        return
    try:
        deleted = db.sweep_expired()
        if deleted:
            log.debug("Cache sweep: removed %d expired row(s).", deleted)
    except Exception:
        log.exception("Periodic cache sweep failed.")


def _vacuum_and_close_cache_db() -> None:
    """Final shutdown: VACUUM and close the cache DB (runs via atexit)."""
    from ai_agent.utils.cache_db import get_cache_db_or_none  # noqa: PLC0415

    db = get_cache_db_or_none()
    if db is None:
        return
    try:
        db.vacuum_and_close()
        log.info("Cache DB shutdown: VACUUM complete.")
    except Exception:
        log.exception("Cache DB shutdown cleanup failed.")


# ---------------------------------------------------------------------------
# Log file rotation helper
# ---------------------------------------------------------------------------

def _purge_old_logs() -> None:
    """Delete log files older than LOG_RETENTION_DAYS inside LOG_DIR.

    Age is determined by each file's modification time (``st_mtime``), which
    reflects when data was last written.  Parsing the date from the filename
    prefix is deliberately avoided: ``TimedRotatingFileHandler`` keeps the
    startup date in the base name when it rotates, so a filename-based age
    would be wrong for rotated files.
    """
    log_dir = Path(os.getenv("LOG_DIR", "logs"))
    if not log_dir.is_dir():
        return

    cutoff = time.time() - LOG_RETENTION_DAYS * 86_400
    removed = 0
    errors = 0
    for entry in log_dir.iterdir():
        if not entry.is_file():
            continue
        if not (entry.name.startswith("app_") and ".log" in entry.name):
            continue
        try:
            if entry.stat().st_mtime < cutoff:
                entry.unlink()
                removed += 1
        except Exception:
            log.exception("Failed to delete old log file: %s", entry)
            errors += 1

    if removed or errors:
        log.info(
            "Log cleanup: removed %d file(s) older than %d day(s)%s.",
            removed,
            LOG_RETENTION_DAYS,
            f", {errors} error(s)" if errors else "",
        )


# ---------------------------------------------------------------------------
# Background cleanup thread
# ---------------------------------------------------------------------------

def _cleanup_loop(interval: int, stop_event: threading.Event) -> None:
    """Run periodic sweeps until *stop_event* is set or the process exits.

    The first sweep runs immediately on startup so stale data is removed
    without waiting for the first interval to elapse.
    """
    while True:
        _sweep_cache_db()
        _purge_old_logs()
        if stop_event.wait(timeout=interval):
            break


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

_stop_event: threading.Event | None = None
_cleanup_thread: threading.Thread | None = None


def _stop_background_cleanup_and_close_cache_db() -> None:
    """Stop the background cleanup thread before closing the shared cache DB."""
    global _stop_event, _cleanup_thread

    if _stop_event is not None:
        _stop_event.set()

    thread = _cleanup_thread
    if (
        thread is not None
        and thread.is_alive()
        and thread is not threading.current_thread()
    ):
        thread.join(timeout=5.0)

    _vacuum_and_close_cache_db()


def register() -> None:
    """Start the background cleanup thread and register the atexit hook.

    Safe to call multiple times (idempotent).
    """
    global _stop_event, _cleanup_thread

    # Stop any previously running background thread before restarting.
    if _stop_event is not None:
        _stop_event.set()

    previous_thread = _cleanup_thread
    if (
        previous_thread is not None
        and previous_thread.is_alive()
        and previous_thread is not threading.current_thread()
    ):
        previous_thread.join(timeout=5.0)

    _stop_event = threading.Event()
    _cleanup_thread = threading.Thread(
        target=_cleanup_loop,
        args=(CLEANUP_INTERVAL_SECONDS, _stop_event),
        name="cache-log-cleanup",
        daemon=True,  # won't block process exit
    )
    _cleanup_thread.start()
    log.info(
        "Background cleanup started (interval: %ds, log retention: %dd).",
        CLEANUP_INTERVAL_SECONDS,
        LOG_RETENTION_DAYS,
    )

    atexit.unregister(_stop_background_cleanup_and_close_cache_db)
    atexit.register(_stop_background_cleanup_and_close_cache_db)