symbiont/symbiont/scheduler.py

"""
Scheduler: Manages task queues and self-waking timers.

When all models are rate-limited, the scheduler sets a system timer
to wake the orchestrator back up when limits expire.
"""

import json
import logging
import subprocess
from datetime import datetime, timedelta
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)

QUEUE_PATH = Path("/data/symbiont/queue.jsonl")
TIMER_NAME = "symbiont-wake"


def enqueue_task(task: str, priority: int = 5, metadata: Optional[dict] = None):
    """Add a task to the persistent queue."""
    QUEUE_PATH.parent.mkdir(parents=True, exist_ok=True)

    entry = {
        "id": datetime.now().strftime("%Y%m%d%H%M%S%f"),
        "task": task,
        "priority": priority,
        "status": "pending",
        "created_at": datetime.now().isoformat(),
        "metadata": metadata or {},
    }

    with open(QUEUE_PATH, "a") as f:
        f.write(json.dumps(entry) + "\n")

    logger.info(f"Enqueued task {entry['id']}: {task[:80]}...")
    return entry["id"]


def get_pending_tasks() -> list[dict]:
    """Read all pending tasks from the queue, sorted by priority."""
    if not QUEUE_PATH.exists():
        return []

    tasks = []
    with open(QUEUE_PATH) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                entry = json.loads(line)
                if entry.get("status") == "pending":
                    tasks.append(entry)
            except json.JSONDecodeError:
                continue

    return sorted(tasks, key=lambda t: t.get("priority", 5))


def mark_task_done(task_id: str, result: Optional[str] = None):
    """Mark a task as completed in the queue file."""
    if not QUEUE_PATH.exists():
        return

    lines = QUEUE_PATH.read_text().strip().split("\n")
    updated = []
    for line in lines:
        if not line.strip():
            continue
        try:
            entry = json.loads(line)
            if entry.get("id") == task_id:
                entry["status"] = "completed"
                entry["completed_at"] = datetime.now().isoformat()
                if result:
                    entry["result_preview"] = result[:500]
            updated.append(json.dumps(entry))
        except json.JSONDecodeError:
            updated.append(line)

    QUEUE_PATH.write_text("\n".join(updated) + "\n")


def schedule_wake(wake_at: datetime):
    """
    Create a systemd transient timer to wake the orchestrator.
    This is how Symbiont sleeps through rate limits instead of busy-waiting.
    """
    delay_seconds = max(1, int((wake_at - datetime.now()).total_seconds()))

    logger.info(f"Scheduling wake in {delay_seconds}s (at {wake_at.isoformat()})")

    # Use systemd-run to create a one-shot timer that runs our wake script
    cmd = [
        "systemd-run",
        "--unit", TIMER_NAME,
        "--on-active", f"{delay_seconds}s",
        "--description", "Symbiont self-wake after rate limit",
        "/usr/bin/python3", "-m", "symbiont.wake",
    ]

    try:
        subprocess.run(cmd, capture_output=True, text=True, check=True)
        logger.info(f"Wake timer set: {TIMER_NAME}")
    except subprocess.CalledProcessError as e:
        logger.error(f"Failed to set wake timer: {e.stderr}")
        # Fallback: write a cron-style at job
        _fallback_schedule(delay_seconds)


def _fallback_schedule(delay_seconds: int):
    """Fallback: use `at` command if systemd-run fails."""
    try:
        proc = subprocess.run(
            ["at", f"now + {max(1, delay_seconds // 60)} minutes"],
            input="cd /data/symbiont && python3 -m symbiont.wake\n",
            capture_output=True,
            text=True,
        )
        if proc.returncode == 0:
            logger.info("Fallback wake scheduled via `at`")
        else:
            logger.error(f"Fallback scheduling failed: {proc.stderr}")
    except FileNotFoundError:
        logger.error("`at` command not available. Cannot schedule wake.")


def cancel_wake():
    """Cancel a pending wake timer."""
    try:
        subprocess.run(
            ["systemctl", "stop", f"{TIMER_NAME}.timer"],
            capture_output=True,
            text=True,
        )
        logger.info("Wake timer cancelled")
    except Exception:
        pass