terminal_inline_assistant.rs

  1use crate::{
  2    acp::AcpThreadHistory,
  3    context::load_context,
  4    inline_prompt_editor::{
  5        CodegenStatus, PromptEditor, PromptEditorEvent, TerminalInlineAssistId,
  6    },
  7    terminal_codegen::{CLEAR_INPUT, CodegenEvent, TerminalCodegen},
  8};
  9use agent::ThreadStore;
 10use agent_settings::AgentSettings;
 11use anyhow::{Context as _, Result};
 12
 13use cloud_llm_client::CompletionIntent;
 14use collections::{HashMap, VecDeque};
 15use editor::{MultiBuffer, actions::SelectAll};
 16use fs::Fs;
 17use gpui::{App, Entity, Focusable, Global, Subscription, Task, UpdateGlobal, WeakEntity};
 18use language::Buffer;
 19use language_model::{
 20    ConfiguredModel, LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage,
 21    Role, report_anthropic_event,
 22};
 23use project::Project;
 24use prompt_store::{PromptBuilder, PromptStore};
 25use std::sync::Arc;
 26use terminal_view::TerminalView;
 27use ui::prelude::*;
 28use util::ResultExt;
 29use uuid::Uuid;
 30use workspace::{Toast, Workspace, notifications::NotificationId};
 31
 32pub fn init(fs: Arc<dyn Fs>, prompt_builder: Arc<PromptBuilder>, cx: &mut App) {
 33    cx.set_global(TerminalInlineAssistant::new(fs, prompt_builder));
 34}
 35
 36const DEFAULT_CONTEXT_LINES: usize = 50;
 37const PROMPT_HISTORY_MAX_LEN: usize = 20;
 38
 39pub struct TerminalInlineAssistant {
 40    next_assist_id: TerminalInlineAssistId,
 41    assists: HashMap<TerminalInlineAssistId, TerminalInlineAssist>,
 42    prompt_history: VecDeque<String>,
 43    fs: Arc<dyn Fs>,
 44    prompt_builder: Arc<PromptBuilder>,
 45}
 46
 47impl Global for TerminalInlineAssistant {}
 48
 49impl TerminalInlineAssistant {
 50    pub fn new(fs: Arc<dyn Fs>, prompt_builder: Arc<PromptBuilder>) -> Self {
 51        Self {
 52            next_assist_id: TerminalInlineAssistId::default(),
 53            assists: HashMap::default(),
 54            prompt_history: VecDeque::default(),
 55            fs,
 56            prompt_builder,
 57        }
 58    }
 59
 60    pub fn assist(
 61        &mut self,
 62        terminal_view: &Entity<TerminalView>,
 63        workspace: WeakEntity<Workspace>,
 64        project: WeakEntity<Project>,
 65        thread_store: Entity<ThreadStore>,
 66        prompt_store: Option<Entity<PromptStore>>,
 67        history: WeakEntity<AcpThreadHistory>,
 68        initial_prompt: Option<String>,
 69        window: &mut Window,
 70        cx: &mut App,
 71    ) {
 72        let terminal = terminal_view.read(cx).terminal().clone();
 73        let assist_id = self.next_assist_id.post_inc();
 74        let session_id = Uuid::new_v4();
 75        let prompt_buffer = cx.new(|cx| {
 76            MultiBuffer::singleton(
 77                cx.new(|cx| Buffer::local(initial_prompt.unwrap_or_default(), cx)),
 78                cx,
 79            )
 80        });
 81        let codegen = cx.new(|_| TerminalCodegen::new(terminal, session_id));
 82
 83        let prompt_editor = cx.new(|cx| {
 84            PromptEditor::new_terminal(
 85                assist_id,
 86                self.prompt_history.clone(),
 87                prompt_buffer.clone(),
 88                codegen,
 89                session_id,
 90                self.fs.clone(),
 91                thread_store.clone(),
 92                prompt_store.clone(),
 93                history,
 94                project.clone(),
 95                workspace.clone(),
 96                window,
 97                cx,
 98            )
 99        });
100        let prompt_editor_render = prompt_editor.clone();
101        let block = terminal_view::BlockProperties {
102            height: 4,
103            render: Box::new(move |_| prompt_editor_render.clone().into_any_element()),
104        };
105        terminal_view.update(cx, |terminal_view, cx| {
106            terminal_view.set_block_below_cursor(block, window, cx);
107        });
108
109        let terminal_assistant = TerminalInlineAssist::new(
110            assist_id,
111            terminal_view,
112            prompt_editor,
113            workspace.clone(),
114            window,
115            cx,
116        );
117
118        self.assists.insert(assist_id, terminal_assistant);
119
120        self.focus_assist(assist_id, window, cx);
121    }
122
123    fn focus_assist(
124        &mut self,
125        assist_id: TerminalInlineAssistId,
126        window: &mut Window,
127        cx: &mut App,
128    ) {
129        let assist = &self.assists[&assist_id];
130        if let Some(prompt_editor) = assist.prompt_editor.as_ref() {
131            prompt_editor.update(cx, |this, cx| {
132                this.editor.update(cx, |editor, cx| {
133                    window.focus(&editor.focus_handle(cx), cx);
134                    editor.select_all(&SelectAll, window, cx);
135                });
136            });
137        }
138    }
139
140    fn handle_prompt_editor_event(
141        &mut self,
142        prompt_editor: Entity<PromptEditor<TerminalCodegen>>,
143        event: &PromptEditorEvent,
144        window: &mut Window,
145        cx: &mut App,
146    ) {
147        let assist_id = prompt_editor.read(cx).id();
148        match event {
149            PromptEditorEvent::StartRequested => {
150                self.start_assist(assist_id, cx);
151            }
152            PromptEditorEvent::StopRequested => {
153                self.stop_assist(assist_id, cx);
154            }
155            PromptEditorEvent::ConfirmRequested { execute } => {
156                self.finish_assist(assist_id, false, *execute, window, cx);
157            }
158            PromptEditorEvent::CancelRequested => {
159                self.finish_assist(assist_id, true, false, window, cx);
160            }
161            PromptEditorEvent::Resized { height_in_lines } => {
162                self.insert_prompt_editor_into_terminal(assist_id, *height_in_lines, window, cx);
163            }
164        }
165    }
166
167    fn start_assist(&mut self, assist_id: TerminalInlineAssistId, cx: &mut App) {
168        let assist = if let Some(assist) = self.assists.get_mut(&assist_id) {
169            assist
170        } else {
171            return;
172        };
173
174        let Some(user_prompt) = assist
175            .prompt_editor
176            .as_ref()
177            .map(|editor| editor.read(cx).prompt(cx))
178        else {
179            return;
180        };
181
182        self.prompt_history.retain(|prompt| *prompt != user_prompt);
183        self.prompt_history.push_back(user_prompt);
184        if self.prompt_history.len() > PROMPT_HISTORY_MAX_LEN {
185            self.prompt_history.pop_front();
186        }
187
188        assist
189            .terminal
190            .update(cx, |terminal, cx| {
191                terminal
192                    .terminal()
193                    .update(cx, |terminal, _| terminal.input(CLEAR_INPUT.as_bytes()));
194            })
195            .log_err();
196
197        let codegen = assist.codegen.clone();
198        let Some(request_task) = self.request_for_inline_assist(assist_id, cx).log_err() else {
199            return;
200        };
201
202        codegen.update(cx, |codegen, cx| codegen.start(request_task, cx));
203    }
204
205    fn stop_assist(&mut self, assist_id: TerminalInlineAssistId, cx: &mut App) {
206        let assist = if let Some(assist) = self.assists.get_mut(&assist_id) {
207            assist
208        } else {
209            return;
210        };
211
212        assist.codegen.update(cx, |codegen, cx| codegen.stop(cx));
213    }
214
215    fn request_for_inline_assist(
216        &self,
217        assist_id: TerminalInlineAssistId,
218        cx: &mut App,
219    ) -> Result<Task<LanguageModelRequest>> {
220        let ConfiguredModel { model, .. } = LanguageModelRegistry::read_global(cx)
221            .inline_assistant_model()
222            .context("No inline assistant model")?;
223
224        let assist = self.assists.get(&assist_id).context("invalid assist")?;
225
226        let shell = std::env::var("SHELL").ok();
227        let (latest_output, working_directory) = assist
228            .terminal
229            .update(cx, |terminal, cx| {
230                let terminal = terminal.entity().read(cx);
231                let latest_output = terminal.last_n_non_empty_lines(DEFAULT_CONTEXT_LINES);
232                let working_directory = terminal
233                    .working_directory()
234                    .map(|path| path.to_string_lossy().into_owned());
235                (latest_output, working_directory)
236            })
237            .ok()
238            .unwrap_or_default();
239
240        let prompt_editor = assist.prompt_editor.clone().context("invalid assist")?;
241
242        let prompt = self.prompt_builder.generate_terminal_assistant_prompt(
243            &prompt_editor.read(cx).prompt(cx),
244            shell.as_deref(),
245            working_directory.as_deref(),
246            &latest_output,
247        )?;
248
249        let temperature = AgentSettings::temperature_for_model(&model, cx);
250
251        let mention_set = prompt_editor.read(cx).mention_set().clone();
252        let load_context_task = load_context(&mention_set, cx);
253
254        Ok(cx.background_spawn(async move {
255            let mut request_message = LanguageModelRequestMessage {
256                role: Role::User,
257                content: vec![],
258                cache: false,
259                reasoning_details: None,
260            };
261
262            if let Some(context) = load_context_task.await {
263                context.add_to_request_message(&mut request_message);
264            }
265
266            request_message.content.push(prompt.into());
267
268            LanguageModelRequest {
269                thread_id: None,
270                prompt_id: None,
271                intent: Some(CompletionIntent::TerminalInlineAssist),
272                messages: vec![request_message],
273                tools: Vec::new(),
274                tool_choice: None,
275                stop: Vec::new(),
276                temperature,
277                thinking_allowed: false,
278            }
279        }))
280    }
281
282    fn finish_assist(
283        &mut self,
284        assist_id: TerminalInlineAssistId,
285        undo: bool,
286        execute: bool,
287        window: &mut Window,
288        cx: &mut App,
289    ) {
290        self.dismiss_assist(assist_id, window, cx);
291
292        if let Some(assist) = self.assists.remove(&assist_id) {
293            assist
294                .terminal
295                .update(cx, |this, cx| {
296                    this.clear_block_below_cursor(cx);
297                    this.focus_handle(cx).focus(window, cx);
298                })
299                .log_err();
300
301            if let Some(ConfiguredModel { model, .. }) =
302                LanguageModelRegistry::read_global(cx).inline_assistant_model()
303            {
304                let codegen = assist.codegen.read(cx);
305                let session_id = codegen.session_id();
306                let message_id = codegen.message_id.clone();
307                let model_telemetry_id = model.telemetry_id();
308                let model_provider_id = model.provider_id().to_string();
309
310                let (phase, event_type, anthropic_event_type) = if undo {
311                    (
312                        "rejected",
313                        "Assistant Response Rejected",
314                        language_model::AnthropicEventType::Reject,
315                    )
316                } else {
317                    (
318                        "accepted",
319                        "Assistant Response Accepted",
320                        language_model::AnthropicEventType::Accept,
321                    )
322                };
323
324                // Fire Zed telemetry
325                telemetry::event!(
326                    event_type,
327                    kind = "inline_terminal",
328                    phase = phase,
329                    model = model_telemetry_id,
330                    model_provider = model_provider_id,
331                    message_id = message_id,
332                    session_id = session_id,
333                );
334
335                report_anthropic_event(
336                    &model,
337                    language_model::AnthropicEventData {
338                        completion_type: language_model::AnthropicCompletionType::Terminal,
339                        event: anthropic_event_type,
340                        language_name: None,
341                        message_id,
342                    },
343                    cx,
344                );
345            }
346
347            assist.codegen.update(cx, |codegen, cx| {
348                if undo {
349                    codegen.undo(cx);
350                } else if execute {
351                    codegen.complete(cx);
352                }
353            });
354        }
355    }
356
357    fn dismiss_assist(
358        &mut self,
359        assist_id: TerminalInlineAssistId,
360        window: &mut Window,
361        cx: &mut App,
362    ) -> bool {
363        let Some(assist) = self.assists.get_mut(&assist_id) else {
364            return false;
365        };
366        if assist.prompt_editor.is_none() {
367            return false;
368        }
369        assist.prompt_editor = None;
370        assist
371            .terminal
372            .update(cx, |this, cx| {
373                this.clear_block_below_cursor(cx);
374                this.focus_handle(cx).focus(window, cx);
375            })
376            .is_ok()
377    }
378
379    fn insert_prompt_editor_into_terminal(
380        &mut self,
381        assist_id: TerminalInlineAssistId,
382        height: u8,
383        window: &mut Window,
384        cx: &mut App,
385    ) {
386        if let Some(assist) = self.assists.get_mut(&assist_id)
387            && let Some(prompt_editor) = assist.prompt_editor.as_ref().cloned()
388        {
389            assist
390                .terminal
391                .update(cx, |terminal, cx| {
392                    terminal.clear_block_below_cursor(cx);
393                    let block = terminal_view::BlockProperties {
394                        height,
395                        render: Box::new(move |_| prompt_editor.clone().into_any_element()),
396                    };
397                    terminal.set_block_below_cursor(block, window, cx);
398                })
399                .log_err();
400        }
401    }
402}
403
404struct TerminalInlineAssist {
405    terminal: WeakEntity<TerminalView>,
406    prompt_editor: Option<Entity<PromptEditor<TerminalCodegen>>>,
407    codegen: Entity<TerminalCodegen>,
408    workspace: WeakEntity<Workspace>,
409    _subscriptions: Vec<Subscription>,
410}
411
412impl TerminalInlineAssist {
413    pub fn new(
414        assist_id: TerminalInlineAssistId,
415        terminal: &Entity<TerminalView>,
416        prompt_editor: Entity<PromptEditor<TerminalCodegen>>,
417        workspace: WeakEntity<Workspace>,
418        window: &mut Window,
419        cx: &mut App,
420    ) -> Self {
421        let codegen = prompt_editor.read(cx).codegen().clone();
422        Self {
423            terminal: terminal.downgrade(),
424            prompt_editor: Some(prompt_editor.clone()),
425            codegen: codegen.clone(),
426            workspace,
427            _subscriptions: vec![
428                window.subscribe(&prompt_editor, cx, |prompt_editor, event, window, cx| {
429                    TerminalInlineAssistant::update_global(cx, |this, cx| {
430                        this.handle_prompt_editor_event(prompt_editor, event, window, cx)
431                    })
432                }),
433                window.subscribe(&codegen, cx, move |codegen, event, window, cx| {
434                    TerminalInlineAssistant::update_global(cx, |this, cx| match event {
435                        CodegenEvent::Finished => {
436                            let assist = if let Some(assist) = this.assists.get(&assist_id) {
437                                assist
438                            } else {
439                                return;
440                            };
441
442                            if let CodegenStatus::Error(error) = &codegen.read(cx).status
443                                && assist.prompt_editor.is_none()
444                                && let Some(workspace) = assist.workspace.upgrade()
445                            {
446                                let error = format!("Terminal inline assistant error: {}", error);
447                                workspace.update(cx, |workspace, cx| {
448                                    struct InlineAssistantError;
449
450                                    let id = NotificationId::composite::<InlineAssistantError>(
451                                        assist_id.0,
452                                    );
453
454                                    workspace.show_toast(Toast::new(id, error), cx);
455                                })
456                            }
457
458                            if assist.prompt_editor.is_none() {
459                                this.finish_assist(assist_id, false, false, window, cx);
460                            }
461                        }
462                    })
463                }),
464            ],
465        }
466    }
467}