1use std::collections::HashMap;
78use std::convert::Infallible;
79use std::sync::Arc;
80
81use axum::body::Body;
82use axum::extract::{Path, Query, State};
83use axum::response::IntoResponse;
84use axum::routing::{get, post};
85use axum::Router;
86use bytes::Bytes;
87
88use crate::handler::{RequestHandler, SendMessageResult};
89use crate::streaming::build_sse_response;
90
91pub struct A2aRouter {
116 handler: Arc<RequestHandler>,
117 config: super::DispatchConfig,
118}
119
120impl A2aRouter {
121 #[must_use]
123 pub fn new(handler: Arc<RequestHandler>) -> Self {
124 Self {
125 handler,
126 config: super::DispatchConfig::default(),
127 }
128 }
129
130 #[must_use]
132 pub const fn with_config(handler: Arc<RequestHandler>, config: super::DispatchConfig) -> Self {
133 Self { handler, config }
134 }
135
136 pub fn into_router(self) -> Router {
141 let max_body = self.config.max_request_body_size;
146 let state = A2aState {
147 handler: self.handler,
148 config: Arc::new(self.config),
149 };
150
151 Router::new()
152 .route("/message:send", post(handle_send_message))
154 .route("/message:stream", post(handle_stream_message))
155 .route("/tasks", get(handle_list_tasks))
157 .route("/tasks/{*rest}", axum::routing::any(handle_tasks_catchall))
162 .route("/extendedAgentCard", get(handle_extended_card))
164 .route("/.well-known/agent-card.json", get(handle_agent_card))
166 .route("/health", get(handle_health))
168 .with_state(state)
169 .layer(axum::extract::DefaultBodyLimit::max(max_body))
170 }
171}
172
173#[derive(Clone)]
176struct A2aState {
177 handler: Arc<RequestHandler>,
178 config: Arc<super::DispatchConfig>,
179}
180
181fn extract_headers(headers: &axum::http::HeaderMap) -> HashMap<String, String> {
184 headers
185 .iter()
186 .filter_map(|(k, v)| {
187 v.to_str()
188 .ok()
189 .map(|val| (k.as_str().to_lowercase(), val.to_owned()))
190 })
191 .collect()
192}
193
194fn a2a_error_to_response(err: &dyn std::fmt::Display, status: u16) -> axum::response::Response {
197 let body = serde_json::json!({ "error": err.to_string() });
198 (
199 axum::http::StatusCode::from_u16(status)
200 .unwrap_or(axum::http::StatusCode::INTERNAL_SERVER_ERROR),
201 axum::Json(body),
202 )
203 .into_response()
204}
205
206const fn server_error_status(err: &crate::error::ServerError) -> u16 {
207 use crate::error::ServerError;
208
209 match err {
210 ServerError::TaskNotFound(_) | ServerError::MethodNotFound(_) => 404,
211 ServerError::InvalidParams(_) | ServerError::Serialization(_) => 400,
212 ServerError::InvalidStateTransition { .. } | ServerError::TaskNotCancelable(_) => 409,
213 ServerError::PushNotSupported => 501,
214 ServerError::PayloadTooLarge(_) => 413,
215 ServerError::Overloaded(_) => 503,
218 _ => 500,
219 }
220}
221
222fn handler_error_to_response(err: &crate::error::ServerError) -> axum::response::Response {
223 a2a_error_to_response(err, server_error_status(err))
224}
225
226fn hyper_sse_to_axum(
231 resp: hyper::Response<http_body_util::combinators::BoxBody<Bytes, Infallible>>,
232) -> axum::response::Response {
233 let (parts, body) = resp.into_parts();
234 let axum_body = Body::new(body);
235 axum::response::Response::from_parts(parts, axum_body)
236}
237
238async fn handle_tasks_catchall(
251 State(state): State<A2aState>,
252 method: axum::http::Method,
253 Path(rest): Path<String>,
254 headers: axum::http::HeaderMap,
255 body: Bytes,
256) -> axum::response::Response {
257 let hdrs = extract_headers(&headers);
258 let segments: Vec<&str> = rest.split('/').filter(|s| !s.is_empty()).collect();
259
260 match (method.as_str(), segments.as_slice()) {
261 ("GET", [id]) if !id.contains(':') => handle_get_task_inner(&state, id, &hdrs).await,
263
264 ("POST", [id_action]) if id_action.ends_with(":cancel") => {
266 let id = &id_action[..id_action.len() - ":cancel".len()];
267 handle_cancel_task_inner(&state, id, &hdrs).await
268 }
269
270 ("GET" | "POST", [id_action]) if id_action.ends_with(":subscribe") => {
272 let id = &id_action[..id_action.len() - ":subscribe".len()];
273 handle_subscribe_inner(&state, id, &hdrs).await
274 }
275
276 ("POST", [task_id, "pushNotificationConfigs"]) => {
278 handle_create_push_config_inner(&state, task_id, &hdrs, body).await
279 }
280
281 ("GET", [task_id, "pushNotificationConfigs"]) => {
283 handle_list_push_configs_inner(&state, task_id, &hdrs).await
284 }
285
286 ("GET", [task_id, "pushNotificationConfigs", config_id]) => {
288 handle_get_push_config_inner(&state, task_id, config_id, &hdrs).await
289 }
290
291 ("DELETE", [task_id, "pushNotificationConfigs", config_id]) => {
293 handle_delete_push_config_inner(&state, task_id, config_id, &hdrs).await
294 }
295
296 _ => a2a_error_to_response(&"not found", 404),
297 }
298}
299
300async fn handle_send_message(
303 State(state): State<A2aState>,
304 headers: axum::http::HeaderMap,
305 body: Bytes,
306) -> axum::response::Response {
307 handle_send_inner(&state, false, &headers, body).await
308}
309
310async fn handle_stream_message(
311 State(state): State<A2aState>,
312 headers: axum::http::HeaderMap,
313 body: Bytes,
314) -> axum::response::Response {
315 handle_send_inner(&state, true, &headers, body).await
316}
317
318async fn handle_list_tasks(
319 State(state): State<A2aState>,
320 Query(query): Query<HashMap<String, String>>,
321 headers: axum::http::HeaderMap,
322) -> axum::response::Response {
323 let hdrs = extract_headers(&headers);
324 let params = a2a_protocol_types::params::ListTasksParams {
325 tenant: None,
326 context_id: query.get("contextId").cloned(),
327 status: query
328 .get("status")
329 .and_then(|s| serde_json::from_value(serde_json::Value::String(s.clone())).ok()),
330 page_size: query.get("pageSize").and_then(|v| v.parse().ok()),
331 page_token: query.get("pageToken").cloned(),
332 status_timestamp_after: query.get("statusTimestampAfter").cloned(),
333 include_artifacts: query.get("includeArtifacts").and_then(|v| v.parse().ok()),
334 history_length: query.get("historyLength").and_then(|v| v.parse().ok()),
335 };
336 match state.handler.on_list_tasks(params, Some(&hdrs)).await {
337 Ok(result) => axum::Json(result).into_response(),
338 Err(e) => handler_error_to_response(&e),
339 }
340}
341
342async fn handle_extended_card(
343 State(state): State<A2aState>,
344 headers: axum::http::HeaderMap,
345) -> axum::response::Response {
346 let hdrs = extract_headers(&headers);
347 match state.handler.on_get_extended_agent_card(Some(&hdrs)).await {
348 Ok(card) => axum::Json(card).into_response(),
349 Err(e) => handler_error_to_response(&e),
350 }
351}
352
353async fn handle_agent_card(State(state): State<A2aState>) -> axum::response::Response {
354 state.handler.agent_card.as_ref().map_or_else(
355 || a2a_error_to_response(&"agent card not configured", 404),
356 |card| axum::Json(card).into_response(),
357 )
358}
359
360async fn handle_health() -> axum::response::Response {
361 axum::Json(serde_json::json!({"status": "ok"})).into_response()
362}
363
364async fn handle_send_inner(
367 state: &A2aState,
368 streaming: bool,
369 headers: &axum::http::HeaderMap,
370 body: Bytes,
371) -> axum::response::Response {
372 let hdrs = extract_headers(headers);
373 let params: a2a_protocol_types::params::MessageSendParams = match serde_json::from_slice(&body)
374 {
375 Ok(p) => p,
376 Err(e) => return a2a_error_to_response(&e, 400),
377 };
378 match state
379 .handler
380 .on_send_message(params, streaming, Some(&hdrs))
381 .await
382 {
383 Ok(SendMessageResult::Response(resp)) => axum::Json(resp).into_response(),
384 Ok(SendMessageResult::Stream(reader)) => hyper_sse_to_axum(build_sse_response(
385 reader,
386 Some(state.config.sse_keep_alive_interval),
387 Some(state.config.sse_channel_capacity),
388 None, )),
390 Err(e) => handler_error_to_response(&e),
391 }
392}
393
394async fn handle_get_task_inner(
395 state: &A2aState,
396 id: &str,
397 hdrs: &HashMap<String, String>,
398) -> axum::response::Response {
399 let params = a2a_protocol_types::params::TaskQueryParams {
400 tenant: None,
401 id: id.to_owned(),
402 history_length: None,
403 };
404 match state.handler.on_get_task(params, Some(hdrs)).await {
405 Ok(task) => axum::Json(task).into_response(),
406 Err(e) => handler_error_to_response(&e),
407 }
408}
409
410async fn handle_cancel_task_inner(
411 state: &A2aState,
412 id: &str,
413 hdrs: &HashMap<String, String>,
414) -> axum::response::Response {
415 let params = a2a_protocol_types::params::CancelTaskParams {
416 tenant: None,
417 id: id.to_owned(),
418 metadata: None,
419 };
420 match state.handler.on_cancel_task(params, Some(hdrs)).await {
421 Ok(task) => axum::Json(task).into_response(),
422 Err(e) => handler_error_to_response(&e),
423 }
424}
425
426async fn handle_subscribe_inner(
427 state: &A2aState,
428 id: &str,
429 hdrs: &HashMap<String, String>,
430) -> axum::response::Response {
431 let params = a2a_protocol_types::params::TaskIdParams {
432 tenant: None,
433 id: id.to_owned(),
434 };
435 match state.handler.on_resubscribe(params, Some(hdrs)).await {
436 Ok(reader) => hyper_sse_to_axum(build_sse_response(
437 reader,
438 Some(state.config.sse_keep_alive_interval),
439 Some(state.config.sse_channel_capacity),
440 None, )),
442 Err(e) => handler_error_to_response(&e),
443 }
444}
445
446async fn handle_create_push_config_inner(
447 state: &A2aState,
448 task_id: &str,
449 hdrs: &HashMap<String, String>,
450 body: Bytes,
451) -> axum::response::Response {
452 let mut value: serde_json::Value = match serde_json::from_slice(&body) {
453 Ok(v) => v,
454 Err(e) => return a2a_error_to_response(&e, 400),
455 };
456 if let Some(obj) = value.as_object_mut() {
457 obj.entry("taskId")
458 .or_insert_with(|| serde_json::Value::String(task_id.to_owned()));
459 }
460 let config: a2a_protocol_types::push::TaskPushNotificationConfig =
461 match serde_json::from_value(value) {
462 Ok(c) => c,
463 Err(e) => return a2a_error_to_response(&e, 400),
464 };
465 match state.handler.on_set_push_config(config, Some(hdrs)).await {
466 Ok(result) => axum::Json(result).into_response(),
467 Err(e) => handler_error_to_response(&e),
468 }
469}
470
471async fn handle_get_push_config_inner(
472 state: &A2aState,
473 task_id: &str,
474 config_id: &str,
475 hdrs: &HashMap<String, String>,
476) -> axum::response::Response {
477 let params = a2a_protocol_types::params::GetPushConfigParams {
478 tenant: None,
479 task_id: task_id.to_owned(),
480 id: config_id.to_owned(),
481 };
482 match state.handler.on_get_push_config(params, Some(hdrs)).await {
483 Ok(config) => axum::Json(config).into_response(),
484 Err(e) => handler_error_to_response(&e),
485 }
486}
487
488async fn handle_list_push_configs_inner(
489 state: &A2aState,
490 task_id: &str,
491 hdrs: &HashMap<String, String>,
492) -> axum::response::Response {
493 match state
494 .handler
495 .on_list_push_configs(task_id, None, Some(hdrs))
496 .await
497 {
498 Ok(configs) => {
499 let resp = a2a_protocol_types::responses::ListPushConfigsResponse {
500 configs,
501 next_page_token: None,
502 };
503 axum::Json(resp).into_response()
504 }
505 Err(e) => handler_error_to_response(&e),
506 }
507}
508
509async fn handle_delete_push_config_inner(
510 state: &A2aState,
511 task_id: &str,
512 config_id: &str,
513 hdrs: &HashMap<String, String>,
514) -> axum::response::Response {
515 let params = a2a_protocol_types::params::DeletePushConfigParams {
516 tenant: None,
517 task_id: task_id.to_owned(),
518 id: config_id.to_owned(),
519 };
520 match state
521 .handler
522 .on_delete_push_config(params, Some(hdrs))
523 .await
524 {
525 Ok(()) => axum::Json(serde_json::json!({})).into_response(),
526 Err(e) => handler_error_to_response(&e),
527 }
528}
529
530#[cfg(test)]
533mod tests {
534 use super::*;
535
536 fn catchall_state() -> A2aState {
552 let handler = Arc::new(
553 crate::builder::RequestHandlerBuilder::new({
554 struct Noop;
555 crate::agent_executor!(Noop, |_ctx, _q| async { Ok(()) });
556 Noop
557 })
558 .build()
559 .unwrap(),
560 );
561 A2aState {
562 handler,
563 config: Arc::new(super::super::DispatchConfig::default()),
564 }
565 }
566
567 async fn seed_task(state: &A2aState, id: &str) {
568 use a2a_protocol_types::task::{ContextId, Task, TaskId, TaskState, TaskStatus};
569 let task = Task {
570 id: TaskId::new(id),
571 context_id: ContextId::new("ctx"),
572 status: TaskStatus::new(TaskState::Submitted),
573 history: None,
574 artifacts: None,
575 metadata: None,
576 };
577 state.handler.task_store.save(&task).await.unwrap();
578 }
579
580 async fn dispatch_tail(state: &A2aState, method: &str, rest: &str) -> axum::http::StatusCode {
581 let response = handle_tasks_catchall(
582 State(state.clone()),
583 axum::http::Method::from_bytes(method.as_bytes()).unwrap(),
584 Path(rest.to_owned()),
585 axum::http::HeaderMap::new(),
586 Bytes::new(),
587 )
588 .await;
589 response.status()
590 }
591
592 #[tokio::test]
596 async fn catchall_routes_cancel_and_strips_the_suffix() {
597 let state = catchall_state();
598 seed_task(&state, "task-abc").await;
599
600 assert_eq!(
602 dispatch_tail(&state, "POST", "task-abc:cancel").await,
603 axum::http::StatusCode::OK,
604 "POST /tasks/task-abc:cancel must cancel task-abc"
605 );
606 assert_eq!(
609 dispatch_tail(&state, "POST", "missing-xyz:cancel").await,
610 axum::http::StatusCode::NOT_FOUND,
611 "an unknown task id must 404 rather than resolve to a truncated one"
612 );
613 }
614
615 #[tokio::test]
619 async fn catchall_routes_subscribe_and_strips_the_suffix() {
620 let state = catchall_state();
621 seed_task(&state, "task-abc").await;
622
623 assert_eq!(
639 dispatch_tail(&state, "GET", "task-abc:subscribe").await,
640 axum::http::StatusCode::OK,
641 "GET /tasks/task-abc:subscribe must subscribe to task-abc"
642 );
643 assert_eq!(
644 dispatch_tail(&state, "GET", "missing-xyz:subscribe").await,
645 axum::http::StatusCode::NOT_FOUND,
646 "subscribe on an unknown id must still 404"
647 );
648 }
649
650 #[tokio::test]
659 async fn catchall_post_without_a_colon_action_falls_through() {
660 let state = catchall_state();
661 seed_task(&state, "tid").await;
662
663 assert_eq!(
665 dispatch_tail(&state, "POST", "tidZZZZZZZ").await,
666 axum::http::StatusCode::NOT_FOUND,
667 "a POST with no colon action must not be routed to CancelTask"
668 );
669 assert_eq!(
671 dispatch_tail(&state, "POST", "tidZZZZZZZZZZ").await,
672 axum::http::StatusCode::NOT_FOUND,
673 "a POST with no colon action must not be routed to SubscribeToTask"
674 );
675 }
676
677 #[tokio::test]
681 async fn catchall_plain_get_does_not_swallow_colon_actions() {
682 let state = catchall_state();
683 seed_task(&state, "task-abc").await;
684
685 assert_eq!(
686 dispatch_tail(&state, "GET", "task-abc").await,
687 axum::http::StatusCode::OK,
688 "GET /tasks/task-abc must fetch the task"
689 );
690 assert_eq!(
694 dispatch_tail(&state, "POST", "task-abc:cancel").await,
695 axum::http::StatusCode::OK,
696 "a colon action must not be captured by the plain `GetTask` arm"
697 );
698 }
699
700 #[test]
701 fn extract_headers_lowercases_names() {
702 let mut map = axum::http::HeaderMap::new();
703 map.insert("X-Request-ID", "abc".parse().unwrap());
704 map.insert("content-type", "application/json".parse().unwrap());
705
706 let result = extract_headers(&map);
707 assert_eq!(result.get("x-request-id").unwrap(), "abc");
708 assert_eq!(result.get("content-type").unwrap(), "application/json");
709 }
710
711 #[test]
712 fn extract_headers_skips_non_utf8_values() {
713 let mut map = axum::http::HeaderMap::new();
714 map.insert("good", "valid".parse().unwrap());
715 let result = extract_headers(&map);
717 assert_eq!(result.len(), 1);
718 assert_eq!(result.get("good").unwrap(), "valid");
719 }
720
721 #[test]
722 fn extract_headers_empty_map() {
723 let map = axum::http::HeaderMap::new();
724 let result = extract_headers(&map);
725 assert!(result.is_empty());
726 }
727
728 #[test]
729 fn a2a_state_is_clone() {
730 fn assert_clone<T: Clone>() {}
731 assert_clone::<A2aState>();
732 }
733
734 #[test]
735 fn server_error_status_task_not_found() {
736 use crate::error::ServerError;
737 assert_eq!(
738 server_error_status(&ServerError::TaskNotFound("t".into())),
739 404
740 );
741 }
742
743 #[test]
744 fn server_error_status_method_not_found() {
745 use crate::error::ServerError;
746 assert_eq!(
747 server_error_status(&ServerError::MethodNotFound("m".into())),
748 404
749 );
750 }
751
752 #[test]
753 fn server_error_status_invalid_params() {
754 use crate::error::ServerError;
755 assert_eq!(
756 server_error_status(&ServerError::InvalidParams("p".into())),
757 400
758 );
759 }
760
761 #[test]
762 fn server_error_status_serialization() {
763 use crate::error::ServerError;
764 let err = ServerError::Serialization(serde_json::from_str::<String>("bad").unwrap_err());
765 assert_eq!(server_error_status(&err), 400);
766 }
767
768 #[test]
769 fn server_error_status_task_not_cancelable() {
770 use crate::error::ServerError;
771 assert_eq!(
772 server_error_status(&ServerError::TaskNotCancelable("t".into())),
773 409
774 );
775 }
776
777 #[test]
778 fn server_error_status_invalid_state_transition() {
779 use crate::error::ServerError;
780 let err = ServerError::InvalidStateTransition {
781 task_id: "t".into(),
782 from: a2a_protocol_types::task::TaskState::Working,
783 to: a2a_protocol_types::task::TaskState::Submitted,
784 };
785 assert_eq!(server_error_status(&err), 409);
786 }
787
788 #[test]
789 fn server_error_status_push_not_supported() {
790 use crate::error::ServerError;
791 assert_eq!(server_error_status(&ServerError::PushNotSupported), 501);
792 }
793
794 #[test]
795 fn server_error_status_payload_too_large() {
796 use crate::error::ServerError;
797 assert_eq!(
798 server_error_status(&ServerError::PayloadTooLarge("big".into())),
799 413
800 );
801 }
802
803 #[test]
804 fn server_error_status_overloaded() {
805 use crate::error::ServerError;
806 assert_eq!(
809 server_error_status(&ServerError::Overloaded("at capacity".into())),
810 503
811 );
812 }
813
814 #[test]
815 fn server_error_status_internal() {
816 use crate::error::ServerError;
817 assert_eq!(
818 server_error_status(&ServerError::Internal("oops".into())),
819 500
820 );
821 }
822
823 #[test]
824 fn a2a_error_to_response_returns_correct_status() {
825 let resp = a2a_error_to_response(&"test error", 400);
826 assert_eq!(resp.status().as_u16(), 400);
827 }
828
829 #[test]
830 fn a2a_error_to_response_returns_json_body() {
831 let resp = a2a_error_to_response(&"not found", 404);
832 assert_eq!(resp.status().as_u16(), 404);
833 }
834
835 #[test]
836 fn a2a_error_to_response_invalid_status_falls_back_to_500() {
837 let resp = a2a_error_to_response(&"bad status", 1000);
839 assert_eq!(resp.status().as_u16(), 500);
840 }
841
842 #[test]
843 fn handler_error_to_response_maps_correctly() {
844 use crate::error::ServerError;
845 let resp = handler_error_to_response(&ServerError::TaskNotFound("t1".into()));
846 assert_eq!(resp.status().as_u16(), 404);
847
848 let resp = handler_error_to_response(&ServerError::InvalidParams("bad".into()));
849 assert_eq!(resp.status().as_u16(), 400);
850
851 let resp = handler_error_to_response(&ServerError::Internal("oops".into()));
852 assert_eq!(resp.status().as_u16(), 500);
853 }
854
855 #[test]
856 fn a2a_router_new_creates_with_defaults() {
857 use crate::builder::RequestHandlerBuilder;
859
860 struct NoopExecutor;
861 impl crate::executor::AgentExecutor for NoopExecutor {
862 fn execute<'a>(
863 &'a self,
864 _ctx: &'a crate::request_context::RequestContext,
865 _queue: &'a dyn crate::streaming::EventQueueWriter,
866 ) -> std::pin::Pin<
867 Box<
868 dyn std::future::Future<Output = a2a_protocol_types::error::A2aResult<()>>
869 + Send
870 + 'a,
871 >,
872 > {
873 Box::pin(async { Ok(()) })
874 }
875 }
876
877 let handler = Arc::new(RequestHandlerBuilder::new(NoopExecutor).build().unwrap());
878 let router = A2aRouter::new(handler);
879 let _axum_router = router.into_router();
881 }
882
883 #[test]
884 fn a2a_router_with_config() {
885 use crate::builder::RequestHandlerBuilder;
886
887 struct NoopExecutor;
888 impl crate::executor::AgentExecutor for NoopExecutor {
889 fn execute<'a>(
890 &'a self,
891 _ctx: &'a crate::request_context::RequestContext,
892 _queue: &'a dyn crate::streaming::EventQueueWriter,
893 ) -> std::pin::Pin<
894 Box<
895 dyn std::future::Future<Output = a2a_protocol_types::error::A2aResult<()>>
896 + Send
897 + 'a,
898 >,
899 > {
900 Box::pin(async { Ok(()) })
901 }
902 }
903
904 let handler = Arc::new(RequestHandlerBuilder::new(NoopExecutor).build().unwrap());
905 let config =
906 super::super::DispatchConfig::default().with_max_request_body_size(8 * 1024 * 1024);
907 let router = A2aRouter::with_config(handler, config);
908 let _axum_router = router.into_router();
909 }
910}