Skip to main content

koda_core/
runtime_env.rs

1//! Thread-safe runtime environment for API keys and config.
2//!
3//! Replaces `unsafe { std::env::set_var() }` with a concurrent map
4//! that is safe to read/write from any tokio task.
5//!
6//! Read priority: runtime map → process environment.
7//!
8//! ## Why not `std::env::set_var`?
9//!
10//! `set_var` is unsafe in multi-threaded programs (undefined behavior in
11//! Rust 2024 edition). Since Koda uses tokio with multiple tasks (main
12//! REPL, background agents, version checker), we need a thread-safe
13//! alternative.
14
15use std::collections::HashMap;
16use std::sync::{OnceLock, RwLock};
17
18static RUNTIME_ENV: OnceLock<RwLock<HashMap<String, String>>> = OnceLock::new();
19
20fn env_map() -> &'static RwLock<HashMap<String, String>> {
21    RUNTIME_ENV.get_or_init(|| RwLock::new(HashMap::new()))
22}
23
24/// Set a runtime environment variable (thread-safe).
25pub fn set(key: impl Into<String>, value: impl Into<String>) {
26    env_map()
27        .write()
28        .unwrap_or_else(|poisoned| poisoned.into_inner())
29        .insert(key.into(), value.into());
30}
31
32/// Get a runtime variable, falling back to `std::env::var`.
33/// Checks our runtime map first, then the real process environment.
34pub fn get(key: &str) -> Option<String> {
35    if let Some(val) = env_map()
36        .read()
37        .unwrap_or_else(|poisoned| poisoned.into_inner())
38        .get(key)
39    {
40        return Some(val.clone());
41    }
42    std::env::var(key).ok()
43}
44
45/// Check if a runtime variable is set (in either runtime map or process env).
46pub fn is_set(key: &str) -> bool {
47    get(key).is_some()
48}
49
50/// Remove a key from the runtime map (does **not** touch process env).
51///
52/// Returns the previous runtime-map value if one was set. After removal
53/// [`get`] will fall back to the process environment as usual.
54///
55/// Primarily intended for tests that need to leave the runtime map clean
56/// for siblings — production code should rarely need this.
57pub fn remove(key: &str) -> Option<String> {
58    env_map()
59        .write()
60        .unwrap_or_else(|poisoned| poisoned.into_inner())
61        .remove(key)
62}
63
64#[cfg(test)]
65mod tests {
66    use super::*;
67
68    #[test]
69    fn test_set_and_get() {
70        set("TEST_RUNTIME_KEY", "hello");
71        assert_eq!(get("TEST_RUNTIME_KEY"), Some("hello".to_string()));
72    }
73
74    #[test]
75    fn test_remove() {
76        set("TEST_REMOVE_KEY", "value");
77        assert!(is_set("TEST_REMOVE_KEY"));
78        env_map().write().unwrap().remove("TEST_REMOVE_KEY");
79        // May still exist in process env, but runtime map entry is gone
80    }
81
82    #[test]
83    fn test_falls_back_to_env() {
84        // PATH should exist in the real environment
85        assert!(get("PATH").is_some());
86    }
87
88    #[test]
89    fn test_runtime_takes_precedence() {
90        set("PATH", "overridden");
91        assert_eq!(get("PATH"), Some("overridden".to_string()));
92        // Clean up
93        env_map().write().unwrap().remove("PATH");
94    }
95
96    #[test]
97    fn test_remove_returns_previous_value() {
98        set("TEST_REMOVE_RETURN", "orig");
99        assert_eq!(remove("TEST_REMOVE_RETURN"), Some("orig".to_string()));
100        // Second remove returns None (already gone from runtime map).
101        assert_eq!(remove("TEST_REMOVE_RETURN"), None);
102    }
103
104    #[test]
105    fn test_remove_does_not_touch_process_env() {
106        // remove() should only clear the runtime map, not std::env.
107        // Use HOME (always set in CI + dev) instead of PATH because the
108        // sibling test_runtime_takes_precedence also touches PATH and
109        // cargo runs unit tests in parallel by default — racing on the
110        // same key would be a flaky-test factory.
111        let _ = remove("HOME"); // ensure runtime map has no override
112        assert!(
113            get("HOME").is_some(),
114            "process env HOME must survive a runtime-map remove"
115        );
116    }
117}