308 lines
9.3 KiB
Rust
308 lines
9.3 KiB
Rust
//! Aggregation: events → 5-hour blocks + 7-day rolling window + per-model.
|
|
//!
|
|
//! The block algorithm is the same one ccusage uses:
|
|
//! - block_start = floor_to_hour(first event's ts)
|
|
//! - block_end = block_start + 5h
|
|
//! - new block when no previous, OR ts >= prev.block_end, OR
|
|
//! gap (ts - prev_event.ts) >= 5h
|
|
//!
|
|
//! Weekly is "past 7 days from now" — Anthropic's actual Max-plan reset day
|
|
//! is buggy and shifts, so a calendar anchor would lie.
|
|
|
|
use chrono::{DateTime, Datelike, Duration, TimeZone, Timelike, Utc};
|
|
use serde::Serialize;
|
|
|
|
use crate::jsonl::{ModelFamily, UsageEvent};
|
|
use crate::settings::Caps;
|
|
|
|
const BLOCK_HOURS: i64 = 5;
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
pub struct ModelBreakdown {
|
|
pub opus: u64,
|
|
pub sonnet: u64,
|
|
pub haiku: u64,
|
|
pub other: u64,
|
|
}
|
|
|
|
impl Default for ModelBreakdown {
|
|
fn default() -> Self {
|
|
Self { opus: 0, sonnet: 0, haiku: 0, other: 0 }
|
|
}
|
|
}
|
|
|
|
impl ModelBreakdown {
|
|
fn add(&mut self, family: ModelFamily, n: u64) {
|
|
match family {
|
|
ModelFamily::Opus => self.opus += n,
|
|
ModelFamily::Sonnet => self.sonnet += n,
|
|
ModelFamily::Haiku => self.haiku += n,
|
|
ModelFamily::Other => self.other += n,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
pub struct BlockSummary {
|
|
pub block_start: DateTime<Utc>,
|
|
pub block_end: DateTime<Utc>,
|
|
pub now: DateTime<Utc>,
|
|
pub seconds_remaining: i64,
|
|
pub total_tokens: u64,
|
|
pub by_family: ModelBreakdown,
|
|
pub message_count: u64,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
pub struct DayBucket {
|
|
/// `YYYY-MM-DD` in the user's local timezone.
|
|
pub date_local: String,
|
|
pub total_tokens: u64,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
pub struct WeeklySummary {
|
|
pub window_start: DateTime<Utc>,
|
|
pub window_end: DateTime<Utc>,
|
|
pub total_tokens: u64,
|
|
pub by_day: [DayBucket; 7],
|
|
pub by_family: ModelBreakdown,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
pub struct UsageSnapshot {
|
|
pub block: Option<BlockSummary>,
|
|
pub weekly: WeeklySummary,
|
|
pub caps: Caps,
|
|
pub generated_at: DateTime<Utc>,
|
|
}
|
|
|
|
pub fn floor_to_hour(ts: DateTime<Utc>) -> DateTime<Utc> {
|
|
Utc.with_ymd_and_hms(ts.year(), ts.month(), ts.day(), ts.hour(), 0, 0)
|
|
.single()
|
|
.unwrap_or(ts)
|
|
}
|
|
|
|
/// Walk a chronologically-sorted slice of events and produce one block per
|
|
/// detected gap or 5h-end boundary. Events MUST be sorted by `ts` ascending.
|
|
pub fn compute_blocks(events: &[UsageEvent]) -> Vec<BlockSummary> {
|
|
let mut out: Vec<BlockSummary> = Vec::new();
|
|
if events.is_empty() {
|
|
return out;
|
|
}
|
|
|
|
let block_window = Duration::hours(BLOCK_HOURS);
|
|
let now = Utc::now();
|
|
|
|
// Working accumulator for the current block.
|
|
struct Acc {
|
|
start: DateTime<Utc>,
|
|
end: DateTime<Utc>,
|
|
total: u64,
|
|
breakdown: ModelBreakdown,
|
|
msgs: u64,
|
|
}
|
|
|
|
let close = |a: Acc, now: DateTime<Utc>| -> BlockSummary {
|
|
BlockSummary {
|
|
block_start: a.start,
|
|
block_end: a.end,
|
|
now,
|
|
seconds_remaining: (a.end - now).num_seconds(),
|
|
total_tokens: a.total,
|
|
by_family: a.breakdown,
|
|
message_count: a.msgs,
|
|
}
|
|
};
|
|
|
|
let mut acc: Option<Acc> = None;
|
|
let mut prev_ts: Option<DateTime<Utc>> = None;
|
|
|
|
for ev in events {
|
|
let starts_new_block = match (&acc, prev_ts) {
|
|
(None, _) => true,
|
|
(Some(a), Some(prev)) => ev.ts >= a.end || (ev.ts - prev) >= block_window,
|
|
(Some(a), None) => ev.ts >= a.end,
|
|
};
|
|
|
|
if starts_new_block {
|
|
if let Some(a) = acc.take() {
|
|
out.push(close(a, now));
|
|
}
|
|
let start = floor_to_hour(ev.ts);
|
|
acc = Some(Acc {
|
|
start,
|
|
end: start + block_window,
|
|
total: 0,
|
|
breakdown: ModelBreakdown::default(),
|
|
msgs: 0,
|
|
});
|
|
}
|
|
|
|
let a = acc.as_mut().expect("acc just initialized");
|
|
a.total += ev.total;
|
|
a.breakdown.add(ev.model_family, ev.total);
|
|
a.msgs += 1;
|
|
prev_ts = Some(ev.ts);
|
|
}
|
|
|
|
if let Some(a) = acc {
|
|
out.push(close(a, now));
|
|
}
|
|
out
|
|
}
|
|
|
|
/// Pick the block that contains `now`. Falls back to the most recent block
|
|
/// whose end is in the future. Returns `None` if there are no blocks.
|
|
pub fn active_block(blocks: &[BlockSummary], now: DateTime<Utc>) -> Option<BlockSummary> {
|
|
blocks
|
|
.iter()
|
|
.rev()
|
|
.find(|b| now >= b.block_start && now < b.block_end)
|
|
.or_else(|| blocks.iter().rev().find(|b| b.block_end > now))
|
|
.cloned()
|
|
}
|
|
|
|
/// Rolling 7-day window from `now`. `by_day[0]` is the oldest day,
|
|
/// `by_day[6]` is today. Day boundaries are in the user's local timezone.
|
|
pub fn compute_weekly(events: &[UsageEvent], now: DateTime<Utc>) -> WeeklySummary {
|
|
let window_start = now - Duration::days(7);
|
|
let mut total: u64 = 0;
|
|
let mut breakdown = ModelBreakdown::default();
|
|
|
|
let today_local = chrono::Local::now().date_naive();
|
|
let mut buckets: [DayBucket; 7] = std::array::from_fn(|i| {
|
|
let date = today_local - chrono::Duration::days(6 - i as i64);
|
|
DayBucket {
|
|
date_local: date.format("%Y-%m-%d").to_string(),
|
|
total_tokens: 0,
|
|
}
|
|
});
|
|
|
|
for ev in events {
|
|
if ev.ts < window_start {
|
|
continue;
|
|
}
|
|
total += ev.total;
|
|
breakdown.add(ev.model_family, ev.total);
|
|
|
|
let ev_local = ev.ts.with_timezone(&chrono::Local).date_naive();
|
|
let day_diff = (today_local - ev_local).num_days();
|
|
if (0..=6).contains(&day_diff) {
|
|
let idx = (6 - day_diff) as usize;
|
|
buckets[idx].total_tokens += ev.total;
|
|
}
|
|
}
|
|
|
|
WeeklySummary {
|
|
window_start,
|
|
window_end: now,
|
|
total_tokens: total,
|
|
by_day: buckets,
|
|
by_family: breakdown,
|
|
}
|
|
}
|
|
|
|
/// One-shot: events → snapshot for the frontend.
|
|
/// `events` must be sorted by `ts` ascending.
|
|
pub fn build_snapshot(
|
|
events: &[UsageEvent],
|
|
caps: &Caps,
|
|
now: DateTime<Utc>,
|
|
) -> UsageSnapshot {
|
|
let blocks = compute_blocks(events);
|
|
let block = active_block(&blocks, now);
|
|
let weekly = compute_weekly(events, now);
|
|
UsageSnapshot {
|
|
block,
|
|
weekly,
|
|
caps: caps.clone(),
|
|
generated_at: now,
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::jsonl::{ModelFamily, UsageEvent};
|
|
use std::path::PathBuf;
|
|
|
|
fn ev(ts: &str, total: u64, family: ModelFamily) -> UsageEvent {
|
|
let ts: DateTime<Utc> = ts.parse().unwrap();
|
|
UsageEvent {
|
|
ts,
|
|
model_raw: "test".into(),
|
|
model_label: "Test".into(),
|
|
model_family: family,
|
|
input: 0,
|
|
output: 0,
|
|
cache_create: 0,
|
|
cache_read: 0,
|
|
total,
|
|
dedupe_key: format!("k-{}", ts.timestamp_nanos_opt().unwrap_or(0)),
|
|
source_file: PathBuf::new(),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn block_starts_at_hour_floor() {
|
|
let events = vec![ev("2026-05-08T22:37:21Z", 100, ModelFamily::Opus)];
|
|
let blocks = compute_blocks(&events);
|
|
assert_eq!(blocks.len(), 1);
|
|
assert_eq!(blocks[0].block_start.hour(), 22);
|
|
assert_eq!(blocks[0].block_start.minute(), 0);
|
|
assert_eq!((blocks[0].block_end - blocks[0].block_start).num_hours(), 5);
|
|
assert_eq!(blocks[0].total_tokens, 100);
|
|
}
|
|
|
|
#[test]
|
|
fn five_hour_gap_starts_new_block() {
|
|
let events = vec![
|
|
ev("2026-05-08T10:00:00Z", 100, ModelFamily::Opus),
|
|
ev("2026-05-08T15:30:00Z", 200, ModelFamily::Sonnet), // 5h30m gap
|
|
];
|
|
let blocks = compute_blocks(&events);
|
|
assert_eq!(blocks.len(), 2);
|
|
assert_eq!(blocks[0].total_tokens, 100);
|
|
assert_eq!(blocks[1].total_tokens, 200);
|
|
assert_eq!(blocks[1].by_family.sonnet, 200);
|
|
}
|
|
|
|
#[test]
|
|
fn within_block_aggregates() {
|
|
let events = vec![
|
|
ev("2026-05-08T10:00:00Z", 100, ModelFamily::Opus),
|
|
ev("2026-05-08T11:00:00Z", 200, ModelFamily::Opus),
|
|
ev("2026-05-08T12:30:00Z", 300, ModelFamily::Sonnet),
|
|
];
|
|
let blocks = compute_blocks(&events);
|
|
assert_eq!(blocks.len(), 1);
|
|
assert_eq!(blocks[0].total_tokens, 600);
|
|
assert_eq!(blocks[0].by_family.opus, 300);
|
|
assert_eq!(blocks[0].by_family.sonnet, 300);
|
|
assert_eq!(blocks[0].message_count, 3);
|
|
}
|
|
|
|
#[test]
|
|
fn block_end_starts_new_block() {
|
|
// Gap < 5h but the previous block_end was reached.
|
|
let events = vec![
|
|
ev("2026-05-08T10:00:00Z", 100, ModelFamily::Opus), // block_end = 15:00
|
|
ev("2026-05-08T15:30:00Z", 200, ModelFamily::Opus), // > 15:00 → new block
|
|
];
|
|
let blocks = compute_blocks(&events);
|
|
assert_eq!(blocks.len(), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn weekly_buckets_today_oldest_first() {
|
|
// We can't pin "today" in tests deterministically, but we can at least
|
|
// assert the array shape and that a recent event lands in the last bucket.
|
|
let now = Utc::now();
|
|
let events = vec![ev(&now.to_rfc3339(), 500, ModelFamily::Opus)];
|
|
let w = compute_weekly(&events, now);
|
|
assert_eq!(w.by_day.len(), 7);
|
|
assert_eq!(w.total_tokens, 500);
|
|
assert_eq!(w.by_day[6].total_tokens, 500);
|
|
}
|
|
}
|