diff --git a/crates/bashkit-python/src/lib.rs b/crates/bashkit-python/src/lib.rs index 7bc57444..f7c7da65 100644 --- a/crates/bashkit-python/src/lib.rs +++ b/crates/bashkit-python/src/lib.rs @@ -491,7 +491,12 @@ impl ScriptedTool { fn execute<'py>(&self, py: Python<'py>, commands: String) -> PyResult> { let mut tool = self.build_rust_tool(); future_into_py(py, async move { - let resp = tool.execute(ToolRequest { commands }).await; + let resp = tool + .execute(ToolRequest { + commands, + timeout_ms: None, + }) + .await; Ok(ExecResult { stdout: resp.stdout, stderr: resp.stderr, @@ -507,7 +512,13 @@ impl ScriptedTool { let rt = tokio::runtime::Runtime::new() .map_err(|e| PyRuntimeError::new_err(format!("Failed to create runtime: {}", e)))?; - let resp = rt.block_on(async move { tool.execute(ToolRequest { commands }).await }); + let resp = rt.block_on(async move { + tool.execute(ToolRequest { + commands, + timeout_ms: None, + }) + .await + }); Ok(ExecResult { stdout: resp.stdout, stderr: resp.stderr, diff --git a/crates/bashkit/examples/scripted_tool.rs b/crates/bashkit/examples/scripted_tool.rs index d9551f35..88fea09a 100644 --- a/crates/bashkit/examples/scripted_tool.rs +++ b/crates/bashkit/examples/scripted_tool.rs @@ -37,6 +37,7 @@ async fn main() -> anyhow::Result<()> { let resp = tool .execute(ToolRequest { commands: "get_user --id 1".to_string(), + timeout_ms: None, }) .await; println!("$ get_user --id 1"); @@ -47,6 +48,7 @@ async fn main() -> anyhow::Result<()> { let resp = tool .execute(ToolRequest { commands: "get_user --id 1 | jq -r '.name'".to_string(), + timeout_ms: None, }) .await; println!("$ get_user --id 1 | jq -r '.name'"); @@ -67,6 +69,7 @@ async fn main() -> anyhow::Result<()> { let resp = tool .execute(ToolRequest { commands: script.to_string(), + timeout_ms: None, }) .await; println!("$ "); @@ -91,6 +94,7 @@ async fn main() -> anyhow::Result<()> { let resp = tool .execute(ToolRequest { commands: script.to_string(), + timeout_ms: None, }) .await; println!("$ "); @@ -113,6 +117,7 @@ async fn main() -> anyhow::Result<()> { let resp = tool .execute(ToolRequest { commands: script.to_string(), + timeout_ms: None, }) .await; println!("$ "); @@ -133,6 +138,7 @@ async fn main() -> anyhow::Result<()> { let resp = tool .execute(ToolRequest { commands: script.to_string(), + timeout_ms: None, }) .await; println!("$ "); diff --git a/crates/bashkit/src/scripted_tool/mod.rs b/crates/bashkit/src/scripted_tool/mod.rs index bd6b80d8..7907fecb 100644 --- a/crates/bashkit/src/scripted_tool/mod.rs +++ b/crates/bashkit/src/scripted_tool/mod.rs @@ -42,6 +42,7 @@ //! //! let resp = tool.execute(ToolRequest { //! commands: "greet --name Alice".to_string(), +//! timeout_ms: None, //! }).await; //! //! assert_eq!(resp.stdout.trim(), "hello Alice"); @@ -415,6 +416,7 @@ mod tests { let resp = tool .execute(ToolRequest { commands: String::new(), + timeout_ms: None, }) .await; assert_eq!(resp.exit_code, 0); @@ -427,6 +429,7 @@ mod tests { let resp = tool .execute(ToolRequest { commands: "get_user --id 42".to_string(), + timeout_ms: None, }) .await; assert_eq!(resp.exit_code, 0); @@ -440,6 +443,7 @@ mod tests { let resp = tool .execute(ToolRequest { commands: "get_user --id=42".to_string(), + timeout_ms: None, }) .await; assert_eq!(resp.exit_code, 0); @@ -452,6 +456,7 @@ mod tests { let resp = tool .execute(ToolRequest { commands: "get_user --id 42 | jq -r '.name'".to_string(), + timeout_ms: None, }) .await; assert_eq!(resp.exit_code, 0); @@ -471,6 +476,7 @@ mod tests { let resp = tool .execute(ToolRequest { commands: script.to_string(), + timeout_ms: None, }) .await; assert_eq!(resp.exit_code, 0); @@ -483,6 +489,7 @@ mod tests { let resp = tool .execute(ToolRequest { commands: "fail_tool".to_string(), + timeout_ms: None, }) .await; assert_ne!(resp.exit_code, 0); @@ -495,6 +502,7 @@ mod tests { let resp = tool .execute(ToolRequest { commands: "fail_tool || echo 'fallback'".to_string(), + timeout_ms: None, }) .await; assert_eq!(resp.exit_code, 0); @@ -507,6 +515,7 @@ mod tests { let resp = tool .execute(ToolRequest { commands: "echo hello | from_stdin".to_string(), + timeout_ms: None, }) .await; assert_eq!(resp.exit_code, 0); @@ -524,6 +533,7 @@ mod tests { let resp = tool .execute(ToolRequest { commands: script.to_string(), + timeout_ms: None, }) .await; assert_eq!(resp.exit_code, 0); @@ -545,6 +555,7 @@ mod tests { let resp = tool .execute(ToolRequest { commands: script.to_string(), + timeout_ms: None, }) .await; assert_eq!(resp.exit_code, 0); @@ -563,6 +574,7 @@ mod tests { let resp = tool .execute(ToolRequest { commands: "echo $API_BASE".to_string(), + timeout_ms: None, }) .await; assert_eq!(resp.exit_code, 0); @@ -581,6 +593,7 @@ mod tests { .execute_with_status( ToolRequest { commands: "get_user --id 1".to_string(), + timeout_ms: None, }, Box::new(move |status| { phases_clone @@ -605,6 +618,7 @@ mod tests { let resp1 = tool .execute(ToolRequest { commands: "get_user --id 1 | jq -r '.name'".to_string(), + timeout_ms: None, }) .await; assert_eq!(resp1.stdout.trim(), "Alice"); @@ -612,6 +626,7 @@ mod tests { let resp2 = tool .execute(ToolRequest { commands: "get_orders --user_id 1 | jq 'length'".to_string(), + timeout_ms: None, }) .await; assert_eq!(resp2.stdout.trim(), "2"); @@ -639,6 +654,7 @@ mod tests { let resp = tool .execute(ToolRequest { commands: "search --verbose --query hello".to_string(), + timeout_ms: None, }) .await; assert_eq!(resp.exit_code, 0); @@ -657,6 +673,7 @@ mod tests { let resp = tool .execute(ToolRequest { commands: "echo_args --name Alice --count 3".to_string(), + timeout_ms: None, }) .await; assert_eq!(resp.exit_code, 0); diff --git a/crates/bashkit/src/tool.rs b/crates/bashkit/src/tool.rs index 227682d8..ddb3b3be 100644 --- a/crates/bashkit/src/tool.rs +++ b/crates/bashkit/src/tool.rs @@ -26,6 +26,7 @@ //! // Execution //! let resp = tool.execute(ToolRequest { //! commands: "echo hello".to_string(), +//! timeout_ms: None, //! }).await; //! assert_eq!(resp.stdout, "hello\n"); //! # }); @@ -38,6 +39,7 @@ use async_trait::async_trait; use schemars::{schema_for, JsonSchema}; use serde::{Deserialize, Serialize}; use std::sync::{Arc, Mutex}; +use std::time::Duration; /// Library version from Cargo.toml pub const VERSION: &str = env!("CARGO_PKG_VERSION"); @@ -136,6 +138,21 @@ SEE ALSO pub struct ToolRequest { /// Bash commands to execute (like `bash -c "commands"`) pub commands: String, + /// Optional per-call timeout in milliseconds. + /// When set, execution is aborted after this duration and a response + /// with `exit_code = 124` is returned (matching the bash `timeout` convention). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub timeout_ms: Option, +} + +impl ToolRequest { + /// Create a request with just commands (no timeout). + pub fn new(commands: impl Into) -> Self { + Self { + commands: commands.into(), + timeout_ms: None, + } + } } /// Response from executing a bash script @@ -621,14 +638,26 @@ impl Tool for BashTool { let mut bash = self.create_bash(); - match bash.exec(&req.commands).await { - Ok(result) => result.into(), - Err(e) => ToolResponse { - stdout: String::new(), - stderr: e.to_string(), - exit_code: 1, - error: Some(error_kind(&e)), - }, + let fut = async { + match bash.exec(&req.commands).await { + Ok(result) => result.into(), + Err(e) => ToolResponse { + stdout: String::new(), + stderr: e.to_string(), + exit_code: 1, + error: Some(error_kind(&e)), + }, + } + }; + + if let Some(ms) = req.timeout_ms { + let dur = Duration::from_millis(ms); + match tokio::time::timeout(dur, fut).await { + Ok(resp) => resp, + Err(_elapsed) => timeout_response(dur), + } + } else { + fut.await } } @@ -669,21 +698,35 @@ impl Tool for BashTool { } }); - let response = match bash.exec_streaming(&req.commands, output_cb).await { - Ok(result) => result.into(), - Err(e) => ToolResponse { - stdout: String::new(), - stderr: e.to_string(), - exit_code: 1, - error: Some(error_kind(&e)), - }, + let timeout_ms = req.timeout_ms; + + let fut = async { + let response = match bash.exec_streaming(&req.commands, output_cb).await { + Ok(result) => result.into(), + Err(e) => ToolResponse { + stdout: String::new(), + stderr: e.to_string(), + exit_code: 1, + error: Some(error_kind(&e)), + }, + }; + + if let Ok(mut cb) = status_cb.lock() { + cb(ToolStatus::new("complete").with_percent(100.0)); + } + + response }; - if let Ok(mut cb) = status_cb.lock() { - cb(ToolStatus::new("complete").with_percent(100.0)); + if let Some(ms) = timeout_ms { + let dur = Duration::from_millis(ms); + match tokio::time::timeout(dur, fut).await { + Ok(resp) => resp, + Err(_elapsed) => timeout_response(dur), + } + } else { + fut.await } - - response } } @@ -701,7 +744,21 @@ fn error_kind(e: &Error) -> String { } } +/// Build a ToolResponse for a timed-out execution (exit code 124, like bash `timeout`). +fn timeout_response(dur: Duration) -> ToolResponse { + ToolResponse { + stdout: String::new(), + stderr: format!( + "bashkit: execution timed out after {:.1}s\n", + dur.as_secs_f64() + ), + exit_code: 124, + error: Some("timeout".to_string()), + } +} + #[cfg(test)] +#[allow(clippy::unwrap_used)] mod tests { use super::*; @@ -793,6 +850,7 @@ mod tests { let mut tool = BashTool::default(); let req = ToolRequest { commands: String::new(), + timeout_ms: None, }; let resp = tool.execute(req).await; assert_eq!(resp.exit_code, 0); @@ -804,6 +862,7 @@ mod tests { let mut tool = BashTool::default(); let req = ToolRequest { commands: "echo hello".to_string(), + timeout_ms: None, }; let resp = tool.execute(req).await; assert_eq!(resp.stdout, "hello\n"); @@ -1019,6 +1078,7 @@ mod tests { let mut tool = BashTool::default(); let req = ToolRequest { commands: "echo test".to_string(), + timeout_ms: None, }; let phases = Arc::new(Mutex::new(Vec::new())); @@ -1047,6 +1107,7 @@ mod tests { let mut tool = BashTool::default(); let req = ToolRequest { commands: "for i in a b c; do echo $i; done".to_string(), + timeout_ms: None, }; let events = Arc::new(Mutex::new(Vec::new())); @@ -1083,6 +1144,7 @@ mod tests { let mut tool = BashTool::default(); let req = ToolRequest { commands: "echo start; echo end".to_string(), + timeout_ms: None, }; let events = Arc::new(Mutex::new(Vec::new())); @@ -1116,6 +1178,7 @@ mod tests { // mix of list + loop: should get 5 distinct events, no duplicates let req = ToolRequest { commands: "echo start; for i in 1 2 3; do echo $i; done; echo end".to_string(), + timeout_ms: None, }; let events = Arc::new(Mutex::new(Vec::new())); @@ -1161,4 +1224,52 @@ mod tests { assert_eq!(status.output.as_deref(), Some("error\n")); assert_eq!(status.stream.as_deref(), Some("stderr")); } + + #[tokio::test] + async fn test_tool_execute_timeout() { + let mut tool = BashTool::default(); + let req = ToolRequest { + commands: "sleep 10".to_string(), + timeout_ms: Some(100), + }; + let resp = tool.execute(req).await; + assert_eq!(resp.exit_code, 124); + assert!(resp.stderr.contains("timed out")); + assert_eq!(resp.error, Some("timeout".to_string())); + } + + #[tokio::test] + async fn test_tool_execute_no_timeout() { + let mut tool = BashTool::default(); + let req = ToolRequest { + commands: "echo fast".to_string(), + timeout_ms: Some(5000), + }; + let resp = tool.execute(req).await; + assert_eq!(resp.exit_code, 0); + assert_eq!(resp.stdout, "fast\n"); + } + + #[test] + fn test_tool_request_new() { + let req = ToolRequest::new("echo test"); + assert_eq!(req.commands, "echo test"); + assert_eq!(req.timeout_ms, None); + } + + #[test] + fn test_tool_request_deserialize_without_timeout() { + let json = r#"{"commands":"echo hello"}"#; + let req: ToolRequest = serde_json::from_str(json).unwrap(); + assert_eq!(req.commands, "echo hello"); + assert_eq!(req.timeout_ms, None); + } + + #[test] + fn test_tool_request_deserialize_with_timeout() { + let json = r#"{"commands":"echo hello","timeout_ms":5000}"#; + let req: ToolRequest = serde_json::from_str(json).unwrap(); + assert_eq!(req.commands, "echo hello"); + assert_eq!(req.timeout_ms, Some(5000)); + } } diff --git a/specs/009-tool-contract.md b/specs/009-tool-contract.md index 69e6f8e9..9d4e6974 100644 --- a/specs/009-tool-contract.md +++ b/specs/009-tool-contract.md @@ -157,7 +157,8 @@ CONFIGURATION ```rust pub struct ToolRequest { - pub commands: String, // Like bash -c "commands" + pub commands: String, // Like bash -c "commands" + pub timeout_ms: Option, // Per-call timeout (exit 124 on expiry) } pub struct ToolResponse { @@ -183,6 +184,7 @@ let mut tool = BashTool::builder() let response = tool.execute(ToolRequest { commands: "echo hello".to_string(), + timeout_ms: None, }).await; ``` @@ -254,9 +256,9 @@ Allows multiple tool implementations to share the same interface. Future tools ( Aligns with `bash -c "commands"` semantics. Clearer that it's inline commands, not a script file. -### Why no `timeout_ms`? +### Why `timeout_ms` on `ToolRequest`? -Use `timeout` builtin in commands: `timeout 5 long_running_cmd`. Keeps the API simple. +Per-call timeouts let orchestrating systems (LangChain, CrewAI) enforce limits without wrapping every command in `timeout`. Returns exit code 124 (matching bash `timeout` convention). The `timeout` builtin still works for per-command granularity. ### Why man-page format for `help()`?