acp_utils/
testing.rs

1//! Duplex-backed test harness for ACP connections.
2//!
3//! [`test_connection`] returns a full `(ConnectionTo<Client>, TestPeer)` pair
4//! over an in-memory duplex transport. Use it for integration-style tests that
5//! need to exercise the full serialize/dispatch path (so wire-format
6//! regressions like extension method-name typos surface in tests).
7//!
8
9use crate::notifications::McpNotification;
10use agent_client_protocol::schema::v1::{
11    CompleteElicitationNotification, CreateElicitationRequest, CreateElicitationResponse, ElicitationFormMode,
12    ElicitationSchema, ElicitationSessionScope, SessionNotification,
13};
14use agent_client_protocol::{
15    self as acp, Agent, Builder, ByteStreams, Client, ConnectionTo, HandleDispatchFrom, NullRun, Responder,
16};
17use std::collections::VecDeque;
18use std::sync::{Arc, Mutex};
19use tokio::io::DuplexStream;
20use tokio::sync::{mpsc, oneshot};
21use tokio::task::spawn_local;
22use tokio_util::compat::{Compat, TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
23
24pub type DuplexByteStreams = ByteStreams<Compat<DuplexStream>, Compat<DuplexStream>>;
25
26pub struct TestPeer {
27    session_notifications: mpsc::UnboundedReceiver<SessionNotification>,
28    mcp_notifications: mpsc::UnboundedReceiver<McpNotification>,
29    elicitation_requests: mpsc::UnboundedReceiver<CreateElicitationRequest>,
30    elicitation_completions: mpsc::UnboundedReceiver<CompleteElicitationNotification>,
31    elicitation_responses: Arc<Mutex<VecDeque<CreateElicitationResponse>>>,
32    responder_capture: Arc<Mutex<Option<oneshot::Sender<Responder<CreateElicitationResponse>>>>>,
33}
34
35impl TestPeer {
36    pub fn new() -> (Self, Builder<Client, impl HandleDispatchFrom<Agent>, NullRun>) {
37        let (sn_tx, sn_rx) = mpsc::unbounded_channel::<SessionNotification>();
38        let (mcp_tx, mcp_rx) = mpsc::unbounded_channel::<McpNotification>();
39        let (el_tx, el_rx) = mpsc::unbounded_channel::<CreateElicitationRequest>();
40        let (complete_tx, complete_rx) = mpsc::unbounded_channel::<CompleteElicitationNotification>();
41        let elicitation_responses: Arc<Mutex<VecDeque<CreateElicitationResponse>>> =
42            Arc::new(Mutex::new(VecDeque::new()));
43        let responder_capture: Arc<Mutex<Option<oneshot::Sender<Responder<CreateElicitationResponse>>>>> =
44            Arc::new(Mutex::new(None));
45
46        let builder = Client
47            .builder()
48            .on_receive_notification(
49                {
50                    let tx = sn_tx;
51                    async move |n: SessionNotification, _cx| {
52                        let _ = tx.send(n);
53                        Ok(())
54                    }
55                },
56                acp::on_receive_notification!(),
57            )
58            .on_receive_notification(
59                {
60                    let tx = mcp_tx;
61                    async move |n: McpNotification, _cx| {
62                        let _ = tx.send(n);
63                        Ok(())
64                    }
65                },
66                acp::on_receive_notification!(),
67            )
68            .on_receive_notification(
69                {
70                    let tx = complete_tx;
71                    async move |notification: CompleteElicitationNotification, _cx| {
72                        let _ = tx.send(notification);
73                        Ok(())
74                    }
75                },
76                acp::on_receive_notification!(),
77            )
78            .on_receive_request(
79                {
80                    let tx = el_tx;
81                    let responses = elicitation_responses.clone();
82                    let capture = responder_capture.clone();
83                    async move |req: CreateElicitationRequest, responder: Responder<CreateElicitationResponse>, _cx| {
84                        if let Some(capture_tx) = capture.lock().unwrap().take() {
85                            return match capture_tx.send(responder) {
86                                Ok(()) => Ok(()),
87                                Err(responder) => responder.respond_with_error(acp::Error::internal_error()),
88                            };
89                        }
90                        let _ = tx.send(req);
91                        let queued = responses.lock().unwrap().pop_front();
92                        match queued {
93                            Some(response) => responder.respond(response),
94                            None => responder.respond_with_error(acp::Error::method_not_found()),
95                        }
96                    }
97                },
98                acp::on_receive_request!(),
99            );
100
101        let peer = Self {
102            session_notifications: sn_rx,
103            mcp_notifications: mcp_rx,
104            elicitation_requests: el_rx,
105            elicitation_completions: complete_rx,
106            elicitation_responses,
107            responder_capture,
108        };
109        (peer, builder)
110    }
111
112    pub async fn next_session_notification(&mut self) -> SessionNotification {
113        self.session_notifications.recv().await.expect("peer channel closed")
114    }
115
116    pub async fn next_mcp_notification(&mut self) -> McpNotification {
117        self.mcp_notifications.recv().await.expect("peer channel closed")
118    }
119
120    pub async fn next_elicitation_request(&mut self) -> CreateElicitationRequest {
121        self.elicitation_requests.recv().await.expect("peer channel closed")
122    }
123
124    pub async fn next_elicitation_completion(&mut self) -> CompleteElicitationNotification {
125        self.elicitation_completions.recv().await.expect("peer channel closed")
126    }
127
128    pub fn queue_elicitation_response(&self, response: CreateElicitationResponse) {
129        self.elicitation_responses.lock().unwrap().push_back(response);
130    }
131
132    pub async fn fake_elicitation(
133        &mut self,
134        cx: &ConnectionTo<Client>,
135    ) -> (Responder<CreateElicitationResponse>, oneshot::Receiver<CreateElicitationResponse>) {
136        let (responder_tx, responder_rx) = oneshot::channel::<Responder<CreateElicitationResponse>>();
137        *self.responder_capture.lock().unwrap() = Some(responder_tx);
138
139        let (response_tx, response_rx) = oneshot::channel::<CreateElicitationResponse>();
140        let cx = cx.clone();
141        spawn_local(async move {
142            if let Ok(resp) = cx.send_request(placeholder_params()).block_task().await {
143                let _ = response_tx.send(resp);
144            }
145        });
146
147        let responder = responder_rx.await.expect("client handler must capture responder");
148        (responder, response_rx)
149    }
150}
151
152/// In-memory ACP transport pair: `(agent_transport, client_transport)`. Hand
153/// each half to a `connect_to` / `connect_with` call on the corresponding
154/// side. Must be used inside a `LocalSet` since the runners are `spawn_local`'d.
155pub fn duplex_pair() -> (DuplexByteStreams, DuplexByteStreams) {
156    let (agent_writer, client_reader) = tokio::io::duplex(4096);
157    let (client_writer, agent_reader) = tokio::io::duplex(4096);
158    let agent_transport = ByteStreams::new(agent_writer.compat_write(), agent_reader.compat());
159    let client_transport = ByteStreams::new(client_writer.compat_write(), client_reader.compat());
160    (agent_transport, client_transport)
161}
162
163/// Build a live `ConnectionTo<Client>` over an in-memory duplex transport with
164/// a peer on the other end. Must be called inside a `LocalSet`.
165pub async fn test_connection() -> (ConnectionTo<Client>, TestPeer) {
166    let (peer, client_builder) = TestPeer::new();
167    let (agent_transport, client_transport) = duplex_pair();
168
169    spawn_local(async move {
170        let _ = client_builder.connect_to(client_transport).await;
171    });
172
173    let (cx_tx, cx_rx) = oneshot::channel::<ConnectionTo<Client>>();
174    spawn_local(async move {
175        let _ = Agent
176            .builder()
177            .connect_with(agent_transport, async move |cx: ConnectionTo<Client>| {
178                let _ = cx_tx.send(cx);
179                std::future::pending::<()>().await;
180                Ok(())
181            })
182            .await;
183    });
184
185    let cx = cx_rx.await.expect("agent side connect_with produced a ConnectionTo");
186    (cx, peer)
187}
188
189fn placeholder_params() -> CreateElicitationRequest {
190    CreateElicitationRequest::new(
191        ElicitationFormMode::new(ElicitationSessionScope::new("test-session"), ElicitationSchema::new()),
192        String::new(),
193    )
194}
acp_utils/testing.rs

acp_utils/
testing.rs
