1use std::collections::HashMap;
53use std::convert::Infallible;
54use std::sync::Arc;
55
56use axum::body::Body;
57use axum::extract::{Path, Query, State};
58use axum::response::IntoResponse;
59use axum::routing::{get, post};
60use axum::Router;
61use bytes::Bytes;
62
63use crate::handler::{RequestHandler, SendMessageResult};
64use crate::streaming::build_sse_response;
65
66pub struct A2aRouter {
91 handler: Arc<RequestHandler>,
92 config: super::DispatchConfig,
93}
94
95impl A2aRouter {
96 #[must_use]
98 pub fn new(handler: Arc<RequestHandler>) -> Self {
99 Self {
100 handler,
101 config: super::DispatchConfig::default(),
102 }
103 }
104
105 #[must_use]
107 pub const fn with_config(handler: Arc<RequestHandler>, config: super::DispatchConfig) -> Self {
108 Self { handler, config }
109 }
110
111 pub fn into_router(self) -> Router {
116 let max_body = self.config.max_request_body_size;
121 let state = A2aState {
122 handler: self.handler,
123 config: Arc::new(self.config),
124 };
125
126 Router::new()
127 .route("/message:send", post(handle_send_message))
129 .route("/message:stream", post(handle_stream_message))
130 .route("/tasks", get(handle_list_tasks))
132 .route("/tasks/{*rest}", axum::routing::any(handle_tasks_catchall))
137 .route("/extendedAgentCard", get(handle_extended_card))
139 .route("/.well-known/agent-card.json", get(handle_agent_card))
141 .route("/health", get(handle_health))
143 .with_state(state)
144 .layer(axum::extract::DefaultBodyLimit::max(max_body))
145 }
146}
147
148#[derive(Clone)]
151struct A2aState {
152 handler: Arc<RequestHandler>,
153 config: Arc<super::DispatchConfig>,
154}
155
156fn extract_headers(headers: &axum::http::HeaderMap) -> HashMap<String, String> {
159 headers
160 .iter()
161 .filter_map(|(k, v)| {
162 v.to_str()
163 .ok()
164 .map(|val| (k.as_str().to_lowercase(), val.to_owned()))
165 })
166 .collect()
167}
168
169fn a2a_error_to_response(err: &dyn std::fmt::Display, status: u16) -> axum::response::Response {
172 let body = serde_json::json!({ "error": err.to_string() });
173 (
174 axum::http::StatusCode::from_u16(status)
175 .unwrap_or(axum::http::StatusCode::INTERNAL_SERVER_ERROR),
176 axum::Json(body),
177 )
178 .into_response()
179}
180
181const fn server_error_status(err: &crate::error::ServerError) -> u16 {
182 use crate::error::ServerError;
183
184 match err {
185 ServerError::TaskNotFound(_) | ServerError::MethodNotFound(_) => 404,
186 ServerError::InvalidParams(_) | ServerError::Serialization(_) => 400,
187 ServerError::InvalidStateTransition { .. } | ServerError::TaskNotCancelable(_) => 409,
188 ServerError::PushNotSupported => 501,
189 ServerError::PayloadTooLarge(_) => 413,
190 ServerError::Overloaded(_) => 503,
193 _ => 500,
194 }
195}
196
197fn handler_error_to_response(err: &crate::error::ServerError) -> axum::response::Response {
198 a2a_error_to_response(err, server_error_status(err))
199}
200
201fn hyper_sse_to_axum(
206 resp: hyper::Response<http_body_util::combinators::BoxBody<Bytes, Infallible>>,
207) -> axum::response::Response {
208 let (parts, body) = resp.into_parts();
209 let axum_body = Body::new(body);
210 axum::response::Response::from_parts(parts, axum_body)
211}
212
213async fn handle_tasks_catchall(
226 State(state): State<A2aState>,
227 method: axum::http::Method,
228 Path(rest): Path<String>,
229 headers: axum::http::HeaderMap,
230 body: Bytes,
231) -> axum::response::Response {
232 let hdrs = extract_headers(&headers);
233 let segments: Vec<&str> = rest.split('/').filter(|s| !s.is_empty()).collect();
234
235 match (method.as_str(), segments.as_slice()) {
236 ("GET", [id]) if !id.contains(':') => handle_get_task_inner(&state, id, &hdrs).await,
238
239 ("POST", [id_action]) if id_action.ends_with(":cancel") => {
241 let id = &id_action[..id_action.len() - ":cancel".len()];
242 handle_cancel_task_inner(&state, id, &hdrs).await
243 }
244
245 ("GET" | "POST", [id_action]) if id_action.ends_with(":subscribe") => {
247 let id = &id_action[..id_action.len() - ":subscribe".len()];
248 handle_subscribe_inner(&state, id, &hdrs).await
249 }
250
251 ("POST", [task_id, "pushNotificationConfigs"]) => {
253 handle_create_push_config_inner(&state, task_id, &hdrs, body).await
254 }
255
256 ("GET", [task_id, "pushNotificationConfigs"]) => {
258 handle_list_push_configs_inner(&state, task_id, &hdrs).await
259 }
260
261 ("GET", [task_id, "pushNotificationConfigs", config_id]) => {
263 handle_get_push_config_inner(&state, task_id, config_id, &hdrs).await
264 }
265
266 ("DELETE", [task_id, "pushNotificationConfigs", config_id]) => {
268 handle_delete_push_config_inner(&state, task_id, config_id, &hdrs).await
269 }
270
271 _ => a2a_error_to_response(&"not found", 404),
272 }
273}
274
275async fn handle_send_message(
278 State(state): State<A2aState>,
279 headers: axum::http::HeaderMap,
280 body: Bytes,
281) -> axum::response::Response {
282 handle_send_inner(&state, false, &headers, body).await
283}
284
285async fn handle_stream_message(
286 State(state): State<A2aState>,
287 headers: axum::http::HeaderMap,
288 body: Bytes,
289) -> axum::response::Response {
290 handle_send_inner(&state, true, &headers, body).await
291}
292
293async fn handle_list_tasks(
294 State(state): State<A2aState>,
295 Query(query): Query<HashMap<String, String>>,
296 headers: axum::http::HeaderMap,
297) -> axum::response::Response {
298 let hdrs = extract_headers(&headers);
299 let params = a2a_protocol_types::params::ListTasksParams {
300 tenant: None,
301 context_id: query.get("contextId").cloned(),
302 status: query
303 .get("status")
304 .and_then(|s| serde_json::from_value(serde_json::Value::String(s.clone())).ok()),
305 page_size: query.get("pageSize").and_then(|v| v.parse().ok()),
306 page_token: query.get("pageToken").cloned(),
307 status_timestamp_after: query.get("statusTimestampAfter").cloned(),
308 include_artifacts: query.get("includeArtifacts").and_then(|v| v.parse().ok()),
309 history_length: query.get("historyLength").and_then(|v| v.parse().ok()),
310 };
311 match state.handler.on_list_tasks(params, Some(&hdrs)).await {
312 Ok(result) => axum::Json(result).into_response(),
313 Err(e) => handler_error_to_response(&e),
314 }
315}
316
317async fn handle_extended_card(
318 State(state): State<A2aState>,
319 headers: axum::http::HeaderMap,
320) -> axum::response::Response {
321 let hdrs = extract_headers(&headers);
322 match state.handler.on_get_extended_agent_card(Some(&hdrs)).await {
323 Ok(card) => axum::Json(card).into_response(),
324 Err(e) => handler_error_to_response(&e),
325 }
326}
327
328async fn handle_agent_card(State(state): State<A2aState>) -> axum::response::Response {
329 state.handler.agent_card.as_ref().map_or_else(
330 || a2a_error_to_response(&"agent card not configured", 404),
331 |card| axum::Json(card).into_response(),
332 )
333}
334
335async fn handle_health() -> axum::response::Response {
336 axum::Json(serde_json::json!({"status": "ok"})).into_response()
337}
338
339async fn handle_send_inner(
342 state: &A2aState,
343 streaming: bool,
344 headers: &axum::http::HeaderMap,
345 body: Bytes,
346) -> axum::response::Response {
347 let hdrs = extract_headers(headers);
348 let params: a2a_protocol_types::params::MessageSendParams = match serde_json::from_slice(&body)
349 {
350 Ok(p) => p,
351 Err(e) => return a2a_error_to_response(&e, 400),
352 };
353 match state
354 .handler
355 .on_send_message(params, streaming, Some(&hdrs))
356 .await
357 {
358 Ok(SendMessageResult::Response(resp)) => axum::Json(resp).into_response(),
359 Ok(SendMessageResult::Stream(reader)) => hyper_sse_to_axum(build_sse_response(
360 reader,
361 Some(state.config.sse_keep_alive_interval),
362 Some(state.config.sse_channel_capacity),
363 None, )),
365 Err(e) => handler_error_to_response(&e),
366 }
367}
368
369async fn handle_get_task_inner(
370 state: &A2aState,
371 id: &str,
372 hdrs: &HashMap<String, String>,
373) -> axum::response::Response {
374 let params = a2a_protocol_types::params::TaskQueryParams {
375 tenant: None,
376 id: id.to_owned(),
377 history_length: None,
378 };
379 match state.handler.on_get_task(params, Some(hdrs)).await {
380 Ok(task) => axum::Json(task).into_response(),
381 Err(e) => handler_error_to_response(&e),
382 }
383}
384
385async fn handle_cancel_task_inner(
386 state: &A2aState,
387 id: &str,
388 hdrs: &HashMap<String, String>,
389) -> axum::response::Response {
390 let params = a2a_protocol_types::params::CancelTaskParams {
391 tenant: None,
392 id: id.to_owned(),
393 metadata: None,
394 };
395 match state.handler.on_cancel_task(params, Some(hdrs)).await {
396 Ok(task) => axum::Json(task).into_response(),
397 Err(e) => handler_error_to_response(&e),
398 }
399}
400
401async fn handle_subscribe_inner(
402 state: &A2aState,
403 id: &str,
404 hdrs: &HashMap<String, String>,
405) -> axum::response::Response {
406 let params = a2a_protocol_types::params::TaskIdParams {
407 tenant: None,
408 id: id.to_owned(),
409 };
410 match state.handler.on_resubscribe(params, Some(hdrs)).await {
411 Ok(reader) => hyper_sse_to_axum(build_sse_response(
412 reader,
413 Some(state.config.sse_keep_alive_interval),
414 Some(state.config.sse_channel_capacity),
415 None, )),
417 Err(e) => handler_error_to_response(&e),
418 }
419}
420
421async fn handle_create_push_config_inner(
422 state: &A2aState,
423 task_id: &str,
424 hdrs: &HashMap<String, String>,
425 body: Bytes,
426) -> axum::response::Response {
427 let mut value: serde_json::Value = match serde_json::from_slice(&body) {
428 Ok(v) => v,
429 Err(e) => return a2a_error_to_response(&e, 400),
430 };
431 if let Some(obj) = value.as_object_mut() {
432 obj.entry("taskId")
433 .or_insert_with(|| serde_json::Value::String(task_id.to_owned()));
434 }
435 let config: a2a_protocol_types::push::TaskPushNotificationConfig =
436 match serde_json::from_value(value) {
437 Ok(c) => c,
438 Err(e) => return a2a_error_to_response(&e, 400),
439 };
440 match state.handler.on_set_push_config(config, Some(hdrs)).await {
441 Ok(result) => axum::Json(result).into_response(),
442 Err(e) => handler_error_to_response(&e),
443 }
444}
445
446async fn handle_get_push_config_inner(
447 state: &A2aState,
448 task_id: &str,
449 config_id: &str,
450 hdrs: &HashMap<String, String>,
451) -> axum::response::Response {
452 let params = a2a_protocol_types::params::GetPushConfigParams {
453 tenant: None,
454 task_id: task_id.to_owned(),
455 id: config_id.to_owned(),
456 };
457 match state.handler.on_get_push_config(params, Some(hdrs)).await {
458 Ok(config) => axum::Json(config).into_response(),
459 Err(e) => handler_error_to_response(&e),
460 }
461}
462
463async fn handle_list_push_configs_inner(
464 state: &A2aState,
465 task_id: &str,
466 hdrs: &HashMap<String, String>,
467) -> axum::response::Response {
468 match state
469 .handler
470 .on_list_push_configs(task_id, None, Some(hdrs))
471 .await
472 {
473 Ok(configs) => {
474 let resp = a2a_protocol_types::responses::ListPushConfigsResponse {
475 configs,
476 next_page_token: None,
477 };
478 axum::Json(resp).into_response()
479 }
480 Err(e) => handler_error_to_response(&e),
481 }
482}
483
484async fn handle_delete_push_config_inner(
485 state: &A2aState,
486 task_id: &str,
487 config_id: &str,
488 hdrs: &HashMap<String, String>,
489) -> axum::response::Response {
490 let params = a2a_protocol_types::params::DeletePushConfigParams {
491 tenant: None,
492 task_id: task_id.to_owned(),
493 id: config_id.to_owned(),
494 };
495 match state
496 .handler
497 .on_delete_push_config(params, Some(hdrs))
498 .await
499 {
500 Ok(()) => axum::Json(serde_json::json!({})).into_response(),
501 Err(e) => handler_error_to_response(&e),
502 }
503}
504
505#[cfg(test)]
508mod tests {
509 use super::*;
510
511 #[test]
512 fn extract_headers_lowercases_names() {
513 let mut map = axum::http::HeaderMap::new();
514 map.insert("X-Request-ID", "abc".parse().unwrap());
515 map.insert("content-type", "application/json".parse().unwrap());
516
517 let result = extract_headers(&map);
518 assert_eq!(result.get("x-request-id").unwrap(), "abc");
519 assert_eq!(result.get("content-type").unwrap(), "application/json");
520 }
521
522 #[test]
523 fn extract_headers_skips_non_utf8_values() {
524 let mut map = axum::http::HeaderMap::new();
525 map.insert("good", "valid".parse().unwrap());
526 let result = extract_headers(&map);
528 assert_eq!(result.len(), 1);
529 assert_eq!(result.get("good").unwrap(), "valid");
530 }
531
532 #[test]
533 fn extract_headers_empty_map() {
534 let map = axum::http::HeaderMap::new();
535 let result = extract_headers(&map);
536 assert!(result.is_empty());
537 }
538
539 #[test]
540 fn a2a_state_is_clone() {
541 fn assert_clone<T: Clone>() {}
542 assert_clone::<A2aState>();
543 }
544
545 #[test]
546 fn server_error_status_task_not_found() {
547 use crate::error::ServerError;
548 assert_eq!(
549 server_error_status(&ServerError::TaskNotFound("t".into())),
550 404
551 );
552 }
553
554 #[test]
555 fn server_error_status_method_not_found() {
556 use crate::error::ServerError;
557 assert_eq!(
558 server_error_status(&ServerError::MethodNotFound("m".into())),
559 404
560 );
561 }
562
563 #[test]
564 fn server_error_status_invalid_params() {
565 use crate::error::ServerError;
566 assert_eq!(
567 server_error_status(&ServerError::InvalidParams("p".into())),
568 400
569 );
570 }
571
572 #[test]
573 fn server_error_status_serialization() {
574 use crate::error::ServerError;
575 let err = ServerError::Serialization(serde_json::from_str::<String>("bad").unwrap_err());
576 assert_eq!(server_error_status(&err), 400);
577 }
578
579 #[test]
580 fn server_error_status_task_not_cancelable() {
581 use crate::error::ServerError;
582 assert_eq!(
583 server_error_status(&ServerError::TaskNotCancelable("t".into())),
584 409
585 );
586 }
587
588 #[test]
589 fn server_error_status_invalid_state_transition() {
590 use crate::error::ServerError;
591 let err = ServerError::InvalidStateTransition {
592 task_id: "t".into(),
593 from: a2a_protocol_types::task::TaskState::Working,
594 to: a2a_protocol_types::task::TaskState::Submitted,
595 };
596 assert_eq!(server_error_status(&err), 409);
597 }
598
599 #[test]
600 fn server_error_status_push_not_supported() {
601 use crate::error::ServerError;
602 assert_eq!(server_error_status(&ServerError::PushNotSupported), 501);
603 }
604
605 #[test]
606 fn server_error_status_payload_too_large() {
607 use crate::error::ServerError;
608 assert_eq!(
609 server_error_status(&ServerError::PayloadTooLarge("big".into())),
610 413
611 );
612 }
613
614 #[test]
615 fn server_error_status_overloaded() {
616 use crate::error::ServerError;
617 assert_eq!(
620 server_error_status(&ServerError::Overloaded("at capacity".into())),
621 503
622 );
623 }
624
625 #[test]
626 fn server_error_status_internal() {
627 use crate::error::ServerError;
628 assert_eq!(
629 server_error_status(&ServerError::Internal("oops".into())),
630 500
631 );
632 }
633
634 #[test]
635 fn a2a_error_to_response_returns_correct_status() {
636 let resp = a2a_error_to_response(&"test error", 400);
637 assert_eq!(resp.status().as_u16(), 400);
638 }
639
640 #[test]
641 fn a2a_error_to_response_returns_json_body() {
642 let resp = a2a_error_to_response(&"not found", 404);
643 assert_eq!(resp.status().as_u16(), 404);
644 }
645
646 #[test]
647 fn a2a_error_to_response_invalid_status_falls_back_to_500() {
648 let resp = a2a_error_to_response(&"bad status", 1000);
650 assert_eq!(resp.status().as_u16(), 500);
651 }
652
653 #[test]
654 fn handler_error_to_response_maps_correctly() {
655 use crate::error::ServerError;
656 let resp = handler_error_to_response(&ServerError::TaskNotFound("t1".into()));
657 assert_eq!(resp.status().as_u16(), 404);
658
659 let resp = handler_error_to_response(&ServerError::InvalidParams("bad".into()));
660 assert_eq!(resp.status().as_u16(), 400);
661
662 let resp = handler_error_to_response(&ServerError::Internal("oops".into()));
663 assert_eq!(resp.status().as_u16(), 500);
664 }
665
666 #[test]
667 fn a2a_router_new_creates_with_defaults() {
668 use crate::builder::RequestHandlerBuilder;
670
671 struct NoopExecutor;
672 impl crate::executor::AgentExecutor for NoopExecutor {
673 fn execute<'a>(
674 &'a self,
675 _ctx: &'a crate::request_context::RequestContext,
676 _queue: &'a dyn crate::streaming::EventQueueWriter,
677 ) -> std::pin::Pin<
678 Box<
679 dyn std::future::Future<Output = a2a_protocol_types::error::A2aResult<()>>
680 + Send
681 + 'a,
682 >,
683 > {
684 Box::pin(async { Ok(()) })
685 }
686 }
687
688 let handler = Arc::new(RequestHandlerBuilder::new(NoopExecutor).build().unwrap());
689 let router = A2aRouter::new(handler);
690 let _axum_router = router.into_router();
692 }
693
694 #[test]
695 fn a2a_router_with_config() {
696 use crate::builder::RequestHandlerBuilder;
697
698 struct NoopExecutor;
699 impl crate::executor::AgentExecutor for NoopExecutor {
700 fn execute<'a>(
701 &'a self,
702 _ctx: &'a crate::request_context::RequestContext,
703 _queue: &'a dyn crate::streaming::EventQueueWriter,
704 ) -> std::pin::Pin<
705 Box<
706 dyn std::future::Future<Output = a2a_protocol_types::error::A2aResult<()>>
707 + Send
708 + 'a,
709 >,
710 > {
711 Box::pin(async { Ok(()) })
712 }
713 }
714
715 let handler = Arc::new(RequestHandlerBuilder::new(NoopExecutor).build().unwrap());
716 let config =
717 super::super::DispatchConfig::default().with_max_request_body_size(8 * 1024 * 1024);
718 let router = A2aRouter::with_config(handler, config);
719 let _axum_router = router.into_router();
720 }
721}