a2a_protocol_client/transport/rest/
mod.rs1mod query;
38mod request;
39mod routing;
40mod streaming;
41
42use std::collections::HashMap;
43use std::future::Future;
44use std::pin::Pin;
45use std::sync::Arc;
46use std::time::Duration;
47
48#[cfg(not(feature = "tls-rustls"))]
49use http_body_util::Full;
50#[cfg(not(feature = "tls-rustls"))]
51use hyper::body::Bytes;
52#[cfg(not(feature = "tls-rustls"))]
53use hyper_util::client::legacy::connect::HttpConnector;
54#[cfg(not(feature = "tls-rustls"))]
55use hyper_util::client::legacy::Client;
56#[cfg(not(feature = "tls-rustls"))]
57use hyper_util::rt::TokioExecutor;
58
59use crate::error::{ClientError, ClientResult};
60use crate::streaming::EventStream;
61use crate::transport::Transport;
62
63#[cfg(not(feature = "tls-rustls"))]
66type HttpClient = Client<HttpConnector, Full<Bytes>>;
67
68#[cfg(feature = "tls-rustls")]
69type HttpClient = crate::tls::HttpsClient;
70
71#[derive(Clone, Debug)]
78pub struct RestTransport {
79 inner: Arc<Inner>,
80}
81
82#[derive(Debug, Clone)]
83struct Inner {
84 client: HttpClient,
85 base_url: String,
86 request_timeout: Duration,
87 stream_connect_timeout: Duration,
88 max_response_size: usize,
89}
90
91impl RestTransport {
92 pub fn new(base_url: impl Into<String>) -> ClientResult<Self> {
98 Self::with_timeout(base_url, Duration::from_secs(30))
99 }
100
101 pub fn with_timeout(
107 base_url: impl Into<String>,
108 request_timeout: Duration,
109 ) -> ClientResult<Self> {
110 Self::with_timeouts(base_url, request_timeout, request_timeout)
111 }
112
113 pub fn with_timeouts(
121 base_url: impl Into<String>,
122 request_timeout: Duration,
123 stream_connect_timeout: Duration,
124 ) -> ClientResult<Self> {
125 Self::with_all_timeouts(
126 base_url,
127 request_timeout,
128 stream_connect_timeout,
129 Duration::from_secs(10),
130 )
131 }
132
133 pub fn with_all_timeouts(
142 base_url: impl Into<String>,
143 request_timeout: Duration,
144 stream_connect_timeout: Duration,
145 connection_timeout: Duration,
146 ) -> ClientResult<Self> {
147 let base_url = base_url.into();
148 if base_url.is_empty()
149 || (!base_url.starts_with("http://") && !base_url.starts_with("https://"))
150 {
151 return Err(ClientError::InvalidEndpoint(format!(
152 "invalid base URL: {base_url}"
153 )));
154 }
155
156 #[cfg(not(feature = "tls-rustls"))]
157 let client = {
158 let mut connector = HttpConnector::new();
159 connector.set_connect_timeout(Some(connection_timeout));
160 connector.set_nodelay(true);
161 Client::builder(TokioExecutor::new())
162 .pool_idle_timeout(Duration::from_secs(90))
163 .build(connector)
164 };
165
166 #[cfg(feature = "tls-rustls")]
167 let client = crate::tls::build_https_client_with_connect_timeout(
168 crate::tls::default_tls_config(),
169 connection_timeout,
170 );
171
172 Ok(Self {
173 inner: Arc::new(Inner {
174 client,
175 base_url: base_url.trim_end_matches('/').to_owned(),
176 request_timeout,
177 stream_connect_timeout,
178 max_response_size: super::DEFAULT_MAX_RESPONSE_SIZE,
179 }),
180 })
181 }
182
183 #[must_use]
189 pub fn with_max_response_size(mut self, max_bytes: usize) -> Self {
190 Arc::make_mut(&mut self.inner).max_response_size = max_bytes;
191 self
192 }
193
194 #[must_use]
196 pub fn base_url(&self) -> &str {
197 &self.inner.base_url
198 }
199}
200
201impl Transport for RestTransport {
202 fn send_request<'a>(
203 &'a self,
204 method: &'a str,
205 params: serde_json::Value,
206 extra_headers: &'a HashMap<String, String>,
207 ) -> Pin<Box<dyn Future<Output = ClientResult<serde_json::Value>> + Send + 'a>> {
208 Box::pin(self.execute_request(method, params, extra_headers))
209 }
210
211 fn send_streaming_request<'a>(
212 &'a self,
213 method: &'a str,
214 params: serde_json::Value,
215 extra_headers: &'a HashMap<String, String>,
216 ) -> Pin<Box<dyn Future<Output = ClientResult<EventStream>> + Send + 'a>> {
217 Box::pin(self.execute_streaming_request(method, params, extra_headers))
218 }
219}
220
221#[cfg(test)]
222mod tests {
223 use super::*;
224
225 #[test]
226 fn rest_transport_rejects_invalid_url() {
227 assert!(RestTransport::new("not-a-url").is_err());
228 }
229
230 #[test]
231 fn rest_transport_stores_base_url() {
232 let t = RestTransport::new("http://localhost:9090").unwrap();
233 assert_eq!(t.base_url(), "http://localhost:9090");
234 }
235
236 #[tokio::test]
238 async fn send_request_via_trait_delegation() {
239 use http_body_util::Full;
240 use hyper::body::Bytes;
241
242 let response_body = r#"{"status":"ok","data":42}"#;
243 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
244 let addr = listener.local_addr().unwrap();
245
246 tokio::spawn(async move {
247 loop {
248 let (stream, _) = listener.accept().await.unwrap();
249 let io = hyper_util::rt::TokioIo::new(stream);
250 let body = response_body.to_owned();
251 tokio::spawn(async move {
252 let service = hyper::service::service_fn(move |_req| {
253 let body = body.clone();
254 async move {
255 Ok::<_, hyper::Error>(
256 hyper::Response::builder()
257 .status(200)
258 .header("content-type", "application/json")
259 .body(Full::new(Bytes::from(body)))
260 .unwrap(),
261 )
262 }
263 });
264 let _ = hyper_util::server::conn::auto::Builder::new(
265 hyper_util::rt::TokioExecutor::new(),
266 )
267 .serve_connection(io, service)
268 .await;
269 });
270 }
271 });
272
273 let url = format!("http://127.0.0.1:{}", addr.port());
274 let transport = RestTransport::new(&url).unwrap();
275 let dyn_transport: &dyn crate::transport::Transport = &transport;
276 let result = dyn_transport
277 .send_request("SendMessage", serde_json::json!({}), &HashMap::new())
278 .await;
279 assert!(result.is_ok(), "send_request via trait should succeed");
280 }
281
282 #[tokio::test]
284 async fn send_streaming_request_via_trait_delegation() {
285 use http_body_util::Full;
286 use hyper::body::Bytes;
287
288 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
289 let addr = listener.local_addr().unwrap();
290
291 tokio::spawn(async move {
292 loop {
293 let (stream, _) = listener.accept().await.unwrap();
294 let io = hyper_util::rt::TokioIo::new(stream);
295 tokio::spawn(async move {
296 let service = hyper::service::service_fn(|_req| async {
297 let sse_body = "data: {\"hello\":\"world\"}\n\n";
298 Ok::<_, hyper::Error>(
299 hyper::Response::builder()
300 .status(200)
301 .header("content-type", "text/event-stream")
302 .body(Full::new(Bytes::from(sse_body)))
303 .unwrap(),
304 )
305 });
306 let _ = hyper_util::server::conn::auto::Builder::new(
307 hyper_util::rt::TokioExecutor::new(),
308 )
309 .serve_connection(io, service)
310 .await;
311 });
312 }
313 });
314
315 let url = format!("http://127.0.0.1:{}", addr.port());
316 let transport = RestTransport::new(&url).unwrap();
317 let dyn_transport: &dyn crate::transport::Transport = &transport;
318 let result = dyn_transport
319 .send_streaming_request(
320 "SendStreamingMessage",
321 serde_json::json!({}),
322 &HashMap::new(),
323 )
324 .await;
325 assert!(
326 result.is_ok(),
327 "send_streaming_request via trait should succeed"
328 );
329 }
330}