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 { execute } => {
 161                self.finish_assist(assist_id, false, *execute, cx);
 162            }
 163            PromptEditorEvent::CancelRequested => {
 164                self.finish_assist(assist_id, true, false, 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            tools: Vec::new(),
 286            stop: Vec::new(),
 287            temperature: 1.0,
 288        })
 289    }
 290
 291    fn finish_assist(
 292        &mut self,
 293        assist_id: TerminalInlineAssistId,
 294        undo: bool,
 295        execute: bool,
 296        cx: &mut WindowContext,
 297    ) {
 298        self.dismiss_assist(assist_id, cx);
 299
 300        if let Some(assist) = self.assists.remove(&assist_id) {
 301            assist
 302                .terminal
 303                .update(cx, |this, cx| {
 304                    this.clear_block_below_cursor(cx);
 305                    this.focus_handle(cx).focus(cx);
 306                })
 307                .log_err();
 308            assist.codegen.update(cx, |codegen, cx| {
 309                if undo {
 310                    codegen.undo(cx);
 311                } else if execute {
 312                    codegen.complete(cx);
 313                }
 314            });
 315        }
 316    }
 317
 318    fn dismiss_assist(
 319        &mut self,
 320        assist_id: TerminalInlineAssistId,
 321        cx: &mut WindowContext,
 322    ) -> bool {
 323        let Some(assist) = self.assists.get_mut(&assist_id) else {
 324            return false;
 325        };
 326        if assist.prompt_editor.is_none() {
 327            return false;
 328        }
 329        assist.prompt_editor = None;
 330        assist
 331            .terminal
 332            .update(cx, |this, cx| {
 333                this.clear_block_below_cursor(cx);
 334                this.focus_handle(cx).focus(cx);
 335            })
 336            .is_ok()
 337    }
 338
 339    fn insert_prompt_editor_into_terminal(
 340        &mut self,
 341        assist_id: TerminalInlineAssistId,
 342        height: u8,
 343        cx: &mut WindowContext,
 344    ) {
 345        if let Some(assist) = self.assists.get_mut(&assist_id) {
 346            if let Some(prompt_editor) = assist.prompt_editor.as_ref().cloned() {
 347                assist
 348                    .terminal
 349                    .update(cx, |terminal, cx| {
 350                        terminal.clear_block_below_cursor(cx);
 351                        let block = terminal_view::BlockProperties {
 352                            height,
 353                            render: Box::new(move |_| prompt_editor.clone().into_any_element()),
 354                        };
 355                        terminal.set_block_below_cursor(block, cx);
 356                    })
 357                    .log_err();
 358            }
 359        }
 360    }
 361}
 362
 363struct TerminalInlineAssist {
 364    terminal: WeakView<TerminalView>,
 365    prompt_editor: Option<View<PromptEditor>>,
 366    codegen: Model<Codegen>,
 367    workspace: Option<WeakView<Workspace>>,
 368    include_context: bool,
 369    _subscriptions: Vec<Subscription>,
 370}
 371
 372impl TerminalInlineAssist {
 373    pub fn new(
 374        assist_id: TerminalInlineAssistId,
 375        terminal: &View<TerminalView>,
 376        include_context: bool,
 377        prompt_editor: View<PromptEditor>,
 378        workspace: Option<WeakView<Workspace>>,
 379        cx: &mut WindowContext,
 380    ) -> Self {
 381        let codegen = prompt_editor.read(cx).codegen.clone();
 382        Self {
 383            terminal: terminal.downgrade(),
 384            prompt_editor: Some(prompt_editor.clone()),
 385            codegen: codegen.clone(),
 386            workspace: workspace.clone(),
 387            include_context,
 388            _subscriptions: vec![
 389                cx.subscribe(&prompt_editor, |prompt_editor, event, cx| {
 390                    TerminalInlineAssistant::update_global(cx, |this, cx| {
 391                        this.handle_prompt_editor_event(prompt_editor, event, cx)
 392                    })
 393                }),
 394                cx.subscribe(&codegen, move |codegen, event, cx| {
 395                    TerminalInlineAssistant::update_global(cx, |this, cx| match event {
 396                        CodegenEvent::Finished => {
 397                            let assist = if let Some(assist) = this.assists.get(&assist_id) {
 398                                assist
 399                            } else {
 400                                return;
 401                            };
 402
 403                            if let CodegenStatus::Error(error) = &codegen.read(cx).status {
 404                                if assist.prompt_editor.is_none() {
 405                                    if let Some(workspace) = assist
 406                                        .workspace
 407                                        .as_ref()
 408                                        .and_then(|workspace| workspace.upgrade())
 409                                    {
 410                                        let error =
 411                                            format!("Terminal inline assistant error: {}", error);
 412                                        workspace.update(cx, |workspace, cx| {
 413                                            struct InlineAssistantError;
 414
 415                                            let id =
 416                                                NotificationId::identified::<InlineAssistantError>(
 417                                                    assist_id.0,
 418                                                );
 419
 420                                            workspace.show_toast(Toast::new(id, error), cx);
 421                                        })
 422                                    }
 423                                }
 424                            }
 425
 426                            if assist.prompt_editor.is_none() {
 427                                this.finish_assist(assist_id, false, false, cx);
 428                            }
 429                        }
 430                    })
 431                }),
 432            ],
 433        }
 434    }
 435}
 436
 437enum PromptEditorEvent {
 438    StartRequested,
 439    StopRequested,
 440    ConfirmRequested { execute: bool },
 441    CancelRequested,
 442    DismissRequested,
 443    Resized { height_in_lines: u8 },
 444}
 445
 446struct PromptEditor {
 447    id: TerminalInlineAssistId,
 448    fs: Arc<dyn Fs>,
 449    height_in_lines: u8,
 450    editor: View<Editor>,
 451    edited_since_done: bool,
 452    prompt_history: VecDeque<String>,
 453    prompt_history_ix: Option<usize>,
 454    pending_prompt: String,
 455    codegen: Model<Codegen>,
 456    _codegen_subscription: Subscription,
 457    editor_subscriptions: Vec<Subscription>,
 458    pending_token_count: Task<Result<()>>,
 459    token_count: Option<usize>,
 460    _token_count_subscriptions: Vec<Subscription>,
 461    workspace: Option<WeakView<Workspace>>,
 462}
 463
 464impl EventEmitter<PromptEditorEvent> for PromptEditor {}
 465
 466impl Render for PromptEditor {
 467    fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
 468        let buttons = match &self.codegen.read(cx).status {
 469            CodegenStatus::Idle => {
 470                vec![
 471                    IconButton::new("cancel", IconName::Close)
 472                        .icon_color(Color::Muted)
 473                        .shape(IconButtonShape::Square)
 474                        .tooltip(|cx| Tooltip::for_action("Cancel Assist", &menu::Cancel, cx))
 475                        .on_click(
 476                            cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
 477                        ),
 478                    IconButton::new("start", IconName::SparkleAlt)
 479                        .icon_color(Color::Muted)
 480                        .shape(IconButtonShape::Square)
 481                        .tooltip(|cx| Tooltip::for_action("Generate", &menu::Confirm, cx))
 482                        .on_click(
 483                            cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::StartRequested)),
 484                        ),
 485                ]
 486            }
 487            CodegenStatus::Pending => {
 488                vec![
 489                    IconButton::new("cancel", IconName::Close)
 490                        .icon_color(Color::Muted)
 491                        .shape(IconButtonShape::Square)
 492                        .tooltip(|cx| Tooltip::text("Cancel Assist", cx))
 493                        .on_click(
 494                            cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
 495                        ),
 496                    IconButton::new("stop", IconName::Stop)
 497                        .icon_color(Color::Error)
 498                        .shape(IconButtonShape::Square)
 499                        .tooltip(|cx| {
 500                            Tooltip::with_meta(
 501                                "Interrupt Generation",
 502                                Some(&menu::Cancel),
 503                                "Changes won't be discarded",
 504                                cx,
 505                            )
 506                        })
 507                        .on_click(
 508                            cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::StopRequested)),
 509                        ),
 510                ]
 511            }
 512            CodegenStatus::Error(_) | CodegenStatus::Done => {
 513                let cancel = IconButton::new("cancel", IconName::Close)
 514                    .icon_color(Color::Muted)
 515                    .shape(IconButtonShape::Square)
 516                    .tooltip(|cx| Tooltip::for_action("Cancel Assist", &menu::Cancel, cx))
 517                    .on_click(cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)));
 518
 519                if self.edited_since_done {
 520                    vec![
 521                        cancel,
 522                        IconButton::new("restart", IconName::RotateCw)
 523                            .icon_color(Color::Info)
 524                            .shape(IconButtonShape::Square)
 525                            .tooltip(|cx| {
 526                                Tooltip::with_meta(
 527                                    "Restart Generation",
 528                                    Some(&menu::Confirm),
 529                                    "Changes will be discarded",
 530                                    cx,
 531                                )
 532                            })
 533                            .on_click(cx.listener(|_, _, cx| {
 534                                cx.emit(PromptEditorEvent::StartRequested);
 535                            })),
 536                    ]
 537                } else {
 538                    vec![
 539                        cancel,
 540                        IconButton::new("accept", IconName::Check)
 541                            .icon_color(Color::Info)
 542                            .shape(IconButtonShape::Square)
 543                            .tooltip(|cx| {
 544                                Tooltip::for_action("Accept Generated Command", &menu::Confirm, cx)
 545                            })
 546                            .on_click(cx.listener(|_, _, cx| {
 547                                cx.emit(PromptEditorEvent::ConfirmRequested { execute: false });
 548                            })),
 549                        IconButton::new("confirm", IconName::Play)
 550                            .icon_color(Color::Info)
 551                            .shape(IconButtonShape::Square)
 552                            .tooltip(|cx| {
 553                                Tooltip::for_action(
 554                                    "Execute Generated Command",
 555                                    &menu::SecondaryConfirm,
 556                                    cx,
 557                                )
 558                            })
 559                            .on_click(cx.listener(|_, _, cx| {
 560                                cx.emit(PromptEditorEvent::ConfirmRequested { execute: true });
 561                            })),
 562                    ]
 563                }
 564            }
 565        };
 566
 567        h_flex()
 568            .bg(cx.theme().colors().editor_background)
 569            .border_y_1()
 570            .border_color(cx.theme().status().info_border)
 571            .py_1p5()
 572            .h_full()
 573            .w_full()
 574            .on_action(cx.listener(Self::confirm))
 575            .on_action(cx.listener(Self::secondary_confirm))
 576            .on_action(cx.listener(Self::cancel))
 577            .on_action(cx.listener(Self::move_up))
 578            .on_action(cx.listener(Self::move_down))
 579            .child(
 580                h_flex()
 581                    .w_12()
 582                    .justify_center()
 583                    .gap_2()
 584                    .child(ModelSelector::new(
 585                        self.fs.clone(),
 586                        IconButton::new("context", IconName::SlidersAlt)
 587                            .shape(IconButtonShape::Square)
 588                            .icon_size(IconSize::Small)
 589                            .icon_color(Color::Muted)
 590                            .tooltip(move |cx| {
 591                                Tooltip::with_meta(
 592                                    format!(
 593                                        "Using {}",
 594                                        LanguageModelRegistry::read_global(cx)
 595                                            .active_model()
 596                                            .map(|model| model.name().0)
 597                                            .unwrap_or_else(|| "No model selected".into()),
 598                                    ),
 599                                    None,
 600                                    "Change Model",
 601                                    cx,
 602                                )
 603                            }),
 604                    ))
 605                    .children(
 606                        if let CodegenStatus::Error(error) = &self.codegen.read(cx).status {
 607                            let error_message = SharedString::from(error.to_string());
 608                            Some(
 609                                div()
 610                                    .id("error")
 611                                    .tooltip(move |cx| Tooltip::text(error_message.clone(), cx))
 612                                    .child(
 613                                        Icon::new(IconName::XCircle)
 614                                            .size(IconSize::Small)
 615                                            .color(Color::Error),
 616                                    ),
 617                            )
 618                        } else {
 619                            None
 620                        },
 621                    ),
 622            )
 623            .child(div().flex_1().child(self.render_prompt_editor(cx)))
 624            .child(
 625                h_flex()
 626                    .gap_1()
 627                    .pr_4()
 628                    .children(self.render_token_count(cx))
 629                    .children(buttons),
 630            )
 631    }
 632}
 633
 634impl FocusableView for PromptEditor {
 635    fn focus_handle(&self, cx: &AppContext) -> FocusHandle {
 636        self.editor.focus_handle(cx)
 637    }
 638}
 639
 640impl PromptEditor {
 641    const MAX_LINES: u8 = 8;
 642
 643    #[allow(clippy::too_many_arguments)]
 644    fn new(
 645        id: TerminalInlineAssistId,
 646        prompt_history: VecDeque<String>,
 647        prompt_buffer: Model<MultiBuffer>,
 648        codegen: Model<Codegen>,
 649        assistant_panel: Option<&View<AssistantPanel>>,
 650        workspace: Option<WeakView<Workspace>>,
 651        fs: Arc<dyn Fs>,
 652        cx: &mut ViewContext<Self>,
 653    ) -> Self {
 654        let prompt_editor = cx.new_view(|cx| {
 655            let mut editor = Editor::new(
 656                EditorMode::AutoHeight {
 657                    max_lines: Self::MAX_LINES as usize,
 658                },
 659                prompt_buffer,
 660                None,
 661                false,
 662                cx,
 663            );
 664            editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
 665            editor.set_placeholder_text("Add a prompt…", cx);
 666            editor
 667        });
 668
 669        let mut token_count_subscriptions = Vec::new();
 670        if let Some(assistant_panel) = assistant_panel {
 671            token_count_subscriptions
 672                .push(cx.subscribe(assistant_panel, Self::handle_assistant_panel_event));
 673        }
 674
 675        let mut this = Self {
 676            id,
 677            height_in_lines: 1,
 678            editor: prompt_editor,
 679            edited_since_done: false,
 680            prompt_history,
 681            prompt_history_ix: None,
 682            pending_prompt: String::new(),
 683            _codegen_subscription: cx.observe(&codegen, Self::handle_codegen_changed),
 684            editor_subscriptions: Vec::new(),
 685            codegen,
 686            fs,
 687            pending_token_count: Task::ready(Ok(())),
 688            token_count: None,
 689            _token_count_subscriptions: token_count_subscriptions,
 690            workspace,
 691        };
 692        this.count_lines(cx);
 693        this.count_tokens(cx);
 694        this.subscribe_to_editor(cx);
 695        this
 696    }
 697
 698    fn subscribe_to_editor(&mut self, cx: &mut ViewContext<Self>) {
 699        self.editor_subscriptions.clear();
 700        self.editor_subscriptions
 701            .push(cx.observe(&self.editor, Self::handle_prompt_editor_changed));
 702        self.editor_subscriptions
 703            .push(cx.subscribe(&self.editor, Self::handle_prompt_editor_events));
 704    }
 705
 706    fn prompt(&self, cx: &AppContext) -> String {
 707        self.editor.read(cx).text(cx)
 708    }
 709
 710    fn count_lines(&mut self, cx: &mut ViewContext<Self>) {
 711        let height_in_lines = cmp::max(
 712            2, // Make the editor at least two lines tall, to account for padding and buttons.
 713            cmp::min(
 714                self.editor
 715                    .update(cx, |editor, cx| editor.max_point(cx).row().0 + 1),
 716                Self::MAX_LINES as u32,
 717            ),
 718        ) as u8;
 719
 720        if height_in_lines != self.height_in_lines {
 721            self.height_in_lines = height_in_lines;
 722            cx.emit(PromptEditorEvent::Resized { height_in_lines });
 723        }
 724    }
 725
 726    fn handle_assistant_panel_event(
 727        &mut self,
 728        _: View<AssistantPanel>,
 729        event: &AssistantPanelEvent,
 730        cx: &mut ViewContext<Self>,
 731    ) {
 732        let AssistantPanelEvent::ContextEdited { .. } = event;
 733        self.count_tokens(cx);
 734    }
 735
 736    fn count_tokens(&mut self, cx: &mut ViewContext<Self>) {
 737        let assist_id = self.id;
 738        let Some(model) = LanguageModelRegistry::read_global(cx).active_model() else {
 739            return;
 740        };
 741        self.pending_token_count = cx.spawn(|this, mut cx| async move {
 742            cx.background_executor().timer(Duration::from_secs(1)).await;
 743            let request =
 744                cx.update_global(|inline_assistant: &mut TerminalInlineAssistant, cx| {
 745                    inline_assistant.request_for_inline_assist(assist_id, cx)
 746                })??;
 747
 748            let token_count = cx.update(|cx| model.count_tokens(request, cx))?.await?;
 749            this.update(&mut cx, |this, cx| {
 750                this.token_count = Some(token_count);
 751                cx.notify();
 752            })
 753        })
 754    }
 755
 756    fn handle_prompt_editor_changed(&mut self, _: View<Editor>, cx: &mut ViewContext<Self>) {
 757        self.count_lines(cx);
 758    }
 759
 760    fn handle_prompt_editor_events(
 761        &mut self,
 762        _: View<Editor>,
 763        event: &EditorEvent,
 764        cx: &mut ViewContext<Self>,
 765    ) {
 766        match event {
 767            EditorEvent::Edited { .. } => {
 768                let prompt = self.editor.read(cx).text(cx);
 769                if self
 770                    .prompt_history_ix
 771                    .map_or(true, |ix| self.prompt_history[ix] != prompt)
 772                {
 773                    self.prompt_history_ix.take();
 774                    self.pending_prompt = prompt;
 775                }
 776
 777                self.edited_since_done = true;
 778                cx.notify();
 779            }
 780            EditorEvent::BufferEdited => {
 781                self.count_tokens(cx);
 782            }
 783            _ => {}
 784        }
 785    }
 786
 787    fn handle_codegen_changed(&mut self, _: Model<Codegen>, cx: &mut ViewContext<Self>) {
 788        match &self.codegen.read(cx).status {
 789            CodegenStatus::Idle => {
 790                self.editor
 791                    .update(cx, |editor, _| editor.set_read_only(false));
 792            }
 793            CodegenStatus::Pending => {
 794                self.editor
 795                    .update(cx, |editor, _| editor.set_read_only(true));
 796            }
 797            CodegenStatus::Done | CodegenStatus::Error(_) => {
 798                self.edited_since_done = false;
 799                self.editor
 800                    .update(cx, |editor, _| editor.set_read_only(false));
 801            }
 802        }
 803    }
 804
 805    fn cancel(&mut self, _: &editor::actions::Cancel, cx: &mut ViewContext<Self>) {
 806        match &self.codegen.read(cx).status {
 807            CodegenStatus::Idle | CodegenStatus::Done | CodegenStatus::Error(_) => {
 808                cx.emit(PromptEditorEvent::CancelRequested);
 809            }
 810            CodegenStatus::Pending => {
 811                cx.emit(PromptEditorEvent::StopRequested);
 812            }
 813        }
 814    }
 815
 816    fn confirm(&mut self, _: &menu::Confirm, cx: &mut ViewContext<Self>) {
 817        match &self.codegen.read(cx).status {
 818            CodegenStatus::Idle => {
 819                if !self.editor.read(cx).text(cx).trim().is_empty() {
 820                    cx.emit(PromptEditorEvent::StartRequested);
 821                }
 822            }
 823            CodegenStatus::Pending => {
 824                cx.emit(PromptEditorEvent::DismissRequested);
 825            }
 826            CodegenStatus::Done => {
 827                if self.edited_since_done {
 828                    cx.emit(PromptEditorEvent::StartRequested);
 829                } else {
 830                    cx.emit(PromptEditorEvent::ConfirmRequested { execute: false });
 831                }
 832            }
 833            CodegenStatus::Error(_) => {
 834                cx.emit(PromptEditorEvent::StartRequested);
 835            }
 836        }
 837    }
 838
 839    fn secondary_confirm(&mut self, _: &menu::SecondaryConfirm, cx: &mut ViewContext<Self>) {
 840        if matches!(self.codegen.read(cx).status, CodegenStatus::Done) {
 841            cx.emit(PromptEditorEvent::ConfirmRequested { execute: true });
 842        }
 843    }
 844
 845    fn move_up(&mut self, _: &MoveUp, cx: &mut ViewContext<Self>) {
 846        if let Some(ix) = self.prompt_history_ix {
 847            if ix > 0 {
 848                self.prompt_history_ix = Some(ix - 1);
 849                let prompt = self.prompt_history[ix - 1].as_str();
 850                self.editor.update(cx, |editor, cx| {
 851                    editor.set_text(prompt, cx);
 852                    editor.move_to_beginning(&Default::default(), cx);
 853                });
 854            }
 855        } else if !self.prompt_history.is_empty() {
 856            self.prompt_history_ix = Some(self.prompt_history.len() - 1);
 857            let prompt = self.prompt_history[self.prompt_history.len() - 1].as_str();
 858            self.editor.update(cx, |editor, cx| {
 859                editor.set_text(prompt, cx);
 860                editor.move_to_beginning(&Default::default(), cx);
 861            });
 862        }
 863    }
 864
 865    fn move_down(&mut self, _: &MoveDown, cx: &mut ViewContext<Self>) {
 866        if let Some(ix) = self.prompt_history_ix {
 867            if ix < self.prompt_history.len() - 1 {
 868                self.prompt_history_ix = Some(ix + 1);
 869                let prompt = self.prompt_history[ix + 1].as_str();
 870                self.editor.update(cx, |editor, cx| {
 871                    editor.set_text(prompt, cx);
 872                    editor.move_to_end(&Default::default(), cx)
 873                });
 874            } else {
 875                self.prompt_history_ix = None;
 876                let prompt = self.pending_prompt.as_str();
 877                self.editor.update(cx, |editor, cx| {
 878                    editor.set_text(prompt, cx);
 879                    editor.move_to_end(&Default::default(), cx)
 880                });
 881            }
 882        }
 883    }
 884
 885    fn render_token_count(&self, cx: &mut ViewContext<Self>) -> Option<impl IntoElement> {
 886        let model = LanguageModelRegistry::read_global(cx).active_model()?;
 887        let token_count = self.token_count?;
 888        let max_token_count = model.max_token_count();
 889
 890        let remaining_tokens = max_token_count as isize - token_count as isize;
 891        let token_count_color = if remaining_tokens <= 0 {
 892            Color::Error
 893        } else if token_count as f32 / max_token_count as f32 >= 0.8 {
 894            Color::Warning
 895        } else {
 896            Color::Muted
 897        };
 898
 899        let mut token_count = h_flex()
 900            .id("token_count")
 901            .gap_0p5()
 902            .child(
 903                Label::new(humanize_token_count(token_count))
 904                    .size(LabelSize::Small)
 905                    .color(token_count_color),
 906            )
 907            .child(Label::new("/").size(LabelSize::Small).color(Color::Muted))
 908            .child(
 909                Label::new(humanize_token_count(max_token_count))
 910                    .size(LabelSize::Small)
 911                    .color(Color::Muted),
 912            );
 913        if let Some(workspace) = self.workspace.clone() {
 914            token_count = token_count
 915                .tooltip(|cx| {
 916                    Tooltip::with_meta(
 917                        "Tokens Used by Inline Assistant",
 918                        None,
 919                        "Click to Open Assistant Panel",
 920                        cx,
 921                    )
 922                })
 923                .cursor_pointer()
 924                .on_mouse_down(gpui::MouseButton::Left, |_, cx| cx.stop_propagation())
 925                .on_click(move |_, cx| {
 926                    cx.stop_propagation();
 927                    workspace
 928                        .update(cx, |workspace, cx| {
 929                            workspace.focus_panel::<AssistantPanel>(cx)
 930                        })
 931                        .ok();
 932                });
 933        } else {
 934            token_count = token_count
 935                .cursor_default()
 936                .tooltip(|cx| Tooltip::text("Tokens Used by Inline Assistant", cx));
 937        }
 938
 939        Some(token_count)
 940    }
 941
 942    fn render_prompt_editor(&self, cx: &mut ViewContext<Self>) -> impl IntoElement {
 943        let settings = ThemeSettings::get_global(cx);
 944        let text_style = TextStyle {
 945            color: if self.editor.read(cx).read_only(cx) {
 946                cx.theme().colors().text_disabled
 947            } else {
 948                cx.theme().colors().text
 949            },
 950            font_family: settings.ui_font.family.clone(),
 951            font_features: settings.ui_font.features.clone(),
 952            font_fallbacks: settings.ui_font.fallbacks.clone(),
 953            font_size: rems(0.875).into(),
 954            font_weight: settings.ui_font.weight,
 955            line_height: relative(1.3),
 956            ..Default::default()
 957        };
 958        EditorElement::new(
 959            &self.editor,
 960            EditorStyle {
 961                background: cx.theme().colors().editor_background,
 962                local_player: cx.theme().players().local(),
 963                text: text_style,
 964                ..Default::default()
 965            },
 966        )
 967    }
 968}
 969
 970#[derive(Debug)]
 971pub enum CodegenEvent {
 972    Finished,
 973}
 974
 975impl EventEmitter<CodegenEvent> for Codegen {}
 976
 977const CLEAR_INPUT: &str = "\x15";
 978const CARRIAGE_RETURN: &str = "\x0d";
 979
 980struct TerminalTransaction {
 981    terminal: Model<Terminal>,
 982}
 983
 984impl TerminalTransaction {
 985    pub fn start(terminal: Model<Terminal>) -> Self {
 986        Self { terminal }
 987    }
 988
 989    pub fn push(&mut self, hunk: String, cx: &mut AppContext) {
 990        // Ensure that the assistant cannot accidentally execute commands that are streamed into the terminal
 991        let input = hunk.replace(CARRIAGE_RETURN, " ");
 992        self.terminal
 993            .update(cx, |terminal, _| terminal.input(input));
 994    }
 995
 996    pub fn undo(&self, cx: &mut AppContext) {
 997        self.terminal
 998            .update(cx, |terminal, _| terminal.input(CLEAR_INPUT.to_string()));
 999    }
1000
1001    pub fn complete(&self, cx: &mut AppContext) {
1002        self.terminal.update(cx, |terminal, _| {
1003            terminal.input(CARRIAGE_RETURN.to_string())
1004        });
1005    }
1006}
1007
1008pub struct Codegen {
1009    status: CodegenStatus,
1010    telemetry: Option<Arc<Telemetry>>,
1011    terminal: Model<Terminal>,
1012    generation: Task<()>,
1013    transaction: Option<TerminalTransaction>,
1014}
1015
1016impl Codegen {
1017    pub fn new(terminal: Model<Terminal>, telemetry: Option<Arc<Telemetry>>) -> Self {
1018        Self {
1019            terminal,
1020            telemetry,
1021            status: CodegenStatus::Idle,
1022            generation: Task::ready(()),
1023            transaction: None,
1024        }
1025    }
1026
1027    pub fn start(&mut self, prompt: LanguageModelRequest, cx: &mut ModelContext<Self>) {
1028        let Some(model) = LanguageModelRegistry::read_global(cx).active_model() else {
1029            return;
1030        };
1031
1032        let telemetry = self.telemetry.clone();
1033        self.status = CodegenStatus::Pending;
1034        self.transaction = Some(TerminalTransaction::start(self.terminal.clone()));
1035        self.generation = cx.spawn(|this, mut cx| async move {
1036            let model_telemetry_id = model.telemetry_id();
1037            let response = model.stream_completion_text(prompt, &cx).await;
1038            let generate = async {
1039                let (mut hunks_tx, mut hunks_rx) = mpsc::channel(1);
1040
1041                let task = cx.background_executor().spawn(async move {
1042                    let mut response_latency = None;
1043                    let request_start = Instant::now();
1044                    let task = async {
1045                        let mut chunks = response?;
1046                        while let Some(chunk) = chunks.next().await {
1047                            if response_latency.is_none() {
1048                                response_latency = Some(request_start.elapsed());
1049                            }
1050                            let chunk = chunk?;
1051                            hunks_tx.send(chunk).await?;
1052                        }
1053
1054                        anyhow::Ok(())
1055                    };
1056
1057                    let result = task.await;
1058
1059                    let error_message = result.as_ref().err().map(|error| error.to_string());
1060                    if let Some(telemetry) = telemetry {
1061                        telemetry.report_assistant_event(
1062                            None,
1063                            telemetry_events::AssistantKind::Inline,
1064                            model_telemetry_id,
1065                            response_latency,
1066                            error_message,
1067                        );
1068                    }
1069
1070                    result?;
1071                    anyhow::Ok(())
1072                });
1073
1074                while let Some(hunk) = hunks_rx.next().await {
1075                    this.update(&mut cx, |this, cx| {
1076                        if let Some(transaction) = &mut this.transaction {
1077                            transaction.push(hunk, cx);
1078                            cx.notify();
1079                        }
1080                    })?;
1081                }
1082
1083                task.await?;
1084                anyhow::Ok(())
1085            };
1086
1087            let result = generate.await;
1088
1089            this.update(&mut cx, |this, cx| {
1090                if let Err(error) = result {
1091                    this.status = CodegenStatus::Error(error);
1092                } else {
1093                    this.status = CodegenStatus::Done;
1094                }
1095                cx.emit(CodegenEvent::Finished);
1096                cx.notify();
1097            })
1098            .ok();
1099        });
1100        cx.notify();
1101    }
1102
1103    pub fn stop(&mut self, cx: &mut ModelContext<Self>) {
1104        self.status = CodegenStatus::Done;
1105        self.generation = Task::ready(());
1106        cx.emit(CodegenEvent::Finished);
1107        cx.notify();
1108    }
1109
1110    pub fn complete(&mut self, cx: &mut ModelContext<Self>) {
1111        if let Some(transaction) = self.transaction.take() {
1112            transaction.complete(cx);
1113        }
1114    }
1115
1116    pub fn undo(&mut self, cx: &mut ModelContext<Self>) {
1117        if let Some(transaction) = self.transaction.take() {
1118            transaction.undo(cx);
1119        }
1120    }
1121}
1122
1123enum CodegenStatus {
1124    Idle,
1125    Pending,
1126    Done,
1127    Error(anyhow::Error),
1128}