diff --git a/README.md b/README.md index 519ca2b69..934d98969 100644 --- a/README.md +++ b/README.md @@ -485,7 +485,7 @@ tokscale usage --json tokscale usage --light ``` -In the TUI, press `2` or navigate to the **Usage** tab to see subscription data. Press `u` or `r` to refresh. +In the TUI, navigate to the **Usage** tab to see subscription data. Press `u` or `r` to refresh. #### Supported Providers diff --git a/crates/tokscale-cli/src/commands/usage/amp.rs b/crates/tokscale-cli/src/commands/usage/amp.rs index c9fc6a6b7..8f3498168 100644 --- a/crates/tokscale-cli/src/commands/usage/amp.rs +++ b/crates/tokscale-cli/src/commands/usage/amp.rs @@ -58,15 +58,19 @@ fn parse_display_text(text: &str) -> Vec { let after = &text[slash_pos + 2..]; if let Some(space_pos) = after.find(|c: char| c.is_ascii_whitespace()) { if let Ok(total) = after[..space_pos].replace(',', "").parse::() { - if total > 0.0 { + if total > 0.0 && total.is_finite() && remaining.is_finite() { let used = (total - remaining).max(0.0); - let used_pct = (used / total * 100.0).clamp(0.0, 100.0); - let remaining_pct = 100.0 - used_pct; + let used_pct = if used.is_finite() { + (used / total * 100.0).clamp(0.0, 100.0) + } else { + 0.0 + }; + let remaining_pct = (100.0 - used_pct).clamp(0.0, 100.0); let mut resets_at = None; // Estimate reset time from hourly replenish rate if let Some(rate) = parse_dollar_after(text, "+$") { - if rate > 0.0 && used > 0.0 { + if rate > 0.0 && used > 0.0 && rate.is_finite() { let secs = (used / rate * 3600.0) as i64; let resets = chrono::Utc::now() + chrono::Duration::seconds(secs); resets_at = Some(resets.to_rfc3339()); diff --git a/crates/tokscale-cli/src/commands/usage/claude.rs b/crates/tokscale-cli/src/commands/usage/claude.rs index 4d2557d0a..7873b1d07 100644 --- a/crates/tokscale-cli/src/commands/usage/claude.rs +++ b/crates/tokscale-cli/src/commands/usage/claude.rs @@ -41,6 +41,13 @@ struct Window { #[derive(Debug, Deserialize)] struct TokenRefresh { access_token: Option, + refresh_token: Option, +} + +#[derive(Debug, Clone, Copy)] +enum CredentialSource { + File, + Keychain, } fn read_keychain() -> Result { @@ -53,15 +60,70 @@ pub fn has_credentials() -> bool { || super::helpers::read_keychain("Claude Code-credentials").is_ok() } -fn read_credentials() -> Result { +fn read_credentials() -> Result<(Credentials, CredentialSource)> { let home = dirs::home_dir().unwrap_or_else(|| std::path::PathBuf::from(".")); let path = home.join(".claude").join(".credentials.json"); - let content = if path.exists() { - std::fs::read_to_string(&path)? + if path.exists() { + let content = std::fs::read_to_string(&path)?; + let creds: Credentials = serde_json::from_str(&content)?; + Ok((creds, CredentialSource::File)) } else { - read_keychain()? + let content = read_keychain()?; + let creds: Credentials = serde_json::from_str(&content)?; + Ok((creds, CredentialSource::Keychain)) + } +} + +fn save_credentials(access_token: &str, refresh_token: &str, subscription_type: Option<&str>, rate_limit_tier: Option<&str>) { + let home = dirs::home_dir().unwrap_or_else(|| std::path::PathBuf::from(".")); + let path = home.join(".claude").join(".credentials.json"); + let mut oauth = serde_json::json!({ + "accessToken": access_token, + "refreshToken": refresh_token, + }); + if let Some(st) = subscription_type { + oauth["subscriptionType"] = serde_json::Value::String(st.to_string()); + } + if let Some(rlt) = rate_limit_tier { + oauth["rateLimitTier"] = serde_json::Value::String(rlt.to_string()); + } + let json = serde_json::json!({ + "claudeAiOauth": oauth + }); + let content = match serde_json::to_string_pretty(&json) { + Ok(c) => c, + Err(e) => { + eprintln!("warning: failed to serialize Claude credentials: {e}"); + return; + } }; - Ok(serde_json::from_str(&content)?) + if let Err(e) = atomic_write_secret(&path, content.as_bytes()) { + eprintln!("warning: failed to save Claude credentials: {e}"); + } +} + +fn atomic_write_secret(path: &std::path::Path, data: &[u8]) -> std::io::Result<()> { + if path.parent().is_none() { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "path has no parent directory", + )); + } + let temp_path = path.with_extension("tmp"); + { + let mut f = std::fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(&temp_path)?; + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + f.set_permissions(std::fs::Permissions::from_mode(0o600))?; + } + std::io::Write::write_all(&mut f, data)?; + } + std::fs::rename(&temp_path, path)?; + Ok(()) } async fn refresh_token(client: &reqwest::Client, rt: &str) -> Result { @@ -115,20 +177,21 @@ fn window_metric(label: &str, w: &Window) -> UsageMetric { pub fn fetch() -> Result { let rt = tokio::runtime::Runtime::new()?; rt.block_on(async { - let creds = read_credentials()?; + let (creds, source) = read_credentials()?; let oauth = creds.claude_ai_oauth.ok_or_else(|| { anyhow::anyhow!("No Claude OAuth credentials. Run 'claude' to log in.") })?; let access_token = oauth .access_token + .clone() .ok_or_else(|| anyhow::anyhow!("No Claude access token."))?; - let plan = oauth.subscription_type.map(|s| { + let plan = oauth.subscription_type.as_ref().map(|s| { let tier = oauth.rate_limit_tier.as_deref().and_then(|t| { t.rsplit('_').next() }); match tier { - Some(mult) => format!("{} {}", capitalize(&s), mult), - None => capitalize(&s), + Some(mult) => format!("{} {}", capitalize(s), mult), + None => capitalize(s), } }); @@ -143,7 +206,18 @@ pub fn fetch() -> Result { let refreshed = refresh_token(&client, rt).await?; let new = refreshed .access_token + .clone() .ok_or_else(|| anyhow::anyhow!("Refresh returned no token."))?; + if matches!(source, CredentialSource::File) { + if let Some(new_rt) = refreshed.refresh_token.as_deref() { + save_credentials( + &new, + new_rt, + oauth.subscription_type.as_deref(), + oauth.rate_limit_tier.as_deref(), + ); + } + } fetch_usage(&client, &new).await? } Err(e) => return Err(e), diff --git a/crates/tokscale-cli/src/commands/usage/codex.rs b/crates/tokscale-cli/src/commands/usage/codex.rs index 5a904493e..883395fd4 100644 --- a/crates/tokscale-cli/src/commands/usage/codex.rs +++ b/crates/tokscale-cli/src/commands/usage/codex.rs @@ -17,6 +17,7 @@ struct Tokens { access_token: Option, refresh_token: Option, account_id: Option, + id_token: Option, } #[derive(Debug, Deserialize)] @@ -44,9 +45,18 @@ struct Window { #[derive(Debug, Deserialize)] struct Refresh { access_token: Option, + refresh_token: Option, + #[allow(dead_code)] + expires_in: Option, +} + +#[derive(Debug, Clone)] +enum CredentialSource { + File(std::path::PathBuf), + Keychain, } -fn read_credentials() -> Result { +fn read_credentials() -> Result<(Auth, CredentialSource)> { let home = dirs::home_dir().unwrap_or_else(|| std::path::PathBuf::from(".")); let mut paths: Vec = Vec::new(); @@ -63,7 +73,7 @@ fn read_credentials() -> Result { if let Ok(auth) = serde_json::from_str::(&content) { // Only accept if tokens contains a usable access_token if auth.tokens.as_ref().and_then(|t| t.access_token.as_ref()).is_some() { - return Ok(auth); + return Ok((auth, CredentialSource::File(p.clone()))); } } } @@ -73,7 +83,7 @@ fn read_credentials() -> Result { if let Ok(raw) = super::helpers::read_keychain("Codex Auth") { if let Ok(auth) = serde_json::from_str::(&raw) { if auth.tokens.as_ref().and_then(|t| t.access_token.as_ref()).is_some() { - return Ok(auth); + return Ok((auth, CredentialSource::Keychain)); } } } @@ -81,6 +91,63 @@ fn read_credentials() -> Result { anyhow::bail!("No Codex credentials found. Run 'codex' to log in.") } +fn save_credentials( + path: &std::path::Path, + access_token: &str, + refresh_token: &str, + account_id: Option<&str>, + id_token: Option<&str>, +) { + let mut tokens = serde_json::json!({ + "access_token": access_token, + "refresh_token": refresh_token, + }); + if let Some(aid) = account_id { + tokens["account_id"] = serde_json::Value::String(aid.to_string()); + } + if let Some(it) = id_token { + tokens["id_token"] = serde_json::Value::String(it.to_string()); + } + let json = serde_json::json!({ + "tokens": tokens, + "last_refresh": chrono::Utc::now().to_rfc3339(), + }); + let content = match serde_json::to_string_pretty(&json) { + Ok(c) => c, + Err(e) => { + eprintln!("warning: failed to serialize Codex credentials: {e}"); + return; + } + }; + if let Err(e) = atomic_write_secret(path, content.as_bytes()) { + eprintln!("warning: failed to save Codex credentials: {e}"); + } +} + +fn atomic_write_secret(path: &std::path::Path, data: &[u8]) -> std::io::Result<()> { + if path.parent().is_none() { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "path has no parent directory", + )); + } + let temp_path = path.with_extension("tmp"); + { + let mut f = std::fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(&temp_path)?; + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + f.set_permissions(std::fs::Permissions::from_mode(0o600))?; + } + std::io::Write::write_all(&mut f, data)?; + } + std::fs::rename(&temp_path, path)?; + Ok(()) +} + pub fn has_credentials() -> bool { let home = dirs::home_dir().unwrap_or_else(|| std::path::PathBuf::from(".")); if let Ok(codex_home) = std::env::var("CODEX_HOME") { @@ -140,12 +207,13 @@ async fn fetch_usage(client: &reqwest::Client, token: &str, account_id: Option<& pub fn fetch() -> Result { let rt = tokio::runtime::Runtime::new()?; rt.block_on(async { - let auth = read_credentials()?; + let (auth, source) = read_credentials()?; let tokens = auth .tokens .ok_or_else(|| anyhow::anyhow!("No Codex tokens."))?; let access_token = tokens .access_token + .clone() .ok_or_else(|| anyhow::anyhow!("No Codex access token."))?; let account_id = tokens.account_id.as_deref(); @@ -153,14 +221,26 @@ pub fn fetch() -> Result { let resp = match fetch_usage(&client, &access_token, account_id).await { Ok(r) => r, Err(e) if e.to_string().contains("NEEDS_AUTH") => { - let rt = tokens + let rt_str = tokens .refresh_token .as_ref() .ok_or_else(|| anyhow::anyhow!("No refresh token."))?; - let refreshed = refresh_token(&client, rt).await?; + let refreshed = refresh_token(&client, rt_str).await?; let new = refreshed .access_token + .clone() .ok_or_else(|| anyhow::anyhow!("Refresh returned no token."))?; + if let CredentialSource::File(ref path) = source { + if let Some(new_rt) = refreshed.refresh_token.as_deref() { + save_credentials( + path, + &new, + new_rt, + tokens.account_id.as_deref(), + tokens.id_token.as_deref(), + ); + } + } fetch_usage(&client, &new, account_id).await? } Err(e) => return Err(e), diff --git a/crates/tokscale-cli/src/commands/usage/copilot.rs b/crates/tokscale-cli/src/commands/usage/copilot.rs index 9cd3a93a3..42d26c2cf 100644 --- a/crates/tokscale-cli/src/commands/usage/copilot.rs +++ b/crates/tokscale-cli/src/commands/usage/copilot.rs @@ -164,7 +164,33 @@ pub fn has_credentials() -> bool { if super::helpers::read_keychain("gh:github.com").is_ok() { return true; } - gh_config_dir().join("hosts.yml").exists() + // Check that hosts.yml exists and contains a github.com: section with an oauth_token + let path = gh_config_dir().join("hosts.yml"); + if !path.exists() { + return false; + } + let Ok(content) = std::fs::read_to_string(&path) else { + return false; + }; + let mut in_github = false; + for line in content.lines() { + let trimmed = line.trim(); + if trimmed == "github.com:" { + in_github = true; + continue; + } + // A non-indented, non-empty, non-comment line starts a new section + if in_github && !line.starts_with(' ') && !line.starts_with('\t') && !trimmed.is_empty() && !trimmed.starts_with('#') { + in_github = false; + } + if in_github && trimmed.starts_with("oauth_token:") { + let token = trimmed.trim_start_matches("oauth_token:").trim(); + if !token.is_empty() { + return true; + } + } + } + false } pub fn fetch() -> Result { @@ -189,14 +215,14 @@ pub fn fetch() -> Result { for (key, value) in snapshots { let pct_remaining = value.get("percent_remaining") - .and_then(|v| v.as_i64()) - .unwrap_or(100) - .clamp(0, 100); + .and_then(|v| v.as_f64()) + .unwrap_or(100.0) + .clamp(0.0, 100.0); let remaining = value.get("remaining").and_then(|v| v.as_i64()); let entitlement = value.get("entitlement").and_then(|v| v.as_i64()); - let used_pct = (100 - pct_remaining) as f64; - let remaining_pct = pct_remaining as f64; + let used_pct = 100.0 - pct_remaining; + let remaining_pct = pct_remaining; let remaining_label = match (remaining, entitlement) { (Some(r), Some(e)) => Some(format!("{r}/{e} left")), diff --git a/crates/tokscale-cli/src/commands/usage/kimi.rs b/crates/tokscale-cli/src/commands/usage/kimi.rs index 1850b8081..8a2caaf58 100644 --- a/crates/tokscale-cli/src/commands/usage/kimi.rs +++ b/crates/tokscale-cli/src/commands/usage/kimi.rs @@ -80,7 +80,40 @@ fn save_credentials(access_token: &str, refresh_token: &str, expires_in: i64) { "scope": "kimi-code", "token_type": "Bearer" }); - let _ = std::fs::write(&path, serde_json::to_string_pretty(&json).unwrap_or_default()); + let content = match serde_json::to_string_pretty(&json) { + Ok(c) => c, + Err(e) => { + eprintln!("warning: failed to serialize Kimi credentials: {e}"); + return; + } + }; + if let Err(e) = atomic_write_secret(&path, content.as_bytes()) { + eprintln!("warning: failed to save Kimi credentials: {e}"); + } +} + +fn atomic_write_secret(path: &std::path::Path, data: &[u8]) -> std::io::Result<()> { + if path.parent().is_none() { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "path has no parent directory", + )); + } + let temp_path = path.with_extension("tmp"); + { + let mut f = std::fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(&temp_path)?; + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + f.set_permissions(std::fs::Permissions::from_mode(0o600))?; + } + std::io::Write::write_all(&mut f, data)?; + } + std::fs::rename(&temp_path, path)?; + Ok(()) } fn needs_refresh(expires_at: Option) -> bool { diff --git a/crates/tokscale-cli/src/commands/usage/minimax.rs b/crates/tokscale-cli/src/commands/usage/minimax.rs index be70a5600..95a441f15 100644 --- a/crates/tokscale-cli/src/commands/usage/minimax.rs +++ b/crates/tokscale-cli/src/commands/usage/minimax.rs @@ -98,7 +98,7 @@ fn infer_plan(total: i64) -> Option { } fn epoch_to_ms(ts: i64) -> i64 { - if ts.abs() < 1_000_000_000 { ts * 1000 } else { ts } + if ts.abs() > 10_000_000_000 { ts } else { ts * 1000 } } fn parse_end_time(ts: i64) -> String { diff --git a/crates/tokscale-cli/src/commands/usage/mod.rs b/crates/tokscale-cli/src/commands/usage/mod.rs index 002e03fef..edbd4dbc2 100644 --- a/crates/tokscale-cli/src/commands/usage/mod.rs +++ b/crates/tokscale-cli/src/commands/usage/mod.rs @@ -2,7 +2,7 @@ mod amp; mod claude; mod codex; mod copilot; -mod helpers; +pub mod helpers; mod kimi; mod minimax; mod zai; @@ -51,6 +51,12 @@ pub fn save_cache(data: &[UsageOutput]) { let _ = std::fs::write(&path, serde_json::to_string(&json).unwrap_or_default()); } +pub fn clear_cache() { + if let Some(path) = cache_path() { + let _ = std::fs::remove_file(&path); + } +} + pub fn load_cache() -> Option> { let path = cache_path()?; let content = std::fs::read_to_string(&path).ok()?; @@ -93,13 +99,10 @@ pub fn fetch_all() -> Vec { std::thread::scope(|s| { active .into_iter() - .map(|(name, _, fetch)| { + .map(|(_, _, fetch)| { s.spawn(move || match fetch() { Ok(o) => Some(o), - Err(e) => { - eprintln!("{name}: {e}"); - None - } + Err(_) => None, }) }) .collect::>() @@ -114,18 +117,32 @@ pub fn fetch_all() -> Vec { const BAR_WIDTH: usize = 12; const CARD_WIDTH: usize = 62; +fn truncate(s: &str, max_len: usize) -> String { + let truncated: String = s.chars().take(max_len).collect(); + if truncated.len() < s.len() { + format!("{}…", truncated) + } else { + truncated + } +} + fn render_light(output: &UsageOutput) { println!("╭{}╮", "─".repeat(CARD_WIDTH)); + // Provider header + println!("│ {: String { - let pct = remaining_percent.clamp(0.0, 100.0) / 100.0; - let filled = (pct * BAR_WIDTH as f64).round() as usize; - let empty = BAR_WIDTH - filled; - format!("[{}{}]", "=".repeat(filled), "-".repeat(empty)) -} - -fn format_reset_time(resets_at: &str) -> String { - use chrono::{DateTime, Duration, Utc}; - let dt = match DateTime::parse_from_rfc3339(resets_at) { - Ok(d) => d.with_timezone(&Utc), - Err(_) => return format!("resets {resets_at}"), - }; - let diff = dt - Utc::now(); - if diff <= Duration::zero() { - return "resets now".into(); - } - let total_mins = diff.num_minutes(); - if total_mins < 60 { - format!("resets in {total_mins}m") - } else if total_mins < 24 * 60 { - let hours = diff.num_hours(); - let mins = (diff - Duration::hours(hours)).num_minutes(); - if mins > 0 { - format!("resets in {hours}h {mins}m") - } else { - format!("resets in {hours}h") - } - } else if diff.num_days() < 7 { - format!("resets {} {}", dt.format("%a"), dt.format("%-I%P")) - } else { - format!("resets {}", dt.format("%b %-d")) - } -} - pub fn render(frame: &mut Frame, app: &mut App, area: Rect) { let block = Block::default() .borders(Borders::ALL) @@ -112,11 +78,11 @@ fn render_loaded(frame: &mut Frame, app: &App, area: Rect, outputs: &[crate::com for m in &output.metrics { let remaining = m.remaining_label.clone().unwrap_or_else(|| format!("{:.0}% left", m.remaining_percent)); - let bar = render_ascii_bar(m.remaining_percent); + let bar = helpers::render_ascii_bar(m.remaining_percent, BAR_WIDTH); let reset = m .resets_at .as_ref() - .map(|r| format_reset_time(r)) + .map(|r| helpers::format_reset_time(r)) .unwrap_or_default(); let label = Span::styled(