Skip to content

Persist usage snapshots across reloads #30

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
SHAREN wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Persist usage snapshots across reloads #30

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
43bb408
Persist usage snapshots across reloads
SHAREN Mar 26, 2026
3bd763d
Persist stale usage snapshots across reloads
SHAREN Mar 27, 2026
80c0904
Align usage refresh labels with persisted timestamps
SHAREN Mar 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
247 changes: 139 additions & 108 deletions src-tauri/src/auth/storage.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,14 @@

use std::fs;
use std::path::PathBuf;
use std::sync::Mutex;

use anyhow::{Context, Result};
use chrono::{DateTime, Utc};

use crate::types::{AccountsStore, AuthData, StoredAccount};
use crate::types::{AccountsStore, AuthData, StoredAccount, UsageInfo};

static STORE_LOCK: Mutex<()> = Mutex::new(());

/// Get the path to the codex-switcher config directory
pub fn get_config_dir() -> Result<PathBuf> {
Expand All @@ -20,6 +24,11 @@ pub fn get_accounts_file() -> Result<PathBuf> {

/// Load the accounts store from disk
pub fn load_accounts() -> Result<AccountsStore> {
let _guard = STORE_LOCK.lock().unwrap();
load_accounts_unlocked()
}

fn load_accounts_unlocked() -> Result<AccountsStore> {
let path = get_accounts_file()?;

if !path.exists() {
Expand All @@ -37,6 +46,11 @@ pub fn load_accounts() -> Result<AccountsStore> {

/// Save the accounts store to disk
pub fn save_accounts(store: &AccountsStore) -> Result<()> {
let _guard = STORE_LOCK.lock().unwrap();
save_accounts_unlocked(store)
}

fn save_accounts_unlocked(store: &AccountsStore) -> Result<()> {
let path = get_accounts_file()?;

// Ensure the config directory exists
Expand All @@ -62,59 +76,62 @@ pub fn save_accounts(store: &AccountsStore) -> Result<()> {
Ok(())
}

fn with_accounts_store_mut<R>(
mutator: impl FnOnce(&mut AccountsStore) -> Result<R>,
) -> Result<R> {
let _guard = STORE_LOCK.lock().unwrap();
let mut store = load_accounts_unlocked()?;
let result = mutator(&mut store)?;
save_accounts_unlocked(&store)?;
Ok(result)
}

/// Add a new account to the store
pub fn add_account(account: StoredAccount) -> Result<StoredAccount> {
let mut store = load_accounts()?;

// Check for duplicate names
if store.accounts.iter().any(|a| a.name == account.name) {
anyhow::bail!("An account with name '{}' already exists", account.name);
}

let account_clone = account.clone();
store.accounts.push(account);
with_accounts_store_mut(move |store| {
if store.accounts.iter().any(|a| a.name == account.name) {
anyhow::bail!("An account with name '{}' already exists", account.name);
}

// If this is the first account, make it active
if store.accounts.len() == 1 {
store.active_account_id = Some(account_clone.id.clone());
}
store.accounts.push(account);

if store.accounts.len() == 1 {
store.active_account_id = Some(account_clone.id.clone());
}

save_accounts(&store)?;
Ok(account_clone)
Ok(account_clone)
})
}

/// Remove an account by ID
pub fn remove_account(account_id: &str) -> Result<()> {
let mut store = load_accounts()?;
with_accounts_store_mut(move |store| {
let initial_len = store.accounts.len();
store.accounts.retain(|a| a.id != account_id);

let initial_len = store.accounts.len();
store.accounts.retain(|a| a.id != account_id);

if store.accounts.len() == initial_len {
anyhow::bail!("Account not found: {account_id}");
}
if store.accounts.len() == initial_len {
anyhow::bail!("Account not found: {account_id}");
}

// If we removed the active account, clear it or set to first available
if store.active_account_id.as_deref() == Some(account_id) {
store.active_account_id = store.accounts.first().map(|a| a.id.clone());
}
if store.active_account_id.as_deref() == Some(account_id) {
store.active_account_id = store.accounts.first().map(|a| a.id.clone());
}

save_accounts(&store)?;
Ok(())
Ok(())
})
}

/// Update the active account ID
pub fn set_active_account(account_id: &str) -> Result<()> {
let mut store = load_accounts()?;

// Verify the account exists
if !store.accounts.iter().any(|a| a.id == account_id) {
anyhow::bail!("Account not found: {account_id}");
}
with_accounts_store_mut(move |store| {
if !store.accounts.iter().any(|a| a.id == account_id) {
anyhow::bail!("Account not found: {account_id}");
}

store.active_account_id = Some(account_id.to_string());
save_accounts(&store)?;
Ok(())
store.active_account_id = Some(account_id.to_string());
Ok(())
})
}

/// Get an account by ID
Expand All @@ -135,14 +152,13 @@ pub fn get_active_account() -> Result<Option<StoredAccount>> {

/// Update an account's last_used_at timestamp
pub fn touch_account(account_id: &str) -> Result<()> {
let mut store = load_accounts()?;

if let Some(account) = store.accounts.iter_mut().find(|a| a.id == account_id) {
account.last_used_at = Some(chrono::Utc::now());
save_accounts(&store)?;
}
with_accounts_store_mut(move |store| {
if let Some(account) = store.accounts.iter_mut().find(|a| a.id == account_id) {
account.last_used_at = Some(chrono::Utc::now());
}

Ok(())
Ok(())
})
}

/// Update an account's metadata (name, email, plan_type)
Expand All @@ -152,40 +168,37 @@ pub fn update_account_metadata(
email: Option<String>,
plan_type: Option<String>,
) -> Result<()> {
let mut store = load_accounts()?;

// Check for duplicate names first (if renaming)
if let Some(ref new_name) = name {
if store
.accounts
.iter()
.any(|a| a.id != account_id && a.name == *new_name)
{
anyhow::bail!("An account with name '{new_name}' already exists");
with_accounts_store_mut(move |store| {
if let Some(ref new_name) = name {
if store
.accounts
.iter()
.any(|a| a.id != account_id && a.name == *new_name)
{
anyhow::bail!("An account with name '{new_name}' already exists");
}
}
}

// Now find and update the account
let account = store
.accounts
.iter_mut()
.find(|a| a.id == account_id)
.context("Account not found")?;
let account = store
.accounts
.iter_mut()
.find(|a| a.id == account_id)
.context("Account not found")?;

if let Some(new_name) = name {
account.name = new_name;
}
if let Some(new_name) = name {
account.name = new_name;
}

if email.is_some() {
account.email = email;
}
if email.is_some() {
account.email = email;
}

if plan_type.is_some() {
account.plan_type = plan_type;
}
if plan_type.is_some() {
account.plan_type = plan_type;
}

save_accounts(&store)?;
Ok(())
Ok(())
})
}

/// Update ChatGPT OAuth tokens for an account and return the updated account.
Expand All @@ -198,44 +211,62 @@ pub fn update_account_chatgpt_tokens(
email: Option<String>,
plan_type: Option<String>,
) -> Result<StoredAccount> {
let mut store = load_accounts()?;

let account = store
.accounts
.iter_mut()
.find(|a| a.id == account_id)
.context("Account not found")?;

match &mut account.auth_data {
AuthData::ChatGPT {
id_token: stored_id_token,
access_token: stored_access_token,
refresh_token: stored_refresh_token,
account_id: stored_account_id,
} => {
*stored_id_token = id_token;
*stored_access_token = access_token;
*stored_refresh_token = refresh_token;
if let Some(new_account_id) = chatgpt_account_id {
*stored_account_id = Some(new_account_id);
with_accounts_store_mut(move |store| {
let account = store
.accounts
.iter_mut()
.find(|a| a.id == account_id)
.context("Account not found")?;

match &mut account.auth_data {
AuthData::ChatGPT {
id_token: stored_id_token,
access_token: stored_access_token,
refresh_token: stored_refresh_token,
account_id: stored_account_id,
} => {
*stored_id_token = id_token;
*stored_access_token = access_token;
*stored_refresh_token = refresh_token;
if let Some(new_account_id) = chatgpt_account_id {
*stored_account_id = Some(new_account_id);
}
}
AuthData::ApiKey { .. } => {
anyhow::bail!("Cannot update OAuth tokens for an API key account");
}
}
AuthData::ApiKey { .. } => {
anyhow::bail!("Cannot update OAuth tokens for an API key account");

if let Some(new_email) = email {
account.email = Some(new_email);
}
}

if let Some(new_email) = email {
account.email = Some(new_email);
}
if let Some(new_plan_type) = plan_type {
account.plan_type = Some(new_plan_type);
}

if let Some(new_plan_type) = plan_type {
account.plan_type = Some(new_plan_type);
}
Ok(account.clone())
})
}

let updated = account.clone();
save_accounts(&store)?;
Ok(updated)
/// Persist the last successful usage snapshot for an account.
pub fn update_account_cached_usage(
account_id: &str,
usage: UsageInfo,
updated_at: DateTime<Utc>,
) -> Result<()> {
with_accounts_store_mut(move |store| {
let account = store
.accounts
.iter_mut()
.find(|a| a.id == account_id)
.context("Account not found")?;

account.cached_usage = Some(usage);
account.cached_usage_updated_at = Some(updated_at);

Ok(())
})
}

/// Get the list of masked account IDs
Expand All @@ -246,8 +277,8 @@ pub fn get_masked_account_ids() -> Result<Vec<String>> {

/// Set the list of masked account IDs
pub fn set_masked_account_ids(ids: Vec<String>) -> Result<()> {
let mut store = load_accounts()?;
store.masked_account_ids = ids;
save_accounts(&store)?;
Ok(())
with_accounts_store_mut(move |store| {
store.masked_account_ids = ids;
Ok(())
})
}
30 changes: 27 additions & 3 deletions src-tauri/src/commands/usage.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,25 +1,49 @@
//! Usage query Tauri commands

use crate::api::usage::{get_account_usage, refresh_all_usage, warmup_account as send_warmup};
use crate::auth::{get_account, load_accounts};
use crate::auth::{get_account, load_accounts, update_account_cached_usage};
use crate::types::{UsageInfo, WarmupSummary};
use chrono::Utc;
use futures::{stream, StreamExt};

fn persist_usage_snapshot(usage: &UsageInfo) {
if usage.error.is_some() {
return;
}

if let Err(error) =
update_account_cached_usage(&usage.account_id, usage.clone(), Utc::now())
{
eprintln!(
"[Usage] Failed to persist usage snapshot for {}: {}",
usage.account_id, error
);
}
}

/// Get usage info for a specific account
#[tauri::command]
pub async fn get_usage(account_id: String) -> Result<UsageInfo, String> {
let account = get_account(&account_id)
.map_err(|e| e.to_string())?
.ok_or_else(|| format!("Account not found: {account_id}"))?;

get_account_usage(&account).await.map_err(|e| e.to_string())
let usage = get_account_usage(&account).await.map_err(|e| e.to_string())?;
persist_usage_snapshot(&usage);
Ok(usage)
}

/// Refresh usage info for all accounts
#[tauri::command]
pub async fn refresh_all_accounts_usage() -> Result<Vec<UsageInfo>, String> {
let store = load_accounts().map_err(|e| e.to_string())?;
Ok(refresh_all_usage(&store.accounts).await)
let usage_list = refresh_all_usage(&store.accounts).await;

for usage in &usage_list {
persist_usage_snapshot(usage);
}

Ok(usage_list)
}

/// Send a minimal warm-up request for one account
Expand Down
Loading