use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use serde_json::Value;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum AITaskType {
AnalyzeContribution,
GenerateCode,
ValidateCode,
RunTests,
PrepareReview,
MergeChanges,
}
impl std::fmt::Display for AITaskType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
AITaskType::AnalyzeContribution => write!(f, "analyze_contribution"),
AITaskType::GenerateCode => write!(f, "generate_code"),
AITaskType::ValidateCode => write!(f, "validate_code"),
AITaskType::RunTests => write!(f, "run_tests"),
AITaskType::PrepareReview => write!(f, "prepare_review"),
AITaskType::MergeChanges => write!(f, "merge_changes"),
}
}
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
#[serde(rename_all = "lowercase")]
pub enum AITaskStatus {
#[default]
Pending,
Running,
Completed,
Failed,
Cancelled,
}
impl std::fmt::Display for AITaskStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
AITaskStatus::Pending => write!(f, "pending"),
AITaskStatus::Running => write!(f, "running"),
AITaskStatus::Completed => write!(f, "completed"),
AITaskStatus::Failed => write!(f, "failed"),
AITaskStatus::Cancelled => write!(f, "cancelled"),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AITask {
#[serde(rename = "_key")]
pub id: String,
pub contribution_id: String,
pub task_type: AITaskType,
pub status: AITaskStatus,
#[serde(default)]
pub priority: i32,
pub created_at: DateTime<Utc>,
#[serde(skip_serializing_if = "Option::is_none")]
pub started_at: Option<DateTime<Utc>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub completed_at: Option<DateTime<Utc>>,
#[serde(default)]
pub retry_count: u32,
#[serde(default = "default_max_retries")]
pub max_retries: u32,
#[serde(skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
#[serde(default)]
pub input: Option<Value>,
#[serde(default)]
pub output: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub agent_id: Option<String>,
}
fn default_max_retries() -> u32 {
3
}
impl AITask {
pub fn new(contribution_id: String, task_type: AITaskType, priority: i32) -> Self {
Self {
id: uuid::Uuid::new_v4().to_string(),
contribution_id,
task_type,
status: AITaskStatus::Pending,
priority,
created_at: Utc::now(),
started_at: None,
completed_at: None,
retry_count: 0,
max_retries: default_max_retries(),
error: None,
input: None,
output: None,
agent_id: None,
}
}
pub fn analyze(contribution_id: String, priority: i32) -> Self {
Self::new(contribution_id, AITaskType::AnalyzeContribution, priority)
}
pub fn start(&mut self, agent_id: String) {
self.status = AITaskStatus::Running;
self.started_at = Some(Utc::now());
self.agent_id = Some(agent_id);
}
pub fn complete(&mut self, output: Option<Value>) {
self.status = AITaskStatus::Completed;
self.completed_at = Some(Utc::now());
self.output = output;
}
pub fn fail(&mut self, error: String) {
self.retry_count += 1;
self.error = Some(error);
if self.retry_count >= self.max_retries {
self.status = AITaskStatus::Failed;
self.completed_at = Some(Utc::now());
} else {
self.status = AITaskStatus::Pending;
self.started_at = None;
self.agent_id = None;
}
}
pub fn can_retry(&self) -> bool {
self.retry_count < self.max_retries
}
}
#[derive(Debug, Serialize)]
pub struct ListAITasksResponse {
pub tasks: Vec<AITask>,
pub total: usize,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_ai_task_creation() {
let task = AITask::analyze("contrib-123".to_string(), 5);
assert_eq!(task.contribution_id, "contrib-123");
assert_eq!(task.task_type, AITaskType::AnalyzeContribution);
assert_eq!(task.status, AITaskStatus::Pending);
assert_eq!(task.priority, 5);
assert!(!task.id.is_empty());
}
#[test]
fn test_ai_task_lifecycle() {
let mut task = AITask::new("contrib-456".to_string(), AITaskType::GenerateCode, 0);
assert_eq!(task.status, AITaskStatus::Pending);
assert!(task.started_at.is_none());
task.start("agent-001".to_string());
assert_eq!(task.status, AITaskStatus::Running);
assert!(task.started_at.is_some());
assert_eq!(task.agent_id, Some("agent-001".to_string()));
task.complete(Some(serde_json::json!({"files_generated": 3})));
assert_eq!(task.status, AITaskStatus::Completed);
assert!(task.completed_at.is_some());
assert!(task.output.is_some());
}
#[test]
fn test_ai_task_retry() {
let mut task = AITask::new("contrib-789".to_string(), AITaskType::ValidateCode, 0);
task.max_retries = 3;
task.fail("Network error".to_string());
assert_eq!(task.status, AITaskStatus::Pending);
assert_eq!(task.retry_count, 1);
assert!(task.can_retry());
task.fail("Timeout".to_string());
assert_eq!(task.status, AITaskStatus::Pending);
assert_eq!(task.retry_count, 2);
assert!(task.can_retry());
task.fail("Service unavailable".to_string());
assert_eq!(task.status, AITaskStatus::Failed);
assert_eq!(task.retry_count, 3);
assert!(!task.can_retry());
}
#[test]
fn test_ai_task_type_serialization() {
let json = serde_json::to_string(&AITaskType::AnalyzeContribution).unwrap();
assert_eq!(json, "\"analyze_contribution\"");
let parsed: AITaskType = serde_json::from_str("\"generate_code\"").unwrap();
assert_eq!(parsed, AITaskType::GenerateCode);
}
}