pub mod engine;
pub mod sandbox;
pub mod types;
pub use engine::{ToolExecutor, ToolOrchestrator, dynamic_to_json};
pub use sandbox::{
DEFAULT_MAX_ARRAY_SIZE,
DEFAULT_MAX_MAP_SIZE,
DEFAULT_MAX_OPERATIONS,
DEFAULT_MAX_STRING_SIZE,
DEFAULT_MAX_TOOL_CALLS,
DEFAULT_TIMEOUT_MS,
EXTENDED_MAX_OPERATIONS,
EXTENDED_MAX_TOOL_CALLS,
EXTENDED_TIMEOUT_MS,
ExecutionLimits,
QUICK_MAX_OPERATIONS,
QUICK_MAX_TOOL_CALLS,
QUICK_TIMEOUT_MS,
};
pub use types::{OrchestratorError, OrchestratorResult, ToolCall};
use serde::Deserialize;
use serde_json::{Value, json};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use brainwires_core::{Tool, ToolContext, ToolInputSchema, ToolResult};
pub struct OrchestratorTool {
orchestrator: Arc<RwLock<ToolOrchestrator>>,
}
impl OrchestratorTool {
pub fn new() -> Self {
Self {
orchestrator: Arc::new(RwLock::new(ToolOrchestrator::new())),
}
}
pub fn get_tools() -> Vec<Tool> {
vec![Self::execute_script_tool()]
}
fn execute_script_tool() -> Tool {
let mut properties = HashMap::new();
properties.insert(
"script".to_string(),
json!({
"type": "string",
"description": "The Rhai script to execute. Use registered tool functions (e.g., read_file(path), search_code(pattern)) and return a final value."
}),
);
properties.insert(
"max_operations".to_string(),
json!({
"type": "integer",
"description": "Maximum operations allowed (default: 100000)",
"default": 100000
}),
);
properties.insert(
"max_tool_calls".to_string(),
json!({
"type": "integer",
"description": "Maximum tool calls allowed (default: 50)",
"default": 50
}),
);
properties.insert(
"timeout_ms".to_string(),
json!({
"type": "integer",
"description": "Timeout in milliseconds (default: 30000)",
"default": 30000
}),
);
Tool {
name: "execute_script".to_string(),
description:
r#"PRIMARY TOOL: Execute a Rhai script for programmatic tool orchestration.
This is the preferred way to interact with tools. Write Rhai scripts to orchestrate
multiple tool calls efficiently, with intermediate results staying out of the context window.
Benefits:
- 37% token reduction vs sequential tool calls
- Loops, conditionals, and data transformation
- Batch operations in a single execution
- Only final result enters context
Available tools can be discovered via `search_tools`. All tools are callable as functions.
Rhai Syntax Quick Reference:
- Variables: `let x = 42;`
- Strings: `let s = "hello";` or template: `` `Hello ${name}` ``
- Arrays: `let arr = [1, 2, 3];`
- Objects: `let obj = #{ key: "value" };`
- Loops: `for item in items { ... }`
- Conditionals: `if condition { ... } else { ... }`
Example - Find and count TODOs:
```rhai
let files = list_directory("src");
let count = 0;
for file in files {
if file.ends_with(".rs") {
let content = read_file(file);
count += content.matches("TODO").len();
}
}
`Found ${count} TODO comments`
```"#
.to_string(),
input_schema: ToolInputSchema::object(properties, vec!["script".to_string()]),
requires_approval: false,
defer_loading: false,
..Default::default()
}
}
pub async fn register_executor<F>(&self, name: impl Into<String>, executor: F)
where
F: Fn(serde_json::Value) -> Result<String, String> + Send + Sync + 'static,
{
let mut orchestrator = self.orchestrator.write().await;
orchestrator.register_executor(name, executor);
}
pub async fn execute(
&self,
tool_use_id: &str,
tool_name: &str,
input: &Value,
_context: &ToolContext,
) -> ToolResult {
let result = match tool_name {
"execute_script" => self.execute_script(input).await,
_ => Err(anyhow::anyhow!("Unknown orchestrator tool: {}", tool_name)),
};
match result {
Ok(output) => ToolResult::success(tool_use_id.to_string(), output),
Err(e) => ToolResult::error(
tool_use_id.to_string(),
format!("Script execution failed: {}", e),
),
}
}
async fn execute_script(&self, input: &Value) -> anyhow::Result<String> {
#[derive(Deserialize)]
struct Input {
script: String,
#[serde(default = "default_max_ops")]
max_operations: u64,
#[serde(default = "default_max_calls")]
max_tool_calls: usize,
#[serde(default = "default_timeout")]
timeout_ms: u64,
}
fn default_max_ops() -> u64 {
100_000
}
fn default_max_calls() -> usize {
50
}
fn default_timeout() -> u64 {
30_000
}
let params: Input = serde_json::from_value(input.clone())?;
let limits = ExecutionLimits::default()
.with_max_operations(params.max_operations)
.with_max_tool_calls(params.max_tool_calls)
.with_timeout_ms(params.timeout_ms);
let orchestrator = self.orchestrator.read().await;
let result = orchestrator.execute(¶ms.script, limits)?;
if result.success {
let mut output = result.output;
if !result.tool_calls.is_empty() {
output.push_str(&format!(
"\n\n--- Script executed {} tool call(s) in {}ms ---",
result.tool_calls.len(),
result.execution_time_ms
));
}
Ok(output)
} else {
Err(anyhow::anyhow!(
result.error.unwrap_or_else(|| "Unknown error".to_string())
))
}
}
pub fn orchestrator(&self) -> Arc<RwLock<ToolOrchestrator>> {
Arc::clone(&self.orchestrator)
}
}
impl Default for OrchestratorTool {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_get_tools() {
let tools = OrchestratorTool::get_tools();
assert_eq!(tools.len(), 1);
assert_eq!(tools[0].name, "execute_script");
assert!(!tools[0].requires_approval);
assert!(!tools[0].defer_loading);
}
#[test]
fn test_execute_script_tool_definition() {
let tool = OrchestratorTool::execute_script_tool();
assert_eq!(tool.name, "execute_script");
assert!(tool.description.contains("Rhai"));
assert!(tool.description.contains("programmatic"));
}
#[tokio::test]
async fn test_orchestrator_creation() {
let orchestrator_tool = OrchestratorTool::new();
let orchestrator = orchestrator_tool.orchestrator.read().await;
assert!(orchestrator.registered_tools().is_empty());
}
#[tokio::test]
async fn test_register_executor() {
let orchestrator_tool = OrchestratorTool::new();
orchestrator_tool
.register_executor("test_tool", |_| Ok("success".to_string()))
.await;
let orchestrator = orchestrator_tool.orchestrator.read().await;
assert!(orchestrator.registered_tools().contains(&"test_tool"));
}
#[tokio::test]
async fn test_execute_simple_script() {
let orchestrator_tool = OrchestratorTool::new();
let context = ToolContext::default();
let input = json!({
"script": "let x = 1 + 2; x"
});
let result = orchestrator_tool
.execute("test-id", "execute_script", &input, &context)
.await;
assert!(!result.is_error);
assert!(result.content.contains("3"));
}
#[tokio::test]
async fn test_execute_with_tool() {
let orchestrator_tool = OrchestratorTool::new();
orchestrator_tool
.register_executor("greet", |input| {
let name = input.as_str().unwrap_or("world");
Ok(format!("Hello, {}!", name))
})
.await;
let context = ToolContext::default();
let input = json!({
"script": r#"greet("Claude")"#
});
let result = orchestrator_tool
.execute("test-id", "execute_script", &input, &context)
.await;
assert!(!result.is_error);
assert!(result.content.contains("Hello, Claude!"));
}
#[tokio::test]
async fn test_execute_unknown_tool_name() {
let orchestrator_tool = OrchestratorTool::new();
let context = ToolContext::default();
let result = orchestrator_tool
.execute("test-id", "unknown_tool", &json!({}), &context)
.await;
assert!(result.is_error);
assert!(result.content.contains("Unknown orchestrator tool"));
}
}