1use super::{idle_notification, initialize_response};
2use crate::notifications::{SessionPreviewParams, SessionPreviewResponse};
3use agent_client_protocol::schema::v2::{
4 AgentCapabilities, CancelSessionNotification, CloseSessionRequest, CloseSessionResponse, CompactionStatus,
5 CompactionUpdate, ContentChunk, Implementation, InitializeRequest, InitializeResponse, ListSessionsRequest,
6 ListSessionsResponse, LoginAuthRequest, LoginAuthResponse, NewSessionRequest, NewSessionResponse, PromptRequest,
7 PromptResponse, ResumeSessionRequest, ResumeSessionResponse, SessionId, SessionInfo, SessionUpdate,
8 SetSessionConfigOptionRequest, SetSessionConfigOptionResponse, UpdateSessionNotification,
9};
10use agent_client_protocol::util::MatchDispatchFrom;
11use agent_client_protocol::{
12 self as acp, Agent, Client, ConnectionTo, Dispatch, HandleDispatchFrom, Handled, NullRun, Responder, V2Builder,
13};
14use tokio::sync::mpsc;
15
16pub struct FakeAgent {
17 initialize: InitializeResponse,
18 new_session: Option<NewSessionResponse>,
19 sessions: Option<Vec<SessionInfo>>,
20 previews: Vec<SessionPreviewResponse>,
21 login_method: Option<String>,
22 hold_config: bool,
23 hold_list_sessions: bool,
24 replay: Vec<UpdateSessionNotification>,
25 live: Vec<UpdateSessionNotification>,
26 capture: Option<Capture>,
27}
28
29pub struct FakeAgentRequests {
30 pub connection: mpsc::UnboundedReceiver<ConnectionTo<Client>>,
31 pub initialize: mpsc::UnboundedReceiver<InitializeRequest>,
32 pub new_session: mpsc::UnboundedReceiver<NewSessionRequest>,
33 pub login: mpsc::UnboundedReceiver<LoginAuthRequest>,
34 pub config: mpsc::UnboundedReceiver<SetSessionConfigOptionRequest>,
35 pub prompt: mpsc::UnboundedReceiver<(PromptRequest, Responder<PromptResponse>)>,
36 pub resume: mpsc::UnboundedReceiver<(ResumeSessionRequest, Responder<ResumeSessionResponse>)>,
37 pub cancel: mpsc::UnboundedReceiver<CancelSessionNotification>,
38 pub pending_config: mpsc::UnboundedReceiver<Responder<SetSessionConfigOptionResponse>>,
39 pub list_sessions: mpsc::UnboundedReceiver<(ListSessionsRequest, Responder<ListSessionsResponse>)>,
40 pub close_session: mpsc::UnboundedReceiver<CloseSessionRequest>,
41}
42
43impl Default for FakeAgent {
44 fn default() -> Self {
45 Self {
46 initialize: initialize_response(),
47 new_session: None,
48 sessions: None,
49 previews: Vec::new(),
50 login_method: None,
51 hold_config: false,
52 hold_list_sessions: false,
53 replay: Vec::new(),
54 live: Vec::new(),
55 capture: None,
56 }
57 }
58}
59
60impl FakeAgent {
61 pub fn remote_server(mut self, info: &crate::notifications::RemoteServerInfo) -> Self {
62 self.initialize.meta = Some(info.to_meta());
63 self
64 }
65
66 pub fn session_preview(mut self, preview: SessionPreviewResponse) -> Self {
67 self.previews.push(preview);
68 self
69 }
70
71 pub fn agent_info(mut self, info: Implementation) -> Self {
72 self.initialize.info = info;
73 self
74 }
75 pub fn capabilities(mut self, capabilities: AgentCapabilities) -> Self {
76 self.initialize.capabilities = capabilities;
77 self
78 }
79 pub fn login_method(mut self, method: &str) -> Self {
80 self.login_method = Some(method.into());
81 self
82 }
83 pub fn hold_config(mut self, hold: bool) -> Self {
84 self.hold_config = hold;
85 self
86 }
87 pub fn hold_list_sessions(mut self, hold: bool) -> Self {
88 self.hold_list_sessions = hold;
89 self
90 }
91 pub fn new_session_response(mut self, response: NewSessionResponse) -> Self {
92 self.new_session = Some(response);
93 self
94 }
95 pub fn sessions(mut self, sessions: Vec<SessionInfo>) -> Self {
96 self.sessions = Some(sessions);
97 self
98 }
99 pub fn replay_message(mut self, session_id: &str, text: &str) -> Self {
100 self.replay.push(message(session_id, text));
101 self
102 }
103 pub fn live_message(mut self, session_id: &str, text: &str) -> Self {
104 self.live.push(message(session_id, text));
105 self
106 }
107 pub fn compaction(mut self, session_id: &str, compaction_id: &str, status: CompactionStatus) -> Self {
108 self.replay.push(UpdateSessionNotification::new(
109 session_id,
110 SessionUpdate::CompactionUpdate(CompactionUpdate::new(compaction_id, status)),
111 ));
112 self
113 }
114
115 pub fn capture(mut self) -> (Self, FakeAgentRequests) {
116 let (connection, connection_rx) = mpsc::unbounded_channel();
117 let (initialize, initialize_rx) = mpsc::unbounded_channel();
118 let (new_session, new_session_rx) = mpsc::unbounded_channel();
119 let (login, login_rx) = mpsc::unbounded_channel();
120 let (config, config_rx) = mpsc::unbounded_channel();
121 let (prompt, prompt_rx) = mpsc::unbounded_channel();
122 let (resume, resume_rx) = mpsc::unbounded_channel();
123 let (cancel, cancel_rx) = mpsc::unbounded_channel();
124 let (pending_config, pending_config_rx) = mpsc::unbounded_channel();
125 let (list_sessions, list_sessions_rx) = mpsc::unbounded_channel();
126 let (close_session, close_session_rx) = mpsc::unbounded_channel();
127 self.capture = Some(Capture {
128 connection,
129 initialize,
130 new_session,
131 login,
132 config,
133 prompt,
134 resume,
135 cancel,
136 pending_config,
137 list_sessions,
138 close_session,
139 });
140 (
141 self,
142 FakeAgentRequests {
143 connection: connection_rx,
144 initialize: initialize_rx,
145 new_session: new_session_rx,
146 login: login_rx,
147 config: config_rx,
148 prompt: prompt_rx,
149 resume: resume_rx,
150 cancel: cancel_rx,
151 pending_config: pending_config_rx,
152 list_sessions: list_sessions_rx,
153 close_session: close_session_rx,
154 },
155 )
156 }
157
158 pub fn agent(self) -> V2Builder<Agent, impl HandleDispatchFrom<Client>, NullRun> {
159 Agent.v2().name("fake-agent").with_handler(self)
160 }
161
162 pub async fn build(self) -> Result<crate::client::AcpClient, crate::client::AcpClientError> {
163 let (agent, client) = super::Channel::duplex();
164 tokio::task::spawn_local(self.agent().connect_to(agent));
165 crate::client::connect_acp_client(client, super::initialize_request()).await
166 }
167}
168
169impl HandleDispatchFrom<Client> for FakeAgent {
170 async fn handle_dispatch_from(
171 &mut self,
172 message: Dispatch,
173 cx: ConnectionTo<Client>,
174 ) -> Result<Handled<Dispatch>, acp::Error> {
175 MatchDispatchFrom::new(message, &cx)
176 .if_request(async |request: InitializeRequest, responder| {
177 if let Some(capture) = &self.capture {
178 let _ = capture.connection.send(cx.clone());
179 let _ = capture.initialize.send(request);
180 }
181 responder.respond(self.initialize.clone())
182 })
183 .await
184 .if_request(async |request: NewSessionRequest, responder| {
185 let Some(response) = &self.new_session else {
186 return Ok(Handled::No { message: (request, responder), retry: false });
187 };
188 if let Some(capture) = &self.capture {
189 let _ = capture.new_session.send(request);
190 }
191 responder.respond(response.clone())?;
192 Ok(Handled::Yes)
193 })
194 .await
195 .if_request(async |request: SessionPreviewParams, responder| {
196 let Some(preview) = self.previews.iter().find(|preview| preview.session_id == request.session_id)
197 else {
198 return Ok(Handled::No { message: (request, responder), retry: false });
199 };
200 responder.respond(preview.clone())?;
201 Ok(Handled::Yes)
202 })
203 .await
204 .if_request(async |request: ListSessionsRequest, responder| {
205 if self.hold_list_sessions
206 && let Some(capture) = &self.capture
207 {
208 let _ = capture.list_sessions.send((request, responder));
209 return Ok(Handled::Yes);
210 }
211 let Some(sessions) = &self.sessions else {
212 return Ok(Handled::No { message: (request, responder), retry: false });
213 };
214 responder.respond(ListSessionsResponse::new(sessions.clone()))?;
215 Ok(Handled::Yes)
216 })
217 .await
218 .if_request(async |request: CloseSessionRequest, responder| {
219 if let Some(capture) = &self.capture {
220 let _ = capture.close_session.send(request);
221 }
222 responder.respond(CloseSessionResponse::new())
223 })
224 .await
225 .if_request(async |request: LoginAuthRequest, responder| {
226 let allowed = self.login_method.as_deref() == Some(request.method_id.0.as_ref());
227 if let Some(capture) = &self.capture {
228 let _ = capture.login.send(request);
229 }
230 if allowed {
231 responder.respond(LoginAuthResponse::new())
232 } else {
233 responder.respond_with_error(acp::Error::invalid_params())
234 }
235 })
236 .await
237 .if_request(async |request: SetSessionConfigOptionRequest, responder| {
238 if let Some(capture) = &self.capture {
239 let _ = capture.config.send(request);
240 if self.hold_config {
241 let _ = capture.pending_config.send(responder);
242 return Ok(());
243 }
244 }
245 responder.respond(SetSessionConfigOptionResponse::new(vec![]))
246 })
247 .await
248 .if_request(async |request: PromptRequest, responder| {
249 if let Some(capture) = &self.capture {
250 let _ = capture.prompt.send((request, responder));
251 } else {
252 responder.respond(PromptResponse::new())?;
253 }
254 Ok(())
255 })
256 .await
257 .if_request(async |request: ResumeSessionRequest, responder| self.resume(request, responder, &cx))
258 .await
259 .if_notification(async |notification: CancelSessionNotification| {
260 if let Some(capture) = &self.capture {
261 let _ = capture.cancel.send(notification);
262 }
263 Ok(())
264 })
265 .await
266 .done()
267 }
268
269 fn describe_chain(&self) -> impl std::fmt::Debug {
270 "FakeAgent"
271 }
272}
273
274impl FakeAgent {
275 fn resume(
276 &self,
277 request: ResumeSessionRequest,
278 responder: Responder<ResumeSessionResponse>,
279 cx: &ConnectionTo<Client>,
280 ) -> Result<(), acp::Error> {
281 if let Some(capture) = &self.capture {
282 let _ = capture.resume.send((request, responder));
283 return Ok(());
284 }
285 if request.replay_from.is_none() {
286 return responder.respond(ResumeSessionResponse::new());
287 }
288 for notification in &self.replay {
289 cx.send_notification(notification.clone())?;
290 }
291 cx.send_notification(idle_notification(request.session_id, None))?;
292 responder.respond(ResumeSessionResponse::new())?;
293 for notification in &self.live {
294 cx.send_notification(notification.clone())?;
295 }
296 Ok(())
297 }
298}
299
300struct Capture {
301 connection: mpsc::UnboundedSender<ConnectionTo<Client>>,
302 initialize: mpsc::UnboundedSender<InitializeRequest>,
303 new_session: mpsc::UnboundedSender<NewSessionRequest>,
304 login: mpsc::UnboundedSender<LoginAuthRequest>,
305 config: mpsc::UnboundedSender<SetSessionConfigOptionRequest>,
306 prompt: mpsc::UnboundedSender<(PromptRequest, Responder<PromptResponse>)>,
307 resume: mpsc::UnboundedSender<(ResumeSessionRequest, Responder<ResumeSessionResponse>)>,
308 cancel: mpsc::UnboundedSender<CancelSessionNotification>,
309 pending_config: mpsc::UnboundedSender<Responder<SetSessionConfigOptionResponse>>,
310 list_sessions: mpsc::UnboundedSender<(ListSessionsRequest, Responder<ListSessionsResponse>)>,
311 close_session: mpsc::UnboundedSender<CloseSessionRequest>,
312}
313
314fn message(session_id: &str, text: &str) -> UpdateSessionNotification {
315 UpdateSessionNotification::new(
316 SessionId::new(session_id),
317 SessionUpdate::AgentMessageChunk(ContentChunk::new(text.into(), "message")),
318 )
319}