From ac8083bd85d415ec558275ea3051693ce08ae33b Mon Sep 17 00:00:00 2001 From: Zaki Date: Fri, 27 Feb 2026 17:58:28 -0800 Subject: [PATCH] feat(ptc): fix WASM wiring, add ptc_script tool, nesting depth, Docker SDK - Fix WASM tool_invoke production wiring: change tool_executor to a shared Arc slot with lazy resolution so WASM tools registered during build_all() can access the executor set afterward - Add nesting_depth field to ToolCallRequest and propagate it through the orchestrator's tool_call_handler into JobContext - Add ptc_script built-in tool: runs Python scripts with ironclaw_tools SDK pre-imported, env-scrubbed subprocess, structured output support - Copy Python SDK into Docker worker image at dist-packages path Co-Authored-By: Claude Opus 4.6 --- Dockerfile.worker | 3 + src/orchestrator/api.rs | 4 +- src/orchestrator/mod.rs | 4 + src/tools/builtin/mod.rs | 2 + src/tools/builtin/ptc_script.rs | 334 ++++++++++++++++++++++++++++++++ src/tools/registry.rs | 81 ++++++-- src/tools/wasm/wrapper.rs | 42 +++- src/worker/api.rs | 4 + 8 files changed, 451 insertions(+), 23 deletions(-) create mode 100644 src/tools/builtin/ptc_script.rs diff --git a/Dockerfile.worker b/Dockerfile.worker index 8f556700..c9bd86ed 100644 --- a/Dockerfile.worker +++ b/Dockerfile.worker @@ -55,6 +55,9 @@ RUN npm install -g @anthropic-ai/claude-code@latest # Copy the binary COPY --from=builder /build/target/release/ironclaw /usr/local/bin/ironclaw +# Install IronClaw Python SDK for programmatic tool calling (PTC) +COPY sdk/python/ironclaw_tools.py /usr/lib/python3/dist-packages/ironclaw_tools.py + # Create non-root user (UID 1000 matches the orchestrator's container config) RUN useradd -m -u 1000 -s /bin/bash sandbox \ && mkdir -p /workspace \ diff --git a/src/orchestrator/api.rs b/src/orchestrator/api.rs index 7c9700ff..e43965b1 100644 --- a/src/orchestrator/api.rs +++ b/src/orchestrator/api.rs @@ -471,11 +471,13 @@ async fn tool_call_handler( ); // Build a minimal JobContext for the tool execution - let ctx = JobContext::with_user( + let mut ctx = JobContext::with_user( state.user_id.clone(), format!("PTC call: {}", req.tool_name), format!("Programmatic tool call from job {}", job_id), ); + // Propagate nesting depth so the executor enforces the global limit + ctx.tool_nesting_depth = req.nesting_depth; // Emit tool_use SSE event if let Some(ref tx) = state.job_event_tx { diff --git a/src/orchestrator/mod.rs b/src/orchestrator/mod.rs index 37c7dc57..d341b767 100644 --- a/src/orchestrator/mod.rs +++ b/src/orchestrator/mod.rs @@ -136,6 +136,10 @@ pub async fn setup_orchestrator( std::time::Duration::from_secs(60), )); + // Wire the executor into the shared slot so WASM tools registered + // during build_all() can resolve it lazily at execution time. + tools.set_tool_executor(Arc::clone(&tool_executor)); + let orchestrator_state = api::OrchestratorState { llm: Arc::clone(llm), job_manager: Arc::clone(&jm), diff --git a/src/tools/builtin/mod.rs b/src/tools/builtin/mod.rs index 8ba8e57b..e06a7d10 100644 --- a/src/tools/builtin/mod.rs +++ b/src/tools/builtin/mod.rs @@ -9,6 +9,7 @@ mod json; mod memory; mod message; pub mod path_utils; +pub mod ptc_script; mod restart; pub mod routine; pub mod secrets_tools; @@ -36,6 +37,7 @@ 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 new file mode 100644 index 00000000..c09a018b --- /dev/null +++ b/src/tools/builtin/ptc_script.rs @@ -0,0 +1,334 @@ +//! PTC script tool for running multi-step Python programs that call tools. +//! +//! Wraps user-provided Python code in a preamble that imports the IronClaw +//! SDK (`ironclaw_tools`), then executes it via `python3 -c`. The script +//! runs in the same environment as the worker container and can call any +//! registered tool through the SDK's `call_tool()` function. + +use std::process::Stdio; +use std::time::Duration; + +use async_trait::async_trait; +use tokio::io::AsyncReadExt; +use tokio::process::Command; + +use crate::context::JobContext; +use crate::tools::tool::{ApprovalRequirement, Tool, ToolDomain, ToolError, ToolOutput, require_str}; + +/// Maximum output size before truncation (64KB). +const MAX_OUTPUT_SIZE: usize = 64 * 1024; + +/// Default script timeout. +const DEFAULT_TIMEOUT_SECS: u64 = 120; + +/// Maximum allowed timeout. +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", +]; + +/// PTC environment variables required by the ironclaw_tools SDK. +const PTC_ENV_VARS: &[&str] = &[ + "IRONCLAW_ORCHESTRATOR_URL", + "IRONCLAW_JOB_ID", + "IRONCLAW_WORKER_TOKEN", +]; + +/// Python preamble injected before the user's script. +const PREAMBLE: &str = r#" +import json, sys, os + +# Import IronClaw SDK +from ironclaw_tools import call_tool, shell, read_file, write_file, http_get + +# Structured output collector +_ptc_outputs = {} + +def ptc_output(key, value): + """Register a named output value for structured results.""" + _ptc_outputs[key] = value + +try: +"#; + +/// Python postamble appended after the user's script. +const POSTAMBLE: &str = r#" +except Exception as _ptc_err: + print(f"SCRIPT_ERROR: {type(_ptc_err).__name__}: {_ptc_err}", file=sys.stderr) + sys.exit(1) + +# Print structured outputs if any were registered +if _ptc_outputs: + print("\n__PTC_OUTPUTS__") + print(json.dumps(_ptc_outputs)) +"#; + +pub struct PtcScriptTool; + +impl Default for PtcScriptTool { + fn default() -> Self { + Self + } +} + +impl PtcScriptTool { + pub fn new() -> Self { + Self + } + + /// 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); + program.push_str(PREAMBLE); + + // Indent user script into the try: block + for line in script.lines() { + program.push_str(" "); + program.push_str(line); + program.push('\n'); + } + + program.push_str(POSTAMBLE); + program + } + + /// Truncate output to MAX_OUTPUT_SIZE with a truncation notice. + fn truncate_output(output: &str) -> String { + if output.len() <= MAX_OUTPUT_SIZE { + output.to_string() + } else { + format!( + "{}\n\n[Output truncated at {} bytes]", + &output[..MAX_OUTPUT_SIZE], + MAX_OUTPUT_SIZE + ) + } + } +} + +#[async_trait] +impl Tool for PtcScriptTool { + fn name(&self) -> &str { + "ptc_script" + } + + fn description(&self) -> &str { + "Execute a Python script that can call IronClaw tools programmatically. \ + The script has access to call_tool(), shell(), read_file(), write_file(), \ + and http_get() from the ironclaw_tools SDK. Use ptc_output(key, value) \ + to return structured results." + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "script": { + "type": "string", + "description": "Python script to execute. Has access to call_tool(), shell(), read_file(), write_file(), http_get(), and ptc_output()." + }, + "timeout_secs": { + "type": "integer", + "description": "Timeout in seconds (default 120, max 300).", + "default": 120, + "minimum": 1, + "maximum": 300 + } + }, + "required": ["script"] + }) + } + + async fn execute( + &self, + params: serde_json::Value, + ctx: &JobContext, + ) -> Result { + let start = std::time::Instant::now(); + + let script = require_str(¶ms, "script")?; + let timeout_secs = params + .get("timeout_secs") + .and_then(|v| v.as_u64()) + .unwrap_or(DEFAULT_TIMEOUT_SECS) + .min(MAX_TIMEOUT_SECS); + let timeout = Duration::from_secs(timeout_secs); + + let program = Self::build_program(script); + + // Build the subprocess command + let mut command = Command::new("python3"); + command.args(["-c", &program]); + + // Scrub environment -- only forward safe vars + PTC vars + extra_env + command.env_clear(); + for var in SAFE_ENV_VARS { + if let Ok(val) = std::env::var(var) { + command.env(var, val); + } + } + for var in PTC_ENV_VARS { + if let Ok(val) = std::env::var(var) { + command.env(var, val); + } + } + // Forward extra_env from JobContext (credentials fetched by worker runtime) + for (k, v) in ctx.extra_env.iter() { + command.env(k, v); + } + + command + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()); + + // Spawn and drain stdout/stderr concurrently + let mut child = command + .spawn() + .map_err(|e| ToolError::ExecutionFailed(format!("Failed to spawn python3: {}", e)))?; + + let stdout_handle = child.stdout.take(); + let stderr_handle = child.stderr.take(); + + let result = tokio::time::timeout(timeout, async { + let stdout_fut = async { + if let Some(mut out) = stdout_handle { + let mut buf = Vec::new(); + (&mut out) + .take(MAX_OUTPUT_SIZE as u64) + .read_to_end(&mut buf) + .await + .ok(); + tokio::io::copy(&mut out, &mut tokio::io::sink()).await.ok(); + String::from_utf8_lossy(&buf).to_string() + } else { + String::new() + } + }; + + let stderr_fut = async { + if let Some(mut err) = stderr_handle { + let mut buf = Vec::new(); + (&mut err) + .take(MAX_OUTPUT_SIZE as u64) + .read_to_end(&mut buf) + .await + .ok(); + tokio::io::copy(&mut err, &mut tokio::io::sink()).await.ok(); + String::from_utf8_lossy(&buf).to_string() + } else { + String::new() + } + }; + + let (stdout, stderr, wait_result) = tokio::join!(stdout_fut, stderr_fut, child.wait()); + let status = wait_result?; + Ok::<_, std::io::Error>((stdout, stderr, status.code().unwrap_or(-1))) + }) + .await; + + let duration = start.elapsed(); + + match result { + Ok(Ok((stdout, stderr, exit_code))) => { + if exit_code != 0 { + let error_msg = if stderr.is_empty() { + format!("Script exited with code {}", exit_code) + } else { + format!( + "Script exited with code {}:\n{}", + exit_code, + Self::truncate_output(&stderr) + ) + }; + return Err(ToolError::ExecutionFailed(error_msg)); + } + + // Combine output + let output = if stderr.is_empty() { + stdout + } else { + format!("{}\n\n--- stderr ---\n{}", stdout, stderr) + }; + + Ok(ToolOutput::text(Self::truncate_output(&output), duration)) + } + Ok(Err(e)) => Err(ToolError::ExecutionFailed(format!( + "Script execution failed: {}", + e + ))), + Err(_) => { + let _ = child.kill().await; + Err(ToolError::Timeout(timeout)) + } + } + } + + fn requires_sanitization(&self) -> bool { + true + } + + fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement { + ApprovalRequirement::Always + } + + fn domain(&self) -> ToolDomain { + ToolDomain::Container + } + + fn execution_timeout(&self) -> Duration { + Duration::from_secs(MAX_TIMEOUT_SECS) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_build_program_indents_script() { + let script = "x = 1\nprint(x)"; + let program = PtcScriptTool::build_program(script); + assert!(program.contains(" x = 1\n")); + assert!(program.contains(" print(x)\n")); + assert!(program.contains("from ironclaw_tools import")); + assert!(program.contains("def ptc_output(")); + } + + #[test] + fn test_build_program_empty_script() { + let program = PtcScriptTool::build_program(""); + // Empty script should still have preamble + postamble + assert!(program.contains("try:")); + assert!(program.contains("except Exception")); + } + + #[test] + fn test_truncate_output() { + let short = "hello"; + assert_eq!(PtcScriptTool::truncate_output(short), "hello"); + + let long = "x".repeat(MAX_OUTPUT_SIZE + 100); + let truncated = PtcScriptTool::truncate_output(&long); + assert!(truncated.len() < long.len()); + assert!(truncated.contains("[Output truncated")); + } + + #[test] + fn test_tool_metadata() { + let tool = PtcScriptTool::new(); + assert_eq!(tool.name(), "ptc_script"); + assert_eq!(tool.domain(), ToolDomain::Container); + assert_eq!( + tool.requires_approval(&serde_json::json!({})), + ApprovalRequirement::Always + ); + assert!(tool.requires_sanitization()); + } +} diff --git a/src/tools/registry.rs b/src/tools/registry.rs index a448c2b2..d32dde9d 100644 --- a/src/tools/registry.rs +++ b/src/tools/registry.rs @@ -19,10 +19,10 @@ use crate::tools::builder::{ use crate::tools::builtin::{ ApplyPatchTool, CancelJobTool, CreateJobTool, EchoTool, ExtensionInfoTool, HttpTool, JobEventsTool, JobPromptTool, JobStatusTool, JsonTool, ListDirTool, ListJobsTool, - MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool, PromptQueue, ReadFileTool, - ShellTool, SkillInstallTool, SkillListTool, SkillRemoveTool, SkillSearchTool, TimeTool, - ToolActivateTool, ToolAuthTool, ToolInstallTool, ToolListTool, ToolRemoveTool, ToolSearchTool, - ToolUpgradeTool, WriteFileTool, + MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool, PtcScriptTool, PromptQueue, + ReadFileTool, ShellTool, SkillInstallTool, SkillListTool, SkillRemoveTool, SkillSearchTool, + TimeTool, ToolActivateTool, ToolAuthTool, ToolInstallTool, ToolListTool, ToolRemoveTool, + ToolSearchTool, ToolUpgradeTool, WriteFileTool, }; use crate::tools::rate_limiter::RateLimiter; use crate::tools::executor::ToolExecutor; @@ -79,6 +79,7 @@ const PROTECTED_TOOL_NAMES: &[&str] = &[ "image_edit", "image_analyze", "tool_info", + "ptc_script", ]; /// Registry of available tools. @@ -94,8 +95,14 @@ pub struct ToolRegistry { rate_limiter: RateLimiter, /// Reference to the message tool for setting context per-turn. message_tool: RwLock>>, - /// Tool executor for injecting into WASM tools (enables PTC via tool_invoke). - tool_executor: RwLock>>, + /// Shared slot for the tool executor (enables PTC via tool_invoke). + /// + /// Uses `std::sync::RwLock` (not tokio) because reads happen inside + /// `spawn_blocking` closures in WASM tool execution. The slot is + /// populated lazily after `AppBuilder::build_all()` completes, so + /// WASM tools registered during startup still get access to the + /// executor when they execute later. + tool_executor_slot: Arc>>>, } impl ToolRegistry { @@ -117,7 +124,7 @@ impl ToolRegistry { secrets_store: None, rate_limiter: RateLimiter::new(), message_tool: RwLock::new(None), - tool_executor: RwLock::new(None), + tool_executor_slot: Arc::new(std::sync::RwLock::new(None)), } } @@ -144,10 +151,21 @@ impl ToolRegistry { /// Set the tool executor for programmatic tool calling (PTC). /// - /// When set, WASM tools registered after this call will have `tool_invoke` - /// enabled, allowing them to call other tools synchronously. - pub async fn set_tool_executor(&self, executor: Arc) { - *self.tool_executor.write().await = Some(executor); + /// Writes the executor into the shared slot so all WASM tools -- + /// including those registered before this call -- can resolve it + /// lazily at execution time. + pub fn set_tool_executor(&self, executor: Arc) { + if let Ok(mut guard) = self.tool_executor_slot.write() { + *guard = Some(executor); + } + } + + /// Get a clone of the shared tool executor slot. + /// + /// WASM wrappers hold this slot and read from it at execution time, + /// allowing the executor to be set after tool registration. + pub fn tool_executor_slot(&self) -> Arc>>> { + Arc::clone(&self.tool_executor_slot) } /// Register a tool. Rejects dynamic tools that try to shadow a protected built-in name. @@ -342,8 +360,9 @@ impl ToolRegistry { self.register_sync(Arc::new(WriteFileTool::new())); self.register_sync(Arc::new(ListDirTool::new())); self.register_sync(Arc::new(ApplyPatchTool::new())); + self.register_sync(Arc::new(PtcScriptTool::new())); - tracing::debug!("Registered 5 development tools"); + tracing::debug!("Registered 6 development tools"); } /// Register memory tools with a workspace. @@ -671,10 +690,10 @@ impl ToolRegistry { wrapper = wrapper.with_oauth_refresh(oauth); } - // Inject tool executor for PTC if available - if let Some(executor) = self.tool_executor.read().await.as_ref() { - wrapper = wrapper.with_tool_executor(Arc::clone(executor)); - } + // Inject shared tool executor slot for PTC (lazy resolution). + // The WASM wrapper reads from this slot at execution time, so the + // executor can be set after tool registration. + wrapper = wrapper.with_tool_executor_slot(Arc::clone(&self.tool_executor_slot)); // Register the tool self.register(Arc::new(wrapper)).await; @@ -906,6 +925,36 @@ mod tests { assert!(def.parameters.get("extra").is_none()); } + #[tokio::test] + async fn test_tool_executor_slot_lazy_resolution() { + let registry = ToolRegistry::new(); + + // Get the slot BEFORE setting the executor (simulates startup order) + let slot = registry.tool_executor_slot(); + + // Slot should be empty + assert!(slot.read().unwrap().is_none()); + + // Set the executor (simulates main.rs wiring after build_all) + let tools = Arc::new(ToolRegistry::new()); + tools.register_builtin_tools(); + let safety = Arc::new(crate::safety::SafetyLayer::new( + &crate::config::SafetyConfig { + max_output_length: 100_000, + injection_check_enabled: true, + }, + )); + let executor = Arc::new(crate::tools::ToolExecutor::new( + tools, + safety, + std::time::Duration::from_secs(60), + )); + registry.set_tool_executor(Arc::clone(&executor)); + + // Slot should now contain the executor + assert!(slot.read().unwrap().is_some()); + } + #[tokio::test] async fn test_builtin_tool_cannot_be_shadowed() { let registry = ToolRegistry::new(); diff --git a/src/tools/wasm/wrapper.rs b/src/tools/wasm/wrapper.rs index 17d85e54..3607deaf 100644 --- a/src/tools/wasm/wrapper.rs +++ b/src/tools/wasm/wrapper.rs @@ -516,8 +516,11 @@ pub struct WasmToolWrapper { secrets_store: Option>, /// OAuth refresh configuration for auto-refreshing expired tokens. oauth_refresh: Option, - /// Tool executor for programmatic tool calling from within WASM tools. + /// Direct tool executor reference (for tests that wire it explicitly). tool_executor: Option>, + /// Shared slot for lazy executor resolution (production path). + /// Reads happen inside `spawn_blocking`, so this uses `std::sync::RwLock`. + tool_executor_slot: Option>>>>, } #[derive(Debug, Clone)] @@ -607,6 +610,7 @@ impl WasmToolWrapper { secrets_store: None, oauth_refresh: None, tool_executor: None, + tool_executor_slot: None, }; wrapper.append_schema_hint_if_permissive(); wrapper @@ -663,15 +667,28 @@ impl WasmToolWrapper { self } - /// Set the tool executor for programmatic tool calling. + /// Set the tool executor for programmatic tool calling (direct reference). /// /// When set, the WASM `tool_invoke` host function can call other /// registered tools synchronously via a bridged resolver closure. + /// Prefer `with_tool_executor_slot()` for production use. pub fn with_tool_executor(mut self, executor: Arc) -> Self { self.tool_executor = Some(executor); self } + /// Set the shared tool executor slot for lazy resolution. + /// + /// The executor is read from this slot at execution time, allowing + /// it to be set after tool registration (production startup order). + pub fn with_tool_executor_slot( + mut self, + slot: Arc>>>, + ) -> Self { + self.tool_executor_slot = Some(slot); + self + } + /// Get the resource limits for this tool. pub fn limits(&self) -> &ResourceLimits { &self.prepared.limits @@ -909,10 +926,22 @@ impl Tool for WasmToolWrapper { // Serialize context for WASM 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(|| { + self.tool_executor_slot + .as_ref() + .and_then(|slot| slot.read().ok()) + .and_then(|guard| guard.clone()) + }); + // Build a tool resolver closure if we have a tool executor. // The resolver creates a single-threaded tokio runtime (same pattern // as http_request) to bridge the sync WASM callback to async tool execution. - let tool_resolver: Option = self.tool_executor.as_ref().map(|executor| { + let tool_resolver: Option = resolved_executor.as_ref().map(|executor| { let executor = Arc::clone(executor); let user_id = ctx.user_id.clone(); Arc::new(move |name: &str, params: serde_json::Value, depth: u32| { @@ -958,9 +987,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 + 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 || { diff --git a/src/worker/api.rs b/src/worker/api.rs index 420a6b05..b56fc03c 100644 --- a/src/worker/api.rs +++ b/src/worker/api.rs @@ -123,6 +123,10 @@ pub struct ToolCallRequest { pub parameters: serde_json::Value, /// Optional timeout in seconds (capped at 300s by the orchestrator). pub timeout_secs: Option, + /// Current nesting depth for tool-invokes-tool chains. + /// Defaults to 0 for top-level calls (backward compatible). + #[serde(default)] + pub nesting_depth: u32, } /// Response from a programmatic tool call.