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