agent-client-protocol 5.0.0

use std::{path::Path, time::Duration};

use agent_client_protocol::V1;
use agent_client_protocol::v1::ConnectionTo;
use agent_client_protocol::v1::{SessionMessage, SessionSetup};
use agent_client_protocol::{
    Agent, Channel, Client, Conductor, RawJsonRpcMessage, Responder, TransportFrame,
    v1::schema::{
        ContentBlock, ContentChunk, NewSessionRequest, NewSessionResponse, PromptRequest,
        PromptResponse, SessionConfigOption, SessionConfigOptionCategory,
        SessionConfigSelectOption, SessionId, SessionNotification, SessionUpdate, StopReason,
        TextContent,
    },
};
use futures::{StreamExt as _, channel::oneshot};

#[cfg(feature = "unstable_protocol_v2")]
use futures::channel::mpsc;

#[cfg(feature = "unstable_protocol_v2")]
use agent_client_protocol::{
    JsonRpcMessage, JsonRpcResponse, Proxy,
    schema::{ProtocolVersion, SuccessorMessage},
    v2::schema as v2,
};

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

/// The JSON value `parse_message` takes for a raw message's `params`.
#[cfg(feature = "unstable_protocol_v2")]
fn params_value(params: Option<&agent_client_protocol::RawJsonRpcParams>) -> serde_json::Value {
    params.cloned().map_or(
        serde_json::Value::Null,
        agent_client_protocol::RawJsonRpcParams::into_value,
    )
}

#[cfg(feature = "unstable_protocol_v2")]
async fn initialize_raw_v2_proxy(
    peer: &mut Channel,
    client_name: &str,
) -> Result<(), agent_client_protocol::Error> {
    let initialize_id = agent_client_protocol::RawJsonRpcId::Number(1);
    let initialize = v2::InitializeProxyRequest::new(v2::InitializeRequest::new(
        ProtocolVersion::V2,
        v2::Implementation::new(client_name, env!("CARGO_PKG_VERSION")),
    ));
    peer.tx
        .unbounded_send(TransportFrame::Single(RawJsonRpcMessage::request(
            "proxy/initialize".to_owned(),
            serde_json::to_value(initialize).expect("initialize request should serialize"),
            initialize_id.clone(),
        )?))
        .expect("proxy should accept initialization");

    let Some(TransportFrame::Single(RawJsonRpcMessage::Response(
        agent_client_protocol::RawJsonRpcResponse::Result { id, result },
    ))) = peer.rx.next().await
    else {
        panic!("expected the proxy initialize response");
    };
    assert_eq!(id, initialize_id);
    let response = v2::InitializeResponse::from_json("proxy/initialize", result)?;
    assert_eq!(response.protocol_version, ProtocolVersion::V2);
    Ok(())
}

// Compile-time regressions for the callback future bounds on the two non-blocking
// session helpers. The callbacks themselves are `'static`, but their future
// types deliberately carry an arbitrary shorter lifetime.
#[allow(dead_code)]
mod callback_future_lifetimes {
    use std::{
        future::Future,
        marker::PhantomData,
        pin::Pin,
        task::{Context, Poll},
    };

    use super::*;

