| 1 | """Cron service for scheduling agent tasks.""" |
| 2 | |
| 3 | import asyncio |
| 4 | import json |
| 5 | import time |
| 6 | import uuid |
| 7 | from dataclasses import asdict |
| 8 | from datetime import datetime |
| 9 | from pathlib import Path |
| 10 | from typing import Any, Callable, Coroutine, Literal |
| 11 | |
| 12 | from filelock import FileLock |
| 13 | from loguru import logger |
| 14 | |
| 15 | from nanobot.cron.types import CronJob, CronJobState, CronPayload, CronRunRecord, CronSchedule, CronStore |
| 16 | |
| 17 | |
| 18 | def _now_ms() -> int: |
| 19 | return int(time.time() * 1000) |
| 20 | |
| 21 | |
| 22 | def _compute_next_run(schedule: CronSchedule, now_ms: int) -> int | None: |
| 23 | """Compute next run time in ms.""" |
| 24 | if schedule.kind == "at": |
| 25 | return schedule.at_ms if schedule.at_ms and schedule.at_ms > now_ms else None |
| 26 | |
| 27 | if schedule.kind == "every": |
| 28 | if not schedule.every_ms or schedule.every_ms <= 0: |
| 29 | return None |
| 30 | # Next interval from now |
| 31 | return now_ms + schedule.every_ms |
| 32 | |
| 33 | if schedule.kind == "cron" and schedule.expr: |
| 34 | try: |
| 35 | from zoneinfo import ZoneInfo |
| 36 | |
| 37 | from croniter import croniter |
| 38 | # Use caller-provided reference time for deterministic scheduling |
| 39 | base_time = now_ms / 1000 |
| 40 | tz = ZoneInfo(schedule.tz) if schedule.tz else datetime.now().astimezone().tzinfo |
| 41 | base_dt = datetime.fromtimestamp(base_time, tz=tz) |
| 42 | cron = croniter(schedule.expr, base_dt) |
| 43 | next_dt = cron.get_next(datetime) |
| 44 | return int(next_dt.timestamp() * 1000) |
| 45 | except Exception: |
| 46 | return None |
| 47 | |
| 48 | return None |
| 49 | |
| 50 | |
| 51 | def _validate_schedule_for_add(schedule: CronSchedule) -> None: |
| 52 | """Validate schedule fields that would otherwise create non-runnable jobs.""" |
| 53 | if schedule.tz and schedule.kind != "cron": |
| 54 | raise ValueError("tz can only be used with cron schedules") |
| 55 | |
| 56 | if schedule.kind == "cron" and schedule.tz: |
| 57 | try: |
| 58 | from zoneinfo import ZoneInfo |
| 59 | |
| 60 | ZoneInfo(schedule.tz) |
| 61 | except Exception: |
| 62 | raise ValueError(f"unknown timezone '{schedule.tz}'") from None |
| 63 | |
| 64 | |
| 65 | class CronService: |
| 66 | """Service for managing and executing scheduled jobs.""" |
| 67 | |
| 68 | _MAX_RUN_HISTORY = 20 |
| 69 | |
| 70 | def __init__( |
| 71 | self, |
| 72 | store_path: Path, |
| 73 | on_job: Callable[[CronJob], Coroutine[Any, Any, str | None]] | None = None, |
| 74 | max_sleep_ms: int = 300_000, # 5 minutes |
| 75 | ): |
| 76 | self.store_path = store_path |
| 77 | self._action_path = store_path.parent / "action.jsonl" |
| 78 | self._lock = FileLock(str(self._action_path.parent) + ".lock") |
| 79 | self.on_job = on_job |
| 80 | self._store: CronStore | None = None |
| 81 | self._timer_task: asyncio.Task | None = None |
| 82 | self._running = False |
| 83 | self._timer_active = False |
| 84 | self.max_sleep_ms = max_sleep_ms |
| 85 | |
| 86 | def _load_jobs(self) -> tuple[list[CronJob], int]: |
| 87 | jobs = [] |
| 88 | version = 1 |
| 89 | if self.store_path.exists(): |
| 90 | try: |
| 91 | data = json.loads(self.store_path.read_text(encoding="utf-8")) |
| 92 | jobs = [] |
| 93 | version = data.get("version", 1) |
| 94 | for j in data.get("jobs", []): |
| 95 | jobs.append(CronJob( |
| 96 | id=j["id"], |
| 97 | name=j["name"], |
| 98 | enabled=j.get("enabled", True), |
| 99 | schedule=CronSchedule( |
| 100 | kind=j["schedule"]["kind"], |
| 101 | at_ms=j["schedule"].get("atMs"), |
| 102 | every_ms=j["schedule"].get("everyMs"), |
| 103 | expr=j["schedule"].get("expr"), |
| 104 | tz=j["schedule"].get("tz"), |
| 105 | ), |
| 106 | payload=CronPayload( |
| 107 | kind=j["payload"].get("kind", "agent_turn"), |
| 108 | message=j["payload"].get("message", ""), |
| 109 | deliver=j["payload"].get("deliver", False), |
| 110 | channel=j["payload"].get("channel"), |
| 111 | to=j["payload"].get("to"), |
| 112 | ), |
| 113 | state=CronJobState( |
| 114 | next_run_at_ms=j.get("state", {}).get("nextRunAtMs"), |
| 115 | last_run_at_ms=j.get("state", {}).get("lastRunAtMs"), |
| 116 | last_status=j.get("state", {}).get("lastStatus"), |
| 117 | last_error=j.get("state", {}).get("lastError"), |
| 118 | run_history=[ |
| 119 | CronRunRecord( |
| 120 | run_at_ms=r["runAtMs"], |
| 121 | status=r["status"], |
| 122 | duration_ms=r.get("durationMs", 0), |
| 123 | error=r.get("error"), |
| 124 | ) |
| 125 | for r in j.get("state", {}).get("runHistory", []) |
| 126 | ], |
| 127 | ), |
| 128 | created_at_ms=j.get("createdAtMs", 0), |
| 129 | updated_at_ms=j.get("updatedAtMs", 0), |
| 130 | delete_after_run=j.get("deleteAfterRun", False), |
| 131 | )) |
| 132 | except Exception as e: |
| 133 | logger.warning("Failed to load cron store: {}", e) |
| 134 | return jobs, version |
| 135 | |
| 136 | def _merge_action(self): |
| 137 | if not self._action_path.exists(): |
| 138 | return |
| 139 | |
| 140 | jobs_map = {j.id: j for j in self._store.jobs} |
| 141 | def _update(params: dict): |
| 142 | j = CronJob.from_dict(params) |
| 143 | jobs_map[j.id] = j |
| 144 | |
| 145 | def _del(params: dict): |
| 146 | if job_id := params.get("job_id"): |
| 147 | jobs_map.pop(job_id) |
| 148 | |
| 149 | with self._lock: |
| 150 | with open(self._action_path, "r", encoding="utf-8") as f: |
| 151 | changed = False |
| 152 | for line in f: |
| 153 | try: |
| 154 | line = line.strip() |
| 155 | action = json.loads(line) |
| 156 | if "action" not in action: |
| 157 | continue |
| 158 | if action["action"] == "del": |
| 159 | _del(action.get("params", {})) |
| 160 | else: |
| 161 | _update(action.get("params", {})) |
| 162 | changed = True |
| 163 | except Exception as exp: |
| 164 | logger.debug(f"load action line error: {exp}") |
| 165 | continue |
| 166 | self._store.jobs = list(jobs_map.values()) |
| 167 | if self._running and changed: |
| 168 | self._action_path.write_text("", encoding="utf-8") |
| 169 | self._save_store() |
| 170 | return |
| 171 | |
| 172 | def _load_store(self) -> CronStore: |
| 173 | """Load jobs from disk. Reloads automatically if file was modified externally. |
| 174 | - Reload every time because it needs to merge operations on the jobs object from other instances. |
| 175 | - During _on_timer execution, return the existing store to prevent concurrent |
| 176 | _load_store calls (e.g. from list_jobs polling) from replacing it mid-execution. |
| 177 | """ |
| 178 | if self._timer_active and self._store: |
| 179 | return self._store |
| 180 | jobs, version = self._load_jobs() |
| 181 | self._store = CronStore(version=version, jobs=jobs) |
| 182 | self._merge_action() |
| 183 | |
| 184 | return self._store |
| 185 | |
| 186 | def _save_store(self) -> None: |
| 187 | """Save jobs to disk.""" |
| 188 | if not self._store: |
| 189 | return |
| 190 | |
| 191 | self.store_path.parent.mkdir(parents=True, exist_ok=True) |
| 192 | |
| 193 | data = { |
| 194 | "version": self._store.version, |
| 195 | "jobs": [ |
| 196 | { |
| 197 | "id": j.id, |
| 198 | "name": j.name, |
| 199 | "enabled": j.enabled, |
| 200 | "schedule": { |
| 201 | "kind": j.schedule.kind, |
| 202 | "atMs": j.schedule.at_ms, |
| 203 | "everyMs": j.schedule.every_ms, |
| 204 | "expr": j.schedule.expr, |
| 205 | "tz": j.schedule.tz, |
| 206 | }, |
| 207 | "payload": { |
| 208 | "kind": j.payload.kind, |
| 209 | "message": j.payload.message, |
| 210 | "deliver": j.payload.deliver, |
| 211 | "channel": j.payload.channel, |
| 212 | "to": j.payload.to, |
| 213 | }, |
| 214 | "state": { |
| 215 | "nextRunAtMs": j.state.next_run_at_ms, |
| 216 | "lastRunAtMs": j.state.last_run_at_ms, |
| 217 | "lastStatus": j.state.last_status, |
| 218 | "lastError": j.state.last_error, |
| 219 | "runHistory": [ |
| 220 | { |
| 221 | "runAtMs": r.run_at_ms, |
| 222 | "status": r.status, |
| 223 | "durationMs": r.duration_ms, |
| 224 | "error": r.error, |
| 225 | } |
| 226 | for r in j.state.run_history |
| 227 | ], |
| 228 | }, |
| 229 | "createdAtMs": j.created_at_ms, |
| 230 | "updatedAtMs": j.updated_at_ms, |
| 231 | "deleteAfterRun": j.delete_after_run, |
| 232 | } |
| 233 | for j in self._store.jobs |
| 234 | ] |
| 235 | } |
| 236 | |
| 237 | self.store_path.write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8") |
| 238 | |
| 239 | async def start(self) -> None: |
| 240 | """Start the cron service.""" |
| 241 | self._running = True |
| 242 | self._load_store() |
| 243 | self._recompute_next_runs() |
| 244 | self._save_store() |
| 245 | self._arm_timer() |
| 246 | logger.info("Cron service started with {} jobs", len(self._store.jobs if self._store else [])) |
| 247 | |
| 248 | def stop(self) -> None: |
| 249 | """Stop the cron service.""" |
| 250 | self._running = False |
| 251 | if self._timer_task: |
| 252 | self._timer_task.cancel() |
| 253 | self._timer_task = None |
| 254 | |
| 255 | def _recompute_next_runs(self) -> None: |
| 256 | """Recompute next run times for all enabled jobs.""" |
| 257 | if not self._store: |
| 258 | return |
| 259 | now = _now_ms() |
| 260 | for job in self._store.jobs: |
| 261 | if job.enabled: |
| 262 | job.state.next_run_at_ms = _compute_next_run(job.schedule, now) |
| 263 | |
| 264 | def _get_next_wake_ms(self) -> int | None: |
| 265 | """Get the earliest next run time across all jobs.""" |
| 266 | if not self._store: |
| 267 | return None |
| 268 | times = [j.state.next_run_at_ms for j in self._store.jobs |
| 269 | if j.enabled and j.state.next_run_at_ms] |
| 270 | return min(times) if times else None |
| 271 | |
| 272 | def _arm_timer(self) -> None: |
| 273 | """Schedule the next timer tick.""" |
| 274 | if self._timer_task: |
| 275 | self._timer_task.cancel() |
| 276 | |
| 277 | if not self._running: |
| 278 | return |
| 279 | |
| 280 | next_wake = self._get_next_wake_ms() |
| 281 | if next_wake is None: |
| 282 | delay_ms = self.max_sleep_ms |
| 283 | else: |
| 284 | delay_ms = min(self.max_sleep_ms, max(0, next_wake - _now_ms())) |
| 285 | delay_s = delay_ms / 1000 |
| 286 | |
| 287 | async def tick(): |
| 288 | await asyncio.sleep(delay_s) |
| 289 | if self._running: |
| 290 | await self._on_timer() |
| 291 | |
| 292 | self._timer_task = asyncio.create_task(tick()) |
| 293 | |
| 294 | async def _on_timer(self) -> None: |
| 295 | """Handle timer tick - run due jobs.""" |
| 296 | self._load_store() |
| 297 | if not self._store: |
| 298 | self._arm_timer() |
| 299 | return |
| 300 | |
| 301 | self._timer_active = True |
| 302 | try: |
| 303 | now = _now_ms() |
| 304 | due_jobs = [ |
| 305 | j for j in self._store.jobs |
| 306 | if j.enabled and j.state.next_run_at_ms and now >= j.state.next_run_at_ms |
| 307 | ] |
| 308 | |
| 309 | for job in due_jobs: |
| 310 | await self._execute_job(job) |
| 311 | |
| 312 | self._save_store() |
| 313 | finally: |
| 314 | self._timer_active = False |
| 315 | self._arm_timer() |
| 316 | |
| 317 | async def _execute_job(self, job: CronJob) -> None: |
| 318 | """Execute a single job.""" |
| 319 | start_ms = _now_ms() |
| 320 | logger.info("Cron: executing job '{}' ({})", job.name, job.id) |
| 321 | |
| 322 | try: |
| 323 | if self.on_job: |
| 324 | await self.on_job(job) |
| 325 | |
| 326 | job.state.last_status = "ok" |
| 327 | job.state.last_error = None |
| 328 | logger.info("Cron: job '{}' completed", job.name) |
| 329 | |
| 330 | except Exception as e: |
| 331 | job.state.last_status = "error" |
| 332 | job.state.last_error = str(e) |
| 333 | logger.error("Cron: job '{}' failed: {}", job.name, e) |
| 334 | |
| 335 | end_ms = _now_ms() |
| 336 | job.state.last_run_at_ms = start_ms |
| 337 | job.updated_at_ms = end_ms |
| 338 | |
| 339 | job.state.run_history.append(CronRunRecord( |
| 340 | run_at_ms=start_ms, |
| 341 | status=job.state.last_status, |
| 342 | duration_ms=end_ms - start_ms, |
| 343 | error=job.state.last_error, |
| 344 | )) |
| 345 | job.state.run_history = job.state.run_history[-self._MAX_RUN_HISTORY:] |
| 346 | |
| 347 | # Handle one-shot jobs |
| 348 | if job.schedule.kind == "at": |
| 349 | if job.delete_after_run: |
| 350 | self._store.jobs = [j for j in self._store.jobs if j.id != job.id] |
| 351 | else: |
| 352 | job.enabled = False |
| 353 | job.state.next_run_at_ms = None |
| 354 | else: |
| 355 | # Compute next run |
| 356 | job.state.next_run_at_ms = _compute_next_run(job.schedule, _now_ms()) |
| 357 | |
| 358 | def _append_action(self, action: Literal["add", "del", "update"], params: dict): |
| 359 | self.store_path.parent.mkdir(parents=True, exist_ok=True) |
| 360 | with self._lock: |
| 361 | with open(self._action_path, "a", encoding="utf-8") as f: |
| 362 | f.write(json.dumps({"action": action, "params": params}, ensure_ascii=False) + "\n") |
| 363 | |
| 364 | |
| 365 | # ========== Public API ========== |
| 366 | |
| 367 | def list_jobs(self, include_disabled: bool = False) -> list[CronJob]: |
| 368 | """List all jobs.""" |
| 369 | store = self._load_store() |
| 370 | jobs = store.jobs if include_disabled else [j for j in store.jobs if j.enabled] |
| 371 | return sorted(jobs, key=lambda j: j.state.next_run_at_ms or float('inf')) |
| 372 | |
| 373 | def add_job( |
| 374 | self, |
| 375 | name: str, |
| 376 | schedule: CronSchedule, |
| 377 | message: str, |
| 378 | deliver: bool = False, |
| 379 | channel: str | None = None, |
| 380 | to: str | None = None, |
| 381 | delete_after_run: bool = False, |
| 382 | ) -> CronJob: |
| 383 | """Add a new job.""" |
| 384 | _validate_schedule_for_add(schedule) |
| 385 | now = _now_ms() |
| 386 | |
| 387 | job = CronJob( |
| 388 | id=str(uuid.uuid4())[:8], |
| 389 | name=name, |
| 390 | enabled=True, |
| 391 | schedule=schedule, |
| 392 | payload=CronPayload( |
| 393 | kind="agent_turn", |
| 394 | message=message, |
| 395 | deliver=deliver, |
| 396 | channel=channel, |
| 397 | to=to, |
| 398 | ), |
| 399 | state=CronJobState(next_run_at_ms=_compute_next_run(schedule, now)), |
| 400 | created_at_ms=now, |
| 401 | updated_at_ms=now, |
| 402 | delete_after_run=delete_after_run, |
| 403 | ) |
| 404 | if self._running: |
| 405 | store = self._load_store() |
| 406 | store.jobs.append(job) |
| 407 | self._save_store() |
| 408 | self._arm_timer() |
| 409 | else: |
| 410 | self._append_action("add", asdict(job)) |
| 411 | |
| 412 | logger.info("Cron: added job '{}' ({})", name, job.id) |
| 413 | return job |
| 414 | |
| 415 | def register_system_job(self, job: CronJob) -> CronJob: |
| 416 | """Register an internal system job (idempotent on restart).""" |
| 417 | store = self._load_store() |
| 418 | now = _now_ms() |
| 419 | job.state = CronJobState(next_run_at_ms=_compute_next_run(job.schedule, now)) |
| 420 | job.created_at_ms = now |
| 421 | job.updated_at_ms = now |
| 422 | store.jobs = [j for j in store.jobs if j.id != job.id] |
| 423 | store.jobs.append(job) |
| 424 | self._save_store() |
| 425 | self._arm_timer() |
| 426 | logger.info("Cron: registered system job '{}' ({})", job.name, job.id) |
| 427 | return job |
| 428 | |
| 429 | def remove_job(self, job_id: str) -> Literal["removed", "protected", "not_found"]: |
| 430 | """Remove a job by ID, unless it is a protected system job.""" |
| 431 | store = self._load_store() |
| 432 | job = next((j for j in store.jobs if j.id == job_id), None) |
| 433 | if job is None: |
| 434 | return "not_found" |
| 435 | if job.payload.kind == "system_event": |
| 436 | logger.info("Cron: refused to remove protected system job {}", job_id) |
| 437 | return "protected" |
| 438 | |
| 439 | before = len(store.jobs) |
| 440 | store.jobs = [j for j in store.jobs if j.id != job_id] |
| 441 | removed = len(store.jobs) < before |
| 442 | |
| 443 | if removed: |
| 444 | if self._running: |
| 445 | self._save_store() |
| 446 | self._arm_timer() |
| 447 | else: |
| 448 | self._append_action("del", {"job_id": job_id}) |
| 449 | logger.info("Cron: removed job {}", job_id) |
| 450 | return "removed" |
| 451 | |
| 452 | return "not_found" |
| 453 | |
| 454 | def enable_job(self, job_id: str, enabled: bool = True) -> CronJob | None: |
| 455 | """Enable or disable a job.""" |
| 456 | store = self._load_store() |
| 457 | for job in store.jobs: |
| 458 | if job.id == job_id: |
| 459 | job.enabled = enabled |
| 460 | job.updated_at_ms = _now_ms() |
| 461 | if enabled: |
| 462 | job.state.next_run_at_ms = _compute_next_run(job.schedule, _now_ms()) |
| 463 | else: |
| 464 | job.state.next_run_at_ms = None |
| 465 | if self._running: |
| 466 | self._save_store() |
| 467 | self._arm_timer() |
| 468 | else: |
| 469 | self._append_action("update", asdict(job)) |
| 470 | return job |
| 471 | return None |
| 472 | |
| 473 | def update_job( |
| 474 | self, |
| 475 | job_id: str, |
| 476 | *, |
| 477 | name: str | None = None, |
| 478 | schedule: CronSchedule | None = None, |
| 479 | message: str | None = None, |
| 480 | deliver: bool | None = None, |
| 481 | channel: str | None = ..., |
| 482 | to: str | None = ..., |
| 483 | delete_after_run: bool | None = None, |
| 484 | ) -> CronJob | Literal["not_found", "protected"]: |
| 485 | """Update mutable fields of an existing job. System jobs cannot be updated. |
| 486 | |
| 487 | For ``channel`` and ``to``, pass an explicit value (including ``None``) |
| 488 | to update; omit (sentinel ``...``) to leave unchanged. |
| 489 | """ |
| 490 | store = self._load_store() |
| 491 | job = next((j for j in store.jobs if j.id == job_id), None) |
| 492 | if job is None: |
| 493 | return "not_found" |
| 494 | if job.payload.kind == "system_event": |
| 495 | return "protected" |
| 496 | |
| 497 | if schedule is not None: |
| 498 | _validate_schedule_for_add(schedule) |
| 499 | job.schedule = schedule |
| 500 | if name is not None: |
| 501 | job.name = name |
| 502 | if message is not None: |
| 503 | job.payload.message = message |
| 504 | if deliver is not None: |
| 505 | job.payload.deliver = deliver |
| 506 | if channel is not ...: |
| 507 | job.payload.channel = channel |
| 508 | if to is not ...: |
| 509 | job.payload.to = to |
| 510 | if delete_after_run is not None: |
| 511 | job.delete_after_run = delete_after_run |
| 512 | |
| 513 | job.updated_at_ms = _now_ms() |
| 514 | if job.enabled: |
| 515 | job.state.next_run_at_ms = _compute_next_run(job.schedule, _now_ms()) |
| 516 | |
| 517 | if self._running: |
| 518 | self._save_store() |
| 519 | self._arm_timer() |
| 520 | else: |
| 521 | self._append_action("update", asdict(job)) |
| 522 | |
| 523 | logger.info("Cron: updated job '{}' ({})", job.name, job.id) |
| 524 | return job |
| 525 | |
| 526 | async def run_job(self, job_id: str, force: bool = False) -> bool: |
| 527 | """Manually run a job without disturbing the service's running state.""" |
| 528 | was_running = self._running |
| 529 | self._running = True |
| 530 | try: |
| 531 | store = self._load_store() |
| 532 | for job in store.jobs: |
| 533 | if job.id == job_id: |
| 534 | if not force and not job.enabled: |
| 535 | return False |
| 536 | await self._execute_job(job) |
| 537 | self._save_store() |
| 538 | return True |
| 539 | return False |
| 540 | finally: |
| 541 | self._running = was_running |
| 542 | if was_running: |
| 543 | self._arm_timer() |
| 544 | |
| 545 | def get_job(self, job_id: str) -> CronJob | None: |
| 546 | """Get a job by ID.""" |
| 547 | store = self._load_store() |
| 548 | return next((j for j in store.jobs if j.id == job_id), None) |
| 549 | |
| 550 | def status(self) -> dict: |
| 551 | """Get service status.""" |
| 552 | store = self._load_store() |
| 553 | return { |
| 554 | "enabled": self._running, |
| 555 | "jobs": len(store.jobs), |
| 556 | "next_wake_at_ms": self._get_next_wake_ms(), |
| 557 | } |
| 558 |