tenant.openagents/omega

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::{
12    Editor, PathKey,
13    display_map::{BlockPlacement, BlockProperties, BlockStyle},
14};
15use futures::StreamExt as _;
16use gpui::{
17    Animation, AnimationExt, App, AppContext as _, Context, Entity, EventEmitter, FocusHandle,
18    Focusable, InteractiveElement as _, IntoElement as _, ParentElement as _, SharedString,
19    Styled as _, Task, TextAlign, Window, actions, div, pulsating_between,
20};
21use multi_buffer::{Anchor, MultiBuffer};
22use project::Project;
23use text::Point;
24use ui::{
25    ButtonCommon, Clickable, Disableable, FluentBuilder as _, IconButton, IconName,
26    StyledTypography as _, h_flex, v_flex,
27};
28
29use edit_prediction::{
30    ContextRetrievalFinishedDebugEvent, ContextRetrievalStartedDebugEvent, DebugEvent,
31    EditPredictionStore,
32};
33use workspace::Item;
34
35pub struct EditPredictionContextView {
36    empty_focus_handle: FocusHandle,
37    project: Entity<Project>,
38    store: Entity<EditPredictionStore>,
39    runs: VecDeque<RetrievalRun>,
40    current_ix: usize,
41    _update_task: Task<Result<()>>,
42}
43
44#[derive(Debug)]
45struct RetrievalRun {
46    editor: Entity<Editor>,
47    started_at: Instant,
48    metadata: Vec<(&'static str, SharedString)>,
49    finished_at: Option<Instant>,
50}
51
52actions!(
53    dev,
54    [
55        /// Go to the previous context retrieval run
56        EditPredictionContextGoBack,
57        /// Go to the next context retrieval run
58        EditPredictionContextGoForward
59    ]
60);
61
62impl EditPredictionContextView {
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 store = EditPredictionStore::global(client, user_store, cx);
71
72        let mut debug_rx = store.update(cx, |store, cx| store.debug_info(&project, cx));
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_store_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            store,
88            _update_task,
89        }
90    }
91
92    fn handle_store_event(
93        &mut self,
94        event: DebugEvent,
95        window: &mut gpui::Window,
96        cx: &mut Context<Self>,
97    ) {
98        match event {
99            DebugEvent::ContextRetrievalStarted(info) => {
100                if info.project_entity_id == self.project.entity_id() {
101                    self.handle_context_retrieval_started(info, window, cx);
102                }
103            }
104            DebugEvent::ContextRetrievalFinished(info) => {
105                if info.project_entity_id == self.project.entity_id() {
106                    self.handle_context_retrieval_finished(info, window, cx);
107                }
108            }
109            DebugEvent::EditPredictionStarted(_) => {}
110            DebugEvent::EditPredictionFinished(_) => {}
111        }
112    }
113
114    fn handle_context_retrieval_started(
115        &mut self,
116        info: ContextRetrievalStartedDebugEvent,
117        window: &mut Window,
118        cx: &mut Context<Self>,
119    ) {
120        if self
121            .runs
122            .back()
123            .is_some_and(|run| run.finished_at.is_none())
124        {
125            self.runs.pop_back();
126        }
127
128        let multibuffer = cx.new(|_| MultiBuffer::new(language::Capability::ReadOnly));
129        let editor = cx
130            .new(|cx| Editor::for_multibuffer(multibuffer, Some(self.project.clone()), window, cx));
131
132        if self.runs.len() == 32 {
133            self.runs.pop_front();
134        }
135
136        self.runs.push_back(RetrievalRun {
137            editor,
138            started_at: info.timestamp,
139            finished_at: None,
140            metadata: Vec::new(),
141        });
142
143        cx.notify();
144    }
145
146    fn handle_context_retrieval_finished(
147        &mut self,
148        info: ContextRetrievalFinishedDebugEvent,
149        window: &mut Window,
150        cx: &mut Context<Self>,
151    ) {
152        let Some(run) = self.runs.back_mut() else {
153            return;
154        };
155
156        run.finished_at = Some(info.timestamp);
157        run.metadata = info.metadata;
158
159        let related_files = self.store.update(cx, |store, cx| {
160            store.context_for_project_with_buffers(&self.project, cx)
161        });
162
163        let editor = run.editor.clone();
164        let multibuffer = run.editor.read(cx).buffer().clone();
165
166        if self.current_ix + 2 == self.runs.len() {
167            self.current_ix += 1;
168        }
169
170        cx.spawn_in(window, async move |this, cx| {
171            let mut paths: Vec<(PathKey, _, Vec<_>, Vec<usize>, usize)> = Vec::new();
172            for (related_file, buffer) in related_files {
173                let orders = related_file
174                    .excerpts
175                    .iter()
176                    .map(|excerpt| excerpt.order)
177                    .collect::<Vec<_>>();
178                let min_order = orders.iter().copied().min().unwrap_or(usize::MAX);
179                let point_ranges = related_file
180                    .excerpts
181                    .iter()
182                    .map(|excerpt| {
183                        Point::new(excerpt.row_range.start, 0)..Point::new(excerpt.row_range.end, 0)
184                    })
185                    .collect::<Vec<_>>();
186                cx.update(|_, cx| {
187                    let path = if let Some(file) = buffer.read(cx).file() {
188                        PathKey::with_sort_prefix(min_order as u64, file.path().clone())
189                    } else {
190                        PathKey::for_buffer(&buffer, cx)
191                    };
192                    paths.push((path, buffer, point_ranges, orders, min_order));
193                })?;
194            }
195
196            paths.sort_by_key(|(_, _, _, _, min_order)| *min_order);
197
198            let mut excerpt_anchors_with_orders: Vec<(Anchor, usize)> = Vec::new();
199
200            multibuffer.update(cx, |multibuffer, cx| {
201                multibuffer.clear(cx);
202
203                for (path, buffer, ranges, orders, _) in paths {
204                    multibuffer.set_excerpts_for_path(path, buffer.clone(), ranges.clone(), 0, cx);
205                    let snapshot = multibuffer.snapshot(cx);
206                    let buffer_snapshot = buffer.read(cx).snapshot();
207                    for (range, order) in ranges.into_iter().zip(orders) {
208                        let text_anchor = buffer_snapshot.anchor_range_inside(range);
209                        if let Some(start) = snapshot.anchor_in_buffer(text_anchor.start) {
210                            excerpt_anchors_with_orders.push((start, order));
211                        }
212                    }
213                }
214            });
215
216            editor.update_in(cx, |editor, window, cx| {
217                let blocks = excerpt_anchors_with_orders
218                    .into_iter()
219                    .map(|(anchor, order)| {
220                        let label = SharedString::from(format!("order: {order}"));
221                        BlockProperties {
222                            placement: BlockPlacement::Above(anchor),
223                            height: Some(1),
224                            style: BlockStyle::Sticky,
225                            render: Arc::new(move |cx| {
226                                div()
227                                    .pl(cx.anchor_x)
228                                    .text_ui_xs(cx)
229                                    .text_color(cx.editor_style.status.info)
230                                    .child(label.clone())
231                                    .into_any_element()
232                            }),
233                            priority: 0,
234                        }
235                    })
236                    .collect::<Vec<_>>();
237                editor.insert_blocks(blocks, None, cx);
238                editor.move_to_beginning(&Default::default(), window, cx);
239            })?;
240
241            this.update(cx, |_, cx| cx.notify())
242        })
243        .detach();
244    }
245
246    fn handle_go_back(
247        &mut self,
248        _: &EditPredictionContextGoBack,
249        window: &mut Window,
250        cx: &mut Context<Self>,
251    ) {
252        self.current_ix = self.current_ix.saturating_sub(1);
253        cx.focus_self(window);
254        cx.notify();
255    }
256
257    fn handle_go_forward(
258        &mut self,
259        _: &EditPredictionContextGoForward,
260        window: &mut Window,
261        cx: &mut Context<Self>,
262    ) {
263        self.current_ix = self
264            .current_ix
265            .add(1)
266            .min(self.runs.len().saturating_sub(1));
267        cx.focus_self(window);
268        cx.notify();
269    }
270
271    fn render_informational_footer(
272        &self,
273        cx: &mut Context<'_, EditPredictionContextView>,
274    ) -> ui::Div {
275        let run = &self.runs[self.current_ix];
276        let new_run_started = self
277            .runs
278            .back()
279            .map_or(false, |latest_run| latest_run.finished_at.is_none());
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(v_flex().h_full().flex_1().child({
289                let t0 = run.started_at;
290                let mut table = ui::Table::new(2).width(ui::px(300.)).no_ui_font();
291                for (key, value) in &run.metadata {
292                    table = table.row(vec![
293                        key.into_any_element(),
294                        value.clone().into_any_element(),
295                    ])
296                }
297                table = table.row(vec![
298                    "Total Time".into_any_element(),
299                    format!("{} ms", (run.finished_at.unwrap_or(t0) - t0).as_millis())
300                        .into_any_element(),
301                ]);
302                table
303            }))
304            .child(
305                v_flex().h_full().text_align(TextAlign::Right).child(
306                    h_flex()
307                        .justify_end()
308                        .child(
309                            IconButton::new("go-back", IconName::ChevronLeft)
310                                .disabled(self.current_ix == 0 || self.runs.len() < 2)
311                                .tooltip(ui::Tooltip::for_action_title(
312                                    "Go to previous run",
313                                    &EditPredictionContextGoBack,
314                                ))
315                                .on_click(cx.listener(|this, _, window, cx| {
316                                    this.handle_go_back(&EditPredictionContextGoBack, window, cx);
317                                })),
318                        )
319                        .child(
320                            div()
321                                .child(format!("{}/{}", self.current_ix + 1, self.runs.len()))
322                                .map(|this| {
323                                    if new_run_started {
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                                    &EditPredictionContextGoBack,
343                                ))
344                                .on_click(cx.listener(|this, _, window, cx| {
345                                    this.handle_go_forward(
346                                        &EditPredictionContextGoForward,
347                                        window,
348                                        cx,
349                                    );
350                                })),
351                        ),
352                ),
353            )
354    }
355}
356
357impl Focusable for EditPredictionContextView {
358    fn focus_handle(&self, cx: &App) -> FocusHandle {
359        self.runs
360            .get(self.current_ix)
361            .map(|run| run.editor.read(cx).focus_handle(cx))
362            .unwrap_or_else(|| self.empty_focus_handle.clone())
363    }
364}
365
366impl EventEmitter<()> for EditPredictionContextView {}
367
368impl Item for EditPredictionContextView {
369    type Event = ();
370
371    fn tab_content_text(&self, _detail: usize, _cx: &App) -> SharedString {
372        "Edit Prediction Context".into()
373    }
374
375    fn buffer_kind(&self, _cx: &App) -> workspace::item::ItemBufferKind {
376        workspace::item::ItemBufferKind::Multibuffer
377    }
378
379    fn act_as_type<'a>(
380        &'a self,
381        type_id: TypeId,
382        self_handle: &'a Entity<Self>,
383        _: &'a App,
384    ) -> Option<gpui::AnyEntity> {
385        if type_id == TypeId::of::<Self>() {
386            Some(self_handle.clone().into())
387        } else if type_id == TypeId::of::<Editor>() {
388            Some(self.runs.get(self.current_ix)?.editor.clone().into())
389        } else {
390            None
391        }
392    }
393}
394
395impl gpui::Render for EditPredictionContextView {
396    fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl ui::IntoElement {
397        v_flex()
398            .key_context("EditPredictionContext")
399            .on_action(cx.listener(Self::handle_go_back))
400            .on_action(cx.listener(Self::handle_go_forward))
401            .size_full()
402            .map(|this| {
403                if self.runs.is_empty() {
404                    this.child(
405                        v_flex()
406                            .size_full()
407                            .justify_center()
408                            .items_center()
409                            .child("No retrieval runs yet"),
410                    )
411                } else {
412                    this.child(self.runs[self.current_ix].editor.clone())
413                        .child(self.render_informational_footer(cx))
414                }
415            })
416    }
417}
418