agent_core/core/auth/
auto_selection.rs

1//! Automatic subscription choice and its secret-free explanation.
2//! Eligibility remains in `quota_policy`; config and the preview use that
3//! same policy so the CLI cannot promise a different account than the broker.
4use super::quota_policy::{
5    self as quota, AccountCapacity, QuotaObservation, RejectReason, Selection, SelectionRequest,
6    Strategy,
7};
8use super::{CredentialRef, OAuthProviderId};
9use serde::{Deserialize, Serialize};
10use std::collections::BTreeMap;
11
12#[derive(Clone, Debug, PartialEq, Eq)]
13pub struct AutoSelectionPolicy {
14    pub strategy: Strategy,
15    pub overrides: BTreeMap<OAuthProviderId, Strategy>,
16    pub urgent_horizon_ms: u64,
17    pub sticky: bool,
18}
19impl Default for AutoSelectionPolicy {
20    fn default() -> Self {
21        Self {
22            strategy: Strategy::SoonestReset,
23            overrides: BTreeMap::new(),
24            urgent_horizon_ms: quota::DEFAULT_URGENT_HORIZON_MS,
25            sticky: true,
26        }
27    }
28}
29impl AutoSelectionPolicy {
30    pub fn strategy_for(&self, provider: OAuthProviderId) -> Strategy {
31        self.overrides
32            .get(&provider)
33            .copied()
34            .unwrap_or(self.strategy)
35    }
36    /// Invalid values warn and retain the documented default, never an
37    /// unvalidated strategy. Keys are relative to `auth.auto.`.
38    pub fn from_config_map(map: &BTreeMap<String, String>) -> (Self, Vec<String>) {
39        let mut policy = Self::default();
40        let mut warnings = Vec::new();
41        for (key, value) in map {
42            let strategy = || match value.trim() {
43                "soonest_reset" => Some(Strategy::SoonestReset),
44                "lowest_utilization" => Some(Strategy::LowestUtilization),
45                "preference_order" => Some(Strategy::PreferenceOrder),
46                _ => None,
47            };
48            let valid = match key.as_str() {
49                "strategy" => strategy().map(|s| policy.strategy = s).is_some(),
50                "sticky" => value
51                    .trim()
52                    .parse::<bool>()
53                    .map(|v| policy.sticky = v)
54                    .is_ok(),
55                "urgent_horizon_hours" => value
56                    .trim()
57                    .parse::<u64>()
58                    .ok()
59                    .filter(|v| (1..=168).contains(v))
60                    .map(|v| policy.urgent_horizon_ms = v * 3_600_000)
61                    .is_some(),
62                _ => {
63                    if let Some(provider) = key
64                        .strip_prefix("strategy.")
65                        .and_then(|p| p.parse::<OAuthProviderId>().ok())
66                    {
67                        // A malformed explicit override must not inherit an
68                        // unrelated global strategy silently.
69                        policy
70                            .overrides
71                            .insert(provider, strategy().unwrap_or(Strategy::SoonestReset));
72                        strategy().is_some()
73                    } else {
74                        false
75                    }
76                }
77            };
78            if !valid {
79                warnings.push(format!("invalid auth.auto.{key}; using documented default"));
80            }
81        }
82        (policy, warnings)
83    }
84}
85
86#[derive(Clone, Debug, Serialize, Deserialize)]
87pub struct PlanRow {
88    pub credential: CredentialRef,
89    pub rank: Option<u32>,
90    pub tier: Option<u8>,
91    pub anchored: Option<bool>,
92    pub budget_reset_ms: Option<u64>,
93    pub utilization: Option<f64>,
94    pub verdict: String,
95    /// The failed observation behind the verdict, when there is one.
96    #[serde(default, skip_serializing_if = "Option::is_none")]
97    pub cause: Option<quota::ObservationFailure>,
98}
99#[derive(Clone, Debug, Serialize, Deserialize)]
100pub struct SelectionPlan {
101    pub provider: String,
102    pub model: Option<String>,
103    pub strategy: Strategy,
104    pub urgent_horizon_ms: u64,
105    pub sticky: bool,
106    pub current: Option<CredentialRef>,
107    pub now_ms: u64,
108    pub rows: Vec<PlanRow>,
109    pub selection: Selection,
110}
111impl SelectionPlan {
112    pub fn from_candidates(req: &SelectionRequest<'_>, candidates: &[AccountCapacity]) -> Self {
113        let selection = quota::select(req, candidates);
114        let rejections = match &selection {
115            Selection::Selected { rejections, .. } | Selection::NoCapacity { rejections, .. } => {
116                rejections
117            }
118        };
119        let mut rows: Vec<_> = candidates
120            .iter()
121            .filter(|c| c.credential.provider == req.provider)
122            .map(|cap| {
123                let eligible = quota::evaluate(req, cap).ok();
124                let rejection = rejections.iter().find(|r| r.credential == cap.credential);
125                let reason = rejection.map(|r| &r.reason);
126                let selected = selection.selected() == Some(&cap.credential);
127                let rank = if selected {
128                    Some(1)
129                } else if let Some(RejectReason::Outranked { rank }) = reason {
130                    Some(*rank)
131                } else {
132                    None
133                };
134                // Show rejected windows too: the operator needs to see WHY the
135                // imminent reset seat cannot serve requests yet.
136                let budget = match &cap.observation {
137                    QuotaObservation::Ok { windows, .. } => {
138                        quota::budget_window(windows, req.model)
139                    }
140                    _ => None,
141                };
142                PlanRow {
143                    credential: cap.credential.clone(),
144                    rank,
145                    tier: eligible.as_ref().map(|e| e.tier),
146                    anchored: eligible.as_ref().and_then(|e| e.anchored),
147                    budget_reset_ms: eligible
148                        .as_ref()
149                        .and_then(|e| e.budget_reset_ms)
150                        .or_else(|| budget.and_then(|w| w.resets_at_ms)),
151                    utilization: eligible
152                        .as_ref()
153                        .and_then(|e| e.utilization)
154                        .or_else(|| budget.and_then(|w| w.used_percent)),
155                    verdict: if selected {
156                        "selected".into()
157                    } else {
158                        rejection
159                            .map(quota::Rejection::describe)
160                            .unwrap_or_else(|| "not considered".into())
161                    },
162                    cause: rejection.and_then(|r| r.cause.clone()),
163                }
164            })
165            .collect();
166        rows.sort_by(|a, b| {
167            a.rank
168                .unwrap_or(u32::MAX)
169                .cmp(&b.rank.unwrap_or(u32::MAX))
170                .then_with(|| a.credential.cmp(&b.credential))
171        });
172        Self {
173            provider: req.provider.to_string(),
174            model: req.model.map(str::to_string),
175            strategy: req.strategy,
176            urgent_horizon_ms: req.urgent_horizon_ms,
177            sticky: req.sticky,
178            current: req.current.cloned(),
179            now_ms: req.now_ms,
180            rows,
181            selection,
182        }
183    }
184}
185
186#[cfg(test)]
187mod tests {
188    use super::*;
189    #[test]
190    fn config_loader_retains_auto_keys_and_surfaces_warnings() {
191        let config = crate::config::load_config_from_str("auth.auto.strategy = preference_order\nauth.auto.strategy.anthropic = lowest_utilization\nauth.auto.sticky = false\nauth.auto.urgent_horizon_hours = 48\n");
192        let (p, warnings) = AutoSelectionPolicy::from_config_map(&config.auth.auto);
193        assert!(warnings.is_empty());
194        assert_eq!(
195            p.strategy_for(OAuthProviderId::Anthropic),
196            Strategy::LowestUtilization
197        );
198        assert_eq!(
199            p.strategy_for(OAuthProviderId::OpenAiCodex),
200            Strategy::PreferenceOrder
201        );
202        assert_eq!(p.urgent_horizon_ms, 48 * 3_600_000);
203        assert!(!p.sticky);
204        let bad = crate::config::load_config_from_str(
205            "auth.auto.sticky = typo\nauth.auto.strategy = typo\n",
206        );
207        assert_eq!(
208            bad.warnings
209                .iter()
210                .filter(|w| w.contains("auth.auto"))
211                .count(),
212            2
213        );
214    }
215
216    #[test]
217    fn defaults_overrides_and_invalid_values() {
218        let p = AutoSelectionPolicy::default();
219        assert_eq!(
220            p.strategy_for(OAuthProviderId::Anthropic),
221            Strategy::SoonestReset
222        );
223        assert!(p.sticky);
224        let input = [
225            ("strategy", "lowest_utilization"),
226            ("strategy.anthropic", "preference_order"),
227            ("urgent_horizon_hours", "48"),
228            ("sticky", "false"),
229        ]
230        .into_iter()
231        .map(|(k, v)| (k.into(), v.into()))
232        .collect();
233        let (p, w) = AutoSelectionPolicy::from_config_map(&input);
234        assert!(w.is_empty());
235        assert_eq!(
236            p.strategy_for(OAuthProviderId::Anthropic),
237            Strategy::PreferenceOrder
238        );
239        assert_eq!(
240            p.strategy_for(OAuthProviderId::OpenAiCodex),
241            Strategy::LowestUtilization
242        );
243        assert_eq!(p.urgent_horizon_ms, 48 * 3_600_000);
244        assert!(!p.sticky);
245        for value in ["0", "169", "-1", "bad"] {
246            let (p, w) = AutoSelectionPolicy::from_config_map(&BTreeMap::from([(
247                "urgent_horizon_hours".into(),
248                value.into(),
249            )]));
250            assert_eq!(p.urgent_horizon_ms, quota::DEFAULT_URGENT_HORIZON_MS);
251            assert_eq!(w.len(), 1);
252        }
253        let (p, w) = AutoSelectionPolicy::from_config_map(&BTreeMap::from([
254            ("strategy.anthropic".into(), "garbage".into()),
255            ("strategy".into(), "lowest_utilization".into()),
256        ]));
257        assert_eq!(
258            p.strategy_for(OAuthProviderId::Anthropic),
259            Strategy::SoonestReset
260        );
261        assert_eq!(w.len(), 1);
262    }
263}
agent_core/core/auth/auto_selection.rs

agent_core/core/auth/
auto_selection.rs
