agent-client-protocol 2.0.0

//! Error handling tests for JSON-RPC layer
//!
//! Tests various error conditions:
//! - Invalid JSON
//! - Unknown methods
//! - Handler-returned errors
//! - Serialization failures
//! - Missing/invalid parameters

use agent_client_protocol::{
    Channel, ConnectionTo, Dispatch, Handled, JsonRpcMessage, JsonRpcRequest, JsonRpcResponse,
    RawJsonRpcMessage, Responder, SentRequest, TransportFrame, role::UntypedRole,
};
use expect_test::expect;
use futures::{AsyncRead, AsyncWrite, StreamExt as _};
use serde::{Deserialize, Serialize};
use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};

/// Test helper to block and wait for a JSON-RPC response.
async fn recv<T: JsonRpcResponse + Send>(
    response: SentRequest<T>,
) -> Result<T, agent_client_protocol::Error> {
    let (tx, rx) = tokio::sync::oneshot::channel();
    response.on_receiving_result(async move |result| {
        tx.send(result)
            .map_err(|_| agent_client_protocol::Error::internal_error())
    })?;
    rx.await
        .map_err(|_| agent_client_protocol::Error::internal_error())?
}

/// Helper to set up test streams.
fn setup_test_streams() -> (
    impl AsyncRead,
    impl AsyncWrite,
    impl AsyncRead,
    impl AsyncWrite,
) {
    let (client_writer, server_reader) = tokio::io::duplex(1024);
    let (server_writer, client_reader) = tokio::io::duplex(1024);

    let server_reader = server_reader.compat();
    let server_writer = server_writer.compat_write();
    let client_reader = client_reader.compat();
    let client_writer = client_writer.compat_write();

    (server_reader, server_writer, client_reader, client_writer)
}

#[tokio::test(flavor = "current_thread")]
async fn response_dispatch_handler_error_reaches_the_local_request_awaiter() {
    tokio::task::LocalSet::new()
        .run_until(async {
            let (transport, mut peer) = Channel::duplex();
            let connection = UntypedRole
                .builder()
                .on_receive_dispatch(
                    async |dispatch: Dispatch, _connection: ConnectionTo<UntypedRole>| {
                        match dispatch {
                            Dispatch::Response(..) => {
                                Err(agent_client_protocol::Error::internal_error()
                                    .data("response interceptor failed"))
                            }
                            message => Ok(Handled::No {
                                message,
                                retry: false,
                            }),
                        }
                    },
                    agent_client_protocol::on_receive_dispatch!(),
                )
                .connect_with(transport, async |connection| {
                    let error = connection
                        .send_request(SimpleRequest {
                            message: "intercept me".into(),
                        })
                        .block_task()
                        .await
                        .expect_err("response handler error should fail the pending request");
                    assert_eq!(
                        error.data,
                        Some(serde_json::json!("response interceptor failed"))
                    );
                    Ok(())
                });

            let peer = async move {
                let frame = peer
                    .rx
                    .next()
                    .await
                    .expect("connection should send one request");
                let TransportFrame::Single(RawJsonRpcMessage::Request(request)) = frame else {
                    panic!("expected one standalone request");
                };
                peer.tx
                    .unbounded_send(TransportFrame::Single(RawJsonRpcMessage::response(
                        request.id,
                        Ok(serde_json::json!({ "result": "ignored" })),
                    )))
                    .expect("connection should accept the test response");
                Ok::<(), agent_client_protocol::Error>(())
            };

            futures::try_join!(connection, peer)?;
            Ok::<(), agent_client_protocol::Error>(())
        })
        .await
        .unwrap();
}

// ============================================================================
// Test types
// ============================================================================

#[derive(Debug, Clone, Serialize, Deserialize)]
struct SimpleRequest {
    message: String,
}

impl JsonRpcMessage for SimpleRequest {
    fn matches_method(method: &str) -> bool {
        method == "simple_method"
    }

    fn method(&self) -> &'static str {
        "simple_method"
    }

    fn to_untyped_message(
        &self,
    ) -> Result<agent_client_protocol::UntypedMessage, agent_client_protocol::Error> {
        agent_client_protocol::UntypedMessage::new(self.method(), self)
    }

    fn parse_message(
        method: &str,
        params: &impl serde::Serialize,
    ) -> Result<Self, agent_client_protocol::Error> {
        if !Self::matches_method(method) {
            return Err(agent_client_protocol::Error::method_not_found());
        }
        agent_client_protocol::util::json_cast_params(params)
    }
}

