forge_workspace/store/
pricing.rs

1//! Cached LiteLLM pricing on the redb `pricing_cache` table.
2//!
3//! One row holds the most-recent fetch: the raw json plus its
4//! `fetched_at`. The usage scan prices from this cache (empty until the
5//! first background fetch lands); the fetch refreshes it about once a
6//! day.
7
8use std::time::SystemTime;
9
10use anyhow::Context;
11use redb::TableDefinition;
12use serde::{Deserialize, Serialize};
13
14use super::Db;
15
16const PRICING_CACHE: TableDefinition<&str, &[u8]> = TableDefinition::new("pricing_cache");
17
18/// The single row key; there is only ever one cached pricing snapshot.
19const KEY: &str = "litellm";
20
21/// A fetched LiteLLM pricing file plus the instant it landed.
22#[derive(Debug, PartialEq, Serialize, Deserialize)]
23pub struct CachedPricing {
24    pub fetched_at: SystemTime,
25    pub json: String,
26}
27
28/// The cached pricing snapshot, or `None` when nothing has been fetched
29/// yet (or the row fails to decode).
30pub fn load(db: &Db) -> anyhow::Result<Option<CachedPricing>> {
31    let txn = db.database().begin_read()?;
32    let table = match txn.open_table(PRICING_CACHE) {
33        Ok(t) => t,
34        Err(redb::TableError::TableDoesNotExist(_)) => return Ok(None),
35        Err(e) => return Err(e.into()),
36    };
37    match table.get(KEY)? {
38        Some(value) => {
39            Ok(Some(serde_json::from_slice(value.value()).context("decode cached pricing")?))
40        }
41        None => Ok(None),
42    }
43}
44
45/// Overwrite the cached pricing snapshot with `entry`.
46pub fn store(db: &Db, entry: &CachedPricing) -> anyhow::Result<()> {
47    let value = serde_json::to_vec(entry).context("serialize cached pricing")?;
48    let txn = db.database().begin_write()?;
49    {
50        let mut table = txn.open_table(PRICING_CACHE)?;
51        table.insert(KEY, value.as_slice())?;
52    }
53    txn.commit()?;
54    Ok(())
55}
56
57#[cfg(test)]
58mod tests {
59    use super::*;
60    use std::time::Duration;
61    use tempfile::tempdir;
62
63    fn open_db() -> (tempfile::TempDir, Db) {
64        let dir = tempdir().expect("tempdir");
65        let db = Db::open(&dir.path().join("db.redb")).expect("open db");
66        (dir, db)
67    }
68
69    #[test]
70    fn load_is_none_before_any_fetch() {
71        let (_dir, db) = open_db();
72        assert!(load(&db).expect("load").is_none());
73    }
74
75    #[test]
76    fn store_then_load_round_trips_and_overwrites() {
77        let (_dir, db) = open_db();
78        let first = CachedPricing {
79            fetched_at: SystemTime::UNIX_EPOCH + Duration::from_secs(1_700_000_000),
80            json: r#"{"m":{}}"#.to_owned(),
81        };
82        store(&db, &first).expect("store");
83        assert_eq!(load(&db).expect("load"), Some(first));
84
85        let second = CachedPricing {
86            fetched_at: SystemTime::UNIX_EPOCH + Duration::from_secs(1_800_000_000),
87            json: r#"{"n":{}}"#.to_owned(),
88        };
89        store(&db, &second).expect("store again");
90        assert_eq!(load(&db).expect("load").expect("present").json, r#"{"n":{}}"#);
91    }
92
93    #[test]
94    fn load_surfaces_a_corrupt_blob_as_error() {
95        let (_dir, db) = open_db();
96        let txn = db.database().begin_write().expect("begin");
97        {
98            let mut table = txn.open_table(PRICING_CACHE).expect("open table");
99            table.insert(KEY, b"not json".as_slice()).expect("insert corrupt");
100        }
101        txn.commit().expect("commit");
102        // A decode failure is a diagnosable error the caller can warn on,
103        // not a silent cache miss.
104        assert!(load(&db).is_err(), "a corrupt row surfaces as an error");
105    }
106}
forge_workspace/store/pricing.rs

forge_workspace/store/
pricing.rs
