diff --git a/.gitignore b/.gitignore index 2577b4a278..4d223cda28 100644 --- a/.gitignore +++ b/.gitignore @@ -17,6 +17,8 @@ target/ # Python __pycache__/ *.pyc +*.pyo +*.pyd # Benchmark results (local runs, not committed) bench-results/ @@ -34,8 +36,5 @@ trace_*.json .claude/settings.local.json .worktrees/ -# Python cache -__pycache__/ -*.pyc -*.pyo -*.pyd +# JetBrains IDE +.idea diff --git a/channels-src/whatsapp/Cargo.lock b/channels-src/whatsapp/Cargo.lock index 0e55d1e532..adefa9aa3b 100644 --- a/channels-src/whatsapp/Cargo.lock +++ b/channels-src/whatsapp/Cargo.lock @@ -269,7 +269,7 @@ dependencies = [ [[package]] name = "whatsapp-channel" -version = "0.1.0" +version = "0.2.0" dependencies = [ "serde", "serde_json", diff --git a/src/app.rs b/src/app.rs index 0ffe782064..5ccfc5f2d9 100644 --- a/src/app.rs +++ b/src/app.rs @@ -400,6 +400,8 @@ impl AppBuilder { let mcp_session_manager = Arc::new(McpSessionManager::new()); let mcp_process_manager = Arc::new(McpProcessManager::new()); + let companion_mcp_server = + crate::tools::mcp::config::derive_nearai_companion_mcp_server(&self.config); // Create WASM tool runtime eagerly so extensions installed after startup // (e.g. via the web UI) can still be activated. The tools directory is only @@ -477,6 +479,7 @@ impl AppBuilder { let mcp_sm = Arc::clone(&mcp_session_manager); let pm = Arc::clone(&mcp_process_manager); let owner_id = self.config.owner_id.clone(); + let companion_mcp_server = companion_mcp_server.clone(); async move { let servers_result = if let Some(ref d) = db { load_mcp_servers_from_db(d.as_ref(), &owner_id).await @@ -484,7 +487,16 @@ impl AppBuilder { crate::tools::mcp::config::load_mcp_servers().await }; match servers_result { - Ok(servers) => { + Ok(mut servers) => { + if let Some(companion) = companion_mcp_server { + let companion_name = companion.name.clone(); + if !servers.insert_if_absent(companion) { + tracing::debug!( + "Skipping derived MCP companion '{}': an existing config with that name is already present", + companion_name + ); + } + } let enabled: Vec<_> = servers.enabled_servers().cloned().collect(); if !enabled.is_empty() { tracing::debug!( @@ -496,6 +508,8 @@ impl AppBuilder { let mut join_set = tokio::task::JoinSet::new(); for server in enabled { let mcp_sm = Arc::clone(&mcp_sm); + let nearai_session = Arc::clone(&self.session); + let nearai_api_key = self.config.llm.nearai.api_key.clone(); let secrets = secrets_store.clone(); let tools = Arc::clone(&tools); let pm = Arc::clone(&pm); @@ -507,6 +521,8 @@ impl AppBuilder { let client = match crate::tools::mcp::create_client_from_config( server, &mcp_sm, + Some(nearai_session), + nearai_api_key, &pm, secrets, &owner_id, @@ -644,6 +660,8 @@ impl AppBuilder { let manager = Arc::new(ExtensionManager::new( Arc::clone(&mcp_session_manager), Arc::clone(&mcp_process_manager), + Some(Arc::clone(&self.session)), + self.config.llm.nearai.api_key.clone(), ext_secrets, Arc::clone(tools), Some(Arc::clone(hooks)), @@ -653,6 +671,7 @@ impl AppBuilder { self.config.tunnel.public_url.clone(), self.config.owner_id.clone(), self.db.clone(), + companion_mcp_server, catalog_entries.clone(), )); tools.register_extension_tools(Arc::clone(&manager)); diff --git a/src/channels/web/handlers/extensions.rs b/src/channels/web/handlers/extensions.rs index 855fba3ed9..429dee130f 100644 --- a/src/channels/web/handlers/extensions.rs +++ b/src/channels/web/handlers/extensions.rs @@ -68,6 +68,7 @@ pub async fn extensions_list_handler( tools: ext.tools, needs_setup: ext.needs_setup, has_auth: ext.has_auth, + derived: ext.derived, activation_status, activation_error: ext.activation_error, version: ext.version, diff --git a/src/channels/web/server.rs b/src/channels/web/server.rs index 27ef7cdce9..5246d0854f 100644 --- a/src/channels/web/server.rs +++ b/src/channels/web/server.rs @@ -1885,6 +1885,7 @@ async fn extensions_list_handler( tools: ext.tools, needs_setup: ext.needs_setup, has_auth: ext.has_auth, + derived: ext.derived, activation_status, activation_error: ext.activation_error, version: ext.version, @@ -2795,6 +2796,7 @@ mod tests { tools: Vec::new(), needs_setup: true, has_auth: false, + derived: false, installed: true, activation_error: None, version: None, @@ -2832,6 +2834,7 @@ mod tests { tools: Vec::new(), needs_setup: true, has_auth: false, + derived: false, installed: true, activation_error: None, version: None, @@ -3490,6 +3493,8 @@ mod tests { let ext_mgr = Arc::new(ExtensionManager::new( mcp_sm, mcp_pm, + None, + None, secrets, tool_registry, None, @@ -3499,6 +3504,7 @@ mod tests { None, "test".to_string(), None, + None, vec![], )); (ext_mgr, wasm_tools_dir, wasm_channels_dir) diff --git a/src/channels/web/types.rs b/src/channels/web/types.rs index 3fad9f3525..70c07ba246 100644 --- a/src/channels/web/types.rs +++ b/src/channels/web/types.rs @@ -462,6 +462,9 @@ pub struct ExtensionInfo { /// Whether this extension has an auth configuration (OAuth or manual token). #[serde(default)] pub has_auth: bool, + /// Whether this extension is derived from runtime/provider state. + #[serde(default)] + pub derived: bool, /// WASM channel activation status. #[serde(skip_serializing_if = "Option::is_none")] pub activation_status: Option, diff --git a/src/cli/mcp.rs b/src/cli/mcp.rs index 2293a6d62c..45c354ed34 100644 --- a/src/cli/mcp.rs +++ b/src/cli/mcp.rs @@ -8,9 +8,10 @@ use std::sync::Arc; use clap::{Args, Subcommand}; -use crate::config::Config; +use crate::config::{Config, LlmConfig}; use crate::db::Database; use crate::secrets::SecretsStore; +use crate::settings::Settings; use crate::tools::mcp::{ McpClient, McpProcessManager, McpServerConfig, McpSessionManager, OAuthConfig, auth::{authorize_mcp_server, is_authenticated}, @@ -173,6 +174,13 @@ async fn add_server(args: McpAddArgs) -> anyhow::Result<()> { description, } = args; + if config::is_nearai_companion_server_name(&name) { + anyhow::bail!( + "Server name '{}' is reserved for the NEAR AI companion MCP server", + name + ); + } + let transport_lower = transport.to_lowercase(); let mut config = match transport_lower.as_str() { @@ -244,7 +252,7 @@ async fn add_server(args: McpAddArgs) -> anyhow::Result<()> { // Save (DB if available, else disk) let db = connect_db().await; - let mut servers = load_servers(db.as_deref()).await?; + let mut servers = load_persisted_servers(db.as_deref()).await?; servers.upsert(config); save_servers(db.as_deref(), &servers).await?; @@ -281,8 +289,15 @@ async fn add_server(args: McpAddArgs) -> anyhow::Result<()> { /// Remove an MCP server. async fn remove_server(name: String) -> anyhow::Result<()> { + if config::is_nearai_companion_server_name(&name) { + anyhow::bail!( + "Server '{}' is derived from the active NEAR AI provider and cannot be removed directly", + name + ); + } + let db = connect_db().await; - let mut servers = load_servers(db.as_deref()).await?; + let mut servers = load_persisted_servers(db.as_deref()).await?; if !servers.remove(&name) { anyhow::bail!("Server '{}' not found", name); } @@ -298,7 +313,7 @@ async fn remove_server(name: String) -> anyhow::Result<()> { /// List configured MCP servers. async fn list_servers(verbose: bool) -> anyhow::Result<()> { let db = connect_db().await; - let servers = load_servers(db.as_deref()).await?; + let servers = load_servers_with_derived(db.as_deref()).await?; if servers.servers.is_empty() { println!(); @@ -404,12 +419,23 @@ async fn list_servers(verbose: bool) -> anyhow::Result<()> { async fn auth_server(name: String, user_id: String) -> anyhow::Result<()> { // Get server config let db = connect_db().await; - let servers = load_servers(db.as_deref()).await?; + let servers = load_servers_with_derived(db.as_deref()).await?; let server = servers .get(&name) .cloned() .ok_or_else(|| anyhow::anyhow!("Server '{}' not found", name))?; + if server.uses_runtime_auth_source() { + println!(); + println!( + " Server '{}' reuses your active NEAR AI authentication and does not support separate MCP OAuth.", + name + ); + println!(" Configure NEAR AI auth (API key or session login) instead."); + println!(); + return Ok(()); + } + // Initialize secrets store let secrets = get_secrets_store().await?; @@ -477,7 +503,7 @@ async fn auth_server(name: String, user_id: String) -> anyhow::Result<()> { async fn test_server(name: String, user_id: String) -> anyhow::Result<()> { // Get server config let db = connect_db().await; - let servers = load_servers(db.as_deref()).await?; + let servers = load_servers_with_derived(db.as_deref()).await?; let server = servers .get(&name) .cloned() @@ -488,35 +514,66 @@ async fn test_server(name: String, user_id: String) -> anyhow::Result<()> { // Create client let session_manager = Arc::new(McpSessionManager::new()); - - // Always check for stored tokens (from either pre-configured OAuth or DCR) - let secrets = get_secrets_store().await?; - let has_tokens = is_authenticated(&server, &secrets, &user_id).await; - - let client = if has_tokens { - // We have stored tokens, use authenticated client - McpClient::new_authenticated(server.clone(), session_manager.clone(), secrets, user_id) - } else if server.requires_auth() { - // OAuth configured but no tokens - need to authenticate - println!(); - println!( - " ✗ Not authenticated. Run 'ironclaw mcp auth {}' first.", - name - ); - println!(); - return Ok(()); - } else { - // Use the factory to dispatch on transport type (HTTP, stdio, unix) + let (client, has_tokens) = if server.uses_runtime_auth_source() { let process_manager = Arc::new(McpProcessManager::new()); - create_client_from_config( - server.clone(), - &session_manager, - &process_manager, - None, - "default", + let llm = resolve_llm_from_env()?; + let nearai_session = crate::llm::create_session_manager(llm.session.clone()).await; + ( + create_client_from_config( + server.clone(), + &session_manager, + Some(nearai_session), + llm.nearai.api_key.clone(), + &process_manager, + None, + "default", + ) + .await + .map_err(|e| anyhow::anyhow!("{}", e))?, + false, ) - .await - .map_err(|e| anyhow::anyhow!("{}", e))? + } else { + // Only initialize the secrets store for non-runtime-auth servers that + // can actually use persisted OAuth/DCR tokens. + let secrets = get_secrets_store().await?; + let has_tokens = is_authenticated(&server, &secrets, &user_id).await; + + if has_tokens { + ( + McpClient::new_authenticated( + server.clone(), + session_manager.clone(), + secrets, + user_id, + ), + true, + ) + } else if server.requires_auth() { + println!(); + println!( + " ✗ Not authenticated. Run 'ironclaw mcp auth {}' first.", + name + ); + println!(); + return Ok(()); + } else { + // Use the factory to dispatch on transport type (HTTP, stdio, unix) + let process_manager = Arc::new(McpProcessManager::new()); + ( + create_client_from_config( + server.clone(), + &session_manager, + None, + None, + &process_manager, + None, + "default", + ) + .await + .map_err(|e| anyhow::anyhow!("{}", e))?, + false, + ) + } }; // Test connection @@ -581,8 +638,15 @@ async fn test_server(name: String, user_id: String) -> anyhow::Result<()> { /// Toggle server enabled/disabled state. async fn toggle_server(name: String, enable: bool, disable: bool) -> anyhow::Result<()> { + if config::is_nearai_companion_server_name(&name) { + anyhow::bail!( + "Server '{}' is derived from the active NEAR AI provider and cannot be toggled directly", + name + ); + } + let db = connect_db().await; - let mut servers = load_servers(db.as_deref()).await?; + let mut servers = load_persisted_servers(db.as_deref()).await?; let server = servers .get_mut(&name) @@ -615,13 +679,30 @@ async fn connect_db() -> Option> { crate::db::connect_from_config(&config.database).await.ok() } -/// Load MCP servers (DB if available, else disk). -async fn load_servers(db: Option<&dyn Database>) -> Result { - if let Some(db) = db { - config::load_mcp_servers_from_db(db, DEFAULT_USER_ID).await +/// Load only persisted MCP servers (DB if available, else disk). +async fn load_persisted_servers( + db: Option<&dyn Database>, +) -> Result { + Ok(if let Some(db) = db { + config::load_mcp_servers_from_db(db, DEFAULT_USER_ID).await? } else { - config::load_mcp_servers().await + config::load_mcp_servers().await? + }) +} + +/// Load MCP servers plus any derived runtime companions. +async fn load_servers_with_derived( + db: Option<&dyn Database>, +) -> Result { + let mut servers = load_persisted_servers(db).await?; + + if let Ok(llm) = resolve_llm_from_env() + && let Some(companion) = config::derive_nearai_companion_mcp_server_from_llm(&llm) + { + servers.insert_if_absent(companion); } + + Ok(servers) } /// Save MCP servers (DB if available, else disk). @@ -629,10 +710,15 @@ async fn save_servers( db: Option<&dyn Database>, servers: &McpServersFile, ) -> Result<(), config::ConfigError> { + let mut persisted = servers.clone(); + persisted + .servers + .retain(|server| !config::is_nearai_companion_server_name(&server.name)); + if let Some(db) = db { - config::save_mcp_servers_to_db(db, DEFAULT_USER_ID, servers).await + config::save_mcp_servers_to_db(db, DEFAULT_USER_ID, &persisted).await } else { - config::save_mcp_servers(servers).await + config::save_mcp_servers(&persisted).await } } @@ -641,6 +727,10 @@ async fn get_secrets_store() -> anyhow::Result Result { + LlmConfig::resolve(&Settings::default()) +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/extensions/manager.rs b/src/extensions/manager.rs index 00d787a5a3..c413668b3e 100644 --- a/src/extensions/manager.rs +++ b/src/extensions/manager.rs @@ -317,6 +317,8 @@ pub struct ExtensionManager { // MCP infrastructure mcp_session_manager: Arc, mcp_process_manager: Arc, + nearai_session_manager: Option>, + nearai_api_key: Option, /// Active MCP clients keyed by server name. mcp_clients: RwLock>>, @@ -340,6 +342,8 @@ pub struct ExtensionManager { user_id: String, /// Optional database store for DB-backed MCP config. store: Option>, + /// Companion MCP server derived from the active provider config. + companion_mcp_server: Option, /// Names of WASM channels that were successfully loaded at startup. active_channel_names: RwLock>, /// Installed channel-relay extensions (no on-disk artifact, tracked in memory). @@ -405,6 +409,8 @@ impl ExtensionManager { pub fn new( mcp_session_manager: Arc, mcp_process_manager: Arc, + nearai_session_manager: Option>, + nearai_api_key: Option, secrets: Arc, tool_registry: Arc, hooks: Option>, @@ -414,6 +420,7 @@ impl ExtensionManager { tunnel_url: Option, user_id: String, store: Option>, + companion_mcp_server: Option, catalog_entries: Vec, ) -> Self { let registry = if catalog_entries.is_empty() { @@ -426,6 +433,8 @@ impl ExtensionManager { discovery: OnlineDiscovery::new(), mcp_session_manager, mcp_process_manager, + nearai_session_manager, + nearai_api_key, mcp_clients: RwLock::new(HashMap::new()), wasm_tool_runtime, wasm_tools_dir, @@ -439,6 +448,7 @@ impl ExtensionManager { tunnel_url, user_id, store, + companion_mcp_server, active_channel_names: RwLock::new(HashSet::new()), installed_relay_extensions: RwLock::new(HashSet::new()), activation_errors: RwLock::new(HashMap::new()), @@ -1009,8 +1019,11 @@ impl ExtensionManager { match self.load_mcp_servers().await { Ok(servers) => { for server in &servers.servers { - let authenticated = - is_authenticated(server, &self.secrets, &self.user_id).await; + let authenticated = if server.uses_runtime_auth_source() { + self.is_runtime_authenticated(server).await + } else { + is_authenticated(server, &self.secrets, &self.user_id).await + }; let clients = self.mcp_clients.read().await; let active = clients.contains_key(&server.name); @@ -1041,7 +1054,10 @@ impl ExtensionManager { active, tools, needs_setup: false, - has_auth: false, + has_auth: server.requires_auth(), + derived: crate::tools::mcp::config::is_nearai_companion_server_name( + &server.name, + ), installed: true, activation_error: None, version: None, @@ -1093,6 +1109,7 @@ impl ExtensionManager { tools: if active { vec![name] } else { Vec::new() }, needs_setup: auth_state == ToolAuthState::NeedsSetup, has_auth: auth_state != ToolAuthState::NoAuth, + derived: false, installed: true, activation_error: None, version, @@ -1149,6 +1166,7 @@ impl ExtensionManager { tools: Vec::new(), needs_setup: auth_state == ToolAuthState::NeedsSetup, has_auth: auth_state != ToolAuthState::NoAuth, + derived: false, installed: true, activation_error, version, @@ -1189,6 +1207,7 @@ impl ExtensionManager { tools: Vec::new(), needs_setup: false, has_auth: true, + derived: false, installed: true, activation_error: None, version: None, @@ -1223,6 +1242,7 @@ impl ExtensionManager { tools: Vec::new(), needs_setup: false, has_auth: false, + derived: false, installed: false, activation_error: None, version: entry.version, @@ -1253,6 +1273,12 @@ impl ExtensionManager { match kind { ExtensionKind::McpServer => { + if crate::tools::mcp::config::is_nearai_companion_server_name(name) { + return Err(ExtensionError::Config( + "This MCP server is derived from the active NEAR AI provider and cannot be removed directly".to_string(), + )); + } + // Unregister tools with this server's prefix let tool_names: Vec = self .tool_registry @@ -1679,10 +1705,40 @@ impl ExtensionManager { &self, ) -> Result { - if let Some(ref store) = self.store { - crate::tools::mcp::config::load_mcp_servers_from_db(store.as_ref(), &self.user_id).await + let mut servers = if let Some(ref store) = self.store { + crate::tools::mcp::config::load_mcp_servers_from_db(store.as_ref(), &self.user_id) + .await? } else { - crate::tools::mcp::config::load_mcp_servers().await + crate::tools::mcp::config::load_mcp_servers().await? + }; + + if let Some(ref companion) = self.companion_mcp_server { + servers.insert_if_absent(companion.clone()); + } + + Ok(servers) + } + + async fn is_runtime_authenticated(&self, server: &McpServerConfig) -> bool { + match server.auth_source { + Some(crate::tools::mcp::config::McpAuthSource::NearAi) => { + if self.nearai_api_key.is_some() { + return true; + } + + if let Ok(key) = std::env::var("NEARAI_API_KEY") + && !key.trim().is_empty() + { + return true; + } + + if let Some(ref session) = self.nearai_session_manager { + return session.has_token().await; + } + + false + } + None => false, } } @@ -2238,6 +2294,20 @@ impl ExtensionManager { .await .map_err(|e| ExtensionError::NotInstalled(e.to_string()))?; + if server.uses_runtime_auth_source() { + if self.is_runtime_authenticated(&server).await { + return Ok(AuthResult::authenticated(name, ExtensionKind::McpServer)); + } + + return Ok(AuthResult::needs_setup( + name, + ExtensionKind::McpServer, + "This MCP server reuses your active NEAR AI authentication. Configure a NEAR AI API key or sign in to NEAR AI first, then try again." + .to_string(), + None, + )); + } + // Check if already authenticated if is_authenticated(&server, &self.secrets, &self.user_id).await { return Ok(AuthResult::authenticated(name, ExtensionKind::McpServer)); @@ -3276,6 +3346,8 @@ impl ExtensionManager { let client = crate::tools::mcp::create_client_from_config( server.clone(), &self.mcp_session_manager, + self.nearai_session_manager.clone(), + self.nearai_api_key.clone(), &self.mcp_process_manager, Some(Arc::clone(&self.secrets)), &self.user_id, @@ -5496,6 +5568,8 @@ mod tests { crate::extensions::manager::ExtensionManager::new( mcp, Arc::new(McpProcessManager::new()), + None, + None, secrets, tools, None, // hooks @@ -5505,6 +5579,7 @@ mod tests { None, // tunnel_url "test".to_string(), None, // db + None, // companion MCP vec![], ) } @@ -5678,6 +5753,8 @@ mod tests { ExtensionManager::new( Arc::new(McpSessionManager::new()), Arc::new(McpProcessManager::new()), + None, + None, Arc::new(InMemorySecretsStore::new(crypto)), Arc::new(ToolRegistry::new()), None, @@ -5687,6 +5764,7 @@ mod tests { None, "test".to_string(), None, + None, Vec::new(), ) } @@ -5828,6 +5906,8 @@ mod tests { ExtensionManager::new( Arc::new(McpSessionManager::new()), Arc::new(McpProcessManager::new()), + None, + None, Arc::new(InMemorySecretsStore::new(crypto)), Arc::new(ToolRegistry::new()), None, @@ -5837,6 +5917,7 @@ mod tests { None, "test".to_string(), Some(db), + None, Vec::new(), ) }; @@ -6086,6 +6167,8 @@ mod tests { let manager = ExtensionManager::new( Arc::new(McpSessionManager::new()), Arc::new(McpProcessManager::new()), + None, + None, Arc::new(InMemorySecretsStore::new(crypto)), Arc::new(ToolRegistry::new()), None, @@ -6095,6 +6178,7 @@ mod tests { None, "test".to_string(), Some(db.clone() as Arc), + None, Vec::new(), ); @@ -6671,6 +6755,8 @@ mod tests { ExtensionManager::new( mcp, Arc::new(McpProcessManager::new()), + None, + None, secrets, tools, None, @@ -6680,6 +6766,7 @@ mod tests { tunnel_url, "test".to_string(), None, + None, vec![], ) } diff --git a/src/extensions/mod.rs b/src/extensions/mod.rs index 2a4d189f8e..4a0284c2b5 100644 --- a/src/extensions/mod.rs +++ b/src/extensions/mod.rs @@ -504,6 +504,10 @@ pub struct InstalledExtension { /// Whether this extension has an auth configuration (OAuth or manual token). #[serde(default)] pub has_auth: bool, + /// Whether this extension is derived from provider/runtime state instead of + /// being a user-managed persisted configuration. + #[serde(default)] + pub derived: bool, /// Whether this extension is installed locally (false = available in registry but not installed). #[serde(default = "default_true")] pub installed: bool, @@ -934,6 +938,7 @@ mod tests { assert!(ext.installed, "installed should default to true"); assert!(!ext.needs_setup, "needs_setup should default to false"); assert!(!ext.has_auth); + assert!(!ext.derived); assert!(ext.tools.is_empty()); assert!(ext.display_name.is_none()); assert!(ext.description.is_none()); @@ -954,6 +959,7 @@ mod tests { tools: vec!["send_email".to_string(), "read_inbox".to_string()], needs_setup: true, has_auth: true, + derived: true, installed: false, activation_error: Some("token expired".to_string()), version: None, @@ -963,6 +969,7 @@ mod tests { assert_eq!(json["description"], "Read and send emails"); assert_eq!(json["url"], "https://gmail.example.com"); assert_eq!(json["needs_setup"], true); + assert_eq!(json["derived"], true); assert_eq!(json["installed"], false); assert_eq!(json["activation_error"], "token expired"); @@ -970,6 +977,7 @@ mod tests { assert_eq!(back.name, "gmail"); assert_eq!(back.tools.len(), 2); assert!(back.needs_setup); + assert!(back.derived); assert!(!back.installed); assert_eq!(back.activation_error.as_deref(), Some("token expired")); } diff --git a/src/llm/mod.rs b/src/llm/mod.rs index 3b6b01c472..b896c80d46 100644 --- a/src/llm/mod.rs +++ b/src/llm/mod.rs @@ -18,6 +18,7 @@ pub mod config; pub mod costs; pub mod error; pub mod failover; +pub mod nearai_auth; mod nearai_chat; pub mod oauth_helpers; mod provider; @@ -42,6 +43,7 @@ pub use config::{ }; pub use error::LlmError; pub use failover::{CooldownConfig, FailoverProvider}; +pub use nearai_auth::resolve_nearai_bearer_token; pub use nearai_chat::{ModelInfo, NearAiChatProvider}; pub use provider::{ ChatMessage, CompletionRequest, CompletionResponse, ContentPart, FinishReason, ImageUrl, diff --git a/src/llm/nearai_auth.rs b/src/llm/nearai_auth.rs new file mode 100644 index 0000000000..cbd4cbee6c --- /dev/null +++ b/src/llm/nearai_auth.rs @@ -0,0 +1,42 @@ +use secrecy::{ExposeSecret, SecretString}; + +use crate::llm::LlmError; +use crate::llm::session::SessionManager; + +/// Resolve the active NEAR AI bearer token. +/// +/// Priority order: +/// 1. Explicit API key from resolved config +/// 2. Existing session token +/// 3. Interactive session authentication +/// 4. `NEARAI_API_KEY` from runtime environment +pub async fn resolve_nearai_bearer_token( + api_key: Option<&SecretString>, + session: &SessionManager, +) -> Result { + if let Some(api_key) = api_key { + return Ok(api_key.expose_secret().to_string()); + } + + if session.has_token().await { + let token = session.get_token().await?; + return Ok(token.expose_secret().to_string()); + } + + session.ensure_authenticated().await?; + + if session.has_token().await { + let token = session.get_token().await?; + return Ok(token.expose_secret().to_string()); + } + + if let Ok(key) = std::env::var("NEARAI_API_KEY") + && !key.is_empty() + { + return Ok(key); + } + + Err(LlmError::AuthFailed { + provider: "nearai".to_string(), + }) +} diff --git a/src/llm/nearai_chat.rs b/src/llm/nearai_chat.rs index 0a9e1fdc66..364252c5cb 100644 --- a/src/llm/nearai_chat.rs +++ b/src/llm/nearai_chat.rs @@ -158,36 +158,7 @@ impl NearAiChatProvider { /// The env var fallback (#3) only triggers after `ensure_authenticated()` /// runs, because `api_key_login()` sets the env var but not a session token. async fn resolve_bearer_token(&self) -> Result { - // 1. Config-level API key takes priority - if let Some(ref api_key) = self.config.api_key { - return Ok(api_key.expose_secret().to_string()); - } - - // 2. Existing session token (OAuth was already completed) - if self.session.has_token().await { - let token = self.session.get_token().await?; - return Ok(token.expose_secret().to_string()); - } - - // No token yet, trigger interactive login - self.session.ensure_authenticated().await?; - - // 3. After login, check if a session token was stored (OAuth path) - if self.session.has_token().await { - let token = self.session.get_token().await?; - return Ok(token.expose_secret().to_string()); - } - - // 4. api_key_login() sets NEARAI_API_KEY env var but not a session token - if let Ok(key) = std::env::var("NEARAI_API_KEY") - && !key.is_empty() - { - return Ok(key); - } - - Err(LlmError::AuthFailed { - provider: "nearai".to_string(), - }) + crate::llm::resolve_nearai_bearer_token(self.config.api_key.as_ref(), &self.session).await } /// Send a single request to the chat completions API. diff --git a/src/tools/builtin/extension_tools.rs b/src/tools/builtin/extension_tools.rs index cb0f71dd72..261eacf539 100644 --- a/src/tools/builtin/extension_tools.rs +++ b/src/tools/builtin/extension_tools.rs @@ -800,6 +800,8 @@ mod tests { Arc::new(ExtensionManager::new( Arc::new(McpSessionManager::new()), Arc::new(crate::tools::mcp::process::McpProcessManager::new()), + None, + None, Arc::new(InMemorySecretsStore::new(crypto)), Arc::new(ToolRegistry::new()), None, @@ -809,6 +811,7 @@ mod tests { None, "test".to_string(), None, + None, Vec::new(), )) } diff --git a/src/tools/mcp/client.rs b/src/tools/mcp/client.rs index c299ac49c9..8d8a00194f 100644 --- a/src/tools/mcp/client.rs +++ b/src/tools/mcp/client.rs @@ -8,12 +8,13 @@ use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use async_trait::async_trait; +use secrecy::SecretString; use tokio::sync::RwLock; use crate::context::JobContext; use crate::secrets::SecretsStore; use crate::tools::mcp::auth::refresh_access_token; -use crate::tools::mcp::config::McpServerConfig; +use crate::tools::mcp::config::{McpAuthSource, McpServerConfig}; use crate::tools::mcp::http_transport::HttpMcpTransport; use crate::tools::mcp::protocol::{ CallToolResult, InitializeResult, ListToolsResult, McpRequest, McpResponse, McpTool, @@ -46,6 +47,13 @@ pub struct McpClient { /// Session manager (shared across clients). session_manager: Option>, + /// NEAR AI auth/session manager for companion MCP servers that reuse the + /// active provider bearer token. + nearai_session_manager: Option>, + + /// Resolved NEAR AI API key for companion MCP servers. + nearai_api_key: Option, + /// Secrets store for retrieving access tokens. secrets: Option>, @@ -80,6 +88,8 @@ impl McpClient { next_id: AtomicU64::new(1), tools_cache: RwLock::new(None), session_manager: None, + nearai_session_manager: None, + nearai_api_key: None, secrets: None, user_id: "default".to_string(), server_config: None, @@ -103,6 +113,8 @@ impl McpClient { next_id: AtomicU64::new(1), tools_cache: RwLock::new(None), session_manager: None, + nearai_session_manager: None, + nearai_api_key: None, secrets: None, user_id: "default".to_string(), server_config: None, @@ -139,6 +151,8 @@ impl McpClient { next_id: AtomicU64::new(1), tools_cache: RwLock::new(None), session_manager: None, + nearai_session_manager: None, + nearai_api_key: None, secrets: None, user_id: "default".to_string(), custom_headers: config.headers.clone(), @@ -170,6 +184,8 @@ impl McpClient { next_id: AtomicU64::new(1), tools_cache: RwLock::new(None), session_manager: Some(session_manager), + nearai_session_manager: None, + nearai_api_key: None, secrets: Some(secrets), user_id: user_id.into(), server_config: Some(config), @@ -206,6 +222,8 @@ impl McpClient { next_id: AtomicU64::new(1), tools_cache: RwLock::new(None), session_manager, + nearai_session_manager: None, + nearai_api_key: None, secrets, user_id: user_id.into(), server_config, @@ -220,6 +238,21 @@ impl McpClient { self } + /// Attach the NEAR AI session manager for companion MCP auth reuse. + pub fn with_nearai_session_manager( + mut self, + nearai_session_manager: Arc, + ) -> Self { + self.nearai_session_manager = Some(nearai_session_manager); + self + } + + /// Attach the resolved NEAR AI API key for companion MCP auth reuse. + pub fn with_nearai_api_key(mut self, nearai_api_key: Option) -> Self { + self.nearai_api_key = nearai_api_key; + self + } + /// Get the server name. pub fn server_name(&self) -> &str { &self.server_name @@ -261,6 +294,37 @@ impl McpClient { } } + /// Resolve a runtime-provided auth token for companion MCP servers. + async fn get_runtime_auth_token(&self) -> Result, ToolError> { + let Some(ref config) = self.server_config else { + return Ok(None); + }; + + match config.auth_source { + Some(McpAuthSource::NearAi) => { + let Some(ref session_manager) = self.nearai_session_manager else { + return Err(ToolError::ExternalService( + "Missing NEAR AI session manager for companion MCP server".to_string(), + )); + }; + + crate::llm::resolve_nearai_bearer_token( + self.nearai_api_key.as_ref(), + session_manager, + ) + .await + .map(Some) + .map_err(|e| { + ToolError::ExternalService(format!( + "Failed to resolve NEAR AI token for MCP server '{}': {}", + self.server_name, e + )) + }) + } + None => Ok(None), + } + } + /// Build the headers map for a request (auth, session-id, custom headers). /// /// Custom headers are applied first. OAuth token injection is skipped if the @@ -274,6 +338,9 @@ impl McpClient { .custom_headers .keys() .any(|k| k.eq_ignore_ascii_case("authorization")); + if !has_custom_auth && let Some(token) = self.get_runtime_auth_token().await? { + headers.insert("Authorization".to_string(), format!("Bearer {}", token)); + } if !has_custom_auth && let Some(token) = self.get_access_token().await? { let trimmed = token.trim(); if !trimmed.is_empty() { @@ -456,13 +523,12 @@ impl McpClient { ))); } - response + let raw_result = response .result - .ok_or_else(|| ToolError::ExternalService("No result in MCP response".to_string())) - .and_then(|r| { - serde_json::from_value(r) - .map_err(|e| ToolError::ExternalService(format!("Invalid tool result: {}", e))) - }) + .ok_or_else(|| ToolError::ExternalService("No result in MCP response".to_string()))?; + + serde_json::from_value(raw_result) + .map_err(|e| ToolError::ExternalService(format!("Invalid tool result: {}", e))) } /// Clear the tools cache. @@ -509,6 +575,8 @@ impl Clone for McpClient { next_id: AtomicU64::new(self.next_id.load(Ordering::SeqCst)), tools_cache: RwLock::new(None), session_manager: self.session_manager.clone(), + nearai_session_manager: self.nearai_session_manager.clone(), + nearai_api_key: self.nearai_api_key.clone(), secrets: self.secrets.clone(), user_id: self.user_id.clone(), server_config: self.server_config.clone(), @@ -556,7 +624,7 @@ impl Tool for McpToolWrapper { // Strip top-level null values before forwarding — LLMs often emit // `"field": null` for optional params, but many MCP servers reject // explicit nulls for fields that should simply be absent. - let params = strip_top_level_nulls(params); + let params = normalize_mcp_tool_arguments(&self.tool.name, strip_top_level_nulls(params)); let result = self.client.call_tool(&self.tool.name, params).await?; let content: String = result @@ -600,6 +668,24 @@ fn strip_top_level_nulls(value: serde_json::Value) -> serde_json::Value { } } +fn normalize_mcp_tool_arguments(tool_name: &str, value: serde_json::Value) -> serde_json::Value { + if tool_name != "web_search" { + return value; + } + + let serde_json::Value::Object(mut map) = value else { + return value; + }; + + // Keep this intentionally narrow: only strip optional fields that the + // model frequently emits as empty strings. Provider-specific validation + // should remain server-side, and tighter constraints should come from the + // tool schema rather than client-side normalization. + map.retain(|_, value| !value.as_str().is_some_and(|s| s.trim().is_empty())); + + serde_json::Value::Object(map) +} + #[cfg(test)] mod tests { use super::*; @@ -767,6 +853,31 @@ mod tests { assert!(client.has_session_manager()); } + #[tokio::test] + async fn test_build_request_headers_with_nearai_runtime_auth() { + use crate::llm::{ + SessionConfig as NearAiSessionConfig, SessionManager as NearAiSessionManager, + }; + use secrecy::SecretString; + + let config = McpServerConfig::new("chat_api", "http://localhost:3000/mcp") + .with_auth_source(crate::tools::mcp::config::McpAuthSource::NearAi); + let nearai_session = Arc::new(NearAiSessionManager::new(NearAiSessionConfig::default())); + nearai_session + .set_token(SecretString::from("sess_test_token")) + .await; + + let client = McpClient::new_with_config(config) + .expect("valid MCP config") + .with_nearai_session_manager(nearai_session); + let headers = client.build_request_headers().await.expect("headers"); + + assert_eq!( + headers.get("Authorization").map(String::as_str), + Some("Bearer sess_test_token") + ); + } + #[test] fn test_next_request_id_monotonically_increasing() { let client = McpClient::new("http://localhost:1234"); @@ -1253,4 +1364,48 @@ mod tests { "Token must be trimmed before use in Authorization header" ); } + + #[test] + fn test_normalize_web_search_arguments_removes_empty_optional_fields() { + let input = serde_json::json!({ + "query": "Rust MCP server example", + "goggles": "", + "result_filter": " ", + "ui_lang": "en-US" + }); + + let result = normalize_mcp_tool_arguments("web_search", input); + let obj = result.as_object().unwrap(); + assert_eq!(obj["query"], "Rust MCP server example"); + assert_eq!(obj["ui_lang"], "en-US"); + assert!(!obj.contains_key("goggles")); + assert!(!obj.contains_key("result_filter")); + } + + #[test] + fn test_normalize_web_search_arguments_strips_any_empty_string_field() { + let input = serde_json::json!({ + "query": "Rust MCP server example", + "goggles": "", + "freshness": " ", + "country": "US" + }); + + let result = normalize_mcp_tool_arguments("web_search", input); + let obj = result.as_object().unwrap(); + assert_eq!(obj["country"], "US"); + assert!(!obj.contains_key("freshness")); + assert!(!obj.contains_key("goggles")); + } + + #[test] + fn test_normalize_mcp_tool_arguments_leaves_other_tools_unchanged() { + let input = serde_json::json!({ + "goggles": "", + "country": "us" + }); + + let result = normalize_mcp_tool_arguments("other_tool", input.clone()); + assert_eq!(result, input); + } } diff --git a/src/tools/mcp/config.rs b/src/tools/mcp/config.rs index 06adbd3dc5..cb9456c2d0 100644 --- a/src/tools/mcp/config.rs +++ b/src/tools/mcp/config.rs @@ -51,6 +51,13 @@ pub struct McpServerConfig { #[serde(skip_serializing_if = "Option::is_none")] pub oauth: Option, + /// Built-in auth source provided by IronClaw at runtime. + /// + /// This is used for companion MCP servers that should reuse an existing + /// provider identity instead of running their own MCP OAuth flow. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub auth_source: Option, + /// Whether this server is enabled. #[serde(default = "default_true")] pub enabled: bool, @@ -60,6 +67,14 @@ pub struct McpServerConfig { pub description: Option, } +/// Runtime-provided auth sources for MCP companion servers. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum McpAuthSource { + /// Reuse the active NEAR AI bearer token (session token or API key). + NearAi, +} + fn default_true() -> bool { true } @@ -73,6 +88,7 @@ impl McpServerConfig { transport: None, headers: HashMap::new(), oauth: None, + auth_source: None, enabled: true, description: None, } @@ -95,6 +111,7 @@ impl McpServerConfig { }), headers: HashMap::new(), oauth: None, + auth_source: None, enabled: true, description: None, } @@ -110,6 +127,7 @@ impl McpServerConfig { }), headers: HashMap::new(), oauth: None, + auth_source: None, enabled: true, description: None, } @@ -121,6 +139,12 @@ impl McpServerConfig { self } + /// Set a runtime-provided auth source. + pub fn with_auth_source(mut self, auth_source: McpAuthSource) -> Self { + self.auth_source = Some(auth_source); + self + } + /// Set description. pub fn with_description(mut self, description: impl Into) -> Self { self.description = Some(description.into()); @@ -222,6 +246,11 @@ impl McpServerConfig { .any(|k| k.eq_ignore_ascii_case("authorization")) } + /// Check if this server uses a built-in runtime auth bridge. + pub fn uses_runtime_auth_source(&self) -> bool { + self.auth_source.is_some() + } + /// Check if this server requires authentication. /// /// Returns true if OAuth is pre-configured OR if this is a remote HTTPS server @@ -234,7 +263,7 @@ impl McpServerConfig { return false; } - if self.oauth.is_some() { + if self.oauth.is_some() || self.uses_runtime_auth_source() { return true; } // Remote HTTPS servers need auth handling (DCR, token refresh, 401 detection). @@ -260,6 +289,50 @@ impl McpServerConfig { } } +/// Reserved name used for the companion chat-api MCP server derived from NEAR AI config. +pub const NEARAI_COMPANION_MCP_NAME: &str = "_nearai_companion_mcp"; + +pub fn is_nearai_companion_server_name(name: &str) -> bool { + name == NEARAI_COMPANION_MCP_NAME +} + +/// Build the companion chat-api MCP server from the active NearAI config. +/// +/// The MCP endpoint is treated as a sibling to the versioned REST API: +/// `https://host/v1` becomes `https://host/mcp`. +pub fn derive_nearai_companion_mcp_server( + config: &crate::config::Config, +) -> Option { + derive_nearai_companion_mcp_server_from_llm(&config.llm) +} + +/// Build the companion chat-api MCP server from an LLM config. +/// +/// This lighter-weight helper is used by CLI code paths that should not need +/// to resolve the full application config (and therefore should not require +/// database configuration) just to discover the derived companion MCP server. +pub fn derive_nearai_companion_mcp_server_from_llm( + llm: &crate::config::LlmConfig, +) -> Option { + if llm.backend != "nearai" { + return None; + } + + let base = llm.nearai.base_url.trim_end_matches('/'); + let mcp_base = base + .strip_suffix("/v1") + .unwrap_or(base) + .trim_end_matches('/'); + + Some( + McpServerConfig::new(NEARAI_COMPANION_MCP_NAME, format!("{mcp_base}/mcp")) + .with_auth_source(McpAuthSource::NearAi) + .with_description( + "Companion chat-api MCP server derived from the active NEAR AI provider", + ), + ) +} + /// OAuth 2.1 configuration for an MCP server. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct OAuthConfig { @@ -356,6 +429,16 @@ impl McpServersFile { } } + /// Insert a server only if no server with the same name already exists. + pub fn insert_if_absent(&mut self, config: McpServerConfig) -> bool { + if self.get(&config.name).is_some() { + false + } else { + self.servers.push(config); + true + } + } + /// Remove a server by name. pub fn remove(&mut self, name: &str) -> bool { let len_before = self.servers.len(); @@ -718,6 +801,23 @@ mod tests { assert!(config.servers.is_empty()); } + #[cfg(feature = "libsql")] + #[test] + fn test_derive_nearai_companion_mcp_server_strips_trailing_v1() { + let mut config = crate::config::Config::for_testing( + std::env::temp_dir().join("ironclaw-test-companion.db"), + std::env::temp_dir().join("ironclaw-test-skills"), + std::env::temp_dir().join("ironclaw-test-installed-skills"), + ); + config.llm.backend = "nearai".to_string(); + config.llm.nearai.base_url = "https://private.near.ai/v1".to_string(); + + let server = derive_nearai_companion_mcp_server(&config).expect("companion server"); + assert_eq!(server.name, NEARAI_COMPANION_MCP_NAME); + assert_eq!(server.url, "https://private.near.ai/mcp"); + assert_eq!(server.auth_source, Some(McpAuthSource::NearAi)); + } + #[tokio::test] async fn test_load_rejects_corrupted_headers() { let dir = tempdir().unwrap(); diff --git a/src/tools/mcp/factory.rs b/src/tools/mcp/factory.rs index c31c50513b..915db1a25f 100644 --- a/src/tools/mcp/factory.rs +++ b/src/tools/mcp/factory.rs @@ -20,6 +20,8 @@ pub enum McpFactoryError { UnixNotSupported { name: String }, #[error("Invalid configuration for MCP server '{name}': {reason}")] InvalidConfig { name: String, reason: String }, + #[error("Missing runtime auth context for MCP server '{name}': {reason}")] + MissingRuntimeAuthContext { name: String, reason: String }, } /// Create an `McpClient` from a server configuration, dispatching on the @@ -27,6 +29,8 @@ pub enum McpFactoryError { pub async fn create_client_from_config( server: McpServerConfig, session_manager: &Arc, + nearai_session_manager: Option>, + nearai_api_key: Option, process_manager: &Arc, secrets: Option>, user_id: &str, @@ -78,6 +82,25 @@ pub async fn create_client_from_config( Err(McpFactoryError::UnixNotSupported { name: server_name }) } EffectiveTransport::Http => { + if server.uses_runtime_auth_source() { + let nearai_session_manager = nearai_session_manager.ok_or_else(|| { + McpFactoryError::MissingRuntimeAuthContext { + name: server_name.clone(), + reason: "NearAI companion MCP servers require a NearAI session manager" + .to_string(), + } + })?; + + return Ok(McpClient::new_with_config(server) + .map_err(|e| McpFactoryError::InvalidConfig { + name: server_name.clone(), + reason: e.to_string(), + })? + .with_nearai_session_manager(nearai_session_manager) + .with_nearai_api_key(nearai_api_key) + .with_session_manager(Arc::clone(session_manager))); + } + if let Some(ref secrets) = secrets { let has_tokens = crate::tools::mcp::is_authenticated(&server, secrets, user_id).await; @@ -122,6 +145,8 @@ mod tests { let client = create_client_from_config( server, &session_manager, + None, + None, &process_manager, None, "test-user", diff --git a/tests/module_init_integration.rs b/tests/module_init_integration.rs index c75ccc6f4c..5f3f3f674e 100644 --- a/tests/module_init_integration.rs +++ b/tests/module_init_integration.rs @@ -203,6 +203,8 @@ async fn extension_manager_with_process_manager_constructs() { let manager = ExtensionManager::new( Arc::new(McpSessionManager::new()), Arc::new(McpProcessManager::new()), + None, + None, secrets, tools, None, @@ -212,6 +214,7 @@ async fn extension_manager_with_process_manager_constructs() { None, "test".to_string(), None, + None, Vec::new(), );