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