Skip to main content

mur_core/remote/
inventory.rs

1//! Machine inventory for tracking available remote machines and their capabilities.
2
3use std::collections::HashMap;
4
5use chrono::{DateTime, Utc};
6use schemars::JsonSchema;
7use serde::{Deserialize, Serialize};
8
9/// Capabilities and resources of a remote machine.
10#[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/// Status of a machine in the inventory.
25#[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/// Transport configuration for reaching a machine.
35#[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/// A registered machine in the inventory.
50#[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/// Inventory of available remote machines.
62#[derive(Debug)]
63pub struct MachineInventory {
64    machines: HashMap<String, Machine>,
65}
66
67impl MachineInventory {
68    /// Create an empty inventory.
69    pub fn new() -> Self {
70        Self {
71            machines: HashMap::new(),
72        }
73    }
74
75    /// Register a machine in the inventory.
76    pub fn register(&mut self, machine: Machine) {
77        self.machines.insert(machine.id.clone(), machine);
78    }
79
80    /// Remove a machine from the inventory.
81    pub fn deregister(&mut self, machine_id: &str) -> Option<Machine> {
82        self.machines.remove(machine_id)
83    }
84
85    /// Get a machine by ID.
86    pub fn get(&self, machine_id: &str) -> Option<&Machine> {
87        self.machines.get(machine_id)
88    }
89
90    /// Get a mutable reference to a machine.
91    pub fn get_mut(&mut self, machine_id: &str) -> Option<&mut Machine> {
92        self.machines.get_mut(machine_id)
93    }
94
95    /// List all machines.
96    pub fn list(&self) -> Vec<&Machine> {
97        self.machines.values().collect()
98    }
99
100    /// List machines matching a given status.
101    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    /// Find machines that have a specific tag.
109    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    /// Find machines that have a specific tool installed.
117    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    /// Update the heartbeat timestamp for a machine.
125    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    /// Mark a machine as offline.
138    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    /// Total number of registered machines.
148    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}