use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use uuid::Uuid; use crate::agents::AgentType; use crate::enums::{TaskResultType, TaskStatus, TestStrategy}; use crate::todo::ToDo; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Task { pub id: Uuid, pub todo_id: Uuid, pub description: String, pub status: TaskStatus, pub assigned_krill: Option, pub assigned_pod: Option, pub result: Option, pub logs: Vec, pub started_at: Option>, pub completed_at: Option>, pub payload: Option, } impl Default for Task { fn default() -> Self { Self { id: Uuid::new_v4(), todo_id: Uuid::nil(), description: String::new(), status: TaskStatus::default(), assigned_krill: None, assigned_pod: None, result: None, logs: Vec::new(), started_at: None, completed_at: None, payload: None, } } } impl Task { pub fn new(todo_id: Uuid, description: String) -> Self { Self { id: Uuid::new_v4(), todo_id, description, status: TaskStatus::default(), assigned_krill: None, assigned_pod: None, result: None, logs: Vec::new(), started_at: None, completed_at: None, payload: None, } } pub fn with_payload(mut self, payload: &T) -> Self { self.payload = serde_json::to_value(payload).ok(); self } pub fn start(&mut self) -> Result<(), String> { match self.status { TaskStatus::Pending => { self.status = TaskStatus::Running; self.started_at = Some(Utc::now()); Ok(()) } _ => Err(format!("cannot start task from status: {:?}", self.status)), } } pub fn complete(&mut self, result: TaskResult) -> Result<(), String> { match self.status { TaskStatus::Running => { self.status = TaskStatus::Completed; self.completed_at = Some(Utc::now()); self.result = Some(result); Ok(()) } _ => Err(format!( "cannot complete task from status: {:?}", self.status )), } } pub fn fail(&mut self, error: String) -> Result<(), String> { match self.status { TaskStatus::Running => { self.status = TaskStatus::Failed; self.completed_at = Some(Utc::now()); self.result = Some(TaskResult::Error(error)); Ok(()) } _ => Err(format!("cannot fail task from status: {:?}", self.status)), } } pub fn retry(&mut self) -> Result<(), String> { match self.status { TaskStatus::Failed | TaskStatus::Pending => { self.status = TaskStatus::Pending; self.started_at = None; self.completed_at = None; self.result = None; self.logs.clear(); Ok(()) } _ => Err(format!("cannot retry task from status: {:?}", self.status)), } } pub fn duration_ms(&self) -> Option { match (self.started_at, self.completed_at) { (Some(start), Some(end)) => Some((end - start).num_milliseconds()), _ => None, } } pub fn add_log(&mut self, log: String) { self.logs.push(log); } } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub enum TaskResult { Success(String), Error(String), Split(Vec), } impl TaskResult { pub fn result_type(&self) -> TaskResultType { match self { TaskResult::Success(_) => TaskResultType::Success, TaskResult::Error(_) => TaskResultType::Error, TaskResult::Split(_) => TaskResultType::Split, } } } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct TestCase { pub name: String, pub command: String, pub expected_output: Option, #[serde(default)] pub test_strategy: Option, #[serde(default)] pub suggested_test_files: Vec, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ResourceLimits { pub cpu_cores: f32, pub memory_mb: u64, pub gpu_memory_mb: Option, pub timeout_seconds: u64, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct SessionBudget { pub max_tool_calls: u32, pub max_tokens: u32, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct FileRange { pub path: String, pub line_start: u32, pub line_end: u32, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ConnectorConfig { pub api_key: String, pub base_url: String, pub model: String, pub max_tokens: Option, pub temperature: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct AgentTaskSpec { pub task_id: Uuid, pub title: String, pub description: String, pub board_snapshot: serde_json::Value, pub assembled_context_id: Uuid, pub selected_files: Vec, pub entry_points: Vec, pub tests: Option>, pub success_criteria: Option, pub token_budget: usize, pub resource_limits: ResourceLimits, pub parent_task_id: Option, pub pre_approved_plan_id: Option, pub execution_token: Option, pub session_budget: Option, pub agent_type: Option, #[serde(default)] pub suggested_test_files: Vec, #[serde(default)] pub project_id: Option, #[serde(default)] pub connector_config: Option, #[serde(default)] pub tools_json: Option, /* Overrides the agent-type default system prompt when set. Takes highest precedence in the resolution chain (above DB overrides and defaults). */ #[serde(default)] pub system_prompt: Option, } impl AgentTaskSpec { pub fn new(task_id: Uuid, title: String, assembled_context_id: Uuid) -> Self { Self { task_id, title, description: String::new(), board_snapshot: serde_json::json!({}), assembled_context_id, selected_files: Vec::new(), entry_points: Vec::new(), tests: None, success_criteria: None, token_budget: 1000000, // Default 1M tokens resource_limits: ResourceLimits { cpu_cores: 1.0, memory_mb: 2048, gpu_memory_mb: None, timeout_seconds: 3600, // 1 hour }, parent_task_id: None, pre_approved_plan_id: None, execution_token: None, session_budget: None, agent_type: None, suggested_test_files: Vec::new(), project_id: None, connector_config: None, tools_json: None, system_prompt: None, } } } #[cfg(test)] mod tests { use super::*; #[test] fn test_task_new() { let todo_id = Uuid::new_v4(); let task = Task::new(todo_id, "Test task description".to_string()); assert!(!task.id.is_nil()); assert_eq!(task.todo_id, todo_id); assert_eq!(task.description, "Test task description"); assert_eq!(task.status, TaskStatus::Pending); } #[test] fn test_task_start() { let mut task = Task::default(); assert!(task.start().is_ok()); assert_eq!(task.status, TaskStatus::Running); assert!(task.started_at.is_some()); } #[test] fn test_task_complete() { let mut task = Task::default(); task.status = TaskStatus::Running; let result = TaskResult::Success("output".to_string()); assert!(task.complete(result.clone()).is_ok()); assert_eq!(task.status, TaskStatus::Completed); assert!(task.completed_at.is_some()); assert_eq!(task.result, Some(result)); } #[test] fn test_task_duration_ms() { let mut task = Task::default(); task.start().expect("task should start"); std::thread::sleep(std::time::Duration::from_millis(10)); task.complete(TaskResult::Success("done".to_string())) .expect("task should complete"); let duration = task.duration_ms(); assert!(duration.is_some()); assert!(duration.unwrap() >= 10); } #[test] fn test_task_retry() { let mut task = Task::default(); task.status = TaskStatus::Failed; task.result = Some(TaskResult::Error("failed".to_string())); assert!(task.retry().is_ok()); assert_eq!(task.status, TaskStatus::Pending); assert!(task.result.is_none()); } #[test] fn test_task_default_values() { let task = Task::default(); assert_eq!(task.status, TaskStatus::Pending); assert!(task.result.is_none()); assert!(task.logs.is_empty()); assert!(task.started_at.is_none()); assert!(task.completed_at.is_none()); } #[test] fn test_task_serialization_roundtrip() { let original = Task::new(Uuid::new_v4(), "Test task".to_string()); let serialized = serde_json::to_string(&original).unwrap(); let deserialized: Task = serde_json::from_str(&serialized).unwrap(); assert_eq!(original.id, deserialized.id); assert_eq!(original.todo_id, deserialized.todo_id); assert_eq!(original.description, deserialized.description); assert_eq!(original.status, deserialized.status); } #[test] fn test_task_fail() { let mut task = Task::default(); task.status = TaskStatus::Running; assert!(task.fail("error message".to_string()).is_ok()); assert_eq!(task.status, TaskStatus::Failed); assert!(task.result.is_some()); } #[test] fn test_task_invalid_transition() { let mut task = Task::default(); // Cannot complete a pending task directly let result = TaskResult::Success("done".to_string()); assert!(task.complete(result).is_err()); } }