1use std::collections::HashMap;
4
5use chrono::{DateTime, Utc};
6use schemars::JsonSchema;
7use serde::{Deserialize, Serialize};
8
9#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
11pub struct MachineCapabilities {
12 pub os: String,
13 pub arch: String,
14 pub cpu_cores: u32,
15 pub memory_gb: f64,
16 #[serde(default)]
17 pub gpu: Option<String>,
18 #[serde(default)]
19 pub tags: Vec<String>,
20 #[serde(default)]
21 pub installed_tools: Vec<String>,
22}
23
24#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Eq)]
26#[serde(rename_all = "snake_case")]
27pub enum MachineStatus {
28 Available,
29 Busy,
30 Offline,
31 Maintenance,
32}
33
34#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
36#[serde(rename_all = "snake_case", tag = "type")]
37pub enum MachineTransport {
38 Ssh {
39 host: String,
40 port: u16,
41 user: String,
42 },
43 Relay {
44 relay_url: String,
45 channel_id: String,
46 },
47}
48
49#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
51pub struct Machine {
52 pub id: String,
53 pub name: String,
54 pub transport: MachineTransport,
55 pub capabilities: MachineCapabilities,
56 pub status: MachineStatus,
57 pub last_heartbeat: Option<DateTime<Utc>>,
58 pub added_at: DateTime<Utc>,
59}
60
61#[derive(Debug)]
63pub struct MachineInventory {
64 machines: HashMap<String, Machine>,
65}
66
67impl MachineInventory {
68 pub fn new() -> Self {
70 Self {
71 machines: HashMap::new(),
72 }
73 }
74
75 pub fn register(&mut self, machine: Machine) {
77 self.machines.insert(machine.id.clone(), machine);
78 }
79
80 pub fn deregister(&mut self, machine_id: &str) -> Option<Machine> {
82 self.machines.remove(machine_id)
83 }
84
85 pub fn get(&self, machine_id: &str) -> Option<&Machine> {
87 self.machines.get(machine_id)
88 }
89
90 pub fn get_mut(&mut self, machine_id: &str) -> Option<&mut Machine> {
92 self.machines.get_mut(machine_id)
93 }
94
95 pub fn list(&self) -> Vec<&Machine> {
97 self.machines.values().collect()
98 }
99
100 pub fn by_status(&self, status: &MachineStatus) -> Vec<&Machine> {
102 self.machines
103 .values()
104 .filter(|m| &m.status == status)
105 .collect()
106 }
107
108 pub fn by_tag(&self, tag: &str) -> Vec<&Machine> {
110 self.machines
111 .values()
112 .filter(|m| m.capabilities.tags.iter().any(|t| t == tag))
113 .collect()
114 }
115
116 pub fn by_tool(&self, tool: &str) -> Vec<&Machine> {
118 self.machines
119 .values()
120 .filter(|m| m.capabilities.installed_tools.iter().any(|t| t == tool))
121 .collect()
122 }
123
124 pub fn heartbeat(&mut self, machine_id: &str) -> anyhow::Result<()> {
126 let machine = self
127 .machines
128 .get_mut(machine_id)
129 .ok_or_else(|| anyhow::anyhow!("unknown machine: {machine_id}"))?;
130 machine.last_heartbeat = Some(Utc::now());
131 if machine.status == MachineStatus::Offline {
132 machine.status = MachineStatus::Available;
133 }
134 Ok(())
135 }
136
137 pub fn mark_offline(&mut self, machine_id: &str) -> anyhow::Result<()> {
139 let machine = self
140 .machines
141 .get_mut(machine_id)
142 .ok_or_else(|| anyhow::anyhow!("unknown machine: {machine_id}"))?;
143 machine.status = MachineStatus::Offline;
144 Ok(())
145 }
146
147 pub fn count(&self) -> usize {
149 self.machines.len()
150 }
151}
152
153impl Default for MachineInventory {
154 fn default() -> Self {
155 Self::new()
156 }
157}
158
159#[cfg(test)]
160mod tests {
161 use super::*;
162
163 fn make_machine(id: &str, name: &str) -> Machine {
164 Machine {
165 id: id.into(),
166 name: name.into(),
167 transport: MachineTransport::Ssh {
168 host: "10.0.0.1".into(),
169 port: 22,
170 user: "deploy".into(),
171 },
172 capabilities: MachineCapabilities {
173 os: "linux".into(),
174 arch: "x86_64".into(),
175 cpu_cores: 8,
176 memory_gb: 32.0,
177 gpu: None,
178 tags: vec!["production".into()],
179 installed_tools: vec!["docker".into(), "cargo".into()],
180 },
181 status: MachineStatus::Available,
182 last_heartbeat: None,
183 added_at: Utc::now(),
184 }
185 }
186
187 #[test]
188 fn test_inventory_new() {
189 let inv = MachineInventory::new();
190 assert_eq!(inv.count(), 0);
191 assert!(inv.list().is_empty());
192 }
193
194 #[test]
195 fn test_register_and_get() {
196 let mut inv = MachineInventory::new();
197 inv.register(make_machine("m1", "prod-1"));
198 assert_eq!(inv.count(), 1);
199 assert_eq!(inv.get("m1").unwrap().name, "prod-1");
200 }
201
202 #[test]
203 fn test_deregister() {
204 let mut inv = MachineInventory::new();
205 inv.register(make_machine("m1", "prod-1"));
206 let removed = inv.deregister("m1");
207 assert!(removed.is_some());
208 assert_eq!(inv.count(), 0);
209 }
210
211 #[test]
212 fn test_by_status() {
213 let mut inv = MachineInventory::new();
214 inv.register(make_machine("m1", "prod-1"));
215 let mut m2 = make_machine("m2", "staging-1");
216 m2.status = MachineStatus::Offline;
217 inv.register(m2);
218
219 let available = inv.by_status(&MachineStatus::Available);
220 assert_eq!(available.len(), 1);
221 assert_eq!(available[0].id, "m1");
222
223 let offline = inv.by_status(&MachineStatus::Offline);
224 assert_eq!(offline.len(), 1);
225 assert_eq!(offline[0].id, "m2");
226 }
227
228 #[test]
229 fn test_by_tag() {
230 let mut inv = MachineInventory::new();
231 inv.register(make_machine("m1", "prod-1"));
232 let mut m2 = make_machine("m2", "staging-1");
233 m2.capabilities.tags = vec!["staging".into()];
234 inv.register(m2);
235
236 let prod = inv.by_tag("production");
237 assert_eq!(prod.len(), 1);
238 assert_eq!(prod[0].id, "m1");
239 }
240
241 #[test]
242 fn test_by_tool() {
243 let mut inv = MachineInventory::new();
244 inv.register(make_machine("m1", "prod-1"));
245
246 let docker_machines = inv.by_tool("docker");
247 assert_eq!(docker_machines.len(), 1);
248
249 let kubectl_machines = inv.by_tool("kubectl");
250 assert!(kubectl_machines.is_empty());
251 }
252
253 #[test]
254 fn test_heartbeat() {
255 let mut inv = MachineInventory::new();
256 let mut m = make_machine("m1", "prod-1");
257 m.status = MachineStatus::Offline;
258 inv.register(m);
259
260 inv.heartbeat("m1").unwrap();
261 let machine = inv.get("m1").unwrap();
262 assert!(machine.last_heartbeat.is_some());
263 assert_eq!(machine.status, MachineStatus::Available);
264 }
265
266 #[test]
267 fn test_heartbeat_unknown() {
268 let mut inv = MachineInventory::new();
269 assert!(inv.heartbeat("unknown").is_err());
270 }
271
272 #[test]
273 fn test_mark_offline() {
274 let mut inv = MachineInventory::new();
275 inv.register(make_machine("m1", "prod-1"));
276 inv.mark_offline("m1").unwrap();
277 assert_eq!(inv.get("m1").unwrap().status, MachineStatus::Offline);
278 }
279
280 #[test]
281 fn test_machine_serialization() {
282 let machine = make_machine("m1", "prod-1");
283 let json = serde_json::to_string(&machine).unwrap();
284 let back: Machine = serde_json::from_str(&json).unwrap();
285 assert_eq!(machine.id, back.id);
286 assert_eq!(machine.name, back.name);
287 }
288
289 #[test]
290 fn test_machine_transport_variants() {
291 let ssh = MachineTransport::Ssh {
292 host: "10.0.0.1".into(),
293 port: 22,
294 user: "root".into(),
295 };
296 let json = serde_json::to_string(&ssh).unwrap();
297 assert!(json.contains("ssh"));
298
299 let relay = MachineTransport::Relay {
300 relay_url: "wss://relay.mur.run".into(),
301 channel_id: "ch-1".into(),
302 };
303 let json = serde_json::to_string(&relay).unwrap();
304 assert!(json.contains("relay"));
305 }
306
307 #[test]
308 fn test_machine_status_serialization() {
309 for status in [
310 MachineStatus::Available,
311 MachineStatus::Busy,
312 MachineStatus::Offline,
313 MachineStatus::Maintenance,
314 ] {
315 let json = serde_json::to_string(&status).unwrap();
316 let back: MachineStatus = serde_json::from_str(&json).unwrap();
317 assert_eq!(status, back);
318 }
319 }
320
321 #[test]
322 fn test_default_inventory() {
323 let inv = MachineInventory::default();
324 assert_eq!(inv.count(), 0);
325 }
326
327 #[test]
328 fn test_get_mut() {
329 let mut inv = MachineInventory::new();
330 inv.register(make_machine("m1", "prod-1"));
331 let machine = inv.get_mut("m1").unwrap();
332 machine.name = "renamed".into();
333 assert_eq!(inv.get("m1").unwrap().name, "renamed");
334 }
335}