437 lines
14 KiB
Rust
437 lines
14 KiB
Rust
use rtx_llm_tools::{
|
|
error::{LlmToolsError, ToolExecutionError},
|
|
tools::{
|
|
ExecutionEnvironment, ParameterType, ResourceLimits, SecurityPolicy, ToolDefinition,
|
|
ToolExecutor, ToolParameter,
|
|
},
|
|
};
|
|
use serde_json::{Value, json};
|
|
use std::collections::HashMap;
|
|
use std::time::Duration;
|
|
|
|
#[tokio::test]
|
|
async fn test_basic_tool_execution() {
|
|
let mut executor = ToolExecutor::new("test_executor");
|
|
|
|
// Register a simple calculator tool
|
|
let calculator = ToolDefinition::new(
|
|
"calculator",
|
|
"Basic arithmetic calculator",
|
|
vec![
|
|
ToolParameter::new("operation", ParameterType::String, true)
|
|
.with_description("The arithmetic operation to perform")
|
|
.with_enum_values(vec!["add", "subtract", "multiply", "divide"]),
|
|
ToolParameter::new("a", ParameterType::Number, true).with_description("First number"),
|
|
ToolParameter::new("b", ParameterType::Number, true).with_description("Second number"),
|
|
],
|
|
);
|
|
|
|
executor.register_tool(calculator).await.unwrap();
|
|
|
|
let mut params = HashMap::new();
|
|
params.insert("operation".to_string(), Value::String("add".to_string()));
|
|
params.insert("a".to_string(), Value::Number(serde_json::Number::from(10)));
|
|
params.insert("b".to_string(), Value::Number(serde_json::Number::from(5)));
|
|
|
|
let result = executor.execute("calculator", params).await.unwrap();
|
|
assert!(result.success);
|
|
assert_eq!(result.output, "15");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_parameter_validation() {
|
|
let mut executor = ToolExecutor::new("validator");
|
|
|
|
let tool = ToolDefinition::new(
|
|
"email_sender",
|
|
"Send emails",
|
|
vec![
|
|
ToolParameter::new("to", ParameterType::String, true)
|
|
.with_validation_regex(r"^[^@]+@[^@]+\.[^@]+$")
|
|
.with_description("Email address"),
|
|
ToolParameter::new("subject", ParameterType::String, true)
|
|
.with_max_length(100)
|
|
.with_description("Email subject"),
|
|
ToolParameter::new("body", ParameterType::String, false)
|
|
.with_max_length(1000)
|
|
.with_description("Email body"),
|
|
],
|
|
);
|
|
|
|
executor.register_tool(tool).await.unwrap();
|
|
|
|
// Test invalid email
|
|
let mut invalid_params = HashMap::new();
|
|
invalid_params.insert("to".to_string(), Value::String("not-an-email".to_string()));
|
|
invalid_params.insert("subject".to_string(), Value::String("Test".to_string()));
|
|
|
|
let result = executor.execute("email_sender", invalid_params).await;
|
|
assert!(result.is_err());
|
|
|
|
match result.unwrap_err() {
|
|
LlmToolsError::ToolExecution(ToolExecutionError::InvalidParameters { .. }) => {
|
|
// Expected
|
|
}
|
|
_ => panic!("Expected parameter validation error"),
|
|
}
|
|
|
|
// Test valid parameters
|
|
let mut valid_params = HashMap::new();
|
|
valid_params.insert(
|
|
"to".to_string(),
|
|
Value::String("[email protected]".to_string()),
|
|
);
|
|
valid_params.insert(
|
|
"subject".to_string(),
|
|
Value::String("Test Subject".to_string()),
|
|
);
|
|
valid_params.insert("body".to_string(), Value::String("Test body".to_string()));
|
|
|
|
let result = executor
|
|
.execute("email_sender", valid_params)
|
|
.await
|
|
.unwrap();
|
|
assert!(result.success);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_security_policy_enforcement() {
|
|
let mut executor = ToolExecutor::new("secure_executor");
|
|
|
|
// Create a security policy that blocks file system access
|
|
let mut policy = SecurityPolicy::new("strict");
|
|
policy.block_file_system_access(true);
|
|
policy.block_network_access(false);
|
|
policy.add_blocked_command("rm");
|
|
policy.add_blocked_command("del");
|
|
policy.set_max_execution_time(Duration::from_secs(5));
|
|
|
|
executor.set_security_policy(policy);
|
|
|
|
// Register a file system tool (should be blocked)
|
|
let file_tool = ToolDefinition::new(
|
|
"file_reader",
|
|
"Read file contents",
|
|
vec![
|
|
ToolParameter::new("path", ParameterType::String, true)
|
|
.with_description("File path to read"),
|
|
],
|
|
);
|
|
|
|
executor.register_tool(file_tool).await.unwrap();
|
|
|
|
let mut params = HashMap::new();
|
|
params.insert("path".to_string(), Value::String("/etc/passwd".to_string()));
|
|
|
|
let result = executor.execute("file_reader", params).await;
|
|
assert!(result.is_err());
|
|
|
|
match result.unwrap_err() {
|
|
LlmToolsError::ToolExecution(ToolExecutionError::SecurityViolation { .. }) => {
|
|
// Expected
|
|
}
|
|
_ => panic!("Expected security violation"),
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_resource_limits() {
|
|
let mut executor = ToolExecutor::new("resource_limited");
|
|
|
|
let mut limits = ResourceLimits::new();
|
|
limits.set_max_memory(1024 * 1024); // 1MB
|
|
limits.set_max_execution_time(Duration::from_millis(100));
|
|
limits.set_max_output_size(1000);
|
|
limits.set_max_cpu_percent(50.0);
|
|
|
|
executor.set_resource_limits(limits);
|
|
|
|
// Register a memory-intensive tool
|
|
let memory_tool = ToolDefinition::new(
|
|
"memory_hog",
|
|
"Allocates large amounts of memory",
|
|
vec![
|
|
ToolParameter::new("size_mb", ParameterType::Number, true)
|
|
.with_description("Memory size in MB"),
|
|
],
|
|
);
|
|
|
|
executor.register_tool(memory_tool).await.unwrap();
|
|
|
|
let mut params = HashMap::new();
|
|
params.insert(
|
|
"size_mb".to_string(),
|
|
Value::Number(serde_json::Number::from(10)),
|
|
); // Try to allocate 10MB
|
|
|
|
let result = executor.execute("memory_hog", params).await;
|
|
assert!(result.is_err());
|
|
|
|
match result.unwrap_err() {
|
|
LlmToolsError::ToolExecution(ToolExecutionError::ResourceExhausted { .. }) => {
|
|
// Expected
|
|
}
|
|
_ => panic!("Expected resource exhaustion"),
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_execution_timeout() {
|
|
let mut executor = ToolExecutor::new("timeout_test");
|
|
|
|
let slow_tool = ToolDefinition::new(
|
|
"slow_operation",
|
|
"A deliberately slow operation",
|
|
vec![
|
|
ToolParameter::new("delay_ms", ParameterType::Number, true)
|
|
.with_description("Delay in milliseconds"),
|
|
],
|
|
);
|
|
|
|
executor.register_tool(slow_tool).await.unwrap();
|
|
executor.set_global_timeout(Duration::from_millis(100));
|
|
|
|
let mut params = HashMap::new();
|
|
params.insert(
|
|
"delay_ms".to_string(),
|
|
Value::Number(serde_json::Number::from(500)),
|
|
); // 500ms delay
|
|
|
|
let result = executor.execute("slow_operation", params).await;
|
|
assert!(result.is_err());
|
|
|
|
match result.unwrap_err() {
|
|
LlmToolsError::ToolExecution(ToolExecutionError::Timeout { .. }) => {
|
|
// Expected
|
|
}
|
|
_ => panic!("Expected timeout error"),
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_sandbox_isolation() {
|
|
let mut executor = ToolExecutor::new("sandboxed");
|
|
executor.enable_sandbox(true);
|
|
|
|
// Create an environment with specific constraints
|
|
let mut env = ExecutionEnvironment::new();
|
|
env.set_working_directory("/tmp/sandbox");
|
|
env.add_environment_variable("SAFE_MODE", "true");
|
|
env.set_user("nobody");
|
|
env.set_group("nogroup");
|
|
|
|
executor.set_execution_environment(env);
|
|
|
|
let env_tool = ToolDefinition::new(
|
|
"env_checker",
|
|
"Check environment variables",
|
|
vec![
|
|
ToolParameter::new("var_name", ParameterType::String, true)
|
|
.with_description("Environment variable name"),
|
|
],
|
|
);
|
|
|
|
executor.register_tool(env_tool).await.unwrap();
|
|
|
|
let mut params = HashMap::new();
|
|
params.insert(
|
|
"var_name".to_string(),
|
|
Value::String("SAFE_MODE".to_string()),
|
|
);
|
|
|
|
let result = executor.execute("env_checker", params).await.unwrap();
|
|
assert!(result.success);
|
|
assert_eq!(result.output, "true");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_tool_chaining() {
|
|
let mut executor = ToolExecutor::new("chainer");
|
|
|
|
// Register multiple tools that can be chained
|
|
let data_generator = ToolDefinition::new(
|
|
"data_gen",
|
|
"Generate sample data",
|
|
vec![
|
|
ToolParameter::new("count", ParameterType::Number, true)
|
|
.with_description("Number of data points"),
|
|
],
|
|
);
|
|
|
|
let data_processor = ToolDefinition::new(
|
|
"data_proc",
|
|
"Process data",
|
|
vec![
|
|
ToolParameter::new("data", ParameterType::Array, true)
|
|
.with_description("Data to process"),
|
|
ToolParameter::new("operation", ParameterType::String, true)
|
|
.with_enum_values(vec!["sum", "avg", "max", "min"]),
|
|
],
|
|
);
|
|
|
|
executor.register_tool(data_generator).await.unwrap();
|
|
executor.register_tool(data_processor).await.unwrap();
|
|
|
|
// Execute first tool
|
|
let mut gen_params = HashMap::new();
|
|
gen_params.insert(
|
|
"count".to_string(),
|
|
Value::Number(serde_json::Number::from(5)),
|
|
);
|
|
|
|
let gen_result = executor.execute("data_gen", gen_params).await.unwrap();
|
|
assert!(gen_result.success);
|
|
|
|
// Use output of first tool as input to second tool
|
|
let generated_data: Vec<i32> = serde_json::from_str(&gen_result.output).unwrap();
|
|
|
|
let mut proc_params = HashMap::new();
|
|
proc_params.insert("data".to_string(), json!(generated_data));
|
|
proc_params.insert("operation".to_string(), Value::String("sum".to_string()));
|
|
|
|
let proc_result = executor.execute("data_proc", proc_params).await.unwrap();
|
|
assert!(proc_result.success);
|
|
|
|
let sum: i32 = proc_result.output.parse().unwrap();
|
|
let expected_sum: i32 = generated_data.iter().sum();
|
|
assert_eq!(sum, expected_sum);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_tool_dependency_validation() {
|
|
let mut executor = ToolExecutor::new("dependency_test");
|
|
|
|
let dependent_tool = ToolDefinition::new("dependent", "Tool that depends on others", vec![])
|
|
.with_dependencies(vec!["required_tool".to_string()]);
|
|
|
|
// Try to register tool without its dependency
|
|
let result = executor.register_tool(dependent_tool).await;
|
|
assert!(result.is_err());
|
|
|
|
match result.unwrap_err() {
|
|
LlmToolsError::ToolExecution(ToolExecutionError::DependencyNotMet { .. }) => {
|
|
// Expected
|
|
}
|
|
_ => panic!("Expected dependency error"),
|
|
}
|
|
|
|
// Register the required dependency first
|
|
let required_tool = ToolDefinition::new("required_tool", "Required dependency", vec![]);
|
|
|
|
executor.register_tool(required_tool).await.unwrap();
|
|
|
|
// Now the dependent tool should register successfully
|
|
let dependent_tool = ToolDefinition::new("dependent", "Tool that depends on others", vec![])
|
|
.with_dependencies(vec!["required_tool".to_string()]);
|
|
|
|
executor.register_tool(dependent_tool).await.unwrap();
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_tool_versioning() {
|
|
let mut executor = ToolExecutor::new("versioned");
|
|
|
|
// Register v1.0 of a tool
|
|
let tool_v1 = ToolDefinition::new(
|
|
"api_tool",
|
|
"API interaction tool v1.0",
|
|
vec![ToolParameter::new("endpoint", ParameterType::String, true)],
|
|
)
|
|
.with_version("1.0.0");
|
|
|
|
executor.register_tool(tool_v1).await.unwrap();
|
|
|
|
// Register v2.0 of the same tool (should coexist)
|
|
let tool_v2 = ToolDefinition::new(
|
|
"api_tool",
|
|
"API interaction tool v2.0 with auth",
|
|
vec![
|
|
ToolParameter::new("endpoint", ParameterType::String, true),
|
|
ToolParameter::new("auth_token", ParameterType::String, false),
|
|
],
|
|
)
|
|
.with_version("2.0.0");
|
|
|
|
executor.register_tool(tool_v2).await.unwrap();
|
|
|
|
// Execute specific version
|
|
let mut params = HashMap::new();
|
|
params.insert(
|
|
"endpoint".to_string(),
|
|
Value::String("https://api.example.com".to_string()),
|
|
);
|
|
|
|
let result_v1 = executor
|
|
.execute_version("api_tool", "1.0.0", params.clone())
|
|
.await
|
|
.unwrap();
|
|
assert!(result_v1.success);
|
|
|
|
// v2.0 execution should also work
|
|
params.insert(
|
|
"auth_token".to_string(),
|
|
Value::String("token123".to_string()),
|
|
);
|
|
let result_v2 = executor
|
|
.execute_version("api_tool", "2.0.0", params)
|
|
.await
|
|
.unwrap();
|
|
assert!(result_v2.success);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_sandbox_violation_detection() {
|
|
let mut executor = ToolExecutor::new("violation_detector");
|
|
executor.enable_sandbox(true);
|
|
|
|
let malicious_tool =
|
|
ToolDefinition::new("malicious", "Tool attempting malicious behavior", vec![]);
|
|
|
|
executor.register_tool(malicious_tool).await.unwrap();
|
|
|
|
// Tool attempts to access restricted resources
|
|
let result = executor.execute("malicious", HashMap::new()).await;
|
|
|
|
match result {
|
|
Err(LlmToolsError::ToolExecution(ToolExecutionError::SandboxViolation { .. })) => {
|
|
// Expected
|
|
}
|
|
Ok(_) => {
|
|
// Tool should be properly sandboxed and complete safely
|
|
}
|
|
Err(e) => panic!("Unexpected error: {:?}", e),
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_tool_audit_logging() {
|
|
let mut executor = ToolExecutor::new("audited");
|
|
executor.enable_audit_logging(true);
|
|
|
|
let audited_tool = ToolDefinition::new(
|
|
"sensitive_op",
|
|
"Sensitive operation requiring audit",
|
|
vec![ToolParameter::new("action", ParameterType::String, true)],
|
|
);
|
|
|
|
executor.register_tool(audited_tool).await.unwrap();
|
|
|
|
let mut params = HashMap::new();
|
|
params.insert(
|
|
"action".to_string(),
|
|
Value::String("delete_user".to_string()),
|
|
);
|
|
|
|
let result = executor.execute("sensitive_op", params).await.unwrap();
|
|
assert!(result.success);
|
|
|
|
// Check audit log was created
|
|
let audit_entries = executor.get_audit_log().await.unwrap();
|
|
assert!(!audit_entries.is_empty());
|
|
|
|
let last_entry = audit_entries.last().unwrap();
|
|
assert_eq!(last_entry.tool_name, "sensitive_op");
|
|
assert!(last_entry.parameters.contains_key("action"));
|
|
assert!(last_entry.timestamp > std::time::Instant::now() - Duration::from_secs(10));
|
|
}
|