Skip to main content

aegis_tools/
system.rs

1//! # SystemStatusTool
2//!
3//! Lets the agent see the machine it runs on: CPU, memory, swap, load average,
4//! disks, uptime and OS info. Backed by `sysinfo` (pure-Rust, cross-platform).
5//! Read-only and light — the cornerstone of Aegis's server-resident scenario.
6
7use crate::registry::{Tool, ToolContext};
8use anyhow::Result;
9use async_trait::async_trait;
10use serde_json::{json, Value};
11
12/// Reports host CPU/memory/disk/load/uptime.
13pub struct SystemStatusTool;
14
15impl SystemStatusTool {
16    /// Create a new `SystemStatusTool`.
17    pub fn new() -> Self {
18        SystemStatusTool
19    }
20}
21
22impl Default for SystemStatusTool {
23    fn default() -> Self {
24        Self::new()
25    }
26}
27
28#[async_trait]
29impl Tool for SystemStatusTool {
30    fn name(&self) -> &str {
31        "system_status"
32    }
33
34    fn description(&self) -> &str {
35        "Report the host's current status: OS, uptime, load average, CPU usage, memory/swap usage, and disk usage. Read-only."
36    }
37
38    fn parameters(&self) -> Value {
39        json!({
40            "type": "object",
41            "properties": {}
42        })
43    }
44
45    async fn execute(&self, _args: Value, _ctx: &ToolContext<'_>) -> Result<String> {
46        // sysinfo is synchronous and CPU sampling needs a short sleep, so run
47        // it on a blocking thread to avoid stalling the async runtime.
48        let out = tokio::task::spawn_blocking(collect_status)
49            .await
50            .map_err(|e| anyhow::anyhow!("system_status task failed: {e}"))?;
51        Ok(out)
52    }
53}
54
55fn human_bytes(bytes: u64) -> String {
56    const UNITS: &[&str] = &["B", "KiB", "MiB", "GiB", "TiB"];
57    let mut v = bytes as f64;
58    let mut i = 0;
59    while v >= 1024.0 && i < UNITS.len() - 1 {
60        v /= 1024.0;
61        i += 1;
62    }
63    format!("{v:.1} {}", UNITS[i])
64}
65
66fn pct(used: u64, total: u64) -> String {
67    if total == 0 {
68        "n/a".to_string()
69    } else {
70        format!("{:.0}%", (used as f64 / total as f64) * 100.0)
71    }
72}
73
74fn collect_status() -> String {
75    use sysinfo::{Disks, System};
76
77    let mut sys = System::new_all();
78    // CPU usage needs two samples separated by a short interval.
79    std::thread::sleep(sysinfo::MINIMUM_CPU_UPDATE_INTERVAL);
80    sys.refresh_cpu_usage();
81
82    let mut out = String::new();
83
84    // Host / OS
85    out.push_str(&format!(
86        "Host: {}\nOS: {} {} (kernel {})\nArch: {}\n",
87        System::host_name().unwrap_or_else(|| "unknown".into()),
88        System::name().unwrap_or_else(|| "unknown".into()),
89        System::os_version().unwrap_or_else(|| "".into()),
90        System::kernel_version().unwrap_or_else(|| "unknown".into()),
91        System::cpu_arch(),
92    ));
93
94    // Uptime
95    let up = System::uptime();
96    let days = up / 86400;
97    let hours = (up % 86400) / 3600;
98    let mins = (up % 3600) / 60;
99    out.push_str(&format!("Uptime: {days}d {hours}h {mins}m\n"));
100
101    // Load average
102    let la = System::load_average();
103    out.push_str(&format!(
104        "Load average: {:.2} {:.2} {:.2} (1m/5m/15m)\n",
105        la.one, la.five, la.fifteen
106    ));
107
108    // CPU
109    out.push_str(&format!(
110        "CPU: {:.1}% used across {} logical core(s)\n",
111        sys.global_cpu_usage(),
112        sys.cpus().len(),
113    ));
114
115    // Memory / swap
116    let (tm, um) = (sys.total_memory(), sys.used_memory());
117    out.push_str(&format!(
118        "Memory: {} / {} ({} used)\n",
119        human_bytes(um),
120        human_bytes(tm),
121        pct(um, tm)
122    ));
123    let (ts, us) = (sys.total_swap(), sys.used_swap());
124    if ts > 0 {
125        out.push_str(&format!(
126            "Swap: {} / {} ({} used)\n",
127            human_bytes(us),
128            human_bytes(ts),
129            pct(us, ts)
130        ));
131    }
132
133    // Disks
134    let disks = Disks::new_with_refreshed_list();
135    if disks.is_empty() {
136        out.push_str("Disks: (none reported)\n");
137    } else {
138        out.push_str("Disks:\n");
139        for d in &disks {
140            let total = d.total_space();
141            let avail = d.available_space();
142            let used = total.saturating_sub(avail);
143            out.push_str(&format!(
144                "  {} at {} — {} / {} ({} used)\n",
145                d.name().to_string_lossy(),
146                d.mount_point().display(),
147                human_bytes(used),
148                human_bytes(total),
149                pct(used, total),
150            ));
151        }
152    }
153
154    out.trim_end().to_string()
155}
156
157#[cfg(test)]
158mod tests {
159    use super::*;
160
161    #[test]
162    fn human_bytes_formats() {
163        assert_eq!(human_bytes(0), "0.0 B");
164        assert_eq!(human_bytes(1024), "1.0 KiB");
165        assert_eq!(human_bytes(1024 * 1024), "1.0 MiB");
166    }
167
168    #[test]
169    fn pct_handles_zero() {
170        assert_eq!(pct(0, 0), "n/a");
171        assert_eq!(pct(1, 2), "50%");
172    }
173}
174
175/// Read-only process listing (top consumers by CPU or memory).
176pub struct ProcessListTool;
177
178impl ProcessListTool {
179    /// Create a new `ProcessListTool`.
180    pub fn new() -> Self {
181        ProcessListTool
182    }
183}
184
185impl Default for ProcessListTool {
186    fn default() -> Self {
187        Self::new()
188    }
189}
190
191#[async_trait]
192impl Tool for ProcessListTool {
193    fn name(&self) -> &str {
194        "process"
195    }
196
197    fn description(&self) -> &str {
198        "List the top running processes by CPU or memory usage (read-only). Optionally filter by name."
199    }
200
201    fn parameters(&self) -> Value {
202        json!({
203            "type": "object",
204            "properties": {
205                "sort_by": { "type": "string", "enum": ["cpu", "memory"], "description": "Sort key (default cpu)" },
206                "limit": { "type": "integer", "description": "How many processes to show (default 15)" },
207                "name": { "type": "string", "description": "Only show processes whose name contains this substring" }
208            }
209        })
210    }
211
212    async fn execute(&self, args: Value, _ctx: &ToolContext<'_>) -> Result<String> {
213        let sort_by = args["sort_by"].as_str().unwrap_or("cpu").to_string();
214        let limit = args["limit"].as_u64().unwrap_or(15) as usize;
215        let name_filter = args["name"].as_str().map(|s| s.to_lowercase());
216
217        let out = tokio::task::spawn_blocking(move || collect_processes(&sort_by, limit, name_filter))
218            .await
219            .map_err(|e| anyhow::anyhow!("process task failed: {e}"))?;
220        Ok(out)
221    }
222}
223
224fn collect_processes(sort_by: &str, limit: usize, name_filter: Option<String>) -> String {
225    use sysinfo::{ProcessesToUpdate, System};
226
227    let mut sys = System::new_all();
228    // Second sample for accurate CPU usage.
229    std::thread::sleep(sysinfo::MINIMUM_CPU_UPDATE_INTERVAL);
230    sys.refresh_processes(ProcessesToUpdate::All, true);
231
232    let mut rows: Vec<(String, u32, f32, u64)> = sys
233        .processes()
234        .iter()
235        .map(|(pid, p)| {
236            (
237                p.name().to_string_lossy().to_string(),
238                pid.as_u32(),
239                p.cpu_usage(),
240                p.memory(),
241            )
242        })
243        .filter(|(name, _, _, _)| {
244            name_filter
245                .as_ref()
246                .map(|f| name.to_lowercase().contains(f))
247                .unwrap_or(true)
248        })
249        .collect();
250
251    if sort_by == "memory" {
252        rows.sort_by(|a, b| b.3.cmp(&a.3));
253    } else {
254        rows.sort_by(|a, b| b.2.partial_cmp(&a.2).unwrap_or(std::cmp::Ordering::Equal));
255    }
256    rows.truncate(limit);
257
258    if rows.is_empty() {
259        return "(no matching processes)".to_string();
260    }
261
262    let mut out = format!("Top {} processes by {}:\n", rows.len(), sort_by);
263    out.push_str(&format!("{:>8}  {:>6}  {:>10}  {}\n", "PID", "CPU%", "MEM", "NAME"));
264    for (name, pid, cpu, mem) in rows {
265        out.push_str(&format!(
266            "{pid:>8}  {cpu:>5.1}  {:>10}  {name}\n",
267            human_bytes(mem)
268        ));
269    }
270    out.trim_end().to_string()
271}