impl JsonRpcRequest for SimpleRequest {
    type Response = SimpleResponse;
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct SimpleResponse {
    result: String,
}

impl JsonRpcResponse for SimpleResponse {
    fn into_json(self, _method: &str) -> Result<serde_json::Value, agent_client_protocol::Error> {
        serde_json::to_value(self).map_err(agent_client_protocol::Error::into_internal_error)
    }

    fn from_value(
        _method: &str,
        value: serde_json::Value,
    ) -> Result<Self, agent_client_protocol::Error> {
        agent_client_protocol::util::json_cast(&value)
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct SimpleNotification {
    message: String,
}

impl JsonRpcMessage for SimpleNotification {
    fn matches_method(method: &str) -> bool {
        method == "simple_notification"
    }

    fn method(&self) -> &'static str {
        "simple_notification"
    }

    fn to_untyped_message(
        &self,
    ) -> Result<agent_client_protocol::UntypedMessage, agent_client_protocol::Error> {
        agent_client_protocol::UntypedMessage::new(self.method(), self)
    }

    fn parse_message(
        method: &str,
        params: &impl serde::Serialize,
    ) -> Result<Self, agent_client_protocol::Error> {
        if !Self::matches_method(method) {
            return Err(agent_client_protocol::Error::method_not_found());
        }
        agent_client_protocol::util::json_cast_params(params)
    }
}

impl agent_client_protocol::JsonRpcNotification for SimpleNotification {}

// ============================================================================
// Test 1: Invalid JSON (complete line with parse error)
// ============================================================================

#[tokio::test(flavor = "current_thread")]
async fn test_invalid_json() {
    use tokio::io::{AsyncReadExt, AsyncWriteExt};
    use tokio::task::LocalSet;

    let local = LocalSet::new();

    local
        .run_until(async {
            // Create duplex streams for bidirectional communication
            let (mut client_writer, server_reader) = tokio::io::duplex(1024);
            let (server_writer, mut client_reader) = tokio::io::duplex(1024);

            let server_reader = server_reader.compat();
            let server_writer = server_writer.compat_write();

            // No handlers - all requests will return errors
            let server_transport =
                agent_client_protocol::ByteStreams::new(server_writer, server_reader);
            let server = UntypedRole.builder();

            // Spawn server
            tokio::task::spawn_local(async move {
                drop(server.connect_to(server_transport).await);
            });

            // Send invalid JSON
            let invalid_json = b"{\"method\": \"test\", \"id\": 1, INVALID}\n";
            client_writer.write_all(invalid_json).await.unwrap();
            client_writer.flush().await.unwrap();

            // Read response
            let mut buffer = vec![0u8; 1024];
            let n = client_reader.read(&mut buffer).await.unwrap();
            let response_str = String::from_utf8_lossy(&buffer[..n]);

            // Parse as JSON and verify structure
            let response: serde_json::Value =
                serde_json::from_str(response_str.trim()).expect("Response should be valid JSON");

            // Use expect_test to verify the exact structure
            expect![[r#"
                {
                  "jsonrpc": "2.0",
                  "id": null,
                  "error": {
                    "code": -32700,
                    "message": "Parse error",
                    "data": {
                      "line": "{\"method\": \"test\", \"id\": 1, INVALID}"
                    }
                  }
                }"#]]
            .assert_eq(&serde_json::to_string_pretty(&response).unwrap());
        })
        .await;
}

// ============================================================================
// Test 1b: Incomplete line (EOF mid-message)
// ============================================================================

#[tokio::test]
async fn test_incomplete_line() {
    use futures::io::Cursor;

    // Incomplete JSON input - no newline, simulates client disconnect
    let incomplete_json = b"{\"method\": \"test\", \"id\": 1";
    let input = Cursor::new(incomplete_json.to_vec());
    let output = Cursor::new(Vec::new());

    // No handlers needed for EOF test
    let transport = agent_client_protocol::ByteStreams::new(output, input);
    let connection = UntypedRole.builder();

    // The server should handle EOF mid-message gracefully
    let result = connection.connect_to(transport).await;

    // Server should terminate cleanly (not hang) when EOF is hit mid-message
    assert!(
        result.is_ok(),
        "expected clean shutdown on EOF, got: {result:?}"
    );
}

#[tokio::test(flavor = "current_thread")]
async fn malformed_standalone_response_is_ignored() {
    use tokio::io::{AsyncBufReadExt as _, AsyncWriteExt as _, BufReader};
    use tokio::task::LocalSet;

    LocalSet::new()
        .run_until(async {
            let (mut client_writer, server_reader) = tokio::io::duplex(4096);
            let (server_writer, client_reader) = tokio::io::duplex(4096);
            let server = UntypedRole.builder().on_receive_request(
                async |request: SimpleRequest,
                       responder: Responder<SimpleResponse>,
                       _cx: ConnectionTo<UntypedRole>| {
                    responder.respond(SimpleResponse {
                        result: format!("echo: {}", request.message),
                    })
                },
                agent_client_protocol::on_receive_request!(),
            );
            let server_task = tokio::task::spawn_local(server.connect_to(
                agent_client_protocol::ByteStreams::new(
                    server_writer.compat_write(),
                    server_reader.compat(),
                ),
            ));

            for message in [
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 99,
                    "result": null,
                    "error": { "code": -32603, "message": "Internal error" }
                }),
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": 41,
                    "method": "simple_method",
                    "params": { "message": "after malformed response" }
                }),
            ] {
                let mut bytes =
                    serde_json::to_vec(&message).expect("test message should serialize");
                bytes.push(b'\n');
                client_writer
                    .write_all(&bytes)
                    .await
                    .expect("test message should be written");
            }
            client_writer
                .flush()
                .await
                .expect("test messages should be flushed");

            let mut client_reader = BufReader::new(client_reader);
            let mut line = String::new();
            tokio::time::timeout(
                std::time::Duration::from_secs(10),
                client_reader.read_line(&mut line),
            )
            .await
            .expect("timed out waiting for the valid request response")
            .expect("response read should succeed");
            let response: serde_json::Value =
                serde_json::from_str(line.trim()).expect("response should be valid JSON");
            assert_eq!(response["id"], 41);
            assert_eq!(
                response["result"]["result"],
                "echo: after malformed response"
            );

            drop(client_writer);
            drop(client_reader);
            tokio::time::timeout(std::time::Duration::from_secs(10), server_task)
                .await
                .expect("server did not stop after EOF")
                .expect("server task panicked")
                .expect("server connection failed");
        })
        .await;
}

