acp_utils/
testing.rs

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

acp_utils/
testing.rs
