scheduled tasks

src/entities/mod.rs

@@ -1 +1,2 @@
 pub mod task;
+pub mod task_run;

src/entities/task.rs

@@ -7,15 +7,20 @@ pub struct Model {
     #[sea_orm(primary_key, auto_increment = false)]
     pub id: Uuid,
     pub goal: String,
-    pub status: String,
-    #[sea_orm(column_type = "Text")]
-    pub logs: String,
-    #[sea_orm(column_type = "Text", nullable)]
-    pub answer: Option<String>,
+    pub cron: Option<String>,
     pub created_at: DateTimeWithTimeZone,
 }
 
 #[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
-pub enum Relation {}
+pub enum Relation {
+    #[sea_orm(has_many = "super::task_run::Entity")]
+    TaskRuns,
+}
+
+impl Related<super::task_run::Entity> for Entity {
+    fn to() -> RelationDef {
+        Relation::TaskRuns.def()
+    }
+}
 
 impl ActiveModelBehavior for ActiveModel {}

src/entities/task_run.rs

@@ -0,0 +1,36 @@
+use sea_orm::entity::prelude::*;
+use serde::{Deserialize, Serialize};
+
+#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)]
+#[sea_orm(table_name = "task_runs")]
+pub struct Model {
+    #[sea_orm(primary_key, auto_increment = false)]
+    pub id: Uuid,
+    pub task_id: Uuid,
+    pub status: String,
+    #[sea_orm(column_type = "Text")]
+    pub logs: String,
+    #[sea_orm(column_type = "Text", nullable)]
+    pub answer: Option<String>,
+    pub created_at: DateTimeWithTimeZone,
+}
+
+#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+pub enum Relation {
+    #[sea_orm(
+        belongs_to = "super::task::Entity",
+        from = "Column::TaskId",
+        to = "super::task::Column::Id",
+        on_update = "NoAction",
+        on_delete = "Cascade"
+    )]
+    Task,
+}
+
+impl Related<super::task::Entity> for Entity {
+    fn to() -> RelationDef {
+        Relation::Task.def()
+    }
+}
+
+impl ActiveModelBehavior for ActiveModel {}

src/main.rs

@@ -1,6 +1,7 @@
 mod agent;
 mod api;
 mod entities;
+mod scheduler;
 mod server;
 mod tools;
 

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(())
+    }
+}

src/server.rs

@@ -1,23 +1,28 @@
 use axum::{
     Json, Router,
     extract::{Path, State},
-    http::{Method, StatusCode},
+    http::StatusCode,
     routing::{get, post},
 };
 use chrono::Utc;
-use sea_orm::{ActiveModelTrait, Database, DatabaseConnection, EntityTrait, QueryOrder, Set};
+use sea_orm::{
+    ActiveModelTrait, Database, DatabaseConnection, EntityTrait, QueryOrder, QuerySelect, Set,
+};
 use serde::{Deserialize, Serialize};
 use std::sync::Arc;
-use tower_http::cors::{Any, CorsLayer};
+use tower_http::cors::CorsLayer;
 use uuid::Uuid;
 
 use crate::agent::Agent;
 use crate::entities::task::{self, Entity as Task};
+use crate::entities::task_run::{self, Entity as TaskRun};
+use crate::scheduler::Scheduler;
 use migration::{Migrator, MigratorTrait};
 
 #[derive(Clone)]
 pub struct AppState {
     pub db: DatabaseConnection,
+    pub scheduler: Arc<Scheduler>,
     pub zen_api_key: Option<String>,
     pub tavily_api_key: Option<String>,
 }
