sweeper.py

import logging
import json
import signal
import threading
import time
from datetime import datetime
from typing import Any, Dict, Optional

from api.health import check_readiness
from config import settings
from queues.redis_queue import RedisRunQueue
from runtime.dispatch import (
    mark_dispatch_enqueued,
    mark_dispatch_enqueue_failed,
    needs_dispatch_recovery,
)
from runtime.errors import ErrorClassification, build_error_record
from runtime.heartbeat import clear_job_heartbeat, is_job_heartbeat_stale
from runtime.lifecycle import append_lifecycle_event
from runtime.metrics import (
    record_background_startup_readiness_failure,
    record_sweeper_dispatch_failed,
    record_sweeper_version_conflict,
)
from runtime.observability import get_or_create_correlation_id, log_structured
from runtime.retry import record_retry_failure, should_retry
from runtime.state import RuntimeState
from runtime.status import update_run_status
from storage.factory import get_workflow_repository
from storage.postgres_repository import RunVersionConflictError
from startup_validation import validate_startup_or_raise


logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
)


def sweep_stale_runs(
    repository=None,
    queue=None,
    stale_after_seconds: Optional[int] = None,
    now: Optional[datetime] = None,
) -> Dict[str, int]:
    repository = repository or get_workflow_repository()
    queue = queue or RedisRunQueue()
    result = {
        "scanned": 0,
        "stale": 0,
        "requeued": 0,
        "failed": 0,
        "processing_reclaimed": 0,
        "dispatch_recovered": 0,
        "dispatch_failed": 0,
        "version_conflicts": 0,
    }
    if hasattr(queue, "reclaim_stale_processing_jobs"):
        result["processing_reclaimed"] = queue.reclaim_stale_processing_jobs(
            stale_after_seconds or settings.worker_heartbeat_stale_seconds
        )

    for run_summary in list_sweeper_candidates(repository):
        run_id = run_summary.get("run_id")
        if not run_id:
            continue

        result["scanned"] += 1
        saved_state = repository.load_run(run_id)
        if saved_state is None:
            continue

        state = RuntimeState(**saved_state)
        try:
            dispatch_result = recover_queued_dispatch(repository, queue, state)
            if dispatch_result == "recovered":
                result["dispatch_recovered"] += 1
                continue
            if dispatch_result == "failed":
                result["dispatch_failed"] += 1
                continue
        except RunVersionConflictError as exc:
            handle_sweeper_run_version_conflict(
                repository=repository,
                run_id=state.run_id,
                action="dispatch_recovery",
                conflict=exc,
            )
            record_sweeper_version_conflict()
            result["version_conflicts"] += 1
            continue

        if not is_job_heartbeat_stale(
            state.data,
            stale_after_seconds=stale_after_seconds,
            now=now,
        ):
            continue

        result["stale"] += 1
        try:
            if requeue_stale_state(repository, queue, state):
                result["requeued"] += 1
            else:
                mark_stale_state_failed(repository, state)
                result["failed"] += 1
        except RunVersionConflictError as exc:
            handle_sweeper_run_version_conflict(
                repository=repository,
                run_id=state.run_id,
                action="stale_requeue",
                conflict=exc,
            )
            record_sweeper_version_conflict()
            result["version_conflicts"] += 1

    return result


def list_sweeper_candidates(repository) -> list:
    if hasattr(repository, "list_active_runs"):
        return repository.list_active_runs()

    return repository.list_runs()


def requeue_stale_state(repository, queue, state: RuntimeState) -> bool:
    retry = record_retry_failure(
        state.data,
        "Worker heartbeat became stale",
    )
    if not should_retry(state.data):
        return False

    state.data["_status"] = "queued"
    clear_job_heartbeat(state.data, status="queued")
    append_lifecycle_event(
        state.data,
        event_type="stale_job_requeued",
        source="sweeper",
        details={
            "retry_count": retry.get("retry_count"),
            "next_delay_seconds": retry.get("next_delay_seconds"),
        },
    )
    repository.save_run(state.run_id, state.model_dump())

    queue.requeue_job(
        build_job_from_state(state),
        delay_seconds=int(retry.get("next_delay_seconds") or 0),
    )
    log_structured(
        logging,
        logging.WARNING,
        event="sweeper_stale_job_requeued",
        service="sweeper",
        run_id=state.run_id,
        job_type=build_job_from_state(state)["type"],
        correlation_id=get_or_create_correlation_id(state.data),
        details={
            "retry_count": retry.get("retry_count"),
            "next_delay_seconds": retry.get("next_delay_seconds"),
        },
    )
    return True


