Skip to main content

oxios_kernel/tools/kernel/
resource_tool.rs

1//! Resource tool — wraps `InfraApi` resource methods behind the `AgentTool` interface.
2//!
3//! Provides agents with system resource monitoring capabilities.
4//! Actions: snapshot, history, overloaded.
5//!
6//! ## Example
7//!
8//! ```json
9//! { "action": "snapshot" }
10//! { "action": "history", "last_n": 10 }
11//! { "action": "overloaded" }
12//! ```
13
14use std::sync::Arc;
15
16use async_trait::async_trait;
17use oxi_sdk::{AgentTool, AgentToolResult, ToolContext};
18use serde_json::{json, Value};
19use tokio::sync::oneshot;
20
21use crate::kernel_handle::KernelHandle;
22use crate::resource_monitor::ResourceMonitor;
23
24/// Agent tool for resource monitoring.
25///
26/// Wraps the resource-related methods of the `InfraApi` domain. Allows agents
27/// to query system resource usage, history, and overload status.
28///
29/// ## Actions
30///
31/// | Action       | Description                     | Required params | Optional params |
32/// |--------------|---------------------------------|-----------------|-----------------|
33/// | `snapshot`   | Get current resource snapshot   | —               | —               |
34/// | `history`    | Get recent resource snapshots   | —               | `last_n`        |
35/// | `overloaded` | Check if system is overloaded   | —               | —               |
36pub struct ResourceTool {
37    resource_monitor: Arc<ResourceMonitor>,
38}
39
40impl ResourceTool {
41    /// Create a new `ResourceTool` from a `KernelHandle`.
42    ///
43    /// Extracts the `ResourceMonitor` Arc from the kernel's Infra API.
44    pub fn from_kernel(kernel: &KernelHandle) -> Self {
45        Self {
46            resource_monitor: kernel.infra.resource_monitor.clone(),
47        }
48    }
49}
50
51impl std::fmt::Debug for ResourceTool {
52    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
53        f.debug_struct("ResourceTool").finish()
54    }
55}
56
57#[async_trait]
58impl AgentTool for ResourceTool {
59    fn name(&self) -> &str {
60        "resource"
61    }
62
63    fn label(&self) -> &str {
64        "Resource"
65    }
66
67    fn description(&self) -> &'static str {
68        "Monitor system resources — CPU, memory, disk, agent count. \
69         Actions: snapshot, history, overloaded."
70    }
71
72    fn parameters_schema(&self) -> Value {
73        json!({
74            "type": "object",
75            "properties": {
76                "action": {
77                    "type": "string",
78                    "enum": ["snapshot", "history", "overloaded"],
79                    "description": "Resource operation to perform"
80                },
81                "last_n": {
82                    "type": "integer",
83                    "description": "Number of historical snapshots to return (history action, default: 10)"
84                }
85            },
86            "required": ["action"]
87        })
88    }
89
90    async fn execute(
91        &self,
92        _tool_call_id: &str,
93        params: Value,
94        _signal: Option<oneshot::Receiver<()>>,
95        _ctx: &ToolContext,
96    ) -> Result<AgentToolResult, String> {
97        let action = params
98            .get("action")
99            .and_then(|v| v.as_str())
100            .ok_or_else(|| "Missing required parameter: action".to_string())?;
101
102        match action {
103            "snapshot" => {
104                let snap = self.resource_monitor.snapshot();
105                Ok(AgentToolResult::success(
106                    serde_json::to_string_pretty(&json!({
107                        "timestamp": snap.timestamp.to_rfc3339(),
108                        "cpu_percent": format!("{:.1}%", snap.cpu_percent),
109                        "memory_used_mb": snap.memory_used_mb,
110                        "memory_total_mb": snap.memory_total_mb,
111                        "memory_percent": format!(
112                            "{:.1}%",
113                            if snap.memory_total_mb > 0 {
114                                (snap.memory_used_mb as f64 / snap.memory_total_mb as f64) * 100.0
115                            } else {
116                                0.0
117                            }
118                        ),
119                        "active_agents": snap.active_agents,
120                        "pending_tasks": snap.pending_tasks,
121                        "total_token_usage": snap.total_token_usage,
122                        "disk_used_gb": format!("{:.2}", snap.disk_used_gb),
123                        "load_avg_1m": format!("{:.2}", snap.load_avg_1m),
124                    }))
125                    .unwrap_or_default(),
126                ))
127            }
128
129            "history" => {
130                let last_n = params["last_n"].as_u64().unwrap_or(10) as usize;
131                let history = self.resource_monitor.history(last_n);
132
133                if history.is_empty() {
134                    return Ok(AgentToolResult::success(
135                        "No resource history available yet.",
136                    ));
137                }
138
139                let display: Vec<Value> = history
140                    .iter()
141                    .map(|snap| {
142                        json!({
143                            "timestamp": snap.timestamp.to_rfc3339(),
144                            "cpu_percent": format!("{:.1}%", snap.cpu_percent),
145                            "memory_mb": format!("{}/{}", snap.memory_used_mb, snap.memory_total_mb),
146                            "active_agents": snap.active_agents,
147                            "load_avg_1m": format!("{:.2}", snap.load_avg_1m),
148                        })
149                    })
150                    .collect();
151
152                Ok(AgentToolResult::success(
153                    serde_json::to_string_pretty(&json!({
154                        "snapshots": display,
155                        "count": display.len(),
156                    }))
157                    .unwrap_or_default(),
158                ))
159            }
160
161            "overloaded" => {
162                let overloaded = self.resource_monitor.is_overloaded();
163                Ok(AgentToolResult::success(
164                    serde_json::to_string(&json!({
165                        "overloaded": overloaded,
166                        "status": if overloaded { "OVERLOADED" } else { "NOMINAL" },
167                    }))
168                    .unwrap_or_default(),
169                ))
170            }
171
172            other => Err(format!(
173                "Unknown resource action '{}'. Valid: snapshot, history, overloaded",
174                other
175            )),
176        }
177    }
178}
179
180#[cfg(test)]
181mod tests {
182    use super::*;
183
184    #[test]
185    fn test_schema_structure() {
186        let schema = json!({
187            "type": "object",
188            "properties": {
189                "action": {
190                    "type": "string",
191                    "enum": ["snapshot", "history", "overloaded"]
192                },
193                "last_n": { "type": "integer" }
194            },
195            "required": ["action"]
196        });
197
198        let actions = schema["properties"]["action"]["enum"].as_array().unwrap();
199        assert_eq!(actions.len(), 3);
200        assert!(actions.iter().any(|a| a == "snapshot"));
201        assert!(actions.iter().any(|a| a == "history"));
202        assert!(actions.iter().any(|a| a == "overloaded"));
203    }
204}