Load tools at launch

This commit is contained in:
Illia Polosukhin
2026-02-03 00:48:22 -08:00
parent 575269546b
commit 343782524f
6 changed files with 117 additions and 24 deletions
+1 -1
View File
@@ -68,7 +68,7 @@ impl ContextMonitor {
/// Estimate the token count for a list of messages. /// Estimate the token count for a list of messages.
pub fn estimate_tokens(&self, messages: &[ChatMessage]) -> usize { 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. /// Check if compaction is needed.
+59 -18
View File
@@ -267,14 +267,36 @@ fn build_wasm_component(source_dir: &Path, release: bool) -> anyhow::Result<Path
} }
// Find the output wasm file // Find the output wasm file
// cargo-component may output to wasm32-wasip1 or wasm32-wasip2 depending on version
let profile = if release { "release" } else { "debug" }; let profile = if release { "release" } else { "debug" };
let target_dir = source_dir let candidates = [
.join("target") source_dir
.join("wasm32-wasip2") .join("target")
.join(profile); .join("wasm32-wasip1")
.join(profile),
source_dir
.join("target")
.join("wasm32-wasip2")
.join(profile),
source_dir
.join("target")
.join("wasm32-unknown-unknown")
.join(profile),
];
let target_dir = candidates.iter().find(|p| p.exists()).ok_or_else(|| {
anyhow::anyhow!(
"No WASM target directory found. Expected one of: {}",
candidates
.iter()
.map(|p| p.display().to_string())
.collect::<Vec<_>>()
.join(", ")
)
})?;
// Look for .wasm files in target dir // 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_map(|e| e.ok())
.filter(|e| { .filter(|e| {
e.path() e.path()
@@ -307,26 +329,45 @@ fn build_wasm_component(source_dir: &Path, release: bool) -> anyhow::Result<Path
/// Find an existing WASM artifact without building. /// Find an existing WASM artifact without building.
fn find_wasm_artifact(source_dir: &Path, name: &str, release: bool) -> anyhow::Result<PathBuf> { fn find_wasm_artifact(source_dir: &Path, name: &str, release: bool) -> anyhow::Result<PathBuf> {
let profile = if release { "release" } else { "debug" }; let profile = if release { "release" } else { "debug" };
let target_dir = source_dir
.join("target")
.join("wasm32-wasip2")
.join(profile);
// Try exact name match first // cargo-component may output to wasm32-wasip1 or wasm32-wasip2 depending on version
let snake_name = name.replace('-', "_"); let target_dirs = [
let candidates = [ source_dir
target_dir.join(format!("{}.wasm", name)), .join("target")
target_dir.join(format!("{}.wasm", snake_name)), .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 { let snake_name = name.replace('-', "_");
if candidate.exists() {
return Ok(candidate.clone()); // 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 // 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(|_| { .map_err(|_| {
anyhow::anyhow!( anyhow::anyhow!(
"Target directory not found: {}. Run without --skip-build.", "Target directory not found: {}. Run without --skip-build.",
+14
View File
@@ -240,6 +240,8 @@ impl SafetyConfig {
pub struct WasmConfig { pub struct WasmConfig {
/// Whether WASM tool execution is enabled. /// Whether WASM tool execution is enabled.
pub enabled: bool, pub enabled: bool,
/// Directory containing installed WASM tools (default: ~/.near-agent/tools/).
pub tools_dir: PathBuf,
/// Default memory limit in bytes (default: 10 MB). /// Default memory limit in bytes (default: 10 MB).
pub default_memory_limit: u64, pub default_memory_limit: u64,
/// Default execution timeout in seconds (default: 60). /// Default execution timeout in seconds (default: 60).
@@ -302,6 +304,7 @@ impl Default for WasmConfig {
fn default() -> Self { fn default() -> Self {
Self { Self {
enabled: true, enabled: true,
tools_dir: default_tools_dir(),
default_memory_limit: 10 * 1024 * 1024, // 10 MB default_memory_limit: 10 * 1024 * 1024, // 10 MB
default_timeout_secs: 60, default_timeout_secs: 60,
default_fuel_limit: 10_000_000, 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 { impl WasmConfig {
fn from_env() -> Result<Self, ConfigError> { fn from_env() -> Result<Self, ConfigError> {
Ok(Self { Ok(Self {
@@ -322,6 +333,9 @@ impl WasmConfig {
message: format!("must be 'true' or 'false': {e}"), message: format!("must be 'true' or 'false': {e}"),
})? })?
.unwrap_or(true), .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( default_memory_limit: parse_optional_env(
"WASM_DEFAULT_MEMORY_LIMIT", "WASM_DEFAULT_MEMORY_LIMIT",
10 * 1024 * 1024, 10 * 1024 * 1024,
+40 -2
View File
@@ -13,7 +13,10 @@ use near_agent::{
history::Store, history::Store,
llm::create_llm_provider, llm::create_llm_provider,
safety::SafetyLayer, safety::SafetyLayer,
tools::ToolRegistry, tools::{
ToolRegistry,
wasm::{WasmToolLoader, WasmToolRuntime},
},
}; };
#[tokio::main] #[tokio::main]
@@ -86,7 +89,42 @@ async fn main() -> anyhow::Result<()> {
// Initialize tool registry // Initialize tool registry
let tools = Arc::new(ToolRegistry::new()); let tools = Arc::new(ToolRegistry::new());
tools.register_builtin_tools(); 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 // Initialize channel manager
let mut channels = ChannelManager::new(); let mut channels = ChannelManager::new();
+1 -1
View File
@@ -258,7 +258,7 @@ fn extract_response(response: &Val) -> Result<(Option<String>, Option<String>),
for (name, val) in fields { for (name, val) in fields {
match name.as_str() { match name.as_str() {
"result" => { "output" => {
if let Val::Option(Some(inner)) = val { if let Val::Option(Some(inner)) = val {
if let Val::String(s) = inner.as_ref() { if let Val::String(s) = inner.as_ref() {
result = Some(s.to_string()); result = Some(s.to_string());
+2 -2
View File
@@ -112,8 +112,8 @@ interface tool {
/// Response from tool execution. /// Response from tool execution.
record response { record response {
/// JSON-encoded result on success. /// JSON-encoded output on success.
result: option<string>, output: option<string>,
/// Error message on failure. /// Error message on failure.
error: option<string>, error: option<string>,
} }