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

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

use crank_core::{CacheScope, CoordinationStateStore, CoordinationStateValue};
use crank_registry::{PostgresRegistry, PublishedAgentCatalog, PublishedAgentTool, RegistryError};
use crank_trace::{DbOperation, ErrorCategory, Stage, StageOutcome};
use serde::{Deserialize, Serialize};
use tokio::sync::{Mutex, RwLock};
use tracing::{Instrument, info, warn};

use crate::manifest::analyze_published_tool_catalog;

const MAX_LOCAL_CATALOGS: usize = 1_024;

#[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, Weak<Mutex<()>>>>>,
}

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

struct CachedCatalog {
    loaded_at: Option<Instant>,
    catalog: PublishedAgentCatalog,
    metrics: CatalogMetrics,
}

#[derive(Clone, Copy, Debug, Default)]
struct CatalogMetrics {
    tool_count: usize,
    estimated_context_tokens: usize,
    warning_count: usize,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
struct CatalogSnapshot {
    catalog: PublishedAgentCatalog,
    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> {
        Ok(self.get_catalog(workspace_slug, agent_slug).await?.tools)
    }

    pub async fn get_catalog(
        &self,
        workspace_slug: &str,
        agent_slug: &str,
    ) -> Result<PublishedAgentCatalog, RegistryError> {
        let span = Stage::McpCatalogLoad.span();
        let result = async {
            self.refresh_if_stale(workspace_slug, agent_slug).await?;
            let guard = self.cached.read().await;
            guard
                .get(&CatalogKey::new(workspace_slug, agent_slug))
                .map(|entry| entry.catalog.clone())
                .ok_or_else(|| RegistryError::PublishedAgentNotFound {
                    workspace_slug: workspace_slug.to_owned(),
                    agent_slug: agent_slug.to_owned(),
                })
        }
        .instrument(span.clone())
        .await;
        match &result {
            Ok(_) => StageOutcome::Success.record(&span),
            Err(_) => {
                StageOutcome::Error.record(&span);
                ErrorCategory::Catalog.record(&span);
            }
        }
        drop(span);
        result
    }

    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;
            locks.retain(|_, lock| lock.strong_count() > 0);
            if let Some(lock) = locks.get(&key).and_then(Weak::upgrade) {
                lock
            } else {
                let lock = Arc::new(Mutex::new(()));
                locks.insert(key.clone(), Arc::downgrade(&lock));
                lock
            }
        };
        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((catalog, age)) = self.load_shared_snapshot(workspace_slug, agent_slug).await {
            let metrics =
                log_catalog_analysis(workspace_slug, agent_slug, "shared_cache", &catalog.tools);
            self.store_local_catalog(
                key,
                CachedCatalog {
                    loaded_at: Instant::now().checked_sub(age),
                    catalog,
                    metrics,
                },
            )
            .await;
            return Ok(());
        }

        let db_span = DbOperation::CatalogLoad.span();
        let catalog_result = self
            .registry
            .get_published_agent_catalog_by_slug(workspace_slug, agent_slug)
            .instrument(db_span.clone())
            .await;
        let catalog = match catalog_result {
            Ok(catalog) => catalog,
            Err(error) => {
                StageOutcome::Error.record(&db_span);
                ErrorCategory::Database.record(&db_span);
                drop(db_span);
                return Err(error);
            }
        };
        StageOutcome::Success.record(&db_span);
        drop(db_span);
        let metrics = log_catalog_analysis(workspace_slug, agent_slug, "postgres", &catalog.tools);
        self.store_shared_snapshot(workspace_slug, agent_slug, &catalog)
            .await;
        let published_tool_count = catalog.tools.len();
        let previous_count = self
            .store_local_catalog(
                key,
                CachedCatalog {
                    loaded_at: Some(Instant::now()),
                    catalog,
                    metrics,
                },
            )
            .await;

        info!(
            name: "mcp.catalog.refreshed",
            workspace_slug,
            agent_slug,
            published_tool_count,
            previous_published_tool_count = previous_count,
            "published agent catalog refreshed"
        );

        Ok(())
    }

    async fn store_local_catalog(&self, key: CatalogKey, entry: CachedCatalog) -> usize {
        let mut guard = self.cached.write().await;
        let previous_count = guard
            .get(&key)
            .map(|current| current.catalog.tools.len())
            .unwrap_or_default();

        if guard.len() >= MAX_LOCAL_CATALOGS && !guard.contains_key(&key) {
            let oldest = guard
                .iter()
                .min_by_key(|(_, current)| current.loaded_at)
                .map(|(candidate, _)| candidate.clone());
            if let Some(oldest) = oldest {
                guard.remove(&oldest);
            }
        }
        guard.insert(key, entry);
        record_catalog_metrics(guard.values().map(|entry| entry.metrics));
        previous_count
    }

    async fn load_shared_snapshot(
        &self,
        workspace_slug: &str,
        agent_slug: &str,
    ) -> Option<(PublishedAgentCatalog, 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.catalog, age))
    }

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

        let payload = match serde_json::to_value(CatalogSnapshot {
            catalog: catalog.clone(),
            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 log_catalog_analysis(
    workspace_slug: &str,
    agent_slug: &str,
    source: &str,
    tools: &[PublishedAgentTool],
) -> CatalogMetrics {
    let analysis = match analyze_published_tool_catalog(tools) {
        Ok(analysis) => analysis,
        Err(_) => {
            warn!(
                name: "mcp.catalog.analysis_failed",
                workspace_slug,
                agent_slug,
                source,
                error_category = "catalog_validation",
                "published catalog analysis failed"
            );
            return CatalogMetrics::default();
        }
    };
    let warning_count = analysis
        .quality
        .findings
        .iter()
        .filter(|finding| finding.severity == crank_core::ToolQualitySeverity::Warning)
        .count();

    info!(
        name: "mcp.catalog.analyzed",
        workspace_slug,
        agent_slug,
        source,
        tool_count = analysis.budget.tool_count,
        serialized_bytes = analysis.budget.serialized_bytes,
        estimated_context_tokens = analysis.budget.estimated_context_tokens,
        largest_tool_estimated_context_tokens =
            analysis.budget.largest_tool_estimated_context_tokens,
        recommended_context_tokens = analysis.budget.recommended_context_tokens,
        exceeds_recommended_budget = analysis.budget.exceeds_recommended_budget,
        catalog_quality_warning_count = warning_count,
        "published agent catalog analyzed"
    );

    CatalogMetrics {
        tool_count: analysis.budget.tool_count,
        estimated_context_tokens: analysis.budget.estimated_context_tokens,
        warning_count,
    }
}

fn record_catalog_metrics(metrics: impl Iterator<Item = CatalogMetrics>) {
    let aggregate = metrics.fold(CatalogMetrics::default(), |mut aggregate, current| {
        aggregate.tool_count = aggregate.tool_count.saturating_add(current.tool_count);
        aggregate.estimated_context_tokens = aggregate
            .estimated_context_tokens
            .saturating_add(current.estimated_context_tokens);
        aggregate.warning_count = aggregate
            .warning_count
            .saturating_add(current.warning_count);
        aggregate
    });

    metrics::gauge!("crank_catalog_tools").set(aggregate.tool_count as f64);
    metrics::gauge!("crank_catalog_estimated_context_tokens")
        .set(aggregate.estimated_context_tokens as f64);
    metrics::gauge!("crank_catalog_warnings").set(aggregate.warning_count as f64);
}

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))
}