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,
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 edit_preview = edit_preview.await;
739
740 Ok(Some(EditPrediction {
741 id: EditPredictionId(request_id),
742 path,
743 excerpt_range: editable_range,
744 cursor_offset,
745 edits,
746 edit_preview,
747 snapshot,
748 input_outline: input_outline.into(),
749 input_events: input_events.into(),
750 input_excerpt: input_excerpt.into(),
751 output_excerpt,
752 buffer_snapshotted_at,
753 response_received_at: Instant::now(),
754 }))
755 })
756 }
757
758 fn parse_edits(
759 output_excerpt: Arc<str>,
760 editable_range: Range<usize>,
761 snapshot: &BufferSnapshot,
762 ) -> Result<Vec<(Range<Anchor>, Arc<str>)>> {
763 let content = output_excerpt.replace(CURSOR_MARKER, "");
764
765 let start_markers = content
766 .match_indices(EDITABLE_REGION_START_MARKER)
767 .collect::<Vec<_>>();
768 anyhow::ensure!(
769 start_markers.len() == 1,
770 "expected exactly one start marker, found {}",
771 start_markers.len()
772 );
773
774 let end_markers = content
775 .match_indices(EDITABLE_REGION_END_MARKER)
776 .collect::<Vec<_>>();
777 anyhow::ensure!(
778 end_markers.len() == 1,
779 "expected exactly one end marker, found {}",
780 end_markers.len()
781 );
782
783 let sof_markers = content
784 .match_indices(START_OF_FILE_MARKER)
785 .collect::<Vec<_>>();
786 anyhow::ensure!(
787 sof_markers.len() <= 1,
788 "expected at most one start-of-file marker, found {}",
789 sof_markers.len()
790 );
791
792 let codefence_start = start_markers[0].0;
793 let content = &content[codefence_start..];
794
795 let newline_ix = content.find('\n').context("could not find newline")?;
796 let content = &content[newline_ix + 1..];
797
798 let codefence_end = content
799 .rfind(&format!("\n{EDITABLE_REGION_END_MARKER}"))
800 .context("could not find end marker")?;
801 let new_text = &content[..codefence_end];
802
803 let old_text = snapshot
804 .text_for_range(editable_range.clone())
805 .collect::<String>();
806
807 Ok(Self::compute_edits(
808 old_text,
809 new_text,
810 editable_range.start,
811 snapshot,
812 ))
813 }
814
815 pub fn compute_edits(
816 old_text: String,
817 new_text: &str,
818 offset: usize,
819 snapshot: &BufferSnapshot,
820 ) -> Vec<(Range<Anchor>, Arc<str>)> {
821 text_diff(&old_text, new_text)
822 .into_iter()
823 .map(|(mut old_range, new_text)| {
824 old_range.start += offset;
825 old_range.end += offset;
826
827 let prefix_len = common_prefix(
828 snapshot.chars_for_range(old_range.clone()),
829 new_text.chars(),
830 );
831 old_range.start += prefix_len;
832
833 let suffix_len = common_prefix(
834 snapshot.reversed_chars_for_range(old_range.clone()),
835 new_text[prefix_len..].chars().rev(),
836 );
837 old_range.end = old_range.end.saturating_sub(suffix_len);
838
839 let new_text = new_text[prefix_len..new_text.len() - suffix_len].into();
840 let range = if old_range.is_empty() {
841 let anchor = snapshot.anchor_after(old_range.start);
842 anchor..anchor
843 } else {
844 snapshot.anchor_after(old_range.start)..snapshot.anchor_before(old_range.end)
845 };
846 (range, new_text)
847 })
848 .collect()
849 }
850
851 pub fn is_completion_rated(&self, completion_id: EditPredictionId) -> bool {
852 self.rated_completions.contains(&completion_id)
853 }
854
855 pub fn completion_shown(&mut self, completion: &EditPrediction, cx: &mut Context<Self>) {
856 self.shown_completions.push_front(completion.clone());
857 if self.shown_completions.len() > 50 {
858 let completion = self.shown_completions.pop_back().unwrap();
859 self.rated_completions.remove(&completion.id);
860 }
861 cx.notify();
862 }
863
864 pub fn rate_completion(
865 &mut self,
866 completion: &EditPrediction,
867 rating: EditPredictionRating,
868 feedback: String,
869 cx: &mut Context<Self>,
870 ) {
871 self.rated_completions.insert(completion.id);
872 telemetry::event!(
873 "Edit Prediction Rated",
874 rating,
875 input_events = completion.input_events,
876 input_excerpt = completion.input_excerpt,
877 input_outline = completion.input_outline,
878 output_excerpt = completion.output_excerpt,
879 feedback
880 );
881 self.client.telemetry().flush_events().detach();
882 cx.notify();
883 }
884
885 pub fn shown_completions(&self) -> impl DoubleEndedIterator<Item = &EditPrediction> {
886 self.shown_completions.iter()
887 }
888
889 pub fn shown_completions_len(&self) -> usize {
890 self.shown_completions.len()
891 }
892
893 fn report_changes_for_buffer(
894 &mut self,
895 buffer: &Entity<Buffer>,
896 project: &Entity<Project>,
897 cx: &mut Context<Self>,
898 ) -> BufferSnapshot {
899 let zeta_project = self.get_or_init_zeta_project(project, cx);
900 let registered_buffer = Self::register_buffer_impl(zeta_project, buffer, project, cx);
901
902 let new_snapshot = buffer.read(cx).snapshot();
903 if new_snapshot.version != registered_buffer.snapshot.version {
904 let old_snapshot = mem::replace(&mut registered_buffer.snapshot, new_snapshot.clone());
905 Self::push_event(
906 zeta_project,
907 Event::BufferChange {
908 old_snapshot,
909 new_snapshot: new_snapshot.clone(),
910 timestamp: Instant::now(),
911 },
912 );
913 }
914
915 new_snapshot
916 }
917
918 fn can_collect_file(&self, file: &Arc<dyn File>, cx: &App) -> bool {
919 self.data_collection_choice.is_enabled() && self.is_file_open_source(file, cx)
920 }
921
922 fn can_collect_events(&self, events: &[Event], cx: &App) -> bool {
923 if !self.data_collection_choice.is_enabled() {
924 return false;
925 }
926 let mut last_checked_file = None;
927 for event in events {
928 match event {
929 Event::BufferChange {
930 old_snapshot,
931 new_snapshot,
932 ..
933 } => {
934 if let Some(old_file) = old_snapshot.file()
935 && let Some(new_file) = new_snapshot.file()
936 {
937 if let Some(last_checked_file) = last_checked_file
938 && Arc::ptr_eq(last_checked_file, old_file)
939 && Arc::ptr_eq(last_checked_file, new_file)
940 {
941 continue;
942 }
943 if !self.can_collect_file(old_file, cx) {
944 return false;
945 }
946 if !Arc::ptr_eq(old_file, new_file) && !self.can_collect_file(new_file, cx)
947 {
948 return false;
949 }
950 last_checked_file = Some(new_file);
951 } else {
952 return false;
953 }
954 }
955 }
956 }
957 true
958 }
959
960 fn is_file_open_source(&self, file: &Arc<dyn File>, cx: &App) -> bool {
961 if !file.is_local() || file.is_private() {
962 return false;
963 }
964 self.license_detection_watchers
965 .get(&file.worktree_id(cx))
966 .is_some_and(|watcher| watcher.is_project_open_source())
967 }
968
969 fn load_data_collection_choice() -> DataCollectionChoice {
970 let choice = KEY_VALUE_STORE
971 .read_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE)
972 .log_err()
973 .flatten();
974
975 match choice.as_deref() {
976 Some("true") => DataCollectionChoice::Enabled,
977 Some("false") => DataCollectionChoice::Disabled,
978 Some(_) => {
979 log::error!("unknown value in '{ZED_PREDICT_DATA_COLLECTION_CHOICE}'");
980 DataCollectionChoice::NotAnswered
981 }
982 None => DataCollectionChoice::NotAnswered,
983 }
984 }
985
986 fn toggle_data_collection_choice(&mut self, cx: &mut Context<Self>) {
987 self.data_collection_choice = self.data_collection_choice.toggle();
988 let new_choice = self.data_collection_choice;
989 db::write_and_log(cx, move || {
990 KEY_VALUE_STORE.write_kvp(
991 ZED_PREDICT_DATA_COLLECTION_CHOICE.into(),
992 new_choice.is_enabled().to_string(),
993 )
994 });
995 }
996}
997
998pub struct PerformPredictEditsParams {
999 pub client: Arc<Client>,
1000 pub llm_token: LlmApiToken,
1001 pub app_version: SemanticVersion,
1002 pub body: PredictEditsBody,
1003}
1004
1005#[derive(Error, Debug)]
1006#[error(
1007 "You must update to Zed version {minimum_version} or higher to continue using edit predictions."
1008)]
1009pub struct ZedUpdateRequiredError {
1010 minimum_version: SemanticVersion,
1011}
1012
1013fn common_prefix<T1: Iterator<Item = char>, T2: Iterator<Item = char>>(a: T1, b: T2) -> usize {
1014 a.zip(b)
1015 .take_while(|(a, b)| a == b)
1016 .map(|(a, _)| a.len_utf8())
1017 .sum()
1018}
1019
1020fn git_info_for_file(
1021 project: &Entity<Project>,
1022 project_path: &ProjectPath,
1023 cx: &App,
1024) -> Option<PredictEditsGitInfo> {
1025 let git_store = project.read(cx).git_store().read(cx);
1026 if let Some((repository, _repo_path)) =
1027 git_store.repository_and_path_for_project_path(project_path, cx)
1028 {
1029 let repository = repository.read(cx);
1030 let head_sha = repository
1031 .head_commit
1032 .as_ref()
1033 .map(|head_commit| head_commit.sha.to_string());
1034 let remote_origin_url = repository.remote_origin_url.clone();
1035 let remote_upstream_url = repository.remote_upstream_url.clone();
1036 if head_sha.is_none() && remote_origin_url.is_none() && remote_upstream_url.is_none() {
1037 return None;
1038 }
1039 Some(PredictEditsGitInfo {
1040 head_sha,
1041 remote_origin_url,
1042 remote_upstream_url,
1043 })
1044 } else {
1045 None
1046 }
1047}
1048
1049pub struct GatherContextOutput {
1050 pub body: PredictEditsBody,
1051 pub editable_range: Range<usize>,
1052 pub included_events_count: usize,
1053}
1054
1055pub fn gather_context(
1056 full_path_str: String,
1057 snapshot: &BufferSnapshot,
1058 cursor_point: language::Point,
1059 prompt_for_events: impl FnOnce() -> (String, usize) + Send + 'static,
1060 cx: &App,
1061) -> Task<Result<GatherContextOutput>> {
1062 cx.background_spawn({
1063 let snapshot = snapshot.clone();
1064 async move {
1065 let input_excerpt = excerpt_for_cursor_position(
1066 cursor_point,
1067 &full_path_str,
1068 &snapshot,
1069 MAX_REWRITE_TOKENS,
1070 MAX_CONTEXT_TOKENS,
1071 );
1072 let (input_events, included_events_count) = prompt_for_events();
1073 let editable_range = input_excerpt.editable_range.to_offset(&snapshot);
1074
1075 let body = PredictEditsBody {
1076 input_events,
1077 input_excerpt: input_excerpt.prompt,
1078 can_collect_data: false,
1079 diagnostic_groups: None,
1080 git_info: None,
1081 outline: None,
1082 speculated_output: None,
1083 };
1084
1085 Ok(GatherContextOutput {
1086 body,
1087 editable_range,
1088 included_events_count,
1089 })
1090 }
1091 })
1092}
1093
1094fn prompt_for_events_impl(events: &[Event], mut remaining_tokens: usize) -> (String, usize) {
1095 let mut result = String::new();
1096 for (ix, event) in events.iter().rev().enumerate() {
1097 let event_string = event.to_prompt();
1098 let event_tokens = guess_token_count(event_string.len());
1099 if event_tokens > remaining_tokens {
1100 return (result, ix);
1101 }
1102
1103 if !result.is_empty() {
1104 result.insert_str(0, "\n\n");
1105 }
1106 result.insert_str(0, &event_string);
1107 remaining_tokens -= event_tokens;
1108 }
1109 return (result, events.len());
1110}
1111
1112struct RegisteredBuffer {
1113 snapshot: BufferSnapshot,
1114 _subscriptions: [gpui::Subscription; 2],
1115}
1116
1117#[derive(Clone)]
1118pub enum Event {
1119 BufferChange {
1120 old_snapshot: BufferSnapshot,
1121 new_snapshot: BufferSnapshot,
1122 timestamp: Instant,
1123 },
1124}
1125
1126impl Event {
1127 fn to_prompt(&self) -> String {
1128 match self {
1129 Event::BufferChange {
1130 old_snapshot,
1131 new_snapshot,
1132 ..
1133 } => {
1134 let mut prompt = String::new();
1135
1136 let old_path = old_snapshot
1137 .file()
1138 .map(|f| f.path().as_ref())
1139 .unwrap_or(RelPath::unix("untitled").unwrap());
1140 let new_path = new_snapshot
1141 .file()
1142 .map(|f| f.path().as_ref())
1143 .unwrap_or(RelPath::unix("untitled").unwrap());
1144 if old_path != new_path {
1145 writeln!(prompt, "User renamed {:?} to {:?}\n", old_path, new_path).unwrap();
1146 }
1147
1148 let diff = language::unified_diff(&old_snapshot.text(), &new_snapshot.text());
1149 if !diff.is_empty() {
1150 write!(
1151 prompt,
1152 "User edited {:?}:\n```diff\n{}\n```",
1153 new_path, diff
1154 )
1155 .unwrap();
1156 }
1157
1158 prompt
1159 }
1160 }
1161 }
1162}
1163
1164#[derive(Debug, Clone)]
1165struct CurrentEditPrediction {
1166 buffer_id: EntityId,
1167 completion: EditPrediction,
1168}
1169
1170impl CurrentEditPrediction {
1171 fn should_replace_completion(&self, old_completion: &Self, snapshot: &BufferSnapshot) -> bool {
1172 if self.buffer_id != old_completion.buffer_id {
1173 return true;
1174 }
1175
1176 let Some(old_edits) = old_completion.completion.interpolate(snapshot) else {
1177 return true;
1178 };
1179 let Some(new_edits) = self.completion.interpolate(snapshot) else {
1180 return false;
1181 };
1182
1183 if old_edits.len() == 1 && new_edits.len() == 1 {
1184 let (old_range, old_text) = &old_edits[0];
1185 let (new_range, new_text) = &new_edits[0];
1186 new_range == old_range && new_text.starts_with(old_text.as_ref())
1187 } else {
1188 true
1189 }
1190 }
1191}
1192
1193struct PendingCompletion {
1194 id: usize,
1195 _task: Task<()>,
1196}
1197
1198#[derive(Debug, Clone, Copy)]
1199pub enum DataCollectionChoice {
1200 NotAnswered,
1201 Enabled,
1202 Disabled,
1203}
1204
1205impl DataCollectionChoice {
1206 pub fn is_enabled(self) -> bool {
1207 match self {
1208 Self::Enabled => true,
1209 Self::NotAnswered | Self::Disabled => false,
1210 }
1211 }
1212
1213 pub fn is_answered(self) -> bool {
1214 match self {
1215 Self::Enabled | Self::Disabled => true,
1216 Self::NotAnswered => false,
1217 }
1218 }
1219
1220 #[must_use]
1221 pub fn toggle(&self) -> DataCollectionChoice {
1222 match self {
1223 Self::Enabled => Self::Disabled,
1224 Self::Disabled => Self::Enabled,
1225 Self::NotAnswered => Self::Enabled,
1226 }
1227 }
1228}
1229
1230impl From<bool> for DataCollectionChoice {
1231 fn from(value: bool) -> Self {
1232 match value {
1233 true => DataCollectionChoice::Enabled,
1234 false => DataCollectionChoice::Disabled,
1235 }
1236 }
1237}
1238
1239async fn llm_token_retry(
1240 llm_token: &LlmApiToken,
1241 client: &Arc<Client>,
1242 build_request: impl Fn(String) -> Result<Request<AsyncBody>>,
1243) -> Result<Response<AsyncBody>> {
1244 let mut did_retry = false;
1245 let http_client = client.http_client();
1246 let mut token = llm_token.acquire(client).await?;
1247 loop {
1248 let request = build_request(token.clone())?;
1249 let response = http_client.send(request).await?;
1250
1251 if !did_retry
1252 && !response.status().is_success()
1253 && response
1254 .headers()
1255 .get(EXPIRED_LLM_TOKEN_HEADER_NAME)
1256 .is_some()
1257 {
1258 did_retry = true;
1259 token = llm_token.refresh(client).await?;
1260 continue;
1261 }
1262
1263 return Ok(response);
1264 }
1265}
1266
1267pub struct ZetaEditPredictionProvider {
1268 zeta: Entity<Zeta>,
1269 singleton_buffer: Option<Entity<Buffer>>,
1270 pending_completions: ArrayVec<PendingCompletion, 2>,
1271 next_pending_completion_id: usize,
1272 current_completion: Option<CurrentEditPrediction>,
1273 last_request_timestamp: Instant,
1274 project: Entity<Project>,
1275}
1276
1277impl ZetaEditPredictionProvider {
1278 pub const THROTTLE_TIMEOUT: Duration = Duration::from_millis(300);
1279
1280 pub fn new(
1281 zeta: Entity<Zeta>,
1282 project: Entity<Project>,
1283 singleton_buffer: Option<Entity<Buffer>>,
1284 ) -> Self {
1285 Self {
1286 zeta,
1287 singleton_buffer,
1288 pending_completions: ArrayVec::new(),
1289 next_pending_completion_id: 0,
1290 current_completion: None,
1291 last_request_timestamp: Instant::now(),
1292 project,
1293 }
1294 }
1295}
1296
1297impl edit_prediction::EditPredictionProvider for ZetaEditPredictionProvider {
1298 fn name() -> &'static str {
1299 "zed-predict"
1300 }
1301
1302 fn display_name() -> &'static str {
1303 "Zed's Edit Predictions"
1304 }
1305
1306 fn show_completions_in_menu() -> bool {
1307 true
1308 }
1309
1310 fn show_tab_accept_marker() -> bool {
1311 true
1312 }
1313
1314 fn data_collection_state(&self, cx: &App) -> DataCollectionState {
1315 if let Some(buffer) = &self.singleton_buffer
1316 && let Some(file) = buffer.read(cx).file()
1317 {
1318 let is_project_open_source = self.zeta.read(cx).is_file_open_source(file, cx);
1319 if self.zeta.read(cx).data_collection_choice.is_enabled() {
1320 DataCollectionState::Enabled {
1321 is_project_open_source,
1322 }
1323 } else {
1324 DataCollectionState::Disabled {
1325 is_project_open_source,
1326 }
1327 }
1328 } else {
1329 return DataCollectionState::Disabled {
1330 is_project_open_source: false,
1331 };
1332 }
1333 }
1334
1335 fn toggle_data_collection(&mut self, cx: &mut App) {
1336 self.zeta
1337 .update(cx, |zeta, cx| zeta.toggle_data_collection_choice(cx));
1338 }
1339
1340 fn usage(&self, cx: &App) -> Option<EditPredictionUsage> {
1341 self.zeta.read(cx).usage(cx)
1342 }
1343
1344 fn is_enabled(
1345 &self,
1346 _buffer: &Entity<Buffer>,
1347 _cursor_position: language::Anchor,
1348 _cx: &App,
1349 ) -> bool {
1350 true
1351 }
1352 fn is_refreshing(&self) -> bool {
1353 !self.pending_completions.is_empty()
1354 }
1355
1356 fn refresh(
1357 &mut self,
1358 buffer: Entity<Buffer>,
1359 position: language::Anchor,
1360 _debounce: bool,
1361 cx: &mut Context<Self>,
1362 ) {
1363 if self.zeta.read(cx).update_required {
1364 return;
1365 }
1366
1367 if self
1368 .zeta
1369 .read(cx)
1370 .user_store
1371 .read_with(cx, |user_store, _cx| {
1372 user_store.account_too_young() || user_store.has_overdue_invoices()
1373 })
1374 {
1375 return;
1376 }
1377
1378 if let Some(current_completion) = self.current_completion.as_ref() {
1379 let snapshot = buffer.read(cx).snapshot();
1380 if current_completion
1381 .completion
1382 .interpolate(&snapshot)
1383 .is_some()
1384 {
1385 return;
1386 }
1387 }
1388
1389 let pending_completion_id = self.next_pending_completion_id;
1390 self.next_pending_completion_id += 1;
1391 let last_request_timestamp = self.last_request_timestamp;
1392
1393 let project = self.project.clone();
1394 let task = cx.spawn(async move |this, cx| {
1395 if let Some(timeout) = (last_request_timestamp + Self::THROTTLE_TIMEOUT)
1396 .checked_duration_since(Instant::now())
1397 {
1398 cx.background_executor().timer(timeout).await;
1399 }
1400
1401 let completion_request = this.update(cx, |this, cx| {
1402 this.last_request_timestamp = Instant::now();
1403 this.zeta.update(cx, |zeta, cx| {
1404 zeta.request_completion(&project, &buffer, position, cx)
1405 })
1406 });
1407
1408 let completion = match completion_request {
1409 Ok(completion_request) => {
1410 let completion_request = completion_request.await;
1411 completion_request.map(|c| {
1412 c.map(|completion| CurrentEditPrediction {
1413 buffer_id: buffer.entity_id(),
1414 completion,
1415 })
1416 })
1417 }
1418 Err(error) => Err(error),
1419 };
1420 let Some(new_completion) = completion
1421 .context("edit prediction failed")
1422 .log_err()
1423 .flatten()
1424 else {
1425 this.update(cx, |this, cx| {
1426 if this.pending_completions[0].id == pending_completion_id {
1427 this.pending_completions.remove(0);
1428 } else {
1429 this.pending_completions.clear();
1430 }
1431
1432 cx.notify();
1433 })
1434 .ok();
1435 return;
1436 };
1437
1438 this.update(cx, |this, cx| {
1439 if this.pending_completions[0].id == pending_completion_id {
1440 this.pending_completions.remove(0);
1441 } else {
1442 this.pending_completions.clear();
1443 }
1444
1445 if let Some(old_completion) = this.current_completion.as_ref() {
1446 let snapshot = buffer.read(cx).snapshot();
1447 if new_completion.should_replace_completion(old_completion, &snapshot) {
1448 this.zeta.update(cx, |zeta, cx| {
1449 zeta.completion_shown(&new_completion.completion, cx);
1450 });
1451 this.current_completion = Some(new_completion);
1452 }
1453 } else {
1454 this.zeta.update(cx, |zeta, cx| {
1455 zeta.completion_shown(&new_completion.completion, cx);
1456 });
1457 this.current_completion = Some(new_completion);
1458 }
1459
1460 cx.notify();
1461 })
1462 .ok();
1463 });
1464
1465 // We always maintain at most two pending completions. When we already
1466 // have two, we replace the newest one.
1467 if self.pending_completions.len() <= 1 {
1468 self.pending_completions.push(PendingCompletion {
1469 id: pending_completion_id,
1470 _task: task,
1471 });
1472 } else if self.pending_completions.len() == 2 {
1473 self.pending_completions.pop();
1474 self.pending_completions.push(PendingCompletion {
1475 id: pending_completion_id,
1476 _task: task,
1477 });
1478 }
1479 }
1480
1481 fn cycle(
1482 &mut self,
1483 _buffer: Entity<Buffer>,
1484 _cursor_position: language::Anchor,
1485 _direction: edit_prediction::Direction,
1486 _cx: &mut Context<Self>,
1487 ) {
1488 // Right now we don't support cycling.
1489 }
1490
1491 fn accept(&mut self, cx: &mut Context<Self>) {
1492 let completion_id = self
1493 .current_completion
1494 .as_ref()
1495 .map(|completion| completion.completion.id);
1496 if let Some(completion_id) = completion_id {
1497 self.zeta
1498 .update(cx, |zeta, cx| {
1499 zeta.accept_edit_prediction(completion_id, cx)
1500 })
1501 .detach();
1502 }
1503 self.pending_completions.clear();
1504 }
1505
1506 fn discard(&mut self, _cx: &mut Context<Self>) {
1507 self.pending_completions.clear();
1508 self.current_completion.take();
1509 }
1510
1511 fn suggest(
1512 &mut self,
1513 buffer: &Entity<Buffer>,
1514 cursor_position: language::Anchor,
1515 cx: &mut Context<Self>,
1516 ) -> Option<edit_prediction::EditPrediction> {
1517 let CurrentEditPrediction {
1518 buffer_id,
1519 completion,
1520 ..
1521 } = self.current_completion.as_mut()?;
1522
1523 // Invalidate previous completion if it was generated for a different buffer.
1524 if *buffer_id != buffer.entity_id() {
1525 self.current_completion.take();
1526 return None;
1527 }
1528
1529 let buffer = buffer.read(cx);
1530 let Some(edits) = completion.interpolate(&buffer.snapshot()) else {
1531 self.current_completion.take();
1532 return None;
1533 };
1534
1535 let cursor_row = cursor_position.to_point(buffer).row;
1536 let (closest_edit_ix, (closest_edit_range, _)) =
1537 edits.iter().enumerate().min_by_key(|(_, (range, _))| {
1538 let distance_from_start = cursor_row.abs_diff(range.start.to_point(buffer).row);
1539 let distance_from_end = cursor_row.abs_diff(range.end.to_point(buffer).row);
1540 cmp::min(distance_from_start, distance_from_end)
1541 })?;
1542
1543 let mut edit_start_ix = closest_edit_ix;
1544 for (range, _) in edits[..edit_start_ix].iter().rev() {
1545 let distance_from_closest_edit =
1546 closest_edit_range.start.to_point(buffer).row - range.end.to_point(buffer).row;
1547 if distance_from_closest_edit <= 1 {
1548 edit_start_ix -= 1;
1549 } else {
1550 break;
1551 }
1552 }
1553
1554 let mut edit_end_ix = closest_edit_ix + 1;
1555 for (range, _) in &edits[edit_end_ix..] {
1556 let distance_from_closest_edit =
1557 range.start.to_point(buffer).row - closest_edit_range.end.to_point(buffer).row;
1558 if distance_from_closest_edit <= 1 {
1559 edit_end_ix += 1;
1560 } else {
1561 break;
1562 }
1563 }
1564
1565 Some(edit_prediction::EditPrediction::Local {
1566 id: Some(completion.id.to_string().into()),
1567 edits: edits[edit_start_ix..edit_end_ix].to_vec(),
1568 edit_preview: Some(completion.edit_preview.clone()),
1569 })
1570 }
1571}
1572
1573/// Typical number of string bytes per token for the purposes of limiting model input. This is
1574/// intentionally low to err on the side of underestimating limits.
1575const BYTES_PER_TOKEN_GUESS: usize = 3;
1576
1577fn guess_token_count(bytes: usize) -> usize {
1578 bytes / BYTES_PER_TOKEN_GUESS
1579}
1580
1581#[cfg(test)]
1582mod tests {
1583 use client::test::FakeServer;
1584 use clock::{FakeSystemClock, ReplicaId};
1585 use cloud_api_types::{CreateLlmTokenResponse, LlmToken};
1586 use gpui::TestAppContext;
1587 use http_client::FakeHttpClient;
1588 use indoc::indoc;
1589 use language::Point;
1590 use parking_lot::Mutex;
1591 use serde_json::json;
1592 use settings::SettingsStore;
1593 use util::{path, rel_path::rel_path};
1594
1595 use super::*;
1596
1597 const BSD_0_TXT: &str = include_str!("../license_examples/0bsd.txt");
1598
1599 #[gpui::test]
1600 async fn test_edit_prediction_basic_interpolation(cx: &mut TestAppContext) {
1601 let buffer = cx.new(|cx| Buffer::local("Lorem ipsum dolor", cx));
1602 let edits: Arc<[(Range<Anchor>, Arc<str>)]> = cx.update(|cx| {
1603 to_completion_edits([(2..5, "REM".into()), (9..11, "".into())], &buffer, cx).into()
1604 });
1605
1606 let edit_preview = cx
1607 .read(|cx| buffer.read(cx).preview_edits(edits.clone(), cx))
1608 .await;
1609
1610 let completion = EditPrediction {
1611 edits,
1612 edit_preview,
1613 path: Path::new("").into(),
1614 snapshot: cx.read(|cx| buffer.read(cx).snapshot()),
1615 id: EditPredictionId(Uuid::new_v4()),
1616 excerpt_range: 0..0,
1617 cursor_offset: 0,
1618 input_outline: "".into(),
1619 input_events: "".into(),
1620 input_excerpt: "".into(),
1621 output_excerpt: "".into(),
1622 buffer_snapshotted_at: Instant::now(),
1623 response_received_at: Instant::now(),
1624 };
1625
1626 cx.update(|cx| {
1627 assert_eq!(
1628 from_completion_edits(
1629 &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
1630 &buffer,
1631 cx
1632 ),
1633 vec![(2..5, "REM".into()), (9..11, "".into())]
1634 );
1635
1636 buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "")], None, cx));
1637 assert_eq!(
1638 from_completion_edits(
1639 &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
1640 &buffer,
1641 cx
1642 ),
1643 vec![(2..2, "REM".into()), (6..8, "".into())]
1644 );
1645
1646 buffer.update(cx, |buffer, cx| buffer.undo(cx));
1647 assert_eq!(
1648 from_completion_edits(
1649 &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
1650 &buffer,
1651 cx
1652 ),
1653 vec![(2..5, "REM".into()), (9..11, "".into())]
1654 );
1655
1656 buffer.update(cx, |buffer, cx| buffer.edit([(2..5, "R")], None, cx));
1657 assert_eq!(
1658 from_completion_edits(
1659 &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
1660 &buffer,
1661 cx
1662 ),
1663 vec![(3..3, "EM".into()), (7..9, "".into())]
1664 );
1665
1666 buffer.update(cx, |buffer, cx| buffer.edit([(3..3, "E")], None, cx));
1667 assert_eq!(
1668 from_completion_edits(
1669 &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
1670 &buffer,
1671 cx
1672 ),
1673 vec![(4..4, "M".into()), (8..10, "".into())]
1674 );
1675
1676 buffer.update(cx, |buffer, cx| buffer.edit([(4..4, "M")], None, cx));
1677 assert_eq!(
1678 from_completion_edits(
1679 &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
1680 &buffer,
1681 cx
1682 ),
1683 vec![(9..11, "".into())]
1684 );
1685
1686 buffer.update(cx, |buffer, cx| buffer.edit([(4..5, "")], None, cx));
1687 assert_eq!(
1688 from_completion_edits(
1689 &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
1690 &buffer,
1691 cx
1692 ),
1693 vec![(4..4, "M".into()), (8..10, "".into())]
1694 );
1695
1696 buffer.update(cx, |buffer, cx| buffer.edit([(8..10, "")], None, cx));
1697 assert_eq!(
1698 from_completion_edits(
1699 &completion.interpolate(&buffer.read(cx).snapshot()).unwrap(),
1700 &buffer,
1701 cx
1702 ),
1703 vec![(4..4, "M".into())]
1704 );
1705
1706 buffer.update(cx, |buffer, cx| buffer.edit([(4..6, "")], None, cx));
1707 assert_eq!(completion.interpolate(&buffer.read(cx).snapshot()), None);
1708 })
1709 }
1710
1711 #[gpui::test]
1712 async fn test_clean_up_diff(cx: &mut TestAppContext) {
1713 init_test(cx);
1714
1715 assert_eq!(
1716 apply_edit_prediction(
1717 indoc! {"
1718 fn main() {
1719 let word_1 = \"lorem\";
1720 let range = word.len()..word.len();
1721 }
1722 "},
1723 indoc! {"
1724 <|editable_region_start|>
1725 fn main() {
1726 let word_1 = \"lorem\";
1727 let range = word_1.len()..word_1.len();
1728 }
1729
1730 <|editable_region_end|>
1731 "},
1732 cx,
1733 )
1734 .await,
1735 indoc! {"
1736 fn main() {
1737 let word_1 = \"lorem\";
1738 let range = word_1.len()..word_1.len();
1739 }
1740 "},
1741 );
1742
1743 assert_eq!(
1744 apply_edit_prediction(
1745 indoc! {"
1746 fn main() {
1747 let story = \"the quick\"
1748 }
1749 "},
1750 indoc! {"
1751 <|editable_region_start|>
1752 fn main() {
1753 let story = \"the quick brown fox jumps over the lazy dog\";
1754 }
1755
1756 <|editable_region_end|>
1757 "},
1758 cx,
1759 )
1760 .await,
1761 indoc! {"
1762 fn main() {
1763 let story = \"the quick brown fox jumps over the lazy dog\";
1764 }
1765 "},
1766 );
1767 }
1768
1769 #[gpui::test]
1770 async fn test_edit_prediction_end_of_buffer(cx: &mut TestAppContext) {
1771 init_test(cx);
1772
1773 let buffer_content = "lorem\n";
1774 let completion_response = indoc! {"
1775 ```animals.js
1776 <|start_of_file|>
1777 <|editable_region_start|>
1778 lorem
1779 ipsum
1780 <|editable_region_end|>
1781 ```"};
1782
1783 assert_eq!(
1784 apply_edit_prediction(buffer_content, completion_response, cx).await,
1785 "lorem\nipsum"
1786 );
1787 }
1788
1789 #[gpui::test]
1790 async fn test_can_collect_data(cx: &mut TestAppContext) {
1791 init_test(cx);
1792
1793 let fs = project::FakeFs::new(cx.executor());
1794 fs.insert_tree(path!("/project"), json!({ "LICENSE": BSD_0_TXT }))
1795 .await;
1796
1797 let project = Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;
1798 let buffer = project
1799 .update(cx, |project, cx| {
1800 project.open_local_buffer(path!("/project/src/main.rs"), cx)
1801 })
1802 .await
1803 .unwrap();
1804
1805 let (zeta, captured_request, _) = make_test_zeta(&project, cx).await;
1806 zeta.update(cx, |zeta, _cx| {
1807 zeta.data_collection_choice = DataCollectionChoice::Enabled
1808 });
1809
1810 run_edit_prediction(&buffer, &project, &zeta, cx).await;
1811 assert_eq!(
1812 captured_request.lock().clone().unwrap().can_collect_data,
1813 true
1814 );
1815
1816 zeta.update(cx, |zeta, _cx| {
1817 zeta.data_collection_choice = DataCollectionChoice::Disabled
1818 });
1819
1820 run_edit_prediction(&buffer, &project, &zeta, cx).await;
1821 assert_eq!(
1822 captured_request.lock().clone().unwrap().can_collect_data,
1823 false
1824 );
1825 }
1826
1827 #[gpui::test]
1828 async fn test_no_data_collection_for_remote_file(cx: &mut TestAppContext) {
1829 init_test(cx);
1830
1831 let fs = project::FakeFs::new(cx.executor());
1832 let project = Project::test(fs.clone(), [], cx).await;
1833
1834 let buffer = cx.new(|_cx| {
1835 Buffer::remote(
1836 language::BufferId::new(1).unwrap(),
1837 ReplicaId::new(1),
1838 language::Capability::ReadWrite,
1839 "fn main() {\n println!(\"Hello\");\n}",
1840 )
1841 });
1842
1843 let (zeta, captured_request, _) = make_test_zeta(&project, cx).await;
1844 zeta.update(cx, |zeta, _cx| {
1845 zeta.data_collection_choice = DataCollectionChoice::Enabled
1846 });
1847
1848 run_edit_prediction(&buffer, &project, &zeta, cx).await;
1849 assert_eq!(
1850 captured_request.lock().clone().unwrap().can_collect_data,
1851 false
1852 );
1853 }
1854
1855 #[gpui::test]
1856 async fn test_no_data_collection_for_private_file(cx: &mut TestAppContext) {
1857 init_test(cx);
1858
1859 let fs = project::FakeFs::new(cx.executor());
1860 fs.insert_tree(
1861 path!("/project"),
1862 json!({
1863 "LICENSE": BSD_0_TXT,
1864 ".env": "SECRET_KEY=secret"
1865 }),
1866 )
1867 .await;
1868
1869 let project = Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;
1870 let buffer = project
1871 .update(cx, |project, cx| {
1872 project.open_local_buffer("/project/.env", cx)
1873 })
1874 .await
1875 .unwrap();
1876
1877 let (zeta, captured_request, _) = make_test_zeta(&project, cx).await;
1878 zeta.update(cx, |zeta, _cx| {
1879 zeta.data_collection_choice = DataCollectionChoice::Enabled
1880 });
1881
1882 run_edit_prediction(&buffer, &project, &zeta, cx).await;
1883 assert_eq!(
1884 captured_request.lock().clone().unwrap().can_collect_data,
1885 false
1886 );
1887 }
1888
1889 #[gpui::test]
1890 async fn test_no_data_collection_for_untitled_buffer(cx: &mut TestAppContext) {
1891 init_test(cx);
1892
1893 let fs = project::FakeFs::new(cx.executor());
1894 let project = Project::test(fs.clone(), [], cx).await;
1895 let buffer = cx.new(|cx| Buffer::local("", cx));
1896
1897 let (zeta, captured_request, _) = make_test_zeta(&project, cx).await;
1898 zeta.update(cx, |zeta, _cx| {
1899 zeta.data_collection_choice = DataCollectionChoice::Enabled
1900 });
1901
1902 run_edit_prediction(&buffer, &project, &zeta, cx).await;
1903 assert_eq!(
1904 captured_request.lock().clone().unwrap().can_collect_data,
1905 false
1906 );
1907 }
1908
1909 #[gpui::test]
1910 async fn test_no_data_collection_when_closed_source(cx: &mut TestAppContext) {
1911 init_test(cx);
1912
1913 let fs = project::FakeFs::new(cx.executor());
1914 fs.insert_tree(path!("/project"), json!({ "main.rs": "fn main() {}" }))
1915 .await;
1916
1917 let project = Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;
1918 let buffer = project
1919 .update(cx, |project, cx| {
1920 project.open_local_buffer("/project/main.rs", cx)
1921 })
1922 .await
1923 .unwrap();
1924
1925 let (zeta, captured_request, _) = make_test_zeta(&project, cx).await;
1926 zeta.update(cx, |zeta, _cx| {
1927 zeta.data_collection_choice = DataCollectionChoice::Enabled
1928 });
1929
1930 run_edit_prediction(&buffer, &project, &zeta, cx).await;
1931 assert_eq!(
1932 captured_request.lock().clone().unwrap().can_collect_data,
1933 false
1934 );
1935 }
1936
1937 #[gpui::test]
1938 async fn test_data_collection_status_changes_on_move(cx: &mut TestAppContext) {
1939 init_test(cx);
1940
1941 let fs = project::FakeFs::new(cx.executor());
1942 fs.insert_tree(
1943 path!("/open_source_worktree"),
1944 json!({ "LICENSE": BSD_0_TXT, "main.rs": "" }),
1945 )
1946 .await;
1947 fs.insert_tree(path!("/closed_source_worktree"), json!({ "main.rs": "" }))
1948 .await;
1949
1950 let project = Project::test(
1951 fs.clone(),
1952 [
1953 path!("/open_source_worktree").as_ref(),
1954 path!("/closed_source_worktree").as_ref(),
1955 ],
1956 cx,
1957 )
1958 .await;
1959 let buffer = project
1960 .update(cx, |project, cx| {
1961 project.open_local_buffer(path!("/open_source_worktree/main.rs"), cx)
1962 })
1963 .await
1964 .unwrap();
1965
1966 let (zeta, captured_request, _) = make_test_zeta(&project, cx).await;
1967 zeta.update(cx, |zeta, _cx| {
1968 zeta.data_collection_choice = DataCollectionChoice::Enabled
1969 });
1970
1971 run_edit_prediction(&buffer, &project, &zeta, cx).await;
1972 assert_eq!(
1973 captured_request.lock().clone().unwrap().can_collect_data,
1974 true
1975 );
1976
1977 let closed_source_file = project
1978 .update(cx, |project, cx| {
1979 let worktree2 = project
1980 .worktree_for_root_name("closed_source_worktree", cx)
1981 .unwrap();
1982 worktree2.update(cx, |worktree2, cx| {
1983 worktree2.load_file(rel_path("main.rs"), cx)
1984 })
1985 })
1986 .await
1987 .unwrap()
1988 .file;
1989
1990 buffer.update(cx, |buffer, cx| {
1991 buffer.file_updated(closed_source_file, cx);
1992 });
1993
1994 run_edit_prediction(&buffer, &project, &zeta, cx).await;
1995 assert_eq!(
1996 captured_request.lock().clone().unwrap().can_collect_data,
1997 false
1998 );
1999 }
2000
2001 #[gpui::test]
2002 async fn test_no_data_collection_for_events_in_uncollectable_buffers(cx: &mut TestAppContext) {
2003 init_test(cx);
2004
2005 let fs = project::FakeFs::new(cx.executor());
2006 fs.insert_tree(
2007 path!("/worktree1"),
2008 json!({ "LICENSE": BSD_0_TXT, "main.rs": "", "other.rs": "" }),
2009 )
2010 .await;
2011 fs.insert_tree(path!("/worktree2"), json!({ "private.rs": "" }))
2012 .await;
2013
2014 let project = Project::test(
2015 fs.clone(),
2016 [path!("/worktree1").as_ref(), path!("/worktree2").as_ref()],
2017 cx,
2018 )
2019 .await;
2020 let buffer = project
2021 .update(cx, |project, cx| {
2022 project.open_local_buffer(path!("/worktree1/main.rs"), cx)
2023 })
2024 .await
2025 .unwrap();
2026 let private_buffer = project
2027 .update(cx, |project, cx| {
2028 project.open_local_buffer(path!("/worktree2/file.rs"), cx)
2029 })
2030 .await
2031 .unwrap();
2032
2033 let (zeta, captured_request, _) = make_test_zeta(&project, cx).await;
2034 zeta.update(cx, |zeta, _cx| {
2035 zeta.data_collection_choice = DataCollectionChoice::Enabled
2036 });
2037
2038 run_edit_prediction(&buffer, &project, &zeta, cx).await;
2039 assert_eq!(
2040 captured_request.lock().clone().unwrap().can_collect_data,
2041 true
2042 );
2043
2044 // this has a side effect of registering the buffer to watch for edits
2045 run_edit_prediction(&private_buffer, &project, &zeta, cx).await;
2046 assert_eq!(
2047 captured_request.lock().clone().unwrap().can_collect_data,
2048 false
2049 );
2050
2051 private_buffer.update(cx, |private_buffer, cx| {
2052 private_buffer.edit([(0..0, "An edit for the history!")], None, cx);
2053 });
2054
2055 run_edit_prediction(&buffer, &project, &zeta, cx).await;
2056 assert_eq!(
2057 captured_request.lock().clone().unwrap().can_collect_data,
2058 false
2059 );
2060
2061 // make an edit that uses too many bytes, causing private_buffer edit to not be able to be
2062 // included
2063 buffer.update(cx, |buffer, cx| {
2064 buffer.edit(
2065 [(0..0, " ".repeat(MAX_EVENT_TOKENS * BYTES_PER_TOKEN_GUESS))],
2066 None,
2067 cx,
2068 );
2069 });
2070
2071 run_edit_prediction(&buffer, &project, &zeta, cx).await;
2072 assert_eq!(
2073 captured_request.lock().clone().unwrap().can_collect_data,
2074 true
2075 );
2076 }
2077
2078 fn init_test(cx: &mut TestAppContext) {
2079 cx.update(|cx| {
2080 let settings_store = SettingsStore::test(cx);
2081 cx.set_global(settings_store);
2082 });
2083 }
2084
2085 async fn apply_edit_prediction(
2086 buffer_content: &str,
2087 completion_response: &str,
2088 cx: &mut TestAppContext,
2089 ) -> String {
2090 let fs = project::FakeFs::new(cx.executor());
2091 let project = Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;
2092 let buffer = cx.new(|cx| Buffer::local(buffer_content, cx));
2093 let (zeta, _, response) = make_test_zeta(&project, cx).await;
2094 *response.lock() = completion_response.to_string();
2095 let edit_prediction = run_edit_prediction(&buffer, &project, &zeta, cx).await;
2096 buffer.update(cx, |buffer, cx| {
2097 buffer.edit(edit_prediction.edits.iter().cloned(), None, cx)
2098 });
2099 buffer.read_with(cx, |buffer, _| buffer.text())
2100 }
2101
2102 async fn run_edit_prediction(
2103 buffer: &Entity<Buffer>,
2104 project: &Entity<Project>,
2105 zeta: &Entity<Zeta>,
2106 cx: &mut TestAppContext,
2107 ) -> EditPrediction {
2108 let cursor = buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(1, 0)));
2109 zeta.update(cx, |zeta, cx| zeta.register_buffer(buffer, &project, cx));
2110 cx.background_executor.run_until_parked();
2111 let completion_task = zeta.update(cx, |zeta, cx| {
2112 zeta.request_completion(&project, buffer, cursor, cx)
2113 });
2114 completion_task.await.unwrap().unwrap()
2115 }
2116
2117 async fn make_test_zeta(
2118 project: &Entity<Project>,
2119 cx: &mut TestAppContext,
2120 ) -> (
2121 Entity<Zeta>,
2122 Arc<Mutex<Option<PredictEditsBody>>>,
2123 Arc<Mutex<String>>,
2124 ) {
2125 let default_response = indoc! {"
2126 ```main.rs
2127 <|start_of_file|>
2128 <|editable_region_start|>
2129 hello world
2130 <|editable_region_end|>
2131 ```"
2132 };
2133 let captured_request: Arc<Mutex<Option<PredictEditsBody>>> = Arc::new(Mutex::new(None));
2134 let completion_response: Arc<Mutex<String>> =
2135 Arc::new(Mutex::new(default_response.to_string()));
2136 let http_client = FakeHttpClient::create({
2137 let captured_request = captured_request.clone();
2138 let completion_response = completion_response.clone();
2139 move |req| {
2140 let captured_request = captured_request.clone();
2141 let completion_response = completion_response.clone();
2142 async move {
2143 match (req.method(), req.uri().path()) {
2144 (&Method::POST, "/client/llm_tokens") => {
2145 Ok(http_client::Response::builder()
2146 .status(200)
2147 .body(
2148 serde_json::to_string(&CreateLlmTokenResponse {
2149 token: LlmToken("the-llm-token".to_string()),
2150 })
2151 .unwrap()
2152 .into(),
2153 )
2154 .unwrap())
2155 }
2156 (&Method::POST, "/predict_edits/v2") => {
2157 let mut request_body = String::new();
2158 req.into_body().read_to_string(&mut request_body).await?;
2159 *captured_request.lock() =
2160 Some(serde_json::from_str(&request_body).unwrap());
2161 Ok(http_client::Response::builder()
2162 .status(200)
2163 .body(
2164 serde_json::to_string(&PredictEditsResponse {
2165 request_id: Uuid::new_v4(),
2166 output_excerpt: completion_response.lock().clone(),
2167 })
2168 .unwrap()
2169 .into(),
2170 )
2171 .unwrap())
2172 }
2173 _ => Ok(http_client::Response::builder()
2174 .status(404)
2175 .body("Not Found".into())
2176 .unwrap()),
2177 }
2178 }
2179 }
2180 });
2181
2182 let client = cx.update(|cx| Client::new(Arc::new(FakeSystemClock::new()), http_client, cx));
2183 cx.update(|cx| {
2184 RefreshLlmTokenListener::register(client.clone(), cx);
2185 });
2186 let _server = FakeServer::for_client(42, &client, cx).await;
2187
2188 let zeta = cx.new(|cx| {
2189 let mut zeta = Zeta::new(client, project.read(cx).user_store(), cx);
2190
2191 let worktrees = project.read(cx).worktrees(cx).collect::<Vec<_>>();
2192 for worktree in worktrees {
2193 let worktree_id = worktree.read(cx).id();
2194 zeta.license_detection_watchers
2195 .entry(worktree_id)
2196 .or_insert_with(|| Rc::new(LicenseDetectionWatcher::new(&worktree, cx)));
2197 }
2198
2199 zeta
2200 });
2201
2202 (zeta, captured_request, completion_response)
2203 }
2204
2205 fn to_completion_edits(
2206 iterator: impl IntoIterator<Item = (Range<usize>, Arc<str>)>,
2207 buffer: &Entity<Buffer>,
2208 cx: &App,
2209 ) -> Vec<(Range<Anchor>, Arc<str>)> {
2210 let buffer = buffer.read(cx);
2211 iterator
2212 .into_iter()
2213 .map(|(range, text)| {
2214 (
2215 buffer.anchor_after(range.start)..buffer.anchor_before(range.end),
2216 text,
2217 )
2218 })
2219 .collect()
2220 }
2221
2222 fn from_completion_edits(
2223 editor_edits: &[(Range<Anchor>, Arc<str>)],
2224 buffer: &Entity<Buffer>,
2225 cx: &App,
2226 ) -> Vec<(Range<usize>, Arc<str>)> {
2227 let buffer = buffer.read(cx);
2228 editor_edits
2229 .iter()
2230 .map(|(range, text)| {
2231 (
2232 range.start.to_offset(buffer)..range.end.to_offset(buffer),
2233 text.clone(),
2234 )
2235 })
2236 .collect()
2237 }
2238
2239 #[ctor::ctor]
2240 fn init_logger() {
2241 zlog::init_test();
2242 }
2243}