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::{
8    Json,
9    extract::State,
10    http::{HeaderMap, StatusCode},
11    response::IntoResponse,
12};
13use futures::{StreamExt, stream::BoxStream};
14use serde_json::Value;
15
16use crate::handler::RequestHandler;
17use crate::middleware::{ServiceParams, extract_service_params};
18use crate::sse;
19
20/// Shared state for the JSON-RPC handler.
21pub struct JsonRpcState<H: RequestHandler> {
22    pub handler: Arc<H>,
23}
24
25impl<H: RequestHandler> Clone for JsonRpcState<H> {
26    fn clone(&self) -> Self {
27        JsonRpcState {
28            handler: self.handler.clone(),
29        }
30    }
31}
32
33/// Largest request body either binding accepts.
34///
35/// Without an explicit bound a single large POST is read into memory. The
36/// limit matches a2a-go's `MaxSSETokenSize` (`internal/sse/sse.go`), so a
37/// payload one SDK accepts is not rejected by the other.
38pub const MAX_REQUEST_BODY_BYTES: usize = 10 * 1024 * 1024;
39
40/// Create an axum router for the JSON-RPC protocol binding.
41///
42/// All requests are dispatched to a single POST endpoint that routes
43/// by the `method` field in the JSON-RPC envelope.
44pub fn jsonrpc_router<H: RequestHandler>(handler: Arc<H>) -> axum::Router {
45    let state = JsonRpcState { handler };
46    axum::Router::new()
47        .route("/", axum::routing::post(handle_jsonrpc::<H>))
48        .layer(axum::extract::DefaultBodyLimit::max(MAX_REQUEST_BODY_BYTES))
49        .with_state(state)
50}
51
52async fn handle_jsonrpc<H: RequestHandler>(
53    State(state): State<JsonRpcState<H>>,
54    headers: HeaderMap,
55    Json(request): Json<JsonRpcRequest>,
56) -> impl IntoResponse {
57    let params = extract_service_params(&headers);
58    let id = request.id.clone();
59    let method = request.method.as_str();
60
61    if request.jsonrpc != "2.0" {
62        return error_response(id, A2AError::invalid_request("invalid jsonrpc version"));
63    }
64
65    if methods::is_streaming(method) {
66        return handle_streaming_request(&state, &params, &request).await;
67    }
68
69    handle_unary_request(&state, &params, &request).await
70}
71
72async fn handle_unary_request<H: RequestHandler>(
73    state: &JsonRpcState<H>,
74    params: &ServiceParams,
75    request: &JsonRpcRequest,
76) -> axum::response::Response {
77    let id = request.id.clone();
78    let raw_params = request.params.clone().unwrap_or(Value::Null);
79
80    let result: Result<Value, A2AError> = match request.method.as_str() {
81        methods::SEND_MESSAGE => match protojson_conv::from_value::<SendMessageRequest>(raw_params)
82        {
83            Ok(req) => state
84                .handler
85                .send_message(params, req)
86                .await
87                .and_then(|r| protojson_value(&r)),
88            Err(e) => Err(parse_error(e)),
89        },
90        methods::GET_TASK => match protojson_conv::from_value::<GetTaskRequest>(raw_params) {
91            Ok(req) => state
92                .handler
93                .get_task(params, req)
94                .await
95                .and_then(|r| protojson_value(&r)),
96            Err(e) => Err(parse_error(e)),
97        },
98        methods::LIST_TASKS => match protojson_conv::from_value::<ListTasksRequest>(raw_params) {
99            Ok(req) => state
100                .handler
101                .list_tasks(params, req)
102                .await
103                .and_then(|r| protojson_value(&r)),
104            Err(e) => Err(parse_error(e)),
105        },
106        methods::CANCEL_TASK => match protojson_conv::from_value::<CancelTaskRequest>(raw_params) {
107            Ok(req) => state
108                .handler
109                .cancel_task(params, req)
110                .await
111                .and_then(|r| protojson_value(&r)),
112            Err(e) => Err(parse_error(e)),
113        },
114        methods::CREATE_PUSH_CONFIG => match parse_create_push_config_request(raw_params) {
115            Ok(req) => state
116                .handler
117                .create_push_config(params, req)
118                .await
119                .and_then(|r| protojson_value(&r)),
120            Err(e) => Err(parse_error(e)),
121        },
122        methods::GET_PUSH_CONFIG => {
123            match protojson_conv::from_value::<GetTaskPushNotificationConfigRequest>(raw_params) {
124                Ok(req) => state
125                    .handler
126                    .get_push_config(params, req)
127                    .await
128                    .and_then(|r| protojson_value(&r)),
129                Err(e) => Err(parse_error(e)),
130            }
131        }
132        methods::LIST_PUSH_CONFIGS => {
133            match protojson_conv::from_value::<ListTaskPushNotificationConfigsRequest>(raw_params) {
134                Ok(req) => state
135                    .handler
136                    .list_push_configs(params, req)
137                    .await
138                    .and_then(|r| protojson_value(&r)),
139                Err(e) => Err(parse_error(e)),
140            }
141        }
142        methods::DELETE_PUSH_CONFIG => {
143            match protojson_conv::from_value::<DeleteTaskPushNotificationConfigRequest>(raw_params)
144            {
145                Ok(req) => state
146                    .handler
147                    .delete_push_config(params, req)
148                    .await
149                    .map(|_| Value::Null),
150                Err(e) => Err(parse_error(e)),
151            }
152        }
153        methods::GET_EXTENDED_AGENT_CARD => {
154            match protojson_conv::from_value::<GetExtendedAgentCardRequest>(raw_params) {
155                Ok(req) => state
156                    .handler
157                    .get_extended_agent_card(params, req)
158                    .await
159                    .and_then(|r| protojson_value(&r)),
160                Err(e) => Err(parse_error(e)),
161            }
162        }
163        "" => Err(A2AError::invalid_request("method is required")),
164        _ => Err(A2AError::method_not_found(&request.method)),
165    };
166
167    match result {
168        Ok(value) => {
169            let resp = JsonRpcResponse::success(id, value);
170            Json(resp).into_response()
171        }
172        Err(e) => error_response(id, e),
173    }
174}
175
176async fn handle_streaming_request<H: RequestHandler>(
177    state: &JsonRpcState<H>,
178    params: &ServiceParams,
179    request: &JsonRpcRequest,
180) -> axum::response::Response {
181    let id = request.id.clone();
182    let raw_params = request.params.clone().unwrap_or(Value::Null);
183
184    match request.method.as_str() {
185        methods::SEND_STREAMING_MESSAGE => {
186            match protojson_conv::from_value::<SendMessageRequest>(raw_params) {
187                Ok(req) => match state.handler.send_streaming_message(params, req).await {
188                    Ok(stream) => {
189                        sse::sse_jsonrpc_stream(id, protojson_stream(stream)).into_response()
190                    }
191                    Err(e) => error_response(id, e),
192                },
193                Err(e) => error_response(id, parse_error(e)),
194            }
195        }
196        methods::SUBSCRIBE_TO_TASK => {
197            match protojson_conv::from_value::<SubscribeToTaskRequest>(raw_params) {
198                Ok(req) => match state.handler.subscribe_to_task(params, req).await {
199                    Ok(stream) => {
200                        sse::sse_jsonrpc_stream(id, protojson_stream(stream)).into_response()
201                    }
202                    Err(e) => error_response(id, e),
203                },
204                Err(e) => error_response(id, parse_error(e)),
205            }
206        }
207        _ => error_response(id, A2AError::method_not_found(&request.method)),
208    }
209}
210
211fn error_response(id: JsonRpcId, err: A2AError) -> axum::response::Response {
212    let resp = JsonRpcResponse::error(id, err.to_jsonrpc_error());
213    (StatusCode::OK, Json(resp)).into_response()
214}
215
216fn protojson_value<T: ProtoJsonPayload>(value: &T) -> Result<Value, A2AError> {
217    protojson_conv::to_value(value)
218        .map_err(|e| A2AError::internal(format!("failed to serialize ProtoJSON payload: {e}")))
219}
220
221fn protojson_stream(
222    stream: BoxStream<'static, Result<StreamResponse, A2AError>>,
223) -> BoxStream<'static, Result<Value, A2AError>> {
224    Box::pin(stream.map(|item| {
225        item.and_then(|value| {
226            protojson_conv::to_value(&value).map_err(|e| {
227                A2AError::internal(format!("failed to serialize ProtoJSON stream payload: {e}"))
228            })
229        })
230    }))
231}
232
233fn parse_create_push_config_request(
234    raw_params: Value,
235) -> Result<TaskPushNotificationConfig, String> {
236    protojson_conv::from_value::<TaskPushNotificationConfig>(raw_params).map_err(|e| e.to_string())
237}
238
239fn parse_error(e: impl std::fmt::Display) -> A2AError {
240    A2AError {
241        code: error_code::PARSE_ERROR,
242        message: format!("invalid params: {e}"),
243        details: None,
244    }
245}
246
247#[cfg(test)]
248mod tests {
249    use super::*;
250    use crate::InMemoryPushConfigStore;
251    use crate::executor::ExecutorContext;
252    use crate::handler::DefaultRequestHandler;
253    use crate::task_store::InMemoryTaskStore;
254    use a2a_pb::protojson_conv;
255    use axum::body::Body;
256    use axum::http::Request;
257
258    use crate::test_util::{CapturingHandler, assert_header_captured, install_crypto_provider};
259    use futures::stream::BoxStream;
260    use http_body_util::BodyExt;
261    use tower::ServiceExt;
262
263    struct EchoExecutor;
264
265    impl crate::AgentExecutor for EchoExecutor {
266        fn execute(
267            &self,
268            ctx: ExecutorContext,
269        ) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
270            let task = Task {
271                id: ctx.task_id.clone(),
272                context_id: ctx.context_id.clone(),
273                status: TaskStatus {
274                    state: TaskState::Completed,
275                    message: ctx.message,
276                    timestamp: None,
277                },
278                artifacts: None,
279                history: None,
280                metadata: None,
281            };
282            Box::pin(futures::stream::once(async move {
283                Ok(StreamResponse::Task(task))
284            }))
285        }
286
287        fn cancel(
288            &self,
289            ctx: ExecutorContext,
290        ) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
291            let task = Task {
292                id: ctx.task_id.clone(),
293                context_id: ctx.context_id.clone(),
294                status: TaskStatus {
295                    state: TaskState::Canceled,
296                    message: None,
297                    timestamp: None,
298                },
299                artifacts: None,
300                history: None,
301                metadata: None,
302            };
303            Box::pin(futures::stream::once(async move {
304                Ok(StreamResponse::Task(task))
305            }))
306        }
307    }
308
309    fn make_app() -> axum::Router {
310        let handler = Arc::new(DefaultRequestHandler::new(
311            EchoExecutor,
312            InMemoryTaskStore::new(),
313        ));
314        jsonrpc_router(handler)
315    }
316
317    fn make_push_app() -> axum::Router {
318        install_crypto_provider();
319        let handler = Arc::new(
320            DefaultRequestHandler::new(EchoExecutor, InMemoryTaskStore::new())
321                .with_push_config_store(InMemoryPushConfigStore::new()),
322        );
323        jsonrpc_router(handler)
324    }
325
326    async fn post_jsonrpc(app: axum::Router, method: &str, params: Value) -> JsonRpcResponse {
327        let rpc = JsonRpcRequest::new(JsonRpcId::Number(1), method, Some(params));
328        let body = serde_json::to_string(&rpc).unwrap();
329        let req = Request::builder()
330            .uri("/")
331            .method("POST")
332            .header("content-type", "application/json")
333            .body(Body::from(body))
334            .unwrap();
335        let resp = app.oneshot(req).await.unwrap();
336        let body = resp.into_body().collect().await.unwrap().to_bytes();
337        serde_json::from_slice(&body).unwrap()
338    }
339
340    /// An unbounded body is read into memory; the limit rejects it at the
341    /// framework boundary instead, before any handler runs.
342    #[tokio::test]
343    async fn test_request_body_over_the_limit_is_rejected() {
344        let app = make_app();
345        let oversized = Body::from("x".repeat(MAX_REQUEST_BODY_BYTES + 1));
346        let req = Request::builder()
347            .uri("/")
348            .method("POST")
349            .header("content-type", "application/json")
350            .body(oversized)
351            .unwrap();
352
353        let resp = app.oneshot(req).await.unwrap();
354        assert_eq!(resp.status(), axum::http::StatusCode::PAYLOAD_TOO_LARGE);
355    }
356
357    #[tokio::test]
358    async fn test_send_message() {
359        let app = make_app();
360        let params = serde_json::json!({
361            "message": {
362                "messageId": "m1",
363                "role": "ROLE_USER",
364                "parts": [{"text": "hello"}]
365            }
366        });
367        let resp = post_jsonrpc(app, methods::SEND_MESSAGE, params).await;
368        assert!(resp.error.is_none(), "unexpected error: {:?}", resp.error);
369        let result =
370            protojson_conv::from_value::<SendMessageResponse>(resp.result.unwrap()).unwrap();
371        assert!(matches!(result, SendMessageResponse::Task(_)));
372    }
373
374    #[tokio::test]
375    async fn test_send_message_legacy_method_rejected() {
376        let app = make_app();
377        let params = serde_json::json!({
378            "message": {
379                "messageId": "m1",
380                "role": "ROLE_USER",
381                "parts": [{"text": "hello"}]
382            }
383        });
384        let resp = post_jsonrpc(app, "message.send", params).await;
385        assert!(resp.error.is_some(), "unexpected result: {:?}", resp.result);
386        assert_eq!(resp.error.unwrap().code, error_code::METHOD_NOT_FOUND);
387    }
388
389    #[tokio::test]
390    async fn test_get_task_not_found() {
391        let app = make_app();
392        let params = serde_json::json!({"id": "nonexistent"});
393        let resp = post_jsonrpc(app, methods::GET_TASK, params).await;
394        assert!(resp.error.is_some());
395        assert_eq!(resp.error.unwrap().code, error_code::TASK_NOT_FOUND);
396    }
397
398    #[tokio::test]
399    async fn test_invalid_method() {
400        let app = make_app();
401        let resp = post_jsonrpc(app, "unknown.method", Value::Null).await;
402        assert!(resp.error.is_some());
403        assert_eq!(resp.error.unwrap().code, error_code::METHOD_NOT_FOUND);
404    }
405
406    #[tokio::test]
407    async fn test_empty_method() {
408        let app = make_app();
409        let resp = post_jsonrpc(app, "", Value::Null).await;
410        assert!(resp.error.is_some());
411    }
412
413    #[tokio::test]
414    async fn test_invalid_jsonrpc_version() {
415        let app = make_app();
416        let rpc = serde_json::json!({
417            "jsonrpc": "1.0",
418            "id": 1,
419            "method": methods::SEND_MESSAGE,
420            "params": {}
421        });
422        let body = serde_json::to_string(&rpc).unwrap();
423        let req = Request::builder()
424            .uri("/")
425            .method("POST")
426            .header("content-type", "application/json")
427            .body(Body::from(body))
428            .unwrap();
429        let resp = app.oneshot(req).await.unwrap();
430        let body = resp.into_body().collect().await.unwrap().to_bytes();
431        let rpc_resp: JsonRpcResponse = serde_json::from_slice(&body).unwrap();
432        assert!(rpc_resp.error.is_some());
433    }
434
435    #[tokio::test]
436    async fn test_invalid_params() {
437        let app = make_app();
438        // Type mismatch: `message` must be an object. A merely-unknown field
439        // is not invalid params; it is ignored per spec §5.7.
440        let params = serde_json::json!({"message": "not-an-object"});
441        let resp = post_jsonrpc(app, methods::SEND_MESSAGE, params).await;
442        assert!(resp.error.is_some());
443        assert_eq!(resp.error.unwrap().code, error_code::PARSE_ERROR);
444    }
445
446    /// Spec §5.7: unrecognized fields in request params are ignored for
447    /// forward compatibility instead of being rejected.
448    #[tokio::test]
449    async fn test_unknown_params_fields_are_ignored() {
450        let app = make_app();
451        let params = serde_json::json!({
452            "message": {
453                "messageId": "m1",
454                "role": "ROLE_USER",
455                "parts": [{"text": "hi"}],
456                "futureField": true
457            },
458            "anotherFutureField": {"nested": 1}
459        });
460        let resp = post_jsonrpc(app, methods::SEND_MESSAGE, params).await;
461        assert!(
462            resp.error.is_none(),
463            "unknown fields must be ignored: {:?}",
464            resp.error
465        );
466        assert!(resp.result.is_some());
467    }
468
469    #[tokio::test]
470    async fn test_list_tasks() {
471        let app = make_app();
472        let params = serde_json::json!({});
473        let resp = post_jsonrpc(app, methods::LIST_TASKS, params).await;
474        let result = protojson_conv::from_value::<ListTasksResponse>(resp.result.unwrap()).unwrap();
475        assert!(result.tasks.is_empty());
476    }
477
478    #[tokio::test]
479    async fn test_cancel_task_not_found() {
480        let app = make_app();
481        let params = serde_json::json!({"id": "nonexistent"});
482        let resp = post_jsonrpc(app, methods::CANCEL_TASK, params).await;
483        assert!(resp.error.is_some());
484    }
485
486    #[tokio::test]
487    async fn test_create_push_config() {
488        let app = make_push_app();
489        let params = serde_json::json!({
490            "taskId": "t1",
491            "id": "cfg1",
492            "url": "http://example.com/callback"
493        });
494        let resp = post_jsonrpc(app, methods::CREATE_PUSH_CONFIG, params).await;
495        assert!(resp.error.is_none(), "unexpected error: {:?}", resp.error);
496        let result =
497            serde_json::from_value::<TaskPushNotificationConfig>(resp.result.unwrap()).unwrap();
498        assert_eq!(result.task_id, "t1");
499        assert_eq!(result.id.as_deref(), Some("cfg1"));
500        assert_eq!(result.url, "http://example.com/callback");
501    }
502
503    #[tokio::test]
504    async fn test_create_push_config_accepts_flat_config_object() {
505        let app = make_push_app();
506        let params = serde_json::json!({
507            "taskId": "t1",
508            "id": "cfg1",
509            "url": "http://example.com/callback"
510        });
511        let resp = post_jsonrpc(app, methods::CREATE_PUSH_CONFIG, params).await;
512        assert!(resp.error.is_none(), "unexpected error: {:?}", resp.error);
513        let result =
514            serde_json::from_value::<TaskPushNotificationConfig>(resp.result.unwrap()).unwrap();
515        assert_eq!(result.task_id, "t1");
516        assert_eq!(result.id.as_deref(), Some("cfg1"));
517        assert_eq!(result.url, "http://example.com/callback");
518    }
519
520    #[tokio::test]
521    async fn test_get_push_config() {
522        let app = make_app();
523        let params = serde_json::json!({
524            "taskId": "t1",
525            "id": "cfg1"
526        });
527        let resp = post_jsonrpc(app, methods::GET_PUSH_CONFIG, params).await;
528        assert!(resp.error.is_some() || resp.result.is_some());
529    }
530
531    #[tokio::test]
532    async fn test_list_push_configs() {
533        let app = make_push_app();
534        let params = serde_json::json!({
535            "taskId": "t1"
536        });
537        let resp = post_jsonrpc(app, methods::LIST_PUSH_CONFIGS, params).await;
538        assert!(resp.error.is_none(), "unexpected error: {:?}", resp.error);
539        let result =
540            serde_json::from_value::<ListTaskPushNotificationConfigsResponse>(resp.result.unwrap())
541                .unwrap();
542        assert!(result.configs.is_empty());
543    }
544
545    #[tokio::test]
546    async fn test_delete_push_config() {
547        let app = make_app();
548        let params = serde_json::json!({
549            "taskId": "t1",
550            "id": "cfg1"
551        });
552        let resp = post_jsonrpc(app, methods::DELETE_PUSH_CONFIG, params).await;
553        assert!(resp.error.is_some() || resp.result.is_some());
554    }
555
556    #[tokio::test]
557    async fn test_get_extended_agent_card() {
558        let app = make_app();
559        let params = serde_json::json!({});
560        let resp = post_jsonrpc(app, methods::GET_EXTENDED_AGENT_CARD, params).await;
561        // DefaultRequestHandler returns NotSupported
562        assert!(resp.error.is_some() || resp.result.is_some());
563    }
564
565    #[tokio::test]
566    async fn test_streaming_send_message() {
567        let app = make_app();
568        let body = serde_json::json!({
569            "message": {
570                "messageId": "m1",
571                "role": "ROLE_USER",
572                "parts": [{"text": "hello"}]
573            }
574        });
575        let rpc = JsonRpcRequest::new(
576            JsonRpcId::Number(1),
577            methods::SEND_STREAMING_MESSAGE,
578            Some(body),
579        );
580        let req = Request::builder()
581            .uri("/")
582            .method("POST")
583            .header("content-type", "application/json")
584            .header("accept", "text/event-stream")
585            .body(Body::from(serde_json::to_string(&rpc).unwrap()))
586            .unwrap();
587        let resp = app.oneshot(req).await.unwrap();
588        // SSE response => 200
589        assert_eq!(resp.status(), StatusCode::OK);
590    }
591
592    #[tokio::test]
593    async fn test_streaming_subscribe_to_task() {
594        let app = make_app();
595        let rpc = JsonRpcRequest::new(
596            JsonRpcId::Number(1),
597            methods::SUBSCRIBE_TO_TASK,
598            Some(serde_json::json!({"id": "t1"})),
599        );
600        let req = Request::builder()
601            .uri("/")
602            .method("POST")
603            .header("content-type", "application/json")
604            .body(Body::from(serde_json::to_string(&rpc).unwrap()))
605            .unwrap();
606        let resp = app.oneshot(req).await.unwrap();
607        // Subscription may fail (task not found), but routing should work
608        assert_eq!(resp.status(), StatusCode::OK);
609    }
610
611    #[tokio::test]
612    async fn test_streaming_send_message_legacy_method_rejected() {
613        let app = make_app();
614        let body = serde_json::json!({
615            "message": {
616                "messageId": "m1",
617                "role": "ROLE_USER",
618                "parts": [{"text": "hello"}]
619            }
620        });
621        let rpc = JsonRpcRequest::new(JsonRpcId::Number(1), "message.stream", Some(body));
622        let req = Request::builder()
623            .uri("/")
624            .method("POST")
625            .header("content-type", "application/json")
626            .header("accept", "text/event-stream")
627            .body(Body::from(serde_json::to_string(&rpc).unwrap()))
628            .unwrap();
629        let resp = app.oneshot(req).await.unwrap();
630        let body = resp.into_body().collect().await.unwrap().to_bytes();
631        let rpc_resp: JsonRpcResponse = serde_json::from_slice(&body).unwrap();
632        assert!(
633            rpc_resp.error.is_some(),
634            "unexpected result: {:?}",
635            rpc_resp.result
636        );
637        assert_eq!(rpc_resp.error.unwrap().code, error_code::METHOD_NOT_FOUND);
638    }
639
640    #[tokio::test]
641    async fn test_send_message_propagates_headers_as_service_params() {
642        let (handler, captured) = CapturingHandler::new();
643        let app = jsonrpc_router(handler);
644
645        let rpc = JsonRpcRequest::new(
646            JsonRpcId::Number(1),
647            methods::SEND_MESSAGE,
648            Some(serde_json::json!({
649                "message": {
650                    "messageId": "m1",
651                    "role": "ROLE_USER",
652                    "parts": [{"text": "hi"}]
653                }
654            })),
655        );
656        let req = Request::builder()
657            .uri("/")
658            .method("POST")
659            .header("content-type", "application/json")
660            .header("authorization", "Bearer jsonrpc-token")
661            .header("x-tenant-id", "acme")
662            .body(Body::from(serde_json::to_string(&rpc).unwrap()))
663            .unwrap();
664        let resp = app.oneshot(req).await.unwrap();
665        assert_eq!(resp.status(), StatusCode::OK);
666
667        assert_header_captured(
668            &captured,
669            "send_message",
670            "authorization",
671            "Bearer jsonrpc-token",
672        );
673        assert_header_captured(&captured, "send_message", "x-tenant-id", "acme");
674    }
675
676    #[tokio::test]
677    async fn test_streaming_send_message_propagates_headers_as_service_params() {
678        let (handler, captured) = CapturingHandler::new();
679        let app = jsonrpc_router(handler);
680
681        let rpc = JsonRpcRequest::new(
682            JsonRpcId::Number(1),
683            methods::SEND_STREAMING_MESSAGE,
684            Some(serde_json::json!({
685                "message": {
686                    "messageId": "m1",
687                    "role": "ROLE_USER",
688                    "parts": [{"text": "hi"}]
689                }
690            })),
691        );
692        let req = Request::builder()
693            .uri("/")
694            .method("POST")
695            .header("content-type", "application/json")
696            .header("accept", "text/event-stream")
697            .header("authorization", "Bearer jsonrpc-streaming-token")
698            .header("x-tenant-id", "acme")
699            .body(Body::from(serde_json::to_string(&rpc).unwrap()))
700            .unwrap();
701        let resp = app.oneshot(req).await.unwrap();
702        assert_eq!(resp.status(), StatusCode::OK);
703        // Drain the SSE body so the handler is driven to completion.
704        let _ = resp.into_body().collect().await.unwrap().to_bytes();
705
706        assert_header_captured(
707            &captured,
708            "send_streaming_message",
709            "authorization",
710            "Bearer jsonrpc-streaming-token",
711        );
712        assert_header_captured(&captured, "send_streaming_message", "x-tenant-id", "acme");
713    }
714
715    #[tokio::test]
716    async fn test_jsonrpc_unary_methods_propagate_headers_as_service_params() {
717        // Exercise every unary JSON-RPC method to lock in header propagation
718        // across the full match arm in handle_unary_request.
719        let cases: &[(&'static str, &'static str, serde_json::Value)] = &[
720            (
721                methods::GET_TASK,
722                "get_task",
723                serde_json::json!({"id": "t1"}),
724            ),
725            (methods::LIST_TASKS, "list_tasks", serde_json::json!({})),
726            (
727                methods::CANCEL_TASK,
728                "cancel_task",
729                serde_json::json!({"id": "t1"}),
730            ),
731            (
732                methods::CREATE_PUSH_CONFIG,
733                "create_push_config",
734                serde_json::json!({"taskId": "t1", "url": "http://example.com/cb"}),
735            ),
736            (
737                methods::GET_PUSH_CONFIG,
738                "get_push_config",
739                serde_json::json!({"taskId": "t1", "id": "cfg1"}),
740            ),
741            (
742                methods::LIST_PUSH_CONFIGS,
743                "list_push_configs",
744                serde_json::json!({"taskId": "t1"}),
745            ),
746            (
747                methods::DELETE_PUSH_CONFIG,
748                "delete_push_config",
749                serde_json::json!({"taskId": "t1", "id": "cfg1"}),
750            ),
751            (
752                methods::GET_EXTENDED_AGENT_CARD,
753                "get_extended_agent_card",
754                serde_json::json!({}),
755            ),
756        ];
757
758        for (rpc_method, handler_method, params) in cases {
759            let (handler, captured) = CapturingHandler::new();
760            let app = jsonrpc_router(handler);
761            let rpc = JsonRpcRequest::new(JsonRpcId::Number(1), *rpc_method, Some(params.clone()));
762            let token = format!("Bearer token-for-{handler_method}");
763            let req = Request::builder()
764                .uri("/")
765                .method("POST")
766                .header("content-type", "application/json")
767                .header("authorization", &token)
768                .body(Body::from(serde_json::to_string(&rpc).unwrap()))
769                .unwrap();
770            let resp = app.oneshot(req).await.unwrap();
771            assert_eq!(resp.status(), StatusCode::OK);
772            assert_header_captured(&captured, handler_method, "authorization", &token);
773        }
774    }
775
776    #[tokio::test]
777    async fn test_jsonrpc_subscribe_to_task_propagates_headers_as_service_params() {
778        let (handler, captured) = CapturingHandler::new();
779        let app = jsonrpc_router(handler);
780        let rpc = JsonRpcRequest::new(
781            JsonRpcId::Number(1),
782            methods::SUBSCRIBE_TO_TASK,
783            Some(serde_json::json!({"id": "t1"})),
784        );
785        let req = Request::builder()
786            .uri("/")
787            .method("POST")
788            .header("content-type", "application/json")
789            .header("authorization", "Bearer subscribe-token")
790            .body(Body::from(serde_json::to_string(&rpc).unwrap()))
791            .unwrap();
792        let resp = app.oneshot(req).await.unwrap();
793        assert_eq!(resp.status(), StatusCode::OK);
794        let _ = resp.into_body().collect().await.unwrap().to_bytes();
795        assert_header_captured(
796            &captured,
797            "subscribe_to_task",
798            "authorization",
799            "Bearer subscribe-token",
800        );
801    }
802
803    #[tokio::test]
804    async fn test_streaming_invalid_method() {
805        let app = make_app();
806        let rpc = serde_json::json!({
807            "jsonrpc": "2.0",
808            "id": 1,
809            "method": "tasks.subscribe",
810            "params": {"bogus": true}
811        });
812        let req = Request::builder()
813            .uri("/")
814            .method("POST")
815            .header("content-type", "application/json")
816            .body(Body::from(serde_json::to_string(&rpc).unwrap()))
817            .unwrap();
818        let resp = app.oneshot(req).await.unwrap();
819        assert_eq!(resp.status(), StatusCode::OK);
820    }
821}
a2a_server/jsonrpc.rs

a2a_server/
jsonrpc.rs
