diff --git a/sdk/python/ironclaw_tools.py b/sdk/python/ironclaw_tools.py index fa6d535d..18a66418 100644 --- a/sdk/python/ironclaw_tools.py +++ b/sdk/python/ironclaw_tools.py @@ -84,7 +84,7 @@ def call_tool(name, params=None, timeout_secs=None): ) try: - with urllib.request.urlopen(req, timeout=max(timeout_secs or 60, 60) + 5) as resp: + with urllib.request.urlopen(req, timeout=(timeout_secs or 60) + 5) as resp: result = json.loads(resp.read().decode("utf-8")) except urllib.error.HTTPError as e: body_text = e.read().decode("utf-8", errors="replace") if e.fp else "" diff --git a/src/tools/builtin/mod.rs b/src/tools/builtin/mod.rs index e06a7d10..6570a261 100644 --- a/src/tools/builtin/mod.rs +++ b/src/tools/builtin/mod.rs @@ -32,12 +32,12 @@ pub use job::{ pub use json::JsonTool; pub use memory::{MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool}; pub use message::MessageTool; +pub use ptc_script::PtcScriptTool; pub use restart::RestartTool; pub use routine::{ EventEmitTool, RoutineCreateTool, RoutineDeleteTool, RoutineFireTool, RoutineHistoryTool, RoutineListTool, RoutineUpdateTool, }; -pub use ptc_script::PtcScriptTool; pub use secrets_tools::{SecretDeleteTool, SecretListTool}; pub use shell::ShellTool; pub use skill_tools::{SkillInstallTool, SkillListTool, SkillRemoveTool, SkillSearchTool}; diff --git a/src/tools/builtin/ptc_script.rs b/src/tools/builtin/ptc_script.rs index c09a018b..205ef603 100644 --- a/src/tools/builtin/ptc_script.rs +++ b/src/tools/builtin/ptc_script.rs @@ -13,7 +13,9 @@ use tokio::io::AsyncReadExt; use tokio::process::Command; use crate::context::JobContext; -use crate::tools::tool::{ApprovalRequirement, Tool, ToolDomain, ToolError, ToolOutput, require_str}; +use crate::tools::tool::{ + ApprovalRequirement, Tool, ToolDomain, ToolError, ToolOutput, require_str, +}; /// Maximum output size before truncation (64KB). const MAX_OUTPUT_SIZE: usize = 64 * 1024; @@ -26,11 +28,23 @@ const MAX_TIMEOUT_SECS: u64 = 300; /// Environment variables safe to forward to the Python subprocess. const SAFE_ENV_VARS: &[&str] = &[ - "PATH", "HOME", "USER", "LOGNAME", "SHELL", "TERM", - "LANG", "LC_ALL", "LC_CTYPE", - "PWD", "TMPDIR", "TMP", "TEMP", - "CARGO_HOME", "RUSTUP_HOME", - "NODE_PATH", "NPM_CONFIG_PREFIX", + "PATH", + "HOME", + "USER", + "LOGNAME", + "SHELL", + "TERM", + "LANG", + "LC_ALL", + "LC_CTYPE", + "PWD", + "TMPDIR", + "TMP", + "TEMP", + "CARGO_HOME", + "RUSTUP_HOME", + "NODE_PATH", + "NPM_CONFIG_PREFIX", ]; /// PTC environment variables required by the ironclaw_tools SDK. @@ -84,7 +98,8 @@ impl PtcScriptTool { /// Build the full Python program from user script + preamble/postamble. fn build_program(script: &str) -> String { - let mut program = String::with_capacity(PREAMBLE.len() + script.len() + POSTAMBLE.len() + 256); + let mut program = + String::with_capacity(PREAMBLE.len() + script.len() + POSTAMBLE.len() + 256); program.push_str(PREAMBLE); // Indent user script into the try: block diff --git a/src/tools/executor.rs b/src/tools/executor.rs index c0a866d3..667edac7 100644 --- a/src/tools/executor.rs +++ b/src/tools/executor.rs @@ -109,14 +109,7 @@ impl ToolExecutor { // Determine timeout: caller override -> tool's own timeout -> default, // capped at MAX_TIMEOUT_SECS. let timeout = timeout_override - .unwrap_or_else(|| { - let tool_timeout = tool.execution_timeout(); - if tool_timeout > Duration::from_secs(MAX_TIMEOUT_SECS) { - self.default_timeout - } else { - tool_timeout - } - }) + .unwrap_or_else(|| tool.execution_timeout()) .min(Duration::from_secs(MAX_TIMEOUT_SECS)); // Execute with timeout @@ -127,12 +120,10 @@ impl ToolExecutor { timeout, })? .map_err(|e| match e { - crate::tools::ToolError::InvalidParameters(reason) => { - PtcError::InvalidParameters { - name: tool_name.to_string(), - reason, - } - } + crate::tools::ToolError::InvalidParameters(reason) => PtcError::InvalidParameters { + name: tool_name.to_string(), + reason, + }, crate::tools::ToolError::RateLimited(_) => PtcError::RateLimited { name: tool_name.to_string(), }, @@ -159,9 +150,7 @@ impl ToolExecutor { // Sanitize output if the tool requires it let (output, was_sanitized) = if tool.requires_sanitization() { let sanitized = self.safety.sanitize_tool_output(tool_name, &raw_output); - if sanitized.was_modified - && sanitized.content.starts_with("[Output blocked") - { + if sanitized.was_modified && sanitized.content.starts_with("[Output blocked") { return Err(PtcError::SafetyBlocked { reason: sanitized.content, }); @@ -249,18 +238,17 @@ mod tests { let ctx = JobContext::new("test", "test"); let result = executor - .execute( - "echo", - serde_json::json!({"message": "hello"}), - &ctx, - None, - ) + .execute("echo", serde_json::json!({"message": "hello"}), &ctx, None) .await; assert!(result.is_ok()); let ptc_result = result.as_ref().ok(); assert!(ptc_result.is_some()); - assert!(ptc_result.map(|r| r.output.contains("hello")).unwrap_or(false)); + assert!( + ptc_result + .map(|r| r.output.contains("hello")) + .unwrap_or(false) + ); } #[tokio::test] @@ -297,10 +285,7 @@ mod tests { .execute("echo", serde_json::json!({"message": "hello"}), &ctx, None) .await; - assert!(matches!( - result, - Err(PtcError::NestingDepthExceeded { .. }) - )); + assert!(matches!(result, Err(PtcError::NestingDepthExceeded { .. }))); } #[tokio::test] diff --git a/src/tools/registry.rs b/src/tools/registry.rs index d32dde9d..8ef333ad 100644 --- a/src/tools/registry.rs +++ b/src/tools/registry.rs @@ -24,7 +24,6 @@ use crate::tools::builtin::{ TimeTool, ToolActivateTool, ToolAuthTool, ToolInstallTool, ToolListTool, ToolRemoveTool, ToolSearchTool, ToolUpgradeTool, WriteFileTool, }; -use crate::tools::rate_limiter::RateLimiter; use crate::tools::executor::ToolExecutor; use crate::tools::tool::{ApprovalRequirement, Tool, ToolDomain}; use crate::tools::wasm::{ diff --git a/src/tools/wasm/wrapper.rs b/src/tools/wasm/wrapper.rs index 3607deaf..8cd770d1 100644 --- a/src/tools/wasm/wrapper.rs +++ b/src/tools/wasm/wrapper.rs @@ -34,10 +34,22 @@ use crate::tools::wasm::runtime::{EPOCH_TICK_INTERVAL, PreparedModule, WasmToolR /// The closure internally creates a tokio runtime to bridge async tool execution. /// Closure that resolves a tool call by name. The `u32` parameter is the current /// nesting depth so the executor can enforce the global depth limit across -/// WASM→executor→WASM chains. +/// WASM->executor->WASM chains. pub type ToolResolver = Arc Result + Send + Sync>; +/// RAII guard that decrements the nesting depth counter on drop, ensuring the +/// counter is restored even if the code between increment and decrement panics. +struct NestingGuard<'a> { + depth: &'a mut u32, +} + +impl Drop for NestingGuard<'_> { + fn drop(&mut self) { + *self.depth -= 1; + } +} + // Generate component model bindings from the WIT file. // // This creates: @@ -480,12 +492,14 @@ impl near::agent::host::Host for StoreData { let params: serde_json::Value = serde_json::from_str(¶ms_json) .map_err(|e| format!("Invalid tool parameters JSON: {}", e))?; - // Increment depth, call resolver with current depth, decrement on return + // Increment depth with RAII guard to ensure decrement even on panic self.tool_nesting_depth += 1; - let result = resolver(&real_name, params, self.tool_nesting_depth); - self.tool_nesting_depth -= 1; - - result + let current_depth = self.tool_nesting_depth; + let _guard = NestingGuard { + depth: &mut self.tool_nesting_depth, + }; + // _guard drops at end of scope (or on panic), decrementing depth + resolver(&real_name, params, current_depth) } fn secret_exists(&mut self, name: String) -> bool { @@ -927,11 +941,8 @@ impl Tool for WasmToolWrapper { let context_json = serde_json::to_string(ctx).ok(); // Resolve the tool executor: direct reference takes priority, then shared slot. - let resolved_executor: Option> = self - .tool_executor - .as_ref() - .cloned() - .or_else(|| { + let resolved_executor: Option> = + self.tool_executor.as_ref().cloned().or_else(|| { self.tool_executor_slot .as_ref() .and_then(|slot| slot.read().ok()) @@ -967,7 +978,8 @@ impl Tool for WasmToolWrapper { .map(|r| r.output) .map_err(|e| e.to_string()) }) - }) as Arc Result + Send + Sync> + }) + as Arc Result + Send + Sync> }); // Clone what we need for the blocking task @@ -987,10 +999,10 @@ impl Tool for WasmToolWrapper { description, schemas, credentials, - secrets_store: None, // Not needed in blocking task - oauth_refresh: None, // Already used above for pre-refresh - tool_executor: None, // Resolver closure captures the executor - tool_executor_slot: None, // Resolver closure captures the executor + secrets_store: None, // Not needed in blocking task + oauth_refresh: None, // Already used above for pre-refresh + tool_executor: None, // Resolver closure captures the executor + tool_executor_slot: None, // Resolver closure captures the executor }; tokio::task::spawn_blocking(move || { @@ -1522,10 +1534,9 @@ mod tests { TEST_GOOGLE_OAUTH_TOKEN, TEST_OAUTH_CLIENT_ID, TEST_OAUTH_CLIENT_SECRET, test_secrets_store, }; - use crate::tools::tool::Tool; + use crate::tools::tool::{Tool, ToolError, ToolOutput}; use crate::tools::wasm::capabilities::Capabilities; use crate::tools::wasm::runtime::{WasmRuntimeConfig, WasmToolRuntime}; - use crate::tools::tool::{ToolError, ToolOutput}; use super::WasmToolWrapper; @@ -2403,8 +2414,8 @@ mod tests { let capabilities = Capabilities::default().with_tool_invoke(aliases); // Create wrapper with executor - let wrapper = WasmToolWrapper::new(runtime, prepared, capabilities) - .with_tool_executor(executor); + let wrapper = + WasmToolWrapper::new(runtime, prepared, capabilities).with_tool_executor(executor); // Execute let ctx = crate::context::JobContext::new("test", "WASM PTC test"); @@ -2461,8 +2472,8 @@ mod tests { aliases.insert("other_alias".to_string(), "echo".to_string()); let capabilities = Capabilities::default().with_tool_invoke(aliases); - let wrapper = WasmToolWrapper::new(runtime, prepared, capabilities) - .with_tool_executor(executor); + let wrapper = + WasmToolWrapper::new(runtime, prepared, capabilities).with_tool_executor(executor); let ctx = crate::context::JobContext::new("test", "WASM PTC test"); let result: Result = wrapper @@ -2512,15 +2523,18 @@ mod tests { // No tool_invoke capability at all let capabilities = Capabilities::default(); - let wrapper = WasmToolWrapper::new(runtime, prepared, capabilities) - .with_tool_executor(executor); + let wrapper = + WasmToolWrapper::new(runtime, prepared, capabilities).with_tool_executor(executor); let ctx = crate::context::JobContext::new("test", "WASM PTC test"); let result: Result = wrapper .execute(serde_json::json!({"message": "hello"}), &ctx) .await; - assert!(result.is_err(), "Should fail when no tool_invoke capability"); + assert!( + result.is_err(), + "Should fail when no tool_invoke capability" + ); let err_msg = format!("{:?}", result.unwrap_err()); assert!( err_msg.contains("not granted") || err_msg.contains("capability"),