cache: share published catalogs across instances

apps/mcp-server/src/catalog.rs

@@ -4,7 +4,9 @@ use std::{
     time::{Duration, Instant},
 };
 
+use crank_core::{CacheScope, CoordinationStateStore, CoordinationStateValue};
 use crank_registry::{PostgresRegistry, PublishedAgentTool, RegistryError};
+use serde::{Deserialize, Serialize};
 use tokio::sync::RwLock;
 use tracing::info;
 
@@ -12,6 +14,7 @@ use tracing::info;
 pub struct PublishedToolCatalog {
     registry: PostgresRegistry,
     refresh_interval: Duration,
+    coordination_store: Arc<dyn CoordinationStateStore>,
     cached: Arc<RwLock<HashMap<CatalogKey, CachedCatalog>>>,
 }
 
@@ -27,11 +30,21 @@ struct CachedCatalog {
     tools: Vec<PublishedAgentTool>,
 }
 
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+struct CatalogSnapshot {
+    tools: Vec<PublishedAgentTool>,
+}
+
 impl PublishedToolCatalog {
-    pub fn new(registry: PostgresRegistry, refresh_interval: Duration) -> Self {
+    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())),
         }
     }
@@ -68,6 +81,18 @@ impl PublishedToolCatalog {
             return Ok(());
         }
 
+        if let Some(tools) = self.load_shared_snapshot(workspace_slug, agent_slug).await {
+            let mut guard = self.cached.write().await;
+            guard.insert(
+                key,
+                CachedCatalog {
+                    loaded_at: Some(Instant::now()),
+                    tools,
+                },
+            );
+            return Ok(());
+        }
+
         let tools = match self
             .registry
             .get_published_agent_tools_by_slug(workspace_slug, agent_slug)
@@ -77,6 +102,8 @@ impl PublishedToolCatalog {
             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)
@@ -104,6 +131,57 @@ impl PublishedToolCatalog {
 
         Ok(())
     }
+
+    async fn load_shared_snapshot(
+        &self,
+        workspace_slug: &str,
+        agent_slug: &str,
+    ) -> Option<Vec<PublishedAgentTool>> {
+        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,
+        };
+        serde_json::from_value::<CatalogSnapshot>(value.payload)
+            .ok()
+            .map(|snapshot| snapshot.tools)
+    }
+
+    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(),
+        }) {
+            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;
+    }
 }
 
 impl CatalogKey {
@@ -114,3 +192,15 @@ impl CatalogKey {
         }
     }
 }
+
+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))
+}