def recover_queued_dispatch(repository, queue, state: RuntimeState) -> Optional[str]:
    if not needs_dispatch_recovery(state.data):
        return None

    job_type = get_dispatch_job_type(state.data)
    tenant_id = state.data.get("tenant_id", "default")
    mark_dispatch_enqueued(
        state.data,
        queue="startup_runs",
        job_type=job_type,
        recovered_by="sweeper",
    )
    append_lifecycle_event(
        state.data,
        event_type="queued_job_recovered",
        source="sweeper",
        details={
            "queue": "startup_runs",
            "job_type": job_type,
        },
    )
    repository.save_run(state.run_id, state.model_dump())

    if job_type == "startup_recovery":
        enqueue = lambda: queue.enqueue_recovery(state.run_id, tenant_id=tenant_id)
    else:
        enqueue = lambda: queue.enqueue_run(state.run_id, tenant_id=tenant_id)

    try:
        enqueue()
    except Exception as exc:
        mark_dispatch_enqueue_failed(
            state.data,
            queue="startup_runs",
            job_type=job_type,
            error=exc,
        )
        append_lifecycle_event(
            state.data,
            event_type="queued_job_recovery_enqueue_failed",
            source="sweeper",
            details={
                "queue": "startup_runs",
                "job_type": job_type,
                "error_type": type(exc).__name__,
            },
        )
        repository.save_run(state.run_id, state.model_dump())
        record_sweeper_dispatch_failed()
        log_structured(
            logging,
            logging.ERROR,
            event="sweeper_queued_job_recovery_enqueue_failed",
            service="sweeper",
            run_id=state.run_id,
            job_type=job_type,
            correlation_id=get_or_create_correlation_id(state.data),
            details={
                "queue": "startup_runs",
                "error": str(exc),
                "error_type": type(exc).__name__,
            },
        )
        return "failed"

    log_structured(
        logging,
        logging.WARNING,
        event="sweeper_queued_job_recovered",
        service="sweeper",
        run_id=state.run_id,
        job_type=job_type,
        correlation_id=get_or_create_correlation_id(state.data),
        details={
            "queue": "startup_runs",
        },
    )
    return "recovered"


def handle_sweeper_run_version_conflict(
    repository,
    run_id: str,
    action: str,
    conflict: RunVersionConflictError,
) -> bool:
    saved_state = repository.load_run(run_id)
    if saved_state is None:
        log_structured(
            logging,
            logging.ERROR,
            event="sweeper_run_version_conflict_missing_run",
            service="sweeper",
            run_id=run_id,
            job_type=action,
            details={
                "expected_version": conflict.expected_version,
                "actual_version": conflict.actual_version,
            },
        )
        return False

    state = RuntimeState(**saved_state)
    details = {
        "action": action,
        "expected_version": conflict.expected_version,
        "actual_version": conflict.actual_version,
        "latest_status": state.data.get("_status"),
    }
    if state.data.get("_status") in {"completed", "failed", "waiting_for_human"}:
        event_type = "sweeper_run_version_conflict_terminal_skip"
        level = logging.INFO
    else:
        event_type = "sweeper_run_version_conflict_rescan_required"
        level = logging.WARNING

    append_lifecycle_event(
        state.data,
        event_type=event_type,
        source="sweeper",
        details=details,
    )
    repository.save_run(state.run_id, state.model_dump())
    log_structured(
        logging,
        level,
        event=event_type,
        service="sweeper",
        run_id=run_id,
        job_type=action,
        correlation_id=state.data.get("_correlation_id"),
        details=details,
    )
    return True


