diff --git a/crates/claude-connector/src/backend/tests.rs b/crates/claude-connector/src/backend/tests.rs index 9f41654..1b14161 100644 --- a/crates/claude-connector/src/backend/tests.rs +++ b/crates/claude-connector/src/backend/tests.rs @@ -2160,7 +2160,7 @@ fn rate_limit_event_updates_quota_cell() { #[test] fn a_rejected_turn_still_reports_its_windows_and_a_bad_window_costs_only_itself() { - let event = r#"{"type":"rate_limit_event","rate_limit_info":{"status":"rejected","unifiedWindows":{"five_hour":{"utilization":0.9996,"resetsAt":4102444800,"status":"rejected"},"seven_day":{"resetsAt":4102531200,"status":"allowed"},"idle":{"utilization":0,"resetsAt":4102444800,"status":"allowed"},"odd":{"utilization":"high","resetsAt":1,"status":"allowed"},"negative":{"utilization":-0.5,"resetsAt":1,"status":"allowed"}}}}"#; + let event = r#"{"type":"rate_limit_event","rate_limit_info":{"status":"rejected","unifiedWindows":{"five_hour":{"utilization":0.9996,"resetsAt":4102444800,"status":"rejected"},"seven_day":{"utilization":0,"resetsAt":4102531200,"status":"allowed"},"seven_day_overage_included":{"utilization":1,"resetsAt":4102531200,"status":"rejected"},"idle":{"utilization":0,"resetsAt":4102444800,"status":"allowed"}}}}"#; let rig = rig(vec![Script::lines(vec![ init_line(MODEL, SESSION), event.to_string(), @@ -2170,12 +2170,12 @@ fn a_rejected_turn_still_reports_its_windows_and_a_bad_window_costs_only_itself( assert_eq!(kind, FailureKind::RateLimited); let mut rejected = window("five_hour", 1000, 4_102_444_800_000); rejected.status = "rejected".to_string(); - // Unknown is never zero: a window missing a field or of another type is - // skipped, not read as empty. - // A window at exactly zero use is a reading, and is kept. + // A window at exactly zero use is a reading, and is kept. Only the two + // windows pacing reads are taken: a per-model bucket that reads rejected + // must not close every lease, and no other name is a window of the plan. assert_eq!( readings(&rig), - vec![rejected, window("idle", 0, 4_102_444_800_000)] + vec![rejected, window("seven_day", 0, 4_102_531_200_000)] ); } diff --git a/crates/claude-connector/src/quota.rs b/crates/claude-connector/src/quota.rs index 0191b09..8d80a1c 100644 --- a/crates/claude-connector/src/quota.rs +++ b/crates/claude-connector/src/quota.rs @@ -11,6 +11,11 @@ use std::time::{SystemTime, UNIX_EPOCH}; use model_connector::{PlanQuota, PlanQuotaSnapshot}; +/// The window the weekly pacing reads, and the one the meter keeps fresh. +pub(crate) const WEEK: &str = "seven_day"; +/// The windows the pacing reads, and so the only ones the cell keeps. +pub(crate) const PACED: [&str; 2] = ["five_hour", WEEK]; + /// The latest observation of each plan window, shared by the backend that /// feeds it and whatever reads it out. Each window carries the time of its /// own observation: an event that names one window says nothing about the diff --git a/crates/claude-connector/src/stream.rs b/crates/claude-connector/src/stream.rs index c5e2e41..039ff75 100644 --- a/crates/claude-connector/src/stream.rs +++ b/crates/claude-connector/src/stream.rs @@ -7,7 +7,7 @@ use model_connector::{FailureKind, PlanQuota, StopReason, Usage}; use serde::Deserialize; use serde_json::Value; -use crate::quota::QuotaCell; +use crate::quota::{PACED, QuotaCell}; /// The tool the CLI gives the model to deliver an answer that must match a /// `--json-schema`. Its call is the structured output, not a tool exchange. @@ -320,12 +320,16 @@ impl<'a> StreamState<'a> { /// `rate_limit_info.unifiedWindows`: utilization as a fraction of the window, /// reset as unix seconds. A window missing any of its three fields is skipped; /// an unknown reading is never zero. +/// Only the windows pacing reads are taken ([`PACED`]): the event also names +/// per-model buckets, and one of those reading `rejected` must not close every +/// lease. fn plan_windows(info: &Value) -> Vec { let Some(windows) = info.get("unifiedWindows").and_then(Value::as_object) else { return Vec::new(); }; windows .iter() + .filter(|(window, _)| PACED.contains(&window.as_str())) .filter_map(|(window, reading)| { let utilization = reading.get("utilization")?.as_f64()?; let resets_at = reading.get("resetsAt")?.as_u64()?; @@ -392,3 +396,50 @@ fn sum_usage(models: impl Iterator) -> Usage { } usage } + +#[cfg(test)] +mod tests { + use super::*; + + fn windows(body: &str) -> Vec<(String, Option, u64)> { + let info: Value = serde_json::from_str(&format!(r#"{{"unifiedWindows":{body}}}"#)).unwrap(); + plan_windows(&info) + .into_iter() + .map(|q| (q.window, q.utilization_permille, q.resets_at_unix_ms)) + .collect() + } + + #[test] + fn only_the_five_hour_and_seven_day_windows_are_taken() { + let taken = windows( + r#"{"five_hour":{"utilization":0.5,"resetsAt":2,"status":"allowed"}, + "seven_day":{"utilization":0.25,"resetsAt":3,"status":"allowed"}, + "seven_day_opus":{"utilization":0.1,"resetsAt":4,"status":"allowed"}, + "seven_day_overage_included":{"utilization":1,"resetsAt":5,"status":"rejected"}}"#, + ); + assert_eq!( + taken, + vec![ + ("five_hour".to_string(), Some(500), 2000), + ("seven_day".to_string(), Some(250), 3000), + ] + ); + } + + #[test] + fn a_window_missing_a_field_or_of_the_wrong_type_costs_only_itself() { + // Unknown is never zero: such a window is skipped, not read as empty. + for broken in [ + r#"{"resetsAt":2,"status":"allowed"}"#, + r#"{"utilization":0.5,"status":"allowed"}"#, + r#"{"utilization":0.5,"resetsAt":2}"#, + r#"{"utilization":"high","resetsAt":2,"status":"allowed"}"#, + r#"{"utilization":-0.5,"resetsAt":2,"status":"allowed"}"#, + ] { + let both = windows(&format!( + r#"{{"five_hour":{broken},"seven_day":{{"utilization":0,"resetsAt":3,"status":"allowed"}}}}"# + )); + assert_eq!(both, vec![("seven_day".to_string(), Some(0), 3000)], "{broken}"); + } + } +}