gitdataai/libs/agent/task/lifecycle.rs

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<Model, 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 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<String>,
    ) -> 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<String>) -> Result<Model, 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 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<String>) -> Result<Model, 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 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<Model, DbErr> {
        // 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<Model, 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()))?;

        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<Model, 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()))?;

        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<Model, 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()))?;

        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
    }
}