mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-06-02 12:05:54 +00:00
0154ea6cd3
* fix: add storage retention cleanup * fix: prune completed tasks on completion * fix: complete storage analysis i18n
284 lines
8.0 KiB
Python
284 lines
8.0 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import typing
|
|
import datetime
|
|
import time
|
|
|
|
from . import app
|
|
from . import entities as core_entities
|
|
|
|
|
|
class TaskContext:
|
|
"""Task tracking context"""
|
|
|
|
current_action: str
|
|
"""Current action being executed"""
|
|
|
|
log: str
|
|
"""Log"""
|
|
|
|
metadata: dict
|
|
"""Structured metadata for progress reporting"""
|
|
|
|
def __init__(self):
|
|
self.current_action = 'default'
|
|
self.log = ''
|
|
self.metadata = {}
|
|
|
|
def _log(self, msg: str):
|
|
self.log += msg + '\n'
|
|
|
|
def set_current_action(self, action: str):
|
|
self.current_action = action
|
|
|
|
def trace(
|
|
self,
|
|
msg: str,
|
|
action: str = None,
|
|
):
|
|
if action is not None:
|
|
self.set_current_action(action)
|
|
|
|
self._log(f'{datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")} | {self.current_action} | {msg}')
|
|
|
|
def to_dict(self) -> dict:
|
|
return {'current_action': self.current_action, 'log': self.log, 'metadata': self.metadata}
|
|
|
|
@staticmethod
|
|
def new() -> TaskContext:
|
|
return TaskContext()
|
|
|
|
@staticmethod
|
|
def placeholder() -> TaskContext:
|
|
global placeholder_context
|
|
|
|
if placeholder_context is None:
|
|
placeholder_context = TaskContext()
|
|
|
|
return placeholder_context
|
|
|
|
|
|
placeholder_context: TaskContext | None = None
|
|
|
|
|
|
class TaskWrapper:
|
|
"""Task wrapper"""
|
|
|
|
_id_index: int = 0
|
|
"""Task ID index"""
|
|
|
|
id: int
|
|
"""Task ID"""
|
|
|
|
task_type: str = 'system' # Task type: system or user
|
|
"""Task type"""
|
|
|
|
kind: str = 'system_task' # Task type determined by the initiator, usually the same task type
|
|
"""Task type"""
|
|
|
|
name: str = ''
|
|
"""Task unique name"""
|
|
|
|
label: str = ''
|
|
"""Task display name"""
|
|
|
|
task_context: TaskContext
|
|
"""Task context"""
|
|
|
|
task: asyncio.Task
|
|
"""Task"""
|
|
|
|
task_stack: list = None
|
|
"""Task stack"""
|
|
|
|
ap: app.Application
|
|
"""Application instance"""
|
|
|
|
scopes: list[core_entities.LifecycleControlScope]
|
|
"""Task scope"""
|
|
|
|
def __init__(
|
|
self,
|
|
ap: app.Application,
|
|
coro: typing.Coroutine,
|
|
task_type: str = 'system',
|
|
kind: str = 'system_task',
|
|
name: str = '',
|
|
label: str = '',
|
|
context: TaskContext = None,
|
|
scopes: list[core_entities.LifecycleControlScope] = [core_entities.LifecycleControlScope.APPLICATION],
|
|
):
|
|
self.id = TaskWrapper._id_index
|
|
TaskWrapper._id_index += 1
|
|
self.ap = ap
|
|
self.task_context = context or TaskContext()
|
|
self.task = self.ap.event_loop.create_task(coro)
|
|
self.task_type = task_type
|
|
self.kind = kind
|
|
self.name = name
|
|
self.label = label if label != '' else name
|
|
self.task.set_name(name)
|
|
self.scopes = scopes
|
|
self.created_at = time.time()
|
|
|
|
def assume_exception(self):
|
|
try:
|
|
exception = self.task.exception()
|
|
if self.task_stack is None:
|
|
self.task_stack = self.task.get_stack()
|
|
return exception
|
|
except Exception:
|
|
return None
|
|
|
|
def assume_result(self):
|
|
try:
|
|
return self.task.result()
|
|
except Exception:
|
|
return None
|
|
|
|
def to_dict(self) -> dict:
|
|
exception_traceback = None
|
|
if self.assume_exception() is not None:
|
|
exception_traceback = 'Traceback (most recent call last):\n'
|
|
|
|
for frame in self.task_stack:
|
|
exception_traceback += (
|
|
f' File "{frame.f_code.co_filename}", line {frame.f_lineno}, in {frame.f_code.co_name}\n'
|
|
)
|
|
|
|
exception_traceback += f' {self.assume_exception().__str__()}\n'
|
|
|
|
return {
|
|
'id': self.id,
|
|
'task_type': self.task_type,
|
|
'kind': self.kind,
|
|
'name': self.name,
|
|
'label': self.label,
|
|
'scopes': [scope.value for scope in self.scopes],
|
|
'created_at': self.created_at,
|
|
'task_context': self.task_context.to_dict(),
|
|
'runtime': {
|
|
'done': self.task.done(),
|
|
'state': self.task._state,
|
|
'exception': self.assume_exception().__str__() if self.assume_exception() is not None else None,
|
|
'exception_traceback': exception_traceback,
|
|
'result': self.assume_result() if self.assume_result() is not None else None,
|
|
},
|
|
}
|
|
|
|
def cancel(self):
|
|
self.task.cancel()
|
|
|
|
|
|
class AsyncTaskManager:
|
|
"""Save all asynchronous tasks in the app
|
|
Include system-level and user-level (plugin installation, update, etc. initiated by users directly)"""
|
|
|
|
ap: app.Application
|
|
|
|
tasks: list[TaskWrapper]
|
|
"""All tasks"""
|
|
|
|
def __init__(self, ap: app.Application):
|
|
self.ap = ap
|
|
self.tasks = []
|
|
|
|
def create_task(
|
|
self,
|
|
coro: typing.Coroutine,
|
|
task_type: str = 'system',
|
|
kind: str = 'system-task',
|
|
name: str = '',
|
|
label: str = '',
|
|
context: TaskContext = None,
|
|
scopes: list[core_entities.LifecycleControlScope] = [core_entities.LifecycleControlScope.APPLICATION],
|
|
) -> TaskWrapper:
|
|
wrapper = TaskWrapper(self.ap, coro, task_type, kind, name, label, context, scopes)
|
|
self.tasks.append(wrapper)
|
|
wrapper.task.add_done_callback(lambda _: self._prune_completed_tasks())
|
|
self._prune_completed_tasks()
|
|
return wrapper
|
|
|
|
def create_user_task(
|
|
self,
|
|
coro: typing.Coroutine,
|
|
kind: str = 'user-task',
|
|
name: str = '',
|
|
label: str = '',
|
|
context: TaskContext = None,
|
|
scopes: list[core_entities.LifecycleControlScope] = [core_entities.LifecycleControlScope.APPLICATION],
|
|
) -> TaskWrapper:
|
|
return self.create_task(coro, 'user', kind, name, label, context, scopes)
|
|
|
|
async def wait_all(self):
|
|
await asyncio.gather(*[t.task for t in self.tasks], return_exceptions=True)
|
|
|
|
def get_all_tasks(self) -> list[TaskWrapper]:
|
|
return self.tasks
|
|
|
|
def get_tasks_dict(
|
|
self,
|
|
type: str = None,
|
|
kind: str = None,
|
|
) -> dict:
|
|
return {
|
|
'tasks': [
|
|
t.to_dict()
|
|
for t in self.tasks
|
|
if (type is None or t.task_type == type) and (kind is None or t.kind == kind)
|
|
],
|
|
'id_index': TaskWrapper._id_index,
|
|
}
|
|
|
|
def get_stats(self) -> dict:
|
|
completed = sum(1 for t in self.tasks if t.task.done())
|
|
return {
|
|
'total': len(self.tasks),
|
|
'running': len(self.tasks) - completed,
|
|
'completed': completed,
|
|
'id_index': TaskWrapper._id_index,
|
|
}
|
|
|
|
def get_task_by_id(self, id: int) -> TaskWrapper | None:
|
|
for t in self.tasks:
|
|
if t.id == id:
|
|
return t
|
|
return None
|
|
|
|
def cancel_by_scope(self, scope: core_entities.LifecycleControlScope):
|
|
for wrapper in self.tasks:
|
|
if not wrapper.task.done() and scope in wrapper.scopes:
|
|
wrapper.task.cancel()
|
|
|
|
def cancel_task(self, task_id: int):
|
|
for wrapper in self.tasks:
|
|
if wrapper.id == task_id:
|
|
if not wrapper.task.done():
|
|
wrapper.task.cancel()
|
|
return
|
|
|
|
def _prune_completed_tasks(self):
|
|
completed_limit = (
|
|
self.ap.instance_config.data.get('system', {})
|
|
.get('task_retention', {})
|
|
.get(
|
|
'completed_limit',
|
|
200,
|
|
)
|
|
)
|
|
try:
|
|
completed_limit = int(completed_limit)
|
|
except (TypeError, ValueError):
|
|
completed_limit = 200
|
|
if completed_limit < 1:
|
|
completed_limit = 1
|
|
|
|
completed_tasks = [wrapper for wrapper in self.tasks if wrapper.task.done()]
|
|
overflow = len(completed_tasks) - completed_limit
|
|
if overflow <= 0:
|
|
return
|
|
|
|
remove_ids = {wrapper.id for wrapper in completed_tasks[:overflow]}
|
|
self.tasks = [wrapper for wrapper in self.tasks if wrapper.id not in remove_ids]
|