ai_usagebar/
config.rs

1//! Config file at `~/.config/ai-usagebar/config.toml`.
2//!
3//! Layout:
4//! ```toml
5//! [anthropic] enabled = true
6//! [openai]    enabled = true   # Codex OAuth from ~/.codex/auth.json
7//! [zai]       enabled = true
8//! [openrouter] enabled = true
9//! ```
10//!
11//! Every field is optional with sensible defaults — missing config file is
12//! treated as "use defaults". API keys are read from env vars (the relevant
13//! `*_api_key_env` field lets the user override which env var name).
14
15use std::path::PathBuf;
16
17use serde::{Deserialize, Serialize};
18
19use crate::error::{AppError, Result};
20use crate::vendor::VendorId;
21
22#[derive(Debug, Clone, Default, Deserialize, Serialize)]
23#[serde(default)]
24pub struct Config {
25    pub ui: UiConfig,
26    pub anthropic: AnthropicConfig,
27    pub openai: OpenAiConfig,
28    pub zai: ZaiConfig,
29    pub openrouter: OpenRouterConfig,
30}
31
32/// UI / dispatch preferences. Currently just `primary` — which vendor the
33/// widget shows when `--vendor` is omitted, and which TUI tab is selected
34/// at startup.
35#[derive(Debug, Clone, Default, Deserialize, Serialize)]
36#[serde(default)]
37pub struct UiConfig {
38    /// `None` → fall back to anthropic for backward compatibility.
39    pub primary: Option<VendorId>,
40}
41
42#[derive(Debug, Clone, Deserialize, Serialize)]
43#[serde(default)]
44pub struct AnthropicConfig {
45    pub enabled: bool,
46    /// Override the credentials file path (defaults to `~/.claude/.credentials.json`).
47    pub credentials_path: Option<PathBuf>,
48}
49
50impl Default for AnthropicConfig {
51    fn default() -> Self {
52        Self {
53            enabled: true,
54            credentials_path: None,
55        }
56    }
57}
58
59#[derive(Debug, Clone, Deserialize, Serialize)]
60#[serde(default)]
61pub struct OpenAiConfig {
62    pub enabled: bool,
63    /// Override the Codex auth file path (defaults to `~/.codex/auth.json`).
64    pub codex_auth_path: Option<PathBuf>,
65    /// Optional admin key env var name for the API-key-only fallback path.
66    pub admin_key_env: String,
67}
68
69impl Default for OpenAiConfig {
70    fn default() -> Self {
71        Self {
72            enabled: true,
73            codex_auth_path: None,
74            admin_key_env: "OPENAI_ADMIN_KEY".to_string(),
75        }
76    }
77}
78
79#[derive(Debug, Clone, Deserialize, Serialize)]
80#[serde(default)]
81pub struct ZaiConfig {
82    pub enabled: bool,
83    /// Env var name to read the key from (env wins over `api_key`).
84    pub api_key_env: String,
85    /// Inline key (fallback when the env var is unset). Chmod 600 your
86    /// config file if you put a real key here.
87    pub api_key: Option<String>,
88    /// Optional plan tier label (lite/pro/max) — display-only.
89    pub plan_tier: Option<String>,
90}
91
92impl Default for ZaiConfig {
93    fn default() -> Self {
94        Self {
95            enabled: true,
96            api_key_env: "ZAI_API_KEY".to_string(),
97            api_key: None,
98            plan_tier: None,
99        }
100    }
101}
102
103#[derive(Debug, Clone, Deserialize, Serialize)]
104#[serde(default)]
105pub struct OpenRouterConfig {
106    pub enabled: bool,
107    pub api_key_env: String,
108    pub api_key: Option<String>,
109}
110
111impl Default for OpenRouterConfig {
112    fn default() -> Self {
113        Self {
114            enabled: true,
115            api_key_env: "OPENROUTER_API_KEY".to_string(),
116            api_key: None,
117        }
118    }
119}
120
121/// Resolve an API key for a vendor: env var wins, then inline config, then
122/// a clear error naming both fields. Used by Z.AI and OpenRouter vendors.
123pub fn resolve_api_key(
124    vendor_label: &str,
125    env_var_name: &str,
126    inline: Option<&str>,
127) -> crate::error::Result<String> {
128    if !env_var_name.is_empty() {
129        if let Ok(v) = std::env::var(env_var_name) {
130            if !v.is_empty() {
131                return Ok(v);
132            }
133        }
134    }
135    if let Some(v) = inline {
136        if !v.is_empty() {
137            return Ok(v.to_string());
138        }
139    }
140    Err(crate::error::AppError::Credentials(format!(
141        "{vendor_label}: no API key. Either export {env_var_name} or set \
142         `api_key` under [{}] in ~/.config/ai-usagebar/config.toml (chmod 600).",
143        vendor_label.to_lowercase()
144    )))
145}
146
147impl Config {
148    /// Load from `~/.config/ai-usagebar/config.toml`. Returns defaults if the
149    /// file doesn't exist; errors only on actual parse failures.
150    pub fn load() -> Result<Self> {
151        let Some(path) = default_path() else {
152            return Ok(Self::default());
153        };
154        Self::load_from(&path)
155    }
156
157    pub fn load_from(path: &std::path::Path) -> Result<Self> {
158        match std::fs::read_to_string(path) {
159            Ok(s) => Ok(toml::from_str(&s)?),
160            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(Self::default()),
161            Err(e) => Err(AppError::io_at(path, e)),
162        }
163    }
164
165    pub fn is_enabled(&self, id: VendorId) -> bool {
166        match id {
167            VendorId::Anthropic => self.anthropic.enabled,
168            VendorId::Openai => self.openai.enabled,
169            VendorId::Zai => self.zai.enabled,
170            VendorId::Openrouter => self.openrouter.enabled,
171        }
172    }
173
174    pub fn enabled_vendors(&self) -> Vec<VendorId> {
175        VendorId::all()
176            .iter()
177            .copied()
178            .filter(|id| self.is_enabled(*id))
179            .collect()
180    }
181}
182
183fn default_path() -> Option<PathBuf> {
184    let proj = directories::ProjectDirs::from("", "", "ai-usagebar")?;
185    Some(proj.config_dir().join("config.toml"))
186}
187
188#[cfg(test)]
189mod tests {
190    use super::*;
191    use std::io::Write;
192    use tempfile::NamedTempFile;
193
194    fn write_toml(s: &str) -> NamedTempFile {
195        let mut f = NamedTempFile::new().unwrap();
196        f.write_all(s.as_bytes()).unwrap();
197        f.flush().unwrap();
198        f
199    }
200
201    #[test]
202    fn defaults_enable_all_vendors() {
203        let c = Config::default();
204        assert!(c.is_enabled(VendorId::Anthropic));
205        assert!(c.is_enabled(VendorId::Openai));
206        assert!(c.is_enabled(VendorId::Zai));
207        assert!(c.is_enabled(VendorId::Openrouter));
208        assert_eq!(c.enabled_vendors().len(), 4);
209    }
210
211    #[test]
212    fn missing_file_uses_defaults() {
213        let path = std::path::Path::new("/tmp/does-not-exist-ai-usagebar-test");
214        let c = Config::load_from(path).unwrap();
215        assert!(c.is_enabled(VendorId::Anthropic));
216    }
217
218    #[test]
219    fn parses_full_config() {
220        let f = write_toml(
221            r#"
222            [anthropic]
223            enabled = true
224
225            [openai]
226            enabled = false
227            admin_key_env = "MY_ADMIN_KEY"
228
229            [zai]
230            enabled = true
231            api_key_env = "MY_ZAI"
232            plan_tier = "pro"
233
234            [openrouter]
235            enabled = false
236            "#,
237        );
238        let c = Config::load_from(f.path()).unwrap();
239        assert!(c.is_enabled(VendorId::Anthropic));
240        assert!(!c.is_enabled(VendorId::Openai));
241        assert!(c.is_enabled(VendorId::Zai));
242        assert!(!c.is_enabled(VendorId::Openrouter));
243        assert_eq!(c.openai.admin_key_env, "MY_ADMIN_KEY");
244        assert_eq!(c.zai.api_key_env, "MY_ZAI");
245        assert_eq!(c.zai.plan_tier.as_deref(), Some("pro"));
246    }
247
248    #[test]
249    fn partial_config_falls_back_to_defaults() {
250        let f = write_toml(
251            r#"[openai]
252enabled = false
253"#,
254        );
255        let c = Config::load_from(f.path()).unwrap();
256        assert!(!c.is_enabled(VendorId::Openai));
257        // Other vendors keep their defaults.
258        assert!(c.is_enabled(VendorId::Anthropic));
259        assert_eq!(c.openai.admin_key_env, "OPENAI_ADMIN_KEY");
260    }
261
262    #[test]
263    fn malformed_toml_returns_error() {
264        let f = write_toml("this is not = = valid");
265        assert!(Config::load_from(f.path()).is_err());
266    }
267
268    // serial guard for env-var manipulation tests so they don't race
269    fn env_guard() -> std::sync::MutexGuard<'static, ()> {
270        static M: std::sync::Mutex<()> = std::sync::Mutex::new(());
271        M.lock().unwrap_or_else(|p| p.into_inner())
272    }
273
274    #[test]
275    fn resolve_api_key_prefers_env_over_inline() {
276        let _g = env_guard();
277        // Use a unique env var name so we don't clobber test parallelism.
278        let var = "AI_USAGEBAR_TEST_ENV_WINS";
279        // SAFETY: tests are single-threaded under env_guard.
280        unsafe { std::env::set_var(var, "from-env") };
281        let got = resolve_api_key("Zai", var, Some("from-inline")).unwrap();
282        unsafe { std::env::remove_var(var) };
283        assert_eq!(got, "from-env");
284    }
285
286    #[test]
287    fn resolve_api_key_falls_back_to_inline() {
288        let _g = env_guard();
289        let var = "AI_USAGEBAR_TEST_INLINE_FALLBACK";
290        unsafe { std::env::remove_var(var) };
291        let got = resolve_api_key("Zai", var, Some("inline-key")).unwrap();
292        assert_eq!(got, "inline-key");
293    }
294
295    #[test]
296    fn resolve_api_key_errors_when_both_missing() {
297        let _g = env_guard();
298        let var = "AI_USAGEBAR_TEST_BOTH_MISSING";
299        unsafe { std::env::remove_var(var) };
300        let err = resolve_api_key("Zai", var, None).unwrap_err();
301        match err {
302            crate::error::AppError::Credentials(msg) => {
303                assert!(msg.contains(var), "error should name env var: {msg}");
304                assert!(
305                    msg.contains("api_key"),
306                    "error should suggest config field: {msg}"
307                );
308            }
309            other => panic!("expected Credentials error, got {other:?}"),
310        }
311    }
312
313    #[test]
314    fn resolve_api_key_treats_empty_env_as_unset() {
315        let _g = env_guard();
316        let var = "AI_USAGEBAR_TEST_EMPTY_ENV";
317        unsafe { std::env::set_var(var, "") };
318        let got = resolve_api_key("OpenRouter", var, Some("inline")).unwrap();
319        unsafe { std::env::remove_var(var) };
320        assert_eq!(got, "inline");
321    }
322
323    #[test]
324    fn config_parses_with_inline_api_key_and_primary() {
325        let f = write_toml(
326            r#"
327            [ui]
328            primary = "openrouter"
329
330            [zai]
331            enabled = true
332            api_key_env = "MY_ZAI"
333            api_key = "sk-zai-inline"
334
335            [openrouter]
336            enabled = true
337            api_key = "sk-or-inline"
338            "#,
339        );
340        let c = Config::load_from(f.path()).unwrap();
341        assert_eq!(c.ui.primary, Some(VendorId::Openrouter));
342        assert_eq!(c.zai.api_key.as_deref(), Some("sk-zai-inline"));
343        assert_eq!(c.openrouter.api_key.as_deref(), Some("sk-or-inline"));
344    }
345
346    #[test]
347    fn enabled_vendors_preserves_canonical_order() {
348        let c = Config::default();
349        assert_eq!(
350            c.enabled_vendors(),
351            vec![
352                VendorId::Anthropic,
353                VendorId::Openai,
354                VendorId::Zai,
355                VendorId::Openrouter
356            ]
357        );
358    }
359}
ai_usagebar/config.rs

ai_usagebar/
config.rs
