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
9mod fake_agent;
10pub use fake_agent::{FakeAgent, FakeAgentRequests};
11
12use crate::notifications::McpNotification;
13use agent_client_protocol::schema::ProtocolVersion;
14use agent_client_protocol::schema::v2::{
15    CompleteElicitationNotification, CreateElicitationRequest, CreateElicitationResponse, ElicitationFormMode,
16    ElicitationSchema, ElicitationSessionScope, IdleStateUpdate, Implementation, InitializeRequest, InitializeResponse,
17    PlanEntry, PlanId, PlanUpdate, PlanUpdateContent, RunningStateUpdate, SessionId, SessionUpdate, StateUpdate,
18    StopReason, UpdateSessionNotification,
19};
20use agent_client_protocol::{
21    self as acp, Agent, ByteStreams, Client, ConnectionTo, HandleConnectionClose, HandleDispatchFrom, NullRun,
22    Responder, RunWithConnectionTo, V2Builder,
23};
24use std::collections::VecDeque;
25use std::sync::{Arc, Mutex};
26use tokio::io::DuplexStream;
27use tokio::sync::{mpsc, oneshot};
28use tokio::task::spawn_local;
29use tokio_util::compat::{Compat, TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
30
31pub type DuplexByteStreams = ByteStreams<Compat<DuplexStream>, Compat<DuplexStream>>;
32
33pub struct TestPeer {
34    session_notifications: mpsc::UnboundedReceiver<UpdateSessionNotification>,
35    mcp_notifications: mpsc::UnboundedReceiver<McpNotification>,
36    elicitation_requests: mpsc::UnboundedReceiver<CreateElicitationRequest>,
37    elicitation_completions: mpsc::UnboundedReceiver<CompleteElicitationNotification>,
38    elicitation_responses: Arc<Mutex<VecDeque<CreateElicitationResponse>>>,
39    responder_capture: Arc<Mutex<Option<oneshot::Sender<Responder<CreateElicitationResponse>>>>>,
40}
41
42impl TestPeer {
43    pub fn new() -> (Self, V2Builder<Client, impl HandleDispatchFrom<Agent>, NullRun>) {
44        let (sn_tx, sn_rx) = mpsc::unbounded_channel::<UpdateSessionNotification>();
45        let (mcp_tx, mcp_rx) = mpsc::unbounded_channel::<McpNotification>();
46        let (el_tx, el_rx) = mpsc::unbounded_channel::<CreateElicitationRequest>();
47        let (complete_tx, complete_rx) = mpsc::unbounded_channel::<CompleteElicitationNotification>();
48        let elicitation_responses: Arc<Mutex<VecDeque<CreateElicitationResponse>>> =
49            Arc::new(Mutex::new(VecDeque::new()));
50        let responder_capture: Arc<Mutex<Option<oneshot::Sender<Responder<CreateElicitationResponse>>>>> =
51            Arc::new(Mutex::new(None));
52
53        let builder = Client
54            .v2()
55            .name("test-client")
56            .on_receive_notification(
57                {
58                    let tx = sn_tx;
59                    async move |n: UpdateSessionNotification, _cx| {
60                        let _ = tx.send(n);
61                        Ok(())
62                    }
63                },
64                acp::on_receive_notification!(),
65            )
66            .on_receive_notification(
67                {
68                    let tx = mcp_tx;
69                    async move |n: McpNotification, _cx| {
70                        let _ = tx.send(n);
71                        Ok(())
72                    }
73                },
74                acp::on_receive_notification!(),
75            )
76            .on_receive_notification(
77                {
78                    let tx = complete_tx;
79                    async move |notification: CompleteElicitationNotification, _cx| {
80                        let _ = tx.send(notification);
81                        Ok(())
82                    }
83                },
84                acp::on_receive_notification!(),
85            )
86            .on_receive_request(
87                {
88                    let tx = el_tx;
89                    let responses = elicitation_responses.clone();
90                    let capture = responder_capture.clone();
91                    async move |req: CreateElicitationRequest, responder: Responder<CreateElicitationResponse>, _cx| {
92                        if let Some(capture_tx) = capture.lock().unwrap().take() {
93                            return match capture_tx.send(responder) {
94                                Ok(()) => Ok(()),
95                                Err(responder) => responder.respond_with_error(acp::Error::internal_error()),
96                            };
97                        }
98                        let _ = tx.send(req);
99                        let queued = responses.lock().unwrap().pop_front();
100                        match queued {
101                            Some(response) => responder.respond(response),
102                            None => responder.respond_with_error(acp::Error::method_not_found()),
103                        }
104                    }
105                },
106                acp::on_receive_request!(),
107            );
108
109        let peer = Self {
110            session_notifications: sn_rx,
111            mcp_notifications: mcp_rx,
112            elicitation_requests: el_rx,
113            elicitation_completions: complete_rx,
114            elicitation_responses,
115            responder_capture,
116        };
117        (peer, builder)
118    }
119
120    pub async fn next_session_notification(&mut self) -> UpdateSessionNotification {
121        self.session_notifications.recv().await.expect("peer channel closed")
122    }
123
124    pub async fn next_mcp_notification(&mut self) -> McpNotification {
125        self.mcp_notifications.recv().await.expect("peer channel closed")
126    }
127
128    pub async fn next_elicitation_request(&mut self) -> CreateElicitationRequest {
129        self.elicitation_requests.recv().await.expect("peer channel closed")
130    }
131
132    pub async fn next_elicitation_completion(&mut self) -> CompleteElicitationNotification {
133        self.elicitation_completions.recv().await.expect("peer channel closed")
134    }
135
136    pub fn queue_elicitation_response(&self, response: CreateElicitationResponse) {
137        self.elicitation_responses.lock().unwrap().push_back(response);
138    }
139
140    pub fn capture_next_elicitation(&self) -> oneshot::Receiver<Responder<CreateElicitationResponse>> {
141        let (sender, receiver) = oneshot::channel();
142        *self.responder_capture.lock().unwrap() = Some(sender);
143        receiver
144    }
145
146    pub async fn fake_elicitation(
147        &mut self,
148        cx: &ConnectionTo<Client>,
149    ) -> (Responder<CreateElicitationResponse>, oneshot::Receiver<CreateElicitationResponse>) {
150        let responder_rx = self.capture_next_elicitation();
151
152        let (response_tx, response_rx) = oneshot::channel::<CreateElicitationResponse>();
153        let cx = cx.clone();
154        spawn_local(async move {
155            if let Ok(resp) = cx.send_request(placeholder_params()).block_task().await {
156                let _ = response_tx.send(resp);
157            }
158        });
159
160        let responder = responder_rx.await.expect("client handler must capture responder");
161        (responder, response_rx)
162    }
163}
164
165/// In-memory ACP transport pair: `(agent_transport, client_transport)`. Hand
166/// each half to a `connect_to` / `connect_with` call on the corresponding
167/// side. Must be used inside a `LocalSet` since the runners are `spawn_local`'d.
168pub fn duplex_pair() -> (DuplexByteStreams, DuplexByteStreams) {
169    let (agent_writer, client_reader) = tokio::io::duplex(4096);
170    let (client_writer, agent_reader) = tokio::io::duplex(4096);
171    let agent_transport = ByteStreams::new(agent_writer.compat_write(), agent_reader.compat());
172    let client_transport = ByteStreams::new(client_writer.compat_write(), client_reader.compat());
173    (agent_transport, client_transport)
174}
175
176/// Build a live `ConnectionTo<Client>` over an in-memory duplex transport with
177/// a peer on the other end. Must be called inside a `LocalSet`.
178pub async fn test_connection() -> (ConnectionTo<Client>, TestPeer) {
179    let (peer, client_builder) = TestPeer::new();
180    let agent = Agent.v2().name("test-agent").on_receive_request(
181        async |_: InitializeRequest, responder: Responder<InitializeResponse>, _cx| {
182            responder.respond(initialize_response())
183        },
184        acp::on_receive_request!(),
185    );
186    let pair = connect_pair(agent, client_builder).await;
187    pair.client.send_request(initialize_request()).block_task().await.expect("initialize test peers");
188    (pair.agent, peer)
189}
190
191pub struct ConnectedPair {
192    pub agent: ConnectionTo<Client>,
193    pub client: ConnectionTo<Agent>,
194    pub agent_task: tokio::task::JoinHandle<Result<(), acp::Error>>,
195    pub client_task: tokio::task::JoinHandle<Result<(), acp::Error>>,
196}
197
198pub async fn connect_pair<T, U, V, X, Y, Z>(
199    agent: V2Builder<Agent, T, U, V>,
200    client: V2Builder<Client, X, Y, Z>,
201) -> ConnectedPair
202where
203    T: HandleDispatchFrom<Client> + 'static,
204    U: RunWithConnectionTo<Client> + 'static,
205    V: HandleConnectionClose<Client> + 'static,
206    X: HandleDispatchFrom<Agent> + 'static,
207    Y: RunWithConnectionTo<Agent> + 'static,
208    Z: HandleConnectionClose<Agent> + 'static,
209{
210    let (agent_transport, client_transport) = duplex_pair();
211    let (agent_tx, agent_rx) = oneshot::channel();
212    let (client_tx, client_rx) = oneshot::channel();
213    let agent_task =
214        spawn_local(async move { agent.with_runner(CaptureConnection(agent_tx)).connect_to(agent_transport).await });
215    let client_task =
216        spawn_local(async move { client.with_runner(CaptureConnection(client_tx)).connect_to(client_transport).await });
217    ConnectedPair {
218        agent: agent_rx.await.expect("agent connection"),
219        client: client_rx.await.expect("client connection"),
220        agent_task,
221        client_task,
222    }
223}
224
225pub struct CaptureConnection<R: acp::Role>(pub oneshot::Sender<ConnectionTo<R>>);
226
227impl<R: acp::Role> RunWithConnectionTo<R> for CaptureConnection<R> {
228    async fn run_with_connection_to(self, cx: ConnectionTo<R>) -> Result<(), acp::Error> {
229        let _ = self.0.send(cx.clone());
230        cx.incoming_closed().await;
231        Ok(())
232    }
233}
234
235/// Initialization request from an in-memory v2 client.
236pub fn initialize_request() -> InitializeRequest {
237    InitializeRequest::new(ProtocolVersion::V2, Implementation::new("test-client", "0.0.0"))
238}
239
240/// Initialization response from an in-memory v2 agent.
241pub fn initialize_response() -> InitializeResponse {
242    InitializeResponse::new(ProtocolVersion::V2, Implementation::new("test-agent", "0.0.0"))
243}
244
245/// A live foreground turn has started.
246pub fn running_notification(session_id: impl Into<SessionId>) -> UpdateSessionNotification {
247    UpdateSessionNotification::new(
248        session_id,
249        SessionUpdate::StateUpdate(StateUpdate::Running(RunningStateUpdate::new())),
250    )
251}
252
253/// Foreground work is idle, optionally with a reported stop reason.
254pub fn idle_notification(
255    session_id: impl Into<SessionId>,
256    stop_reason: Option<StopReason>,
257) -> UpdateSessionNotification {
258    UpdateSessionNotification::new(
259        session_id,
260        SessionUpdate::StateUpdate(StateUpdate::Idle(IdleStateUpdate::new().stop_reason(stop_reason))),
261    )
262}
263
264/// Replace the entries of an agent-owned plan.
265pub fn plan_notification(
266    session_id: impl Into<SessionId>,
267    plan_id: impl Into<PlanId>,
268    entries: Vec<PlanEntry>,
269) -> UpdateSessionNotification {
270    UpdateSessionNotification::new(
271        session_id,
272        SessionUpdate::PlanUpdate(PlanUpdate::new(PlanUpdateContent::items(plan_id, entries))),
273    )
274}
275
276fn placeholder_params() -> CreateElicitationRequest {
277    CreateElicitationRequest::new(
278        ElicitationFormMode::new(ElicitationSessionScope::new("test-session"), ElicitationSchema::new()),
279        String::new(),
280    )
281}
acp_utils/testing.rs

acp_utils/
testing.rs
