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
157 .store
158 .read(cx)
159 .context_for_project_with_buffers(&self.project, cx)
160 .map_or(Vec::new(), |files| files.collect());
161
162 let editor = run.editor.clone();
163 let multibuffer = run.editor.read(cx).buffer().clone();
164
165 if self.current_ix + 2 == self.runs.len() {
166 self.current_ix += 1;
167 }
168
169 cx.spawn_in(window, async move |this, cx| {
170 let mut paths = Vec::new();
171 for (related_file, buffer) in related_files {
172 let point_ranges = related_file
173 .excerpts
174 .iter()
175 .map(|excerpt| {
176 Point::new(excerpt.row_range.start, 0)..Point::new(excerpt.row_range.end, 0)
177 })
178 .collect::<Vec<_>>();
179 cx.update(|_, cx| {
180 let path = PathKey::for_buffer(&buffer, cx);
181 paths.push((path, buffer, point_ranges));
182 })?;
183 }
184
185 multibuffer.update(cx, |multibuffer, cx| {
186 multibuffer.clear(cx);
187
188 for (path, buffer, ranges) in paths {
189 multibuffer.set_excerpts_for_path(path, buffer, ranges, 0, cx);
190 }
191 })?;
192
193 editor.update_in(cx, |editor, window, cx| {
194 editor.move_to_beginning(&Default::default(), window, cx);
195 })?;
196
197 this.update(cx, |_, cx| cx.notify())
198 })
199 .detach();
200 }
201
202 fn handle_go_back(
203 &mut self,
204 _: &EditPredictionContextGoBack,
205 window: &mut Window,
206 cx: &mut Context<Self>,
207 ) {
208 self.current_ix = self.current_ix.saturating_sub(1);
209 cx.focus_self(window);
210 cx.notify();
211 }
212
213 fn handle_go_forward(
214 &mut self,
215 _: &EditPredictionContextGoForward,
216 window: &mut Window,
217 cx: &mut Context<Self>,
218 ) {
219 self.current_ix = self
220 .current_ix
221 .add(1)
222 .min(self.runs.len().saturating_sub(1));
223 cx.focus_self(window);
224 cx.notify();
225 }
226
227 fn render_informational_footer(
228 &self,
229 cx: &mut Context<'_, EditPredictionContextView>,
230 ) -> ui::Div {
231 let run = &self.runs[self.current_ix];
232 let new_run_started = self
233 .runs
234 .back()
235 .map_or(false, |latest_run| latest_run.finished_at.is_none());
236
237 h_flex()
238 .p_2()
239 .w_full()
240 .font_buffer(cx)
241 .text_xs()
242 .border_t_1()
243 .gap_2()
244 .child(v_flex().h_full().flex_1().child({
245 let t0 = run.started_at;
246 let mut table = ui::Table::<2>::new().width(ui::px(300.)).no_ui_font();
247 for (key, value) in &run.metadata {
248 table = table.row([key.into_any_element(), value.clone().into_any_element()])
249 }
250 table = table.row([
251 "Total Time".into_any_element(),
252 format!("{} ms", (run.finished_at.unwrap_or(t0) - t0).as_millis())
253 .into_any_element(),
254 ]);
255 table
256 }))
257 .child(
258 v_flex().h_full().text_align(TextAlign::Right).child(
259 h_flex()
260 .justify_end()
261 .child(
262 IconButton::new("go-back", IconName::ChevronLeft)
263 .disabled(self.current_ix == 0 || self.runs.len() < 2)
264 .tooltip(ui::Tooltip::for_action_title(
265 "Go to previous run",
266 &EditPredictionContextGoBack,
267 ))
268 .on_click(cx.listener(|this, _, window, cx| {
269 this.handle_go_back(&EditPredictionContextGoBack, window, cx);
270 })),
271 )
272 .child(
273 div()
274 .child(format!("{}/{}", self.current_ix + 1, self.runs.len()))
275 .map(|this| {
276 if new_run_started {
277 this.with_animation(
278 "pulsating-count",
279 Animation::new(Duration::from_secs(2))
280 .repeat()
281 .with_easing(pulsating_between(0.4, 0.8)),
282 |label, delta| label.opacity(delta),
283 )
284 .into_any_element()
285 } else {
286 this.into_any_element()
287 }
288 }),
289 )
290 .child(
291 IconButton::new("go-forward", IconName::ChevronRight)
292 .disabled(self.current_ix + 1 == self.runs.len())
293 .tooltip(ui::Tooltip::for_action_title(
294 "Go to next run",
295 &EditPredictionContextGoBack,
296 ))
297 .on_click(cx.listener(|this, _, window, cx| {
298 this.handle_go_forward(
299 &EditPredictionContextGoForward,
300 window,
301 cx,
302 );
303 })),
304 ),
305 ),
306 )
307 }
308}
309
310impl Focusable for EditPredictionContextView {
311 fn focus_handle(&self, cx: &App) -> FocusHandle {
312 self.runs
313 .get(self.current_ix)
314 .map(|run| run.editor.read(cx).focus_handle(cx))
315 .unwrap_or_else(|| self.empty_focus_handle.clone())
316 }
317}
318
319impl EventEmitter<()> for EditPredictionContextView {}
320
321impl Item for EditPredictionContextView {
322 type Event = ();
323
324 fn tab_content_text(&self, _detail: usize, _cx: &App) -> SharedString {
325 "Edit Prediction Context".into()
326 }
327
328 fn buffer_kind(&self, _cx: &App) -> workspace::item::ItemBufferKind {
329 workspace::item::ItemBufferKind::Multibuffer
330 }
331
332 fn act_as_type<'a>(
333 &'a self,
334 type_id: TypeId,
335 self_handle: &'a Entity<Self>,
336 _: &'a App,
337 ) -> Option<gpui::AnyEntity> {
338 if type_id == TypeId::of::<Self>() {
339 Some(self_handle.clone().into())
340 } else if type_id == TypeId::of::<Editor>() {
341 Some(self.runs.get(self.current_ix)?.editor.clone().into())
342 } else {
343 None
344 }
345 }
346}
347
348impl gpui::Render for EditPredictionContextView {
349 fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl ui::IntoElement {
350 v_flex()
351 .key_context("EditPredictionContext")
352 .on_action(cx.listener(Self::handle_go_back))
353 .on_action(cx.listener(Self::handle_go_forward))
354 .size_full()
355 .map(|this| {
356 if self.runs.is_empty() {
357 this.child(
358 v_flex()
359 .size_full()
360 .justify_center()
361 .items_center()
362 .child("No retrieval runs yet"),
363 )
364 } else {
365 this.child(self.runs[self.current_ix].editor.clone())
366 .child(self.render_informational_footer(cx))
367 }
368 })
369 }
370}