@@ -25,18 +30,42 @@ pub struct AppState {
 #[derive(Deserialize)]
 pub struct CreateTaskRequest {
     pub goal: String,
+    pub cron: Option<String>,
+}
+
+#[derive(Deserialize)]
+pub struct UpdateTaskRequest {
+    pub goal: String,
+    pub cron: Option<String>,
 }
 
 #[derive(Serialize)]
 pub struct TaskResponse {
     pub id: Uuid,
     pub goal: String,
+    pub cron: Option<String>,
+    pub created_at: chrono::DateTime<chrono::FixedOffset>,
+    pub runs: Vec<TaskRunResponse>,
+}
+
+#[derive(Serialize)]
+pub struct TaskRunResponse {
+    pub id: Uuid,
     pub status: String,
     pub logs: String,
     pub answer: Option<String>,
     pub created_at: chrono::DateTime<chrono::FixedOffset>,
 }
 
+#[derive(Serialize)]
+pub struct RecentRunResponse {
+    pub id: Uuid,
+    pub task_id: Uuid,
+    pub goal: String,
+    pub status: String,
+    pub created_at: chrono::DateTime<chrono::FixedOffset>,
+}
+
 pub async fn start(db_url: &str) -> Result<(), Box<dyn std::error::Error>> {
     let db = Database::connect(db_url).await?;
     Migrator::up(&db, None).await?;
@@ -44,20 +73,31 @@ pub async fn start(db_url: &str) -> Result<(), Box<dyn std::error::Error>> {
     let zen_api_key = std::env::var("ZEN_API_KEY").ok();
     let tavily_api_key = std::env::var("TAVILY_API_KEY").ok();
 
+    let scheduler =
+        Arc::new(Scheduler::new(db.clone(), zen_api_key.clone(), tavily_api_key.clone()).await?);
+
+    // Load existing scheduled tasks
+    let existing_tasks = Task::find().all(&db).await?;
+    for task in existing_tasks {
+        if let Some(cron) = task.cron {
+            let _ = scheduler.add_task_job(task.id, &cron).await;
+        }
+    }
+
     let state = Arc::new(AppState {
         db,
+        scheduler,
         zen_api_key,
         tavily_api_key,
     });
 
-    let cors = CorsLayer::new()
-        .allow_origin(Any)
-        .allow_methods([Method::GET, Method::POST])
-        .allow_headers(Any);
+    let cors = CorsLayer::permissive();
 
     let app = Router::new()
         .route("/tasks", post(create_task).get(list_tasks))
-        .route("/tasks/:id", get(get_task))
+        .route("/tasks/:id", get(get_task).put(update_task))
+        .route("/tasks/:id/runs", post(rerun_task))
+        .route("/runs/recent", get(get_recent_runs))
         .layer(cors)
         .with_state(state);
 
@@ -72,6 +112,7 @@ async fn list_tasks(
     State(state): State<Arc<AppState>>,
 ) -> Result<Json<Vec<TaskResponse>>, (StatusCode, String)> {
     let tasks = Task::find()
+        .find_with_related(TaskRun)
         .order_by_desc(task::Column::CreatedAt)
         .all(&state.db)
         .await
@@ -79,13 +120,21 @@ async fn list_tasks(
 
     let response = tasks
         .into_iter()
-        .map(|t| TaskResponse {
+        .map(|(t, runs)| TaskResponse {
             id: t.id,
             goal: t.goal,
-            status: t.status,
-            logs: t.logs,
-            answer: t.answer,
+            cron: t.cron,
             created_at: t.created_at,
+            runs: runs
+                .into_iter()
+                .map(|r| TaskRunResponse {
+                    id: r.id,
+                    status: r.status,
+                    logs: r.logs,
+                    answer: r.answer,
+                    created_at: r.created_at,
+                })
+                .collect(),
         })
         .collect();
 
@@ -98,13 +147,11 @@ async fn create_task(
 ) -> Result<Json<TaskResponse>, (StatusCode, String)> {
     let task_id = Uuid::new_v4();
 
-    // Initial save
+    // Initial task save
     let new_task = task::ActiveModel {
         id: Set(task_id),
         goal: Set(payload.goal.clone()),
-        status: Set("running".to_string()),
-        logs: Set(String::new()),
-        answer: Set(None),
+        cron: Set(payload.cron.clone()),
         created_at: Set(Utc::now().into()),
     };
 
@@ -113,14 +160,54 @@ async fn create_task(
         .await
         .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
 
-    // Run agent in the background or synchronously for now as requested
-    // "it can be created and executed on the agent with an api. the resulting logs after execution will get returned by the endpoint"
-    // This implies we wait for it to finish.
+    if let Some(cron) = &payload.cron {
+        let _ = state.scheduler.add_task_job(task_id, cron).await;
+    } else {
+        let _ = state.scheduler.remove_task_job(task_id).await;
+    }
+
+    execute_agent_run(state, task_id, payload.goal).await
+}
+
+async fn rerun_task(
+    State(state): State<Arc<AppState>>,
+    Path(id): Path<Uuid>,
+) -> Result<Json<TaskResponse>, (StatusCode, String)> {
+    let task = Task::find_by_id(id)
+        .one(&state.db)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
+        .ok_or((StatusCode::NOT_FOUND, "Task not found".to_string()))?;
+
+    execute_agent_run(state, task.id, task.goal).await
+}
+
+async fn execute_agent_run(
+    state: Arc<AppState>,
+    task_id: Uuid,
+    goal: String,
+) -> Result<Json<TaskResponse>, (StatusCode, String)> {
+    let run_id = Uuid::new_v4();
+
+    // Initial run save
+    let new_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(Utc::now().into()),
+    };
+
+    new_run
+        .insert(&state.db)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
 
     let mut agent = Agent::new(
         state.zen_api_key.clone(),
         state.tavily_api_key.clone(),
-        payload.goal.clone(),
+        goal.clone(),
     );
 
     let (logs, answer) = match agent.run().await {
@@ -129,51 +216,111 @@ async fn create_task(
     };
 
     // Update with final logs and status
-    let mut task: task::ActiveModel = Task::find_by_id(task_id)
+    let mut run: task_run::ActiveModel = TaskRun::find_by_id(run_id)
         .one(&state.db)
         .await
         .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
         .ok_or((
             StatusCode::NOT_FOUND,
-            "Task not found after insert".to_string(),
+            "Run not found after insert".to_string(),
         ))?
         .into();
 
-    task.logs = Set(logs.clone());
-    task.answer = Set(answer.clone());
-    task.status = Set("completed".to_string());
+    run.logs = Set(logs.clone());
+    run.answer = Set(answer.clone());
+    run.status = Set("completed".to_string());
 
-    let updated_task = task
-        .update(&state.db)
+    run.update(&state.db)
         .await
         .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
 
-    Ok(Json(TaskResponse {
-        id: updated_task.id,
-        goal: updated_task.goal,
-        status: updated_task.status,
-        logs: updated_task.logs,
-        answer: updated_task.answer,
-        created_at: updated_task.created_at,
-    }))
+    get_task(Path(task_id), State(state)).await
+}
+
+async fn update_task(
+    State(state): State<Arc<AppState>>,
+    Path(id): Path<Uuid>,
+    Json(payload): Json<UpdateTaskRequest>,
+) -> Result<Json<TaskResponse>, (StatusCode, String)> {
+    let mut task: task::ActiveModel = Task::find_by_id(id)
+        .one(&state.db)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
+        .ok_or((StatusCode::NOT_FOUND, "Task not found".to_string()))?
+        .into();
+
+    task.goal = Set(payload.goal.clone());
+    task.cron = Set(payload.cron.clone());
+
+    task.update(&state.db)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
+
+    if let Some(cron) = &payload.cron {
+        let _ = state.scheduler.add_task_job(id, cron).await;
+    } else {
+        let _ = state.scheduler.remove_task_job(id).await;
+    }
+
+    get_task(Path(id), State(state)).await
 }
 
 async fn get_task(
     Path(id): Path<Uuid>,
     State(state): State<Arc<AppState>>,
 ) -> Result<Json<TaskResponse>, (StatusCode, String)> {
-    let task = Task::find_by_id(id)
-        .one(&state.db)
+    let results = Task::find_by_id(id)
+        .find_with_related(TaskRun)
+        .all(&state.db)
         .await
-        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
+
+    let (task, runs) = results
+        .into_iter()
+        .next()
         .ok_or((StatusCode::NOT_FOUND, "Task not found".to_string()))?;
 
     Ok(Json(TaskResponse {
         id: task.id,
         goal: task.goal,
-        status: task.status,
-        logs: task.logs,
-        answer: task.answer,
+        cron: task.cron,
         created_at: task.created_at,
+        runs: runs
+            .into_iter()
+            .map(|r| TaskRunResponse {
+                id: r.id,
+                status: r.status,
+                logs: r.logs,
+                answer: r.answer,
+                created_at: r.created_at,
+            })
+            .collect(),
     }))
 }
+
+async fn get_recent_runs(
+    State(state): State<Arc<AppState>>,
+) -> Result<Json<Vec<RecentRunResponse>>, (StatusCode, String)> {
+    let results = TaskRun::find()
+        .find_also_related(Task)
+        .order_by_desc(task_run::Column::CreatedAt)
+        .limit(50)
+        .all(&state.db)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
+
+    let response = results
+        .into_iter()
+        .filter_map(|(run, task_opt)| {
+            task_opt.map(|task| RecentRunResponse {
+                id: run.id,
+                task_id: run.task_id,
+                goal: task.goal,
+                status: run.status,
+                created_at: run.created_at,
+            })
+        })
+        .collect();
+
+    Ok(Json(response))
+}

src/tools.rs

@@ -20,28 +20,6 @@ pub fn get_tools() -> Vec<Tool> {
                 }),
             },
         },
-        Tool {
-            tool_type: "function".to_string(),
-            function: FunctionDefinition {
-                name: "write_file".to_string(),
-                description: "Write content to a file. Ensure parent directories exist."
-                    .to_string(),
-                parameters: serde_json::json!({
-                    "type": "object",
-                    "properties": {
-                        "path": {
-                            "type": "string",
-                            "description": "The path to the file to write"
-                        },
-                        "content": {
-                            "type": "string",
-                            "description": "The content to write to the file"
-                        }
-                    },
-                    "required": ["path", "content"]
-                }),
-            },
-        },
         Tool {
             tool_type: "function".to_string(),
             function: FunctionDefinition {