kcode_intelligence_router/
receipts.rs1use std::{
2 collections::BTreeMap,
3 fs::{self, OpenOptions},
4 io::Write,
5 path::{Path, PathBuf},
6 sync::{Arc, Mutex},
7};
8
9use chrono::{DateTime, NaiveDate, Utc};
10use serde::{Deserialize, Serialize};
11use uuid::Uuid;
12
13use crate::{CostEstimate, Error, Result, estimate_cost};
14
15const RECEIPT_VERSION: u32 = 3;
16
17#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
19#[serde(rename_all = "camelCase")]
20pub struct TokenUsage {
21 pub input_tokens: u64,
22 pub cached_input_tokens: u64,
23 pub thinking_tokens: u64,
24 pub output_tokens: u64,
25}
26
27impl TokenUsage {
28 fn add_assign(&mut self, other: Self) {
29 self.input_tokens = self.input_tokens.saturating_add(other.input_tokens);
30 self.cached_input_tokens = self
31 .cached_input_tokens
32 .saturating_add(other.cached_input_tokens);
33 self.thinking_tokens = self.thinking_tokens.saturating_add(other.thinking_tokens);
34 self.output_tokens = self.output_tokens.saturating_add(other.output_tokens);
35 }
36}
37
38#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
40#[serde(tag = "kind", rename_all = "snake_case")]
41pub enum Metering {
42 Tokens(TokenUsage),
43 DurationSeconds { seconds: f64 },
44 Unavailable,
45}
46
47#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
49#[serde(rename_all = "camelCase")]
50pub struct UsageReceipt {
51 pub version: u32,
52 pub id: Uuid,
53 pub recorded_at: DateTime<Utc>,
54 pub user_id: String,
55 #[serde(default, skip_serializing_if = "Option::is_none")]
57 pub operation_id: Option<Uuid>,
58 #[serde(default, skip_serializing_if = "Option::is_none")]
60 pub parent_operation_id: Option<Uuid>,
61 pub operation: String,
62 pub requested_model: String,
63 pub actual_model: String,
64 #[serde(skip_serializing_if = "Option::is_none")]
65 pub provider_request_id: Option<String>,
66 #[serde(skip_serializing_if = "Option::is_none")]
67 pub provider_thread_id: Option<String>,
68 pub metering: Metering,
69 #[serde(default, skip_serializing_if = "Option::is_none")]
72 pub cost: Option<CostEstimate>,
73}
74
75impl UsageReceipt {
76 pub(crate) fn new(
77 user_id: impl Into<String>,
78 operation_id: Uuid,
79 parent_operation_id: Option<Uuid>,
80 operation: impl Into<String>,
81 requested_model: impl Into<String>,
82 actual_model: impl Into<String>,
83 metering: Metering,
84 ) -> Self {
85 let actual_model = actual_model.into();
86 let cost = estimate_cost(&actual_model, &metering);
87 Self {
88 version: RECEIPT_VERSION,
89 id: Uuid::new_v4(),
90 recorded_at: Utc::now(),
91 user_id: user_id.into(),
92 operation_id: Some(operation_id),
93 parent_operation_id,
94 operation: operation.into(),
95 requested_model: requested_model.into(),
96 actual_model,
97 provider_request_id: None,
98 provider_thread_id: None,
99 metering,
100 cost,
101 }
102 }
103}
104
105#[derive(Clone, Debug, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)]
107#[serde(rename_all = "camelCase")]
108pub struct DailyUsageKey {
109 pub user_id: String,
110 pub model: String,
111}
112
113#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
115#[serde(rename_all = "camelCase")]
116pub struct DailyUsage {
117 pub calls: u64,
118 pub token_calls: u64,
119 pub duration_calls: u64,
120 pub unmetered_calls: u64,
121 pub tokens: TokenUsage,
122 pub duration_seconds: f64,
123 pub estimated_cost_usd_nanos: u64,
124 pub unpriced_calls: u64,
125}
126
127#[derive(Clone, Debug)]
129pub(crate) struct ReceiptStore {
130 root: Arc<PathBuf>,
131 writer: Arc<Mutex<()>>,
132}
133
134impl ReceiptStore {
135 pub(crate) fn open(root: impl Into<PathBuf>) -> Result<Self> {
136 let root = root.into();
137 fs::create_dir_all(&root).map_err(|error| receipt_io("create", &root, error))?;
138 #[cfg(unix)]
139 {
140 use std::os::unix::fs::PermissionsExt;
141 fs::set_permissions(&root, fs::Permissions::from_mode(0o700))
142 .map_err(|error| receipt_io("protect", &root, error))?;
143 }
144 Ok(Self {
145 root: Arc::new(root),
146 writer: Arc::new(Mutex::new(())),
147 })
148 }
149
150 pub(crate) fn record(&self, receipt: &UsageReceipt) -> Result<()> {
152 let _guard = self.writer.lock().map_err(|_| {
153 Error::internal(
154 "receipt_store_unavailable",
155 "The usage receipt writer is unavailable.",
156 )
157 })?;
158 let bytes = serde_json::to_vec_pretty(receipt).map_err(|error| {
159 Error::internal(
160 "receipt_serialization_failed",
161 format!("Could not serialize usage receipt: {error}"),
162 )
163 })?;
164 let final_path = self.root.join(format!("{}.json", receipt.id));
165 let temporary_path = self.root.join(format!(".{}.tmp", receipt.id));
166 let result = (|| -> std::io::Result<()> {
167 let mut options = OpenOptions::new();
168 options.write(true).create_new(true);
169 #[cfg(unix)]
170 {
171 use std::os::unix::fs::OpenOptionsExt;
172 options.mode(0o600);
173 }
174 let mut file = options.open(&temporary_path)?;
175 file.write_all(&bytes)?;
176 file.sync_all()?;
177 fs::rename(&temporary_path, &final_path)?;
178 sync_directory(&self.root)?;
179 Ok(())
180 })();
181 if result.is_err() {
182 let _ = fs::remove_file(&temporary_path);
183 }
184 result.map_err(|error| receipt_io("write", &final_path, error))
185 }
186
187 pub(crate) fn receipts(&self) -> Result<Vec<UsageReceipt>> {
189 let mut receipts = Vec::new();
190 let entries =
191 fs::read_dir(&*self.root).map_err(|error| receipt_io("list", &self.root, error))?;
192 for entry in entries {
193 let entry = entry.map_err(|error| receipt_io("list", &self.root, error))?;
194 let path = entry.path();
195 if path.extension().and_then(|value| value.to_str()) != Some("json") {
196 continue;
197 }
198 let bytes = fs::read(&path).map_err(|error| receipt_io("read", &path, error))?;
199 let receipt = serde_json::from_slice::<UsageReceipt>(&bytes).map_err(|error| {
200 Error::internal(
201 "invalid_usage_receipt",
202 format!("Invalid usage receipt {}: {error}", path.display()),
203 )
204 })?;
205 receipts.push(receipt);
206 }
207 receipts.sort_by_key(|receipt| (receipt.recorded_at, receipt.id));
208 Ok(receipts)
209 }
210
211 pub(crate) fn daily_usage(
213 &self,
214 day: NaiveDate,
215 ) -> Result<BTreeMap<DailyUsageKey, DailyUsage>> {
216 let mut groups = BTreeMap::new();
217 for receipt in self.receipts()? {
218 if receipt.recorded_at.date_naive() != day {
219 continue;
220 }
221 let totals = groups
222 .entry(DailyUsageKey {
223 user_id: receipt.user_id,
224 model: receipt.actual_model,
225 })
226 .or_insert_with(DailyUsage::default);
227 totals.calls = totals.calls.saturating_add(1);
228 if let Some(cost) = receipt.cost {
229 totals.estimated_cost_usd_nanos = totals
230 .estimated_cost_usd_nanos
231 .saturating_add(cost.usd_nanos);
232 } else {
233 totals.unpriced_calls = totals.unpriced_calls.saturating_add(1);
234 }
235 match receipt.metering {
236 Metering::Tokens(tokens) => {
237 totals.token_calls = totals.token_calls.saturating_add(1);
238 totals.tokens.add_assign(tokens);
239 }
240 Metering::DurationSeconds { seconds } => {
241 totals.duration_calls = totals.duration_calls.saturating_add(1);
242 totals.duration_seconds += seconds;
243 }
244 Metering::Unavailable => {
245 totals.unmetered_calls = totals.unmetered_calls.saturating_add(1);
246 }
247 }
248 }
249 Ok(groups)
250 }
251}
252
253fn receipt_io(operation: &str, path: &Path, error: std::io::Error) -> Error {
254 Error::internal(
255 "receipt_io_failed",
256 format!(
257 "Could not {operation} usage receipts at {}: {error}",
258 path.display()
259 ),
260 )
261}
262
263#[cfg(unix)]
264fn sync_directory(path: &Path) -> std::io::Result<()> {
265 std::fs::File::open(path)?.sync_all()
266}
267
268#[cfg(not(unix))]
269fn sync_directory(_path: &Path) -> std::io::Result<()> {
270 Ok(())
271}
272
273#[cfg(test)]
274mod tests {
275 use super::*;
276
277 fn temporary_store() -> (PathBuf, ReceiptStore) {
278 let root = std::env::temp_dir().join(format!("kennedy-usage-{}", Uuid::new_v4()));
279 let store = ReceiptStore::open(&root).unwrap();
280 (root, store)
281 }
282
283 #[test]
284 fn one_file_is_written_for_each_call_and_daily_totals_scan_them() {
285 let (root, store) = temporary_store();
286 let mut first = UsageReceipt::new(
287 "user-a",
288 Uuid::new_v4(),
289 None,
290 "search",
291 "gemini-2.5-flash",
292 "gemini-2.5-flash",
293 Metering::Tokens(TokenUsage {
294 input_tokens: 10,
295 cached_input_tokens: 4,
296 thinking_tokens: 3,
297 output_tokens: 2,
298 }),
299 );
300 first.recorded_at = DateTime::parse_from_rfc3339("2026-07-26T01:00:00Z")
301 .unwrap()
302 .with_timezone(&Utc);
303 assert!(first.operation_id.is_some());
304 let mut second = first.clone();
305 second.id = Uuid::new_v4();
306 second.user_id = "user-b".into();
307 store.record(&first).unwrap();
308 store.record(&second).unwrap();
309
310 assert_eq!(
311 fs::read_dir(&root).unwrap().flatten().count(),
312 2,
313 "there should be exactly one durable file per call"
314 );
315 let totals = store
316 .daily_usage(NaiveDate::from_ymd_opt(2026, 7, 26).unwrap())
317 .unwrap();
318 assert_eq!(totals.len(), 2);
319 assert_eq!(
320 totals[&DailyUsageKey {
321 user_id: "user-a".into(),
322 model: "gemini-2.5-flash".into(),
323 }]
324 .tokens
325 .cached_input_tokens,
326 4
327 );
328 fs::remove_dir_all(root).unwrap();
329 }
330}