Skip to main content

mermaid_cli/utils/
auth.rs

1//! Provider API-key resolution.
2//!
3//! Mermaid's auth surface is uniform across providers: an API key lives in
4//! an environment variable, with the option to override the variable name
5//! per-provider in `config.toml`. There's no in-config secret storage —
6//! keys never sit on disk in plaintext. The Ollama cloud key follows the same
7//! rule: it is read from `OLLAMA_API_KEY` and never persisted (#88).
8
9/// Resolve an API key from the environment.
10///
11/// `default_env` is the env var the built-in registry expects (e.g.
12/// `"GROQ_API_KEY"`). `override_env` is an optional per-provider
13/// override from `config.toml`'s `[providers.<name>] api_key_env = ...`.
14/// When set, it takes precedence — a user who's already standardized on
15/// `LLM_API_KEY` for everything can point all their providers at it.
16///
17/// Empty values are treated as unset (matches the existing
18/// `get_cloud_api_key` semantics).
19pub fn resolve_api_key(default_env: &str, override_env: Option<&str>) -> Option<String> {
20    let env_var = override_env.unwrap_or(default_env);
21    match std::env::var(env_var) {
22        Ok(key) if !key.is_empty() => Some(key),
23        _ => None,
24    }
25}
26
27/// Resolve an API key with a legacy fallback env var.
28///
29/// If `override_env` is set, it remains authoritative and no fallback
30/// is attempted. Without an override, `default_env` is checked first
31/// and `fallback_env` is accepted only when the default is unset.
32pub fn resolve_api_key_with_fallback(
33    default_env: &str,
34    fallback_env: &str,
35    override_env: Option<&str>,
36) -> Option<String> {
37    if override_env.is_some() {
38        return resolve_api_key(default_env, override_env);
39    }
40    resolve_api_key(default_env, None).or_else(|| resolve_api_key(fallback_env, None))
41}
42
43#[cfg(test)]
44mod tests {
45    use super::*;
46
47    /// Generate a unique env var name per test so concurrent test runs
48    /// don't step on each other's process-global env state. `temp_env`
49    /// restores the prior value after the closure, but a collision with
50    /// a concurrent test from another module on a shared name would
51    /// still race — unique names are belt-and-braces.
52    fn unique_env(prefix: &str) -> String {
53        use std::sync::atomic::{AtomicUsize, Ordering};
54        static N: AtomicUsize = AtomicUsize::new(0);
55        format!(
56            "{}_{}_{}",
57            prefix,
58            std::process::id(),
59            N.fetch_add(1, Ordering::SeqCst)
60        )
61    }
62
63    #[test]
64    fn returns_none_when_env_var_unset() {
65        let var = unique_env("MERMAID_TEST_AUTH_UNSET");
66        temp_env::with_var_unset(&var, || {
67            assert_eq!(resolve_api_key(&var, None), None);
68        });
69    }
70
71    #[test]
72    fn returns_value_when_env_var_set() {
73        let var = unique_env("MERMAID_TEST_AUTH_SET");
74        temp_env::with_var(&var, Some("secret-value"), || {
75            assert_eq!(
76                resolve_api_key(&var, None),
77                Some("secret-value".to_string())
78            );
79        });
80    }
81
82    #[test]
83    fn empty_string_treated_as_unset() {
84        let var = unique_env("MERMAID_TEST_AUTH_EMPTY");
85        temp_env::with_var(&var, Some(""), || {
86            assert_eq!(resolve_api_key(&var, None), None);
87        });
88    }
89
90    #[test]
91    fn override_env_takes_precedence_over_default() {
92        let default_var = unique_env("MERMAID_TEST_AUTH_DEFAULT");
93        let override_var = unique_env("MERMAID_TEST_AUTH_OVERRIDE");
94        temp_env::with_vars(
95            [
96                (default_var.as_str(), Some("default-key")),
97                (override_var.as_str(), Some("override-key")),
98            ],
99            || {
100                let resolved = resolve_api_key(&default_var, Some(&override_var));
101                assert_eq!(resolved, Some("override-key".to_string()));
102            },
103        );
104    }
105
106    #[test]
107    fn override_env_unset_falls_through_to_none() {
108        // When the override env name is provided but unset, we DON'T fall
109        // back to the default — the user explicitly asked for a different
110        // var and got nothing. Better to fail loudly than silently use a
111        // key the user thought they'd disabled.
112        let default_var = unique_env("MERMAID_TEST_AUTH_DEFAULT2");
113        let override_var = unique_env("MERMAID_TEST_AUTH_OVERRIDE2");
114        temp_env::with_vars(
115            [
116                (default_var.as_str(), Some("default-key")),
117                (override_var.as_str(), None),
118            ],
119            || {
120                let resolved = resolve_api_key(&default_var, Some(&override_var));
121                assert_eq!(resolved, None);
122            },
123        );
124    }
125
126    #[test]
127    fn fallback_env_used_only_when_default_is_unset() {
128        let default_var = unique_env("MERMAID_TEST_AUTH_FALLBACK_DEFAULT");
129        let fallback_var = unique_env("MERMAID_TEST_AUTH_FALLBACK_LEGACY");
130        temp_env::with_vars(
131            [
132                (default_var.as_str(), None),
133                (fallback_var.as_str(), Some("legacy-key")),
134            ],
135            || {
136                let resolved = resolve_api_key_with_fallback(&default_var, &fallback_var, None);
137                assert_eq!(resolved, Some("legacy-key".to_string()));
138            },
139        );
140
141        temp_env::with_vars(
142            [
143                (default_var.as_str(), Some("default-key")),
144                (fallback_var.as_str(), Some("legacy-key")),
145            ],
146            || {
147                let resolved = resolve_api_key_with_fallback(&default_var, &fallback_var, None);
148                assert_eq!(resolved, Some("default-key".to_string()));
149            },
150        );
151    }
152
153    #[test]
154    fn fallback_env_is_ignored_when_override_is_set() {
155        let default_var = unique_env("MERMAID_TEST_AUTH_OVERRIDE_DEFAULT3");
156        let fallback_var = unique_env("MERMAID_TEST_AUTH_OVERRIDE_FALLBACK3");
157        let override_var = unique_env("MERMAID_TEST_AUTH_OVERRIDE3");
158        temp_env::with_vars(
159            [
160                (default_var.as_str(), Some("default-key")),
161                (fallback_var.as_str(), Some("legacy-key")),
162                (override_var.as_str(), None),
163            ],
164            || {
165                let resolved =
166                    resolve_api_key_with_fallback(&default_var, &fallback_var, Some(&override_var));
167                assert_eq!(resolved, None);
168            },
169        );
170    }
171}