a2a-server-lf 0.5.2

// Copyright AGNTCY Contributors (https://github.com/agntcy)
// Copyright A2A Contributors (https://github.com/a2aproject)
// SPDX-License-Identifier: Apache-2.0
use std::sync::Arc;

use a2a::*;
use a2a_pb::protojson_conv::{self, ProtoJsonPayload};
use axum::{
    Json,
    body::Bytes,
    extract::State,
    http::{HeaderMap, StatusCode},
    response::IntoResponse,
};
use futures::{StreamExt, stream::BoxStream};
use serde_json::Value;

use crate::handler::RequestHandler;
use crate::middleware::{ServiceParams, extract_service_params};
use crate::sse;

/// Shared state for the JSON-RPC handler.
pub struct JsonRpcState<H: RequestHandler> {
    pub handler: Arc<H>,
}

impl<H: RequestHandler> Clone for JsonRpcState<H> {
    fn clone(&self) -> Self {
        JsonRpcState {
            handler: self.handler.clone(),
        }
    }
}

/// Largest request body either binding accepts.
///
/// Without an explicit bound a single large POST is read into memory. The
/// limit matches a2a-go's `MaxSSETokenSize` (`internal/sse/sse.go`), so a
/// payload one SDK accepts is not rejected by the other.
pub const MAX_REQUEST_BODY_BYTES: usize = 10 * 1024 * 1024;

/// Create an axum router for the JSON-RPC protocol binding.
///
/// All requests are dispatched to a single POST endpoint that routes
/// by the `method` field in the JSON-RPC envelope.
pub fn jsonrpc_router<H: RequestHandler>(handler: Arc<H>) -> axum::Router {
    let state = JsonRpcState { handler };
    axum::Router::new()
        .route("/", axum::routing::post(handle_jsonrpc::<H>))
        .layer(axum::extract::DefaultBodyLimit::max(MAX_REQUEST_BODY_BYTES))
        .with_state(state)
}

async fn handle_jsonrpc<H: RequestHandler>(
    State(state): State<JsonRpcState<H>>,
    headers: HeaderMap,
    body: Bytes,
) -> impl IntoResponse {
    let params = extract_service_params(&headers);

    // A `Json<JsonRpcRequest>` extractor would reject a malformed body with
    // axum's own error response before this handler runs at all, which is
    // the wrong shape for JSON-RPC over HTTP: the status stays 200 and the
    // error belongs inside the envelope regardless of what failed.
    let request: JsonRpcRequest = match serde_json::from_slice(&body) {
        Ok(request) => request,
        Err(error) => return malformed_request_response(&body, error),
    };
    let id = request.id.clone();
    let method = request.method.as_str();

    if request.jsonrpc != "2.0" {
        return error_response(id, A2AError::invalid_request("invalid jsonrpc version"));
    }

    if let Err(error) = check_a2a_version(&params) {
        return error_response(id, error);
    }

    if methods::is_streaming(method) {
        return handle_streaming_request(&state, &params, &request).await;
    }

    handle_unary_request(&state, &params, &request).await
}

/// A body that failed to deserialize into [`JsonRpcRequest`]: invalid JSON
/// syntax is `ParseError` (-32700); well-formed JSON with the wrong shape
/// (e.g. a missing `method`) is `InvalidRequestError` (-32600). The id is
/// recovered on a best-effort basis for the latter, since JSON-RPC 2.0 asks
/// for it to be echoed back when it can be determined; a syntax error, by
/// definition, cannot.
fn malformed_request_response(body: &[u8], error: serde_json::Error) -> axum::response::Response {
    if !error.is_data() {
        return error_response(
            JsonRpcId::Null,
            A2AError {
                code: error_code::PARSE_ERROR,
                message: format!("parse error: {error}"),
                details: None,
            },
        );
    }
    let id = serde_json::from_slice::<Value>(body)
        .ok()
        .and_then(|v| v.get("id").cloned())
        .and_then(|v| serde_json::from_value(v).ok())
        .unwrap_or(JsonRpcId::Null);
    error_response(
        id,
        A2AError::invalid_request(format!("invalid request: {error}")),
    )
}