    struct LifetimeTaggedFuture<'a>(PhantomData<&'a ()>);

    impl Future for LifetimeTaggedFuture<'_> {
        type Output = Result<(), agent_client_protocol::Error>;

        fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
            Poll::Ready(Ok(()))
        }
    }

    fn session_callback<'a>()
    -> impl FnOnce(SessionSetup<'static, Agent, NewSessionResponse>) -> LifetimeTaggedFuture<'a>
    + Send
    + 'static {
        |_setup| LifetimeTaggedFuture(PhantomData)
    }

    fn proxy_session_callback<'a>()
    -> impl FnOnce(NewSessionResponse) -> LifetimeTaggedFuture<'a> + Send + 'static {
        |_response| LifetimeTaggedFuture(PhantomData)
    }

    #[cfg(feature = "unstable_protocol_v2")]
    fn v2_proxy_session_callback<'a>() -> impl FnOnce(
        agent_client_protocol::v2::SessionSetup<Conductor, v2::NewSessionResponse>,
    ) -> LifetimeTaggedFuture<'a>
    + Send
    + 'static {
        |_setup| LifetimeTaggedFuture(PhantomData)
    }

    #[cfg(feature = "unstable_protocol_v2")]
    fn v2_proxy_resume_callback<'a>() -> impl FnOnce(
        agent_client_protocol::v2::SessionSetup<Conductor, v2::ResumeSessionResponse>,
    ) -> LifetimeTaggedFuture<'a>
    + Send
    + 'static {
        |_setup| LifetimeTaggedFuture(PhantomData)
    }

    #[cfg(all(feature = "unstable_protocol_v2", feature = "unstable_session_fork"))]
    fn v2_proxy_fork_callback<'a>() -> impl FnOnce(
        agent_client_protocol::v2::SessionSetup<Conductor, v2::ForkSessionResponse>,
    ) -> LifetimeTaggedFuture<'a>
    + Send
    + 'static {
        |_setup| LifetimeTaggedFuture(PhantomData)
    }

    fn on_session_start_accepts_non_static_callback_future<'a>(
        connection: &ConnectionTo<Agent>,
        _scope: &'a str,
    ) -> Result<(), agent_client_protocol::Error> {
        connection
            .build_session_cwd()?
            .on_session_start(session_callback::<'a>())
    }

    fn on_proxy_session_start_accepts_non_static_callback_future<'a>(
        connection: &ConnectionTo<Conductor>,
        request: NewSessionRequest,
        responder: Responder<NewSessionResponse, V1>,
        _scope: &'a str,
    ) -> Result<(), agent_client_protocol::Error> {
        connection
            .build_session_from(request)
            .on_proxy_session_start(responder, proxy_session_callback::<'a>())
    }

    #[cfg(feature = "unstable_protocol_v2")]
    fn v2_on_proxy_session_start_accepts_non_static_callback_future<'a>(
        connection: &agent_client_protocol::v2::ConnectionTo<Conductor>,
        request: v2::NewSessionRequest,
        responder: Responder<v2::NewSessionResponse, agent_client_protocol::V2>,
        _scope: &'a str,
    ) -> Result<(), agent_client_protocol::Error> {
        connection
            .build_session_from(request)
            .on_proxy_session_start(responder, v2_proxy_session_callback::<'a>())
    }

    #[cfg(feature = "unstable_protocol_v2")]
    fn v2_on_proxy_resume_accepts_non_static_callback_future<'a>(
        connection: &agent_client_protocol::v2::ConnectionTo<Conductor>,
        request: v2::ResumeSessionRequest,
        responder: Responder<v2::ResumeSessionResponse, agent_client_protocol::V2>,
        _scope: &'a str,
    ) -> Result<(), agent_client_protocol::Error> {
        connection
            .resume_session_from(request)
            .on_proxy_session_start(responder, v2_proxy_resume_callback::<'a>())
    }

    #[cfg(all(feature = "unstable_protocol_v2", feature = "unstable_session_fork"))]
    fn v2_on_proxy_fork_accepts_non_static_callback_future<'a>(
        connection: &agent_client_protocol::v2::ConnectionTo<Conductor>,
        request: v2::ForkSessionRequest,
        responder: Responder<v2::ForkSessionResponse, agent_client_protocol::V2>,
        _scope: &'a str,
    ) -> Result<(), agent_client_protocol::Error> {
        connection
            .fork_session_from(request)
            .on_proxy_session_start(responder, v2_proxy_fork_callback::<'a>())
    }
}

#[tokio::test(flavor = "current_thread")]
async fn active_session_preserves_config_options_from_new_session_response() {
    let config_options = vec![
        SessionConfigOption::select(
            "model",
            "Model",
            "sonnet",
            vec![
                SessionConfigSelectOption::new("sonnet", "Sonnet"),
                SessionConfigSelectOption::new("opus", "Opus"),
            ],
        )
        .category(SessionConfigOptionCategory::Model),
    ];
    let expected_response =
        NewSessionResponse::new("config-options-session").config_options(config_options.clone());
    let agent_response = expected_response.clone();

    let agent = Agent.v1().on_receive_request(
        async move |_request: NewSessionRequest,
                    responder: Responder<NewSessionResponse, V1>,
                    _connection: ConnectionTo<Client>| {
            responder.respond(agent_response.clone())
        },
        agent_client_protocol::on_receive_request!(),
    );

    let client = Client.v1().connect_with(agent, async move |connection| {
        let setup = connection
            .build_session_cwd()?
            .block_task()
            .start_session()
            .await?;

        assert_eq!(
            setup.response.config_options.as_deref(),
            Some(config_options.as_slice())
        );
        assert_eq!(setup.response, expected_response);
        Ok(())
    });

    tokio::time::timeout(TIMEOUT, client)
        .await
        .expect("session setup timed out")
        .expect("session connection failed");
}

