File size: 4,074 Bytes
d90101d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | use std::sync::Arc;
use dashmap::DashMap;
use forge_app::domain::AgentId;
use forge_app::{AgentRepository, EnvironmentInfra};
use forge_domain::{Agent, AgentInfo};
use tokio::sync::RwLock;
/// AgentRegistryService manages the active-agent ID and a registry of runtime
/// Agents in-memory. It lazily loads agents from AgentRepository on first
/// access.
pub struct ForgeAgentRegistryService<R> {
// Infrastructure dependency for loading agents
repository: Arc<R>,
// In-memory storage for agents keyed by AgentId string
// Lazily initialized on first access
// Wrapped in RwLock to allow invalidation
agents: RwLock<Option<DashMap<String, Agent>>>,
// In-memory storage for the active agent ID
active_agent_id: RwLock<Option<AgentId>>,
}
impl<R> ForgeAgentRegistryService<R> {
/// Creates a new AgentRegistryService with the given repository
pub fn new(repository: Arc<R>) -> Self {
Self {
repository,
agents: RwLock::new(None),
active_agent_id: RwLock::new(None),
}
}
}
impl<R: AgentRepository + EnvironmentInfra<Config = forge_config::ForgeConfig>>
ForgeAgentRegistryService<R>
{
/// Lazily initializes and returns the agents map
/// Loads agents from repository on first call, subsequent calls return
/// cached value
async fn ensure_agents_loaded(&self) -> anyhow::Result<DashMap<String, Agent>> {
// Check if already loaded
{
let agents_read = self.agents.read().await;
if let Some(agents) = agents_read.as_ref() {
return Ok(agents.clone());
}
}
// Not loaded yet, acquire write lock and load
let mut agents_write = self.agents.write().await;
// Double-check in case another task loaded while we were waiting for write
// lock
if let Some(agents) = agents_write.as_ref() {
return Ok(agents.clone());
}
// Load agents
let agents_map = self.load_agents().await?;
// Store and return
*agents_write = Some(agents_map.clone());
Ok(agents_map)
}
/// Load agents from repository and populate the in-memory map.
///
/// Reads the default provider and model from the latest [`ForgeConfig`]
/// (via `get_config()`) and passes them to the repository so agents that
/// do not specify their own provider/model receive the session-level
/// defaults.
async fn load_agents(&self) -> anyhow::Result<DashMap<String, Agent>> {
let agents = self.repository.get_agents().await?;
let agents_map = DashMap::new();
for agent in agents {
agents_map.insert(agent.id.as_str().to_string(), agent);
}
Ok(agents_map)
}
}
#[async_trait::async_trait]
impl<R: AgentRepository + EnvironmentInfra<Config = forge_config::ForgeConfig> + Send + Sync>
forge_app::AgentRegistry for ForgeAgentRegistryService<R>
{
async fn get_active_agent_id(&self) -> anyhow::Result<Option<AgentId>> {
let agent_id = self.active_agent_id.read().await;
Ok(agent_id.clone())
}
async fn set_active_agent_id(&self, agent_id: AgentId) -> anyhow::Result<()> {
let mut active_agent = self.active_agent_id.write().await;
*active_agent = Some(agent_id);
Ok(())
}
async fn get_agents(&self) -> anyhow::Result<Vec<Agent>> {
let agents = self.ensure_agents_loaded().await?;
Ok(agents.iter().map(|entry| entry.value().clone()).collect())
}
async fn get_agent_infos(&self) -> anyhow::Result<Vec<AgentInfo>> {
self.repository.get_agent_infos().await
}
async fn get_agent(&self, agent_id: &AgentId) -> anyhow::Result<Option<Agent>> {
let agents = self.ensure_agents_loaded().await?;
Ok(agents.get(agent_id.as_str()).map(|v| v.value().clone()))
}
async fn reload_agents(&self) -> anyhow::Result<()> {
*self.agents.write().await = None;
self.ensure_agents_loaded().await?;
Ok(())
}
}
|