// ============================================================================
// Test 2: Unknown method (no handler claims)
// ============================================================================

#[tokio::test(flavor = "current_thread")]
async fn test_unknown_method() {
    use tokio::task::LocalSet;

    let local = LocalSet::new();

    local
        .run_until(async {
            let (server_reader, server_writer, client_reader, client_writer) = setup_test_streams();

            // No handlers - all requests will be "method not found"
            let server_transport =
                agent_client_protocol::ByteStreams::new(server_writer, server_reader);
            let server = UntypedRole.builder();
            let client_transport =
                agent_client_protocol::ByteStreams::new(client_writer, client_reader);
            let client = UntypedRole.builder();

            // Spawn server
            tokio::task::spawn_local(async move {
                server.connect_to(server_transport).await.ok();
            });

            // Send request from client
            let result = client
                .connect_with(
                    client_transport,
                    async |cx| -> Result<(), agent_client_protocol::Error> {
                        let request = SimpleRequest {
                            message: "test".to_string(),
                        };

                        let result: Result<SimpleResponse, _> =
                            recv(cx.send_request(request)).await;

                        // Should get an error because no handler claims the method
                        assert!(result.is_err());
                        if let Err(err) = result {
                            // Should be "method not found" or similar error
                            assert!(matches!(
                                err.code,
                                agent_client_protocol::ErrorCode::MethodNotFound
                            ));
                        }
                        Ok(())
                    },
                )
                .await;

            assert!(result.is_ok(), "Test failed: {result:?}");
        })
        .await;
}

// ============================================================================
// Test 3: Handler returns error
// ============================================================================

