返回 CodeWhale
model_catalog.rs
根目录 / crates / tui / src / model_catalog.rs
1 //! Offline model metadata catalog (#3072).
2 //!
3 //! This module adds a secret-free metadata layer in front of the legacy model
4 //! tables. It is intentionally conservative: startup reads a local cache plus a
5 //! bundled snapshot, never performs a network refresh, and only overrides a
6 //! legacy fact when the active catalog entry actually carries that field.
7
8 use std::collections::BTreeMap;
9 use std::path::PathBuf;
10 use std::sync::{OnceLock, RwLock};
11
12 use anyhow::Result;
13 use chrono::{DateTime, Duration, Utc};
14 use serde::{Deserialize, Serialize};
15
16 const BUNDLED_CATALOG_JSON: &str = include_str!("../assets/model_catalog.bundled.json");
17 const OPENROUTER_CACHE_FILE: &str = "openrouter.json";
18
19 #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
20 #[serde(rename_all = "snake_case")]
21 pub enum MetadataProvenance {
22 ProviderApi,
23 Bundled,
24 UserOverride,
25 #[default]
26 Unknown,
27 }
28
29 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
30 pub struct CatalogEntry {
31 pub id: String,
32 #[serde(default, skip_serializing_if = "Option::is_none")]
33 pub context_window: Option<u32>,
34 #[serde(default, skip_serializing_if = "Option::is_none")]
35 pub max_output: Option<u32>,
36 #[serde(default, skip_serializing_if = "Option::is_none")]
37 pub supports_reasoning: Option<bool>,
38 #[serde(default, skip_serializing_if = "Option::is_none")]
39 pub input_usd_per_million: Option<f64>,
40 #[serde(default, skip_serializing_if = "Option::is_none")]
41 pub output_usd_per_million: Option<f64>,
42 #[serde(default, skip_serializing_if = "Vec::is_empty")]
43 pub modalities: Vec<String>,
44 #[serde(default, skip_serializing_if = "Vec::is_empty")]
45 pub supported_parameters: Vec<String>,
46 #[serde(default, skip_serializing_if = "Option::is_none")]
47 pub provider_model_id: Option<String>,
48 #[serde(default)]
49 pub provenance: MetadataProvenance,
50 }
51
52 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
53 pub struct CatalogCache {
54 pub schema_version: u32,
55 pub source: String,
56 pub fetched_at: DateTime<Utc>,
57 pub ttl_secs: u64,
58 #[serde(default)]
59 pub entries: BTreeMap<String, CatalogEntry>,
60 }
61
62 impl CatalogCache {
63 #[must_use]
64 pub fn is_stale(&self, now: DateTime<Utc>) -> bool {
65 if now <= self.fetched_at {
66 return false;
67 }
68 let ttl = Duration::seconds(self.ttl_secs.min(i64::MAX as u64) as i64);
69 now.signed_duration_since(self.fetched_at) > ttl
70 }
71 }
72
73 #[derive(Debug, Clone)]
74 pub(crate) struct MergedCatalog {
75 user_overrides: BTreeMap<String, CatalogEntry>,
76 provider_cache: Option<CatalogCache>,
77 bundled: CatalogCache,
78 now: DateTime<Utc>,
79 }
80
81 impl MergedCatalog {
82 pub(crate) fn from_sources(
83 user_overrides: BTreeMap<String, CatalogEntry>,
84 provider_cache: Option<CatalogCache>,
85 bundled: CatalogCache,
86 now: DateTime<Utc>,
87 ) -> Self {
88 Self {
89 user_overrides,
90 provider_cache,
91 bundled,
92 now,
93 }
94 }
95
96 #[must_use]
97 pub(crate) fn resolve(&self, model: &str) -> Option<&CatalogEntry> {
98 if let Some(entry) = entry_for(&self.user_overrides, model) {
99 return Some(entry);
100 }
101 if let Some(provider_cache) = self
102 .provider_cache
103 .as_ref()
104 .filter(|cache| !cache.is_stale(self.now))
105 && let Some(entry) = entry_for(&provider_cache.entries, model)
106 {
107 return Some(entry);
108 }
109 entry_for(&self.bundled.entries, model)
110 }
111 }
112
113 fn entry_for<'a>(
114 entries: &'a BTreeMap<String, CatalogEntry>,
115 model: &str,
116 ) -> Option<&'a CatalogEntry> {
117 entries.get(model).or_else(|| {
118 let lower = model.to_lowercase();
119 (lower != model).then(|| entries.get(&lower)).flatten()
120 })
121 }
122
123 fn active_catalog() -> &'static RwLock<MergedCatalog> {
124 static ACTIVE: OnceLock<RwLock<MergedCatalog>> = OnceLock::new();
125 ACTIVE.get_or_init(|| {
126 RwLock::new(MergedCatalog::from_sources(
127 BTreeMap::new(),
128 load_cached(),
129 bundled_catalog(),
130 Utc::now(),
131 ))
132 })
133 }
134
135 #[must_use]
136 pub fn resolved_entry(model: &str) -> Option<CatalogEntry> {
137 active_catalog()
138 .read()
139 .ok()
140 .and_then(|catalog| catalog.resolve(model).cloned())
141 }
142
143 #[must_use]
144 pub fn resolved_context_window(model: &str) -> Option<u32> {
145 resolved_entry(model).and_then(|entry| entry.context_window)
146 }
147
148 #[must_use]
149 pub fn resolved_max_output(model: &str) -> Option<u32> {
150 resolved_entry(model).and_then(|entry| entry.max_output)
151 }
152
153 #[must_use]
154 pub fn resolved_supports_reasoning(model: &str) -> Option<bool> {
155 resolved_entry(model).and_then(|entry| entry.supports_reasoning)
156 }
157
158 #[must_use]
159 #[cfg_attr(test, allow(dead_code))]
160 pub fn resolved_usd_pricing(model: &str) -> Option<(f64, f64)> {
161 let entry = resolved_entry(model)?;
162 Some((entry.input_usd_per_million?, entry.output_usd_per_million?))
163 }
164
165 pub fn bundled_catalog() -> CatalogCache {
166 serde_json::from_str(BUNDLED_CATALOG_JSON).expect("bundled model catalog must parse")
167 }
168
169 fn catalog_cache_read_path() -> Result<PathBuf> {
170 Ok(codewhale_config::resolve_state_dir("catalog")?.join(OPENROUTER_CACHE_FILE))
171 }
172
173 pub fn load_cached() -> Option<CatalogCache> {
174 let path = catalog_cache_read_path().ok()?;
175 let raw = std::fs::read_to_string(path).ok()?;
176 serde_json::from_str(&raw).ok()
177 }
178
179 #[cfg(test)]
180 static TEST_CATALOG_LOCK: std::sync::LazyLock<std::sync::Mutex<()>> =
181 std::sync::LazyLock::new(|| std::sync::Mutex::new(()));
182
183 #[cfg(test)]
184 pub(crate) fn test_catalog_lock() -> std::sync::MutexGuard<'static, ()> {
185 TEST_CATALOG_LOCK.lock().expect("model catalog test lock")
186 }
187
188 #[cfg(test)]
189 pub(crate) struct ActiveCatalogGuard {
190 previous: MergedCatalog,
191 }
192
193 #[cfg(test)]
194 impl Drop for ActiveCatalogGuard {
195 fn drop(&mut self) {
196 let mut active = active_catalog().write().expect("active catalog write lock");
197 *active = self.previous.clone();
198 }
199 }
200
201 #[cfg(test)]
202 pub(crate) fn replace_active_catalog_for_test(catalog: MergedCatalog) -> ActiveCatalogGuard {
203 let mut active = active_catalog().write().expect("active catalog write lock");
204 let previous = active.clone();
205 *active = catalog;
206 ActiveCatalogGuard { previous }
207 }
208
209 #[cfg(test)]
210 mod tests {
211 use super::*;
212
213 fn entry(id: &str, context_window: u32, provenance: MetadataProvenance) -> CatalogEntry {
214 CatalogEntry {
215 id: id.to_string(),
216 context_window: Some(context_window),
217 max_output: Some(context_window / 2),
218 supports_reasoning: Some(false),
219 input_usd_per_million: None,
220 output_usd_per_million: None,
221 modalities: Vec::new(),
222 supported_parameters: Vec::new(),
223 provider_model_id: None,
224 provenance,
225 }
226 }
227
228 fn cache(
229 fetched_at: DateTime<Utc>,
230 ttl_secs: u64,
231 entries: BTreeMap<String, CatalogEntry>,
232 ) -> CatalogCache {
233 CatalogCache {
234 schema_version: 1,
235 source: "test".to_string(),
236 fetched_at,
237 ttl_secs,
238 entries,
239 }
240 }
241
242 #[test]
243 fn bundled_snapshot_parses_and_is_nonempty() {
244 let bundled = bundled_catalog();
245 assert_eq!(bundled.schema_version, 1);
246 assert!(!bundled.entries.is_empty());
247 assert_eq!(
248 bundled.entries["deepseek-v4-pro"].provenance,
249 MetadataProvenance::Bundled
250 );
251 }
252
253 #[test]
254 fn merge_order_is_user_override_then_provider_then_bundled() {
255 let now = Utc::now();
256 let mut bundled_entries = BTreeMap::new();
257 bundled_entries.insert(
258 "sample/model".to_string(),
259 entry("sample/model", 1_000, MetadataProvenance::Bundled),
260 );
261 let bundled = cache(now, 3600, bundled_entries);
262
263 let mut provider_entries = BTreeMap::new();
264 provider_entries.insert(
265 "sample/model".to_string(),
266 entry("sample/model", 2_000, MetadataProvenance::ProviderApi),
267 );
268 let provider_cache = cache(now, 3600, provider_entries);
269
270 let mut override_entries = BTreeMap::new();
271 override_entries.insert(
272 "sample/model".to_string(),
273 entry("sample/model", 3_000, MetadataProvenance::UserOverride),
274 );
275
276 let merged =
277 MergedCatalog::from_sources(override_entries, Some(provider_cache), bundled, now);
278 let resolved = merged.resolve("sample/model").expect("resolved");
279 assert_eq!(resolved.context_window, Some(3_000));
280 assert_eq!(resolved.provenance, MetadataProvenance::UserOverride);
281 }
282
283 #[test]
284 fn stale_cache_is_ignored_for_facts() {
285 let now = Utc::now();
286 let mut bundled_entries = BTreeMap::new();
287 bundled_entries.insert(
288 "sample/model".to_string(),
289 entry("sample/model", 1_000, MetadataProvenance::Bundled),
290 );
291 let bundled = cache(now, 3600, bundled_entries);
292
293 let mut provider_entries = BTreeMap::new();
294 provider_entries.insert(
295 "sample/model".to_string(),
296 entry("sample/model", 9_000, MetadataProvenance::ProviderApi),
297 );
298 let provider_cache = cache(now - Duration::seconds(10), 1, provider_entries);
299 assert!(provider_cache.is_stale(now));
300
301 let merged =
302 MergedCatalog::from_sources(BTreeMap::new(), Some(provider_cache), bundled, now);
303 let resolved = merged.resolve("sample/model").expect("resolved");
304 assert_eq!(resolved.context_window, Some(1_000));
305 assert_eq!(resolved.provenance, MetadataProvenance::Bundled);
306 }
307
308 #[test]
309 fn cache_roundtrip_serializes_no_secret_fields() {
310 let mut entries = BTreeMap::new();
311 entries.insert(
312 "sample/model".to_string(),
313 CatalogEntry {
314 input_usd_per_million: Some(0.25),
315 output_usd_per_million: Some(1.25),
316 ..entry("sample/model", 32_000, MetadataProvenance::ProviderApi)
317 },
318 );
319 let cache = cache(Utc::now(), 60, entries);
320 let json = serde_json::to_string_pretty(&cache).expect("serialize");
321 let lowered = json.to_lowercase();
322 for forbidden in ["api_key", "authorization", "token", "secret"] {
323 assert!(
324 !lowered.contains(forbidden),
325 "cache JSON must not contain auth field {forbidden}: {json}"
326 );
327 }
328 let parsed: CatalogCache = serde_json::from_str(&json).expect("roundtrip");
329 assert_eq!(parsed.entries.len(), 1);
330 }
331 }
332
332 lines RUST