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

acp_utils/
testing.rs
