From 92863bf8601ab03b25dce5b72edc854a4855a5a9 Mon Sep 17 00:00:00 2001 From: Illia Polosukhin Date: Fri, 13 Feb 2026 21:46:19 -0800 Subject: [PATCH] fix: resolve libSQL onboarding crash, keychain double-prompt, and setup audit findings Three bugs fixed: 1. libSQL onboarding crash ("Missing required setting 'database_url'"): DatabaseConfig::resolve() only checked DATABASE_BACKEND env var, falling back to Postgres default. Now reads settings.database_backend, plus settings.libsql_path and settings.libsql_url as fallbacks. 2. OS keychain prompts twice during startup: Config::from_env() and Config::from_db() both called get_master_key(). Now caches the key in SECRETS_MASTER_KEY env var after first read so from_db() skips keychain. 3. "Path not found: nearai.session" warning: from_db_map() tried to apply app-specific DB keys (nearai.session_token) to the Settings struct. Now skips keys that don't map to known Settings fields. Also fixed bootstrap migration key mismatch (nearai.session -> nearai.session_token). Setup module audit fixes (14 findings): - Replace unreachable!() with proper error in provider match - Extract setup_api_key_provider() to deduplicate setup_anthropic/setup_openai - Add SAFETY comments to all unsafe std::env::set_var blocks - Fix .unwrap() calls with proper error handling - Remove incorrect #[allow(dead_code)] on used TelegramUpdate::update_id - Log warnings instead of silently discarding HTTP errors in Telegram binding - Guard select_many against empty options, fix mask_api_key for non-ASCII - Update stale doc comment in mod.rs, rename misleading variable - Add 7 new tests (model fetcher fallbacks, channel discovery, secret gen) Co-Authored-By: Claude Opus 4.6 --- .gitignore | 1 + src/bootstrap.rs | 2 +- src/config.rs | 78 ++++- src/main.rs | 131 +++++--- src/settings.rs | 52 ++- src/setup/channels.rs | 31 +- src/setup/mod.rs | 2 +- src/setup/prompts.rs | 5 + src/setup/wizard.rs | 757 +++++++++++++++++++++++++++++++++++++----- 9 files changed, 910 insertions(+), 149 deletions(-) diff --git a/.gitignore b/.gitignore index 0f80f04c..e1863e31 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,7 @@ .env .env.local +.env.* target/ diff --git a/src/bootstrap.rs b/src/bootstrap.rs index e24b996e..752505e7 100644 --- a/src/bootstrap.rs +++ b/src/bootstrap.rs @@ -211,7 +211,7 @@ pub async fn migrate_disk_to_db( Ok(content) => match serde_json::from_str::(&content) { Ok(value) => { store - .set_setting(user_id, "nearai.session", &value) + .set_setting(user_id, "nearai.session_token", &value) .await .map_err(|e| { MigrationError::Database(format!( diff --git a/src/config.rs b/src/config.rs index a6ddb158..82befc5e 100644 --- a/src/config.rs +++ b/src/config.rs @@ -74,7 +74,7 @@ impl Config { settings: &Settings, ) -> Result { Ok(Self { - database: DatabaseConfig::resolve(bootstrap)?, + database: DatabaseConfig::resolve(bootstrap, settings)?, llm: LlmConfig::resolve(settings)?, embeddings: EmbeddingsConfig::resolve(settings)?, tunnel: TunnelConfig::resolve(settings)?, @@ -179,12 +179,18 @@ pub struct DatabaseConfig { } impl DatabaseConfig { - fn resolve(bootstrap: &crate::bootstrap::BootstrapConfig) -> Result { + fn resolve( + bootstrap: &crate::bootstrap::BootstrapConfig, + settings: &Settings, + ) -> Result { + // Priority: env var > settings > default let backend: DatabaseBackend = if let Some(b) = optional_env("DATABASE_BACKEND")? { b.parse().map_err(|e| ConfigError::InvalidValue { key: "DATABASE_BACKEND".to_string(), message: e, })? + } else if let Some(ref b) = settings.database_backend { + b.parse().unwrap_or(DatabaseBackend::default()) } else { DatabaseBackend::default() }; @@ -215,15 +221,19 @@ impl DatabaseConfig { .or(bootstrap.database_pool_size) .unwrap_or(10); - let libsql_path = optional_env("LIBSQL_PATH")?.map(PathBuf::from).or_else(|| { - if backend == DatabaseBackend::LibSql { - Some(default_libsql_path()) - } else { - None - } - }); + // Priority: env var > settings > default (if libsql backend) + let libsql_path = optional_env("LIBSQL_PATH")? + .map(PathBuf::from) + .or_else(|| settings.libsql_path.as_ref().map(PathBuf::from)) + .or_else(|| { + if backend == DatabaseBackend::LibSql { + Some(default_libsql_path()) + } else { + None + } + }); - let libsql_url = optional_env("LIBSQL_URL")?; + let libsql_url = optional_env("LIBSQL_URL")?.or_else(|| settings.libsql_url.clone()); let libsql_auth_token = optional_env("LIBSQL_AUTH_TOKEN")?.map(SecretString::from); if libsql_url.is_some() && libsql_auth_token.is_none() { @@ -401,12 +411,14 @@ pub struct NearAiConfig { impl LlmConfig { fn resolve(settings: &Settings) -> Result { - // Determine backend (default: NearAi) + // Determine backend: env var > settings > default (NearAi) let backend: LlmBackend = if let Some(b) = optional_env("LLM_BACKEND")? { b.parse().map_err(|e| ConfigError::InvalidValue { key: "LLM_BACKEND".to_string(), message: e, })? + } else if let Some(ref b) = settings.llm_backend { + b.parse().unwrap_or(LlmBackend::NearAi) } else { LlmBackend::NearAi }; @@ -473,6 +485,7 @@ impl LlmConfig { let ollama = if backend == LlmBackend::Ollama { let base_url = optional_env("OLLAMA_BASE_URL")? + .or_else(|| settings.ollama_base_url.clone()) .unwrap_or_else(|| "http://localhost:11434".to_string()); let model = optional_env("OLLAMA_MODEL")?.unwrap_or_else(|| "llama3".to_string()); Some(OllamaConfig { base_url, model }) @@ -481,8 +494,9 @@ impl LlmConfig { }; let openai_compatible = if backend == LlmBackend::OpenAiCompatible { - let base_url = - optional_env("LLM_BASE_URL")?.ok_or_else(|| ConfigError::MissingRequired { + let base_url = optional_env("LLM_BASE_URL")? + .or_else(|| settings.openai_compatible_base_url.clone()) + .ok_or_else(|| ConfigError::MissingRequired { key: "LLM_BASE_URL".to_string(), hint: "Set LLM_BASE_URL when LLM_BACKEND=openai_compatible".to_string(), })?; @@ -1351,6 +1365,44 @@ impl ClaudeCodeConfig { } } +/// Load API keys from the encrypted secrets store into process env vars. +/// +/// This bridges the gap between secrets stored during onboarding and the +/// env-var-first resolution in `LlmConfig::resolve()`. Only sets env vars +/// that aren't already present, so explicit env vars always win. +pub async fn inject_llm_keys_from_secrets( + secrets: &dyn crate::secrets::SecretsStore, + user_id: &str, +) { + let mappings = [ + ("llm_openai_api_key", "OPENAI_API_KEY"), + ("llm_anthropic_api_key", "ANTHROPIC_API_KEY"), + ("llm_compatible_api_key", "LLM_API_KEY"), + ]; + + for (secret_name, env_var) in mappings { + if std::env::var(env_var).is_ok() { + continue; + } + match secrets.get_decrypted(user_id, secret_name).await { + Ok(decrypted) => { + // SAFETY: single-threaded at this point in startup + unsafe { + std::env::set_var(env_var, decrypted.expose()); + } + tracing::debug!( + "Injected secret '{}' into env var '{}'", + secret_name, + env_var + ); + } + Err(_) => { + // Secret doesn't exist, that's fine + } + } + } +} + // Helper functions fn optional_env(key: &str) -> Result, ConfigError> { diff --git a/src/main.rs b/src/main.rs index 97355aaa..c566085e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -45,6 +45,7 @@ use ironclaw::secrets::PostgresSecretsStore; use ironclaw::secrets::SecretsCrypto; #[cfg(any(feature = "postgres", feature = "libsql"))] use ironclaw::setup::{SetupConfig, SetupWizard}; +use secrecy::ExposeSecret as _; #[tokio::main] async fn main() -> anyhow::Result<()> { @@ -299,6 +300,19 @@ async fn main() -> anyhow::Result<()> { Err(e) => return Err(e.into()), }; + // If the master key was loaded from the OS keychain, cache it in the process + // env so that later config re-resolution (Config::from_db) doesn't prompt + // the user for keychain access again. + if config.secrets.source == ironclaw::settings::KeySource::Keychain + && std::env::var("SECRETS_MASTER_KEY").is_err() + && let Some(key) = config.secrets.master_key() + { + // SAFETY: Single-threaded at this point in startup. + unsafe { + std::env::set_var("SECRETS_MASTER_KEY", key.expose_secret()); + } + } + // Initialize session manager and authenticate before channel setup let session_config = SessionConfig { auth_base_url: config.llm.nearai.auth_base_url.clone(), @@ -443,6 +457,72 @@ async fn main() -> anyhow::Result<()> { tracing::warn!("Failed to cleanup stale sandbox jobs: {}", e); } } + + // Create secrets store early: needed for injecting LLM API keys from encrypted + // storage before creating the LLM provider, and later for MCP auth + WASM channels. + // + // When both `postgres` and `libsql` features are compiled, the runtime-selected + // backend determines which store is created: whichever DB init branch ran will + // have set its handle (pg_pool or libsql_db), and the or_else chain picks it up. + let secrets_store: Option> = + if let Some(master_key) = config.secrets.master_key() { + match SecretsCrypto::new(master_key.clone()) { + Ok(crypto) => { + let crypto = Arc::new(crypto); + let store: Option> = None; + + #[cfg(feature = "libsql")] + let store = store.or_else(|| { + libsql_db.take().map(|db| { + Arc::new(LibSqlSecretsStore::new(db, Arc::clone(&crypto))) + as Arc + }) + }); + + #[cfg(feature = "postgres")] + let store = store.or_else(|| { + pg_pool.as_ref().map(|pool| { + Arc::new(PostgresSecretsStore::new(pool.clone(), Arc::clone(&crypto))) + as Arc + }) + }); + + store + } + Err(e) => { + tracing::warn!("Failed to initialize secrets crypto: {}", e); + #[cfg(feature = "libsql")] + let _ = libsql_db.take(); + None + } + } + } else { + #[cfg(feature = "libsql")] + let _ = libsql_db.take(); + None + }; + + // Inject LLM API keys from the encrypted secrets store into env vars so that + // LlmConfig::resolve() picks them up. Then re-resolve LlmConfig with the + // newly available keys (backend may have been set during onboarding but the + // API key is in the secrets store, not in env vars). + if let Some(ref secrets) = secrets_store { + ironclaw::config::inject_llm_keys_from_secrets(secrets.as_ref(), "default").await; + + // Re-resolve LlmConfig now that env vars may have been populated + if let Some(ref db_ref) = db { + match Config::from_db(db_ref.as_ref(), "default", &bootstrap).await { + Ok(refreshed) => { + config = refreshed; + tracing::debug!("LlmConfig re-resolved after secret injection"); + } + Err(e) => { + tracing::warn!("Failed to re-resolve config after secret injection: {}", e); + } + } + } + } + // Initialize LLM provider (clone session so we can reuse it for embeddings) let llm = create_llm_provider(&config.llm, session.clone())?; tracing::info!("LLM provider initialized: {}", llm.model_name()); @@ -520,49 +600,6 @@ async fn main() -> anyhow::Result<()> { tracing::info!("Builder mode enabled"); } - // Create secrets store if master key is configured (needed for MCP auth and WASM channels). - // - // When both `postgres` and `libsql` features are compiled, the runtime-selected - // backend determines which store is created: whichever DB init branch ran will - // have set its handle (pg_pool or libsql_db), and the or_else chain picks it up. - let secrets_store: Option> = - if let Some(master_key) = config.secrets.master_key() { - match SecretsCrypto::new(master_key.clone()) { - Ok(crypto) => { - let crypto = Arc::new(crypto); - let store: Option> = None; - - #[cfg(feature = "libsql")] - let store = store.or_else(|| { - libsql_db.take().map(|db| { - Arc::new(LibSqlSecretsStore::new(db, Arc::clone(&crypto))) - as Arc - }) - }); - - #[cfg(feature = "postgres")] - let store = store.or_else(|| { - pg_pool.as_ref().map(|pool| { - Arc::new(PostgresSecretsStore::new(pool.clone(), Arc::clone(&crypto))) - as Arc - }) - }); - - store - } - Err(e) => { - tracing::warn!("Failed to initialize secrets crypto: {}", e); - #[cfg(feature = "libsql")] - let _ = libsql_db.take(); - None - } - } - } else { - #[cfg(feature = "libsql")] - let _ = libsql_db.take(); - None - }; - let mcp_session_manager = Arc::new(McpSessionManager::new()); // Create WASM tool runtime (sync, just builds the wasmtime engine) @@ -1193,9 +1230,13 @@ async fn check_onboard_needed() -> Option<&'static str> { // For now, we don't require it for first run } - // First run (onboarding never completed and no session) + // First run (onboarding never completed and no provider configured) + let settings = ironclaw::settings::Settings::load(); let session_path = ironclaw::llm::session::default_session_path(); - if !bootstrap.onboard_completed && !session_path.exists() { + let has_provider = std::env::var("LLM_BACKEND").is_ok() + || settings.llm_backend.is_some() + || session_path.exists(); + if !bootstrap.onboard_completed && !has_provider { return Some("First run"); } diff --git a/src/settings.rs b/src/settings.rs index 08cc6b6a..48ec1e40 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -40,8 +40,18 @@ pub struct Settings { #[serde(default)] pub secrets_master_key_source: KeySource, - // === Step 3: NEAR AI Auth === - // Session stored separately in session.json + // === Step 3: Inference Provider === + /// LLM backend: "nearai", "anthropic", "openai", "ollama", "openai_compatible". + #[serde(default)] + pub llm_backend: Option, + + /// Ollama base URL (when llm_backend = "ollama"). + #[serde(default)] + pub ollama_base_url: Option, + + /// OpenAI-compatible endpoint base URL (when llm_backend = "openai_compatible"). + #[serde(default)] + pub openai_compatible_base_url: Option, // === Step 4: Model Selection === /// Currently selected model. @@ -512,16 +522,25 @@ impl Settings { /// Each key is a dotted path (e.g., "agent.name"), value is a JSONB value. /// Missing keys get their default value. pub fn from_db_map(map: &std::collections::HashMap) -> Self { - // Start with defaults, then overlay each DB setting + // Start with defaults, then overlay each DB setting. + // + // The settings table stores both Settings struct fields and app-specific + // data (e.g. nearai.session_token). Skip keys that don't correspond to + // a known Settings path. let mut settings = Self::default(); for (key, value) in map { + // Check if this key maps to a known Settings field + if settings.get(key).is_none() { + continue; + } + // Convert the JSONB value to a string for the existing set() method let value_str = match value { serde_json::Value::String(s) => s.clone(), serde_json::Value::Bool(b) => b.to_string(), serde_json::Value::Number(n) => n.to_string(), - serde_json::Value::Null => "null".to_string(), + serde_json::Value::Null => continue, // null means default, skip other => other.to_string(), }; @@ -912,4 +931,29 @@ mod tests { .unwrap(); assert_eq!(settings.channels.telegram_owner_id, Some(987654321)); } + + #[test] + fn test_llm_backend_round_trip() { + let dir = tempdir().unwrap(); + let path = dir.path().join("settings.json"); + + let settings = Settings { + llm_backend: Some("anthropic".to_string()), + ollama_base_url: Some("http://localhost:11434".to_string()), + openai_compatible_base_url: Some("http://my-vllm:8000/v1".to_string()), + ..Default::default() + }; + settings.save_to(&path).unwrap(); + + let loaded = Settings::load_from(&path); + assert_eq!(loaded.llm_backend, Some("anthropic".to_string())); + assert_eq!( + loaded.ollama_base_url, + Some("http://localhost:11434".to_string()) + ); + assert_eq!( + loaded.openai_compatible_base_url, + Some("http://my-vllm:8000/v1".to_string()) + ); + } } diff --git a/src/setup/channels.rs b/src/setup/channels.rs index 77fcc38d..7ac84cbf 100644 --- a/src/setup/channels.rs +++ b/src/setup/channels.rs @@ -107,7 +107,6 @@ struct TelegramGetUpdatesResponse { #[derive(Debug, Deserialize)] struct TelegramUpdate { - #[allow(dead_code)] update_id: i64, message: Option, } @@ -231,7 +230,9 @@ async fn bind_telegram_owner(token: &SecretString) -> Result, String "https://api.telegram.org/bot{}/deleteWebhook", token.expose_secret() ); - let _ = client.post(&delete_url).send().await; + if let Err(e) = client.post(&delete_url).send().await { + tracing::warn!("Failed to delete webhook (getUpdates may not work): {e}"); + } let updates_url = format!( "https://api.telegram.org/bot{}/getUpdates", @@ -282,11 +283,14 @@ async fn bind_telegram_owner(token: &SecretString) -> Result, String "https://api.telegram.org/bot{}/getUpdates", token.expose_secret() ); - let _ = client + if let Err(e) = client .get(&ack_url) .query(&[("offset", &(update.update_id + 1).to_string())]) .send() - .await; + .await + { + tracing::warn!("Failed to acknowledge Telegram update: {e}"); + } return Ok(Some(from.id)); } @@ -593,14 +597,11 @@ pub async fn setup_wasm_channel( print_success(&format!("{} saved to database", secret_config.name)); } - // Optionally validate the configuration + // TODO(#XX): Substitute secrets into the validation URL and make a + // GET request to verify the configured credentials actually work. if let Some(ref validation_endpoint) = setup.validation_endpoint { - print_info("Validating configuration..."); - // The validation endpoint may contain placeholders like {telegram_bot_token} - // For now, we skip validation since we'd need to substitute secrets - // A full implementation would fetch secrets and substitute them print_info(&format!( - "Validation endpoint configured: {} (validation skipped)", + "Validation endpoint configured: {} (validation not yet implemented)", validation_endpoint )); } @@ -631,4 +632,14 @@ mod tests { let secret = generate_webhook_secret(); assert_eq!(secret.len(), 64); // 32 bytes = 64 hex chars } + + #[test] + fn test_generate_secret_with_length() { + let s = generate_secret_with_length(16); + assert_eq!(s.len(), 32); // 16 bytes = 32 hex chars + assert!(s.chars().all(|c| c.is_ascii_hexdigit())); + + let s2 = generate_secret_with_length(1); + assert_eq!(s2.len(), 2); + } } diff --git a/src/setup/mod.rs b/src/setup/mod.rs index ca4d4c56..ef57884a 100644 --- a/src/setup/mod.rs +++ b/src/setup/mod.rs @@ -3,7 +3,7 @@ //! Provides a guided setup experience for: //! 1. Database connection //! 2. Security (secrets master key) -//! 3. NEAR AI authentication +//! 3. Inference provider selection //! 4. Model selection //! 5. Embeddings //! 6. Channel configuration (HTTP, Telegram, etc.) diff --git a/src/setup/prompts.rs b/src/setup/prompts.rs index 7d77c245..ce075572 100644 --- a/src/setup/prompts.rs +++ b/src/setup/prompts.rs @@ -21,6 +21,7 @@ use secrecy::SecretString; /// Display a numbered menu and get user selection. /// /// Returns the index (0-based) of the selected option. +/// Pressing Enter without input selects the first option (index 0). /// /// # Example /// @@ -84,6 +85,10 @@ pub fn select_one(prompt: &str, options: &[&str]) -> io::Result { /// ])?; /// ``` pub fn select_many(prompt: &str, options: &[(&str, bool)]) -> io::Result> { + if options.is_empty() { + return Ok(vec![]); + } + let mut stdout = io::stdout(); let mut selected: Vec = options.iter().map(|(_, s)| *s).collect(); let mut cursor_pos = 0; diff --git a/src/setup/wizard.rs b/src/setup/wizard.rs index 688c3bc6..2b289ccd 100644 --- a/src/setup/wizard.rs +++ b/src/setup/wizard.rs @@ -3,7 +3,7 @@ //! The wizard guides users through: //! 1. Database connection //! 2. Security (secrets master key) -//! 3. NEAR AI authentication +//! 3. Inference provider (NEAR AI, Anthropic, OpenAI, Ollama, OpenAI-compatible) //! 4. Model selection //! 5. Embeddings //! 6. Channel configuration @@ -14,7 +14,7 @@ use std::sync::Arc; #[cfg(feature = "postgres")] use deadpool_postgres::{Config as PoolConfig, Runtime}; -use secrecy::SecretString; +use secrecy::{ExposeSecret, SecretString}; #[cfg(feature = "postgres")] use tokio_postgres::NoTls; @@ -29,7 +29,7 @@ use crate::setup::channels::{ }; use crate::setup::prompts::{ confirm, input, optional_input, print_error, print_header, print_info, print_step, - print_success, select_many, select_one, + print_success, secret_input, select_many, select_one, }; /// Setup wizard error. @@ -132,12 +132,12 @@ impl SetupWizard { print_step(2, total_steps, "Security"); self.step_security().await?; - // Step 3: Authentication (unless skipped) + // Step 3: Inference provider selection (unless skipped) if !self.config.skip_auth { - print_step(3, total_steps, "NEAR AI Authentication"); - self.step_authentication().await?; + print_step(3, total_steps, "Inference Provider"); + self.step_inference_provider().await?; } else { - print_info("Skipping authentication (using existing session)"); + print_info("Skipping inference provider setup (using existing config)"); } // Step 4: Model selection @@ -561,8 +561,67 @@ impl SetupWizard { Ok(()) } - /// Step 3: NEAR AI authentication. - async fn step_authentication(&mut self) -> Result<(), SetupError> { + /// Step 3: Inference provider selection. + /// + /// Lets the user pick from all supported LLM backends, then runs the + /// provider-specific auth sub-flow (API key entry, NEAR AI login, etc.). + async fn step_inference_provider(&mut self) -> Result<(), SetupError> { + // Show current provider if already configured + if let Some(ref current) = self.settings.llm_backend { + let display = match current.as_str() { + "nearai" => "NEAR AI", + "anthropic" => "Anthropic (Claude)", + "openai" => "OpenAI", + "ollama" => "Ollama (local)", + "openai_compatible" => "OpenAI-compatible endpoint", + other => other, + }; + print_info(&format!("Current provider: {}", display)); + println!(); + + if confirm("Keep current provider?", true).map_err(SetupError::Io)? { + // Still run the auth sub-flow in case they need to update keys + match current.as_str() { + "nearai" => return self.setup_nearai().await, + "anthropic" => return self.setup_anthropic().await, + "openai" => return self.setup_openai().await, + "ollama" => return self.setup_ollama(), + "openai_compatible" => return self.setup_openai_compatible().await, + _ => {} + } + return Ok(()); + } + } + + print_info("Select your inference provider:"); + println!(); + + let options = &[ + "NEAR AI - multi-model access via NEAR account", + "Anthropic - Claude models (direct API key)", + "OpenAI - GPT models (direct API key)", + "Ollama - local models, no API key needed", + "OpenAI-compatible - custom endpoint (vLLM, LiteLLM, Together, etc.)", + ]; + + let choice = select_one("Provider:", options).map_err(SetupError::Io)?; + + match choice { + 0 => self.setup_nearai().await?, + 1 => self.setup_anthropic().await?, + 2 => self.setup_openai().await?, + 3 => self.setup_ollama()?, + 4 => self.setup_openai_compatible().await?, + _ => return Err(SetupError::Config("Invalid provider selection".to_string())), + } + + Ok(()) + } + + /// NEAR AI provider setup (extracted from the old step_authentication). + async fn setup_nearai(&mut self) -> Result<(), SetupError> { + self.settings.llm_backend = Some("nearai".to_string()); + // Check if we already have a session if let Some(ref session) = self.session_manager && session.has_token().await @@ -570,7 +629,7 @@ impl SetupWizard { print_info("Existing session found. Validating..."); match session.ensure_authenticated().await { Ok(()) => { - print_success("Session valid"); + print_success("NEAR AI session valid"); return Ok(()); } Err(e) => { @@ -594,10 +653,191 @@ impl SetupWizard { .map_err(|e| SetupError::Auth(e.to_string()))?; self.session_manager = Some(session); + print_success("NEAR AI configured"); + Ok(()) + } + + /// Anthropic provider setup: collect API key and store in secrets. + async fn setup_anthropic(&mut self) -> Result<(), SetupError> { + self.setup_api_key_provider( + "anthropic", + "ANTHROPIC_API_KEY", + "llm_anthropic_api_key", + "Anthropic API key", + "https://console.anthropic.com/settings/keys", + ) + .await + } + + /// OpenAI provider setup: collect API key and store in secrets. + async fn setup_openai(&mut self) -> Result<(), SetupError> { + self.setup_api_key_provider( + "openai", + "OPENAI_API_KEY", + "llm_openai_api_key", + "OpenAI API key", + "https://platform.openai.com/api-keys", + ) + .await + } + + /// Shared setup flow for API-key-based providers (Anthropic, OpenAI). + async fn setup_api_key_provider( + &mut self, + backend: &str, + env_var: &str, + secret_name: &str, + prompt_label: &str, + hint_url: &str, + ) -> Result<(), SetupError> { + let display_name = match backend { + "anthropic" => "Anthropic", + "openai" => "OpenAI", + other => other, + }; + + self.settings.llm_backend = Some(backend.to_string()); + if self.settings.selected_model.is_some() { + self.settings.selected_model = None; + } + + // Check env var first + if let Ok(existing) = std::env::var(env_var) { + print_info(&format!("{env_var} found: {}", mask_api_key(&existing))); + if confirm("Use this key?", true).map_err(SetupError::Io)? { + print_success(&format!("{display_name} configured (from env)")); + return Ok(()); + } + } + + println!(); + print_info(&format!("Get your API key from: {hint_url}")); + println!(); + + let key = secret_input(prompt_label).map_err(SetupError::Io)?; + let key_str = key.expose_secret(); + + if key_str.is_empty() { + return Err(SetupError::Config("API key cannot be empty".to_string())); + } + + // Store in secrets if available + if let Ok(ctx) = self.init_secrets_context().await { + ctx.save_secret(secret_name, &key) + .await + .map_err(|e| SetupError::Config(format!("Failed to save API key: {e}")))?; + print_success("API key encrypted and saved"); + } else { + print_info(&format!( + "Secrets not available. Set {env_var} in your environment." + )); + } + + // SAFETY: Onboarding runs single-threaded before the async runtime spawns workers. + unsafe { + std::env::set_var(env_var, key_str); + } + + print_success(&format!("{display_name} configured")); + Ok(()) + } + + /// Ollama provider setup: just needs a base URL, no API key. + fn setup_ollama(&mut self) -> Result<(), SetupError> { + self.settings.llm_backend = Some("ollama".to_string()); + if self.settings.selected_model.is_some() { + self.settings.selected_model = None; + } + + let default_url = self + .settings + .ollama_base_url + .as_deref() + .unwrap_or("http://localhost:11434"); + + let url_input = optional_input( + "Ollama base URL", + Some(&format!("default: {}", default_url)), + ) + .map_err(SetupError::Io)?; + + let url = url_input.unwrap_or_else(|| default_url.to_string()); + self.settings.ollama_base_url = Some(url.clone()); + + // SAFETY: Onboarding runs single-threaded before the async runtime spawns workers. + unsafe { + std::env::set_var("OLLAMA_BASE_URL", &url); + } + + print_success(&format!("Ollama configured ({})", url)); + Ok(()) + } + + /// OpenAI-compatible provider setup: base URL + optional API key. + async fn setup_openai_compatible(&mut self) -> Result<(), SetupError> { + self.settings.llm_backend = Some("openai_compatible".to_string()); + if self.settings.selected_model.is_some() { + self.settings.selected_model = None; + } + + let existing_url = self + .settings + .openai_compatible_base_url + .clone() + .or_else(|| std::env::var("LLM_BASE_URL").ok()); + + let url = if let Some(ref u) = existing_url { + let url_input = optional_input("Base URL", Some(&format!("current: {}", u))) + .map_err(SetupError::Io)?; + url_input.unwrap_or_else(|| u.clone()) + } else { + input("Base URL (e.g., http://localhost:8000/v1)").map_err(SetupError::Io)? + }; + + if url.is_empty() { + return Err(SetupError::Config( + "Base URL is required for OpenAI-compatible provider".to_string(), + )); + } + + self.settings.openai_compatible_base_url = Some(url.clone()); + // SAFETY: Onboarding runs single-threaded before the async runtime spawns workers. + unsafe { + std::env::set_var("LLM_BASE_URL", &url); + } + + // Optional API key + if confirm("Does this endpoint require an API key?", false).map_err(SetupError::Io)? { + let key = secret_input("API key").map_err(SetupError::Io)?; + let key_str = key.expose_secret(); + + if !key_str.is_empty() { + if let Ok(ctx) = self.init_secrets_context().await { + ctx.save_secret("llm_compatible_api_key", &key) + .await + .map_err(|e| { + SetupError::Config(format!("Failed to save API key: {}", e)) + })?; + print_success("API key encrypted and saved"); + } else { + print_info("Secrets not available. Set LLM_API_KEY in your environment."); + } + + // SAFETY: Onboarding runs single-threaded before the async runtime spawns workers. + unsafe { + std::env::set_var("LLM_API_KEY", key_str); + } + } + } + + print_success(&format!("OpenAI-compatible configured ({})", url)); Ok(()) } /// Step 4: Model selection. + /// + /// Branches on the selected LLM backend and fetches models from the + /// appropriate provider API, with static defaults as fallback. async fn step_model_selection(&mut self) -> Result<(), SetupError> { // Show current model if already configured if let Some(ref current) = self.settings.selected_model { @@ -614,58 +854,97 @@ impl SetupWizard { } } - // Try to fetch available models - let models = if let Some(ref session) = self.session_manager { - self.fetch_available_models(session).await - } else { - vec![] - }; + let backend = self.settings.llm_backend.as_deref().unwrap_or("nearai"); - // Default models if we couldn't fetch - let default_models = [ - ( - "fireworks::accounts/fireworks/models/llama4-maverick-instruct-basic", - "Llama 4 Maverick (default, fast)", - ), - ( - "anthropic::claude-sonnet-4-20250514", - "Claude Sonnet 4 (best quality)", - ), - ("openai::gpt-4o", "GPT-4o"), - ]; + match backend { + "anthropic" => { + let models = fetch_anthropic_models().await; + self.select_from_model_list(&models)?; + } + "openai" => { + let models = fetch_openai_models().await; + self.select_from_model_list(&models)?; + } + "ollama" => { + let base_url = self + .settings + .ollama_base_url + .as_deref() + .unwrap_or("http://localhost:11434"); + let models = fetch_ollama_models(base_url).await; + if models.is_empty() { + print_info("No models found. Pull one first: ollama pull llama3"); + } + self.select_from_model_list(&models)?; + } + "openai_compatible" => { + // No standard API for listing models on arbitrary endpoints + let model_id = input("Model name (e.g., meta-llama/Llama-3-8b-chat-hf)") + .map_err(SetupError::Io)?; + if model_id.is_empty() { + return Err(SetupError::Config("Model name is required".to_string())); + } + self.settings.selected_model = Some(model_id.clone()); + print_success(&format!("Selected {}", model_id)); + } + _ => { + // NEAR AI: use existing provider list_models() + let fetched = self.fetch_nearai_models().await; + let default_models: Vec<(String, String)> = vec![ + ( + "fireworks::accounts/fireworks/models/llama4-maverick-instruct-basic" + .into(), + "Llama 4 Maverick (default, fast)".into(), + ), + ( + "anthropic::claude-sonnet-4-20250514".into(), + "Claude Sonnet 4 (best quality)".into(), + ), + ("openai::gpt-4o".into(), "GPT-4o".into()), + ]; - println!("Available models:"); - println!(); - - let options: Vec<&str> = if models.is_empty() { - default_models.iter().map(|(_, desc)| *desc).collect() - } else { - models.iter().map(|m| m.as_str()).collect() - }; - - // Add custom option - let mut all_options = options.clone(); - all_options.push("Custom model ID"); - - let choice = select_one("Select a model:", &all_options).map_err(SetupError::Io)?; - - let selected_model = if choice == all_options.len() - 1 { - // Custom model - input("Enter model ID").map_err(SetupError::Io)? - } else if models.is_empty() { - default_models[choice].0.to_string() - } else { - models[choice].clone() - }; - - self.settings.selected_model = Some(selected_model.clone()); - print_success(&format!("Selected {}", selected_model)); + let models = if fetched.is_empty() { + default_models + } else { + fetched.iter().map(|m| (m.clone(), m.clone())).collect() + }; + self.select_from_model_list(&models)?; + } + } Ok(()) } - /// Fetch available models from the API. - async fn fetch_available_models(&self, session: &Arc) -> Vec { + /// Present a model list to the user, with a "Custom model ID" escape hatch. + /// + /// Each entry is `(model_id, display_label)`. + fn select_from_model_list(&mut self, models: &[(String, String)]) -> Result<(), SetupError> { + println!("Available models:"); + println!(); + + let mut options: Vec<&str> = models.iter().map(|(_, desc)| desc.as_str()).collect(); + options.push("Custom model ID"); + + let choice = select_one("Select a model:", &options).map_err(SetupError::Io)?; + + let selected = if choice == options.len() - 1 { + input("Enter model ID").map_err(SetupError::Io)? + } else { + models[choice].0.clone() + }; + + self.settings.selected_model = Some(selected.clone()); + print_success(&format!("Selected {}", selected)); + Ok(()) + } + + /// Fetch available models from the NEAR AI API. + async fn fetch_nearai_models(&self) -> Vec { + let session = match self.session_manager { + Some(ref s) => Arc::clone(s), + None => return vec![], + }; + use crate::config::LlmConfig; use crate::llm::create_llm_provider; @@ -690,7 +969,7 @@ impl SetupWizard { openai_compatible: None, }; - match create_llm_provider(&config, Arc::clone(session)) { + match create_llm_provider(&config, session) { Ok(provider) => match provider.list_models().await { Ok(models) => models, Err(e) => { @@ -719,23 +998,51 @@ impl SetupWizard { return Ok(()); } - let options = [ - "NEAR AI (uses same auth, no extra cost)", - "OpenAI (requires API key)", - ]; + let backend = self.settings.llm_backend.as_deref().unwrap_or("nearai"); + let has_openai_key = std::env::var("OPENAI_API_KEY").is_ok(); + let has_nearai = backend == "nearai" || self.session_manager.is_some(); + + // If the LLM backend is OpenAI and we already have a key, default to OpenAI embeddings + if backend == "openai" && has_openai_key { + self.settings.embeddings.enabled = true; + self.settings.embeddings.provider = "openai".to_string(); + self.settings.embeddings.model = "text-embedding-3-small".to_string(); + print_success("Embeddings enabled via OpenAI (using existing API key)"); + return Ok(()); + } + + // If no NEAR AI session and no OpenAI key, only OpenAI is viable + if !has_nearai && !has_openai_key { + print_info("No NEAR AI session or OpenAI key found for embeddings."); + print_info("Set OPENAI_API_KEY in your environment to enable embeddings."); + self.settings.embeddings.enabled = false; + return Ok(()); + } + + let mut options = Vec::new(); + if has_nearai { + options.push("NEAR AI (uses same auth, no extra cost)"); + } + options.push("OpenAI (requires API key)"); let choice = select_one("Select embeddings provider:", &options).map_err(SetupError::Io)?; - match choice { - 0 => { + // Map choice back to provider name + let provider = if has_nearai && choice == 0 { + "nearai" + } else { + "openai" + }; + + match provider { + "nearai" => { self.settings.embeddings.enabled = true; self.settings.embeddings.provider = "nearai".to_string(); self.settings.embeddings.model = "text-embedding-3-small".to_string(); print_success("Embeddings enabled via NEAR AI"); } - 1 => { - // Check if API key is set - if std::env::var("OPENAI_API_KEY").is_err() { + _ => { + if !has_openai_key { print_info("OPENAI_API_KEY not set in environment."); print_info("Add it to your .env file or environment to enable embeddings."); } @@ -744,7 +1051,6 @@ impl SetupWizard { self.settings.embeddings.model = "text-embedding-3-small".to_string(); print_success("Embeddings configured for OpenAI"); } - _ => unreachable!(), } Ok(()) @@ -767,10 +1073,12 @@ impl SetupWizard { )); }; - let crypto = SecretsCrypto::new(SecretString::from(key)) - .map_err(|e| SetupError::Config(e.to_string()))?; - self.secrets_crypto = Some(Arc::new(crypto)); - Arc::clone(self.secrets_crypto.as_ref().unwrap()) + let crypto = Arc::new( + SecretsCrypto::new(SecretString::from(key)) + .map_err(|e| SetupError::Config(e.to_string()))?, + ); + self.secrets_crypto = Some(Arc::clone(&crypto)); + crypto }; // Create backend-appropriate secrets store @@ -813,7 +1121,14 @@ impl SetupWizard { if let Some(url) = url { self.test_database_connection_postgres(&url).await?; self.run_migrations_postgres().await?; - self.db_pool.clone().unwrap() + match self.db_pool.clone() { + Some(pool) => pool, + None => { + return Err(SetupError::Database( + "Database pool not initialized after connection test".to_string(), + )); + } + } } else { return Ok(None); } @@ -1091,6 +1406,18 @@ impl SetupWizard { KeySource::None => println!(" Security: disabled"), } + if let Some(ref provider) = self.settings.llm_backend { + let display = match provider.as_str() { + "nearai" => "NEAR AI", + "anthropic" => "Anthropic", + "openai" => "OpenAI", + "ollama" => "Ollama", + "openai_compatible" => "OpenAI-compatible", + other => other, + }; + println!(" Provider: {}", display); + } + if let Some(ref model) = self.settings.selected_model { // Truncate long model names let display = if model.len() > 40 { @@ -1191,6 +1518,186 @@ fn mask_password_in_url(url: &str) -> String { format!("{}{}:****{}", scheme, username, after_at) } +/// Fetch models from the Anthropic API. +/// +/// Returns `(model_id, display_label)` pairs. Falls back to static defaults on error. +async fn fetch_anthropic_models() -> Vec<(String, String)> { + let static_defaults = vec![ + ("claude-sonnet-4-20250514".into(), "Claude Sonnet 4".into()), + ("claude-opus-4-20250514".into(), "Claude Opus 4".into()), + ( + "claude-3-5-haiku-20241022".into(), + "Claude 3.5 Haiku (fast)".into(), + ), + ]; + + let api_key = match std::env::var("ANTHROPIC_API_KEY") { + Ok(k) if !k.is_empty() => k, + _ => return static_defaults, + }; + + let client = reqwest::Client::new(); + let resp = match client + .get("https://api.anthropic.com/v1/models") + .header("x-api-key", &api_key) + .header("anthropic-version", "2023-06-01") + .timeout(std::time::Duration::from_secs(5)) + .send() + .await + { + Ok(r) if r.status().is_success() => r, + _ => return static_defaults, + }; + + #[derive(serde::Deserialize)] + struct ModelEntry { + id: String, + } + #[derive(serde::Deserialize)] + struct ModelsResponse { + data: Vec, + } + + match resp.json::().await { + Ok(body) => { + let mut models: Vec<(String, String)> = body + .data + .into_iter() + .filter(|m| !m.id.contains("embedding") && !m.id.contains("audio")) + .map(|m| { + let label = m.id.clone(); + (m.id, label) + }) + .collect(); + if models.is_empty() { + return static_defaults; + } + models.sort_by(|a, b| a.0.cmp(&b.0)); + models + } + Err(_) => static_defaults, + } +} + +/// Fetch models from the OpenAI API. +/// +/// Returns `(model_id, display_label)` pairs. Falls back to static defaults on error. +async fn fetch_openai_models() -> Vec<(String, String)> { + let static_defaults = vec![ + ("gpt-4o".into(), "GPT-4o".into()), + ("gpt-4o-mini".into(), "GPT-4o Mini (fast)".into()), + ("o3".into(), "o3 (reasoning)".into()), + ]; + + let api_key = match std::env::var("OPENAI_API_KEY") { + Ok(k) if !k.is_empty() => k, + _ => return static_defaults, + }; + + let client = reqwest::Client::new(); + let resp = match client + .get("https://api.openai.com/v1/models") + .bearer_auth(&api_key) + .timeout(std::time::Duration::from_secs(5)) + .send() + .await + { + Ok(r) if r.status().is_success() => r, + _ => return static_defaults, + }; + + #[derive(serde::Deserialize)] + struct ModelEntry { + id: String, + } + #[derive(serde::Deserialize)] + struct ModelsResponse { + data: Vec, + } + + // Prefixes that indicate chat-relevant models + let chat_prefixes = ["gpt-4", "gpt-3.5", "o1", "o3", "o4", "chatgpt"]; + + match resp.json::().await { + Ok(body) => { + let mut models: Vec<(String, String)> = body + .data + .into_iter() + .filter(|m| { + chat_prefixes.iter().any(|p| m.id.starts_with(p)) + && !m.id.contains("realtime") + && !m.id.contains("audio") + }) + .map(|m| { + let label = m.id.clone(); + (m.id, label) + }) + .collect(); + if models.is_empty() { + return static_defaults; + } + models.sort_by(|a, b| a.0.cmp(&b.0)); + models + } + Err(_) => static_defaults, + } +} + +/// Fetch installed models from a local Ollama instance. +/// +/// Returns `(model_name, display_label)` pairs. Falls back to static defaults on error. +async fn fetch_ollama_models(base_url: &str) -> Vec<(String, String)> { + let static_defaults = vec![ + ("llama3".into(), "llama3".into()), + ("mistral".into(), "mistral".into()), + ("codellama".into(), "codellama".into()), + ]; + + let url = format!("{}/api/tags", base_url.trim_end_matches('/')); + let client = reqwest::Client::new(); + + let resp = match client + .get(&url) + .timeout(std::time::Duration::from_secs(5)) + .send() + .await + { + Ok(r) if r.status().is_success() => r, + Ok(_) => return static_defaults, + Err(_) => { + print_info("Could not connect to Ollama. Is it running?"); + return static_defaults; + } + }; + + #[derive(serde::Deserialize)] + struct ModelEntry { + name: String, + } + #[derive(serde::Deserialize)] + struct TagsResponse { + models: Vec, + } + + match resp.json::().await { + Ok(body) => { + let models: Vec<(String, String)> = body + .models + .into_iter() + .map(|m| { + let label = m.name.clone(); + (m.name, label) + }) + .collect(); + if models.is_empty() { + return static_defaults; + } + models + } + Err(_) => static_defaults, + } +} + /// Discover WASM channels in a directory. /// /// Returns a list of (channel_name, capabilities_file) pairs. @@ -1210,14 +1717,14 @@ async fn discover_wasm_channels(dir: &std::path::Path) -> Vec<(String, ChannelCa let path = entry.path(); // Look for .capabilities.json files - let extension = path.file_name().and_then(|n| n.to_str()).unwrap_or(""); + let filename = path.file_name().and_then(|n| n.to_str()).unwrap_or(""); - if !extension.ends_with(".capabilities.json") { + if !filename.ends_with(".capabilities.json") { continue; } // Extract channel name - let name = extension.trim_end_matches(".capabilities.json").to_string(); + let name = filename.trim_end_matches(".capabilities.json").to_string(); if name.is_empty() { continue; } @@ -1257,6 +1764,20 @@ async fn discover_wasm_channels(dir: &std::path::Path) -> Vec<(String, ChannelCa channels } +/// Mask an API key for display: show first 6 + last 4 chars. +/// +/// Uses char-based indexing to avoid panicking on multi-byte UTF-8. +fn mask_api_key(key: &str) -> String { + let chars: Vec = key.chars().collect(); + if chars.len() < 12 { + let prefix: String = chars.iter().take(4).collect(); + return format!("{prefix}..."); + } + let prefix: String = chars[..6].iter().collect(); + let suffix: String = chars[chars.len() - 4..].iter().collect(); + format!("{prefix}...{suffix}") +} + /// Capitalize the first letter of a string. fn capitalize_first(s: &str) -> String { let mut chars = s.chars(); @@ -1374,6 +1895,20 @@ mod tests { assert_eq!(capitalize_first(""), ""); } + #[test] + fn test_mask_api_key() { + assert_eq!( + mask_api_key("sk-ant-api03-abcdef1234567890"), + "sk-ant...7890" + ); + assert_eq!(mask_api_key("short"), "shor..."); + assert_eq!(mask_api_key("exactly12ch"), "exac..."); + assert_eq!(mask_api_key("exactly12chr"), "exactl...2chr"); + assert_eq!(mask_api_key(""), "..."); + // Multi-byte chars should not panic + assert_eq!(mask_api_key("日本語キー"), "日本語キ..."); + } + #[tokio::test] async fn test_install_missing_bundled_channels_installs_telegram() { use crate::channels::wasm::available_channel_names; @@ -1422,4 +1957,76 @@ mod tests { "telegram should not be duplicated" ); } + + #[tokio::test] + async fn test_fetch_anthropic_models_static_fallback() { + // With no API key, should return static defaults + let _guard = EnvGuard::clear("ANTHROPIC_API_KEY"); + let models = fetch_anthropic_models().await; + assert!(!models.is_empty()); + assert!( + models.iter().any(|(id, _)| id.contains("claude")), + "static defaults should include a Claude model" + ); + } + + #[tokio::test] + async fn test_fetch_openai_models_static_fallback() { + let _guard = EnvGuard::clear("OPENAI_API_KEY"); + let models = fetch_openai_models().await; + assert!(!models.is_empty()); + assert!( + models.iter().any(|(id, _)| id.contains("gpt")), + "static defaults should include a GPT model" + ); + } + + #[tokio::test] + async fn test_fetch_ollama_models_unreachable_fallback() { + // Point at a port nothing listens on + let models = fetch_ollama_models("http://127.0.0.1:1").await; + assert!(!models.is_empty(), "should fall back to static defaults"); + } + + #[tokio::test] + async fn test_discover_wasm_channels_empty_dir() { + let dir = tempdir().unwrap(); + let channels = discover_wasm_channels(dir.path()).await; + assert!(channels.is_empty()); + } + + #[tokio::test] + async fn test_discover_wasm_channels_nonexistent_dir() { + let channels = + discover_wasm_channels(std::path::Path::new("/tmp/ironclaw_nonexistent_dir")).await; + assert!(channels.is_empty()); + } + + /// RAII guard that sets/clears an env var for the duration of a test. + struct EnvGuard { + key: &'static str, + original: Option, + } + + impl EnvGuard { + fn clear(key: &'static str) -> Self { + let original = std::env::var(key).ok(); + unsafe { + std::env::remove_var(key); + } + Self { key, original } + } + } + + impl Drop for EnvGuard { + fn drop(&mut self) { + unsafe { + if let Some(ref val) = self.original { + std::env::set_var(self.key, val); + } else { + std::env::remove_var(self.key); + } + } + } + } }