1use std::sync::atomic::{AtomicU64, Ordering};
21
22use super::protocol::{A2AErrorData, A2AMessage, A2ARequest, A2AResponse, A2ATask, AgentCard};
23
24#[derive(Debug, thiserror::Error)]
26pub enum A2AError {
27 #[error("HTTP error: {0}")]
29 Http(String),
30
31 #[error("Parse error: {0}")]
33 Parse(String),
34
35 #[error("API error [{code}]: {message}")]
37 Api { code: i32, message: String },
38
39 #[error("Timeout: {0}")]
41 Timeout(String),
42}
43
44impl From<reqwest::Error> for A2AError {
45 fn from(err: reqwest::Error) -> Self {
46 if err.is_timeout() {
47 A2AError::Timeout(err.to_string())
48 } else {
49 A2AError::Http(err.to_string())
50 }
51 }
52}
53
54impl From<A2AErrorData> for A2AError {
55 fn from(err: A2AErrorData) -> Self {
56 A2AError::Api {
57 code: err.code,
58 message: err.message,
59 }
60 }
61}
62
63pub struct A2AClient {
65 base_url: String,
67 http: reqwest::Client,
69 next_id: AtomicU64,
71}
72
73impl A2AClient {
74 pub fn new(base_url: String) -> Self {
76 Self {
77 base_url: base_url.trim_end_matches('/').to_string(),
78 http: reqwest::Client::new(),
79 next_id: AtomicU64::new(1),
80 }
81 }
82
83 pub fn with_http_client(base_url: String, http: reqwest::Client) -> Self {
85 Self {
86 base_url: base_url.trim_end_matches('/').to_string(),
87 http,
88 next_id: AtomicU64::new(1),
89 }
90 }
91
92 fn alloc_id(&self) -> u64 {
94 self.next_id.fetch_add(1, Ordering::SeqCst)
95 }
96
97 pub async fn get_agent_card(&self) -> Result<AgentCard, A2AError> {
99 let url = format!("{}/.well-known/agent.json", self.base_url);
100 let resp = self.http.get(&url).send().await?;
101 let status = resp.status();
102 if !status.is_success() {
103 return Err(A2AError::Http(format!(
104 "Agent card request failed with status {}",
105 status
106 )));
107 }
108 let card: AgentCard = resp
109 .json()
110 .await
111 .map_err(|e| A2AError::Parse(format!("Failed to parse agent card: {}", e)))?;
112 Ok(card)
113 }
114
115 pub async fn send_task(&self, message: A2AMessage) -> Result<A2ATask, A2AError> {
117 let id = self.alloc_id();
118 let req = A2ARequest::send_task(id, &message);
119 let resp = self.post_request(req).await?;
120
121 let result = resp.into_result().map_err(A2AError::from)?;
123 let task: A2ATask = result
124 .get("task")
125 .ok_or_else(|| A2AError::Parse("Missing 'task' in response".to_string()))
126 .and_then(|v| {
127 serde_json::from_value(v.clone())
128 .map_err(|e| A2AError::Parse(format!("Failed to parse task: {}", e)))
129 })?;
130 Ok(task)
131 }
132
133 pub async fn get_task(&self, task_id: &str) -> Result<A2ATask, A2AError> {
135 let id = self.alloc_id();
136 let req = A2ARequest::get_task(id, task_id);
137 let resp = self.post_request(req).await?;
138
139 let result = resp.into_result().map_err(A2AError::from)?;
140 let task: A2ATask = result
141 .get("task")
142 .ok_or_else(|| A2AError::Parse("Missing 'task' in response".to_string()))
143 .and_then(|v| {
144 serde_json::from_value(v.clone())
145 .map_err(|e| A2AError::Parse(format!("Failed to parse task: {}", e)))
146 })?;
147 Ok(task)
148 }
149
150 pub async fn cancel_task(&self, task_id: &str) -> Result<A2ATask, A2AError> {
152 let id = self.alloc_id();
153 let req = A2ARequest::cancel_task(id, task_id);
154 let resp = self.post_request(req).await?;
155
156 let result = resp.into_result().map_err(A2AError::from)?;
157 let task: A2ATask = result
158 .get("task")
159 .ok_or_else(|| A2AError::Parse("Missing 'task' in response".to_string()))
160 .and_then(|v| {
161 serde_json::from_value(v.clone())
162 .map_err(|e| A2AError::Parse(format!("Failed to parse task: {}", e)))
163 })?;
164 Ok(task)
165 }
166
167 pub async fn post_request(&self, req: A2ARequest) -> Result<A2AResponse, A2AError> {
169 let url = format!("{}/", self.base_url);
170 let resp = self.http.post(&url).json(&req).send().await?;
171 let status = resp.status();
172 if !status.is_success() {
173 return Err(A2AError::Http(format!(
174 "A2A request failed with status {}",
175 status
176 )));
177 }
178 let a2a_resp: A2AResponse = resp
179 .json()
180 .await
181 .map_err(|e| A2AError::Parse(format!("Failed to parse A2A response: {}", e)))?;
182 Ok(a2a_resp)
183 }
184}
185
186#[cfg(test)]
187mod tests {
188 use super::*;
189
190 #[test]
191 fn a2a_error_from_reqwest_timeout() {
192 let err = A2AError::Timeout("connection timed out".to_string());
194 assert!(err.to_string().contains("Timeout"));
195 }
196
197 #[test]
198 fn a2a_error_from_reqwest_http() {
199 let err = A2AError::Http("404 not found".to_string());
200 assert!(err.to_string().contains("HTTP error"));
201 }
202
203 #[test]
204 fn a2a_error_from_error_data() {
205 let data = A2AErrorData::method_not_found();
206 let err: A2AError = data.into();
207 match err {
208 A2AError::Api { code, message } => {
209 assert_eq!(code, -32601);
210 assert!(message.contains("Method not found"));
211 }
212 _ => panic!("Expected Api variant"),
213 }
214 }
215
216 #[test]
217 fn a2a_error_parse() {
218 let err = A2AError::Parse("bad json".to_string());
219 assert!(err.to_string().contains("Parse error"));
220 }
221
222 #[test]
223 fn client_new_trims_trailing_slash() {
224 let client = A2AClient::new("http://localhost:8080/".to_string());
225 assert_eq!(client.base_url, "http://localhost:8080");
226 }
227
228 #[test]
229 fn client_alloc_id_increments() {
230 let client = A2AClient::new("http://localhost:8080".to_string());
231 assert_eq!(client.alloc_id(), 1);
232 assert_eq!(client.alloc_id(), 2);
233 assert_eq!(client.alloc_id(), 3);
234 }
235
236 #[test]
237 fn client_with_custom_http() {
238 let http = reqwest::Client::builder()
239 .timeout(std::time::Duration::from_secs(30))
240 .build()
241 .unwrap();
242 let client = A2AClient::with_http_client("http://localhost:8080".to_string(), http);
243 assert_eq!(client.base_url, "http://localhost:8080");
244 }
245
246 #[tokio::test]
247 async fn get_agent_card_invalid_url() {
248 let client = A2AClient::new("http://localhost:19999".to_string());
249 let result = client.get_agent_card().await;
250 assert!(result.is_err());
251 match result.unwrap_err() {
252 A2AError::Http(_) | A2AError::Timeout(_) => {} other => panic!("Expected Http or Timeout error, got: {:?}", other),
254 }
255 }
256
257 #[tokio::test]
258 async fn send_task_invalid_url() {
259 let client = A2AClient::new("http://localhost:19999".to_string());
260 let result = client.send_task(A2AMessage::user("hello")).await;
261 assert!(result.is_err());
262 }
263
264 #[tokio::test]
265 async fn get_task_invalid_url() {
266 let client = A2AClient::new("http://localhost:19999".to_string());
267 let result = client.get_task("task-123").await;
268 assert!(result.is_err());
269 }
270
271 #[tokio::test]
272 async fn cancel_task_invalid_url() {
273 let client = A2AClient::new("http://localhost:19999".to_string());
274 let result = client.cancel_task("task-123").await;
275 assert!(result.is_err());
276 }
277
278 #[tokio::test]
279 async fn post_request_invalid_url() {
280 let client = A2AClient::new("http://localhost:19999".to_string());
281 let req = A2ARequest::new(1, "test", None);
282 let result = client.post_request(req).await;
283 assert!(result.is_err());
284 }
285}