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