crank/apps/mcp-server/src/catalog.rs

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<RwLock<CachedCatalog>>,
}

#[derive(Default)]
struct CachedCatalog {
    loaded_at: Option<Instant>,
    tools: Vec<RegistryOperation>,
    tools_by_name: HashMap<String, RegistryOperation>,
}

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<Vec<RegistryOperation>, 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<Option<RegistryOperation>, 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::<HashMap<_, _>>();
        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(())
    }
}