scheduled tasks

src/scheduler.rs

@@ -0,0 +1,118 @@
+use crate::agent::Agent;
+use crate::entities::task::Entity as Task;
+use crate::entities::task_run;
+use dashmap::DashMap;
+use sea_orm::{DatabaseConnection, EntityTrait, Set};
+use tokio_cron_scheduler::{Job, JobScheduler};
+use uuid::Uuid;
+
+pub struct Scheduler {
+    scheduler: JobScheduler,
+    db: DatabaseConnection,
+    tasks_to_jobs: DashMap<Uuid, Uuid>,
+    zen_api_key: Option<String>,
+    tavily_api_key: Option<String>,
+}
+
+impl Scheduler {
+    pub async fn new(
+        db: DatabaseConnection,
+        zen_api_key: Option<String>,
+        tavily_api_key: Option<String>,
+    ) -> Result<Self, Box<dyn std::error::Error>> {
+        let scheduler = JobScheduler::new().await?;
+        scheduler.start().await?;
+        Ok(Self {
+            scheduler,
+            db,
+            tasks_to_jobs: DashMap::new(),
+            zen_api_key,
+            tavily_api_key,
+        })
+    }
+
+    pub async fn add_task_job(
+        &self,
+        task_id: Uuid,
+        cron_expr: &str,
+    ) -> Result<(), Box<dyn std::error::Error>> {
+        // Remove existing job if any
+        if let Some((_, old_job_id)) = self.tasks_to_jobs.remove(&task_id) {
+            let _ = self.scheduler.remove(&old_job_id).await;
+        }
+
+        let db = self.db.clone();
+        let zen_key = self.zen_api_key.clone();
+        let tavily_key = self.tavily_api_key.clone();
+
+        let job = Job::new_async(cron_expr, move |_uuid, _l| {
+            let db = db.clone();
+            let zen_key = zen_key.clone();
+            let tavily_key = tavily_key.clone();
+            Box::pin(async move {
+                if let Err(e) = Self::run_task(db, zen_key, tavily_key, task_id).await {
+                    eprintln!("Error in scheduled task {}: {}", task_id, e);
+                }
+            })
+        })?;
+
+        let job_id = self.scheduler.add(job).await?;
+        self.tasks_to_jobs.insert(task_id, job_id);
+
+        Ok(())
+    }
+
+    pub async fn remove_task_job(&self, task_id: Uuid) -> Result<(), Box<dyn std::error::Error>> {
+        if let Some((_, job_id)) = self.tasks_to_jobs.remove(&task_id) {
+            self.scheduler.remove(&job_id).await?;
+        }
+        Ok(())
+    }
+
+    async fn run_task(
+        db: DatabaseConnection,
+        zen_key: Option<String>,
+        tavily_key: Option<String>,
+        task_id: Uuid,
+    ) -> Result<(), Box<dyn std::error::Error>> {
+        let task = Task::find_by_id(task_id)
+            .one(&db)
+            .await?
+            .ok_or("Task not found")?;
+
+        // Create a new run entry
+        let run_id = Uuid::new_v4();
+        let run = task_run::ActiveModel {
+            id: Set(run_id),
+            task_id: Set(task_id),
+            status: Set("running".to_string()),
+            logs: Set(String::new()),
+            answer: Set(None),
+            created_at: Set(chrono::Utc::now().into()),
+        };
+
+        use sea_orm::ActiveModelTrait;
+        run.insert(&db).await?;
+
+        // Start agent in background
+        let mut agent = Agent::new(zen_key, tavily_key, task.goal.clone());
+
+        tokio::spawn(async move {
+            let (logs, answer) = match agent.run().await {
+                Ok(res) => res,
+                Err(e) => (format!("Scheduled run failed: {}", e), None),
+            };
+
+            let run_complete = task_run::ActiveModel {
+                id: Set(run_id),
+                status: Set("completed".to_string()),
+                logs: Set(logs),
+                answer: Set(answer),
+                ..Default::default()
+            };
+            let _ = run_complete.update(&db).await;
+        });
+
+        Ok(())
+    }
+}