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

a2a_server/
jsonrpc.rs
