a2a-server-lf 0.4.3

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

use a2a::*;
use a2a_pb::protojson_conv::{self, ProtoJsonPayload};
use axum::{
    Json,
    extract::{Path, Query, State},
    http::{HeaderMap, StatusCode},
    response::IntoResponse,
};
use chrono::{DateTime, Utc};
use futures::{StreamExt, stream::BoxStream};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;

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

const REST_SEND_MESSAGE_PATH: &str = "/message:send";
const REST_SEND_MESSAGE_LEGACY_PATH: &str = "/message/send";
const REST_STREAM_MESSAGE_PATH: &str = "/message:stream";
const REST_STREAM_MESSAGE_LEGACY_PATH: &str = "/message/stream";
const REST_EXTENDED_AGENT_CARD_PATH: &str = "/extendedAgentCard";
const REST_EXTENDED_AGENT_CARD_LEGACY_PATH: &str = "/agent-card/extended";
const REST_PUSH_CONFIGS_PATH: &str = "/tasks/{id}/pushNotificationConfigs";
const REST_PUSH_CONFIGS_LEGACY_PATH: &str = "/tasks/{id}/push-configs";
const REST_PUSH_CONFIG_PATH: &str = "/tasks/{id}/pushNotificationConfigs/{config_id}";
const REST_PUSH_CONFIG_LEGACY_PATH: &str = "/tasks/{id}/push-configs/{config_id}";
#[derive(Serialize)]
struct RestErrorEnvelope {
    error: RestErrorStatus,
}

#[derive(Serialize)]
struct RestErrorStatus {
    code: u16,
    status: &'static str,
    message: String,
    details: Vec<TypedDetail>,
}

/// Shared state for REST handlers.
pub struct RestState<H: RequestHandler> {
    pub handler: Arc<H>,
}

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

/// Create an axum router for the REST/HTTP+JSON protocol binding.
pub fn rest_router<H: RequestHandler>(handler: Arc<H>) -> axum::Router {
    let state = RestState { handler };
    axum::Router::new()
        .route(
            REST_SEND_MESSAGE_PATH,
            axum::routing::post(handle_send_message::<H>),
        )
        .route(
            REST_SEND_MESSAGE_LEGACY_PATH,
            axum::routing::post(handle_send_message::<H>),
        )
        .route(
            REST_STREAM_MESSAGE_PATH,
            axum::routing::post(handle_stream_message::<H>),
        )
        .route(
            REST_STREAM_MESSAGE_LEGACY_PATH,
            axum::routing::post(handle_stream_message::<H>),
        )
        .route(
            "/tasks/{id}",
            axum::routing::get(handle_get_task_or_subscribe::<H>)
                .post(handle_post_task_action::<H>),
        )
        .route("/tasks", axum::routing::get(handle_list_tasks::<H>))
        .route(
            "/tasks/{id}/cancel",
            axum::routing::post(handle_cancel_task_legacy::<H>),
        )
        .route(
            "/tasks/{id}/subscribe",
            axum::routing::get(handle_subscribe_to_task_legacy::<H>)
                .post(handle_subscribe_to_task_legacy::<H>),
        )
        .route(
            REST_PUSH_CONFIGS_PATH,
            axum::routing::post(handle_create_push_config::<H>).get(handle_list_push_configs::<H>),
        )
        .route(
            REST_PUSH_CONFIGS_LEGACY_PATH,
            axum::routing::post(handle_create_push_config::<H>).get(handle_list_push_configs::<H>),
        )
        .route(
            REST_PUSH_CONFIG_PATH,
            axum::routing::get(handle_get_push_config::<H>).delete(handle_delete_push_config::<H>),
        )
        .route(
            REST_PUSH_CONFIG_LEGACY_PATH,
            axum::routing::get(handle_get_push_config::<H>).delete(handle_delete_push_config::<H>),
        )
        .route(
            REST_EXTENDED_AGENT_CARD_PATH,
            axum::routing::get(handle_get_extended_agent_card::<H>),
        )
        .route(
            REST_EXTENDED_AGENT_CARD_LEGACY_PATH,
            axum::routing::get(handle_get_extended_agent_card::<H>),
        )
        .with_state(state)
}

async fn handle_send_message<H: RequestHandler>(
    State(state): State<RestState<H>>,
    headers: HeaderMap,
    Json(raw_req): Json<Value>,
) -> impl IntoResponse {
    let req = match protojson_conv::from_value::<SendMessageRequest>(raw_req) {
        Ok(req) => req,
        Err(e) => return rest_payload_error_response(e),
    };
    let params = extract_service_params(&headers);
    match state.handler.send_message(&params, req).await {
        Ok(resp) => protojson_json_response(&resp),
        Err(e) => rest_error_response(e),
    }
}

