Skip to repository content4401 lines · 145.2 KB · rust
tenant.openagents/omega
No repository description is available.
OpenAgents Git authority 2026-07-28T04:03:47.843Z Public web read
NIP-34 coordinate
30617:7649603503856e5148d571eac2766b288a8ff1e9e35d380337a1d2b0015b4f92:omegaMaintainersHidden in public view
References2 branches · 1 tag
Read-only clone
git clone https://openagents.com/git/tenant.openagents/omega.gitBrowse files
edit_prediction_tests.rs
1use client::{RefreshLlmTokenListener, UserStore, test::FakeServer};
2use clock::FakeSystemClock;
3use clock::ReplicaId;
4use cloud_api_types::{
5 CreateLlmTokenResponse, LlmToken, Organization, OrganizationConfiguration,
6 OrganizationEditPredictionConfiguration, OrganizationId, SettledEditPrediction,
7 SubmitEditPredictionSettledBatchBody, SubmitEditPredictionSettledResponse,
8};
9use cloud_llm_client::{
10 EditPredictionRejectReason, EditPredictionRejection, PredictEditsRequestTrigger,
11 RejectEditPredictionsBody,
12 predict_edits_v3::{PredictEditsV3Request, PredictEditsV3Response},
13 predict_edits_v4::{PredictEditsV4Request, PredictEditsV4Response},
14};
15use db::AppDatabase;
16use edit_prediction_types::EditPredictionRequestTrigger;
17use feature_flags::{FeatureFlag as _, FeatureFlagAppExt as _, FeatureFlagsSettings};
18use futures::{
19 AsyncReadExt, FutureExt, StreamExt,
20 channel::{mpsc, oneshot},
21};
22use gpui::App;
23use gpui::{
24 Entity, TestAppContext, UpdateGlobal,
25 http_client::{FakeHttpClient, Response},
26};
27use indoc::indoc;
28use language::{
29 Anchor, Buffer, Capability, Diagnostic, DiagnosticEntry, DiagnosticSet, DiagnosticSeverity,
30 Point, unified_diff_with_offsets,
31};
32use lsp::LanguageServerId;
33use parking_lot::Mutex;
34use pretty_assertions::{assert_eq, assert_matches};
35use project::{FakeFs, Project};
36use serde_json::json;
37use settings::EditPredictionDataCollectionChoice;
38use settings::SettingsStore;
39use std::{ops::Range, path::Path, sync::Arc, time::Duration};
40use util::{
41 path,
42 test::{TextRangeMarker, marked_text_ranges_by},
43};
44use uuid::Uuid;
45use workspace::{AppState, CollaboratorId, MultiWorkspace};
46use zeta_prompt::Zeta2PromptInput;
47
48use crate::prediction::EditPredictionInputs;
49use crate::udiff::apply_diff_to_string;
50use crate::{
51 BufferEditPrediction, EDIT_PREDICTION_SETTLED_QUIESCENCE, EditPredictionId,
52 EditPredictionStore, REJECT_REQUEST_DEBOUNCE, REQUEST_TIMEOUT_BACKOFF,
53};
54
55use super::*;
56
57#[gpui::test]
58async fn test_current_state(cx: &mut TestAppContext) {
59 let (ep_store, mut requests) = init_test_with_fake_client(cx);
60 cx.update(|cx| set_jumps_feature_flag_override(cx, "on"));
61 let fs = FakeFs::new(cx.executor());
62 fs.insert_tree(
63 "/root",
64 json!({
65 "1.txt": "Hello!\nHow\nBye\n",
66 "2.txt": "Hola!\nComo\nAdios\n"
67 }),
68 )
69 .await;
70 let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
71
72 let buffer1 = project
73 .update(cx, |project, cx| {
74 let path = project.find_project_path(path!("/root/1.txt"), cx).unwrap();
75 project.set_active_path(Some(path.clone()), cx);
76 project.open_buffer(path, cx)
77 })
78 .await
79 .unwrap();
80 let snapshot1 = buffer1.read_with(cx, |buffer, _cx| buffer.snapshot());
81 let position = snapshot1.anchor_before(language::Point::new(1, 3));
82
83 ep_store.update(cx, |ep_store, cx| {
84 ep_store.register_project(&project, cx);
85 ep_store.register_buffer(&buffer1, &project, cx);
86 });
87
88 // Prediction for current file
89
90 ep_store.update(cx, |ep_store, cx| {
91 ep_store.refresh_prediction_from_buffer(
92 project.clone(),
93 buffer1.clone(),
94 position,
95 EditPredictionRequestTrigger::Other,
96 cx,
97 )
98 });
99 let (_request, respond_tx) = requests.predict_v4.next().await.unwrap();
100
101 respond_tx
102 .send(PredictEditsV4Response {
103 request_id: Uuid::new_v4().to_string(),
104 patch: indoc! {r"
105 --- a/root/1.txt
106 +++ b/root/1.txt
107 @@ ... @@
108 Hello!
109 -How
110 +How are you?
111 Bye
112 "}
113 .to_string(),
114 model_version: None,
115 })
116 .unwrap();
117
118 cx.run_until_parked();
119
120 ep_store.update(cx, |ep_store, cx| {
121 let prediction = ep_store
122 .prediction_at(&buffer1, None, &project, cx)
123 .unwrap();
124 assert_matches!(prediction, BufferEditPrediction::Local { .. });
125 });
126
127 ep_store.update(cx, |ep_store, cx| {
128 ep_store.reject_current_prediction(EditPredictionRejectReason::Discarded, &project, cx);
129 });
130}
131
132#[gpui::test]
133async fn test_refresh_prediction_from_buffer_suppressed_while_following(cx: &mut TestAppContext) {
134 let (ep_store, mut requests) = init_test_with_fake_client(cx);
135 let fs = FakeFs::new(cx.executor());
136 fs.insert_tree(
137 "/root",
138 json!({
139 "foo.md": "Hello!\nHow\nBye\n"
140 }),
141 )
142 .await;
143 let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
144
145 let app_state = cx.update(|cx| {
146 let app_state = AppState::test(cx);
147 AppState::set_global(app_state.clone(), cx);
148 app_state
149 });
150 let multi_workspace =
151 cx.add_window(|window, cx| MultiWorkspace::test_new(project.clone(), window, cx));
152 let workspace = multi_workspace
153 .read_with(cx, |multi_workspace, _| multi_workspace.workspace().clone())
154 .unwrap();
155 cx.update(|cx| {
156 AppState::set_global(workspace.read(cx).app_state().clone(), cx);
157 });
158 drop(app_state);
159
160 let buffer = project
161 .update(cx, |project, cx| {
162 let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
163 project.open_buffer(path, cx)
164 })
165 .await
166 .unwrap();
167 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
168 let position = snapshot.anchor_before(language::Point::new(1, 3));
169
170 multi_workspace
171 .update(cx, |multi_workspace, window, cx| {
172 multi_workspace.workspace().update(cx, |workspace, cx| {
173 workspace.start_following(CollaboratorId::Agent, window, cx);
174 });
175 })
176 .unwrap();
177 cx.run_until_parked();
178
179 ep_store.update(cx, |ep_store, cx| {
180 ep_store.register_project(&project, cx);
181 ep_store.register_buffer(&buffer, &project, cx);
182 ep_store.refresh_prediction_from_buffer(
183 project.clone(),
184 buffer.clone(),
185 position,
186 EditPredictionRequestTrigger::Other,
187 cx,
188 );
189 });
190 cx.run_until_parked();
191
192 assert_no_predict_request_ready(&mut requests.predict);
193}
194
195#[gpui::test]
196async fn test_simple_request(cx: &mut TestAppContext) {
197 let (ep_store, mut requests) = init_test_with_fake_client(cx);
198 let fs = FakeFs::new(cx.executor());
199 fs.insert_tree(
200 "/root",
201 json!({
202 "foo.md": "Hello!\nHow\nBye\n"
203 }),
204 )
205 .await;
206 let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
207
208 let buffer = project
209 .update(cx, |project, cx| {
210 let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
211 project.open_buffer(path, cx)
212 })
213 .await
214 .unwrap();
215 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
216 let position = snapshot.anchor_before(language::Point::new(1, 3));
217
218 let prediction_task = ep_store.update(cx, |ep_store, cx| {
219 ep_store.request_prediction(
220 &project,
221 &buffer,
222 position,
223 PredictEditsRequestTrigger::Other,
224 cx,
225 )
226 });
227
228 let (request, respond_tx) = requests.predict.next().await.unwrap();
229
230 // TODO Put back when we have a structured request again
231 // assert_eq!(
232 // request.excerpt_path.as_ref(),
233 // Path::new(path!("root/foo.md"))
234 // );
235 // assert_eq!(
236 // request.cursor_point,
237 // Point {
238 // line: Line(1),
239 // column: 3
240 // }
241 // );
242
243 respond_tx
244 .send(model_response(
245 &request,
246 indoc! { r"
247 --- a/root/foo.md
248 +++ b/root/foo.md
249 @@ ... @@
250 Hello!
251 -How
252 +How are you?
253 Bye
254 "},
255 ))
256 .unwrap();
257
258 let prediction = prediction_task.await.unwrap().unwrap().prediction;
259
260 assert_eq!(prediction.edits.len(), 1);
261 assert_eq!(
262 prediction.edits[0].0.to_point(&snapshot).start,
263 language::Point::new(1, 3)
264 );
265 assert_eq!(prediction.edits[0].1.as_ref(), " are you?");
266}
267
268#[gpui::test]
269async fn test_zeta_request_sends_settled_body_when_data_collection_is_disabled(
270 cx: &mut TestAppContext,
271) {
272 let (ep_store, mut requests) = init_test_with_fake_client(cx);
273 let fs = FakeFs::new(cx.executor());
274 fs.insert_tree(
275 "/root",
276 json!({
277 "foo.md": "Hello!\nHow\nBye\n"
278 }),
279 )
280 .await;
281 let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
282
283 let buffer = project
284 .update(cx, |project, cx| {
285 let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
286 project.open_buffer(path, cx)
287 })
288 .await
289 .unwrap();
290 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
291 let position = snapshot.anchor_before(language::Point::new(1, 3));
292
293 ep_store.update(cx, |ep_store, cx| {
294 ep_store.register_buffer(&buffer, &project, cx);
295 });
296
297 let prediction_task = ep_store.update(cx, |ep_store, cx| {
298 ep_store.request_prediction(
299 &project,
300 &buffer,
301 position,
302 PredictEditsRequestTrigger::Other,
303 cx,
304 )
305 });
306
307 let (request, respond_tx) = requests.predict.next().await.unwrap();
308 assert!(!request.input.can_collect_data);
309 respond_tx
310 .send(model_response(&request, SIMPLE_DIFF))
311 .unwrap();
312
313 prediction_task.await.unwrap().unwrap();
314 cx.run_until_parked();
315 cx.executor()
316 .advance_clock(EDIT_PREDICTION_SETTLED_QUIESCENCE);
317 cx.run_until_parked();
318
319 let settled_request = requests
320 .settled
321 .next()
322 .now_or_never()
323 .flatten()
324 .expect("settled request should be sent");
325 assert!(!settled_request.can_collect_data);
326 assert_eq!(settled_request.settled_editable_region, None);
327 assert_eq!(settled_request.sample_data, None);
328}
329
330#[gpui::test]
331async fn test_request_events(cx: &mut TestAppContext) {
332 let (ep_store, mut requests) = init_test_with_fake_client(cx);
333 let fs = FakeFs::new(cx.executor());
334 fs.insert_tree(
335 "/root",
336 json!({
337 "foo.md": "Hello!\n\nBye\n"
338 }),
339 )
340 .await;
341 let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
342
343 let buffer = project
344 .update(cx, |project, cx| {
345 let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
346 project.open_buffer(path, cx)
347 })
348 .await
349 .unwrap();
350
351 ep_store.update(cx, |ep_store, cx| {
352 ep_store.register_buffer(&buffer, &project, cx);
353 });
354
355 buffer.update(cx, |buffer, cx| {
356 buffer.edit(vec![(7..7, "How")], None, cx);
357 });
358
359 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
360 let position = snapshot.anchor_before(language::Point::new(1, 3));
361
362 let prediction_task = ep_store.update(cx, |ep_store, cx| {
363 ep_store.request_prediction(
364 &project,
365 &buffer,
366 position,
367 PredictEditsRequestTrigger::Other,
368 cx,
369 )
370 });
371
372 let (request, respond_tx) = requests.predict.next().await.unwrap();
373
374 let prompt = prompt_from_request(&request);
375 assert!(
376 prompt.contains(indoc! {"
377 --- a/root/foo.md
378 +++ b/root/foo.md
379 @@ -1,3 +1,3 @@
380 Hello!
381 -
382 +How
383 Bye
384 "}),
385 "{prompt}"
386 );
387
388 respond_tx
389 .send(model_response(
390 &request,
391 indoc! {r#"
392 --- a/root/foo.md
393 +++ b/root/foo.md
394 @@ ... @@
395 Hello!
396 -How
397 +How are you?
398 Bye
399 "#},
400 ))
401 .unwrap();
402
403 let prediction = prediction_task.await.unwrap().unwrap().prediction;
404
405 assert_eq!(prediction.edits.len(), 1);
406 assert_eq!(prediction.edits[0].1.as_ref(), " are you?");
407}
408
409#[gpui::test]
410async fn test_edit_history_getter_pause_splits_last_event(cx: &mut TestAppContext) {
411 let (ep_store, _requests) = init_test_with_fake_client(cx);
412 let fs = FakeFs::new(cx.executor());
413 fs.insert_tree(
414 "/root",
415 json!({
416 "foo.md": "Hello!\n\nBye\n"
417 }),
418 )
419 .await;
420 let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
421
422 let buffer = project
423 .update(cx, |project, cx| {
424 let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
425 project.open_buffer(path, cx)
426 })
427 .await
428 .unwrap();
429
430 ep_store.update(cx, |ep_store, cx| {
431 ep_store.register_buffer(&buffer, &project, cx);
432 });
433
434 // First burst: insert "How"
435 buffer.update(cx, |buffer, cx| {
436 buffer.edit(vec![(7..7, "How")], None, cx);
437 });
438
439 // Simulate a pause longer than the grouping threshold (e.g. 500ms).
440 cx.executor().advance_clock(LAST_CHANGE_GROUPING_TIME * 2);
441 cx.run_until_parked();
442
443 // Second burst: append " are you?" immediately after "How" on the same line.
444 //
445 // Keeping both bursts on the same line ensures the existing line-span coalescing logic
446 // groups them into a single `LastEvent`, allowing the pause-split getter to return two diffs.
447 buffer.update(cx, |buffer, cx| {
448 buffer.edit(vec![(10..10, " are you?")], None, cx);
449 });
450
451 // A second edit shortly after the first post-pause edit ensures the last edit timestamp is
452 // advanced after the pause boundary is recorded, making pause-splitting deterministic.
453 buffer.update(cx, |buffer, cx| {
454 buffer.edit(vec![(19..19, "!")], None, cx);
455 });
456
457 // With time-based splitting, there are two distinct events.
458 let events = ep_store.update(cx, |ep_store, cx| {
459 ep_store.edit_history_for_project(&project, cx)
460 });
461 assert_eq!(events.len(), 2);
462
463 let first_total_edit_range = buffer.read_with(cx, |buffer, _| {
464 events[0].total_edit_range.to_point(&buffer.snapshot())
465 });
466 assert_eq!(first_total_edit_range, Point::new(1, 0)..Point::new(1, 3));
467
468 let zeta_prompt::Event::BufferChange { diff, .. } = events[0].event.as_ref();
469 assert_eq!(
470 diff.as_str(),
471 indoc! {"
472 @@ -1,3 +1,3 @@
473 Hello!
474 -
475 +How
476 Bye
477 "}
478 );
479
480 let second_total_edit_range = buffer.read_with(cx, |buffer, _| {
481 events[1].total_edit_range.to_point(&buffer.snapshot())
482 });
483 assert_eq!(second_total_edit_range, Point::new(1, 3)..Point::new(1, 13));
484
485 let zeta_prompt::Event::BufferChange { diff, .. } = events[1].event.as_ref();
486 assert_eq!(
487 diff.as_str(),
488 indoc! {"
489 @@ -1,3 +1,3 @@
490 Hello!
491 -How
492 +How are you?!
493 Bye
494 "}
495 );
496}
497
498#[gpui::test]
499async fn test_predicted_edits_are_separated_in_edit_history(cx: &mut TestAppContext) {
500 let (ep_store, _requests) = init_test_with_fake_client(cx);
501 let fs = FakeFs::new(cx.executor());
502
503 // Create a file with 30 lines to test line-based coalescing
504 let content = (1..=30)
505 .map(|i| format!("Line {}\n", i))
506 .collect::<String>();
507 fs.insert_tree(
508 "/root",
509 json!({
510 "foo.md": content
511 }),
512 )
513 .await;
514 let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
515
516 let buffer = project
517 .update(cx, |project, cx| {
518 let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
519 project.open_buffer(path, cx)
520 })
521 .await
522 .unwrap();
523
524 ep_store.update(cx, |ep_store, cx| {
525 ep_store.register_buffer(&buffer, &project, cx);
526 });
527
528 // First edit: multi-line edit spanning rows 10-12 (replacing lines 11-13)
529 buffer.update(cx, |buffer, cx| {
530 let start = Point::new(10, 0).to_offset(buffer);
531 let end = Point::new(13, 0).to_offset(buffer);
532 buffer.edit(vec![(start..end, "Middle A\nMiddle B\n")], None, cx);
533 });
534
535 let events = ep_store.update(cx, |ep_store, cx| {
536 ep_store.edit_history_for_project(&project, cx)
537 });
538 assert_eq!(
539 render_events(&events),
540 indoc! {"
541 @@ -8,9 +8,8 @@
542 Line 8
543 Line 9
544 Line 10
545 -Line 11
546 -Line 12
547 -Line 13
548 +Middle A
549 +Middle B
550 Line 14
551 Line 15
552 Line 16
553 "},
554 "After first edit"
555 );
556
557 // Second edit: insert ABOVE the first edit's range (row 5, within 8 lines of row 10)
558 // This tests that coalescing considers the START of the existing range
559 buffer.update(cx, |buffer, cx| {
560 let offset = Point::new(5, 0).to_offset(buffer);
561 buffer.edit(vec![(offset..offset, "Above\n")], None, cx);
562 });
563
564 let events = ep_store.update(cx, |ep_store, cx| {
565 ep_store.edit_history_for_project(&project, cx)
566 });
567 assert_eq!(
568 render_events(&events),
569 indoc! {"
570 @@ -3,14 +3,14 @@
571 Line 3
572 Line 4
573 Line 5
574 +Above
575 Line 6
576 Line 7
577 Line 8
578 Line 9
579 Line 10
580 -Line 11
581 -Line 12
582 -Line 13
583 +Middle A
584 +Middle B
585 Line 14
586 Line 15
587 Line 16
588 "},
589 "After inserting above (should coalesce)"
590 );
591
592 // Third edit: insert BELOW the first edit's range (row 14 in current buffer, within 8 lines of row 12)
593 // This tests that coalescing considers the END of the existing range
594 buffer.update(cx, |buffer, cx| {
595 let offset = Point::new(14, 0).to_offset(buffer);
596 buffer.edit(vec![(offset..offset, "Below\n")], None, cx);
597 });
598
599 let events = ep_store.update(cx, |ep_store, cx| {
600 ep_store.edit_history_for_project(&project, cx)
601 });
602 assert_eq!(
603 render_events(&events),
604 indoc! {"
605 @@ -3,15 +3,16 @@
606 Line 3
607 Line 4
608 Line 5
609 +Above
610 Line 6
611 Line 7
612 Line 8
613 Line 9
614 Line 10
615 -Line 11
616 -Line 12
617 -Line 13
618 +Middle A
619 +Middle B
620 Line 14
621 +Below
622 Line 15
623 Line 16
624 Line 17
625 "},
626 "After inserting below (should coalesce)"
627 );
628
629 // Fourth edit: insert FAR BELOW (row 25, beyond 8 lines from the current range end ~row 15)
630 // This should NOT coalesce - creates a new event
631 buffer.update(cx, |buffer, cx| {
632 let offset = Point::new(25, 0).to_offset(buffer);
633 buffer.edit(vec![(offset..offset, "Far below\n")], None, cx);
634 });
635
636 let events = ep_store.update(cx, |ep_store, cx| {
637 ep_store.edit_history_for_project(&project, cx)
638 });
639 assert_eq!(
640 render_events(&events),
641 indoc! {"
642 @@ -3,15 +3,16 @@
643 Line 3
644 Line 4
645 Line 5
646 +Above
647 Line 6
648 Line 7
649 Line 8
650 Line 9
651 Line 10
652 -Line 11
653 -Line 12
654 -Line 13
655 +Middle A
656 +Middle B
657 Line 14
658 +Below
659 Line 15
660 Line 16
661 Line 17
662
663 ---
664 @@ -23,6 +23,7 @@
665 Line 22
666 Line 23
667 Line 24
668 +Far below
669 Line 25
670 Line 26
671 Line 27
672 "},
673 "After inserting far below (should NOT coalesce)"
674 );
675}
676
677fn render_events(events: &[StoredEvent]) -> String {
678 events
679 .iter()
680 .map(|e| {
681 let zeta_prompt::Event::BufferChange { diff, .. } = e.event.as_ref();
682 diff.as_str()
683 })
684 .collect::<Vec<_>>()
685 .join("\n---\n")
686}
687
688fn render_events_with_predicted(events: &[StoredEvent]) -> Vec<String> {
689 events
690 .iter()
691 .map(|e| {
692 let zeta_prompt::Event::BufferChange {
693 diff, predicted, ..
694 } = e.event.as_ref();
695 let prefix = if *predicted { "predicted" } else { "manual" };
696 format!("{}\n{}", prefix, diff)
697 })
698 .collect()
699}
700
701fn make_collaborator_replica(
702 buffer: &Entity<Buffer>,
703 cx: &mut TestAppContext,
704) -> (Entity<Buffer>, clock::Global) {
705 let (state, version) =
706 buffer.read_with(cx, |buffer, _cx| (buffer.to_proto(_cx), buffer.version()));
707 let collaborator = cx.new(|_cx| {
708 Buffer::from_proto(ReplicaId::new(1), Capability::ReadWrite, state, None).unwrap()
709 });
710 (collaborator, version)
711}
712
713async fn apply_collaborator_edit(
714 collaborator: &Entity<Buffer>,
715 buffer: &Entity<Buffer>,
716 since_version: &mut clock::Global,
717 edit_range: Range<usize>,
718 new_text: &str,
719 cx: &mut TestAppContext,
720) {
721 collaborator.update(cx, |collaborator, cx| {
722 collaborator.edit([(edit_range, new_text)], None, cx);
723 });
724
725 let serialize_task = collaborator.read_with(cx, |collaborator, cx| {
726 collaborator.serialize_ops(Some(since_version.clone()), cx)
727 });
728 let ops = serialize_task.await;
729 *since_version = collaborator.read_with(cx, |collaborator, _cx| collaborator.version());
730
731 buffer.update(cx, |buffer, cx| {
732 buffer.apply_ops(
733 ops.into_iter()
734 .map(|op| language::proto::deserialize_operation(op).unwrap()),
735 cx,
736 );
737 });
738}
739
740#[gpui::test]
741async fn test_nearby_collaborator_edits_are_kept_in_history(cx: &mut TestAppContext) {
742 let (ep_store, _requests) = init_test_with_fake_client(cx);
743 let fs = FakeFs::new(cx.executor());
744 fs.insert_tree(
745 "/root",
746 json!({
747 "foo.rs": "line 0\nline 1\nline 2\nline 3\nline 4\nline 5\nline 6\nline 7\nline 8\nline 9\nline 10\nline 11\nline 12\nline 13\nline 14\n"
748 }),
749 )
750 .await;
751 let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
752
753 let buffer = project
754 .update(cx, |project, cx| {
755 let path = project.find_project_path(path!("root/foo.rs"), cx).unwrap();
756 project.set_active_path(Some(path.clone()), cx);
757 project.open_buffer(path, cx)
758 })
759 .await
760 .unwrap();
761
762 let cursor = buffer.read_with(cx, |buffer, _cx| buffer.anchor_before(Point::new(1, 0)));
763
764 ep_store.update(cx, |ep_store, cx| {
765 ep_store.register_buffer(&buffer, &project, cx);
766 let _ = ep_store.prediction_at(&buffer, Some(cursor), &project, cx);
767 });
768
769 buffer.update(cx, |buffer, cx| {
770 buffer.edit(vec![(0..6, "LOCAL ZERO")], None, cx);
771 });
772
773 let (collaborator, mut collaborator_version) = make_collaborator_replica(&buffer, cx);
774
775 let (line_one_start, line_one_len) = collaborator.read_with(cx, |buffer, _cx| {
776 (Point::new(1, 0).to_offset(buffer), buffer.line_len(1))
777 });
778
779 apply_collaborator_edit(
780 &collaborator,
781 &buffer,
782 &mut collaborator_version,
783 line_one_start..line_one_start + line_one_len as usize,
784 "REMOTE ONE",
785 cx,
786 )
787 .await;
788
789 let events = ep_store.update(cx, |ep_store, cx| {
790 ep_store.edit_history_for_project(&project, cx)
791 });
792
793 assert_eq!(
794 render_events_with_predicted(&events),
795 vec![indoc! {"
796 manual
797 @@ -1,5 +1,5 @@
798 -line 0
799 -line 1
800 +LOCAL ZERO
801 +REMOTE ONE
802 line 2
803 line 3
804 line 4
805 "}]
806 );
807}
808
809#[gpui::test]
810async fn test_distant_collaborator_edits_are_omitted_from_history(cx: &mut TestAppContext) {
811 let (ep_store, _requests) = init_test_with_fake_client(cx);
812 let fs = FakeFs::new(cx.executor());
813 fs.insert_tree(
814 "/root",
815 json!({
816 "foo.rs": (0..1000)
817 .map(|i| format!("line {i}\n"))
818 .collect::<String>()
819 }),
820 )
821 .await;
822 let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
823
824 let buffer = project
825 .update(cx, |project, cx| {
826 let path = project.find_project_path(path!("root/foo.rs"), cx).unwrap();
827 project.set_active_path(Some(path.clone()), cx);
828 project.open_buffer(path, cx)
829 })
830 .await
831 .unwrap();
832
833 let cursor = buffer.read_with(cx, |buffer, _cx| buffer.anchor_before(Point::new(1, 0)));
834
835 ep_store.update(cx, |ep_store, cx| {
836 ep_store.register_buffer(&buffer, &project, cx);
837 let _ = ep_store.prediction_at(&buffer, Some(cursor), &project, cx);
838 });
839
840 buffer.update(cx, |buffer, cx| {
841 buffer.edit(vec![(0..6, "LOCAL ZERO")], None, cx);
842 });
843
844 let (collaborator, mut collaborator_version) = make_collaborator_replica(&buffer, cx);
845
846 let far_line_start = buffer.read_with(cx, |buffer, _cx| Point::new(900, 0).to_offset(buffer));
847
848 apply_collaborator_edit(
849 &collaborator,
850 &buffer,
851 &mut collaborator_version,
852 far_line_start..far_line_start + 7,
853 "REMOTE FAR",
854 cx,
855 )
856 .await;
857
858 let events = ep_store.update(cx, |ep_store, cx| {
859 ep_store.edit_history_for_project(&project, cx)
860 });
861
862 assert_eq!(
863 render_events_with_predicted(&events),
864 vec![indoc! {"
865 manual
866 @@ -1,4 +1,4 @@
867 -line 0
868 +LOCAL ZERO
869 line 1
870 line 2
871 line 3
872 "}]
873 );
874}
875
876#[gpui::test]
877async fn test_irrelevant_collaborator_edits_in_different_files_are_omitted_from_history(
878 cx: &mut TestAppContext,
879) {
880 let (ep_store, _requests) = init_test_with_fake_client(cx);
881 let fs = FakeFs::new(cx.executor());
882 fs.insert_tree(
883 "/root",
884 json!({
885 "foo.rs": "line 0\nline 1\nline 2\nline 3\n",
886 "bar.rs": "line 0\nline 1\nline 2\nline 3\n"
887 }),
888 )
889 .await;
890 let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
891
892 let foo_buffer = project
893 .update(cx, |project, cx| {
894 let path = project.find_project_path(path!("root/foo.rs"), cx).unwrap();
895 project.set_active_path(Some(path.clone()), cx);
896 project.open_buffer(path, cx)
897 })
898 .await
899 .unwrap();
900 let bar_buffer = project
901 .update(cx, |project, cx| {
902 let path = project.find_project_path(path!("root/bar.rs"), cx).unwrap();
903 project.open_buffer(path, cx)
904 })
905 .await
906 .unwrap();
907
908 let foo_cursor = foo_buffer.read_with(cx, |buffer, _cx| buffer.anchor_before(Point::new(1, 0)));
909
910 ep_store.update(cx, |ep_store, cx| {
911 ep_store.register_buffer(&foo_buffer, &project, cx);
912 ep_store.register_buffer(&bar_buffer, &project, cx);
913 let _ = ep_store.prediction_at(&foo_buffer, Some(foo_cursor), &project, cx);
914 });
915
916 let (bar_collaborator, mut bar_version) = make_collaborator_replica(&bar_buffer, cx);
917
918 apply_collaborator_edit(
919 &bar_collaborator,
920 &bar_buffer,
921 &mut bar_version,
922 0..6,
923 "REMOTE BAR",
924 cx,
925 )
926 .await;
927
928 let events = ep_store.update(cx, |ep_store, cx| {
929 ep_store.edit_history_for_project(&project, cx)
930 });
931
932 assert!(events.is_empty());
933}
934
935#[gpui::test]
936async fn test_large_edits_are_omitted_from_history(cx: &mut TestAppContext) {
937 let (ep_store, _requests) = init_test_with_fake_client(cx);
938 let fs = FakeFs::new(cx.executor());
939 fs.insert_tree(
940 "/root",
941 json!({
942 "foo.rs": (0..20)
943 .map(|i| format!("line {i}\n"))
944 .collect::<String>()
945 }),
946 )
947 .await;
948 let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
949
950 let buffer = project
951 .update(cx, |project, cx| {
952 let path = project.find_project_path(path!("root/foo.rs"), cx).unwrap();
953 project.set_active_path(Some(path.clone()), cx);
954 project.open_buffer(path, cx)
955 })
956 .await
957 .unwrap();
958
959 let cursor = buffer.read_with(cx, |buffer, _cx| buffer.anchor_before(Point::new(1, 0)));
960
961 ep_store.update(cx, |ep_store, cx| {
962 ep_store.register_buffer(&buffer, &project, cx);
963 let _ = ep_store.prediction_at(&buffer, Some(cursor), &project, cx);
964 });
965
966 buffer.update(cx, |buffer, cx| {
967 buffer.edit(vec![(0..6, "LOCAL ZERO")], None, cx);
968 });
969
970 let (collaborator, mut collaborator_version) = make_collaborator_replica(&buffer, cx);
971
972 let (line_three_start, line_three_len) = collaborator.read_with(cx, |buffer, _cx| {
973 (Point::new(3, 0).to_offset(buffer), buffer.line_len(3))
974 });
975 let large_edit = "X".repeat(EDIT_HISTORY_DIFF_SIZE_LIMIT + 1);
976
977 apply_collaborator_edit(
978 &collaborator,
979 &buffer,
980 &mut collaborator_version,
981 line_three_start..line_three_start + line_three_len as usize,
982 &large_edit,
983 cx,
984 )
985 .await;
986
987 buffer.update(cx, |buffer, cx| {
988 let line_seven_start = Point::new(7, 0).to_offset(buffer);
989 let line_seven_end = Point::new(7, 6).to_offset(buffer);
990 buffer.edit(
991 vec![(line_seven_start..line_seven_end, "LOCAL SEVEN")],
992 None,
993 cx,
994 );
995 });
996
997 let events = ep_store.update(cx, |ep_store, cx| {
998 ep_store.edit_history_for_project(&project, cx)
999 });
1000
1001 let rendered_events = render_events_with_predicted(&events);
1002
1003 assert_eq!(rendered_events.len(), 2);
1004 assert!(rendered_events[0].contains("+LOCAL ZERO"));
1005 assert!(!rendered_events[0].contains(&large_edit));
1006 assert!(rendered_events[1].contains("+LOCAL SEVEN"));
1007 assert!(!rendered_events[1].contains(&large_edit));
1008}
1009
1010#[gpui::test]
1011async fn test_predicted_flag_coalescing(cx: &mut TestAppContext) {
1012 let (ep_store, _requests) = init_test_with_fake_client(cx);
1013 let fs = FakeFs::new(cx.executor());
1014 fs.insert_tree(
1015 "/root",
1016 json!({
1017 "foo.rs": "line 0\nline 1\nline 2\nline 3\nline 4\nline 5\nline 6\nline 7\nline 8\nline 9\nline 10\nline 11\nline 12\nline 13\nline 14\n"
1018 }),
1019 )
1020 .await;
1021 let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
1022
1023 let buffer = project
1024 .update(cx, |project, cx| {
1025 let path = project.find_project_path(path!("root/foo.rs"), cx).unwrap();
1026 project.open_buffer(path, cx)
1027 })
1028 .await
1029 .unwrap();
1030
1031 ep_store.update(cx, |ep_store, cx| {
1032 ep_store.register_buffer(&buffer, &project, cx);
1033 });
1034
1035 // Case 1: Manual edits have `predicted` set to false.
1036 buffer.update(cx, |buffer, cx| {
1037 buffer.edit(vec![(0..6, "LINE ZERO")], None, cx);
1038 });
1039
1040 let events = ep_store.update(cx, |ep_store, cx| {
1041 ep_store.edit_history_for_project(&project, cx)
1042 });
1043
1044 assert_eq!(
1045 render_events_with_predicted(&events),
1046 vec![indoc! {"
1047 manual
1048 @@ -1,4 +1,4 @@
1049 -line 0
1050 +LINE ZERO
1051 line 1
1052 line 2
1053 line 3
1054 "}]
1055 );
1056
1057 // Case 2: Multiple successive manual edits near each other are merged into one
1058 // event with `predicted` set to false.
1059 buffer.update(cx, |buffer, cx| {
1060 let offset = Point::new(1, 0).to_offset(buffer);
1061 let end = Point::new(1, 6).to_offset(buffer);
1062 buffer.edit(vec![(offset..end, "LINE ONE")], None, cx);
1063 });
1064
1065 let events = ep_store.update(cx, |ep_store, cx| {
1066 ep_store.edit_history_for_project(&project, cx)
1067 });
1068 assert_eq!(
1069 render_events_with_predicted(&events),
1070 vec![indoc! {"
1071 manual
1072 @@ -1,5 +1,5 @@
1073 -line 0
1074 -line 1
1075 +LINE ZERO
1076 +LINE ONE
1077 line 2
1078 line 3
1079 line 4
1080 "}]
1081 );
1082
1083 // Case 3: Accepted predictions have `predicted` set to true.
1084 // Case 5: A manual edit that follows a predicted edit is not merged with the
1085 // predicted edit, even if it is nearby.
1086 ep_store.update(cx, |ep_store, cx| {
1087 buffer.update(cx, |buffer, cx| {
1088 let offset = Point::new(2, 0).to_offset(buffer);
1089 let end = Point::new(2, 6).to_offset(buffer);
1090 buffer.edit(vec![(offset..end, "LINE TWO")], None, cx);
1091 });
1092 ep_store.report_changes_for_buffer(&buffer, &project, true, true, cx);
1093 });
1094
1095 let events = ep_store.update(cx, |ep_store, cx| {
1096 ep_store.edit_history_for_project(&project, cx)
1097 });
1098 assert_eq!(
1099 render_events_with_predicted(&events),
1100 vec![
1101 indoc! {"
1102 manual
1103 @@ -1,5 +1,5 @@
1104 -line 0
1105 -line 1
1106 +LINE ZERO
1107 +LINE ONE
1108 line 2
1109 line 3
1110 line 4
1111 "},
1112 indoc! {"
1113 predicted
1114 @@ -1,6 +1,6 @@
1115 LINE ZERO
1116 LINE ONE
1117 -line 2
1118 +LINE TWO
1119 line 3
1120 line 4
1121 line 5
1122 "}
1123 ]
1124 );
1125
1126 // Case 4: Multiple successive accepted predictions near each other are merged
1127 // into one event with `predicted` set to true.
1128 ep_store.update(cx, |ep_store, cx| {
1129 buffer.update(cx, |buffer, cx| {
1130 let offset = Point::new(3, 0).to_offset(buffer);
1131 let end = Point::new(3, 6).to_offset(buffer);
1132 buffer.edit(vec![(offset..end, "LINE THREE")], None, cx);
1133 });
1134 ep_store.report_changes_for_buffer(&buffer, &project, true, true, cx);
1135 });
1136
1137 let events = ep_store.update(cx, |ep_store, cx| {
1138 ep_store.edit_history_for_project(&project, cx)
1139 });
1140 assert_eq!(
1141 render_events_with_predicted(&events),
1142 vec![
1143 indoc! {"
1144 manual
1145 @@ -1,5 +1,5 @@
1146 -line 0
1147 -line 1
1148 +LINE ZERO
1149 +LINE ONE
1150 line 2
1151 line 3
1152 line 4
1153 "},
1154 indoc! {"
1155 predicted
1156 @@ -1,7 +1,7 @@
1157 LINE ZERO
1158 LINE ONE
1159 -line 2
1160 -line 3
1161 +LINE TWO
1162 +LINE THREE
1163 line 4
1164 line 5
1165 line 6
1166 "}
1167 ]
1168 );
1169
1170 // Case 5 (continued): A manual edit that follows a predicted edit is not merged
1171 // with the predicted edit, even if it is nearby.
1172 buffer.update(cx, |buffer, cx| {
1173 let offset = Point::new(4, 0).to_offset(buffer);
1174 let end = Point::new(4, 6).to_offset(buffer);
1175 buffer.edit(vec![(offset..end, "LINE FOUR")], None, cx);
1176 });
1177
1178 let events = ep_store.update(cx, |ep_store, cx| {
1179 ep_store.edit_history_for_project(&project, cx)
1180 });
1181 assert_eq!(
1182 render_events_with_predicted(&events),
1183 vec![
1184 indoc! {"
1185 manual
1186 @@ -1,5 +1,5 @@
1187 -line 0
1188 -line 1
1189 +LINE ZERO
1190 +LINE ONE
1191 line 2
1192 line 3
1193 line 4
1194 "},
1195 indoc! {"
1196 predicted
1197 @@ -1,7 +1,7 @@
1198 LINE ZERO
1199 LINE ONE
1200 -line 2
1201 -line 3
1202 +LINE TWO
1203 +LINE THREE
1204 line 4
1205 line 5
1206 line 6
1207 "},
1208 indoc! {"
1209 manual
1210 @@ -2,7 +2,7 @@
1211 LINE ONE
1212 LINE TWO
1213 LINE THREE
1214 -line 4
1215 +LINE FOUR
1216 line 5
1217 line 6
1218 line 7
1219 "}
1220 ]
1221 );
1222
1223 // Case 6: If we then perform a manual edit at a *different* location (more than
1224 // 8 lines away), then the edits at the prior location can be merged with each
1225 // other, even if some are predicted and some are not. `predicted` means all
1226 // constituent edits were predicted.
1227 buffer.update(cx, |buffer, cx| {
1228 let offset = Point::new(14, 0).to_offset(buffer);
1229 let end = Point::new(14, 7).to_offset(buffer);
1230 buffer.edit(vec![(offset..end, "LINE FOURTEEN")], None, cx);
1231 });
1232
1233 let events = ep_store.update(cx, |ep_store, cx| {
1234 ep_store.edit_history_for_project(&project, cx)
1235 });
1236 assert_eq!(
1237 render_events_with_predicted(&events),
1238 vec![
1239 indoc! {"
1240 manual
1241 @@ -1,8 +1,8 @@
1242 -line 0
1243 -line 1
1244 -line 2
1245 -line 3
1246 -line 4
1247 +LINE ZERO
1248 +LINE ONE
1249 +LINE TWO
1250 +LINE THREE
1251 +LINE FOUR
1252 line 5
1253 line 6
1254 line 7
1255 "},
1256 indoc! {"
1257 manual
1258 @@ -12,4 +12,4 @@
1259 line 11
1260 line 12
1261 line 13
1262 -line 14
1263 +LINE FOURTEEN
1264 "}
1265 ]
1266 );
1267}
1268
1269#[gpui::test]
1270async fn test_empty_prediction(cx: &mut TestAppContext) {
1271 let (ep_store, mut requests) = init_test_with_fake_client(cx);
1272 let fs = FakeFs::new(cx.executor());
1273 fs.insert_tree(
1274 "/root",
1275 json!({
1276 "foo.md": "Hello!\nHow\nBye\n"
1277 }),
1278 )
1279 .await;
1280 let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
1281
1282 let buffer = project
1283 .update(cx, |project, cx| {
1284 let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
1285 project.open_buffer(path, cx)
1286 })
1287 .await
1288 .unwrap();
1289 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
1290 let position = snapshot.anchor_before(language::Point::new(1, 3));
1291
1292 ep_store.update(cx, |ep_store, cx| {
1293 ep_store.register_buffer(&buffer, &project, cx);
1294 ep_store.refresh_prediction_from_buffer(
1295 project.clone(),
1296 buffer.clone(),
1297 position,
1298 EditPredictionRequestTrigger::Explicit,
1299 cx,
1300 );
1301 });
1302
1303 let (request, respond_tx) = requests.predict.next().await.unwrap();
1304 let mut response = model_response(&request, "");
1305 response.model_version = Some("zeta2:test-empty".to_string());
1306 let id = response.request_id.clone();
1307 respond_tx.send(response).unwrap();
1308
1309 cx.run_until_parked();
1310
1311 ep_store.update(cx, |ep_store, cx| {
1312 assert!(
1313 ep_store
1314 .prediction_at(&buffer, None, &project, cx)
1315 .is_none()
1316 );
1317 let shown_predictions = ep_store.rateable_predictions().collect::<Vec<_>>();
1318 assert_eq!(shown_predictions.len(), 1);
1319 assert_eq!(shown_predictions[0].id.to_string(), id);
1320 assert!(shown_predictions[0].edits.is_empty());
1321 assert!(shown_predictions[0].editable_range.is_some());
1322 assert!(matches!(
1323 shown_predictions[0].trigger,
1324 PredictEditsRequestTrigger::Explicit
1325 ));
1326 });
1327
1328 // prediction is reported as rejected
1329 let (reject_request, _) = requests.reject.next().await.unwrap();
1330
1331 assert_eq!(
1332 &reject_request.rejections,
1333 &[EditPredictionRejection {
1334 request_id: id.clone(),
1335 reason: EditPredictionRejectReason::Empty,
1336 was_shown: false,
1337 model_version: Some("zeta2:test-empty".to_string()),
1338 e2e_latency_ms: Some(0),
1339 }]
1340 );
1341 cx.executor()
1342 .advance_clock(EDIT_PREDICTION_SETTLED_QUIESCENCE);
1343 cx.run_until_parked();
1344
1345 let settled_request = requests
1346 .settled
1347 .next()
1348 .now_or_never()
1349 .flatten()
1350 .expect("empty prediction should still send settled request");
1351 assert_eq!(settled_request.request_id, id);
1352 assert_eq!(settled_request.settled_editable_region, None);
1353 assert_eq!(settled_request.sample_data, None);
1354}
1355
1356#[gpui::test]
1357async fn test_interpolated_empty(cx: &mut TestAppContext) {
1358 let (ep_store, mut requests) = init_test_with_fake_client(cx);
1359 let fs = FakeFs::new(cx.executor());
1360 fs.insert_tree(
1361 "/root",
1362 json!({
1363 "foo.md": "Hello!\nHow\nBye\n"
1364 }),
1365 )
1366 .await;
1367 let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
1368
1369 let buffer = project
1370 .update(cx, |project, cx| {
1371 let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
1372 project.open_buffer(path, cx)
1373 })
1374 .await
1375 .unwrap();
1376 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
1377 let position = snapshot.anchor_before(language::Point::new(1, 3));
1378
1379 ep_store.update(cx, |ep_store, cx| {
1380 ep_store.refresh_prediction_from_buffer(
1381 project.clone(),
1382 buffer.clone(),
1383 position,
1384 EditPredictionRequestTrigger::Other,
1385 cx,
1386 );
1387 });
1388
1389 let (request, respond_tx) = requests.predict.next().await.unwrap();
1390
1391 buffer.update(cx, |buffer, cx| {
1392 buffer.edit([(10..10, " are you?")], None, cx);
1393 });
1394
1395 let mut response = model_response(&request, SIMPLE_DIFF);
1396 response.model_version = Some("zeta2:test-interpolated-empty".to_string());
1397 let id = response.request_id.clone();
1398 respond_tx.send(response).unwrap();
1399
1400 cx.run_until_parked();
1401
1402 ep_store.update(cx, |ep_store, cx| {
1403 assert!(
1404 ep_store
1405 .prediction_at(&buffer, None, &project, cx)
1406 .is_none()
1407 );
1408 let shown_predictions = ep_store.rateable_predictions().collect::<Vec<_>>();
1409 assert_eq!(shown_predictions.len(), 1);
1410 assert_eq!(shown_predictions[0].id.to_string(), id);
1411 assert!(shown_predictions[0].edits.is_empty());
1412 assert!(shown_predictions[0].editable_range.is_some());
1413 });
1414
1415 let (reject_request, _) = requests.reject.next().await.unwrap();
1416
1417 assert_eq!(
1418 &reject_request.rejections,
1419 &[EditPredictionRejection {
1420 request_id: id,
1421 reason: EditPredictionRejectReason::InterpolatedEmpty,
1422 was_shown: false,
1423 model_version: Some("zeta2:test-interpolated-empty".to_string()),
1424 e2e_latency_ms: Some(0),
1425 }]
1426 );
1427}
1428
1429#[gpui::test]
1430async fn test_interpolate_failed(cx: &mut TestAppContext) {
1431 let (ep_store, mut requests) = init_test_with_fake_client(cx);
1432 let fs = FakeFs::new(cx.executor());
1433 fs.insert_tree(
1434 "/root",
1435 json!({
1436 "foo.md": "Hello!\nHow\nBye\n"
1437 }),
1438 )
1439 .await;
1440 let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
1441
1442 let buffer = project
1443 .update(cx, |project, cx| {
1444 let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
1445 project.open_buffer(path, cx)
1446 })
1447 .await
1448 .unwrap();
1449 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
1450 let position = snapshot.anchor_before(language::Point::new(1, 3));
1451
1452 ep_store.update(cx, |ep_store, cx| {
1453 ep_store.refresh_prediction_from_buffer(
1454 project.clone(),
1455 buffer.clone(),
1456 position,
1457 EditPredictionRequestTrigger::Other,
1458 cx,
1459 );
1460 });
1461
1462 let (request, respond_tx) = requests.predict.next().await.unwrap();
1463
1464 buffer.update(cx, |buffer, cx| {
1465 buffer.edit([(10..10, " is it?")], None, cx);
1466 });
1467
1468 let mut response = model_response(&request, SIMPLE_DIFF);
1469 response.model_version = Some("zeta2:test-interpolate-failed".to_string());
1470 let id = response.request_id.clone();
1471 respond_tx.send(response).unwrap();
1472
1473 cx.run_until_parked();
1474
1475 ep_store.update(cx, |ep_store, cx| {
1476 assert!(
1477 ep_store
1478 .prediction_at(&buffer, None, &project, cx)
1479 .is_none()
1480 );
1481 assert!(ep_store.rateable_predictions().next().is_none());
1482 });
1483
1484 let (reject_request, _) = requests.reject.next().await.unwrap();
1485
1486 assert_eq!(
1487 &reject_request.rejections,
1488 &[EditPredictionRejection {
1489 request_id: id,
1490 reason: EditPredictionRejectReason::InterpolateFailed,
1491 was_shown: false,
1492 model_version: Some("zeta2:test-interpolate-failed".to_string()),
1493 e2e_latency_ms: Some(0),
1494 }]
1495 );
1496}
1497
1498const SIMPLE_DIFF: &str = indoc! { r"
1499 --- a/root/foo.md
1500 +++ b/root/foo.md
1501 @@ ... @@
1502 Hello!
1503 -How
1504 +How are you?
1505 Bye
1506"};
1507
1508#[gpui::test]
1509async fn test_replace_current(cx: &mut TestAppContext) {
1510 let (ep_store, mut requests) = init_test_with_fake_client(cx);
1511 let fs = FakeFs::new(cx.executor());
1512 fs.insert_tree(
1513 "/root",
1514 json!({
1515 "foo.md": "Hello!\nHow\nBye\n"
1516 }),
1517 )
1518 .await;
1519 let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
1520
1521 let buffer = project
1522 .update(cx, |project, cx| {
1523 let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
1524 project.open_buffer(path, cx)
1525 })
1526 .await
1527 .unwrap();
1528 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
1529 let position = snapshot.anchor_before(language::Point::new(1, 3));
1530
1531 ep_store.update(cx, |ep_store, cx| {
1532 ep_store.refresh_prediction_from_buffer(
1533 project.clone(),
1534 buffer.clone(),
1535 position,
1536 EditPredictionRequestTrigger::Other,
1537 cx,
1538 );
1539 });
1540
1541 let (request, respond_tx) = requests.predict.next().await.unwrap();
1542 let first_response = model_response(&request, SIMPLE_DIFF);
1543 let first_id = first_response.request_id.clone();
1544 respond_tx.send(first_response).unwrap();
1545
1546 cx.run_until_parked();
1547
1548 ep_store.update(cx, |ep_store, cx| {
1549 assert_eq!(
1550 ep_store
1551 .prediction_at(&buffer, None, &project, cx)
1552 .unwrap()
1553 .id
1554 .0,
1555 first_id
1556 );
1557 });
1558
1559 // a second request is triggered
1560 ep_store.update(cx, |ep_store, cx| {
1561 ep_store.refresh_prediction_from_buffer(
1562 project.clone(),
1563 buffer.clone(),
1564 position,
1565 EditPredictionRequestTrigger::Other,
1566 cx,
1567 );
1568 });
1569
1570 let (request, respond_tx) = requests.predict.next().await.unwrap();
1571 let second_response = model_response(&request, SIMPLE_DIFF);
1572 let second_id = second_response.request_id.clone();
1573 respond_tx.send(second_response).unwrap();
1574
1575 cx.run_until_parked();
1576
1577 ep_store.update(cx, |ep_store, cx| {
1578 // second replaces first
1579 assert_eq!(
1580 ep_store
1581 .prediction_at(&buffer, None, &project, cx)
1582 .unwrap()
1583 .id
1584 .0,
1585 second_id
1586 );
1587 });
1588
1589 // first is reported as replaced
1590 let (reject_request, _) = requests.reject.next().await.unwrap();
1591
1592 assert_eq!(
1593 &reject_request.rejections,
1594 &[EditPredictionRejection {
1595 request_id: first_id,
1596 reason: EditPredictionRejectReason::Replaced,
1597 was_shown: false,
1598 model_version: None,
1599 e2e_latency_ms: Some(0),
1600 }]
1601 );
1602}
1603
1604#[gpui::test]
1605async fn test_current_preferred(cx: &mut TestAppContext) {
1606 let (ep_store, mut requests) = init_test_with_fake_client(cx);
1607 let fs = FakeFs::new(cx.executor());
1608 fs.insert_tree(
1609 "/root",
1610 json!({
1611 "foo.md": "Hello!\nHow\nBye\n"
1612 }),
1613 )
1614 .await;
1615 let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
1616
1617 let buffer = project
1618 .update(cx, |project, cx| {
1619 let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
1620 project.open_buffer(path, cx)
1621 })
1622 .await
1623 .unwrap();
1624 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
1625 let position = snapshot.anchor_before(language::Point::new(1, 3));
1626
1627 ep_store.update(cx, |ep_store, cx| {
1628 ep_store.refresh_prediction_from_buffer(
1629 project.clone(),
1630 buffer.clone(),
1631 position,
1632 EditPredictionRequestTrigger::Other,
1633 cx,
1634 );
1635 });
1636
1637 let (request, respond_tx) = requests.predict.next().await.unwrap();
1638 let first_response = model_response(&request, SIMPLE_DIFF);
1639 let first_id = first_response.request_id.clone();
1640 respond_tx.send(first_response).unwrap();
1641
1642 cx.run_until_parked();
1643
1644 ep_store.update(cx, |ep_store, cx| {
1645 assert_eq!(
1646 ep_store
1647 .prediction_at(&buffer, None, &project, cx)
1648 .unwrap()
1649 .id
1650 .0,
1651 first_id
1652 );
1653 });
1654
1655 // a second request is triggered
1656 ep_store.update(cx, |ep_store, cx| {
1657 ep_store.refresh_prediction_from_buffer(
1658 project.clone(),
1659 buffer.clone(),
1660 position,
1661 EditPredictionRequestTrigger::Other,
1662 cx,
1663 );
1664 });
1665
1666 let (request, respond_tx) = requests.predict.next().await.unwrap();
1667 // worse than current prediction
1668 let mut second_response = model_response(
1669 &request,
1670 indoc! { r"
1671 --- a/root/foo.md
1672 +++ b/root/foo.md
1673 @@ ... @@
1674 Hello!
1675 -How
1676 +How are
1677 Bye
1678 "},
1679 );
1680 second_response.model_version = Some("zeta2:test-current-preferred".to_string());
1681 let second_id = second_response.request_id.clone();
1682 respond_tx.send(second_response).unwrap();
1683
1684 cx.run_until_parked();
1685
1686 ep_store.update(cx, |ep_store, cx| {
1687 // first is preferred over second
1688 assert_eq!(
1689 ep_store
1690 .prediction_at(&buffer, None, &project, cx)
1691 .unwrap()
1692 .id
1693 .0,
1694 first_id
1695 );
1696 let shown_prediction_ids = ep_store
1697 .rateable_predictions()
1698 .map(|prediction| prediction.id.to_string())
1699 .collect::<Vec<_>>();
1700 assert!(shown_prediction_ids.is_empty());
1701 });
1702
1703 // second is reported as rejected
1704 let (reject_request, _) = requests.reject.next().await.unwrap();
1705
1706 assert_eq!(
1707 &reject_request.rejections,
1708 &[EditPredictionRejection {
1709 request_id: second_id,
1710 reason: EditPredictionRejectReason::CurrentPreferred,
1711 was_shown: false,
1712 model_version: Some("zeta2:test-current-preferred".to_string()),
1713 e2e_latency_ms: Some(0),
1714 }]
1715 );
1716}
1717
1718#[gpui::test]
1719async fn test_cancel_earlier_pending_requests(cx: &mut TestAppContext) {
1720 let (ep_store, mut requests) = init_test_with_fake_client(cx);
1721 let fs = FakeFs::new(cx.executor());
1722 fs.insert_tree(
1723 "/root",
1724 json!({
1725 "foo.md": "Hello!\nHow\nBye\n"
1726 }),
1727 )
1728 .await;
1729 let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
1730
1731 let buffer = project
1732 .update(cx, |project, cx| {
1733 let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
1734 project.open_buffer(path, cx)
1735 })
1736 .await
1737 .unwrap();
1738 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
1739 let position = snapshot.anchor_before(language::Point::new(1, 3));
1740
1741 // start two refresh tasks
1742 ep_store.update(cx, |ep_store, cx| {
1743 ep_store.refresh_prediction_from_buffer(
1744 project.clone(),
1745 buffer.clone(),
1746 position,
1747 EditPredictionRequestTrigger::Other,
1748 cx,
1749 );
1750 });
1751
1752 let (request1, respond_first) = requests.predict.next().await.unwrap();
1753
1754 ep_store.update(cx, |ep_store, cx| {
1755 ep_store.refresh_prediction_from_buffer(
1756 project.clone(),
1757 buffer.clone(),
1758 position,
1759 EditPredictionRequestTrigger::Other,
1760 cx,
1761 );
1762 });
1763
1764 let (request, respond_second) = requests.predict.next().await.unwrap();
1765
1766 // wait for throttle
1767 cx.run_until_parked();
1768
1769 // second responds first
1770 let second_response = model_response(&request, SIMPLE_DIFF);
1771 let second_id = second_response.request_id.clone();
1772 respond_second.send(second_response).unwrap();
1773
1774 cx.run_until_parked();
1775
1776 ep_store.update(cx, |ep_store, cx| {
1777 // current prediction is second
1778 assert_eq!(
1779 ep_store
1780 .prediction_at(&buffer, None, &project, cx)
1781 .unwrap()
1782 .id
1783 .0,
1784 second_id
1785 );
1786 });
1787
1788 let mut first_response = model_response(&request1, SIMPLE_DIFF);
1789 first_response.model_version = Some("zeta2:test-canceled".to_string());
1790 let first_id = first_response.request_id.clone();
1791 respond_first.send(first_response).unwrap();
1792
1793 cx.run_until_parked();
1794
1795 ep_store.update(cx, |ep_store, cx| {
1796 // current prediction is still second, since first was cancelled
1797 assert_eq!(
1798 ep_store
1799 .prediction_at(&buffer, None, &project, cx)
1800 .unwrap()
1801 .id
1802 .0,
1803 second_id
1804 );
1805 });
1806
1807 // first is reported as rejected
1808 let (reject_request, _) = requests.reject.next().await.unwrap();
1809
1810 cx.run_until_parked();
1811
1812 assert_eq!(
1813 &reject_request.rejections,
1814 &[EditPredictionRejection {
1815 request_id: first_id,
1816 reason: EditPredictionRejectReason::Canceled,
1817 was_shown: false,
1818 model_version: Some("zeta2:test-canceled".to_string()),
1819 e2e_latency_ms: None,
1820 }]
1821 );
1822}
1823
1824#[gpui::test]
1825async fn test_cancel_second_on_third_request(cx: &mut TestAppContext) {
1826 let (ep_store, mut requests) = init_test_with_fake_client(cx);
1827 let fs = FakeFs::new(cx.executor());
1828 fs.insert_tree(
1829 "/root",
1830 json!({
1831 "foo.md": "Hello!\nHow\nBye\n"
1832 }),
1833 )
1834 .await;
1835 let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
1836
1837 let buffer = project
1838 .update(cx, |project, cx| {
1839 let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
1840 project.open_buffer(path, cx)
1841 })
1842 .await
1843 .unwrap();
1844 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
1845 let position = snapshot.anchor_before(language::Point::new(1, 3));
1846
1847 // start two refresh tasks
1848 ep_store.update(cx, |ep_store, cx| {
1849 ep_store.refresh_prediction_from_buffer(
1850 project.clone(),
1851 buffer.clone(),
1852 position,
1853 EditPredictionRequestTrigger::Other,
1854 cx,
1855 );
1856 });
1857
1858 let (request1, respond_first) = requests.predict.next().await.unwrap();
1859
1860 ep_store.update(cx, |ep_store, cx| {
1861 ep_store.refresh_prediction_from_buffer(
1862 project.clone(),
1863 buffer.clone(),
1864 position,
1865 EditPredictionRequestTrigger::Other,
1866 cx,
1867 );
1868 });
1869
1870 let (request2, respond_second) = requests.predict.next().await.unwrap();
1871
1872 // wait for throttle, so requests are sent
1873 cx.run_until_parked();
1874
1875 ep_store.update(cx, |ep_store, cx| {
1876 // start a third request
1877 ep_store.refresh_prediction_from_buffer(
1878 project.clone(),
1879 buffer.clone(),
1880 position,
1881 EditPredictionRequestTrigger::Other,
1882 cx,
1883 );
1884
1885 // 2 are pending, so 2nd is cancelled
1886 assert_eq!(
1887 ep_store
1888 .get_or_init_project(&project, cx)
1889 .cancelled_predictions
1890 .iter()
1891 .copied()
1892 .collect::<Vec<_>>(),
1893 [1]
1894 );
1895 });
1896
1897 // wait for throttle
1898 cx.run_until_parked();
1899
1900 let (request3, respond_third) = requests.predict.next().await.unwrap();
1901
1902 let first_response = model_response(&request1, SIMPLE_DIFF);
1903 let first_id = first_response.request_id.clone();
1904 respond_first.send(first_response).unwrap();
1905
1906 cx.run_until_parked();
1907
1908 ep_store.update(cx, |ep_store, cx| {
1909 // current prediction is first
1910 assert_eq!(
1911 ep_store
1912 .prediction_at(&buffer, None, &project, cx)
1913 .unwrap()
1914 .id
1915 .0,
1916 first_id
1917 );
1918 });
1919
1920 let mut cancelled_response = model_response(&request2, SIMPLE_DIFF);
1921 cancelled_response.model_version = Some("zeta2:test-canceled-second".to_string());
1922 let cancelled_id = cancelled_response.request_id.clone();
1923 respond_second.send(cancelled_response).unwrap();
1924
1925 cx.run_until_parked();
1926
1927 ep_store.update(cx, |ep_store, cx| {
1928 // current prediction is still first, since second was cancelled
1929 assert_eq!(
1930 ep_store
1931 .prediction_at(&buffer, None, &project, cx)
1932 .unwrap()
1933 .id
1934 .0,
1935 first_id
1936 );
1937 });
1938
1939 let third_response = model_response(&request3, SIMPLE_DIFF);
1940 let third_response_id = third_response.request_id.clone();
1941 respond_third.send(third_response).unwrap();
1942
1943 cx.run_until_parked();
1944
1945 ep_store.update(cx, |ep_store, cx| {
1946 // third completes and replaces first
1947 assert_eq!(
1948 ep_store
1949 .prediction_at(&buffer, None, &project, cx)
1950 .unwrap()
1951 .id
1952 .0,
1953 third_response_id
1954 );
1955 });
1956
1957 // second is reported as rejected
1958 let (reject_request, _) = requests.reject.next().await.unwrap();
1959
1960 cx.run_until_parked();
1961
1962 assert_eq!(
1963 &reject_request.rejections,
1964 &[
1965 EditPredictionRejection {
1966 request_id: cancelled_id,
1967 reason: EditPredictionRejectReason::Canceled,
1968 was_shown: false,
1969 model_version: Some("zeta2:test-canceled-second".to_string()),
1970 e2e_latency_ms: None,
1971 },
1972 EditPredictionRejection {
1973 request_id: first_id,
1974 reason: EditPredictionRejectReason::Replaced,
1975 was_shown: false,
1976 model_version: None,
1977 // 2 throttle waits (for 2nd and 3rd requests) elapsed
1978 // between this request's start and response.
1979 e2e_latency_ms: Some(2 * EditPredictionStore::THROTTLE_TIMEOUT.as_millis()),
1980 }
1981 ]
1982 );
1983}
1984
1985#[gpui::test]
1986async fn test_cloud_timeout_backs_off_zeta_requests(cx: &mut TestAppContext) {
1987 let (ep_store, mut requests) = init_test_with_fake_client(cx);
1988 let fs = FakeFs::new(cx.executor());
1989 fs.insert_tree(
1990 "/root",
1991 json!({
1992 "foo.md": "Hello!\nHow\nBye\n"
1993 }),
1994 )
1995 .await;
1996 let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
1997
1998 let buffer = project
1999 .update(cx, |project, cx| {
2000 let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
2001 project.open_buffer(path, cx)
2002 })
2003 .await
2004 .unwrap();
2005 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
2006 let position = snapshot.anchor_before(language::Point::new(1, 3));
2007
2008 ep_store.update(cx, |ep_store, cx| {
2009 ep_store.register_project(&project, cx);
2010 ep_store.register_buffer(&buffer, &project, cx);
2011 });
2012
2013 ep_store.update(cx, |ep_store, cx| {
2014 ep_store.refresh_prediction_from_buffer(
2015 project.clone(),
2016 buffer.clone(),
2017 position,
2018 EditPredictionRequestTrigger::Other,
2019 cx,
2020 );
2021 });
2022 let (_request, respond_tx) = requests.predict.next().await.unwrap();
2023 respond_tx.send(request_timeout_response()).unwrap();
2024 cx.run_until_parked();
2025
2026 ep_store.update(cx, |ep_store, cx| {
2027 ep_store.refresh_prediction_from_buffer(
2028 project.clone(),
2029 buffer.clone(),
2030 position,
2031 EditPredictionRequestTrigger::Other,
2032 cx,
2033 );
2034 });
2035 cx.background_executor
2036 .advance_clock(EditPredictionStore::THROTTLE_TIMEOUT);
2037 cx.background_executor.run_until_parked();
2038 cx.run_until_parked();
2039 assert_no_predict_request_ready(&mut requests.predict);
2040
2041 cx.background_executor
2042 .advance_clock(REQUEST_TIMEOUT_BACKOFF);
2043 cx.background_executor.run_until_parked();
2044 cx.run_until_parked();
2045
2046 ep_store.update(cx, |ep_store, cx| {
2047 ep_store.refresh_prediction_from_buffer(
2048 project.clone(),
2049 buffer.clone(),
2050 position,
2051 EditPredictionRequestTrigger::Other,
2052 cx,
2053 );
2054 });
2055 let (_request, respond_tx) = requests.predict.next().await.unwrap();
2056 respond_tx.send(empty_response()).unwrap();
2057 cx.run_until_parked();
2058}
2059
2060#[gpui::test]
2061async fn test_same_frame_duplicate_requests_deduplicated(cx: &mut TestAppContext) {
2062 let (ep_store, mut requests) = init_test_with_fake_client(cx);
2063 let fs = FakeFs::new(cx.executor());
2064 fs.insert_tree(
2065 "/root",
2066 json!({
2067 "foo.md": "Hello!\nHow\nBye\n"
2068 }),
2069 )
2070 .await;
2071 let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
2072
2073 let buffer = project
2074 .update(cx, |project, cx| {
2075 let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
2076 project.open_buffer(path, cx)
2077 })
2078 .await
2079 .unwrap();
2080 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
2081 let position = snapshot.anchor_before(language::Point::new(1, 3));
2082
2083 // Enqueue two refresh calls in the same synchronous frame (no yielding).
2084 // Both `cx.spawn` tasks are created before either executes, so they both
2085 // capture the same `proceed_count_at_enqueue`. Only the first task should
2086 // pass the deduplication gate; the second should be skipped.
2087 ep_store.update(cx, |ep_store, cx| {
2088 ep_store.refresh_prediction_from_buffer(
2089 project.clone(),
2090 buffer.clone(),
2091 position,
2092 EditPredictionRequestTrigger::Other,
2093 cx,
2094 );
2095 ep_store.refresh_prediction_from_buffer(
2096 project.clone(),
2097 buffer.clone(),
2098 position,
2099 EditPredictionRequestTrigger::Other,
2100 cx,
2101 );
2102 });
2103
2104 // Let both spawned tasks run to completion (including any throttle waits).
2105 cx.run_until_parked();
2106
2107 // Exactly one prediction request should have been sent.
2108 let (request, respond_tx) = requests.predict.next().await.unwrap();
2109 respond_tx
2110 .send(model_response(&request, SIMPLE_DIFF))
2111 .unwrap();
2112 cx.run_until_parked();
2113
2114 // No second request should be pending.
2115 assert_no_predict_request_ready(&mut requests.predict);
2116}
2117
2118#[gpui::test]
2119async fn test_rejections_flushing(cx: &mut TestAppContext) {
2120 let (ep_store, mut requests) = init_test_with_fake_client(cx);
2121
2122 ep_store.update(cx, |ep_store, cx| {
2123 ep_store.reject_prediction(
2124 EditPredictionId("test-1".into()),
2125 EditPredictionRejectReason::Discarded,
2126 false,
2127 None,
2128 None,
2129 cx,
2130 );
2131 ep_store.reject_prediction(
2132 EditPredictionId("test-2".into()),
2133 EditPredictionRejectReason::Canceled,
2134 true,
2135 None,
2136 None,
2137 cx,
2138 );
2139 });
2140
2141 cx.executor().advance_clock(REJECT_REQUEST_DEBOUNCE);
2142 cx.run_until_parked();
2143
2144 let (reject_request, respond_tx) = requests.reject.next().await.unwrap();
2145 respond_tx.send(()).unwrap();
2146
2147 // batched
2148 assert_eq!(reject_request.rejections.len(), 2);
2149 assert_eq!(
2150 reject_request.rejections[0],
2151 EditPredictionRejection {
2152 request_id: "test-1".to_string(),
2153 reason: EditPredictionRejectReason::Discarded,
2154 was_shown: false,
2155 model_version: None,
2156 e2e_latency_ms: None
2157 }
2158 );
2159 assert_eq!(
2160 reject_request.rejections[1],
2161 EditPredictionRejection {
2162 request_id: "test-2".to_string(),
2163 reason: EditPredictionRejectReason::Canceled,
2164 was_shown: true,
2165 model_version: None,
2166 e2e_latency_ms: None
2167 }
2168 );
2169
2170 // Reaching batch size limit sends without debounce
2171 ep_store.update(cx, |ep_store, cx| {
2172 for i in 0..70 {
2173 ep_store.reject_prediction(
2174 EditPredictionId(format!("batch-{}", i).into()),
2175 EditPredictionRejectReason::Discarded,
2176 false,
2177 None,
2178 None,
2179 cx,
2180 );
2181 }
2182 });
2183
2184 // First MAX/2 items are sent immediately
2185 cx.run_until_parked();
2186 let (reject_request, respond_tx) = requests.reject.next().await.unwrap();
2187 respond_tx.send(()).unwrap();
2188
2189 assert_eq!(reject_request.rejections.len(), 50);
2190 assert_eq!(reject_request.rejections[0].request_id, "batch-0");
2191 assert_eq!(reject_request.rejections[49].request_id, "batch-49");
2192
2193 // Remaining items are debounced with the next batch
2194 cx.executor().advance_clock(Duration::from_secs(15));
2195 cx.run_until_parked();
2196
2197 let (reject_request, respond_tx) = requests.reject.next().await.unwrap();
2198 respond_tx.send(()).unwrap();
2199
2200 assert_eq!(reject_request.rejections.len(), 20);
2201 assert_eq!(reject_request.rejections[0].request_id, "batch-50");
2202 assert_eq!(reject_request.rejections[19].request_id, "batch-69");
2203
2204 // Request failure
2205 ep_store.update(cx, |ep_store, cx| {
2206 ep_store.reject_prediction(
2207 EditPredictionId("retry-1".into()),
2208 EditPredictionRejectReason::Discarded,
2209 false,
2210 None,
2211 None,
2212 cx,
2213 );
2214 });
2215
2216 cx.executor().advance_clock(REJECT_REQUEST_DEBOUNCE);
2217 cx.run_until_parked();
2218
2219 let (reject_request, _respond_tx) = requests.reject.next().await.unwrap();
2220 assert_eq!(reject_request.rejections.len(), 1);
2221 assert_eq!(reject_request.rejections[0].request_id, "retry-1");
2222 // Simulate failure
2223 drop(_respond_tx);
2224
2225 // Add another rejection
2226 ep_store.update(cx, |ep_store, cx| {
2227 ep_store.reject_prediction(
2228 EditPredictionId("retry-2".into()),
2229 EditPredictionRejectReason::Discarded,
2230 false,
2231 None,
2232 None,
2233 cx,
2234 );
2235 });
2236
2237 cx.executor().advance_clock(REJECT_REQUEST_DEBOUNCE);
2238 cx.run_until_parked();
2239
2240 // Retry should include both the failed item and the new one
2241 let (reject_request, respond_tx) = requests.reject.next().await.unwrap();
2242 respond_tx.send(()).unwrap();
2243
2244 assert_eq!(reject_request.rejections.len(), 2);
2245 assert_eq!(reject_request.rejections[0].request_id, "retry-1");
2246 assert_eq!(reject_request.rejections[1].request_id, "retry-2");
2247}
2248
2249#[gpui::test]
2250fn test_active_buffer_diagnostics_fetching(cx: &mut TestAppContext) {
2251 let diagnostic_marker: TextRangeMarker = ('«', '»').into();
2252 let search_range_marker: TextRangeMarker = ('[', ']').into();
2253
2254 let (text, mut ranges) = marked_text_ranges_by(
2255 indoc! {r#"
2256 fn alpha() {
2257 let «first_value» = 1;
2258 }
2259
2260 [fn beta() {
2261 let «second_value» = 2;
2262 let third_value = second_value + missing_symbol;
2263 }ˇ]
2264
2265 fn gamma() {
2266 let «fourth_value» = missing_other_symbol;
2267 }
2268 "#},
2269 vec![diagnostic_marker.clone(), search_range_marker.clone()],
2270 );
2271
2272 let diagnostic_ranges = ranges.remove(&diagnostic_marker).unwrap_or_default();
2273 let search_ranges = ranges.remove(&search_range_marker).unwrap_or_default();
2274
2275 let buffer = cx.new(|cx| Buffer::local(&text, cx));
2276
2277 buffer.update(cx, |buffer, cx| {
2278 let snapshot = buffer.snapshot();
2279 let diagnostics = DiagnosticSet::new(
2280 diagnostic_ranges
2281 .iter()
2282 .enumerate()
2283 .map(|(index, range)| DiagnosticEntry {
2284 range: snapshot.offset_to_point_utf16(range.start)
2285 ..snapshot.offset_to_point_utf16(range.end),
2286 diagnostic: Diagnostic {
2287 severity: match index {
2288 0 => DiagnosticSeverity::WARNING,
2289 1 => DiagnosticSeverity::ERROR,
2290 _ => DiagnosticSeverity::HINT,
2291 },
2292 message: match index {
2293 0 => "first warning".to_string(),
2294 1 => "second error".to_string(),
2295 _ => "third hint".to_string(),
2296 },
2297 group_id: index + 1,
2298 is_primary: true,
2299 source_kind: language::DiagnosticSourceKind::Pushed,
2300 ..Diagnostic::default()
2301 },
2302 }),
2303 &snapshot,
2304 );
2305 buffer.update_diagnostics(LanguageServerId(0), diagnostics, cx);
2306 });
2307
2308 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
2309 let search_range = snapshot.offset_to_point(search_ranges[0].start)
2310 ..snapshot.offset_to_point(search_ranges[0].end);
2311
2312 let active_buffer_diagnostics = zeta::active_buffer_diagnostics(&snapshot, search_range, 5, 0);
2313
2314 assert_eq!(
2315 active_buffer_diagnostics,
2316 vec![zeta_prompt::ActiveBufferDiagnostic {
2317 severity: Some(1),
2318 message: "second error".to_string(),
2319 snippet: " let second_value = 2;".to_string(),
2320 snippet_buffer_row_range: 5..5,
2321 diagnostic_range_in_snippet: 8..20,
2322 }]
2323 );
2324
2325 let active_buffer_diagnostics =
2326 zeta::active_buffer_diagnostics(&snapshot, Point::new(0, 0)..snapshot.max_point(), 5, 100);
2327 assert_eq!(
2328 active_buffer_diagnostics,
2329 vec![
2330 zeta_prompt::ActiveBufferDiagnostic {
2331 severity: Some(1),
2332 message: "second error".to_string(),
2333 snippet: String::new(),
2334 snippet_buffer_row_range: 5..5,
2335 diagnostic_range_in_snippet: 0..0,
2336 },
2337 zeta_prompt::ActiveBufferDiagnostic {
2338 severity: Some(2),
2339 message: "first warning".to_string(),
2340 snippet: String::new(),
2341 snippet_buffer_row_range: 1..1,
2342 diagnostic_range_in_snippet: 0..0,
2343 },
2344 zeta_prompt::ActiveBufferDiagnostic {
2345 severity: Some(4),
2346 message: "third hint".to_string(),
2347 snippet: String::new(),
2348 snippet_buffer_row_range: 10..10,
2349 diagnostic_range_in_snippet: 0..0,
2350 },
2351 ]
2352 );
2353
2354 let buffer = cx.new(|cx| {
2355 Buffer::local(
2356 indoc! {"
2357 one
2358 two
2359 three
2360 four
2361 five
2362 "},
2363 cx,
2364 )
2365 });
2366
2367 buffer.update(cx, |buffer, cx| {
2368 let snapshot = buffer.snapshot();
2369 let diagnostics = DiagnosticSet::new(
2370 vec![
2371 DiagnosticEntry {
2372 range: text::PointUtf16::new(0, 0)..text::PointUtf16::new(0, 3),
2373 diagnostic: Diagnostic {
2374 severity: DiagnosticSeverity::ERROR,
2375 message: "row zero".to_string(),
2376 group_id: 1,
2377 is_primary: true,
2378 source_kind: language::DiagnosticSourceKind::Pushed,
2379 ..Diagnostic::default()
2380 },
2381 },
2382 DiagnosticEntry {
2383 range: text::PointUtf16::new(2, 0)..text::PointUtf16::new(2, 5),
2384 diagnostic: Diagnostic {
2385 severity: DiagnosticSeverity::WARNING,
2386 message: "row two".to_string(),
2387 group_id: 2,
2388 is_primary: true,
2389 source_kind: language::DiagnosticSourceKind::Pushed,
2390 ..Diagnostic::default()
2391 },
2392 },
2393 DiagnosticEntry {
2394 range: text::PointUtf16::new(4, 0)..text::PointUtf16::new(4, 4),
2395 diagnostic: Diagnostic {
2396 severity: DiagnosticSeverity::INFORMATION,
2397 message: "row four".to_string(),
2398 group_id: 3,
2399 is_primary: true,
2400 source_kind: language::DiagnosticSourceKind::Pushed,
2401 ..Diagnostic::default()
2402 },
2403 },
2404 ],
2405 &snapshot,
2406 );
2407 buffer.update_diagnostics(LanguageServerId(0), diagnostics, cx);
2408 });
2409
2410 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
2411
2412 let active_buffer_diagnostics =
2413 zeta::active_buffer_diagnostics(&snapshot, Point::new(2, 0)..Point::new(4, 0), 3, 0);
2414
2415 assert_eq!(
2416 active_buffer_diagnostics
2417 .iter()
2418 .map(|diagnostic| (
2419 diagnostic.severity,
2420 diagnostic.message.clone(),
2421 diagnostic.snippet.clone(),
2422 diagnostic.snippet_buffer_row_range.clone(),
2423 diagnostic.diagnostic_range_in_snippet.clone(),
2424 ))
2425 .collect::<Vec<_>>(),
2426 vec![
2427 (
2428 Some(2),
2429 "row two".to_string(),
2430 "three".to_string(),
2431 2..2,
2432 0..5,
2433 ),
2434 (
2435 Some(3),
2436 "row four".to_string(),
2437 "five".to_string(),
2438 4..4,
2439 0..4,
2440 ),
2441 ]
2442 );
2443}
2444
2445#[gpui::test]
2446fn test_active_buffer_diagnostics_collection_limits(cx: &mut TestAppContext) {
2447 let text = (0..25)
2448 .map(|row| format!("line {row}\n"))
2449 .collect::<String>();
2450 let buffer = cx.new(|cx| Buffer::local(&text, cx));
2451
2452 buffer.update(cx, |buffer, cx| {
2453 let snapshot = buffer.snapshot();
2454 let diagnostics = DiagnosticSet::new(
2455 (0..25)
2456 .map(|row| DiagnosticEntry {
2457 range: text::PointUtf16::new(row, 0)..text::PointUtf16::new(row, 4),
2458 diagnostic: Diagnostic {
2459 severity: DiagnosticSeverity::ERROR,
2460 message: format!("row {row}"),
2461 group_id: row as usize,
2462 is_primary: true,
2463 source_kind: language::DiagnosticSourceKind::Pushed,
2464 ..Diagnostic::default()
2465 },
2466 })
2467 .collect::<Vec<_>>(),
2468 &snapshot,
2469 );
2470 buffer.update_diagnostics(LanguageServerId(0), diagnostics, cx);
2471 });
2472
2473 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
2474 let active_buffer_diagnostics =
2475 zeta::active_buffer_diagnostics(&snapshot, Point::new(0, 0)..Point::new(25, 0), 12, 0);
2476
2477 assert_eq!(active_buffer_diagnostics.len(), 20);
2478 assert!(
2479 active_buffer_diagnostics
2480 .iter()
2481 .any(|diagnostic| diagnostic.message == "row 12")
2482 );
2483 assert!(
2484 active_buffer_diagnostics
2485 .iter()
2486 .all(|diagnostic| diagnostic.message != "row 0" && diagnostic.message != "row 24")
2487 );
2488
2489 let text = (0..300)
2490 .map(|row| format!("line {row} has some diagnostic context\n"))
2491 .collect::<String>();
2492 let long_message = "diagnostic message ".repeat(1000);
2493 let buffer = cx.new(|cx| Buffer::local(&text, cx));
2494
2495 buffer.update(cx, |buffer, cx| {
2496 let snapshot = buffer.snapshot();
2497 let diagnostics = DiagnosticSet::new(
2498 vec![DiagnosticEntry {
2499 range: text::PointUtf16::new(150, 0)..text::PointUtf16::new(150, 4),
2500 diagnostic: Diagnostic {
2501 severity: DiagnosticSeverity::ERROR,
2502 message: long_message.clone(),
2503 group_id: 1,
2504 is_primary: true,
2505 source_kind: language::DiagnosticSourceKind::Pushed,
2506 ..Diagnostic::default()
2507 },
2508 }],
2509 &snapshot,
2510 );
2511 buffer.update_diagnostics(LanguageServerId(0), diagnostics, cx);
2512 });
2513
2514 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
2515 let active_buffer_diagnostics = zeta::active_buffer_diagnostics(
2516 &snapshot,
2517 Point::new(100, 0)..Point::new(200, 0),
2518 150,
2519 2000,
2520 );
2521
2522 assert_eq!(active_buffer_diagnostics.len(), 1);
2523 assert!(
2524 active_buffer_diagnostics[0].message.len()
2525 <= crate::zeta::MAX_ACTIVE_BUFFER_DIAGNOSTIC_MESSAGE_TOKENS_TO_COLLECT * 3 + 2
2526 );
2527 assert!(active_buffer_diagnostics[0].message.len() < long_message.len());
2528 assert!(
2529 active_buffer_diagnostics[0].snippet.len()
2530 <= crate::zeta::MAX_ACTIVE_BUFFER_DIAGNOSTIC_SNIPPET_TOKENS_TO_COLLECT * 3 + 2
2531 );
2532 assert!(active_buffer_diagnostics[0].snippet.len() < text.len());
2533}
2534
2535// Generate a model response that would apply the given diff to the active file.
2536fn model_response(request: &PredictEditsV3Request, diff_to_apply: &str) -> PredictEditsV3Response {
2537 let editable_range =
2538 zeta_prompt::excerpt_range_for_format(Default::default(), &request.input.excerpt_ranges).1;
2539 let excerpt = request.input.cursor_excerpt[editable_range.clone()].to_string();
2540 let new_excerpt = apply_diff_to_string(diff_to_apply, &excerpt).unwrap();
2541
2542 PredictEditsV3Response {
2543 request_id: Uuid::new_v4().to_string(),
2544 editable_range,
2545 output: new_excerpt,
2546 cursor_offset: None,
2547 model_version: None,
2548 }
2549}
2550
2551fn empty_response() -> PredictEditsV3Response {
2552 PredictEditsV3Response {
2553 request_id: Uuid::new_v4().to_string(),
2554 editable_range: 0..0,
2555 output: String::new(),
2556 cursor_offset: None,
2557 model_version: None,
2558 }
2559}
2560
2561const REQUEST_TIMEOUT_RESPONSE_ID: &str = "__request_timeout__";
2562
2563fn request_timeout_response() -> PredictEditsV3Response {
2564 PredictEditsV3Response {
2565 request_id: REQUEST_TIMEOUT_RESPONSE_ID.to_string(),
2566 editable_range: 0..0,
2567 output: String::new(),
2568 cursor_offset: None,
2569 model_version: None,
2570 }
2571}
2572
2573fn prompt_from_request(request: &PredictEditsV3Request) -> String {
2574 zeta_prompt::format_zeta_prompt(&request.input, zeta_prompt::ZetaFormat::default())
2575 .expect("default zeta prompt formatting should succeed in edit prediction tests")
2576}
2577
2578fn assert_no_predict_request_ready<Request, Response>(
2579 requests: &mut mpsc::UnboundedReceiver<(Request, oneshot::Sender<Response>)>,
2580) {
2581 if requests.next().now_or_never().flatten().is_some() {
2582 panic!("Unexpected prediction request while throttled.");
2583 }
2584}
2585
2586struct RequestChannels {
2587 predict: mpsc::UnboundedReceiver<(
2588 PredictEditsV3Request,
2589 oneshot::Sender<PredictEditsV3Response>,
2590 )>,
2591 predict_v4: mpsc::UnboundedReceiver<(
2592 PredictEditsV4Request,
2593 oneshot::Sender<PredictEditsV4Response>,
2594 )>,
2595 reject: mpsc::UnboundedReceiver<(RejectEditPredictionsBody, oneshot::Sender<()>)>,
2596 settled: mpsc::UnboundedReceiver<SettledEditPrediction>,
2597}
2598
2599fn init_test_with_fake_client(
2600 cx: &mut TestAppContext,
2601) -> (Entity<EditPredictionStore>, RequestChannels) {
2602 init_test_with_fake_client_and_legacy_data_collection(cx, None)
2603}
2604
2605fn init_test_with_fake_client_and_legacy_data_collection(
2606 cx: &mut TestAppContext,
2607 legacy_data_collection_choice: Option<&str>,
2608) -> (Entity<EditPredictionStore>, RequestChannels) {
2609 let result = cx.update(move |cx| {
2610 cx.set_global(AppDatabase::test_new());
2611 let settings_store = SettingsStore::test(cx);
2612 cx.set_global(settings_store);
2613 disable_jumps_feature_flag(cx);
2614 zlog::init_test();
2615
2616 if let Some(legacy_data_collection_choice) = legacy_data_collection_choice {
2617 KeyValueStore::global(cx)
2618 .write_kvp(
2619 ZED_PREDICT_DATA_COLLECTION_CHOICE.into(),
2620 legacy_data_collection_choice.to_string(),
2621 )
2622 .now_or_never()
2623 .expect("legacy data collection write should complete immediately")
2624 .expect("legacy data collection write should succeed");
2625 }
2626
2627 let (predict_req_tx, predict_req_rx) = mpsc::unbounded();
2628 let (predict_v4_req_tx, predict_v4_req_rx) = mpsc::unbounded();
2629 let (reject_req_tx, reject_req_rx) = mpsc::unbounded();
2630 let (settled_req_tx, settled_req_rx) = mpsc::unbounded();
2631
2632 let http_client = FakeHttpClient::create({
2633 move |req| {
2634 let uri = req.uri().path().to_string();
2635 let content_encoding = req
2636 .headers()
2637 .get("Content-Encoding")
2638 .and_then(|value| value.to_str().ok())
2639 .map(str::to_owned);
2640 let mut body = req.into_body();
2641 let predict_req_tx = predict_req_tx.clone();
2642 let predict_v4_req_tx = predict_v4_req_tx.clone();
2643 let reject_req_tx = reject_req_tx.clone();
2644 let settled_req_tx = settled_req_tx.clone();
2645 async move {
2646 let resp = match uri.as_str() {
2647 "/client/llm_tokens" => serde_json::to_string(&json!({
2648 "token": "test"
2649 }))
2650 .unwrap(),
2651 "/predict_edits/v3" => {
2652 let mut buf = Vec::new();
2653 body.read_to_end(&mut buf).await.ok();
2654 let decompressed = zstd::decode_all(&buf[..]).unwrap();
2655 let req = serde_json::from_slice(&decompressed).unwrap();
2656
2657 let (res_tx, res_rx) = oneshot::channel::<PredictEditsV3Response>();
2658 predict_req_tx.unbounded_send((req, res_tx)).unwrap();
2659 let response = res_rx.await?;
2660 if response.request_id == REQUEST_TIMEOUT_RESPONSE_ID {
2661 return Ok(Response::builder()
2662 .status(http_client::http::StatusCode::REQUEST_TIMEOUT)
2663 .body(
2664 http_client::http::StatusCode::REQUEST_TIMEOUT
2665 .as_str()
2666 .into(),
2667 )
2668 .unwrap());
2669 }
2670 serde_json::to_string(&response).unwrap()
2671 }
2672 "/predict_edits/v4" => {
2673 let mut buf = Vec::new();
2674 body.read_to_end(&mut buf).await.ok();
2675 let decompressed = zstd::decode_all(&buf[..]).unwrap();
2676 let req = serde_json::from_slice(&decompressed).unwrap();
2677
2678 let (res_tx, res_rx) = oneshot::channel::<PredictEditsV4Response>();
2679 predict_v4_req_tx.unbounded_send((req, res_tx)).unwrap();
2680 let response = res_rx.await?;
2681 if response.request_id == REQUEST_TIMEOUT_RESPONSE_ID {
2682 return Ok(Response::builder()
2683 .status(http_client::http::StatusCode::REQUEST_TIMEOUT)
2684 .body(
2685 http_client::http::StatusCode::REQUEST_TIMEOUT
2686 .as_str()
2687 .into(),
2688 )
2689 .unwrap());
2690 }
2691 serde_json::to_string(&response).unwrap()
2692 }
2693 "/predict_edits/reject" => {
2694 let mut buf = Vec::new();
2695 body.read_to_end(&mut buf).await.ok();
2696 let req = serde_json::from_slice(&buf).unwrap();
2697
2698 let (res_tx, res_rx) = oneshot::channel();
2699 reject_req_tx.unbounded_send((req, res_tx)).unwrap();
2700 serde_json::to_string(&res_rx.await?).unwrap()
2701 }
2702 "/predict_edits/settled" => {
2703 let mut buf = Vec::new();
2704 body.read_to_end(&mut buf).await.ok();
2705 let body = if content_encoding.as_deref() == Some("zstd") {
2706 zstd::decode_all(&buf[..]).unwrap()
2707 } else {
2708 buf
2709 };
2710 let req: SubmitEditPredictionSettledBatchBody =
2711 serde_json::from_slice(&body).unwrap();
2712 for prediction in req.predictions {
2713 settled_req_tx.unbounded_send(prediction).unwrap();
2714 }
2715 serde_json::to_string(&SubmitEditPredictionSettledResponse {}).unwrap()
2716 }
2717 _ => {
2718 panic!("Unexpected path: {}", uri)
2719 }
2720 };
2721
2722 Ok(Response::builder().body(resp.into()).unwrap())
2723 }
2724 }
2725 });
2726
2727 let client = client::Client::new(Arc::new(FakeSystemClock::new()), http_client, cx);
2728 client.cloud_client().set_credentials(1, "test".into());
2729
2730 let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
2731 language_model::init(cx);
2732 RefreshLlmTokenListener::register(client.clone(), user_store.clone(), cx);
2733 let ep_store = EditPredictionStore::global(&client, &user_store, cx);
2734
2735 (
2736 ep_store,
2737 user_store,
2738 RequestChannels {
2739 predict: predict_req_rx,
2740 predict_v4: predict_v4_req_rx,
2741 reject: reject_req_rx,
2742 settled: settled_req_rx,
2743 },
2744 )
2745 });
2746
2747 let (ep_store, user_store, channels) = result;
2748 set_test_organization(&user_store, cx);
2749 (ep_store, channels)
2750}
2751
2752/// Configures a current organization on the given `UserStore` for tests.
2753///
2754/// The test client starts out signed out, which causes `UserStore` to clear the
2755/// current organization once that initial status is processed. This waits for
2756/// that to happen before configuring the organization, so it isn't subsequently
2757/// wiped out.
2758fn set_test_organization(user_store: &Entity<UserStore>, cx: &mut TestAppContext) {
2759 cx.run_until_parked();
2760 cx.update(|cx| {
2761 user_store.update(cx, |store, cx| {
2762 store.set_current_organization_configuration_for_test(
2763 Arc::new(Organization {
2764 id: OrganizationId("org_1".into()),
2765 name: "Organization 1".into(),
2766 is_personal: false,
2767 }),
2768 OrganizationConfiguration {
2769 is_zed_model_provider_enabled: true,
2770 is_agent_thread_feedback_enabled: true,
2771 is_collaboration_enabled: true,
2772 edit_prediction: OrganizationEditPredictionConfiguration {
2773 is_enabled: true,
2774 is_feedback_enabled: true,
2775 },
2776 },
2777 cx,
2778 )
2779 });
2780 });
2781}
2782
2783#[gpui::test]
2784async fn test_edit_prediction_basic_interpolation(cx: &mut TestAppContext) {
2785 let buffer = cx.new(|cx| Buffer::local("Lorem ipsum dolor", cx));
2786 let edits: Arc<[(Range<Anchor>, Arc<str>)]> = cx.update(|cx| {
2787 to_completion_edits([(2..5, "REM".into()), (9..11, "".into())], &buffer, cx).into()
2788 });
2789
2790 let edit_preview = cx
2791 .read(|cx| buffer.read(cx).preview_edits(edits.clone(), cx))
2792 .await;
2793
2794 let prediction = EditPrediction {
2795 edits,
2796 cursor_position: None,
2797 editable_range: None,
2798 edit_preview,
2799 buffer: buffer.clone(),
2800 snapshot: cx.read(|cx| buffer.read(cx).snapshot()),
2801 id: EditPredictionId("the-id".into()),
2802 inputs: EditPredictionInputs::V2(Zeta2PromptInput {
2803 events: Default::default(),
2804 related_files: Default::default(),
2805 active_buffer_diagnostics: vec![],
2806 cursor_path: Path::new("").into(),
2807 cursor_excerpt: "".into(),
2808 cursor_offset_in_excerpt: 0,
2809 excerpt_start_row: None,
2810 excerpt_ranges: Default::default(),
2811 syntax_ranges: None,
2812 in_open_source_repo: false,
2813 can_collect_data: false,
2814 repo_url: None,
2815 }),
2816 model_version: None,
2817 trigger: PredictEditsRequestTrigger::Other,
2818 };
2819
2820 cx.update(|cx| {
2821 assert_eq!(
2822 from_completion_edits(
2823 &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
2824 &buffer,
2825 cx
2826 ),
2827 vec![(2..5, "REM".into()), (9..11, "".into())]
2828 );
2829
2830 buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "")], None, cx));
2831 assert_eq!(
2832 from_completion_edits(
2833 &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
2834 &buffer,
2835 cx
2836 ),
2837 vec![(2..2, "REM".into()), (6..8, "".into())]
2838 );
2839
2840 buffer.update(cx, |buffer, cx| buffer.undo(cx));
2841 assert_eq!(
2842 from_completion_edits(
2843 &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
2844 &buffer,
2845 cx
2846 ),
2847 vec![(2..5, "REM".into()), (9..11, "".into())]
2848 );
2849
2850 buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "R")], None, cx));
2851 assert_eq!(
2852 from_completion_edits(
2853 &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
2854 &buffer,
2855 cx
2856 ),
2857 vec![(3..3, "EM".into()), (7..9, "".into())]
2858 );
2859
2860 buffer.update(cx, |buffer, cx| buffer.edit([(3..3, "E")], None, cx));
2861 assert_eq!(
2862 from_completion_edits(
2863 &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
2864 &buffer,
2865 cx
2866 ),
2867 vec![(4..4, "M".into()), (8..10, "".into())]
2868 );
2869
2870 buffer.update(cx, |buffer, cx| buffer.edit([(4..4, "M")], None, cx));
2871 assert_eq!(
2872 from_completion_edits(
2873 &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
2874 &buffer,
2875 cx
2876 ),
2877 vec![(9..11, "".into())]
2878 );
2879
2880 buffer.update(cx, |buffer, cx| buffer.edit([(4..5, "")], None, cx));
2881 assert_eq!(
2882 from_completion_edits(
2883 &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
2884 &buffer,
2885 cx
2886 ),
2887 vec![(4..4, "M".into()), (8..10, "".into())]
2888 );
2889
2890 buffer.update(cx, |buffer, cx| buffer.edit([(8..10, "")], None, cx));
2891 assert_eq!(
2892 from_completion_edits(
2893 &prediction.interpolate(&buffer.read(cx).snapshot()).unwrap(),
2894 &buffer,
2895 cx
2896 ),
2897 vec![(4..4, "M".into())]
2898 );
2899
2900 buffer.update(cx, |buffer, cx| buffer.edit([(4..6, "")], None, cx));
2901 assert_eq!(prediction.interpolate(&buffer.read(cx).snapshot()), None);
2902 })
2903}
2904
2905#[gpui::test]
2906async fn test_clean_up_diff(cx: &mut TestAppContext) {
2907 init_test(cx);
2908
2909 assert_eq!(
2910 apply_edit_prediction(
2911 indoc! {"
2912 fn main() {
2913 let word_1 = \"lorem\";
2914 let range = word.len()..word.len();
2915 }
2916 "},
2917 indoc! {"
2918 fn main() {
2919 let word_1 = \"lorem\";
2920 let range = word_1.len()..word_1.len();
2921 }
2922 "},
2923 cx,
2924 )
2925 .await,
2926 indoc! {"
2927 fn main() {
2928 let word_1 = \"lorem\";
2929 let range = word_1.len()..word_1.len();
2930 }
2931 "},
2932 );
2933
2934 assert_eq!(
2935 apply_edit_prediction(
2936 indoc! {"
2937 fn main() {
2938 let story = \"the quick\"
2939 }
2940 "},
2941 indoc! {"
2942 fn main() {
2943 let story = \"the quick brown fox jumps over the lazy dog\";
2944 }
2945 "},
2946 cx,
2947 )
2948 .await,
2949 indoc! {"
2950 fn main() {
2951 let story = \"the quick brown fox jumps over the lazy dog\";
2952 }
2953 "},
2954 );
2955}
2956
2957#[gpui::test]
2958async fn test_edit_prediction_end_of_buffer(cx: &mut TestAppContext) {
2959 init_test(cx);
2960
2961 let buffer_content = "lorem\n";
2962 let completion_response = "lorem\nipsum\n";
2963
2964 assert_eq!(
2965 apply_edit_prediction(buffer_content, completion_response, cx).await,
2966 "lorem\nipsum\n"
2967 );
2968}
2969
2970#[gpui::test]
2971async fn test_edit_prediction_v4_end_of_buffer(cx: &mut TestAppContext) {
2972 init_test(cx);
2973
2974 let fs = FakeFs::new(cx.executor());
2975 fs.insert_tree(
2976 "/project",
2977 json!({
2978 "file.txt": "lorem\n"
2979 }),
2980 )
2981 .await;
2982 let project = Project::test(fs, [path!("/project").as_ref()], cx).await;
2983 let buffer = project
2984 .update(cx, |project, cx| {
2985 let path = project
2986 .find_project_path(path!("/project/file.txt"), cx)
2987 .unwrap();
2988 project.open_buffer(path, cx)
2989 })
2990 .await
2991 .unwrap();
2992 let (ep_store, response) = make_test_ep_store(&project, cx).await;
2993 *response.lock() = "lorem\nipsum\n".to_string();
2994
2995 let position = buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(1, 0)));
2996 ep_store.update(cx, |ep_store, cx| {
2997 ep_store.register_project(&project, cx);
2998 ep_store.register_buffer(&buffer, &project, cx);
2999 ep_store.refresh_prediction_from_buffer(
3000 project.clone(),
3001 buffer.clone(),
3002 position,
3003 EditPredictionRequestTrigger::Other,
3004 cx,
3005 );
3006 });
3007 cx.run_until_parked();
3008
3009 let edits = ep_store.update(cx, |ep_store, cx| {
3010 let prediction = ep_store
3011 .prediction_at(&buffer, None, &project, cx)
3012 .expect("should have prediction");
3013 let prediction = match prediction {
3014 BufferEditPrediction::Local { prediction }
3015 | BufferEditPrediction::Jump { prediction } => prediction,
3016 };
3017 assert!(prediction.editable_range.is_some());
3018 prediction.edits.iter().cloned().collect::<Vec<_>>()
3019 });
3020 buffer.update(cx, |buffer, cx| buffer.edit(edits, None, cx));
3021
3022 buffer.read_with(cx, |buffer, _| {
3023 assert_eq!(buffer.text(), "lorem\nipsum\n");
3024 });
3025}
3026
3027#[gpui::test]
3028async fn test_edit_prediction_no_spurious_trailing_newline(cx: &mut TestAppContext) {
3029 // Test that zeta2's newline normalization logic doesn't insert spurious newlines.
3030 // When the buffer ends without a trailing newline, but the model returns output
3031 // with a trailing newline, zeta2 should normalize both sides before diffing
3032 // so no spurious newline is inserted.
3033 let (ep_store, mut requests) = init_test_with_fake_client(cx);
3034 let fs = FakeFs::new(cx.executor());
3035
3036 // Single line buffer with no trailing newline
3037 fs.insert_tree(
3038 "/root",
3039 json!({
3040 "foo.txt": "hello"
3041 }),
3042 )
3043 .await;
3044 let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
3045
3046 let buffer = project
3047 .update(cx, |project, cx| {
3048 let path = project
3049 .find_project_path(path!("root/foo.txt"), cx)
3050 .unwrap();
3051 project.open_buffer(path, cx)
3052 })
3053 .await
3054 .unwrap();
3055
3056 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
3057 let position = snapshot.anchor_before(language::Point::new(0, 5));
3058
3059 ep_store.update(cx, |ep_store, cx| {
3060 ep_store.refresh_prediction_from_buffer(
3061 project.clone(),
3062 buffer.clone(),
3063 position,
3064 EditPredictionRequestTrigger::Other,
3065 cx,
3066 );
3067 });
3068
3069 let (request, respond_tx) = requests.predict.next().await.unwrap();
3070
3071 // Model returns output WITH a trailing newline, even though the buffer doesn't have one.
3072 // Zeta2 should normalize both sides before diffing, so no spurious newline is inserted.
3073 let excerpt_length = request.input.cursor_excerpt.len();
3074 let response = PredictEditsV3Response {
3075 request_id: Uuid::new_v4().to_string(),
3076 output: "hello world\n".to_string(),
3077 editable_range: 0..excerpt_length,
3078 model_version: None,
3079 cursor_offset: None,
3080 };
3081 respond_tx.send(response).unwrap();
3082
3083 cx.run_until_parked();
3084
3085 // The prediction should insert " world" without adding a newline
3086 ep_store.update(cx, |ep_store, cx| {
3087 let prediction = ep_store
3088 .prediction_at(&buffer, None, &project, cx)
3089 .expect("should have prediction");
3090 let edits: Vec<_> = prediction
3091 .edits
3092 .iter()
3093 .map(|(range, text)| {
3094 let snapshot = buffer.read(cx).snapshot();
3095 (range.to_offset(&snapshot), text.clone())
3096 })
3097 .collect();
3098 assert_eq!(edits, vec![(5..5, " world".into())]);
3099 });
3100}
3101
3102#[gpui::test]
3103async fn test_v3_prediction_strips_cursor_marker_from_edit_text(cx: &mut TestAppContext) {
3104 let (ep_store, mut requests) = init_test_with_fake_client(cx);
3105 let fs = FakeFs::new(cx.executor());
3106
3107 fs.insert_tree(
3108 "/root",
3109 json!({
3110 "foo.txt": "hello"
3111 }),
3112 )
3113 .await;
3114 let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
3115
3116 let buffer = project
3117 .update(cx, |project, cx| {
3118 let path = project
3119 .find_project_path(path!("root/foo.txt"), cx)
3120 .unwrap();
3121 project.open_buffer(path, cx)
3122 })
3123 .await
3124 .unwrap();
3125
3126 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
3127 let position = snapshot.anchor_before(language::Point::new(0, 5));
3128
3129 ep_store.update(cx, |ep_store, cx| {
3130 ep_store.refresh_prediction_from_buffer(
3131 project.clone(),
3132 buffer.clone(),
3133 position,
3134 EditPredictionRequestTrigger::Other,
3135 cx,
3136 );
3137 });
3138
3139 let (request, respond_tx) = requests.predict.next().await.unwrap();
3140 let excerpt_length = request.input.cursor_excerpt.len();
3141 respond_tx
3142 .send(PredictEditsV3Response {
3143 request_id: Uuid::new_v4().to_string(),
3144 output: "hello world".to_string(),
3145 editable_range: 0..excerpt_length,
3146 model_version: None,
3147 cursor_offset: Some(5),
3148 })
3149 .unwrap();
3150
3151 cx.run_until_parked();
3152
3153 ep_store.update(cx, |ep_store, cx| {
3154 let prediction = ep_store
3155 .prediction_at(&buffer, None, &project, cx)
3156 .expect("should have prediction");
3157 let snapshot = buffer.read(cx).snapshot();
3158 let edits: Vec<_> = prediction
3159 .edits
3160 .iter()
3161 .map(|(range, text)| (range.to_offset(&snapshot), text.clone()))
3162 .collect();
3163
3164 assert_eq!(edits, vec![(5..5, " world".into())]);
3165 });
3166}
3167
3168fn init_test(cx: &mut TestAppContext) {
3169 cx.update(|cx| {
3170 cx.set_global(AppDatabase::test_new());
3171 let settings_store = SettingsStore::test(cx);
3172 cx.set_global(settings_store);
3173 disable_jumps_feature_flag(cx);
3174 });
3175}
3176
3177fn disable_jumps_feature_flag(cx: &mut App) {
3178 SettingsStore::update_global(cx, |store, _| {
3179 store.register_setting::<FeatureFlagsSettings>();
3180 });
3181 set_jumps_feature_flag_override(cx, "off");
3182 cx.update_flags(false, vec![]);
3183}
3184
3185fn set_jumps_feature_flag_override(cx: &mut App, value: &str) {
3186 SettingsStore::update_global(cx, |store, cx| {
3187 store.update_user_settings(cx, |content| {
3188 content.feature_flags.get_or_insert_default().insert(
3189 EditPredictionJumpsFeatureFlag::NAME.to_string(),
3190 value.to_string(),
3191 );
3192 });
3193 });
3194}
3195
3196async fn apply_edit_prediction(
3197 buffer_content: &str,
3198 completion_response: &str,
3199 cx: &mut TestAppContext,
3200) -> String {
3201 let fs = project::FakeFs::new(cx.executor());
3202 let project = Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;
3203 let buffer = cx.new(|cx| Buffer::local(buffer_content, cx));
3204 let (ep_store, response) = make_test_ep_store(&project, cx).await;
3205 *response.lock() = completion_response.to_string();
3206 let edit_prediction = run_edit_prediction(&buffer, &project, &ep_store, cx).await;
3207 buffer.update(cx, |buffer, cx| {
3208 buffer.edit(edit_prediction.edits.iter().cloned(), None, cx)
3209 });
3210 buffer.read_with(cx, |buffer, _| buffer.text())
3211}
3212
3213async fn run_edit_prediction(
3214 buffer: &Entity<Buffer>,
3215 project: &Entity<Project>,
3216 ep_store: &Entity<EditPredictionStore>,
3217 cx: &mut TestAppContext,
3218) -> EditPrediction {
3219 let cursor = buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(1, 0)));
3220 ep_store.update(cx, |ep_store, cx| {
3221 ep_store.register_buffer(buffer, &project, cx)
3222 });
3223 cx.background_executor.run_until_parked();
3224 let prediction_task = ep_store.update(cx, |ep_store, cx| {
3225 ep_store.request_prediction(
3226 &project,
3227 buffer,
3228 cursor,
3229 PredictEditsRequestTrigger::Other,
3230 cx,
3231 )
3232 });
3233 prediction_task.await.unwrap().unwrap().prediction
3234}
3235
3236async fn make_test_ep_store(
3237 project: &Entity<Project>,
3238 cx: &mut TestAppContext,
3239) -> (Entity<EditPredictionStore>, Arc<Mutex<String>>) {
3240 let default_response = "hello world\n".to_string();
3241 let completion_response: Arc<Mutex<String>> = Arc::new(Mutex::new(default_response));
3242 let http_client = FakeHttpClient::create({
3243 let completion_response = completion_response.clone();
3244 let mut next_request_id = 0;
3245 move |req| {
3246 let completion_response = completion_response.clone();
3247 let method = req.method().clone();
3248 let uri = req.uri().path().to_string();
3249 let mut body = req.into_body();
3250 async move {
3251 match (method, uri.as_str()) {
3252 (Method::POST, "/client/llm_tokens") => Ok(http_client::Response::builder()
3253 .status(200)
3254 .body(
3255 serde_json::to_string(&CreateLlmTokenResponse {
3256 token: LlmToken("the-llm-token".to_string()),
3257 })
3258 .unwrap()
3259 .into(),
3260 )
3261 .unwrap()),
3262 (Method::POST, "/predict_edits/v3") => {
3263 let mut buf = Vec::new();
3264 body.read_to_end(&mut buf).await.ok();
3265 let decompressed = zstd::decode_all(&buf[..]).unwrap();
3266 let req: PredictEditsV3Request =
3267 serde_json::from_slice(&decompressed).unwrap();
3268
3269 next_request_id += 1;
3270 Ok(http_client::Response::builder()
3271 .status(200)
3272 .body(
3273 serde_json::to_string(&PredictEditsV3Response {
3274 request_id: format!("request-{next_request_id}"),
3275 editable_range: 0..req.input.cursor_excerpt.len(),
3276 output: completion_response.lock().clone(),
3277 model_version: None,
3278 cursor_offset: None,
3279 })
3280 .unwrap()
3281 .into(),
3282 )
3283 .unwrap())
3284 }
3285 (Method::POST, "/predict_edits/v4") => {
3286 let mut buf = Vec::new();
3287 body.read_to_end(&mut buf).await.ok();
3288 let decompressed = zstd::decode_all(&buf[..]).unwrap();
3289 let req: PredictEditsV4Request =
3290 serde_json::from_slice(&decompressed).unwrap();
3291
3292 next_request_id += 1;
3293 let output = completion_response.lock().clone();
3294 let file = req
3295 .input
3296 .editable_context
3297 .iter()
3298 .find(|file| file.path == req.input.cursor_path)
3299 .or_else(|| req.input.editable_context.first())
3300 .expect("V4 requests should include editable context");
3301 let excerpt = file
3302 .excerpts
3303 .first()
3304 .expect("V4 editable context should include an excerpt");
3305 let diff = unified_diff_with_offsets(
3306 &excerpt.text,
3307 &output,
3308 excerpt.row_range.start,
3309 excerpt.row_range.start,
3310 );
3311 let path = file.path.to_string_lossy();
3312 let response = PredictEditsV4Response {
3313 request_id: format!("request-{next_request_id}"),
3314 patch: format!("--- a/{path}\n+++ b/{path}\n{diff}"),
3315 model_version: None,
3316 };
3317 Ok(http_client::Response::builder()
3318 .status(200)
3319 .body(serde_json::to_string(&response).unwrap().into())
3320 .unwrap())
3321 }
3322 _ => Ok(http_client::Response::builder()
3323 .status(404)
3324 .body("Not Found".to_string().into())
3325 .unwrap()),
3326 }
3327 }
3328 }
3329 });
3330
3331 let client = cx.update(|cx| Client::new(Arc::new(FakeSystemClock::new()), http_client, cx));
3332 let user_store = cx.update(|cx| cx.new(|cx| client::UserStore::new(client.clone(), cx)));
3333 cx.update(|cx| {
3334 RefreshLlmTokenListener::register(client.clone(), user_store.clone(), cx);
3335 });
3336 let _server = FakeServer::for_client(42, &client, cx).await;
3337
3338 let project_user_store = cx.update(|cx| project.read(cx).user_store());
3339 set_test_organization(&project_user_store, cx);
3340
3341 let ep_store = cx.new(|cx| {
3342 let mut ep_store = EditPredictionStore::new(client, project.read(cx).user_store(), cx);
3343 ep_store.set_edit_prediction_model(EditPredictionModel::Zeta);
3344
3345 let worktrees = project.read(cx).worktrees(cx).collect::<Vec<_>>();
3346 for worktree in worktrees {
3347 let worktree_id = worktree.read(cx).id();
3348 ep_store
3349 .get_or_init_project(project, cx)
3350 .license_detection_watchers
3351 .entry(worktree_id)
3352 .or_insert_with(|| Rc::new(LicenseDetectionWatcher::new(&worktree, cx)));
3353 }
3354
3355 ep_store
3356 });
3357
3358 (ep_store, completion_response)
3359}
3360
3361fn to_completion_edits(
3362 iterator: impl IntoIterator<Item = (Range<usize>, Arc<str>)>,
3363 buffer: &Entity<Buffer>,
3364 cx: &App,
3365) -> Vec<(Range<Anchor>, Arc<str>)> {
3366 let buffer = buffer.read(cx);
3367 iterator
3368 .into_iter()
3369 .map(|(range, text)| {
3370 (
3371 buffer.anchor_after(range.start)..buffer.anchor_before(range.end),
3372 text,
3373 )
3374 })
3375 .collect()
3376}
3377
3378fn from_completion_edits(
3379 editor_edits: &[(Range<Anchor>, Arc<str>)],
3380 buffer: &Entity<Buffer>,
3381 cx: &App,
3382) -> Vec<(Range<usize>, Arc<str>)> {
3383 let buffer = buffer.read(cx);
3384 editor_edits
3385 .iter()
3386 .map(|(range, text)| {
3387 (
3388 range.start.to_offset(buffer)..range.end.to_offset(buffer),
3389 text.clone(),
3390 )
3391 })
3392 .collect()
3393}
3394
3395#[gpui::test]
3396async fn test_unauthenticated_without_custom_url_blocks_prediction_impl(cx: &mut TestAppContext) {
3397 init_test(cx);
3398
3399 let fs = FakeFs::new(cx.executor());
3400 fs.insert_tree(
3401 "/project",
3402 serde_json::json!({
3403 "main.rs": "fn main() {\n \n}\n"
3404 }),
3405 )
3406 .await;
3407
3408 let project = Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;
3409
3410 let request_count = Arc::new(std::sync::atomic::AtomicUsize::default());
3411 let http_client = FakeHttpClient::create({
3412 let request_count = request_count.clone();
3413 move |_req| {
3414 request_count.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
3415 async move {
3416 Ok(gpui::http_client::Response::builder()
3417 .status(401)
3418 .body("Unauthorized".into())
3419 .unwrap())
3420 }
3421 }
3422 });
3423
3424 let client =
3425 cx.update(|cx| client::Client::new(Arc::new(FakeSystemClock::new()), http_client, cx));
3426 let user_store = cx.update(|cx| cx.new(|cx| client::UserStore::new(client.clone(), cx)));
3427 cx.update(|cx| {
3428 RefreshLlmTokenListener::register(client.clone(), user_store.clone(), cx);
3429 });
3430
3431 let ep_store = cx.new(|cx| EditPredictionStore::new(client, project.read(cx).user_store(), cx));
3432
3433 let buffer = project
3434 .update(cx, |project, cx| {
3435 let path = project
3436 .find_project_path(path!("/project/main.rs"), cx)
3437 .unwrap();
3438 project.open_buffer(path, cx)
3439 })
3440 .await
3441 .unwrap();
3442
3443 let cursor = buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(1, 4)));
3444 ep_store.update(cx, |ep_store, cx| {
3445 ep_store.register_buffer(&buffer, &project, cx)
3446 });
3447 cx.background_executor.run_until_parked();
3448
3449 let completion_task = ep_store.update(cx, |ep_store, cx| {
3450 ep_store.set_edit_prediction_model(EditPredictionModel::Zeta);
3451 ep_store.request_prediction(
3452 &project,
3453 &buffer,
3454 cursor,
3455 PredictEditsRequestTrigger::Other,
3456 cx,
3457 )
3458 });
3459
3460 assert!(completion_task.await.unwrap().is_none());
3461 assert_eq!(request_count.load(std::sync::atomic::Ordering::SeqCst), 0);
3462}
3463
3464#[gpui::test]
3465async fn test_edit_prediction_settled(cx: &mut TestAppContext) {
3466 let (ep_store, _requests) = init_test_with_fake_client(cx);
3467 let fs = FakeFs::new(cx.executor());
3468
3469 // Buffer with two clearly separated regions:
3470 // Region A = lines 0-9 (offsets 0..50)
3471 // Region B = lines 20-29 (offsets 105..155)
3472 // A big gap in between so edits in one region never overlap the other.
3473 let mut content = String::new();
3474 for i in 0..30 {
3475 content.push_str(&format!("line {i:02}\n"));
3476 }
3477
3478 fs.insert_tree(
3479 "/root",
3480 json!({
3481 "foo.md": content.clone()
3482 }),
3483 )
3484 .await;
3485 let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
3486
3487 let buffer = project
3488 .update(cx, |project, cx| {
3489 let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
3490 project.open_buffer(path, cx)
3491 })
3492 .await
3493 .unwrap();
3494
3495 type SettledEventRecord = (EditPredictionId, String);
3496 let settled_events: Arc<Mutex<Vec<SettledEventRecord>>> = Arc::new(Mutex::new(Vec::new()));
3497
3498 ep_store.update(cx, |ep_store, cx| {
3499 ep_store.register_buffer(&buffer, &project, cx);
3500
3501 let settled_events = settled_events.clone();
3502 ep_store.settled_event_callback = Some(Box::new(move |id, text| {
3503 settled_events.lock().push((id, text));
3504 }));
3505 });
3506
3507 // --- Phase 1: edit in region A and enqueue prediction A ---
3508
3509 buffer.update(cx, |buffer, cx| {
3510 // Edit at the start of line 0.
3511 buffer.edit(vec![(0..0, "ADDED ")], None, cx);
3512 });
3513 cx.run_until_parked();
3514
3515 let snapshot_a = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
3516 let empty_edits: Arc<[(Range<Anchor>, Arc<str>)]> = Vec::new().into();
3517 let edit_preview_a = buffer
3518 .read_with(cx, |buffer, cx| {
3519 buffer.preview_edits(empty_edits.clone(), cx)
3520 })
3521 .await;
3522
3523 // Region A: first 10 lines of the buffer.
3524 let editable_region_a = 0..snapshot_a.point_to_offset(Point::new(10, 0));
3525
3526 ep_store.update(cx, |ep_store, cx| {
3527 ep_store.enqueue_settled_prediction(
3528 EditPredictionId("prediction-a".into()),
3529 &project,
3530 &buffer,
3531 &snapshot_a,
3532 editable_region_a.clone(),
3533 &edit_preview_a,
3534 None,
3535 None,
3536 None,
3537 Duration::from_secs(0),
3538 cx,
3539 );
3540 });
3541
3542 // --- Phase 2: repeatedly edit in region A to keep it unsettled ---
3543
3544 // Let the worker process the channel message before we start advancing.
3545 cx.run_until_parked();
3546
3547 for region_a_edit_offset in (5..).take(3) {
3548 // Edit inside region A (not at the boundary) so `last_edit_at` is
3549 // updated before the worker's next wake.
3550 buffer.update(cx, |buffer, cx| {
3551 buffer.edit(
3552 vec![(region_a_edit_offset..region_a_edit_offset, "x")],
3553 None,
3554 cx,
3555 );
3556 });
3557 cx.run_until_parked();
3558
3559 cx.executor()
3560 .advance_clock(EDIT_PREDICTION_SETTLED_QUIESCENCE / 2);
3561 cx.run_until_parked();
3562 assert!(
3563 settled_events.lock().is_empty(),
3564 "no settled events should fire while region A is still being edited"
3565 );
3566 }
3567
3568 // Still nothing settled.
3569 assert!(settled_events.lock().is_empty());
3570
3571 // --- Phase 3: edit in distinct region B, enqueue prediction B ---
3572 // Advance a small amount so B's quiescence window starts later than A's,
3573 // but not so much that A settles (A's last edit was at the start of
3574 // iteration 3, and it needs a full Q to settle).
3575 cx.executor()
3576 .advance_clock(EDIT_PREDICTION_SETTLED_QUIESCENCE / 4);
3577 cx.run_until_parked();
3578 assert!(settled_events.lock().is_empty());
3579
3580 let snapshot_b = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
3581 let line_20_offset = snapshot_b.point_to_offset(Point::new(20, 0));
3582
3583 buffer.update(cx, |buffer, cx| {
3584 buffer.edit(vec![(line_20_offset..line_20_offset, "NEW ")], None, cx);
3585 });
3586 cx.run_until_parked();
3587
3588 let snapshot_b2 = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
3589 let edit_preview_b = buffer
3590 .read_with(cx, |buffer, cx| buffer.preview_edits(empty_edits, cx))
3591 .await;
3592 let editable_region_b = line_20_offset..snapshot_b2.point_to_offset(Point::new(25, 0));
3593
3594 ep_store.update(cx, |ep_store, cx| {
3595 ep_store.enqueue_settled_prediction(
3596 EditPredictionId("prediction-b".into()),
3597 &project,
3598 &buffer,
3599 &snapshot_b2,
3600 editable_region_b.clone(),
3601 &edit_preview_b,
3602 None,
3603 None,
3604 None,
3605 Duration::from_secs(0),
3606 cx,
3607 );
3608 });
3609
3610 cx.run_until_parked();
3611 assert!(
3612 settled_events.lock().is_empty(),
3613 "neither prediction should have settled yet"
3614 );
3615
3616 // --- Phase 4: let enough time pass for region A to settle ---
3617 // A's last edit was at T_a (during the last loop iteration). The worker is
3618 // sleeping until T_a + Q. We advance just enough to reach that wake time
3619 // (Q/4 since we already advanced Q/4 in phase 3 on top of the loop's
3620 // 3*Q/2). At that point A has been quiet for Q and settles, but B was
3621 // enqueued only Q/4 ago and stays pending.
3622 cx.executor()
3623 .advance_clock(EDIT_PREDICTION_SETTLED_QUIESCENCE / 4);
3624 cx.run_until_parked();
3625
3626 {
3627 let events = settled_events.lock().clone();
3628 assert_eq!(
3629 events.len(),
3630 1,
3631 "prediction and capture_sample for A should have settled, got: {events:?}"
3632 );
3633 assert_eq!(events[0].0, EditPredictionId("prediction-a".into()));
3634 }
3635
3636 // --- Phase 5: let more time pass for region B to settle ---
3637 // B's last edit was Q/4 before A settled. The worker rescheduled to
3638 // B's last_edit_at + Q, which is 3Q/4 from now.
3639 cx.executor()
3640 .advance_clock(EDIT_PREDICTION_SETTLED_QUIESCENCE * 3 / 4);
3641 cx.run_until_parked();
3642
3643 {
3644 let events = settled_events.lock().clone();
3645 assert_eq!(
3646 events.len(),
3647 2,
3648 "both prediction and capture_sample settled events should be emitted for each request, got: {events:?}"
3649 );
3650 assert_eq!(events[1].0, EditPredictionId("prediction-b".into()));
3651 }
3652}
3653
3654const MIT_LICENSE: &str = indoc! {r#"
3655 MIT License
3656
3657 Permission is hereby granted, free of charge, to any person obtaining a copy
3658 of this software and associated documentation files (the "Software"), to deal
3659 in the Software without restriction, including without limitation the rights
3660 to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
3661 copies of the Software, and to permit persons to whom the Software is
3662 furnished to do so, subject to the following conditions:
3663
3664 The above copyright notice and this permission notice shall be included in all
3665 copies or substantial portions of the Software.
3666
3667 THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
3668 IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
3669 FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
3670 AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
3671 LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
3672 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
3673 SOFTWARE.
3674"#};
3675
3676async fn init_sample_capture_test(
3677 tree: serde_json::Value,
3678 cx: &mut TestAppContext,
3679) -> (
3680 Entity<EditPredictionStore>,
3681 RequestChannels,
3682 Entity<Project>,
3683 Entity<Buffer>,
3684) {
3685 let (ep_store, requests) = init_test_with_fake_client(cx);
3686 let fs = FakeFs::new(cx.executor());
3687 fs.insert_tree("/root", tree).await;
3688 let project = Project::test(fs, vec![path!("/root").as_ref()], cx).await;
3689 let buffer = project
3690 .update(cx, |project, cx| {
3691 let path = project.find_project_path(path!("root/foo.md"), cx).unwrap();
3692 project.open_buffer(path, cx)
3693 })
3694 .await
3695 .unwrap();
3696 ep_store.update(cx, |ep_store, cx| {
3697 ep_store.register_buffer(&buffer, &project, cx);
3698 });
3699 cx.run_until_parked();
3700 (ep_store, requests, project, buffer)
3701}
3702
3703/// Enqueues a settled prediction capture with injected sample data, as the
3704/// real request path would when data collection is enabled. Returns the
3705/// editable offset range.
3706async fn enqueue_sample_capture(
3707 ep_store: &Entity<EditPredictionStore>,
3708 project: &Entity<Project>,
3709 buffer: &Entity<Buffer>,
3710 id: &str,
3711 editable_point_range: Range<Point>,
3712 context: CapturedPredictionContext,
3713 prompt_history_boundary: Option<PromptHistoryBoundary>,
3714 navigation_history: VecDeque<RecentFile>,
3715 cx: &mut TestAppContext,
3716) -> Range<usize> {
3717 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot());
3718 let editable_offset_range = snapshot.point_to_offset(editable_point_range.start)
3719 ..snapshot.point_to_offset(editable_point_range.end);
3720 let empty_edits: Arc<[(Range<Anchor>, Arc<str>)]> = Vec::new().into();
3721 let edit_preview = buffer
3722 .read_with(cx, |buffer, cx| buffer.preview_edits(empty_edits, cx))
3723 .await;
3724 ep_store.update(cx, |ep_store, cx| {
3725 ep_store.enqueue_settled_prediction(
3726 EditPredictionId(id.to_string().into()),
3727 project,
3728 buffer,
3729 &snapshot,
3730 editable_offset_range.clone(),
3731 &edit_preview,
3732 None,
3733 None,
3734 None,
3735 Duration::from_secs(0),
3736 cx,
3737 );
3738 let pending_capture = ep_store
3739 .projects
3740 .get_mut(&project.entity_id())
3741 .unwrap()
3742 .pending_prediction_captures
3743 .last_mut()
3744 .unwrap();
3745 pending_capture.can_collect_data = true;
3746 pending_capture.sample_data = Some(PendingPredictionCaptureSampleData {
3747 context_task: Task::ready(Ok(context)),
3748 editable_path: snapshot.file().unwrap().path().as_std_path().into(),
3749 editable_offset_range: editable_offset_range.clone(),
3750 next_edit_cursor_offset: None,
3751 future_edit_history_events: Vec::new(),
3752 navigation_history,
3753 edit_events_before_quiescence: 0,
3754 prompt_history_boundary,
3755 });
3756 });
3757 editable_offset_range
3758}
3759
3760#[gpui::test]
3761async fn test_edit_prediction_settled_sends_sample_data_after_quiescence(cx: &mut TestAppContext) {
3762 let (ep_store, mut requests, project, buffer) = init_sample_capture_test(
3763 json!({
3764 "LICENSE": MIT_LICENSE,
3765 "foo.md": (0..60).map(|ix| format!("line {ix}\n")).collect::<String>(),
3766 }),
3767 cx,
3768 )
3769 .await;
3770
3771 buffer.update(cx, |buffer, cx| {
3772 let offset = Point::new(1, 0).to_offset(buffer);
3773 buffer.edit(vec![(offset..offset, "prompted ")], None, cx);
3774 });
3775 cx.run_until_parked();
3776
3777 let boundary = ep_store.update(cx, |ep_store, _cx| {
3778 let project_state = ep_store.projects.get(&project.entity_id()).unwrap();
3779 PromptHistoryBoundary {
3780 first_event_seq: project_state
3781 .last_event
3782 .as_ref()
3783 .map_or(project_state.next_last_event_seq, |last_event| {
3784 last_event.seq
3785 }),
3786 snapshot: project_state
3787 .last_event
3788 .as_ref()
3789 .map(|last_event| last_event.new_snapshot.clone()),
3790 }
3791 });
3792 let editable_offset_range = enqueue_sample_capture(
3793 &ep_store,
3794 &project,
3795 &buffer,
3796 "prediction-sample",
3797 Point::new(2, 0)..Point::new(3, 0),
3798 CapturedPredictionContext {
3799 repository_url: Some("https://example.com/repo.git".to_string()),
3800 revision: Some("abc123".to_string()),
3801 uncommitted_diff: Some("--- a/foo.md\n+++ b/foo.md\n".to_string()),
3802 buffer_diagnostics: vec![zeta_prompt::ActiveBufferDiagnostic {
3803 severity: Some(1),
3804 message: "sample diagnostic".to_string(),
3805 snippet: String::new(),
3806 snippet_buffer_row_range: 0..0,
3807 diagnostic_range_in_snippet: 0..0,
3808 }],
3809 editable_context: vec![zeta_prompt::RelatedFile {
3810 path: Path::new("foo.md").into(),
3811 max_row: 60,
3812 excerpts: vec![zeta_prompt::RelatedExcerpt {
3813 row_range: 0..2,
3814 text: "line 0\nline 1\n".into(),
3815 order: 0,
3816 context_source: zeta_prompt::ContextSource::CurrentFile,
3817 }],
3818 in_open_source_repo: true,
3819 }],
3820 },
3821 Some(boundary),
3822 VecDeque::from([RecentFile {
3823 path: Path::new("foo.md").into(),
3824 cursor_position: Some(3),
3825 }]),
3826 cx,
3827 )
3828 .await;
3829 let expected_next_edit_cursor_offset = editable_offset_range.start;
3830
3831 // A second capture whose user has data collection disabled; its sample
3832 // and settled region must be redacted at send time.
3833 let boundary = ep_store.update(cx, |ep_store, _cx| {
3834 let project_state = ep_store.projects.get(&project.entity_id()).unwrap();
3835 PromptHistoryBoundary {
3836 first_event_seq: project_state
3837 .last_event
3838 .as_ref()
3839 .map_or(project_state.next_last_event_seq, |last_event| {
3840 last_event.seq
3841 }),
3842 snapshot: project_state
3843 .last_event
3844 .as_ref()
3845 .map(|last_event| last_event.new_snapshot.clone()),
3846 }
3847 });
3848 enqueue_sample_capture(
3849 &ep_store,
3850 &project,
3851 &buffer,
3852 "prediction-redacted",
3853 Point::new(2, 0)..Point::new(3, 0),
3854 CapturedPredictionContext {
3855 repository_url: None,
3856 revision: None,
3857 uncommitted_diff: None,
3858 buffer_diagnostics: Vec::new(),
3859 editable_context: Vec::new(),
3860 },
3861 Some(boundary),
3862 VecDeque::new(),
3863 cx,
3864 )
3865 .await;
3866 ep_store.update(cx, |ep_store, _cx| {
3867 ep_store
3868 .projects
3869 .get_mut(&project.entity_id())
3870 .unwrap()
3871 .pending_prediction_captures
3872 .last_mut()
3873 .unwrap()
3874 .can_collect_data = false;
3875 });
3876
3877 cx.executor().advance_clock(LAST_CHANGE_GROUPING_TIME * 2);
3878 cx.run_until_parked();
3879
3880 for (ix, row) in [2, 20, 30, 40, 50].into_iter().enumerate() {
3881 buffer.update(cx, |buffer, cx| {
3882 let start = Point::new(row, 0).to_offset(buffer);
3883 let end = Point::new(row, 4).to_offset(buffer);
3884 buffer.edit(vec![(start..end, format!("future {ix}"))], None, cx);
3885 });
3886 cx.run_until_parked();
3887 }
3888
3889 cx.executor()
3890 .advance_clock(EDIT_PREDICTION_SETTLED_QUIESCENCE);
3891 cx.run_until_parked();
3892
3893 let mut settled_by_id = std::collections::HashMap::new();
3894 for _ in 0..2 {
3895 let request = requests
3896 .settled
3897 .next()
3898 .await
3899 .expect("settled request should be sent");
3900 settled_by_id.insert(request.request_id.clone(), request);
3901 }
3902
3903 let redacted_request = settled_by_id.remove("prediction-redacted").unwrap();
3904 assert!(!redacted_request.can_collect_data);
3905 assert_eq!(redacted_request.settled_editable_region, None);
3906 assert_eq!(redacted_request.sample_data, None);
3907
3908 let settled_request = settled_by_id.remove("prediction-sample").unwrap();
3909 let sample_data = settled_request
3910 .sample_data
3911 .expect("sample data should be sent after quiescence");
3912 assert_eq!(
3913 sample_data.repository_url.as_deref(),
3914 Some("https://example.com/repo.git")
3915 );
3916 assert_eq!(sample_data.revision.as_deref(), Some("abc123"));
3917 assert_eq!(
3918 sample_data.uncommitted_diff.as_deref(),
3919 Some("--- a/foo.md\n+++ b/foo.md\n")
3920 );
3921 assert_eq!(sample_data.editable_path.as_ref(), Path::new("foo.md"));
3922 assert_eq!(sample_data.editable_offset_range, editable_offset_range);
3923 assert_eq!(sample_data.buffer_diagnostics.len(), 1);
3924 assert_eq!(sample_data.editable_context.len(), 1);
3925 let editable_context = &sample_data.editable_context[0];
3926 assert_eq!(editable_context.path.as_ref(), Path::new("foo.md"));
3927 assert_eq!(editable_context.excerpts.len(), 1);
3928 assert_eq!(
3929 editable_context.excerpts[0].context_source,
3930 zeta_prompt::ContextSource::CurrentFile
3931 );
3932 assert_eq!(sample_data.future_edit_history_events.len(), 4);
3933 assert_eq!(sample_data.navigation_history.len(), 1);
3934 assert_eq!(sample_data.edit_events_before_quiescence, 5);
3935 assert_eq!(
3936 sample_data.next_edit_cursor_offset,
3937 Some(expected_next_edit_cursor_offset)
3938 );
3939
3940 let future_event_diffs = sample_data
3941 .future_edit_history_events
3942 .iter()
3943 .map(|event| match event.as_ref() {
3944 zeta_prompt::Event::BufferChange { diff, .. } => diff.as_str(),
3945 })
3946 .collect::<Vec<_>>();
3947 assert!(future_event_diffs.iter().all(|diff| {
3948 diff.lines()
3949 .all(|line| !line.starts_with("+prompted ") && line != "-line 1")
3950 }));
3951 assert!(
3952 future_event_diffs
3953 .iter()
3954 .any(|diff| diff.contains("future 0"))
3955 );
3956 assert!(
3957 !future_event_diffs
3958 .iter()
3959 .any(|diff| diff.contains("future 4"))
3960 );
3961}
3962
3963#[gpui::test]
3964async fn test_edit_prediction_settled_sample_data_requires_observing_all_events_since_request(
3965 cx: &mut TestAppContext,
3966) {
3967 let (ep_store, mut requests, project, buffer) = init_sample_capture_test(
3968 json!({
3969 "LICENSE": MIT_LICENSE,
3970 "foo.md": (0..30).map(|ix| format!("line {ix}\n")).collect::<String>(),
3971 }),
3972 cx,
3973 )
3974 .await;
3975
3976 // Two predictions are requested while no event is pending.
3977 let (boundary_observed, boundary_missed) = ep_store.update(cx, |ep_store, _cx| {
3978 let project_state = ep_store.projects.get(&project.entity_id()).unwrap();
3979 let first_event_seq = project_state
3980 .last_event
3981 .as_ref()
3982 .map_or(project_state.next_last_event_seq, |last_event| {
3983 last_event.seq
3984 });
3985 let snapshot = project_state
3986 .last_event
3987 .as_ref()
3988 .map(|last_event| last_event.new_snapshot.clone());
3989 (
3990 PromptHistoryBoundary {
3991 first_event_seq,
3992 snapshot: snapshot.clone(),
3993 },
3994 PromptHistoryBoundary {
3995 first_event_seq,
3996 snapshot,
3997 },
3998 )
3999 });
4000 assert!(boundary_observed.snapshot.is_none());
4001
4002 // The first capture is enqueued immediately, so it observes both
4003 // subsequent events.
4004 enqueue_sample_capture(
4005 &ep_store,
4006 &project,
4007 &buffer,
4008 "prediction-observed",
4009 Point::new(1, 0)..Point::new(2, 0),
4010 CapturedPredictionContext {
4011 repository_url: None,
4012 revision: None,
4013 uncommitted_diff: None,
4014 buffer_diagnostics: Vec::new(),
4015 editable_context: Vec::new(),
4016 },
4017 Some(boundary_observed),
4018 VecDeque::new(),
4019 cx,
4020 )
4021 .await;
4022
4023 buffer.update(cx, |buffer, cx| {
4024 let offset = Point::new(1, 0).to_offset(buffer);
4025 buffer.edit(vec![(offset..offset, "first ")], None, cx);
4026 });
4027 cx.run_until_parked();
4028 buffer.update(cx, |buffer, cx| {
4029 let offset = Point::new(20, 0).to_offset(buffer);
4030 buffer.edit(vec![(offset..offset, "second ")], None, cx);
4031 });
4032 cx.run_until_parked();
4033
4034 // The second capture is enqueued only after the first event was
4035 // finalized, so its future history has a gap and it must be dropped.
4036 enqueue_sample_capture(
4037 &ep_store,
4038 &project,
4039 &buffer,
4040 "prediction-missed",
4041 Point::new(1, 0)..Point::new(2, 0),
4042 CapturedPredictionContext {
4043 repository_url: None,
4044 revision: None,
4045 uncommitted_diff: None,
4046 buffer_diagnostics: Vec::new(),
4047 editable_context: Vec::new(),
4048 },
4049 Some(boundary_missed),
4050 VecDeque::new(),
4051 cx,
4052 )
4053 .await;
4054
4055 cx.executor()
4056 .advance_clock(EDIT_PREDICTION_SETTLED_QUIESCENCE);
4057 cx.run_until_parked();
4058
4059 let mut settled_by_id = std::collections::HashMap::new();
4060 for _ in 0..2 {
4061 let request = requests
4062 .settled
4063 .next()
4064 .await
4065 .expect("settled request should be sent");
4066 settled_by_id.insert(request.request_id.clone(), request);
4067 }
4068
4069 let observed_request = settled_by_id.remove("prediction-observed").unwrap();
4070 let sample_data = observed_request
4071 .sample_data
4072 .expect("sample data should be sent");
4073 assert_eq!(sample_data.edit_events_before_quiescence, 2);
4074 assert_eq!(sample_data.future_edit_history_events.len(), 2);
4075
4076 let missed_request = settled_by_id.remove("prediction-missed").unwrap();
4077 assert_eq!(missed_request.sample_data, None);
4078}
4079
4080#[gpui::test]
4081async fn test_edit_prediction_settled_drops_future_events_when_their_oss_status_is_unknown(
4082 cx: &mut TestAppContext,
4083) {
4084 let (ep_store, mut requests, project, buffer) = init_sample_capture_test(
4085 json!({ "foo.md": (0..30).map(|ix| format!("line {ix}\n")).collect::<String>() }),
4086 cx,
4087 )
4088 .await;
4089
4090 enqueue_sample_capture(
4091 &ep_store,
4092 &project,
4093 &buffer,
4094 "prediction-non-oss",
4095 Point::new(1, 0)..Point::new(2, 0),
4096 CapturedPredictionContext {
4097 repository_url: None,
4098 revision: None,
4099 uncommitted_diff: None,
4100 buffer_diagnostics: Vec::new(),
4101 editable_context: Vec::new(),
4102 },
4103 None,
4104 VecDeque::new(),
4105 cx,
4106 )
4107 .await;
4108
4109 // There is no LICENSE file, so this event's open-source status is unknown
4110 // and the sample must be dropped.
4111 buffer.update(cx, |buffer, cx| {
4112 let offset = Point::new(20, 0).to_offset(buffer);
4113 buffer.edit(vec![(offset..offset, "outside ")], None, cx);
4114 });
4115 cx.run_until_parked();
4116
4117 cx.executor()
4118 .advance_clock(EDIT_PREDICTION_SETTLED_QUIESCENCE);
4119 cx.run_until_parked();
4120
4121 let settled_request = requests
4122 .settled
4123 .next()
4124 .await
4125 .expect("settled request should be sent");
4126 assert_eq!(settled_request.sample_data, None);
4127}
4128
4129#[gpui::test]
4130fn test_buffer_path_with_id_fallback_for_untitled_buffers(cx: &mut TestAppContext) {
4131 let buffer_1 = cx.new(|cx| Buffer::local("one", cx));
4132 let buffer_2 = cx.new(|cx| Buffer::local("two", cx));
4133
4134 let snapshot_1 = buffer_1.read_with(cx, |buffer, _| buffer.text_snapshot());
4135 let snapshot_2 = buffer_2.read_with(cx, |buffer, _| buffer.text_snapshot());
4136
4137 let path_1 = cx.read(|cx| buffer_path_with_id_fallback(None, &snapshot_1, cx));
4138 let path_2 = cx.read(|cx| buffer_path_with_id_fallback(None, &snapshot_2, cx));
4139
4140 assert_eq!(
4141 path_1.as_ref(),
4142 Path::new(&format!("untitled-{}", snapshot_1.remote_id()))
4143 );
4144 assert_eq!(
4145 path_2.as_ref(),
4146 Path::new(&format!("untitled-{}", snapshot_2.remote_id()))
4147 );
4148 assert_ne!(path_1.as_ref(), path_2.as_ref());
4149}
4150
4151#[gpui::test]
4152async fn test_data_collection_disabled_by_default(cx: &mut TestAppContext) {
4153 let (ep_store, _channels) = init_test_with_fake_client(cx);
4154
4155 cx.update(|cx| {
4156 assert!(!ep_store.read(cx).is_data_collection_enabled(cx));
4157 });
4158}
4159
4160#[gpui::test]
4161async fn test_data_collection_enabled_via_legacy_kv_store(cx: &mut TestAppContext) {
4162 let (ep_store, _channels) =
4163 init_test_with_fake_client_and_legacy_data_collection(cx, Some("true"));
4164
4165 cx.update(|cx| {
4166 assert!(ep_store.read(cx).is_data_collection_enabled(cx));
4167 });
4168}
4169
4170#[gpui::test]
4171async fn test_data_collection_default_uses_cached_legacy_value(cx: &mut TestAppContext) {
4172 let (ep_store, _channels) =
4173 init_test_with_fake_client_and_legacy_data_collection(cx, Some("true"));
4174
4175 cx.update(|cx| {
4176 assert!(ep_store.read(cx).is_data_collection_enabled(cx));
4177 });
4178
4179 cx.update(|cx| KeyValueStore::global(cx))
4180 .delete_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE.into())
4181 .await
4182 .unwrap();
4183
4184 cx.update(|cx| {
4185 assert!(ep_store.read(cx).is_data_collection_enabled(cx));
4186 });
4187}
4188
4189#[gpui::test]
4190async fn test_data_collection_setting_overrides_kv_store(cx: &mut TestAppContext) {
4191 let (ep_store, _channels) =
4192 init_test_with_fake_client_and_legacy_data_collection(cx, Some("true"));
4193
4194 // An explicit false in settings.json wins over the KV store.
4195 cx.update_global::<SettingsStore, _>(|settings, cx| {
4196 settings.update_user_settings(cx, |content| {
4197 content
4198 .project
4199 .all_languages
4200 .edit_predictions
4201 .get_or_insert_default()
4202 .allow_data_collection = Some(EditPredictionDataCollectionChoice::No);
4203 });
4204 });
4205
4206 cx.update(|cx| {
4207 assert!(!ep_store.read(cx).is_data_collection_enabled(cx));
4208 });
4209}
4210
4211#[gpui::test]
4212async fn test_data_collection_enabled_via_setting(cx: &mut TestAppContext) {
4213 let (ep_store, _channels) = init_test_with_fake_client(cx);
4214
4215 cx.update_global::<SettingsStore, _>(|settings, cx| {
4216 settings.update_user_settings(cx, |content| {
4217 content
4218 .project
4219 .all_languages
4220 .edit_predictions
4221 .get_or_insert_default()
4222 .allow_data_collection = Some(EditPredictionDataCollectionChoice::Yes);
4223 });
4224 });
4225
4226 cx.update(|cx| {
4227 assert!(ep_store.read(cx).is_data_collection_enabled(cx));
4228 });
4229}
4230
4231#[gpui::test]
4232async fn test_data_collection_always_enabled_for_staff(cx: &mut TestAppContext) {
4233 let (ep_store, _channels) = init_test_with_fake_client(cx);
4234
4235 cx.update(|cx| {
4236 cx.set_staff(true);
4237 assert!(ep_store.read(cx).is_data_collection_enabled(cx));
4238 });
4239}
4240
4241#[gpui::test]
4242async fn test_data_collection_disabled_by_organization_configuration(cx: &mut TestAppContext) {
4243 let (ep_store, _channels) = init_test_with_fake_client(cx);
4244
4245 cx.update_global::<SettingsStore, _>(|settings, cx| {
4246 settings.update_user_settings(cx, |content| {
4247 content
4248 .project
4249 .all_languages
4250 .edit_predictions
4251 .get_or_insert_default()
4252 .allow_data_collection = Some(EditPredictionDataCollectionChoice::Yes);
4253 });
4254 });
4255
4256 let user_store = cx.update(|cx| ep_store.read(cx).user_store.clone());
4257 cx.update(|cx| {
4258 user_store.update(cx, |user_store, cx| {
4259 user_store.set_current_organization_configuration_for_test(
4260 Arc::new(Organization {
4261 id: OrganizationId("org-1".into()),
4262 name: "Org 1".into(),
4263 is_personal: false,
4264 }),
4265 OrganizationConfiguration {
4266 is_zed_model_provider_enabled: true,
4267 is_agent_thread_feedback_enabled: true,
4268 is_collaboration_enabled: true,
4269 edit_prediction: OrganizationEditPredictionConfiguration {
4270 is_enabled: true,
4271 is_feedback_enabled: false,
4272 },
4273 },
4274 cx,
4275 );
4276 });
4277
4278 assert!(!ep_store.read(cx).is_data_collection_enabled(cx));
4279 });
4280}
4281
4282// When a user had data collection enabled via the legacy KV store (with no explicit
4283// setting in settings.json), toggle_data_collection must read the *resolved* state
4284// (true) and write Some(false).
4285#[gpui::test]
4286async fn test_toggle_data_collection_from_kv_enabled_state(cx: &mut TestAppContext) {
4287 let (ep_store, _channels) =
4288 init_test_with_fake_client_and_legacy_data_collection(cx, Some("true"));
4289
4290 cx.update(|cx| {
4291 assert!(
4292 ep_store.read(cx).is_data_collection_enabled(cx),
4293 "data collection should be enabled via KV store before toggle"
4294 );
4295 });
4296
4297 // Simulate what toggle_data_collection does: capture the resolved current
4298 // state, then write its inverse.
4299 let is_currently_enabled = cx.update(|cx| ep_store.read(cx).is_data_collection_enabled(cx));
4300 cx.update_global::<SettingsStore, _>(|settings, cx| {
4301 settings.update_user_settings(cx, |content| {
4302 content
4303 .project
4304 .all_languages
4305 .edit_predictions
4306 .get_or_insert_default()
4307 .allow_data_collection = Some(if is_currently_enabled {
4308 EditPredictionDataCollectionChoice::No
4309 } else {
4310 EditPredictionDataCollectionChoice::Yes
4311 });
4312 });
4313 });
4314
4315 cx.update(|cx| {
4316 assert!(
4317 !ep_store.read(cx).is_data_collection_enabled(cx),
4318 "data collection should be disabled after toggling off from KV-enabled state"
4319 );
4320 });
4321}
4322
4323#[gpui::test]
4324async fn test_upsell_shown_by_default(cx: &mut TestAppContext) {
4325 init_test(cx);
4326 let kvp = cx.update(|cx| KeyValueStore::global(cx));
4327 kvp.delete_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE.into())
4328 .await
4329 .ok();
4330 kvp.delete_kvp(ZedPredictUpsell::KEY.into()).await.ok();
4331
4332 cx.update(|cx| assert!(should_show_upsell_modal(cx)));
4333}
4334
4335#[gpui::test]
4336async fn test_upsell_dismissed_when_data_collection_choice_in_kv_store(cx: &mut TestAppContext) {
4337 init_test(cx);
4338
4339 // Any value for the data collection key means the old upsell was already
4340 // shown, regardless of whether data collection was accepted or declined.
4341 for value in &["true", "false"] {
4342 cx.update(|cx| KeyValueStore::global(cx))
4343 .write_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE.into(), value.to_string())
4344 .await
4345 .unwrap();
4346
4347 cx.update(|cx| {
4348 assert!(
4349 !should_show_upsell_modal(cx),
4350 "upsell should be suppressed when data collection choice is '{value}'"
4351 );
4352 });
4353 }
4354
4355 cx.update(|cx| KeyValueStore::global(cx))
4356 .delete_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE.into())
4357 .await
4358 .unwrap();
4359}
4360
4361#[gpui::test]
4362async fn test_upsell_dismissed_when_dismissed_key_set(cx: &mut TestAppContext) {
4363 init_test(cx);
4364 let kvp = cx.update(|cx| KeyValueStore::global(cx));
4365 kvp.delete_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE.into())
4366 .await
4367 .ok();
4368 kvp.write_kvp(ZedPredictUpsell::KEY.into(), "1".into())
4369 .await
4370 .unwrap();
4371
4372 cx.update(|cx| assert!(!should_show_upsell_modal(cx)));
4373
4374 kvp.delete_kvp(ZedPredictUpsell::KEY.into()).await.unwrap();
4375}
4376
4377#[gpui::test]
4378async fn test_upsell_dismissed_via_dismissable_api(cx: &mut TestAppContext) {
4379 init_test(cx);
4380 let kvp = cx.update(|cx| KeyValueStore::global(cx));
4381 kvp.delete_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE.into())
4382 .await
4383 .ok();
4384 kvp.delete_kvp(ZedPredictUpsell::KEY.into()).await.ok();
4385
4386 cx.update(|cx| {
4387 assert!(should_show_upsell_modal(cx));
4388 ZedPredictUpsell::set_dismissed(true, cx);
4389 });
4390 cx.run_until_parked();
4391
4392 cx.update(|cx| assert!(!should_show_upsell_modal(cx)));
4393
4394 kvp.delete_kvp(ZedPredictUpsell::KEY.into()).await.unwrap();
4395}
4396
4397#[ctor::ctor(unsafe)]
4398fn init_logger() {
4399 zlog::init_test();
4400}
4401