Coverage for open_webui/utils/timers.py: 20%
208 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-07 05:07 +0000
1"""Durable one-shot timers backed by internal child chats."""
3from __future__ import annotations
5import asyncio
6import copy
7import logging
8import re
9import time
10from datetime import datetime, timezone
11from typing import Literal
12from uuid import uuid4
14from fastapi import Request
15from open_webui.internal.db import get_async_db
16from open_webui.models.chat_messages import ChatMessages
17from open_webui.models.chats import Chat, ChatForm, Chats
18from open_webui.models.users import UserModel, Users
19from open_webui.tasks import has_active_tasks
20from open_webui.utils.auth import VERIFIED_USER_ROLES
21from open_webui.utils.json_codec import JSONCodec
22from open_webui.utils.misc import get_message_list
23from sqlalchemy import select
24from starlette.datastructures import Headers
26log = logging.getLogger(__name__)
28_RELATIVE_TIME = re.compile(r'^(?:\+|in\s+)?(\d+)\s*(s|sec(?:onds?)?|m|min(?:utes?)?|h|hours?|d|days?)$')
29_RFC3339_TIME = re.compile(r'^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(?:\.\d+)?(?:Z|[+-]\d{2}:\d{2})$')
30_TIME_UNITS_NS = {
31 's': 1_000_000_000,
32 'm': 60 * 1_000_000_000,
33 'h': 60 * 60 * 1_000_000_000,
34 'd': 24 * 60 * 60 * 1_000_000_000,
35}
36_timer_locks: dict[str, asyncio.Lock] = {}
39def parse_timer_at(value: str) -> int:
40 """Normalize a relative offset or timezone-aware RFC 3339 timestamp."""
41 raw = value.strip()
42 now = time.time_ns()
43 relative = _RELATIVE_TIME.fullmatch(raw.lower())
44 if relative:
45 count = int(relative.group(1))
46 if count <= 0:
47 raise ValueError('at must be in the future.')
48 return now + count * _TIME_UNITS_NS[relative.group(2)[0]]
50 if not _RFC3339_TIME.fullmatch(raw):
51 raise ValueError(
52 'at must be a relative time such as 10s or in 10 seconds, or an RFC 3339 timestamp with a timezone.'
53 )
54 try:
55 parsed = datetime.fromisoformat(raw.replace('Z', '+00:00'))
56 except ValueError as exc:
57 raise ValueError(
58 'at must be a relative time such as 10s or in 10 seconds, or an RFC 3339 timestamp with a timezone.'
59 ) from exc
60 if parsed.tzinfo is None or parsed.utcoffset() is None:
61 raise ValueError('absolute at values must include an explicit timezone.')
63 due_at = int(parsed.timestamp() * 1_000_000_000)
64 if due_at <= now:
65 raise ValueError('at must be in the future.')
66 return due_at
69async def create_timer(
70 *,
71 prompt: str,
72 at: str,
73 cancel_on: list[Literal['chat.read', 'chat.user_message']] | None,
74 request: Request,
75 user_data: dict,
76 metadata: dict,
77 parent_chat_id: str,
78 parent_message_id: str | None,
79) -> str:
80 prompt = prompt.strip()
81 if not prompt:
82 return 'Error: prompt must not be empty.'
83 if not parent_chat_id or not user_data.get('id'):
84 return 'Error: chat and user context are required.'
86 try:
87 due_at = parse_timer_at(at)
88 except ValueError as exc:
89 return f'Error: {exc}'
91 selected_events = cancel_on or []
92 allowed_events = {'chat.read', 'chat.user_message'}
93 if any(event not in allowed_events for event in selected_events):
94 return 'Error: cancel_on accepts only chat.read and chat.user_message.'
95 selected_events = list(dict.fromkeys(selected_events))
97 model_id = metadata.get('model_id') or (metadata.get('model') or {}).get('id')
98 if not model_id:
99 return 'Error: model context is required.'
100 if metadata.get('direct'):
101 return 'Error: timers are unavailable for direct connections.'
103 chat_id = str(uuid4())
104 user_message_id = str(uuid4())
105 user = UserModel(**user_data)
106 run = {
107 'model_id': model_id,
108 'session_id': metadata.get('session_id'),
109 'tool_ids': copy.deepcopy(metadata.get('tool_ids') or []),
110 'skill_ids': copy.deepcopy(metadata.get('skill_ids') or []),
111 'system_prompt': metadata.get('system_prompt'),
112 'filter_ids': copy.deepcopy(metadata.get('filter_ids') or []),
113 'terminal_id': metadata.get('terminal_id'),
114 'features': copy.deepcopy(metadata.get('features') or {}),
115 'files': copy.deepcopy(metadata.get('files') or []),
116 'variables': copy.deepcopy(metadata.get('variables') or {}),
117 }
119 chat = await Chats.insert_new_chat(
120 chat_id,
121 user.id,
122 ChatForm(
123 chat={
124 'id': chat_id,
125 'title': f'Timer: {prompt[:60]}',
126 'models': [model_id],
127 'history': {
128 'currentId': user_message_id,
129 'messages': {
130 user_message_id: {
131 'id': user_message_id,
132 'parentId': None,
133 'childrenIds': [],
134 'role': 'user',
135 'content': prompt,
136 'timestamp': int(time.time()),
137 'models': [model_id],
138 },
139 },
140 },
141 'messages': [{'role': 'user', 'content': prompt}],
142 }
143 ),
144 internal_meta={
145 'internal': True,
146 'type': 'timer',
147 'parent_chat_id': parent_chat_id,
148 'parent_message_id': parent_message_id,
149 'timer_at': due_at,
150 'status': 'pending',
151 'timer_model_id': model_id,
152 'timer_task_message_id': user_message_id,
153 'cancel_on': selected_events,
154 'run': run,
155 },
156 timer_at=due_at,
157 )
158 if not chat:
159 return 'Error: failed to create timer.'
161 return JSONCodec.dumps(
162 {
163 'status': 'set',
164 'at': datetime.fromtimestamp(due_at / 1_000_000_000, timezone.utc).isoformat().replace('+00:00', 'Z'),
165 'cancel_on': selected_events,
166 },
167 ensure_ascii=False,
168 )
171async def claim_due_timers(now_ns: int, limit: int = 10) -> list[tuple[str, str]]:
172 """Claim due timers by moving them from pending to running."""
173 async with get_async_db() as db:
174 stmt = (
175 select(Chat)
176 .where(Chat.timer_at <= now_ns)
177 .where(Chat.meta['status'].as_string() == 'pending')
178 .order_by(Chat.timer_at)
179 .limit(limit)
180 )
181 if db.bind.dialect.name == 'postgresql': 181 ↛ 182line 181 didn't jump to line 182 because the condition on line 181 was never true
182 stmt = stmt.with_for_update(skip_locked=True)
184 result = await db.execute(stmt)
186 claimed = []
187 for row in result.scalars().all(): 187 ↛ 188line 187 didn't jump to line 188 because the loop on line 187 never started
188 claim_id = str(uuid4())
189 row.meta = {
190 **(row.meta or {}),
191 'status': 'running',
192 'timer_started_at': now_ns,
193 'timer_claim_id': claim_id,
194 }
195 row.timer_at = None
196 row.updated_at = int(time.time())
197 claimed.append((row.id, claim_id))
198 await db.commit()
199 return claimed
202async def cancel_timers_for_chat(
203 parent_chat_id: str, event: Literal['chat.read', 'chat.user_message'], user_id: str
204) -> None:
205 """Owner-scoped: without the user filter, reading a chat cancels every user's timers on it."""
206 async with get_async_db() as db:
207 result = await db.execute(
208 select(Chat)
209 .where(Chat.timer_at.isnot(None))
210 .where(Chat.user_id == user_id)
211 .where(Chat.meta['parent_chat_id'].as_string() == parent_chat_id)
212 .where(Chat.meta['status'].as_string() == 'pending')
213 )
214 now_ns = int(time.time_ns())
215 for row in result.scalars().all():
216 meta = row.meta or {}
217 if event not in (meta.get('cancel_on') or []):
218 continue
219 row.meta = {
220 **meta,
221 'status': 'cancelled',
222 'timer_cancelled_at': now_ns,
223 'timer_cancelled_by': event,
224 }
225 row.timer_at = None
226 row.updated_at = int(time.time())
227 await db.commit()
230async def execute_due_timer(app, timer_id: str, claim_id: str | None = None) -> None:
231 lock = _timer_locks.setdefault(timer_id, asyncio.Lock())
232 async with lock:
233 from open_webui.socket.main import sio
234 from open_webui.utils.subagents import _parent_locks
236 timer = await Chats.get_chat_by_id(timer_id)
237 if not timer:
238 return
239 meta = timer.meta or {}
240 if meta.get('status') != 'running':
241 return
242 if claim_id is not None and meta.get('timer_claim_id') != claim_id:
243 return
245 parent_chat_id = meta.get('parent_chat_id') or ''
246 parent = await Chats.get_chat_by_id_and_user_id(parent_chat_id, timer.user_id)
247 if not parent:
248 await _set_timer_state(timer_id, 'error', timer_error='parent chat no longer exists')
249 return
251 prompt_message_id = meta.get('timer_task_message_id')
252 prompt_message = await Chats.get_message_by_id_and_message_id(timer_id, prompt_message_id)
253 if not prompt_message:
254 await _set_timer_state(timer_id, 'error', timer_error='timer task message is missing')
255 return
257 user = await Users.get_user_by_id(timer.user_id)
258 if not user:
259 await _set_timer_state(timer_id, 'error', timer_error='timer user no longer exists')
260 return
262 # Re-gate the rehydrated owner: a demoted owner must not run.
263 if user.role not in VERIFIED_USER_ROLES:
264 await _set_timer_state(timer_id, 'error', timer_error='owner no longer permitted to run timers')
265 return
267 run = meta.get('run') or {}
268 model_id = run.get('model_id') or meta.get('timer_model_id')
269 if not model_id:
270 await _set_timer_state(timer_id, 'error', timer_error='model context is missing')
271 return
273 prompt = prompt_message.get('content') or ''
274 if isinstance(prompt, list):
275 prompt = ''.join(
276 str(part.get('text', '')) for part in prompt if isinstance(part, dict) and part.get('type') == 'text'
277 )
279 user_message_id = str(uuid4())
280 assistant_message_id = str(uuid4())
281 user_message = None
282 assistant_message = None
283 message_list = []
284 parent_lock = _parent_locks.setdefault(parent_chat_id, asyncio.Lock())
285 async with parent_lock:
286 async with get_async_db() as db:
287 stmt = select(Chat).where(Chat.id == parent_chat_id, Chat.user_id == timer.user_id)
288 if db.bind.dialect.name == 'postgresql':
289 stmt = stmt.with_for_update()
290 result = await db.execute(stmt)
291 parent = result.scalar_one_or_none()
292 if not parent:
293 await _set_timer_state(timer_id, 'error', timer_error='parent chat no longer exists')
294 return
295 if await has_active_tasks(app.state.redis, parent_chat_id):
296 timer_row = await db.get(Chat, timer_id)
297 if timer_row: 297 ↛ anywhereline 297 didn't jump anywhere: it always raised an exception.
298 timer_meta = timer_row.meta or {}
299 timer_row.meta = {
300 **timer_meta,
301 'status': 'pending',
302 'timer_claim_id': None,
303 'timer_started_at': None,
304 }
305 timer_row.timer_at = timer_meta.get('timer_at')
306 timer_row.updated_at = int(time.time())
307 await db.commit()
308 return
310 parent_chat = copy.deepcopy(parent.chat or {})
311 history = parent_chat.setdefault('history', {})
312 messages = history.setdefault('messages', {})
313 done_assistants = [
314 message
315 for message in messages.values()
316 if message.get('role') == 'assistant' and message.get('done') is not False
317 ]
318 parent_id = (
319 max(done_assistants, key=lambda message: message.get('timestamp', 0)).get('id')
320 if done_assistants
321 else meta.get('parent_message_id')
322 )
323 message_list = get_message_list(messages, parent_id)
325 user_message = {
326 'id': user_message_id,
327 'parentId': parent_id,
328 'childrenIds': [assistant_message_id],
329 'role': 'user',
330 'content': prompt,
331 'model': model_id,
332 'meta': {'internal': True, 'type': 'timer', 'timer_id': timer_id},
333 'timestamp': int(time.time()),
334 }
335 assistant_message = {
336 'id': assistant_message_id,
337 'parentId': user_message_id,
338 'childrenIds': [],
339 'role': 'assistant',
340 'content': '',
341 'done': False,
342 'model': model_id,
343 'timestamp': int(time.time()),
344 }
345 messages[user_message_id] = user_message
346 messages[assistant_message_id] = assistant_message
347 if parent_id and parent_id in messages:
348 children = messages[parent_id].setdefault('childrenIds', [])
349 if user_message_id not in children:
350 children.append(user_message_id)
352 parent.chat = parent_chat
353 history['currentId'] = assistant_message_id
354 parent.updated_at = int(time.time())
355 timer_row = await db.get(Chat, timer_id)
356 if timer_row:
357 timer_row.meta = {
358 **(timer_row.meta or {}),
359 'status': 'completed',
360 'timer_completed_at': int(time.time_ns()),
361 }
362 timer_row.updated_at = int(time.time())
363 await db.commit()
364 await ChatMessages.upsert_message(user_message_id, parent_chat_id, timer.user_id, user_message)
365 await ChatMessages.upsert_message(assistant_message_id, parent_chat_id, timer.user_id, assistant_message)
367 await sio.emit(
368 'events',
369 {
370 'chat_id': parent_chat_id,
371 'message_id': assistant_message_id,
372 'data': {'type': 'chat:reload'},
373 },
374 room=f'user:{timer.user_id}',
375 )
376 form_data = {
377 'model': model_id,
378 'messages': [
379 *([{'role': 'system', 'content': run.get('system_prompt')}] if run.get('system_prompt') else []),
380 *message_list,
381 {'role': 'user', 'content': prompt},
382 ],
383 'stream': True,
384 'chat_id': parent_chat_id,
385 'id': assistant_message_id,
386 'parent_id': user_message.get('parentId'),
387 'user_message': user_message,
388 'session_id': run.get('session_id') or f'timer:{parent_chat_id}',
389 'background_tasks': {},
390 'tool_ids': run.get('tool_ids') or [],
391 'skill_ids': run.get('skill_ids') or [],
392 'filter_ids': run.get('filter_ids') or [],
393 'features': run.get('features') or {},
394 'files': run.get('files') or [],
395 'variables': run.get('variables') or {},
396 }
397 if run.get('terminal_id'):
398 form_data['terminal_id'] = run['terminal_id']
399 request = Request(
400 {
401 'type': 'http',
402 'asgi': {'version': '3.0', 'spec_version': '2.0'},
403 'method': 'POST',
404 'path': '/api/v1/timers/internal',
405 'query_string': b'',
406 'headers': Headers({}).raw,
407 'client': ('127.0.0.1', 0),
408 'server': ('127.0.0.1', 80),
409 'scheme': 'http',
410 'app': app,
411 }
412 )
413 request.state.token = None
414 request.state.enable_api_keys = False
415 try:
416 await app.state.CHAT_COMPLETION_HANDLER(request, form_data, user=user)
417 except Exception as exc:
418 log.exception(f'Timer {timer_id} completion failed')
419 await _set_timer_state(timer_id, 'error', timer_error=str(exc)[:500])
422async def _set_timer_state(timer_id: str, status: str, **fields) -> None:
423 async with get_async_db() as db:
424 row = await db.get(Chat, timer_id)
425 if not row:
426 return
427 row.meta = {**(row.meta or {}), 'status': status, **fields}
428 row.updated_at = int(time.time())
429 await db.commit()