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