bub-schedule 0.0.1__tar.gz

bub_schedule-0.0.1/PKG-INFO

@@ -0,0 +1,9 @@
+Metadata-Version: 2.3
+Name: bub-schedule
+Version: 0.0.1
+Summary: Scheduling tools for Bub
+Author: Frost Ming
+Author-email: Frost Ming <me@frostming.com>
+Requires-Dist: apscheduler>=3.11.2
+Requires-Dist: republic>=0.5.6
+Requires-Python: >=3.12

bub_schedule-0.0.1/pyproject.toml

@@ -0,0 +1,28 @@
+[project]
+name = "bub-schedule"
+version = "0.0.1"
+description = "Scheduling tools for Bub"
+authors = [
+    { name = "Frost Ming", email = "me@frostming.com" }
+]
+requires-python = ">=3.12"
+dependencies = [
+    "apscheduler>=3.11.2",
+    "republic>=0.5.6",
+]
+
+[project.entry-points.bub]
+schedule = "bub_schedule.plugin:ScheduleImpl"
+
+[build-system]
+requires = ["uv_build>=0.9.7,<0.10.0"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "pytest>=9.0.3",
+    "pytest-asyncio>=1.3.0",
+]
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"

bub_schedule-0.0.1/src/bub_schedule/__init__.py

@@ -0,0 +1,2 @@
+def hello() -> str:
+    return "Hello from bub-schedule!"

bub_schedule-0.0.1/src/bub_schedule/channel.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+import asyncio
+from asyncio import Event
+
+from apscheduler.schedulers.base import BaseScheduler
+from bub.channels import Lifecycle
+from bub.framework import BubFramework
+from loguru import logger
+
+
+class ScheduleChannel(Lifecycle):
+    name = "schedule"
+
+    # Class-level runtime state (singleton per process)
+    _framework: BubFramework | None = None
+
+    def __init__(self, scheduler: BaseScheduler, *, framework: BubFramework) -> None:
+        self.scheduler = scheduler
+        self._instance_framework = framework
+
+    @classmethod
+    def current_framework(cls) -> BubFramework:
+        """Return the live framework bound to the current gateway process."""
+        if cls._framework is None:
+            raise RuntimeError(
+                "no live schedule framework available, cannot deliver scheduled message"
+            )
+        return cls._framework
+
+    async def start(self, stop_event: Event) -> None:
+        ScheduleChannel._framework = self._instance_framework
+
+        loop = asyncio.get_running_loop()
+        loop.call_soon_threadsafe(self.scheduler.start)
+        logger.info("schedule.start complete")
+
+    async def stop(self) -> None:
+        loop = asyncio.get_running_loop()
+        loop.call_soon_threadsafe(self.scheduler.shutdown)
+
+        ScheduleChannel._framework = None
+        logger.info("schedule.stop complete")

bub_schedule-0.0.1/src/bub_schedule/jobs.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from bub.channels.message import ChannelMessage
+
+SCHEDULE_SUBPROCESS_TIMEOUT_SECONDS = 300
+
+
+async def run_scheduled_reminder(
+    message: str, session_id: str, workspace: str | None = None
+) -> None:
+    from bub_schedule.channel import ScheduleChannel
+
+    if ":" in session_id:
+        channel, chat_id = session_id.split(":", 1)
+    else:
+        channel = "schedule"
+        chat_id = "default"
+
+    payload = ChannelMessage(
+        content=message,
+        session_id=session_id,
+        channel=channel,
+        chat_id=chat_id,
+    )
+    framework = ScheduleChannel.current_framework()
+    await framework.process_inbound(payload)

bub_schedule-0.0.1/src/bub_schedule/jobstore.py

