1use std::{
2 any::TypeId,
3 collections::VecDeque,
4 ops::Add,
5 sync::Arc,
6 time::{Duration, Instant},
7};
8
9use anyhow::Result;
10use client::{Client, UserStore};
11use editor::{Editor, PathKey};
12use futures::StreamExt as _;
13use gpui::{
14 Animation, AnimationExt, App, AppContext as _, Context, Entity, EventEmitter, FocusHandle,
15 Focusable, InteractiveElement as _, IntoElement as _, ParentElement as _, SharedString,
16 Styled as _, Task, TextAlign, Window, actions, div, pulsating_between,
17};
18use multi_buffer::MultiBuffer;
19use project::Project;
20use text::Point;
21use ui::{
22 ButtonCommon, Clickable, Disableable, FluentBuilder as _, IconButton, IconName,
23 StyledTypography as _, h_flex, v_flex,
24};
25
26use edit_prediction::{
27 ContextRetrievalFinishedDebugEvent, ContextRetrievalStartedDebugEvent, DebugEvent,
28 EditPredictionStore,
29};
30use workspace::Item;
31
32pub struct EditPredictionContextView {
33 empty_focus_handle: FocusHandle,
34 project: Entity<Project>,
35 store: Entity<EditPredictionStore>,
36 runs: VecDeque<RetrievalRun>,
37 current_ix: usize,
38 _update_task: Task<Result<()>>,
39}
40
41#[derive(Debug)]
42struct RetrievalRun {
43 editor: Entity<Editor>,
44 started_at: Instant,
45 metadata: Vec<(&'static str, SharedString)>,
46 finished_at: Option<Instant>,
47}
48
49actions!(
50 dev,
51 [
52 /// Go to the previous context retrieval run
53 EditPredictionContextGoBack,
54 /// Go to the next context retrieval run
55 EditPredictionContextGoForward
56 ]
57);
58
59impl EditPredictionContextView {
60 pub fn new(
61 project: Entity<Project>,
62 client: &Arc<Client>,
63 user_store: &Entity<UserStore>,
64 window: &mut gpui::Window,
65 cx: &mut Context<Self>,
66 ) -> Self {
67 let store = EditPredictionStore::global(client, user_store, cx);
68
69 let mut debug_rx = store.update(cx, |store, cx| store.debug_info(&project, cx));
70 let _update_task = cx.spawn_in(window, async move |this, cx| {
71 while let Some(event) = debug_rx.next().await {
72 this.update_in(cx, |this, window, cx| {
73 this.handle_store_event(event, window, cx)
74 })?;
75 }
76 Ok(())
77 });
78
79 Self {
80 empty_focus_handle: cx.focus_handle(),
81 project,
82 runs: VecDeque::new(),
83 current_ix: 0,
84 store,
85 _update_task,
86 }
87 }
88
89 fn handle_store_event(
90 &mut self,
91 event: DebugEvent,
92 window: &mut gpui::Window,
93 cx: &mut Context<Self>,
94 ) {
95 match event {
96 DebugEvent::ContextRetrievalStarted(info) => {
97 if info.project_entity_id == self.project.entity_id() {
98 self.handle_context_retrieval_started(info, window, cx);
99 }
100 }
101 DebugEvent::ContextRetrievalFinished(info) => {
102 if info.project_entity_id == self.project.entity_id() {
103 self.handle_context_retrieval_finished(info, window, cx);
104 }
105 }
106 DebugEvent::EditPredictionStarted(_) => {}
107 DebugEvent::EditPredictionFinished(_) => {}
108 }
109 }
110
111 fn handle_context_retrieval_started(
112 &mut self,
113 info: ContextRetrievalStartedDebugEvent,
114 window: &mut Window,
115 cx: &mut Context<Self>,
116 ) {
117 if self
118 .runs
119 .back()
120 .is_some_and(|run| run.finished_at.is_none())
121 {
122 self.runs.pop_back();
123 }
124
125 let multibuffer = cx.new(|_| MultiBuffer::new(language::Capability::ReadOnly));
126 let editor = cx
127 .new(|cx| Editor::for_multibuffer(multibuffer, Some(self.project.clone()), window, cx));
128
129 if self.runs.len() == 32 {
130 self.runs.pop_front();
131 }
132
133 self.runs.push_back(RetrievalRun {
134 editor,
135 started_at: info.timestamp,
136 finished_at: None,
137 metadata: Vec::new(),
138 });
139
140 cx.notify();
141 }
142
143 fn handle_context_retrieval_finished(
144 &mut self,
145 info: ContextRetrievalFinishedDebugEvent,
146 window: &mut Window,
147 cx: &mut Context<Self>,
148 ) {
149 let Some(run) = self.runs.back_mut() else {
150 return;
151 };
152
153 run.finished_at = Some(info.timestamp);
154 run.metadata = info.metadata;
155
156 let related_files = self.store.update(cx, |store, cx| {
157 store.context_for_project_with_buffers(&self.project, cx)
158 });
159
160 let editor = run.editor.clone();
161 let multibuffer = run.editor.read(cx).buffer().clone();
162
163 if self.current_ix + 2 == self.runs.len() {
164 self.current_ix += 1;
165 }
166
167 cx.spawn_in(window, async move |this, cx| {
168 let mut paths = Vec::new();
169 for (related_file, buffer) in related_files {
170 let point_ranges = related_file
171 .excerpts
172 .iter()
173 .map(|excerpt| {
174 Point::new(excerpt.row_range.start, 0)..Point::new(excerpt.row_range.end, 0)
175 })
176 .collect::<Vec<_>>();
177 cx.update(|_, cx| {
178 let path = PathKey::for_buffer(&buffer, cx);
179 paths.push((path, buffer, point_ranges));
180 })?;
181 }
182
183 multibuffer.update(cx, |multibuffer, cx| {
184 multibuffer.clear(cx);
185
186 for (path, buffer, ranges) in paths {
187 multibuffer.set_excerpts_for_path(path, buffer, ranges, 0, cx);
188 }
189 });
190
191 editor.update_in(cx, |editor, window, cx| {
192 editor.move_to_beginning(&Default::default(), window, cx);
193 })?;
194
195 this.update(cx, |_, cx| cx.notify())
196 })
197 .detach();
198 }
199
200 fn handle_go_back(
201 &mut self,
202 _: &EditPredictionContextGoBack,
203 window: &mut Window,
204 cx: &mut Context<Self>,
205 ) {
206 self.current_ix = self.current_ix.saturating_sub(1);
207 cx.focus_self(window);
208 cx.notify();
209 }
210
211 fn handle_go_forward(
212 &mut self,
213 _: &EditPredictionContextGoForward,
214 window: &mut Window,
215 cx: &mut Context<Self>,
216 ) {
217 self.current_ix = self
218 .current_ix
219 .add(1)
220 .min(self.runs.len().saturating_sub(1));
221 cx.focus_self(window);
222 cx.notify();
223 }
224
225 fn render_informational_footer(
226 &self,
227 cx: &mut Context<'_, EditPredictionContextView>,
228 ) -> ui::Div {
229 let run = &self.runs[self.current_ix];
230 let new_run_started = self
231 .runs
232 .back()
233 .map_or(false, |latest_run| latest_run.finished_at.is_none());
234
235 h_flex()
236 .p_2()
237 .w_full()
238 .font_buffer(cx)
239 .text_xs()
240 .border_t_1()
241 .gap_2()
242 .child(v_flex().h_full().flex_1().child({
243 let t0 = run.started_at;
244 let mut table = ui::Table::new(2).width(ui::px(300.)).no_ui_font();
245 for (key, value) in &run.metadata {
246 table = table.row(vec![
247 key.into_any_element(),
248 value.clone().into_any_element(),
249 ])
250 }
251 table = table.row(vec![
252 "Total Time".into_any_element(),
253 format!("{} ms", (run.finished_at.unwrap_or(t0) - t0).as_millis())
254 .into_any_element(),
255 ]);
256 table
257 }))
258 .child(
259 v_flex().h_full().text_align(TextAlign::Right).child(
260 h_flex()
261 .justify_end()
262 .child(
263 IconButton::new("go-back", IconName::ChevronLeft)
264 .disabled(self.current_ix == 0 || self.runs.len() < 2)
265 .tooltip(ui::Tooltip::for_action_title(
266 "Go to previous run",
267 &EditPredictionContextGoBack,
268 ))
269 .on_click(cx.listener(|this, _, window, cx| {
270 this.handle_go_back(&EditPredictionContextGoBack, window, cx);
271 })),
272 )
273 .child(
274 div()
275 .child(format!("{}/{}", self.current_ix + 1, self.runs.len()))
276 .map(|this| {
277 if new_run_started {
278 this.with_animation(
279 "pulsating-count",
280 Animation::new(Duration::from_secs(2))
281 .repeat()
282 .with_easing(pulsating_between(0.4, 0.8)),
283 |label, delta| label.opacity(delta),
284 )
285 .into_any_element()
286 } else {
287 this.into_any_element()
288 }
289 }),
290 )
291 .child(
292 IconButton::new("go-forward", IconName::ChevronRight)
293 .disabled(self.current_ix + 1 == self.runs.len())
294 .tooltip(ui::Tooltip::for_action_title(
295 "Go to next run",
296 &EditPredictionContextGoBack,
297 ))
298 .on_click(cx.listener(|this, _, window, cx| {
299 this.handle_go_forward(
300 &EditPredictionContextGoForward,
301 window,
302 cx,
303 );
304 })),
305 ),
306 ),
307 )
308 }
309}
310
311impl Focusable for EditPredictionContextView {
312 fn focus_handle(&self, cx: &App) -> FocusHandle {
313 self.runs
314 .get(self.current_ix)
315 .map(|run| run.editor.read(cx).focus_handle(cx))
316 .unwrap_or_else(|| self.empty_focus_handle.clone())
317 }
318}
319
320impl EventEmitter<()> for EditPredictionContextView {}
321
322impl Item for EditPredictionContextView {
323 type Event = ();
324
325 fn tab_content_text(&self, _detail: usize, _cx: &App) -> SharedString {
326 "Edit Prediction Context".into()
327 }
328
329 fn buffer_kind(&self, _cx: &App) -> workspace::item::ItemBufferKind {
330 workspace::item::ItemBufferKind::Multibuffer
331 }
332
333 fn act_as_type<'a>(
334 &'a self,
335 type_id: TypeId,
336 self_handle: &'a Entity<Self>,
337 _: &'a App,
338 ) -> Option<gpui::AnyEntity> {
339 if type_id == TypeId::of::<Self>() {
340 Some(self_handle.clone().into())
341 } else if type_id == TypeId::of::<Editor>() {
342 Some(self.runs.get(self.current_ix)?.editor.clone().into())
343 } else {
344 None
345 }
346 }
347}
348
349impl gpui::Render for EditPredictionContextView {
350 fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl ui::IntoElement {
351 v_flex()
352 .key_context("EditPredictionContext")
353 .on_action(cx.listener(Self::handle_go_back))
354 .on_action(cx.listener(Self::handle_go_forward))
355 .size_full()
356 .map(|this| {
357 if self.runs.is_empty() {
358 this.child(
359 v_flex()
360 .size_full()
361 .justify_center()
362 .items_center()
363 .child("No retrieval runs yet"),
364 )
365 } else {
366 this.child(self.runs[self.current_ix].editor.clone())
367 .child(self.render_informational_footer(cx))
368 }
369 })
370 }
371}