From 19445b4630f21816f4314853298b58163c648629 Mon Sep 17 00:00:00 2001 From: Andrei Hasna Date: Thu, 13 Aug 2026 00:34:24 +0300 Subject: [PATCH] fix(core): classify auth-profile usage per rate-limit lane usage_profile_health selected a single snapshot with an exact match on the `codex` limit id and discarded the rest, so a provider that bills several lanes against one profile was judged entirely on its default lane. Two consequences, only one of which is closed here. Auto-switch was silently disabled for any other lane. A usage limit reached on `codex_bengalfox` failed `usage_limit_matches_auto_switch_config`, which returns false rather than falling back to something safer, so the switch never fired at all. The gate now classifies against the lane the limit actually came from, and per-lane health is kept lane-first so a profile healthy on one lane is never compared against a profile healthy on another. Cooldowns are scoped the same way: a profile capped on one lane stays a legitimate target for a lane it has not spent. Health reported to the brokers is deliberately unchanged. usage_health_by_lane and usage_lane_availability are added so a caller can ask which lanes remain, but they return their own type rather than UsageProfileHealth, because a free sibling lane does not make the profile usable for work bound to the spent one. Wiring the display and dispatch paths to that answer needs the selection and apply types to carry a model, which they do not, and is not attempted here. Lanes are never merged and never maxed against each other; scoring the best lane would turn a false Exhausted into a false Healthy, which routes work into a refusal instead of away from one. Agent: Augustus --- .../src/session/auth_profile_auto_switch.rs | 116 +++++-- codex-rs/core/src/usage_profile_health.rs | 290 +++++++++++++++++- 2 files changed, 382 insertions(+), 24 deletions(-) diff --git a/codex-rs/core/src/session/auth_profile_auto_switch.rs b/codex-rs/core/src/session/auth_profile_auto_switch.rs index d418ac20e2..d3c941d1b6 100644 --- a/codex-rs/core/src/session/auth_profile_auto_switch.rs +++ b/codex-rs/core/src/session/auth_profile_auto_switch.rs @@ -3,14 +3,16 @@ use std::collections::HashSet; use std::sync::Arc; use crate::session::session::Session; +use crate::usage_profile_health::DEFAULT_USAGE_LANE; use crate::usage_profile_health::UsageProfileAutoSwitchWindow; use crate::usage_profile_health::UsageProfileCooldownKey; use crate::usage_profile_health::UsageProfileHealth; use crate::usage_profile_health::UsageProfileRateLimitSnapshot; use crate::usage_profile_health::UsageProfileRateLimitWindow; use crate::usage_profile_health::choose_profile_for_auto_switch; -use crate::usage_profile_health::exhausted_auto_switch_window; -use crate::usage_profile_health::usage_health_for_snapshots; +use crate::usage_profile_health::exhausted_auto_switch_window_for_lane; +use crate::usage_profile_health::usage_health_for_lane; +use crate::usage_profile_health::usage_lane_id; use crate::usage_profile_health::usage_limit_matches_auto_switch_config; use codex_login::AuthProfile; use codex_login::AuthProfileSubscriptionProvider; @@ -21,7 +23,13 @@ use codex_protocol::protocol::RateLimitSnapshot; #[derive(Debug, Default)] pub(crate) struct AuthProfileAutoSwitchTurnState { attempted_profiles: HashSet>, - known_health_by_profile: BTreeMap, + /// Lane -> profile -> health. + /// + /// Health is only comparable within one lane: a profile healthy on the Spark lane and a + /// profile healthy on the default lane are not interchangeable switch targets. Keying + /// the lane outermost keeps every selection round scoped to a single lane, so the + /// per-profile map handed to the selector stays a like-for-like comparison. + known_health_by_profile: BTreeMap>, exhausted_profile_cooldowns: HashSet, } @@ -79,13 +87,18 @@ impl AuthProfileAutoSwitchTurnState { } }; let ordered = ordered_auth_profiles(&auto_switch_config.profiles, &profiles); + let lane = rate_limit_snapshot + .as_ref() + .map(usage_lane_id) + .unwrap_or(DEFAULT_USAGE_LANE); next_profile_from_known_health( &auto_switch_config, current_profile.as_deref(), &ordered, + lane, &self.attempted_profiles, &self.exhausted_profile_cooldowns, - &self.known_health_by_profile, + self.known_health_by_profile.get(lane), ) } @@ -98,14 +111,19 @@ impl AuthProfileAutoSwitchTurnState { let (Some(profile), Some(snapshot)) = (profile, snapshot) else { return; }; - let trigger_window = exhausted_auto_switch_window(snapshot, config); - let health = usage_health_for_snapshots( + let lane = usage_lane_id(snapshot); + let trigger_window = exhausted_auto_switch_window_for_lane(snapshot, config, lane); + let health = usage_health_for_lane( &[*snapshot], config, trigger_window.map(|window| window.label), /*is_fresh*/ true, + lane, ); - self.known_health_by_profile.insert(profile.clone(), health); + self.known_health_by_profile + .entry(lane.to_string()) + .or_default() + .insert(profile.clone(), health); if let Some(window) = trigger_window { self.exhausted_profile_cooldowns .insert(profile_cooldown_key(profile, snapshot, window)); @@ -137,11 +155,7 @@ fn profile_cooldown_key( snapshot: &UsageProfileRateLimitSnapshot<'_>, window: UsageProfileAutoSwitchWindow, ) -> UsageProfileCooldownKey { - UsageProfileCooldownKey::new( - Some(profile.to_string()), - snapshot.limit_id.unwrap_or("codex"), - window, - ) + UsageProfileCooldownKey::new(Some(profile.to_string()), usage_lane_id(snapshot), window) } fn ordered_auth_profiles( @@ -187,28 +201,34 @@ fn next_profile_from_known_health( config: &crate::config::AuthProfileAutoSwitchConfig, current_profile: Option<&str>, ordered: &[String], + lane: &str, attempted_profiles: &HashSet>, exhausted_profile_cooldowns: &HashSet, - known_health_by_profile: &BTreeMap, + known_health_by_profile: Option<&BTreeMap>, ) -> Option { let candidates = auth_profile_candidates( current_profile, ordered, + lane, attempted_profiles, exhausted_profile_cooldowns, ); - choose_profile_for_auto_switch(config, &candidates, known_health_by_profile).selected_profile + let known_health_by_profile = known_health_by_profile.cloned().unwrap_or_default(); + choose_profile_for_auto_switch(config, &candidates, &known_health_by_profile).selected_profile } fn auth_profile_candidates( current_profile: Option<&str>, ordered: &[String], + lane: &str, attempted_profiles: &HashSet>, exhausted_profile_cooldowns: &HashSet, ) -> Vec { next_untried_profiles(current_profile, ordered, attempted_profiles) .into_iter() - .filter(|profile| !profile_has_exhausted_cooldown(profile, exhausted_profile_cooldowns)) + .filter(|profile| { + !profile_has_exhausted_cooldown(profile, lane, exhausted_profile_cooldowns) + }) .collect() } @@ -238,13 +258,16 @@ fn next_untried_profiles( .collect() } +/// A cooldown only excludes the profile from the lane it was recorded against. A profile +/// capped on one lane is still a legitimate switch target for a lane it has not spent. fn profile_has_exhausted_cooldown( profile: &str, + lane: &str, exhausted_profile_cooldowns: &HashSet, ) -> bool { - exhausted_profile_cooldowns - .iter() - .any(|cooldown| cooldown.profile.as_deref() == Some(profile)) + exhausted_profile_cooldowns.iter().any(|cooldown| { + cooldown.profile.as_deref() == Some(profile) && cooldown.limit_id.eq_ignore_ascii_case(lane) + }) } #[cfg(test)] @@ -372,6 +395,60 @@ mod tests { usage_limit_matches_auto_switch_config(config, snapshot) } + fn spark_snapshot(window_minutes: i64) -> RateLimitSnapshot { + RateLimitSnapshot { + limit_id: Some("codex_bengalfox".to_string()), + ..codex_snapshot(window_minutes) + } + } + + #[test] + fn profile_health_is_recorded_under_the_lane_the_limit_came_from() { + // A limit reached on the Spark lane says nothing about the default lane, and must + // not be filed against it: the two lanes have independent quotas. + let mut state = AuthProfileAutoSwitchTurnState::default(); + let config = AuthProfileAutoSwitchConfig { + enabled: true, + ..Default::default() + }; + let spark = spark_snapshot(7 * 24 * 60); + let snapshot = core_rate_limit_snapshot(&spark); + + state.record_profile_health(Some(&"account003".to_string()), Some(&snapshot), &config); + + assert_eq!( + Some(&UsageProfileHealth::Exhausted { + retry_at: Some(123) + }), + state + .known_health_by_profile + .get("codex_bengalfox") + .and_then(|by_profile| by_profile.get("account003")) + ); + assert_eq!(None, state.known_health_by_profile.get("codex")); + } + + #[test] + fn cooldown_on_one_lane_does_not_exclude_the_profile_from_another_lane() { + let cooldowns = HashSet::from([UsageProfileCooldownKey { + profile: Some("account002".to_string()), + limit_id: "codex_bengalfox".to_string(), + window_label: FIVE_HOUR_LIMIT_LABEL.to_string(), + resets_at: Some(123), + }]); + + assert!(!profile_has_exhausted_cooldown( + "account002", + "codex", + &cooldowns + )); + assert!(profile_has_exhausted_cooldown( + "account002", + "codex_bengalfox", + &cooldowns + )); + } + #[test] fn highest_available_skips_known_exhausted_profile_and_prefers_known_healthy_profile() { let mut config = AuthProfileAutoSwitchConfig { @@ -408,9 +485,10 @@ mod tests { &config, Some("account001"), &ordered, + DEFAULT_USAGE_LANE, &attempted, &exhausted_profile_cooldowns, - &known_health_by_profile, + Some(&known_health_by_profile), ) ); } diff --git a/codex-rs/core/src/usage_profile_health.rs b/codex-rs/core/src/usage_profile_health.rs index 563d0aea14..f998cdbd82 100644 --- a/codex-rs/core/src/usage_profile_health.rs +++ b/codex-rs/core/src/usage_profile_health.rs @@ -4,6 +4,14 @@ use std::time::Duration; use crate::config::AuthProfileAutoSwitchConfig; use crate::config::AuthProfileAutoSwitchStrategy; +/// The rate-limit lane every caller means when it does not say otherwise. +/// +/// A provider bills several lanes against one auth profile and returns one snapshot per +/// lane (`codex`, `codex_bengalfox`, ...). Lanes are matched exactly, never by prefix: +/// `codex_model` is a different product, not a codex lane, and matching it by prefix +/// would let whichever lane happened to sort first decide a profile's health. +pub const DEFAULT_USAGE_LANE: &str = "codex"; + pub const PRIMARY_LIMIT_FALLBACK_LABEL: &str = "usage"; pub const SECONDARY_LIMIT_FALLBACK_LABEL: &str = "secondary usage"; pub const FIVE_HOUR_LIMIT_LABEL: &str = "5h"; @@ -53,6 +61,29 @@ pub enum UsageProfileHealth { Unknown, } +/// One lane's health, carrying the lane it was computed for. +#[derive(Clone, Debug, PartialEq)] +pub struct UsageProfileLaneHealth { + pub lane: String, + pub health: UsageProfileHealth, +} + +/// Whether an auth profile has any usable lane left, and which. +/// +/// Deliberately a separate type from [`UsageProfileHealth`]: it must not be substitutable +/// for one. `Usable` says a lane is free, not that the profile will serve the model the +/// caller is about to use — routing on it as if it were `Healthy` sends default-model work +/// into a lane that is measurably refusing requests. +#[derive(Clone, Debug, PartialEq)] +pub enum UsageProfileLaneAvailability { + /// At least one lane is healthy. `usable_lanes` names them, in snapshot order. + Usable { usable_lanes: Vec }, + /// Every classified lane is exhausted; `retry_at` is the earliest known reset. + Exhausted { retry_at: Option }, + /// No lane could be classified, or the read was stale. + Unknown, +} + #[derive(Clone, Debug, PartialEq, Eq)] pub struct UsageProfileSelection { pub selected_profile: Option, @@ -120,14 +151,33 @@ pub fn usage_limit_matches_auto_switch_config( let Some(snapshot) = snapshot else { return true; }; - exhausted_auto_switch_window(snapshot, config).is_some() + // Classify against the lane the limit actually came from. Hardcoding the default lane + // here made a usage limit on any other lane fail to match, which does not fall back to + // some safer behaviour: it returns `false` and silently disables auto-switch entirely. + exhausted_auto_switch_window_for_lane(snapshot, config, usage_lane_id(snapshot)).is_some() +} + +/// The lane a snapshot describes: its `limit_id`, else its `limit_name`, else the default. +pub fn usage_lane_id<'a>(snapshot: &UsageProfileRateLimitSnapshot<'a>) -> &'a str { + snapshot + .limit_id + .or(snapshot.limit_name) + .unwrap_or(DEFAULT_USAGE_LANE) } pub fn exhausted_auto_switch_window( snapshot: &UsageProfileRateLimitSnapshot<'_>, config: &AuthProfileAutoSwitchConfig, ) -> Option { - if !config.enabled || !is_codex_limit(snapshot) { + exhausted_auto_switch_window_for_lane(snapshot, config, DEFAULT_USAGE_LANE) +} + +pub fn exhausted_auto_switch_window_for_lane( + snapshot: &UsageProfileRateLimitSnapshot<'_>, + config: &AuthProfileAutoSwitchConfig, + lane: &str, +) -> Option { + if !config.enabled || !snapshot_matches_lane(snapshot, lane) { return None; } @@ -182,12 +232,37 @@ pub fn usage_health_for_snapshots( config: &AuthProfileAutoSwitchConfig, trigger_window_label: Option<&str>, is_fresh: bool, +) -> UsageProfileHealth { + usage_health_for_lane( + snapshots, + config, + trigger_window_label, + is_fresh, + DEFAULT_USAGE_LANE, + ) +} + +/// Health of one named `lane` across `snapshots`. +/// +/// Lanes are never merged and never maxed against each other. A profile whose default lane +/// is spent is genuinely unusable for work on that lane, however much headroom a sibling +/// lane has; taking the best lane's score would convert a false `Exhausted` into a false +/// `Healthy`, which is the more expensive direction because it routes work into a refusal. +pub fn usage_health_for_lane( + snapshots: &[UsageProfileRateLimitSnapshot<'_>], + config: &AuthProfileAutoSwitchConfig, + trigger_window_label: Option<&str>, + is_fresh: bool, + lane: &str, ) -> UsageProfileHealth { if !is_fresh { return UsageProfileHealth::Unknown; } - let Some(snapshot) = snapshots.iter().find(|snapshot| is_codex_limit(snapshot)) else { + let Some(snapshot) = snapshots + .iter() + .find(|snapshot| snapshot_matches_lane(snapshot, lane)) + else { return UsageProfileHealth::Unknown; }; @@ -238,6 +313,69 @@ pub fn usage_health_for_snapshots( }) } +/// Classify every lane in `snapshots` independently, in snapshot order. +/// +/// The single-lane query answers "is this profile usable on the default lane". This answers +/// "which lanes is it usable on at all" — a strictly wider question, and the only one that +/// can distinguish a profile with nothing left from a profile whose default lane is spent +/// while a sibling lane is free. +pub fn usage_health_by_lane( + snapshots: &[UsageProfileRateLimitSnapshot<'_>], + config: &AuthProfileAutoSwitchConfig, + trigger_window_label: Option<&str>, + is_fresh: bool, +) -> Vec { + snapshots + .iter() + .map(|snapshot| { + let lane = usage_lane_id(snapshot); + UsageProfileLaneHealth { + lane: lane.to_string(), + health: usage_health_for_lane( + std::slice::from_ref(snapshot), + config, + trigger_window_label, + is_fresh, + lane, + ), + } + }) + .collect() +} + +/// Reduce per-lane health to whether the profile has anything usable left. +/// +/// `Exhausted` requires that every classified lane is exhausted. A lane that could not be +/// classified leaves the answer `Unknown` rather than `Exhausted`, so an incomplete read is +/// never published as a profile with nothing left. +pub fn usage_lane_availability(lanes: &[UsageProfileLaneHealth]) -> UsageProfileLaneAvailability { + let usable_lanes = lanes + .iter() + .filter(|lane| matches!(lane.health, UsageProfileHealth::Healthy(_))) + .map(|lane| lane.lane.clone()) + .collect::>(); + if !usable_lanes.is_empty() { + return UsageProfileLaneAvailability::Usable { usable_lanes }; + } + + let mut retry_at = None; + let mut every_lane_exhausted = !lanes.is_empty(); + for lane in lanes { + match lane.health { + UsageProfileHealth::Exhausted { + retry_at: lane_retry_at, + } => merge_retry_at(&mut retry_at, lane_retry_at), + _ => every_lane_exhausted = false, + } + } + + if every_lane_exhausted { + UsageProfileLaneAvailability::Exhausted { retry_at } + } else { + UsageProfileLaneAvailability::Unknown + } +} + pub fn choose_profile_for_auto_switch( config: &AuthProfileAutoSwitchConfig, candidates: &[String], @@ -428,13 +566,19 @@ fn exhausted_auto_switch_window_for_limit( } fn is_codex_limit(snapshot: &UsageProfileRateLimitSnapshot<'_>) -> bool { + snapshot_matches_lane(snapshot, DEFAULT_USAGE_LANE) +} + +/// Exact, case-insensitive lane match. `limit_id` is authoritative when present: a snapshot +/// labelled `codex_model` is not the `codex` lane even when its display name says otherwise. +fn snapshot_matches_lane(snapshot: &UsageProfileRateLimitSnapshot<'_>, lane: &str) -> bool { if let Some(limit_id) = snapshot.limit_id { - return limit_id.eq_ignore_ascii_case("codex"); + return limit_id.eq_ignore_ascii_case(lane); } snapshot .limit_name - .is_none_or(|limit_name| limit_name.eq_ignore_ascii_case("codex")) + .is_none_or(|limit_name| limit_name.eq_ignore_ascii_case(lane)) } fn usage_profile_score_is_better( @@ -916,6 +1060,142 @@ mod tests { ); } + /// A snapshot for `lane` whose primary weekly window sits at `used_percent`. + fn lane_snapshot( + lane: &'static str, + used_percent: f64, + ) -> UsageProfileRateLimitSnapshot<'static> { + UsageProfileRateLimitSnapshot { + limit_id: Some(lane), + limit_name: None, + primary: Some(window(used_percent, MINUTES_PER_WEEK, Some(900))), + secondary: None, + } + } + + /// A provider that bills two lanes: `codex` spent, `codex_bengalfox` untouched. + fn spent_codex_free_spark() -> Vec> { + vec![ + lane_snapshot("codex", /*used_percent*/ 100.0), + lane_snapshot("codex_bengalfox", /*used_percent*/ 0.0), + ] + } + + #[test] + fn usage_limit_matches_auto_switch_config_for_a_non_default_lane() { + // A usage limit reached on the Spark lane is still a usage limit: refusing to + // recognise it silently disables auto-switch instead of moving to another profile. + let snapshot = lane_snapshot("codex_bengalfox", /*used_percent*/ 100.0); + + assert!(usage_limit_matches_auto_switch_config( + &config(), + Some(&snapshot) + )); + } + + #[test] + fn usage_health_for_snapshots_scores_only_the_default_lane() { + // Guard against over-correcting: the default query must keep answering for the + // `codex` lane alone. Reporting Healthy here would route default-model work into + // a lane that is measurably refusing requests. + assert_eq!( + UsageProfileHealth::Exhausted { + retry_at: Some(900) + }, + usage_health_for_snapshots( + &spent_codex_free_spark(), + &config(), + Some(WEEKLY_LIMIT_LABEL), + /*is_fresh*/ true, + ) + ); + } + + #[test] + fn usage_health_by_lane_names_the_lane_that_is_still_usable() { + let lanes = usage_health_by_lane( + &spent_codex_free_spark(), + &config(), + Some(WEEKLY_LIMIT_LABEL), + /*is_fresh*/ true, + ); + + assert_eq!( + UsageProfileLaneAvailability::Usable { + usable_lanes: vec!["codex_bengalfox".to_string()], + }, + usage_lane_availability(&lanes) + ); + } + + #[test] + fn usage_lane_availability_reports_exhausted_only_when_every_lane_is_exhausted() { + // The negative control. It cannot be run against a live profile on this fleet + // because no account currently has every lane at 100%, so it lives as a fixture. + let lanes = usage_health_by_lane( + &[ + lane_snapshot("codex", /*used_percent*/ 100.0), + lane_snapshot("codex_bengalfox", /*used_percent*/ 100.0), + ], + &config(), + Some(WEEKLY_LIMIT_LABEL), + /*is_fresh*/ true, + ); + + assert_eq!( + UsageProfileLaneAvailability::Exhausted { + retry_at: Some(900) + }, + usage_lane_availability(&lanes) + ); + } + + #[test] + fn usage_health_by_lane_leaves_a_single_lane_profile_unchanged() { + // Accounts that only ever return one lane must classify exactly as before. + let snapshots = [lane_snapshot("codex", /*used_percent*/ 40.0)]; + let lanes = usage_health_by_lane( + &snapshots, + &config(), + Some(WEEKLY_LIMIT_LABEL), + /*is_fresh*/ true, + ); + + assert_eq!( + vec![UsageProfileLaneHealth { + lane: "codex".to_string(), + health: usage_health_for_snapshots( + &snapshots, + &config(), + Some(WEEKLY_LIMIT_LABEL), + /*is_fresh*/ true, + ), + }], + lanes + ); + assert_eq!( + UsageProfileLaneAvailability::Usable { + usable_lanes: vec!["codex".to_string()], + }, + usage_lane_availability(&lanes) + ); + } + + #[test] + fn usage_health_by_lane_reports_unknown_for_a_stale_read() { + let lanes = usage_health_by_lane( + &spent_codex_free_spark(), + &config(), + Some(WEEKLY_LIMIT_LABEL), + /*is_fresh*/ false, + ); + + assert_eq!( + UsageProfileLaneAvailability::Unknown, + usage_lane_availability(&lanes) + ); + } + #[test] fn earliest_exhausted_reset_at_ignores_non_codex_limits() { let snapshot = UsageProfileRateLimitSnapshot {