apollo/
cron_scheduler.rs

1//! SurrealDB-backed Cron Scheduler — persistent scheduled tasks.
2//! Stores jobs in SurrealDB, ticks every 60s, spawns agent sessions for due jobs.
3
4use crate::memory::surreal::SurrealMemory;
5use serde::{Deserialize, Serialize};
6use std::str::FromStr;
7use std::sync::Arc;
8
9/// A cron job stored in SurrealDB.
10#[derive(Debug, Clone, Serialize, Deserialize)]
11pub struct CronJob {
12    pub id: Option<String>,
13    pub name: String,
14    pub schedule: String,
15    #[serde(default)]
16    pub one_shot: bool,
17    pub task: String,
18    pub channel: String,
19    #[serde(default)]
20    pub chat_id: String,
21    pub model: String,
22    pub enabled: bool,
23    pub last_run: Option<String>,
24    pub next_run: Option<String>,
25    #[serde(default = "default_job_status")]
26    pub status: String,
27    #[serde(default)]
28    pub retry_count: u32,
29    #[serde(default = "default_max_retries")]
30    pub max_retries: u32,
31    #[serde(default)]
32    pub last_error: Option<String>,
33    #[serde(default)]
34    pub lease_until: Option<String>,
35    #[serde(default)]
36    pub run_token: Option<String>,
37}
38
39fn default_job_status() -> String {
40    "active".to_string()
41}
42
43fn default_max_retries() -> u32 {
44    3
45}
46
47/// SurrealDB-backed cron scheduler.
48pub struct CronScheduler {
49    memory: Option<Arc<SurrealMemory>>,
50    /// In-memory fallback for noop/testing — keyed by name
51    noop_jobs: std::sync::Mutex<Vec<CronJob>>,
52}
53
54impl CronScheduler {
55    /// Create with a SurrealDB memory backend.
56    pub fn new(memory: Arc<SurrealMemory>) -> Self {
57        Self::new_maybe(Some(memory))
58    }
59
60    /// Create without memory backend (in-memory only).
61    pub fn new_noop() -> Self {
62        Self::new_maybe(None)
63    }
64
65    fn new_maybe(memory: Option<Arc<SurrealMemory>>) -> Self {
66        Self {
67            memory,
68            noop_jobs: std::sync::Mutex::new(Vec::new()),
69        }
70    }
71
72    /// Add a new cron job. Returns the job ID.
73    pub async fn add(
74        &self,
75        name: &str,
76        schedule: &str,
77        task: &str,
78        channel: &str,
79        chat_id: &str,
80        model: &str,
81    ) -> anyhow::Result<String> {
82        // Validate cron expression
83        let parsed = cron::Schedule::from_str(schedule)
84            .map_err(|e| anyhow::anyhow!("Invalid cron expression: {}", e))?;
85
86        // Compute next run
87        let next_run = parsed.upcoming(chrono::Utc).next().map(|t| t.to_rfc3339());
88
89        let job = CronJob {
90            id: Some(name.to_string()),
91            name: name.to_string(),
92            schedule: schedule.to_string(),
93            one_shot: false,
94            task: task.to_string(),
95            channel: channel.to_string(),
96            chat_id: chat_id.to_string(),
97            model: model.to_string(),
98            enabled: true,
99            last_run: None,
100            next_run,
101            status: default_job_status(),
102            retry_count: 0,
103            max_retries: default_max_retries(),
104            last_error: None,
105            lease_until: None,
106            run_token: None,
107        };
108
109        if let Some(ref memory) = self.memory {
110            let db = memory.db();
111            let created: Option<CronJob> = db.create("cron_jobs").content(job).await?;
112            let created = created.ok_or_else(|| anyhow::anyhow!("Failed to create cron job"))?;
113            Ok(created
114                .id
115                .ok_or_else(|| anyhow::anyhow!("Created cron job is missing an id"))?
116                .to_string())
117        } else {
118            let mut jobs = self.noop_jobs.lock().unwrap();
119            jobs.push(job);
120            Ok(name.to_string())
121        }
122    }
123
124    pub async fn add_once(
125        &self,
126        name: &str,
127        run_at: chrono::DateTime<chrono::Utc>,
128        task: &str,
129        channel: &str,
130        chat_id: &str,
131        model: &str,
132    ) -> anyhow::Result<String> {
133        if run_at <= chrono::Utc::now() {
134            anyhow::bail!("run_at must be in the future");
135        }
136        let job = CronJob {
137            id: Some(name.to_string()),
138            name: name.to_string(),
139            schedule: String::new(),
140            one_shot: true,
141            task: task.to_string(),
142            channel: channel.to_string(),
143            chat_id: chat_id.to_string(),
144            model: model.to_string(),
145            enabled: true,
146            last_run: None,
147            next_run: Some(run_at.to_rfc3339()),
148            status: default_job_status(),
149            retry_count: 0,
150            max_retries: default_max_retries(),
151            last_error: None,
152            lease_until: None,
153            run_token: None,
154        };
155        if let Some(ref memory) = self.memory {
156            let created: Option<CronJob> = memory.db().create("cron_jobs").content(job).await?;
157            Ok(created
158                .and_then(|job| job.id)
159                .ok_or_else(|| anyhow::anyhow!("Created one-shot job is missing an id"))?)
160        } else {
161            self.noop_jobs.lock().unwrap().push(job);
162            Ok(name.to_string())
163        }
164    }
165
166    /// List all cron jobs.
167    pub async fn list(&self) -> anyhow::Result<Vec<CronJob>> {
168        if let Some(ref memory) = self.memory {
169            let db = memory.db();
170            let mut result: surrealdb::Response =
171                db.query("SELECT * FROM cron_jobs ORDER BY name").await?;
172            let jobs: Vec<CronJob> = result.take(0)?;
173            Ok(jobs)
174        } else {
175            let jobs = self.noop_jobs.lock().unwrap();
176            Ok(jobs.clone())
177        }
178    }
179
180    /// Remove a cron job by ID or name.
181    pub async fn remove(&self, id_or_name: &str) -> anyhow::Result<bool> {
182        if let Some(ref memory) = self.memory {
183            let db = memory.db();
184            let mut result: surrealdb::Response = db
185                .query("DELETE FROM cron_jobs WHERE id = $target OR name = $target")
186                .bind(("target", id_or_name.to_string()))
187                .await?;
188            let deleted: Vec<CronJob> = result.take(0)?;
189            Ok(!deleted.is_empty())
190        } else {
191            let mut jobs = self.noop_jobs.lock().unwrap();
192            let before = jobs.len();
193            jobs.retain(|j| j.name != id_or_name);
194            Ok(jobs.len() < before)
195        }
196    }
197
198    /// Enable a cron job.
199    pub async fn enable(&self, id_or_name: &str) -> anyhow::Result<bool> {
200        if let Some(ref memory) = self.memory {
201            let db = memory.db();
202            let mut result: surrealdb::Response = db
203                .query("UPDATE cron_jobs SET enabled = true WHERE id = $target OR name = $target")
204                .bind(("target", id_or_name.to_string()))
205                .await?;
206            let updated: Vec<CronJob> = result.take(0)?;
207            Ok(!updated.is_empty())
208        } else {
209            let mut jobs = self.noop_jobs.lock().unwrap();
210            if let Some(job) = jobs.iter_mut().find(|j| j.name == id_or_name) {
211                job.enabled = true;
212                Ok(true)
213            } else {
214                Ok(false)
215            }
216        }
217    }
218
219    /// Disable a cron job.
220    pub async fn disable(&self, id_or_name: &str) -> anyhow::Result<bool> {
221        if let Some(ref memory) = self.memory {
222            let db = memory.db();
223            let mut result: surrealdb::Response = db
224                .query("UPDATE cron_jobs SET enabled = false WHERE id = $target OR name = $target")
225                .bind(("target", id_or_name.to_string()))
226                .await?;
227            let updated: Vec<CronJob> = result.take(0)?;
228            Ok(!updated.is_empty())
229        } else {
230            let mut jobs = self.noop_jobs.lock().unwrap();
231            if let Some(job) = jobs.iter_mut().find(|j| j.name == id_or_name) {
232                job.enabled = false;
233                Ok(true)
234            } else {
235                Ok(false)
236            }
237        }
238    }
239
240    /// Get due jobs (next_run <= now, enabled).
241    pub async fn due_jobs(&self) -> anyhow::Result<Vec<CronJob>> {
242        let now = chrono::Utc::now().to_rfc3339();
243        if let Some(ref memory) = self.memory {
244            let db = memory.db();
245            let mut result: surrealdb::Response = db
246                .query("SELECT * FROM cron_jobs WHERE enabled = true AND ((status = 'active' OR status = NONE) AND next_run != NONE AND next_run <= $now OR status = 'running' AND lease_until != NONE AND lease_until <= $now)")
247                .bind(("now", now))
248                .await?;
249            let jobs: Vec<CronJob> = result.take(0)?;
250            Ok(jobs)
251        } else {
252            let jobs = self.noop_jobs.lock().unwrap();
253            Ok(jobs
254                .iter()
255                .filter(|j| {
256                    j.enabled
257                        && ((j.status == "active" && j.next_run.as_deref() <= Some(&now))
258                            || (j.status == "running" && j.lease_until.as_deref() <= Some(&now)))
259                })
260                .cloned()
261                .collect())
262        }
263    }
264
265    pub async fn claim_due_jobs(&self, channel: &str) -> anyhow::Result<Vec<CronJob>> {
266        let due = self.due_jobs().await?;
267        let claim_now = chrono::Utc::now().to_rfc3339();
268        let lease_until = (chrono::Utc::now() + chrono::Duration::minutes(10)).to_rfc3339();
269        let mut claimed = Vec::new();
270        for job in due {
271            if job.channel != channel {
272                continue;
273            }
274            let Some(job_id) = job.id.as_deref() else {
275                continue;
276            };
277            let run_token = uuid::Uuid::new_v4().to_string();
278            if let Some(ref memory) = self.memory {
279                let mut result = memory
280                    .db()
281                    .query("UPDATE cron_jobs SET status = 'running', lease_until = $lease, run_token = $run_token WHERE id = $id AND enabled = true AND (status = 'active' OR status = NONE OR status = 'running' AND lease_until <= $now) RETURN AFTER")
282                    .bind(("id", job_id.to_string()))
283                    .bind(("lease", lease_until.clone()))
284                    .bind(("run_token", run_token.clone()))
285                    .bind(("now", claim_now.clone()))
286                    .await?;
287                let updated: Vec<CronJob> = result.take(0)?;
288                claimed.extend(updated);
289            } else {
290                let mut jobs = self.noop_jobs.lock().unwrap();
291                if let Some(stored) = jobs.iter_mut().find(|stored| {
292                    stored.id.as_deref() == Some(job_id)
293                        && (stored.status == "active"
294                            || stored.lease_until.as_deref() <= Some(claim_now.as_str()))
295                }) {
296                    stored.status = "running".to_string();
297                    stored.lease_until = Some(lease_until.clone());
298                    stored.run_token = Some(run_token);
299                    claimed.push(stored.clone());
300                }
301            }
302        }
303        Ok(claimed)
304    }
305
306    /// Mark a job as just run and compute next_run.
307    pub async fn mark_run(
308        &self,
309        job_id: &str,
310        run_token: &str,
311        schedule: &str,
312    ) -> anyhow::Result<()> {
313        let now = chrono::Utc::now();
314        let next_run = cron::Schedule::from_str(schedule)
315            .ok()
316            .and_then(|s| s.upcoming(chrono::Utc).next())
317            .map(|t| t.to_rfc3339());
318        let status = if schedule.is_empty() {
319            "completed"
320        } else {
321            "active"
322        };
323
324        if let Some(ref memory) = self.memory {
325            let db = memory.db();
326            let _ = db
327                .query("UPDATE cron_jobs SET last_run = $last, next_run = $next, status = $status, retry_count = 0, last_error = NONE, lease_until = NONE, run_token = NONE WHERE id = $id AND status = 'running' AND run_token = $run_token")
328                .bind(("last", now.to_rfc3339()))
329                .bind(("next", next_run))
330                .bind(("status", status.to_string()))
331                .bind(("id", job_id.to_string()))
332                .bind(("run_token", run_token.to_string()))
333                .await?;
334        } else {
335            let mut jobs = self.noop_jobs.lock().unwrap();
336            if let Some(job) = jobs.iter_mut().find(|job| {
337                job.id.as_deref() == Some(job_id) && job.run_token.as_deref() == Some(run_token)
338            }) {
339                job.last_run = Some(now.to_rfc3339());
340                job.next_run = next_run;
341                job.status = status.to_string();
342                job.retry_count = 0;
343                job.last_error = None;
344                job.lease_until = None;
345                job.run_token = None;
346            }
347        }
348        Ok(())
349    }
350
351    pub async fn release_run(&self, job_id: &str, run_token: &str) -> anyhow::Result<()> {
352        if let Some(ref memory) = self.memory {
353            let _ = memory
354                .db()
355                .query("UPDATE cron_jobs SET status = 'active', lease_until = NONE, run_token = NONE WHERE id = $id AND status = 'running' AND run_token = $run_token")
356                .bind(("id", job_id.to_string()))
357                .bind(("run_token", run_token.to_string()))
358                .await?;
359        } else {
360            let mut jobs = self.noop_jobs.lock().unwrap();
361            if let Some(job) = jobs.iter_mut().find(|job| {
362                job.id.as_deref() == Some(job_id) && job.run_token.as_deref() == Some(run_token)
363            }) {
364                job.status = "active".to_string();
365                job.lease_until = None;
366                job.run_token = None;
367            }
368        }
369        Ok(())
370    }
371
372    pub async fn fail_run(&self, job_id: &str, run_token: &str, error: &str) -> anyhow::Result<()> {
373        let redacted = crate::redaction::redact_text(error);
374        if let Some(ref memory) = self.memory {
375            let db = memory.db();
376            let Some(job) = self.list().await?.into_iter().find(|job| {
377                job.id.as_deref() == Some(job_id) && job.run_token.as_deref() == Some(run_token)
378            }) else {
379                return Ok(());
380            };
381            let retry_count = job.retry_count.saturating_add(1);
382            let status = if retry_count >= job.max_retries {
383                "failed"
384            } else {
385                "active"
386            };
387            let delay = 2_i64.saturating_pow(retry_count.min(10)) * 60;
388            let next_run = (chrono::Utc::now() + chrono::Duration::seconds(delay)).to_rfc3339();
389            let _ = db
390                .query("UPDATE cron_jobs SET status = $status, retry_count = $retry_count, last_error = $error, next_run = $next_run, lease_until = NONE, run_token = NONE WHERE id = $id AND status = 'running' AND run_token = $run_token")
391                .bind(("status", status.to_string()))
392                .bind(("retry_count", retry_count))
393                .bind(("error", redacted))
394                .bind(("next_run", next_run))
395                .bind(("id", job_id.to_string()))
396                .bind(("run_token", run_token.to_string()))
397                .await?;
398        } else {
399            let mut jobs = self.noop_jobs.lock().unwrap();
400            if let Some(job) = jobs.iter_mut().find(|job| {
401                job.id.as_deref() == Some(job_id) && job.run_token.as_deref() == Some(run_token)
402            }) {
403                job.retry_count = job.retry_count.saturating_add(1);
404                job.status = if job.retry_count >= job.max_retries {
405                    "failed".to_string()
406                } else {
407                    "active".to_string()
408                };
409                job.last_error = Some(redacted);
410                job.lease_until = None;
411                job.run_token = None;
412            }
413        }
414        Ok(())
415    }
416}
417
418/// A due job ready to execute (returned by the ticker).
419#[derive(Debug, Clone)]
420pub struct DueJob {
421    pub job: CronJob,
422}
423
424/// Start the cron ticker as a background task. Returns a receiver for due jobs.
425pub fn start_cron_ticker(
426    scheduler: Arc<CronScheduler>,
427    channel: String,
428) -> (
429    tokio::sync::mpsc::Receiver<DueJob>,
430    Arc<tokio::sync::Notify>,
431) {
432    let (tx, rx) = tokio::sync::mpsc::channel(32);
433    let shutdown = Arc::new(tokio::sync::Notify::new());
434    let shutdown_clone = shutdown.clone();
435
436    tokio::spawn(async move {
437        let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(60));
438        // Skip immediate first tick
439        interval.tick().await;
440
441        loop {
442            tokio::select! {
443                _ = interval.tick() => {
444                    match scheduler.claim_due_jobs(&channel).await {
445                        Ok(jobs) => {
446                            for job in jobs {
447                                tracing::info!("Cron: job '{}' is due", job.name);
448
449                                if tx.send(DueJob { job }).await.is_err() {
450                                    return; // Receiver dropped
451                                }
452                            }
453                        }
454                        Err(e) => {
455                            tracing::error!("Cron ticker error: {}", e);
456                        }
457                    }
458                }
459                _ = shutdown_clone.notified() => {
460                    tracing::info!("Cron ticker: shutting down");
461                    break;
462                }
463            }
464        }
465    });
466
467    (rx, shutdown)
468}
469
470#[cfg(test)]
471mod tests {
472    use super::*;
473
474    fn due_job(channel: &str) -> CronJob {
475        CronJob {
476            id: Some("job-1".to_string()),
477            name: "job".to_string(),
478            schedule: "0 0 0 * * * *".to_string(),
479            one_shot: false,
480            task: "task".to_string(),
481            channel: channel.to_string(),
482            chat_id: "chat".to_string(),
483            model: "model".to_string(),
484            enabled: true,
485            last_run: None,
486            next_run: Some((chrono::Utc::now() - chrono::Duration::minutes(1)).to_rfc3339()),
487            status: "active".to_string(),
488            retry_count: 0,
489            max_retries: 3,
490            last_error: None,
491            lease_until: None,
492            run_token: None,
493        }
494    }
495
496    #[tokio::test]
497    async fn claims_only_active_channel() {
498        let scheduler = CronScheduler::new_noop();
499        scheduler
500            .noop_jobs
501            .lock()
502            .unwrap()
503            .push(due_job("telegram"));
504        assert!(scheduler.claim_due_jobs("cli").await.unwrap().is_empty());
505        let claimed = scheduler.claim_due_jobs("telegram").await.unwrap();
506        assert_eq!(claimed.len(), 1);
507        assert!(claimed[0].run_token.is_some());
508    }
509
510    #[tokio::test]
511    async fn stale_run_token_cannot_complete_job() {
512        let scheduler = CronScheduler::new_noop();
513        scheduler.noop_jobs.lock().unwrap().push(due_job("cli"));
514        let job = scheduler.claim_due_jobs("cli").await.unwrap().remove(0);
515        scheduler
516            .mark_run("job-1", "stale-token", &job.schedule)
517            .await
518            .unwrap();
519        assert_eq!(scheduler.list().await.unwrap()[0].status, "running");
520        scheduler
521            .mark_run("job-1", job.run_token.as_deref().unwrap(), &job.schedule)
522            .await
523            .unwrap();
524        assert_eq!(scheduler.list().await.unwrap()[0].status, "active");
525    }
526}
apollo/cron_scheduler.rs

apollo/
cron_scheduler.rs
