ag_store/
usage.rs

1//! Session-usage persistence adapters and query helpers.
2
3use std::sync::Arc;
4
5use ag_agent::SessionStats;
6use async_trait::async_trait;
7use sqlx::SqlitePool;
8
9use crate::DbError;
10use crate::timestamp::TimestampSource;
11
12/// Row returned when loading per-model token usage from the `session_usage`
13/// table.
14pub struct SessionUsageRow {
15    /// Row creation timestamp in Unix seconds.
16    pub created_at: i64,
17    /// Accumulated input-token count.
18    pub input_tokens: i64,
19    /// Number of agent invocations included in the totals.
20    pub invocation_count: i64,
21    /// Provider model identifier.
22    pub model: String,
23    /// Accumulated output-token count.
24    pub output_tokens: i64,
25    /// Owning session identifier, when present.
26    pub session_id: Option<String>,
27}
28
29/// Session-usage persistence boundary used by app orchestration and tests.
30#[cfg_attr(test, mockall::automock)]
31#[async_trait]
32pub trait UsageRepository: Send + Sync {
33    /// Loads per-model token usage rows for a session, ordered by model name.
34    async fn load_session_usage(&self, session_id: &str) -> Result<Vec<SessionUsageRow>, DbError>;
35
36    /// Accumulates per-model token usage for a session.
37    async fn upsert_session_usage(
38        &self,
39        session_id: &str,
40        model: &str,
41        stats: &SessionStats,
42    ) -> Result<(), DbError>;
43}
44
45/// `SQLite` implementation of [`UsageRepository`].
46#[derive(Clone)]
47pub(crate) struct SqliteUsageRepository {
48    pool: SqlitePool,
49    timestamp_source: Arc<dyn TimestampSource>,
50}
51
52impl SqliteUsageRepository {
53    /// Creates a usage repository backed by the provided pool and timestamp
54    /// source.
55    pub(crate) fn new(pool: SqlitePool, timestamp_source: Arc<dyn TimestampSource>) -> Self {
56        Self {
57            pool,
58            timestamp_source,
59        }
60    }
61
62    fn now(&self) -> i64 {
63        self.timestamp_source.now_timestamp_seconds()
64    }
65}
66
67#[async_trait]
68impl UsageRepository for SqliteUsageRepository {
69    async fn load_session_usage(&self, session_id: &str) -> Result<Vec<SessionUsageRow>, DbError> {
70        let rows = sqlx::query_as!(
71            SessionUsageRow,
72            r#"
73SELECT session_id, model, created_at, input_tokens, invocation_count, output_tokens
74FROM session_usage
75WHERE session_id = ?
76ORDER BY model
77            "#,
78            session_id
79        )
80        .fetch_all(&self.pool)
81        .await?;
82
83        Ok(rows)
84    }
85
86    async fn upsert_session_usage(
87        &self,
88        session_id: &str,
89        model: &str,
90        stats: &SessionStats,
91    ) -> Result<(), DbError> {
92        if stats.input_tokens == 0 && stats.output_tokens == 0 {
93            return Ok(());
94        }
95
96        let now = self.now();
97
98        sqlx::query!(
99            r"
100INSERT INTO session_usage (
101    session_id, model, created_at, input_tokens, output_tokens, invocation_count
102)
103VALUES (?, ?, ?, ?, ?, 1)
104ON CONFLICT(session_id, model) DO UPDATE SET
105    input_tokens = input_tokens + excluded.input_tokens,
106    output_tokens = output_tokens + excluded.output_tokens,
107    invocation_count = invocation_count + 1
108",
109            session_id,
110            model,
111            now,
112            stats.input_tokens.cast_signed(),
113            stats.output_tokens.cast_signed()
114        )
115        .execute(&self.pool)
116        .await?;
117
118        Ok(())
119    }
120}
ag_store/usage.rs

ag_store/
usage.rs
