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, ParentElement as _, SharedString, Styled as _, Task, TextAlign, Window, actions,
16 pulsating_between,
17};
18use multi_buffer::MultiBuffer;
19use project::Project;
20use text::OffsetRangeExt;
21use ui::{
22 ButtonCommon, Clickable, Color, Disableable, FluentBuilder as _, Icon, IconButton, IconName,
23 IconSize, InteractiveElement, IntoElement, ListHeader, ListItem, StyledTypography, div, h_flex,
24 v_flex,
25};
26use workspace::{Item, ItemHandle as _};
27use zeta2::{Zeta, ZetaContextRetrievalDebugInfo, ZetaDebugInfo, ZetaSearchQueryDebugInfo};
28
29pub struct Zeta2ContextView {
30 empty_focus_handle: FocusHandle,
31 project: Entity<Project>,
32 zeta: Entity<Zeta>,
33 runs: VecDeque<RetrievalRun>,
34 current_ix: usize,
35 _update_task: Task<Result<()>>,
36}
37
38#[derive(Debug)]
39struct RetrievalRun {
40 editor: Entity<Editor>,
41 search_queries: Vec<GlobQueries>,
42 started_at: Instant,
43 search_results_generated_at: Option<Instant>,
44 search_results_executed_at: Option<Instant>,
45 search_results_filtered_at: Option<Instant>,
46 finished_at: Option<Instant>,
47}
48
49#[derive(Debug)]
50struct GlobQueries {
51 glob: String,
52 alternations: Vec<String>,
53}
54
55actions!(
56 dev,
57 [
58 /// Go to the previous context retrieval run
59 Zeta2ContextGoBack,
60 /// Go to the next context retrieval run
61 Zeta2ContextGoForward
62 ]
63);
64
65impl Zeta2ContextView {
66 pub fn new(
67 project: Entity<Project>,
68 client: &Arc<Client>,
69 user_store: &Entity<UserStore>,
70 window: &mut gpui::Window,
71 cx: &mut Context<Self>,
72 ) -> Self {
73 let zeta = Zeta::global(client, user_store, cx);
74
75 let mut debug_rx = zeta.update(cx, |zeta, _| zeta.debug_info());
76 let _update_task = cx.spawn_in(window, async move |this, cx| {
77 while let Some(event) = debug_rx.next().await {
78 this.update_in(cx, |this, window, cx| {
79 this.handle_zeta_event(event, window, cx)
80 })?;
81 }
82 Ok(())
83 });
84
85 Self {
86 empty_focus_handle: cx.focus_handle(),
87 project,
88 runs: VecDeque::new(),
89 current_ix: 0,
90 zeta,
91 _update_task,
92 }
93 }
94
95 fn handle_zeta_event(
96 &mut self,
97 event: ZetaDebugInfo,
98 window: &mut gpui::Window,
99 cx: &mut Context<Self>,
100 ) {
101 match event {
102 ZetaDebugInfo::ContextRetrievalStarted(info) => {
103 if info.project == self.project {
104 self.handle_context_retrieval_started(info, window, cx);
105 }
106 }
107 ZetaDebugInfo::SearchQueriesGenerated(info) => {
108 if info.project == self.project {
109 self.handle_search_queries_generated(info, window, cx);
110 }
111 }
112 ZetaDebugInfo::SearchQueriesExecuted(info) => {
113 if info.project == self.project {
114 self.handle_search_queries_executed(info, window, cx);
115 }
116 }
117 ZetaDebugInfo::SearchResultsFiltered(info) => {
118 if info.project == self.project {
119 self.handle_search_results_filtered(info, window, cx);
120 }
121 }
122 ZetaDebugInfo::ContextRetrievalFinished(info) => {
123 if info.project == self.project {
124 self.handle_context_retrieval_finished(info, window, cx);
125 }
126 }
127 ZetaDebugInfo::EditPredicted(_) => {}
128 }
129 }
130
131 fn handle_context_retrieval_started(
132 &mut self,
133 info: ZetaContextRetrievalDebugInfo,
134 window: &mut Window,
135 cx: &mut Context<Self>,
136 ) {
137 if self
138 .runs
139 .back()
140 .is_some_and(|run| run.search_results_executed_at.is_none())
141 {
142 self.runs.pop_back();
143 }
144
145 let multibuffer = cx.new(|_| MultiBuffer::new(language::Capability::ReadOnly));
146 let editor = cx
147 .new(|cx| Editor::for_multibuffer(multibuffer, Some(self.project.clone()), window, cx));
148
149 if self.runs.len() == 32 {
150 self.runs.pop_front();
151 }
152
153 self.runs.push_back(RetrievalRun {
154 editor,
155 search_queries: Vec::new(),
156 started_at: info.timestamp,
157 search_results_generated_at: None,
158 search_results_executed_at: None,
159 search_results_filtered_at: None,
160 finished_at: None,
161 });
162
163 cx.notify();
164 }
165
166 fn handle_context_retrieval_finished(
167 &mut self,
168 info: ZetaContextRetrievalDebugInfo,
169 window: &mut Window,
170 cx: &mut Context<Self>,
171 ) {
172 let Some(run) = self.runs.back_mut() else {
173 return;
174 };
175
176 run.finished_at = Some(info.timestamp);
177
178 let multibuffer = run.editor.read(cx).buffer().clone();
179 multibuffer.update(cx, |multibuffer, cx| {
180 multibuffer.clear(cx);
181
182 let context = self.zeta.read(cx).context_for_project(&self.project);
183 let mut paths = Vec::new();
184 for (buffer, ranges) in context {
185 let path = PathKey::for_buffer(&buffer, cx);
186 let snapshot = buffer.read(cx).snapshot();
187 let ranges = ranges
188 .iter()
189 .map(|range| range.to_point(&snapshot))
190 .collect::<Vec<_>>();
191 paths.push((path, buffer, ranges));
192 }
193
194 for (path, buffer, ranges) in paths {
195 multibuffer.set_excerpts_for_path(path, buffer, ranges, 0, cx);
196 }
197 });
198
199 run.editor.update(cx, |editor, cx| {
200 editor.move_to_beginning(&Default::default(), window, cx);
201 });
202
203 cx.notify();
204 }
205
206 fn handle_search_queries_generated(
207 &mut self,
208 info: ZetaSearchQueryDebugInfo,
209 _window: &mut Window,
210 cx: &mut Context<Self>,
211 ) {
212 let Some(run) = self.runs.back_mut() else {
213 return;
214 };
215
216 run.search_results_generated_at = Some(info.timestamp);
217 run.search_queries = info
218 .queries
219 .into_iter()
220 .map(|query| {
221 let mut regex_parser = regex_syntax::ast::parse::Parser::new();
222
223 GlobQueries {
224 glob: query.glob,
225 alternations: match regex_parser.parse(&query.regex) {
226 Ok(regex_syntax::ast::Ast::Alternation(ref alt)) => {
227 alt.asts.iter().map(|ast| ast.to_string()).collect()
228 }
229 _ => vec![query.regex],
230 },
231 }
232 })
233 .collect();
234 cx.notify();
235 }
236
237 fn handle_search_queries_executed(
238 &mut self,
239 info: ZetaContextRetrievalDebugInfo,
240 _window: &mut Window,
241 cx: &mut Context<Self>,
242 ) {
243 if self.current_ix + 2 == self.runs.len() {
244 // Switch to latest when the queries are executed
245 self.current_ix += 1;
246 }
247
248 let Some(run) = self.runs.back_mut() else {
249 return;
250 };
251
252 run.search_results_executed_at = Some(info.timestamp);
253 cx.notify();
254 }
255
256 fn handle_search_results_filtered(
257 &mut self,
258 info: ZetaContextRetrievalDebugInfo,
259 _window: &mut Window,
260 cx: &mut Context<Self>,
261 ) {
262 let Some(run) = self.runs.back_mut() else {
263 return;
264 };
265
266 run.search_results_filtered_at = Some(info.timestamp);
267 cx.notify();
268 }
269
270 fn handle_go_back(
271 &mut self,
272 _: &Zeta2ContextGoBack,
273 window: &mut Window,
274 cx: &mut Context<Self>,
275 ) {
276 self.current_ix = self.current_ix.saturating_sub(1);
277 cx.focus_self(window);
278 cx.notify();
279 }
280
281 fn handle_go_forward(
282 &mut self,
283 _: &Zeta2ContextGoForward,
284 window: &mut Window,
285 cx: &mut Context<Self>,
286 ) {
287 self.current_ix = self
288 .current_ix
289 .add(1)
290 .min(self.runs.len().saturating_sub(1));
291 cx.focus_self(window);
292 cx.notify();
293 }
294
295 fn render_informational_footer(&self, cx: &mut Context<'_, Zeta2ContextView>) -> ui::Div {
296 let is_latest = self.runs.len() == self.current_ix + 1;
297 let run = &self.runs[self.current_ix];
298
299 h_flex()
300 .p_2()
301 .w_full()
302 .font_buffer(cx)
303 .text_xs()
304 .border_t_1()
305 .gap_2()
306 .child(
307 v_flex().h_full().flex_1().children(
308 run.search_queries
309 .iter()
310 .enumerate()
311 .flat_map(|(ix, query)| {
312 std::iter::once(ListHeader::new(query.glob.clone()).into_any_element())
313 .chain(query.alternations.iter().enumerate().map(
314 move |(alt_ix, alt)| {
315 ListItem::new(ix * 100 + alt_ix)
316 .start_slot(
317 Icon::new(IconName::MagnifyingGlass)
318 .color(Color::Muted)
319 .size(IconSize::Small),
320 )
321 .child(alt.clone())
322 .into_any_element()
323 },
324 ))
325 }),
326 ),
327 )
328 .child(
329 v_flex()
330 .h_full()
331 .text_align(TextAlign::Right)
332 .child(
333 h_flex()
334 .justify_end()
335 .child(
336 IconButton::new("go-back", IconName::ChevronLeft)
337 .disabled(self.current_ix == 0 || self.runs.len() < 2)
338 .tooltip(ui::Tooltip::for_action_title(
339 "Go to previous run",
340 &Zeta2ContextGoBack,
341 ))
342 .on_click(cx.listener(|this, _, window, cx| {
343 this.handle_go_back(&Zeta2ContextGoBack, window, cx);
344 })),
345 )
346 .child(
347 div()
348 .child(format!("{}/{}", self.current_ix + 1, self.runs.len()))
349 .map(|this| {
350 if self.runs.back().is_some_and(|back| {
351 back.search_results_executed_at.is_none()
352 }) {
353 this.with_animation(
354 "pulsating-count",
355 Animation::new(Duration::from_secs(2))
356 .repeat()
357 .with_easing(pulsating_between(0.4, 0.8)),
358 |label, delta| label.opacity(delta),
359 )
360 .into_any_element()
361 } else {
362 this.into_any_element()
363 }
364 }),
365 )
366 .child(
367 IconButton::new("go-forward", IconName::ChevronRight)
368 .disabled(self.current_ix + 1 == self.runs.len())
369 .tooltip(ui::Tooltip::for_action_title(
370 "Go to next run",
371 &Zeta2ContextGoBack,
372 ))
373 .on_click(cx.listener(|this, _, window, cx| {
374 this.handle_go_forward(&Zeta2ContextGoForward, window, cx);
375 })),
376 ),
377 )
378 .map(|mut div| {
379 let pending_message = |div: ui::Div, msg: &'static str| {
380 if is_latest {
381 return div.child(msg);
382 } else {
383 return div.child("Canceled");
384 }
385 };
386
387 let t0 = run.started_at;
388 let Some(t1) = run.search_results_generated_at else {
389 return pending_message(div, "Planning search...");
390 };
391 div = div.child(format!("Planned search: {:>5} ms", (t1 - t0).as_millis()));
392
393 let Some(t2) = run.search_results_executed_at else {
394 return pending_message(div, "Running search...");
395 };
396 div = div.child(format!("Ran search: {:>5} ms", (t2 - t1).as_millis()));
397
398 let Some(t3) = run.search_results_filtered_at else {
399 return pending_message(div, "Filtering results...");
400 };
401 div =
402 div.child(format!("Filtered results: {:>5} ms", (t3 - t2).as_millis()));
403
404 let Some(t4) = run.finished_at else {
405 return pending_message(div, "Building excerpts");
406 };
407 div = div
408 .child(format!("Build excerpts: {:>5} µs", (t4 - t3).as_micros()))
409 .child(format!("Total: {:>5} ms", (t4 - t0).as_millis()));
410 div
411 }),
412 )
413 }
414}
415
416impl Focusable for Zeta2ContextView {
417 fn focus_handle(&self, cx: &App) -> FocusHandle {
418 self.runs
419 .get(self.current_ix)
420 .map(|run| run.editor.read(cx).focus_handle(cx))
421 .unwrap_or_else(|| self.empty_focus_handle.clone())
422 }
423}
424
425impl EventEmitter<()> for Zeta2ContextView {}
426
427impl Item for Zeta2ContextView {
428 type Event = ();
429
430 fn tab_content_text(&self, _detail: usize, _cx: &App) -> SharedString {
431 "Edit Prediction Context".into()
432 }
433
434 fn buffer_kind(&self, _cx: &App) -> workspace::item::ItemBufferKind {
435 workspace::item::ItemBufferKind::Multibuffer
436 }
437
438 fn act_as_type<'a>(
439 &'a self,
440 type_id: TypeId,
441 self_handle: &'a Entity<Self>,
442 _: &'a App,
443 ) -> Option<gpui::AnyView> {
444 if type_id == TypeId::of::<Self>() {
445 Some(self_handle.to_any())
446 } else if type_id == TypeId::of::<Editor>() {
447 Some(self.runs.get(self.current_ix)?.editor.to_any())
448 } else {
449 None
450 }
451 }
452}
453
454impl gpui::Render for Zeta2ContextView {
455 fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl ui::IntoElement {
456 v_flex()
457 .key_context("Zeta2Context")
458 .on_action(cx.listener(Self::handle_go_back))
459 .on_action(cx.listener(Self::handle_go_forward))
460 .size_full()
461 .map(|this| {
462 if self.runs.is_empty() {
463 this.child(
464 v_flex()
465 .size_full()
466 .justify_center()
467 .items_center()
468 .child("No retrieval runs yet"),
469 )
470 } else {
471 this.child(self.runs[self.current_ix].editor.clone())
472 .child(self.render_informational_footer(cx))
473 }
474 })
475 }
476}