Skip to main content

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::{CostEstimate, Error, Result, estimate_cost};
14
15const RECEIPT_VERSION: u32 = 3;
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    /// Stable identity assigned by the application to this provider call.
56    #[serde(default, skip_serializing_if = "Option::is_none")]
57    pub operation_id: Option<Uuid>,
58    /// Causal parent operation, when this call is descendant inference.
59    #[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    /// Locally calculated standard-tier provider cost, when the model and
70    /// provider usage can be priced.
71    #[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/// One grouping key in a linearly calculated daily report.
106#[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/// Totals calculated by opening every receipt for the requested UTC day.
114#[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/// Intentionally simple one-file-per-call receipt storage.
128#[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    /// Atomically writes exactly one new JSON receipt.
151    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    /// Opens and decodes every receipt file. There are deliberately no indexes yet.
188    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    /// Calculates per-user, per-model totals for one UTC day by scanning all receipts.
212    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}