1use std::{
2 path::{Path, PathBuf},
3 sync::Arc,
4 time::Duration,
5};
6
7use crate::{AgentServer, AgentServerSettings, AllAgentServersSettings};
8use acp_thread::{AcpThread, AgentThreadEntry, ToolCall, ToolCallStatus};
9use agent_client_protocol as acp;
10
11use futures::{FutureExt, StreamExt, channel::mpsc, select};
12use gpui::{Entity, TestAppContext};
13use indoc::indoc;
14use project::{FakeFs, Project};
15use settings::{Settings, SettingsStore};
16use util::path;
17
18pub async fn test_basic(server: impl AgentServer + 'static, cx: &mut TestAppContext) {
19 let fs = init_test(cx).await;
20 let project = Project::test(fs, [], cx).await;
21 let thread = new_test_thread(server, project.clone(), "/private/tmp", cx).await;
22
23 thread
24 .update(cx, |thread, cx| thread.send_raw("Hello from Zed!", cx))
25 .await
26 .unwrap();
27
28 thread.read_with(cx, |thread, _| {
29 assert!(
30 thread.entries().len() >= 2,
31 "Expected at least 2 entries. Got: {:?}",
32 thread.entries()
33 );
34 assert!(matches!(
35 thread.entries()[0],
36 AgentThreadEntry::UserMessage(_)
37 ));
38 assert!(matches!(
39 thread.entries()[1],
40 AgentThreadEntry::AssistantMessage(_)
41 ));
42 });
43}
44
45pub async fn test_path_mentions(server: impl AgentServer + 'static, cx: &mut TestAppContext) {
46 let _fs = init_test(cx).await;
47
48 let tempdir = tempfile::tempdir().unwrap();
49 std::fs::write(
50 tempdir.path().join("foo.rs"),
51 indoc! {"
52 fn main() {
53 println!(\"Hello, world!\");
54 }
55 "},
56 )
57 .expect("failed to write file");
58 let project = Project::example([tempdir.path()], &mut cx.to_async()).await;
59 let thread = new_test_thread(server, project.clone(), tempdir.path(), cx).await;
60 thread
61 .update(cx, |thread, cx| {
62 thread.send(
63 vec![
64 acp::ContentBlock::Text(acp::TextContent {
65 text: "Read the file ".into(),
66 annotations: None,
67 }),
68 acp::ContentBlock::ResourceLink(acp::ResourceLink {
69 uri: "foo.rs".into(),
70 name: "foo.rs".into(),
71 annotations: None,
72 description: None,
73 mime_type: None,
74 size: None,
75 title: None,
76 }),
77 acp::ContentBlock::Text(acp::TextContent {
78 text: " and tell me what the content of the println! is".into(),
79 annotations: None,
80 }),
81 ],
82 cx,
83 )
84 })
85 .await
86 .unwrap();
87
88 thread.read_with(cx, |thread, cx| {
89 assert!(matches!(
90 thread.entries()[0],
91 AgentThreadEntry::UserMessage(_)
92 ));
93 let assistant_message = &thread
94 .entries()
95 .iter()
96 .rev()
97 .find_map(|entry| match entry {
98 AgentThreadEntry::AssistantMessage(msg) => Some(msg),
99 _ => None,
100 })
101 .unwrap();
102
103 assert!(
104 assistant_message.to_markdown(cx).contains("Hello, world!"),
105 "unexpected assistant message: {:?}",
106 assistant_message.to_markdown(cx)
107 );
108 });
109
110 drop(tempdir);
111}
112
113pub async fn test_tool_call(server: impl AgentServer + 'static, cx: &mut TestAppContext) {
114 let _fs = init_test(cx).await;
115
116 let tempdir = tempfile::tempdir().unwrap();
117 let foo_path = tempdir.path().join("foo");
118 std::fs::write(&foo_path, "Lorem ipsum dolor").expect("failed to write file");
119
120 let project = Project::example([tempdir.path()], &mut cx.to_async()).await;
121 let thread = new_test_thread(server, project.clone(), "/private/tmp", cx).await;
122
123 thread
124 .update(cx, |thread, cx| {
125 thread.send_raw(
126 &format!("Read {} and tell me what you see.", foo_path.display()),
127 cx,
128 )
129 })
130 .await
131 .unwrap();
132 thread.read_with(cx, |thread, _cx| {
133 assert!(thread.entries().iter().any(|entry| {
134 matches!(
135 entry,
136 AgentThreadEntry::ToolCall(ToolCall {
137 status: ToolCallStatus::Allowed { .. },
138 ..
139 })
140 )
141 }));
142 assert!(
143 thread
144 .entries()
145 .iter()
146 .any(|entry| { matches!(entry, AgentThreadEntry::AssistantMessage(_)) })
147 );
148 });
149
150 drop(tempdir);
151}
152
153pub async fn test_tool_call_with_confirmation(
154 server: impl AgentServer + 'static,
155 allow_option_id: acp::PermissionOptionId,
156 cx: &mut TestAppContext,
157) {
158 let fs = init_test(cx).await;
159 let project = Project::test(fs, [path!("/private/tmp").as_ref()], cx).await;
160 let thread = new_test_thread(server, project.clone(), "/private/tmp", cx).await;
161 let full_turn = thread.update(cx, |thread, cx| {
162 thread.send_raw(
163 r#"Run exactly `touch hello.txt && echo "Hello, world!" | tee hello.txt` in the terminal."#,
164 cx,
165 )
166 });
167
168 run_until_first_tool_call(
169 &thread,
170 |entry| {
171 matches!(
172 entry,
173 AgentThreadEntry::ToolCall(ToolCall {
174 status: ToolCallStatus::WaitingForConfirmation { .. },
175 ..
176 })
177 )
178 },
179 cx,
180 )
181 .await;
182
183 let tool_call_id = thread.read_with(cx, |thread, cx| {
184 let AgentThreadEntry::ToolCall(ToolCall {
185 id,
186 label,
187 status: ToolCallStatus::WaitingForConfirmation { .. },
188 ..
189 }) = &thread
190 .entries()
191 .iter()
192 .find(|entry| matches!(entry, AgentThreadEntry::ToolCall(_)))
193 .unwrap()
194 else {
195 panic!();
196 };
197
198 let label = label.read(cx).source();
199 assert!(label.contains("touch"), "Got: {}", label);
200
201 id.clone()
202 });
203
204 thread.update(cx, |thread, cx| {
205 thread.authorize_tool_call(
206 tool_call_id,
207 allow_option_id,
208 acp::PermissionOptionKind::AllowOnce,
209 cx,
210 );
211
212 assert!(thread.entries().iter().any(|entry| matches!(
213 entry,
214 AgentThreadEntry::ToolCall(ToolCall {
215 status: ToolCallStatus::Allowed { .. },
216 ..
217 })
218 )));
219 });
220
221 full_turn.await.unwrap();
222
223 thread.read_with(cx, |thread, cx| {
224 let AgentThreadEntry::ToolCall(ToolCall {
225 content,
226 status: ToolCallStatus::Allowed { .. },
227 ..
228 }) = thread
229 .entries()
230 .iter()
231 .find(|entry| matches!(entry, AgentThreadEntry::ToolCall(_)))
232 .unwrap()
233 else {
234 panic!();
235 };
236
237 assert!(
238 content.iter().any(|c| c.to_markdown(cx).contains("Hello")),
239 "Expected content to contain 'Hello'"
240 );
241 });
242}
243
244pub async fn test_cancel(server: impl AgentServer + 'static, cx: &mut TestAppContext) {
245 let fs = init_test(cx).await;
246
247 let project = Project::test(fs, [path!("/private/tmp").as_ref()], cx).await;
248 let thread = new_test_thread(server, project.clone(), "/private/tmp", cx).await;
249 let full_turn = thread.update(cx, |thread, cx| {
250 thread.send_raw(
251 r#"Run exactly `touch hello.txt && echo "Hello, world!" | tee hello.txt` in the terminal."#,
252 cx,
253 )
254 });
255
256 let first_tool_call_ix = run_until_first_tool_call(
257 &thread,
258 |entry| {
259 matches!(
260 entry,
261 AgentThreadEntry::ToolCall(ToolCall {
262 status: ToolCallStatus::WaitingForConfirmation { .. },
263 ..
264 })
265 )
266 },
267 cx,
268 )
269 .await;
270
271 thread.read_with(cx, |thread, cx| {
272 let AgentThreadEntry::ToolCall(ToolCall {
273 id,
274 label,
275 status: ToolCallStatus::WaitingForConfirmation { .. },
276 ..
277 }) = &thread.entries()[first_tool_call_ix]
278 else {
279 panic!("{:?}", thread.entries()[1]);
280 };
281
282 let label = label.read(cx).source();
283 assert!(label.contains("touch"), "Got: {}", label);
284
285 id.clone()
286 });
287
288 let _ = thread.update(cx, |thread, cx| thread.cancel(cx));
289 full_turn.await.unwrap();
290 thread.read_with(cx, |thread, _| {
291 let AgentThreadEntry::ToolCall(ToolCall {
292 status: ToolCallStatus::Canceled,
293 ..
294 }) = &thread.entries()[first_tool_call_ix]
295 else {
296 panic!();
297 };
298 });
299
300 thread
301 .update(cx, |thread, cx| {
302 thread.send_raw(r#"Stop running and say goodbye to me."#, cx)
303 })
304 .await
305 .unwrap();
306 thread.read_with(cx, |thread, _| {
307 assert!(matches!(
308 &thread.entries().last().unwrap(),
309 AgentThreadEntry::AssistantMessage(..),
310 ))
311 });
312}
313
314#[macro_export]
315macro_rules! common_e2e_tests {
316 ($server:expr, allow_option_id = $allow_option_id:expr) => {
317 mod common_e2e {
318 use super::*;
319
320 #[::gpui::test]
321 #[cfg_attr(not(feature = "e2e"), ignore)]
322 async fn basic(cx: &mut ::gpui::TestAppContext) {
323 $crate::e2e_tests::test_basic($server, cx).await;
324 }
325
326 #[::gpui::test]
327 #[cfg_attr(not(feature = "e2e"), ignore)]
328 async fn path_mentions(cx: &mut ::gpui::TestAppContext) {
329 $crate::e2e_tests::test_path_mentions($server, cx).await;
330 }
331
332 #[::gpui::test]
333 #[cfg_attr(not(feature = "e2e"), ignore)]
334 async fn tool_call(cx: &mut ::gpui::TestAppContext) {
335 $crate::e2e_tests::test_tool_call($server, cx).await;
336 }
337
338 #[::gpui::test]
339 #[cfg_attr(not(feature = "e2e"), ignore)]
340 async fn tool_call_with_confirmation(cx: &mut ::gpui::TestAppContext) {
341 $crate::e2e_tests::test_tool_call_with_confirmation(
342 $server,
343 ::agent_client_protocol::PermissionOptionId($allow_option_id.into()),
344 cx,
345 )
346 .await;
347 }
348
349 #[::gpui::test]
350 #[cfg_attr(not(feature = "e2e"), ignore)]
351 async fn cancel(cx: &mut ::gpui::TestAppContext) {
352 $crate::e2e_tests::test_cancel($server, cx).await;
353 }
354 }
355 };
356}
357
358// Helpers
359
360pub async fn init_test(cx: &mut TestAppContext) -> Arc<FakeFs> {
361 env_logger::try_init().ok();
362
363 cx.update(|cx| {
364 let settings_store = SettingsStore::test(cx);
365 cx.set_global(settings_store);
366 Project::init_settings(cx);
367 language::init(cx);
368 crate::settings::init(cx);
369
370 crate::AllAgentServersSettings::override_global(
371 AllAgentServersSettings {
372 claude: Some(AgentServerSettings {
373 command: crate::claude::tests::local_command(),
374 }),
375 gemini: Some(AgentServerSettings {
376 command: crate::gemini::tests::local_command(),
377 }),
378 codex: Some(AgentServerSettings {
379 command: crate::codex::tests::local_command(),
380 }),
381 },
382 cx,
383 );
384 });
385
386 cx.executor().allow_parking();
387
388 FakeFs::new(cx.executor())
389}
390
391pub async fn new_test_thread(
392 server: impl AgentServer + 'static,
393 project: Entity<Project>,
394 current_dir: impl AsRef<Path>,
395 cx: &mut TestAppContext,
396) -> Entity<AcpThread> {
397 let connection = cx
398 .update(|cx| server.connect(current_dir.as_ref(), &project, cx))
399 .await
400 .unwrap();
401
402 let thread = connection
403 .new_thread(project.clone(), current_dir.as_ref(), &mut cx.to_async())
404 .await
405 .unwrap();
406
407 thread
408}
409
410pub async fn run_until_first_tool_call(
411 thread: &Entity<AcpThread>,
412 wait_until: impl Fn(&AgentThreadEntry) -> bool + 'static,
413 cx: &mut TestAppContext,
414) -> usize {
415 let (mut tx, mut rx) = mpsc::channel::<usize>(1);
416
417 let subscription = cx.update(|cx| {
418 cx.subscribe(thread, move |thread, _, cx| {
419 for (ix, entry) in thread.read(cx).entries().iter().enumerate() {
420 if wait_until(entry) {
421 return tx.try_send(ix).unwrap();
422 }
423 }
424 })
425 });
426
427 select! {
428 // We have to use a smol timer here because
429 // cx.background_executor().timer isn't real in the test context
430 _ = futures::FutureExt::fuse(smol::Timer::after(Duration::from_secs(20))) => {
431 panic!("Timeout waiting for tool call")
432 }
433 ix = rx.next().fuse() => {
434 drop(subscription);
435 ix.unwrap()
436 }
437 }
438}
439
440pub fn get_zed_path() -> PathBuf {
441 let mut zed_path = std::env::current_exe().unwrap();
442
443 while zed_path
444 .file_name()
445 .map_or(true, |name| name.to_string_lossy() != "debug")
446 {
447 if !zed_path.pop() {
448 panic!("Could not find target directory");
449 }
450 }
451
452 zed_path.push("zed");
453
454 if !zed_path.exists() {
455 panic!("\n🚨 Run `cargo build` at least once before running e2e tests\n\n");
456 }
457
458 zed_path
459}