1mod response;
13
14use std::collections::HashMap;
15use std::convert::Infallible;
16use std::sync::Arc;
17
18use bytes::Bytes;
19use http_body_util::combinators::BoxBody;
20use hyper::body::Incoming;
21
22use a2a_protocol_types::jsonrpc::{
23 JsonRpcError, JsonRpcErrorResponse, JsonRpcId, JsonRpcRequest, JsonRpcSuccessResponse,
24 JsonRpcVersion,
25};
26
27use crate::agent_card::StaticAgentCardHandler;
28use crate::dispatch::cors::CorsConfig;
29use crate::error::ServerError;
30use crate::handler::{RequestHandler, SendMessageResult};
31use crate::serve::Dispatcher;
32use crate::streaming::build_sse_response;
33
34use response::{
35 error_response, error_response_bytes, extract_headers, json_response, parse_error_response,
36 parse_params, read_body_limited, success_response, success_response_bytes,
37};
38
39pub struct JsonRpcDispatcher {
47 handler: Arc<RequestHandler>,
48 card_handler: Option<StaticAgentCardHandler>,
49 cors: Option<CorsConfig>,
50 config: super::DispatchConfig,
51}
52
53impl JsonRpcDispatcher {
54 #[must_use]
57 pub fn new(handler: Arc<RequestHandler>) -> Self {
58 Self::with_config(handler, super::DispatchConfig::default())
59 }
60
61 #[must_use]
63 pub fn with_config(handler: Arc<RequestHandler>, config: super::DispatchConfig) -> Self {
64 let card_handler = handler
65 .agent_card
66 .as_ref()
67 .and_then(|card| StaticAgentCardHandler::new(card).ok());
68 Self {
69 handler,
70 card_handler,
71 cors: None,
72 config,
73 }
74 }
75
76 #[must_use]
81 pub fn with_cors(mut self, cors: CorsConfig) -> Self {
82 self.cors = Some(cors);
83 self
84 }
85
86 pub async fn dispatch(
93 &self,
94 req: hyper::Request<Incoming>,
95 ) -> hyper::Response<BoxBody<Bytes, Infallible>> {
96 if req.method() == "OPTIONS" {
98 if let Some(ref cors) = self.cors {
99 return cors.preflight_response();
100 }
101 return json_response(204, Vec::new());
102 }
103
104 if req.method() == "GET" && req.uri().path() == "/.well-known/agent-card.json" {
107 let mut resp = self.card_handler.as_ref().map_or_else(
108 || json_response(404, br#"{"error":"agent card not configured"}"#.to_vec()),
109 |h| h.handle(&req).map(http_body_util::BodyExt::boxed),
110 );
111 if let Some(ref cors) = self.cors {
112 cors.apply_headers(&mut resp);
113 }
114 return resp;
115 }
116
117 let requested_extensions = req
122 .headers()
123 .get(a2a_protocol_types::A2A_EXTENSIONS_HEADER)
124 .and_then(|v| v.to_str().ok())
125 .map(str::to_owned);
126
127 let mut resp = self.dispatch_inner(req).await;
128 if let Some(hval) = self
129 .handler
130 .activated_extensions_header_value(requested_extensions.as_deref())
131 {
132 if let Ok(v) = hyper::header::HeaderValue::from_str(&hval) {
133 resp.headers_mut()
134 .insert(a2a_protocol_types::A2A_EXTENSIONS_HEADER, v);
135 }
136 }
137 if let Some(ref cors) = self.cors {
138 cors.apply_headers(&mut resp);
139 }
140 resp
141 }
142
143 #[allow(clippy::too_many_lines)]
145 async fn dispatch_inner(
146 &self,
147 req: hyper::Request<Incoming>,
148 ) -> hyper::Response<BoxBody<Bytes, Infallible>> {
149 if let Some(ct) = req.headers().get("content-type") {
151 let ct_str = ct.to_str().unwrap_or("");
152 if !ct_str.starts_with("application/json")
153 && !ct_str.starts_with(a2a_protocol_types::A2A_CONTENT_TYPE)
154 {
155 return parse_error_response(
156 None,
157 &format!("unsupported Content-Type: {ct_str}; expected application/json or application/a2a+json"),
158 );
159 }
160 }
161
162 let version_value = req
166 .headers()
167 .get(a2a_protocol_types::A2A_VERSION_HEADER)
168 .and_then(|v| v.to_str().ok());
169 if let Err(err) =
170 super::validate_version_header(version_value, self.config.require_version_header)
171 {
172 return error_response(None, &ServerError::Protocol(err));
173 }
174
175 let headers = extract_headers(req.headers());
177
178 let body_bytes = match read_body_limited(
180 req.into_body(),
181 self.config.max_request_body_size,
182 self.config.body_read_timeout,
183 )
184 .await
185 {
186 Ok(bytes) => bytes,
187 Err(msg) => return parse_error_response(None, &msg),
188 };
189
190 let raw: serde_json::Value = match serde_json::from_slice(&body_bytes) {
192 Ok(v) => v,
193 Err(e) => return parse_error_response(None, &e.to_string()),
194 };
195
196 if raw.is_array() {
197 let serde_json::Value::Array(items) = raw else {
199 unreachable!()
200 };
201 if items.is_empty() {
202 return parse_error_response(None, "empty batch request");
203 }
204 if items.len() > self.config.max_batch_size {
206 return parse_error_response(
207 None,
208 &format!(
209 "batch too large: {} requests exceeds {} limit",
210 items.len(),
211 self.config.max_batch_size
212 ),
213 );
214 }
215 let mut responses: Vec<serde_json::Value> = Vec::with_capacity(items.len());
216 for item in items {
217 let rpc_req: JsonRpcRequest = match serde_json::from_value(item) {
218 Ok(r) => r,
219 Err(e) => {
220 let err_resp = JsonRpcErrorResponse::new(
222 None,
223 JsonRpcError::new(
224 a2a_protocol_types::error::ErrorCode::ParseError.as_i32(),
225 format!("Parse error: {e}"),
226 ),
227 );
228 if let Ok(v) = serde_json::to_value(&err_resp) {
229 responses.push(v);
230 }
231 continue;
232 }
233 };
234 let resp_body = self.dispatch_single_request(&rpc_req, &headers).await;
235 if let Ok(v) = serde_json::from_slice::<serde_json::Value>(&resp_body) {
236 responses.push(v);
237 }
238 }
239 let body = serde_json::to_vec(&responses).unwrap_or_default();
240 json_response(200, body)
241 } else {
242 let rpc_req: JsonRpcRequest = match serde_json::from_value(raw) {
244 Ok(r) => r,
245 Err(e) => return parse_error_response(None, &e.to_string()),
246 };
247 self.dispatch_single_request_http(&rpc_req, &headers).await
248 }
249 }
250
251 #[allow(clippy::too_many_lines)]
255 async fn dispatch_single_request_http(
256 &self,
257 rpc_req: &JsonRpcRequest,
258 headers: &HashMap<String, String>,
259 ) -> hyper::Response<BoxBody<Bytes, Infallible>> {
260 let id = rpc_req.id.to_response_id();
261 trace_info!(method = %rpc_req.method, "dispatching JSON-RPC request");
262
263 match rpc_req.method.as_str() {
265 "SendStreamingMessage" => {
266 return self.dispatch_send_message(id, rpc_req, true, headers).await;
267 }
268 "SubscribeToTask" => {
269 return match parse_params::<a2a_protocol_types::params::TaskIdParams>(rpc_req) {
270 Ok(p) => match self.handler.on_resubscribe(p, Some(headers)).await {
271 Ok(reader) => build_sse_response(
272 reader,
273 Some(self.config.sse_keep_alive_interval),
274 Some(self.config.sse_channel_capacity),
275 Some(id.clone()),
278 ),
279 Err(e) => error_response(id, &e),
280 },
281 Err(e) => error_response(id, &e),
282 };
283 }
284 _ => {}
285 }
286
287 let body = self.dispatch_single_request(rpc_req, headers).await;
288 json_response(200, body)
289 }
290
291 #[allow(clippy::too_many_lines)]
295 async fn dispatch_single_request(
296 &self,
297 rpc_req: &JsonRpcRequest,
298 headers: &HashMap<String, String>,
299 ) -> Vec<u8> {
300 let id = rpc_req.id.to_response_id();
301
302 match rpc_req.method.as_str() {
303 "SendMessage" => {
304 match self
305 .dispatch_send_message_inner(id.clone(), rpc_req, false, headers)
306 .await
307 {
308 Ok(resp) => serde_json::to_vec(&resp).unwrap_or_default(),
309 Err(body) => body,
310 }
311 }
312 "SendStreamingMessage" => {
313 let err = ServerError::InvalidParams(
315 "SendStreamingMessage not supported in batch requests".into(),
316 );
317 let a2a_err = err.to_a2a_error();
318 let resp = JsonRpcErrorResponse::new(
319 id,
320 JsonRpcError::new(a2a_err.code.as_i32(), a2a_err.message),
321 );
322 serde_json::to_vec(&resp).unwrap_or_default()
323 }
324 "GetTask" => {
325 match parse_params::<a2a_protocol_types::params::TaskQueryParams>(rpc_req) {
326 Ok(p) => match self.handler.on_get_task(p, Some(headers)).await {
327 Ok(r) => success_response_bytes(id, &r),
328 Err(e) => error_response_bytes(id, &e),
329 },
330 Err(e) => error_response_bytes(id, &e),
331 }
332 }
333 "ListTasks" => {
334 match parse_params::<a2a_protocol_types::params::ListTasksParams>(rpc_req) {
335 Ok(p) => match self.handler.on_list_tasks(p, Some(headers)).await {
336 Ok(r) => success_response_bytes(id, &r),
337 Err(e) => error_response_bytes(id, &e),
338 },
339 Err(e) => error_response_bytes(id, &e),
340 }
341 }
342 "CancelTask" => {
343 match parse_params::<a2a_protocol_types::params::CancelTaskParams>(rpc_req) {
344 Ok(p) => match self.handler.on_cancel_task(p, Some(headers)).await {
345 Ok(r) => success_response_bytes(id, &r),
346 Err(e) => error_response_bytes(id, &e),
347 },
348 Err(e) => error_response_bytes(id, &e),
349 }
350 }
351 "SubscribeToTask" => {
352 let err = ServerError::InvalidParams(
353 "SubscribeToTask not supported in batch requests".into(),
354 );
355 error_response_bytes(id, &err)
356 }
357 "CreateTaskPushNotificationConfig" => {
358 match parse_params::<a2a_protocol_types::push::TaskPushNotificationConfig>(rpc_req)
359 {
360 Ok(p) => match self.handler.on_set_push_config(p, Some(headers)).await {
361 Ok(r) => success_response_bytes(id, &r),
362 Err(e) => error_response_bytes(id, &e),
363 },
364 Err(e) => error_response_bytes(id, &e),
365 }
366 }
367 "GetTaskPushNotificationConfig" => {
368 match parse_params::<a2a_protocol_types::params::GetPushConfigParams>(rpc_req) {
369 Ok(p) => match self.handler.on_get_push_config(p, Some(headers)).await {
370 Ok(r) => success_response_bytes(id, &r),
371 Err(e) => error_response_bytes(id, &e),
372 },
373 Err(e) => error_response_bytes(id, &e),
374 }
375 }
376 "ListTaskPushNotificationConfigs" => {
377 match parse_params::<a2a_protocol_types::params::ListPushConfigsParams>(rpc_req) {
378 Ok(p) => match self
379 .handler
380 .on_list_push_configs(&p.task_id, p.tenant.as_deref(), Some(headers))
381 .await
382 {
383 Ok(configs) => {
384 let resp = a2a_protocol_types::responses::ListPushConfigsResponse {
385 configs,
386 next_page_token: None,
387 };
388 success_response_bytes(id, &resp)
389 }
390 Err(e) => error_response_bytes(id, &e),
391 },
392 Err(e) => error_response_bytes(id, &e),
393 }
394 }
395 "DeleteTaskPushNotificationConfig" => {
396 match parse_params::<a2a_protocol_types::params::DeletePushConfigParams>(rpc_req) {
397 Ok(p) => match self.handler.on_delete_push_config(p, Some(headers)).await {
398 Ok(()) => success_response_bytes(id, &serde_json::json!({})),
399 Err(e) => error_response_bytes(id, &e),
400 },
401 Err(e) => error_response_bytes(id, &e),
402 }
403 }
404 "GetExtendedAgentCard" => {
405 match self.handler.on_get_extended_agent_card(Some(headers)).await {
406 Ok(r) => success_response_bytes(id, &r),
407 Err(e) => error_response_bytes(id, &e),
408 }
409 }
410 other => {
411 let err = ServerError::MethodNotFound(other.to_owned());
412 error_response_bytes(id, &err)
413 }
414 }
415 }
416
417 async fn dispatch_send_message_inner(
420 &self,
421 id: JsonRpcId,
422 rpc_req: &JsonRpcRequest,
423 streaming: bool,
424 headers: &HashMap<String, String>,
425 ) -> Result<JsonRpcSuccessResponse<serde_json::Value>, Vec<u8>> {
426 let params = match parse_params::<a2a_protocol_types::params::MessageSendParams>(rpc_req) {
427 Ok(p) => p,
428 Err(e) => return Err(error_response_bytes(id, &e)),
429 };
430 match self
431 .handler
432 .on_send_message(params, streaming, Some(headers))
433 .await
434 {
435 Ok(SendMessageResult::Response(resp)) => {
436 let result = serde_json::to_value(&resp).unwrap_or(serde_json::Value::Null);
437 Ok(JsonRpcSuccessResponse {
438 jsonrpc: JsonRpcVersion,
439 id,
440 result,
441 })
442 }
443 Ok(SendMessageResult::Stream(_)) => {
444 let err = ServerError::Internal("unexpected stream response".into());
446 Err(error_response_bytes(id, &err))
447 }
448 Err(e) => Err(error_response_bytes(id, &e)),
449 }
450 }
451
452 async fn dispatch_send_message(
453 &self,
454 id: JsonRpcId,
455 rpc_req: &JsonRpcRequest,
456 streaming: bool,
457 headers: &HashMap<String, String>,
458 ) -> hyper::Response<BoxBody<Bytes, Infallible>> {
459 let params = match parse_params::<a2a_protocol_types::params::MessageSendParams>(rpc_req) {
460 Ok(p) => p,
461 Err(e) => return error_response(id, &e),
462 };
463 match self
464 .handler
465 .on_send_message(params, streaming, Some(headers))
466 .await
467 {
468 Ok(SendMessageResult::Response(resp)) => success_response(id, &resp),
469 Ok(SendMessageResult::Stream(reader)) => build_sse_response(
470 reader,
471 Some(self.config.sse_keep_alive_interval),
472 Some(self.config.sse_channel_capacity),
473 Some(id.clone()),
475 ),
476 Err(e) => error_response(id, &e),
477 }
478 }
479}
480
481impl std::fmt::Debug for JsonRpcDispatcher {
482 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
483 f.debug_struct("JsonRpcDispatcher").finish()
484 }
485}
486
487impl Dispatcher for JsonRpcDispatcher {
490 fn dispatch(
491 &self,
492 req: hyper::Request<Incoming>,
493 ) -> std::pin::Pin<
494 Box<dyn std::future::Future<Output = crate::serve::DispatchResponse> + Send + '_>,
495 > {
496 Box::pin(self.dispatch(req))
497 }
498}