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

ruvswarm_persistence/
models.rs
