1159 lines · 41.8 KB · Rust
RawDownload
use crate::context_store::ContextStore;use agent::HistoryStore;use collections::VecDeque;use editor::actions::Paste;use editor::display_map::EditorMargins;use editor::{    ContextMenuOptions, Editor, EditorElement, EditorEvent, EditorMode, EditorStyle, MultiBuffer,    actions::{MoveDown, MoveUp},};use fs::Fs;use gpui::{    AnyElement, App, ClipboardEntry, Context, CursorStyle, Entity, EventEmitter, FocusHandle,    Focusable, Subscription, TextStyle, WeakEntity, Window,};use language_model::{LanguageModel, LanguageModelRegistry};use parking_lot::Mutex;use prompt_store::PromptStore;use settings::Settings;use std::cmp;use std::rc::Rc;use std::sync::Arc;use theme::ThemeSettings;use ui::utils::WithRemSize;use ui::{IconButtonShape, KeyBinding, PopoverMenuHandle, Tooltip, prelude::*};use workspace::Workspace;use zed_actions::agent::ToggleModelSelector; use crate::agent_model_selector::AgentModelSelector;use crate::buffer_codegen::BufferCodegen;use crate::context_picker::{ContextPicker, ContextPickerCompletionProvider};use crate::context_strip::{ContextStrip, ContextStripEvent, SuggestContextKind};use crate::message_editor::{ContextCreasesAddon, extract_message_creases, insert_message_creases};use crate::terminal_codegen::TerminalCodegen;use crate::{CycleNextInlineAssist, CyclePreviousInlineAssist, ModelUsageContext};use crate::{RemoveAllContext, ToggleContextPicker}; 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<AgentModelSelector>,    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();         const RIGHT_PADDING: Pixels = px(9.);         let (left_gutter_width, right_padding) = match &self.mode {            PromptEditorMode::Buffer {                id: _,                codegen,                editor_margins,            } => {                let codegen = codegen.read(cx);                 if codegen.alternative_count(cx) > 1 {                    buttons.push(self.render_cycle_controls(codegen, cx));                }                 let editor_margins = editor_margins.lock();                let gutter = editor_margins.gutter;                 let left_gutter_width = gutter.full_width() + (gutter.margin / 2.0);                let right_padding = editor_margins.right + RIGHT_PADDING;                 (left_gutter_width, right_padding)            }            PromptEditorMode::Terminal { .. } => {                // Give the equivalent of the same left-padding that we're using on the right                (Pixels::from(40.0), Pixels::from(24.))            }        };         let bottom_padding = match &self.mode {            PromptEditorMode::Buffer { .. } => rems_from_px(2.0),            PromptEditorMode::Terminal { .. } => rems_from_px(8.0),        };         buttons.extend(self.render_buttons(window, cx));         v_flex()            .key_context("PromptEditor")            .capture_action(cx.listener(Self::paste))            .bg(cx.theme().colors().editor_background)            .block_mouse_except_scroll()            .gap_0p5()            .border_y_1()            .border_color(cx.theme().status().info_border)            .size_full()            .pt_0p5()            .pb(bottom_padding)            .pr(right_padding)            .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());                                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 existing_creases = self.editor.update(cx, extract_message_creases);         let focus = self.editor.focus_handle(cx).contains_focused(window, cx);        self.editor = cx.new(|cx| {            let mut editor = Editor::auto_height(1, Self::MAX_LINES as usize, window, cx);            editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);            editor.set_placeholder_text("Add a prompt…", window, cx);            editor.set_text(prompt, window, cx);            insert_message_creases(                &mut editor,                &existing_creases,                &self.context_store,                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 agent_panel_keybinding =            ui::text_for_action(&zed_actions::assistant::ToggleFocus, window, cx)                .map(|keybinding| format!("{keybinding} to chat"))                .unwrap_or_default();         format!("{action}… ({agent_panel_keybinding} ― ↓↑ for history — @ to include context)")    }     pub fn prompt(&self, cx: &App) -> String {        self.editor.read(cx).text(cx)    }     fn paste(&mut self, _: &Paste, _window: &mut Window, cx: &mut Context<Self>) {        let images = cx            .read_from_clipboard()            .map(|item| {                item.into_entries()                    .filter_map(|entry| {                        if let ClipboardEntry::Image(image) = entry {                            Some(image)                        } else {                            None                        }                    })                    .collect::<Vec<_>>()            })            .unwrap_or_default();         if images.is_empty() {            return;        }        cx.stop_propagation();         self.context_store.update(cx, |store, cx| {            for image in images {                store.add_image_instance(Arc::new(image), cx);            }        });    }     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().read(cx).is_via_remote_server();                         workspace                            .client()                            .telemetry()                            .log_edit_event("inline assist", is_via_ssh);                    });                }                let prompt = self.editor.read(cx).text(cx);                if self                    .prompt_history_ix                    .is_none_or(|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));        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 => {}            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 if self.context_strip.read(cx).has_context_items(cx) {            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",                            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",                                    cx,                                )                            })                            .on_click(cx.listener(|_, _, _, cx| {                                cx.emit(PromptEditorEvent::StartRequested);                            }))                            .into_any_element(),

Lines 1–500 of 1159.