#[tokio::test(flavor = "current_thread")]
async fn on_session_start_callback_can_consume_later_session_messages() {
    let session_id = SessionId::new("ordered-session");
    let new_session_id = session_id.clone();
    let prompt_session_id = session_id.clone();

    let agent = Agent
        .v1()
        .on_receive_request(
            async move |_request: NewSessionRequest,
                        responder: Responder<NewSessionResponse, V1>,
                        _connection: ConnectionTo<Client>| {
                responder.respond(NewSessionResponse::new(new_session_id.clone()))
            },
            agent_client_protocol::on_receive_request!(),
        )
        .on_receive_request(
            async move |request: PromptRequest,
                        responder: Responder<PromptResponse, V1>,
                        connection: ConnectionTo<Client>| {
                assert_eq!(request.session_id, prompt_session_id);
                connection.send_notification(SessionNotification::new(
                    request.session_id,
                    SessionUpdate::AgentMessageChunk(ContentChunk::new(ContentBlock::Text(
                        TextContent::new("ordered response"),
                    ))),
                ))?;
                responder.respond(PromptResponse::new(StopReason::EndTurn))
            },
            agent_client_protocol::on_receive_request!(),
        );

    let (result_tx, result_rx) = oneshot::channel();
    let client = Client.v1().connect_with(agent, async move |connection| {
        connection.build_session_cwd()?.on_session_start(
            async move |SessionSetup { mut session, .. }| {
                session.send_prompt("test ordering")?;
                let text = session.read_to_string().await?;
                result_tx
                    .send(text)
                    .map_err(|_| agent_client_protocol::Error::internal_error())
            },
        )?;

        let text = result_rx
            .await
            .map_err(|_| agent_client_protocol::Error::internal_error())?;
        assert_eq!(text, "ordered response");
        Ok(())
    });

    tokio::time::timeout(TIMEOUT, client)
        .await
        .expect("session callback deadlocked the incoming dispatch loop")
        .expect("session connection failed");
}

#[tokio::test(flavor = "current_thread")]
async fn prompt_error_reaches_the_session_stream_and_keeps_the_connection() {
    let session_id = SessionId::new("prompt-error-session");
    let new_session_id = session_id.clone();
    let other_session_id = SessionId::new("other-session");
    let other_new_session_id = other_session_id.clone();
    let meta = agent_client_protocol::v1::schema::Meta::from_iter([(
        "turn".to_owned(),
        serde_json::json!({ "kept": true }),
    )]);
    let expected_response = PromptResponse::new(StopReason::EndTurn).meta(meta.clone());

    let agent = Agent
        .v1()
        .on_receive_request(
            {
                let created = std::sync::Mutex::new(0_u8);
                async move |_request: NewSessionRequest,
                            responder: Responder<NewSessionResponse, V1>,
                            _connection: ConnectionTo<Client>| {
                    let id = {
                        let mut created = created.lock().expect("lock poisoned");
                        *created += 1;
                        if *created == 1 {
                            new_session_id.clone()
                        } else {
                            other_new_session_id.clone()
                        }
                    };
                    responder.respond(NewSessionResponse::new(id))
                }
            },
            agent_client_protocol::on_receive_request!(),
        )
        .on_receive_request(
            async move |request: PromptRequest,
                        responder: Responder<PromptResponse, V1>,
                        connection: ConnectionTo<Client>| {
                let ContentBlock::Text(text) = &request.prompt[0] else {
                    panic!("expected a text prompt");
                };
                if text.text == "fail" {
                    return responder.respond_with_error(
                        agent_client_protocol::Error::invalid_params().data("rejected prompt"),
                    );
                }
                connection.send_notification(SessionNotification::new(
                    request.session_id,
                    SessionUpdate::AgentMessageChunk(ContentChunk::new(ContentBlock::Text(
                        TextContent::new("ok"),
                    ))),
                ))?;
                responder.respond(PromptResponse::new(StopReason::EndTurn).meta(meta.clone()))
            },
            agent_client_protocol::on_receive_request!(),
        );

    let client = Client.v1().connect_with(agent, async move |connection| {
        let mut session = connection
            .build_session_cwd()?
            .block_task()
            .start_session()
            .await?
            .session;
        let mut other = connection
            .build_session_cwd()?
            .block_task()
            .start_session()
            .await?
            .session;
        assert_eq!(session.session_id(), &session_id);
        assert_eq!(other.session_id(), &other_session_id);

        session.send_prompt("fail")?;
        let error = session
            .read_to_string()
            .await
            .expect_err("the prompt error should be returned");
        assert_eq!(error.code, agent_client_protocol::ErrorCode::INVALID_PARAMS);

        // The error did not close the connection or the other session.
        other.send_prompt("hello")?;
        assert_eq!(other.read_to_string().await?, "ok");

        // The same session keeps working, and the complete response arrives.
        session.send_prompt_blocks(vec![ContentBlock::Text(TextContent::new("hello"))])?;
        let SessionMessage::Message(_) = session.read_update().await? else {
            panic!("expected the agent message chunk first");
        };
        let SessionMessage::PromptResponse(response) = session.read_update().await? else {
            panic!("expected the prompt response");
        };
        assert_eq!(response?, expected_response);
        Ok(())
    });

    tokio::time::timeout(TIMEOUT, client)
        .await
        .expect("prompt error test timed out")
        .expect("session connection failed");
}

