Add auth mode, fix MCP token handling, and parallelize startup loading

Auth mode: when a tool requires an API key, the thread enters a special
mode where the next user message is routed directly to the credential
store, bypassing logs, turns, history, and compaction entirely. This
prevents tokens from leaking into debug output or persistent storage.

Fix MCP auth: auth_mcp now actually uses the token parameter (was
ignored as _token) and falls back to manual token entry when OAuth
and DCR are both unsupported.

Parallel loading: WASM tools, WASM channels, and MCP servers now load
concurrently at startup. Within each loader, individual items also
load in parallel (join_all for WASM, JoinSet for MCP servers).

Co-Authored-By: Claude Opus 4.6 <[email protected]>
This commit is contained in:
Illia Polosukhin
2026-02-06 13:35:17 -08:00
co-authored by Claude Opus 4.6
parent 9b729795fb
commit 2cdd04a359
7 changed files with 533 additions and 142 deletions
+233 -12
View File
@@ -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<SafetyLayer>,
pub tools: Arc<ToolRegistry>,
pub workspace: Option<Arc<Workspace>>,
pub extension_manager: Option<Arc<ExtensionManager>>,
}
/// The main agent that coordinates all components.
@@ -374,13 +376,6 @@ impl Agent {
}
async fn handle_message(&self, message: &IncomingMessage) -> Result<Option<String>, 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<Mutex<Session>>,
thread_id: Uuid,
) -> Result<Option<String>, 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<String, Error>,
) -> 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<String, Error> = 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<String, Error> = 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<String, Error> = 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<String, Error> =
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<String, Error> = 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.");
}
}
+1 -1
View File
@@ -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};
+82
View File
@@ -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<PendingApproval>,
/// Pending auth token request (thread is in auth mode).
#[serde(default)]
pub pending_auth: Option<PendingAuth>,
}
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<PendingAuth> {
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());
}
}
+15 -9
View File
@@ -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);
}
+63 -5
View File
@@ -460,12 +460,35 @@ impl ExtensionManager {
async fn auth_mcp(
&self,
name: &str,
_token: Option<&str>,
token: Option<&str>,
) -> Result<AuthResult, ExtensionError> {
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())),
}
}
+124 -106
View File
@@ -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<Arc<WasmToolRuntime>> = 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<Arc<dyn SecretsStore + Send + Sync>> =
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<Arc<WasmToolRuntime>> =
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(),
+15 -9
View File
@@ -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);
}