diff --git a/backend/open_webui/env.py b/backend/open_webui/env.py index 68c4272872..5936ad9939 100644 --- a/backend/open_webui/env.py +++ b/backend/open_webui/env.py @@ -388,6 +388,14 @@ try: except ValueError: REDIS_RESPONSE_STREAM_TTL = 3600 +# Seconds a task survives without a heartbeat. 0 disables expiry. +try: + REDIS_TASK_TTL = int(os.getenv('REDIS_TASK_TTL', '300')) + if REDIS_TASK_TTL != 0 and REDIS_TASK_TTL < 60: + REDIS_TASK_TTL = 300 +except ValueError: + REDIS_TASK_TTL = 300 + REDIS_SENTINEL_HOSTS = os.getenv('REDIS_SENTINEL_HOSTS', '') REDIS_SENTINEL_PORT = os.getenv('REDIS_SENTINEL_PORT', '26379') diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py index 0e77bd4fe0..e7c2915130 100644 --- a/backend/open_webui/main.py +++ b/backend/open_webui/main.py @@ -74,6 +74,7 @@ from open_webui.config import ( seed_registered_defaults, ) from open_webui.constants import ERROR_MESSAGES, TASKS +from open_webui.utils.recurrence import RecurrenceEvaluationTimeout from open_webui.env import ( USE_SLIM, AIOHTTP_CLIENT_SESSION_SSL, @@ -107,6 +108,7 @@ from open_webui.env import ( MAX_BODY_LOG_SIZE, # Redis REDIS_KEY_PREFIX, + REDIS_TASK_TTL, REDIS_URL, RESET_CONFIG_ON_START, SAFE_MODE, @@ -202,6 +204,7 @@ from open_webui.tasks import ( list_task_ids_by_item_id, list_tasks, redis_task_command_listener, + redis_task_heartbeat, stop_item_tasks, stop_task, ) # Import from tasks.py @@ -388,6 +391,8 @@ async def lifespan(app: FastAPI): if app.state.redis is not None: app.state.redis_task_command_listener = asyncio.create_task(redis_task_command_listener(app)) + if REDIS_TASK_TTL > 0: + app.state.redis_task_heartbeat = asyncio.create_task(redis_task_heartbeat(app)) if WEBSOCKET_MANAGER == 'redis': app.state.redis_event_listener = asyncio.create_task(redis_event_listener()) @@ -478,6 +483,9 @@ async def lifespan(app: FastAPI): if hasattr(app.state, 'redis_task_command_listener'): app.state.redis_task_command_listener.cancel() + if hasattr(app.state, 'redis_task_heartbeat'): + app.state.redis_task_heartbeat.cancel() + if hasattr(app.state, 'redis_event_listener'): app.state.redis_event_listener.cancel() @@ -503,6 +511,12 @@ app = FastAPI( lifespan=lifespan, ) + +@app.exception_handler(RecurrenceEvaluationTimeout) +async def recurrence_timeout_handler(request: Request, exc: RecurrenceEvaluationTimeout): + return JSONResponse(status_code=400, content={'detail': str(exc)}) + + # Used by readiness checks to gate traffic until startup work is done. app.state.startup_complete = False diff --git a/backend/open_webui/models/automations.py b/backend/open_webui/models/automations.py index 11a906670b..8c7f6089e3 100644 --- a/backend/open_webui/models/automations.py +++ b/backend/open_webui/models/automations.py @@ -312,6 +312,7 @@ class AutomationTable: rows = result.scalars().all() from open_webui.utils.automations import next_run_ns + from open_webui.utils.recurrence import RecurrenceEvaluationTimeout # Batch-fetch user timezones so rescheduling respects each # user's local timezone instead of falling back to server time. @@ -323,13 +324,20 @@ class AutomationTable: tz_result = await db.execute(select(User.id, User.timezone).where(User.id.in_(user_ids))) timezone_by_user_id = {uid: tz for uid, tz in tz_result.all()} + claimed = [] for row in rows: + try: + next_run_at = await next_run_ns(row.data.get('rrule', ''), tz=timezone_by_user_id.get(row.user_id)) + except RecurrenceEvaluationTimeout: + log.warning('Skipping automation %s: recurrence evaluation timed out', row.id) + continue row.last_run_at = now_ns - row.next_run_at = next_run_ns(row.data.get('rrule', ''), tz=timezone_by_user_id.get(row.user_id)) + row.next_run_at = next_run_at + claimed.append(row) await db.commit() - return [AutomationModel.model_validate(r) for r in rows] + return [AutomationModel.model_validate(r) for r in claimed] #################### diff --git a/backend/open_webui/models/calendar.py b/backend/open_webui/models/calendar.py index 3a3147a7ad..fd7390fcea 100644 --- a/backend/open_webui/models/calendar.py +++ b/backend/open_webui/models/calendar.py @@ -8,7 +8,7 @@ from open_webui.constants import ERROR_MESSAGES from open_webui.models.access_grants import AccessGrantModel, AccessGrants from open_webui.models.groups import Groups from open_webui.models.users import User, UserModel, UserResponse -from pydantic import BaseModel, ConfigDict, Field, field_validator +from pydantic import BaseModel, ConfigDict, Field from sqlalchemy import ( JSON, BigInteger, @@ -179,6 +179,20 @@ class CalendarUpdateForm(BaseModel): access_grants: Optional[list[dict]] = None +async def validate_calendar_rrule(value: Optional[str]) -> None: + if value: + from open_webui.utils.recurrence import rrule_interval_seconds + + try: + interval = await rrule_interval_seconds(value) + except ValueError: + raise + except Exception as e: + raise ValueError(ERROR_MESSAGES.AUTOMATION_INVALID_RRULE(e)) from e + if interval is not None and interval < MIN_CALENDAR_RRULE_INTERVAL_SECONDS: + raise ValueError(ERROR_MESSAGES.CALENDAR_RRULE_TOO_FREQUENT) + + class CalendarEventForm(BaseModel): calendar_id: str title: str @@ -193,22 +207,6 @@ class CalendarEventForm(BaseModel): meta: Optional[dict] = None attendees: Optional[list[dict]] = None - @field_validator('rrule') - @classmethod - def reject_sub_daily_rrule(cls, value: Optional[str]) -> Optional[str]: - if value: - from open_webui.utils.automations import rrule_interval_seconds - - try: - interval = rrule_interval_seconds(value) - except ValueError: - raise - except Exception as e: - raise ValueError(ERROR_MESSAGES.AUTOMATION_INVALID_RRULE(e)) - if interval is not None and interval < MIN_CALENDAR_RRULE_INTERVAL_SECONDS: - raise ValueError(ERROR_MESSAGES.CALENDAR_RRULE_TOO_FREQUENT) - return value - class CalendarEventUpdateForm(BaseModel): calendar_id: Optional[str] = None @@ -225,22 +223,6 @@ class CalendarEventUpdateForm(BaseModel): is_cancelled: Optional[bool] = None attendees: Optional[list[dict]] = None - @field_validator('rrule') - @classmethod - def reject_sub_daily_rrule(cls, value: Optional[str]) -> Optional[str]: - if value: - from open_webui.utils.automations import rrule_interval_seconds - - try: - interval = rrule_interval_seconds(value) - except ValueError: - raise - except Exception as e: - raise ValueError(ERROR_MESSAGES.AUTOMATION_INVALID_RRULE(e)) - if interval is not None and interval < MIN_CALENDAR_RRULE_INTERVAL_SECONDS: - raise ValueError(ERROR_MESSAGES.CALENDAR_RRULE_TOO_FREQUENT) - return value - class RSVPForm(BaseModel): status: str # 'accepted' | 'declined' | 'tentative' | 'pending' @@ -465,6 +447,7 @@ class CalendarEventTable: async def insert_new_event( self, user_id: str, form_data: CalendarEventForm, db: Optional[AsyncSession] = None ) -> Optional[CalendarEventModel]: + await validate_calendar_rrule(form_data.rrule) async with get_async_db_context(db) as db: now = int(time.time_ns()) event = CalendarEvent( @@ -695,6 +678,7 @@ class CalendarEventTable: async def update_event_by_id( self, id: str, form_data: CalendarEventUpdateForm, db: Optional[AsyncSession] = None ) -> Optional[CalendarEventModel]: + await validate_calendar_rrule(form_data.rrule) async with get_async_db_context(db) as db: result = await db.execute(select(CalendarEvent).filter(CalendarEvent.id == id)) event = result.scalars().first() diff --git a/backend/open_webui/models/chats.py b/backend/open_webui/models/chats.py index 46131dc8e9..2af71fd944 100644 --- a/backend/open_webui/models/chats.py +++ b/backend/open_webui/models/chats.py @@ -2536,18 +2536,15 @@ class ChatTable: except Exception: return False - async def move_chats_by_user_id_and_folder_id( + async def move_chats_by_folder_id( self, - user_id: str, folder_id: str, new_folder_id: str | None, db: AsyncSession | None = None, ) -> bool: try: async with get_async_db_context(db) as session: - await session.execute( - update(Chat).filter_by(user_id=user_id, folder_id=folder_id).values(folder_id=new_folder_id) - ) + await session.execute(update(Chat).filter_by(folder_id=folder_id).values(folder_id=new_folder_id)) await session.commit() return True diff --git a/backend/open_webui/routers/automations.py b/backend/open_webui/routers/automations.py index caf8879e28..4b44a8b8ed 100644 --- a/backend/open_webui/routers/automations.py +++ b/backend/open_webui/routers/automations.py @@ -87,7 +87,7 @@ async def check_automation_limits(request, user, rrule_str: str, db, is_create: if min_interval: min_interval = int(min_interval) if min_interval > 0: - interval = rrule_interval_seconds(rrule_str) + interval = await rrule_interval_seconds(rrule_str) if interval is not None and interval < min_interval: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -150,7 +150,7 @@ async def enrich_automation(automation: AutomationModel, db: AsyncSession, tz: s return AutomationResponse( **automation.model_dump(), last_run=last_run, - next_runs=next_n_runs_ns(automation.data['rrule'], tz=tz), + next_runs=await next_n_runs_ns(automation.data['rrule'], tz=tz), ) @@ -216,7 +216,7 @@ async def create_new_automation( await check_automation_folder_access(form_data.folder_id, user, db) await check_automation_channel_access(form_data, user, db) try: - validate_rrule(form_data.data.rrule, tz=user.timezone) + await validate_rrule(form_data.data.rrule, tz=user.timezone) except ValueError as e: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -226,7 +226,7 @@ async def create_new_automation( await check_automation_limits(request, user, form_data.data.rrule, db, is_create=True) tz = user.timezone - automation = await Automations.insert(user.id, form_data, next_run_ns(form_data.data.rrule, tz=tz), db=db) + automation = await Automations.insert(user.id, form_data, await next_run_ns(form_data.data.rrule, tz=tz), db=db) response = await enrich_automation(automation, db, tz=tz) await publish_event( request, @@ -276,7 +276,7 @@ async def update_automation_by_id( await check_automation_channel_access(form_data, user, db) try: - validate_rrule(form_data.data.rrule, tz=user.timezone) + await validate_rrule(form_data.data.rrule, tz=user.timezone) except ValueError as e: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -286,7 +286,7 @@ async def update_automation_by_id( await check_automation_limits(request, user, form_data.data.rrule, db, is_create=False) tz = user.timezone - updated = await Automations.update_by_id(id, form_data, next_run_ns(form_data.data.rrule, tz=tz), db=db) + updated = await Automations.update_by_id(id, form_data, await next_run_ns(form_data.data.rrule, tz=tz), db=db) response = await enrich_automation(updated, db, tz=tz) await publish_event( request, @@ -313,7 +313,7 @@ async def toggle_automation_by_id( await check_automations_permission(request, user) automation = await Automations.get_by_id(id, db=db) check_automation_access(automation, user) - toggled = await Automations.toggle(id, next_run_ns(automation.data['rrule'], tz=user.timezone), db=db) + toggled = await Automations.toggle(id, await next_run_ns(automation.data['rrule'], tz=user.timezone), db=db) response = await enrich_automation(toggled, db, tz=user.timezone) await publish_event( request, diff --git a/backend/open_webui/routers/calendar.py b/backend/open_webui/routers/calendar.py index f95be48e84..86d9824803 100644 --- a/backend/open_webui/routers/calendar.py +++ b/backend/open_webui/routers/calendar.py @@ -273,7 +273,10 @@ async def get_events( async def create_event(request: Request, form_data: CalendarEventForm, user: UserModel = Depends(get_verified_user)): await check_calendar_permission(request, user) await _check_calendar_access(form_data.calendar_id, user, 'write') - event = await CalendarEvents.insert_new_event(user.id, form_data) + try: + event = await CalendarEvents.insert_new_event(user.id, form_data) + except ValueError as e: + raise HTTPException(status_code=422, detail=str(e)) from e await publish_event( request, EVENTS.CALENDAR_EVENT_CREATED, @@ -325,7 +328,10 @@ async def update_event( if form_data.calendar_id is not None and form_data.calendar_id != event.calendar_id: await _check_calendar_access(form_data.calendar_id, user, 'write') - updated = await CalendarEvents.update_event_by_id(event_id, form_data) + try: + updated = await CalendarEvents.update_event_by_id(event_id, form_data) + except ValueError as e: + raise HTTPException(status_code=422, detail=str(e)) from e if not updated: raise HTTPException(status_code=500, detail='Failed to update') await publish_event( diff --git a/backend/open_webui/routers/folders.py b/backend/open_webui/routers/folders.py index 9653a6d9b8..046a2d8834 100644 --- a/backend/open_webui/routers/folders.py +++ b/backend/open_webui/routers/folders.py @@ -704,8 +704,8 @@ async def delete_folder_by_id( for folder_id in folder_ids: if delete_contents: await Chats.delete_chats_by_user_id_and_folder_id(folder_owner_id, folder_id, db=db) - else: - await Chats.move_chats_by_user_id_and_folder_id(folder_owner_id, folder_id, None, db=db) + + await Chats.move_chats_by_folder_id(folder_id, None, db=db) # Clean up access grants for this folder await AccessGrants.revoke_all_access('folder', folder_id, db=db) diff --git a/backend/open_webui/tasks.py b/backend/open_webui/tasks.py index 322cd32e60..05f622ef06 100644 --- a/backend/open_webui/tasks.py +++ b/backend/open_webui/tasks.py @@ -6,7 +6,7 @@ from uuid import uuid4 from redis.asyncio import Redis -from open_webui.env import REDIS_KEY_PREFIX, REDIS_RESPONSE_STREAM_TTL +from open_webui.env import REDIS_KEY_PREFIX, REDIS_RESPONSE_STREAM_TTL, REDIS_TASK_TTL from open_webui.utils.json_codec import JSONCodec, dumps_bytes log = logging.getLogger(__name__) @@ -66,13 +66,28 @@ async def redis_task_command_listener(app): reconnect_interval = min(reconnect_interval * 2, REDIS_PUBSUB_MAX_RECONNECT_INTERVAL) +async def redis_task_heartbeat(app): + redis: Redis = app.state.redis + while True: + await asyncio.sleep(REDIS_TASK_TTL / 4) + try: + pipe = redis.pipeline(transaction=False) + for task_id in list(tasks): + # EXPIRE cannot recreate a task already removed by cleanup. + pipe.expire(f'{REDIS_TASKS_KEY}:{task_id}', REDIS_TASK_TTL) + await pipe.execute() + except Exception: + log.exception('Redis task heartbeat failed') + + ### ------------------------------ ### REDIS-ENABLED HANDLERS ### ------------------------------ async def redis_save_task(redis: Redis, task_id: str, item_id: str | None): - pipe = redis.pipeline() + pipe = redis.pipeline(transaction=False) + pipe.set(f'{REDIS_TASKS_KEY}:{task_id}', '1', ex=REDIS_TASK_TTL or None) pipe.hset(REDIS_TASKS_KEY, task_id, item_id or '') if item_id: pipe.sadd(f'{REDIS_ITEM_TASKS_KEY}:{item_id}', task_id) @@ -80,25 +95,36 @@ async def redis_save_task(redis: Redis, task_id: str, item_id: str | None): async def redis_cleanup_task(redis: Redis, task_id: str, item_id: str | None): - pipe = redis.pipeline() + pipe = redis.pipeline(transaction=False) + pipe.delete(f'{REDIS_TASKS_KEY}:{task_id}') pipe.hdel(REDIS_TASKS_KEY, task_id) pipe.hdel(REDIS_RESPONSE_STREAMS_KEY, task_id) if item_id: pipe.srem(f'{REDIS_ITEM_TASKS_KEY}:{item_id}', task_id) - await pipe.execute() - # Remove the set key entirely if no tasks remain for this item - if await redis.scard(f'{REDIS_ITEM_TASKS_KEY}:{item_id}') == 0: - await redis.delete(f'{REDIS_ITEM_TASKS_KEY}:{item_id}') - else: - await pipe.execute() + await pipe.execute() -async def redis_list_tasks(redis: Redis) -> list[str]: - return list(await redis.hkeys(REDIS_TASKS_KEY)) +async def redis_list_tasks(redis: Redis, item_id: str | None = None) -> list[str]: + task_ids = list( + await redis.smembers(f'{REDIS_ITEM_TASKS_KEY}:{item_id}') + if item_id is not None + else await redis.hkeys(REDIS_TASKS_KEY) + ) + if not task_ids or REDIS_TASK_TTL == 0: + return task_ids + pipe = redis.pipeline(transaction=False) + for task_id in task_ids: + pipe.exists(f'{REDIS_TASKS_KEY}:{task_id}') -async def redis_list_item_tasks(redis: Redis, item_id: str) -> list[str]: - return list(await redis.smembers(f'{REDIS_ITEM_TASKS_KEY}:{item_id}')) + active = [] + for task_id, exists in zip(task_ids, await pipe.execute()): + if exists: + active.append(task_id) + else: + task_item_id = item_id if item_id is not None else await redis.hget(REDIS_TASKS_KEY, task_id) + await redis_cleanup_task(redis, task_id, task_item_id or None) + return active async def redis_send_command(redis: Redis, command: dict): @@ -166,7 +192,7 @@ async def list_task_ids_by_item_id(redis, id): List all tasks associated with a specific ID. """ if redis: - return await redis_list_item_tasks(redis, id) + return await redis_list_tasks(redis, id) return list(item_tasks.get(id, [])) diff --git a/backend/open_webui/tools/builtin.py b/backend/open_webui/tools/builtin.py index 8dedd75c13..d80884d1cd 100644 --- a/backend/open_webui/tools/builtin.py +++ b/backend/open_webui/tools/builtin.py @@ -3768,7 +3768,7 @@ async def create_automation( # Validate the RRULE try: - validate_rrule(rrule, tz=user.timezone) + await validate_rrule(rrule, tz=user.timezone) except ValueError as e: return JSONCodec.dumps({'error': f'Invalid schedule: {e}'}) @@ -3794,7 +3794,7 @@ async def create_automation( is_active=True, ) - automation = await Automations.insert(user_id, form, next_run_ns(rrule, tz=tz)) + automation = await Automations.insert(user_id, form, await next_run_ns(rrule, tz=tz)) return JSONCodec.dumps( { @@ -3805,7 +3805,7 @@ async def create_automation( 'model_id': model_id, 'target': automation.data.get('target'), 'is_active': automation.is_active, - 'next_runs': next_n_runs_ns(rrule, tz=tz), + 'next_runs': await next_n_runs_ns(rrule, tz=tz), }, ensure_ascii=False, ) @@ -3876,7 +3876,7 @@ async def update_automation( # Validate RRULE if changed if rrule is not None: try: - validate_rrule(new_rrule, tz=user.timezone) + await validate_rrule(new_rrule, tz=user.timezone) except ValueError as e: return JSONCodec.dumps({'error': f'Invalid schedule: {e}'}) @@ -3898,7 +3898,7 @@ async def update_automation( is_active=automation.is_active, ) - updated = await Automations.update_by_id(automation_id, form, next_run_ns(new_rrule, tz=tz)) + updated = await Automations.update_by_id(automation_id, form, await next_run_ns(new_rrule, tz=tz)) return JSONCodec.dumps( { @@ -3909,7 +3909,7 @@ async def update_automation( 'model_id': new_model_id, 'target': updated.data.get('target'), 'is_active': updated.is_active, - 'next_runs': next_n_runs_ns(new_rrule, tz=tz), + 'next_runs': await next_n_runs_ns(new_rrule, tz=tz), }, ensure_ascii=False, ) @@ -3977,7 +3977,7 @@ async def list_automations( 'rrule': rrule, 'is_active': item.is_active, 'last_run_at': item.last_run_at, - 'next_runs': next_n_runs_ns(rrule, tz=user.timezone if user else None), + 'next_runs': await next_n_runs_ns(rrule, tz=user.timezone if user else None), } ) @@ -4024,7 +4024,7 @@ async def toggle_automation( rrule = automation.data.get('rrule', '') toggled = await Automations.toggle( automation_id, - next_run_ns(rrule, tz=user.timezone if user else None), + await next_run_ns(rrule, tz=user.timezone if user else None), ) return JSONCodec.dumps( diff --git a/backend/open_webui/utils/automations.py b/backend/open_webui/utils/automations.py index 557e84ce13..54887467c0 100644 --- a/backend/open_webui/utils/automations.py +++ b/backend/open_webui/utils/automations.py @@ -20,12 +20,10 @@ import logging import os import random import time -from datetime import datetime, timedelta +from datetime import timedelta from typing import Optional from uuid import uuid4 -from zoneinfo import ZoneInfo -from dateutil.rrule import HOURLY, MINUTELY, SECONDLY, rruleset, rrulestr from fastapi import Request from fastapi.security import HTTPAuthorizationCredentials from open_webui.constants import ERROR_MESSAGES @@ -39,6 +37,13 @@ from open_webui.models.messages import MessageForm from open_webui.models.users import Users from open_webui.utils.auth import create_token from open_webui.utils.misc import parse_duration +from open_webui.utils.recurrence import ( + _resolve_tz, + next_n_runs_ns, + next_run_ns, + rrule_interval_seconds, + validate_rrule, +) from open_webui.utils.task import prompt_template from open_webui.utils.terminals import get_terminal_server_url from starlette.datastructures import Headers @@ -50,151 +55,6 @@ TIMER_POLL_INTERVAL = int(os.getenv('TIMER_POLL_INTERVAL', '1')) CALENDAR_ALERT_LOOKAHEAD_MINUTES = int(os.getenv('CALENDAR_ALERT_LOOKAHEAD_MINUTES', '10')) -#################### -# RRULE Helpers -#################### - - -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 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: - rule = _parse_rule(s, now) - except Exception as e: - raise ValueError(ERROR_MESSAGES.AUTOMATION_INVALID_RRULE(e)) - if rule.after(now) is None: - raise ValueError(ERROR_MESSAGES.AUTOMATION_NO_FUTURE_RUNS) - - -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) - dt = _parse_rule(s, now_naive).after(now_naive) - if dt is None: - return None - if zi: - dt = dt.replace(tzinfo=zi) - return int(dt.timestamp() * 1_000_000_000) - - -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() - rule = _parse_rule(s, now) - dt = now - for _ in range(n): - dt = rule.after(dt) - if not dt: - break - 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 - - -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() - rule = _parse_rule(s, now) - first = rule.after(now) - if first is None: - return None - second = rule.after(first) - if second is None: - return None - return int((second - first).total_seconds()) - - ############################ # Worker Loop ############################ diff --git a/backend/open_webui/utils/recurrence.py b/backend/open_webui/utils/recurrence.py new file mode 100644 index 0000000000..d5c2b6b1f2 --- /dev/null +++ b/backend/open_webui/utils/recurrence.py @@ -0,0 +1,171 @@ +"""Recurrence calculations isolated from application/DB imports for worker processes.""" + +import logging +from datetime import datetime, timedelta +from typing import Optional +from zoneinfo import ZoneInfo + +from anyio import fail_after, to_process +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 + + +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())