From cf1fca4e5724dbc4012f4affc925aff02ecc5e37 Mon Sep 17 00:00:00 2001 From: Michelle Tilley Date: Wed, 8 Jul 2026 11:21:53 -0700 Subject: [PATCH 1/6] Add cached UsageSnapshot --- .../20260707000000_create_ai_usage.sql | 10 ++ crates/atuin-ai/src/commands/inline.rs | 35 +++++- crates/atuin-ai/src/driver.rs | 40 +++++- crates/atuin-ai/src/lib.rs | 1 + crates/atuin-ai/src/session.rs | 19 +++ crates/atuin-ai/src/store.rs | 49 ++++++++ crates/atuin-ai/src/stream.rs | 13 +- crates/atuin-ai/src/usage.rs | 115 ++++++++++++++++++ 8 files changed, 277 insertions(+), 5 deletions(-) create mode 100644 crates/atuin-ai/migrations/20260707000000_create_ai_usage.sql create mode 100644 crates/atuin-ai/src/usage.rs diff --git a/crates/atuin-ai/migrations/20260707000000_create_ai_usage.sql b/crates/atuin-ai/migrations/20260707000000_create_ai_usage.sql new file mode 100644 index 00000000000..37e02320ef9 --- /dev/null +++ b/crates/atuin-ai/migrations/20260707000000_create_ai_usage.sql @@ -0,0 +1,10 @@ +-- Cache of the server-reported credit usage snapshot, keyed by a hash of +-- the auth token (the client never learns its hub user id). One row per +-- key; snapshot is the JSON `credits` object from the hub. +CREATE TABLE IF NOT EXISTS usage ( + user_key TEXT NOT NULL, + snapshot TEXT NOT NULL, + updated_at INTEGER NOT NULL, + + PRIMARY KEY (user_key) +); diff --git a/crates/atuin-ai/src/commands/inline.rs b/crates/atuin-ai/src/commands/inline.rs index 2701744b085..17093218095 100644 --- a/crates/atuin-ai/src/commands/inline.rs +++ b/crates/atuin-ai/src/commands/inline.rs @@ -163,6 +163,23 @@ async fn run_inline_tui( .await .context("failed to open AI session database")?; + // Cached usage renders immediately; a background fetch (spawned below, + // once the event channel exists) replaces it unless it's fresh. + let usage_key = crate::usage::cache_key(&ctx.token); + let (cached_usage, usage_is_fresh) = match service.get_cached_usage(&usage_key).await { + Ok(Some((json, updated_at))) => { + let snapshot = serde_json::from_str::(&json).ok(); + let age = time::OffsetDateTime::now_utc().unix_timestamp() - updated_at; + let fresh = snapshot.is_some() && age < crate::usage::REFRESH_AFTER.as_secs() as i64; + (snapshot, fresh) + } + Ok(None) => (None, false), + Err(e) => { + debug!("failed to read usage cache: {e}"); + (None, false) + } + }; + let cwd = std::env::current_dir() .ok() .map(|p| p.to_string_lossy().into_owned()); @@ -266,7 +283,8 @@ async fn run_inline_tui( let skill_registry = crate::skills::SkillRegistry::discover(project_root.as_deref()).await; // ─── Build initial ViewState from FSM ─────────────────────── - let initial_view = build_view_state(&fsm, in_git_project, &skill_registry); + let mut initial_view = build_view_state(&fsm, in_git_project, &skill_registry); + initial_view.usage = cached_usage; // ─── Build IoContext ──────────────────────────────────────── let io = IoContext { @@ -288,6 +306,20 @@ async fn run_inline_tui( // Wrap sender for components: they send AiTuiEvent, we wrap it let tui_tx = DriverEventSender(tx.clone()); + if !usage_is_fresh { + let endpoint = ctx.endpoint.clone(); + let token = ctx.token.clone(); + let usage_tx = tx.clone(); + tokio::spawn(async move { + match crate::usage::fetch_usage(&endpoint, &token).await { + Ok(snapshot) => { + let _ = usage_tx.send(DriverEvent::Usage(snapshot)); + } + Err(e) => debug!("background usage fetch failed: {e}"), + } + }); + } + println!(); if let Some(prompt) = initial_prompt { @@ -412,6 +444,7 @@ fn build_view_state( archived_events, model_picker: fsm.ctx.model_picker.clone(), model: fsm.ctx.model.clone(), + usage: None, turns, has_command, committed_turn_count: 0, diff --git a/crates/atuin-ai/src/driver.rs b/crates/atuin-ai/src/driver.rs index 13ea6c9b548..f3a7cfbf6ab 100644 --- a/crates/atuin-ai/src/driver.rs +++ b/crates/atuin-ai/src/driver.rs @@ -43,6 +43,9 @@ pub(crate) enum DriverEvent { Tui(AiTuiEvent), /// Internal FSM event (from spawned stream/tool tasks) Fsm(Event), + /// Fresh credit-usage snapshot (from the done event or a background + /// fetch). Handled by the driver directly — the FSM never sees it. + Usage(crate::usage::UsageSnapshot), } // ============================================================================ @@ -87,6 +90,9 @@ pub(crate) struct ViewState { pub model_picker: Option, /// Model alias currently in effect (`None` = server default). pub model: Option, + /// Latest known credit usage: cached at startup, refreshed from the + /// done event and background fetches. + pub usage: Option, // ─── Pre-computed for rendering ──────────────────────────── pub turns: Vec, @@ -187,6 +193,10 @@ pub(crate) fn run_driver( tracing::trace!(?tui_event, state = ?fsm.state, "TUI event"); translate_tui_event(tui_event, &handle, &io) } + DriverEvent::Usage(snapshot) => { + update_usage(&handle, &io, snapshot); + None + } }; if let Some(event) = fsm_event { @@ -929,6 +939,26 @@ fn execute_effect(effect: &Effect, ctx: DriverContext) { // Persistence // ============================================================================ +/// Apply a fresh usage snapshot: sync it to the view and write it to the +/// local cache so the next TUI open can render it immediately. +fn update_usage(handle: &Handle, io: &IoContext, snapshot: crate::usage::UsageSnapshot) { + handle.update({ + let snapshot = snapshot.clone(); + move |vs| vs.usage = Some(snapshot) + }); + + match serde_json::to_string(&snapshot) { + Ok(json) => { + let key = crate::usage::cache_key(&io.app_ctx.token); + let rt = tokio::runtime::Handle::current(); + if let Err(e) = rt.block_on(io.session_mgr.set_cached_usage(&key, &json)) { + tracing::warn!("Failed to persist usage cache: {e}"); + } + } + Err(e) => tracing::warn!("Failed to serialize usage snapshot: {e}"), + } +} + fn persist(fsm: &AgentFsm, io: &mut IoContext) { let start = std::time::Instant::now(); let rt = tokio::runtime::Handle::current(); @@ -1063,7 +1093,15 @@ async fn run_stream_bridge( }, Ok(StreamFrame::Control(control)) => match control { StreamControl::StatusChanged(status) => Some(Event::StreamStatusChanged(status)), - StreamControl::Done { session_id } => Some(Event::StreamDone { session_id }), + StreamControl::Done { + session_id, + credits, + } => { + if let Some(snapshot) = credits { + let _ = tx.send(DriverEvent::Usage(snapshot)); + } + Some(Event::StreamDone { session_id }) + } StreamControl::Error(msg) => Some(Event::StreamError(msg)), }, Ok(StreamFrame::SessionIdentity(session_id)) => { diff --git a/crates/atuin-ai/src/lib.rs b/crates/atuin-ai/src/lib.rs index 7f95dccee2e..6e3e19123b8 100644 --- a/crates/atuin-ai/src/lib.rs +++ b/crates/atuin-ai/src/lib.rs @@ -18,4 +18,5 @@ pub(crate) mod store; pub(crate) mod stream; pub(crate) mod tools; pub(crate) mod tui; +pub(crate) mod usage; pub(crate) mod user_context; diff --git a/crates/atuin-ai/src/session.rs b/crates/atuin-ai/src/session.rs index 848330fc044..591d1ba4c9a 100644 --- a/crates/atuin-ai/src/session.rs +++ b/crates/atuin-ai/src/session.rs @@ -54,6 +54,11 @@ pub(crate) trait SessionService: Send + Sync { async fn get_metadata(&self, session_id: &str, key: &str) -> Result>; async fn set_metadata(&self, session_id: &str, key: &str, value: &str) -> Result<()>; + + /// Read the cached usage snapshot (JSON, written-at unix timestamp) for + /// a user key. Not session-scoped: usage is per hub account. + async fn get_cached_usage(&self, user_key: &str) -> Result>; + async fn set_cached_usage(&self, user_key: &str, snapshot_json: &str) -> Result<()>; } // --------------------------------------------------------------------------- @@ -139,6 +144,14 @@ impl SessionService for LocalSessionService { async fn set_metadata(&self, session_id: &str, key: &str, value: &str) -> Result<()> { self.store.set_metadata(session_id, key, value).await } + + async fn get_cached_usage(&self, user_key: &str) -> Result> { + self.store.get_usage(user_key).await + } + + async fn set_cached_usage(&self, user_key: &str, snapshot_json: &str) -> Result<()> { + self.store.set_usage(user_key, snapshot_json).await + } } // --------------------------------------------------------------------------- @@ -337,6 +350,12 @@ impl SessionManager { .set_metadata(&self.session_id, key, value) .await } + + /// Write the usage cache for a user key. Not tied to the current + /// session, so no session row is created. + pub async fn set_cached_usage(&self, user_key: &str, snapshot_json: &str) -> Result<()> { + self.service.set_cached_usage(user_key, snapshot_json).await + } } #[cfg(test)] diff --git a/crates/atuin-ai/src/store.rs b/crates/atuin-ai/src/store.rs index 20b9e881492..06dbc778d95 100644 --- a/crates/atuin-ai/src/store.rs +++ b/crates/atuin-ai/src/store.rs @@ -331,6 +331,35 @@ impl AiSessionStore { .await?; Ok(()) } + + // ── Usage cache (server credit totals, one row per user key) ── + + /// Read the cached usage snapshot for a user key. Returns the snapshot + /// JSON and the unix timestamp it was written at. + pub async fn get_usage(&self, user_key: &str) -> Result> { + let row: Option<(String, i64)> = + sqlx::query_as("SELECT snapshot, updated_at FROM usage WHERE user_key = ?1") + .bind(user_key) + .fetch_optional(&self.pool) + .await?; + Ok(row) + } + + /// Write the usage snapshot for a user key (upsert). + pub async fn set_usage(&self, user_key: &str, snapshot_json: &str) -> Result<()> { + let now = OffsetDateTime::now_utc().unix_timestamp(); + sqlx::query( + "INSERT INTO usage (user_key, snapshot, updated_at) + VALUES (?1, ?2, ?3) + ON CONFLICT (user_key) DO UPDATE SET snapshot = ?2, updated_at = ?3", + ) + .bind(user_key) + .bind(snapshot_json) + .bind(now) + .execute(&self.pool) + .await?; + Ok(()) + } } #[cfg(test)] @@ -524,6 +553,26 @@ mod tests { assert_eq!(before.updated_at, after.updated_at); } + #[tokio::test] + async fn test_usage_cache_roundtrip() { + let store = new_test_store().await; + + assert!(store.get_usage("key-a").await.unwrap().is_none()); + + store.set_usage("key-a", r#"{"used":1}"#).await.unwrap(); + let (json, updated_at) = store.get_usage("key-a").await.unwrap().unwrap(); + assert_eq!(json, r#"{"used":1}"#); + assert!(updated_at > 0); + + // Upsert replaces the snapshot for the same key + store.set_usage("key-a", r#"{"used":2}"#).await.unwrap(); + let (json, _) = store.get_usage("key-a").await.unwrap().unwrap(); + assert_eq!(json, r#"{"used":2}"#); + + // Other keys are independent + assert!(store.get_usage("key-b").await.unwrap().is_none()); + } + #[tokio::test] async fn test_events_ordered_chronologically() { let store = new_test_store().await; diff --git a/crates/atuin-ai/src/stream.rs b/crates/atuin-ai/src/stream.rs index 467695f3024..62a37b83001 100644 --- a/crates/atuin-ai/src/stream.rs +++ b/crates/atuin-ai/src/stream.rs @@ -21,7 +21,11 @@ pub(crate) static APP_USER_AGENT: &str = concat!("atuin/", env!("CARGO_PKG_VERSI /// Frames that alter the stream lifecycle — terminal or state-changing. #[derive(Debug, Clone)] pub(crate) enum StreamControl { - Done { session_id: String }, + Done { + session_id: String, + /// Period credit totals from the server, when it sends them. + credits: Option, + }, Error(String), StatusChanged(String), } @@ -257,9 +261,12 @@ pub(crate) fn create_chat_stream( .and_then(|v| v.as_str()) .unwrap_or("") .to_string(); - yield Ok(StreamFrame::Control(StreamControl::Done { session_id })); + let credits = json.get("credits") + .cloned() + .and_then(|v| serde_json::from_value(v).ok()); + yield Ok(StreamFrame::Control(StreamControl::Done { session_id, credits })); } else { - yield Ok(StreamFrame::Control(StreamControl::Done { session_id: String::new() })); + yield Ok(StreamFrame::Control(StreamControl::Done { session_id: String::new(), credits: None })); } break; } diff --git a/crates/atuin-ai/src/usage.rs b/crates/atuin-ai/src/usage.rs new file mode 100644 index 00000000000..1fe21d3eedb --- /dev/null +++ b/crates/atuin-ai/src/usage.rs @@ -0,0 +1,115 @@ +//! Server-side credit usage: fetching and the done-event snapshot type. +//! +//! The hub reports the user's period credit totals two ways: a `credits` +//! object on the chat `done` event, and `GET /api/cli/usage` for reading it +//! outside a chat. Both share the same shape, deserialized here as +//! [`UsageSnapshot`]. Snapshots are cached in ai.db (see `store`) so the TUI +//! can render usage immediately on open, then refreshed in the background. + +use std::time::Duration; + +use eyre::{Context, Result}; +use reqwest::header::USER_AGENT; +use serde::{Deserialize, Serialize}; + +/// Cached usage older than this triggers a background refresh on TUI open. +pub(crate) const REFRESH_AFTER: Duration = Duration::from_secs(60); + +/// Used/limit pair in credits (billable tokens × model multiplier). +/// Limits use the server's sentinels: -1 unlimited, 0 disabled. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct UsageBucket { + pub used: i64, + pub limit: i64, +} + +/// The user's credit totals against their limits for the current period. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct UsageSnapshot { + /// e.g. "calendar_monthly" + pub period: String, + /// RFC 3339 timestamp of the next period reset. + pub resets_at: String, + pub requests: UsageBucket, + pub input: UsageBucket, + pub output: UsageBucket, +} + +/// Key for the local usage cache. The client never learns its hub user id, +/// so rows are keyed by a hash of the auth token: a different login (or a +/// rotated token) simply misses the cache and refetches. +pub(crate) fn cache_key(token: &str) -> String { + format!("{:016x}", xxhash_rust::xxh3::xxh3_64(token.as_bytes())) +} + +/// Fetch current usage from the hub. Mirrors the `credits` object on the +/// chat `done` event, for refreshing without starting a chat. +pub(crate) async fn fetch_usage(endpoint: &str, token: &str) -> Result { + atuin_common::tls::ensure_crypto_provider(); + let url = crate::stream::hub_url(endpoint, "/api/cli/usage")?; + + let response = reqwest::Client::new() + .get(url) + .header(USER_AGENT, crate::stream::APP_USER_AGENT) + .bearer_auth(token) + .timeout(Duration::from_secs(10)) + .send() + .await + .context("failed to fetch usage")?; + + let status = response.status(); + if !status.is_success() { + eyre::bail!("usage request failed ({status})"); + } + + response + .json::() + .await + .context("failed to parse usage response") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn deserializes_server_payload() { + // Shape documented in the hub's CliUsageController / credits_payload. + let json = r#"{ + "period": "calendar_monthly", + "resets_at": "2026-08-01T00:00:00Z", + "requests": {"used": 3, "limit": -1}, + "input": {"used": 12345, "limit": 5000000}, + "output": {"used": 678, "limit": 1000000} + }"#; + + let snapshot: UsageSnapshot = serde_json::from_str(json).unwrap(); + assert_eq!(snapshot.period, "calendar_monthly"); + assert_eq!(snapshot.requests.limit, -1); + assert_eq!(snapshot.input.used, 12345); + assert_eq!(snapshot.output.limit, 1_000_000); + } + + #[test] + fn snapshot_roundtrips_through_json() { + let snapshot = UsageSnapshot { + period: "calendar_monthly".into(), + resets_at: "2026-08-01T00:00:00Z".into(), + requests: UsageBucket { used: 1, limit: 10 }, + input: UsageBucket { used: 2, limit: 20 }, + output: UsageBucket { used: 3, limit: 0 }, + }; + + let json = serde_json::to_string(&snapshot).unwrap(); + assert_eq!( + serde_json::from_str::(&json).unwrap(), + snapshot + ); + } + + #[test] + fn cache_key_distinguishes_tokens() { + assert_ne!(cache_key("token-a"), cache_key("token-b")); + assert_eq!(cache_key("token-a"), cache_key("token-a")); + } +} From 70feef4aa36f334033c6352fde1f5e0640c72b74 Mon Sep 17 00:00:00 2001 From: Michelle Tilley Date: Wed, 8 Jul 2026 11:45:20 -0700 Subject: [PATCH 2/6] Show current model and usage in status bar --- crates/atuin-ai/src/tui/view/mod.rs | 67 +++++++++++++++++++++++++ crates/atuin-ai/src/usage.rs | 78 +++++++++++++++++++++++++++++ 2 files changed, 145 insertions(+) diff --git a/crates/atuin-ai/src/tui/view/mod.rs b/crates/atuin-ai/src/tui/view/mod.rs index 14d178e6f6e..86bb0f9cd65 100644 --- a/crates/atuin-ai/src/tui/view/mod.rs +++ b/crates/atuin-ai/src/tui/view/mod.rs @@ -179,11 +179,78 @@ fn input_view(state: &ViewState) -> Elements { }) }) + + #(status_bar_view(state)) } }) } } +/// Usage below this percentage isn't worth a status-bar warning. +const USAGE_BAR_THRESHOLD_PCT: f64 = 50.0; + +/// Width of the usage bar in cells. +const USAGE_BAR_WIDTH: usize = 5; + +/// One-line status bar under the input box: current model on the left; +/// on the right, once usage crosses the threshold, a small bar chart with +/// the percentage and time until the period resets. +fn status_bar_view(state: &ViewState) -> Elements { + let model_label = format!(" Model: {}", state.model.as_deref().unwrap_or("default")); + + let usage = state.usage.as_ref().and_then(|snapshot| { + let pct = snapshot.as_percentage()?; + if pct < USAGE_BAR_THRESHOLD_PCT { + return None; + } + Some((pct, snapshot.resets_in())) + }); + + element! { + HStack(key: "status-bar") { + View(width: WidthConstraint::Fill) { + Text { + Span(text: model_label, style: Style::default().fg(Color::DarkGray)) + } + } + #(if let Some((pct, resets_in)) = usage { + #({ + let filled = ((pct / 100.0).clamp(0.0, 1.0) * USAGE_BAR_WIDTH as f64).round() as usize; + let bar_filled = "█".repeat(filled); + let bar_empty = "░".repeat(USAGE_BAR_WIDTH - filled); + let pct_text = format!(" {}%", pct.round() as i64); + let resets_text = resets_in + .map(|d| format!(" · resets in {} ", crate::usage::format_reset_delta(d))) + .unwrap_or_default(); + + let bar_color = if pct >= 90.0 { + Color::Red + } else if pct >= 70.0 { + Color::Yellow + } else { + Color::Green + }; + + let width = (USAGE_BAR_WIDTH + + pct_text.chars().count() + + resets_text.chars().count()) as u16; + + element! { + View(width: WidthConstraint::Fixed(width)) { + Text { + Span(text: bar_filled, style: Style::default().fg(bar_color)) + Span(text: bar_empty, style: Style::default().fg(Color::DarkGray)) + Span(text: pct_text, style: Style::default().fg(Color::Gray)) + Span(text: resets_text, style: Style::default().fg(Color::DarkGray)) + } + } + } + }) + }) + } + } +} + /// Render the /model picker: one row per model, the in-use model marked. /// `current` is the session's explicit selection; when unset, the server /// default is what's actually in use, so mark that row instead. diff --git a/crates/atuin-ai/src/usage.rs b/crates/atuin-ai/src/usage.rs index 1fe21d3eedb..750b82ec41e 100644 --- a/crates/atuin-ai/src/usage.rs +++ b/crates/atuin-ai/src/usage.rs @@ -35,6 +35,50 @@ pub(crate) struct UsageSnapshot { pub output: UsageBucket, } +impl UsageSnapshot { + pub(crate) fn resets_in(&self) -> Option { + let reset_time = chrono::DateTime::parse_from_rfc3339(&self.resets_at).ok()?; + let now = chrono::Utc::now().fixed_offset(); + let duration = reset_time - now; + duration.to_std().ok() + } + + pub(crate) fn as_percentage(&self) -> Option { + let input_percentage = if self.input.limit > 0 { + Some(self.input.used as f64 / self.input.limit as f64 * 100.0) + } else { + None + }; + + let output_percentage = if self.output.limit > 0 { + Some(self.output.used as f64 / self.output.limit as f64 * 100.0) + } else { + None + }; + + match (input_percentage, output_percentage) { + (Some(input), Some(output)) if input > output => Some(input), + (Some(_), Some(output)) => Some(output), + (Some(input), None) => Some(input), + (None, Some(output)) => Some(output), + (None, None) => None, + } + } +} + +/// Format a reset delta as its largest sensible unit: "4d", "23h", or "56m". +/// Sub-minute deltas render as "1m" — "0m" would read as already reset. +pub(crate) fn format_reset_delta(delta: Duration) -> String { + let minutes = delta.as_secs() / 60; + if minutes >= 24 * 60 { + format!("{}d", minutes / (24 * 60)) + } else if minutes >= 60 { + format!("{}h", minutes / 60) + } else { + format!("{}m", minutes.max(1)) + } +} + /// Key for the local usage cache. The client never learns its hub user id, /// so rows are keyed by a hash of the auth token: a different login (or a /// rotated token) simply misses the cache and refetches. @@ -107,6 +151,40 @@ mod tests { ); } + #[test] + fn as_percentage_averages_limited_buckets() { + let mut snapshot = UsageSnapshot { + period: "calendar_monthly".into(), + resets_at: "2026-08-01T00:00:00Z".into(), + requests: UsageBucket { used: 3, limit: -1 }, + input: UsageBucket { + used: 50, + limit: 100, + }, + output: UsageBucket { + used: 90, + limit: 100, + }, + }; + assert_eq!(snapshot.as_percentage(), Some(70.0)); + + // Unlimited/disabled buckets drop out of the average + snapshot.output.limit = -1; + assert_eq!(snapshot.as_percentage(), Some(50.0)); + + snapshot.input.limit = 0; + assert_eq!(snapshot.as_percentage(), None); + } + + #[test] + fn formats_reset_deltas() { + let mins = |m: u64| Duration::from_secs(m * 60); + assert_eq!(format_reset_delta(mins(4 * 24 * 60 + 300)), "4d"); + assert_eq!(format_reset_delta(mins(23 * 60 + 59)), "23h"); + assert_eq!(format_reset_delta(mins(56)), "56m"); + assert_eq!(format_reset_delta(Duration::from_secs(30)), "1m"); + } + #[test] fn cache_key_distinguishes_tokens() { assert_ne!(cache_key("token-a"), cache_key("token-b")); From 9fb5057a3c5e8e44986c1d200e152b983bb8afa4 Mon Sep 17 00:00:00 2001 From: Michelle Tilley Date: Wed, 8 Jul 2026 11:54:28 -0700 Subject: [PATCH 3/6] Fix as_percentage test --- crates/atuin-ai/src/usage.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/atuin-ai/src/usage.rs b/crates/atuin-ai/src/usage.rs index 750b82ec41e..cf37fa4e57a 100644 --- a/crates/atuin-ai/src/usage.rs +++ b/crates/atuin-ai/src/usage.rs @@ -152,7 +152,7 @@ mod tests { } #[test] - fn as_percentage_averages_limited_buckets() { + fn as_percentage_uses_higher_limited_bucket() { let mut snapshot = UsageSnapshot { period: "calendar_monthly".into(), resets_at: "2026-08-01T00:00:00Z".into(), @@ -166,7 +166,7 @@ mod tests { limit: 100, }, }; - assert_eq!(snapshot.as_percentage(), Some(70.0)); + assert_eq!(snapshot.as_percentage(), Some(90.0)); // Unlimited/disabled buckets drop out of the average snapshot.output.limit = -1; From ef16e0966c012df89dd5afa916fc0ff576038db5 Mon Sep 17 00:00:00 2001 From: Michelle Tilley Date: Wed, 8 Jul 2026 13:56:58 -0700 Subject: [PATCH 4/6] Return CachedUsageSnapshot from get_cached_usage --- crates/atuin-ai/src/commands/inline.rs | 9 ++++---- crates/atuin-ai/src/driver.rs | 13 ++++-------- crates/atuin-ai/src/session.rs | 29 +++++++++++++++++++------- 3 files changed, 29 insertions(+), 22 deletions(-) diff --git a/crates/atuin-ai/src/commands/inline.rs b/crates/atuin-ai/src/commands/inline.rs index 17093218095..7068fbb3398 100644 --- a/crates/atuin-ai/src/commands/inline.rs +++ b/crates/atuin-ai/src/commands/inline.rs @@ -167,11 +167,10 @@ async fn run_inline_tui( // once the event channel exists) replaces it unless it's fresh. let usage_key = crate::usage::cache_key(&ctx.token); let (cached_usage, usage_is_fresh) = match service.get_cached_usage(&usage_key).await { - Ok(Some((json, updated_at))) => { - let snapshot = serde_json::from_str::(&json).ok(); - let age = time::OffsetDateTime::now_utc().unix_timestamp() - updated_at; - let fresh = snapshot.is_some() && age < crate::usage::REFRESH_AFTER.as_secs() as i64; - (snapshot, fresh) + Ok(Some(cached_snapshot)) => { + let age = time::OffsetDateTime::now_utc().unix_timestamp() - cached_snapshot.written_at; + let fresh = age < crate::usage::REFRESH_AFTER.as_secs() as i64; + (Some(cached_snapshot.snapshot), fresh) } Ok(None) => (None, false), Err(e) => { diff --git a/crates/atuin-ai/src/driver.rs b/crates/atuin-ai/src/driver.rs index f3a7cfbf6ab..c3b84e0cff7 100644 --- a/crates/atuin-ai/src/driver.rs +++ b/crates/atuin-ai/src/driver.rs @@ -947,15 +947,10 @@ fn update_usage(handle: &Handle, io: &IoContext, snapshot: crate::usa move |vs| vs.usage = Some(snapshot) }); - match serde_json::to_string(&snapshot) { - Ok(json) => { - let key = crate::usage::cache_key(&io.app_ctx.token); - let rt = tokio::runtime::Handle::current(); - if let Err(e) = rt.block_on(io.session_mgr.set_cached_usage(&key, &json)) { - tracing::warn!("Failed to persist usage cache: {e}"); - } - } - Err(e) => tracing::warn!("Failed to serialize usage snapshot: {e}"), + let key = crate::usage::cache_key(&io.app_ctx.token); + let rt = tokio::runtime::Handle::current(); + if let Err(e) = rt.block_on(io.session_mgr.set_cached_usage(&key, &snapshot)) { + tracing::warn!("Failed to persist usage cache: {e}"); } } diff --git a/crates/atuin-ai/src/session.rs b/crates/atuin-ai/src/session.rs index 591d1ba4c9a..d278d5a4968 100644 --- a/crates/atuin-ai/src/session.rs +++ b/crates/atuin-ai/src/session.rs @@ -11,11 +11,17 @@ use eyre::Result; use crate::event_serde; use crate::store::{AiSessionStore, StoredEvent, StoredSession}; use crate::tui::ConversationEvent; +use crate::usage::UsageSnapshot; // --------------------------------------------------------------------------- // Trait // --------------------------------------------------------------------------- +pub(crate) struct CachedUsageSnapshot { + pub snapshot: UsageSnapshot, + pub written_at: i64, +} + #[async_trait] pub(crate) trait SessionService: Send + Sync { async fn create_session( @@ -57,8 +63,8 @@ pub(crate) trait SessionService: Send + Sync { /// Read the cached usage snapshot (JSON, written-at unix timestamp) for /// a user key. Not session-scoped: usage is per hub account. - async fn get_cached_usage(&self, user_key: &str) -> Result>; - async fn set_cached_usage(&self, user_key: &str, snapshot_json: &str) -> Result<()>; + async fn get_cached_usage(&self, user_key: &str) -> Result>; + async fn set_cached_usage(&self, user_key: &str, snapshot: &UsageSnapshot) -> Result<()>; } // --------------------------------------------------------------------------- @@ -145,12 +151,19 @@ impl SessionService for LocalSessionService { self.store.set_metadata(session_id, key, value).await } - async fn get_cached_usage(&self, user_key: &str) -> Result> { - self.store.get_usage(user_key).await + async fn get_cached_usage(&self, user_key: &str) -> Result> { + match self.store.get_usage(user_key).await? { + Some((json, written_at)) => Ok(Some(CachedUsageSnapshot { + snapshot: serde_json::from_str(&json)?, + written_at, + })), + None => return Ok(None), + } } - async fn set_cached_usage(&self, user_key: &str, snapshot_json: &str) -> Result<()> { - self.store.set_usage(user_key, snapshot_json).await + async fn set_cached_usage(&self, user_key: &str, snapshot: &UsageSnapshot) -> Result<()> { + let snapshot_json = serde_json::to_string(snapshot)?; + self.store.set_usage(user_key, &snapshot_json).await } } @@ -353,8 +366,8 @@ impl SessionManager { /// Write the usage cache for a user key. Not tied to the current /// session, so no session row is created. - pub async fn set_cached_usage(&self, user_key: &str, snapshot_json: &str) -> Result<()> { - self.service.set_cached_usage(user_key, snapshot_json).await + pub async fn set_cached_usage(&self, user_key: &str, snapshot: &UsageSnapshot) -> Result<()> { + self.service.set_cached_usage(user_key, snapshot).await } } From 1ef21236bb408c9bee22305a45ae2b5039d34c4a Mon Sep 17 00:00:00 2001 From: Michelle Tilley Date: Wed, 8 Jul 2026 14:25:56 -0700 Subject: [PATCH 5/6] Return CachedUsageSnapshot from atuin_ai::store::AiSessionStore --- crates/atuin-ai/src/session.rs | 8 +----- crates/atuin-ai/src/store.rs | 47 +++++++++++++++++++++++++++------- 2 files changed, 39 insertions(+), 16 deletions(-) diff --git a/crates/atuin-ai/src/session.rs b/crates/atuin-ai/src/session.rs index d278d5a4968..d8abef69810 100644 --- a/crates/atuin-ai/src/session.rs +++ b/crates/atuin-ai/src/session.rs @@ -152,13 +152,7 @@ impl SessionService for LocalSessionService { } async fn get_cached_usage(&self, user_key: &str) -> Result> { - match self.store.get_usage(user_key).await? { - Some((json, written_at)) => Ok(Some(CachedUsageSnapshot { - snapshot: serde_json::from_str(&json)?, - written_at, - })), - None => return Ok(None), - } + self.store.get_usage(user_key).await } async fn set_cached_usage(&self, user_key: &str, snapshot: &UsageSnapshot) -> Result<()> { diff --git a/crates/atuin-ai/src/store.rs b/crates/atuin-ai/src/store.rs index 06dbc778d95..edd908cfd3a 100644 --- a/crates/atuin-ai/src/store.rs +++ b/crates/atuin-ai/src/store.rs @@ -6,6 +6,8 @@ use eyre::{Result, eyre}; use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePool, SqlitePoolOptions}; use time::OffsetDateTime; +use crate::session::CachedUsageSnapshot; + // Database row mappings — all columns are kept even if not yet read in // non-test code, since they're part of the schema and used in tests. #[derive(Debug)] @@ -336,13 +338,22 @@ impl AiSessionStore { /// Read the cached usage snapshot for a user key. Returns the snapshot /// JSON and the unix timestamp it was written at. - pub async fn get_usage(&self, user_key: &str) -> Result> { + pub async fn get_usage(&self, user_key: &str) -> Result> { let row: Option<(String, i64)> = sqlx::query_as("SELECT snapshot, updated_at FROM usage WHERE user_key = ?1") .bind(user_key) .fetch_optional(&self.pool) .await?; - Ok(row) + + let Some((snapshot, written_at)) = row else { + return Ok(None); + }; + + let snapshot = serde_json::from_str(&snapshot)?; + Ok(Some(CachedUsageSnapshot { + snapshot, + written_at, + })) } /// Write the usage snapshot for a user key (upsert). @@ -365,6 +376,7 @@ impl AiSessionStore { #[cfg(test)] mod tests { use super::*; + use crate::usage::UsageSnapshot; async fn new_test_store() -> AiSessionStore { AiSessionStore::new("sqlite::memory:", 2.0).await.unwrap() @@ -555,19 +567,36 @@ mod tests { #[tokio::test] async fn test_usage_cache_roundtrip() { + use crate::usage::UsageBucket; + let store = new_test_store().await; assert!(store.get_usage("key-a").await.unwrap().is_none()); - store.set_usage("key-a", r#"{"used":1}"#).await.unwrap(); - let (json, updated_at) = store.get_usage("key-a").await.unwrap().unwrap(); - assert_eq!(json, r#"{"used":1}"#); - assert!(updated_at > 0); + let snapshot = UsageSnapshot { + period: "calendar_monthly".into(), + resets_at: "2026-08-01T00:00:00Z".into(), + requests: UsageBucket { used: 1, limit: 10 }, + input: UsageBucket { used: 2, limit: 20 }, + output: UsageBucket { used: 3, limit: 0 }, + }; + let json = serde_json::to_string(&snapshot).unwrap(); + store.set_usage("key-a", &json).await.unwrap(); + + let cached = store.get_usage("key-a").await.unwrap().unwrap(); + assert!(cached.written_at > 0); + assert_eq!(cached.snapshot, snapshot); // Upsert replaces the snapshot for the same key - store.set_usage("key-a", r#"{"used":2}"#).await.unwrap(); - let (json, _) = store.get_usage("key-a").await.unwrap().unwrap(); - assert_eq!(json, r#"{"used":2}"#); + let updated = UsageSnapshot { + requests: UsageBucket { used: 4, limit: 10 }, + ..snapshot + }; + let json = serde_json::to_string(&updated).unwrap(); + store.set_usage("key-a", &json).await.unwrap(); + + let cached = store.get_usage("key-a").await.unwrap().unwrap(); + assert_eq!(cached.snapshot, updated); // Other keys are independent assert!(store.get_usage("key-b").await.unwrap().is_none()); From ee4ec82914e3c691bedf4dc2007412c1bd2dec26 Mon Sep 17 00:00:00 2001 From: Michelle Tilley Date: Wed, 8 Jul 2026 14:32:04 -0700 Subject: [PATCH 6/6] Take &UsageSnapshot in set_usage --- crates/atuin-ai/src/session.rs | 3 +-- crates/atuin-ai/src/store.rs | 11 +++++------ 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/crates/atuin-ai/src/session.rs b/crates/atuin-ai/src/session.rs index d8abef69810..216560bac68 100644 --- a/crates/atuin-ai/src/session.rs +++ b/crates/atuin-ai/src/session.rs @@ -156,8 +156,7 @@ impl SessionService for LocalSessionService { } async fn set_cached_usage(&self, user_key: &str, snapshot: &UsageSnapshot) -> Result<()> { - let snapshot_json = serde_json::to_string(snapshot)?; - self.store.set_usage(user_key, &snapshot_json).await + self.store.set_usage(user_key, snapshot).await } } diff --git a/crates/atuin-ai/src/store.rs b/crates/atuin-ai/src/store.rs index edd908cfd3a..e7fe23939d8 100644 --- a/crates/atuin-ai/src/store.rs +++ b/crates/atuin-ai/src/store.rs @@ -6,7 +6,7 @@ use eyre::{Result, eyre}; use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePool, SqlitePoolOptions}; use time::OffsetDateTime; -use crate::session::CachedUsageSnapshot; +use crate::{session::CachedUsageSnapshot, usage::UsageSnapshot}; // Database row mappings — all columns are kept even if not yet read in // non-test code, since they're part of the schema and used in tests. @@ -357,7 +357,8 @@ impl AiSessionStore { } /// Write the usage snapshot for a user key (upsert). - pub async fn set_usage(&self, user_key: &str, snapshot_json: &str) -> Result<()> { + pub async fn set_usage(&self, user_key: &str, snapshot: &UsageSnapshot) -> Result<()> { + let snapshot_json = serde_json::to_string(snapshot)?; let now = OffsetDateTime::now_utc().unix_timestamp(); sqlx::query( "INSERT INTO usage (user_key, snapshot, updated_at) @@ -580,8 +581,7 @@ mod tests { input: UsageBucket { used: 2, limit: 20 }, output: UsageBucket { used: 3, limit: 0 }, }; - let json = serde_json::to_string(&snapshot).unwrap(); - store.set_usage("key-a", &json).await.unwrap(); + store.set_usage("key-a", &snapshot).await.unwrap(); let cached = store.get_usage("key-a").await.unwrap().unwrap(); assert!(cached.written_at > 0); @@ -592,8 +592,7 @@ mod tests { requests: UsageBucket { used: 4, limit: 10 }, ..snapshot }; - let json = serde_json::to_string(&updated).unwrap(); - store.set_usage("key-a", &json).await.unwrap(); + store.set_usage("key-a", &updated).await.unwrap(); let cached = store.get_usage("key-a").await.unwrap().unwrap(); assert_eq!(cached.snapshot, updated);