mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-30 16:19:21 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
10468ce5be | ||
|
|
3c46dd74b8 |
@@ -55,9 +55,6 @@ RUN npm install -g @anthropic-ai/claude-code@latest
|
|||||||
# Copy the binary
|
# Copy the binary
|
||||||
COPY --from=builder /build/target/release/ironclaw /usr/local/bin/ironclaw
|
COPY --from=builder /build/target/release/ironclaw /usr/local/bin/ironclaw
|
||||||
|
|
||||||
# Install IronClaw Python SDK for programmatic tool calling (PTC)
|
|
||||||
COPY sdk/python/ironclaw_tools.py /usr/lib/python3/dist-packages/ironclaw_tools.py
|
|
||||||
|
|
||||||
# Create non-root user (UID 1000 matches the orchestrator's container config)
|
# Create non-root user (UID 1000 matches the orchestrator's container config)
|
||||||
RUN useradd -m -u 1000 -s /bin/bash sandbox \
|
RUN useradd -m -u 1000 -s /bin/bash sandbox \
|
||||||
&& mkdir -p /workspace \
|
&& mkdir -p /workspace \
|
||||||
|
|||||||
@@ -1,158 +0,0 @@
|
|||||||
"""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=60):
|
|
||||||
"""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: Timeout in seconds (default 60, max 300).
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
Tool output as a string.
|
|
||||||
|
|
||||||
Raises:
|
|
||||||
RuntimeError: If the tool call fails.
|
|
||||||
"""
|
|
||||||
url = f"{_base_url()}/tools/call"
|
|
||||||
server_timeout = min(int(timeout_secs), 300)
|
|
||||||
body = {
|
|
||||||
"tool_name": name,
|
|
||||||
"parameters": params or {},
|
|
||||||
"timeout_secs": server_timeout,
|
|
||||||
}
|
|
||||||
|
|
||||||
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:
|
|
||||||
# Client-side timeout slightly longer than server-side to account
|
|
||||||
# for network latency, preventing premature client timeouts.
|
|
||||||
client_timeout = server_timeout + 5
|
|
||||||
with urllib.request.urlopen(req, timeout=client_timeout) 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)
|
|
||||||
@@ -1,148 +0,0 @@
|
|||||||
"""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()
|
|
||||||
+72
-1
@@ -176,6 +176,28 @@ pub(crate) fn parse_string_env(
|
|||||||
Ok(optional_env(key)?.unwrap_or_else(|| default.into()))
|
Ok(optional_env(key)?.unwrap_or_else(|| default.into()))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Redact userinfo (username:password) from a parsed URL for safe inclusion in error messages.
|
||||||
|
fn redact_parsed_url(parsed: &reqwest::Url) -> String {
|
||||||
|
if !parsed.username().is_empty() || parsed.password().is_some() {
|
||||||
|
let mut redacted = parsed.clone();
|
||||||
|
let _ = redacted.set_username("");
|
||||||
|
let _ = redacted.set_password(None);
|
||||||
|
format!("{} (credentials redacted)", redacted)
|
||||||
|
} else {
|
||||||
|
parsed.to_string()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Redact userinfo (username:password) from a URL string for safe inclusion in error messages.
|
||||||
|
///
|
||||||
|
/// Use this when you only have the raw string and no parsed URL (e.g., on a parse-error path).
|
||||||
|
fn redact_url(url: &str) -> String {
|
||||||
|
match reqwest::Url::parse(url) {
|
||||||
|
Ok(parsed) => redact_parsed_url(&parsed),
|
||||||
|
Err(_) => "<invalid URL>".to_string(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Validate a user-configurable base URL to prevent SSRF attacks (#1103).
|
/// Validate a user-configurable base URL to prevent SSRF attacks (#1103).
|
||||||
///
|
///
|
||||||
/// Rejects:
|
/// Rejects:
|
||||||
@@ -188,9 +210,10 @@ pub(crate) fn parse_string_env(
|
|||||||
pub(crate) fn validate_base_url(url: &str, field_name: &str) -> Result<(), ConfigError> {
|
pub(crate) fn validate_base_url(url: &str, field_name: &str) -> Result<(), ConfigError> {
|
||||||
use std::net::{IpAddr, Ipv4Addr};
|
use std::net::{IpAddr, Ipv4Addr};
|
||||||
|
|
||||||
|
// Parse once and reuse — on parse failure, fall back to `redact_url()` for the error message.
|
||||||
let parsed = reqwest::Url::parse(url).map_err(|e| ConfigError::InvalidValue {
|
let parsed = reqwest::Url::parse(url).map_err(|e| ConfigError::InvalidValue {
|
||||||
key: field_name.to_string(),
|
key: field_name.to_string(),
|
||||||
message: format!("invalid URL '{}': {}", url, e),
|
message: format!("invalid URL '{}': {}", redact_url(url), e),
|
||||||
})?;
|
})?;
|
||||||
|
|
||||||
let scheme = parsed.scheme();
|
let scheme = parsed.scheme();
|
||||||
@@ -489,4 +512,52 @@ mod tests {
|
|||||||
"Expected DNS resolution failure, got: {err}"
|
"Expected DNS resolution failure, got: {err}"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn validate_base_url_credentials_redacted_in_errors() {
|
||||||
|
// Verify that error messages do not leak credentials from URLs
|
||||||
|
let err = validate_base_url("ftp://user:[email protected]", "TEST")
|
||||||
|
.unwrap_err()
|
||||||
|
.to_string();
|
||||||
|
assert!(
|
||||||
|
!err.contains("secret"),
|
||||||
|
"error message should not contain the password: {err}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn redact_url_strips_credentials() {
|
||||||
|
let result = redact_url("https://user:[email protected]/v1");
|
||||||
|
assert!(
|
||||||
|
!result.contains("user:pass"),
|
||||||
|
"credentials should be redacted: {result}"
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
result.contains("example.com"),
|
||||||
|
"host should be preserved: {result}"
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
result.contains("redacted"),
|
||||||
|
"should indicate redaction: {result}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn redact_url_no_credentials_unchanged() {
|
||||||
|
let result = redact_url("https://example.com/v1");
|
||||||
|
assert!(
|
||||||
|
!result.contains("redacted"),
|
||||||
|
"no-credential URL should not mention redaction: {result}"
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
result.contains("example.com"),
|
||||||
|
"host should be preserved: {result}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn redact_url_invalid_url_returns_placeholder() {
|
||||||
|
let result = redact_url("not-a-url");
|
||||||
|
assert_eq!(result, "<invalid URL>");
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -196,12 +196,6 @@ pub struct JobContext {
|
|||||||
pub tool_output_stash: Arc<tokio::sync::RwLock<HashMap<String, String>>>,
|
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".
|
/// User's preferred timezone (IANA name, e.g. "America/New_York"). Defaults to "UTC".
|
||||||
pub user_timezone: String,
|
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 {
|
impl JobContext {
|
||||||
@@ -243,7 +237,6 @@ impl JobContext {
|
|||||||
metadata: serde_json::Value::Null,
|
metadata: serde_json::Value::Null,
|
||||||
tool_output_stash: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
|
tool_output_stash: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
|
||||||
user_timezone: "UTC".to_string(),
|
user_timezone: "UTC".to_string(),
|
||||||
tool_nesting_depth: 0,
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -134,7 +134,6 @@ impl JobStore for LibSqlBackend {
|
|||||||
// TODO(#661): persist user_timezone in agent_jobs table so
|
// TODO(#661): persist user_timezone in agent_jobs table so
|
||||||
// background/routine jobs retain the session's timezone context.
|
// background/routine jobs retain the session's timezone context.
|
||||||
user_timezone: "UTC".to_string(),
|
user_timezone: "UTC".to_string(),
|
||||||
tool_nesting_depth: 0,
|
|
||||||
}))
|
}))
|
||||||
}
|
}
|
||||||
None => Ok(None),
|
None => Ok(None),
|
||||||
|
|||||||
@@ -258,7 +258,6 @@ impl Store {
|
|||||||
// TODO(#661): persist user_timezone in agent_jobs table so
|
// TODO(#661): persist user_timezone in agent_jobs table so
|
||||||
// background/routine jobs retain the session's timezone context.
|
// background/routine jobs retain the session's timezone context.
|
||||||
user_timezone: "UTC".to_string(),
|
user_timezone: "UTC".to_string(),
|
||||||
tool_nesting_depth: 0,
|
|
||||||
}))
|
}))
|
||||||
}
|
}
|
||||||
None => Ok(None),
|
None => Ok(None),
|
||||||
|
|||||||
@@ -1,5 +1,7 @@
|
|||||||
//! Shared test helpers for OpenAI Codex provider tests.
|
//! Shared test helpers for OpenAI Codex provider tests.
|
||||||
|
|
||||||
|
#![cfg(test)]
|
||||||
|
|
||||||
use crate::config::OpenAiCodexConfig;
|
use crate::config::OpenAiCodexConfig;
|
||||||
|
|
||||||
/// Build a minimal JWT for testing (header.payload.signature).
|
/// Build a minimal JWT for testing (header.payload.signature).
|
||||||
|
|||||||
@@ -306,8 +306,6 @@ async fn async_main() -> anyhow::Result<()> {
|
|||||||
&components.llm,
|
&components.llm,
|
||||||
components.db.as_ref(),
|
components.db.as_ref(),
|
||||||
components.secrets_store.as_ref(),
|
components.secrets_store.as_ref(),
|
||||||
&components.tools,
|
|
||||||
&components.safety,
|
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
let container_job_manager = orch.container_job_manager;
|
let container_job_manager = orch.container_job_manager;
|
||||||
|
|||||||
+14
-552
@@ -15,18 +15,15 @@ use tokio::sync::{Mutex, broadcast};
|
|||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
use crate::channels::web::types::SseEvent;
|
use crate::channels::web::types::SseEvent;
|
||||||
use crate::context::JobContext;
|
|
||||||
use crate::db::Database;
|
use crate::db::Database;
|
||||||
use crate::llm::{CompletionRequest, LlmProvider, ToolCompletionRequest};
|
use crate::llm::{CompletionRequest, LlmProvider, ToolCompletionRequest};
|
||||||
use crate::orchestrator::auth::{TokenStore, worker_auth_middleware};
|
use crate::orchestrator::auth::{TokenStore, worker_auth_middleware};
|
||||||
use crate::orchestrator::job_manager::ContainerJobManager;
|
use crate::orchestrator::job_manager::ContainerJobManager;
|
||||||
use crate::secrets::SecretsStore;
|
use crate::secrets::SecretsStore;
|
||||||
use crate::tools::ToolExecutor;
|
|
||||||
use crate::worker::api::JobEventPayload;
|
use crate::worker::api::JobEventPayload;
|
||||||
use crate::worker::api::{
|
use crate::worker::api::{
|
||||||
CompletionReport, CredentialResponse, JobDescription, ProxyCompletionRequest,
|
CompletionReport, CredentialResponse, JobDescription, ProxyCompletionRequest,
|
||||||
ProxyCompletionResponse, ProxyToolCompletionRequest, ProxyToolCompletionResponse, StatusUpdate,
|
ProxyCompletionResponse, ProxyToolCompletionRequest, ProxyToolCompletionResponse, StatusUpdate,
|
||||||
ToolCallRequest, ToolCallResponse,
|
|
||||||
};
|
};
|
||||||
|
|
||||||
/// A follow-up prompt queued for a Claude Code bridge.
|
/// A follow-up prompt queued for a Claude Code bridge.
|
||||||
@@ -52,8 +49,6 @@ pub struct OrchestratorState {
|
|||||||
pub secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>>,
|
pub secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>>,
|
||||||
/// User ID for secret lookups (single-tenant, typically "default").
|
/// User ID for secret lookups (single-tenant, typically "default").
|
||||||
pub user_id: String,
|
pub user_id: String,
|
||||||
/// Tool executor for programmatic tool calling (PTC).
|
|
||||||
pub tool_executor: Option<Arc<ToolExecutor>>,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// The orchestrator's internal API server.
|
/// The orchestrator's internal API server.
|
||||||
@@ -75,7 +70,6 @@ impl OrchestratorApi {
|
|||||||
.route("/worker/{job_id}/event", post(job_event_handler))
|
.route("/worker/{job_id}/event", post(job_event_handler))
|
||||||
.route("/worker/{job_id}/prompt", get(get_prompt_handler))
|
.route("/worker/{job_id}/prompt", get(get_prompt_handler))
|
||||||
.route("/worker/{job_id}/credentials", get(get_credentials_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(
|
.route_layer(axum::middleware::from_fn_with_state(
|
||||||
state.token_store.clone(),
|
state.token_store.clone(),
|
||||||
worker_auth_middleware,
|
worker_auth_middleware,
|
||||||
@@ -297,26 +291,20 @@ async fn job_event_handler(
|
|||||||
.unwrap_or("")
|
.unwrap_or("")
|
||||||
.to_string(),
|
.to_string(),
|
||||||
},
|
},
|
||||||
"tool_use" => {
|
"tool_use" => SseEvent::JobToolUse {
|
||||||
// Redact raw parameters from worker-reported tool_use events
|
job_id: job_id_str,
|
||||||
// before broadcasting via SSE. Workers are untrusted and may
|
tool_name: payload
|
||||||
// include sensitive data (API keys, passwords, PII) in the
|
.data
|
||||||
// input payload. We replace it with a placeholder to prevent
|
.get("tool_name")
|
||||||
// leaking secrets to the web UI.
|
.and_then(|v| v.as_str())
|
||||||
let redacted_input = serde_json::json!({
|
.unwrap_or("unknown")
|
||||||
"_note": "parameters redacted for security"
|
.to_string(),
|
||||||
});
|
input: payload
|
||||||
SseEvent::JobToolUse {
|
.data
|
||||||
job_id: job_id_str,
|
.get("input")
|
||||||
tool_name: payload
|
.cloned()
|
||||||
.data
|
.unwrap_or(serde_json::Value::Null),
|
||||||
.get("tool_name")
|
},
|
||||||
.and_then(|v| v.as_str())
|
|
||||||
.unwrap_or("unknown")
|
|
||||||
.to_string(),
|
|
||||||
input: redacted_input,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
"tool_result" => SseEvent::JobToolResult {
|
"tool_result" => SseEvent::JobToolResult {
|
||||||
job_id: job_id_str,
|
job_id: job_id_str,
|
||||||
tool_name: payload
|
tool_name: payload
|
||||||
@@ -455,106 +443,6 @@ 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 mut ctx = JobContext::with_user(
|
|
||||||
state.user_id.clone(),
|
|
||||||
format!("PTC call: {}", req.tool_name),
|
|
||||||
format!("Programmatic tool call from job {}", job_id),
|
|
||||||
);
|
|
||||||
// Do not trust client-provided nesting_depth — a malicious worker
|
|
||||||
// could send any value to bypass the limit. The orchestrator must
|
|
||||||
// increment the depth server-side: each hop through the orchestrator
|
|
||||||
// adds 1. This way even if a worker always sends 0, the depth still
|
|
||||||
// increases with each real nesting level.
|
|
||||||
ctx.tool_nesting_depth = req.nesting_depth.saturating_add(1);
|
|
||||||
|
|
||||||
// Emit tool_use SSE event with redacted parameters to avoid leaking
|
|
||||||
// sensitive data (API keys, passwords, PII) to the web UI.
|
|
||||||
if let Some(ref tx) = state.job_event_tx {
|
|
||||||
let redacted_params = serde_json::json!({
|
|
||||||
"_note": "parameters redacted for security"
|
|
||||||
});
|
|
||||||
let _ = tx.send((
|
|
||||||
job_id,
|
|
||||||
SseEvent::JobToolUse {
|
|
||||||
job_id: job_id.to_string(),
|
|
||||||
tool_name: req.tool_name.clone(),
|
|
||||||
input: redacted_params,
|
|
||||||
},
|
|
||||||
));
|
|
||||||
}
|
|
||||||
|
|
||||||
// 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 {
|
fn format_finish_reason(reason: crate::llm::FinishReason) -> String {
|
||||||
match reason {
|
match reason {
|
||||||
crate::llm::FinishReason::Stop => "stop".to_string(),
|
crate::llm::FinishReason::Stop => "stop".to_string(),
|
||||||
@@ -592,7 +480,6 @@ mod tests {
|
|||||||
store: None,
|
store: None,
|
||||||
secrets_store: None,
|
secrets_store: None,
|
||||||
user_id: "default".to_string(),
|
user_id: "default".to_string(),
|
||||||
tool_executor: None,
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -822,7 +709,6 @@ mod tests {
|
|||||||
store: None,
|
store: None,
|
||||||
secrets_store: Some(secrets_store),
|
secrets_store: Some(secrets_store),
|
||||||
user_id: "default".to_string(),
|
user_id: "default".to_string(),
|
||||||
tool_executor: None,
|
|
||||||
};
|
};
|
||||||
|
|
||||||
let router = OrchestratorApi::router(state);
|
let router = OrchestratorApi::router(state);
|
||||||
@@ -858,7 +744,6 @@ mod tests {
|
|||||||
store: None,
|
store: None,
|
||||||
secrets_store: None,
|
secrets_store: None,
|
||||||
user_id: "default".to_string(),
|
user_id: "default".to_string(),
|
||||||
tool_executor: None,
|
|
||||||
};
|
};
|
||||||
|
|
||||||
let job_id = Uuid::new_v4();
|
let job_id = Uuid::new_v4();
|
||||||
@@ -914,7 +799,6 @@ mod tests {
|
|||||||
store: None,
|
store: None,
|
||||||
secrets_store: None,
|
secrets_store: None,
|
||||||
user_id: "default".to_string(),
|
user_id: "default".to_string(),
|
||||||
tool_executor: None,
|
|
||||||
};
|
};
|
||||||
|
|
||||||
let job_id = Uuid::new_v4();
|
let job_id = Uuid::new_v4();
|
||||||
@@ -963,7 +847,6 @@ mod tests {
|
|||||||
store: None,
|
store: None,
|
||||||
secrets_store: None,
|
secrets_store: None,
|
||||||
user_id: "default".to_string(),
|
user_id: "default".to_string(),
|
||||||
tool_executor: None,
|
|
||||||
};
|
};
|
||||||
|
|
||||||
let job_id = Uuid::new_v4();
|
let job_id = Uuid::new_v4();
|
||||||
@@ -1043,425 +926,4 @@ mod tests {
|
|||||||
assert_eq!(handle.worker_iteration, 5);
|
assert_eq!(handle.worker_iteration, 5);
|
||||||
assert_eq!(handle.last_worker_status.as_deref(), Some("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);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn tool_call_nesting_depth_incremented_server_side() {
|
|
||||||
// A worker sending nesting_depth=4 should get depth=5 after the
|
|
||||||
// orchestrator increments it. With MAX_NESTING_DEPTH=5, this
|
|
||||||
// should be rejected (depth >= max).
|
|
||||||
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"},
|
|
||||||
"nesting_depth": 4,
|
|
||||||
});
|
|
||||||
|
|
||||||
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| s.to_lowercase().contains("nesting"))
|
|
||||||
.unwrap_or(false),
|
|
||||||
"error should mention nesting depth, got: {:?}",
|
|
||||||
json["error"]
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn job_event_tool_use_redacts_input() {
|
|
||||||
// Worker-reported tool_use events must have their input redacted
|
|
||||||
// before SSE broadcast to prevent leaking sensitive parameters.
|
|
||||||
let (tx, mut rx) = broadcast::channel(16);
|
|
||||||
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: Some(tx),
|
|
||||||
prompt_queue: Arc::new(Mutex::new(HashMap::new())),
|
|
||||||
store: None,
|
|
||||||
secrets_store: None,
|
|
||||||
user_id: "default".to_string(),
|
|
||||||
tool_executor: None,
|
|
||||||
};
|
|
||||||
|
|
||||||
let job_id = Uuid::new_v4();
|
|
||||||
let token = token_store.create_token(job_id).await;
|
|
||||||
let router = OrchestratorApi::router(state);
|
|
||||||
|
|
||||||
// Worker sends a tool_use event with sensitive data in input
|
|
||||||
let payload = serde_json::json!({
|
|
||||||
"event_type": "tool_use",
|
|
||||||
"data": {
|
|
||||||
"tool_name": "shell",
|
|
||||||
"input": {"command": "curl -H 'Authorization: Bearer sk-secret-key' https://api.example.com"}
|
|
||||||
}
|
|
||||||
});
|
|
||||||
|
|
||||||
let req = Request::builder()
|
|
||||||
.method("POST")
|
|
||||||
.uri(format!("/worker/{}/event", 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 (_recv_id, event) = rx.recv().await.unwrap();
|
|
||||||
match event {
|
|
||||||
SseEvent::JobToolUse {
|
|
||||||
tool_name, input, ..
|
|
||||||
} => {
|
|
||||||
assert_eq!(tool_name, "shell");
|
|
||||||
// The input must be redacted, not the raw worker payload
|
|
||||||
assert!(
|
|
||||||
input.get("_note").is_some(),
|
|
||||||
"input should be redacted placeholder, got: {}",
|
|
||||||
input
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
!input.to_string().contains("sk-secret-key"),
|
|
||||||
"input must not contain sensitive data"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
other => panic!("Expected JobToolUse, got {:?}", other),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -49,9 +49,7 @@ use uuid::Uuid;
|
|||||||
use crate::channels::web::types::SseEvent;
|
use crate::channels::web::types::SseEvent;
|
||||||
use crate::db::Database;
|
use crate::db::Database;
|
||||||
use crate::llm::LlmProvider;
|
use crate::llm::LlmProvider;
|
||||||
use crate::safety::SafetyLayer;
|
|
||||||
use crate::secrets::SecretsStore;
|
use crate::secrets::SecretsStore;
|
||||||
use crate::tools::{ToolExecutor, ToolRegistry};
|
|
||||||
|
|
||||||
/// Resolve the orchestrator port from the `ORCHESTRATOR_PORT` environment
|
/// Resolve the orchestrator port from the `ORCHESTRATOR_PORT` environment
|
||||||
/// variable, falling back to 50051.
|
/// variable, falling back to 50051.
|
||||||
@@ -77,8 +75,6 @@ pub async fn setup_orchestrator(
|
|||||||
llm: &Arc<dyn LlmProvider>,
|
llm: &Arc<dyn LlmProvider>,
|
||||||
db: Option<&Arc<dyn Database>>,
|
db: Option<&Arc<dyn Database>>,
|
||||||
secrets_store: Option<&Arc<dyn SecretsStore + Send + Sync>>,
|
secrets_store: Option<&Arc<dyn SecretsStore + Send + Sync>>,
|
||||||
tools: &Arc<ToolRegistry>,
|
|
||||||
safety: &Arc<SafetyLayer>,
|
|
||||||
) -> OrchestratorSetup {
|
) -> OrchestratorSetup {
|
||||||
let prompt_queue = Arc::new(Mutex::new(
|
let prompt_queue = Arc::new(Mutex::new(
|
||||||
HashMap::<Uuid, VecDeque<api::PendingPrompt>>::new(),
|
HashMap::<Uuid, VecDeque<api::PendingPrompt>>::new(),
|
||||||
@@ -129,17 +125,6 @@ pub async fn setup_orchestrator(
|
|||||||
};
|
};
|
||||||
let jm = Arc::new(ContainerJobManager::new(job_config, token_store.clone()));
|
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),
|
|
||||||
));
|
|
||||||
|
|
||||||
// Wire the executor into the shared slot so WASM tools registered
|
|
||||||
// during build_all() can resolve it lazily at execution time.
|
|
||||||
tools.set_tool_executor(Arc::clone(&tool_executor));
|
|
||||||
|
|
||||||
let orchestrator_state = api::OrchestratorState {
|
let orchestrator_state = api::OrchestratorState {
|
||||||
llm: Arc::clone(llm),
|
llm: Arc::clone(llm),
|
||||||
job_manager: Arc::clone(&jm),
|
job_manager: Arc::clone(&jm),
|
||||||
@@ -149,7 +134,6 @@ pub async fn setup_orchestrator(
|
|||||||
store: db.cloned(),
|
store: db.cloned(),
|
||||||
secrets_store: secrets_store.cloned(),
|
secrets_store: secrets_store.cloned(),
|
||||||
user_id: "default".to_string(),
|
user_id: "default".to_string(),
|
||||||
tool_executor: Some(tool_executor),
|
|
||||||
};
|
};
|
||||||
|
|
||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
|
|||||||
@@ -9,7 +9,6 @@ mod json;
|
|||||||
mod memory;
|
mod memory;
|
||||||
mod message;
|
mod message;
|
||||||
pub mod path_utils;
|
pub mod path_utils;
|
||||||
pub mod ptc_script;
|
|
||||||
mod restart;
|
mod restart;
|
||||||
pub mod routine;
|
pub mod routine;
|
||||||
pub mod secrets_tools;
|
pub mod secrets_tools;
|
||||||
@@ -32,7 +31,6 @@ pub use job::{
|
|||||||
pub use json::JsonTool;
|
pub use json::JsonTool;
|
||||||
pub use memory::{MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool};
|
pub use memory::{MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool};
|
||||||
pub use message::MessageTool;
|
pub use message::MessageTool;
|
||||||
pub use ptc_script::PtcScriptTool;
|
|
||||||
pub use restart::RestartTool;
|
pub use restart::RestartTool;
|
||||||
pub use routine::{
|
pub use routine::{
|
||||||
EventEmitTool, RoutineCreateTool, RoutineDeleteTool, RoutineFireTool, RoutineHistoryTool,
|
EventEmitTool, RoutineCreateTool, RoutineDeleteTool, RoutineFireTool, RoutineHistoryTool,
|
||||||
|
|||||||
@@ -1,371 +0,0 @@
|
|||||||
//! PTC script tool for running multi-step Python programs that call tools.
|
|
||||||
//!
|
|
||||||
//! Wraps user-provided Python code in a preamble that imports the IronClaw
|
|
||||||
//! SDK (`ironclaw_tools`), then executes it via `python3 -c`. The script
|
|
||||||
//! runs in the same environment as the worker container and can call any
|
|
||||||
//! registered tool through the SDK's `call_tool()` function.
|
|
||||||
|
|
||||||
use std::process::Stdio;
|
|
||||||
use std::time::Duration;
|
|
||||||
|
|
||||||
use async_trait::async_trait;
|
|
||||||
use tokio::io::AsyncReadExt;
|
|
||||||
use tokio::process::Command;
|
|
||||||
|
|
||||||
use crate::context::JobContext;
|
|
||||||
use crate::tools::tool::{
|
|
||||||
ApprovalRequirement, Tool, ToolDomain, ToolError, ToolOutput, require_str,
|
|
||||||
};
|
|
||||||
|
|
||||||
/// Maximum output size before truncation (64KB).
|
|
||||||
const MAX_OUTPUT_SIZE: usize = 64 * 1024;
|
|
||||||
|
|
||||||
/// Default script timeout.
|
|
||||||
const DEFAULT_TIMEOUT_SECS: u64 = 120;
|
|
||||||
|
|
||||||
/// Maximum allowed timeout.
|
|
||||||
const MAX_TIMEOUT_SECS: u64 = 300;
|
|
||||||
|
|
||||||
/// Environment variables safe to forward to the Python subprocess.
|
|
||||||
const SAFE_ENV_VARS: &[&str] = &[
|
|
||||||
"PATH",
|
|
||||||
"HOME",
|
|
||||||
"USER",
|
|
||||||
"LOGNAME",
|
|
||||||
"SHELL",
|
|
||||||
"TERM",
|
|
||||||
"LANG",
|
|
||||||
"LC_ALL",
|
|
||||||
"LC_CTYPE",
|
|
||||||
"PWD",
|
|
||||||
"TMPDIR",
|
|
||||||
"TMP",
|
|
||||||
"TEMP",
|
|
||||||
"CARGO_HOME",
|
|
||||||
"RUSTUP_HOME",
|
|
||||||
"NODE_PATH",
|
|
||||||
"NPM_CONFIG_PREFIX",
|
|
||||||
];
|
|
||||||
|
|
||||||
/// PTC environment variables required by the ironclaw_tools SDK.
|
|
||||||
const PTC_ENV_VARS: &[&str] = &[
|
|
||||||
"IRONCLAW_ORCHESTRATOR_URL",
|
|
||||||
"IRONCLAW_JOB_ID",
|
|
||||||
"IRONCLAW_WORKER_TOKEN",
|
|
||||||
];
|
|
||||||
|
|
||||||
/// Python preamble injected before the user's script.
|
|
||||||
const PREAMBLE: &str = r#"
|
|
||||||
import json, sys, os
|
|
||||||
|
|
||||||
# Import IronClaw SDK
|
|
||||||
from ironclaw_tools import call_tool, shell, read_file, write_file, http_get
|
|
||||||
|
|
||||||
# Structured output collector
|
|
||||||
_ptc_outputs = {}
|
|
||||||
|
|
||||||
def ptc_output(key, value):
|
|
||||||
"""Register a named output value for structured results."""
|
|
||||||
_ptc_outputs[key] = value
|
|
||||||
|
|
||||||
try:
|
|
||||||
"#;
|
|
||||||
|
|
||||||
/// Python postamble appended after the user's script.
|
|
||||||
const POSTAMBLE: &str = r#"
|
|
||||||
except Exception as _ptc_err:
|
|
||||||
print(f"SCRIPT_ERROR: {type(_ptc_err).__name__}: {_ptc_err}", file=sys.stderr)
|
|
||||||
sys.exit(1)
|
|
||||||
|
|
||||||
# Print structured outputs if any were registered
|
|
||||||
if _ptc_outputs:
|
|
||||||
print("\n__PTC_OUTPUTS__")
|
|
||||||
print(json.dumps(_ptc_outputs))
|
|
||||||
"#;
|
|
||||||
|
|
||||||
pub struct PtcScriptTool;
|
|
||||||
|
|
||||||
impl Default for PtcScriptTool {
|
|
||||||
fn default() -> Self {
|
|
||||||
Self
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl PtcScriptTool {
|
|
||||||
pub fn new() -> Self {
|
|
||||||
Self
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Build the full Python program from user script + preamble/postamble.
|
|
||||||
fn build_program(script: &str) -> String {
|
|
||||||
let mut program =
|
|
||||||
String::with_capacity(PREAMBLE.len() + script.len() + POSTAMBLE.len() + 256);
|
|
||||||
program.push_str(PREAMBLE);
|
|
||||||
|
|
||||||
// Indent user script into the try: block
|
|
||||||
for line in script.lines() {
|
|
||||||
program.push_str(" ");
|
|
||||||
program.push_str(line);
|
|
||||||
program.push('\n');
|
|
||||||
}
|
|
||||||
|
|
||||||
program.push_str(POSTAMBLE);
|
|
||||||
program
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Truncate output to MAX_OUTPUT_SIZE with a truncation notice.
|
|
||||||
fn truncate_output(output: &str) -> String {
|
|
||||||
if output.len() <= MAX_OUTPUT_SIZE {
|
|
||||||
output.to_string()
|
|
||||||
} else {
|
|
||||||
let mut i = MAX_OUTPUT_SIZE;
|
|
||||||
while i > 0 && !output.is_char_boundary(i) {
|
|
||||||
i -= 1;
|
|
||||||
}
|
|
||||||
format!(
|
|
||||||
"{}\n\n[Output truncated at {} bytes]",
|
|
||||||
&output[..i],
|
|
||||||
MAX_OUTPUT_SIZE
|
|
||||||
)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[async_trait]
|
|
||||||
impl Tool for PtcScriptTool {
|
|
||||||
fn name(&self) -> &str {
|
|
||||||
"ptc_script"
|
|
||||||
}
|
|
||||||
|
|
||||||
fn description(&self) -> &str {
|
|
||||||
"Execute a Python script that can call IronClaw tools programmatically. \
|
|
||||||
The script has access to call_tool(), shell(), read_file(), write_file(), \
|
|
||||||
and http_get() from the ironclaw_tools SDK. Use ptc_output(key, value) \
|
|
||||||
to return structured results."
|
|
||||||
}
|
|
||||||
|
|
||||||
fn parameters_schema(&self) -> serde_json::Value {
|
|
||||||
serde_json::json!({
|
|
||||||
"type": "object",
|
|
||||||
"properties": {
|
|
||||||
"script": {
|
|
||||||
"type": "string",
|
|
||||||
"description": "Python script to execute. Has access to call_tool(), shell(), read_file(), write_file(), http_get(), and ptc_output()."
|
|
||||||
},
|
|
||||||
"timeout_secs": {
|
|
||||||
"type": "integer",
|
|
||||||
"description": "Timeout in seconds (default 120, max 300).",
|
|
||||||
"default": 120,
|
|
||||||
"minimum": 1,
|
|
||||||
"maximum": 300
|
|
||||||
}
|
|
||||||
},
|
|
||||||
"required": ["script"]
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn execute(
|
|
||||||
&self,
|
|
||||||
params: serde_json::Value,
|
|
||||||
ctx: &JobContext,
|
|
||||||
) -> Result<ToolOutput, ToolError> {
|
|
||||||
let start = std::time::Instant::now();
|
|
||||||
|
|
||||||
let script = require_str(¶ms, "script")?;
|
|
||||||
let timeout_secs = params
|
|
||||||
.get("timeout_secs")
|
|
||||||
.and_then(|v| v.as_u64())
|
|
||||||
.unwrap_or(DEFAULT_TIMEOUT_SECS)
|
|
||||||
.min(MAX_TIMEOUT_SECS);
|
|
||||||
let timeout = Duration::from_secs(timeout_secs);
|
|
||||||
|
|
||||||
let program = Self::build_program(script);
|
|
||||||
|
|
||||||
// Build the subprocess command
|
|
||||||
let mut command = Command::new("python3");
|
|
||||||
command.args(["-c", &program]);
|
|
||||||
|
|
||||||
// Scrub environment -- only forward safe vars + PTC vars + extra_env
|
|
||||||
command.env_clear();
|
|
||||||
for var in SAFE_ENV_VARS {
|
|
||||||
if let Ok(val) = std::env::var(var) {
|
|
||||||
command.env(var, val);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
for var in PTC_ENV_VARS {
|
|
||||||
if let Ok(val) = std::env::var(var) {
|
|
||||||
command.env(var, val);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
// Forward extra_env from JobContext (credentials fetched by worker runtime)
|
|
||||||
for (k, v) in ctx.extra_env.iter() {
|
|
||||||
command.env(k, v);
|
|
||||||
}
|
|
||||||
|
|
||||||
command
|
|
||||||
.stdin(Stdio::null())
|
|
||||||
.stdout(Stdio::piped())
|
|
||||||
.stderr(Stdio::piped());
|
|
||||||
|
|
||||||
// Spawn and drain stdout/stderr concurrently
|
|
||||||
let mut child = command
|
|
||||||
.spawn()
|
|
||||||
.map_err(|e| ToolError::ExecutionFailed(format!("Failed to spawn python3: {}", e)))?;
|
|
||||||
|
|
||||||
let stdout_handle = child.stdout.take();
|
|
||||||
let stderr_handle = child.stderr.take();
|
|
||||||
|
|
||||||
let result = tokio::time::timeout(timeout, async {
|
|
||||||
let stdout_fut = async {
|
|
||||||
if let Some(mut out) = stdout_handle {
|
|
||||||
let mut buf = Vec::new();
|
|
||||||
(&mut out)
|
|
||||||
.take(MAX_OUTPUT_SIZE as u64)
|
|
||||||
.read_to_end(&mut buf)
|
|
||||||
.await
|
|
||||||
.ok();
|
|
||||||
tokio::io::copy(&mut out, &mut tokio::io::sink()).await.ok();
|
|
||||||
String::from_utf8_lossy(&buf).to_string()
|
|
||||||
} else {
|
|
||||||
String::new()
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
let stderr_fut = async {
|
|
||||||
if let Some(mut err) = stderr_handle {
|
|
||||||
let mut buf = Vec::new();
|
|
||||||
(&mut err)
|
|
||||||
.take(MAX_OUTPUT_SIZE as u64)
|
|
||||||
.read_to_end(&mut buf)
|
|
||||||
.await
|
|
||||||
.ok();
|
|
||||||
tokio::io::copy(&mut err, &mut tokio::io::sink()).await.ok();
|
|
||||||
String::from_utf8_lossy(&buf).to_string()
|
|
||||||
} else {
|
|
||||||
String::new()
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
let (stdout, stderr, wait_result) = tokio::join!(stdout_fut, stderr_fut, child.wait());
|
|
||||||
let status = wait_result?;
|
|
||||||
Ok::<_, std::io::Error>((stdout, stderr, status.code().unwrap_or(-1)))
|
|
||||||
})
|
|
||||||
.await;
|
|
||||||
|
|
||||||
let duration = start.elapsed();
|
|
||||||
|
|
||||||
match result {
|
|
||||||
Ok(Ok((stdout, stderr, exit_code))) => {
|
|
||||||
if exit_code != 0 {
|
|
||||||
let error_msg = if stderr.is_empty() {
|
|
||||||
format!("Script exited with code {}", exit_code)
|
|
||||||
} else {
|
|
||||||
format!(
|
|
||||||
"Script exited with code {}:\n{}",
|
|
||||||
exit_code,
|
|
||||||
Self::truncate_output(&stderr)
|
|
||||||
)
|
|
||||||
};
|
|
||||||
return Err(ToolError::ExecutionFailed(error_msg));
|
|
||||||
}
|
|
||||||
|
|
||||||
// Combine output
|
|
||||||
let output = if stderr.is_empty() {
|
|
||||||
stdout
|
|
||||||
} else {
|
|
||||||
format!("{}\n\n--- stderr ---\n{}", stdout, stderr)
|
|
||||||
};
|
|
||||||
|
|
||||||
Ok(ToolOutput::text(Self::truncate_output(&output), duration))
|
|
||||||
}
|
|
||||||
Ok(Err(e)) => Err(ToolError::ExecutionFailed(format!(
|
|
||||||
"Script execution failed: {}",
|
|
||||||
e
|
|
||||||
))),
|
|
||||||
Err(_) => {
|
|
||||||
let _ = child.kill().await;
|
|
||||||
Err(ToolError::Timeout(timeout))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn requires_sanitization(&self) -> bool {
|
|
||||||
true
|
|
||||||
}
|
|
||||||
|
|
||||||
fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement {
|
|
||||||
ApprovalRequirement::Always
|
|
||||||
}
|
|
||||||
|
|
||||||
fn domain(&self) -> ToolDomain {
|
|
||||||
ToolDomain::Container
|
|
||||||
}
|
|
||||||
|
|
||||||
fn execution_timeout(&self) -> Duration {
|
|
||||||
Duration::from_secs(MAX_TIMEOUT_SECS)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod tests {
|
|
||||||
use super::*;
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_build_program_indents_script() {
|
|
||||||
let script = "x = 1\nprint(x)";
|
|
||||||
let program = PtcScriptTool::build_program(script);
|
|
||||||
assert!(program.contains(" x = 1\n"));
|
|
||||||
assert!(program.contains(" print(x)\n"));
|
|
||||||
assert!(program.contains("from ironclaw_tools import"));
|
|
||||||
assert!(program.contains("def ptc_output("));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_build_program_empty_script() {
|
|
||||||
let program = PtcScriptTool::build_program("");
|
|
||||||
// Empty script should still have preamble + postamble
|
|
||||||
assert!(program.contains("try:"));
|
|
||||||
assert!(program.contains("except Exception"));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_truncate_output() {
|
|
||||||
let short = "hello";
|
|
||||||
assert_eq!(PtcScriptTool::truncate_output(short), "hello");
|
|
||||||
|
|
||||||
let long = "x".repeat(MAX_OUTPUT_SIZE + 100);
|
|
||||||
let truncated = PtcScriptTool::truncate_output(&long);
|
|
||||||
assert!(truncated.len() < long.len());
|
|
||||||
assert!(truncated.contains("[Output truncated"));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_truncate_output_multibyte_boundary() {
|
|
||||||
// Build a string of multi-byte chars (emoji = 4 bytes each) that crosses MAX_OUTPUT_SIZE
|
|
||||||
let emoji = "\u{1F600}"; // 4 bytes
|
|
||||||
let count = MAX_OUTPUT_SIZE / emoji.len() + 10;
|
|
||||||
let long: String = emoji.repeat(count);
|
|
||||||
assert!(long.len() > MAX_OUTPUT_SIZE);
|
|
||||||
|
|
||||||
let truncated = PtcScriptTool::truncate_output(&long);
|
|
||||||
// Must not panic and must contain valid UTF-8
|
|
||||||
assert!(truncated.contains("[Output truncated"));
|
|
||||||
// The kept portion must end on a char boundary (valid UTF-8 guaranteed by compilation)
|
|
||||||
let kept = truncated.split("\n\n[Output truncated").next().unwrap();
|
|
||||||
assert!(kept.len() <= MAX_OUTPUT_SIZE);
|
|
||||||
// Every char should be complete (no partial emoji)
|
|
||||||
assert!(kept.chars().all(|c| c == '\u{1F600}'));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_tool_metadata() {
|
|
||||||
let tool = PtcScriptTool::new();
|
|
||||||
assert_eq!(tool.name(), "ptc_script");
|
|
||||||
assert_eq!(tool.domain(), ToolDomain::Container);
|
|
||||||
assert_eq!(
|
|
||||||
tool.requires_approval(&serde_json::json!({})),
|
|
||||||
ApprovalRequirement::Always
|
|
||||||
);
|
|
||||||
assert!(tool.requires_sanitization());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -1,490 +0,0 @@
|
|||||||
//! 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;
|
|
||||||
use crate::tools::tool::ToolDomain;
|
|
||||||
|
|
||||||
/// 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 },
|
|
||||||
|
|
||||||
#[error("Tool {name} has domain Container and cannot be executed on the orchestrator")]
|
|
||||||
DomainBlocked { name: String },
|
|
||||||
}
|
|
||||||
|
|
||||||
/// 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(),
|
|
||||||
})?;
|
|
||||||
|
|
||||||
// Reject Container-domain tools — they must run inside a sandbox,
|
|
||||||
// not on the orchestrator host. Without this check a compromised
|
|
||||||
// worker could invoke shell/file tools on the host (sandbox escape).
|
|
||||||
if tool.domain() == ToolDomain::Container {
|
|
||||||
return Err(PtcError::DomainBlocked {
|
|
||||||
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(|| tool.execution_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, ToolDomain, 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),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// A tool that declares Container domain — must be blocked by the executor.
|
|
||||||
struct ContainerDomainTool;
|
|
||||||
|
|
||||||
#[async_trait::async_trait]
|
|
||||||
impl Tool for ContainerDomainTool {
|
|
||||||
fn name(&self) -> &str {
|
|
||||||
"container_tool"
|
|
||||||
}
|
|
||||||
fn description(&self) -> &str {
|
|
||||||
"Simulates a container-domain tool"
|
|
||||||
}
|
|
||||||
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> {
|
|
||||||
Ok(ToolOutput::text(
|
|
||||||
"should not reach here",
|
|
||||||
Duration::from_millis(1),
|
|
||||||
))
|
|
||||||
}
|
|
||||||
fn domain(&self) -> ToolDomain {
|
|
||||||
ToolDomain::Container
|
|
||||||
}
|
|
||||||
fn requires_sanitization(&self) -> bool {
|
|
||||||
false
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_container_domain_blocked() {
|
|
||||||
let tools = Arc::new(ToolRegistry::new());
|
|
||||||
tools.register(Arc::new(ContainerDomainTool)).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("container_tool", serde_json::json!({}), &ctx, None)
|
|
||||||
.await;
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
matches!(result, Err(PtcError::DomainBlocked { .. })),
|
|
||||||
"Container-domain tools must be rejected: {:?}",
|
|
||||||
result
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[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
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -18,7 +18,6 @@ pub mod redaction;
|
|||||||
pub mod schema_validator;
|
pub mod schema_validator;
|
||||||
pub mod wasm;
|
pub mod wasm;
|
||||||
|
|
||||||
mod executor;
|
|
||||||
mod registry;
|
mod registry;
|
||||||
mod tool;
|
mod tool;
|
||||||
|
|
||||||
@@ -32,7 +31,6 @@ pub use builder::{
|
|||||||
TestCase, TestHarness, TestResult, TestSuite, ValidationError, ValidationResult, WasmValidator,
|
TestCase, TestHarness, TestResult, TestSuite, ValidationError, ValidationResult, WasmValidator,
|
||||||
};
|
};
|
||||||
pub(crate) use coercion::prepare_tool_params;
|
pub(crate) use coercion::prepare_tool_params;
|
||||||
pub use executor::{PtcError, PtcToolResult, ToolExecutor};
|
|
||||||
pub use rate_limiter::RateLimiter;
|
pub use rate_limiter::RateLimiter;
|
||||||
pub use registry::ToolRegistry;
|
pub use registry::ToolRegistry;
|
||||||
pub use tool::{
|
pub use tool::{
|
||||||
|
|||||||
+5
-73
@@ -19,12 +19,11 @@ use crate::tools::builder::{
|
|||||||
use crate::tools::builtin::{
|
use crate::tools::builtin::{
|
||||||
ApplyPatchTool, CancelJobTool, CreateJobTool, EchoTool, ExtensionInfoTool, HttpTool,
|
ApplyPatchTool, CancelJobTool, CreateJobTool, EchoTool, ExtensionInfoTool, HttpTool,
|
||||||
JobEventsTool, JobPromptTool, JobStatusTool, JsonTool, ListDirTool, ListJobsTool,
|
JobEventsTool, JobPromptTool, JobStatusTool, JsonTool, ListDirTool, ListJobsTool,
|
||||||
MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool, PromptQueue, PtcScriptTool,
|
MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool, PromptQueue, ReadFileTool,
|
||||||
ReadFileTool, ShellTool, SkillInstallTool, SkillListTool, SkillRemoveTool, SkillSearchTool,
|
ShellTool, SkillInstallTool, SkillListTool, SkillRemoveTool, SkillSearchTool, TimeTool,
|
||||||
TimeTool, ToolActivateTool, ToolAuthTool, ToolInstallTool, ToolListTool, ToolRemoveTool,
|
ToolActivateTool, ToolAuthTool, ToolInstallTool, ToolListTool, ToolRemoveTool, ToolSearchTool,
|
||||||
ToolSearchTool, ToolUpgradeTool, WriteFileTool,
|
ToolUpgradeTool, WriteFileTool,
|
||||||
};
|
};
|
||||||
use crate::tools::executor::ToolExecutor;
|
|
||||||
use crate::tools::rate_limiter::RateLimiter;
|
use crate::tools::rate_limiter::RateLimiter;
|
||||||
use crate::tools::tool::{ApprovalRequirement, Tool, ToolDomain};
|
use crate::tools::tool::{ApprovalRequirement, Tool, ToolDomain};
|
||||||
use crate::tools::wasm::{
|
use crate::tools::wasm::{
|
||||||
@@ -79,7 +78,6 @@ const PROTECTED_TOOL_NAMES: &[&str] = &[
|
|||||||
"image_edit",
|
"image_edit",
|
||||||
"image_analyze",
|
"image_analyze",
|
||||||
"tool_info",
|
"tool_info",
|
||||||
"ptc_script",
|
|
||||||
];
|
];
|
||||||
|
|
||||||
/// Registry of available tools.
|
/// Registry of available tools.
|
||||||
@@ -95,14 +93,6 @@ pub struct ToolRegistry {
|
|||||||
rate_limiter: RateLimiter,
|
rate_limiter: RateLimiter,
|
||||||
/// Reference to the message tool for setting context per-turn.
|
/// Reference to the message tool for setting context per-turn.
|
||||||
message_tool: RwLock<Option<Arc<crate::tools::builtin::MessageTool>>>,
|
message_tool: RwLock<Option<Arc<crate::tools::builtin::MessageTool>>>,
|
||||||
/// Shared slot for the tool executor (enables PTC via tool_invoke).
|
|
||||||
///
|
|
||||||
/// Uses `std::sync::RwLock` (not tokio) because reads happen inside
|
|
||||||
/// `spawn_blocking` closures in WASM tool execution. The slot is
|
|
||||||
/// populated lazily after `AppBuilder::build_all()` completes, so
|
|
||||||
/// WASM tools registered during startup still get access to the
|
|
||||||
/// executor when they execute later.
|
|
||||||
tool_executor_slot: Arc<std::sync::RwLock<Option<Arc<ToolExecutor>>>>,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
impl ToolRegistry {
|
impl ToolRegistry {
|
||||||
@@ -124,7 +114,6 @@ impl ToolRegistry {
|
|||||||
secrets_store: None,
|
secrets_store: None,
|
||||||
rate_limiter: RateLimiter::new(),
|
rate_limiter: RateLimiter::new(),
|
||||||
message_tool: RwLock::new(None),
|
message_tool: RwLock::new(None),
|
||||||
tool_executor_slot: Arc::new(std::sync::RwLock::new(None)),
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -149,27 +138,6 @@ impl ToolRegistry {
|
|||||||
&self.rate_limiter
|
&self.rate_limiter
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Set the tool executor for programmatic tool calling (PTC).
|
|
||||||
///
|
|
||||||
/// Writes the executor into the shared slot so all WASM tools --
|
|
||||||
/// including those registered before this call -- can resolve it
|
|
||||||
/// lazily at execution time.
|
|
||||||
pub fn set_tool_executor(&self, executor: Arc<ToolExecutor>) {
|
|
||||||
if let Ok(mut guard) = self.tool_executor_slot.write() {
|
|
||||||
*guard = Some(executor);
|
|
||||||
} else {
|
|
||||||
tracing::error!("tool_executor_slot RwLock is poisoned; PTC will be unavailable");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Get a clone of the shared tool executor slot.
|
|
||||||
///
|
|
||||||
/// WASM wrappers hold this slot and read from it at execution time,
|
|
||||||
/// allowing the executor to be set after tool registration.
|
|
||||||
pub fn tool_executor_slot(&self) -> Arc<std::sync::RwLock<Option<Arc<ToolExecutor>>>> {
|
|
||||||
Arc::clone(&self.tool_executor_slot)
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Register a tool. Rejects dynamic tools that try to shadow a protected built-in name.
|
/// Register a tool. Rejects dynamic tools that try to shadow a protected built-in name.
|
||||||
pub async fn register(&self, tool: Arc<dyn Tool>) {
|
pub async fn register(&self, tool: Arc<dyn Tool>) {
|
||||||
let name = tool.name().to_string();
|
let name = tool.name().to_string();
|
||||||
@@ -362,9 +330,8 @@ impl ToolRegistry {
|
|||||||
self.register_sync(Arc::new(WriteFileTool::new()));
|
self.register_sync(Arc::new(WriteFileTool::new()));
|
||||||
self.register_sync(Arc::new(ListDirTool::new()));
|
self.register_sync(Arc::new(ListDirTool::new()));
|
||||||
self.register_sync(Arc::new(ApplyPatchTool::new()));
|
self.register_sync(Arc::new(ApplyPatchTool::new()));
|
||||||
self.register_sync(Arc::new(PtcScriptTool::new()));
|
|
||||||
|
|
||||||
tracing::debug!("Registered 6 development tools");
|
tracing::debug!("Registered 5 development tools");
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Register memory tools with a workspace.
|
/// Register memory tools with a workspace.
|
||||||
@@ -692,11 +659,6 @@ impl ToolRegistry {
|
|||||||
wrapper = wrapper.with_oauth_refresh(oauth);
|
wrapper = wrapper.with_oauth_refresh(oauth);
|
||||||
}
|
}
|
||||||
|
|
||||||
// Inject shared tool executor slot for PTC (lazy resolution).
|
|
||||||
// The WASM wrapper reads from this slot at execution time, so the
|
|
||||||
// executor can be set after tool registration.
|
|
||||||
wrapper = wrapper.with_tool_executor_slot(Arc::clone(&self.tool_executor_slot));
|
|
||||||
|
|
||||||
// Register the tool
|
// Register the tool
|
||||||
self.register(Arc::new(wrapper)).await;
|
self.register(Arc::new(wrapper)).await;
|
||||||
|
|
||||||
@@ -927,36 +889,6 @@ mod tests {
|
|||||||
assert!(def.parameters.get("extra").is_none());
|
assert!(def.parameters.get("extra").is_none());
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn test_tool_executor_slot_lazy_resolution() {
|
|
||||||
let registry = ToolRegistry::new();
|
|
||||||
|
|
||||||
// Get the slot BEFORE setting the executor (simulates startup order)
|
|
||||||
let slot = registry.tool_executor_slot();
|
|
||||||
|
|
||||||
// Slot should be empty
|
|
||||||
assert!(slot.read().unwrap().is_none());
|
|
||||||
|
|
||||||
// Set the executor (simulates main.rs wiring after build_all)
|
|
||||||
let tools = Arc::new(ToolRegistry::new());
|
|
||||||
tools.register_builtin_tools();
|
|
||||||
let safety = Arc::new(crate::safety::SafetyLayer::new(
|
|
||||||
&crate::config::SafetyConfig {
|
|
||||||
max_output_length: 100_000,
|
|
||||||
injection_check_enabled: true,
|
|
||||||
},
|
|
||||||
));
|
|
||||||
let executor = Arc::new(crate::tools::ToolExecutor::new(
|
|
||||||
tools,
|
|
||||||
safety,
|
|
||||||
std::time::Duration::from_secs(60),
|
|
||||||
));
|
|
||||||
registry.set_tool_executor(Arc::clone(&executor));
|
|
||||||
|
|
||||||
// Slot should now contain the executor
|
|
||||||
assert!(slot.read().unwrap().is_some());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_builtin_tool_cannot_be_shadowed() {
|
async fn test_builtin_tool_cannot_be_shadowed() {
|
||||||
let registry = ToolRegistry::new();
|
let registry = ToolRegistry::new();
|
||||||
|
|||||||
+9
-383
@@ -19,7 +19,6 @@ use wasmtime_wasi::{ResourceTable, WasiCtx, WasiCtxBuilder, WasiView};
|
|||||||
use crate::context::JobContext;
|
use crate::context::JobContext;
|
||||||
use crate::safety::LeakDetector;
|
use crate::safety::LeakDetector;
|
||||||
use crate::secrets::SecretsStore;
|
use crate::secrets::SecretsStore;
|
||||||
use crate::tools::ToolExecutor;
|
|
||||||
use crate::tools::tool::{Tool, ToolError, ToolOutput};
|
use crate::tools::tool::{Tool, ToolError, ToolOutput};
|
||||||
use crate::tools::wasm::capabilities::Capabilities;
|
use crate::tools::wasm::capabilities::Capabilities;
|
||||||
use crate::tools::wasm::credential_injector::{
|
use crate::tools::wasm::credential_injector::{
|
||||||
@@ -30,26 +29,6 @@ use crate::tools::wasm::host::{HostState, LogLevel};
|
|||||||
use crate::tools::wasm::limits::{ResourceLimits, WasmResourceLimiter};
|
use crate::tools::wasm::limits::{ResourceLimits, WasmResourceLimiter};
|
||||||
use crate::tools::wasm::runtime::{EPOCH_TICK_INTERVAL, PreparedModule, WasmToolRuntime};
|
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>;
|
|
||||||
|
|
||||||
/// RAII guard that decrements the nesting depth counter on drop, ensuring the
|
|
||||||
/// counter is restored even if the code between increment and decrement panics.
|
|
||||||
struct NestingGuard<'a> {
|
|
||||||
depth: &'a mut u32,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl Drop for NestingGuard<'_> {
|
|
||||||
fn drop(&mut self) {
|
|
||||||
*self.depth = self.depth.saturating_sub(1);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Generate component model bindings from the WIT file.
|
// Generate component model bindings from the WIT file.
|
||||||
//
|
//
|
||||||
// This creates:
|
// This creates:
|
||||||
@@ -120,11 +99,6 @@ struct StoreData {
|
|||||||
/// Dedicated tokio runtime for HTTP requests, lazily initialized.
|
/// Dedicated tokio runtime for HTTP requests, lazily initialized.
|
||||||
/// Reused across multiple `http_request` calls within one execution.
|
/// Reused across multiple `http_request` calls within one execution.
|
||||||
http_runtime: Option<tokio::runtime::Runtime>,
|
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 {
|
impl StoreData {
|
||||||
@@ -133,7 +107,6 @@ impl StoreData {
|
|||||||
capabilities: Capabilities,
|
capabilities: Capabilities,
|
||||||
credentials: HashMap<String, String>,
|
credentials: HashMap<String, String>,
|
||||||
host_credentials: Vec<ResolvedHostCredential>,
|
host_credentials: Vec<ResolvedHostCredential>,
|
||||||
tool_resolver: Option<ToolResolver>,
|
|
||||||
) -> Self {
|
) -> Self {
|
||||||
// Minimal WASI context: no filesystem, no env vars (security)
|
// Minimal WASI context: no filesystem, no env vars (security)
|
||||||
let wasi = WasiCtxBuilder::new().build();
|
let wasi = WasiCtxBuilder::new().build();
|
||||||
@@ -146,8 +119,6 @@ impl StoreData {
|
|||||||
credentials,
|
credentials,
|
||||||
host_credentials,
|
host_credentials,
|
||||||
http_runtime: None,
|
http_runtime: None,
|
||||||
tool_resolver,
|
|
||||||
tool_nesting_depth: 0,
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -467,39 +438,14 @@ impl near::agent::host::Host for StoreData {
|
|||||||
result.map_err(|e| self.redact_credentials(&e))
|
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
|
// 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()?;
|
self.host_state.record_tool_invoke()?;
|
||||||
|
|
||||||
// Check nesting depth
|
// Tool invocation requires async context and access to the tool registry,
|
||||||
if self.tool_nesting_depth >= MAX_NESTING_DEPTH {
|
// which aren't available inside a synchronous WASM callback.
|
||||||
return Err(format!(
|
Err("Tool invocation from WASM tools is not yet supported".to_string())
|
||||||
"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 with RAII guard to ensure decrement even on panic
|
|
||||||
self.tool_nesting_depth += 1;
|
|
||||||
let current_depth = self.tool_nesting_depth;
|
|
||||||
let _guard = NestingGuard {
|
|
||||||
depth: &mut self.tool_nesting_depth,
|
|
||||||
};
|
|
||||||
// _guard drops at end of scope (or on panic), decrementing depth
|
|
||||||
resolver(&real_name, params, current_depth)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn secret_exists(&mut self, name: String) -> bool {
|
fn secret_exists(&mut self, name: String) -> bool {
|
||||||
@@ -530,11 +476,6 @@ pub struct WasmToolWrapper {
|
|||||||
secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>>,
|
secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>>,
|
||||||
/// OAuth refresh configuration for auto-refreshing expired tokens.
|
/// OAuth refresh configuration for auto-refreshing expired tokens.
|
||||||
oauth_refresh: Option<OAuthRefreshConfig>,
|
oauth_refresh: Option<OAuthRefreshConfig>,
|
||||||
/// Direct tool executor reference (for tests that wire it explicitly).
|
|
||||||
tool_executor: Option<Arc<ToolExecutor>>,
|
|
||||||
/// Shared slot for lazy executor resolution (production path).
|
|
||||||
/// Reads happen inside `spawn_blocking`, so this uses `std::sync::RwLock`.
|
|
||||||
tool_executor_slot: Option<Arc<std::sync::RwLock<Option<Arc<ToolExecutor>>>>>,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
@@ -623,8 +564,6 @@ impl WasmToolWrapper {
|
|||||||
credentials: HashMap::new(),
|
credentials: HashMap::new(),
|
||||||
secrets_store: None,
|
secrets_store: None,
|
||||||
oauth_refresh: None,
|
oauth_refresh: None,
|
||||||
tool_executor: None,
|
|
||||||
tool_executor_slot: None,
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -679,28 +618,6 @@ impl WasmToolWrapper {
|
|||||||
self
|
self
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Set the tool executor for programmatic tool calling (direct reference).
|
|
||||||
///
|
|
||||||
/// When set, the WASM `tool_invoke` host function can call other
|
|
||||||
/// registered tools synchronously via a bridged resolver closure.
|
|
||||||
/// Prefer `with_tool_executor_slot()` for production use.
|
|
||||||
pub fn with_tool_executor(mut self, executor: Arc<ToolExecutor>) -> Self {
|
|
||||||
self.tool_executor = Some(executor);
|
|
||||||
self
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Set the shared tool executor slot for lazy resolution.
|
|
||||||
///
|
|
||||||
/// The executor is read from this slot at execution time, allowing
|
|
||||||
/// it to be set after tool registration (production startup order).
|
|
||||||
pub fn with_tool_executor_slot(
|
|
||||||
mut self,
|
|
||||||
slot: Arc<std::sync::RwLock<Option<Arc<ToolExecutor>>>>,
|
|
||||||
) -> Self {
|
|
||||||
self.tool_executor_slot = Some(slot);
|
|
||||||
self
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Get the resource limits for this tool.
|
/// Get the resource limits for this tool.
|
||||||
pub fn limits(&self) -> &ResourceLimits {
|
pub fn limits(&self) -> &ResourceLimits {
|
||||||
&self.prepared.limits
|
&self.prepared.limits
|
||||||
@@ -729,7 +646,6 @@ impl WasmToolWrapper {
|
|||||||
params: serde_json::Value,
|
params: serde_json::Value,
|
||||||
context_json: Option<String>,
|
context_json: Option<String>,
|
||||||
host_credentials: Vec<ResolvedHostCredential>,
|
host_credentials: Vec<ResolvedHostCredential>,
|
||||||
tool_resolver: Option<ToolResolver>,
|
|
||||||
) -> Result<(String, Vec<crate::tools::wasm::host::LogEntry>), WasmError> {
|
) -> Result<(String, Vec<crate::tools::wasm::host::LogEntry>), WasmError> {
|
||||||
let engine = self.runtime.engine();
|
let engine = self.runtime.engine();
|
||||||
let limits = &self.prepared.limits;
|
let limits = &self.prepared.limits;
|
||||||
@@ -740,7 +656,6 @@ impl WasmToolWrapper {
|
|||||||
self.capabilities.clone(),
|
self.capabilities.clone(),
|
||||||
self.credentials.clone(),
|
self.credentials.clone(),
|
||||||
host_credentials,
|
host_credentials,
|
||||||
tool_resolver,
|
|
||||||
);
|
);
|
||||||
let mut store = Store::new(engine, store_data);
|
let mut store = Store::new(engine, store_data);
|
||||||
|
|
||||||
@@ -839,7 +754,6 @@ pub(super) fn extract_wasm_metadata(
|
|||||||
Capabilities::default(),
|
Capabilities::default(),
|
||||||
HashMap::new(),
|
HashMap::new(),
|
||||||
vec![],
|
vec![],
|
||||||
None,
|
|
||||||
);
|
);
|
||||||
let mut store = Store::new(engine, store_data);
|
let mut store = Store::new(engine, store_data);
|
||||||
|
|
||||||
@@ -939,48 +853,6 @@ impl Tool for WasmToolWrapper {
|
|||||||
// Serialize context for WASM
|
// Serialize context for WASM
|
||||||
let context_json = serde_json::to_string(ctx).ok();
|
let context_json = serde_json::to_string(ctx).ok();
|
||||||
|
|
||||||
// Resolve the tool executor: direct reference takes priority, then shared slot.
|
|
||||||
let resolved_executor: Option<Arc<ToolExecutor>> =
|
|
||||||
self.tool_executor.as_ref().cloned().or_else(|| {
|
|
||||||
self.tool_executor_slot
|
|
||||||
.as_ref()
|
|
||||||
.and_then(|slot| slot.read().ok())
|
|
||||||
.and_then(|guard| guard.clone())
|
|
||||||
});
|
|
||||||
|
|
||||||
// Build a tool resolver closure if we have a tool executor.
|
|
||||||
// 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> = resolved_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
|
// Clone what we need for the blocking task
|
||||||
let runtime = Arc::clone(&self.runtime);
|
let runtime = Arc::clone(&self.runtime);
|
||||||
let prepared = Arc::clone(&self.prepared);
|
let prepared = Arc::clone(&self.prepared);
|
||||||
@@ -998,14 +870,12 @@ impl Tool for WasmToolWrapper {
|
|||||||
description,
|
description,
|
||||||
schemas,
|
schemas,
|
||||||
credentials,
|
credentials,
|
||||||
secrets_store: None, // Not needed in blocking task
|
secrets_store: None, // Not needed in blocking task
|
||||||
oauth_refresh: None, // Already used above for pre-refresh
|
oauth_refresh: None, // Already used above for pre-refresh
|
||||||
tool_executor: None, // Resolver closure captures the executor
|
|
||||||
tool_executor_slot: None, // Resolver closure captures the executor
|
|
||||||
};
|
};
|
||||||
|
|
||||||
tokio::task::spawn_blocking(move || {
|
tokio::task::spawn_blocking(move || {
|
||||||
wrapper.execute_sync(params, context_json, host_credentials, tool_resolver)
|
wrapper.execute_sync(params, context_json, host_credentials)
|
||||||
})
|
})
|
||||||
.await
|
.await
|
||||||
.map_err(|e| WasmError::ExecutionPanicked(e.to_string()))?
|
.map_err(|e| WasmError::ExecutionPanicked(e.to_string()))?
|
||||||
@@ -1516,7 +1386,6 @@ fn build_tool_usage_hint(tool_name: &str, schema: &serde_json::Value) -> String
|
|||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use std::collections::HashMap;
|
|
||||||
use std::sync::{Arc, Mutex};
|
use std::sync::{Arc, Mutex};
|
||||||
|
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
@@ -1533,12 +1402,10 @@ mod tests {
|
|||||||
TEST_GOOGLE_OAUTH_TOKEN, TEST_OAUTH_CLIENT_ID, TEST_OAUTH_CLIENT_SECRET,
|
TEST_GOOGLE_OAUTH_TOKEN, TEST_OAUTH_CLIENT_ID, TEST_OAUTH_CLIENT_SECRET,
|
||||||
test_secrets_store,
|
test_secrets_store,
|
||||||
};
|
};
|
||||||
use crate::tools::tool::{Tool, ToolError, ToolOutput};
|
use crate::tools::tool::Tool;
|
||||||
use crate::tools::wasm::capabilities::Capabilities;
|
use crate::tools::wasm::capabilities::Capabilities;
|
||||||
use crate::tools::wasm::runtime::{WasmRuntimeConfig, WasmToolRuntime};
|
use crate::tools::wasm::runtime::{WasmRuntimeConfig, WasmToolRuntime};
|
||||||
|
|
||||||
use super::WasmToolWrapper;
|
|
||||||
|
|
||||||
struct RecordingSecretsStore {
|
struct RecordingSecretsStore {
|
||||||
inner: InMemorySecretsStore,
|
inner: InMemorySecretsStore,
|
||||||
get_decrypted_lookups: Mutex<Vec<(String, String)>>,
|
get_decrypted_lookups: Mutex<Vec<(String, String)>>,
|
||||||
@@ -1766,7 +1633,6 @@ mod tests {
|
|||||||
Capabilities::default(),
|
Capabilities::default(),
|
||||||
HashMap::new(),
|
HashMap::new(),
|
||||||
host_credentials,
|
host_credentials,
|
||||||
None,
|
|
||||||
);
|
);
|
||||||
|
|
||||||
// Should inject for matching host
|
// Should inject for matching host
|
||||||
@@ -1806,7 +1672,6 @@ mod tests {
|
|||||||
Capabilities::default(),
|
Capabilities::default(),
|
||||||
HashMap::new(),
|
HashMap::new(),
|
||||||
host_credentials,
|
host_credentials,
|
||||||
None,
|
|
||||||
);
|
);
|
||||||
|
|
||||||
let mut headers = HashMap::new();
|
let mut headers = HashMap::new();
|
||||||
@@ -1833,7 +1698,6 @@ mod tests {
|
|||||||
Capabilities::default(),
|
Capabilities::default(),
|
||||||
HashMap::new(),
|
HashMap::new(),
|
||||||
host_credentials,
|
host_credentials,
|
||||||
None,
|
|
||||||
);
|
);
|
||||||
|
|
||||||
let text = "Error: request to https://api.example.com?key=super-secret-token failed";
|
let text = "Error: request to https://api.example.com?key=super-secret-token failed";
|
||||||
@@ -2320,244 +2184,6 @@ mod tests {
|
|||||||
assert!(hint.contains("native JSON arrays/objects")); // safety: test-only assertion
|
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 = crate::tools::coercion::prepare_params_for_schema(¶ms, &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 = crate::tools::coercion::prepare_params_for_schema(¶ms, &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;
|
|
||||||
|
|
||||||
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 = crate::tools::coercion::prepare_params_for_schema(¶ms, &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
|
/// Regression test: leak scan must run on raw headers (before credential
|
||||||
/// injection), not after. If it ran post-injection, the host-injected
|
/// injection), not after. If it ran post-injection, the host-injected
|
||||||
/// Slack bot token (`xoxb-...`) would trigger a Block and reject the
|
/// Slack bot token (`xoxb-...`) would trigger a Block and reject the
|
||||||
|
|||||||
@@ -114,36 +114,6 @@ pub struct CredentialResponse {
|
|||||||
pub value: String,
|
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>,
|
|
||||||
/// Current nesting depth for tool-invokes-tool chains.
|
|
||||||
/// Defaults to 0 for top-level calls (backward compatible).
|
|
||||||
#[serde(default)]
|
|
||||||
pub nesting_depth: u32,
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Response from a programmatic tool call.
|
|
||||||
#[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 {
|
impl WorkerHttpClient {
|
||||||
/// Create a new client from environment.
|
/// Create a new client from environment.
|
||||||
///
|
///
|
||||||
@@ -429,14 +399,6 @@ 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.
|
/// Signal job completion to the orchestrator.
|
||||||
pub async fn report_complete(&self, report: &CompletionReport) -> Result<(), WorkerError> {
|
pub async fn report_complete(&self, report: &CompletionReport) -> Result<(), WorkerError> {
|
||||||
let _: serde_json::Value = self
|
let _: serde_json::Value = self
|
||||||
|
|||||||
@@ -65,7 +65,6 @@ async fn core_registration_covers_expected_tools() {
|
|||||||
"http",
|
"http",
|
||||||
"json",
|
"json",
|
||||||
"list_dir",
|
"list_dir",
|
||||||
"ptc_script",
|
|
||||||
"read_file",
|
"read_file",
|
||||||
"shell",
|
"shell",
|
||||||
"time",
|
"time",
|
||||||
|
|||||||
@@ -1,20 +0,0 @@
|
|||||||
[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]
|
|
||||||
@@ -1,52 +0,0 @@
|
|||||||
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);
|
|
||||||
Reference in New Issue
Block a user