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