diff --git a/backend/open_webui/constants.py b/backend/open_webui/constants.py index a1d1bc73ca..5d1dfc7446 100644 --- a/backend/open_webui/constants.py +++ b/backend/open_webui/constants.py @@ -118,6 +118,7 @@ class ERROR_MESSAGES(str, Enum): AUTOMATION_TOO_FREQUENT = lambda interval='': f'Schedule too frequent. Minimum interval is {interval} seconds.' AUTOMATION_INVALID_RRULE = lambda err='': f'Invalid RRULE: {err}' AUTOMATION_NO_FUTURE_RUNS = 'RRULE has no future occurrences' + AUTOMATION_RRULE_TIMEOUT = 'RRULE took too long to resolve' FEATURE_DISABLED = lambda name='': f'{name} is disabled' INPUT_TOO_LONG = lambda size='': f'Input prompt exceeds maximum length of {size}' diff --git a/backend/open_webui/models/automations.py b/backend/open_webui/models/automations.py index d3c0649455..2c94cba68b 100644 --- a/backend/open_webui/models/automations.py +++ b/backend/open_webui/models/automations.py @@ -1,3 +1,4 @@ +import asyncio import logging import time from typing import Optional @@ -258,14 +259,15 @@ class AutomationTable: self, id: str, next_run_at: Optional[int], + is_active: bool, db: Optional[AsyncSession] = None, ) -> Optional[AutomationModel]: async with get_async_db_context(db) as db: row = await db.get(Automation, id) if not row: return None - row.is_active = not row.is_active - row.next_run_at = next_run_at if row.is_active else None + row.is_active = is_active + row.next_run_at = next_run_at row.updated_at = int(time.time_ns()) await db.commit() return AutomationModel.model_validate(row) @@ -316,13 +318,25 @@ 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()} - for row in rows: + next_runs = await asyncio.gather( + *[next_run_ns(row.data.get('rrule', ''), tz=timezone_by_user_id.get(row.user_id)) for row in rows], + return_exceptions=True, + ) + + claimed = [] + for row, next_run_at in zip(rows, next_runs): + if isinstance(next_run_at, BaseException): + # Look again in five minutes rather than run it blind or re-walk it every poll + log.warning('Skipping automation %s: %s', row.id, next_run_at) + row.next_run_at = now_ns + 300 * 1_000_000_000 + 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/routers/automations.py b/backend/open_webui/routers/automations.py index e01c5c255b..c8021c9d70 100644 --- a/backend/open_webui/routers/automations.py +++ b/backend/open_webui/routers/automations.py @@ -85,7 +85,13 @@ 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) + try: + interval = await rrule_interval_seconds(rrule_str) + except ValueError as e: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=str(e), + ) if interval is not None and interval < min_interval: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -110,7 +116,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), ) @@ -174,8 +180,10 @@ async def create_new_automation( ): await check_automations_permission(request, user) await check_automation_folder_access(form_data.folder_id, user, db) + tz = user.timezone try: - validate_rrule(form_data.data.rrule, tz=user.timezone) + await validate_rrule(form_data.data.rrule, tz=tz) + next_run_at = await next_run_ns(form_data.data.rrule, tz=tz) except ValueError as e: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -184,8 +192,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, next_run_at, db=db) response = await enrich_automation(automation, db, tz=tz) await publish_event( request, @@ -233,8 +240,10 @@ async def update_automation_by_id( check_automation_access(automation, user) await check_automation_folder_access(form_data.folder_id, user, db) + tz = user.timezone try: - validate_rrule(form_data.data.rrule, tz=user.timezone) + await validate_rrule(form_data.data.rrule, tz=tz) + next_run_at = await next_run_ns(form_data.data.rrule, tz=tz) except ValueError as e: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -243,8 +252,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, next_run_at, db=db) response = await enrich_automation(updated, db, tz=tz) await publish_event( request, @@ -271,7 +279,15 @@ 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) + is_active = not automation.is_active + try: + next_run_at = await next_run_ns(automation.data['rrule'], tz=user.timezone) if is_active else None + except ValueError as e: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=str(e), + ) + toggled = await Automations.toggle(id, next_run_at, is_active, db=db) response = await enrich_automation(toggled, db, tz=user.timezone) await publish_event( request, diff --git a/backend/open_webui/tools/builtin.py b/backend/open_webui/tools/builtin.py index 136d2001dc..11c1fb2eb8 100644 --- a/backend/open_webui/tools/builtin.py +++ b/backend/open_webui/tools/builtin.py @@ -3627,7 +3627,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}'}) @@ -3648,7 +3648,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( { @@ -3658,7 +3658,7 @@ async def create_automation( 'folder_id': automation.folder_id, 'model_id': model_id, '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, ) @@ -3727,7 +3727,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}'}) @@ -3748,7 +3748,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( { @@ -3758,7 +3758,7 @@ async def update_automation( 'folder_id': updated.folder_id, 'model_id': new_model_id, '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, ) @@ -3825,7 +3825,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), } ) @@ -3870,9 +3870,11 @@ async def toggle_automation( return JSONCodec.dumps({'error': 'Access denied'}) rrule = automation.data.get('rrule', '') + is_active = not automation.is_active 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) if is_active else None, + is_active, ) return JSONCodec.dumps( diff --git a/backend/open_webui/utils/automations.py b/backend/open_webui/utils/automations.py index 712ea45096..9254515972 100644 --- a/backend/open_webui/utils/automations.py +++ b/backend/open_webui/utils/automations.py @@ -20,6 +20,7 @@ import logging import os import random import time +from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timedelta from typing import Optional from uuid import uuid4 @@ -117,7 +118,7 @@ def _parse_rule(s: str, now: Optional[datetime] = None): return rrulestr(s, ignoretz=True) -def validate_rrule(s: str, tz: str = None) -> None: +def _validate_rrule_blocking(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 @@ -130,11 +131,14 @@ def validate_rrule(s: str, tz: str = None) -> None: rule = _parse_rule(s, now) except Exception as e: raise ValueError(ERROR_MESSAGES.AUTOMATION_INVALID_RRULE(e)) - if rule.after(now) is None: + first = rule.after(now) + if first is None: raise ValueError(ERROR_MESSAGES.AUTOMATION_NO_FUTURE_RUNS) + # The scheduler needs the occurrence after that one, which costs more, so prove it is affordable too + rule.after(first) -def next_run_ns(s: str, tz: str = None) -> Optional[int]: +def _next_run_ns_blocking(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() @@ -147,7 +151,7 @@ def next_run_ns(s: str, tz: str = None) -> Optional[int]: return int(dt.timestamp() * 1_000_000_000) -def next_n_runs_ns(s: str, n: int = 5, tz: str = None) -> list[int]: +def _next_n_runs_ns_blocking(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 @@ -170,7 +174,7 @@ def next_n_runs_ns(s: str, n: int = 5, tz: str = None) -> list[int]: return result -def rrule_interval_seconds(s: str) -> Optional[int]: +def _rrule_interval_seconds_blocking(s: str) -> Optional[int]: """Approximate interval between recurrences in seconds. Returns None for one-shot (COUNT=1) schedules or rules @@ -189,6 +193,46 @@ def rrule_interval_seconds(s: str) -> Optional[int]: return int((second - first).total_seconds()) +# A walk already under way cannot be stopped, so a single thread caps what one can cost. +_rrule_pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix='rrule') + + +async def _run_bounded(fn, *args): + """Resolve occurrences off the event loop, raising ValueError on rules that take too long.""" + try: + return await asyncio.wait_for(asyncio.wrap_future(_rrule_pool.submit(fn, *args)), timeout=1) + except TimeoutError: + log.warning('Gave up resolving recurrence rule %.200r', args[0]) + raise ValueError(ERROR_MESSAGES.AUTOMATION_RRULE_TIMEOUT) from None + except ValueError: + raise + except Exception as e: + raise ValueError(ERROR_MESSAGES.AUTOMATION_INVALID_RRULE(e)) from None + + +async def validate_rrule(s: str, tz: str = None) -> None: + """Raise ValueError if the RRULE is malformed, exhausted, or too slow to resolve.""" + await _run_bounded(_validate_rrule_blocking, s, tz) + + +async def next_run_ns(s: str, tz: str = None) -> Optional[int]: + """Next occurrence as epoch nanoseconds. Raises ValueError if the rule cannot be resolved.""" + return await _run_bounded(_next_run_ns_blocking, s, tz) + + +async def next_n_runs_ns(s: str, n: int = 5, tz: str = None) -> list[int]: + """Next N occurrences for UI preview, empty when they cannot be resolved.""" + try: + return await _run_bounded(_next_n_runs_ns_blocking, s, n, tz) + except ValueError: + return [] + + +async def rrule_interval_seconds(s: str) -> Optional[int]: + """Approximate interval between recurrences. Raises ValueError if the rule cannot be resolved.""" + return await _run_bounded(_rrule_interval_seconds_blocking, s) + + ############################ # Worker Loop ############################