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