@@ -0,0 +1,145 @@
+import base64
+import json
+import pickle
+import threading
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+from apscheduler.job import Job
+from apscheduler.jobstores.base import BaseJobStore, ConflictingIdError, JobLookupError
+from loguru import logger
+
+
+class JSONJobStore(BaseJobStore):
+    """
+    A simple JSON-based job store for APScheduler.
+
+    Jobs are serialized using pickle and stored as base64-encoded strings in a JSON file.
+    This provides persistence across restarts without requiring SQLAlchemy.
+    """
+
+    def __init__(self, file_path: str | Path):
+        super().__init__()
+        self.file_path = Path(file_path)
+        self._lock = threading.RLock()
+        self._jobs: dict[str, dict[str, Any]] = self._load()
+
+    def _load(self) -> dict[str, dict[str, Any]]:
+        """Load jobs from JSON file."""
+        if self.file_path.exists():
+            try:
+                with open(self.file_path, encoding="utf-8") as f:
+                    loaded_jobs = json.load(f)
+                    return loaded_jobs  # type: ignore[no-any-return]
+            except (OSError, json.JSONDecodeError) as e:
+                logger.error(f"Error loading job store: {e}")
+        return {}
+
+    def _save(self):
+        """Save jobs to JSON file."""
+        self.file_path.parent.mkdir(parents=True, exist_ok=True)
+        try:
+            with open(self.file_path, "w", encoding="utf-8") as f:
+                json.dump(self._jobs, f, ensure_ascii=False, indent=2)
+        except OSError as e:
+            logger.error(f"Error saving job store: {e}")
+
+    def _serialize_job(self, job: Job) -> dict[str, Any]:
+        """Serialize a job to a storable format."""
+        return {
+            "id": job.id,
+            "data": base64.b64encode(pickle.dumps(job)).decode("ascii"),
+            "next_run_time": (
+                job.next_run_time.isoformat() if job.next_run_time else None
+            ),
+        }
+
+    def _deserialize_job(self, job_data: dict[str, Any]) -> Job | None:
+        """Deserialize a job from stored format."""
+        try:
+            job = pickle.loads(base64.b64decode(job_data["data"]))  # noqa: S301
+            job._scheduler = self._scheduler
+            job._jobstore_alias = self._alias
+        except Exception as e:
+            logger.error(f"Error deserializing job {job_data.get('id')}: {e}")
+            return None
+        else:
+            return job
+
+    def shutdown(self):
+        """Called when the scheduler shuts down."""
+        with self._lock:
+            self._save()
+
+    def lookup_job(self, job_id: str) -> Job | None:
+        """Look up a job by its ID."""
+        with self._lock:
+            job_data = self._jobs.get(job_id)
+            if job_data:
+                return self._deserialize_job(job_data)
+            return None
+
+    def get_due_jobs(self, now: datetime) -> list[Job]:
+        """Get jobs that are due to be run."""
+        with self._lock:
+            due_jobs = []
+            for job_data in self._jobs.values():
+                next_run_time_str = job_data.get("next_run_time")
+                if next_run_time_str:
+                    next_run_time = datetime.fromisoformat(next_run_time_str)
+                    if next_run_time <= now:
+                        job = self._deserialize_job(job_data)
+                        if job:
+                            due_jobs.append(job)
+            return due_jobs
+
+    def get_next_run_time(self) -> datetime | None:
+        """Get the earliest next run time among all jobs."""
+        with self._lock:
+            next_times = []
+            for job_data in self._jobs.values():
+                next_run_time_str = job_data.get("next_run_time")
+                if next_run_time_str:
+                    next_times.append(datetime.fromisoformat(next_run_time_str))
+            return min(next_times) if next_times else None
+
+    def get_all_jobs(self) -> list[Job]:
+        """Get all jobs in the store."""
+        with self._lock:
+            jobs = []
+            for job_data in self._jobs.values():
+                job = self._deserialize_job(job_data)
+                if job:
+                    jobs.append(job)
+            return jobs
+
+    def add_job(self, job: Job):
+        """Add a job to the store."""
+        with self._lock:
+            if job.id in self._jobs:
+                raise ConflictingIdError(job.id)
+            self._jobs[job.id] = self._serialize_job(job)
+            self._save()
+
+    def update_job(self, job: Job):
+        """Update a job in the store."""
+        with self._lock:
+            if job.id not in self._jobs:
+                raise JobLookupError(job.id)
+            self._jobs[job.id] = self._serialize_job(job)
+            self._save()
+
+    def remove_job(self, job_id: str):
+        """Remove a job from the store."""
+        with self._lock:
+            if job_id not in self._jobs:
+                raise JobLookupError(job_id)
+            del self._jobs[job_id]
+            self._save()
+
+    def remove_all_jobs(self):
+        """Remove all jobs from the store."""
+        with self._lock:
+            self._jobs.clear()
+            self._save()

bub_schedule-0.0.1/src/bub_schedule/plugin.py

@@ -0,0 +1,33 @@
+import bub
+from apscheduler.schedulers.asyncio import AsyncIOScheduler
+from apscheduler.schedulers.base import BaseScheduler
+from bub import hookimpl
+from bub.channels import Channel
+from bub.framework import BubFramework
+from bub.types import Envelope, MessageHandler, State
+
+from bub_schedule.jobstore import JSONJobStore
+
+
+def default_scheduler() -> BaseScheduler:
+    job_file = bub.home / "jobs.json"
+    job_store = JSONJobStore(job_file)
+    return AsyncIOScheduler(jobstores={"default": job_store})
+
+
+class ScheduleImpl:
+    def __init__(self, framework: BubFramework | None = None) -> None:
+        from bub_schedule import tools  # noqa: F401
+
+        self.framework = framework
+        self.scheduler = default_scheduler()
+
+    @hookimpl
+    def load_state(self, message: Envelope, session_id: str) -> State:
+        return {"scheduler": self.scheduler}
+
+    @hookimpl
+    def provide_channels(self, message_handler: MessageHandler) -> list[Channel]:
+        from bub_schedule.channel import ScheduleChannel
+
+        return [ScheduleChannel(self.scheduler, framework=self.framework)]

bub_schedule-0.0.1/src/bub_schedule/tools.py