async fn handle_stream_message<H: RequestHandler>(
    State(state): State<RestState<H>>,
    headers: HeaderMap,
    Json(raw_req): Json<Value>,
) -> impl IntoResponse {
    let req = match protojson_conv::from_value::<SendMessageRequest>(raw_req) {
        Ok(req) => req,
        Err(e) => return rest_payload_error_response(e),
    };
    let params = extract_service_params(&headers);
    match state.handler.send_streaming_message(&params, req).await {
        Ok(stream) => sse::sse_from_stream(protojson_stream(stream)).into_response(),
        Err(e) => rest_error_response(e),
    }
}

#[derive(Deserialize)]
pub struct GetTaskQuery {
    #[serde(default, rename = "historyLength")]
    pub history_length: Option<i32>,
}

async fn handle_get_task_inner<H: RequestHandler>(
    state: RestState<H>,
    params: ServiceParams,
    id: String,
    query: GetTaskQuery,
) -> axum::response::Response {
    let req = GetTaskRequest {
        id,
        history_length: query.history_length,
        tenant: None,
    };
    match state.handler.get_task(&params, req).await {
        Ok(task) => protojson_json_response(&task),
        Err(e) => rest_error_response(e),
    }
}

async fn handle_get_task_or_subscribe<H: RequestHandler>(
    State(state): State<RestState<H>>,
    Path(id): Path<String>,
    Query(query): Query<GetTaskQuery>,
    headers: HeaderMap,
) -> impl IntoResponse {
    let params = extract_service_params(&headers);
    if let Some(task_id) = id.strip_suffix(":subscribe") {
        return handle_subscribe_to_task_inner(state, params, task_id.to_string()).await;
    }

    handle_get_task_inner(state, params, id, query).await
}

#[derive(Deserialize)]
pub struct ListTasksQuery {
    #[serde(default, rename = "contextId")]
    pub context_id: Option<String>,
    #[serde(default)]
    pub status: Option<String>,
    #[serde(default, rename = "pageSize")]
    pub page_size: Option<i32>,
    #[serde(default, rename = "pageToken")]
    pub page_token: Option<String>,
    #[serde(default, rename = "historyLength")]
    pub history_length: Option<i32>,

    #[serde(default, rename = "statusTimestampAfter")]
    pub status_timestamp_after: Option<DateTime<Utc>>,

    #[serde(default, rename = "includeArtifacts")]
    pub include_artifacts: Option<bool>,
}

async fn handle_list_tasks<H: RequestHandler>(
    State(state): State<RestState<H>>,
    Query(query): Query<ListTasksQuery>,
    headers: HeaderMap,
) -> impl IntoResponse {
    let params = extract_service_params(&headers);
    let status = query
        .status
        .and_then(|s| serde_json::from_value::<TaskState>(serde_json::Value::String(s)).ok());
    let req = ListTasksRequest {
        context_id: query.context_id,
        status,
        page_size: query.page_size,
        page_token: query.page_token,
        history_length: query.history_length,
        status_timestamp_after: query.status_timestamp_after,
        include_artifacts: query.include_artifacts,
        tenant: None,
    };
    match state.handler.list_tasks(&params, req).await {
        Ok(resp) => protojson_json_response(&resp),
        Err(e) => rest_error_response(e),
    }
}

async fn handle_cancel_task_inner<H: RequestHandler>(
    state: RestState<H>,
    params: ServiceParams,
    id: String,
) -> axum::response::Response {
    let req = CancelTaskRequest {
        id,
        metadata: None,
        tenant: None,
    };
    match state.handler.cancel_task(&params, req).await {
        Ok(task) => protojson_json_response(&task),
        Err(e) => rest_error_response(e),
    }
}

async fn handle_post_task_action<H: RequestHandler>(
    State(state): State<RestState<H>>,
    Path(id): Path<String>,
    headers: HeaderMap,
) -> impl IntoResponse {
    let params = extract_service_params(&headers);
    if let Some(task_id) = id.strip_suffix(":cancel") {
        return handle_cancel_task_inner(state, params, task_id.to_string()).await;
    }
    if let Some(task_id) = id.strip_suffix(":subscribe") {
        return handle_subscribe_to_task_inner(state, params, task_id.to_string()).await;
    }

    rest_error_response(A2AError::invalid_request("unsupported task action"))
}

async fn handle_cancel_task_legacy<H: RequestHandler>(
    State(state): State<RestState<H>>,
    Path(id): Path<String>,
    headers: HeaderMap,
) -> impl IntoResponse {
    let params = extract_service_params(&headers);
    handle_cancel_task_inner(state, params, id).await
}

