Skip to main content

telos_agent/tools/builtin/
powershell.rs

1//! `PowerShell` tool — run PowerShell commands in the workspace cwd.
2
3use 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}