a.tolmachev
117 lines
3.1 KiB
Rust
117 lines
3.1 KiB
Rust
use std::{
|
|
collections::HashMap,
|
|
sync::Arc,
|
|
time::{Duration, Instant},
|
|
};
|
|
|
|
use crank_registry::{PostgresRegistry, PublishedAgentTool, RegistryError};
|
|
use tokio::sync::RwLock;
|
|
use tracing::info;
|
|
|
|
#[derive(Clone)]
|
|
pub struct PublishedToolCatalog {
|
|
registry: PostgresRegistry,
|
|
refresh_interval: Duration,
|
|
cached: Arc<RwLock<HashMap<CatalogKey, CachedCatalog>>>,
|
|
}
|
|
|
|
#[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>,
|
|
}
|
|
|
|
impl PublishedToolCatalog {
|
|
pub fn new(registry: PostgresRegistry, refresh_interval: Duration) -> Self {
|
|
Self {
|
|
registry,
|
|
refresh_interval,
|
|
cached: Arc::new(RwLock::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 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),
|
|
};
|
|
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(())
|
|
}
|
|
}
|
|
|
|
impl CatalogKey {
|
|
fn new(workspace_slug: &str, agent_slug: &str) -> Self {
|
|
Self {
|
|
workspace_slug: workspace_slug.to_owned(),
|
|
agent_slug: agent_slug.to_owned(),
|
|
}
|
|
}
|
|
}
|