1604 lines · 65.9 KB · Rust
RawDownload
use std::collections::BTreeMap;use std::rc::Rc;use std::sync::Arc; use crate::agent_model_selector::AgentModelSelector;use crate::context::{AgentContextKey, ContextCreasesAddon, ContextLoadResult, load_context};use crate::tool_compatibility::{IncompatibleToolsState, IncompatibleToolsTooltip};use crate::ui::{    MaxModeTooltip,    preview::{AgentPreview, UsageCallout},};use agent_settings::{AgentSettings, CompletionMode};use assistant_context_editor::language_model_selector::ToggleModelSelector;use buffer_diff::BufferDiff;use client::UserStore;use collections::{HashMap, HashSet};use editor::actions::{MoveUp, Paste};use editor::{    AnchorRangeExt, ContextMenuOptions, ContextMenuPlacement, Editor, EditorElement, EditorEvent,    EditorMode, EditorStyle, MultiBuffer,};use file_icons::FileIcons;use fs::Fs;use futures::future::Shared;use futures::{FutureExt as _, future};use gpui::{    Animation, AnimationExt, App, Entity, EventEmitter, Focusable, Subscription, Task, TextStyle,    WeakEntity, linear_color_stop, linear_gradient, point, pulsating_between,};use language::{Buffer, Language, Point};use language_model::{    ConfiguredModel, LanguageModelRequestMessage, MessageContent, RequestUsage,    ZED_CLOUD_PROVIDER_ID,};use multi_buffer;use project::Project;use prompt_store::PromptStore;use proto::Plan;use settings::Settings;use std::time::Duration;use theme::ThemeSettings;use ui::{Disclosure, KeyBinding, PopoverMenuHandle, Tooltip, prelude::*};use util::{ResultExt as _, maybe};use workspace::{CollaboratorId, Workspace};use zed_llm_client::CompletionIntent; use crate::context_picker::{ContextPicker, ContextPickerCompletionProvider, crease_for_mention};use crate::context_store::ContextStore;use crate::context_strip::{ContextStrip, ContextStripEvent, SuggestContextKind};use crate::profile_selector::ProfileSelector;use crate::thread::{MessageCrease, Thread, TokenUsageRatio};use crate::thread_store::{TextThreadStore, ThreadStore};use crate::{    ActiveThread, AgentDiffPane, Chat, ChatWithFollow, ExpandMessageEditor, Follow, KeepAll,    ModelUsageContext, NewThread, OpenAgentDiff, RejectAll, RemoveAllContext, ToggleBurnMode,    ToggleContextPicker, ToggleProfileSelector, register_agent_preview,}; #[derive(RegisterComponent)]pub struct MessageEditor {    thread: Entity<Thread>,    incompatible_tools_state: Entity<IncompatibleToolsState>,    editor: Entity<Editor>,    workspace: WeakEntity<Workspace>,    project: Entity<Project>,    user_store: Entity<UserStore>,    context_store: Entity<ContextStore>,    prompt_store: Option<Entity<PromptStore>>,    context_strip: Entity<ContextStrip>,    context_picker_menu_handle: PopoverMenuHandle<ContextPicker>,    model_selector: Entity<AgentModelSelector>,    last_loaded_context: Option<ContextLoadResult>,    load_context_task: Option<Shared<Task<()>>>,    profile_selector: Entity<ProfileSelector>,    edits_expanded: bool,    editor_is_expanded: bool,    last_estimated_token_count: Option<usize>,    update_token_count_task: Option<Task<()>>,    _subscriptions: Vec<Subscription>,} const MAX_EDITOR_LINES: usize = 8; pub(crate) fn create_editor(    workspace: WeakEntity<Workspace>,    context_store: WeakEntity<ContextStore>,    thread_store: WeakEntity<ThreadStore>,    text_thread_store: WeakEntity<TextThreadStore>,    window: &mut Window,    cx: &mut App,) -> Entity<Editor> {    let language = Language::new(        language::LanguageConfig {            completion_query_characters: HashSet::from_iter(['.', '-', '_', '@']),            ..Default::default()        },        None,    );     let editor = cx.new(|cx| {        let buffer = cx.new(|cx| Buffer::local("", cx).with_language(Arc::new(language), cx));        let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));        let mut editor = Editor::new(            editor::EditorMode::AutoHeight {                max_lines: MAX_EDITOR_LINES,            },            buffer,            None,            window,            cx,        );        editor.set_placeholder_text("Message the agent – @ to include context", cx);        editor.set_show_indent_guides(false, cx);        editor.set_soft_wrap();        editor.set_use_modal_editing(true);        editor.set_context_menu_options(ContextMenuOptions {            min_entries_visible: 12,            max_entries_visible: 12,            placement: Some(ContextMenuPlacement::Above),        });        editor.register_addon(ContextCreasesAddon::new());        editor    });     let editor_entity = editor.downgrade();    editor.update(cx, |editor, _| {        editor.set_completion_provider(Some(Rc::new(ContextPickerCompletionProvider::new(            workspace,            context_store,            Some(thread_store),            Some(text_thread_store),            editor_entity,            None,        ))));    });    editor} impl MessageEditor {    pub fn new(        fs: Arc<dyn Fs>,        workspace: WeakEntity<Workspace>,        user_store: Entity<UserStore>,        context_store: Entity<ContextStore>,        prompt_store: Option<Entity<PromptStore>>,        thread_store: WeakEntity<ThreadStore>,        text_thread_store: WeakEntity<TextThreadStore>,        thread: Entity<Thread>,        window: &mut Window,        cx: &mut Context<Self>,    ) -> Self {        let context_picker_menu_handle = PopoverMenuHandle::default();        let model_selector_menu_handle = PopoverMenuHandle::default();         let editor = create_editor(            workspace.clone(),            context_store.downgrade(),            thread_store.clone(),            text_thread_store.clone(),            window,            cx,        );         let context_strip = cx.new(|cx| {            ContextStrip::new(                context_store.clone(),                workspace.clone(),                Some(thread_store.clone()),                Some(text_thread_store.clone()),                context_picker_menu_handle.clone(),                SuggestContextKind::File,                ModelUsageContext::Thread(thread.clone()),                window,                cx,            )        });         let incompatible_tools =            cx.new(|cx| IncompatibleToolsState::new(thread.read(cx).tools().clone(), cx));         let subscriptions = vec![            cx.subscribe_in(&context_strip, window, Self::handle_context_strip_event),            cx.subscribe(&editor, |this, _, event, cx| match event {                EditorEvent::BufferEdited => this.handle_message_changed(cx),                _ => {}            }),            cx.observe(&context_store, |this, _, cx| {                // When context changes, reload it for token counting.                let _ = this.reload_context(cx);            }),            cx.observe(&thread.read(cx).action_log().clone(), |_, _, cx| {                cx.notify()            }),        ];         let model_selector = cx.new(|cx| {            AgentModelSelector::new(                fs.clone(),                model_selector_menu_handle,                editor.focus_handle(cx),                ModelUsageContext::Thread(thread.clone()),                window,                cx,            )        });         let profile_selector = cx.new(|cx| {            ProfileSelector::new(                fs,                thread.clone(),                thread_store,                editor.focus_handle(cx),                cx,            )        });         Self {            editor: editor.clone(),            project: thread.read(cx).project().clone(),            user_store,            thread,            incompatible_tools_state: incompatible_tools.clone(),            workspace,            context_store,            prompt_store,            context_strip,            context_picker_menu_handle,            load_context_task: None,            last_loaded_context: None,            model_selector,            edits_expanded: false,            editor_is_expanded: false,            profile_selector,            last_estimated_token_count: None,            update_token_count_task: None,            _subscriptions: subscriptions,        }    }     pub fn context_store(&self) -> &Entity<ContextStore> {        &self.context_store    }     pub fn expand_message_editor(        &mut self,        _: &ExpandMessageEditor,        _window: &mut Window,        cx: &mut Context<Self>,    ) {        self.set_editor_is_expanded(!self.editor_is_expanded, cx);    }     fn set_editor_is_expanded(&mut self, is_expanded: bool, cx: &mut Context<Self>) {        self.editor_is_expanded = is_expanded;        self.editor.update(cx, |editor, _| {            if self.editor_is_expanded {                editor.set_mode(EditorMode::Full {                    scale_ui_elements_with_buffer_font_size: false,                    show_active_line_background: false,                    sized_by_content: false,                })            } else {                editor.set_mode(EditorMode::AutoHeight {                    max_lines: MAX_EDITOR_LINES,                })            }        });        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 chat(&mut self, _: &Chat, window: &mut Window, cx: &mut Context<Self>) {        if self.is_editor_empty(cx) {            return;        }         self.thread.update(cx, |thread, cx| {            thread.cancel_editing(cx);        });         if self.thread.read(cx).is_generating() {            self.stop_current_and_send_new_message(window, cx);            return;        }         self.set_editor_is_expanded(false, cx);        self.send_to_model(window, cx);         cx.notify();    }     fn chat_with_follow(        &mut self,        _: &ChatWithFollow,        window: &mut Window,        cx: &mut Context<Self>,    ) {        self.workspace            .update(cx, |this, cx| {                this.follow(CollaboratorId::Agent, window, cx)            })            .log_err();         self.chat(&Chat, window, cx);    }     fn is_editor_empty(&self, cx: &App) -> bool {        self.editor.read(cx).text(cx).trim().is_empty()    }     pub fn is_editor_fully_empty(&self, cx: &App) -> bool {        self.editor.read(cx).is_empty(cx)    }     fn send_to_model(&mut self, window: &mut Window, cx: &mut Context<Self>) {        let Some(ConfiguredModel { model, provider }) = self            .thread            .update(cx, |thread, cx| thread.get_or_init_configured_model(cx))        else {            return;        };         if provider.must_accept_terms(cx) {            cx.notify();            return;        }         let (user_message, user_message_creases) = self.editor.update(cx, |editor, cx| {            let creases = extract_message_creases(editor, cx);            let text = editor.text(cx);            editor.clear(window, cx);            (text, creases)        });         self.last_estimated_token_count.take();        cx.emit(MessageEditorEvent::EstimatedTokenCount);         let thread = self.thread.clone();        let git_store = self.project.read(cx).git_store().clone();        let checkpoint = git_store.update(cx, |git_store, cx| git_store.checkpoint(cx));        let context_task = self.reload_context(cx);        let window_handle = window.window_handle();         cx.spawn(async move |_this, cx| {            let (checkpoint, loaded_context) = future::join(checkpoint, context_task).await;            let loaded_context = loaded_context.unwrap_or_default();             thread                .update(cx, |thread, cx| {                    thread.insert_user_message(                        user_message,                        loaded_context,                        checkpoint.ok(),                        user_message_creases,                        cx,                    );                })                .log_err();             thread                .update(cx, |thread, cx| {                    thread.advance_prompt_id();                    thread.send_to_model(                        model,                        CompletionIntent::UserPrompt,                        Some(window_handle),                        cx,                    );                })                .log_err();        })        .detach();    }     fn stop_current_and_send_new_message(&mut self, window: &mut Window, cx: &mut Context<Self>) {        self.thread.update(cx, |thread, cx| {            thread.cancel_editing(cx);        });         let cancelled = self.thread.update(cx, |thread, cx| {            thread.cancel_last_completion(Some(window.window_handle()), cx)        });         if cancelled {            self.set_editor_is_expanded(false, cx);            self.send_to_model(window, cx);        }    }     fn handle_context_strip_event(        &mut self,        _context_strip: &Entity<ContextStrip>,        event: &ContextStripEvent,        window: &mut Window,        cx: &mut Context<Self>,    ) {        match event {            ContextStripEvent::PickerDismissed            | ContextStripEvent::BlurredEmpty            | ContextStripEvent::BlurredDown => {                let editor_focus_handle = self.editor.focus_handle(cx);                window.focus(&editor_focus_handle);            }            ContextStripEvent::BlurredUp => {}        }    }     fn move_up(&mut self, _: &MoveUp, window: &mut Window, cx: &mut Context<Self>) {        if self.context_picker_menu_handle.is_deployed() {            cx.propagate();        } else if self.context_strip.read(cx).has_context_items(cx) {            self.context_strip.focus_handle(cx).focus(window);        }    }     fn paste(&mut self, _: &Paste, _: &mut Window, cx: &mut Context<Self>) {        crate::active_thread::attach_pasted_images_as_context(&self.context_store, cx);    }     fn handle_review_click(&mut self, window: &mut Window, cx: &mut Context<Self>) {        if self.thread.read(cx).has_pending_edit_tool_uses() {            return;        }         self.edits_expanded = true;        AgentDiffPane::deploy(self.thread.clone(), self.workspace.clone(), window, cx).log_err();        cx.notify();    }     fn handle_edit_bar_expand(&mut self, cx: &mut Context<Self>) {        self.edits_expanded = !self.edits_expanded;        cx.notify();    }     fn handle_file_click(        &self,        buffer: Entity<Buffer>,        window: &mut Window,        cx: &mut Context<Self>,    ) {        if let Ok(diff) =            AgentDiffPane::deploy(self.thread.clone(), self.workspace.clone(), window, cx)        {            let path_key = multi_buffer::PathKey::for_buffer(&buffer, cx);            diff.update(cx, |diff, cx| diff.move_to_path(path_key, window, cx));        }    }     pub fn toggle_burn_mode(        &mut self,        _: &ToggleBurnMode,        _window: &mut Window,        cx: &mut Context<Self>,    ) {        self.thread.update(cx, |thread, _cx| {            let active_completion_mode = thread.completion_mode();             thread.set_completion_mode(match active_completion_mode {                CompletionMode::Burn => CompletionMode::Normal,                CompletionMode::Normal => CompletionMode::Burn,            });        });    }     fn handle_accept_all(&mut self, _window: &mut Window, cx: &mut Context<Self>) {        if self.thread.read(cx).has_pending_edit_tool_uses() {            return;        }         self.thread.update(cx, |thread, cx| {            thread.keep_all_edits(cx);        });        cx.notify();    }     fn handle_reject_all(&mut self, _window: &mut Window, cx: &mut Context<Self>) {        if self.thread.read(cx).has_pending_edit_tool_uses() {            return;        }         // Since there's no reject_all_edits method in the thread API,        // we need to iterate through all buffers and reject their edits

Lines 1–500 of 1604.