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