hanzoai/dev
2
0
Fork 0
mirror of https://github.com/hanzoai/dev.git synced 2026-10-25 03:32:32 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(rate-limits): persist reset state and dedupe warnings

Browse Source
This commit is contained in:
James Peter
2025-09-27 06:27:20 +10:00
parent f0ec0dec00
commit 51fd6015d3
2 changed files with 306 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/account_usage.rs
+252 -20
View File
@@ -14,6 +14,7 @@ const USAGE_VERSION: u32 = 1;
const USAGE_SUBDIR: &str = "usage";
const HOURLY_HISTORY_DAYS: i64 = 183; // retain ~6 months of hourly usage for history views
const UNKNOWN_RESET_RELOG_INTERVAL: Duration = Duration::hours(24);
const RESET_PASSED_TOLERANCE: Duration = Duration::seconds(5);
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum RateLimitWarningScope {
@@ -375,18 +376,6 @@ pub fn record_usage_limit_hint(
})
}
fn reset_changed(prev: Option<DateTime<Utc>>, next: Option<DateTime<Utc>>) -> bool {
match (prev, next) {
(Some(a), Some(b)) => a
.signed_duration_since(b)
.num_seconds()
.abs()
> 5,
(None, None) => false,
_ => true,
}
}
fn record_threshold_log(
logs: &mut Vec<RateLimitWarningRecord>,
threshold: f64,
@@ -396,20 +385,52 @@ fn record_threshold_log(
if let Some(existing) = logs.iter_mut().find(|entry| {
(entry.threshold - threshold).abs() < f64::EPSILON
}) {
if reset_changed(existing.reset_at, reset_at) {
existing.logged_at = None;
} else if reset_at.is_none()
&& existing
.logged_at
.is_some_and(|logged| observed_at.signed_duration_since(logged) >= UNKNOWN_RESET_RELOG_INTERVAL)
{
let previous_reset = existing.reset_at;
let previous_logged = existing.logged_at;
let new_reset = reset_at;
let reset_moved_earlier = match (previous_reset, new_reset) {
(Some(prev), Some(next)) => next + RESET_PASSED_TOLERANCE < prev,
_ => false,
};
let logged_after_prev_reset = previous_reset
.and_then(|prev| {
previous_logged.map(|logged| {
logged + RESET_PASSED_TOLERANCE >= prev
})
})
.unwrap_or(false);
let prev_reset_elapsed = previous_reset
.map(|prev| observed_at + RESET_PASSED_TOLERANCE >= prev)
.unwrap_or(false);
let unknown_reset_elapsed = (previous_reset.is_none() || new_reset.is_none())
&& previous_logged
.is_some_and(|logged| observed_at.signed_duration_since(logged) >= UNKNOWN_RESET_RELOG_INTERVAL);
let mut should_clear = false;
if reset_moved_earlier {
should_clear = true;
} else if prev_reset_elapsed && !logged_after_prev_reset {
should_clear = true;
} else if unknown_reset_elapsed {
should_clear = true;
}
existing.reset_at = new_reset;
if should_clear {
existing.logged_at = None;
}
existing.reset_at = reset_at;
if existing.logged_at.is_none() {
existing.logged_at = Some(observed_at);
return true;
}
return false;
}
@@ -537,6 +558,25 @@ pub fn load_account_usage(
#[cfg(test)]
mod tests {
//! Regression coverage for rate-limit warning relogging.
//!
//! These cases enforce the desired behaviour:
//! - **No duplicate within a window**: once a threshold logs, subsequent polls before
//! the stored reset timestamp must remain silent even if the backend repeats or
//! extends the reset time.
//! - **Relog after reset passes**: the first poll at or after the recorded reset may
//! emit again, regardless of whether the backend has already advanced the window.
//! - **Relog on earlier reset**: if the backend moves the reset earlier (window
//! shrinks), we allow an immediate relog even before the previously stored reset.
//! - **Unknown reset fallback**: when reset metadata disappears, we rely on the
//! 24-hour `UNKNOWN_RESET_RELOG_INTERVAL` to unblock further warnings. When the
//! backend begins reporting timestamps again, we should also allow a relog provided
//! the fallback window has elapsed.
//! - **Missing metadata alone is not enough**: before the fallback timer elapses we
//! must keep warnings muted even if new snapshots omit reset times.
//!
//! The helper tests below construct scenarios targeting each rule so the state
//! machine in `record_threshold_log` can be refactored confidently.
use super::*;
use std::fs::File;
use crate::protocol::TokenUsage;
@@ -552,6 +592,198 @@ mod tests {
}
}
#[test]
fn rate_limit_warning_only_logs_once_per_reset() {
let home = TempDir::new().expect("tempdir");
let now = Utc::now();
let reset_at = now + Duration::days(7);
let first = record_rate_limit_warning(
home.path(),
"acct-1",
Some("Team"),
RateLimitWarningScope::Secondary,
75.0,
Some(reset_at),
now,
"Secondary usage exceeded 75% of the limit. Run /limits for detailed usage.",
)
.expect("first record succeeds");
assert!(first, "first logging should emit");
let second = record_rate_limit_warning(
home.path(),
"acct-1",
Some("Team"),
RateLimitWarningScope::Secondary,
75.0,
Some(now + Duration::days(7) + Duration::hours(6)),
now + Duration::hours(6),
"Secondary usage exceeded 75% of the limit. Run /limits for detailed usage.",
)
.expect("second record succeeds");
assert!(!second, "duplicate logging before reset should be suppressed");
let third = record_rate_limit_warning(
home.path(),
"acct-1",
Some("Team"),
RateLimitWarningScope::Secondary,
75.0,
Some(now + Duration::days(15)),
now + Duration::days(8),
"Secondary usage exceeded 75% of the limit. Run /limits for detailed usage.",
)
.expect("third record succeeds");
assert!(third, "after reset passes we should emit again");
}
#[test]
fn rate_limit_warning_relogs_after_reset_with_new_timestamp() {
let home = TempDir::new().expect("tempdir");
let now = Utc::now();
let msg = "Secondary usage exceeded 75% of the limit. Run /limits for detailed usage.";
let first = record_rate_limit_warning(
home.path(),
"acct-1",
Some("Team"),
RateLimitWarningScope::Secondary,
75.0,
Some(now + Duration::hours(1)),
now,
msg,
)
.expect("first record succeeds");
assert!(first);
// Backend extends reset window beyond the old reset time; we should re-emit now that
// the prior window has expired and a new one started.
let second = record_rate_limit_warning(
home.path(),
"acct-1",
Some("Team"),
RateLimitWarningScope::Secondary,
75.0,
Some(now + Duration::hours(2)),
now + Duration::minutes(65),
msg,
)
.expect("second record succeeds");
assert!(second, "after reset we should log again even if next window is later");
// Subsequent updates inside the new window should remain suppressed until that reset passes.
let third = record_rate_limit_warning(
home.path(),
"acct-1",
Some("Team"),
RateLimitWarningScope::Secondary,
75.0,
Some(now + Duration::hours(2)),
now + Duration::minutes(70),
msg,
)
.expect("third record succeeds");
assert!(!third, "duplicate logging inside the same window should stay muted");
}
#[test]
fn rate_limit_warning_relogs_after_unknown_reset_interval() {
let home = TempDir::new().expect("tempdir");
let now = Utc::now();
let msg = "Secondary usage exceeded 75% of the limit. Run /limits for detailed usage.";
let first = record_rate_limit_warning(
home.path(),
"acct-1",
Some("Team"),
RateLimitWarningScope::Secondary,
75.0,
Some(now + Duration::hours(1)),
now,
msg,
)
.expect("first record succeeds");
assert!(first);
// Backend stops providing reset info — still within backoff window.
let second = record_rate_limit_warning(
home.path(),
"acct-1",
Some("Team"),
RateLimitWarningScope::Secondary,
75.0,
None,
now + Duration::minutes(20),
msg,
)
.expect("second record succeeds");
assert!(!second, "dropping reset info should keep warning muted initially");
// After the unknown interval we should allow another log.
let third = record_rate_limit_warning(
home.path(),
"acct-1",
Some("Team"),
RateLimitWarningScope::Secondary,
75.0,
None,
now + Duration::hours(25),
msg,
)
.expect("third record succeeds");
assert!(third, "after backoff expires we should re-emit");
}
#[test]
fn rate_limit_warning_relogs_when_reset_info_returns() {
let home = TempDir::new().expect("tempdir");
let now = Utc::now();
let msg = "Secondary usage exceeded 75% of the limit. Run /limits for detailed usage.";
let first = record_rate_limit_warning(
home.path(),
"acct-1",
Some("Team"),
RateLimitWarningScope::Secondary,
75.0,
Some(now + Duration::hours(1)),
now,
msg,
)
.expect("first record succeeds");
assert!(first);
let second = record_rate_limit_warning(
home.path(),
"acct-1",
Some("Team"),
RateLimitWarningScope::Secondary,
75.0,
None,
now + Duration::minutes(20),
msg,
)
.expect("second record succeeds");
assert!(!second);
let third = record_rate_limit_warning(
home.path(),
"acct-1",
Some("Team"),
RateLimitWarningScope::Secondary,
75.0,
Some(now + Duration::hours(30)),
now + Duration::hours(25),
msg,
)
.expect("third record succeeds");
assert!(third, "restored reset metadata after fallback window should re-log");
}
#[test]
fn creates_usage_file_and_accumulates_tokens() {
let home = TempDir::new().expect("tempdir");
codex-rs/tui/src/chatwidget.rs
+54 -14
View File
@@ -338,10 +338,14 @@ impl RateLimitWarningState {
) -> Vec<RateLimitWarning> {
let mut warnings = Vec::new();
while self.weekly_index < RATE_LIMIT_WARNING_THRESHOLDS.len()
&& secondary_used_percent >= RATE_LIMIT_WARNING_THRESHOLDS[self.weekly_index]
let mut next_weekly_index = self.weekly_index;
while next_weekly_index < RATE_LIMIT_WARNING_THRESHOLDS.len()
&& secondary_used_percent >= RATE_LIMIT_WARNING_THRESHOLDS[next_weekly_index]
{
let threshold = RATE_LIMIT_WARNING_THRESHOLDS[self.weekly_index];
next_weekly_index += 1;
}
if next_weekly_index > self.weekly_index {
let threshold = RATE_LIMIT_WARNING_THRESHOLDS[next_weekly_index - 1];
warnings.push(RateLimitWarning {
scope: RateLimitWarningScope::Secondary,
threshold,
@@ -349,13 +353,17 @@ impl RateLimitWarningState {
"Secondary usage exceeded {threshold:.0}% of the limit. Run /limits for detailed usage."
),
});
self.weekly_index += 1;
self.weekly_index = next_weekly_index;
}
while self.hourly_index < RATE_LIMIT_WARNING_THRESHOLDS.len()
&& primary_used_percent >= RATE_LIMIT_WARNING_THRESHOLDS[self.hourly_index]
let mut next_hourly_index = self.hourly_index;
while next_hourly_index < RATE_LIMIT_WARNING_THRESHOLDS.len()
&& primary_used_percent >= RATE_LIMIT_WARNING_THRESHOLDS[next_hourly_index]
{
let threshold = RATE_LIMIT_WARNING_THRESHOLDS[self.hourly_index];
next_hourly_index += 1;
}
if next_hourly_index > self.hourly_index {
let threshold = RATE_LIMIT_WARNING_THRESHOLDS[next_hourly_index - 1];
warnings.push(RateLimitWarning {
scope: RateLimitWarningScope::Primary,
threshold,
@@ -363,7 +371,7 @@ impl RateLimitWarningState {
"Hourly usage exceeded {threshold:.0}% of the limit. Run /limits for detailed usage."
),
});
self.hourly_index += 1;
self.hourly_index = next_hourly_index;
}
warnings
@@ -7661,12 +7669,18 @@ impl ChatWidget<'_> {
fn update_rate_limit_resets(&mut self, current: &RateLimitSnapshotEvent) {
let now = Utc::now();
self.rate_limit_primary_next_reset_at = current
.primary_reset_after_seconds
.map(|secs| now + ChronoDuration::seconds(secs as i64));
self.rate_limit_secondary_next_reset_at = current
.secondary_reset_after_seconds
.map(|secs| now + ChronoDuration::seconds(secs as i64));
if let Some(secs) = current.primary_reset_after_seconds {
self.rate_limit_primary_next_reset_at =
Some(now + ChronoDuration::seconds(secs as i64));
} else {
self.rate_limit_primary_next_reset_at = None;
}
if let Some(secs) = current.secondary_reset_after_seconds {
self.rate_limit_secondary_next_reset_at =
Some(now + ChronoDuration::seconds(secs as i64));
} else {
self.rate_limit_secondary_next_reset_at = None;
}
}
pub(crate) fn handle_update_command(&mut self) {
@@ -15880,6 +15894,32 @@ mod tests {
);
}
#[test]
fn rate_limit_warnings_emit_only_highest_threshold() {
let mut state = RateLimitWarningState::default();
// Jump straight past multiple weekly thresholds; emit just the highest crossed one.
let weekly = state.take_warnings(80.0, 10.0);
assert_eq!(weekly.len(), 1);
assert_eq!(weekly[0].threshold, 75.0);
assert!(matches!(weekly[0].scope, RateLimitWarningScope::Secondary));
// Calling again without crossing new thresholds should produce nothing.
assert!(state.take_warnings(80.0, 10.0).is_empty());
// Crossing the final threshold should produce it exactly once.
let weekly_final = state.take_warnings(95.0, 10.0);
assert_eq!(weekly_final.len(), 1);
assert_eq!(weekly_final[0].threshold, 90.0);
assert!(matches!(weekly_final[0].scope, RateLimitWarningScope::Secondary));
// Primary scope follows the same rule.
let primary = state.take_warnings(95.0, 80.0);
assert_eq!(primary.len(), 1);
assert_eq!(primary[0].threshold, 75.0);
assert!(matches!(primary[0].scope, RateLimitWarningScope::Primary));
}
#[tokio::test(flavor = "current_thread")]
async fn spinner_stays_while_any_agent_running() {
let mut chat = make_widget();