hanzoai/dev
2
0
Fork 0
mirror of https://github.com/hanzoai/dev.git synced 2026-10-25 03:32:32 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(usage): compact persisted stats and harden writes

Browse Source
This commit is contained in:
James Peter
2025-09-27 11:00:58 +10:00
parent 0c36744a4c
commit 7573b40576
2 changed files with 333 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/account_usage.rs
+203 -5
View File
@@ -1,8 +1,9 @@
use std::collections::BTreeMap;
use std::fs::{self, OpenOptions};
use std::io::{ErrorKind, Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use chrono::{DateTime, Duration, Utc};
use chrono::{DateTime, Datelike, Duration, NaiveDate, TimeZone, Timelike, Utc};
use crate::protocol::RateLimitSnapshotEvent;
use fs2::FileExt;
use serde::{Deserialize, Serialize};
@@ -83,6 +84,12 @@ struct TokenWindowEntry {
tokens: TokenTotals,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
struct AggregatedUsageEntry {
period_start: DateTime<Utc>,
tokens: TokenTotals,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
struct RateLimitInfo {
#[serde(default)]
@@ -114,6 +121,12 @@ struct AccountUsageData {
#[serde(default)]
hourly_entries: Vec<TokenWindowEntry>,
#[serde(default)]
hourly_buckets: Vec<AggregatedUsageEntry>,
#[serde(default)]
daily_buckets: Vec<AggregatedUsageEntry>,
#[serde(default)]
monthly_buckets: Vec<AggregatedUsageEntry>,
#[serde(default)]
tokens_last_hour: TokenTotals,
#[serde(default)]
rate_limit: Option<RateLimitInfo>,
@@ -128,6 +141,9 @@ impl AccountUsageData {
last_updated: Utc::now(),
totals: TokenTotals::default(),
hourly_entries: Vec::new(),
hourly_buckets: Vec::new(),
daily_buckets: Vec::new(),
monthly_buckets: Vec::new(),
tokens_last_hour: TokenTotals::default(),
rate_limit: None,
}
@@ -142,6 +158,7 @@ impl AccountUsageData {
}
fn update_last_hour(&mut self, now: DateTime<Utc>) {
self.compact_usage(now);
let hourly_cutoff = now - Duration::hours(1);
let history_cutoff = now - Duration::days(HOURLY_HISTORY_DAYS);
self.hourly_entries
@@ -156,6 +173,137 @@ impl AccountUsageData {
}
self.tokens_last_hour = totals;
}
fn compact_usage(&mut self, now: DateTime<Utc>) {
let recent_cutoff = now - Duration::hours(1);
let mut rollover: BTreeMap<DateTime<Utc>, TokenTotals> = BTreeMap::new();
let mut recent: Vec<TokenWindowEntry> = Vec::new();
for entry in self.hourly_entries.drain(..) {
if entry.timestamp >= recent_cutoff {
recent.push(entry);
} else {
let bucket = truncate_to_hour(entry.timestamp);
rollover
.entry(bucket)
.or_insert_with(TokenTotals::default)
.add_totals(&entry.tokens);
}
}
self.hourly_entries = recent;
for (period_start, tokens) in rollover {
add_to_bucket(&mut self.hourly_buckets, period_start, tokens);
}
self.compact_hourly_buckets(now);
self.compact_daily_buckets(now);
}
fn compact_hourly_buckets(&mut self, now: DateTime<Utc>) {
if self.hourly_buckets.is_empty() {
return;
}
let current_hour = truncate_to_hour(now);
let cutoff = current_hour - Duration::hours(24);
let mut remaining: Vec<AggregatedUsageEntry> = Vec::new();
let mut daily_rollover: BTreeMap<DateTime<Utc>, TokenTotals> = BTreeMap::new();
for entry in self.hourly_buckets.drain(..) {
if entry.period_start < cutoff {
let day_key = truncate_to_day(entry.period_start);
daily_rollover
.entry(day_key)
.or_insert_with(TokenTotals::default)
.add_totals(&entry.tokens);
} else {
remaining.push(entry);
}
}
remaining.sort_by_key(|item| item.period_start);
self.hourly_buckets = remaining;
for (period_start, tokens) in daily_rollover {
add_to_bucket(&mut self.daily_buckets, period_start, tokens);
}
}
fn compact_daily_buckets(&mut self, now: DateTime<Utc>) {
if self.daily_buckets.is_empty() {
return;
}
let today = truncate_to_day(now);
let cutoff = today - Duration::days(30);
let mut remaining: Vec<AggregatedUsageEntry> = Vec::new();
let mut monthly_rollover: BTreeMap<DateTime<Utc>, TokenTotals> = BTreeMap::new();
for entry in self.daily_buckets.drain(..) {
if entry.period_start < cutoff {
let month_key = truncate_to_month(entry.period_start);
monthly_rollover
.entry(month_key)
.or_insert_with(TokenTotals::default)
.add_totals(&entry.tokens);
} else {
remaining.push(entry);
}
}
remaining.sort_by_key(|item| item.period_start);
self.daily_buckets = remaining;
for (period_start, tokens) in monthly_rollover {
add_to_bucket(&mut self.monthly_buckets, period_start, tokens);
}
}
}
fn add_to_bucket(
buckets: &mut Vec<AggregatedUsageEntry>,
period_start: DateTime<Utc>,
tokens: TokenTotals,
) {
match buckets.binary_search_by(|entry| entry.period_start.cmp(&period_start)) {
Ok(idx) => buckets[idx].tokens.add_totals(&tokens),
Err(idx) => {
buckets.insert(
idx,
AggregatedUsageEntry {
period_start,
tokens,
},
);
}
}
}
fn truncate_to_hour(ts: DateTime<Utc>) -> DateTime<Utc> {
let naive = ts.naive_utc();
let trimmed = naive
.with_minute(0)
.and_then(|dt| dt.with_second(0))
.and_then(|dt| dt.with_nanosecond(0))
.expect("valid hour truncation");
Utc.from_utc_datetime(&trimmed)
}
fn truncate_to_day(ts: DateTime<Utc>) -> DateTime<Utc> {
let date = ts.date_naive();
let start = date.and_hms_opt(0, 0, 0).expect("valid day truncation");
Utc.from_utc_datetime(&start)
}
fn truncate_to_month(ts: DateTime<Utc>) -> DateTime<Utc> {
let date = ts.date_naive();
let month_start = NaiveDate::from_ymd_opt(date.year(), date.month(), 1)
.expect("valid month truncation")
.and_hms_opt(0, 0, 0)
.expect("valid month start time");
Utc.from_utc_datetime(&month_start)
}
#[derive(Debug, Clone)]
@@ -175,6 +323,12 @@ pub struct StoredUsageEntry {
pub tokens: TokenTotals,
}
#[derive(Debug, Clone)]
pub struct StoredUsageBucket {
pub period_start: DateTime<Utc>,
pub tokens: TokenTotals,
}
#[derive(Debug, Clone)]
pub struct StoredUsageSummary {
pub account_id: String,
@@ -182,6 +336,9 @@ pub struct StoredUsageSummary {
pub totals: TokenTotals,
pub last_updated: DateTime<Utc>,
pub hourly_entries: Vec<StoredUsageEntry>,
pub hourly_buckets: Vec<StoredUsageBucket>,
pub daily_buckets: Vec<StoredUsageBucket>,
pub monthly_buckets: Vec<StoredUsageBucket>,
}
fn usage_dir(codex_home: &Path) -> PathBuf {
@@ -242,10 +399,21 @@ where
update(&mut data);
let json = serde_json::to_string_pretty(&data)?;
file.seek(SeekFrom::Start(0))?;
file.set_len(0)?;
file.write_all(json.as_bytes())?;
file.flush()?;
let tmp_path = usage_dir.join(format!("{account_id}.json.tmp"));
{
let mut tmp = OpenOptions::new()
.write(true)
.create(true)
.truncate(true)
.open(&tmp_path)?;
tmp.write_all(json.as_bytes())?;
tmp.sync_all()?;
}
if let Err(err) = fs::rename(&tmp_path, &path) {
let _ = fs::remove_file(&tmp_path);
file.unlock()?;
return Err(err);
}
file.unlock()?;
Ok(())
}
@@ -544,12 +712,42 @@ pub fn load_account_usage(
})
.collect();
let hourly_buckets = data
.hourly_buckets
.into_iter()
.map(|entry| StoredUsageBucket {
period_start: entry.period_start,
tokens: entry.tokens,
})
.collect();
let daily_buckets = data
.daily_buckets
.into_iter()
.map(|entry| StoredUsageBucket {
period_start: entry.period_start,
tokens: entry.tokens,
})
.collect();
let monthly_buckets = data
.monthly_buckets
.into_iter()
.map(|entry| StoredUsageBucket {
period_start: entry.period_start,
tokens: entry.tokens,
})
.collect();
Ok(Some(StoredUsageSummary {
account_id: data.account_id,
plan: data.plan,
totals: data.totals,
last_updated: data.last_updated,
hourly_entries,
hourly_buckets,
daily_buckets,
monthly_buckets,
}))
}
codex-rs/tui/src/chatwidget.rs
+130 -78
View File
@@ -250,7 +250,7 @@ use codex_core::protocol::RateLimitSnapshotEvent;
use codex_core::protocol::ValidationGroup;
use crate::rate_limits_view::{build_limits_view, RateLimitResetInfo, DEFAULT_GRID_CONFIG};
use codex_core::review_format::format_review_findings_block;
use chrono::{DateTime, Datelike, Duration as ChronoDuration, Local, Timelike, Utc};
use chrono::{DateTime, Datelike, Duration as ChronoDuration, Local, TimeZone, Timelike, Utc};
use crossterm::event::KeyCode;
use crossterm::event::KeyModifiers;
use ratatui::style::Stylize;
@@ -2306,22 +2306,29 @@ impl ChatWidget<'_> {
let record = self.history_record_from_cell(self.history_cells[idx].as_ref());
if let Some(record) = record {
if let HistoryMutation::Replaced { id, .. } = self
.history_state
.apply_event(HistoryEvent::Replace { index: idx, record })
{
if let Some(cell) = self.history_cells.get_mut(idx) {
if let Some(plain_cell) = cell
.as_any_mut()
.downcast_mut::<history_cell::PlainHistoryCell>()
{
let state = plain_cell.state_mut();
state.id = id;
if let Some(record_index) = self.record_index_for_cell(idx) {
if let HistoryMutation::Replaced { id, .. } = self
.history_state
.apply_event(HistoryEvent::Replace { index: record_index, record })
{
if let Some(cell) = self.history_cells.get_mut(idx) {
if let Some(plain_cell) = cell
.as_any_mut()
.downcast_mut::<history_cell::PlainHistoryCell>()
{
let state = plain_cell.state_mut();
state.id = id;
}
}
if idx < self.history_cell_ids.len() {
self.history_cell_ids[idx] = Some(id);
}
}
if idx < self.history_cell_ids.len() {
self.history_cell_ids[idx] = Some(id);
}
} else {
tracing::warn!(
"history-state mismatch: no persisted index for patch cell {}",
idx
);
}
}
@@ -3844,35 +3851,20 @@ impl ChatWidget<'_> {
- ChronoDuration::seconds(now.second() as i64)
- ChronoDuration::nanoseconds(now.nanosecond() as i64);
let mut hours: Vec<(DateTime<Local>, TokenTotals)> = (0..12)
.map(|offset| {
(
anchor - ChronoDuration::hours(offset as i64),
TokenTotals::default(),
)
let hourly_totals = Self::aggregate_hourly_totals(summary);
let series: Vec<(DateTime<Local>, TokenTotals)> = (0..12)
.map(|offset| anchor - ChronoDuration::hours(offset as i64))
.map(|dt| {
let utc_key = Self::truncate_utc_hour(dt.with_timezone(&Utc));
let totals = hourly_totals
.get(&utc_key)
.cloned()
.unwrap_or_default();
(dt, totals)
})
.collect();
let mut lookup: HashMap<i64, usize> = HashMap::new();
for (idx, (dt, _)) in hours.iter().enumerate() {
lookup.insert(dt.timestamp(), idx);
}
if let Some(summary) = summary {
for entry in &summary.hourly_entries {
let local_time = entry.timestamp.with_timezone(&Local);
let truncated = local_time
- ChronoDuration::minutes(local_time.minute() as i64)
- ChronoDuration::seconds(local_time.second() as i64)
- ChronoDuration::nanoseconds(local_time.nanosecond() as i64);
if let Some(&pos) = lookup.get(&truncated.timestamp()) {
let (_, totals) = &mut hours[pos];
Self::accumulate_token_totals(totals, &entry.tokens);
}
}
}
let max_total = hours
let max_total = series
.iter()
.map(|(_, totals)| totals.total_tokens)
.max()
@@ -3885,12 +3877,12 @@ impl ChatWidget<'_> {
)]));
let prefix = status_content_prefix();
let tokens_width = hours
let tokens_width = series
.iter()
.map(|(_, totals)| format_with_separators(totals.total_tokens).len())
.max()
.unwrap_or(0);
for (dt, totals) in hours.iter() {
for (dt, totals) in series.iter() {
let label = Self::format_hour_label(*dt);
let bar = Self::bar_segment(totals.total_tokens, max_total, WIDTH);
let tokens = format_with_separators(totals.total_tokens);
@@ -3919,27 +3911,15 @@ impl ChatWidget<'_> {
fn daily_usage_lines(summary: Option<&StoredUsageSummary>) -> Vec<RtLine<'static>> {
const WIDTH: usize = 14;
let today = Local::now().date_naive();
let mut daily: Vec<(chrono::NaiveDate, TokenTotals)> = (0..7)
.map(|offset| {
(
today - ChronoDuration::days(offset as i64),
TokenTotals::default(),
)
let day_totals = Self::aggregate_daily_totals(summary);
let daily: Vec<(chrono::NaiveDate, TokenTotals)> = (0..7)
.map(|offset| today - ChronoDuration::days(offset as i64))
.map(|day| {
let totals = day_totals.get(&day).cloned().unwrap_or_default();
(day, totals)
})
.collect();
if let Some(summary) = summary {
for entry in &summary.hourly_entries {
let entry_date = entry.timestamp.with_timezone(&Local).date_naive();
let diff = today.signed_duration_since(entry_date).num_days();
if (0..=6).contains(&diff) {
let idx = diff as usize;
let (_, totals) = &mut daily[idx];
Self::accumulate_token_totals(totals, &entry.tokens);
}
}
}
let max_total = daily
.iter()
.map(|(_, totals)| totals.total_tokens)
@@ -4022,11 +4002,17 @@ impl ChatWidget<'_> {
let mut year = today.year();
let mut month = today.month();
let month_totals = Self::aggregate_monthly_totals(summary);
let mut months: Vec<(chrono::NaiveDate, TokenTotals)> = Vec::with_capacity(MONTHS);
for _ in 0..MONTHS {
let start = chrono::NaiveDate::from_ymd_opt(year, month, 1)
.expect("valid month start");
months.push((start, TokenTotals::default()));
let key = (start.year(), start.month());
let totals = month_totals
.get(&key)
.cloned()
.unwrap_or_default();
months.push((start, totals));
if month == 1 {
month = 12;
year -= 1;
@@ -4035,22 +4021,6 @@ impl ChatWidget<'_> {
}
}
let mut lookup: HashMap<(i32, u32), usize> = HashMap::new();
for (idx, (start, _)) in months.iter().enumerate() {
lookup.insert((start.year(), start.month()), idx);
}
if let Some(summary) = summary {
for entry in &summary.hourly_entries {
let local_time = entry.timestamp.with_timezone(&Local);
let date = local_time.date_naive();
if let Some(&pos) = lookup.get(&(date.year(), date.month())) {
let (_, totals) = &mut months[pos];
Self::accumulate_token_totals(totals, &entry.tokens);
}
}
}
let max_total = months
.iter()
.map(|(_, totals)| totals.total_tokens)
@@ -4120,6 +4090,88 @@ impl ChatWidget<'_> {
)])
}
fn truncate_utc_hour(ts: DateTime<Utc>) -> DateTime<Utc> {
let naive = ts.naive_utc();
let trimmed = naive
.with_minute(0)
.and_then(|dt| dt.with_second(0))
.and_then(|dt| dt.with_nanosecond(0))
.expect("valid hour truncation");
Utc.from_utc_datetime(&trimmed)
}
fn aggregate_hourly_totals(
summary: Option<&StoredUsageSummary>,
) -> HashMap<DateTime<Utc>, TokenTotals> {
let mut totals = HashMap::new();
if let Some(summary) = summary {
for entry in &summary.hourly_entries {
let key = Self::truncate_utc_hour(entry.timestamp);
let slot = totals.entry(key).or_insert_with(TokenTotals::default);
Self::accumulate_token_totals(slot, &entry.tokens);
}
for bucket in &summary.hourly_buckets {
let slot = totals
.entry(bucket.period_start)
.or_insert_with(TokenTotals::default);
Self::accumulate_token_totals(slot, &bucket.tokens);
}
}
totals
}
fn aggregate_daily_totals(
summary: Option<&StoredUsageSummary>,
) -> HashMap<chrono::NaiveDate, TokenTotals> {
let mut totals = HashMap::new();
if let Some(summary) = summary {
for bucket in &summary.daily_buckets {
let key = bucket.period_start.date_naive();
let slot = totals.entry(key).or_insert_with(TokenTotals::default);
Self::accumulate_token_totals(slot, &bucket.tokens);
}
for bucket in &summary.hourly_buckets {
let key = bucket.period_start.date_naive();
let slot = totals.entry(key).or_insert_with(TokenTotals::default);
Self::accumulate_token_totals(slot, &bucket.tokens);
}
for entry in &summary.hourly_entries {
let key = entry.timestamp.date_naive();
let slot = totals.entry(key).or_insert_with(TokenTotals::default);
Self::accumulate_token_totals(slot, &entry.tokens);
}
}
totals
}
fn aggregate_monthly_totals(
summary: Option<&StoredUsageSummary>,
) -> HashMap<(i32, u32), TokenTotals> {
let mut totals = HashMap::new();
if let Some(summary) = summary {
let mut accumulate = |dt: DateTime<Utc>, tokens: &TokenTotals| {
let date = dt.date_naive();
let key = (date.year(), date.month());
let slot = totals.entry(key).or_insert_with(TokenTotals::default);
Self::accumulate_token_totals(slot, tokens);
};
for bucket in &summary.monthly_buckets {
accumulate(bucket.period_start, &bucket.tokens);
}
for bucket in &summary.daily_buckets {
accumulate(bucket.period_start, &bucket.tokens);
}
for bucket in &summary.hourly_buckets {
accumulate(bucket.period_start, &bucket.tokens);
}
for entry in &summary.hourly_entries {
accumulate(entry.timestamp, &entry.tokens);
}
}
totals
}
// dispatch_command() removed — command routing is handled at the App layer via AppEvent::DispatchCommand
pub(crate) fn handle_paste(&mut self, text: String) {