from __future__ import annotations from datetime import datetime from typing import Any from uuid import uuid4 from apscheduler.schedulers.asyncio import AsyncIOScheduler from apscheduler.triggers.cron import CronTrigger from apscheduler.triggers.date import DateTrigger from apscheduler.triggers.interval import IntervalTrigger from tzlocal import get_localzone from traderai.memory import MemoryStore, iso_now, time_since UEX_NOTIFICATION_JOB_ID = "uex-notification-poll" class WakeScheduler: def __init__(self, memory: MemoryStore) -> None: self.memory = memory self.scheduler = AsyncIOScheduler(timezone=get_localzone()) self.agent = None self.uex = None self.notification_poll_seconds = 60 def bind_agent(self, agent: Any) -> None: self.agent = agent def bind_uex_notifications(self, uex: Any, poll_seconds: int = 60) -> None: self.uex = uex self.notification_poll_seconds = max(15, poll_seconds) def start(self) -> None: if not self.scheduler.running: self.scheduler.start() self._schedule_notification_poll() for job in self.memory.list_jobs(): self._schedule_existing(job) def shutdown(self) -> None: if self.scheduler.running: self.scheduler.shutdown(wait=False) def schedule_date(self, run_at: str, prompt: str, job_id: str | None = None) -> dict[str, Any]: parsed = datetime.fromisoformat(run_at) job_id = job_id or f"wake-{uuid4()}" trigger = DateTrigger(run_date=parsed) self.scheduler.add_job(self._run_job, trigger=trigger, id=job_id, args=[job_id, prompt], replace_existing=True) return self.memory.add_job(job_id, prompt, "date", run_at, parsed.isoformat()) def schedule_cron(self, cron: str, prompt: str, job_id: str | None = None) -> dict[str, Any]: job_id = job_id or f"wake-{uuid4()}" trigger = CronTrigger.from_crontab(cron) self.scheduler.add_job(self._run_job, trigger=trigger, id=job_id, args=[job_id, prompt], replace_existing=True) next_run = self.scheduler.get_job(job_id).next_run_time return self.memory.add_job(job_id, prompt, "cron", cron, next_run.isoformat() if next_run else None) def list_jobs(self) -> list[dict[str, Any]]: return self.memory.list_jobs() def _schedule_existing(self, job: dict[str, Any]) -> None: if job["trigger_type"] == "cron": trigger = CronTrigger.from_crontab(job["trigger_value"]) elif job["trigger_type"] == "date": trigger = DateTrigger(run_date=datetime.fromisoformat(job["trigger_value"])) else: return self.scheduler.add_job( self._run_job, trigger=trigger, id=job["id"], args=[job["id"], job["prompt"]], replace_existing=True, ) async def _run_job(self, job_id: str, prompt: str) -> None: last = self.memory.last_interaction() last_text = f"{last['created_at']} ({time_since(last['created_at'])})" if last else "never" wake_message = ( f"Scheduled wake job fired. Current time is {iso_now()}. " f"The last chat interaction was {last_text}. Job instruction: {prompt}" ) if self.agent is None: self.memory.add_outbox(wake_message) self._mark_job_finished(job_id) return try: text = await self.agent.generate_wake_response(wake_message) except Exception as exc: text = f"Wake job failed: {exc}. Job instruction: {prompt}" self.memory.add_outbox(text) self._mark_job_finished(job_id) def _mark_job_finished(self, job_id: str) -> None: job = self.scheduler.get_job(job_id) next_run = job.next_run_time if job else None self.memory.mark_job_run(job_id, next_run.isoformat() if next_run else None, enabled=bool(next_run)) def _schedule_notification_poll(self) -> None: if self.uex is None: return self.scheduler.add_job( self.poll_uex_notifications, trigger=IntervalTrigger(seconds=self.notification_poll_seconds), id=UEX_NOTIFICATION_JOB_ID, replace_existing=True, next_run_time=datetime.now(), ) async def poll_uex_notifications(self) -> list[dict[str, Any]]: if self.uex is None: return [] try: response = await self.uex.get_user_notifications() except Exception as exc: self.memory.add_outbox(f"UEX notification poll failed: {exc}") self.memory.set_profile("uex_last_notification_error", str(exc)) return [] notifications = response.get("notifications") or [] pending = [item for item in notifications if not item.get("date_read")] profile = self.memory.get_profile() seen = set(profile.get("uex_seen_notification_keys") or []) new_pending = [item for item in pending if self._notification_key(item) not in seen] if new_pending: for item in new_pending: self.memory.add_outbox(self._notification_text(item)) seen.update(self._notification_key(item) for item in new_pending) self.memory.set_profile("uex_seen_notification_keys", sorted(seen)) self.memory.set_profile("uex_last_notification_check", iso_now()) elif notifications: seen.update(self._notification_key(item) for item in pending) self.memory.set_profile("uex_seen_notification_keys", sorted(seen)) self.memory.set_profile("uex_last_notification_check", iso_now()) return new_pending @staticmethod def _notification_key(item: dict[str, Any]) -> str: for key in ("code", "id"): value = item.get(key) if value not in (None, ""): return f"{key}:{value}" return f"notification:{item.get('date_added')}:{item.get('message')}" @staticmethod def _notification_text(item: dict[str, Any]) -> str: message = item.get("message") or "You have a pending UEX notification." redir = item.get("redir") code = item.get("code") details = [] if code: details.append(f"code `{code}`") if redir: details.append(f"path `{redir}`") suffix = f" ({', '.join(details)})" if details else "" return f"UEX notification: {message}{suffix}"