use dashmap::DashMap;
use std::sync::Arc;
use crate::error::{Error, Result};
use tap_agent::TapAgent;
#[derive(Debug)]
pub struct AgentRegistry {
max_agents: Option<usize>,
agents: DashMap<String, Arc<TapAgent>>,
}
impl AgentRegistry {
pub fn new(max_agents: Option<usize>) -> Self {
Self {
max_agents,
agents: DashMap::new(),
}
}
pub fn agent_count(&self) -> usize {
self.agents.len()
}
pub fn has_agent(&self, did: &str) -> bool {
self.agents.contains_key(did)
}
pub async fn get_agent(&self, did: &str) -> Result<Arc<TapAgent>> {
self.agents
.get(did)
.map(|agent| agent.clone())
.ok_or_else(|| Error::AgentNotFound(did.to_string()))
}
pub async fn register_agent(&self, did: String, agent: Arc<TapAgent>) -> Result<()> {
if let Some(max) = self.max_agents {
if self.agent_count() >= max {
return Err(Error::AgentRegistration(format!(
"Maximum number of agents ({}) reached",
max
)));
}
}
if self.has_agent(&did) {
return Err(Error::AgentRegistration(format!(
"Agent with DID {} is already registered",
did
)));
}
self.agents.insert(did, agent);
Ok(())
}
pub async fn unregister_agent(&self, did: &str) -> Result<()> {
if !self.has_agent(did) {
return Err(Error::AgentNotFound(did.to_string()));
}
self.agents.remove(did);
Ok(())
}
pub fn get_all_dids(&self) -> Vec<String> {
self.agents
.iter()
.map(|entry| entry.key().clone())
.collect()
}
}