#[derive(Debug, Clone, Serialize, Deserialize)]
struct ErrorRequest {
    value: String,
}

impl JsonRpcMessage for ErrorRequest {
    fn matches_method(method: &str) -> bool {
        method == "error_method"
    }

    fn method(&self) -> &'static str {
        "error_method"
    }

    fn to_untyped_message(
        &self,
    ) -> Result<agent_client_protocol::UntypedMessage, agent_client_protocol::Error> {
        agent_client_protocol::UntypedMessage::new(self.method(), self)
    }

    fn parse_message(
        method: &str,
        params: &impl serde::Serialize,
    ) -> Result<Self, agent_client_protocol::Error> {
        if !Self::matches_method(method) {
            return Err(agent_client_protocol::Error::method_not_found());
        }
        agent_client_protocol::util::json_cast_params(params)
    }
}

impl JsonRpcRequest for ErrorRequest {
    type Response = SimpleResponse;
}

#[tokio::test(flavor = "current_thread")]
async fn test_handler_returns_error() {
    use tokio::task::LocalSet;

    let local = LocalSet::new();

    local
        .run_until(async {
            let (server_reader, server_writer, client_reader, client_writer) = setup_test_streams();

            let server_transport =
                agent_client_protocol::ByteStreams::new(server_writer, server_reader);
            let server = UntypedRole.builder().on_receive_request(
                async |_request: ErrorRequest,
                       responder: Responder<SimpleResponse>,
                       _connection: ConnectionTo<UntypedRole>| {
                    // Explicitly return an error
                    responder.respond_with_error(agent_client_protocol::Error::internal_error())
                },
                agent_client_protocol::on_receive_request!(),
            );

            let client_transport =
                agent_client_protocol::ByteStreams::new(client_writer, client_reader);
            let client = UntypedRole.builder();

            tokio::task::spawn_local(async move {
                server.connect_to(server_transport).await.ok();
            });

            let result = client
                .connect_with(
                    client_transport,
                    async |cx| -> Result<(), agent_client_protocol::Error> {
                        let request = ErrorRequest {
                            value: "trigger error".to_string(),
                        };

                        let result: Result<SimpleResponse, _> =
                            recv(cx.send_request(request)).await;

                        // Should get the error the handler returned
                        assert!(result.is_err());
                        if let Err(err) = result {
                            assert!(matches!(
                                err.code,
                                agent_client_protocol::ErrorCode::InternalError
                            ));
                        }
                        Ok(())
                    },
                )
                .await;

            assert!(result.is_ok(), "Test failed: {result:?}");
        })
        .await;
}

// ============================================================================
// Test 4: Request without required params
// ============================================================================

#[derive(Debug, Clone, Serialize, Deserialize)]
struct EmptyRequest;

impl JsonRpcMessage for EmptyRequest {
    fn matches_method(method: &str) -> bool {
        method == "strict_method"
    }

    fn method(&self) -> &'static str {
        "strict_method"
    }

    fn to_untyped_message(
        &self,
    ) -> Result<agent_client_protocol::UntypedMessage, agent_client_protocol::Error> {
        agent_client_protocol::UntypedMessage::new(self.method(), self)
    }

    fn parse_message(
        method: &str,
        _params: &impl serde::Serialize,
    ) -> Result<Self, agent_client_protocol::Error> {
        if !Self::matches_method(method) {
            return Err(agent_client_protocol::Error::method_not_found());
        }
        Ok(EmptyRequest)
    }
}

impl JsonRpcRequest for EmptyRequest {
    type Response = SimpleResponse;
}

