Skip to main content

microclaw_engine/internal/storage/db/
tasks.rs

1use super::*;
2
3#[derive(Debug, Clone)]
4#[allow(dead_code)]
5pub struct TaskRunLog {
6    pub id: i64,
7    pub task_id: i64,
8    pub chat_id: i64,
9    pub started_at: String,
10    pub finished_at: String,
11    pub duration_ms: i64,
12    pub success: bool,
13    pub result_summary: Option<String>,
14}
15
16#[derive(Debug, Clone)]
17#[allow(dead_code)]
18pub struct ScheduledTask {
19    pub id: i64,
20    pub chat_id: i64,
21    pub prompt: String,
22    pub schedule_type: String,  // "cron" or "once"
23    pub schedule_value: String, // cron expression or ISO timestamp
24    pub timezone: String,       // IANA timezone; empty means "use app default"
25    pub next_run: String,       // ISO timestamp
26    pub last_run: Option<String>,
27    pub status: String, // "active", "paused", "completed", "cancelled"
28    pub created_at: String,
29    /// Optional completion contract: JSON array of exit criteria, verified
30    /// after each run (see src/completion_contract.rs).
31    pub exit_criteria: Option<String>,
32    /// Number of times this task has run (success or failure).
33    pub run_count: i64,
34    /// Retire as `completed` after this many runs. NULL = unlimited.
35    pub max_runs: Option<i64>,
36    /// Retire as `completed` once the next firing would pass this instant.
37    pub not_after: Option<String>,
38}
39
40#[derive(Debug, Clone)]
41pub struct ScheduledTaskDlqEntry {
42    pub id: i64,
43    pub task_id: i64,
44    pub chat_id: i64,
45    pub failed_at: String,
46    pub started_at: String,
47    pub finished_at: String,
48    pub duration_ms: i64,
49    pub error_summary: Option<String>,
50    pub replayed_at: Option<String>,
51    pub replay_note: Option<String>,
52}
53
54impl Database {
55    pub fn create_scheduled_task(
56        &self,
57        chat_id: i64,
58        prompt: &str,
59        schedule_type: &str,
60        schedule_value: &str,
61        next_run: &str,
62    ) -> Result<i64, MicroClawError> {
63        self.create_scheduled_task_with_timezone(
64            chat_id,
65            prompt,
66            schedule_type,
67            schedule_value,
68            "",
69            next_run,
70        )
71    }
72
73    pub fn create_scheduled_task_with_timezone(
74        &self,
75        chat_id: i64,
76        prompt: &str,
77        schedule_type: &str,
78        schedule_value: &str,
79        timezone: &str,
80        next_run: &str,
81    ) -> Result<i64, MicroClawError> {
82        self.create_scheduled_task_full(
83            chat_id,
84            prompt,
85            schedule_type,
86            schedule_value,
87            timezone,
88            next_run,
89            None,
90        )
91    }
92
93    #[allow(clippy::too_many_arguments)]
94    pub fn create_scheduled_task_full(
95        &self,
96        chat_id: i64,
97        prompt: &str,
98        schedule_type: &str,
99        schedule_value: &str,
100        timezone: &str,
101        next_run: &str,
102        exit_criteria: Option<&str>,
103    ) -> Result<i64, MicroClawError> {
104        self.create_scheduled_task_lifecycle(
105            chat_id,
106            prompt,
107            schedule_type,
108            schedule_value,
109            timezone,
110            next_run,
111            exit_criteria,
112            None,
113            None,
114        )
115    }
116
117    #[allow(clippy::too_many_arguments)]
118    pub fn create_scheduled_task_lifecycle(
119        &self,
120        chat_id: i64,
121        prompt: &str,
122        schedule_type: &str,
123        schedule_value: &str,
124        timezone: &str,
125        next_run: &str,
126        exit_criteria: Option<&str>,
127        max_runs: Option<i64>,
128        not_after: Option<&str>,
129    ) -> Result<i64, MicroClawError> {
130        let conn = self.lock_conn();
131        let now = chrono::Utc::now().to_rfc3339();
132        conn.execute(
133            "INSERT INTO scheduled_tasks (chat_id, prompt, schedule_type, schedule_value, timezone, next_run, status, created_at, exit_criteria, max_runs, not_after)
134             VALUES (?1, ?2, ?3, ?4, ?5, ?6, 'active', ?7, ?8, ?9, ?10)",
135            params![
136                chat_id,
137                prompt,
138                schedule_type,
139                schedule_value,
140                timezone,
141                next_run,
142                now,
143                exit_criteria,
144                max_runs,
145                not_after
146            ],
147        )?;
148        Ok(conn.last_insert_rowid())
149    }
150
151    pub fn get_due_tasks(&self, now: &str) -> Result<Vec<ScheduledTask>, MicroClawError> {
152        let conn = self.lock_conn();
153        let mut stmt = conn.prepare(
154            "SELECT id, chat_id, prompt, schedule_type, schedule_value, timezone, next_run, last_run, status, created_at, exit_criteria, run_count, max_runs, not_after
155             FROM scheduled_tasks
156             WHERE status = 'active' AND next_run <= ?1",
157        )?;
158        let tasks = stmt
159            .query_map(params![now], |row| {
160                Ok(ScheduledTask {
161                    id: row.get(0)?,
162                    chat_id: row.get(1)?,
163                    prompt: row.get(2)?,
164                    schedule_type: row.get(3)?,
165                    schedule_value: row.get(4)?,
166                    timezone: row.get(5)?,
167                    next_run: row.get(6)?,
168                    last_run: row.get(7)?,
169                    status: row.get(8)?,
170                    created_at: row.get(9)?,
171                    exit_criteria: row.get(10)?,
172                    run_count: row.get(11)?,
173                    max_runs: row.get(12)?,
174                    not_after: row.get(13)?,
175                })
176            })?
177            .collect::<Result<Vec<_>, _>>()?;
178        Ok(tasks)
179    }
180
181    pub fn get_tasks_for_chat(&self, chat_id: i64) -> Result<Vec<ScheduledTask>, MicroClawError> {
182        let conn = self.lock_conn();
183        let mut stmt = conn.prepare(
184            "SELECT id, chat_id, prompt, schedule_type, schedule_value, timezone, next_run, last_run, status, created_at, exit_criteria, run_count, max_runs, not_after
185             FROM scheduled_tasks
186             WHERE chat_id = ?1 AND status IN ('active', 'paused')
187             ORDER BY id",
188        )?;
189        let tasks = stmt
190            .query_map(params![chat_id], |row| {
191                Ok(ScheduledTask {
192                    id: row.get(0)?,
193                    chat_id: row.get(1)?,
194                    prompt: row.get(2)?,
195                    schedule_type: row.get(3)?,
196                    schedule_value: row.get(4)?,
197                    timezone: row.get(5)?,
198                    next_run: row.get(6)?,
199                    last_run: row.get(7)?,
200                    status: row.get(8)?,
201                    created_at: row.get(9)?,
202                    exit_criteria: row.get(10)?,
203                    run_count: row.get(11)?,
204                    max_runs: row.get(12)?,
205                    not_after: row.get(13)?,
206                })
207            })?
208            .collect::<Result<Vec<_>, _>>()?;
209        Ok(tasks)
210    }
211
212    /// List scheduled tasks across all chats for management views.
213    /// `status` filters to one status; `None` returns every task. Newest first.
214    pub fn list_scheduled_tasks(
215        &self,
216        status: Option<&str>,
217        limit: usize,
218    ) -> Result<Vec<ScheduledTask>, MicroClawError> {
219        let conn = self.lock_conn();
220        let mut stmt = conn.prepare(
221            "SELECT id, chat_id, prompt, schedule_type, schedule_value, timezone, next_run, last_run, status, created_at, exit_criteria, run_count, max_runs, not_after
222             FROM scheduled_tasks
223             WHERE (?1 IS NULL OR status = ?1)
224             ORDER BY id DESC
225             LIMIT ?2",
226        )?;
227        let tasks = stmt
228            .query_map(params![status, limit as i64], |row| {
229                Ok(ScheduledTask {
230                    id: row.get(0)?,
231                    chat_id: row.get(1)?,
232                    prompt: row.get(2)?,
233                    schedule_type: row.get(3)?,
234                    schedule_value: row.get(4)?,
235                    timezone: row.get(5)?,
236                    next_run: row.get(6)?,
237                    last_run: row.get(7)?,
238                    status: row.get(8)?,
239                    created_at: row.get(9)?,
240                    exit_criteria: row.get(10)?,
241                    run_count: row.get(11)?,
242                    max_runs: row.get(12)?,
243                    not_after: row.get(13)?,
244                })
245            })?
246            .collect::<Result<Vec<_>, _>>()?;
247        Ok(tasks)
248    }
249
250    pub fn get_task_by_id(&self, task_id: i64) -> Result<Option<ScheduledTask>, MicroClawError> {
251        let conn = self.lock_conn();
252        let result = conn.query_row(
253            "SELECT id, chat_id, prompt, schedule_type, schedule_value, timezone, next_run, last_run, status, created_at, exit_criteria, run_count, max_runs, not_after
254             FROM scheduled_tasks
255             WHERE id = ?1",
256            params![task_id],
257            |row| {
258                Ok(ScheduledTask {
259                    id: row.get(0)?,
260                    chat_id: row.get(1)?,
261                    prompt: row.get(2)?,
262                    schedule_type: row.get(3)?,
263                    schedule_value: row.get(4)?,
264                    timezone: row.get(5)?,
265                    next_run: row.get(6)?,
266                    last_run: row.get(7)?,
267                    status: row.get(8)?,
268                    created_at: row.get(9)?,
269                    exit_criteria: row.get(10)?,
270                    run_count: row.get(11)?,
271                    max_runs: row.get(12)?,
272                    not_after: row.get(13)?,
273                })
274            },
275        );
276        match result {
277            Ok(task) => Ok(Some(task)),
278            Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
279            Err(e) => Err(e.into()),
280        }
281    }
282
283    /// Fetch one scheduled task by id regardless of status (management/
284    /// detail views and tests need retired tasks too).
285    pub fn get_scheduled_task(
286        &self,
287        task_id: i64,
288    ) -> Result<Option<ScheduledTask>, MicroClawError> {
289        let conn = self.lock_conn();
290        let mut stmt = conn.prepare(
291            "SELECT id, chat_id, prompt, schedule_type, schedule_value, timezone, next_run, last_run, status, created_at, exit_criteria, run_count, max_runs, not_after
292             FROM scheduled_tasks WHERE id = ?1",
293        )?;
294        let task = stmt
295            .query_map(params![task_id], |row| {
296                Ok(ScheduledTask {
297                    id: row.get(0)?,
298                    chat_id: row.get(1)?,
299                    prompt: row.get(2)?,
300                    schedule_type: row.get(3)?,
301                    schedule_value: row.get(4)?,
302                    timezone: row.get(5)?,
303                    next_run: row.get(6)?,
304                    last_run: row.get(7)?,
305                    status: row.get(8)?,
306                    created_at: row.get(9)?,
307                    exit_criteria: row.get(10)?,
308                    run_count: row.get(11)?,
309                    max_runs: row.get(12)?,
310                    not_after: row.get(13)?,
311                })
312            })?
313            .next()
314            .transpose()?;
315        Ok(task)
316    }
317
318    pub fn update_task_status(&self, task_id: i64, status: &str) -> Result<bool, MicroClawError> {
319        let conn = self.lock_conn();
320        let rows = conn.execute(
321            "UPDATE scheduled_tasks SET status = ?1 WHERE id = ?2",
322            params![status, task_id],
323        )?;
324        Ok(rows > 0)
325    }
326
327    pub fn requeue_scheduled_task(
328        &self,
329        task_id: i64,
330        next_run: &str,
331    ) -> Result<bool, MicroClawError> {
332        let conn = self.lock_conn();
333        let rows = conn.execute(
334            "UPDATE scheduled_tasks
335             SET status = 'active', next_run = ?1
336             WHERE id = ?2",
337            params![next_run, task_id],
338        )?;
339        Ok(rows > 0)
340    }
341
342    pub fn update_task_after_run(
343        &self,
344        task_id: i64,
345        last_run: &str,
346        next_run: Option<&str>,
347        success: bool,
348    ) -> Result<(), MicroClawError> {
349        self.update_task_after_run_lifecycle(task_id, last_run, next_run, success, false)
350    }
351
352    /// `lifecycle_finished` marks a RECURRING task that just retired on
353    /// purpose (max_runs reached / not_after passed): it becomes `completed`
354    /// regardless of this run's outcome, so DLQ auto-replay never resurrects
355    /// a task whose lifecycle has ended. `run_count` increments on every run.
356    pub fn update_task_after_run_lifecycle(
357        &self,
358        task_id: i64,
359        last_run: &str,
360        next_run: Option<&str>,
361        success: bool,
362        lifecycle_finished: bool,
363    ) -> Result<(), MicroClawError> {
364        let conn = self.lock_conn();
365        match next_run {
366            Some(next) => {
367                // Recurring task: reschedule and stay active regardless of this
368                // run's outcome (a transient failure retries on the next tick).
369                conn.execute(
370                    "UPDATE scheduled_tasks
371                     SET last_run = ?1, next_run = ?2, status = 'active',
372                         run_count = run_count + 1
373                     WHERE id = ?3",
374                    params![last_run, next, task_id],
375                )?;
376            }
377            None => {
378                // One-shot task: reflect the actual outcome. A failed one-shot
379                // becomes 'failed' (and is recorded in the DLQ) rather than
380                // masquerading as 'completed'. A lifecycle retirement is always
381                // 'completed' — the task did all the running it was asked to.
382                let status = if lifecycle_finished || success {
383                    "completed"
384                } else {
385                    "failed"
386                };
387                conn.execute(
388                    "UPDATE scheduled_tasks
389                     SET last_run = ?1, status = ?2, run_count = run_count + 1
390                     WHERE id = ?3",
391                    params![last_run, status, task_id],
392                )?;
393            }
394        }
395        Ok(())
396    }
397
398    pub fn recover_running_tasks(&self) -> Result<usize, MicroClawError> {
399        let conn = self.lock_conn();
400        let rows = conn.execute(
401            "UPDATE scheduled_tasks
402             SET status = 'active'
403             WHERE status = 'running'",
404            [],
405        )?;
406        Ok(rows)
407    }
408
409    // --- Task run logs ---
410
411    #[allow(clippy::too_many_arguments)]
412    pub fn log_task_run(
413        &self,
414        task_id: i64,
415        chat_id: i64,
416        started_at: &str,
417        finished_at: &str,
418        duration_ms: i64,
419        success: bool,
420        result_summary: Option<&str>,
421    ) -> Result<i64, MicroClawError> {
422        let conn = self.lock_conn();
423        conn.execute(
424            "INSERT INTO task_run_logs (task_id, chat_id, started_at, finished_at, duration_ms, success, result_summary)
425             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
426            params![
427                task_id,
428                chat_id,
429                started_at,
430                finished_at,
431                duration_ms,
432                success as i32,
433                result_summary,
434            ],
435        )?;
436        Ok(conn.last_insert_rowid())
437    }
438
439    pub fn get_task_run_logs(
440        &self,
441        task_id: i64,
442        limit: usize,
443    ) -> Result<Vec<TaskRunLog>, MicroClawError> {
444        let conn = self.lock_conn();
445        let mut stmt = conn.prepare(
446            "SELECT id, task_id, chat_id, started_at, finished_at, duration_ms, success, result_summary
447             FROM task_run_logs
448             WHERE task_id = ?1
449             ORDER BY id DESC
450             LIMIT ?2",
451        )?;
452        let logs = stmt
453            .query_map(params![task_id, limit as i64], |row| {
454                Ok(TaskRunLog {
455                    id: row.get(0)?,
456                    task_id: row.get(1)?,
457                    chat_id: row.get(2)?,
458                    started_at: row.get(3)?,
459                    finished_at: row.get(4)?,
460                    duration_ms: row.get(5)?,
461                    success: row.get::<_, i32>(6)? != 0,
462                    result_summary: row.get(7)?,
463                })
464            })?
465            .collect::<Result<Vec<_>, _>>()?;
466        Ok(logs)
467    }
468
469    pub fn get_task_run_summary_since(
470        &self,
471        since: Option<&str>,
472    ) -> Result<(i64, i64), MicroClawError> {
473        let conn = self.lock_conn();
474        if let Some(since) = since {
475            let (total, success): (i64, i64) = conn.query_row(
476                "SELECT
477                    COUNT(*) AS total_runs,
478                    COALESCE(SUM(CASE WHEN success != 0 THEN 1 ELSE 0 END), 0) AS success_runs
479                 FROM task_run_logs
480                 WHERE started_at >= ?1",
481                params![since],
482                |row| Ok((row.get(0)?, row.get(1)?)),
483            )?;
484            Ok((total, success))
485        } else {
486            let (total, success): (i64, i64) = conn.query_row(
487                "SELECT
488                    COUNT(*) AS total_runs,
489                    COALESCE(SUM(CASE WHEN success != 0 THEN 1 ELSE 0 END), 0) AS success_runs
490                 FROM task_run_logs",
491                [],
492                |row| Ok((row.get(0)?, row.get(1)?)),
493            )?;
494            Ok((total, success))
495        }
496    }
497
498    pub fn insert_scheduled_task_dlq(
499        &self,
500        task_id: i64,
501        chat_id: i64,
502        started_at: &str,
503        finished_at: &str,
504        duration_ms: i64,
505        error_summary: Option<&str>,
506    ) -> Result<i64, MicroClawError> {
507        let conn = self.lock_conn();
508        let failed_at = chrono::Utc::now().to_rfc3339();
509        conn.execute(
510            "INSERT INTO scheduled_task_dlq (
511                task_id, chat_id, failed_at, started_at, finished_at, duration_ms, error_summary
512             ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
513            params![
514                task_id,
515                chat_id,
516                failed_at,
517                started_at,
518                finished_at,
519                duration_ms,
520                error_summary
521            ],
522        )?;
523        Ok(conn.last_insert_rowid())
524    }
525
526    pub fn list_scheduled_task_dlq(
527        &self,
528        chat_id: Option<i64>,
529        task_id: Option<i64>,
530        include_replayed: bool,
531        limit: usize,
532    ) -> Result<Vec<ScheduledTaskDlqEntry>, MicroClawError> {
533        let conn = self.lock_conn();
534        let replay_filter = if include_replayed {
535            ""
536        } else {
537            " AND replayed_at IS NULL"
538        };
539        let mapper = |row: &rusqlite::Row<'_>| {
540            Ok(ScheduledTaskDlqEntry {
541                id: row.get(0)?,
542                task_id: row.get(1)?,
543                chat_id: row.get(2)?,
544                failed_at: row.get(3)?,
545                started_at: row.get(4)?,
546                finished_at: row.get(5)?,
547                duration_ms: row.get(6)?,
548                error_summary: row.get(7)?,
549                replayed_at: row.get(8)?,
550                replay_note: row.get(9)?,
551            })
552        };
553        let query = match (chat_id, task_id) {
554            (Some(_), Some(_)) => format!(
555                "SELECT id, task_id, chat_id, failed_at, started_at, finished_at, duration_ms,
556                        error_summary, replayed_at, replay_note
557                 FROM scheduled_task_dlq
558                 WHERE chat_id = ?1 AND task_id = ?2{replay_filter}
559                 ORDER BY failed_at DESC LIMIT ?3"
560            ),
561            (Some(_), None) => format!(
562                "SELECT id, task_id, chat_id, failed_at, started_at, finished_at, duration_ms,
563                        error_summary, replayed_at, replay_note
564                 FROM scheduled_task_dlq
565                 WHERE chat_id = ?1{replay_filter}
566                 ORDER BY failed_at DESC LIMIT ?2"
567            ),
568            (None, Some(_)) => format!(
569                "SELECT id, task_id, chat_id, failed_at, started_at, finished_at, duration_ms,
570                        error_summary, replayed_at, replay_note
571                 FROM scheduled_task_dlq
572                 WHERE task_id = ?1{replay_filter}
573                 ORDER BY failed_at DESC LIMIT ?2"
574            ),
575            (None, None) => format!(
576                "SELECT id, task_id, chat_id, failed_at, started_at, finished_at, duration_ms,
577                        error_summary, replayed_at, replay_note
578                 FROM scheduled_task_dlq
579                 WHERE 1=1{replay_filter}
580                 ORDER BY failed_at DESC LIMIT ?1"
581            ),
582        };
583        let mut stmt = conn.prepare(&query)?;
584        match (chat_id, task_id) {
585            (Some(c), Some(t)) => stmt
586                .query_map(params![c, t, limit as i64], mapper)?
587                .collect::<Result<Vec<_>, _>>()
588                .map_err(Into::into),
589            (Some(c), None) => stmt
590                .query_map(params![c, limit as i64], mapper)?
591                .collect::<Result<Vec<_>, _>>()
592                .map_err(Into::into),
593            (None, Some(t)) => stmt
594                .query_map(params![t, limit as i64], mapper)?
595                .collect::<Result<Vec<_>, _>>()
596                .map_err(Into::into),
597            (None, None) => stmt
598                .query_map(params![limit as i64], mapper)?
599                .collect::<Result<Vec<_>, _>>()
600                .map_err(Into::into),
601        }
602    }
603
604    pub fn mark_scheduled_task_dlq_replayed(
605        &self,
606        dlq_id: i64,
607        note: Option<&str>,
608    ) -> Result<bool, MicroClawError> {
609        let conn = self.lock_conn();
610        let replayed_at = chrono::Utc::now().to_rfc3339();
611        let rows = conn.execute(
612            "UPDATE scheduled_task_dlq
613             SET replayed_at = ?1, replay_note = ?2
614             WHERE id = ?3",
615            params![replayed_at, note, dlq_id],
616        )?;
617        Ok(rows > 0)
618    }
619
620    #[allow(dead_code)]
621    pub fn delete_task(&self, task_id: i64) -> Result<bool, MicroClawError> {
622        let conn = self.lock_conn();
623        let rows = conn.execute(
624            "DELETE FROM scheduled_tasks WHERE id = ?1",
625            params![task_id],
626        )?;
627        Ok(rows > 0)
628    }
629
630    // --- Sessions ---
631
632    /// Exact depth of the scheduled-task dead-letter queue (unlike
633    /// `list_scheduled_task_dlq`, this is not clamped by a LIMIT).
634    pub fn count_scheduled_task_dlq(&self, include_replayed: bool) -> Result<i64, MicroClawError> {
635        let conn = self.lock_conn();
636        let replay_filter = if include_replayed {
637            ""
638        } else {
639            " WHERE replayed_at IS NULL"
640        };
641        let count: i64 = conn.query_row(
642            &format!("SELECT COUNT(*) FROM scheduled_task_dlq{replay_filter}"),
643            [],
644            |row| row.get(0),
645        )?;
646        Ok(count)
647    }
648}
649
650impl Database {
651    pub fn claim_due_tasks(
652        &self,
653        now: &str,
654        limit: usize,
655    ) -> Result<Vec<ScheduledTask>, MicroClawError> {
656        let conn = self.lock_conn();
657        let tx = conn.unchecked_transaction()?;
658
659        let mut stmt = tx.prepare(
660            "SELECT id, chat_id, prompt, schedule_type, schedule_value, timezone, next_run, last_run, status, created_at, exit_criteria, run_count, max_runs, not_after
661             FROM scheduled_tasks
662             WHERE status = 'active' AND next_run <= ?1
663             ORDER BY next_run ASC, id ASC
664             LIMIT ?2",
665        )?;
666        let candidates = stmt
667            .query_map(params![now, limit as i64], |row| {
668                Ok(ScheduledTask {
669                    id: row.get(0)?,
670                    chat_id: row.get(1)?,
671                    prompt: row.get(2)?,
672                    schedule_type: row.get(3)?,
673                    schedule_value: row.get(4)?,
674                    timezone: row.get(5)?,
675                    next_run: row.get(6)?,
676                    last_run: row.get(7)?,
677                    status: row.get(8)?,
678                    created_at: row.get(9)?,
679                    exit_criteria: row.get(10)?,
680                    run_count: row.get(11)?,
681                    max_runs: row.get(12)?,
682                    not_after: row.get(13)?,
683                })
684            })?
685            .collect::<Result<Vec<_>, _>>()?;
686        drop(stmt);
687
688        let mut claimed = Vec::new();
689        for task in candidates {
690            let rows = tx.execute(
691                "UPDATE scheduled_tasks
692                 SET status = 'running'
693                 WHERE id = ?1 AND status = 'active' AND next_run <= ?2",
694                params![task.id, now],
695            )?;
696            if rows > 0 {
697                let mut claimed_task = task;
698                claimed_task.status = "running".to_string();
699                claimed.push(claimed_task);
700            }
701        }
702
703        tx.commit()?;
704        Ok(claimed)
705    }
706}
707
708#[cfg(test)]
709mod tests {
710    use super::*;
711    use crate::internal::storage::db::test_support::*;
712
713    #[test]
714    fn test_create_and_get_scheduled_task() {
715        let (db, dir) = test_db();
716        let id = db
717            .create_scheduled_task(
718                100,
719                "say hello",
720                "cron",
721                "0 */5 * * * *",
722                "2024-06-01T00:05:00Z",
723            )
724            .unwrap();
725        assert!(id > 0);
726
727        let tasks = db.get_tasks_for_chat(100).unwrap();
728        assert_eq!(tasks.len(), 1);
729        assert_eq!(tasks[0].prompt, "say hello");
730        assert_eq!(tasks[0].schedule_type, "cron");
731        assert_eq!(tasks[0].status, "active");
732        cleanup(&dir);
733    }
734
735    #[test]
736    fn test_get_due_tasks() {
737        let (db, dir) = test_db();
738        db.create_scheduled_task(100, "task1", "cron", "0 * * * * *", "2024-01-01T00:00:00Z")
739            .unwrap();
740        db.create_scheduled_task(
741            100,
742            "task2",
743            "once",
744            "2099-12-31T00:00:00Z",
745            "2099-12-31T00:00:00Z",
746        )
747        .unwrap();
748
749        // Only task1 is due
750        let due = db.get_due_tasks("2024-06-01T00:00:00Z").unwrap();
751        assert_eq!(due.len(), 1);
752        assert_eq!(due[0].prompt, "task1");
753
754        // Both are due in the far future
755        let due = db.get_due_tasks("2100-01-01T00:00:00Z").unwrap();
756        assert_eq!(due.len(), 2);
757        cleanup(&dir);
758    }
759
760    #[test]
761    fn test_get_tasks_for_chat_filters_status() {
762        let (db, dir) = test_db();
763        let id1 = db
764            .create_scheduled_task(
765                100,
766                "active task",
767                "cron",
768                "0 * * * * *",
769                "2024-01-01T00:00:00Z",
770            )
771            .unwrap();
772        let id2 = db
773            .create_scheduled_task(
774                100,
775                "to cancel",
776                "once",
777                "2024-01-01T00:00:00Z",
778                "2024-01-01T00:00:00Z",
779            )
780            .unwrap();
781        db.update_task_status(id2, "cancelled").unwrap();
782
783        // Only active/paused tasks should be returned
784        let tasks = db.get_tasks_for_chat(100).unwrap();
785        assert_eq!(tasks.len(), 1);
786        assert_eq!(tasks[0].id, id1);
787
788        // Pause the active one
789        db.update_task_status(id1, "paused").unwrap();
790        let tasks = db.get_tasks_for_chat(100).unwrap();
791        assert_eq!(tasks.len(), 1);
792        assert_eq!(tasks[0].status, "paused");
793        cleanup(&dir);
794    }
795
796    #[test]
797    fn test_update_task_status() {
798        let (db, dir) = test_db();
799        let id = db
800            .create_scheduled_task(100, "test", "cron", "0 * * * * *", "2024-01-01T00:00:00Z")
801            .unwrap();
802
803        assert!(db.update_task_status(id, "paused").unwrap());
804        assert!(db.update_task_status(id, "active").unwrap());
805        assert!(db.update_task_status(id, "cancelled").unwrap());
806
807        // Non-existent task
808        assert!(!db.update_task_status(9999, "paused").unwrap());
809        cleanup(&dir);
810    }
811
812    #[test]
813    fn test_requeue_scheduled_task() {
814        let (db, dir) = test_db();
815        let id = db
816            .create_scheduled_task(100, "test", "cron", "0 * * * * *", "2024-01-01T00:00:00Z")
817            .unwrap();
818        db.update_task_status(id, "paused").unwrap();
819
820        assert!(db
821            .requeue_scheduled_task(id, "2099-01-01T00:00:00Z")
822            .unwrap());
823        let task = db.get_task_by_id(id).unwrap().unwrap();
824        assert_eq!(task.status, "active");
825        assert_eq!(task.next_run, "2099-01-01T00:00:00Z");
826
827        assert!(!db
828            .requeue_scheduled_task(9999, "2099-01-01T00:00:00Z")
829            .unwrap());
830        cleanup(&dir);
831    }
832
833    #[test]
834    fn test_update_task_after_run_cron() {
835        let (db, dir) = test_db();
836        let id = db
837            .create_scheduled_task(100, "test", "cron", "0 * * * * *", "2024-01-01T00:00:00Z")
838            .unwrap();
839
840        db.update_task_after_run(
841            id,
842            "2024-01-01T00:01:00Z",
843            Some("2024-01-01T00:02:00Z"),
844            true,
845        )
846        .unwrap();
847
848        let tasks = db.get_tasks_for_chat(100).unwrap();
849        assert_eq!(tasks[0].last_run.as_deref(), Some("2024-01-01T00:01:00Z"));
850        assert_eq!(tasks[0].next_run, "2024-01-01T00:02:00Z");
851        assert_eq!(tasks[0].status, "active");
852        cleanup(&dir);
853    }
854
855    #[test]
856    fn test_recover_running_tasks() {
857        let (db, dir) = test_db();
858        let id = db
859            .create_scheduled_task(100, "test", "cron", "0 * * * * *", "2024-01-01T00:00:00Z")
860            .unwrap();
861        assert_eq!(
862            db.claim_due_tasks("2024-01-01T00:00:00Z", 10)
863                .unwrap()
864                .len(),
865            1
866        );
867
868        let recovered = db.recover_running_tasks().unwrap();
869        assert_eq!(recovered, 1);
870        let task = db.get_task_by_id(id).unwrap().unwrap();
871        assert_eq!(task.status, "active");
872        cleanup(&dir);
873    }
874
875    #[test]
876    fn test_update_task_after_run_one_shot() {
877        let (db, dir) = test_db();
878        let id = db
879            .create_scheduled_task(
880                100,
881                "test",
882                "once",
883                "2024-01-01T00:00:00Z",
884                "2024-01-01T00:00:00Z",
885            )
886            .unwrap();
887
888        // One-shot success: no next_run, should mark as completed
889        db.update_task_after_run(id, "2024-01-01T00:00:00Z", None, true)
890            .unwrap();
891
892        // Should not appear in active/paused list
893        let tasks = db.get_tasks_for_chat(100).unwrap();
894        assert!(tasks.is_empty());
895        assert_eq!(db.get_task_by_id(id).unwrap().unwrap().status, "completed");
896        cleanup(&dir);
897    }
898
899    #[test]
900    fn test_update_task_after_run_one_shot_failure_marked_failed() {
901        let (db, dir) = test_db();
902        let id = db
903            .create_scheduled_task(
904                100,
905                "test",
906                "once",
907                "2024-01-01T00:00:00Z",
908                "2024-01-01T00:00:00Z",
909            )
910            .unwrap();
911
912        // A failed one-shot must be recorded as 'failed', not 'completed'.
913        db.update_task_after_run(id, "2024-01-01T00:00:00Z", None, false)
914            .unwrap();
915
916        assert_eq!(db.get_task_by_id(id).unwrap().unwrap().status, "failed");
917        cleanup(&dir);
918    }
919
920    #[test]
921    fn test_delete_task() {
922        let (db, dir) = test_db();
923        let id = db
924            .create_scheduled_task(100, "test", "cron", "0 * * * * *", "2024-01-01T00:00:00Z")
925            .unwrap();
926
927        assert!(db.delete_task(id).unwrap());
928        assert!(!db.delete_task(id).unwrap()); // already deleted
929
930        let tasks = db.get_tasks_for_chat(100).unwrap();
931        assert!(tasks.is_empty());
932        cleanup(&dir);
933    }
934
935    #[test]
936    fn test_log_task_run() {
937        let (db, dir) = test_db();
938        let task_id = db
939            .create_scheduled_task(100, "test", "cron", "0 * * * * *", "2024-01-01T00:00:00Z")
940            .unwrap();
941
942        let log_id = db
943            .log_task_run(
944                task_id,
945                100,
946                "2024-01-01T00:00:00Z",
947                "2024-01-01T00:00:05Z",
948                5000,
949                true,
950                Some("Success"),
951            )
952            .unwrap();
953        assert!(log_id > 0);
954
955        let logs = db.get_task_run_logs(task_id, 10).unwrap();
956        assert_eq!(logs.len(), 1);
957        assert_eq!(logs[0].task_id, task_id);
958        assert_eq!(logs[0].duration_ms, 5000);
959        assert!(logs[0].success);
960        assert_eq!(logs[0].result_summary.as_deref(), Some("Success"));
961        cleanup(&dir);
962    }
963
964    #[test]
965    fn test_get_task_run_logs_ordering_and_limit() {
966        let (db, dir) = test_db();
967        let task_id = db
968            .create_scheduled_task(100, "test", "cron", "0 * * * * *", "2024-01-01T00:00:00Z")
969            .unwrap();
970
971        for i in 0..5 {
972            db.log_task_run(
973                task_id,
974                100,
975                &format!("2024-01-01T00:0{i}:00Z"),
976                &format!("2024-01-01T00:0{i}:05Z"),
977                5000,
978                true,
979                Some(&format!("Run {i}")),
980            )
981            .unwrap();
982        }
983
984        // Limit to 3, most recent first
985        let logs = db.get_task_run_logs(task_id, 3).unwrap();
986        assert_eq!(logs.len(), 3);
987        assert_eq!(logs[0].result_summary.as_deref(), Some("Run 4")); // most recent
988        assert_eq!(logs[2].result_summary.as_deref(), Some("Run 2"));
989        cleanup(&dir);
990    }
991
992    #[test]
993    fn test_get_task_run_summary_since() {
994        let (db, dir) = test_db();
995        let task_id = db
996            .create_scheduled_task(100, "test", "cron", "0 * * * * *", "2024-01-01T00:00:00Z")
997            .unwrap();
998        db.log_task_run(
999            task_id,
1000            100,
1001            "2024-01-01T00:00:00Z",
1002            "2024-01-01T00:00:05Z",
1003            5000,
1004            true,
1005            Some("ok"),
1006        )
1007        .unwrap();
1008        db.log_task_run(
1009            task_id,
1010            100,
1011            "2024-01-01T00:10:00Z",
1012            "2024-01-01T00:10:05Z",
1013            5000,
1014            false,
1015            Some("fail"),
1016        )
1017        .unwrap();
1018
1019        let (total_all, success_all) = db.get_task_run_summary_since(None).unwrap();
1020        assert_eq!(total_all, 2);
1021        assert_eq!(success_all, 1);
1022
1023        let (total_since, success_since) = db
1024            .get_task_run_summary_since(Some("2024-01-01T00:05:00Z"))
1025            .unwrap();
1026        assert_eq!(total_since, 1);
1027        assert_eq!(success_since, 0);
1028        cleanup(&dir);
1029    }
1030
1031    #[test]
1032    fn test_count_scheduled_task_dlq_is_not_limit_clamped() {
1033        let (db, dir) = test_db();
1034        let task_id = db
1035            .create_scheduled_task(100, "test", "cron", "0 * * * * *", "2024-01-01T00:00:00Z")
1036            .unwrap();
1037        assert_eq!(db.count_scheduled_task_dlq(false).unwrap(), 0);
1038        for i in 0..3 {
1039            db.insert_scheduled_task_dlq(
1040                task_id,
1041                100,
1042                "2024-01-01T00:00:00Z",
1043                "2024-01-01T00:00:05Z",
1044                5000 + i,
1045                Some("Error: timeout"),
1046            )
1047            .unwrap();
1048        }
1049        assert_eq!(db.count_scheduled_task_dlq(false).unwrap(), 3);
1050        assert_eq!(db.count_scheduled_task_dlq(true).unwrap(), 3);
1051        cleanup(&dir);
1052    }
1053
1054    #[test]
1055    fn test_scheduled_task_dlq_insert_list_and_mark_replayed() {
1056        let (db, dir) = test_db();
1057        let task_id = db
1058            .create_scheduled_task(100, "test", "cron", "0 * * * * *", "2024-01-01T00:00:00Z")
1059            .unwrap();
1060
1061        let dlq_id = db
1062            .insert_scheduled_task_dlq(
1063                task_id,
1064                100,
1065                "2024-01-01T00:00:00Z",
1066                "2024-01-01T00:00:05Z",
1067                5000,
1068                Some("Error: timeout"),
1069            )
1070            .unwrap();
1071        assert!(dlq_id > 0);
1072
1073        let pending = db
1074            .list_scheduled_task_dlq(Some(100), Some(task_id), false, 10)
1075            .unwrap();
1076        assert_eq!(pending.len(), 1);
1077        assert_eq!(pending[0].task_id, task_id);
1078        assert_eq!(pending[0].replayed_at, None);
1079
1080        assert!(db
1081            .mark_scheduled_task_dlq_replayed(dlq_id, Some("queued replay"))
1082            .unwrap());
1083
1084        let pending_after = db
1085            .list_scheduled_task_dlq(Some(100), Some(task_id), false, 10)
1086            .unwrap();
1087        assert!(pending_after.is_empty());
1088
1089        let all = db
1090            .list_scheduled_task_dlq(Some(100), Some(task_id), true, 10)
1091            .unwrap();
1092        assert_eq!(all.len(), 1);
1093        assert!(all[0].replayed_at.is_some());
1094        assert_eq!(all[0].replay_note.as_deref(), Some("queued replay"));
1095        cleanup(&dir);
1096    }
1097
1098    #[test]
1099    fn test_clear_chat_context_removes_session_messages_and_scheduled_tasks() {
1100        let (db, dir) = test_db();
1101        db.upsert_chat(100, Some("chat-100"), "private").unwrap();
1102        db.save_session(100, r#"[{"role":"user","content":"hi"}]"#)
1103            .unwrap();
1104        db.store_message(&StoredMessage {
1105            id: "m1".into(),
1106            chat_id: 100,
1107            sender_name: "alice".into(),
1108            content: "hello".into(),
1109            is_from_bot: false,
1110            timestamp: "2024-01-01T00:00:01Z".into(),
1111        })
1112        .unwrap();
1113        db.insert_memory(Some(100), "User likes Rust", "PROFILE")
1114            .unwrap();
1115        let task_id = db
1116            .create_scheduled_task(
1117                100,
1118                "daily summary",
1119                "cron",
1120                "0 0 8 * * *",
1121                "2099-01-01T08:00:00Z",
1122            )
1123            .unwrap();
1124        db.log_task_run(
1125            task_id,
1126            100,
1127            "2024-01-01T08:00:00Z",
1128            "2024-01-01T08:00:01Z",
1129            1000,
1130            true,
1131            Some("ok"),
1132        )
1133        .unwrap();
1134        db.insert_scheduled_task_dlq(
1135            task_id,
1136            100,
1137            "2024-01-01T09:00:00Z",
1138            "2024-01-01T09:00:01Z",
1139            1000,
1140            Some("failure"),
1141        )
1142        .unwrap();
1143
1144        assert!(db.clear_chat_context(100).unwrap());
1145        assert!(db.load_session(100).unwrap().is_none());
1146        assert!(db.get_recent_messages(100, 10).unwrap().is_empty());
1147        assert!(db.get_tasks_for_chat(100).unwrap().is_empty());
1148        assert!(db.get_task_run_logs(task_id, 10).unwrap().is_empty());
1149        assert!(db
1150            .list_scheduled_task_dlq(Some(100), Some(task_id), true, 10)
1151            .unwrap()
1152            .is_empty());
1153        assert!(!db.search_memories(100, "Rust", 10).unwrap().is_empty());
1154        assert!(db.get_chat_type(100).unwrap().is_some());
1155
1156        cleanup(&dir);
1157    }
1158
1159    #[test]
1160    fn test_clear_chat_conversation_keeps_scheduled_tasks() {
1161        let (db, dir) = test_db();
1162        db.upsert_chat(100, Some("chat-100"), "private").unwrap();
1163        db.save_session(100, r#"[{"role":"user","content":"hi"}]"#)
1164            .unwrap();
1165        db.store_message(&StoredMessage {
1166            id: "m1".into(),
1167            chat_id: 100,
1168            sender_name: "alice".into(),
1169            content: "hello".into(),
1170            is_from_bot: false,
1171            timestamp: "2024-01-01T00:00:01Z".into(),
1172        })
1173        .unwrap();
1174        db.insert_memory(Some(100), "User likes Rust", "PROFILE")
1175            .unwrap();
1176        let task_id = db
1177            .create_scheduled_task(
1178                100,
1179                "daily summary",
1180                "cron",
1181                "0 0 8 * * *",
1182                "2099-01-01T08:00:00Z",
1183            )
1184            .unwrap();
1185        db.log_task_run(
1186            task_id,
1187            100,
1188            "2024-01-01T08:00:00Z",
1189            "2024-01-01T08:00:01Z",
1190            1000,
1191            true,
1192            Some("ok"),
1193        )
1194        .unwrap();
1195        db.insert_scheduled_task_dlq(
1196            task_id,
1197            100,
1198            "2024-01-01T09:00:00Z",
1199            "2024-01-01T09:00:01Z",
1200            1000,
1201            Some("failure"),
1202        )
1203        .unwrap();
1204
1205        assert!(db.clear_chat_conversation(100).unwrap());
1206        assert!(db.load_session(100).unwrap().is_none());
1207        assert!(db.get_recent_messages(100, 10).unwrap().is_empty());
1208        assert_eq!(db.get_tasks_for_chat(100).unwrap().len(), 1);
1209        assert_eq!(db.get_task_run_logs(task_id, 10).unwrap().len(), 1);
1210        assert_eq!(
1211            db.list_scheduled_task_dlq(Some(100), Some(task_id), true, 10)
1212                .unwrap()
1213                .len(),
1214            1
1215        );
1216        assert!(!db.search_memories(100, "Rust", 10).unwrap().is_empty());
1217        assert!(db.get_chat_type(100).unwrap().is_some());
1218
1219        cleanup(&dir);
1220    }
1221}