pforge_runtime/handlers/
cli.rs

1use crate::{Error, Result};
2use schemars::JsonSchema;
3use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
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: HashMap<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: HashMap<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: HashMap<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!("Failed to execute command '{}': {}", self.command, e))
83            })?;
84
85            Ok::<_, Error>(CliOutput {
86                stdout: String::from_utf8_lossy(&output.stdout).to_string(),
87                stderr: String::from_utf8_lossy(&output.stderr).to_string(),
88                exit_code: output.status.code().unwrap_or(-1),
89            })
90        };
91
92        if let Some(timeout_ms) = self.timeout_ms {
93            timeout(Duration::from_millis(timeout_ms), exec_future)
94                .await
95                .map_err(|_| Error::Timeout)?
96        } else {
97            exec_future.await
98        }
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105
106    #[tokio::test]
107    async fn test_cli_handler_new() {
108        let handler = CliHandler::new(
109            "echo".to_string(),
110            vec!["hello".to_string()],
111            None,
112            HashMap::new(),
113            None,
114            false,
115        );
116
117        assert_eq!(handler.command, "echo");
118        assert_eq!(handler.args.len(), 1);
119        assert_eq!(handler.args[0], "hello");
120        assert!(handler.cwd.is_none());
121        assert!(handler.env.is_empty());
122        assert!(handler.timeout_ms.is_none());
123        assert!(!handler.stream);
124    }
125
126    #[tokio::test]
127    async fn test_cli_handler_execute_simple() {
128        let handler = CliHandler::new(
129            "echo".to_string(),
130            vec!["hello".to_string()],
131            None,
132            HashMap::new(),
133            None,
134            false,
135        );
136
137        let input = CliInput {
138            args: vec![],
139            env: HashMap::new(),
140        };
141
142        let result = handler.execute(input).await;
143        assert!(result.is_ok());
144
145        let output = result.unwrap();
146        assert!(output.stdout.contains("hello"));
147        assert_eq!(output.exit_code, 0);
148    }
149
150    #[tokio::test]
151    async fn test_cli_handler_execute_with_input_args() {
152        let handler = CliHandler::new(
153            "echo".to_string(),
154            vec![],
155            None,
156            HashMap::new(),
157            None,
158            false,
159        );
160
161        let input = CliInput {
162            args: vec!["test".to_string(), "message".to_string()],
163            env: HashMap::new(),
164        };
165
166        let result = handler.execute(input).await;
167        assert!(result.is_ok());
168
169        let output = result.unwrap();
170        assert!(output.stdout.contains("test"));
171        assert!(output.stdout.contains("message"));
172    }
173
174    #[tokio::test]
175    async fn test_cli_handler_execute_with_timeout() {
176        let handler = CliHandler::new(
177            "sleep".to_string(),
178            vec!["2".to_string()],
179            None,
180            HashMap::new(),
181            Some(100), // 100ms timeout
182            false,
183        );
184
185        let input = CliInput {
186            args: vec![],
187            env: HashMap::new(),
188        };
189
190        let result = handler.execute(input).await;
191        assert!(result.is_err());
192        assert!(matches!(result.unwrap_err(), Error::Timeout));
193    }
194
195    #[tokio::test]
196    async fn test_cli_handler_execute_invalid_command() {
197        let handler = CliHandler::new(
198            "nonexistent_command_that_should_fail".to_string(),
199            vec![],
200            None,
201            HashMap::new(),
202            None,
203            false,
204        );
205
206        let input = CliInput {
207            args: vec![],
208            env: HashMap::new(),
209        };
210
211        let result = handler.execute(input).await;
212        assert!(result.is_err());
213        assert!(matches!(result.unwrap_err(), Error::Handler(_)));
214    }
215
216    #[tokio::test]
217    async fn test_cli_handler_with_env() {
218        let mut env = HashMap::new();
219        env.insert("TEST_VAR".to_string(), "test_value".to_string());
220
221        let handler = CliHandler::new(
222            "sh".to_string(),
223            vec!["-c".to_string(), "echo $TEST_VAR".to_string()],
224            None,
225            env,
226            None,
227            false,
228        );
229
230        let input = CliInput {
231            args: vec![],
232            env: HashMap::new(),
233        };
234
235        let result = handler.execute(input).await;
236        assert!(result.is_ok());
237
238        let output = result.unwrap();
239        assert!(output.stdout.contains("test_value"));
240    }
241}
pforge_runtime/handlers/cli.rs

pforge_runtime/handlers/
cli.rs
