mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-31 16:49:34 +00:00
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<std::sync::RwLock> 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 <[email protected]>
This commit is contained in:
@@ -55,6 +55,9 @@ RUN npm install -g @anthropic-ai/claude-code@latest
|
|||||||
# Copy the binary
|
# Copy the binary
|
||||||
COPY --from=builder /build/target/release/ironclaw /usr/local/bin/ironclaw
|
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)
|
# Create non-root user (UID 1000 matches the orchestrator's container config)
|
||||||
RUN useradd -m -u 1000 -s /bin/bash sandbox \
|
RUN useradd -m -u 1000 -s /bin/bash sandbox \
|
||||||
&& mkdir -p /workspace \
|
&& mkdir -p /workspace \
|
||||||
|
|||||||
@@ -471,11 +471,13 @@ async fn tool_call_handler(
|
|||||||
);
|
);
|
||||||
|
|
||||||
// Build a minimal JobContext for the tool execution
|
// Build a minimal JobContext for the tool execution
|
||||||
let ctx = JobContext::with_user(
|
let mut ctx = JobContext::with_user(
|
||||||
state.user_id.clone(),
|
state.user_id.clone(),
|
||||||
format!("PTC call: {}", req.tool_name),
|
format!("PTC call: {}", req.tool_name),
|
||||||
format!("Programmatic tool call from job {}", job_id),
|
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
|
// Emit tool_use SSE event
|
||||||
if let Some(ref tx) = state.job_event_tx {
|
if let Some(ref tx) = state.job_event_tx {
|
||||||
|
|||||||
@@ -136,6 +136,10 @@ pub async fn setup_orchestrator(
|
|||||||
std::time::Duration::from_secs(60),
|
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 {
|
let orchestrator_state = api::OrchestratorState {
|
||||||
llm: Arc::clone(llm),
|
llm: Arc::clone(llm),
|
||||||
job_manager: Arc::clone(&jm),
|
job_manager: Arc::clone(&jm),
|
||||||
|
|||||||
@@ -9,6 +9,7 @@ mod json;
|
|||||||
mod memory;
|
mod memory;
|
||||||
mod message;
|
mod message;
|
||||||
pub mod path_utils;
|
pub mod path_utils;
|
||||||
|
pub mod ptc_script;
|
||||||
mod restart;
|
mod restart;
|
||||||
pub mod routine;
|
pub mod routine;
|
||||||
pub mod secrets_tools;
|
pub mod secrets_tools;
|
||||||
@@ -36,6 +37,7 @@ pub use routine::{
|
|||||||
EventEmitTool, RoutineCreateTool, RoutineDeleteTool, RoutineFireTool, RoutineHistoryTool,
|
EventEmitTool, RoutineCreateTool, RoutineDeleteTool, RoutineFireTool, RoutineHistoryTool,
|
||||||
RoutineListTool, RoutineUpdateTool,
|
RoutineListTool, RoutineUpdateTool,
|
||||||
};
|
};
|
||||||
|
pub use ptc_script::PtcScriptTool;
|
||||||
pub use secrets_tools::{SecretDeleteTool, SecretListTool};
|
pub use secrets_tools::{SecretDeleteTool, SecretListTool};
|
||||||
pub use shell::ShellTool;
|
pub use shell::ShellTool;
|
||||||
pub use skill_tools::{SkillInstallTool, SkillListTool, SkillRemoveTool, SkillSearchTool};
|
pub use skill_tools::{SkillInstallTool, SkillListTool, SkillRemoveTool, SkillSearchTool};
|
||||||
|
|||||||
@@ -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<ToolOutput, ToolError> {
|
||||||
|
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());
|
||||||
|
}
|
||||||
|
}
|
||||||
+65
-16
@@ -19,10 +19,10 @@ use crate::tools::builder::{
|
|||||||
use crate::tools::builtin::{
|
use crate::tools::builtin::{
|
||||||
ApplyPatchTool, CancelJobTool, CreateJobTool, EchoTool, ExtensionInfoTool, HttpTool,
|
ApplyPatchTool, CancelJobTool, CreateJobTool, EchoTool, ExtensionInfoTool, HttpTool,
|
||||||
JobEventsTool, JobPromptTool, JobStatusTool, JsonTool, ListDirTool, ListJobsTool,
|
JobEventsTool, JobPromptTool, JobStatusTool, JsonTool, ListDirTool, ListJobsTool,
|
||||||
MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool, PromptQueue, ReadFileTool,
|
MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool, PtcScriptTool, PromptQueue,
|
||||||
ShellTool, SkillInstallTool, SkillListTool, SkillRemoveTool, SkillSearchTool, TimeTool,
|
ReadFileTool, ShellTool, SkillInstallTool, SkillListTool, SkillRemoveTool, SkillSearchTool,
|
||||||
ToolActivateTool, ToolAuthTool, ToolInstallTool, ToolListTool, ToolRemoveTool, ToolSearchTool,
|
TimeTool, ToolActivateTool, ToolAuthTool, ToolInstallTool, ToolListTool, ToolRemoveTool,
|
||||||
ToolUpgradeTool, WriteFileTool,
|
ToolSearchTool, ToolUpgradeTool, WriteFileTool,
|
||||||
};
|
};
|
||||||
use crate::tools::rate_limiter::RateLimiter;
|
use crate::tools::rate_limiter::RateLimiter;
|
||||||
use crate::tools::executor::ToolExecutor;
|
use crate::tools::executor::ToolExecutor;
|
||||||
@@ -79,6 +79,7 @@ const PROTECTED_TOOL_NAMES: &[&str] = &[
|
|||||||
"image_edit",
|
"image_edit",
|
||||||
"image_analyze",
|
"image_analyze",
|
||||||
"tool_info",
|
"tool_info",
|
||||||
|
"ptc_script",
|
||||||
];
|
];
|
||||||
|
|
||||||
/// Registry of available tools.
|
/// Registry of available tools.
|
||||||
@@ -94,8 +95,14 @@ pub struct ToolRegistry {
|
|||||||
rate_limiter: RateLimiter,
|
rate_limiter: RateLimiter,
|
||||||
/// Reference to the message tool for setting context per-turn.
|
/// Reference to the message tool for setting context per-turn.
|
||||||
message_tool: RwLock<Option<Arc<crate::tools::builtin::MessageTool>>>,
|
message_tool: RwLock<Option<Arc<crate::tools::builtin::MessageTool>>>,
|
||||||
/// Tool executor for injecting into WASM tools (enables PTC via tool_invoke).
|
/// Shared slot for the tool executor (enables PTC via tool_invoke).
|
||||||
tool_executor: RwLock<Option<Arc<ToolExecutor>>>,
|
///
|
||||||
|
/// 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<std::sync::RwLock<Option<Arc<ToolExecutor>>>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl ToolRegistry {
|
impl ToolRegistry {
|
||||||
@@ -117,7 +124,7 @@ impl ToolRegistry {
|
|||||||
secrets_store: None,
|
secrets_store: None,
|
||||||
rate_limiter: RateLimiter::new(),
|
rate_limiter: RateLimiter::new(),
|
||||||
message_tool: RwLock::new(None),
|
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).
|
/// Set the tool executor for programmatic tool calling (PTC).
|
||||||
///
|
///
|
||||||
/// When set, WASM tools registered after this call will have `tool_invoke`
|
/// Writes the executor into the shared slot so all WASM tools --
|
||||||
/// enabled, allowing them to call other tools synchronously.
|
/// including those registered before this call -- can resolve it
|
||||||
pub async fn set_tool_executor(&self, executor: Arc<ToolExecutor>) {
|
/// lazily at execution time.
|
||||||
*self.tool_executor.write().await = Some(executor);
|
pub fn set_tool_executor(&self, executor: Arc<ToolExecutor>) {
|
||||||
|
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<std::sync::RwLock<Option<Arc<ToolExecutor>>>> {
|
||||||
|
Arc::clone(&self.tool_executor_slot)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Register a tool. Rejects dynamic tools that try to shadow a protected built-in name.
|
/// 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(WriteFileTool::new()));
|
||||||
self.register_sync(Arc::new(ListDirTool::new()));
|
self.register_sync(Arc::new(ListDirTool::new()));
|
||||||
self.register_sync(Arc::new(ApplyPatchTool::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.
|
/// Register memory tools with a workspace.
|
||||||
@@ -671,10 +690,10 @@ impl ToolRegistry {
|
|||||||
wrapper = wrapper.with_oauth_refresh(oauth);
|
wrapper = wrapper.with_oauth_refresh(oauth);
|
||||||
}
|
}
|
||||||
|
|
||||||
// Inject tool executor for PTC if available
|
// Inject shared tool executor slot for PTC (lazy resolution).
|
||||||
if let Some(executor) = self.tool_executor.read().await.as_ref() {
|
// The WASM wrapper reads from this slot at execution time, so the
|
||||||
wrapper = wrapper.with_tool_executor(Arc::clone(executor));
|
// executor can be set after tool registration.
|
||||||
}
|
wrapper = wrapper.with_tool_executor_slot(Arc::clone(&self.tool_executor_slot));
|
||||||
|
|
||||||
// Register the tool
|
// Register the tool
|
||||||
self.register(Arc::new(wrapper)).await;
|
self.register(Arc::new(wrapper)).await;
|
||||||
@@ -906,6 +925,36 @@ mod tests {
|
|||||||
assert!(def.parameters.get("extra").is_none());
|
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]
|
#[tokio::test]
|
||||||
async fn test_builtin_tool_cannot_be_shadowed() {
|
async fn test_builtin_tool_cannot_be_shadowed() {
|
||||||
let registry = ToolRegistry::new();
|
let registry = ToolRegistry::new();
|
||||||
|
|||||||
@@ -516,8 +516,11 @@ pub struct WasmToolWrapper {
|
|||||||
secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>>,
|
secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>>,
|
||||||
/// OAuth refresh configuration for auto-refreshing expired tokens.
|
/// OAuth refresh configuration for auto-refreshing expired tokens.
|
||||||
oauth_refresh: Option<OAuthRefreshConfig>,
|
oauth_refresh: Option<OAuthRefreshConfig>,
|
||||||
/// Tool executor for programmatic tool calling from within WASM tools.
|
/// Direct tool executor reference (for tests that wire it explicitly).
|
||||||
tool_executor: Option<Arc<ToolExecutor>>,
|
tool_executor: Option<Arc<ToolExecutor>>,
|
||||||
|
/// Shared slot for lazy executor resolution (production path).
|
||||||
|
/// Reads happen inside `spawn_blocking`, so this uses `std::sync::RwLock`.
|
||||||
|
tool_executor_slot: Option<Arc<std::sync::RwLock<Option<Arc<ToolExecutor>>>>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
@@ -607,6 +610,7 @@ impl WasmToolWrapper {
|
|||||||
secrets_store: None,
|
secrets_store: None,
|
||||||
oauth_refresh: None,
|
oauth_refresh: None,
|
||||||
tool_executor: None,
|
tool_executor: None,
|
||||||
|
tool_executor_slot: None,
|
||||||
};
|
};
|
||||||
wrapper.append_schema_hint_if_permissive();
|
wrapper.append_schema_hint_if_permissive();
|
||||||
wrapper
|
wrapper
|
||||||
@@ -663,15 +667,28 @@ impl WasmToolWrapper {
|
|||||||
self
|
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
|
/// When set, the WASM `tool_invoke` host function can call other
|
||||||
/// registered tools synchronously via a bridged resolver closure.
|
/// 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<ToolExecutor>) -> Self {
|
pub fn with_tool_executor(mut self, executor: Arc<ToolExecutor>) -> Self {
|
||||||
self.tool_executor = Some(executor);
|
self.tool_executor = Some(executor);
|
||||||
self
|
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<std::sync::RwLock<Option<Arc<ToolExecutor>>>>,
|
||||||
|
) -> Self {
|
||||||
|
self.tool_executor_slot = Some(slot);
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
/// Get the resource limits for this tool.
|
/// Get the resource limits for this tool.
|
||||||
pub fn limits(&self) -> &ResourceLimits {
|
pub fn limits(&self) -> &ResourceLimits {
|
||||||
&self.prepared.limits
|
&self.prepared.limits
|
||||||
@@ -909,10 +926,22 @@ impl Tool for WasmToolWrapper {
|
|||||||
// Serialize context for WASM
|
// Serialize context for WASM
|
||||||
let context_json = serde_json::to_string(ctx).ok();
|
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(|| {
|
||||||
|
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.
|
// Build a tool resolver closure if we have a tool executor.
|
||||||
// The resolver creates a single-threaded tokio runtime (same pattern
|
// The resolver creates a single-threaded tokio runtime (same pattern
|
||||||
// as http_request) to bridge the sync WASM callback to async tool execution.
|
// 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 tool_resolver: Option<ToolResolver> = resolved_executor.as_ref().map(|executor| {
|
||||||
let executor = Arc::clone(executor);
|
let executor = Arc::clone(executor);
|
||||||
let user_id = ctx.user_id.clone();
|
let user_id = ctx.user_id.clone();
|
||||||
Arc::new(move |name: &str, params: serde_json::Value, depth: u32| {
|
Arc::new(move |name: &str, params: serde_json::Value, depth: u32| {
|
||||||
@@ -958,9 +987,10 @@ impl Tool for WasmToolWrapper {
|
|||||||
description,
|
description,
|
||||||
schemas,
|
schemas,
|
||||||
credentials,
|
credentials,
|
||||||
secrets_store: None, // Not needed in blocking task
|
secrets_store: None, // Not needed in blocking task
|
||||||
oauth_refresh: None, // Already used above for pre-refresh
|
oauth_refresh: None, // Already used above for pre-refresh
|
||||||
tool_executor: None, // Resolver closure captures the executor
|
tool_executor: None, // Resolver closure captures the executor
|
||||||
|
tool_executor_slot: None, // Resolver closure captures the executor
|
||||||
};
|
};
|
||||||
|
|
||||||
tokio::task::spawn_blocking(move || {
|
tokio::task::spawn_blocking(move || {
|
||||||
|
|||||||
@@ -123,6 +123,10 @@ pub struct ToolCallRequest {
|
|||||||
pub parameters: serde_json::Value,
|
pub parameters: serde_json::Value,
|
||||||
/// Optional timeout in seconds (capped at 300s by the orchestrator).
|
/// Optional timeout in seconds (capped at 300s by the orchestrator).
|
||||||
pub timeout_secs: Option<u64>,
|
pub timeout_secs: Option<u64>,
|
||||||
|
/// 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.
|
/// Response from a programmatic tool call.
|
||||||
|
|||||||
Reference in New Issue
Block a user