ruv_swarm_persistence/
models.rs

1//! Data models for persistence layer
2
3use chrono::{DateTime, Utc};
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6use uuid::Uuid;
7
8/// Agent model representing a swarm agent
9#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
10pub struct AgentModel {
11    pub id: String,
12    pub name: String,
13    pub agent_type: String,
14    pub status: AgentStatus,
15    pub capabilities: Vec<String>,
16    pub metadata: HashMap<String, serde_json::Value>,
17    pub heartbeat: DateTime<Utc>,
18    pub created_at: DateTime<Utc>,
19    pub updated_at: DateTime<Utc>,
20}
21
22impl AgentModel {
23    pub fn new(name: String, agent_type: String, capabilities: Vec<String>) -> Self {
24        let now = Utc::now();
25        Self {
26            id: Uuid::new_v4().to_string(),
27            name,
28            agent_type,
29            status: AgentStatus::Initializing,
30            capabilities,
31            metadata: HashMap::new(),
32            heartbeat: now,
33            created_at: now,
34            updated_at: now,
35        }
36    }
37    
38    pub fn update_heartbeat(&mut self) {
39        self.heartbeat = Utc::now();
40        self.updated_at = Utc::now();
41    }
42    
43    pub fn set_status(&mut self, status: AgentStatus) {
44        self.status = status;
45        self.updated_at = Utc::now();
46    }
47}
48
49/// Agent status enumeration
50#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
51#[serde(rename_all = "snake_case")]
52pub enum AgentStatus {
53    Initializing,
54    Active,
55    Idle,
56    Busy,
57    Paused,
58    Error,
59    Shutdown,
60}
61
62impl ToString for AgentStatus {
63    fn to_string(&self) -> String {
64        match self {
65            AgentStatus::Initializing => "initializing",
66            AgentStatus::Active => "active",
67            AgentStatus::Idle => "idle",
68            AgentStatus::Busy => "busy",
69            AgentStatus::Paused => "paused",
70            AgentStatus::Error => "error",
71            AgentStatus::Shutdown => "shutdown",
72        }.to_string()
73    }
74}
75
76/// Task model representing work items in the swarm
77#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
78pub struct TaskModel {
79    pub id: String,
80    pub task_type: String,
81    pub priority: TaskPriority,
82    pub status: TaskStatus,
83    pub assigned_to: Option<String>,
84    pub payload: serde_json::Value,
85    pub result: Option<serde_json::Value>,
86    pub error: Option<String>,
87    pub retry_count: u32,
88    pub max_retries: u32,
89    pub dependencies: Vec<String>,
90    pub created_at: DateTime<Utc>,
91    pub updated_at: DateTime<Utc>,
92    pub started_at: Option<DateTime<Utc>>,
93    pub completed_at: Option<DateTime<Utc>>,
94}
95
96impl TaskModel {
97    pub fn new(task_type: String, payload: serde_json::Value, priority: TaskPriority) -> Self {
98        let now = Utc::now();
99        Self {
100            id: Uuid::new_v4().to_string(),
101            task_type,
102            priority,
103            status: TaskStatus::Pending,
104            assigned_to: None,
105            payload,
106            result: None,
107            error: None,
108            retry_count: 0,
109            max_retries: 3,
110            dependencies: Vec::new(),
111            created_at: now,
112            updated_at: now,
113            started_at: None,
114            completed_at: None,
115        }
116    }
117    
118    pub fn assign_to(&mut self, agent_id: &str) {
119        self.assigned_to = Some(agent_id.to_string());
120        self.status = TaskStatus::Assigned;
121        self.updated_at = Utc::now();
122    }
123    
124    pub fn start(&mut self) {
125        self.status = TaskStatus::Running;
126        self.started_at = Some(Utc::now());
127        self.updated_at = Utc::now();
128    }
129    
130    pub fn complete(&mut self, result: serde_json::Value) {
131        self.status = TaskStatus::Completed;
132        self.result = Some(result);
133        self.completed_at = Some(Utc::now());
134        self.updated_at = Utc::now();
135    }
136    
137    pub fn fail(&mut self, error: String) {
138        self.retry_count += 1;
139        if self.retry_count >= self.max_retries {
140            self.status = TaskStatus::Failed;
141        } else {
142            self.status = TaskStatus::Pending;
143            self.assigned_to = None;
144        }
145        self.error = Some(error);
146        self.updated_at = Utc::now();
147    }
148}
149
150/// Task priority levels
151#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
152#[serde(rename_all = "snake_case")]
153pub enum TaskPriority {
154    Low = 0,
155    Medium = 1,
156    High = 2,
157    Critical = 3,
158}
159
160/// Task status enumeration
161#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
162#[serde(rename_all = "snake_case")]
163pub enum TaskStatus {
164    Pending,
165    Assigned,
166    Running,
167    Completed,
168    Failed,
169    Cancelled,
170}
171
172/// Event model for event sourcing
173#[derive(Debug, Clone, Serialize, Deserialize)]
174pub struct EventModel {
175    pub id: String,
176    pub event_type: String,
177    pub agent_id: Option<String>,
178    pub task_id: Option<String>,
179    pub payload: serde_json::Value,
180    pub metadata: HashMap<String, serde_json::Value>,
181    pub timestamp: DateTime<Utc>,
182    pub sequence: u64,
183}
184
185impl EventModel {
186    pub fn new(event_type: String, payload: serde_json::Value) -> Self {
187        Self {
188            id: Uuid::new_v4().to_string(),
189            event_type,
190            agent_id: None,
191            task_id: None,
192            payload,
193            metadata: HashMap::new(),
194            timestamp: Utc::now(),
195            sequence: 0, // Should be set by storage layer
196        }
197    }
198    
199    pub fn with_agent(mut self, agent_id: String) -> Self {
200        self.agent_id = Some(agent_id);
201        self
202    }
203    
204    pub fn with_task(mut self, task_id: String) -> Self {
205        self.task_id = Some(task_id);
206        self
207    }
208}
209
210/// Message model for inter-agent communication
211#[derive(Debug, Clone, Serialize, Deserialize)]
212pub struct MessageModel {
213    pub id: String,
214    pub from_agent: String,
215    pub to_agent: String,
216    pub message_type: String,
217    pub content: serde_json::Value,
218    pub priority: MessagePriority,
219    pub read: bool,
220    pub created_at: DateTime<Utc>,
221    pub read_at: Option<DateTime<Utc>>,
222}
223
224impl MessageModel {
225    pub fn new(
226        from_agent: String,
227        to_agent: String,
228        message_type: String,
229        content: serde_json::Value,
230    ) -> Self {
231        Self {
232            id: Uuid::new_v4().to_string(),
233            from_agent,
234            to_agent,
235            message_type,
236            content,
237            priority: MessagePriority::Normal,
238            read: false,
239            created_at: Utc::now(),
240            read_at: None,
241        }
242    }
243    
244    pub fn with_priority(mut self, priority: MessagePriority) -> Self {
245        self.priority = priority;
246        self
247    }
248    
249    pub fn mark_read(&mut self) {
250        self.read = true;
251        self.read_at = Some(Utc::now());
252    }
253}
254
255/// Message priority levels
256#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
257#[serde(rename_all = "snake_case")]
258pub enum MessagePriority {
259    Low,
260    Normal,
261    High,
262    Urgent,
263}
264
265/// Metric model for performance tracking
266#[derive(Debug, Clone, Serialize, Deserialize)]
267pub struct MetricModel {
268    pub id: String,
269    pub metric_type: String,
270    pub agent_id: Option<String>,
271    pub value: f64,
272    pub unit: String,
273    pub tags: HashMap<String, String>,
274    pub timestamp: DateTime<Utc>,
275}
276
277impl MetricModel {
278    pub fn new(metric_type: String, value: f64, unit: String) -> Self {
279        Self {
280            id: Uuid::new_v4().to_string(),
281            metric_type,
282            agent_id: None,
283            value,
284            unit,
285            tags: HashMap::new(),
286            timestamp: Utc::now(),
287        }
288    }
289    
290    pub fn with_agent(mut self, agent_id: String) -> Self {
291        self.agent_id = Some(agent_id);
292        self
293    }
294    
295    pub fn with_tag(mut self, key: String, value: String) -> Self {
296        self.tags.insert(key, value);
297        self
298    }
299}
300
301#[cfg(test)]
302mod tests {
303    use super::*;
304    
305    #[test]
306    fn test_agent_model() {
307        let mut agent = AgentModel::new(
308            "test-agent".to_string(),
309            "worker".to_string(),
310            vec!["compute".to_string(), "storage".to_string()],
311        );
312        
313        assert_eq!(agent.status, AgentStatus::Initializing);
314        agent.set_status(AgentStatus::Active);
315        assert_eq!(agent.status, AgentStatus::Active);
316    }
317    
318    #[test]
319    fn test_task_model() {
320        let mut task = TaskModel::new(
321            "process".to_string(),
322            serde_json::json!({"data": "test"}),
323            TaskPriority::High,
324        );
325        
326        assert_eq!(task.status, TaskStatus::Pending);
327        task.assign_to("agent-1");
328        assert_eq!(task.status, TaskStatus::Assigned);
329        assert_eq!(task.assigned_to, Some("agent-1".to_string()));
330    }
331}
ruv_swarm_persistence/models.rs

ruv_swarm_persistence/
models.rs
