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::OffsetRangeExt;
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, _| store.debug_info());
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::EditPredictionRequested(_) => {}
107 }
108 }
109
110 fn handle_context_retrieval_started(
111 &mut self,
112 info: ContextRetrievalStartedDebugEvent,
113 window: &mut Window,
114 cx: &mut Context<Self>,
115 ) {
116 if self
117 .runs
118 .back()
119 .is_some_and(|run| run.finished_at.is_none())
120 {
121 self.runs.pop_back();
122 }
123
124 let multibuffer = cx.new(|_| MultiBuffer::new(language::Capability::ReadOnly));
125 let editor = cx
126 .new(|cx| Editor::for_multibuffer(multibuffer, Some(self.project.clone()), window, cx));
127
128 if self.runs.len() == 32 {
129 self.runs.pop_front();
130 }
131
132 self.runs.push_back(RetrievalRun {
133 editor,
134 started_at: info.timestamp,
135 finished_at: None,
136 metadata: Vec::new(),
137 });
138
139 cx.notify();
140 }
141
142 fn handle_context_retrieval_finished(
143 &mut self,
144 info: ContextRetrievalFinishedDebugEvent,
145 window: &mut Window,
146 cx: &mut Context<Self>,
147 ) {
148 let Some(run) = self.runs.back_mut() else {
149 return;
150 };
151
152 run.finished_at = Some(info.timestamp);
153 run.metadata = info.metadata;
154
155 let project = self.project.clone();
156 let related_files = self
157 .store
158 .read(cx)
159 .context_for_project(&self.project, cx)
160 .to_vec();
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 in related_files {
172 let (buffer, point_ranges): (_, Vec<_>) =
173 if let Some(buffer) = related_file.buffer.upgrade() {
174 let snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot())?;
175
176 (
177 buffer,
178 related_file
179 .excerpts
180 .iter()
181 .map(|excerpt| excerpt.anchor_range.to_point(&snapshot))
182 .collect(),
183 )
184 } else {
185 (
186 project
187 .update(cx, |project, cx| {
188 project.open_buffer(related_file.path.clone(), cx)
189 })?
190 .await?,
191 related_file
192 .excerpts
193 .iter()
194 .map(|excerpt| excerpt.point_range.clone())
195 .collect(),
196 )
197 };
198 cx.update(|_, cx| {
199 let path = PathKey::for_buffer(&buffer, cx);
200 paths.push((path, buffer, point_ranges));
201 })?;
202 }
203
204 multibuffer.update(cx, |multibuffer, cx| {
205 multibuffer.clear(cx);
206
207 for (path, buffer, ranges) in paths {
208 multibuffer.set_excerpts_for_path(path, buffer, ranges, 0, cx);
209 }
210 })?;
211
212 editor.update_in(cx, |editor, window, cx| {
213 editor.move_to_beginning(&Default::default(), window, cx);
214 })?;
215
216 this.update(cx, |_, cx| cx.notify())
217 })
218 .detach();
219 }
220
221 fn handle_go_back(
222 &mut self,
223 _: &EditPredictionContextGoBack,
224 window: &mut Window,
225 cx: &mut Context<Self>,
226 ) {
227 self.current_ix = self.current_ix.saturating_sub(1);
228 cx.focus_self(window);
229 cx.notify();
230 }
231
232 fn handle_go_forward(
233 &mut self,
234 _: &EditPredictionContextGoForward,
235 window: &mut Window,
236 cx: &mut Context<Self>,
237 ) {
238 self.current_ix = self
239 .current_ix
240 .add(1)
241 .min(self.runs.len().saturating_sub(1));
242 cx.focus_self(window);
243 cx.notify();
244 }
245
246 fn render_informational_footer(
247 &self,
248 cx: &mut Context<'_, EditPredictionContextView>,
249 ) -> ui::Div {
250 let run = &self.runs[self.current_ix];
251 let new_run_started = self
252 .runs
253 .back()
254 .map_or(false, |latest_run| latest_run.finished_at.is_none());
255
256 h_flex()
257 .p_2()
258 .w_full()
259 .font_buffer(cx)
260 .text_xs()
261 .border_t_1()
262 .gap_2()
263 .child(v_flex().h_full().flex_1().child({
264 let t0 = run.started_at;
265 let mut table = ui::Table::<2>::new().width(ui::px(300.)).no_ui_font();
266 for (key, value) in &run.metadata {
267 table = table.row([key.into_any_element(), value.clone().into_any_element()])
268 }
269 table = table.row([
270 "Total Time".into_any_element(),
271 format!("{} ms", (run.finished_at.unwrap_or(t0) - t0).as_millis())
272 .into_any_element(),
273 ]);
274 table
275 }))
276 .child(
277 v_flex().h_full().text_align(TextAlign::Right).child(
278 h_flex()
279 .justify_end()
280 .child(
281 IconButton::new("go-back", IconName::ChevronLeft)
282 .disabled(self.current_ix == 0 || self.runs.len() < 2)
283 .tooltip(ui::Tooltip::for_action_title(
284 "Go to previous run",
285 &EditPredictionContextGoBack,
286 ))
287 .on_click(cx.listener(|this, _, window, cx| {
288 this.handle_go_back(&EditPredictionContextGoBack, window, cx);
289 })),
290 )
291 .child(
292 div()
293 .child(format!("{}/{}", self.current_ix + 1, self.runs.len()))
294 .map(|this| {
295 if new_run_started {
296 this.with_animation(
297 "pulsating-count",
298 Animation::new(Duration::from_secs(2))
299 .repeat()
300 .with_easing(pulsating_between(0.4, 0.8)),
301 |label, delta| label.opacity(delta),
302 )
303 .into_any_element()
304 } else {
305 this.into_any_element()
306 }
307 }),
308 )
309 .child(
310 IconButton::new("go-forward", IconName::ChevronRight)
311 .disabled(self.current_ix + 1 == self.runs.len())
312 .tooltip(ui::Tooltip::for_action_title(
313 "Go to next run",
314 &EditPredictionContextGoBack,
315 ))
316 .on_click(cx.listener(|this, _, window, cx| {
317 this.handle_go_forward(
318 &EditPredictionContextGoForward,
319 window,
320 cx,
321 );
322 })),
323 ),
324 ),
325 )
326 }
327}
328
329impl Focusable for EditPredictionContextView {
330 fn focus_handle(&self, cx: &App) -> FocusHandle {
331 self.runs
332 .get(self.current_ix)
333 .map(|run| run.editor.read(cx).focus_handle(cx))
334 .unwrap_or_else(|| self.empty_focus_handle.clone())
335 }
336}
337
338impl EventEmitter<()> for EditPredictionContextView {}
339
340impl Item for EditPredictionContextView {
341 type Event = ();
342
343 fn tab_content_text(&self, _detail: usize, _cx: &App) -> SharedString {
344 "Edit Prediction Context".into()
345 }
346
347 fn buffer_kind(&self, _cx: &App) -> workspace::item::ItemBufferKind {
348 workspace::item::ItemBufferKind::Multibuffer
349 }
350
351 fn act_as_type<'a>(
352 &'a self,
353 type_id: TypeId,
354 self_handle: &'a Entity<Self>,
355 _: &'a App,
356 ) -> Option<gpui::AnyEntity> {
357 if type_id == TypeId::of::<Self>() {
358 Some(self_handle.clone().into())
359 } else if type_id == TypeId::of::<Editor>() {
360 Some(self.runs.get(self.current_ix)?.editor.clone().into())
361 } else {
362 None
363 }
364 }
365}
366
367impl gpui::Render for EditPredictionContextView {
368 fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl ui::IntoElement {
369 v_flex()
370 .key_context("EditPredictionContext")
371 .on_action(cx.listener(Self::handle_go_back))
372 .on_action(cx.listener(Self::handle_go_forward))
373 .size_full()
374 .map(|this| {
375 if self.runs.is_empty() {
376 this.child(
377 v_flex()
378 .size_full()
379 .justify_center()
380 .items_center()
381 .child("No retrieval runs yet"),
382 )
383 } else {
384 this.child(self.runs[self.current_ix].editor.clone())
385 .child(self.render_informational_footer(cx))
386 }
387 })
388 }
389}