tenant.openagents/omega

1use std::{cmp, sync::Arc};
2
3use client::{Client, UserStore};
4use cloud_llm_client::EditPredictionRejectReason;
5use edit_prediction_types::{
6    DataCollectionState, EditPredictionDelegate, EditPredictionDiscardReason,
7    EditPredictionIconSet, EditPredictionRequestTrigger, SuggestionDisplayType,
8};
9use feature_flags::FeatureFlagAppExt;
10use fs::Fs;
11use gpui::{App, Entity, prelude::*};
12use language::{Buffer, ToPoint as _};
13use project::Project;
14use settings::{EditPredictionDataCollectionChoice, update_settings_file};
15
16use crate::{BufferEditPrediction, EditPredictionStore};
17
18pub struct ZedEditPredictionDelegate {
19    store: Entity<EditPredictionStore>,
20    project: Entity<Project>,
21    singleton_buffer: Option<Entity<Buffer>>,
22}
23
24impl ZedEditPredictionDelegate {
25    pub fn new(
26        project: Entity<Project>,
27        singleton_buffer: Option<Entity<Buffer>>,
28        client: &Arc<Client>,
29        user_store: &Entity<UserStore>,
30        cx: &mut Context<Self>,
31    ) -> Self {
32        let store = EditPredictionStore::global(client, user_store, cx);
33        store.update(cx, |store, cx| {
34            store.register_project(&project, cx);
35        });
36
37        cx.observe(&store, |_this, _ep_store, cx| {
38            cx.notify();
39        })
40        .detach();
41
42        Self {
43            project: project,
44            store: store,
45            singleton_buffer,
46        }
47    }
48}
49
50impl EditPredictionDelegate for ZedEditPredictionDelegate {
51    fn name() -> &'static str {
52        "zed-predict"
53    }
54
55    fn display_name() -> &'static str {
56        "Zed's Edit Predictions"
57    }
58
59    fn show_predictions_in_menu() -> bool {
60        true
61    }
62
63    fn show_tab_accept_marker() -> bool {
64        true
65    }
66
67    fn icons(&self, cx: &App) -> EditPredictionIconSet {
68        self.store.read(cx).icons(cx)
69    }
70
71    fn data_collection_state(&self, cx: &App) -> DataCollectionState {
72        if let Some(buffer) = &self.singleton_buffer
73            && let Some(file) = buffer.read(cx).file()
74        {
75            let is_project_open_source =
76                self.store
77                    .read(cx)
78                    .is_file_open_source(&self.project, file, cx);
79
80            if self.store.read(cx).is_data_collection_enabled(cx) {
81                DataCollectionState::Enabled {
82                    is_project_open_source,
83                }
84            } else {
85                DataCollectionState::Disabled {
86                    is_project_open_source,
87                }
88            }
89        } else {
90            DataCollectionState::Disabled {
91                is_project_open_source: false,
92            }
93        }
94    }
95
96    fn can_toggle_data_collection(&self, cx: &App) -> bool {
97        if cx.is_staff() {
98            return false;
99        }
100
101        self.store
102            .read(cx)
103            .is_data_collection_allowed_by_organization(cx)
104    }
105
106    fn toggle_data_collection(&mut self, cx: &mut App) {
107        let fs = <dyn Fs>::global(cx);
108        let is_currently_enabled = self.store.read(cx).is_data_collection_enabled(cx);
109        update_settings_file(fs, cx, move |settings, _| {
110            let edit_predictions = settings
111                .project
112                .all_languages
113                .edit_predictions
114                .get_or_insert_default();
115
116            edit_predictions.allow_data_collection = Some(if is_currently_enabled {
117                EditPredictionDataCollectionChoice::No
118            } else {
119                EditPredictionDataCollectionChoice::Yes
120            });
121        });
122    }
123
124    fn usage(&self, cx: &App) -> Option<client::EditPredictionUsage> {
125        self.store.read(cx).usage(cx)
126    }
127
128    fn is_enabled(
129        &self,
130        _buffer: &Entity<language::Buffer>,
131        _cursor_position: language::Anchor,
132        _cx: &App,
133    ) -> bool {
134        true
135    }
136
137    fn is_refreshing(&self, cx: &App) -> bool {
138        self.store.read(cx).is_refreshing(&self.project)
139    }
140
141    fn refresh(
142        &mut self,
143        buffer: Entity<language::Buffer>,
144        cursor_position: language::Anchor,
145        _debounce: bool,
146        trigger: EditPredictionRequestTrigger,
147        cx: &mut Context<Self>,
148    ) {
149        let store = self.store.read(cx);
150
151        if store.user_store.read_with(cx, |user_store, _cx| {
152            user_store.account_too_young() || user_store.has_overdue_invoices()
153        }) {
154            return;
155        }
156
157        self.store.update(cx, |store, cx| {
158            if let Some(current) =
159                store.prediction_at(&buffer, Some(cursor_position), &self.project, cx)
160                && let BufferEditPrediction::Local { prediction } = current
161                && prediction.interpolate(buffer.read(cx)).is_some()
162            {
163                return;
164            }
165
166            store.refresh_context(&self.project, &buffer, cursor_position, cx);
167            store.refresh_prediction_from_buffer(
168                self.project.clone(),
169                buffer,
170                cursor_position,
171                trigger,
172                cx,
173            )
174        });
175    }
176
177    fn accept(&mut self, cx: &mut Context<Self>) {
178        self.store.update(cx, |store, cx| {
179            store.accept_current_prediction(&self.project, cx);
180        });
181    }
182
183    fn discard(&mut self, reason: EditPredictionDiscardReason, cx: &mut Context<Self>) {
184        let reject_reason = match reason {
185            EditPredictionDiscardReason::Rejected => EditPredictionRejectReason::Rejected,
186            EditPredictionDiscardReason::Ignored => EditPredictionRejectReason::Discarded,
187        };
188        self.store.update(cx, |store, cx| {
189            store.reject_current_prediction(reject_reason, &self.project, cx);
190        });
191    }
192
193    fn did_show(&mut self, display_type: SuggestionDisplayType, cx: &mut Context<Self>) {
194        self.store.update(cx, |store, cx| {
195            store.did_show_current_prediction(&self.project, display_type, cx);
196        });
197    }
198
199    fn suggest(
200        &mut self,
201        buffer: &Entity<language::Buffer>,
202        cursor_position: language::Anchor,
203        cx: &mut Context<Self>,
204    ) -> Option<edit_prediction_types::EditPrediction> {
205        self.store.update(cx, |store, cx| {
206            let prediction =
207                store.prediction_at(buffer, Some(cursor_position), &self.project, cx)?;
208
209            let prediction = match prediction {
210                BufferEditPrediction::Local { prediction } => prediction,
211                BufferEditPrediction::Jump { prediction } => {
212                    return Some(edit_prediction_types::EditPrediction::Jump {
213                        id: Some(prediction.id.0.clone()),
214                        snapshot: prediction.snapshot.clone(),
215                        target: prediction.edits.first().unwrap().0.start,
216                    });
217                }
218            };
219
220            let buffer = buffer.read(cx);
221            let snapshot = buffer.snapshot();
222
223            let Some(edits) = prediction.interpolate(&snapshot) else {
224                store.reject_current_prediction(
225                    EditPredictionRejectReason::InterpolateFailed,
226                    &self.project,
227                    cx,
228                );
229                return None;
230            };
231
232            if edits.is_empty() {
233                store.reject_current_prediction(
234                    EditPredictionRejectReason::InterpolatedEmpty,
235                    &self.project,
236                    cx,
237                );
238                return None;
239            }
240
241            let cursor_row = cursor_position.to_point(&snapshot).row;
242            let (closest_edit_ix, (closest_edit_range, _)) =
243                edits.iter().enumerate().min_by_key(|(_, (range, _))| {
244                    let distance_from_start =
245                        cursor_row.abs_diff(range.start.to_point(&snapshot).row);
246                    let distance_from_end = cursor_row.abs_diff(range.end.to_point(&snapshot).row);
247                    cmp::min(distance_from_start, distance_from_end)
248                })?;
249
250            let mut edit_start_ix = closest_edit_ix;
251            for (range, _) in edits[..edit_start_ix].iter().rev() {
252                let distance_from_closest_edit = closest_edit_range.start.to_point(&snapshot).row
253                    - range.end.to_point(&snapshot).row;
254                if distance_from_closest_edit <= 1 {
255                    edit_start_ix -= 1;
256                } else {
257                    break;
258                }
259            }
260
261            let mut edit_end_ix = closest_edit_ix + 1;
262            for (range, _) in &edits[edit_end_ix..] {
263                let distance_from_closest_edit = range.start.to_point(buffer).row
264                    - closest_edit_range.end.to_point(&snapshot).row;
265                if distance_from_closest_edit <= 1 {
266                    edit_end_ix += 1;
267                } else {
268                    break;
269                }
270            }
271
272            Some(edit_prediction_types::EditPrediction::Local {
273                id: Some(prediction.id.0.clone()),
274                edits: edits[edit_start_ix..edit_end_ix].to_vec(),
275                cursor_position: prediction.cursor_position,
276                edit_preview: Some(prediction.edit_preview.clone()),
277            })
278        })
279    }
280}
281