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 d418ac20e..d3c941d1b 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 563d0aea1..f998cdbd8 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 {