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