diff --git a/codex-rs/core/src/auth_profile_usage.rs b/codex-rs/core/src/auth_profile_usage.rs index d154ba358..ff4542880 100644 --- a/codex-rs/core/src/auth_profile_usage.rs +++ b/codex-rs/core/src/auth_profile_usage.rs @@ -254,10 +254,9 @@ pub fn recommend_auth_profile( }; } Some(AuthProfileUsageHealth::Unknown) => { - return AuthProfileUsageRecommendation { - profile: (*candidate).clone(), - reason: AuthProfileUsageRecommendationReason::SelectedUnknownFallback, - }; + if first_unknown.is_none() { + first_unknown = Some((*candidate).clone()); + } } Some(AuthProfileUsageHealth::Exhausted { .. }) | None => {} } @@ -676,6 +675,41 @@ mod tests { ); } + #[test] + fn ordered_recommendation_prefers_healthy_over_earlier_unknown() { + let health = vec![ + ( + Some("work".to_string()), + AuthProfileUsageHealth::Exhausted { retry_at: Some(1) }, + ), + (Some("second".to_string()), AuthProfileUsageHealth::Unknown), + ( + Some("third".to_string()), + AuthProfileUsageHealth::Healthy { + remaining_percent: 20.0, + resets_at: None, + }, + ), + ]; + + assert_eq!( + AuthProfileUsageRecommendation { + profile: Some("third".to_string()), + reason: AuthProfileUsageRecommendationReason::SelectedOrderedHealthy, + }, + recommend_auth_profile( + Some("work"), + AuthProfileAutoSwitchStrategy::Ordered, + &[ + Some("work".to_string()), + Some("second".to_string()), + Some("third".to_string()), + ], + &health, + ) + ); + } + #[test] fn recommendation_keeps_current_profile_when_current_usage_is_unknown() { let health = vec![ diff --git a/codex-rs/tui/src/chatwidget/tests/popups_and_settings.rs b/codex-rs/tui/src/chatwidget/tests/popups_and_settings.rs index 5632c4290..3daab5593 100644 --- a/codex-rs/tui/src/chatwidget/tests/popups_and_settings.rs +++ b/codex-rs/tui/src/chatwidget/tests/popups_and_settings.rs @@ -3457,6 +3457,185 @@ async fn usage_heartbeat_suppressed_for_exhausted_profile_until_reset() { ); } +#[tokio::test] +async fn same_name_relogin_clears_exhausted_profile_heartbeat_suppression() { + let (mut chat, _rx, _op_rx) = make_chatwidget_manual(Some("gpt-5.2")).await; + chat.thread_id = Some(ThreadId::new()); + chat.config.selected_auth_profile = Some("work".to_string()); + save_popup_chatgpt_auth_profile(&chat, "work", "work@example.com"); + + let stale_exhausted = rate_limit_snapshot_display_for_limit( + &profile_usage_snapshot( + /*secondary_used_percent*/ 100, /*primary_used_percent*/ 100, + ), + "codex".to_string(), + Local::now() - chrono::Duration::minutes(RATE_LIMIT_STALE_THRESHOLD_MINUTES + 1), + ); + chat.auth_profile_rate_limit_snapshots_by_profile.insert( + Some("work".to_string()), + BTreeMap::from([("codex".to_string(), stale_exhausted)]), + ); + chat.auth_profile_usage_exhausted_reset_at_by_profile + .insert( + Some("work".to_string()), + chrono::Utc::now().timestamp() + 3600, + ); + assert!(chat.auth_profile_usage_refresh_targets().is_empty()); + + chat.begin_selected_auth_profile_credential_mutation("work"); + chat.finish_selected_auth_profile_credential_mutation( + "work", /*credentials_changed*/ true, + ); + + assert!( + chat.auth_profile_usage_exhausted_reset_at_by_profile + .is_empty() + ); + assert!(chat.auth_profile_rate_limit_snapshots_by_profile.is_empty()); + assert_eq!( + chat.auth_profile_usage_refresh_targets(), + vec![RateLimitRefreshTarget::Named("work".to_string())] + ); +} + +#[tokio::test] +async fn relogin_invalidates_captured_profile_after_selection_changes() { + let (mut chat, _rx, _op_rx) = make_chatwidget_manual(Some("gpt-5.2")).await; + chat.thread_id = Some(ThreadId::new()); + chat.config.selected_auth_profile = Some("work".to_string()); + save_popup_chatgpt_auth_profile(&chat, "work", "work@example.com"); + save_popup_chatgpt_auth_profile(&chat, "personal", "personal@example.com"); + + let stale_exhausted = rate_limit_snapshot_display_for_limit( + &profile_usage_snapshot( + /*secondary_used_percent*/ 100, /*primary_used_percent*/ 100, + ), + "codex".to_string(), + Local::now() - chrono::Duration::minutes(RATE_LIMIT_STALE_THRESHOLD_MINUTES + 1), + ); + chat.auth_profile_rate_limit_snapshots_by_profile.insert( + Some("work".to_string()), + BTreeMap::from([("codex".to_string(), stale_exhausted)]), + ); + chat.auth_profile_usage_exhausted_reset_at_by_profile + .insert( + Some("work".to_string()), + chrono::Utc::now().timestamp() + 3600, + ); + + chat.begin_selected_auth_profile_credential_mutation("work"); + chat.config.selected_auth_profile = Some("personal".to_string()); + chat.finish_selected_auth_profile_credential_mutation( + "work", /*credentials_changed*/ true, + ); + + assert!( + !chat + .auth_profile_rate_limit_snapshots_by_profile + .contains_key(&Some("work".to_string())) + ); + assert!( + !chat + .auth_profile_usage_exhausted_reset_at_by_profile + .contains_key(&Some("work".to_string())) + ); + + chat.config.selected_auth_profile = Some("work".to_string()); + assert!( + chat.auth_profile_usage_refresh_targets() + .contains(&RateLimitRefreshTarget::Named("work".to_string())) + ); +} + +#[tokio::test] +async fn profile_relogin_invalidation_preserves_unrelated_profile_usage_state() { + let (mut chat, _rx, _op_rx) = make_chatwidget_manual(Some("gpt-5.2")).await; + chat.thread_id = Some(ThreadId::new()); + chat.config.selected_auth_profile = Some("work".to_string()); + save_popup_chatgpt_auth_profile(&chat, "work", "work@example.com"); + save_popup_chatgpt_auth_profile(&chat, "personal", "personal@example.com"); + + for profile in ["work", "personal"] { + let snapshot = rate_limit_snapshot_display_for_limit( + &profile_usage_snapshot( + /*secondary_used_percent*/ 80, /*primary_used_percent*/ 60, + ), + "codex".to_string(), + Local::now(), + ); + let profile = Some(profile.to_string()); + chat.auth_profile_rate_limit_snapshots_by_profile.insert( + profile.clone(), + BTreeMap::from([("codex".to_string(), snapshot)]), + ); + chat.auth_profile_usage_heartbeat_requested_at_by_profile + .insert(profile.clone(), std::time::Instant::now()); + chat.auth_profile_usage_heartbeat_failed_at_by_profile + .insert(profile.clone(), std::time::Instant::now()); + chat.auth_profile_usage_exhausted_reset_at_by_profile + .insert(profile, chrono::Utc::now().timestamp() + 3600); + } + chat.rate_limit_snapshots_by_limit_id.insert( + "codex".to_string(), + rate_limit_snapshot_display_for_limit( + &profile_usage_snapshot( + /*secondary_used_percent*/ 20, /*primary_used_percent*/ 10, + ), + "codex".to_string(), + Local::now(), + ), + ); + + chat.begin_selected_auth_profile_credential_mutation("work"); + chat.config.selected_auth_profile = Some("personal".to_string()); + chat.finish_selected_auth_profile_credential_mutation( + "work", /*credentials_changed*/ true, + ); + + let work = Some("work".to_string()); + let personal = Some("personal".to_string()); + assert!( + !chat + .auth_profile_rate_limit_snapshots_by_profile + .contains_key(&work) + ); + assert!( + chat.auth_profile_rate_limit_snapshots_by_profile + .contains_key(&personal) + ); + assert!( + !chat + .auth_profile_usage_heartbeat_requested_at_by_profile + .contains_key(&work) + ); + assert!( + chat.auth_profile_usage_heartbeat_requested_at_by_profile + .contains_key(&personal) + ); + assert!( + !chat + .auth_profile_usage_heartbeat_failed_at_by_profile + .contains_key(&work) + ); + assert!( + chat.auth_profile_usage_heartbeat_failed_at_by_profile + .contains_key(&personal) + ); + assert!( + !chat + .auth_profile_usage_exhausted_reset_at_by_profile + .contains_key(&work) + ); + assert!( + chat.auth_profile_usage_exhausted_reset_at_by_profile + .contains_key(&personal) + ); + assert!( + !chat.rate_limit_snapshots_by_limit_id.is_empty(), + "relogin for the previously selected profile must preserve the current profile snapshot" + ); +} + fn assert_no_rate_limit_refresh_event(rx: &mut tokio::sync::mpsc::UnboundedReceiver) { while let Ok(event) = rx.try_recv() { assert!( diff --git a/codex-rs/tui/src/chatwidget/usage_limit_reset/state.rs b/codex-rs/tui/src/chatwidget/usage_limit_reset/state.rs index 4378cdb0a..9de581182 100644 --- a/codex-rs/tui/src/chatwidget/usage_limit_reset/state.rs +++ b/codex-rs/tui/src/chatwidget/usage_limit_reset/state.rs @@ -1,5 +1,10 @@ use super::*; +enum AuthProfileUsageInvalidation<'a> { + AllProfiles, + Profile(&'a str), +} + impl ChatWidget { pub(crate) fn advance_rate_limit_reset_generation(&mut self) -> u64 { let Some(next_generation) = self.rate_limit_reset_generation.checked_add(1) else { @@ -63,8 +68,8 @@ impl ChatWidget { self.auth_profile_credential_mutations_in_flight .remove(profile); } - if credentials_changed && self.config.selected_auth_profile.as_deref() == Some(profile) { - self.invalidate_rate_limit_reset_state_after_account_update(); + if credentials_changed { + self.invalidate_rate_limit_reset_state(AuthProfileUsageInvalidation::Profile(profile)); } } @@ -205,6 +210,45 @@ impl ChatWidget { } pub(crate) fn invalidate_rate_limit_reset_state_after_account_update(&mut self) { + self.invalidate_rate_limit_reset_state(AuthProfileUsageInvalidation::AllProfiles); + } + + fn invalidate_rate_limit_reset_state(&mut self, scope: AuthProfileUsageInvalidation<'_>) { + // Per-profile usage caches always go, for every profile in scope: their + // credentials just changed, so any cached usage snapshot, heartbeat + // bookkeeping, or exhausted-reset suppression for them is stale. + let scope_covers_selected_profile = match scope { + AuthProfileUsageInvalidation::AllProfiles => { + self.auth_profile_rate_limit_snapshots_by_profile.clear(); + self.auth_profile_usage_heartbeat_requested_at_by_profile + .clear(); + self.auth_profile_usage_heartbeat_failed_at_by_profile + .clear(); + self.auth_profile_usage_exhausted_reset_at_by_profile + .clear(); + true + } + AuthProfileUsageInvalidation::Profile(profile) => { + let covers_selected_profile = + self.config.selected_auth_profile.as_deref() == Some(profile); + let profile = Some(profile.to_string()); + self.auth_profile_rate_limit_snapshots_by_profile + .remove(&profile); + self.auth_profile_usage_heartbeat_requested_at_by_profile + .remove(&profile); + self.auth_profile_usage_heartbeat_failed_at_by_profile + .remove(&profile); + self.auth_profile_usage_exhausted_reset_at_by_profile + .remove(&profile); + covers_selected_profile + } + }; + if !scope_covers_selected_profile { + // A credential mutation for a profile that is not selected must not + // disturb the selected profile's live snapshot or cancel its + // in-flight usage-limit reset state. + return; + } let automatic_reset_owned_failed_turn = self.automatic_usage_limit_reset_owns_failed_turn(); self.advance_rate_limit_reset_generation(); self.rate_limit_reset_credits = None; @@ -219,6 +263,7 @@ impl ChatWidget { self.pending_post_reset_refresh = None; self.automatic_reset_opted_out_generation = None; self.usage_limit_auto_reset_key = None; + self.rate_limit_snapshots_by_limit_id.clear(); self.auth_profile_auto_switch_snapshots_by_limit_id.clear(); self.prepare_for_usage_limit_reset(); if automatic_reset_owned_failed_turn {