Skip to main content

a2a_rust/client/
discovery.rs

1use std::collections::BTreeMap;
2use std::sync::RwLock;
3use std::time::{Duration, Instant};
4
5use reqwest::Url;
6
7use crate::A2AError;
8use crate::jsonrpc::PROTOCOL_VERSION;
9use crate::types::AgentCard;
10
11/// Discovery cache configuration for remote agent cards.
12#[derive(Debug, Clone, Copy)]
13pub struct AgentCardDiscoveryConfig {
14    /// Maximum time to reuse a cached discovery response.
15    pub ttl: Duration,
16}
17
18impl Default for AgentCardDiscoveryConfig {
19    fn default() -> Self {
20        Self {
21            ttl: Duration::from_secs(300),
22        }
23    }
24}
25
26#[derive(Debug, Clone)]
27struct CachedAgentCard {
28    card: AgentCard,
29    fetched_at: Instant,
30}
31
32/// Discovers and caches remote A2A agent cards.
33#[derive(Debug)]
34pub struct AgentCardDiscovery {
35    client: reqwest::Client,
36    config: AgentCardDiscoveryConfig,
37    cache: RwLock<BTreeMap<String, CachedAgentCard>>,
38}
39
40impl Default for AgentCardDiscovery {
41    fn default() -> Self {
42        Self::new()
43    }
44}
45
46impl AgentCardDiscovery {
47    /// Create a discovery client with default caching behavior.
48    pub fn new() -> Self {
49        Self::with_config(AgentCardDiscoveryConfig::default())
50    }
51
52    /// Create a discovery client with explicit cache settings.
53    pub fn with_config(config: AgentCardDiscoveryConfig) -> Self {
54        Self::with_http_client(reqwest::Client::new(), config)
55    }
56
57    /// Create a discovery client with a caller-provided HTTP client.
58    pub fn with_http_client(client: reqwest::Client, config: AgentCardDiscoveryConfig) -> Self {
59        Self {
60            client,
61            config,
62            cache: RwLock::new(BTreeMap::new()),
63        }
64    }
65
66    /// Discover an agent card, using the cache when still fresh.
67    pub async fn discover(&self, base_url: &str) -> Result<AgentCard, A2AError> {
68        let base_url = normalize_base_url(base_url)?;
69        let cache_key = cache_key(&base_url);
70
71        if let Some(card) = self.cached_card(&cache_key)? {
72            return Ok(card);
73        }
74
75        self.fetch_and_store(cache_key, base_url).await
76    }
77
78    /// Force a fresh agent-card fetch and replace any cached entry.
79    pub async fn refresh(&self, base_url: &str) -> Result<AgentCard, A2AError> {
80        let base_url = normalize_base_url(base_url)?;
81        self.fetch_and_store(cache_key(&base_url), base_url).await
82    }
83
84    fn cached_card(&self, cache_key: &str) -> Result<Option<AgentCard>, A2AError> {
85        let cache = self
86            .cache
87            .read()
88            .map_err(|_| A2AError::Internal("discovery cache lock poisoned".to_owned()))?;
89
90        let Some(cached) = cache.get(cache_key) else {
91            return Ok(None);
92        };
93
94        if cached.fetched_at.elapsed() >= self.config.ttl {
95            return Ok(None);
96        }
97
98        Ok(Some(cached.card.clone()))
99    }
100
101    async fn fetch_and_store(
102        &self,
103        cache_key: String,
104        base_url: Url,
105    ) -> Result<AgentCard, A2AError> {
106        let discovery_url = well_known_agent_card_url(&base_url)?;
107        let response = self
108            .client
109            .get(discovery_url)
110            .header("A2A-Version", PROTOCOL_VERSION)
111            .send()
112            .await?;
113
114        let status = response.status();
115        let bytes = response.bytes().await?;
116        if !status.is_success() {
117            return Err(A2AError::InvalidAgentResponse(format!(
118                "agent discovery returned HTTP {}",
119                status
120            )));
121        }
122
123        let card: AgentCard = serde_json::from_slice(&bytes)
124            .map_err(|error| A2AError::InvalidAgentResponse(error.to_string()))?;
125        let mut cache = self
126            .cache
127            .write()
128            .map_err(|_| A2AError::Internal("discovery cache lock poisoned".to_owned()))?;
129        cache.insert(
130            cache_key,
131            CachedAgentCard {
132                card: card.clone(),
133                fetched_at: Instant::now(),
134            },
135        );
136
137        Ok(card)
138    }
139}
140
141pub(crate) fn normalize_base_url(base_url: &str) -> Result<Url, A2AError> {
142    let mut url =
143        Url::parse(base_url).map_err(|error| A2AError::InvalidRequest(error.to_string()))?;
144    url.set_query(None);
145    url.set_fragment(None);
146    Ok(url)
147}
148
149pub(crate) fn resolve_interface_url(base_url: &Url, interface_url: &str) -> Result<Url, A2AError> {
150    Url::parse(interface_url)
151        .or_else(|_| base_url.join(interface_url))
152        .map_err(|error| A2AError::InvalidAgentResponse(error.to_string()))
153}
154
155pub(crate) fn ensure_trailing_slash(mut url: Url) -> Url {
156    if !url.path().ends_with('/') {
157        let path = format!("{}/", url.path());
158        url.set_path(&path);
159    }
160
161    url
162}
163
164fn cache_key(base_url: &Url) -> String {
165    let mut normalized = base_url.clone();
166    if normalized.path() != "/" {
167        let trimmed = normalized.path().trim_end_matches('/').to_owned();
168        normalized.set_path(&trimmed);
169    }
170
171    normalized.to_string()
172}
173
174fn well_known_agent_card_url(base_url: &Url) -> Result<Url, A2AError> {
175    ensure_trailing_slash(base_url.clone())
176        .join(".well-known/agent-card.json")
177        .map_err(|error| A2AError::InvalidRequest(error.to_string()))
178}