zclaw_openfang/crates/zclaw-kernel/src/registry.rs

//! Agent registry

use dashmap::DashMap;
use zclaw_types::{AgentConfig, AgentId, AgentInfo, AgentState};
use chrono::Utc;

/// In-memory registry of active agents
pub struct AgentRegistry {
    agents: DashMap<AgentId, AgentConfig>,
    states: DashMap<AgentId, AgentState>,
    created_at: DashMap<AgentId, chrono::DateTime<Utc>>,
    message_counts: DashMap<AgentId, u64>,
}

impl AgentRegistry {
    pub fn new() -> Self {
        Self {
            agents: DashMap::new(),
            states: DashMap::new(),
            created_at: DashMap::new(),
            message_counts: DashMap::new(),
        }
    }

    /// Register an agent
    pub fn register(&self, config: AgentConfig) {
        let id = config.id;
        self.agents.insert(id, config);
        self.states.insert(id, AgentState::Running);
        self.created_at.insert(id, Utc::now());
    }

    /// Register an agent with persisted runtime state
    pub fn register_with_runtime(
        &self,
        config: AgentConfig,
        state: AgentState,
        message_count: u64,
    ) {
        let id = config.id;
        self.agents.insert(id, config);
        self.states.insert(id, state);
        self.created_at.insert(id, Utc::now());
        if message_count > 0 {
            self.message_counts.insert(id, message_count);
        }
    }

    /// Unregister an agent
    pub fn unregister(&self, id: &AgentId) {
        self.agents.remove(id);
        self.states.remove(id);
        self.created_at.remove(id);
        self.message_counts.remove(id);
    }

    /// Update an agent's configuration (preserves state and message count)
    pub fn update(&self, config: AgentConfig) {
        let id = config.id;
        self.agents.insert(id, config);
    }

    /// Get an agent by ID
    pub fn get(&self, id: &AgentId) -> Option<AgentConfig> {
        self.agents.get(id).map(|r| r.clone())
    }

    /// Get agent info
    pub fn get_info(&self, id: &AgentId) -> Option<AgentInfo> {
        let config = self.agents.get(id)?;
        let state = self.states.get(id).map(|s| *s).unwrap_or(AgentState::Terminated);
        let created_at = self.created_at.get(id).map(|t| *t).unwrap_or_else(Utc::now);

        Some(AgentInfo {
            id: *id,
            name: config.name.clone(),
            description: config.description.clone(),
            model: config.model.model.clone(),
            provider: config.model.provider.clone(),
            state,
            message_count: self.message_counts.get(id).map(|c| *c as usize).unwrap_or(0),
            created_at,
            updated_at: Utc::now(),
            soul: config.soul.clone(),
            system_prompt: config.system_prompt.clone(),
            temperature: config.temperature,
            max_tokens: config.max_tokens,
            user_profile: None,
        })
    }

    /// List all agents
    pub fn list(&self) -> Vec<AgentInfo> {
        self.agents.iter()
            .filter_map(|entry| {
                let id = entry.key();
                self.get_info(id)
            })
            .collect()
    }

    /// Update agent state
    pub fn set_state(&self, id: &AgentId, state: AgentState) {
        self.states.insert(*id, state);
    }

    /// Get agent state
    pub fn get_state(&self, id: &AgentId) -> AgentState {
        self.states.get(id).map(|s| *s).unwrap_or(AgentState::Terminated)
    }

    /// Count active agents
    pub fn count(&self) -> usize {
        self.agents.len()
    }

    /// Increment message count for an agent
    pub fn increment_message_count(&self, id: &AgentId) {
        self.message_counts.entry(*id).and_modify(|c| *c += 1).or_insert(1);
    }
}

impl Default for AgentRegistry {
    fn default() -> Self {
        Self::new()
    }
}