async fn handle_unary_request<H: RequestHandler>(
    state: &JsonRpcState<H>,
    params: &ServiceParams,
    request: &JsonRpcRequest,
) -> axum::response::Response {
    let id = request.id.clone();
    let raw_params = request.params.clone().unwrap_or(Value::Null);

    let result: Result<Value, A2AError> = match request.method.as_str() {
        methods::SEND_MESSAGE => match protojson_conv::from_value::<SendMessageRequest>(raw_params)
        {
            Ok(req) => state
                .handler
                .send_message(params, req)
                .await
                .and_then(|r| protojson_value(&r)),
            Err(e) => Err(invalid_params(e)),
        },
        methods::GET_TASK => match protojson_conv::from_value::<GetTaskRequest>(raw_params) {
            Ok(req) => state
                .handler
                .get_task(params, req)
                .await
                .and_then(|r| protojson_value(&r)),
            Err(e) => Err(invalid_params(e)),
        },
        methods::LIST_TASKS => match protojson_conv::from_value::<ListTasksRequest>(raw_params) {
            Ok(req) => state
                .handler
                .list_tasks(params, req)
                .await
                .and_then(|r| protojson_value(&r)),
            Err(e) => Err(invalid_params(e)),
        },
        methods::CANCEL_TASK => match protojson_conv::from_value::<CancelTaskRequest>(raw_params) {
            Ok(req) => state
                .handler
                .cancel_task(params, req)
                .await
                .and_then(|r| protojson_value(&r)),
            Err(e) => Err(invalid_params(e)),
        },
        methods::CREATE_PUSH_CONFIG => match parse_create_push_config_request(raw_params) {
            Ok(req) => state
                .handler
                .create_push_config(params, req)
                .await
                .and_then(|r| protojson_value(&r)),
            Err(e) => Err(invalid_params(e)),
        },
        methods::GET_PUSH_CONFIG => {
            match protojson_conv::from_value::<GetTaskPushNotificationConfigRequest>(raw_params) {
                Ok(req) => state
                    .handler
                    .get_push_config(params, req)
                    .await
                    .and_then(|r| protojson_value(&r)),
                Err(e) => Err(invalid_params(e)),
            }
        }
        methods::LIST_PUSH_CONFIGS => {
            match protojson_conv::from_value::<ListTaskPushNotificationConfigsRequest>(raw_params) {
                Ok(req) => state
                    .handler
                    .list_push_configs(params, req)
                    .await
                    .and_then(|r| protojson_value(&r)),
                Err(e) => Err(invalid_params(e)),
            }
        }
        methods::DELETE_PUSH_CONFIG => {
            match protojson_conv::from_value::<DeleteTaskPushNotificationConfigRequest>(raw_params)
            {
                Ok(req) => state
                    .handler
                    .delete_push_config(params, req)
                    .await
                    .map(|_| Value::Null),
                Err(e) => Err(invalid_params(e)),
            }
        }
        methods::GET_EXTENDED_AGENT_CARD => {
            match protojson_conv::from_value::<GetExtendedAgentCardRequest>(raw_params) {
                Ok(req) => state
                    .handler
                    .get_extended_agent_card(params, req)
                    .await
                    .and_then(|r| protojson_value(&r)),
                Err(e) => Err(invalid_params(e)),
            }
        }
        "" => Err(A2AError::invalid_request("method is required")),
        _ => Err(A2AError::method_not_found(&request.method)),
    };

    match result {
        Ok(value) => {
            let resp = JsonRpcResponse::success(id, value);
            Json(resp).into_response()
        }
        Err(e) => error_response(id, e),
    }
}

async fn handle_streaming_request<H: RequestHandler>(
    state: &JsonRpcState<H>,
    params: &ServiceParams,
    request: &JsonRpcRequest,
) -> axum::response::Response {
    let id = request.id.clone();
    let raw_params = request.params.clone().unwrap_or(Value::Null);

    match request.method.as_str() {
        methods::SEND_STREAMING_MESSAGE => {
            match protojson_conv::from_value::<SendMessageRequest>(raw_params) {
                Ok(req) => match state.handler.send_streaming_message(params, req).await {
                    Ok(stream) => {
                        sse::sse_jsonrpc_stream(id, protojson_stream(stream)).into_response()
                    }
                    Err(e) => error_response(id, e),
                },
                Err(e) => error_response(id, invalid_params(e)),
            }
        }
        methods::SUBSCRIBE_TO_TASK => {
            match protojson_conv::from_value::<SubscribeToTaskRequest>(raw_params) {
                Ok(req) => match state.handler.subscribe_to_task(params, req).await {
                    Ok(stream) => {
                        sse::sse_jsonrpc_stream(id, protojson_stream(stream)).into_response()
                    }
                    Err(e) => error_response(id, e),
                },
                Err(e) => error_response(id, invalid_params(e)),
            }
        }
        _ => error_response(id, A2AError::method_not_found(&request.method)),
    }
}

fn error_response(id: JsonRpcId, err: A2AError) -> axum::response::Response {
    let resp = JsonRpcResponse::error(id, err.to_jsonrpc_error());
    (StatusCode::OK, Json(resp)).into_response()
}

fn protojson_value<T: ProtoJsonPayload>(value: &T) -> Result<Value, A2AError> {
    protojson_conv::to_value(value)
        .map_err(|e| A2AError::internal(format!("failed to serialize ProtoJSON payload: {e}")))
}

fn protojson_stream(
    stream: BoxStream<'static, Result<StreamResponse, A2AError>>,
) -> BoxStream<'static, Result<Value, A2AError>> {
    Box::pin(stream.map(|item| {
        item.and_then(|value| {
            protojson_conv::to_value(&value).map_err(|e| {
                A2AError::internal(format!("failed to serialize ProtoJSON stream payload: {e}"))
            })
        })
    }))
}

