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