mj_controller/database/
usage_report.rs

1//! One consistent store snapshot of parent usage and retained descendants.
2use super::*;
3use mj_core::usage::{
4    ModelUsageEstimate, SessionUsageEstimate, SessionUsageQuery, SessionUsageReport, UsageScope,
5};
6
7mod pricing;
8
9pub fn load_session_usage_report(query: &SessionUsageQuery) -> Result<SessionUsageReport> {
10    load_from(
11        &database_path(),
12        query,
13        chrono::Utc::now().timestamp_millis(),
14    )
15}
16
17fn load_from(path: &Path, query: &SessionUsageQuery, now_ms: i64) -> Result<SessionUsageReport> {
18    let mut connection = open_reader(path)?;
19    let tx = connection.transaction()?;
20    let since_ms = now_ms.saturating_sub(query.period.duration_ms());
21    let mut report = SessionUsageReport {
22        since_ms,
23        until_ms: now_ms,
24        ..Default::default()
25    };
26    let mut roots = tx.prepare("SELECT session_id, COALESCE(session_title_override, acp_session_title, title), last_profile, state
27        FROM sessions s WHERE NOT EXISTS(SELECT 1 FROM subagent_accounting a WHERE a.child_session_id=s.session_id) ORDER BY session_id")?;
28    let roots = roots
29        .query_map([], |row| {
30            Ok((
31                row.get::<_, String>(0)?,
32                row.get::<_, String>(1)?,
33                row.get::<_, String>(2)?,
34                row.get::<_, String>(3)?,
35            ))
36        })?
37        .collect::<rusqlite::Result<Vec<_>>>()?;
38    report.profiles = roots
39        .iter()
40        .map(|(_, _, profile, _)| profile.clone())
41        .collect::<BTreeSet<_>>()
42        .into_iter()
43        .collect();
44    for (session_id, title, profile, state) in roots {
45        let suspended = !SessionState::from_stored(&state)
46            .context("unknown usage session state")?
47            .is_active();
48        if (suspended && !query.include_suspended)
49            || query
50                .profile
51                .as_ref()
52                .is_some_and(|wanted| wanted != &profile)
53        {
54            continue;
55        }
56        let mut tree = tx.prepare("WITH RECURSIVE tree(id) AS (
57            SELECT ?1 UNION SELECT a.child_session_id FROM subagent_accounting a JOIN tree t ON a.parent_session_id=t.id)
58            SELECT id FROM tree ORDER BY id")?;
59        let members = tree
60            .query_map([&session_id], |row| row.get::<_, String>(0))?
61            .collect::<rusqlite::Result<Vec<_>>>()?;
62        let mut estimate = SessionUsageEstimate {
63            session_id,
64            title,
65            profile,
66            suspended,
67            subagent_count: members.len().saturating_sub(1),
68            ..Default::default()
69        };
70        let mut groups: BTreeMap<Option<String>, ModelUsageEstimate> = BTreeMap::new();
71        for member in members {
72            let mut turns = tx.prepare(
73                "SELECT u.body, s.model, s.effort FROM session_turn_usage u
74                LEFT JOIN session_turn_selections s USING(session_id, command_id)
75                WHERE u.session_id=?1 AND json_extract(u.body, '$.completed_at_ms') >= ?2
76                AND json_extract(u.body, '$.completed_at_ms') <= ?3 ORDER BY u.completed_ordinal",
77            )?;
78            let rows = turns.query_map(params![member, since_ms, now_ms], |row| {
79                Ok((
80                    row.get::<_, String>(0)?,
81                    UsageSelection {
82                        model: row.get(1)?,
83                        effort: row.get(2)?,
84                    },
85                ))
86            })?;
87            for row in rows {
88                let (body, selection) = row?;
89                let turn = deserialize_turn_outcome(&body)?;
90                estimate.coverage.recorded_turns += 1;
91                match turn.usage {
92                    Some(usage) if usage.scope == UsageScope::Turn => {
93                        estimate.coverage.full_turn_reports += 1;
94                        let mut attributed = BTreeMap::new();
95                        super::usage::attribute(&usage, selection, &mut attributed)?;
96                        for (selection, totals) in attributed {
97                            // An unattributed cache counter is unknown, not zero.
98                            // Pricing the input as uncached would overstate its cost.
99                            let caches_known = (usage.cached_read_tokens.is_none()
100                                || totals.contains_key("cached_read_tokens"))
101                                && (usage.cached_write_tokens.is_none()
102                                    || totals.contains_key("cached_write_tokens"));
103                            let cost = if caches_known {
104                                pricing::estimate(selection.model.as_deref(), &totals)?
105                            } else {
106                                None
107                            };
108                            let group =
109                                groups.entry(selection.model.clone()).or_insert_with(|| {
110                                    ModelUsageEstimate {
111                                        model: selection.model,
112                                        cost_usd_ticks: Some(0),
113                                        ..Default::default()
114                                    }
115                                });
116                            group.cost_usd_ticks = match (group.cost_usd_ticks, cost) {
117                                (Some(a), Some(b)) => {
118                                    Some(a.checked_add(b).context("model cost overflow")?)
119                                }
120                                _ => None,
121                            };
122                            super::usage::merge_totals(&mut group.totals, &totals)?;
123                        }
124                    }
125                    Some(usage) if usage.scope == UsageScope::LastRequest => {
126                        estimate.coverage.last_request_reports += 1
127                    }
128                    Some(_) => estimate.coverage.unspecified_reports += 1,
129                    None => estimate.coverage.missing_reports += 1,
130                }
131            }
132            // A turn already running at the start of the window can still consume in it.
133            let unfinished: bool = tx.query_row("SELECT EXISTS(SELECT 1 FROM materialized_sessions WHERE session_id=?1
134                AND active_turn_json IS NOT NULL AND json_extract(active_turn_json, '$.started_at_ms') <= ?2)", params![member, now_ms], |row| row.get(0))?;
135            estimate.coverage.unfinished_turns += u64::from(unfinished);
136        }
137        estimate.models = groups.into_values().collect();
138        estimate.models.sort_by(|a, b| {
139            b.cost_usd_ticks
140                .cmp(&a.cost_usd_ticks)
141                .then_with(|| a.model.cmp(&b.model))
142        });
143        estimate.cost_usd_ticks = estimate
144            .models
145            .iter()
146            .filter_map(|model| model.cost_usd_ticks)
147            .try_fold(0u64, |total, cost| {
148                total.checked_add(cost).context("session cost overflow")
149            })?;
150        report.cost_usd_ticks = report
151            .cost_usd_ticks
152            .checked_add(estimate.cost_usd_ticks)
153            .context("report cost overflow")?;
154        report.sessions.push(estimate);
155    }
156    report.sessions.sort_by(|a, b| {
157        b.cost_usd_ticks
158            .cmp(&a.cost_usd_ticks)
159            .then_with(|| a.session_id.cmp(&b.session_id))
160    });
161    Ok(report)
162}
163
164#[cfg(test)]
165mod tests {
166    use super::*;
167    use mj_core::usage::UsagePeriod;
168
169    // Store/schema integration: retained descendants, inclusive time boundaries,
170    // model switches and incomplete reports must agree with the report reader.
171    #[test]
172    fn usage_report_preserves_retained_trees_and_filters_completed_usage() -> Result<()> {
173        let dir = tempfile::tempdir()?;
174        let path = dir.path().join("usage-report.sqlite");
175        let now = 31 * 86_400_000;
176        let mut state = State::default();
177        for id in ["root", "child", "grandchild", "suspended"] {
178            let mut session = super::super::tests::session(id, "project");
179            session.last_profile = if id == "suspended" { "claude" } else { "codex" }.into();
180            session.state = if id == "suspended" {
181                SessionState::Stopped
182            } else {
183                SessionState::Running
184            };
185            state.sessions.insert(id.into(), session);
186        }
187        save_state_to(&path, &state)?;
188        let connection = open(&path)?;
189        connection.execute_batch("INSERT INTO subagent_accounting VALUES ('child','root','child'),('grandchild','child','grandchild')")?;
190        let insert = |session: &str,
191                      seq: u64,
192                      timestamp: i64,
193                      model: &str,
194                      usage: serde_json::Value|
195         -> Result<()> {
196            let command = format!("command-{seq}");
197            let body = serde_json::json!({"command_id":command,"completed_ordinal":seq,"completed_at_ms":timestamp,
198                "outcome":{"kind":"completed","stop_reason":"EndTurn"},"usage":usage});
199            let turn: MaterializedTurnOutcome = serde_json::from_value(body)?;
200            connection.execute("INSERT INTO session_turn_usage(session_id,command_id,completed_ordinal,body) VALUES (?1,?2,?3,?4)", params![session, command, seq, serialize_turn_outcome(&turn)?])?;
201            connection.execute(
202                "INSERT INTO session_turn_selections VALUES (?1,?2,?3,'high')",
203                params![session, command, model],
204            )?;
205            Ok(())
206        };
207        let inclusive = serde_json::json!({"scope":"turn","total_tokens":1_100_000,"input_tokens":1_000_000,"output_tokens":100_000,"cached_read_tokens":500_000,"cached_write_tokens":0});
208        insert(
209            "root",
210            1,
211            now - 3_600_000,
212            "global.openai.gpt-6-luna",
213            inclusive.clone(),
214        )?;
215        insert("root", 2, now - 3_600_001, "gpt-6-sol", inclusive.clone())?;
216        insert("root", 3, now + 1, "gpt-6-astra", inclusive.clone())?;
217        insert(
218            "child",
219            1,
220            now,
221            "claude-opus-5-5-20260922",
222            serde_json::json!({"scope":"turn","total_tokens":1_800_000,"input_tokens":1_000_000,"output_tokens":100_000,"cached_read_tokens":500_000,"cached_write_tokens":200_000}),
223        )?;
224        insert("grandchild", 1, now, "custom-model", inclusive.clone())?;
225        let mut missing_cache_breakdown = inclusive.clone();
226        missing_cache_breakdown["provider_details"] = serde_json::json!({"model_usage":{"gpt-6.1-sol":{
227            "scope":"turn","total_tokens":1_100_000,"input_tokens":1_000_000,"output_tokens":100_000
228        }}});
229        insert(
230            "grandchild",
231            2,
232            now,
233            "custom-model",
234            missing_cache_breakdown,
235        )?;
236        insert(
237            "root",
238            4,
239            now,
240            "gpt-6-astra",
241            serde_json::json!({"scope":"last_request","total_tokens":1_100_000,"input_tokens":1_000_000,"output_tokens":100_000}),
242        )?;
243        insert("root", 5, now, "gpt-6-astra", serde_json::Value::Null)?;
244        insert(
245            "suspended",
246            1,
247            now - 2 * 86_400_000,
248            "gpt-6-astra",
249            inclusive.clone(),
250        )?;
251        insert(
252            "suspended",
253            2,
254            now - 8 * 86_400_000,
255            "gpt-6-sol",
256            inclusive.clone(),
257        )?;
258        insert(
259            "suspended",
260            3,
261            now - 31 * 86_400_000,
262            "gpt-6-astra",
263            inclusive,
264        )?;
265        drop(connection);
266        delete_session_from(&path, "grandchild")?;
267        delete_session_from(&path, "child")?;
268
269        let hourly = load_from(&path, &SessionUsageQuery::default(), now)?;
270        assert_eq!(hourly.sessions.len(), 1);
271        let root = &hourly.sessions[0];
272        assert_eq!(root.subagent_count, 2);
273        assert_eq!(root.cost_usd_ticks, 72_050_000_000); // $7.10 + $0.105, with caches counted once.
274        assert_eq!(root.models.len(), 5);
275        assert!(
276            root.models
277                .iter()
278                .find(|m| m.model.as_deref() == Some("gpt-6.1-sol"))
279                .unwrap()
280                .cost_usd_ticks
281                .is_none()
282        );
283        assert!(root.models.last().unwrap().cost_usd_ticks.is_none());
284        assert_eq!(root.coverage.full_turn_reports, 4);
285        assert_eq!(root.coverage.last_request_reports, 1);
286        assert_eq!(root.coverage.missing_reports, 1);
287        assert!(root.is_partial());
288        assert_eq!(hourly.profiles, ["claude", "codex"]);
289        assert_eq!(hourly.cost_usd_ticks, root.cost_usd_ticks);
290
291        let weekly_query = SessionUsageQuery {
292            period: UsagePeriod::Weekly,
293            include_suspended: true,
294            profile: None,
295        };
296        let weekly = load_from(&path, &weekly_query, now)?;
297        assert_eq!(weekly.sessions.len(), 2);
298        // Cache-discounted Astra: $5 + $0.50 + $5 = $10.50, greater than root.
299        assert_eq!(weekly.sessions[0].session_id, "suspended");
300        assert_eq!(weekly.sessions[0].cost_usd_ticks, 105_000_000_000);
301        assert!(weekly.sessions[0].suspended);
302        assert_eq!(weekly.sessions[1].cost_usd_ticks, 93_050_000_000);
303        assert_eq!(weekly.cost_usd_ticks, 198_050_000_000);
304        let filtered = load_from(
305            &path,
306            &SessionUsageQuery {
307                profile: Some("codex".into()),
308                ..weekly_query.clone()
309            },
310            now,
311        )?;
312        assert_eq!(filtered.sessions, [weekly.sessions[1].clone()]);
313        let monthly = load_from(
314            &path,
315            &SessionUsageQuery {
316                period: UsagePeriod::Monthly,
317                profile: Some("claude".into()),
318                ..weekly_query
319            },
320            now,
321        )?;
322        assert_eq!(monthly.sessions[0].cost_usd_ticks, 126_000_000_000);
323        assert_eq!(monthly.sessions[0].coverage.recorded_turns, 2);
324        Ok(())
325    }
326}
mj_controller/database/usage_report.rs

mj_controller/database/
usage_report.rs