@@ -0,0 +1,142 @@
+import uuid
+from datetime import UTC, datetime, timedelta
+from typing import cast
+
+from apscheduler.jobstores.base import ConflictingIdError, JobLookupError
+from apscheduler.schedulers.base import BaseScheduler
+from apscheduler.triggers.cron import CronTrigger
+from apscheduler.triggers.date import DateTrigger
+from apscheduler.triggers.interval import IntervalTrigger
+from bub import tool
+from pydantic import BaseModel, Field
+from republic import ToolContext
+
+from bub_schedule.jobs import run_scheduled_reminder
+
+
+def _ensure_scheduler(state: dict) -> BaseScheduler:
+    if "scheduler" not in state:
+        raise RuntimeError(
+            "scheduler not found in state, is ScheduleImpl plugin loaded?"
+        )
+    return cast(BaseScheduler, state["scheduler"])
+
+
+class ScheduleAddInput(BaseModel):
+    after_seconds: int | None = Field(
+        None, description="If set, schedule to run after this many seconds from now"
+    )
+    interval_seconds: int | None = Field(
+        None, description="If set, repeat at this interval"
+    )
+    cron: str | None = Field(
+        None,
+        description="If set, run with cron expression in crontab format: minute hour day month day_of_week",
+    )
+    message: str = Field(
+        ...,
+        description="Reminder message to send, prefix the message with ',' to run a bash command instead",
+    )
+
+
+@tool(name="schedule.add", context=True, model=ScheduleAddInput)
+def schedule_add(params: ScheduleAddInput, context: ToolContext) -> str:
+    """Schedule a reminder message to be sent to current session in the future."""
+    job_id = str(uuid.uuid4())[:8]
+    if params.after_seconds is not None:
+        trigger = DateTrigger(
+            run_date=datetime.now(UTC) + timedelta(seconds=params.after_seconds)
+        )
+    elif params.interval_seconds is not None:
+        trigger = IntervalTrigger(seconds=params.interval_seconds)
+    else:
+        try:
+            trigger = CronTrigger.from_crontab(params.cron)
+        except ValueError as exc:
+            raise RuntimeError(f"invalid cron expression: {params.cron}") from exc
+    scheduler = _ensure_scheduler(context.state)
+    workspace = context.state.get("_runtime_workspace")
+    try:
+        job = scheduler.add_job(
+            run_scheduled_reminder,
+            trigger=trigger,
+            id=job_id,
+            kwargs={
+                "message": params.message,
+                "session_id": context.state.get("session_id", ""),
+                "workspace": str(workspace) if workspace else None,
+            },
+            coalesce=True,
+            max_instances=1,
+        )
+    except ConflictingIdError as exc:
+        raise RuntimeError(f"job id already exists: {job_id}") from exc
+
+    next_run = "-"
+    if isinstance(job.next_run_time, datetime):
+        next_run = job.next_run_time.isoformat()
+    return f"scheduled: {job.id} next={next_run}"
+
+
+@tool(name="schedule.remove", context=True)
+def schedule_remove(job_id: str, context: ToolContext) -> str:
+    """Remove one scheduled job by id."""
+    scheduler = _ensure_scheduler(context.state)
+    try:
+        scheduler.remove_job(job_id)
+    except JobLookupError as exc:
+        raise RuntimeError(f"job not found: {job_id}") from exc
+    return f"removed: {job_id}"
+
+
+@tool(name="schedule.list", context=True)
+def schedule_list(context: ToolContext) -> str:
+    """List scheduled jobs for current workspace."""
+    scheduler = _ensure_scheduler(context.state)
+    jobs = scheduler.get_jobs()
+    rows: list[str] = []
+    for job in jobs:
+        next_run = "-"
+        if isinstance(job.next_run_time, datetime):
+            next_run = job.next_run_time.isoformat()
+        message = str(job.kwargs.get("message", ""))
+        job_session = job.kwargs.get("session_id")
+        if job_session and job_session != context.state.get("session_id", ""):
+            continue
+        rows.append(f"{job.id} next={next_run} msg={message}")
+
+    if not rows:
+        return "(no scheduled jobs)"
+
+    return "\n".join(rows)
+
+
+@tool(name="schedule.trigger", context=True)
+async def schedule_trigger(job_id: str, context: ToolContext) -> str:
+    """Manually trigger a scheduled job to run immediately.
+    
+    Executes the job function directly without modifying the schedule.
+    The next scheduled run remains unchanged.
+    """
+    import inspect
+    
+    scheduler = _ensure_scheduler(context.state)
+    try:
+        job = scheduler.get_job(job_id)
+        if job is None:
+            raise RuntimeError(f"job not found: {job_id}")
+        
+        # Execute job function directly, preserving original schedule
+        result = job.func(*job.args, **job.kwargs)
+        
+        # Handle async job functions
+        if inspect.iscoroutine(result):
+            await result
+        
+        next_run = "-"
+        if isinstance(job.next_run_time, datetime):
+            next_run = job.next_run_time.isoformat()
+        
+        return f"triggered: {job_id} (next scheduled run: {next_run})"
+    except JobLookupError as exc:
+        raise RuntimeError(f"job not found: {job_id}") from exc