Skip to repository content3552 lines · 132.9 KB · rust
tenant.openagents/omega
No repository description is available.
OpenAgents Git authority 2026-07-28T01:50:57.769Z Public web read
NIP-34 coordinate
30617:7649603503856e5148d571eac2766b288a8ff1e9e35d380337a1d2b0015b4f92:omegaMaintainersHidden in public view
References2 branches · 1 tag
Read-only clone
git clone https://openagents.com/git/tenant.openagents/omega.gitBrowse files
edit_prediction.rs
1use anyhow::{Context as _, Result, anyhow};
2use buffer_diff::BufferDiff;
3use client::{Client, EditPredictionUsage, UserStore, global_llm_token};
4use cloud_api_client::LlmApiToken;
5use cloud_api_types::{
6 EditPredictionRecentFile, EditPredictionSettledKeptChars,
7 MAX_EDIT_PREDICTION_SETTLED_PER_REQUEST, OrganizationId, SettledEditPrediction,
8 SettledEditPredictionSampleData, SubmitEditPredictionFeedbackBody,
9 SubmitEditPredictionSettledBatchBody, SubmitEditPredictionSettledResponse,
10};
11use cloud_llm_client::predict_edits_v3::{
12 PREDICT_EDITS_MODE_HEADER_NAME, PREDICT_EDITS_REQUEST_ID_HEADER_NAME,
13 PREDICT_EDITS_TRIGGER_HEADER_NAME, PredictEditsMode, PredictEditsV3Request,
14 PredictEditsV3Response, RawCompletionRequest, RawCompletionResponse,
15};
16use cloud_llm_client::predict_edits_v4::{PredictEditsV4Request, PredictEditsV4Response};
17use cloud_llm_client::{
18 EditPredictionRejectReason, EditPredictionRejection,
19 MAX_EDIT_PREDICTION_REJECTIONS_PER_REQUEST, MINIMUM_REQUIRED_VERSION_HEADER_NAME,
20 PREFERRED_EXPERIMENT_HEADER_NAME, PredictEditsRequestTrigger, RejectEditPredictionsBodyRef,
21 ZED_VERSION_HEADER_NAME,
22};
23use collections::{HashMap, HashSet};
24use copilot::{Copilot, Reinstall, SignIn, SignOut};
25use credentials_provider::CredentialsProvider;
26use db::kvp::{Dismissable, KeyValueStore};
27use edit_prediction_context::{RelatedExcerptStore, RelatedExcerptStoreEvent, RelatedFile};
28use edit_prediction_types::EditPredictionRequestTrigger;
29use feature_flags::{FeatureFlag, FeatureFlagAppExt as _, PresenceFlag, register_feature_flag};
30use futures::{
31 AsyncReadExt as _, FutureExt as _, StreamExt as _,
32 channel::mpsc::{self, UnboundedReceiver},
33 select_biased,
34};
35use git::repository::FileHistoryChangedFileSets;
36use gpui::BackgroundExecutor;
37use gpui::TaskExt;
38use gpui::http_client::Url;
39use gpui::{
40 App, AsyncApp, Context, Entity, EntityId, Global, SharedString, Task, WeakEntity, actions,
41 http_client::{self, AsyncBody, Method},
42 prelude::*,
43};
44use heapless::Vec as ArrayVec;
45use language::{
46 Anchor, Buffer, BufferEditSource, BufferSnapshot, EditPredictionPromptFormat,
47 EditPredictionsMode, EditPreview, File, OffsetRangeExt, Point, TextBufferSnapshot, ToOffset,
48 ToPoint, language_settings::all_language_settings,
49};
50use project::{DisableAiSettings, Project, ProjectPath, WorktreeId};
51use release_channel::AppVersion;
52use semver::Version;
53use serde::de::DeserializeOwned;
54use settings::{
55 EditPredictionDataCollectionChoice, EditPredictionProvider, Settings as _, update_settings_file,
56};
57use std::collections::{VecDeque, hash_map};
58use std::env;
59use std::rc::Rc;
60use text::{AnchorRangeExt, Edit};
61use workspace::{AppState, Workspace};
62use zeta_prompt::ContextSource;
63use zeta_prompt::{Zeta2PromptInput, Zeta3PromptInput, ZetaFormat};
64
65use std::mem;
66use std::ops::Range;
67use std::path::Path;
68use std::str::FromStr as _;
69use std::sync::Arc;
70use std::time::{Duration, Instant};
71
72use thiserror::Error;
73use util::ResultExt as _;
74
75pub mod cursor_excerpt;
76pub mod data_collection;
77pub mod example_spec;
78pub mod fim;
79mod license_detection;
80pub mod mercury;
81pub mod metrics;
82pub mod ollama;
83mod onboarding_modal;
84pub mod open_ai_response;
85mod prediction;
86
87pub mod udiff;
88
89pub mod open_ai_compatible;
90mod zed_edit_prediction_delegate;
91pub mod zeta;
92
93#[cfg(test)]
94mod edit_prediction_tests;
95
96use crate::cursor_excerpt::expand_context_syntactically_then_linewise;
97use crate::data_collection::{CapturedPredictionContext, capture_prediction_context};
98use crate::example_spec::RecentFile;
99use crate::license_detection::LicenseDetectionWatcher;
100use crate::mercury::Mercury;
101pub use crate::metrics::{KeptRateResult, compute_kept_rate};
102use crate::onboarding_modal::ZedPredictModal;
103use crate::prediction::EditPredictionResult;
104pub use crate::prediction::{EditPrediction, EditPredictionId, EditPredictionInputs};
105pub use language_model::ApiKeyState;
106pub use telemetry_events::EditPredictionRating;
107pub use zed_edit_prediction_delegate::ZedEditPredictionDelegate;
108
109actions!(
110 edit_prediction,
111 [
112 /// Resets the edit prediction onboarding state.
113 ResetOnboarding,
114 /// Clears the edit prediction history.
115 ClearHistory,
116 ]
117);
118
119/// Maximum number of events to track.
120const EVENT_COUNT_MAX: usize = 10;
121const RECENT_PATH_COUNT_MAX: usize = 20;
122const CHANGE_GROUPING_LINE_SPAN: u32 = 8;
123const EDIT_HISTORY_DIFF_SIZE_LIMIT: usize = 2048 * 3; // ~2048 tokens or ~50% of typical prompt budget
124const COLLABORATOR_EDIT_LOCALITY_CONTEXT_TOKENS: usize = 512;
125const GIT_CHANGED_FILE_SETS_COMMIT_LIMIT: usize = 100;
126const LAST_CHANGE_GROUPING_TIME: Duration = Duration::from_secs(1);
127const ZED_PREDICT_DATA_COLLECTION_CHOICE: &str = "zed_predict_data_collection_choice";
128const REJECT_REQUEST_DEBOUNCE: Duration = Duration::from_secs(15);
129const REQUEST_TIMEOUT_BACKOFF: Duration = Duration::from_secs(10);
130
131const EDIT_PREDICTION_SETTLED_TTL: Duration = Duration::from_secs(60 * 5);
132const EDIT_PREDICTION_SETTLED_QUIESCENCE: Duration = Duration::from_secs(10);
133const EDIT_PREDICTION_CAPTURE_MAX_FUTURE_EVENTS: usize = 4;
134const EDIT_PREDICTION_SETTLED_MAX_EDITABLE_REGION_BYTES: usize = 4 * 1024;
135
136pub struct EditPredictionJumpsFeatureFlag;
137
138impl FeatureFlag for EditPredictionJumpsFeatureFlag {
139 const NAME: &'static str = "edit_prediction_jumps";
140 type Value = PresenceFlag;
141}
142register_feature_flag!(EditPredictionJumpsFeatureFlag);
143
144#[derive(Clone)]
145struct EditPredictionStoreGlobal(Entity<EditPredictionStore>);
146
147impl Global for EditPredictionStoreGlobal {}
148
149/// Configuration for using the raw Zeta2 endpoint.
150/// When set, the client uses the raw endpoint and constructs the prompt itself.
151/// The version is also used as the Baseten environment name (lowercased).
152#[derive(Clone)]
153pub struct Zeta2RawConfig {
154 pub model_id: Option<String>,
155 pub environment: Option<String>,
156 pub format: ZetaFormat,
157}
158
159pub struct EditPredictionStore {
160 client: Arc<Client>,
161 user_store: Entity<UserStore>,
162 llm_token: LlmApiToken,
163 _fetch_experiments_task: Task<()>,
164 projects: HashMap<EntityId, ProjectState>,
165 update_required: bool,
166 edit_prediction_model: EditPredictionModel,
167 zeta2_raw_config: Option<Zeta2RawConfig>,
168 request_backoff_until: Option<Instant>,
169 preferred_experiment: Option<String>,
170 available_experiments: Vec<String>,
171 pub mercury: Mercury,
172 legacy_data_collection_enabled: bool,
173 reject_predictions_tx: mpsc::UnboundedSender<EditPredictionRejectionPayload>,
174 settled_predictions_tx: mpsc::UnboundedSender<Instant>,
175 rateable_predictions: VecDeque<EditPrediction>,
176 rated_predictions: HashSet<EditPredictionId>,
177 #[cfg(test)]
178 settled_event_callback: Option<Box<dyn Fn(EditPredictionId, String)>>,
179 credentials_provider: Arc<dyn CredentialsProvider>,
180}
181
182pub(crate) struct EditPredictionRejectionPayload {
183 rejection: EditPredictionRejection,
184 organization_id: Option<OrganizationId>,
185}
186
187#[derive(Copy, Clone, PartialEq, Eq)]
188pub enum EditPredictionModel {
189 Zeta,
190 Fim { format: EditPredictionPromptFormat },
191 Mercury,
192}
193
194pub struct EditPredictionModelInput {
195 project: Entity<Project>,
196 buffer: Entity<Buffer>,
197 snapshot: BufferSnapshot,
198 position: Anchor,
199 events: Vec<Arc<zeta_prompt::Event>>,
200 related_files: Vec<RelatedFile>,
201 editable_context: Option<Task<anyhow::Result<Vec<RelatedFile>>>>,
202 mode: PredictEditsMode,
203 trigger: PredictEditsRequestTrigger,
204 diagnostic_search_range: Range<Point>,
205 debug_tx: Option<mpsc::UnboundedSender<DebugEvent>>,
206 can_collect_data: bool,
207 is_open_source: bool,
208 allow_jump: bool,
209}
210
211#[derive(Debug)]
212pub enum DebugEvent {
213 ContextRetrievalStarted(ContextRetrievalStartedDebugEvent),
214 ContextRetrievalFinished(ContextRetrievalFinishedDebugEvent),
215 EditPredictionStarted(EditPredictionStartedDebugEvent),
216 EditPredictionFinished(EditPredictionFinishedDebugEvent),
217}
218
219#[derive(Debug)]
220pub struct ContextRetrievalStartedDebugEvent {
221 pub project_entity_id: EntityId,
222 pub timestamp: Instant,
223 pub search_prompt: String,
224}
225
226#[derive(Debug)]
227pub struct ContextRetrievalFinishedDebugEvent {
228 pub project_entity_id: EntityId,
229 pub timestamp: Instant,
230 pub metadata: Vec<(&'static str, SharedString)>,
231}
232
233#[derive(Debug)]
234pub struct EditPredictionStartedDebugEvent {
235 pub buffer: WeakEntity<Buffer>,
236 pub position: Anchor,
237 pub prompt: Option<String>,
238}
239
240#[derive(Debug)]
241pub struct EditPredictionFinishedDebugEvent {
242 pub buffer: WeakEntity<Buffer>,
243 pub position: Anchor,
244 pub model_output: Option<String>,
245}
246
247/// An event with associated metadata for reconstructing buffer state.
248#[derive(Clone)]
249pub struct StoredEvent {
250 pub event: Arc<zeta_prompt::Event>,
251 pub old_snapshot: TextBufferSnapshot,
252 pub new_snapshot_version: clock::Global,
253 pub total_edit_range: Range<Anchor>,
254 pub(crate) file_context: Option<Entity<StoredFileContext>>,
255}
256
257pub(crate) struct StoredFileContext {
258 pub(crate) uncommitted_diff: Option<Entity<BufferDiff>>,
259 pub(crate) git_changed_file_sets: Option<Arc<FileHistoryChangedFileSets>>,
260 pub(crate) git_changed_file_sets_task: Option<Task<()>>,
261}
262
263impl StoredEvent {
264 fn can_merge(
265 &self,
266 next_old_event: &StoredEvent,
267 latest_snapshot: &TextBufferSnapshot,
268 latest_edit_range: &Range<Anchor>,
269 ) -> bool {
270 // Events must be for the same buffer and be contiguous across included snapshots to be mergeable.
271 if self.old_snapshot.remote_id() != next_old_event.old_snapshot.remote_id() {
272 return false;
273 }
274 if self.old_snapshot.remote_id() != latest_snapshot.remote_id() {
275 return false;
276 }
277 if self.new_snapshot_version != next_old_event.old_snapshot.version {
278 return false;
279 }
280 if !latest_snapshot
281 .version
282 .observed_all(&next_old_event.new_snapshot_version)
283 {
284 return false;
285 }
286
287 let a_is_predicted = matches!(
288 self.event.as_ref(),
289 zeta_prompt::Event::BufferChange {
290 predicted: true,
291 ..
292 }
293 );
294 let b_is_predicted = matches!(
295 next_old_event.event.as_ref(),
296 zeta_prompt::Event::BufferChange {
297 predicted: true,
298 ..
299 }
300 );
301
302 // If events come from the same source (both predicted or both manual) then
303 // we would have coalesced them already.
304 if a_is_predicted == b_is_predicted {
305 return false;
306 }
307
308 let left_range = self.total_edit_range.to_point(latest_snapshot);
309 let right_range = next_old_event.total_edit_range.to_point(latest_snapshot);
310 let latest_range = latest_edit_range.to_point(latest_snapshot);
311
312 // Events near to the latest edit are not merged if their sources differ.
313 if lines_between_ranges(&left_range, &latest_range)
314 .min(lines_between_ranges(&right_range, &latest_range))
315 <= CHANGE_GROUPING_LINE_SPAN
316 {
317 return false;
318 }
319
320 // Events that are distant from each other are not merged.
321 if lines_between_ranges(&left_range, &right_range) > CHANGE_GROUPING_LINE_SPAN {
322 return false;
323 }
324
325 true
326 }
327}
328
329fn lines_between_ranges(left: &Range<Point>, right: &Range<Point>) -> u32 {
330 if left.start > right.end {
331 return left.start.row - right.end.row;
332 }
333 if right.start > left.end {
334 return right.start.row - left.end.row;
335 }
336 0
337}
338
339fn push_recent_file(files: &mut VecDeque<RecentFile>, mut file: RecentFile) {
340 if let Some(ix) = files.iter().position(|probe| probe.path == file.path)
341 && let Some(previous) = files.remove(ix)
342 && file.cursor_position.is_none()
343 {
344 file.cursor_position = previous.cursor_position;
345 }
346 files.push_front(file);
347 files.truncate(RECENT_PATH_COUNT_MAX);
348}
349
350struct ProjectState {
351 events: VecDeque<StoredEvent>,
352 last_event: Option<LastEvent>,
353 next_last_event_seq: u64,
354 recently_viewed_files: VecDeque<RecentFile>,
355 recently_opened_files: VecDeque<RecentFile>,
356 registered_buffers: HashMap<gpui::EntityId, RegisteredBuffer>,
357 file_contexts: HashMap<ProjectPath, WeakEntity<StoredFileContext>>,
358 current_prediction: Option<CurrentEditPrediction>,
359 last_edit_source: Option<BufferEditSource>,
360 next_pending_prediction_id: usize,
361 pending_predictions: ArrayVec<PendingPrediction, 2, u8>,
362 pending_prediction_captures: Vec<PendingPredictionCapture>,
363 debug_tx: Option<mpsc::UnboundedSender<DebugEvent>>,
364 last_edit_prediction_refresh: Option<(EntityId, Instant)>,
365 cancelled_predictions: HashSet<usize>,
366 context: Entity<RelatedExcerptStore>,
367 license_detection_watchers: HashMap<WorktreeId, Rc<LicenseDetectionWatcher>>,
368 _subscriptions: [gpui::Subscription; 2],
369 copilot: Option<Entity<Copilot>>,
370}
371
372impl ProjectState {
373 pub fn events(&self, cx: &App) -> Vec<StoredEvent> {
374 self.events
375 .iter()
376 .cloned()
377 .chain(self.last_event.as_ref().iter().flat_map(|event| {
378 let (one, two) = event.split_by_pause();
379 let one = one.finalize(&self.license_detection_watchers, cx);
380 let two = two.and_then(|two| two.finalize(&self.license_detection_watchers, cx));
381 one.into_iter().chain(two)
382 }))
383 .collect()
384 }
385
386 fn cancel_pending_prediction(
387 &mut self,
388 pending_prediction: PendingPrediction,
389 cx: &mut Context<EditPredictionStore>,
390 ) {
391 self.cancelled_predictions.insert(pending_prediction.id);
392
393 if pending_prediction.drop_on_cancel {
394 drop(pending_prediction.task);
395 } else {
396 cx.spawn(async move |this, cx| {
397 let Some((prediction_id, model_version)) = pending_prediction.task.await else {
398 return;
399 };
400
401 this.update(cx, |this, cx| {
402 this.reject_prediction(
403 prediction_id,
404 EditPredictionRejectReason::Canceled,
405 false,
406 model_version,
407 None,
408 cx,
409 );
410 })
411 .ok();
412 })
413 .detach()
414 }
415 }
416
417 fn active_buffer(
418 &self,
419 project: &Entity<Project>,
420 cx: &App,
421 ) -> Option<(Entity<Buffer>, Option<Anchor>)> {
422 let project = project.read(cx);
423 let active_path = project.path_for_entry(project.active_entry()?, cx)?;
424 let active_buffer = project.buffer_store().read(cx).get_by_path(&active_path)?;
425 let registered_buffer = self.registered_buffers.get(&active_buffer.entity_id())?;
426 Some((active_buffer, registered_buffer.last_position))
427 }
428
429 fn file_context_for_path(
430 &mut self,
431 path: ProjectPath,
432 cx: &mut Context<EditPredictionStore>,
433 ) -> Entity<StoredFileContext> {
434 if let Some(context) = self
435 .file_contexts
436 .get_mut(&path)
437 .and_then(|entry| entry.upgrade())
438 {
439 context
440 } else {
441 let context = cx.new(|_| StoredFileContext {
442 uncommitted_diff: None,
443 git_changed_file_sets: None,
444 git_changed_file_sets_task: None,
445 });
446 self.file_contexts.insert(path, context.downgrade());
447 context
448 }
449 }
450
451 fn update_recent_file_cursor(&mut self, path: &Path, cursor_position: usize) {
452 for file in &mut self.recently_opened_files {
453 if file.path.as_ref() == path && file.cursor_position.is_none() {
454 file.cursor_position = Some(cursor_position);
455 }
456 }
457 for file in &mut self.recently_viewed_files {
458 if file.path.as_ref() == path {
459 file.cursor_position = Some(cursor_position);
460 }
461 }
462 }
463
464 fn finalize_last_event(&mut self, cx: &mut Context<EditPredictionStore>) {
465 let Some(last_event) = self.last_event.take() else {
466 return;
467 };
468 let event = last_event.finalize(&self.license_detection_watchers, cx);
469
470 for capture in &mut self.pending_prediction_captures {
471 capture.try_record_future_event(
472 &last_event,
473 event.as_ref(),
474 &self.license_detection_watchers,
475 cx,
476 );
477 }
478
479 let Some(event) = event else {
480 return;
481 };
482 if self.events.len() + 1 >= EVENT_COUNT_MAX {
483 self.events.pop_front();
484 }
485 self.events.push_back(event);
486 }
487
488 fn clear_history(&mut self) {
489 self.events.clear();
490 self.last_event.take();
491 for capture in &mut self.pending_prediction_captures {
492 capture.sample_data = None;
493 }
494 }
495}
496
497#[derive(Debug, Clone)]
498struct CurrentEditPrediction {
499 pub requested_by: EntityId,
500 pub prediction: EditPrediction,
501 pub was_shown: bool,
502 pub shown_with: Option<edit_prediction_types::SuggestionDisplayType>,
503 pub e2e_latency: std::time::Duration,
504}
505
506impl CurrentEditPrediction {
507 fn should_replace_prediction(&self, old_prediction: &Self, cx: &App) -> bool {
508 let Some(new_edits) = self
509 .prediction
510 .interpolate(&self.prediction.buffer.read(cx))
511 else {
512 return false;
513 };
514
515 if self.prediction.buffer != old_prediction.prediction.buffer {
516 return true;
517 }
518
519 let Some(old_edits) = old_prediction
520 .prediction
521 .interpolate(&old_prediction.prediction.buffer.read(cx))
522 else {
523 return true;
524 };
525
526 // This reduces the occurrence of UI thrash from replacing edits
527 //
528 // TODO: This is fairly arbitrary - should have a more general heuristic that handles multiple edits.
529 if self.requested_by == self.prediction.buffer.entity_id()
530 && self.requested_by == old_prediction.prediction.buffer.entity_id()
531 && old_edits.len() == 1
532 && new_edits.len() == 1
533 {
534 let (old_range, old_text) = &old_edits[0];
535 let (new_range, new_text) = &new_edits[0];
536 new_range == old_range && new_text.starts_with(old_text.as_ref())
537 } else {
538 true
539 }
540 }
541}
542
543const DIAGNOSTIC_LINES_RANGE: u32 = 20;
544
545#[derive(Debug)]
546struct PendingPrediction {
547 id: usize,
548 task: Task<Option<(EditPredictionId, Option<String>)>>,
549 /// If true, the task is dropped immediately on cancel (cancelling the HTTP request).
550 /// If false, the task is awaited to completion so rejection can be reported.
551 drop_on_cancel: bool,
552}
553
554/// A prediction from the perspective of a buffer.
555#[derive(Debug)]
556enum BufferEditPrediction<'a> {
557 Local { prediction: &'a EditPrediction },
558 Jump { prediction: &'a EditPrediction },
559}
560
561#[cfg(test)]
562impl std::ops::Deref for BufferEditPrediction<'_> {
563 type Target = EditPrediction;
564
565 fn deref(&self) -> &Self::Target {
566 match self {
567 BufferEditPrediction::Local { prediction } => prediction,
568 BufferEditPrediction::Jump { prediction } => prediction,
569 }
570 }
571}
572
573struct PendingPredictionCapture {
574 request_id: EditPredictionId,
575 edited_buffer_id: EntityId,
576 editable_anchor_range: Range<Anchor>,
577 editable_region_before_prediction: String,
578 predicted_editable_region: String,
579 ts_error_count_before_prediction: usize,
580 ts_error_count_after_prediction: usize,
581 organization_id: Option<OrganizationId>,
582 can_collect_data: bool,
583 is_in_open_source_repo: bool,
584 sample_data: Option<PendingPredictionCaptureSampleData>,
585 model_version: Option<String>,
586 enqueued_at: Instant,
587 last_edit_at: Instant,
588 e2e_latency: std::time::Duration,
589}
590
591struct PendingPredictionCaptureSampleData {
592 context_task: Task<Result<CapturedPredictionContext>>,
593 editable_path: Arc<Path>,
594 editable_offset_range: Range<usize>,
595 next_edit_cursor_offset: Option<usize>,
596 future_edit_history_events: Vec<Arc<zeta_prompt::Event>>,
597 navigation_history: VecDeque<RecentFile>,
598 edit_events_before_quiescence: u32,
599 prompt_history_boundary: Option<PromptHistoryBoundary>,
600}
601
602/// Marks where the prompt's edit history ended. Sample data may only include
603/// content the user produced after this point.
604struct PromptHistoryBoundary {
605 /// The seq of the first event this capture is expected to observe: the
606 /// event that was pending when the prediction was requested, or the next
607 /// event to be created if none was pending. Observing a later seq first
608 /// means events were lost while the prediction request was in flight.
609 first_event_seq: u64,
610 /// The prompt's end snapshot within the event that was pending when the
611 /// prediction was requested, if any. The first observed event is trimmed
612 /// to its suffix after this snapshot.
613 snapshot: Option<TextBufferSnapshot>,
614}
615
616impl PendingPredictionCapture {
617 /// Records the project's last event (pending or finalizing) into this
618 /// sample's future edit history. Returns false if the sample must be
619 /// dropped because its future history can't be captured accurately.
620 fn try_record_future_event(
621 &mut self,
622 last_event: &LastEvent,
623 finalized_event: Option<&StoredEvent>,
624 license_detection_watchers: &HashMap<WorktreeId, Rc<LicenseDetectionWatcher>>,
625 cx: &App,
626 ) {
627 let Some(sample) = &mut self.sample_data else {
628 return;
629 };
630 let boundary = sample.prompt_history_boundary.take();
631 let suffix_snapshot = match &boundary {
632 Some(boundary) => {
633 if last_event.seq != boundary.first_event_seq {
634 // Events were finalized before this capture was enqueued,
635 // so events are missing from the future history.
636 self.sample_data.take();
637 return;
638 }
639 boundary.snapshot.as_ref()
640 }
641 None => None,
642 };
643
644 let event = match suffix_snapshot {
645 Some(snapshot) => {
646 let suffix = last_event
647 .suffix_after(snapshot)
648 .and_then(|suffix| suffix.finalize(license_detection_watchers, cx));
649 let Some(suffix) = suffix else {
650 return;
651 };
652 suffix.event
653 }
654 None => match finalized_event {
655 Some(event) => event.event.clone(),
656 None => return,
657 },
658 };
659
660 if !event.in_open_source_repo() {
661 self.sample_data.take();
662 return;
663 }
664 sample.edit_events_before_quiescence += 1;
665 if sample.future_edit_history_events.len() < EDIT_PREDICTION_CAPTURE_MAX_FUTURE_EVENTS {
666 sample.future_edit_history_events.push(event);
667 }
668 }
669}
670
671struct RegisteredBuffer {
672 file: Option<Arc<dyn File>>,
673 snapshot: TextBufferSnapshot,
674 last_position: Option<Anchor>,
675 _subscriptions: [gpui::Subscription; 2],
676}
677
678#[derive(Clone)]
679struct LastEvent {
680 /// Project-wide monotonic sequence number identifying this event.
681 seq: u64,
682 old_snapshot: TextBufferSnapshot,
683 new_snapshot: TextBufferSnapshot,
684 old_file: Option<Arc<dyn File>>,
685 new_file: Option<Arc<dyn File>>,
686 latest_edit_range: Range<Anchor>,
687 total_edit_range: Range<Anchor>,
688 total_edit_range_at_last_pause_boundary: Option<Range<Anchor>>,
689 predicted: bool,
690 snapshot_after_last_editing_pause: Option<TextBufferSnapshot>,
691 last_edit_time: Option<Instant>,
692 file_context: Option<Entity<StoredFileContext>>,
693}
694
695impl LastEvent {
696 pub fn finalize(
697 &self,
698 license_detection_watchers: &HashMap<WorktreeId, Rc<LicenseDetectionWatcher>>,
699 cx: &App,
700 ) -> Option<StoredEvent> {
701 let path = buffer_path_with_id_fallback(self.new_file.as_ref(), &self.new_snapshot, cx);
702 let old_path = buffer_path_with_id_fallback(self.old_file.as_ref(), &self.old_snapshot, cx);
703
704 let in_open_source_repo =
705 [self.new_file.as_ref(), self.old_file.as_ref()]
706 .iter()
707 .all(|file| {
708 file.is_some_and(|file| {
709 license_detection_watchers
710 .get(&file.worktree_id(cx))
711 .is_some_and(|watcher| watcher.is_project_open_source())
712 })
713 });
714
715 let (diff, old_range, new_range) = compute_diff_between_snapshots_in_range(
716 &self.old_snapshot,
717 &self.new_snapshot,
718 &self.total_edit_range,
719 )?;
720
721 if path == old_path && diff.is_empty() {
722 None
723 } else {
724 Some(StoredEvent {
725 event: Arc::new(zeta_prompt::Event::BufferChange {
726 old_path,
727 path,
728 diff,
729 old_range,
730 new_range: new_range.clone(),
731 in_open_source_repo,
732 predicted: self.predicted,
733 }),
734 old_snapshot: self.old_snapshot.clone(),
735 new_snapshot_version: self.new_snapshot.version.clone(),
736 total_edit_range: self.new_snapshot.anchor_before(new_range.start)
737 ..self.new_snapshot.anchor_before(new_range.end),
738 file_context: self.file_context.clone(),
739 })
740 }
741 }
742
743 pub fn split_by_pause(&self) -> (LastEvent, Option<LastEvent>) {
744 let Some(boundary_snapshot) = self.snapshot_after_last_editing_pause.as_ref() else {
745 return (self.clone(), None);
746 };
747
748 let Some(after) = self.suffix_after(boundary_snapshot) else {
749 return (self.clone(), None);
750 };
751
752 let total_edit_range_before_pause = self
753 .total_edit_range_at_last_pause_boundary
754 .clone()
755 .unwrap_or_else(|| self.total_edit_range.clone());
756
757 let before = LastEvent {
758 new_snapshot: boundary_snapshot.clone(),
759 latest_edit_range: total_edit_range_before_pause.clone(),
760 total_edit_range: total_edit_range_before_pause,
761 total_edit_range_at_last_pause_boundary: None,
762 snapshot_after_last_editing_pause: None,
763 ..self.clone()
764 };
765
766 (before, Some(after))
767 }
768
769 /// The portion of this event that happened after `boundary_snapshot`, or
770 /// None if the buffer hasn't changed since.
771 pub fn suffix_after(&self, boundary_snapshot: &TextBufferSnapshot) -> Option<LastEvent> {
772 let total_edit_range =
773 compute_total_edit_range_between_snapshots(boundary_snapshot, &self.new_snapshot)?;
774 Some(LastEvent {
775 old_snapshot: boundary_snapshot.clone(),
776 latest_edit_range: total_edit_range.clone(),
777 total_edit_range,
778 total_edit_range_at_last_pause_boundary: None,
779 snapshot_after_last_editing_pause: None,
780 ..self.clone()
781 })
782 }
783}
784
785fn compute_total_edit_range_between_snapshots(
786 old_snapshot: &TextBufferSnapshot,
787 new_snapshot: &TextBufferSnapshot,
788) -> Option<Range<Anchor>> {
789 let edits: Vec<Edit<usize>> = new_snapshot
790 .edits_since::<usize>(&old_snapshot.version)
791 .collect();
792
793 let (first_edit, last_edit) = edits.first().zip(edits.last())?;
794 let new_start_point = new_snapshot.offset_to_point(first_edit.new.start);
795 let new_end_point = new_snapshot.offset_to_point(last_edit.new.end);
796
797 Some(new_snapshot.anchor_before(new_start_point)..new_snapshot.anchor_before(new_end_point))
798}
799
800fn compute_old_range_for_new_range(
801 old_snapshot: &TextBufferSnapshot,
802 new_snapshot: &TextBufferSnapshot,
803 total_edit_range: &Range<Anchor>,
804) -> Option<Range<Point>> {
805 let new_start_offset = total_edit_range.start.to_offset(new_snapshot);
806 let new_end_offset = total_edit_range.end.to_offset(new_snapshot);
807
808 let edits: Vec<Edit<usize>> = new_snapshot
809 .edits_since::<usize>(&old_snapshot.version)
810 .collect();
811 let mut old_start_offset = None;
812 let mut old_end_offset = None;
813 let mut delta: isize = 0;
814
815 for edit in &edits {
816 if old_start_offset.is_none() && new_start_offset <= edit.new.end {
817 old_start_offset = Some(if new_start_offset < edit.new.start {
818 new_start_offset.checked_add_signed(-delta)?
819 } else {
820 edit.old.start
821 });
822 }
823
824 if old_end_offset.is_none() && new_end_offset <= edit.new.end {
825 old_end_offset = Some(if new_end_offset < edit.new.start {
826 new_end_offset.checked_add_signed(-delta)?
827 } else {
828 edit.old.end
829 });
830 }
831
832 delta += edit.new.len() as isize - edit.old.len() as isize;
833 }
834
835 let old_start_offset =
836 old_start_offset.unwrap_or_else(|| new_start_offset.saturating_add_signed(-delta));
837 let old_end_offset =
838 old_end_offset.unwrap_or_else(|| new_end_offset.saturating_add_signed(-delta));
839
840 Some(
841 old_snapshot.offset_to_point(old_start_offset)
842 ..old_snapshot.offset_to_point(old_end_offset),
843 )
844}
845
846fn compute_diff_between_snapshots_in_range(
847 old_snapshot: &TextBufferSnapshot,
848 new_snapshot: &TextBufferSnapshot,
849 total_edit_range: &Range<Anchor>,
850) -> Option<(String, Range<usize>, Range<usize>)> {
851 let new_start_offset = total_edit_range.start.to_offset(new_snapshot);
852 let new_end_offset = total_edit_range.end.to_offset(new_snapshot);
853 let new_start_point = new_snapshot.offset_to_point(new_start_offset);
854 let new_end_point = new_snapshot.offset_to_point(new_end_offset);
855 let old_range = compute_old_range_for_new_range(old_snapshot, new_snapshot, total_edit_range)?;
856 let old_start_point = old_range.start;
857 let old_end_point = old_range.end;
858 let old_start_offset = old_snapshot.point_to_offset(old_start_point);
859 let old_end_offset = old_snapshot.point_to_offset(old_end_point);
860
861 const CONTEXT_LINES: u32 = 3;
862
863 let old_context_start_row = old_start_point.row.saturating_sub(CONTEXT_LINES);
864 let new_context_start_row = new_start_point.row.saturating_sub(CONTEXT_LINES);
865 let old_context_end_row =
866 (old_end_point.row + 1 + CONTEXT_LINES).min(old_snapshot.max_point().row);
867 let new_context_end_row =
868 (new_end_point.row + 1 + CONTEXT_LINES).min(new_snapshot.max_point().row);
869
870 let old_start_line_offset = old_snapshot.point_to_offset(Point::new(old_context_start_row, 0));
871 let new_start_line_offset = new_snapshot.point_to_offset(Point::new(new_context_start_row, 0));
872 let old_end_line_offset = old_snapshot
873 .point_to_offset(Point::new(old_context_end_row + 1, 0).min(old_snapshot.max_point()));
874 let new_end_line_offset = new_snapshot
875 .point_to_offset(Point::new(new_context_end_row + 1, 0).min(new_snapshot.max_point()));
876 let old_edit_range = old_start_line_offset..old_end_line_offset;
877 let new_edit_range = new_start_line_offset..new_end_line_offset;
878
879 if new_edit_range.len() > EDIT_HISTORY_DIFF_SIZE_LIMIT
880 || old_edit_range.len() > EDIT_HISTORY_DIFF_SIZE_LIMIT
881 {
882 return None;
883 }
884
885 let old_region_text: String = old_snapshot.text_for_range(old_edit_range).collect();
886 let new_region_text: String = new_snapshot.text_for_range(new_edit_range).collect();
887
888 let diff = language::unified_diff_with_offsets(
889 &old_region_text,
890 &new_region_text,
891 old_context_start_row,
892 new_context_start_row,
893 );
894
895 Some((
896 diff,
897 old_start_offset..old_end_offset,
898 new_start_offset..new_end_offset,
899 ))
900}
901
902pub(crate) fn buffer_path_with_id_fallback(
903 file: Option<&Arc<dyn File>>,
904 snapshot: &TextBufferSnapshot,
905 cx: &App,
906) -> Arc<Path> {
907 if let Some(file) = file {
908 file.full_path(cx).into()
909 } else {
910 Path::new(&format!("untitled-{}", snapshot.remote_id())).into()
911 }
912}
913
914fn predict_edits_request_trigger_from_editor_trigger(
915 trigger: EditPredictionRequestTrigger,
916) -> PredictEditsRequestTrigger {
917 match trigger {
918 EditPredictionRequestTrigger::DiagnosticNavigation => {
919 PredictEditsRequestTrigger::DiagnosticNavigation
920 }
921 EditPredictionRequestTrigger::Explicit => PredictEditsRequestTrigger::Explicit,
922 EditPredictionRequestTrigger::BufferEdit => PredictEditsRequestTrigger::BufferEdit,
923 EditPredictionRequestTrigger::LSPCompletionAccepted => {
924 PredictEditsRequestTrigger::LSPCompletionAccepted
925 }
926 EditPredictionRequestTrigger::PredictionAccepted => {
927 PredictEditsRequestTrigger::PredictionAccepted
928 }
929 EditPredictionRequestTrigger::PredictionPartiallyAccepted => {
930 PredictEditsRequestTrigger::PredictionPartiallyAccepted
931 }
932 EditPredictionRequestTrigger::EditorCreated => PredictEditsRequestTrigger::EditorCreated,
933 EditPredictionRequestTrigger::ProviderChanged => {
934 PredictEditsRequestTrigger::ProviderChanged
935 }
936 EditPredictionRequestTrigger::UserInfoChanged => {
937 PredictEditsRequestTrigger::UserInfoChanged
938 }
939 EditPredictionRequestTrigger::VimModeChanged => PredictEditsRequestTrigger::VimModeChanged,
940 EditPredictionRequestTrigger::SettingsChanged => {
941 PredictEditsRequestTrigger::SettingsChanged
942 }
943 EditPredictionRequestTrigger::Other => PredictEditsRequestTrigger::Other,
944 }
945}
946
947impl EditPredictionStore {
948 pub fn try_global(cx: &App) -> Option<Entity<Self>> {
949 cx.try_global::<EditPredictionStoreGlobal>()
950 .map(|global| global.0.clone())
951 }
952
953 pub fn global(
954 client: &Arc<Client>,
955 user_store: &Entity<UserStore>,
956 cx: &mut App,
957 ) -> Entity<Self> {
958 cx.try_global::<EditPredictionStoreGlobal>()
959 .map(|global| global.0.clone())
960 .unwrap_or_else(|| {
961 let ep_store = cx.new(|cx| Self::new(client.clone(), user_store.clone(), cx));
962 cx.set_global(EditPredictionStoreGlobal(ep_store.clone()));
963 ep_store
964 })
965 }
966
967 pub fn new(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut Context<Self>) -> Self {
968 let llm_token = global_llm_token(cx);
969 let legacy_data_collection_enabled = Self::load_legacy_data_collection_enabled(cx);
970
971 let (reject_tx, reject_rx) = mpsc::unbounded();
972 cx.background_spawn({
973 let client = client.clone();
974 let llm_token = llm_token.clone();
975 let app_version = AppVersion::global(cx);
976 let background_executor = cx.background_executor().clone();
977 async move {
978 Self::handle_rejected_predictions(
979 reject_rx,
980 client,
981 llm_token,
982 app_version,
983 background_executor,
984 )
985 .await
986 }
987 })
988 .detach();
989
990 let (settled_predictions_tx, settled_predictions_rx) = mpsc::unbounded();
991 cx.spawn({
992 let client = client.clone();
993 let llm_token = llm_token.clone();
994 let app_version = AppVersion::global(cx);
995 async move |this, cx| {
996 Self::run_settled_predictions_worker(
997 this,
998 settled_predictions_rx,
999 client,
1000 llm_token,
1001 app_version,
1002 cx,
1003 )
1004 .await;
1005 }
1006 })
1007 .detach();
1008
1009 let mut current_user = user_store.read(cx).watch_current_user();
1010 let fetch_experiments_task = cx.spawn(async move |this, cx| {
1011 while current_user.borrow().is_none() {
1012 current_user.next().await;
1013 }
1014
1015 this.update(cx, |this, cx| {
1016 if cx.is_staff() {
1017 this.refresh_available_experiments(cx);
1018 }
1019 })
1020 .log_err();
1021 });
1022
1023 let credentials_provider = zed_credentials_provider::global(cx);
1024
1025 let this = Self {
1026 projects: HashMap::default(),
1027 client,
1028 user_store,
1029 llm_token,
1030 _fetch_experiments_task: fetch_experiments_task,
1031 update_required: false,
1032 edit_prediction_model: EditPredictionModel::Zeta,
1033 zeta2_raw_config: Self::zeta2_raw_config_from_env(),
1034 request_backoff_until: None,
1035 preferred_experiment: None,
1036 available_experiments: Vec::new(),
1037 mercury: Mercury::new(cx),
1038 legacy_data_collection_enabled,
1039
1040 reject_predictions_tx: reject_tx,
1041 settled_predictions_tx,
1042 rated_predictions: Default::default(),
1043 rateable_predictions: Default::default(),
1044 #[cfg(test)]
1045 settled_event_callback: None,
1046
1047 credentials_provider,
1048 };
1049
1050 this
1051 }
1052
1053 fn zeta2_raw_config_from_env() -> Option<Zeta2RawConfig> {
1054 let version_str = env::var("ZED_ZETA_FORMAT").ok()?;
1055 let format = ZetaFormat::parse(&version_str).ok()?;
1056 let model_id = env::var("ZED_ZETA_MODEL").ok();
1057 let environment = env::var("ZED_ZETA_ENVIRONMENT").ok();
1058 Some(Zeta2RawConfig {
1059 model_id,
1060 environment,
1061 format,
1062 })
1063 }
1064
1065 pub fn set_edit_prediction_model(&mut self, model: EditPredictionModel) {
1066 self.edit_prediction_model = model;
1067 }
1068
1069 pub fn set_zeta2_raw_config(&mut self, config: Zeta2RawConfig) {
1070 self.zeta2_raw_config = Some(config);
1071 }
1072
1073 pub fn zeta2_raw_config(&self) -> Option<&Zeta2RawConfig> {
1074 self.zeta2_raw_config.as_ref()
1075 }
1076
1077 pub(crate) fn back_off_requests_after_timeout(&mut self, cx: &mut Context<Self>) {
1078 self.request_backoff_until = Some(cx.background_executor().now() + REQUEST_TIMEOUT_BACKOFF);
1079 log::info!(
1080 "Backing off edit prediction requests for {:?} after Cloud timeout",
1081 REQUEST_TIMEOUT_BACKOFF
1082 );
1083 }
1084
1085 fn request_backoff_active(&mut self, cx: &App) -> bool {
1086 let Some(backoff_until) = self.request_backoff_until else {
1087 return false;
1088 };
1089
1090 if cx.background_executor().now() < backoff_until {
1091 true
1092 } else {
1093 self.request_backoff_until = None;
1094 false
1095 }
1096 }
1097
1098 pub fn preferred_experiment(&self) -> Option<&str> {
1099 self.preferred_experiment.as_deref()
1100 }
1101
1102 pub fn set_preferred_experiment(&mut self, experiment: Option<String>) {
1103 self.preferred_experiment = experiment;
1104 }
1105
1106 pub fn available_experiments(&self) -> &[String] {
1107 &self.available_experiments
1108 }
1109
1110 pub fn active_experiment(&self) -> Option<&str> {
1111 self.preferred_experiment.as_deref().or_else(|| {
1112 self.rateable_predictions
1113 .iter()
1114 .find_map(|p| p.model_version.as_ref())
1115 .and_then(|model_version| model_version.strip_prefix("zeta2:"))
1116 })
1117 }
1118
1119 pub fn refresh_available_experiments(&mut self, cx: &mut Context<Self>) {
1120 let client = self.client.clone();
1121 let llm_token = self.llm_token.clone();
1122 let app_version = AppVersion::global(cx);
1123 let is_jumps_api = cx.has_flag::<EditPredictionJumpsFeatureFlag>();
1124 let organization_id = self
1125 .user_store
1126 .read(cx)
1127 .current_organization()
1128 .map(|organization| organization.id.clone());
1129
1130 cx.spawn(async move |this, cx| {
1131 let experiments = cx
1132 .background_spawn(async move {
1133 let organization_id =
1134 organization_id.ok_or_else(|| anyhow!("No organization selected."))?;
1135 let url = client.http_client().build_zed_llm_url(
1136 "/edit_prediction_experiments",
1137 &[("is_jumps_api", if is_jumps_api { "true" } else { "false" })],
1138 )?;
1139 let mut response = client
1140 .authenticated_llm_request(&llm_token, organization_id, |token| {
1141 Ok(http_client::Request::builder()
1142 .method(Method::GET)
1143 .uri(url.as_ref())
1144 .header("Authorization", format!("Bearer {token}"))
1145 .header(ZED_VERSION_HEADER_NAME, app_version.to_string())
1146 .body(Default::default())?)
1147 })
1148 .await?;
1149 if response.status().is_success() {
1150 let mut body = Vec::new();
1151 response.body_mut().read_to_end(&mut body).await?;
1152 let experiments: Vec<String> = serde_json::from_slice(&body)?;
1153 Ok(experiments)
1154 } else {
1155 let mut body = String::new();
1156 response.body_mut().read_to_string(&mut body).await?;
1157 anyhow::bail!(
1158 "Failed to fetch experiments: {:?}\nBody: {}",
1159 response.status(),
1160 body
1161 );
1162 }
1163 })
1164 .await?;
1165 this.update(cx, |this, cx| {
1166 this.available_experiments = experiments;
1167 cx.notify();
1168 })?;
1169 anyhow::Ok(())
1170 })
1171 .detach_and_log_err(cx);
1172 }
1173
1174 pub fn icons(&self, cx: &App) -> edit_prediction_types::EditPredictionIconSet {
1175 use ui::IconName;
1176 match self.edit_prediction_model {
1177 EditPredictionModel::Mercury => {
1178 edit_prediction_types::EditPredictionIconSet::new(IconName::Inception)
1179 }
1180 EditPredictionModel::Zeta => {
1181 edit_prediction_types::EditPredictionIconSet::new(IconName::OmegaPredict)
1182 .with_disabled(IconName::OmegaPredictDisabled)
1183 .with_up(IconName::OmegaPredictUp)
1184 .with_down(IconName::OmegaPredictDown)
1185 .with_error(IconName::OmegaPredictError)
1186 }
1187 EditPredictionModel::Fim { .. } => {
1188 let settings = &all_language_settings(None, cx).edit_predictions;
1189 match settings.provider {
1190 EditPredictionProvider::Ollama => {
1191 edit_prediction_types::EditPredictionIconSet::new(IconName::AiOllama)
1192 }
1193 _ => {
1194 edit_prediction_types::EditPredictionIconSet::new(IconName::AiOpenAiCompat)
1195 }
1196 }
1197 }
1198 }
1199 }
1200
1201 pub fn has_mercury_api_token(&self, cx: &App) -> bool {
1202 self.mercury.api_token.read(cx).has_key()
1203 }
1204
1205 pub fn mercury_has_payment_required_error(&self) -> bool {
1206 self.mercury.has_payment_required_error()
1207 }
1208
1209 pub fn clear_history(&mut self) {
1210 for project_state in self.projects.values_mut() {
1211 project_state.clear_history();
1212 }
1213 }
1214
1215 pub fn clear_history_for_project(&mut self, project: &Entity<Project>) {
1216 if let Some(project_state) = self.projects.get_mut(&project.entity_id()) {
1217 project_state.clear_history();
1218 }
1219 }
1220
1221 pub fn edit_history_for_project(
1222 &self,
1223 project: &Entity<Project>,
1224 cx: &App,
1225 ) -> Vec<StoredEvent> {
1226 self.projects
1227 .get(&project.entity_id())
1228 .map(|project_state| project_state.events(cx))
1229 .unwrap_or_default()
1230 }
1231
1232 pub fn context_for_project<'a>(
1233 &'a self,
1234 project: &Entity<Project>,
1235 cx: &'a mut App,
1236 ) -> Vec<RelatedFile> {
1237 self.projects
1238 .get(&project.entity_id())
1239 .map(|project_state| {
1240 project_state.context.update(cx, |context, cx| {
1241 context
1242 .related_files_with_buffers(cx)
1243 .map(|(mut related_file, buffer)| {
1244 related_file.in_open_source_repo = buffer
1245 .read(cx)
1246 .file()
1247 .map_or(false, |file| self.is_file_open_source(&project, file, cx));
1248 related_file
1249 })
1250 .collect()
1251 })
1252 })
1253 .unwrap_or_default()
1254 }
1255
1256 pub fn copilot_for_project(&self, project: &Entity<Project>) -> Option<Entity<Copilot>> {
1257 self.projects
1258 .get(&project.entity_id())
1259 .and_then(|project| project.copilot.clone())
1260 }
1261
1262 pub fn start_copilot_for_project(
1263 &mut self,
1264 project: &Entity<Project>,
1265 cx: &mut Context<Self>,
1266 ) -> Option<Entity<Copilot>> {
1267 if DisableAiSettings::get(None, cx).disable_ai {
1268 return None;
1269 }
1270 let state = self.get_or_init_project(project, cx);
1271
1272 if state.copilot.is_some() {
1273 return state.copilot.clone();
1274 }
1275 let _project = project.clone();
1276 let project = project.read(cx);
1277
1278 let node = project.node_runtime().cloned();
1279 if let Some(node) = node {
1280 let next_id = project.languages().next_language_server_id();
1281 let fs = project.fs().clone();
1282
1283 let copilot = cx.new(|cx| Copilot::new(Some(_project), next_id, fs, node, cx));
1284 state.copilot = Some(copilot.clone());
1285 Some(copilot)
1286 } else {
1287 None
1288 }
1289 }
1290
1291 pub fn context_for_project_with_buffers<'a>(
1292 &'a self,
1293 project: &Entity<Project>,
1294 cx: &'a mut App,
1295 ) -> Vec<(RelatedFile, Entity<Buffer>)> {
1296 self.projects
1297 .get(&project.entity_id())
1298 .map(|project| {
1299 project.context.update(cx, |context, cx| {
1300 context.related_files_with_buffers(cx).collect()
1301 })
1302 })
1303 .unwrap_or_default()
1304 }
1305
1306 pub fn usage(&self, cx: &App) -> Option<EditPredictionUsage> {
1307 if matches!(self.edit_prediction_model, EditPredictionModel::Zeta) {
1308 self.user_store.read(cx).edit_prediction_usage()
1309 } else {
1310 None
1311 }
1312 }
1313
1314 pub fn register_project(&mut self, project: &Entity<Project>, cx: &mut Context<Self>) {
1315 self.get_or_init_project(project, cx);
1316 }
1317
1318 pub fn register_buffer(
1319 &mut self,
1320 buffer: &Entity<Buffer>,
1321 project: &Entity<Project>,
1322 cx: &mut Context<Self>,
1323 ) {
1324 let opened_path = buffer
1325 .read(cx)
1326 .file()
1327 .map(|file| ProjectPath::from_file(file.as_ref(), cx));
1328 let project_state = self.get_or_init_project(project, cx);
1329 if let Some(path) = opened_path {
1330 push_recent_file(
1331 &mut project_state.recently_opened_files,
1332 RecentFile {
1333 path: path.path.as_std_path().into(),
1334 cursor_position: None,
1335 },
1336 );
1337 }
1338 Self::register_buffer_impl(project_state, buffer, project, cx);
1339 }
1340
1341 fn ensure_git_changed_file_sets_loading(
1342 file_context: &Entity<StoredFileContext>,
1343 project: &Entity<Project>,
1344 project_path: &ProjectPath,
1345 cx: &mut Context<Self>,
1346 ) {
1347 let should_start = file_context.update(cx, |file_context, _| {
1348 file_context.git_changed_file_sets.is_none()
1349 && file_context.git_changed_file_sets_task.is_none()
1350 });
1351 if !should_start {
1352 return;
1353 }
1354
1355 let Some((repository, repo_path)) = project
1356 .read(cx)
1357 .git_store()
1358 .read(cx)
1359 .repository_and_path_for_project_path(project_path, cx)
1360 else {
1361 file_context.update(cx, |file_context, _| {
1362 file_context.git_changed_file_sets = Some(Arc::default());
1363 });
1364 return;
1365 };
1366
1367 let receiver = repository.update(cx, |repository, _| {
1368 repository
1369 .file_history_changed_files(vec![repo_path], GIT_CHANGED_FILE_SETS_COMMIT_LIMIT)
1370 });
1371 let task = cx.spawn({
1372 let file_context = file_context.downgrade();
1373 async move |_, cx| {
1374 let result = receiver.await;
1375 let Some(file_context) = file_context.upgrade() else {
1376 return;
1377 };
1378 file_context.update(cx, |file_context, _| {
1379 file_context.git_changed_file_sets = result
1380 .context("failed to receive git changed file sets")
1381 .flatten()
1382 .log_with_level(log::Level::Trace)
1383 .map(|mut file_sets| file_sets.pop().unwrap_or_default())
1384 .context("failed to load git changed file sets")
1385 .map(Arc::new)
1386 .log_err();
1387 file_context.git_changed_file_sets_task = None;
1388 });
1389 }
1390 });
1391 file_context.update(cx, |file_context, _| {
1392 file_context.git_changed_file_sets_task = Some(task);
1393 });
1394 }
1395
1396 fn get_or_init_project(
1397 &mut self,
1398 project: &Entity<Project>,
1399 cx: &mut Context<Self>,
1400 ) -> &mut ProjectState {
1401 let entity_id = project.entity_id();
1402 self.projects
1403 .entry(entity_id)
1404 .or_insert_with(|| ProjectState {
1405 context: {
1406 let related_excerpt_store = cx.new(|cx| RelatedExcerptStore::new(project, cx));
1407 cx.subscribe(&related_excerpt_store, move |this, _, event, _| {
1408 this.handle_excerpt_store_event(entity_id, event);
1409 })
1410 .detach();
1411 related_excerpt_store
1412 },
1413 events: VecDeque::new(),
1414 last_event: None,
1415 next_last_event_seq: 0,
1416 recently_viewed_files: VecDeque::new(),
1417 recently_opened_files: VecDeque::new(),
1418 debug_tx: None,
1419 registered_buffers: HashMap::default(),
1420 file_contexts: HashMap::default(),
1421 current_prediction: None,
1422 last_edit_source: None,
1423 cancelled_predictions: HashSet::default(),
1424 pending_predictions: ArrayVec::new(),
1425 pending_prediction_captures: Vec::new(),
1426 next_pending_prediction_id: 0,
1427 last_edit_prediction_refresh: None,
1428 license_detection_watchers: HashMap::default(),
1429 _subscriptions: [
1430 cx.subscribe(&project, Self::handle_project_event),
1431 cx.observe_release(&project, move |this, _, cx| {
1432 this.projects.remove(&entity_id);
1433 cx.notify();
1434 }),
1435 ],
1436 copilot: None,
1437 })
1438 }
1439
1440 pub fn remove_project(&mut self, project: &Entity<Project>) {
1441 self.projects.remove(&project.entity_id());
1442 }
1443
1444 fn handle_excerpt_store_event(
1445 &mut self,
1446 project_entity_id: EntityId,
1447 event: &RelatedExcerptStoreEvent,
1448 ) {
1449 if let Some(project_state) = self.projects.get(&project_entity_id) {
1450 if let Some(debug_tx) = project_state.debug_tx.clone() {
1451 match event {
1452 RelatedExcerptStoreEvent::StartedRefresh => {
1453 debug_tx
1454 .unbounded_send(DebugEvent::ContextRetrievalStarted(
1455 ContextRetrievalStartedDebugEvent {
1456 project_entity_id: project_entity_id,
1457 timestamp: Instant::now(),
1458 search_prompt: String::new(),
1459 },
1460 ))
1461 .ok();
1462 }
1463 RelatedExcerptStoreEvent::FinishedRefresh {
1464 cache_hit_count,
1465 cache_miss_count,
1466 mean_definition_latency,
1467 max_definition_latency,
1468 } => {
1469 debug_tx
1470 .unbounded_send(DebugEvent::ContextRetrievalFinished(
1471 ContextRetrievalFinishedDebugEvent {
1472 project_entity_id: project_entity_id,
1473 timestamp: Instant::now(),
1474 metadata: vec![
1475 (
1476 "Cache Hits",
1477 format!(
1478 "{}/{}",
1479 cache_hit_count,
1480 cache_hit_count + cache_miss_count
1481 )
1482 .into(),
1483 ),
1484 (
1485 "Max LSP Time",
1486 format!("{} ms", max_definition_latency.as_millis())
1487 .into(),
1488 ),
1489 (
1490 "Mean LSP Time",
1491 format!("{} ms", mean_definition_latency.as_millis())
1492 .into(),
1493 ),
1494 ],
1495 },
1496 ))
1497 .ok();
1498 }
1499 }
1500 }
1501 }
1502 }
1503
1504 pub fn debug_info(
1505 &mut self,
1506 project: &Entity<Project>,
1507 cx: &mut Context<Self>,
1508 ) -> mpsc::UnboundedReceiver<DebugEvent> {
1509 let project_state = self.get_or_init_project(project, cx);
1510 let (debug_watch_tx, debug_watch_rx) = mpsc::unbounded();
1511 project_state.debug_tx = Some(debug_watch_tx);
1512 debug_watch_rx
1513 }
1514
1515 fn handle_project_event(
1516 &mut self,
1517 project: Entity<Project>,
1518 event: &project::Event,
1519 cx: &mut Context<Self>,
1520 ) {
1521 if !is_ep_store_provider(all_language_settings(None, cx).edit_predictions.provider) {
1522 return;
1523 }
1524 // TODO [zeta2] init with recent paths
1525 match event {
1526 project::Event::BufferEdited { source } => {
1527 self.get_or_init_project(&project, cx).last_edit_source = Some(*source);
1528 }
1529 project::Event::ActiveEntryChanged(Some(active_entry_id)) => {
1530 let Some(project_state) = self.projects.get_mut(&project.entity_id()) else {
1531 return;
1532 };
1533 let path = project.read(cx).path_for_entry(*active_entry_id, cx);
1534 if let Some(path) = path {
1535 let cursor_position = project
1536 .read(cx)
1537 .buffer_store()
1538 .read(cx)
1539 .get_by_path(&path)
1540 .and_then(|buffer| {
1541 let position = project_state
1542 .registered_buffers
1543 .get(&buffer.entity_id())?
1544 .last_position?;
1545 Some(position.to_offset(&buffer.read(cx).snapshot()))
1546 });
1547
1548 let recent_file = RecentFile {
1549 path: path.path.as_std_path().into(),
1550 cursor_position,
1551 };
1552 let can_collect_navigation = project_state
1553 .license_detection_watchers
1554 .get(&path.worktree_id)
1555 .is_some_and(|watcher| watcher.is_project_open_source());
1556 for capture in &mut project_state.pending_prediction_captures {
1557 if let Some(sample_data) = capture.sample_data.as_mut() {
1558 if can_collect_navigation {
1559 push_recent_file(
1560 &mut sample_data.navigation_history,
1561 recent_file.clone(),
1562 );
1563 } else {
1564 capture.sample_data = None;
1565 }
1566 }
1567 }
1568 push_recent_file(&mut project_state.recently_viewed_files, recent_file);
1569 }
1570 }
1571 _ => (),
1572 }
1573 }
1574
1575 fn register_buffer_impl<'a>(
1576 project_state: &'a mut ProjectState,
1577 buffer: &Entity<Buffer>,
1578 project: &Entity<Project>,
1579 cx: &mut Context<Self>,
1580 ) -> &'a mut RegisteredBuffer {
1581 let buffer_id = buffer.entity_id();
1582
1583 if let Some(file) = buffer.read(cx).file() {
1584 let worktree_id = file.worktree_id(cx);
1585 if let Some(worktree) = project.read(cx).worktree_for_id(worktree_id, cx) {
1586 project_state
1587 .license_detection_watchers
1588 .entry(worktree_id)
1589 .or_insert_with(|| {
1590 let project_entity_id = project.entity_id();
1591 cx.observe_release(&worktree, move |this, _worktree, _cx| {
1592 let Some(project_state) = this.projects.get_mut(&project_entity_id)
1593 else {
1594 return;
1595 };
1596 project_state
1597 .license_detection_watchers
1598 .remove(&worktree_id);
1599 })
1600 .detach();
1601 Rc::new(LicenseDetectionWatcher::new(&worktree, cx))
1602 });
1603 }
1604 }
1605
1606 match project_state.registered_buffers.entry(buffer_id) {
1607 hash_map::Entry::Occupied(entry) => entry.into_mut(),
1608 hash_map::Entry::Vacant(entry) => {
1609 let buf = buffer.read(cx);
1610 let snapshot = buf.text_snapshot();
1611 let file = buf.file().cloned();
1612 let project_entity_id = project.entity_id();
1613 entry.insert(RegisteredBuffer {
1614 snapshot,
1615 file,
1616 last_position: None,
1617 _subscriptions: [
1618 cx.subscribe(buffer, {
1619 let project = project.downgrade();
1620 move |this, buffer, event, cx| {
1621 if let language::BufferEvent::Edited { source } = event
1622 && let Some(project) = project.upgrade()
1623 {
1624 let project_state = this.get_or_init_project(&project, cx);
1625 project_state.last_edit_source = Some(*source);
1626 this.report_changes_for_buffer(
1627 &buffer,
1628 &project,
1629 false,
1630 source.is_local(),
1631 cx,
1632 );
1633 }
1634 }
1635 }),
1636 cx.observe_release(buffer, move |this, _buffer, _cx| {
1637 let Some(project_state) = this.projects.get_mut(&project_entity_id)
1638 else {
1639 return;
1640 };
1641 project_state.registered_buffers.remove(&buffer_id);
1642 }),
1643 ],
1644 })
1645 }
1646 }
1647 }
1648
1649 fn report_changes_for_buffer(
1650 &mut self,
1651 buffer: &Entity<Buffer>,
1652 project: &Entity<Project>,
1653 is_predicted: bool,
1654 is_local: bool,
1655 cx: &mut Context<Self>,
1656 ) {
1657 let project_state = self.get_or_init_project(project, cx);
1658 let registered_buffer = Self::register_buffer_impl(project_state, buffer, project, cx);
1659
1660 let buf = buffer.read(cx);
1661 let new_file = buf.file().cloned();
1662 let new_snapshot = buf.text_snapshot();
1663 if new_snapshot.version == registered_buffer.snapshot.version {
1664 return;
1665 }
1666 let old_file = mem::replace(&mut registered_buffer.file, new_file.clone());
1667 let old_snapshot = mem::replace(&mut registered_buffer.snapshot, new_snapshot.clone());
1668 let mut edit_range: Option<Range<Anchor>> = None;
1669 let now = cx.background_executor().now();
1670
1671 for (_edit, anchor_range) in
1672 new_snapshot.anchored_edits_since::<usize>(&old_snapshot.version)
1673 {
1674 edit_range = Some(match edit_range {
1675 None => anchor_range,
1676 Some(acc) => acc.start..anchor_range.end,
1677 });
1678 }
1679
1680 let Some(edit_range) = edit_range else {
1681 return;
1682 };
1683
1684 for pending_capture in &mut project_state.pending_prediction_captures {
1685 if pending_capture.edited_buffer_id == buffer.entity_id()
1686 && edit_range.overlaps(&pending_capture.editable_anchor_range, &new_snapshot)
1687 {
1688 pending_capture.last_edit_at = now;
1689 if is_local
1690 && !is_predicted
1691 && let Some(sample_data) = pending_capture.sample_data.as_mut()
1692 && sample_data.next_edit_cursor_offset.is_none()
1693 {
1694 sample_data.next_edit_cursor_offset =
1695 Some(edit_range.start.to_offset(&new_snapshot));
1696 }
1697 }
1698 }
1699
1700 let include_in_history = is_local
1701 || collaborator_edit_overlaps_locality_region(
1702 project_state,
1703 project,
1704 buffer,
1705 &buf.snapshot(),
1706 &edit_range,
1707 cx,
1708 );
1709
1710 if !include_in_history {
1711 return;
1712 }
1713
1714 let is_recordable_history_edit =
1715 compute_diff_between_snapshots_in_range(&old_snapshot, &new_snapshot, &edit_range)
1716 .is_some();
1717
1718 if !is_recordable_history_edit {
1719 project_state.finalize_last_event(cx);
1720 return;
1721 }
1722
1723 if let Some(last_event) = project_state.last_event.as_mut() {
1724 let is_next_snapshot_of_same_buffer = old_snapshot.remote_id()
1725 == last_event.new_snapshot.remote_id()
1726 && old_snapshot.version == last_event.new_snapshot.version;
1727
1728 let prediction_source_changed = is_predicted != last_event.predicted;
1729
1730 let should_coalesce = is_next_snapshot_of_same_buffer
1731 && !prediction_source_changed
1732 && lines_between_ranges(
1733 &edit_range.to_point(&new_snapshot),
1734 &last_event.latest_edit_range.to_point(&new_snapshot),
1735 ) <= CHANGE_GROUPING_LINE_SPAN;
1736
1737 if should_coalesce {
1738 let pause_elapsed = last_event
1739 .last_edit_time
1740 .map(|t| now.duration_since(t) >= LAST_CHANGE_GROUPING_TIME)
1741 .unwrap_or(false);
1742 if pause_elapsed {
1743 last_event.snapshot_after_last_editing_pause =
1744 Some(last_event.new_snapshot.clone());
1745 last_event.total_edit_range_at_last_pause_boundary =
1746 Some(last_event.total_edit_range.clone());
1747 }
1748
1749 last_event.latest_edit_range = edit_range.clone();
1750 last_event.total_edit_range =
1751 merge_anchor_ranges(&last_event.total_edit_range, &edit_range, &new_snapshot);
1752 last_event.new_snapshot = new_snapshot;
1753 last_event.last_edit_time = Some(now);
1754 return;
1755 }
1756 }
1757
1758 project_state.finalize_last_event(cx);
1759
1760 merge_trailing_events_if_needed(
1761 &mut project_state.events,
1762 &old_snapshot,
1763 &new_snapshot,
1764 &edit_range,
1765 );
1766
1767 let file_context = new_file.as_ref().map(|file| {
1768 let project_path = ProjectPath::from_file(file.as_ref(), cx);
1769 let file_context = project_state.file_context_for_path(project_path.clone(), cx);
1770 Self::ensure_git_changed_file_sets_loading(&file_context, project, &project_path, cx);
1771 file_context
1772 });
1773
1774 let seq = project_state.next_last_event_seq;
1775 project_state.next_last_event_seq += 1;
1776 project_state.last_event = Some(LastEvent {
1777 seq,
1778 old_file,
1779 new_file,
1780 old_snapshot,
1781 new_snapshot,
1782 latest_edit_range: edit_range.clone(),
1783 total_edit_range: edit_range,
1784 total_edit_range_at_last_pause_boundary: None,
1785 predicted: is_predicted,
1786 snapshot_after_last_editing_pause: None,
1787 last_edit_time: Some(now),
1788 file_context,
1789 });
1790 }
1791
1792 fn prediction_at(
1793 &mut self,
1794 buffer: &Entity<Buffer>,
1795 position: Option<language::Anchor>,
1796 project: &Entity<Project>,
1797 cx: &App,
1798 ) -> Option<BufferEditPrediction<'_>> {
1799 let project_state = self.projects.get_mut(&project.entity_id())?;
1800 if let Some(position) = position {
1801 let snapshot = buffer.read(cx).snapshot();
1802 let cursor_position = position.to_offset(&snapshot);
1803 if let Some(file) = snapshot.file() {
1804 project_state.update_recent_file_cursor(file.path().as_std_path(), cursor_position);
1805 }
1806 if let Some(buffer) = project_state
1807 .registered_buffers
1808 .get_mut(&buffer.entity_id())
1809 {
1810 buffer.last_position = Some(position);
1811 }
1812 }
1813
1814 let CurrentEditPrediction {
1815 requested_by,
1816 prediction,
1817 ..
1818 } = project_state.current_prediction.as_ref()?;
1819
1820 if prediction.targets_buffer(buffer.read(cx)) {
1821 Some(BufferEditPrediction::Local { prediction })
1822 } else if requested_by == &buffer.entity_id() {
1823 Some(BufferEditPrediction::Jump { prediction })
1824 } else {
1825 None
1826 }
1827 }
1828
1829 fn accept_current_prediction(&mut self, project: &Entity<Project>, cx: &mut Context<Self>) {
1830 let Some(current_prediction) = self
1831 .projects
1832 .get_mut(&project.entity_id())
1833 .and_then(|project_state| project_state.current_prediction.take())
1834 else {
1835 return;
1836 };
1837
1838 self.report_changes_for_buffer(
1839 ¤t_prediction.prediction.buffer,
1840 project,
1841 true,
1842 true,
1843 cx,
1844 );
1845
1846 // can't hold &mut project_state ref across report_changes_for_buffer_call
1847 let Some(project_state) = self.projects.get_mut(&project.entity_id()) else {
1848 return;
1849 };
1850
1851 for pending_prediction in mem::take(&mut project_state.pending_predictions) {
1852 project_state.cancel_pending_prediction(pending_prediction, cx);
1853 }
1854
1855 match self.edit_prediction_model {
1856 EditPredictionModel::Mercury => {
1857 mercury::edit_prediction_accepted(
1858 current_prediction.prediction.id,
1859 self.client.http_client(),
1860 cx,
1861 );
1862 }
1863 EditPredictionModel::Zeta => {
1864 let is_cloud = !matches!(
1865 all_language_settings(None, cx).edit_predictions.provider,
1866 EditPredictionProvider::Ollama | EditPredictionProvider::OpenAiCompatibleApi
1867 );
1868 if is_cloud {
1869 zeta::edit_prediction_accepted(self, current_prediction, cx)
1870 }
1871 }
1872 EditPredictionModel::Fim { .. } => {}
1873 }
1874 }
1875
1876 async fn handle_rejected_predictions(
1877 rx: UnboundedReceiver<EditPredictionRejectionPayload>,
1878 client: Arc<Client>,
1879 llm_token: LlmApiToken,
1880 app_version: Version,
1881 background_executor: BackgroundExecutor,
1882 ) {
1883 let mut rx = std::pin::pin!(rx.peekable());
1884 let mut batched = Vec::new();
1885
1886 while let Some(EditPredictionRejectionPayload {
1887 rejection,
1888 organization_id,
1889 }) = rx.next().await
1890 {
1891 batched.push(rejection);
1892
1893 if batched.len() < MAX_EDIT_PREDICTION_REJECTIONS_PER_REQUEST / 2 {
1894 select_biased! {
1895 next = rx.as_mut().peek().fuse() => {
1896 if next.is_some() {
1897 continue;
1898 }
1899 }
1900 () = background_executor.timer(REJECT_REQUEST_DEBOUNCE).fuse() => {},
1901 }
1902 }
1903
1904 let url = client
1905 .http_client()
1906 .build_zed_llm_url("/predict_edits/reject", &[])
1907 .unwrap();
1908
1909 let flush_count = batched
1910 .len()
1911 // in case items have accumulated after failure
1912 .min(MAX_EDIT_PREDICTION_REJECTIONS_PER_REQUEST);
1913 let start = batched.len() - flush_count;
1914
1915 let body = RejectEditPredictionsBodyRef {
1916 rejections: &batched[start..],
1917 };
1918
1919 let result = Self::send_api_request::<()>(
1920 |builder| {
1921 let req = builder
1922 .uri(url.as_ref())
1923 .body(serde_json::to_string(&body)?.into());
1924 anyhow::Ok(req?)
1925 },
1926 client.clone(),
1927 llm_token.clone(),
1928 organization_id,
1929 app_version.clone(),
1930 )
1931 .await;
1932
1933 if result.log_err().is_some() {
1934 batched.drain(start..);
1935 }
1936 }
1937 }
1938
1939 async fn run_settled_predictions_worker(
1940 this: WeakEntity<Self>,
1941 mut rx: UnboundedReceiver<Instant>,
1942 client: Arc<Client>,
1943 llm_token: LlmApiToken,
1944 app_version: Version,
1945 cx: &mut AsyncApp,
1946 ) {
1947 let mut next_wake_time: Option<Instant> = None;
1948 loop {
1949 let now = cx.background_executor().now();
1950 if let Some(wake_time) = next_wake_time.take() {
1951 cx.background_executor()
1952 .timer(wake_time.duration_since(now))
1953 .await;
1954 } else {
1955 let Some(new_enqueue_time) = rx.next().await else {
1956 break;
1957 };
1958 next_wake_time = Some(new_enqueue_time + EDIT_PREDICTION_SETTLED_QUIESCENCE);
1959 while rx.next().now_or_never().flatten().is_some() {}
1960 continue;
1961 }
1962
1963 let Some(this) = this.upgrade() else {
1964 break;
1965 };
1966
1967 let now = cx.background_executor().now();
1968 let mut oldest_edited_at = None;
1969 let mut ready_predictions = Vec::new();
1970
1971 this.update(cx, |this, cx| {
1972 for project_state in this.projects.values_mut() {
1973 let ProjectState {
1974 last_event,
1975 registered_buffers,
1976 license_detection_watchers,
1977 pending_prediction_captures,
1978 ..
1979 } = project_state;
1980 let pending_last_event = last_event.as_ref().map(|last_event| {
1981 (
1982 last_event,
1983 last_event.finalize(license_detection_watchers, cx),
1984 )
1985 });
1986 let mut pending_index = 0;
1987 while pending_index < pending_prediction_captures.len() {
1988 let pending_capture = &pending_prediction_captures[pending_index];
1989 let age = now.saturating_duration_since(pending_capture.enqueued_at);
1990 if age >= EDIT_PREDICTION_SETTLED_TTL {
1991 pending_prediction_captures.remove(pending_index);
1992 continue;
1993 }
1994
1995 let quiet_for = now.saturating_duration_since(pending_capture.last_edit_at);
1996 if quiet_for >= EDIT_PREDICTION_SETTLED_QUIESCENCE {
1997 let Some(registered_buffer) =
1998 registered_buffers.get(&pending_capture.edited_buffer_id)
1999 else {
2000 pending_prediction_captures.remove(pending_index);
2001 continue;
2002 };
2003 let editable_offset_range = pending_capture
2004 .editable_anchor_range
2005 .to_offset(®istered_buffer.snapshot);
2006 if editable_offset_range.len()
2007 > EDIT_PREDICTION_SETTLED_MAX_EDITABLE_REGION_BYTES
2008 {
2009 // The prediction was obliterated by a huge edit;
2010 // kept-rate against it would be meaningless and the
2011 // region would blow the body size cap.
2012 pending_prediction_captures.remove(pending_index);
2013 continue;
2014 }
2015 let settled_editable_region = registered_buffer
2016 .snapshot
2017 .text_for_range(editable_offset_range)
2018 .collect::<String>();
2019 let mut pending_capture =
2020 pending_prediction_captures.remove(pending_index);
2021 if let Some((last_event, finalized_event)) = pending_last_event.as_ref()
2022 {
2023 pending_capture.try_record_future_event(
2024 last_event,
2025 finalized_event.as_ref(),
2026 license_detection_watchers,
2027 cx,
2028 );
2029 }
2030 ready_predictions.push((pending_capture, settled_editable_region));
2031 continue;
2032 }
2033
2034 if oldest_edited_at.is_none_or(|time| pending_capture.last_edit_at < time) {
2035 oldest_edited_at = Some(pending_capture.last_edit_at);
2036 }
2037 pending_index += 1;
2038 }
2039 }
2040 });
2041
2042 let mut ready_predictions_by_organization_id: HashMap<_, Vec<_>> = HashMap::default();
2043 for (pending_capture, settled_editable_region) in ready_predictions {
2044 #[cfg(test)]
2045 {
2046 let request_id = pending_capture.request_id.clone();
2047 let settled_editable_region = settled_editable_region.clone();
2048 this.update(cx, |this, _| {
2049 if let Some(callback) = &this.settled_event_callback {
2050 callback(request_id, settled_editable_region);
2051 }
2052 });
2053 }
2054 ready_predictions_by_organization_id
2055 .entry(pending_capture.organization_id.clone())
2056 .or_default()
2057 .push((pending_capture, settled_editable_region));
2058 }
2059
2060 cx.background_spawn({
2061 let client = client.clone();
2062 let llm_token = llm_token.clone();
2063 let app_version = app_version.clone();
2064 async move {
2065 send_settled_batches(
2066 client,
2067 llm_token,
2068 app_version,
2069 ready_predictions_by_organization_id,
2070 )
2071 .await;
2072 }
2073 })
2074 .detach();
2075
2076 next_wake_time = oldest_edited_at.map(|time| time + EDIT_PREDICTION_SETTLED_QUIESCENCE);
2077 }
2078 }
2079
2080 pub(crate) fn enqueue_settled_prediction(
2081 &mut self,
2082 request_id: EditPredictionId,
2083 project: &Entity<Project>,
2084 edited_buffer: &Entity<Buffer>,
2085 edited_buffer_snapshot: &BufferSnapshot,
2086 editable_offset_range: Range<usize>,
2087 edit_preview: &EditPreview,
2088 context_task: Option<Task<Result<CapturedPredictionContext>>>,
2089 prompt_history_boundary: Option<PromptHistoryBoundary>,
2090 model_version: Option<String>,
2091 e2e_latency: std::time::Duration,
2092 cx: &mut Context<Self>,
2093 ) {
2094 let this = &mut *self;
2095 let is_in_open_source_repo = edited_buffer_snapshot
2096 .file()
2097 .map_or(false, |file| this.is_file_open_source(project, file, cx));
2098 let can_collect_data = !cfg!(test)
2099 && is_in_open_source_repo
2100 && this.is_data_collection_enabled(cx)
2101 && matches!(this.edit_prediction_model, EditPredictionModel::Zeta);
2102
2103 let organization_id = this
2104 .user_store
2105 .read(cx)
2106 .current_organization()
2107 .map(|organization| organization.id.clone());
2108 let project_state = this.get_or_init_project(project, cx);
2109 if !project_state
2110 .registered_buffers
2111 .contains_key(&edited_buffer.entity_id())
2112 {
2113 return;
2114 }
2115
2116 let editable_region_before_prediction = edited_buffer_snapshot
2117 .text_for_range(editable_offset_range.clone())
2118 .collect::<String>();
2119 let editable_anchor_range_for_result =
2120 edited_buffer_snapshot.anchor_range_inside(editable_offset_range.clone());
2121 let predicted_editable_region = edit_preview
2122 .result_text_snapshot()
2123 .text_for_range(editable_anchor_range_for_result.clone())
2124 .collect();
2125 let ts_error_count_before_prediction = crate::metrics::count_tree_sitter_errors(
2126 edited_buffer_snapshot
2127 .syntax_layers_for_range(editable_anchor_range_for_result.clone(), true),
2128 );
2129 let ts_error_count_after_prediction = crate::metrics::count_tree_sitter_errors(
2130 edit_preview.result_syntax_snapshot().layers_for_range(
2131 editable_anchor_range_for_result,
2132 edit_preview.result_text_snapshot(),
2133 true,
2134 ),
2135 );
2136 let editable_anchor_range =
2137 edited_buffer_snapshot.anchor_range_inside(editable_offset_range.clone());
2138 let now = cx.background_executor().now();
2139 let sample_data = if can_collect_data
2140 && let Some(context_task) = context_task
2141 && let Some(file) = edited_buffer_snapshot.file()
2142 {
2143 Some(PendingPredictionCaptureSampleData {
2144 context_task,
2145 editable_path: file.path().as_std_path().into(),
2146 editable_offset_range,
2147 next_edit_cursor_offset: None,
2148 future_edit_history_events: Vec::new(),
2149 navigation_history: VecDeque::new(),
2150 edit_events_before_quiescence: 0,
2151 prompt_history_boundary,
2152 })
2153 } else {
2154 None
2155 };
2156 project_state
2157 .pending_prediction_captures
2158 .push(PendingPredictionCapture {
2159 request_id,
2160 edited_buffer_id: edited_buffer.entity_id(),
2161 editable_anchor_range,
2162 editable_region_before_prediction,
2163 predicted_editable_region,
2164 ts_error_count_before_prediction,
2165 ts_error_count_after_prediction,
2166 organization_id,
2167 can_collect_data,
2168 is_in_open_source_repo,
2169 sample_data,
2170 model_version,
2171 e2e_latency,
2172 enqueued_at: now,
2173 last_edit_at: now,
2174 });
2175 this.settled_predictions_tx.unbounded_send(now).ok();
2176 }
2177
2178 fn reject_current_prediction(
2179 &mut self,
2180 reason: EditPredictionRejectReason,
2181 project: &Entity<Project>,
2182 cx: &App,
2183 ) {
2184 if let Some(project_state) = self.projects.get_mut(&project.entity_id()) {
2185 project_state.pending_predictions.clear();
2186 if let Some(prediction) = project_state.current_prediction.take() {
2187 let model_version = prediction.prediction.model_version.clone();
2188 self.reject_prediction(
2189 prediction.prediction.id,
2190 reason,
2191 prediction.was_shown,
2192 model_version,
2193 Some(prediction.e2e_latency),
2194 cx,
2195 );
2196 }
2197 };
2198 }
2199
2200 fn did_show_current_prediction(
2201 &mut self,
2202 project: &Entity<Project>,
2203 display_type: edit_prediction_types::SuggestionDisplayType,
2204 _cx: &mut Context<Self>,
2205 ) {
2206 let Some(project_state) = self.projects.get_mut(&project.entity_id()) else {
2207 return;
2208 };
2209
2210 let Some(current_prediction) = project_state.current_prediction.as_mut() else {
2211 return;
2212 };
2213
2214 let is_jump = display_type == edit_prediction_types::SuggestionDisplayType::Jump;
2215 let previous_shown_with = current_prediction.shown_with;
2216
2217 if previous_shown_with.is_none() || !is_jump {
2218 current_prediction.shown_with = Some(display_type);
2219 }
2220
2221 let is_first_non_jump_show = !current_prediction.was_shown && !is_jump;
2222
2223 if is_first_non_jump_show {
2224 current_prediction.was_shown = true;
2225 }
2226
2227 if is_first_non_jump_show {
2228 self.rateable_predictions
2229 .push_front(current_prediction.prediction.clone());
2230 if self.rateable_predictions.len() > 50 {
2231 let completion = self.rateable_predictions.pop_back().unwrap();
2232 self.rated_predictions.remove(&completion.id);
2233 }
2234 }
2235 }
2236
2237 fn reject_prediction(
2238 &mut self,
2239 prediction_id: EditPredictionId,
2240 reason: EditPredictionRejectReason,
2241 was_shown: bool,
2242 model_version: Option<String>,
2243 e2e_latency: Option<std::time::Duration>,
2244 cx: &App,
2245 ) {
2246 match self.edit_prediction_model {
2247 EditPredictionModel::Zeta => {
2248 let is_cloud = !matches!(
2249 all_language_settings(None, cx).edit_predictions.provider,
2250 EditPredictionProvider::Ollama | EditPredictionProvider::OpenAiCompatibleApi
2251 );
2252
2253 if is_cloud {
2254 let organization_id = self
2255 .user_store
2256 .read(cx)
2257 .current_organization()
2258 .map(|organization| organization.id.clone());
2259
2260 self.reject_predictions_tx
2261 .unbounded_send(EditPredictionRejectionPayload {
2262 rejection: EditPredictionRejection {
2263 request_id: prediction_id.to_string(),
2264 reason,
2265 was_shown,
2266 model_version,
2267 e2e_latency_ms: e2e_latency.map(|latency| latency.as_millis()),
2268 },
2269 organization_id,
2270 })
2271 .log_err();
2272 }
2273 }
2274 EditPredictionModel::Mercury => {
2275 mercury::edit_prediction_rejected(
2276 prediction_id,
2277 was_shown,
2278 reason,
2279 self.client.http_client(),
2280 cx,
2281 );
2282 }
2283 EditPredictionModel::Fim { .. } => {}
2284 }
2285 }
2286
2287 fn is_refreshing(&self, project: &Entity<Project>) -> bool {
2288 self.projects
2289 .get(&project.entity_id())
2290 .is_some_and(|project_state| !project_state.pending_predictions.is_empty())
2291 }
2292
2293 pub fn refresh_prediction_from_buffer(
2294 &mut self,
2295 project: Entity<Project>,
2296 buffer: Entity<Buffer>,
2297 position: language::Anchor,
2298 trigger: EditPredictionRequestTrigger,
2299 cx: &mut Context<Self>,
2300 ) {
2301 if currently_following(&project, cx) {
2302 return;
2303 }
2304
2305 let trigger = predict_edits_request_trigger_from_editor_trigger(trigger);
2306
2307 self.queue_prediction_refresh(project.clone(), buffer.entity_id(), cx, move |this, cx| {
2308 let Some(request_task) = this
2309 .update(cx, |this, cx| {
2310 this.request_prediction_internal(
2311 project.clone(),
2312 buffer.clone(),
2313 position,
2314 trigger,
2315 cx,
2316 )
2317 })
2318 .log_err()
2319 else {
2320 return Task::ready(anyhow::Ok(None));
2321 };
2322
2323 cx.spawn(async move |_cx| {
2324 request_task.await.map(|prediction_result| {
2325 prediction_result
2326 .map(|prediction_result| (prediction_result, buffer.entity_id()))
2327 })
2328 })
2329 })
2330 }
2331
2332 pub const THROTTLE_TIMEOUT: Duration = Duration::from_millis(300);
2333}
2334
2335async fn send_settled_batches(
2336 client: Arc<Client>,
2337 llm_token: LlmApiToken,
2338 app_version: Version,
2339 ready_predictions_by_organization_id: hash_map::HashMap<
2340 Option<OrganizationId>,
2341 Vec<(PendingPredictionCapture, String)>,
2342 collections::FxBuildHasher,
2343 >,
2344) {
2345 let Some(url) = client
2346 .http_client()
2347 .build_zed_llm_url("/predict_edits/settled", &[])
2348 .context("failed to build edit predictions settled url")
2349 .log_err()
2350 else {
2351 return;
2352 };
2353
2354 for (organization_id, ready_predictions) in ready_predictions_by_organization_id {
2355 let mut ready_predictions = ready_predictions.into_iter();
2356 loop {
2357 let done_batch = ready_predictions
2358 .by_ref()
2359 .take(MAX_EDIT_PREDICTION_SETTLED_PER_REQUEST);
2360 let mut batch = Vec::with_capacity(MAX_EDIT_PREDICTION_SETTLED_PER_REQUEST);
2361 for (pending_capture, settled_editable_region) in done_batch {
2362 let PendingPredictionCapture {
2363 request_id,
2364 editable_region_before_prediction,
2365 predicted_editable_region,
2366 ts_error_count_before_prediction,
2367 ts_error_count_after_prediction,
2368 can_collect_data,
2369 is_in_open_source_repo,
2370 sample_data,
2371 model_version,
2372 e2e_latency,
2373 ..
2374 } = pending_capture;
2375 let kept_rate_result = compute_kept_rate(
2376 &editable_region_before_prediction,
2377 &predicted_editable_region,
2378 &settled_editable_region,
2379 );
2380
2381 let sample_data = if can_collect_data
2382 && let Some(sample_data) = sample_data
2383 && let Ok(context) = sample_data.context_task.await
2384 {
2385 Some(SettledEditPredictionSampleData {
2386 repository_url: context.repository_url,
2387 revision: context.revision,
2388 uncommitted_diff: context.uncommitted_diff,
2389 editable_path: sample_data.editable_path,
2390 editable_offset_range: sample_data.editable_offset_range,
2391 buffer_diagnostics: context.buffer_diagnostics,
2392 editable_context: context.editable_context,
2393 future_edit_history_events: sample_data.future_edit_history_events,
2394 navigation_history: sample_data
2395 .navigation_history
2396 .into_iter()
2397 .map(|file| EditPredictionRecentFile {
2398 path: file.path,
2399 cursor_position: file.cursor_position,
2400 })
2401 .collect(),
2402 edit_events_before_quiescence: sample_data.edit_events_before_quiescence,
2403 next_edit_cursor_offset: sample_data.next_edit_cursor_offset,
2404 })
2405 } else {
2406 None
2407 };
2408
2409 batch.push(SettledEditPrediction {
2410 request_id: request_id.0.to_string(),
2411 settled_editable_region: can_collect_data.then_some(settled_editable_region),
2412 ts_error_count_before_prediction,
2413 ts_error_count_after_prediction,
2414 can_collect_data,
2415 is_in_open_source_repo,
2416 sample_data,
2417 kept_chars: EditPredictionSettledKeptChars {
2418 candidate_new: kept_rate_result.candidate_new_chars,
2419 reference_new: kept_rate_result.reference_new_chars,
2420 candidate_deleted: kept_rate_result.candidate_deleted_chars,
2421 reference_deleted: kept_rate_result.reference_deleted_chars,
2422 kept: kept_rate_result.kept_chars,
2423 correctly_deleted: kept_rate_result.correctly_deleted_chars,
2424 discarded: kept_rate_result.discarded_chars,
2425 context: kept_rate_result.context_chars,
2426 kept_rate: kept_rate_result.kept_rate,
2427 recall_rate: kept_rate_result.recall_rate,
2428 },
2429 example: None,
2430 model_version,
2431 e2e_latency_ms: e2e_latency.as_millis().min(u128::from(u64::MAX)) as u64,
2432 });
2433 }
2434
2435 if batch.is_empty() {
2436 break;
2437 }
2438
2439 let result = async {
2440 let body = SubmitEditPredictionSettledBatchBody { predictions: batch };
2441 let compressed = zstd::encode_all(&serde_json::to_vec(&body)?[..], 3)?;
2442 EditPredictionStore::send_api_request::<SubmitEditPredictionSettledResponse>(
2443 |builder| {
2444 Ok(builder
2445 .uri(url.as_ref())
2446 .header("Content-Encoding", "zstd")
2447 .body(compressed.clone().into())?)
2448 },
2449 client.clone(),
2450 llm_token.clone(),
2451 organization_id.clone(),
2452 app_version.clone(),
2453 )
2454 .await?;
2455 anyhow::Ok(())
2456 }
2457 .await;
2458
2459 if let Err(error) = result {
2460 log::error!("failed to submit edit predictions settled: {error:?}");
2461 }
2462 }
2463 }
2464}
2465
2466fn currently_following(project: &Entity<Project>, cx: &App) -> bool {
2467 let Some(app_state) = AppState::try_global(cx) else {
2468 return false;
2469 };
2470
2471 app_state
2472 .workspace_store
2473 .read(cx)
2474 .workspaces()
2475 .filter_map(|workspace| workspace.upgrade())
2476 .any(|workspace| {
2477 workspace.read(cx).project().entity_id() == project.entity_id()
2478 && workspace
2479 .read(cx)
2480 .leader_for_pane(workspace.read(cx).active_pane())
2481 .is_some()
2482 })
2483}
2484
2485fn is_ep_store_provider(provider: EditPredictionProvider) -> bool {
2486 match provider {
2487 EditPredictionProvider::Zed
2488 | EditPredictionProvider::Mercury
2489 | EditPredictionProvider::Ollama
2490 | EditPredictionProvider::OpenAiCompatibleApi => true,
2491 EditPredictionProvider::None
2492 | EditPredictionProvider::Copilot
2493 | EditPredictionProvider::Codestral => false,
2494 }
2495}
2496
2497impl EditPredictionStore {
2498 fn queue_prediction_refresh(
2499 &mut self,
2500 project: Entity<Project>,
2501 throttle_entity: EntityId,
2502 cx: &mut Context<Self>,
2503 do_refresh: impl FnOnce(
2504 WeakEntity<Self>,
2505 &mut AsyncApp,
2506 ) -> Task<Result<Option<(EditPredictionResult, EntityId)>>>
2507 + 'static,
2508 ) {
2509 let (needs_acceptance_tracking, max_pending_predictions) =
2510 match all_language_settings(None, cx).edit_predictions.provider {
2511 EditPredictionProvider::Zed | EditPredictionProvider::Mercury => (true, 2),
2512 EditPredictionProvider::Ollama => (false, 1),
2513 EditPredictionProvider::OpenAiCompatibleApi => (false, 2),
2514 EditPredictionProvider::None
2515 | EditPredictionProvider::Copilot
2516 | EditPredictionProvider::Codestral => {
2517 log::error!("queue_prediction_refresh called with non-store provider");
2518 return;
2519 }
2520 };
2521
2522 let drop_on_cancel = !needs_acceptance_tracking;
2523 let throttle_timeout = Self::THROTTLE_TIMEOUT;
2524 let project_state = self.get_or_init_project(&project, cx);
2525 let pending_prediction_id = project_state.next_pending_prediction_id;
2526 project_state.next_pending_prediction_id += 1;
2527 let throttle_at_enqueue = project_state.last_edit_prediction_refresh;
2528
2529 let task = cx.spawn(async move |this, cx| {
2530 let throttle_wait = this
2531 .update(cx, |this, cx| {
2532 let project_state = this.get_or_init_project(&project, cx);
2533 let throttle = project_state.last_edit_prediction_refresh;
2534
2535 let now = cx.background_executor().now();
2536 throttle.and_then(|(last_entity, last_timestamp)| {
2537 if throttle_entity != last_entity {
2538 return None;
2539 }
2540 (last_timestamp + throttle_timeout).checked_duration_since(now)
2541 })
2542 })
2543 .ok()
2544 .flatten();
2545
2546 if let Some(timeout) = throttle_wait {
2547 cx.background_executor().timer(timeout).await;
2548 }
2549
2550 // If this task was cancelled before the throttle timeout expired,
2551 // do not perform a request. Also skip if another task already
2552 // proceeded since we were enqueued (duplicate).
2553 let mut is_cancelled = true;
2554 this.update(cx, |this, cx| {
2555 let project_state = this.get_or_init_project(&project, cx);
2556 let was_cancelled = project_state
2557 .cancelled_predictions
2558 .remove(&pending_prediction_id);
2559 if was_cancelled {
2560 return;
2561 }
2562
2563 // Another request has been already sent since this was enqueued
2564 if project_state.last_edit_prediction_refresh != throttle_at_enqueue {
2565 return;
2566 }
2567
2568 let new_refresh = (throttle_entity, cx.background_executor().now());
2569 project_state.last_edit_prediction_refresh = Some(new_refresh);
2570 is_cancelled = false;
2571 })
2572 .ok();
2573 if is_cancelled {
2574 return None;
2575 }
2576
2577 let new_prediction_result = do_refresh(this.clone(), cx).await.log_err().flatten();
2578 let new_prediction_metadata = new_prediction_result.as_ref().map(|(result, _)| {
2579 (
2580 result.prediction.id.clone(),
2581 result.prediction.model_version.clone(),
2582 )
2583 });
2584
2585 // When a prediction completes, remove it from the pending list, and cancel
2586 // any pending predictions that were enqueued before it.
2587 this.update(cx, |this, cx| {
2588 let project_state = this.get_or_init_project(&project, cx);
2589
2590 let is_cancelled = project_state
2591 .cancelled_predictions
2592 .remove(&pending_prediction_id);
2593
2594 let new_current_prediction = if !is_cancelled
2595 && let Some((prediction_result, requested_by)) = new_prediction_result
2596 {
2597 let EditPredictionResult {
2598 prediction,
2599 reject_reason,
2600 e2e_latency,
2601 } = prediction_result;
2602
2603 if let Some(reject_reason) = reject_reason {
2604 let should_allow_rating_prediction = matches!(
2605 reject_reason,
2606 EditPredictionRejectReason::Empty
2607 | EditPredictionRejectReason::InterpolatedEmpty
2608 );
2609 let prediction_id = prediction.id.clone();
2610 let model_version = prediction.model_version.clone();
2611
2612 this.reject_prediction(
2613 prediction_id,
2614 reject_reason,
2615 false,
2616 model_version,
2617 Some(e2e_latency),
2618 cx,
2619 );
2620
2621 if should_allow_rating_prediction {
2622 this.rateable_predictions.push_front(prediction);
2623 if this.rateable_predictions.len() > 50
2624 && let Some(completion) = this.rateable_predictions.pop_back()
2625 {
2626 this.rated_predictions.remove(&completion.id);
2627 }
2628 }
2629
2630 None
2631 } else {
2632 let new_prediction = CurrentEditPrediction {
2633 requested_by,
2634 prediction,
2635 was_shown: false,
2636 shown_with: None,
2637 e2e_latency,
2638 };
2639
2640 if let Some(current_prediction) = project_state.current_prediction.as_ref()
2641 {
2642 if new_prediction.should_replace_prediction(¤t_prediction, cx) {
2643 this.reject_current_prediction(
2644 EditPredictionRejectReason::Replaced,
2645 &project,
2646 cx,
2647 );
2648
2649 Some(new_prediction)
2650 } else {
2651 this.reject_prediction(
2652 new_prediction.prediction.id,
2653 EditPredictionRejectReason::CurrentPreferred,
2654 false,
2655 new_prediction.prediction.model_version,
2656 Some(new_prediction.e2e_latency),
2657 cx,
2658 );
2659 None
2660 }
2661 } else {
2662 Some(new_prediction)
2663 }
2664 }
2665 } else {
2666 None
2667 };
2668
2669 let project_state = this.get_or_init_project(&project, cx);
2670
2671 if let Some(new_prediction) = new_current_prediction {
2672 project_state.current_prediction = Some(new_prediction);
2673 }
2674
2675 let mut pending_predictions = mem::take(&mut project_state.pending_predictions);
2676 for (ix, pending_prediction) in pending_predictions.iter().enumerate() {
2677 if pending_prediction.id == pending_prediction_id {
2678 pending_predictions.remove(ix);
2679 for pending_prediction in pending_predictions.drain(0..ix) {
2680 project_state.cancel_pending_prediction(pending_prediction, cx)
2681 }
2682 break;
2683 }
2684 }
2685 this.get_or_init_project(&project, cx).pending_predictions = pending_predictions;
2686 cx.notify();
2687 })
2688 .ok();
2689
2690 new_prediction_metadata
2691 });
2692
2693 if project_state.pending_predictions.len() < max_pending_predictions {
2694 project_state
2695 .pending_predictions
2696 .push(PendingPrediction {
2697 id: pending_prediction_id,
2698 task,
2699 drop_on_cancel,
2700 })
2701 .unwrap();
2702 } else {
2703 let pending_prediction = project_state.pending_predictions.pop().unwrap();
2704 project_state
2705 .pending_predictions
2706 .push(PendingPrediction {
2707 id: pending_prediction_id,
2708 task,
2709 drop_on_cancel,
2710 })
2711 .unwrap();
2712 project_state.cancel_pending_prediction(pending_prediction, cx);
2713 }
2714 }
2715
2716 pub fn request_prediction(
2717 &mut self,
2718 project: &Entity<Project>,
2719 active_buffer: &Entity<Buffer>,
2720 position: language::Anchor,
2721 trigger: PredictEditsRequestTrigger,
2722 cx: &mut Context<Self>,
2723 ) -> Task<Result<Option<EditPredictionResult>>> {
2724 self.request_prediction_internal(
2725 project.clone(),
2726 active_buffer.clone(),
2727 position,
2728 trigger,
2729 cx,
2730 )
2731 }
2732
2733 fn request_prediction_internal(
2734 &mut self,
2735 project: Entity<Project>,
2736 active_buffer: Entity<Buffer>,
2737 position: language::Anchor,
2738 trigger: PredictEditsRequestTrigger,
2739 cx: &mut Context<Self>,
2740 ) -> Task<Result<Option<EditPredictionResult>>> {
2741 let is_cloud_zeta = matches!(self.edit_prediction_model, EditPredictionModel::Zeta)
2742 && !matches!(
2743 all_language_settings(None, cx).edit_predictions.provider,
2744 EditPredictionProvider::Ollama | EditPredictionProvider::OpenAiCompatibleApi
2745 );
2746 if is_cloud_zeta && !self.client.cloud_client().has_credentials() {
2747 return Task::ready(Ok(None));
2748 }
2749
2750 if is_cloud_zeta && self.request_backoff_active(cx) {
2751 log::debug!(
2752 "Skipping Zeta edit prediction request while backing off after Cloud timeout"
2753 );
2754 return Task::ready(Ok(None));
2755 }
2756
2757 self.get_or_init_project(&project, cx);
2758 let (stored_events, prompt_history_boundary, debug_tx) = {
2759 let project_state = self.projects.get(&project.entity_id()).unwrap();
2760 (
2761 project_state.events(cx),
2762 Some(PromptHistoryBoundary {
2763 first_event_seq: project_state
2764 .last_event
2765 .as_ref()
2766 .map_or(project_state.next_last_event_seq, |last_event| {
2767 last_event.seq
2768 }),
2769 snapshot: project_state
2770 .last_event
2771 .as_ref()
2772 .map(|last_event| last_event.new_snapshot.clone()),
2773 }),
2774 project_state.debug_tx.clone(),
2775 )
2776 };
2777 let events: Vec<Arc<zeta_prompt::Event>> =
2778 stored_events.iter().map(|e| e.event.clone()).collect();
2779
2780 let snapshot = active_buffer.read(cx).snapshot();
2781 let cursor_point = position.to_point(&snapshot);
2782 let diagnostic_search_start = cursor_point.row.saturating_sub(DIAGNOSTIC_LINES_RANGE);
2783 let diagnostic_search_end = cursor_point.row + DIAGNOSTIC_LINES_RANGE;
2784 let diagnostic_search_range =
2785 Point::new(diagnostic_search_start, 0)..Point::new(diagnostic_search_end, 0);
2786
2787 let related_files = self.context_for_project(&project, cx);
2788 let allow_jump = is_cloud_zeta && cx.has_flag::<EditPredictionJumpsFeatureFlag>();
2789 let mode = match all_language_settings(snapshot.file(), cx).edit_predictions_mode() {
2790 EditPredictionsMode::Eager => PredictEditsMode::Eager,
2791 EditPredictionsMode::Subtle => PredictEditsMode::Subtle,
2792 };
2793
2794 let buffer_id = active_buffer.read(cx).remote_id();
2795 let (repository_url, revision) = project
2796 .read(cx)
2797 .git_store()
2798 .read(cx)
2799 .repository_and_path_for_buffer_id(buffer_id, cx)
2800 .map(|(repository, _)| {
2801 let snapshot = repository.read(cx).snapshot();
2802 (
2803 snapshot
2804 .remote_origin_url
2805 .clone()
2806 .or_else(|| snapshot.remote_upstream_url.clone()),
2807 snapshot
2808 .head_commit
2809 .as_ref()
2810 .map(|commit| commit.sha.to_string()),
2811 )
2812 })
2813 .unwrap_or_default();
2814
2815 let is_staff_zed_repo = cx.is_staff()
2816 && repository_url
2817 .as_ref()
2818 .is_some_and(|url| is_zed_industries_repo(url));
2819 let is_open_source = is_staff_zed_repo
2820 || (snapshot
2821 .file()
2822 .map_or(false, |file| self.is_file_open_source(&project, file, cx))
2823 && events.iter().all(|event| event.in_open_source_repo())
2824 && related_files.iter().all(|file| file.in_open_source_repo));
2825
2826 let can_collect_data = !cfg!(test)
2827 && is_open_source
2828 && self.is_data_collection_enabled(cx)
2829 && matches!(self.edit_prediction_model, EditPredictionModel::Zeta);
2830 let editable_context = allow_jump.then(|| {
2831 self.collect_editable_context(
2832 project.clone(),
2833 active_buffer.clone(),
2834 position,
2835 Vec::new(),
2836 vec![ContextSource::CurrentFile, ContextSource::EditHistory],
2837 cx,
2838 )
2839 });
2840 let inputs = EditPredictionModelInput {
2841 project: project.clone(),
2842 buffer: active_buffer,
2843 snapshot,
2844 position,
2845 events,
2846 related_files,
2847 editable_context,
2848 mode,
2849 trigger,
2850 diagnostic_search_range,
2851 debug_tx,
2852 can_collect_data,
2853 is_open_source,
2854 allow_jump,
2855 };
2856
2857 let task = match self.edit_prediction_model {
2858 EditPredictionModel::Zeta => {
2859 let context_task = can_collect_data
2860 .then(|| {
2861 let editable_context_task = self.collect_editable_context(
2862 inputs.project.clone(),
2863 inputs.buffer.clone(),
2864 inputs.position,
2865 Vec::new(),
2866 vec![ContextSource::CurrentFile, ContextSource::EditHistory],
2867 cx,
2868 );
2869 capture_prediction_context(
2870 inputs.project.clone(),
2871 inputs.buffer.clone(),
2872 inputs.position,
2873 stored_events,
2874 repository_url.clone(),
2875 revision,
2876 editable_context_task,
2877 cx,
2878 )
2879 })
2880 .flatten();
2881 zeta::request_prediction_with_zeta(
2882 self,
2883 inputs,
2884 context_task,
2885 prompt_history_boundary,
2886 repository_url,
2887 cx,
2888 )
2889 }
2890 EditPredictionModel::Fim { format } => fim::request_prediction(inputs, format, cx),
2891 EditPredictionModel::Mercury => {
2892 self.mercury
2893 .request_prediction(inputs, self.credentials_provider.clone(), cx)
2894 }
2895 };
2896
2897 task
2898 }
2899
2900 async fn send_raw_llm_request(
2901 request: RawCompletionRequest,
2902 client: Arc<Client>,
2903 custom_url: Option<Arc<Url>>,
2904 llm_token: LlmApiToken,
2905 organization_id: Option<OrganizationId>,
2906 app_version: Version,
2907 ) -> Result<(RawCompletionResponse, Option<EditPredictionUsage>)> {
2908 let url = if let Some(custom_url) = custom_url {
2909 custom_url.as_ref().clone()
2910 } else {
2911 client
2912 .http_client()
2913 .build_zed_llm_url("/predict_edits/raw", &[])?
2914 };
2915
2916 Self::send_api_request(
2917 |builder| {
2918 let req = builder
2919 .uri(url.as_ref())
2920 .body(serde_json::to_string(&request)?.into());
2921 Ok(req?)
2922 },
2923 client,
2924 llm_token,
2925 organization_id,
2926 app_version,
2927 )
2928 .await
2929 }
2930
2931 pub(crate) async fn send_v3_request(
2932 input: Zeta2PromptInput,
2933 preferred_experiment: Option<String>,
2934 client: Arc<Client>,
2935 llm_token: LlmApiToken,
2936 organization_id: Option<OrganizationId>,
2937 app_version: Version,
2938 trigger: PredictEditsRequestTrigger,
2939 mode: PredictEditsMode,
2940 ) -> Result<(PredictEditsV3Response, Option<EditPredictionUsage>)> {
2941 let request = PredictEditsV3Request { input };
2942 Self::send_predict_edits_request(
2943 "/predict_edits/v3",
2944 request,
2945 preferred_experiment,
2946 client,
2947 llm_token,
2948 organization_id,
2949 app_version,
2950 trigger,
2951 mode,
2952 )
2953 .await
2954 }
2955
2956 pub(crate) async fn send_v4_request(
2957 input: Zeta3PromptInput,
2958 preferred_experiment: Option<String>,
2959 client: Arc<Client>,
2960 llm_token: LlmApiToken,
2961 organization_id: Option<OrganizationId>,
2962 app_version: Version,
2963 trigger: PredictEditsRequestTrigger,
2964 mode: PredictEditsMode,
2965 ) -> Result<(PredictEditsV4Response, Option<EditPredictionUsage>)> {
2966 let request = PredictEditsV4Request { input };
2967 Self::send_predict_edits_request(
2968 "/predict_edits/v4",
2969 request,
2970 preferred_experiment,
2971 client,
2972 llm_token,
2973 organization_id,
2974 app_version,
2975 trigger,
2976 mode,
2977 )
2978 .await
2979 }
2980
2981 async fn send_predict_edits_request<Req, Res>(
2982 path: &str,
2983 request: Req,
2984 preferred_experiment: Option<String>,
2985 client: Arc<Client>,
2986 llm_token: LlmApiToken,
2987 organization_id: Option<OrganizationId>,
2988 app_version: Version,
2989 trigger: PredictEditsRequestTrigger,
2990 mode: PredictEditsMode,
2991 ) -> Result<(Res, Option<EditPredictionUsage>)>
2992 where
2993 Req: serde::Serialize,
2994 Res: serde::de::DeserializeOwned,
2995 {
2996 let url = client.http_client().build_zed_llm_url(path, &[])?;
2997 let request_id = uuid::Uuid::new_v4().to_string();
2998
2999 let json_bytes = serde_json::to_vec(&request)?;
3000 let compressed = zstd::encode_all(&json_bytes[..], 3)?;
3001
3002 Self::send_api_request(
3003 |builder| {
3004 let builder = builder
3005 .uri(url.as_ref())
3006 .header("Content-Encoding", "zstd")
3007 .header(PREDICT_EDITS_MODE_HEADER_NAME, mode.as_ref())
3008 .header(PREDICT_EDITS_REQUEST_ID_HEADER_NAME, request_id.as_str())
3009 .header(PREDICT_EDITS_TRIGGER_HEADER_NAME, trigger.as_ref());
3010 let builder = if let Some(preferred_experiment) = preferred_experiment.as_deref() {
3011 builder.header(PREFERRED_EXPERIMENT_HEADER_NAME, preferred_experiment)
3012 } else {
3013 builder
3014 };
3015 let req = builder.body(compressed.clone().into());
3016 Ok(req?)
3017 },
3018 client,
3019 llm_token,
3020 organization_id,
3021 app_version,
3022 )
3023 .await
3024 }
3025
3026 async fn send_api_request<Res>(
3027 build: impl Fn(http_client::http::request::Builder) -> Result<http_client::Request<AsyncBody>>,
3028 client: Arc<Client>,
3029 llm_token: LlmApiToken,
3030 organization_id: Option<OrganizationId>,
3031 app_version: Version,
3032 ) -> Result<(Res, Option<EditPredictionUsage>)>
3033 where
3034 Res: DeserializeOwned,
3035 {
3036 let organization_id =
3037 organization_id.ok_or_else(|| anyhow!("No organization selected."))?;
3038
3039 let response = client
3040 .authenticated_llm_request(&llm_token, organization_id, |token| {
3041 build(
3042 http_client::Request::builder()
3043 .method(Method::POST)
3044 .header("Content-Type", "application/json")
3045 .header(ZED_VERSION_HEADER_NAME, app_version.to_string())
3046 .header("Authorization", format!("Bearer {token}")),
3047 )
3048 })
3049 .await?;
3050
3051 Self::process_api_response(response, &app_version).await
3052 }
3053
3054 async fn process_api_response<Res>(
3055 mut response: http_client::Response<AsyncBody>,
3056 app_version: &Version,
3057 ) -> Result<(Res, Option<EditPredictionUsage>)>
3058 where
3059 Res: DeserializeOwned,
3060 {
3061 if let Some(minimum_required_version) = response
3062 .headers()
3063 .get(MINIMUM_REQUIRED_VERSION_HEADER_NAME)
3064 .and_then(|version| Version::from_str(version.to_str().ok()?).ok())
3065 {
3066 anyhow::ensure!(
3067 *app_version >= minimum_required_version,
3068 ZedUpdateRequiredError {
3069 minimum_version: minimum_required_version
3070 }
3071 );
3072 }
3073
3074 if response.status().is_success() {
3075 let usage = EditPredictionUsage::from_headers(response.headers()).ok();
3076 let mut body = Vec::new();
3077 response.body_mut().read_to_end(&mut body).await?;
3078 Ok((serde_json::from_slice(&body)?, usage))
3079 } else {
3080 let status = response.status();
3081 let mut body = String::new();
3082 response.body_mut().read_to_string(&mut body).await?;
3083 if status == http_client::http::StatusCode::REQUEST_TIMEOUT {
3084 return Err(anyhow::Error::new(CloudRequestTimeoutError));
3085 }
3086 anyhow::bail!("Request failed with status: {status:?}\nBody: {body}");
3087 }
3088 }
3089
3090 pub fn refresh_context(
3091 &mut self,
3092 project: &Entity<Project>,
3093 buffer: &Entity<language::Buffer>,
3094 cursor_position: language::Anchor,
3095 cx: &mut Context<Self>,
3096 ) {
3097 self.get_or_init_project(project, cx)
3098 .context
3099 .update(cx, |store, cx| {
3100 store.refresh(buffer.clone(), cursor_position, cx);
3101 });
3102 }
3103
3104 pub fn collect_editable_context(
3105 &mut self,
3106 project: Entity<Project>,
3107 buffer: Entity<language::Buffer>,
3108 cursor_position: language::Anchor,
3109 oracle_targets: Vec<edit_prediction_context::OracleTarget>,
3110 context_sources: Vec<ContextSource>,
3111 cx: &mut Context<Self>,
3112 ) -> Task<anyhow::Result<Vec<RelatedFile>>> {
3113 use edit_prediction_context::{EditHistoryContextEntry, collect_editable_context};
3114
3115 let buffers_by_id = project.read(cx).opened_buffers(cx).into_iter().fold(
3116 HashMap::default(),
3117 |mut buffers_by_id, buffer| {
3118 buffers_by_id.insert(buffer.read(cx).remote_id(), buffer.clone());
3119 buffers_by_id
3120 },
3121 );
3122 let edit_history = self
3123 .edit_history_for_project(&project, cx)
3124 .into_iter()
3125 .filter_map(|event| {
3126 let buffer = buffers_by_id.get(&event.old_snapshot.remote_id())?.clone();
3127 Some(EditHistoryContextEntry {
3128 buffer,
3129 edited_range: event.total_edit_range,
3130 })
3131 })
3132 .collect();
3133
3134 cx.spawn(async move |_, cx| {
3135 collect_editable_context(
3136 project,
3137 buffer,
3138 cursor_position,
3139 edit_history,
3140 oracle_targets,
3141 context_sources,
3142 cx,
3143 )
3144 .await
3145 })
3146 }
3147
3148 #[cfg(feature = "cli-support")]
3149 pub fn set_context_for_buffer(
3150 &mut self,
3151 project: &Entity<Project>,
3152 related_files: Vec<RelatedFile>,
3153 cx: &mut Context<Self>,
3154 ) {
3155 self.get_or_init_project(project, cx)
3156 .context
3157 .update(cx, |store, cx| {
3158 store.set_related_files(related_files, cx);
3159 });
3160 }
3161
3162 #[cfg(feature = "cli-support")]
3163 pub fn set_recent_paths_for_project(
3164 &mut self,
3165 project: &Entity<Project>,
3166 paths: impl IntoIterator<Item = project::ProjectPath>,
3167 cx: &mut Context<Self>,
3168 ) {
3169 let project_state = self.get_or_init_project(project, cx);
3170 project_state.recently_viewed_files = paths
3171 .into_iter()
3172 .map(|path| RecentFile {
3173 path: path.path.as_std_path().into(),
3174 cursor_position: None,
3175 })
3176 .collect();
3177 }
3178
3179 pub fn recently_opened_files_for_project(&self, project: &Entity<Project>) -> Vec<RecentFile> {
3180 self.projects
3181 .get(&project.entity_id())
3182 .map(|project_state| {
3183 project_state
3184 .recently_opened_files
3185 .iter()
3186 .cloned()
3187 .collect()
3188 })
3189 .unwrap_or_default()
3190 }
3191
3192 pub fn recently_viewed_files_for_project(&self, project: &Entity<Project>) -> Vec<RecentFile> {
3193 self.projects
3194 .get(&project.entity_id())
3195 .map(|project_state| {
3196 project_state
3197 .recently_viewed_files
3198 .iter()
3199 .cloned()
3200 .collect()
3201 })
3202 .unwrap_or_default()
3203 }
3204
3205 fn is_file_open_source(
3206 &self,
3207 project: &Entity<Project>,
3208 file: &Arc<dyn File>,
3209 cx: &App,
3210 ) -> bool {
3211 if !file.is_local() || file.is_private() {
3212 return false;
3213 }
3214 let Some(project_state) = self.projects.get(&project.entity_id()) else {
3215 return false;
3216 };
3217 project_state
3218 .license_detection_watchers
3219 .get(&file.worktree_id(cx))
3220 .as_ref()
3221 .is_some_and(|watcher| watcher.is_project_open_source())
3222 }
3223
3224 pub(crate) fn is_data_collection_enabled(&self, cx: &App) -> bool {
3225 if !self.is_data_collection_allowed_by_organization(cx) {
3226 return false;
3227 }
3228
3229 if cx.is_staff() {
3230 return true;
3231 }
3232
3233 match all_language_settings(None, cx)
3234 .edit_predictions
3235 .allow_data_collection
3236 {
3237 EditPredictionDataCollectionChoice::Yes => true,
3238 EditPredictionDataCollectionChoice::No => false,
3239 // Fall back to the legacy KV entry captured when the store was
3240 // created, preserving existing users' choices without per-request
3241 // database reads.
3242 EditPredictionDataCollectionChoice::Default => self.legacy_data_collection_enabled,
3243 }
3244 }
3245
3246 fn load_legacy_data_collection_enabled(cx: &App) -> bool {
3247 KeyValueStore::global(cx)
3248 .read_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE)
3249 .log_err()
3250 .flatten()
3251 .as_deref()
3252 == Some("true")
3253 }
3254
3255 pub(crate) fn is_data_collection_allowed_by_organization(&self, cx: &App) -> bool {
3256 self.user_store
3257 .read(cx)
3258 .current_organization_configuration()
3259 .is_none_or(|organization_configuration| {
3260 organization_configuration
3261 .edit_prediction
3262 .is_feedback_enabled
3263 })
3264 }
3265
3266 pub fn rateable_predictions(&self) -> impl DoubleEndedIterator<Item = &EditPrediction> {
3267 self.rateable_predictions.iter()
3268 }
3269
3270 pub fn rateable_predictions_count(&self) -> usize {
3271 self.rateable_predictions.len()
3272 }
3273
3274 pub fn is_prediction_rated(&self, id: &EditPredictionId) -> bool {
3275 self.rated_predictions.contains(id)
3276 }
3277
3278 pub fn rate_prediction(
3279 &mut self,
3280 prediction: &EditPrediction,
3281 rating: EditPredictionRating,
3282 feedback: String,
3283 expected_output: Option<String>,
3284 cx: &mut Context<Self>,
3285 ) {
3286 let organization = self.user_store.read(cx).current_organization();
3287
3288 self.rated_predictions.insert(prediction.id.clone());
3289
3290 cx.background_spawn({
3291 let client = self.client.clone();
3292 let prediction_id = prediction.id.to_string();
3293 let inputs = serde_json::to_value(&prediction.inputs);
3294 let output = prediction
3295 .edit_preview
3296 .as_unified_diff(prediction.snapshot.file(), &prediction.edits);
3297 async move {
3298 client
3299 .cloud_client()
3300 .submit_edit_prediction_feedback(SubmitEditPredictionFeedbackBody {
3301 organization_id: organization.map(|organization| organization.id.clone()),
3302 request_id: prediction_id,
3303 rating: match rating {
3304 EditPredictionRating::Positive => "positive".to_string(),
3305 EditPredictionRating::Negative => "negative".to_string(),
3306 },
3307 inputs: inputs?,
3308 output,
3309 expected_output,
3310 feedback,
3311 })
3312 .await?;
3313
3314 anyhow::Ok(())
3315 }
3316 })
3317 .detach_and_log_err(cx);
3318
3319 cx.notify();
3320 }
3321}
3322
3323fn collaborator_edit_overlaps_locality_region(
3324 project_state: &ProjectState,
3325 project: &Entity<Project>,
3326 buffer: &Entity<Buffer>,
3327 snapshot: &BufferSnapshot,
3328 edit_range: &Range<Anchor>,
3329 cx: &App,
3330) -> bool {
3331 let Some((active_buffer, Some(position))) = project_state.active_buffer(project, cx) else {
3332 return false;
3333 };
3334
3335 if active_buffer.entity_id() != buffer.entity_id() {
3336 return false;
3337 }
3338
3339 let locality_point_range = expand_context_syntactically_then_linewise(
3340 snapshot,
3341 (position..position).to_point(snapshot),
3342 COLLABORATOR_EDIT_LOCALITY_CONTEXT_TOKENS,
3343 );
3344 let locality_anchor_range = snapshot.anchor_range_inside(locality_point_range);
3345
3346 edit_range.overlaps(&locality_anchor_range, snapshot)
3347}
3348
3349fn merge_trailing_events_if_needed(
3350 events: &mut VecDeque<StoredEvent>,
3351 end_snapshot: &TextBufferSnapshot,
3352 latest_snapshot: &TextBufferSnapshot,
3353 latest_edit_range: &Range<Anchor>,
3354) {
3355 if let Some(last_event) = events.back() {
3356 if last_event.old_snapshot.remote_id() != latest_snapshot.remote_id() {
3357 return;
3358 }
3359 if !latest_snapshot
3360 .version
3361 .observed_all(&last_event.new_snapshot_version)
3362 {
3363 return;
3364 }
3365 }
3366
3367 let mut next_old_event = None;
3368 let mut mergeable_count = 0;
3369 for old_event in events.iter().rev() {
3370 if let Some(next_old_event) = next_old_event
3371 && !old_event.can_merge(next_old_event, latest_snapshot, latest_edit_range)
3372 {
3373 break;
3374 }
3375 mergeable_count += 1;
3376 next_old_event = Some(old_event);
3377 }
3378
3379 if mergeable_count <= 1 {
3380 return;
3381 }
3382
3383 let merge_start = events.len() - mergeable_count;
3384 let oldest_event = &events[merge_start];
3385 let oldest_snapshot = oldest_event.old_snapshot.clone();
3386 let newest_snapshot = end_snapshot;
3387 let mut merged_edit_range = oldest_event.total_edit_range.clone();
3388
3389 for event in events.range(events.len() - mergeable_count + 1..) {
3390 merged_edit_range =
3391 merge_anchor_ranges(&merged_edit_range, &event.total_edit_range, latest_snapshot);
3392 }
3393
3394 if let Some((diff, old_range, new_range)) = compute_diff_between_snapshots_in_range(
3395 &oldest_snapshot,
3396 newest_snapshot,
3397 &merged_edit_range,
3398 ) {
3399 let merged_event = match oldest_event.event.as_ref() {
3400 zeta_prompt::Event::BufferChange {
3401 old_path,
3402 path,
3403 in_open_source_repo,
3404 ..
3405 } => StoredEvent {
3406 event: Arc::new(zeta_prompt::Event::BufferChange {
3407 old_path: old_path.clone(),
3408 path: path.clone(),
3409 diff,
3410 old_range,
3411 new_range: new_range.clone(),
3412 in_open_source_repo: *in_open_source_repo,
3413 predicted: events.range(merge_start..).all(|event| {
3414 matches!(
3415 event.event.as_ref(),
3416 zeta_prompt::Event::BufferChange {
3417 predicted: true,
3418 ..
3419 }
3420 )
3421 }),
3422 }),
3423 old_snapshot: oldest_snapshot.clone(),
3424 new_snapshot_version: newest_snapshot.version.clone(),
3425 total_edit_range: newest_snapshot.anchor_before(new_range.start)
3426 ..newest_snapshot.anchor_before(new_range.end),
3427 file_context: oldest_event.file_context.clone(),
3428 },
3429 };
3430 events.truncate(events.len() - mergeable_count);
3431 events.push_back(merged_event);
3432 }
3433}
3434
3435fn merge_anchor_ranges(
3436 left: &Range<Anchor>,
3437 right: &Range<Anchor>,
3438 snapshot: &TextBufferSnapshot,
3439) -> Range<Anchor> {
3440 let start = if left.start.cmp(&right.start, snapshot).is_le() {
3441 left.start
3442 } else {
3443 right.start
3444 };
3445 let end = if left.end.cmp(&right.end, snapshot).is_ge() {
3446 left.end
3447 } else {
3448 right.end
3449 };
3450 start..end
3451}
3452
3453#[derive(Error, Debug)]
3454#[error(
3455 "You must update to Omega version {minimum_version} or higher to continue using edit predictions."
3456)]
3457pub struct ZedUpdateRequiredError {
3458 minimum_version: Version,
3459}
3460
3461#[derive(Error, Debug)]
3462#[error("Cloud request timed out")]
3463pub(crate) struct CloudRequestTimeoutError;
3464
3465struct ZedPredictUpsell;
3466
3467fn is_upsell_dismissed(cx: &App) -> bool {
3468 // To make this backwards compatible with older versions of Zed, we
3469 // check if the user has seen the previous Edit Prediction Onboarding
3470 // before, by checking the data collection choice which was written to
3471 // the database once the user clicked on "Accept and Enable"
3472 let kvp = KeyValueStore::global(cx);
3473 if kvp
3474 .read_kvp(ZED_PREDICT_DATA_COLLECTION_CHOICE)
3475 .log_err()
3476 .is_some_and(|s| s.is_some())
3477 {
3478 return true;
3479 }
3480
3481 kvp.read_kvp(ZedPredictUpsell::KEY)
3482 .log_err()
3483 .is_some_and(|s| s.is_some())
3484}
3485
3486impl Dismissable for ZedPredictUpsell {
3487 const KEY: &'static str = "dismissed-edit-predict-upsell";
3488
3489 fn dismissed(cx: &App) -> bool {
3490 is_upsell_dismissed(cx)
3491 }
3492}
3493
3494pub fn should_show_upsell_modal(cx: &App) -> bool {
3495 !is_upsell_dismissed(cx)
3496}
3497
3498pub fn init(cx: &mut App) {
3499 cx.observe_new(move |workspace: &mut Workspace, _, _cx| {
3500 workspace.register_action(
3501 move |workspace, _: &zed_actions::OpenOmegaPredictOnboarding, window, cx| {
3502 ZedPredictModal::toggle(
3503 workspace,
3504 window,
3505 cx,
3506 )
3507 },
3508 );
3509
3510 workspace.register_action(|workspace, _: &ResetOnboarding, _window, cx| {
3511 update_settings_file(workspace.app_state().fs.clone(), cx, move |settings, _| {
3512 settings
3513 .project
3514 .all_languages
3515 .edit_predictions
3516 .get_or_insert_default()
3517 .provider = Some(EditPredictionProvider::None)
3518 });
3519 });
3520 fn copilot_for_project(project: &Entity<Project>, cx: &mut App) -> Option<Entity<Copilot>> {
3521 EditPredictionStore::try_global(cx).and_then(|store| {
3522 store.update(cx, |this, cx| this.start_copilot_for_project(project, cx))
3523 })
3524 }
3525
3526 workspace.register_action(|workspace, _: &SignIn, window, cx| {
3527 if let Some(copilot) = copilot_for_project(workspace.project(), cx) {
3528 copilot_ui::initiate_sign_in(copilot, window, cx);
3529 }
3530 });
3531 workspace.register_action(|workspace, _: &Reinstall, window, cx| {
3532 if let Some(copilot) = copilot_for_project(workspace.project(), cx) {
3533 copilot_ui::reinstall_and_sign_in(copilot, window, cx);
3534 }
3535 });
3536 workspace.register_action(|workspace, _: &SignOut, window, cx| {
3537 if let Some(copilot) = copilot_for_project(workspace.project(), cx) {
3538 copilot_ui::initiate_sign_out(copilot, window, cx);
3539 }
3540 });
3541 })
3542 .detach();
3543}
3544
3545fn is_zed_industries_repo(url: &str) -> bool {
3546 url.strip_prefix("https://github.com/zed-industries/")
3547 .or_else(|| url.strip_prefix("http://github.com/zed-industries/"))
3548 .or_else(|| url.strip_prefix("git@github.com:zed-industries/"))
3549 .or_else(|| url.strip_prefix("ssh://git@github.com/zed-industries/"))
3550 .is_some_and(|repo| !repo.is_empty())
3551}
3552