a2a_server/
jsonrpc.rs

1// Copyright AGNTCY Contributors (https://github.com/agntcy)
2// SPDX-License-Identifier: Apache-2.0
3use std::sync::Arc;
4
5use a2a::*;
6use a2a_pb::protojson_conv::{self, ProtoJsonPayload};
7use axum::{Json, extract::State, http::StatusCode, response::IntoResponse};
8use futures::{StreamExt, stream::BoxStream};
9use serde_json::Value;
10
11use crate::handler::RequestHandler;
12use crate::middleware::ServiceParams;
13use crate::push_config_compat::json_value as compat_json_value;
14use crate::sse;
15
16/// Shared state for the JSON-RPC handler.
17pub struct JsonRpcState<H: RequestHandler> {
18    pub handler: Arc<H>,
19}
20
21impl<H: RequestHandler> Clone for JsonRpcState<H> {
22    fn clone(&self) -> Self {
23        JsonRpcState {
24            handler: self.handler.clone(),
25        }
26    }
27}
28
29/// Create an axum router for the JSON-RPC protocol binding.
30///
31/// All requests are dispatched to a single POST endpoint that routes
32/// by the `method` field in the JSON-RPC envelope.
33pub fn jsonrpc_router<H: RequestHandler>(handler: Arc<H>) -> axum::Router {
34    let state = JsonRpcState { handler };
35    axum::Router::new()
36        .route("/", axum::routing::post(handle_jsonrpc::<H>))
37        .with_state(state)
38}
39
40async fn handle_jsonrpc<H: RequestHandler>(
41    State(state): State<JsonRpcState<H>>,
42    Json(request): Json<JsonRpcRequest>,
43) -> impl IntoResponse {
44    let params = ServiceParams::new();
45    let id = request.id.clone();
46    let method = request.method.as_str();
47
48    if request.jsonrpc != "2.0" {
49        return error_response(id, A2AError::invalid_request("invalid jsonrpc version"));
50    }
51
52    if methods::is_streaming(method) {
53        return handle_streaming_request(&state, &params, &request).await;
54    }
55
56    handle_unary_request(&state, &params, &request).await
57}
58
59async fn handle_unary_request<H: RequestHandler>(
60    state: &JsonRpcState<H>,
61    params: &ServiceParams,
62    request: &JsonRpcRequest,
63) -> axum::response::Response {
64    let id = request.id.clone();
65    let raw_params = request.params.clone().unwrap_or(Value::Null);
66
67    let result: Result<Value, A2AError> = match request.method.as_str() {
68        methods::SEND_MESSAGE => match protojson_conv::from_value::<SendMessageRequest>(raw_params)
69        {
70            Ok(req) => state
71                .handler
72                .send_message(params, req)
73                .await
74                .and_then(|r| protojson_value(&r)),
75            Err(e) => Err(parse_error(e)),
76        },
77        methods::GET_TASK => match protojson_conv::from_value::<GetTaskRequest>(raw_params) {
78            Ok(req) => state
79                .handler
80                .get_task(params, req)
81                .await
82                .and_then(|r| protojson_value(&r)),
83            Err(e) => Err(parse_error(e)),
84        },
85        methods::LIST_TASKS => match protojson_conv::from_value::<ListTasksRequest>(raw_params) {
86            Ok(req) => state
87                .handler
88                .list_tasks(params, req)
89                .await
90                .and_then(|r| protojson_value(&r)),
91            Err(e) => Err(parse_error(e)),
92        },
93        methods::CANCEL_TASK => match protojson_conv::from_value::<CancelTaskRequest>(raw_params) {
94            Ok(req) => state
95                .handler
96                .cancel_task(params, req)
97                .await
98                .and_then(|r| protojson_value(&r)),
99            Err(e) => Err(parse_error(e)),
100        },
101        methods::CREATE_PUSH_CONFIG => match parse_create_push_config_request(raw_params) {
102            Ok(req) => state
103                .handler
104                .create_push_config(params, req)
105                .await
106                .and_then(|r| compat_json_value(&r)),
107            Err(e) => Err(parse_error(e)),
108        },
109        methods::GET_PUSH_CONFIG => {
110            match protojson_conv::from_value::<GetTaskPushNotificationConfigRequest>(raw_params) {
111                Ok(req) => state
112                    .handler
113                    .get_push_config(params, req)
114                    .await
115                    .and_then(|r| compat_json_value(&r)),
116                Err(e) => Err(parse_error(e)),
117            }
118        }
119        methods::LIST_PUSH_CONFIGS => {
120            match protojson_conv::from_value::<ListTaskPushNotificationConfigsRequest>(raw_params) {
121                Ok(req) => state
122                    .handler
123                    .list_push_configs(params, req)
124                    .await
125                    .and_then(|r| compat_json_value(&r.configs)),
126                Err(e) => Err(parse_error(e)),
127            }
128        }
129        methods::DELETE_PUSH_CONFIG => {
130            match protojson_conv::from_value::<DeleteTaskPushNotificationConfigRequest>(raw_params)
131            {
132                Ok(req) => state
133                    .handler
134                    .delete_push_config(params, req)
135                    .await
136                    .map(|_| Value::Null),
137                Err(e) => Err(parse_error(e)),
138            }
139        }
140        methods::GET_EXTENDED_AGENT_CARD => {
141            match protojson_conv::from_value::<GetExtendedAgentCardRequest>(raw_params) {
142                Ok(req) => state
143                    .handler
144                    .get_extended_agent_card(params, req)
145                    .await
146                    .and_then(|r| protojson_value(&r)),
147                Err(e) => Err(parse_error(e)),
148            }
149        }
150        "" => Err(A2AError::invalid_request("method is required")),
151        _ => Err(A2AError::method_not_found(&request.method)),
152    };
153
154    match result {
155        Ok(value) => {
156            let resp = JsonRpcResponse::success(id, value);
157            Json(resp).into_response()
158        }
159        Err(e) => error_response(id, e),
160    }
161}
162
163async fn handle_streaming_request<H: RequestHandler>(
164    state: &JsonRpcState<H>,
165    params: &ServiceParams,
166    request: &JsonRpcRequest,
167) -> axum::response::Response {
168    let id = request.id.clone();
169    let raw_params = request.params.clone().unwrap_or(Value::Null);
170
171    match request.method.as_str() {
172        methods::SEND_STREAMING_MESSAGE => {
173            match protojson_conv::from_value::<SendMessageRequest>(raw_params) {
174                Ok(req) => match state.handler.send_streaming_message(params, req).await {
175                    Ok(stream) => {
176                        sse::sse_jsonrpc_stream(id, protojson_stream(stream)).into_response()
177                    }
178                    Err(e) => error_response(id, e),
179                },
180                Err(e) => error_response(id, parse_error(e)),
181            }
182        }
183        methods::SUBSCRIBE_TO_TASK => {
184            match protojson_conv::from_value::<SubscribeToTaskRequest>(raw_params) {
185                Ok(req) => match state.handler.subscribe_to_task(params, req).await {
186                    Ok(stream) => {
187                        sse::sse_jsonrpc_stream(id, protojson_stream(stream)).into_response()
188                    }
189                    Err(e) => error_response(id, e),
190                },
191                Err(e) => error_response(id, parse_error(e)),
192            }
193        }
194        _ => error_response(id, A2AError::method_not_found(&request.method)),
195    }
196}
197
198fn error_response(id: JsonRpcId, err: A2AError) -> axum::response::Response {
199    let resp = JsonRpcResponse::error(id, err.to_jsonrpc_error());
200    (StatusCode::OK, Json(resp)).into_response()
201}
202
203fn protojson_value<T: ProtoJsonPayload>(value: &T) -> Result<Value, A2AError> {
204    protojson_conv::to_value(value)
205        .map_err(|e| A2AError::internal(format!("failed to serialize ProtoJSON payload: {e}")))
206}
207
208fn protojson_stream(
209    stream: BoxStream<'static, Result<StreamResponse, A2AError>>,
210) -> BoxStream<'static, Result<Value, A2AError>> {
211    Box::pin(stream.map(|item| {
212        item.and_then(|value| {
213            protojson_conv::to_value(&value).map_err(|e| {
214                A2AError::internal(format!("failed to serialize ProtoJSON stream payload: {e}"))
215            })
216        })
217    }))
218}
219
220fn parse_create_push_config_request(
221    raw_params: Value,
222) -> Result<CreateTaskPushNotificationConfigRequest, String> {
223    match protojson_conv::from_value::<CreateTaskPushNotificationConfigRequest>(raw_params.clone())
224    {
225        Ok(req) => Ok(req),
226        Err(protojson_error) => serde_json::from_value::<CreateTaskPushNotificationConfigRequest>(
227            raw_params,
228        )
229        .map_err(|serde_error| {
230            format!("{protojson_error}; nested request parse failed: {serde_error}")
231        }),
232    }
233}
234
235fn parse_error(e: impl std::fmt::Display) -> A2AError {
236    A2AError {
237        code: error_code::PARSE_ERROR,
238        message: format!("invalid params: {e}"),
239        details: None,
240    }
241}
242
243#[cfg(test)]
244mod tests {
245    use super::*;
246    use crate::InMemoryPushConfigStore;
247    use crate::executor::ExecutorContext;
248    use crate::handler::DefaultRequestHandler;
249    use crate::task_store::InMemoryTaskStore;
250    use a2a_pb::protojson_conv;
251    use axum::body::Body;
252    use axum::http::Request;
253    use futures::stream::BoxStream;
254    use http_body_util::BodyExt;
255    use tower::ServiceExt;
256
257    struct EchoExecutor;
258
259    impl crate::AgentExecutor for EchoExecutor {
260        fn execute(
261            &self,
262            ctx: ExecutorContext,
263        ) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
264            let task = Task {
265                id: ctx.task_id.clone(),
266                context_id: ctx.context_id.clone(),
267                status: TaskStatus {
268                    state: TaskState::Completed,
269                    message: ctx.message,
270                    timestamp: None,
271                },
272                artifacts: None,
273                history: None,
274                metadata: None,
275            };
276            Box::pin(futures::stream::once(async move {
277                Ok(StreamResponse::Task(task))
278            }))
279        }
280
281        fn cancel(
282            &self,
283            ctx: ExecutorContext,
284        ) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
285            let task = Task {
286                id: ctx.task_id.clone(),
287                context_id: ctx.context_id.clone(),
288                status: TaskStatus {
289                    state: TaskState::Canceled,
290                    message: None,
291                    timestamp: None,
292                },
293                artifacts: None,
294                history: None,
295                metadata: None,
296            };
297            Box::pin(futures::stream::once(async move {
298                Ok(StreamResponse::Task(task))
299            }))
300        }
301    }
302
303    fn make_app() -> axum::Router {
304        let handler = Arc::new(DefaultRequestHandler::new(
305            EchoExecutor,
306            InMemoryTaskStore::new(),
307        ));
308        jsonrpc_router(handler)
309    }
310
311    fn make_push_app() -> axum::Router {
312        let handler = Arc::new(
313            DefaultRequestHandler::new(EchoExecutor, InMemoryTaskStore::new())
314                .with_push_config_store(InMemoryPushConfigStore::new()),
315        );
316        jsonrpc_router(handler)
317    }
318
319    async fn post_jsonrpc(app: axum::Router, method: &str, params: Value) -> JsonRpcResponse {
320        let rpc = JsonRpcRequest::new(JsonRpcId::Number(1), method, Some(params));
321        let body = serde_json::to_string(&rpc).unwrap();
322        let req = Request::builder()
323            .uri("/")
324            .method("POST")
325            .header("content-type", "application/json")
326            .body(Body::from(body))
327            .unwrap();
328        let resp = app.oneshot(req).await.unwrap();
329        let body = resp.into_body().collect().await.unwrap().to_bytes();
330        serde_json::from_slice(&body).unwrap()
331    }
332
333    #[tokio::test]
334    async fn test_send_message() {
335        let app = make_app();
336        let params = serde_json::json!({
337            "message": {
338                "messageId": "m1",
339                "role": "ROLE_USER",
340                "parts": [{"text": "hello"}]
341            }
342        });
343        let resp = post_jsonrpc(app, methods::SEND_MESSAGE, params).await;
344        assert!(resp.error.is_none(), "unexpected error: {:?}", resp.error);
345        let result =
346            protojson_conv::from_value::<SendMessageResponse>(resp.result.unwrap()).unwrap();
347        assert!(matches!(result, SendMessageResponse::Task(_)));
348    }
349
350    #[tokio::test]
351    async fn test_send_message_legacy_method_rejected() {
352        let app = make_app();
353        let params = serde_json::json!({
354            "message": {
355                "messageId": "m1",
356                "role": "ROLE_USER",
357                "parts": [{"text": "hello"}]
358            }
359        });
360        let resp = post_jsonrpc(app, "message.send", params).await;
361        assert!(resp.error.is_some(), "unexpected result: {:?}", resp.result);
362        assert_eq!(resp.error.unwrap().code, error_code::METHOD_NOT_FOUND);
363    }
364
365    #[tokio::test]
366    async fn test_get_task_not_found() {
367        let app = make_app();
368        let params = serde_json::json!({"id": "nonexistent"});
369        let resp = post_jsonrpc(app, methods::GET_TASK, params).await;
370        assert!(resp.error.is_some());
371        assert_eq!(resp.error.unwrap().code, error_code::TASK_NOT_FOUND);
372    }
373
374    #[tokio::test]
375    async fn test_invalid_method() {
376        let app = make_app();
377        let resp = post_jsonrpc(app, "unknown.method", Value::Null).await;
378        assert!(resp.error.is_some());
379        assert_eq!(resp.error.unwrap().code, error_code::METHOD_NOT_FOUND);
380    }
381
382    #[tokio::test]
383    async fn test_empty_method() {
384        let app = make_app();
385        let resp = post_jsonrpc(app, "", Value::Null).await;
386        assert!(resp.error.is_some());
387    }
388
389    #[tokio::test]
390    async fn test_invalid_jsonrpc_version() {
391        let app = make_app();
392        let rpc = serde_json::json!({
393            "jsonrpc": "1.0",
394            "id": 1,
395            "method": methods::SEND_MESSAGE,
396            "params": {}
397        });
398        let body = serde_json::to_string(&rpc).unwrap();
399        let req = Request::builder()
400            .uri("/")
401            .method("POST")
402            .header("content-type", "application/json")
403            .body(Body::from(body))
404            .unwrap();
405        let resp = app.oneshot(req).await.unwrap();
406        let body = resp.into_body().collect().await.unwrap().to_bytes();
407        let rpc_resp: JsonRpcResponse = serde_json::from_slice(&body).unwrap();
408        assert!(rpc_resp.error.is_some());
409    }
410
411    #[tokio::test]
412    async fn test_invalid_params() {
413        let app = make_app();
414        let params = serde_json::json!({"bogus": true});
415        let resp = post_jsonrpc(app, methods::SEND_MESSAGE, params).await;
416        assert!(resp.error.is_some());
417        assert_eq!(resp.error.unwrap().code, error_code::PARSE_ERROR);
418    }
419
420    #[tokio::test]
421    async fn test_list_tasks() {
422        let app = make_app();
423        let params = serde_json::json!({});
424        let resp = post_jsonrpc(app, methods::LIST_TASKS, params).await;
425        let result = protojson_conv::from_value::<ListTasksResponse>(resp.result.unwrap()).unwrap();
426        assert!(result.tasks.is_empty());
427    }
428
429    #[tokio::test]
430    async fn test_cancel_task_not_found() {
431        let app = make_app();
432        let params = serde_json::json!({"id": "nonexistent"});
433        let resp = post_jsonrpc(app, methods::CANCEL_TASK, params).await;
434        assert!(resp.error.is_some());
435    }
436
437    #[tokio::test]
438    async fn test_create_push_config() {
439        let app = make_push_app();
440        let params = serde_json::json!({
441            "taskId": "t1",
442            "id": "cfg1",
443            "url": "http://example.com/callback"
444        });
445        let resp = post_jsonrpc(app, methods::CREATE_PUSH_CONFIG, params).await;
446        assert!(resp.error.is_none(), "unexpected error: {:?}", resp.error);
447        let result =
448            serde_json::from_value::<TaskPushNotificationConfig>(resp.result.unwrap()).unwrap();
449        assert_eq!(result.task_id, "t1");
450        assert_eq!(result.config.id.as_deref(), Some("cfg1"));
451        assert_eq!(result.config.url, "http://example.com/callback");
452    }
453
454    #[tokio::test]
455    async fn test_create_push_config_accepts_nested_config_object() {
456        let app = make_push_app();
457        let params = serde_json::json!({
458            "taskId": "t1",
459            "config": {
460                "id": "cfg1",
461                "url": "http://example.com/callback"
462            }
463        });
464        let resp = post_jsonrpc(app, methods::CREATE_PUSH_CONFIG, params).await;
465        assert!(resp.error.is_none(), "unexpected error: {:?}", resp.error);
466        let result =
467            serde_json::from_value::<TaskPushNotificationConfig>(resp.result.unwrap()).unwrap();
468        assert_eq!(result.task_id, "t1");
469        assert_eq!(result.config.id.as_deref(), Some("cfg1"));
470        assert_eq!(result.config.url, "http://example.com/callback");
471    }
472
473    #[tokio::test]
474    async fn test_get_push_config() {
475        let app = make_app();
476        let params = serde_json::json!({
477            "taskId": "t1",
478            "id": "cfg1"
479        });
480        let resp = post_jsonrpc(app, methods::GET_PUSH_CONFIG, params).await;
481        assert!(resp.error.is_some() || resp.result.is_some());
482    }
483
484    #[tokio::test]
485    async fn test_list_push_configs() {
486        let app = make_push_app();
487        let params = serde_json::json!({
488            "taskId": "t1"
489        });
490        let resp = post_jsonrpc(app, methods::LIST_PUSH_CONFIGS, params).await;
491        assert!(resp.error.is_none(), "unexpected error: {:?}", resp.error);
492        let result =
493            serde_json::from_value::<Vec<TaskPushNotificationConfig>>(resp.result.unwrap())
494                .unwrap();
495        assert!(result.is_empty());
496    }
497
498    #[tokio::test]
499    async fn test_delete_push_config() {
500        let app = make_app();
501        let params = serde_json::json!({
502            "taskId": "t1",
503            "id": "cfg1"
504        });
505        let resp = post_jsonrpc(app, methods::DELETE_PUSH_CONFIG, params).await;
506        assert!(resp.error.is_some() || resp.result.is_some());
507    }
508
509    #[tokio::test]
510    async fn test_get_extended_agent_card() {
511        let app = make_app();
512        let params = serde_json::json!({});
513        let resp = post_jsonrpc(app, methods::GET_EXTENDED_AGENT_CARD, params).await;
514        // DefaultRequestHandler returns NotSupported
515        assert!(resp.error.is_some() || resp.result.is_some());
516    }
517
518    #[tokio::test]
519    async fn test_streaming_send_message() {
520        let app = make_app();
521        let body = serde_json::json!({
522            "message": {
523                "messageId": "m1",
524                "role": "ROLE_USER",
525                "parts": [{"text": "hello"}]
526            }
527        });
528        let rpc = JsonRpcRequest::new(
529            JsonRpcId::Number(1),
530            methods::SEND_STREAMING_MESSAGE,
531            Some(body),
532        );
533        let req = Request::builder()
534            .uri("/")
535            .method("POST")
536            .header("content-type", "application/json")
537            .header("accept", "text/event-stream")
538            .body(Body::from(serde_json::to_string(&rpc).unwrap()))
539            .unwrap();
540        let resp = app.oneshot(req).await.unwrap();
541        // SSE response => 200
542        assert_eq!(resp.status(), StatusCode::OK);
543    }
544
545    #[tokio::test]
546    async fn test_streaming_subscribe_to_task() {
547        let app = make_app();
548        let rpc = JsonRpcRequest::new(
549            JsonRpcId::Number(1),
550            methods::SUBSCRIBE_TO_TASK,
551            Some(serde_json::json!({"id": "t1"})),
552        );
553        let req = Request::builder()
554            .uri("/")
555            .method("POST")
556            .header("content-type", "application/json")
557            .body(Body::from(serde_json::to_string(&rpc).unwrap()))
558            .unwrap();
559        let resp = app.oneshot(req).await.unwrap();
560        // Subscription may fail (task not found), but routing should work
561        assert_eq!(resp.status(), StatusCode::OK);
562    }
563
564    #[tokio::test]
565    async fn test_streaming_send_message_legacy_method_rejected() {
566        let app = make_app();
567        let body = serde_json::json!({
568            "message": {
569                "messageId": "m1",
570                "role": "ROLE_USER",
571                "parts": [{"text": "hello"}]
572            }
573        });
574        let rpc = JsonRpcRequest::new(JsonRpcId::Number(1), "message.stream", Some(body));
575        let req = Request::builder()
576            .uri("/")
577            .method("POST")
578            .header("content-type", "application/json")
579            .header("accept", "text/event-stream")
580            .body(Body::from(serde_json::to_string(&rpc).unwrap()))
581            .unwrap();
582        let resp = app.oneshot(req).await.unwrap();
583        let body = resp.into_body().collect().await.unwrap().to_bytes();
584        let rpc_resp: JsonRpcResponse = serde_json::from_slice(&body).unwrap();
585        assert!(
586            rpc_resp.error.is_some(),
587            "unexpected result: {:?}",
588            rpc_resp.result
589        );
590        assert_eq!(rpc_resp.error.unwrap().code, error_code::METHOD_NOT_FOUND);
591    }
592
593    #[tokio::test]
594    async fn test_streaming_invalid_method() {
595        let app = make_app();
596        let rpc = serde_json::json!({
597            "jsonrpc": "2.0",
598            "id": 1,
599            "method": "tasks.subscribe",
600            "params": {"bogus": true}
601        });
602        let req = Request::builder()
603            .uri("/")
604            .method("POST")
605            .header("content-type", "application/json")
606            .body(Body::from(serde_json::to_string(&rpc).unwrap()))
607            .unwrap();
608        let resp = app.oneshot(req).await.unwrap();
609        assert_eq!(resp.status(), StatusCode::OK);
610    }
611}
a2a_server/jsonrpc.rs

a2a_server/
jsonrpc.rs
