terminal_inline_assistant.rs

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