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