#[tokio::test(flavor = "current_thread")]
async fn on_session_start_installs_routing_before_later_message() {
    let session_id = SessionId::new("queued-session");
    let response_session_id = session_id.clone();
    let notification_session_id = session_id.clone();
    let (transport, mut peer) = Channel::duplex();
    let (result_tx, result_rx) = oneshot::channel();

    let client = Client
        .v1()
        .connect_with(transport, async move |connection| {
            connection.build_session_cwd()?.on_session_start(
                async move |SessionSetup { mut session, .. }| {
                    let update = session.read_update().await?;
                    assert!(matches!(update, SessionMessage::Message(_)));
                    result_tx
                        .send(())
                        .map_err(|()| agent_client_protocol::Error::internal_error())
                },
            )?;

            result_rx
                .await
                .map_err(|_| agent_client_protocol::Error::internal_error())?;
            Ok(())
        });

    let peer = async move {
        let Some(TransportFrame::Single(RawJsonRpcMessage::Request(request))) =
            peer.rx.next().await
        else {
            panic!("expected a session/new request");
        };
        assert_eq!(request.method.as_ref(), "session/new");

        let response = RawJsonRpcMessage::response(
            request.id,
            Ok(
                serde_json::to_value(NewSessionResponse::new(response_session_id))
                    .expect("session response should serialize"),
            ),
        );
        let notification = RawJsonRpcMessage::notification(
            "session/update".into(),
            serde_json::to_value(SessionNotification::new(
                notification_session_id,
                SessionUpdate::AgentMessageChunk(ContentChunk::new(ContentBlock::Text(
                    TextContent::new("queued"),
                ))),
            ))
            .expect("session notification should serialize"),
        )
        .expect("session notification should form valid JSON-RPC parameters");
        for message in [response, notification] {
            peer.tx
                .unbounded_send(TransportFrame::Single(message))
                .expect("client should accept the queued messages");
        }

        while peer.rx.next().await.is_some() {}
        Ok::<(), agent_client_protocol::Error>(())
    };

    tokio::time::timeout(TIMEOUT, async { futures::try_join!(client, peer) })
        .await
        .expect("queued session update was not routed")
        .expect("session connection failed");
}

#[tokio::test(flavor = "current_thread")]
async fn start_session_installs_routing_before_later_message() {
    let session_id = SessionId::new("queued-blocking-session");
    let response_session_id = session_id.clone();
    let notification_session_id = session_id.clone();
    let (transport, mut peer) = Channel::duplex();

    let client = Client
        .v1()
        .connect_with(transport, async move |connection| {
            let mut session = connection
                .build_session_cwd()?
                .block_task()
                .start_session()
                .await?
                .session;
            let update = session.read_update().await?;
            assert!(matches!(update, SessionMessage::Message(_)));
            Ok(())
        });

    let peer = async move {
        let Some(TransportFrame::Single(RawJsonRpcMessage::Request(request))) =
            peer.rx.next().await
        else {
            panic!("expected a session/new request");
        };
        assert_eq!(request.method.as_ref(), "session/new");

        let response = RawJsonRpcMessage::response(
            request.id,
            Ok(
                serde_json::to_value(NewSessionResponse::new(response_session_id))
                    .expect("session response should serialize"),
            ),
        );
        let notification = RawJsonRpcMessage::notification(
            "session/update".into(),
            serde_json::to_value(SessionNotification::new(
                notification_session_id,
                SessionUpdate::AgentMessageChunk(ContentChunk::new(ContentBlock::Text(
                    TextContent::new("queued"),
                ))),
            ))
            .expect("session notification should serialize"),
        )
        .expect("session notification should form valid JSON-RPC parameters");
        for message in [response, notification] {
            peer.tx
                .unbounded_send(TransportFrame::Single(message))
                .expect("client should accept the queued messages");
        }

        while peer.rx.next().await.is_some() {}
        Ok::<(), agent_client_protocol::Error>(())
    };

    tokio::time::timeout(TIMEOUT, async { futures::try_join!(client, peer) })
        .await
        .expect("queued session update was not routed")
        .expect("session connection failed");
}

