fix: PTC timeout handling and nesting depth panic safety

- Python SDK: remove 60s minimum timeout enforcement, respect
  requested timeout with 5s network buffer
- Rust executor: cap timeout at MAX_TIMEOUT_SECS instead of
  falling back to default when exceeded
- WASM wrapper: use RAII guard for nesting depth to prevent
  leak on panic

Addresses Gemini review feedback on PR #408.

Co-Authored-By: Claude Opus 4.6 <[email protected]>
This commit is contained in:
Zaki
2026-03-21 08:12:37 +00:00
committed by Claude
co-authored by Claude Opus 4.6
parent ac8083bd85
commit cd23380a66
6 changed files with 76 additions and 63 deletions
+1 -1
View File
@@ -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 ""
+1 -1
View File
@@ -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};
+22 -7
View File
@@ -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
+13 -28
View File
@@ -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]
-1
View File
@@ -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::{
+39 -25
View File
@@ -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
/// WASMexecutorWASM chains.
/// WASM->executor->WASM chains.
pub type ToolResolver =
Arc<dyn Fn(&str, serde_json::Value, u32) -> Result<String, String> + 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(&params_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<Arc<ToolExecutor>> = self
.tool_executor
.as_ref()
.cloned()
.or_else(|| {
let resolved_executor: Option<Arc<ToolExecutor>> =
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<dyn Fn(&str, serde_json::Value, u32) -> Result<String, String> + Send + Sync>
})
as Arc<dyn Fn(&str, serde_json::Value, u32) -> Result<String, String> + 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<ToolOutput, ToolError> = 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<ToolOutput, ToolError> = 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"),