openclaw-mission-control/backend/app/services/queue.py

"""Generic Redis-backed queue helpers for RQ-backed background workloads."""

from __future__ import annotations

import json
import time
from dataclasses import dataclass
from datetime import UTC, datetime
from typing import Any, cast

import redis

from app.core.config import settings
from app.core.logging import get_logger

logger = get_logger(__name__)

_SCHEDULED_SUFFIX = ":scheduled"
_DRY_RUN_BATCH_SIZE = 100


@dataclass(frozen=True)
class QueuedTask:
    """Generic queued task envelope."""

    task_type: str
    payload: dict[str, Any]
    created_at: datetime
    attempts: int = 0

    def to_json(self) -> str:
        return json.dumps(
            {
                "task_type": self.task_type,
                "payload": self.payload,
                "created_at": self.created_at.isoformat(),
                "attempts": self.attempts,
            },
            sort_keys=True,
        )


def _redis_client(redis_url: str | None = None) -> redis.Redis:
    return redis.Redis.from_url(redis_url or settings.rq_redis_url)


def _scheduled_queue_name(queue_name: str) -> str:
    return f"{queue_name}{_SCHEDULED_SUFFIX}"


def _now_seconds() -> float:
    return time.time()


def _drain_ready_scheduled_tasks(
    client: redis.Redis,
    queue_name: str,
    *,
    max_items: int = _DRY_RUN_BATCH_SIZE,
) -> float | None:
    scheduled_queue = _scheduled_queue_name(queue_name)
    now = _now_seconds()

    ready_items = cast(
        list[str | bytes],
        client.zrangebyscore(
            scheduled_queue,
            "-inf",
            now,
            start=0,
            num=max_items,
        ),
    )
    if ready_items:
        ready_values = tuple(ready_items)
        client.lpush(queue_name, *ready_values)
        client.zrem(scheduled_queue, *ready_values)
        logger.debug(
            "rq.queue.drain_ready_scheduled",
            extra={
                "queue_name": queue_name,
                "count": len(ready_items),
            },
        )

    next_item = cast(
        list[tuple[str | bytes, float]],
        client.zrangebyscore(
            scheduled_queue,
            now,
            "+inf",
            start=0,
            num=1,
            withscores=True,
        ),
    )
    if not next_item:
        return None

    next_score = float(next_item[0][1])
    return max(0.0, next_score - now)


def _schedule_for_later(
    task: QueuedTask,
    queue_name: str,
    delay_seconds: float,
    *,
    redis_url: str | None = None,
) -> bool:
    client = _redis_client(redis_url=redis_url)
    scheduled_queue = _scheduled_queue_name(queue_name)
    score = _now_seconds() + delay_seconds
    client.zadd(scheduled_queue, {task.to_json(): score})
    logger.info(
        "rq.queue.scheduled",
        extra={
            "task_type": task.task_type,
            "queue_name": queue_name,
            "delay_seconds": delay_seconds,
        },
    )
    return True


def enqueue_task(
    task: QueuedTask,
    queue_name: str,
    *,
    redis_url: str | None = None,
) -> bool:
    """Persist a task envelope in a Redis list-backed queue."""
    try:
        client = _redis_client(redis_url=redis_url)
        client.lpush(queue_name, task.to_json())
        logger.info(
            "rq.queue.enqueued",
            extra={
                "task_type": task.task_type,
                "queue_name": queue_name,
                "attempt": task.attempts,
            },
        )
        return True
    except Exception as exc:
        logger.warning(
            "rq.queue.enqueue_failed",
            extra={"task_type": task.task_type, "queue_name": queue_name, "error": str(exc)},
        )
        return False


def _coerce_datetime(raw: object | None) -> datetime:
    if raw is None:
        return datetime.now(UTC)
    if isinstance(raw, str):
        try:
            return datetime.fromisoformat(raw)
        except ValueError:
            return datetime.now(UTC)
    if isinstance(raw, (int, float)):
        try:
            return datetime.fromtimestamp(raw, tz=UTC)
        except (TypeError, ValueError, OverflowError):
            return datetime.now(UTC)
    return datetime.now(UTC)


def dequeue_task(
    queue_name: str,
    *,
    redis_url: str | None = None,
    block: bool = False,
    block_timeout: float = 0,
) -> QueuedTask | None:
    """Pop one task envelope from the queue."""
    client = _redis_client(redis_url=redis_url)
    timeout = max(0.0, float(block_timeout))
    raw: str | bytes | None
    if block:
        next_delay = _drain_ready_scheduled_tasks(client, queue_name)
        if timeout == 0:
            timeout = next_delay if next_delay is not None else 0
        else:
            timeout = min(timeout, next_delay) if next_delay is not None else timeout
        raw_result = cast(
            tuple[bytes | str, bytes | str] | None,
            client.brpop([queue_name], timeout=timeout),
        )
        if raw_result is None:
            _drain_ready_scheduled_tasks(client, queue_name)
            return None
        raw = raw_result[1]
    else:
        raw = cast(str | bytes | None, client.rpop(queue_name))
    if raw is None:
        _drain_ready_scheduled_tasks(client, queue_name)
        return None
    return _decode_task(raw, queue_name)


def _decode_task(raw: str | bytes, queue_name: str) -> QueuedTask:
    if isinstance(raw, bytes):
        raw = raw.decode("utf-8")

    try:
        payload: dict[str, Any] = json.loads(raw)
        if "task_type" not in payload and "payload" not in payload:
            return QueuedTask(
                task_type="legacy",
                payload=payload,
                created_at=_coerce_datetime(
                    payload.get("created_at") or payload.get("received_at")
                ),
                attempts=int(payload.get("attempts", 0)),
            )
        return QueuedTask(
            task_type=str(payload["task_type"]),
            payload=payload["payload"],
            created_at=datetime.fromisoformat(payload["created_at"]),
            attempts=int(payload.get("attempts", 0)),
        )
    except Exception as exc:
        logger.error(
            "rq.queue.dequeue_failed",
            extra={"queue_name": queue_name, "raw_payload": str(raw), "error": str(exc)},
        )
        raise


def _requeue_with_attempt(task: QueuedTask) -> QueuedTask:
    return QueuedTask(
        task_type=task.task_type,
        payload=task.payload,
        created_at=task.created_at,
        attempts=task.attempts + 1,
    )


def requeue_if_failed(
    task: QueuedTask,
    queue_name: str,
    *,
    max_retries: int,
    redis_url: str | None = None,
    delay_seconds: float = 0,
) -> bool:
    """Requeue a failed task with capped retries.

    Returns True if requeued.
    """
    requeued_task = _requeue_with_attempt(task)
    if requeued_task.attempts > max_retries:
        logger.warning(
            "rq.queue.drop_failed_task",
            extra={
                "task_type": task.task_type,
                "queue_name": queue_name,
                "attempts": requeued_task.attempts,
            },
        )
        return False
    if delay_seconds > 0:
        return _schedule_for_later(
            requeued_task,
            queue_name,
            delay_seconds,
            redis_url=redis_url,
        )
    return enqueue_task(
        requeued_task,
        queue_name,
        redis_url=redis_url,
    )