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