diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index bc905e32..bac4df94 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -42,6 +42,7 @@ use crate::config::{AgentConfig, HeartbeatConfig}; use crate::context::ContextManager; use crate::context::JobContext; use crate::error::Error; +use crate::extensions::ExtensionManager; use crate::history::Store; use crate::llm::{ChatMessage, LlmProvider, Reasoning, ReasoningContext, RespondResult}; use crate::safety::SafetyLayer; @@ -68,6 +69,7 @@ pub struct AgentDeps { pub safety: Arc, pub tools: Arc, pub workspace: Option>, + pub extension_manager: Option>, } /// The main agent that coordinates all components. @@ -374,13 +376,6 @@ impl Agent { } async fn handle_message(&self, message: &IncomingMessage) -> Result, Error> { - tracing::debug!( - "Received message from {} on {}: {}", - message.user_id, - message.channel, - truncate(&message.content, 100) - ); - // Parse submission type first let submission = SubmissionParser::parse(&message.content); @@ -394,6 +389,41 @@ impl Agent { ) .await; + // Auth mode interception: if the thread is awaiting a token, route + // the message directly to the credential store. Nothing touches + // logs, turns, history, or compaction. + let pending_auth = { + let sess = session.lock().await; + sess.threads + .get(&thread_id) + .and_then(|t| t.pending_auth.clone()) + }; + + if let Some(pending) = pending_auth { + match &submission { + Submission::UserInput { content } => { + return self + .process_auth_token(message, &pending, content, session, thread_id) + .await; + } + _ => { + // Any control submission (interrupt, undo, etc.) cancels auth mode + let mut sess = session.lock().await; + if let Some(thread) = sess.threads.get_mut(&thread_id) { + thread.pending_auth = None; + } + // Fall through to normal handling + } + } + } + + tracing::debug!( + "Received message from {} on {} ({} chars)", + message.user_id, + message.channel, + message.content.len() + ); + // Process based on submission type let result = match submission { Submission::UserInput { content } => { @@ -867,6 +897,19 @@ impl Agent { } } + // If tool_auth returned awaiting_token, enter auth mode + // and short-circuit: return the instructions directly so + // the LLM doesn't get a chance to hallucinate tool calls. + if let Some((ext_name, instructions)) = + detect_auth_awaiting(&tc.name, &tool_result) + { + let mut sess = session.lock().await; + if let Some(thread) = sess.threads.get_mut(&thread_id) { + thread.enter_auth_mode(ext_name); + } + return Ok(AgenticLoopResult::Response(instructions)); + } + // Add tool result to context for next LLM call let result_content = match tool_result { Ok(output) => { @@ -1236,6 +1279,29 @@ impl Agent { } } + // If tool_auth returned awaiting_token, enter auth mode and + // return instructions directly (skip agentic loop continuation). + if let Some((ext_name, instructions)) = + detect_auth_awaiting(&pending.tool_name, &tool_result) + { + { + let mut sess = session.lock().await; + if let Some(thread) = sess.threads.get_mut(&thread_id) { + thread.enter_auth_mode(ext_name); + thread.complete_turn(&instructions); + } + } + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::Status("Awaiting token".into()), + &message.metadata, + ) + .await; + return Ok(SubmissionResult::response(instructions)); + } + // Add tool result to context let result_content = match tool_result { Ok(output) => { @@ -1336,6 +1402,75 @@ impl Agent { } } + /// Handle an auth token submitted while the thread is in auth mode. + /// + /// The token goes directly to the extension manager's credential store, + /// completely bypassing logging, turn creation, history, and compaction. + async fn process_auth_token( + &self, + message: &IncomingMessage, + pending: &crate::agent::session::PendingAuth, + token: &str, + session: Arc>, + thread_id: Uuid, + ) -> Result, Error> { + let token = token.trim(); + + // Clear auth mode regardless of outcome + { + let mut sess = session.lock().await; + if let Some(thread) = sess.threads.get_mut(&thread_id) { + thread.pending_auth = None; + } + } + + let ext_mgr = match self.deps.extension_manager.as_ref() { + Some(mgr) => mgr, + None => return Ok(Some("Extension manager not available.".to_string())), + }; + + match ext_mgr.auth(&pending.extension_name, Some(token)).await { + Ok(result) if result.status == "authenticated" => { + tracing::info!( + "Extension '{}' authenticated via auth mode", + pending.extension_name + ); + + // Notify via channel status so the response doesn't echo the token + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::Status("Authenticated".into()), + &message.metadata, + ) + .await; + + Ok(Some(format!( + "{} authenticated successfully.", + pending.extension_name + ))) + } + Ok(result) => { + // Unexpected state, re-enter auth mode + { + let mut sess = session.lock().await; + if let Some(thread) = sess.threads.get_mut(&thread_id) { + thread.enter_auth_mode(pending.extension_name.clone()); + } + } + let msg = result + .instructions + .unwrap_or_else(|| "Invalid token. Please try again.".to_string()); + Ok(Some(msg)) + } + Err(e) => Ok(Some(format!( + "Authentication failed for {}: {}", + pending.extension_name, e + ))), + } + } + async fn process_new_thread( &self, message: &IncomingMessage, @@ -1737,10 +1872,96 @@ impl Agent { } } -fn truncate(s: &str, max_len: usize) -> String { - if s.len() <= max_len { - s.to_string() - } else { - format!("{}...", &s[..max_len]) +/// Check if a tool_auth result indicates the extension is awaiting a token. +/// +/// Returns `Some((extension_name, instructions))` if the tool result contains +/// `awaiting_token: true`, meaning the thread should enter auth mode. +fn detect_auth_awaiting( + tool_name: &str, + result: &Result, +) -> Option<(String, String)> { + if tool_name != "tool_auth" { + return None; + } + let output = result.as_ref().ok()?; + let parsed: serde_json::Value = serde_json::from_str(output).ok()?; + if parsed.get("awaiting_token") != Some(&serde_json::Value::Bool(true)) { + return None; + } + let name = parsed.get("name")?.as_str()?.to_string(); + let instructions = parsed + .get("instructions") + .and_then(|v| v.as_str()) + .unwrap_or("Please provide your API token/key.") + .to_string(); + Some((name, instructions)) +} + +#[cfg(test)] +mod tests { + use crate::error::Error; + + use super::detect_auth_awaiting; + + #[test] + fn test_detect_auth_awaiting_positive() { + let result: Result = Ok(serde_json::json!({ + "name": "telegram", + "kind": "WasmTool", + "awaiting_token": true, + "status": "awaiting_token", + "instructions": "Please provide your Telegram Bot API token." + }) + .to_string()); + + let detected = detect_auth_awaiting("tool_auth", &result); + assert!(detected.is_some()); + let (name, instructions) = detected.unwrap(); + assert_eq!(name, "telegram"); + assert!(instructions.contains("Telegram Bot API")); + } + + #[test] + fn test_detect_auth_awaiting_not_awaiting() { + let result: Result = Ok(serde_json::json!({ + "name": "telegram", + "kind": "WasmTool", + "awaiting_token": false, + "status": "authenticated" + }) + .to_string()); + + assert!(detect_auth_awaiting("tool_auth", &result).is_none()); + } + + #[test] + fn test_detect_auth_awaiting_wrong_tool() { + let result: Result = Ok(serde_json::json!({ + "name": "telegram", + "awaiting_token": true, + }) + .to_string()); + + assert!(detect_auth_awaiting("tool_list", &result).is_none()); + } + + #[test] + fn test_detect_auth_awaiting_error_result() { + let result: Result = + Err(crate::error::ToolError::NotFound { name: "x".into() }.into()); + assert!(detect_auth_awaiting("tool_auth", &result).is_none()); + } + + #[test] + fn test_detect_auth_awaiting_default_instructions() { + let result: Result = Ok(serde_json::json!({ + "name": "custom_tool", + "awaiting_token": true, + "status": "awaiting_token" + }) + .to_string()); + + let (_, instructions) = detect_auth_awaiting("tool_auth", &result).unwrap(); + assert_eq!(instructions, "Please provide your API token/key."); } } diff --git a/src/agent/mod.rs b/src/agent/mod.rs index 2467eb24..7c3d7685 100644 --- a/src/agent/mod.rs +++ b/src/agent/mod.rs @@ -30,7 +30,7 @@ pub use heartbeat::{HeartbeatConfig, HeartbeatResult, HeartbeatRunner, spawn_hea pub use router::{MessageIntent, Router}; pub use scheduler::Scheduler; pub use self_repair::{BrokenTool, RepairResult, RepairTask, SelfRepair, StuckJob}; -pub use session::{PendingApproval, Session, Thread, ThreadState, Turn, TurnState}; +pub use session::{PendingApproval, PendingAuth, Session, Thread, ThreadState, Turn, TurnState}; pub use session_manager::SessionManager; pub use submission::{Submission, SubmissionParser, SubmissionResult}; pub use task::{Task, TaskContext, TaskHandler, TaskOutput, TaskStatus}; diff --git a/src/agent/session.rs b/src/agent/session.rs index 3cf9faf0..a40d3fbc 100644 --- a/src/agent/session.rs +++ b/src/agent/session.rs @@ -121,6 +121,18 @@ pub enum ThreadState { Interrupted, } +/// Pending auth token request. +/// +/// When `tool_auth` returns `awaiting_token`, the thread enters auth mode. +/// The next user message is intercepted before entering the normal pipeline +/// (no logging, no turn creation, no history) and routed directly to the +/// credential store. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PendingAuth { + /// Extension name to authenticate. + pub extension_name: String, +} + /// Pending tool approval request stored on a thread. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct PendingApproval { @@ -158,6 +170,9 @@ pub struct Thread { /// Pending approval request (when state is AwaitingApproval). #[serde(default)] pub pending_approval: Option, + /// Pending auth token request (thread is in auth mode). + #[serde(default)] + pub pending_auth: Option, } impl Thread { @@ -173,6 +188,7 @@ impl Thread { updated_at: now, metadata: serde_json::Value::Null, pending_approval: None, + pending_auth: None, } } @@ -238,6 +254,18 @@ impl Thread { self.updated_at = Utc::now(); } + /// Enter auth mode: next user message will be routed directly to + /// the credential store, bypassing the normal pipeline entirely. + pub fn enter_auth_mode(&mut self, extension_name: String) { + self.pending_auth = Some(PendingAuth { extension_name }); + self.updated_at = Utc::now(); + } + + /// Take the pending auth (clearing auth mode). + pub fn take_pending_auth(&mut self) -> Option { + self.pending_auth.take() + } + /// Interrupt the current turn. pub fn interrupt(&mut self) { if let Some(turn) = self.turns.last_mut() { @@ -511,4 +539,58 @@ mod tests { assert_eq!(thread.turns[1].user_input, "How are you?"); assert!(thread.turns[1].response.is_none()); } + + #[test] + fn test_enter_auth_mode() { + let mut thread = Thread::new(Uuid::new_v4()); + assert!(thread.pending_auth.is_none()); + + thread.enter_auth_mode("telegram".to_string()); + assert!(thread.pending_auth.is_some()); + assert_eq!( + thread.pending_auth.as_ref().unwrap().extension_name, + "telegram" + ); + } + + #[test] + fn test_take_pending_auth() { + let mut thread = Thread::new(Uuid::new_v4()); + thread.enter_auth_mode("notion".to_string()); + + let pending = thread.take_pending_auth(); + assert!(pending.is_some()); + assert_eq!(pending.unwrap().extension_name, "notion"); + + // Should be cleared after take + assert!(thread.pending_auth.is_none()); + assert!(thread.take_pending_auth().is_none()); + } + + #[test] + fn test_pending_auth_serialization() { + let mut thread = Thread::new(Uuid::new_v4()); + thread.enter_auth_mode("openai".to_string()); + + let json = serde_json::to_string(&thread).expect("should serialize"); + assert!(json.contains("pending_auth")); + assert!(json.contains("openai")); + + let restored: Thread = serde_json::from_str(&json).expect("should deserialize"); + assert!(restored.pending_auth.is_some()); + assert_eq!(restored.pending_auth.unwrap().extension_name, "openai"); + } + + #[test] + fn test_pending_auth_default_none() { + // Deserialization of old data without pending_auth should default to None + let mut thread = Thread::new(Uuid::new_v4()); + thread.pending_auth = None; + let json = serde_json::to_string(&thread).expect("serialize"); + + // Remove the pending_auth field to simulate old data + let json = json.replace(",\"pending_auth\":null", ""); + let restored: Thread = serde_json::from_str(&json).expect("should deserialize"); + assert!(restored.pending_auth.is_none()); + } } diff --git a/src/channels/wasm/loader.rs b/src/channels/wasm/loader.rs index 6ab145f2..00dfcd80 100644 --- a/src/channels/wasm/loader.rs +++ b/src/channels/wasm/loader.rs @@ -151,17 +151,18 @@ impl WasmChannelLoader { } let mut results = LoadResults::default(); + + // Collect all .wasm entries first, then load in parallel + let mut channel_entries = Vec::new(); let mut entries = fs::read_dir(dir).await?; while let Some(entry) = entries.next_entry().await? { let path = entry.path(); - // Only process .wasm files if path.extension().and_then(|e| e.to_str()) != Some("wasm") { continue; } - // Extract channel name from filename let name = match path.file_stem().and_then(|s| s.to_str()) { Some(n) => n.to_string(), None => { @@ -173,15 +174,20 @@ impl WasmChannelLoader { } }; - // Look for sidecar capabilities file let cap_path = path.with_extension("capabilities.json"); - let cap_path_option = if cap_path.exists() { - Some(cap_path.as_path()) - } else { - None - }; + let has_cap = cap_path.exists(); + channel_entries.push((name, path, if has_cap { Some(cap_path) } else { None })); + } - match self.load_from_files(&name, &path, cap_path_option).await { + // Load all channels in parallel (file I/O + WASM compilation) + let load_futures = channel_entries + .iter() + .map(|(name, path, cap_path)| self.load_from_files(name, path, cap_path.as_deref())); + + let load_results = futures::future::join_all(load_futures).await; + + for ((name, path, _), result) in channel_entries.into_iter().zip(load_results) { + match result { Ok(loaded) => { results.loaded.push(loaded); } diff --git a/src/extensions/manager.rs b/src/extensions/manager.rs index 3b2fffe0..89c30426 100644 --- a/src/extensions/manager.rs +++ b/src/extensions/manager.rs @@ -460,12 +460,35 @@ impl ExtensionManager { async fn auth_mcp( &self, name: &str, - _token: Option<&str>, + token: Option<&str>, ) -> Result { let server = get_mcp_server(name) .await .map_err(|e| ExtensionError::NotInstalled(e.to_string()))?; + // If a token was provided directly, store it and we're done. + if let Some(token_value) = token { + let secret_name = server.token_secret_name(); + let params = + CreateSecretParams::new(&secret_name, token_value).with_provider(name.to_string()); + self.secrets + .create(&self.user_id, params) + .await + .map_err(|e| ExtensionError::AuthFailed(e.to_string()))?; + + tracing::info!("MCP server '{}' authenticated via manual token", name); + return Ok(AuthResult { + name: name.to_string(), + kind: ExtensionKind::McpServer, + auth_url: None, + callback_type: None, + instructions: None, + setup_url: None, + awaiting_token: false, + status: "authenticated".to_string(), + }); + } + // Check if already authenticated if is_authenticated(&server, &self.secrets, &self.user_id).await { return Ok(AuthResult { @@ -483,7 +506,7 @@ impl ExtensionManager { // Run the full OAuth flow (opens browser, waits for callback) match authorize_mcp_server(&server, &self.secrets, &self.user_id).await { Ok(_token) => { - tracing::info!("MCP server '{}' authenticated successfully", name); + tracing::info!("MCP server '{}' authenticated via OAuth", name); Ok(AuthResult { name: name.to_string(), kind: ExtensionKind::McpServer, @@ -496,10 +519,45 @@ impl ExtensionManager { }) } Err(crate::tools::mcp::auth::AuthError::NotSupported) => { - // Server doesn't support OAuth at all, try to build a non-interactive auth URL - self.auth_mcp_build_url(name, &server).await + // Server doesn't support OAuth, try building a URL first + match self.auth_mcp_build_url(name, &server).await { + Ok(result) => Ok(result), + Err(_) => { + // No OAuth, no DCR: fall back to manual token entry + Ok(AuthResult { + name: name.to_string(), + kind: ExtensionKind::McpServer, + auth_url: None, + callback_type: None, + instructions: Some(format!( + "Server '{}' does not support OAuth. \ + Please provide an API token/key for this server.", + name + )), + setup_url: None, + awaiting_token: true, + status: "awaiting_token".to_string(), + }) + } + } + } + Err(e) => { + // OAuth failed for some other reason, fall back to manual token + Ok(AuthResult { + name: name.to_string(), + kind: ExtensionKind::McpServer, + auth_url: None, + callback_type: None, + instructions: Some(format!( + "OAuth failed for '{}': {}. \ + Please provide an API token/key manually.", + name, e + )), + setup_url: None, + awaiting_token: true, + status: "awaiting_token".to_string(), + }) } - Err(e) => Err(ExtensionError::AuthFailed(e.to_string())), } } diff --git a/src/main.rs b/src/main.rs index 491f92aa..913e9717 100644 --- a/src/main.rs +++ b/src/main.rs @@ -301,44 +301,6 @@ async fn main() -> anyhow::Result<()> { tracing::info!("Builder mode enabled"); } - // Load installed WASM tools (save runtime handle for extension manager) - let wasm_tool_runtime: Option> = 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); - } - } - - Some(runtime) - } - Err(e) => { - tracing::warn!("Failed to initialize WASM runtime: {}", e); - None - } - } - } else { - None - }; - // Create secrets store if master key is configured (needed for MCP auth and WASM channels) let secrets_store: Option> = if let (Some(store), Some(master_key)) = (&store, config.secrets.master_key()) { @@ -356,91 +318,146 @@ async fn main() -> anyhow::Result<()> { None }; - // Load configured MCP servers let mcp_session_manager = Arc::new(McpSessionManager::new()); - if let Some(ref secrets) = secrets_store { - match load_mcp_servers().await { - Ok(servers) => { - let enabled_count = servers.servers.iter().filter(|s| s.enabled).count(); - if enabled_count > 0 { - tracing::info!("Loading {} configured MCP server(s)...", enabled_count); + + // Create WASM tool runtime (sync, just builds the wasmtime engine) + let wasm_tool_runtime: Option> = + if config.wasm.enabled && config.wasm.tools_dir.exists() { + match WasmToolRuntime::new(config.wasm.to_runtime_config()) { + Ok(runtime) => Some(Arc::new(runtime)), + Err(e) => { + tracing::warn!("Failed to initialize WASM runtime: {}", e); + None } + } + } else { + None + }; - for server in servers.enabled_servers() { - tracing::debug!( - "Checking authentication for MCP server '{}'...", - server.name - ); - // Check for stored tokens (from either pre-configured OAuth or DCR) - let has_tokens = is_authenticated(server, secrets, "default").await; - tracing::debug!("MCP server '{}' has_tokens={}", server.name, has_tokens); + // Load WASM tools and MCP servers concurrently. + // Both register into the shared ToolRegistry (RwLock-based) so concurrent writes are safe. + let wasm_tools_future = async { + if let Some(ref runtime) = wasm_tool_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); + } + } + } + }; - let client = if has_tokens || server.requires_auth() { - // Use authenticated client if we have tokens or OAuth is configured - McpClient::new_authenticated( - server.clone(), - Arc::clone(&mcp_session_manager), - Arc::clone(secrets), - "default", - ) - } else { - // No tokens and no OAuth - try unauthenticated - McpClient::new_with_name(&server.name, &server.url) - }; + let mcp_servers_future = async { + if let Some(ref secrets) = secrets_store { + match load_mcp_servers().await { + Ok(servers) => { + let enabled: Vec<_> = servers.enabled_servers().cloned().collect(); + if !enabled.is_empty() { + tracing::info!("Loading {} configured MCP server(s)...", enabled.len()); + } - tracing::debug!("Fetching tools from MCP server '{}'...", server.name); - match client.list_tools().await { - Ok(mcp_tools) => { + let mut join_set = tokio::task::JoinSet::new(); + for server in enabled { + let mcp_sm = Arc::clone(&mcp_session_manager); + let secrets = Arc::clone(secrets); + let tools = Arc::clone(&tools); + + join_set.spawn(async move { + let server_name = server.name.clone(); tracing::debug!( - "Got {} tools from MCP server '{}'", - mcp_tools.len(), - server.name + "Checking authentication for MCP server '{}'...", + server_name ); - match client.create_tools().await { - Ok(tool_impls) => { - for tool in tool_impls { - tools.register(tool).await; - } - tracing::info!( - "Loaded {} tools from MCP server '{}'", - mcp_tools.len(), - server.name + let has_tokens = is_authenticated(&server, &secrets, "default").await; + tracing::debug!( + "MCP server '{}' has_tokens={}", + server_name, + has_tokens + ); + + let client = if has_tokens || server.requires_auth() { + McpClient::new_authenticated(server, mcp_sm, secrets, "default") + } else { + McpClient::new_with_name(&server_name, &server.url) + }; + + tracing::debug!("Fetching tools from MCP server '{}'...", server_name); + match client.list_tools().await { + Ok(mcp_tools) => { + let tool_count = mcp_tools.len(); + tracing::debug!( + "Got {} tools from MCP server '{}'", + tool_count, + server_name ); + match client.create_tools().await { + Ok(tool_impls) => { + for tool in tool_impls { + tools.register(tool).await; + } + tracing::info!( + "Loaded {} tools from MCP server '{}'", + tool_count, + server_name + ); + } + Err(e) => { + tracing::warn!( + "Failed to create tools from MCP server '{}': {}", + server_name, + e + ); + } + } } Err(e) => { - tracing::warn!( - "Failed to create tools from MCP server '{}': {}", - server.name, - e - ); + let err_str = e.to_string(); + if err_str.contains("401") || err_str.contains("authentication") + { + tracing::warn!( + "MCP server '{}' requires authentication. \ + Run: ironclaw mcp auth {}", + server_name, + server_name + ); + } else { + tracing::warn!( + "Failed to connect to MCP server '{}': {}", + server_name, + e + ); + } } } - } - Err(e) => { - // Check if it's an auth error - let err_str = e.to_string(); - if err_str.contains("401") || err_str.contains("authentication") { - tracing::warn!( - "MCP server '{}' requires authentication. Run: ironclaw mcp auth {}", - server.name, - server.name - ); - } else { - tracing::warn!( - "Failed to connect to MCP server '{}': {}", - server.name, - e - ); - } + }); + } + + while let Some(result) = join_set.join_next().await { + if let Err(e) = result { + tracing::warn!("MCP server loading task panicked: {}", e); } } } - } - Err(e) => { - tracing::debug!("No MCP servers configured ({})", e); + Err(e) => { + tracing::debug!("No MCP servers configured ({})", e); + } } } - } + }; + + tokio::join!(wasm_tools_future, mcp_servers_future); // Create extension manager for in-chat discovery/install/auth/activate let extension_manager = if let Some(ref secrets) = secrets_store { @@ -728,6 +745,7 @@ async fn main() -> anyhow::Result<()> { safety, tools, workspace, + extension_manager, }; let agent = Agent::new( config.agent.clone(), diff --git a/src/tools/wasm/loader.rs b/src/tools/wasm/loader.rs index 0029b4c3..ade7661c 100644 --- a/src/tools/wasm/loader.rs +++ b/src/tools/wasm/loader.rs @@ -165,17 +165,18 @@ impl WasmToolLoader { } let mut results = LoadResults::default(); + + // Collect all .wasm entries first, then load in parallel + let mut tool_entries = Vec::new(); let mut entries = fs::read_dir(dir).await?; while let Some(entry) = entries.next_entry().await? { let path = entry.path(); - // Only process .wasm files if path.extension().and_then(|e| e.to_str()) != Some("wasm") { continue; } - // Extract tool name from filename let name = match path.file_stem().and_then(|s| s.to_str()) { Some(n) => n.to_string(), None => { @@ -187,15 +188,20 @@ impl WasmToolLoader { } }; - // Look for sidecar capabilities file let cap_path = path.with_extension("capabilities.json"); - let cap_path_option = if cap_path.exists() { - Some(cap_path.as_path()) - } else { - None - }; + let has_cap = cap_path.exists(); + tool_entries.push((name, path, if has_cap { Some(cap_path) } else { None })); + } - match self.load_from_files(&name, &path, cap_path_option).await { + // Load all tools in parallel (file I/O + WASM compilation + registration) + let load_futures = tool_entries + .iter() + .map(|(name, path, cap_path)| self.load_from_files(name, path, cap_path.as_deref())); + + let load_results = futures::future::join_all(load_futures).await; + + for ((name, path, _), result) in tool_entries.into_iter().zip(load_results) { + match result { Ok(()) => { results.loaded.push(name); }