telos_agent/tools/builtin/
powershell.rs1use async_trait::async_trait;
4use base64::Engine;
5use serde_json::{Value, json};
6#[cfg(windows)]
7use std::os::windows::process::CommandExt;
8use tokio::io::{AsyncBufReadExt, AsyncReadExt, BufReader};
9use tokio::process::Command;
10use tokio::time::{Duration, timeout};
11
12use crate::error::AgentError;
13use crate::tools::api::{PermissionDecision, Tool, ToolContext, ToolDefinition, ToolOutput};
14
15use super::{optional_usize_any, required_string};
16
17pub struct PowerShellTool;
18
19#[allow(dead_code)]
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub enum PowerShellEdition {
22 Core,
23 Desktop,
24}
25
26impl PowerShellEdition {
27 #[allow(dead_code)]
28 pub fn from_path(path: &str) -> Self {
29 let base = path.rsplit(['/', '\\']).next().unwrap_or(path).to_ascii_lowercase();
30 if base.trim_end_matches(".exe") == "pwsh" { Self::Core } else { Self::Desktop }
31 }
32}
33
34#[allow(dead_code)]
35pub fn encode_powershell_command(command: &str) -> String {
36 let mut bytes = Vec::with_capacity(command.len() * 2);
37 for unit in command.encode_utf16() {
38 bytes.extend_from_slice(&unit.to_le_bytes());
39 }
40 base64::engine::general_purpose::STANDARD.encode(bytes)
41}
42
43pub fn build_powershell_args(command: &str) -> Vec<String> {
44 let wrapped = format!(
45 concat!(
46 "$OutputEncoding = [System.Text.UTF8Encoding]::new($false); ",
47 "[Console]::InputEncoding = [System.Text.UTF8Encoding]::new($false); ",
48 "[Console]::OutputEncoding = [System.Text.UTF8Encoding]::new($false); ",
49 "{command}"
50 ),
51 command = command
52 );
53 vec!["-NoProfile".into(), "-NonInteractive".into(), "-Command".into(), wrapped]
54}
55
56pub fn find_powershell_executable() -> Option<String> {
57 if let Ok(path) = std::env::var("TELOS_POWERSHELL_PATH")
58 && !path.trim().is_empty()
59 {
60 return Some(path);
61 }
62 let candidates: &[&str] =
63 if cfg!(windows) { &["pwsh.exe", "powershell.exe"] } else { &["pwsh", "powershell"] };
64 candidates.iter().find(|candidate| executable_exists(candidate)).map(|s| (*s).into())
65}
66
67fn executable_exists(candidate: &str) -> bool {
68 let mut command = std::process::Command::new(candidate);
69 command
70 .arg("-NoProfile")
71 .arg("-NonInteractive")
72 .arg("-Command")
73 .arg("$PSVersionTable.PSVersion");
74 hide_console_window_std(&mut command);
75 command.output().is_ok()
76}
77
78#[async_trait]
79impl Tool for PowerShellTool {
80 fn definition(&self) -> ToolDefinition {
81 ToolDefinition {
82 name: "PowerShell".into(),
83 description: "Run a PowerShell command in the current working directory. Prefer Read/Edit/Write/Glob/Grep for file operations.".into(),
84 input_schema: json!({
85 "type": "object",
86 "properties": {
87 "command": { "type": "string" },
88 "description": { "type": "string" },
89 "timeout_ms": { "type": "integer", "description": "Maximum runtime in milliseconds. Defaults to 120000." }
90 },
91 "required": ["command"]
92 }),
93 }
94 }
95
96 fn prompt_text(&self) -> Option<&'static str> {
97 Some(
98 "Use the PowerShell tool for shell commands in this environment. \
99Prefer Read, Edit, Write, Glob, or Grep for file operations. \
100Use PowerShell syntax, not Bash syntax. \
101Provide a short `description` summarizing the command's intent. \
102\
103When chaining commands with `&&`, be aware that cmdlets like `cmdkey` misparse \
104target names containing `@` or `/` — the entire chain aborts including earlier \
105commands that already succeeded. Split such commands into separate calls. \
106\
107For Windows Credential Manager operations with special characters in target names, \
108use P/Invoke via Add-Type rather than cmdkey. \
109Always use fully-qualified .NET namespace paths in P/Invoke code. \
110\
111For PowerShell script/module/GUI development, use the `powershell-use` skill \
112(best practices, Windows Forms/WPF, PowerShell Gallery, PSResourceGet).",
113 )
114 }
115
116 async fn validate(&self, arguments: &Value, _context: &ToolContext) -> Result<(), AgentError> {
117 required_string(arguments, "command").map(|_| ())
118 }
119
120 async fn check_permission(
121 &self,
122 arguments: &Value,
123 _context: &ToolContext,
124 ) -> Result<PermissionDecision, AgentError> {
125 let command = required_string(arguments, "command")?;
126 match crate::tools::command_security::powershell::analyze(command) {
127 crate::tools::command_security::powershell::CommandSafety::Safe => {
128 Ok(PermissionDecision::Allow)
129 }
130 crate::tools::command_security::powershell::CommandSafety::NeedsReview { reason } => {
131 Ok(PermissionDecision::Ask {
132 reason: format!("PowerShell command needs review: {reason}"),
133 })
134 }
135 }
136 }
137
138 async fn invoke(
139 &self,
140 arguments: Value,
141 context: ToolContext,
142 ) -> Result<ToolOutput, AgentError> {
143 let command = required_string(&arguments, "command")?;
144 let timeout_ms =
145 optional_usize_any(&arguments, &["timeout_ms"]).unwrap_or(120_000).max(1) as u64;
146 if let Some(tx) = &context.progress {
147 let _ = tx.send(crate::tools::api::ToolProgress {
148 tool_call_id: None,
149 message: format!("running PowerShell command with {timeout_ms}ms timeout"),
150 data: Some(json!({ "command": command, "timeout_ms": timeout_ms })),
151 });
152 }
153 let executable = find_powershell_executable().ok_or_else(|| AgentError::ToolExecution {
154 tool: "PowerShell".into(),
155 message: "PowerShell executable not found; install pwsh or powershell".into(),
156 })?;
157 let mut child = Command::new(executable);
158 child
159 .args(build_powershell_args(command))
160 .current_dir(&context.cwd)
161 .env_clear()
162 .envs(context.env.iter());
163 hide_console_window(&mut child);
164 child.kill_on_drop(true);
165 let progress = context.progress.clone();
166 let output = timeout(
167 Duration::from_millis(timeout_ms),
168 run_powershell_child(child, progress, context.tool_call_id.clone()),
169 )
170 .await
171 .map_err(|_| AgentError::ToolExecution {
172 tool: "PowerShell".into(),
173 message: format!("Command timed out after {timeout_ms}ms"),
174 })?
175 .map_err(|err| AgentError::ToolExecution {
176 tool: "PowerShell".into(),
177 message: err.to_string(),
178 })?;
179
180 let stdout = String::from_utf8_lossy(&output.stdout).to_string();
181 let stderr = String::from_utf8_lossy(&output.stderr).to_string();
182 if !output.status.success() {
183 return Err(AgentError::ToolExecution {
184 tool: "PowerShell".into(),
185 message: format!(
186 "Command failed with exit code {:?}\nstdout:\n{}\nstderr:\n{}",
187 output.status.code(),
188 trim_large_output(&stdout),
189 trim_large_output(&stderr)
190 ),
191 });
192 }
193 Ok(ToolOutput::json(json!({
194 "status": output.status.code(),
195 "success": output.status.success(),
196 "stdout": trim_large_output(&stdout),
197 "stderr": trim_large_output(&stderr),
198 })))
199 }
200}
201
202async fn run_powershell_child(
203 mut command: Command,
204 progress: Option<tokio::sync::mpsc::UnboundedSender<crate::tools::api::ToolProgress>>,
205 tool_call_id: Option<String>,
206) -> std::io::Result<std::process::Output> {
207 command.stdout(std::process::Stdio::piped()).stderr(std::process::Stdio::piped());
208 let mut child = command.spawn()?;
209 let stdout = child.stdout.take().expect("stdout was piped");
210 let stderr = child.stderr.take().expect("stderr was piped");
211
212 let stdout_progress = progress.clone();
213 let stdout_tool_call_id = tool_call_id.clone();
214 let stdout_task = tokio::spawn(async move {
215 read_stream_with_progress(stdout, stdout_progress, stdout_tool_call_id, "stdout").await
216 });
217 let stderr_task = tokio::spawn(async move {
218 read_stream_with_progress(stderr, progress, tool_call_id, "stderr").await
219 });
220
221 let status = child.wait().await?;
222 let stdout = stdout_task.await.map_err(std::io::Error::other)??;
223 let stderr = stderr_task.await.map_err(std::io::Error::other)??;
224 Ok(std::process::Output { status, stdout, stderr })
225}
226
227fn trim_large_output(output: &str) -> String {
228 const MAX_CHARS: usize = 20_000;
229 if output.chars().count() <= MAX_CHARS {
230 return output.to_string();
231 }
232 let preview = output.chars().take(MAX_CHARS).collect::<String>();
233 format!("{preview}\n<truncated output after {MAX_CHARS} chars>")
234}
235
236#[cfg_attr(not(windows), allow(unused_variables))]
237fn hide_console_window(command: &mut Command) {
238 #[cfg(windows)]
239 {
240 const CREATE_NO_WINDOW: u32 = 0x0800_0000;
241 command.creation_flags(CREATE_NO_WINDOW);
242 }
243}
244
245#[cfg_attr(not(windows), allow(unused_variables))]
246fn hide_console_window_std(command: &mut std::process::Command) {
247 #[cfg(windows)]
248 {
249 const CREATE_NO_WINDOW: u32 = 0x0800_0000;
250 command.creation_flags(CREATE_NO_WINDOW);
251 }
252}
253
254async fn read_stream_with_progress(
255 stream: impl tokio::io::AsyncRead + Unpin,
256 progress: Option<tokio::sync::mpsc::UnboundedSender<crate::tools::api::ToolProgress>>,
257 tool_call_id: Option<String>,
258 stream_name: &'static str,
259) -> std::io::Result<Vec<u8>> {
260 let mut reader = BufReader::new(stream);
261 let mut buf = Vec::new();
262 let mut line = String::new();
263
264 loop {
265 line.clear();
266 let bytes_read = reader.read_line(&mut line).await?;
267 if bytes_read == 0 {
268 break;
269 }
270 buf.extend_from_slice(line.as_bytes());
271 if let Some(tx) = &progress {
272 let _ = tx.send(crate::tools::api::ToolProgress {
273 tool_call_id: tool_call_id.clone(),
274 message: format!("{stream_name} update"),
275 data: Some(json!({
276 "stream": stream_name,
277 "output": line,
278 })),
279 });
280 }
281 }
282
283 let mut tail = Vec::new();
284 reader.read_to_end(&mut tail).await?;
285 if !tail.is_empty() {
286 if let Ok(text) = String::from_utf8(tail.clone())
287 && let Some(tx) = &progress
288 {
289 let _ = tx.send(crate::tools::api::ToolProgress {
290 tool_call_id,
291 message: format!("{stream_name} update"),
292 data: Some(json!({
293 "stream": stream_name,
294 "output": text,
295 })),
296 });
297 }
298 buf.extend_from_slice(&tail);
299 }
300 Ok(buf)
301}
302
303#[cfg(test)]
304mod tests {
305 use super::*;
306
307 #[test]
308 fn infers_powershell_edition_from_path() {
309 assert_eq!(PowerShellEdition::from_path("pwsh"), PowerShellEdition::Core);
310 assert_eq!(
311 PowerShellEdition::from_path("C:\\Program Files\\PowerShell\\7\\pwsh.exe"),
312 PowerShellEdition::Core
313 );
314 assert_eq!(PowerShellEdition::from_path("powershell.exe"), PowerShellEdition::Desktop);
315 }
316
317 #[test]
318 fn encoded_command_uses_utf16le_base64() {
319 assert_eq!(
320 encode_powershell_command("Write-Output hi"),
321 "VwByAGkAdABlAC0ATwB1AHQAcAB1AHQAIABoAGkA"
322 );
323 }
324
325 #[test]
326 fn build_args_use_noninteractive_no_profile_command() {
327 let args = build_powershell_args("Get-Process");
328 assert_eq!(&args[..3], ["-NoProfile", "-NonInteractive", "-Command"]);
329 assert!(args[3].contains("[Console]::OutputEncoding"));
330 assert!(args[3].ends_with("Get-Process"));
331 }
332}