|
5 | 5 | //! ## Available Tools |
6 | 6 | //! |
7 | 7 | //! - `docker_ps` - List running containers |
| 8 | +//! - `docker_stats` - Get container resource usage (CPU, memory, I/O) |
8 | 9 | //! - `docker_build` - Build images |
9 | 10 | //! - `docker_run` - Run containers |
10 | 11 | //! - `docker_logs` - Get container logs |
@@ -34,6 +35,7 @@ impl DockerTools { |
34 | 35 | pub fn all() -> Vec<Box<dyn Tool>> { |
35 | 36 | vec![ |
36 | 37 | Box::new(DockerPsTool::new()), |
| 38 | + Box::new(DockerStatsTool::new()), |
37 | 39 | Box::new(DockerBuildTool::new()), |
38 | 40 | Box::new(DockerRunTool::new()), |
39 | 41 | Box::new(DockerLogsTool::new()), |
@@ -152,6 +154,120 @@ impl Tool for DockerPsTool { |
152 | 154 | } |
153 | 155 | } |
154 | 156 |
|
| 157 | +// ============================================================================ |
| 158 | +// Docker Stats Tool |
| 159 | +// ============================================================================ |
| 160 | + |
| 161 | +/// Get container resource usage statistics |
| 162 | +pub struct DockerStatsTool { |
| 163 | + config: ToolConfig, |
| 164 | +} |
| 165 | + |
| 166 | +impl DockerStatsTool { |
| 167 | + pub fn new() -> Self { |
| 168 | + let parameters = create_schema( |
| 169 | + serde_json::json!({ |
| 170 | + "all": { |
| 171 | + "type": "boolean", |
| 172 | + "description": "Show all containers (default shows just running)", |
| 173 | + "default": false |
| 174 | + }, |
| 175 | + "no_trunc": { |
| 176 | + "type": "boolean", |
| 177 | + "description": "Don't truncate output", |
| 178 | + "default": false |
| 179 | + }, |
| 180 | + "format": { |
| 181 | + "type": "string", |
| 182 | + "description": "Format output (json for structured data)", |
| 183 | + "default": "table" |
| 184 | + } |
| 185 | + }), |
| 186 | + vec![], |
| 187 | + ); |
| 188 | + |
| 189 | + Self { |
| 190 | + config: tool_config_with_timeout( |
| 191 | + "docker_stats", |
| 192 | + "Get resource usage statistics for containers (CPU, memory, network, I/O). Returns a one-time snapshot.", |
| 193 | + parameters, |
| 194 | + 30, |
| 195 | + ), |
| 196 | + } |
| 197 | + } |
| 198 | +} |
| 199 | + |
| 200 | +impl Default for DockerStatsTool { |
| 201 | + fn default() -> Self { |
| 202 | + Self::new() |
| 203 | + } |
| 204 | +} |
| 205 | + |
| 206 | +#[async_trait] |
| 207 | +impl Tool for DockerStatsTool { |
| 208 | + async fn execute(&self, input: ToolInput) -> AofResult<ToolResult> { |
| 209 | + let all: bool = input.get_arg("all").unwrap_or(false); |
| 210 | + let no_trunc: bool = input.get_arg("no_trunc").unwrap_or(false); |
| 211 | + let format: String = input.get_arg("format").unwrap_or_else(|_| "table".to_string()); |
| 212 | + |
| 213 | + // CRITICAL: Always use --no-stream to get a one-time snapshot |
| 214 | + // Without this flag, docker stats runs continuously like 'top' |
| 215 | + let mut args = vec!["stats".to_string(), "--no-stream".to_string()]; |
| 216 | + |
| 217 | + if all { |
| 218 | + args.push("-a".to_string()); |
| 219 | + } |
| 220 | + |
| 221 | + if no_trunc { |
| 222 | + args.push("--no-trunc".to_string()); |
| 223 | + } |
| 224 | + |
| 225 | + if format == "json" { |
| 226 | + args.push("--format={{json .}}".to_string()); |
| 227 | + } |
| 228 | + |
| 229 | + debug!(args = ?args, "Executing docker stats"); |
| 230 | + |
| 231 | + let args_str: Vec<&str> = args.iter().map(|s| s.as_str()).collect(); |
| 232 | + let result = execute_command("docker", &args_str, None, 30).await; |
| 233 | + |
| 234 | + match result { |
| 235 | + Ok(output) => { |
| 236 | + if output.success { |
| 237 | + if format == "json" { |
| 238 | + // Parse JSON lines output |
| 239 | + let stats: Vec<serde_json::Value> = output |
| 240 | + .stdout |
| 241 | + .lines() |
| 242 | + .filter_map(|line| serde_json::from_str(line).ok()) |
| 243 | + .collect(); |
| 244 | + |
| 245 | + Ok(ToolResult::success(serde_json::json!({ |
| 246 | + "stats": stats, |
| 247 | + "count": stats.len() |
| 248 | + }))) |
| 249 | + } else { |
| 250 | + // Return table format as-is |
| 251 | + Ok(ToolResult::success(serde_json::json!({ |
| 252 | + "output": output.stdout |
| 253 | + }))) |
| 254 | + } |
| 255 | + } else { |
| 256 | + Ok(ToolResult::error(format!( |
| 257 | + "docker stats failed: {}", |
| 258 | + output.stderr |
| 259 | + ))) |
| 260 | + } |
| 261 | + } |
| 262 | + Err(e) => Ok(ToolResult::error(e)), |
| 263 | + } |
| 264 | + } |
| 265 | + |
| 266 | + fn config(&self) -> &ToolConfig { |
| 267 | + &self.config |
| 268 | + } |
| 269 | +} |
| 270 | + |
155 | 271 | // ============================================================================ |
156 | 272 | // Docker Build Tool |
157 | 273 | // ============================================================================ |
|
0 commit comments