terminal_inline_assistant.rs

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