mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-09-02 01:29:23 +00:00
feat(ptc): programmatic tool calling -- executor, SDK, and E2E tests
Add ToolExecutor for standalone tool dispatch used by both the
orchestrator HTTP RPC endpoint and the WASM tool_invoke host function.
Includes Python SDK for container scripts, WASM test fixture, and
comprehensive E2E test coverage across all PTC paths.
Implementation:
- ToolExecutor with timeout, nesting depth limit, safety sanitization
- Orchestrator POST /worker/{job_id}/tools/call endpoint with SSE events
- WASM tool_invoke host function with alias resolution
- Python SDK (stdlib-only) with call_tool + convenience wrappers
Tests (16 new):
- 6 orchestrator HTTP RPC tests (auth, echo, not-found, timeout, SSE, no-executor)
- 3 executor integration tests (sanitization, invalid params, sequential)
- 4 Python SDK tests (env vars, request format, HTTP error, wrappers)
- 3 WASM E2E tests (echo via alias, alias not granted, no capability)
Refs #407
Co-Authored-By: Claude Opus 4.6 <[email protected]>
This commit is contained in:
+339
-7
@@ -19,6 +19,7 @@ use wasmtime_wasi::{ResourceTable, WasiCtx, WasiCtxBuilder, WasiView};
|
||||
use crate::context::JobContext;
|
||||
use crate::safety::LeakDetector;
|
||||
use crate::secrets::SecretsStore;
|
||||
use crate::tools::ToolExecutor;
|
||||
use crate::tools::tool::{Tool, ToolError, ToolOutput};
|
||||
use crate::tools::wasm::capabilities::Capabilities;
|
||||
use crate::tools::wasm::credential_injector::{
|
||||
@@ -29,6 +30,14 @@ use crate::tools::wasm::host::{HostState, LogLevel};
|
||||
use crate::tools::wasm::limits::{ResourceLimits, WasmResourceLimiter};
|
||||
use crate::tools::wasm::runtime::{EPOCH_TICK_INTERVAL, PreparedModule, WasmToolRuntime};
|
||||
|
||||
/// Synchronous tool resolver callable from within a WASM host function.
|
||||
/// 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.
|
||||
pub type ToolResolver =
|
||||
Arc<dyn Fn(&str, serde_json::Value, u32) -> Result<String, String> + Send + Sync>;
|
||||
|
||||
// Generate component model bindings from the WIT file.
|
||||
//
|
||||
// This creates:
|
||||
@@ -99,6 +108,11 @@ struct StoreData {
|
||||
/// Dedicated tokio runtime for HTTP requests, lazily initialized.
|
||||
/// Reused across multiple `http_request` calls within one execution.
|
||||
http_runtime: Option<tokio::runtime::Runtime>,
|
||||
/// Optional tool resolver for programmatic tool calling (PTC).
|
||||
/// When set, WASM tools can invoke other tools via the `tool_invoke` host function.
|
||||
tool_resolver: Option<ToolResolver>,
|
||||
/// Current nesting depth for tool_invoke calls. Prevents infinite recursion.
|
||||
tool_nesting_depth: u32,
|
||||
}
|
||||
|
||||
impl StoreData {
|
||||
@@ -107,6 +121,7 @@ impl StoreData {
|
||||
capabilities: Capabilities,
|
||||
credentials: HashMap<String, String>,
|
||||
host_credentials: Vec<ResolvedHostCredential>,
|
||||
tool_resolver: Option<ToolResolver>,
|
||||
) -> Self {
|
||||
// Minimal WASI context: no filesystem, no env vars (security)
|
||||
let wasi = WasiCtxBuilder::new().build();
|
||||
@@ -119,6 +134,8 @@ impl StoreData {
|
||||
credentials,
|
||||
host_credentials,
|
||||
http_runtime: None,
|
||||
tool_resolver,
|
||||
tool_nesting_depth: 0,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -438,14 +455,37 @@ impl near::agent::host::Host for StoreData {
|
||||
result.map_err(|e| self.redact_credentials(&e))
|
||||
}
|
||||
|
||||
fn tool_invoke(&mut self, alias: String, _params_json: String) -> Result<String, String> {
|
||||
fn tool_invoke(&mut self, alias: String, params_json: String) -> Result<String, String> {
|
||||
use crate::tools::executor::MAX_NESTING_DEPTH;
|
||||
|
||||
// Validate capability and resolve alias
|
||||
let _real_name = self.host_state.check_tool_invoke_allowed(&alias)?;
|
||||
let real_name = self.host_state.check_tool_invoke_allowed(&alias)?;
|
||||
self.host_state.record_tool_invoke()?;
|
||||
|
||||
// Tool invocation requires async context and access to the tool registry,
|
||||
// which aren't available inside a synchronous WASM callback.
|
||||
Err("Tool invocation from WASM tools is not yet supported".to_string())
|
||||
// Check nesting depth
|
||||
if self.tool_nesting_depth >= MAX_NESTING_DEPTH {
|
||||
return Err(format!(
|
||||
"Tool invoke nesting depth exceeded (max {})",
|
||||
MAX_NESTING_DEPTH
|
||||
));
|
||||
}
|
||||
|
||||
// Get the resolver
|
||||
let resolver = self
|
||||
.tool_resolver
|
||||
.as_ref()
|
||||
.ok_or("Tool invocation not available: no tool executor configured")?;
|
||||
|
||||
// Parse parameters
|
||||
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
|
||||
self.tool_nesting_depth += 1;
|
||||
let result = resolver(&real_name, params, self.tool_nesting_depth);
|
||||
self.tool_nesting_depth -= 1;
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
fn secret_exists(&mut self, name: String) -> bool {
|
||||
@@ -476,6 +516,8 @@ pub struct WasmToolWrapper {
|
||||
secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>>,
|
||||
/// OAuth refresh configuration for auto-refreshing expired tokens.
|
||||
oauth_refresh: Option<OAuthRefreshConfig>,
|
||||
/// Tool executor for programmatic tool calling from within WASM tools.
|
||||
tool_executor: Option<Arc<ToolExecutor>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -564,7 +606,10 @@ impl WasmToolWrapper {
|
||||
credentials: HashMap::new(),
|
||||
secrets_store: None,
|
||||
oauth_refresh: None,
|
||||
}
|
||||
tool_executor: None,
|
||||
};
|
||||
wrapper.append_schema_hint_if_permissive();
|
||||
wrapper
|
||||
}
|
||||
|
||||
/// Override the tool description.
|
||||
@@ -618,6 +663,15 @@ impl WasmToolWrapper {
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the tool executor for programmatic tool calling.
|
||||
///
|
||||
/// When set, the WASM `tool_invoke` host function can call other
|
||||
/// registered tools synchronously via a bridged resolver closure.
|
||||
pub fn with_tool_executor(mut self, executor: Arc<ToolExecutor>) -> Self {
|
||||
self.tool_executor = Some(executor);
|
||||
self
|
||||
}
|
||||
|
||||
/// Get the resource limits for this tool.
|
||||
pub fn limits(&self) -> &ResourceLimits {
|
||||
&self.prepared.limits
|
||||
@@ -646,6 +700,7 @@ impl WasmToolWrapper {
|
||||
params: serde_json::Value,
|
||||
context_json: Option<String>,
|
||||
host_credentials: Vec<ResolvedHostCredential>,
|
||||
tool_resolver: Option<ToolResolver>,
|
||||
) -> Result<(String, Vec<crate::tools::wasm::host::LogEntry>), WasmError> {
|
||||
let engine = self.runtime.engine();
|
||||
let limits = &self.prepared.limits;
|
||||
@@ -656,6 +711,7 @@ impl WasmToolWrapper {
|
||||
self.capabilities.clone(),
|
||||
self.credentials.clone(),
|
||||
host_credentials,
|
||||
tool_resolver,
|
||||
);
|
||||
let mut store = Store::new(engine, store_data);
|
||||
|
||||
@@ -853,6 +909,38 @@ impl Tool for WasmToolWrapper {
|
||||
// Serialize context for WASM
|
||||
let context_json = serde_json::to_string(ctx).ok();
|
||||
|
||||
// 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<ToolResolver> = self.tool_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| {
|
||||
let executor = Arc::clone(&executor);
|
||||
let name = name.to_string();
|
||||
let mut ctx = JobContext::with_user(
|
||||
user_id.clone(),
|
||||
format!("WASM PTC: {}", name),
|
||||
"Programmatic tool call from WASM tool".to_string(),
|
||||
);
|
||||
// Propagate depth so the executor enforces the global limit
|
||||
ctx.tool_nesting_depth = depth;
|
||||
|
||||
let rt = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.map_err(|e| format!("Failed to create runtime: {}", e))?;
|
||||
|
||||
rt.block_on(async {
|
||||
executor
|
||||
.execute(&name, params, &ctx, None)
|
||||
.await
|
||||
.map(|r| r.output)
|
||||
.map_err(|e| e.to_string())
|
||||
})
|
||||
}) as Arc<dyn Fn(&str, serde_json::Value, u32) -> Result<String, String> + Send + Sync>
|
||||
});
|
||||
|
||||
// Clone what we need for the blocking task
|
||||
let runtime = Arc::clone(&self.runtime);
|
||||
let prepared = Arc::clone(&self.prepared);
|
||||
@@ -872,10 +960,11 @@ impl Tool for WasmToolWrapper {
|
||||
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
|
||||
};
|
||||
|
||||
tokio::task::spawn_blocking(move || {
|
||||
wrapper.execute_sync(params, context_json, host_credentials)
|
||||
wrapper.execute_sync(params, context_json, host_credentials, tool_resolver)
|
||||
})
|
||||
.await
|
||||
.map_err(|e| WasmError::ExecutionPanicked(e.to_string()))?
|
||||
@@ -1386,6 +1475,7 @@ fn build_tool_usage_hint(tool_name: &str, schema: &serde_json::Value) -> String
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use async_trait::async_trait;
|
||||
@@ -1405,6 +1495,9 @@ mod tests {
|
||||
use crate::tools::tool::Tool;
|
||||
use crate::tools::wasm::capabilities::Capabilities;
|
||||
use crate::tools::wasm::runtime::{WasmRuntimeConfig, WasmToolRuntime};
|
||||
use crate::tools::tool::{ToolError, ToolOutput};
|
||||
|
||||
use super::WasmToolWrapper;
|
||||
|
||||
struct RecordingSecretsStore {
|
||||
inner: InMemorySecretsStore,
|
||||
@@ -1633,6 +1726,7 @@ mod tests {
|
||||
Capabilities::default(),
|
||||
HashMap::new(),
|
||||
host_credentials,
|
||||
None,
|
||||
);
|
||||
|
||||
// Should inject for matching host
|
||||
@@ -1672,6 +1766,7 @@ mod tests {
|
||||
Capabilities::default(),
|
||||
HashMap::new(),
|
||||
host_credentials,
|
||||
None,
|
||||
);
|
||||
|
||||
let mut headers = HashMap::new();
|
||||
@@ -1698,6 +1793,7 @@ mod tests {
|
||||
Capabilities::default(),
|
||||
HashMap::new(),
|
||||
host_credentials,
|
||||
None,
|
||||
);
|
||||
|
||||
let text = "Error: request to https://api.example.com?key=super-secret-token failed";
|
||||
@@ -2184,6 +2280,242 @@ mod tests {
|
||||
assert!(hint.contains("native JSON arrays/objects")); // safety: test-only assertion
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_coerce_params_already_correct_type() {
|
||||
let schema = serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"count": { "type": "number" }
|
||||
}
|
||||
});
|
||||
let params = serde_json::json!({"count": 5});
|
||||
let result = super::coerce_params_to_schema(params, &schema);
|
||||
assert_eq!(result["count"], serde_json::json!(5));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_coerce_params_invalid_string_not_coerced() {
|
||||
let schema = serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"count": { "type": "number" }
|
||||
}
|
||||
});
|
||||
let params = serde_json::json!({"count": "not-a-number"});
|
||||
let result = super::coerce_params_to_schema(params, &schema);
|
||||
// Should remain as string since it can't be parsed
|
||||
assert_eq!(result["count"], serde_json::json!("not-a-number"));
|
||||
}
|
||||
|
||||
// === Programmatic Tool Calling (PTC) integration tests ===
|
||||
//
|
||||
// These tests require the test-ptc WASM binary to be pre-built:
|
||||
// cargo build --target wasm32-wasip2 --release --manifest-path tools-src/test-ptc/Cargo.toml
|
||||
|
||||
use crate::config::SafetyConfig;
|
||||
use crate::tools::executor::ToolExecutor;
|
||||
use crate::tools::tool::Tool;
|
||||
|
||||
fn wasm_binary_path() -> std::path::PathBuf {
|
||||
let manifest_dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR"));
|
||||
manifest_dir.join("tools-src/test-ptc/target/wasm32-wasip2/release/test_ptc_tool.wasm")
|
||||
}
|
||||
|
||||
fn load_wasm_binary() -> Option<Vec<u8>> {
|
||||
let path = wasm_binary_path();
|
||||
if !path.exists() {
|
||||
eprintln!(
|
||||
"WASM test binary not found at {:?}. Build with: \
|
||||
cargo build --target wasm32-wasip2 --release --manifest-path tools-src/test-ptc/Cargo.toml",
|
||||
path
|
||||
);
|
||||
return None;
|
||||
}
|
||||
Some(std::fs::read(&path).expect("failed to read WASM binary"))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore]
|
||||
async fn test_wasm_tool_invoke_echo() {
|
||||
let wasm_bytes = match load_wasm_binary() {
|
||||
Some(b) => b,
|
||||
None => return, // Skip if binary not built
|
||||
};
|
||||
|
||||
// Set up runtime
|
||||
let runtime = Arc::new(
|
||||
WasmToolRuntime::new(WasmRuntimeConfig::default())
|
||||
.expect("failed to create WASM runtime"),
|
||||
);
|
||||
|
||||
// Set up tool registry with echo
|
||||
let tools = Arc::new(crate::tools::registry::ToolRegistry::new());
|
||||
tools.register_builtin_tools();
|
||||
let safety = Arc::new(crate::safety::SafetyLayer::new(&SafetyConfig {
|
||||
max_output_length: 100_000,
|
||||
injection_check_enabled: true,
|
||||
}));
|
||||
let executor = Arc::new(ToolExecutor::new(
|
||||
tools,
|
||||
safety,
|
||||
std::time::Duration::from_secs(60),
|
||||
));
|
||||
|
||||
// Prepare WASM module
|
||||
let prepared = runtime
|
||||
.prepare("test_ptc", &wasm_bytes, None)
|
||||
.await
|
||||
.expect("failed to prepare WASM module");
|
||||
|
||||
// Build capabilities with echo_alias -> echo
|
||||
let mut aliases = HashMap::new();
|
||||
aliases.insert("echo_alias".to_string(), "echo".to_string());
|
||||
let capabilities = Capabilities::default().with_tool_invoke(aliases);
|
||||
|
||||
// Create wrapper with executor
|
||||
let wrapper = WasmToolWrapper::new(runtime, prepared, capabilities)
|
||||
.with_tool_executor(executor);
|
||||
|
||||
// Execute
|
||||
let ctx = crate::context::JobContext::new("test", "WASM PTC test");
|
||||
let result: Result<ToolOutput, ToolError> = wrapper
|
||||
.execute(serde_json::json!({"message": "hello"}), &ctx)
|
||||
.await;
|
||||
let result = result.expect("WASM tool execution should succeed");
|
||||
|
||||
let output = result.result.as_str().unwrap_or("");
|
||||
assert!(
|
||||
output.contains("via_wasm:"),
|
||||
"Output should contain 'via_wasm:' prefix, got: {}",
|
||||
output
|
||||
);
|
||||
assert!(
|
||||
output.contains("hello"),
|
||||
"Output should contain 'hello', got: {}",
|
||||
output
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore]
|
||||
async fn test_wasm_tool_invoke_alias_not_granted() {
|
||||
let wasm_bytes = match load_wasm_binary() {
|
||||
Some(b) => b,
|
||||
None => return,
|
||||
};
|
||||
|
||||
let runtime = Arc::new(
|
||||
WasmToolRuntime::new(WasmRuntimeConfig::default())
|
||||
.expect("failed to create WASM runtime"),
|
||||
);
|
||||
|
||||
let tools = Arc::new(crate::tools::registry::ToolRegistry::new());
|
||||
tools.register_builtin_tools();
|
||||
let safety = Arc::new(crate::safety::SafetyLayer::new(&SafetyConfig {
|
||||
max_output_length: 100_000,
|
||||
injection_check_enabled: true,
|
||||
}));
|
||||
let executor = Arc::new(ToolExecutor::new(
|
||||
tools,
|
||||
safety,
|
||||
std::time::Duration::from_secs(60),
|
||||
));
|
||||
|
||||
let prepared = runtime
|
||||
.prepare("test_ptc", &wasm_bytes, None)
|
||||
.await
|
||||
.expect("failed to prepare WASM module");
|
||||
|
||||
// Only grant a DIFFERENT alias, not "echo_alias"
|
||||
let mut aliases = HashMap::new();
|
||||
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 ctx = crate::context::JobContext::new("test", "WASM PTC test");
|
||||
let result: Result<ToolOutput, ToolError> = wrapper
|
||||
.execute(serde_json::json!({"message": "hello"}), &ctx)
|
||||
.await;
|
||||
|
||||
// Should fail because "echo_alias" is not in the aliases
|
||||
assert!(result.is_err(), "Should fail when alias not granted");
|
||||
let err_msg = format!("{:?}", result.unwrap_err());
|
||||
assert!(
|
||||
err_msg.contains("Unknown tool alias") || err_msg.contains("echo_alias"),
|
||||
"Error should mention unknown alias, got: {}",
|
||||
err_msg
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore]
|
||||
async fn test_wasm_tool_invoke_no_capability() {
|
||||
let wasm_bytes = match load_wasm_binary() {
|
||||
Some(b) => b,
|
||||
None => return,
|
||||
};
|
||||
|
||||
let runtime = Arc::new(
|
||||
WasmToolRuntime::new(WasmRuntimeConfig::default())
|
||||
.expect("failed to create WASM runtime"),
|
||||
);
|
||||
|
||||
let tools = Arc::new(crate::tools::registry::ToolRegistry::new());
|
||||
tools.register_builtin_tools();
|
||||
let safety = Arc::new(crate::safety::SafetyLayer::new(&SafetyConfig {
|
||||
max_output_length: 100_000,
|
||||
injection_check_enabled: true,
|
||||
}));
|
||||
let executor = Arc::new(ToolExecutor::new(
|
||||
tools,
|
||||
safety,
|
||||
std::time::Duration::from_secs(60),
|
||||
));
|
||||
|
||||
let prepared = runtime
|
||||
.prepare("test_ptc", &wasm_bytes, None)
|
||||
.await
|
||||
.expect("failed to prepare WASM module");
|
||||
|
||||
// No tool_invoke capability at all
|
||||
let capabilities = Capabilities::default();
|
||||
|
||||
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");
|
||||
let err_msg = format!("{:?}", result.unwrap_err());
|
||||
assert!(
|
||||
err_msg.contains("not granted") || err_msg.contains("capability"),
|
||||
"Error should mention capability not granted, got: {}",
|
||||
err_msg
|
||||
);
|
||||
}
|
||||
|
||||
/// Regression: permissive fallback schema (empty properties) must NOT coerce.
|
||||
/// This documents the bug where WASM tools with no sidecar `parameters` field
|
||||
/// got the permissive fallback, causing coercion to be a no-op and LLM-provided
|
||||
/// string integers to reach the WASM tool un-coerced.
|
||||
#[test]
|
||||
fn test_coerce_noop_with_permissive_schema() {
|
||||
let permissive = serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"additionalProperties": true
|
||||
});
|
||||
let params = serde_json::json!({"query": "test", "count": "10"});
|
||||
let result = super::coerce_params_to_schema(params, &permissive);
|
||||
// With empty properties, no coercion happens — string stays string
|
||||
assert_eq!(result["count"], serde_json::json!("10"));
|
||||
}
|
||||
|
||||
/// Regression test: leak scan must run on raw headers (before credential
|
||||
/// injection), not after. If it ran post-injection, the host-injected
|
||||
/// Slack bot token (`xoxb-...`) would trigger a Block and reject the
|
||||
|
||||
Reference in New Issue
Block a user