4132 lines · 154.1 KB · Rust
RawDownload
use crate::{    Assistant, AssistantPanel, AssistantPanelEvent, CycleNextInlineAssist,    CyclePreviousInlineAssist,};use anyhow::{Context as _, Result, anyhow};use assistant_context_editor::{RequestType, humanize_token_count};use assistant_settings::AssistantSettings;use client::{ErrorExt, telemetry::Telemetry};use collections::{HashMap, HashSet, VecDeque, hash_map};use editor::{    Anchor, AnchorRangeExt, CodeActionProvider, Editor, EditorElement, EditorEvent, EditorMode,    EditorStyle, ExcerptId, ExcerptRange, GutterDimensions, MultiBuffer, MultiBufferSnapshot,    ToOffset as _, ToPoint,    actions::{MoveDown, MoveUp, SelectAll},    display_map::{        BlockContext, BlockPlacement, BlockProperties, BlockStyle, CustomBlockId, RenderBlock,        ToDisplayPoint,    },};use feature_flags::{    Assistant2FeatureFlag, FeatureFlagAppExt as _, FeatureFlagViewExt as _, ZedPro,};use fs::Fs;use futures::{    SinkExt, Stream, StreamExt,    channel::mpsc,    future::{BoxFuture, LocalBoxFuture},    join,};use gpui::{    AnyElement, App, ClickEvent, Context, CursorStyle, Entity, EventEmitter, FocusHandle,    Focusable, FontWeight, Global, HighlightStyle, Subscription, Task, TextStyle, UpdateGlobal,    WeakEntity, Window, anchored, deferred, point,};use language::{Buffer, IndentKind, Point, Selection, TransactionId, line_diff};use language_model::{    ConfiguredModel, LanguageModel, LanguageModelRegistry, LanguageModelRequest,    LanguageModelRequestMessage, LanguageModelTextStream, Role, report_assistant_event,};use language_model_selector::{LanguageModelSelector, LanguageModelSelectorPopoverMenu};use multi_buffer::MultiBufferRow;use parking_lot::Mutex;use project::{CodeAction, LspAction, ProjectTransaction};use prompt_store::PromptBuilder;use rope::Rope;use settings::{Settings, SettingsStore, update_settings_file};use smol::future::FutureExt;use std::{    cmp,    future::{self, Future},    iter, mem,    ops::{Range, RangeInclusive},    pin::Pin,    rc::Rc,    sync::Arc,    task::{self, Poll},    time::{Duration, Instant},};use streaming_diff::{CharOperation, LineDiff, LineOperation, StreamingDiff};use telemetry_events::{AssistantEventData, AssistantKind, AssistantPhase};use terminal_view::terminal_panel::TerminalPanel;use text::{OffsetRangeExt, ToPoint as _};use theme::ThemeSettings;use ui::{    CheckboxWithLabel, IconButtonShape, KeyBinding, Popover, Tooltip, prelude::*, text_for_action,};use util::{RangeExt, ResultExt};use workspace::{ItemHandle, Toast, Workspace, notifications::NotificationId}; pub fn init(    fs: Arc<dyn Fs>,    prompt_builder: Arc<PromptBuilder>,    telemetry: Arc<Telemetry>,    cx: &mut App,) {    cx.set_global(InlineAssistant::new(fs, prompt_builder, telemetry));    cx.observe_new(|_, window, cx| {        let Some(window) = window else {            return;        };        let workspace = cx.entity().clone();        InlineAssistant::update_global(cx, |inline_assistant, cx| {            inline_assistant.register_workspace(&workspace, window, cx)        });         cx.observe_flag::<Assistant2FeatureFlag, _>(window, {            |is_assistant2_enabled, _workspace, _window, cx| {                InlineAssistant::update_global(cx, |inline_assistant, _cx| {                    inline_assistant.is_assistant2_enabled = is_assistant2_enabled;                });            }        })        .detach();    })    .detach();} const PROMPT_HISTORY_MAX_LEN: usize = 20; pub struct InlineAssistant {    next_assist_id: InlineAssistId,    next_assist_group_id: InlineAssistGroupId,    assists: HashMap<InlineAssistId, InlineAssist>,    assists_by_editor: HashMap<WeakEntity<Editor>, EditorInlineAssists>,    assist_groups: HashMap<InlineAssistGroupId, InlineAssistGroup>,    confirmed_assists: HashMap<InlineAssistId, Entity<CodegenAlternative>>,    prompt_history: VecDeque<String>,    prompt_builder: Arc<PromptBuilder>,    telemetry: Arc<Telemetry>,    fs: Arc<dyn Fs>,    is_assistant2_enabled: bool,} impl Global for InlineAssistant {} impl InlineAssistant {    pub fn new(        fs: Arc<dyn Fs>,        prompt_builder: Arc<PromptBuilder>,        telemetry: Arc<Telemetry>,    ) -> Self {        Self {            next_assist_id: InlineAssistId::default(),            next_assist_group_id: InlineAssistGroupId::default(),            assists: HashMap::default(),            assists_by_editor: HashMap::default(),            assist_groups: HashMap::default(),            confirmed_assists: HashMap::default(),            prompt_history: VecDeque::default(),            prompt_builder,            telemetry,            fs,            is_assistant2_enabled: false,        }    }     pub fn register_workspace(        &mut self,        workspace: &Entity<Workspace>,        window: &mut Window,        cx: &mut App,    ) {        window            .subscribe(workspace, cx, |workspace, event, window, cx| {                Self::update_global(cx, |this, cx| {                    this.handle_workspace_event(workspace, event, window, cx)                });            })            .detach();         let workspace = workspace.downgrade();        cx.observe_global::<SettingsStore>(move |cx| {            let Some(workspace) = workspace.upgrade() else {                return;            };            let Some(terminal_panel) = workspace.read(cx).panel::<TerminalPanel>(cx) else {                return;            };            let enabled = AssistantSettings::get_global(cx).enabled;            terminal_panel.update(cx, |terminal_panel, cx| {                terminal_panel.set_assistant_enabled(enabled, cx)            });        })        .detach();    }     fn handle_workspace_event(        &mut self,        workspace: Entity<Workspace>,        event: &workspace::Event,        window: &mut Window,        cx: &mut App,    ) {        match event {            workspace::Event::UserSavedItem { item, .. } => {                // When the user manually saves an editor, automatically accepts all finished transformations.                if let Some(editor) = item.upgrade().and_then(|item| item.act_as::<Editor>(cx)) {                    if let Some(editor_assists) = self.assists_by_editor.get(&editor.downgrade()) {                        for assist_id in editor_assists.assist_ids.clone() {                            let assist = &self.assists[&assist_id];                            if let CodegenStatus::Done = assist.codegen.read(cx).status(cx) {                                self.finish_assist(assist_id, false, window, cx)                            }                        }                    }                }            }            workspace::Event::ItemAdded { item } => {                self.register_workspace_item(&workspace, item.as_ref(), window, cx);            }            _ => (),        }    }     fn register_workspace_item(        &mut self,        workspace: &Entity<Workspace>,        item: &dyn ItemHandle,        window: &mut Window,        cx: &mut App,    ) {        let is_assistant2_enabled = self.is_assistant2_enabled;         if let Some(editor) = item.act_as::<Editor>(cx) {            editor.update(cx, |editor, cx| {                if is_assistant2_enabled {                    editor.remove_code_action_provider(                        ASSISTANT_CODE_ACTION_PROVIDER_ID.into(),                        window,                        cx,                    );                } else {                    editor.add_code_action_provider(                        Rc::new(AssistantCodeActionProvider {                            editor: cx.entity().downgrade(),                            workspace: workspace.downgrade(),                        }),                        window,                        cx,                    );                }            });        }    }     pub fn assist(        &mut self,        editor: &Entity<Editor>,        workspace: Option<WeakEntity<Workspace>>,        assistant_panel: Option<&Entity<AssistantPanel>>,        initial_prompt: Option<String>,        window: &mut Window,        cx: &mut App,    ) {        let (snapshot, initial_selections) = editor.update(cx, |editor, cx| {            (                editor.snapshot(window, cx),                editor.selections.all::<Point>(cx),            )        });         let mut selections = Vec::<Selection<Point>>::new();        let mut newest_selection = None;        for mut selection in initial_selections {            if selection.end > selection.start {                selection.start.column = 0;                // If the selection ends at the start of the line, we don't want to include it.                if selection.end.column == 0 {                    selection.end.row -= 1;                }                selection.end.column = snapshot                    .buffer_snapshot                    .line_len(MultiBufferRow(selection.end.row));            } else if let Some(fold) =                snapshot.crease_for_buffer_row(MultiBufferRow(selection.end.row))            {                selection.start = fold.range().start;                selection.end = fold.range().end;                if MultiBufferRow(selection.end.row) < snapshot.buffer_snapshot.max_row() {                    let chars = snapshot                        .buffer_snapshot                        .chars_at(Point::new(selection.end.row + 1, 0));                     for c in chars {                        if c == '\n' {                            break;                        }                        if c.is_whitespace() {                            continue;                        }                        if snapshot                            .language_at(selection.end)                            .is_some_and(|language| language.config().brackets.is_closing_brace(c))                        {                            selection.end.row += 1;                            selection.end.column = snapshot                                .buffer_snapshot                                .line_len(MultiBufferRow(selection.end.row));                        }                    }                }            }             if let Some(prev_selection) = selections.last_mut() {                if selection.start <= prev_selection.end {                    prev_selection.end = selection.end;                    continue;                }            }             let latest_selection = newest_selection.get_or_insert_with(|| selection.clone());            if selection.id > latest_selection.id {                *latest_selection = selection.clone();            }            selections.push(selection);        }        let snapshot = &snapshot.buffer_snapshot;        let newest_selection = newest_selection.unwrap();         let mut codegen_ranges = Vec::new();        for (buffer, buffer_range, excerpt_id) in            snapshot.ranges_to_buffer_ranges(selections.iter().map(|selection| {                snapshot.anchor_before(selection.start)..snapshot.anchor_after(selection.end)            }))        {            let start = buffer.anchor_before(buffer_range.start);            let end = buffer.anchor_after(buffer_range.end);             codegen_ranges.push(Anchor::range_in_buffer(                excerpt_id,                buffer.remote_id(),                start..end,            ));             if let Some(ConfiguredModel { model, .. }) =                LanguageModelRegistry::read_global(cx).default_model()            {                self.telemetry.report_assistant_event(AssistantEventData {                    conversation_id: None,                    kind: AssistantKind::Inline,                    phase: AssistantPhase::Invoked,                    message_id: None,                    model: model.telemetry_id(),                    model_provider: model.provider_id().to_string(),                    response_latency: None,                    error_message: None,                    language_name: buffer.language().map(|language| language.name().to_proto()),                });            }        }         let assist_group_id = self.next_assist_group_id.post_inc();        let prompt_buffer = cx.new(|cx| Buffer::local(initial_prompt.unwrap_or_default(), cx));        let prompt_buffer = cx.new(|cx| MultiBuffer::singleton(prompt_buffer, cx));         let mut assists = Vec::new();        let mut assist_to_focus = None;        for range in codegen_ranges {            let assist_id = self.next_assist_id.post_inc();            let codegen = cx.new(|cx| {                Codegen::new(                    editor.read(cx).buffer().clone(),                    range.clone(),                    None,                    self.telemetry.clone(),                    self.prompt_builder.clone(),                    cx,                )            });             let gutter_dimensions = Arc::new(Mutex::new(GutterDimensions::default()));            let prompt_editor = cx.new(|cx| {                PromptEditor::new(                    assist_id,                    gutter_dimensions.clone(),                    self.prompt_history.clone(),                    prompt_buffer.clone(),                    codegen.clone(),                    editor,                    assistant_panel,                    workspace.clone(),                    self.fs.clone(),                    window,                    cx,                )            });             if assist_to_focus.is_none() {                let focus_assist = if newest_selection.reversed {                    range.start.to_point(&snapshot) == newest_selection.start                } else {                    range.end.to_point(&snapshot) == newest_selection.end                };                if focus_assist {                    assist_to_focus = Some(assist_id);                }            }             let [prompt_block_id, end_block_id] =                self.insert_assist_blocks(editor, &range, &prompt_editor, cx);             assists.push((                assist_id,                range,                prompt_editor,                prompt_block_id,                end_block_id,            ));        }         let editor_assists = self            .assists_by_editor            .entry(editor.downgrade())            .or_insert_with(|| EditorInlineAssists::new(&editor, window, cx));        let mut assist_group = InlineAssistGroup::new();        for (assist_id, range, prompt_editor, prompt_block_id, end_block_id) in assists {            self.assists.insert(                assist_id,                InlineAssist::new(                    assist_id,                    assist_group_id,                    assistant_panel.is_some(),                    editor,                    &prompt_editor,                    prompt_block_id,                    end_block_id,                    range,                    prompt_editor.read(cx).codegen.clone(),                    workspace.clone(),                    window,                    cx,                ),            );            assist_group.assist_ids.push(assist_id);            editor_assists.assist_ids.push(assist_id);        }        self.assist_groups.insert(assist_group_id, assist_group);         if let Some(assist_id) = assist_to_focus {            self.focus_assist(assist_id, window, cx);        }    }     pub fn suggest_assist(        &mut self,        editor: &Entity<Editor>,        mut range: Range<Anchor>,        initial_prompt: String,        initial_transaction_id: Option<TransactionId>,        focus: bool,        workspace: Option<WeakEntity<Workspace>>,        assistant_panel: Option<&Entity<AssistantPanel>>,        window: &mut Window,        cx: &mut App,    ) -> InlineAssistId {        let assist_group_id = self.next_assist_group_id.post_inc();        let prompt_buffer = cx.new(|cx| Buffer::local(&initial_prompt, cx));        let prompt_buffer = cx.new(|cx| MultiBuffer::singleton(prompt_buffer, cx));         let assist_id = self.next_assist_id.post_inc();         let buffer = editor.read(cx).buffer().clone();        {            let snapshot = buffer.read(cx).read(cx);            range.start = range.start.bias_left(&snapshot);            range.end = range.end.bias_right(&snapshot);        }         let codegen = cx.new(|cx| {            Codegen::new(                editor.read(cx).buffer().clone(),                range.clone(),                initial_transaction_id,                self.telemetry.clone(),                self.prompt_builder.clone(),                cx,            )        });         let gutter_dimensions = Arc::new(Mutex::new(GutterDimensions::default()));        let prompt_editor = cx.new(|cx| {            PromptEditor::new(                assist_id,                gutter_dimensions.clone(),                self.prompt_history.clone(),                prompt_buffer.clone(),                codegen.clone(),                editor,                assistant_panel,                workspace.clone(),                self.fs.clone(),                window,                cx,            )        });         let [prompt_block_id, end_block_id] =            self.insert_assist_blocks(editor, &range, &prompt_editor, cx);         let editor_assists = self            .assists_by_editor            .entry(editor.downgrade())            .or_insert_with(|| EditorInlineAssists::new(&editor, window, cx));         let mut assist_group = InlineAssistGroup::new();        self.assists.insert(            assist_id,            InlineAssist::new(                assist_id,                assist_group_id,                assistant_panel.is_some(),                editor,                &prompt_editor,                prompt_block_id,                end_block_id,                range,                prompt_editor.read(cx).codegen.clone(),                workspace.clone(),                window,                cx,

Lines 1–500 of 4132.