use std::{ collections::HashMap, sync::Arc, time::{Duration, Instant}, }; use crank_registry::{PostgresRegistry, RegistryError, RegistryOperation}; use tokio::sync::RwLock; use tracing::info; #[derive(Clone)] pub struct PublishedToolCatalog { registry: PostgresRegistry, refresh_interval: Duration, cached: Arc>, } #[derive(Default)] struct CachedCatalog { loaded_at: Option, tools: Vec, tools_by_name: HashMap, } impl PublishedToolCatalog { pub fn new(registry: PostgresRegistry, refresh_interval: Duration) -> Self { Self { registry, refresh_interval, cached: Arc::new(RwLock::new(CachedCatalog::default())), } } pub async fn list_tools(&self) -> Result, RegistryError> { self.refresh_if_stale().await?; let guard = self.cached.read().await; Ok(guard.tools.clone()) } pub async fn get_tool( &self, tool_name: &str, ) -> Result, RegistryError> { self.refresh_if_stale().await?; let guard = self.cached.read().await; Ok(guard.tools_by_name.get(tool_name).cloned()) } async fn refresh_if_stale(&self) -> Result<(), RegistryError> { let should_refresh = { let guard = self.cached.read().await; match guard.loaded_at { Some(loaded_at) => loaded_at.elapsed() >= self.refresh_interval, None => true, } }; if !should_refresh { return Ok(()); } let tools = self.registry.list_published_operations().await?; let tools_by_name = tools .iter() .cloned() .map(|operation| (operation.name.clone(), operation)) .collect::>(); let mut guard = self.cached.write().await; let previous_count = guard.tools.len(); guard.loaded_at = Some(Instant::now()); guard.tools = tools; guard.tools_by_name = tools_by_name; info!( published_tool_count = guard.tools.len(), previous_published_tool_count = previous_count, "published tool catalog refreshed" ); Ok(()) } }