#[tokio::test(flavor = "current_thread")]
async fn test_missing_required_params() {
    use tokio::task::LocalSet;

    let local = LocalSet::new();

    local
        .run_until(async {
            let (server_reader, server_writer, client_reader, client_writer) = setup_test_streams();

            // Handler that validates params - since EmptyRequest has no params but we're checking
            // against SimpleRequest which requires a message field, this will fail
            let server_transport =
                agent_client_protocol::ByteStreams::new(server_writer, server_reader);
            let server = UntypedRole.builder().on_receive_request(
                async |_request: EmptyRequest,
                       responder: Responder<SimpleResponse>,
                       _connection: ConnectionTo<UntypedRole>| {
                    // This will be called, but EmptyRequest parsing already succeeded
                    // The test is actually checking if EmptyRequest (no params) fails to parse as SimpleRequest
                    // But with the new API, EmptyRequest parses successfully since it expects no params
                    // We need to manually check - but actually the parse_request for EmptyRequest
                    // accepts anything for "strict_method", so the error must come from somewhere else
                    responder.respond_with_error(agent_client_protocol::Error::invalid_params())
                },
                agent_client_protocol::on_receive_request!(),
            );

            let client_transport =
                agent_client_protocol::ByteStreams::new(client_writer, client_reader);
            let client = UntypedRole.builder();

            tokio::task::spawn_local(async move {
                server.connect_to(server_transport).await.ok();
            });

            let result = client
                .connect_with(
                    client_transport,
                    async |cx| -> Result<(), agent_client_protocol::Error> {
                        // Send request with no params (EmptyRequest has no fields)
                        let request = EmptyRequest;

                        let result: Result<SimpleResponse, _> =
                            recv(cx.send_request(request)).await;

                        // Should get invalid_params error
                        assert!(result.is_err());
                        if let Err(err) = result {
                            assert!(matches!(
                                err.code,
                                agent_client_protocol::ErrorCode::InvalidParams
                            )); // JSONRPC_INVALID_PARAMS
                        }
                        Ok(())
                    },
                )
                .await;

            assert!(result.is_ok(), "Test failed: {result:?}");
        })
        .await;
}

// ============================================================================
// Test 5: Invalid params returns error but connection stays alive (issue #131)
// ============================================================================

#[tokio::test(flavor = "current_thread")]
async fn test_invalid_params_keeps_connection_alive() {
    use tokio::io::{AsyncReadExt, AsyncWriteExt};
    use tokio::task::LocalSet;

    let local = LocalSet::new();

    local
        .run_until(async {
            let (mut client_writer, server_reader) = tokio::io::duplex(4096);
            let (server_writer, mut client_reader) = tokio::io::duplex(4096);

            let server_reader = server_reader.compat();
            let server_writer = server_writer.compat_write();

            // Register a handler for SimpleRequest (requires "message" field)
            let server_transport =
                agent_client_protocol::ByteStreams::new(server_writer, server_reader);
            let server = UntypedRole.builder().on_receive_request(
                async |request: SimpleRequest,
                       responder: Responder<SimpleResponse>,
                       _connection: ConnectionTo<UntypedRole>| {
                    responder.respond(SimpleResponse {
                        result: format!("echo: {}", request.message),
                    })
                },
                agent_client_protocol::on_receive_request!(),
            );

            tokio::task::spawn_local(async move {
                drop(server.connect_to(server_transport).await);
            });

            // 1) Send a request with WRONG params (missing "message" field)
            let bad_request =
                b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"simple_method\",\"params\":{\"wrong_field\":\"hello\"}}\n";
            client_writer.write_all(bad_request).await.unwrap();
            client_writer.flush().await.unwrap();

            // Read the error response
            let mut buffer = vec![0u8; 4096];
            let n = client_reader.read(&mut buffer).await.unwrap();
            let response_str = String::from_utf8_lossy(&buffer[..n]);
            let response: serde_json::Value =
                serde_json::from_str(response_str.trim()).expect("Response should be valid JSON");

            // Verify it's an error response with the correct id and error code
            assert_eq!(response["id"], 1);
            assert!(response["error"].is_object(), "Expected error object");
            assert_eq!(
                response["error"]["code"], -32602,
                "Expected invalid params (-32602)"
            );

            // 2) Send a VALID request to prove the connection is still alive
            let good_request =
                b"{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"simple_method\",\"params\":{\"message\":\"hello\"}}\n";
            client_writer.write_all(good_request).await.unwrap();
            client_writer.flush().await.unwrap();

            // Read the success response
            let n = client_reader.read(&mut buffer).await.unwrap();
            let response_str = String::from_utf8_lossy(&buffer[..n]);
            let response: serde_json::Value =
                serde_json::from_str(response_str.trim()).expect("Response should be valid JSON");

            // Verify it's a success response
            assert_eq!(response["id"], 2);
            assert_eq!(response["result"]["result"], "echo: hello");
        })
        .await;
}

