diff --git a/src/bootstrap.rs b/src/bootstrap.rs index e186adc1..f8a283f3 100644 --- a/src/bootstrap.rs +++ b/src/bootstrap.rs @@ -116,9 +116,18 @@ pub fn load_ironclaw_env() { .join(".ironclaw") .join("ironclaw.db"); if default_db.exists() { - // SAFETY: `load_ironclaw_env` is called from a synchronous `fn main()` - // before the Tokio runtime is started, so no other threads exist yet. - unsafe { std::env::set_var("DATABASE_BACKEND", "libsql") }; + if tokio::runtime::Handle::try_current().is_ok() { + // Tokio runtime is active (multi-threaded); std::env::set_var is UB here. + // Fall back to the thread-safe runtime overlay so the value is always set. + tracing::warn!( + "load_ironclaw_env called with active Tokio runtime; \ + using runtime env overlay for DATABASE_BACKEND" + ); + crate::config::set_runtime_env("DATABASE_BACKEND", "libsql"); + } else { + // SAFETY: No Tokio runtime = no other threads = safe to call set_var. + unsafe { std::env::set_var("DATABASE_BACKEND", "libsql") }; + } } } } diff --git a/src/cli/doctor.rs b/src/cli/doctor.rs index aa47b6bf..f6e221fb 100644 --- a/src/cli/doctor.rs +++ b/src/cli/doctor.rs @@ -220,7 +220,7 @@ async fn check_nearai_session() -> CheckResult { let session_path = crate::config::llm::default_session_path(); if !session_path.exists() { // Check for API key mode - if std::env::var("NEARAI_API_KEY").is_ok() { + if crate::config::helpers::env_or_override("NEARAI_API_KEY").is_some() { return CheckResult::Pass("API key configured".into()); } return CheckResult::Fail(format!( diff --git a/src/config/helpers.rs b/src/config/helpers.rs index d6521b38..ce6ce092 100644 --- a/src/config/helpers.rs +++ b/src/config/helpers.rs @@ -1,6 +1,9 @@ +use std::collections::HashMap; +use std::sync::{Mutex, OnceLock}; + use crate::error::ConfigError; -use super::INJECTED_VARS; +use crate::config::INJECTED_VARS; /// Crate-wide mutex for tests that mutate process environment variables. /// @@ -11,6 +14,73 @@ use super::INJECTED_VARS; #[cfg(test)] pub(crate) static ENV_MUTEX: std::sync::Mutex<()> = std::sync::Mutex::new(()); +/// Thread-safe mutable overlay for env vars set at runtime. +/// +/// Unlike `INJECTED_VARS` (which is set once at startup from the secrets +/// store), this map supports writes at any point during the process +/// lifetime. It replaces unsafe `std::env::set_var` calls that would +/// otherwise be UB in multi-threaded programs (Rust 1.82+). +/// +/// Priority: real env vars > `RUNTIME_ENV_OVERRIDES` > `INJECTED_VARS`. +static RUNTIME_ENV_OVERRIDES: OnceLock>> = OnceLock::new(); + +fn runtime_overrides() -> &'static Mutex> { + RUNTIME_ENV_OVERRIDES.get_or_init(|| Mutex::new(HashMap::new())) +} + +/// Set a runtime environment override (thread-safe alternative to `std::env::set_var`). +/// +/// Values set here are visible to `optional_env()`, `env_or_override()`, and +/// all config resolution that goes through those helpers. This avoids the UB +/// of `std::env::set_var` in multi-threaded programs. +pub fn set_runtime_env(key: &str, value: &str) { + runtime_overrides() + .lock() + .unwrap_or_else(|e| e.into_inner()) + .insert(key.to_string(), value.to_string()); +} + +/// Read an env var, checking the real environment first, then runtime overrides. +/// +/// Priority: real env vars > runtime overrides > `INJECTED_VARS`. +/// Empty values are treated as unset at every layer for consistency with +/// `optional_env()`. +/// +/// Use this instead of `std::env::var()` when the value might have been set +/// via `set_runtime_env()` (e.g., `NEARAI_API_KEY` during interactive login). +pub fn env_or_override(key: &str) -> Option { + // Real env vars always win + if let Ok(val) = std::env::var(key) + && !val.is_empty() + { + return Some(val); + } + + // Check runtime overrides (skip empty values for consistency with optional_env) + if let Some(val) = runtime_overrides() + .lock() + .unwrap_or_else(|e| e.into_inner()) + .get(key) + .filter(|v| !v.is_empty()) + .cloned() + { + return Some(val); + } + + // Check INJECTED_VARS (secrets from DB, set once at startup) + if let Some(val) = INJECTED_VARS + .lock() + .unwrap_or_else(|e| e.into_inner()) + .get(key) + .filter(|v| !v.is_empty()) + .cloned() + { + return Some(val); + } + + None +} + pub(crate) fn optional_env(key: &str) -> Result, ConfigError> { // Check real env vars first (always win over injected secrets) match std::env::var(key) { @@ -24,6 +94,17 @@ pub(crate) fn optional_env(key: &str) -> Result, ConfigError> { } } + // Fall back to runtime overrides (set via set_runtime_env) + if let Some(val) = runtime_overrides() + .lock() + .unwrap_or_else(|e| e.into_inner()) + .get(key) + .filter(|v| !v.is_empty()) + .cloned() + { + return Ok(Some(val)); + } + // Fall back to thread-safe overlay (secrets injected from DB) if let Some(val) = INJECTED_VARS .lock() @@ -94,3 +175,55 @@ pub(crate) fn parse_string_env( ) -> Result { Ok(optional_env(key)?.unwrap_or_else(|| default.into())) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn runtime_env_override_is_visible_to_env_or_override() { + // Use a unique key that won't collide with real env vars. + let key = "IRONCLAW_TEST_RUNTIME_OVERRIDE_42"; + + // Not set initially + assert!(env_or_override(key).is_none()); + + // Set via the thread-safe overlay + set_runtime_env(key, "test_value"); + + // Now visible + assert_eq!(env_or_override(key), Some("test_value".to_string())); + } + + #[test] + fn runtime_env_override_is_visible_to_optional_env() { + let key = "IRONCLAW_TEST_OPTIONAL_ENV_OVERRIDE_42"; + + assert_eq!(optional_env(key).unwrap(), None); + + set_runtime_env(key, "hello"); + + assert_eq!(optional_env(key).unwrap(), Some("hello".to_string())); + } + + #[test] + fn real_env_var_takes_priority_over_runtime_override() { + let _guard = ENV_MUTEX.lock().unwrap(); + let key = "IRONCLAW_TEST_ENV_PRIORITY_42"; + + // Set runtime override + set_runtime_env(key, "override_value"); + + // Set real env var (should win) + // SAFETY: test runs under ENV_MUTEX + unsafe { std::env::set_var(key, "real_value") }; + + assert_eq!(env_or_override(key), Some("real_value".to_string())); + + // Clean up + unsafe { std::env::remove_var(key) }; + + // Now the runtime override is visible again + assert_eq!(env_or_override(key), Some("override_value".to_string())); + } +} diff --git a/src/config/mod.rs b/src/config/mod.rs index 1bdd446e..c6952897 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -54,6 +54,10 @@ pub use crate::llm::config::{ }; pub use crate::llm::session::SessionConfig; +// Thread-safe env var override helpers (replaces unsafe `std::env::set_var` +// for mid-process env mutations in multi-threaded contexts). +pub use self::helpers::{env_or_override, set_runtime_env}; + /// Thread-safe overlay for injected env vars (secrets loaded from DB). /// /// Used by `inject_llm_keys_from_secrets()` to make API keys available to diff --git a/src/llm/session.rs b/src/llm/session.rs index 1cb858a1..49f7cb7a 100644 --- a/src/llm/session.rs +++ b/src/llm/session.rs @@ -373,9 +373,10 @@ impl SessionManager { /// NEAR AI Cloud API key entry flow. /// /// Prompts the user to enter a NEAR AI Cloud API key from - /// cloud.near.ai. The key is set as `NEARAI_API_KEY` env var so - /// `LlmConfig::resolve()` auto-selects ChatCompletions mode, and - /// saved to `~/.ironclaw/.env` for persistence across restarts. + /// cloud.near.ai. The key is stored in the thread-safe runtime + /// env overlay (via `set_runtime_env`) so `LlmConfig::resolve()` + /// auto-selects ChatCompletions mode, and persisted to + /// `~/.ironclaw/.env` for survival across restarts. /// No session token is saved and no `/v1/users/me` validation is /// performed (different auth model). async fn api_key_login(&self) -> Result<(), LlmError> { @@ -403,15 +404,11 @@ impl SessionManager { }); } - // Set env var so Config picks it up immediately - // (LlmConfig::resolve() auto-selects ChatCompletions mode when - // NEARAI_API_KEY is present). - // - // SAFETY: called during single-threaded interactive login flow. - #[allow(unused_unsafe)] - unsafe { - std::env::set_var("NEARAI_API_KEY", &key); - } + // Make the key visible to Config resolution and `env_or_override()` + // callers for the remainder of this process. Uses a thread-safe + // overlay instead of `std::env::set_var`, which is UB in + // multi-threaded programs (Rust 1.82+). + crate::config::helpers::set_runtime_env("NEARAI_API_KEY", &key); // Persist to ~/.ironclaw/.env so the key survives restarts // (bootstrap layer — available before DB is connected). diff --git a/src/setup/wizard.rs b/src/setup/wizard.rs index 38066d25..89d1e5be 100644 --- a/src/setup/wizard.rs +++ b/src/setup/wizard.rs @@ -1183,9 +1183,9 @@ impl SetupWizard { self.persist_session_to_db().await; // If the user chose the API key path, NEARAI_API_KEY is now set - // in the environment. Persist it to the encrypted secrets store - // so inject_llm_keys_from_secrets() can load it on future runs. - if let Ok(api_key) = std::env::var("NEARAI_API_KEY") + // in the runtime env overlay. Persist it to the encrypted secrets + // store so inject_llm_keys_from_secrets() can load it on future runs. + if let Some(api_key) = crate::config::helpers::env_or_override("NEARAI_API_KEY") && !api_key.is_empty() && let Ok(ctx) = self.init_secrets_context().await { @@ -2613,8 +2613,9 @@ impl SetupWizard { env_vars.push((base_url_env.clone(), base_url.clone())); } - // Preserve NEARAI_API_KEY if present (set by API key auth flow) - if let Ok(api_key) = std::env::var("NEARAI_API_KEY") + // Preserve NEARAI_API_KEY if present (set by API key auth flow + // via the thread-safe runtime env overlay). + if let Some(api_key) = crate::config::helpers::env_or_override("NEARAI_API_KEY") && !api_key.is_empty() { env_vars.push(("NEARAI_API_KEY".to_string(), api_key));