1mod completion_diff_element;
2mod init;
3mod input_excerpt;
4mod license_detection;
5mod onboarding_modal;
6mod onboarding_telemetry;
7mod rate_completion_modal;
8
9pub(crate) use completion_diff_element::*;
10use db::kvp::{Dismissable, KEY_VALUE_STORE};
11use edit_prediction::DataCollectionState;
12pub use init::*;
13use license_detection::LicenseDetectionWatcher;
14pub use rate_completion_modal::*;
15
16use anyhow::{Context as _, Result, anyhow};
17use arrayvec::ArrayVec;
18use client::{Client, EditPredictionUsage, UserStore};
19use cloud_llm_client::{
20 AcceptEditPredictionBody, EXPIRED_LLM_TOKEN_HEADER_NAME, MINIMUM_REQUIRED_VERSION_HEADER_NAME,
21 PredictEditsBody, PredictEditsGitInfo, PredictEditsResponse, ZED_VERSION_HEADER_NAME,
22};
23use collections::{HashMap, HashSet, VecDeque};
24use futures::AsyncReadExt;
25use gpui::{
26 App, AppContext as _, AsyncApp, Context, Entity, EntityId, Global, SemanticVersion,
27 SharedString, Subscription, Task, actions,
28};
29use http_client::{AsyncBody, HttpClient, Method, Request, Response};
30use input_excerpt::excerpt_for_cursor_position;
31use language::{
32 Anchor, Buffer, BufferSnapshot, EditPreview, File, OffsetRangeExt, ToOffset, ToPoint, text_diff,
33};
34use language_model::{LlmApiToken, RefreshLlmTokenListener};
35use project::{Project, ProjectPath};
36use release_channel::AppVersion;
37use settings::WorktreeId;
38use std::collections::hash_map;
39use std::mem;
40use std::str::FromStr;
41use std::{
42 cmp,
43 fmt::Write,
44 future::Future,
45 ops::Range,
46 path::Path,
47 rc::Rc,
48 sync::Arc,
49 time::{Duration, Instant},
50};
51use telemetry_events::EditPredictionRating;
52use thiserror::Error;
53use util::ResultExt;
54use util::rel_path::RelPath;
55use uuid::Uuid;
56use workspace::notifications::{ErrorMessagePrompt, NotificationId, show_app_notification};
57use worktree::Worktree;
58
59const CURSOR_MARKER: &str = "<|user_cursor_is_here|>";
60const START_OF_FILE_MARKER: &str = "<|start_of_file|>";
61const EDITABLE_REGION_START_MARKER: &str = "<|editable_region_start|>";
62const EDITABLE_REGION_END_MARKER: &str = "<|editable_region_end|>";
63const BUFFER_CHANGE_GROUPING_INTERVAL: Duration = Duration::from_secs(1);
64const ZED_PREDICT_DATA_COLLECTION_CHOICE: &str = "zed_predict_data_collection_choice";
65
66const MAX_CONTEXT_TOKENS: usize = 150;
67const MAX_REWRITE_TOKENS: usize = 350;
68const MAX_EVENT_TOKENS: usize = 500;
69
70/// Maximum number of events to track.
71const MAX_EVENT_COUNT: usize = 16;
72
73actions!(
74 edit_prediction,
75 [
76 /// Clears the edit prediction history.
77 ClearHistory
78 ]
79);
80
81#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
82pub struct EditPredictionId(Uuid);
83
84impl From<EditPredictionId> for gpui::ElementId {
85 fn from(value: EditPredictionId) -> Self {
86 gpui::ElementId::Uuid(value.0)
87 }
88}
89
90impl std::fmt::Display for EditPredictionId {
91 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
92 write!(f, "{}", self.0)
93 }
94}
95
96struct ZedPredictUpsell;
97
98impl Dismissable for ZedPredictUpsell {
99 const KEY: &'static str = "dismissed-edit-predict-upsell";
100
101 fn dismissed() -> bool {
102 // To make this backwards compatible with older versions of Zed, we
103 // check if the user has seen the previous Edit Prediction Onboarding
104 // before, by checking the data collection choice which was written to
105 // the database once the user clicked on "Accept and Enable"
106 if KEY_VALUE_STORE
107 .read_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE)
108 .log_err()
109 .is_some_and(|s| s.is_some())
110 {
111 return true;
112 }
113
114 KEY_VALUE_STORE
115 .read_kvp(Self::KEY)
116 .log_err()
117 .is_some_and(|s| s.is_some())
118 }
119}
120
121pub fn should_show_upsell_modal() -> bool {
122 !ZedPredictUpsell::dismissed()
123}
124
125#[derive(Clone)]
126struct ZetaGlobal(Entity<Zeta>);
127
128impl Global for ZetaGlobal {}
129
130#[derive(Clone)]
131pub struct EditPrediction {
132 id: EditPredictionId,
133 path: Arc<Path>,
134 excerpt_range: Range<usize>,
135 cursor_offset: usize,
136 edits: Arc<[(Range<Anchor>, Arc<str>)]>,
137 snapshot: BufferSnapshot,
138 edit_preview: EditPreview,
139 input_outline: Arc<str>,
140 input_events: Arc<str>,
141 input_excerpt: Arc<str>,
142 output_excerpt: Arc<str>,
143 buffer_snapshotted_at: Instant,
144 response_received_at: Instant,
145}
146
147impl EditPrediction {
148 fn latency(&self) -> Duration {
149 self.response_received_at
150 .duration_since(self.buffer_snapshotted_at)
151 }
152
153 fn interpolate(&self, new_snapshot: &BufferSnapshot) -> Option<Vec<(Range<Anchor>, Arc<str>)>> {
154 edit_prediction::interpolate_edits(&self.snapshot, new_snapshot, &self.edits)
155 }
156}
157
158impl std::fmt::Debug for EditPrediction {
159 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
160 f.debug_struct("EditPrediction")
161 .field("id", &self.id)
162 .field("path", &self.path)
163 .field("edits", &self.edits)
164 .finish_non_exhaustive()
165 }
166}
167
168pub struct Zeta {
169 projects: HashMap<EntityId, ZetaProject>,
170 client: Arc<Client>,
171 shown_completions: VecDeque<EditPrediction>,
172 rated_completions: HashSet<EditPredictionId>,
173 data_collection_choice: DataCollectionChoice,
174 llm_token: LlmApiToken,
175 _llm_token_subscription: Subscription,
176 /// Whether an update to a newer version of Zed is required to continue using Zeta.
177 update_required: bool,
178 user_store: Entity<UserStore>,
179 license_detection_watchers: HashMap<WorktreeId, Rc<LicenseDetectionWatcher>>,
180}
181
182struct ZetaProject {
183 events: VecDeque<Event>,
184 registered_buffers: HashMap<gpui::EntityId, RegisteredBuffer>,
185}
186
187impl Zeta {
188 pub fn global(cx: &mut App) -> Option<Entity<Self>> {
189 cx.try_global::<ZetaGlobal>().map(|global| global.0.clone())
190 }
191
192 pub fn register(
193 worktree: Option<Entity<Worktree>>,
194 client: Arc<Client>,
195 user_store: Entity<UserStore>,
196 cx: &mut App,
197 ) -> Entity<Self> {
198 let this = Self::global(cx).unwrap_or_else(|| {
199 let entity = cx.new(|cx| Self::new(client, user_store, cx));
200 cx.set_global(ZetaGlobal(entity.clone()));
201 entity
202 });
203
204 this.update(cx, move |this, cx| {
205 if let Some(worktree) = worktree {
206 let worktree_id = worktree.read(cx).id();
207 this.license_detection_watchers
208 .entry(worktree_id)
209 .or_insert_with(|| Rc::new(LicenseDetectionWatcher::new(&worktree, cx)));
210 }
211 });
212
213 this
214 }
215
216 pub fn clear_history(&mut self) {
217 for zeta_project in self.projects.values_mut() {
218 zeta_project.events.clear();
219 }
220 }
221
222 pub fn usage(&self, cx: &App) -> Option<EditPredictionUsage> {
223 self.user_store.read(cx).edit_prediction_usage()
224 }
225
226 fn new(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut Context<Self>) -> Self {
227 let refresh_llm_token_listener = RefreshLlmTokenListener::global(cx);
228 let data_collection_choice = Self::load_data_collection_choice();
229 Self {
230 projects: HashMap::default(),
231 client,
232 shown_completions: VecDeque::new(),
233 rated_completions: HashSet::default(),
234 data_collection_choice,
235 llm_token: LlmApiToken::default(),
236 _llm_token_subscription: cx.subscribe(
237 &refresh_llm_token_listener,
238 |this, _listener, _event, cx| {
239 let client = this.client.clone();
240 let llm_token = this.llm_token.clone();
241 cx.spawn(async move |_this, _cx| {
242 llm_token.refresh(&client).await?;
243 anyhow::Ok(())
244 })
245 .detach_and_log_err(cx);
246 },
247 ),
248 update_required: false,
249 license_detection_watchers: HashMap::default(),
250 user_store,
251 }
252 }
253
254 fn get_or_init_zeta_project(
255 &mut self,
256 project: &Entity<Project>,
257 cx: &mut Context<Self>,
258 ) -> &mut ZetaProject {
259 let project_id = project.entity_id();
260 match self.projects.entry(project_id) {
261 hash_map::Entry::Occupied(entry) => entry.into_mut(),
262 hash_map::Entry::Vacant(entry) => {
263 cx.observe_release(project, move |this, _, _cx| {
264 this.projects.remove(&project_id);
265 })
266 .detach();
267 entry.insert(ZetaProject {
268 events: VecDeque::with_capacity(MAX_EVENT_COUNT),
269 registered_buffers: HashMap::default(),
270 })
271 }
272 }
273 }
274
275 fn push_event(zeta_project: &mut ZetaProject, event: Event) {
276 let events = &mut zeta_project.events;
277
278 if let Some(Event::BufferChange {
279 new_snapshot: last_new_snapshot,
280 timestamp: last_timestamp,
281 ..
282 }) = events.back_mut()
283 {
284 // Coalesce edits for the same buffer when they happen one after the other.
285 let Event::BufferChange {
286 old_snapshot,
287 new_snapshot,
288 timestamp,
289 } = &event;
290
291 if timestamp.duration_since(*last_timestamp) <= BUFFER_CHANGE_GROUPING_INTERVAL
292 && old_snapshot.remote_id() == last_new_snapshot.remote_id()
293 && old_snapshot.version == last_new_snapshot.version
294 {
295 *last_new_snapshot = new_snapshot.clone();
296 *last_timestamp = *timestamp;
297 return;
298 }
299 }
300
301 if events.len() >= MAX_EVENT_COUNT {
302 // These are halved instead of popping to improve prompt caching.
303 events.drain(..MAX_EVENT_COUNT / 2);
304 }
305
306 events.push_back(event);
307 }
308
309 pub fn register_buffer(
310 &mut self,
311 buffer: &Entity<Buffer>,
312 project: &Entity<Project>,
313 cx: &mut Context<Self>,
314 ) {
315 let zeta_project = self.get_or_init_zeta_project(project, cx);
316 Self::register_buffer_impl(zeta_project, buffer, project, cx);
317 }
318
319 fn register_buffer_impl<'a>(
320 zeta_project: &'a mut ZetaProject,
321 buffer: &Entity<Buffer>,
322 project: &Entity<Project>,
323 cx: &mut Context<Self>,
324 ) -> &'a mut RegisteredBuffer {
325 let buffer_id = buffer.entity_id();
326 match zeta_project.registered_buffers.entry(buffer_id) {
327 hash_map::Entry::Occupied(entry) => entry.into_mut(),
328 hash_map::Entry::Vacant(entry) => {
329 let snapshot = buffer.read(cx).snapshot();
330 let project_entity_id = project.entity_id();
331 entry.insert(RegisteredBuffer {
332 snapshot,
333 _subscriptions: [
334 cx.subscribe(buffer, {
335 let project = project.downgrade();
336 move |this, buffer, event, cx| {
337 if let language::BufferEvent::Edited = event
338 && let Some(project) = project.upgrade()
339 {
340 this.report_changes_for_buffer(&buffer, &project, cx);
341 }
342 }
343 }),
344 cx.observe_release(buffer, move |this, _buffer, _cx| {
345 let Some(zeta_project) = this.projects.get_mut(&project_entity_id)
346 else {
347 return;
348 };
349 zeta_project.registered_buffers.remove(&buffer_id);
350 }),
351 ],
352 })
353 }
354 }
355 }
356
357 fn request_completion_impl<F, R>(
358 &mut self,
359 project: &Entity<Project>,
360 buffer: &Entity<Buffer>,
361 cursor: language::Anchor,
362 cx: &mut Context<Self>,
363 perform_predict_edits: F,
364 ) -> Task<Result<Option<EditPrediction>>>
365 where
366 F: FnOnce(PerformPredictEditsParams) -> R + 'static,
367 R: Future<Output = Result<(PredictEditsResponse, Option<EditPredictionUsage>)>>
368 + Send
369 + 'static,
370 {
371 let buffer = buffer.clone();
372 let buffer_snapshotted_at = Instant::now();
373 let snapshot = self.report_changes_for_buffer(&buffer, project, cx);
374 let zeta = cx.entity();
375 let client = self.client.clone();
376 let llm_token = self.llm_token.clone();
377 let app_version = AppVersion::global(cx);
378
379 let zeta_project = self.get_or_init_zeta_project(project, cx);
380 let mut events = Vec::with_capacity(zeta_project.events.len());
381 events.extend(zeta_project.events.iter().cloned());
382 let events = Arc::new(events);
383
384 let (git_info, can_collect_file) = if let Some(file) = snapshot.file() {
385 let can_collect_file = self.can_collect_file(file, cx);
386 let git_info = if can_collect_file {
387 git_info_for_file(project, &ProjectPath::from_file(file.as_ref(), cx), cx)
388 } else {
389 None
390 };
391 (git_info, can_collect_file)
392 } else {
393 (None, false)
394 };
395
396 let full_path: Arc<Path> = snapshot
397 .file()
398 .map(|f| Arc::from(f.full_path(cx).as_path()))
399 .unwrap_or_else(|| Arc::from(Path::new("untitled")));
400 let full_path_str = full_path.to_string_lossy().into_owned();
401 let cursor_point = cursor.to_point(&snapshot);
402 let cursor_offset = cursor_point.to_offset(&snapshot);
403 let prompt_for_events = {
404 let events = events.clone();
405 move || prompt_for_events_impl(&events, MAX_EVENT_TOKENS)
406 };
407 let gather_task = gather_context(
408 full_path_str,
409 &snapshot,
410 cursor_point,
411 prompt_for_events,
412 cx,
413 );
414
415 cx.spawn(async move |this, cx| {
416 let GatherContextOutput {
417 mut body,
418 editable_range,
419 included_events_count,
420 } = gather_task.await?;
421 let done_gathering_context_at = Instant::now();
422
423 let included_events = &events[events.len() - included_events_count..events.len()];
424 body.can_collect_data = can_collect_file
425 && this
426 .read_with(cx, |this, cx| this.can_collect_events(included_events, cx))
427 .unwrap_or(false);
428 if body.can_collect_data {
429 body.git_info = git_info;
430 }
431
432 log::debug!(
433 "Events:\n{}\nExcerpt:\n{:?}",
434 body.input_events,
435 body.input_excerpt
436 );
437
438 let input_outline = body.outline.clone().unwrap_or_default();
439 let input_events = body.input_events.clone();
440 let input_excerpt = body.input_excerpt.clone();
441
442 let response = perform_predict_edits(PerformPredictEditsParams {
443 client,
444 llm_token,
445 app_version,
446 body,
447 })
448 .await;
449 let (response, usage) = match response {
450 Ok(response) => response,
451 Err(err) => {
452 if err.is::<ZedUpdateRequiredError>() {
453 cx.update(|cx| {
454 zeta.update(cx, |zeta, _cx| {
455 zeta.update_required = true;
456 });
457
458 let error_message: SharedString = err.to_string().into();
459 show_app_notification(
460 NotificationId::unique::<ZedUpdateRequiredError>(),
461 cx,
462 move |cx| {
463 cx.new(|cx| {
464 ErrorMessagePrompt::new(error_message.clone(), cx)
465 .with_link_button(
466 "Update Zed",
467 "https://zed.dev/releases",
468 )
469 })
470 },
471 );
472 })
473 .ok();
474 }
475
476 return Err(err);
477 }
478 };
479
480 let received_response_at = Instant::now();
481 log::debug!("completion response: {}", &response.output_excerpt);
482
483 if let Some(usage) = usage {
484 this.update(cx, |this, cx| {
485 this.user_store.update(cx, |user_store, cx| {
486 user_store.update_edit_prediction_usage(usage, cx);
487 });
488 })
489 .ok();
490 }
491
492 let edit_prediction = Self::process_completion_response(
493 response,
494 buffer,
495 &snapshot,
496 editable_range,
497 cursor_offset,
498 full_path,
499 input_outline,
500 input_events,
501 input_excerpt,
502 buffer_snapshotted_at,
503 cx,
504 )
505 .await;
506
507 let finished_at = Instant::now();
508
509 // record latency for ~1% of requests
510 if rand::random::<u8>() <= 2 {
511 telemetry::event!(
512 "Edit Prediction Request",
513 context_latency = done_gathering_context_at
514 .duration_since(buffer_snapshotted_at)
515 .as_millis(),
516 request_latency = received_response_at
517 .duration_since(done_gathering_context_at)
518 .as_millis(),
519 process_latency = finished_at.duration_since(received_response_at).as_millis()
520 );
521 }
522
523 edit_prediction
524 })
525 }
526
527 #[cfg(any(test, feature = "test-support"))]
528 pub fn fake_completion(
529 &mut self,
530 project: &Entity<Project>,
531 buffer: &Entity<Buffer>,
532 position: language::Anchor,
533 response: PredictEditsResponse,
534 cx: &mut Context<Self>,
535 ) -> Task<Result<Option<EditPrediction>>> {
536 self.request_completion_impl(project, buffer, position, cx, |_params| {
537 std::future::ready(Ok((response, None)))
538 })
539 }
540
541 pub fn request_completion(
542 &mut self,
543 project: &Entity<Project>,
544 buffer: &Entity<Buffer>,
545 position: language::Anchor,
546 cx: &mut Context<Self>,
547 ) -> Task<Result<Option<EditPrediction>>> {
548 self.request_completion_impl(project, buffer, position, cx, Self::perform_predict_edits)
549 }
550
551 pub fn perform_predict_edits(
552 params: PerformPredictEditsParams,
553 ) -> impl Future<Output = Result<(PredictEditsResponse, Option<EditPredictionUsage>)>> {
554 async move {
555 let PerformPredictEditsParams {
556 client,
557 llm_token,
558 app_version,
559 body,
560 ..
561 } = params;
562
563 let http_client = client.http_client();
564 let mut token = llm_token.acquire(&client).await?;
565 let mut did_retry = false;
566
567 loop {
568 let request_builder = http_client::Request::builder().method(Method::POST);
569 let request_builder =
570 if let Ok(predict_edits_url) = std::env::var("ZED_PREDICT_EDITS_URL") {
571 request_builder.uri(predict_edits_url)
572 } else {
573 request_builder.uri(
574 http_client
575 .build_zed_llm_url("/predict_edits/v2", &[])?
576 .as_ref(),
577 )
578 };
579 let request = request_builder
580 .header("Content-Type", "application/json")
581 .header("Authorization", format!("Bearer {}", token))
582 .header(ZED_VERSION_HEADER_NAME, app_version.to_string())
583 .body(serde_json::to_string(&body)?.into())?;
584
585 let mut response = http_client.send(request).await?;
586
587 if let Some(minimum_required_version) = response
588 .headers()
589 .get(MINIMUM_REQUIRED_VERSION_HEADER_NAME)
590 .and_then(|version| SemanticVersion::from_str(version.to_str().ok()?).ok())
591 {
592 anyhow::ensure!(
593 app_version >= minimum_required_version,
594 ZedUpdateRequiredError {
595 minimum_version: minimum_required_version
596 }
597 );
598 }
599
600 if response.status().is_success() {
601 let usage = EditPredictionUsage::from_headers(response.headers()).ok();
602
603 let mut body = String::new();
604 response.body_mut().read_to_string(&mut body).await?;
605 return Ok((serde_json::from_str(&body)?, usage));
606 } else if !did_retry
607 && response
608 .headers()
609 .get(EXPIRED_LLM_TOKEN_HEADER_NAME)
610 .is_some()
611 {
612 did_retry = true;
613 token = llm_token.refresh(&client).await?;
614 } else {
615 let mut body = String::new();
616 response.body_mut().read_to_string(&mut body).await?;
617 anyhow::bail!(
618 "error predicting edits.\nStatus: {:?}\nBody: {}",
619 response.status(),
620 body
621 );
622 }
623 }
624 }
625 }
626
627 fn accept_edit_prediction(
628 &mut self,
629 request_id: EditPredictionId,
630 cx: &mut Context<Self>,
631 ) -> Task<Result<()>> {
632 let client = self.client.clone();
633 let llm_token = self.llm_token.clone();
634 let app_version = AppVersion::global(cx);
635 cx.spawn(async move |this, cx| {
636 let http_client = client.http_client();
637 let mut response = llm_token_retry(&llm_token, &client, |token| {
638 let request_builder = http_client::Request::builder().method(Method::POST);
639 let request_builder =
640 if let Ok(accept_prediction_url) = std::env::var("ZED_ACCEPT_PREDICTION_URL") {
641 request_builder.uri(accept_prediction_url)
642 } else {
643 request_builder.uri(
644 http_client
645 .build_zed_llm_url("/predict_edits/accept", &[])?
646 .as_ref(),
647 )
648 };
649 Ok(request_builder
650 .header("Content-Type", "application/json")
651 .header("Authorization", format!("Bearer {}", token))
652 .header(ZED_VERSION_HEADER_NAME, app_version.to_string())
653 .body(
654 serde_json::to_string(&AcceptEditPredictionBody {
655 request_id: request_id.0.to_string(),
656 })?
657 .into(),
658 )?)
659 })
660 .await?;
661
662 if let Some(minimum_required_version) = response
663 .headers()
664 .get(MINIMUM_REQUIRED_VERSION_HEADER_NAME)
665 .and_then(|version| SemanticVersion::from_str(version.to_str().ok()?).ok())
666 && app_version < minimum_required_version
667 {
668 return Err(anyhow!(ZedUpdateRequiredError {
669 minimum_version: minimum_required_version
670 }));
671 }
672
673 if response.status().is_success() {
674 if let Some(usage) = EditPredictionUsage::from_headers(response.headers()).ok() {
675 this.update(cx, |this, cx| {
676 this.user_store.update(cx, |user_store, cx| {
677 user_store.update_edit_prediction_usage(usage, cx);
678 });
679 })?;
680 }
681
682 Ok(())
683 } else {
684 let mut body = String::new();
685 response.body_mut().read_to_string(&mut body).await?;
686 Err(anyhow!(
687 "error accepting edit prediction.\nStatus: {:?}\nBody: {}",
688 response.status(),
689 body
690 ))
691 }
692 })
693 }
694
695 fn process_completion_response(
696 prediction_response: PredictEditsResponse,
697 buffer: Entity<Buffer>,
698 snapshot: &BufferSnapshot,
699 editable_range: Range<usize>,
700 cursor_offset: usize,
701 path: Arc<Path>,
702 input_outline: String,
703 input_events: String,
704 input_excerpt: String,
705 buffer_snapshotted_at: Instant,
706 cx: &AsyncApp,
707 ) -> Task<Result<Option<EditPrediction>>> {
708 let snapshot = snapshot.clone();
709 let request_id = prediction_response.request_id;
710 let output_excerpt = prediction_response.output_excerpt;
711 cx.spawn(async move |cx| {
712 let output_excerpt: Arc<str> = output_excerpt.into();
713
714 let edits: Arc<[(Range<Anchor>, Arc<str>)]> = cx
715 .background_spawn({
716 let output_excerpt = output_excerpt.clone();
717 let editable_range = editable_range.clone();
718 let snapshot = snapshot.clone();
719 async move { Self::parse_edits(output_excerpt, editable_range, &snapshot) }
720 })
721 .await?
722 .into();
723
724 let Some((edits, snapshot, edit_preview)) = buffer.read_with(cx, {
725 let edits = edits.clone();
726 move |buffer, cx| {
727 let new_snapshot = buffer.snapshot();
728 let edits: Arc<[(Range<Anchor>, Arc<str>)]> =
729 edit_prediction::interpolate_edits(&snapshot, &new_snapshot, &edits)?
730 .into();
731 Some((edits.clone(), new_snapshot, buffer.preview_edits(edits, cx)))
732 }
733 })?
734 else {
735 return anyhow::Ok(None);
736 };
737
738 let request_id = Uuid::from_str(&request_id).context("failed to parse request id")?;
739
740 let edit_preview = edit_preview.await;
741
742 Ok(Some(EditPrediction {
743 id: EditPredictionId(request_id),
744 path,
745 excerpt_range: editable_range,
746 cursor_offset,
747 edits,
748 edit_preview,
749 snapshot,
750 input_outline: input_outline.into(),
751 input_events: input_events.into(),
752 input_excerpt: input_excerpt.into(),
753 output_excerpt,
754 buffer_snapshotted_at,
755 response_received_at: Instant::now(),
756 }))
757 })
758 }
759
760 fn parse_edits(
761 output_excerpt: Arc<str>,
762 editable_range: Range<usize>,
763 snapshot: &BufferSnapshot,
764 ) -> Result<Vec<(Range<Anchor>, Arc<str>)>> {
765 let content = output_excerpt.replace(CURSOR_MARKER, "");
766
767 let start_markers = content
768 .match_indices(EDITABLE_REGION_START_MARKER)
769 .collect::<Vec<_>>();
770 anyhow::ensure!(
771 start_markers.len() == 1,
772 "expected exactly one start marker, found {}",
773 start_markers.len()
774 );
775
776 let end_markers = content
777 .match_indices(EDITABLE_REGION_END_MARKER)
778 .collect::<Vec<_>>();
779 anyhow::ensure!(
780 end_markers.len() == 1,
781 "expected exactly one end marker, found {}",
782 end_markers.len()
783 );
784
785 let sof_markers = content
786 .match_indices(START_OF_FILE_MARKER)
787 .collect::<Vec<_>>();
788 anyhow::ensure!(
789 sof_markers.len() <= 1,
790 "expected at most one start-of-file marker, found {}",
791 sof_markers.len()
792 );
793
794 let codefence_start = start_markers[0].0;
795 let content = &content[codefence_start..];
796
797 let newline_ix = content.find('\n').context("could not find newline")?;
798 let content = &content[newline_ix + 1..];
799
800 let codefence_end = content
801 .rfind(&format!("\n{EDITABLE_REGION_END_MARKER}"))
802 .context("could not find end marker")?;
803 let new_text = &content[..codefence_end];
804
805 let old_text = snapshot
806 .text_for_range(editable_range.clone())
807 .collect::<String>();
808
809 Ok(Self::compute_edits(
810 old_text,
811 new_text,
812 editable_range.start,
813 snapshot,
814 ))
815 }
816
817 pub fn compute_edits(
818 old_text: String,
819 new_text: &str,
820 offset: usize,
821 snapshot: &BufferSnapshot,
822 ) -> Vec<(Range<Anchor>, Arc<str>)> {
823 text_diff(&old_text, new_text)
824 .into_iter()
825 .map(|(mut old_range, new_text)| {
826 old_range.start += offset;
827 old_range.end += offset;
828
829 let prefix_len = common_prefix(
830 snapshot.chars_for_range(old_range.clone()),
831 new_text.chars(),
832 );
833 old_range.start += prefix_len;
834
835 let suffix_len = common_prefix(
836 snapshot.reversed_chars_for_range(old_range.clone()),
837 new_text[prefix_len..].chars().rev(),
838 );
839 old_range.end = old_range.end.saturating_sub(suffix_len);
840
841 let new_text = new_text[prefix_len..new_text.len() - suffix_len].into();
842 let range = if old_range.is_empty() {
843 let anchor = snapshot.anchor_after(old_range.start);
844 anchor..anchor
845 } else {
846 snapshot.anchor_after(old_range.start)..snapshot.anchor_before(old_range.end)
847 };
848 (range, new_text)
849 })
850 .collect()
851 }
852
853 pub fn is_completion_rated(&self, completion_id: EditPredictionId) -> bool {
854 self.rated_completions.contains(&completion_id)
855 }
856
857 pub fn completion_shown(&mut self, completion: &EditPrediction, cx: &mut Context<Self>) {
858 self.shown_completions.push_front(completion.clone());
859 if self.shown_completions.len() > 50 {
860 let completion = self.shown_completions.pop_back().unwrap();
861 self.rated_completions.remove(&completion.id);
862 }
863 cx.notify();
864 }
865
866 pub fn rate_completion(
867 &mut self,
868 completion: &EditPrediction,
869 rating: EditPredictionRating,
870 feedback: String,
871 cx: &mut Context<Self>,
872 ) {
873 self.rated_completions.insert(completion.id);
874 telemetry::event!(
875 "Edit Prediction Rated",
876 rating,
877 input_events = completion.input_events,
878 input_excerpt = completion.input_excerpt,
879 input_outline = completion.input_outline,
880 output_excerpt = completion.output_excerpt,
881 feedback
882 );
883 self.client.telemetry().flush_events().detach();
884 cx.notify();
885 }
886
887 pub fn shown_completions(&self) -> impl DoubleEndedIterator<Item = &EditPrediction> {
888 self.shown_completions.iter()
889 }
890
891 pub fn shown_completions_len(&self) -> usize {
892 self.shown_completions.len()
893 }
894
895 fn report_changes_for_buffer(
896 &mut self,
897 buffer: &Entity<Buffer>,
898 project: &Entity<Project>,
899 cx: &mut Context<Self>,
900 ) -> BufferSnapshot {
901 let zeta_project = self.get_or_init_zeta_project(project, cx);
902 let registered_buffer = Self::register_buffer_impl(zeta_project, buffer, project, cx);
903
904 let new_snapshot = buffer.read(cx).snapshot();
905 if new_snapshot.version != registered_buffer.snapshot.version {
906 let old_snapshot = mem::replace(&mut registered_buffer.snapshot, new_snapshot.clone());
907 Self::push_event(
908 zeta_project,
909 Event::BufferChange {
910 old_snapshot,
911 new_snapshot: new_snapshot.clone(),
912 timestamp: Instant::now(),
913 },
914 );
915 }
916
917 new_snapshot
918 }
919
920 fn can_collect_file(&self, file: &Arc<dyn File>, cx: &App) -> bool {
921 self.data_collection_choice.is_enabled() && self.is_file_open_source(file, cx)
922 }
923
924 fn can_collect_events(&self, events: &[Event], cx: &App) -> bool {
925 if !self.data_collection_choice.is_enabled() {
926 return false;
927 }
928 let mut last_checked_file = None;
929 for event in events {
930 match event {
931 Event::BufferChange {
932 old_snapshot,
933 new_snapshot,
934 ..
935 } => {
936 if let Some(old_file) = old_snapshot.file()
937 && let Some(new_file) = new_snapshot.file()
938 {
939 if let Some(last_checked_file) = last_checked_file
940 && Arc::ptr_eq(last_checked_file, old_file)
941 && Arc::ptr_eq(last_checked_file, new_file)
942 {
943 continue;
944 }
945 if !self.can_collect_file(old_file, cx) {
946 return false;
947 }
948 if !Arc::ptr_eq(old_file, new_file) && !self.can_collect_file(new_file, cx)
949 {
950 return false;
951 }
952 last_checked_file = Some(new_file);
953 } else {
954 return false;
955 }
956 }
957 }
958 }
959 true
960 }
961
962 fn is_file_open_source(&self, file: &Arc<dyn File>, cx: &App) -> bool {
963 if !file.is_local() || file.is_private() {
964 return false;
965 }
966 self.license_detection_watchers
967 .get(&file.worktree_id(cx))
968 .is_some_and(|watcher| watcher.is_project_open_source())
969 }
970
971 fn load_data_collection_choice() -> DataCollectionChoice {
972 let choice = KEY_VALUE_STORE
973 .read_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE)
974 .log_err()
975 .flatten();
976
977 match choice.as_deref() {
978 Some("true") => DataCollectionChoice::Enabled,
979 Some("false") => DataCollectionChoice::Disabled,
980 Some(_) => {
981 log::error!("unknown value in '{ZED_PREDICT_DATA_COLLECTION_CHOICE}'");
982 DataCollectionChoice::NotAnswered
983 }
984 None => DataCollectionChoice::NotAnswered,
985 }
986 }
987
988 fn toggle_data_collection_choice(&mut self, cx: &mut Context<Self>) {
989 self.data_collection_choice = self.data_collection_choice.toggle();
990 let new_choice = self.data_collection_choice;
991 db::write_and_log(cx, move || {
992 KEY_VALUE_STORE.write_kvp(
993 ZED_PREDICT_DATA_COLLECTION_CHOICE.into(),
994 new_choice.is_enabled().to_string(),
995 )
996 });
997 }
998}
999
1000pub struct PerformPredictEditsParams {
1001 pub client: Arc<Client>,
1002 pub llm_token: LlmApiToken,
1003 pub app_version: SemanticVersion,
1004 pub body: PredictEditsBody,
1005}
1006
1007#[derive(Error, Debug)]
1008#[error(
1009 "You must update to Zed version {minimum_version} or higher to continue using edit predictions."
1010)]
1011pub struct ZedUpdateRequiredError {
1012 minimum_version: SemanticVersion,
1013}
1014
1015fn common_prefix<T1: Iterator<Item = char>, T2: Iterator<Item = char>>(a: T1, b: T2) -> usize {
1016 a.zip(b)
1017 .take_while(|(a, b)| a == b)
1018 .map(|(a, _)| a.len_utf8())
1019 .sum()
1020}
1021
1022fn git_info_for_file(
1023 project: &Entity<Project>,
1024 project_path: &ProjectPath,
1025 cx: &App,
1026) -> Option<PredictEditsGitInfo> {
1027 let git_store = project.read(cx).git_store().read(cx);
1028 if let Some((repository, _repo_path)) =
1029 git_store.repository_and_path_for_project_path(project_path, cx)
1030 {
1031 let repository = repository.read(cx);
1032 let head_sha = repository
1033 .head_commit
1034 .as_ref()
1035 .map(|head_commit| head_commit.sha.to_string());
1036 let remote_origin_url = repository.remote_origin_url.clone();
1037 let remote_upstream_url = repository.remote_upstream_url.clone();
1038 if head_sha.is_none() && remote_origin_url.is_none() && remote_upstream_url.is_none() {
1039 return None;
1040 }
1041 Some(PredictEditsGitInfo {
1042 head_sha,
1043 remote_origin_url,
1044 remote_upstream_url,
1045 })
1046 } else {
1047 None
1048 }
1049}
1050
1051pub struct GatherContextOutput {
1052 pub body: PredictEditsBody,
1053 pub editable_range: Range<usize>,
1054 pub included_events_count: usize,
1055}
1056
1057pub fn gather_context(
1058 full_path_str: String,
1059 snapshot: &BufferSnapshot,
1060 cursor_point: language::Point,
1061 prompt_for_events: impl FnOnce() -> (String, usize) + Send + 'static,
1062 cx: &App,
1063) -> Task<Result<GatherContextOutput>> {
1064 cx.background_spawn({
1065 let snapshot = snapshot.clone();
1066 async move {
1067 let input_excerpt = excerpt_for_cursor_position(
1068 cursor_point,
1069 &full_path_str,
1070 &snapshot,
1071 MAX_REWRITE_TOKENS,
1072 MAX_CONTEXT_TOKENS,
1073 );
1074 let (input_events, included_events_count) = prompt_for_events();
1075 let editable_range = input_excerpt.editable_range.to_offset(&snapshot);
1076
1077 let body = PredictEditsBody {
1078 input_events,
1079 input_excerpt: input_excerpt.prompt,
1080 can_collect_data: false,
1081 diagnostic_groups: None,
1082 git_info: None,
1083 outline: None,
1084 speculated_output: None,
1085 };
1086
1087 Ok(GatherContextOutput {
1088 body,
1089 editable_range,
1090 included_events_count,
1091 })
1092 }
1093 })
1094}
1095
1096fn prompt_for_events_impl(events: &[Event], mut remaining_tokens: usize) -> (String, usize) {
1097 let mut result = String::new();
1098 for (ix, event) in events.iter().rev().enumerate() {
1099 let event_string = event.to_prompt();
1100 let event_tokens = guess_token_count(event_string.len());
1101 if event_tokens > remaining_tokens {
1102 return (result, ix);
1103 }
1104
1105 if !result.is_empty() {
1106 result.insert_str(0, "\n\n");
1107 }
1108 result.insert_str(0, &event_string);
1109 remaining_tokens -= event_tokens;
1110 }
1111 return (result, events.len());
1112}
1113
1114struct RegisteredBuffer {
1115 snapshot: BufferSnapshot,
1116 _subscriptions: [gpui::Subscription; 2],
1117}
1118
1119#[derive(Clone)]
1120pub enum Event {
1121 BufferChange {
1122 old_snapshot: BufferSnapshot,
1123 new_snapshot: BufferSnapshot,
1124 timestamp: Instant,
1125 },
1126}
1127
1128impl Event {
1129 fn to_prompt(&self) -> String {
1130 match self {
1131 Event::BufferChange {
1132 old_snapshot,
1133 new_snapshot,
1134 ..
1135 } => {
1136 let mut prompt = String::new();
1137
1138 let old_path = old_snapshot
1139 .file()
1140 .map(|f| f.path().as_ref())
1141 .unwrap_or(RelPath::unix("untitled").unwrap());
1142 let new_path = new_snapshot
1143 .file()
1144 .map(|f| f.path().as_ref())
1145 .unwrap_or(RelPath::unix("untitled").unwrap());
1146 if old_path != new_path {
1147 writeln!(prompt, "User renamed {:?} to {:?}\n", old_path, new_path).unwrap();
1148 }
1149
1150 let diff = language::unified_diff(&old_snapshot.text(), &new_snapshot.text());
1151 if !diff.is_empty() {
1152 write!(
1153 prompt,
1154 "User edited {:?}:\n```diff\n{}\n```",
1155 new_path, diff
1156 )
1157 .unwrap();
1158 }
1159
1160 prompt
1161 }
1162 }
1163 }
1164}
1165
1166#[derive(Debug, Clone)]
1167struct CurrentEditPrediction {
1168 buffer_id: EntityId,
1169 completion: EditPrediction,
1170}
1171
1172impl CurrentEditPrediction {
1173 fn should_replace_completion(&self, old_completion: &Self, snapshot: &BufferSnapshot) -> bool {
1174 if self.buffer_id != old_completion.buffer_id {
1175 return true;
1176 }
1177
1178 let Some(old_edits) = old_completion.completion.interpolate(snapshot) else {
1179 return true;
1180 };
1181 let Some(new_edits) = self.completion.interpolate(snapshot) else {
1182 return false;
1183 };
1184
1185 if old_edits.len() == 1 && new_edits.len() == 1 {
1186 let (old_range, old_text) = &old_edits[0];
1187 let (new_range, new_text) = &new_edits[0];
1188 new_range == old_range && new_text.starts_with(old_text.as_ref())
1189 } else {
1190 true
1191 }
1192 }
1193}
1194
1195struct PendingCompletion {
1196 id: usize,
1197 _task: Task<()>,
1198}
1199
1200#[derive(Debug, Clone, Copy)]
1201pub enum DataCollectionChoice {
1202 NotAnswered,
1203 Enabled,
1204 Disabled,
1205}
1206
1207impl DataCollectionChoice {
1208 pub fn is_enabled(self) -> bool {
1209 match self {
1210 Self::Enabled => true,
1211 Self::NotAnswered | Self::Disabled => false,
1212 }
1213 }
1214
1215 pub fn is_answered(self) -> bool {
1216 match self {
1217 Self::Enabled | Self::Disabled => true,
1218 Self::NotAnswered => false,
1219 }
1220 }
1221
1222 #[must_use]
1223 pub fn toggle(&self) -> DataCollectionChoice {
1224 match self {
1225 Self::Enabled => Self::Disabled,
1226 Self::Disabled => Self::Enabled,
1227 Self::NotAnswered => Self::Enabled,
1228 }
1229 }
1230}
1231
1232impl From<bool> for DataCollectionChoice {
1233 fn from(value: bool) -> Self {
1234 match value {
1235 true => DataCollectionChoice::Enabled,
1236 false => DataCollectionChoice::Disabled,
1237 }
1238 }
1239}
1240
1241async fn llm_token_retry(
1242 llm_token: &LlmApiToken,
1243 client: &Arc<Client>,
1244 build_request: impl Fn(String) -> Result<Request<AsyncBody>>,
1245) -> Result<Response<AsyncBody>> {
1246 let mut did_retry = false;
1247 let http_client = client.http_client();
1248 let mut token = llm_token.acquire(client).await?;
1249 loop {
1250 let request = build_request(token.clone())?;
1251 let response = http_client.send(request).await?;
1252
1253 if !did_retry
1254 && !response.status().is_success()
1255 && response
1256 .headers()
1257 .get(EXPIRED_LLM_TOKEN_HEADER_NAME)
1258 .is_some()
1259 {
1260 did_retry = true;
1261 token = llm_token.refresh(client).await?;
1262 continue;
1263 }
1264
1265 return Ok(response);
1266 }
1267}
1268
1269pub struct ZetaEditPredictionProvider {
1270 zeta: Entity<Zeta>,
1271 singleton_buffer: Option<Entity<Buffer>>,
1272 pending_completions: ArrayVec<PendingCompletion, 2>,
1273 next_pending_completion_id: usize,
1274 current_completion: Option<CurrentEditPrediction>,
1275 last_request_timestamp: Instant,
1276 project: Entity<Project>,
1277}
1278
1279impl ZetaEditPredictionProvider {
1280 pub const THROTTLE_TIMEOUT: Duration = Duration::from_millis(300);
1281
1282 pub fn new(
1283 zeta: Entity<Zeta>,
1284 project: Entity<Project>,
1285 singleton_buffer: Option<Entity<Buffer>>,
1286 ) -> Self {
1287 Self {
1288 zeta,
1289 singleton_buffer,
1290 pending_completions: ArrayVec::new(),
1291 next_pending_completion_id: 0,
1292 current_completion: None,
1293 last_request_timestamp: Instant::now(),
1294 project,
1295 }
1296 }
1297}
1298
1299impl edit_prediction::EditPredictionProvider for ZetaEditPredictionProvider {
1300 fn name() -> &'static str {
1301 "zed-predict"
1302 }
1303
1304 fn display_name() -> &'static str {
1305 "Zed's Edit Predictions"
1306 }
1307
1308 fn show_completions_in_menu() -> bool {
1309 true
1310 }
1311
1312 fn show_tab_accept_marker() -> bool {
1313 true
1314 }
1315
1316 fn data_collection_state(&self, cx: &App) -> DataCollectionState {
1317 if let Some(buffer) = &self.singleton_buffer
1318 && let Some(file) = buffer.read(cx).file()
1319 {
1320 let is_project_open_source = self.zeta.read(cx).is_file_open_source(file, cx);
1321 if self.zeta.read(cx).data_collection_choice.is_enabled() {
1322 DataCollectionState::Enabled {
1323 is_project_open_source,
1324 }
1325 } else {
1326 DataCollectionState::Disabled {
1327 is_project_open_source,
1328 }
1329 }
1330 } else {
1331 return DataCollectionState::Disabled {
1332 is_project_open_source: false,
1333 };
1334 }
1335 }
1336
1337 fn toggle_data_collection(&mut self, cx: &mut App) {
1338 self.zeta
1339 .update(cx, |zeta, cx| zeta.toggle_data_collection_choice(cx));
1340 }
1341
1342 fn usage(&self, cx: &App) -> Option<EditPredictionUsage> {
1343 self.zeta.read(cx).usage(cx)
1344 }
1345
1346 fn is_enabled(
1347 &self,
1348 _buffer: &Entity<Buffer>,
1349 _cursor_position: language::Anchor,
1350 _cx: &App,
1351 ) -> bool {
1352 true
1353 }
1354 fn is_refreshing(&self) -> bool {
1355 !self.pending_completions.is_empty()
1356 }
1357
1358 fn refresh(
1359 &mut self,
1360 buffer: Entity<Buffer>,
1361 position: language::Anchor,
1362 _debounce: bool,
1363 cx: &mut Context<Self>,
1364 ) {
1365 if self.zeta.read(cx).update_required {
1366 return;
1367 }
1368
1369 if self
1370 .zeta
1371 .read(cx)
1372 .user_store
1373 .read_with(cx, |user_store, _cx| {
1374 user_store.account_too_young() || user_store.has_overdue_invoices()
1375 })
1376 {
1377 return;
1378 }
1379
1380 if let Some(current_completion) = self.current_completion.as_ref() {
1381 let snapshot = buffer.read(cx).snapshot();
1382 if current_completion
1383 .completion
1384 .interpolate(&snapshot)
1385 .is_some()
1386 {
1387 return;
1388 }
1389 }
1390
1391 let pending_completion_id = self.next_pending_completion_id;
1392 self.next_pending_completion_id += 1;
1393 let last_request_timestamp = self.last_request_timestamp;
1394
1395 let project = self.project.clone();
1396 let task = cx.spawn(async move |this, cx| {
1397 if let Some(timeout) = (last_request_timestamp + Self::THROTTLE_TIMEOUT)
1398 .checked_duration_since(Instant::now())
1399 {
1400 cx.background_executor().timer(timeout).await;
1401 }
1402
1403 let completion_request = this.update(cx, |this, cx| {
1404 this.last_request_timestamp = Instant::now();
1405 this.zeta.update(cx, |zeta, cx| {
1406 zeta.request_completion(&project, &buffer, position, cx)
1407 })
1408 });
1409
1410 let completion = match completion_request {
1411 Ok(completion_request) => {
1412 let completion_request = completion_request.await;
1413 completion_request.map(|c| {
1414 c.map(|completion| CurrentEditPrediction {
1415 buffer_id: buffer.entity_id(),
1416 completion,
1417 })
1418 })
1419 }
1420 Err(error) => Err(error),
1421 };
1422 let Some(new_completion) = completion
1423 .context("edit prediction failed")
1424 .log_err()
1425 .flatten()
1426 else {
1427 this.update(cx, |this, cx| {
1428 if this.pending_completions[0].id == pending_completion_id {
1429 this.pending_completions.remove(0);
1430 } else {
1431 this.pending_completions.clear();
1432 }
1433
1434 cx.notify();
1435 })
1436 .ok();
1437 return;
1438 };
1439
1440 this.update(cx, |this, cx| {
1441 if this.pending_completions[0].id == pending_completion_id {
1442 this.pending_completions.remove(0);
1443 } else {
1444 this.pending_completions.clear();
1445 }
1446
1447 if let Some(old_completion) = this.current_completion.as_ref() {
1448 let snapshot = buffer.read(cx).snapshot();
1449 if new_completion.should_replace_completion(old_completion, &snapshot) {
1450 this.zeta.update(cx, |zeta, cx| {
1451 zeta.completion_shown(&new_completion.completion, cx);
1452 });
1453 this.current_completion = Some(new_completion);
1454 }
1455 } else {
1456 this.zeta.update(cx, |zeta, cx| {
1457 zeta.completion_shown(&new_completion.completion, cx);
1458 });
1459 this.current_completion = Some(new_completion);
1460 }
1461
1462 cx.notify();
1463 })
1464 .ok();
1465 });
1466
1467 // We always maintain at most two pending completions. When we already
1468 // have two, we replace the newest one.
1469 if self.pending_completions.len() <= 1 {
1470 self.pending_completions.push(PendingCompletion {
1471 id: pending_completion_id,
1472 _task: task,
1473 });
1474 } else if self.pending_completions.len() == 2 {
1475 self.pending_completions.pop();
1476 self.pending_completions.push(PendingCompletion {
1477 id: pending_completion_id,
1478 _task: task,
1479 });
1480 }
1481 }
1482
1483 fn cycle(
1484 &mut self,
1485 _buffer: Entity<Buffer>,
1486 _cursor_position: language::Anchor,
1487 _direction: edit_prediction::Direction,
1488 _cx: &mut Context<Self>,
1489 ) {
1490 // Right now we don't support cycling.
1491 }
1492
1493 fn accept(&mut self, cx: &mut Context<Self>) {
1494 let completion_id = self
1495 .current_completion
1496 .as_ref()
1497 .map(|completion| completion.completion.id);
1498 if let Some(completion_id) = completion_id {
1499 self.zeta
1500 .update(cx, |zeta, cx| {
1501 zeta.accept_edit_prediction(completion_id, cx)
1502 })
1503 .detach();
1504 }
1505 self.pending_completions.clear();
1506 }
1507
1508 fn discard(&mut self, _cx: &mut Context<Self>) {
1509 self.pending_completions.clear();
1510 self.current_completion.take();
1511 }
1512
1513 fn suggest(
1514 &mut self,
1515 buffer: &Entity<Buffer>,
1516 cursor_position: language::Anchor,
1517 cx: &mut Context<Self>,
1518 ) -> Option<edit_prediction::EditPrediction> {
1519 let CurrentEditPrediction {
1520 buffer_id,
1521 completion,
1522 ..
1523 } = self.current_completion.as_mut()?;
1524
1525 // Invalidate previous completion if it was generated for a different buffer.
1526 if *buffer_id != buffer.entity_id() {
1527 self.current_completion.take();
1528 return None;
1529 }
1530
1531 let buffer = buffer.read(cx);
1532 let Some(edits) = completion.interpolate(&buffer.snapshot()) else {
1533 self.current_completion.take();
1534 return None;
1535 };
1536
1537 let cursor_row = cursor_position.to_point(buffer).row;
1538 let (closest_edit_ix, (closest_edit_range, _)) =
1539 edits.iter().enumerate().min_by_key(|(_, (range, _))| {
1540 let distance_from_start = cursor_row.abs_diff(range.start.to_point(buffer).row);
1541 let distance_from_end = cursor_row.abs_diff(range.end.to_point(buffer).row);
1542 cmp::min(distance_from_start, distance_from_end)
1543 })?;
1544
1545 let mut edit_start_ix = closest_edit_ix;
1546 for (range, _) in edits[..edit_start_ix].iter().rev() {
1547 let distance_from_closest_edit =
1548 closest_edit_range.start.to_point(buffer).row - range.end.to_point(buffer).row;
1549 if distance_from_closest_edit <= 1 {
1550 edit_start_ix -= 1;
1551 } else {
1552 break;
1553 }
1554 }
1555
1556 let mut edit_end_ix = closest_edit_ix + 1;
1557 for (range, _) in &edits[edit_end_ix..] {
1558 let distance_from_closest_edit =
1559 range.start.to_point(buffer).row - closest_edit_range.end.to_point(buffer).row;
1560 if distance_from_closest_edit <= 1 {
1561 edit_end_ix += 1;
1562 } else {
1563 break;
1564 }
1565 }
1566
1567 Some(edit_prediction::EditPrediction::Local {
1568 id: Some(completion.id.to_string().into()),
1569 edits: edits[edit_start_ix..edit_end_ix].to_vec(),
1570 edit_preview: Some(completion.edit_preview.clone()),
1571 })
1572 }
1573}
1574
1575/// Typical number of string bytes per token for the purposes of limiting model input. This is
1576/// intentionally low to err on the side of underestimating limits.
1577const BYTES_PER_TOKEN_GUESS: usize = 3;
1578
1579fn guess_token_count(bytes: usize) -> usize {
1580 bytes / BYTES_PER_TOKEN_GUESS
1581}
1582
1583#[cfg(test)]
1584mod tests {
1585 use client::test::FakeServer;
1586 use clock::{FakeSystemClock, ReplicaId};
1587 use cloud_api_types::{CreateLlmTokenResponse, LlmToken};
1588 use gpui::TestAppContext;
1589 use http_client::FakeHttpClient;
1590 use indoc::indoc;
1591 use language::Point;
1592 use parking_lot::Mutex;
1593 use serde_json::json;
1594 use settings::SettingsStore;
1595 use util::{path, rel_path::rel_path};
1596
1597 use super::*;
1598
1599 const BSD_0_TXT: &str = include_str!("../license_examples/0bsd.txt");
1600
1601 #[gpui::test]
1602 async fn test_edit_prediction_basic_interpolation(cx: &mut TestAppContext) {
1603 let buffer = cx.new(|cx| Buffer::local("Lorem ipsum dolor", cx));
1604 let edits: Arc<[(Range<Anchor>, Arc<str>)]> = cx.update(|cx| {
1605 to_completion_edits([(2..5, "REM".into()), (9..11, "".into())], &buffer, cx).into()
1606 });
1607
1608 let edit_preview = cx
1609 .read(|cx| buffer.read(cx).preview_edits(edits.clone(), cx))
1610 .await;
1611
1612 let completion = EditPrediction {
1613 edits,
1614 edit_preview,
1615 path: Path::new("").into(),
1616 snapshot: cx.read(|cx| buffer.read(cx).snapshot()),
1617 id: EditPredictionId(Uuid::new_v4()),
1618 excerpt_range: 0..0,
1619 cursor_offset: 0,
1620 input_outline: "".into(),
1621 input_events: "".into(),
1622 input_excerpt: "".into(),
1623 output_excerpt: "".into(),
1624 buffer_snapshotted_at: Instant::now(),
1625 response_received_at: Instant::now(),
1626 };
1627
1628 cx.update(|cx| {
1629 assert_eq!(
1630 from_completion_edits(
1631 &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
1632 &buffer,
1633 cx
1634 ),
1635 vec![(2..5, "REM".into()), (9..11, "".into())]
1636 );
1637
1638 buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "")], None, cx));
1639 assert_eq!(
1640 from_completion_edits(
1641 &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
1642 &buffer,
1643 cx
1644 ),
1645 vec![(2..2, "REM".into()), (6..8, "".into())]
1646 );
1647
1648 buffer.update(cx, |buffer, cx| buffer.undo(cx));
1649 assert_eq!(
1650 from_completion_edits(
1651 &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
1652 &buffer,
1653 cx
1654 ),
1655 vec![(2..5, "REM".into()), (9..11, "".into())]
1656 );
1657
1658 buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "R")], None, cx));
1659 assert_eq!(
1660 from_completion_edits(
1661 &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
1662 &buffer,
1663 cx
1664 ),
1665 vec![(3..3, "EM".into()), (7..9, "".into())]
1666 );
1667
1668 buffer.update(cx, |buffer, cx| buffer.edit([(3..3, "E")], None, cx));
1669 assert_eq!(
1670 from_completion_edits(
1671 &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
1672 &buffer,
1673 cx
1674 ),
1675 vec![(4..4, "M".into()), (8..10, "".into())]
1676 );
1677
1678 buffer.update(cx, |buffer, cx| buffer.edit([(4..4, "M")], None, cx));
1679 assert_eq!(
1680 from_completion_edits(
1681 &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
1682 &buffer,
1683 cx
1684 ),
1685 vec![(9..11, "".into())]
1686 );
1687
1688 buffer.update(cx, |buffer, cx| buffer.edit([(4..5, "")], None, cx));
1689 assert_eq!(
1690 from_completion_edits(
1691 &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
1692 &buffer,
1693 cx
1694 ),
1695 vec![(4..4, "M".into()), (8..10, "".into())]
1696 );
1697
1698 buffer.update(cx, |buffer, cx| buffer.edit([(8..10, "")], None, cx));
1699 assert_eq!(
1700 from_completion_edits(
1701 &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
1702 &buffer,
1703 cx
1704 ),
1705 vec![(4..4, "M".into())]
1706 );
1707
1708 buffer.update(cx, |buffer, cx| buffer.edit([(4..6, "")], None, cx));
1709 assert_eq!(completion.interpolate(&buffer.read(cx).snapshot()), None);
1710 })
1711 }
1712
1713 #[gpui::test]
1714 async fn test_clean_up_diff(cx: &mut TestAppContext) {
1715 init_test(cx);
1716
1717 assert_eq!(
1718 apply_edit_prediction(
1719 indoc! {"
1720 fn main() {
1721 let word_1 = \"lorem\";
1722 let range = word.len()..word.len();
1723 }
1724 "},
1725 indoc! {"
1726 <|editable_region_start|>
1727 fn main() {
1728 let word_1 = \"lorem\";
1729 let range = word_1.len()..word_1.len();
1730 }
1731
1732 <|editable_region_end|>
1733 "},
1734 cx,
1735 )
1736 .await,
1737 indoc! {"
1738 fn main() {
1739 let word_1 = \"lorem\";
1740 let range = word_1.len()..word_1.len();
1741 }
1742 "},
1743 );
1744
1745 assert_eq!(
1746 apply_edit_prediction(
1747 indoc! {"
1748 fn main() {
1749 let story = \"the quick\"
1750 }
1751 "},
1752 indoc! {"
1753 <|editable_region_start|>
1754 fn main() {
1755 let story = \"the quick brown fox jumps over the lazy dog\";
1756 }
1757
1758 <|editable_region_end|>
1759 "},
1760 cx,
1761 )
1762 .await,
1763 indoc! {"
1764 fn main() {
1765 let story = \"the quick brown fox jumps over the lazy dog\";
1766 }
1767 "},
1768 );
1769 }
1770
1771 #[gpui::test]
1772 async fn test_edit_prediction_end_of_buffer(cx: &mut TestAppContext) {
1773 init_test(cx);
1774
1775 let buffer_content = "lorem\n";
1776 let completion_response = indoc! {"
1777 ```animals.js
1778 <|start_of_file|>
1779 <|editable_region_start|>
1780 lorem
1781 ipsum
1782 <|editable_region_end|>
1783 ```"};
1784
1785 assert_eq!(
1786 apply_edit_prediction(buffer_content, completion_response, cx).await,
1787 "lorem\nipsum"
1788 );
1789 }
1790
1791 #[gpui::test]
1792 async fn test_can_collect_data(cx: &mut TestAppContext) {
1793 init_test(cx);
1794
1795 let fs = project::FakeFs::new(cx.executor());
1796 fs.insert_tree(path!("/project"), json!({ "LICENSE": BSD_0_TXT }))
1797 .await;
1798
1799 let project = Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;
1800 let buffer = project
1801 .update(cx, |project, cx| {
1802 project.open_local_buffer(path!("/project/src/main.rs"), cx)
1803 })
1804 .await
1805 .unwrap();
1806
1807 let (zeta, captured_request, _) = make_test_zeta(&project, cx).await;
1808 zeta.update(cx, |zeta, _cx| {
1809 zeta.data_collection_choice = DataCollectionChoice::Enabled
1810 });
1811
1812 run_edit_prediction(&buffer, &project, &zeta, cx).await;
1813 assert_eq!(
1814 captured_request.lock().clone().unwrap().can_collect_data,
1815 true
1816 );
1817
1818 zeta.update(cx, |zeta, _cx| {
1819 zeta.data_collection_choice = DataCollectionChoice::Disabled
1820 });
1821
1822 run_edit_prediction(&buffer, &project, &zeta, cx).await;
1823 assert_eq!(
1824 captured_request.lock().clone().unwrap().can_collect_data,
1825 false
1826 );
1827 }
1828
1829 #[gpui::test]
1830 async fn test_no_data_collection_for_remote_file(cx: &mut TestAppContext) {
1831 init_test(cx);
1832
1833 let fs = project::FakeFs::new(cx.executor());
1834 let project = Project::test(fs.clone(), [], cx).await;
1835
1836 let buffer = cx.new(|_cx| {
1837 Buffer::remote(
1838 language::BufferId::new(1).unwrap(),
1839 ReplicaId::new(1),
1840 language::Capability::ReadWrite,
1841 "fn main() {\n println!(\"Hello\");\n}",
1842 )
1843 });
1844
1845 let (zeta, captured_request, _) = make_test_zeta(&project, cx).await;
1846 zeta.update(cx, |zeta, _cx| {
1847 zeta.data_collection_choice = DataCollectionChoice::Enabled
1848 });
1849
1850 run_edit_prediction(&buffer, &project, &zeta, cx).await;
1851 assert_eq!(
1852 captured_request.lock().clone().unwrap().can_collect_data,
1853 false
1854 );
1855 }
1856
1857 #[gpui::test]
1858 async fn test_no_data_collection_for_private_file(cx: &mut TestAppContext) {
1859 init_test(cx);
1860
1861 let fs = project::FakeFs::new(cx.executor());
1862 fs.insert_tree(
1863 path!("/project"),
1864 json!({
1865 "LICENSE": BSD_0_TXT,
1866 ".env": "SECRET_KEY=secret"
1867 }),
1868 )
1869 .await;
1870
1871 let project = Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;
1872 let buffer = project
1873 .update(cx, |project, cx| {
1874 project.open_local_buffer("/project/.env", cx)
1875 })
1876 .await
1877 .unwrap();
1878
1879 let (zeta, captured_request, _) = make_test_zeta(&project, cx).await;
1880 zeta.update(cx, |zeta, _cx| {
1881 zeta.data_collection_choice = DataCollectionChoice::Enabled
1882 });
1883
1884 run_edit_prediction(&buffer, &project, &zeta, cx).await;
1885 assert_eq!(
1886 captured_request.lock().clone().unwrap().can_collect_data,
1887 false
1888 );
1889 }
1890
1891 #[gpui::test]
1892 async fn test_no_data_collection_for_untitled_buffer(cx: &mut TestAppContext) {
1893 init_test(cx);
1894
1895 let fs = project::FakeFs::new(cx.executor());
1896 let project = Project::test(fs.clone(), [], cx).await;
1897 let buffer = cx.new(|cx| Buffer::local("", cx));
1898
1899 let (zeta, captured_request, _) = make_test_zeta(&project, cx).await;
1900 zeta.update(cx, |zeta, _cx| {
1901 zeta.data_collection_choice = DataCollectionChoice::Enabled
1902 });
1903
1904 run_edit_prediction(&buffer, &project, &zeta, cx).await;
1905 assert_eq!(
1906 captured_request.lock().clone().unwrap().can_collect_data,
1907 false
1908 );
1909 }
1910
1911 #[gpui::test]
1912 async fn test_no_data_collection_when_closed_source(cx: &mut TestAppContext) {
1913 init_test(cx);
1914
1915 let fs = project::FakeFs::new(cx.executor());
1916 fs.insert_tree(path!("/project"), json!({ "main.rs": "fn main() {}" }))
1917 .await;
1918
1919 let project = Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;
1920 let buffer = project
1921 .update(cx, |project, cx| {
1922 project.open_local_buffer("/project/main.rs", cx)
1923 })
1924 .await
1925 .unwrap();
1926
1927 let (zeta, captured_request, _) = make_test_zeta(&project, cx).await;
1928 zeta.update(cx, |zeta, _cx| {
1929 zeta.data_collection_choice = DataCollectionChoice::Enabled
1930 });
1931
1932 run_edit_prediction(&buffer, &project, &zeta, cx).await;
1933 assert_eq!(
1934 captured_request.lock().clone().unwrap().can_collect_data,
1935 false
1936 );
1937 }
1938
1939 #[gpui::test]
1940 async fn test_data_collection_status_changes_on_move(cx: &mut TestAppContext) {
1941 init_test(cx);
1942
1943 let fs = project::FakeFs::new(cx.executor());
1944 fs.insert_tree(
1945 path!("/open_source_worktree"),
1946 json!({ "LICENSE": BSD_0_TXT, "main.rs": "" }),
1947 )
1948 .await;
1949 fs.insert_tree(path!("/closed_source_worktree"), json!({ "main.rs": "" }))
1950 .await;
1951
1952 let project = Project::test(
1953 fs.clone(),
1954 [
1955 path!("/open_source_worktree").as_ref(),
1956 path!("/closed_source_worktree").as_ref(),
1957 ],
1958 cx,
1959 )
1960 .await;
1961 let buffer = project
1962 .update(cx, |project, cx| {
1963 project.open_local_buffer(path!("/open_source_worktree/main.rs"), cx)
1964 })
1965 .await
1966 .unwrap();
1967
1968 let (zeta, captured_request, _) = make_test_zeta(&project, cx).await;
1969 zeta.update(cx, |zeta, _cx| {
1970 zeta.data_collection_choice = DataCollectionChoice::Enabled
1971 });
1972
1973 run_edit_prediction(&buffer, &project, &zeta, cx).await;
1974 assert_eq!(
1975 captured_request.lock().clone().unwrap().can_collect_data,
1976 true
1977 );
1978
1979 let closed_source_file = project
1980 .update(cx, |project, cx| {
1981 let worktree2 = project
1982 .worktree_for_root_name("closed_source_worktree", cx)
1983 .unwrap();
1984 worktree2.update(cx, |worktree2, cx| {
1985 worktree2.load_file(rel_path("main.rs"), cx)
1986 })
1987 })
1988 .await
1989 .unwrap()
1990 .file;
1991
1992 buffer.update(cx, |buffer, cx| {
1993 buffer.file_updated(closed_source_file, cx);
1994 });
1995
1996 run_edit_prediction(&buffer, &project, &zeta, cx).await;
1997 assert_eq!(
1998 captured_request.lock().clone().unwrap().can_collect_data,
1999 false
2000 );
2001 }
2002
2003 #[gpui::test]
2004 async fn test_no_data_collection_for_events_in_uncollectable_buffers(cx: &mut TestAppContext) {
2005 init_test(cx);
2006
2007 let fs = project::FakeFs::new(cx.executor());
2008 fs.insert_tree(
2009 path!("/worktree1"),
2010 json!({ "LICENSE": BSD_0_TXT, "main.rs": "", "other.rs": "" }),
2011 )
2012 .await;
2013 fs.insert_tree(path!("/worktree2"), json!({ "private.rs": "" }))
2014 .await;
2015
2016 let project = Project::test(
2017 fs.clone(),
2018 [path!("/worktree1").as_ref(), path!("/worktree2").as_ref()],
2019 cx,
2020 )
2021 .await;
2022 let buffer = project
2023 .update(cx, |project, cx| {
2024 project.open_local_buffer(path!("/worktree1/main.rs"), cx)
2025 })
2026 .await
2027 .unwrap();
2028 let private_buffer = project
2029 .update(cx, |project, cx| {
2030 project.open_local_buffer(path!("/worktree2/file.rs"), cx)
2031 })
2032 .await
2033 .unwrap();
2034
2035 let (zeta, captured_request, _) = make_test_zeta(&project, cx).await;
2036 zeta.update(cx, |zeta, _cx| {
2037 zeta.data_collection_choice = DataCollectionChoice::Enabled
2038 });
2039
2040 run_edit_prediction(&buffer, &project, &zeta, cx).await;
2041 assert_eq!(
2042 captured_request.lock().clone().unwrap().can_collect_data,
2043 true
2044 );
2045
2046 // this has a side effect of registering the buffer to watch for edits
2047 run_edit_prediction(&private_buffer, &project, &zeta, cx).await;
2048 assert_eq!(
2049 captured_request.lock().clone().unwrap().can_collect_data,
2050 false
2051 );
2052
2053 private_buffer.update(cx, |private_buffer, cx| {
2054 private_buffer.edit([(0..0, "An edit for the history!")], None, cx);
2055 });
2056
2057 run_edit_prediction(&buffer, &project, &zeta, cx).await;
2058 assert_eq!(
2059 captured_request.lock().clone().unwrap().can_collect_data,
2060 false
2061 );
2062
2063 // make an edit that uses too many bytes, causing private_buffer edit to not be able to be
2064 // included
2065 buffer.update(cx, |buffer, cx| {
2066 buffer.edit(
2067 [(0..0, " ".repeat(MAX_EVENT_TOKENS * BYTES_PER_TOKEN_GUESS))],
2068 None,
2069 cx,
2070 );
2071 });
2072
2073 run_edit_prediction(&buffer, &project, &zeta, cx).await;
2074 assert_eq!(
2075 captured_request.lock().clone().unwrap().can_collect_data,
2076 true
2077 );
2078 }
2079
2080 fn init_test(cx: &mut TestAppContext) {
2081 cx.update(|cx| {
2082 let settings_store = SettingsStore::test(cx);
2083 cx.set_global(settings_store);
2084 });
2085 }
2086
2087 async fn apply_edit_prediction(
2088 buffer_content: &str,
2089 completion_response: &str,
2090 cx: &mut TestAppContext,
2091 ) -> String {
2092 let fs = project::FakeFs::new(cx.executor());
2093 let project = Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;
2094 let buffer = cx.new(|cx| Buffer::local(buffer_content, cx));
2095 let (zeta, _, response) = make_test_zeta(&project, cx).await;
2096 *response.lock() = completion_response.to_string();
2097 let edit_prediction = run_edit_prediction(&buffer, &project, &zeta, cx).await;
2098 buffer.update(cx, |buffer, cx| {
2099 buffer.edit(edit_prediction.edits.iter().cloned(), None, cx)
2100 });
2101 buffer.read_with(cx, |buffer, _| buffer.text())
2102 }
2103
2104 async fn run_edit_prediction(
2105 buffer: &Entity<Buffer>,
2106 project: &Entity<Project>,
2107 zeta: &Entity<Zeta>,
2108 cx: &mut TestAppContext,
2109 ) -> EditPrediction {
2110 let cursor = buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(1, 0)));
2111 zeta.update(cx, |zeta, cx| zeta.register_buffer(buffer, &project, cx));
2112 cx.background_executor.run_until_parked();
2113 let completion_task = zeta.update(cx, |zeta, cx| {
2114 zeta.request_completion(&project, buffer, cursor, cx)
2115 });
2116 completion_task.await.unwrap().unwrap()
2117 }
2118
2119 async fn make_test_zeta(
2120 project: &Entity<Project>,
2121 cx: &mut TestAppContext,
2122 ) -> (
2123 Entity<Zeta>,
2124 Arc<Mutex<Option<PredictEditsBody>>>,
2125 Arc<Mutex<String>>,
2126 ) {
2127 let default_response = indoc! {"
2128 ```main.rs
2129 <|start_of_file|>
2130 <|editable_region_start|>
2131 hello world
2132 <|editable_region_end|>
2133 ```"
2134 };
2135 let captured_request: Arc<Mutex<Option<PredictEditsBody>>> = Arc::new(Mutex::new(None));
2136 let completion_response: Arc<Mutex<String>> =
2137 Arc::new(Mutex::new(default_response.to_string()));
2138 let http_client = FakeHttpClient::create({
2139 let captured_request = captured_request.clone();
2140 let completion_response = completion_response.clone();
2141 move |req| {
2142 let captured_request = captured_request.clone();
2143 let completion_response = completion_response.clone();
2144 async move {
2145 match (req.method(), req.uri().path()) {
2146 (&Method::POST, "/client/llm_tokens") => {
2147 Ok(http_client::Response::builder()
2148 .status(200)
2149 .body(
2150 serde_json::to_string(&CreateLlmTokenResponse {
2151 token: LlmToken("the-llm-token".to_string()),
2152 })
2153 .unwrap()
2154 .into(),
2155 )
2156 .unwrap())
2157 }
2158 (&Method::POST, "/predict_edits/v2") => {
2159 let mut request_body = String::new();
2160 req.into_body().read_to_string(&mut request_body).await?;
2161 *captured_request.lock() =
2162 Some(serde_json::from_str(&request_body).unwrap());
2163 Ok(http_client::Response::builder()
2164 .status(200)
2165 .body(
2166 serde_json::to_string(&PredictEditsResponse {
2167 request_id: Uuid::new_v4().to_string(),
2168 output_excerpt: completion_response.lock().clone(),
2169 })
2170 .unwrap()
2171 .into(),
2172 )
2173 .unwrap())
2174 }
2175 _ => Ok(http_client::Response::builder()
2176 .status(404)
2177 .body("Not Found".into())
2178 .unwrap()),
2179 }
2180 }
2181 }
2182 });
2183
2184 let client = cx.update(|cx| Client::new(Arc::new(FakeSystemClock::new()), http_client, cx));
2185 cx.update(|cx| {
2186 RefreshLlmTokenListener::register(client.clone(), cx);
2187 });
2188 let _server = FakeServer::for_client(42, &client, cx).await;
2189
2190 let zeta = cx.new(|cx| {
2191 let mut zeta = Zeta::new(client, project.read(cx).user_store(), cx);
2192
2193 let worktrees = project.read(cx).worktrees(cx).collect::<Vec<_>>();
2194 for worktree in worktrees {
2195 let worktree_id = worktree.read(cx).id();
2196 zeta.license_detection_watchers
2197 .entry(worktree_id)
2198 .or_insert_with(|| Rc::new(LicenseDetectionWatcher::new(&worktree, cx)));
2199 }
2200
2201 zeta
2202 });
2203
2204 (zeta, captured_request, completion_response)
2205 }
2206
2207 fn to_completion_edits(
2208 iterator: impl IntoIterator<Item = (Range<usize>, Arc<str>)>,
2209 buffer: &Entity<Buffer>,
2210 cx: &App,
2211 ) -> Vec<(Range<Anchor>, Arc<str>)> {
2212 let buffer = buffer.read(cx);
2213 iterator
2214 .into_iter()
2215 .map(|(range, text)| {
2216 (
2217 buffer.anchor_after(range.start)..buffer.anchor_before(range.end),
2218 text,
2219 )
2220 })
2221 .collect()
2222 }
2223
2224 fn from_completion_edits(
2225 editor_edits: &[(Range<Anchor>, Arc<str>)],
2226 buffer: &Entity<Buffer>,
2227 cx: &App,
2228 ) -> Vec<(Range<usize>, Arc<str>)> {
2229 let buffer = buffer.read(cx);
2230 editor_edits
2231 .iter()
2232 .map(|(range, text)| {
2233 (
2234 range.start.to_offset(buffer)..range.end.to_offset(buffer),
2235 text.clone(),
2236 )
2237 })
2238 .collect()
2239 }
2240
2241 #[ctor::ctor]
2242 fn init_logger() {
2243 zlog::init_test();
2244 }
2245}