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 let (anchor_ranges, _) =
205 multibuffer.set_excerpts_for_path(path, buffer, ranges, 0, cx);
206 for (anchor_range, order) in anchor_ranges.into_iter().zip(orders) {
207 excerpt_anchors_with_orders.push((anchor_range.start, order));
208 }
209 }
210 });
211
212 editor.update_in(cx, |editor, window, cx| {
213 let blocks = excerpt_anchors_with_orders
214 .into_iter()
215 .map(|(anchor, order)| {
216 let label = SharedString::from(format!("order: {order}"));
217 BlockProperties {
218 placement: BlockPlacement::Above(anchor),
219 height: Some(1),
220 style: BlockStyle::Sticky,
221 render: Arc::new(move |cx| {
222 div()
223 .pl(cx.anchor_x)
224 .text_ui_xs(cx)
225 .text_color(cx.editor_style.status.info)
226 .child(label.clone())
227 .into_any_element()
228 }),
229 priority: 0,
230 }
231 })
232 .collect::<Vec<_>>();
233 editor.insert_blocks(blocks, None, cx);
234 editor.move_to_beginning(&Default::default(), window, cx);
235 })?;
236
237 this.update(cx, |_, cx| cx.notify())
238 })
239 .detach();
240 }
241
242 fn handle_go_back(
243 &mut self,
244 _: &EditPredictionContextGoBack,
245 window: &mut Window,
246 cx: &mut Context<Self>,
247 ) {
248 self.current_ix = self.current_ix.saturating_sub(1);
249 cx.focus_self(window);
250 cx.notify();
251 }
252
253 fn handle_go_forward(
254 &mut self,
255 _: &EditPredictionContextGoForward,
256 window: &mut Window,
257 cx: &mut Context<Self>,
258 ) {
259 self.current_ix = self
260 .current_ix
261 .add(1)
262 .min(self.runs.len().saturating_sub(1));
263 cx.focus_self(window);
264 cx.notify();
265 }
266
267 fn render_informational_footer(
268 &self,
269 cx: &mut Context<'_, EditPredictionContextView>,
270 ) -> ui::Div {
271 let run = &self.runs[self.current_ix];
272 let new_run_started = self
273 .runs
274 .back()
275 .map_or(false, |latest_run| latest_run.finished_at.is_none());
276
277 h_flex()
278 .p_2()
279 .w_full()
280 .font_buffer(cx)
281 .text_xs()
282 .border_t_1()
283 .gap_2()
284 .child(v_flex().h_full().flex_1().child({
285 let t0 = run.started_at;
286 let mut table = ui::Table::new(2).width(ui::px(300.)).no_ui_font();
287 for (key, value) in &run.metadata {
288 table = table.row(vec![
289 key.into_any_element(),
290 value.clone().into_any_element(),
291 ])
292 }
293 table = table.row(vec![
294 "Total Time".into_any_element(),
295 format!("{} ms", (run.finished_at.unwrap_or(t0) - t0).as_millis())
296 .into_any_element(),
297 ]);
298 table
299 }))
300 .child(
301 v_flex().h_full().text_align(TextAlign::Right).child(
302 h_flex()
303 .justify_end()
304 .child(
305 IconButton::new("go-back", IconName::ChevronLeft)
306 .disabled(self.current_ix == 0 || self.runs.len() < 2)
307 .tooltip(ui::Tooltip::for_action_title(
308 "Go to previous run",
309 &EditPredictionContextGoBack,
310 ))
311 .on_click(cx.listener(|this, _, window, cx| {
312 this.handle_go_back(&EditPredictionContextGoBack, window, cx);
313 })),
314 )
315 .child(
316 div()
317 .child(format!("{}/{}", self.current_ix + 1, self.runs.len()))
318 .map(|this| {
319 if new_run_started {
320 this.with_animation(
321 "pulsating-count",
322 Animation::new(Duration::from_secs(2))
323 .repeat()
324 .with_easing(pulsating_between(0.4, 0.8)),
325 |label, delta| label.opacity(delta),
326 )
327 .into_any_element()
328 } else {
329 this.into_any_element()
330 }
331 }),
332 )
333 .child(
334 IconButton::new("go-forward", IconName::ChevronRight)
335 .disabled(self.current_ix + 1 == self.runs.len())
336 .tooltip(ui::Tooltip::for_action_title(
337 "Go to next run",
338 &EditPredictionContextGoBack,
339 ))
340 .on_click(cx.listener(|this, _, window, cx| {
341 this.handle_go_forward(
342 &EditPredictionContextGoForward,
343 window,
344 cx,
345 );
346 })),
347 ),
348 ),
349 )
350 }
351}
352
353impl Focusable for EditPredictionContextView {
354 fn focus_handle(&self, cx: &App) -> FocusHandle {
355 self.runs
356 .get(self.current_ix)
357 .map(|run| run.editor.read(cx).focus_handle(cx))
358 .unwrap_or_else(|| self.empty_focus_handle.clone())
359 }
360}
361
362impl EventEmitter<()> for EditPredictionContextView {}
363
364impl Item for EditPredictionContextView {
365 type Event = ();
366
367 fn tab_content_text(&self, _detail: usize, _cx: &App) -> SharedString {
368 "Edit Prediction Context".into()
369 }
370
371 fn buffer_kind(&self, _cx: &App) -> workspace::item::ItemBufferKind {
372 workspace::item::ItemBufferKind::Multibuffer
373 }
374
375 fn act_as_type<'a>(
376 &'a self,
377 type_id: TypeId,
378 self_handle: &'a Entity<Self>,
379 _: &'a App,
380 ) -> Option<gpui::AnyEntity> {
381 if type_id == TypeId::of::<Self>() {
382 Some(self_handle.clone().into())
383 } else if type_id == TypeId::of::<Editor>() {
384 Some(self.runs.get(self.current_ix)?.editor.clone().into())
385 } else {
386 None
387 }
388 }
389}
390
391impl gpui::Render for EditPredictionContextView {
392 fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl ui::IntoElement {
393 v_flex()
394 .key_context("EditPredictionContext")
395 .on_action(cx.listener(Self::handle_go_back))
396 .on_action(cx.listener(Self::handle_go_forward))
397 .size_full()
398 .map(|this| {
399 if self.runs.is_empty() {
400 this.child(
401 v_flex()
402 .size_full()
403 .justify_center()
404 .items_center()
405 .child("No retrieval runs yet"),
406 )
407 } else {
408 this.child(self.runs[self.current_ix].editor.clone())
409 .child(self.render_informational_footer(cx))
410 }
411 })
412 }
413}