fn parse_create_push_config_request(
    raw_params: Value,
) -> Result<TaskPushNotificationConfig, String> {
    protojson_conv::from_value::<TaskPushNotificationConfig>(raw_params).map_err(|e| e.to_string())
}

/// A2A §9.5: `params` that fail schema validation are
/// `InvalidParamsError` (-32602). `PARSE_ERROR` (-32700) is reserved for a
/// body that is not valid JSON, which `malformed_request_response` handles
/// before this point.
fn invalid_params(e: impl std::fmt::Display) -> A2AError {
    A2AError::invalid_params(format!("invalid params: {e}"))
}

/// The major version this build implements (§3.6.2). Minor is accepted
/// unconditionally since nothing here branches on it.
const SUPPORTED_MAJOR_VERSION: u32 = 1;

/// Mirrors a2acli's own `parse_protocol_version`: tolerant of a bare major,
/// extra trailing components, and surrounding whitespace.
fn parse_major_version(value: &str) -> Option<u32> {
    value.trim().split('.').next()?.parse().ok()
}

/// §3.6.2: reject a request whose `A2A-Version` header names an
/// unsupported major version. §3.6.1/§3.6.2 both require an absent header
/// to be interpreted as "0.3", not as "whatever this interface happens to
/// serve" -- so it's rejected here exactly like any other unsupported
/// version, rather than silently processed as this build's version.
fn check_a2a_version(params: &ServiceParams) -> Result<(), A2AError> {
    let requested = params
        .get("a2a-version")
        .and_then(|v| v.first())
        .map(String::as_str)
        .unwrap_or("0.3");
    match parse_major_version(requested) {
        Some(SUPPORTED_MAJOR_VERSION) => Ok(()),
        _ => Err(A2AError::version_not_supported(requested)),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::InMemoryPushConfigStore;
    use crate::executor::ExecutorContext;
    use crate::handler::DefaultRequestHandler;
    use crate::task_store::InMemoryTaskStore;
    use a2a_pb::protojson_conv;
    use axum::body::Body;
    use axum::http::Request;

    use crate::test_util::{CapturingHandler, assert_header_captured, install_crypto_provider};
    use futures::stream::BoxStream;
    use http_body_util::BodyExt;
    use tower::ServiceExt;

    struct EchoExecutor;

    impl crate::AgentExecutor for EchoExecutor {
        fn execute(
            &self,
            ctx: ExecutorContext,
        ) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
            let task = Task {
                id: ctx.task_id.clone(),
                context_id: ctx.context_id.clone(),
                status: TaskStatus {
                    state: TaskState::Completed,
                    message: ctx.message,
                    timestamp: None,
                },
                artifacts: None,
                history: None,
                metadata: None,
            };
            Box::pin(futures::stream::once(async move {
                Ok(StreamResponse::Task(task))
            }))
        }

        fn cancel(
            &self,
            ctx: ExecutorContext,
        ) -> BoxStream<'static, Result<StreamResponse, A2AError>> {
            let task = Task {
                id: ctx.task_id.clone(),
                context_id: ctx.context_id.clone(),
                status: TaskStatus {
                    state: TaskState::Canceled,
                    message: None,
                    timestamp: None,
                },
                artifacts: None,
                history: None,
                metadata: None,
            };
            Box::pin(futures::stream::once(async move {
                Ok(StreamResponse::Task(task))
            }))
        }
    }

    fn make_app() -> axum::Router {
        let handler = Arc::new(DefaultRequestHandler::new(
            EchoExecutor,
            InMemoryTaskStore::new(),
        ));
        jsonrpc_router(handler)
    }

    fn make_push_app() -> axum::Router {
        install_crypto_provider();
        let handler = Arc::new(
            DefaultRequestHandler::new(EchoExecutor, InMemoryTaskStore::new())
                .with_push_config_store(InMemoryPushConfigStore::new()),
        );
        jsonrpc_router(handler)
    }

    async fn post_jsonrpc(app: axum::Router, method: &str, params: Value) -> JsonRpcResponse {
        let rpc = JsonRpcRequest::new(JsonRpcId::Number(1), method, Some(params));
        let body = serde_json::to_string(&rpc).unwrap();
        let req = Request::builder()
            .uri("/")
            .method("POST")
            .header("content-type", "application/json")
            // A supported version, so tests using this helper for something
            // other than version negotiation itself aren't confounded by
            // §3.6.2's now-enforced absent-header rejection.
            .header("a2a-version", "1.0")
            .body(Body::from(body))
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        let body = resp.into_body().collect().await.unwrap().to_bytes();
        serde_json::from_slice(&body).unwrap()
    }

    /// An unbounded body is read into memory; the limit rejects it at the
    /// framework boundary instead, before any handler runs.
    #[tokio::test]
    async fn test_request_body_over_the_limit_is_rejected() {
        let app = make_app();
        let oversized = Body::from("x".repeat(MAX_REQUEST_BODY_BYTES + 1));
        let req = Request::builder()
            .uri("/")
            .method("POST")
            .header("content-type", "application/json")
            .body(oversized)
            .unwrap();

        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), axum::http::StatusCode::PAYLOAD_TOO_LARGE);
    }

    #[tokio::test]
    async fn test_send_message() {
        let app = make_app();
        let params = serde_json::json!({
            "message": {
                "messageId": "m1",
                "role": "ROLE_USER",
                "parts": [{"text": "hello"}]
            }
        });
        let resp = post_jsonrpc(app, methods::SEND_MESSAGE, params).await;
        assert!(resp.error.is_none(), "unexpected error: {:?}", resp.error);
        let result =
            protojson_conv::from_value::<SendMessageResponse>(resp.result.unwrap()).unwrap();
        assert!(matches!(result, SendMessageResponse::Task(_)));
    }

    #[tokio::test]
    async fn test_send_message_legacy_method_rejected() {
        let app = make_app();
        let params = serde_json::json!({
            "message": {
                "messageId": "m1",
                "role": "ROLE_USER",
                "parts": [{"text": "hello"}]
            }
        });
        let resp = post_jsonrpc(app, "message.send", params).await;
        assert!(resp.error.is_some(), "unexpected result: {:?}", resp.result);
        assert_eq!(resp.error.unwrap().code, error_code::METHOD_NOT_FOUND);
    }

    #[tokio::test]
    async fn test_get_task_not_found() {
        let app = make_app();
        let params = serde_json::json!({"id": "nonexistent"});
        let resp = post_jsonrpc(app, methods::GET_TASK, params).await;
        assert!(resp.error.is_some());
        assert_eq!(resp.error.unwrap().code, error_code::TASK_NOT_FOUND);
    }

    #[tokio::test]
    async fn test_invalid_method() {
        let app = make_app();
        let resp = post_jsonrpc(app, "unknown.method", Value::Null).await;
        assert!(resp.error.is_some());
        assert_eq!(resp.error.unwrap().code, error_code::METHOD_NOT_FOUND);
    }

    #[tokio::test]
    async fn test_error_response_boundary_does_not_rewrite_internal_errors() {
        // The boundary no longer sanitizes: an INTERNAL_ERROR keeps its
        // message, since the boundary cannot tell an executor's own failure
        // report from a server fault. Sanitization happens at the raise site
        // via `sanitized_internal_error`.
        let resp = error_response(JsonRpcId::Number(1), A2AError::internal("boom details"));
        let body = resp.into_body().collect().await.unwrap().to_bytes();
        let rpc_resp: JsonRpcResponse = serde_json::from_slice(&body).unwrap();
        let error = rpc_resp.error.unwrap();
        assert_eq!(error.code, error_code::INTERNAL_ERROR);
        assert_eq!(error.message, "boom details");

        // Non-internal errors keep their message (client-validation feedback).
        let resp = error_response(JsonRpcId::Number(1), A2AError::task_not_found("t1"));
        let body = resp.into_body().collect().await.unwrap().to_bytes();
        let rpc_resp: JsonRpcResponse = serde_json::from_slice(&body).unwrap();
        assert_eq!(rpc_resp.error.unwrap().message, "task not found: t1");
    }

    #[tokio::test]
    async fn test_empty_method() {
        let app = make_app();
        let resp = post_jsonrpc(app, "", Value::Null).await;
        assert!(resp.error.is_some());
    }

    #[tokio::test]
    async fn test_invalid_jsonrpc_version() {
        let app = make_app();
        let rpc = serde_json::json!({
            "jsonrpc": "1.0",
            "id": 1,
            "method": methods::SEND_MESSAGE,
            "params": {}
        });
        let body = serde_json::to_string(&rpc).unwrap();
        let req = Request::builder()
            .uri("/")
            .method("POST")
            .header("content-type", "application/json")
            .body(Body::from(body))
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        let body = resp.into_body().collect().await.unwrap().to_bytes();
        let rpc_resp: JsonRpcResponse = serde_json::from_slice(&body).unwrap();
        assert!(rpc_resp.error.is_some());
    }

    /// A `Json<T>` extractor would have axum reject this with a 422 before
    /// the handler runs; JSON-RPC over HTTP always answers 200 with the
    /// error inside the envelope, regardless of what about the body failed.
    #[tokio::test]
    async fn test_missing_method_field_is_invalid_request_not_a_422() {
        let app = make_app();
        let req = Request::builder()
            .uri("/")
            .method("POST")
            .header("content-type", "application/json")
            .body(Body::from(r#"{"jsonrpc":"2.0","id":1,"params":{}}"#))
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
        let body = resp.into_body().collect().await.unwrap().to_bytes();
        let rpc_resp: JsonRpcResponse = serde_json::from_slice(&body).unwrap();
        let code = rpc_resp.error.expect("expected an error").code;
        assert!(
            code == error_code::INVALID_REQUEST || code == error_code::INVALID_PARAMS,
            "got {code}"
        );
        assert_eq!(rpc_resp.id, JsonRpcId::Number(1), "id should be recovered");
    }

    #[tokio::test]
    async fn test_invalid_json_syntax_is_parse_error_not_a_400() {
        let app = make_app();
        let req = Request::builder()
            .uri("/")
            .method("POST")
            .header("content-type", "application/json")
            .body(Body::from("{this is not valid json"))
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
        let body = resp.into_body().collect().await.unwrap().to_bytes();
        let rpc_resp: JsonRpcResponse = serde_json::from_slice(&body).unwrap();
        assert_eq!(
            rpc_resp.error.expect("expected an error").code,
            error_code::PARSE_ERROR
        );
        assert_eq!(rpc_resp.id, JsonRpcId::Null, "id cannot be recovered");
    }

    /// A2A §3.3.2 / §9.5: parameters that fail to decode are
    /// `InvalidParamsError` (-32602). `JSONParseError` (-32700) is reserved
    /// for a body that is not valid JSON at all, which
    /// `test_invalid_json_syntax_is_parse_error_not_a_400` covers.
    #[tokio::test]
    async fn test_invalid_params() {
        let app = make_app();
        // Type mismatch: `message` must be an object. A merely-unknown field
        // is not invalid params; it is ignored per spec §5.7.
        let params = serde_json::json!({"message": "not-an-object"});
        let resp = post_jsonrpc(app, methods::SEND_MESSAGE, params).await;
        let error = resp.error.expect("expected an error");
        assert_eq!(error.code, error_code::INVALID_PARAMS);
    }

    /// A2A §9.5: a `params` payload that fails to decode is
    /// `InvalidParamsError` (-32602) for *every* method, not just
    /// `message/send`. Each arm decodes its own request type, so guard the
    /// whole table rather than a single representative method.
    #[tokio::test]
    async fn test_invalid_params_is_reported_for_every_method() {
        // Each payload is a well-formed JSON object whose one recognized
        // field has the wrong type, so decoding fails before the handler.
        let cases = [
            (methods::GET_TASK, serde_json::json!({ "id": 123 })),
            (
                methods::LIST_TASKS,
                serde_json::json!({ "pageSize": "not-a-number" }),
            ),
            (methods::CANCEL_TASK, serde_json::json!({ "id": 123 })),
            (
                methods::CREATE_PUSH_CONFIG,
                serde_json::json!({ "taskId": 123 }),
            ),
            (
                methods::GET_PUSH_CONFIG,
                serde_json::json!({ "taskId": 123 }),
            ),
            (
                methods::LIST_PUSH_CONFIGS,
                serde_json::json!({ "taskId": 123 }),
            ),
            (
                methods::DELETE_PUSH_CONFIG,
                serde_json::json!({ "taskId": 123 }),
            ),
            (
                methods::GET_EXTENDED_AGENT_CARD,
                serde_json::json!({ "tenant": 123 }),
            ),
        ];
        for (method, params) in cases {
            let resp = post_jsonrpc(make_app(), method, params).await;
            let error = resp
                .error
                .unwrap_or_else(|| panic!("{method}: expected an error, got {:?}", resp.result));
            assert_eq!(error.code, error_code::INVALID_PARAMS, "method {method}");
        }
    }

    /// Spec §5.7: unrecognized fields in request params are ignored for
    /// forward compatibility instead of being rejected.
    #[tokio::test]
    async fn test_unknown_params_fields_are_ignored() {
        let app = make_app();
        let params = serde_json::json!({
            "message": {
                "messageId": "m1",
                "role": "ROLE_USER",
                "parts": [{"text": "hi"}],
                "futureField": true
            },
            "anotherFutureField": {"nested": 1}
        });
        let resp = post_jsonrpc(app, methods::SEND_MESSAGE, params).await;
        assert!(
            resp.error.is_none(),
            "unknown fields must be ignored: {:?}",
            resp.error
        );
        assert!(resp.result.is_some());
    }

    #[tokio::test]
    async fn test_list_tasks() {
        let app = make_app();
        let params = serde_json::json!({});
        let resp = post_jsonrpc(app, methods::LIST_TASKS, params).await;
        let result = protojson_conv::from_value::<ListTasksResponse>(resp.result.unwrap()).unwrap();
        assert!(result.tasks.is_empty());
    }

    #[tokio::test]
    async fn test_cancel_task_not_found() {
        let app = make_app();
        let params = serde_json::json!({"id": "nonexistent"});
        let resp = post_jsonrpc(app, methods::CANCEL_TASK, params).await;
        assert!(resp.error.is_some());
    }

    #[tokio::test]
    async fn test_create_push_config() {
        let app = make_push_app();
        let params = serde_json::json!({
            "taskId": "t1",
            "id": "cfg1",
            "url": "http://example.com/callback"
        });
        let resp = post_jsonrpc(app, methods::CREATE_PUSH_CONFIG, params).await;
        assert!(resp.error.is_none(), "unexpected error: {:?}", resp.error);
        let result =
            serde_json::from_value::<TaskPushNotificationConfig>(resp.result.unwrap()).unwrap();
        assert_eq!(result.task_id, "t1");
        assert_eq!(result.id.as_deref(), Some("cfg1"));
        assert_eq!(result.url, "http://example.com/callback");
    }

    #[tokio::test]
    async fn test_create_push_config_accepts_flat_config_object() {
        let app = make_push_app();
        let params = serde_json::json!({
            "taskId": "t1",
            "id": "cfg1",
            "url": "http://example.com/callback"
        });
        let resp = post_jsonrpc(app, methods::CREATE_PUSH_CONFIG, params).await;
        assert!(resp.error.is_none(), "unexpected error: {:?}", resp.error);
        let result =
            serde_json::from_value::<TaskPushNotificationConfig>(resp.result.unwrap()).unwrap();
        assert_eq!(result.task_id, "t1");
        assert_eq!(result.id.as_deref(), Some("cfg1"));
        assert_eq!(result.url, "http://example.com/callback");
    }

    #[tokio::test]
    async fn test_get_push_config() {
        let app = make_app();
        let params = serde_json::json!({
            "taskId": "t1",
            "id": "cfg1"
        });
        let resp = post_jsonrpc(app, methods::GET_PUSH_CONFIG, params).await;
        assert!(resp.error.is_some() || resp.result.is_some());
    }

    #[tokio::test]
    async fn test_list_push_configs() {
        let app = make_push_app();
        let params = serde_json::json!({
            "taskId": "t1"
        });
        let resp = post_jsonrpc(app, methods::LIST_PUSH_CONFIGS, params).await;
        assert!(resp.error.is_none(), "unexpected error: {:?}", resp.error);
        let result =
            serde_json::from_value::<ListTaskPushNotificationConfigsResponse>(resp.result.unwrap())
                .unwrap();
        assert!(result.configs.is_empty());
    }

    #[tokio::test]
    async fn test_delete_push_config() {
        let app = make_app();
        let params = serde_json::json!({
            "taskId": "t1",
            "id": "cfg1"
        });
        let resp = post_jsonrpc(app, methods::DELETE_PUSH_CONFIG, params).await;
        assert!(resp.error.is_some() || resp.result.is_some());
    }

    #[tokio::test]
    async fn test_get_extended_agent_card() {
        let app = make_app();
        let params = serde_json::json!({});
        let resp = post_jsonrpc(app, methods::GET_EXTENDED_AGENT_CARD, params).await;
        // DefaultRequestHandler returns NotSupported
        assert!(resp.error.is_some() || resp.result.is_some());
    }

    #[tokio::test]
    async fn test_streaming_send_message() {
        let app = make_app();
        let body = serde_json::json!({
            "message": {
                "messageId": "m1",
                "role": "ROLE_USER",
                "parts": [{"text": "hello"}]
            }
        });
        let rpc = JsonRpcRequest::new(
            JsonRpcId::Number(1),
            methods::SEND_STREAMING_MESSAGE,
            Some(body),
        );
        let req = Request::builder()
            .uri("/")
            .method("POST")
            .header("content-type", "application/json")
            .header("accept", "text/event-stream")
            .body(Body::from(serde_json::to_string(&rpc).unwrap()))
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        // SSE response => 200
        assert_eq!(resp.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn test_streaming_subscribe_to_task() {
        let app = make_app();
        let rpc = JsonRpcRequest::new(
            JsonRpcId::Number(1),
            methods::SUBSCRIBE_TO_TASK,
            Some(serde_json::json!({"id": "t1"})),
        );
        let req = Request::builder()
            .uri("/")
            .method("POST")
            .header("content-type", "application/json")
            .body(Body::from(serde_json::to_string(&rpc).unwrap()))
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        // Subscription may fail (task not found), but routing should work
        assert_eq!(resp.status(), StatusCode::OK);
    }

    /// A params-decode failure on the streaming binding must carry the same
    /// code as the unary binding: `InvalidParamsError` (-32602), for both
    /// streaming methods.
    #[tokio::test]
    async fn test_streaming_invalid_params_use_invalid_params_error_code() {
        let cases = [
            (
                methods::SEND_STREAMING_MESSAGE,
                serde_json::json!({"message": "not-an-object"}),
            ),
            (methods::SUBSCRIBE_TO_TASK, serde_json::json!({"id": 123})),
        ];
        for (method, params) in cases {
            let app = make_app();
            let rpc = JsonRpcRequest::new(JsonRpcId::Number(1), method, Some(params));
            let req = Request::builder()
                .uri("/")
                .method("POST")
                .header("content-type", "application/json")
                // A supported version, so this test exercises params decoding
                // rather than §3.6.2's absent-header rejection.
                .header("a2a-version", "1.0")
                .body(Body::from(serde_json::to_string(&rpc).unwrap()))
                .unwrap();
            let resp = app.oneshot(req).await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK);
            let body = resp.into_body().collect().await.unwrap().to_bytes();
            let parsed: JsonRpcResponse = serde_json::from_slice(&body).unwrap();
            let error = parsed.error.expect("expected an error");
            assert_eq!(error.code, error_code::INVALID_PARAMS, "method {method}");
        }
    }

    #[tokio::test]
    async fn test_streaming_send_message_legacy_method_rejected() {
        let app = make_app();
        let body = serde_json::json!({
            "message": {
                "messageId": "m1",
                "role": "ROLE_USER",
                "parts": [{"text": "hello"}]
            }
        });
        let rpc = JsonRpcRequest::new(JsonRpcId::Number(1), "message.stream", Some(body));
        let req = Request::builder()
            .uri("/")
            .method("POST")
            .header("content-type", "application/json")
            .header("accept", "text/event-stream")
            .header("a2a-version", "1.0")
            .body(Body::from(serde_json::to_string(&rpc).unwrap()))
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        let body = resp.into_body().collect().await.unwrap().to_bytes();
        let rpc_resp: JsonRpcResponse = serde_json::from_slice(&body).unwrap();
        assert!(
            rpc_resp.error.is_some(),
            "unexpected result: {:?}",
            rpc_resp.result
        );
        assert_eq!(rpc_resp.error.unwrap().code, error_code::METHOD_NOT_FOUND);
    }

    #[tokio::test]
    async fn test_send_message_propagates_headers_as_service_params() {
        let (handler, captured) = CapturingHandler::new();
        let app = jsonrpc_router(handler);

        let rpc = JsonRpcRequest::new(
            JsonRpcId::Number(1),
            methods::SEND_MESSAGE,
            Some(serde_json::json!({
                "message": {
                    "messageId": "m1",
                    "role": "ROLE_USER",
                    "parts": [{"text": "hi"}]
                }
            })),
        );
        let req = Request::builder()
            .uri("/")
            .method("POST")
            .header("content-type", "application/json")
            .header("authorization", "Bearer jsonrpc-token")
            .header("x-tenant-id", "acme")
            .header("a2a-version", "1.0")
            .body(Body::from(serde_json::to_string(&rpc).unwrap()))
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);

        assert_header_captured(
            &captured,
            "send_message",
            "authorization",
            "Bearer jsonrpc-token",
        );
        assert_header_captured(&captured, "send_message", "x-tenant-id", "acme");
    }

    #[tokio::test]
    async fn test_unsupported_a2a_version_is_rejected() {
        let app = make_app();
        let rpc = JsonRpcRequest::new(
            JsonRpcId::Number(1),
            methods::SEND_MESSAGE,
            Some(serde_json::json!({
                "message": {"messageId": "m1", "role": "ROLE_USER", "parts": [{"text": "hi"}]}
            })),
        );
        let req = Request::builder()
            .uri("/")
            .method("POST")
            .header("content-type", "application/json")
            .header("a2a-version", "99.0")
            .body(Body::from(serde_json::to_string(&rpc).unwrap()))
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
        let body = resp.into_body().collect().await.unwrap().to_bytes();
        let parsed: JsonRpcResponse = serde_json::from_slice(&body).unwrap();
        assert_eq!(
            parsed.error.expect("expected an error").code,
            error_code::VERSION_NOT_SUPPORTED
        );
    }

    #[tokio::test]
    async fn test_missing_a2a_version_is_rejected() {
        // §3.6.1/§3.6.2: an absent A2A-Version header MUST be interpreted
        // as "0.3", not as this interface's own version -- so this server,
        // which only speaks 1.0, rejects it the same as any other
        // unsupported version. No a2a-version header at all, unlike
        // `post_jsonrpc`, which sets one so other tests aren't confounded
        // by this same check.
        let app = make_app();
        let rpc = JsonRpcRequest::new(
            JsonRpcId::Number(1),
            methods::SEND_MESSAGE,
            Some(serde_json::json!({
                "message": {"messageId": "m1", "role": "ROLE_USER", "parts": [{"text": "hi"}]}
            })),
        );
        let req = Request::builder()
            .uri("/")
            .method("POST")
            .header("content-type", "application/json")
            .body(Body::from(serde_json::to_string(&rpc).unwrap()))
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
        let body = resp.into_body().collect().await.unwrap().to_bytes();
        let parsed: JsonRpcResponse = serde_json::from_slice(&body).unwrap();
        assert_eq!(
            parsed.error.expect("expected an error").code,
            error_code::VERSION_NOT_SUPPORTED
        );
    }

    #[tokio::test]
    async fn test_supported_a2a_version_is_accepted() {
        let app = make_app();
        let rpc = JsonRpcRequest::new(
            JsonRpcId::Number(1),
            methods::SEND_MESSAGE,
            Some(serde_json::json!({
                "message": {"messageId": "m1", "role": "ROLE_USER", "parts": [{"text": "hi"}]}
            })),
        );
        let req = Request::builder()
            .uri("/")
            .method("POST")
            .header("content-type", "application/json")
            .header("a2a-version", "1.0")
            .body(Body::from(serde_json::to_string(&rpc).unwrap()))
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
        let body = resp.into_body().collect().await.unwrap().to_bytes();
        let parsed: JsonRpcResponse = serde_json::from_slice(&body).unwrap();
        assert!(
            parsed.error.is_none(),
            "expected success, got {:?}",
            parsed.error
        );
    }

    #[tokio::test]
    async fn test_streaming_send_message_propagates_headers_as_service_params() {
        let (handler, captured) = CapturingHandler::new();
        let app = jsonrpc_router(handler);

        let rpc = JsonRpcRequest::new(
            JsonRpcId::Number(1),
            methods::SEND_STREAMING_MESSAGE,
            Some(serde_json::json!({
                "message": {
                    "messageId": "m1",
                    "role": "ROLE_USER",
                    "parts": [{"text": "hi"}]
                }
            })),
        );
        let req = Request::builder()
            .uri("/")
            .method("POST")
            .header("content-type", "application/json")
            .header("accept", "text/event-stream")
            .header("authorization", "Bearer jsonrpc-streaming-token")
            .header("x-tenant-id", "acme")
            .header("a2a-version", "1.0")
            .body(Body::from(serde_json::to_string(&rpc).unwrap()))
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
        // Drain the SSE body so the handler is driven to completion.
        let _ = resp.into_body().collect().await.unwrap().to_bytes();

        assert_header_captured(
            &captured,
            "send_streaming_message",
            "authorization",
            "Bearer jsonrpc-streaming-token",
        );
        assert_header_captured(&captured, "send_streaming_message", "x-tenant-id", "acme");
    }

    #[tokio::test]
    async fn test_jsonrpc_unary_methods_propagate_headers_as_service_params() {
        // Exercise every unary JSON-RPC method to lock in header propagation
        // across the full match arm in handle_unary_request.
        let cases: &[(&'static str, &'static str, serde_json::Value)] = &[
            (
                methods::GET_TASK,
                "get_task",
                serde_json::json!({"id": "t1"}),
            ),
            (methods::LIST_TASKS, "list_tasks", serde_json::json!({})),
            (
                methods::CANCEL_TASK,
                "cancel_task",
                serde_json::json!({"id": "t1"}),
            ),
            (
                methods::CREATE_PUSH_CONFIG,
                "create_push_config",
                serde_json::json!({"taskId": "t1", "url": "http://example.com/cb"}),
            ),
            (
                methods::GET_PUSH_CONFIG,
                "get_push_config",
                serde_json::json!({"taskId": "t1", "id": "cfg1"}),
            ),
            (
                methods::LIST_PUSH_CONFIGS,
                "list_push_configs",
                serde_json::json!({"taskId": "t1"}),
            ),
            (
                methods::DELETE_PUSH_CONFIG,
                "delete_push_config",
                serde_json::json!({"taskId": "t1", "id": "cfg1"}),
            ),
            (
                methods::GET_EXTENDED_AGENT_CARD,
                "get_extended_agent_card",
                serde_json::json!({}),
            ),
        ];

        for (rpc_method, handler_method, params) in cases {
            let (handler, captured) = CapturingHandler::new();
            let app = jsonrpc_router(handler);
            let rpc = JsonRpcRequest::new(JsonRpcId::Number(1), *rpc_method, Some(params.clone()));
            let token = format!("Bearer token-for-{handler_method}");
            let req = Request::builder()
                .uri("/")
                .method("POST")
                .header("content-type", "application/json")
                .header("authorization", &token)
                .header("a2a-version", "1.0")
                .body(Body::from(serde_json::to_string(&rpc).unwrap()))
                .unwrap();
            let resp = app.oneshot(req).await.unwrap();
            assert_eq!(resp.status(), StatusCode::OK);
            assert_header_captured(&captured, handler_method, "authorization", &token);
        }
    }

    #[tokio::test]
    async fn test_jsonrpc_subscribe_to_task_propagates_headers_as_service_params() {
        let (handler, captured) = CapturingHandler::new();
        let app = jsonrpc_router(handler);
        let rpc = JsonRpcRequest::new(
            JsonRpcId::Number(1),
            methods::SUBSCRIBE_TO_TASK,
            Some(serde_json::json!({"id": "t1"})),
        );
        let req = Request::builder()
            .uri("/")
            .method("POST")
            .header("content-type", "application/json")
            .header("authorization", "Bearer subscribe-token")
            .header("a2a-version", "1.0")
            .body(Body::from(serde_json::to_string(&rpc).unwrap()))
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
        let _ = resp.into_body().collect().await.unwrap().to_bytes();
        assert_header_captured(
            &captured,
            "subscribe_to_task",
            "authorization",
            "Bearer subscribe-token",
        );
    }

    #[tokio::test]
    async fn test_streaming_invalid_method() {
        let app = make_app();
        let rpc = serde_json::json!({
            "jsonrpc": "2.0",
            "id": 1,
            "method": "tasks.subscribe",
            "params": {"bogus": true}
        });
        let req = Request::builder()
            .uri("/")
            .method("POST")
            .header("content-type", "application/json")
            .body(Body::from(serde_json::to_string(&rpc).unwrap()))
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
    }
}