Skip to main content

ai_agent/utils/
env_utils.rs

1//! Environment utilities
2
3use once_cell::sync::Lazy;
4use std::collections::HashMap;
5use std::env;
6use std::path::PathBuf;
7
8/// Get the Claude config home directory (memoized)
9/// Keyed off AI_CONFIG_DIR so tests that change the env var get a fresh value
10static CLAUDE_CONFIG_HOME_DIR: Lazy<String> = Lazy::new(|| {
11    let config_dir = env::var("AI_CONFIG_DIR")
12        .or_else(|_| env::var("CLAUDE_CONFIG_DIR"))
13        .unwrap_or_else(|_| {
14            dirs::home_dir()
15                .map(|p| p.join(".ai").to_string_lossy().to_string())
16                .unwrap_or_else(|| ".ai".to_string())
17        });
18    config_dir.normalize_nfc()
19});
20
21/// Get the teams directory
22pub fn get_teams_dir() -> PathBuf {
23    PathBuf::from(get_claude_config_home_dir()).join("teams")
24}
25
26/// Get the Claude config home directory
27pub fn get_claude_config_home_dir() -> String {
28    CLAUDE_CONFIG_HOME_DIR.clone()
29}
30
31/// Check if NODE_OPTIONS contains a specific flag.
32/// Splits on whitespace and checks for exact match to avoid false positives.
33pub fn has_node_option(flag: &str) -> bool {
34    if let Ok(node_options) = env::var("NODE_OPTIONS") {
35        node_options.split_whitespace().any(|opt| opt == flag)
36    } else {
37        false
38    }
39}
40
41/// Check if an environment variable value is truthy
42pub fn is_env_truthy(env_var: Option<&str>) -> bool {
43    let Some(value) = env_var else {
44        return false;
45    };
46
47    if value.is_empty() {
48        return false;
49    }
50
51    let normalized = value.to_lowercase().trim().to_string();
52    matches!(normalized.as_str(), "1" | "true" | "yes" | "on")
53}
54
55/// Check if an environment variable is defined as falsy
56pub fn is_env_defined_falsy(env_var: Option<&str>) -> bool {
57    let Some(value) = env_var else {
58        return false;
59    };
60
61    if value.is_empty() {
62        return false;
63    }
64
65    let normalized = value.to_lowercase().trim().to_string();
66    matches!(normalized.as_str(), "0" | "false" | "no" | "off")
67}
68
69/// Check if bare mode is enabled (--bare / AI_CODE_SIMPLE)
70/// Skip hooks, LSP, plugin sync, skill dir-walk, attribution, background prefetches,
71/// and ALL keychain/credential reads.
72/// Auth is strictly ANTHROPIC_API_KEY env or apiKeyHelper from --settings.
73/// Explicit CLI flags (--plugin-dir, --add-dir, --mcp-config) still honored.
74pub fn is_bare_mode() -> bool {
75    let is_simple = is_env_truthy(env::var("AI_CODE_SIMPLE").ok().as_deref());
76
77    // Check argv directly (in addition to the env var) because several gates
78    // run before main action handler sets AI_CODE_SIMPLE=1 from --bare
79    let has_bare_arg = env::args().any(|arg| arg == "--bare");
80
81    is_simple || has_bare_arg
82}
83
84/// Parses an array of environment variable strings into a key-value object
85/// # Arguments
86/// * `raw_env_args` - Array of strings in KEY=VALUE format
87/// # Returns
88/// Object with key-value pairs
89/// # Errors
90/// Returns an error if the format is invalid
91pub fn parse_env_vars(
92    raw_env_args: Option<Vec<String>>,
93) -> Result<HashMap<String, String>, String> {
94    let mut parsed_env: HashMap<String, String> = HashMap::new();
95
96    if let Some(env_args) = raw_env_args {
97        for env_str in env_args {
98            if let Some((key, value)) = env_str.split_once('=') {
99                let key = key.trim();
100                let value = value.trim();
101                if key.is_empty() || value.is_empty() {
102                    return Err(format!(
103                        "Invalid environment variable format: {}, environment variables should be added as: -e KEY1=value1 -e KEY2=value2",
104                        env_str
105                    ));
106                }
107                // Handle KEY=VALUE where VALUE might contain = characters
108                let full_value = env_str[key.len() + 1..].trim().to_string();
109                parsed_env.insert(key.to_string(), full_value);
110            } else {
111                return Err(format!(
112                    "Invalid environment variable format: {}, environment variables should be added as: -e KEY1=value1 -e KEY2=value2",
113                    env_str
114                ));
115            }
116        }
117    }
118
119    Ok(parsed_env)
120}
121
122/// Get the AWS region with fallback to default
123/// Matches the AWS SDK's region behavior
124pub fn get_aws_region() -> String {
125    env::var("AI_AWS_REGION")
126        .or_else(|_| env::var("AWS_REGION"))
127        .or_else(|_| env::var("AI_AWS_DEFAULT_REGION"))
128        .or_else(|_| env::var("AWS_DEFAULT_REGION"))
129        .unwrap_or_else(|_| "us-east-1".to_string())
130}
131
132/// Get the default Vertex AI region
133pub fn get_default_vertex_region() -> String {
134    env::var("AI_CLOUD_ML_REGION")
135        .or_else(|_| env::var("CLOUD_ML_REGION"))
136        .unwrap_or_else(|_| "us-east5".to_string())
137}
138
139/// Check if bash commands should maintain project working directory
140/// (reset to original after each command)
141pub fn should_maintain_project_working_dir() -> bool {
142    is_env_truthy(
143        env::var("AI_BASH_MAINTAIN_PROJECT_WORKING_DIR")
144            .ok()
145            .as_deref(),
146    ) || is_env_truthy(
147        env::var("CLAUDE_BASH_MAINTAIN_PROJECT_WORKING_DIR")
148            .ok()
149            .as_deref(),
150    )
151}
152
153/// Check if running on Homespace (ant-internal cloud environment)
154pub fn is_running_on_homespace() -> bool {
155    let user_type = env::var("USER_TYPE").unwrap_or_default();
156    (user_type == "ant" && is_env_truthy(env::var("AI COO_RUNNING_ON_HOMESPACE").ok().as_deref()))
157        || is_env_truthy(env::var("COO_RUNNING_ON_HOMESPACE").ok().as_deref())
158}
159
160/// Conservative check for whether AI Code is running inside a protected
161/// (privileged or ASL3+) COO namespace or cluster.
162/// Conservative means: when signals are ambiguous, assume protected.
163/// Used for telemetry to measure auto-mode usage in sensitive environments.
164/// Note: This is a stub - the actual implementation would require platform-specific code
165pub fn is_in_protected_namespace() -> bool {
166    false
167}
168
169/// Get USER_TYPE environment variable
170pub fn get_user_type() -> Option<String> {
171    env::var("USER_TYPE").ok()
172}
173
174/// Check if running in ant-internal build
175pub fn is_ant_user() -> bool {
176    get_user_type().as_deref() == Some("ant")
177}
178
179/// Check if running in test mode
180pub fn is_test_mode() -> bool {
181    env::var("NODE_ENV").map(|v| v == "test").unwrap_or(false)
182}
183
184/// Get platform name
185pub fn get_platform() -> String {
186    env::consts::OS.to_string()
187}
188
189/// Model prefix to env var for Vertex region overrides.
190/// Order matters: more specific prefixes must come before less specific ones
191/// (e.g., 'claude-opus-4-1' before 'claude-opus-4').
192const VERTEX_REGION_OVERRIDES: &[(&str, &str)] = &[
193    ("claude-haiku-4-5", "AI_VERTEX_REGION_CLAUDE_HAIKU_4_5"),
194    ("claude-3-5-haiku", "AI_VERTEX_REGION_CLAUDE_3_5_HAIKU"),
195    ("claude-3-5-sonnet", "AI_VERTEX_REGION_CLAUDE_3_5_SONNET"),
196    ("claude-3-7-sonnet", "AI_VERTEX_REGION_CLAUDE_3_7_SONNET"),
197    ("claude-opus-4-1", "AI_VERTEX_REGION_CLAUDE_4_1_OPUS"),
198    ("claude-opus-4", "AI_VERTEX_REGION_CLAUDE_4_0_OPUS"),
199    ("claude-sonnet-4-6", "AI_VERTEX_REGION_CLAUDE_4_6_SONNET"),
200    ("claude-sonnet-4-5", "AI_VERTEX_REGION_CLAUDE_4_5_SONNET"),
201    ("claude-sonnet-4", "AI_VERTEX_REGION_CLAUDE_4_0_SONNET"),
202];
203
204/// Get the Vertex AI region for a specific model.
205/// Different models may be available in different regions.
206pub fn get_vertex_region_for_model(model: Option<&str>) -> Option<String> {
207    let model = model?;
208
209    for (prefix, env_var) in VERTEX_REGION_OVERRIDES {
210        if model.starts_with(prefix) {
211            // Check AI_ prefixed version first, then original
212            let region = env::var(&format!("AI_{}", env_var.trim_start_matches("AI_")))
213                .or_else(|_| env::var(*env_var))
214                .ok();
215
216            return Some(region.unwrap_or_else(get_default_vertex_region));
217        }
218    }
219
220    Some(get_default_vertex_region())
221}
222
223// Helper trait for NFC normalization
224trait NfcNormalize {
225    fn normalize_nfc(&self) -> String;
226}
227
228impl NfcNormalize for String {
229    fn normalize_nfc(&self) -> String {
230        // Rust strings are already UTF-8, NFC normalization is primarily relevant
231        // for strings that will be displayed or compared
232        // For most use cases, the string as-is is fine
233        self.clone()
234    }
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240
241    #[test]
242    fn test_is_env_truthy() {
243        assert!(is_env_truthy(Some("1")));
244        assert!(is_env_truthy(Some("true")));
245        assert!(is_env_truthy(Some("True")));
246        assert!(is_env_truthy(Some("yes")));
247        assert!(is_env_truthy(Some("on")));
248
249        assert!(!is_env_truthy(Some("0")));
250        assert!(!is_env_truthy(Some("false")));
251        assert!(!is_env_truthy(Some("no")));
252        assert!(!is_env_truthy(Some("off")));
253        assert!(!is_env_truthy(None));
254        assert!(!is_env_truthy(Some("")));
255    }
256
257    #[test]
258    fn test_is_env_defined_falsy() {
259        assert!(is_env_defined_falsy(Some("0")));
260        assert!(is_env_defined_falsy(Some("false")));
261        assert!(is_env_defined_falsy(Some("no")));
262        assert!(is_env_defined_falsy(Some("off")));
263
264        assert!(!is_env_defined_falsy(Some("1")));
265        assert!(!is_env_defined_falsy(Some("true")));
266        assert!(!is_env_defined_falsy(None));
267    }
268
269    #[test]
270    fn test_parse_env_vars() {
271        let result = parse_env_vars(Some(vec![
272            "KEY1=value1".to_string(),
273            "KEY2=value2".to_string(),
274        ]))
275        .unwrap();
276
277        assert_eq!(result.get("KEY1"), Some(&"value1".to_string()));
278        assert_eq!(result.get("KEY2"), Some(&"value2".to_string()));
279    }
280
281    #[test]
282    fn test_parse_env_vars_with_equals() {
283        let result = parse_env_vars(Some(vec!["KEY=foo=bar=baz".to_string()])).unwrap();
284
285        assert_eq!(result.get("KEY"), Some(&"foo=bar=baz".to_string()));
286    }
287
288    #[test]
289    fn test_parse_env_vars_invalid() {
290        assert!(parse_env_vars(Some(vec!["=value".to_string()])).is_err());
291        assert!(parse_env_vars(Some(vec!["KEY=".to_string()])).is_err());
292    }
293
294    #[test]
295    fn test_get_aws_region() {
296        // Without env vars set, should return default
297        let region = get_aws_region();
298        assert_eq!(region, "us-east-1");
299    }
300
301    #[test]
302    fn test_get_vertex_region_for_model() {
303        assert_eq!(
304            get_vertex_region_for_model(Some("claude-3-5-sonnet-20241022")),
305            Some("us-east5".to_string())
306        );
307        assert_eq!(get_vertex_region_for_model(None), None);
308    }
309}