vtcode_acp_client/
discovery.rs

1//! Agent discovery and registry functionality
2
3use crate::error::{AcpError, AcpResult};
4use serde::{Deserialize, Serialize};
5use serde_json::Value;
6use std::collections::HashMap;
7use std::sync::Arc;
8use tokio::sync::RwLock;
9
10/// Information about a registered agent
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct AgentInfo {
13    /// Unique agent identifier
14    pub id: String,
15
16    /// Agent display name
17    pub name: String,
18
19    /// Base URL for agent communication
20    pub base_url: String,
21
22    /// Agent description
23    pub description: Option<String>,
24
25    /// Supported actions/tools
26    pub capabilities: Vec<String>,
27
28    /// Agent metadata (version, tags, etc.)
29    #[serde(default)]
30    pub metadata: HashMap<String, Value>,
31
32    /// Whether agent is currently online
33    #[serde(default = "default_online")]
34    pub online: bool,
35
36    /// Last heartbeat/update timestamp
37    pub last_seen: Option<String>,
38}
39
40fn default_online() -> bool {
41    true
42}
43
44/// Agent registry for discovery and lookup
45#[derive(Clone)]
46pub struct AgentRegistry {
47    agents: Arc<RwLock<HashMap<String, AgentInfo>>>,
48}
49
50impl AgentRegistry {
51    /// Create a new agent registry
52    pub fn new() -> Self {
53        Self {
54            agents: Arc::new(RwLock::new(HashMap::new())),
55        }
56    }
57
58    /// Register an agent
59    pub async fn register(&self, agent: AgentInfo) -> AcpResult<()> {
60        let mut agents = self.agents.write().await;
61        agents.insert(agent.id.clone(), agent);
62        Ok(())
63    }
64
65    /// Unregister an agent
66    pub async fn unregister(&self, agent_id: &str) -> AcpResult<()> {
67        let mut agents = self.agents.write().await;
68        agents.remove(agent_id);
69        Ok(())
70    }
71
72    /// Find agent by ID
73    pub async fn find(&self, agent_id: &str) -> AcpResult<AgentInfo> {
74        let agents = self.agents.read().await;
75        agents
76            .get(agent_id)
77            .cloned()
78            .ok_or_else(|| AcpError::AgentNotFound(agent_id.to_string()))
79    }
80
81    /// Find agents by capability
82    pub async fn find_by_capability(&self, capability: &str) -> AcpResult<Vec<AgentInfo>> {
83        let agents = self.agents.read().await;
84        let matching = agents
85            .values()
86            .filter(|a| a.online && a.capabilities.contains(&capability.to_string()))
87            .cloned()
88            .collect();
89        Ok(matching)
90    }
91
92    /// List all registered agents
93    pub async fn list_all(&self) -> AcpResult<Vec<AgentInfo>> {
94        let agents = self.agents.read().await;
95        Ok(agents.values().cloned().collect())
96    }
97
98    /// List online agents
99    pub async fn list_online(&self) -> AcpResult<Vec<AgentInfo>> {
100        let agents = self.agents.read().await;
101        Ok(agents.values().filter(|a| a.online).cloned().collect())
102    }
103
104    /// Update agent status
105    pub async fn update_status(&self, agent_id: &str, online: bool) -> AcpResult<()> {
106        let mut agents = self.agents.write().await;
107        if let Some(agent) = agents.get_mut(agent_id) {
108            agent.online = online;
109            agent.last_seen = Some(chrono::Utc::now().to_rfc3339());
110            Ok(())
111        } else {
112            Err(AcpError::AgentNotFound(agent_id.to_string()))
113        }
114    }
115
116    /// Get agent count
117    pub async fn count(&self) -> usize {
118        self.agents.read().await.len()
119    }
120
121    /// Clear all agents
122    pub async fn clear(&self) {
123        self.agents.write().await.clear();
124    }
125}
126
127impl Default for AgentRegistry {
128    fn default() -> Self {
129        Self::new()
130    }
131}
132
133#[cfg(test)]
134mod tests {
135    use super::*;
136
137    #[tokio::test]
138    async fn test_agent_registry() {
139        let registry = AgentRegistry::new();
140
141        let agent = AgentInfo {
142            id: "test-agent".to_string(),
143            name: "Test Agent".to_string(),
144            base_url: "http://localhost:8080".to_string(),
145            description: Some("A test agent".to_string()),
146            capabilities: vec!["bash".to_string(), "python".to_string()],
147            metadata: HashMap::new(),
148            online: true,
149            last_seen: None,
150        };
151
152        registry.register(agent.clone()).await.unwrap();
153
154        let found = registry.find("test-agent").await.unwrap();
155        assert_eq!(found.id, "test-agent");
156
157        assert_eq!(registry.count().await, 1);
158
159        registry.unregister("test-agent").await.unwrap();
160        assert_eq!(registry.count().await, 0);
161    }
162
163    #[tokio::test]
164    async fn test_find_by_capability() {
165        let registry = AgentRegistry::new();
166
167        let agent1 = AgentInfo {
168            id: "agent-1".to_string(),
169            name: "Agent 1".to_string(),
170            base_url: "http://localhost:8080".to_string(),
171            description: None,
172            capabilities: vec!["bash".to_string()],
173            metadata: HashMap::new(),
174            online: true,
175            last_seen: None,
176        };
177
178        let agent2 = AgentInfo {
179            id: "agent-2".to_string(),
180            name: "Agent 2".to_string(),
181            base_url: "http://localhost:8081".to_string(),
182            description: None,
183            capabilities: vec!["bash".to_string(), "python".to_string()],
184            metadata: HashMap::new(),
185            online: true,
186            last_seen: None,
187        };
188
189        registry.register(agent1).await.unwrap();
190        registry.register(agent2).await.unwrap();
191
192        let bash_agents = registry.find_by_capability("bash").await.unwrap();
193        assert_eq!(bash_agents.len(), 2);
194
195        let python_agents = registry.find_by_capability("python").await.unwrap();
196        assert_eq!(python_agents.len(), 1);
197    }
198}
vtcode_acp_client/discovery.rs

vtcode_acp_client/
discovery.rs
