From 065a7498d53acdbd77f48da421fa1f805511a7d0 Mon Sep 17 00:00:00 2001 From: Coffee Date: Mon, 16 Mar 2026 16:12:02 +0800 Subject: [PATCH] feat: nearai mcp for web search --- channels-src/whatsapp/Cargo.lock | 2 +- src/app.rs | 14 +- src/channels/web/server.rs | 2 + src/cli/mcp.rs | 21 ++- src/extensions/manager.rs | 28 +++- src/llm/mod.rs | 2 + src/llm/nearai_chat.rs | 31 +--- src/tools/builtin/extension_tools.rs | 2 + src/tools/mcp/client.rs | 204 ++++++++++++++++++++++++++- src/tools/mcp/config.rs | 77 +++++++++- src/tools/mcp/factory.rs | 15 ++ 11 files changed, 359 insertions(+), 39 deletions(-) diff --git a/channels-src/whatsapp/Cargo.lock b/channels-src/whatsapp/Cargo.lock index 0e55d1e5..adefa9aa 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 00804de1..6e041cc5 100644 --- a/src/app.rs +++ b/src/app.rs @@ -393,6 +393,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 @@ -469,6 +471,7 @@ impl AppBuilder { let tools = Arc::clone(tools); let mcp_sm = Arc::clone(&mcp_session_manager); let pm = Arc::clone(&mcp_process_manager); + 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(), "default").await @@ -476,7 +479,10 @@ 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 { + servers.upsert(companion); + } let enabled: Vec<_> = servers.enabled_servers().cloned().collect(); if !enabled.is_empty() { tracing::debug!( @@ -488,6 +494,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); @@ -498,6 +506,8 @@ impl AppBuilder { let client = match crate::tools::mcp::create_client_from_config( server, &mcp_sm, + Some(nearai_session), + nearai_api_key, &pm, secrets, "default", @@ -635,6 +645,7 @@ impl AppBuilder { let manager = Arc::new(ExtensionManager::new( Arc::clone(&mcp_session_manager), Arc::clone(&mcp_process_manager), + Some(Arc::clone(&self.session)), ext_secrets, Arc::clone(tools), Some(Arc::clone(hooks)), @@ -644,6 +655,7 @@ impl AppBuilder { self.config.tunnel.public_url.clone(), "default".to_string(), self.db.clone(), + companion_mcp_server, catalog_entries.clone(), )); tools.register_extension_tools(Arc::clone(&manager)); diff --git a/src/channels/web/server.rs b/src/channels/web/server.rs index e8cb33c2..c7916c2d 100644 --- a/src/channels/web/server.rs +++ b/src/channels/web/server.rs @@ -3365,6 +3365,7 @@ mod tests { let ext_mgr = Arc::new(ExtensionManager::new( mcp_sm, mcp_pm, + None, secrets, tool_registry, None, @@ -3374,6 +3375,7 @@ mod tests { None, "test".to_string(), None, + None, vec![], )); (ext_mgr, wasm_tools_dir, wasm_channels_dir) diff --git a/src/cli/mcp.rs b/src/cli/mcp.rs index 2293a6d6..dddfbae1 100644 --- a/src/cli/mcp.rs +++ b/src/cli/mcp.rs @@ -493,7 +493,22 @@ async fn test_server(name: String, user_id: String) -> anyhow::Result<()> { let secrets = get_secrets_store().await?; let has_tokens = is_authenticated(&server, &secrets, &user_id).await; - let client = if has_tokens { + let client = if server.uses_runtime_auth_source() { + let process_manager = Arc::new(McpProcessManager::new()); + let config = crate::config::Config::from_env().await?; + let nearai_session = crate::llm::create_session_manager(config.llm.session.clone()).await; + create_client_from_config( + server.clone(), + &session_manager, + Some(nearai_session), + config.llm.nearai.api_key.clone(), + &process_manager, + None, + "default", + ) + .await + .map_err(|e| anyhow::anyhow!("{}", e))? + } else 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() { @@ -508,9 +523,13 @@ async fn test_server(name: String, user_id: String) -> anyhow::Result<()> { } else { // Use the factory to dispatch on transport type (HTTP, stdio, unix) let process_manager = Arc::new(McpProcessManager::new()); + let config = crate::config::Config::from_env().await?; + let nearai_session = crate::llm::create_session_manager(config.llm.session.clone()).await; create_client_from_config( server.clone(), &session_manager, + Some(nearai_session), + config.llm.nearai.api_key.clone(), &process_manager, None, "default", diff --git a/src/extensions/manager.rs b/src/extensions/manager.rs index 680c4dfc..62561b25 100644 --- a/src/extensions/manager.rs +++ b/src/extensions/manager.rs @@ -64,6 +64,7 @@ pub struct ExtensionManager { // MCP infrastructure mcp_session_manager: Arc, mcp_process_manager: Arc, + nearai_session_manager: Option>, /// Active MCP clients keyed by server name. mcp_clients: RwLock>>, @@ -87,6 +88,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). @@ -147,6 +150,7 @@ impl ExtensionManager { pub fn new( mcp_session_manager: Arc, mcp_process_manager: Arc, + nearai_session_manager: Option>, secrets: Arc, tool_registry: Arc, hooks: Option>, @@ -156,6 +160,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() { @@ -168,6 +173,7 @@ impl ExtensionManager { discovery: OnlineDiscovery::new(), mcp_session_manager, mcp_process_manager, + nearai_session_manager, mcp_clients: RwLock::new(HashMap::new()), wasm_tool_runtime, wasm_tools_dir, @@ -181,6 +187,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()), @@ -1257,11 +1264,18 @@ 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.upsert(companion.clone()); } + + Ok(servers) } async fn get_mcp_server( @@ -2854,6 +2868,8 @@ impl ExtensionManager { let client = crate::tools::mcp::create_client_from_config( server.clone(), &self.mcp_session_manager, + self.nearai_session_manager.clone(), + None, &self.mcp_process_manager, Some(Arc::clone(&self.secrets)), &self.user_id, @@ -4595,6 +4611,7 @@ mod tests { crate::extensions::manager::ExtensionManager::new( mcp, Arc::new(McpProcessManager::new()), + None, secrets, tools, None, // hooks @@ -4604,6 +4621,7 @@ mod tests { None, // tunnel_url "test".to_string(), None, // db + None, // companion MCP vec![], ) } @@ -4774,6 +4792,7 @@ mod tests { ExtensionManager::new( Arc::new(McpSessionManager::new()), Arc::new(McpProcessManager::new()), + None, Arc::new(InMemorySecretsStore::new(crypto)), Arc::new(ToolRegistry::new()), None, @@ -4783,6 +4802,7 @@ mod tests { None, "test".to_string(), None, + None, Vec::new(), ) } @@ -5344,6 +5364,7 @@ mod tests { ExtensionManager::new( mcp, Arc::new(McpProcessManager::new()), + None, secrets, tools, None, @@ -5353,6 +5374,7 @@ mod tests { tunnel_url, "test".to_string(), None, + None, vec![], ) } diff --git a/src/llm/mod.rs b/src/llm/mod.rs index 3c9de369..3dccf221 100644 --- a/src/llm/mod.rs +++ b/src/llm/mod.rs @@ -16,6 +16,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; @@ -40,6 +41,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_chat.rs b/src/llm/nearai_chat.rs index bf2b8738..57c351ae 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 cb0f71dd..91c0bae4 100644 --- a/src/tools/builtin/extension_tools.rs +++ b/src/tools/builtin/extension_tools.rs @@ -800,6 +800,7 @@ mod tests { Arc::new(ExtensionManager::new( Arc::new(McpSessionManager::new()), Arc::new(crate::tools::mcp::process::McpProcessManager::new()), + None, Arc::new(InMemorySecretsStore::new(crypto)), Arc::new(ToolRegistry::new()), None, @@ -809,6 +810,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 c299ac49..206fdbc5 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() { @@ -509,6 +576,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 +625,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 +669,62 @@ 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; + }; + + for key in ["goggles", "result_filter"] { + if map + .get(key) + .and_then(serde_json::Value::as_str) + .is_some_and(|s| s.trim().is_empty()) + { + map.remove(key); + } + } + + if let Some(country) = map + .get("country") + .and_then(serde_json::Value::as_str) + .map(str::trim) + .filter(|s| !s.is_empty()) + { + map.insert( + "country".to_string(), + serde_json::Value::String(country.to_ascii_uppercase()), + ); + } + + if let Some(ui_lang) = map + .get("ui_lang") + .and_then(serde_json::Value::as_str) + .map(str::trim) + { + if ui_lang.is_empty() || !ui_lang.contains('-') { + map.remove("ui_lang"); + } + } + + if let Some(freshness) = map + .get("freshness") + .and_then(serde_json::Value::as_str) + .map(str::trim) + { + let valid = + matches!(freshness, "pd" | "pw" | "pm" | "py") || freshness.split_once("to").is_some(); + if freshness.is_empty() || !valid { + map.remove("freshness"); + } + } + + serde_json::Value::Object(map) +} + #[cfg(test)] mod tests { use super::*; @@ -767,6 +892,29 @@ 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).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 +1401,56 @@ 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_uppercases_country_and_drops_short_ui_lang() { + let input = serde_json::json!({ + "country": "us", + "ui_lang": "en" + }); + + let result = normalize_mcp_tool_arguments("web_search", input); + let obj = result.as_object().unwrap(); + assert_eq!(obj["country"], "US"); + assert!(!obj.contains_key("ui_lang")); + } + + #[test] + fn test_normalize_web_search_arguments_drops_invalid_freshness() { + let input = serde_json::json!({ + "freshness": "365d" + }); + + let result = normalize_mcp_tool_arguments("web_search", input); + let obj = result.as_object().unwrap(); + assert!(!obj.contains_key("freshness")); + } + + #[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 06adbd3d..ad392eed 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,35 @@ impl McpServerConfig { } } +/// Default name used for the companion chat-api MCP server derived from NEAR AI config. +pub const NEARAI_COMPANION_MCP_NAME: &str = "chat_api"; + +/// 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 { + if config.llm.backend != "nearai" { + return None; + } + + let base = config.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 { @@ -718,6 +776,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 c31c5051..52d91bf6 100644 --- a/src/tools/mcp/factory.rs +++ b/src/tools/mcp/factory.rs @@ -27,6 +27,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 +80,17 @@ pub async fn create_client_from_config( Err(McpFactoryError::UnixNotSupported { name: server_name }) } EffectiveTransport::Http => { + if server.uses_runtime_auth_source() { + return Ok(McpClient::new_with_config(server) + .with_nearai_session_manager( + nearai_session_manager.expect( + "NearAI companion MCP servers require a 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 +135,8 @@ mod tests { let client = create_client_from_config( server, &session_manager, + None, + None, &process_manager, None, "test-user",