agent-client-protocol 5.0.0

//! Tests for the universal connection termination event.

use std::{future, io, time::Duration};

use agent_client_protocol::{
    Channel, ConnectionTermination, ConnectionTo, Error, Lines, RawJsonRpcMessage, TransportFrame,
    is_connection_terminated, is_incoming_transport_closed, role::UntypedRole,
};
use agent_client_protocol_test::MyRequest;
use futures::{StreamExt as _, channel::oneshot, stream};

const TIMEOUT: Duration = Duration::from_secs(2);

type Escaped = ConnectionTo<UntypedRole>;

async fn within<T>(future: impl Future<Output = T>) -> T {
    tokio::time::timeout(TIMEOUT, future)
        .await
        .expect("operation timed out")
}

#[tokio::test]
async fn clean_eof_terminates_with_incoming_closed() {
    let (transport, peer) = Channel::duplex();
    let (escaped_tx, escaped_rx) = oneshot::channel::<Escaped>();

    let driver = UntypedRole
        .builder()
        .with_spawned(async move |cx| {
            escaped_tx.send(cx).expect("observer is alive");
            future::pending().await
        })
        .connect_to(transport);
    let observer = async {
        let cx = escaped_rx.await.expect("connection escaped");
        assert_eq!(cx.termination(), None);
        drop(peer);
        cx.terminated().await
    };

    let (result, termination) = within(futures::future::join(driver, observer)).await;
    result.expect("clean EOF is not an error");
    assert_eq!(termination, ConnectionTermination::IncomingClosed);
    assert!(termination.is_clean());
}

#[tokio::test]
async fn foreground_success_terminates_with_local_shutdown() {
    let (transport, _peer) = Channel::duplex();

    let cx = within(
        UntypedRole
            .builder()
            .connect_with(transport, async |cx| Ok(cx)),
    )
    .await
    .expect("foreground succeeds");

    assert_eq!(cx.termination(), Some(ConnectionTermination::LocalShutdown));
    assert_eq!(cx.terminated().await, ConnectionTermination::LocalShutdown);
}

#[tokio::test]
async fn foreground_error_is_reported() {
    let (transport, _peer) = Channel::duplex();
    let (escaped_tx, escaped_rx) = oneshot::channel::<Escaped>();

    let error = within(UntypedRole.builder().connect_with(transport, async |cx| {
        escaped_tx.send(cx).expect("observer is alive");
        Err::<(), _>(Error::invalid_params().data("foreground failed"))
    }))
    .await
    .expect_err("foreground error is returned");

    let cx = escaped_rx.await.expect("connection escaped");
    assert_eq!(
        cx.termination(),
        Some(ConnectionTermination::ForegroundError(error))
    );
}

#[tokio::test]
async fn runner_error_is_a_handler_error() {
    let (transport, _peer) = Channel::duplex();
    let (escaped_tx, escaped_rx) = oneshot::channel::<Escaped>();

    let error = within(
        UntypedRole
            .builder()
            .with_spawned(async move |cx| {
                escaped_tx.send(cx).expect("observer is alive");
                Err(Error::internal_error().data("runner failed"))
            })
            .connect_with(transport, async |_cx| {
                future::pending::<Result<(), Error>>().await
            }),
    )
    .await
    .expect_err("runner error is returned");

    let cx = escaped_rx.await.expect("connection escaped");
    assert_eq!(
        cx.terminated().await,
        ConnectionTermination::HandlerError(error)
    );
}

#[tokio::test]
async fn close_callback_error_is_a_handler_error() {
    let (transport, peer) = Channel::duplex();
    drop(peer);
    let (escaped_tx, escaped_rx) = oneshot::channel::<Escaped>();

    let error = within(
        UntypedRole
            .builder()
            .on_close(async move |cx| {
                escaped_tx.send(cx).expect("observer is alive");
                Err(Error::internal_error().data("close failed"))
            })
            .connect_with(transport, async |_cx| {
                future::pending::<Result<(), Error>>().await
            }),
    )
    .await
    .expect_err("close callback error is returned");

    let cx = escaped_rx.await.expect("connection escaped");
    assert_eq!(
        cx.termination(),
        Some(ConnectionTermination::HandlerError(error))
    );
}

#[tokio::test]
async fn transport_read_error_is_a_transport_error() {
    let outgoing = futures::sink::unfold((), |(), _line: String| async { Ok::<_, io::Error>(()) });
    let (escaped_tx, escaped_rx) = oneshot::channel::<Escaped>();
    let (fail_tx, fail_rx) = oneshot::channel::<()>();
    // Fail the read only after the foreground has escaped the connection.
    let incoming = stream::once(async {
        let _ = fail_rx.await;
        Err(io::Error::other("read failed"))
    })
    .chain(stream::pending());

    let error = within(UntypedRole.builder().connect_with(
        Lines::new(outgoing, incoming),
        async |cx| {
            escaped_tx.send(cx).expect("observer is alive");
            let _ = fail_tx.send(());
            future::pending::<Result<(), Error>>().await
        },
    ))
    .await
    .expect_err("read error is returned");

    let cx = escaped_rx.await.expect("connection escaped");
    assert_eq!(
        cx.termination(),
        Some(ConnectionTermination::TransportError(error))
    );
}

