pforge_runtime/handlers/
cli.rs

1use crate::{Error, Result};
2use rustc_hash::FxHashMap;
3use schemars::JsonSchema;
4use serde::{Deserialize, Serialize};
5use std::process::Stdio;
6use tokio::process::Command;
7use tokio::time::{timeout, Duration};
8
9#[derive(Debug, Clone)]
10pub struct CliHandler {
11    pub command: String,
12    pub args: Vec<String>,
13    pub cwd: Option<String>,
14    pub env: FxHashMap<String, String>,
15    pub timeout_ms: Option<u64>,
16    pub stream: bool,
17}
18
19#[derive(Debug, Deserialize, JsonSchema)]
20pub struct CliInput {
21    #[serde(default)]
22    pub args: Vec<String>,
23    #[serde(default)]
24    pub env: FxHashMap<String, String>,
25}
26
27#[derive(Debug, Serialize, JsonSchema)]
28pub struct CliOutput {
29    pub stdout: String,
30    pub stderr: String,
31    pub exit_code: i32,
32}
33
34impl CliHandler {
35    pub fn new(
36        command: String,
37        args: Vec<String>,
38        cwd: Option<String>,
39        env: FxHashMap<String, String>,
40        timeout_ms: Option<u64>,
41        stream: bool,
42    ) -> Self {
43        Self {
44            command,
45            args,
46            cwd,
47            env,
48            timeout_ms,
49            stream,
50        }
51    }
52
53    pub async fn execute(&self, input: CliInput) -> Result<CliOutput> {
54        let mut cmd = Command::new(&self.command);
55
56        // Add base args
57        cmd.args(&self.args);
58
59        // Add input args
60        cmd.args(&input.args);
61
62        // Set working directory
63        if let Some(cwd) = &self.cwd {
64            cmd.current_dir(cwd);
65        }
66
67        // Set environment variables (base + input)
68        for (k, v) in &self.env {
69            cmd.env(k, v);
70        }
71        for (k, v) in &input.env {
72            cmd.env(k, v);
73        }
74
75        // Configure stdio
76        cmd.stdout(Stdio::piped());
77        cmd.stderr(Stdio::piped());
78
79        // Execute with timeout
80        let exec_future = async {
81            let output = cmd.output().await.map_err(|e| {
82                Error::Handler(format!(
83                    "Failed to execute command '{}': {}",
84                    self.command, e
85                ))
86            })?;
87
88            Ok::<_, Error>(CliOutput {
89                stdout: String::from_utf8_lossy(&output.stdout).to_string(),
90                stderr: String::from_utf8_lossy(&output.stderr).to_string(),
91                exit_code: output.status.code().unwrap_or(-1),
92            })
93        };
94
95        if let Some(timeout_ms) = self.timeout_ms {
96            timeout(Duration::from_millis(timeout_ms), exec_future)
97                .await
98                .map_err(|_| Error::Timeout)?
99        } else {
100            exec_future.await
101        }
102    }
103}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108
109    #[tokio::test]
110    async fn test_cli_handler_new() {
111        let handler = CliHandler::new(
112            "echo".to_string(),
113            vec!["hello".to_string()],
114            None,
115            FxHashMap::default(),
116            None,
117            false,
118        );
119
120        assert_eq!(handler.command, "echo");
121        assert_eq!(handler.args.len(), 1);
122        assert_eq!(handler.args[0], "hello");
123        assert!(handler.cwd.is_none());
124        assert!(handler.env.is_empty());
125        assert!(handler.timeout_ms.is_none());
126        assert!(!handler.stream);
127    }
128
129    #[tokio::test]
130    async fn test_cli_handler_execute_simple() {
131        let handler = CliHandler::new(
132            "echo".to_string(),
133            vec!["hello".to_string()],
134            None,
135            FxHashMap::default(),
136            None,
137            false,
138        );
139
140        let input = CliInput {
141            args: vec![],
142            env: FxHashMap::default(),
143        };
144
145        let result = handler.execute(input).await;
146        assert!(result.is_ok());
147
148        let output = result.unwrap();
149        assert!(output.stdout.contains("hello"));
150        assert_eq!(output.exit_code, 0);
151    }
152
153    #[tokio::test]
154    async fn test_cli_handler_execute_with_input_args() {
155        let handler = CliHandler::new(
156            "echo".to_string(),
157            vec![],
158            None,
159            FxHashMap::default(),
160            None,
161            false,
162        );
163
164        let input = CliInput {
165            args: vec!["test".to_string(), "message".to_string()],
166            env: FxHashMap::default(),
167        };
168
169        let result = handler.execute(input).await;
170        assert!(result.is_ok());
171
172        let output = result.unwrap();
173        assert!(output.stdout.contains("test"));
174        assert!(output.stdout.contains("message"));
175    }
176
177    #[tokio::test]
178    async fn test_cli_handler_execute_with_timeout() {
179        let handler = CliHandler::new(
180            "sleep".to_string(),
181            vec!["2".to_string()],
182            None,
183            FxHashMap::default(),
184            Some(100), // 100ms timeout
185            false,
186        );
187
188        let input = CliInput {
189            args: vec![],
190            env: FxHashMap::default(),
191        };
192
193        let result = handler.execute(input).await;
194        assert!(result.is_err());
195        assert!(matches!(result.unwrap_err(), Error::Timeout));
196    }
197
198    #[tokio::test]
199    async fn test_cli_handler_execute_invalid_command() {
200        let handler = CliHandler::new(
201            "nonexistent_command_that_should_fail".to_string(),
202            vec![],
203            None,
204            FxHashMap::default(),
205            None,
206            false,
207        );
208
209        let input = CliInput {
210            args: vec![],
211            env: FxHashMap::default(),
212        };
213
214        let result = handler.execute(input).await;
215        assert!(result.is_err());
216        assert!(matches!(result.unwrap_err(), Error::Handler(_)));
217    }
218
219    #[tokio::test]
220    async fn test_cli_handler_with_env() {
221        let mut env = FxHashMap::default();
222        env.insert("TEST_VAR".to_string(), "test_value".to_string());
223
224        let handler = CliHandler::new(
225            "sh".to_string(),
226            vec!["-c".to_string(), "echo $TEST_VAR".to_string()],
227            None,
228            env,
229            None,
230            false,
231        );
232
233        let input = CliInput {
234            args: vec![],
235            env: FxHashMap::default(),
236        };
237
238        let result = handler.execute(input).await;
239        assert!(result.is_ok());
240
241        let output = result.unwrap();
242        assert!(output.stdout.contains("test_value"));
243    }
244}
pforge_runtime/handlers/cli.rs

pforge_runtime/handlers/
cli.rs
