From cd617500a8c9517950a1a33f9fd25cb247c444ac Mon Sep 17 00:00:00 2001 From: Coffee Date: Tue, 24 Mar 2026 17:50:23 +0800 Subject: [PATCH] Align NEAR AI companion auth with runtime config --- src/cli/mcp.rs | 200 +++++++++++++++++++++++++++++++++++++- src/extensions/manager.rs | 91 ++++++++++++++++- src/llm/nearai_auth.rs | 44 ++++++++- src/tools/mcp/client.rs | 110 ++++++++++++++++++++- src/tools/mcp/config.rs | 6 +- 5 files changed, 438 insertions(+), 13 deletions(-) diff --git a/src/cli/mcp.rs b/src/cli/mcp.rs index d9c87aeb..095bcc69 100644 --- a/src/cli/mcp.rs +++ b/src/cli/mcp.rs @@ -515,7 +515,7 @@ async fn test_server(name: String, user_id: String) -> anyhow::Result<()> { let session_manager = Arc::new(McpSessionManager::new()); let (client, has_tokens) = if server.uses_runtime_auth_source() { let process_manager = Arc::new(McpProcessManager::new()); - let llm = resolve_llm_from_env()?; + let llm = resolve_llm_for_cli(as_settings_store(db.as_deref())).await?; let nearai_session = crate::llm::create_session_manager(llm.session.clone()).await; ( create_client_from_config( @@ -695,7 +695,7 @@ async fn load_servers_with_derived( ) -> Result { let mut servers = load_persisted_servers(db).await?; - if let Ok(llm) = resolve_llm_from_env() + if let Ok(llm) = resolve_llm_for_cli(as_settings_store(db)).await && let Some(companion) = config::derive_nearai_companion_mcp_server_from_llm(&llm) { servers.insert_if_absent(companion); @@ -726,15 +726,86 @@ async fn get_secrets_store() -> anyhow::Result Result { - let settings = crate::config::load_bootstrap_settings(None)?; +fn as_settings_store( + db: Option<&dyn Database>, +) -> Option<&(dyn crate::db::SettingsStore + Sync)> { + db.map(|db| db as &(dyn crate::db::SettingsStore + Sync)) +} + +async fn resolve_llm_for_cli( + store: Option<&(dyn crate::db::SettingsStore + Sync)>, +) -> Result { + resolve_llm_for_cli_with_toml(store, None).await +} + +async fn resolve_llm_for_cli_with_toml( + store: Option<&(dyn crate::db::SettingsStore + Sync)>, + toml_path: Option<&std::path::Path>, +) -> Result { + if let Some(store) = store { + let _ = dotenvy::dotenv(); + crate::bootstrap::load_ironclaw_env(); + + let mut settings = match store.get_all_settings(DEFAULT_USER_ID).await { + Ok(map) => crate::settings::Settings::from_db_map(&map), + Err(e) => { + tracing::warn!( + "Failed to load CLI settings from DB, falling back to defaults before env/TOML resolution: {}", + e + ); + crate::settings::Settings::default() + } + }; + + apply_cli_toml_overlay(&mut settings, toml_path)?; + return LlmConfig::resolve(&settings); + } + + let settings = crate::config::load_bootstrap_settings(toml_path)?; LlmConfig::resolve(&settings) } +fn apply_cli_toml_overlay( + settings: &mut crate::settings::Settings, + explicit_path: Option<&std::path::Path>, +) -> Result<(), crate::error::ConfigError> { + let path = explicit_path + .map(std::path::PathBuf::from) + .unwrap_or_else(crate::settings::Settings::default_toml_path); + + match crate::settings::Settings::load_toml(&path) { + Ok(Some(toml_settings)) => { + settings.merge_from(&toml_settings); + } + Ok(None) => { + if explicit_path.is_some() { + return Err(crate::error::ConfigError::ParseError(format!( + "Config file not found: {}", + path.display() + ))); + } + } + Err(e) => { + return Err(crate::error::ConfigError::ParseError(e)); + } + } + + Ok(()) +} + #[cfg(test)] mod tests { use super::*; + use std::collections::HashMap; + + use async_trait::async_trait; + + use crate::error::DatabaseError; + use crate::history::SettingRow; + #[cfg(feature = "libsql")] + use tempfile::NamedTempFile; + #[test] fn test_mcp_command_parsing() { // Just verify the command structure is valid @@ -791,4 +862,125 @@ mod tests { assert!(result.is_err()); assert!(result.unwrap_err().contains("invalid env var format")); } + + #[cfg(feature = "libsql")] + #[tokio::test] + async fn test_resolve_llm_for_cli_uses_db_backed_selected_model() { + struct MockSettingsStore { + settings: HashMap, + } + + #[async_trait] + impl crate::db::SettingsStore for MockSettingsStore { + async fn get_setting( + &self, + _user_id: &str, + key: &str, + ) -> Result, DatabaseError> { + Ok(self.settings.get(key).cloned()) + } + + async fn get_setting_full( + &self, + _user_id: &str, + _key: &str, + ) -> Result, DatabaseError> { + Ok(None) + } + + async fn set_setting( + &self, + _user_id: &str, + _key: &str, + _value: &serde_json::Value, + ) -> Result<(), DatabaseError> { + Err(DatabaseError::Query("unused in test".to_string())) + } + + async fn delete_setting( + &self, + _user_id: &str, + _key: &str, + ) -> Result { + Err(DatabaseError::Query("unused in test".to_string())) + } + + async fn list_settings(&self, _user_id: &str) -> Result, DatabaseError> { + Ok(Vec::new()) + } + + async fn get_all_settings( + &self, + _user_id: &str, + ) -> Result, DatabaseError> { + Ok(self.settings.clone()) + } + + async fn set_all_settings( + &self, + _user_id: &str, + _settings: &HashMap, + ) -> Result<(), DatabaseError> { + Err(DatabaseError::Query("unused in test".to_string())) + } + + async fn has_settings(&self, _user_id: &str) -> Result { + Ok(!self.settings.is_empty()) + } + } + + struct EnvGuard(&'static str, Option); + + impl Drop for EnvGuard { + fn drop(&mut self) { + // SAFETY: Protected by ENV_MUTEX for the duration of the test. + unsafe { + match &self.1 { + Some(value) => std::env::set_var(self.0, value), + None => std::env::remove_var(self.0), + } + } + } + } + + let _mutex = crate::config::helpers::ENV_MUTEX.lock().expect("env mutex"); + let prev_backend = std::env::var("LLM_BACKEND").ok(); + let prev_base_url = std::env::var("NEARAI_BASE_URL").ok(); + let prev_auth_url = std::env::var("NEARAI_AUTH_URL").ok(); + let prev_model = std::env::var("NEARAI_MODEL").ok(); + + // SAFETY: Protected by ENV_MUTEX for the duration of the test. + unsafe { + std::env::set_var("LLM_BACKEND", ""); + std::env::set_var("NEARAI_BASE_URL", "http://127.0.0.1:11434/v1"); + std::env::set_var("NEARAI_AUTH_URL", "http://127.0.0.1:11435"); + std::env::set_var("NEARAI_MODEL", ""); + } + + let _backend_guard = EnvGuard("LLM_BACKEND", prev_backend); + let _base_url_guard = EnvGuard("NEARAI_BASE_URL", prev_base_url); + let _auth_url_guard = EnvGuard("NEARAI_AUTH_URL", prev_auth_url); + let _model_guard = EnvGuard("NEARAI_MODEL", prev_model); + + let empty_toml = NamedTempFile::new().expect("temp toml"); + let store = MockSettingsStore { + settings: HashMap::from([ + ("llm_backend".to_string(), serde_json::json!("nearai")), + ( + "selected_model".to_string(), + serde_json::json!("db-backed-nearai-model"), + ), + ]), + }; + + let llm = resolve_llm_for_cli_with_toml(Some(&store), Some(empty_toml.path())) + .await + .expect("resolve llm"); + assert_eq!(llm.backend, "nearai"); + assert_eq!(llm.nearai.model, "db-backed-nearai-model"); + + let companion = config::derive_nearai_companion_mcp_server_from_llm(&llm) + .expect("derived companion"); + assert_eq!(companion.url, "http://127.0.0.1:11434/mcp"); + } } diff --git a/src/extensions/manager.rs b/src/extensions/manager.rs index fa8ef1a5..2bc97f2a 100644 --- a/src/extensions/manager.rs +++ b/src/extensions/manager.rs @@ -1960,7 +1960,7 @@ impl ExtensionManager { return true; } - if let Ok(key) = std::env::var("NEARAI_API_KEY") + if let Some(key) = crate::config::helpers::env_or_override("NEARAI_API_KEY") && !key.trim().is_empty() { return true; @@ -4792,6 +4792,12 @@ impl ExtensionManager { .get_mcp_server(name) .await .map_err(|e| ExtensionError::NotInstalled(e.to_string()))?; + if server.uses_runtime_auth_source() { + return Err(ExtensionError::Other(format!( + "Server '{}' reuses your active NEAR AI authentication and does not accept manually configured MCP tokens", + name + ))); + } let mut names = std::collections::HashSet::new(); names.insert(server.token_secret_name()); names @@ -5844,6 +5850,51 @@ mod tests { ); } + #[tokio::test] + async fn test_runtime_auth_detects_runtime_nearai_api_key_override() { + struct EnvGuard(&'static str, Option); + + impl Drop for EnvGuard { + fn drop(&mut self) { + // SAFETY: Protected by ENV_MUTEX for the duration of the test. + unsafe { + match &self.1 { + Some(value) => std::env::set_var(self.0, value), + None => std::env::remove_var(self.0), + } + } + crate::config::helpers::set_runtime_env(self.0, ""); + } + } + + let _mutex = crate::config::helpers::ENV_MUTEX.lock().expect("env mutex"); + let prev = std::env::var("NEARAI_API_KEY").ok(); + // SAFETY: Protected by ENV_MUTEX for the duration of the test. + unsafe { std::env::remove_var("NEARAI_API_KEY") }; + let _env_guard = EnvGuard("NEARAI_API_KEY", prev); + + crate::config::helpers::set_runtime_env("NEARAI_API_KEY", "runtime-overlay-key"); + + let dir = tempfile::tempdir().expect("temp dir"); + let companion = crate::tools::mcp::config::McpServerConfig::new( + crate::tools::mcp::config::NEARAI_COMPANION_MCP_NAME, + "https://private.near.ai/mcp", + ) + .with_auth_source(crate::tools::mcp::config::McpAuthSource::NearAi); + let manager = make_test_manager_with_dirs( + None, + dir.path().join("tools"), + dir.path().join("channels"), + Some(companion.clone()), + None, + ); + + assert!( + manager.is_runtime_authenticated(&companion).await, + "runtime NEARAI_API_KEY override should count as authenticated" + ); + } + async fn start_runtime_auth_mock_mcp_server() -> (String, tokio::task::JoinHandle<()>) { use axum::extract::State; use axum::http::{HeaderMap, StatusCode}; @@ -7448,6 +7499,44 @@ mod tests { ); } + #[tokio::test] + async fn test_configure_token_rejects_runtime_auth_companion() { + let dir = tempfile::tempdir().expect("temp dir"); + let companion = crate::tools::mcp::config::McpServerConfig::new( + crate::tools::mcp::config::NEARAI_COMPANION_MCP_NAME, + "https://private.near.ai/mcp", + ) + .with_auth_source(crate::tools::mcp::config::McpAuthSource::NearAi); + let token_secret_name = companion.token_secret_name(); + let mgr = make_test_manager_with_dirs( + None, + dir.path().join("tools"), + dir.path().join("channels"), + Some(companion), + None, + ); + + let err = mgr + .configure_token( + crate::tools::mcp::config::NEARAI_COMPANION_MCP_NAME, + "manual-token", + ) + .await + .expect_err("runtime-auth companion should reject manual token configuration"); + + assert!( + err.to_string().contains("active NEAR AI authentication"), + "expected runtime-auth rejection message, got: {err}" + ); + assert!( + !mgr.secrets + .exists("test", &token_secret_name) + .await + .unwrap_or(false), + "configure_token must not persist a manual MCP token for the runtime-auth companion" + ); + } + #[tokio::test] async fn test_auth_is_read_only_for_wasm_channel() { // Regression: auth() must be a pure status check — it must not store diff --git a/src/llm/nearai_auth.rs b/src/llm/nearai_auth.rs index f049bd4d..3f184566 100644 --- a/src/llm/nearai_auth.rs +++ b/src/llm/nearai_auth.rs @@ -20,9 +20,7 @@ pub async fn resolve_nearai_bearer_token_if_available( return Ok(Some(token.expose_secret().to_string())); } - if let Ok(key) = std::env::var("NEARAI_API_KEY") - && !key.is_empty() - { + if let Some(key) = crate::config::helpers::env_or_override("NEARAI_API_KEY") { return Ok(Some(key)); } @@ -54,3 +52,43 @@ pub async fn resolve_nearai_bearer_token( provider: "nearai".to_string(), }) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::helpers::{ENV_MUTEX, set_runtime_env}; + use crate::llm::session::SessionConfig; + + struct EnvGuard(&'static str, Option); + + impl Drop for EnvGuard { + fn drop(&mut self) { + // SAFETY: tests hold ENV_MUTEX while mutating the process environment. + unsafe { + match &self.1 { + Some(value) => std::env::set_var(self.0, value), + None => std::env::remove_var(self.0), + } + } + set_runtime_env(self.0, ""); + } + } + + #[tokio::test] + async fn test_resolve_bearer_token_if_available_uses_runtime_env_override() { + let _guard = ENV_MUTEX.lock().expect("env mutex"); + let prev = std::env::var("NEARAI_API_KEY").ok(); + // SAFETY: tests hold ENV_MUTEX while mutating the process environment. + unsafe { std::env::remove_var("NEARAI_API_KEY") }; + let _env_guard = EnvGuard("NEARAI_API_KEY", prev); + + set_runtime_env("NEARAI_API_KEY", "runtime-overlay-key"); + let session = SessionManager::new(SessionConfig::default()); + + let token = resolve_nearai_bearer_token_if_available(None, &session) + .await + .expect("resolve token"); + + assert_eq!(token.as_deref(), Some("runtime-overlay-key")); + } +} diff --git a/src/tools/mcp/client.rs b/src/tools/mcp/client.rs index 329131bd..5f2fa8de 100644 --- a/src/tools/mcp/client.rs +++ b/src/tools/mcp/client.rs @@ -130,6 +130,9 @@ impl McpClient { /// /// Returns an error if the config uses a non-HTTP transport. pub fn new_with_config(config: McpServerConfig) -> Result { + config + .validate() + .map_err(|e| ToolError::InvalidParameters(e.to_string()))?; if !matches!( config.effective_transport(), crate::tools::mcp::config::EffectiveTransport::Http @@ -281,6 +284,9 @@ impl McpClient { let Some(ref config) = self.server_config else { return Ok(None); }; + if config.uses_runtime_auth_source() { + return Ok(None); + } match secrets .get_decrypted(&self.user_id, &config.token_secret_name()) .await @@ -904,7 +910,10 @@ mod tests { }; use secrecy::SecretString; - let config = McpServerConfig::new("chat_api", "http://localhost:3000/mcp") + let config = McpServerConfig::new( + crate::tools::mcp::config::NEARAI_COMPANION_MCP_NAME, + "http://localhost:3000/mcp", + ) .with_auth_source(crate::tools::mcp::config::McpAuthSource::NearAi); let nearai_session = Arc::new(NearAiSessionManager::new(NearAiSessionConfig::default())); nearai_session @@ -928,7 +937,10 @@ mod tests { SessionConfig as NearAiSessionConfig, SessionManager as NearAiSessionManager, }; - let config = McpServerConfig::new("chat_api", "http://localhost:3000/mcp") + let config = McpServerConfig::new( + crate::tools::mcp::config::NEARAI_COMPANION_MCP_NAME, + "http://localhost:3000/mcp", + ) .with_auth_source(crate::tools::mcp::config::McpAuthSource::NearAi); let nearai_session = Arc::new(NearAiSessionManager::new(NearAiSessionConfig::default())); @@ -943,6 +955,86 @@ mod tests { ); } + #[tokio::test] + async fn test_build_request_headers_runtime_auth_ignores_persisted_mcp_token() { + use crate::llm::{ + SessionConfig as NearAiSessionConfig, SessionManager as NearAiSessionManager, + }; + use crate::secrets::{CreateSecretParams, DecryptedSecret, Secret, SecretError, SecretRef}; + use secrecy::SecretString; + use uuid::Uuid; + + struct PersistedTokenStore; + + #[async_trait] + impl crate::secrets::SecretsStore for PersistedTokenStore { + async fn create( + &self, + _user_id: &str, + _params: CreateSecretParams, + ) -> Result { + unimplemented!() + } + async fn get(&self, _user_id: &str, _name: &str) -> Result { + unimplemented!() + } + async fn get_decrypted( + &self, + _user_id: &str, + _name: &str, + ) -> Result { + DecryptedSecret::from_bytes(b"persisted-mcp-token".to_vec()) + } + async fn exists(&self, _user_id: &str, _name: &str) -> Result { + Ok(true) + } + async fn delete(&self, _user_id: &str, _name: &str) -> Result { + Ok(true) + } + async fn list(&self, _user_id: &str) -> Result, SecretError> { + Ok(Vec::new()) + } + async fn record_usage(&self, _secret_id: Uuid) -> Result<(), SecretError> { + Ok(()) + } + async fn is_accessible( + &self, + _user_id: &str, + _secret_name: &str, + _allowed_secrets: &[String], + ) -> Result { + Ok(true) + } + } + + let config = McpServerConfig::new( + crate::tools::mcp::config::NEARAI_COMPANION_MCP_NAME, + "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_runtime_token")) + .await; + + let secrets: Arc = + Arc::new(PersistedTokenStore); + let client = McpClient::new_authenticated( + config, + Arc::new(McpSessionManager::new()), + secrets, + "test-user", + ) + .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_runtime_token"), + "runtime auth must win even if a persisted MCP token exists" + ); + } + #[test] fn test_next_request_id_monotonically_increasing() { let client = McpClient::new("http://localhost:1234"); @@ -1324,6 +1416,20 @@ mod tests { ); } + #[test] + fn test_new_with_config_rejects_invalid_runtime_auth_name() { + let config = McpServerConfig::new("chat_api", "http://localhost:3000/mcp") + .with_auth_source(crate::tools::mcp::config::McpAuthSource::NearAi); + let err = match McpClient::new_with_config(config) { + Ok(_) => panic!("invalid runtime-auth config must be rejected"), + Err(err) => err.to_string(), + }; + assert!( + err.contains(crate::tools::mcp::config::NEARAI_COMPANION_MCP_NAME), + "error should mention reserved companion requirement: {err}" + ); + } + // --- Issue 13: McpToolWrapper unit tests --- fn make_test_mcp_tool(destructive: bool) -> McpTool { diff --git a/src/tools/mcp/config.rs b/src/tools/mcp/config.rs index eb1a7daf..cd5a18ac 100644 --- a/src/tools/mcp/config.rs +++ b/src/tools/mcp/config.rs @@ -301,7 +301,7 @@ impl McpServerConfig { } } -/// Reserved name used for the companion chat-api MCP server derived from NEAR AI config. +/// Reserved name used for the companion MCP server derived from active NEAR AI config. pub const NEARAI_COMPANION_MCP_NAME: &str = "_nearai_companion_mcp"; pub fn is_nearai_companion_server_name(name: &str) -> bool { @@ -326,7 +326,7 @@ fn strip_reserved_nearai_companion_servers(config: &mut McpServersFile, source: removed } -/// Build the companion chat-api MCP server from the active NearAI config. +/// Build the companion MCP server from the active NEAR AI config. /// /// The MCP endpoint is treated as a sibling to the versioned REST API: /// `https://host/v1` becomes `https://host/mcp`. @@ -336,7 +336,7 @@ pub fn derive_nearai_companion_mcp_server( derive_nearai_companion_mcp_server_from_llm(&config.llm) } -/// Build the companion chat-api MCP server from an LLM config. +/// Build the companion 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