tap_node/agent/
mod.rs

1//! Agent management for TAP Node
2//!
3//! This module provides utilities for managing multiple TAP agents within a TAP Node.
4
5use dashmap::DashMap;
6use std::sync::Arc;
7
8use crate::error::{Error, Result};
9use tap_agent::TapAgent;
10
11/// Registry of TAP agents
12#[derive(Debug)]
13pub struct AgentRegistry {
14    /// Maximum number of agents that can be registered
15    max_agents: Option<usize>,
16    /// Mapping of agent DIDs to agent instances
17    agents: DashMap<String, Arc<TapAgent>>,
18}
19
20impl AgentRegistry {
21    /// Create a new agent registry with the specified maximum number of agents
22    pub fn new(max_agents: Option<usize>) -> Self {
23        Self {
24            max_agents,
25            agents: DashMap::new(),
26        }
27    }
28
29    /// Get the current number of registered agents
30    pub fn agent_count(&self) -> usize {
31        self.agents.len()
32    }
33
34    /// Check if the registry has an agent with the given DID
35    pub fn has_agent(&self, did: &str) -> bool {
36        self.agents.contains_key(did)
37    }
38
39    /// Get an agent by DID
40    pub async fn get_agent(&self, did: &str) -> Result<Arc<TapAgent>> {
41        self.agents
42            .get(did)
43            .map(|agent| agent.clone())
44            .ok_or_else(|| Error::AgentNotFound(did.to_string()))
45    }
46
47    /// Register a new agent
48    pub async fn register_agent(&self, did: String, agent: Arc<TapAgent>) -> Result<()> {
49        // Check if we've reached the maximum number of agents
50        if let Some(max) = self.max_agents {
51            if self.agent_count() >= max {
52                return Err(Error::AgentRegistration(format!(
53                    "Maximum number of agents ({}) reached",
54                    max
55                )));
56            }
57        }
58
59        // Check if the agent is already registered
60        if self.has_agent(&did) {
61            return Err(Error::AgentRegistration(format!(
62                "Agent with DID {} is already registered",
63                did
64            )));
65        }
66
67        // Register the agent
68        self.agents.insert(did, agent);
69
70        Ok(())
71    }
72
73    /// Unregister an agent
74    pub async fn unregister_agent(&self, did: &str) -> Result<()> {
75        // Check if the agent exists
76        if !self.has_agent(did) {
77            return Err(Error::AgentNotFound(did.to_string()));
78        }
79
80        // Unregister the agent
81        self.agents.remove(did);
82
83        Ok(())
84    }
85
86    /// Get all registered agent DIDs
87    pub fn get_all_dids(&self) -> Vec<String> {
88        self.agents
89            .iter()
90            .map(|entry| entry.key().clone())
91            .collect()
92    }
93}
tap_node/agent/mod.rs

tap_node/agent/
mod.rs