#[tokio::test(flavor = "current_thread")]
async fn start_session_returns_session_new_error_without_closing_connection() {
    let agent = Agent.v1().on_receive_request(
        async |request: NewSessionRequest,
               responder: Responder<NewSessionResponse, V1>,
               _connection: ConnectionTo<Client>| {
            if request.cwd == Path::new("/rejected") {
                responder.respond_with_error(
                    agent_client_protocol::Error::invalid_params().data("session rejected"),
                )
            } else {
                responder.respond(NewSessionResponse::new("accepted-session"))
            }
        },
        agent_client_protocol::on_receive_request!(),
    );

    let client = Client.v1().connect_with(agent, async move |connection| {
        let error = connection
            .build_session("/rejected")
            .block_task()
            .start_session()
            .await
            .expect_err("the session/new error must reach the caller");
        assert_eq!(
            error,
            agent_client_protocol::Error::invalid_params().data("session rejected"),
            "the helper must preserve the peer's JSON-RPC error"
        );

        let session = connection
            .build_session("/accepted")
            .block_task()
            .start_session()
            .await?
            .session;
        assert_eq!(session.session_id(), &SessionId::new("accepted-session"));
        Ok(())
    });

    tokio::time::timeout(TIMEOUT, client)
        .await
        .expect("session/new error test timed out")
        .expect("session/new error closed the connection");
}

#[cfg(feature = "unstable_protocol_v2")]
#[tokio::test(flavor = "current_thread")]
async fn v2_proxy_session_start_installs_routing_before_later_message() {
    let session_id = v2::SessionId::new("queued-v2-session");
    let setup_response = v2::NewSessionResponse::new(session_id.clone()).config_options(vec![
        v2::SessionConfigOption::boolean("thinking", "Thinking", true),
    ]);
    let callback_response = setup_response.clone();
    let response_session_id = session_id.clone();
    let notification_session_id = session_id.clone();
    let (transport, mut peer) = Channel::duplex();
    let (callback_tx, mut callback_rx) = mpsc::unbounded();
    let (peer_done_tx, peer_done_rx) = oneshot::channel();

    let proxy = Proxy
        .v2()
        .on_receive_request_from(
            Client,
            async |request: v2::InitializeProxyRequest, responder, _connection| {
                responder.respond(v2::InitializeResponse::new(
                    request.initialize.protocol_version,
                    v2::Implementation::new("queued-proxy", env!("CARGO_PKG_VERSION")),
                ).into())
            },
            agent_client_protocol::on_receive_request!(),
        )
        .on_receive_request_from(
            Client,
            async move |request: v2::NewSessionRequest,
                        responder,
                        connection: agent_client_protocol::v2::ConnectionTo<Conductor>| {
                let callback_response = callback_response.clone();
                let response_session_id = response_session_id.clone();
                let callback_tx = callback_tx.clone();
                connection
                    .build_session_from(request)
                    .on_proxy_session_start(responder, move |opened| async move {
                        assert_eq!(opened.session.session_id(), &response_session_id);
                        assert_eq!(&opened.response, &callback_response);
                        callback_tx.unbounded_send(()).map_err(|_| {
                            agent_client_protocol::Error::internal_error()
                                .data("v2 proxy callback receiver was dropped")
                        })
                    })
            },
            agent_client_protocol::on_receive_request!(),
        )
        .connect_with(transport, async move |_connection| {
            callback_rx.next().await.ok_or_else(|| {
                agent_client_protocol::Error::internal_error().data("v2 proxy callback did not run")
            })?;
            peer_done_rx.await.map_err(|_| {
                agent_client_protocol::Error::internal_error().data("raw peer stopped early")
            })
        });

    let peer = async move {
        initialize_raw_v2_proxy(&mut peer, "queued-client").await?;

        let upstream_id = agent_client_protocol::RawJsonRpcId::Number(2);
        peer.tx
            .unbounded_send(TransportFrame::Single(RawJsonRpcMessage::request(
                "session/new".to_owned(),
                serde_json::to_value(v2::NewSessionRequest::new("/queued-v2-session"))
                    .expect("session request should serialize"),
                upstream_id.clone(),
            )?))
            .expect("proxy should accept session/new");

        let Some(TransportFrame::Single(RawJsonRpcMessage::Request(forwarded))) =
            peer.rx.next().await
        else {
            panic!("expected a forwarded session/new request");
        };
        let successor = SuccessorMessage::<v2::NewSessionRequest>::parse_message(
            forwarded.method.as_ref(),
            &params_value(forwarded.params.as_ref()),
        )?;
        assert_eq!(
            successor.message.cwd,
            v2::AbsolutePath::new("/queued-v2-session")
        );

        let response = RawJsonRpcMessage::response(
            forwarded.id,
            Ok(serde_json::to_value(setup_response).expect("session response should serialize")),
        );
        let update = SuccessorMessage::new(v2::UpdateSessionNotification::new(
            notification_session_id,
            v2::SessionUpdate::StateUpdate(v2::StateUpdate::Running(v2::RunningStateUpdate::new())),
        ))
        .to_untyped_message()?;
        let (method, params) = update.into_parts();
        let notification = RawJsonRpcMessage::notification(method, params)?;
        for message in [response, notification] {
            peer.tx
                .unbounded_send(TransportFrame::Single(message))
                .expect("proxy should accept the queued messages");
        }

        let mut saw_response = false;
        let mut saw_update = false;
        for _ in 0..2 {
            let Some(TransportFrame::Single(message)) = peer.rx.next().await else {
                panic!("expected a forwarded session response and update");
            };
            match message {
                RawJsonRpcMessage::Response(
                    agent_client_protocol::RawJsonRpcResponse::Result { id, result },
                ) => {
                    assert_eq!(id, upstream_id);
                    let response = v2::NewSessionResponse::from_json("session/new", result)?;
                    assert_eq!(response.session_id, session_id);
                    saw_response = true;
                }
                RawJsonRpcMessage::Notification(notification) => {
                    let update = v2::UpdateSessionNotification::parse_message(
                        notification.method.as_ref(),
                        &params_value(notification.params.as_ref()),
                    )?;
                    assert_eq!(update.session_id, session_id);
                    assert!(matches!(
                        update.update,
                        v2::SessionUpdate::StateUpdate(v2::StateUpdate::Running(_))
                    ));
                    saw_update = true;
                }
                message => panic!("unexpected proxy output: {message:?}"),
            }
        }
        assert!(saw_response);
        assert!(saw_update);
        peer_done_tx
            .send(())
            .map_err(|()| agent_client_protocol::Error::internal_error())
    };

    tokio::time::timeout(TIMEOUT, async { futures::try_join!(proxy, peer) })
        .await
        .expect("queued v2 session update was not routed")
        .expect("v2 proxy session connection failed");
}

