feat: implement live model registry cache and background sync

src/lib.rs

@@ -3,7 +3,9 @@
 pub mod budget;
 pub mod client;
 pub mod persona;
+pub mod registry;
 
 pub use budget::{BudgetExceeded, BudgetStatus, BudgetTracker, CreditTier};
 pub use client::{ClientError, ModelClient, OpenRouterClient};
 pub use persona::{DispatchResult, ModelAttempt, ModelCandidate, Persona, PersonaRegistry, Tier};
+pub use registry::{ModelInfo, ModelPricing, ModelRegistry};

src/registry.rs

@@ -0,0 +1,223 @@
+use crate::persona::{ModelCandidate, Tier};
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+use tokio::sync::RwLock;
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ModelPricing {
+    #[serde(default)]
+    pub prompt: String,
+    #[serde(default)]
+    pub completion: String,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ModelInfo {
+    pub id: String,
+    #[serde(default)]
+    pub name: Option<String>,
+    #[serde(default)]
+    pub context_length: Option<u64>,
+    #[serde(default)]
+    pub pricing: Option<ModelPricing>,
+}
+
+#[derive(Debug, Deserialize)]
+struct OpenRouterModelsResponse {
+    data: Vec<ModelInfo>,
+}
+
+#[derive(Clone)]
+pub struct ModelRegistry {
+    models: Arc<RwLock<HashMap<String, ModelInfo>>>,
+    last_refreshed: Arc<RwLock<Option<Instant>>>,
+}
+
+impl Default for ModelRegistry {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl ModelRegistry {
+    pub fn new() -> Self {
+        Self {
+            models: Arc::new(RwLock::new(HashMap::new())),
+            last_refreshed: Arc::new(RwLock::new(None)),
+        }
+    }
+
+    pub async fn model_count(&self) -> usize {
+        let read = self.models.read().await;
+        read.len()
+    }
+
+    pub async fn last_refreshed_at(&self) -> Option<Instant> {
+        let read = self.last_refreshed.read().await;
+        *read
+    }
+
+    pub async fn update_from_json(&self, json_str: &str) -> Result<usize, String> {
+        let resp: OpenRouterModelsResponse = serde_json::from_str(json_str)
+            .map_err(|e| format!("Failed to parse OpenRouter models response: {e}"))?;
+
+        let count = resp.data.len();
+        let mut map = HashMap::with_capacity(count);
+        for m in resp.data {
+            map.insert(m.id.clone(), m);
+        }
+
+        {
+            let mut write = self.models.write().await;
+            *write = map;
+        }
+        {
+            let mut time_write = self.last_refreshed.write().await;
+            *time_write = Some(Instant::now());
+        }
+
+        Ok(count)
+    }
+
+    /// Fetches live models catalog from OpenRouter `/v1/models`
+    pub async fn fetch_live(&self, base_url: &str) -> Result<usize, String> {
+        let endpoint = format!("{}/models", base_url.trim_end_matches('/'));
+        let client = reqwest::Client::builder()
+            .timeout(Duration::from_secs(15))
+            .build()
+            .map_err(|e| e.to_string())?;
+
+        let res = client
+            .get(&endpoint)
+            .header("HTTP-Referer", "https://github.com/brokecrew/brokecrew")
+            .header("X-Title", "BrokeCrew")
+            .send()
+            .await
+            .map_err(|e| format!("Network error fetching models catalog: {e}"))?;
+
+        if !res.status().is_success() {
+            return Err(format!("OpenRouter models endpoint returned HTTP {}", res.status()));
+        }
+
+        let body = res
+            .text()
+            .await
+            .map_err(|e| format!("Failed to read models response body: {e}"))?;
+
+        self.update_from_json(&body).await
+    }
+
+    /// Starts a background task refreshing the models registry periodically
+    pub fn start_background_sync(self: Arc<Self>, base_url: String, interval: Duration) {
+        tokio::spawn(async move {
+            let mut ticker = tokio::time::interval(interval);
+            loop {
+                ticker.tick().await;
+                if let Err(e) = self.fetch_live(&base_url).await {
+                    tracing::warn!("Periodic models registry sync failed: {}", e);
+                } else {
+                    tracing::debug!("Periodic models registry sync succeeded");
+                }
+            }
+        });
+    }
+
+    pub async fn is_model_available(&self, model_id: &str) -> bool {
+        let read = self.models.read().await;
+        read.contains_key(model_id)
+    }
+
+    /// Filters and orders candidates according to live availability, stealth permissions, and escalation.
+    /// If the live registry is empty (e.g. offline mode / uninitialized), candidate presence checks are relaxed.
+    pub async fn filter_eligible<'a>(
+        &self,
+        candidates: &'a [ModelCandidate],
+        allow_stealth: bool,
+        require_paid: bool,
+    ) -> Vec<&'a ModelCandidate> {
+        let read = self.models.read().await;
+        let is_empty = read.is_empty();
+
+        candidates
+            .iter()
+            .filter(|c| {
+                if !allow_stealth && c.is_stealth {
+                    return false;
+                }
+                if require_paid && c.tier != Tier::Paid {
+                    return false;
+                }
+                // Skip unlisted models if we have a populated registry
+                if !is_empty && !read.contains_key(&c.id) {
+                    return false;
+                }
+                true
+            })
+            .collect()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_registry_update_and_filtering() {
+        let registry = ModelRegistry::new();
+
+        let json_data = r#"{
+            "data": [
+                {
+                    "id": "qwen/qwen-2.5-coder-32b-instruct:free",
+                    "name": "Qwen 2.5 Coder 32B",
+                    "context_length": 32768,
+                    "pricing": { "prompt": "0", "completion": "0" }
+                },
+                {
+                    "id": "google/gemini-2.0-flash-001",
+                    "name": "Gemini 2.0 Flash",
+                    "context_length": 1048576,
+                    "pricing": { "prompt": "0.0000001", "completion": "0.0000004" }
+                }
+            ]
+        }"#;
+
+        let count = registry.update_from_json(json_data).await.expect("update failed");
+        assert_eq!(count, 2);
+
+        assert!(registry.is_model_available("qwen/qwen-2.5-coder-32b-instruct:free").await);
+        assert!(!registry.is_model_available("unlisted/model").await);
+
+        let candidates = vec![
+            ModelCandidate {
+                id: "stealth/experiment".into(),
+                tier: Tier::Free,
+                is_stealth: true,
+            },
+            ModelCandidate {
+                id: "qwen/qwen-2.5-coder-32b-instruct:free".into(),
+                tier: Tier::Free,
+                is_stealth: false,
+            },
+            ModelCandidate {
+                id: "google/gemini-2.0-flash-001".into(),
+                tier: Tier::Paid,
+                is_stealth: false,
+            },
+            ModelCandidate {
+                id: "unlisted/dropped-model:free".into(),
+                tier: Tier::Free,
+                is_stealth: false,
+            },
+        ];
+
+        // Eligible without stealth, free-tier
+        let filtered = registry.filter_eligible(&candidates, false, false).await;
+        assert_eq!(filtered.len(), 2);
+        assert_eq!(filtered[0].id, "qwen/qwen-2.5-coder-32b-instruct:free");
+        assert_eq!(filtered[1].id, "google/gemini-2.0-flash-001");
+        // "unlisted/dropped-model:free" is excluded
+    }
+}