diff --git a/crates/waku-client/src/client.rs b/crates/waku-client/src/client.rs index 6da56159..5822fd11 100644 --- a/crates/waku-client/src/client.rs +++ b/crates/waku-client/src/client.rs @@ -237,6 +237,32 @@ impl DaemonClient { pub fn shutdown(&self) { let _ = self.inner.outgoing.send(Outgoing::Shutdown); } + + #[cfg(test)] + pub(crate) fn disconnected_for_test(last_sequences: Vec) -> Self { + let sequences: HashMap<(Uuid, Uuid), LastSequence> = last_sequences + .into_iter() + .map(|cursor| { + ( + (cursor.session_id, cursor.runtime_id), + LastSequence { + epoch: cursor.epoch, + sequence: cursor.sequence, + }, + ) + }) + .collect(); + let inner = Arc::new(ClientInner { + outgoing: unbounded().0, + pending: Mutex::new(HashMap::new()), + sessions: Mutex::new(HashMap::new()), + pending_events: Mutex::new(HashMap::new()), + task_state_subscribers: Mutex::new(Vec::new()), + last_sequences: Mutex::new(sequences), + disconnected: AtomicBool::new(true), + }); + Self { inner } + } } fn daemon_url(address: &str) -> anyhow::Result { diff --git a/crates/waku-client/src/process.rs b/crates/waku-client/src/process.rs index 60d7d534..80be302d 100644 --- a/crates/waku-client/src/process.rs +++ b/crates/waku-client/src/process.rs @@ -17,7 +17,7 @@ use uuid::Uuid; use crate::DaemonClient; use waku_protocol::{ APP_EXECUTABLE_ENV, Command, DAEMON_TOKEN_ENV, DaemonReady, DaemonSettings, PROTOCOL_VERSION, - ResponsePayload, + ReplayCursor, ResponsePayload, }; const START_TIMEOUT: Duration = Duration::from_secs(15); const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(1); @@ -129,6 +129,7 @@ pub fn parse_allowed_origins(text: &str) -> anyhow::Result> { pub struct DaemonProcess { client: DaemonClient, child: Child, + address: String, } impl DaemonProcess { @@ -232,13 +233,43 @@ impl DaemonProcess { return Err(error); } }; - Ok(Self { client, child }) + Ok(Self { + client, + child, + address: client_address, + }) } pub fn client(&self) -> DaemonClient { self.client.clone() } + pub fn address(&self) -> &str { + &self.address + } + + pub(crate) fn replace_client(&mut self, client: DaemonClient) { + self.client = client; + } + + #[cfg(test)] + fn for_test(client: DaemonClient, address: String) -> Self { + #[cfg(unix)] + let child = std::process::Command::new("true") + .spawn() + .expect("`true` is required for tests"); + #[cfg(windows)] + let child = std::process::Command::new("cmd") + .args(["/c", "exit", "0"]) + .spawn() + .expect("cmd is required for tests"); + Self { + client, + child, + address, + } + } + fn has_exited(&mut self) -> bool { !matches!(self.child.try_wait(), Ok(None)) } @@ -310,7 +341,11 @@ struct SupervisorInner { } enum DaemonTarget { - Local(DaemonProcess), + Local { + process: DaemonProcess, + address: String, + token: String, + }, Restarting(DaemonClient), Remote { client: DaemonClient, @@ -322,7 +357,7 @@ enum DaemonTarget { impl DaemonTarget { fn client(&self) -> DaemonClient { match self { - Self::Local(process) => process.client(), + Self::Local { process, .. } => process.client(), Self::Restarting(client) => client.clone(), Self::Remote { client, .. } => client.clone(), } @@ -354,8 +389,14 @@ impl DaemonSupervisor { let process = DaemonProcess::spawn_configured(executable, exposure.clone())?; let settings = read_settings(&process.client())?; let initial_stamp = ExecutableStamp::read(executable)?; + let address = process.address().to_owned(); + let token = exposure.token.clone(); let supervisor = Self::from_target( - DaemonTarget::Local(process), + DaemonTarget::Local { + process, + address, + token, + }, Some(executable.to_owned()), Some(exposure), settings, @@ -503,6 +544,40 @@ impl Drop for DaemonSupervisor { } } +fn local_reconnect_params(target: &DaemonTarget) -> Option<(String, String, Vec)> { + match target { + DaemonTarget::Local { + process, + address, + token, + } if process.client().is_disconnected() => Some(( + address.clone(), + token.clone(), + process.client().last_sequences(), + )), + _ => None, + } +} + +fn replace_local_client(target: &mut DaemonTarget, replacement: DaemonClient) -> bool { + match target { + DaemonTarget::Local { process, .. } => { + process.replace_client(replacement); + true + } + _ => false, + } +} + +fn try_local_reconnect( + target: &mut DaemonTarget, + connect: &mut impl FnMut(&str, String, Vec) -> anyhow::Result, +) -> Option { + let (address, token, resume_from) = local_reconnect_params(target)?; + let replacement = connect(&address, token, resume_from).ok()?; + replace_local_client(target, replacement.clone()).then_some(replacement) +} + fn monitor_daemon( weak_inner: std::sync::Weak, mut active_stamp: Option, @@ -558,8 +633,41 @@ fn monitor_daemon( .retain(|subscriber| subscriber.send(replacement.clone()).is_ok()); continue; } + let local_reconnect = { + let _restart = inner.restart.lock(); + let mut target = inner.target.lock(); + let disconnected = match &*target { + DaemonTarget::Local { process, .. } if process.client().is_disconnected() => { + Some(process.client().clone()) + } + _ => None, + }; + let Some(disconnected) = disconnected else { + continue; + }; + let still_current = matches!( + &*target, + DaemonTarget::Local { process, .. } + if process.client().same_connection(&disconnected) + && process.client().is_disconnected() + ); + if !still_current { + continue; + } + let mut connect = |address: &str, token: String, resume_from: Vec| { + DaemonClient::connect_with_resume(address, token, resume_from) + }; + try_local_reconnect(&mut *target, &mut connect) + }; + if let Some(replacement) = local_reconnect { + inner + .client_updates + .lock() + .retain(|subscriber| subscriber.send(replacement.clone()).is_ok()); + continue; + } let process_exited = match &mut *inner.target.lock() { - DaemonTarget::Local(process) => process.has_exited(), + DaemonTarget::Local { process, .. } => process.has_exited(), DaemonTarget::Restarting(_) => true, DaemonTarget::Remote { .. } => continue, }; @@ -604,12 +712,12 @@ fn replace_local_daemon( bail!("the connected daemon is managed outside Waku Desktop") } DaemonTarget::Restarting(_) => None, - DaemonTarget::Local(process) => { + DaemonTarget::Local { process, .. } => { let disconnected = process.client(); let previous = std::mem::replace(&mut *target, DaemonTarget::Restarting(disconnected)); match previous { - DaemonTarget::Local(process) => Some(process), + DaemonTarget::Local { process, .. } => Some(process), _ => unreachable!("local daemon target changed while locked"), } } @@ -620,7 +728,13 @@ fn replace_local_daemon( drop(previous); let replacement = DaemonProcess::spawn_configured(executable, exposure.clone())?; let client = replacement.client(); - *inner.target.lock() = DaemonTarget::Local(replacement); + let address = replacement.address().to_owned(); + let token = exposure.token.clone(); + *inner.target.lock() = DaemonTarget::Local { + process: replacement, + address, + token, + }; inner .client_updates .lock() @@ -711,4 +825,114 @@ mod tests { ); assert_eq!(desktop_client_address("[::]:34123").unwrap(), "[::1]:34123"); } + + #[test] + fn local_reconnect_params_requires_disconnected_client() { + let disconnected = DaemonClient::disconnected_for_test(vec![ReplayCursor { + session_id: Uuid::from_u128(1), + runtime_id: Uuid::from_u128(2), + epoch: Uuid::from_u128(3), + sequence: 5, + }]); + let local = DaemonTarget::Local { + process: DaemonProcess::for_test(disconnected.clone(), "127.0.0.1:1234".into()), + address: "127.0.0.1:1234".into(), + token: "token-a".into(), + }; + let (address, token, resume_from) = local_reconnect_params(&local).unwrap(); + assert_eq!(address, "127.0.0.1:1234"); + assert_eq!(token, "token-a"); + assert_eq!(resume_from.len(), 1); + assert_eq!(resume_from[0].session_id, Uuid::from_u128(1)); + assert_eq!(resume_from[0].runtime_id, Uuid::from_u128(2)); + assert_eq!(resume_from[0].epoch, Uuid::from_u128(3)); + assert_eq!(resume_from[0].sequence, 5); + } + + #[test] + fn local_reconnect_params_skips_connected_and_remote_targets() { + // Remote target is never considered for local reconnect. + let remote = DaemonTarget::Remote { + client: DaemonClient::disconnected_for_test(Vec::new()), + address: "127.0.0.1:1234".into(), + token: "token-a".into(), + }; + assert!(local_reconnect_params(&remote).is_none()); + } + + #[test] + fn replace_local_client_swaps_the_client() { + let original = DaemonClient::disconnected_for_test(Vec::new()); + let replacement = DaemonClient::disconnected_for_test(Vec::new()); + let mut local = DaemonTarget::Local { + process: DaemonProcess::for_test(original.clone(), "127.0.0.1:1234".into()), + address: "127.0.0.1:1234".into(), + token: "token-a".into(), + }; + assert!(replace_local_client(&mut local, replacement.clone())); + match local { + DaemonTarget::Local { process, .. } => { + assert!(process.client().same_connection(&replacement)); + assert!(!process.client().same_connection(&original)); + } + _ => panic!("expected local target"), + } + } + + #[test] + fn try_local_reconnect_uses_injected_connector() { + let session_id = Uuid::from_u128(1); + let runtime_id = Uuid::from_u128(2); + let epoch = Uuid::from_u128(3); + let original = DaemonClient::disconnected_for_test(vec![ReplayCursor { + session_id, + runtime_id, + epoch, + sequence: 5, + }]); + let replacement = DaemonClient::disconnected_for_test(Vec::new()); + let mut local = DaemonTarget::Local { + process: DaemonProcess::for_test(original.clone(), "127.0.0.1:1234".into()), + address: "127.0.0.1:1234".into(), + token: "token-a".into(), + }; + + let mut calls: Vec<(String, String, Vec)> = Vec::new(); + let replacement_to_return = replacement.clone(); + let mut connect = |address: &str, token: String, resume_from: Vec| { + calls.push((address.to_string(), token, resume_from)); + Ok(replacement_to_return.clone()) + }; + + let result = try_local_reconnect(&mut local, &mut connect).unwrap(); + assert!(result.same_connection(&replacement)); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].0, "127.0.0.1:1234"); + assert_eq!(calls[0].1, "token-a"); + assert_eq!(calls[0].2.len(), 1); + assert_eq!(calls[0].2[0].session_id, session_id); + assert_eq!(calls[0].2[0].sequence, 5); + + match local { + DaemonTarget::Local { process, .. } => { + assert!(process.client().same_connection(&replacement)); + } + _ => panic!("expected local target"), + } + } + + #[test] + fn try_local_reconnect_returns_none_when_already_connected() { + // We cannot create a genuinely connected client without a server, but we + // can simulate the "no reconnect needed" case by using a Remote target. + let mut remote = DaemonTarget::Remote { + client: DaemonClient::disconnected_for_test(Vec::new()), + address: "127.0.0.1:1234".into(), + token: "token-a".into(), + }; + let mut connect = |_address: &str, _token: String, _resume_from: Vec| { + panic!("connector should not be called") + }; + assert!(try_local_reconnect(&mut remote, &mut connect).is_none()); + } } diff --git a/crates/waku-core/src/usage_history.rs b/crates/waku-core/src/usage_history.rs index f6ddeae3..14a4c4d4 100644 --- a/crates/waku-core/src/usage_history.rs +++ b/crates/waku-core/src/usage_history.rs @@ -61,6 +61,10 @@ fn might_carry_usage(line: &str, provider: UsageProvider) -> bool { match provider { UsageProvider::Claude => line.contains("\"usage\""), UsageProvider::Codex => line.contains("\"token_count\""), + // OpenCode/OpenCode2 are read from their SQLite database, never from + // JSONL transcripts, so this JSONL pre-filter is never consulted for + // them; `false` keeps the signature honest. + UsageProvider::OpenCode | UsageProvider::OpenCode2 => false, } } @@ -558,7 +562,9 @@ pub struct FileCacheEntry { pub type ScanCache = HashMap; -/// The transcript root scanned for one provider. +/// The transcript root scanned for one provider. OpenCode/OpenCode2 keep their +/// sessions in a SQLite database (see [`provider_db_path`]), not JSONL, so they +/// have no transcript directory here. fn provider_root(provider: UsageProvider) -> Option { match provider { UsageProvider::Claude => match std::env::var_os("CLAUDE_CONFIG_DIR") { @@ -569,9 +575,34 @@ fn provider_root(provider: UsageProvider) -> Option { Some(dir) if !dir.is_empty() => Some(PathBuf::from(dir).join("sessions")), _ => dirs::home_dir().map(|home| home.join(".codex/sessions")), }, + UsageProvider::OpenCode | UsageProvider::OpenCode2 => None, } } +/// The OpenCode on-disk database for one provider, if it can be located. Both +/// OpenCode and its `opencode2` sibling store sessions under an `opencode` +/// (or `opencode2`) data directory; we probe the common locations and return +/// the first that exists. A missing database is not an error — it simply means +/// that provider has never been used on this machine, mirroring how Claude/Codex +/// are skipped when their transcript directory is absent. +fn provider_db_path(provider: UsageProvider) -> Option { + let name = match provider { + UsageProvider::OpenCode => "opencode", + UsageProvider::OpenCode2 => "opencode2", + UsageProvider::Claude | UsageProvider::Codex => return None, + }; + let home = dirs::home_dir()?; + let candidates = [ + // XDG data home (Linux; also where the desktop app lands on this box). + Some(home.join(".local/share").join(name).join("opencode.db")), + // Cross-platform data dir (Windows %LOCALAPPDATA%, macOS Application Support). + dirs::data_dir().map(|data| data.join(name).join("opencode.db")), + // Legacy config location. + Some(home.join(".config").join(name).join("opencode.db")), + ]; + candidates.into_iter().flatten().find(|path| path.is_file()) +} + /// Lists `.jsonl` transcripts under `root` modified at or after `since_ms`. /// Errors on individual entries are swallowed: session files rotate and get /// removed while the walk is in flight, and a partial listing beats failing @@ -667,11 +698,97 @@ fn read_transcript_records(path: &Path, provider: UsageProvider) -> Option return None, } } Some(records) } +/// Reads OpenCode (or `opencode2`) session totals straight from its SQLite +/// database. OpenCode stores per-session token and cost aggregates in the +/// `session`/`session_v2` tables (the newer `session_v2` is the live one; we +/// union both so pre-migration history is not lost), so unlike Claude/Codex we +/// never parse JSONL. Each row becomes one [`UsageRecord`] keyed by its session +/// id — the aggregator's global de-duplication drops the same session if it +/// appears in both tables, and a re-scan rebuilds the snapshot from scratch so +/// there is no cross-scan double counting. +fn read_opencode_db( + path: &Path, + provider: UsageProvider, + since_ms: i64, +) -> Result, String> { + use rusqlite::OpenFlags; + let conn = rusqlite::Connection::open_with_flags( + path, + OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX, + ) + .map_err(|cause| format!("open {}: {cause}", path.display()))?; + + // `model` is a JSON object like {"id":"glm-5.3-flash","providerID":"opencode-go"}; + // the pricing table keys on the bare model id, so we unwrap it. + let mut statement = conn + .prepare( + "SELECT id, directory, model, cost, tokens_input, tokens_output, \ + tokens_reasoning, tokens_cache_read, tokens_cache_write, time_created \ + FROM session_v2 WHERE time_created >= ?1 \ + UNION ALL \ + SELECT id, directory, model, cost, tokens_input, tokens_output, \ + tokens_reasoning, tokens_cache_read, tokens_cache_write, time_created \ + FROM session WHERE time_created >= ?1", + ) + .map_err(|cause| format!("query {}: {cause}", path.display()))?; + + let rows = statement + .query_map(rusqlite::params![since_ms], |row| { + Ok(( + row.get::<_, String>(0)?, // id + row.get::<_, String>(1)?, // directory + row.get::<_, String>(2)?, // model (JSON) + row.get::<_, f64>(3)?, // cost + row.get::<_, i64>(4)?, // tokens_input + row.get::<_, i64>(5)?, // tokens_output + row.get::<_, i64>(6)?, // tokens_reasoning + row.get::<_, i64>(7)?, // tokens_cache_read + row.get::<_, i64>(8)?, // tokens_cache_write + row.get::<_, i64>(9)?, // time_created (ms) + )) + }) + .map_err(|cause| format!("read {}: {cause}", path.display()))?; + + let mut records = Vec::new(); + for row in rows { + let (id, directory, model_json, cost, input, output, reasoning, cache_read, cache_write, created) = + row.map_err(|cause| format!("row {}: {cause}", path.display()))?; + let model = serde_json::from_str::(&model_json) + .ok() + .and_then(|value| value.get("id").and_then(Value::as_str).map(str::to_owned)) + .filter(|id| !id.is_empty()) + .unwrap_or_else(|| model_json.clone()); + if model.is_empty() { + continue; + } + records.push(UsageRecord { + provider, + timestamp_ms: created, + model, + session_id: id.clone(), + project: directory, + totals: TokenTotals { + uncached_input: input.max(0) as u64, + cached_input: cache_read.max(0) as u64, + cache_creation: cache_write.max(0) as u64, + output: output.max(0) as u64, + reasoning: reasoning.max(0) as u64, + }, + reported_cost_usd: Some(cost), + dedupe_key: Some(id), + }); + } + Ok(records) +} + /* ------------------------------------------------------------------------- */ /* Aggregation */ /* ------------------------------------------------------------------------- */ @@ -698,7 +815,7 @@ struct Bucket { struct ProjectAccumulator { cost_usd: f64, total_tokens: u64, - by_provider: [ProviderDay; 2], + by_provider: [ProviderDay; UsageProvider::ALL.len()], sessions: HashSet<(UsageProvider, String)>, /// Cost per model, for the row's "top models" caption. models: HashMap, @@ -855,6 +972,22 @@ pub fn scan( let mut errors = Vec::new(); for provider in UsageProvider::ALL { + // OpenCode/OpenCode2 persist sessions in a SQLite database rather than + // JSONL transcripts: read the database directly and fold its rows in. + if let Some(db) = provider_db_path(provider) { + if db.is_file() { + scanned_files += 1; + match read_opencode_db(&db, provider, mtime_cutoff_ms) { + Ok(records) => { + for record in &records { + aggregator.add(record, rates); + } + } + Err(message) => errors.push(message), + } + } + continue; + } let Some(root) = provider_root(provider) else { continue; }; @@ -968,7 +1101,7 @@ fn derive_history( day: *day, cost_usd: 0.0, total_tokens: 0, - by_provider: [ProviderDay::default(); 2], + by_provider: [ProviderDay::default(); UsageProvider::ALL.len()], }); day_entry.cost_usd += bucket.cost_usd; day_entry.total_tokens += tokens; @@ -1028,7 +1161,7 @@ fn derive_history( first_day: first_of_month(day.day), cost_usd: 0.0, total_tokens: 0, - by_provider: [ProviderDay::default(); 2], + by_provider: [ProviderDay::default(); UsageProvider::ALL.len()], sessions: 0, active_days: 0, top_models: Vec::new(), diff --git a/crates/waku-protocol/src/usage_history.rs b/crates/waku-protocol/src/usage_history.rs index 5964c0f5..2545b3e1 100644 --- a/crates/waku-protocol/src/usage_history.rs +++ b/crates/waku-protocol/src/usage_history.rs @@ -93,15 +93,27 @@ use chrono::Datelike as _; pub enum UsageProvider { Claude, Codex, + /// OpenCode (and its `opencode2` sibling) persist sessions in a SQLite + /// database rather than JSONL transcripts, so the usage scan reads their + /// `session`/`session_v2` tables directly instead of walking `.jsonl` files. + OpenCode, + OpenCode2, } impl UsageProvider { - pub const ALL: [UsageProvider; 2] = [UsageProvider::Claude, UsageProvider::Codex]; + pub const ALL: [UsageProvider; 4] = [ + UsageProvider::Claude, + UsageProvider::Codex, + UsageProvider::OpenCode, + UsageProvider::OpenCode2, + ]; pub fn label(self) -> &'static str { match self { UsageProvider::Claude => "Claude Code", UsageProvider::Codex => "Codex", + UsageProvider::OpenCode => "OpenCode", + UsageProvider::OpenCode2 => "OpenCode 2", } } @@ -109,6 +121,8 @@ impl UsageProvider { match self { UsageProvider::Claude => 0, UsageProvider::Codex => 1, + UsageProvider::OpenCode => 2, + UsageProvider::OpenCode2 => 3, } } } @@ -178,7 +192,7 @@ pub struct DaySlice { pub day: NaiveDate, pub cost_usd: f64, pub total_tokens: u64, - pub by_provider: [ProviderDay; 2], + pub by_provider: [ProviderDay; 4], } #[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, TS)] @@ -196,7 +210,7 @@ pub struct MonthSlice { pub first_day: NaiveDate, pub cost_usd: f64, pub total_tokens: u64, - pub by_provider: [ProviderDay; 2], + pub by_provider: [ProviderDay; 4], pub sessions: u64, pub active_days: u32, pub top_models: Vec<(String, f64)>, @@ -208,7 +222,7 @@ pub struct ProjectSlice { pub path: String, pub cost_usd: f64, pub total_tokens: u64, - pub by_provider: [ProviderDay; 2], + pub by_provider: [ProviderDay; 4], pub sessions: u64, pub cost_share: f64, pub last_day: Option, diff --git a/src/app/usage_page.rs b/src/app/usage_page.rs index c1d7b874..ebb8a942 100644 --- a/src/app/usage_page.rs +++ b/src/app/usage_page.rs @@ -32,6 +32,8 @@ fn provider_kind(provider: UsageProvider) -> ProviderKind { match provider { UsageProvider::Claude => ProviderKind::Claude, UsageProvider::Codex => ProviderKind::Codex, + UsageProvider::OpenCode => ProviderKind::OpenCode, + UsageProvider::OpenCode2 => ProviderKind::OpenCode2, } } @@ -698,7 +700,7 @@ impl Waku { // One column per day, per provider in ALL order. The chart paths and // the hover readout both consume this, so the number under the cursor // is by construction the number that was plotted. - let series: Vec<[f64; 2]> = days + let series: Vec<[f64; UsageProvider::ALL.len()]> = days .iter() .map(|day| { let slice = history.day(*day); @@ -713,7 +715,7 @@ impl Waku { }) .unwrap_or(0.0) }; - [value(UsageProvider::Claude), value(UsageProvider::Codex)] + UsageProvider::ALL.map(value) }) .collect(); // The scale tops out at the largest single provider-day, not the @@ -760,10 +762,7 @@ impl Waku { } let hover = self.usage_chart_hover.filter(|index| *index < day_count); - let colors = [ - provider_color(theme, ProviderKind::Claude), - provider_color(theme, ProviderKind::Codex), - ]; + let colors = UsageProvider::ALL.map(|provider| provider_color(theme, provider_kind(provider))); let bounds_cell = self.usage_chart_bounds.clone(); let paint_series = series.clone(); let paint_ticks = ticks.clone(); @@ -1343,7 +1342,7 @@ impl Waku { ) .child(div().mt(px(9.0)).child(usage_split_bar( &theme, - colors, + &colors, if peak <= 0.0 { 0.0 } else { @@ -2060,11 +2059,8 @@ fn rank_by_cost(history: &UsageHistory) -> bool { history.cost_usd > 0.0 } -fn usage_provider_colors(theme: &Theme) -> [Hsla; 2] { - [ - provider_color(theme, ProviderKind::Claude), - provider_color(theme, ProviderKind::Codex), - ] +fn usage_provider_colors(theme: &Theme) -> [Hsla; UsageProvider::ALL.len()] { + UsageProvider::ALL.map(|provider| provider_color(theme, provider_kind(provider))) } /// The period total in the ranking unit. @@ -2141,15 +2137,15 @@ fn usage_list_empty_row(theme: &Theme, message: String) -> Div { /// one glance carries both size and mix. fn usage_split_bar( theme: &Theme, - colors: [Hsla; 2], + colors: &[Hsla], length: f32, - by_provider: &[ProviderDay; 2], + by_provider: &[ProviderDay], by_cost: bool, ) -> Div { - let values = [ - usage_provider_value(&by_provider[0], by_cost), - usage_provider_value(&by_provider[1], by_cost), - ]; + let values: Vec = by_provider + .iter() + .map(|entry| usage_provider_value(entry, by_cost)) + .collect(); let sum = values[0] + values[1]; let length = if length > 0.0 { length.clamp(0.02, 1.0) @@ -2192,7 +2188,7 @@ fn usage_provider_value(entry: &ProviderDay, by_cost: bool) -> f64 { /// Per-provider amounts with their marks, skipping providers absent from /// the row. -fn usage_provider_values(theme: &Theme, by_provider: &[ProviderDay; 2], by_cost: bool) -> Div { +fn usage_provider_values(theme: &Theme, by_provider: &[ProviderDay], by_cost: bool) -> Div { let mut row = div().flex().items_center().gap(px(14.0)); for provider in UsageProvider::ALL { let entry = by_provider[provider.index()]; @@ -2318,21 +2314,21 @@ fn usage_month_strip( first_day: NaiveDate, peak: f64, by_cost: bool, - colors: [Hsla; 2], + colors: [Hsla; UsageProvider::ALL.len()], ) -> impl IntoElement { let day_count = usage_history::days_in_month(first_day); - let values: Vec<[f64; 2]> = (0..day_count) + let values: Vec> = (0..day_count) .map(|offset| { let day = first_day + chrono::Days::new(u64::from(offset)); history .day(day) .map(|slice| { - [ - usage_provider_value(&slice.by_provider[0], by_cost), - usage_provider_value(&slice.by_provider[1], by_cost), - ] + UsageProvider::ALL + .iter() + .map(|provider| usage_provider_value(&slice.by_provider[provider.index()], by_cost)) + .collect() }) - .unwrap_or([0.0, 0.0]) + .unwrap_or_else(Vec::new) }) .collect(); canvas( @@ -2496,7 +2492,7 @@ fn usage_month_row( history: &UsageHistory, month: &MonthSlice, theme: &Theme, - colors: [Hsla; 2], + colors: [Hsla; UsageProvider::ALL.len()], by_cost: bool, peak: f64, day_peak: f64, @@ -2580,7 +2576,7 @@ fn usage_month_row( ) .child(div().mt(px(9.0)).child(usage_split_bar( theme, - colors, + &colors, if peak <= 0.0 { 0.0 } else {