#[cfg(all(feature = "unstable_protocol_v2", feature = "unstable_session_fork"))]
#[tokio::test(flavor = "current_thread")]
async fn v2_proxy_fork_installs_response_id_routing_before_later_message() {
    let source_session_id = v2::SessionId::new("queued-v2-source");
    let forked_session_id = v2::SessionId::new("queued-v2-fork");
    let setup_response =
        v2::ForkSessionResponse::new(forked_session_id.clone()).config_options(vec![
            v2::SessionConfigOption::boolean("thinking", "Thinking", true),
        ]);
    let callback_response = setup_response.clone();
    let callback_session_id = forked_session_id.clone();
    let notification_session_id = forked_session_id.clone();
    let expected_source_session_id = source_session_id.clone();
    let (transport, mut peer) = Channel::duplex();
    let (callback_tx, mut callback_rx) = mpsc::unbounded();
    let (peer_done_tx, peer_done_rx) = oneshot::channel();

    let proxy = Proxy
        .v2()
        .on_receive_request_from(
            Client,
            async |request: v2::InitializeProxyRequest, responder, _connection| {
                responder.respond(v2::InitializeResponse::new(
                    request.initialize.protocol_version,
                    v2::Implementation::new("queued-fork-proxy", env!("CARGO_PKG_VERSION")),
                ).into())
            },
            agent_client_protocol::on_receive_request!(),
        )
        .on_receive_request_from(
            Client,
            async move |request: v2::ForkSessionRequest,
                        responder,
                        connection: agent_client_protocol::v2::ConnectionTo<Conductor>| {
                let callback_response = callback_response.clone();
                let callback_session_id = callback_session_id.clone();
                let callback_tx = callback_tx.clone();
                connection
                    .fork_session_from(request)
                    .on_proxy_session_start(responder, move |opened| async move {
                        assert_eq!(opened.session.session_id(), &callback_session_id);
                        assert_eq!(&opened.response, &callback_response);
                        callback_tx.unbounded_send(()).map_err(|_| {
                            agent_client_protocol::Error::internal_error()
                                .data("v2 fork callback receiver was dropped")
                        })
                    })
            },
            agent_client_protocol::on_receive_request!(),
        )
        .connect_with(transport, async move |_connection| {
            callback_rx.next().await.ok_or_else(|| {
                agent_client_protocol::Error::internal_error().data("v2 fork callback did not run")
            })?;
            peer_done_rx.await.map_err(|_| {
                agent_client_protocol::Error::internal_error().data("raw peer stopped early")
            })
        });

    let peer = async move {
        initialize_raw_v2_proxy(&mut peer, "queued-fork-client").await?;

        let upstream_id = agent_client_protocol::RawJsonRpcId::Number(2);
        peer.tx
            .unbounded_send(TransportFrame::Single(RawJsonRpcMessage::request(
                "session/fork".to_owned(),
                serde_json::to_value(v2::ForkSessionRequest::new(
                    source_session_id,
                    "/queued-v2-fork",
                ))
                .expect("fork request should serialize"),
                upstream_id.clone(),
            )?))
            .expect("proxy should accept session/fork");

        let Some(TransportFrame::Single(RawJsonRpcMessage::Request(forwarded))) =
            peer.rx.next().await
        else {
            panic!("expected a forwarded session/fork request");
        };
        let successor = SuccessorMessage::<v2::ForkSessionRequest>::parse_message(
            forwarded.method.as_ref(),
            &params_value(forwarded.params.as_ref()),
        )?;
        assert_eq!(successor.message.session_id, expected_source_session_id);
        assert_eq!(
            successor.message.cwd,
            v2::AbsolutePath::new("/queued-v2-fork")
        );

        let response = RawJsonRpcMessage::response(
            forwarded.id,
            Ok(serde_json::to_value(setup_response).expect("fork response should serialize")),
        );
        let update = SuccessorMessage::new(v2::UpdateSessionNotification::new(
            notification_session_id,
            v2::SessionUpdate::StateUpdate(v2::StateUpdate::Running(v2::RunningStateUpdate::new())),
        ))
        .to_untyped_message()?;
        let (method, params) = update.into_parts();
        let notification = RawJsonRpcMessage::notification(method, params)?;
        for message in [response, notification] {
            peer.tx
                .unbounded_send(TransportFrame::Single(message))
                .expect("proxy should accept the queued messages");
        }

        let mut saw_response = false;
        let mut saw_update = false;
        for _ in 0..2 {
            let Some(TransportFrame::Single(message)) = peer.rx.next().await else {
                panic!("expected a forwarded fork response and update");
            };
            match message {
                RawJsonRpcMessage::Response(
                    agent_client_protocol::RawJsonRpcResponse::Result { id, result },
                ) => {
                    assert_eq!(id, upstream_id);
                    let response = v2::ForkSessionResponse::from_json("session/fork", result)?;
                    assert_eq!(response.session_id, forked_session_id);
                    assert_eq!(response.config_options.len(), 1);
                    saw_response = true;
                }
                RawJsonRpcMessage::Notification(notification) => {
                    let update = v2::UpdateSessionNotification::parse_message(
                        notification.method.as_ref(),
                        &params_value(notification.params.as_ref()),
                    )?;
                    assert_eq!(update.session_id, forked_session_id);
                    assert!(matches!(
                        update.update,
                        v2::SessionUpdate::StateUpdate(v2::StateUpdate::Running(_))
                    ));
                    saw_update = true;
                }
                message => panic!("unexpected proxy output: {message:?}"),
            }
        }
        assert!(saw_response);
        assert!(saw_update);
        peer_done_tx
            .send(())
            .map_err(|()| agent_client_protocol::Error::internal_error())
    };

    tokio::time::timeout(TIMEOUT, async { futures::try_join!(proxy, peer) })
        .await
        .expect("queued v2 fork update was not routed")
        .expect("v2 proxy fork connection failed");
}

