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