diff --git a/crates/claude-connector/src/meter.rs b/crates/claude-connector/src/meter.rs index 2f572a2..51ef749 100644 --- a/crates/claude-connector/src/meter.rs +++ b/crates/claude-connector/src/meter.rs @@ -18,7 +18,7 @@ use std::sync::Arc; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use model_connector::Usage; +use model_connector::{PlanQuota, Usage}; use crate::backend::ClaudeCodeBackend; use crate::quota::{EXPIRED, QuotaCell, WEEK}; @@ -42,10 +42,16 @@ impl Probe for Arc { } } -/// Why the meter is waiting; logged once per change. +/// Why the meter is waiting. A refresh is logged every time, with the reading +/// it saw; the other states are logged once per change. #[derive(Debug, Clone, PartialEq, Eq)] enum State { - Read, + /// The week window as the refresh left it; a window with no reading is + /// `None`. + Read { + permille: Option, + resets_at_unix_ms: u64, + }, Idle, /// The call failed; the kind is the failure's kind, never its message. Failed(String), @@ -55,7 +61,15 @@ enum State { impl State { fn line(&self) -> String { match self { - Self::Read => "meter: weekly reading refreshed".to_string(), + Self::Read { + permille, + resets_at_unix_ms, + } => { + let permille = permille.map_or("unknown".to_string(), |p| p.to_string()); + format!( + "meter: weekly reading refreshed, permille {permille} resets_at_unix_ms {resets_at_unix_ms}" + ) + } Self::Idle => "meter: no lease asked lately, making no call".to_string(), Self::Failed(kind) => format!("meter: the call failed ({kind}), backing off"), Self::Unread => { @@ -105,9 +119,12 @@ impl Meter { // reports its windows and then errors. Only the cell says whether the // week was read. let called = self.probe.call(); - let state = if self.refreshed_since(now_ms) { + let state = if let Some(week) = self.week_refreshed_since(now_ms) { self.misses = 0; - State::Read + State::Read { + permille: week.utilization_permille, + resets_at_unix_ms: week.resets_at_unix_ms, + } } else { self.misses = self.misses.saturating_add(1); match called { @@ -139,16 +156,18 @@ impl Meter { stale_at.max(self.not_before_ms) } - fn refreshed_since(&self, since_ms: u64) -> bool { + /// The week window if it was observed at or after `since_ms`, from one + /// snapshot, so the logged value is the one that decided the refresh. + fn week_refreshed_since(&self, since_ms: u64) -> Option { self.cell .snapshot() .windows - .iter() - .any(|quota| quota.window == WEEK && quota.observed_at_unix_ms >= since_ms) + .into_iter() + .find(|quota| quota.window == WEEK && quota.observed_at_unix_ms >= since_ms) } fn report(&mut self, state: State) { - if self.last_state.as_ref() != Some(&state) { + if matches!(state, State::Read { .. }) || self.last_state.as_ref() != Some(&state) { (self.log)(&state.line()); self.last_state = Some(state); } @@ -181,7 +200,8 @@ fn unix_ms() -> u64 { } /// Starts the meter thread. It writes only through the calls it makes, and -/// logs one line to stderr when its state changes. From here a lease asked of +/// logs one line to stderr per refresh of the weekly reading and one per +/// change of any other state. From here a lease asked of /// a cell with no weekly reading waits for the meter's first read. pub(crate) fn start( cell: Arc, @@ -202,7 +222,6 @@ pub(crate) fn start( mod tests { use super::*; use crate::quota::Asking; - use model_connector::PlanQuota; use std::sync::Mutex; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -227,6 +246,10 @@ mod tests { reading_at: Option, status: String, fails_with: Option, + /// The used thousandths the reading reports; 60 when unset. + permille: Option, + /// The reset the reading reports; `RESET` when unset. + resets_at: Option, } /// A call that answers as scripted and feeds the cell the way the stream @@ -242,7 +265,10 @@ mod tests { self.made.fetch_add(1, Ordering::SeqCst); let script = self.script.lock().unwrap().clone(); if let Some(at) = script.reading_at { - self.cell.observe_at(at, reading(&script.status)); + let mut week = reading(&script.status); + week[0].utilization_permille = Some(script.permille.unwrap_or(60)); + week[0].resets_at_unix_ms = script.resets_at.unwrap_or(RESET); + self.cell.observe_at(at, week); } if let Some(kind) = script.fails_with { return Err(kind); @@ -299,6 +325,7 @@ mod tests { reading_at: Some(at), status: "allowed".to_string(), fails_with: None, + ..Script::default() }; } fn reads_nothing(&self) { @@ -398,7 +425,11 @@ mod tests { rig.reads(at); rig.meter.tick(at); } - assert_eq!(rig.calls(), 17, "one read per period, the lease never lapsing"); + assert_eq!( + rig.calls(), + 17, + "one read per period, the lease never lapsing" + ); } #[test] @@ -477,10 +508,16 @@ mod tests { reading_at: Some(NOW), status: "rejected".to_string(), fails_with: Some("RateLimited".to_string()), + ..Script::default() }; assert_eq!(rig.meter.tick(NOW), STALE, "the next read is a period away"); assert_eq!(rig.meter.misses, 0); - assert_eq!(rig.lines(), vec![State::Read.line()]); + assert_eq!( + rig.lines(), + vec![format!( + "meter: weekly reading refreshed, permille 60 resets_at_unix_ms {RESET}" + )] + ); rig.reads(NOW + 15 * MIN); rig.meter.tick(NOW + 15 * MIN); assert_eq!(rig.calls(), 2, "a rejected reading is kept fresh"); @@ -495,13 +532,83 @@ mod tests { assert_eq!(rig.meter.misses, 1); } + #[test] + fn a_refresh_logs_the_reading_and_an_unchanged_value_logs_again() { + let mut rig = rig(); + rig.reads(NOW); + rig.meter.tick(NOW); + rig.reads(NOW + 15 * MIN); + rig.meter.tick(NOW + 15 * MIN); + let line = + format!("meter: weekly reading refreshed, permille 60 resets_at_unix_ms {RESET}"); + assert_eq!(rig.lines(), vec![line.clone(), line]); + } + + #[test] + fn the_logged_value_follows_the_reading_and_an_expired_week_is_unknown() { + let mut rig = rig(); + for (at, permille) in [(NOW, 60), (NOW + 15 * MIN, 75)] { + *rig.script.lock().unwrap() = Script { + reading_at: Some(at), + status: "allowed".to_string(), + permille: Some(permille), + resets_at: Some(RESET + permille as u64), + ..Script::default() + }; + rig.meter.tick(at); + } + let mut expired = self::rig(); + let now = unix_ms(); + *expired.script.lock().unwrap() = Script { + reading_at: Some(now), + status: "allowed".to_string(), + resets_at: Some(1), + ..Script::default() + }; + expired.meter.tick(now); + assert_eq!( + expired.lines(), + vec!["meter: weekly reading refreshed, permille unknown resets_at_unix_ms 1"] + ); + assert_eq!( + rig.lines(), + vec![ + format!( + "meter: weekly reading refreshed, permille 60 resets_at_unix_ms {}", + RESET + 60 + ), + format!( + "meter: weekly reading refreshed, permille 75 resets_at_unix_ms {}", + RESET + 75 + ), + ] + ); + } + + #[test] + fn idle_logs_once_per_change_however_often_it_repeats() { + let mut rig = rig(); + rig.lease = None; + for step in 0..3 { + rig.meter.tick(NOW + 60 * MIN * (step + 1)); + } + assert_eq!(rig.lines(), vec![State::Idle.line()]); + } + #[test] fn the_log_says_what_the_meter_is_doing_in_words() { assert_eq!( State::Idle.line(), "meter: no lease asked lately, making no call" ); - assert_eq!(State::Read.line(), "meter: weekly reading refreshed"); + assert_eq!( + State::Read { + permille: None, + resets_at_unix_ms: 7 + } + .line(), + "meter: weekly reading refreshed, permille unknown resets_at_unix_ms 7" + ); assert_eq!( State::Unread.line(), "meter: the call completed without a weekly reading, backing off" @@ -527,6 +634,7 @@ mod tests { reading_at: Some(unix_ms() + 60 * MIN), status: "allowed".to_string(), fails_with: None, + ..Script::default() })), }; start(cell.clone(), Box::new(probe)).unwrap();