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}