From 343782524fbea8bf2ff769e79e09444f00b4f7b4 Mon Sep 17 00:00:00 2001 From: Illia Polosukhin Date: Tue, 3 Feb 2026 00:48:22 -0800 Subject: [PATCH] Load tools at launch --- src/agent/context_monitor.rs | 2 +- src/cli/tool.rs | 77 +++++++++++++++++++++++++++--------- src/config.rs | 14 +++++++ src/main.rs | 42 +++++++++++++++++++- src/tools/wasm/wrapper.rs | 2 +- wit/tool.wit | 4 +- 6 files changed, 117 insertions(+), 24 deletions(-) diff --git a/src/agent/context_monitor.rs b/src/agent/context_monitor.rs index c236a520..03677cdc 100644 --- a/src/agent/context_monitor.rs +++ b/src/agent/context_monitor.rs @@ -68,7 +68,7 @@ impl ContextMonitor { /// Estimate the token count for a list of messages. pub fn estimate_tokens(&self, messages: &[ChatMessage]) -> usize { - messages.iter().map(|m| estimate_message_tokens(m)).sum() + messages.iter().map(estimate_message_tokens).sum() } /// Check if compaction is needed. diff --git a/src/cli/tool.rs b/src/cli/tool.rs index 3161eb98..bee90001 100644 --- a/src/cli/tool.rs +++ b/src/cli/tool.rs @@ -267,14 +267,36 @@ fn build_wasm_component(source_dir: &Path, release: bool) -> anyhow::Result>() + .join(", ") + ) + })?; // Look for .wasm files in target dir - let entries: Vec<_> = std::fs::read_dir(&target_dir)? + let entries: Vec<_> = std::fs::read_dir(target_dir)? .filter_map(|e| e.ok()) .filter(|e| { e.path() @@ -307,26 +329,45 @@ fn build_wasm_component(source_dir: &Path, release: bool) -> anyhow::Result anyhow::Result { let profile = if release { "release" } else { "debug" }; - let target_dir = source_dir - .join("target") - .join("wasm32-wasip2") - .join(profile); - // Try exact name match first - let snake_name = name.replace('-', "_"); - let candidates = [ - target_dir.join(format!("{}.wasm", name)), - target_dir.join(format!("{}.wasm", snake_name)), + // cargo-component may output to wasm32-wasip1 or wasm32-wasip2 depending on version + let target_dirs = [ + source_dir + .join("target") + .join("wasm32-wasip1") + .join(profile), + source_dir + .join("target") + .join("wasm32-wasip2") + .join(profile), + source_dir + .join("target") + .join("wasm32-unknown-unknown") + .join(profile), ]; - for candidate in &candidates { - if candidate.exists() { - return Ok(candidate.clone()); + let snake_name = name.replace('-', "_"); + + // Try exact name match in any target dir first + for target_dir in &target_dirs { + let candidates = [ + target_dir.join(format!("{}.wasm", name)), + target_dir.join(format!("{}.wasm", snake_name)), + ]; + for candidate in &candidates { + if candidate.exists() { + return Ok(candidate.clone()); + } } } + // Find a target dir that exists + let target_dir = target_dirs.iter().find(|p| p.exists()).ok_or_else(|| { + anyhow::anyhow!("No target directory found. Run without --skip-build to build first.") + })?; + // Fall back to any .wasm file - let entries: Vec<_> = std::fs::read_dir(&target_dir) + let entries: Vec<_> = std::fs::read_dir(target_dir) .map_err(|_| { anyhow::anyhow!( "Target directory not found: {}. Run without --skip-build.", diff --git a/src/config.rs b/src/config.rs index 20d7d2dc..c3ebf3de 100644 --- a/src/config.rs +++ b/src/config.rs @@ -240,6 +240,8 @@ impl SafetyConfig { pub struct WasmConfig { /// Whether WASM tool execution is enabled. pub enabled: bool, + /// Directory containing installed WASM tools (default: ~/.near-agent/tools/). + pub tools_dir: PathBuf, /// Default memory limit in bytes (default: 10 MB). pub default_memory_limit: u64, /// Default execution timeout in seconds (default: 60). @@ -302,6 +304,7 @@ impl Default for WasmConfig { fn default() -> Self { Self { enabled: true, + tools_dir: default_tools_dir(), default_memory_limit: 10 * 1024 * 1024, // 10 MB default_timeout_secs: 60, default_fuel_limit: 10_000_000, @@ -311,6 +314,14 @@ impl Default for WasmConfig { } } +/// Get the default tools directory (~/.near-agent/tools/). +fn default_tools_dir() -> PathBuf { + dirs::home_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join(".near-agent") + .join("tools") +} + impl WasmConfig { fn from_env() -> Result { Ok(Self { @@ -322,6 +333,9 @@ impl WasmConfig { message: format!("must be 'true' or 'false': {e}"), })? .unwrap_or(true), + tools_dir: optional_env("WASM_TOOLS_DIR")? + .map(PathBuf::from) + .unwrap_or_else(default_tools_dir), default_memory_limit: parse_optional_env( "WASM_DEFAULT_MEMORY_LIMIT", 10 * 1024 * 1024, diff --git a/src/main.rs b/src/main.rs index 0c411ed9..4daa555d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -13,7 +13,10 @@ use near_agent::{ history::Store, llm::create_llm_provider, safety::SafetyLayer, - tools::ToolRegistry, + tools::{ + ToolRegistry, + wasm::{WasmToolLoader, WasmToolRuntime}, + }, }; #[tokio::main] @@ -86,7 +89,42 @@ async fn main() -> anyhow::Result<()> { // Initialize tool registry let tools = Arc::new(ToolRegistry::new()); tools.register_builtin_tools(); - tracing::info!("Tool registry initialized with {} tools", tools.count()); + tracing::info!("Registered {} built-in tools", tools.count()); + + // Load installed WASM tools + if config.wasm.enabled && config.wasm.tools_dir.exists() { + match WasmToolRuntime::new(config.wasm.to_runtime_config()) { + Ok(runtime) => { + let runtime = Arc::new(runtime); + let loader = WasmToolLoader::new(Arc::clone(&runtime), Arc::clone(&tools)); + + match loader.load_from_dir(&config.wasm.tools_dir).await { + Ok(results) => { + if !results.loaded.is_empty() { + tracing::info!( + "Loaded {} WASM tools from {}", + results.loaded.len(), + config.wasm.tools_dir.display() + ); + } + for (path, err) in &results.errors { + tracing::warn!("Failed to load WASM tool {}: {}", path.display(), err); + } + } + Err(e) => { + tracing::warn!("Failed to scan WASM tools directory: {}", e); + } + } + } + Err(e) => { + tracing::warn!("Failed to initialize WASM runtime: {}", e); + } + } + } + tracing::info!( + "Tool registry initialized with {} total tools", + tools.count() + ); // Initialize channel manager let mut channels = ChannelManager::new(); diff --git a/src/tools/wasm/wrapper.rs b/src/tools/wasm/wrapper.rs index e8ac80b2..b499cc3b 100644 --- a/src/tools/wasm/wrapper.rs +++ b/src/tools/wasm/wrapper.rs @@ -258,7 +258,7 @@ fn extract_response(response: &Val) -> Result<(Option, Option), for (name, val) in fields { match name.as_str() { - "result" => { + "output" => { if let Val::Option(Some(inner)) = val { if let Val::String(s) = inner.as_ref() { result = Some(s.to_string()); diff --git a/wit/tool.wit b/wit/tool.wit index 37c2a0a3..a6a45a3d 100644 --- a/wit/tool.wit +++ b/wit/tool.wit @@ -112,8 +112,8 @@ interface tool { /// Response from tool execution. record response { - /// JSON-encoded result on success. - result: option, + /// JSON-encoded output on success. + output: option, /// Error message on failure. error: option, }