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#[cfg(test)]
51mod tests {
52    use super::*;
53
54    #[test]
55    fn test_set_and_get() {
56        set("TEST_RUNTIME_KEY", "hello");
57        assert_eq!(get("TEST_RUNTIME_KEY"), Some("hello".to_string()));
58    }
59
60    #[test]
61    fn test_remove() {
62        set("TEST_REMOVE_KEY", "value");
63        assert!(is_set("TEST_REMOVE_KEY"));
64        env_map().write().unwrap().remove("TEST_REMOVE_KEY");
65        // May still exist in process env, but runtime map entry is gone
66    }
67
68    #[test]
69    fn test_falls_back_to_env() {
70        // PATH should exist in the real environment
71        assert!(get("PATH").is_some());
72    }
73
74    #[test]
75    fn test_runtime_takes_precedence() {
76        set("PATH", "overridden");
77        assert_eq!(get("PATH"), Some("overridden".to_string()));
78        // Clean up
79        env_map().write().unwrap().remove("PATH");
80    }
81
82    #[test]
83    fn test_missing_key() {
84        assert!(get("DEFINITELY_NOT_SET_12345").is_none());
85        assert!(!is_set("DEFINITELY_NOT_SET_12345"));
86    }
87}