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:
Zaki
2026-03-21 08:11:31 +00:00
committed by Claude
co-authored by Claude Opus 4.6
parent 6232609080
commit 42e6650ab8
15 changed files with 1649 additions and 7 deletions
+155
View File
@@ -0,0 +1,155 @@
"""IronClaw Programmatic Tool Calling SDK for container scripts.
Thin wrapper using only Python stdlib. Reads connection details from
environment variables injected by the orchestrator:
IRONCLAW_ORCHESTRATOR_URL - Base URL of the orchestrator API
IRONCLAW_JOB_ID - UUID of the current job
IRONCLAW_WORKER_TOKEN - Bearer token scoped to this job
Usage:
from ironclaw_tools import call_tool, shell, read_file, write_file, http_get
# Call any registered tool by name
result = call_tool("echo", {"message": "hello"})
print(result) # "hello"
# Convenience wrappers
output = shell("ls -la")
content = read_file("/workspace/README.md")
write_file("/workspace/output.txt", "results here")
body = http_get("https://api.example.com/data")
"""
import json
import os
import urllib.request
import urllib.error
def _env(name):
"""Get a required environment variable."""
value = os.environ.get(name)
if not value:
raise RuntimeError(
f"Missing required environment variable: {name}. "
"This SDK must be run inside an IronClaw container."
)
return value
def _base_url():
"""Build the base URL for tool call requests."""
orchestrator = _env("IRONCLAW_ORCHESTRATOR_URL").rstrip("/")
job_id = _env("IRONCLAW_JOB_ID")
return f"{orchestrator}/worker/{job_id}"
def _token():
"""Get the bearer token."""
return _env("IRONCLAW_WORKER_TOKEN")
def call_tool(name, params=None, timeout_secs=None):
"""Call a tool on the orchestrator by name.
Args:
name: Tool name (e.g., "echo", "shell", "read_file").
params: Dictionary of parameters to pass to the tool.
timeout_secs: Optional timeout in seconds (max 300).
Returns:
Tool output as a string.
Raises:
RuntimeError: If the tool call fails.
"""
url = f"{_base_url()}/tools/call"
body = {
"tool_name": name,
"parameters": params or {},
}
if timeout_secs is not None:
body["timeout_secs"] = min(int(timeout_secs), 300)
data = json.dumps(body).encode("utf-8")
req = urllib.request.Request(
url,
data=data,
headers={
"Content-Type": "application/json",
"Authorization": f"Bearer {_token()}",
},
method="POST",
)
try:
with urllib.request.urlopen(req, timeout=max(timeout_secs or 60, 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 ""
raise RuntimeError(
f"Tool call failed: HTTP {e.code}: {body_text}"
) from None
except urllib.error.URLError as e:
raise RuntimeError(f"Connection to orchestrator failed: {e.reason}") from None
if not result.get("success"):
raise RuntimeError(f"Tool '{name}' failed: {result.get('error', 'unknown error')}")
return result.get("output", "")
def shell(command, timeout_secs=60):
"""Execute a shell command via the orchestrator.
Args:
command: Shell command string to execute.
timeout_secs: Timeout in seconds (default 60).
Returns:
Command output as a string.
"""
return call_tool("shell", {"command": command}, timeout_secs=timeout_secs)
def read_file(path):
"""Read a file via the orchestrator.
Args:
path: Absolute path to the file.
Returns:
File contents as a string.
"""
return call_tool("read_file", {"path": path})
def write_file(path, content):
"""Write a file via the orchestrator.
Args:
path: Absolute path to write to.
content: String content to write.
Returns:
Write confirmation message.
"""
return call_tool("write_file", {"path": path, "content": content})
def http_get(url, headers=None, timeout_secs=30):
"""Make an HTTP GET request via the orchestrator's HTTP tool.
Args:
url: URL to fetch.
headers: Optional dictionary of headers.
timeout_secs: Timeout in seconds (default 30).
Returns:
Response body as a string.
"""
params = {"url": url, "method": "GET"}
if headers:
params["headers"] = headers
return call_tool("http", params, timeout_secs=timeout_secs)
+148
View File
@@ -0,0 +1,148 @@
"""Tests for the IronClaw Programmatic Tool Calling Python SDK."""
import json
import os
import sys
import unittest
from unittest.mock import patch, MagicMock
import urllib.error
# Ensure ironclaw_tools is importable regardless of working directory.
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
class TestEnvMissing(unittest.TestCase):
"""Test that missing env vars produce clear errors."""
def setUp(self):
# Clear all relevant env vars
for var in ["IRONCLAW_ORCHESTRATOR_URL", "IRONCLAW_JOB_ID", "IRONCLAW_WORKER_TOKEN"]:
os.environ.pop(var, None)
def test_env_missing(self):
from ironclaw_tools import call_tool
with self.assertRaises(RuntimeError) as ctx:
call_tool("echo", {"message": "hello"})
# Should mention the missing variable
self.assertIn("IRONCLAW_ORCHESTRATOR_URL", str(ctx.exception))
class TestCallToolRequestFormat(unittest.TestCase):
"""Test that call_tool sends correctly formatted requests."""
def setUp(self):
os.environ["IRONCLAW_ORCHESTRATOR_URL"] = "http://localhost:50051"
os.environ["IRONCLAW_JOB_ID"] = "550e8400-e29b-41d4-a716-446655440000"
os.environ["IRONCLAW_WORKER_TOKEN"] = "test-token-123"
def tearDown(self):
for var in ["IRONCLAW_ORCHESTRATOR_URL", "IRONCLAW_JOB_ID", "IRONCLAW_WORKER_TOKEN"]:
os.environ.pop(var, None)
@patch("ironclaw_tools.urllib.request.urlopen")
def test_call_tool_request_format(self, mock_urlopen):
from ironclaw_tools import call_tool
# Mock successful response
mock_response = MagicMock()
mock_response.read.return_value = json.dumps({
"success": True,
"output": "hello",
"duration_ms": 5,
"was_sanitized": False,
}).encode("utf-8")
mock_response.__enter__ = lambda s: s
mock_response.__exit__ = MagicMock(return_value=False)
mock_urlopen.return_value = mock_response
result = call_tool("echo", {"message": "hello"}, timeout_secs=30)
# Verify the request was made
mock_urlopen.assert_called_once()
call_args = mock_urlopen.call_args
req = call_args[0][0] # First positional arg is the Request object
# Check URL
self.assertIn("/worker/550e8400-e29b-41d4-a716-446655440000/tools/call", req.full_url)
# Check headers
self.assertEqual(req.get_header("Content-type"), "application/json")
self.assertEqual(req.get_header("Authorization"), "Bearer test-token-123")
# Check body
body = json.loads(req.data.decode("utf-8"))
self.assertEqual(body["tool_name"], "echo")
self.assertEqual(body["parameters"], {"message": "hello"})
self.assertEqual(body["timeout_secs"], 30)
# Check return value
self.assertEqual(result, "hello")
class TestCallToolHttpError(unittest.TestCase):
"""Test HTTP error handling."""
def setUp(self):
os.environ["IRONCLAW_ORCHESTRATOR_URL"] = "http://localhost:50051"
os.environ["IRONCLAW_JOB_ID"] = "550e8400-e29b-41d4-a716-446655440000"
os.environ["IRONCLAW_WORKER_TOKEN"] = "test-token-123"
def tearDown(self):
for var in ["IRONCLAW_ORCHESTRATOR_URL", "IRONCLAW_JOB_ID", "IRONCLAW_WORKER_TOKEN"]:
os.environ.pop(var, None)
@patch("ironclaw_tools.urllib.request.urlopen")
def test_call_tool_http_error(self, mock_urlopen):
from ironclaw_tools import call_tool
mock_urlopen.side_effect = urllib.error.HTTPError(
url="http://localhost:50051/worker/test/tools/call",
code=500,
msg="Internal Server Error",
hdrs=None,
fp=None,
)
with self.assertRaises(RuntimeError) as ctx:
call_tool("echo", {"message": "hello"})
self.assertIn("500", str(ctx.exception))
class TestConvenienceWrappers(unittest.TestCase):
"""Test that convenience wrappers call call_tool correctly."""
def setUp(self):
os.environ["IRONCLAW_ORCHESTRATOR_URL"] = "http://localhost:50051"
os.environ["IRONCLAW_JOB_ID"] = "550e8400-e29b-41d4-a716-446655440000"
os.environ["IRONCLAW_WORKER_TOKEN"] = "test-token-123"
def tearDown(self):
for var in ["IRONCLAW_ORCHESTRATOR_URL", "IRONCLAW_JOB_ID", "IRONCLAW_WORKER_TOKEN"]:
os.environ.pop(var, None)
@patch("ironclaw_tools.call_tool")
def test_convenience_wrappers(self, mock_call_tool):
from ironclaw_tools import shell, read_file, write_file, http_get
mock_call_tool.return_value = "output"
# Test shell
shell("ls -la")
mock_call_tool.assert_called_with("shell", {"command": "ls -la"}, timeout_secs=60)
# Test read_file
read_file("/workspace/README.md")
mock_call_tool.assert_called_with("read_file", {"path": "/workspace/README.md"})
# Test write_file
write_file("/workspace/out.txt", "content")
mock_call_tool.assert_called_with("write_file", {"path": "/workspace/out.txt", "content": "content"})
# Test http_get
http_get("https://api.example.com/data")
mock_call_tool.assert_called_with("http", {"url": "https://api.example.com/data", "method": "GET"}, timeout_secs=30)
if __name__ == "__main__":
unittest.main()
+7
View File
@@ -196,6 +196,12 @@ pub struct JobContext {
pub tool_output_stash: Arc<tokio::sync::RwLock<HashMap<String, String>>>,
/// User's preferred timezone (IANA name, e.g. "America/New_York"). Defaults to "UTC".
pub user_timezone: String,
/// Current nesting depth for programmatic tool calling (PTC).
///
/// Tracks how deep we are in a tool-invokes-tool chain so the executor
/// can enforce MAX_NESTING_DEPTH globally, even across WASM→executor→WASM chains.
#[serde(skip)]
pub tool_nesting_depth: u32,
}
impl JobContext {
@@ -237,6 +243,7 @@ impl JobContext {
metadata: serde_json::Value::Null,
tool_output_stash: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
user_timezone: "UTC".to_string(),
tool_nesting_depth: 0,
}
}
+1
View File
@@ -134,6 +134,7 @@ impl JobStore for LibSqlBackend {
// TODO(#661): persist user_timezone in agent_jobs table so
// background/routine jobs retain the session's timezone context.
user_timezone: "UTC".to_string(),
tool_nesting_depth: 0,
}))
}
None => Ok(None),
+1
View File
@@ -258,6 +258,7 @@ impl Store {
// TODO(#661): persist user_timezone in agent_jobs table so
// background/routine jobs retain the session's timezone context.
user_timezone: "UTC".to_string(),
tool_nesting_depth: 0,
}))
}
None => Ok(None),
+2
View File
@@ -306,6 +306,8 @@ async fn async_main() -> anyhow::Result<()> {
&components.llm,
components.db.as_ref(),
components.secrets_store.as_ref(),
&components.tools,
&components.safety,
)
.await;
let container_job_manager = orch.container_job_manager;
+418
View File
@@ -15,15 +15,18 @@ use tokio::sync::{Mutex, broadcast};
use uuid::Uuid;
use crate::channels::web::types::SseEvent;
use crate::context::JobContext;
use crate::db::Database;
use crate::llm::{CompletionRequest, LlmProvider, ToolCompletionRequest};
use crate::orchestrator::auth::{TokenStore, worker_auth_middleware};
use crate::orchestrator::job_manager::ContainerJobManager;
use crate::secrets::SecretsStore;
use crate::tools::ToolExecutor;
use crate::worker::api::JobEventPayload;
use crate::worker::api::{
CompletionReport, CredentialResponse, JobDescription, ProxyCompletionRequest,
ProxyCompletionResponse, ProxyToolCompletionRequest, ProxyToolCompletionResponse, StatusUpdate,
ToolCallRequest, ToolCallResponse,
};
/// A follow-up prompt queued for a Claude Code bridge.
@@ -49,6 +52,8 @@ pub struct OrchestratorState {
pub secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>>,
/// User ID for secret lookups (single-tenant, typically "default").
pub user_id: String,
/// Tool executor for programmatic tool calling (PTC).
pub tool_executor: Option<Arc<ToolExecutor>>,
}
/// The orchestrator's internal API server.
@@ -70,6 +75,7 @@ impl OrchestratorApi {
.route("/worker/{job_id}/event", post(job_event_handler))
.route("/worker/{job_id}/prompt", get(get_prompt_handler))
.route("/worker/{job_id}/credentials", get(get_credentials_handler))
.route("/worker/{job_id}/tools/call", post(tool_call_handler))
.route_layer(axum::middleware::from_fn_with_state(
state.token_store.clone(),
worker_auth_middleware,
@@ -443,6 +449,96 @@ async fn get_credentials_handler(
))
}
/// Execute a tool programmatically on behalf of a container worker (PTC).
///
/// Builds a minimal `JobContext` from the job metadata and delegates to
/// `ToolExecutor::execute`. Emits SSE events for tool_use/tool_result so
/// the web UI can observe PTC calls.
async fn tool_call_handler(
State(state): State<OrchestratorState>,
Path(job_id): Path<Uuid>,
Json(req): Json<ToolCallRequest>,
) -> Result<Json<ToolCallResponse>, StatusCode> {
let executor = state
.tool_executor
.as_ref()
.ok_or(StatusCode::SERVICE_UNAVAILABLE)?;
tracing::debug!(
job_id = %job_id,
tool = %req.tool_name,
"PTC tool call request"
);
// Build a minimal JobContext for the tool execution
let ctx = JobContext::with_user(
state.user_id.clone(),
format!("PTC call: {}", req.tool_name),
format!("Programmatic tool call from job {}", job_id),
);
// Emit tool_use SSE event
if let Some(ref tx) = state.job_event_tx {
let _ = tx.send((
job_id,
SseEvent::JobToolUse {
job_id: job_id.to_string(),
tool_name: req.tool_name.clone(),
input: req.parameters.clone(),
},
));
}
// Determine timeout override
let timeout_override = req
.timeout_secs
.map(|s| std::time::Duration::from_secs(s.min(300)));
// Execute the tool
match executor
.execute(&req.tool_name, req.parameters, &ctx, timeout_override)
.await
{
Ok(result) => {
// Emit tool_result SSE event
if let Some(ref tx) = state.job_event_tx {
let _ = tx.send((
job_id,
SseEvent::JobToolResult {
job_id: job_id.to_string(),
tool_name: req.tool_name.clone(),
output: result.output.clone(),
},
));
}
Ok(Json(ToolCallResponse {
success: true,
output: Some(result.output),
error: None,
duration_ms: result.duration.as_millis() as u64,
was_sanitized: result.was_sanitized,
}))
}
Err(e) => {
tracing::warn!(
job_id = %job_id,
tool = %req.tool_name,
error = %e,
"PTC tool call failed"
);
Ok(Json(ToolCallResponse {
success: false,
output: None,
error: Some(e.to_string()),
duration_ms: 0,
was_sanitized: false,
}))
}
}
}
fn format_finish_reason(reason: crate::llm::FinishReason) -> String {
match reason {
crate::llm::FinishReason::Stop => "stop".to_string(),
@@ -480,6 +576,7 @@ mod tests {
store: None,
secrets_store: None,
user_id: "default".to_string(),
tool_executor: None,
}
}
@@ -709,6 +806,7 @@ mod tests {
store: None,
secrets_store: Some(secrets_store),
user_id: "default".to_string(),
tool_executor: None,
};
let router = OrchestratorApi::router(state);
@@ -744,6 +842,7 @@ mod tests {
store: None,
secrets_store: None,
user_id: "default".to_string(),
tool_executor: None,
};
let job_id = Uuid::new_v4();
@@ -799,6 +898,7 @@ mod tests {
store: None,
secrets_store: None,
user_id: "default".to_string(),
tool_executor: None,
};
let job_id = Uuid::new_v4();
@@ -847,6 +947,7 @@ mod tests {
store: None,
secrets_store: None,
user_id: "default".to_string(),
tool_executor: None,
};
let job_id = Uuid::new_v4();
@@ -926,4 +1027,321 @@ mod tests {
assert_eq!(handle.worker_iteration, 5);
assert_eq!(handle.last_worker_status.as_deref(), Some("Iteration 5"));
}
// -- Programmatic tool calling (PTC) tests --
use std::time::Duration;
use crate::config::SafetyConfig;
use crate::context::JobContext;
use crate::safety::SafetyLayer;
use crate::tools::{Tool, ToolError, ToolExecutor, ToolOutput, ToolRegistry};
/// A tool that sleeps for 10 seconds (used to test timeout enforcement).
struct SlowTool;
#[async_trait::async_trait]
impl Tool for SlowTool {
fn name(&self) -> &str {
"slow_tool"
}
fn description(&self) -> &str {
"A tool that sleeps"
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({"type": "object"})
}
async fn execute(
&self,
_params: serde_json::Value,
_ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
tokio::time::sleep(Duration::from_secs(10)).await;
Ok(ToolOutput::text("done", Duration::from_secs(10)))
}
fn requires_sanitization(&self) -> bool {
false
}
}
/// Build an `OrchestratorState` with a real `ToolExecutor` wired in.
///
/// Also returns the broadcast receiver when `with_broadcast` is true,
/// so SSE-related tests can observe emitted events.
fn test_state_with_executor(
with_broadcast: bool,
) -> (
OrchestratorState,
Option<broadcast::Receiver<(Uuid, SseEvent)>>,
) {
let tools = Arc::new(ToolRegistry::new());
tools.register_builtin_tools();
let safety = Arc::new(SafetyLayer::new(&SafetyConfig {
max_output_length: 100_000,
injection_check_enabled: true,
}));
let executor = ToolExecutor::new(Arc::clone(&tools), safety, Duration::from_secs(60));
let token_store = TokenStore::new();
let jm = ContainerJobManager::new(ContainerJobConfig::default(), token_store.clone());
let (tx, rx) = if with_broadcast {
let (tx, rx) = broadcast::channel(16);
(Some(tx), Some(rx))
} else {
(None, None)
};
let state = OrchestratorState {
llm: Arc::new(StubLlm::default()),
job_manager: Arc::new(jm),
token_store,
job_event_tx: tx,
prompt_queue: Arc::new(Mutex::new(HashMap::new())),
store: None,
secrets_store: None,
user_id: "default".to_string(),
tool_executor: Some(Arc::new(executor)),
};
(state, rx)
}
#[tokio::test]
async fn tool_call_echo_success() {
let (state, _) = test_state_with_executor(false);
let job_id = Uuid::new_v4();
let token = state.token_store.create_token(job_id).await;
let router = OrchestratorApi::router(state);
let payload = serde_json::json!({
"tool_name": "echo",
"parameters": {"message": "hello"},
});
let req = Request::builder()
.method("POST")
.uri(format!("/worker/{}/tools/call", job_id))
.header("Authorization", format!("Bearer {}", token))
.header("Content-Type", "application/json")
.body(Body::from(serde_json::to_vec(&payload).unwrap()))
.unwrap();
let resp = router.oneshot(req).await.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let body = axum::body::to_bytes(resp.into_body(), 4096).await.unwrap();
let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert_eq!(json["success"], true);
assert!(
json["output"]
.as_str()
.map(|s| s.contains("hello"))
.unwrap_or(false),
"output should contain 'hello', got: {:?}",
json["output"]
);
assert!(
json["duration_ms"].is_u64(),
"duration_ms should be present as a number"
);
}
#[tokio::test]
async fn tool_call_not_found() {
let (state, _) = test_state_with_executor(false);
let job_id = Uuid::new_v4();
let token = state.token_store.create_token(job_id).await;
let router = OrchestratorApi::router(state);
let payload = serde_json::json!({
"tool_name": "nonexistent_tool",
"parameters": {},
});
let req = Request::builder()
.method("POST")
.uri(format!("/worker/{}/tools/call", job_id))
.header("Authorization", format!("Bearer {}", token))
.header("Content-Type", "application/json")
.body(Body::from(serde_json::to_vec(&payload).unwrap()))
.unwrap();
let resp = router.oneshot(req).await.unwrap();
// Handler returns Ok(Json(...)) even on tool failure
assert_eq!(resp.status(), StatusCode::OK);
let body = axum::body::to_bytes(resp.into_body(), 4096).await.unwrap();
let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert_eq!(json["success"], false);
assert!(
json["error"]
.as_str()
.map(|s| s.to_lowercase().contains("not found"))
.unwrap_or(false),
"error should mention 'not found', got: {:?}",
json["error"]
);
}
#[tokio::test]
async fn tool_call_no_executor() {
// Use regular test_state() which has tool_executor: None
let state = test_state();
let job_id = Uuid::new_v4();
let token = state.token_store.create_token(job_id).await;
let router = OrchestratorApi::router(state);
let payload = serde_json::json!({
"tool_name": "echo",
"parameters": {"message": "hello"},
});
let req = Request::builder()
.method("POST")
.uri(format!("/worker/{}/tools/call", job_id))
.header("Authorization", format!("Bearer {}", token))
.header("Content-Type", "application/json")
.body(Body::from(serde_json::to_vec(&payload).unwrap()))
.unwrap();
let resp = router.oneshot(req).await.unwrap();
assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE);
}
#[tokio::test]
async fn tool_call_with_sse_events() {
let (state, rx) = test_state_with_executor(true);
let mut rx = rx.expect("broadcast receiver should be present");
let job_id = Uuid::new_v4();
let token = state.token_store.create_token(job_id).await;
let router = OrchestratorApi::router(state);
let payload = serde_json::json!({
"tool_name": "echo",
"parameters": {"message": "hello"},
});
let req = Request::builder()
.method("POST")
.uri(format!("/worker/{}/tools/call", job_id))
.header("Authorization", format!("Bearer {}", token))
.header("Content-Type", "application/json")
.body(Body::from(serde_json::to_vec(&payload).unwrap()))
.unwrap();
let resp = router.oneshot(req).await.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
// Collect events from broadcast channel
let mut saw_tool_use = false;
let mut saw_tool_result = false;
while let Ok((recv_id, event)) = rx.try_recv() {
assert_eq!(recv_id, job_id);
match event {
SseEvent::JobToolUse { tool_name, .. } => {
assert_eq!(tool_name, "echo");
saw_tool_use = true;
}
SseEvent::JobToolResult { tool_name, .. } => {
assert_eq!(tool_name, "echo");
saw_tool_result = true;
}
_ => {}
}
}
assert!(saw_tool_use, "should have emitted JobToolUse event");
assert!(saw_tool_result, "should have emitted JobToolResult event");
}
#[tokio::test]
async fn tool_call_with_timeout() {
// Build a registry that includes our SlowTool
let tools = Arc::new(ToolRegistry::new());
tools.register(Arc::new(SlowTool)).await;
let safety = Arc::new(SafetyLayer::new(&SafetyConfig {
max_output_length: 100_000,
injection_check_enabled: true,
}));
let executor = ToolExecutor::new(Arc::clone(&tools), safety, Duration::from_secs(60));
let token_store = TokenStore::new();
let jm = ContainerJobManager::new(ContainerJobConfig::default(), token_store.clone());
let state = OrchestratorState {
llm: Arc::new(StubLlm::default()),
job_manager: Arc::new(jm),
token_store: token_store.clone(),
job_event_tx: None,
prompt_queue: Arc::new(Mutex::new(HashMap::new())),
store: None,
secrets_store: None,
user_id: "default".to_string(),
tool_executor: Some(Arc::new(executor)),
};
let job_id = Uuid::new_v4();
let token = token_store.create_token(job_id).await;
let router = OrchestratorApi::router(state);
let payload = serde_json::json!({
"tool_name": "slow_tool",
"parameters": {},
"timeout_secs": 1,
});
let req = Request::builder()
.method("POST")
.uri(format!("/worker/{}/tools/call", job_id))
.header("Authorization", format!("Bearer {}", token))
.header("Content-Type", "application/json")
.body(Body::from(serde_json::to_vec(&payload).unwrap()))
.unwrap();
let resp = router.oneshot(req).await.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let body = axum::body::to_bytes(resp.into_body(), 4096).await.unwrap();
let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
assert_eq!(json["success"], false);
assert!(
json["error"]
.as_str()
.map(|s| {
let lower = s.to_lowercase();
lower.contains("timed out") || lower.contains("timeout")
})
.unwrap_or(false),
"error should mention timeout, got: {:?}",
json["error"]
);
}
#[tokio::test]
async fn tool_call_auth_required() {
let (state, _) = test_state_with_executor(false);
let job_id = Uuid::new_v4();
// Do NOT create a token -- request should be rejected
let router = OrchestratorApi::router(state);
let payload = serde_json::json!({
"tool_name": "echo",
"parameters": {"message": "hello"},
});
let req = Request::builder()
.method("POST")
.uri(format!("/worker/{}/tools/call", job_id))
// No Authorization header
.header("Content-Type", "application/json")
.body(Body::from(serde_json::to_vec(&payload).unwrap()))
.unwrap();
let resp = router.oneshot(req).await.unwrap();
assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
}
}
+12
View File
@@ -49,7 +49,9 @@ use uuid::Uuid;
use crate::channels::web::types::SseEvent;
use crate::db::Database;
use crate::llm::LlmProvider;
use crate::safety::SafetyLayer;
use crate::secrets::SecretsStore;
use crate::tools::{ToolExecutor, ToolRegistry};
/// Resolve the orchestrator port from the `ORCHESTRATOR_PORT` environment
/// variable, falling back to 50051.
@@ -75,6 +77,8 @@ pub async fn setup_orchestrator(
llm: &Arc<dyn LlmProvider>,
db: Option<&Arc<dyn Database>>,
secrets_store: Option<&Arc<dyn SecretsStore + Send + Sync>>,
tools: &Arc<ToolRegistry>,
safety: &Arc<SafetyLayer>,
) -> OrchestratorSetup {
let prompt_queue = Arc::new(Mutex::new(
HashMap::<Uuid, VecDeque<api::PendingPrompt>>::new(),
@@ -125,6 +129,13 @@ pub async fn setup_orchestrator(
};
let jm = Arc::new(ContainerJobManager::new(job_config, token_store.clone()));
// Build ToolExecutor for programmatic tool calling (PTC)
let tool_executor = Arc::new(ToolExecutor::new(
Arc::clone(tools),
Arc::clone(safety),
std::time::Duration::from_secs(60),
));
let orchestrator_state = api::OrchestratorState {
llm: Arc::clone(llm),
job_manager: Arc::clone(&jm),
@@ -134,6 +145,7 @@ pub async fn setup_orchestrator(
store: db.cloned(),
secrets_store: secrets_store.cloned(),
user_id: "default".to_string(),
tool_executor: Some(tool_executor),
};
tokio::spawn(async move {
+441
View File
@@ -0,0 +1,441 @@
//! Tool executor for programmatic tool calling (PTC).
//!
//! Provides a standalone execution engine that can be used by both the
//! Docker HTTP RPC path (orchestrator endpoint) and the WASM host function
//! path (tool_invoke). Extracts the tool dispatch flow into a reusable struct.
use std::sync::Arc;
use std::time::{Duration, Instant};
use crate::context::JobContext;
use crate::safety::SafetyLayer;
use crate::tools::registry::ToolRegistry;
/// Maximum allowed nesting depth for tool-invokes-tool chains.
pub const MAX_NESTING_DEPTH: u32 = 5;
/// Maximum per-call timeout (5 minutes).
const MAX_TIMEOUT_SECS: u64 = 300;
/// Result of a programmatic tool call.
#[derive(Debug, Clone)]
pub struct PtcToolResult {
/// Tool output (potentially sanitized).
pub output: String,
/// Whether the output was modified by the safety layer.
pub was_sanitized: bool,
/// Wall-clock duration of the tool execution.
pub duration: Duration,
}
/// Errors that can occur during programmatic tool execution.
#[derive(Debug, thiserror::Error)]
pub enum PtcError {
#[error("Tool not found: {name}")]
NotFound { name: String },
#[error("Tool execution failed: {name}: {reason}")]
ExecutionFailed { name: String, reason: String },
#[error("Tool execution timed out: {name} (timeout: {timeout:?})")]
Timeout { name: String, timeout: Duration },
#[error("Invalid parameters for tool {name}: {reason}")]
InvalidParameters { name: String, reason: String },
#[error("Tool {name} is rate limited")]
RateLimited { name: String },
#[error("Tool output blocked by safety layer: {reason}")]
SafetyBlocked { reason: String },
#[error("Nesting depth exceeded (max {max})")]
NestingDepthExceeded { max: u32 },
}
/// Standalone tool execution engine for programmatic tool calling.
///
/// Used by:
/// - The orchestrator's `POST /worker/{job_id}/tools/call` endpoint
/// - The WASM `tool_invoke` host function
pub struct ToolExecutor {
tools: Arc<ToolRegistry>,
safety: Arc<SafetyLayer>,
default_timeout: Duration,
}
impl ToolExecutor {
/// Create a new tool executor.
pub fn new(
tools: Arc<ToolRegistry>,
safety: Arc<SafetyLayer>,
default_timeout: Duration,
) -> Self {
Self {
tools,
safety,
default_timeout,
}
}
/// Execute a tool by name with the given parameters.
///
/// Flow: lookup -> execute with timeout -> sanitize output -> return.
pub async fn execute(
&self,
tool_name: &str,
params: serde_json::Value,
ctx: &JobContext,
timeout_override: Option<Duration>,
) -> Result<PtcToolResult, PtcError> {
// Enforce global nesting depth limit
if ctx.tool_nesting_depth >= MAX_NESTING_DEPTH {
return Err(PtcError::NestingDepthExceeded {
max: MAX_NESTING_DEPTH,
});
}
let start = Instant::now();
// Look up the tool
let tool = self
.tools
.get(tool_name)
.await
.ok_or_else(|| PtcError::NotFound {
name: tool_name.to_string(),
})?;
// 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
}
})
.min(Duration::from_secs(MAX_TIMEOUT_SECS));
// Execute with timeout
let tool_result = tokio::time::timeout(timeout, tool.execute(params, ctx))
.await
.map_err(|_| PtcError::Timeout {
name: tool_name.to_string(),
timeout,
})?
.map_err(|e| match e {
crate::tools::ToolError::InvalidParameters(reason) => {
PtcError::InvalidParameters {
name: tool_name.to_string(),
reason,
}
}
crate::tools::ToolError::RateLimited(_) => PtcError::RateLimited {
name: tool_name.to_string(),
},
other => PtcError::ExecutionFailed {
name: tool_name.to_string(),
reason: other.to_string(),
},
})?;
// Get output string
let raw_output = tool_result
.raw
.as_deref()
.or_else(|| tool_result.result.as_str())
.unwrap_or("")
.to_string();
let raw_output = if raw_output.is_empty() {
serde_json::to_string(&tool_result.result).unwrap_or_default()
} else {
raw_output
};
// 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")
{
return Err(PtcError::SafetyBlocked {
reason: sanitized.content,
});
}
(sanitized.content, sanitized.was_modified)
} else {
(raw_output, false)
};
Ok(PtcToolResult {
output,
was_sanitized,
duration: start.elapsed(),
})
}
}
impl std::fmt::Debug for ToolExecutor {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ToolExecutor")
.field("default_timeout", &self.default_timeout)
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::safety::SafetyLayer;
use crate::tools::tool::{Tool, ToolError, ToolOutput};
fn test_safety_config() -> crate::config::SafetyConfig {
crate::config::SafetyConfig {
max_output_length: 100_000,
injection_check_enabled: true,
}
}
struct SlowTool;
#[async_trait::async_trait]
impl Tool for SlowTool {
fn name(&self) -> &str {
"slow_tool"
}
fn description(&self) -> &str {
"A tool that sleeps"
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({"type": "object"})
}
async fn execute(
&self,
_params: serde_json::Value,
_ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
tokio::time::sleep(Duration::from_secs(10)).await;
Ok(ToolOutput::text("done", Duration::from_secs(10)))
}
fn requires_sanitization(&self) -> bool {
false
}
}
#[tokio::test]
async fn test_execute_not_found() {
let tools = Arc::new(ToolRegistry::new());
let safety = Arc::new(SafetyLayer::new(&test_safety_config()));
let executor = ToolExecutor::new(tools, safety, Duration::from_secs(60));
let ctx = JobContext::new("test", "test");
let result = executor
.execute("nonexistent", serde_json::json!({}), &ctx, None)
.await;
assert!(matches!(result, Err(PtcError::NotFound { .. })));
}
#[tokio::test]
async fn test_execute_echo() {
let tools = Arc::new(ToolRegistry::new());
tools.register_builtin_tools();
let safety = Arc::new(SafetyLayer::new(&test_safety_config()));
let executor = ToolExecutor::new(tools, safety, Duration::from_secs(60));
let ctx = JobContext::new("test", "test");
let result = executor
.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));
}
#[tokio::test]
async fn test_execute_timeout() {
let tools = Arc::new(ToolRegistry::new());
tools.register(Arc::new(SlowTool)).await;
let safety = Arc::new(SafetyLayer::new(&test_safety_config()));
let executor = ToolExecutor::new(tools, safety, Duration::from_secs(60));
let ctx = JobContext::new("test", "test");
let result = executor
.execute(
"slow_tool",
serde_json::json!({}),
&ctx,
Some(Duration::from_millis(50)),
)
.await;
assert!(matches!(result, Err(PtcError::Timeout { .. })));
}
#[tokio::test]
async fn test_nesting_depth_exceeded() {
let tools = Arc::new(ToolRegistry::new());
tools.register_builtin_tools();
let safety = Arc::new(SafetyLayer::new(&test_safety_config()));
let executor = ToolExecutor::new(tools, safety, Duration::from_secs(60));
let mut ctx = JobContext::new("test", "test");
ctx.tool_nesting_depth = MAX_NESTING_DEPTH; // already at max
let result = executor
.execute("echo", serde_json::json!({"message": "hello"}), &ctx, None)
.await;
assert!(matches!(
result,
Err(PtcError::NestingDepthExceeded { .. })
));
}
#[tokio::test]
async fn test_nesting_depth_within_limit() {
let tools = Arc::new(ToolRegistry::new());
tools.register_builtin_tools();
let safety = Arc::new(SafetyLayer::new(&test_safety_config()));
let executor = ToolExecutor::new(tools, safety, Duration::from_secs(60));
let mut ctx = JobContext::new("test", "test");
ctx.tool_nesting_depth = MAX_NESTING_DEPTH - 1; // one below max
let result = executor
.execute("echo", serde_json::json!({"message": "hello"}), &ctx, None)
.await;
assert!(result.is_ok());
}
struct LeakyTool;
#[async_trait::async_trait]
impl Tool for LeakyTool {
fn name(&self) -> &str {
"leaky_tool"
}
fn description(&self) -> &str {
"Returns output with fake bearer token"
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({"type": "object"})
}
async fn execute(
&self,
_params: serde_json::Value,
_ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
// Bearer token pattern triggers LeakAction::Redact (not Block),
// so the safety layer redacts it and returns sanitized output.
let output =
"Here is some data: Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9_longtokenvalue end";
Ok(ToolOutput::text(output, Duration::from_millis(1)))
}
fn requires_sanitization(&self) -> bool {
true
}
}
struct InvalidParamsTool;
#[async_trait::async_trait]
impl Tool for InvalidParamsTool {
fn name(&self) -> &str {
"invalid_params_tool"
}
fn description(&self) -> &str {
"Always fails with InvalidParameters"
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({"type": "object"})
}
async fn execute(
&self,
_params: serde_json::Value,
_ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
Err(ToolError::InvalidParameters("bad params".to_string()))
}
fn requires_sanitization(&self) -> bool {
false
}
}
#[tokio::test]
async fn test_execute_safety_sanitization() {
let tools = Arc::new(ToolRegistry::new());
tools.register(Arc::new(LeakyTool)).await;
let safety = Arc::new(SafetyLayer::new(&test_safety_config()));
let executor = ToolExecutor::new(tools, safety, Duration::from_secs(60));
let ctx = JobContext::new("test", "test");
let result = executor
.execute("leaky_tool", serde_json::json!({}), &ctx, None)
.await;
// The safety layer should detect the API key pattern and modify the output
assert!(result.is_ok());
let ptc_result = result.unwrap();
assert!(
ptc_result.was_sanitized,
"Output with API key should be sanitized"
);
}
#[tokio::test]
async fn test_execute_invalid_params() {
let tools = Arc::new(ToolRegistry::new());
tools.register(Arc::new(InvalidParamsTool)).await;
let safety = Arc::new(SafetyLayer::new(&test_safety_config()));
let executor = ToolExecutor::new(tools, safety, Duration::from_secs(60));
let ctx = JobContext::new("test", "test");
let result = executor
.execute("invalid_params_tool", serde_json::json!({}), &ctx, None)
.await;
match result {
Err(PtcError::InvalidParameters { name, reason }) => {
assert_eq!(name, "invalid_params_tool");
assert!(reason.contains("bad params"));
}
other => panic!("Expected InvalidParameters, got {:?}", other),
}
}
#[tokio::test]
async fn test_execute_sequential_calls() {
let tools = Arc::new(ToolRegistry::new());
tools.register_builtin_tools();
let safety = Arc::new(SafetyLayer::new(&test_safety_config()));
let executor = ToolExecutor::new(tools, safety, Duration::from_secs(60));
let ctx = JobContext::new("test", "test");
let messages = ["alpha", "beta", "gamma"];
for msg in &messages {
let result = executor
.execute("echo", serde_json::json!({"message": msg}), &ctx, None)
.await
.expect("echo should succeed");
assert!(
result.output.contains(msg),
"Output should contain '{}'",
msg
);
}
}
}
+2
View File
@@ -18,6 +18,7 @@ pub mod redaction;
pub mod schema_validator;
pub mod wasm;
mod executor;
mod registry;
mod tool;
@@ -25,6 +26,7 @@ pub use autonomy::{
AUTONOMOUS_TOOL_DENYLIST, autonomous_allowed_tool_names, autonomous_unavailable_error,
autonomous_unavailable_message, is_autonomous_tool_denylisted,
};
pub use executor::{PtcError, PtcToolResult, ToolExecutor};
pub use builder::{
BuildPhase, BuildRequirement, BuildResult, BuildSoftwareTool, BuilderConfig, Language,
LlmSoftwareBuilder, SoftwareBuilder, SoftwareType, Template, TemplateEngine, TemplateType,
+17
View File
@@ -25,6 +25,7 @@ use crate::tools::builtin::{
ToolUpgradeTool, WriteFileTool,
};
use crate::tools::rate_limiter::RateLimiter;
use crate::tools::executor::ToolExecutor;
use crate::tools::tool::{ApprovalRequirement, Tool, ToolDomain};
use crate::tools::wasm::{
Capabilities, OAuthRefreshConfig, ResourceLimits, SharedCredentialRegistry, WasmError,
@@ -93,6 +94,8 @@ pub struct ToolRegistry {
rate_limiter: RateLimiter,
/// Reference to the message tool for setting context per-turn.
message_tool: RwLock<Option<Arc<crate::tools::builtin::MessageTool>>>,
/// Tool executor for injecting into WASM tools (enables PTC via tool_invoke).
tool_executor: RwLock<Option<Arc<ToolExecutor>>>,
}
impl ToolRegistry {
@@ -114,6 +117,7 @@ impl ToolRegistry {
secrets_store: None,
rate_limiter: RateLimiter::new(),
message_tool: RwLock::new(None),
tool_executor: RwLock::new(None),
}
}
@@ -138,6 +142,14 @@ impl ToolRegistry {
&self.rate_limiter
}
/// 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<ToolExecutor>) {
*self.tool_executor.write().await = Some(executor);
}
/// Register a tool. Rejects dynamic tools that try to shadow a protected built-in name.
pub async fn register(&self, tool: Arc<dyn Tool>) {
let name = tool.name().to_string();
@@ -659,6 +671,11 @@ 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));
}
// Register the tool
self.register(Arc::new(wrapper)).await;
+339 -7
View File
@@ -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(&params_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
+34
View File
@@ -114,6 +114,32 @@ pub struct CredentialResponse {
pub value: String,
}
/// Request to call a tool programmatically via the orchestrator.
#[derive(Debug, Serialize, Deserialize)]
pub struct ToolCallRequest {
/// Name of the tool to invoke.
pub tool_name: String,
/// JSON parameters to pass to the tool.
pub parameters: serde_json::Value,
/// Optional timeout in seconds (capped at 300s by the orchestrator).
pub timeout_secs: Option<u64>,
}
/// Response from a programmatic tool call.
#[derive(Debug, Serialize, Deserialize)]
pub struct ToolCallResponse {
/// Whether the tool call succeeded.
pub success: bool,
/// Tool output (present on success).
pub output: Option<String>,
/// Error message (present on failure).
pub error: Option<String>,
/// Execution duration in milliseconds.
pub duration_ms: u64,
/// Whether the output was modified by the safety layer.
pub was_sanitized: bool,
}
impl WorkerHttpClient {
/// Create a new client from environment.
///
@@ -399,6 +425,14 @@ impl WorkerHttpClient {
})
}
/// Call a tool programmatically via the orchestrator (PTC).
///
/// This bypasses the LLM round-trip and invokes a tool directly on the
/// orchestrator side. Useful for scripted multi-step sequences.
pub async fn call_tool(&self, req: &ToolCallRequest) -> Result<ToolCallResponse, WorkerError> {
self.post_json("tools/call", req, "tool call").await
}
/// Signal job completion to the orchestrator.
pub async fn report_complete(&self, report: &CompletionReport) -> Result<(), WorkerError> {
let _: serde_json::Value = self
+20
View File
@@ -0,0 +1,20 @@
[package]
name = "test-ptc-tool"
version = "0.1.0"
edition = "2021"
publish = false
[dependencies]
wit-bindgen = "0.41.0"
serde_json = "1.0"
[lib]
crate-type = ["cdylib"]
[profile.release]
opt-level = "s"
lto = true
strip = true
codegen-units = 1
[workspace]
+52
View File
@@ -0,0 +1,52 @@
wit_bindgen::generate!({
world: "sandboxed-tool",
path: "../../wit/tool.wit",
});
struct TestPtcTool;
impl exports::near::agent::tool::Guest for TestPtcTool {
fn execute(req: exports::near::agent::tool::Request) -> exports::near::agent::tool::Response {
match execute_inner(&req.params) {
Ok(result) => exports::near::agent::tool::Response {
output: Some(result),
error: None,
},
Err(e) => exports::near::agent::tool::Response {
output: None,
error: Some(e),
},
}
}
fn schema() -> String {
r#"{"type":"object","properties":{"message":{"type":"string"}},"required":["message"]}"#.to_string()
}
fn description() -> String {
"Test tool for PTC: calls echo via tool_invoke".to_string()
}
}
fn execute_inner(params: &str) -> Result<String, String> {
let parsed: serde_json::Value = serde_json::from_str(params)
.map_err(|e| format!("Invalid params: {}", e))?;
let message = parsed.get("message")
.and_then(|v| v.as_str())
.ok_or("Missing 'message' parameter")?;
// Build the parameters for the echo tool
let echo_params = serde_json::json!({"message": message});
// Call tool_invoke with alias "echo_alias" which should resolve to "echo"
let result = near::agent::host::tool_invoke(
"echo_alias",
&echo_params.to_string(),
)?;
// Prefix to prove it went through WASM
Ok(format!("via_wasm:{}", result))
}
export!(TestPtcTool);