terminal_inline_assistant.rs

   1use crate::{
   2    humanize_token_count, prompts::generate_terminal_assistant_prompt, AssistantPanel,
   3    AssistantPanelEvent, LanguageModelCompletionProvider, 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::{LanguageModelRequest, LanguageModelRequestMessage, Role};
  20use settings::Settings;
  21use std::{
  22    cmp,
  23    sync::Arc,
  24    time::{Duration, Instant},
  25};
  26use terminal::Terminal;
  27use terminal_view::TerminalView;
  28use theme::ThemeSettings;
  29use ui::{prelude::*, IconButtonShape, Tooltip};
  30use util::ResultExt;
  31use workspace::{notifications::NotificationId, Toast, Workspace};
  32
  33pub fn init(fs: Arc<dyn Fs>, telemetry: Arc<Telemetry>, cx: &mut AppContext) {
  34    cx.set_global(TerminalInlineAssistant::new(fs, telemetry));
  35}
  36
  37const PROMPT_HISTORY_MAX_LEN: usize = 20;
  38
  39#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, Hash)]
  40struct TerminalInlineAssistId(usize);
  41
  42impl TerminalInlineAssistId {
  43    fn post_inc(&mut self) -> TerminalInlineAssistId {
  44        let id = *self;
  45        self.0 += 1;
  46        id
  47    }
  48}
  49
  50pub struct TerminalInlineAssistant {
  51    next_assist_id: TerminalInlineAssistId,
  52    assists: HashMap<TerminalInlineAssistId, TerminalInlineAssist>,
  53    prompt_history: VecDeque<String>,
  54    telemetry: Option<Arc<Telemetry>>,
  55    fs: Arc<dyn Fs>,
  56}
  57
  58impl Global for TerminalInlineAssistant {}
  59
  60impl TerminalInlineAssistant {
  61    pub fn new(fs: Arc<dyn Fs>, telemetry: Arc<Telemetry>) -> Self {
  62        Self {
  63            next_assist_id: TerminalInlineAssistId::default(),
  64            assists: HashMap::default(),
  65            prompt_history: VecDeque::default(),
  66            telemetry: Some(telemetry),
  67            fs,
  68        }
  69    }
  70
  71    pub fn assist(
  72        &mut self,
  73        terminal_view: &View<TerminalView>,
  74        workspace: Option<WeakView<Workspace>>,
  75        assistant_panel: Option<&View<AssistantPanel>>,
  76        initial_prompt: Option<String>,
  77        cx: &mut WindowContext,
  78    ) {
  79        let terminal = terminal_view.read(cx).terminal().clone();
  80        let assist_id = self.next_assist_id.post_inc();
  81        let prompt_buffer =
  82            cx.new_model(|cx| Buffer::local(initial_prompt.unwrap_or_default(), cx));
  83        let prompt_buffer = cx.new_model(|cx| MultiBuffer::singleton(prompt_buffer, cx));
  84        let codegen = cx.new_model(|_| Codegen::new(terminal, self.telemetry.clone()));
  85
  86        let prompt_editor = cx.new_view(|cx| {
  87            PromptEditor::new(
  88                assist_id,
  89                self.prompt_history.clone(),
  90                prompt_buffer.clone(),
  91                codegen,
  92                assistant_panel,
  93                workspace.clone(),
  94                self.fs.clone(),
  95                cx,
  96            )
  97        });
  98        let prompt_editor_render = prompt_editor.clone();
  99        let block = terminal_view::BlockProperties {
 100            height: 2,
 101            render: Box::new(move |_| prompt_editor_render.clone().into_any_element()),
 102        };
 103        terminal_view.update(cx, |terminal_view, cx| {
 104            terminal_view.set_block_below_cursor(block, cx);
 105        });
 106
 107        let terminal_assistant = TerminalInlineAssist::new(
 108            assist_id,
 109            terminal_view,
 110            assistant_panel.is_some(),
 111            prompt_editor,
 112            workspace.clone(),
 113            cx,
 114        );
 115
 116        self.assists.insert(assist_id, terminal_assistant);
 117
 118        self.focus_assist(assist_id, cx);
 119    }
 120
 121    fn focus_assist(&mut self, assist_id: TerminalInlineAssistId, cx: &mut WindowContext) {
 122        let assist = &self.assists[&assist_id];
 123        if let Some(prompt_editor) = assist.prompt_editor.as_ref() {
 124            prompt_editor.update(cx, |this, cx| {
 125                this.editor.update(cx, |editor, cx| {
 126                    editor.focus(cx);
 127                    editor.select_all(&SelectAll, cx);
 128                });
 129            });
 130        }
 131    }
 132
 133    fn handle_prompt_editor_event(
 134        &mut self,
 135        prompt_editor: View<PromptEditor>,
 136        event: &PromptEditorEvent,
 137        cx: &mut WindowContext,
 138    ) {
 139        let assist_id = prompt_editor.read(cx).id;
 140        match event {
 141            PromptEditorEvent::StartRequested => {
 142                self.start_assist(assist_id, cx);
 143            }
 144            PromptEditorEvent::StopRequested => {
 145                self.stop_assist(assist_id, cx);
 146            }
 147            PromptEditorEvent::ConfirmRequested => {
 148                self.finish_assist(assist_id, false, cx);
 149            }
 150            PromptEditorEvent::CancelRequested => {
 151                self.finish_assist(assist_id, true, cx);
 152            }
 153            PromptEditorEvent::DismissRequested => {
 154                self.dismiss_assist(assist_id, cx);
 155            }
 156            PromptEditorEvent::Resized { height_in_lines } => {
 157                self.insert_prompt_editor_into_terminal(assist_id, *height_in_lines, cx);
 158            }
 159        }
 160    }
 161
 162    fn start_assist(&mut self, assist_id: TerminalInlineAssistId, cx: &mut WindowContext) {
 163        let assist = if let Some(assist) = self.assists.get_mut(&assist_id) {
 164            assist
 165        } else {
 166            return;
 167        };
 168
 169        let Some(user_prompt) = assist
 170            .prompt_editor
 171            .as_ref()
 172            .map(|editor| editor.read(cx).prompt(cx))
 173        else {
 174            return;
 175        };
 176
 177        self.prompt_history.retain(|prompt| *prompt != user_prompt);
 178        self.prompt_history.push_back(user_prompt.clone());
 179        if self.prompt_history.len() > PROMPT_HISTORY_MAX_LEN {
 180            self.prompt_history.pop_front();
 181        }
 182
 183        assist
 184            .terminal
 185            .update(cx, |terminal, cx| {
 186                terminal
 187                    .terminal()
 188                    .update(cx, |terminal, _| terminal.input(CLEAR_INPUT.to_string()));
 189            })
 190            .log_err();
 191
 192        let codegen = assist.codegen.clone();
 193        let Some(request) = self.request_for_inline_assist(assist_id, cx).log_err() else {
 194            return;
 195        };
 196
 197        codegen.update(cx, |codegen, cx| codegen.start(request, cx));
 198    }
 199
 200    fn stop_assist(&mut self, assist_id: TerminalInlineAssistId, cx: &mut WindowContext) {
 201        let assist = if let Some(assist) = self.assists.get_mut(&assist_id) {
 202            assist
 203        } else {
 204            return;
 205        };
 206
 207        assist.codegen.update(cx, |codegen, cx| codegen.stop(cx));
 208    }
 209
 210    fn request_for_inline_assist(
 211        &self,
 212        assist_id: TerminalInlineAssistId,
 213        cx: &mut WindowContext,
 214    ) -> Result<LanguageModelRequest> {
 215        let assist = self.assists.get(&assist_id).context("invalid assist")?;
 216
 217        let shell = std::env::var("SHELL").ok();
 218        let working_directory = assist
 219            .terminal
 220            .update(cx, |terminal, cx| {
 221                terminal
 222                    .model()
 223                    .read(cx)
 224                    .working_directory()
 225                    .map(|path| path.to_string_lossy().to_string())
 226            })
 227            .ok()
 228            .flatten();
 229
 230        let context_request = if assist.include_context {
 231            assist.workspace.as_ref().and_then(|workspace| {
 232                let workspace = workspace.upgrade()?.read(cx);
 233                let assistant_panel = workspace.panel::<AssistantPanel>(cx)?;
 234                Some(
 235                    assistant_panel
 236                        .read(cx)
 237                        .active_context(cx)?
 238                        .read(cx)
 239                        .to_completion_request(cx),
 240                )
 241            })
 242        } else {
 243            None
 244        };
 245
 246        let prompt = generate_terminal_assistant_prompt(
 247            &assist
 248                .prompt_editor
 249                .clone()
 250                .context("invalid assist")?
 251                .read(cx)
 252                .prompt(cx),
 253            shell.as_deref(),
 254            working_directory.as_deref(),
 255        );
 256
 257        let mut messages = Vec::new();
 258        if let Some(context_request) = context_request {
 259            messages = context_request.messages;
 260        }
 261
 262        messages.push(LanguageModelRequestMessage {
 263            role: Role::User,
 264            content: prompt,
 265        });
 266
 267        Ok(LanguageModelRequest {
 268            messages,
 269            stop: Vec::new(),
 270            temperature: 1.0,
 271        })
 272    }
 273
 274    fn finish_assist(
 275        &mut self,
 276        assist_id: TerminalInlineAssistId,
 277        undo: bool,
 278        cx: &mut WindowContext,
 279    ) {
 280        self.dismiss_assist(assist_id, cx);
 281
 282        if let Some(assist) = self.assists.remove(&assist_id) {
 283            assist
 284                .terminal
 285                .update(cx, |this, cx| {
 286                    this.clear_block_below_cursor(cx);
 287                    this.focus_handle(cx).focus(cx);
 288                })
 289                .log_err();
 290            assist.codegen.update(cx, |codegen, cx| {
 291                if undo {
 292                    codegen.undo(cx);
 293                } else {
 294                    codegen.complete(cx);
 295                }
 296            });
 297        }
 298    }
 299
 300    fn dismiss_assist(
 301        &mut self,
 302        assist_id: TerminalInlineAssistId,
 303        cx: &mut WindowContext,
 304    ) -> bool {
 305        let Some(assist) = self.assists.get_mut(&assist_id) else {
 306            return false;
 307        };
 308        if assist.prompt_editor.is_none() {
 309            return false;
 310        }
 311        assist.prompt_editor = None;
 312        assist
 313            .terminal
 314            .update(cx, |this, cx| {
 315                this.clear_block_below_cursor(cx);
 316                this.focus_handle(cx).focus(cx);
 317            })
 318            .is_ok()
 319    }
 320
 321    fn insert_prompt_editor_into_terminal(
 322        &mut self,
 323        assist_id: TerminalInlineAssistId,
 324        height: u8,
 325        cx: &mut WindowContext,
 326    ) {
 327        if let Some(assist) = self.assists.get_mut(&assist_id) {
 328            if let Some(prompt_editor) = assist.prompt_editor.as_ref().cloned() {
 329                assist
 330                    .terminal
 331                    .update(cx, |terminal, cx| {
 332                        terminal.clear_block_below_cursor(cx);
 333                        let block = terminal_view::BlockProperties {
 334                            height,
 335                            render: Box::new(move |_| prompt_editor.clone().into_any_element()),
 336                        };
 337                        terminal.set_block_below_cursor(block, cx);
 338                    })
 339                    .log_err();
 340            }
 341        }
 342    }
 343}
 344
 345struct TerminalInlineAssist {
 346    terminal: WeakView<TerminalView>,
 347    prompt_editor: Option<View<PromptEditor>>,
 348    codegen: Model<Codegen>,
 349    workspace: Option<WeakView<Workspace>>,
 350    include_context: bool,
 351    _subscriptions: Vec<Subscription>,
 352}
 353
 354impl TerminalInlineAssist {
 355    pub fn new(
 356        assist_id: TerminalInlineAssistId,
 357        terminal: &View<TerminalView>,
 358        include_context: bool,
 359        prompt_editor: View<PromptEditor>,
 360        workspace: Option<WeakView<Workspace>>,
 361        cx: &mut WindowContext,
 362    ) -> Self {
 363        let codegen = prompt_editor.read(cx).codegen.clone();
 364        Self {
 365            terminal: terminal.downgrade(),
 366            prompt_editor: Some(prompt_editor.clone()),
 367            codegen: codegen.clone(),
 368            workspace: workspace.clone(),
 369            include_context,
 370            _subscriptions: vec![
 371                cx.subscribe(&prompt_editor, |prompt_editor, event, cx| {
 372                    TerminalInlineAssistant::update_global(cx, |this, cx| {
 373                        this.handle_prompt_editor_event(prompt_editor, event, cx)
 374                    })
 375                }),
 376                cx.subscribe(&codegen, move |codegen, event, cx| {
 377                    TerminalInlineAssistant::update_global(cx, |this, cx| match event {
 378                        CodegenEvent::Finished => {
 379                            let assist = if let Some(assist) = this.assists.get(&assist_id) {
 380                                assist
 381                            } else {
 382                                return;
 383                            };
 384
 385                            if let CodegenStatus::Error(error) = &codegen.read(cx).status {
 386                                if assist.prompt_editor.is_none() {
 387                                    if let Some(workspace) = assist
 388                                        .workspace
 389                                        .as_ref()
 390                                        .and_then(|workspace| workspace.upgrade())
 391                                    {
 392                                        let error =
 393                                            format!("Terminal inline assistant error: {}", error);
 394                                        workspace.update(cx, |workspace, cx| {
 395                                            struct InlineAssistantError;
 396
 397                                            let id =
 398                                                NotificationId::identified::<InlineAssistantError>(
 399                                                    assist_id.0,
 400                                                );
 401
 402                                            workspace.show_toast(Toast::new(id, error), cx);
 403                                        })
 404                                    }
 405                                }
 406                            }
 407
 408                            if assist.prompt_editor.is_none() {
 409                                this.finish_assist(assist_id, false, cx);
 410                            }
 411                        }
 412                    })
 413                }),
 414            ],
 415        }
 416    }
 417}
 418
 419enum PromptEditorEvent {
 420    StartRequested,
 421    StopRequested,
 422    ConfirmRequested,
 423    CancelRequested,
 424    DismissRequested,
 425    Resized { height_in_lines: u8 },
 426}
 427
 428struct PromptEditor {
 429    id: TerminalInlineAssistId,
 430    fs: Arc<dyn Fs>,
 431    height_in_lines: u8,
 432    editor: View<Editor>,
 433    edited_since_done: bool,
 434    prompt_history: VecDeque<String>,
 435    prompt_history_ix: Option<usize>,
 436    pending_prompt: String,
 437    codegen: Model<Codegen>,
 438    _codegen_subscription: Subscription,
 439    editor_subscriptions: Vec<Subscription>,
 440    pending_token_count: Task<Result<()>>,
 441    token_count: Option<usize>,
 442    _token_count_subscriptions: Vec<Subscription>,
 443    workspace: Option<WeakView<Workspace>>,
 444}
 445
 446impl EventEmitter<PromptEditorEvent> for PromptEditor {}
 447
 448impl Render for PromptEditor {
 449    fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
 450        let buttons = match &self.codegen.read(cx).status {
 451            CodegenStatus::Idle => {
 452                vec![
 453                    IconButton::new("cancel", IconName::Close)
 454                        .icon_color(Color::Muted)
 455                        .shape(IconButtonShape::Square)
 456                        .tooltip(|cx| Tooltip::for_action("Cancel Assist", &menu::Cancel, cx))
 457                        .on_click(
 458                            cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
 459                        ),
 460                    IconButton::new("start", IconName::SparkleAlt)
 461                        .icon_color(Color::Muted)
 462                        .shape(IconButtonShape::Square)
 463                        .tooltip(|cx| Tooltip::for_action("Generate", &menu::Confirm, cx))
 464                        .on_click(
 465                            cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::StartRequested)),
 466                        ),
 467                ]
 468            }
 469            CodegenStatus::Pending => {
 470                vec![
 471                    IconButton::new("cancel", IconName::Close)
 472                        .icon_color(Color::Muted)
 473                        .shape(IconButtonShape::Square)
 474                        .tooltip(|cx| Tooltip::text("Cancel Assist", cx))
 475                        .on_click(
 476                            cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
 477                        ),
 478                    IconButton::new("stop", IconName::Stop)
 479                        .icon_color(Color::Error)
 480                        .shape(IconButtonShape::Square)
 481                        .tooltip(|cx| {
 482                            Tooltip::with_meta(
 483                                "Interrupt Generation",
 484                                Some(&menu::Cancel),
 485                                "Changes won't be discarded",
 486                                cx,
 487                            )
 488                        })
 489                        .on_click(
 490                            cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::StopRequested)),
 491                        ),
 492                ]
 493            }
 494            CodegenStatus::Error(_) | CodegenStatus::Done => {
 495                vec![
 496                    IconButton::new("cancel", IconName::Close)
 497                        .icon_color(Color::Muted)
 498                        .shape(IconButtonShape::Square)
 499                        .tooltip(|cx| Tooltip::for_action("Cancel Assist", &menu::Cancel, cx))
 500                        .on_click(
 501                            cx.listener(|_, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
 502                        ),
 503                    if self.edited_since_done {
 504                        IconButton::new("restart", IconName::RotateCw)
 505                            .icon_color(Color::Info)
 506                            .shape(IconButtonShape::Square)
 507                            .tooltip(|cx| {
 508                                Tooltip::with_meta(
 509                                    "Restart Generation",
 510                                    Some(&menu::Confirm),
 511                                    "Changes will be discarded",
 512                                    cx,
 513                                )
 514                            })
 515                            .on_click(cx.listener(|_, _, cx| {
 516                                cx.emit(PromptEditorEvent::StartRequested);
 517                            }))
 518                    } else {
 519                        IconButton::new("confirm", IconName::Play)
 520                            .icon_color(Color::Info)
 521                            .shape(IconButtonShape::Square)
 522                            .tooltip(|cx| {
 523                                Tooltip::for_action("Execute generated command", &menu::Confirm, cx)
 524                            })
 525                            .on_click(cx.listener(|_, _, cx| {
 526                                cx.emit(PromptEditorEvent::ConfirmRequested);
 527                            }))
 528                    },
 529                ]
 530            }
 531        };
 532
 533        h_flex()
 534            .bg(cx.theme().colors().editor_background)
 535            .border_y_1()
 536            .border_color(cx.theme().status().info_border)
 537            .py_1p5()
 538            .h_full()
 539            .w_full()
 540            .on_action(cx.listener(Self::confirm))
 541            .on_action(cx.listener(Self::cancel))
 542            .on_action(cx.listener(Self::move_up))
 543            .on_action(cx.listener(Self::move_down))
 544            .child(
 545                h_flex()
 546                    .w_12()
 547                    .justify_center()
 548                    .gap_2()
 549                    .child(ModelSelector::new(
 550                        self.fs.clone(),
 551                        IconButton::new("context", IconName::Settings)
 552                            .shape(IconButtonShape::Square)
 553                            .icon_size(IconSize::Small)
 554                            .icon_color(Color::Muted)
 555                            .tooltip(move |cx| {
 556                                Tooltip::with_meta(
 557                                    format!(
 558                                        "Using {}",
 559                                        LanguageModelCompletionProvider::read_global(cx)
 560                                            .active_model()
 561                                            .map(|model| model.name().0)
 562                                            .unwrap_or_else(|| "No model selected".into()),
 563                                    ),
 564                                    None,
 565                                    "Change Model",
 566                                    cx,
 567                                )
 568                            }),
 569                    ))
 570                    .children(
 571                        if let CodegenStatus::Error(error) = &self.codegen.read(cx).status {
 572                            let error_message = SharedString::from(error.to_string());
 573                            Some(
 574                                div()
 575                                    .id("error")
 576                                    .tooltip(move |cx| Tooltip::text(error_message.clone(), cx))
 577                                    .child(
 578                                        Icon::new(IconName::XCircle)
 579                                            .size(IconSize::Small)
 580                                            .color(Color::Error),
 581                                    ),
 582                            )
 583                        } else {
 584                            None
 585                        },
 586                    ),
 587            )
 588            .child(div().flex_1().child(self.render_prompt_editor(cx)))
 589            .child(
 590                h_flex()
 591                    .gap_2()
 592                    .pr_4()
 593                    .children(self.render_token_count(cx))
 594                    .children(buttons),
 595            )
 596    }
 597}
 598
 599impl FocusableView for PromptEditor {
 600    fn focus_handle(&self, cx: &AppContext) -> FocusHandle {
 601        self.editor.focus_handle(cx)
 602    }
 603}
 604
 605impl PromptEditor {
 606    const MAX_LINES: u8 = 8;
 607
 608    #[allow(clippy::too_many_arguments)]
 609    fn new(
 610        id: TerminalInlineAssistId,
 611        prompt_history: VecDeque<String>,
 612        prompt_buffer: Model<MultiBuffer>,
 613        codegen: Model<Codegen>,
 614        assistant_panel: Option<&View<AssistantPanel>>,
 615        workspace: Option<WeakView<Workspace>>,
 616        fs: Arc<dyn Fs>,
 617        cx: &mut ViewContext<Self>,
 618    ) -> Self {
 619        let prompt_editor = cx.new_view(|cx| {
 620            let mut editor = Editor::new(
 621                EditorMode::AutoHeight {
 622                    max_lines: Self::MAX_LINES as usize,
 623                },
 624                prompt_buffer,
 625                None,
 626                false,
 627                cx,
 628            );
 629            editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
 630            editor.set_placeholder_text("Add a prompt…", cx);
 631            editor
 632        });
 633
 634        let mut token_count_subscriptions = Vec::new();
 635        if let Some(assistant_panel) = assistant_panel {
 636            token_count_subscriptions
 637                .push(cx.subscribe(assistant_panel, Self::handle_assistant_panel_event));
 638        }
 639
 640        let mut this = Self {
 641            id,
 642            height_in_lines: 1,
 643            editor: prompt_editor,
 644            edited_since_done: false,
 645            prompt_history,
 646            prompt_history_ix: None,
 647            pending_prompt: String::new(),
 648            _codegen_subscription: cx.observe(&codegen, Self::handle_codegen_changed),
 649            editor_subscriptions: Vec::new(),
 650            codegen,
 651            fs,
 652            pending_token_count: Task::ready(Ok(())),
 653            token_count: None,
 654            _token_count_subscriptions: token_count_subscriptions,
 655            workspace,
 656        };
 657        this.count_lines(cx);
 658        this.count_tokens(cx);
 659        this.subscribe_to_editor(cx);
 660        this
 661    }
 662
 663    fn subscribe_to_editor(&mut self, cx: &mut ViewContext<Self>) {
 664        self.editor_subscriptions.clear();
 665        self.editor_subscriptions
 666            .push(cx.observe(&self.editor, Self::handle_prompt_editor_changed));
 667        self.editor_subscriptions
 668            .push(cx.subscribe(&self.editor, Self::handle_prompt_editor_events));
 669    }
 670
 671    fn prompt(&self, cx: &AppContext) -> String {
 672        self.editor.read(cx).text(cx)
 673    }
 674
 675    fn count_lines(&mut self, cx: &mut ViewContext<Self>) {
 676        let height_in_lines = cmp::max(
 677            2, // Make the editor at least two lines tall, to account for padding and buttons.
 678            cmp::min(
 679                self.editor
 680                    .update(cx, |editor, cx| editor.max_point(cx).row().0 + 1),
 681                Self::MAX_LINES as u32,
 682            ),
 683        ) as u8;
 684
 685        if height_in_lines != self.height_in_lines {
 686            self.height_in_lines = height_in_lines;
 687            cx.emit(PromptEditorEvent::Resized { height_in_lines });
 688        }
 689    }
 690
 691    fn handle_assistant_panel_event(
 692        &mut self,
 693        _: View<AssistantPanel>,
 694        event: &AssistantPanelEvent,
 695        cx: &mut ViewContext<Self>,
 696    ) {
 697        let AssistantPanelEvent::ContextEdited { .. } = event;
 698        self.count_tokens(cx);
 699    }
 700
 701    fn count_tokens(&mut self, cx: &mut ViewContext<Self>) {
 702        let assist_id = self.id;
 703        self.pending_token_count = cx.spawn(|this, mut cx| async move {
 704            cx.background_executor().timer(Duration::from_secs(1)).await;
 705            let request =
 706                cx.update_global(|inline_assistant: &mut TerminalInlineAssistant, cx| {
 707                    inline_assistant.request_for_inline_assist(assist_id, cx)
 708                })??;
 709
 710            if let Some(token_count) = cx.update(|cx| {
 711                LanguageModelCompletionProvider::read_global(cx).count_tokens(request, cx)
 712            })? {
 713                let token_count = token_count.await?;
 714
 715                this.update(&mut cx, |this, cx| {
 716                    this.token_count = Some(token_count);
 717                    cx.notify();
 718                })
 719            } else {
 720                Ok(())
 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 = LanguageModelCompletionProvider::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        self.status = CodegenStatus::Pending;
 989        self.transaction = Some(TerminalTransaction::start(self.terminal.clone()));
 990
 991        let telemetry = self.telemetry.clone();
 992        let model_telemetry_id = LanguageModelCompletionProvider::read_global(cx)
 993            .active_model()
 994            .map(|m| m.telemetry_id())
 995            .unwrap_or_default();
 996        let response =
 997            LanguageModelCompletionProvider::read_global(cx).stream_completion(prompt, cx);
 998
 999        self.generation = cx.spawn(|this, mut cx| async move {
1000            let response = response.await;
1001            let generate = async {
1002                let (mut hunks_tx, mut hunks_rx) = mpsc::channel(1);
1003
1004                let task = cx.background_executor().spawn(async move {
1005                    let mut response_latency = None;
1006                    let request_start = Instant::now();
1007                    let task = async {
1008                        let mut chunks = response?;
1009                        while let Some(chunk) = chunks.next().await {
1010                            if response_latency.is_none() {
1011                                response_latency = Some(request_start.elapsed());
1012                            }
1013                            let chunk = chunk?;
1014                            hunks_tx.send(chunk).await?;
1015                        }
1016
1017                        anyhow::Ok(())
1018                    };
1019
1020                    let result = task.await;
1021
1022                    let error_message = result.as_ref().err().map(|error| error.to_string());
1023                    if let Some(telemetry) = telemetry {
1024                        telemetry.report_assistant_event(
1025                            None,
1026                            telemetry_events::AssistantKind::Inline,
1027                            model_telemetry_id,
1028                            response_latency,
1029                            error_message,
1030                        );
1031                    }
1032
1033                    result?;
1034                    anyhow::Ok(())
1035                });
1036
1037                while let Some(hunk) = hunks_rx.next().await {
1038                    this.update(&mut cx, |this, cx| {
1039                        if let Some(transaction) = &mut this.transaction {
1040                            transaction.push(hunk, cx);
1041                            cx.notify();
1042                        }
1043                    })?;
1044                }
1045
1046                task.await?;
1047                anyhow::Ok(())
1048            };
1049
1050            let result = generate.await;
1051
1052            this.update(&mut cx, |this, cx| {
1053                if let Err(error) = result {
1054                    this.status = CodegenStatus::Error(error);
1055                } else {
1056                    this.status = CodegenStatus::Done;
1057                }
1058                cx.emit(CodegenEvent::Finished);
1059                cx.notify();
1060            })
1061            .ok();
1062        });
1063        cx.notify();
1064    }
1065
1066    pub fn stop(&mut self, cx: &mut ModelContext<Self>) {
1067        self.status = CodegenStatus::Done;
1068        self.generation = Task::ready(());
1069        cx.emit(CodegenEvent::Finished);
1070        cx.notify();
1071    }
1072
1073    pub fn complete(&mut self, cx: &mut ModelContext<Self>) {
1074        if let Some(transaction) = self.transaction.take() {
1075            transaction.complete(cx);
1076        }
1077    }
1078
1079    pub fn undo(&mut self, cx: &mut ModelContext<Self>) {
1080        if let Some(transaction) = self.transaction.take() {
1081            transaction.undo(cx);
1082        }
1083    }
1084}
1085
1086enum CodegenStatus {
1087    Idle,
1088    Pending,
1089    Done,
1090    Error(anyhow::Error),
1091}