Skip to main content

ai_usagebar/
active.rs

1//! Active-vendor state file. Set by `--cycle-next` / `--cycle-prev` (which
2//! Waybar's `on-scroll-up`/`on-scroll-down` invoke), read by the widget on
3//! every tick. The TUI does NOT consult this — it has its own tab state.
4//!
5//! On-disk shape: a single line with the vendor slug (e.g. `openai`). Located
6//! at `<cache-dir>/active_vendor`.
7
8use std::fs;
9use std::path::{Path, PathBuf};
10
11use crate::cache::atomic_write;
12use crate::error::{AppError, Result};
13use crate::vendor::VendorId;
14
15fn state_dir() -> Result<PathBuf> {
16    let base = directories::BaseDirs::new()
17        .ok_or_else(|| AppError::Other("could not resolve XDG cache dir".into()))?;
18    Ok(base.cache_dir().join("ai-usagebar"))
19}
20
21fn state_path() -> Result<PathBuf> {
22    Ok(state_dir()?.join("active_vendor"))
23}
24
25/// Read the persisted active vendor, if any. `None` means "no override —
26/// callers fall back to [ui] primary or anthropic".
27pub fn read() -> Option<VendorId> {
28    read_from(&state_path().ok()?)
29}
30
31/// Read the persisted active vendor from an explicit path. The real-path
32/// [`read`] is a thin wrapper over this; tests use this directly with a
33/// `TempDir`-backed path so they never touch `~/.cache/ai-usagebar`.
34pub fn read_from(path: &Path) -> Option<VendorId> {
35    let raw = fs::read_to_string(path).ok()?;
36    parse_slug(raw.trim())
37}
38
39/// Persist `vendor` as the active one. Atomic.
40pub fn write(vendor: VendorId) -> Result<()> {
41    write_to(&state_path()?, vendor)
42}
43
44/// Persist `vendor` to an explicit path. Atomic. Test-friendly counterpart
45/// to [`write`] (mirrors [`crate::cache::Cache::at`] vs `for_vendor`).
46pub fn write_to(path: &Path, vendor: VendorId) -> Result<()> {
47    atomic_write(path, vendor.slug().as_bytes())
48}
49
50/// Cycle the active vendor by `delta` positions through `enabled` (which
51/// preserves canonical order). Wraps. If no state exists, starts at `start`
52/// (usually `[ui] primary` or anthropic).
53pub fn cycle(enabled: &[VendorId], start: VendorId, delta: i32) -> Result<VendorId> {
54    cycle_at(&state_path()?, enabled, start, delta)
55}
56
57/// Cycle using an explicit state-file path. The real-path [`cycle`] is a thin
58/// wrapper over this; tests drive this with a `TempDir` path so the cycle +
59/// persistence logic is covered without reading or writing the real cache.
60pub fn cycle_at(
61    path: &Path,
62    enabled: &[VendorId],
63    start: VendorId,
64    delta: i32,
65) -> Result<VendorId> {
66    if enabled.is_empty() {
67        return Err(AppError::Other("no enabled vendors to cycle".into()));
68    }
69    let current = read_from(path)
70        .filter(|v| enabled.contains(v))
71        .unwrap_or(start);
72    let cur_idx = enabled.iter().position(|v| *v == current).unwrap_or(0);
73    let n = enabled.len() as i32;
74    let next_idx = ((cur_idx as i32 + delta).rem_euclid(n)) as usize;
75    let next = enabled[next_idx];
76    write_to(path, next)?;
77    Ok(next)
78}
79
80fn parse_slug(s: &str) -> Option<VendorId> {
81    match s {
82        "anthropic" => Some(VendorId::Anthropic),
83        "openai" => Some(VendorId::Openai),
84        "zai" => Some(VendorId::Zai),
85        "openrouter" => Some(VendorId::Openrouter),
86        "deepseek" => Some(VendorId::Deepseek),
87        "kimi" => Some(VendorId::Kimi),
88        _ => None,
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95    use tempfile::TempDir;
96
97    // Deliberately excludes Deepseek so the not-in-enabled-set fallback test
98    // below has a real vendor to persist that is outside this cycle set.
99    const CYCLE_SET: [VendorId; 4] = [
100        VendorId::Anthropic,
101        VendorId::Openai,
102        VendorId::Zai,
103        VendorId::Openrouter,
104    ];
105
106    #[test]
107    fn parse_slug_round_trip() {
108        for id in VendorId::all() {
109            assert_eq!(parse_slug(id.slug()), Some(*id));
110        }
111    }
112
113    #[test]
114    fn parse_slug_unknown_returns_none() {
115        assert!(parse_slug("not-a-vendor").is_none());
116        assert!(parse_slug("").is_none());
117    }
118
119    #[test]
120    fn read_from_missing_or_garbage_returns_none() {
121        let td = TempDir::new().unwrap();
122        // Missing file → None.
123        assert!(read_from(&td.path().join("active_vendor")).is_none());
124        // Round-trip a real slug.
125        let path = td.path().join("active_vendor");
126        write_to(&path, VendorId::Zai).unwrap();
127        assert_eq!(read_from(&path), Some(VendorId::Zai));
128        // Garbage content → None (not a known slug).
129        fs::write(&path, "not-a-vendor").unwrap();
130        assert!(read_from(&path).is_none());
131    }
132
133    #[test]
134    fn cycle_at_persists_state_across_calls() {
135        let td = TempDir::new().unwrap();
136        let path = td.path().join("active_vendor");
137
138        // No state yet → starts at `start`, steps forward to Openai.
139        let v = cycle_at(&path, &CYCLE_SET, VendorId::Anthropic, 1).unwrap();
140        assert_eq!(v, VendorId::Openai);
141        assert_eq!(read_from(&path), Some(VendorId::Openai));
142
143        // Next forward step reads the persisted Openai → Zai.
144        let v = cycle_at(&path, &CYCLE_SET, VendorId::Anthropic, 1).unwrap();
145        assert_eq!(v, VendorId::Zai);
146        assert_eq!(read_from(&path), Some(VendorId::Zai));
147    }
148
149    #[test]
150    fn cycle_at_wraps_forward_and_backward() {
151        let td = TempDir::new().unwrap();
152        let path = td.path().join("active_vendor");
153        write_to(&path, VendorId::Anthropic).unwrap();
154
155        // backward from Anthropic wraps to Openrouter
156        assert_eq!(
157            cycle_at(&path, &CYCLE_SET, VendorId::Anthropic, -1).unwrap(),
158            VendorId::Openrouter
159        );
160        // forward from Openrouter wraps back to Anthropic
161        assert_eq!(
162            cycle_at(&path, &CYCLE_SET, VendorId::Anthropic, 1).unwrap(),
163            VendorId::Anthropic
164        );
165    }
166
167    #[test]
168    fn cycle_at_ignores_persisted_vendor_not_in_enabled_set() {
169        let td = TempDir::new().unwrap();
170        let path = td.path().join("active_vendor");
171        // Persist a vendor that isn't in the enabled set → fall back to `start`.
172        write_to(&path, VendorId::Deepseek).unwrap();
173        let enabled = [VendorId::Anthropic, VendorId::Openai];
174        // start=Openai (idx 1), +1 wraps to idx 0 = Anthropic.
175        let v = cycle_at(&path, &enabled, VendorId::Openai, 1).unwrap();
176        assert_eq!(v, VendorId::Anthropic);
177    }
178
179    #[test]
180    fn cycle_at_errors_on_empty_enabled() {
181        let td = TempDir::new().unwrap();
182        let path = td.path().join("active_vendor");
183        let res = cycle_at(&path, &[], VendorId::Anthropic, 1);
184        assert!(matches!(res, Err(AppError::Other(_))));
185    }
186}