""" Scheduler module for Rex AI Assistant. This file is a merge-resolved, compatibility-first scheduler that supports both older "JobDefinition" style jobs and newer "ScheduledJob" style jobs. Supported APIs (both): - Legacy: scheduler.register_handler(name, handler) job = scheduler.add_job(name, interval_seconds, handler_name, job_id=..., enabled=..., metadata=...) scheduler.list_jobs() -> list[JobDefinition/ScheduledJob] scheduler.run_job(job_id, manual=True|False) -> bool scheduler.run_due_jobs() -> list[str] scheduler.get_job(job_id) -> JobDefinition|None - Newer: scheduler.register_callback(name, callback) job = scheduler.add_job(job_id=..., name=..., schedule="interval:600", callback_name=..., workflow_id=..., ...) scheduler.run_job(job_id, force=True|False) -> bool scheduler.start() / scheduler.stop() for background loop (optional) scheduler.update_job(job_id, **updates) Persistence: - Defaults to: data/scheduler/jobs.json """ from __future__ import annotations import json import logging import threading import time import uuid from collections.abc import Callable from dataclasses import dataclass, field from datetime import UTC, datetime, timedelta from pathlib import Path from typing import Any logger = logging.getLogger(__name__) # Handler signature used by both styles: # We pass the job object itself to the callback. JobCallback = Callable[["ScheduledJob"], None] @dataclass class SchedulerMetrics: """Cumulative metrics for scheduler execution.""" total_runs: int = 0 successful_runs: int = 0 failed_runs: int = 0 last_run_at: datetime | None = None last_failure_at: datetime | None = None def _utc_now() -> datetime: return datetime.now() def _ensure_tz(dt: datetime) -> datetime: if dt.tzinfo is None: return dt return dt.astimezone(UTC) def _parse_dt(value: Any) -> datetime | None: if value is None: return None if isinstance(value, datetime): return _ensure_tz(value) if isinstance(value, str) and value.strip(): try: return _ensure_tz(datetime.fromisoformat(value)) except ValueError: return None return None def _schedule_from_interval(interval_seconds: int) -> str: return f"interval:{int(interval_seconds)}" def _interval_from_schedule(schedule: str) -> int: s = (schedule or "").strip().lower() if s.startswith("interval:"): try: return int(s.split(":", 1)[1].strip()) except ValueError: return 3600 return 3600 def _calculate_next_run(schedule: str, from_time: datetime | None = None) -> datetime: base = _ensure_tz(from_time or _utc_now()) seconds = _interval_from_schedule(schedule) return base + timedelta(seconds=seconds) @dataclass class ScheduledJob: """ Unified job model. Fields cover the newer scheduler, but we provide legacy-compatible aliases: - interval_seconds property (derived from schedule) - handler_name alias for callback_name - next_run_at alias for next_run - last_run_at used by legacy systems """ # Required job_id: str name: str # Scheduling schedule: str = field(default_factory=lambda: "interval:3600") enabled: bool = True next_run: datetime = field(default_factory=_utc_now) last_run_at: datetime | None = None last_scheduled_run: datetime | None = None # Execution callback_name: str | None = None workflow_id: str | None = None # Limits and stats max_runs: int | None = None run_count: int = 0 # Extra metadata: dict[str, Any] = field(default_factory=dict) # ---------- Legacy compatibility aliases ---------- @property def interval_seconds(self) -> int: return _interval_from_schedule(self.schedule) @property def handler_name(self) -> str: # Legacy expects a string handler_name. Prefer callback_name; fallback to empty. return self.callback_name or "" @property def next_run_at(self) -> datetime: return self.next_run @next_run_at.setter def next_run_at(self, value: datetime) -> None: self.next_run = _ensure_tz(value) # ---------- Serialization ---------- def to_dict(self) -> dict[str, Any]: return { "job_id": self.job_id, "name": self.name, "schedule": self.schedule, "enabled": self.enabled, "next_run": _ensure_tz(self.next_run).isoformat(), "last_run_at": _ensure_tz(self.last_run_at).isoformat() if self.last_run_at else None, "last_scheduled_run": ( _ensure_tz(self.last_scheduled_run).isoformat() if self.last_scheduled_run else None ), "max_runs": self.max_runs, "run_count": self.run_count, "callback_name": self.callback_name, "workflow_id": self.workflow_id, "metadata": self.metadata or {}, } @classmethod def from_dict(cls, data: dict[str, Any]) -> ScheduledJob: # Accept both legacy and newer keys. job_id = data.get("job_id") or data.get("id") or str(uuid.uuid4()) name = data.get("name") or "Unnamed Job" # Legacy: interval_seconds + handler_name + next_run_at interval_seconds = data.get("interval_seconds") handler_name = data.get("handler_name") next_run_at = data.get("next_run_at") # New: schedule + callback_name + next_run schedule = data.get("schedule") callback_name = data.get("callback_name") enabled = bool(data.get("enabled", True)) metadata = dict(data.get("metadata", {}) or {}) # Decide schedule if not schedule: if interval_seconds is not None: schedule = _schedule_from_interval(int(interval_seconds)) else: schedule = "interval:3600" # Decide callback_name if not callback_name: if handler_name: callback_name = str(handler_name) # Parse times next_run = _parse_dt(data.get("next_run")) or _parse_dt(next_run_at) last_run = _parse_dt(data.get("last_run_at")) last_scheduled_run = _parse_dt(data.get("last_scheduled_run")) # If next_run missing, compute it if next_run is None: next_run = _calculate_next_run(schedule) return cls( job_id=str(job_id), name=str(name), schedule=str(schedule), enabled=enabled, next_run=_ensure_tz(next_run), last_run_at=_ensure_tz(last_run) if last_run else None, last_scheduled_run=_ensure_tz(last_scheduled_run) if last_scheduled_run else None, max_runs=data.get("max_runs"), run_count=int(data.get("run_count", 0) or 0), callback_name=callback_name, workflow_id=data.get("workflow_id"), metadata=metadata, ) # Backwards-compatible export name used in older code JobDefinition = ScheduledJob class Scheduler: """ Job scheduler with persistent storage. You can use it in two ways: - Manually: call run_due_jobs() periodically. - Background: call start() to run due jobs in a daemon thread. Note: Workflow triggering is stubbed. If a job has workflow_id but no callback, it will log and return success, leaving actual workflow integration to the workflow runner layer. """ def __init__( self, jobs_file: Path | None = None, *, storage_path: Path | None = None, now_func: Callable[[], datetime] | None = None, poll_interval_seconds: float = 1.0, ) -> None: if jobs_file is None and storage_path is not None: jobs_file = storage_path self.jobs_file = jobs_file or Path("data/scheduler/jobs.json") self.jobs_file.parent.mkdir(parents=True, exist_ok=True) self._now = now_func or _utc_now self._poll_interval = float(poll_interval_seconds) self._lock = threading.RLock() self._jobs: dict[str, ScheduledJob] = {} self._callbacks: dict[str, JobCallback] = {} self._metrics = SchedulerMetrics() self._running = False self._thread: threading.Thread | None = None self._load_jobs() # ------------------------- # Persistence # ------------------------- def _load_jobs(self) -> None: if not self.jobs_file.exists(): return try: raw = json.loads(self.jobs_file.read_text(encoding="utf-8")) except Exception as e: logger.error("Failed to read jobs file %s: %s", self.jobs_file, e) return # Accept either: # - legacy: {"jobs": [ ... ]} # - newer: [ ... ] items: list[dict[str, Any]] if isinstance(raw, dict) and isinstance(raw.get("jobs"), list): items = raw["jobs"] elif isinstance(raw, list): items = raw else: items = [] loaded = 0 with self._lock: self._jobs.clear() for item in items: if not isinstance(item, dict): continue job = ScheduledJob.from_dict(item) # Ensure next_run is sane if job.next_run is None: job.next_run = _calculate_next_run(job.schedule, from_time=self._now()) self._jobs[job.job_id] = job loaded += 1 logger.info("Loaded %d job(s) from %s", loaded, self.jobs_file) def _save_jobs(self) -> None: with self._lock: payload = [job.to_dict() for job in self._jobs.values()] try: self.jobs_file.parent.mkdir(parents=True, exist_ok=True) self.jobs_file.write_text(json.dumps(payload, indent=2), encoding="utf-8") except Exception as e: logger.error("Failed to save jobs to %s: %s", self.jobs_file, e) # ------------------------- # Registration # ------------------------- def register_handler(self, name: str, handler: JobCallback) -> None: """Legacy name for registering callbacks.""" self.register_callback(name, handler) def register_callback(self, name: str, callback: JobCallback) -> None: with self._lock: self._callbacks[str(name)] = callback logger.debug("Registered scheduler callback: %s", name) @property def callbacks(self) -> dict[str, JobCallback]: """Expose registered callbacks for compatibility with older tests.""" return self._callbacks def _calculate_next_run(self, schedule: str, *, from_time: datetime | None = None) -> datetime: """Compatibility wrapper for test suites expecting this method.""" return _calculate_next_run(schedule, from_time=from_time) def update_next_run(self, job_id: str) -> None: """Recalculate and update a job's next_run based on its schedule.""" job = self._jobs.get(job_id) if job is None: return job.next_run = _calculate_next_run(job.schedule, from_time=self._now()) self._save_jobs() # ------------------------- # CRUD # ------------------------- def add_job( self, *args, **kwargs, ) -> ScheduledJob: """ Supports both call styles. Legacy style: add_job(name, interval_seconds, handler_name, job_id=None, enabled=True, metadata=None) New style: add_job( job_id="...", name="...", schedule="interval:600", enabled=True, callback_name="...", workflow_id=None, max_runs=None, metadata=None, ) """ # Detect legacy positional signature if ( len(args) >= 3 and isinstance(args[0], str) and isinstance(args[1], (int, float)) and isinstance(args[2], str) ): name = args[0] interval_seconds = int(args[1]) handler_name = args[2] job_id = kwargs.get("job_id") or str(uuid.uuid4()) enabled = bool(kwargs.get("enabled", True)) metadata = kwargs.get("metadata") or {} schedule = _schedule_from_interval(interval_seconds) now = _ensure_tz(self._now()) job = ScheduledJob( job_id=str(job_id), name=str(name), schedule=schedule, enabled=enabled, next_run=_calculate_next_run(schedule, from_time=now), last_run_at=None, callback_name=str(handler_name) if handler_name else None, workflow_id=None, max_runs=None, run_count=0, metadata=dict(metadata), ) with self._lock: self._jobs[job.job_id] = job self._save_jobs() return job # New style kwargs job_id = str(kwargs.get("job_id") or str(uuid.uuid4())) name = str(kwargs.get("name") or "Unnamed Job") schedule = str(kwargs.get("schedule") or "interval:3600") enabled = bool(kwargs.get("enabled", True)) callback_name = kwargs.get("callback_name") workflow_id = kwargs.get("workflow_id") max_runs = kwargs.get("max_runs") metadata = kwargs.get("metadata") or {} now = _ensure_tz(self._now()) job = ScheduledJob( job_id=job_id, name=name, schedule=schedule, enabled=enabled, next_run=_calculate_next_run(schedule, from_time=now), last_run_at=None, callback_name=str(callback_name) if callback_name else None, workflow_id=str(workflow_id) if workflow_id else None, max_runs=max_runs, run_count=0, metadata=dict(metadata), ) with self._lock: self._jobs[job.job_id] = job self._save_jobs() return job def remove_job(self, job_id: str) -> bool: with self._lock: existed = job_id in self._jobs if existed: del self._jobs[job_id] if existed: self._save_jobs() return existed def list_jobs(self) -> list[ScheduledJob]: with self._lock: return list(self._jobs.values()) def get_job(self, job_id: str) -> ScheduledJob | None: with self._lock: return self._jobs.get(job_id) def find_job_by_name(self, name: str) -> ScheduledJob | None: with self._lock: for job in self._jobs.values(): if job.name == name: return job return None def update_job(self, job_id: str, **updates: Any) -> ScheduledJob | None: with self._lock: job = self._jobs.get(job_id) if not job: return None # Apply updates conservatively to known fields for key, value in updates.items(): if not hasattr(job, key): continue if key in {"next_run", "last_run_at"} and isinstance(value, str): parsed = _parse_dt(value) if parsed: setattr(job, key, parsed) continue setattr(job, key, value) # If schedule changed, recompute next_run from now if "schedule" in updates: job.next_run = _calculate_next_run(job.schedule, from_time=self._now()) self._save_jobs() return job # ------------------------- # Execution # ------------------------- def run_job(self, job_id: str, *, manual: bool = False, force: bool = False) -> bool: """ Execute a job immediately. Compatibility: - Legacy callers pass manual=True. - New callers pass force=True. Rules: - If force or manual, ignore enabled and max_runs checks. - Otherwise, skip if disabled or max_runs reached. """ with self._lock: job = self._jobs.get(job_id) if not job: logger.warning("Job not found: %s", job_id) return False bypass = bool(force or manual) if not bypass: if not job.enabled: return False if job.max_runs is not None and job.run_count >= job.max_runs: return False scheduled_for = _ensure_tz(job.next_run) if job.last_scheduled_run and _ensure_tz(job.last_scheduled_run) == scheduled_for: logger.info( "Skipping job %s: already ran for scheduled time %s", job.job_id, scheduled_for.isoformat(), ) return False else: scheduled_for = _ensure_tz(job.next_run) callback_name = job.callback_name # Execute outside lock ok = True try: if callback_name and callback_name in self._callbacks: self._callbacks[callback_name](job) elif job.workflow_id: logger.info( "Job %s would trigger workflow %s (integration stub)", job.job_id, job.workflow_id, ) else: logger.warning("Job %s has no callback_name or workflow_id", job.job_id) except Exception as e: ok = False logger.error("Error running job %s: %s", job_id, e, exc_info=True) # Update job timing and counts with self._lock: now = _ensure_tz(self._now()) job.last_run_at = now if not bypass: job.last_scheduled_run = scheduled_for job.run_count += 1 base_time = now if scheduled_for and now <= scheduled_for: base_time = scheduled_for job.next_run = _calculate_next_run(job.schedule, from_time=base_time) self._metrics.total_runs += 1 if ok: self._metrics.successful_runs += 1 self._metrics.last_run_at = now else: self._metrics.failed_runs += 1 self._metrics.last_failure_at = now self._save_jobs() return ok def run_due_jobs(self) -> list[str]: """Run all jobs whose next_run is due. Returns list of job_ids that were executed.""" now = _ensure_tz(self._now()) due: list[str] = [] with self._lock: jobs_snapshot = list(self._jobs.values()) for job in jobs_snapshot: if not job.enabled: continue if job.max_runs is not None and job.run_count >= job.max_runs: continue if _ensure_tz(job.next_run) <= now: if self.run_job(job.job_id): due.append(job.job_id) return due def get_metrics(self) -> dict[str, Any]: """Return scheduler metrics and queue depth.""" with self._lock: now = _ensure_tz(self._now()) due_jobs = [ job for job in self._jobs.values() if job.enabled and (job.max_runs is None or job.run_count < job.max_runs) and _ensure_tz(job.next_run) <= now ] return { "total_jobs": len(self._jobs), "enabled_jobs": sum(1 for job in self._jobs.values() if job.enabled), "due_jobs": len(due_jobs), "total_runs": self._metrics.total_runs, "successful_runs": self._metrics.successful_runs, "failed_runs": self._metrics.failed_runs, "last_run_at": ( self._metrics.last_run_at.isoformat() if self._metrics.last_run_at else None ), "last_failure_at": ( self._metrics.last_failure_at.isoformat() if self._metrics.last_failure_at else None ), } # ------------------------- # Background loop (optional) # ------------------------- def start(self) -> None: if self._running: logger.warning("Scheduler already running") return self._running = True self._thread = threading.Thread(target=self._run_loop, daemon=True, name="RexScheduler") self._thread.start() logger.info("Scheduler started") def stop(self) -> None: if not self._running: return self._running = False t = self._thread if t: t.join(timeout=5.0) logger.info("Scheduler stopped") def _run_loop(self) -> None: logger.info("Scheduler loop starting") while self._running: try: self.run_due_jobs() time.sleep(self._poll_interval) except Exception as e: logger.error("Scheduler loop error: %s", e, exc_info=True) time.sleep(max(self._poll_interval, 5.0)) logger.info("Scheduler loop exiting") # Global scheduler instance _SCHEDULER: Scheduler | None = None _SCHEDULER_LOCK = threading.Lock() def get_scheduler() -> Scheduler: """Get the global scheduler instance.""" global _SCHEDULER if _SCHEDULER is None: with _SCHEDULER_LOCK: if _SCHEDULER is None: _SCHEDULER = Scheduler() return _SCHEDULER def set_scheduler(scheduler: Scheduler) -> None: """Set the global scheduler instance (for testing).""" global _SCHEDULER with _SCHEDULER_LOCK: _SCHEDULER = scheduler __all__ = ["ScheduledJob", "JobDefinition", "Scheduler", "get_scheduler", "set_scheduler"]