mur_core/a2a/

discovery.rs

//! A2A Discovery — well-known URLs and local agent registry.

use super::protocol::AgentCard;
use anyhow::{Context, Result};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::{Path, PathBuf};

/// Local registry of known A2A agents.
pub struct AgentRegistry {
    /// Path to the registry JSON file.
    registry_path: PathBuf,
}

/// A registered agent entry.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RegisteredAgent {
    /// Base URL of the agent.
    pub url: String,
    /// Agent name (from card).
    pub name: String,
    /// Agent description.
    pub description: String,
    /// Agent version.
    pub version: String,
    /// Skills offered.
    pub skills: Vec<String>,
    /// When the agent was registered.
    pub registered_at: DateTime<Utc>,
    /// When the agent was last seen healthy.
    pub last_seen: Option<DateTime<Utc>>,
    /// Whether the agent is currently reachable.
    pub healthy: bool,
    /// Optional tags for categorization.
    #[serde(default)]
    pub tags: Vec<String>,
}

impl AgentRegistry {
    /// Create a new agent registry.
    pub fn new(registry_path: &Path) -> Self {
        Self {
            registry_path: registry_path.to_path_buf(),
        }
    }

    /// Default registry path in the commander data directory.
    pub fn default_path() -> PathBuf {
        directories::BaseDirs::new()
            .map(|d| d.home_dir().join(".mur").join("commander").join("agents.json"))
            .expect("Cannot determine home directory for agent registry")
    }

    /// List all registered agents.
    pub fn list(&self) -> Result<Vec<RegisteredAgent>> {
        if !self.registry_path.exists() {
            return Ok(Vec::new());
        }
        let content = std::fs::read_to_string(&self.registry_path)
            .context("Reading agent registry")?;
        serde_json::from_str(&content).context("Parsing agent registry")
    }

    /// Register an agent from its card.
    pub fn register(&self, card: &AgentCard) -> Result<RegisteredAgent> {
        let mut agents = self.list()?;

        // Update existing or add new
        let entry = RegisteredAgent {
            url: card.url.clone(),
            name: card.name.clone(),
            description: card.description.clone(),
            version: card.version.clone(),
            skills: card.skills.iter().map(|s| s.id.clone()).collect(),
            registered_at: Utc::now(),
            last_seen: Some(Utc::now()),
            healthy: true,
            tags: vec![],
        };

        if let Some(existing) = agents.iter_mut().find(|a| a.url == card.url) {
            *existing = entry.clone();
        } else {
            agents.push(entry.clone());
        }

        self.save(&agents)?;
        Ok(entry)
    }

    /// Remove an agent by URL.
    pub fn remove(&self, url: &str) -> Result<bool> {
        let mut agents = self.list()?;
        let before = agents.len();
        agents.retain(|a| a.url != url);
        if agents.len() == before {
            return Ok(false);
        }
        self.save(&agents)?;
        Ok(true)
    }

    /// Update the health status of an agent.
    pub fn update_health(&self, url: &str, healthy: bool) -> Result<()> {
        let mut agents = self.list()?;
        if let Some(agent) = agents.iter_mut().find(|a| a.url == url) {
            agent.healthy = healthy;
            if healthy {
                agent.last_seen = Some(Utc::now());
            }
        }
        self.save(&agents)?;
        Ok(())
    }

    /// Find agents that have a specific skill.
    pub fn find_by_skill(&self, skill_id: &str) -> Result<Vec<RegisteredAgent>> {
        let agents = self.list()?;
        Ok(agents
            .into_iter()
            .filter(|a| a.skills.iter().any(|s| s == skill_id))
            .collect())
    }

    /// Find healthy agents only.
    pub fn healthy_agents(&self) -> Result<Vec<RegisteredAgent>> {
        let agents = self.list()?;
        Ok(agents.into_iter().filter(|a| a.healthy).collect())
    }

    /// Save the registry to disk.
    fn save(&self, agents: &[RegisteredAgent]) -> Result<()> {
        if let Some(parent) = self.registry_path.parent() {
            std::fs::create_dir_all(parent)?;
        }
        let json = serde_json::to_string_pretty(agents)?;
        std::fs::write(&self.registry_path, json).context("Writing agent registry")?;
        Ok(())
    }
}

