vct_core/quota/

provider.rs

//! Generic background quota worker, shared by every quota provider.
//!
//! Each provider gets its own thread (so one provider's slow HTTP never stalls
//! the others), but they all run the same panic-isolated poll loop here (cadence
//! from `config.usage.quota.refresh_interval`). A provider supplies a stateful
//! `resolve` closure returning a [`QuotaOutcome`]; the worker applies it to the
//! shared snapshot:
//!
//! - [`QuotaOutcome::Data`] — store the fresh snapshot + persist it.
//! - [`QuotaOutcome::NeedsLogin`] — flip `needs_login` on the *current*
//!   snapshot (keep its data), so a refresh failure never blanks out
//!   still-valid last-known-good numbers (S3).
//! - [`QuotaOutcome::Transient`] — keep last-known-good indefinitely, so a rate
//!   limit or network blip never blanks a panel; the next success overwrites it.

use crate::models::{
    ClaudeQuotaSnapshot, CodexQuotaSnapshot, CopilotQuotaSnapshot, CursorQuotaSnapshot,
    GrokQuotaSnapshot, QuotaSource,
};
use serde::Serialize;
use std::panic::{AssertUnwindSafe, catch_unwind};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::thread::JoinHandle;
use std::time::Duration;

/// A normalized quota snapshot the worker can store and flag.
pub trait QuotaSnapshot: Clone + Default + Send + Serialize + 'static {
    /// Unix seconds when this snapshot was produced.
    fn fetched_at(&self) -> i64;
    /// Whether this snapshot carries anything to show (data or a login hint).
    fn is_present(&self) -> bool;
    /// Sets the `needs_login` flag without touching the data.
    fn set_needs_login(&mut self, value: bool);
}

impl QuotaSnapshot for ClaudeQuotaSnapshot {
    fn fetched_at(&self) -> i64 {
        self.fetched_at
    }
    fn is_present(&self) -> bool {
        self.five_hour.is_some()
            || self.seven_day.is_some()
            // Mirror the render gate: a scoped row needs both the window and its
            // model label, so a label-less scoped window is not "present" alone.
            || (self.scoped_weekly.is_some() && self.scoped_label.is_some())
            || self.needs_login
    }
    fn set_needs_login(&mut self, value: bool) {
        self.needs_login = value;
    }
}

impl QuotaSnapshot for CodexQuotaSnapshot {
    fn fetched_at(&self) -> i64 {
        self.fetched_at
    }
    fn is_present(&self) -> bool {
        self.source != QuotaSource::None || self.needs_login
    }
    fn set_needs_login(&mut self, value: bool) {
        self.needs_login = value;
    }
}

impl QuotaSnapshot for CopilotQuotaSnapshot {
    fn fetched_at(&self) -> i64 {
        self.fetched_at
    }
    fn is_present(&self) -> bool {
        // Mirror the render gate: a premium gauge, an unlimited flag, the plan
        // label, or a login hint all give the panel something to show. (Chat /
        // completions unlimited flags are not rendered, so they must not gate
        // visibility either, or the panel would show "no Copilot quota".)
        self.premium.is_some()
            || self.premium_unlimited
            || self.plan_type.is_some()
            || self.needs_login
    }
    fn set_needs_login(&mut self, value: bool) {
        self.needs_login = value;
    }
}

impl QuotaSnapshot for CursorQuotaSnapshot {
    fn fetched_at(&self) -> i64 {
        self.fetched_at
    }
    fn is_present(&self) -> bool {
        self.total.is_some()
            || self.auto.is_some()
            || self.api.is_some()
            || self.plan_type.is_some()
            || self.needs_login
    }
    fn set_needs_login(&mut self, value: bool) {
        self.needs_login = value;
    }
}

impl QuotaSnapshot for GrokQuotaSnapshot {
    fn fetched_at(&self) -> i64 {
        self.fetched_at
    }
    fn is_present(&self) -> bool {
        // Mirror the render gate: the included gauge, the plan label, or a login
        // hint each give the panel something to show. The money fields never
        // stand alone — a billing body that maps at all yields `included`.
        self.included.is_some() || self.plan_type.is_some() || self.needs_login
    }
    fn set_needs_login(&mut self, value: bool) {
        self.needs_login = value;
    }
}

/// The result of one provider `resolve` tick.
pub enum QuotaOutcome<T> {
    /// A fresh snapshot to store (may itself carry `needs_login`, e.g. Codex
    /// session-fallback data + login hint).
    Data(T),
    /// Auth failed and there is no fallback data: flag the current snapshot for
    /// re-login but keep whatever it is already showing.
    NeedsLogin,
    /// Transient failure (network / rate limit): keep last-known-good.
    Transient,
}

/// Spawns a detached background worker that refreshes `shared` (and the on-disk
/// cache via `save`) every `refresh_secs` until `shutdown` is set.
///
/// The worker is panic-isolated and holds the mutex only for the assignment, so
/// it can never poison the lock. It is not joined on quit — `shutdown` is set as
/// a courtesy and the OS reclaims the thread on process exit.
pub fn spawn_quota_worker<T, R, S>(
    label: &'static str,
    shared: Arc<Mutex<T>>,
    shutdown: Arc<AtomicBool>,
    refresh_secs: u64,
    mut resolve: R,
    save: S,
) -> JoinHandle<()>
where
    T: QuotaSnapshot,
    R: FnMut() -> QuotaOutcome<T> + Send + 'static,
    S: Fn(&T) + Send + 'static,
{
    std::thread::spawn(move || {
        loop {
            if shutdown.load(Ordering::Relaxed) {
                break;
            }
            match catch_unwind(AssertUnwindSafe(&mut resolve)) {
                Ok(QuotaOutcome::Data(snap)) => {
                    if let Ok(mut guard) = shared.lock() {
                        *guard = snap.clone();
                    }
                    save(&snap);
                }
                Ok(QuotaOutcome::NeedsLogin) => {
                    let mut updated = None;
                    if let Ok(mut guard) = shared.lock() {
                        guard.set_needs_login(true);
                        updated = Some(guard.clone());
                    }
                    if let Some(snap) = updated {
                        save(&snap);
                    }
                }
                // A transient failure (network error, 429 rate limit) leaves the
                // last-known-good snapshot untouched — the panel keeps showing it
                // with a growing staleness marker until the next success
                // overwrites it, rather than blanking out.
                Ok(QuotaOutcome::Transient) => {}
                Err(_) => log::warn!("{label} quota worker panicked; keeping last snapshot"),
            }
            // Sleep in 200ms slices so shutdown stays responsive.
            for _ in 0..(refresh_secs * 5) {
                if shutdown.load(Ordering::Relaxed) {
                    break;
                }
                std::thread::sleep(Duration::from_millis(200));
            }
        }
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::models::CodexQuotaSnapshot;

    #[test]
    fn needs_login_snapshot_is_present() {
        let snap = CodexQuotaSnapshot {
            needs_login: true,
            ..Default::default()
        };
        assert!(snap.is_present());
    }
}