Skip to main content

beam_core/
schedule_store.rs

1use std::collections::BTreeMap;
2use std::fs;
3use std::path::PathBuf;
4
5use anyhow::Result;
6use chrono::SecondsFormat;
7use serde::{Deserialize, Serialize};
8use sha2::{Digest, Sha256};
9use thiserror::Error;
10use uuid::Uuid;
11
12use crate::BeamPaths;
13
14#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
15#[serde(rename_all = "camelCase")]
16pub enum ParsedScheduleKind {
17    Once,
18    Interval,
19    Cron,
20}
21
22#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
23#[serde(rename_all = "camelCase")]
24pub struct ParsedSchedule {
25    pub kind: ParsedScheduleKind,
26    #[serde(default, skip_serializing_if = "Option::is_none")]
27    pub run_at: Option<String>,
28    #[serde(default, skip_serializing_if = "Option::is_none")]
29    pub minutes: Option<u64>,
30    #[serde(default, skip_serializing_if = "Option::is_none")]
31    pub expr: Option<String>,
32    pub display: String,
33}
34
35#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
36#[serde(rename_all = "camelCase")]
37pub struct ScheduleRepeat {
38    pub times: Option<u64>,
39    #[serde(default)]
40    pub completed: u64,
41}
42
43#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
44#[serde(rename_all = "camelCase")]
45pub enum ScheduleChatType {
46    Group,
47    P2p,
48    TopicGroup,
49}
50
51#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
52#[serde(rename_all = "camelCase")]
53pub enum ScheduleDeliver {
54    Origin,
55    Local,
56}
57
58impl Default for ScheduleDeliver {
59    fn default() -> Self {
60        Self::Origin
61    }
62}
63
64#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
65#[serde(rename_all = "camelCase")]
66pub struct ScheduledTask {
67    pub id: String,
68    pub name: String,
69    pub schedule: String,
70    pub parsed: ParsedSchedule,
71    pub prompt: String,
72    pub working_dir: String,
73    pub chat_id: String,
74    #[serde(default, skip_serializing_if = "Option::is_none")]
75    pub root_message_id: Option<String>,
76    #[serde(default, skip_serializing_if = "Option::is_none")]
77    pub scope: Option<String>,
78    #[serde(default, skip_serializing_if = "Option::is_none")]
79    pub chat_type: Option<ScheduleChatType>,
80    #[serde(default, skip_serializing_if = "Option::is_none")]
81    pub lark_app_id: Option<String>,
82    #[serde(default, skip_serializing_if = "Option::is_none")]
83    pub creator_chat_id: Option<String>,
84    #[serde(default, skip_serializing_if = "Option::is_none")]
85    pub creator_root_message_id: Option<String>,
86    #[serde(default, skip_serializing_if = "Option::is_none")]
87    pub creator_lark_app_id: Option<String>,
88    #[serde(default = "default_true")]
89    pub enabled: bool,
90    pub created_at: String,
91    #[serde(default, skip_serializing_if = "Option::is_none")]
92    pub last_run_at: Option<String>,
93    #[serde(default, skip_serializing_if = "Option::is_none")]
94    pub next_run_at: Option<String>,
95    #[serde(default, skip_serializing_if = "Option::is_none")]
96    pub last_status: Option<String>,
97    #[serde(default, skip_serializing_if = "Option::is_none")]
98    pub last_error: Option<String>,
99    #[serde(default, skip_serializing_if = "Option::is_none")]
100    pub last_delivery_error: Option<String>,
101    #[serde(default, skip_serializing_if = "Option::is_none")]
102    pub repeat: Option<ScheduleRepeat>,
103    #[serde(default)]
104    pub deliver: ScheduleDeliver,
105}
106
107#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
108#[serde(rename_all = "camelCase")]
109pub struct CreateTaskInput {
110    #[serde(default, skip_serializing_if = "Option::is_none")]
111    pub id: Option<String>,
112    pub name: String,
113    pub schedule: String,
114    pub parsed: ParsedSchedule,
115    pub prompt: String,
116    pub working_dir: String,
117    pub chat_id: String,
118    #[serde(default, skip_serializing_if = "Option::is_none")]
119    pub root_message_id: Option<String>,
120    #[serde(default, skip_serializing_if = "Option::is_none")]
121    pub scope: Option<String>,
122    #[serde(default, skip_serializing_if = "Option::is_none")]
123    pub chat_type: Option<ScheduleChatType>,
124    #[serde(default, skip_serializing_if = "Option::is_none")]
125    pub lark_app_id: Option<String>,
126    #[serde(default, skip_serializing_if = "Option::is_none")]
127    pub creator_chat_id: Option<String>,
128    #[serde(default, skip_serializing_if = "Option::is_none")]
129    pub creator_root_message_id: Option<String>,
130    #[serde(default, skip_serializing_if = "Option::is_none")]
131    pub creator_lark_app_id: Option<String>,
132    #[serde(default, skip_serializing_if = "Option::is_none")]
133    pub next_run_at: Option<String>,
134    #[serde(default, skip_serializing_if = "Option::is_none")]
135    pub repeat: Option<ScheduleRepeat>,
136    #[serde(default, skip_serializing_if = "Option::is_none")]
137    pub deliver: Option<ScheduleDeliver>,
138}
139
140#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
141#[serde(rename_all = "camelCase")]
142pub struct ScheduleTaskUpdate {
143    #[serde(default, skip_serializing_if = "Option::is_none")]
144    pub enabled: Option<bool>,
145    #[serde(default, skip_serializing_if = "Option::is_none")]
146    pub last_run_at: Option<String>,
147    #[serde(default, skip_serializing_if = "Option::is_none")]
148    pub next_run_at: Option<Option<String>>,
149    #[serde(default, skip_serializing_if = "Option::is_none")]
150    pub last_status: Option<Option<String>>,
151    #[serde(default, skip_serializing_if = "Option::is_none")]
152    pub last_error: Option<Option<String>>,
153    #[serde(default, skip_serializing_if = "Option::is_none")]
154    pub last_delivery_error: Option<Option<String>>,
155    #[serde(default, skip_serializing_if = "Option::is_none")]
156    pub repeat: Option<Option<ScheduleRepeat>>,
157    #[serde(default, skip_serializing_if = "Option::is_none")]
158    pub root_message_id: Option<Option<String>>,
159    #[serde(default, skip_serializing_if = "Option::is_none")]
160    pub chat_type: Option<Option<ScheduleChatType>>,
161}
162
163#[derive(Debug, Error)]
164pub enum ScheduleStoreError {
165    #[error(
166        "IdempotencyConflict: schedule task {task_id} exists with different canonical input (existing={existing_input_hash}…, incoming={incoming_input_hash}…)"
167    )]
168    IdempotencyConflict {
169        task_id: String,
170        existing_input_hash: String,
171        incoming_input_hash: String,
172    },
173    #[error(transparent)]
174    Io(#[from] std::io::Error),
175    #[error(transparent)]
176    Serde(#[from] serde_json::Error),
177}
178
179pub fn create_task(
180    paths: &BeamPaths,
181    input: CreateTaskInput,
182) -> Result<ScheduledTask, ScheduleStoreError> {
183    let mut tasks = load_tasks(paths)?;
184    if let Some(id) = &input.id
185        && let Some(existing) = tasks.get(id)
186    {
187        let existing_hash = compute_input_hash(&canonical_schedule_input(existing))?;
188        let incoming_hash = compute_input_hash(&canonical_schedule_input_input(&input))?;
189        if existing_hash == incoming_hash {
190            return Ok(existing.clone());
191        }
192        return Err(ScheduleStoreError::IdempotencyConflict {
193            task_id: id.clone(),
194            existing_input_hash: existing_hash,
195            incoming_input_hash: incoming_hash,
196        });
197    }
198
199    let id = input.id.unwrap_or_else(|| {
200        Uuid::new_v4()
201            .simple()
202            .to_string()
203            .chars()
204            .take(8)
205            .collect()
206    });
207    let task = ScheduledTask {
208        id: id.clone(),
209        name: input.name,
210        schedule: input.schedule,
211        parsed: input.parsed,
212        prompt: input.prompt,
213        working_dir: input.working_dir,
214        chat_id: input.chat_id,
215        root_message_id: input.root_message_id,
216        scope: input.scope,
217        chat_type: input.chat_type,
218        lark_app_id: input.lark_app_id,
219        creator_chat_id: input.creator_chat_id,
220        creator_root_message_id: input.creator_root_message_id,
221        creator_lark_app_id: input.creator_lark_app_id,
222        enabled: true,
223        created_at: UtcNow::now(),
224        last_run_at: None,
225        next_run_at: input.next_run_at,
226        last_status: None,
227        last_error: None,
228        last_delivery_error: None,
229        repeat: input.repeat,
230        deliver: input.deliver.unwrap_or_default(),
231    };
232    tasks.insert(id, task.clone());
233    save_tasks(paths, &tasks)?;
234    Ok(task)
235}
236
237pub fn get_task(
238    paths: &BeamPaths,
239    id: &str,
240) -> Result<Option<ScheduledTask>, ScheduleStoreError> {
241    let tasks = load_tasks(paths)?;
242    Ok(tasks.get(id).cloned())
243}
244
245pub fn remove_task(paths: &BeamPaths, id: &str) -> Result<bool, ScheduleStoreError> {
246    let mut tasks = load_tasks(paths)?;
247    let existed = tasks.remove(id).is_some();
248    if existed {
249        save_tasks(paths, &tasks)?;
250    }
251    Ok(existed)
252}
253
254pub fn update_task(
255    paths: &BeamPaths,
256    id: &str,
257    updates: ScheduleTaskUpdate,
258) -> Result<(), ScheduleStoreError> {
259    let mut tasks = load_tasks(paths)?;
260    if let Some(task) = tasks.get_mut(id) {
261        if let Some(enabled) = updates.enabled {
262            task.enabled = enabled;
263        }
264        if let Some(last_run_at) = updates.last_run_at {
265            task.last_run_at = Some(last_run_at);
266        }
267        if let Some(next_run_at) = updates.next_run_at {
268            task.next_run_at = next_run_at;
269        }
270        if let Some(last_status) = updates.last_status {
271            task.last_status = last_status;
272        }
273        if let Some(last_error) = updates.last_error {
274            task.last_error = last_error;
275        }
276        if let Some(last_delivery_error) = updates.last_delivery_error {
277            task.last_delivery_error = last_delivery_error;
278        }
279        if let Some(repeat) = updates.repeat {
280            task.repeat = repeat;
281        }
282        if let Some(root_message_id) = updates.root_message_id {
283            task.root_message_id = root_message_id;
284        }
285        if let Some(chat_type) = updates.chat_type {
286            task.chat_type = chat_type;
287        }
288        save_tasks(paths, &tasks)?;
289    }
290    Ok(())
291}
292
293pub fn mark_run(
294    paths: &BeamPaths,
295    id: &str,
296    success: bool,
297    error: Option<&str>,
298    delivery_error: Option<&str>,
299) -> Result<(), ScheduleStoreError> {
300    let mut tasks = load_tasks(paths)?;
301    let Some(task) = tasks.get_mut(id) else {
302        return Ok(());
303    };
304
305    task.last_run_at = Some(UtcNow::now());
306    task.last_status = Some(if success {
307        "ok".to_string()
308    } else {
309        "error".to_string()
310    });
311    task.last_error = if success {
312        None
313    } else {
314        error.map(|s| s.to_string())
315    };
316    task.last_delivery_error = delivery_error.map(|s| s.to_string());
317
318    if let Some(repeat) = task.repeat.as_mut() {
319        repeat.completed = repeat.completed.saturating_add(1);
320        if matches!(repeat.times, Some(times) if times > 0 && repeat.completed >= times) {
321            tasks.remove(id);
322            save_tasks(paths, &tasks)?;
323            return Ok(());
324        }
325    }
326
327    if matches!(task.parsed.kind, ParsedScheduleKind::Once) {
328        task.enabled = false;
329        task.next_run_at = None;
330    }
331
332    save_tasks(paths, &tasks)?;
333    Ok(())
334}
335
336pub fn list_tasks(paths: &BeamPaths) -> Result<Vec<ScheduledTask>, ScheduleStoreError> {
337    let tasks = load_tasks(paths)?;
338    Ok(tasks.values().cloned().collect())
339}
340
341pub fn append_output_log(
342    paths: &BeamPaths,
343    task_id: &str,
344    content: &str,
345) -> Result<PathBuf, ScheduleStoreError> {
346    let dir = task_output_dir(paths, task_id);
347    fs::create_dir_all(&dir)?;
348    let fname = format!("{}.md", UtcNow::now().replace(':', "-").replace('.', "-"));
349    let path = dir.join(fname);
350    fs::write(&path, content)?;
351    Ok(path)
352}
353
354fn load_tasks(
355    paths: &BeamPaths,
356) -> Result<BTreeMap<String, ScheduledTask>, ScheduleStoreError> {
357    let path = paths.schedules_json();
358    if !path.exists() {
359        return Ok(BTreeMap::new());
360    }
361    let raw = fs::read_to_string(&path)?;
362    if raw.trim().is_empty() {
363        return Ok(BTreeMap::new());
364    }
365    let tasks = serde_json::from_str(&raw)?;
366    Ok(tasks)
367}
368
369fn save_tasks(
370    paths: &BeamPaths,
371    tasks: &BTreeMap<String, ScheduledTask>,
372) -> Result<(), ScheduleStoreError> {
373    let path = paths.schedules_json();
374    if let Some(parent) = path.parent() {
375        fs::create_dir_all(parent)?;
376    }
377    let tmp = path.with_extension("json.tmp");
378    fs::write(&tmp, serde_json::to_vec_pretty(tasks)?)?;
379    fs::rename(&tmp, &path)?;
380    Ok(())
381}
382
383fn canonical_schedule_input(task: &ScheduledTask) -> serde_json::Value {
384    serde_json::json!({
385        "name": task.name,
386        "schedule": task.schedule,
387        "parsed": {
388            "kind": task.parsed.kind,
389            "runAt": task.parsed.run_at,
390            "minutes": task.parsed.minutes,
391            "expr": task.parsed.expr,
392        },
393        "prompt": task.prompt,
394        "workingDir": task.working_dir,
395        "chatId": task.chat_id,
396        "rootMessageId": task.root_message_id,
397        "scope": task.scope,
398        "larkAppId": task.lark_app_id,
399        "repeat": task.repeat.as_ref().map(|repeat| serde_json::json!({ "times": repeat.times })),
400        "deliver": match task.deliver {
401            ScheduleDeliver::Origin => "origin",
402            ScheduleDeliver::Local => "local",
403        }
404    })
405}
406
407fn canonical_schedule_input_input(input: &CreateTaskInput) -> serde_json::Value {
408    serde_json::json!({
409        "name": input.name,
410        "schedule": input.schedule,
411        "parsed": {
412            "kind": input.parsed.kind,
413            "runAt": input.parsed.run_at,
414            "minutes": input.parsed.minutes,
415            "expr": input.parsed.expr,
416        },
417        "prompt": input.prompt,
418        "workingDir": input.working_dir,
419        "chatId": input.chat_id,
420        "rootMessageId": input.root_message_id,
421        "scope": input.scope,
422        "larkAppId": input.lark_app_id,
423        "repeat": input.repeat.as_ref().map(|repeat| serde_json::json!({ "times": repeat.times })),
424        "deliver": match input.deliver.clone().unwrap_or_default() {
425            ScheduleDeliver::Origin => "origin",
426            ScheduleDeliver::Local => "local",
427        }
428    })
429}
430
431fn compute_input_hash(value: &serde_json::Value) -> Result<String, ScheduleStoreError> {
432    let canonical = canonical_json(value);
433    let mut hasher = Sha256::new();
434    hasher.update(canonical.as_bytes());
435    Ok(format!("sha256:{:x}", hasher.finalize()))
436}
437
438fn canonical_json(value: &serde_json::Value) -> String {
439    match value {
440        serde_json::Value::Null => "null".to_string(),
441        serde_json::Value::Bool(v) => if *v { "true" } else { "false" }.to_string(),
442        serde_json::Value::Number(v) => v.to_string(),
443        serde_json::Value::String(v) => serde_json::to_string(v).expect("string serializable"),
444        serde_json::Value::Array(items) => {
445            let mut out = String::from("[");
446            let mut first = true;
447            for item in items {
448                if !first {
449                    out.push(',');
450                }
451                first = false;
452                out.push_str(&canonical_json(item));
453            }
454            out.push(']');
455            out
456        }
457        serde_json::Value::Object(map) => {
458            let mut keys: Vec<_> = map.keys().collect();
459            keys.sort();
460            let mut out = String::from("{");
461            let mut first = true;
462            for key in keys {
463                if !first {
464                    out.push(',');
465                }
466                first = false;
467                out.push_str(&serde_json::to_string(key).expect("key serializable"));
468                out.push(':');
469                out.push_str(&canonical_json(&map[key]));
470            }
471            out.push('}');
472            out
473        }
474    }
475}
476
477fn task_output_dir(paths: &BeamPaths, task_id: &str) -> PathBuf {
478    paths.schedules_output_dir().join(task_id)
479}
480
481struct UtcNow;
482
483impl UtcNow {
484    fn now() -> String {
485        chrono::Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true)
486    }
487}
488
489fn default_true() -> bool {
490    true
491}
492
493#[cfg(test)]
494mod tests {
495    use super::*;
496    use std::time::{SystemTime, UNIX_EPOCH};
497
498    fn temp_paths(label: &str) -> BeamPaths {
499        let nanos = SystemTime::now()
500            .duration_since(UNIX_EPOCH)
501            .unwrap_or_default()
502            .as_nanos();
503        BeamPaths::from_root(std::env::temp_dir().join(format!(
504            "beam-schedule-store-{label}-{nanos}-{}",
505            std::process::id()
506        )))
507    }
508
509    fn parsed_cron() -> ParsedSchedule {
510        ParsedSchedule {
511            kind: ParsedScheduleKind::Cron,
512            run_at: None,
513            minutes: None,
514            expr: Some("0 9 * * *".to_string()),
515            display: "0 9 * * *".to_string(),
516        }
517    }
518
519    #[test]
520    fn create_task_returns_existing_when_canonical_input_matches() {
521        let paths = temp_paths("identical");
522        let input = CreateTaskInput {
523            id: Some("wf_task".to_string()),
524            name: "schedule-demo daily 9am".to_string(),
525            schedule: "0 9 * * *".to_string(),
526            parsed: parsed_cron(),
527            prompt: "Schedule demo: run workflow self-check.".to_string(),
528            working_dir: "/tmp/beam-schedule-demo".to_string(),
529            chat_id: "oc_workflow_demo".to_string(),
530            root_message_id: None,
531            scope: Some("thread".to_string()),
532            chat_type: None,
533            lark_app_id: None,
534            creator_chat_id: None,
535            creator_root_message_id: None,
536            creator_lark_app_id: None,
537            next_run_at: None,
538            repeat: None,
539            deliver: None,
540        };
541        let created = create_task(&paths, input.clone()).expect("create");
542        let returned = create_task(&paths, input).expect("create existing");
543        assert_eq!(created.id, "wf_task");
544        assert_eq!(returned.id, "wf_task");
545        assert_eq!(list_tasks(&paths).expect("list").len(), 1);
546        let _ = std::fs::remove_dir_all(paths.root());
547    }
548
549    #[test]
550    fn create_task_conflicts_when_canonical_input_differs() {
551        let paths = temp_paths("conflict");
552        let input = CreateTaskInput {
553            id: Some("wf_task".to_string()),
554            name: "schedule-demo daily 9am".to_string(),
555            schedule: "0 9 * * *".to_string(),
556            parsed: parsed_cron(),
557            prompt: "Schedule demo: run workflow self-check.".to_string(),
558            working_dir: "/tmp/beam-schedule-demo".to_string(),
559            chat_id: "oc_workflow_demo".to_string(),
560            root_message_id: None,
561            scope: Some("thread".to_string()),
562            chat_type: None,
563            lark_app_id: None,
564            creator_chat_id: None,
565            creator_root_message_id: None,
566            creator_lark_app_id: None,
567            next_run_at: None,
568            repeat: None,
569            deliver: None,
570        };
571        let _ = create_task(&paths, input).expect("create");
572        let changed = CreateTaskInput {
573            id: Some("wf_task".to_string()),
574            name: "schedule-demo daily 9am".to_string(),
575            schedule: "0 9 * * *".to_string(),
576            parsed: parsed_cron(),
577            prompt: "changed prompt".to_string(),
578            working_dir: "/tmp/beam-schedule-demo".to_string(),
579            chat_id: "oc_workflow_demo".to_string(),
580            root_message_id: None,
581            scope: Some("thread".to_string()),
582            chat_type: None,
583            lark_app_id: None,
584            creator_chat_id: None,
585            creator_root_message_id: None,
586            creator_lark_app_id: None,
587            next_run_at: None,
588            repeat: None,
589            deliver: None,
590        };
591        let err = create_task(&paths, changed).expect_err("conflict");
592        assert!(matches!(
593            err,
594            ScheduleStoreError::IdempotencyConflict { .. }
595        ));
596        let _ = std::fs::remove_dir_all(paths.root());
597    }
598
599    #[test]
600    fn mark_run_removes_finite_repeat_after_completion() {
601        let paths = temp_paths("mark");
602        let input = CreateTaskInput {
603            id: Some("wf_task".to_string()),
604            name: "schedule-demo daily 9am".to_string(),
605            schedule: "0 9 * * *".to_string(),
606            parsed: parsed_cron(),
607            prompt: "Schedule demo: run workflow self-check.".to_string(),
608            working_dir: "/tmp/beam-schedule-demo".to_string(),
609            chat_id: "oc_workflow_demo".to_string(),
610            root_message_id: None,
611            scope: Some("thread".to_string()),
612            chat_type: None,
613            lark_app_id: None,
614            creator_chat_id: None,
615            creator_root_message_id: None,
616            creator_lark_app_id: None,
617            next_run_at: None,
618            repeat: Some(ScheduleRepeat {
619                times: Some(1),
620                completed: 0,
621            }),
622            deliver: None,
623        };
624        let _ = create_task(&paths, input).expect("create");
625        mark_run(&paths, "wf_task", true, None, None).expect("mark run");
626        assert!(get_task(&paths, "wf_task").expect("get").is_none());
627        let _ = std::fs::remove_dir_all(paths.root());
628    }
629}