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