terminal_inline_assistant.rs

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