otelite_api/
pricing_cache.rs

1//! Server-side cache of the LiteLLM pricing database.
2//!
3//! On server startup we spawn a background task that fetches the LiteLLM JSON
4//! from GitHub and swaps it into an `ArcSwap`-style `RwLock`. If the fetch
5//! fails (offline, rate-limit, upstream 5xx) the cache stays at whatever it
6//! was — `PricingDatabase::empty()` initially, meaning the Claude fallback
7//! table in `otelite-core::pricing` handles every lookup until the first
8//! successful refresh.
9//!
10//! The background task then sleeps for [`REFRESH_INTERVAL`] and tries again.
11//! Failures are logged but never fatal.
12
13use otelite_core::pricing::{PricingDatabase, LITELLM_RAW_URL};
14use std::sync::Arc;
15use std::time::Duration;
16use tokio::sync::RwLock;
17use tracing::{debug, info, warn};
18
19/// Refresh the pricing database once per day. LiteLLM updates the JSON when new
20/// models are added or prices change, typically weekly, so daily is ample.
21pub const REFRESH_INTERVAL: Duration = Duration::from_secs(24 * 60 * 60);
22
23/// Initial fetch timeout — if GitHub is slow we'd rather start the server with
24/// an empty cache and refresh in the background than block startup.
25pub const INITIAL_FETCH_TIMEOUT: Duration = Duration::from_secs(10);
26
27/// Ongoing refresh timeout can be more generous.
28pub const REFRESH_FETCH_TIMEOUT: Duration = Duration::from_secs(30);
29
30/// Thread-safe handle to the current pricing database.
31#[derive(Clone)]
32pub struct PricingCache {
33    inner: Arc<RwLock<PricingState>>,
34}
35
36#[derive(Clone)]
37pub struct PricingState {
38    pub db: PricingDatabase,
39    /// Unix millis of the last successful LiteLLM fetch, or None if never.
40    pub last_fetched_unix_ms: Option<i64>,
41    /// Unix millis of the last failed LiteLLM fetch (for surfacing staleness).
42    pub last_failed_unix_ms: Option<i64>,
43}
44
45impl PricingCache {
46    pub fn new() -> Self {
47        Self {
48            inner: Arc::new(RwLock::new(PricingState {
49                db: PricingDatabase::empty(),
50                last_fetched_unix_ms: None,
51                last_failed_unix_ms: None,
52            })),
53        }
54    }
55
56    pub async fn snapshot(&self) -> PricingState {
57        self.inner.read().await.clone()
58    }
59
60    /// Fetch LiteLLM JSON and swap into the cache. Returns Ok even if we had
61    /// to fall back to the empty DB — the only error is a panicked lock.
62    async fn refresh(&self, timeout: Duration) {
63        let client = match reqwest::Client::builder().timeout(timeout).build() {
64            Ok(c) => c,
65            Err(e) => {
66                warn!("pricing: failed to build HTTP client: {e}");
67                return;
68            },
69        };
70        match client.get(LITELLM_RAW_URL).send().await {
71            Ok(resp) if resp.status().is_success() => {
72                let body = match resp.text().await {
73                    Ok(b) => b,
74                    Err(e) => {
75                        warn!("pricing: failed to read LiteLLM response: {e}");
76                        self.record_failure().await;
77                        return;
78                    },
79                };
80                match PricingDatabase::from_litellm_json(&body) {
81                    Ok(db) => {
82                        let count = db.len();
83                        let mut state = self.inner.write().await;
84                        state.db = db;
85                        state.last_fetched_unix_ms = Some(now_unix_ms());
86                        info!("pricing: loaded {count} LiteLLM entries");
87                    },
88                    Err(e) => {
89                        warn!("pricing: failed to parse LiteLLM JSON: {e}");
90                        self.record_failure().await;
91                    },
92                }
93            },
94            Ok(resp) => {
95                warn!("pricing: LiteLLM fetch returned HTTP {}", resp.status());
96                self.record_failure().await;
97            },
98            Err(e) => {
99                debug!("pricing: LiteLLM fetch failed: {e}");
100                self.record_failure().await;
101            },
102        }
103    }
104
105    async fn record_failure(&self) {
106        let mut state = self.inner.write().await;
107        state.last_failed_unix_ms = Some(now_unix_ms());
108    }
109
110    /// Spawn a background task that does the initial fetch and then refreshes
111    /// on a timer. Returns a clone of the cache for the caller to retain.
112    pub fn spawn_refresher(self) -> Self {
113        let task_cache = self.clone();
114        tokio::spawn(async move {
115            task_cache.refresh(INITIAL_FETCH_TIMEOUT).await;
116            let mut ticker = tokio::time::interval(REFRESH_INTERVAL);
117            // The interval's first tick fires immediately — skip past it since
118            // we just fetched.
119            ticker.tick().await;
120            loop {
121                ticker.tick().await;
122                task_cache.refresh(REFRESH_FETCH_TIMEOUT).await;
123            }
124        });
125        self
126    }
127}
128
129impl Default for PricingCache {
130    fn default() -> Self {
131        Self::new()
132    }
133}
134
135fn now_unix_ms() -> i64 {
136    use std::time::{SystemTime, UNIX_EPOCH};
137    SystemTime::now()
138        .duration_since(UNIX_EPOCH)
139        .map(|d| d.as_millis() as i64)
140        .unwrap_or(0)
141}
otelite_api/pricing_cache.rs

otelite_api/
pricing_cache.rs
