| 1 | import asyncio |
| 2 | from datetime import datetime, timezone, timedelta |
| 3 | import os |
| 4 | import random |
| 5 | import threading |
| 6 | from urllib.parse import urlparse |
| 7 | import uuid |
| 8 | from enum import Enum |
| 9 | from os.path import exists |
| 10 | from typing import Any, Callable, Dict, Literal, Optional, Type, TypeVar, Union, cast, ClassVar |
| 11 | |
| 12 | import nest_asyncio |
| 13 | nest_asyncio.apply() |
| 14 | |
| 15 | from crontab import CronTab |
| 16 | from pydantic import BaseModel, Field, PrivateAttr |
| 17 | |
| 18 | from agent import Agent, AgentContext, UserMessage |
| 19 | from initialize import initialize_agent |
| 20 | from helpers.persist_chat import save_tmp_chat |
| 21 | from helpers.print_style import PrintStyle |
| 22 | from helpers.defer import DeferredTask |
| 23 | from helpers.files import get_abs_path, make_dirs, read_file, write_file |
| 24 | from helpers.localization import Localization |
| 25 | from helpers import projects, guids |
| 26 | import pytz |
| 27 | from typing import Annotated |
| 28 | |
| 29 | SCHEDULER_FOLDER = "usr/scheduler" |
| 30 | LOCAL_TIMEZONE_ALIASES = {"local", "user", "default", "current", "current_timezone"} |
| 31 | |
| 32 | |
| 33 | def normalize_schedule_timezone(timezone_name: str | None) -> str: |
| 34 | name = str(timezone_name or "").strip() |
| 35 | if not name or name.lower() in LOCAL_TIMEZONE_ALIASES: |
| 36 | return Localization.get().get_timezone() |
| 37 | try: |
| 38 | pytz.timezone(name) |
| 39 | except pytz.exceptions.UnknownTimeZoneError: |
| 40 | PrintStyle.error(f"Unknown task schedule timezone: {name}, using current user timezone") |
| 41 | return Localization.get().get_timezone() |
| 42 | return name |
| 43 | |
| 44 | |
| 45 | def _now() -> datetime: |
| 46 | localization = Localization.get() |
| 47 | now = getattr(localization, "now", None) |
| 48 | if callable(now): |
| 49 | return now() |
| 50 | |
| 51 | try: |
| 52 | tzinfo = pytz.timezone(localization.get_timezone()) |
| 53 | except Exception: |
| 54 | tzinfo = pytz.timezone("UTC") |
| 55 | return datetime.now(tzinfo) |
| 56 | |
| 57 | |
| 58 | def _localize_task_datetime(dt: datetime) -> datetime: |
| 59 | if dt.tzinfo is not None: |
| 60 | return dt |
| 61 | |
| 62 | localization = Localization.get() |
| 63 | localize = getattr(localization, "localize_naive_datetime", None) |
| 64 | if callable(localize): |
| 65 | return localize(dt) |
| 66 | |
| 67 | try: |
| 68 | tzinfo = pytz.timezone(localization.get_timezone()) |
| 69 | except Exception: |
| 70 | tzinfo = pytz.timezone("UTC") |
| 71 | return tzinfo.localize(dt) |
| 72 | |
| 73 | # ---------------------- |
| 74 | # Task Models |
| 75 | # ---------------------- |
| 76 | |
| 77 | |
| 78 | class TaskState(str, Enum): |
| 79 | IDLE = "idle" |
| 80 | RUNNING = "running" |
| 81 | DISABLED = "disabled" |
| 82 | ERROR = "error" |
| 83 | |
| 84 | |
| 85 | class TaskType(str, Enum): |
| 86 | AD_HOC = "adhoc" |
| 87 | SCHEDULED = "scheduled" |
| 88 | PLANNED = "planned" |
| 89 | |
| 90 | |
| 91 | class TaskSchedule(BaseModel): |
| 92 | minute: str |
| 93 | hour: str |
| 94 | day: str |
| 95 | month: str |
| 96 | weekday: str |
| 97 | timezone: str = Field(default_factory=lambda: Localization.get().get_timezone()) |
| 98 | |
| 99 | def to_crontab(self) -> str: |
| 100 | return f"{self.minute} {self.hour} {self.day} {self.month} {self.weekday}" |
| 101 | |
| 102 | |
| 103 | class TaskPlan(BaseModel): |
| 104 | todo: list[datetime] = Field(default_factory=list) |
| 105 | in_progress: datetime | None = None |
| 106 | done: list[datetime] = Field(default_factory=list) |
| 107 | |
| 108 | @classmethod |
| 109 | def create( |
| 110 | cls, |
| 111 | todo: list[datetime] | None = None, |
| 112 | in_progress: datetime | None = None, |
| 113 | done: list[datetime] | None = None, |
| 114 | ): |
| 115 | todo = list(todo or []) |
| 116 | done = list(done or []) |
| 117 | if todo: |
| 118 | for idx, dt in enumerate(todo): |
| 119 | todo[idx] = _localize_task_datetime(dt) |
| 120 | if in_progress: |
| 121 | in_progress = _localize_task_datetime(in_progress) |
| 122 | if done: |
| 123 | for idx, dt in enumerate(done): |
| 124 | done[idx] = _localize_task_datetime(dt) |
| 125 | return cls(todo=todo, in_progress=in_progress, done=done) |
| 126 | |
| 127 | def add_todo(self, launch_time: datetime): |
| 128 | launch_time = _localize_task_datetime(launch_time) |
| 129 | self.todo.append(launch_time) |
| 130 | self.todo = sorted(self.todo) |
| 131 | |
| 132 | def set_in_progress(self, launch_time: datetime): |
| 133 | launch_time = _localize_task_datetime(launch_time) |
| 134 | if launch_time not in self.todo: |
| 135 | raise ValueError(f"Launch time {launch_time} not in todo list") |
| 136 | self.todo.remove(launch_time) |
| 137 | self.todo = sorted(self.todo) |
| 138 | self.in_progress = launch_time |
| 139 | |
| 140 | def set_done(self, launch_time: datetime): |
| 141 | launch_time = _localize_task_datetime(launch_time) |
| 142 | if launch_time != self.in_progress: |
| 143 | raise ValueError(f"Launch time {launch_time} is not the same as in progress time {self.in_progress}") |
| 144 | if launch_time in self.done: |
| 145 | raise ValueError(f"Launch time {launch_time} already in done list") |
| 146 | self.in_progress = None |
| 147 | self.done.append(launch_time) |
| 148 | self.done = sorted(self.done) |
| 149 | |
| 150 | def get_next_launch_time(self) -> datetime | None: |
| 151 | return self.todo[0] if self.todo else None |
| 152 | |
| 153 | def should_launch(self) -> datetime | None: |
| 154 | next_launch_time = self.get_next_launch_time() |
| 155 | if next_launch_time is None: |
| 156 | return None |
| 157 | if _now() > next_launch_time: |
| 158 | return next_launch_time |
| 159 | return None |
| 160 | |
| 161 | |
| 162 | class BaseTask(BaseModel): |
| 163 | uuid: str = Field(default_factory=lambda: guids.generate_id()) |
| 164 | context_id: Optional[str] = Field(default=None) |
| 165 | state: TaskState = Field(default=TaskState.IDLE) |
| 166 | name: str = Field() |
| 167 | system_prompt: str |
| 168 | prompt: str |
| 169 | attachments: list[str] = Field(default_factory=list) |
| 170 | project_name: str | None = Field(default=None) |
| 171 | project_color: str | None = Field(default=None) |
| 172 | created_at: datetime = Field(default_factory=_now) |
| 173 | updated_at: datetime = Field(default_factory=_now) |
| 174 | last_run: datetime | None = None |
| 175 | last_result: str | None = None |
| 176 | |
| 177 | def __init__(self, *args, **kwargs): |
| 178 | super().__init__(*args, **kwargs) |
| 179 | if not self.context_id: |
| 180 | self.context_id = self.uuid |
| 181 | self._lock = threading.RLock() |
| 182 | |
| 183 | def update(self, |
| 184 | name: str | None = None, |
| 185 | state: TaskState | None = None, |
| 186 | system_prompt: str | None = None, |
| 187 | prompt: str | None = None, |
| 188 | attachments: list[str] | None = None, |
| 189 | last_run: datetime | None = None, |
| 190 | last_result: str | None = None, |
| 191 | context_id: str | None = None, |
| 192 | **kwargs): |
| 193 | with self._lock: |
| 194 | if name is not None: |
| 195 | self.name = name |
| 196 | self.updated_at = _now() |
| 197 | if state is not None: |
| 198 | self.state = state |
| 199 | self.updated_at = _now() |
| 200 | if system_prompt is not None: |
| 201 | self.system_prompt = system_prompt |
| 202 | self.updated_at = _now() |
| 203 | if prompt is not None: |
| 204 | self.prompt = prompt |
| 205 | self.updated_at = _now() |
| 206 | if attachments is not None: |
| 207 | self.attachments = attachments |
| 208 | self.updated_at = _now() |
| 209 | if last_run is not None: |
| 210 | self.last_run = last_run |
| 211 | self.updated_at = _now() |
| 212 | if last_result is not None: |
| 213 | self.last_result = last_result |
| 214 | self.updated_at = _now() |
| 215 | if context_id is not None: |
| 216 | self.context_id = context_id |
| 217 | self.updated_at = _now() |
| 218 | for key, value in kwargs.items(): |
| 219 | if value is not None: |
| 220 | setattr(self, key, value) |
| 221 | self.updated_at = _now() |
| 222 | |
| 223 | def check_schedule(self, frequency_seconds: float = 60.0) -> bool: |
| 224 | return False |
| 225 | |
| 226 | def get_next_run(self) -> datetime | None: |
| 227 | return None |
| 228 | |
| 229 | def is_dedicated(self) -> bool: |
| 230 | return self.context_id == self.uuid |
| 231 | |
| 232 | def get_next_run_minutes(self) -> int | None: |
| 233 | next_run = self.get_next_run() |
| 234 | if next_run is None: |
| 235 | return None |
| 236 | return int((next_run - _now()).total_seconds() / 60) |
| 237 | |
| 238 | async def on_run(self): |
| 239 | pass |
| 240 | |
| 241 | async def on_finish(self): |
| 242 | # Ensure that updated_at is refreshed to reflect completion time |
| 243 | # This helps track when the task actually finished, regardless of success/error |
| 244 | await TaskScheduler.get().update_task( |
| 245 | self.uuid, |
| 246 | updated_at=_now() |
| 247 | ) |
| 248 | |
| 249 | async def on_error(self, error: str): |
| 250 | # Update task state to ERROR and set last result |
| 251 | scheduler = TaskScheduler.get() |
| 252 | await scheduler.reload() # Ensure we have the latest state |
| 253 | updated_task = await scheduler.update_task( |
| 254 | self.uuid, |
| 255 | state=TaskState.ERROR, |
| 256 | last_run=_now(), |
| 257 | last_result=f"ERROR: {error}" |
| 258 | ) |
| 259 | if not updated_task: |
| 260 | PrintStyle.error( |
| 261 | f"Failed to update task {self.uuid} state to ERROR after error: {error}" |
| 262 | ) |
| 263 | await scheduler.save() # Force save after update |
| 264 | |
| 265 | async def on_success(self, result: str): |
| 266 | # Update task state to IDLE and set last result |
| 267 | scheduler = TaskScheduler.get() |
| 268 | await scheduler.reload() # Ensure we have the latest state |
| 269 | updated_task = await scheduler.update_task( |
| 270 | self.uuid, |
| 271 | state=TaskState.IDLE, |
| 272 | last_run=_now(), |
| 273 | last_result=result |
| 274 | ) |
| 275 | if not updated_task: |
| 276 | PrintStyle.error( |
| 277 | f"Failed to update task {self.uuid} state to IDLE after success" |
| 278 | ) |
| 279 | await scheduler.save() # Force save after update |
| 280 | |
| 281 | |
| 282 | class AdHocTask(BaseTask): |
| 283 | type: Literal[TaskType.AD_HOC] = TaskType.AD_HOC |
| 284 | token: str = Field(default_factory=lambda: str(random.randint(1000000000000000000, 9999999999999999999))) |
| 285 | |
| 286 | @classmethod |
| 287 | def create( |
| 288 | cls, |
| 289 | name: str, |
| 290 | system_prompt: str, |
| 291 | prompt: str, |
| 292 | token: str, |
| 293 | attachments: list[str] | None = None, |
| 294 | context_id: str | None = None, |
| 295 | project_name: str | None = None, |
| 296 | project_color: str | None = None |
| 297 | ): |
| 298 | return cls(name=name, |
| 299 | system_prompt=system_prompt, |
| 300 | prompt=prompt, |
| 301 | attachments=list(attachments or []), |
| 302 | token=token, |
| 303 | context_id=context_id, |
| 304 | project_name=project_name, |
| 305 | project_color=project_color) |
| 306 | |
| 307 | def update(self, |
| 308 | name: str | None = None, |
| 309 | state: TaskState | None = None, |
| 310 | system_prompt: str | None = None, |
| 311 | prompt: str | None = None, |
| 312 | attachments: list[str] | None = None, |
| 313 | last_run: datetime | None = None, |
| 314 | last_result: str | None = None, |
| 315 | context_id: str | None = None, |
| 316 | token: str | None = None, |
| 317 | **kwargs): |
| 318 | super().update(name=name, |
| 319 | state=state, |
| 320 | system_prompt=system_prompt, |
| 321 | prompt=prompt, |
| 322 | attachments=attachments, |
| 323 | last_run=last_run, |
| 324 | last_result=last_result, |
| 325 | context_id=context_id, |
| 326 | token=token, |
| 327 | **kwargs) |
| 328 | |
| 329 | |
| 330 | class ScheduledTask(BaseTask): |
| 331 | type: Literal[TaskType.SCHEDULED] = TaskType.SCHEDULED |
| 332 | schedule: TaskSchedule |
| 333 | |
| 334 | @classmethod |
| 335 | def create( |
| 336 | cls, |
| 337 | name: str, |
| 338 | system_prompt: str, |
| 339 | prompt: str, |
| 340 | schedule: TaskSchedule, |
| 341 | attachments: list[str] | None = None, |
| 342 | context_id: str | None = None, |
| 343 | timezone: str | None = None, |
| 344 | project_name: str | None = None, |
| 345 | project_color: str | None = None, |
| 346 | ): |
| 347 | # Set timezone in schedule if provided |
| 348 | if timezone is not None: |
| 349 | schedule.timezone = normalize_schedule_timezone(timezone) |
| 350 | else: |
| 351 | schedule.timezone = normalize_schedule_timezone(schedule.timezone) |
| 352 | |
| 353 | return cls(name=name, |
| 354 | system_prompt=system_prompt, |
| 355 | prompt=prompt, |
| 356 | attachments=list(attachments or []), |
| 357 | schedule=schedule, |
| 358 | context_id=context_id, |
| 359 | project_name=project_name, |
| 360 | project_color=project_color) |
| 361 | |
| 362 | def update(self, |
| 363 | name: str | None = None, |
| 364 | state: TaskState | None = None, |
| 365 | system_prompt: str | None = None, |
| 366 | prompt: str | None = None, |
| 367 | attachments: list[str] | None = None, |
| 368 | last_run: datetime | None = None, |
| 369 | last_result: str | None = None, |
| 370 | context_id: str | None = None, |
| 371 | schedule: TaskSchedule | None = None, |
| 372 | **kwargs): |
| 373 | super().update(name=name, |
| 374 | state=state, |
| 375 | system_prompt=system_prompt, |
| 376 | prompt=prompt, |
| 377 | attachments=attachments, |
| 378 | last_run=last_run, |
| 379 | last_result=last_result, |
| 380 | context_id=context_id, |
| 381 | schedule=schedule, |
| 382 | **kwargs) |
| 383 | |
| 384 | def check_schedule(self, frequency_seconds: float = 60.0) -> bool: |
| 385 | with self._lock: |
| 386 | crontab = CronTab(crontab=self.schedule.to_crontab()) # type: ignore |
| 387 | |
| 388 | # Get the timezone from the schedule or use UTC as fallback |
| 389 | self.schedule.timezone = normalize_schedule_timezone(self.schedule.timezone) |
| 390 | task_timezone = pytz.timezone(self.schedule.timezone) |
| 391 | |
| 392 | # Get reference time in task's timezone (by default now - frequency_seconds) |
| 393 | reference_time = (_now() - timedelta(seconds=frequency_seconds)).astimezone(task_timezone) |
| 394 | |
| 395 | # Get next run time as seconds until next execution |
| 396 | next_run_seconds: Optional[float] = crontab.next( # type: ignore |
| 397 | now=reference_time, |
| 398 | return_datetime=False |
| 399 | ) # type: ignore |
| 400 | |
| 401 | if next_run_seconds is None: |
| 402 | return False |
| 403 | |
| 404 | return next_run_seconds < frequency_seconds |
| 405 | |
| 406 | def get_next_run(self) -> datetime | None: |
| 407 | with self._lock: |
| 408 | crontab = CronTab(crontab=self.schedule.to_crontab()) # type: ignore |
| 409 | self.schedule.timezone = normalize_schedule_timezone(self.schedule.timezone) |
| 410 | task_timezone = pytz.timezone(self.schedule.timezone) |
| 411 | now_in_task_timezone = datetime.now(timezone.utc).astimezone(task_timezone) |
| 412 | next_run = crontab.next(now=now_in_task_timezone, return_datetime=True) # type: ignore |
| 413 | if next_run is None: |
| 414 | return None |
| 415 | if next_run.tzinfo is None: |
| 416 | next_run = task_timezone.localize(next_run) |
| 417 | return next_run.astimezone(timezone.utc) |
| 418 | |
| 419 | |
| 420 | class PlannedTask(BaseTask): |
| 421 | type: Literal[TaskType.PLANNED] = TaskType.PLANNED |
| 422 | plan: TaskPlan |
| 423 | |
| 424 | @classmethod |
| 425 | def create( |
| 426 | cls, |
| 427 | name: str, |
| 428 | system_prompt: str, |
| 429 | prompt: str, |
| 430 | plan: TaskPlan, |
| 431 | attachments: list[str] | None = None, |
| 432 | context_id: str | None = None, |
| 433 | project_name: str | None = None, |
| 434 | project_color: str | None = None |
| 435 | ): |
| 436 | return cls(name=name, |
| 437 | system_prompt=system_prompt, |
| 438 | prompt=prompt, |
| 439 | plan=plan, |
| 440 | attachments=list(attachments or []), |
| 441 | context_id=context_id, |
| 442 | project_name=project_name, |
| 443 | project_color=project_color) |
| 444 | |
| 445 | def update(self, |
| 446 | name: str | None = None, |
| 447 | state: TaskState | None = None, |
| 448 | system_prompt: str | None = None, |
| 449 | prompt: str | None = None, |
| 450 | attachments: list[str] | None = None, |
| 451 | last_run: datetime | None = None, |
| 452 | last_result: str | None = None, |
| 453 | context_id: str | None = None, |
| 454 | plan: TaskPlan | None = None, |
| 455 | **kwargs): |
| 456 | super().update(name=name, |
| 457 | state=state, |
| 458 | system_prompt=system_prompt, |
| 459 | prompt=prompt, |
| 460 | attachments=attachments, |
| 461 | last_run=last_run, |
| 462 | last_result=last_result, |
| 463 | context_id=context_id, |
| 464 | plan=plan, |
| 465 | **kwargs) |
| 466 | |
| 467 | def check_schedule(self, frequency_seconds: float = 60.0) -> bool: |
| 468 | with self._lock: |
| 469 | return self.plan.should_launch() is not None |
| 470 | |
| 471 | def get_next_run(self) -> datetime | None: |
| 472 | with self._lock: |
| 473 | return self.plan.get_next_launch_time() |
| 474 | |
| 475 | async def on_run(self): |
| 476 | with self._lock: |
| 477 | # Get the next launch time and set it as in_progress |
| 478 | next_launch_time = self.plan.should_launch() |
| 479 | if next_launch_time is not None: |
| 480 | self.plan.set_in_progress(next_launch_time) |
| 481 | await super().on_run() |
| 482 | |
| 483 | async def on_finish(self): |
| 484 | # Handle plan item progression regardless of success or error |
| 485 | plan_updated = False |
| 486 | |
| 487 | with self._lock: |
| 488 | # If there's an in_progress time, mark it as done |
| 489 | if self.plan.in_progress is not None: |
| 490 | self.plan.set_done(self.plan.in_progress) |
| 491 | plan_updated = True |
| 492 | |
| 493 | # If we updated the plan, make sure to persist it |
| 494 | if plan_updated: |
| 495 | scheduler = TaskScheduler.get() |
| 496 | await scheduler.reload() |
| 497 | await scheduler.update_task(self.uuid, plan=self.plan) |
| 498 | await scheduler.save() # Force save |
| 499 | |
| 500 | # Call the parent implementation for any additional cleanup |
| 501 | await super().on_finish() |
| 502 | |
| 503 | async def on_success(self, result: str): |
| 504 | # Call parent implementation to update state, etc. |
| 505 | await super().on_success(result) |
| 506 | |
| 507 | async def on_error(self, error: str): |
| 508 | # Call parent implementation to update state, etc. |
| 509 | await super().on_error(error) |
| 510 | |
| 511 | |
| 512 | class SchedulerTaskList(BaseModel): |
| 513 | tasks: list[Annotated[Union[ScheduledTask, AdHocTask, PlannedTask], Field(discriminator="type")]] = Field(default_factory=list) |
| 514 | # Singleton instance |
| 515 | __instance: ClassVar[Optional["SchedulerTaskList"]] = PrivateAttr(default=None) |
| 516 | |
| 517 | # lock: threading.Lock = Field(exclude=True, default=threading.Lock()) |
| 518 | |
| 519 | @classmethod |
| 520 | def get(cls) -> "SchedulerTaskList": |
| 521 | path = get_abs_path(SCHEDULER_FOLDER, "tasks.json") |
| 522 | if cls.__instance is None: |
| 523 | if not exists(path): |
| 524 | make_dirs(path) |
| 525 | cls.__instance = asyncio.run(cls(tasks=[]).save()) |
| 526 | else: |
| 527 | cls.__instance = cls.model_validate_json(read_file(path)) |
| 528 | else: |
| 529 | asyncio.run(cls.__instance.reload()) |
| 530 | return cls.__instance |
| 531 | |
| 532 | def __init__(self, *args, **kwargs): |
| 533 | super().__init__(*args, **kwargs) |
| 534 | self._lock = threading.RLock() |
| 535 | |
| 536 | async def reload(self) -> "SchedulerTaskList": |
| 537 | path = get_abs_path(SCHEDULER_FOLDER, "tasks.json") |
| 538 | if exists(path): |
| 539 | with self._lock: |
| 540 | data = self.__class__.model_validate_json(read_file(path)) |
| 541 | self.tasks.clear() |
| 542 | self.tasks.extend(data.tasks) |
| 543 | return self |
| 544 | |
| 545 | async def add_task(self, task: Union[ScheduledTask, AdHocTask, PlannedTask]) -> "SchedulerTaskList": |
| 546 | with self._lock: |
| 547 | self.tasks.append(task) |
| 548 | await self.save() |
| 549 | return self |
| 550 | |
| 551 | async def save(self) -> "SchedulerTaskList": |
| 552 | with self._lock: |
| 553 | # Debug: check for AdHocTasks with null tokens before saving |
| 554 | for task in self.tasks: |
| 555 | if isinstance(task, AdHocTask): |
| 556 | if task.token is None or task.token == "": |
| 557 | PrintStyle.warning( |
| 558 | f"WARNING: AdHocTask {task.name} ({task.uuid}) has a null or empty token before saving: '{task.token}'" |
| 559 | ) |
| 560 | # Generate a new token to prevent errors |
| 561 | task.token = str(random.randint(1000000000000000000, 9999999999999999999)) |
| 562 | PrintStyle.info( |
| 563 | f"Fixed: Generated new token '{task.token}' for task {task.name}" |
| 564 | ) |
| 565 | |
| 566 | path = get_abs_path(SCHEDULER_FOLDER, "tasks.json") |
| 567 | if not exists(path): |
| 568 | make_dirs(path) |
| 569 | |
| 570 | # Get the JSON string before writing |
| 571 | json_data = self.model_dump_json() |
| 572 | |
| 573 | # Debug: check if 'null' appears as token value in JSON |
| 574 | if '"type": "adhoc"' in json_data and '"token": null' in json_data: |
| 575 | PrintStyle.error( |
| 576 | "ERROR: Found null token in JSON output for an adhoc task" |
| 577 | ) |
| 578 | |
| 579 | write_file(path, json_data) |
| 580 | |
| 581 | # Debug: Verify after saving |
| 582 | if exists(path): |
| 583 | loaded_json = read_file(path) |
| 584 | if '"type": "adhoc"' in loaded_json and '"token": null' in loaded_json: |
| 585 | PrintStyle.error( |
| 586 | "ERROR: Null token persisted in JSON file for an adhoc task" |
| 587 | ) |
| 588 | |
| 589 | return self |
| 590 | |
| 591 | async def update_task_by_uuid( |
| 592 | self, |
| 593 | task_uuid: str, |
| 594 | updater_func: Callable[[Union[ScheduledTask, AdHocTask, PlannedTask]], None], |
| 595 | verify_func: Callable[[Union[ScheduledTask, AdHocTask, PlannedTask]], bool] = lambda task: True |
| 596 | ) -> Union[ScheduledTask, AdHocTask, PlannedTask] | None: |
| 597 | """ |
| 598 | Atomically update a task by UUID using the provided updater function. |
| 599 | |
| 600 | The updater_func should take the task as an argument and perform any necessary updates. |
| 601 | This method ensures that the task is updated and saved atomically, preventing race conditions. |
| 602 | |
| 603 | Returns the updated task or None if not found. |
| 604 | """ |
| 605 | with self._lock: |
| 606 | # Reload to ensure we have the latest state |
| 607 | await self.reload() |
| 608 | |
| 609 | # Find the task |
| 610 | task = next((task for task in self.tasks if task.uuid == task_uuid and verify_func(task)), None) |
| 611 | if task is None: |
| 612 | return None |
| 613 | |
| 614 | # Apply the updates via the provided function |
| 615 | updater_func(task) |
| 616 | |
| 617 | # Save the changes |
| 618 | await self.save() |
| 619 | |
| 620 | return task |
| 621 | |
| 622 | def get_tasks(self) -> list[Union[ScheduledTask, AdHocTask, PlannedTask]]: |
| 623 | with self._lock: |
| 624 | return self.tasks |
| 625 | |
| 626 | def get_tasks_by_context_id(self, context_id: str, only_running: bool = False) -> list[Union[ScheduledTask, AdHocTask, PlannedTask]]: |
| 627 | with self._lock: |
| 628 | return [ |
| 629 | task for task in self.tasks |
| 630 | if task.context_id == context_id |
| 631 | and (not only_running or task.state == TaskState.RUNNING) |
| 632 | ] |
| 633 | |
| 634 | async def get_due_tasks(self) -> list[Union[ScheduledTask, AdHocTask, PlannedTask]]: |
| 635 | with self._lock: |
| 636 | await self.reload() |
| 637 | return [ |
| 638 | task for task in self.tasks |
| 639 | if task.check_schedule() and task.state == TaskState.IDLE |
| 640 | ] |
| 641 | |
| 642 | def get_task_by_uuid(self, task_uuid: str) -> Union[ScheduledTask, AdHocTask, PlannedTask] | None: |
| 643 | with self._lock: |
| 644 | return next((task for task in self.tasks if task.uuid == task_uuid), None) |
| 645 | |
| 646 | def get_task_by_name(self, name: str) -> Union[ScheduledTask, AdHocTask, PlannedTask] | None: |
| 647 | with self._lock: |
| 648 | return next((task for task in self.tasks if task.name == name), None) |
| 649 | |
| 650 | def find_task_by_name(self, name: str) -> list[Union[ScheduledTask, AdHocTask, PlannedTask]]: |
| 651 | with self._lock: |
| 652 | return [task for task in self.tasks if name.lower() in task.name.lower()] |
| 653 | |
| 654 | async def remove_task_by_uuid(self, task_uuid: str) -> "SchedulerTaskList": |
| 655 | with self._lock: |
| 656 | self.tasks = [task for task in self.tasks if task.uuid != task_uuid] |
| 657 | await self.save() |
| 658 | return self |
| 659 | |
| 660 | async def remove_task_by_name(self, name: str) -> "SchedulerTaskList": |
| 661 | with self._lock: |
| 662 | self.tasks = [task for task in self.tasks if task.name != name] |
| 663 | await self.save() |
| 664 | return self |
| 665 | |
| 666 | |
| 667 | class TaskScheduler: |
| 668 | |
| 669 | _tasks: SchedulerTaskList |
| 670 | _printer: PrintStyle |
| 671 | _instance = None |
| 672 | _running_deferred_tasks: Dict[str, DeferredTask] |
| 673 | _running_tasks_lock: threading.RLock |
| 674 | |
| 675 | @classmethod |
| 676 | def get(cls) -> "TaskScheduler": |
| 677 | if cls._instance is None: |
| 678 | cls._instance = cls() |
| 679 | return cls._instance |
| 680 | |
| 681 | def __init__(self): |
| 682 | # Only initialize if this is a new instance |
| 683 | if not hasattr(self, '_initialized'): |
| 684 | self._tasks = SchedulerTaskList.get() |
| 685 | self._printer = PrintStyle(italic=True, font_color="green", padding=False) |
| 686 | self._running_deferred_tasks = {} |
| 687 | self._running_tasks_lock = threading.RLock() |
| 688 | self._initialized = True |
| 689 | |
| 690 | def _register_running_task(self, task_uuid: str, deferred_task: DeferredTask) -> None: |
| 691 | with self._running_tasks_lock: |
| 692 | self._running_deferred_tasks[task_uuid] = deferred_task |
| 693 | |
| 694 | def _unregister_running_task(self, task_uuid: str) -> None: |
| 695 | with self._running_tasks_lock: |
| 696 | self._running_deferred_tasks.pop(task_uuid, None) |
| 697 | |
| 698 | def cancel_running_task(self, task_uuid: str, terminate_thread: bool = False) -> bool: |
| 699 | with self._running_tasks_lock: |
| 700 | deferred_task = self._running_deferred_tasks.get(task_uuid) |
| 701 | if not deferred_task: |
| 702 | return False |
| 703 | PrintStyle.info(f"Scheduler cancelling task {task_uuid}") |
| 704 | deferred_task.kill(terminate_thread=terminate_thread) |
| 705 | return True |
| 706 | |
| 707 | def cancel_tasks_by_context(self, context_id: str, terminate_thread: bool = False) -> bool: |
| 708 | cancelled_any = False |
| 709 | with self._running_tasks_lock: |
| 710 | running_tasks = list(self._running_deferred_tasks.keys()) |
| 711 | for task_uuid in running_tasks: |
| 712 | task = self.get_task_by_uuid(task_uuid) |
| 713 | if task and task.context_id == context_id: |
| 714 | if self.cancel_running_task(task_uuid, terminate_thread=terminate_thread): |
| 715 | cancelled_any = True |
| 716 | return cancelled_any |
| 717 | |
| 718 | async def reload(self): |
| 719 | await self._tasks.reload() |
| 720 | |
| 721 | def get_tasks(self) -> list[Union[ScheduledTask, AdHocTask, PlannedTask]]: |
| 722 | return self._tasks.get_tasks() |
| 723 | |
| 724 | def get_tasks_by_context_id(self, context_id: str, only_running: bool = False) -> list[Union[ScheduledTask, AdHocTask, PlannedTask]]: |
| 725 | return self._tasks.get_tasks_by_context_id(context_id, only_running) |
| 726 | |
| 727 | async def add_task(self, task: Union[ScheduledTask, AdHocTask, PlannedTask]) -> "TaskScheduler": |
| 728 | await self._tasks.add_task(task) |
| 729 | ctx = await self._get_chat_context(task) # invoke context creation |
| 730 | from helpers.state_monitor_integration import mark_dirty_all |
| 731 | mark_dirty_all(reason="task_scheduler.TaskScheduler.add_task") |
| 732 | return self |
| 733 | |
| 734 | async def remove_task_by_uuid(self, task_uuid: str) -> "TaskScheduler": |
| 735 | await self._tasks.remove_task_by_uuid(task_uuid) |
| 736 | from helpers.state_monitor_integration import mark_dirty_all |
| 737 | mark_dirty_all(reason="task_scheduler.TaskScheduler.remove_task_by_uuid") |
| 738 | return self |
| 739 | |
| 740 | async def remove_task_by_name(self, name: str) -> "TaskScheduler": |
| 741 | await self._tasks.remove_task_by_name(name) |
| 742 | from helpers.state_monitor_integration import mark_dirty_all |
| 743 | mark_dirty_all(reason="task_scheduler.TaskScheduler.remove_task_by_name") |
| 744 | return self |
| 745 | |
| 746 | def get_task_by_uuid(self, task_uuid: str) -> Union[ScheduledTask, AdHocTask, PlannedTask] | None: |
| 747 | return self._tasks.get_task_by_uuid(task_uuid) |
| 748 | |
| 749 | def get_task_by_name(self, name: str) -> Union[ScheduledTask, AdHocTask, PlannedTask] | None: |
| 750 | return self._tasks.get_task_by_name(name) |
| 751 | |
| 752 | def find_task_by_name(self, name: str) -> list[Union[ScheduledTask, AdHocTask, PlannedTask]]: |
| 753 | return self._tasks.find_task_by_name(name) |
| 754 | |
| 755 | async def tick(self): |
| 756 | for task in await self._tasks.get_due_tasks(): |
| 757 | await self._run_task(task) |
| 758 | |
| 759 | async def run_task_by_uuid(self, task_uuid: str, task_context: str | None = None): |
| 760 | # First reload tasks to ensure we have the latest state |
| 761 | await self._tasks.reload() |
| 762 | |
| 763 | # Get the task to run |
| 764 | task = self.get_task_by_uuid(task_uuid) |
| 765 | if not task: |
| 766 | raise ValueError(f"Task with UUID '{task_uuid}' not found") |
| 767 | |
| 768 | # If the task is already running, raise an error |
| 769 | if task.state == TaskState.RUNNING: |
| 770 | raise ValueError(f"Task '{task.name}' is already running") |
| 771 | |
| 772 | # If the task is disabled, raise an error |
| 773 | if task.state == TaskState.DISABLED: |
| 774 | raise ValueError(f"Task '{task.name}' is disabled") |
| 775 | |
| 776 | # If the task is in error state, reset it to IDLE first |
| 777 | if task.state == TaskState.ERROR: |
| 778 | PrintStyle.info(f"Resetting task '{task.name}' from ERROR to IDLE state before running") |
| 779 | await self.update_task(task_uuid, state=TaskState.IDLE) |
| 780 | # Force a reload to ensure we have the updated state |
| 781 | await self._tasks.reload() |
| 782 | task = self.get_task_by_uuid(task_uuid) |
| 783 | if not task: |
| 784 | raise ValueError(f"Task with UUID '{task_uuid}' not found after state reset") |
| 785 | |
| 786 | # Run the task |
| 787 | await self._run_task(task, task_context) |
| 788 | |
| 789 | async def run_task_by_name(self, name: str, task_context: str | None = None): |
| 790 | task = self._tasks.get_task_by_name(name) |
| 791 | if task is None: |
| 792 | raise ValueError(f"Task with name {name} not found") |
| 793 | await self._run_task(task, task_context) |
| 794 | |
| 795 | async def save(self): |
| 796 | await self._tasks.save() |
| 797 | |
| 798 | async def update_task_checked( |
| 799 | self, |
| 800 | task_uuid: str, |
| 801 | verify_func: Callable[[Union[ScheduledTask, AdHocTask, PlannedTask]], bool] = lambda task: True, |
| 802 | **update_params |
| 803 | ) -> Union[ScheduledTask, AdHocTask, PlannedTask] | None: |
| 804 | """ |
| 805 | Atomically update a task by UUID with the provided parameters. |
| 806 | This prevents race conditions when multiple processes update tasks concurrently. |
| 807 | |
| 808 | Returns the updated task or None if not found. |
| 809 | """ |
| 810 | def _update_task(task): |
| 811 | task.update(**update_params) |
| 812 | |
| 813 | updated = await self._tasks.update_task_by_uuid(task_uuid, _update_task, verify_func) |
| 814 | if updated is not None: |
| 815 | from helpers.state_monitor_integration import mark_dirty_all |
| 816 | mark_dirty_all(reason="task_scheduler.TaskScheduler.update_task_checked") |
| 817 | return updated |
| 818 | |
| 819 | async def update_task(self, task_uuid: str, **update_params) -> Union[ScheduledTask, AdHocTask, PlannedTask] | None: |
| 820 | return await self.update_task_checked(task_uuid, lambda task: True, **update_params) |
| 821 | |
| 822 | async def __new_context(self, task: Union[ScheduledTask, AdHocTask, PlannedTask]) -> AgentContext: |
| 823 | if not task.context_id: |
| 824 | raise ValueError(f"Task {task.name} has no context ID") |
| 825 | |
| 826 | config = initialize_agent() |
| 827 | context: AgentContext = AgentContext(config, id=task.context_id, name=task.name) |
| 828 | # context.id = task.context_id |
| 829 | # initial name before renaming is same as task name |
| 830 | # context.name = task.name |
| 831 | |
| 832 | # Activate project if set |
| 833 | if task.project_name: |
| 834 | projects.activate_project(context.id, task.project_name) |
| 835 | |
| 836 | # Save the context |
| 837 | save_tmp_chat(context) |
| 838 | return context |
| 839 | |
| 840 | async def _get_chat_context(self, task: Union[ScheduledTask, AdHocTask, PlannedTask]) -> AgentContext: |
| 841 | context = AgentContext.get(task.context_id) if task.context_id else None |
| 842 | |
| 843 | if context: |
| 844 | assert isinstance(context, AgentContext) |
| 845 | PrintStyle.info( |
| 846 | f"Scheduler Task {task.name} loaded from task {task.uuid}, context ok" |
| 847 | ) |
| 848 | save_tmp_chat(context) |
| 849 | return context |
| 850 | else: |
| 851 | message = ( |
| 852 | f"Scheduler Task {task.name} loaded from task {task.uuid} but context not found" |
| 853 | ) |
| 854 | if task.is_dedicated(): |
| 855 | PrintStyle.info(f"{message}; creating dedicated context") |
| 856 | else: |
| 857 | PrintStyle.warning(message) |
| 858 | return await self.__new_context(task) |
| 859 | |
| 860 | async def _persist_chat(self, task: Union[ScheduledTask, AdHocTask, PlannedTask], context: AgentContext): |
| 861 | if context.id != task.context_id: |
| 862 | raise ValueError(f"Context ID mismatch for task {task.name}: context {context.id} != task {task.context_id}") |
| 863 | save_tmp_chat(context) |
| 864 | |
| 865 | async def _run_task(self, task: Union[ScheduledTask, AdHocTask, PlannedTask], task_context: str | None = None): |
| 866 | |
| 867 | async def _run_task_wrapper(task_uuid: str, task_context: str | None = None): |
| 868 | |
| 869 | # preflight checks with a snapshot of the task |
| 870 | task_snapshot: Union[ScheduledTask, AdHocTask, PlannedTask] | None = self.get_task_by_uuid(task_uuid) |
| 871 | if task_snapshot is None: |
| 872 | PrintStyle.error(f"Scheduler Task with UUID '{task_uuid}' not found") |
| 873 | self._unregister_running_task(task_uuid) |
| 874 | return |
| 875 | if task_snapshot.state == TaskState.RUNNING: |
| 876 | PrintStyle.warning(f"Scheduler Task '{task_snapshot.name}' already running, skipping") |
| 877 | self._unregister_running_task(task_uuid) |
| 878 | return |
| 879 | |
| 880 | # Atomically fetch and check the task's current state |
| 881 | current_task = await self.update_task_checked(task_uuid, lambda task: task.state != TaskState.RUNNING, state=TaskState.RUNNING) |
| 882 | if not current_task: |
| 883 | PrintStyle.error(f"Scheduler Task with UUID '{task_uuid}' not found or updated by another process") |
| 884 | self._unregister_running_task(task_uuid) |
| 885 | return |
| 886 | if current_task.state != TaskState.RUNNING: |
| 887 | # This means the update failed due to state conflict |
| 888 | PrintStyle.warning(f"Scheduler Task '{current_task.name}' state is '{current_task.state}', skipping") |
| 889 | self._unregister_running_task(task_uuid) |
| 890 | return |
| 891 | |
| 892 | await current_task.on_run() |
| 893 | |
| 894 | # the agent instance - init in try block |
| 895 | agent = None |
| 896 | |
| 897 | try: |
| 898 | PrintStyle.info(f"Scheduler Task '{current_task.name}' started") |
| 899 | |
| 900 | context = await self._get_chat_context(current_task) |
| 901 | AgentContext.use(context.id) |
| 902 | |
| 903 | # Ensure the context is properly registered in the AgentContext._contexts |
| 904 | # This is critical for the polling mechanism to find and stream logs |
| 905 | # Dict operations are atomic |
| 906 | # AgentContext._contexts[context.id] = context |
| 907 | agent = context.streaming_agent or context.agent0 |
| 908 | |
| 909 | # Prepare attachment filenames for logging |
| 910 | attachment_filenames = [] |
| 911 | if current_task.attachments: |
| 912 | for attachment in current_task.attachments: |
| 913 | if os.path.exists(attachment): |
| 914 | attachment_filenames.append(attachment) |
| 915 | else: |
| 916 | try: |
| 917 | url = urlparse(attachment) |
| 918 | if url.scheme in ["http", "https", "ftp", "ftps", "sftp"]: |
| 919 | attachment_filenames.append(attachment) |
| 920 | else: |
| 921 | PrintStyle.warning(f"Skipping attachment: [{attachment}]") |
| 922 | except Exception: |
| 923 | PrintStyle.warning(f"Skipping attachment: [{attachment}]") |
| 924 | |
| 925 | self._printer.print("User message:") |
| 926 | self._printer.print(f"> {current_task.prompt}") |
| 927 | if attachment_filenames: |
| 928 | self._printer.print("Attachments:") |
| 929 | for filename in attachment_filenames: |
| 930 | self._printer.print(f"- {filename}") |
| 931 | |
| 932 | task_prompt = f"# Starting scheduler task '{current_task.name}' ({current_task.uuid})" |
| 933 | if task_context: |
| 934 | task_prompt = f"## Context:\n{task_context}\n\n## Task:\n{current_task.prompt}" |
| 935 | else: |
| 936 | task_prompt = f"## Task:\n{current_task.prompt}" |
| 937 | |
| 938 | # Log the message with message_id and attachments |
| 939 | msg_id = str(uuid.uuid4()) |
| 940 | context.log.log( |
| 941 | type="user", |
| 942 | heading="", |
| 943 | content=task_prompt, |
| 944 | kvps={"attachments": attachment_filenames}, |
| 945 | id=msg_id, |
| 946 | ) |
| 947 | |
| 948 | agent.hist_add_user_message( |
| 949 | UserMessage( |
| 950 | message=task_prompt, |
| 951 | system_message=[current_task.system_prompt], |
| 952 | attachments=attachment_filenames, |
| 953 | id=msg_id)) |
| 954 | |
| 955 | # Persist after setting up the context but before running the agent |
| 956 | # This ensures the task context is saved and can be found by polling |
| 957 | await self._persist_chat(current_task, context) |
| 958 | |
| 959 | result = await agent.monologue() |
| 960 | |
| 961 | # Success |
| 962 | PrintStyle.success(f"Scheduler Task '{current_task.name}' completed: {result}") |
| 963 | await self._persist_chat(current_task, context) |
| 964 | await current_task.on_success(result) |
| 965 | |
| 966 | # Explicitly verify task was updated in storage after success |
| 967 | await self._tasks.reload() |
| 968 | updated_task = self.get_task_by_uuid(task_uuid) |
| 969 | if updated_task and updated_task.state != TaskState.IDLE: |
| 970 | PrintStyle.warning(f"Fixing task state consistency: '{current_task.name}' state is not IDLE after success") |
| 971 | await self.update_task(task_uuid, state=TaskState.IDLE) |
| 972 | |
| 973 | except asyncio.CancelledError: |
| 974 | PrintStyle.warning(f"Scheduler Task '{current_task.name}' cancelled by user") |
| 975 | try: |
| 976 | await asyncio.shield(self.update_task(task_uuid, state=TaskState.IDLE)) |
| 977 | except Exception: |
| 978 | pass |
| 979 | raise |
| 980 | except Exception as e: |
| 981 | # Error |
| 982 | PrintStyle.error(f"Scheduler Task '{current_task.name}' failed: {e}") |
| 983 | await current_task.on_error(str(e)) |
| 984 | |
| 985 | # Explicitly verify task was updated in storage after error |
| 986 | await self._tasks.reload() |
| 987 | updated_task = self.get_task_by_uuid(task_uuid) |
| 988 | if updated_task and updated_task.state != TaskState.ERROR: |
| 989 | PrintStyle.warning(f"Fixing task state consistency: '{current_task.name}' state is not ERROR after failure") |
| 990 | await self.update_task(task_uuid, state=TaskState.ERROR) |
| 991 | |
| 992 | # if agent: |
| 993 | # await agent.handle_exception("scheduler", e) |
| 994 | finally: |
| 995 | # Call on_finish for task-specific cleanup |
| 996 | try: |
| 997 | await asyncio.shield(current_task.on_finish()) |
| 998 | except asyncio.CancelledError: |
| 999 | pass |
| 1000 | except Exception: |
| 1001 | pass |
| 1002 | |
| 1003 | # Make one final save to ensure all states are persisted |
| 1004 | try: |
| 1005 | await asyncio.shield(self._tasks.save()) |
| 1006 | except asyncio.CancelledError: |
| 1007 | pass |
| 1008 | except Exception: |
| 1009 | pass |
| 1010 | |
| 1011 | self._unregister_running_task(task_uuid) |
| 1012 | |
| 1013 | deferred_task = DeferredTask(thread_name=self.__class__.__name__) |
| 1014 | self._register_running_task(task.uuid, deferred_task) |
| 1015 | deferred_task.start_task(_run_task_wrapper, task.uuid, task_context) |
| 1016 | |
| 1017 | # Ensure background execution doesn't exit immediately on async await, especially in script contexts. |
| 1018 | # Yielding briefly keeps callers like CLI scripts alive long enough for the DeferredTask thread to spin up |
| 1019 | # without leaving stray pending tasks that trigger \"Task was destroyed\" warnings when the loop shuts down. |
| 1020 | await asyncio.sleep(0.1) |
| 1021 | |
| 1022 | def serialize_all_tasks(self) -> list[Dict[str, Any]]: |
| 1023 | """ |
| 1024 | Serialize all tasks in the scheduler to a list of dictionaries. |
| 1025 | """ |
| 1026 | return serialize_tasks(self.get_tasks()) |
| 1027 | |
| 1028 | def serialize_task(self, task_id: str) -> Optional[Dict[str, Any]]: |
| 1029 | """ |
| 1030 | Serialize a specific task in the scheduler by UUID. |
| 1031 | Returns None if task is not found. |
| 1032 | """ |
| 1033 | # Get task without locking, as get_task_by_uuid() is already thread-safe |
| 1034 | task = self.get_task_by_uuid(task_id) |
| 1035 | if task: |
| 1036 | return serialize_task(task) |
| 1037 | return None |
| 1038 | |
| 1039 | |
| 1040 | # ---------------------- |
| 1041 | # Task Serialization Helpers |
| 1042 | # ---------------------- |
| 1043 | |
| 1044 | def serialize_datetime(dt: Optional[datetime]) -> Optional[str]: |
| 1045 | """ |
| 1046 | Serialize a datetime object to ISO format string in the user's timezone. |
| 1047 | |
| 1048 | This uses the Localization singleton to convert the datetime to the user's timezone |
| 1049 | before serializing it to an ISO format string for frontend display. |
| 1050 | |
| 1051 | Returns None if the input is None. |
| 1052 | """ |
| 1053 | # Use the Localization singleton for timezone conversion and serialization |
| 1054 | return Localization.get().serialize_datetime(dt) |
| 1055 | |
| 1056 | |
| 1057 | def parse_datetime(dt_str: Optional[str]) -> Optional[datetime]: |
| 1058 | """ |
| 1059 | Parse ISO format datetime string with timezone awareness. |
| 1060 | |
| 1061 | This converts from the localized ISO format returned by serialize_datetime |
| 1062 | back to a datetime object with proper timezone handling. |
| 1063 | |
| 1064 | Returns None if dt_str is None. |
| 1065 | """ |
| 1066 | if not dt_str: |
| 1067 | return None |
| 1068 | |
| 1069 | try: |
| 1070 | # Use the Localization singleton for consistent timezone handling |
| 1071 | return Localization.get().localtime_str_to_utc_dt(dt_str) |
| 1072 | except ValueError as e: |
| 1073 | raise ValueError(f"Invalid datetime format: {dt_str}. Expected ISO format. Error: {e}") |
| 1074 | |
| 1075 | |
| 1076 | def serialize_task_schedule(schedule: TaskSchedule) -> Dict[str, str]: |
| 1077 | """Convert TaskSchedule to a standardized dictionary format.""" |
| 1078 | schedule.timezone = normalize_schedule_timezone(schedule.timezone) |
| 1079 | return { |
| 1080 | 'minute': schedule.minute, |
| 1081 | 'hour': schedule.hour, |
| 1082 | 'day': schedule.day, |
| 1083 | 'month': schedule.month, |
| 1084 | 'weekday': schedule.weekday, |
| 1085 | 'timezone': schedule.timezone |
| 1086 | } |
| 1087 | |
| 1088 | |
| 1089 | def parse_task_schedule(schedule_data: Dict[str, str]) -> TaskSchedule: |
| 1090 | """Parse dictionary into TaskSchedule with validation.""" |
| 1091 | try: |
| 1092 | return TaskSchedule( |
| 1093 | minute=schedule_data.get('minute', '*'), |
| 1094 | hour=schedule_data.get('hour', '*'), |
| 1095 | day=schedule_data.get('day', '*'), |
| 1096 | month=schedule_data.get('month', '*'), |
| 1097 | weekday=schedule_data.get('weekday', '*'), |
| 1098 | timezone=normalize_schedule_timezone(schedule_data.get('timezone')) |
| 1099 | ) |
| 1100 | except Exception as e: |
| 1101 | raise ValueError(f"Invalid schedule format: {e}") from e |
| 1102 | |
| 1103 | |
| 1104 | def serialize_task_plan(plan: TaskPlan) -> Dict[str, Any]: |
| 1105 | """Convert TaskPlan to a standardized dictionary format.""" |
| 1106 | return { |
| 1107 | 'todo': [serialize_datetime(dt) for dt in plan.todo], |
| 1108 | 'in_progress': serialize_datetime(plan.in_progress) if plan.in_progress else None, |
| 1109 | 'done': [serialize_datetime(dt) for dt in plan.done] |
| 1110 | } |
| 1111 | |
| 1112 | |
| 1113 | def parse_task_plan(plan_data: Dict[str, Any]) -> TaskPlan: |
| 1114 | """Parse dictionary into TaskPlan with validation.""" |
| 1115 | try: |
| 1116 | # Handle case where plan_data might be None or empty |
| 1117 | if not plan_data: |
| 1118 | return TaskPlan(todo=[], in_progress=None, done=[]) |
| 1119 | |
| 1120 | # Parse todo items with careful validation |
| 1121 | todo_dates = [] |
| 1122 | for dt_str in plan_data.get('todo', []): |
| 1123 | if dt_str: |
| 1124 | parsed_dt = parse_datetime(dt_str) |
| 1125 | if parsed_dt: |
| 1126 | # Ensure datetime is timezone-aware (use the user's timezone if not specified) |
| 1127 | if parsed_dt.tzinfo is None: |
| 1128 | parsed_dt = _localize_task_datetime(parsed_dt) |
| 1129 | todo_dates.append(parsed_dt) |
| 1130 | |
| 1131 | # Parse in_progress with validation |
| 1132 | in_progress = None |
| 1133 | if plan_data.get('in_progress'): |
| 1134 | in_progress = parse_datetime(plan_data.get('in_progress')) |
| 1135 | # Ensure datetime is timezone-aware |
| 1136 | if in_progress and in_progress.tzinfo is None: |
| 1137 | in_progress = _localize_task_datetime(in_progress) |
| 1138 | |
| 1139 | # Parse done items with validation |
| 1140 | done_dates = [] |
| 1141 | for dt_str in plan_data.get('done', []): |
| 1142 | if dt_str: |
| 1143 | parsed_dt = parse_datetime(dt_str) |
| 1144 | if parsed_dt: |
| 1145 | # Ensure datetime is timezone-aware |
| 1146 | if parsed_dt.tzinfo is None: |
| 1147 | parsed_dt = _localize_task_datetime(parsed_dt) |
| 1148 | done_dates.append(parsed_dt) |
| 1149 | |
| 1150 | # Sort dates for better usability |
| 1151 | todo_dates.sort() |
| 1152 | done_dates.sort(reverse=True) # Most recent first for done items |
| 1153 | |
| 1154 | # Cast to ensure type safety |
| 1155 | todo_dates_cast: list[datetime] = cast(list[datetime], todo_dates) |
| 1156 | done_dates_cast: list[datetime] = cast(list[datetime], done_dates) |
| 1157 | |
| 1158 | return TaskPlan.create( |
| 1159 | todo=todo_dates_cast, |
| 1160 | in_progress=in_progress, |
| 1161 | done=done_dates_cast |
| 1162 | ) |
| 1163 | except Exception as e: |
| 1164 | PrintStyle.error( |
| 1165 | f"Error parsing task plan: {e}" |
| 1166 | ) |
| 1167 | # Return empty plan instead of failing |
| 1168 | return TaskPlan(todo=[], in_progress=None, done=[]) |
| 1169 | |
| 1170 | |
| 1171 | T = TypeVar('T', bound=Union[ScheduledTask, AdHocTask, PlannedTask]) |
| 1172 | |
| 1173 | |
| 1174 | def serialize_task(task: Union[ScheduledTask, AdHocTask, PlannedTask]) -> Dict[str, Any]: |
| 1175 | """ |
| 1176 | Standardized serialization for task objects with proper handling of all complex types. |
| 1177 | """ |
| 1178 | # Start with a basic dictionary |
| 1179 | task_dict = { |
| 1180 | "uuid": task.uuid, |
| 1181 | "name": task.name, |
| 1182 | "state": task.state, |
| 1183 | "system_prompt": task.system_prompt, |
| 1184 | "prompt": task.prompt, |
| 1185 | "attachments": task.attachments, |
| 1186 | "project_name": task.project_name, |
| 1187 | "project_color": task.project_color, |
| 1188 | "created_at": serialize_datetime(task.created_at), |
| 1189 | "updated_at": serialize_datetime(task.updated_at), |
| 1190 | "last_run": serialize_datetime(task.last_run), |
| 1191 | "next_run": serialize_datetime(task.get_next_run()), |
| 1192 | "last_result": task.last_result, |
| 1193 | "context_id": task.context_id, |
| 1194 | "dedicated_context": task.is_dedicated(), |
| 1195 | "project": { |
| 1196 | "name": task.project_name, |
| 1197 | "color": task.project_color, |
| 1198 | }, |
| 1199 | } |
| 1200 | |
| 1201 | # Add type-specific fields |
| 1202 | if isinstance(task, ScheduledTask): |
| 1203 | task_dict['type'] = 'scheduled' |
| 1204 | task_dict['schedule'] = serialize_task_schedule(task.schedule) # type: ignore |
| 1205 | elif isinstance(task, AdHocTask): |
| 1206 | task_dict['type'] = 'adhoc' |
| 1207 | adhoc_task = cast(AdHocTask, task) |
| 1208 | task_dict['token'] = adhoc_task.token |
| 1209 | else: |
| 1210 | task_dict['type'] = 'planned' |
| 1211 | planned_task = cast(PlannedTask, task) |
| 1212 | task_dict['plan'] = serialize_task_plan(planned_task.plan) # type: ignore |
| 1213 | |
| 1214 | return task_dict |
| 1215 | |
| 1216 | |
| 1217 | def serialize_tasks(tasks: list[Union[ScheduledTask, AdHocTask, PlannedTask]]) -> list[Dict[str, Any]]: |
| 1218 | """ |
| 1219 | Serialize a list of tasks to a list of dictionaries. |
| 1220 | """ |
| 1221 | return [serialize_task(task) for task in tasks] |
| 1222 | |
| 1223 | |
| 1224 | def deserialize_task(task_data: Dict[str, Any], task_class: Optional[Type[T]] = None) -> T: |
| 1225 | """ |
| 1226 | Deserialize dictionary into appropriate task object with validation. |
| 1227 | If task_class is provided, uses that type. Otherwise determines type from data. |
| 1228 | """ |
| 1229 | task_type_str = task_data.get('type', '') |
| 1230 | determined_class = None |
| 1231 | |
| 1232 | if not task_class: |
| 1233 | # Determine task class from data |
| 1234 | if task_type_str == 'scheduled': |
| 1235 | determined_class = cast(Type[T], ScheduledTask) |
| 1236 | elif task_type_str == 'adhoc': |
| 1237 | determined_class = cast(Type[T], AdHocTask) |
| 1238 | # Ensure token is a valid non-empty string |
| 1239 | if not task_data.get('token'): |
| 1240 | task_data['token'] = str(random.randint(1000000000000000000, 9999999999999999999)) |
| 1241 | elif task_type_str == 'planned': |
| 1242 | determined_class = cast(Type[T], PlannedTask) |
| 1243 | else: |
| 1244 | raise ValueError(f"Unknown task type: {task_type_str}") |
| 1245 | else: |
| 1246 | determined_class = task_class |
| 1247 | # If this is an AdHocTask, ensure token is valid |
| 1248 | if determined_class == AdHocTask and not task_data.get('token'): # type: ignore |
| 1249 | task_data['token'] = str(random.randint(1000000000000000000, 9999999999999999999)) |
| 1250 | |
| 1251 | common_args = { |
| 1252 | "uuid": task_data.get("uuid"), |
| 1253 | "name": task_data.get("name"), |
| 1254 | "state": TaskState(task_data.get("state", TaskState.IDLE)), |
| 1255 | "system_prompt": task_data.get("system_prompt", ""), |
| 1256 | "prompt": task_data.get("prompt", ""), |
| 1257 | "attachments": task_data.get("attachments", []), |
| 1258 | "project_name": task_data.get("project_name"), |
| 1259 | "project_color": task_data.get("project_color"), |
| 1260 | "created_at": parse_datetime(task_data.get("created_at")), |
| 1261 | "updated_at": parse_datetime(task_data.get("updated_at")), |
| 1262 | "last_run": parse_datetime(task_data.get("last_run")), |
| 1263 | "last_result": task_data.get("last_result"), |
| 1264 | "context_id": task_data.get("context_id"), |
| 1265 | } |
| 1266 | |
| 1267 | # Add type-specific fields |
| 1268 | if determined_class == ScheduledTask: # type: ignore |
| 1269 | schedule_data = task_data.get("schedule", {}) |
| 1270 | common_args["schedule"] = parse_task_schedule(schedule_data) |
| 1271 | return ScheduledTask(**common_args) # type: ignore |
| 1272 | elif determined_class == AdHocTask: # type: ignore |
| 1273 | common_args["token"] = task_data.get("token", "") |
| 1274 | return AdHocTask(**common_args) # type: ignore |
| 1275 | else: |
| 1276 | plan_data = task_data.get("plan", {}) |
| 1277 | common_args["plan"] = parse_task_plan(plan_data) |
| 1278 | return PlannedTask(**common_args) # type: ignore |