From ae3c86a7ea61f77e7fb00c119e0bc9042f25c97a Mon Sep 17 00:00:00 2001 From: Illia Polosukhin Date: Thu, 5 Feb 2026 20:11:29 -0800 Subject: [PATCH] Add in-chat extension discovery, auth, and activation system Introduces a unified extension abstraction over MCP servers and WASM tools with six agent-callable tools (tool_search, tool_install, tool_auth, tool_activate, tool_list, tool_remove) so users can add capabilities conversationally without CLI commands. Includes built-in registry of 11 MCP servers, online discovery via URL probing and GitHub search, OAuth 2.1 flows for MCP servers, and manual token auth for WASM tools. Co-Authored-By: Claude Opus 4.6 --- src/channels/wasm/router.rs | 94 ++- src/extensions/discovery.rs | 326 ++++++++++ src/extensions/manager.rs | 892 +++++++++++++++++++++++++++ src/extensions/mod.rs | 224 +++++++ src/extensions/registry.rs | 545 ++++++++++++++++ src/lib.rs | 1 + src/main.rs | 42 +- src/tools/builtin/extension_tools.rs | 523 ++++++++++++++++ src/tools/builtin/mod.rs | 4 + src/tools/mcp/auth.rs | 12 +- src/tools/registry.rs | 17 +- 11 files changed, 2664 insertions(+), 16 deletions(-) create mode 100644 src/extensions/discovery.rs create mode 100644 src/extensions/manager.rs create mode 100644 src/extensions/mod.rs create mode 100644 src/extensions/registry.rs create mode 100644 src/tools/builtin/extension_tools.rs diff --git a/src/channels/wasm/router.rs b/src/channels/wasm/router.rs index 892f9e4a..44f81f99 100644 --- a/src/channels/wasm/router.rs +++ b/src/channels/wasm/router.rs @@ -173,11 +173,23 @@ impl Default for WasmChannelRouter { #[derive(Clone)] pub struct RouterState { router: Arc, + extension_manager: Option>, } impl RouterState { pub fn new(router: Arc) -> Self { - Self { router } + Self { + router, + extension_manager: None, + } + } + + pub fn with_extension_manager( + mut self, + manager: Arc, + ) -> Self { + self.extension_manager = Some(manager); + self } } @@ -384,14 +396,73 @@ async fn webhook_handler( } } +/// OAuth callback handler for extension authentication. +/// +/// Handles OAuth redirect callbacks at /oauth/callback?code=xxx&state=yyy. +/// This is used when authenticating MCP servers or WASM tool OAuth flows +/// via a tunnel URL (remote callback). +#[allow(dead_code)] +async fn oauth_callback_handler( + State(_state): State, + Query(params): Query>, +) -> impl IntoResponse { + let code = params.get("code").cloned().unwrap_or_default(); + let _state = params.get("state").cloned().unwrap_or_default(); + + if code.is_empty() { + let error = params + .get("error") + .cloned() + .unwrap_or_else(|| "unknown".to_string()); + return ( + StatusCode::BAD_REQUEST, + axum::response::Html(format!( + "\ +
\ +

Authorization Failed

\ +

Error: {}

\ +
", + error + )), + ); + } + + // TODO: In a future iteration, use the state nonce to look up the pending auth + // and complete the token exchange. For now, the OAuth flow uses local callbacks + // via authorize_mcp_server() which handles the full flow synchronously. + + ( + StatusCode::OK, + axum::response::Html( + "\ +
\ +

Connected!

\ +

You can close this window and return to IronClaw.

\ +
" + .to_string(), + ), + ) +} + /// Create an Axum router for WASM channel webhooks. /// /// This router can be merged with the existing HTTP channel router. -pub fn create_wasm_channel_router(router: Arc) -> Router { - let state = RouterState::new(router); +pub fn create_wasm_channel_router( + router: Arc, + extension_manager: Option>, +) -> Router { + let mut state = RouterState::new(router); + if let Some(manager) = extension_manager { + state = state.with_extension_manager(manager); + } Router::new() .route("/wasm-channels/health", get(health_handler)) + .route("/oauth/callback", get(oauth_callback_handler)) // Catch-all for webhook paths .route("/webhook/{*path}", get(webhook_handler)) .route("/webhook/{*path}", post(webhook_handler)) @@ -401,12 +472,25 @@ pub fn create_wasm_channel_router(router: Arc) -> Router { /// HTTP server for WASM channel webhooks. pub struct WasmChannelServer { router: Arc, + extension_manager: Option>, } impl WasmChannelServer { /// Create a new server. pub fn new(router: Arc) -> Self { - Self { router } + Self { + router, + extension_manager: None, + } + } + + /// Set the extension manager for OAuth callback handling. + pub fn with_extension_manager( + mut self, + manager: Arc, + ) -> Self { + self.extension_manager = Some(manager); + self } /// Start the HTTP server. @@ -416,7 +500,7 @@ impl WasmChannelServer { &self, addr: SocketAddr, ) -> Result, std::io::Error> { - let app = create_wasm_channel_router(self.router.clone()); + let app = create_wasm_channel_router(self.router.clone(), self.extension_manager.clone()); let listener = tokio::net::TcpListener::bind(addr).await?; diff --git a/src/extensions/discovery.rs b/src/extensions/discovery.rs new file mode 100644 index 00000000..b40815e7 --- /dev/null +++ b/src/extensions/discovery.rs @@ -0,0 +1,326 @@ +//! Online extension discovery for finding MCP servers not in the built-in registry. +//! +//! Multi-tier search strategy: +//! 1. Probe well-known URL patterns (mcp.{service}.com, {service}.com/mcp) +//! 2. Search GitHub for MCP server repositories +//! 3. Validate discovered URLs via .well-known/oauth-protected-resource +//! +//! All sources run concurrently with per-source timeouts. + +use std::time::Duration; + +use serde::Deserialize; + +use crate::extensions::{AuthHint, ExtensionKind, ExtensionSource, RegistryEntry}; + +/// Handles online discovery of MCP servers. +pub struct OnlineDiscovery { + http_client: reqwest::Client, +} + +impl OnlineDiscovery { + pub fn new() -> Self { + let http_client = reqwest::Client::builder() + .timeout(Duration::from_secs(10)) + .user_agent("IronClaw/1.0") + .build() + .unwrap_or_else(|_| reqwest::Client::new()); + + Self { http_client } + } + + /// Run the full discovery pipeline for a query. + /// + /// Searches multiple sources concurrently, deduplicates, validates, + /// and returns only confirmed MCP servers. + pub async fn discover(&self, query: &str) -> Vec { + let query_clean = query.trim().to_lowercase(); + if query_clean.is_empty() { + return Vec::new(); + } + + // Run all discovery sources concurrently + let (patterns, github) = tokio::join!( + self.probe_common_patterns(&query_clean), + with_timeout(self.search_github(&query_clean), Duration::from_secs(8)), + ); + + // Collect and deduplicate by URL + let mut seen_urls = std::collections::HashSet::new(); + let mut candidates: Vec = Vec::new(); + + for entry in patterns { + let url = extract_url(&entry.source); + if seen_urls.insert(url) { + candidates.push(entry); + } + } + + for entry in github.unwrap_or_default() { + let url = extract_url(&entry.source); + if seen_urls.insert(url) { + candidates.push(entry); + } + } + + candidates + } + + /// Probe common URL patterns for MCP servers. + /// + /// Tries patterns like: + /// - https://mcp.{query}.com + /// - https://mcp.{query}.app + /// - https://{query}.com/mcp + pub async fn probe_common_patterns(&self, query: &str) -> Vec { + // Extract a clean service name (no spaces, lowercase) + let service = query + .split_whitespace() + .next() + .unwrap_or(query) + .replace('-', ""); + + let patterns = vec![ + format!("https://mcp.{}.com", service), + format!("https://mcp.{}.app", service), + format!("https://mcp.{}.dev", service), + format!("https://{}.com/mcp", service), + ]; + + let mut results = Vec::new(); + let futures: Vec<_> = patterns + .into_iter() + .map(|url| { + let client = self.http_client.clone(); + let query_owned = query.to_string(); + async move { + if validate_mcp_url_with_client(&client, &url).await { + Some(RegistryEntry { + name: query_owned.replace(' ', "-"), + display_name: titlecase(&query_owned), + kind: ExtensionKind::McpServer, + description: format!("MCP server discovered at {}", url), + keywords: vec![], + source: ExtensionSource::McpUrl { + url: url.to_string(), + }, + auth_hint: AuthHint::Dcr, + }) + } else { + None + } + } + }) + .collect(); + + let probe_results = futures::future::join_all(futures).await; + for result in probe_results.into_iter().flatten() { + results.push(result); + } + + results + } + + /// Search GitHub for MCP server repositories. + /// + /// Uses the GitHub search API (no auth needed for low-rate public queries). + pub async fn search_github(&self, query: &str) -> Vec { + let search_url = format!( + "https://api.github.com/search/repositories?q={}+topic:mcp-server&per_page=5&sort=stars", + urlencoding::encode(query) + ); + + let response = match self.http_client.get(&search_url).send().await { + Ok(r) => r, + Err(e) => { + tracing::debug!("GitHub search failed: {}", e); + return Vec::new(); + } + }; + + if !response.status().is_success() { + tracing::debug!("GitHub search returned {}", response.status()); + return Vec::new(); + } + + let body: GitHubSearchResponse = match response.json().await { + Ok(b) => b, + Err(e) => { + tracing::debug!("Failed to parse GitHub search response: {}", e); + return Vec::new(); + } + }; + + body.items + .into_iter() + .filter_map(|item| { + // Only include repos that look like MCP servers + let has_mcp_topic = item + .topics + .iter() + .any(|t| t.contains("mcp") || t.contains("model-context-protocol")); + if !has_mcp_topic { + return None; + } + + // Try to extract a homepage URL (which might be the MCP endpoint) + let url = item.homepage.filter(|h| !h.is_empty()).unwrap_or_else(|| { + // Fall back to repo URL as a reference + item.html_url.clone() + }); + + Some(RegistryEntry { + name: item.name.clone(), + display_name: titlecase(&item.name.replace('-', " ")), + kind: ExtensionKind::McpServer, + description: item + .description + .unwrap_or_else(|| format!("MCP server from GitHub: {}", item.full_name)), + keywords: item.topics, + source: ExtensionSource::Discovered { url }, + auth_hint: AuthHint::Dcr, + }) + }) + .collect() + } + + /// Validate a URL is a real MCP server. + pub async fn validate_mcp_url(&self, url: &str) -> bool { + validate_mcp_url_with_client(&self.http_client, url).await + } +} + +impl Default for OnlineDiscovery { + fn default() -> Self { + Self::new() + } +} + +/// Validate that a URL is a real MCP server by checking .well-known endpoints. +/// +/// Tries: +/// 1. GET {origin}/.well-known/oauth-protected-resource -> 200 with JSON = confirmed +/// 2. Fallback: HEAD/GET the URL itself to check if it's alive +async fn validate_mcp_url_with_client(client: &reqwest::Client, url: &str) -> bool { + let parsed = match reqwest::Url::parse(url) { + Ok(u) => u, + Err(_) => return false, + }; + let origin = parsed.origin().ascii_serialization(); + + // Check .well-known/oauth-protected-resource + let well_known_url = format!("{}/.well-known/oauth-protected-resource", origin); + match client.get(&well_known_url).send().await { + Ok(resp) if resp.status().is_success() => { + // Try to parse as JSON to confirm it's a real MCP endpoint + if let Ok(text) = resp.text().await { + return serde_json::from_str::(&text).is_ok(); + } + } + _ => {} + } + + // Fallback: try a HEAD request on the URL itself to check if it's alive + match client.head(url).send().await { + Ok(resp) => { + // Accept various status codes that indicate the server exists + let status = resp.status().as_u16(); + // 401/403 means it exists but needs auth, which is fine for MCP + matches!(status, 200..=299 | 401 | 403 | 405) + } + Err(_) => false, + } +} + +/// Run a future with a timeout, returning None if it times out. +async fn with_timeout( + future: impl std::future::Future, + duration: Duration, +) -> Option { + tokio::time::timeout(duration, future).await.ok() +} + +fn extract_url(source: &ExtensionSource) -> String { + match source { + ExtensionSource::McpUrl { url } => url.clone(), + ExtensionSource::Discovered { url } => url.clone(), + ExtensionSource::WasmDownload { wasm_url, .. } => wasm_url.clone(), + ExtensionSource::WasmBuildable { repo_url, .. } => repo_url.clone(), + } +} + +fn titlecase(s: &str) -> String { + s.split_whitespace() + .map(|word| { + let mut chars = word.chars(); + match chars.next() { + Some(c) => format!("{}{}", c.to_uppercase(), chars.as_str()), + None => String::new(), + } + }) + .collect::>() + .join(" ") +} + +#[derive(Debug, Deserialize)] +struct GitHubSearchResponse { + #[serde(default)] + items: Vec, +} + +#[derive(Debug, Deserialize)] +struct GitHubRepo { + name: String, + full_name: String, + html_url: String, + description: Option, + #[serde(default)] + homepage: Option, + #[serde(default)] + topics: Vec, +} + +#[cfg(test)] +mod tests { + use crate::extensions::ExtensionSource; + use crate::extensions::discovery::{ + OnlineDiscovery, extract_url, titlecase, validate_mcp_url_with_client, + }; + + #[test] + fn test_titlecase() { + assert_eq!(titlecase("google calendar"), "Google Calendar"); + assert_eq!(titlecase("notion"), "Notion"); + assert_eq!(titlecase(""), ""); + } + + #[test] + fn test_extract_url() { + let mcp = ExtensionSource::McpUrl { + url: "https://mcp.notion.com".to_string(), + }; + assert_eq!(extract_url(&mcp), "https://mcp.notion.com"); + + let discovered = ExtensionSource::Discovered { + url: "https://example.com".to_string(), + }; + assert_eq!(extract_url(&discovered), "https://example.com"); + } + + #[tokio::test] + async fn test_validate_invalid_url() { + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(3)) + .build() + .unwrap(); + + // Invalid URL should fail + assert!(!validate_mcp_url_with_client(&client, "not-a-url").await); + } + + #[test] + fn test_discovery_new() { + // Just make sure it constructs without panicking + let _discovery = OnlineDiscovery::new(); + } +} diff --git a/src/extensions/manager.rs b/src/extensions/manager.rs new file mode 100644 index 00000000..3b2fffe0 --- /dev/null +++ b/src/extensions/manager.rs @@ -0,0 +1,892 @@ +//! Central extension manager that dispatches operations by ExtensionKind. +//! +//! Holds references to MCP infrastructure, WASM tool runtime, secrets store, +//! and tool registry. All extension operations (search, install, auth, activate, +//! list, remove) flow through here. + +use std::collections::HashMap; +use std::path::PathBuf; +use std::sync::Arc; + +use tokio::sync::RwLock; + +use crate::extensions::discovery::OnlineDiscovery; +use crate::extensions::registry::ExtensionRegistry; +use crate::extensions::{ + ActivateResult, AuthResult, ExtensionError, ExtensionKind, ExtensionSource, InstallResult, + InstalledExtension, RegistryEntry, ResultSource, SearchResult, +}; +use crate::secrets::{CreateSecretParams, SecretsStore}; +use crate::tools::ToolRegistry; +use crate::tools::mcp::McpClient; +use crate::tools::mcp::auth::{ + PkceChallenge, authorize_mcp_server, build_authorization_url, discover_full_oauth_metadata, + find_available_port, is_authenticated, register_client, +}; +use crate::tools::mcp::config::{ + McpServerConfig, add_mcp_server, get_mcp_server, load_mcp_servers, remove_mcp_server, +}; +use crate::tools::mcp::session::McpSessionManager; +use crate::tools::wasm::{WasmToolLoader, WasmToolRuntime, discover_tools}; + +/// Pending OAuth authorization state. +struct PendingAuth { + _name: String, + _kind: ExtensionKind, + created_at: std::time::Instant, +} + +/// Central manager for extension lifecycle operations. +pub struct ExtensionManager { + registry: ExtensionRegistry, + discovery: OnlineDiscovery, + + // MCP infrastructure + mcp_session_manager: Arc, + /// Active MCP clients keyed by server name. + mcp_clients: RwLock>>, + + // WASM tool infrastructure + wasm_tool_runtime: Option>, + wasm_tools_dir: PathBuf, + wasm_channels_dir: PathBuf, + + // Shared + secrets: Arc, + tool_registry: Arc, + pending_auth: RwLock>, + /// Tunnel URL for remote OAuth callbacks (used in future iterations). + _tunnel_url: Option, + user_id: String, +} + +impl ExtensionManager { + #[allow(clippy::too_many_arguments)] + pub fn new( + mcp_session_manager: Arc, + secrets: Arc, + tool_registry: Arc, + wasm_tool_runtime: Option>, + wasm_tools_dir: PathBuf, + wasm_channels_dir: PathBuf, + tunnel_url: Option, + user_id: String, + ) -> Self { + Self { + registry: ExtensionRegistry::new(), + discovery: OnlineDiscovery::new(), + mcp_session_manager, + mcp_clients: RwLock::new(HashMap::new()), + wasm_tool_runtime, + wasm_tools_dir, + wasm_channels_dir, + secrets, + tool_registry, + pending_auth: RwLock::new(HashMap::new()), + _tunnel_url: tunnel_url, + user_id, + } + } + + /// Search for extensions. If `discover` is true, also searches online. + pub async fn search( + &self, + query: &str, + discover: bool, + ) -> Result, ExtensionError> { + let mut results = self.registry.search(query).await; + + if discover && results.is_empty() { + tracing::info!("No built-in results for '{}', searching online...", query); + let discovered = self.discovery.discover(query).await; + + if !discovered.is_empty() { + // Cache for future lookups + self.registry.cache_discovered(discovered.clone()).await; + + // Add to results + for entry in discovered { + results.push(SearchResult { + entry, + source: ResultSource::Discovered, + validated: true, + }); + } + } + } + + Ok(results) + } + + /// Install an extension by name (from registry) or by explicit URL. + pub async fn install( + &self, + name: &str, + url: Option<&str>, + kind_hint: Option, + ) -> Result { + // If we have a registry entry, use it + if let Some(entry) = self.registry.get(name).await { + return self.install_from_entry(&entry).await; + } + + // If a URL was provided, determine kind and install + if let Some(url) = url { + let kind = kind_hint.unwrap_or_else(|| infer_kind_from_url(url)); + return match kind { + ExtensionKind::McpServer => self.install_mcp_from_url(name, url).await, + ExtensionKind::WasmTool => self.install_wasm_tool_from_url(name, url).await, + ExtensionKind::WasmChannel => { + Err(ExtensionError::InstallFailed( + "WASM channel installation from URL not yet supported. \ + Place the .wasm and .capabilities.json files in ~/.ironclaw/channels/ and restart." + .to_string(), + )) + } + }; + } + + Err(ExtensionError::NotFound(format!( + "'{}' not found in registry. Try searching with discover:true or provide a URL.", + name + ))) + } + + /// Authenticate an installed extension. + pub async fn auth( + &self, + name: &str, + token: Option<&str>, + ) -> Result { + // Clean up expired pending auths + self.cleanup_expired_auths().await; + + // Determine what kind of extension this is + let kind = self.determine_installed_kind(name).await?; + + match kind { + ExtensionKind::McpServer => self.auth_mcp(name, token).await, + ExtensionKind::WasmTool => self.auth_wasm_tool(name, token).await, + ExtensionKind::WasmChannel => self.auth_wasm_tool(name, token).await, + } + } + + /// Activate an installed (and optionally authenticated) extension. + pub async fn activate(&self, name: &str) -> Result { + let kind = self.determine_installed_kind(name).await?; + + match kind { + ExtensionKind::McpServer => self.activate_mcp(name).await, + ExtensionKind::WasmTool => self.activate_wasm_tool(name).await, + ExtensionKind::WasmChannel => Err(ExtensionError::ChannelNeedsRestart), + } + } + + /// List all installed extensions with their status. + pub async fn list( + &self, + kind_filter: Option, + ) -> Result, ExtensionError> { + let mut extensions = Vec::new(); + + // List MCP servers + if kind_filter.is_none() || kind_filter == Some(ExtensionKind::McpServer) { + match load_mcp_servers().await { + Ok(servers) => { + for server in &servers.servers { + let authenticated = + is_authenticated(server, &self.secrets, &self.user_id).await; + let clients = self.mcp_clients.read().await; + let active = clients.contains_key(&server.name); + + // Get tool names if active + let tools = if active { + self.tool_registry + .list() + .await + .into_iter() + .filter(|t| t.starts_with(&format!("{}_", server.name))) + .collect() + } else { + Vec::new() + }; + + extensions.push(InstalledExtension { + name: server.name.clone(), + kind: ExtensionKind::McpServer, + description: server.description.clone(), + authenticated, + active, + tools, + }); + } + } + Err(e) => { + tracing::debug!("Failed to load MCP servers for listing: {}", e); + } + } + } + + // List WASM tools + if (kind_filter.is_none() || kind_filter == Some(ExtensionKind::WasmTool)) + && self.wasm_tools_dir.exists() + { + match discover_tools(&self.wasm_tools_dir).await { + Ok(tools) => { + for (name, _discovered) in tools { + let active = self.tool_registry.has(&name).await; + + extensions.push(InstalledExtension { + name: name.clone(), + kind: ExtensionKind::WasmTool, + description: None, + authenticated: true, // WASM tools don't always need auth + active, + tools: if active { vec![name] } else { Vec::new() }, + }); + } + } + Err(e) => { + tracing::debug!("Failed to discover WASM tools for listing: {}", e); + } + } + } + + // List WASM channels + if (kind_filter.is_none() || kind_filter == Some(ExtensionKind::WasmChannel)) + && self.wasm_channels_dir.exists() + { + match crate::channels::wasm::discover_channels(&self.wasm_channels_dir).await { + Ok(channels) => { + for (name, _discovered) in channels { + extensions.push(InstalledExtension { + name, + kind: ExtensionKind::WasmChannel, + description: None, + authenticated: true, + active: true, // If loaded at startup, they're active + tools: Vec::new(), + }); + } + } + Err(e) => { + tracing::debug!("Failed to discover WASM channels for listing: {}", e); + } + } + } + + Ok(extensions) + } + + /// Remove an installed extension. + pub async fn remove(&self, name: &str) -> Result { + let kind = self.determine_installed_kind(name).await?; + + match kind { + ExtensionKind::McpServer => { + // Unregister tools with this server's prefix + let tool_names: Vec = self + .tool_registry + .list() + .await + .into_iter() + .filter(|t| t.starts_with(&format!("{}_", name))) + .collect(); + + for tool_name in &tool_names { + self.tool_registry.unregister(tool_name).await; + } + + // Remove MCP client + self.mcp_clients.write().await.remove(name); + + // Remove from config + remove_mcp_server(name) + .await + .map_err(|e| ExtensionError::Config(e.to_string()))?; + + Ok(format!( + "Removed MCP server '{}' and {} tool(s)", + name, + tool_names.len() + )) + } + ExtensionKind::WasmTool => { + // Unregister from tool registry + self.tool_registry.unregister(name).await; + + // Delete files + let wasm_path = self.wasm_tools_dir.join(format!("{}.wasm", name)); + let cap_path = self + .wasm_tools_dir + .join(format!("{}.capabilities.json", name)); + + if wasm_path.exists() { + tokio::fs::remove_file(&wasm_path) + .await + .map_err(|e| ExtensionError::Other(e.to_string()))?; + } + if cap_path.exists() { + let _ = tokio::fs::remove_file(&cap_path).await; + } + + Ok(format!("Removed WASM tool '{}'", name)) + } + ExtensionKind::WasmChannel => Err(ExtensionError::Other( + "Channel removal requires restart. Delete the .wasm file from ~/.ironclaw/channels/ and restart." + .to_string(), + )), + } + } + + // ── Private helpers ────────────────────────────────────────────────── + + async fn install_from_entry( + &self, + entry: &RegistryEntry, + ) -> Result { + match entry.kind { + ExtensionKind::McpServer => { + let url = match &entry.source { + ExtensionSource::McpUrl { url } => url.clone(), + ExtensionSource::Discovered { url } => url.clone(), + _ => { + return Err(ExtensionError::InstallFailed( + "Registry entry for MCP server has no URL".to_string(), + )); + } + }; + self.install_mcp_from_url(&entry.name, &url).await + } + ExtensionKind::WasmTool => match &entry.source { + ExtensionSource::WasmDownload { wasm_url, .. } => { + self.install_wasm_tool_from_url(&entry.name, wasm_url).await + } + _ => Err(ExtensionError::InstallFailed( + "WASM tool entry has no download URL".to_string(), + )), + }, + ExtensionKind::WasmChannel => Err(ExtensionError::InstallFailed( + "WASM channel installation not yet supported via this flow".to_string(), + )), + } + } + + async fn install_mcp_from_url( + &self, + name: &str, + url: &str, + ) -> Result { + // Check if already installed + if get_mcp_server(name).await.is_ok() { + return Err(ExtensionError::AlreadyInstalled(name.to_string())); + } + + let config = McpServerConfig::new(name, url); + config + .validate() + .map_err(|e| ExtensionError::InvalidUrl(e.to_string()))?; + + add_mcp_server(config) + .await + .map_err(|e| ExtensionError::Config(e.to_string()))?; + + tracing::info!("Installed MCP server '{}' at {}", name, url); + + Ok(InstallResult { + name: name.to_string(), + kind: ExtensionKind::McpServer, + message: format!( + "MCP server '{}' installed. Run auth next to authenticate.", + name + ), + }) + } + + async fn install_wasm_tool_from_url( + &self, + name: &str, + url: &str, + ) -> Result { + // Download the WASM binary + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(60)) + .build() + .map_err(|e| ExtensionError::DownloadFailed(e.to_string()))?; + + let response = client + .get(url) + .send() + .await + .map_err(|e| ExtensionError::DownloadFailed(e.to_string()))?; + + if !response.status().is_success() { + return Err(ExtensionError::DownloadFailed(format!( + "HTTP {}", + response.status() + ))); + } + + let bytes = response + .bytes() + .await + .map_err(|e| ExtensionError::DownloadFailed(e.to_string()))?; + + // Ensure tools directory exists + tokio::fs::create_dir_all(&self.wasm_tools_dir) + .await + .map_err(|e| ExtensionError::InstallFailed(e.to_string()))?; + + // Write the WASM file + let wasm_path = self.wasm_tools_dir.join(format!("{}.wasm", name)); + tokio::fs::write(&wasm_path, &bytes) + .await + .map_err(|e| ExtensionError::InstallFailed(e.to_string()))?; + + tracing::info!( + "Installed WASM tool '{}' ({} bytes) to {}", + name, + bytes.len(), + wasm_path.display() + ); + + Ok(InstallResult { + name: name.to_string(), + kind: ExtensionKind::WasmTool, + message: format!("WASM tool '{}' installed. Run activate to load it.", name), + }) + } + + async fn auth_mcp( + &self, + name: &str, + _token: Option<&str>, + ) -> Result { + let server = get_mcp_server(name) + .await + .map_err(|e| ExtensionError::NotInstalled(e.to_string()))?; + + // Check if already authenticated + if is_authenticated(&server, &self.secrets, &self.user_id).await { + 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(), + }); + } + + // 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); + 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(), + }) + } + 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 + } + Err(e) => Err(ExtensionError::AuthFailed(e.to_string())), + } + } + + /// Build an auth URL for cases where non-interactive auth is needed + /// (e.g., running via Telegram where we can't open a browser). + async fn auth_mcp_build_url( + &self, + name: &str, + server: &McpServerConfig, + ) -> Result { + // Try to discover OAuth metadata and build a URL the user can open manually + let metadata = discover_full_oauth_metadata(&server.url) + .await + .map_err(|e| ExtensionError::AuthFailed(e.to_string()))?; + + // Try DCR if no client_id configured + let (client_id, redirect_uri) = if let Some(ref oauth) = server.oauth { + let port = find_available_port() + .await + .map_err(|e| ExtensionError::AuthFailed(e.to_string()))?; + let redirect = format!("http://localhost:{}/callback", port.1); + (oauth.client_id.clone(), redirect) + } else if let Some(ref reg_endpoint) = metadata.registration_endpoint { + let port = find_available_port() + .await + .map_err(|e| ExtensionError::AuthFailed(e.to_string()))?; + let redirect = format!("http://localhost:{}/callback", port.1); + + let registration = register_client(reg_endpoint, &redirect) + .await + .map_err(|e| ExtensionError::AuthFailed(e.to_string()))?; + + (registration.client_id, redirect) + } else { + return Err(ExtensionError::AuthFailed( + "Server doesn't support OAuth or Dynamic Client Registration".to_string(), + )); + }; + + let pkce = PkceChallenge::generate(); + let auth_url = build_authorization_url( + &metadata.authorization_endpoint, + &client_id, + &redirect_uri, + &metadata.scopes_supported, + Some(&pkce), + &std::collections::HashMap::new(), + ); + + // Store pending auth for later callback handling + self.pending_auth.write().await.insert( + name.to_string(), + PendingAuth { + _name: name.to_string(), + _kind: ExtensionKind::McpServer, + created_at: std::time::Instant::now(), + }, + ); + + Ok(AuthResult { + name: name.to_string(), + kind: ExtensionKind::McpServer, + auth_url: Some(auth_url), + callback_type: Some("local".to_string()), + instructions: None, + setup_url: None, + awaiting_token: false, + status: "awaiting_authorization".to_string(), + }) + } + + async fn auth_wasm_tool( + &self, + name: &str, + token: Option<&str>, + ) -> Result { + // Read the capabilities file to get auth config + let cap_path = self + .wasm_tools_dir + .join(format!("{}.capabilities.json", name)); + + if !cap_path.exists() { + // No capabilities = no auth needed + return Ok(AuthResult { + name: name.to_string(), + kind: ExtensionKind::WasmTool, + auth_url: None, + callback_type: None, + instructions: None, + setup_url: None, + awaiting_token: false, + status: "no_auth_required".to_string(), + }); + } + + let cap_bytes = tokio::fs::read(&cap_path) + .await + .map_err(|e| ExtensionError::Other(e.to_string()))?; + + let cap_file = crate::tools::wasm::CapabilitiesFile::from_bytes(&cap_bytes) + .map_err(|e| ExtensionError::Other(e.to_string()))?; + + let auth = match cap_file.auth { + Some(auth) => auth, + None => { + return Ok(AuthResult { + name: name.to_string(), + kind: ExtensionKind::WasmTool, + auth_url: None, + callback_type: None, + instructions: None, + setup_url: None, + awaiting_token: false, + status: "no_auth_required".to_string(), + }); + } + }; + + // Check env var first + if let Some(ref env_var) = auth.env_var { + if let Ok(value) = std::env::var(env_var) { + // Store the env var value as a secret + let params = CreateSecretParams::new(&auth.secret_name, &value) + .with_provider(name.to_string()); + self.secrets + .create(&self.user_id, params) + .await + .map_err(|e| ExtensionError::AuthFailed(e.to_string()))?; + + return Ok(AuthResult { + name: name.to_string(), + kind: ExtensionKind::WasmTool, + auth_url: None, + callback_type: None, + instructions: None, + setup_url: None, + awaiting_token: false, + status: "authenticated".to_string(), + }); + } + } + + // Check if already authenticated + if self + .secrets + .exists(&self.user_id, &auth.secret_name) + .await + .unwrap_or(false) + { + return Ok(AuthResult { + name: name.to_string(), + kind: ExtensionKind::WasmTool, + auth_url: None, + callback_type: None, + instructions: None, + setup_url: None, + awaiting_token: false, + status: "authenticated".to_string(), + }); + } + + // If a token was provided, store it + if let Some(token_value) = token { + let params = CreateSecretParams::new(&auth.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()))?; + + return Ok(AuthResult { + name: name.to_string(), + kind: ExtensionKind::WasmTool, + auth_url: None, + callback_type: None, + instructions: None, + setup_url: None, + awaiting_token: false, + status: "authenticated".to_string(), + }); + } + + // Return instructions for manual token entry + let display = auth.display_name.unwrap_or_else(|| name.to_string()); + let instructions = auth + .instructions + .unwrap_or_else(|| format!("Please provide your {} API token/key.", display)); + + Ok(AuthResult { + name: name.to_string(), + kind: ExtensionKind::WasmTool, + auth_url: None, + callback_type: None, + instructions: Some(instructions), + setup_url: auth.setup_url, + awaiting_token: true, + status: "awaiting_token".to_string(), + }) + } + + async fn activate_mcp(&self, name: &str) -> Result { + // Check if already activated + { + let clients = self.mcp_clients.read().await; + if clients.contains_key(name) { + // Already connected, just return the tool names + let tools: Vec = self + .tool_registry + .list() + .await + .into_iter() + .filter(|t| t.starts_with(&format!("{}_", name))) + .collect(); + + return Ok(ActivateResult { + name: name.to_string(), + kind: ExtensionKind::McpServer, + tools_loaded: tools, + message: format!("MCP server '{}' already active", name), + }); + } + } + + let server = get_mcp_server(name) + .await + .map_err(|e| ExtensionError::NotInstalled(e.to_string()))?; + + let has_tokens = is_authenticated(&server, &self.secrets, &self.user_id).await; + + let client = if has_tokens || server.requires_auth() { + McpClient::new_authenticated( + server.clone(), + Arc::clone(&self.mcp_session_manager), + Arc::clone(&self.secrets), + &self.user_id, + ) + } else { + McpClient::new_with_name(&server.name, &server.url) + }; + + // Try to list and create tools + let mcp_tools = client + .list_tools() + .await + .map_err(|e| ExtensionError::ActivationFailed(e.to_string()))?; + + let tool_impls = client + .create_tools() + .await + .map_err(|e| ExtensionError::ActivationFailed(e.to_string()))?; + + let tool_names: Vec = mcp_tools + .iter() + .map(|t| format!("{}_{}", name, t.name)) + .collect(); + + for tool in tool_impls { + self.tool_registry.register(tool).await; + } + + // Store the client + self.mcp_clients + .write() + .await + .insert(name.to_string(), Arc::new(client)); + + tracing::info!( + "Activated MCP server '{}' with {} tools", + name, + tool_names.len() + ); + + Ok(ActivateResult { + name: name.to_string(), + kind: ExtensionKind::McpServer, + tools_loaded: tool_names, + message: format!("Connected to '{}' and loaded tools", name), + }) + } + + async fn activate_wasm_tool(&self, name: &str) -> Result { + // Check if already active + if self.tool_registry.has(name).await { + return Ok(ActivateResult { + name: name.to_string(), + kind: ExtensionKind::WasmTool, + tools_loaded: vec![name.to_string()], + message: format!("WASM tool '{}' already active", name), + }); + } + + let runtime = self.wasm_tool_runtime.as_ref().ok_or_else(|| { + ExtensionError::ActivationFailed("WASM runtime not available".to_string()) + })?; + + let wasm_path = self.wasm_tools_dir.join(format!("{}.wasm", name)); + if !wasm_path.exists() { + return Err(ExtensionError::NotInstalled(format!( + "WASM tool '{}' not found at {}", + name, + wasm_path.display() + ))); + } + + let cap_path = self + .wasm_tools_dir + .join(format!("{}.capabilities.json", name)); + let cap_path_option = if cap_path.exists() { + Some(cap_path.as_path()) + } else { + None + }; + + let loader = WasmToolLoader::new(Arc::clone(runtime), Arc::clone(&self.tool_registry)); + loader + .load_from_files(name, &wasm_path, cap_path_option) + .await + .map_err(|e| ExtensionError::ActivationFailed(e.to_string()))?; + + tracing::info!("Activated WASM tool '{}'", name); + + Ok(ActivateResult { + name: name.to_string(), + kind: ExtensionKind::WasmTool, + tools_loaded: vec![name.to_string()], + message: format!("WASM tool '{}' loaded and ready", name), + }) + } + + /// Determine what kind of installed extension this is. + async fn determine_installed_kind(&self, name: &str) -> Result { + // Check MCP servers first + if get_mcp_server(name).await.is_ok() { + return Ok(ExtensionKind::McpServer); + } + + // Check WASM tools + let wasm_path = self.wasm_tools_dir.join(format!("{}.wasm", name)); + if wasm_path.exists() { + return Ok(ExtensionKind::WasmTool); + } + + // Check WASM channels + let channel_path = self.wasm_channels_dir.join(format!("{}.wasm", name)); + if channel_path.exists() { + return Ok(ExtensionKind::WasmChannel); + } + + Err(ExtensionError::NotInstalled(format!( + "'{}' is not installed as an MCP server, WASM tool, or WASM channel", + name + ))) + } + + async fn cleanup_expired_auths(&self) { + let mut pending = self.pending_auth.write().await; + pending.retain(|_, auth| auth.created_at.elapsed() < std::time::Duration::from_secs(300)); + } +} + +/// Infer the extension kind from a URL. +fn infer_kind_from_url(url: &str) -> ExtensionKind { + if url.ends_with(".wasm") { + ExtensionKind::WasmTool + } else { + ExtensionKind::McpServer + } +} + +#[cfg(test)] +mod tests { + use crate::extensions::ExtensionKind; + use crate::extensions::manager::infer_kind_from_url; + + #[test] + fn test_infer_kind_from_url() { + assert_eq!( + infer_kind_from_url("https://example.com/tool.wasm"), + ExtensionKind::WasmTool + ); + assert_eq!( + infer_kind_from_url("https://mcp.notion.com"), + ExtensionKind::McpServer + ); + assert_eq!( + infer_kind_from_url("https://example.com/mcp"), + ExtensionKind::McpServer + ); + } +} diff --git a/src/extensions/mod.rs b/src/extensions/mod.rs new file mode 100644 index 00000000..4d3b8512 --- /dev/null +++ b/src/extensions/mod.rs @@ -0,0 +1,224 @@ +//! Unified extension system for discovering, installing, authenticating, and activating +//! MCP servers and WASM tools through conversational agent interactions. +//! +//! Extensions are the user-facing abstraction over MCP servers and WASM tools. The agent +//! can search a built-in registry (or discover online), install, authenticate, and activate +//! extensions at runtime without CLI commands. +//! +//! ```text +//! User: "add notion" +//! -> tool_search("notion") -> finds MCP server in registry +//! -> tool_install("notion") -> saves config to mcp-servers.json +//! -> tool_auth("notion") -> OAuth 2.1 flow, returns URL +//! -> tool_activate("notion") -> connects, registers tools +//! ``` + +pub mod discovery; +pub mod manager; +pub mod registry; + +pub use discovery::OnlineDiscovery; +pub use manager::ExtensionManager; +pub use registry::ExtensionRegistry; + +use serde::{Deserialize, Serialize}; + +/// The kind of extension, determining how it's installed, authenticated, and activated. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ExtensionKind { + /// Hosted MCP server, HTTP transport, OAuth 2.1 auth. + McpServer, + /// Sandboxed WASM module, file-based, capabilities auth. + WasmTool, + /// WASM channel module (future: dynamic activation, currently needs restart). + WasmChannel, +} + +impl std::fmt::Display for ExtensionKind { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ExtensionKind::McpServer => write!(f, "mcp_server"), + ExtensionKind::WasmTool => write!(f, "wasm_tool"), + ExtensionKind::WasmChannel => write!(f, "wasm_channel"), + } + } +} + +/// A registry entry describing a known or discovered extension. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RegistryEntry { + /// Unique identifier (e.g., "notion", "weather", "telegram"). + pub name: String, + /// Human-readable name (e.g., "Notion", "Weather Tool"). + pub display_name: String, + /// What kind of extension this is. + pub kind: ExtensionKind, + /// Short description of what this extension does. + pub description: String, + /// Search keywords beyond the name. + #[serde(default)] + pub keywords: Vec, + /// Where to get this extension. + pub source: ExtensionSource, + /// How authentication works. + pub auth_hint: AuthHint, +} + +/// Where the extension binary or server lives. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum ExtensionSource { + /// URL to a hosted MCP server. + McpUrl { url: String }, + /// Downloadable WASM binary. + WasmDownload { + wasm_url: String, + #[serde(default)] + capabilities_url: Option, + }, + /// Build from source repository. + WasmBuildable { + repo_url: String, + #[serde(default)] + build_dir: Option, + }, + /// Discovered online (not yet validated for a specific source type). + Discovered { url: String }, +} + +/// Hint about what authentication method is needed. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum AuthHint { + /// MCP server supports Dynamic Client Registration (zero-config OAuth). + Dcr, + /// MCP server needs a pre-configured OAuth client_id. + OAuthPreConfigured { + /// URL where the user can create an OAuth app. + setup_url: String, + }, + /// WASM tool has auth defined in its capabilities.json file. + CapabilitiesAuth, + /// No authentication needed. + None, +} + +/// Where a search result came from. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ResultSource { + /// From the built-in curated registry. + Registry, + /// From online discovery (validated). + Discovered, +} + +/// Result of searching for extensions. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SearchResult { + /// The registry entry. + #[serde(flatten)] + pub entry: RegistryEntry, + /// Where this result came from. + pub source: ResultSource, + /// Whether the endpoint was validated (for discovered entries). + #[serde(default)] + pub validated: bool, +} + +/// Result of installing an extension. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct InstallResult { + pub name: String, + pub kind: ExtensionKind, + pub message: String, +} + +/// Result of authenticating an extension. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AuthResult { + pub name: String, + pub kind: ExtensionKind, + /// OAuth URL to open (for OAuth flows). + #[serde(skip_serializing_if = "Option::is_none")] + pub auth_url: Option, + /// Whether using local or remote callback. + #[serde(skip_serializing_if = "Option::is_none")] + pub callback_type: Option, + /// Instructions for manual token entry (for WASM tools). + #[serde(skip_serializing_if = "Option::is_none")] + pub instructions: Option, + /// URL for manual token setup. + #[serde(skip_serializing_if = "Option::is_none")] + pub setup_url: Option, + /// Whether the tool is waiting for a token from the user. + #[serde(default)] + pub awaiting_token: bool, + /// Current auth status. + pub status: String, +} + +/// Result of activating an extension. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ActivateResult { + pub name: String, + pub kind: ExtensionKind, + /// Names of tools that were loaded/registered. + pub tools_loaded: Vec, + pub message: String, +} + +/// An installed extension with its current status. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct InstalledExtension { + pub name: String, + pub kind: ExtensionKind, + #[serde(skip_serializing_if = "Option::is_none")] + pub description: Option, + pub authenticated: bool, + pub active: bool, + /// Tool names if active. + #[serde(default)] + pub tools: Vec, +} + +/// Error type for extension operations. +#[derive(Debug, thiserror::Error)] +pub enum ExtensionError { + #[error("Extension not found: {0}")] + NotFound(String), + + #[error("Extension already installed: {0}")] + AlreadyInstalled(String), + + #[error("Extension not installed: {0}")] + NotInstalled(String), + + #[error("Authentication failed: {0}")] + AuthFailed(String), + + #[error("Activation failed: {0}")] + ActivationFailed(String), + + #[error("Installation failed: {0}")] + InstallFailed(String), + + #[error("Discovery failed: {0}")] + DiscoveryFailed(String), + + #[error("Invalid URL: {0}")] + InvalidUrl(String), + + #[error("Download failed: {0}")] + DownloadFailed(String), + + #[error("Config error: {0}")] + Config(String), + + #[error("Channels require restart to activate")] + ChannelNeedsRestart, + + #[error("{0}")] + Other(String), +} diff --git a/src/extensions/registry.rs b/src/extensions/registry.rs new file mode 100644 index 00000000..8d136b08 --- /dev/null +++ b/src/extensions/registry.rs @@ -0,0 +1,545 @@ +//! Curated in-memory catalog of known extensions with fuzzy search. +//! +//! The registry holds well-known MCP servers and WASM tools that can be installed +//! via conversational commands. Online discoveries are cached here too. + +use tokio::sync::RwLock; + +use crate::extensions::{ + AuthHint, ExtensionKind, ExtensionSource, RegistryEntry, ResultSource, SearchResult, +}; + +/// Curated extension registry with fuzzy search. +pub struct ExtensionRegistry { + /// Built-in curated entries. + entries: Vec, + /// Cached entries from online discovery (session-lived). + discovery_cache: RwLock>, +} + +impl ExtensionRegistry { + /// Create a new registry populated with known extensions. + pub fn new() -> Self { + Self { + entries: builtin_entries(), + discovery_cache: RwLock::new(Vec::new()), + } + } + + /// Search the registry by query string. Returns results sorted by relevance. + /// + /// Splits the query into lowercase tokens and scores each entry by matches + /// in name, keywords, and description. + pub async fn search(&self, query: &str) -> Vec { + let tokens: Vec = query + .to_lowercase() + .split_whitespace() + .map(|s| s.to_string()) + .collect(); + + if tokens.is_empty() { + // Return all entries when query is empty + return self + .entries + .iter() + .map(|e| SearchResult { + entry: e.clone(), + source: ResultSource::Registry, + validated: true, + }) + .collect(); + } + + let mut scored: Vec<(SearchResult, u32)> = Vec::new(); + + // Score built-in entries + for entry in &self.entries { + let score = score_entry(entry, &tokens); + if score > 0 { + scored.push(( + SearchResult { + entry: entry.clone(), + source: ResultSource::Registry, + validated: true, + }, + score, + )); + } + } + + // Score cached discoveries + let cache = self.discovery_cache.read().await; + for entry in cache.iter() { + let score = score_entry(entry, &tokens); + if score > 0 { + scored.push(( + SearchResult { + entry: entry.clone(), + source: ResultSource::Discovered, + validated: true, + }, + score, + )); + } + } + + scored.sort_by(|a, b| b.1.cmp(&a.1)); + scored.into_iter().map(|(r, _)| r).collect() + } + + /// Look up an entry by exact name. + pub async fn get(&self, name: &str) -> Option { + if let Some(entry) = self.entries.iter().find(|e| e.name == name) { + return Some(entry.clone()); + } + let cache = self.discovery_cache.read().await; + cache.iter().find(|e| e.name == name).cloned() + } + + /// Add discovered entries to the cache. + pub async fn cache_discovered(&self, entries: Vec) { + let mut cache = self.discovery_cache.write().await; + for entry in entries { + // Deduplicate by name + if !cache.iter().any(|e| e.name == entry.name) { + cache.push(entry); + } + } + } +} + +impl Default for ExtensionRegistry { + fn default() -> Self { + Self::new() + } +} + +/// Score an entry against search tokens. Higher = better match. +fn score_entry(entry: &RegistryEntry, tokens: &[String]) -> u32 { + let mut score = 0u32; + let name_lower = entry.name.to_lowercase(); + let display_lower = entry.display_name.to_lowercase(); + let desc_lower = entry.description.to_lowercase(); + let keywords_lower: Vec = entry.keywords.iter().map(|k| k.to_lowercase()).collect(); + + for token in tokens { + // Exact name match is the strongest signal + if name_lower == *token { + score += 100; + } else if name_lower.contains(token.as_str()) { + score += 50; + } + + // Display name match + if display_lower.contains(token.as_str()) { + score += 30; + } + + // Keyword match + for kw in &keywords_lower { + if kw == token { + score += 40; + } else if kw.contains(token.as_str()) { + score += 20; + } + } + + // Description match (weakest signal) + if desc_lower.contains(token.as_str()) { + score += 10; + } + } + + score +} + +/// Well-known extensions that ship with ironclaw. +fn builtin_entries() -> Vec { + vec![ + // -- MCP Servers -- + RegistryEntry { + name: "notion".to_string(), + display_name: "Notion".to_string(), + kind: ExtensionKind::McpServer, + description: "Connect to Notion for reading and writing pages, databases, and comments" + .to_string(), + keywords: vec![ + "notes".into(), + "wiki".into(), + "docs".into(), + "pages".into(), + "database".into(), + ], + source: ExtensionSource::McpUrl { + url: "https://mcp.notion.com/mcp".to_string(), + }, + auth_hint: AuthHint::Dcr, + }, + RegistryEntry { + name: "linear".to_string(), + display_name: "Linear".to_string(), + kind: ExtensionKind::McpServer, + description: + "Connect to Linear for issue tracking, project management, and team workflows" + .to_string(), + keywords: vec![ + "issues".into(), + "tickets".into(), + "project".into(), + "tracking".into(), + "bugs".into(), + ], + source: ExtensionSource::McpUrl { + url: "https://mcp.linear.app".to_string(), + }, + auth_hint: AuthHint::Dcr, + }, + RegistryEntry { + name: "google-calendar".to_string(), + display_name: "Google Calendar".to_string(), + kind: ExtensionKind::McpServer, + description: "Connect to Google Calendar for managing events, schedules, and reminders" + .to_string(), + keywords: vec![ + "calendar".into(), + "events".into(), + "schedule".into(), + "meetings".into(), + "google".into(), + ], + source: ExtensionSource::McpUrl { + url: "https://mcp.google.com/calendar".to_string(), + }, + auth_hint: AuthHint::Dcr, + }, + RegistryEntry { + name: "google-drive".to_string(), + display_name: "Google Drive".to_string(), + kind: ExtensionKind::McpServer, + description: "Connect to Google Drive for file management, search, and document access" + .to_string(), + keywords: vec![ + "drive".into(), + "files".into(), + "documents".into(), + "storage".into(), + "google".into(), + ], + source: ExtensionSource::McpUrl { + url: "https://mcp.google.com/drive".to_string(), + }, + auth_hint: AuthHint::Dcr, + }, + RegistryEntry { + name: "github".to_string(), + display_name: "GitHub".to_string(), + kind: ExtensionKind::McpServer, + description: + "Connect to GitHub for repository management, issues, PRs, and code search" + .to_string(), + keywords: vec![ + "git".into(), + "repos".into(), + "code".into(), + "pull-request".into(), + "issues".into(), + ], + source: ExtensionSource::McpUrl { + url: "https://mcp.github.com".to_string(), + }, + auth_hint: AuthHint::Dcr, + }, + RegistryEntry { + name: "slack".to_string(), + display_name: "Slack".to_string(), + kind: ExtensionKind::McpServer, + description: + "Connect to Slack for messaging, channel management, and team communication" + .to_string(), + keywords: vec![ + "messaging".into(), + "chat".into(), + "channels".into(), + "team".into(), + "communication".into(), + ], + source: ExtensionSource::McpUrl { + url: "https://mcp.slack.com".to_string(), + }, + auth_hint: AuthHint::Dcr, + }, + RegistryEntry { + name: "sentry".to_string(), + display_name: "Sentry".to_string(), + kind: ExtensionKind::McpServer, + description: + "Connect to Sentry for error tracking, performance monitoring, and debugging" + .to_string(), + keywords: vec![ + "errors".into(), + "monitoring".into(), + "debugging".into(), + "crashes".into(), + "performance".into(), + ], + source: ExtensionSource::McpUrl { + url: "https://mcp.sentry.dev/sse".to_string(), + }, + auth_hint: AuthHint::Dcr, + }, + RegistryEntry { + name: "stripe".to_string(), + display_name: "Stripe".to_string(), + kind: ExtensionKind::McpServer, + description: + "Connect to Stripe for payment processing, subscriptions, and financial data" + .to_string(), + keywords: vec![ + "payments".into(), + "billing".into(), + "subscriptions".into(), + "invoices".into(), + "finance".into(), + ], + source: ExtensionSource::McpUrl { + url: "https://mcp.stripe.com".to_string(), + }, + auth_hint: AuthHint::Dcr, + }, + RegistryEntry { + name: "cloudflare".to_string(), + display_name: "Cloudflare".to_string(), + kind: ExtensionKind::McpServer, + description: + "Connect to Cloudflare for DNS, Workers, KV, and infrastructure management" + .to_string(), + keywords: vec![ + "cdn".into(), + "dns".into(), + "workers".into(), + "hosting".into(), + "infrastructure".into(), + ], + source: ExtensionSource::McpUrl { + url: "https://mcp.cloudflare.com/sse".to_string(), + }, + auth_hint: AuthHint::Dcr, + }, + RegistryEntry { + name: "asana".to_string(), + display_name: "Asana".to_string(), + kind: ExtensionKind::McpServer, + description: "Connect to Asana for task management, projects, and team coordination" + .to_string(), + keywords: vec![ + "tasks".into(), + "projects".into(), + "management".into(), + "team".into(), + ], + source: ExtensionSource::McpUrl { + url: "https://mcp.asana.com".to_string(), + }, + auth_hint: AuthHint::Dcr, + }, + RegistryEntry { + name: "intercom".to_string(), + display_name: "Intercom".to_string(), + kind: ExtensionKind::McpServer, + description: "Connect to Intercom for customer messaging, support, and engagement" + .to_string(), + keywords: vec![ + "support".into(), + "customers".into(), + "messaging".into(), + "chat".into(), + "helpdesk".into(), + ], + source: ExtensionSource::McpUrl { + url: "https://mcp.intercom.com".to_string(), + }, + auth_hint: AuthHint::Dcr, + }, + ] +} + +#[cfg(test)] +mod tests { + use crate::extensions::registry::{ExtensionRegistry, score_entry}; + use crate::extensions::{AuthHint, ExtensionKind, ExtensionSource, RegistryEntry}; + + #[test] + fn test_score_exact_name_match() { + let entry = RegistryEntry { + name: "notion".to_string(), + display_name: "Notion".to_string(), + kind: ExtensionKind::McpServer, + description: "Workspace tool".to_string(), + keywords: vec!["notes".into()], + source: ExtensionSource::McpUrl { + url: "https://example.com".to_string(), + }, + auth_hint: AuthHint::Dcr, + }; + + let score = score_entry(&entry, &["notion".to_string()]); + assert!( + score >= 100, + "Exact name match should score >= 100, got {}", + score + ); + } + + #[test] + fn test_score_partial_name_match() { + let entry = RegistryEntry { + name: "google-calendar".to_string(), + display_name: "Google Calendar".to_string(), + kind: ExtensionKind::McpServer, + description: "Calendar management".to_string(), + keywords: vec!["events".into()], + source: ExtensionSource::McpUrl { + url: "https://example.com".to_string(), + }, + auth_hint: AuthHint::Dcr, + }; + + let score = score_entry(&entry, &["calendar".to_string()]); + assert!( + score > 0, + "Partial name match should score > 0, got {}", + score + ); + } + + #[test] + fn test_score_keyword_match() { + let entry = RegistryEntry { + name: "notion".to_string(), + display_name: "Notion".to_string(), + kind: ExtensionKind::McpServer, + description: "Workspace tool".to_string(), + keywords: vec!["wiki".into(), "notes".into()], + source: ExtensionSource::McpUrl { + url: "https://example.com".to_string(), + }, + auth_hint: AuthHint::Dcr, + }; + + let score = score_entry(&entry, &["wiki".to_string()]); + assert!( + score >= 40, + "Exact keyword match should score >= 40, got {}", + score + ); + } + + #[test] + fn test_score_no_match() { + let entry = RegistryEntry { + name: "notion".to_string(), + display_name: "Notion".to_string(), + kind: ExtensionKind::McpServer, + description: "Workspace tool".to_string(), + keywords: vec!["notes".into()], + source: ExtensionSource::McpUrl { + url: "https://example.com".to_string(), + }, + auth_hint: AuthHint::Dcr, + }; + + let score = score_entry(&entry, &["xyzfoobar".to_string()]); + assert_eq!(score, 0, "No match should score 0"); + } + + #[tokio::test] + async fn test_search_returns_sorted() { + let registry = ExtensionRegistry::new(); + let results = registry.search("notion").await; + + assert!(!results.is_empty(), "Should find notion in registry"); + assert_eq!(results[0].entry.name, "notion"); + } + + #[tokio::test] + async fn test_search_empty_query_returns_all() { + let registry = ExtensionRegistry::new(); + let results = registry.search("").await; + + assert!(results.len() > 5, "Empty query should return all entries"); + } + + #[tokio::test] + async fn test_search_by_keyword() { + let registry = ExtensionRegistry::new(); + let results = registry.search("issues tickets").await; + + assert!( + !results.is_empty(), + "Should find entries matching 'issues tickets'" + ); + // Linear should be near the top since it has both keywords + let linear_pos = results.iter().position(|r| r.entry.name == "linear"); + assert!(linear_pos.is_some(), "Linear should appear in results"); + } + + #[tokio::test] + async fn test_get_exact_name() { + let registry = ExtensionRegistry::new(); + + let entry = registry.get("notion").await; + assert!(entry.is_some()); + assert_eq!(entry.unwrap().display_name, "Notion"); + + let missing = registry.get("nonexistent").await; + assert!(missing.is_none()); + } + + #[tokio::test] + async fn test_cache_discovered() { + let registry = ExtensionRegistry::new(); + + let discovered = RegistryEntry { + name: "custom-mcp".to_string(), + display_name: "Custom MCP".to_string(), + kind: ExtensionKind::McpServer, + description: "A custom MCP server".to_string(), + keywords: vec![], + source: ExtensionSource::McpUrl { + url: "https://custom.example.com".to_string(), + }, + auth_hint: AuthHint::Dcr, + }; + + registry.cache_discovered(vec![discovered]).await; + + let entry = registry.get("custom-mcp").await; + assert!(entry.is_some()); + + let results = registry.search("custom").await; + assert!(!results.is_empty()); + } + + #[tokio::test] + async fn test_cache_deduplication() { + let registry = ExtensionRegistry::new(); + + let entry = RegistryEntry { + name: "dup".to_string(), + display_name: "Dup".to_string(), + kind: ExtensionKind::McpServer, + description: "Test".to_string(), + keywords: vec![], + source: ExtensionSource::McpUrl { + url: "https://example.com".to_string(), + }, + auth_hint: AuthHint::None, + }; + + registry.cache_discovered(vec![entry.clone()]).await; + registry.cache_discovered(vec![entry]).await; + + let results = registry.search("dup").await; + assert_eq!(results.len(), 1, "Should not duplicate cached entries"); + } +} diff --git a/src/lib.rs b/src/lib.rs index 7cabc3f5..d78cbff7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -46,6 +46,7 @@ pub mod context; pub mod error; pub mod estimation; pub mod evaluation; +pub mod extensions; pub mod history; pub mod llm; pub mod safety; diff --git a/src/main.rs b/src/main.rs index c93f6477..4e177154 100644 --- a/src/main.rs +++ b/src/main.rs @@ -19,6 +19,7 @@ use ironclaw::{ }, config::Config, context::ContextManager, + extensions::ExtensionManager, history::Store, llm::{SessionConfig, create_llm_provider, create_session_manager}, safety::SafetyLayer, @@ -356,8 +357,10 @@ async fn main() -> anyhow::Result<()> { tracing::info!("Builder mode enabled"); } - // Load installed WASM tools - if config.wasm.enabled && config.wasm.tools_dir.exists() { + // 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); @@ -380,12 +383,17 @@ async fn main() -> anyhow::Result<()> { 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> = @@ -490,6 +498,29 @@ async fn main() -> anyhow::Result<()> { } } + // Create extension manager for in-chat discovery/install/auth/activate + let extension_manager = if let Some(ref secrets) = secrets_store { + let manager = Arc::new(ExtensionManager::new( + Arc::clone(&mcp_session_manager), + Arc::clone(secrets), + Arc::clone(&tools), + wasm_tool_runtime.clone(), + config.wasm.tools_dir.clone(), + config.channels.wasm_channels_dir.clone(), + config.tunnel.public_url.clone(), + "default".to_string(), + )); + tools.register_extension_tools(Arc::clone(&manager)); + tracing::info!("Extension manager initialized with in-chat discovery tools"); + Some(manager) + } else { + tracing::debug!( + "Extension manager not available (no secrets store). \ + Extension tools won't be registered." + ); + None + }; + tracing::info!( "Tool registry initialized with {} total tools", tools.count() @@ -657,7 +688,10 @@ async fn main() -> anyhow::Result<()> { // Start WASM channel webhook server if we have channels with webhooks if has_webhook_channels && config.tunnel.public_url.is_some() { - let server = WasmChannelServer::new(wasm_router); + let mut server = WasmChannelServer::new(wasm_router); + if let Some(ref ext_mgr) = extension_manager { + server = server.with_extension_manager(Arc::clone(ext_mgr)); + } let addr = std::net::SocketAddr::from(([0, 0, 0, 0], 8080)); match server.start(addr).await { Ok(_handle) => { diff --git a/src/tools/builtin/extension_tools.rs b/src/tools/builtin/extension_tools.rs new file mode 100644 index 00000000..41c9d708 --- /dev/null +++ b/src/tools/builtin/extension_tools.rs @@ -0,0 +1,523 @@ +//! Agent-callable tools for managing extensions (MCP servers and WASM tools). +//! +//! These six tools let the LLM search, install, authenticate, activate, list, +//! and remove extensions entirely through conversation. + +use std::sync::Arc; + +use async_trait::async_trait; + +use crate::context::JobContext; +use crate::extensions::{ExtensionKind, ExtensionManager}; +use crate::tools::tool::{Tool, ToolError, ToolOutput}; + +// ── tool_search ────────────────────────────────────────────────────────── + +pub struct ToolSearchTool { + manager: Arc, +} + +impl ToolSearchTool { + pub fn new(manager: Arc) -> Self { + Self { manager } + } +} + +#[async_trait] +impl Tool for ToolSearchTool { + fn name(&self) -> &str { + "tool_search" + } + + fn description(&self) -> &str { + "Search for available extensions (MCP servers, WASM tools) to add. \ + Use discover:true to search online if the built-in registry has no results." + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Search query (name, keyword, or description fragment)" + }, + "discover": { + "type": "boolean", + "description": "If true, also search online (slower, 5-15s). Try without first.", + "default": false + } + }, + "required": ["query"] + }) + } + + async fn execute( + &self, + params: serde_json::Value, + _ctx: &JobContext, + ) -> Result { + let start = std::time::Instant::now(); + + let query = params.get("query").and_then(|v| v.as_str()).unwrap_or(""); + let discover = params + .get("discover") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + + let results = self + .manager + .search(query, discover) + .await + .map_err(|e| ToolError::ExecutionFailed(e.to_string()))?; + + let output = serde_json::json!({ + "results": results, + "count": results.len(), + "searched_online": discover, + }); + + Ok(ToolOutput::success(output, start.elapsed())) + } +} + +// ── tool_install ───────────────────────────────────────────────────────── + +pub struct ToolInstallTool { + manager: Arc, +} + +impl ToolInstallTool { + pub fn new(manager: Arc) -> Self { + Self { manager } + } +} + +#[async_trait] +impl Tool for ToolInstallTool { + fn name(&self) -> &str { + "tool_install" + } + + fn description(&self) -> &str { + "Install an extension (MCP server or WASM tool). \ + Use the name from tool_search results, or provide an explicit URL." + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "name": { + "type": "string", + "description": "Extension name (from search results or custom)" + }, + "url": { + "type": "string", + "description": "Explicit URL (for extensions not in the registry)" + }, + "kind": { + "type": "string", + "enum": ["mcp_server", "wasm_tool"], + "description": "Extension type (auto-detected if omitted)" + } + }, + "required": ["name"] + }) + } + + async fn execute( + &self, + params: serde_json::Value, + _ctx: &JobContext, + ) -> Result { + let start = std::time::Instant::now(); + + let name = params + .get("name") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::InvalidParameters("name is required".to_string()))?; + + let url = params.get("url").and_then(|v| v.as_str()); + + let kind_hint = params + .get("kind") + .and_then(|v| v.as_str()) + .and_then(|k| match k { + "mcp_server" => Some(ExtensionKind::McpServer), + "wasm_tool" => Some(ExtensionKind::WasmTool), + _ => None, + }); + + let result = self + .manager + .install(name, url, kind_hint) + .await + .map_err(|e| ToolError::ExecutionFailed(e.to_string()))?; + + let output = serde_json::to_value(&result) + .unwrap_or_else(|_| serde_json::json!({"error": "serialization failed"})); + + Ok(ToolOutput::success(output, start.elapsed())) + } + + fn requires_approval(&self) -> bool { + true + } +} + +// ── tool_auth ──────────────────────────────────────────────────────────── + +pub struct ToolAuthTool { + manager: Arc, +} + +impl ToolAuthTool { + pub fn new(manager: Arc) -> Self { + Self { manager } + } +} + +#[async_trait] +impl Tool for ToolAuthTool { + fn name(&self) -> &str { + "tool_auth" + } + + fn description(&self) -> &str { + "Authenticate an installed extension. For MCP servers, starts OAuth flow. \ + For WASM tools with manual auth, returns instructions; call again with token param to complete." + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "name": { + "type": "string", + "description": "Extension name to authenticate" + }, + "token": { + "type": "string", + "description": "API token/key for manual auth (WASM tools). Provide after user gives you the token." + } + }, + "required": ["name"] + }) + } + + async fn execute( + &self, + params: serde_json::Value, + _ctx: &JobContext, + ) -> Result { + let start = std::time::Instant::now(); + + let name = params + .get("name") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::InvalidParameters("name is required".to_string()))?; + + let token = params.get("token").and_then(|v| v.as_str()); + + let result = self + .manager + .auth(name, token) + .await + .map_err(|e| ToolError::ExecutionFailed(e.to_string()))?; + + let output = serde_json::to_value(&result) + .unwrap_or_else(|_| serde_json::json!({"error": "serialization failed"})); + + Ok(ToolOutput::success(output, start.elapsed())) + } + + fn requires_approval(&self) -> bool { + true + } +} + +// ── tool_activate ──────────────────────────────────────────────────────── + +pub struct ToolActivateTool { + manager: Arc, +} + +impl ToolActivateTool { + pub fn new(manager: Arc) -> Self { + Self { manager } + } +} + +#[async_trait] +impl Tool for ToolActivateTool { + fn name(&self) -> &str { + "tool_activate" + } + + fn description(&self) -> &str { + "Activate an installed extension, connecting to MCP servers or loading WASM tools into the runtime." + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "name": { + "type": "string", + "description": "Extension name to activate" + } + }, + "required": ["name"] + }) + } + + async fn execute( + &self, + params: serde_json::Value, + _ctx: &JobContext, + ) -> Result { + let start = std::time::Instant::now(); + + let name = params + .get("name") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::InvalidParameters("name is required".to_string()))?; + + let result = self + .manager + .activate(name) + .await + .map_err(|e| ToolError::ExecutionFailed(e.to_string()))?; + + let output = serde_json::to_value(&result) + .unwrap_or_else(|_| serde_json::json!({"error": "serialization failed"})); + + Ok(ToolOutput::success(output, start.elapsed())) + } +} + +// ── tool_list ──────────────────────────────────────────────────────────── + +pub struct ToolListTool { + manager: Arc, +} + +impl ToolListTool { + pub fn new(manager: Arc) -> Self { + Self { manager } + } +} + +#[async_trait] +impl Tool for ToolListTool { + fn name(&self) -> &str { + "tool_list" + } + + fn description(&self) -> &str { + "List all installed extensions with their authentication and activation status." + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "kind": { + "type": "string", + "enum": ["mcp_server", "wasm_tool", "wasm_channel"], + "description": "Filter by extension type (omit to list all)" + } + } + }) + } + + async fn execute( + &self, + params: serde_json::Value, + _ctx: &JobContext, + ) -> Result { + let start = std::time::Instant::now(); + + let kind_filter = params + .get("kind") + .and_then(|v| v.as_str()) + .and_then(|k| match k { + "mcp_server" => Some(ExtensionKind::McpServer), + "wasm_tool" => Some(ExtensionKind::WasmTool), + "wasm_channel" => Some(ExtensionKind::WasmChannel), + _ => None, + }); + + let extensions = self + .manager + .list(kind_filter) + .await + .map_err(|e| ToolError::ExecutionFailed(e.to_string()))?; + + let output = serde_json::json!({ + "extensions": extensions, + "count": extensions.len(), + }); + + Ok(ToolOutput::success(output, start.elapsed())) + } +} + +// ── tool_remove ────────────────────────────────────────────────────────── + +pub struct ToolRemoveTool { + manager: Arc, +} + +impl ToolRemoveTool { + pub fn new(manager: Arc) -> Self { + Self { manager } + } +} + +#[async_trait] +impl Tool for ToolRemoveTool { + fn name(&self) -> &str { + "tool_remove" + } + + fn description(&self) -> &str { + "Remove an installed extension (MCP server or WASM tool). \ + Unregisters tools and deletes configuration." + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "name": { + "type": "string", + "description": "Extension name to remove" + } + }, + "required": ["name"] + }) + } + + async fn execute( + &self, + params: serde_json::Value, + _ctx: &JobContext, + ) -> Result { + let start = std::time::Instant::now(); + + let name = params + .get("name") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::InvalidParameters("name is required".to_string()))?; + + let message = self + .manager + .remove(name) + .await + .map_err(|e| ToolError::ExecutionFailed(e.to_string()))?; + + let output = serde_json::json!({ + "name": name, + "message": message, + }); + + Ok(ToolOutput::success(output, start.elapsed())) + } + + fn requires_approval(&self) -> bool { + true + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_tool_search_schema() { + let tool = ToolSearchTool { + manager: test_manager_stub(), + }; + assert_eq!(tool.name(), "tool_search"); + let schema = tool.parameters_schema(); + assert!(schema.get("properties").is_some()); + assert!(schema["properties"].get("query").is_some()); + } + + #[test] + fn test_tool_install_schema() { + let tool = ToolInstallTool { + manager: test_manager_stub(), + }; + assert_eq!(tool.name(), "tool_install"); + assert!(tool.requires_approval()); + let schema = tool.parameters_schema(); + assert!(schema["properties"].get("name").is_some()); + assert!(schema["properties"].get("url").is_some()); + } + + #[test] + fn test_tool_auth_schema() { + let tool = ToolAuthTool { + manager: test_manager_stub(), + }; + assert_eq!(tool.name(), "tool_auth"); + assert!(tool.requires_approval()); + let schema = tool.parameters_schema(); + assert!(schema["properties"].get("name").is_some()); + assert!(schema["properties"].get("token").is_some()); + } + + #[test] + fn test_tool_activate_schema() { + let tool = ToolActivateTool { + manager: test_manager_stub(), + }; + assert_eq!(tool.name(), "tool_activate"); + assert!(!tool.requires_approval()); + } + + #[test] + fn test_tool_list_schema() { + let tool = ToolListTool { + manager: test_manager_stub(), + }; + assert_eq!(tool.name(), "tool_list"); + assert!(!tool.requires_approval()); + let schema = tool.parameters_schema(); + assert!(schema["properties"].get("kind").is_some()); + } + + #[test] + fn test_tool_remove_schema() { + let tool = ToolRemoveTool { + manager: test_manager_stub(), + }; + assert_eq!(tool.name(), "tool_remove"); + assert!(tool.requires_approval()); + } + + /// Create a stub manager for schema tests (these don't call execute). + fn test_manager_stub() -> Arc { + use crate::secrets::{InMemorySecretsStore, SecretsCrypto}; + use crate::tools::ToolRegistry; + use crate::tools::mcp::session::McpSessionManager; + + let master_key = + secrecy::SecretString::from("0123456789abcdef0123456789abcdef".to_string()); + let crypto = Arc::new(SecretsCrypto::new(master_key).unwrap()); + + Arc::new(ExtensionManager::new( + Arc::new(McpSessionManager::new()), + Arc::new(InMemorySecretsStore::new(crypto)), + Arc::new(ToolRegistry::new()), + None, + std::path::PathBuf::from("/tmp/ironclaw-test-tools"), + std::path::PathBuf::from("/tmp/ironclaw-test-channels"), + None, + "test".to_string(), + )) + } +} diff --git a/src/tools/builtin/mod.rs b/src/tools/builtin/mod.rs index b922ca66..46aee56b 100644 --- a/src/tools/builtin/mod.rs +++ b/src/tools/builtin/mod.rs @@ -2,6 +2,7 @@ mod echo; mod ecommerce; +pub mod extension_tools; mod file; mod http; mod job; @@ -15,6 +16,9 @@ mod time; pub use echo::EchoTool; pub use ecommerce::EcommerceTool; +pub use extension_tools::{ + ToolActivateTool, ToolAuthTool, ToolInstallTool, ToolListTool, ToolRemoveTool, ToolSearchTool, +}; pub use file::{ApplyPatchTool, ListDirTool, ReadFileTool, WriteFileTool}; pub use http::HttpTool; pub use job::{CancelJobTool, CreateJobTool, JobStatusTool, ListJobsTool}; diff --git a/src/tools/mcp/auth.rs b/src/tools/mcp/auth.rs index 98b61afb..20766dd1 100644 --- a/src/tools/mcp/auth.rs +++ b/src/tools/mcp/auth.rs @@ -467,7 +467,7 @@ pub async fn authorize_mcp_server( } /// Find an available port for the OAuth callback. -async fn find_available_port() -> Result<(TcpListener, u16), AuthError> { +pub async fn find_available_port() -> Result<(TcpListener, u16), AuthError> { for port in 9876..=9886 { if let Ok(listener) = TcpListener::bind(format!("127.0.0.1:{}", port)).await { return Ok((listener, port)); @@ -477,7 +477,7 @@ async fn find_available_port() -> Result<(TcpListener, u16), AuthError> { } /// Build the authorization URL with all required parameters. -fn build_authorization_url( +pub fn build_authorization_url( base_url: &str, client_id: &str, redirect_uri: &str, @@ -518,7 +518,7 @@ fn build_authorization_url( } /// Wait for the authorization callback and extract the code. -async fn wait_for_authorization_callback( +pub async fn wait_for_authorization_callback( listener: TcpListener, server_name: &str, ) -> Result { @@ -590,7 +590,7 @@ async fn wait_for_authorization_callback( } /// Exchange the authorization code for an access token. -async fn exchange_code_for_token( +pub async fn exchange_code_for_token( token_url: &str, client_id: &str, code: &str, @@ -644,7 +644,7 @@ async fn exchange_code_for_token( } /// Store access and refresh tokens securely. -async fn store_tokens( +pub async fn store_tokens( secrets: &Arc, user_id: &str, server_config: &McpServerConfig, @@ -675,7 +675,7 @@ async fn store_tokens( } /// Store the DCR client ID for future token refresh. -async fn store_client_id( +pub async fn store_client_id( secrets: &Arc, user_id: &str, server_config: &McpServerConfig, diff --git a/src/tools/registry.rs b/src/tools/registry.rs index 9e87b751..acc038a7 100644 --- a/src/tools/registry.rs +++ b/src/tools/registry.rs @@ -6,13 +6,15 @@ use std::sync::Arc; use tokio::sync::RwLock; use crate::context::ContextManager; +use crate::extensions::ExtensionManager; use crate::llm::{LlmProvider, ToolDefinition}; use crate::safety::SafetyLayer; use crate::tools::builder::{BuildSoftwareTool, BuilderConfig, LlmSoftwareBuilder}; use crate::tools::builtin::{ ApplyPatchTool, CancelJobTool, CreateJobTool, EchoTool, HttpTool, JobStatusTool, JsonTool, ListDirTool, ListJobsTool, MemoryReadTool, MemorySearchTool, MemoryTreeTool, MemoryWriteTool, - ReadFileTool, ShellTool, TimeTool, WriteFileTool, + ReadFileTool, ShellTool, TimeTool, ToolActivateTool, ToolAuthTool, ToolInstallTool, + ToolListTool, ToolRemoveTool, ToolSearchTool, WriteFileTool, }; use crate::tools::tool::Tool; use crate::tools::wasm::{ @@ -159,6 +161,19 @@ impl ToolRegistry { tracing::info!("Registered 4 job management tools"); } + /// Register extension management tools (search, install, auth, activate, list, remove). + /// + /// These allow the LLM to manage MCP servers and WASM tools through conversation. + pub fn register_extension_tools(&self, manager: Arc) { + self.register_sync(Arc::new(ToolSearchTool::new(Arc::clone(&manager)))); + self.register_sync(Arc::new(ToolInstallTool::new(Arc::clone(&manager)))); + self.register_sync(Arc::new(ToolAuthTool::new(Arc::clone(&manager)))); + self.register_sync(Arc::new(ToolActivateTool::new(Arc::clone(&manager)))); + self.register_sync(Arc::new(ToolListTool::new(Arc::clone(&manager)))); + self.register_sync(Arc::new(ToolRemoveTool::new(manager))); + tracing::info!("Registered 6 extension management tools"); + } + /// Register the software builder tool. /// /// The builder tool allows the agent to create new software including WASM tools,