async fn handle_subscribe_to_task_inner<H: RequestHandler>(
    state: RestState<H>,
    params: ServiceParams,
    id: String,
) -> axum::response::Response {
    let req = SubscribeToTaskRequest { id, tenant: None };
    match state.handler.subscribe_to_task(&params, req).await {
        Ok(stream) => sse::sse_from_stream(protojson_stream(stream)).into_response(),
        Err(e) => rest_error_response(e),
    }
}

async fn handle_subscribe_to_task_legacy<H: RequestHandler>(
    State(state): State<RestState<H>>,
    Path(id): Path<String>,
    headers: HeaderMap,
) -> impl IntoResponse {
    let params = extract_service_params(&headers);
    handle_subscribe_to_task_inner(state, params, id).await
}

#[derive(Deserialize)]
pub struct ListPushConfigsQuery {
    #[serde(default, rename = "pageSize")]
    pub page_size: Option<i32>,

    #[serde(default, rename = "pageToken")]
    pub page_token: Option<String>,
}

async fn handle_create_push_config<H: RequestHandler>(
    State(state): State<RestState<H>>,
    Path(id): Path<String>,
    headers: HeaderMap,
    Json(raw_config): Json<Value>,
) -> impl IntoResponse {
    let req = match parse_create_push_config_request(id, raw_config) {
        Ok(req) => req,
        Err(e) => return rest_payload_error_response(e),
    };
    let params = extract_service_params(&headers);
    match state.handler.create_push_config(&params, req).await {
        Ok(resp) => protojson_json_response(&resp),
        Err(e) => rest_error_response(e),
    }
}

async fn handle_get_push_config<H: RequestHandler>(
    State(state): State<RestState<H>>,
    Path((id, config_id)): Path<(String, String)>,
    headers: HeaderMap,
) -> impl IntoResponse {
    let params = extract_service_params(&headers);
    let req = GetTaskPushNotificationConfigRequest {
        task_id: id,
        id: config_id,
        tenant: None,
    };
    match state.handler.get_push_config(&params, req).await {
        Ok(resp) => protojson_json_response(&resp),
        Err(e) => rest_error_response(e),
    }
}

async fn handle_list_push_configs<H: RequestHandler>(
    State(state): State<RestState<H>>,
    Path(id): Path<String>,
    Query(query): Query<ListPushConfigsQuery>,
    headers: HeaderMap,
) -> impl IntoResponse {
    let params = extract_service_params(&headers);
    let req = ListTaskPushNotificationConfigsRequest {
        task_id: id,
        page_size: query.page_size,
        page_token: query.page_token,
        tenant: None,
    };
    match state.handler.list_push_configs(&params, req).await {
        Ok(resp) => protojson_json_response(&resp),
        Err(e) => rest_error_response(e),
    }
}

async fn handle_delete_push_config<H: RequestHandler>(
    State(state): State<RestState<H>>,
    Path((id, config_id)): Path<(String, String)>,
    headers: HeaderMap,
) -> impl IntoResponse {
    let params = extract_service_params(&headers);
    let req = DeleteTaskPushNotificationConfigRequest {
        task_id: id,
        id: config_id,
        tenant: None,
    };
    match state.handler.delete_push_config(&params, req).await {
        Ok(()) => StatusCode::OK.into_response(),
        Err(e) => rest_error_response(e),
    }
}

async fn handle_get_extended_agent_card<H: RequestHandler>(
    State(state): State<RestState<H>>,
    headers: HeaderMap,
) -> impl IntoResponse {
    let params = extract_service_params(&headers);
    let req = GetExtendedAgentCardRequest { tenant: None };
    match state.handler.get_extended_agent_card(&params, req).await {
        Ok(card) => protojson_json_response(&card),
        Err(e) => rest_error_response(e),
    }
}

