cmd_usage/arkuse/
db.rs

1// opencode's own message store, read-only.
2//
3// Every assistant row opencode writes carries a provider, model and `tokens`,
4// plus a timestamp — a complete, backfilled per-request record for everything
5// the agent ran on this machine. ARK's plan has its own currency (AFP) and
6// opencode cannot price it, so only the tokens are used; the coefficient table
7// in `core/ark.json` turns those into AFP.
8//
9// Two layouts: v1 `message` keeps them flat (`role` / `modelID` / `providerID`),
10// v2 `session_message` puts the role in the row's `type` and nests the model
11// (`model.id` / `model.providerID`). v2 stopped appending to `message` at the
12// migration, so reading only the legacy table reports zero for every recent
13// session; `read_rows` picks by table and keeps `message` as the fallback.
14//
15// Only `provider = "agent-plan"` rows are plan usage. `opencode-go/*` and
16// `opencode/*` rows are other services and are excluded here, never added to
17// AFP — they may appear in the same store.
18//
19// Read-only by construction: `OpenFlags::SQLITE_OPEN_READ_ONLY`, so a running
20// opencode is unaffected.
21use rusqlite::{Connection, OpenFlags};
22use serde::Deserialize;
23
24#[derive(Debug, Clone, Default)]
25pub struct Row {
26    pub at_ms: i64,
27    pub session: String,
28    pub provider: String,
29    pub model: String,
30    pub input: u64,
31    pub output: u64,
32    pub cache_read: u64,
33    pub cache_write: u64,
34}
35
36#[derive(Deserialize)]
37struct V2AssistantData {
38    model: Option<V2Model>,
39    tokens: Option<Tokens>,
40}
41
42#[derive(Deserialize)]
43struct V2Model {
44    id: Option<String>,
45    #[serde(rename = "providerID")]
46    provider: Option<String>,
47}
48
49#[derive(Deserialize)]
50struct AssistantData {
51    role: Option<String>,
52    #[serde(rename = "providerID")]
53    provider: Option<String>,
54    #[serde(rename = "modelID")]
55    model: Option<String>,
56    tokens: Option<Tokens>,
57}
58
59/// `reasoning` is inside `output` in opencode's own accounting (the docs say so,
60/// and the store's `cost` bills it at the output rate), so it is not added
61/// again here — counting it twice would inflate AFP.
62#[derive(Deserialize)]
63struct Tokens {
64    input: Option<u64>,
65    output: Option<u64>,
66    cache: Option<Cache>,
67}
68
69#[derive(Deserialize)]
70struct Cache {
71    read: Option<u64>,
72    write: Option<u64>,
73}
74
75/// The only provider that consumes this plan. `agent-plan` is what opencode
76/// writes for ARK's plan base URL.
77pub const PROVIDERS: [&str; 1] = ["agent-plan"];
78
79/// opencode's data directory: `$XDG_DATA_HOME/opencode`, else the platform
80/// default.
81pub fn opencode_home() -> String {
82    if let Ok(dir) = std::env::var("XDG_DATA_HOME") {
83        if !dir.is_empty() {
84            return format!("{dir}/opencode");
85        }
86    }
87    crate::paths::home()
88        .join(".local/share/opencode")
89        .display()
90        .to_string()
91}
92
93/// `$OPENCODE_DB`, else the store under [`opencode_home`].
94pub fn default_path() -> String {
95    if let Ok(path) = std::env::var("OPENCODE_DB") {
96        if !path.is_empty() {
97            return path;
98        }
99    }
100    format!("{}/opencode.db", opencode_home())
101}
102
103/// Read every assistant row the plan billed, at or after `since_ms`.
104///
105/// The table picked here is the difference between a real figure and a silent
106/// zero: v2 stopped appending to `message` at the migration, and every report
107/// was zero until this branch existed. `time_created >= ?` keeps either scan
108/// bounded; rows are filtered by provider in Rust because the JSON payload is
109/// not indexed.
110pub fn read_rows(path: &str, since_ms: i64) -> Result<Vec<Row>, String> {
111    let connection = Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_ONLY)
112        .map_err(|e| format!("open {path}: {e}"))?;
113    if has_v2(&connection) {
114        read_v2(&connection, since_ms)
115    } else {
116        read_legacy(&connection, since_ms)
117    }
118}
119
120/// Does this store use opencode v2's `session_message` table?
121fn has_v2(connection: &Connection) -> bool {
122    connection
123        .query_row(
124            "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'session_message'",
125            [],
126            |_| Ok(()),
127        )
128        .is_ok()
129}
130
131/// v2 rows: the row's `type` names it (no `$.role`), the model ref is nested,
132/// and only a completed turn carries `tokens` — an in-flight one is skipped
133/// rather than counted as a zero-token request.
134fn read_v2(connection: &Connection, since_ms: i64) -> Result<Vec<Row>, String> {
135    let mut statement = connection
136        .prepare(
137            "SELECT data, time_created, session_id FROM session_message
138             WHERE type = 'assistant' AND time_created >= ?",
139        )
140        .map_err(|e| format!("prepare: {e}"))?;
141    let mut rows = statement
142        .query([since_ms])
143        .map_err(|e| format!("query: {e}"))?;
144
145    let mut out = Vec::new();
146    while let Some(row) = rows.next().map_err(|e| format!("row: {e}"))? {
147        let data: String = row.get(0).map_err(|e| format!("data: {e}"))?;
148        let at_ms: i64 = row.get(1).map_err(|e| format!("time: {e}"))?;
149        let session: String = row.get(2).unwrap_or_default();
150        let Ok(parsed) = serde_json::from_str::<V2AssistantData>(&data) else {
151            continue;
152        };
153        let Some(tokens) = parsed.tokens else {
154            continue; // still streaming: not billed yet
155        };
156        let model = parsed.model.unwrap_or(V2Model {
157            id: None,
158            provider: None,
159        });
160        let provider = model.provider.unwrap_or_default();
161        if !PROVIDERS.contains(&provider.as_str()) {
162            continue;
163        }
164        out.push(Row {
165            at_ms,
166            session,
167            provider,
168            model: model.id.unwrap_or_default(),
169            input: tokens.input.unwrap_or(0),
170            output: tokens.output.unwrap_or(0),
171            cache_read: tokens.cache.as_ref().and_then(|c| c.read).unwrap_or(0),
172            cache_write: tokens.cache.and_then(|c| c.write).unwrap_or(0),
173        });
174    }
175    Ok(out)
176}
177
178/// Pre-v2 rows: `role` in the payload, flat `modelID` / `providerID`.
179fn read_legacy(connection: &Connection, since_ms: i64) -> Result<Vec<Row>, String> {
180    let mut statement = connection
181        .prepare(
182            "SELECT data, time_created, session_id FROM message
183             WHERE data LIKE '%\"tokens\"%' AND time_created >= ?",
184        )
185        .map_err(|e| format!("prepare: {e}"))?;
186    let mut rows = statement
187        .query([since_ms])
188        .map_err(|e| format!("query: {e}"))?;
189
190    let mut out = Vec::new();
191    while let Some(row) = rows.next().map_err(|e| format!("row: {e}"))? {
192        let data: String = row.get(0).map_err(|e| format!("data: {e}"))?;
193        let at_ms: i64 = row.get(1).map_err(|e| format!("time: {e}"))?;
194        let session: String = row.get(2).unwrap_or_default();
195        let Ok(parsed) = serde_json::from_str::<AssistantData>(&data) else {
196            continue;
197        };
198        if parsed.role.as_deref() != Some("assistant") {
199            continue;
200        }
201        let provider = parsed.provider.unwrap_or_default();
202        if !PROVIDERS.contains(&provider.as_str()) {
203            continue;
204        }
205        let tokens = parsed.tokens.unwrap_or(Tokens {
206            input: None,
207            output: None,
208            cache: None,
209        });
210        out.push(Row {
211            at_ms,
212            session,
213            provider,
214            model: parsed.model.unwrap_or_default(),
215            input: tokens.input.unwrap_or(0),
216            output: tokens.output.unwrap_or(0),
217            cache_read: tokens.cache.as_ref().and_then(|c| c.read).unwrap_or(0),
218            cache_write: tokens.cache.and_then(|c| c.write).unwrap_or(0),
219        });
220    }
221    Ok(out)
222}
223
224#[cfg(test)]
225mod tests {
226    use super::*;
227
228    /// A v2 store: `session_message`, role in `type`, model ref nested.
229    fn v2_fixture() -> (tempfile::TempDir, String) {
230        let dir = tempfile::tempdir().unwrap();
231        let path = dir.path().join("opencode.db");
232        let connection = Connection::open(&path).unwrap();
233        connection
234            .execute_batch(
235                "CREATE TABLE session_message (\
236                 id TEXT, session_id TEXT, type TEXT, seq INTEGER, \
237                 time_created INTEGER, time_updated INTEGER, data TEXT);",
238            )
239            .unwrap();
240        let insert = |kind: &str, data: &str, at: i64| {
241            connection
242                .execute(
243                    "INSERT INTO session_message \
244                     (id, session_id, type, seq, time_created, time_updated, data) \
245                     VALUES (?, 'ses_v2', ?, 0, ?, ?, ?)",
246                    (format!("msg_{at}"), kind, at, at, data),
247                )
248                .unwrap();
249        };
250        insert(
251            "assistant",
252            r#"{"model":{"id":"deepseek-v4.1-flash","providerID":"agent-plan"},"tokens":{"input":10,"output":2,"cache":{"read":7,"write":0}}}"#,
253            1_000,
254        );
255        insert(
256            "assistant",
257            r#"{"model":{"id":"gpt-5.5","providerID":"opencode-go"},"tokens":{"input":1,"output":1}}"#,
258            2_000,
259        );
260        // In-flight: a model ref but no tokens yet.
261        insert(
262            "assistant",
263            r#"{"model":{"id":"deepseek-v4.1-flash","providerID":"agent-plan"}}"#,
264            3_000,
265        );
266        // The role lives in `type`, so this user row must not count.
267        insert("user", r#"{"text":"hi"}"#, 4_000);
268        (dir, path.to_string_lossy().into_owned())
269    }
270
271    #[test]
272    fn reads_only_plan_assistant_rows_in_window() {
273        let (_dir, path) = v2_fixture();
274        let rows = read_rows(&path, 0).unwrap();
275        assert_eq!(rows.len(), 1, "{rows:#?}");
276        assert_eq!(rows[0].provider, "agent-plan");
277        assert_eq!(rows[0].model, "deepseek-v4.1-flash");
278        assert_eq!(rows[0].cache_read, 7);
279
280        let later = read_rows(&path, 1_500).unwrap();
281        assert!(later.is_empty(), "the only plan row is at 1s");
282    }
283
284    #[test]
285    fn missing_database_is_an_error_not_a_panic() {
286        let error = read_rows("/nonexistent/opencode.db", 0).unwrap_err();
287        assert!(error.contains("open"), "{error}");
288    }
289}
cmd_usage/arkuse/db.rs

cmd_usage/arkuse/
db.rs