#[tokio::test]
async fn pending_requests_fail_with_termination_after_local_shutdown() {
    let (transport, mut peer) = Channel::duplex();
    let (escaped_tx, escaped_rx) = oneshot::channel::<Escaped>();

    // An escaped handle keeps a request alive past the foreground.
    let driver = UntypedRole.builder().connect_with(transport, async |cx| {
        escaped_tx.send(cx.clone()).expect("observer is alive");
        // Wait until the request reaches the peer, then finish.
        assert!(matches!(
            peer.rx.next().await,
            Some(TransportFrame::Single(RawJsonRpcMessage::Request(_)))
        ));
        Ok(())
    });
    let requester = async {
        let cx = escaped_rx.await.expect("connection escaped");
        let response = cx.send_request(MyRequest {}).block_task().await;
        (cx, response)
    };

    let (result, (cx, response)) = within(futures::future::join(driver, requester)).await;
    result.expect("foreground succeeds");

    let error = response.expect_err("request cannot be answered after termination");
    assert!(is_connection_terminated(&error), "{error:?}");
    assert!(!is_incoming_transport_closed(&error));
    let data = error.data.as_ref().expect("termination error data");
    assert_eq!(data["method"], "myRequest");
    assert_eq!(data["termination"], "local_shutdown");

    // Requests sent afterwards fail the same way.
    let late = cx
        .send_request(MyRequest {})
        .block_task()
        .await
        .expect_err("late request fails");
    assert!(is_connection_terminated(&late), "{late:?}");
}

#[tokio::test]
async fn pending_requests_keep_the_eof_error_after_clean_eof() {
    let (transport, mut peer) = Channel::duplex();

    let cx = within(UntypedRole.builder().connect_with(transport, async |cx| {
        let request = cx.send_request(MyRequest {});
        assert!(matches!(
            peer.rx.next().await,
            Some(TransportFrame::Single(RawJsonRpcMessage::Request(_)))
        ));
        drop(peer);
        let error = request
            .block_task()
            .await
            .expect_err("EOF fails pending requests");
        assert!(is_incoming_transport_closed(&error), "{error:?}");
        assert_eq!(cx.termination(), None, "foreground is still running");
        Ok(cx)
    }))
    .await
    .expect("foreground succeeds");

    assert_eq!(
        cx.termination(),
        Some(ConnectionTermination::IncomingClosed)
    );
}

#[tokio::test]
async fn dropping_the_driver_publishes_dropped() {
    let (transport, _peer) = Channel::duplex();
    let (escaped_tx, escaped_rx) = oneshot::channel::<Escaped>();

    let mut driver = Box::pin(UntypedRole.builder().connect_with(transport, async |cx| {
        escaped_tx.send(cx).expect("observer is alive");
        future::pending::<Result<(), Error>>().await
    }));
    // Poll once so the foreground starts and escapes the connection.
    assert!(futures::poll!(driver.as_mut()).is_pending());
    let cx = escaped_rx.await.expect("connection escaped");

    // A request awaiting a response when the driver is dropped.
    let request = cx.send_request(MyRequest {});
    assert!(futures::poll!(driver.as_mut()).is_pending());
    drop(driver);

    assert_eq!(cx.termination(), Some(ConnectionTermination::Dropped));
    let error = within(request.block_task())
        .await
        .expect_err("dropped driver fails pending requests");
    assert!(is_connection_terminated(&error), "{error:?}");
    assert_eq!(error.data.as_ref().unwrap()["termination"], "dropped");
}

/// A peer's error response carrying the SDK's termination reasons in its
/// data must not make the local connection look terminated or closed.
#[tokio::test]
async fn peer_errors_with_termination_reasons_do_not_match() {
    let (transport, mut peer) = Channel::duplex();

    let cx = within(UntypedRole.builder().connect_with(transport, async |cx| {
        for (message, reason, termination) in [
            (
                "Connection terminated",
                agent_client_protocol::CONNECTION_TERMINATED_REASON,
                Some("local_shutdown"),
            ),
            (
                "Incoming transport closed",
                agent_client_protocol::INCOMING_TRANSPORT_CLOSED_REASON,
                None,
            ),
        ] {
            let request = cx.send_request(MyRequest {});
            let Some(TransportFrame::Single(RawJsonRpcMessage::Request(sent))) =
                peer.rx.next().await
            else {
                panic!("expected the request");
            };
            // The exact error the SDK produces locally, sent from the wire.
            let mut data = serde_json::json!({ "reason": reason, "method": "myRequest" });
            if let Some(termination) = termination {
                data["termination"] = termination.into();
            }
            let mut forged = Error::internal_error().data(data);
            forged.message = message.to_string();
            peer.tx
                .unbounded_send(TransportFrame::Single(RawJsonRpcMessage::response(
                    sent.id,
                    Err(forged.clone()),
                )))
                .expect("connection is alive");

            let error = request
                .block_task()
                .await
                .expect_err("the peer answered with an error");
            assert_eq!(error, forged, "the peer's error is delivered unchanged");
            assert!(!is_connection_terminated(&error), "{error:?}");
            assert!(!is_incoming_transport_closed(&error), "{error:?}");
            // Building the same error locally does not match either.
            assert!(!is_connection_terminated(&forged));
            assert!(!is_incoming_transport_closed(&forged));
        }
        assert_eq!(cx.termination(), None, "the connection is still running");
        Ok(cx)
    }))
    .await
    .expect("foreground succeeds");

    assert_eq!(cx.termination(), Some(ConnectionTermination::LocalShutdown));
}