FabledCurator/backend/app/celery_signals.py

"""FC-3i: task_run lifecycle via Celery signals.

Subscribes to task_prerun / task_postrun / task_failure / task_retry
and persists one task_run row per task attempt. Drop-in for every
existing and future Celery task — no per-task instrumentation.

Signal handlers run inside the worker process (sync context); DB
writes go through the existing shared sync engine
(backend.app.tasks._sync_engine.sync_session_factory) — one engine
per worker process, not per-task, so we don't blow Postgres
max_connections under load (the reason FC-3g shared-engine fix
existed).

Failure-mode discipline (operator-pressed point): every handler is
wrapped in try/except that swallows + logs. If the DB is down or the
handler has a bug, the real task still runs — the dashboard goes
dark for that interval. Monitoring NEVER breaks the thing it's
monitoring.
"""

import functools
import logging
from datetime import UTC, datetime

from celery.exceptions import SoftTimeLimitExceeded
from celery.signals import task_failure, task_postrun, task_prerun, task_retry

from .models import TaskRun
from .tasks._sync_engine import sync_session_factory

log = logging.getLogger(__name__)

# Celery-internal tasks that would generate dashboard noise without
# operational value. Conservative list; extend only when a specific
# task proves noisy.
_UNTRACKED_TASK_NAMES = frozenset({
    "celery.chord_unlock",
    "celery.backend_cleanup",
    "celery.chunks",
})

_MAX_ERROR_MESSAGE_LEN = 2000
_MAX_ARGS_SUMMARY_LEN = 255
_MAX_WORKER_HOSTNAME_LEN = 128

# PostgreSQL Integer is signed 32-bit. Tasks called with a first-arg
# int outside this range (e.g. an absurdly large mock value, or a
# string-of-digits coercible to int but bigger than 2^31-1) would crash
# the INSERT with NumericValueOutOfRange. Bound the recorded value to
# the column's range; values outside become None (target_id is
# nullable, so this is safe).
_INT32_MAX = 2_147_483_647
_INT32_MIN = -2_147_483_648


def _queue_for(task) -> str:
    """The queue Celery routes this task to — asked of the router itself.

    This was a hand-kept copy of `celery_app.task_routes`, and it drifted
    twice (the 2026-06-02 audit, then #4432). Long-lane jobs were recorded as
    `maintenance`, and translation and gpu_queue runs as `default`, where the
    5-minute stall sweep failed healthy 35-minute translation runs. The router
    answers from the same table the broker uses, so the two cannot disagree.
    """
    name = getattr(task, "name", "") or ""
    app = getattr(task, "app", None)
    if app is None:
        from .celery_app import celery as app
    return _routed_queue(app, name)


@functools.lru_cache(maxsize=1024)
def _routed_queue(app, name: str) -> str:
    try:
        queue = app.amqp.router.route({}, name).get("queue")
    except Exception:  # noqa: BLE001 — monitoring never breaks the task
        log.warning("task_run: could not resolve the queue for %s", name)
        return "default"
    return getattr(queue, "name", None) or (queue if isinstance(queue, str) else "default")


def _target_id_from_args(args) -> int | None:
    """Best-effort: if the first positional arg parses as int AND fits
    in the column's signed-32-bit range, record it as target_id
    (image_id, source_id, etc.). Never raises."""
    if not args:
        return None
    try:
        value = int(args[0])
    except (TypeError, ValueError):
        return None
    if value < _INT32_MIN or value > _INT32_MAX:
        return None
    return value


def _truncate(s, limit: int) -> str | None:
    if s is None:
        return None
    text = str(s)
    return text if len(text) <= limit else text[:limit]


def _is_tracked(task_name: str | None) -> bool:
    return bool(task_name) and task_name not in _UNTRACKED_TASK_NAMES


@task_prerun.connect
def _on_prerun(sender=None, task_id=None, task=None, args=None,
               kwargs=None, **_):
    if not _is_tracked(getattr(task, "name", None)):
        return
    try:
        Session = sync_session_factory()
        with Session() as session:
            session.add(TaskRun(
                celery_task_id=task_id or "",
                queue=_queue_for(task),
                task_name=task.name,
                target_id=_target_id_from_args(args),
                started_at=datetime.now(UTC),
                status="running",
                args_summary=_truncate(repr(args), _MAX_ARGS_SUMMARY_LEN),
                worker_hostname=_truncate(
                    getattr(sender, "hostname", None),
                    _MAX_WORKER_HOSTNAME_LEN,
                ),
            ))
            session.commit()
    except Exception:  # noqa: BLE001 — never break the worker
        log.exception("task_run prerun insert failed (task=%s)",
                      getattr(task, "name", "?"))


def _finalize(task_id: str, *, status: str,
              error_type: str | None = None,
              error_message: str | None = None,
              retry_count: int | None = None) -> None:
    """Shared write path for postrun/failure/retry. Picks the most-
    recent task_run row for this celery_task_id that's still 'running'
    (retries reuse the same celery_task_id; each new attempt's prerun
    inserts a fresh row, so finalize targets the latest running row)."""
    try:
        from sqlalchemy import select

        Session = sync_session_factory()
        now = datetime.now(UTC)
        with Session() as session:
            row = session.execute(
                select(TaskRun)
                .where(TaskRun.celery_task_id == task_id)
                .where(TaskRun.status == "running")
                .order_by(TaskRun.id.desc())
                .limit(1)
            ).scalar_one_or_none()
            if row is None:
                return  # no prerun row (untracked, insert failed, or already finalized)
            row.finished_at = now
            row.duration_ms = int(
                (now - row.started_at).total_seconds() * 1000
            )
            row.status = status
            if error_type is not None:
                row.error_type = _truncate(error_type, 128)
            if error_message is not None:
                row.error_message = _truncate(error_message, _MAX_ERROR_MESSAGE_LEN)
            if retry_count is not None:
                row.retry_count = retry_count
            session.commit()
    except Exception:  # noqa: BLE001
        log.exception("task_run finalize failed (task_id=%s)", task_id)


@task_postrun.connect
def _on_postrun(sender=None, task_id=None, task=None, args=None,
                kwargs=None, retval=None, state=None, **_):
    if not _is_tracked(getattr(task, "name", None)):
        return
    # state is one of SUCCESS/FAILURE/RETRY/etc. Only handle SUCCESS;
    # task_failure handles FAILURE explicitly (with the exception).
    if state != "SUCCESS":
        return
    _finalize(task_id, status="ok")


@task_failure.connect
def _on_failure(sender=None, task_id=None, exception=None,
                args=None, kwargs=None, einfo=None, **_):
    if not _is_tracked(getattr(sender, "name", None)):
        return
    status = ("timeout"
              if isinstance(exception, SoftTimeLimitExceeded)
              else "error")
    _finalize(
        task_id, status=status,
        error_type=type(exception).__name__ if exception else "Unknown",
        error_message=str(exception) if exception else None,
    )


@task_retry.connect
def _on_retry(sender=None, request=None, reason=None, einfo=None, **_):
    if not _is_tracked(getattr(sender, "name", None)):
        return
    task_id = getattr(request, "id", None)
    if not task_id:
        return
    # Mark current attempt's row as 'retry' (terminal for this row).
    # The next attempt's task_prerun inserts a fresh row.
    _finalize(
        task_id, status="retry",
        error_type=type(reason).__name__ if reason else "Retry",
        error_message=str(reason) if reason else None,
        retry_count=getattr(request, "retries", 0),
    )