tai/tai_pkg/scheduler.py

"""Schedule parsing, display, and the scheduler loop."""

import re
import sqlite3
import threading
from datetime import datetime, timedelta, timezone

from tai_pkg.agents import REGISTRY
from tai_pkg.consts import Consts
from tai_pkg.seal import Seal
from tai_pkg.seal import SealError
from tai_pkg.text import Text

__all__ = ["Scheduler"]


class Scheduler:
    """Scheduler."""

    @staticmethod
    def parse_schedule_at(raw):
        text = str(raw or "").strip()
        if not text:
            raise ValueError("empty datetime")
        if text[-1:] in ("Z", "z"):
            text = text[:-1] + "+00:00"
        try:
            moment = Text._fromisoformat(text)
        except ValueError:
            raise ValueError("invalid datetime, use ISO like 2026-10-08T09:00")
        return moment.astimezone(timezone.utc).isoformat()

    @staticmethod
    def parse_duration(raw):
        match = re.fullmatch(r"(\d+)\s*([smhdw])?", str(raw or "").strip().lower())
        if not match:
            return None
        return int(match.group(1)) * {"s": 1, "m": 60, "h": 3600, "d": 86400, "w": 604800}[match.group(2) or "s"]

    @staticmethod
    def format_delay(seconds):
        seconds = int(seconds)
        if seconds < 0:
            return "overdue"
        if seconds < 60:
            return "%ds" % seconds
        if seconds < 3600:
            return "%dm" % (seconds // 60)
        if seconds < 86400:
            return "%dh" % (seconds // 3600)
        return "%dd" % (seconds // 86400)

    @staticmethod
    def local_display(utc_iso):
        try:
            return Text._fromisoformat(utc_iso).astimezone().strftime("%Y-%m-%d %H:%M")
        except ValueError:
            return utc_iso

    @staticmethod
    def format_schedule_lines(items, now):
        if not items:
            return "(no schedules)"
        lines = []
        for item in items:
            try:
                due_in = (Text._fromisoformat(item["next_run"]) - now).total_seconds()
            except ValueError:
                due_in = -1
            when = "every %s" % Scheduler.format_delay(item["every"]) if item["every"] else "at %s" % Scheduler.local_display(item["next_run"])
            marked = " [%s]" % " ".join(item["tags"]) if item["tags"] else ""
            lines.append("#%d [%s] %s (%s) %s%s :: %s" % (item["id"], item["status"], when, Scheduler.format_delay(due_in), item["name"] or "(no name)", marked, item["prompt"][:80]))
            if item["last"]:
                extra = " :: %s" % item["result"][:120] if item["result"] else ""
                lines.append("    last: %s%s" % (item["last"], extra))
        return "\n".join(lines)

    @staticmethod
    def note_schedule_result(db_path, row_id, status, result=""):
        try:
            db = sqlite3.connect(db_path, timeout=30)
            db.execute("PRAGMA journal_mode=WAL")
            db.execute("PRAGMA busy_timeout=30000")
            try:
                db.execute("UPDATE schedules SET last_status = ?, last_result = ?, updated = ? WHERE id = ?", (status, result[:2000], datetime.now(timezone.utc).isoformat(), row_id))
                db.commit()
            finally:
                db.close()
        except sqlite3.Error:
            pass

    @staticmethod
    def watch_schedule(db_path, row_id, agent_id, timeout):
        status, text = REGISTRY.poll_agent(agent_id, wait=timeout + 30)
        Scheduler.note_schedule_result(db_path, row_id, status, text)

    @staticmethod
    def scheduler_tick(config, seal, runner=None):
        fired = 0
        now = datetime.now(timezone.utc)
        db = sqlite3.connect(config.db_path, timeout=30)
        db.execute("PRAGMA journal_mode=WAL")
        db.execute("PRAGMA busy_timeout=30000")
        try:
            rows = db.execute("SELECT id, prompt, profile, every_sec, next_run, timeout FROM schedules WHERE status = 'pending' AND next_run <= ? ORDER BY next_run", (now.isoformat(),)).fetchall()
            for row_id, prompt, profile, every_sec, next_run, timeout in rows:
                stamp = datetime.now(timezone.utc).isoformat()
                if every_sec:
                    try:
                        claimed_next = Text._fromisoformat(next_run)
                    except ValueError:
                        continue
                    while claimed_next <= datetime.now(timezone.utc):
                        claimed_next += timedelta(seconds=every_sec)
                    done = db.execute("UPDATE schedules SET next_run = ?, updated = ? WHERE id = ? AND status = 'pending' AND next_run = ?", (claimed_next.isoformat(), stamp, row_id, next_run))
                else:
                    done = db.execute("UPDATE schedules SET status = 'done', updated = ? WHERE id = ? AND status = 'pending'", (stamp, row_id))
                db.commit()
                if done.rowcount != 1:
                    continue
                if prompt.startswith(Seal.PREFIX):
                    try:
                        prompt = seal.unlock(prompt)
                    except SealError as exc:
                        Scheduler.note_schedule_result(config.db_path, row_id, "error: " + Text.short_error(exc))
                        continue
                agent_id = REGISTRY.spawn_agent(prompt, profile, timeout, config, seal, 0, runner)
                watcher = threading.Thread(target=Scheduler.watch_schedule, args=(config.db_path, row_id, agent_id, timeout), daemon=True)
                watcher.start()
                fired += 1
        finally:
            db.close()
        return fired

    @staticmethod
    def run_scheduler_loop(config, seal, stop_flag):
        while not stop_flag.is_set():
            try:
                Scheduler.scheduler_tick(config, seal)
            except Exception as exc:
                print("scheduler error: %s" % Text.short_error(exc))
            stop_flag.wait(Consts.SCHEDULER_INTERVAL)

    @staticmethod
    def start_scheduler(config, seal):
        stop_flag = threading.Event()
        worker = threading.Thread(target=Scheduler.run_scheduler_loop, args=(config, seal, stop_flag), daemon=True)
        worker.start()
        return stop_flag