From 42e6650ab86d8a77ee2494c0745874ea87036ab4 Mon Sep 17 00:00:00 2001 From: Zaki Date: Fri, 27 Feb 2026 17:02:29 -0800 Subject: [PATCH] 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 --- sdk/python/ironclaw_tools.py | 155 +++++++++++ sdk/python/test_ironclaw_tools.py | 148 ++++++++++ src/context/state.rs | 7 + src/db/libsql/jobs.rs | 1 + src/history/store.rs | 1 + src/main.rs | 2 + src/orchestrator/api.rs | 418 ++++++++++++++++++++++++++++ src/orchestrator/mod.rs | 12 + src/tools/executor.rs | 441 ++++++++++++++++++++++++++++++ src/tools/mod.rs | 2 + src/tools/registry.rs | 17 ++ src/tools/wasm/wrapper.rs | 346 ++++++++++++++++++++++- src/worker/api.rs | 34 +++ tools-src/test-ptc/Cargo.toml | 20 ++ tools-src/test-ptc/src/lib.rs | 52 ++++ 15 files changed, 1649 insertions(+), 7 deletions(-) create mode 100644 sdk/python/ironclaw_tools.py create mode 100644 sdk/python/test_ironclaw_tools.py create mode 100644 src/tools/executor.rs create mode 100644 tools-src/test-ptc/Cargo.toml create mode 100644 tools-src/test-ptc/src/lib.rs diff --git a/sdk/python/ironclaw_tools.py b/sdk/python/ironclaw_tools.py new file mode 100644 index 00000000..fa6d535d --- /dev/null +++ b/sdk/python/ironclaw_tools.py @@ -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) diff --git a/sdk/python/test_ironclaw_tools.py b/sdk/python/test_ironclaw_tools.py new file mode 100644 index 00000000..004ef3b4 --- /dev/null +++ b/sdk/python/test_ironclaw_tools.py @@ -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() diff --git a/src/context/state.rs b/src/context/state.rs index f5307947..2d306693 100644 --- a/src/context/state.rs +++ b/src/context/state.rs @@ -196,6 +196,12 @@ pub struct JobContext { pub tool_output_stash: Arc>>, /// 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, } } diff --git a/src/db/libsql/jobs.rs b/src/db/libsql/jobs.rs index 208d348b..8a075659 100644 --- a/src/db/libsql/jobs.rs +++ b/src/db/libsql/jobs.rs @@ -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), diff --git a/src/history/store.rs b/src/history/store.rs index f0b593c2..2768b9c0 100644 --- a/src/history/store.rs +++ b/src/history/store.rs @@ -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), diff --git a/src/main.rs b/src/main.rs index af310fc4..dfa27475 100644 --- a/src/main.rs +++ b/src/main.rs @@ -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; diff --git a/src/orchestrator/api.rs b/src/orchestrator/api.rs index 8d77c581..7c9700ff 100644 --- a/src/orchestrator/api.rs +++ b/src/orchestrator/api.rs @@ -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>, /// User ID for secret lookups (single-tenant, typically "default"). pub user_id: String, + /// Tool executor for programmatic tool calling (PTC). + pub tool_executor: Option>, } /// 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, + Path(job_id): Path, + Json(req): Json, +) -> Result, 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 { + 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>, + ) { + 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); + } } diff --git a/src/orchestrator/mod.rs b/src/orchestrator/mod.rs index b72f90ee..37c7dc57 100644 --- a/src/orchestrator/mod.rs +++ b/src/orchestrator/mod.rs @@ -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, db: Option<&Arc>, secrets_store: Option<&Arc>, + tools: &Arc, + safety: &Arc, ) -> OrchestratorSetup { let prompt_queue = Arc::new(Mutex::new( HashMap::>::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 { diff --git a/src/tools/executor.rs b/src/tools/executor.rs new file mode 100644 index 00000000..c0a866d3 --- /dev/null +++ b/src/tools/executor.rs @@ -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, + safety: Arc, + default_timeout: Duration, +} + +impl ToolExecutor { + /// Create a new tool executor. + pub fn new( + tools: Arc, + safety: Arc, + 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, + ) -> Result { + // 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 { + 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 { + // 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 { + 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 + ); + } + } +} diff --git a/src/tools/mod.rs b/src/tools/mod.rs index 653544fd..83438b76 100644 --- a/src/tools/mod.rs +++ b/src/tools/mod.rs @@ -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, diff --git a/src/tools/registry.rs b/src/tools/registry.rs index 4564de7c..a448c2b2 100644 --- a/src/tools/registry.rs +++ b/src/tools/registry.rs @@ -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>>, + /// Tool executor for injecting into WASM tools (enables PTC via tool_invoke). + tool_executor: RwLock>>, } 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) { + *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) { 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; diff --git a/src/tools/wasm/wrapper.rs b/src/tools/wasm/wrapper.rs index be089dd8..17d85e54 100644 --- a/src/tools/wasm/wrapper.rs +++ b/src/tools/wasm/wrapper.rs @@ -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 Result + 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, + /// 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, + /// 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, host_credentials: Vec, + tool_resolver: Option, ) -> 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 { + fn tool_invoke(&mut self, alias: String, params_json: String) -> Result { + 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>, /// OAuth refresh configuration for auto-refreshing expired tokens. oauth_refresh: Option, + /// Tool executor for programmatic tool calling from within WASM tools. + tool_executor: Option>, } #[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) -> 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, host_credentials: Vec, + tool_resolver: Option, ) -> Result<(String, Vec), 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 = 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 Result + 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> { + 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 = 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 = 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 = 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 diff --git a/src/worker/api.rs b/src/worker/api.rs index 43fda2dd..420a6b05 100644 --- a/src/worker/api.rs +++ b/src/worker/api.rs @@ -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, +} + +/// 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, + /// Error message (present on failure). + pub error: Option, + /// 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 { + 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 diff --git a/tools-src/test-ptc/Cargo.toml b/tools-src/test-ptc/Cargo.toml new file mode 100644 index 00000000..a79f4f0e --- /dev/null +++ b/tools-src/test-ptc/Cargo.toml @@ -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] diff --git a/tools-src/test-ptc/src/lib.rs b/tools-src/test-ptc/src/lib.rs new file mode 100644 index 00000000..947f61a9 --- /dev/null +++ b/tools-src/test-ptc/src/lib.rs @@ -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 { + 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);