feat(usage): track cloud quotas and expose :limits
Acceptance Criteria:\n- header shows hourly/weekly usage with colored thresholds\n- :limits command prints persisted usage data and quotas\n- token usage survives restarts and emits 80%/95% toasts Test Notes:\n- cargo test -p owlen-core usage
This commit is contained in:
332
crates/owlen-core/src/usage.rs
Normal file
332
crates/owlen-core/src/usage.rs
Normal file
@@ -0,0 +1,332 @@
|
||||
use crate::{Error, Result};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
use tokio::fs;
|
||||
|
||||
const LEDGER_VERSION: u32 = 1;
|
||||
const SECONDS_PER_HOUR: i64 = 60 * 60;
|
||||
const SECONDS_PER_WEEK: i64 = 7 * 24 * 60 * 60;
|
||||
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
struct UsageRecord {
|
||||
timestamp: i64,
|
||||
prompt_tokens: u32,
|
||||
completion_tokens: u32,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Deserialize)]
|
||||
struct LedgerFile {
|
||||
version: u32,
|
||||
providers: HashMap<String, VecDeque<UsageRecord>>,
|
||||
}
|
||||
|
||||
impl Default for LedgerFile {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
version: LEDGER_VERSION,
|
||||
providers: HashMap::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct UsageLedger {
|
||||
path: PathBuf,
|
||||
providers: HashMap<String, VecDeque<UsageRecord>>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct UsageQuota {
|
||||
pub hourly_quota_tokens: Option<u64>,
|
||||
pub weekly_quota_tokens: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
|
||||
pub enum UsageWindow {
|
||||
Hour,
|
||||
Week,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
|
||||
pub enum UsageBand {
|
||||
Normal = 0,
|
||||
Warning = 1,
|
||||
Critical = 2,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct WindowMetrics {
|
||||
pub prompt_tokens: u64,
|
||||
pub completion_tokens: u64,
|
||||
pub total_tokens: u64,
|
||||
pub quota_tokens: Option<u64>,
|
||||
}
|
||||
|
||||
impl WindowMetrics {
|
||||
pub fn percent_of_quota(&self) -> Option<f64> {
|
||||
let quota = self.quota_tokens?;
|
||||
if quota == 0 {
|
||||
return None;
|
||||
}
|
||||
Some(self.total_tokens as f64 / quota as f64)
|
||||
}
|
||||
|
||||
pub fn band(&self) -> UsageBand {
|
||||
match self.percent_of_quota() {
|
||||
Some(p) if p >= 0.95_f64 => UsageBand::Critical,
|
||||
Some(p) if p >= 0.80_f64 => UsageBand::Warning,
|
||||
_ => UsageBand::Normal,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct UsageSnapshot {
|
||||
pub provider: String,
|
||||
pub hourly: WindowMetrics,
|
||||
pub weekly: WindowMetrics,
|
||||
pub last_updated: Option<SystemTime>,
|
||||
}
|
||||
|
||||
impl UsageSnapshot {
|
||||
pub fn window(&self, window: UsageWindow) -> &WindowMetrics {
|
||||
match window {
|
||||
UsageWindow::Hour => &self.hourly,
|
||||
UsageWindow::Week => &self.weekly,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl UsageLedger {
|
||||
pub fn empty(path: PathBuf) -> Self {
|
||||
Self {
|
||||
path,
|
||||
providers: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn load_or_default(path: PathBuf) -> Result<Self> {
|
||||
if !path.exists() {
|
||||
return Ok(Self {
|
||||
path,
|
||||
providers: HashMap::new(),
|
||||
});
|
||||
}
|
||||
|
||||
let contents = fs::read_to_string(&path)
|
||||
.await
|
||||
.map_err(|err| Error::Storage(format!("Failed to read usage ledger: {err}")))?;
|
||||
|
||||
let file: LedgerFile = match serde_json::from_str(&contents) {
|
||||
Ok(file) => file,
|
||||
Err(err) => {
|
||||
return Err(Error::Storage(format!(
|
||||
"Failed to parse usage ledger at {}: {err}",
|
||||
path.display()
|
||||
)));
|
||||
}
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
path,
|
||||
providers: file.providers,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn persist(&self) -> Result<()> {
|
||||
if let Some(parent) = self.path.parent() {
|
||||
fs::create_dir_all(parent)
|
||||
.await
|
||||
.map_err(|err| Error::Storage(format!("Failed to create data directory: {err}")))?;
|
||||
}
|
||||
|
||||
let serialized = self.serialize()?;
|
||||
|
||||
fs::write(&self.path, serialized)
|
||||
.await
|
||||
.map_err(|err| Error::Storage(format!("Failed to write usage ledger: {err}")))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn record(
|
||||
&mut self,
|
||||
provider: &str,
|
||||
usage: &crate::types::TokenUsage,
|
||||
timestamp: SystemTime,
|
||||
) {
|
||||
let total_tokens = usage.total_tokens;
|
||||
if total_tokens == 0 {
|
||||
return;
|
||||
}
|
||||
|
||||
let ts = match timestamp.duration_since(UNIX_EPOCH) {
|
||||
Ok(duration) => duration.as_secs() as i64,
|
||||
Err(_) => 0,
|
||||
};
|
||||
|
||||
let entry = self
|
||||
.providers
|
||||
.entry(provider.to_string())
|
||||
.or_insert_with(VecDeque::new);
|
||||
|
||||
entry.push_back(UsageRecord {
|
||||
timestamp: ts,
|
||||
prompt_tokens: usage.prompt_tokens,
|
||||
completion_tokens: usage.completion_tokens,
|
||||
});
|
||||
|
||||
self.prune_old(provider, ts);
|
||||
}
|
||||
|
||||
pub fn provider_keys(&self) -> impl Iterator<Item = &String> {
|
||||
self.providers.keys()
|
||||
}
|
||||
|
||||
pub fn serialize(&self) -> Result<String> {
|
||||
let file = LedgerFile {
|
||||
version: LEDGER_VERSION,
|
||||
providers: self.providers.clone(),
|
||||
};
|
||||
|
||||
serde_json::to_string_pretty(&file)
|
||||
.map_err(|err| Error::Storage(format!("Failed to serialize usage ledger: {err}")))
|
||||
}
|
||||
|
||||
pub fn path(&self) -> &Path {
|
||||
&self.path
|
||||
}
|
||||
|
||||
pub fn snapshot(&self, provider: &str, quotas: UsageQuota, now: SystemTime) -> UsageSnapshot {
|
||||
let now_secs = now
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_else(|_| Duration::from_secs(0))
|
||||
.as_secs() as i64;
|
||||
|
||||
let mut snapshot = UsageSnapshot {
|
||||
provider: provider.to_string(),
|
||||
hourly: WindowMetrics {
|
||||
quota_tokens: quotas.hourly_quota_tokens,
|
||||
..Default::default()
|
||||
},
|
||||
weekly: WindowMetrics {
|
||||
quota_tokens: quotas.weekly_quota_tokens,
|
||||
..Default::default()
|
||||
},
|
||||
last_updated: None,
|
||||
};
|
||||
|
||||
if let Some(records) = self.providers.get(provider) {
|
||||
for record in records {
|
||||
if now_secs - record.timestamp <= SECONDS_PER_HOUR {
|
||||
snapshot.hourly.prompt_tokens += record.prompt_tokens as u64;
|
||||
snapshot.hourly.completion_tokens += record.completion_tokens as u64;
|
||||
}
|
||||
|
||||
if now_secs - record.timestamp <= SECONDS_PER_WEEK {
|
||||
snapshot.weekly.prompt_tokens += record.prompt_tokens as u64;
|
||||
snapshot.weekly.completion_tokens += record.completion_tokens as u64;
|
||||
}
|
||||
}
|
||||
|
||||
snapshot.hourly.total_tokens =
|
||||
snapshot.hourly.prompt_tokens + snapshot.hourly.completion_tokens;
|
||||
snapshot.weekly.total_tokens =
|
||||
snapshot.weekly.prompt_tokens + snapshot.weekly.completion_tokens;
|
||||
|
||||
snapshot.last_updated = records.back().and_then(|record| {
|
||||
UNIX_EPOCH.checked_add(Duration::from_secs(record.timestamp as u64))
|
||||
});
|
||||
}
|
||||
|
||||
snapshot
|
||||
}
|
||||
|
||||
pub fn prune_old(&mut self, provider: &str, now_secs: i64) {
|
||||
if let Some(records) = self.providers.get_mut(provider) {
|
||||
while let Some(front) = records.front() {
|
||||
if now_secs - front.timestamp > SECONDS_PER_WEEK {
|
||||
records.pop_front();
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn prune_all(&mut self, now: SystemTime) {
|
||||
let now_secs = now
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_else(|_| Duration::from_secs(0))
|
||||
.as_secs() as i64;
|
||||
let provider_keys: Vec<String> = self.providers.keys().cloned().collect();
|
||||
for provider in provider_keys {
|
||||
self.prune_old(&provider, now_secs);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::types::TokenUsage;
|
||||
use std::time::{Duration, UNIX_EPOCH};
|
||||
use tempfile::tempdir;
|
||||
|
||||
fn make_usage(prompt: u32, completion: u32) -> TokenUsage {
|
||||
TokenUsage {
|
||||
prompt_tokens: prompt,
|
||||
completion_tokens: completion,
|
||||
total_tokens: prompt.saturating_add(completion),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn records_and_summarizes_usage() {
|
||||
let temp = tempdir().expect("tempdir");
|
||||
let path = temp.path().join("ledger.json");
|
||||
let mut ledger = UsageLedger::empty(path);
|
||||
|
||||
let usage = make_usage(40, 10);
|
||||
let timestamp = UNIX_EPOCH + Duration::from_secs(1);
|
||||
ledger.record("ollama_cloud", &usage, timestamp);
|
||||
|
||||
let quotas = UsageQuota {
|
||||
hourly_quota_tokens: Some(100),
|
||||
weekly_quota_tokens: Some(1000),
|
||||
};
|
||||
|
||||
let snapshot = ledger.snapshot("ollama_cloud", quotas, UNIX_EPOCH + Duration::from_secs(2));
|
||||
|
||||
assert_eq!(snapshot.hourly.total_tokens, 50);
|
||||
assert_eq!(snapshot.weekly.total_tokens, 50);
|
||||
assert_eq!(snapshot.hourly.quota_tokens, Some(100));
|
||||
assert_eq!(snapshot.weekly.quota_tokens, Some(1000));
|
||||
assert_eq!(snapshot.hourly.band(), UsageBand::Normal);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prunes_records_outside_week() {
|
||||
let temp = tempdir().expect("tempdir");
|
||||
let path = temp.path().join("ledger.json");
|
||||
let mut ledger = UsageLedger::empty(path);
|
||||
|
||||
let old_usage = make_usage(30, 5);
|
||||
let recent_usage = make_usage(20, 5);
|
||||
|
||||
let base = UNIX_EPOCH;
|
||||
ledger.record("ollama_cloud", &old_usage, base);
|
||||
|
||||
// Advance beyond a week for the second record.
|
||||
let later = UNIX_EPOCH + Duration::from_secs(SECONDS_PER_WEEK as u64 + 120);
|
||||
ledger.record("ollama_cloud", &recent_usage, later);
|
||||
|
||||
let quotas = UsageQuota::default();
|
||||
let snapshot = ledger.snapshot("ollama_cloud", quotas, later);
|
||||
|
||||
assert_eq!(snapshot.hourly.total_tokens, 25);
|
||||
assert_eq!(snapshot.weekly.total_tokens, 25);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user