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            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
 734        });
 735
 736        let mut token_count_subscriptions = Vec::new();
 737        if let Some(assistant_panel) = assistant_panel {
 738            token_count_subscriptions.push(cx.subscribe_in(
 739                assistant_panel,
 740                window,
 741                Self::handle_assistant_panel_event,
 742            ));
 743        }
 744
 745        let mut this = Self {
 746            id,
 747            height_in_lines: 1,
 748            editor: prompt_editor,
 749            language_model_selector: cx.new(|cx| {
 750                let fs = fs.clone();
 751                LanguageModelSelector::new(
 752                    |cx| LanguageModelRegistry::read_global(cx).default_model(),
 753                    move |model, cx| {
 754                        update_settings_file::<AssistantSettings>(
 755                            fs.clone(),
 756                            cx,
 757                            move |settings, _| settings.set_model(model.clone()),
 758                        );
 759                    },
 760                    window,
 761                    cx,
 762                )
 763            }),
 764            edited_since_done: false,
 765            prompt_history,
 766            prompt_history_ix: None,
 767            pending_prompt: String::new(),
 768            _codegen_subscription: cx.observe_in(&codegen, window, Self::handle_codegen_changed),
 769            editor_subscriptions: Vec::new(),
 770            codegen,
 771            pending_token_count: Task::ready(Ok(())),
 772            token_count: None,
 773            _token_count_subscriptions: token_count_subscriptions,
 774            workspace,
 775        };
 776        this.count_lines(cx);
 777        this.count_tokens(cx);
 778        this.subscribe_to_editor(cx);
 779        this
 780    }
 781
 782    fn placeholder_text(window: &Window, cx: &App) -> String {
 783        let context_keybinding = text_for_action(&zed_actions::assistant::ToggleFocus, window, cx)
 784            .map(|keybinding| format!("{keybinding} for context"))
 785            .unwrap_or_default();
 786
 787        format!("Generate…{context_keybinding} • ↓↑ for history")
 788    }
 789
 790    fn subscribe_to_editor(&mut self, cx: &mut Context<Self>) {
 791        self.editor_subscriptions.clear();
 792        self.editor_subscriptions
 793            .push(cx.observe(&self.editor, Self::handle_prompt_editor_changed));
 794        self.editor_subscriptions
 795            .push(cx.subscribe(&self.editor, Self::handle_prompt_editor_events));
 796    }
 797
 798    fn prompt(&self, cx: &App) -> String {
 799        self.editor.read(cx).text(cx)
 800    }
 801
 802    fn count_lines(&mut self, cx: &mut Context<Self>) {
 803        let height_in_lines = cmp::max(
 804            2, // Make the editor at least two lines tall, to account for padding and buttons.
 805            cmp::min(
 806                self.editor
 807                    .update(cx, |editor, cx| editor.max_point(cx).row().0 + 1),
 808                Self::MAX_LINES as u32,
 809            ),
 810        ) as u8;
 811
 812        if height_in_lines != self.height_in_lines {
 813            self.height_in_lines = height_in_lines;
 814            cx.emit(PromptEditorEvent::Resized { height_in_lines });
 815        }
 816    }
 817
 818    fn handle_assistant_panel_event(
 819        &mut self,
 820        _: &Entity<AssistantPanel>,
 821        event: &AssistantPanelEvent,
 822        _: &mut Window,
 823        cx: &mut Context<Self>,
 824    ) {
 825        let AssistantPanelEvent::ContextEdited { .. } = event;
 826        self.count_tokens(cx);
 827    }
 828
 829    fn count_tokens(&mut self, cx: &mut Context<Self>) {
 830        let assist_id = self.id;
 831        let Some(ConfiguredModel { model, .. }) =
 832            LanguageModelRegistry::read_global(cx).inline_assistant_model()
 833        else {
 834            return;
 835        };
 836        self.pending_token_count = cx.spawn(async move |this, cx| {
 837            cx.background_executor().timer(Duration::from_secs(1)).await;
 838            let request =
 839                cx.update_global(|inline_assistant: &mut TerminalInlineAssistant, cx| {
 840                    inline_assistant.request_for_inline_assist(assist_id, cx)
 841                })??;
 842
 843            let token_count = cx.update(|cx| model.count_tokens(request, cx))?.await?;
 844            this.update(cx, |this, cx| {
 845                this.token_count = Some(token_count);
 846                cx.notify();
 847            })
 848        })
 849    }
 850
 851    fn handle_prompt_editor_changed(&mut self, _: Entity<Editor>, cx: &mut Context<Self>) {
 852        self.count_lines(cx);
 853    }
 854
 855    fn handle_prompt_editor_events(
 856        &mut self,
 857        _: Entity<Editor>,
 858        event: &EditorEvent,
 859        cx: &mut Context<Self>,
 860    ) {
 861        match event {
 862            EditorEvent::Edited { .. } => {
 863                let prompt = self.editor.read(cx).text(cx);
 864                if self
 865                    .prompt_history_ix
 866                    .map_or(true, |ix| self.prompt_history[ix] != prompt)
 867                {
 868                    self.prompt_history_ix.take();
 869                    self.pending_prompt = prompt;
 870                }
 871
 872                self.edited_since_done = true;
 873                cx.notify();
 874            }
 875            EditorEvent::BufferEdited => {
 876                self.count_tokens(cx);
 877            }
 878            _ => {}
 879        }
 880    }
 881
 882    fn handle_codegen_changed(
 883        &mut self,
 884        _: Entity<Codegen>,
 885        _: &mut Window,
 886        cx: &mut Context<Self>,
 887    ) {
 888        match &self.codegen.read(cx).status {
 889            CodegenStatus::Idle => {
 890                self.editor
 891                    .update(cx, |editor, _| editor.set_read_only(false));
 892            }
 893            CodegenStatus::Pending => {
 894                self.editor
 895                    .update(cx, |editor, _| editor.set_read_only(true));
 896            }
 897            CodegenStatus::Done | CodegenStatus::Error(_) => {
 898                self.edited_since_done = false;
 899                self.editor
 900                    .update(cx, |editor, _| editor.set_read_only(false));
 901            }
 902        }
 903    }
 904
 905    fn cancel(&mut self, _: &editor::actions::Cancel, _: &mut Window, cx: &mut Context<Self>) {
 906        match &self.codegen.read(cx).status {
 907            CodegenStatus::Idle | CodegenStatus::Done | CodegenStatus::Error(_) => {
 908                cx.emit(PromptEditorEvent::CancelRequested);
 909            }
 910            CodegenStatus::Pending => {
 911                cx.emit(PromptEditorEvent::StopRequested);
 912            }
 913        }
 914    }
 915
 916    fn confirm(&mut self, _: &menu::Confirm, _: &mut Window, cx: &mut Context<Self>) {
 917        match &self.codegen.read(cx).status {
 918            CodegenStatus::Idle => {
 919                if !self.editor.read(cx).text(cx).trim().is_empty() {
 920                    cx.emit(PromptEditorEvent::StartRequested);
 921                }
 922            }
 923            CodegenStatus::Pending => {
 924                cx.emit(PromptEditorEvent::DismissRequested);
 925            }
 926            CodegenStatus::Done => {
 927                if self.edited_since_done {
 928                    cx.emit(PromptEditorEvent::StartRequested);
 929                } else {
 930                    cx.emit(PromptEditorEvent::ConfirmRequested { execute: false });
 931                }
 932            }
 933            CodegenStatus::Error(_) => {
 934                cx.emit(PromptEditorEvent::StartRequested);
 935            }
 936        }
 937    }
 938
 939    fn secondary_confirm(
 940        &mut self,
 941        _: &menu::SecondaryConfirm,
 942        _: &mut Window,
 943        cx: &mut Context<Self>,
 944    ) {
 945        if matches!(self.codegen.read(cx).status, CodegenStatus::Done) {
 946            cx.emit(PromptEditorEvent::ConfirmRequested { execute: true });
 947        }
 948    }
 949
 950    fn move_up(&mut self, _: &MoveUp, window: &mut Window, cx: &mut Context<Self>) {
 951        if let Some(ix) = self.prompt_history_ix {
 952            if ix > 0 {
 953                self.prompt_history_ix = Some(ix - 1);
 954                let prompt = self.prompt_history[ix - 1].as_str();
 955                self.editor.update(cx, |editor, cx| {
 956                    editor.set_text(prompt, window, cx);
 957                    editor.move_to_beginning(&Default::default(), window, cx);
 958                });
 959            }
 960        } else if !self.prompt_history.is_empty() {
 961            self.prompt_history_ix = Some(self.prompt_history.len() - 1);
 962            let prompt = self.prompt_history[self.prompt_history.len() - 1].as_str();
 963            self.editor.update(cx, |editor, cx| {
 964                editor.set_text(prompt, window, cx);
 965                editor.move_to_beginning(&Default::default(), window, cx);
 966            });
 967        }
 968    }
 969
 970    fn move_down(&mut self, _: &MoveDown, window: &mut Window, cx: &mut Context<Self>) {
 971        if let Some(ix) = self.prompt_history_ix {
 972            if ix < self.prompt_history.len() - 1 {
 973                self.prompt_history_ix = Some(ix + 1);
 974                let prompt = self.prompt_history[ix + 1].as_str();
 975                self.editor.update(cx, |editor, cx| {
 976                    editor.set_text(prompt, window, cx);
 977                    editor.move_to_end(&Default::default(), window, cx)
 978                });
 979            } else {
 980                self.prompt_history_ix = None;
 981                let prompt = self.pending_prompt.as_str();
 982                self.editor.update(cx, |editor, cx| {
 983                    editor.set_text(prompt, window, cx);
 984                    editor.move_to_end(&Default::default(), window, cx)
 985                });
 986            }
 987        }
 988    }
 989
 990    fn render_token_count(&self, cx: &mut Context<Self>) -> Option<impl IntoElement> {
 991        let model = LanguageModelRegistry::read_global(cx)
 992            .inline_assistant_model()?
 993            .model;
 994        let token_count = self.token_count?;
 995        let max_token_count = model.max_token_count();
 996
 997        let remaining_tokens = max_token_count as isize - token_count as isize;
 998        let token_count_color = if remaining_tokens <= 0 {
 999            Color::Error
1000        } else if token_count as f32 / max_token_count as f32 >= 0.8 {
1001            Color::Warning
1002        } else {
1003            Color::Muted
1004        };
1005
1006        let mut token_count = h_flex()
1007            .id("token_count")
1008            .gap_0p5()
1009            .child(
1010                Label::new(humanize_token_count(token_count))
1011                    .size(LabelSize::Small)
1012                    .color(token_count_color),
1013            )
1014            .child(Label::new("/").size(LabelSize::Small).color(Color::Muted))
1015            .child(
1016                Label::new(humanize_token_count(max_token_count))
1017                    .size(LabelSize::Small)
1018                    .color(Color::Muted),
1019            );
1020        if let Some(workspace) = self.workspace.clone() {
1021            token_count = token_count
1022                .tooltip(|window, cx| {
1023                    Tooltip::with_meta(
1024                        "Tokens Used by Inline Assistant",
1025                        None,
1026                        "Click to Open Assistant Panel",
1027                        window,
1028                        cx,
1029                    )
1030                })
1031                .cursor_pointer()
1032                .on_mouse_down(gpui::MouseButton::Left, |_, _, cx| cx.stop_propagation())
1033                .on_click(move |_, window, cx| {
1034                    cx.stop_propagation();
1035                    workspace
1036                        .update(cx, |workspace, cx| {
1037                            workspace.focus_panel::<AssistantPanel>(window, cx)
1038                        })
1039                        .ok();
1040                });
1041        } else {
1042            token_count = token_count
1043                .cursor_default()
1044                .tooltip(Tooltip::text("Tokens Used by Inline Assistant"));
1045        }
1046
1047        Some(token_count)
1048    }
1049
1050    fn render_prompt_editor(&self, cx: &mut Context<Self>) -> impl IntoElement {
1051        let settings = ThemeSettings::get_global(cx);
1052        let text_style = TextStyle {
1053            color: if self.editor.read(cx).read_only(cx) {
1054                cx.theme().colors().text_disabled
1055            } else {
1056                cx.theme().colors().text
1057            },
1058            font_family: settings.buffer_font.family.clone(),
1059            font_fallbacks: settings.buffer_font.fallbacks.clone(),
1060            font_size: settings.buffer_font_size(cx).into(),
1061            font_weight: settings.buffer_font.weight,
1062            line_height: relative(settings.buffer_line_height.value()),
1063            ..Default::default()
1064        };
1065        EditorElement::new(
1066            &self.editor,
1067            EditorStyle {
1068                background: cx.theme().colors().editor_background,
1069                local_player: cx.theme().players().local(),
1070                text: text_style,
1071                ..Default::default()
1072            },
1073        )
1074    }
1075}
1076
1077#[derive(Debug)]
1078pub enum CodegenEvent {
1079    Finished,
1080}
1081
1082impl EventEmitter<CodegenEvent> for Codegen {}
1083
1084#[cfg(not(target_os = "windows"))]
1085const CLEAR_INPUT: &str = "\x15";
1086#[cfg(target_os = "windows")]
1087const CLEAR_INPUT: &str = "\x03";
1088const CARRIAGE_RETURN: &str = "\x0d";
1089
1090struct TerminalTransaction {
1091    terminal: Entity<Terminal>,
1092}
1093
1094impl TerminalTransaction {
1095    pub fn start(terminal: Entity<Terminal>) -> Self {
1096        Self { terminal }
1097    }
1098
1099    pub fn push(&mut self, hunk: String, cx: &mut App) {
1100        // Ensure that the assistant cannot accidentally execute commands that are streamed into the terminal
1101        let input = Self::sanitize_input(hunk);
1102        self.terminal
1103            .update(cx, |terminal, _| terminal.input(input));
1104    }
1105
1106    pub fn undo(&self, cx: &mut App) {
1107        self.terminal
1108            .update(cx, |terminal, _| terminal.input(CLEAR_INPUT.to_string()));
1109    }
1110
1111    pub fn complete(&self, cx: &mut App) {
1112        self.terminal.update(cx, |terminal, _| {
1113            terminal.input(CARRIAGE_RETURN.to_string())
1114        });
1115    }
1116
1117    fn sanitize_input(input: String) -> String {
1118        input.replace(['\r', '\n'], "")
1119    }
1120}
1121
1122pub struct Codegen {
1123    status: CodegenStatus,
1124    telemetry: Option<Arc<Telemetry>>,
1125    terminal: Entity<Terminal>,
1126    generation: Task<()>,
1127    message_id: Option<String>,
1128    transaction: Option<TerminalTransaction>,
1129}
1130
1131impl Codegen {
1132    pub fn new(terminal: Entity<Terminal>, telemetry: Option<Arc<Telemetry>>) -> Self {
1133        Self {
1134            terminal,
1135            telemetry,
1136            status: CodegenStatus::Idle,
1137            generation: Task::ready(()),
1138            message_id: None,
1139            transaction: None,
1140        }
1141    }
1142
1143    pub fn start(&mut self, prompt: LanguageModelRequest, cx: &mut Context<Self>) {
1144        let Some(ConfiguredModel { model, .. }) =
1145            LanguageModelRegistry::read_global(cx).inline_assistant_model()
1146        else {
1147            return;
1148        };
1149
1150        let model_api_key = model.api_key(cx);
1151        let http_client = cx.http_client();
1152        let telemetry = self.telemetry.clone();
1153        self.status = CodegenStatus::Pending;
1154        self.transaction = Some(TerminalTransaction::start(self.terminal.clone()));
1155        self.generation = cx.spawn(async move |this, cx| {
1156            let model_telemetry_id = model.telemetry_id();
1157            let model_provider_id = model.provider_id();
1158            let response = model.stream_completion_text(prompt, &cx).await;
1159            let generate = async {
1160                let message_id = response
1161                    .as_ref()
1162                    .ok()
1163                    .and_then(|response| response.message_id.clone());
1164
1165                let (mut hunks_tx, mut hunks_rx) = mpsc::channel(1);
1166
1167                let task = cx.background_spawn({
1168                    let message_id = message_id.clone();
1169                    let executor = cx.background_executor().clone();
1170                    async move {
1171                        let mut response_latency = None;
1172                        let request_start = Instant::now();
1173                        let task = async {
1174                            let mut chunks = response?.stream;
1175                            while let Some(chunk) = chunks.next().await {
1176                                if response_latency.is_none() {
1177                                    response_latency = Some(request_start.elapsed());
1178                                }
1179                                let chunk = chunk?;
1180                                hunks_tx.send(chunk).await?;
1181                            }
1182
1183                            anyhow::Ok(())
1184                        };
1185
1186                        let result = task.await;
1187
1188                        let error_message = result.as_ref().err().map(|error| error.to_string());
1189                        report_assistant_event(
1190                            AssistantEventData {
1191                                conversation_id: None,
1192                                kind: AssistantKind::InlineTerminal,
1193                                message_id,
1194                                phase: AssistantPhase::Response,
1195                                model: model_telemetry_id,
1196                                model_provider: model_provider_id.to_string(),
1197                                response_latency,
1198                                error_message,
1199                                language_name: None,
1200                            },
1201                            telemetry,
1202                            http_client,
1203                            model_api_key,
1204                            &executor,
1205                        );
1206
1207                        result?;
1208                        anyhow::Ok(())
1209                    }
1210                });
1211
1212                this.update(cx, |this, _| {
1213                    this.message_id = message_id;
1214                })?;
1215
1216                while let Some(hunk) = hunks_rx.next().await {
1217                    this.update(cx, |this, cx| {
1218                        if let Some(transaction) = &mut this.transaction {
1219                            transaction.push(hunk, cx);
1220                            cx.notify();
1221                        }
1222                    })?;
1223                }
1224
1225                task.await?;
1226                anyhow::Ok(())
1227            };
1228
1229            let result = generate.await;
1230
1231            this.update(cx, |this, cx| {
1232                if let Err(error) = result {
1233                    this.status = CodegenStatus::Error(error);
1234                } else {
1235                    this.status = CodegenStatus::Done;
1236                }
1237                cx.emit(CodegenEvent::Finished);
1238                cx.notify();
1239            })
1240            .ok();
1241        });
1242        cx.notify();
1243    }
1244
1245    pub fn stop(&mut self, cx: &mut Context<Self>) {
1246        self.status = CodegenStatus::Done;
1247        self.generation = Task::ready(());
1248        cx.emit(CodegenEvent::Finished);
1249        cx.notify();
1250    }
1251
1252    pub fn complete(&mut self, cx: &mut Context<Self>) {
1253        if let Some(transaction) = self.transaction.take() {
1254            transaction.complete(cx);
1255        }
1256    }
1257
1258    pub fn undo(&mut self, cx: &mut Context<Self>) {
1259        if let Some(transaction) = self.transaction.take() {
1260            transaction.undo(cx);
1261        }
1262    }
1263}
1264
1265enum CodegenStatus {
1266    Idle,
1267    Pending,
1268    Done,
1269    Error(anyhow::Error),
1270}