crank/crates/crank-community-mcp/src/catalog.rs

use std::{
    collections::HashMap,
    sync::Arc,
    time::{Duration, Instant, SystemTime, UNIX_EPOCH},
};

use crank_core::{CacheScope, CoordinationStateStore, CoordinationStateValue};
use crank_registry::{PostgresRegistry, PublishedAgentTool, RegistryError};
use serde::{Deserialize, Serialize};
use tokio::sync::{Mutex, RwLock};
use tracing::info;

#[derive(Clone)]
pub struct PublishedToolCatalog {
    registry: PostgresRegistry,
    refresh_interval: Duration,
    coordination_store: Arc<dyn CoordinationStateStore>,
    cached: Arc<RwLock<HashMap<CatalogKey, CachedCatalog>>>,
    refresh_locks: Arc<Mutex<HashMap<CatalogKey, Arc<Mutex<()>>>>>,
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
struct CatalogKey {
    workspace_slug: String,
    agent_slug: String,
}

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

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
struct CatalogSnapshot {
    tools: Vec<PublishedAgentTool>,
    generated_at_ms: u64,
}

impl PublishedToolCatalog {
    pub fn new(
        registry: PostgresRegistry,
        refresh_interval: Duration,
        coordination_store: Arc<dyn CoordinationStateStore>,
    ) -> Self {
        Self {
            registry,
            refresh_interval,
            coordination_store,
            cached: Arc::new(RwLock::new(HashMap::new())),
            refresh_locks: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    pub async fn list_tools(
        &self,
        workspace_slug: &str,
        agent_slug: &str,
    ) -> Result<Vec<PublishedAgentTool>, RegistryError> {
        self.refresh_if_stale(workspace_slug, agent_slug).await?;
        let guard = self.cached.read().await;
        Ok(guard
            .get(&CatalogKey::new(workspace_slug, agent_slug))
            .map(|entry| entry.tools.clone())
            .unwrap_or_default())
    }

    async fn refresh_if_stale(
        &self,
        workspace_slug: &str,
        agent_slug: &str,
    ) -> Result<(), RegistryError> {
        let key = CatalogKey::new(workspace_slug, agent_slug);
        let should_refresh = {
            let guard = self.cached.read().await;

            match guard.get(&key).and_then(|entry| entry.loaded_at) {
                Some(loaded_at) => loaded_at.elapsed() >= self.refresh_interval,
                None => true,
            }
        };

        if !should_refresh {
            return Ok(());
        }

        let refresh_lock = {
            let mut locks = self.refresh_locks.lock().await;
            Arc::clone(
                locks
                    .entry(key.clone())
                    .or_insert_with(|| Arc::new(Mutex::new(()))),
            )
        };
        let _refresh_guard = refresh_lock.lock().await;
        let still_stale = {
            let guard = self.cached.read().await;
            match guard.get(&key).and_then(|entry| entry.loaded_at) {
                Some(loaded_at) => loaded_at.elapsed() >= self.refresh_interval,
                None => true,
            }
        };
        if !still_stale {
            return Ok(());
        }

        if let Some((tools, age)) = self.load_shared_snapshot(workspace_slug, agent_slug).await {
            let mut guard = self.cached.write().await;
            guard.insert(
                key,
                CachedCatalog {
                    loaded_at: Instant::now().checked_sub(age),
                    tools,
                },
            );
            return Ok(());
        }

        let tools = match self
            .registry
            .get_published_agent_tools_by_slug(workspace_slug, agent_slug)
            .await
        {
            Ok(tools) => tools,
            Err(RegistryError::PublishedAgentNotFound { .. }) => Vec::new(),
            Err(error) => return Err(error),
        };
        self.store_shared_snapshot(workspace_slug, agent_slug, &tools)
            .await;
        let mut guard = self.cached.write().await;
        let previous_count = guard
            .get(&key)
            .map(|entry| entry.tools.len())
            .unwrap_or_default();

        guard.insert(
            key,
            CachedCatalog {
                loaded_at: Some(Instant::now()),
                tools,
            },
        );

        info!(
            workspace_slug,
            agent_slug,
            published_tool_count = guard
                .get(&CatalogKey::new(workspace_slug, agent_slug))
                .map(|entry| entry.tools.len())
                .unwrap_or_default(),
            previous_published_tool_count = previous_count,
            "published agent catalog refreshed"
        );

        Ok(())
    }

    async fn load_shared_snapshot(
        &self,
        workspace_slug: &str,
        agent_slug: &str,
    ) -> Option<(Vec<PublishedAgentTool>, Duration)> {
        if self.refresh_interval.is_zero() {
            return None;
        }

        let key = catalog_snapshot_key(workspace_slug, agent_slug);
        let value = match self
            .coordination_store
            .get_value(CacheScope::Coordination, &key)
            .await
        {
            Ok(value) => value?,
            Err(_) => return None,
        };
        let snapshot = serde_json::from_value::<CatalogSnapshot>(value.payload).ok()?;
        let age = Duration::from_millis(now_unix_ms().saturating_sub(snapshot.generated_at_ms));
        (age < self.refresh_interval).then_some((snapshot.tools, age))
    }

    async fn store_shared_snapshot(
        &self,
        workspace_slug: &str,
        agent_slug: &str,
        tools: &[PublishedAgentTool],
    ) {
        let Some(ttl) = catalog_snapshot_ttl(self.refresh_interval) else {
            return;
        };

        let payload = match serde_json::to_value(CatalogSnapshot {
            tools: tools.to_vec(),
            generated_at_ms: now_unix_ms(),
        }) {
            Ok(payload) => payload,
            Err(_) => return,
        };
        let key = catalog_snapshot_key(workspace_slug, agent_slug);
        let _ = self
            .coordination_store
            .put_value(
                CacheScope::Coordination,
                &key,
                CoordinationStateValue { payload },
                ttl,
            )
            .await;
    }
}

fn now_unix_ms() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|duration| u64::try_from(duration.as_millis()).unwrap_or(u64::MAX))
        .unwrap_or_default()
}

impl CatalogKey {
    fn new(workspace_slug: &str, agent_slug: &str) -> Self {
        Self {
            workspace_slug: workspace_slug.to_owned(),
            agent_slug: agent_slug.to_owned(),
        }
    }
}

fn catalog_snapshot_key(workspace_slug: &str, agent_slug: &str) -> String {
    format!("published_catalog:{workspace_slug}:{agent_slug}")
}

fn catalog_snapshot_ttl(refresh_interval: Duration) -> Option<Duration> {
    if refresh_interval.is_zero() {
        return None;
    }

    refresh_interval.checked_mul(2).or(Some(refresh_interval))
}