From c1cc349b19a4d1acf4fbbd541753826f171f8499 Mon Sep 17 00:00:00 2001 From: Meta Vens Date: Sun, 1 Nov 2026 02:49:16 +0100 Subject: [PATCH] usage-meter-poll: the cell decides how a reading combines; a poll is as old as its question QuotaCell::observe_at is the one decision for both sensors: a reading older than the held one is dropped, and a rejection stands against a reading that only reports utilization (status polled) until its window resets or a reading that carries its own verdict replaces it. The poll stamps its answer with the moment it asked. Also: Retry-After in HTTP-date form is honoured; a refused token is not sent again until it reads differently (the poll thread no longer exits); the client's timeout is a parameter so a silent endpoint is tested; tests pin redirects(0), the two-window mapping, the slow-server race, and the late poll against a rejection. Co-Authored-By: Claude Sonnet 5.5 --- crates/claude-connector/src/quota.rs | 98 +++++- crates/claude-connector/src/usage_poll.rs | 402 ++++++++++++++++------ 2 files changed, 398 insertions(+), 102 deletions(-) diff --git a/crates/claude-connector/src/quota.rs b/crates/claude-connector/src/quota.rs index 1b3d27e..b319d4c 100644 --- a/crates/claude-connector/src/quota.rs +++ b/crates/claude-connector/src/quota.rs @@ -23,18 +23,35 @@ impl QuotaCell { Self::default() } - /// Records the windows one event reported, stamping each with now. Each - /// replaces that window's earlier observation, and windows the event did - /// not mention stay as they were, stamp included. + /// Records the windows one event reported, stamping each with now. See + /// [`QuotaCell::observe_at`] for how a reading combines with what is held. pub fn observe(&self, windows: impl IntoIterator) { self.observe_at(unix_ms(), windows); } - fn observe_at(&self, now: u64, windows: impl IntoIterator) { + /// Records readings taken at `at` (unix ms): the moment the sensor asked, + /// not the moment its answer arrived. The cell alone decides how a reading + /// combines with the one it holds for that window, and every sensor goes + /// through this decision: + /// + /// - a reading older than the held one is dropped, so a slow answer cannot + /// replace what a faster sensor saw meanwhile; + /// - the provider's rejection stands until its window resets, or until a + /// reading that itself carries a verdict from the provider replaces it: + /// a [`POLLED`] reading reports utilization only, and says nothing about + /// whether the provider accepts calls. + /// + /// Windows the readings did not mention stay as they were, stamp included. + pub fn observe_at(&self, at: u64, windows: impl IntoIterator) { let mut held = self.0.lock().unwrap_or_else(|e| e.into_inner()); for mut quota in windows { - quota.observed_at_unix_ms = now; - held.insert(quota.window.clone(), quota); + quota.observed_at_unix_ms = at; + if held + .get("a.window) + .is_none_or(|current| supersedes(current, "a)) + { + held.insert(quota.window.clone(), quota); + } } } @@ -65,6 +82,23 @@ impl QuotaCell { /// The status of a window whose reset has passed since it was observed. const EXPIRED: &str = "expired"; +/// The status of a window the provider refuses calls on. +pub(crate) const REJECTED: &str = "rejected"; +/// The status of a window read from the usage endpoint: utilization and reset +/// time, with no verdict on whether calls are accepted. +pub(crate) const POLLED: &str = "polled"; + +/// Whether `incoming` (already stamped with its own reading time) replaces +/// `current` for the same window. +fn supersedes(current: &PlanQuota, incoming: &PlanQuota) -> bool { + if incoming.observed_at_unix_ms < current.observed_at_unix_ms { + return false; + } + let rejection_stands = current.status == REJECTED + && incoming.status == POLLED + && incoming.observed_at_unix_ms < current.resets_at_unix_ms; + !rejection_stands +} fn unix_ms() -> u64 { SystemTime::now() @@ -135,6 +169,58 @@ mod tests { assert_eq!(expired[0].resets_at_unix_ms, RESET); } + fn with_status(mut quota: PlanQuota, status: &str) -> PlanQuota { + quota.status = status.to_string(); + quota + } + + fn held(cell: &QuotaCell) -> (String, Option, u64) { + let window = cell.snapshot_at(250).windows.remove(0); + ( + window.status, + window.utilization_permille, + window.observed_at_unix_ms, + ) + } + + #[test] + fn a_reading_older_than_the_held_one_is_dropped_and_an_equal_one_replaces() { + let cell = QuotaCell::new(); + cell.observe_at(200, [quota("five_hour", 80, 0)]); + cell.observe_at(100, [quota("five_hour", 10, 0)]); + assert_eq!(held(&cell), ("allowed".to_string(), Some(80), 200)); + cell.observe_at(200, [quota("five_hour", 90, 0)]); + assert_eq!(held(&cell), ("allowed".to_string(), Some(90), 200)); + } + + #[test] + fn a_rejection_stands_against_a_poll_until_the_window_resets() { + let cell = QuotaCell::new(); + let polled = with_status(quota("five_hour", 500, 0), POLLED); + cell.observe_at(100, [with_status(quota("five_hour", 1000, 0), REJECTED)]); + cell.observe_at(200, [polled.clone()]); + assert_eq!(held(&cell), (REJECTED.to_string(), Some(1000), 100)); + // The window the rejection described is over: a poll of the next one lands. + cell.observe_at(RESET, [polled]); + assert_eq!(cell.snapshot_at(RESET - 1).windows[0].status, POLLED); + } + + #[test] + fn a_newer_reading_that_carries_its_own_verdict_lifts_a_rejection() { + let cell = QuotaCell::new(); + cell.observe_at(100, [with_status(quota("five_hour", 1000, 0), REJECTED)]); + cell.observe_at(200, [quota("five_hour", 300, 0)]); + assert_eq!(held(&cell), ("allowed".to_string(), Some(300), 200)); + } + + #[test] + fn a_poll_replaces_what_is_not_a_rejection() { + let cell = QuotaCell::new(); + cell.observe_at(100, [quota("five_hour", 300, 0)]); + cell.observe_at(200, [with_status(quota("five_hour", 400, 0), POLLED)]); + assert_eq!(held(&cell), (POLLED.to_string(), Some(400), 200)); + } + #[test] fn the_live_clock_stamps_and_expires() { let cell = QuotaCell::new(); diff --git a/crates/claude-connector/src/usage_poll.rs b/crates/claude-connector/src/usage_poll.rs index 013269f..e23a74f 100644 --- a/crates/claude-connector/src/usage_poll.rs +++ b/crates/claude-connector/src/usage_poll.rs @@ -3,7 +3,8 @@ //! quiet the weekly reading goes stale; this thread asks the subscription's //! usage endpoint when the `seven_day` reading is `STALE` or older, and //! records the answer in the same [`QuotaCell`] the stream feeds. Neither -//! sensor owns the reading: the newest observation of a window wins. A failed +//! sensor owns the reading: the cell decides how a reading combines with the +//! one it holds, and a poll is stamped with the moment it asked. A failed //! poll never touches the cell; only the cell's own reset rule retires a //! window. //! @@ -17,10 +18,11 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; use model_connector::PlanQuota; use serde_json::Value; +use sha2::{Digest, Sha256}; use zeroize::Zeroizing; use crate::backend::TokenSource; -use crate::quota::QuotaCell; +use crate::quota::{POLLED, QuotaCell}; /// The usage endpoint. const ENDPOINT: &str = "https://api.anthropic.com/api/oauth/usage"; @@ -31,8 +33,6 @@ const CEILING: Duration = Duration::from_secs(6 * 60 * 60); const TIMEOUT: Duration = Duration::from_secs(10); /// A usage body is a few hundred bytes; this is a ceiling, not an expectation. const MAX_BODY: u64 = 1 << 20; -/// The status every polled window carries. -const POLLED: &str = "polled"; /// What the endpoint said, reduced to what the schedule needs. It never holds /// the token. @@ -63,7 +63,7 @@ impl State { Self::NoToken => "usage poll: no token to ask with", Self::RateLimited => "usage poll: rate limited, waiting", Self::Refused => { - "usage poll: refused (401/403), the token likely lacks the user:profile scope; polling stopped until restart" + "usage poll: refused (401/403), the token likely lacks the user:profile scope; not asking again until the token changes" } Self::NotFound => "usage poll: endpoint not found, backing off", Self::Unreadable => "usage poll: answer not readable, backing off", @@ -73,40 +73,35 @@ impl State { } } -#[derive(Debug, PartialEq, Eq)] -enum Next { - Wait(Duration), - Stop, -} - #[derive(Debug, PartialEq)] struct Verdict { windows: Vec, - next: Next, + /// How long until the next poll may happen. + wait: Duration, state: State, } /// What one answer means: the windows to record (only a readable body has /// any), how long to wait, and the state to report. fn judge(answer: &Answer) -> Verdict { - let (windows, next, state) = match answer { + let (windows, wait, state) = match answer { Answer::Body(body) => match parse_usage(body) { - Some(windows) => (windows, Next::Wait(STALE), State::Polled), - None => (Vec::new(), Next::Wait(CEILING), State::Unreadable), + Some(windows) => (windows, STALE, State::Polled), + None => (Vec::new(), CEILING, State::Unreadable), }, Answer::Status(429, retry_after) => ( Vec::new(), - Next::Wait(retry_after.unwrap_or_default().clamp(STALE, CEILING)), + retry_after.unwrap_or_default().clamp(STALE, CEILING), State::RateLimited, ), - Answer::Status(401 | 403, _) => (Vec::new(), Next::Stop, State::Refused), - Answer::Status(404, _) => (Vec::new(), Next::Wait(CEILING), State::NotFound), - Answer::Status(..) => (Vec::new(), Next::Wait(STALE), State::Unexpected), - Answer::Unreachable => (Vec::new(), Next::Wait(STALE), State::Unreachable), + Answer::Status(401 | 403, _) => (Vec::new(), STALE, State::Refused), + Answer::Status(404, _) => (Vec::new(), CEILING, State::NotFound), + Answer::Status(..) => (Vec::new(), STALE, State::Unexpected), + Answer::Unreachable => (Vec::new(), STALE, State::Unreachable), }; Verdict { windows, - next, + wait, state, } } @@ -153,15 +148,22 @@ struct UsageClient { impl UsageClient { fn new(url: impl Into) -> Self { + Self::with_timeout(url, TIMEOUT) + } + + fn with_timeout(url: impl Into, timeout: Duration) -> Self { Self { agent: ureq::AgentBuilder::new() - .timeout(TIMEOUT) + .timeout(timeout) .redirects(0) .build(), url: url.into(), } } + /// Asks the endpoint. Our copy of the header value is zeroized when it + /// drops; ureq copies it into its own request buffers, which are freed + /// without zeroing. fn ask(&self, token: &Zeroizing) -> Answer { let bearer = Zeroizing::new(format!("Bearer {}", token.as_str())); let sent = self @@ -184,11 +186,9 @@ impl UsageClient { } Ok(response) => Answer::Status(response.status(), None), Err(ureq::Error::Status(code, response)) => { - // Whole seconds only; the HTTP-date form counts as absent. let retry_after = response .header("retry-after") - .and_then(|seconds| seconds.trim().parse::().ok()) - .map(Duration::from_secs); + .and_then(|value| retry_after(value, unix_ms())); Answer::Status(code, retry_after) } Err(ureq::Error::Transport(_)) => Answer::Unreachable, @@ -196,6 +196,19 @@ impl UsageClient { } } +/// A `Retry-After` value as a wait from `now_ms`: delay-seconds, or an +/// HTTP-date (RFC 9110 10.2.3) measured from now, a date already past being a +/// zero wait. `None` when it is neither. +fn retry_after(value: &str, now_ms: u64) -> Option { + let value = value.trim(); + if let Ok(seconds) = value.parse::() { + return Some(Duration::from_secs(seconds)); + } + let date = chrono::DateTime::parse_from_rfc2822(value).ok()?; + let until = u64::try_from(date.timestamp_millis()).ok()?; + Some(Duration::from_millis(until.saturating_sub(now_ms))) +} + /// The schedule and the thing it drives. Time comes in as an argument so the /// schedule is testable without waiting. struct Poller { @@ -206,7 +219,9 @@ struct Poller { /// No poll before this instant, unix ms: the last answer's own wait. not_before_ms: u64, started: bool, - stopped: bool, + /// The digest of the token the endpoint last refused: while the token + /// reads the same, the poller does not ask again. + refused_token: Option<[u8; 32]>, last_state: Option, } @@ -224,51 +239,58 @@ impl Poller { log, not_before_ms: 0, started: false, - stopped: false, + refused_token: None, last_state: None, } } /// Polls if a poll is due at `now_ms`; returns how long to wait before - /// asking again, or `None` once polling has stopped for good. + /// asking again. /// /// The first call polls. After that a poll is due only when the cell's /// `seven_day` observation is `STALE` old or older (or absent), and never - /// before the last answer's own wait has passed. - fn tick(&mut self, now_ms: u64) -> Option { - if self.stopped { - return None; - } + /// before the last answer's own wait has passed. A token the endpoint + /// refused is not sent again until it reads differently. + fn tick(&mut self, now_ms: u64) -> Duration { let due_at = self.due_at_ms(); if now_ms < due_at { - return Some(Duration::from_millis(due_at - now_ms)); + return Duration::from_millis(due_at - now_ms); } self.started = true; - let verdict = match self.tokens.read() { - Some(token) => judge(&self.client.ask(&token)), + let token = self.tokens.read(); + if let (Some(refused), Some(token)) = (&self.refused_token, &token) + && digest(token) != *refused + { + self.refused_token = None; + } + let verdict = match token { + _ if self.refused_token.is_some() => Verdict { + windows: Vec::new(), + wait: STALE, + state: State::Refused, + }, + Some(token) => { + let verdict = judge(&self.client.ask(&token)); + if verdict.state == State::Refused { + self.refused_token = Some(digest(&token)); + } + verdict + } None => Verdict { windows: Vec::new(), - next: Next::Wait(STALE), + wait: STALE, state: State::NoToken, }, }; - self.cell.observe(verdict.windows); + // The reading is as old as the question: an answer that arrives after + // the stream spoke does not outrank it. + self.cell.observe_at(now_ms, verdict.windows); if self.last_state != Some(verdict.state) { (self.log)(verdict.state.line()); self.last_state = Some(verdict.state); } - match verdict.next { - Next::Stop => { - self.stopped = true; - None - } - Next::Wait(wait) => { - self.not_before_ms = now_ms.saturating_add(duration_ms(wait)); - Some(Duration::from_millis( - self.due_at_ms().saturating_sub(now_ms), - )) - } - } + self.not_before_ms = now_ms.saturating_add(duration_ms(verdict.wait)); + Duration::from_millis(self.due_at_ms().saturating_sub(now_ms)) } /// The earliest instant the next poll may happen. @@ -289,12 +311,17 @@ impl Poller { } fn run(mut self) { - while let Some(wait) = self.tick(unix_ms()) { + loop { + let wait = self.tick(unix_ms()); std::thread::sleep(wait); } } } +fn digest(token: &Zeroizing) -> [u8; 32] { + Sha256::digest(token.as_bytes()).into() +} + fn duration_ms(duration: Duration) -> u64 { u64::try_from(duration.as_millis()).unwrap_or(u64::MAX) } @@ -408,38 +435,32 @@ mod tests { // ---- judgement ---- - fn wait_of(answer: &Answer) -> Next { - judge(answer).next + fn wait_of(answer: &Answer) -> Duration { + judge(answer).wait } #[test] fn rate_limit_waits_the_longer_of_retry_after_and_the_floor_capped() { let limited = |after| wait_of(&Answer::Status(429, after)); - assert_eq!(limited(Some(secs(3600))), Next::Wait(secs(3600))); - assert_eq!(limited(Some(secs(20 * 60))), Next::Wait(secs(20 * 60))); - assert_eq!(limited(Some(secs(10))), Next::Wait(STALE)); - assert_eq!(limited(None), Next::Wait(STALE)); - assert_eq!(limited(Some(secs(100_000))), Next::Wait(CEILING)); - assert_eq!(limited(Some(secs(6 * 3600))), Next::Wait(CEILING)); + assert_eq!(limited(Some(secs(3600))), secs(3600)); + assert_eq!(limited(Some(secs(20 * 60))), secs(20 * 60)); + assert_eq!(limited(Some(secs(10))), STALE); + assert_eq!(limited(None), STALE); + assert_eq!(limited(Some(secs(100_000))), CEILING); + assert_eq!(limited(Some(secs(6 * 3600))), CEILING); assert_eq!(judge(&Answer::Status(429, None)).state, State::RateLimited); } #[test] - fn refusal_stops_polling_and_the_rest_back_off() { - assert_eq!(wait_of(&Answer::Status(401, None)), Next::Stop); - assert_eq!(wait_of(&Answer::Status(403, Some(secs(5)))), Next::Stop); + fn refusal_backs_off_like_the_rest() { + assert_eq!(wait_of(&Answer::Status(401, None)), STALE); + assert_eq!(wait_of(&Answer::Status(403, Some(secs(5)))), STALE); assert_eq!(judge(&Answer::Status(403, None)).state, State::Refused); - assert_eq!(wait_of(&Answer::Status(404, None)), Next::Wait(CEILING)); - assert_eq!( - wait_of(&Answer::Body("".to_string())), - Next::Wait(CEILING) - ); - assert_eq!(wait_of(&Answer::Unreachable), Next::Wait(STALE)); - assert_eq!(wait_of(&Answer::Status(503, None)), Next::Wait(STALE)); - assert_eq!( - wait_of(&Answer::Body(body(&reading(1.0), "null"))), - Next::Wait(STALE) - ); + assert_eq!(wait_of(&Answer::Status(404, None)), CEILING); + assert_eq!(wait_of(&Answer::Body("".to_string())), CEILING); + assert_eq!(wait_of(&Answer::Unreachable), STALE); + assert_eq!(wait_of(&Answer::Status(503, None)), STALE); + assert_eq!(wait_of(&Answer::Body(body(&reading(1.0), "null"))), STALE); } // ---- a fake server, a poller, and the cell ---- @@ -530,7 +551,7 @@ mod tests { )]); let cell = Arc::new(QuotaCell::new()); let (mut poller, lines) = make_poller(&cell, &url, Some(CANARY)); - let wait = poller.tick(unix_ms()).unwrap(); + let wait = poller.tick(unix_ms()); let windows = cell.snapshot().windows; assert_eq!(windows.len(), 2); assert_eq!(windows[0].status, "polled"); @@ -555,7 +576,7 @@ mod tests { let observed = window(&cell, "seven_day").unwrap().observed_at_unix_ms; // 5 min old: no request, and the wait is the time left. - let wait = poller.tick(observed + 5 * MIN).unwrap(); + let wait = poller.tick(observed + 5 * MIN); assert_eq!(requests(&seen), 1); assert_eq!(wait, Duration::from_millis(10 * MIN)); // One millisecond short of 15 min: still none. @@ -590,7 +611,7 @@ mod tests { // The stream reports seven_day afresh, after the poll. cell.observe([stream_reading("seven_day", 400)]); let observed = window(&cell, "seven_day").unwrap().observed_at_unix_ms; - let wait = poller.tick(observed + 14 * MIN).unwrap(); + let wait = poller.tick(observed + 14 * MIN); assert_eq!(requests(&seen), 1); assert_eq!(wait, Duration::from_millis(MIN)); } @@ -599,19 +620,23 @@ mod tests { fn failures_never_touch_the_cell() { let limited = |headers: &str| http("429 Too Many Requests", headers, "{}"); let cases = [ - (limited("Retry-After: 3600\r\n"), Some(60)), - (limited(""), Some(15)), + (limited("Retry-After: 3600\r\n"), 60), + (limited(""), 15), ( limited("Retry-After: Wed, 21 Oct 2026 07:28:00 GMT\r\n"), - Some(15), + 15, ), - (http("403 Forbidden", "", "{}"), None), - (http("401 Unauthorized", "", "{}"), None), - (http("404 Not Found", "", "{}"), Some(360)), - (http("200 OK", "", "not json"), Some(360)), + ( + limited("Retry-After: Fri, 31 Dec 2100 23:59:59 GMT\r\n"), + 360, + ), + (http("403 Forbidden", "", "{}"), 15), + (http("401 Unauthorized", "", "{}"), 15), + (http("404 Not Found", "", "{}"), 360), + (http("200 OK", "", "not json"), 360), ( http("302 Found", "Location: http://127.0.0.1:9/\r\n", ""), - Some(15), + 15, ), ]; for (reply, minutes) in cases { @@ -625,9 +650,8 @@ mod tests { let (mut poller, lines) = make_poller(&cell, &url, Some(CANARY)); let wait = poller.tick(unix_ms()); assert_eq!(cell.snapshot(), before, "{reply}"); - assert_eq!(wait.map(|w| w.as_secs() / 60), minutes, "{reply}"); + assert_eq!(wait.as_secs() / 60, minutes, "{reply}"); assert_eq!(lines.lock().unwrap().len(), 1, "{reply}"); - assert_eq!(poller.tick(unix_ms()).is_none(), minutes.is_none()); } } @@ -642,13 +666,13 @@ mod tests { let cell = Arc::new(QuotaCell::new()); let (mut poller, lines) = make_poller(&cell, &format!("http://127.0.0.1:{port}/u"), Some(CANARY)); - assert_eq!(poller.tick(unix_ms()), Some(STALE)); + assert_eq!(poller.tick(unix_ms()), STALE); assert!(cell.snapshot().windows.is_empty()); assert_eq!(lines.lock().unwrap().len(), 1); let (url, seen) = serve(vec![http("200 OK", "", "{}")]); let (mut poller, lines) = make_poller(&cell, &url, None); - assert_eq!(poller.tick(unix_ms()), Some(STALE)); + assert_eq!(poller.tick(unix_ms()), STALE); assert_eq!(requests(&seen), 0); assert_eq!(lines.lock().unwrap()[0], State::NoToken.line()); } @@ -690,7 +714,7 @@ mod tests { let t = unix_ms(); poller.tick(t); poller.tick(t + 16 * MIN); - assert!(poller.tick(t + 32 * MIN).is_none()); + poller.tick(t + 32 * MIN); let lines = lines.lock().unwrap(); assert_eq!(lines.len(), 3, "one line per change of state"); let seen = format!("{lines:?} {:?}", cell.snapshot()); @@ -727,16 +751,202 @@ mod tests { ); } + /// A token the test can rotate. + struct Rotating(Arc>>); + impl TokenSource for Rotating { + fn read(&self) -> Option> { + self.0.lock().unwrap().clone().map(Zeroizing::new) + } + } + #[test] - fn run_polls_until_the_endpoint_refuses_then_returns() { - let (url, seen) = serve(vec![http("403 Forbidden", "", "{}")]); + fn a_refused_token_is_not_sent_again_until_it_changes() { + let (url, seen) = serve(vec![ + http("403 Forbidden", "", "{}"), + http("200 OK", "", &body("null", &reading(35.0))), + ]); + let token = Arc::new(Mutex::new(Some("first".to_string()))); + let lines = Arc::new(Mutex::new(Vec::new())); + let sink = lines.clone(); let cell = Arc::new(QuotaCell::new()); - let (poller, lines) = make_poller(&cell, &url, Some(CANARY)); - poller.run(); + let mut poller = Poller::new( + cell.clone(), + Box::new(Rotating(token.clone())), + UsageClient::new(&url), + Box::new(move |line| sink.lock().unwrap().push(line.to_string())), + ); + let t = unix_ms(); + poller.tick(t); assert_eq!(requests(&seen), 1); - let lines = lines.lock().unwrap(); - assert_eq!(lines.len(), 1); - assert!(lines[0].contains("user:profile"), "{lines:?}"); + assert!(lines.lock().unwrap()[0].contains("user:profile")); + // The same token, a quarter hour on and again: the endpoint is not asked. + assert_eq!(poller.tick(t + 16 * MIN), STALE); + poller.tick(t + 32 * MIN); + assert_eq!(requests(&seen), 1); + assert_eq!(lines.lock().unwrap().len(), 1, "one state, one line"); + // A credential that vanishes does not count as a new one. + *token.lock().unwrap() = None; + poller.tick(t + 48 * MIN); + assert_eq!(requests(&seen), 1); + // A rotated credential is tried at the next due time. + *token.lock().unwrap() = Some("second".to_string()); + poller.tick(t + 64 * MIN); + assert_eq!(requests(&seen), 2); + assert!(window(&cell, "seven_day").is_some()); + assert_eq!(lines.lock().unwrap()[1], State::Polled.line()); + } + + #[test] + fn retry_after_reads_seconds_and_http_dates() { + const DATE: &str = "Wed, 21 Oct 2026 07:28:00 GMT"; + let at = chrono::DateTime::parse_from_rfc2822(DATE) + .unwrap() + .timestamp_millis() as u64; + let three_hours = 3 * 3600 * 1000; + assert_eq!(retry_after(DATE, at - three_hours), Some(secs(3 * 3600))); + assert_eq!(retry_after(DATE, at), Some(Duration::ZERO)); + assert_eq!(retry_after(DATE, at + 1), Some(Duration::ZERO), "past"); + assert_eq!(retry_after("3600", at), Some(secs(3600))); + assert_eq!(retry_after(" 3600 ", at), Some(secs(3600))); + for bad in ["", "soon", "-5", "3600.5", "1, 2", "Wed, 21 Oct 2026"] { + assert_eq!(retry_after(bad, at), None, "{bad:?}"); + } + } + + #[test] + fn keys_beyond_the_two_plan_windows_are_not_windows() { + let populated = reading(77.0); + let body = format!( + r#"{{"five_hour":{},"seven_day":{},"seven_day_opus":{populated},"seven_day_sonnet":{populated},"extra_usage":{populated}}}"#, + reading(6.0), + reading(35.0) + ); + let names: Vec<_> = parse_usage(&body) + .unwrap() + .into_iter() + .map(|quota| quota.window) + .collect(); + assert_eq!(names, ["five_hour", "seven_day"]); + } + + #[test] + fn a_redirect_is_not_followed_with_the_token() { + let (target, target_seen) = serve(vec![http("200 OK", "", &body("null", &reading(35.0)))]); + let (url, _) = serve(vec![http( + "302 Found", + &format!("Location: {target}\r\n"), + "", + )]); + let cell = Arc::new(QuotaCell::new()); + let (mut poller, lines) = make_poller(&cell, &url, Some(CANARY)); + poller.tick(unix_ms()); + assert_eq!(requests(&target_seen), 0, "the redirect target was asked"); + assert!(cell.snapshot().windows.is_empty()); + assert_eq!(lines.lock().unwrap()[0], State::Unexpected.line()); + } + + #[test] + fn a_silent_endpoint_frees_the_poll_at_its_timeout() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let url = format!("http://{}/u", listener.local_addr().unwrap()); + let holder = std::thread::spawn(move || { + // Accept, say nothing, keep the connection open. + let held = listener.accept(); + std::thread::sleep(Duration::from_secs(3)); + drop(held); + }); + let started = std::time::Instant::now(); + let answer = UsageClient::with_timeout(url, Duration::from_millis(300)) + .ask(&Zeroizing::new(CANARY.to_string())); + assert_eq!(answer, Answer::Unreachable); + assert!(started.elapsed() < Duration::from_secs(2)); + holder.join().unwrap(); + } + + #[test] + fn the_real_timeout_frees_the_thread_within_a_minute() { + assert!(TIMEOUT >= secs(1) && TIMEOUT <= secs(60), "{TIMEOUT:?}"); + } + + /// Answers one connection after `delay`, as a slow endpoint does. + fn serve_slowly(delay: Duration, reply: String) -> String { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let url = format!("http://{}/api/oauth/usage", listener.local_addr().unwrap()); + std::thread::spawn(move || { + let Ok((mut stream, _)) = listener.accept() else { + return; + }; + let mut request = vec![0u8; 4096]; + let _ = stream.read(&mut request); + std::thread::sleep(delay); + let _ = stream.write_all(reply.as_bytes()); + }); + url + } + + fn stream_event(window: &str, permille: u32, status: &str) -> PlanQuota { + PlanQuota { + status: status.to_string(), + ..stream_reading(window, permille) + } + } + + /// The stream speaks while a poll is in flight; the poll then answers + /// with a reading older than the stream's. Returns the cell after both. + fn stream_speaks_while_the_poll_is_in_flight(status: &str) -> Arc { + let url = serve_slowly( + Duration::from_millis(500), + http("200 OK", "", &body(&reading(50.0), &reading(35.0))), + ); + let cell = Arc::new(QuotaCell::new()); + let (mut poller, _) = make_poller(&cell, &url, Some(CANARY)); + std::thread::scope(|scope| { + let polling = scope.spawn(|| poller.tick(unix_ms())); + std::thread::sleep(Duration::from_millis(150)); + cell.observe([stream_event("five_hour", 1000, status)]); + polling.join().unwrap(); + }); + cell + } + + #[test] + fn a_poll_answered_late_does_not_replace_what_the_stream_saw_meanwhile() { + let cell = stream_speaks_while_the_poll_is_in_flight("allowed"); + let five = window(&cell, "five_hour").unwrap(); + assert_eq!( + (five.status.as_str(), five.utilization_permille), + ("allowed", Some(1000)) + ); + let seven = window(&cell, "seven_day").unwrap(); + assert_eq!( + (seven.status.as_str(), seven.utilization_permille), + ("polled", Some(350)), + "the window the stream did not name takes the poll's reading" + ); + } + + #[test] + fn a_slow_poll_cannot_lift_the_rejection_the_stream_reported_meanwhile() { + let cell = stream_speaks_while_the_poll_is_in_flight("rejected"); + let five = window(&cell, "five_hour").unwrap(); + assert_eq!( + (five.status.as_str(), five.utilization_permille), + ("rejected", Some(1000)) + ); + } + + #[test] + fn a_later_poll_does_not_lift_a_rejection_the_provider_reported_earlier() { + let (url, _) = serve(vec![http("200 OK", "", &body(&reading(50.0), "null"))]); + let cell = Arc::new(QuotaCell::new()); + cell.observe([stream_event("five_hour", 1000, "rejected")]); + let (mut poller, _) = make_poller(&cell, &url, Some(CANARY)); + poller.tick(unix_ms() + MIN); + let five = window(&cell, "five_hour").unwrap(); + assert_eq!( + (five.status.as_str(), five.utilization_permille), + ("rejected", Some(1000)) + ); } #[test]