gitops/optimclaw
1
0
Fork 0
mirror of https://github.com/outbackdingo/optimclaw.git synced 2026-08-25 14:53:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add in-chat extension discovery, auth, and activation system

Browse Source 
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 <[email protected]>
This commit is contained in:
Illia Polosukhin
2026-02-05 20:11:29 -08:00
co-authored by Claude Opus 4.6
parent e0016a95e8
commit ae3c86a7ea
11 changed files with 2664 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/channels/wasm/router.rs
+89 -5
View File
@@ -173,11 +173,23 @@ impl Default for WasmChannelRouter {
#[derive(Clone)]
pub struct RouterState {
router: Arc<WasmChannelRouter>,
extension_manager: Option<Arc<crate::extensions::ExtensionManager>>,
}
impl RouterState {
pub fn new(router: Arc<WasmChannelRouter>) -> Self {
Self { router }
Self {
router,
extension_manager: None,
}
}
pub fn with_extension_manager(
mut self,
manager: Arc<crate::extensions::ExtensionManager>,
) -> 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<RouterState>,
Query(params): Query<HashMap<String, String>>,
) -> 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!(
"<!DOCTYPE html><html><body style=\"font-family: sans-serif; \
display: flex; justify-content: center; align-items: center; \
height: 100vh; margin: 0; background: #191919; color: white;\">\
<div style=\"text-align: center;\">\
<h1>Authorization Failed</h1>\
<p>Error: {}</p>\
</div></body></html>",
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(
"<!DOCTYPE html><html><body style=\"font-family: sans-serif; \
display: flex; justify-content: center; align-items: center; \
height: 100vh; margin: 0; background: #191919; color: white;\">\
<div style=\"text-align: center;\">\
<h1>Connected!</h1>\
<p>You can close this window and return to IronClaw.</p>\
</div></body></html>"
.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<WasmChannelRouter>) -> Router {
let state = RouterState::new(router);
pub fn create_wasm_channel_router(
router: Arc<WasmChannelRouter>,
extension_manager: Option<Arc<crate::extensions::ExtensionManager>>,
) -> 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<WasmChannelRouter>) -> Router {
/// HTTP server for WASM channel webhooks.
pub struct WasmChannelServer {
router: Arc<WasmChannelRouter>,
extension_manager: Option<Arc<crate::extensions::ExtensionManager>>,
}
impl WasmChannelServer {
/// Create a new server.
pub fn new(router: Arc<WasmChannelRouter>) -> 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<crate::extensions::ExtensionManager>,
) -> Self {
self.extension_manager = Some(manager);
self
}
/// Start the HTTP server.
@@ -416,7 +500,7 @@ impl WasmChannelServer {
&self,
addr: SocketAddr,
) -> Result<tokio::task::JoinHandle<()>, 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?;
src/extensions/discovery.rs
+326
View File
@@ -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<RegistryEntry> {
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<RegistryEntry> = 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<RegistryEntry> {
// 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<RegistryEntry> {
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::<serde_json::Value>(&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<T>(
future: impl std::future::Future<Output = T>,
duration: Duration,
) -> Option<T> {
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::<Vec<_>>()
.join(" ")
}
#[derive(Debug, Deserialize)]
struct GitHubSearchResponse {
#[serde(default)]
items: Vec<GitHubRepo>,
}
#[derive(Debug, Deserialize)]
struct GitHubRepo {
name: String,
full_name: String,
html_url: String,
description: Option<String>,
#[serde(default)]
homepage: Option<String>,
#[serde(default)]
topics: Vec<String>,
}
#[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();
}
}
src/extensions/manager.rs
+892
View File
@@ -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<McpSessionManager>,
/// Active MCP clients keyed by server name.
mcp_clients: RwLock<HashMap<String, Arc<McpClient>>>,
// WASM tool infrastructure
wasm_tool_runtime: Option<Arc<WasmToolRuntime>>,
wasm_tools_dir: PathBuf,
wasm_channels_dir: PathBuf,
// Shared
secrets: Arc<dyn SecretsStore + Send + Sync>,
tool_registry: Arc<ToolRegistry>,
pending_auth: RwLock<HashMap<String, PendingAuth>>,
/// Tunnel URL for remote OAuth callbacks (used in future iterations).
_tunnel_url: Option<String>,
user_id: String,
}
impl ExtensionManager {
#[allow(clippy::too_many_arguments)]
pub fn new(
mcp_session_manager: Arc<McpSessionManager>,
secrets: Arc<dyn SecretsStore + Send + Sync>,
tool_registry: Arc<ToolRegistry>,
wasm_tool_runtime: Option<Arc<WasmToolRuntime>>,
wasm_tools_dir: PathBuf,
wasm_channels_dir: PathBuf,
tunnel_url: Option<String>,
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<Vec<SearchResult>, 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<ExtensionKind>,
) -> Result<InstallResult, ExtensionError> {
// 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<AuthResult, ExtensionError> {
// 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<ActivateResult, ExtensionError> {
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<ExtensionKind>,
) -> Result<Vec<InstalledExtension>, 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<String, ExtensionError> {
let kind = self.determine_installed_kind(name).await?;
match kind {
ExtensionKind::McpServer => {
// Unregister tools with this server's prefix
let tool_names: Vec<String> = 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<InstallResult, ExtensionError> {
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<InstallResult, ExtensionError> {
// 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<InstallResult, ExtensionError> {
// 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<AuthResult, ExtensionError> {
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<AuthResult, ExtensionError> {
// 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<AuthResult, ExtensionError> {
// 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<ActivateResult, ExtensionError> {
// 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<String> = 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<String> = 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<ActivateResult, ExtensionError> {
// 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<ExtensionKind, ExtensionError> {
// 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
);
}
}
src/extensions/mod.rs
+224
View File
@@ -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<String>,
/// 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<String>,
},
/// Build from source repository.
WasmBuildable {
repo_url: String,
#[serde(default)]
build_dir: Option<String>,
},
/// 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<String>,
/// Whether using local or remote callback.
#[serde(skip_serializing_if = "Option::is_none")]
pub callback_type: Option<String>,
/// Instructions for manual token entry (for WASM tools).
#[serde(skip_serializing_if = "Option::is_none")]
pub instructions: Option<String>,
/// URL for manual token setup.
#[serde(skip_serializing_if = "Option::is_none")]
pub setup_url: Option<String>,
/// 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<String>,
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<String>,
pub authenticated: bool,
pub active: bool,
/// Tool names if active.
#[serde(default)]
pub tools: Vec<String>,
}
/// 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),
}
src/extensions/registry.rs
+545
View File
@@ -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<RegistryEntry>,
/// Cached entries from online discovery (session-lived).
discovery_cache: RwLock<Vec<RegistryEntry>>,
}
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<SearchResult> {
let tokens: Vec<String> = 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<RegistryEntry> {
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<RegistryEntry>) {
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<String> = 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<RegistryEntry> {
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");
}
}
src/lib.rs
+1
View File
@@ -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;
src/main.rs
+38 -4
View File
@@ -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<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);
@@ -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<Arc<dyn SecretsStore + Send + Sync>> =
@@ -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) => {
src/tools/builtin/extension_tools.rs
+523
View File
@@ -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<ExtensionManager>,
}
impl ToolSearchTool {
pub fn new(manager: Arc<ExtensionManager>) -> 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<ToolOutput, ToolError> {
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<ExtensionManager>,
}
impl ToolInstallTool {
pub fn new(manager: Arc<ExtensionManager>) -> 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<ToolOutput, ToolError> {
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<ExtensionManager>,
}
impl ToolAuthTool {
pub fn new(manager: Arc<ExtensionManager>) -> 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<ToolOutput, ToolError> {
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<ExtensionManager>,
}
impl ToolActivateTool {
pub fn new(manager: Arc<ExtensionManager>) -> 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<ToolOutput, ToolError> {
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<ExtensionManager>,
}
impl ToolListTool {
pub fn new(manager: Arc<ExtensionManager>) -> 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<ToolOutput, ToolError> {
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<ExtensionManager>,
}
impl ToolRemoveTool {
pub fn new(manager: Arc<ExtensionManager>) -> 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<ToolOutput, ToolError> {
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<ExtensionManager> {
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(),
))
}
}
src/tools/builtin/mod.rs
+4
View File
@@ -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};
src/tools/mcp/auth.rs
+6 -6
View File
@@ -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<String, AuthError> {
@@ -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<dyn SecretsStore + Send + Sync>,
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<dyn SecretsStore + Send + Sync>,
user_id: &str,
server_config: &McpServerConfig,
src/tools/registry.rs
+16 -1
View File
@@ -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<ExtensionManager>) {
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,