1207 lines · 44.8 KB · Rust
RawDownload
use crate::assistant_model_selector::AssistantModelSelector;use crate::buffer_codegen::BufferCodegen;use crate::context_picker::ContextPicker;use crate::context_store::ContextStore;use crate::context_strip::{ContextStrip, ContextStripEvent, SuggestContextKind};use crate::terminal_codegen::TerminalCodegen;use crate::thread_store::ThreadStore;use crate::{CycleNextInlineAssist, CyclePreviousInlineAssist};use crate::{RemoveAllContext, ToggleContextPicker};use client::ErrorExt;use collections::VecDeque;use editor::{    Editor, EditorElement, EditorEvent, EditorMode, EditorStyle, GutterDimensions, MultiBuffer,    actions::{MoveDown, MoveUp},};use feature_flags::{FeatureFlagAppExt as _, ZedPro};use fs::Fs;use gpui::{    AnyElement, App, ClickEvent, Context, CursorStyle, Entity, EventEmitter, FocusHandle,    Focusable, FontWeight, Subscription, TextStyle, WeakEntity, Window, anchored, deferred, point,};use language_model::{LanguageModel, LanguageModelRegistry};use language_model_selector::{ModelType, ToggleModelSelector};use parking_lot::Mutex;use settings::Settings;use std::cmp;use std::sync::Arc;use theme::ThemeSettings;use ui::utils::WithRemSize;use ui::{    CheckboxWithLabel, IconButtonShape, KeyBinding, Popover, PopoverMenuHandle, Tooltip, prelude::*,};use util::ResultExt;use workspace::Workspace; pub struct PromptEditor<T> {    pub editor: Entity<Editor>,    mode: PromptEditorMode,    context_store: Entity<ContextStore>,    context_strip: Entity<ContextStrip>,    context_picker_menu_handle: PopoverMenuHandle<ContextPicker>,    model_selector: Entity<AssistantModelSelector>,    edited_since_done: bool,    prompt_history: VecDeque<String>,    prompt_history_ix: Option<usize>,    pending_prompt: String,    _codegen_subscription: Subscription,    editor_subscriptions: Vec<Subscription>,    _context_strip_subscription: Subscription,    show_rate_limit_notice: bool,    _phantom: std::marker::PhantomData<T>,} impl<T: 'static> EventEmitter<PromptEditorEvent> for PromptEditor<T> {} impl<T: 'static> Render for PromptEditor<T> {    fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {        let ui_font_size = ThemeSettings::get_global(cx).ui_font_size(cx);        let mut buttons = Vec::new();         let left_gutter_width = match &self.mode {            PromptEditorMode::Buffer {                id: _,                codegen,                gutter_dimensions,            } => {                let codegen = codegen.read(cx);                 if codegen.alternative_count(cx) > 1 {                    buttons.push(self.render_cycle_controls(&codegen, cx));                }                 let gutter_dimensions = gutter_dimensions.lock();                 gutter_dimensions.full_width() + (gutter_dimensions.margin / 2.0)            }            PromptEditorMode::Terminal { .. } => {                // Give the equivalent of the same left-padding that we're using on the right                Pixels::from(40.0)            }        };         let bottom_padding = match &self.mode {            PromptEditorMode::Buffer { .. } => Pixels::from(0.),            PromptEditorMode::Terminal { .. } => Pixels::from(8.0),        };         buttons.extend(self.render_buttons(window, cx));         v_flex()            .key_context("PromptEditor")            .bg(cx.theme().colors().editor_background)            .block_mouse_down()            .gap_0p5()            .border_y_1()            .border_color(cx.theme().status().info_border)            .size_full()            .pt_0p5()            .pb(bottom_padding)            .pr_6()            .child(                h_flex()                    .items_start()                    .cursor(CursorStyle::Arrow)                    .on_action(cx.listener(Self::toggle_context_picker))                    .on_action(cx.listener(|this, _: &ToggleModelSelector, window, cx| {                        this.model_selector                            .update(cx, |model_selector, cx| model_selector.toggle(window, cx));                    }))                    .on_action(cx.listener(Self::confirm))                    .on_action(cx.listener(Self::cancel))                    .on_action(cx.listener(Self::move_up))                    .on_action(cx.listener(Self::move_down))                    .on_action(cx.listener(Self::remove_all_context))                    .capture_action(cx.listener(Self::cycle_prev))                    .capture_action(cx.listener(Self::cycle_next))                    .child(                        WithRemSize::new(ui_font_size)                            .flex()                            .flex_row()                            .flex_shrink_0()                            .items_center()                            .h_full()                            .w(left_gutter_width)                            .justify_center()                            .gap_2()                            .child(self.render_close_button(cx))                            .map(|el| {                                let CodegenStatus::Error(error) = self.codegen_status(cx) else {                                    return el;                                };                                 let error_message = SharedString::from(error.to_string());                                if error.error_code() == proto::ErrorCode::RateLimitExceeded                                    && cx.has_flag::<ZedPro>()                                {                                    el.child(                                        v_flex()                                            .child(                                                IconButton::new(                                                    "rate-limit-error",                                                    IconName::XCircle,                                                )                                                .toggle_state(self.show_rate_limit_notice)                                                .shape(IconButtonShape::Square)                                                .icon_size(IconSize::Small)                                                .on_click(                                                    cx.listener(Self::toggle_rate_limit_notice),                                                ),                                            )                                            .children(self.show_rate_limit_notice.then(|| {                                                deferred(                                                    anchored()                                                        .position_mode(                                                            gpui::AnchoredPositionMode::Local,                                                        )                                                        .position(point(px(0.), px(24.)))                                                        .anchor(gpui::Corner::TopLeft)                                                        .child(self.render_rate_limit_notice(cx)),                                                )                                            })),                                    )                                } else {                                    el.child(                                        div()                                            .id("error")                                            .tooltip(Tooltip::text(error_message))                                            .child(                                                Icon::new(IconName::XCircle)                                                    .size(IconSize::Small)                                                    .color(Color::Error),                                            ),                                    )                                }                            }),                    )                    .child(                        h_flex()                            .w_full()                            .justify_between()                            .child(div().flex_1().child(self.render_editor(window, cx)))                            .child(                                WithRemSize::new(ui_font_size)                                    .flex()                                    .flex_row()                                    .items_center()                                    .gap_1()                                    .children(buttons),                            ),                    ),            )            .child(                WithRemSize::new(ui_font_size)                    .flex()                    .flex_row()                    .items_center()                    .child(h_flex().flex_shrink_0().w(left_gutter_width))                    .child(                        h_flex()                            .w_full()                            .pl_1()                            .items_start()                            .justify_between()                            .child(self.context_strip.clone())                            .child(self.model_selector.clone()),                    ),            )    }} impl<T: 'static> Focusable for PromptEditor<T> {    fn focus_handle(&self, cx: &App) -> FocusHandle {        self.editor.focus_handle(cx)    }} impl<T: 'static> PromptEditor<T> {    const MAX_LINES: u8 = 8;     fn codegen_status<'a>(&'a self, cx: &'a App) -> &'a CodegenStatus {        match &self.mode {            PromptEditorMode::Buffer { codegen, .. } => codegen.read(cx).status(cx),            PromptEditorMode::Terminal { codegen, .. } => &codegen.read(cx).status,        }    }     fn subscribe_to_editor(&mut self, window: &mut Window, cx: &mut Context<Self>) {        self.editor_subscriptions.clear();        self.editor_subscriptions.push(cx.subscribe_in(            &self.editor,            window,            Self::handle_prompt_editor_events,        ));    }     pub fn set_show_cursor_when_unfocused(        &mut self,        show_cursor_when_unfocused: bool,        cx: &mut Context<Self>,    ) {        self.editor.update(cx, |editor, cx| {            editor.set_show_cursor_when_unfocused(show_cursor_when_unfocused, cx)        });    }     pub fn unlink(&mut self, window: &mut Window, cx: &mut Context<Self>) {        let prompt = self.prompt(cx);        let focus = self.editor.focus_handle(cx).contains_focused(window, cx);        self.editor = cx.new(|cx| {            let mut editor = Editor::auto_height(Self::MAX_LINES as usize, window, cx);            editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);            editor.set_placeholder_text(Self::placeholder_text(&self.mode, window, cx), cx);            editor.set_placeholder_text("Add a prompt…", cx);            editor.set_text(prompt, window, cx);            if focus {                window.focus(&editor.focus_handle(cx));            }            editor        });        self.subscribe_to_editor(window, cx);    }     pub fn placeholder_text(mode: &PromptEditorMode, window: &mut Window, cx: &mut App) -> String {        let action = match mode {            PromptEditorMode::Buffer { codegen, .. } => {                if codegen.read(cx).is_insertion {                    "Generate"                } else {                    "Transform"                }            }            PromptEditorMode::Terminal { .. } => "Generate",        };         let assistant_panel_keybinding =            ui::text_for_action(&zed_actions::assistant::ToggleFocus, window, cx)                .map(|keybinding| format!("{keybinding} to chat ― "))                .unwrap_or_default();         format!("{action}… ({assistant_panel_keybinding}↓↑ for history)")    }     pub fn prompt(&self, cx: &App) -> String {        self.editor.read(cx).text(cx)    }     fn toggle_rate_limit_notice(        &mut self,        _: &ClickEvent,        window: &mut Window,        cx: &mut Context<Self>,    ) {        self.show_rate_limit_notice = !self.show_rate_limit_notice;        if self.show_rate_limit_notice {            window.focus(&self.editor.focus_handle(cx));        }        cx.notify();    }     fn handle_prompt_editor_events(        &mut self,        _: &Entity<Editor>,        event: &EditorEvent,        window: &mut Window,        cx: &mut Context<Self>,    ) {        match event {            EditorEvent::Edited { .. } => {                if let Some(workspace) = window.root::<Workspace>().flatten() {                    workspace.update(cx, |workspace, cx| {                        let is_via_ssh = workspace                            .project()                            .update(cx, |project, _| project.is_via_ssh());                         workspace                            .client()                            .telemetry()                            .log_edit_event("inline assist", is_via_ssh);                    });                }                let prompt = self.editor.read(cx).text(cx);                if self                    .prompt_history_ix                    .map_or(true, |ix| self.prompt_history[ix] != prompt)                {                    self.prompt_history_ix.take();                    self.pending_prompt = prompt;                }                 self.edited_since_done = true;                cx.notify();            }            EditorEvent::Blurred => {                if self.show_rate_limit_notice {                    self.show_rate_limit_notice = false;                    cx.notify();                }            }            _ => {}        }    }     fn toggle_context_picker(        &mut self,        _: &ToggleContextPicker,        window: &mut Window,        cx: &mut Context<Self>,    ) {        self.context_picker_menu_handle.toggle(window, cx);    }     pub fn remove_all_context(        &mut self,        _: &RemoveAllContext,        _window: &mut Window,        cx: &mut Context<Self>,    ) {        self.context_store.update(cx, |store, _cx| store.clear());        cx.notify();    }     fn cancel(        &mut self,        _: &editor::actions::Cancel,        _window: &mut Window,        cx: &mut Context<Self>,    ) {        match self.codegen_status(cx) {            CodegenStatus::Idle | CodegenStatus::Done | CodegenStatus::Error(_) => {                cx.emit(PromptEditorEvent::CancelRequested);            }            CodegenStatus::Pending => {                cx.emit(PromptEditorEvent::StopRequested);            }        }    }     fn confirm(&mut self, _: &menu::Confirm, _window: &mut Window, cx: &mut Context<Self>) {        match self.codegen_status(cx) {            CodegenStatus::Idle => {                cx.emit(PromptEditorEvent::StartRequested);            }            CodegenStatus::Pending => {                cx.emit(PromptEditorEvent::DismissRequested);            }            CodegenStatus::Done => {                if self.edited_since_done {                    cx.emit(PromptEditorEvent::StartRequested);                } else {                    cx.emit(PromptEditorEvent::ConfirmRequested { execute: false });                }            }            CodegenStatus::Error(_) => {                cx.emit(PromptEditorEvent::StartRequested);            }        }    }     fn move_up(&mut self, _: &MoveUp, window: &mut Window, cx: &mut Context<Self>) {        if let Some(ix) = self.prompt_history_ix {            if ix > 0 {                self.prompt_history_ix = Some(ix - 1);                let prompt = self.prompt_history[ix - 1].as_str();                self.editor.update(cx, |editor, cx| {                    editor.set_text(prompt, window, cx);                    editor.move_to_beginning(&Default::default(), window, cx);                });            }        } else if !self.prompt_history.is_empty() {            self.prompt_history_ix = Some(self.prompt_history.len() - 1);            let prompt = self.prompt_history[self.prompt_history.len() - 1].as_str();            self.editor.update(cx, |editor, cx| {                editor.set_text(prompt, window, cx);                editor.move_to_beginning(&Default::default(), window, cx);            });        }    }     fn move_down(&mut self, _: &MoveDown, window: &mut Window, cx: &mut Context<Self>) {        if let Some(ix) = self.prompt_history_ix {            if ix < self.prompt_history.len() - 1 {                self.prompt_history_ix = Some(ix + 1);                let prompt = self.prompt_history[ix + 1].as_str();                self.editor.update(cx, |editor, cx| {                    editor.set_text(prompt, window, cx);                    editor.move_to_end(&Default::default(), window, cx)                });            } else {                self.prompt_history_ix = None;                let prompt = self.pending_prompt.as_str();                self.editor.update(cx, |editor, cx| {                    editor.set_text(prompt, window, cx);                    editor.move_to_end(&Default::default(), window, cx)                });            }        } else {            self.context_strip.focus_handle(cx).focus(window);        }    }     fn render_buttons(&self, _window: &mut Window, cx: &mut Context<Self>) -> Vec<AnyElement> {        let mode = match &self.mode {            PromptEditorMode::Buffer { codegen, .. } => {                let codegen = codegen.read(cx);                if codegen.is_insertion {                    GenerationMode::Generate                } else {                    GenerationMode::Transform                }            }            PromptEditorMode::Terminal { .. } => GenerationMode::Generate,        };         let codegen_status = self.codegen_status(cx);         match codegen_status {            CodegenStatus::Idle => {                vec![                    Button::new("start", mode.start_label())                        .label_size(LabelSize::Small)                        .icon(IconName::Return)                        .icon_size(IconSize::XSmall)                        .icon_color(Color::Muted)                        .on_click(                            cx.listener(|_, _, _, cx| cx.emit(PromptEditorEvent::StartRequested)),                        )                        .into_any_element(),                ]            }            CodegenStatus::Pending => vec![                IconButton::new("stop", IconName::Stop)                    .icon_color(Color::Error)                    .shape(IconButtonShape::Square)                    .tooltip(move |window, cx| {                        Tooltip::with_meta(                            mode.tooltip_interrupt(),                            Some(&menu::Cancel),                            "Changes won't be discarded",                            window,                            cx,                        )                    })                    .on_click(cx.listener(|_, _, _, cx| cx.emit(PromptEditorEvent::StopRequested)))                    .into_any_element(),            ],            CodegenStatus::Done | CodegenStatus::Error(_) => {                let has_error = matches!(codegen_status, CodegenStatus::Error(_));                if has_error || self.edited_since_done {                    vec![                        IconButton::new("restart", IconName::RotateCw)                            .icon_color(Color::Info)                            .shape(IconButtonShape::Square)                            .tooltip(move |window, cx| {                                Tooltip::with_meta(                                    mode.tooltip_restart(),                                    Some(&menu::Confirm),                                    "Changes will be discarded",                                    window,                                    cx,                                )

Lines 1–500 of 1207.