1use std::collections::BTreeMap;
15use std::path::PathBuf;
16
17use rpi_ai::model::Model;
18use serde::{Deserialize, Serialize};
19
20pub const DEFAULT_CATALOG_BASE_URL: &str = "https://pi.dev";
22
23pub const REMOTE_CATALOG_REFRESH_INTERVAL_MS: i64 = 4 * 60 * 60 * 1000;
25
26#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
28#[serde(rename_all = "camelCase")]
29pub struct CatalogEntry {
30 #[serde(default)]
31 pub models: Vec<Model>,
32 #[serde(default)]
34 pub checked_at: i64,
35 #[serde(default)]
37 pub last_modified: i64,
38 #[serde(default, skip_serializing_if = "Option::is_none")]
39 pub etag: Option<String>,
40}
41
42#[derive(Debug, Clone, Default, Serialize, Deserialize)]
44pub struct CatalogStore {
45 #[serde(flatten)]
46 pub providers: BTreeMap<String, CatalogEntry>,
47}
48
49impl CatalogStore {
50 pub fn path() -> Result<PathBuf, String> {
51 let dir = crate::config::agent_dir().map_err(|e| e.to_string())?;
52 Ok(dir.join("model-catalog.json"))
53 }
54
55 pub fn load() -> Self {
58 Self::load_from(&Self::path().ok())
59 }
60
61 pub fn load_from(path: &Option<PathBuf>) -> Self {
62 let Some(path) = path else {
63 return Self::default();
64 };
65 match std::fs::read_to_string(path) {
66 Ok(text) => serde_json::from_str(&text).unwrap_or_default(),
67 Err(_) => Self::default(),
68 }
69 }
70
71 pub fn save(&self) -> Result<(), String> {
72 let path = Self::path()?;
73 if let Some(parent) = path.parent() {
74 std::fs::create_dir_all(parent).map_err(|e| e.to_string())?;
75 }
76 let text = serde_json::to_string_pretty(self).map_err(|e| e.to_string())?;
77 std::fs::write(&path, text).map_err(|e| e.to_string())
78 }
79}
80
81pub fn parse_catalog(provider_id: &str, value: &serde_json::Value) -> Result<Vec<Model>, String> {
85 let entries: Vec<serde_json::Value> = if let Some(array) = value.as_array() {
86 array.clone()
87 } else if let Some(models) = value.get("models").and_then(|m| m.as_array()) {
88 models.clone()
89 } else if let Some(object) = value.as_object() {
90 object.values().cloned().collect()
91 } else {
92 return Err(format!(
93 "Invalid model catalog for provider \"{provider_id}\""
94 ));
95 };
96
97 let mut out = Vec::with_capacity(entries.len());
98 for entry in entries {
99 if entry.get("id").is_none() {
100 continue;
101 }
102 match serde_json::from_value::<Model>(entry) {
103 Ok(mut model) => {
104 model.provider = rpi_ai::types::ProviderId::from(provider_id);
105 out.push(model);
106 }
107 Err(_) => continue,
108 }
109 }
110 Ok(out)
111}
112
113pub fn remote_models(entry: Option<&CatalogEntry>, local_generated_at: Option<i64>) -> Vec<Model> {
117 let Some(entry) = entry else {
118 return Vec::new();
119 };
120 if let Some(local) = local_generated_at {
121 if entry.last_modified == 0 || entry.last_modified <= local {
122 return Vec::new();
123 }
124 }
125 entry.models.clone()
126}
127
128pub fn merge_models(baseline: &[Model], dynamic: &[Model]) -> Vec<Model> {
130 let mut merged = baseline.to_vec();
131 for model in dynamic {
132 match merged.iter().position(|m| m.id == model.id) {
133 Some(index) => merged[index] = model.clone(),
134 None => merged.push(model.clone()),
135 }
136 }
137 merged
138}
139
140#[derive(Debug, Clone, PartialEq)]
142pub enum RefreshOutcome {
143 Skipped,
145 NotModified,
147 Updated(usize),
149}
150
151pub async fn refresh_provider_catalog(
156 provider_id: &str,
157 base_url: &str,
158 allow_network: bool,
159 force: bool,
160 local_generated_at: Option<i64>,
161 store: &mut CatalogStore,
162) -> Result<Vec<Model>, String> {
163 let stored = store.providers.get(provider_id).cloned();
164
165 if !allow_network {
168 return Ok(remote_models(stored.as_ref(), local_generated_at));
169 }
170 if !force {
171 if let Some(entry) = &stored {
172 let age = now_ms() - entry.checked_at;
173 if entry.checked_at > 0 && age >= 0 && age < REMOTE_CATALOG_REFRESH_INTERVAL_MS {
174 return Ok(remote_models(Some(entry), local_generated_at));
175 }
176 }
177 }
178
179 let url = format!(
180 "{}/api/models/providers/{}",
181 base_url.trim_end_matches('/'),
182 urlencode(provider_id)
183 );
184 let validator = stored
185 .as_ref()
186 .filter(|e| !e.models.is_empty())
187 .and_then(|e| e.etag.clone());
188
189 let client = reqwest::Client::builder()
190 .timeout(std::time::Duration::from_secs(15))
191 .build()
192 .map_err(|e| format!("catalog HTTP client: {e}"))?;
193 let mut request = client
194 .get(&url)
195 .header("accept", "application/json")
196 .header("User-Agent", format!("rpi/{}", crate::VERSION));
197 if let Some(etag) = &validator {
198 request = request.header("if-none-match", etag);
199 }
200
201 let response = match request.send().await {
202 Ok(r) => r,
203 Err(error) => {
204 if let Some(entry) = stored.clone() {
206 store.providers.insert(
207 provider_id.to_string(),
208 CatalogEntry {
209 checked_at: now_ms(),
210 ..entry
211 },
212 );
213 let _ = store.save();
214 }
215 return Err(format!(
216 "Model catalog request failed for {provider_id}: {error}"
217 ));
218 }
219 };
220
221 let status = response.status();
222 let checked_at = now_ms();
223
224 if status.as_u16() == 304 {
225 if let Some(mut entry) = stored.clone() {
226 entry.checked_at = checked_at;
227 store
228 .providers
229 .insert(provider_id.to_string(), entry.clone());
230 let _ = store.save();
231 return Ok(remote_models(Some(&entry), local_generated_at));
232 }
233 }
234
235 if status.as_u16() == 404 || status.as_u16() == 501 {
236 let mut entry = stored.unwrap_or_default();
237 entry.checked_at = checked_at;
238 entry.last_modified = 0;
239 entry.etag = None;
240 entry.models.clear();
241 store.providers.insert(provider_id.to_string(), entry);
242 let _ = store.save();
243 return Ok(Vec::new());
244 }
245
246 if !status.is_success() {
247 if let Some(entry) = stored {
248 store.providers.insert(
249 provider_id.to_string(),
250 CatalogEntry {
251 checked_at,
252 ..entry
253 },
254 );
255 let _ = store.save();
256 }
257 return Err(format!(
258 "Model catalog request failed for {provider_id}: {}",
259 status.as_u16()
260 ));
261 }
262
263 let last_modified = response
264 .headers()
265 .get("last-modified")
266 .and_then(|v| v.to_str().ok())
267 .map(parse_http_date_ms)
268 .unwrap_or(0);
269 let etag = response
270 .headers()
271 .get("etag")
272 .and_then(|v| v.to_str().ok())
273 .map(str::to_string);
274 let body: serde_json::Value = response
275 .json()
276 .await
277 .map_err(|e| format!("invalid catalog JSON for {provider_id}: {e}"))?;
278 let models = parse_catalog(provider_id, &body)?;
279
280 let entry = CatalogEntry {
281 models: models.clone(),
282 checked_at,
283 last_modified,
284 etag,
285 };
286 let overlay = remote_models(Some(&entry), local_generated_at);
287 let count = overlay.len();
288 store.providers.insert(provider_id.to_string(), entry);
289 let _ = store.save();
290 let _ = count;
291 Ok(overlay)
292}
293
294pub async fn merge_into_catalog(
299 baseline: Vec<Model>,
300 allow_network: bool,
301 local_generated_at: Option<i64>,
302) -> Vec<Model> {
303 use std::collections::BTreeSet;
304 let providers: BTreeSet<String> = baseline.iter().map(|m| m.provider.to_string()).collect();
305 if providers.is_empty() {
306 return baseline;
307 }
308 let mut store = CatalogStore::load();
309 let mut result = baseline;
310 for provider_id in providers {
311 let overlay = match refresh_provider_catalog(
312 &provider_id,
313 DEFAULT_CATALOG_BASE_URL,
314 allow_network,
315 false,
316 local_generated_at,
317 &mut store,
318 )
319 .await
320 {
321 Ok(models) => models,
322 Err(_) => Vec::new(),
323 };
324 if !overlay.is_empty() {
325 result = merge_models(&result, &overlay);
326 }
327 }
328 result
329}
330
331fn now_ms() -> i64 {
332 std::time::SystemTime::now()
333 .duration_since(std::time::UNIX_EPOCH)
334 .map(|d| d.as_millis() as i64)
335 .unwrap_or(0)
336}
337
338fn urlencode(value: &str) -> String {
340 let mut out = String::with_capacity(value.len());
341 for byte in value.bytes() {
342 match byte {
343 b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
344 out.push(byte as char)
345 }
346 other => out.push_str(&format!("%{other:02X}")),
347 }
348 }
349 out
350}
351
352fn parse_http_date_ms(value: &str) -> i64 {
355 const MONTHS: [&str; 12] = [
357 "Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec",
358 ];
359 let parts: Vec<&str> = value.split_whitespace().collect();
360 if parts.len() < 5 {
361 return 0;
362 }
363 let day: i64 = match parts[1].parse() {
364 Ok(d) => d,
365 Err(_) => return 0,
366 };
367 let month = match MONTHS.iter().position(|m| *m == parts[2]) {
368 Some(m) => m as i64,
369 None => return 0,
370 };
371 let year: i64 = match parts[3].parse() {
372 Ok(y) => y,
373 Err(_) => return 0,
374 };
375 let time: Vec<&str> = parts[4].split(':').collect();
376 if time.len() != 3 {
377 return 0;
378 }
379 let (h, m, s): (i64, i64, i64) = match (time[0].parse(), time[1].parse(), time[2].parse()) {
380 (Ok(h), Ok(m), Ok(s)) => (h, m, s),
381 _ => return 0,
382 };
383 let month1 = month + 1;
386 let y = if month1 <= 2 { year - 1 } else { year };
387 let era = if y >= 0 { y } else { y - 399 } / 400;
388 let yoe = y - era * 400;
389 let mp = (month1 + 9) % 12;
390 let doy = (153 * mp + 2) / 5 + day - 1;
391 let doe = yoe * 365 + yoe / 4 - yoe / 100 + doy;
392 let days = era * 146097 + doe - 719468;
393 (days * 86_400 + h * 3600 + m * 60 + s) * 1000
394}
395
396#[cfg(test)]
397mod tests {
398 use super::*;
399 use rpi_ai::types::{Api, ModelCost, ProviderId};
400
401 fn model(id: &str, provider: &str) -> Model {
402 Model {
403 id: id.to_string(),
404 name: id.to_string(),
405 api: Api::AnthropicMessages,
406 provider: ProviderId::from(provider),
407 base_url: String::new(),
408 reasoning: false,
409 thinking_level_map: None,
410 input: vec![],
411 cost: ModelCost::default(),
412 context_window: 0,
413 max_tokens: 0,
414 sampling_params: None,
415 headers: None,
416 compat: None,
417 }
418 }
419
420 #[test]
421 fn parse_array_and_object() {
422 let array = serde_json::json!([{ "id": "m1", "name": "M1", "api": "anthropic-messages",
423 "provider": "x", "baseUrl": "", "reasoning": false, "input": [],
424 "cost": { "input": 0.0, "output": 0.0, "cacheRead": 0.0, "cacheWrite": 0.0 },
425 "contextWindow": 0, "maxTokens": 0 }]);
426 let parsed = parse_catalog("anthropic", &array).unwrap();
427 assert_eq!(parsed.len(), 1);
428 assert_eq!(parsed[0].provider.to_string(), "anthropic");
429
430 let wrapped = serde_json::json!({ "models": array });
431 assert_eq!(parse_catalog("anthropic", &wrapped).unwrap().len(), 1);
432 }
433
434 #[test]
435 fn merge_overlay_wins_by_id() {
436 let base = vec![model("a", "p"), model("b", "p")];
437 let overlay = vec![model("b", "p"), model("c", "p")];
438 let merged = merge_models(&base, &overlay);
439 assert_eq!(merged.len(), 3);
440 assert_eq!(
441 merged.iter().map(|m| m.id.as_str()).collect::<Vec<_>>(),
442 vec!["a", "b", "c"]
443 );
444 }
445
446 #[test]
447 fn remote_models_respects_local_generation() {
448 let entry = CatalogEntry {
449 models: vec![model("a", "p")],
450 checked_at: 0,
451 last_modified: 1_000,
452 etag: None,
453 };
454 assert!(remote_models(Some(&entry), Some(2_000)).is_empty());
455 assert_eq!(remote_models(Some(&entry), Some(500)).len(), 1);
456 assert_eq!(remote_models(Some(&entry), None).len(), 1);
457 assert!(remote_models(None, None).is_empty());
458 }
459
460 #[test]
461 fn http_date_parses_imf_fixdate() {
462 assert_eq!(
464 parse_http_date_ms("Sun, 06 Nov 1994 08:49:37 GMT"),
465 784_111_777_000
466 );
467 assert_eq!(parse_http_date_ms("not a date"), 0);
468 }
469}