kernex_memory/store/
usage.rs

1//! Token usage recording and cost tracking.
2
3use super::Store;
4use kernex_core::error::KernexError;
5use kernex_core::pricing::pricing_for;
6
7/// Aggregated token usage for a session or sender.
8#[derive(Debug, Clone, Default)]
9pub struct UsageSummary {
10    /// Total tokens consumed across all recorded requests.
11    pub total_tokens: i64,
12    /// Estimated total cost in USD.
13    pub total_cost_usd: f64,
14    /// Number of API requests recorded.
15    pub request_count: i64,
16}
17
18impl Store {
19    /// Record token usage for a completed API request.
20    ///
21    /// Cost is estimated using known per-model pricing. If the model is
22    /// unrecognized, cost is recorded as 0.0.
23    pub async fn record_usage(
24        &self,
25        sender_id: &str,
26        session_id: &str,
27        tokens: u64,
28        model: &str,
29    ) -> Result<(), KernexError> {
30        let cost = pricing_for(model)
31            .map(|p| p.estimate_cost(tokens))
32            .unwrap_or(0.0);
33
34        sqlx::query(
35            "INSERT INTO token_usage (sender_id, session_id, model, tokens, cost_usd)
36             VALUES (?, ?, ?, ?, ?)",
37        )
38        .bind(sender_id)
39        .bind(session_id)
40        .bind(model)
41        .bind(tokens as i64)
42        .bind(cost)
43        .execute(&self.pool)
44        .await
45        .map_err(|e| KernexError::Store(format!("failed to record token usage: {e}")))?;
46
47        Ok(())
48    }
49
50    /// Get aggregated token usage for a session.
51    pub async fn get_session_usage(&self, session_id: &str) -> Result<UsageSummary, KernexError> {
52        let row: Option<(i64, f64, i64)> = sqlx::query_as(
53            "SELECT COALESCE(SUM(tokens), 0), COALESCE(SUM(cost_usd), 0.0), COUNT(*)
54             FROM token_usage WHERE session_id = ?",
55        )
56        .bind(session_id)
57        .fetch_optional(&self.pool)
58        .await
59        .map_err(|e| KernexError::Store(format!("failed to query session usage: {e}")))?;
60
61        let (total_tokens, total_cost_usd, request_count) = row.unwrap_or((0, 0.0, 0));
62
63        Ok(UsageSummary {
64            total_tokens,
65            total_cost_usd,
66            request_count,
67        })
68    }
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74    use kernex_core::config::MemoryConfig;
75
76    async fn make_store() -> Store {
77        let config = MemoryConfig {
78            db_path: ":memory:".to_string(),
79            ..Default::default()
80        };
81        Store::new(&config).await.unwrap()
82    }
83
84    #[tokio::test]
85    async fn test_record_and_get_usage() {
86        let store = make_store().await;
87        store
88            .record_usage("user-1", "sess-abc", 1000, "claude-sonnet-4-6")
89            .await
90            .unwrap();
91        store
92            .record_usage("user-1", "sess-abc", 500, "claude-sonnet-4-6")
93            .await
94            .unwrap();
95
96        let summary = store.get_session_usage("sess-abc").await.unwrap();
97        assert_eq!(summary.total_tokens, 1500);
98        assert_eq!(summary.request_count, 2);
99        assert!(summary.total_cost_usd > 0.0);
100    }
101
102    #[tokio::test]
103    async fn test_get_usage_empty_session() {
104        let store = make_store().await;
105        let summary = store.get_session_usage("sess-nonexistent").await.unwrap();
106        assert_eq!(summary.total_tokens, 0);
107        assert_eq!(summary.request_count, 0);
108        assert_eq!(summary.total_cost_usd, 0.0);
109    }
110
111    #[tokio::test]
112    async fn test_record_usage_unknown_model_zero_cost() {
113        let store = make_store().await;
114        store
115            .record_usage("user-1", "sess-local", 2000, "llama3.2")
116            .await
117            .unwrap();
118
119        let summary = store.get_session_usage("sess-local").await.unwrap();
120        assert_eq!(summary.total_tokens, 2000);
121        assert_eq!(summary.total_cost_usd, 0.0);
122    }
123
124    #[tokio::test]
125    async fn test_usage_isolated_by_session() {
126        let store = make_store().await;
127        store
128            .record_usage("user-1", "sess-1", 100, "gpt-4o")
129            .await
130            .unwrap();
131        store
132            .record_usage("user-1", "sess-2", 200, "gpt-4o")
133            .await
134            .unwrap();
135
136        let s1 = store.get_session_usage("sess-1").await.unwrap();
137        let s2 = store.get_session_usage("sess-2").await.unwrap();
138        assert_eq!(s1.total_tokens, 100);
139        assert_eq!(s2.total_tokens, 200);
140    }
141}
kernex_memory/store/usage.rs

kernex_memory/store/
usage.rs
