mj_controller/database/
usage.rs

1use super::*;
2#[cfg(test)]
3use mj_core::usage::ProviderCost;
4use mj_core::usage::{TokenUsage, UsageScope};
5
6pub(super) fn record_child_accounting(
7    connection: &Connection,
8    child: &SubagentRecord,
9) -> Result<()> {
10    connection.execute(
11        "INSERT INTO subagent_accounting(child_session_id, parent_session_id, task_name)
12        VALUES (?1, ?2, ?3) ON CONFLICT(child_session_id) DO NOTHING",
13        params![
14            child.child_session_id,
15            child.parent_session_id,
16            child.task_name
17        ],
18    )?;
19    let stored: (String, String) = connection.query_row(
20        "SELECT parent_session_id, task_name FROM subagent_accounting WHERE child_session_id=?1",
21        [&child.child_session_id],
22        |row| Ok((row.get(0)?, row.get(1)?)),
23    )?;
24    ensure!(
25        stored == (child.parent_session_id.clone(), child.task_name.clone()),
26        "accounting identity changed for child {}",
27        child.child_session_id
28    );
29    Ok(())
30}
31
32pub(super) fn record_turn_selection(
33    connection: &Connection,
34    session: &str,
35    command: &str,
36    configuration: &mj_core::state::SessionConfiguration,
37) -> Result<()> {
38    let surface =
39        mj_core::acp::surface::AcpSessionSurface::from_configuration(&configuration.values);
40    connection.execute(
41        "INSERT INTO session_turn_selections(session_id, command_id, model, effort)
42        VALUES (?1, ?2, ?3, ?4) ON CONFLICT(session_id, command_id) DO NOTHING",
43        params![
44            session,
45            command,
46            surface.current_model(),
47            surface.current_effort()
48        ],
49    )?;
50    Ok(())
51}
52
53pub fn load_session_usage(
54    session_id: &str,
55    after_seq: u64,
56    limit: usize,
57) -> Result<Option<UsagePage>> {
58    load_session_usage_from(&database_path(), session_id, after_seq, limit)
59}
60
61fn context_exists(connection: &Connection, session: &str) -> Result<bool> {
62    Ok(connection.query_row(
63        "SELECT EXISTS(SELECT 1 FROM session_contexts WHERE session_id=?1)",
64        [session],
65        |row| row.get(0),
66    )?)
67}
68
69fn load_session_usage_from(
70    path: &Path,
71    session_id: &str,
72    after_seq: u64,
73    limit: usize,
74) -> Result<Option<UsagePage>> {
75    let mut connection = open_reader(path)?;
76    let tx = connection.transaction()?;
77    if !context_exists(&tx, session_id)? {
78        return Ok(None);
79    }
80    let latest_seq: u64 = tx.query_row(
81        "SELECT COALESCE(MAX(completed_ordinal), 0) FROM session_turn_usage WHERE session_id=?1",
82        [session_id],
83        |row| row.get(0),
84    )?;
85    let mut statement = tx.prepare("SELECT body FROM session_turn_usage WHERE session_id=?1 AND completed_ordinal > ?2 ORDER BY completed_ordinal LIMIT ?3")?;
86    let turns = statement
87        .query_map(
88            params![session_id, after_seq, limit.clamp(1, 1000) as i64],
89            |row| row.get::<_, String>(0),
90        )?
91        .map(|row| deserialize_turn_outcome(&row?))
92        .collect::<Result<Vec<_>>>()?;
93    let mut turn_selections = BTreeMap::new();
94    for turn in &turns {
95        turn_selections.insert(
96            turn.command_id.clone(),
97            selection_for(&tx, session_id, &turn.command_id)?,
98        );
99    }
100    let summary = session_summary(&tx, session_id)?;
101    Ok(Some(UsagePage {
102        session_id: session_id.into(),
103        next_after_seq: turns
104            .last()
105            .map_or(latest_seq.max(after_seq), |turn| turn.completed_ordinal),
106        latest_seq,
107        turns,
108        totals: summary.totals,
109        coverage: summary.coverage,
110        provider_session_cost: summary.provider_session_cost,
111        turn_selections,
112        by_model: summary.by_model,
113    }))
114}
115
116fn selection_for(connection: &Connection, session: &str, command: &str) -> Result<UsageSelection> {
117    Ok(connection.query_row("SELECT model, effort FROM session_turn_selections WHERE session_id=?1 AND command_id=?2",
118        params![session, command], |row| Ok(UsageSelection { model: row.get(0)?, effort: row.get(1)? })).optional()?.unwrap_or_default())
119}
120
121fn counters(usage: &TokenUsage) -> BTreeMap<String, u64> {
122    let mut values = BTreeMap::from([
123        ("total_tokens".into(), usage.total_tokens),
124        ("input_tokens".into(), usage.input_tokens),
125        ("output_tokens".into(), usage.output_tokens),
126    ]);
127    for (name, value) in [
128        ("thought_tokens", usage.thought_tokens),
129        ("cached_read_tokens", usage.cached_read_tokens),
130        ("cached_write_tokens", usage.cached_write_tokens),
131    ] {
132        if let Some(value) = value {
133            values.insert(name.into(), value);
134        }
135    }
136    values
137}
138
139fn add_counters(
140    totals: &mut BTreeMap<String, UsageCounterTotal>,
141    values: &BTreeMap<String, u64>,
142) -> Result<()> {
143    for (name, tokens) in values {
144        let total = totals.entry(name.clone()).or_insert(UsageCounterTotal {
145            tokens: 0,
146            reported_turns: 0,
147        });
148        total.tokens = total
149            .tokens
150            .checked_add(*tokens)
151            .context("usage token total overflow")?;
152        total.reported_turns = total
153            .reported_turns
154            .checked_add(1)
155            .context("usage turn total overflow")?;
156    }
157    Ok(())
158}
159
160/// Attribute each counter once. Partial or inconsistent breakdowns retain an
161/// unknown remainder instead of assigning it to the configured model.
162fn attribute(
163    usage: &TokenUsage,
164    selected: UsageSelection,
165    groups: &mut BTreeMap<UsageSelection, BTreeMap<String, UsageCounterTotal>>,
166) -> Result<()> {
167    let values = counters(usage);
168    let breakdown = usage
169        .provider_details
170        .as_ref()
171        .map(|details| &details.model_usage)
172        .filter(|models| !models.is_empty());
173    let Some(models) = breakdown else {
174        return add_counters(groups.entry(selected).or_default(), &values);
175    };
176    let mut assigned: BTreeMap<UsageSelection, BTreeMap<String, u64>> = BTreeMap::new();
177    for (counter, total) in values {
178        let parts = models
179            .iter()
180            .map(|(model, report)| (model, report.scope, counters(report).get(&counter).copied()))
181            .collect::<Vec<_>>();
182        let sum = parts.iter().try_fold(0u64, |sum, (_, scope, count)| {
183            if *scope != UsageScope::Turn {
184                return None;
185            }
186            sum.checked_add((*count)?)
187        });
188        if let Some(sum) = sum.filter(|sum| *sum <= total) {
189            for (model, _, count) in parts {
190                let selection = UsageSelection {
191                    model: Some(model.clone()),
192                    effort: (selected.model.as_ref() == Some(model))
193                        .then(|| selected.effort.clone())
194                        .flatten(),
195                };
196                assigned
197                    .entry(selection)
198                    .or_default()
199                    .insert(counter.clone(), count.expect("validated counter"));
200            }
201            if sum < total {
202                assigned
203                    .entry(UsageSelection::default())
204                    .or_default()
205                    .insert(counter, total - sum);
206            }
207        } else {
208            assigned
209                .entry(UsageSelection::default())
210                .or_default()
211                .insert(counter, total);
212        }
213    }
214    for (selection, values) in assigned {
215        add_counters(groups.entry(selection).or_default(), &values)?;
216    }
217    Ok(())
218}
219
220fn session_summary(connection: &Connection, session: &str) -> Result<UsageSummary> {
221    let mut summary = UsageSummary::default();
222    let mut groups = BTreeMap::new();
223    let mut statement = connection.prepare("SELECT u.body, s.model, s.effort FROM session_turn_usage u
224        LEFT JOIN session_turn_selections s USING(session_id, command_id) WHERE u.session_id=?1 ORDER BY u.completed_ordinal")?;
225    let rows = statement.query_map([session], |row| {
226        Ok((
227            row.get::<_, String>(0)?,
228            UsageSelection {
229                model: row.get(1)?,
230                effort: row.get(2)?,
231            },
232        ))
233    })?;
234    for row in rows {
235        let (body, selection) = row?;
236        let turn = deserialize_turn_outcome(&body)?;
237        summary.coverage.recorded_turns += 1;
238        match turn.usage {
239            Some(usage) if usage.scope == UsageScope::Turn => {
240                summary.coverage.full_turn_reports += 1;
241                add_counters(&mut summary.totals, &counters(&usage))?;
242                attribute(&usage, selection, &mut groups)?;
243            }
244            Some(usage) if usage.scope == UsageScope::LastRequest => {
245                summary.coverage.last_request_reports += 1
246            }
247            Some(_) => summary.coverage.unspecified_reports += 1,
248            None => summary.coverage.missing_reports += 1,
249        }
250    }
251    summary.coverage.unfinished_turns = connection.query_row("SELECT COUNT(*) FROM session_turn_selections s
252        WHERE session_id=?1 AND NOT EXISTS(SELECT 1 FROM session_turn_usage u WHERE u.session_id=s.session_id AND u.command_id=s.command_id)",
253        [session], |row| row.get(0))?;
254    summary.by_model = groups
255        .into_iter()
256        .map(|(selection, totals)| UsageModelTotal { selection, totals })
257        .collect();
258    let cost: Option<String> = connection
259        .query_row(
260            "SELECT body FROM session_provider_cost WHERE session_id=?1",
261            [session],
262            |row| row.get(0),
263        )
264        .optional()?;
265    summary.provider_session_cost = cost.map(|body| serde_json::from_str(&body)).transpose()?;
266    Ok(summary)
267}
268
269pub fn load_usage_tree(parent: &str) -> Result<Option<UsageTree>> {
270    load_usage_tree_from(&database_path(), parent)
271}
272
273fn load_usage_tree_from(path: &Path, parent: &str) -> Result<Option<UsageTree>> {
274    let mut connection = open_reader(path)?;
275    let tx = connection.transaction()?;
276    if !context_exists(&tx, parent)? {
277        return Ok(None);
278    }
279    let mut statement = tx.prepare("WITH RECURSIVE tree(id) AS (
280        SELECT ?1 UNION SELECT a.child_session_id FROM subagent_accounting a JOIN tree t ON a.parent_session_id=t.id)
281        SELECT t.id, a.parent_session_id, a.task_name, EXISTS(SELECT 1 FROM sessions s WHERE s.session_id=t.id)
282        FROM tree t LEFT JOIN subagent_accounting a ON a.child_session_id=t.id ORDER BY t.id")?;
283    let identities = statement
284        .query_map([parent], |row| {
285            Ok((
286                row.get::<_, String>(0)?,
287                row.get::<_, Option<String>>(1)?,
288                row.get::<_, Option<String>>(2)?,
289                row.get::<_, bool>(3)?,
290            ))
291        })?
292        .collect::<rusqlite::Result<Vec<_>>>()?;
293    let mut summary = UsageSummary::default();
294    let mut groups: BTreeMap<UsageSelection, BTreeMap<String, UsageCounterTotal>> = BTreeMap::new();
295    let mut sessions = Vec::new();
296    for (session_id, parent_session_id, task_name, operational_session_present) in identities {
297        let member = session_summary(&tx, &session_id)?;
298        merge_summary(&mut summary, &member)?;
299        for group in &member.by_model {
300            merge_totals(
301                groups.entry(group.selection.clone()).or_default(),
302                &group.totals,
303            )?;
304        }
305        sessions.push(UsageTreeSession {
306            session_id,
307            parent_session_id,
308            task_name,
309            operational_session_present,
310            summary: member,
311        });
312    }
313    summary.by_model = groups
314        .into_iter()
315        .map(|(selection, totals)| UsageModelTotal { selection, totals })
316        .collect();
317    Ok(Some(UsageTree {
318        parent_session_id: parent.into(),
319        sessions,
320        summary,
321    }))
322}
323
324fn merge_totals(
325    target: &mut BTreeMap<String, UsageCounterTotal>,
326    source: &BTreeMap<String, UsageCounterTotal>,
327) -> Result<()> {
328    for (name, value) in source {
329        let total = target.entry(name.clone()).or_insert(UsageCounterTotal {
330            tokens: 0,
331            reported_turns: 0,
332        });
333        total.tokens = total
334            .tokens
335            .checked_add(value.tokens)
336            .context("usage tree token overflow")?;
337        total.reported_turns = total
338            .reported_turns
339            .checked_add(value.reported_turns)
340            .context("usage tree turn overflow")?;
341    }
342    Ok(())
343}
344
345fn merge_summary(target: &mut UsageSummary, source: &UsageSummary) -> Result<()> {
346    merge_totals(&mut target.totals, &source.totals)?;
347    let a = &mut target.coverage;
348    let b = &source.coverage;
349    for (target, value) in [
350        (&mut a.recorded_turns, b.recorded_turns),
351        (&mut a.full_turn_reports, b.full_turn_reports),
352        (&mut a.last_request_reports, b.last_request_reports),
353        (&mut a.unspecified_reports, b.unspecified_reports),
354        (&mut a.missing_reports, b.missing_reports),
355        (&mut a.unfinished_turns, b.unfinished_turns),
356    ] {
357        *target = target
358            .checked_add(value)
359            .context("usage coverage overflow")?;
360    }
361    Ok(())
362}
363
364#[cfg(test)]
365mod tests {
366    use super::*;
367    use mj_core::state::TurnOutcomeKind;
368    use mj_core::usage::{TokenUsage, UsageScope};
369
370    fn completed(command: &str, ordinal: u64, input: u64) -> MaterializedTurnOutcome {
371        serde_json::from_value(serde_json::json!({
372            "command_id":command, "completed_ordinal":ordinal, "completed_at_ms":100,
373            "outcome":{"kind":"completed","stop_reason":"EndTurn"},
374            "usage":{"scope":"turn","total_tokens":input+10,"input_tokens":input,"output_tokens":10}
375        }))
376        .unwrap()
377    }
378
379    fn configuration(model: &str, effort: &str) -> mj_core::state::SessionConfiguration {
380        serde_json::from_value(serde_json::json!({"model":model,"effort":effort})).unwrap()
381    }
382
383    #[test]
384    fn semantic_turn_migration_preserves_all_billed_turns_and_exact_last_turns() -> Result<()> {
385        let directory = tempfile::tempdir()?;
386        let path = directory.path().join("semantic-turns.sqlite3");
387        let turns = super::super::turn_outcome::tests::legacy_turns();
388        for turn in &turns {
389            save_session_to(
390                &path,
391                &super::super::tests::session(&turn.command_id, "project"),
392            )?;
393            let mut session = MaterializedSession::empty(&turn.command_id);
394            session.last_turn_outcome = Some(turn.clone());
395            save_materialized_session_to(&path, &session)?;
396        }
397        let connection = Connection::open(&path)?;
398        for turn in &turns {
399            connection.execute(
400                "INSERT INTO session_turn_usage VALUES (?1,?1,?2,?3,?4)",
401                params![
402                    turn.command_id,
403                    turn.completed_ordinal,
404                    turn.turn_start_position,
405                    serialize_turn_outcome(turn)?
406                ],
407            )?;
408        }
409        let before = turns
410            .iter()
411            .map(|turn| session_summary(&connection, &turn.command_id))
412            .collect::<Result<Vec<_>>>()?;
413        for turn in &turns {
414            let body = serde_json::to_string(turn)?;
415            connection.execute(
416                "UPDATE session_turn_usage SET body=?2 WHERE session_id=?1",
417                params![turn.command_id, body],
418            )?;
419            connection.execute(
420                "UPDATE materialized_sessions SET last_turn_outcome_json=?2 WHERE session_id=?1",
421                params![turn.command_id, body],
422            )?;
423        }
424        connection.execute_batch("DELETE FROM schema_migrations WHERE version=82; UPDATE schema_compatibility SET minimum_compatible_version=81; PRAGMA user_version=81;")?;
425        drop(connection);
426        super::super::schema::forget_verified_schema(&path);
427        let connection = open(&path)?;
428        let state = super::super::schema::read_schema_state(&connection)?;
429        assert_eq!(state.revision, 82);
430        assert_eq!(state.minimum_compatible, Some(82));
431        assert!(state.ensure_supported_by(81).is_err());
432        for (turn, before) in turns.iter().zip(before) {
433            assert_eq!(session_summary(&connection, &turn.command_id)?, before);
434            let page = load_session_usage_from(&path, &turn.command_id, 0, 100)?.unwrap();
435            assert_eq!(page.turns.as_slice(), std::slice::from_ref(turn));
436            assert_eq!(page.coverage.recorded_turns, 1);
437            assert_eq!(page.totals["input_tokens"].tokens, 42);
438            assert_eq!(
439                read_materialized_turn_state(&connection, &turn.command_id)?
440                    .unwrap()
441                    .2,
442                Some(turn.clone())
443            );
444            let body: String = connection.query_row(
445                "SELECT last_turn_outcome_json FROM materialized_sessions WHERE session_id=?1",
446                [&turn.command_id],
447                |row| row.get(0),
448            )?;
449            assert_eq!(
450                serde_json::from_str::<mj_core::event_outcome::ApiTurnOutcome>(&body)?,
451                mj_core::event_outcome::ApiTurnOutcome::from(turn)
452            );
453        }
454        Ok(())
455    }
456
457    #[test]
458    fn turn_selections_follow_event_order_and_survive_cleanup_and_replay() -> Result<()> {
459        let dir = tempfile::tempdir()?;
460        let path = dir.path().join("usage.sqlite");
461        save_session_to(&path, &super::super::tests::session("root", "project"))?;
462        save_session_to(&path, &super::super::tests::session("child", "project"))?;
463        let connection = open(&path)?;
464        connection.execute(
465            "INSERT INTO subagent_accounting VALUES ('child','root','investigate')",
466            [],
467        )?;
468        drop(connection);
469        let events = [
470            MaterializedSessionMutation {
471                configuration: Some(configuration("sol", "high")),
472                active_turn: Some(Some(MaterializedTurn {
473                    command_id: "one".into(),
474                    accepted_ordinal: None,
475                    turn_start_position: 1,
476                    started_at_ms: 1,
477                    steered_into: None,
478                })),
479                ..Default::default()
480            },
481            MaterializedSessionMutation {
482                active_turn: Some(None),
483                last_turn_outcome: Some(completed("one", 2, 20)),
484                ..Default::default()
485            },
486            MaterializedSessionMutation {
487                configuration: Some(configuration("luna", "low")),
488                active_turn: Some(Some(MaterializedTurn {
489                    command_id: "two".into(),
490                    accepted_ordinal: None,
491                    turn_start_position: 3,
492                    started_at_ms: 3,
493                    steered_into: None,
494                })),
495                ..Default::default()
496            },
497            MaterializedSessionMutation {
498                active_turn: Some(None),
499                last_turn_outcome: Some(completed("two", 4, 40)),
500                ..Default::default()
501            },
502            MaterializedSessionMutation {
503                active_turn: Some(Some(MaterializedTurn {
504                    command_id: "crashed".into(),
505                    accepted_ordinal: None,
506                    turn_start_position: 5,
507                    started_at_ms: 5,
508                    steered_into: None,
509                })),
510                ..Default::default()
511            },
512        ];
513        for _ in 0..2 {
514            apply_projection_page_to(&path, "child", |page| {
515                for (i, event) in events.iter().enumerate() {
516                    page.apply(
517                        i as u64 + 1,
518                        &format!("{:064x}", i),
519                        &format!("{:064x}", i + 1),
520                        event,
521                    )?;
522                }
523                Ok(())
524            })?;
525        }
526        let before = load_session_usage_from(&path, "child", 0, 1)?.unwrap();
527        assert_eq!(before.turn_selections["one"].model.as_deref(), Some("sol"));
528        assert_eq!(before.by_model.len(), 2);
529        assert_eq!(before.coverage.unfinished_turns, 1);
530        assert_eq!(before.by_model[0].selection.effort.as_deref(), Some("low"));
531        let later = load_session_usage_from(&path, "child", 2, 1)?.unwrap();
532        assert_eq!(later.turn_selections["two"].model.as_deref(), Some("luna"));
533        assert_eq!(before.totals["input_tokens"].tokens, 60);
534        delete_session_from(&path, "child")?;
535        delete_session_from(&path, "root")?;
536        assert_eq!(
537            before,
538            load_session_usage_from(&path, "child", 0, 1)?.unwrap()
539        );
540        let tree = load_usage_tree_from(&path, "root")?.unwrap();
541        assert_eq!(tree.sessions.len(), 2);
542        assert_eq!(tree.summary.totals, before.totals);
543        assert_eq!(tree.summary.by_model, before.by_model);
544        assert_eq!(tree.summary.coverage.unfinished_turns, 1);
545        assert!(tree.sessions.iter().all(|s| !s.operational_session_present));
546        assert_eq!(tree.sessions[0].task_name.as_deref(), Some("investigate"));
547        Ok(())
548    }
549
550    #[test]
551    fn tree_keeps_registered_descendants_when_operational_state_is_removed() -> Result<()> {
552        let dir = tempfile::tempdir()?;
553        let path = dir.path().join("tree.sqlite");
554        let mut state = State::default();
555        for id in ["root", "child", "grandchild"] {
556            state
557                .sessions
558                .insert(id.into(), super::super::tests::session(id, "project"));
559        }
560        {
561            let (child, parent) = ("child", "root");
562            let relation: SubagentRecord = serde_json::from_value(serde_json::json!({
563                "child_session_id":child,"parent_session_id":parent,"task_name":child,
564                "profile_id":"codex","working_directory":"","initial_prompt":"probe",
565                "request_key":child,"created_at":"2026-09-30T00:00:00Z"
566            }))?;
567            state.subagents.insert(child.into(), relation);
568        }
569        save_state_to(&path, &state)?;
570        // Operational children cannot delegate today; retained accounting
571        // still traverses deeper historical/imported relationships.
572        open(&path)?.execute(
573            "INSERT INTO subagent_accounting VALUES ('grandchild','child','nested')",
574            [],
575        )?;
576        let before = load_usage_tree_from(&path, "root")?.unwrap();
577        assert_eq!(before.sessions.len(), 3);
578        save_state_to(&path, &State::default())?;
579        let after = load_usage_tree_from(&path, "root")?.unwrap();
580        assert_eq!(after.sessions.len(), 3);
581        assert!(
582            after
583                .sessions
584                .iter()
585                .all(|s| !s.operational_session_present)
586        );
587        Ok(())
588    }
589
590    #[test]
591    fn provider_breakdowns_partition_counters_without_double_counting_or_guessing() -> Result<()> {
592        use mj_core::usage::ProviderTurnUsage;
593        let mut usage = completed("one", 1, 100).usage.unwrap();
594        let mut first = usage.clone();
595        first.input_tokens = 30;
596        first.output_tokens = 4;
597        first.total_tokens = 34;
598        let mut second = first.clone();
599        second.input_tokens = 60;
600        second.output_tokens = 6;
601        second.total_tokens = 66;
602        usage.provider_details = Some(Box::new(ProviderTurnUsage {
603            model_usage: BTreeMap::from([("a".into(), first), ("b".into(), second)]),
604            ..Default::default()
605        }));
606        let selected = UsageSelection {
607            model: Some("a".into()),
608            effort: Some("high".into()),
609        };
610        let mut groups = BTreeMap::new();
611        attribute(&usage, selected.clone(), &mut groups)?;
612        assert_eq!(groups[&selected]["input_tokens"].tokens, 30);
613        assert_eq!(
614            groups[&UsageSelection {
615                model: Some("b".into()),
616                effort: None
617            }]["input_tokens"]
618                .tokens,
619            60
620        );
621        assert_eq!(
622            groups[&UsageSelection::default()]["input_tokens"].tokens,
623            10
624        );
625        for (counter, total) in counters(&usage) {
626            assert_eq!(
627                groups
628                    .values()
629                    .filter_map(|v| v.get(&counter))
630                    .map(|v| v.tokens)
631                    .sum::<u64>(),
632                total
633            );
634        }
635        usage
636            .provider_details
637            .as_mut()
638            .unwrap()
639            .model_usage
640            .get_mut("b")
641            .unwrap()
642            .scope = UsageScope::LastRequest;
643        groups.clear();
644        attribute(&usage, selected, &mut groups)?;
645        assert_eq!(groups.len(), 1);
646        assert_eq!(
647            groups[&UsageSelection::default()]["input_tokens"].tokens,
648            100
649        );
650        Ok(())
651    }
652
653    #[test]
654    fn usage_survives_reopening_and_replay_with_honest_coverage() -> Result<()> {
655        let dir = tempfile::tempdir()?;
656        let path = dir.path().join("usage.sqlite");
657        let conn = open(&path)?;
658        // Use the same projection setup as the materialized database tests.
659        drop(conn);
660        save_session_to(&path, &super::super::tests::session("usage", "project"))?;
661
662        let mut turns = Vec::new();
663        for (i, scope) in [
664            Some(UsageScope::Turn),
665            Some(UsageScope::Turn),
666            Some(UsageScope::LastRequest),
667            None,
668            Some(UsageScope::Unspecified),
669        ]
670        .into_iter()
671        .enumerate()
672        {
673            let turn = MaterializedTurnOutcome {
674                diagnostic: Some(mj_core::diagnostic::TurnDiagnostic {
675                    message: "Usage limit exceeded".into(),
676                    code: Some("provider.auth_error".into()),
677                    http_status: Some(403),
678                    reset_at: None,
679                }),
680                command_id: format!("p{i}"),
681                accepted_ordinal: None,
682                turn_start_position: Some(i as u64 + 1),
683                completed_ordinal: i as u64 + 1,
684                completed_at_ms: 100,
685                outcome: TurnOutcomeKind::Completed {
686                    stop_reason: "EndTurn".into(),
687                },
688                usage: scope.map(|scope| TokenUsage {
689                    provider_details: Some(Box::new(mj_core::usage::ProviderTurnUsage {
690                        cost: Some(mj_core::usage::ProviderTurnCost::from_usd_ticks(
691                            88_767_200, false,
692                        )),
693                        model_calls: Some(2),
694                        api_duration_ms: Some(1700),
695                        elapsed_ms: Some(2377),
696                        model_usage: BTreeMap::new(),
697                    })),
698                    scope,
699                    total_tokens: 30,
700                    input_tokens: 20,
701                    output_tokens: 10,
702                    thought_tokens: None,
703                    cached_read_tokens: Some(0),
704                    cached_write_tokens: None,
705                }),
706            };
707            turns.push(turn);
708        }
709        for _ in 0..2 {
710            apply_projection_page_to(&path, "usage", |page| {
711                for turn in &turns {
712                    let prior = if turn.completed_ordinal == 1 {
713                        mj_core::relay::RELAY_EVENT_GENESIS_DIGEST.into()
714                    } else {
715                        format!("{:064x}", turn.completed_ordinal - 1)
716                    };
717                    page.apply(
718                        turn.completed_ordinal,
719                        &prior,
720                        &format!("{:064x}", turn.completed_ordinal),
721                        &MaterializedSessionMutation {
722                            last_turn_outcome: Some(turn.clone()),
723                            provider_cost: Some(ProviderCost {
724                                amount: 1.25,
725                                currency: "USD".into(),
726                                observed_at_ms: 100,
727                            }),
728                            ..Default::default()
729                        },
730                    )?;
731                }
732                Ok(())
733            })?;
734        }
735        let first = load_session_usage_from(&path, "usage", 0, 2)?.unwrap();
736        assert_eq!(first.turns.len(), 2);
737        let details = first.turns[0]
738            .usage
739            .as_ref()
740            .unwrap()
741            .provider_details
742            .as_ref()
743            .unwrap();
744        assert_eq!(details.cost.as_ref().unwrap().usd, "0.0088767200");
745        assert_eq!(details.elapsed_ms, Some(2377));
746        assert_eq!(
747            first.turns[0].diagnostic.as_ref().unwrap().http_status,
748            Some(403)
749        );
750        assert_eq!(first.provider_session_cost.as_ref().unwrap().amount, 1.25);
751        assert_eq!(first.coverage.recorded_turns, 5);
752        assert_eq!(first.coverage.full_turn_reports, 2);
753        assert_eq!(first.coverage.last_request_reports, 1);
754        assert_eq!(first.coverage.unspecified_reports, 1);
755        assert_eq!(first.coverage.missing_reports, 1);
756        // Only the two whole-turn reports are summed: the last-request and
757        // unspecified turns stay out of the totals.
758        assert_eq!(first.totals["total_tokens"].tokens, 60);
759        assert_eq!(first.totals["total_tokens"].reported_turns, 2);
760        assert_eq!(first.totals["cached_read_tokens"].tokens, 0);
761        assert!(!first.totals.contains_key("thought_tokens"));
762        let next = load_session_usage_from(&path, "usage", first.next_after_seq, 2)?.unwrap();
763        assert_eq!(next.turns.len(), 2);
764        assert_eq!(next.latest_seq, 5);
765        assert_eq!(next.totals, first.totals);
766        let last = load_session_usage_from(&path, "usage", next.next_after_seq, 2)?.unwrap();
767        assert_eq!(last.turns.len(), 1);
768        assert_eq!(
769            last.turns[0].usage.as_ref().unwrap().scope,
770            UsageScope::Unspecified
771        );
772        assert_eq!(last.next_after_seq, last.latest_seq);
773        assert_eq!(last.totals, first.totals);
774        Ok(())
775    }
776}
mj_controller/database/usage.rs

mj_controller/database/
usage.rs
