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};
10use std::time::Duration;
11
12use crate::cache::{acquire_lock, atomic_write};
13use crate::error::{AppError, Result};
14use crate::vendor::VendorId;
15
16/// Much shorter than the vendors' fetch locks (15–45s): the critical section
17/// here is one small read plus a rename, so a wait this long already means a
18/// wedged holder — and a scroll event that blocks for seconds is worse than a
19/// dropped one.
20const LOCK_TIMEOUT: Duration = Duration::from_secs(5);
21
22fn state_dir() -> Result<PathBuf> {
23    let base = directories::BaseDirs::new()
24        .ok_or_else(|| AppError::Other("could not resolve XDG cache dir".into()))?;
25    Ok(base.cache_dir().join("ai-usagebar"))
26}
27
28fn state_path() -> Result<PathBuf> {
29    Ok(state_dir()?.join("active_vendor"))
30}
31
32/// Read the persisted active vendor, if any. `None` means "no override —
33/// callers fall back to [ui] primary or anthropic".
34pub fn read() -> Option<VendorId> {
35    read_from(&state_path().ok()?)
36}
37
38/// Read the persisted active vendor from an explicit path. The real-path
39/// [`read`] is a thin wrapper over this; tests use this directly with a
40/// `TempDir`-backed path so they never touch `~/.cache/ai-usagebar`.
41pub fn read_from(path: &Path) -> Option<VendorId> {
42    let raw = fs::read_to_string(path).ok()?;
43    parse_slug(raw.trim())
44}
45
46/// Persist `vendor` as the active one. Atomic.
47pub fn write(vendor: VendorId) -> Result<()> {
48    write_to(&state_path()?, vendor)
49}
50
51/// Persist `vendor` to an explicit path. Atomic. Test-friendly counterpart
52/// to [`write`] (mirrors [`crate::cache::Cache::at`] vs `for_vendor`).
53pub fn write_to(path: &Path, vendor: VendorId) -> Result<()> {
54    atomic_write(path, vendor.slug().as_bytes())
55}
56
57/// Cycle the active vendor by `delta` positions through `enabled` (which
58/// preserves canonical order). Wraps. If no state exists, starts at `start`
59/// (usually `[ui] primary` or anthropic).
60pub fn cycle(enabled: &[VendorId], start: VendorId, delta: i32) -> Result<VendorId> {
61    cycle_at(&state_path()?, enabled, start, delta)
62}
63
64/// The flock guarding a state file's read-modify-write, as a sibling of the
65/// state file itself (mirroring `Cache::lock_path`'s `.fetch.lock`).
66fn lock_path_for(state: &Path) -> PathBuf {
67    let mut p = state.as_os_str().to_os_string();
68    p.push(".lock");
69    PathBuf::from(p)
70}
71
72/// Cycle using an explicit state-file path. The real-path [`cycle`] is a thin
73/// wrapper over this; tests drive this with a `TempDir` path so the cycle +
74/// persistence logic is covered without reading or writing the real cache.
75pub fn cycle_at(
76    path: &Path,
77    enabled: &[VendorId],
78    start: VendorId,
79    delta: i32,
80) -> Result<VendorId> {
81    if enabled.is_empty() {
82        return Err(AppError::Other("no enabled vendors to cycle".into()));
83    }
84    // One flick of a scroll wheel fires several `--cycle-next` processes at
85    // once. An atomic *write* only guarantees no torn file — it does not stop
86    // two of them reading the same current vendor and both persisting the same
87    // next one, silently eating a step. The lock has to span read→compute→write.
88    let _lock = acquire_lock(&lock_path_for(path), LOCK_TIMEOUT)?;
89    let current = read_from(path)
90        .filter(|v| enabled.contains(v))
91        .unwrap_or(start);
92    let cur_idx = enabled.iter().position(|v| *v == current).unwrap_or(0);
93    let n = enabled.len() as i32;
94    let next_idx = ((cur_idx as i32 + delta).rem_euclid(n)) as usize;
95    let next = enabled[next_idx];
96    write_to(path, next)?;
97    Ok(next)
98}
99
100fn parse_slug(s: &str) -> Option<VendorId> {
101    match s {
102        "anthropic" => Some(VendorId::Anthropic),
103        "anthropic_api" => Some(VendorId::AnthropicApi),
104        "openai" => Some(VendorId::Openai),
105        "zai" => Some(VendorId::Zai),
106        "openrouter" => Some(VendorId::Openrouter),
107        "deepseek" => Some(VendorId::Deepseek),
108        "kimi" => Some(VendorId::Kimi),
109        "kilo" => Some(VendorId::Kilo),
110        "novita" => Some(VendorId::Novita),
111        "moonshot" => Some(VendorId::Moonshot),
112        "grok" => Some(VendorId::Grok),
113        "supergrok" => Some(VendorId::Supergrok),
114        "antigravity" => Some(VendorId::Antigravity),
115        "cursor" => Some(VendorId::Cursor),
116        "minimax" => Some(VendorId::Minimax),
117        "kiro" => Some(VendorId::Kiro),
118        "nous" => Some(VendorId::NousResearch),
119        "opencode-go" => Some(VendorId::OpenCodeGo),
120        _ => None,
121    }
122}
123
124#[cfg(test)]
125mod tests {
126    use super::*;
127    use tempfile::TempDir;
128
129    // Deliberately excludes Deepseek so the not-in-enabled-set fallback test
130    // below has a real vendor to persist that is outside this cycle set.
131    const CYCLE_SET: [VendorId; 4] = [
132        VendorId::Anthropic,
133        VendorId::Openai,
134        VendorId::Zai,
135        VendorId::Openrouter,
136    ];
137
138    /// One flick of a scroll wheel fires several `--cycle-next` processes at
139    /// once. Before the lock spanned the read-modify-write, two of them could
140    /// read the same current vendor and both persist the same next one, so N
141    /// events advanced fewer than N steps. With the lock held across the whole
142    /// operation, every cycle observes its predecessor's write.
143    #[test]
144    fn concurrent_cycles_do_not_lose_a_step() {
145        let td = TempDir::new().unwrap();
146        let path = td.path().join("active_vendor");
147        let start = VendorId::Anthropic;
148
149        // Four threads, each doing one +1 step, over a 4-vendor set: if no step
150        // is lost the value returns exactly to `start`.
151        const THREADS: usize = 4;
152        std::thread::scope(|s| {
153            for _ in 0..THREADS {
154                s.spawn(|| {
155                    let _ = cycle_at(&path, &CYCLE_SET, start, 1);
156                });
157            }
158        });
159
160        let landed = read_from(&path).expect("a vendor must have been persisted");
161        assert_eq!(
162            landed,
163            start,
164            "{THREADS} single steps over {} vendors must return to the start; \
165             landing on {landed:?} means a step was lost to a race",
166            CYCLE_SET.len()
167        );
168    }
169
170    #[test]
171    fn parse_slug_round_trip() {
172        for id in VendorId::all() {
173            assert_eq!(parse_slug(id.slug()), Some(*id));
174        }
175    }
176
177    #[test]
178    fn parse_slug_unknown_returns_none() {
179        assert!(parse_slug("not-a-vendor").is_none());
180        assert!(parse_slug("").is_none());
181    }
182
183    #[test]
184    fn read_from_missing_or_garbage_returns_none() {
185        let td = TempDir::new().unwrap();
186        // Missing file → None.
187        assert!(read_from(&td.path().join("active_vendor")).is_none());
188        // Round-trip a real slug.
189        let path = td.path().join("active_vendor");
190        write_to(&path, VendorId::Zai).unwrap();
191        assert_eq!(read_from(&path), Some(VendorId::Zai));
192        // Garbage content → None (not a known slug).
193        fs::write(&path, "not-a-vendor").unwrap();
194        assert!(read_from(&path).is_none());
195    }
196
197    #[test]
198    fn cycle_at_persists_state_across_calls() {
199        let td = TempDir::new().unwrap();
200        let path = td.path().join("active_vendor");
201
202        // No state yet → starts at `start`, steps forward to Openai.
203        let v = cycle_at(&path, &CYCLE_SET, VendorId::Anthropic, 1).unwrap();
204        assert_eq!(v, VendorId::Openai);
205        assert_eq!(read_from(&path), Some(VendorId::Openai));
206
207        // Next forward step reads the persisted Openai → Zai.
208        let v = cycle_at(&path, &CYCLE_SET, VendorId::Anthropic, 1).unwrap();
209        assert_eq!(v, VendorId::Zai);
210        assert_eq!(read_from(&path), Some(VendorId::Zai));
211    }
212
213    #[test]
214    fn cycle_at_wraps_forward_and_backward() {
215        let td = TempDir::new().unwrap();
216        let path = td.path().join("active_vendor");
217        write_to(&path, VendorId::Anthropic).unwrap();
218
219        // backward from Anthropic wraps to Openrouter
220        assert_eq!(
221            cycle_at(&path, &CYCLE_SET, VendorId::Anthropic, -1).unwrap(),
222            VendorId::Openrouter
223        );
224        // forward from Openrouter wraps back to Anthropic
225        assert_eq!(
226            cycle_at(&path, &CYCLE_SET, VendorId::Anthropic, 1).unwrap(),
227            VendorId::Anthropic
228        );
229    }
230
231    #[test]
232    fn cycle_at_ignores_persisted_vendor_not_in_enabled_set() {
233        let td = TempDir::new().unwrap();
234        let path = td.path().join("active_vendor");
235        // Persist a vendor that isn't in the enabled set → fall back to `start`.
236        write_to(&path, VendorId::Deepseek).unwrap();
237        let enabled = [VendorId::Anthropic, VendorId::Openai];
238        // start=Openai (idx 1), +1 wraps to idx 0 = Anthropic.
239        let v = cycle_at(&path, &enabled, VendorId::Openai, 1).unwrap();
240        assert_eq!(v, VendorId::Anthropic);
241    }
242
243    #[test]
244    fn cycle_at_errors_on_empty_enabled() {
245        let td = TempDir::new().unwrap();
246        let path = td.path().join("active_vendor");
247        let res = cycle_at(&path, &[], VendorId::Anthropic, 1);
248        assert!(matches!(res, Err(AppError::Other(_))));
249    }
250}
ai_usagebar/active.rs

ai_usagebar/
active.rs
