diff options
Diffstat (limited to 'makima/src/db/models.rs')
| -rw-r--r-- | makima/src/db/models.rs | 201 |
1 files changed, 201 insertions, 0 deletions
diff --git a/makima/src/db/models.rs b/makima/src/db/models.rs index 636d81a..cc30465 100644 --- a/makima/src/db/models.rs +++ b/makima/src/db/models.rs @@ -2339,6 +2339,111 @@ pub struct CheckpointPatchInfo { // Red Team Types // ============================================================================ +// ============================================================================= +// Supervisor State and Heartbeat Types +// ============================================================================= + +/// Supervisor state for contract supervisor tasks. +/// Captures detailed activity state for monitoring and restoration. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, ToSchema)] +#[serde(rename_all = "snake_case")] +pub enum SupervisorStateEnum { + /// Supervisor is starting up + Initializing, + /// Supervisor is idle, waiting for work + Idle, + /// Supervisor is actively working + Working, + /// Supervisor is waiting for user input/confirmation + WaitingForUser, + /// Supervisor is waiting for spawned tasks to complete + WaitingForTasks, + /// Supervisor is blocked (external dependency, error, etc.) + Blocked, + /// Supervisor has completed its contract + Completed, + /// Supervisor has failed + Failed, + /// Supervisor was interrupted (daemon crash, etc.) + Interrupted, +} + +impl Default for SupervisorStateEnum { + fn default() -> Self { + SupervisorStateEnum::Initializing + } +} + +impl std::fmt::Display for SupervisorStateEnum { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + SupervisorStateEnum::Initializing => write!(f, "initializing"), + SupervisorStateEnum::Idle => write!(f, "idle"), + SupervisorStateEnum::Working => write!(f, "working"), + SupervisorStateEnum::WaitingForUser => write!(f, "waiting_for_user"), + SupervisorStateEnum::WaitingForTasks => write!(f, "waiting_for_tasks"), + SupervisorStateEnum::Blocked => write!(f, "blocked"), + SupervisorStateEnum::Completed => write!(f, "completed"), + SupervisorStateEnum::Failed => write!(f, "failed"), + SupervisorStateEnum::Interrupted => write!(f, "interrupted"), + } + } +} + +impl std::str::FromStr for SupervisorStateEnum { + type Err = String; + + fn from_str(s: &str) -> Result<Self, Self::Err> { + match s.to_lowercase().as_str() { + "initializing" => Ok(SupervisorStateEnum::Initializing), + "idle" => Ok(SupervisorStateEnum::Idle), + "working" => Ok(SupervisorStateEnum::Working), + "waiting_for_user" | "waitingforuser" => Ok(SupervisorStateEnum::WaitingForUser), + "waiting_for_tasks" | "waitingfortasks" => Ok(SupervisorStateEnum::WaitingForTasks), + "blocked" => Ok(SupervisorStateEnum::Blocked), + "completed" => Ok(SupervisorStateEnum::Completed), + "failed" => Ok(SupervisorStateEnum::Failed), + "interrupted" => Ok(SupervisorStateEnum::Interrupted), + _ => Err(format!("Unknown supervisor state: {}", s)), + } + } +} + +/// Enhanced heartbeat record for supervisor task monitoring. +/// Stored in the database for historical analysis and dead supervisor detection. +#[derive(Debug, Clone, FromRow, Serialize, ToSchema)] +#[serde(rename_all = "camelCase")] +pub struct SupervisorHeartbeatRecord { + pub id: Uuid, + pub supervisor_task_id: Uuid, + pub contract_id: Uuid, + pub state: String, + pub phase: String, + pub current_activity: Option<String>, + pub progress: i32, + #[sqlx(try_from = "Vec<Uuid>")] + pub pending_task_ids: Vec<Uuid>, + pub timestamp: DateTime<Utc>, +} + +/// Request payload for sending a supervisor heartbeat. +#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)] +#[serde(rename_all = "camelCase")] +pub struct SupervisorHeartbeatRequest { + pub task_id: Uuid, + pub contract_id: Uuid, + pub state: SupervisorStateEnum, + pub phase: String, + pub current_activity: Option<String>, + /// Progress percentage (0-100) + pub progress: u8, + pub pending_task_ids: Vec<Uuid>, +} + +// ============================================================================= +// Red Team Types +// ============================================================================= + /// Red Team notification record #[derive(Debug, Clone, FromRow, Serialize, ToSchema)] #[serde(rename_all = "camelCase")] @@ -2395,3 +2500,99 @@ impl std::str::FromStr for NotificationSeverity { } } } + +// ============================================================================= +// Unit Tests +// ============================================================================= + +#[cfg(test)] +mod tests { + use super::*; + use uuid::Uuid; + + #[test] + fn test_supervisor_state_enum_display() { + assert_eq!(SupervisorStateEnum::Initializing.to_string(), "initializing"); + assert_eq!(SupervisorStateEnum::Idle.to_string(), "idle"); + assert_eq!(SupervisorStateEnum::Working.to_string(), "working"); + assert_eq!(SupervisorStateEnum::WaitingForUser.to_string(), "waiting_for_user"); + assert_eq!(SupervisorStateEnum::WaitingForTasks.to_string(), "waiting_for_tasks"); + assert_eq!(SupervisorStateEnum::Blocked.to_string(), "blocked"); + assert_eq!(SupervisorStateEnum::Completed.to_string(), "completed"); + assert_eq!(SupervisorStateEnum::Failed.to_string(), "failed"); + assert_eq!(SupervisorStateEnum::Interrupted.to_string(), "interrupted"); + } + + #[test] + fn test_supervisor_state_enum_from_str() { + // Standard lowercase + assert_eq!("initializing".parse::<SupervisorStateEnum>().unwrap(), SupervisorStateEnum::Initializing); + assert_eq!("idle".parse::<SupervisorStateEnum>().unwrap(), SupervisorStateEnum::Idle); + assert_eq!("working".parse::<SupervisorStateEnum>().unwrap(), SupervisorStateEnum::Working); + assert_eq!("waiting_for_user".parse::<SupervisorStateEnum>().unwrap(), SupervisorStateEnum::WaitingForUser); + assert_eq!("waiting_for_tasks".parse::<SupervisorStateEnum>().unwrap(), SupervisorStateEnum::WaitingForTasks); + assert_eq!("blocked".parse::<SupervisorStateEnum>().unwrap(), SupervisorStateEnum::Blocked); + assert_eq!("completed".parse::<SupervisorStateEnum>().unwrap(), SupervisorStateEnum::Completed); + assert_eq!("failed".parse::<SupervisorStateEnum>().unwrap(), SupervisorStateEnum::Failed); + assert_eq!("interrupted".parse::<SupervisorStateEnum>().unwrap(), SupervisorStateEnum::Interrupted); + + // Case insensitive + assert_eq!("WORKING".parse::<SupervisorStateEnum>().unwrap(), SupervisorStateEnum::Working); + assert_eq!("Working".parse::<SupervisorStateEnum>().unwrap(), SupervisorStateEnum::Working); + + // Alternative formats + assert_eq!("waitingforuser".parse::<SupervisorStateEnum>().unwrap(), SupervisorStateEnum::WaitingForUser); + assert_eq!("waitingfortasks".parse::<SupervisorStateEnum>().unwrap(), SupervisorStateEnum::WaitingForTasks); + + // Invalid state + assert!("invalid_state".parse::<SupervisorStateEnum>().is_err()); + } + + #[test] + fn test_supervisor_state_enum_serialization() { + // Test JSON serialization + let state = SupervisorStateEnum::Working; + let json = serde_json::to_string(&state).unwrap(); + assert_eq!(json, "\"working\""); + + // Test JSON deserialization + let deserialized: SupervisorStateEnum = serde_json::from_str("\"working\"").unwrap(); + assert_eq!(deserialized, SupervisorStateEnum::Working); + + // Test underscore variants + let json = "\"waiting_for_user\""; + let deserialized: SupervisorStateEnum = serde_json::from_str(json).unwrap(); + assert_eq!(deserialized, SupervisorStateEnum::WaitingForUser); + } + + #[test] + fn test_supervisor_state_enum_default() { + let default_state = SupervisorStateEnum::default(); + assert_eq!(default_state, SupervisorStateEnum::Initializing); + } + + #[test] + fn test_supervisor_heartbeat_request_serialization() { + let request = SupervisorHeartbeatRequest { + task_id: Uuid::nil(), + contract_id: Uuid::nil(), + state: SupervisorStateEnum::Working, + phase: "execute".to_string(), + current_activity: Some("Implementing feature".to_string()), + progress: 50, + pending_task_ids: vec![Uuid::nil()], + }; + + let json = serde_json::to_string(&request).unwrap(); + assert!(json.contains("\"state\":\"working\"")); + assert!(json.contains("\"phase\":\"execute\"")); + assert!(json.contains("\"progress\":50")); + assert!(json.contains("\"currentActivity\":\"Implementing feature\"")); + + // Test deserialization + let deserialized: SupervisorHeartbeatRequest = serde_json::from_str(&json).unwrap(); + assert_eq!(deserialized.state, SupervisorStateEnum::Working); + assert_eq!(deserialized.phase, "execute"); + assert_eq!(deserialized.progress, 50); + } +} |