use crate::memory::surreal::SurrealMemory;
use serde::{Deserialize, Serialize};
use std::str::FromStr;
use std::sync::Arc;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CronJob {
pub id: Option<String>,
pub name: String,
pub schedule: String,
#[serde(default)]
pub one_shot: bool,
pub task: String,
pub channel: String,
#[serde(default)]
pub chat_id: String,
pub model: String,
pub enabled: bool,
pub last_run: Option<String>,
pub next_run: Option<String>,
#[serde(default = "default_job_status")]
pub status: String,
#[serde(default)]
pub retry_count: u32,
#[serde(default = "default_max_retries")]
pub max_retries: u32,
#[serde(default)]
pub last_error: Option<String>,
#[serde(default)]
pub lease_until: Option<String>,
#[serde(default)]
pub run_token: Option<String>,
}
fn default_job_status() -> String {
"active".to_string()
}
fn default_max_retries() -> u32 {
3
}
pub struct CronScheduler {
memory: Option<Arc<SurrealMemory>>,
noop_jobs: std::sync::Mutex<Vec<CronJob>>,
}
impl CronScheduler {
pub fn new(memory: Arc<SurrealMemory>) -> Self {
Self::new_maybe(Some(memory))
}
pub fn new_noop() -> Self {
Self::new_maybe(None)
}
fn new_maybe(memory: Option<Arc<SurrealMemory>>) -> Self {
Self {
memory,
noop_jobs: std::sync::Mutex::new(Vec::new()),
}
}
pub async fn add(
&self,
name: &str,
schedule: &str,
task: &str,
channel: &str,
chat_id: &str,
model: &str,
) -> anyhow::Result<String> {
let parsed = cron::Schedule::from_str(schedule)
.map_err(|e| anyhow::anyhow!("Invalid cron expression: {}", e))?;
let next_run = parsed
.upcoming(chrono::Utc)
.next()
.map(|t| t.to_rfc3339())
.ok_or_else(|| {
anyhow::anyhow!("cron expression {schedule:?} has no future occurrences")
})?;
let next_run = Some(next_run);
let job = CronJob {
id: Some(name.to_string()),
name: name.to_string(),
schedule: schedule.to_string(),
one_shot: false,
task: task.to_string(),
channel: channel.to_string(),
chat_id: chat_id.to_string(),
model: model.to_string(),
enabled: true,
last_run: None,
next_run,
status: default_job_status(),
retry_count: 0,
max_retries: default_max_retries(),
last_error: None,
lease_until: None,
run_token: None,
};
if let Some(ref memory) = self.memory {
let db = memory.db().await?;
let created: Option<CronJob> = db.create("cron_jobs").content(job).await?;
let created = created.ok_or_else(|| anyhow::anyhow!("Failed to create cron job"))?;
Ok(created
.id
.ok_or_else(|| anyhow::anyhow!("Created cron job is missing an id"))?
.to_string())
} else {
let mut jobs = self.noop_jobs.lock().unwrap();
jobs.push(job);
Ok(name.to_string())
}
}
pub async fn add_once(
&self,
name: &str,
run_at: chrono::DateTime<chrono::Utc>,
task: &str,
channel: &str,
chat_id: &str,
model: &str,
) -> anyhow::Result<String> {
if run_at <= chrono::Utc::now() {
anyhow::bail!("run_at must be in the future");
}
let job = CronJob {
id: Some(name.to_string()),
name: name.to_string(),
schedule: String::new(),
one_shot: true,
task: task.to_string(),
channel: channel.to_string(),
chat_id: chat_id.to_string(),
model: model.to_string(),
enabled: true,
last_run: None,
next_run: Some(run_at.to_rfc3339()),
status: default_job_status(),
retry_count: 0,
max_retries: default_max_retries(),
last_error: None,
lease_until: None,
run_token: None,
};
if let Some(ref memory) = self.memory {
let created: Option<CronJob> =
memory.db().await?.create("cron_jobs").content(job).await?;
Ok(created
.and_then(|job| job.id)
.ok_or_else(|| anyhow::anyhow!("Created one-shot job is missing an id"))?)
} else {
self.noop_jobs.lock().unwrap().push(job);
Ok(name.to_string())
}
}
pub async fn list(&self) -> anyhow::Result<Vec<CronJob>> {
if let Some(ref memory) = self.memory {
let db = memory.db().await?;
let mut result: surrealdb::Response =
db.query("SELECT * FROM cron_jobs ORDER BY name").await?;
let jobs: Vec<CronJob> = result.take(0)?;
Ok(jobs)
} else {
let jobs = self.noop_jobs.lock().unwrap();
Ok(jobs.clone())
}
}
pub async fn remove(&self, id_or_name: &str) -> anyhow::Result<bool> {
if let Some(ref memory) = self.memory {
let db = memory.db().await?;
let mut result: surrealdb::Response = db
.query("DELETE FROM cron_jobs WHERE id = $target OR name = $target")
.bind(("target", id_or_name.to_string()))
.await?;
let deleted: Vec<CronJob> = result.take(0)?;
Ok(!deleted.is_empty())
} else {
let mut jobs = self.noop_jobs.lock().unwrap();
let before = jobs.len();
jobs.retain(|j| j.name != id_or_name);
Ok(jobs.len() < before)
}
}
async fn find(&self, id_or_name: &str) -> anyhow::Result<Option<CronJob>> {
Ok(self
.list()
.await?
.into_iter()
.find(|job| job.id.as_deref() == Some(id_or_name) || job.name == id_or_name))
}
pub async fn enable(&self, id_or_name: &str) -> anyhow::Result<bool> {
let Some(job) = self.find(id_or_name).await? else {
return Ok(false);
};
let revived = if job.status == "running" {
None
} else if job.one_shot || job.schedule.is_empty() {
let pending = job
.next_run
.as_deref()
.and_then(|t| chrono::DateTime::parse_from_rfc3339(t).ok())
.is_some_and(|t| t > chrono::Utc::now());
if !pending {
anyhow::bail!(
"cannot enable one-shot job {id_or_name:?}: its run time has already passed"
);
}
Some(job.next_run.clone())
} else {
let parsed = cron::Schedule::from_str(&job.schedule).map_err(|e| {
anyhow::anyhow!(
"cannot enable {id_or_name:?}: invalid cron expression {:?}: {e}",
job.schedule
)
})?;
let next = parsed.upcoming(chrono::Utc).next().ok_or_else(|| {
anyhow::anyhow!(
"cannot enable {id_or_name:?}: cron expression {:?} has no future occurrences",
job.schedule
)
})?;
Some(Some(next.to_rfc3339()))
};
if let Some(ref memory) = self.memory {
let db = memory.db().await?;
let query = if revived.is_some() {
"UPDATE cron_jobs SET enabled = true, status = 'active', next_run = $next, retry_count = 0, last_error = NONE WHERE id = $target OR name = $target"
} else {
"UPDATE cron_jobs SET enabled = true WHERE id = $target OR name = $target"
};
let mut result: surrealdb::Response = db
.query(query)
.bind(("target", id_or_name.to_string()))
.bind(("next", revived.clone().unwrap_or_default()))
.await?;
let updated: Vec<CronJob> = result.take(0)?;
Ok(!updated.is_empty())
} else {
let mut jobs = self.noop_jobs.lock().unwrap();
let Some(job) = jobs
.iter_mut()
.find(|j| j.id.as_deref() == Some(id_or_name) || j.name == id_or_name)
else {
return Ok(false);
};
job.enabled = true;
if let Some(next_run) = revived {
job.next_run = next_run;
job.status = "active".to_string();
job.retry_count = 0;
job.last_error = None;
}
Ok(true)
}
}
pub async fn disable(&self, id_or_name: &str) -> anyhow::Result<bool> {
if let Some(ref memory) = self.memory {
let db = memory.db().await?;
let mut result: surrealdb::Response = db
.query("UPDATE cron_jobs SET enabled = false WHERE id = $target OR name = $target")
.bind(("target", id_or_name.to_string()))
.await?;
let updated: Vec<CronJob> = result.take(0)?;
Ok(!updated.is_empty())
} else {
let mut jobs = self.noop_jobs.lock().unwrap();
if let Some(job) = jobs
.iter_mut()
.find(|j| j.id.as_deref() == Some(id_or_name) || j.name == id_or_name)
{
job.enabled = false;
Ok(true)
} else {
Ok(false)
}
}
}
pub async fn due_jobs(&self) -> anyhow::Result<Vec<CronJob>> {
let now = chrono::Utc::now().to_rfc3339();
if let Some(ref memory) = self.memory {
let db = memory.db().await?;
let mut result: surrealdb::Response = db
.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)")
.bind(("now", now))
.await?;
let jobs: Vec<CronJob> = result.take(0)?;
Ok(jobs)
} else {
let jobs = self.noop_jobs.lock().unwrap();
Ok(jobs
.iter()
.filter(|j| {
j.enabled
&& ((j.status == "active" && j.next_run.as_deref() <= Some(&now))
|| (j.status == "running" && j.lease_until.as_deref() <= Some(&now)))
})
.cloned()
.collect())
}
}
pub async fn claim_due_jobs(&self, channel: &str) -> anyhow::Result<Vec<CronJob>> {
let due = self.due_jobs().await?;
let claim_now = chrono::Utc::now().to_rfc3339();
let lease_until = (chrono::Utc::now() + chrono::Duration::minutes(10)).to_rfc3339();
let mut claimed = Vec::new();
for job in due {
if job.channel != channel {
continue;
}
let Some(job_id) = job.id.as_deref() else {
continue;
};
let run_token = uuid::Uuid::new_v4().to_string();
if let Some(ref memory) = self.memory {
let mut result = memory
.db().await?
.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")
.bind(("id", job_id.to_string()))
.bind(("lease", lease_until.clone()))
.bind(("run_token", run_token.clone()))
.bind(("now", claim_now.clone()))
.await?;
let updated: Vec<CronJob> = result.take(0)?;
claimed.extend(updated);
} else {
let mut jobs = self.noop_jobs.lock().unwrap();
if let Some(stored) = jobs.iter_mut().find(|stored| {
stored.id.as_deref() == Some(job_id)
&& (stored.status == "active"
|| stored.lease_until.as_deref() <= Some(claim_now.as_str()))
}) {
stored.status = "running".to_string();
stored.lease_until = Some(lease_until.clone());
stored.run_token = Some(run_token);
claimed.push(stored.clone());
}
}
}
Ok(claimed)
}
pub async fn mark_run(
&self,
job_id: &str,
run_token: &str,
schedule: &str,
) -> anyhow::Result<()> {
let now = chrono::Utc::now();
let (next_run, status, last_error) = if schedule.is_empty() {
(None, "completed", None)
} else {
match cron::Schedule::from_str(schedule) {
Err(e) => {
let reason = format!("invalid cron expression {schedule:?}: {e}");
tracing::error!("cron job {}: {}", job_id, reason);
(None, "invalid_schedule", Some(reason))
}
Ok(parsed) => match parsed.upcoming(chrono::Utc).next() {
Some(next) => (Some(next.to_rfc3339()), "active", None),
None => {
let reason =
format!("cron expression {schedule:?} has no future occurrences");
tracing::info!("cron job {}: {}", job_id, reason);
(None, "exhausted", Some(reason))
}
},
}
};
if let Some(ref memory) = self.memory {
let db = memory.db().await?;
let mut result = db
.query("UPDATE cron_jobs SET last_run = $last, next_run = $next, status = $status, retry_count = 0, last_error = $last_error, lease_until = NONE, run_token = NONE WHERE id = $id AND status = 'running' AND run_token = $run_token")
.bind(("last", now.to_rfc3339()))
.bind(("next", next_run))
.bind(("status", status.to_string()))
.bind(("last_error", last_error.clone()))
.bind(("id", job_id.to_string()))
.bind(("run_token", run_token.to_string()))
.await?;
let _: Vec<CronJob> = result.take(0)?;
} else {
let mut jobs = self.noop_jobs.lock().unwrap();
if let Some(job) = jobs.iter_mut().find(|job| {
job.id.as_deref() == Some(job_id) && job.run_token.as_deref() == Some(run_token)
}) {
job.last_run = Some(now.to_rfc3339());
job.next_run = next_run;
job.status = status.to_string();
job.retry_count = 0;
job.last_error = last_error.clone();
job.lease_until = None;
job.run_token = None;
}
}
Ok(())
}
pub async fn release_run(&self, job_id: &str, run_token: &str) -> anyhow::Result<()> {
if let Some(ref memory) = self.memory {
let mut result = memory
.db().await?
.query("UPDATE cron_jobs SET status = 'active', lease_until = NONE, run_token = NONE WHERE id = $id AND status = 'running' AND run_token = $run_token")
.bind(("id", job_id.to_string()))
.bind(("run_token", run_token.to_string()))
.await?;
let _: Vec<CronJob> = result.take(0)?;
} else {
let mut jobs = self.noop_jobs.lock().unwrap();
if let Some(job) = jobs.iter_mut().find(|job| {
job.id.as_deref() == Some(job_id) && job.run_token.as_deref() == Some(run_token)
}) {
job.status = "active".to_string();
job.lease_until = None;
job.run_token = None;
}
}
Ok(())
}
pub async fn fail_run(&self, job_id: &str, run_token: &str, error: &str) -> anyhow::Result<()> {
let redacted = crate::redaction::redact_text(error);
if let Some(ref memory) = self.memory {
let db = memory.db().await?;
let Some(job) = self.list().await?.into_iter().find(|job| {
job.id.as_deref() == Some(job_id) && job.run_token.as_deref() == Some(run_token)
}) else {
return Ok(());
};
let retry_count = job.retry_count.saturating_add(1);
let status = if retry_count >= job.max_retries {
"failed"
} else {
"active"
};
let delay = 2_i64.saturating_pow(retry_count.min(10)) * 60;
let next_run = (chrono::Utc::now() + chrono::Duration::seconds(delay)).to_rfc3339();
let mut result = db
.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")
.bind(("status", status.to_string()))
.bind(("retry_count", retry_count))
.bind(("error", redacted))
.bind(("next_run", next_run))
.bind(("id", job_id.to_string()))
.bind(("run_token", run_token.to_string()))
.await?;
let _: Vec<CronJob> = result.take(0)?;
} else {
let mut jobs = self.noop_jobs.lock().unwrap();
if let Some(job) = jobs.iter_mut().find(|job| {
job.id.as_deref() == Some(job_id) && job.run_token.as_deref() == Some(run_token)
}) {
job.retry_count = job.retry_count.saturating_add(1);
job.status = if job.retry_count >= job.max_retries {
"failed".to_string()
} else {
"active".to_string()
};
job.last_error = Some(redacted);
job.lease_until = None;
job.run_token = None;
}
}
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct DueJob {
pub job: CronJob,
}
pub fn start_cron_ticker(
scheduler: Arc<CronScheduler>,
channel: String,
) -> (
tokio::sync::mpsc::Receiver<DueJob>,
Arc<tokio::sync::Notify>,
) {
let (tx, rx) = tokio::sync::mpsc::channel(32);
let shutdown = Arc::new(tokio::sync::Notify::new());
let shutdown_clone = shutdown.clone();
tokio::spawn(async move {
let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(60));
interval.tick().await;
loop {
tokio::select! {
_ = interval.tick() => {
match scheduler.claim_due_jobs(&channel).await {
Ok(jobs) => {
for job in jobs {
tracing::info!("Cron: job '{}' is due", job.name);
if tx.send(DueJob { job }).await.is_err() {
return; }
}
}
Err(e) => {
tracing::error!("Cron ticker error: {}", e);
}
}
}
_ = shutdown_clone.notified() => {
tracing::info!("Cron ticker: shutting down");
break;
}
}
}
});
(rx, shutdown)
}
#[cfg(test)]
mod tests {
use super::*;
fn due_job(channel: &str) -> CronJob {
CronJob {
id: Some("job-1".to_string()),
name: "job".to_string(),
schedule: "0 0 0 * * * *".to_string(),
one_shot: false,
task: "task".to_string(),
channel: channel.to_string(),
chat_id: "chat".to_string(),
model: "model".to_string(),
enabled: true,
last_run: None,
next_run: Some((chrono::Utc::now() - chrono::Duration::minutes(1)).to_rfc3339()),
status: "active".to_string(),
retry_count: 0,
max_retries: 3,
last_error: None,
lease_until: None,
run_token: None,
}
}
#[tokio::test]
async fn claims_only_active_channel() {
let scheduler = CronScheduler::new_noop();
scheduler
.noop_jobs
.lock()
.unwrap()
.push(due_job("telegram"));
assert!(scheduler.claim_due_jobs("cli").await.unwrap().is_empty());
let claimed = scheduler.claim_due_jobs("telegram").await.unwrap();
assert_eq!(claimed.len(), 1);
assert!(claimed[0].run_token.is_some());
}
#[tokio::test]
async fn stale_run_token_cannot_complete_job() {
let scheduler = CronScheduler::new_noop();
scheduler.noop_jobs.lock().unwrap().push(due_job("cli"));
let job = scheduler.claim_due_jobs("cli").await.unwrap().remove(0);
scheduler
.mark_run("job-1", "stale-token", &job.schedule)
.await
.unwrap();
assert_eq!(scheduler.list().await.unwrap()[0].status, "running");
scheduler
.mark_run("job-1", job.run_token.as_deref().unwrap(), &job.schedule)
.await
.unwrap();
assert_eq!(scheduler.list().await.unwrap()[0].status, "active");
}
async fn mark_run_with_schedule(schedule: &str) -> CronJob {
let scheduler = CronScheduler::new_noop();
let mut job = due_job("cli");
job.schedule = schedule.to_string();
scheduler.noop_jobs.lock().unwrap().push(job);
let claimed = scheduler.claim_due_jobs("cli").await.unwrap().remove(0);
scheduler
.mark_run("job-1", claimed.run_token.as_deref().unwrap(), schedule)
.await
.unwrap();
scheduler.list().await.unwrap().remove(0)
}
#[tokio::test]
async fn exhausted_schedule_is_not_left_active() {
let job = mark_run_with_schedule("0 0 12 1 1 * 2020").await;
assert_eq!(job.status, "exhausted");
assert!(job.next_run.is_none());
assert!(job.last_error.is_some(), "the reason must be recorded");
}
#[tokio::test]
async fn unparseable_schedule_is_not_left_active() {
let job = mark_run_with_schedule("not a cron expression").await;
assert_eq!(job.status, "invalid_schedule");
assert!(job.next_run.is_none());
assert!(job.last_error.is_some());
}
#[tokio::test]
async fn add_refuses_a_schedule_with_no_future_occurrence() {
let scheduler = CronScheduler::new_noop();
let error = scheduler
.add(
"yearly",
"0 0 12 1 1 * 2020",
"task",
"cli",
"chat",
"model",
)
.await
.expect_err("a job that can never run must not be reported as scheduled");
assert!(
error.to_string().contains("no future occurrences"),
"unexpected error: {error}"
);
assert!(scheduler.list().await.unwrap().is_empty());
}
#[tokio::test]
async fn enable_revives_a_job_left_with_a_dead_status() {
let scheduler = CronScheduler::new_noop();
let mut job = due_job("cli");
job.schedule = "0 0 * * * * *".to_string();
job.status = "exhausted".to_string();
job.next_run = None;
job.last_error = Some("dead".to_string());
scheduler.noop_jobs.lock().unwrap().push(job);
assert!(scheduler.disable("job-1").await.unwrap());
assert!(scheduler.enable("job-1").await.unwrap());
let job = scheduler.list().await.unwrap().remove(0);
assert!(job.enabled);
assert_eq!(
job.status, "active",
"enable must clear a recoverable status"
);
assert!(job.next_run.is_some(), "enable must recompute next_run");
assert!(job.last_error.is_none());
}
#[tokio::test]
async fn enable_fails_loudly_when_the_schedule_cannot_produce_a_next_run() {
let scheduler = CronScheduler::new_noop();
let mut job = due_job("cli");
job.schedule = "0 0 12 1 1 * 2020".to_string();
job.status = "exhausted".to_string();
job.next_run = None;
scheduler.noop_jobs.lock().unwrap().push(job);
let error = scheduler
.enable("job-1")
.await
.expect_err("a job that cannot be revived must not report success");
assert!(
error.to_string().contains("no future occurrences"),
"unexpected error: {error}"
);
}
#[tokio::test]
async fn enable_leaves_a_running_job_alone() {
let scheduler = CronScheduler::new_noop();
scheduler.noop_jobs.lock().unwrap().push(due_job("cli"));
let claimed = scheduler.claim_due_jobs("cli").await.unwrap().remove(0);
assert!(scheduler.enable("job-1").await.unwrap());
let job = scheduler.list().await.unwrap().remove(0);
assert_eq!(job.status, "running");
assert_eq!(job.run_token, claimed.run_token);
}
#[tokio::test]
async fn valid_schedule_stays_active_with_a_next_run() {
let job = mark_run_with_schedule("0 0 * * * * *").await;
assert_eq!(job.status, "active");
assert!(job.next_run.is_some());
assert!(job.last_error.is_none());
}
}