rpi_tools/tools/
powershell.rs1use 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}