terminal_inline_assistant.rs

   1use crate::{
   2    humanize_token_count, prompts::PromptBuilder, AssistantPanel, AssistantPanelEvent,
   3    ModelSelector, DEFAULT_CONTEXT_LINES,
   4};
   5use anyhow::{Context as _, Result};
   6use client::telemetry::Telemetry;
   7use collections::{HashMap, VecDeque};
   8use editor::{
   9    actions::{MoveDown, MoveUp, SelectAll},
  10    Editor, EditorElement, EditorEvent, EditorMode, EditorStyle, MultiBuffer,
  11};
  12use fs::Fs;
  13use futures::{channel::mpsc, SinkExt, StreamExt};
  14use gpui::{
  15    AppContext, Context, EventEmitter, FocusHandle, FocusableView, Global, Model, ModelContext,
  16    Subscription, Task, TextStyle, UpdateGlobal, View, WeakView,
  17};
  18use language::Buffer;
  19use language_model::{
  20    LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage, Role,
  21};
  22use settings::Settings;
  23use std::{
  24    cmp,
  25    sync::Arc,
  26    time::{Duration, Instant},
  27};
  28use terminal::Terminal;
  29use terminal_view::TerminalView;
  30use theme::ThemeSettings;
  31use ui::{prelude::*, IconButtonShape, Tooltip};
  32use util::ResultExt;
  33use workspace::{notifications::NotificationId, Toast, Workspace};
  34
  35pub fn init(
  36    fs: Arc<dyn Fs>,
  37    prompt_builder: Arc<PromptBuilder>,
  38    telemetry: Arc<Telemetry>,
  39    cx: &mut AppContext,
  40) {
  41    cx.set_global(TerminalInlineAssistant::new(fs, prompt_builder, telemetry));
  42}
  43
  44const PROMPT_HISTORY_MAX_LEN: usize = 20;
  45
  46#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
  47struct TerminalInlineAssistId(usize);
  48
  49impl TerminalInlineAssistId {
  50    fn post_inc(&mut self) -> TerminalInlineAssistId {
  51        let id = *self;
  52        self.0 += 1;
  53        id
  54    }
  55}
  56
  57pub struct TerminalInlineAssistant {
  58    next_assist_id: TerminalInlineAssistId,
  59    assists: HashMap<TerminalInlineAssistId, TerminalInlineAssist>,
  60    prompt_history: VecDeque<String>,
  61    telemetry: Option<Arc<Telemetry>>,
  62    fs: Arc<dyn Fs>,
  63    prompt_builder: Arc<PromptBuilder>,
  64}
  65
  66impl Global for TerminalInlineAssistant {}
  67
  68impl TerminalInlineAssistant {
  69    pub fn new(
  70        fs: Arc<dyn Fs>,
  71        prompt_builder: Arc<PromptBuilder>,
  72        telemetry: Arc<Telemetry>,
  73    ) -> Self {
  74        Self {
  75            next_assist_id: TerminalInlineAssistId::default(),
  76            assists: HashMap::default(),
  77            prompt_history: VecDeque::default(),
  78            telemetry: Some(telemetry),
  79            fs,
  80            prompt_builder,
  81        }
  82    }
  83
  84    pub fn assist(
  85        &mut self,
  86        terminal_view: &View<TerminalView>,
  87        workspace: Option<WeakView<Workspace>>,
  88        assistant_panel: Option<&View<AssistantPanel>>,
  89        initial_prompt: Option<String>,
  90        cx: &mut WindowContext,
  91    ) {
  92        let terminal = terminal_view.read(cx).terminal().clone();
  93        let assist_id = self.next_assist_id.post_inc();
  94        let prompt_buffer =
  95            cx.new_model(|cx| Buffer::local(initial_prompt.unwrap_or_default(), cx));
  96        let prompt_buffer = cx.new_model(|cx| MultiBuffer::singleton(prompt_buffer, cx));
  97        let codegen = cx.new_model(|_| Codegen::new(terminal, self.telemetry.clone()));
  98
  99        let prompt_editor = cx.new_view(|cx| {
 100            PromptEditor::new(
 101                assist_id,
 102                self.prompt_history.clone(),
 103                prompt_buffer.clone(),
 104                codegen,
 105                assistant_panel,
 106                workspace.clone(),
 107                self.fs.clone(),
 108                cx,
 109            )
 110        });
 111        let prompt_editor_render = prompt_editor.clone();
 112        let block = terminal_view::BlockProperties {
 113            height: 2,
 114            render: Box::new(move |_| prompt_editor_render.clone().into_any_element()),
 115        };
 116        terminal_view.update(cx, |terminal_view, cx| {
 117            terminal_view.set_block_below_cursor(block, cx);
 118        });
 119
 120        let terminal_assistant = TerminalInlineAssist::new(
 121            assist_id,
 122            terminal_view,
 123            assistant_panel.is_some(),
 124            prompt_editor,
 125            workspace.clone(),
 126            cx,
 127        );
 128
 129        self.assists.insert(assist_id, terminal_assistant);
 130
 131        self.focus_assist(assist_id, cx);
 132    }
 133
 134    fn focus_assist(&mut self, assist_id: TerminalInlineAssistId, cx: &mut WindowContext) {
 135        let assist = &self.assists[&assist_id];
 136        if let Some(prompt_editor) = assist.prompt_editor.as_ref() {
 137            prompt_editor.update(cx, |this, cx| {
 138                this.editor.update(cx, |editor, cx| {
 139                    editor.focus(cx);
 140                    editor.select_all(&SelectAll, cx);
 141                });
 142            });
 143        }
 144    }
 145
 146    fn handle_prompt_editor_event(
 147        &mut self,
 148        prompt_editor: View<PromptEditor>,
 149        event: &PromptEditorEvent,
 150        cx: &mut WindowContext,
 151    ) {
 152        let assist_id = prompt_editor.read(cx).id;
 153        match event {
 154            PromptEditorEvent::StartRequested => {
 155                self.start_assist(assist_id, cx);
 156            }
 157            PromptEditorEvent::StopRequested => {
 158                self.stop_assist(assist_id, cx);
 159            }
 160            PromptEditorEvent::ConfirmRequested => {
 161                self.finish_assist(assist_id, false, cx);
 162            }
 163            PromptEditorEvent::CancelRequested => {
 164                self.finish_assist(assist_id, true, cx);
 165            }
 166            PromptEditorEvent::DismissRequested => {
 167                self.dismiss_assist(assist_id, cx);
 168            }
 169            PromptEditorEvent::Resized { height_in_lines } => {
 170                self.insert_prompt_editor_into_terminal(assist_id, *height_in_lines, cx);
 171            }
 172        }
 173    }
 174
 175    fn start_assist(&mut self, assist_id: TerminalInlineAssistId, cx: &mut WindowContext) {
 176        let assist = if let Some(assist) = self.assists.get_mut(&assist_id) {
 177            assist
 178        } else {
 179            return;
 180        };
 181
 182        let Some(user_prompt) = assist
 183            .prompt_editor
 184            .as_ref()
 185            .map(|editor| editor.read(cx).prompt(cx))
 186        else {
 187            return;
 188        };
 189
 190        self.prompt_history.retain(|prompt| *prompt != user_prompt);
 191        self.prompt_history.push_back(user_prompt.clone());
 192        if self.prompt_history.len() > PROMPT_HISTORY_MAX_LEN {
 193            self.prompt_history.pop_front();
 194        }
 195
 196        assist
 197            .terminal
 198            .update(cx, |terminal, cx| {
 199                terminal
 200                    .terminal()
 201                    .update(cx, |terminal, _| terminal.input(CLEAR_INPUT.to_string()));
 202            })
 203            .log_err();
 204
 205        let codegen = assist.codegen.clone();
 206        let Some(request) = self.request_for_inline_assist(assist_id, cx).log_err() else {
 207            return;
 208        };
 209
 210        codegen.update(cx, |codegen, cx| codegen.start(request, cx));
 211    }
 212
 213    fn stop_assist(&mut self, assist_id: TerminalInlineAssistId, cx: &mut WindowContext) {
 214        let assist = if let Some(assist) = self.assists.get_mut(&assist_id) {
 215            assist
 216        } else {
 217            return;
 218        };
 219
 220        assist.codegen.update(cx, |codegen, cx| codegen.stop(cx));
 221    }
 222
 223    fn request_for_inline_assist(
 224        &self,
 225        assist_id: TerminalInlineAssistId,
 226        cx: &mut WindowContext,
 227    ) -> Result<LanguageModelRequest> {
 228        let assist = self.assists.get(&assist_id).context("invalid assist")?;
 229
 230        let shell = std::env::var("SHELL").ok();
 231        let (latest_output, working_directory) = assist
 232            .terminal
 233            .update(cx, |terminal, cx| {
 234                let terminal = terminal.model().read(cx);
 235                let latest_output = terminal.last_n_non_empty_lines(DEFAULT_CONTEXT_LINES);
 236                let working_directory = terminal
 237                    .working_directory()
 238                    .map(|path| path.to_string_lossy().to_string());
 239                (latest_output, working_directory)
 240            })
 241            .ok()
 242            .unwrap_or_default();
 243
 244        let context_request = if assist.include_context {
 245            assist.workspace.as_ref().and_then(|workspace| {
 246                let workspace = workspace.upgrade()?.read(cx);
 247                let assistant_panel = workspace.panel::<AssistantPanel>(cx)?;
 248                Some(
 249                    assistant_panel
 250                        .read(cx)
 251                        .active_context(cx)?
 252                        .read(cx)
 253                        .to_completion_request(cx),
 254                )
 255            })
 256        } else {
 257            None
 258        };
 259
 260        let prompt = self.prompt_builder.generate_terminal_assistant_prompt(
 261            &assist
 262                .prompt_editor
 263                .clone()
 264                .context("invalid assist")?
 265                .read(cx)
 266                .prompt(cx),
 267            shell.as_deref(),
 268            working_directory.as_deref(),
 269            &latest_output,
 270        )?;
 271
 272        let mut messages = Vec::new();
 273        if let Some(context_request) = context_request {
 274            messages = context_request.messages;
 275        }
 276
 277        messages.push(LanguageModelRequestMessage {
 278            role: Role::User,
 279            content: vec![prompt.into()],
 280            cache: false,
 281        });
 282
 283        Ok(LanguageModelRequest {
 284            messages,
 285            stop: Vec::new(),
 286            temperature: 1.0,
 287        })
 288    }
 289
 290    fn finish_assist(
 291        &mut self,
 292        assist_id: TerminalInlineAssistId,
 293        undo: bool,
 294        cx: &mut WindowContext,
 295    ) {
 296        self.dismiss_assist(assist_id, cx);
 297
 298        if let Some(assist) = self.assists.remove(&assist_id) {
 299            assist
 300                .terminal
 301                .update(cx, |this, cx| {
 302                    this.clear_block_below_cursor(cx);
 303                    this.focus_handle(cx).focus(cx);
 304                })
 305                .log_err();
 306            assist.codegen.update(cx, |codegen, cx| {
 307                if undo {
 308                    codegen.undo(cx);
 309                } else {
 310                    codegen.complete(cx);
 311                }
 312            });
 313        }
 314    }
 315
 316    fn dismiss_assist(
 317        &mut self,
 318        assist_id: TerminalInlineAssistId,
 319        cx: &mut WindowContext,
 320    ) -> bool {
 321        let Some(assist) = self.assists.get_mut(&assist_id) else {
 322            return false;
 323        };
 324        if assist.prompt_editor.is_none() {
 325            return false;
 326        }
 327        assist.prompt_editor = None;
 328        assist
 329            .terminal
 330            .update(cx, |this, cx| {
 331                this.clear_block_below_cursor(cx);
 332                this.focus_handle(cx).focus(cx);
 333            })
 334            .is_ok()
 335    }
 336
 337    fn insert_prompt_editor_into_terminal(
 338        &mut self,
 339        assist_id: TerminalInlineAssistId,
 340        height: u8,
 341        cx: &mut WindowContext,
 342    ) {
 343        if let Some(assist) = self.assists.get_mut(&assist_id) {
 344            if let Some(prompt_editor) = assist.prompt_editor.as_ref().cloned() {
 345                assist
 346                    .terminal
 347                    .update(cx, |terminal, cx| {
 348                        terminal.clear_block_below_cursor(cx);
 349                        let block = terminal_view::BlockProperties {
 350                            height,
 351                            render: Box::new(move |_| prompt_editor.clone().into_any_element()),
 352                        };
 353                        terminal.set_block_below_cursor(block, cx);
 354                    })
 355                    .log_err();
 356            }
 357        }
 358    }
 359}
 360
 361struct TerminalInlineAssist {
 362    terminal: WeakView<TerminalView>,
 363    prompt_editor: Option<View<PromptEditor>>,
 364    codegen: Model<Codegen>,
 365    workspace: Option<WeakView<Workspace>>,
 366    include_context: bool,
 367    _subscriptions: Vec<Subscription>,
 368}
 369
 370impl TerminalInlineAssist {
 371    pub fn new(
 372        assist_id: TerminalInlineAssistId,
 373        terminal: &View<TerminalView>,
 374        include_context: bool,
 375        prompt_editor: View<PromptEditor>,
 376        workspace: Option<WeakView<Workspace>>,
 377        cx: &mut WindowContext,
 378    ) -> Self {
 379        let codegen = prompt_editor.read(cx).codegen.clone();
 380        Self {
 381            terminal: terminal.downgrade(),
 382            prompt_editor: Some(prompt_editor.clone()),
 383            codegen: codegen.clone(),
 384            workspace: workspace.clone(),
 385            include_context,
 386            _subscriptions: vec![
 387                cx.subscribe(&prompt_editor, |prompt_editor, event, cx| {
 388                    TerminalInlineAssistant::update_global(cx, |this, cx| {
 389                        this.handle_prompt_editor_event(prompt_editor, event, cx)
 390                    })
 391                }),
 392                cx.subscribe(&codegen, move |codegen, event, cx| {
 393                    TerminalInlineAssistant::update_global(cx, |this, cx| match event {
 394                        CodegenEvent::Finished => {
 395                            let assist = if let Some(assist) = this.assists.get(&assist_id) {
 396                                assist
 397                            } else {
 398                                return;
 399                            };
 400
 401                            if let CodegenStatus::Error(error) = &codegen.read(cx).status {
 402                                if assist.prompt_editor.is_none() {
 403                                    if let Some(workspace) = assist
 404                                        .workspace
 405                                        .as_ref()
 406                                        .and_then(|workspace| workspace.upgrade())
 407                                    {
 408                                        let error =
 409                                            format!("Terminal inline assistant error: {}", error);
 410                                        workspace.update(cx, |workspace, cx| {
 411                                            struct InlineAssistantError;
 412
 413                                            let id =
 414                                                NotificationId::identified::<InlineAssistantError>(
 415                                                    assist_id.0,
 416                                                );
 417
 418                                            workspace.show_toast(Toast::new(id, error), cx);
 419                                        })
 420                                    }
 421                                }
 422                            }
 423
 424                            if assist.prompt_editor.is_none() {
 425                                this.finish_assist(assist_id, false, cx);
 426                            }
 427                        }
 428                    })
 429                }),
 430            ],
 431        }
 432    }
 433}
 434
 435enum PromptEditorEvent {
 436    StartRequested,
 437    StopRequested,
 438    ConfirmRequested,
 439    CancelRequested,
 440    DismissRequested,
 441    Resized { height_in_lines: u8 },
 442}
 443
 444struct PromptEditor {
 445    id: TerminalInlineAssistId,
 446    fs: Arc<dyn Fs>,
 447    height_in_lines: u8,
 448    editor: View<Editor>,
 449    edited_since_done: bool,
 450    prompt_history: VecDeque<String>,
 451    prompt_history_ix: Option<usize>,
 452    pending_prompt: String,
 453    codegen: Model<Codegen>,
 454    _codegen_subscription: Subscription,
 455    editor_subscriptions: Vec<Subscription>,
 456    pending_token_count: Task<Result<()>>,
 457    token_count: Option<usize>,
 458    _token_count_subscriptions: Vec<Subscription>,
 459    workspace: Option<WeakView<Workspace>>,
 460}
 461
 462impl EventEmitter<PromptEditorEvent> for PromptEditor {}
 463
 464impl Render for PromptEditor {
 465    fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
 466        let buttons = match &self.codegen.read(cx).status {
 467            CodegenStatus::Idle => {
 468                vec![
 469                    IconButton::new("cancel", IconName::Close)
 470                        .icon_color(Color::Muted)
 471                        .shape(IconButtonShape::Square)
 472                        .tooltip(|cx| Tooltip::for_action("Cancel Assist", &menu::Cancel, cx))
 473                        .on_click(
 474                            cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
 475                        ),
 476                    IconButton::new("start", IconName::SparkleAlt)
 477                        .icon_color(Color::Muted)
 478                        .shape(IconButtonShape::Square)
 479                        .tooltip(|cx| Tooltip::for_action("Generate", &menu::Confirm, cx))
 480                        .on_click(
 481                            cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::StartRequested)),
 482                        ),
 483                ]
 484            }
 485            CodegenStatus::Pending => {
 486                vec![
 487                    IconButton::new("cancel", IconName::Close)
 488                        .icon_color(Color::Muted)
 489                        .shape(IconButtonShape::Square)
 490                        .tooltip(|cx| Tooltip::text("Cancel Assist", cx))
 491                        .on_click(
 492                            cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
 493                        ),
 494                    IconButton::new("stop", IconName::Stop)
 495                        .icon_color(Color::Error)
 496                        .shape(IconButtonShape::Square)
 497                        .tooltip(|cx| {
 498                            Tooltip::with_meta(
 499                                "Interrupt Generation",
 500                                Some(&menu::Cancel),
 501                                "Changes won't be discarded",
 502                                cx,
 503                            )
 504                        })
 505                        .on_click(
 506                            cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::StopRequested)),
 507                        ),
 508                ]
 509            }
 510            CodegenStatus::Error(_) | CodegenStatus::Done => {
 511                vec![
 512                    IconButton::new("cancel", IconName::Close)
 513                        .icon_color(Color::Muted)
 514                        .shape(IconButtonShape::Square)
 515                        .tooltip(|cx| Tooltip::for_action("Cancel Assist", &menu::Cancel, cx))
 516                        .on_click(
 517                            cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
 518                        ),
 519                    if self.edited_since_done {
 520                        IconButton::new("restart", IconName::RotateCw)
 521                            .icon_color(Color::Info)
 522                            .shape(IconButtonShape::Square)
 523                            .tooltip(|cx| {
 524                                Tooltip::with_meta(
 525                                    "Restart Generation",
 526                                    Some(&menu::Confirm),
 527                                    "Changes will be discarded",
 528                                    cx,
 529                                )
 530                            })
 531                            .on_click(cx.listener(|_, _, cx| {
 532                                cx.emit(PromptEditorEvent::StartRequested);
 533                            }))
 534                    } else {
 535                        IconButton::new("confirm", IconName::Play)
 536                            .icon_color(Color::Info)
 537                            .shape(IconButtonShape::Square)
 538                            .tooltip(|cx| {
 539                                Tooltip::for_action("Execute generated command", &menu::Confirm, cx)
 540                            })
 541                            .on_click(cx.listener(|_, _, cx| {
 542                                cx.emit(PromptEditorEvent::ConfirmRequested);
 543                            }))
 544                    },
 545                ]
 546            }
 547        };
 548
 549        h_flex()
 550            .bg(cx.theme().colors().editor_background)
 551            .border_y_1()
 552            .border_color(cx.theme().status().info_border)
 553            .py_1p5()
 554            .h_full()
 555            .w_full()
 556            .on_action(cx.listener(Self::confirm))
 557            .on_action(cx.listener(Self::cancel))
 558            .on_action(cx.listener(Self::move_up))
 559            .on_action(cx.listener(Self::move_down))
 560            .child(
 561                h_flex()
 562                    .w_12()
 563                    .justify_center()
 564                    .gap_2()
 565                    .child(ModelSelector::new(
 566                        self.fs.clone(),
 567                        IconButton::new("context", IconName::SlidersAlt)
 568                            .shape(IconButtonShape::Square)
 569                            .icon_size(IconSize::Small)
 570                            .icon_color(Color::Muted)
 571                            .tooltip(move |cx| {
 572                                Tooltip::with_meta(
 573                                    format!(
 574                                        "Using {}",
 575                                        LanguageModelRegistry::read_global(cx)
 576                                            .active_model()
 577                                            .map(|model| model.name().0)
 578                                            .unwrap_or_else(|| "No model selected".into()),
 579                                    ),
 580                                    None,
 581                                    "Change Model",
 582                                    cx,
 583                                )
 584                            }),
 585                    ))
 586                    .children(
 587                        if let CodegenStatus::Error(error) = &self.codegen.read(cx).status {
 588                            let error_message = SharedString::from(error.to_string());
 589                            Some(
 590                                div()
 591                                    .id("error")
 592                                    .tooltip(move |cx| Tooltip::text(error_message.clone(), cx))
 593                                    .child(
 594                                        Icon::new(IconName::XCircle)
 595                                            .size(IconSize::Small)
 596                                            .color(Color::Error),
 597                                    ),
 598                            )
 599                        } else {
 600                            None
 601                        },
 602                    ),
 603            )
 604            .child(div().flex_1().child(self.render_prompt_editor(cx)))
 605            .child(
 606                h_flex()
 607                    .gap_2()
 608                    .pr_4()
 609                    .children(self.render_token_count(cx))
 610                    .children(buttons),
 611            )
 612    }
 613}
 614
 615impl FocusableView for PromptEditor {
 616    fn focus_handle(&self, cx: &AppContext) -> FocusHandle {
 617        self.editor.focus_handle(cx)
 618    }
 619}
 620
 621impl PromptEditor {
 622    const MAX_LINES: u8 = 8;
 623
 624    #[allow(clippy::too_many_arguments)]
 625    fn new(
 626        id: TerminalInlineAssistId,
 627        prompt_history: VecDeque<String>,
 628        prompt_buffer: Model<MultiBuffer>,
 629        codegen: Model<Codegen>,
 630        assistant_panel: Option<&View<AssistantPanel>>,
 631        workspace: Option<WeakView<Workspace>>,
 632        fs: Arc<dyn Fs>,
 633        cx: &mut ViewContext<Self>,
 634    ) -> Self {
 635        let prompt_editor = cx.new_view(|cx| {
 636            let mut editor = Editor::new(
 637                EditorMode::AutoHeight {
 638                    max_lines: Self::MAX_LINES as usize,
 639                },
 640                prompt_buffer,
 641                None,
 642                false,
 643                cx,
 644            );
 645            editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
 646            editor.set_placeholder_text("Add a prompt…", cx);
 647            editor
 648        });
 649
 650        let mut token_count_subscriptions = Vec::new();
 651        if let Some(assistant_panel) = assistant_panel {
 652            token_count_subscriptions
 653                .push(cx.subscribe(assistant_panel, Self::handle_assistant_panel_event));
 654        }
 655
 656        let mut this = Self {
 657            id,
 658            height_in_lines: 1,
 659            editor: prompt_editor,
 660            edited_since_done: false,
 661            prompt_history,
 662            prompt_history_ix: None,
 663            pending_prompt: String::new(),
 664            _codegen_subscription: cx.observe(&codegen, Self::handle_codegen_changed),
 665            editor_subscriptions: Vec::new(),
 666            codegen,
 667            fs,
 668            pending_token_count: Task::ready(Ok(())),
 669            token_count: None,
 670            _token_count_subscriptions: token_count_subscriptions,
 671            workspace,
 672        };
 673        this.count_lines(cx);
 674        this.count_tokens(cx);
 675        this.subscribe_to_editor(cx);
 676        this
 677    }
 678
 679    fn subscribe_to_editor(&mut self, cx: &mut ViewContext<Self>) {
 680        self.editor_subscriptions.clear();
 681        self.editor_subscriptions
 682            .push(cx.observe(&self.editor, Self::handle_prompt_editor_changed));
 683        self.editor_subscriptions
 684            .push(cx.subscribe(&self.editor, Self::handle_prompt_editor_events));
 685    }
 686
 687    fn prompt(&self, cx: &AppContext) -> String {
 688        self.editor.read(cx).text(cx)
 689    }
 690
 691    fn count_lines(&mut self, cx: &mut ViewContext<Self>) {
 692        let height_in_lines = cmp::max(
 693            2, // Make the editor at least two lines tall, to account for padding and buttons.
 694            cmp::min(
 695                self.editor
 696                    .update(cx, |editor, cx| editor.max_point(cx).row().0 + 1),
 697                Self::MAX_LINES as u32,
 698            ),
 699        ) as u8;
 700
 701        if height_in_lines != self.height_in_lines {
 702            self.height_in_lines = height_in_lines;
 703            cx.emit(PromptEditorEvent::Resized { height_in_lines });
 704        }
 705    }
 706
 707    fn handle_assistant_panel_event(
 708        &mut self,
 709        _: View<AssistantPanel>,
 710        event: &AssistantPanelEvent,
 711        cx: &mut ViewContext<Self>,
 712    ) {
 713        let AssistantPanelEvent::ContextEdited { .. } = event;
 714        self.count_tokens(cx);
 715    }
 716
 717    fn count_tokens(&mut self, cx: &mut ViewContext<Self>) {
 718        let assist_id = self.id;
 719        let Some(model) = LanguageModelRegistry::read_global(cx).active_model() else {
 720            return;
 721        };
 722        self.pending_token_count = cx.spawn(|this, mut cx| async move {
 723            cx.background_executor().timer(Duration::from_secs(1)).await;
 724            let request =
 725                cx.update_global(|inline_assistant: &mut TerminalInlineAssistant, cx| {
 726                    inline_assistant.request_for_inline_assist(assist_id, cx)
 727                })??;
 728
 729            let token_count = cx.update(|cx| model.count_tokens(request, cx))?.await?;
 730            this.update(&mut cx, |this, cx| {
 731                this.token_count = Some(token_count);
 732                cx.notify();
 733            })
 734        })
 735    }
 736
 737    fn handle_prompt_editor_changed(&mut self, _: View<Editor>, cx: &mut ViewContext<Self>) {
 738        self.count_lines(cx);
 739    }
 740
 741    fn handle_prompt_editor_events(
 742        &mut self,
 743        _: View<Editor>,
 744        event: &EditorEvent,
 745        cx: &mut ViewContext<Self>,
 746    ) {
 747        match event {
 748            EditorEvent::Edited { .. } => {
 749                let prompt = self.editor.read(cx).text(cx);
 750                if self
 751                    .prompt_history_ix
 752                    .map_or(true, |ix| self.prompt_history[ix] != prompt)
 753                {
 754                    self.prompt_history_ix.take();
 755                    self.pending_prompt = prompt;
 756                }
 757
 758                self.edited_since_done = true;
 759                cx.notify();
 760            }
 761            EditorEvent::BufferEdited => {
 762                self.count_tokens(cx);
 763            }
 764            _ => {}
 765        }
 766    }
 767
 768    fn handle_codegen_changed(&mut self, _: Model<Codegen>, cx: &mut ViewContext<Self>) {
 769        match &self.codegen.read(cx).status {
 770            CodegenStatus::Idle => {
 771                self.editor
 772                    .update(cx, |editor, _| editor.set_read_only(false));
 773            }
 774            CodegenStatus::Pending => {
 775                self.editor
 776                    .update(cx, |editor, _| editor.set_read_only(true));
 777            }
 778            CodegenStatus::Done | CodegenStatus::Error(_) => {
 779                self.edited_since_done = false;
 780                self.editor
 781                    .update(cx, |editor, _| editor.set_read_only(false));
 782            }
 783        }
 784    }
 785
 786    fn cancel(&mut self, _: &editor::actions::Cancel, cx: &mut ViewContext<Self>) {
 787        match &self.codegen.read(cx).status {
 788            CodegenStatus::Idle | CodegenStatus::Done | CodegenStatus::Error(_) => {
 789                cx.emit(PromptEditorEvent::CancelRequested);
 790            }
 791            CodegenStatus::Pending => {
 792                cx.emit(PromptEditorEvent::StopRequested);
 793            }
 794        }
 795    }
 796
 797    fn confirm(&mut self, _: &menu::Confirm, cx: &mut ViewContext<Self>) {
 798        match &self.codegen.read(cx).status {
 799            CodegenStatus::Idle => {
 800                if !self.editor.read(cx).text(cx).trim().is_empty() {
 801                    cx.emit(PromptEditorEvent::StartRequested);
 802                }
 803            }
 804            CodegenStatus::Pending => {
 805                cx.emit(PromptEditorEvent::DismissRequested);
 806            }
 807            CodegenStatus::Done | CodegenStatus::Error(_) => {
 808                if self.edited_since_done {
 809                    cx.emit(PromptEditorEvent::StartRequested);
 810                } else {
 811                    cx.emit(PromptEditorEvent::ConfirmRequested);
 812                }
 813            }
 814        }
 815    }
 816
 817    fn move_up(&mut self, _: &MoveUp, cx: &mut ViewContext<Self>) {
 818        if let Some(ix) = self.prompt_history_ix {
 819            if ix > 0 {
 820                self.prompt_history_ix = Some(ix - 1);
 821                let prompt = self.prompt_history[ix - 1].as_str();
 822                self.editor.update(cx, |editor, cx| {
 823                    editor.set_text(prompt, cx);
 824                    editor.move_to_beginning(&Default::default(), cx);
 825                });
 826            }
 827        } else if !self.prompt_history.is_empty() {
 828            self.prompt_history_ix = Some(self.prompt_history.len() - 1);
 829            let prompt = self.prompt_history[self.prompt_history.len() - 1].as_str();
 830            self.editor.update(cx, |editor, cx| {
 831                editor.set_text(prompt, cx);
 832                editor.move_to_beginning(&Default::default(), cx);
 833            });
 834        }
 835    }
 836
 837    fn move_down(&mut self, _: &MoveDown, cx: &mut ViewContext<Self>) {
 838        if let Some(ix) = self.prompt_history_ix {
 839            if ix < self.prompt_history.len() - 1 {
 840                self.prompt_history_ix = Some(ix + 1);
 841                let prompt = self.prompt_history[ix + 1].as_str();
 842                self.editor.update(cx, |editor, cx| {
 843                    editor.set_text(prompt, cx);
 844                    editor.move_to_end(&Default::default(), cx)
 845                });
 846            } else {
 847                self.prompt_history_ix = None;
 848                let prompt = self.pending_prompt.as_str();
 849                self.editor.update(cx, |editor, cx| {
 850                    editor.set_text(prompt, cx);
 851                    editor.move_to_end(&Default::default(), cx)
 852                });
 853            }
 854        }
 855    }
 856
 857    fn render_token_count(&self, cx: &mut ViewContext<Self>) -> Option<impl IntoElement> {
 858        let model = LanguageModelRegistry::read_global(cx).active_model()?;
 859        let token_count = self.token_count?;
 860        let max_token_count = model.max_token_count();
 861
 862        let remaining_tokens = max_token_count as isize - token_count as isize;
 863        let token_count_color = if remaining_tokens <= 0 {
 864            Color::Error
 865        } else if token_count as f32 / max_token_count as f32 >= 0.8 {
 866            Color::Warning
 867        } else {
 868            Color::Muted
 869        };
 870
 871        let mut token_count = h_flex()
 872            .id("token_count")
 873            .gap_0p5()
 874            .child(
 875                Label::new(humanize_token_count(token_count))
 876                    .size(LabelSize::Small)
 877                    .color(token_count_color),
 878            )
 879            .child(Label::new("/").size(LabelSize::Small).color(Color::Muted))
 880            .child(
 881                Label::new(humanize_token_count(max_token_count))
 882                    .size(LabelSize::Small)
 883                    .color(Color::Muted),
 884            );
 885        if let Some(workspace) = self.workspace.clone() {
 886            token_count = token_count
 887                .tooltip(|cx| {
 888                    Tooltip::with_meta(
 889                        "Tokens Used by Inline Assistant",
 890                        None,
 891                        "Click to Open Assistant Panel",
 892                        cx,
 893                    )
 894                })
 895                .cursor_pointer()
 896                .on_mouse_down(gpui::MouseButton::Left, |_, cx| cx.stop_propagation())
 897                .on_click(move |_, cx| {
 898                    cx.stop_propagation();
 899                    workspace
 900                        .update(cx, |workspace, cx| {
 901                            workspace.focus_panel::<AssistantPanel>(cx)
 902                        })
 903                        .ok();
 904                });
 905        } else {
 906            token_count = token_count
 907                .cursor_default()
 908                .tooltip(|cx| Tooltip::text("Tokens Used by Inline Assistant", cx));
 909        }
 910
 911        Some(token_count)
 912    }
 913
 914    fn render_prompt_editor(&self, cx: &mut ViewContext<Self>) -> impl IntoElement {
 915        let settings = ThemeSettings::get_global(cx);
 916        let text_style = TextStyle {
 917            color: if self.editor.read(cx).read_only(cx) {
 918                cx.theme().colors().text_disabled
 919            } else {
 920                cx.theme().colors().text
 921            },
 922            font_family: settings.ui_font.family.clone(),
 923            font_features: settings.ui_font.features.clone(),
 924            font_fallbacks: settings.ui_font.fallbacks.clone(),
 925            font_size: rems(0.875).into(),
 926            font_weight: settings.ui_font.weight,
 927            line_height: relative(1.3),
 928            ..Default::default()
 929        };
 930        EditorElement::new(
 931            &self.editor,
 932            EditorStyle {
 933                background: cx.theme().colors().editor_background,
 934                local_player: cx.theme().players().local(),
 935                text: text_style,
 936                ..Default::default()
 937            },
 938        )
 939    }
 940}
 941
 942#[derive(Debug)]
 943pub enum CodegenEvent {
 944    Finished,
 945}
 946
 947impl EventEmitter<CodegenEvent> for Codegen {}
 948
 949const CLEAR_INPUT: &str = "\x15";
 950const CARRIAGE_RETURN: &str = "\x0d";
 951
 952struct TerminalTransaction {
 953    terminal: Model<Terminal>,
 954}
 955
 956impl TerminalTransaction {
 957    pub fn start(terminal: Model<Terminal>) -> Self {
 958        Self { terminal }
 959    }
 960
 961    pub fn push(&mut self, hunk: String, cx: &mut AppContext) {
 962        // Ensure that the assistant cannot accidentally execute commands that are streamed into the terminal
 963        let input = hunk.replace(CARRIAGE_RETURN, " ");
 964        self.terminal
 965            .update(cx, |terminal, _| terminal.input(input));
 966    }
 967
 968    pub fn undo(&self, cx: &mut AppContext) {
 969        self.terminal
 970            .update(cx, |terminal, _| terminal.input(CLEAR_INPUT.to_string()));
 971    }
 972
 973    pub fn complete(&self, cx: &mut AppContext) {
 974        self.terminal.update(cx, |terminal, _| {
 975            terminal.input(CARRIAGE_RETURN.to_string())
 976        });
 977    }
 978}
 979
 980pub struct Codegen {
 981    status: CodegenStatus,
 982    telemetry: Option<Arc<Telemetry>>,
 983    terminal: Model<Terminal>,
 984    generation: Task<()>,
 985    transaction: Option<TerminalTransaction>,
 986}
 987
 988impl Codegen {
 989    pub fn new(terminal: Model<Terminal>, telemetry: Option<Arc<Telemetry>>) -> Self {
 990        Self {
 991            terminal,
 992            telemetry,
 993            status: CodegenStatus::Idle,
 994            generation: Task::ready(()),
 995            transaction: None,
 996        }
 997    }
 998
 999    pub fn start(&mut self, prompt: LanguageModelRequest, cx: &mut ModelContext<Self>) {
1000        let Some(model) = LanguageModelRegistry::read_global(cx).active_model() else {
1001            return;
1002        };
1003
1004        let telemetry = self.telemetry.clone();
1005        self.status = CodegenStatus::Pending;
1006        self.transaction = Some(TerminalTransaction::start(self.terminal.clone()));
1007        self.generation = cx.spawn(|this, mut cx| async move {
1008            let model_telemetry_id = model.telemetry_id();
1009            let response = model.stream_completion(prompt, &cx).await;
1010            let generate = async {
1011                let (mut hunks_tx, mut hunks_rx) = mpsc::channel(1);
1012
1013                let task = cx.background_executor().spawn(async move {
1014                    let mut response_latency = None;
1015                    let request_start = Instant::now();
1016                    let task = async {
1017                        let mut chunks = response?;
1018                        while let Some(chunk) = chunks.next().await {
1019                            if response_latency.is_none() {
1020                                response_latency = Some(request_start.elapsed());
1021                            }
1022                            let chunk = chunk?;
1023                            hunks_tx.send(chunk).await?;
1024                        }
1025
1026                        anyhow::Ok(())
1027                    };
1028
1029                    let result = task.await;
1030
1031                    let error_message = result.as_ref().err().map(|error| error.to_string());
1032                    if let Some(telemetry) = telemetry {
1033                        telemetry.report_assistant_event(
1034                            None,
1035                            telemetry_events::AssistantKind::Inline,
1036                            model_telemetry_id,
1037                            response_latency,
1038                            error_message,
1039                        );
1040                    }
1041
1042                    result?;
1043                    anyhow::Ok(())
1044                });
1045
1046                while let Some(hunk) = hunks_rx.next().await {
1047                    this.update(&mut cx, |this, cx| {
1048                        if let Some(transaction) = &mut this.transaction {
1049                            transaction.push(hunk, cx);
1050                            cx.notify();
1051                        }
1052                    })?;
1053                }
1054
1055                task.await?;
1056                anyhow::Ok(())
1057            };
1058
1059            let result = generate.await;
1060
1061            this.update(&mut cx, |this, cx| {
1062                if let Err(error) = result {
1063                    this.status = CodegenStatus::Error(error);
1064                } else {
1065                    this.status = CodegenStatus::Done;
1066                }
1067                cx.emit(CodegenEvent::Finished);
1068                cx.notify();
1069            })
1070            .ok();
1071        });
1072        cx.notify();
1073    }
1074
1075    pub fn stop(&mut self, cx: &mut ModelContext<Self>) {
1076        self.status = CodegenStatus::Done;
1077        self.generation = Task::ready(());
1078        cx.emit(CodegenEvent::Finished);
1079        cx.notify();
1080    }
1081
1082    pub fn complete(&mut self, cx: &mut ModelContext<Self>) {
1083        if let Some(transaction) = self.transaction.take() {
1084            transaction.complete(cx);
1085        }
1086    }
1087
1088    pub fn undo(&mut self, cx: &mut ModelContext<Self>) {
1089        if let Some(transaction) = self.transaction.take() {
1090            transaction.undo(cx);
1091        }
1092    }
1093}
1094
1095enum CodegenStatus {
1096    Idle,
1097    Pending,
1098    Done,
1099    Error(anyhow::Error),
1100}