zeph_memory/store/
usage_records.rs

1// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4//! Durable per-LLM-call usage ledger (`usage_records` table, issue #6549).
5
6use zeph_db::ActiveDialect;
7use zeph_db::sql;
8
9use super::SqliteStore;
10use crate::error::MemoryError;
11use crate::types::{ConversationId, MessageId, UsageRecord, UsageSource};
12
13/// Row shape shared by every `usage_records` SELECT in this module.
14type UsageRow = (
15    Option<MessageId>,
16    Option<ConversationId>,
17    String,
18    String,
19    String,
20    i64,
21    i64,
22    i64,
23    i64,
24    Option<i64>,
25    f64,
26    i64,
27    Option<i64>,
28    Option<f64>,
29);
30
31fn to_i64(v: u64) -> i64 {
32    i64::try_from(v).unwrap_or(i64::MAX)
33}
34
35fn to_u64(v: i64) -> u64 {
36    u64::try_from(v).unwrap_or(0)
37}
38
39fn row_to_usage_record(row: UsageRow) -> UsageRecord {
40    let (
41        message_id,
42        conversation_id,
43        source_str,
44        provider_name,
45        model_name,
46        input_tokens,
47        output_tokens,
48        cache_read_tokens,
49        cache_write_tokens,
50        reasoning_tokens,
51        cost_cents,
52        latency_ms,
53        ttft_ms,
54        tokens_per_sec,
55    ) = row;
56    let source = source_str.parse().unwrap_or_else(|_| {
57        tracing::warn!(value = %source_str, "unrecognized usage_records.source, defaulting to conversation");
58        UsageSource::Conversation
59    });
60    UsageRecord {
61        message_id,
62        conversation_id,
63        source,
64        provider_name,
65        model_name,
66        input_tokens: to_u64(input_tokens),
67        output_tokens: to_u64(output_tokens),
68        cache_read_tokens: to_u64(cache_read_tokens),
69        cache_write_tokens: to_u64(cache_write_tokens),
70        reasoning_tokens: reasoning_tokens.map(to_u64),
71        cost_cents,
72        latency_ms: to_u64(latency_ms),
73        ttft_ms: ttft_ms.map(to_u64),
74        tokens_per_sec,
75    }
76}
77
78impl SqliteStore {
79    /// Insert one durable usage row.
80    ///
81    /// Callers write exactly one row per production `CostTracker::record_usage`-feeding
82    /// call site — see spec `082-per-message-usage-cost-tracking` §3 for the completeness
83    /// invariant. `record.cost_cents` must come from `CostTracker::price_of` so the row's
84    /// cost matches the live daily aggregate for the same call.
85    ///
86    /// # Errors
87    ///
88    /// Returns an error if the insert fails.
89    pub async fn record_usage_row(&self, record: &UsageRecord) -> Result<(), MemoryError> {
90        zeph_db::query(sql!(
91            "INSERT INTO usage_records \
92             (message_id, conversation_id, source, provider_name, model_name, \
93              input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, \
94              reasoning_tokens, cost_cents, latency_ms, ttft_ms, tokens_per_sec) \
95             VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"
96        ))
97        .bind(record.message_id)
98        .bind(record.conversation_id)
99        .bind(record.source.as_str())
100        .bind(&record.provider_name)
101        .bind(&record.model_name)
102        .bind(to_i64(record.input_tokens))
103        .bind(to_i64(record.output_tokens))
104        .bind(to_i64(record.cache_read_tokens))
105        .bind(to_i64(record.cache_write_tokens))
106        .bind(record.reasoning_tokens.map(to_i64))
107        .bind(record.cost_cents)
108        .bind(to_i64(record.latency_ms))
109        .bind(record.ttft_ms.map(to_i64))
110        .bind(record.tokens_per_sec)
111        .execute(&self.pool)
112        .await?;
113        Ok(())
114    }
115
116    /// Fetch the usage row for a single conversational message, if one was recorded.
117    ///
118    /// # Errors
119    ///
120    /// Returns an error if the query fails.
121    pub async fn message_usage(
122        &self,
123        message_id: MessageId,
124    ) -> Result<Option<UsageRecord>, MemoryError> {
125        let row: Option<UsageRow> = zeph_db::query_as(sql!(
126            "SELECT message_id, conversation_id, source, provider_name, model_name, \
127                    input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, \
128                    reasoning_tokens, cost_cents, latency_ms, ttft_ms, tokens_per_sec \
129             FROM usage_records WHERE message_id = ?"
130        ))
131        .bind(message_id)
132        .fetch_optional(&self.pool)
133        .await?;
134        Ok(row.map(row_to_usage_record))
135    }
136
137    /// Fetch every conversational usage row for a conversation, ordered by message id.
138    ///
139    /// Background/orchestration rows (planner, aggregator, ensemble member) carry no
140    /// `message_id` and are excluded — use a direct `usage_records` query if those are
141    /// needed alongside conversational rows.
142    ///
143    /// # Errors
144    ///
145    /// Returns an error if the query fails.
146    pub async fn conversation_usage(
147        &self,
148        conversation_id: ConversationId,
149    ) -> Result<Vec<UsageRecord>, MemoryError> {
150        let rows: Vec<UsageRow> = zeph_db::query_as(sql!(
151            "SELECT message_id, conversation_id, source, provider_name, model_name, \
152                    input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, \
153                    reasoning_tokens, cost_cents, latency_ms, ttft_ms, tokens_per_sec \
154             FROM usage_records \
155             WHERE conversation_id = ? AND message_id IS NOT NULL \
156             ORDER BY message_id ASC"
157        ))
158        .bind(conversation_id)
159        .fetch_all(&self.pool)
160        .await?;
161        Ok(rows.into_iter().map(row_to_usage_record).collect())
162    }
163
164    /// Sum `cost_cents` across every usage row created at or after `since_epoch_secs`
165    /// (Unix epoch seconds, UTC).
166    ///
167    /// Used for the current-day reconciliation invariant: `usage_cost_since(utc_midnight)`
168    /// must equal `CostTracker::current_spend()` (spec `082` US-001 AC).
169    ///
170    /// # Errors
171    ///
172    /// Returns an error if the query fails.
173    pub async fn usage_cost_since(&self, since_epoch_secs: i64) -> Result<f64, MemoryError> {
174        let epoch_expr = <ActiveDialect as zeph_db::dialect::Dialect>::epoch_from_col("created_at");
175        let raw = format!(
176            "SELECT COALESCE(SUM(cost_cents), 0.0) FROM usage_records WHERE {epoch_expr} >= ?"
177        );
178        let sql = zeph_db::rewrite_placeholders(&raw);
179        let total: f64 = zeph_db::query_scalar(sqlx::AssertSqlSafe(sql))
180            .bind(since_epoch_secs)
181            .fetch_one(&self.pool)
182            .await?;
183        Ok(total)
184    }
185}
186
187#[cfg(test)]
188mod tests {
189    use super::*;
190
191    fn sample(source: UsageSource) -> UsageRecord {
192        UsageRecord {
193            message_id: None,
194            conversation_id: None,
195            source,
196            provider_name: "quality".to_string(),
197            model_name: "claude-sonnet-5".to_string(),
198            input_tokens: 100,
199            output_tokens: 50,
200            cache_read_tokens: 10,
201            cache_write_tokens: 5,
202            reasoning_tokens: Some(20),
203            cost_cents: 0.42,
204            latency_ms: 800,
205            ttft_ms: Some(120),
206            tokens_per_sec: Some(63.5),
207        }
208    }
209
210    #[tokio::test]
211    async fn record_and_fetch_conversational_row() {
212        let store = SqliteStore::new(":memory:").await.expect("store");
213        let cid = store.create_conversation().await.expect("conversation");
214        let mid = store
215            .save_message(cid, "assistant", "hello")
216            .await
217            .expect("message");
218
219        let mut record = sample(UsageSource::Conversation);
220        record.message_id = Some(mid);
221        record.conversation_id = Some(cid);
222        store.record_usage_row(&record).await.expect("insert");
223
224        let fetched = store
225            .message_usage(mid)
226            .await
227            .expect("query")
228            .expect("row exists");
229        assert_eq!(fetched.source, UsageSource::Conversation);
230        assert_eq!(fetched.input_tokens, 100);
231        assert_eq!(fetched.output_tokens, 50);
232        assert_eq!(fetched.cache_read_tokens, 10);
233        assert_eq!(fetched.cache_write_tokens, 5);
234        assert_eq!(fetched.reasoning_tokens, Some(20));
235        assert!((fetched.cost_cents - 0.42).abs() < 1e-9);
236        assert_eq!(fetched.ttft_ms, Some(120));
237
238        let conv_rows = store.conversation_usage(cid).await.expect("conv query");
239        assert_eq!(conv_rows.len(), 1);
240        assert_eq!(conv_rows[0].message_id, Some(mid));
241    }
242
243    #[tokio::test]
244    async fn background_rows_have_no_message_id_and_are_excluded_from_conversation_usage() {
245        let store = SqliteStore::new(":memory:").await.expect("store");
246        let cid = store.create_conversation().await.expect("conversation");
247
248        let mut planner_row = sample(UsageSource::Planner);
249        planner_row.conversation_id = Some(cid);
250        store.record_usage_row(&planner_row).await.expect("insert");
251
252        let conv_rows = store.conversation_usage(cid).await.expect("conv query");
253        assert!(conv_rows.is_empty(), "background rows carry no message_id");
254    }
255
256    #[tokio::test]
257    async fn message_usage_none_when_unrecorded() {
258        let store = SqliteStore::new(":memory:").await.expect("store");
259        let cid = store.create_conversation().await.expect("conversation");
260        let mid = store
261            .save_message(cid, "assistant", "no usage row")
262            .await
263            .expect("message");
264        assert!(store.message_usage(mid).await.expect("query").is_none());
265    }
266
267    #[tokio::test]
268    async fn usage_cost_since_sums_rows_in_window() {
269        let store = SqliteStore::new(":memory:").await.expect("store");
270        store
271            .record_usage_row(&sample(UsageSource::Aggregator))
272            .await
273            .expect("insert 1");
274        store
275            .record_usage_row(&sample(UsageSource::EnsembleMember))
276            .await
277            .expect("insert 2");
278
279        let total = store.usage_cost_since(0).await.expect("sum");
280        assert!((total - 0.84).abs() < 1e-6, "total={total}");
281
282        let future = std::time::SystemTime::now()
283            .duration_since(std::time::UNIX_EPOCH)
284            .unwrap()
285            .as_secs()
286            .cast_signed()
287            + 3600;
288        let none_yet = store.usage_cost_since(future).await.expect("sum future");
289        assert!((none_yet - 0.0).abs() < 1e-9);
290    }
291}
zeph_memory/store/usage_records.rs

zeph_memory/store/
usage_records.rs
