mesh_client/
lib.rs

1use anyhow::Context;
2use async_nats::Client as NatsClient;
3use mesh_types::AgentCard;
4use reqwest::Client as HttpClient;
5use serde_json::json;
6use std::time::Duration;
7use uuid::Uuid;
8use futures::StreamExt;
9use futures::future::BoxFuture;
10use tokio::sync::watch;
11
12#[derive(Clone)]
13pub struct MeshClientConfig {
14    pub nats_url: String,
15    pub registry_url: String,
16    pub did: String,
17    pub region: Option<String>,
18}
19
20pub struct RegistryClient {
21    base: String,
22    http: HttpClient,
23}
24
25impl RegistryClient {
26    pub fn new(base: impl Into<String>) -> Self {
27        Self { base: base.into(), http: HttpClient::new() }
28    }
29
30    pub async fn get_agent(&self, did: &str) -> Result<AgentCard, anyhow::Error> {
31        let url = format!("{}/agents/{}", self.base.trim_end_matches('/'), did);
32        let res = self
33            .http
34            .get(&url)
35            .send()
36            .await
37            .with_context(|| format!("failed to GET {}", url))?;
38        let status = res.status();
39        let body = res.text().await?;
40        if !status.is_success() {
41            return Err(anyhow::anyhow!("registry GET failed: {} {}", status, body));
42        }
43        let card: AgentCard = serde_json::from_str(&body)?;
44        Ok(card)
45    }
46}
47
48pub struct MeshClient {
49    cfg: MeshClientConfig,
50    nats: NatsClient,
51    registry: RegistryClient,
52}
53
54pub struct RequestOptions {
55    pub domain: Vec<String>,
56    pub capability_id: String,
57    pub description: Option<String>,
58    pub timeout_ms: Option<u64>,
59}
60
61pub type MatchHandler = Box<dyn Fn(mesh_types::MatchOrReject) -> BoxFuture<'static, ()> + Send + Sync + 'static>;
62
63pub struct ListenSubscription {
64    stop_tx: watch::Sender<bool>,
65}
66
67impl ListenSubscription {
68    pub async fn unsubscribe(self) {
69        let _ = self.stop_tx.send(true);
70    }
71}
72
73pub struct HeartbeatHandle {
74    stop_tx: watch::Sender<bool>,
75}
76
77impl HeartbeatHandle {
78    pub async fn stop(self) {
79        let _ = self.stop_tx.send(true);
80    }
81}
82
83impl MeshClient {
84    pub async fn new(cfg: MeshClientConfig) -> Result<Self, anyhow::Error> {
85        let nats = async_nats::connect(&cfg.nats_url)
86            .await
87            .with_context(|| format!("connect to nats at {}", &cfg.nats_url))?;
88        let registry = RegistryClient::new(cfg.registry_url.clone());
89        Ok(Self { cfg, nats, registry })
90    }
91
92    pub async fn register(&self, card: &AgentCard) -> Result<(), anyhow::Error> {
93        // Validate AgentCard locally before sending to registry
94        mesh_types::validate_agent_card(card).map_err(|e| anyhow::anyhow!("invalid agent card: {}", e))?;
95
96        let url = format!("{}/agents", self.cfg.registry_url.trim_end_matches('/'));
97        let res = self
98            .registry
99            .http
100            .post(&url)
101            .json(card)
102            .send()
103            .await
104            .context("failed to POST agent card")?;
105        let status = res.status();
106        if !status.is_success() {
107            let body = res.text().await.unwrap_or_default();
108            return Err(anyhow::anyhow!("register failed: {} {}", status, body));
109        }
110        Ok(())
111    }
112
113    /// Publish a capability request on the control plane and await a reply.
114    /// Returns the raw JSON reply as serde_json::Value for now.
115    pub async fn request(&self, opts: RequestOptions) -> Result<mesh_types::MatchOrReject, anyhow::Error> {
116        let region = self.cfg.region.clone().unwrap_or_else(|| "global".into());
117        let subject = format!("mesh.requests.{}.{}", opts.domain.join("."), region);
118
119        let event_id = Uuid::new_v4().to_string();
120        let cloud_event = json!({
121            "specversion": "1.0",
122            "type": "amp.capability.request",
123            "source": self.cfg.did,
124            "id": event_id,
125            "time": chrono::Utc::now().to_rfc3339(),
126            "data": {
127                "capability_id": opts.capability_id,
128                "description": opts.description,
129            }
130        });
131
132        let payload = serde_json::to_vec(&cloud_event)?;
133
134        let timeout = Duration::from_millis(opts.timeout_ms.unwrap_or(30_000));
135
136        let req_fut = self.nats.request(subject, payload.into());
137        let msg = tokio::time::timeout(timeout, req_fut)
138            .await
139            .map_err(|_| anyhow::anyhow!("request timeout"))??;
140
141        let reply = String::from_utf8_lossy(&msg.payload).to_string();
142        let json: serde_json::Value = serde_json::from_str(&reply).context("parse reply JSON")?;
143        let parsed = mesh_types::parse_match_or_reject(&json).map_err(|e| anyhow::anyhow!("parse reply: {}", e))?;
144        Ok(parsed)
145    }
146
147    fn sanitize_did(did: &str) -> String {
148        did.replace(':', "_")
149    }
150
151    /// Listen for matches directed to this agent DID.
152    /// Handler will be invoked for each incoming match/reject.
153    pub async fn listen(&self, handler: MatchHandler) -> Result<ListenSubscription, anyhow::Error> {
154        let subj = format!("mesh.matches.{}", Self::sanitize_did(&self.cfg.did));
155        let mut sub = self.nats.subscribe(subj).await.context("subscribe failed")?;
156
157        let (stop_tx, mut stop_rx) = watch::channel(false);
158
159        tokio::spawn(async move {
160            loop {
161                tokio::select! {
162                    maybe = sub.next() => {
163                        match maybe {
164                            Some(msg) => {
165                                let payload = String::from_utf8_lossy(&msg.payload).to_string();
166                                if let Ok(json) = serde_json::from_str::<serde_json::Value>(&payload) {
167                                    if let Ok(parsed) = mesh_types::parse_match_or_reject(&json) {
168                                        (handler)(parsed).await;
169                                    }
170                                }
171                            }
172                            None => break,
173                        }
174                    }
175                    _ = stop_rx.changed() => {
176                        if *stop_rx.borrow() { break; }
177                    }
178                }
179            }
180        });
181
182        Ok(ListenSubscription { stop_tx })
183    }
184
185    /// Start periodic heartbeats to the control plane.
186    pub async fn start_heartbeat(&self, interval_ms: u64) -> Result<HeartbeatHandle, anyhow::Error> {
187        let subj = format!("mesh.agents.heartbeat.{}", Self::sanitize_did(&self.cfg.did));
188        let (stop_tx, mut stop_rx) = watch::channel(false);
189        let nats = self.nats.clone();
190        let did = self.cfg.did.clone();
191
192        tokio::spawn(async move {
193            let mut interval = tokio::time::interval(Duration::from_millis(interval_ms));
194            loop {
195                tokio::select! {
196                    _ = interval.tick() => {
197                        let payload = json!({ "did": did, "time": chrono::Utc::now().to_rfc3339() });
198                        let _ = nats.publish(subj.clone(), serde_json::to_vec(&payload).unwrap().into()).await;
199                    }
200                    _ = stop_rx.changed() => {
201                        if *stop_rx.borrow() { break; }
202                    }
203                }
204            }
205        });
206
207        Ok(HeartbeatHandle { stop_tx })
208    }
209
210    pub async fn close(self) -> Result<(), anyhow::Error> {
211        // NATS client flushes on drop
212        Ok(())
213    }
214}
215
mesh_client/lib.rs

mesh_client/
lib.rs
