open-webui/backend/open_webui/utils/recurrence.py

"""Recurrence calculations isolated from application/DB imports for worker processes."""

import asyncio
import logging
from datetime import datetime, timedelta
from functools import partial
from typing import Optional
from zoneinfo import ZoneInfo

import anyio
from anyio import fail_after, to_process, to_thread
from dateutil.rrule import HOURLY, MINUTELY, SECONDLY, rruleset, rrulestr
from open_webui.constants import ERROR_MESSAGES

log = logging.getLogger(__name__)
RRULE_TIMEOUT_SECONDS = 2


class RecurrenceEvaluationTimeout(ValueError):
    """The evaluation budget expired; the schedule may still have occurrences."""


def _resolve_tz(tz: str = None) -> Optional[ZoneInfo]:
    """Safely resolve a timezone string to ZoneInfo.

    Returns None (→ server-local fallback) when *tz* is empty, None,
    or an unrecognised IANA key.  Logs a warning on bad keys so
    misconfiguration is visible in the server logs.
    """
    if not tz:
        return None
    try:
        return ZoneInfo(tz)
    except (KeyError, Exception):
        log.warning('Unknown timezone %r — falling back to server time', tz)
        return None


def _parse_rule(s: str, now: Optional[datetime] = None):
    """Parse RRULE with clock-aligned DTSTART for sub-daily frequencies.

    SECONDLY/MINUTELY/HOURLY rules use a fixed epoch DTSTART (2000-01-01 00:00)
    so intervals snap to clock boundaries (e.g. every 5min = :00, :05, :10).
    """
    upper = s.upper()
    if 'EXRULE' in upper:
        raise ValueError('EXRULE is not supported in recurrence rules')

    parsed = rrulestr(s, ignoretz=True)
    rules = parsed._rrule if isinstance(parsed, rruleset) else [parsed]
    if len(rules) > 1:
        raise ValueError('only one RRULE is supported per recurrence rule')

    rule = rules[0]
    start = rule._dtstart.replace(tzinfo=None)
    anchor = now or datetime.now()
    parts = s.split()
    stripped = '\n'.join(part for part in parts if not part.upper().startswith('DTSTART')) or s
    has_dtstart = any(part.upper().startswith('DTSTART') for part in parts)
    step = {
        SECONDLY: timedelta(seconds=rule._interval),
        MINUTELY: timedelta(minutes=rule._interval),
        HOURLY: timedelta(hours=rule._interval),
    }.get(rule._freq)

    if step is None:
        if not rule._dtstart.tzinfo:
            return parsed
        return rrulestr(stripped, dtstart=start, ignoretz=True)

    if rule._interval < 1:
        raise ValueError('RRULE INTERVAL must be a positive integer')
    dtstart = None
    if has_dtstart:
        emitted = ((anchor - start) // step) if anchor > start else 0
        emitted *= len(rule._byminute or (0,)) * len(rule._bysecond or (0,))
        if emitted <= 100_000:
            if rule._dtstart.tzinfo:
                dtstart = start
            else:
                return parsed
    if not has_dtstart or dtstart is None:
        epoch = datetime(2000, 1, 1)
        dtstart = epoch + ((anchor - epoch) // step) * step

    return rrulestr(stripped, dtstart=dtstart, ignoretz=True)


def _next_occurrences(s: str, now: datetime, n: int) -> list[datetime]:
    rule = _parse_rule(s, now)
    occurrences = []
    for _ in range(n):
        now = rule.after(now)
        if now is None:
            break
        occurrences.append(now)
    return occurrences


async def _get_next_occurrences(s: str, now: datetime, n: int) -> list[datetime]:
    # A result-count or date limit cannot bound work before the first match.
    try:
        with fail_after(RRULE_TIMEOUT_SECONDS):
            return await to_process.run_sync(_next_occurrences, s, now, n, cancellable=True)
    except TimeoutError as e:
        raise RecurrenceEvaluationTimeout('Schedule took too long to evaluate; simplify its recurrence rule.') from e
    except NotImplementedError:
        # Windows' SelectorEventLoop (required by psycopg) cannot spawn subprocesses.
        run_on_proactor_loop = partial(
            anyio.run, _get_next_occurrences, s, now, n, backend_options={'loop_factory': asyncio.ProactorEventLoop}
        )
        return await to_thread.run_sync(run_on_proactor_loop)


async def validate_rrule(s: str, tz: str = None) -> None:
    """Raise ValueError if the RRULE is malformed or exhausted.

    When *tz* is provided the "now" reference uses the user's local
    clock so that near-future schedules are not incorrectly rejected
    on servers whose system clock is ahead (e.g. UTC vs US timezones).
    """
    upper = s.upper()
    if 'COUNT=' in upper and 'DTSTART' not in upper:
        raise ValueError(ERROR_MESSAGES.AUTOMATION_COUNT_REQUIRES_DTSTART)
    zi = _resolve_tz(tz)
    now = datetime.now(zi).replace(tzinfo=None) if zi else datetime.now()
    try:
        occurrences = await _get_next_occurrences(s, now, 1)
    except RecurrenceEvaluationTimeout:
        raise
    except Exception as e:
        raise ValueError(ERROR_MESSAGES.AUTOMATION_INVALID_RRULE(e))
    if not occurrences:
        raise ValueError(ERROR_MESSAGES.AUTOMATION_NO_FUTURE_RUNS)


async def next_run_ns(s: str, tz: str = None) -> Optional[int]:
    """Next occurrence as epoch nanoseconds, respecting user timezone."""
    zi = _resolve_tz(tz)
    now = datetime.now(zi) if zi else datetime.now()
    now_naive = now.replace(tzinfo=None)
    occurrences = await _get_next_occurrences(s, now_naive, 1)
    if not occurrences:
        return None
    dt = occurrences[0]
    if zi:
        dt = dt.replace(tzinfo=zi)
    return int(dt.timestamp() * 1_000_000_000)


async def next_n_runs_ns(s: str, n: int = 5, tz: str = None) -> list[int]:
    """Compute next N occurrences for UI preview.

    Uses the user's timezone for the starting "now" so that the
    preview matches the user's local clock (same as next_run_ns).
    """
    zi = _resolve_tz(tz)
    result = []
    now = datetime.now(zi).replace(tzinfo=None) if zi else datetime.now()
    for dt in await _get_next_occurrences(s, now, n):
        if zi:
            dt_tz = dt.replace(tzinfo=zi)
            result.append(int(dt_tz.timestamp() * 1_000_000_000))
        else:
            result.append(int(dt.timestamp() * 1_000_000_000))
    return result


async def rrule_interval_seconds(s: str) -> Optional[int]:
    """Approximate interval between recurrences in seconds.

    Returns None for one-shot (COUNT=1) schedules or rules
    with fewer than two future occurrences.
    """
    s = '\n'.join(part for part in s.split() if not part.upper().startswith('DTSTART')) or s
    now = datetime.now()
    occurrences = await _get_next_occurrences(s, now, 2)
    if len(occurrences) < 2:
        return None
    return int((occurrences[1] - occurrences[0]).total_seconds())