use std::collections::HashMap;
use std::sync::RwLock;
use std::time::Instant;
#[derive(Debug, Clone)]
struct AgentEntry {
id: String,
last_used: Option<Instant>,
}
#[derive(Debug, Default)]
pub struct AgentCapabilityRegistry {
by_capability: RwLock<HashMap<String, Vec<AgentEntry>>>,
}
impl AgentCapabilityRegistry {
pub fn new() -> Self {
Self::default()
}
pub fn register(&self, capability: impl Into<String>, agent_id: impl Into<String>) {
let capability = capability.into();
let agent_id = agent_id.into();
let mut g = self.by_capability.write().expect("registry poisoned");
let entries = g.entry(capability).or_default();
if !entries.iter().any(|e| e.id == agent_id) {
entries.push(AgentEntry {
id: agent_id,
last_used: None,
});
}
}
pub fn unregister(&self, capability: &str, agent_id: &str) {
let mut g = self.by_capability.write().expect("registry poisoned");
if let Some(entries) = g.get_mut(capability) {
entries.retain(|e| e.id != agent_id);
if entries.is_empty() {
g.remove(capability);
}
}
}
pub fn agents_for(&self, capability: &str) -> Vec<String> {
let g = self.by_capability.read().expect("registry poisoned");
g.get(capability)
.map(|v| v.iter().map(|e| e.id.clone()).collect())
.unwrap_or_default()
}
pub fn all_agents(&self) -> Vec<String> {
let g = self.by_capability.read().expect("registry poisoned");
let mut seen: std::collections::BTreeSet<String> = std::collections::BTreeSet::new();
for entries in g.values() {
for e in entries {
seen.insert(e.id.clone());
}
}
seen.into_iter().collect()
}
pub fn select(&self, capability: &str, hint: Option<&str>) -> Option<String> {
let g = self.by_capability.read().expect("registry poisoned");
let entries = g.get(capability)?;
if entries.is_empty() {
return None;
}
if let Some(h) = hint {
if let Some(e) = entries.iter().find(|e| e.id == h) {
return Some(e.id.clone());
}
}
let mut best: Option<&AgentEntry> = None;
for e in entries.iter() {
best = match best {
None => Some(e),
Some(b) if e.last_used > b.last_used => Some(e),
Some(b) => Some(b),
};
}
best.map(|e| e.id.clone())
}
pub fn note_used(&self, capability: &str, agent_id: &str) {
let mut g = self.by_capability.write().expect("registry poisoned");
if let Some(entries) = g.get_mut(capability) {
if let Some(e) = entries.iter_mut().find(|e| e.id == agent_id) {
e.last_used = Some(Instant::now());
}
}
}
pub fn capability_count(&self) -> usize {
self.by_capability.read().expect("registry poisoned").len()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn register_and_list() {
let r = AgentCapabilityRegistry::new();
r.register("summarize", "alpha");
r.register("summarize", "beta");
r.register("transcribe-audio", "alpha");
assert_eq!(r.agents_for("summarize"), vec!["alpha", "beta"]);
assert_eq!(r.agents_for("transcribe-audio"), vec!["alpha"]);
assert!(r.agents_for("nonexistent").is_empty());
assert_eq!(r.all_agents(), vec!["alpha", "beta"]);
assert_eq!(r.capability_count(), 2);
}
#[test]
fn register_is_idempotent() {
let r = AgentCapabilityRegistry::new();
r.register("summarize", "alpha");
r.register("summarize", "alpha");
r.register("summarize", "alpha");
assert_eq!(r.agents_for("summarize"), vec!["alpha"]);
}
#[test]
fn select_honors_hint() {
let r = AgentCapabilityRegistry::new();
r.register("summarize", "alpha");
r.register("summarize", "beta");
r.register("summarize", "gamma");
assert_eq!(r.select("summarize", Some("beta")), Some("beta".into()));
assert_eq!(
r.select("summarize", Some("nonexistent")),
Some("alpha".into())
);
}
#[test]
fn select_uses_mru_when_no_hint() {
let r = AgentCapabilityRegistry::new();
r.register("summarize", "alpha");
r.register("summarize", "beta");
assert_eq!(r.select("summarize", None), Some("alpha".into()));
r.note_used("summarize", "beta");
assert_eq!(r.select("summarize", None), Some("beta".into()));
std::thread::sleep(std::time::Duration::from_millis(2));
r.note_used("summarize", "alpha");
assert_eq!(r.select("summarize", None), Some("alpha".into()));
}
#[test]
fn unregister_removes() {
let r = AgentCapabilityRegistry::new();
r.register("summarize", "alpha");
r.register("summarize", "beta");
r.unregister("summarize", "alpha");
assert_eq!(r.agents_for("summarize"), vec!["beta"]);
r.unregister("summarize", "beta");
assert!(r.agents_for("summarize").is_empty());
assert_eq!(r.capability_count(), 0);
}
#[test]
fn select_returns_none_for_unknown_capability() {
let r = AgentCapabilityRegistry::new();
assert_eq!(r.select("nope", None), None);
}
#[test]
fn select_three_agent_tiebreak() {
let r = AgentCapabilityRegistry::new();
r.register("summarize", "alpha");
r.register("summarize", "beta");
r.register("summarize", "gamma");
assert_eq!(r.select("summarize", None), Some("alpha".into()));
r.note_used("summarize", "gamma");
assert_eq!(r.select("summarize", None), Some("gamma".into()));
std::thread::sleep(std::time::Duration::from_millis(2));
r.note_used("summarize", "alpha");
assert_eq!(r.select("summarize", None), Some("alpha".into()));
assert_eq!(r.select("summarize", Some("beta")), Some("beta".into()));
}
}