/// Construct the well-known agent card URL from a base URL.
pub fn well_known_url(base_url: &str) -> String {
    format!(
        "{}/.well-known/agent.json",
        base_url.trim_end_matches('/')
    )
}

/// Discover agents from a list of well-known URLs.
pub async fn discover_many(
    client: &reqwest::Client,
    urls: &[String],
) -> HashMap<String, Result<AgentCard>> {
    let mut results = HashMap::new();

    for url in urls {
        let well_known = well_known_url(url);
        let result = async {
            let resp = client
                .get(&well_known)
                .send()
                .await
                .with_context(|| format!("Fetching {}", well_known))?;

            if !resp.status().is_success() {
                anyhow::bail!("HTTP {}", resp.status());
            }

            resp.json::<AgentCard>()
                .await
                .context("Parsing agent card")
        }
        .await;

        results.insert(url.clone(), result);
    }

    results
}

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

    fn test_registry() -> (TempDir, AgentRegistry) {
        let dir = TempDir::new().unwrap();
        let path = dir.path().join("agents.json");
        (dir, AgentRegistry::new(&path))
    }

    fn test_card() -> AgentCard {
        use super::super::protocol::*;
        AgentCard {
            name: "Test Agent".into(),
            description: "A test agent".into(),
            url: "http://localhost:8080/a2a".into(),
            version: "1.0.0".into(),
            protocol_version: "0.1".into(),
            capabilities: AgentCapabilities {
                streaming: false,
                push_notifications: false,
                state_management: false,
            },
            skills: vec![AgentSkill {
                id: "test-skill".into(),
                name: "Test".into(),
                description: "A test skill".into(),
                tags: vec![],
                input_schema: None,
                output_schema: None,
            }],
            authentication: None,
        }
    }

    #[test]
    fn test_empty_registry() {
        let (_dir, registry) = test_registry();
        let agents = registry.list().unwrap();
        assert!(agents.is_empty());
    }

    #[test]
    fn test_register_agent() {
        let (_dir, registry) = test_registry();
        let card = test_card();
        let entry = registry.register(&card).unwrap();
        assert_eq!(entry.name, "Test Agent");
        assert!(entry.healthy);

        let agents = registry.list().unwrap();
        assert_eq!(agents.len(), 1);
    }

    #[test]
    fn test_register_updates_existing() {
        let (_dir, registry) = test_registry();
        let card = test_card();
        registry.register(&card).unwrap();
        registry.register(&card).unwrap();

        let agents = registry.list().unwrap();
        assert_eq!(agents.len(), 1); // Should not duplicate
    }

    #[test]
    fn test_remove_agent() {
        let (_dir, registry) = test_registry();
        let card = test_card();
        registry.register(&card).unwrap();

        let removed = registry.remove("http://localhost:8080/a2a").unwrap();
        assert!(removed);

        let removed_again = registry.remove("http://localhost:8080/a2a").unwrap();
        assert!(!removed_again);

        assert!(registry.list().unwrap().is_empty());
    }

    #[test]
    fn test_update_health() {
        let (_dir, registry) = test_registry();
        let card = test_card();
        registry.register(&card).unwrap();

        registry
            .update_health("http://localhost:8080/a2a", false)
            .unwrap();

        let agents = registry.list().unwrap();
        assert!(!agents[0].healthy);
    }

    #[test]
    fn test_find_by_skill() {
        let (_dir, registry) = test_registry();
        let card = test_card();
        registry.register(&card).unwrap();

        let found = registry.find_by_skill("test-skill").unwrap();
        assert_eq!(found.len(), 1);

        let not_found = registry.find_by_skill("nonexistent").unwrap();
        assert!(not_found.is_empty());
    }

    #[test]
    fn test_healthy_agents() {
        let (_dir, registry) = test_registry();
        let card = test_card();
        registry.register(&card).unwrap();

        let healthy = registry.healthy_agents().unwrap();
        assert_eq!(healthy.len(), 1);

        registry
            .update_health("http://localhost:8080/a2a", false)
            .unwrap();
        let healthy = registry.healthy_agents().unwrap();
        assert!(healthy.is_empty());
    }

    #[test]
    fn test_well_known_url() {
        assert_eq!(
            well_known_url("http://localhost:3939"),
            "http://localhost:3939/.well-known/agent.json"
        );
        assert_eq!(
            well_known_url("http://localhost:3939/"),
            "http://localhost:3939/.well-known/agent.json"
        );
    }
}