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    ContextMenuOptions, 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            mode: None,
 298            messages,
 299            tools: Vec::new(),
 300            stop: Vec::new(),
 301            temperature: None,
 302        })
 303    }
 304
 305    fn finish_assist(
 306        &mut self,
 307        assist_id: TerminalInlineAssistId,
 308        undo: bool,
 309        execute: bool,
 310        window: &mut Window,
 311        cx: &mut App,
 312    ) {
 313        self.dismiss_assist(assist_id, window, cx);
 314
 315        if let Some(assist) = self.assists.remove(&assist_id) {
 316            assist
 317                .terminal
 318                .update(cx, |this, cx| {
 319                    this.clear_block_below_cursor(cx);
 320                    this.focus_handle(cx).focus(window);
 321                })
 322                .log_err();
 323
 324            if let Some(ConfiguredModel { model, .. }) =
 325                LanguageModelRegistry::read_global(cx).inline_assistant_model()
 326            {
 327                let codegen = assist.codegen.read(cx);
 328                let executor = cx.background_executor().clone();
 329                report_assistant_event(
 330                    AssistantEventData {
 331                        conversation_id: None,
 332                        kind: AssistantKind::InlineTerminal,
 333                        message_id: codegen.message_id.clone(),
 334                        phase: if undo {
 335                            AssistantPhase::Rejected
 336                        } else {
 337                            AssistantPhase::Accepted
 338                        },
 339                        model: model.telemetry_id(),
 340                        model_provider: model.provider_id().to_string(),
 341                        response_latency: None,
 342                        error_message: None,
 343                        language_name: None,
 344                    },
 345                    codegen.telemetry.clone(),
 346                    cx.http_client(),
 347                    model.api_key(cx),
 348                    &executor,
 349                );
 350            }
 351
 352            assist.codegen.update(cx, |codegen, cx| {
 353                if undo {
 354                    codegen.undo(cx);
 355                } else if execute {
 356                    codegen.complete(cx);
 357                }
 358            });
 359        }
 360    }
 361
 362    fn dismiss_assist(
 363        &mut self,
 364        assist_id: TerminalInlineAssistId,
 365        window: &mut Window,
 366        cx: &mut App,
 367    ) -> bool {
 368        let Some(assist) = self.assists.get_mut(&assist_id) else {
 369            return false;
 370        };
 371        if assist.prompt_editor.is_none() {
 372            return false;
 373        }
 374        assist.prompt_editor = None;
 375        assist
 376            .terminal
 377            .update(cx, |this, cx| {
 378                this.clear_block_below_cursor(cx);
 379                this.focus_handle(cx).focus(window);
 380            })
 381            .is_ok()
 382    }
 383
 384    fn insert_prompt_editor_into_terminal(
 385        &mut self,
 386        assist_id: TerminalInlineAssistId,
 387        height: u8,
 388        window: &mut Window,
 389        cx: &mut App,
 390    ) {
 391        if let Some(assist) = self.assists.get_mut(&assist_id) {
 392            if let Some(prompt_editor) = assist.prompt_editor.as_ref().cloned() {
 393                assist
 394                    .terminal
 395                    .update(cx, |terminal, cx| {
 396                        terminal.clear_block_below_cursor(cx);
 397                        let block = terminal_view::BlockProperties {
 398                            height,
 399                            render: Box::new(move |_| prompt_editor.clone().into_any_element()),
 400                        };
 401                        terminal.set_block_below_cursor(block, window, cx);
 402                    })
 403                    .log_err();
 404            }
 405        }
 406    }
 407}
 408
 409struct TerminalInlineAssist {
 410    terminal: WeakEntity<TerminalView>,
 411    prompt_editor: Option<Entity<PromptEditor>>,
 412    codegen: Entity<Codegen>,
 413    workspace: Option<WeakEntity<Workspace>>,
 414    include_context: bool,
 415    _subscriptions: Vec<Subscription>,
 416}
 417
 418impl TerminalInlineAssist {
 419    pub fn new(
 420        assist_id: TerminalInlineAssistId,
 421        terminal: &Entity<TerminalView>,
 422        include_context: bool,
 423        prompt_editor: Entity<PromptEditor>,
 424        workspace: Option<WeakEntity<Workspace>>,
 425        window: &mut Window,
 426        cx: &mut App,
 427    ) -> Self {
 428        let codegen = prompt_editor.read(cx).codegen.clone();
 429        Self {
 430            terminal: terminal.downgrade(),
 431            prompt_editor: Some(prompt_editor.clone()),
 432            codegen: codegen.clone(),
 433            workspace: workspace.clone(),
 434            include_context,
 435            _subscriptions: vec![
 436                window.subscribe(&prompt_editor, cx, |prompt_editor, event, window, cx| {
 437                    TerminalInlineAssistant::update_global(cx, |this, cx| {
 438                        this.handle_prompt_editor_event(prompt_editor, event, window, cx)
 439                    })
 440                }),
 441                window.subscribe(&codegen, cx, move |codegen, event, window, cx| {
 442                    TerminalInlineAssistant::update_global(cx, |this, cx| match event {
 443                        CodegenEvent::Finished => {
 444                            let assist = if let Some(assist) = this.assists.get(&assist_id) {
 445                                assist
 446                            } else {
 447                                return;
 448                            };
 449
 450                            if let CodegenStatus::Error(error) = &codegen.read(cx).status {
 451                                if assist.prompt_editor.is_none() {
 452                                    if let Some(workspace) = assist
 453                                        .workspace
 454                                        .as_ref()
 455                                        .and_then(|workspace| workspace.upgrade())
 456                                    {
 457                                        let error =
 458                                            format!("Terminal inline assistant error: {}", error);
 459                                        workspace.update(cx, |workspace, cx| {
 460                                            struct InlineAssistantError;
 461
 462                                            let id =
 463                                                NotificationId::composite::<InlineAssistantError>(
 464                                                    assist_id.0,
 465                                                );
 466
 467                                            workspace.show_toast(Toast::new(id, error), cx);
 468                                        })
 469                                    }
 470                                }
 471                            }
 472
 473                            if assist.prompt_editor.is_none() {
 474                                this.finish_assist(assist_id, false, false, window, cx);
 475                            }
 476                        }
 477                    })
 478                }),
 479            ],
 480        }
 481    }
 482}
 483
 484enum PromptEditorEvent {
 485    StartRequested,
 486    StopRequested,
 487    ConfirmRequested { execute: bool },
 488    CancelRequested,
 489    DismissRequested,
 490    Resized { height_in_lines: u8 },
 491}
 492
 493struct PromptEditor {
 494    id: TerminalInlineAssistId,
 495    height_in_lines: u8,
 496    editor: Entity<Editor>,
 497    language_model_selector: Entity<LanguageModelSelector>,
 498    edited_since_done: bool,
 499    prompt_history: VecDeque<String>,
 500    prompt_history_ix: Option<usize>,
 501    pending_prompt: String,
 502    codegen: Entity<Codegen>,
 503    _codegen_subscription: Subscription,
 504    editor_subscriptions: Vec<Subscription>,
 505    pending_token_count: Task<Result<()>>,
 506    token_count: Option<usize>,
 507    _token_count_subscriptions: Vec<Subscription>,
 508    workspace: Option<WeakEntity<Workspace>>,
 509}
 510
 511impl EventEmitter<PromptEditorEvent> for PromptEditor {}
 512
 513impl Render for PromptEditor {
 514    fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
 515        let status = &self.codegen.read(cx).status;
 516        let buttons = match status {
 517            CodegenStatus::Idle => {
 518                vec![
 519                    IconButton::new("cancel", IconName::Close)
 520                        .icon_color(Color::Muted)
 521                        .shape(IconButtonShape::Square)
 522                        .tooltip(|window, cx| {
 523                            Tooltip::for_action("Cancel Assist", &menu::Cancel, window, cx)
 524                        })
 525                        .on_click(
 526                            cx.listener(|_, _, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
 527                        ),
 528                    IconButton::new("start", IconName::SparkleAlt)
 529                        .icon_color(Color::Muted)
 530                        .shape(IconButtonShape::Square)
 531                        .tooltip(|window, cx| {
 532                            Tooltip::for_action("Generate", &menu::Confirm, window, cx)
 533                        })
 534                        .on_click(
 535                            cx.listener(|_, _, _, cx| cx.emit(PromptEditorEvent::StartRequested)),
 536                        ),
 537                ]
 538            }
 539            CodegenStatus::Pending => {
 540                vec![
 541                    IconButton::new("cancel", IconName::Close)
 542                        .icon_color(Color::Muted)
 543                        .shape(IconButtonShape::Square)
 544                        .tooltip(Tooltip::text("Cancel Assist"))
 545                        .on_click(
 546                            cx.listener(|_, _, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
 547                        ),
 548                    IconButton::new("stop", IconName::Stop)
 549                        .icon_color(Color::Error)
 550                        .shape(IconButtonShape::Square)
 551                        .tooltip(|window, cx| {
 552                            Tooltip::with_meta(
 553                                "Interrupt Generation",
 554                                Some(&menu::Cancel),
 555                                "Changes won't be discarded",
 556                                window,
 557                                cx,
 558                            )
 559                        })
 560                        .on_click(
 561                            cx.listener(|_, _, _, cx| cx.emit(PromptEditorEvent::StopRequested)),
 562                        ),
 563                ]
 564            }
 565            CodegenStatus::Error(_) | CodegenStatus::Done => {
 566                let cancel = IconButton::new("cancel", IconName::Close)
 567                    .icon_color(Color::Muted)
 568                    .shape(IconButtonShape::Square)
 569                    .tooltip(|window, cx| {
 570                        Tooltip::for_action("Cancel Assist", &menu::Cancel, window, cx)
 571                    })
 572                    .on_click(
 573                        cx.listener(|_, _, _, cx| cx.emit(PromptEditorEvent::CancelRequested)),
 574                    );
 575
 576                let has_error = matches!(status, CodegenStatus::Error(_));
 577                if has_error || self.edited_since_done {
 578                    vec![
 579                        cancel,
 580                        IconButton::new("restart", IconName::RotateCw)
 581                            .icon_color(Color::Info)
 582                            .shape(IconButtonShape::Square)
 583                            .tooltip(|window, cx| {
 584                                Tooltip::with_meta(
 585                                    "Restart Generation",
 586                                    Some(&menu::Confirm),
 587                                    "Changes will be discarded",
 588                                    window,
 589                                    cx,
 590                                )
 591                            })
 592                            .on_click(cx.listener(|_, _, _, cx| {
 593                                cx.emit(PromptEditorEvent::StartRequested);
 594                            })),
 595                    ]
 596                } else {
 597                    vec![
 598                        cancel,
 599                        IconButton::new("accept", IconName::Check)
 600                            .icon_color(Color::Info)
 601                            .shape(IconButtonShape::Square)
 602                            .tooltip(|window, cx| {
 603                                Tooltip::for_action(
 604                                    "Accept Generated Command",
 605                                    &menu::Confirm,
 606                                    window,
 607                                    cx,
 608                                )
 609                            })
 610                            .on_click(cx.listener(|_, _, _, cx| {
 611                                cx.emit(PromptEditorEvent::ConfirmRequested { execute: false });
 612                            })),
 613                        IconButton::new("confirm", IconName::Play)
 614                            .icon_color(Color::Info)
 615                            .shape(IconButtonShape::Square)
 616                            .tooltip(|window, cx| {
 617                                Tooltip::for_action(
 618                                    "Execute Generated Command",
 619                                    &menu::SecondaryConfirm,
 620                                    window,
 621                                    cx,
 622                                )
 623                            })
 624                            .on_click(cx.listener(|_, _, _, cx| {
 625                                cx.emit(PromptEditorEvent::ConfirmRequested { execute: true });
 626                            })),
 627                    ]
 628                }
 629            }
 630        };
 631
 632        h_flex()
 633            .bg(cx.theme().colors().editor_background)
 634            .border_y_1()
 635            .border_color(cx.theme().status().info_border)
 636            .py_2()
 637            .h_full()
 638            .w_full()
 639            .on_action(cx.listener(Self::confirm))
 640            .on_action(cx.listener(Self::secondary_confirm))
 641            .on_action(cx.listener(Self::cancel))
 642            .on_action(cx.listener(Self::move_up))
 643            .on_action(cx.listener(Self::move_down))
 644            .child(
 645                h_flex()
 646                    .w_12()
 647                    .justify_center()
 648                    .gap_2()
 649                    .child(LanguageModelSelectorPopoverMenu::new(
 650                        self.language_model_selector.clone(),
 651                        IconButton::new("change-model", IconName::SettingsAlt)
 652                            .shape(IconButtonShape::Square)
 653                            .icon_size(IconSize::Small)
 654                            .icon_color(Color::Muted),
 655                        move |window, cx| {
 656                            Tooltip::with_meta(
 657                                format!(
 658                                    "Using {}",
 659                                    LanguageModelRegistry::read_global(cx)
 660                                        .inline_assistant_model()
 661                                        .map(|inline_assistant| inline_assistant.model.name().0)
 662                                        .unwrap_or_else(|| "No model selected".into()),
 663                                ),
 664                                None,
 665                                "Change Model",
 666                                window,
 667                                cx,
 668                            )
 669                        },
 670                        gpui::Corner::TopRight,
 671                    ))
 672                    .children(
 673                        if let CodegenStatus::Error(error) = &self.codegen.read(cx).status {
 674                            let error_message = SharedString::from(error.to_string());
 675                            Some(
 676                                div()
 677                                    .id("error")
 678                                    .tooltip(Tooltip::text(error_message))
 679                                    .child(
 680                                        Icon::new(IconName::XCircle)
 681                                            .size(IconSize::Small)
 682                                            .color(Color::Error),
 683                                    ),
 684                            )
 685                        } else {
 686                            None
 687                        },
 688                    ),
 689            )
 690            .child(div().flex_1().child(self.render_prompt_editor(cx)))
 691            .child(
 692                h_flex()
 693                    .gap_1()
 694                    .pr_4()
 695                    .children(self.render_token_count(cx))
 696                    .children(buttons),
 697            )
 698    }
 699}
 700
 701impl Focusable for PromptEditor {
 702    fn focus_handle(&self, cx: &App) -> FocusHandle {
 703        self.editor.focus_handle(cx)
 704    }
 705}
 706
 707impl PromptEditor {
 708    const MAX_LINES: u8 = 8;
 709
 710    fn new(
 711        id: TerminalInlineAssistId,
 712        prompt_history: VecDeque<String>,
 713        prompt_buffer: Entity<MultiBuffer>,
 714        codegen: Entity<Codegen>,
 715        assistant_panel: Option<&Entity<AssistantPanel>>,
 716        workspace: Option<WeakEntity<Workspace>>,
 717        fs: Arc<dyn Fs>,
 718        window: &mut Window,
 719        cx: &mut Context<Self>,
 720    ) -> Self {
 721        let prompt_editor = cx.new(|cx| {
 722            let mut editor = Editor::new(
 723                EditorMode::AutoHeight {
 724                    max_lines: Self::MAX_LINES as usize,
 725                },
 726                prompt_buffer,
 727                None,
 728                window,
 729                cx,
 730            );
 731            editor.set_soft_wrap_mode(language::language_settings::SoftWrap::EditorWidth, cx);
 732            editor.set_placeholder_text(Self::placeholder_text(window, cx), cx);
 733            editor.set_context_menu_options(ContextMenuOptions {
 734                min_entries_visible: 12,
 735                max_entries_visible: 12,
 736                placement: None,
 737            });
 738            editor
 739        });
 740
 741        let mut token_count_subscriptions = Vec::new();
 742        if let Some(assistant_panel) = assistant_panel {
 743            token_count_subscriptions.push(cx.subscribe_in(
 744                assistant_panel,
 745                window,
 746                Self::handle_assistant_panel_event,
 747            ));
 748        }
 749
 750        let mut this = Self {
 751            id,
 752            height_in_lines: 1,
 753            editor: prompt_editor,
 754            language_model_selector: cx.new(|cx| {
 755                let fs = fs.clone();
 756                LanguageModelSelector::new(
 757                    |cx| LanguageModelRegistry::read_global(cx).default_model(),
 758                    move |model, cx| {
 759                        update_settings_file::<AssistantSettings>(
 760                            fs.clone(),
 761                            cx,
 762                            move |settings, _| settings.set_model(model.clone()),
 763                        );
 764                    },
 765                    window,
 766                    cx,
 767                )
 768            }),
 769            edited_since_done: false,
 770            prompt_history,
 771            prompt_history_ix: None,
 772            pending_prompt: String::new(),
 773            _codegen_subscription: cx.observe_in(&codegen, window, Self::handle_codegen_changed),
 774            editor_subscriptions: Vec::new(),
 775            codegen,
 776            pending_token_count: Task::ready(Ok(())),
 777            token_count: None,
 778            _token_count_subscriptions: token_count_subscriptions,
 779            workspace,
 780        };
 781        this.count_lines(cx);
 782        this.count_tokens(cx);
 783        this.subscribe_to_editor(cx);
 784        this
 785    }
 786
 787    fn placeholder_text(window: &Window, cx: &App) -> String {
 788        let context_keybinding = text_for_action(&zed_actions::assistant::ToggleFocus, window, cx)
 789            .map(|keybinding| format!("{keybinding} for context"))
 790            .unwrap_or_default();
 791
 792        format!("Generate…{context_keybinding} • ↓↑ for history")
 793    }
 794
 795    fn subscribe_to_editor(&mut self, cx: &mut Context<Self>) {
 796        self.editor_subscriptions.clear();
 797        self.editor_subscriptions
 798            .push(cx.observe(&self.editor, Self::handle_prompt_editor_changed));
 799        self.editor_subscriptions
 800            .push(cx.subscribe(&self.editor, Self::handle_prompt_editor_events));
 801    }
 802
 803    fn prompt(&self, cx: &App) -> String {
 804        self.editor.read(cx).text(cx)
 805    }
 806
 807    fn count_lines(&mut self, cx: &mut Context<Self>) {
 808        let height_in_lines = cmp::max(
 809            2, // Make the editor at least two lines tall, to account for padding and buttons.
 810            cmp::min(
 811                self.editor
 812                    .update(cx, |editor, cx| editor.max_point(cx).row().0 + 1),
 813                Self::MAX_LINES as u32,
 814            ),
 815        ) as u8;
 816
 817        if height_in_lines != self.height_in_lines {
 818            self.height_in_lines = height_in_lines;
 819            cx.emit(PromptEditorEvent::Resized { height_in_lines });
 820        }
 821    }
 822
 823    fn handle_assistant_panel_event(
 824        &mut self,
 825        _: &Entity<AssistantPanel>,
 826        event: &AssistantPanelEvent,
 827        _: &mut Window,
 828        cx: &mut Context<Self>,
 829    ) {
 830        let AssistantPanelEvent::ContextEdited { .. } = event;
 831        self.count_tokens(cx);
 832    }
 833
 834    fn count_tokens(&mut self, cx: &mut Context<Self>) {
 835        let assist_id = self.id;
 836        let Some(ConfiguredModel { model, .. }) =
 837            LanguageModelRegistry::read_global(cx).inline_assistant_model()
 838        else {
 839            return;
 840        };
 841        self.pending_token_count = cx.spawn(async move |this, cx| {
 842            cx.background_executor().timer(Duration::from_secs(1)).await;
 843            let request =
 844                cx.update_global(|inline_assistant: &mut TerminalInlineAssistant, cx| {
 845                    inline_assistant.request_for_inline_assist(assist_id, cx)
 846                })??;
 847
 848            let token_count = cx.update(|cx| model.count_tokens(request, cx))?.await?;
 849            this.update(cx, |this, cx| {
 850                this.token_count = Some(token_count);
 851                cx.notify();
 852            })
 853        })
 854    }
 855
 856    fn handle_prompt_editor_changed(&mut self, _: Entity<Editor>, cx: &mut Context<Self>) {
 857        self.count_lines(cx);
 858    }
 859
 860    fn handle_prompt_editor_events(
 861        &mut self,
 862        _: Entity<Editor>,
 863        event: &EditorEvent,
 864        cx: &mut Context<Self>,
 865    ) {
 866        match event {
 867            EditorEvent::Edited { .. } => {
 868                let prompt = self.editor.read(cx).text(cx);
 869                if self
 870                    .prompt_history_ix
 871                    .map_or(true, |ix| self.prompt_history[ix] != prompt)
 872                {
 873                    self.prompt_history_ix.take();
 874                    self.pending_prompt = prompt;
 875                }
 876
 877                self.edited_since_done = true;
 878                cx.notify();
 879            }
 880            EditorEvent::BufferEdited => {
 881                self.count_tokens(cx);
 882            }
 883            _ => {}
 884        }
 885    }
 886
 887    fn handle_codegen_changed(
 888        &mut self,
 889        _: Entity<Codegen>,
 890        _: &mut Window,
 891        cx: &mut Context<Self>,
 892    ) {
 893        match &self.codegen.read(cx).status {
 894            CodegenStatus::Idle => {
 895                self.editor
 896                    .update(cx, |editor, _| editor.set_read_only(false));
 897            }
 898            CodegenStatus::Pending => {
 899                self.editor
 900                    .update(cx, |editor, _| editor.set_read_only(true));
 901            }
 902            CodegenStatus::Done | CodegenStatus::Error(_) => {
 903                self.edited_since_done = false;
 904                self.editor
 905                    .update(cx, |editor, _| editor.set_read_only(false));
 906            }
 907        }
 908    }
 909
 910    fn cancel(&mut self, _: &editor::actions::Cancel, _: &mut Window, cx: &mut Context<Self>) {
 911        match &self.codegen.read(cx).status {
 912            CodegenStatus::Idle | CodegenStatus::Done | CodegenStatus::Error(_) => {
 913                cx.emit(PromptEditorEvent::CancelRequested);
 914            }
 915            CodegenStatus::Pending => {
 916                cx.emit(PromptEditorEvent::StopRequested);
 917            }
 918        }
 919    }
 920
 921    fn confirm(&mut self, _: &menu::Confirm, _: &mut Window, cx: &mut Context<Self>) {
 922        match &self.codegen.read(cx).status {
 923            CodegenStatus::Idle => {
 924                if !self.editor.read(cx).text(cx).trim().is_empty() {
 925                    cx.emit(PromptEditorEvent::StartRequested);
 926                }
 927            }
 928            CodegenStatus::Pending => {
 929                cx.emit(PromptEditorEvent::DismissRequested);
 930            }
 931            CodegenStatus::Done => {
 932                if self.edited_since_done {
 933                    cx.emit(PromptEditorEvent::StartRequested);
 934                } else {
 935                    cx.emit(PromptEditorEvent::ConfirmRequested { execute: false });
 936                }
 937            }
 938            CodegenStatus::Error(_) => {
 939                cx.emit(PromptEditorEvent::StartRequested);
 940            }
 941        }
 942    }
 943
 944    fn secondary_confirm(
 945        &mut self,
 946        _: &menu::SecondaryConfirm,
 947        _: &mut Window,
 948        cx: &mut Context<Self>,
 949    ) {
 950        if matches!(self.codegen.read(cx).status, CodegenStatus::Done) {
 951            cx.emit(PromptEditorEvent::ConfirmRequested { execute: true });
 952        }
 953    }
 954
 955    fn move_up(&mut self, _: &MoveUp, window: &mut Window, cx: &mut Context<Self>) {
 956        if let Some(ix) = self.prompt_history_ix {
 957            if ix > 0 {
 958                self.prompt_history_ix = Some(ix - 1);
 959                let prompt = self.prompt_history[ix - 1].as_str();
 960                self.editor.update(cx, |editor, cx| {
 961                    editor.set_text(prompt, window, cx);
 962                    editor.move_to_beginning(&Default::default(), window, cx);
 963                });
 964            }
 965        } else if !self.prompt_history.is_empty() {
 966            self.prompt_history_ix = Some(self.prompt_history.len() - 1);
 967            let prompt = self.prompt_history[self.prompt_history.len() - 1].as_str();
 968            self.editor.update(cx, |editor, cx| {
 969                editor.set_text(prompt, window, cx);
 970                editor.move_to_beginning(&Default::default(), window, cx);
 971            });
 972        }
 973    }
 974
 975    fn move_down(&mut self, _: &MoveDown, window: &mut Window, cx: &mut Context<Self>) {
 976        if let Some(ix) = self.prompt_history_ix {
 977            if ix < self.prompt_history.len() - 1 {
 978                self.prompt_history_ix = Some(ix + 1);
 979                let prompt = self.prompt_history[ix + 1].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            } else {
 985                self.prompt_history_ix = None;
 986                let prompt = self.pending_prompt.as_str();
 987                self.editor.update(cx, |editor, cx| {
 988                    editor.set_text(prompt, window, cx);
 989                    editor.move_to_end(&Default::default(), window, cx)
 990                });
 991            }
 992        }
 993    }
 994
 995    fn render_token_count(&self, cx: &mut Context<Self>) -> Option<impl IntoElement> {
 996        let model = LanguageModelRegistry::read_global(cx)
 997            .inline_assistant_model()?
 998            .model;
 999        let token_count = self.token_count?;
1000        let max_token_count = model.max_token_count();
1001
1002        let remaining_tokens = max_token_count as isize - token_count as isize;
1003        let token_count_color = if remaining_tokens <= 0 {
1004            Color::Error
1005        } else if token_count as f32 / max_token_count as f32 >= 0.8 {
1006            Color::Warning
1007        } else {
1008            Color::Muted
1009        };
1010
1011        let mut token_count = h_flex()
1012            .id("token_count")
1013            .gap_0p5()
1014            .child(
1015                Label::new(humanize_token_count(token_count))
1016                    .size(LabelSize::Small)
1017                    .color(token_count_color),
1018            )
1019            .child(Label::new("/").size(LabelSize::Small).color(Color::Muted))
1020            .child(
1021                Label::new(humanize_token_count(max_token_count))
1022                    .size(LabelSize::Small)
1023                    .color(Color::Muted),
1024            );
1025        if let Some(workspace) = self.workspace.clone() {
1026            token_count = token_count
1027                .tooltip(|window, cx| {
1028                    Tooltip::with_meta(
1029                        "Tokens Used by Inline Assistant",
1030                        None,
1031                        "Click to Open Assistant Panel",
1032                        window,
1033                        cx,
1034                    )
1035                })
1036                .cursor_pointer()
1037                .on_mouse_down(gpui::MouseButton::Left, |_, _, cx| cx.stop_propagation())
1038                .on_click(move |_, window, cx| {
1039                    cx.stop_propagation();
1040                    workspace
1041                        .update(cx, |workspace, cx| {
1042                            workspace.focus_panel::<AssistantPanel>(window, cx)
1043                        })
1044                        .ok();
1045                });
1046        } else {
1047            token_count = token_count
1048                .cursor_default()
1049                .tooltip(Tooltip::text("Tokens Used by Inline Assistant"));
1050        }
1051
1052        Some(token_count)
1053    }
1054
1055    fn render_prompt_editor(&self, cx: &mut Context<Self>) -> impl IntoElement {
1056        let settings = ThemeSettings::get_global(cx);
1057        let text_style = TextStyle {
1058            color: if self.editor.read(cx).read_only(cx) {
1059                cx.theme().colors().text_disabled
1060            } else {
1061                cx.theme().colors().text
1062            },
1063            font_family: settings.buffer_font.family.clone(),
1064            font_fallbacks: settings.buffer_font.fallbacks.clone(),
1065            font_size: settings.buffer_font_size(cx).into(),
1066            font_weight: settings.buffer_font.weight,
1067            line_height: relative(settings.buffer_line_height.value()),
1068            ..Default::default()
1069        };
1070        EditorElement::new(
1071            &self.editor,
1072            EditorStyle {
1073                background: cx.theme().colors().editor_background,
1074                local_player: cx.theme().players().local(),
1075                text: text_style,
1076                ..Default::default()
1077            },
1078        )
1079    }
1080}
1081
1082#[derive(Debug)]
1083pub enum CodegenEvent {
1084    Finished,
1085}
1086
1087impl EventEmitter<CodegenEvent> for Codegen {}
1088
1089#[cfg(not(target_os = "windows"))]
1090const CLEAR_INPUT: &str = "\x15";
1091#[cfg(target_os = "windows")]
1092const CLEAR_INPUT: &str = "\x03";
1093const CARRIAGE_RETURN: &str = "\x0d";
1094
1095struct TerminalTransaction {
1096    terminal: Entity<Terminal>,
1097}
1098
1099impl TerminalTransaction {
1100    pub fn start(terminal: Entity<Terminal>) -> Self {
1101        Self { terminal }
1102    }
1103
1104    pub fn push(&mut self, hunk: String, cx: &mut App) {
1105        // Ensure that the assistant cannot accidentally execute commands that are streamed into the terminal
1106        let input = Self::sanitize_input(hunk);
1107        self.terminal
1108            .update(cx, |terminal, _| terminal.input(input));
1109    }
1110
1111    pub fn undo(&self, cx: &mut App) {
1112        self.terminal
1113            .update(cx, |terminal, _| terminal.input(CLEAR_INPUT.to_string()));
1114    }
1115
1116    pub fn complete(&self, cx: &mut App) {
1117        self.terminal.update(cx, |terminal, _| {
1118            terminal.input(CARRIAGE_RETURN.to_string())
1119        });
1120    }
1121
1122    fn sanitize_input(input: String) -> String {
1123        input.replace(['\r', '\n'], "")
1124    }
1125}
1126
1127pub struct Codegen {
1128    status: CodegenStatus,
1129    telemetry: Option<Arc<Telemetry>>,
1130    terminal: Entity<Terminal>,
1131    generation: Task<()>,
1132    message_id: Option<String>,
1133    transaction: Option<TerminalTransaction>,
1134}
1135
1136impl Codegen {
1137    pub fn new(terminal: Entity<Terminal>, telemetry: Option<Arc<Telemetry>>) -> Self {
1138        Self {
1139            terminal,
1140            telemetry,
1141            status: CodegenStatus::Idle,
1142            generation: Task::ready(()),
1143            message_id: None,
1144            transaction: None,
1145        }
1146    }
1147
1148    pub fn start(&mut self, prompt: LanguageModelRequest, cx: &mut Context<Self>) {
1149        let Some(ConfiguredModel { model, .. }) =
1150            LanguageModelRegistry::read_global(cx).inline_assistant_model()
1151        else {
1152            return;
1153        };
1154
1155        let model_api_key = model.api_key(cx);
1156        let http_client = cx.http_client();
1157        let telemetry = self.telemetry.clone();
1158        self.status = CodegenStatus::Pending;
1159        self.transaction = Some(TerminalTransaction::start(self.terminal.clone()));
1160        self.generation = cx.spawn(async move |this, cx| {
1161            let model_telemetry_id = model.telemetry_id();
1162            let model_provider_id = model.provider_id();
1163            let response = model.stream_completion_text(prompt, &cx).await;
1164            let generate = async {
1165                let message_id = response
1166                    .as_ref()
1167                    .ok()
1168                    .and_then(|response| response.message_id.clone());
1169
1170                let (mut hunks_tx, mut hunks_rx) = mpsc::channel(1);
1171
1172                let task = cx.background_spawn({
1173                    let message_id = message_id.clone();
1174                    let executor = cx.background_executor().clone();
1175                    async move {
1176                        let mut response_latency = None;
1177                        let request_start = Instant::now();
1178                        let task = async {
1179                            let mut chunks = response?.stream;
1180                            while let Some(chunk) = chunks.next().await {
1181                                if response_latency.is_none() {
1182                                    response_latency = Some(request_start.elapsed());
1183                                }
1184                                let chunk = chunk?;
1185                                hunks_tx.send(chunk).await?;
1186                            }
1187
1188                            anyhow::Ok(())
1189                        };
1190
1191                        let result = task.await;
1192
1193                        let error_message = result.as_ref().err().map(|error| error.to_string());
1194                        report_assistant_event(
1195                            AssistantEventData {
1196                                conversation_id: None,
1197                                kind: AssistantKind::InlineTerminal,
1198                                message_id,
1199                                phase: AssistantPhase::Response,
1200                                model: model_telemetry_id,
1201                                model_provider: model_provider_id.to_string(),
1202                                response_latency,
1203                                error_message,
1204                                language_name: None,
1205                            },
1206                            telemetry,
1207                            http_client,
1208                            model_api_key,
1209                            &executor,
1210                        );
1211
1212                        result?;
1213                        anyhow::Ok(())
1214                    }
1215                });
1216
1217                this.update(cx, |this, _| {
1218                    this.message_id = message_id;
1219                })?;
1220
1221                while let Some(hunk) = hunks_rx.next().await {
1222                    this.update(cx, |this, cx| {
1223                        if let Some(transaction) = &mut this.transaction {
1224                            transaction.push(hunk, cx);
1225                            cx.notify();
1226                        }
1227                    })?;
1228                }
1229
1230                task.await?;
1231                anyhow::Ok(())
1232            };
1233
1234            let result = generate.await;
1235
1236            this.update(cx, |this, cx| {
1237                if let Err(error) = result {
1238                    this.status = CodegenStatus::Error(error);
1239                } else {
1240                    this.status = CodegenStatus::Done;
1241                }
1242                cx.emit(CodegenEvent::Finished);
1243                cx.notify();
1244            })
1245            .ok();
1246        });
1247        cx.notify();
1248    }
1249
1250    pub fn stop(&mut self, cx: &mut Context<Self>) {
1251        self.status = CodegenStatus::Done;
1252        self.generation = Task::ready(());
1253        cx.emit(CodegenEvent::Finished);
1254        cx.notify();
1255    }
1256
1257    pub fn complete(&mut self, cx: &mut Context<Self>) {
1258        if let Some(transaction) = self.transaction.take() {
1259            transaction.complete(cx);
1260        }
1261    }
1262
1263    pub fn undo(&mut self, cx: &mut Context<Self>) {
1264        if let Some(transaction) = self.transaction.take() {
1265            transaction.undo(cx);
1266        }
1267    }
1268}
1269
1270enum CodegenStatus {
1271    Idle,
1272    Pending,
1273    Done,
1274    Error(anyhow::Error),
1275}