// ============================================================================
// Helpers for raw-wire tests
// ============================================================================

async fn read_jsonrpc_response_line(
    reader: &mut tokio::io::BufReader<tokio::io::DuplexStream>,
) -> serde_json::Value {
    use tokio::io::AsyncBufReadExt as _;

    let mut line = String::new();
    match tokio::time::timeout(
        tokio::time::Duration::from_secs(1),
        reader.read_line(&mut line),
    )
    .await
    {
        Ok(Ok(0)) | Err(_) => panic!("timed out waiting for JSON-RPC response"),
        Ok(Ok(_)) => serde_json::from_str(line.trim()).expect("response should be valid JSON"),
        Ok(Err(e)) => panic!("failed to read JSON-RPC response line: {e}"),
    }
}

// ============================================================================
// Test 6: Bad request params returns -32602 and connection stays alive (from Ben's branch)
// ============================================================================

#[tokio::test(flavor = "current_thread")]
async fn test_bad_request_params_return_invalid_params_and_connection_stays_alive() {
    use tokio::io::{AsyncWriteExt, BufReader};
    use tokio::task::LocalSet;

    let local = LocalSet::new();

    local
        .run_until(async {
            let (mut client_writer, server_reader) = tokio::io::duplex(2048);
            let (server_writer, client_reader) = tokio::io::duplex(2048);

            let server_reader = server_reader.compat();
            let server_writer = server_writer.compat_write();

            let server_transport =
                agent_client_protocol::ByteStreams::new(server_writer, server_reader);
            let server = UntypedRole.builder().on_receive_request(
                async |request: SimpleRequest,
                       responder: Responder<SimpleResponse>,
                       _connection: ConnectionTo<UntypedRole>| {
                    responder.respond(SimpleResponse {
                        result: format!("echo: {}", request.message),
                    })
                },
                agent_client_protocol::on_receive_request!(),
            );

            tokio::task::spawn_local(async move {
                if let Err(err) = server.connect_to(server_transport).await {
                    panic!("server should stay alive: {err:?}");
                }
            });

            let mut client_reader = BufReader::new(client_reader);

            client_writer
                .write_all(
                    br#"{"jsonrpc":"2.0","id":3,"method":"simple_method","params":{"content":"hello"}}
"#,
                )
                .await
                .unwrap();
            client_writer.flush().await.unwrap();

            let invalid_response = read_jsonrpc_response_line(&mut client_reader).await;
            expect![[r#"
                {
                  "jsonrpc": "2.0",
                  "id": 3,
                  "error": {
                    "code": -32602,
                    "message": "Invalid params",
                    "data": {
                      "error": "missing field `message`",
                      "json": {
                        "content": "hello"
                      },
                      "phase": "deserialization"
                    }
                  }
                }"#]]
            .assert_eq(&serde_json::to_string_pretty(&invalid_response).unwrap());

            client_writer
                .write_all(
                    br#"{"jsonrpc":"2.0","id":4,"method":"simple_method","params":{"message":"hello"}}
"#,
                )
                .await
                .unwrap();
            client_writer.flush().await.unwrap();

            let ok_response = read_jsonrpc_response_line(&mut client_reader).await;
            expect![[r#"
                {
                  "jsonrpc": "2.0",
                  "id": 4,
                  "result": {
                    "result": "echo: hello"
                  }
                }"#]]
            .assert_eq(&serde_json::to_string_pretty(&ok_response).unwrap());
        })
        .await;
}

// ============================================================================
// Test 7: Notification errors
// ============================================================================

#[tokio::test(flavor = "current_thread")]
async fn test_notification_errors_are_ignored_and_connection_stays_alive() {
    use tokio::io::{AsyncWriteExt, BufReader};
    use tokio::task::LocalSet;

    let local = LocalSet::new();

    local
        .run_until(async {
            let (mut client_writer, server_reader) = tokio::io::duplex(2048);
            let (server_writer, client_reader) = tokio::io::duplex(2048);

            let server_reader = server_reader.compat();
            let server_writer = server_writer.compat_write();

            let server_transport =
                agent_client_protocol::ByteStreams::new(server_writer, server_reader);
            let server = UntypedRole
                .builder()
                .on_receive_notification(
                    async |notif: SimpleNotification,
                           _connection: ConnectionTo<UntypedRole>| {
                        assert_eq!(notif.message, "handler error");
                        Err::<(), _>(agent_client_protocol::Error::internal_error())
                    },
                    agent_client_protocol::on_receive_notification!(),
                )
                .on_receive_request(
                    async |request: SimpleRequest,
                           responder: Responder<SimpleResponse>,
                           _connection: ConnectionTo<UntypedRole>| {
                        responder.respond(SimpleResponse {
                            result: format!("echo: {}", request.message),
                        })
                    },
                    agent_client_protocol::on_receive_request!(),
                );

            tokio::task::spawn_local(async move {
                if let Err(err) = server.connect_to(server_transport).await {
                    panic!("server should stay alive: {err:?}");
                }
            });

            let mut client_reader = BufReader::new(client_reader);

            // Send notifications that fail during parsing and inside the handler,
            // followed by a request that acts as an ordering barrier.
            client_writer
                .write_all(
                    br#"{"jsonrpc":"2.0","method":"simple_notification","params":{"wrong_field":"hello"}}
"#,
                )
                .await
                .unwrap();
            client_writer
                .write_all(
                    br#"{"jsonrpc":"2.0","method":"simple_notification","params":{"message":"handler error"}}
"#,
                )
                .await
                .unwrap();
            client_writer.flush().await.unwrap();

            client_writer
                .write_all(
                    br#"{"jsonrpc":"2.0","id":10,"method":"simple_method","params":{"message":"after bad notification"}}
"#,
                )
                .await
                .unwrap();
            client_writer.flush().await.unwrap();

            // The first line must be the request response. Any reply to either
            // failing notification would arrive before it.
            let ok_response = read_jsonrpc_response_line(&mut client_reader).await;
            expect![[r#"
                {
                  "jsonrpc": "2.0",
                  "id": 10,
                  "result": {
                    "result": "echo: after bad notification"
                  }
                }"#]]
            .assert_eq(&serde_json::to_string_pretty(&ok_response).unwrap());
        })
        .await;
}

#[tokio::test(flavor = "current_thread")]
async fn dispatch_handler_can_ignore_notification_and_connection_stays_alive() {
    use tokio::io::{AsyncWriteExt, BufReader};
    use tokio::task::LocalSet;

    LocalSet::new()
        .run_until(async {
            let (mut client_writer, server_reader) = tokio::io::duplex(2048);
            let (server_writer, client_reader) = tokio::io::duplex(2048);
            let server_transport = agent_client_protocol::ByteStreams::new(
                server_writer.compat_write(),
                server_reader.compat(),
            );
            let server = UntypedRole
                .builder()
                .on_receive_request(
                    async |request: SimpleRequest,
                           responder: Responder<SimpleResponse>,
                           _connection: ConnectionTo<UntypedRole>| {
                        responder.respond(SimpleResponse {
                            result: format!("echo: {}", request.message),
                        })
                    },
                    agent_client_protocol::on_receive_request!(),
                )
                .on_receive_dispatch(
                    async |message: Dispatch, _connection: ConnectionTo<UntypedRole>| match message
                    {
                        Dispatch::Request(_, responder) => responder
                            .respond_with_error(agent_client_protocol::Error::method_not_found()),
                        Dispatch::Notification(_) => Ok(()),
                        Dispatch::Response(result, router) => router.route_with_result(result),
                    },
                    agent_client_protocol::on_receive_dispatch!(),
                );

            tokio::task::spawn_local(async move {
                if let Err(error) = server.connect_to(server_transport).await {
                    panic!("server should stay alive: {error:?}");
                }
            });

            client_writer
                .write_all(
                    br#"{"jsonrpc":"2.0","method":"unknown/notification","params":{}}
{"jsonrpc":"2.0","id":11,"method":"simple_method","params":{"message":"after ignored notification"}}
"#,
                )
                .await
                .unwrap();
            client_writer.flush().await.unwrap();

            let mut client_reader = BufReader::new(client_reader);
            let response = read_jsonrpc_response_line(&mut client_reader).await;
            assert_eq!(response["id"], 11);
            assert_eq!(
                response["result"]["result"],
                "echo: after ignored notification"
            );
        })
        .await;
}