agent_core/core/auth/
usage_reads.rs

1//! Process-wide table of Anthropic/OAuth usage reads, one entry per account
2//! slot, shared by every [`LocalBroker`](super::LocalBroker) in the process
3//! (a daemon can hold several brokers, one per runtime `TokenCache`).
4//!
5//! * **Single-flight.** One network read per slot at a time; a reader that
6//!   waited for a read in flight takes that read's result (success or
7//!   failure) and sends nothing. The per-slot gate is never removed, so
8//!   invalidating a cached reading can never open a second flight.
9//! * **Negative cache.** A transient failure (usage-endpoint 429, 5xx,
10//!   timeout, transport) holds the slot: reads inside the hold return the
11//!   same failure with no request. Holds back off 2, 4, 8, then 10 minutes;
12//!   a longer `Retry-After` wins, capped at 10 minutes. A success clears the
13//!   hold. Credential rejections are never held (their path is unchanged).
14//! * **Last good.** The slot's last successful snapshot and the seat it was
15//!   read from, which automatic selection may use in place of a read that
16//!   failed transiently (see [`USAGE_READ_GRACE`],
17//!   [`UsageReads::last_good_for`] and `quota_policy::SelectionRequest::grace`).
18//!
19//! Keys are the stable slot identity (credential file + storage key + usage
20//! endpoint), never the seat digest, which changes on token rotation. Last
21//! good is bound to the account it was read from: when the credential now in
22//! the slot carries an account id, the reading counts only for that same id
23//! (a refresh rotation keeps it; a re-login to another account drops it).
24//! Credentials without an account id (an Anthropic login whose profile
25//! lookup failed, or one made before ids were recorded) cannot tell a
26//! rotation from a re-login, so their reading counts only for the exact
27//! credential it was read with: a rotation loses grace until the next good
28//! read, and a re-login to another account is never graced on the previous
29//! occupant's reading.
30
31use std::collections::HashMap;
32use std::future::Future;
33use std::sync::{Arc, Mutex, MutexGuard, OnceLock};
34use std::time::Duration;
35
36use super::account::SeatIdentity;
37use super::broker::BrokerError;
38use super::quota_policy::ObservationFailure;
39use super::usage::{UsageError, UsageSnapshot};
40
41/// First hold after a transient usage-read failure.
42pub const NEGATIVE_BACKOFF_START: Duration = Duration::from_secs(2 * 60);
43/// Longest hold, including one asked for by `Retry-After`.
44pub const NEGATIVE_BACKOFF_CAP: Duration = Duration::from_secs(10 * 60);
45/// A last successful reading younger than this may stand in for a read that
46/// failed transiently. Never for an exhausted window, an auth rejection or
47/// an active cooldown.
48pub const USAGE_READ_GRACE: Duration = Duration::from_secs(30 * 60);
49/// Tracing target of the per-read INFO line (passes the daemon log filter).
50pub const USAGE_READ_TARGET: &str = "agent_core::auth::usage_read";
51
52/// A usage read that failed for a reason that says nothing about the
53/// account: rate limit, server error, timeout or transport. Secret-free.
54#[derive(Debug, Clone, PartialEq, Eq)]
55pub struct TransientReadFailure {
56    /// The usage error kind (`rate_limited`, `upstream_status`, `timeout`,
57    /// `transport`).
58    pub kind: &'static str,
59    pub status: Option<u16>,
60    pub retry_after_s: Option<u64>,
61    /// The failed read itself, so a held read reports exactly what the
62    /// network read did (broker error text and typed cause).
63    error: UsageError,
64}
65
66impl TransientReadFailure {
67    /// `Some` only for the failures the negative cache holds.
68    pub fn from_usage_error(e: &UsageError) -> Option<Self> {
69        let (status, retry_after_s) = match e {
70            UsageError::RateLimited { retry_after_s } => (Some(429), *retry_after_s),
71            UsageError::UpstreamStatus {
72                status,
73                retry_after_s,
74            } if *status >= 500 => (Some(*status), *retry_after_s),
75            UsageError::Timeout | UsageError::Transport { .. } => (None, None),
76            _ => return None,
77        };
78        Some(Self {
79            kind: e.kind(),
80            status,
81            retry_after_s,
82            error: e.clone(),
83        })
84    }
85
86    pub fn is_rate_limited(&self) -> bool {
87        self.status == Some(429)
88    }
89
90    /// The broker error callers saw for this failure before the table.
91    pub fn to_broker_error(&self) -> BrokerError {
92        self.error.clone().into_broker_error()
93    }
94
95    /// The typed cause the selection report keeps for this failure.
96    pub fn cause(&self) -> ObservationFailure {
97        self.error.failure()
98    }
99}
100
101/// What one network read produced, before the table classifies it.
102// One value per network read, moved straight into the table: boxing the
103// snapshot would buy nothing.
104#[allow(clippy::large_enum_variant)]
105pub enum Fetched {
106    Ok(UsageSnapshot, SeatIdentity),
107    /// The usage request itself failed.
108    Usage(UsageError),
109    /// The read failed before a usage request (slot check, token refresh),
110    /// with its typed cause.
111    Failed {
112        error: BrokerError,
113        cause: ObservationFailure,
114    },
115}
116
117#[derive(Debug, Clone)]
118pub enum ReadError {
119    /// Held: no request for this slot before `retry_at_ms` (epoch ms).
120    Transient {
121        failure: TransientReadFailure,
122        retry_at_ms: u64,
123    },
124    Other {
125        error: BrokerError,
126        cause: ObservationFailure,
127    },
128}
129
130impl ReadError {
131    pub fn to_broker_error(&self) -> BrokerError {
132        match self {
133            Self::Transient { failure, .. } => failure.to_broker_error(),
134            Self::Other { error, .. } => error.clone(),
135        }
136    }
137
138    /// The typed cause the selection report keeps for this failure.
139    pub fn cause(&self) -> ObservationFailure {
140        match self {
141            Self::Transient { failure, .. } => failure.cause(),
142            Self::Other { cause, .. } => cause.clone(),
143        }
144    }
145}
146
147pub type ReadResult = Result<(UsageSnapshot, SeatIdentity), ReadError>;
148
149/// How a read was answered.
150#[derive(Debug, Clone, Copy, PartialEq, Eq)]
151pub enum ReadPath {
152    Network,
153    Joined,
154    NegativeCacheHit,
155}
156
157#[derive(Debug, Clone)]
158struct Hold {
159    failure: TransientReadFailure,
160    until_ms: u64,
161    backoff_ms: u64,
162}
163
164#[derive(Default)]
165struct Entry {
166    /// Bumped by every recorded network read.
167    generation: u64,
168    last: Option<ReadResult>,
169    hold: Option<Hold>,
170    /// Last successful snapshot and the seat whose token read it.
171    last_good: Option<(UsageSnapshot, SeatIdentity)>,
172}
173
174/// Whether a reading taken from seat `read_from` may stand for the
175/// credential now stored in the slot (`now`): the same account id when
176/// either carries one; without ids on both sides, only the exact same
177/// credential (see the module docs).
178fn same_account(read_from: &SeatIdentity, now: &SeatIdentity) -> bool {
179    match (read_from.seat_fingerprint(), now.seat_fingerprint()) {
180        (Some(a), Some(b)) => a == b,
181        (None, None) => read_from == now,
182        _ => false,
183    }
184}
185
186/// See the module docs. [`process_table`] is the one production instance.
187#[derive(Default)]
188pub struct UsageReads {
189    entries: Mutex<HashMap<String, Entry>>,
190    gates: Mutex<HashMap<String, Arc<tokio::sync::Mutex<()>>>>,
191}
192
193/// The table shared by every broker in this process.
194pub fn process_table() -> Arc<UsageReads> {
195    static TABLE: OnceLock<Arc<UsageReads>> = OnceLock::new();
196    TABLE.get_or_init(|| Arc::new(UsageReads::default())).clone()
197}
198
199fn ms(d: Duration) -> u64 {
200    d.as_millis() as u64
201}
202
203impl UsageReads {
204    fn entries(&self) -> MutexGuard<'_, HashMap<String, Entry>> {
205        self.entries.lock().unwrap_or_else(|e| e.into_inner())
206    }
207
208    /// The active hold for `key`: its failure and when it ends.
209    pub fn held_at(&self, key: &str, now_ms: u64) -> Option<(TransientReadFailure, u64)> {
210        self.entries()
211            .get(key)
212            .and_then(|e| e.hold.as_ref())
213            .filter(|h| h.until_ms > now_ms)
214            .map(|h| (h.failure.clone(), h.until_ms))
215    }
216
217    /// Record one network read for `key` at `now_ms` and classify it.
218    // ReadError carries the broker's own error (as `usage()` returns it);
219    // one per read, so its size is no cost worth a box.
220    #[allow(clippy::result_large_err)]
221    pub fn record_at(&self, key: &str, fetched: Fetched, now_ms: u64) -> ReadResult {
222        let mut entries = self.entries();
223        let entry = entries.entry(key.to_string()).or_default();
224        let result = match fetched {
225            Fetched::Ok(snapshot, seat) => {
226                entry.hold = None;
227                entry.last_good = Some((snapshot.clone(), seat.clone()));
228                Ok((snapshot, seat))
229            }
230            Fetched::Usage(e) => match TransientReadFailure::from_usage_error(&e) {
231                Some(failure) => {
232                    let cap = ms(NEGATIVE_BACKOFF_CAP);
233                    let backoff_ms = match &entry.hold {
234                        Some(h) => h.backoff_ms.saturating_mul(2).min(cap),
235                        None => ms(NEGATIVE_BACKOFF_START),
236                    };
237                    let asked = failure.retry_after_s.map_or(0, |s| s.saturating_mul(1000));
238                    let until_ms = now_ms.saturating_add(backoff_ms.max(asked).min(cap));
239                    entry.hold = Some(Hold {
240                        failure: failure.clone(),
241                        until_ms,
242                        backoff_ms,
243                    });
244                    Err(ReadError::Transient {
245                        failure,
246                        retry_at_ms: until_ms,
247                    })
248                }
249                None => {
250                    if e.is_auth_failure() {
251                        entry.last_good = None;
252                    }
253                    entry.hold = None;
254                    Err(ReadError::Other {
255                        cause: e.failure(),
256                        error: e.into_broker_error(),
257                    })
258                }
259            },
260            Fetched::Failed { error, cause } => {
261                // A refused or missing login ends any grace for the slot.
262                if cause.kind.needs_login() {
263                    entry.last_good = None;
264                }
265                entry.hold = None;
266                Err(ReadError::Other { error, cause })
267            }
268        };
269        entry.generation = entry.generation.wrapping_add(1);
270        entry.last = Some(result.clone());
271        result
272    }
273
274    /// The slot's last successful snapshot, if it was read from the account
275    /// now stored in the slot (`current`): never another account's reading.
276    pub fn last_good_for(&self, key: &str, current: &SeatIdentity) -> Option<UsageSnapshot> {
277        self.entries()
278            .get(key)
279            .and_then(|e| e.last_good.as_ref())
280            .filter(|(_, read_from)| same_account(read_from, current))
281            .map(|(snapshot, _)| snapshot.clone())
282    }
283
284    /// Drop the slot's last good reading (e.g. after a reported limit).
285    pub fn forget_last_good(&self, key: &str) {
286        if let Some(e) = self.entries().get_mut(key) {
287            e.last_good = None;
288        }
289    }
290
291    /// Drop the slot's hold and last good reading (`synaps auth state
292    /// --clear`): the next read goes to the network. The gate stays, so a
293    /// read in flight keeps its single flight; if one is in flight, its
294    /// result is still recorded when it lands.
295    pub fn forget(&self, key: &str) {
296        if let Some(e) = self.entries().get_mut(key) {
297            e.hold = None;
298            e.last_good = None;
299        }
300    }
301
302    fn generation(&self, key: &str) -> u64 {
303        self.entries().get(key).map_or(0, |e| e.generation)
304    }
305
306    fn gate(&self, key: &str) -> Arc<tokio::sync::Mutex<()>> {
307        self.gates
308            .lock()
309            .unwrap_or_else(|e| e.into_inner())
310            .entry(key.to_string())
311            .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())))
312            .clone()
313    }
314
315    /// Read `key`: a held failure is returned without a request; otherwise
316    /// one flight per key, and a reader that waited for a flight takes its
317    /// result.
318    pub async fn read<F, Fut>(&self, key: &str, fetch: F) -> (ReadResult, ReadPath)
319    where
320        F: FnOnce() -> Fut,
321        Fut: Future<Output = Fetched>,
322    {
323        if let Some((failure, retry_at_ms)) = self.held_at(key, crate::epoch_millis()) {
324            return (
325                Err(ReadError::Transient {
326                    failure,
327                    retry_at_ms,
328                }),
329                ReadPath::NegativeCacheHit,
330            );
331        }
332        let seen = self.generation(key);
333        let gate = self.gate(key);
334        let _flight = gate.lock().await;
335        let joined = self
336            .entries()
337            .get(key)
338            .filter(|e| e.generation != seen)
339            .and_then(|e| e.last.clone());
340        if let Some(result) = joined {
341            return (result, ReadPath::Joined);
342        }
343        let fetched = fetch().await;
344        (
345            self.record_at(key, fetched, crate::epoch_millis()),
346            ReadPath::Network,
347        )
348    }
349}
350
351#[cfg(test)]
352mod tests {
353    use super::*;
354
355    const MIN: u64 = 60_000;
356
357    fn limited(retry_after_s: Option<u64>) -> Fetched {
358        Fetched::Usage(UsageError::RateLimited { retry_after_s })
359    }
360
361    fn held_for(t: &UsageReads, fetched: Fetched, now: u64) -> u64 {
362        match t.record_at("k", fetched, now) {
363            Err(ReadError::Transient { retry_at_ms, .. }) => {
364                assert_eq!(
365                    t.held_at("k", now + 1).map(|(_, until)| until),
366                    Some(retry_at_ms),
367                    "a read inside the hold is answered from the table"
368                );
369                assert!(t.held_at("k", retry_at_ms).is_none(), "the hold ends on time");
370                retry_at_ms - now
371            }
372            other => panic!("expected a held failure, got {other:?}"),
373        }
374    }
375
376    fn seat(refresh: &str, account_id: Option<&str>) -> SeatIdentity {
377        let creds = super::super::OAuthCredentials {
378            auth_type: "oauth".into(),
379            refresh: refresh.into(),
380            access: "a".into(),
381            expires: 0,
382            account_id: account_id.map(Into::into),
383        };
384        SeatIdentity::of(super::super::OAuthProviderId::Anthropic, &creds)
385    }
386
387    fn good_from(read_from: SeatIdentity) -> Fetched {
388        let snapshot = super::super::usage::parse_anthropic_usage(
389            r#"{"seven_day":{"utilization":10,"resets_at":"2099-01-01T00:00:00Z"}}"#,
390            "a",
391            1,
392        )
393        .unwrap();
394        Fetched::Ok(snapshot, read_from)
395    }
396
397    fn good() -> Fetched {
398        good_from(seat("r", None))
399    }
400
401    #[test]
402    fn failed_reads_back_off_two_four_eight_ten_minutes_and_a_success_clears() {
403        let t = UsageReads::default();
404        let mut now = 1_000_000;
405        let mut holds = Vec::new();
406        for _ in 0..5 {
407            let held = held_for(&t, limited(None), now);
408            holds.push(held / MIN);
409            now += held;
410        }
411        assert_eq!(holds, vec![2, 4, 8, 10, 10]);
412        assert!(t.record_at("k", good(), now).is_ok());
413        assert!(t.held_at("k", now).is_none(), "a success clears the hold");
414        assert!(t.last_good_for("k", &seat("r", None)).is_some());
415        assert_eq!(held_for(&t, limited(None), now), 2 * MIN, "and restarts the backoff");
416
417        // A longer Retry-After wins, capped at 10 minutes.
418        let t = UsageReads::default();
419        assert_eq!(held_for(&t, limited(Some(300)), now), 5 * MIN);
420        assert_eq!(held_for(&t, limited(Some(3600)), now + 5 * MIN), 10 * MIN);
421        let t = UsageReads::default();
422        assert_eq!(held_for(&t, limited(Some(30)), now), 2 * MIN, "a shorter one does not");
423
424        // 5xx, timeouts and transport failures are held too; a 5xx keeps its
425        // Retry-After like a 429.
426        for e in [
427            UsageError::UpstreamStatus {
428                status: 503,
429                retry_after_s: None,
430            },
431            UsageError::Timeout,
432            UsageError::Transport { detail: "connect".into() },
433        ] {
434            let t = UsageReads::default();
435            assert_eq!(held_for(&t, Fetched::Usage(e), now), 2 * MIN);
436        }
437        let t = UsageReads::default();
438        let slow = UsageError::UpstreamStatus {
439            status: 503,
440            retry_after_s: Some(300),
441        };
442        assert_eq!(held_for(&t, Fetched::Usage(slow), now), 5 * MIN);
443        // Credential rejections and other answers are not, and a refused
444        // login drops the last good reading.
445        let t = UsageReads::default();
446        assert!(t.record_at("k", good(), now).is_ok());
447        for e in [
448            UsageError::Unauthorized { status: 401 },
449            UsageError::UpstreamStatus {
450                status: 404,
451                retry_after_s: None,
452            },
453            UsageError::Malformed { detail: "x".into() },
454        ] {
455            assert!(matches!(
456                t.record_at("k", Fetched::Usage(e), now),
457                Err(ReadError::Other { .. })
458            ));
459            assert!(t.held_at("k", now).is_none());
460        }
461        assert!(
462            t.last_good_for("k", &seat("r", None)).is_none(),
463            "the 401 dropped the last good reading"
464        );
465    }
466
467    /// Last good is bound to the account it was read from: the same account
468    /// id after a token rotation keeps it, another account id (or a switch
469    /// between id and no id) never sees it. Without ids only the exact
470    /// credential the reading came from does (see the module docs).
471    #[test]
472    fn last_good_counts_only_for_the_account_it_was_read_from() {
473        let t = UsageReads::default();
474        assert!(t.record_at("k", good_from(seat("r1", Some("acct-1"))), 1).is_ok());
475        assert!(t.last_good_for("k", &seat("r1", Some("acct-1"))).is_some());
476        assert!(
477            t.last_good_for("k", &seat("r2-rotated", Some("acct-1"))).is_some(),
478            "a refresh rotation keeps the account"
479        );
480        assert!(
481            t.last_good_for("k", &seat("r1", Some("acct-2"))).is_none(),
482            "another account in the slot is never graced on this reading"
483        );
484        assert!(t.last_good_for("k", &seat("r1", None)).is_none());
485        assert!(t.last_good_for("other", &seat("r1", Some("acct-1"))).is_none());
486
487        let t = UsageReads::default();
488        assert!(t.record_at("k", good_from(seat("r1", None)), 1).is_ok());
489        assert!(t.last_good_for("k", &seat("r1", None)).is_some());
490        assert!(
491            t.last_good_for("k", &seat("r2", None)).is_none(),
492            "without ids, other refresh material may be another account"
493        );
494        assert!(t.last_good_for("k", &seat("r1", Some("acct-1"))).is_none());
495    }
496}
agent_core/core/auth/usage_reads.rs

agent_core/core/auth/
usage_reads.rs