#[cfg(feature = "unstable_protocol_v2")]
#[tokio::test(flavor = "current_thread")]
async fn v2_proxy_resume_forwards_replay_before_its_response() {
    let session_id = v2::SessionId::new("queued-v2-resume");
    let setup_response =
        v2::ResumeSessionResponse::new().config_options(vec![v2::SessionConfigOption::boolean(
            "thinking", "Thinking", false,
        )]);
    let callback_response = setup_response.clone();
    let upstream_response = setup_response.clone();
    let callback_session_id = session_id.clone();
    let notification_session_id = session_id.clone();
    let (transport, mut peer) = Channel::duplex();
    let (callback_tx, mut callback_rx) = mpsc::unbounded();
    let (peer_done_tx, peer_done_rx) = oneshot::channel();

    let proxy = Proxy
        .v2()
        .on_receive_request_from(
            Client,
            async |request: v2::InitializeProxyRequest, responder, _connection| {
                responder.respond(v2::InitializeResponse::new(
                    request.initialize.protocol_version,
                    v2::Implementation::new("queued-resume-proxy", env!("CARGO_PKG_VERSION")),
                ).into())
            },
            agent_client_protocol::on_receive_request!(),
        )
        .on_receive_request_from(
            Client,
            async move |request: v2::ResumeSessionRequest,
                        responder,
                        connection: agent_client_protocol::v2::ConnectionTo<Conductor>| {
                let callback_response = callback_response.clone();
                let callback_session_id = callback_session_id.clone();
                let callback_tx = callback_tx.clone();
                connection
                    .resume_session_from(request)
                    .on_proxy_session_start(responder, move |opened| async move {
                        assert_eq!(opened.session.session_id(), &callback_session_id);
                        assert_eq!(&opened.response, &callback_response);
                        callback_tx.unbounded_send(()).map_err(|_| {
                            agent_client_protocol::Error::internal_error()
                                .data("v2 resume callback receiver was dropped")
                        })
                    })
            },
            agent_client_protocol::on_receive_request!(),
        )
        .connect_with(transport, async move |_connection| {
            callback_rx.next().await.ok_or_else(|| {
                agent_client_protocol::Error::internal_error()
                    .data("v2 resume callback did not run")
            })?;
            peer_done_rx.await.map_err(|_| {
                agent_client_protocol::Error::internal_error().data("raw peer stopped early")
            })
        });

    let peer = async move {
        initialize_raw_v2_proxy(&mut peer, "queued-resume-client").await?;

        let upstream_id = agent_client_protocol::RawJsonRpcId::Number(2);
        peer.tx
            .unbounded_send(TransportFrame::Single(RawJsonRpcMessage::request(
                "session/resume".to_owned(),
                serde_json::to_value(v2::ResumeSessionRequest::new(
                    session_id.clone(),
                    "/queued-v2-resume",
                ))
                .expect("resume request should serialize"),
                upstream_id.clone(),
            )?))
            .expect("proxy should accept session/resume");

        let Some(TransportFrame::Single(RawJsonRpcMessage::Request(forwarded))) =
            peer.rx.next().await
        else {
            panic!("expected a forwarded session/resume request");
        };
        let successor = SuccessorMessage::<v2::ResumeSessionRequest>::parse_message(
            forwarded.method.as_ref(),
            &params_value(forwarded.params.as_ref()),
        )?;
        assert_eq!(successor.message.session_id, session_id);

        let update = SuccessorMessage::new(v2::UpdateSessionNotification::new(
            notification_session_id,
            v2::SessionUpdate::StateUpdate(v2::StateUpdate::Running(v2::RunningStateUpdate::new())),
        ))
        .to_untyped_message()?;
        let (method, params) = update.into_parts();
        let notification = RawJsonRpcMessage::notification(method, params)?;
        let response = RawJsonRpcMessage::response(
            forwarded.id,
            Ok(serde_json::to_value(setup_response).expect("resume response should serialize")),
        );
        for message in [notification, response] {
            peer.tx
                .unbounded_send(TransportFrame::Single(message))
                .expect("proxy should accept the queued messages");
        }

        let Some(TransportFrame::Single(RawJsonRpcMessage::Notification(notification))) =
            peer.rx.next().await
        else {
            panic!("expected replay to be forwarded before the resume response");
        };
        let replay = v2::UpdateSessionNotification::parse_message(
            notification.method.as_ref(),
            &params_value(notification.params.as_ref()),
        )?;
        assert_eq!(replay.session_id, session_id);
        assert!(matches!(
            replay.update,
            v2::SessionUpdate::StateUpdate(v2::StateUpdate::Running(_))
        ));

        let Some(TransportFrame::Single(RawJsonRpcMessage::Response(
            agent_client_protocol::RawJsonRpcResponse::Result { id, result },
        ))) = peer.rx.next().await
        else {
            panic!("expected the resume response after replay");
        };
        assert_eq!(id, upstream_id);
        let response = v2::ResumeSessionResponse::from_json("session/resume", result)?;
        assert_eq!(response, upstream_response);

        peer_done_tx
            .send(())
            .map_err(|()| agent_client_protocol::Error::internal_error())
    };

    tokio::time::timeout(TIMEOUT, async { futures::try_join!(proxy, peer) })
        .await
        .expect("queued v2 resume replay was not routed before its response")
        .expect("v2 proxy resume connection failed");
}