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    /// Which local client ran the call — `opencode`, `dsh` or `pi`. Only rows on
27    /// [`PROVIDERS`] reach this struct, so a client that used the plan is named in
28    /// the report rather than silently inflating it.
29    pub client: &'static str,
30    pub at_ms: i64,
31    pub session: String,
32    pub provider: String,
33    pub model: String,
34    pub input: u64,
35    pub output: u64,
36    pub cache_read: u64,
37    pub cache_write: u64,
38}
39
40#[derive(Deserialize)]
41struct V2AssistantData {
42    model: Option<V2Model>,
43    tokens: Option<Tokens>,
44}
45
46#[derive(Deserialize)]
47struct V2Model {
48    id: Option<String>,
49    #[serde(rename = "providerID")]
50    provider: Option<String>,
51}
52
53#[derive(Deserialize)]
54struct AssistantData {
55    role: Option<String>,
56    #[serde(rename = "providerID")]
57    provider: Option<String>,
58    #[serde(rename = "modelID")]
59    model: Option<String>,
60    tokens: Option<Tokens>,
61}
62
63/// `reasoning` is inside `output` in opencode's own accounting (the docs say so,
64/// and the store's `cost` bills it at the output rate), so it is not added
65/// again here — counting it twice would inflate AFP.
66#[derive(Deserialize)]
67struct Tokens {
68    input: Option<u64>,
69    output: Option<u64>,
70    cache: Option<Cache>,
71}
72
73#[derive(Deserialize)]
74struct Cache {
75    read: Option<u64>,
76    write: Option<u64>,
77}
78
79/// The only provider that consumes this plan. `agent-plan` is what opencode
80/// writes for ARK's plan base URL.
81pub const PROVIDERS: [&str; 1] = ["agent-plan"];
82
83/// opencode's data directory: `$XDG_DATA_HOME/opencode`, else the platform
84/// default.
85pub fn opencode_home() -> String {
86    if let Ok(dir) = std::env::var("XDG_DATA_HOME") {
87        if !dir.is_empty() {
88            return format!("{dir}/opencode");
89        }
90    }
91    crate::paths::home()
92        .join(".local/share/opencode")
93        .display()
94        .to_string()
95}
96
97/// `$OPENCODE_DB`, else the store under [`opencode_home`].
98pub fn default_path() -> String {
99    if let Ok(path) = std::env::var("OPENCODE_DB") {
100        if !path.is_empty() {
101            return path;
102        }
103    }
104    format!("{}/opencode.db", opencode_home())
105}
106
107/// Read every assistant row the plan billed, at or after `since_ms`.
108///
109/// The table picked here is the difference between a real figure and a silent
110/// zero: v2 stopped appending to `message` at the migration, and every report
111/// was zero until this branch existed. `time_created >= ?` keeps either scan
112/// bounded; rows are filtered by provider in Rust because the JSON payload is
113/// not indexed.
114pub fn read_rows(path: &str, since_ms: i64) -> Result<Vec<Row>, String> {
115    let connection = Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_ONLY)
116        .map_err(|e| format!("open {path}: {e}"))?;
117    if has_v2(&connection) {
118        read_v2(&connection, since_ms)
119    } else {
120        read_legacy(&connection, since_ms)
121    }
122}
123
124/// opencode's store plus the other local clients the plan billed.
125///
126/// `pi` and `dsh` keep their own session logs (see [`crate::stores`]); a client
127/// pointed at the plan's `agent-plan` provider is the same subscription, so its
128/// calls count. Anything on another gateway is not this plan's usage.
129pub fn read_all(path: &str, since_ms: i64) -> Result<Vec<Row>, String> {
130    let mut rows = read_rows(path, since_ms)?;
131    rows.extend(external(&crate::stores::Dirs::from_env(), since_ms));
132    Ok(rows)
133}
134
135/// Map the extra clients' turns onto this module's rows, dropping any model call
136/// that did not run on this plan's provider.
137fn external(dirs: &crate::stores::Dirs, since_ms: i64) -> Vec<Row> {
138    crate::stores::read(dirs, since_ms)
139        .into_iter()
140        .filter(|turn| PROVIDERS.contains(&turn.provider.as_str()))
141        .map(|turn| Row {
142            client: turn.client,
143            at_ms: turn.at_ms,
144            session: turn.session,
145            provider: turn.provider,
146            model: turn.model,
147            input: turn.input,
148            output: turn.output,
149            cache_read: turn.cache_read,
150            cache_write: turn.cache_write,
151        })
152        .collect()
153}
154
155/// Does this store use opencode v2's `session_message` table?
156fn has_v2(connection: &Connection) -> bool {
157    connection
158        .query_row(
159            "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'session_message'",
160            [],
161            |_| Ok(()),
162        )
163        .is_ok()
164}
165
166/// v2 rows: the row's `type` names it (no `$.role`), the model ref is nested,
167/// and only a completed turn carries `tokens` — an in-flight one is skipped
168/// rather than counted as a zero-token request.
169fn read_v2(connection: &Connection, since_ms: i64) -> Result<Vec<Row>, String> {
170    let mut statement = connection
171        .prepare(
172            "SELECT data, time_created, session_id FROM session_message
173             WHERE type = 'assistant' AND time_created >= ?",
174        )
175        .map_err(|e| format!("prepare: {e}"))?;
176    let mut rows = statement
177        .query([since_ms])
178        .map_err(|e| format!("query: {e}"))?;
179
180    let mut out = Vec::new();
181    while let Some(row) = rows.next().map_err(|e| format!("row: {e}"))? {
182        let data: String = row.get(0).map_err(|e| format!("data: {e}"))?;
183        let at_ms: i64 = row.get(1).map_err(|e| format!("time: {e}"))?;
184        let session: String = row.get(2).unwrap_or_default();
185        let Ok(parsed) = serde_json::from_str::<V2AssistantData>(&data) else {
186            continue;
187        };
188        let Some(tokens) = parsed.tokens else {
189            continue; // still streaming: not billed yet
190        };
191        let model = parsed.model.unwrap_or(V2Model {
192            id: None,
193            provider: None,
194        });
195        let provider = model.provider.unwrap_or_default();
196        if !PROVIDERS.contains(&provider.as_str()) {
197            continue;
198        }
199        out.push(Row {
200            client: crate::stores::OPENCODE,
201            at_ms,
202            session,
203            provider,
204            model: model.id.unwrap_or_default(),
205            input: tokens.input.unwrap_or(0),
206            output: tokens.output.unwrap_or(0),
207            cache_read: tokens.cache.as_ref().and_then(|c| c.read).unwrap_or(0),
208            cache_write: tokens.cache.and_then(|c| c.write).unwrap_or(0),
209        });
210    }
211    Ok(out)
212}
213
214/// Pre-v2 rows: `role` in the payload, flat `modelID` / `providerID`.
215fn read_legacy(connection: &Connection, since_ms: i64) -> Result<Vec<Row>, String> {
216    let mut statement = connection
217        .prepare(
218            "SELECT data, time_created, session_id FROM message
219             WHERE data LIKE '%\"tokens\"%' AND time_created >= ?",
220        )
221        .map_err(|e| format!("prepare: {e}"))?;
222    let mut rows = statement
223        .query([since_ms])
224        .map_err(|e| format!("query: {e}"))?;
225
226    let mut out = Vec::new();
227    while let Some(row) = rows.next().map_err(|e| format!("row: {e}"))? {
228        let data: String = row.get(0).map_err(|e| format!("data: {e}"))?;
229        let at_ms: i64 = row.get(1).map_err(|e| format!("time: {e}"))?;
230        let session: String = row.get(2).unwrap_or_default();
231        let Ok(parsed) = serde_json::from_str::<AssistantData>(&data) else {
232            continue;
233        };
234        if parsed.role.as_deref() != Some("assistant") {
235            continue;
236        }
237        let provider = parsed.provider.unwrap_or_default();
238        if !PROVIDERS.contains(&provider.as_str()) {
239            continue;
240        }
241        let tokens = parsed.tokens.unwrap_or(Tokens {
242            input: None,
243            output: None,
244            cache: None,
245        });
246        out.push(Row {
247            client: crate::stores::OPENCODE,
248            at_ms,
249            session,
250            provider,
251            model: parsed.model.unwrap_or_default(),
252            input: tokens.input.unwrap_or(0),
253            output: tokens.output.unwrap_or(0),
254            cache_read: tokens.cache.as_ref().and_then(|c| c.read).unwrap_or(0),
255            cache_write: tokens.cache.and_then(|c| c.write).unwrap_or(0),
256        });
257    }
258    Ok(out)
259}
260
261#[cfg(test)]
262mod tests {
263    use super::*;
264
265    /// A v2 store: `session_message`, role in `type`, model ref nested.
266    fn v2_fixture() -> (tempfile::TempDir, String) {
267        let dir = tempfile::tempdir().unwrap();
268        let path = dir.path().join("opencode.db");
269        let connection = Connection::open(&path).unwrap();
270        connection
271            .execute_batch(
272                "CREATE TABLE session_message (\
273                 id TEXT, session_id TEXT, type TEXT, seq INTEGER, \
274                 time_created INTEGER, time_updated INTEGER, data TEXT);",
275            )
276            .unwrap();
277        let insert = |kind: &str, data: &str, at: i64| {
278            connection
279                .execute(
280                    "INSERT INTO session_message \
281                     (id, session_id, type, seq, time_created, time_updated, data) \
282                     VALUES (?, 'ses_v2', ?, 0, ?, ?, ?)",
283                    (format!("msg_{at}"), kind, at, at, data),
284                )
285                .unwrap();
286        };
287        insert(
288            "assistant",
289            r#"{"model":{"id":"deepseek-v4.1-flash","providerID":"agent-plan"},"tokens":{"input":10,"output":2,"cache":{"read":7,"write":0}}}"#,
290            1_000,
291        );
292        insert(
293            "assistant",
294            r#"{"model":{"id":"gpt-5.5","providerID":"opencode-go"},"tokens":{"input":1,"output":1}}"#,
295            2_000,
296        );
297        // In-flight: a model ref but no tokens yet.
298        insert(
299            "assistant",
300            r#"{"model":{"id":"deepseek-v4.1-flash","providerID":"agent-plan"}}"#,
301            3_000,
302        );
303        // The role lives in `type`, so this user row must not count.
304        insert("user", r#"{"text":"hi"}"#, 4_000);
305        (dir, path.to_string_lossy().into_owned())
306    }
307
308    #[test]
309    fn reads_only_plan_assistant_rows_in_window() {
310        let (_dir, path) = v2_fixture();
311        let rows = read_rows(&path, 0).unwrap();
312        assert_eq!(rows.len(), 1, "{rows:#?}");
313        assert_eq!(rows[0].provider, "agent-plan");
314        assert_eq!(rows[0].model, "deepseek-v4.1-flash");
315        assert_eq!(rows[0].cache_read, 7);
316
317        let later = read_rows(&path, 1_500).unwrap();
318        assert!(later.is_empty(), "the only plan row is at 1s");
319    }
320
321    #[test]
322    fn missing_database_is_an_error_not_a_panic() {
323        let error = read_rows("/nonexistent/opencode.db", 0).unwrap_err();
324        assert!(error.contains("open"), "{error}");
325    }
326
327    /// `pi` and `dsh` are the same plan only where the provider says so: a client
328    /// pointed at another gateway is not this subscription's usage.
329    #[test]
330    fn other_clients_keep_only_plan_rows() {
331        let dir = tempfile::tempdir().unwrap();
332        let dsh = dir.path().join("dsh/--p--/session-1");
333        std::fs::create_dir_all(&dsh).unwrap();
334        let log = concat!(
335            r#"{"type":"assistant/message","time":1000,"data":{"message":{"source":{"provider":"agent-plan","model":"deepseek-v4.1-flash"}},"usage":{"inputTokens":10,"outputTokens":2}}}"#,
336            "\n",
337            r#"{"type":"assistant/message","time":1000,"data":{"message":{"source":{"provider":"opencode-go","model":"deepseek-v4.1-flash"}},"usage":{"inputTokens":10,"outputTokens":2}}}"#,
338            "\n",
339        );
340        std::fs::write(dsh.join("session.v4.jsonl"), log.as_bytes()).unwrap();
341        let dirs = crate::stores::Dirs {
342            dsh: dir.path().join("dsh"),
343            pi: dir.path().join("pi"),
344        };
345        let rows = external(&dirs, 0);
346        assert_eq!(rows.len(), 1, "{rows:#?}");
347        assert_eq!(rows[0].client, crate::stores::DSH);
348        assert_eq!(rows[0].provider, "agent-plan");
349    }
350}
cmd_usage/arkuse/db.rs

cmd_usage/arkuse/
db.rs
