agentic-soc-platform/backend/apps/common/worker_runner.py

import logging
import time
from dataclasses import dataclass

from django.core.management.base import CommandError

from apps.common.logging import configure_process_file_logging

SLEEP_ALWAYS = "always"
SLEEP_WHEN_IDLE = "when_idle"

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class WorkerIterationResult:
    processed: bool = False
    message: str = ""


def add_worker_arguments(parser, *, interval_help="Seconds between worker iterations."):
    parser.add_argument("--once", action="store_true", help="Run one worker iteration and exit.")
    parser.add_argument("--interval", type=float, help=interval_help)


def _coerce_result(result):
    if isinstance(result, WorkerIterationResult):
        return result
    if isinstance(result, bool):
        return WorkerIterationResult(processed=result)
    return WorkerIterationResult(processed=bool(result))


def _resolve_interval(options, default_interval):
    interval = options.get("interval")
    if interval is None:
        interval = default_interval() if callable(default_interval) else default_interval
    if interval <= 0:
        raise CommandError("--interval must be greater than 0.")
    return interval


def _styled(command, style_name, message):
    styler = getattr(command.style, style_name, None)
    return styler(message) if styler else message


def _worker_label(worker_name):
    return worker_name if worker_name.endswith("worker") else f"{worker_name} worker"


def _refresh_runtime_config_cache():
    from apps.settings.runtime_config import invalidate

    invalidate()


def _run_once_or_raise(worker_label, run_once):
    try:
        _refresh_runtime_config_cache()
        return _coerce_result(run_once())
    except CommandError:
        raise
    except Exception as exc:
        raise CommandError(f"{worker_label} iteration failed: {type(exc).__name__}: {exc}") from exc


def _should_sleep(result, sleep_policy):
    if sleep_policy == SLEEP_ALWAYS:
        return True
    if sleep_policy == SLEEP_WHEN_IDLE:
        return not result.processed
    raise ValueError(f"unsupported worker sleep policy: {sleep_policy}")


def run_worker(
    command,
    *,
    options,
    worker_name,
    run_once,
    default_interval,
    sleep_policy=SLEEP_WHEN_IDLE,
    started_message=None,
    stopped_message=None,
    sleep_seconds=None,
    log_role=None,
):
    worker_label = _worker_label(worker_name)
    if sleep_policy not in {SLEEP_ALWAYS, SLEEP_WHEN_IDLE}:
        raise ValueError(f"unsupported worker sleep policy: {sleep_policy}")
    if log_role and not configure_process_file_logging(log_role):
        raise CommandError(f"unsupported worker log role: {log_role}")
    interval = _resolve_interval(options, default_interval)

    if options["once"]:
        result = _run_once_or_raise(worker_label, run_once)
        command.stdout.write(result.message or f"{worker_label} completed one iteration")
        return

    command.stdout.write(_styled(command, "SUCCESS", started_message or f"{worker_label} started"))
    try:
        while True:
            try:
                result = _run_once_or_raise(worker_label, run_once)
                if result.message:
                    command.stdout.write(result.message)
                if _should_sleep(result, sleep_policy):
                    current_sleep_seconds = sleep_seconds() if sleep_seconds else interval
                    if current_sleep_seconds <= 0:
                        raise CommandError("worker sleep interval must be greater than 0.")
                    time.sleep(current_sleep_seconds)
            except Exception as exc:
                logger.exception("%s iteration failed", worker_label)
                command.stderr.write(_styled(command, "ERROR", f"{worker_label} failed: {type(exc).__name__}: {exc}"))
                time.sleep(interval)
    except KeyboardInterrupt:
        command.stdout.write(_styled(command, "WARNING", stopped_message or f"{worker_label} stopped."))