diff --git a/backend/open_webui/utils/automations.py b/backend/open_webui/utils/automations.py index 4f2ec3b2c9..7823c0207b 100644 --- a/backend/open_webui/utils/automations.py +++ b/backend/open_webui/utils/automations.py @@ -11,6 +11,7 @@ The scheduler_worker_loop handles all time-based background work: Environment: SCHEDULER_POLL_INTERVAL – seconds between polls (default: 10) + TIMER_POLL_INTERVAL – seconds between timer polls (default: 1) CALENDAR_ALERT_LOOKAHEAD_MINUTES – default alert window (default: 5) """ @@ -43,6 +44,7 @@ from starlette.datastructures import Headers log = logging.getLogger(__name__) SCHEDULER_POLL_INTERVAL = int(os.getenv('SCHEDULER_POLL_INTERVAL', os.getenv('AUTOMATION_POLL_INTERVAL', '10'))) +TIMER_POLL_INTERVAL = int(os.getenv('TIMER_POLL_INTERVAL', '1')) CALENDAR_ALERT_LOOKAHEAD_MINUTES = int(os.getenv('CALENDAR_ALERT_LOOKAHEAD_MINUTES', '10')) @@ -190,10 +192,15 @@ async def scheduler_worker_loop(app) -> None: Runs on every instance. Poll interval is configurable via SCHEDULER_POLL_INTERVAL env var (default: 10 seconds). """ - log.info(f'Scheduler worker started (poll interval: {SCHEDULER_POLL_INTERVAL}s)') + log.info( + f'Scheduler worker started (timer poll interval: {TIMER_POLL_INTERVAL}s, ' + f'scheduler poll interval: {SCHEDULER_POLL_INTERVAL}s)' + ) + next_scheduler_poll = 0.0 while True: try: + now = time.monotonic() # ── Timers ── try: from open_webui.utils.timers import claim_due_timers, execute_due_timer @@ -203,6 +210,12 @@ async def scheduler_worker_loop(app) -> None: except Exception: log.exception('Scheduler: timer error') + if now < next_scheduler_poll: + await asyncio.sleep(max(1, TIMER_POLL_INTERVAL)) + continue + # Jitter to spread automation/calendar load across instances; timers keep a tight poll. + next_scheduler_poll = now + SCHEDULER_POLL_INTERVAL + random.uniform(0, 2) + # ── Automations ── if await Config.get('automations.enable'): try: @@ -225,8 +238,7 @@ async def scheduler_worker_loop(app) -> None: except Exception: log.exception('Scheduler worker error') - # Jitter to spread load across instances - await asyncio.sleep(SCHEDULER_POLL_INTERVAL + random.uniform(0, 2)) + await asyncio.sleep(max(1, TIMER_POLL_INTERVAL)) ########################## diff --git a/backend/open_webui/utils/subagents.py b/backend/open_webui/utils/subagents.py index 5b0475fc6f..efb562f8cc 100644 --- a/backend/open_webui/utils/subagents.py +++ b/backend/open_webui/utils/subagents.py @@ -9,13 +9,15 @@ from uuid import uuid4 from fastapi import Request from fastapi.security import HTTPAuthorizationCredentials +from open_webui.internal.db import get_async_db from open_webui.models.chat_messages import ChatMessages -from open_webui.models.chats import ChatForm, Chats +from open_webui.models.chats import Chat, ChatForm, Chats from open_webui.models.config import Config from open_webui.models.users import UserModel, Users from open_webui.tasks import create_task, has_active_tasks from open_webui.utils.auth import create_token from open_webui.utils.misc import get_message_list +from sqlalchemy import select from starlette.datastructures import Headers DEFAULT_SUBAGENT_SYSTEM_PROMPT = """You are a sub-agent working on a specific task assigned by the lead agent. @@ -81,145 +83,143 @@ async def process_pending_internal_messages( if await has_active_tasks(source_request.app.state.redis, parent_chat_id): return - chat = await Chats.get_chat_by_id_and_user_id(parent_chat_id, user_id) user = await Users.get_user_by_id(user_id) - if not chat or not user: + if not user: return - history = copy.deepcopy(chat.chat.get('history') or {}) - messages = history.get('messages') or {} - pending = [ - message - for message in messages.values() - if message.get('role') == 'user' - and not message.get('childrenIds') - and ( - (message.get('meta') or {}).get('async_subagent_result') is True - or ( - (message.get('meta') or {}).get('internal') is True - and (message.get('meta') or {}).get('type') == 'timer' + async with get_async_db() as db: + stmt = select(Chat).where(Chat.id == parent_chat_id, Chat.user_id == user_id) + if db.bind.dialect.name == 'postgresql': + stmt = stmt.with_for_update() + result = await db.execute(stmt) + chat = result.scalar_one_or_none() + if not chat: + return + + history = copy.deepcopy((chat.chat or {}).get('history') or {}) + messages = history.get('messages') or {} + pending = [ + message + for message in messages.values() + if message.get('role') == 'user' + and not message.get('childrenIds') + and ( + (message.get('meta') or {}).get('async_subagent_result') is True + or ( + (message.get('meta') or {}).get('internal') is True + and (message.get('meta') or {}).get('type') == 'timer' + ) ) - ) - ] - if not pending: - return + ] + if not pending: + return - first = pending[0] - first_meta = first.get('meta') or {} - kind = 'subagent' if first_meta.get('async_subagent_result') is True else 'timer' - parent_id = first.get('parentId') - if kind == 'timer' and first_meta.get('timer_id'): - timer = await Chats.get_chat_by_id(first_meta['timer_id']) - run = {**run, **(((timer.meta or {}).get('run') if timer else None) or {})} - model_id = first.get('model') or run['model_id'] - if kind == 'timer': - batch = [ - message - for message in pending - if message.get('parentId') == parent_id - and (message.get('model') or model_id) == model_id - and (message.get('meta') or {}).get('internal') is True - and (message.get('meta') or {}).get('type') == 'timer' - ] - else: - batch = [ - message - for message in pending - if message.get('parentId') == parent_id - and (message.get('model') or model_id) == model_id - and (message.get('meta') or {}).get('async_subagent_result') is True - ] - combined_content = '\n\n'.join(message.get('content', '') for message in batch if message.get('content')) - if kind == 'timer': - timer_ids = [message['meta']['timer_id'] for message in batch if (message.get('meta') or {}).get('timer_id')] - combined_meta = {'internal': True, 'type': 'timer'} - if len(timer_ids) == 1: - combined_meta['timer_id'] = timer_ids[0] - elif timer_ids: - combined_meta['timer_ids'] = timer_ids - else: - delegation_ids = [ - message['meta']['delegation_id'] - for message in batch - if (message.get('meta') or {}).get('delegation_id') - ] - subagent_chat_ids = [ - message['meta']['subagent_chat_id'] - for message in batch - if (message.get('meta') or {}).get('subagent_chat_id') - ] - combined_meta = {'async_subagent_result': True} - if len(delegation_ids) == 1: - combined_meta['delegation_id'] = delegation_ids[0] - elif delegation_ids: - combined_meta['delegation_ids'] = delegation_ids - if len(subagent_chat_ids) == 1: - combined_meta['subagent_chat_id'] = subagent_chat_ids[0] - elif subagent_chat_ids: - combined_meta['subagent_chat_ids'] = subagent_chat_ids - - pending_flag = 'timer_pending' if kind == 'timer' else 'async_subagent_pending' - reuse_message = len(batch) == 1 and not (first.get('meta') or {}).get(pending_flag) - user_message_id = first['id'] if reuse_message else str(uuid4()) - if not reuse_message: - removed_ids = {message['id'] for message in batch} - for message_id in removed_ids: - messages.pop(message_id, None) - if parent_id and parent_id in messages: - messages[parent_id]['childrenIds'] = [ - child_id for child_id in messages[parent_id].get('childrenIds', []) if child_id not in removed_ids + first = pending[0] + first_meta = first.get('meta') or {} + kind = 'subagent' if first_meta.get('async_subagent_result') is True else 'timer' + parent_id = first.get('parentId') + if kind == 'timer' and first_meta.get('timer_id'): + timer = await Chats.get_chat_by_id(first_meta['timer_id']) + run = {**run, **(((timer.meta or {}).get('run') if timer else None) or {})} + model_id = first.get('model') or run['model_id'] + if kind == 'timer': + batch = [first] + else: + batch = [ + message + for message in pending + if message.get('parentId') == parent_id + and (message.get('model') or model_id) == model_id + and (message.get('meta') or {}).get('async_subagent_result') is True ] + combined_content = '\n\n'.join(message.get('content', '') for message in batch if message.get('content')) + if kind == 'timer': + timer_ids = [ + message['meta']['timer_id'] for message in batch if (message.get('meta') or {}).get('timer_id') + ] + combined_meta = {'internal': True, 'type': 'timer'} + if len(timer_ids) == 1: + combined_meta['timer_id'] = timer_ids[0] + elif timer_ids: + combined_meta['timer_ids'] = timer_ids + else: + delegation_ids = [ + message['meta']['delegation_id'] + for message in batch + if (message.get('meta') or {}).get('delegation_id') + ] + subagent_chat_ids = [ + message['meta']['subagent_chat_id'] + for message in batch + if (message.get('meta') or {}).get('subagent_chat_id') + ] + combined_meta = {'async_subagent_result': True} + if len(delegation_ids) == 1: + combined_meta['delegation_id'] = delegation_ids[0] + elif delegation_ids: + combined_meta['delegation_ids'] = delegation_ids + if len(subagent_chat_ids) == 1: + combined_meta['subagent_chat_id'] = subagent_chat_ids[0] + elif subagent_chat_ids: + combined_meta['subagent_chat_ids'] = subagent_chat_ids + + pending_flag = 'timer_pending' if kind == 'timer' else 'async_subagent_pending' + reuse_message = len(batch) == 1 and not (first.get('meta') or {}).get(pending_flag) + user_message_id = first['id'] if reuse_message else str(uuid4()) + removed_ids = set() + if not reuse_message: + removed_ids = {message['id'] for message in batch} + for message_id in removed_ids: + messages.pop(message_id, None) + if parent_id and parent_id in messages: + messages[parent_id]['childrenIds'] = [ + child_id + for child_id in messages[parent_id].get('childrenIds', []) + if child_id not in removed_ids + ] + + assistant_message_id = str(uuid4()) + message_list = get_message_list(messages, parent_id) + system_prompt = run.get('system_prompt') + user_message = { + 'id': user_message_id, + 'parentId': parent_id, + 'childrenIds': [assistant_message_id], + 'role': 'user', + 'content': combined_content, + 'model': model_id, + 'meta': combined_meta, + 'timestamp': int(time.time()), + } + assistant_message = { + 'id': assistant_message_id, + 'parentId': user_message_id, + 'childrenIds': [], + 'role': 'assistant', + 'content': '', + 'done': False, + 'model': model_id, + 'timestamp': int(time.time()), + } + + if parent_id and parent_id in messages: + parent_children = [ + child_id for child_id in messages[parent_id].get('childrenIds', []) if child_id != user_message_id + ] + parent_children.append(user_message_id) + messages[parent_id]['childrenIds'] = parent_children + messages[user_message_id] = {**messages.get(user_message_id, {}), **user_message} + messages[assistant_message_id] = assistant_message history['messages'] = messages - history['currentId'] = parent_id - updated_chat = copy.deepcopy(chat.chat) - updated_chat['history'] = history - await Chats.update_chat_by_id(parent_chat_id, updated_chat) + history['currentId'] = assistant_message_id + chat.chat = {**(chat.chat or {}), 'history': history} + chat.updated_at = int(time.time()) + await db.commit() + + if removed_ids: await ChatMessages.delete_message_ids_by_chat_id(parent_chat_id, removed_ids) - - assistant_message_id = str(uuid4()) - message_list = get_message_list(messages, parent_id) - system_prompt = run.get('system_prompt') - user_message = { - 'id': user_message_id, - 'parentId': parent_id, - 'childrenIds': [assistant_message_id], - 'role': 'user', - 'content': combined_content, - 'model': model_id, - 'meta': combined_meta, - 'timestamp': int(time.time()), - } - assistant_message = { - 'id': assistant_message_id, - 'parentId': user_message_id, - 'childrenIds': [], - 'role': 'assistant', - 'content': '', - 'done': False, - 'model': model_id, - 'timestamp': int(time.time()), - } - - if parent_id and parent_id in messages: - parent_children = [ - child_id for child_id in messages[parent_id].get('childrenIds', []) if child_id != user_message_id - ] - parent_children.append(user_message_id) - await Chats.upsert_message_to_chat_by_id_and_message_id( - parent_chat_id, - parent_id, - {'childrenIds': parent_children}, - ) - await Chats.upsert_message_to_chat_by_id_and_message_id( - parent_chat_id, - user_message_id, - user_message, - ) - await Chats.upsert_message_to_chat_by_id_and_message_id( - parent_chat_id, - assistant_message_id, - assistant_message, - ) + await ChatMessages.upsert_message(user_message_id, parent_chat_id, user_id, user_message) + await ChatMessages.upsert_message(assistant_message_id, parent_chat_id, user_id, assistant_message) from open_webui.socket.main import sio diff --git a/backend/open_webui/utils/timers.py b/backend/open_webui/utils/timers.py index d33d063382..bd14841350 100644 --- a/backend/open_webui/utils/timers.py +++ b/backend/open_webui/utils/timers.py @@ -275,45 +275,61 @@ async def execute_due_timer(app, timer_id: str, claim_id: str | None = None) -> user_message_id = str(uuid4()) parent_lock = _parent_locks.setdefault(parent_chat_id, asyncio.Lock()) async with parent_lock: - parent = await Chats.get_chat_by_id_and_user_id(parent_chat_id, timer.user_id) - if not parent: - await _set_timer_status(timer_id, 'error', timer_error='parent chat no longer exists') - return - parent_chat = copy.deepcopy(parent.chat or {}) - history = parent_chat.setdefault('history', {}) - messages = history.setdefault('messages', {}) - done_assistants = [ - message - for message in messages.values() - if message.get('role') == 'assistant' and message.get('done') is not False - ] - parent_id = ( - max(done_assistants, key=lambda message: message.get('timestamp', 0)).get('id') - if done_assistants - else meta.get('parent_message_id') - ) - pending_meta = {'internal': True, 'type': 'timer', 'timer_id': timer_id} - if await has_active_tasks(app.state.redis, parent_chat_id): - pending_meta['timer_pending'] = True + async with get_async_db() as db: + stmt = select(Chat).where(Chat.id == parent_chat_id, Chat.user_id == timer.user_id) + if db.bind.dialect.name == 'postgresql': + stmt = stmt.with_for_update() + result = await db.execute(stmt) + parent = result.scalar_one_or_none() + if not parent: + await _set_timer_status(timer_id, 'error', timer_error='parent chat no longer exists') + return - user_message = { - 'id': user_message_id, - 'parentId': parent_id, - 'childrenIds': [], - 'role': 'user', - 'content': prompt, - 'model': model_id, - 'meta': pending_meta, - 'timestamp': int(time.time()), - } - messages[user_message_id] = user_message - if parent_id and parent_id in messages: - children = messages[parent_id].setdefault('childrenIds', []) - if user_message_id not in children: - children.append(user_message_id) - await Chats.update_chat_by_id(parent_chat_id, parent_chat) + parent_chat = copy.deepcopy(parent.chat or {}) + history = parent_chat.setdefault('history', {}) + messages = history.setdefault('messages', {}) + done_assistants = [ + message + for message in messages.values() + if message.get('role') == 'assistant' and message.get('done') is not False + ] + parent_id = ( + max(done_assistants, key=lambda message: message.get('timestamp', 0)).get('id') + if done_assistants + else meta.get('parent_message_id') + ) + pending_meta = {'internal': True, 'type': 'timer', 'timer_id': timer_id} + if await has_active_tasks(app.state.redis, parent_chat_id): + pending_meta['timer_pending'] = True + + user_message = { + 'id': user_message_id, + 'parentId': parent_id, + 'childrenIds': [], + 'role': 'user', + 'content': prompt, + 'model': model_id, + 'meta': pending_meta, + 'timestamp': int(time.time()), + } + messages[user_message_id] = user_message + if parent_id and parent_id in messages: + children = messages[parent_id].setdefault('childrenIds', []) + if user_message_id not in children: + children.append(user_message_id) + + parent.chat = parent_chat + parent.updated_at = int(time.time()) + timer_row = await db.get(Chat, timer_id) + if timer_row: + timer_row.meta = { + **(timer_row.meta or {}), + 'timer_status': 'dispatched', + 'timer_dispatched_at': int(time.time_ns()), + } + timer_row.updated_at = int(time.time()) + await db.commit() await ChatMessages.upsert_message(user_message_id, parent_chat_id, timer.user_id, user_message) - await _set_timer_status(timer_id, 'dispatched', timer_dispatched_at=int(time.time_ns())) if user_message['meta'].get('timer_pending') is True: await sio.emit(