Skip to main content

rpi_tools/tools/
powershell.rs

1//! Optional Windows PowerShell tool, matching upstream's `powershell` tool.
2//! The command is passed through `-EncodedCommand` so quotes and newlines in
3//! user input cannot be reinterpreted by the outer Bash transport.
4
5use std::sync::Arc;
6
7use async_trait::async_trait;
8use base64::{engine::general_purpose::STANDARD, Engine as _};
9use rpi_agent::agent_tool::AgentTool;
10use rpi_agent::error::AgentError;
11use rpi_agent::types::{AgentToolResult, TextContentOrImage, ToolResultPartial};
12use rpi_ai::types::Tool;
13use schemars::JsonSchema;
14use serde::Deserialize;
15use tokio_util::sync::CancellationToken;
16
17use crate::env::ShellOverride;
18use crate::shell_output::{execute_shell_with_capture, ShellCaptureOptions};
19use crate::tools::tool_context::ExecutionToolContext;
20
21#[derive(Debug, Clone, JsonSchema, Deserialize)]
22pub struct PowerShellInput {
23    pub command: String,
24    #[serde(default)]
25    pub timeout: Option<f64>,
26}
27
28pub struct PowerShellTool {
29    schema: Tool,
30    env: Arc<dyn crate::env::ExecutionEnv>,
31}
32
33impl PowerShellTool {
34    fn schema() -> Tool {
35        Tool {
36            name: "powershell".into(),
37            description: "Execute a PowerShell command on Windows. Returns stdout and stderr."
38                .into(),
39            parameters: rpi_ai::types::Schema::new(
40                serde_json::to_value(schemars::schema_for!(PowerShellInput)).unwrap_or_default(),
41            ),
42            constrained_sampling: None,
43        }
44    }
45}
46
47pub fn create_powershell_tool(context: &ExecutionToolContext) -> Arc<dyn AgentTool> {
48    Arc::new(PowerShellTool {
49        schema: PowerShellTool::schema(),
50        env: context.env().clone(),
51    })
52}
53
54#[async_trait]
55impl AgentTool for PowerShellTool {
56    fn schema(&self) -> &Tool {
57        &self.schema
58    }
59    fn label(&self) -> &str {
60        "powershell"
61    }
62
63    async fn execute(
64        &self,
65        _tool_call_id: &str,
66        params: serde_json::Value,
67        signal: CancellationToken,
68        _on_update: Arc<dyn Fn(ToolResultPartial) + Send + Sync>,
69    ) -> Result<AgentToolResult, AgentError> {
70        if !cfg!(windows) {
71            return Err(AgentError::Tool(
72                "PowerShell tool is only available on Windows".into(),
73            ));
74        }
75        let input: PowerShellInput = serde_json::from_value(params)
76            .map_err(|e| AgentError::Validation(format!("powershell input invalid: {e}")))?;
77        let utf16: Vec<u8> = input
78            .command
79            .encode_utf16()
80            .flat_map(u16::to_le_bytes)
81            .collect();
82        let encoded = STANDARD.encode(utf16);
83        let command = encoded;
84        let program = if cfg!(windows) {
85            std::env::var_os("ProgramFiles")
86                .map(|p| std::path::PathBuf::from(p).join("PowerShell/7/pwsh.exe"))
87                .filter(|p| p.exists())
88                .unwrap_or_else(|| std::path::PathBuf::from("powershell.exe"))
89        } else {
90            std::path::PathBuf::from("powershell.exe")
91        };
92        let capture = execute_shell_with_capture(
93            &self.env,
94            &command,
95            ShellCaptureOptions {
96                shell: Some(ShellOverride {
97                    program,
98                    args: vec![
99                        "-NoLogo".into(),
100                        "-NoProfile".into(),
101                        "-NonInteractive".into(),
102                        "-ExecutionPolicy".into(),
103                        "Bypass".into(),
104                        "-EncodedCommand".into(),
105                    ],
106                }),
107                cwd: Some(self.env.cwd().to_path_buf()),
108                env: None,
109                inherit_env: true,
110                timeout: input.timeout,
111                cancel: Some(&signal),
112                on_chunk: None,
113                return_execution_errors: true,
114            },
115        )
116        .await
117        .map_err(|e| AgentError::Tool(e.to_string()))?;
118        let mut output = capture.output;
119        if capture.cancelled {
120            return Err(AgentError::Tool("PowerShell command aborted".into()));
121        }
122        if let Some(error) = capture.execution_error {
123            return Err(AgentError::Tool(error.to_string()));
124        }
125        if let Some(code) = capture.exit_code.filter(|code| *code != 0) {
126            output.push_str(&format!("\n\nPowerShell exited with code {code}"));
127            return Err(AgentError::Tool(output));
128        }
129        if output.trim().is_empty() {
130            output = "(no output)".into();
131        }
132        Ok(AgentToolResult {
133            content: vec![TextContentOrImage::text(output)],
134            ..Default::default()
135        })
136    }
137}