assistant.rs

   1use crate::{
   2    assistant_settings::{AssistantDockPosition, AssistantSettings, OpenAIModel},
   3    codegen::{self, Codegen, CodegenKind, OpenAICompletionProvider},
   4    stream_completion, MessageId, MessageMetadata, MessageStatus, OpenAIRequest, RequestMessage,
   5    Role, SavedConversation, SavedConversationMetadata, SavedMessage, OPENAI_API_URL,
   6};
   7use anyhow::{anyhow, Result};
   8use chrono::{DateTime, Local};
   9use collections::{hash_map, HashMap, HashSet, VecDeque};
  10use editor::{
  11    display_map::{
  12        BlockContext, BlockDisposition, BlockId, BlockProperties, BlockStyle, ToDisplayPoint,
  13    },
  14    scroll::autoscroll::{Autoscroll, AutoscrollStrategy},
  15    Anchor, Editor, MoveDown, MoveUp, MultiBufferSnapshot, ToOffset,
  16};
  17use fs::Fs;
  18use futures::StreamExt;
  19use gpui::{
  20    actions,
  21    elements::{
  22        ChildView, Component, Empty, Flex, Label, MouseEventHandler, ParentElement, SafeStylable,
  23        Stack, Svg, Text, UniformList, UniformListState,
  24    },
  25    fonts::HighlightStyle,
  26    geometry::vector::{vec2f, Vector2F},
  27    platform::{CursorStyle, MouseButton},
  28    Action, AnyElement, AppContext, AsyncAppContext, ClipboardItem, Element, Entity, ModelContext,
  29    ModelHandle, SizeConstraint, Subscription, Task, View, ViewContext, ViewHandle, WeakViewHandle,
  30    WindowContext,
  31};
  32use language::{language_settings::SoftWrap, Buffer, LanguageRegistry, ToOffset as _};
  33use search::BufferSearchBar;
  34use settings::SettingsStore;
  35use std::{
  36    cell::{Cell, RefCell},
  37    cmp, env,
  38    fmt::Write,
  39    iter,
  40    ops::Range,
  41    path::{Path, PathBuf},
  42    rc::Rc,
  43    sync::Arc,
  44    time::Duration,
  45};
  46use theme::{
  47    components::{action_button::Button, ComponentExt},
  48    AssistantStyle,
  49};
  50use util::{paths::CONVERSATIONS_DIR, post_inc, ResultExt, TryFutureExt};
  51use workspace::{
  52    dock::{DockPosition, Panel},
  53    searchable::Direction,
  54    Save, Toast, ToggleZoom, Toolbar, Workspace,
  55};
  56
  57actions!(
  58    assistant,
  59    [
  60        NewConversation,
  61        Assist,
  62        Split,
  63        CycleMessageRole,
  64        QuoteSelection,
  65        ToggleFocus,
  66        ResetKey,
  67        InlineAssist,
  68        ToggleIncludeConversation,
  69    ]
  70);
  71
  72pub fn init(cx: &mut AppContext) {
  73    settings::register::<AssistantSettings>(cx);
  74    cx.add_action(
  75        |this: &mut AssistantPanel,
  76         _: &workspace::NewFile,
  77         cx: &mut ViewContext<AssistantPanel>| {
  78            this.new_conversation(cx);
  79        },
  80    );
  81    cx.add_action(ConversationEditor::assist);
  82    cx.capture_action(ConversationEditor::cancel_last_assist);
  83    cx.capture_action(ConversationEditor::save);
  84    cx.add_action(ConversationEditor::quote_selection);
  85    cx.capture_action(ConversationEditor::copy);
  86    cx.add_action(ConversationEditor::split);
  87    cx.capture_action(ConversationEditor::cycle_message_role);
  88    cx.add_action(AssistantPanel::save_api_key);
  89    cx.add_action(AssistantPanel::reset_api_key);
  90    cx.add_action(AssistantPanel::toggle_zoom);
  91    cx.add_action(AssistantPanel::deploy);
  92    cx.add_action(AssistantPanel::select_next_match);
  93    cx.add_action(AssistantPanel::select_prev_match);
  94    cx.add_action(AssistantPanel::handle_editor_cancel);
  95    cx.add_action(
  96        |workspace: &mut Workspace, _: &ToggleFocus, cx: &mut ViewContext<Workspace>| {
  97            workspace.toggle_panel_focus::<AssistantPanel>(cx);
  98        },
  99    );
 100    cx.add_action(AssistantPanel::inline_assist);
 101    cx.add_action(AssistantPanel::cancel_last_inline_assist);
 102    cx.add_action(InlineAssistant::confirm);
 103    cx.add_action(InlineAssistant::cancel);
 104    cx.add_action(InlineAssistant::toggle_include_conversation);
 105    cx.add_action(InlineAssistant::move_up);
 106    cx.add_action(InlineAssistant::move_down);
 107}
 108
 109#[derive(Debug)]
 110pub enum AssistantPanelEvent {
 111    ZoomIn,
 112    ZoomOut,
 113    Focus,
 114    Close,
 115    DockPositionChanged,
 116}
 117
 118pub struct AssistantPanel {
 119    workspace: WeakViewHandle<Workspace>,
 120    width: Option<f32>,
 121    height: Option<f32>,
 122    active_editor_index: Option<usize>,
 123    prev_active_editor_index: Option<usize>,
 124    editors: Vec<ViewHandle<ConversationEditor>>,
 125    saved_conversations: Vec<SavedConversationMetadata>,
 126    saved_conversations_list_state: UniformListState,
 127    zoomed: bool,
 128    has_focus: bool,
 129    toolbar: ViewHandle<Toolbar>,
 130    api_key: Rc<RefCell<Option<String>>>,
 131    api_key_editor: Option<ViewHandle<Editor>>,
 132    has_read_credentials: bool,
 133    languages: Arc<LanguageRegistry>,
 134    fs: Arc<dyn Fs>,
 135    subscriptions: Vec<Subscription>,
 136    next_inline_assist_id: usize,
 137    pending_inline_assists: HashMap<usize, PendingInlineAssist>,
 138    pending_inline_assist_ids_by_editor: HashMap<WeakViewHandle<Editor>, Vec<usize>>,
 139    include_conversation_in_next_inline_assist: bool,
 140    inline_prompt_history: VecDeque<String>,
 141    _watch_saved_conversations: Task<Result<()>>,
 142}
 143
 144impl AssistantPanel {
 145    const INLINE_PROMPT_HISTORY_MAX_LEN: usize = 20;
 146
 147    pub fn load(
 148        workspace: WeakViewHandle<Workspace>,
 149        cx: AsyncAppContext,
 150    ) -> Task<Result<ViewHandle<Self>>> {
 151        cx.spawn(|mut cx| async move {
 152            let fs = workspace.read_with(&cx, |workspace, _| workspace.app_state().fs.clone())?;
 153            let saved_conversations = SavedConversationMetadata::list(fs.clone())
 154                .await
 155                .log_err()
 156                .unwrap_or_default();
 157
 158            // TODO: deserialize state.
 159            let workspace_handle = workspace.clone();
 160            workspace.update(&mut cx, |workspace, cx| {
 161                cx.add_view::<Self, _>(|cx| {
 162                    const CONVERSATION_WATCH_DURATION: Duration = Duration::from_millis(100);
 163                    let _watch_saved_conversations = cx.spawn(move |this, mut cx| async move {
 164                        let mut events = fs
 165                            .watch(&CONVERSATIONS_DIR, CONVERSATION_WATCH_DURATION)
 166                            .await;
 167                        while events.next().await.is_some() {
 168                            let saved_conversations = SavedConversationMetadata::list(fs.clone())
 169                                .await
 170                                .log_err()
 171                                .unwrap_or_default();
 172                            this.update(&mut cx, |this, cx| {
 173                                this.saved_conversations = saved_conversations;
 174                                cx.notify();
 175                            })
 176                            .ok();
 177                        }
 178
 179                        anyhow::Ok(())
 180                    });
 181
 182                    let toolbar = cx.add_view(|cx| {
 183                        let mut toolbar = Toolbar::new();
 184                        toolbar.set_can_navigate(false, cx);
 185                        toolbar.add_item(cx.add_view(|cx| BufferSearchBar::new(cx)), cx);
 186                        toolbar
 187                    });
 188                    let mut this = Self {
 189                        workspace: workspace_handle,
 190                        active_editor_index: Default::default(),
 191                        prev_active_editor_index: Default::default(),
 192                        editors: Default::default(),
 193                        saved_conversations,
 194                        saved_conversations_list_state: Default::default(),
 195                        zoomed: false,
 196                        has_focus: false,
 197                        toolbar,
 198                        api_key: Rc::new(RefCell::new(None)),
 199                        api_key_editor: None,
 200                        has_read_credentials: false,
 201                        languages: workspace.app_state().languages.clone(),
 202                        fs: workspace.app_state().fs.clone(),
 203                        width: None,
 204                        height: None,
 205                        subscriptions: Default::default(),
 206                        next_inline_assist_id: 0,
 207                        pending_inline_assists: Default::default(),
 208                        pending_inline_assist_ids_by_editor: Default::default(),
 209                        include_conversation_in_next_inline_assist: false,
 210                        inline_prompt_history: Default::default(),
 211                        _watch_saved_conversations,
 212                    };
 213
 214                    let mut old_dock_position = this.position(cx);
 215                    this.subscriptions =
 216                        vec![cx.observe_global::<SettingsStore, _>(move |this, cx| {
 217                            let new_dock_position = this.position(cx);
 218                            if new_dock_position != old_dock_position {
 219                                old_dock_position = new_dock_position;
 220                                cx.emit(AssistantPanelEvent::DockPositionChanged);
 221                            }
 222                            cx.notify();
 223                        })];
 224
 225                    this
 226                })
 227            })
 228        })
 229    }
 230
 231    pub fn inline_assist(
 232        workspace: &mut Workspace,
 233        _: &InlineAssist,
 234        cx: &mut ViewContext<Workspace>,
 235    ) {
 236        let this = if let Some(this) = workspace.panel::<AssistantPanel>(cx) {
 237            if this
 238                .update(cx, |assistant, cx| assistant.load_api_key(cx))
 239                .is_some()
 240            {
 241                this
 242            } else {
 243                workspace.focus_panel::<AssistantPanel>(cx);
 244                return;
 245            }
 246        } else {
 247            return;
 248        };
 249
 250        let active_editor = if let Some(active_editor) = workspace
 251            .active_item(cx)
 252            .and_then(|item| item.act_as::<Editor>(cx))
 253        {
 254            active_editor
 255        } else {
 256            return;
 257        };
 258
 259        this.update(cx, |assistant, cx| {
 260            assistant.new_inline_assist(&active_editor, cx)
 261        });
 262    }
 263
 264    fn new_inline_assist(&mut self, editor: &ViewHandle<Editor>, cx: &mut ViewContext<Self>) {
 265        let api_key = if let Some(api_key) = self.api_key.borrow().clone() {
 266            api_key
 267        } else {
 268            return;
 269        };
 270
 271        let inline_assist_id = post_inc(&mut self.next_inline_assist_id);
 272        let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx);
 273        let provider = Arc::new(OpenAICompletionProvider::new(
 274            api_key,
 275            cx.background().clone(),
 276        ));
 277        let selection = editor.read(cx).selections.newest_anchor().clone();
 278        let codegen_kind = if editor.read(cx).selections.newest::<usize>(cx).is_empty() {
 279            CodegenKind::Generate {
 280                position: selection.start,
 281            }
 282        } else {
 283            CodegenKind::Transform {
 284                range: selection.start..selection.end,
 285            }
 286        };
 287        let codegen = cx.add_model(|cx| {
 288            Codegen::new(editor.read(cx).buffer().clone(), codegen_kind, provider, cx)
 289        });
 290
 291        let measurements = Rc::new(Cell::new(BlockMeasurements::default()));
 292        let inline_assistant = cx.add_view(|cx| {
 293            let assistant = InlineAssistant::new(
 294                inline_assist_id,
 295                measurements.clone(),
 296                self.include_conversation_in_next_inline_assist,
 297                self.inline_prompt_history.clone(),
 298                codegen.clone(),
 299                cx,
 300            );
 301            cx.focus_self();
 302            assistant
 303        });
 304        let block_id = editor.update(cx, |editor, cx| {
 305            editor.change_selections(None, cx, |selections| {
 306                selections.select_anchor_ranges([selection.head()..selection.head()])
 307            });
 308            editor.insert_blocks(
 309                [BlockProperties {
 310                    style: BlockStyle::Flex,
 311                    position: selection.head().bias_left(&snapshot),
 312                    height: 2,
 313                    render: Arc::new({
 314                        let inline_assistant = inline_assistant.clone();
 315                        move |cx: &mut BlockContext| {
 316                            measurements.set(BlockMeasurements {
 317                                anchor_x: cx.anchor_x,
 318                                gutter_width: cx.gutter_width,
 319                            });
 320                            ChildView::new(&inline_assistant, cx).into_any()
 321                        }
 322                    }),
 323                    disposition: if selection.reversed {
 324                        BlockDisposition::Above
 325                    } else {
 326                        BlockDisposition::Below
 327                    },
 328                }],
 329                Some(Autoscroll::Strategy(AutoscrollStrategy::Newest)),
 330                cx,
 331            )[0]
 332        });
 333
 334        self.pending_inline_assists.insert(
 335            inline_assist_id,
 336            PendingInlineAssist {
 337                editor: editor.downgrade(),
 338                inline_assistant: Some((block_id, inline_assistant.clone())),
 339                codegen: codegen.clone(),
 340                _subscriptions: vec![
 341                    cx.subscribe(&inline_assistant, Self::handle_inline_assistant_event),
 342                    cx.subscribe(editor, {
 343                        let inline_assistant = inline_assistant.downgrade();
 344                        move |_, editor, event, cx| {
 345                            if let Some(inline_assistant) = inline_assistant.upgrade(cx) {
 346                                if let editor::Event::SelectionsChanged { local } = event {
 347                                    if *local && inline_assistant.read(cx).has_focus {
 348                                        cx.focus(&editor);
 349                                    }
 350                                }
 351                            }
 352                        }
 353                    }),
 354                    cx.observe(&codegen, {
 355                        let editor = editor.downgrade();
 356                        move |this, _, cx| {
 357                            if let Some(editor) = editor.upgrade(cx) {
 358                                this.update_highlights_for_editor(&editor, cx);
 359                            }
 360                        }
 361                    }),
 362                    cx.subscribe(&codegen, move |this, codegen, event, cx| match event {
 363                        codegen::Event::Undone => {
 364                            this.finish_inline_assist(inline_assist_id, false, cx)
 365                        }
 366                        codegen::Event::Finished => {
 367                            let pending_assist = if let Some(pending_assist) =
 368                                this.pending_inline_assists.get(&inline_assist_id)
 369                            {
 370                                pending_assist
 371                            } else {
 372                                return;
 373                            };
 374
 375                            let error = codegen
 376                                .read(cx)
 377                                .error()
 378                                .map(|error| format!("Inline assistant error: {}", error));
 379                            if let Some(error) = error {
 380                                if pending_assist.inline_assistant.is_none() {
 381                                    if let Some(workspace) = this.workspace.upgrade(cx) {
 382                                        workspace.update(cx, |workspace, cx| {
 383                                            workspace.show_toast(
 384                                                Toast::new(inline_assist_id, error),
 385                                                cx,
 386                                            );
 387                                        })
 388                                    }
 389                                }
 390                            }
 391
 392                            this.finish_inline_assist(inline_assist_id, false, cx);
 393                        }
 394                    }),
 395                ],
 396            },
 397        );
 398        self.pending_inline_assist_ids_by_editor
 399            .entry(editor.downgrade())
 400            .or_default()
 401            .push(inline_assist_id);
 402        self.update_highlights_for_editor(&editor, cx);
 403    }
 404
 405    fn handle_inline_assistant_event(
 406        &mut self,
 407        inline_assistant: ViewHandle<InlineAssistant>,
 408        event: &InlineAssistantEvent,
 409        cx: &mut ViewContext<Self>,
 410    ) {
 411        let assist_id = inline_assistant.read(cx).id;
 412        match event {
 413            InlineAssistantEvent::Confirmed {
 414                prompt,
 415                include_conversation,
 416            } => {
 417                self.confirm_inline_assist(assist_id, prompt, *include_conversation, cx);
 418            }
 419            InlineAssistantEvent::Canceled => {
 420                self.finish_inline_assist(assist_id, true, cx);
 421            }
 422            InlineAssistantEvent::Dismissed => {
 423                self.hide_inline_assist(assist_id, cx);
 424            }
 425            InlineAssistantEvent::IncludeConversationToggled {
 426                include_conversation,
 427            } => {
 428                self.include_conversation_in_next_inline_assist = *include_conversation;
 429            }
 430        }
 431    }
 432
 433    fn cancel_last_inline_assist(
 434        workspace: &mut Workspace,
 435        _: &editor::Cancel,
 436        cx: &mut ViewContext<Workspace>,
 437    ) {
 438        if let Some(panel) = workspace.panel::<AssistantPanel>(cx) {
 439            if let Some(editor) = workspace
 440                .active_item(cx)
 441                .and_then(|item| item.downcast::<Editor>())
 442            {
 443                let handled = panel.update(cx, |panel, cx| {
 444                    if let Some(assist_id) = panel
 445                        .pending_inline_assist_ids_by_editor
 446                        .get(&editor.downgrade())
 447                        .and_then(|assist_ids| assist_ids.last().copied())
 448                    {
 449                        panel.finish_inline_assist(assist_id, true, cx);
 450                        true
 451                    } else {
 452                        false
 453                    }
 454                });
 455                if handled {
 456                    return;
 457                }
 458            }
 459        }
 460
 461        cx.propagate_action();
 462    }
 463
 464    fn finish_inline_assist(&mut self, assist_id: usize, undo: bool, cx: &mut ViewContext<Self>) {
 465        self.hide_inline_assist(assist_id, cx);
 466
 467        if let Some(pending_assist) = self.pending_inline_assists.remove(&assist_id) {
 468            if let hash_map::Entry::Occupied(mut entry) = self
 469                .pending_inline_assist_ids_by_editor
 470                .entry(pending_assist.editor)
 471            {
 472                entry.get_mut().retain(|id| *id != assist_id);
 473                if entry.get().is_empty() {
 474                    entry.remove();
 475                }
 476            }
 477
 478            if let Some(editor) = pending_assist.editor.upgrade(cx) {
 479                self.update_highlights_for_editor(&editor, cx);
 480
 481                if undo {
 482                    pending_assist
 483                        .codegen
 484                        .update(cx, |codegen, cx| codegen.undo(cx));
 485                }
 486            }
 487        }
 488    }
 489
 490    fn hide_inline_assist(&mut self, assist_id: usize, cx: &mut ViewContext<Self>) {
 491        if let Some(pending_assist) = self.pending_inline_assists.get_mut(&assist_id) {
 492            if let Some(editor) = pending_assist.editor.upgrade(cx) {
 493                if let Some((block_id, _)) = pending_assist.inline_assistant.take() {
 494                    editor.update(cx, |editor, cx| {
 495                        editor.remove_blocks(HashSet::from_iter([block_id]), None, cx);
 496                    });
 497                }
 498            }
 499        }
 500    }
 501
 502    fn confirm_inline_assist(
 503        &mut self,
 504        inline_assist_id: usize,
 505        user_prompt: &str,
 506        include_conversation: bool,
 507        cx: &mut ViewContext<Self>,
 508    ) {
 509        let conversation = if include_conversation {
 510            self.active_editor()
 511                .map(|editor| editor.read(cx).conversation.clone())
 512        } else {
 513            None
 514        };
 515
 516        let pending_assist =
 517            if let Some(pending_assist) = self.pending_inline_assists.get_mut(&inline_assist_id) {
 518                pending_assist
 519            } else {
 520                return;
 521            };
 522
 523        let editor = if let Some(editor) = pending_assist.editor.upgrade(cx) {
 524            editor
 525        } else {
 526            return;
 527        };
 528
 529        self.inline_prompt_history
 530            .retain(|prompt| prompt != user_prompt);
 531        self.inline_prompt_history.push_back(user_prompt.into());
 532        if self.inline_prompt_history.len() > Self::INLINE_PROMPT_HISTORY_MAX_LEN {
 533            self.inline_prompt_history.pop_front();
 534        }
 535
 536        let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx);
 537        let range = pending_assist.codegen.read(cx).range();
 538        let selected_text = snapshot.text_for_range(range.clone()).collect::<String>();
 539
 540        let language = snapshot.language_at(range.start);
 541        let language_name = if let Some(language) = language.as_ref() {
 542            if Arc::ptr_eq(language, &language::PLAIN_TEXT) {
 543                None
 544            } else {
 545                Some(language.name())
 546            }
 547        } else {
 548            None
 549        };
 550        let language_name = language_name.as_deref();
 551
 552        let mut prompt = String::new();
 553        if let Some(language_name) = language_name {
 554            writeln!(prompt, "You're an expert {language_name} engineer.").unwrap();
 555        }
 556        match pending_assist.codegen.read(cx).kind() {
 557            CodegenKind::Transform { .. } => {
 558                writeln!(
 559                    prompt,
 560                    "You're currently working inside an editor on this file:"
 561                )
 562                .unwrap();
 563                if let Some(language_name) = language_name {
 564                    writeln!(prompt, "```{language_name}").unwrap();
 565                } else {
 566                    writeln!(prompt, "```").unwrap();
 567                }
 568                for chunk in snapshot.text_for_range(Anchor::min()..Anchor::max()) {
 569                    write!(prompt, "{chunk}").unwrap();
 570                }
 571                writeln!(prompt, "```").unwrap();
 572
 573                writeln!(
 574                    prompt,
 575                    "In particular, the user has selected the following text:"
 576                )
 577                .unwrap();
 578                if let Some(language_name) = language_name {
 579                    writeln!(prompt, "```{language_name}").unwrap();
 580                } else {
 581                    writeln!(prompt, "```").unwrap();
 582                }
 583                writeln!(prompt, "{selected_text}").unwrap();
 584                writeln!(prompt, "```").unwrap();
 585                writeln!(prompt).unwrap();
 586                writeln!(
 587                    prompt,
 588                    "Modify the selected text given the user prompt: {user_prompt}"
 589                )
 590                .unwrap();
 591                writeln!(
 592                    prompt,
 593                    "You MUST reply only with the edited selected text, not the entire file."
 594                )
 595                .unwrap();
 596            }
 597            CodegenKind::Generate { .. } => {
 598                writeln!(
 599                    prompt,
 600                    "You're currently working inside an editor on this file:"
 601                )
 602                .unwrap();
 603                if let Some(language_name) = language_name {
 604                    writeln!(prompt, "```{language_name}").unwrap();
 605                } else {
 606                    writeln!(prompt, "```").unwrap();
 607                }
 608                for chunk in snapshot.text_for_range(Anchor::min()..range.start) {
 609                    write!(prompt, "{chunk}").unwrap();
 610                }
 611                write!(prompt, "<|>").unwrap();
 612                for chunk in snapshot.text_for_range(range.start..Anchor::max()) {
 613                    write!(prompt, "{chunk}").unwrap();
 614                }
 615                writeln!(prompt).unwrap();
 616                writeln!(prompt, "```").unwrap();
 617                writeln!(
 618                    prompt,
 619                    "Assume the cursor is located where the `<|>` marker is."
 620                )
 621                .unwrap();
 622                writeln!(
 623                    prompt,
 624                    "Text can't be replaced, so assume your answer will be inserted at the cursor."
 625                )
 626                .unwrap();
 627                writeln!(
 628                    prompt,
 629                    "Complete the text given the user prompt: {user_prompt}"
 630                )
 631                .unwrap();
 632            }
 633        }
 634        if let Some(language_name) = language_name {
 635            writeln!(prompt, "Your answer MUST always be valid {language_name}.").unwrap();
 636        }
 637        writeln!(prompt, "Always wrap your response in a Markdown codeblock.").unwrap();
 638        writeln!(prompt, "Never make remarks about the output.").unwrap();
 639
 640        let mut messages = Vec::new();
 641        let mut model = settings::get::<AssistantSettings>(cx)
 642            .default_open_ai_model
 643            .clone();
 644        if let Some(conversation) = conversation {
 645            let conversation = conversation.read(cx);
 646            let buffer = conversation.buffer.read(cx);
 647            messages.extend(
 648                conversation
 649                    .messages(cx)
 650                    .map(|message| message.to_open_ai_message(buffer)),
 651            );
 652            model = conversation.model.clone();
 653        }
 654
 655        messages.push(RequestMessage {
 656            role: Role::User,
 657            content: prompt,
 658        });
 659        let request = OpenAIRequest {
 660            model: model.full_name().into(),
 661            messages,
 662            stream: true,
 663        };
 664        pending_assist
 665            .codegen
 666            .update(cx, |codegen, cx| codegen.start(request, cx));
 667    }
 668
 669    fn update_highlights_for_editor(
 670        &self,
 671        editor: &ViewHandle<Editor>,
 672        cx: &mut ViewContext<Self>,
 673    ) {
 674        let mut background_ranges = Vec::new();
 675        let mut foreground_ranges = Vec::new();
 676        let empty_inline_assist_ids = Vec::new();
 677        let inline_assist_ids = self
 678            .pending_inline_assist_ids_by_editor
 679            .get(&editor.downgrade())
 680            .unwrap_or(&empty_inline_assist_ids);
 681
 682        for inline_assist_id in inline_assist_ids {
 683            if let Some(pending_assist) = self.pending_inline_assists.get(inline_assist_id) {
 684                let codegen = pending_assist.codegen.read(cx);
 685                background_ranges.push(codegen.range());
 686                foreground_ranges.extend(codegen.last_equal_ranges().iter().cloned());
 687            }
 688        }
 689
 690        let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx);
 691        merge_ranges(&mut background_ranges, &snapshot);
 692        merge_ranges(&mut foreground_ranges, &snapshot);
 693        editor.update(cx, |editor, cx| {
 694            if background_ranges.is_empty() {
 695                editor.clear_background_highlights::<PendingInlineAssist>(cx);
 696            } else {
 697                editor.highlight_background::<PendingInlineAssist>(
 698                    background_ranges,
 699                    |theme| theme.assistant.inline.pending_edit_background,
 700                    cx,
 701                );
 702            }
 703
 704            if foreground_ranges.is_empty() {
 705                editor.clear_text_highlights::<PendingInlineAssist>(cx);
 706            } else {
 707                editor.highlight_text::<PendingInlineAssist>(
 708                    foreground_ranges,
 709                    HighlightStyle {
 710                        fade_out: Some(0.6),
 711                        ..Default::default()
 712                    },
 713                    cx,
 714                );
 715            }
 716        });
 717    }
 718
 719    fn new_conversation(&mut self, cx: &mut ViewContext<Self>) -> ViewHandle<ConversationEditor> {
 720        let editor = cx.add_view(|cx| {
 721            ConversationEditor::new(
 722                self.api_key.clone(),
 723                self.languages.clone(),
 724                self.fs.clone(),
 725                cx,
 726            )
 727        });
 728        self.add_conversation(editor.clone(), cx);
 729        editor
 730    }
 731
 732    fn add_conversation(
 733        &mut self,
 734        editor: ViewHandle<ConversationEditor>,
 735        cx: &mut ViewContext<Self>,
 736    ) {
 737        self.subscriptions
 738            .push(cx.subscribe(&editor, Self::handle_conversation_editor_event));
 739
 740        let conversation = editor.read(cx).conversation.clone();
 741        self.subscriptions
 742            .push(cx.observe(&conversation, |_, _, cx| cx.notify()));
 743
 744        let index = self.editors.len();
 745        self.editors.push(editor);
 746        self.set_active_editor_index(Some(index), cx);
 747    }
 748
 749    fn set_active_editor_index(&mut self, index: Option<usize>, cx: &mut ViewContext<Self>) {
 750        self.prev_active_editor_index = self.active_editor_index;
 751        self.active_editor_index = index;
 752        if let Some(editor) = self.active_editor() {
 753            let editor = editor.read(cx).editor.clone();
 754            self.toolbar.update(cx, |toolbar, cx| {
 755                toolbar.set_active_item(Some(&editor), cx);
 756            });
 757            if self.has_focus(cx) {
 758                cx.focus(&editor);
 759            }
 760        } else {
 761            self.toolbar.update(cx, |toolbar, cx| {
 762                toolbar.set_active_item(None, cx);
 763            });
 764        }
 765
 766        cx.notify();
 767    }
 768
 769    fn handle_conversation_editor_event(
 770        &mut self,
 771        _: ViewHandle<ConversationEditor>,
 772        event: &ConversationEditorEvent,
 773        cx: &mut ViewContext<Self>,
 774    ) {
 775        match event {
 776            ConversationEditorEvent::TabContentChanged => cx.notify(),
 777        }
 778    }
 779
 780    fn save_api_key(&mut self, _: &menu::Confirm, cx: &mut ViewContext<Self>) {
 781        if let Some(api_key) = self
 782            .api_key_editor
 783            .as_ref()
 784            .map(|editor| editor.read(cx).text(cx))
 785        {
 786            if !api_key.is_empty() {
 787                cx.platform()
 788                    .write_credentials(OPENAI_API_URL, "Bearer", api_key.as_bytes())
 789                    .log_err();
 790                *self.api_key.borrow_mut() = Some(api_key);
 791                self.api_key_editor.take();
 792                cx.focus_self();
 793                cx.notify();
 794            }
 795        } else {
 796            cx.propagate_action();
 797        }
 798    }
 799
 800    fn reset_api_key(&mut self, _: &ResetKey, cx: &mut ViewContext<Self>) {
 801        cx.platform().delete_credentials(OPENAI_API_URL).log_err();
 802        self.api_key.take();
 803        self.api_key_editor = Some(build_api_key_editor(cx));
 804        cx.focus_self();
 805        cx.notify();
 806    }
 807
 808    fn toggle_zoom(&mut self, _: &workspace::ToggleZoom, cx: &mut ViewContext<Self>) {
 809        if self.zoomed {
 810            cx.emit(AssistantPanelEvent::ZoomOut)
 811        } else {
 812            cx.emit(AssistantPanelEvent::ZoomIn)
 813        }
 814    }
 815
 816    fn deploy(&mut self, action: &search::buffer_search::Deploy, cx: &mut ViewContext<Self>) {
 817        let mut propagate_action = true;
 818        if let Some(search_bar) = self.toolbar.read(cx).item_of_type::<BufferSearchBar>() {
 819            search_bar.update(cx, |search_bar, cx| {
 820                if search_bar.show(cx) {
 821                    search_bar.search_suggested(cx);
 822                    if action.focus {
 823                        search_bar.select_query(cx);
 824                        cx.focus_self();
 825                    }
 826                    propagate_action = false
 827                }
 828            });
 829        }
 830        if propagate_action {
 831            cx.propagate_action();
 832        }
 833    }
 834
 835    fn handle_editor_cancel(&mut self, _: &editor::Cancel, cx: &mut ViewContext<Self>) {
 836        if let Some(search_bar) = self.toolbar.read(cx).item_of_type::<BufferSearchBar>() {
 837            if !search_bar.read(cx).is_dismissed() {
 838                search_bar.update(cx, |search_bar, cx| {
 839                    search_bar.dismiss(&Default::default(), cx)
 840                });
 841                return;
 842            }
 843        }
 844        cx.propagate_action();
 845    }
 846
 847    fn select_next_match(&mut self, _: &search::SelectNextMatch, cx: &mut ViewContext<Self>) {
 848        if let Some(search_bar) = self.toolbar.read(cx).item_of_type::<BufferSearchBar>() {
 849            search_bar.update(cx, |bar, cx| bar.select_match(Direction::Next, 1, cx));
 850        }
 851    }
 852
 853    fn select_prev_match(&mut self, _: &search::SelectPrevMatch, cx: &mut ViewContext<Self>) {
 854        if let Some(search_bar) = self.toolbar.read(cx).item_of_type::<BufferSearchBar>() {
 855            search_bar.update(cx, |bar, cx| bar.select_match(Direction::Prev, 1, cx));
 856        }
 857    }
 858
 859    fn active_editor(&self) -> Option<&ViewHandle<ConversationEditor>> {
 860        self.editors.get(self.active_editor_index?)
 861    }
 862
 863    fn render_hamburger_button(cx: &mut ViewContext<Self>) -> impl Element<Self> {
 864        enum History {}
 865        let theme = theme::current(cx);
 866        let tooltip_style = theme::current(cx).tooltip.clone();
 867        MouseEventHandler::new::<History, _>(0, cx, |state, _| {
 868            let style = theme.assistant.hamburger_button.style_for(state);
 869            Svg::for_style(style.icon.clone())
 870                .contained()
 871                .with_style(style.container)
 872        })
 873        .with_cursor_style(CursorStyle::PointingHand)
 874        .on_click(MouseButton::Left, |_, this: &mut Self, cx| {
 875            if this.active_editor().is_some() {
 876                this.set_active_editor_index(None, cx);
 877            } else {
 878                this.set_active_editor_index(this.prev_active_editor_index, cx);
 879            }
 880        })
 881        .with_tooltip::<History>(1, "History", None, tooltip_style, cx)
 882    }
 883
 884    fn render_editor_tools(&self, cx: &mut ViewContext<Self>) -> Vec<AnyElement<Self>> {
 885        if self.active_editor().is_some() {
 886            vec![
 887                Self::render_split_button(cx).into_any(),
 888                Self::render_quote_button(cx).into_any(),
 889                Self::render_assist_button(cx).into_any(),
 890            ]
 891        } else {
 892            Default::default()
 893        }
 894    }
 895
 896    fn render_split_button(cx: &mut ViewContext<Self>) -> impl Element<Self> {
 897        let theme = theme::current(cx);
 898        let tooltip_style = theme::current(cx).tooltip.clone();
 899        MouseEventHandler::new::<Split, _>(0, cx, |state, _| {
 900            let style = theme.assistant.split_button.style_for(state);
 901            Svg::for_style(style.icon.clone())
 902                .contained()
 903                .with_style(style.container)
 904        })
 905        .with_cursor_style(CursorStyle::PointingHand)
 906        .on_click(MouseButton::Left, |_, this: &mut Self, cx| {
 907            if let Some(active_editor) = this.active_editor() {
 908                active_editor.update(cx, |editor, cx| editor.split(&Default::default(), cx));
 909            }
 910        })
 911        .with_tooltip::<Split>(
 912            1,
 913            "Split Message",
 914            Some(Box::new(Split)),
 915            tooltip_style,
 916            cx,
 917        )
 918    }
 919
 920    fn render_assist_button(cx: &mut ViewContext<Self>) -> impl Element<Self> {
 921        let theme = theme::current(cx);
 922        let tooltip_style = theme::current(cx).tooltip.clone();
 923        MouseEventHandler::new::<Assist, _>(0, cx, |state, _| {
 924            let style = theme.assistant.assist_button.style_for(state);
 925            Svg::for_style(style.icon.clone())
 926                .contained()
 927                .with_style(style.container)
 928        })
 929        .with_cursor_style(CursorStyle::PointingHand)
 930        .on_click(MouseButton::Left, |_, this: &mut Self, cx| {
 931            if let Some(active_editor) = this.active_editor() {
 932                active_editor.update(cx, |editor, cx| editor.assist(&Default::default(), cx));
 933            }
 934        })
 935        .with_tooltip::<Assist>(1, "Assist", Some(Box::new(Assist)), tooltip_style, cx)
 936    }
 937
 938    fn render_quote_button(cx: &mut ViewContext<Self>) -> impl Element<Self> {
 939        let theme = theme::current(cx);
 940        let tooltip_style = theme::current(cx).tooltip.clone();
 941        MouseEventHandler::new::<QuoteSelection, _>(0, cx, |state, _| {
 942            let style = theme.assistant.quote_button.style_for(state);
 943            Svg::for_style(style.icon.clone())
 944                .contained()
 945                .with_style(style.container)
 946        })
 947        .with_cursor_style(CursorStyle::PointingHand)
 948        .on_click(MouseButton::Left, |_, this: &mut Self, cx| {
 949            if let Some(workspace) = this.workspace.upgrade(cx) {
 950                cx.window_context().defer(move |cx| {
 951                    workspace.update(cx, |workspace, cx| {
 952                        ConversationEditor::quote_selection(workspace, &Default::default(), cx)
 953                    });
 954                });
 955            }
 956        })
 957        .with_tooltip::<QuoteSelection>(
 958            1,
 959            "Quote Selection",
 960            Some(Box::new(QuoteSelection)),
 961            tooltip_style,
 962            cx,
 963        )
 964    }
 965
 966    fn render_plus_button(cx: &mut ViewContext<Self>) -> impl Element<Self> {
 967        let theme = theme::current(cx);
 968        let tooltip_style = theme::current(cx).tooltip.clone();
 969        MouseEventHandler::new::<NewConversation, _>(0, cx, |state, _| {
 970            let style = theme.assistant.plus_button.style_for(state);
 971            Svg::for_style(style.icon.clone())
 972                .contained()
 973                .with_style(style.container)
 974        })
 975        .with_cursor_style(CursorStyle::PointingHand)
 976        .on_click(MouseButton::Left, |_, this: &mut Self, cx| {
 977            this.new_conversation(cx);
 978        })
 979        .with_tooltip::<NewConversation>(
 980            1,
 981            "New Conversation",
 982            Some(Box::new(NewConversation)),
 983            tooltip_style,
 984            cx,
 985        )
 986    }
 987
 988    fn render_zoom_button(&self, cx: &mut ViewContext<Self>) -> impl Element<Self> {
 989        enum ToggleZoomButton {}
 990
 991        let theme = theme::current(cx);
 992        let tooltip_style = theme::current(cx).tooltip.clone();
 993        let style = if self.zoomed {
 994            &theme.assistant.zoom_out_button
 995        } else {
 996            &theme.assistant.zoom_in_button
 997        };
 998
 999        MouseEventHandler::new::<ToggleZoomButton, _>(0, cx, |state, _| {
1000            let style = style.style_for(state);
1001            Svg::for_style(style.icon.clone())
1002                .contained()
1003                .with_style(style.container)
1004        })
1005        .with_cursor_style(CursorStyle::PointingHand)
1006        .on_click(MouseButton::Left, |_, this, cx| {
1007            this.toggle_zoom(&ToggleZoom, cx);
1008        })
1009        .with_tooltip::<ToggleZoom>(
1010            0,
1011            if self.zoomed { "Zoom Out" } else { "Zoom In" },
1012            Some(Box::new(ToggleZoom)),
1013            tooltip_style,
1014            cx,
1015        )
1016    }
1017
1018    fn render_saved_conversation(
1019        &mut self,
1020        index: usize,
1021        cx: &mut ViewContext<Self>,
1022    ) -> impl Element<Self> {
1023        let conversation = &self.saved_conversations[index];
1024        let path = conversation.path.clone();
1025        MouseEventHandler::new::<SavedConversationMetadata, _>(index, cx, move |state, cx| {
1026            let style = &theme::current(cx).assistant.saved_conversation;
1027            Flex::row()
1028                .with_child(
1029                    Label::new(
1030                        conversation.mtime.format("%F %I:%M%p").to_string(),
1031                        style.saved_at.text.clone(),
1032                    )
1033                    .aligned()
1034                    .contained()
1035                    .with_style(style.saved_at.container),
1036                )
1037                .with_child(
1038                    Label::new(conversation.title.clone(), style.title.text.clone())
1039                        .aligned()
1040                        .contained()
1041                        .with_style(style.title.container),
1042                )
1043                .contained()
1044                .with_style(*style.container.style_for(state))
1045        })
1046        .with_cursor_style(CursorStyle::PointingHand)
1047        .on_click(MouseButton::Left, move |_, this, cx| {
1048            this.open_conversation(path.clone(), cx)
1049                .detach_and_log_err(cx)
1050        })
1051    }
1052
1053    fn open_conversation(&mut self, path: PathBuf, cx: &mut ViewContext<Self>) -> Task<Result<()>> {
1054        if let Some(ix) = self.editor_index_for_path(&path, cx) {
1055            self.set_active_editor_index(Some(ix), cx);
1056            return Task::ready(Ok(()));
1057        }
1058
1059        let fs = self.fs.clone();
1060        let api_key = self.api_key.clone();
1061        let languages = self.languages.clone();
1062        cx.spawn(|this, mut cx| async move {
1063            let saved_conversation = fs.load(&path).await?;
1064            let saved_conversation = serde_json::from_str(&saved_conversation)?;
1065            let conversation = cx.add_model(|cx| {
1066                Conversation::deserialize(saved_conversation, path.clone(), api_key, languages, cx)
1067            });
1068            this.update(&mut cx, |this, cx| {
1069                // If, by the time we've loaded the conversation, the user has already opened
1070                // the same conversation, we don't want to open it again.
1071                if let Some(ix) = this.editor_index_for_path(&path, cx) {
1072                    this.set_active_editor_index(Some(ix), cx);
1073                } else {
1074                    let editor = cx
1075                        .add_view(|cx| ConversationEditor::for_conversation(conversation, fs, cx));
1076                    this.add_conversation(editor, cx);
1077                }
1078            })?;
1079            Ok(())
1080        })
1081    }
1082
1083    fn editor_index_for_path(&self, path: &Path, cx: &AppContext) -> Option<usize> {
1084        self.editors
1085            .iter()
1086            .position(|editor| editor.read(cx).conversation.read(cx).path.as_deref() == Some(path))
1087    }
1088
1089    fn load_api_key(&mut self, cx: &mut ViewContext<Self>) -> Option<String> {
1090        if self.api_key.borrow().is_none() && !self.has_read_credentials {
1091            self.has_read_credentials = true;
1092            let api_key = if let Ok(api_key) = env::var("OPENAI_API_KEY") {
1093                Some(api_key)
1094            } else if let Some((_, api_key)) = cx
1095                .platform()
1096                .read_credentials(OPENAI_API_URL)
1097                .log_err()
1098                .flatten()
1099            {
1100                String::from_utf8(api_key).log_err()
1101            } else {
1102                None
1103            };
1104            if let Some(api_key) = api_key {
1105                *self.api_key.borrow_mut() = Some(api_key);
1106            } else if self.api_key_editor.is_none() {
1107                self.api_key_editor = Some(build_api_key_editor(cx));
1108                cx.notify();
1109            }
1110        }
1111
1112        self.api_key.borrow().clone()
1113    }
1114}
1115
1116fn build_api_key_editor(cx: &mut ViewContext<AssistantPanel>) -> ViewHandle<Editor> {
1117    cx.add_view(|cx| {
1118        let mut editor = Editor::single_line(
1119            Some(Arc::new(|theme| theme.assistant.api_key_editor.clone())),
1120            cx,
1121        );
1122        editor.set_placeholder_text("sk-000000000000000000000000000000000000000000000000", cx);
1123        editor
1124    })
1125}
1126
1127impl Entity for AssistantPanel {
1128    type Event = AssistantPanelEvent;
1129}
1130
1131impl View for AssistantPanel {
1132    fn ui_name() -> &'static str {
1133        "AssistantPanel"
1134    }
1135
1136    fn render(&mut self, cx: &mut ViewContext<Self>) -> AnyElement<Self> {
1137        let theme = &theme::current(cx);
1138        let style = &theme.assistant;
1139        if let Some(api_key_editor) = self.api_key_editor.as_ref() {
1140            Flex::column()
1141                .with_child(
1142                    Text::new(
1143                        "Paste your OpenAI API key and press Enter to use the assistant",
1144                        style.api_key_prompt.text.clone(),
1145                    )
1146                    .aligned(),
1147                )
1148                .with_child(
1149                    ChildView::new(api_key_editor, cx)
1150                        .contained()
1151                        .with_style(style.api_key_editor.container)
1152                        .aligned(),
1153                )
1154                .contained()
1155                .with_style(style.api_key_prompt.container)
1156                .aligned()
1157                .into_any()
1158        } else {
1159            let title = self.active_editor().map(|editor| {
1160                Label::new(editor.read(cx).title(cx), style.title.text.clone())
1161                    .contained()
1162                    .with_style(style.title.container)
1163                    .aligned()
1164                    .left()
1165                    .flex(1., false)
1166            });
1167            let mut header = Flex::row()
1168                .with_child(Self::render_hamburger_button(cx).aligned())
1169                .with_children(title);
1170            if self.has_focus {
1171                header.add_children(
1172                    self.render_editor_tools(cx)
1173                        .into_iter()
1174                        .map(|tool| tool.aligned().flex_float()),
1175                );
1176                header.add_child(Self::render_plus_button(cx).aligned().flex_float());
1177                header.add_child(self.render_zoom_button(cx).aligned());
1178            }
1179
1180            Flex::column()
1181                .with_child(
1182                    header
1183                        .contained()
1184                        .with_style(theme.workspace.tab_bar.container)
1185                        .expanded()
1186                        .constrained()
1187                        .with_height(theme.workspace.tab_bar.height),
1188                )
1189                .with_children(if self.toolbar.read(cx).hidden() {
1190                    None
1191                } else {
1192                    Some(ChildView::new(&self.toolbar, cx).expanded())
1193                })
1194                .with_child(if let Some(editor) = self.active_editor() {
1195                    ChildView::new(editor, cx).flex(1., true).into_any()
1196                } else {
1197                    UniformList::new(
1198                        self.saved_conversations_list_state.clone(),
1199                        self.saved_conversations.len(),
1200                        cx,
1201                        |this, range, items, cx| {
1202                            for ix in range {
1203                                items.push(this.render_saved_conversation(ix, cx).into_any());
1204                            }
1205                        },
1206                    )
1207                    .flex(1., true)
1208                    .into_any()
1209                })
1210                .into_any()
1211        }
1212    }
1213
1214    fn focus_in(&mut self, _: gpui::AnyViewHandle, cx: &mut ViewContext<Self>) {
1215        self.has_focus = true;
1216        self.toolbar
1217            .update(cx, |toolbar, cx| toolbar.focus_changed(true, cx));
1218        cx.notify();
1219        if cx.is_self_focused() {
1220            if let Some(editor) = self.active_editor() {
1221                cx.focus(editor);
1222            } else if let Some(api_key_editor) = self.api_key_editor.as_ref() {
1223                cx.focus(api_key_editor);
1224            }
1225        }
1226    }
1227
1228    fn focus_out(&mut self, _: gpui::AnyViewHandle, cx: &mut ViewContext<Self>) {
1229        self.has_focus = false;
1230        self.toolbar
1231            .update(cx, |toolbar, cx| toolbar.focus_changed(false, cx));
1232        cx.notify();
1233    }
1234}
1235
1236impl Panel for AssistantPanel {
1237    fn position(&self, cx: &WindowContext) -> DockPosition {
1238        match settings::get::<AssistantSettings>(cx).dock {
1239            AssistantDockPosition::Left => DockPosition::Left,
1240            AssistantDockPosition::Bottom => DockPosition::Bottom,
1241            AssistantDockPosition::Right => DockPosition::Right,
1242        }
1243    }
1244
1245    fn position_is_valid(&self, _: DockPosition) -> bool {
1246        true
1247    }
1248
1249    fn set_position(&mut self, position: DockPosition, cx: &mut ViewContext<Self>) {
1250        settings::update_settings_file::<AssistantSettings>(self.fs.clone(), cx, move |settings| {
1251            let dock = match position {
1252                DockPosition::Left => AssistantDockPosition::Left,
1253                DockPosition::Bottom => AssistantDockPosition::Bottom,
1254                DockPosition::Right => AssistantDockPosition::Right,
1255            };
1256            settings.dock = Some(dock);
1257        });
1258    }
1259
1260    fn size(&self, cx: &WindowContext) -> f32 {
1261        let settings = settings::get::<AssistantSettings>(cx);
1262        match self.position(cx) {
1263            DockPosition::Left | DockPosition::Right => {
1264                self.width.unwrap_or_else(|| settings.default_width)
1265            }
1266            DockPosition::Bottom => self.height.unwrap_or_else(|| settings.default_height),
1267        }
1268    }
1269
1270    fn set_size(&mut self, size: Option<f32>, cx: &mut ViewContext<Self>) {
1271        match self.position(cx) {
1272            DockPosition::Left | DockPosition::Right => self.width = size,
1273            DockPosition::Bottom => self.height = size,
1274        }
1275        cx.notify();
1276    }
1277
1278    fn should_zoom_in_on_event(event: &AssistantPanelEvent) -> bool {
1279        matches!(event, AssistantPanelEvent::ZoomIn)
1280    }
1281
1282    fn should_zoom_out_on_event(event: &AssistantPanelEvent) -> bool {
1283        matches!(event, AssistantPanelEvent::ZoomOut)
1284    }
1285
1286    fn is_zoomed(&self, _: &WindowContext) -> bool {
1287        self.zoomed
1288    }
1289
1290    fn set_zoomed(&mut self, zoomed: bool, cx: &mut ViewContext<Self>) {
1291        self.zoomed = zoomed;
1292        cx.notify();
1293    }
1294
1295    fn set_active(&mut self, active: bool, cx: &mut ViewContext<Self>) {
1296        if active {
1297            self.load_api_key(cx);
1298
1299            if self.editors.is_empty() {
1300                self.new_conversation(cx);
1301            }
1302        }
1303    }
1304
1305    fn icon_path(&self, cx: &WindowContext) -> Option<&'static str> {
1306        settings::get::<AssistantSettings>(cx)
1307            .button
1308            .then(|| "icons/ai.svg")
1309    }
1310
1311    fn icon_tooltip(&self) -> (String, Option<Box<dyn Action>>) {
1312        ("Assistant Panel".into(), Some(Box::new(ToggleFocus)))
1313    }
1314
1315    fn should_change_position_on_event(event: &Self::Event) -> bool {
1316        matches!(event, AssistantPanelEvent::DockPositionChanged)
1317    }
1318
1319    fn should_activate_on_event(_: &Self::Event) -> bool {
1320        false
1321    }
1322
1323    fn should_close_on_event(event: &AssistantPanelEvent) -> bool {
1324        matches!(event, AssistantPanelEvent::Close)
1325    }
1326
1327    fn has_focus(&self, _: &WindowContext) -> bool {
1328        self.has_focus
1329    }
1330
1331    fn is_focus_event(event: &Self::Event) -> bool {
1332        matches!(event, AssistantPanelEvent::Focus)
1333    }
1334}
1335
1336enum ConversationEvent {
1337    MessagesEdited,
1338    SummaryChanged,
1339    StreamedCompletion,
1340}
1341
1342#[derive(Default)]
1343struct Summary {
1344    text: String,
1345    done: bool,
1346}
1347
1348struct Conversation {
1349    buffer: ModelHandle<Buffer>,
1350    message_anchors: Vec<MessageAnchor>,
1351    messages_metadata: HashMap<MessageId, MessageMetadata>,
1352    next_message_id: MessageId,
1353    summary: Option<Summary>,
1354    pending_summary: Task<Option<()>>,
1355    completion_count: usize,
1356    pending_completions: Vec<PendingCompletion>,
1357    model: OpenAIModel,
1358    token_count: Option<usize>,
1359    max_token_count: usize,
1360    pending_token_count: Task<Option<()>>,
1361    api_key: Rc<RefCell<Option<String>>>,
1362    pending_save: Task<Result<()>>,
1363    path: Option<PathBuf>,
1364    _subscriptions: Vec<Subscription>,
1365}
1366
1367impl Entity for Conversation {
1368    type Event = ConversationEvent;
1369}
1370
1371impl Conversation {
1372    fn new(
1373        api_key: Rc<RefCell<Option<String>>>,
1374        language_registry: Arc<LanguageRegistry>,
1375        cx: &mut ModelContext<Self>,
1376    ) -> Self {
1377        let markdown = language_registry.language_for_name("Markdown");
1378        let buffer = cx.add_model(|cx| {
1379            let mut buffer = Buffer::new(0, cx.model_id() as u64, "");
1380            buffer.set_language_registry(language_registry);
1381            cx.spawn_weak(|buffer, mut cx| async move {
1382                let markdown = markdown.await?;
1383                let buffer = buffer
1384                    .upgrade(&cx)
1385                    .ok_or_else(|| anyhow!("buffer was dropped"))?;
1386                buffer.update(&mut cx, |buffer: &mut Buffer, cx| {
1387                    buffer.set_language(Some(markdown), cx)
1388                });
1389                anyhow::Ok(())
1390            })
1391            .detach_and_log_err(cx);
1392            buffer
1393        });
1394
1395        let settings = settings::get::<AssistantSettings>(cx);
1396        let model = settings.default_open_ai_model.clone();
1397
1398        let mut this = Self {
1399            message_anchors: Default::default(),
1400            messages_metadata: Default::default(),
1401            next_message_id: Default::default(),
1402            summary: None,
1403            pending_summary: Task::ready(None),
1404            completion_count: Default::default(),
1405            pending_completions: Default::default(),
1406            token_count: None,
1407            max_token_count: tiktoken_rs::model::get_context_size(&model.full_name()),
1408            pending_token_count: Task::ready(None),
1409            model: model.clone(),
1410            _subscriptions: vec![cx.subscribe(&buffer, Self::handle_buffer_event)],
1411            pending_save: Task::ready(Ok(())),
1412            path: None,
1413            api_key,
1414            buffer,
1415        };
1416        let message = MessageAnchor {
1417            id: MessageId(post_inc(&mut this.next_message_id.0)),
1418            start: language::Anchor::MIN,
1419        };
1420        this.message_anchors.push(message.clone());
1421        this.messages_metadata.insert(
1422            message.id,
1423            MessageMetadata {
1424                role: Role::User,
1425                sent_at: Local::now(),
1426                status: MessageStatus::Done,
1427            },
1428        );
1429
1430        this.count_remaining_tokens(cx);
1431        this
1432    }
1433
1434    fn serialize(&self, cx: &AppContext) -> SavedConversation {
1435        SavedConversation {
1436            zed: "conversation".into(),
1437            version: SavedConversation::VERSION.into(),
1438            text: self.buffer.read(cx).text(),
1439            message_metadata: self.messages_metadata.clone(),
1440            messages: self
1441                .messages(cx)
1442                .map(|message| SavedMessage {
1443                    id: message.id,
1444                    start: message.offset_range.start,
1445                })
1446                .collect(),
1447            summary: self
1448                .summary
1449                .as_ref()
1450                .map(|summary| summary.text.clone())
1451                .unwrap_or_default(),
1452            model: self.model.clone(),
1453        }
1454    }
1455
1456    fn deserialize(
1457        saved_conversation: SavedConversation,
1458        path: PathBuf,
1459        api_key: Rc<RefCell<Option<String>>>,
1460        language_registry: Arc<LanguageRegistry>,
1461        cx: &mut ModelContext<Self>,
1462    ) -> Self {
1463        let model = saved_conversation.model;
1464        let markdown = language_registry.language_for_name("Markdown");
1465        let mut message_anchors = Vec::new();
1466        let mut next_message_id = MessageId(0);
1467        let buffer = cx.add_model(|cx| {
1468            let mut buffer = Buffer::new(0, cx.model_id() as u64, saved_conversation.text);
1469            for message in saved_conversation.messages {
1470                message_anchors.push(MessageAnchor {
1471                    id: message.id,
1472                    start: buffer.anchor_before(message.start),
1473                });
1474                next_message_id = cmp::max(next_message_id, MessageId(message.id.0 + 1));
1475            }
1476            buffer.set_language_registry(language_registry);
1477            cx.spawn_weak(|buffer, mut cx| async move {
1478                let markdown = markdown.await?;
1479                let buffer = buffer
1480                    .upgrade(&cx)
1481                    .ok_or_else(|| anyhow!("buffer was dropped"))?;
1482                buffer.update(&mut cx, |buffer: &mut Buffer, cx| {
1483                    buffer.set_language(Some(markdown), cx)
1484                });
1485                anyhow::Ok(())
1486            })
1487            .detach_and_log_err(cx);
1488            buffer
1489        });
1490
1491        let mut this = Self {
1492            message_anchors,
1493            messages_metadata: saved_conversation.message_metadata,
1494            next_message_id,
1495            summary: Some(Summary {
1496                text: saved_conversation.summary,
1497                done: true,
1498            }),
1499            pending_summary: Task::ready(None),
1500            completion_count: Default::default(),
1501            pending_completions: Default::default(),
1502            token_count: None,
1503            max_token_count: tiktoken_rs::model::get_context_size(&model.full_name()),
1504            pending_token_count: Task::ready(None),
1505            model,
1506            _subscriptions: vec![cx.subscribe(&buffer, Self::handle_buffer_event)],
1507            pending_save: Task::ready(Ok(())),
1508            path: Some(path),
1509            api_key,
1510            buffer,
1511        };
1512        this.count_remaining_tokens(cx);
1513        this
1514    }
1515
1516    fn handle_buffer_event(
1517        &mut self,
1518        _: ModelHandle<Buffer>,
1519        event: &language::Event,
1520        cx: &mut ModelContext<Self>,
1521    ) {
1522        match event {
1523            language::Event::Edited => {
1524                self.count_remaining_tokens(cx);
1525                cx.emit(ConversationEvent::MessagesEdited);
1526            }
1527            _ => {}
1528        }
1529    }
1530
1531    fn count_remaining_tokens(&mut self, cx: &mut ModelContext<Self>) {
1532        let messages = self
1533            .messages(cx)
1534            .into_iter()
1535            .filter_map(|message| {
1536                Some(tiktoken_rs::ChatCompletionRequestMessage {
1537                    role: match message.role {
1538                        Role::User => "user".into(),
1539                        Role::Assistant => "assistant".into(),
1540                        Role::System => "system".into(),
1541                    },
1542                    content: self
1543                        .buffer
1544                        .read(cx)
1545                        .text_for_range(message.offset_range)
1546                        .collect(),
1547                    name: None,
1548                })
1549            })
1550            .collect::<Vec<_>>();
1551        let model = self.model.clone();
1552        self.pending_token_count = cx.spawn_weak(|this, mut cx| {
1553            async move {
1554                cx.background().timer(Duration::from_millis(200)).await;
1555                let token_count = cx
1556                    .background()
1557                    .spawn(async move {
1558                        tiktoken_rs::num_tokens_from_messages(&model.full_name(), &messages)
1559                    })
1560                    .await?;
1561
1562                this.upgrade(&cx)
1563                    .ok_or_else(|| anyhow!("conversation was dropped"))?
1564                    .update(&mut cx, |this, cx| {
1565                        this.max_token_count =
1566                            tiktoken_rs::model::get_context_size(&this.model.full_name());
1567                        this.token_count = Some(token_count);
1568                        cx.notify()
1569                    });
1570                anyhow::Ok(())
1571            }
1572            .log_err()
1573        });
1574    }
1575
1576    fn remaining_tokens(&self) -> Option<isize> {
1577        Some(self.max_token_count as isize - self.token_count? as isize)
1578    }
1579
1580    fn set_model(&mut self, model: OpenAIModel, cx: &mut ModelContext<Self>) {
1581        self.model = model;
1582        self.count_remaining_tokens(cx);
1583        cx.notify();
1584    }
1585
1586    fn assist(
1587        &mut self,
1588        selected_messages: HashSet<MessageId>,
1589        cx: &mut ModelContext<Self>,
1590    ) -> Vec<MessageAnchor> {
1591        let mut user_messages = Vec::new();
1592
1593        let last_message_id = if let Some(last_message_id) =
1594            self.message_anchors.iter().rev().find_map(|message| {
1595                message
1596                    .start
1597                    .is_valid(self.buffer.read(cx))
1598                    .then_some(message.id)
1599            }) {
1600            last_message_id
1601        } else {
1602            return Default::default();
1603        };
1604
1605        let mut should_assist = false;
1606        for selected_message_id in selected_messages {
1607            let selected_message_role =
1608                if let Some(metadata) = self.messages_metadata.get(&selected_message_id) {
1609                    metadata.role
1610                } else {
1611                    continue;
1612                };
1613
1614            if selected_message_role == Role::Assistant {
1615                if let Some(user_message) = self.insert_message_after(
1616                    selected_message_id,
1617                    Role::User,
1618                    MessageStatus::Done,
1619                    cx,
1620                ) {
1621                    user_messages.push(user_message);
1622                }
1623            } else {
1624                should_assist = true;
1625            }
1626        }
1627
1628        if should_assist {
1629            let Some(api_key) = self.api_key.borrow().clone() else {
1630                return Default::default();
1631            };
1632
1633            let request = OpenAIRequest {
1634                model: self.model.full_name().to_string(),
1635                messages: self
1636                    .messages(cx)
1637                    .filter(|message| matches!(message.status, MessageStatus::Done))
1638                    .map(|message| message.to_open_ai_message(self.buffer.read(cx)))
1639                    .collect(),
1640                stream: true,
1641            };
1642
1643            let stream = stream_completion(api_key, cx.background().clone(), request);
1644            let assistant_message = self
1645                .insert_message_after(last_message_id, Role::Assistant, MessageStatus::Pending, cx)
1646                .unwrap();
1647
1648            // Queue up the user's next reply.
1649            let user_message = self
1650                .insert_message_after(assistant_message.id, Role::User, MessageStatus::Done, cx)
1651                .unwrap();
1652            user_messages.push(user_message);
1653
1654            let task = cx.spawn_weak({
1655                |this, mut cx| async move {
1656                    let assistant_message_id = assistant_message.id;
1657                    let stream_completion = async {
1658                        let mut messages = stream.await?;
1659
1660                        while let Some(message) = messages.next().await {
1661                            let mut message = message?;
1662                            if let Some(choice) = message.choices.pop() {
1663                                this.upgrade(&cx)
1664                                    .ok_or_else(|| anyhow!("conversation was dropped"))?
1665                                    .update(&mut cx, |this, cx| {
1666                                        let text: Arc<str> = choice.delta.content?.into();
1667                                        let message_ix =
1668                                            this.message_anchors.iter().position(|message| {
1669                                                message.id == assistant_message_id
1670                                            })?;
1671                                        this.buffer.update(cx, |buffer, cx| {
1672                                            let offset = this.message_anchors[message_ix + 1..]
1673                                                .iter()
1674                                                .find(|message| message.start.is_valid(buffer))
1675                                                .map_or(buffer.len(), |message| {
1676                                                    message
1677                                                        .start
1678                                                        .to_offset(buffer)
1679                                                        .saturating_sub(1)
1680                                                });
1681                                            buffer.edit([(offset..offset, text)], None, cx);
1682                                        });
1683                                        cx.emit(ConversationEvent::StreamedCompletion);
1684
1685                                        Some(())
1686                                    });
1687                            }
1688                            smol::future::yield_now().await;
1689                        }
1690
1691                        this.upgrade(&cx)
1692                            .ok_or_else(|| anyhow!("conversation was dropped"))?
1693                            .update(&mut cx, |this, cx| {
1694                                this.pending_completions
1695                                    .retain(|completion| completion.id != this.completion_count);
1696                                this.summarize(cx);
1697                            });
1698
1699                        anyhow::Ok(())
1700                    };
1701
1702                    let result = stream_completion.await;
1703                    if let Some(this) = this.upgrade(&cx) {
1704                        this.update(&mut cx, |this, cx| {
1705                            if let Some(metadata) =
1706                                this.messages_metadata.get_mut(&assistant_message.id)
1707                            {
1708                                match result {
1709                                    Ok(_) => {
1710                                        metadata.status = MessageStatus::Done;
1711                                    }
1712                                    Err(error) => {
1713                                        metadata.status =
1714                                            MessageStatus::Error(error.to_string().trim().into());
1715                                    }
1716                                }
1717                                cx.notify();
1718                            }
1719                        });
1720                    }
1721                }
1722            });
1723
1724            self.pending_completions.push(PendingCompletion {
1725                id: post_inc(&mut self.completion_count),
1726                _task: task,
1727            });
1728        }
1729
1730        user_messages
1731    }
1732
1733    fn cancel_last_assist(&mut self) -> bool {
1734        self.pending_completions.pop().is_some()
1735    }
1736
1737    fn cycle_message_roles(&mut self, ids: HashSet<MessageId>, cx: &mut ModelContext<Self>) {
1738        for id in ids {
1739            if let Some(metadata) = self.messages_metadata.get_mut(&id) {
1740                metadata.role.cycle();
1741                cx.emit(ConversationEvent::MessagesEdited);
1742                cx.notify();
1743            }
1744        }
1745    }
1746
1747    fn insert_message_after(
1748        &mut self,
1749        message_id: MessageId,
1750        role: Role,
1751        status: MessageStatus,
1752        cx: &mut ModelContext<Self>,
1753    ) -> Option<MessageAnchor> {
1754        if let Some(prev_message_ix) = self
1755            .message_anchors
1756            .iter()
1757            .position(|message| message.id == message_id)
1758        {
1759            // Find the next valid message after the one we were given.
1760            let mut next_message_ix = prev_message_ix + 1;
1761            while let Some(next_message) = self.message_anchors.get(next_message_ix) {
1762                if next_message.start.is_valid(self.buffer.read(cx)) {
1763                    break;
1764                }
1765                next_message_ix += 1;
1766            }
1767
1768            let start = self.buffer.update(cx, |buffer, cx| {
1769                let offset = self
1770                    .message_anchors
1771                    .get(next_message_ix)
1772                    .map_or(buffer.len(), |message| message.start.to_offset(buffer) - 1);
1773                buffer.edit([(offset..offset, "\n")], None, cx);
1774                buffer.anchor_before(offset + 1)
1775            });
1776            let message = MessageAnchor {
1777                id: MessageId(post_inc(&mut self.next_message_id.0)),
1778                start,
1779            };
1780            self.message_anchors
1781                .insert(next_message_ix, message.clone());
1782            self.messages_metadata.insert(
1783                message.id,
1784                MessageMetadata {
1785                    role,
1786                    sent_at: Local::now(),
1787                    status,
1788                },
1789            );
1790            cx.emit(ConversationEvent::MessagesEdited);
1791            Some(message)
1792        } else {
1793            None
1794        }
1795    }
1796
1797    fn split_message(
1798        &mut self,
1799        range: Range<usize>,
1800        cx: &mut ModelContext<Self>,
1801    ) -> (Option<MessageAnchor>, Option<MessageAnchor>) {
1802        let start_message = self.message_for_offset(range.start, cx);
1803        let end_message = self.message_for_offset(range.end, cx);
1804        if let Some((start_message, end_message)) = start_message.zip(end_message) {
1805            // Prevent splitting when range spans multiple messages.
1806            if start_message.id != end_message.id {
1807                return (None, None);
1808            }
1809
1810            let message = start_message;
1811            let role = message.role;
1812            let mut edited_buffer = false;
1813
1814            let mut suffix_start = None;
1815            if range.start > message.offset_range.start && range.end < message.offset_range.end - 1
1816            {
1817                if self.buffer.read(cx).chars_at(range.end).next() == Some('\n') {
1818                    suffix_start = Some(range.end + 1);
1819                } else if self.buffer.read(cx).reversed_chars_at(range.end).next() == Some('\n') {
1820                    suffix_start = Some(range.end);
1821                }
1822            }
1823
1824            let suffix = if let Some(suffix_start) = suffix_start {
1825                MessageAnchor {
1826                    id: MessageId(post_inc(&mut self.next_message_id.0)),
1827                    start: self.buffer.read(cx).anchor_before(suffix_start),
1828                }
1829            } else {
1830                self.buffer.update(cx, |buffer, cx| {
1831                    buffer.edit([(range.end..range.end, "\n")], None, cx);
1832                });
1833                edited_buffer = true;
1834                MessageAnchor {
1835                    id: MessageId(post_inc(&mut self.next_message_id.0)),
1836                    start: self.buffer.read(cx).anchor_before(range.end + 1),
1837                }
1838            };
1839
1840            self.message_anchors
1841                .insert(message.index_range.end + 1, suffix.clone());
1842            self.messages_metadata.insert(
1843                suffix.id,
1844                MessageMetadata {
1845                    role,
1846                    sent_at: Local::now(),
1847                    status: MessageStatus::Done,
1848                },
1849            );
1850
1851            let new_messages =
1852                if range.start == range.end || range.start == message.offset_range.start {
1853                    (None, Some(suffix))
1854                } else {
1855                    let mut prefix_end = None;
1856                    if range.start > message.offset_range.start
1857                        && range.end < message.offset_range.end - 1
1858                    {
1859                        if self.buffer.read(cx).chars_at(range.start).next() == Some('\n') {
1860                            prefix_end = Some(range.start + 1);
1861                        } else if self.buffer.read(cx).reversed_chars_at(range.start).next()
1862                            == Some('\n')
1863                        {
1864                            prefix_end = Some(range.start);
1865                        }
1866                    }
1867
1868                    let selection = if let Some(prefix_end) = prefix_end {
1869                        cx.emit(ConversationEvent::MessagesEdited);
1870                        MessageAnchor {
1871                            id: MessageId(post_inc(&mut self.next_message_id.0)),
1872                            start: self.buffer.read(cx).anchor_before(prefix_end),
1873                        }
1874                    } else {
1875                        self.buffer.update(cx, |buffer, cx| {
1876                            buffer.edit([(range.start..range.start, "\n")], None, cx)
1877                        });
1878                        edited_buffer = true;
1879                        MessageAnchor {
1880                            id: MessageId(post_inc(&mut self.next_message_id.0)),
1881                            start: self.buffer.read(cx).anchor_before(range.end + 1),
1882                        }
1883                    };
1884
1885                    self.message_anchors
1886                        .insert(message.index_range.end + 1, selection.clone());
1887                    self.messages_metadata.insert(
1888                        selection.id,
1889                        MessageMetadata {
1890                            role,
1891                            sent_at: Local::now(),
1892                            status: MessageStatus::Done,
1893                        },
1894                    );
1895                    (Some(selection), Some(suffix))
1896                };
1897
1898            if !edited_buffer {
1899                cx.emit(ConversationEvent::MessagesEdited);
1900            }
1901            new_messages
1902        } else {
1903            (None, None)
1904        }
1905    }
1906
1907    fn summarize(&mut self, cx: &mut ModelContext<Self>) {
1908        if self.message_anchors.len() >= 2 && self.summary.is_none() {
1909            let api_key = self.api_key.borrow().clone();
1910            if let Some(api_key) = api_key {
1911                let messages = self
1912                    .messages(cx)
1913                    .take(2)
1914                    .map(|message| message.to_open_ai_message(self.buffer.read(cx)))
1915                    .chain(Some(RequestMessage {
1916                        role: Role::User,
1917                        content:
1918                            "Summarize the conversation into a short title without punctuation"
1919                                .into(),
1920                    }));
1921                let request = OpenAIRequest {
1922                    model: self.model.full_name().to_string(),
1923                    messages: messages.collect(),
1924                    stream: true,
1925                };
1926
1927                let stream = stream_completion(api_key, cx.background().clone(), request);
1928                self.pending_summary = cx.spawn(|this, mut cx| {
1929                    async move {
1930                        let mut messages = stream.await?;
1931
1932                        while let Some(message) = messages.next().await {
1933                            let mut message = message?;
1934                            if let Some(choice) = message.choices.pop() {
1935                                let text = choice.delta.content.unwrap_or_default();
1936                                this.update(&mut cx, |this, cx| {
1937                                    this.summary
1938                                        .get_or_insert(Default::default())
1939                                        .text
1940                                        .push_str(&text);
1941                                    cx.emit(ConversationEvent::SummaryChanged);
1942                                });
1943                            }
1944                        }
1945
1946                        this.update(&mut cx, |this, cx| {
1947                            if let Some(summary) = this.summary.as_mut() {
1948                                summary.done = true;
1949                                cx.emit(ConversationEvent::SummaryChanged);
1950                            }
1951                        });
1952
1953                        anyhow::Ok(())
1954                    }
1955                    .log_err()
1956                });
1957            }
1958        }
1959    }
1960
1961    fn message_for_offset(&self, offset: usize, cx: &AppContext) -> Option<Message> {
1962        self.messages_for_offsets([offset], cx).pop()
1963    }
1964
1965    fn messages_for_offsets(
1966        &self,
1967        offsets: impl IntoIterator<Item = usize>,
1968        cx: &AppContext,
1969    ) -> Vec<Message> {
1970        let mut result = Vec::new();
1971
1972        let mut messages = self.messages(cx).peekable();
1973        let mut offsets = offsets.into_iter().peekable();
1974        let mut current_message = messages.next();
1975        while let Some(offset) = offsets.next() {
1976            // Locate the message that contains the offset.
1977            while current_message.as_ref().map_or(false, |message| {
1978                !message.offset_range.contains(&offset) && messages.peek().is_some()
1979            }) {
1980                current_message = messages.next();
1981            }
1982            let Some(message) = current_message.as_ref() else {
1983                break;
1984            };
1985
1986            // Skip offsets that are in the same message.
1987            while offsets.peek().map_or(false, |offset| {
1988                message.offset_range.contains(offset) || messages.peek().is_none()
1989            }) {
1990                offsets.next();
1991            }
1992
1993            result.push(message.clone());
1994        }
1995        result
1996    }
1997
1998    fn messages<'a>(&'a self, cx: &'a AppContext) -> impl 'a + Iterator<Item = Message> {
1999        let buffer = self.buffer.read(cx);
2000        let mut message_anchors = self.message_anchors.iter().enumerate().peekable();
2001        iter::from_fn(move || {
2002            while let Some((start_ix, message_anchor)) = message_anchors.next() {
2003                let metadata = self.messages_metadata.get(&message_anchor.id)?;
2004                let message_start = message_anchor.start.to_offset(buffer);
2005                let mut message_end = None;
2006                let mut end_ix = start_ix;
2007                while let Some((_, next_message)) = message_anchors.peek() {
2008                    if next_message.start.is_valid(buffer) {
2009                        message_end = Some(next_message.start);
2010                        break;
2011                    } else {
2012                        end_ix += 1;
2013                        message_anchors.next();
2014                    }
2015                }
2016                let message_end = message_end
2017                    .unwrap_or(language::Anchor::MAX)
2018                    .to_offset(buffer);
2019                return Some(Message {
2020                    index_range: start_ix..end_ix,
2021                    offset_range: message_start..message_end,
2022                    id: message_anchor.id,
2023                    anchor: message_anchor.start,
2024                    role: metadata.role,
2025                    sent_at: metadata.sent_at,
2026                    status: metadata.status.clone(),
2027                });
2028            }
2029            None
2030        })
2031    }
2032
2033    fn save(
2034        &mut self,
2035        debounce: Option<Duration>,
2036        fs: Arc<dyn Fs>,
2037        cx: &mut ModelContext<Conversation>,
2038    ) {
2039        self.pending_save = cx.spawn(|this, mut cx| async move {
2040            if let Some(debounce) = debounce {
2041                cx.background().timer(debounce).await;
2042            }
2043
2044            let (old_path, summary) = this.read_with(&cx, |this, _| {
2045                let path = this.path.clone();
2046                let summary = if let Some(summary) = this.summary.as_ref() {
2047                    if summary.done {
2048                        Some(summary.text.clone())
2049                    } else {
2050                        None
2051                    }
2052                } else {
2053                    None
2054                };
2055                (path, summary)
2056            });
2057
2058            if let Some(summary) = summary {
2059                let conversation = this.read_with(&cx, |this, cx| this.serialize(cx));
2060                let path = if let Some(old_path) = old_path {
2061                    old_path
2062                } else {
2063                    let mut discriminant = 1;
2064                    let mut new_path;
2065                    loop {
2066                        new_path = CONVERSATIONS_DIR.join(&format!(
2067                            "{} - {}.zed.json",
2068                            summary.trim(),
2069                            discriminant
2070                        ));
2071                        if fs.is_file(&new_path).await {
2072                            discriminant += 1;
2073                        } else {
2074                            break;
2075                        }
2076                    }
2077                    new_path
2078                };
2079
2080                fs.create_dir(CONVERSATIONS_DIR.as_ref()).await?;
2081                fs.atomic_write(path.clone(), serde_json::to_string(&conversation).unwrap())
2082                    .await?;
2083                this.update(&mut cx, |this, _| this.path = Some(path));
2084            }
2085
2086            Ok(())
2087        });
2088    }
2089}
2090
2091struct PendingCompletion {
2092    id: usize,
2093    _task: Task<()>,
2094}
2095
2096enum ConversationEditorEvent {
2097    TabContentChanged,
2098}
2099
2100#[derive(Copy, Clone, Debug, PartialEq)]
2101struct ScrollPosition {
2102    offset_before_cursor: Vector2F,
2103    cursor: Anchor,
2104}
2105
2106struct ConversationEditor {
2107    conversation: ModelHandle<Conversation>,
2108    fs: Arc<dyn Fs>,
2109    editor: ViewHandle<Editor>,
2110    blocks: HashSet<BlockId>,
2111    scroll_position: Option<ScrollPosition>,
2112    _subscriptions: Vec<Subscription>,
2113}
2114
2115impl ConversationEditor {
2116    fn new(
2117        api_key: Rc<RefCell<Option<String>>>,
2118        language_registry: Arc<LanguageRegistry>,
2119        fs: Arc<dyn Fs>,
2120        cx: &mut ViewContext<Self>,
2121    ) -> Self {
2122        let conversation = cx.add_model(|cx| Conversation::new(api_key, language_registry, cx));
2123        Self::for_conversation(conversation, fs, cx)
2124    }
2125
2126    fn for_conversation(
2127        conversation: ModelHandle<Conversation>,
2128        fs: Arc<dyn Fs>,
2129        cx: &mut ViewContext<Self>,
2130    ) -> Self {
2131        let editor = cx.add_view(|cx| {
2132            let mut editor = Editor::for_buffer(conversation.read(cx).buffer.clone(), None, cx);
2133            editor.set_soft_wrap_mode(SoftWrap::EditorWidth, cx);
2134            editor.set_show_gutter(false, cx);
2135            editor.set_show_wrap_guides(false, cx);
2136            editor
2137        });
2138
2139        let _subscriptions = vec![
2140            cx.observe(&conversation, |_, _, cx| cx.notify()),
2141            cx.subscribe(&conversation, Self::handle_conversation_event),
2142            cx.subscribe(&editor, Self::handle_editor_event),
2143        ];
2144
2145        let mut this = Self {
2146            conversation,
2147            editor,
2148            blocks: Default::default(),
2149            scroll_position: None,
2150            fs,
2151            _subscriptions,
2152        };
2153        this.update_message_headers(cx);
2154        this
2155    }
2156
2157    fn assist(&mut self, _: &Assist, cx: &mut ViewContext<Self>) {
2158        let cursors = self.cursors(cx);
2159
2160        let user_messages = self.conversation.update(cx, |conversation, cx| {
2161            let selected_messages = conversation
2162                .messages_for_offsets(cursors, cx)
2163                .into_iter()
2164                .map(|message| message.id)
2165                .collect();
2166            conversation.assist(selected_messages, cx)
2167        });
2168        let new_selections = user_messages
2169            .iter()
2170            .map(|message| {
2171                let cursor = message
2172                    .start
2173                    .to_offset(self.conversation.read(cx).buffer.read(cx));
2174                cursor..cursor
2175            })
2176            .collect::<Vec<_>>();
2177        if !new_selections.is_empty() {
2178            self.editor.update(cx, |editor, cx| {
2179                editor.change_selections(
2180                    Some(Autoscroll::Strategy(AutoscrollStrategy::Fit)),
2181                    cx,
2182                    |selections| selections.select_ranges(new_selections),
2183                );
2184            });
2185            // Avoid scrolling to the new cursor position so the assistant's output is stable.
2186            cx.defer(|this, _| this.scroll_position = None);
2187        }
2188    }
2189
2190    fn cancel_last_assist(&mut self, _: &editor::Cancel, cx: &mut ViewContext<Self>) {
2191        if !self
2192            .conversation
2193            .update(cx, |conversation, _| conversation.cancel_last_assist())
2194        {
2195            cx.propagate_action();
2196        }
2197    }
2198
2199    fn cycle_message_role(&mut self, _: &CycleMessageRole, cx: &mut ViewContext<Self>) {
2200        let cursors = self.cursors(cx);
2201        self.conversation.update(cx, |conversation, cx| {
2202            let messages = conversation
2203                .messages_for_offsets(cursors, cx)
2204                .into_iter()
2205                .map(|message| message.id)
2206                .collect();
2207            conversation.cycle_message_roles(messages, cx)
2208        });
2209    }
2210
2211    fn cursors(&self, cx: &AppContext) -> Vec<usize> {
2212        let selections = self.editor.read(cx).selections.all::<usize>(cx);
2213        selections
2214            .into_iter()
2215            .map(|selection| selection.head())
2216            .collect()
2217    }
2218
2219    fn handle_conversation_event(
2220        &mut self,
2221        _: ModelHandle<Conversation>,
2222        event: &ConversationEvent,
2223        cx: &mut ViewContext<Self>,
2224    ) {
2225        match event {
2226            ConversationEvent::MessagesEdited => {
2227                self.update_message_headers(cx);
2228                self.conversation.update(cx, |conversation, cx| {
2229                    conversation.save(Some(Duration::from_millis(500)), self.fs.clone(), cx);
2230                });
2231            }
2232            ConversationEvent::SummaryChanged => {
2233                cx.emit(ConversationEditorEvent::TabContentChanged);
2234                self.conversation.update(cx, |conversation, cx| {
2235                    conversation.save(None, self.fs.clone(), cx);
2236                });
2237            }
2238            ConversationEvent::StreamedCompletion => {
2239                self.editor.update(cx, |editor, cx| {
2240                    if let Some(scroll_position) = self.scroll_position {
2241                        let snapshot = editor.snapshot(cx);
2242                        let cursor_point = scroll_position.cursor.to_display_point(&snapshot);
2243                        let scroll_top =
2244                            cursor_point.row() as f32 - scroll_position.offset_before_cursor.y();
2245                        editor.set_scroll_position(
2246                            vec2f(scroll_position.offset_before_cursor.x(), scroll_top),
2247                            cx,
2248                        );
2249                    }
2250                });
2251            }
2252        }
2253    }
2254
2255    fn handle_editor_event(
2256        &mut self,
2257        _: ViewHandle<Editor>,
2258        event: &editor::Event,
2259        cx: &mut ViewContext<Self>,
2260    ) {
2261        match event {
2262            editor::Event::ScrollPositionChanged { autoscroll, .. } => {
2263                let cursor_scroll_position = self.cursor_scroll_position(cx);
2264                if *autoscroll {
2265                    self.scroll_position = cursor_scroll_position;
2266                } else if self.scroll_position != cursor_scroll_position {
2267                    self.scroll_position = None;
2268                }
2269            }
2270            editor::Event::SelectionsChanged { .. } => {
2271                self.scroll_position = self.cursor_scroll_position(cx);
2272            }
2273            _ => {}
2274        }
2275    }
2276
2277    fn cursor_scroll_position(&self, cx: &mut ViewContext<Self>) -> Option<ScrollPosition> {
2278        self.editor.update(cx, |editor, cx| {
2279            let snapshot = editor.snapshot(cx);
2280            let cursor = editor.selections.newest_anchor().head();
2281            let cursor_row = cursor.to_display_point(&snapshot.display_snapshot).row() as f32;
2282            let scroll_position = editor
2283                .scroll_manager
2284                .anchor()
2285                .scroll_position(&snapshot.display_snapshot);
2286
2287            let scroll_bottom = scroll_position.y() + editor.visible_line_count().unwrap_or(0.);
2288            if (scroll_position.y()..scroll_bottom).contains(&cursor_row) {
2289                Some(ScrollPosition {
2290                    cursor,
2291                    offset_before_cursor: vec2f(
2292                        scroll_position.x(),
2293                        cursor_row - scroll_position.y(),
2294                    ),
2295                })
2296            } else {
2297                None
2298            }
2299        })
2300    }
2301
2302    fn update_message_headers(&mut self, cx: &mut ViewContext<Self>) {
2303        self.editor.update(cx, |editor, cx| {
2304            let buffer = editor.buffer().read(cx).snapshot(cx);
2305            let excerpt_id = *buffer.as_singleton().unwrap().0;
2306            let old_blocks = std::mem::take(&mut self.blocks);
2307            let new_blocks = self
2308                .conversation
2309                .read(cx)
2310                .messages(cx)
2311                .map(|message| BlockProperties {
2312                    position: buffer.anchor_in_excerpt(excerpt_id, message.anchor),
2313                    height: 2,
2314                    style: BlockStyle::Sticky,
2315                    render: Arc::new({
2316                        let conversation = self.conversation.clone();
2317                        // let metadata = message.metadata.clone();
2318                        // let message = message.clone();
2319                        move |cx| {
2320                            enum Sender {}
2321                            enum ErrorTooltip {}
2322
2323                            let theme = theme::current(cx);
2324                            let style = &theme.assistant;
2325                            let message_id = message.id;
2326                            let sender = MouseEventHandler::new::<Sender, _>(
2327                                message_id.0,
2328                                cx,
2329                                |state, _| match message.role {
2330                                    Role::User => {
2331                                        let style = style.user_sender.style_for(state);
2332                                        Label::new("You", style.text.clone())
2333                                            .contained()
2334                                            .with_style(style.container)
2335                                    }
2336                                    Role::Assistant => {
2337                                        let style = style.assistant_sender.style_for(state);
2338                                        Label::new("Assistant", style.text.clone())
2339                                            .contained()
2340                                            .with_style(style.container)
2341                                    }
2342                                    Role::System => {
2343                                        let style = style.system_sender.style_for(state);
2344                                        Label::new("System", style.text.clone())
2345                                            .contained()
2346                                            .with_style(style.container)
2347                                    }
2348                                },
2349                            )
2350                            .with_cursor_style(CursorStyle::PointingHand)
2351                            .on_down(MouseButton::Left, {
2352                                let conversation = conversation.clone();
2353                                move |_, _, cx| {
2354                                    conversation.update(cx, |conversation, cx| {
2355                                        conversation.cycle_message_roles(
2356                                            HashSet::from_iter(Some(message_id)),
2357                                            cx,
2358                                        )
2359                                    })
2360                                }
2361                            });
2362
2363                            Flex::row()
2364                                .with_child(sender.aligned())
2365                                .with_child(
2366                                    Label::new(
2367                                        message.sent_at.format("%I:%M%P").to_string(),
2368                                        style.sent_at.text.clone(),
2369                                    )
2370                                    .contained()
2371                                    .with_style(style.sent_at.container)
2372                                    .aligned(),
2373                                )
2374                                .with_children(
2375                                    if let MessageStatus::Error(error) = &message.status {
2376                                        Some(
2377                                            Svg::new("icons/circle_x_mark_12.svg")
2378                                                .with_color(style.error_icon.color)
2379                                                .constrained()
2380                                                .with_width(style.error_icon.width)
2381                                                .contained()
2382                                                .with_style(style.error_icon.container)
2383                                                .with_tooltip::<ErrorTooltip>(
2384                                                    message_id.0,
2385                                                    error.to_string(),
2386                                                    None,
2387                                                    theme.tooltip.clone(),
2388                                                    cx,
2389                                                )
2390                                                .aligned(),
2391                                        )
2392                                    } else {
2393                                        None
2394                                    },
2395                                )
2396                                .aligned()
2397                                .left()
2398                                .contained()
2399                                .with_style(style.message_header)
2400                                .into_any()
2401                        }
2402                    }),
2403                    disposition: BlockDisposition::Above,
2404                })
2405                .collect::<Vec<_>>();
2406
2407            editor.remove_blocks(old_blocks, None, cx);
2408            let ids = editor.insert_blocks(new_blocks, None, cx);
2409            self.blocks = HashSet::from_iter(ids);
2410        });
2411    }
2412
2413    fn quote_selection(
2414        workspace: &mut Workspace,
2415        _: &QuoteSelection,
2416        cx: &mut ViewContext<Workspace>,
2417    ) {
2418        let Some(panel) = workspace.panel::<AssistantPanel>(cx) else {
2419            return;
2420        };
2421        let Some(editor) = workspace
2422            .active_item(cx)
2423            .and_then(|item| item.act_as::<Editor>(cx))
2424        else {
2425            return;
2426        };
2427
2428        let text = editor.read_with(cx, |editor, cx| {
2429            let range = editor.selections.newest::<usize>(cx).range();
2430            let buffer = editor.buffer().read(cx).snapshot(cx);
2431            let start_language = buffer.language_at(range.start);
2432            let end_language = buffer.language_at(range.end);
2433            let language_name = if start_language == end_language {
2434                start_language.map(|language| language.name())
2435            } else {
2436                None
2437            };
2438            let language_name = language_name.as_deref().unwrap_or("").to_lowercase();
2439
2440            let selected_text = buffer.text_for_range(range).collect::<String>();
2441            if selected_text.is_empty() {
2442                None
2443            } else {
2444                Some(if language_name == "markdown" {
2445                    selected_text
2446                        .lines()
2447                        .map(|line| format!("> {}", line))
2448                        .collect::<Vec<_>>()
2449                        .join("\n")
2450                } else {
2451                    format!("```{language_name}\n{selected_text}\n```")
2452                })
2453            }
2454        });
2455
2456        // Activate the panel
2457        if !panel.read(cx).has_focus(cx) {
2458            workspace.toggle_panel_focus::<AssistantPanel>(cx);
2459        }
2460
2461        if let Some(text) = text {
2462            panel.update(cx, |panel, cx| {
2463                let conversation = panel
2464                    .active_editor()
2465                    .cloned()
2466                    .unwrap_or_else(|| panel.new_conversation(cx));
2467                conversation.update(cx, |conversation, cx| {
2468                    conversation
2469                        .editor
2470                        .update(cx, |editor, cx| editor.insert(&text, cx))
2471                });
2472            });
2473        }
2474    }
2475
2476    fn copy(&mut self, _: &editor::Copy, cx: &mut ViewContext<Self>) {
2477        let editor = self.editor.read(cx);
2478        let conversation = self.conversation.read(cx);
2479        if editor.selections.count() == 1 {
2480            let selection = editor.selections.newest::<usize>(cx);
2481            let mut copied_text = String::new();
2482            let mut spanned_messages = 0;
2483            for message in conversation.messages(cx) {
2484                if message.offset_range.start >= selection.range().end {
2485                    break;
2486                } else if message.offset_range.end >= selection.range().start {
2487                    let range = cmp::max(message.offset_range.start, selection.range().start)
2488                        ..cmp::min(message.offset_range.end, selection.range().end);
2489                    if !range.is_empty() {
2490                        spanned_messages += 1;
2491                        write!(&mut copied_text, "## {}\n\n", message.role).unwrap();
2492                        for chunk in conversation.buffer.read(cx).text_for_range(range) {
2493                            copied_text.push_str(&chunk);
2494                        }
2495                        copied_text.push('\n');
2496                    }
2497                }
2498            }
2499
2500            if spanned_messages > 1 {
2501                cx.platform()
2502                    .write_to_clipboard(ClipboardItem::new(copied_text));
2503                return;
2504            }
2505        }
2506
2507        cx.propagate_action();
2508    }
2509
2510    fn split(&mut self, _: &Split, cx: &mut ViewContext<Self>) {
2511        self.conversation.update(cx, |conversation, cx| {
2512            let selections = self.editor.read(cx).selections.disjoint_anchors();
2513            for selection in selections.into_iter() {
2514                let buffer = self.editor.read(cx).buffer().read(cx).snapshot(cx);
2515                let range = selection
2516                    .map(|endpoint| endpoint.to_offset(&buffer))
2517                    .range();
2518                conversation.split_message(range, cx);
2519            }
2520        });
2521    }
2522
2523    fn save(&mut self, _: &Save, cx: &mut ViewContext<Self>) {
2524        self.conversation.update(cx, |conversation, cx| {
2525            conversation.save(None, self.fs.clone(), cx)
2526        });
2527    }
2528
2529    fn cycle_model(&mut self, cx: &mut ViewContext<Self>) {
2530        self.conversation.update(cx, |conversation, cx| {
2531            let new_model = conversation.model.cycle();
2532            conversation.set_model(new_model, cx);
2533        });
2534    }
2535
2536    fn title(&self, cx: &AppContext) -> String {
2537        self.conversation
2538            .read(cx)
2539            .summary
2540            .as_ref()
2541            .map(|summary| summary.text.clone())
2542            .unwrap_or_else(|| "New Conversation".into())
2543    }
2544
2545    fn render_current_model(
2546        &self,
2547        style: &AssistantStyle,
2548        cx: &mut ViewContext<Self>,
2549    ) -> impl Element<Self> {
2550        enum Model {}
2551
2552        MouseEventHandler::new::<Model, _>(0, cx, |state, cx| {
2553            let style = style.model.style_for(state);
2554            let model_display_name = self.conversation.read(cx).model.short_name();
2555            Label::new(model_display_name, style.text.clone())
2556                .contained()
2557                .with_style(style.container)
2558        })
2559        .with_cursor_style(CursorStyle::PointingHand)
2560        .on_click(MouseButton::Left, |_, this, cx| this.cycle_model(cx))
2561    }
2562
2563    fn render_remaining_tokens(
2564        &self,
2565        style: &AssistantStyle,
2566        cx: &mut ViewContext<Self>,
2567    ) -> Option<impl Element<Self>> {
2568        let remaining_tokens = self.conversation.read(cx).remaining_tokens()?;
2569        let remaining_tokens_style = if remaining_tokens <= 0 {
2570            &style.no_remaining_tokens
2571        } else if remaining_tokens <= 500 {
2572            &style.low_remaining_tokens
2573        } else {
2574            &style.remaining_tokens
2575        };
2576        Some(
2577            Label::new(
2578                remaining_tokens.to_string(),
2579                remaining_tokens_style.text.clone(),
2580            )
2581            .contained()
2582            .with_style(remaining_tokens_style.container),
2583        )
2584    }
2585}
2586
2587impl Entity for ConversationEditor {
2588    type Event = ConversationEditorEvent;
2589}
2590
2591impl View for ConversationEditor {
2592    fn ui_name() -> &'static str {
2593        "ConversationEditor"
2594    }
2595
2596    fn render(&mut self, cx: &mut ViewContext<Self>) -> AnyElement<Self> {
2597        let theme = &theme::current(cx).assistant;
2598        Stack::new()
2599            .with_child(
2600                ChildView::new(&self.editor, cx)
2601                    .contained()
2602                    .with_style(theme.container),
2603            )
2604            .with_child(
2605                Flex::row()
2606                    .with_child(self.render_current_model(theme, cx))
2607                    .with_children(self.render_remaining_tokens(theme, cx))
2608                    .aligned()
2609                    .top()
2610                    .right(),
2611            )
2612            .into_any()
2613    }
2614
2615    fn focus_in(&mut self, _: gpui::AnyViewHandle, cx: &mut ViewContext<Self>) {
2616        if cx.is_self_focused() {
2617            cx.focus(&self.editor);
2618        }
2619    }
2620}
2621
2622#[derive(Clone, Debug)]
2623struct MessageAnchor {
2624    id: MessageId,
2625    start: language::Anchor,
2626}
2627
2628#[derive(Clone, Debug)]
2629pub struct Message {
2630    offset_range: Range<usize>,
2631    index_range: Range<usize>,
2632    id: MessageId,
2633    anchor: language::Anchor,
2634    role: Role,
2635    sent_at: DateTime<Local>,
2636    status: MessageStatus,
2637}
2638
2639impl Message {
2640    fn to_open_ai_message(&self, buffer: &Buffer) -> RequestMessage {
2641        let content = buffer
2642            .text_for_range(self.offset_range.clone())
2643            .collect::<String>();
2644        RequestMessage {
2645            role: self.role,
2646            content: content.trim_end().into(),
2647        }
2648    }
2649}
2650
2651enum InlineAssistantEvent {
2652    Confirmed {
2653        prompt: String,
2654        include_conversation: bool,
2655    },
2656    Canceled,
2657    Dismissed,
2658    IncludeConversationToggled {
2659        include_conversation: bool,
2660    },
2661}
2662
2663struct InlineAssistant {
2664    id: usize,
2665    prompt_editor: ViewHandle<Editor>,
2666    confirmed: bool,
2667    has_focus: bool,
2668    include_conversation: bool,
2669    measurements: Rc<Cell<BlockMeasurements>>,
2670    prompt_history: VecDeque<String>,
2671    prompt_history_ix: Option<usize>,
2672    pending_prompt: String,
2673    codegen: ModelHandle<Codegen>,
2674    _subscriptions: Vec<Subscription>,
2675}
2676
2677impl Entity for InlineAssistant {
2678    type Event = InlineAssistantEvent;
2679}
2680
2681impl View for InlineAssistant {
2682    fn ui_name() -> &'static str {
2683        "InlineAssistant"
2684    }
2685
2686    fn render(&mut self, cx: &mut ViewContext<Self>) -> AnyElement<Self> {
2687        enum ErrorIcon {}
2688        let theme = theme::current(cx);
2689
2690        Flex::row()
2691            .with_child(
2692                Flex::row()
2693                    .with_child(
2694                        Button::action(ToggleIncludeConversation)
2695                            .with_tooltip("Include Conversation", theme.tooltip.clone())
2696                            .with_id(self.id)
2697                            .with_contents(theme::components::svg::Svg::new("icons/ai.svg"))
2698                            .toggleable(self.include_conversation)
2699                            .with_style(theme.assistant.inline.include_conversation.clone())
2700                            .element()
2701                            .aligned(),
2702                    )
2703                    .with_children(if let Some(error) = self.codegen.read(cx).error() {
2704                        Some(
2705                            Svg::new("icons/circle_x_mark_12.svg")
2706                                .with_color(theme.assistant.error_icon.color)
2707                                .constrained()
2708                                .with_width(theme.assistant.error_icon.width)
2709                                .contained()
2710                                .with_style(theme.assistant.error_icon.container)
2711                                .with_tooltip::<ErrorIcon>(
2712                                    self.id,
2713                                    error.to_string(),
2714                                    None,
2715                                    theme.tooltip.clone(),
2716                                    cx,
2717                                )
2718                                .aligned(),
2719                        )
2720                    } else {
2721                        None
2722                    })
2723                    .aligned()
2724                    .constrained()
2725                    .dynamically({
2726                        let measurements = self.measurements.clone();
2727                        move |constraint, _, _| {
2728                            let measurements = measurements.get();
2729                            SizeConstraint {
2730                                min: vec2f(measurements.gutter_width, constraint.min.y()),
2731                                max: vec2f(measurements.gutter_width, constraint.max.y()),
2732                            }
2733                        }
2734                    }),
2735            )
2736            .with_child(Empty::new().constrained().dynamically({
2737                let measurements = self.measurements.clone();
2738                move |constraint, _, _| {
2739                    let measurements = measurements.get();
2740                    SizeConstraint {
2741                        min: vec2f(
2742                            measurements.anchor_x - measurements.gutter_width,
2743                            constraint.min.y(),
2744                        ),
2745                        max: vec2f(
2746                            measurements.anchor_x - measurements.gutter_width,
2747                            constraint.max.y(),
2748                        ),
2749                    }
2750                }
2751            }))
2752            .with_child(
2753                ChildView::new(&self.prompt_editor, cx)
2754                    .aligned()
2755                    .left()
2756                    .flex(1., true),
2757            )
2758            .contained()
2759            .with_style(theme.assistant.inline.container)
2760            .into_any()
2761            .into_any()
2762    }
2763
2764    fn focus_in(&mut self, _: gpui::AnyViewHandle, cx: &mut ViewContext<Self>) {
2765        cx.focus(&self.prompt_editor);
2766        self.has_focus = true;
2767    }
2768
2769    fn focus_out(&mut self, _: gpui::AnyViewHandle, _: &mut ViewContext<Self>) {
2770        self.has_focus = false;
2771    }
2772}
2773
2774impl InlineAssistant {
2775    fn new(
2776        id: usize,
2777        measurements: Rc<Cell<BlockMeasurements>>,
2778        include_conversation: bool,
2779        prompt_history: VecDeque<String>,
2780        codegen: ModelHandle<Codegen>,
2781        cx: &mut ViewContext<Self>,
2782    ) -> Self {
2783        let prompt_editor = cx.add_view(|cx| {
2784            let mut editor = Editor::single_line(
2785                Some(Arc::new(|theme| theme.assistant.inline.editor.clone())),
2786                cx,
2787            );
2788            let placeholder = match codegen.read(cx).kind() {
2789                CodegenKind::Transform { .. } => "Enter transformation prompt…",
2790                CodegenKind::Generate { .. } => "Enter generation prompt…",
2791            };
2792            editor.set_placeholder_text(placeholder, cx);
2793            editor
2794        });
2795        let subscriptions = vec![
2796            cx.observe(&codegen, Self::handle_codegen_changed),
2797            cx.subscribe(&prompt_editor, Self::handle_prompt_editor_events),
2798        ];
2799        Self {
2800            id,
2801            prompt_editor,
2802            confirmed: false,
2803            has_focus: false,
2804            include_conversation,
2805            measurements,
2806            prompt_history,
2807            prompt_history_ix: None,
2808            pending_prompt: String::new(),
2809            codegen,
2810            _subscriptions: subscriptions,
2811        }
2812    }
2813
2814    fn handle_prompt_editor_events(
2815        &mut self,
2816        _: ViewHandle<Editor>,
2817        event: &editor::Event,
2818        cx: &mut ViewContext<Self>,
2819    ) {
2820        if let editor::Event::Edited = event {
2821            self.pending_prompt = self.prompt_editor.read(cx).text(cx);
2822            cx.notify();
2823        }
2824    }
2825
2826    fn handle_codegen_changed(&mut self, _: ModelHandle<Codegen>, cx: &mut ViewContext<Self>) {
2827        let is_read_only = !self.codegen.read(cx).idle();
2828        self.prompt_editor.update(cx, |editor, cx| {
2829            let was_read_only = editor.read_only();
2830            if was_read_only != is_read_only {
2831                if is_read_only {
2832                    editor.set_read_only(true);
2833                    editor.set_field_editor_style(
2834                        Some(Arc::new(|theme| {
2835                            theme.assistant.inline.disabled_editor.clone()
2836                        })),
2837                        cx,
2838                    );
2839                } else {
2840                    editor.set_read_only(false);
2841                    editor.set_field_editor_style(
2842                        Some(Arc::new(|theme| theme.assistant.inline.editor.clone())),
2843                        cx,
2844                    );
2845                }
2846            }
2847        });
2848        cx.notify();
2849    }
2850
2851    fn cancel(&mut self, _: &editor::Cancel, cx: &mut ViewContext<Self>) {
2852        cx.emit(InlineAssistantEvent::Canceled);
2853    }
2854
2855    fn confirm(&mut self, _: &menu::Confirm, cx: &mut ViewContext<Self>) {
2856        if self.confirmed {
2857            cx.emit(InlineAssistantEvent::Dismissed);
2858        } else {
2859            let prompt = self.prompt_editor.read(cx).text(cx);
2860            self.prompt_editor.update(cx, |editor, cx| {
2861                editor.set_read_only(true);
2862                editor.set_field_editor_style(
2863                    Some(Arc::new(|theme| {
2864                        theme.assistant.inline.disabled_editor.clone()
2865                    })),
2866                    cx,
2867                );
2868            });
2869            cx.emit(InlineAssistantEvent::Confirmed {
2870                prompt,
2871                include_conversation: self.include_conversation,
2872            });
2873            self.confirmed = true;
2874            cx.notify();
2875        }
2876    }
2877
2878    fn toggle_include_conversation(
2879        &mut self,
2880        _: &ToggleIncludeConversation,
2881        cx: &mut ViewContext<Self>,
2882    ) {
2883        self.include_conversation = !self.include_conversation;
2884        cx.emit(InlineAssistantEvent::IncludeConversationToggled {
2885            include_conversation: self.include_conversation,
2886        });
2887        cx.notify();
2888    }
2889
2890    fn move_up(&mut self, _: &MoveUp, cx: &mut ViewContext<Self>) {
2891        if let Some(ix) = self.prompt_history_ix {
2892            if ix > 0 {
2893                self.prompt_history_ix = Some(ix - 1);
2894                let prompt = self.prompt_history[ix - 1].clone();
2895                self.set_prompt(&prompt, cx);
2896            }
2897        } else if !self.prompt_history.is_empty() {
2898            self.prompt_history_ix = Some(self.prompt_history.len() - 1);
2899            let prompt = self.prompt_history[self.prompt_history.len() - 1].clone();
2900            self.set_prompt(&prompt, cx);
2901        }
2902    }
2903
2904    fn move_down(&mut self, _: &MoveDown, cx: &mut ViewContext<Self>) {
2905        if let Some(ix) = self.prompt_history_ix {
2906            if ix < self.prompt_history.len() - 1 {
2907                self.prompt_history_ix = Some(ix + 1);
2908                let prompt = self.prompt_history[ix + 1].clone();
2909                self.set_prompt(&prompt, cx);
2910            } else {
2911                self.prompt_history_ix = None;
2912                let pending_prompt = self.pending_prompt.clone();
2913                self.set_prompt(&pending_prompt, cx);
2914            }
2915        }
2916    }
2917
2918    fn set_prompt(&mut self, prompt: &str, cx: &mut ViewContext<Self>) {
2919        self.prompt_editor.update(cx, |editor, cx| {
2920            editor.buffer().update(cx, |buffer, cx| {
2921                let len = buffer.len(cx);
2922                buffer.edit([(0..len, prompt)], None, cx);
2923            });
2924        });
2925    }
2926}
2927
2928// This wouldn't need to exist if we could pass parameters when rendering child views.
2929#[derive(Copy, Clone, Default)]
2930struct BlockMeasurements {
2931    anchor_x: f32,
2932    gutter_width: f32,
2933}
2934
2935struct PendingInlineAssist {
2936    editor: WeakViewHandle<Editor>,
2937    inline_assistant: Option<(BlockId, ViewHandle<InlineAssistant>)>,
2938    codegen: ModelHandle<Codegen>,
2939    _subscriptions: Vec<Subscription>,
2940}
2941
2942fn merge_ranges(ranges: &mut Vec<Range<Anchor>>, buffer: &MultiBufferSnapshot) {
2943    ranges.sort_unstable_by(|a, b| {
2944        a.start
2945            .cmp(&b.start, buffer)
2946            .then_with(|| b.end.cmp(&a.end, buffer))
2947    });
2948
2949    let mut ix = 0;
2950    while ix + 1 < ranges.len() {
2951        let b = ranges[ix + 1].clone();
2952        let a = &mut ranges[ix];
2953        if a.end.cmp(&b.start, buffer).is_gt() {
2954            if a.end.cmp(&b.end, buffer).is_lt() {
2955                a.end = b.end;
2956            }
2957            ranges.remove(ix + 1);
2958        } else {
2959            ix += 1;
2960        }
2961    }
2962}
2963
2964#[cfg(test)]
2965mod tests {
2966    use super::*;
2967    use crate::MessageId;
2968    use gpui::AppContext;
2969
2970    #[gpui::test]
2971    fn test_inserting_and_removing_messages(cx: &mut AppContext) {
2972        cx.set_global(SettingsStore::test(cx));
2973        init(cx);
2974        let registry = Arc::new(LanguageRegistry::test());
2975        let conversation = cx.add_model(|cx| Conversation::new(Default::default(), registry, cx));
2976        let buffer = conversation.read(cx).buffer.clone();
2977
2978        let message_1 = conversation.read(cx).message_anchors[0].clone();
2979        assert_eq!(
2980            messages(&conversation, cx),
2981            vec![(message_1.id, Role::User, 0..0)]
2982        );
2983
2984        let message_2 = conversation.update(cx, |conversation, cx| {
2985            conversation
2986                .insert_message_after(message_1.id, Role::Assistant, MessageStatus::Done, cx)
2987                .unwrap()
2988        });
2989        assert_eq!(
2990            messages(&conversation, cx),
2991            vec![
2992                (message_1.id, Role::User, 0..1),
2993                (message_2.id, Role::Assistant, 1..1)
2994            ]
2995        );
2996
2997        buffer.update(cx, |buffer, cx| {
2998            buffer.edit([(0..0, "1"), (1..1, "2")], None, cx)
2999        });
3000        assert_eq!(
3001            messages(&conversation, cx),
3002            vec![
3003                (message_1.id, Role::User, 0..2),
3004                (message_2.id, Role::Assistant, 2..3)
3005            ]
3006        );
3007
3008        let message_3 = conversation.update(cx, |conversation, cx| {
3009            conversation
3010                .insert_message_after(message_2.id, Role::User, MessageStatus::Done, cx)
3011                .unwrap()
3012        });
3013        assert_eq!(
3014            messages(&conversation, cx),
3015            vec![
3016                (message_1.id, Role::User, 0..2),
3017                (message_2.id, Role::Assistant, 2..4),
3018                (message_3.id, Role::User, 4..4)
3019            ]
3020        );
3021
3022        let message_4 = conversation.update(cx, |conversation, cx| {
3023            conversation
3024                .insert_message_after(message_2.id, Role::User, MessageStatus::Done, cx)
3025                .unwrap()
3026        });
3027        assert_eq!(
3028            messages(&conversation, cx),
3029            vec![
3030                (message_1.id, Role::User, 0..2),
3031                (message_2.id, Role::Assistant, 2..4),
3032                (message_4.id, Role::User, 4..5),
3033                (message_3.id, Role::User, 5..5),
3034            ]
3035        );
3036
3037        buffer.update(cx, |buffer, cx| {
3038            buffer.edit([(4..4, "C"), (5..5, "D")], None, cx)
3039        });
3040        assert_eq!(
3041            messages(&conversation, cx),
3042            vec![
3043                (message_1.id, Role::User, 0..2),
3044                (message_2.id, Role::Assistant, 2..4),
3045                (message_4.id, Role::User, 4..6),
3046                (message_3.id, Role::User, 6..7),
3047            ]
3048        );
3049
3050        // Deleting across message boundaries merges the messages.
3051        buffer.update(cx, |buffer, cx| buffer.edit([(1..4, "")], None, cx));
3052        assert_eq!(
3053            messages(&conversation, cx),
3054            vec![
3055                (message_1.id, Role::User, 0..3),
3056                (message_3.id, Role::User, 3..4),
3057            ]
3058        );
3059
3060        // Undoing the deletion should also undo the merge.
3061        buffer.update(cx, |buffer, cx| buffer.undo(cx));
3062        assert_eq!(
3063            messages(&conversation, cx),
3064            vec![
3065                (message_1.id, Role::User, 0..2),
3066                (message_2.id, Role::Assistant, 2..4),
3067                (message_4.id, Role::User, 4..6),
3068                (message_3.id, Role::User, 6..7),
3069            ]
3070        );
3071
3072        // Redoing the deletion should also redo the merge.
3073        buffer.update(cx, |buffer, cx| buffer.redo(cx));
3074        assert_eq!(
3075            messages(&conversation, cx),
3076            vec![
3077                (message_1.id, Role::User, 0..3),
3078                (message_3.id, Role::User, 3..4),
3079            ]
3080        );
3081
3082        // Ensure we can still insert after a merged message.
3083        let message_5 = conversation.update(cx, |conversation, cx| {
3084            conversation
3085                .insert_message_after(message_1.id, Role::System, MessageStatus::Done, cx)
3086                .unwrap()
3087        });
3088        assert_eq!(
3089            messages(&conversation, cx),
3090            vec![
3091                (message_1.id, Role::User, 0..3),
3092                (message_5.id, Role::System, 3..4),
3093                (message_3.id, Role::User, 4..5)
3094            ]
3095        );
3096    }
3097
3098    #[gpui::test]
3099    fn test_message_splitting(cx: &mut AppContext) {
3100        cx.set_global(SettingsStore::test(cx));
3101        init(cx);
3102        let registry = Arc::new(LanguageRegistry::test());
3103        let conversation = cx.add_model(|cx| Conversation::new(Default::default(), registry, cx));
3104        let buffer = conversation.read(cx).buffer.clone();
3105
3106        let message_1 = conversation.read(cx).message_anchors[0].clone();
3107        assert_eq!(
3108            messages(&conversation, cx),
3109            vec![(message_1.id, Role::User, 0..0)]
3110        );
3111
3112        buffer.update(cx, |buffer, cx| {
3113            buffer.edit([(0..0, "aaa\nbbb\nccc\nddd\n")], None, cx)
3114        });
3115
3116        let (_, message_2) =
3117            conversation.update(cx, |conversation, cx| conversation.split_message(3..3, cx));
3118        let message_2 = message_2.unwrap();
3119
3120        // We recycle newlines in the middle of a split message
3121        assert_eq!(buffer.read(cx).text(), "aaa\nbbb\nccc\nddd\n");
3122        assert_eq!(
3123            messages(&conversation, cx),
3124            vec![
3125                (message_1.id, Role::User, 0..4),
3126                (message_2.id, Role::User, 4..16),
3127            ]
3128        );
3129
3130        let (_, message_3) =
3131            conversation.update(cx, |conversation, cx| conversation.split_message(3..3, cx));
3132        let message_3 = message_3.unwrap();
3133
3134        // We don't recycle newlines at the end of a split message
3135        assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\nccc\nddd\n");
3136        assert_eq!(
3137            messages(&conversation, cx),
3138            vec![
3139                (message_1.id, Role::User, 0..4),
3140                (message_3.id, Role::User, 4..5),
3141                (message_2.id, Role::User, 5..17),
3142            ]
3143        );
3144
3145        let (_, message_4) =
3146            conversation.update(cx, |conversation, cx| conversation.split_message(9..9, cx));
3147        let message_4 = message_4.unwrap();
3148        assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\nccc\nddd\n");
3149        assert_eq!(
3150            messages(&conversation, cx),
3151            vec![
3152                (message_1.id, Role::User, 0..4),
3153                (message_3.id, Role::User, 4..5),
3154                (message_2.id, Role::User, 5..9),
3155                (message_4.id, Role::User, 9..17),
3156            ]
3157        );
3158
3159        let (_, message_5) =
3160            conversation.update(cx, |conversation, cx| conversation.split_message(9..9, cx));
3161        let message_5 = message_5.unwrap();
3162        assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\n\nccc\nddd\n");
3163        assert_eq!(
3164            messages(&conversation, cx),
3165            vec![
3166                (message_1.id, Role::User, 0..4),
3167                (message_3.id, Role::User, 4..5),
3168                (message_2.id, Role::User, 5..9),
3169                (message_4.id, Role::User, 9..10),
3170                (message_5.id, Role::User, 10..18),
3171            ]
3172        );
3173
3174        let (message_6, message_7) = conversation.update(cx, |conversation, cx| {
3175            conversation.split_message(14..16, cx)
3176        });
3177        let message_6 = message_6.unwrap();
3178        let message_7 = message_7.unwrap();
3179        assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\n\nccc\ndd\nd\n");
3180        assert_eq!(
3181            messages(&conversation, cx),
3182            vec![
3183                (message_1.id, Role::User, 0..4),
3184                (message_3.id, Role::User, 4..5),
3185                (message_2.id, Role::User, 5..9),
3186                (message_4.id, Role::User, 9..10),
3187                (message_5.id, Role::User, 10..14),
3188                (message_6.id, Role::User, 14..17),
3189                (message_7.id, Role::User, 17..19),
3190            ]
3191        );
3192    }
3193
3194    #[gpui::test]
3195    fn test_messages_for_offsets(cx: &mut AppContext) {
3196        cx.set_global(SettingsStore::test(cx));
3197        init(cx);
3198        let registry = Arc::new(LanguageRegistry::test());
3199        let conversation = cx.add_model(|cx| Conversation::new(Default::default(), registry, cx));
3200        let buffer = conversation.read(cx).buffer.clone();
3201
3202        let message_1 = conversation.read(cx).message_anchors[0].clone();
3203        assert_eq!(
3204            messages(&conversation, cx),
3205            vec![(message_1.id, Role::User, 0..0)]
3206        );
3207
3208        buffer.update(cx, |buffer, cx| buffer.edit([(0..0, "aaa")], None, cx));
3209        let message_2 = conversation
3210            .update(cx, |conversation, cx| {
3211                conversation.insert_message_after(message_1.id, Role::User, MessageStatus::Done, cx)
3212            })
3213            .unwrap();
3214        buffer.update(cx, |buffer, cx| buffer.edit([(4..4, "bbb")], None, cx));
3215
3216        let message_3 = conversation
3217            .update(cx, |conversation, cx| {
3218                conversation.insert_message_after(message_2.id, Role::User, MessageStatus::Done, cx)
3219            })
3220            .unwrap();
3221        buffer.update(cx, |buffer, cx| buffer.edit([(8..8, "ccc")], None, cx));
3222
3223        assert_eq!(buffer.read(cx).text(), "aaa\nbbb\nccc");
3224        assert_eq!(
3225            messages(&conversation, cx),
3226            vec![
3227                (message_1.id, Role::User, 0..4),
3228                (message_2.id, Role::User, 4..8),
3229                (message_3.id, Role::User, 8..11)
3230            ]
3231        );
3232
3233        assert_eq!(
3234            message_ids_for_offsets(&conversation, &[0, 4, 9], cx),
3235            [message_1.id, message_2.id, message_3.id]
3236        );
3237        assert_eq!(
3238            message_ids_for_offsets(&conversation, &[0, 1, 11], cx),
3239            [message_1.id, message_3.id]
3240        );
3241
3242        let message_4 = conversation
3243            .update(cx, |conversation, cx| {
3244                conversation.insert_message_after(message_3.id, Role::User, MessageStatus::Done, cx)
3245            })
3246            .unwrap();
3247        assert_eq!(buffer.read(cx).text(), "aaa\nbbb\nccc\n");
3248        assert_eq!(
3249            messages(&conversation, cx),
3250            vec![
3251                (message_1.id, Role::User, 0..4),
3252                (message_2.id, Role::User, 4..8),
3253                (message_3.id, Role::User, 8..12),
3254                (message_4.id, Role::User, 12..12)
3255            ]
3256        );
3257        assert_eq!(
3258            message_ids_for_offsets(&conversation, &[0, 4, 8, 12], cx),
3259            [message_1.id, message_2.id, message_3.id, message_4.id]
3260        );
3261
3262        fn message_ids_for_offsets(
3263            conversation: &ModelHandle<Conversation>,
3264            offsets: &[usize],
3265            cx: &AppContext,
3266        ) -> Vec<MessageId> {
3267            conversation
3268                .read(cx)
3269                .messages_for_offsets(offsets.iter().copied(), cx)
3270                .into_iter()
3271                .map(|message| message.id)
3272                .collect()
3273        }
3274    }
3275
3276    #[gpui::test]
3277    fn test_serialization(cx: &mut AppContext) {
3278        cx.set_global(SettingsStore::test(cx));
3279        init(cx);
3280        let registry = Arc::new(LanguageRegistry::test());
3281        let conversation =
3282            cx.add_model(|cx| Conversation::new(Default::default(), registry.clone(), cx));
3283        let buffer = conversation.read(cx).buffer.clone();
3284        let message_0 = conversation.read(cx).message_anchors[0].id;
3285        let message_1 = conversation.update(cx, |conversation, cx| {
3286            conversation
3287                .insert_message_after(message_0, Role::Assistant, MessageStatus::Done, cx)
3288                .unwrap()
3289        });
3290        let message_2 = conversation.update(cx, |conversation, cx| {
3291            conversation
3292                .insert_message_after(message_1.id, Role::System, MessageStatus::Done, cx)
3293                .unwrap()
3294        });
3295        buffer.update(cx, |buffer, cx| {
3296            buffer.edit([(0..0, "a"), (1..1, "b\nc")], None, cx);
3297            buffer.finalize_last_transaction();
3298        });
3299        let _message_3 = conversation.update(cx, |conversation, cx| {
3300            conversation
3301                .insert_message_after(message_2.id, Role::System, MessageStatus::Done, cx)
3302                .unwrap()
3303        });
3304        buffer.update(cx, |buffer, cx| buffer.undo(cx));
3305        assert_eq!(buffer.read(cx).text(), "a\nb\nc\n");
3306        assert_eq!(
3307            messages(&conversation, cx),
3308            [
3309                (message_0, Role::User, 0..2),
3310                (message_1.id, Role::Assistant, 2..6),
3311                (message_2.id, Role::System, 6..6),
3312            ]
3313        );
3314
3315        let deserialized_conversation = cx.add_model(|cx| {
3316            Conversation::deserialize(
3317                conversation.read(cx).serialize(cx),
3318                Default::default(),
3319                Default::default(),
3320                registry.clone(),
3321                cx,
3322            )
3323        });
3324        let deserialized_buffer = deserialized_conversation.read(cx).buffer.clone();
3325        assert_eq!(deserialized_buffer.read(cx).text(), "a\nb\nc\n");
3326        assert_eq!(
3327            messages(&deserialized_conversation, cx),
3328            [
3329                (message_0, Role::User, 0..2),
3330                (message_1.id, Role::Assistant, 2..6),
3331                (message_2.id, Role::System, 6..6),
3332            ]
3333        );
3334    }
3335
3336    fn messages(
3337        conversation: &ModelHandle<Conversation>,
3338        cx: &AppContext,
3339    ) -> Vec<(MessageId, Role, Range<usize>)> {
3340        conversation
3341            .read(cx)
3342            .messages(cx)
3343            .map(|message| (message.id, message.role, message.offset_range))
3344            .collect()
3345    }
3346}