gitdataai/libs/models/agent_task/mod.rs

//! Agent task model — sub-agents and sub-tasks as a unified concept.
//!
//! An `agent_task` represents either:
//! - A **root task** (parent_id = NULL): initiated by a user or system event.
//!   The parent agent (Supervisor) spawns sub-tasks and coordinates their results.
//! - A **sub-task** (parent_id = set): a unit of work executed by a sub-agent.
//!
//! Status lifecycle: `pending` → `running` → `done` | `failed`
//!
//! Sub-agents are represented as `agent_task` records with a parent reference,
//! allowing hierarchical task trees and result aggregation.

use crate::{DateTimeUtc, IssueId, ProjectId, UserId};
use sea_orm::entity::prelude::*;
use serde::{Deserialize, Serialize};

/// Agent task type — the kind of agent that executes this task.
#[derive(
    Clone, Debug, PartialEq, Eq, EnumIter, Serialize, Deserialize, sea_orm::DeriveActiveEnum,
)]
#[sea_orm(rs_type = "String", db_type = "String(StringLen::None)")]
pub enum AgentType {
    #[sea_orm(string_value = "React")]
    React,
    #[sea_orm(string_value = "Chat")]
    Chat,
}

impl Default for AgentType {
    fn default() -> Self {
        AgentType::React
    }
}

impl std::fmt::Display for AgentType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            AgentType::React => write!(f, "React"),
            AgentType::Chat => write!(f, "Chat"),
        }
    }
}

/// Task status lifecycle.
#[derive(
    Clone, Debug, PartialEq, Eq, EnumIter, Serialize, Deserialize, sea_orm::DeriveActiveEnum,
)]
#[sea_orm(rs_type = "String", db_type = "String(StringLen::None)")]
pub enum TaskStatus {
    #[sea_orm(string_value = "Pending")]
    Pending,
    #[sea_orm(string_value = "Running")]
    Running,
    #[sea_orm(string_value = "Paused")]
    Paused,
    #[sea_orm(string_value = "Done")]
    Done,
    #[sea_orm(string_value = "Failed")]
    Failed,
    #[sea_orm(string_value = "Cancelled")]
    Cancelled,
}

impl Default for TaskStatus {
    fn default() -> Self {
        TaskStatus::Pending
    }
}

impl std::fmt::Display for TaskStatus {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            TaskStatus::Pending => write!(f, "Pending"),
            TaskStatus::Running => write!(f, "Running"),
            TaskStatus::Paused => write!(f, "Paused"),
            TaskStatus::Done => write!(f, "Done"),
            TaskStatus::Failed => write!(f, "Failed"),
            TaskStatus::Cancelled => write!(f, "Cancelled"),
        }
    }
}

/// Agent task record — represents both root tasks and sub-tasks.
#[derive(Clone, Debug, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "agent_task")]
pub struct Model {
    #[sea_orm(primary_key)]
    pub id: i64,

    /// Project this task belongs to.
    pub project_uuid: ProjectId,

    /// Parent task (NULL for root tasks, set for sub-tasks).
    #[sea_orm(nullable)]
    pub parent_id: Option<i64>,

    /// Issue this task is bound to (optional).
    #[sea_orm(nullable)]
    pub issue_id: Option<IssueId>,

    /// Agent type that executes this task.
    #[sea_orm(column_type = "String(StringLen::None)", default = "React")]
    pub agent_type: AgentType,

    /// Current task status.
    #[sea_orm(column_type = "String(StringLen::None)", default = "Pending")]
    pub status: TaskStatus,

    /// Human-readable task title / goal description.
    #[sea_orm(nullable)]
    pub title: Option<String>,

    /// Task input — the prompt or goal text.
    pub input: String,

    /// Task output — populated when status = done.
    #[sea_orm(nullable)]
    pub output: Option<String>,

    /// Error message — populated when status = failed.
    #[sea_orm(nullable)]
    pub error: Option<String>,

    /// User who initiated this task.
    #[sea_orm(nullable)]
    pub created_by: Option<UserId>,

    pub created_at: DateTimeUtc,
    pub updated_at: DateTimeUtc,

    /// When execution started (status → running).
    #[sea_orm(nullable)]
    pub started_at: Option<DateTimeUtc>,

    /// When execution completed (status → done | failed).
    #[sea_orm(nullable)]
    pub done_at: Option<DateTimeUtc>,

    /// Current progress description (e.g., "step 2/5: analyzing code").
    #[sea_orm(nullable)]
    pub progress: Option<String>,

    /// Number of times this task has been retried.
    #[sea_orm(nullable)]
    pub retry_count: Option<i32>,
}

#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {
    #[sea_orm(belongs_to = "Entity", from = "Column::ParentId", to = "Column::Id")]
    ParentTask,
}

impl ActiveModelBehavior for ActiveModel {}

impl Model {
    pub fn is_root_task(&self) -> bool {
        self.parent_id.is_none()
    }

    pub fn is_done(&self) -> bool {
        matches!(
            self.status,
            TaskStatus::Done | TaskStatus::Failed | TaskStatus::Cancelled
        )
    }

    pub fn is_running(&self) -> bool {
        matches!(
            self.status,
            TaskStatus::Running | TaskStatus::Pending | TaskStatus::Paused
        )
    }
}