kcode_intelligence_router/
receipts.rs

1use 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::{Error, Result};
14
15const RECEIPT_VERSION: u32 = 1;
16
17/// The four token categories Kennedy accounts independently.
18#[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/// Provider metering returned by one call.
39#[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/// Durable evidence for one completed provider call.
48#[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    pub operation: String,
56    pub requested_model: String,
57    pub actual_model: String,
58    #[serde(skip_serializing_if = "Option::is_none")]
59    pub provider_request_id: Option<String>,
60    #[serde(skip_serializing_if = "Option::is_none")]
61    pub provider_thread_id: Option<String>,
62    pub metering: Metering,
63}
64
65impl UsageReceipt {
66    pub(crate) fn new(
67        user_id: impl Into<String>,
68        operation: impl Into<String>,
69        requested_model: impl Into<String>,
70        actual_model: impl Into<String>,
71        metering: Metering,
72    ) -> Self {
73        Self {
74            version: RECEIPT_VERSION,
75            id: Uuid::new_v4(),
76            recorded_at: Utc::now(),
77            user_id: user_id.into(),
78            operation: operation.into(),
79            requested_model: requested_model.into(),
80            actual_model: actual_model.into(),
81            provider_request_id: None,
82            provider_thread_id: None,
83            metering,
84        }
85    }
86}
87
88/// One grouping key in a linearly calculated daily report.
89#[derive(Clone, Debug, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)]
90#[serde(rename_all = "camelCase")]
91pub struct DailyUsageKey {
92    pub user_id: String,
93    pub model: String,
94}
95
96/// Totals calculated by opening every receipt for the requested UTC day.
97#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
98#[serde(rename_all = "camelCase")]
99pub struct DailyUsage {
100    pub calls: u64,
101    pub token_calls: u64,
102    pub duration_calls: u64,
103    pub unmetered_calls: u64,
104    pub tokens: TokenUsage,
105    pub duration_seconds: f64,
106}
107
108/// Intentionally simple one-file-per-call receipt storage.
109#[derive(Clone, Debug)]
110pub(crate) struct ReceiptStore {
111    root: Arc<PathBuf>,
112    writer: Arc<Mutex<()>>,
113}
114
115impl ReceiptStore {
116    pub(crate) fn open(root: impl Into<PathBuf>) -> Result<Self> {
117        let root = root.into();
118        fs::create_dir_all(&root).map_err(|error| receipt_io("create", &root, error))?;
119        #[cfg(unix)]
120        {
121            use std::os::unix::fs::PermissionsExt;
122            fs::set_permissions(&root, fs::Permissions::from_mode(0o700))
123                .map_err(|error| receipt_io("protect", &root, error))?;
124        }
125        Ok(Self {
126            root: Arc::new(root),
127            writer: Arc::new(Mutex::new(())),
128        })
129    }
130
131    /// Atomically writes exactly one new JSON receipt.
132    pub(crate) fn record(&self, receipt: &UsageReceipt) -> Result<()> {
133        let _guard = self.writer.lock().map_err(|_| {
134            Error::internal(
135                "receipt_store_unavailable",
136                "The usage receipt writer is unavailable.",
137            )
138        })?;
139        let bytes = serde_json::to_vec_pretty(receipt).map_err(|error| {
140            Error::internal(
141                "receipt_serialization_failed",
142                format!("Could not serialize usage receipt: {error}"),
143            )
144        })?;
145        let final_path = self.root.join(format!("{}.json", receipt.id));
146        let temporary_path = self.root.join(format!(".{}.tmp", receipt.id));
147        let result = (|| -> std::io::Result<()> {
148            let mut options = OpenOptions::new();
149            options.write(true).create_new(true);
150            #[cfg(unix)]
151            {
152                use std::os::unix::fs::OpenOptionsExt;
153                options.mode(0o600);
154            }
155            let mut file = options.open(&temporary_path)?;
156            file.write_all(&bytes)?;
157            file.sync_all()?;
158            fs::rename(&temporary_path, &final_path)?;
159            sync_directory(&self.root)?;
160            Ok(())
161        })();
162        if result.is_err() {
163            let _ = fs::remove_file(&temporary_path);
164        }
165        result.map_err(|error| receipt_io("write", &final_path, error))
166    }
167
168    /// Opens and decodes every receipt file. There are deliberately no indexes yet.
169    pub(crate) fn receipts(&self) -> Result<Vec<UsageReceipt>> {
170        let mut receipts = Vec::new();
171        let entries =
172            fs::read_dir(&*self.root).map_err(|error| receipt_io("list", &self.root, error))?;
173        for entry in entries {
174            let entry = entry.map_err(|error| receipt_io("list", &self.root, error))?;
175            let path = entry.path();
176            if path.extension().and_then(|value| value.to_str()) != Some("json") {
177                continue;
178            }
179            let bytes = fs::read(&path).map_err(|error| receipt_io("read", &path, error))?;
180            let receipt = serde_json::from_slice::<UsageReceipt>(&bytes).map_err(|error| {
181                Error::internal(
182                    "invalid_usage_receipt",
183                    format!("Invalid usage receipt {}: {error}", path.display()),
184                )
185            })?;
186            receipts.push(receipt);
187        }
188        receipts.sort_by_key(|receipt| (receipt.recorded_at, receipt.id));
189        Ok(receipts)
190    }
191
192    /// Calculates per-user, per-model totals for one UTC day by scanning all receipts.
193    pub(crate) fn daily_usage(
194        &self,
195        day: NaiveDate,
196    ) -> Result<BTreeMap<DailyUsageKey, DailyUsage>> {
197        let mut groups = BTreeMap::new();
198        for receipt in self.receipts()? {
199            if receipt.recorded_at.date_naive() != day {
200                continue;
201            }
202            let totals = groups
203                .entry(DailyUsageKey {
204                    user_id: receipt.user_id,
205                    model: receipt.actual_model,
206                })
207                .or_insert_with(DailyUsage::default);
208            totals.calls = totals.calls.saturating_add(1);
209            match receipt.metering {
210                Metering::Tokens(tokens) => {
211                    totals.token_calls = totals.token_calls.saturating_add(1);
212                    totals.tokens.add_assign(tokens);
213                }
214                Metering::DurationSeconds { seconds } => {
215                    totals.duration_calls = totals.duration_calls.saturating_add(1);
216                    totals.duration_seconds += seconds;
217                }
218                Metering::Unavailable => {
219                    totals.unmetered_calls = totals.unmetered_calls.saturating_add(1);
220                }
221            }
222        }
223        Ok(groups)
224    }
225}
226
227fn receipt_io(operation: &str, path: &Path, error: std::io::Error) -> Error {
228    Error::internal(
229        "receipt_io_failed",
230        format!(
231            "Could not {operation} usage receipts at {}: {error}",
232            path.display()
233        ),
234    )
235}
236
237#[cfg(unix)]
238fn sync_directory(path: &Path) -> std::io::Result<()> {
239    std::fs::File::open(path)?.sync_all()
240}
241
242#[cfg(not(unix))]
243fn sync_directory(_path: &Path) -> std::io::Result<()> {
244    Ok(())
245}
246
247#[cfg(test)]
248mod tests {
249    use super::*;
250
251    fn temporary_store() -> (PathBuf, ReceiptStore) {
252        let root = std::env::temp_dir().join(format!("kennedy-usage-{}", Uuid::new_v4()));
253        let store = ReceiptStore::open(&root).unwrap();
254        (root, store)
255    }
256
257    #[test]
258    fn one_file_is_written_for_each_call_and_daily_totals_scan_them() {
259        let (root, store) = temporary_store();
260        let mut first = UsageReceipt::new(
261            "user-a",
262            "search",
263            "gemini-2.5-flash",
264            "gemini-2.5-flash",
265            Metering::Tokens(TokenUsage {
266                input_tokens: 10,
267                cached_input_tokens: 4,
268                thinking_tokens: 3,
269                output_tokens: 2,
270            }),
271        );
272        first.recorded_at = DateTime::parse_from_rfc3339("2026-07-26T01:00:00Z")
273            .unwrap()
274            .with_timezone(&Utc);
275        let mut second = first.clone();
276        second.id = Uuid::new_v4();
277        second.user_id = "user-b".into();
278        store.record(&first).unwrap();
279        store.record(&second).unwrap();
280
281        assert_eq!(
282            fs::read_dir(&root).unwrap().flatten().count(),
283            2,
284            "there should be exactly one durable file per call"
285        );
286        let totals = store
287            .daily_usage(NaiveDate::from_ymd_opt(2026, 7, 26).unwrap())
288            .unwrap();
289        assert_eq!(totals.len(), 2);
290        assert_eq!(
291            totals[&DailyUsageKey {
292                user_id: "user-a".into(),
293                model: "gemini-2.5-flash".into(),
294            }]
295                .tokens
296                .cached_input_tokens,
297            4
298        );
299        fs::remove_dir_all(root).unwrap();
300    }
301}
kcode_intelligence_router/receipts.rs

kcode_intelligence_router/
receipts.rs
