Adding builder capability

This commit is contained in:
Illia Polosukhin
2026-02-03 08:36:55 -08:00
parent 343782524f
commit c9ebb117ab
19 changed files with 3689 additions and 210 deletions
Generated
+16
View File
@@ -1303,6 +1303,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1"
dependencies = [
"ahash 0.8.12",
"serde",
]
[[package]]
@@ -1983,6 +1984,7 @@ dependencies = [
"tracing",
"tracing-subscriber",
"uuid",
"wasmparser 0.220.1",
"wasmtime",
]
@@ -4074,6 +4076,20 @@ dependencies = [
"wasmparser 0.244.0",
]
[[package]]
name = "wasmparser"
version = "0.220.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8d07b6a3b550fefa1a914b6d54fc175dd11c3392da11eee604e6ffc759805d25"
dependencies = [
"ahash 0.8.12",
"bitflags 2.10.0",
"hashbrown 0.14.5",
"indexmap 2.13.0",
"semver",
"serde",
]
[[package]]
name = "wasmparser"
version = "0.221.3"
+1
View File
@@ -73,6 +73,7 @@ pgvector = { version = "0.4", features = ["postgres"] }
# WASM sandbox for untrusted tool execution
wasmtime = { version = "28", features = ["component-model"] }
wasmparser = "0.220" # WASM binary parsing for validation
# Cryptography for secrets management
aes-gcm = "0.10"
+69
View File
@@ -17,6 +17,7 @@ pub struct Config {
pub safety: SafetyConfig,
pub wasm: WasmConfig,
pub secrets: SecretsConfig,
pub builder: BuilderModeConfig,
}
impl Config {
@@ -33,6 +34,7 @@ impl Config {
safety: SafetyConfig::from_env()?,
wasm: WasmConfig::from_env()?,
secrets: SecretsConfig::from_env()?,
builder: BuilderModeConfig::from_env()?,
})
}
}
@@ -376,6 +378,73 @@ impl WasmConfig {
}
}
/// Builder mode configuration.
#[derive(Debug, Clone)]
pub struct BuilderModeConfig {
/// Whether the software builder tool is enabled.
pub enabled: bool,
/// Directory for build artifacts (default: temp dir).
pub build_dir: Option<PathBuf>,
/// Maximum iterations for the build loop.
pub max_iterations: u32,
/// Build timeout in seconds.
pub timeout_secs: u64,
/// Whether to automatically register built WASM tools.
pub auto_register: bool,
}
impl Default for BuilderModeConfig {
fn default() -> Self {
Self {
enabled: false,
build_dir: None,
max_iterations: 20,
timeout_secs: 600,
auto_register: true,
}
}
}
impl BuilderModeConfig {
fn from_env() -> Result<Self, ConfigError> {
Ok(Self {
enabled: optional_env("BUILDER_ENABLED")?
.map(|s| s.parse())
.transpose()
.map_err(|e| ConfigError::InvalidValue {
key: "BUILDER_ENABLED".to_string(),
message: format!("must be 'true' or 'false': {e}"),
})?
.unwrap_or(false),
build_dir: optional_env("BUILDER_DIR")?.map(PathBuf::from),
max_iterations: parse_optional_env("BUILDER_MAX_ITERATIONS", 20)?,
timeout_secs: parse_optional_env("BUILDER_TIMEOUT_SECS", 600)?,
auto_register: optional_env("BUILDER_AUTO_REGISTER")?
.map(|s| s.parse())
.transpose()
.map_err(|e| ConfigError::InvalidValue {
key: "BUILDER_AUTO_REGISTER".to_string(),
message: format!("must be 'true' or 'false': {e}"),
})?
.unwrap_or(true),
})
}
/// Convert to BuilderConfig for the builder tool.
pub fn to_builder_config(&self) -> crate::tools::BuilderConfig {
crate::tools::BuilderConfig {
build_dir: self.build_dir.clone().unwrap_or_else(std::env::temp_dir),
max_iterations: self.max_iterations,
timeout: Duration::from_secs(self.timeout_secs),
cleanup_on_failure: true,
validate_wasm: true,
run_tests: true,
auto_register: self.auto_register,
wasm_output_dir: None,
}
}
}
// Helper functions
fn required_env(key: &str) -> Result<String, ConfigError> {
+13 -1
View File
@@ -97,6 +97,8 @@ pub struct JobContext {
pub job_id: Uuid,
/// Current state.
pub state: JobState,
/// User ID that owns this job (for workspace scoping).
pub user_id: String,
/// Conversation ID if linked to a conversation.
pub conversation_id: Option<Uuid>,
/// Job title.
@@ -134,9 +136,19 @@ pub struct JobContext {
impl JobContext {
/// Create a new job context.
pub fn new(title: impl Into<String>, description: impl Into<String>) -> Self {
Self::with_user("default", title, description)
}
/// Create a new job context with a specific user ID.
pub fn with_user(
user_id: impl Into<String>,
title: impl Into<String>,
description: impl Into<String>,
) -> Self {
Self {
job_id: Uuid::new_v4(),
state: JobState::Pending,
user_id: user_id.into(),
conversation_id: None,
title: title.into(),
description: description.into(),
@@ -224,7 +236,7 @@ impl JobContext {
impl Default for JobContext {
fn default() -> Self {
Self::new("Untitled", "No description")
Self::with_user("default", "Untitled", "No description")
}
}
+9 -1
View File
@@ -47,6 +47,13 @@ impl Store {
Ok(self.pool.get().await?)
}
/// Get a clone of the database pool.
///
/// Useful for sharing the pool with other components like Workspace.
pub fn pool(&self) -> Pool {
self.pool.clone()
}
// ==================== Conversations ====================
/// Create a new conversation.
@@ -173,11 +180,12 @@ impl Store {
Ok(Some(JobContext {
job_id: row.get("id"),
state,
user_id: "default".to_string(), // Not stored in DB yet
conversation_id: row.get("conversation_id"),
title: row.get("title"),
description: row.get("description"),
category: row.get("category"),
state,
budget: row.get("budget_amount"),
budget_token: row.get("budget_token"),
bid_amount: row.get("bid_amount"),
+82 -1
View File
@@ -238,7 +238,9 @@ Respond in JSON format:
.with_temperature(0.7);
let response = self.llm.complete(request).await?;
Ok(response.content)
// Strip any internal thinking tags before returning to user
Ok(strip_thinking_tags(&response.content))
}
fn build_planning_prompt(&self, context: &ReasoningContext) -> String {
@@ -339,6 +341,43 @@ fn extract_json(text: &str) -> Option<&str> {
}
}
/// Strip `<thinking>...</thinking>` blocks from LLM output.
///
/// Some models (especially Claude with extended thinking) include internal
/// reasoning in thinking tags. We strip these before showing to users.
fn strip_thinking_tags(text: &str) -> String {
let mut result = String::with_capacity(text.len());
let mut remaining = text;
while let Some(start) = remaining.find("<thinking>") {
// Add everything before the tag
result.push_str(&remaining[..start]);
// Find the closing tag
if let Some(end_offset) = remaining[start..].find("</thinking>") {
// Skip past the closing tag (start + offset + tag length)
let end = start + end_offset + "</thinking>".len();
remaining = &remaining[end..];
} else {
// No closing tag found, discard everything from here
// (malformed, but handle gracefully by not including the unclosed tag)
remaining = "";
break;
}
}
// Add any remaining content after the last thinking block
result.push_str(remaining);
// Clean up any double newlines left behind
let mut cleaned = result.trim().to_string();
while cleaned.contains("\n\n\n") {
cleaned = cleaned.replace("\n\n\n", "\n\n");
}
cleaned
}
#[cfg(test)]
mod tests {
use super::*;
@@ -363,4 +402,46 @@ That's my plan."#;
assert_eq!(context.messages.len(), 1);
assert!(context.job_description.is_some());
}
#[test]
fn test_strip_thinking_tags_basic() {
let input = "<thinking>Let me think about this...</thinking>Hello, user!";
let output = strip_thinking_tags(input);
assert_eq!(output, "Hello, user!");
}
#[test]
fn test_strip_thinking_tags_multiple() {
let input =
"<thinking>First thought</thinking>Hello<thinking>Second thought</thinking> world!";
let output = strip_thinking_tags(input);
assert_eq!(output, "Hello world!");
}
#[test]
fn test_strip_thinking_tags_multiline() {
let input = r#"<thinking>
I need to consider:
1. What the user wants
2. How to respond
</thinking>
Here is my response to your question."#;
let output = strip_thinking_tags(input);
assert_eq!(output, "Here is my response to your question.");
}
#[test]
fn test_strip_thinking_tags_no_tags() {
let input = "Just a normal response without thinking tags.";
let output = strip_thinking_tags(input);
assert_eq!(output, "Just a normal response without thinking tags.");
}
#[test]
fn test_strip_thinking_tags_unclosed() {
// Malformed: unclosed tag should strip from there to end
let input = "Hello <thinking>this never closes";
let output = strip_thinking_tags(input);
assert_eq!(output, "Hello");
}
}
+19
View File
@@ -17,6 +17,7 @@ use near_agent::{
ToolRegistry,
wasm::{WasmToolLoader, WasmToolRuntime},
},
workspace::Workspace,
};
#[tokio::main]
@@ -91,6 +92,24 @@ async fn main() -> anyhow::Result<()> {
tools.register_builtin_tools();
tracing::info!("Registered {} built-in tools", tools.count());
// Register memory tools if database is available
if let Some(ref store) = store {
let workspace = Arc::new(Workspace::new("default", store.pool()));
tools.register_memory_tools(workspace);
}
// Register builder tool if enabled
if config.builder.enabled {
tools
.register_builder_tool(
llm.clone(),
safety.clone(),
Some(config.builder.to_builder_config()),
)
.await;
tracing::info!("Builder mode enabled");
}
// Load installed WASM tools
if config.wasm.enabled && config.wasm.tools_dir.exists() {
match WasmToolRuntime::new(config.wasm.to_runtime_config()) {
-201
View File
@@ -1,201 +0,0 @@
//! Dynamic tool builder for creating tools at runtime.
use std::time::Duration;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use crate::context::JobContext;
use crate::error::ToolError as AgentToolError;
use crate::tools::tool::{Tool, ToolError, ToolOutput};
/// Requirement specification for a new tool.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolRequirement {
/// Name for the new tool.
pub name: String,
/// Description of what the tool should do.
pub description: String,
/// Expected input parameters.
pub input_description: String,
/// Expected output format.
pub output_description: String,
/// Any external services or APIs needed.
pub dependencies: Vec<String>,
/// Security requirements.
pub security_requirements: Vec<String>,
}
/// Configuration for the tool sandbox.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SandboxConfig {
/// Maximum execution time.
pub max_execution_time: Duration,
/// Maximum memory in bytes.
pub max_memory_bytes: u64,
/// Allowed network hosts (empty = no network).
pub allowed_hosts: Vec<String>,
/// Allowed filesystem paths (empty = no filesystem).
pub allowed_paths: Vec<String>,
/// Environment variables to pass.
pub env_vars: Vec<(String, String)>,
}
impl Default for SandboxConfig {
fn default() -> Self {
Self {
max_execution_time: Duration::from_secs(30),
max_memory_bytes: 128 * 1024 * 1024, // 128 MB
allowed_hosts: vec![],
allowed_paths: vec![],
env_vars: vec![],
}
}
}
/// A dynamically created tool.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DynamicTool {
/// Tool name.
pub name: String,
/// Tool description.
pub description: String,
/// Generated code for the tool.
pub code: String,
/// Language of the generated code.
pub language: String,
/// Parameter schema.
pub parameters_schema: serde_json::Value,
/// Sandbox configuration.
pub sandbox_config: SandboxConfig,
/// When the tool was created.
pub created_at: DateTime<Utc>,
/// Job that created this tool (if any).
pub created_by_job_id: Option<uuid::Uuid>,
}
/// Trait for building tools dynamically.
#[async_trait]
pub trait ToolBuilder: Send + Sync {
/// Analyze a requirement and determine if a tool can be built.
async fn analyze_requirement(
&self,
description: &str,
) -> Result<ToolRequirement, AgentToolError>;
/// Build a tool from a requirement.
async fn build_tool(
&self,
requirement: &ToolRequirement,
) -> Result<DynamicTool, AgentToolError>;
/// Attempt to repair a broken tool.
async fn repair_tool(
&self,
tool: &DynamicTool,
error: &ToolError,
) -> Result<DynamicTool, AgentToolError>;
}
/// Default tool builder that uses LLM to generate tools.
pub struct LlmToolBuilder {
// TODO: Add LLM provider reference
}
impl LlmToolBuilder {
/// Create a new LLM-based tool builder.
pub fn new() -> Self {
Self {}
}
}
impl Default for LlmToolBuilder {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl ToolBuilder for LlmToolBuilder {
async fn analyze_requirement(
&self,
description: &str,
) -> Result<ToolRequirement, AgentToolError> {
// TODO: Use LLM to analyze the description and extract requirements
// For now, return a basic requirement
Ok(ToolRequirement {
name: "custom_tool".to_string(),
description: description.to_string(),
input_description: "JSON object with parameters".to_string(),
output_description: "JSON result".to_string(),
dependencies: vec![],
security_requirements: vec![],
})
}
async fn build_tool(
&self,
_requirement: &ToolRequirement,
) -> Result<DynamicTool, AgentToolError> {
// TODO: Use LLM to generate tool code
// For now, return a placeholder
Err(AgentToolError::BuilderFailed(
"Tool building not yet implemented".to_string(),
))
}
async fn repair_tool(
&self,
_tool: &DynamicTool,
error: &ToolError,
) -> Result<DynamicTool, AgentToolError> {
// TODO: Use LLM to analyze error and fix the tool
Err(AgentToolError::BuilderFailed(format!(
"Tool repair not yet implemented: {}",
error
)))
}
}
/// Wrapper to execute dynamic tools.
pub struct DynamicToolExecutor {
tool: DynamicTool,
}
impl DynamicToolExecutor {
/// Create an executor for a dynamic tool.
pub fn new(tool: DynamicTool) -> Self {
Self { tool }
}
}
#[async_trait]
impl Tool for DynamicToolExecutor {
fn name(&self) -> &str {
&self.tool.name
}
fn description(&self) -> &str {
&self.tool.description
}
fn parameters_schema(&self) -> serde_json::Value {
self.tool.parameters_schema.clone()
}
async fn execute(
&self,
_params: serde_json::Value,
_ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
// TODO: Execute the tool code in a sandbox
Err(ToolError::ExecutionFailed(
"Dynamic tool execution not yet implemented".to_string(),
))
}
fn requires_sanitization(&self) -> bool {
true // Dynamic tools always need sanitization
}
}
+898
View File
@@ -0,0 +1,898 @@
//! Software builder for creating programs and tools using LLM-driven code generation.
//!
//! This module provides a general-purpose software building capability that:
//! - Uses an agent loop similar to Codex for iterative development
//! - Can build any software (binaries, libraries, scripts)
//! - Has special context injection when building WASM tools
//! - Integrates with existing tool loading infrastructure
//!
//! # Architecture
//!
//! ```text
//! ┌─────────────────────────────────────────────────────────────────────────────┐
//! │ Software Build Loop │
//! │ │
//! │ 1. Analyze requirement ─▶ Determine project type, language, structure │
//! │ 2. Generate scaffold ─▶ Create initial project files │
//! │ 3. Implement code ─▶ Write the actual implementation │
//! │ 4. Build/compile ─▶ Run build commands (cargo, npm, etc.) │
//! │ 5. Fix errors ─▶ Parse errors, modify code, retry │
//! │ 6. Test ─▶ Run tests, fix failures │
//! │ 7. Package ─▶ Produce final artifact │
//! └─────────────────────────────────────────────────────────────────────────────┘
//! ```
//!
//! For WASM tools specifically:
//! - Injects Tool trait interface documentation
//! - Injects WASM host function documentation
//! - Compiles to wasm32-wasip2 target
//! - Validates against tool interface
//! - Registers with ToolRegistry
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use uuid::Uuid;
use crate::context::JobContext;
use crate::error::ToolError as AgentToolError;
use crate::llm::{ChatMessage, LlmProvider, Reasoning, ReasoningContext, ToolDefinition};
use crate::safety::SafetyLayer;
use crate::tools::ToolRegistry;
use crate::tools::tool::{Tool, ToolError, ToolOutput};
/// Requirement specification for building software.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BuildRequirement {
/// Name for the software.
pub name: String,
/// Description of what it should do.
pub description: String,
/// Type of software to build.
pub software_type: SoftwareType,
/// Target language/runtime.
pub language: Language,
/// Expected input format (for tools/CLIs).
pub input_spec: Option<String>,
/// Expected output format.
pub output_spec: Option<String>,
/// External dependencies needed.
pub dependencies: Vec<String>,
/// Security/capability requirements (for WASM tools).
pub capabilities: Vec<String>,
}
/// Type of software being built.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum SoftwareType {
/// A WASM tool for the agent.
WasmTool,
/// A standalone CLI application.
CliBinary,
/// A library/crate.
Library,
/// A script (Python, Bash, etc.).
Script,
/// A web service/API.
WebService,
}
/// Programming language for the build.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum Language {
Rust,
Python,
TypeScript,
JavaScript,
Go,
Bash,
}
impl Language {
/// Get the file extension for this language.
pub fn extension(&self) -> &'static str {
match self {
Language::Rust => "rs",
Language::Python => "py",
Language::TypeScript => "ts",
Language::JavaScript => "js",
Language::Go => "go",
Language::Bash => "sh",
}
}
/// Get the build command for this language.
pub fn build_command(&self, project_dir: &str) -> Option<String> {
match self {
Language::Rust => Some(format!("cd {} && cargo build --release", project_dir)),
Language::TypeScript => Some(format!("cd {} && npm run build", project_dir)),
Language::Go => Some(format!("cd {} && go build ./...", project_dir)),
Language::Python | Language::JavaScript | Language::Bash => None, // Interpreted
}
}
/// Get the test command for this language.
pub fn test_command(&self, project_dir: &str) -> String {
match self {
Language::Rust => format!("cd {} && cargo test", project_dir),
Language::Python => format!("cd {} && python -m pytest", project_dir),
Language::TypeScript | Language::JavaScript => {
format!("cd {} && npm test", project_dir)
}
Language::Go => format!("cd {} && go test ./...", project_dir),
Language::Bash => format!("cd {} && shellcheck *.sh", project_dir),
}
}
}
/// Result of a build operation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BuildResult {
/// Unique ID for this build.
pub build_id: Uuid,
/// The requirement that was built.
pub requirement: BuildRequirement,
/// Path to the output artifact.
pub artifact_path: PathBuf,
/// Build logs.
pub logs: Vec<BuildLog>,
/// Whether the build succeeded.
pub success: bool,
/// Error message if failed.
pub error: Option<String>,
/// When the build started.
pub started_at: DateTime<Utc>,
/// When the build completed.
pub completed_at: DateTime<Utc>,
/// Number of iterations to complete.
pub iterations: u32,
/// Validation warnings (for WASM tools).
#[serde(default)]
pub validation_warnings: Vec<String>,
/// Test results summary.
#[serde(default)]
pub tests_passed: u32,
/// Number of tests that failed.
#[serde(default)]
pub tests_failed: u32,
/// Whether the tool was auto-registered (for WASM tools).
#[serde(default)]
pub registered: bool,
}
/// A log entry from the build process.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BuildLog {
pub timestamp: DateTime<Utc>,
pub phase: BuildPhase,
pub message: String,
pub details: Option<String>,
}
/// Phases of the build process.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum BuildPhase {
Analyzing,
Scaffolding,
Implementing,
Building,
Testing,
Fixing,
Validating,
Registering,
Packaging,
Complete,
Failed,
}
/// Configuration for the software builder.
#[derive(Debug, Clone)]
pub struct BuilderConfig {
/// Directory where builds happen.
pub build_dir: PathBuf,
/// Maximum iterations before giving up.
pub max_iterations: u32,
/// Timeout for the entire build.
pub timeout: Duration,
/// Whether to clean up failed builds.
pub cleanup_on_failure: bool,
/// Whether to validate WASM tools after building.
pub validate_wasm: bool,
/// Whether to run tests after building.
pub run_tests: bool,
/// Whether to auto-register successful WASM tool builds.
pub auto_register: bool,
/// Directory to copy successful WASM tools for persistence.
pub wasm_output_dir: Option<PathBuf>,
}
impl Default for BuilderConfig {
fn default() -> Self {
Self {
build_dir: std::env::temp_dir().join("near-agent-builds"),
max_iterations: 10,
timeout: Duration::from_secs(600), // 10 minutes
cleanup_on_failure: false, // Keep for debugging
validate_wasm: true,
run_tests: true,
auto_register: true,
wasm_output_dir: None,
}
}
}
/// Trait for building software.
#[async_trait]
pub trait SoftwareBuilder: Send + Sync {
/// Analyze a natural language description and extract a structured requirement.
async fn analyze(&self, description: &str) -> Result<BuildRequirement, AgentToolError>;
/// Build software from a requirement.
async fn build(&self, requirement: &BuildRequirement) -> Result<BuildResult, AgentToolError>;
/// Attempt to repair a failed build.
async fn repair(
&self,
result: &BuildResult,
error: &str,
) -> Result<BuildResult, AgentToolError>;
}
/// LLM-powered software builder.
pub struct LlmSoftwareBuilder {
config: BuilderConfig,
llm: Arc<dyn LlmProvider>,
safety: Arc<SafetyLayer>,
tools: Arc<ToolRegistry>,
}
impl LlmSoftwareBuilder {
/// Create a new LLM-based software builder.
pub fn new(
config: BuilderConfig,
llm: Arc<dyn LlmProvider>,
safety: Arc<SafetyLayer>,
tools: Arc<ToolRegistry>,
) -> Self {
// Ensure build directory exists
if let Err(e) = std::fs::create_dir_all(&config.build_dir) {
tracing::warn!("Failed to create build directory: {}", e);
}
Self {
config,
llm,
safety,
tools,
}
}
/// Get the build tools available for the build loop.
async fn get_build_tools(&self) -> Vec<ToolDefinition> {
// Only include tools useful for building software
self.tools
.tool_definitions_for(&[
"shell",
"read_file",
"write_file",
"list_dir",
"apply_patch",
"http", // For fetching docs/deps
])
.await
}
/// Create the system prompt for the build agent.
fn build_system_prompt(&self, requirement: &BuildRequirement) -> String {
let mut prompt = format!(
r#"You are a software developer building a program.
## Task
Build: {name}
Description: {description}
Type: {software_type:?}
Language: {language:?}
## Process
1. Create the project structure with necessary files
2. Implement the code based on the requirements
3. Build/compile if needed
4. Run tests to verify correctness
5. Fix any errors and iterate
## Guidelines
- Write clean, well-structured code
- Handle errors appropriately
- Add minimal but useful comments
- Follow idiomatic patterns for the language
- Test edge cases
## Tools Available
- shell: Run build commands, tests, install dependencies
- read_file: Read existing files
- write_file: Create new files
- apply_patch: Edit existing files surgically
- list_dir: Explore project structure
"#,
name = requirement.name,
description = requirement.description,
software_type = requirement.software_type,
language = requirement.language,
);
// Add tool-specific context when building WASM tools
if requirement.software_type == SoftwareType::WasmTool {
prompt.push_str(&self.wasm_tool_context());
}
prompt
}
/// Get additional context for building WASM tools.
fn wasm_tool_context(&self) -> String {
r#"
## WASM Tool Requirements
You are building a WASM tool for an autonomous agent. The tool must:
1. **Implement the guest interface** - Export a `run` function that takes JSON input and returns JSON output
2. **Use only available host functions**:
- `host_log(level, message)` - Log messages (levels: debug, info, warn, error)
- `host_time()` - Get current Unix timestamp
- `host_http_request(method, url, headers, body)` - Make HTTP requests (if capability granted)
- `host_workspace_read(path)` - Read from workspace (if capability granted)
- `host_workspace_write(path, content)` - Write to workspace (if capability granted)
- `host_get_secret(name)` - Get injected secret (if capability granted)
3. **Handle errors gracefully** - Return error results, never panic
4. **Be deterministic** - Same input should produce same output (except for time/HTTP)
## WASM Tool Template (Rust)
```rust
// Cargo.toml
[package]
name = "tool_name"
version = "0.1.0"
edition = "2024"
[lib]
crate-type = ["cdylib"]
[dependencies]
serde = { version = "1", features = ["derive"] }
serde_json = "1"
// src/lib.rs
use serde::{Deserialize, Serialize};
#[derive(Deserialize)]
struct Input {
// Define your input parameters
}
#[derive(Serialize)]
struct Output {
// Define your output structure
}
// Host function imports
extern "C" {
fn host_log(level: i32, ptr: *const u8, len: usize);
}
fn log_info(msg: &str) {
unsafe { host_log(1, msg.as_ptr(), msg.len()); }
}
#[no_mangle]
pub extern "C" fn run(input_ptr: *const u8, input_len: usize) -> *mut u8 {
// Parse input
let input_bytes = unsafe { std::slice::from_raw_parts(input_ptr, input_len) };
let input: Input = match serde_json::from_slice(input_bytes) {
Ok(i) => i,
Err(e) => return error_response(&format!("Invalid input: {}", e)),
};
// Your implementation here
let output = Output { /* ... */ };
// Return output
let json = serde_json::to_vec(&output).unwrap();
let ptr = json.as_ptr() as *mut u8;
std::mem::forget(json);
ptr
}
fn error_response(msg: &str) -> *mut u8 {
let json = serde_json::json!({"error": msg}).to_string();
let ptr = json.as_ptr() as *mut u8;
std::mem::forget(json);
ptr
}
```
## Build Commands for WASM
```bash
# Add WASM target
rustup target add wasm32-wasip2
# Build
cargo build --target wasm32-wasip2 --release
# Output will be at: target/wasm32-wasip2/release/tool_name.wasm
```
## Tool Capabilities
When defining capabilities for your tool, specify which host functions it needs:
- `http`: Allows HTTP requests to specified endpoints
- `workspace`: Allows reading/writing workspace files
- `secrets`: Allows accessing injected secrets
"#
.to_string()
}
/// Execute the build loop.
async fn execute_build_loop(
&self,
requirement: &BuildRequirement,
project_dir: &Path,
) -> Result<BuildResult, AgentToolError> {
let build_id = Uuid::new_v4();
let started_at = Utc::now();
let mut logs = Vec::new();
let mut iteration = 0;
// Create reasoning engine
let reasoning = Reasoning::new(self.llm.clone(), self.safety.clone());
// Build initial context
let tool_defs = self.get_build_tools().await;
let mut reason_ctx = ReasoningContext::new().with_tools(tool_defs);
// Add system prompt
reason_ctx
.messages
.push(ChatMessage::system(self.build_system_prompt(requirement)));
// Add initial user message
reason_ctx.messages.push(ChatMessage::user(format!(
"Build the {} in directory: {}\n\nRequirements:\n- {}\n\nStart by creating the project structure.",
requirement.name,
project_dir.display(),
requirement.description
)));
logs.push(BuildLog {
timestamp: Utc::now(),
phase: BuildPhase::Analyzing,
message: "Starting build process".into(),
details: None,
});
// Main build loop
let mut current_phase = BuildPhase::Scaffolding;
let mut last_error: Option<String> = None;
loop {
iteration += 1;
if iteration > self.config.max_iterations {
logs.push(BuildLog {
timestamp: Utc::now(),
phase: BuildPhase::Failed,
message: "Maximum iterations exceeded".into(),
details: last_error.clone(),
});
return Ok(BuildResult {
build_id,
requirement: requirement.clone(),
artifact_path: project_dir.to_path_buf(),
logs,
success: false,
error: Some("Maximum iterations exceeded".into()),
started_at,
completed_at: Utc::now(),
iterations: iteration,
validation_warnings: Vec::new(),
tests_passed: 0,
tests_failed: 0,
registered: false,
});
}
// Get next action from LLM
let selections = reasoning.select_tools(&reason_ctx).await.map_err(|e| {
AgentToolError::BuilderFailed(format!("LLM tool selection failed: {}", e))
})?;
if selections.is_empty() {
// No tools selected - get response and check if done
let response = reasoning.respond(&reason_ctx).await.map_err(|e| {
AgentToolError::BuilderFailed(format!("LLM response failed: {}", e))
})?;
reason_ctx.messages.push(ChatMessage::assistant(&response));
// Check for completion signals
let response_lower = response.to_lowercase();
if response_lower.contains("build complete")
|| response_lower.contains("successfully built")
|| response_lower.contains("all tests pass")
{
logs.push(BuildLog {
timestamp: Utc::now(),
phase: BuildPhase::Complete,
message: "Build completed successfully".into(),
details: Some(response),
});
// Determine artifact path
let artifact_path = self.find_artifact(requirement, project_dir).await;
return Ok(BuildResult {
build_id,
requirement: requirement.clone(),
artifact_path,
logs,
success: true,
error: None,
started_at,
completed_at: Utc::now(),
iterations: iteration,
validation_warnings: Vec::new(),
tests_passed: 0,
tests_failed: 0,
registered: false,
});
}
// Ask for next steps
reason_ctx
.messages
.push(ChatMessage::user("Continue with the next step."));
continue;
}
// Execute selected tools
for selection in &selections {
logs.push(BuildLog {
timestamp: Utc::now(),
phase: current_phase,
message: format!("Executing: {}", selection.tool_name),
details: Some(selection.reasoning.clone()),
});
// Execute tool
let tool_result = self
.execute_build_tool(&selection.tool_name, &selection.parameters, project_dir)
.await;
match tool_result {
Ok(output) => {
let output_str =
serde_json::to_string_pretty(&output.result).unwrap_or_default();
// Add to context
reason_ctx.messages.push(ChatMessage::tool_result(
"tool_call",
&selection.tool_name,
output_str.clone(),
));
// Update phase based on tool
current_phase = match selection.tool_name.as_str() {
"write_file" => BuildPhase::Implementing,
"shell" if selection.parameters.to_string().contains("build") => {
BuildPhase::Building
}
"shell" if selection.parameters.to_string().contains("test") => {
BuildPhase::Testing
}
_ => current_phase,
};
// Check for build/test errors in output
if output_str.contains("error") || output_str.contains("failed") {
last_error = Some(output_str);
current_phase = BuildPhase::Fixing;
}
}
Err(e) => {
let error_msg = format!("Tool error: {}", e);
last_error = Some(error_msg.clone());
reason_ctx.messages.push(ChatMessage::tool_result(
"tool_call",
&selection.tool_name,
format!("Error: {}", e),
));
logs.push(BuildLog {
timestamp: Utc::now(),
phase: BuildPhase::Fixing,
message: "Tool execution failed".into(),
details: Some(error_msg),
});
current_phase = BuildPhase::Fixing;
}
}
}
}
}
/// Execute a build tool.
async fn execute_build_tool(
&self,
tool_name: &str,
params: &serde_json::Value,
_project_dir: &Path,
) -> Result<ToolOutput, ToolError> {
let tool =
self.tools.get(tool_name).await.ok_or_else(|| {
ToolError::ExecutionFailed(format!("Tool not found: {}", tool_name))
})?;
// Execute with a dummy context (build tools don't need job context)
let ctx = JobContext::default();
tool.execute(params.clone(), &ctx).await
}
/// Find the build artifact based on project type.
async fn find_artifact(&self, requirement: &BuildRequirement, project_dir: &Path) -> PathBuf {
match (&requirement.software_type, &requirement.language) {
(SoftwareType::WasmTool, Language::Rust) => {
// WASM output location
project_dir.join(format!(
"target/wasm32-wasip2/release/{}.wasm",
requirement.name.replace('-', "_")
))
}
(SoftwareType::CliBinary, Language::Rust) => project_dir.join(format!(
"target/release/{}",
requirement.name.replace('-', "_")
)),
(SoftwareType::Script, Language::Python) => {
project_dir.join(format!("{}.py", requirement.name))
}
(SoftwareType::Script, Language::Bash) => {
project_dir.join(format!("{}.sh", requirement.name))
}
_ => project_dir.to_path_buf(),
}
}
}
#[async_trait]
impl SoftwareBuilder for LlmSoftwareBuilder {
async fn analyze(&self, description: &str) -> Result<BuildRequirement, AgentToolError> {
// Use LLM to parse the description
let reasoning = Reasoning::new(self.llm.clone(), self.safety.clone());
let prompt = format!(
r#"Analyze this software requirement and extract structured information.
Description: {}
Respond with a JSON object containing:
- name: A short identifier (snake_case)
- description: What the software should do
- software_type: One of "wasm_tool", "cli_binary", "library", "script", "web_service"
- language: One of "rust", "python", "typescript", "javascript", "go", "bash"
- input_spec: Expected input format (optional)
- output_spec: Expected output format (optional)
- dependencies: List of external dependencies needed
- capabilities: For WASM tools, list needed capabilities (http, workspace, secrets)
JSON:"#,
description
);
let ctx = ReasoningContext::new().with_message(ChatMessage::user(&prompt));
let response = reasoning
.respond(&ctx)
.await
.map_err(|e| AgentToolError::BuilderFailed(format!("Analysis failed: {}", e)))?;
// Extract JSON from response
let json_start = response.find('{').unwrap_or(0);
let json_end = response.rfind('}').map(|i| i + 1).unwrap_or(response.len());
let json_str = &response[json_start..json_end];
serde_json::from_str(json_str).map_err(|e| {
AgentToolError::BuilderFailed(format!("Failed to parse requirement: {}", e))
})
}
async fn build(&self, requirement: &BuildRequirement) -> Result<BuildResult, AgentToolError> {
// Create project directory
let project_dir = self.config.build_dir.join(&requirement.name);
if project_dir.exists() {
std::fs::remove_dir_all(&project_dir).map_err(|e| {
AgentToolError::BuilderFailed(format!("Failed to clean project dir: {}", e))
})?;
}
std::fs::create_dir_all(&project_dir).map_err(|e| {
AgentToolError::BuilderFailed(format!("Failed to create project dir: {}", e))
})?;
// Run the build loop with timeout
let result = tokio::time::timeout(
self.config.timeout,
self.execute_build_loop(requirement, &project_dir),
)
.await;
match result {
Ok(Ok(build_result)) => Ok(build_result),
Ok(Err(e)) => Err(e),
Err(_) => Err(AgentToolError::BuilderFailed("Build timed out".into())),
}
}
async fn repair(
&self,
result: &BuildResult,
error: &str,
) -> Result<BuildResult, AgentToolError> {
// Create a new requirement with repair context
let mut requirement = result.requirement.clone();
requirement.description = format!(
"{}\n\nPrevious build failed with error:\n{}\n\nFix the issues and rebuild.",
requirement.description, error
);
// Rebuild (preserving project directory if it exists)
self.build(&requirement).await
}
}
/// Tool that allows the agent to build software on demand.
pub struct BuildSoftwareTool {
builder: Arc<dyn SoftwareBuilder>,
}
impl BuildSoftwareTool {
pub fn new(builder: Arc<dyn SoftwareBuilder>) -> Self {
Self { builder }
}
}
#[async_trait]
impl Tool for BuildSoftwareTool {
fn name(&self) -> &str {
"build_software"
}
fn description(&self) -> &str {
"Build software from a description. Can create WASM tools, CLI applications, scripts, \
and more. The builder will scaffold, implement, compile, and test the software iteratively."
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {
"description": {
"type": "string",
"description": "Natural language description of what to build"
},
"type": {
"type": "string",
"enum": ["wasm_tool", "cli_binary", "library", "script"],
"description": "Type of software to build (optional, will be inferred)"
},
"language": {
"type": "string",
"enum": ["rust", "python", "typescript", "bash"],
"description": "Programming language to use (optional, will be inferred)"
}
},
"required": ["description"]
})
}
async fn execute(
&self,
params: serde_json::Value,
_ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
let description = params
.get("description")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("missing 'description'".into()))?;
let start = std::time::Instant::now();
// Analyze the requirement
let mut requirement = self
.builder
.analyze(description)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("Analysis failed: {}", e)))?;
// Override type/language if specified
if let Some(type_str) = params.get("type").and_then(|v| v.as_str()) {
requirement.software_type = match type_str {
"wasm_tool" => SoftwareType::WasmTool,
"cli_binary" => SoftwareType::CliBinary,
"library" => SoftwareType::Library,
"script" => SoftwareType::Script,
_ => requirement.software_type,
};
}
if let Some(lang_str) = params.get("language").and_then(|v| v.as_str()) {
requirement.language = match lang_str {
"rust" => Language::Rust,
"python" => Language::Python,
"typescript" => Language::TypeScript,
"bash" => Language::Bash,
_ => requirement.language,
};
}
// Build
let result = self
.builder
.build(&requirement)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("Build failed: {}", e)))?;
let output = serde_json::json!({
"build_id": result.build_id.to_string(),
"name": result.requirement.name,
"success": result.success,
"artifact_path": result.artifact_path.display().to_string(),
"iterations": result.iterations,
"error": result.error,
"phases": result.logs.iter().map(|l| format!("{:?}: {}", l.phase, l.message)).collect::<Vec<_>>()
});
Ok(ToolOutput::success(output, start.elapsed()))
}
fn requires_approval(&self) -> bool {
true // Building software should require approval
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_language_extensions() {
assert_eq!(Language::Rust.extension(), "rs");
assert_eq!(Language::Python.extension(), "py");
assert_eq!(Language::TypeScript.extension(), "ts");
}
#[test]
fn test_build_commands() {
assert!(Language::Rust.build_command("/tmp/project").is_some());
assert!(Language::Python.build_command("/tmp/project").is_none());
}
#[test]
fn test_software_type_serialization() {
let json = serde_json::to_string(&SoftwareType::WasmTool).unwrap();
assert_eq!(json, "\"wasm_tool\"");
}
}
+37
View File
@@ -0,0 +1,37 @@
//! Software builder for creating programs and tools using LLM-driven code generation.
//!
//! This module provides a general-purpose software building capability that:
//! - Uses an agent loop similar to Codex for iterative development
//! - Can build any software (binaries, libraries, scripts)
//! - Has special context injection when building WASM tools
//! - Integrates with existing tool loading infrastructure
//!
//! # Architecture
//!
//! ```text
//! ┌─────────────────────────────────────────────────────────────────────────────┐
//! │ Software Build Loop │
//! │ │
//! │ 1. Analyze requirement ─▶ Determine project type, language, structure │
//! │ 2. Generate scaffold ─▶ Create initial project files │
//! │ 3. Implement code ─▶ Write the actual implementation │
//! │ 4. Build/compile ─▶ Run build commands (cargo, npm, etc.) │
//! │ 5. Fix errors ─▶ Parse errors, modify code, retry │
//! │ 6. Test ─▶ Run tests, fix failures │
//! │ 7. Validate ─▶ For WASM tools, verify interface compliance │
//! │ 8. Package ─▶ Produce final artifact │
//! └─────────────────────────────────────────────────────────────────────────────┘
//! ```
mod core;
mod templates;
mod testing;
mod validation;
pub use core::{
BuildLog, BuildPhase, BuildRequirement, BuildResult, BuildSoftwareTool, BuilderConfig,
Language, LlmSoftwareBuilder, SoftwareBuilder, SoftwareType,
};
pub use templates::{Template, TemplateEngine, TemplateType};
pub use testing::{TestCase, TestHarness, TestResult, TestSuite};
pub use validation::{ValidationError, ValidationResult, WasmValidator};
+501
View File
@@ -0,0 +1,501 @@
//! Code templates for common tool patterns.
//!
//! Templates provide scaffolding that the LLM fills in, reducing the chance
//! of structural errors and ensuring consistent patterns.
use std::collections::HashMap;
use serde::{Deserialize, Serialize};
/// Type of template.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum TemplateType {
/// WASM tool with HTTP capability.
WasmHttpTool,
/// WASM tool for data transformation.
WasmTransformTool,
/// WASM tool for computation.
WasmComputeTool,
/// CLI application.
CliBinary,
/// Python script.
PythonScript,
/// Bash script.
BashScript,
}
/// A code template with placeholders.
#[derive(Debug, Clone)]
pub struct Template {
pub template_type: TemplateType,
pub name: &'static str,
pub description: &'static str,
pub files: Vec<TemplateFile>,
}
/// A file within a template.
#[derive(Debug, Clone)]
pub struct TemplateFile {
pub path: &'static str,
pub content: &'static str,
pub is_required: bool,
}
/// Engine for rendering templates with variable substitution.
#[derive(Debug, Default)]
pub struct TemplateEngine {
variables: HashMap<String, String>,
}
impl TemplateEngine {
pub fn new() -> Self {
Self::default()
}
/// Set a template variable.
pub fn set(&mut self, key: impl Into<String>, value: impl Into<String>) -> &mut Self {
self.variables.insert(key.into(), value.into());
self
}
/// Render a template string, replacing {{variable}} placeholders.
pub fn render(&self, template: &str) -> String {
let mut result = template.to_string();
for (key, value) in &self.variables {
let placeholder = format!("{{{{{}}}}}", key);
result = result.replace(&placeholder, value);
}
result
}
/// Render all files in a template.
pub fn render_template(&self, template: &Template) -> Vec<(String, String)> {
template
.files
.iter()
.map(|f| (self.render(f.path), self.render(f.content)))
.collect()
}
}
impl Template {
/// Get template by type.
pub fn get(template_type: TemplateType) -> Self {
match template_type {
TemplateType::WasmHttpTool => Self::wasm_http_tool(),
TemplateType::WasmTransformTool => Self::wasm_transform_tool(),
TemplateType::WasmComputeTool => Self::wasm_compute_tool(),
TemplateType::CliBinary => Self::cli_binary(),
TemplateType::PythonScript => Self::python_script(),
TemplateType::BashScript => Self::bash_script(),
}
}
fn wasm_http_tool() -> Self {
Self {
template_type: TemplateType::WasmHttpTool,
name: "WASM HTTP Tool",
description: "A WASM tool that makes HTTP requests to external APIs",
files: vec![
TemplateFile {
path: "Cargo.toml",
content: WASM_CARGO_TOML,
is_required: true,
},
TemplateFile {
path: "src/lib.rs",
content: WASM_HTTP_LIB_RS,
is_required: true,
},
],
}
}
fn wasm_transform_tool() -> Self {
Self {
template_type: TemplateType::WasmTransformTool,
name: "WASM Transform Tool",
description: "A WASM tool that transforms data (JSON, text, etc.)",
files: vec![
TemplateFile {
path: "Cargo.toml",
content: WASM_CARGO_TOML,
is_required: true,
},
TemplateFile {
path: "src/lib.rs",
content: WASM_TRANSFORM_LIB_RS,
is_required: true,
},
],
}
}
fn wasm_compute_tool() -> Self {
Self {
template_type: TemplateType::WasmComputeTool,
name: "WASM Compute Tool",
description: "A WASM tool for pure computation (no I/O)",
files: vec![
TemplateFile {
path: "Cargo.toml",
content: WASM_CARGO_TOML,
is_required: true,
},
TemplateFile {
path: "src/lib.rs",
content: WASM_COMPUTE_LIB_RS,
is_required: true,
},
],
}
}
fn cli_binary() -> Self {
Self {
template_type: TemplateType::CliBinary,
name: "CLI Binary",
description: "A command-line application with argument parsing",
files: vec![
TemplateFile {
path: "Cargo.toml",
content: CLI_CARGO_TOML,
is_required: true,
},
TemplateFile {
path: "src/main.rs",
content: CLI_MAIN_RS,
is_required: true,
},
],
}
}
fn python_script() -> Self {
Self {
template_type: TemplateType::PythonScript,
name: "Python Script",
description: "A Python script with argument parsing",
files: vec![TemplateFile {
path: "{{name}}.py",
content: PYTHON_SCRIPT,
is_required: true,
}],
}
}
fn bash_script() -> Self {
Self {
template_type: TemplateType::BashScript,
name: "Bash Script",
description: "A Bash script with argument handling",
files: vec![TemplateFile {
path: "{{name}}.sh",
content: BASH_SCRIPT,
is_required: true,
}],
}
}
}
// =============================================================================
// WASM Templates
// =============================================================================
const WASM_CARGO_TOML: &str = r##"[package]
name = "{{name}}"
version = "0.1.0"
edition = "2021"
[lib]
crate-type = ["cdylib"]
[dependencies]
serde = { version = "1", features = ["derive"] }
serde_json = "1"
[profile.release]
opt-level = "s"
lto = true
"##;
const WASM_HTTP_LIB_RS: &str = r##"//! {{description}}
//!
//! This WASM tool makes HTTP requests to external APIs.
use serde::{Deserialize, Serialize};
// Host function imports
#[link(wasm_import_module = "env")]
extern "C" {
fn host_log(level: i32, ptr: *const u8, len: usize);
fn host_http_request(
method_ptr: *const u8, method_len: usize,
url_ptr: *const u8, url_len: usize,
headers_ptr: *const u8, headers_len: usize,
body_ptr: *const u8, body_len: usize,
response_ptr: *mut u8, response_max_len: usize,
) -> i32;
}
fn log_info(msg: &str) {
unsafe { host_log(1, msg.as_ptr(), msg.len()); }
}
fn http_get(url: &str) -> Result<String, String> {
let method = "GET";
let mut response_buf = vec![0u8; 65536];
let result = unsafe {
host_http_request(
method.as_ptr(), method.len(),
url.as_ptr(), url.len(),
std::ptr::null(), 0,
std::ptr::null(), 0,
response_buf.as_mut_ptr(), response_buf.len(),
)
};
if result < 0 { return Err(format!("HTTP error: {}", result)); }
response_buf.truncate(result as usize);
String::from_utf8(response_buf).map_err(|e| e.to_string())
}
#[derive(Deserialize)]
struct Input {
{{input_fields}}
}
#[derive(Serialize)]
struct Output {
{{output_fields}}
}
#[no_mangle]
pub extern "C" fn run(input_ptr: *const u8, input_len: usize) -> u64 {
let result = run_inner(input_ptr, input_len);
let json = match result {
Ok(output) => serde_json::to_string(&output).unwrap_or_else(|e| {
format!("{{\"error\":\"serialize: {}\"}}", e)
}),
Err(e) => format!("{{\"error\":\"{}\"}}", e.replace('"', "'")),
};
let bytes = json.into_bytes();
let ptr = bytes.as_ptr() as u64;
let len = bytes.len() as u64;
std::mem::forget(bytes);
(len << 32) | ptr
}
fn run_inner(input_ptr: *const u8, input_len: usize) -> Result<Output, String> {
let input_bytes = unsafe { std::slice::from_raw_parts(input_ptr, input_len) };
let input: Input = serde_json::from_slice(input_bytes)
.map_err(|e| format!("Invalid input: {}", e))?;
log_info("Processing request...");
{{implementation}}
Ok(Output {
{{output_construction}}
})
}
"##;
const WASM_TRANSFORM_LIB_RS: &str = r##"//! {{description}}
//!
//! This WASM tool transforms input data.
use serde::{Deserialize, Serialize};
#[link(wasm_import_module = "env")]
extern "C" {
fn host_log(level: i32, ptr: *const u8, len: usize);
}
fn log_info(msg: &str) {
unsafe { host_log(1, msg.as_ptr(), msg.len()); }
}
#[derive(Deserialize)]
struct Input {
{{input_fields}}
}
#[derive(Serialize)]
struct Output {
{{output_fields}}
}
#[no_mangle]
pub extern "C" fn run(input_ptr: *const u8, input_len: usize) -> u64 {
let result = run_inner(input_ptr, input_len);
let json = match result {
Ok(output) => serde_json::to_string(&output).unwrap_or_else(|e| {
format!("{{\"error\":\"serialize: {}\"}}", e)
}),
Err(e) => format!("{{\"error\":\"{}\"}}", e.replace('"', "'")),
};
let bytes = json.into_bytes();
let ptr = bytes.as_ptr() as u64;
let len = bytes.len() as u64;
std::mem::forget(bytes);
(len << 32) | ptr
}
fn run_inner(input_ptr: *const u8, input_len: usize) -> Result<Output, String> {
let input_bytes = unsafe { std::slice::from_raw_parts(input_ptr, input_len) };
let input: Input = serde_json::from_slice(input_bytes)
.map_err(|e| format!("Invalid input: {}", e))?;
log_info("Transforming data...");
{{implementation}}
Ok(Output {
{{output_construction}}
})
}
"##;
const WASM_COMPUTE_LIB_RS: &str = r##"//! {{description}}
//!
//! This WASM tool performs pure computation.
use serde::{Deserialize, Serialize};
#[derive(Deserialize)]
struct Input {
{{input_fields}}
}
#[derive(Serialize)]
struct Output {
{{output_fields}}
}
#[no_mangle]
pub extern "C" fn run(input_ptr: *const u8, input_len: usize) -> u64 {
let result = run_inner(input_ptr, input_len);
let json = match result {
Ok(output) => serde_json::to_string(&output).unwrap_or_else(|e| {
format!("{{\"error\":\"serialize: {}\"}}", e)
}),
Err(e) => format!("{{\"error\":\"{}\"}}", e.replace('"', "'")),
};
let bytes = json.into_bytes();
let ptr = bytes.as_ptr() as u64;
let len = bytes.len() as u64;
std::mem::forget(bytes);
(len << 32) | ptr
}
fn run_inner(input_ptr: *const u8, input_len: usize) -> Result<Output, String> {
let input_bytes = unsafe { std::slice::from_raw_parts(input_ptr, input_len) };
let input: Input = serde_json::from_slice(input_bytes)
.map_err(|e| format!("Invalid input: {}", e))?;
{{implementation}}
Ok(Output {
{{output_construction}}
})
}
"##;
// =============================================================================
// CLI Templates
// =============================================================================
const CLI_CARGO_TOML: &str = r##"[package]
name = "{{name}}"
version = "0.1.0"
edition = "2021"
[dependencies]
clap = { version = "4", features = ["derive"] }
anyhow = "1"
serde = { version = "1", features = ["derive"] }
serde_json = "1"
"##;
const CLI_MAIN_RS: &str = r##"//! {{description}}
use clap::Parser;
use anyhow::Result;
#[derive(Parser, Debug)]
#[command(name = "{{name}}")]
#[command(about = "{{description}}")]
struct Args {
{{cli_args}}
}
fn main() -> Result<()> {
let args = Args::parse();
{{implementation}}
Ok(())
}
"##;
// =============================================================================
// Script Templates
// =============================================================================
const PYTHON_SCRIPT: &str = r##"#!/usr/bin/env python3
"""{{description}}"""
import argparse
import json
import sys
def main():
parser = argparse.ArgumentParser(description="{{description}}")
{{python_args}}
args = parser.parse_args()
{{implementation}}
if __name__ == "__main__":
main()
"##;
const BASH_SCRIPT: &str = r##"#!/bin/bash
# {{description}}
set -euo pipefail
usage() {
echo "Usage: $0 {{bash_usage}}"
exit 1
}
{{bash_arg_parsing}}
{{implementation}}
"##;
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_template_engine() {
let mut engine = TemplateEngine::new();
engine.set("name", "my_tool");
engine.set("description", "A cool tool");
let result = engine.render("Name: {{name}}, Desc: {{description}}");
assert_eq!(result, "Name: my_tool, Desc: A cool tool");
}
#[test]
fn test_get_template() {
let template = Template::get(TemplateType::WasmHttpTool);
assert_eq!(template.name, "WASM HTTP Tool");
assert!(!template.files.is_empty());
}
}
+527
View File
@@ -0,0 +1,527 @@
//! Testing harness for built tools.
//!
//! Provides automated testing of generated tools before registration,
//! ensuring they work correctly with various inputs.
use std::path::Path;
use std::sync::Arc;
use std::time::{Duration, Instant};
use serde::{Deserialize, Serialize};
use thiserror::Error;
use crate::context::JobContext;
use crate::tools::tool::Tool;
use crate::tools::wasm::{Capabilities, WasmError, WasmToolRuntime, WasmToolWrapper};
/// Errors during testing.
#[derive(Debug, Error)]
pub enum TestError {
#[error("Failed to load WASM module: {0}")]
LoadError(#[from] WasmError),
#[error("Test execution failed: {0}")]
ExecutionFailed(String),
#[error("Test timed out after {0:?}")]
Timeout(Duration),
#[error("Output mismatch: expected {expected}, got {actual}")]
OutputMismatch { expected: String, actual: String },
#[error("Test assertion failed: {0}")]
AssertionFailed(String),
#[error("IO error: {0}")]
IoError(#[from] std::io::Error),
}
/// A single test case.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TestCase {
/// Name of the test.
pub name: String,
/// Description of what this test verifies.
pub description: Option<String>,
/// Input JSON to pass to the tool.
pub input: serde_json::Value,
/// Expected output (if exact match required).
pub expected_output: Option<serde_json::Value>,
/// Expected fields in output (partial match).
pub expected_fields: Option<Vec<ExpectedField>>,
/// Whether the tool should return an error.
pub expect_error: bool,
/// Expected error message substring (if expect_error is true).
pub error_contains: Option<String>,
/// Timeout for this specific test.
pub timeout_ms: Option<u64>,
}
/// An expected field in the output.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExpectedField {
/// JSON path to the field (e.g., "result.value" or "data[0].name").
pub path: String,
/// Expected value at that path.
pub value: Option<serde_json::Value>,
/// Just check that the field exists (if value is None).
pub exists: bool,
}
/// Result of running a single test.
#[derive(Debug, Clone)]
pub struct TestResult {
/// Name of the test.
pub name: String,
/// Whether the test passed.
pub passed: bool,
/// Duration of the test.
pub duration: Duration,
/// Error message if failed.
pub error: Option<String>,
/// Actual output from the tool.
pub actual_output: Option<serde_json::Value>,
}
/// A suite of tests for a tool.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TestSuite {
/// Name of the test suite.
pub name: String,
/// Description of the suite.
pub description: Option<String>,
/// Test cases in the suite.
pub tests: Vec<TestCase>,
/// Default timeout for tests in milliseconds.
pub default_timeout_ms: u64,
}
impl Default for TestSuite {
fn default() -> Self {
Self {
name: "default".to_string(),
description: None,
tests: Vec::new(),
default_timeout_ms: 5000,
}
}
}
impl TestSuite {
pub fn new(name: impl Into<String>) -> Self {
Self {
name: name.into(),
..Default::default()
}
}
/// Add a test case.
pub fn add_test(&mut self, test: TestCase) -> &mut Self {
self.tests.push(test);
self
}
/// Add a simple input/output test.
pub fn add_io_test(
&mut self,
name: impl Into<String>,
input: serde_json::Value,
expected: serde_json::Value,
) -> &mut Self {
self.tests.push(TestCase {
name: name.into(),
description: None,
input,
expected_output: Some(expected),
expected_fields: None,
expect_error: false,
error_contains: None,
timeout_ms: None,
});
self
}
/// Add a test that expects an error.
pub fn add_error_test(
&mut self,
name: impl Into<String>,
input: serde_json::Value,
error_contains: impl Into<String>,
) -> &mut Self {
self.tests.push(TestCase {
name: name.into(),
description: None,
input,
expected_output: None,
expected_fields: None,
expect_error: true,
error_contains: Some(error_contains.into()),
timeout_ms: None,
});
self
}
}
/// Harness for running tests against WASM tools.
pub struct TestHarness {
runtime: Arc<WasmToolRuntime>,
capabilities: Capabilities,
default_timeout: Duration,
}
impl TestHarness {
pub fn new(runtime: Arc<WasmToolRuntime>) -> Self {
Self {
runtime,
capabilities: Capabilities::none(),
default_timeout: Duration::from_secs(5),
}
}
/// Set capabilities for test execution.
pub fn with_capabilities(mut self, caps: Capabilities) -> Self {
self.capabilities = caps;
self
}
/// Set default timeout.
pub fn with_timeout(mut self, timeout: Duration) -> Self {
self.default_timeout = timeout;
self
}
/// Run a test suite against a WASM file.
pub async fn run_suite_file(
&self,
wasm_path: &Path,
suite: &TestSuite,
) -> Result<Vec<TestResult>, TestError> {
let bytes = tokio::fs::read(wasm_path).await?;
self.run_suite_bytes(&bytes, suite).await
}
/// Run a test suite against WASM bytes.
pub async fn run_suite_bytes(
&self,
wasm_bytes: &[u8],
suite: &TestSuite,
) -> Result<Vec<TestResult>, TestError> {
// Prepare the module
let prepared = self.runtime.prepare(&suite.name, wasm_bytes, None).await?;
// Create a tool wrapper for execution
let tool = WasmToolWrapper::new(
Arc::clone(&self.runtime),
prepared,
self.capabilities.clone(),
);
let mut results = Vec::with_capacity(suite.tests.len());
for test in &suite.tests {
let result = self.run_test(&tool, test, suite.default_timeout_ms).await;
results.push(result);
}
Ok(results)
}
/// Run a single test case.
async fn run_test(
&self,
tool: &WasmToolWrapper,
test: &TestCase,
default_timeout_ms: u64,
) -> TestResult {
let timeout = Duration::from_millis(test.timeout_ms.unwrap_or(default_timeout_ms));
let start = Instant::now();
let ctx = JobContext::default();
// Execute with timeout
let exec_result = tokio::time::timeout(timeout, async {
tool.execute(test.input.clone(), &ctx).await
})
.await;
let duration = start.elapsed();
match exec_result {
Err(_) => TestResult {
name: test.name.clone(),
passed: false,
duration,
error: Some(format!("Test timed out after {:?}", timeout)),
actual_output: None,
},
Ok(Err(e)) => {
// Execution error
if test.expect_error {
let error_str = e.to_string();
let matches = test
.error_contains
.as_ref()
.is_none_or(|expected| error_str.contains(expected));
TestResult {
name: test.name.clone(),
passed: matches,
duration,
error: if matches {
None
} else {
Some(format!(
"Expected error containing '{}', got: {}",
test.error_contains.as_deref().unwrap_or(""),
error_str
))
},
actual_output: None,
}
} else {
TestResult {
name: test.name.clone(),
passed: false,
duration,
error: Some(format!("Unexpected error: {}", e)),
actual_output: None,
}
}
}
Ok(Ok(output)) => {
let actual = output.result;
// Check if output contains an error field
if let Some(error_val) = actual.get("error") {
if test.expect_error {
let error_str = error_val.as_str().unwrap_or("");
let matches = test
.error_contains
.as_ref()
.is_none_or(|expected| error_str.contains(expected));
return TestResult {
name: test.name.clone(),
passed: matches,
duration,
error: if matches {
None
} else {
Some(format!(
"Expected error containing '{}', got: {}",
test.error_contains.as_deref().unwrap_or(""),
error_str
))
},
actual_output: Some(actual),
};
} else {
return TestResult {
name: test.name.clone(),
passed: false,
duration,
error: Some(format!("Unexpected error in output: {}", error_val)),
actual_output: Some(actual),
};
}
}
// Verify expected output
if let Some(ref expected) = test.expected_output {
if &actual != expected {
return TestResult {
name: test.name.clone(),
passed: false,
duration,
error: Some(format!(
"Output mismatch:\nExpected: {}\nActual: {}",
serde_json::to_string_pretty(expected).unwrap_or_default(),
serde_json::to_string_pretty(&actual).unwrap_or_default()
)),
actual_output: Some(actual),
};
}
}
// Verify expected fields
if let Some(ref fields) = test.expected_fields {
for field in fields {
let field_value = get_json_path(&actual, &field.path);
if field.exists && field_value.is_none() {
return TestResult {
name: test.name.clone(),
passed: false,
duration,
error: Some(format!("Missing expected field: {}", field.path)),
actual_output: Some(actual),
};
}
if let Some(ref expected_value) = field.value {
if field_value != Some(expected_value) {
return TestResult {
name: test.name.clone(),
passed: false,
duration,
error: Some(format!(
"Field '{}' mismatch: expected {:?}, got {:?}",
field.path, expected_value, field_value
)),
actual_output: Some(actual),
};
}
}
}
}
TestResult {
name: test.name.clone(),
passed: true,
duration,
error: None,
actual_output: Some(actual),
}
}
}
}
}
/// Get a value from a JSON object by path (e.g., "foo.bar[0].baz").
fn get_json_path<'a>(value: &'a serde_json::Value, path: &str) -> Option<&'a serde_json::Value> {
let mut current = value;
for segment in path.split('.') {
// Handle array indexing like "items[0]"
if let Some(bracket_pos) = segment.find('[') {
let key = &segment[..bracket_pos];
let index_str = &segment[bracket_pos + 1..segment.len() - 1];
if !key.is_empty() {
current = current.get(key)?;
}
let index: usize = index_str.parse().ok()?;
current = current.get(index)?;
} else {
current = current.get(segment)?;
}
}
Some(current)
}
/// Generate basic test cases for a tool based on its schema.
pub fn generate_basic_tests(name: &str, input_schema: &serde_json::Value) -> TestSuite {
let mut suite = TestSuite::new(format!("{}_basic_tests", name));
suite.description = Some("Auto-generated basic tests".to_string());
// Test with empty input
suite.add_error_test("empty_input", serde_json::json!({}), "");
// Test with null values for required fields
if let Some(required) = input_schema.get("required").and_then(|r| r.as_array()) {
let mut null_input = serde_json::Map::new();
for req in required {
if let Some(field_name) = req.as_str() {
null_input.insert(field_name.to_string(), serde_json::Value::Null);
}
}
suite.add_error_test(
"null_required_fields",
serde_json::Value::Object(null_input),
"",
);
}
// Test with valid minimal input (if we can construct it)
if let Some(properties) = input_schema.get("properties").and_then(|p| p.as_object()) {
let mut minimal_input = serde_json::Map::new();
for (name, prop) in properties {
if let Some(prop_type) = prop.get("type").and_then(|t| t.as_str()) {
let value = match prop_type {
"string" => serde_json::Value::String("test".to_string()),
"integer" | "number" => serde_json::Value::Number(0.into()),
"boolean" => serde_json::Value::Bool(false),
"array" => serde_json::Value::Array(vec![]),
"object" => serde_json::Value::Object(serde_json::Map::new()),
_ => continue,
};
minimal_input.insert(name.clone(), value);
}
}
suite.tests.push(TestCase {
name: "minimal_valid_input".to_string(),
description: Some("Test with minimal valid input".to_string()),
input: serde_json::Value::Object(minimal_input),
expected_output: None,
expected_fields: None,
expect_error: false,
error_contains: None,
timeout_ms: None,
});
}
suite
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_get_json_path() {
let json = serde_json::json!({
"foo": {
"bar": [1, 2, 3],
"baz": "hello"
}
});
assert_eq!(
get_json_path(&json, "foo.baz"),
Some(&serde_json::json!("hello"))
);
assert_eq!(
get_json_path(&json, "foo.bar[0]"),
Some(&serde_json::json!(1))
);
assert_eq!(
get_json_path(&json, "foo.bar[2]"),
Some(&serde_json::json!(3))
);
assert_eq!(get_json_path(&json, "foo.missing"), None);
}
#[test]
fn test_test_suite_builder() {
let mut suite = TestSuite::new("my_tests");
suite
.add_io_test(
"basic",
serde_json::json!({"x": 1}),
serde_json::json!({"y": 2}),
)
.add_error_test("invalid", serde_json::json!({}), "required");
assert_eq!(suite.tests.len(), 2);
assert!(!suite.tests[0].expect_error);
assert!(suite.tests[1].expect_error);
}
#[test]
fn test_generate_basic_tests() {
let schema = serde_json::json!({
"type": "object",
"properties": {
"name": {"type": "string"},
"count": {"type": "integer"}
},
"required": ["name"]
});
let suite = generate_basic_tests("my_tool", &schema);
assert!(!suite.tests.is_empty());
}
}
+319
View File
@@ -0,0 +1,319 @@
//! WASM tool validation.
//!
//! Validates that built WASM modules conform to the expected tool interface
//! before they can be registered with the agent.
use std::path::Path;
use thiserror::Error;
/// Errors during WASM validation.
#[derive(Debug, Error)]
pub enum ValidationError {
#[error("Failed to read WASM file: {0}")]
IoError(#[from] std::io::Error),
#[error("Invalid WASM module: {0}")]
InvalidModule(String),
#[error("Missing required export: {0}")]
MissingExport(String),
#[error("Invalid export signature for '{name}': expected {expected}, got {actual}")]
InvalidSignature {
name: String,
expected: String,
actual: String,
},
#[error("Module uses disallowed import: {module}::{name}")]
DisallowedImport { module: String, name: String },
#[error("Module exceeds size limit: {size} bytes (max: {max} bytes)")]
TooLarge { size: u64, max: u64 },
#[error("Validation failed: {0}")]
Other(String),
}
/// Result of WASM validation.
#[derive(Debug)]
pub struct ValidationResult {
/// Whether the module is valid.
pub is_valid: bool,
/// List of validation errors (empty if valid).
pub errors: Vec<ValidationError>,
/// List of warnings (non-fatal issues).
pub warnings: Vec<String>,
/// Detected exports.
pub exports: Vec<ExportInfo>,
/// Detected imports.
pub imports: Vec<ImportInfo>,
/// Module size in bytes.
pub size_bytes: u64,
}
/// Information about an exported function.
#[derive(Debug, Clone)]
pub struct ExportInfo {
pub name: String,
pub kind: ExportKind,
}
/// Kind of export.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ExportKind {
Function,
Memory,
Table,
Global,
}
/// Information about an imported function.
#[derive(Debug, Clone)]
pub struct ImportInfo {
pub module: String,
pub name: String,
pub kind: ImportKind,
}
/// Kind of import.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ImportKind {
Function,
Memory,
Table,
Global,
}
/// Validator for WASM tool modules.
pub struct WasmValidator {
/// Maximum module size in bytes.
max_size: u64,
/// Required exports that must be present.
required_exports: Vec<String>,
/// Allowed import modules.
allowed_import_modules: Vec<String>,
}
impl Default for WasmValidator {
fn default() -> Self {
Self {
max_size: 10 * 1024 * 1024, // 10 MB
required_exports: vec!["run".to_string()],
allowed_import_modules: vec![
"env".to_string(),
"wasi_snapshot_preview1".to_string(),
"wasi".to_string(),
],
}
}
}
impl WasmValidator {
pub fn new() -> Self {
Self::default()
}
/// Set maximum module size.
pub fn with_max_size(mut self, max_bytes: u64) -> Self {
self.max_size = max_bytes;
self
}
/// Add a required export.
pub fn with_required_export(mut self, name: impl Into<String>) -> Self {
self.required_exports.push(name.into());
self
}
/// Add an allowed import module.
pub fn with_allowed_import(mut self, module: impl Into<String>) -> Self {
self.allowed_import_modules.push(module.into());
self
}
/// Validate a WASM file.
pub async fn validate_file(&self, path: &Path) -> Result<ValidationResult, ValidationError> {
let bytes = tokio::fs::read(path).await?;
self.validate_bytes(&bytes)
}
/// Validate WASM bytes.
pub fn validate_bytes(&self, bytes: &[u8]) -> Result<ValidationResult, ValidationError> {
let mut errors = Vec::new();
let mut warnings = Vec::new();
let mut exports = Vec::new();
let mut imports = Vec::new();
let size_bytes = bytes.len() as u64;
// Check size
if size_bytes > self.max_size {
errors.push(ValidationError::TooLarge {
size: size_bytes,
max: self.max_size,
});
}
// Parse WASM module
let parser = wasmparser::Parser::new(0);
for payload in parser.parse_all(bytes) {
match payload {
Ok(wasmparser::Payload::ExportSection(reader)) => {
for export in reader {
match export {
Ok(exp) => {
let kind = match exp.kind {
wasmparser::ExternalKind::Func => ExportKind::Function,
wasmparser::ExternalKind::Memory => ExportKind::Memory,
wasmparser::ExternalKind::Table => ExportKind::Table,
wasmparser::ExternalKind::Global => ExportKind::Global,
wasmparser::ExternalKind::Tag => continue,
};
exports.push(ExportInfo {
name: exp.name.to_string(),
kind,
});
}
Err(e) => {
errors.push(ValidationError::InvalidModule(format!(
"Failed to parse export: {}",
e
)));
}
}
}
}
Ok(wasmparser::Payload::ImportSection(reader)) => {
for import in reader {
match import {
Ok(imp) => {
let kind = match imp.ty {
wasmparser::TypeRef::Func(_) => ImportKind::Function,
wasmparser::TypeRef::Memory(_) => ImportKind::Memory,
wasmparser::TypeRef::Table(_) => ImportKind::Table,
wasmparser::TypeRef::Global(_) => ImportKind::Global,
wasmparser::TypeRef::Tag(_) => continue,
};
imports.push(ImportInfo {
module: imp.module.to_string(),
name: imp.name.to_string(),
kind,
});
// Check if import module is allowed
if !self
.allowed_import_modules
.contains(&imp.module.to_string())
{
errors.push(ValidationError::DisallowedImport {
module: imp.module.to_string(),
name: imp.name.to_string(),
});
}
}
Err(e) => {
errors.push(ValidationError::InvalidModule(format!(
"Failed to parse import: {}",
e
)));
}
}
}
}
Ok(_) => {
// Other sections are OK
}
Err(e) => {
errors.push(ValidationError::InvalidModule(format!(
"Failed to parse WASM: {}",
e
)));
break;
}
}
}
// Check required exports
for required in &self.required_exports {
if !exports.iter().any(|e| &e.name == required) {
errors.push(ValidationError::MissingExport(required.clone()));
}
}
// Check for common issues (warnings)
if !exports
.iter()
.any(|e| e.name == "memory" && e.kind == ExportKind::Memory)
{
warnings
.push("Module does not export memory - host cannot read/write data".to_string());
}
// Check for potentially dangerous imports
for import in &imports {
if import.module == "wasi_snapshot_preview1" {
match import.name.as_str() {
"fd_write" | "fd_read" | "path_open" | "path_create_directory" => {
warnings.push(format!(
"Module uses WASI filesystem function '{}' - ensure this is intended",
import.name
));
}
"sock_send" | "sock_recv" | "sock_accept" => {
warnings.push(format!(
"Module uses WASI socket function '{}' - ensure this is intended",
import.name
));
}
_ => {}
}
}
}
Ok(ValidationResult {
is_valid: errors.is_empty(),
errors,
warnings,
exports,
imports,
size_bytes,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_validator_default() {
let validator = WasmValidator::new();
assert_eq!(validator.max_size, 10 * 1024 * 1024);
assert!(validator.required_exports.contains(&"run".to_string()));
}
#[test]
fn test_validator_builder() {
let validator = WasmValidator::new()
.with_max_size(1024)
.with_required_export("custom_export")
.with_allowed_import("custom_module");
assert_eq!(validator.max_size, 1024);
assert!(
validator
.required_exports
.contains(&"custom_export".to_string())
);
assert!(
validator
.allowed_import_modules
.contains(&"custom_module".to_string())
);
}
// Note: Full WASM parsing tests would require actual WASM binaries
}
+739
View File
@@ -0,0 +1,739 @@
//! File operation tools for reading, writing, and navigating the filesystem.
//!
//! These tools provide controlled access to the filesystem with:
//! - Path validation and sandboxing
//! - Size limits on read/write operations
//! - Support for common development tasks
use std::path::{Path, PathBuf};
use async_trait::async_trait;
use tokio::fs;
use crate::context::JobContext;
use crate::tools::tool::{Tool, ToolError, ToolOutput};
/// Maximum file size for reading (1MB).
const MAX_READ_SIZE: u64 = 1024 * 1024;
/// Maximum file size for writing (5MB).
const MAX_WRITE_SIZE: usize = 5 * 1024 * 1024;
/// Maximum directory listing entries.
const MAX_DIR_ENTRIES: usize = 500;
/// Validate that a path is safe (no traversal attacks).
fn validate_path(path_str: &str, base_dir: Option<&Path>) -> Result<PathBuf, ToolError> {
let path = PathBuf::from(path_str);
// Reject paths with suspicious components (validation only, no action needed)
for component in path.components() {
match component {
std::path::Component::ParentDir => {
// Allow .. but validate final path is within sandbox
}
std::path::Component::Normal(s) => {
let s = s.to_string_lossy();
if s.starts_with('.') && s != "." && s != ".." && !s.starts_with(".git") {
// Hidden files are OK for .git, .gitignore, etc.
}
}
_ => {}
}
}
// Resolve to absolute path
let resolved = if path.is_absolute() {
path.canonicalize().unwrap_or_else(|_| path.clone())
} else if let Some(base) = base_dir {
base.join(&path)
.canonicalize()
.unwrap_or_else(|_| base.join(&path))
} else {
std::env::current_dir()
.unwrap_or_else(|_| PathBuf::from("."))
.join(&path)
};
// If base_dir is set, ensure path is within it
if let Some(base) = base_dir {
// Canonicalize the base to handle symlinks (e.g., /var -> /private/var on macOS)
let base_canonical = base.canonicalize().unwrap_or_else(|_| base.to_path_buf());
// For files that don't exist yet, we need to check the parent directory
// and ensure the resolved path would be within the base
let check_path = if resolved.exists() {
resolved.canonicalize().unwrap_or_else(|_| resolved.clone())
} else {
// For non-existent files, canonicalize the parent and append the filename
if let Some(parent) = resolved.parent() {
if parent.exists() {
let canonical_parent = parent
.canonicalize()
.unwrap_or_else(|_| parent.to_path_buf());
if let Some(filename) = resolved.file_name() {
canonical_parent.join(filename)
} else {
resolved.clone()
}
} else {
resolved.clone()
}
} else {
resolved.clone()
}
};
if !check_path.starts_with(&base_canonical) {
return Err(ToolError::NotAuthorized(format!(
"Path escapes sandbox: {}",
path_str
)));
}
}
Ok(resolved)
}
/// Read file contents tool.
#[derive(Debug, Default)]
pub struct ReadFileTool {
base_dir: Option<PathBuf>,
}
impl ReadFileTool {
pub fn new() -> Self {
Self::default()
}
pub fn with_base_dir(mut self, dir: PathBuf) -> Self {
self.base_dir = Some(dir);
self
}
}
#[async_trait]
impl Tool for ReadFileTool {
fn name(&self) -> &str {
"read_file"
}
fn description(&self) -> &str {
"Read the contents of a file. Returns the file content as text. \
For large files, you can specify offset and limit to read a portion."
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {
"path": {
"type": "string",
"description": "Path to the file to read"
},
"offset": {
"type": "integer",
"description": "Line number to start reading from (1-indexed, optional)"
},
"limit": {
"type": "integer",
"description": "Maximum number of lines to read (optional)"
}
},
"required": ["path"]
})
}
async fn execute(
&self,
params: serde_json::Value,
_ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
let path_str = params
.get("path")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("missing 'path' parameter".into()))?;
let offset = params.get("offset").and_then(|v| v.as_u64()).unwrap_or(0) as usize;
let limit = params.get("limit").and_then(|v| v.as_u64());
let start = std::time::Instant::now();
let path = validate_path(path_str, self.base_dir.as_deref())?;
// Check file size
let metadata = fs::metadata(&path)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("Cannot access file: {}", e)))?;
if metadata.len() > MAX_READ_SIZE {
return Err(ToolError::ExecutionFailed(format!(
"File too large ({} bytes). Maximum is {} bytes. Use offset/limit for partial reads.",
metadata.len(),
MAX_READ_SIZE
)));
}
// Read file
let content = fs::read_to_string(&path)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("Failed to read file: {}", e)))?;
// Apply offset and limit
let lines: Vec<&str> = content.lines().collect();
let total_lines = lines.len();
let start_line = if offset > 0 {
offset.saturating_sub(1)
} else {
0
};
let end_line = if let Some(lim) = limit {
(start_line + lim as usize).min(total_lines)
} else {
total_lines
};
let selected_lines: Vec<String> = lines[start_line..end_line]
.iter()
.enumerate()
.map(|(i, line)| format!("{:>6}{}", start_line + i + 1, line))
.collect();
let result = serde_json::json!({
"content": selected_lines.join("\n"),
"total_lines": total_lines,
"lines_shown": end_line - start_line,
"path": path.display().to_string()
});
Ok(ToolOutput::success(result, start.elapsed()))
}
fn requires_sanitization(&self) -> bool {
true // File content could contain anything
}
}
/// Write file contents tool.
#[derive(Debug, Default)]
pub struct WriteFileTool {
base_dir: Option<PathBuf>,
}
impl WriteFileTool {
pub fn new() -> Self {
Self::default()
}
pub fn with_base_dir(mut self, dir: PathBuf) -> Self {
self.base_dir = Some(dir);
self
}
}
#[async_trait]
impl Tool for WriteFileTool {
fn name(&self) -> &str {
"write_file"
}
fn description(&self) -> &str {
"Write content to a file. Creates the file if it doesn't exist, overwrites if it does. \
Parent directories are created automatically. Use apply_patch for targeted edits to existing files."
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {
"path": {
"type": "string",
"description": "Path to the file to write"
},
"content": {
"type": "string",
"description": "Content to write to the file"
}
},
"required": ["path", "content"]
})
}
async fn execute(
&self,
params: serde_json::Value,
_ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
let path_str = params
.get("path")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("missing 'path' parameter".into()))?;
let content = params
.get("content")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("missing 'content' parameter".into()))?;
let start = std::time::Instant::now();
// Check content size
if content.len() > MAX_WRITE_SIZE {
return Err(ToolError::InvalidParameters(format!(
"Content too large ({} bytes). Maximum is {} bytes.",
content.len(),
MAX_WRITE_SIZE
)));
}
let path = validate_path(path_str, self.base_dir.as_deref())?;
// Create parent directories
if let Some(parent) = path.parent() {
fs::create_dir_all(parent).await.map_err(|e| {
ToolError::ExecutionFailed(format!("Failed to create directories: {}", e))
})?;
}
// Write file
fs::write(&path, content)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("Failed to write file: {}", e)))?;
let result = serde_json::json!({
"path": path.display().to_string(),
"bytes_written": content.len(),
"success": true
});
Ok(ToolOutput::success(result, start.elapsed()))
}
fn requires_approval(&self) -> bool {
true // File writes should require approval
}
fn requires_sanitization(&self) -> bool {
false // We're writing, not reading external data
}
}
/// List directory contents tool.
#[derive(Debug, Default)]
pub struct ListDirTool {
base_dir: Option<PathBuf>,
}
impl ListDirTool {
pub fn new() -> Self {
Self::default()
}
pub fn with_base_dir(mut self, dir: PathBuf) -> Self {
self.base_dir = Some(dir);
self
}
}
#[async_trait]
impl Tool for ListDirTool {
fn name(&self) -> &str {
"list_dir"
}
fn description(&self) -> &str {
"List contents of a directory. Shows files and subdirectories with their sizes. \
Use for exploring project structure."
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {
"path": {
"type": "string",
"description": "Path to the directory to list (defaults to current directory)"
},
"recursive": {
"type": "boolean",
"description": "If true, list contents recursively (default false)"
},
"max_depth": {
"type": "integer",
"description": "Maximum depth for recursive listing (default 3)"
}
},
"required": []
})
}
async fn execute(
&self,
params: serde_json::Value,
_ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
let path_str = params.get("path").and_then(|v| v.as_str()).unwrap_or(".");
let recursive = params
.get("recursive")
.and_then(|v| v.as_bool())
.unwrap_or(false);
let max_depth = params
.get("max_depth")
.and_then(|v| v.as_u64())
.unwrap_or(3) as usize;
let start = std::time::Instant::now();
let path = validate_path(path_str, self.base_dir.as_deref())?;
let mut entries = Vec::new();
list_dir_inner(&path, &path, recursive, max_depth, 0, &mut entries).await?;
// Sort entries
entries.sort_by(|a, b| {
let a_is_dir = a.ends_with('/');
let b_is_dir = b.ends_with('/');
match (a_is_dir, b_is_dir) {
(true, false) => std::cmp::Ordering::Less,
(false, true) => std::cmp::Ordering::Greater,
_ => a.cmp(b),
}
});
let truncated = entries.len() > MAX_DIR_ENTRIES;
if truncated {
entries.truncate(MAX_DIR_ENTRIES);
}
let result = serde_json::json!({
"path": path.display().to_string(),
"entries": entries,
"count": entries.len(),
"truncated": truncated
});
Ok(ToolOutput::success(result, start.elapsed()))
}
fn requires_sanitization(&self) -> bool {
false // Directory listings are safe
}
}
/// Recursively list directory contents.
async fn list_dir_inner(
base: &Path,
path: &Path,
recursive: bool,
max_depth: usize,
current_depth: usize,
entries: &mut Vec<String>,
) -> Result<(), ToolError> {
if entries.len() >= MAX_DIR_ENTRIES {
return Ok(());
}
let mut dir = fs::read_dir(path)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("Failed to read directory: {}", e)))?;
while let Some(entry) = dir
.next_entry()
.await
.map_err(|e| ToolError::ExecutionFailed(format!("Failed to read entry: {}", e)))?
{
if entries.len() >= MAX_DIR_ENTRIES {
break;
}
let entry_path = entry.path();
let relative = entry_path
.strip_prefix(base)
.unwrap_or(&entry_path)
.to_string_lossy();
let metadata = entry.metadata().await.ok();
let is_dir = metadata.as_ref().is_some_and(|m| m.is_dir());
let display = if is_dir {
format!("{}/", relative)
} else {
let size = metadata.as_ref().map(|m| m.len()).unwrap_or(0);
format!("{} ({})", relative, format_size(size))
};
entries.push(display);
if recursive && is_dir && current_depth < max_depth {
// Skip common non-essential directories
let name = entry.file_name();
let name_str = name.to_string_lossy();
if !matches!(
name_str.as_ref(),
"node_modules" | "target" | ".git" | "__pycache__" | "venv" | ".venv"
) {
Box::pin(list_dir_inner(
base,
&entry_path,
recursive,
max_depth,
current_depth + 1,
entries,
))
.await?;
}
}
}
Ok(())
}
/// Format file size in human-readable form.
fn format_size(bytes: u64) -> String {
const KB: u64 = 1024;
const MB: u64 = KB * 1024;
const GB: u64 = MB * 1024;
if bytes >= GB {
format!("{:.1}GB", bytes as f64 / GB as f64)
} else if bytes >= MB {
format!("{:.1}MB", bytes as f64 / MB as f64)
} else if bytes >= KB {
format!("{:.1}KB", bytes as f64 / KB as f64)
} else {
format!("{}B", bytes)
}
}
/// Apply patch tool for targeted file edits.
#[derive(Debug, Default)]
pub struct ApplyPatchTool {
base_dir: Option<PathBuf>,
}
impl ApplyPatchTool {
pub fn new() -> Self {
Self::default()
}
pub fn with_base_dir(mut self, dir: PathBuf) -> Self {
self.base_dir = Some(dir);
self
}
}
#[async_trait]
impl Tool for ApplyPatchTool {
fn name(&self) -> &str {
"apply_patch"
}
fn description(&self) -> &str {
"Apply targeted edits to a file using search/replace. Finds the exact 'old_string' \
and replaces it with 'new_string'. Use for surgical code changes without rewriting entire files. \
The old_string must match exactly (including whitespace and indentation)."
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {
"path": {
"type": "string",
"description": "Path to the file to edit"
},
"old_string": {
"type": "string",
"description": "The exact string to find and replace"
},
"new_string": {
"type": "string",
"description": "The string to replace it with"
},
"replace_all": {
"type": "boolean",
"description": "If true, replace all occurrences (default false, replaces first only)"
}
},
"required": ["path", "old_string", "new_string"]
})
}
async fn execute(
&self,
params: serde_json::Value,
_ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
let path_str = params
.get("path")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("missing 'path' parameter".into()))?;
let old_string = params
.get("old_string")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("missing 'old_string' parameter".into()))?;
let new_string = params
.get("new_string")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("missing 'new_string' parameter".into()))?;
let replace_all = params
.get("replace_all")
.and_then(|v| v.as_bool())
.unwrap_or(false);
let start = std::time::Instant::now();
let path = validate_path(path_str, self.base_dir.as_deref())?;
// Read current content
let content = fs::read_to_string(&path)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("Failed to read file: {}", e)))?;
// Check if old_string exists
if !content.contains(old_string) {
return Err(ToolError::ExecutionFailed(format!(
"Could not find the specified text in {}. Make sure old_string matches exactly.",
path.display()
)));
}
// Apply replacement
let new_content = if replace_all {
content.replace(old_string, new_string)
} else {
content.replacen(old_string, new_string, 1)
};
// Count replacements
let replacements = if replace_all {
content.matches(old_string).count()
} else {
1
};
// Write back
fs::write(&path, &new_content)
.await
.map_err(|e| ToolError::ExecutionFailed(format!("Failed to write file: {}", e)))?;
let result = serde_json::json!({
"path": path.display().to_string(),
"replacements": replacements,
"success": true
});
Ok(ToolOutput::success(result, start.elapsed()))
}
fn requires_approval(&self) -> bool {
true // File edits should require approval
}
fn requires_sanitization(&self) -> bool {
false // We're writing, not reading external data
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
#[tokio::test]
async fn test_read_file() {
let dir = TempDir::new().unwrap();
let file_path = dir.path().join("test.txt");
std::fs::write(&file_path, "line 1\nline 2\nline 3\n").unwrap();
let tool = ReadFileTool::new().with_base_dir(dir.path().to_path_buf());
let ctx = JobContext::default();
let result = tool
.execute(
serde_json::json!({"path": file_path.to_str().unwrap()}),
&ctx,
)
.await
.unwrap();
let content = result.result.get("content").unwrap().as_str().unwrap();
assert!(content.contains("line 1"));
assert!(content.contains("line 2"));
}
#[tokio::test]
async fn test_write_file() {
let dir = TempDir::new().unwrap();
let file_path = dir.path().join("new_file.txt");
let tool = WriteFileTool::new().with_base_dir(dir.path().to_path_buf());
let ctx = JobContext::default();
let result = tool
.execute(
serde_json::json!({
"path": file_path.to_str().unwrap(),
"content": "hello world"
}),
&ctx,
)
.await
.unwrap();
assert!(result.result.get("success").unwrap().as_bool().unwrap());
assert_eq!(std::fs::read_to_string(&file_path).unwrap(), "hello world");
}
#[tokio::test]
async fn test_apply_patch() {
let dir = TempDir::new().unwrap();
let file_path = dir.path().join("code.rs");
std::fs::write(&file_path, "fn main() {\n println!(\"old\");\n}\n").unwrap();
let tool = ApplyPatchTool::new().with_base_dir(dir.path().to_path_buf());
let ctx = JobContext::default();
let result = tool
.execute(
serde_json::json!({
"path": file_path.to_str().unwrap(),
"old_string": "println!(\"old\")",
"new_string": "println!(\"new\")"
}),
&ctx,
)
.await
.unwrap();
assert!(result.result.get("success").unwrap().as_bool().unwrap());
let content = std::fs::read_to_string(&file_path).unwrap();
assert!(content.contains("println!(\"new\")"));
}
#[tokio::test]
async fn test_list_dir() {
let dir = TempDir::new().unwrap();
std::fs::write(dir.path().join("file1.txt"), "content").unwrap();
std::fs::create_dir(dir.path().join("subdir")).unwrap();
let tool = ListDirTool::new();
let ctx = JobContext::default();
let result = tool
.execute(
serde_json::json!({"path": dir.path().to_str().unwrap()}),
&ctx,
)
.await
.unwrap();
let entries = result.result.get("entries").unwrap().as_array().unwrap();
assert!(entries.len() >= 2);
}
}
+5 -1
View File
@@ -2,20 +2,24 @@
mod echo;
mod ecommerce;
mod file;
mod http;
mod json;
mod marketplace;
mod memory;
mod restaurant;
mod shell;
mod taskrabbit;
mod time;
pub use echo::EchoTool;
pub use ecommerce::EcommerceTool;
pub use file::{ApplyPatchTool, ListDirTool, ReadFileTool, WriteFileTool};
pub use http::HttpTool;
pub use json::JsonTool;
pub use marketplace::MarketplaceTool;
pub use memory::{MemoryReadTool, MemorySearchTool, MemoryWriteTool};
pub use memory::{MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool};
pub use restaurant::RestaurantTool;
pub use shell::ShellTool;
pub use taskrabbit::TaskRabbitTool;
pub use time::TimeTool;
+351
View File
@@ -0,0 +1,351 @@
//! Shell execution tool for running commands in a sandboxed environment.
//!
//! Provides controlled command execution with:
//! - Working directory isolation
//! - Timeout enforcement
//! - Output capture and truncation
//! - Blocked command patterns for safety
use std::collections::HashSet;
use std::path::PathBuf;
use std::process::Stdio;
use std::sync::LazyLock;
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::{Tool, ToolError, ToolOutput};
/// Maximum output size before truncation (64KB).
const MAX_OUTPUT_SIZE: usize = 64 * 1024;
/// Default command timeout.
const DEFAULT_TIMEOUT: Duration = Duration::from_secs(120);
/// Commands that are always blocked for safety.
static BLOCKED_COMMANDS: LazyLock<HashSet<&'static str>> = LazyLock::new(|| {
HashSet::from([
"rm -rf /",
"rm -rf /*",
":(){ :|:& };:", // Fork bomb
"dd if=/dev/zero",
"mkfs",
"chmod -R 777 /",
"> /dev/sda",
"curl | sh",
"wget | sh",
"curl | bash",
"wget | bash",
])
});
/// Patterns that indicate potentially dangerous commands.
static DANGEROUS_PATTERNS: LazyLock<Vec<&'static str>> = LazyLock::new(|| {
vec![
"sudo ",
"doas ",
" | sh",
" | bash",
" | zsh",
"eval ",
"$(curl",
"$(wget",
"/etc/passwd",
"/etc/shadow",
"~/.ssh",
".bash_history",
"id_rsa",
]
});
/// Shell command execution tool.
#[derive(Debug)]
pub struct ShellTool {
/// Working directory for commands (if None, uses job's working dir or cwd).
working_dir: Option<PathBuf>,
/// Command timeout.
timeout: Duration,
/// Whether to allow potentially dangerous commands (requires explicit approval).
allow_dangerous: bool,
}
impl ShellTool {
/// Create a new shell tool with default settings.
pub fn new() -> Self {
Self {
working_dir: None,
timeout: DEFAULT_TIMEOUT,
allow_dangerous: false,
}
}
/// Set the working directory.
pub fn with_working_dir(mut self, dir: PathBuf) -> Self {
self.working_dir = Some(dir);
self
}
/// Set the command timeout.
pub fn with_timeout(mut self, timeout: Duration) -> Self {
self.timeout = timeout;
self
}
/// Check if a command is blocked.
fn is_blocked(&self, cmd: &str) -> Option<&'static str> {
let normalized = cmd.to_lowercase();
for blocked in BLOCKED_COMMANDS.iter() {
if normalized.contains(blocked) {
return Some("Command contains blocked pattern");
}
}
if !self.allow_dangerous {
for pattern in DANGEROUS_PATTERNS.iter() {
if normalized.contains(pattern) {
return Some("Command contains potentially dangerous pattern");
}
}
}
None
}
/// Execute a command and capture output.
async fn execute_command(
&self,
cmd: &str,
workdir: Option<&str>,
timeout: Option<u64>,
) -> Result<(String, i32), ToolError> {
// Check for blocked commands
if let Some(reason) = self.is_blocked(cmd) {
return Err(ToolError::NotAuthorized(format!(
"{}: {}",
reason,
truncate_for_error(cmd)
)));
}
// Determine working directory
let cwd = workdir
.map(PathBuf::from)
.or_else(|| self.working_dir.clone())
.unwrap_or_else(|| std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")));
// Build command
let mut command = if cfg!(target_os = "windows") {
let mut c = Command::new("cmd");
c.args(["/C", cmd]);
c
} else {
let mut c = Command::new("sh");
c.args(["-c", cmd]);
c
};
command
.current_dir(&cwd)
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped());
// Spawn process
let mut child = command
.spawn()
.map_err(|e| ToolError::ExecutionFailed(format!("Failed to spawn command: {}", e)))?;
// Determine timeout
let timeout_duration = timeout.map(Duration::from_secs).unwrap_or(self.timeout);
// Wait with timeout
let result = tokio::time::timeout(timeout_duration, async {
let status = child.wait().await?;
// Read stdout
let mut stdout = String::new();
if let Some(mut out) = child.stdout.take() {
let mut buf = vec![0u8; MAX_OUTPUT_SIZE];
let n = out.read(&mut buf).await.unwrap_or(0);
stdout = String::from_utf8_lossy(&buf[..n]).to_string();
}
// Read stderr
let mut stderr = String::new();
if let Some(mut err) = child.stderr.take() {
let mut buf = vec![0u8; MAX_OUTPUT_SIZE];
let n = err.read(&mut buf).await.unwrap_or(0);
stderr = String::from_utf8_lossy(&buf[..n]).to_string();
}
// Combine output
let output = if stderr.is_empty() {
stdout
} else if stdout.is_empty() {
stderr
} else {
format!("{}\n\n--- stderr ---\n{}", stdout, stderr)
};
Ok::<_, std::io::Error>((output, status.code().unwrap_or(-1)))
})
.await;
match result {
Ok(Ok((output, code))) => Ok((truncate_output(&output), code)),
Ok(Err(e)) => Err(ToolError::ExecutionFailed(format!(
"Command execution failed: {}",
e
))),
Err(_) => {
// Timeout - try to kill the process
let _ = child.kill().await;
Err(ToolError::Timeout(timeout_duration))
}
}
}
}
impl Default for ShellTool {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl Tool for ShellTool {
fn name(&self) -> &str {
"shell"
}
fn description(&self) -> &str {
"Execute shell commands. Use for running builds, tests, git operations, and other CLI tasks. \
Commands run in a subprocess with captured output. Long-running commands have a timeout."
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {
"command": {
"type": "string",
"description": "The shell command to execute"
},
"workdir": {
"type": "string",
"description": "Working directory for the command (optional)"
},
"timeout": {
"type": "integer",
"description": "Timeout in seconds (optional, default 120)"
}
},
"required": ["command"]
})
}
async fn execute(
&self,
params: serde_json::Value,
_ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
let command = params
.get("command")
.and_then(|v| v.as_str())
.ok_or_else(|| ToolError::InvalidParameters("missing 'command' parameter".into()))?;
let workdir = params.get("workdir").and_then(|v| v.as_str());
let timeout = params.get("timeout").and_then(|v| v.as_u64());
let start = std::time::Instant::now();
let (output, exit_code) = self.execute_command(command, workdir, timeout).await?;
let duration = start.elapsed();
let result = serde_json::json!({
"output": output,
"exit_code": exit_code,
"success": exit_code == 0
});
Ok(ToolOutput::success(result, duration))
}
fn requires_approval(&self) -> bool {
true // Shell commands should require approval
}
fn requires_sanitization(&self) -> bool {
true // Shell output could contain anything
}
}
/// Truncate output to fit within limits.
fn truncate_output(s: &str) -> String {
if s.len() <= MAX_OUTPUT_SIZE {
s.to_string()
} else {
let half = MAX_OUTPUT_SIZE / 2;
format!(
"{}\n\n... [truncated {} bytes] ...\n\n{}",
&s[..half],
s.len() - MAX_OUTPUT_SIZE,
&s[s.len() - half..]
)
}
}
/// Truncate command for error messages.
fn truncate_for_error(s: &str) -> String {
if s.len() <= 100 {
s.to_string()
} else {
format!("{}...", &s[..100])
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_echo_command() {
let tool = ShellTool::new();
let ctx = JobContext::default();
let result = tool
.execute(serde_json::json!({"command": "echo hello"}), &ctx)
.await
.unwrap();
let output = result.result.get("output").unwrap().as_str().unwrap();
assert!(output.contains("hello"));
assert_eq!(result.result.get("exit_code").unwrap().as_i64().unwrap(), 0);
}
#[test]
fn test_blocked_commands() {
let tool = ShellTool::new();
assert!(tool.is_blocked("rm -rf /").is_some());
assert!(tool.is_blocked("sudo rm file").is_some());
assert!(tool.is_blocked("curl http://x | sh").is_some());
assert!(tool.is_blocked("echo hello").is_none());
assert!(tool.is_blocked("cargo build").is_none());
}
#[tokio::test]
async fn test_command_timeout() {
let tool = ShellTool::new().with_timeout(Duration::from_millis(100));
let ctx = JobContext::default();
let result = tool
.execute(serde_json::json!({"command": "sleep 10"}), &ctx)
.await;
assert!(matches!(result, Err(ToolError::Timeout(_))));
}
}
+7 -2
View File
@@ -5,17 +5,22 @@
//! - Interact with the marketplace
//! - Execute sandboxed code (via WASM sandbox)
//! - Delegate tasks to other services
//! - Build new software and tools
pub mod builder;
pub mod builtin;
pub mod mcp;
pub mod wasm;
mod builder;
mod registry;
mod sandbox;
mod tool;
pub use builder::{DynamicTool, SandboxConfig, ToolBuilder, ToolRequirement};
pub use builder::{
BuildPhase, BuildRequirement, BuildResult, BuildSoftwareTool, BuilderConfig, Language,
LlmSoftwareBuilder, SoftwareBuilder, SoftwareType, Template, TemplateEngine, TemplateType,
TestCase, TestHarness, TestResult, TestSuite, ValidationError, ValidationResult, WasmValidator,
};
pub use registry::ToolRegistry;
pub use sandbox::ToolSandbox;
pub use tool::{Tool, ToolError, ToolOutput};
+66 -2
View File
@@ -5,13 +5,19 @@ use std::sync::Arc;
use tokio::sync::RwLock;
use crate::llm::ToolDefinition;
use crate::tools::builtin::{EchoTool, HttpTool, JsonTool, TimeTool};
use crate::llm::{LlmProvider, ToolDefinition};
use crate::safety::SafetyLayer;
use crate::tools::builder::{BuildSoftwareTool, BuilderConfig, LlmSoftwareBuilder};
use crate::tools::builtin::{
ApplyPatchTool, EchoTool, HttpTool, JsonTool, ListDirTool, MemoryReadTool, MemorySearchTool,
MemoryTreeTool, MemoryWriteTool, ReadFileTool, ShellTool, TimeTool, WriteFileTool,
};
use crate::tools::tool::Tool;
use crate::tools::wasm::{
Capabilities, ResourceLimits, WasmError, WasmStorageError, WasmToolRuntime, WasmToolStore,
WasmToolWrapper,
};
use crate::workspace::Workspace;
/// Registry of available tools.
pub struct ToolRegistry {
@@ -110,6 +116,64 @@ impl ToolRegistry {
tracing::info!("Registered {} built-in tools", self.count());
}
/// Register development tools for building software.
///
/// These tools provide shell access, file operations, and code editing
/// capabilities needed for the software builder. Call this after
/// `register_builtin_tools()` to enable code generation features.
pub fn register_dev_tools(&self) {
self.register_sync(Arc::new(ShellTool::new()));
self.register_sync(Arc::new(ReadFileTool::new()));
self.register_sync(Arc::new(WriteFileTool::new()));
self.register_sync(Arc::new(ListDirTool::new()));
self.register_sync(Arc::new(ApplyPatchTool::new()));
tracing::info!("Registered 5 development tools");
}
/// Register memory tools with a workspace.
///
/// Memory tools require a workspace for persistence. Call this after
/// `register_builtin_tools()` if you have a workspace available.
pub fn register_memory_tools(&self, workspace: Arc<Workspace>) {
self.register_sync(Arc::new(MemorySearchTool::new(Arc::clone(&workspace))));
self.register_sync(Arc::new(MemoryWriteTool::new(Arc::clone(&workspace))));
self.register_sync(Arc::new(MemoryReadTool::new(Arc::clone(&workspace))));
self.register_sync(Arc::new(MemoryTreeTool::new(workspace)));
tracing::info!("Registered 4 memory tools");
}
/// Register the software builder tool.
///
/// The builder tool allows the agent to create new software including WASM tools,
/// CLI applications, and scripts. It uses an LLM-driven iterative build loop.
///
/// This also registers the dev tools (shell, file operations) needed by the builder.
pub async fn register_builder_tool(
self: &Arc<Self>,
llm: Arc<dyn LlmProvider>,
safety: Arc<SafetyLayer>,
config: Option<BuilderConfig>,
) {
// First register dev tools needed by the builder
self.register_dev_tools();
// Create the builder (arg order: config, llm, safety, tools)
let builder = Arc::new(LlmSoftwareBuilder::new(
config.unwrap_or_default(),
llm,
safety,
Arc::clone(self),
));
// Register the build_software tool
self.register(Arc::new(BuildSoftwareTool::new(builder)))
.await;
tracing::info!("Registered software builder tool");
}
/// Register a WASM tool from bytes.
///
/// This validates and compiles the WASM component, then registers it as a tool.
+30 -1
View File
@@ -1,10 +1,39 @@
//! Sandboxed tool execution environment.
//!
//! NOTE: For WASM-based sandboxing with full security, use the `wasm` module instead.
//! This module provides a simpler process-based sandbox for scripts.
use std::time::Duration;
use crate::tools::builder::SandboxConfig;
use crate::tools::tool::ToolError;
/// Configuration for the sandbox.
#[derive(Debug, Clone)]
pub struct SandboxConfig {
/// Maximum execution time.
pub max_execution_time: Duration,
/// Maximum memory in bytes.
pub max_memory_bytes: u64,
/// Allowed network hosts (empty = no network).
pub allowed_hosts: Vec<String>,
/// Allowed filesystem paths (empty = no filesystem).
pub allowed_paths: Vec<String>,
/// Environment variables to pass.
pub env_vars: Vec<(String, String)>,
}
impl Default for SandboxConfig {
fn default() -> Self {
Self {
max_execution_time: Duration::from_secs(30),
max_memory_bytes: 128 * 1024 * 1024, // 128 MB
allowed_hosts: vec![],
allowed_paths: vec![],
env_vars: vec![],
}
}
}
/// Result of a sandboxed execution.
#[derive(Debug)]
pub struct SandboxResult {