From 655550489660a6bb8024138d8299bfe5d399054d Mon Sep 17 00:00:00 2001 From: yearth Date: Mon, 23 Mar 2026 14:49:35 +0800 Subject: [PATCH 1/3] feat(usage): read rate_limits from stdin, keep API as fallback --- src/config/types.rs | 16 +++ src/core/segments/usage.rs | 241 +++++++++++++++++++++++++++---------- 2 files changed, 195 insertions(+), 62 deletions(-) diff --git a/src/config/types.rs b/src/config/types.rs index e5a78dc1..92a1dcb6 100644 --- a/src/config/types.rs +++ b/src/config/types.rs @@ -110,6 +110,20 @@ pub struct OutputStyle { pub name: String, } +#[derive(Deserialize)] +pub struct RateLimitPeriod { + /// Usage as a percentage (0-100). Provided directly by Claude Code via stdin. + pub used_percentage: Option, + /// Unix timestamp (seconds) when this window resets. + pub resets_at: Option, +} + +#[derive(Deserialize)] +pub struct RateLimits { + pub five_hour: Option, + pub seven_day: Option, +} + #[derive(Deserialize)] pub struct InputData { pub model: Model, @@ -117,6 +131,8 @@ pub struct InputData { pub transcript_path: String, pub cost: Option, pub output_style: Option, + /// Rate limit data passed directly by Claude Code — no API call needed. + pub rate_limits: Option, } // OpenAI-style nested token details diff --git a/src/core/segments/usage.rs b/src/core/segments/usage.rs index d5dd9bde..afe09f4b 100644 --- a/src/core/segments/usage.rs +++ b/src/core/segments/usage.rs @@ -1,18 +1,21 @@ use super::{Segment, SegmentData}; -use crate::config::{InputData, SegmentId}; -use crate::utils::credentials; -use chrono::{DateTime, Datelike, Duration, Local, Timelike, Utc}; +use crate::config::{InputData, RateLimitPeriod, SegmentId}; +use chrono::{DateTime, Datelike, Duration, Local, TimeZone, Timelike, Utc}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; +// --------------------------------------------------------------------------- +// API fallback types (used only when stdin does not carry rate_limits) +// --------------------------------------------------------------------------- + #[derive(Debug, Deserialize)] struct ApiUsageResponse { - five_hour: UsagePeriod, - seven_day: UsagePeriod, + five_hour: ApiUsagePeriod, + seven_day: ApiUsagePeriod, } #[derive(Debug, Deserialize)] -struct UsagePeriod { +struct ApiUsagePeriod { utilization: f64, resets_at: Option, } @@ -25,6 +28,10 @@ struct ApiUsageCache { cached_at: String, } +// --------------------------------------------------------------------------- +// Segment +// --------------------------------------------------------------------------- + #[derive(Default)] pub struct UsageSegment; @@ -33,21 +40,37 @@ impl UsageSegment { Self } - fn get_circle_icon(utilization: f64) -> String { + fn get_circle_icon(utilization: f64) -> &'static str { let percent = (utilization * 100.0) as u8; match percent { - 0..=12 => "\u{f0a9e}".to_string(), // circle_slice_1 - 13..=25 => "\u{f0a9f}".to_string(), // circle_slice_2 - 26..=37 => "\u{f0aa0}".to_string(), // circle_slice_3 - 38..=50 => "\u{f0aa1}".to_string(), // circle_slice_4 - 51..=62 => "\u{f0aa2}".to_string(), // circle_slice_5 - 63..=75 => "\u{f0aa3}".to_string(), // circle_slice_6 - 76..=87 => "\u{f0aa4}".to_string(), // circle_slice_7 - _ => "\u{f0aa5}".to_string(), // circle_slice_8 + 0..=12 => "\u{f0a9e}", // circle_slice_1 + 13..=25 => "\u{f0a9f}", // circle_slice_2 + 26..=37 => "\u{f0aa0}", // circle_slice_3 + 38..=50 => "\u{f0aa1}", // circle_slice_4 + 51..=62 => "\u{f0aa2}", // circle_slice_5 + 63..=75 => "\u{f0aa3}", // circle_slice_6 + 76..=87 => "\u{f0aa4}", // circle_slice_7 + _ => "\u{f0aa5}", // circle_slice_8 } } - fn format_reset_time(reset_time_str: Option<&str>) -> String { + /// Format a Unix timestamp (seconds) into a human-readable reset time string. + fn format_reset_time_unix(ts: i64) -> String { + let dt = Utc.timestamp_opt(ts, 0).single(); + match dt { + Some(utc_dt) => { + let mut local_dt = utc_dt.with_timezone(&Local); + if local_dt.minute() > 45 { + local_dt += Duration::hours(1); + } + format!("{}-{}-{}", local_dt.month(), local_dt.day(), local_dt.hour()) + } + None => "?".to_string(), + } + } + + /// Format an RFC 3339 reset time string (used by the API fallback path). + fn format_reset_time_rfc3339(reset_time_str: Option<&str>) -> String { if let Some(time_str) = reset_time_str { if let Ok(dt) = DateTime::parse_from_rfc3339(time_str) { let mut local_dt = dt.with_timezone(&Local); @@ -65,6 +88,18 @@ impl UsageSegment { "?".to_string() } + /// Extract utilization percentage from a stdin RateLimitPeriod. + /// Clamps to [0, 100] and returns None if the value is missing or non-finite. + fn parse_utilization(period: &RateLimitPeriod) -> Option { + let v = period.used_percentage?; + if !v.is_finite() { + return None; + } + Some(v.clamp(0.0, 100.0)) + } + + // --- API fallback helpers (unchanged from original) --- + fn get_cache_path() -> Option { let home = dirs::home_dir()?; Some( @@ -79,7 +114,6 @@ impl UsageSegment { if !cache_path.exists() { return None; } - let content = std::fs::read_to_string(&cache_path).ok()?; serde_json::from_str(&content).ok() } @@ -97,9 +131,9 @@ impl UsageSegment { fn is_cache_valid(&self, cache: &ApiUsageCache, cache_duration: u64) -> bool { if let Ok(cached_at) = DateTime::parse_from_rfc3339(&cache.cached_at) { - let now = Utc::now(); - let elapsed = now.signed_duration_since(cached_at.with_timezone(&Utc)); - elapsed.num_seconds() < cache_duration as i64 + let elapsed = Utc::now().signed_duration_since(cached_at.with_timezone(&Utc)); + let secs = elapsed.num_seconds(); + secs >= 0 && secs < cache_duration as i64 } else { false } @@ -107,11 +141,9 @@ impl UsageSegment { fn get_claude_code_version() -> String { use std::process::Command; - let output = Command::new("npm") .args(["view", "@anthropic-ai/claude-code", "version"]) .output(); - match output { Ok(output) if output.status.success() => { let version = String::from_utf8_lossy(&output.stdout).trim().to_string(); @@ -121,7 +153,6 @@ impl UsageSegment { } _ => {} } - "claude-code".to_string() } @@ -130,11 +161,8 @@ impl UsageSegment { .or_else(|_| std::env::var("USERPROFILE")) .ok()?; let settings_path = format!("{}/.claude/settings.json", home); - let content = std::fs::read_to_string(&settings_path).ok()?; let settings: serde_json::Value = serde_json::from_str(&content).ok()?; - - // Try HTTPS_PROXY first, then HTTP_PROXY settings .get("env")? .get("HTTPS_PROXY") @@ -149,6 +177,8 @@ impl UsageSegment { token: &str, timeout_secs: u64, ) -> Option { + use crate::utils::credentials; + let _ = credentials::get_oauth_token(); // ensure token is accessible let url = format!("{}/api/oauth/usage", api_base_url); let user_agent = Self::get_claude_code_version(); @@ -178,13 +208,35 @@ impl UsageSegment { response.into_body().read_json().ok() } -} -impl Segment for UsageSegment { - fn collect(&self, _input: &InputData) -> Option { + /// Collect from stdin rate_limits — no network, no credentials, always fresh. + fn collect_from_stdin(&self, input: &InputData) -> Option { + let rate_limits = input.rate_limits.as_ref()?; + let five_hour = rate_limits.five_hour.as_ref()?; + let five_hour_util = Self::parse_utilization(five_hour)?; + + let seven_day_util = rate_limits + .seven_day + .as_ref() + .and_then(|p| Self::parse_utilization(p)) + .unwrap_or(0.0); + + let reset_str = rate_limits + .five_hour + .as_ref() + .and_then(|p| p.resets_at) + .map(Self::format_reset_time_unix) + .unwrap_or_else(|| "?".to_string()); + + Some(self.build_segment_data(five_hour_util, seven_day_util, reset_str)) + } + + /// Fallback: collect via Anthropic API (for older Claude Code versions + /// that do not include rate_limits in stdin). + fn collect_from_api(&self, input: &InputData) -> Option { + use crate::utils::credentials; let token = credentials::get_oauth_token()?; - // Load config from file to get segment options let config = crate::config::Config::load().ok()?; let segment_config = config.segments.iter().find(|s| s.id == SegmentId::Usage); @@ -206,16 +258,12 @@ impl Segment for UsageSegment { let cached_data = self.load_cache(); let use_cached = cached_data .as_ref() - .map(|cache| self.is_cache_valid(cache, cache_duration)) + .map(|c| self.is_cache_valid(c, cache_duration)) .unwrap_or(false); let (five_hour_util, seven_day_util, resets_at) = if use_cached { - let cache = cached_data.unwrap(); - ( - cache.five_hour_utilization, - cache.seven_day_utilization, - cache.resets_at, - ) + let c = cached_data.unwrap(); + (c.five_hour_utilization, c.seven_day_utilization, c.resets_at) } else { match self.fetch_api_usage(api_base_url, &token, timeout) { Some(response) => { @@ -233,43 +281,112 @@ impl Segment for UsageSegment { ) } None => { - if let Some(cache) = cached_data { - ( - cache.five_hour_utilization, - cache.seven_day_utilization, - cache.resets_at, - ) - } else { - return None; - } + let c = cached_data?; + (c.five_hour_utilization, c.seven_day_utilization, c.resets_at) } } }; + let reset_str = Self::format_reset_time_rfc3339(resets_at.as_deref()); + Some(self.build_segment_data(five_hour_util, seven_day_util, reset_str)) + } + + fn build_segment_data( + &self, + five_hour_util: f64, + seven_day_util: f64, + reset_str: String, + ) -> SegmentData { let dynamic_icon = Self::get_circle_icon(seven_day_util / 100.0); let five_hour_percent = five_hour_util.round() as u8; let primary = format!("{}%", five_hour_percent); - let secondary = format!("· {}", Self::format_reset_time(resets_at.as_deref())); + let secondary = format!("· {}", reset_str); let mut metadata = HashMap::new(); - metadata.insert("dynamic_icon".to_string(), dynamic_icon); - metadata.insert( - "five_hour_utilization".to_string(), - five_hour_util.to_string(), - ); - metadata.insert( - "seven_day_utilization".to_string(), - seven_day_util.to_string(), - ); - - Some(SegmentData { - primary, - secondary, - metadata, - }) + metadata.insert("dynamic_icon".to_string(), dynamic_icon.to_string()); + metadata.insert("five_hour_utilization".to_string(), five_hour_util.to_string()); + metadata.insert("seven_day_utilization".to_string(), seven_day_util.to_string()); + + SegmentData { primary, secondary, metadata } + } +} + +impl Segment for UsageSegment { + fn collect(&self, input: &InputData) -> Option { + // Prefer stdin data: zero latency, no credentials, always up to date. + // Fall back to the API path for older Claude Code versions that do not + // include rate_limits in stdin. + self.collect_from_stdin(input) + .or_else(|| self.collect_from_api(input)) } fn id(&self) -> SegmentId { SegmentId::Usage } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::{RateLimitPeriod, RateLimits}; + + fn make_input(five_pct: Option, five_resets_at: Option, seven_pct: Option) -> InputData { + InputData { + model: crate::config::Model { id: "".into(), display_name: "".into() }, + workspace: crate::config::Workspace { current_dir: "".into() }, + transcript_path: "".into(), + cost: None, + output_style: None, + rate_limits: Some(RateLimits { + five_hour: Some(RateLimitPeriod { + used_percentage: five_pct, + resets_at: five_resets_at, + }), + seven_day: Some(RateLimitPeriod { + used_percentage: seven_pct, + resets_at: None, + }), + }), + } + } + + #[test] + fn test_collect_from_stdin_basic() { + let seg = UsageSegment::new(); + let input = make_input(Some(75.0), None, Some(45.0)); + let data = seg.collect_from_stdin(&input).unwrap(); + assert_eq!(data.primary, "75%"); + assert_eq!(data.metadata["five_hour_utilization"], "75"); + assert_eq!(data.metadata["seven_day_utilization"], "45"); + } + + #[test] + fn test_collect_from_stdin_clamps_over_100() { + let seg = UsageSegment::new(); + let input = make_input(Some(120.0), None, Some(0.0)); + let data = seg.collect_from_stdin(&input).unwrap(); + assert_eq!(data.primary, "100%"); + } + + #[test] + fn test_collect_from_stdin_no_rate_limits() { + let seg = UsageSegment::new(); + let mut input = make_input(Some(50.0), None, None); + input.rate_limits = None; + assert!(seg.collect_from_stdin(&input).is_none()); + } + + #[test] + fn test_parse_utilization_non_finite() { + let period = RateLimitPeriod { used_percentage: Some(f64::NAN), resets_at: None }; + assert!(UsageSegment::parse_utilization(&period).is_none()); + } + + #[test] + fn test_circle_icon_boundaries() { + assert_eq!(UsageSegment::get_circle_icon(0.0), "\u{f0a9e}"); // 0% -> slice_1 + assert_eq!(UsageSegment::get_circle_icon(0.50), "\u{f0aa1}"); // 50% -> slice_4 (38..=50) + assert_eq!(UsageSegment::get_circle_icon(0.51), "\u{f0aa2}"); // 51% -> slice_5 + assert_eq!(UsageSegment::get_circle_icon(1.0), "\u{f0aa5}"); // 100% -> slice_8 + } +} From 0f3e4f2e96cf57cf20b889916f9cb3b63831cf22 Mon Sep 17 00:00:00 2001 From: yearth Date: Mon, 23 Mar 2026 14:56:38 +0800 Subject: [PATCH 2/3] fix: remove redundant credentials call and unused input param in API fallback --- src/core/segments/usage.rs | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/core/segments/usage.rs b/src/core/segments/usage.rs index afe09f4b..b4d74f54 100644 --- a/src/core/segments/usage.rs +++ b/src/core/segments/usage.rs @@ -177,8 +177,6 @@ impl UsageSegment { token: &str, timeout_secs: u64, ) -> Option { - use crate::utils::credentials; - let _ = credentials::get_oauth_token(); // ensure token is accessible let url = format!("{}/api/oauth/usage", api_base_url); let user_agent = Self::get_claude_code_version(); @@ -233,7 +231,7 @@ impl UsageSegment { /// Fallback: collect via Anthropic API (for older Claude Code versions /// that do not include rate_limits in stdin). - fn collect_from_api(&self, input: &InputData) -> Option { + fn collect_from_api(&self) -> Option { use crate::utils::credentials; let token = credentials::get_oauth_token()?; @@ -317,7 +315,7 @@ impl Segment for UsageSegment { // Fall back to the API path for older Claude Code versions that do not // include rate_limits in stdin. self.collect_from_stdin(input) - .or_else(|| self.collect_from_api(input)) + .or_else(|| self.collect_from_api()) } fn id(&self) -> SegmentId { From 3dcd3c8eb500e3db87e336deb1650c574c3c2743 Mon Sep 17 00:00:00 2001 From: yearth Date: Mon, 23 Mar 2026 15:32:26 +0800 Subject: [PATCH 3/3] fix: handle seven_day-only stdin data, add module doc comment --- src/core/segments/usage.rs | 50 +++++++++++++++++++++++++++++++++++--- 1 file changed, 47 insertions(+), 3 deletions(-) diff --git a/src/core/segments/usage.rs b/src/core/segments/usage.rs index b4d74f54..f0f2f12d 100644 --- a/src/core/segments/usage.rs +++ b/src/core/segments/usage.rs @@ -1,3 +1,12 @@ +//! UsageSegment displays Claude Code rate limit utilization (5-hour and 7-day windows). +//! +//! Data source priority: +//! 1. `rate_limits` field in stdin — provided directly by Claude Code, zero latency, +//! no credentials required, always current. Available in Claude Code >= 2.x. +//! 2. Anthropic API (`/api/oauth/usage`) — legacy fallback for older Claude Code versions +//! that do not include `rate_limits` in stdin. Uses a file-based cache to reduce +//! network calls. + use super::{Segment, SegmentData}; use crate::config::{InputData, RateLimitPeriod, SegmentId}; use chrono::{DateTime, Datelike, Duration, Local, TimeZone, Timelike, Utc}; @@ -208,10 +217,16 @@ impl UsageSegment { } /// Collect from stdin rate_limits — no network, no credentials, always fresh. + /// + /// Prefers five_hour data for the primary display value. If only seven_day is + /// available, falls back to that so the segment is still shown. fn collect_from_stdin(&self, input: &InputData) -> Option { let rate_limits = input.rate_limits.as_ref()?; - let five_hour = rate_limits.five_hour.as_ref()?; - let five_hour_util = Self::parse_utilization(five_hour)?; + + let five_hour_util = rate_limits + .five_hour + .as_ref() + .and_then(|p| Self::parse_utilization(p)); let seven_day_util = rate_limits .seven_day @@ -219,14 +234,18 @@ impl UsageSegment { .and_then(|p| Self::parse_utilization(p)) .unwrap_or(0.0); + // Require at least one valid utilization value to show the segment + let display_util = five_hour_util.or(if seven_day_util > 0.0 { Some(seven_day_util) } else { None })?; + let reset_str = rate_limits .five_hour .as_ref() .and_then(|p| p.resets_at) + .or_else(|| rate_limits.seven_day.as_ref().and_then(|p| p.resets_at)) .map(Self::format_reset_time_unix) .unwrap_or_else(|| "?".to_string()); - Some(self.build_segment_data(five_hour_util, seven_day_util, reset_str)) + Some(self.build_segment_data(display_util, seven_day_util, reset_str)) } /// Fallback: collect via Anthropic API (for older Claude Code versions @@ -374,6 +393,31 @@ mod tests { assert!(seg.collect_from_stdin(&input).is_none()); } + #[test] + fn test_collect_from_stdin_only_seven_day() { + // If only seven_day is present, segment should still render using it + let seg = UsageSegment::new(); + let mut input = make_input(None, None, Some(60.0)); + // Remove five_hour + if let Some(ref mut rl) = input.rate_limits { + rl.five_hour = None; + } + let data = seg.collect_from_stdin(&input).unwrap(); + assert_eq!(data.primary, "60%"); + } + + #[test] + fn test_collect_from_stdin_both_missing() { + // Both periods missing -> None + let seg = UsageSegment::new(); + let mut input = make_input(None, None, None); + if let Some(ref mut rl) = input.rate_limits { + rl.five_hour = None; + rl.seven_day = None; + } + assert!(seg.collect_from_stdin(&input).is_none()); + } + #[test] fn test_parse_utilization_non_finite() { let period = RateLimitPeriod { used_percentage: Some(f64::NAN), resets_at: None };