fn protojson_json_response<T: ProtoJsonPayload>(value: &T) -> axum::response::Response {
    match protojson_conv::to_value(value) {
        Ok(payload) => Json(payload).into_response(),
        Err(e) => rest_error_response(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(
    task_id: String,
    raw_config: Value,
) -> Result<TaskPushNotificationConfig, String> {
    protojson_conv::from_value::<TaskPushNotificationConfig>(raw_config)
        .map_err(|e| e.to_string())
        .and_then(|mut config| {
            if !config.task_id.is_empty() && config.task_id != task_id {
                Err(format!(
                    "taskId in request body ({}) did not match task id in path ({task_id})",
                    config.task_id
                ))
            } else {
                config.task_id = task_id;
                Ok(config)
            }
        })
}

fn rest_payload_error_response(error: impl std::fmt::Display) -> axum::response::Response {
    let violation = FieldViolation {
        field: String::new(),
        description: error.to_string(),
    };
    rest_error_response(
        A2AError::invalid_params("invalid request body")
            .with_details(vec![TypedDetail::bad_request(vec![violation])]),
    )
}

fn rest_error_response(err: A2AError) -> axum::response::Response {
    let status =
        StatusCode::from_u16(err.http_status_code()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
    let reason = error_reason(err.code);
    let grpc_status = rest_grpc_status(err.code);

    let mut metadata = HashMap::from([("timestamp".to_string(), Utc::now().to_rfc3339())]);
    let mut details = Vec::new();

    if let Some(extra) = err.details {
        for detail in extra {
            if detail.type_url == errordetails::ERROR_INFO_TYPE {
                if let Some(Value::Object(meta)) = detail.value.get("metadata") {
                    for (k, v) in meta {
                        if let Some(s) = v.as_str() {
                            metadata.entry(k.clone()).or_insert_with(|| s.to_string());
                        }
                    }
                }
            } else {
                details.push(detail);
            }
        }
    }

    details.push(TypedDetail::error_info(
        reason,
        errordetails::PROTOCOL_DOMAIN,
        Some(metadata),
    ));

    let body = RestErrorEnvelope {
        error: RestErrorStatus {
            code: status.as_u16(),
            status: grpc_status,
            message: err.message,
            details,
        },
    };

    (status, Json(body)).into_response()
}

fn rest_grpc_status(code: i32) -> &'static str {
    match code {
        error_code::TASK_NOT_FOUND => "NOT_FOUND",
        error_code::METHOD_NOT_FOUND => "UNIMPLEMENTED",
        error_code::TASK_NOT_CANCELABLE
        | error_code::EXTENDED_CARD_NOT_CONFIGURED
        | error_code::EXTENSION_SUPPORT_REQUIRED
        | error_code::PUSH_NOTIFICATION_NOT_SUPPORTED
        | error_code::UNSUPPORTED_OPERATION
        | error_code::VERSION_NOT_SUPPORTED => "FAILED_PRECONDITION",
        error_code::CONTENT_TYPE_NOT_SUPPORTED
        | error_code::PARSE_ERROR
        | error_code::INVALID_REQUEST
        | error_code::INVALID_PARAMS => "INVALID_ARGUMENT",
        _ => "INTERNAL",
    }
}

#[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 futures::stream::BoxStream;
    use http_body_util::BodyExt;

    use crate::test_util::{CapturingHandler, assert_header_captured, install_crypto_provider};
    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(),
        ));
        rest_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()),
        );
        rest_router(handler)
    }

    #[tokio::test]
    async fn test_send_message() {
        let app = make_app();
        let body = serde_json::json!({
            "message": {
                "messageId": "m1",
                "role": "ROLE_USER",
                "parts": [{"text": "hello"}]
            }
        });
        let req = Request::builder()
            .uri(REST_SEND_MESSAGE_PATH)
            .method("POST")
            .header("content-type", "application/json")
            .body(Body::from(serde_json::to_string(&body).unwrap()))
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn test_send_message_legacy_alias() {
        let app = make_app();
        let body = serde_json::json!({
            "message": {
                "messageId": "m1",
                "role": "ROLE_USER",
                "parts": [{"text": "hello"}]
            }
        });
        let req = Request::builder()
            .uri(REST_SEND_MESSAGE_LEGACY_PATH)
            .method("POST")
            .header("content-type", "application/json")
            .body(Body::from(serde_json::to_string(&body).unwrap()))
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn test_get_task_not_found() {
        let app = make_app();
        let req = Request::builder()
            .uri("/tasks/nonexistent")
            .method("GET")
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::NOT_FOUND);
    }

    #[tokio::test]
    async fn test_list_tasks_empty() {
        let app = make_app();
        let req = Request::builder()
            .uri("/tasks")
            .method("GET")
            .body(Body::empty())
            .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 list_resp =
            protojson_conv::from_str::<ListTasksResponse>(std::str::from_utf8(&body).unwrap())
                .unwrap();
        assert!(list_resp.tasks.is_empty());
    }

    #[tokio::test]
    async fn test_cancel_task_not_found() {
        let app = make_app();
        let req = Request::builder()
            .uri("/tasks/nonexistent:cancel")
            .method("POST")
            .header("content-type", "application/json")
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::NOT_FOUND);
    }

    #[tokio::test]
    async fn test_create_push_config_not_supported() {
        let app = make_app();
        let body = serde_json::json!({
            "url": "http://example.com/callback"
        });
        let req = Request::builder()
            .uri("/tasks/t1/pushNotificationConfigs")
            .method("POST")
            .header("content-type", "application/json")
            .body(Body::from(serde_json::to_string(&body).unwrap()))
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
    }

    #[tokio::test]
    async fn test_create_push_config_accepts_raw_config_body() {
        let app = make_push_app();
        let body = serde_json::json!({
            "id": "cfg1",
            "url": "http://example.com/callback"
        });
        let req = Request::builder()
            .uri("/tasks/t1/pushNotificationConfigs")
            .method("POST")
            .header("content-type", "application/json")
            .body(Body::from(serde_json::to_string(&body).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 payload = serde_json::from_slice::<TaskPushNotificationConfig>(&body).unwrap();
        assert_eq!(payload.task_id, "t1");
        assert_eq!(payload.id.as_deref(), Some("cfg1"));
        assert_eq!(payload.url, "http://example.com/callback");
    }

    #[tokio::test]
    async fn test_create_push_config_accepts_flat_request_body() {
        let app = make_push_app();
        let body = serde_json::json!({
            "taskId": "t1",
            "id": "cfg1",
            "url": "http://example.com/callback"
        });
        let req = Request::builder()
            .uri("/tasks/t1/pushNotificationConfigs")
            .method("POST")
            .header("content-type", "application/json")
            .body(Body::from(serde_json::to_string(&body).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 payload = serde_json::from_slice::<TaskPushNotificationConfig>(&body).unwrap();
        assert_eq!(payload.task_id, "t1");
        assert_eq!(payload.id.as_deref(), Some("cfg1"));
        assert_eq!(payload.url, "http://example.com/callback");
    }

    #[tokio::test]
    async fn test_create_push_config_rejects_mismatched_task_id_in_body() {
        let app = make_push_app();
        let body = serde_json::json!({
            "taskId": "other-task",
            "url": "http://example.com/callback"
        });
        let req = Request::builder()
            .uri("/tasks/t1/pushNotificationConfigs")
            .method("POST")
            .header("content-type", "application/json")
            .body(Body::from(serde_json::to_string(&body).unwrap()))
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
    }

    #[tokio::test]
    async fn test_get_extended_agent_card() {
        let app = make_app();
        let req = Request::builder()
            .uri(REST_EXTENDED_AGENT_CARD_PATH)
            .method("GET")
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
    }

    #[tokio::test]
    async fn test_rest_error_response_status_codes() {
        let resp = rest_error_response(A2AError::task_not_found("t1"));
        assert_eq!(resp.status(), StatusCode::NOT_FOUND);

        let body = resp.into_body().collect().await.unwrap().to_bytes();
        let payload: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert_eq!(payload["error"]["code"], 404);
        assert_eq!(payload["error"]["status"], "NOT_FOUND");
        assert_eq!(payload["error"]["details"][0]["reason"], "TASK_NOT_FOUND");

        let resp = rest_error_response(A2AError::internal("boom"));
        assert_eq!(resp.status(), StatusCode::INTERNAL_SERVER_ERROR);

        let resp = rest_error_response(A2AError::content_type_not_supported());
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
    }

    #[tokio::test]
    async fn test_rest_error_response_merges_error_info_metadata() {
        let existing_info = TypedDetail::error_info(
            "TASK_NOT_FOUND",
            errordetails::PROTOCOL_DOMAIN,
            Some(HashMap::from([("taskId".to_string(), "t99".to_string())])),
        );
        let err = A2AError::task_not_found("t99").with_details(vec![existing_info]);
        let resp = rest_error_response(err);
        assert_eq!(resp.status(), StatusCode::NOT_FOUND);

        let body = resp.into_body().collect().await.unwrap().to_bytes();
        let payload: serde_json::Value = serde_json::from_slice(&body).unwrap();
        let details = payload["error"]["details"].as_array().unwrap();
        // Only one ErrorInfo should be present (merged)
        assert_eq!(details.len(), 1);
        assert_eq!(details[0]["@type"], errordetails::ERROR_INFO_TYPE);
        assert_eq!(details[0]["reason"], "TASK_NOT_FOUND");
        // Merged metadata should contain both taskId and timestamp
        assert_eq!(details[0]["metadata"]["taskId"], "t99");
        assert!(details[0]["metadata"]["timestamp"].is_string());
    }

    #[tokio::test]
    async fn test_rest_error_response_preserves_non_error_info_details() {
        let bad_req = TypedDetail::bad_request(vec![FieldViolation {
            field: "message.parts".into(),
            description: "required".into(),
        }]);
        let err = A2AError::invalid_params("bad request").with_details(vec![bad_req]);
        let resp = rest_error_response(err);
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);

        let body = resp.into_body().collect().await.unwrap().to_bytes();
        let payload: serde_json::Value = serde_json::from_slice(&body).unwrap();
        let details = payload["error"]["details"].as_array().unwrap();
        // BadRequest detail + auto-generated ErrorInfo
        assert_eq!(details.len(), 2);
        assert_eq!(details[0]["@type"], errordetails::BAD_REQUEST_TYPE);
        assert_eq!(details[0]["fieldViolations"][0]["field"], "message.parts");
        assert_eq!(details[1]["@type"], errordetails::ERROR_INFO_TYPE);
        assert_eq!(details[1]["reason"], "INVALID_PARAMS");
    }

    #[tokio::test]
    async fn test_rest_payload_error_response() {
        let resp = rest_payload_error_response("missing field 'parts'");
        assert_eq!(resp.status(), StatusCode::BAD_REQUEST);

        let body = resp.into_body().collect().await.unwrap().to_bytes();
        let payload: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert_eq!(
            payload["error"]["message"].as_str().unwrap(),
            "invalid request body"
        );
        let details = payload["error"]["details"].as_array().unwrap();
        // BadRequest detail + ErrorInfo
        assert_eq!(details.len(), 2);
        assert_eq!(details[0]["@type"], errordetails::BAD_REQUEST_TYPE);
        assert_eq!(
            details[0]["fieldViolations"][0]["description"],
            "missing field 'parts'"
        );
        assert_eq!(details[1]["@type"], errordetails::ERROR_INFO_TYPE);
        assert_eq!(details[1]["reason"], "INVALID_PARAMS");
    }

    #[tokio::test]
    async fn test_subscribe_to_task_not_found() {
        let app = make_app();
        let req = Request::builder()
            .uri("/tasks/nonexistent:subscribe")
            .method("GET")
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert!(resp.status().is_client_error());
    }

    #[tokio::test]
    async fn test_stream_message() {
        let app = make_app();
        let body = serde_json::json!({
            "message": {
                "messageId": "m1",
                "role": "ROLE_USER",
                "parts": [{"text": "hello"}]
            }
        });
        let req = Request::builder()
            .uri(REST_STREAM_MESSAGE_PATH)
            .method("POST")
            .header("content-type", "application/json")
            .header("accept", "text/event-stream")
            .body(Body::from(serde_json::to_string(&body).unwrap()))
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn test_get_push_config_not_found() {
        let app = make_app();
        let req = Request::builder()
            .uri("/tasks/t1/pushNotificationConfigs/cfg1")
            .method("GET")
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        // Push not supported → bad request
        assert!(resp.status().is_client_error());
    }

    #[tokio::test]
    async fn test_list_push_configs() {
        let app = make_push_app();
        let req = Request::builder()
            .uri("/tasks/t1/pushNotificationConfigs")
            .method("GET")
            .body(Body::empty())
            .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 payload =
            serde_json::from_slice::<ListTaskPushNotificationConfigsResponse>(&body).unwrap();
        assert!(payload.configs.is_empty());
    }

    #[tokio::test]
    async fn test_delete_push_config() {
        let app = make_app();
        let req = Request::builder()
            .uri("/tasks/t1/pushNotificationConfigs/cfg1")
            .method("DELETE")
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert!(resp.status().is_client_error());
    }

    #[tokio::test]
    async fn test_delete_push_config_with_store_returns_ok() {
        let app = make_push_app();
        let create_body = serde_json::json!({
            "id": "cfg1",
            "url": "http://example.com/callback"
        });
        let create_req = Request::builder()
            .uri("/tasks/t1/pushNotificationConfigs")
            .method("POST")
            .header("content-type", "application/json")
            .body(Body::from(serde_json::to_string(&create_body).unwrap()))
            .unwrap();
        let create_resp = app.clone().oneshot(create_req).await.unwrap();
        assert_eq!(create_resp.status(), StatusCode::OK);

        let delete_req = Request::builder()
            .uri("/tasks/t1/pushNotificationConfigs/cfg1")
            .method("DELETE")
            .body(Body::empty())
            .unwrap();
        let delete_resp = app.clone().oneshot(delete_req).await.unwrap();
        assert_eq!(delete_resp.status(), StatusCode::OK);

        let list_req = Request::builder()
            .uri("/tasks/t1/pushNotificationConfigs")
            .method("GET")
            .body(Body::empty())
            .unwrap();
        let list_resp = app.oneshot(list_req).await.unwrap();
        assert_eq!(list_resp.status(), StatusCode::OK);
        let body = list_resp.into_body().collect().await.unwrap().to_bytes();
        let payload =
            serde_json::from_slice::<ListTaskPushNotificationConfigsResponse>(&body).unwrap();
        assert!(payload.configs.is_empty());
    }

    #[tokio::test]
    async fn test_list_tasks_with_query_params() {
        let app = make_app();
        let req = Request::builder()
            .uri("/tasks?contextId=c1&pageSize=5&pageToken=tok&historyLength=3&includeArtifacts=true&statusTimestampAfter=2025-01-01T00:00:00Z")
            .method("GET")
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
    }

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

        let body = serde_json::json!({
            "message": {
                "messageId": "m1",
                "role": "ROLE_USER",
                "parts": [{"text": "hi"}]
            }
        });
        let req = Request::builder()
            .uri(REST_SEND_MESSAGE_PATH)
            .method("POST")
            .header("content-type", "application/json")
            .header("authorization", "Bearer test-token-abc")
            .header("x-tenant-id", "acme")
            .body(Body::from(serde_json::to_string(&body).unwrap()))
            .unwrap();
        let _ = app.oneshot(req).await.unwrap();

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

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

        let body = serde_json::json!({
            "message": {
                "messageId": "m1",
                "role": "ROLE_USER",
                "parts": [{"text": "hi"}]
            }
        });
        let req = Request::builder()
            .uri(REST_STREAM_MESSAGE_PATH)
            .method("POST")
            .header("content-type", "application/json")
            .header("accept", "text/event-stream")
            .header("authorization", "Bearer streaming-token")
            .body(Body::from(serde_json::to_string(&body).unwrap()))
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        // Drain the SSE body to ensure the handler has been fully driven.
        let _ = resp.into_body().collect().await.unwrap().to_bytes();

        assert_header_captured(
            &captured,
            "send_streaming_message",
            "authorization",
            "Bearer streaming-token",
        );
    }

    #[tokio::test]
    async fn test_get_task_propagates_headers_as_service_params() {
        let (handler, captured) = CapturingHandler::new();
        let app = rest_router(handler);
        let req = Request::builder()
            .uri("/tasks/some-id")
            .method("GET")
            .header("authorization", "Bearer get-task-token")
            .header("x-tenant-id", "acme")
            .body(Body::empty())
            .unwrap();
        let _ = app.oneshot(req).await.unwrap();
        assert_header_captured(
            &captured,
            "get_task",
            "authorization",
            "Bearer get-task-token",
        );
        assert_header_captured(&captured, "get_task", "x-tenant-id", "acme");
    }

    #[tokio::test]
    async fn test_list_tasks_propagates_headers_as_service_params() {
        let (handler, captured) = CapturingHandler::new();
        let app = rest_router(handler);
        let req = Request::builder()
            .uri("/tasks")
            .method("GET")
            .header("authorization", "Bearer list-tasks-token")
            .body(Body::empty())
            .unwrap();
        let _ = app.oneshot(req).await.unwrap();
        assert_header_captured(
            &captured,
            "list_tasks",
            "authorization",
            "Bearer list-tasks-token",
        );
    }

    #[tokio::test]
    async fn test_cancel_task_propagates_headers_as_service_params() {
        let (handler, captured) = CapturingHandler::new();
        let app = rest_router(handler);
        let req = Request::builder()
            .uri("/tasks/some-id:cancel")
            .method("POST")
            .header("authorization", "Bearer cancel-token")
            .body(Body::empty())
            .unwrap();
        let _ = app.oneshot(req).await.unwrap();
        assert_header_captured(
            &captured,
            "cancel_task",
            "authorization",
            "Bearer cancel-token",
        );
    }

    #[tokio::test]
    async fn test_cancel_task_legacy_propagates_headers_as_service_params() {
        // /tasks/{id}/cancel routes through handle_cancel_task_legacy.
        let (handler, captured) = CapturingHandler::new();
        let app = rest_router(handler);
        let req = Request::builder()
            .uri("/tasks/some-id/cancel")
            .method("POST")
            .header("authorization", "Bearer cancel-legacy-token")
            .body(Body::empty())
            .unwrap();
        let _ = app.oneshot(req).await.unwrap();
        assert_header_captured(
            &captured,
            "cancel_task",
            "authorization",
            "Bearer cancel-legacy-token",
        );
    }

    #[tokio::test]
    async fn test_create_push_config_propagates_headers_as_service_params() {
        let (handler, captured) = CapturingHandler::new();
        let app = rest_router(handler);
        let body = serde_json::json!({
            "id": "cfg1",
            "url": "http://example.com/callback"
        });
        let req = Request::builder()
            .uri("/tasks/t1/pushNotificationConfigs")
            .method("POST")
            .header("content-type", "application/json")
            .header("authorization", "Bearer push-create-token")
            .body(Body::from(serde_json::to_string(&body).unwrap()))
            .unwrap();
        let _ = app.oneshot(req).await.unwrap();
        assert_header_captured(
            &captured,
            "create_push_config",
            "authorization",
            "Bearer push-create-token",
        );
    }

    #[tokio::test]
    async fn test_list_push_configs_propagates_headers_as_service_params() {
        let (handler, captured) = CapturingHandler::new();
        let app = rest_router(handler);
        let req = Request::builder()
            .uri("/tasks/t1/pushNotificationConfigs")
            .method("GET")
            .header("authorization", "Bearer push-list-token")
            .body(Body::empty())
            .unwrap();
        let _ = app.oneshot(req).await.unwrap();
        assert_header_captured(
            &captured,
            "list_push_configs",
            "authorization",
            "Bearer push-list-token",
        );
    }

    #[tokio::test]
    async fn test_get_push_config_propagates_headers_as_service_params() {
        let (handler, captured) = CapturingHandler::new();
        let app = rest_router(handler);
        let req = Request::builder()
            .uri("/tasks/t1/pushNotificationConfigs/cfg1")
            .method("GET")
            .header("authorization", "Bearer push-get-token")
            .body(Body::empty())
            .unwrap();
        let _ = app.oneshot(req).await.unwrap();
        assert_header_captured(
            &captured,
            "get_push_config",
            "authorization",
            "Bearer push-get-token",
        );
    }

    #[tokio::test]
    async fn test_delete_push_config_propagates_headers_as_service_params() {
        let (handler, captured) = CapturingHandler::new();
        let app = rest_router(handler);
        let req = Request::builder()
            .uri("/tasks/t1/pushNotificationConfigs/cfg1")
            .method("DELETE")
            .header("authorization", "Bearer push-delete-token")
            .body(Body::empty())
            .unwrap();
        let _ = app.oneshot(req).await.unwrap();
        assert_header_captured(
            &captured,
            "delete_push_config",
            "authorization",
            "Bearer push-delete-token",
        );
    }

    #[tokio::test]
    async fn test_get_extended_agent_card_propagates_headers_as_service_params() {
        let (handler, captured) = CapturingHandler::new();
        let app = rest_router(handler);
        let req = Request::builder()
            .uri(REST_EXTENDED_AGENT_CARD_PATH)
            .method("GET")
            .header("authorization", "Bearer card-token")
            .body(Body::empty())
            .unwrap();
        let _ = app.oneshot(req).await.unwrap();
        assert_header_captured(
            &captured,
            "get_extended_agent_card",
            "authorization",
            "Bearer card-token",
        );
    }

    #[tokio::test]
    async fn test_subscribe_to_task_propagates_headers_as_service_params() {
        let (handler, captured) = CapturingHandler::new();
        let app = rest_router(handler);
        let req = Request::builder()
            .uri("/tasks/some-id:subscribe")
            .method("GET")
            .header("authorization", "Bearer subscribe-token")
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        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_subscribe_to_task_legacy_propagates_headers_as_service_params() {
        // /tasks/{id}/subscribe routes through handle_subscribe_to_task_legacy.
        let (handler, captured) = CapturingHandler::new();
        let app = rest_router(handler);
        let req = Request::builder()
            .uri("/tasks/some-id/subscribe")
            .method("GET")
            .header("authorization", "Bearer subscribe-legacy-token")
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        let _ = resp.into_body().collect().await.unwrap().to_bytes();
        assert_header_captured(
            &captured,
            "subscribe_to_task",
            "authorization",
            "Bearer subscribe-legacy-token",
        );
    }

    #[tokio::test]
    async fn test_get_task_with_history_length() {
        // First create a task, then get it with historyLength
        let app = make_app();
        // Create task via send_message
        let body = serde_json::json!({
            "message": {
                "messageId": "m1",
                "role": "ROLE_USER",
                "parts": [{"text": "hello"}]
            }
        });
        let req = Request::builder()
            .uri(REST_SEND_MESSAGE_PATH)
            .method("POST")
            .header("content-type", "application/json")
            .body(Body::from(serde_json::to_string(&body).unwrap()))
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
        let resp_body = resp.into_body().collect().await.unwrap().to_bytes();
        let send_resp: serde_json::Value = serde_json::from_slice(&resp_body).unwrap();
        let task_id = send_resp["task"]["id"].as_str().unwrap();

        let app = make_app();
        let req = Request::builder()
            .uri(format!("/tasks/{}?historyLength=5", task_id))
            .method("GET")
            .body(Body::empty())
            .unwrap();
        let resp = app.oneshot(req).await.unwrap();
        // Each make_app() creates a new store, so this will be not found
        assert!(resp.status() == StatusCode::OK || resp.status() == StatusCode::NOT_FOUND);
    }
}