terminal_inline_assistant.rs

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