1use super::auth::OpencodeGoAuth;
2use super::types::{OpencodeGoUsageSnapshot, OpencodeGoUsageSource, OpencodeGoUsageWindow};
3use chrono::{DateTime, Duration, Utc};
4use rusqlite::Connection;
5use serde::Deserialize;
6use std::path::{Path, PathBuf};
7use tempfile::TempDir;
8use thiserror::Error;
9
10const FIVE_HOUR_LIMIT_USD: f64 = 12.0;
11const WEEKLY_LIMIT_USD: f64 = 30.0;
12const MONTHLY_LIMIT_USD: f64 = 60.0;
13
14const WINDOW_SPECS: [WindowSpec; 3] = [
15 WindowSpec {
16 entry_type: "five_hour_spend",
17 window_seconds: 5 * 60 * 60,
18 limit_usd: FIVE_HOUR_LIMIT_USD,
19 },
20 WindowSpec {
21 entry_type: "weekly_spend",
22 window_seconds: 7 * 24 * 60 * 60,
23 limit_usd: WEEKLY_LIMIT_USD,
24 },
25 WindowSpec {
26 entry_type: "monthly_spend",
27 window_seconds: 30 * 24 * 60 * 60,
28 limit_usd: MONTHLY_LIMIT_USD,
29 },
30];
31
32const LIMIT_EPSILON: f64 = 1e-9;
33
34#[derive(Debug, Error)]
35pub enum OpencodeGoUsageError {
36 #[error("could not determine home directory for opencode.db")]
37 HomeDirNotFound,
38
39 #[error("failed to read OpenCode usage database: {0}")]
40 Io(#[from] std::io::Error),
41
42 #[error("failed to query OpenCode usage database: {0}")]
43 Sql(#[from] rusqlite::Error),
44
45 #[error("failed to parse OpenCode message row: {0}")]
46 Parse(#[from] serde_json::Error),
47}
48
49#[derive(Debug, Clone, PartialEq)]
50struct UsageRecord {
51 completed_at: DateTime<Utc>,
52 cost_usd: f64,
53}
54
55#[derive(Debug, Clone, Copy)]
56struct WindowSpec {
57 entry_type: &'static str,
58 window_seconds: i64,
59 limit_usd: f64,
60}
61
62#[derive(Debug, Deserialize)]
63struct MessageRow {
64 role: String,
65 #[serde(rename = "providerID")]
66 provider_id: Option<String>,
67 cost: Option<f64>,
68 time: Option<MessageTime>,
69}
70
71#[derive(Debug, Deserialize)]
72struct MessageTime {
73 completed: Option<i64>,
74}
75
76pub struct OpencodeGoUsageStore;
77
78impl OpencodeGoUsageStore {
79 pub fn fetch_usage() -> Result<OpencodeGoUsageSnapshot, OpencodeGoUsageError> {
84 Self::fetch_usage_with_paths_at(None, None, Utc::now())
85 }
86
87 pub fn fetch_usage_from_path_at(
92 db_path: &Path,
93 now: DateTime<Utc>,
94 ) -> Result<OpencodeGoUsageSnapshot, OpencodeGoUsageError> {
95 Self::fetch_usage_with_paths_at(Some(db_path), None, now)
96 }
97
98 pub fn fetch_usage_with_paths_at(
103 db_path: Option<&Path>,
104 auth_path: Option<&Path>,
105 now: DateTime<Utc>,
106 ) -> Result<OpencodeGoUsageSnapshot, OpencodeGoUsageError> {
107 let credentials_available = auth_path
108 .map_or_else(
109 OpencodeGoAuth::read_api_key,
110 OpencodeGoAuth::read_api_key_from,
111 )
112 .is_ok();
113 let db_path = match db_path {
114 Some(path) => path.to_path_buf(),
115 None => Self::default_db_path()?,
116 };
117 let records = if db_path.exists() {
118 Self::load_records(&db_path)?
119 } else {
120 Vec::new()
121 };
122
123 Ok(Self::snapshot_from_records(
124 now,
125 &records,
126 credentials_available,
127 ))
128 }
129
130 fn default_db_path() -> Result<PathBuf, OpencodeGoUsageError> {
131 let home = dirs::home_dir().ok_or(OpencodeGoUsageError::HomeDirNotFound)?;
132 Ok(home.join(".local/share/opencode/opencode.db"))
133 }
134
135 fn load_records(db_path: &Path) -> Result<Vec<UsageRecord>, OpencodeGoUsageError> {
136 let (temp_dir, temp_db_path) = Self::copy_sqlite_database(db_path)?;
137 let conn = Connection::open(&temp_db_path)?;
138 let records = Self::query_records(&conn)?;
139 drop(conn);
140 drop(temp_dir);
141 Ok(records)
142 }
143
144 fn copy_sqlite_database(db_path: &Path) -> Result<(TempDir, PathBuf), OpencodeGoUsageError> {
145 let temp_dir = tempfile::tempdir()?;
146 let file_name = db_path
147 .file_name()
148 .ok_or_else(|| std::io::Error::other("opencode.db path has no file name"))?;
149 let temp_db_path = temp_dir.path().join(file_name);
150 std::fs::copy(db_path, &temp_db_path)?;
151
152 for suffix in ["-wal", "-shm"] {
153 let sidecar_name = format!("{}{}", file_name.to_string_lossy(), suffix);
154 let src = db_path.with_file_name(&sidecar_name);
155 if src.exists() {
156 let dst = temp_dir.path().join(sidecar_name);
157 std::fs::copy(src, dst)?;
158 }
159 }
160
161 Ok((temp_dir, temp_db_path))
162 }
163
164 fn query_records(conn: &Connection) -> Result<Vec<UsageRecord>, OpencodeGoUsageError> {
165 let mut stmt = match conn.prepare("SELECT data FROM message") {
166 Ok(stmt) => stmt,
167 Err(rusqlite::Error::SqliteFailure(_, Some(message)))
168 if message.contains("no such table: message") =>
169 {
170 return Ok(Vec::new());
171 }
172 Err(err) => return Err(err.into()),
173 };
174 let rows = stmt.query_map([], |row| row.get::<_, String>(0))?;
175 let mut records = Vec::new();
176
177 for row in rows {
178 let row = serde_json::from_str::<MessageRow>(&row?)?;
179 if row.role != "assistant" || row.provider_id.as_deref() != Some("opencode-go") {
180 continue;
181 }
182
183 let Some(cost_usd) = row.cost else {
184 continue;
185 };
186 if cost_usd <= 0.0 {
187 continue;
188 }
189
190 let completed = row.time.and_then(|time| time.completed);
191 let Some(completed_at) = completed.and_then(DateTime::from_timestamp_millis) else {
192 continue;
193 };
194
195 records.push(UsageRecord {
196 completed_at,
197 cost_usd,
198 });
199 }
200
201 records.sort_by_key(|record| record.completed_at);
202 Ok(records)
203 }
204
205 fn snapshot_from_records(
206 now: DateTime<Utc>,
207 records: &[UsageRecord],
208 credentials_available: bool,
209 ) -> OpencodeGoUsageSnapshot {
210 let windows = WINDOW_SPECS
211 .iter()
212 .map(|spec| Self::window_from_records(now, records, *spec))
213 .collect();
214
215 OpencodeGoUsageSnapshot {
216 source: OpencodeGoUsageSource::LocalDatabase,
217 credentials_available,
218 total_messages: records.len(),
219 windows,
220 }
221 }
222
223 fn window_from_records(
224 now: DateTime<Utc>,
225 records: &[UsageRecord],
226 spec: WindowSpec,
227 ) -> OpencodeGoUsageWindow {
228 let duration = Duration::seconds(spec.window_seconds);
229 let window_start = now - duration;
230 let active_records: Vec<&UsageRecord> = records
231 .iter()
232 .filter(|record| record.completed_at >= window_start)
233 .collect();
234 let spent_usd = active_records
235 .iter()
236 .map(|record| record.cost_usd)
237 .sum::<f64>();
238
239 let resets_at = if spent_usd + LIMIT_EPSILON >= spec.limit_usd {
240 let mut remaining = spent_usd;
241 active_records.iter().find_map(|record| {
242 remaining -= record.cost_usd;
243 if remaining + LIMIT_EPSILON < spec.limit_usd {
244 Some(record.completed_at + duration)
245 } else {
246 None
247 }
248 })
249 } else {
250 None
251 };
252
253 OpencodeGoUsageWindow {
254 entry_type: spec.entry_type,
255 spent_usd,
256 limit_usd: spec.limit_usd,
257 resets_at,
258 }
259 }
260}
261
262#[cfg(test)]
263#[expect(clippy::unwrap_used)]
264mod tests {
265 use super::*;
266 use chrono::TimeZone;
267
268 type TestResult = Result<(), Box<dyn std::error::Error>>;
269
270 fn usage_record(ts: i64, cost_usd: f64) -> UsageRecord {
271 UsageRecord {
272 completed_at: Utc.timestamp_millis_opt(ts).single().unwrap(),
273 cost_usd,
274 }
275 }
276
277 #[test]
278 fn snapshot_is_empty_when_no_messages_exist() {
279 let now = Utc.timestamp_millis_opt(1_000_000).single().unwrap();
280 let snapshot = OpencodeGoUsageStore::snapshot_from_records(now, &[], false);
281
282 assert_eq!(snapshot.total_messages, 0);
283 assert_eq!(snapshot.windows.len(), 3);
284 assert!(snapshot.windows.iter().all(|window| !window.is_limited()));
285 assert!(
286 snapshot
287 .windows
288 .iter()
289 .all(|window| window.spent_usd == 0.0)
290 );
291 }
292
293 #[test]
294 fn computes_five_hour_reset_from_oldest_blocking_message() {
295 let now = Utc
296 .timestamp_millis_opt(20 * 60 * 60 * 1000)
297 .single()
298 .unwrap();
299 let records = vec![
300 usage_record(15 * 60 * 60 * 1000, 4.0),
301 usage_record(16 * 60 * 60 * 1000, 5.0),
302 usage_record(19 * 60 * 60 * 1000, 4.0),
303 ];
304
305 let snapshot = OpencodeGoUsageStore::snapshot_from_records(now, &records, true);
306 let five_hour = snapshot
307 .windows
308 .iter()
309 .find(|window| window.entry_type == "five_hour_spend")
310 .unwrap();
311
312 assert!(five_hour.is_limited());
313 assert_eq!(
314 five_hour.resets_at,
315 Some(records[0].completed_at + Duration::hours(5))
316 );
317 assert!((five_hour.spent_usd - 13.0).abs() < LIMIT_EPSILON);
318 }
319
320 #[test]
321 fn computes_longer_windows_independently() {
322 let now = Utc
323 .timestamp_millis_opt(40 * 24 * 60 * 60 * 1000)
324 .single()
325 .unwrap();
326 let records = vec![
327 usage_record(10 * 24 * 60 * 60 * 1000, 31.0),
328 usage_record(34 * 24 * 60 * 60 * 1000, 11.0),
329 usage_record(35 * 24 * 60 * 60 * 1000, 10.0),
330 usage_record(39 * 24 * 60 * 60 * 1000, 10.0),
331 ];
332
333 let snapshot = OpencodeGoUsageStore::snapshot_from_records(now, &records, true);
334 let weekly = snapshot
335 .windows
336 .iter()
337 .find(|window| window.entry_type == "weekly_spend")
338 .unwrap();
339 let monthly = snapshot
340 .windows
341 .iter()
342 .find(|window| window.entry_type == "monthly_spend")
343 .unwrap();
344
345 assert!(weekly.is_limited());
346 assert_eq!(
347 weekly.resets_at,
348 Some(records[1].completed_at + Duration::days(7))
349 );
350 assert!(monthly.is_limited());
351 assert_eq!(
352 monthly.resets_at,
353 Some(records[0].completed_at + Duration::days(30))
354 );
355 }
356
357 #[test]
358 fn reads_only_opencode_go_assistant_messages_from_sqlite() -> TestResult {
359 let tmp = tempfile::NamedTempFile::new()?;
360 let conn = Connection::open(tmp.path())?;
361 conn.execute("CREATE TABLE message (data TEXT NOT NULL)", [])?;
362 conn.execute(
363 "INSERT INTO message (data) VALUES (?1)",
364 [r#"{"role":"assistant","providerID":"opencode-go","cost":1.5,"time":{"completed":3600000}}"#],
365 )?;
366 conn.execute(
367 "INSERT INTO message (data) VALUES (?1)",
368 [r#"{"role":"assistant","providerID":"opencode","cost":9.0,"time":{"completed":3600000}}"#],
369 )?;
370 conn.execute(
371 "INSERT INTO message (data) VALUES (?1)",
372 [r#"{"role":"user","providerID":"opencode-go","cost":9.0,"time":{"completed":3600000}}"#],
373 )?;
374 drop(conn);
375
376 let snapshot = OpencodeGoUsageStore::fetch_usage_from_path_at(
377 tmp.path(),
378 Utc.timestamp_millis_opt(10 * 60 * 60 * 1000)
379 .single()
380 .unwrap(),
381 )?;
382
383 assert_eq!(snapshot.total_messages, 1);
384 let five_hour = snapshot
385 .windows
386 .iter()
387 .find(|window| window.entry_type == "five_hour_spend")
388 .ok_or("missing five_hour window")?;
389 assert!((five_hour.spent_usd - 0.0).abs() < LIMIT_EPSILON);
390 let weekly = snapshot
391 .windows
392 .iter()
393 .find(|window| window.entry_type == "weekly_spend")
394 .ok_or("missing weekly window")?;
395 assert!((weekly.spent_usd - 1.5).abs() < LIMIT_EPSILON);
396 Ok(())
397 }
398}