//! 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 models::IssueId; use sea_orm::{ entity::EntityTrait, query::{QueryFilter, QueryOrder, QuerySelect}, ActiveModelTrait, ColumnTrait, DbErr, }; use serde::Serialize; use std::sync::Arc; /// Event payload published to WebSocket clients via Redis Pub/Sub. #[derive(Debug, Clone, Serialize)] pub struct TaskEvent { pub task_id: i64, pub project_id: uuid::Uuid, pub parent_id: Option, pub event: String, pub message: Option, pub output: Option, pub error: Option, pub status: String, } impl TaskEvent { pub fn started(task_id: i64, project_id: uuid::Uuid, parent_id: Option) -> Self { Self { task_id, project_id, parent_id, event: "started".to_string(), message: None, output: None, error: None, status: TaskStatus::Running.to_string(), } } pub fn progress( task_id: i64, project_id: uuid::Uuid, parent_id: Option, msg: String, ) -> Self { Self { task_id, project_id, parent_id, event: "progress".to_string(), message: Some(msg), output: None, error: None, status: TaskStatus::Running.to_string(), } } pub fn completed( task_id: i64, project_id: uuid::Uuid, parent_id: Option, output: String, ) -> Self { Self { task_id, project_id, parent_id, event: "done".to_string(), message: None, output: Some(output), error: None, status: TaskStatus::Done.to_string(), } } pub fn failed( task_id: i64, project_id: uuid::Uuid, parent_id: Option, error: String, ) -> Self { Self { task_id, project_id, parent_id, event: "failed".to_string(), message: None, output: None, error: Some(error), status: TaskStatus::Failed.to_string(), } } pub fn cancelled(task_id: i64, project_id: uuid::Uuid, parent_id: Option) -> Self { Self { task_id, project_id, parent_id, event: "cancelled".to_string(), message: None, output: None, error: None, status: TaskStatus::Cancelled.to_string(), } } } /// Helper trait for publishing task lifecycle events via Redis Pub/Sub. /// /// Callers inject a suitable `publish_fn` at construction time via /// `TaskEvents::new(...)`. If no publisher is supplied events are silently /// dropped (graceful degradation on startup). pub trait TaskEventPublisher: Send + Sync { fn publish(&self, project_id: uuid::Uuid, event: TaskEvent); } /// No-op publisher used when no Redis Pub/Sub connection is available. #[derive(Clone, Default)] pub struct NoOpPublisher; impl TaskEventPublisher for NoOpPublisher { fn publish(&self, _: uuid::Uuid, _: TaskEvent) {} } #[derive(Clone)] pub struct TaskEvents { publisher: Arc, } impl TaskEvents { pub fn new(publisher: impl TaskEventPublisher + 'static) -> Self { Self { publisher: Arc::new(publisher), } } pub fn noop() -> Self { Self::new(NoOpPublisher) } fn emit(&self, task: &Model, event: TaskEvent) { self.publisher.publish(task.project_uuid, event); } pub fn emit_started(&self, task: &Model) { self.emit( task, TaskEvent::started(task.id, task.project_uuid, task.parent_id), ); } pub fn emit_progress(&self, task: &Model, msg: String) { self.emit( task, TaskEvent::progress(task.id, task.project_uuid, task.parent_id, msg), ); } pub fn emit_completed(&self, task: &Model, output: String) { self.emit( task, TaskEvent::completed(task.id, task.project_uuid, task.parent_id, output), ); } pub fn emit_failed(&self, task: &Model, error: String) { self.emit( task, TaskEvent::failed(task.id, task.project_uuid, task.parent_id, error), ); } pub fn emit_cancelled(&self, task: &Model) { self.emit( task, TaskEvent::cancelled(task.id, task.project_uuid, task.parent_id), ); } } /// Builder for TaskService so that the events publisher can be set independently /// of the database connection. #[derive(Clone, Default)] pub struct TaskServiceBuilder { events: Option, } impl TaskServiceBuilder { pub fn with_events(mut self, events: TaskEvents) -> Self { self.events = Some(events); self } pub async fn build(self, db: AppDatabase) -> TaskService { TaskService { db, events: self.events.unwrap_or_else(TaskEvents::noop), } } } /// Service for managing agent tasks (root tasks and sub-tasks). #[derive(Clone)] pub struct TaskService { db: AppDatabase, events: TaskEvents, } impl TaskService { pub fn new(db: AppDatabase) -> Self { Self { db, events: TaskEvents::noop(), } } pub fn with_events(db: AppDatabase, events: TaskEvents) -> Self { Self { db, events } } /// Create a new task (root or sub-task) with status = pending. pub async fn create( &self, project_uuid: impl Into, input: impl Into, agent_type: AgentType, ) -> Result { self.create_with_parent(project_uuid, None, input, agent_type, None, None) .await } /// Create a new task bound to an issue. pub async fn create_for_issue( &self, project_uuid: impl Into, issue_id: IssueId, input: impl Into, agent_type: AgentType, ) -> Result { self.create_with_parent(project_uuid, None, input, agent_type, None, Some(issue_id)) .await } /// Create a new sub-task with a parent reference. pub async fn create_subtask( &self, project_uuid: impl Into, parent_id: i64, input: impl Into, agent_type: AgentType, title: Option, ) -> Result { self.create_with_parent( project_uuid, Some(parent_id), input, agent_type, title, None, ) .await } async fn create_with_parent( &self, project_uuid: impl Into, parent_id: Option, input: impl Into, agent_type: AgentType, title: Option, issue_id: Option, ) -> Result { let model = ActiveModel { project_uuid: sea_orm::Set(project_uuid.into()), parent_id: sea_orm::Set(parent_id), issue_id: sea_orm::Set(issue_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 { 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 } /// Propagate child task status up the tree. /// /// When a child task reaches a terminal state, checks whether all its /// siblings are also terminal. If so, marks the parent as failed so that /// a stuck parent is never left in the `Running` state. pub async fn propagate_to_parent(&self, task_id: i64) -> Result, DbErr> { let model = self .get(task_id) .await? .ok_or_else(|| DbErr::RecordNotFound("agent_task not found".to_string()))?; let Some(parent_id) = model.parent_id else { return Ok(None); }; let siblings = self.children(parent_id).await?; if siblings.iter().all(|s| s.is_done()) { let parent = self.get(parent_id).await?.ok_or_else(|| { DbErr::RecordNotFound(format!("parent task {} not found", parent_id)) })?; if parent.is_running() { let mut active: ActiveModel = parent.into(); active.status = sea_orm::Set(TaskStatus::Failed); active.error = sea_orm::Set(Some("All sub-tasks failed or were cancelled".to_string())); 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?; return Ok(Some(updated)); } } Ok(None) } /// Get a task by ID. pub async fn get(&self, task_id: i64) -> Result, 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, 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, ) -> Result, DbErr> { let uuid: uuid::Uuid = project_uuid.into(); Entity::find() .filter(C::ProjectUuid.eq(uuid)) .filter(C::Status.is_in([TaskStatus::Pending, TaskStatus::Running, TaskStatus::Paused])) .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, limit: u64, ) -> Result, 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 in a terminal state. /// Returns true if there are no children (empty tree counts as done). pub async fn are_children_done(&self, parent_id: i64) -> Result { let children = self.children(parent_id).await?; Ok(children.is_empty() || children.iter().all(|c| c.is_done())) } }