use models::agent_task::{ActiveModel, Column as C, Entity, Model, TaskStatus}; use sea_orm::{ActiveModelTrait, ColumnTrait, DbErr, EntityTrait, QueryFilter}; pub struct TaskLifecycle; impl super::TaskService { /// Mark a task as running and record the start time. pub async fn start(&self, task_id: i64) -> Result { let model = Entity::find_by_id(task_id).one(self.db()).await?; let model = model.ok_or_else(|| DbErr::RecordNotFound("agent_task not found".to_string()))?; let mut active: ActiveModel = model.into(); active.status = sea_orm::Set(TaskStatus::Running); active.started_at = sea_orm::Set(Some(chrono::Utc::now().into())); active.updated_at = sea_orm::Set(chrono::Utc::now().into()); let updated = active.update(self.db()).await?; self.events().emit_started(&updated); Ok(updated) } /// Update progress text (e.g., "step 2/5: analyzing PR"). pub async fn update_progress( &self, task_id: i64, progress: impl Into, ) -> Result<(), DbErr> { let model = Entity::find_by_id(task_id).one(self.db()).await?; let model = model.ok_or_else(|| DbErr::RecordNotFound("agent_task not found".to_string()))?; let progress_str = progress.into(); let mut active: ActiveModel = model.into(); active.progress = sea_orm::Set(Some(progress_str.clone())); active.updated_at = sea_orm::Set(chrono::Utc::now().into()); let updated = active.update(self.db()).await?; self.events().emit_progress(&updated, progress_str); Ok(()) } /// Mark a task as completed with the output text. pub async fn complete(&self, task_id: i64, output: impl Into) -> Result { let model = Entity::find_by_id(task_id).one(self.db()).await?; let model = model.ok_or_else(|| DbErr::RecordNotFound("agent_task not found".to_string()))?; let mut active: ActiveModel = model.into(); active.status = sea_orm::Set(TaskStatus::Done); let out = output.into(); active.output = sea_orm::Set(Some(out.clone())); active.done_at = sea_orm::Set(Some(chrono::Utc::now().into())); active.updated_at = sea_orm::Set(chrono::Utc::now().into()); let updated = active.update(self.db()).await?; self.events().emit_completed(&updated, out); Ok(updated) } /// Mark a task as failed with an error message. pub async fn fail(&self, task_id: i64, error: impl Into) -> Result { let model = Entity::find_by_id(task_id).one(self.db()).await?; let model = model.ok_or_else(|| DbErr::RecordNotFound("agent_task not found".to_string()))?; let mut active: ActiveModel = model.into(); active.status = sea_orm::Set(TaskStatus::Failed); let err = error.into(); active.error = sea_orm::Set(Some(err.clone())); active.done_at = sea_orm::Set(Some(chrono::Utc::now().into())); active.updated_at = sea_orm::Set(chrono::Utc::now().into()); let updated = active.update(self.db()).await?; self.events().emit_failed(&updated, err); Ok(updated) } /// Propagate child task status up the tree. /// /// Only allows cancelling tasks that are not yet in a terminal state /// (Pending / Running / Paused). /// /// Cancelled children are marked done so that `are_children_done()` returns /// true for the parent after cancellation. pub async fn cancel(&self, task_id: i64) -> Result { // Collect all task IDs (parent + descendants) using an explicit stack. let mut stack = vec![task_id]; let mut idx = 0; while idx < stack.len() { let current = stack[idx]; let children = Entity::find() .filter(C::ParentId.eq(current)) .all(self.db()) .await?; for child in children { stack.push(child.id); } idx += 1; } // Mark every collected task as cancelled (terminal state). for id in &stack { let model = Entity::find_by_id(*id).one(self.db()).await?; if let Some(m) = model { if !m.is_done() { let mut active: ActiveModel = m.into(); active.status = sea_orm::Set(TaskStatus::Cancelled); active.done_at = sea_orm::Set(Some(chrono::Utc::now().into())); active.updated_at = sea_orm::Set(chrono::Utc::now().into()); active.update(self.db()).await?; } } } let final_model = Entity::find_by_id(task_id) .one(self.db()) .await? .ok_or_else(|| DbErr::RecordNotFound("agent_task not found".to_string()))?; self.events().emit_cancelled(&final_model); Ok(final_model) } /// Pause a running or pending task. /// /// Pausing a task that is not Pending/Running is a no-op that returns /// the current model (same behaviour as `start` on an already-running task). pub async fn pause(&self, task_id: i64) -> Result { let model = Entity::find_by_id(task_id).one(self.db()).await?; let model = model.ok_or_else(|| DbErr::RecordNotFound("agent_task not found".to_string()))?; if !model.is_running() { // Already in a terminal or paused state — return unchanged. return Ok(model); } let mut active: ActiveModel = model.into(); active.status = sea_orm::Set(TaskStatus::Paused); active.updated_at = sea_orm::Set(chrono::Utc::now().into()); active.update(self.db()).await } /// Resume a paused task back to Running. /// /// Returns an error if the task is not currently Paused. pub async fn resume(&self, task_id: i64) -> Result { let model = Entity::find_by_id(task_id).one(self.db()).await?; let model = model.ok_or_else(|| DbErr::RecordNotFound("agent_task not found".to_string()))?; if model.status != TaskStatus::Paused { return Err(DbErr::Custom(format!( "cannot resume task {}: expected status Paused, got {}", task_id, model.status ))); } let mut active: ActiveModel = model.into(); active.status = sea_orm::Set(TaskStatus::Running); active.updated_at = sea_orm::Set(chrono::Utc::now().into()); active.update(self.db()).await } /// Retry a failed or cancelled task by resetting it to Pending. /// /// Clears `output`, `error`, and `done_at`; increments `retry_count`. /// Only tasks in Failed or Cancelled state can be retried. pub async fn retry(&self, task_id: i64) -> Result { let model = Entity::find_by_id(task_id).one(self.db()).await?; let model = model.ok_or_else(|| DbErr::RecordNotFound("agent_task not found".to_string()))?; match model.status { TaskStatus::Failed | TaskStatus::Cancelled | TaskStatus::Done => {} _ => { return Err(DbErr::Custom(format!( "cannot retry task {}: only Failed/Cancelled/Done tasks can be retried (got {})", task_id, model.status ))); } } let retry_count = model.retry_count.map(|c| c + 1).unwrap_or(1); let mut active: ActiveModel = model.into(); active.status = sea_orm::Set(TaskStatus::Pending); active.output = sea_orm::Set(None); active.error = sea_orm::Set(None); active.done_at = sea_orm::Set(None); active.started_at = sea_orm::Set(None); active.retry_count = sea_orm::Set(Some(retry_count)); active.updated_at = sea_orm::Set(chrono::Utc::now().into()); active.update(self.db()).await } }