fix(core): classify auth-profile usage per rate-limit lane

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<Option<String>>,
-    known_health_by_profile: BTreeMap<String, UsageProfileHealth>,
+    /// 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<String, BTreeMap<String, UsageProfileHealth>>,
     exhausted_profile_cooldowns: HashSet<UsageProfileCooldownKey>,
 }
 
@@ -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<Option<String>>,
     exhausted_profile_cooldowns: &HashSet<UsageProfileCooldownKey>,
-    known_health_by_profile: &BTreeMap<String, UsageProfileHealth>,
+    known_health_by_profile: Option<&BTreeMap<String, UsageProfileHealth>>,
 ) -> Option<String> {
     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<Option<String>>,
     exhausted_profile_cooldowns: &HashSet<UsageProfileCooldownKey>,
 ) -> Vec<String> {
     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<UsageProfileCooldownKey>,
 ) -> 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),
             )
         );
     }