1use serde::{Deserialize, Serialize};
8use serde_json::Value;
9use std::collections::HashMap;
10use std::time::Duration;
11use tokio::io::AsyncWriteExt;
12
13#[derive(Debug, Serialize)]
15struct JsonRpcRequest {
16 jsonrpc: &'static str,
17 method: String,
18 params: Value,
19 id: u64,
20}
21
22#[derive(Debug, Deserialize)]
24struct JsonRpcResponse {
25 #[allow(dead_code)]
26 jsonrpc: Option<String>,
27 result: Option<Value>,
28 error: Option<JsonRpcError>,
29 #[allow(dead_code)]
30 id: Option<u64>,
31}
32
33#[derive(Debug, Deserialize)]
34struct JsonRpcError {
35 #[allow(dead_code)]
36 code: Option<i64>,
37 message: String,
38}
39
40#[derive(Debug, Clone)]
42pub struct SubprocessTool {
43 pub command: String,
45 pub args: Vec<String>,
47 pub cwd: Option<String>,
49 pub env: HashMap<String, String>,
51 pub timeout: Duration,
53}
54
55impl SubprocessTool {
56 pub fn new(command: &str) -> Self {
57 Self {
58 command: command.to_string(),
59 args: Vec::new(),
60 cwd: None,
61 env: HashMap::new(),
62 timeout: Duration::from_secs(30),
63 }
64 }
65
66 pub fn with_args(mut self, args: Vec<String>) -> Self {
67 self.args = args;
68 self
69 }
70
71 pub fn with_cwd(mut self, cwd: &str) -> Self {
72 self.cwd = Some(cwd.to_string());
73 self
74 }
75
76 pub fn with_timeout(mut self, timeout: Duration) -> Self {
77 self.timeout = timeout;
78 self
79 }
80
81 pub fn with_env(mut self, key: &str, value: &str) -> Self {
82 self.env.insert(key.to_string(), value.to_string());
83 self
84 }
85}
86
87pub struct SubprocessToolExecutor {
89 tools: HashMap<String, SubprocessTool>,
90 fallback: Option<std::sync::Arc<dyn super::ToolExecutor>>,
92 next_id: std::sync::atomic::AtomicU64,
93}
94
95impl SubprocessToolExecutor {
96 pub fn new() -> Self {
97 Self {
98 tools: HashMap::new(),
99 fallback: None,
100 next_id: std::sync::atomic::AtomicU64::new(1),
101 }
102 }
103
104 pub fn register(&mut self, name: &str, tool: SubprocessTool) {
106 self.tools.insert(name.to_string(), tool);
107 }
108
109 pub fn with_fallback(mut self, fallback: std::sync::Arc<dyn super::ToolExecutor>) -> Self {
111 self.fallback = Some(fallback);
112 self
113 }
114
115 async fn execute_subprocess(
116 &self,
117 tool_name: &str,
118 tool: &SubprocessTool,
119 params: &Value,
120 timeout_ms: Option<u64>,
121 ) -> Result<Value, String> {
122 let effective_timeout = match timeout_ms {
130 Some(ms) => tool.timeout.max(Duration::from_millis(ms)),
131 None => tool.timeout,
132 };
133 let request = JsonRpcRequest {
134 jsonrpc: "2.0",
135 method: tool_name.to_string(),
136 params: params.clone(),
137 id: self
138 .next_id
139 .fetch_add(1, std::sync::atomic::Ordering::Relaxed),
140 };
141
142 let request_json = serde_json::to_string(&request)
143 .map_err(|e| format!("failed to serialize request: {}", e))?;
144
145 let mut cmd = crate::spawn::program_command(&tool.command);
148 cmd.args(&tool.args)
149 .stdin(std::process::Stdio::piped())
150 .stdout(std::process::Stdio::piped())
151 .stderr(std::process::Stdio::piped());
152
153 if let Some(ref cwd) = tool.cwd {
154 cmd.current_dir(cwd);
155 }
156 for (k, v) in &tool.env {
157 cmd.env(k, v);
158 }
159
160 let mut child = cmd
161 .spawn()
162 .map_err(|e| format!("failed to spawn subprocess '{}': {}", tool.command, e))?;
163
164 if let Some(mut stdin) = child.stdin.take() {
166 stdin
167 .write_all(request_json.as_bytes())
168 .await
169 .map_err(|e| format!("failed to write to subprocess stdin: {}", e))?;
170 stdin
171 .write_all(b"\n")
172 .await
173 .map_err(|e| format!("failed to write newline to stdin: {}", e))?;
174 }
176
177 let output = match tokio::time::timeout(effective_timeout, child.wait_with_output()).await {
179 Ok(Ok(output)) => {
180 let stdout = String::from_utf8_lossy(&output.stdout).to_string();
181 if !output.status.success() && stdout.trim().is_empty() {
182 let stderr = String::from_utf8_lossy(&output.stderr);
183 return Err(format!(
184 "subprocess exited with status {}: {}",
185 output.status,
186 stderr.trim()
187 ));
188 }
189 stdout
190 }
191 Ok(Err(e)) => {
192 return Err(format!("failed to read subprocess output: {}", e));
193 }
194 Err(_) => {
195 return Err(format!(
197 "subprocess '{}' timed out after {:?}",
198 tool.command, effective_timeout
199 ));
200 }
201 };
202
203 let response: JsonRpcResponse = serde_json::from_str(&output).map_err(|e| {
205 format!(
206 "invalid JSON-RPC response from '{}': {} (raw: {})",
207 tool.command,
208 e,
209 output.trim()
210 )
211 })?;
212
213 if let Some(error) = response.error {
214 return Err(format!("subprocess tool error: {}", error.message));
215 }
216
217 response
218 .result
219 .ok_or_else(|| "subprocess returned no result".to_string())
220 }
221}
222
223impl Default for SubprocessToolExecutor {
224 fn default() -> Self {
225 Self::new()
226 }
227}
228
229#[async_trait::async_trait]
230impl super::ToolExecutor for SubprocessToolExecutor {
231 async fn execute(&self, tool: &str, params: &Value) -> Result<Value, String> {
232 self.execute_with_action(tool, params, "", None).await
233 }
234
235 async fn execute_with_action(
236 &self,
237 tool: &str,
238 params: &Value,
239 action_id: &str,
240 timeout_ms: Option<u64>,
241 ) -> Result<Value, String> {
242 self.execute_with_action_in_session(tool, params, action_id, timeout_ms, None, 1)
243 .await
244 }
245
246 async fn execute_with_action_in_session(
247 &self,
248 tool: &str,
249 params: &Value,
250 action_id: &str,
251 timeout_ms: Option<u64>,
252 session_id: Option<&str>,
253 attempt: u32,
254 ) -> Result<Value, String> {
255 if let Some(subprocess_tool) = self.tools.get(tool) {
256 self.execute_subprocess(tool, subprocess_tool, params, timeout_ms)
260 .await
261 } else if let Some(ref fallback) = self.fallback {
262 fallback
263 .execute_with_action_in_session(
264 tool, params, action_id, timeout_ms, session_id, attempt,
265 )
266 .await
267 } else {
268 Err(format!(
269 "unknown subprocess tool: '{}' (no fallback configured)",
270 tool
271 ))
272 }
273 }
274}
275
276#[cfg(test)]
277mod tests {
278 #[cfg(unix)]
281 use super::*;
282 #[cfg(unix)]
283 use crate::ToolExecutor;
284
285 #[cfg(unix)]
292 #[tokio::test]
293 async fn action_budget_lifts_short_subprocess_timeout() {
294 let mut exec = SubprocessToolExecutor::new();
295 let tool = SubprocessTool::new("sh")
298 .with_args(vec![
299 "-c".to_string(),
300 "sleep 1; printf '{\"jsonrpc\":\"2.0\",\"result\":{\"ok\":true},\"id\":1}'"
301 .to_string(),
302 ])
303 .with_timeout(Duration::from_millis(100));
304 exec.register("slow", tool);
305
306 let reaped = exec.execute("slow", &Value::Null).await;
308 assert!(
309 reaped.is_err() && reaped.as_ref().unwrap_err().contains("timed out"),
310 "expected the short internal timeout to reap: {reaped:?}"
311 );
312
313 let ok = exec
316 .execute_with_action("slow", &Value::Null, "a0", Some(5_000))
317 .await;
318 assert!(
319 ok.is_ok(),
320 "action budget must lift the short timeout: {ok:?}"
321 );
322 assert_eq!(ok.unwrap(), serde_json::json!({ "ok": true }));
323 }
324}