1use std::collections::HashMap;
32use std::sync::{Arc, OnceLock, RwLock};
33
34use serde::{Deserialize, Serialize};
35
36use super::model_pricing::ModelCost;
37
38const REFRESH_INTERVAL_SECS: u64 = 12 * 60 * 60;
41
42const CACHE_FILE: &str = "model-prices.json";
44
45const MODELS_URL: &str = "https://openrouter.ai/api/v1/models";
46
47const LITELLM_URL: &str =
50 "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json";
51
52#[derive(Debug, Clone, Default, Serialize, Deserialize)]
55pub struct LivePriceTable {
56 pub fetched_at: u64,
58 pub models: HashMap<String, ModelCost>,
59}
60
61impl LivePriceTable {
62 #[must_use]
64 pub fn len(&self) -> usize {
65 self.models.len()
66 }
67
68 #[must_use]
69 pub fn is_empty(&self) -> bool {
70 self.models.is_empty()
71 }
72}
73
74fn snapshot() -> &'static RwLock<Option<Arc<LivePriceTable>>> {
75 static SNAP: OnceLock<RwLock<Option<Arc<LivePriceTable>>>> = OnceLock::new();
76 SNAP.get_or_init(|| RwLock::new(None))
77}
78
79fn enabled() -> bool {
81 let v = std::env::var("LEAN_CTX_LIVE_PRICING").unwrap_or_default();
82 !matches!(
83 v.trim().to_ascii_lowercase().as_str(),
84 "off" | "0" | "false" | "no"
85 )
86}
87
88#[must_use]
92pub fn lookup(model: &str) -> Option<(String, ModelCost)> {
93 if !enabled() {
94 return None;
95 }
96 let guard = snapshot()
97 .read()
98 .unwrap_or_else(std::sync::PoisonError::into_inner);
99 let table = guard.as_ref()?;
100 for key in lookup_candidates(model) {
101 if let Some(cost) = table.models.get(&key) {
102 return Some((key, *cost));
103 }
104 }
105 None
106}
107
108pub fn ensure_loaded() -> usize {
112 if !enabled() {
113 return 0;
114 }
115 {
116 let guard = snapshot()
117 .read()
118 .unwrap_or_else(std::sync::PoisonError::into_inner);
119 if let Some(t) = guard.as_ref() {
120 return t.len();
121 }
122 }
123 let Some(table) = load_cache_file() else {
124 return 0;
125 };
126 let len = table.len();
127 install(table);
128 len
129}
130
131pub fn install(table: LivePriceTable) {
133 let mut guard = snapshot()
134 .write()
135 .unwrap_or_else(std::sync::PoisonError::into_inner);
136 *guard = Some(Arc::new(table));
137}
138
139#[must_use]
142pub fn status() -> Option<(u64, usize)> {
143 if !enabled() {
144 return None;
145 }
146 let guard = snapshot()
147 .read()
148 .unwrap_or_else(std::sync::PoisonError::into_inner);
149 guard.as_ref().map(|t| (t.fetched_at, t.len()))
150}
151
152#[cfg(test)]
154pub fn clear_for_tests() {
155 let mut guard = snapshot()
156 .write()
157 .unwrap_or_else(std::sync::PoisonError::into_inner);
158 *guard = None;
159}
160
161fn cache_path() -> Option<std::path::PathBuf> {
162 crate::core::paths::cache_dir()
163 .ok()
164 .map(|d| d.join(CACHE_FILE))
165}
166
167fn load_cache_file() -> Option<LivePriceTable> {
168 let path = cache_path()?;
169 let raw = std::fs::read(path).ok()?;
170 let table: LivePriceTable = serde_json::from_slice(&raw).ok()?;
171 if table.is_empty() { None } else { Some(table) }
172}
173
174fn store_cache_file(table: &LivePriceTable) {
176 let Some(path) = cache_path() else { return };
177 if let Some(dir) = path.parent()
178 && std::fs::create_dir_all(dir).is_err()
179 {
180 return;
181 }
182 let Ok(json) = serde_json::to_vec(table) else {
183 return;
184 };
185 let tmp = path.with_extension("json.tmp");
186 if std::fs::write(&tmp, json).is_ok() {
187 let _ = std::fs::rename(&tmp, &path);
188 }
189}
190
191async fn fetch_catalog(
193 client: &reqwest::Client,
194 url: &str,
195 parse: fn(&serde_json::Value) -> HashMap<String, ModelCost>,
196) -> anyhow::Result<HashMap<String, ModelCost>> {
197 let body = client
198 .get(url)
199 .timeout(std::time::Duration::from_secs(30))
200 .send()
201 .await?
202 .error_for_status()?
203 .bytes()
204 .await?;
205 let json: serde_json::Value = serde_json::from_slice(&body)?;
206 let map = parse(&json);
207 anyhow::ensure!(!map.is_empty(), "catalog {url} parsed empty");
208 Ok(map)
209}
210
211pub async fn refresh_now(client: &reqwest::Client) -> anyhow::Result<usize> {
221 let (openrouter, litellm) = tokio::join!(
222 fetch_catalog(client, MODELS_URL, parse_openrouter_models),
223 fetch_catalog(client, LITELLM_URL, parse_litellm_models),
224 );
225
226 let mut models = match &openrouter {
227 Ok(map) => map.clone(),
228 Err(e) => {
229 tracing::warn!("OpenRouter price catalog unavailable: {e:#}");
230 HashMap::new()
231 }
232 };
233 match &litellm {
234 Ok(map) => {
236 for (k, v) in map {
237 models.entry(k.clone()).or_insert(*v);
238 }
239 }
240 Err(e) => tracing::warn!("LiteLLM price catalog unavailable: {e:#}"),
241 }
242 anyhow::ensure!(
243 !models.is_empty(),
244 "no price catalog reachable (OpenRouter: {}, LiteLLM: {})",
245 openrouter
246 .as_ref()
247 .map_or_else(ToString::to_string, |m| format!("{} keys", m.len())),
248 litellm
249 .as_ref()
250 .map_or_else(ToString::to_string, |m| format!("{} keys", m.len())),
251 );
252
253 let table = LivePriceTable {
254 fetched_at: std::time::SystemTime::now()
255 .duration_since(std::time::UNIX_EPOCH)
256 .map_or(0, |d| d.as_secs()),
257 models,
258 };
259 let len = table.len();
260 store_cache_file(&table);
261 install(table);
262 Ok(len)
263}
264
265pub fn spawn_background_refresh() {
270 static SPAWNED: OnceLock<()> = OnceLock::new();
271 if !enabled() {
272 return;
273 }
274 ensure_loaded();
275 if SPAWNED.set(()).is_err() {
276 return;
277 }
278 tokio::spawn(async {
279 let client = reqwest::Client::new();
280 loop {
281 let stale = {
282 let guard = snapshot()
283 .read()
284 .unwrap_or_else(std::sync::PoisonError::into_inner);
285 guard.as_ref().is_none_or(|t| {
286 let now = std::time::SystemTime::now()
287 .duration_since(std::time::UNIX_EPOCH)
288 .map_or(0, |d| d.as_secs());
289 now.saturating_sub(t.fetched_at) >= REFRESH_INTERVAL_SECS
290 })
291 };
292 if stale {
293 match refresh_now(&client).await {
294 Ok(n) => tracing::info!("live model pricing refreshed ({n} lookup keys)"),
295 Err(e) => tracing::warn!(
296 "live model pricing refresh failed (keeping previous table): {e:#}"
297 ),
298 }
299 }
300 tokio::time::sleep(std::time::Duration::from_secs(REFRESH_INTERVAL_SECS / 12)).await;
301 }
302 });
303}
304
305fn canon(s: &str) -> String {
308 s.trim().to_lowercase().replace([' ', '.'], "-")
309}
310
311fn strip_date_suffix(s: &str) -> Option<&str> {
313 let (base, tail) = s.rsplit_once('-')?;
314 if tail.len() == 8 && tail.bytes().all(|b| b.is_ascii_digit()) {
315 Some(base)
316 } else {
317 None
318 }
319}
320
321fn lookup_candidates(model: &str) -> Vec<String> {
324 let full = canon(model);
325 if full.is_empty() {
326 return Vec::new();
327 }
328 let mut out = vec![full.clone()];
329 let mut push = |s: String| {
330 if !s.is_empty() && !out.contains(&s) {
331 out.push(s);
332 }
333 };
334 let no_vendor = full.split_once('/').map(|(_, m)| m.to_string());
335 if let Some(nv) = &no_vendor {
336 push(nv.clone());
337 }
338 for base in [Some(full.as_str()), no_vendor.as_deref()]
339 .into_iter()
340 .flatten()
341 {
342 let no_variant = base.split(':').next().unwrap_or(base);
343 push(no_variant.to_string());
344 if let Some(no_date) = strip_date_suffix(no_variant) {
345 push(no_date.to_string());
346 }
347 }
348 out
349}
350
351fn index_keys(id: &str) -> (Vec<String>, bool) {
355 let full = canon(id);
356 let has_variant = full.contains(':');
357 let mut keys = vec![full.clone()];
358 let mut push = |s: String| {
359 if !s.is_empty() && !keys.contains(&s) {
360 keys.push(s);
361 }
362 };
363 let no_vendor = full.split_once('/').map(|(_, m)| m.to_string());
364 if let Some(nv) = &no_vendor {
365 push(nv.clone());
366 }
367 for base in [Some(full.as_str()), no_vendor.as_deref()]
368 .into_iter()
369 .flatten()
370 {
371 let no_variant = base.split(':').next().unwrap_or(base);
372 push(no_variant.to_string());
373 if let Some(no_date) = strip_date_suffix(no_variant) {
374 push(no_date.to_string());
375 }
376 }
377 (keys, has_variant)
378}
379
380fn per_mtok(pricing: &serde_json::Value, field: &str) -> Option<f64> {
382 let v = pricing.get(field)?;
383 let n = v
384 .as_str()
385 .map_or_else(|| v.as_f64(), |s| s.trim().parse::<f64>().ok())?;
386 if n.is_finite() && n >= 0.0 {
387 Some(n * 1_000_000.0)
388 } else {
389 None
390 }
391}
392
393fn parse_openrouter_models(json: &serde_json::Value) -> HashMap<String, ModelCost> {
400 let mut map: HashMap<String, ModelCost> = HashMap::new();
401 let Some(data) = json.get("data").and_then(serde_json::Value::as_array) else {
402 return map;
403 };
404
405 let mut deferred: Vec<(&serde_json::Value, &str)> = Vec::new();
408 let absorb = |map: &mut HashMap<String, ModelCost>, m: &serde_json::Value, id: &str| {
409 let Some(pricing) = m.get("pricing") else {
410 return;
411 };
412 let (Some(input), Some(output)) =
413 (per_mtok(pricing, "prompt"), per_mtok(pricing, "completion"))
414 else {
415 return;
416 };
417 let cost = ModelCost {
418 input_per_m: input,
419 output_per_m: output,
420 cache_write_per_m: per_mtok(pricing, "input_cache_write").unwrap_or(input),
421 cache_read_per_m: per_mtok(pricing, "input_cache_read").unwrap_or(input),
422 };
423 let (keys, _) = index_keys(id);
424 for key in keys {
425 map.entry(key).or_insert(cost);
426 }
427 if let Some(slug) = m.get("canonical_slug").and_then(serde_json::Value::as_str) {
430 let (slug_keys, _) = index_keys(slug);
431 for key in slug_keys {
432 map.entry(key).or_insert(cost);
433 }
434 }
435 };
436
437 for m in data {
438 let Some(id) = m.get("id").and_then(serde_json::Value::as_str) else {
439 continue;
440 };
441 if canon(id).contains(':') {
442 deferred.push((m, id));
443 } else {
444 absorb(&mut map, m, id);
445 }
446 }
447 for (m, id) in deferred {
448 absorb(&mut map, m, id);
449 }
450 map
451}
452
453fn litellm_per_mtok(entry: &serde_json::Value, field: &str) -> Option<f64> {
456 let n = entry.get(field)?.as_f64()?;
457 if n.is_finite() && n >= 0.0 {
458 Some(n * 1_000_000.0)
459 } else {
460 None
461 }
462}
463
464fn parse_litellm_models(json: &serde_json::Value) -> HashMap<String, ModelCost> {
474 let mut map: HashMap<String, ModelCost> = HashMap::new();
475 let Some(entries) = json.as_object() else {
476 return map;
477 };
478 for (key, entry) in entries {
479 if key == "sample_spec" || !entry.is_object() {
480 continue;
481 }
482 let Some(input) = litellm_per_mtok(entry, "input_cost_per_token") else {
483 continue;
484 };
485 let output = match litellm_per_mtok(entry, "output_cost_per_token") {
488 Some(o) => o,
489 None if entry.get("mode").and_then(serde_json::Value::as_str) == Some("embedding") => {
490 0.0
491 }
492 None => continue,
493 };
494 let cost = ModelCost {
495 input_per_m: input,
496 output_per_m: output,
497 cache_write_per_m: litellm_per_mtok(entry, "cache_creation_input_token_cost")
498 .unwrap_or(input),
499 cache_read_per_m: litellm_per_mtok(entry, "cache_read_input_token_cost")
500 .unwrap_or(input),
501 };
502 let (keys, _) = index_keys(key);
503 for k in keys {
504 map.entry(k).or_insert(cost);
505 }
506 }
507 map
508}
509
510#[cfg(test)]
511mod tests {
512 use super::*;
513
514 fn fixture() -> serde_json::Value {
515 serde_json::json!({
516 "data": [
517 {
518 "id": "deepseek/deepseek-v4-flash",
519 "canonical_slug": "deepseek/deepseek-v4-flash-20260423",
520 "pricing": {"prompt": "0.00000007", "completion": "0.00000028",
521 "input_cache_read": "0.000000007"}
522 },
523 {
524 "id": "anthropic/claude-sonnet-5",
525 "canonical_slug": "anthropic/claude-sonnet-5-20260630",
526 "pricing": {"prompt": "0.000002", "completion": "0.00001",
527 "web_search": "0.01",
528 "input_cache_read": "0.0000002",
529 "input_cache_write": "0.0000025"}
530 },
531 {
532 "id": "poolside/laguna-xs-2.1:free",
533 "canonical_slug": "poolside/laguna-xs-2.1-20260625",
534 "pricing": {"prompt": "0", "completion": "0"}
535 },
536 {
537 "id": "poolside/laguna-xs-2.1",
538 "canonical_slug": "poolside/laguna-xs-2.1-20260625",
539 "pricing": {"prompt": "0.00000006", "completion": "0.00000012"}
540 },
541 {"id": "broken/no-pricing"}
542 ]
543 })
544 }
545
546 #[test]
547 fn parses_usd_per_token_strings_into_per_mtok() {
548 let map = parse_openrouter_models(&fixture());
549 let flash = map.get("deepseek/deepseek-v4-flash").expect("indexed");
550 assert!((flash.input_per_m - 0.07).abs() < 1e-9);
551 assert!((flash.output_per_m - 0.28).abs() < 1e-9);
552 assert!((flash.cache_read_per_m - 0.007).abs() < 1e-9);
553 assert!((flash.cache_write_per_m - 0.07).abs() < 1e-9);
555 assert!(!map.contains_key("broken/no-pricing"));
556 }
557
558 #[test]
559 fn date_stamped_and_vendor_prefixed_names_resolve() {
560 let map = parse_openrouter_models(&fixture());
561 for name in [
563 "deepseek/deepseek-v4-flash-20260423",
564 "deepseek-v4-flash-20260423",
565 "deepseek-v4-flash",
566 ] {
567 let mut found = false;
568 for key in lookup_candidates(name) {
569 if map.contains_key(&key) {
570 found = true;
571 break;
572 }
573 }
574 assert!(found, "{name} must resolve against the live table");
575 }
576 }
577
578 #[test]
579 fn free_variant_never_hijacks_the_paid_model_key() {
580 let map = parse_openrouter_models(&fixture());
581 let paid = map
582 .get("poolside/laguna-xs-2-1")
583 .expect("paid model indexed");
584 assert!(
585 paid.input_per_m > 0.0,
586 "canonical key must carry the paid price"
587 );
588 let free = map
589 .get("poolside/laguna-xs-2-1:free")
590 .expect("variant indexed");
591 assert_eq!(
592 free.input_per_m, 0.0,
593 "the :free variant stays free under its full name"
594 );
595 }
596
597 #[test]
598 fn dot_dash_and_case_unify() {
599 assert_eq!(canon("Claude-Opus-4.5"), "claude-opus-4-5");
600 assert_eq!(
601 strip_date_suffix("deepseek-v4-flash-20260423"),
602 Some("deepseek-v4-flash")
603 );
604 assert_eq!(
605 strip_date_suffix("claude-opus-4-5"),
606 None,
607 "short numeric tails are versions"
608 );
609 assert_eq!(strip_date_suffix("no-date"), None);
610 }
611
612 fn litellm_fixture() -> serde_json::Value {
613 serde_json::json!({
614 "sample_spec": {
615 "input_cost_per_token": 0.0,
616 "output_cost_per_token": 0.0,
617 "mode": "one of: chat, embedding, completion, …"
618 },
619 "azure/gpt-4o": {
620 "input_cost_per_token": 2.5e-6,
621 "output_cost_per_token": 1e-5,
622 "cache_read_input_token_cost": 1.25e-6,
623 "mode": "chat",
624 "litellm_provider": "azure"
625 },
626 "bedrock/anthropic.claude-sonnet-4-5": {
627 "input_cost_per_token": 3e-6,
628 "output_cost_per_token": 1.5e-5,
629 "cache_creation_input_token_cost": 3.75e-6,
630 "cache_read_input_token_cost": 3e-7,
631 "mode": "chat"
632 },
633 "text-embedding-3-small": {
634 "input_cost_per_token": 2e-8,
635 "mode": "embedding"
636 },
637 "vertex_ai/imagegeneration": {
638 "output_cost_per_image": 0.02,
639 "mode": "image_generation"
640 }
641 })
642 }
643
644 #[test]
645 fn litellm_map_parses_prefixes_embeddings_and_skips_junk() {
646 let map = parse_litellm_models(&litellm_fixture());
647
648 let azure = map.get("azure/gpt-4o").expect("prefixed key");
650 assert!((azure.input_per_m - 2.5).abs() < 1e-9);
651 assert!((azure.output_per_m - 10.0).abs() < 1e-9);
652 assert!((azure.cache_read_per_m - 1.25).abs() < 1e-9);
653 assert!(map.contains_key("gpt-4o"), "bare key indexed too");
654
655 let bedrock = map
657 .get("bedrock/anthropic-claude-sonnet-4-5")
658 .expect("bedrock key (canon: dots→dashes)");
659 assert!((bedrock.cache_write_per_m - 3.75).abs() < 1e-9);
660
661 let emb = map.get("text-embedding-3-small").expect("embedding");
663 assert!((emb.input_per_m - 0.02).abs() < 1e-9);
664 assert!((emb.output_per_m - 0.0).abs() < f64::EPSILON);
665
666 assert!(!map.contains_key("sample_spec"));
668 assert!(!map.contains_key("vertex_ai/imagegeneration"));
669 }
670
671 #[test]
672 fn merged_table_lets_openrouter_win_and_litellm_fill_gaps() {
673 let mut merged = parse_openrouter_models(&fixture());
676 for (k, v) in parse_litellm_models(&litellm_fixture()) {
677 merged.entry(k).or_insert(v);
678 }
679
680 assert!(merged.contains_key("azure/gpt-4o"));
682 let flash = merged
684 .get("deepseek/deepseek-v4-flash")
685 .expect("openrouter");
686 assert!((flash.input_per_m - 0.07).abs() < 1e-9);
687 }
688
689 #[test]
690 fn snapshot_lookup_respects_kill_switch_and_install() {
691 let _lock = crate::core::data_dir::test_env_lock();
692 clear_for_tests();
693 assert!(
694 lookup("deepseek/deepseek-v4-flash").is_none(),
695 "empty snapshot"
696 );
697
698 install(LivePriceTable {
699 fetched_at: 1,
700 models: parse_openrouter_models(&fixture()),
701 });
702 let (_, cost) = lookup("deepseek/deepseek-v4-flash-20260423").expect("live hit");
703 assert!((cost.input_per_m - 0.07).abs() < 1e-9);
704
705 crate::test_env::set_var("LEAN_CTX_LIVE_PRICING", "off");
706 assert!(
707 lookup("deepseek/deepseek-v4-flash").is_none(),
708 "kill switch"
709 );
710 crate::test_env::remove_var("LEAN_CTX_LIVE_PRICING");
711 clear_for_tests();
712 }
713}