diff --git a/.env.example b/.env.example index c3a81008..1a80fc02 100644 --- a/.env.example +++ b/.env.example @@ -1,5 +1,5 @@ # Database Configuration -DATABASE_URL=postgres://near_agent:password@localhost:5432/near_agent +DATABASE_URL=postgres://ironclaw:password@localhost:5432/ironclaw DATABASE_POOL_SIZE=10 # LLM Provider (NEAR AI) @@ -7,7 +7,7 @@ DATABASE_POOL_SIZE=10 # Session token is stored in ~/.near-agent/session.json and managed automatically. # On first run, the agent will open a browser for OAuth authentication. NEARAI_MODEL=claude-3-5-sonnet-20241022 -NEARAI_BASE_URL=https://api.near.ai +NEARAI_BASE_URL=https://cloud-api.near.ai NEARAI_AUTH_URL=https://private.near.ai # NEARAI_SESSION_PATH=~/.near-agent/session.json # optional, default shown @@ -28,7 +28,7 @@ HTTP_PORT=8080 HTTP_WEBHOOK_SECRET=your-webhook-secret # Agent Settings -AGENT_NAME=near-agent +AGENT_NAME=ironclaw AGENT_MAX_PARALLEL_JOBS=5 AGENT_JOB_TIMEOUT_SECS=3600 AGENT_STUCK_THRESHOLD_SECS=300 @@ -51,4 +51,4 @@ SAFETY_MAX_OUTPUT_LENGTH=100000 SAFETY_INJECTION_CHECK_ENABLED=true # Logging -RUST_LOG=near_agent=debug,tower_http=debug +RUST_LOG=ironclaw=debug,tower_http=debug diff --git a/channels-src/telegram/src/lib.rs b/channels-src/telegram/src/lib.rs index 146b1e11..984385bc 100644 --- a/channels-src/telegram/src/lib.rs +++ b/channels-src/telegram/src/lib.rs @@ -31,7 +31,7 @@ use serde::{Deserialize, Serialize}; // Re-export generated types use exports::near::agent::channel::{ AgentResponse, ChannelConfig, Guest, HttpEndpointConfig, IncomingHttpRequest, - OutgoingHttpResponse, PollConfig, + OutgoingHttpResponse, PollConfig, StatusType, StatusUpdate, }; use near::agent::channel_host::{self, EmittedMessage}; @@ -157,6 +157,9 @@ struct SentMessage { /// Workspace path for storing polling state. const POLLING_STATE_PATH: &str = "state/last_update_id"; +/// Workspace path for persisting owner_id across WASM callbacks. +const OWNER_ID_PATH: &str = "state/owner_id"; + // ============================================================================ // Channel Metadata // ============================================================================ @@ -188,6 +191,11 @@ struct TelegramConfig { #[serde(default)] bot_username: Option, + /// Telegram user ID of the bot owner. When set, only messages from this + /// user are processed. All others are silently dropped. + #[serde(default)] + owner_id: Option, + /// Whether to respond to all group messages (not just mentions). #[serde(default)] respond_to_all_group_messages: bool, @@ -228,6 +236,27 @@ impl Guest for TelegramChannel { ); } + // Persist owner_id so subsequent callbacks (on_http_request, on_poll) can read it + if let Some(owner_id) = config.owner_id { + if let Err(e) = channel_host::workspace_write(OWNER_ID_PATH, &owner_id.to_string()) { + channel_host::log( + channel_host::LogLevel::Error, + &format!("Failed to persist owner_id: {}", e), + ); + } + channel_host::log( + channel_host::LogLevel::Info, + &format!("Owner restriction enabled: user {}", owner_id), + ); + } else { + // Clear any stale owner_id from a previous config + let _ = channel_host::workspace_write(OWNER_ID_PATH, ""); + channel_host::log( + channel_host::LogLevel::Warn, + "No owner_id configured, bot is open to all users", + ); + } + // Mode is determined by whether the host injected a tunnel_url // If tunnel is configured, use webhooks. Otherwise, use polling. let webhook_mode = config.tunnel_url.is_some(); @@ -501,6 +530,54 @@ impl Guest for TelegramChannel { } } + fn on_status(update: StatusUpdate) { + // Only send typing indicator for Thinking status + if !matches!(update.status, StatusType::Thinking) { + return; + } + + // Parse chat_id from metadata + let metadata: TelegramMessageMetadata = match serde_json::from_str(&update.metadata_json) { + Ok(m) => m, + Err(_) => { + channel_host::log( + channel_host::LogLevel::Debug, + "on_status: no valid Telegram metadata, skipping typing indicator", + ); + return; + } + }; + + // POST /sendChatAction with action "typing" + let payload = serde_json::json!({ + "chat_id": metadata.chat_id, + "action": "typing" + }); + + let payload_bytes = match serde_json::to_vec(&payload) { + Ok(b) => b, + Err(_) => return, + }; + + let headers = serde_json::json!({ + "Content-Type": "application/json" + }); + + let result = channel_host::http_request( + "POST", + "https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/sendChatAction", + &headers.to_string(), + Some(&payload_bytes), + ); + + if let Err(e) = result { + channel_host::log( + channel_host::LogLevel::Debug, + &format!("sendChatAction failed: {}", e), + ); + } + } + fn on_shutdown() { channel_host::log( channel_host::LogLevel::Info, @@ -536,14 +613,15 @@ fn delete_webhook() -> Result<(), String> { return Err(format!("HTTP {}: {}", response.status, body_str)); } - let api_response: TelegramApiResponse = - serde_json::from_slice(&response.body) - .map_err(|e| format!("Failed to parse response: {}", e))?; + let api_response: TelegramApiResponse = serde_json::from_slice(&response.body) + .map_err(|e| format!("Failed to parse response: {}", e))?; if !api_response.ok { return Err(format!( "Telegram API error: {}", - api_response.description.unwrap_or_else(|| "unknown".to_string()) + api_response + .description + .unwrap_or_else(|| "unknown".to_string()) )); } @@ -574,7 +652,8 @@ fn register_webhook(tunnel_url: &str, webhook_secret: Option<&str>) -> Result<() body["secret_token"] = serde_json::Value::String(secret.to_string()); } - let body_bytes = serde_json::to_vec(&body).map_err(|e| format!("Failed to serialize body: {}", e))?; + let body_bytes = + serde_json::to_vec(&body).map_err(|e| format!("Failed to serialize body: {}", e))?; let headers = serde_json::json!({ "Content-Type": "application/json" @@ -657,6 +736,24 @@ fn handle_message(message: TelegramMessage) { return; } + // Owner validation: silently drop messages from non-owner users + if let Some(owner_id_str) = channel_host::workspace_read(OWNER_ID_PATH) { + if !owner_id_str.is_empty() { + if let Ok(owner_id) = owner_id_str.parse::() { + if from.id != owner_id { + channel_host::log( + channel_host::LogLevel::Debug, + &format!( + "Dropping message from non-owner user {} (owner: {})", + from.id, owner_id + ), + ); + return; + } + } + } + } + let is_private = message.chat.chat_type == "private"; // For group chats, check if the bot was mentioned @@ -783,6 +880,40 @@ mod tests { assert_eq!(clean_message_text(" spaced "), "spaced"); } + #[test] + fn test_config_with_owner_id() { + let json = r#"{"owner_id": 123456789}"#; + let config: TelegramConfig = serde_json::from_str(json).unwrap(); + assert_eq!(config.owner_id, Some(123456789)); + } + + #[test] + fn test_config_without_owner_id() { + let json = r#"{}"#; + let config: TelegramConfig = serde_json::from_str(json).unwrap(); + assert_eq!(config.owner_id, None); + } + + #[test] + fn test_config_with_null_owner_id() { + let json = r#"{"owner_id": null}"#; + let config: TelegramConfig = serde_json::from_str(json).unwrap(); + assert_eq!(config.owner_id, None); + } + + #[test] + fn test_config_full() { + let json = r#"{ + "bot_username": "my_bot", + "owner_id": 42, + "respond_to_all_group_messages": true + }"#; + let config: TelegramConfig = serde_json::from_str(json).unwrap(); + assert_eq!(config.bot_username, Some("my_bot".to_string())); + assert_eq!(config.owner_id, Some(42)); + assert!(config.respond_to_all_group_messages); + } + #[test] fn test_parse_update() { let json = r#"{ diff --git a/channels-src/telegram/telegram.capabilities.json b/channels-src/telegram/telegram.capabilities.json index 23a4af65..70f56e01 100644 --- a/channels-src/telegram/telegram.capabilities.json +++ b/channels-src/telegram/telegram.capabilities.json @@ -35,6 +35,7 @@ }, "config": { "bot_username": null, + "owner_id": null, "respond_to_all_group_messages": false, "polling_enabled": false, "poll_interval_ms": 30000 diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index 72397b9a..bd179d56 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -32,13 +32,11 @@ use uuid::Uuid; use crate::agent::compaction::ContextCompactor; use crate::agent::context_monitor::ContextMonitor; use crate::agent::heartbeat::spawn_heartbeat; -use crate::agent::self_repair::DefaultSelfRepair; +use crate::agent::self_repair::{DefaultSelfRepair, RepairResult, SelfRepair}; use crate::agent::session::{PendingApproval, Session, ThreadState}; use crate::agent::session_manager::SessionManager; use crate::agent::submission::{Submission, SubmissionParser, SubmissionResult}; -use crate::agent::{ - HeartbeatConfig as AgentHeartbeatConfig, MessageIntent, RepairTask, Router, Scheduler, -}; +use crate::agent::{HeartbeatConfig as AgentHeartbeatConfig, MessageIntent, Router, Scheduler}; use crate::channels::{ChannelManager, IncomingMessage, OutgoingResponse, StatusUpdate}; use crate::config::{AgentConfig, HeartbeatConfig}; use crate::context::ContextManager; @@ -148,16 +146,98 @@ impl Agent { // Start channels let mut message_stream = self.channels.start_all().await?; - // Start self-repair task + // Start self-repair task with notification forwarding let repair = Arc::new(DefaultSelfRepair::new( self.context_manager.clone(), self.config.stuck_threshold, self.config.max_repair_attempts, )); - let repair_task = RepairTask::new(repair, self.config.repair_check_interval); - + let repair_interval = self.config.repair_check_interval; + let repair_channels = self.channels.clone(); let repair_handle = tokio::spawn(async move { - repair_task.run().await; + loop { + tokio::time::sleep(repair_interval).await; + + // Check stuck jobs + let stuck_jobs = repair.detect_stuck_jobs().await; + for job in stuck_jobs { + tracing::info!("Attempting to repair stuck job {}", job.job_id); + let result = repair.repair_stuck_job(&job).await; + let notification = match &result { + Ok(RepairResult::Success { message }) => { + tracing::info!("Repair succeeded: {}", message); + Some(format!( + "Job {} was stuck for {}s, recovery succeeded: {}", + job.job_id, + job.stuck_duration.as_secs(), + message + )) + } + Ok(RepairResult::Failed { message }) => { + tracing::error!("Repair failed: {}", message); + Some(format!( + "Job {} was stuck for {}s, recovery failed permanently: {}", + job.job_id, + job.stuck_duration.as_secs(), + message + )) + } + Ok(RepairResult::ManualRequired { message }) => { + tracing::warn!("Manual intervention needed: {}", message); + Some(format!( + "Job {} needs manual intervention: {}", + job.job_id, message + )) + } + Ok(RepairResult::Retry { message }) => { + tracing::warn!("Repair needs retry: {}", message); + None // Don't spam the user on retries + } + Err(e) => { + tracing::error!("Repair error: {}", e); + None + } + }; + + if let Some(msg) = notification { + let response = OutgoingResponse::text(format!("Self-Repair: {}", msg)); + let _ = repair_channels.broadcast_all("default", response).await; + } + } + + // Check broken tools + let broken_tools = repair.detect_broken_tools().await; + for tool in broken_tools { + tracing::info!("Attempting to repair broken tool: {}", tool.name); + match repair.repair_broken_tool(&tool).await { + Ok(RepairResult::Success { message }) => { + let response = OutgoingResponse::text(format!( + "Self-Repair: Tool '{}' repaired: {}", + tool.name, message + )); + let _ = repair_channels.broadcast_all("default", response).await; + } + Ok(result) => { + tracing::info!("Tool repair result: {:?}", result); + } + Err(e) => { + tracing::error!("Tool repair error: {}", e); + } + } + } + } + }); + + // Spawn session pruning task + let session_mgr = self.session_manager.clone(); + let session_idle_timeout = self.config.session_idle_timeout; + let pruning_handle = tokio::spawn(async move { + let mut interval = tokio::time::interval(std::time::Duration::from_secs(600)); // Every 10 min + interval.tick().await; // Skip immediate first tick + loop { + interval.tick().await; + session_mgr.prune_stale_sessions(session_idle_timeout).await; + } }); // Spawn heartbeat if enabled @@ -280,6 +360,7 @@ impl Agent { // Cleanup tracing::info!("Agent shutting down..."); repair_handle.abort(); + pruning_handle.abort(); if let Some(handle) = heartbeat_handle { handle.abort(); } @@ -322,6 +403,9 @@ impl Agent { Submission::Compact => self.process_compact(session, thread_id).await, Submission::Clear => self.process_clear(session, thread_id).await, Submission::NewThread => self.process_new_thread(message).await, + Submission::Heartbeat => self.process_heartbeat().await, + Submission::Summarize => self.process_summarize(session, thread_id).await, + Submission::Suggest => self.process_suggest(session, thread_id).await, Submission::SwitchThread { thread_id: target } => { self.process_switch_thread(message, target).await } @@ -480,10 +564,22 @@ impl Agent { let messages = thread.messages(); if let Some(strategy) = self.context_monitor.suggest_compaction(&messages) { - tracing::info!( - "Context at {:.1}% capacity, auto-compacting", - self.context_monitor.usage_percent(&messages) - ); + let pct = self.context_monitor.usage_percent(&messages); + tracing::info!("Context at {:.1}% capacity, auto-compacting", pct); + + // Notify the user that compaction is happening + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::Status(format!( + "Context at {:.0}% capacity, compacting...", + pct + )), + &message.metadata, + ) + .await; + let compactor = ContextCompactor::new(self.llm().clone()); if let Err(e) = compactor .compact(thread, strategy, self.workspace().map(|w| w.as_ref())) @@ -528,6 +624,7 @@ impl Agent { .send_status( &message.channel, StatusUpdate::Thinking("Processing...".into()), + &message.metadata, ) .await; @@ -546,7 +643,11 @@ impl Agent { if thread.state == ThreadState::Interrupted { let _ = self .channels - .send_status(&message.channel, StatusUpdate::Status("Interrupted".into())) + .send_status( + &message.channel, + StatusUpdate::Status("Interrupted".into()), + &message.metadata, + ) .await; return Ok(SubmissionResult::Interrupted); } @@ -557,7 +658,11 @@ impl Agent { thread.complete_turn(&response); let _ = self .channels - .send_status(&message.channel, StatusUpdate::Status("Done".into())) + .send_status( + &message.channel, + StatusUpdate::Status("Done".into()), + &message.metadata, + ) .await; Ok(SubmissionResult::response(response)) } @@ -573,6 +678,7 @@ impl Agent { .send_status( &message.channel, StatusUpdate::Status("Awaiting approval".into()), + &message.metadata, ) .await; Ok(SubmissionResult::NeedApproval { @@ -694,6 +800,7 @@ impl Agent { "Executing {} tool(s)...", tool_calls.len() )), + &message.metadata, ) .await; @@ -851,6 +958,7 @@ impl Agent { .send_status( &message.channel, StatusUpdate::Thinking("Processing...".into()), + &message.metadata, ) .await; } @@ -1163,7 +1271,11 @@ impl Agent { thread.complete_turn(&response); let _ = self .channels - .send_status(&message.channel, StatusUpdate::Status("Done".into())) + .send_status( + &message.channel, + StatusUpdate::Status("Done".into()), + &message.metadata, + ) .await; Ok(SubmissionResult::response(response)) } @@ -1180,6 +1292,7 @@ impl Agent { .send_status( &message.channel, StatusUpdate::Status("Awaiting approval".into()), + &message.metadata, ) .await; Ok(SubmissionResult::NeedApproval { @@ -1205,7 +1318,11 @@ impl Agent { let _ = self .channels - .send_status(&message.channel, StatusUpdate::Status("Rejected".into())) + .send_status( + &message.channel, + StatusUpdate::Status("Rejected".into()), + &message.metadata, + ) .await; Ok(SubmissionResult::response(format!( @@ -1435,6 +1552,134 @@ impl Agent { } } + /// Trigger a manual heartbeat check. + async fn process_heartbeat(&self) -> Result { + let Some(workspace) = self.workspace() else { + return Ok(SubmissionResult::error( + "Heartbeat requires a workspace (database must be connected).", + )); + }; + + let runner = crate::agent::HeartbeatRunner::new( + crate::agent::HeartbeatConfig::default(), + workspace.clone(), + self.llm().clone(), + ); + + match runner.check_heartbeat().await { + crate::agent::HeartbeatResult::Ok => Ok(SubmissionResult::ok_with_message( + "Heartbeat: all clear, nothing needs attention.", + )), + crate::agent::HeartbeatResult::NeedsAttention(msg) => Ok(SubmissionResult::response( + format!("Heartbeat findings:\n\n{}", msg), + )), + crate::agent::HeartbeatResult::Skipped => Ok(SubmissionResult::ok_with_message( + "Heartbeat skipped: no HEARTBEAT.md checklist found in workspace.", + )), + crate::agent::HeartbeatResult::Failed(err) => Ok(SubmissionResult::error(format!( + "Heartbeat failed: {}", + err + ))), + } + } + + /// Summarize the current thread's conversation. + async fn process_summarize( + &self, + session: Arc>, + thread_id: Uuid, + ) -> Result { + let messages = { + let sess = session.lock().await; + let thread = sess + .threads + .get(&thread_id) + .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; + thread.messages() + }; + + if messages.is_empty() { + return Ok(SubmissionResult::ok_with_message( + "Nothing to summarize (empty thread).", + )); + } + + // Build a summary prompt with the conversation + let mut context = Vec::new(); + context.push(ChatMessage::system( + "Summarize the conversation so far in 3-5 concise bullet points. \ + Focus on decisions made, actions taken, and key outcomes. \ + Be brief and factual.", + )); + // Include the conversation messages (truncate to last 20 to avoid context overflow) + let start = if messages.len() > 20 { + messages.len() - 20 + } else { + 0 + }; + context.extend_from_slice(&messages[start..]); + context.push(ChatMessage::user("Summarize this conversation.")); + + let request = crate::llm::CompletionRequest::new(context) + .with_max_tokens(512) + .with_temperature(0.3); + + match self.llm().complete(request).await { + Ok(response) => Ok(SubmissionResult::response(format!( + "Thread Summary:\n\n{}", + response.content.trim() + ))), + Err(e) => Ok(SubmissionResult::error(format!("Summarize failed: {}", e))), + } + } + + /// Suggest next steps based on the current thread. + async fn process_suggest( + &self, + session: Arc>, + thread_id: Uuid, + ) -> Result { + let messages = { + let sess = session.lock().await; + let thread = sess + .threads + .get(&thread_id) + .ok_or_else(|| Error::from(crate::error::JobError::NotFound { id: thread_id }))?; + thread.messages() + }; + + if messages.is_empty() { + return Ok(SubmissionResult::ok_with_message( + "Nothing to suggest from (empty thread).", + )); + } + + let mut context = Vec::new(); + context.push(ChatMessage::system( + "Based on the conversation so far, suggest 2-4 concrete next steps the user could take. \ + Be actionable and specific. Format as a numbered list.", + )); + let start = if messages.len() > 20 { + messages.len() - 20 + } else { + 0 + }; + context.extend_from_slice(&messages[start..]); + context.push(ChatMessage::user("What should I do next?")); + + let request = crate::llm::CompletionRequest::new(context) + .with_max_tokens(512) + .with_temperature(0.5); + + match self.llm().complete(request).await { + Ok(response) => Ok(SubmissionResult::response(format!( + "Suggested Next Steps:\n\n{}", + response.content.trim() + ))), + Err(e) => Ok(SubmissionResult::error(format!("Suggest failed: {}", e))), + } + } + async fn handle_command( &self, command: &str, @@ -1458,6 +1703,10 @@ impl Agent { /thread - Switch thread /resume - Resume checkpoint + /heartbeat - Run heartbeat check now + /summarize - Summarize current thread + /suggest - Suggest next steps + /quit - Exit"# .to_string(), )), diff --git a/src/agent/session_manager.rs b/src/agent/session_manager.rs index 69787c9e..dcf46343 100644 --- a/src/agent/session_manager.rs +++ b/src/agent/session_manager.rs @@ -131,6 +131,81 @@ impl SessionManager { managers.insert(thread_id, Arc::clone(&mgr)); mgr } + + /// Remove sessions that have been idle for longer than the given duration. + /// + /// Returns the number of sessions pruned. + pub async fn prune_stale_sessions(&self, max_idle: std::time::Duration) -> usize { + let cutoff = chrono::Utc::now() - chrono::TimeDelta::seconds(max_idle.as_secs() as i64); + + // Find stale session user_ids + let stale_users: Vec = { + let sessions = self.sessions.read().await; + sessions + .iter() + .filter_map(|(user_id, session)| { + // Try to lock; skip if contended (someone is actively using it) + let sess = session.try_lock().ok()?; + if sess.last_active_at < cutoff { + Some(user_id.clone()) + } else { + None + } + }) + .collect() + }; + + if stale_users.is_empty() { + return 0; + } + + // Collect thread IDs from stale sessions for cleanup + let mut stale_thread_ids: Vec = Vec::new(); + { + let sessions = self.sessions.read().await; + for user_id in &stale_users { + if let Some(session) = sessions.get(user_id) { + if let Ok(sess) = session.try_lock() { + stale_thread_ids.extend(sess.threads.keys()); + } + } + } + } + + // Remove sessions + let count = { + let mut sessions = self.sessions.write().await; + let before = sessions.len(); + for user_id in &stale_users { + sessions.remove(user_id); + } + before - sessions.len() + }; + + // Clean up thread mappings that point to stale sessions + { + let mut thread_map = self.thread_map.write().await; + thread_map.retain(|key, _| !stale_users.contains(&key.user_id)); + } + + // Clean up undo managers for stale threads + { + let mut undo_managers = self.undo_managers.write().await; + for thread_id in &stale_thread_ids { + undo_managers.remove(thread_id); + } + } + + if count > 0 { + tracing::info!( + "Pruned {} stale session(s) (idle > {}s)", + count, + max_idle.as_secs() + ); + } + + count + } } impl Default for SessionManager { @@ -183,4 +258,42 @@ mod tests { assert!(Arc::ptr_eq(&undo1, &undo2)); } + + #[tokio::test] + async fn test_prune_stale_sessions() { + let manager = SessionManager::new(); + + // Create two sessions and resolve threads (which updates last_active_at) + let (_, _thread_id) = manager.resolve_thread("user-active", "cli", None).await; + let (s2, _thread_id) = manager.resolve_thread("user-stale", "cli", None).await; + + // Backdate the stale session's last_active_at AFTER thread creation + { + let mut sess = s2.lock().await; + sess.last_active_at = chrono::Utc::now() - chrono::TimeDelta::seconds(86400 * 10); // 10 days ago + } + + // Prune with 7-day timeout + let pruned = manager + .prune_stale_sessions(std::time::Duration::from_secs(86400 * 7)) + .await; + assert_eq!(pruned, 1); + + // Active session should still exist + let sessions = manager.sessions.read().await; + assert!(sessions.contains_key("user-active")); + assert!(!sessions.contains_key("user-stale")); + } + + #[tokio::test] + async fn test_prune_no_stale_sessions() { + let manager = SessionManager::new(); + let _s1 = manager.get_or_create_session("user-1").await; + + // Nothing should be pruned when timeout is long + let pruned = manager + .prune_stale_sessions(std::time::Duration::from_secs(86400 * 365)) + .await; + assert_eq!(pruned, 0); + } } diff --git a/src/agent/submission.rs b/src/agent/submission.rs index f79c0b03..f1c36d68 100644 --- a/src/agent/submission.rs +++ b/src/agent/submission.rs @@ -31,6 +31,15 @@ impl SubmissionParser { if lower == "/clear" { return Submission::Clear; } + if lower == "/heartbeat" { + return Submission::Heartbeat; + } + if lower == "/summarize" || lower == "/summary" { + return Submission::Summarize; + } + if lower == "/suggest" { + return Submission::Suggest; + } if lower == "/thread new" || lower == "/new" { return Submission::NewThread; } @@ -139,6 +148,15 @@ pub enum Submission { /// Create a new thread. NewThread, + + /// Trigger a manual heartbeat check. + Heartbeat, + + /// Summarize the current thread. + Summarize, + + /// Suggest next steps based on the current thread. + Suggest, } impl Submission { @@ -202,6 +220,9 @@ impl Submission { | Self::Redo | Self::Clear | Self::NewThread + | Self::Heartbeat + | Self::Summarize + | Self::Suggest ) } } @@ -355,6 +376,27 @@ mod tests { ); } + #[test] + fn test_parser_heartbeat() { + let submission = SubmissionParser::parse("/heartbeat"); + assert!(matches!(submission, Submission::Heartbeat)); + } + + #[test] + fn test_parser_summarize() { + let submission = SubmissionParser::parse("/summarize"); + assert!(matches!(submission, Submission::Summarize)); + + let submission = SubmissionParser::parse("/summary"); + assert!(matches!(submission, Submission::Summarize)); + } + + #[test] + fn test_parser_suggest() { + let submission = SubmissionParser::parse("/suggest"); + assert!(matches!(submission, Submission::Suggest)); + } + #[test] fn test_parser_invalid_commands_become_user_input() { // Invalid UUID should become user input diff --git a/src/channels/channel.rs b/src/channels/channel.rs index 68ec714d..c94140a4 100644 --- a/src/channels/channel.rs +++ b/src/channels/channel.rs @@ -141,8 +141,15 @@ pub trait Channel: Send + Sync { /// Send a status update (thinking, tool execution, etc.). /// + /// The metadata contains channel-specific routing info (e.g., Telegram chat_id) + /// needed to deliver the status to the correct destination. + /// /// Default implementation does nothing (for channels that don't support status). - async fn send_status(&self, _status: StatusUpdate) -> Result<(), ChannelError> { + async fn send_status( + &self, + _status: StatusUpdate, + _metadata: &serde_json::Value, + ) -> Result<(), ChannelError> { Ok(()) } diff --git a/src/channels/cli/mod.rs b/src/channels/cli/mod.rs index dfb60adf..afeb590e 100644 --- a/src/channels/cli/mod.rs +++ b/src/channels/cli/mod.rs @@ -116,7 +116,11 @@ impl Channel for TuiChannel { Ok(()) } - async fn send_status(&self, status: StatusUpdate) -> Result<(), ChannelError> { + async fn send_status( + &self, + status: StatusUpdate, + _metadata: &serde_json::Value, + ) -> Result<(), ChannelError> { let event = match status { StatusUpdate::Thinking(msg) => AppEvent::ThinkingMessage(format!("πŸ€” {}", msg)), StatusUpdate::ToolStarted { name } => AppEvent::ToolStarted { name }, diff --git a/src/channels/manager.rs b/src/channels/manager.rs index 6d8f230c..6d9b99a6 100644 --- a/src/channels/manager.rs +++ b/src/channels/manager.rs @@ -83,14 +83,18 @@ impl ChannelManager { } /// Send a status update to a specific channel. + /// + /// The metadata contains channel-specific routing info (e.g., Telegram chat_id) + /// needed to deliver the status to the correct destination. pub async fn send_status( &self, channel_name: &str, status: StatusUpdate, + metadata: &serde_json::Value, ) -> Result<(), ChannelError> { let channels = self.channels.read().await; if let Some(channel) = channels.get(channel_name) { - channel.send_status(status).await + channel.send_status(status, metadata).await } else { // Silently ignore if channel not found (status is best-effort) Ok(()) diff --git a/src/channels/repl.rs b/src/channels/repl.rs index 1046460e..a498ff92 100644 --- a/src/channels/repl.rs +++ b/src/channels/repl.rs @@ -183,7 +183,11 @@ impl Channel for ReplChannel { Ok(()) } - async fn send_status(&self, status: StatusUpdate) -> Result<(), ChannelError> { + async fn send_status( + &self, + status: StatusUpdate, + _metadata: &serde_json::Value, + ) -> Result<(), ChannelError> { let debug = self.is_debug(); match status { 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/channels/wasm/wrapper.rs b/src/channels/wasm/wrapper.rs index b0927616..550e6932 100644 --- a/src/channels/wasm/wrapper.rs +++ b/src/channels/wasm/wrapper.rs @@ -420,6 +420,10 @@ pub struct WasmChannel { /// Keys are placeholder names like "TELEGRAM_BOT_TOKEN". /// Wrapped in Arc for sharing with the polling task. credentials: Arc>>, + + /// Background task that repeats typing indicators every 4 seconds. + /// Telegram's "typing..." indicator expires after ~5s, so we refresh it. + typing_task: RwLock>>, } impl WasmChannel { @@ -447,6 +451,7 @@ impl WasmChannel { poll_shutdown_tx: RwLock::new(None), endpoints: RwLock::new(Vec::new()), credentials: Arc::new(RwLock::new(HashMap::new())), + typing_task: RwLock::new(None), } } @@ -1012,6 +1017,214 @@ impl WasmChannel { } } + /// Execute the on_status callback. + /// + /// Called to notify the WASM channel of agent status changes (e.g., typing). + pub async fn call_on_status( + &self, + status: &StatusUpdate, + metadata: &serde_json::Value, + ) -> Result<(), WasmChannelError> { + // If no WASM bytes, do nothing (for testing) + if self.prepared.component_bytes.is_empty() { + return Ok(()); + } + + let runtime = Arc::clone(&self.runtime); + let prepared = Arc::clone(&self.prepared); + let capabilities = self.capabilities.clone(); + let timeout = self.runtime.config().callback_timeout; + let channel_name = self.name.clone(); + let credentials = self.get_credentials().await; + + let wit_update = status_to_wit(status, metadata); + + let result = tokio::time::timeout(timeout, async move { + tokio::task::spawn_blocking(move || { + let mut store = + Self::create_store(&runtime, &prepared, &capabilities, credentials)?; + let instance = Self::instantiate_component(&runtime, &prepared, &mut store)?; + + let channel_iface = instance.near_agent_channel(); + channel_iface + .call_on_status(&mut store, &wit_update) + .map_err(|e| Self::map_wasm_error(e, &prepared.name, prepared.limits.fuel))?; + + Ok(()) + }) + .await + .map_err(|e| WasmChannelError::ExecutionPanicked { + name: channel_name.clone(), + reason: e.to_string(), + })? + }) + .await; + + match result { + Ok(Ok(())) => { + tracing::debug!( + channel = %self.name, + "WASM channel on_status completed" + ); + Ok(()) + } + Ok(Err(e)) => Err(e), + Err(_) => Err(WasmChannelError::Timeout { + name: self.name.clone(), + callback: "on_status".to_string(), + }), + } + } + + /// Execute a single on_status callback with a fresh WASM instance. + /// + /// Static method for use by the background typing repeat task (which + /// doesn't have access to `&self`). + async fn execute_status( + channel_name: &str, + runtime: &Arc, + prepared: &Arc, + capabilities: &ChannelCapabilities, + credentials: &RwLock>, + timeout: Duration, + wit_update: wit_channel::StatusUpdate, + ) -> Result<(), WasmChannelError> { + if prepared.component_bytes.is_empty() { + return Ok(()); + } + + let runtime = Arc::clone(runtime); + let prepared = Arc::clone(prepared); + let capabilities = capabilities.clone(); + let credentials_snapshot = credentials.read().await.clone(); + let channel_name_owned = channel_name.to_string(); + + let result = tokio::time::timeout(timeout, async move { + tokio::task::spawn_blocking(move || { + let mut store = + Self::create_store(&runtime, &prepared, &capabilities, credentials_snapshot)?; + let instance = Self::instantiate_component(&runtime, &prepared, &mut store)?; + + let channel_iface = instance.near_agent_channel(); + channel_iface + .call_on_status(&mut store, &wit_update) + .map_err(|e| Self::map_wasm_error(e, &prepared.name, prepared.limits.fuel))?; + + Ok(()) + }) + .await + .map_err(|e| WasmChannelError::ExecutionPanicked { + name: channel_name_owned.clone(), + reason: e.to_string(), + })? + }) + .await; + + match result { + Ok(Ok(())) => Ok(()), + Ok(Err(e)) => Err(e), + Err(_) => Err(WasmChannelError::Timeout { + name: channel_name.to_string(), + callback: "on_status".to_string(), + }), + } + } + + /// Cancel the background typing indicator task if running. + async fn cancel_typing_task(&self) { + if let Some(handle) = self.typing_task.write().await.take() { + handle.abort(); + } + } + + /// Handle a status update, managing the typing repeat timer. + /// + /// On Thinking: fires on_status once, then spawns a background task + /// that repeats the call every 4 seconds (Telegram's typing indicator + /// expires after ~5s). + /// + /// On Done/Interrupted/Status: cancels the repeat task, fires on_status once. + /// On StreamChunk: no-op (too noisy). + async fn handle_status_update( + &self, + status: StatusUpdate, + metadata: &serde_json::Value, + ) -> Result<(), ChannelError> { + match &status { + StatusUpdate::Thinking(_) => { + // Cancel any existing typing task + self.cancel_typing_task().await; + + // Fire once immediately + if let Err(e) = self.call_on_status(&status, metadata).await { + tracing::debug!( + channel = %self.name, + error = %e, + "on_status(Thinking) failed (best-effort)" + ); + } + + // Spawn background repeater + let channel_name = self.name.clone(); + let runtime = Arc::clone(&self.runtime); + let prepared = Arc::clone(&self.prepared); + let capabilities = self.capabilities.clone(); + let credentials = self.credentials.clone(); + let callback_timeout = self.runtime.config().callback_timeout; + let wit_update = status_to_wit(&status, metadata); + + let handle = tokio::spawn(async move { + let mut interval = tokio::time::interval(Duration::from_secs(4)); + // Skip the first tick (we already fired above) + interval.tick().await; + + loop { + interval.tick().await; + + let wit_update_clone = clone_wit_status_update(&wit_update); + + if let Err(e) = Self::execute_status( + &channel_name, + &runtime, + &prepared, + &capabilities, + &credentials, + callback_timeout, + wit_update_clone, + ) + .await + { + tracing::debug!( + channel = %channel_name, + error = %e, + "Typing repeat on_status failed (best-effort)" + ); + } + } + }); + + *self.typing_task.write().await = Some(handle); + } + StatusUpdate::StreamChunk(_) => { + // No-op, too noisy + } + _ => { + // Done, Interrupted, Status, ToolStarted, ToolCompleted: cancel and fire once + self.cancel_typing_task().await; + + if let Err(e) = self.call_on_status(&status, metadata).await { + tracing::debug!( + channel = %self.name, + error = %e, + "on_status failed (best-effort)" + ); + } + } + } + + Ok(()) + } + /// Process emitted messages from a callback. async fn process_emitted_messages( &self, @@ -1403,6 +1616,9 @@ impl Channel for WasmChannel { msg: &IncomingMessage, response: OutgoingResponse, ) -> Result<(), ChannelError> { + // Stop the typing indicator, we're about to send the actual response + self.cancel_typing_task().await; + // Check if there's a pending synchronous response waiter if let Some(tx) = self.pending_responses.write().await.remove(&msg.id) { let _ = tx.send(response.content.clone()); @@ -1428,11 +1644,13 @@ impl Channel for WasmChannel { Ok(()) } - async fn send_status(&self, status: StatusUpdate) -> Result<(), ChannelError> { - // WASM channels don't support status updates by default - // Could be extended with an optional on_status callback - let _ = status; - Ok(()) + async fn send_status( + &self, + status: StatusUpdate, + metadata: &serde_json::Value, + ) -> Result<(), ChannelError> { + // Delegate to the typing indicator implementation + self.handle_status_update(status, metadata).await } async fn health_check(&self) -> Result<(), ChannelError> { @@ -1447,6 +1665,9 @@ impl Channel for WasmChannel { } async fn shutdown(&self) -> Result<(), ChannelError> { + // Cancel typing indicator + self.cancel_typing_task().await; + // Send shutdown signal if let Some(tx) = self.shutdown_tx.write().await.take() { let _ = tx.send(()); @@ -1458,8 +1679,6 @@ impl Channel for WasmChannel { // Clear the message sender *self.message_tx.write().await = None; - // TODO: Call WASM on_shutdown if we add that callback - tracing::info!( channel = %self.name, "WASM channel shut down" @@ -1529,6 +1748,14 @@ impl Channel for SharedWasmChannel { self.inner.respond(msg, response).await } + async fn send_status( + &self, + status: StatusUpdate, + metadata: &serde_json::Value, + ) -> Result<(), ChannelError> { + self.inner.send_status(status, metadata).await + } + async fn health_check(&self) -> Result<(), ChannelError> { self.inner.health_check().await } @@ -1579,6 +1806,62 @@ fn convert_http_response(wit: wit_channel::OutgoingHttpResponse) -> HttpResponse } } +/// Convert a StatusUpdate + metadata into the WIT StatusUpdate type. +fn status_to_wit(status: &StatusUpdate, metadata: &serde_json::Value) -> wit_channel::StatusUpdate { + let metadata_json = serde_json::to_string(metadata).unwrap_or_default(); + + match status { + StatusUpdate::Thinking(msg) => wit_channel::StatusUpdate { + status: wit_channel::StatusType::Thinking, + message: msg.clone(), + metadata_json, + }, + StatusUpdate::ToolStarted { name } => wit_channel::StatusUpdate { + status: wit_channel::StatusType::ToolStarted, + message: name.clone(), + metadata_json, + }, + StatusUpdate::ToolCompleted { name, success } => wit_channel::StatusUpdate { + status: wit_channel::StatusType::ToolCompleted, + message: format!("{}: {}", name, if *success { "ok" } else { "failed" }), + metadata_json, + }, + StatusUpdate::StreamChunk(chunk) => wit_channel::StatusUpdate { + status: wit_channel::StatusType::Thinking, + message: chunk.clone(), + metadata_json, + }, + StatusUpdate::Status(msg) => { + // Map well-known status strings to WIT types + let status_type = match msg.as_str() { + "Done" => wit_channel::StatusType::Done, + "Interrupted" => wit_channel::StatusType::Interrupted, + _ => wit_channel::StatusType::Thinking, + }; + wit_channel::StatusUpdate { + status: status_type, + message: msg.clone(), + metadata_json, + } + } + } +} + +/// Clone a WIT StatusUpdate (the generated type doesn't derive Clone). +fn clone_wit_status_update(update: &wit_channel::StatusUpdate) -> wit_channel::StatusUpdate { + wit_channel::StatusUpdate { + status: match update.status { + wit_channel::StatusType::Thinking => wit_channel::StatusType::Thinking, + wit_channel::StatusType::Done => wit_channel::StatusType::Done, + wit_channel::StatusType::Interrupted => wit_channel::StatusType::Interrupted, + wit_channel::StatusType::ToolStarted => wit_channel::StatusType::ToolStarted, + wit_channel::StatusType::ToolCompleted => wit_channel::StatusType::ToolCompleted, + }, + message: update.message.clone(), + metadata_json: update.metadata_json.clone(), + } +} + /// HTTP response from a WASM channel callback. #[derive(Debug, Clone)] pub struct HttpResponse { @@ -1844,4 +2127,213 @@ mod tests { channel.shutdown().await.expect("Shutdown should succeed"); } + + #[tokio::test] + async fn test_typing_task_starts_on_thinking() { + let channel = create_test_channel(); + let _stream = channel.start().await.expect("Channel should start"); + + let metadata = serde_json::json!({"chat_id": 123}); + + // Sending Thinking should succeed (no-op for no WASM) + let result = channel + .send_status( + crate::channels::StatusUpdate::Thinking("Processing...".into()), + &metadata, + ) + .await; + assert!(result.is_ok()); + + // A typing task should have been spawned + assert!(channel.typing_task.read().await.is_some()); + + // Shutdown should cancel the typing task + channel.shutdown().await.expect("Shutdown should succeed"); + assert!(channel.typing_task.read().await.is_none()); + } + + #[tokio::test] + async fn test_typing_task_cancelled_on_done() { + let channel = create_test_channel(); + let _stream = channel.start().await.expect("Channel should start"); + + let metadata = serde_json::json!({"chat_id": 123}); + + // Start typing + let _ = channel + .send_status( + crate::channels::StatusUpdate::Thinking("Processing...".into()), + &metadata, + ) + .await; + assert!(channel.typing_task.read().await.is_some()); + + // Send Done status + let _ = channel + .send_status( + crate::channels::StatusUpdate::Status("Done".into()), + &metadata, + ) + .await; + + // Typing task should be cancelled + assert!(channel.typing_task.read().await.is_none()); + + channel.shutdown().await.expect("Shutdown should succeed"); + } + + #[tokio::test] + async fn test_typing_task_replaced_on_new_thinking() { + let channel = create_test_channel(); + let _stream = channel.start().await.expect("Channel should start"); + + let metadata = serde_json::json!({"chat_id": 123}); + + // Start typing + let _ = channel + .send_status( + crate::channels::StatusUpdate::Thinking("First...".into()), + &metadata, + ) + .await; + + // Get handle of first task + let first_handle = { + let guard = channel.typing_task.read().await; + guard.as_ref().map(|h| h.id()) + }; + assert!(first_handle.is_some()); + + // Start typing again (should replace the previous task) + let _ = channel + .send_status( + crate::channels::StatusUpdate::Thinking("Second...".into()), + &metadata, + ) + .await; + + // Should still have a typing task, but it's a new one + let second_handle = { + let guard = channel.typing_task.read().await; + guard.as_ref().map(|h| h.id()) + }; + assert!(second_handle.is_some()); + // The task IDs should differ (old one was aborted, new one spawned) + assert_ne!(first_handle, second_handle); + + channel.shutdown().await.expect("Shutdown should succeed"); + } + + #[tokio::test] + async fn test_respond_cancels_typing_task() { + use crate::channels::IncomingMessage; + + let channel = create_test_channel(); + let _stream = channel.start().await.expect("Channel should start"); + + let metadata = serde_json::json!({"chat_id": 123}); + + // Start typing + let _ = channel + .send_status( + crate::channels::StatusUpdate::Thinking("Processing...".into()), + &metadata, + ) + .await; + assert!(channel.typing_task.read().await.is_some()); + + // Respond should cancel the typing task + let msg = IncomingMessage::new("test", "user1", "hello").with_metadata(metadata); + let _ = channel + .respond(&msg, crate::channels::OutgoingResponse::text("response")) + .await; + + // Typing task should be gone + assert!(channel.typing_task.read().await.is_none()); + + channel.shutdown().await.expect("Shutdown should succeed"); + } + + #[tokio::test] + async fn test_stream_chunk_is_noop() { + let channel = create_test_channel(); + let _stream = channel.start().await.expect("Channel should start"); + + let metadata = serde_json::json!({"chat_id": 123}); + + // StreamChunk should not start a typing task + let result = channel + .send_status( + crate::channels::StatusUpdate::StreamChunk("chunk".into()), + &metadata, + ) + .await; + assert!(result.is_ok()); + assert!(channel.typing_task.read().await.is_none()); + + channel.shutdown().await.expect("Shutdown should succeed"); + } + + #[test] + fn test_status_to_wit_thinking() { + use super::status_to_wit; + + let metadata = serde_json::json!({"chat_id": 42}); + let wit = status_to_wit( + &crate::channels::StatusUpdate::Thinking("Processing...".into()), + &metadata, + ); + + assert!(matches!( + wit.status, + super::wit_channel::StatusType::Thinking + )); + assert_eq!(wit.message, "Processing..."); + assert!(wit.metadata_json.contains("42")); + } + + #[test] + fn test_status_to_wit_done() { + use super::status_to_wit; + + let metadata = serde_json::json!(null); + let wit = status_to_wit( + &crate::channels::StatusUpdate::Status("Done".into()), + &metadata, + ); + + assert!(matches!(wit.status, super::wit_channel::StatusType::Done)); + } + + #[test] + fn test_status_to_wit_interrupted() { + use super::status_to_wit; + + let metadata = serde_json::json!(null); + let wit = status_to_wit( + &crate::channels::StatusUpdate::Status("Interrupted".into()), + &metadata, + ); + + assert!(matches!( + wit.status, + super::wit_channel::StatusType::Interrupted + )); + } + + #[test] + fn test_clone_wit_status_update() { + use super::{clone_wit_status_update, wit_channel}; + + let original = wit_channel::StatusUpdate { + status: wit_channel::StatusType::Thinking, + message: "hello".to_string(), + metadata_json: "{\"a\":1}".to_string(), + }; + + let cloned = clone_wit_status_update(&original); + assert!(matches!(cloned.status, wit_channel::StatusType::Thinking)); + assert_eq!(cloned.message, "hello"); + assert_eq!(cloned.metadata_json, "{\"a\":1}"); + } } diff --git a/src/cli/memory.rs b/src/cli/memory.rs new file mode 100644 index 00000000..e10cf44f --- /dev/null +++ b/src/cli/memory.rs @@ -0,0 +1,268 @@ +//! Memory/workspace CLI commands. +//! +//! Exposes the workspace system for direct CLI use without starting the agent. + +use std::io::Read; +use std::sync::Arc; + +use clap::Subcommand; + +use crate::workspace::{EmbeddingProvider, SearchConfig, Workspace}; + +#[derive(Subcommand, Debug, Clone)] +pub enum MemoryCommand { + /// Search workspace memory (hybrid full-text + semantic) + Search { + /// Search query + query: String, + + /// Maximum number of results + #[arg(short, long, default_value = "5")] + limit: usize, + }, + + /// Read a file from the workspace + Read { + /// File path (e.g., "MEMORY.md", "daily/2024-01-15.md") + path: String, + }, + + /// Write content to a workspace file + Write { + /// File path (e.g., "notes/idea.md") + path: String, + + /// Content to write (omit to read from stdin) + content: Option, + + /// Append instead of overwrite + #[arg(short, long)] + append: bool, + }, + + /// Show workspace directory tree + Tree { + /// Root path to start from + #[arg(default_value = "")] + path: String, + + /// Maximum depth to traverse + #[arg(short, long, default_value = "3")] + depth: usize, + }, + + /// Show workspace status (document count, index health) + Status, +} + +/// Run a memory command. +pub async fn run_memory_command( + cmd: MemoryCommand, + pool: deadpool_postgres::Pool, + embeddings: Option>, +) -> anyhow::Result<()> { + let mut workspace = Workspace::new("default", pool); + if let Some(emb) = embeddings { + workspace = workspace.with_embeddings(emb); + } + + match cmd { + MemoryCommand::Search { query, limit } => search(&workspace, &query, limit).await, + MemoryCommand::Read { path } => read(&workspace, &path).await, + MemoryCommand::Write { + path, + content, + append, + } => write(&workspace, &path, content, append).await, + MemoryCommand::Tree { path, depth } => tree(&workspace, &path, depth).await, + MemoryCommand::Status => status(&workspace).await, + } +} + +async fn search(workspace: &Workspace, query: &str, limit: usize) -> anyhow::Result<()> { + let config = SearchConfig::default().with_limit(limit.min(50)); + let results = workspace.search_with_config(query, config).await?; + + if results.is_empty() { + println!("No results found for: {}", query); + return Ok(()); + } + + println!("Found {} result(s) for \"{}\":\n", results.len(), query); + + for (i, result) in results.iter().enumerate() { + let score_bar = score_indicator(result.score); + println!("{}. [{}] (score: {:.3})", i + 1, score_bar, result.score); + + // Show a content preview (first 200 chars) + let preview = truncate_content(&result.content, 200); + for line in preview.lines() { + println!(" {}", line); + } + println!(); + } + + Ok(()) +} + +async fn read(workspace: &Workspace, path: &str) -> anyhow::Result<()> { + match workspace.read(path).await { + Ok(doc) => { + println!("{}", doc.content); + } + Err(crate::error::WorkspaceError::DocumentNotFound { .. }) => { + anyhow::bail!("File not found: {}", path); + } + Err(e) => return Err(e.into()), + } + Ok(()) +} + +async fn write( + workspace: &Workspace, + path: &str, + content: Option, + append: bool, +) -> anyhow::Result<()> { + let content = match content { + Some(c) => c, + None => { + // Read from stdin + let mut buf = String::new(); + std::io::stdin().read_to_string(&mut buf)?; + buf + } + }; + + if append { + workspace.append(path, &content).await?; + println!("Appended to {}", path); + } else { + workspace.write(path, &content).await?; + println!("Wrote to {}", path); + } + + Ok(()) +} + +async fn tree(workspace: &Workspace, path: &str, max_depth: usize) -> anyhow::Result<()> { + let root = if path.is_empty() { "." } else { path }; + println!("{}/", root); + print_tree(workspace, path, "", max_depth, 0).await?; + Ok(()) +} + +async fn print_tree( + workspace: &Workspace, + path: &str, + prefix: &str, + max_depth: usize, + current_depth: usize, +) -> anyhow::Result<()> { + if current_depth >= max_depth { + return Ok(()); + } + + let entries = workspace.list(path).await?; + let count = entries.len(); + + for (i, entry) in entries.iter().enumerate() { + let is_last = i == count - 1; + let connector = if is_last { "└── " } else { "β”œβ”€β”€ " }; + let child_prefix = if is_last { " " } else { "β”‚ " }; + + if entry.is_directory { + println!("{}{}{}/", prefix, connector, entry.name()); + Box::pin(print_tree( + workspace, + &entry.path, + &format!("{}{}", prefix, child_prefix), + max_depth, + current_depth + 1, + )) + .await?; + } else { + println!("{}{}{}", prefix, connector, entry.name()); + } + } + + Ok(()) +} + +async fn status(workspace: &Workspace) -> anyhow::Result<()> { + let all_paths = workspace.list_all().await?; + let file_count = all_paths.len(); + + // Count directories by collecting unique parent paths + let mut dirs: std::collections::HashSet = std::collections::HashSet::new(); + for path in &all_paths { + if let Some(parent) = path.rsplit_once('/') { + dirs.insert(parent.0.to_string()); + } + } + + println!("Workspace Status"); + println!(" User: {}", workspace.user_id()); + println!(" Files: {}", file_count); + println!(" Directories: {}", dirs.len()); + + // Check key files + let key_files = [ + "MEMORY.md", + "HEARTBEAT.md", + "IDENTITY.md", + "SOUL.md", + "AGENTS.md", + "USER.md", + ]; + println!("\n Identity files:"); + for path in &key_files { + let exists = workspace.exists(path).await.unwrap_or(false); + let marker = if exists { "+" } else { "-" }; + println!(" [{}] {}", marker, path); + } + + Ok(()) +} + +fn truncate_content(s: &str, max_len: usize) -> String { + if s.len() <= max_len { + s.to_string() + } else { + format!("{}...", &s[..max_len]) + } +} + +fn score_indicator(score: f32) -> &'static str { + if score > 0.8_f32 { + "=====>" + } else if score > 0.5_f32 { + "====>" + } else if score > 0.3_f32 { + "===>" + } else if score > 0.1_f32 { + "==>" + } else { + "=>" + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_score_indicator() { + assert_eq!(score_indicator(0.9_f32), "=====>"); + assert_eq!(score_indicator(0.6_f32), "====>"); + assert_eq!(score_indicator(0.4_f32), "===>"); + assert_eq!(score_indicator(0.2_f32), "==>"); + assert_eq!(score_indicator(0.05_f32), "=>"); + } + + #[test] + fn test_truncate_content() { + assert_eq!(truncate_content("hello", 10), "hello"); + assert_eq!(truncate_content("hello world", 5), "hello..."); + } +} diff --git a/src/cli/mod.rs b/src/cli/mod.rs index 4d863059..fa3e9db1 100644 --- a/src/cli/mod.rs +++ b/src/cli/mod.rs @@ -6,13 +6,19 @@ //! - Managing configuration (`config list`, `config get`, `config set`) //! - Managing WASM tools (`tool install`, `tool list`, `tool remove`) //! - Managing MCP servers (`mcp add`, `mcp auth`, `mcp list`, `mcp test`) +//! - Querying workspace memory (`memory search`, `memory read`, `memory write`) +//! - Checking system health (`status`) mod config; mod mcp; +pub mod memory; +pub mod status; mod tool; pub use config::{ConfigCommand, run_config_command}; pub use mcp::{McpCommand, run_mcp_command}; +pub use memory::{MemoryCommand, run_memory_command}; +pub use status::run_status_command; pub use tool::{ToolCommand, run_tool_command}; use clap::{Parser, Subcommand}; @@ -79,6 +85,13 @@ pub enum Command { /// Manage MCP servers (hosted tool providers) #[command(subcommand)] Mcp(McpCommand), + + /// Query and manage workspace memory + #[command(subcommand)] + Memory(MemoryCommand), + + /// Show system health and diagnostics + Status, } impl Cli { diff --git a/src/cli/status.rs b/src/cli/status.rs new file mode 100644 index 00000000..0b4656c3 --- /dev/null +++ b/src/cli/status.rs @@ -0,0 +1,193 @@ +//! System health and diagnostics CLI command. +//! +//! Checks database connectivity, session validity, embeddings, +//! WASM runtime, tool count, and channel availability. + +use std::path::PathBuf; + +use crate::settings::Settings; + +/// Run the status command, printing system health info. +pub async fn run_status_command() -> anyhow::Result<()> { + let settings = Settings::load(); + + println!("IronClaw Status"); + println!("===============\n"); + + // Version + println!( + " Version: {} v{}", + env!("CARGO_PKG_NAME"), + env!("CARGO_PKG_VERSION") + ); + + // Database + let db_url_set = settings.database_url.is_some() || std::env::var("DATABASE_URL").is_ok(); + print!(" Database: "); + if db_url_set { + // Try to connect + match check_database().await { + Ok(()) => println!("connected"), + Err(e) => println!("error ({})", e), + } + } else { + println!("not configured"); + } + + // Session / Auth + print!(" Session: "); + let session_path = crate::llm::session::default_session_path(); + if session_path.exists() { + println!("found ({})", session_path.display()); + } else { + println!("not found (run `ironclaw setup`)"); + } + + // Secrets + print!(" Secrets: "); + let secrets_configured = settings.secrets_master_key_source != crate::settings::KeySource::None + || std::env::var("SECRETS_MASTER_KEY").is_ok() + || crate::secrets::keychain::has_master_key(); + if secrets_configured { + println!("configured ({:?})", settings.secrets_master_key_source); + } else { + println!("not configured"); + } + + // Embeddings + print!(" Embeddings: "); + let emb_enabled = settings.embeddings.enabled + || std::env::var("OPENAI_API_KEY").is_ok() + || std::env::var("EMBEDDING_ENABLED") + .map(|v| v == "true") + .unwrap_or(false); + if emb_enabled { + println!( + "enabled (provider: {}, model: {})", + settings.embeddings.provider, settings.embeddings.model + ); + } else { + println!("disabled"); + } + + // WASM tools + print!(" WASM Tools: "); + let tools_dir = settings + .wasm + .tools_dir + .clone() + .unwrap_or_else(default_tools_dir); + if tools_dir.exists() { + let count = count_wasm_files(&tools_dir); + println!("{} installed ({})", count, tools_dir.display()); + } else { + println!("directory not found ({})", tools_dir.display()); + } + + // WASM channels + print!(" Channels: "); + let channels_dir = settings + .channels + .wasm_channels_dir + .clone() + .unwrap_or_else(default_channels_dir); + let mut channel_info = vec!["cli".to_string()]; + if settings.channels.http_enabled { + channel_info.push(format!( + "http:{}", + settings.channels.http_port.unwrap_or(3000) + )); + } + if channels_dir.exists() { + let wasm_count = count_wasm_files(&channels_dir); + if wasm_count > 0 { + channel_info.push(format!("{} wasm", wasm_count)); + } + } + println!("{}", channel_info.join(", ")); + + // Heartbeat + print!(" Heartbeat: "); + let hb_enabled = settings.heartbeat.enabled + || std::env::var("HEARTBEAT_ENABLED") + .map(|v| v == "true") + .unwrap_or(false); + if hb_enabled { + println!("enabled (interval: {}s)", settings.heartbeat.interval_secs); + } else { + println!("disabled"); + } + + // MCP servers + print!(" MCP Servers: "); + match crate::tools::mcp::config::load_mcp_servers().await { + Ok(servers) => { + let enabled = servers.servers.iter().filter(|s| s.enabled).count(); + let total = servers.servers.len(); + println!("{} enabled / {} configured", enabled, total); + } + Err(_) => println!("none configured"), + } + + // Settings path + println!("\n Settings: {}", Settings::default_path().display()); + + Ok(()) +} + +async fn check_database() -> anyhow::Result<()> { + let _ = dotenvy::dotenv(); + let settings = Settings::load(); + let url = std::env::var("DATABASE_URL") + .ok() + .or(settings.database_url) + .ok_or_else(|| anyhow::anyhow!("no URL"))?; + + let config: deadpool_postgres::Config = deadpool_postgres::Config { + url: Some(url), + ..Default::default() + }; + let pool = config + .create_pool( + Some(deadpool_postgres::Runtime::Tokio1), + tokio_postgres::NoTls, + ) + .map_err(|e| anyhow::anyhow!("pool error: {}", e))?; + + let client = tokio::time::timeout(std::time::Duration::from_secs(5), pool.get()) + .await + .map_err(|_| anyhow::anyhow!("timeout"))? + .map_err(|e| anyhow::anyhow!("{}", e))?; + + client + .execute("SELECT 1", &[]) + .await + .map_err(|e| anyhow::anyhow!("{}", e))?; + + Ok(()) +} + +fn count_wasm_files(dir: &std::path::Path) -> usize { + std::fs::read_dir(dir) + .map(|entries| { + entries + .filter_map(|e| e.ok()) + .filter(|e| e.path().extension().is_some_and(|ext| ext == "wasm")) + .count() + }) + .unwrap_or(0) +} + +fn default_tools_dir() -> PathBuf { + dirs::home_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join(".ironclaw") + .join("tools") +} + +fn default_channels_dir() -> PathBuf { + dirs::home_dir() + .unwrap_or_else(|| PathBuf::from(".")) + .join(".ironclaw") + .join("channels") +} diff --git a/src/config.rs b/src/config.rs index 797d058c..cc2a7225 100644 --- a/src/config.rs +++ b/src/config.rs @@ -415,6 +415,8 @@ pub struct AgentConfig { pub max_repair_attempts: u32, /// Whether to use planning before tool execution. pub use_planning: bool, + /// Session idle timeout. Sessions inactive longer than this are pruned. + pub session_idle_timeout: Duration, } impl AgentConfig { @@ -478,6 +480,16 @@ impl AgentConfig { message: format!("must be 'true' or 'false': {e}"), })? .unwrap_or(settings.agent.use_planning), + session_idle_timeout: Duration::from_secs( + optional_env("SESSION_IDLE_TIMEOUT_SECS")? + .map(|s| s.parse()) + .transpose() + .map_err(|e| ConfigError::InvalidValue { + key: "SESSION_IDLE_TIMEOUT_SECS".to_string(), + message: format!("must be a positive integer: {e}"), + })? + .unwrap_or(settings.agent.session_idle_timeout_secs), + ), }) } } 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 0261fad2..4e177154 100644 --- a/src/main.rs +++ b/src/main.rs @@ -14,9 +14,12 @@ use ironclaw::{ WasmChannelRuntime, WasmChannelRuntimeConfig, WasmChannelServer, }, }, - cli::{Cli, Command, run_mcp_command, run_tool_command}, + cli::{ + Cli, Command, run_mcp_command, run_memory_command, run_status_command, run_tool_command, + }, config::Config, context::ContextManager, + extensions::ExtensionManager, history::Store, llm::{SessionConfig, create_llm_provider, create_session_manager}, safety::SafetyLayer, @@ -62,6 +65,69 @@ async fn main() -> anyhow::Result<()> { return run_mcp_command(mcp_cmd.clone()).await; } + Some(Command::Memory(mem_cmd)) => { + tracing_subscriber::fmt() + .with_env_filter( + EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")), + ) + .init(); + + // Memory commands need database (and optionally embeddings) + let _ = dotenvy::dotenv(); + let config = Config::from_env().map_err(|e| anyhow::anyhow!("{}", e))?; + let store = ironclaw::history::Store::new(&config.database).await?; + store.run_migrations().await?; + + // Set up embeddings if available + let session = ironclaw::llm::create_session_manager(ironclaw::llm::SessionConfig { + auth_base_url: config.llm.nearai.auth_base_url.clone(), + session_path: config.llm.nearai.session_path.clone(), + ..Default::default() + }) + .await; + + let embeddings: Option> = + if config.embeddings.enabled { + match config.embeddings.provider.as_str() { + "nearai" => Some(Arc::new( + ironclaw::workspace::NearAiEmbeddings::new( + &config.llm.nearai.base_url, + session, + ) + .with_model(&config.embeddings.model, 1536), + )), + _ => { + if let Some(api_key) = config.embeddings.openai_api_key() { + let dim = match config.embeddings.model.as_str() { + "text-embedding-3-large" => 3072, + _ => 1536, + }; + Some(Arc::new(ironclaw::workspace::OpenAiEmbeddings::with_model( + api_key, + &config.embeddings.model, + dim, + ))) + } else { + None + } + } + } + } else { + None + }; + + return run_memory_command(mem_cmd.clone(), store.pool(), embeddings).await; + } + Some(Command::Status) => { + let _ = dotenvy::dotenv(); + tracing_subscriber::fmt() + .with_env_filter( + EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn")), + ) + .init(); + + return run_status_command().await; + } Some(Command::Setup { skip_auth, channels_only, @@ -291,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); @@ -315,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> = @@ -425,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() @@ -592,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/settings.rs b/src/settings.rs index bb594a19..adad7d3b 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -229,6 +229,11 @@ pub struct AgentSettings { /// Maximum repair attempts. #[serde(default = "default_max_repair_attempts")] pub max_repair_attempts: u32, + + /// Session idle timeout in seconds (default: 7 days). Sessions inactive + /// longer than this are pruned from memory. + #[serde(default = "default_session_idle_timeout")] + pub session_idle_timeout_secs: u64, } fn default_agent_name() -> String { @@ -251,6 +256,10 @@ fn default_repair_interval() -> u64 { 60 // 1 minute } +fn default_session_idle_timeout() -> u64 { + 7 * 24 * 3600 // 7 days +} + fn default_max_repair_attempts() -> u32 { 3 } @@ -269,6 +278,7 @@ impl Default for AgentSettings { use_planning: true, repair_check_interval_secs: default_repair_interval(), max_repair_attempts: default_max_repair_attempts(), + session_idle_timeout_secs: default_session_idle_timeout(), } } } 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, diff --git a/wit/channel.wit b/wit/channel.wit index f2514c5a..99ea8160 100644 --- a/wit/channel.wit +++ b/wit/channel.wit @@ -16,9 +16,9 @@ // β”‚ β–Ό β”‚ // β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ // β”‚ β–Ό β–Ό β–Ό β”‚ -// β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ -// β”‚ β”‚ on-http-req β”‚ β”‚ on-poll β”‚ β”‚ on-respond β”‚ WASM Exports β”‚ -// β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ +// β”‚ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”‚ +// β”‚ β”‚on-http-reqβ”‚ β”‚on-poll β”‚ β”‚on-respondβ”‚ β”‚on-status β”‚ WASM Exports β”‚ +// β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ // β”‚ β”‚ β”‚ β”‚ β”‚ // β”‚ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β”‚ // β”‚ β”‚ β”‚ @@ -215,6 +215,32 @@ interface channel { metadata-json: string, } + // ==================== Status Types ==================== + + /// Types of status updates the agent can send to channels. + enum status-type { + /// Agent is thinking/processing a response. + thinking, + /// Agent finished processing (response sent or about to be sent). + done, + /// Agent processing was interrupted. + interrupted, + /// A tool execution started. + tool-started, + /// A tool execution completed. + tool-completed, + } + + /// A status update from the agent. + record status-update { + /// The type of status change. + status: status-type, + /// Human-readable description of the status. + message: string, + /// Channel-specific metadata as JSON string (e.g., contains chat_id for routing). + metadata-json: string, + } + // ==================== Lifecycle Callbacks ==================== /// Initialize the channel. @@ -261,6 +287,15 @@ interface channel { /// - Err(string): Delivery failure message on-respond: func(response: agent-response) -> result<_, string>; + /// Notify the channel of agent status changes. + /// + /// Called when the agent starts thinking, finishes, or changes state. + /// Channels can use this to show typing indicators or status messages. + /// + /// Arguments: + /// - update: The status update + on-status: func(update: status-update); + /// Clean up channel resources. /// /// Called when the channel is being unloaded.