gitdataai/libs/agent/task/service.rs

//! Task service for creating, tracking, and executing agent tasks.
//!
//! All methods are async and interact with the database directly.
//! Execution of the task logic (running the ReAct loop, etc.) is delegated
//! to the caller — this service only manages task lifecycle and state.

use db::database::AppDatabase;
use models::agent_task::{
    ActiveModel, AgentType, Column as C, Entity, Model, TaskStatus,
};
use sea_orm::{
    entity::EntityTrait, query::{QueryFilter, QueryOrder, QuerySelect}, ActiveModelTrait,
    ColumnTrait, DbErr,
};

/// Service for managing agent tasks (root tasks and sub-tasks).
#[derive(Clone)]
pub struct TaskService {
    db: AppDatabase,
}

impl TaskService {
    pub fn new(db: AppDatabase) -> Self {
        Self { db }
    }

    /// Create a new task (root or sub-task) with status = pending.
    pub async fn create(
        &self,
        project_uuid: impl Into<uuid::Uuid>,
        input: impl Into<String>,
        agent_type: AgentType,
    ) -> Result<Model, DbErr> {
        self.create_with_parent(project_uuid, None, input, agent_type, None).await
    }

    /// Create a new sub-task with a parent reference.
    pub async fn create_subtask(
        &self,
        project_uuid: impl Into<uuid::Uuid>,
        parent_id: i64,
        input: impl Into<String>,
        agent_type: AgentType,
        title: Option<String>,
    ) -> Result<Model, DbErr> {
        self.create_with_parent(project_uuid, Some(parent_id), input, agent_type, title)
            .await
    }

    async fn create_with_parent(
        &self,
        project_uuid: impl Into<uuid::Uuid>,
        parent_id: Option<i64>,
        input: impl Into<String>,
        agent_type: AgentType,
        title: Option<String>,
    ) -> Result<Model, DbErr> {
        let model = ActiveModel {
            project_uuid: sea_orm::Set(project_uuid.into()),
            parent_id: sea_orm::Set(parent_id),
            agent_type: sea_orm::Set(agent_type),
            status: sea_orm::Set(TaskStatus::Pending),
            title: sea_orm::Set(title),
            input: sea_orm::Set(input.into()),
            ..Default::default()
        };
        model.insert(&self.db).await
    }

    /// 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());
        active.update(&self.db).await
    }

    /// 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 mut active: ActiveModel = model.into();
        active.progress = sea_orm::Set(Some(progress.into()));
        active.updated_at = sea_orm::Set(chrono::Utc::now().into());
        active.update(&self.db).await?;
        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);
        active.output = sea_orm::Set(Some(output.into()));
        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
    }

    /// 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);
        active.error = sea_orm::Set(Some(error.into()));
        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
    }

    /// Get a task by ID.
    pub async fn get(&self, task_id: i64) -> Result<Option<Model>, DbErr> {
        Entity::find_by_id(task_id).one(&self.db).await
    }

    /// List all sub-tasks for a parent task.
    pub async fn children(&self, parent_id: i64) -> Result<Vec<Model>, DbErr> {
        Entity::find()
            .filter(C::ParentId.eq(parent_id))
            .order_by_asc(C::CreatedAt)
            .all(&self.db)
            .await
    }

    /// List all active (non-terminal) tasks for a project.
    pub async fn active_tasks(&self, project_uuid: impl Into<uuid::Uuid>) -> Result<Vec<Model>, DbErr> {
        let uuid: uuid::Uuid = project_uuid.into();
        Entity::find()
            .filter(C::ProjectUuid.eq(uuid))
            .filter(C::Status.is_in([TaskStatus::Pending, TaskStatus::Running]))
            .order_by_desc(C::CreatedAt)
            .all(&self.db)
            .await
    }

    /// List all tasks (root only) for a project.
    pub async fn list(
        &self,
        project_uuid: impl Into<uuid::Uuid>,
        limit: u64,
    ) -> Result<Vec<Model>, DbErr> {
        let uuid: uuid::Uuid = project_uuid.into();
        Entity::find()
            .filter(C::ProjectUuid.eq(uuid))
            .filter(C::ParentId.is_null())
            .order_by_desc(C::CreatedAt)
            .limit(limit)
            .all(&self.db)
            .await
    }

    /// Delete a task and all its sub-tasks recursively.
    /// Only allows deletion of root tasks.
    pub async fn delete(&self, task_id: i64) -> Result<(), DbErr> {
        self.delete_recursive(task_id).await
    }

    async fn delete_recursive(&self, task_id: i64) -> Result<(), DbErr> {
        // Collect all task IDs to delete using an explicit stack (avoiding async recursion).
        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;
        }

        for task_id in stack {
            let model = Entity::find_by_id(task_id).one(&self.db).await?;
            if let Some(m) = model {
                let active: ActiveModel = m.into();
                active.delete(&self.db).await?;
            }
        }
        Ok(())
    }

    /// Check if all sub-tasks of a given parent are done.
    pub async fn are_children_done(&self, parent_id: i64) -> Result<bool, DbErr> {
        let children = self.children(parent_id).await?;
        if children.is_empty() {
            return Ok(true);
        }
        Ok(children.iter().all(|c| c.is_done()))
    }
}