Skip to content

Commit 4e682ea

Browse files
fix(grok): classify cadence from full billing cycle (#419)
* fix(grok): classify Weekly/Monthly from reset cycle for pace Match upstream docs/grok.md: the bar stays Credits, and Weekly vs Monthly comes from whether resetsAt matches a common cycle (4-12 days weekly, 20-45 monthly). Set window_minutes so tray pace works. Near-reset untyped windows keep Weekly (steipete#2929) instead of hard-coding 7 days on every path. * Fix Grok cadence classification structure --------- Co-authored-by: Manu Cuervo <169969406+manucuervodj@users.noreply.github.com>
1 parent a008d9f commit 4e682ea

4 files changed

Lines changed: 209 additions & 33 deletions

File tree

‎apps/desktop-tauri/src-tauri/src/commands/bridge.rs‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -326,7 +326,12 @@ impl ProviderUsageSnapshot {
326326
provider_id: id.cli_name().to_string(),
327327
display_name: id.display_name().to_string(),
328328
primary: primary_snap,
329-
primary_label: Some(metadata.session_label.to_string()),
329+
primary_label: Some(
330+
usage
331+
.primary_label
332+
.clone()
333+
.unwrap_or_else(|| metadata.session_label.to_string()),
334+
),
330335
secondary: secondary_snap,
331336
secondary_label: usage
332337
.secondary

‎rust/src/cli/usage.rs‎

Lines changed: 31 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -554,11 +554,18 @@ fn append_usage_window_lines(
554554
metadata: &crate::core::ProviderMetadata,
555555
use_color: bool,
556556
) {
557-
append_window_line(lines, metadata.session_label, &usage.primary, use_color);
558-
// Upstream 0.50.1 #2957: pace for the 5-hour session window.
559-
if usage.primary.window_minutes == Some(crate::core::SESSION_WINDOW_MINUTES)
560-
&& let Some(pace) =
561-
UsagePace::weekly(&usage.primary, None, crate::core::SESSION_WINDOW_MINUTES)
557+
let primary_label = usage
558+
.primary_label
559+
.as_deref()
560+
.unwrap_or(metadata.session_label);
561+
append_window_line(lines, primary_label, &usage.primary, use_color);
562+
// Pace for primary windows whose provider published a common cadence.
563+
let pace_minutes = usage.primary.window_minutes.filter(|minutes| {
564+
*minutes == crate::core::SESSION_WINDOW_MINUTES
565+
|| *minutes >= crate::core::WEEKLY_WINDOW_MINUTES
566+
});
567+
if let Some(minutes) = pace_minutes
568+
&& let Some(pace) = UsagePace::weekly(&usage.primary, None, minutes)
562569
{
563570
lines.push(format!(
564571
" Pace: {} {}",
@@ -643,14 +650,18 @@ fn append_model_specific_line(
643650
pub fn render_brief_text(provider: ProviderId, result: &ProviderFetchResult) -> String {
644651
let metadata = instantiate_provider(provider).metadata().clone();
645652
let usage = &result.usage;
653+
let primary_label = usage
654+
.primary_label
655+
.as_deref()
656+
.unwrap_or(metadata.session_label);
646657
let mut parts = Vec::new();
647658
let reset = if usage.primary.is_informational {
648-
parts.push(format!("{} unavailable", metadata.session_label));
659+
parts.push(format!("{primary_label} unavailable"));
649660
usage.secondary.as_ref().unwrap_or(&usage.primary)
650661
} else {
651662
parts.push(format!(
652663
"{} {}",
653-
metadata.session_label,
664+
primary_label,
654665
format_percent(usage.primary.used_percent)
655666
));
656667
&usage.primary
@@ -799,6 +810,19 @@ mod tests {
799810
);
800811
}
801812

813+
#[test]
814+
fn primary_label_override_is_shared_by_full_and_brief_renderers() {
815+
let result =
816+
fetch_result(UsageSnapshot::new(RateWindow::new(42.0)).with_primary_label("Monthly"));
817+
818+
let full = render_text_with_status(ProviderId::Grok, &result, None, false);
819+
let brief = render_brief_text(ProviderId::Grok, &result);
820+
821+
assert!(full.contains("Monthly:"));
822+
assert!(brief.contains("Grok: Monthly 42%"));
823+
assert!(!brief.contains("Credits 42%"));
824+
}
825+
802826
#[test]
803827
fn gemini_plan_preserves_acronym_casing() {
804828
let result = fetch_result(

‎rust/src/core/usage_snapshot.rs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,11 @@ pub struct UsageSnapshot {
7676
/// Primary rate window (usually session-based, e.g., 5-hour for Claude)
7777
pub primary: RateWindow,
7878

79+
/// Provider-resolved label for the primary rate window when the stable
80+
/// provider metadata label is not specific enough for this snapshot.
81+
#[serde(default, skip_serializing_if = "Option::is_none")]
82+
pub primary_label: Option<String>,
83+
7984
/// Secondary rate window (usually weekly/monthly)
8085
#[serde(skip_serializing_if = "Option::is_none")]
8186
pub secondary: Option<RateWindow>,
@@ -113,6 +118,7 @@ impl UsageSnapshot {
113118
pub fn new(primary: RateWindow) -> Self {
114119
Self {
115120
primary,
121+
primary_label: None,
116122
secondary: None,
117123
model_specific: None,
118124
tertiary: None,
@@ -124,6 +130,12 @@ impl UsageSnapshot {
124130
}
125131
}
126132

133+
/// Builder pattern: override the primary window label for this snapshot.
134+
pub fn with_primary_label(mut self, label: impl Into<String>) -> Self {
135+
self.primary_label = Some(label.into());
136+
self
137+
}
138+
127139
/// Builder pattern: set secondary window
128140
pub fn with_secondary(mut self, secondary: RateWindow) -> Self {
129141
self.secondary = Some(secondary);

‎rust/src/providers/grok/mod.rs‎

Lines changed: 160 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ impl GrokProvider {
3333
metadata: ProviderMetadata {
3434
id: ProviderId::Grok,
3535
display_name: "Grok",
36-
session_label: "Monthly",
36+
session_label: "Credits",
3737
weekly_label: "On-demand",
3838
supports_opus: false,
3939
supports_credits: false,
@@ -434,6 +434,25 @@ fn text_field(value: &Value, key: &str) -> Option<String> {
434434
struct GrokBillingSnapshot {
435435
used_percent: Option<f64>,
436436
resets_at: Option<DateTime<Utc>>,
437+
window_minutes: Option<u32>,
438+
}
439+
440+
/// Classify Grok from the full billing-cycle duration, not time remaining.
441+
/// This preserves the upstream #2431/#2566 invariant that a monthly plan near
442+
/// its reset must not become a weekly plan.
443+
fn primary_label_for_cycle_minutes(minutes: u32) -> Option<&'static str> {
444+
const DAY_MINUTES: u32 = 24 * 60;
445+
if minutes <= 60 {
446+
return None;
447+
}
448+
let days = (minutes + DAY_MINUTES / 2) / DAY_MINUTES;
449+
if (4..=12).contains(&days) {
450+
Some("Weekly")
451+
} else if (20..=45).contains(&days) {
452+
Some("Monthly")
453+
} else {
454+
None
455+
}
437456
}
438457

439458
fn result_from_billing(
@@ -443,17 +462,30 @@ fn result_from_billing(
443462
team_id: Option<String>,
444463
login_method: Option<String>,
445464
) -> ProviderFetchResult {
446-
// Upstream #2431 / #2566: do not infer windowMinutes from time-until-reset.
447-
// A monthly quota near its reset would otherwise be misclassified as weekly.
465+
// Dynamic cadence is provider-owned and comes only from a complete billing
466+
// cycle. A reset timestamp alone is insufficient because monthly quotas can
467+
// have only a few days remaining.
468+
let primary_label = billing
469+
.window_minutes
470+
.and_then(primary_label_for_cycle_minutes);
448471
let primary = match billing.used_percent {
449-
Some(used_percent) => RateWindow::with_details(used_percent, None, billing.resets_at, None),
472+
Some(used_percent) => RateWindow::with_details(
473+
used_percent,
474+
billing.window_minutes,
475+
billing.resets_at,
476+
None,
477+
),
450478
None => {
451479
let mut window = RateWindow::informational("Usage unavailable");
452480
window.resets_at = billing.resets_at;
481+
window.window_minutes = billing.window_minutes;
453482
window
454483
}
455484
};
456485
let mut usage = UsageSnapshot::new(primary);
486+
if let Some(label) = primary_label {
487+
usage = usage.with_primary_label(label);
488+
}
457489
usage.account_email = email;
458490
usage.account_organization = team_id;
459491
usage.login_method = login_method;
@@ -532,29 +564,50 @@ fn parse_grpc_web_response(data: &[u8]) -> Result<GrokBillingSnapshot, ProviderE
532564
})
533565
.map(|field| field.value as f64);
534566

567+
let now = Utc::now();
535568
let resets_at = scan
536569
.varints
537570
.iter()
538-
.filter_map(|field| {
539-
// Varint timestamps are Unix seconds inside the range checked below.
540-
#[allow(
541-
clippy::cast_possible_wrap,
542-
reason = "varint timestamps are bounded to the Unix-seconds range checked below"
543-
)]
544-
let seconds = field.value as i64;
545-
(1_700_000_000..=2_100_000_000)
546-
.contains(&field.value)
547-
.then(|| Utc.timestamp_opt(seconds, 0).single())
548-
.flatten()
549-
})
550-
.filter(|dt| *dt > Utc::now())
571+
.filter_map(varint_timestamp)
572+
.filter(|dt| *dt > now)
551573
.min();
552574
Ok(GrokBillingSnapshot {
553575
used_percent,
554576
resets_at,
577+
window_minutes: current_period_window_minutes(&scan, now),
555578
})
556579
}
557580

581+
fn varint_timestamp(field: &VarintField) -> Option<DateTime<Utc>> {
582+
// Varint timestamps are Unix seconds inside the range checked below.
583+
#[allow(
584+
clippy::cast_possible_wrap,
585+
reason = "varint timestamps are bounded to the Unix-seconds range checked below"
586+
)]
587+
let seconds = field.value as i64;
588+
(1_700_000_000..=2_100_000_000)
589+
.contains(&field.value)
590+
.then(|| Utc.timestamp_opt(seconds, 0).single())
591+
.flatten()
592+
}
593+
594+
fn current_period_window_minutes(scan: &ProtoScan, now: DateTime<Utc>) -> Option<u32> {
595+
let timestamp_at = |path: &[u64]| {
596+
scan.varints
597+
.iter()
598+
.find(|field| field.path.as_slice() == path)
599+
.and_then(varint_timestamp)
600+
};
601+
let start = timestamp_at(&[1, 8, 2, 1])?;
602+
let end = timestamp_at(&[1, 8, 3, 1])?;
603+
if start > now || end <= now || end <= start {
604+
return None;
605+
}
606+
u32::try_from((end - start).num_minutes())
607+
.ok()
608+
.filter(|minutes| *minutes > 0)
609+
}
610+
558611
fn grpc_web_data_frames(data: &[u8]) -> Vec<Vec<u8>> {
559612
let mut frames = Vec::new();
560613
let mut index = 0;
@@ -591,6 +644,7 @@ struct Fixed32Field {
591644
}
592645

593646
struct VarintField {
647+
path: Vec<u64>,
594648
value: u64,
595649
}
596650

@@ -624,17 +678,20 @@ impl ProtoScan {
624678
wire: u64,
625679
) -> Option<usize> {
626680
match wire {
627-
0 => self.scan_varint(data, i),
681+
0 => self.scan_varint(data, i, path),
628682
2 => self.scan_length_delimited(data, i, path, depth),
629683
5 => self.scan_fixed32(data, i, path),
630684
1 => Some(i.saturating_add(8)),
631685
_ => None,
632686
}
633687
}
634688

635-
fn scan_varint(&mut self, data: &[u8], i: usize) -> Option<usize> {
689+
fn scan_varint(&mut self, data: &[u8], i: usize, path: &[u64]) -> Option<usize> {
636690
let (value, next) = read_varint(data, i)?;
637-
self.varints.push(VarintField { value });
691+
self.varints.push(VarintField {
692+
path: path.to_vec(),
693+
value,
694+
});
638695
Some(next)
639696
}
640697

@@ -794,22 +851,99 @@ mod tests {
794851
}
795852

796853
#[test]
797-
fn billing_snapshot_leaves_window_minutes_unset() {
798-
// A monthly quota with six days left must not be reported as weekly.
799-
let resets = Utc::now() + chrono::Duration::days(6);
854+
fn billing_snapshot_uses_full_weekly_cycle_for_pace() {
855+
let now = Utc::now();
856+
let resets = now + chrono::Duration::days(2);
800857
let result = result_from_billing(
801858
GrokBillingSnapshot {
802859
used_percent: Some(12.0),
803860
resets_at: Some(resets),
861+
window_minutes: Some(crate::core::WEEKLY_WINDOW_MINUTES),
862+
},
863+
"web",
864+
None,
865+
None,
866+
Some("SuperGrok".into()),
867+
);
868+
assert_eq!(
869+
result.usage.primary.window_minutes,
870+
Some(crate::core::WEEKLY_WINDOW_MINUTES)
871+
);
872+
assert_eq!(result.usage.primary_label.as_deref(), Some("Weekly"));
873+
let pace = crate::core::UsagePace::weekly(
874+
&result.usage.primary,
875+
Some(now),
876+
crate::core::WEEKLY_WINDOW_MINUTES,
877+
);
878+
assert!(pace.is_some(), "weekly window + reset must yield pace");
879+
}
880+
881+
#[test]
882+
fn monthly_cycle_stays_monthly_with_six_days_remaining() {
883+
let now = Utc::now();
884+
let resets = now + chrono::Duration::days(6);
885+
let monthly_minutes = 31 * 24 * 60;
886+
let result = result_from_billing(
887+
GrokBillingSnapshot {
888+
used_percent: Some(40.0),
889+
resets_at: Some(resets),
890+
window_minutes: Some(monthly_minutes),
891+
},
892+
"cli",
893+
None,
894+
None,
895+
Some("SuperGrok Heavy".into()),
896+
);
897+
assert_eq!(result.usage.primary_label.as_deref(), Some("Monthly"));
898+
assert_eq!(result.usage.primary.window_minutes, Some(monthly_minutes));
899+
}
900+
901+
#[test]
902+
fn reset_distance_alone_does_not_invent_a_cadence() {
903+
let resets = Utc::now() + chrono::Duration::days(6);
904+
let result = result_from_billing(
905+
GrokBillingSnapshot {
906+
used_percent: Some(80.0),
907+
resets_at: Some(resets),
908+
window_minutes: None,
804909
},
805910
"web",
806911
None,
807912
None,
808913
Some("SuperGrok".into()),
809914
);
810915
assert_eq!(result.usage.primary.window_minutes, None);
811-
assert_eq!(result.usage.primary.resets_at, Some(resets));
812-
assert_eq!(result.usage.login_method.as_deref(), Some("SuperGrok"));
916+
assert_eq!(result.usage.primary_label, None);
917+
}
918+
919+
#[test]
920+
fn current_period_paths_define_the_full_cycle() {
921+
let now = Utc.timestamp_opt(1_800_000_000, 0).single().unwrap();
922+
let start = now - chrono::Duration::days(25);
923+
let end = now + chrono::Duration::days(6);
924+
let scan = ProtoScan {
925+
fixed32: Vec::new(),
926+
varints: vec![
927+
VarintField {
928+
path: vec![1, 8, 2, 1],
929+
value: u64::try_from(start.timestamp()).unwrap(),
930+
},
931+
VarintField {
932+
path: vec![1, 8, 3, 1],
933+
value: u64::try_from(end.timestamp()).unwrap(),
934+
},
935+
],
936+
order: 0,
937+
};
938+
939+
assert_eq!(
940+
current_period_window_minutes(&scan, now),
941+
Some(31 * 24 * 60)
942+
);
943+
assert_eq!(
944+
primary_label_for_cycle_minutes(current_period_window_minutes(&scan, now).unwrap()),
945+
Some("Monthly")
946+
);
813947
}
814948

815949
#[test]
@@ -819,6 +953,7 @@ mod tests {
819953
GrokBillingSnapshot {
820954
used_percent: None,
821955
resets_at: Some(resets),
956+
window_minutes: None,
822957
},
823958
"cli",
824959
Some("user@example.com".into()),

0 commit comments

Comments
 (0)