symbi-runtime 1.21.1

//! Agent registry for multi-agent composition
//!
//! Maps agent names to their configurations and IDs, enabling
//! runtime agent spawning and lifecycle management.

use crate::reasoning::inference::InferenceProvider;
use crate::types::AgentId;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

/// Configuration for a registered agent.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RegisteredAgent {
    /// Unique identifier.
    pub agent_id: AgentId,
    /// Human-readable name.
    pub name: String,
    /// System prompt for this agent.
    pub system_prompt: String,
    /// Tool names this agent has access to.
    pub tools: Vec<String>,
    /// Optional response format (e.g., "json", "text").
    pub response_format: Option<String>,
    /// When this agent was registered.
    pub created_at: chrono::DateTime<chrono::Utc>,
    /// Retained canonical requirements. Prompt-only routes cannot execute these.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub definition: Option<dsl::ConversationalAgent>,
}

/// Thread-safe registry of named agents.
#[derive(Clone)]
pub struct AgentRegistry {
    agents: Arc<RwLock<HashMap<String, RegisteredAgent>>>,
}

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

impl AgentRegistry {
    /// Create a new empty registry.
    pub fn new() -> Self {
        Self {
            agents: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Spawn (register) a new agent.
    pub async fn spawn_agent(
        &self,
        name: impl Into<String>,
        system_prompt: impl Into<String>,
        tools: Vec<String>,
        response_format: Option<String>,
    ) -> AgentId {
        let name = name.into();
        let agent_id = AgentId::new();

        let agent = RegisteredAgent {
            agent_id,
            name: name.clone(),
            system_prompt: system_prompt.into(),
            tools,
            response_format,
            created_at: chrono::Utc::now(),
            definition: None,
        };

        self.agents.write().await.insert(name, agent);
        agent_id
    }

    /// Register an operator-loaded source snapshot, without reparsing a mutable path.
    pub async fn register_canonical(
        &self,
        definition: dsl::ConversationalAgent,
        system_prompt: String,
        tools: Vec<String>,
    ) -> AgentId {
        let agent_id = AgentId::new();
        let name = definition.settings().agent_name.clone();
        let agent = RegisteredAgent {
            agent_id,
            name: name.clone(),
            system_prompt,
            tools,
            response_format: None,
            created_at: chrono::Utc::now(),
            definition: Some(definition),
        };
        self.agents.write().await.insert(name, agent);
        agent_id
    }

    /// Dynamic DSL registration may not replace an operator's canonical contract.
    pub async fn spawn_prompt_agent(
        &self,
        name: String,
        system_prompt: String,
        tools: Vec<String>,
        response_format: Option<String>,
    ) -> Result<AgentId, AgentRegistryError> {
        let mut agents = self.agents.write().await;
        if agents
            .get(&name)
            .is_some_and(|agent| agent.definition.is_some())
        {
            return Err(AgentRegistryError::UnsupportedExecution {
                name,
                reason: "dynamic registration cannot replace a canonical source contract".into(),
            });
        }
        let agent_id = AgentId::new();
        agents.insert(
            name.clone(),
            RegisteredAgent {
                agent_id,
                name,
                system_prompt,
                tools,
                response_format,
                created_at: chrono::Utc::now(),
                definition: None,
            },
        );
        Ok(agent_id)
    }

    /// Get a registered agent by name.
    pub async fn get_agent(&self, name: &str) -> Option<RegisteredAgent> {
        self.agents.read().await.get(name).cloned()
    }

    /// List all registered agents.
    pub async fn list_agents(&self) -> Vec<RegisteredAgent> {
        self.agents.read().await.values().cloned().collect()
    }

    /// Remove an agent by name.
    pub async fn remove_agent(&self, name: &str) -> bool {
        self.agents.write().await.remove(name).is_some()
    }

    /// Check if an agent exists.
    pub async fn has_agent(&self, name: &str) -> bool {
        self.agents.read().await.contains_key(name)
    }

    /// Send a message to an agent and get a response.
    ///
    /// Uses the agent's system prompt and the provided inference provider
    /// to run a single-turn conversation.
    pub async fn ask_agent(
        &self,
        name: &str,
        message: &str,
        provider: &dyn InferenceProvider,
    ) -> Result<String, AgentRegistryError> {
        let agent = self
            .get_agent(name)
            .await
            .ok_or_else(|| AgentRegistryError::NotFound {
                name: name.to_string(),
            })?;
        if agent.definition.is_some() {
            return Err(AgentRegistryError::UnsupportedExecution {
                name: name.into(), reason: "canonical agents require a governed executor with their retained source contract".into(),
            });
        }

        use crate::reasoning::conversation::{Conversation, ConversationMessage};
        use crate::reasoning::inference::InferenceOptions;

        let mut conv = Conversation::with_system(&agent.system_prompt);
        conv.push(ConversationMessage::user(message));

        let options = InferenceOptions::default();
        let response = provider.complete(&conv, &options).await.map_err(|e| {
            AgentRegistryError::InferenceError {
                agent_name: name.to_string(),
                message: e.to_string(),
            }
        })?;

        Ok(response.content)
    }
}

/// Errors from the agent registry.
#[derive(Debug, thiserror::Error)]
pub enum AgentRegistryError {
    #[error("Agent '{name}' cannot use this execution route: {reason}")]
    UnsupportedExecution { name: String, reason: String },
    #[error("Agent '{name}' not found")]
    NotFound { name: String },

    #[error("Inference error for agent '{agent_name}': {message}")]
    InferenceError { agent_name: String, message: String },
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_spawn_and_get_agent() {
        let registry = AgentRegistry::new();

        let id = registry
            .spawn_agent(
                "researcher",
                "You are a researcher.",
                vec!["search".into()],
                None,
            )
            .await;

        let agent = registry.get_agent("researcher").await.unwrap();
        assert_eq!(agent.agent_id, id);
        assert_eq!(agent.name, "researcher");
        assert_eq!(agent.tools, vec!["search"]);
    }

    #[tokio::test]
    async fn test_list_agents() {
        let registry = AgentRegistry::new();

        registry.spawn_agent("a", "Agent A", vec![], None).await;
        registry.spawn_agent("b", "Agent B", vec![], None).await;

        let agents = registry.list_agents().await;
        assert_eq!(agents.len(), 2);
    }

    #[tokio::test]
    async fn test_remove_agent() {
        let registry = AgentRegistry::new();

        registry
            .spawn_agent("temp", "Temporary", vec![], None)
            .await;
        assert!(registry.has_agent("temp").await);

        assert!(registry.remove_agent("temp").await);
        assert!(!registry.has_agent("temp").await);
    }

    #[tokio::test]
    async fn test_get_nonexistent() {
        let registry = AgentRegistry::new();
        assert!(registry.get_agent("nope").await.is_none());
    }

    #[tokio::test]
    async fn test_spawn_replaces_existing() {
        let registry = AgentRegistry::new();

        let id1 = registry.spawn_agent("agent", "v1", vec![], None).await;
        let id2 = registry.spawn_agent("agent", "v2", vec![], None).await;

        assert_ne!(id1, id2);
        let agent = registry.get_agent("agent").await.unwrap();
        assert_eq!(agent.system_prompt, "v2");
    }
}