use chrono::Utc; use sea_orm::{ActiveModelTrait, DatabaseConnection, EntityTrait, Set}; use serde::{Deserialize, Serialize}; use std::sync::Arc; use uuid::Uuid; use crate::config::Config; use crate::domain::agent::Agent; use crate::entities::task::Entity as Task; use crate::entities::task_run::{self, Entity as TaskRun}; use crate::scheduler::Scheduler; #[derive(Deserialize)] pub struct CreateTaskRequest { pub goal: String, pub cron: Option, } #[derive(Deserialize)] pub struct UpdateTaskRequest { pub goal: String, pub cron: Option, } #[derive(Serialize, Clone, Debug)] pub struct TaskResponse { pub id: Uuid, pub goal: String, pub cron: Option, pub created_at: chrono::DateTime, pub runs: Vec, } #[derive(Serialize, Clone, Debug)] pub struct TaskRunResponse { pub id: Uuid, pub status: String, pub logs: String, pub answer: Option, pub created_at: chrono::DateTime, } #[derive(Serialize, Clone, Debug)] pub struct RecentRunResponse { pub id: Uuid, pub task_id: Uuid, pub goal: String, pub status: String, pub created_at: chrono::DateTime, } use crate::error::AppResult; pub async fn execute_agent_run( db: &DatabaseConnection, _scheduler: &Arc, config: &Arc, task_id: Uuid, goal: String, ) -> AppResult { let run_id = Uuid::new_v4(); tracing::info!(%task_id, %run_id, "Starting agent execution run"); 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(db) .await .map_err(crate::error::AppError::Database)?; let _ = _scheduler .tx .send(crate::server::notifications::WsEvent::RunStarted( RecentRunResponse { id: run_id, task_id, goal: goal.clone(), status: "running".to_string(), created_at: Utc::now().into(), }, )); let mut agent = Agent::new( db.clone(), config.zen_api_key.clone(), config.tavily_api_key.clone(), goal.clone(), )?; let (logs, answer, status) = match agent.run(config).await { Ok((logs, answer)) => { tracing::info!(%task_id, %run_id, "Agent execution completed successfully"); (logs, answer, "completed".to_string()) } Err(e) => { tracing::error!(%task_id, %run_id, error = %e, "Agent execution failed"); ( format!("Execution failed: {}", e), None, "failed".to_string(), ) } }; let run: task_run::ActiveModel = TaskRun::find_by_id(run_id) .one(db) .await .map_err(crate::error::AppError::Database)? .ok_or_else(|| crate::error::AppError::NotFound("Run not found after insert".into()))? .into(); let mut run = run; run.logs = Set(logs.clone()); run.answer = Set(answer.clone()); run.status = Set(status); run.update(db) .await .map_err(crate::error::AppError::Database)?; let task_response = get_task_inner(task_id, db).await?; let _ = _scheduler .tx .send(crate::server::notifications::WsEvent::RunFinished( task_response.clone(), )); Ok(task_response) } pub async fn get_task_inner(id: Uuid, db: &DatabaseConnection) -> AppResult { let results = Task::find_by_id(id) .find_with_related(TaskRun) .all(db) .await .map_err(crate::error::AppError::Database)?; let (t, mut runs) = results .into_iter() .next() .ok_or_else(|| crate::error::AppError::NotFound("Task not found".into()))?; runs.sort_by(|a, b| b.created_at.cmp(&a.created_at)); Ok(TaskResponse { id: t.id, goal: t.goal, 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(), }) }