def get_dispatch_job_type(state_data: Dict[str, Any]) -> str:
    dispatch = state_data.get("_dispatch") or {}
    job_type = dispatch.get("job_type")
    if job_type:
        return job_type

    if state_data.get("_recovery_job", {}).get("status") == "queued":
        return "startup_recovery"

    return "startup_run"


def mark_stale_state_failed(repository, state: RuntimeState) -> None:
    state.pending_agents = []
    state.failed_agents = state.failed_agents or ["worker"]
    clear_job_heartbeat(state.data, status="failed")
    state.errors.append(
        build_error_record(
            agent_name="worker",
            classification=ErrorClassification(
                retryable=False,
                code="stale_job_timeout",
                message="Worker heartbeat became stale and retries were exhausted",
            ),
            source="sweeper",
            retry_count=state.data.get("_retry", {}).get("retry_count"),
            max_retries=state.data.get("_retry", {}).get("max_retries"),
        )
    )
    append_lifecycle_event(
        state.data,
        event_type="stale_job_failed",
        source="sweeper",
    )
    update_run_status(state)
    repository.save_run(state.run_id, state.model_dump())
    log_structured(
        logging,
        logging.ERROR,
        event="sweeper_stale_job_failed",
        service="sweeper",
        run_id=state.run_id,
        job_type=build_job_from_state(state)["type"],
        correlation_id=get_or_create_correlation_id(state.data),
        details={
            "error_code": "stale_job_timeout",
        },
    )


def build_job_from_state(state: RuntimeState) -> Dict[str, Any]:
    heartbeat = state.data.get("_heartbeat", {})
    job_type = heartbeat.get("job_type") or "startup_run"
    return {
        "type": job_type,
        "run_id": state.run_id,
        "tenant_id": state.data.get("tenant_id", "default"),
    }


def create_sweeper_shutdown_event() -> threading.Event:
    return threading.Event()


def request_sweeper_shutdown(stop_event: threading.Event) -> None:
    stop_event.set()


def install_sweeper_shutdown_handlers(stop_event: threading.Event) -> None:
    for signal_name in ("SIGTERM", "SIGINT"):
        shutdown_signal = getattr(signal, signal_name, None)
        if shutdown_signal is None:
            continue
        signal.signal(
            shutdown_signal,
            lambda *_: request_sweeper_shutdown(stop_event),
        )


def sleep_until_next_scan(seconds: int, stop_event: threading.Event) -> None:
    stop_event.wait(seconds)


def run_sweeper_loop(stop_event: Optional[threading.Event] = None) -> None:
    validate_sweeper_startup()
    stop_event = stop_event or create_sweeper_shutdown_event()
    install_sweeper_shutdown_handlers(stop_event)
    log_structured(
        logging,
        logging.INFO,
        event="sweeper_started",
        service="sweeper",
        run_id="-",
        job_type="sweeper_loop",
    )
    while not stop_event.is_set():
        result = sweep_stale_runs()
        log_structured(
            logging,
            logging.INFO,
            event="sweeper_scan_completed",
            service="sweeper",
            run_id="-",
            job_type="sweeper_loop",
            details=result,
        )
        if stop_event.is_set():
            break
        sleep_until_next_scan(settings.worker_heartbeat_interval_seconds, stop_event)
    log_structured(
        logging,
        logging.INFO,
        event="sweeper_shutdown_completed",
        service="sweeper",
        run_id="-",
        job_type="sweeper_loop",
        details={
            "signal": "SIGTERM",
            "graceful_shutdown": True,
        },
    )


def validate_sweeper_startup() -> None:
    validate_startup_or_raise()
    readiness = check_readiness()
    if not readiness["ready"]:
        record_background_startup_readiness_failure("sweeper")
        log_structured(
            logging,
            logging.ERROR,
            event="sweeper_startup_readiness_failed",
            service="sweeper",
            details={
                "dependencies": readiness["dependencies"],
            },
        )
        raise RuntimeError(
            "sweeper startup readiness failed: "
            + json.dumps(readiness["dependencies"], ensure_ascii=False)
        )


if __name__ == "__main__":
    run_sweeper_loop()