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