1use crate::{AssistantPanel, Context, InlineAssistId, InlineAssistant};
2use anyhow::{anyhow, Error, Result};
3use collections::HashMap;
4use editor::Editor;
5use futures::future;
6use gpui::{Model, ModelContext, Task, UpdateGlobal as _, View, WeakView, WindowContext};
7use language::{Anchor, Buffer, BufferSnapshot};
8use language_model::{LanguageModelRegistry, LanguageModelRequestMessage, Role};
9use project::Project;
10use rope::Point;
11use serde::{Deserialize, Serialize};
12use smol::stream::StreamExt;
13use std::{cmp, ops::Range, sync::Arc};
14use text::{AnchorRangeExt as _, OffsetRangeExt as _};
15use util::ResultExt as _;
16use workspace::Workspace;
17
18pub struct WorkflowStepResolution {
19 tagged_range: Range<Anchor>,
20 output: String,
21 pub result: Option<Result<ResolvedWorkflowStep, Arc<Error>>>,
22}
23
24#[derive(Clone, Debug, Eq, PartialEq)]
25pub struct ResolvedWorkflowStep {
26 pub title: String,
27 pub suggestions: HashMap<Model<Buffer>, Vec<WorkflowSuggestionGroup>>,
28}
29
30#[derive(Clone, Debug, Eq, PartialEq)]
31pub struct WorkflowSuggestionGroup {
32 pub context_range: Range<language::Anchor>,
33 pub suggestions: Vec<WorkflowSuggestion>,
34}
35
36#[derive(Clone, Debug, Eq, PartialEq)]
37pub enum WorkflowSuggestion {
38 Update {
39 range: Range<language::Anchor>,
40 description: String,
41 },
42 CreateFile {
43 description: String,
44 },
45 InsertSiblingBefore {
46 position: language::Anchor,
47 description: String,
48 },
49 InsertSiblingAfter {
50 position: language::Anchor,
51 description: String,
52 },
53 PrependChild {
54 position: language::Anchor,
55 description: String,
56 },
57 AppendChild {
58 position: language::Anchor,
59 description: String,
60 },
61 Delete {
62 range: Range<language::Anchor>,
63 },
64}
65
66impl WorkflowStepResolution {
67 pub fn new(range: Range<Anchor>) -> Self {
68 Self {
69 tagged_range: range,
70 output: String::new(),
71 result: None,
72 }
73 }
74
75 pub fn resolve(
76 &mut self,
77 context: &Context,
78 cx: &mut ModelContext<WorkflowStepResolution>,
79 ) -> Option<Task<()>> {
80 let project = context.project()?;
81 let context_buffer = context.buffer().clone();
82 let prompt_builder = context.prompt_builder();
83 let mut request = context.to_completion_request(cx);
84 let model = LanguageModelRegistry::read_global(cx).active_model();
85 let step_text = context_buffer
86 .read(cx)
87 .text_for_range(self.tagged_range.clone())
88 .collect::<String>();
89
90 Some(cx.spawn(|this, mut cx| async move {
91 let result = async {
92 let Some(model) = model else {
93 return Err(anyhow!("no model selected"));
94 };
95
96 this.update(&mut cx, |this, cx| {
97 this.output.clear();
98 this.result = None;
99 cx.notify();
100 })?;
101
102 let mut prompt = prompt_builder.generate_step_resolution_prompt()?;
103 prompt.push_str(&step_text);
104 request.messages.push(LanguageModelRequestMessage {
105 role: Role::User,
106 content: vec![prompt.into()],
107 });
108
109 // Invoke the model to get its edit suggestions for this workflow step.
110 let mut stream = model
111 .use_tool_stream::<tool::WorkflowStepResolution>(request, &cx)
112 .await?;
113 while let Some(chunk) = stream.next().await {
114 let chunk = chunk?;
115 this.update(&mut cx, |this, cx| {
116 this.output.push_str(&chunk);
117 cx.notify();
118 })?;
119 }
120
121 let resolution = this.update(&mut cx, |this, _| {
122 serde_json::from_str::<tool::WorkflowStepResolution>(&this.output)
123 })??;
124
125 // Translate the parsed suggestions to our internal types, which anchor the suggestions to locations in the code.
126 let suggestion_tasks: Vec<_> = resolution
127 .suggestions
128 .iter()
129 .map(|suggestion| suggestion.resolve(project.clone(), cx.clone()))
130 .collect();
131
132 // Expand the context ranges of each suggestion and group suggestions with overlapping context ranges.
133 let suggestions = future::join_all(suggestion_tasks)
134 .await
135 .into_iter()
136 .filter_map(|task| task.log_err())
137 .collect::<Vec<_>>();
138
139 let mut suggestions_by_buffer = HashMap::default();
140 for (buffer, suggestion) in suggestions {
141 suggestions_by_buffer
142 .entry(buffer)
143 .or_insert_with(Vec::new)
144 .push(suggestion);
145 }
146
147 let mut suggestion_groups_by_buffer = HashMap::default();
148 for (buffer, mut suggestions) in suggestions_by_buffer {
149 let mut suggestion_groups = Vec::<WorkflowSuggestionGroup>::new();
150 let snapshot = buffer.update(&mut cx, |buffer, _| buffer.snapshot())?;
151 // Sort suggestions by their range so that earlier, larger ranges come first
152 suggestions.sort_by(|a, b| a.range().cmp(&b.range(), &snapshot));
153
154 // Merge overlapping suggestions
155 suggestions.dedup_by(|a, b| b.try_merge(&a, &snapshot));
156
157 // Create context ranges for each suggestion
158 for suggestion in suggestions {
159 let context_range = {
160 let suggestion_point_range = suggestion.range().to_point(&snapshot);
161 let start_row = suggestion_point_range.start.row.saturating_sub(5);
162 let end_row = cmp::min(
163 suggestion_point_range.end.row + 5,
164 snapshot.max_point().row,
165 );
166 let start = snapshot.anchor_before(Point::new(start_row, 0));
167 let end = snapshot
168 .anchor_after(Point::new(end_row, snapshot.line_len(end_row)));
169 start..end
170 };
171
172 if let Some(last_group) = suggestion_groups.last_mut() {
173 if last_group
174 .context_range
175 .end
176 .cmp(&context_range.start, &snapshot)
177 .is_ge()
178 {
179 // Merge with the previous group if context ranges overlap
180 last_group.context_range.end = context_range.end;
181 last_group.suggestions.push(suggestion);
182 } else {
183 // Create a new group
184 suggestion_groups.push(WorkflowSuggestionGroup {
185 context_range,
186 suggestions: vec![suggestion],
187 });
188 }
189 } else {
190 // Create the first group
191 suggestion_groups.push(WorkflowSuggestionGroup {
192 context_range,
193 suggestions: vec![suggestion],
194 });
195 }
196 }
197
198 suggestion_groups_by_buffer.insert(buffer, suggestion_groups);
199 }
200
201 Ok((resolution.step_title, suggestion_groups_by_buffer))
202 };
203
204 let result = result.await;
205 this.update(&mut cx, |this, cx| {
206 this.result = Some(match result {
207 Ok((title, suggestions)) => Ok(ResolvedWorkflowStep { title, suggestions }),
208 Err(error) => Err(Arc::new(error)),
209 });
210 cx.notify();
211 })
212 .ok();
213 }))
214 }
215}
216
217impl WorkflowSuggestion {
218 pub fn range(&self) -> Range<language::Anchor> {
219 match self {
220 WorkflowSuggestion::Update { range, .. } => range.clone(),
221 WorkflowSuggestion::CreateFile { .. } => language::Anchor::MIN..language::Anchor::MAX,
222 WorkflowSuggestion::InsertSiblingBefore { position, .. }
223 | WorkflowSuggestion::InsertSiblingAfter { position, .. }
224 | WorkflowSuggestion::PrependChild { position, .. }
225 | WorkflowSuggestion::AppendChild { position, .. } => *position..*position,
226 WorkflowSuggestion::Delete { range } => range.clone(),
227 }
228 }
229
230 pub fn description(&self) -> Option<&str> {
231 match self {
232 WorkflowSuggestion::Update { description, .. }
233 | WorkflowSuggestion::CreateFile { description }
234 | WorkflowSuggestion::InsertSiblingBefore { description, .. }
235 | WorkflowSuggestion::InsertSiblingAfter { description, .. }
236 | WorkflowSuggestion::PrependChild { description, .. }
237 | WorkflowSuggestion::AppendChild { description, .. } => Some(description),
238 WorkflowSuggestion::Delete { .. } => None,
239 }
240 }
241
242 fn description_mut(&mut self) -> Option<&mut String> {
243 match self {
244 WorkflowSuggestion::Update { description, .. }
245 | WorkflowSuggestion::CreateFile { description }
246 | WorkflowSuggestion::InsertSiblingBefore { description, .. }
247 | WorkflowSuggestion::InsertSiblingAfter { description, .. }
248 | WorkflowSuggestion::PrependChild { description, .. }
249 | WorkflowSuggestion::AppendChild { description, .. } => Some(description),
250 WorkflowSuggestion::Delete { .. } => None,
251 }
252 }
253
254 fn try_merge(&mut self, other: &Self, buffer: &BufferSnapshot) -> bool {
255 let range = self.range();
256 let other_range = other.range();
257
258 // Don't merge if we don't contain the other suggestion.
259 if range.start.cmp(&other_range.start, buffer).is_gt()
260 || range.end.cmp(&other_range.end, buffer).is_lt()
261 {
262 return false;
263 }
264
265 if let Some(description) = self.description_mut() {
266 if let Some(other_description) = other.description() {
267 description.push('\n');
268 description.push_str(other_description);
269 }
270 }
271 true
272 }
273
274 pub fn show(
275 &self,
276 editor: &View<Editor>,
277 excerpt_id: editor::ExcerptId,
278 workspace: &WeakView<Workspace>,
279 assistant_panel: &View<AssistantPanel>,
280 cx: &mut WindowContext,
281 ) -> Option<InlineAssistId> {
282 let mut initial_transaction_id = None;
283 let initial_prompt;
284 let suggestion_range;
285 let buffer = editor.read(cx).buffer().clone();
286 let snapshot = buffer.read(cx).snapshot(cx);
287
288 match self {
289 WorkflowSuggestion::Update { range, description } => {
290 initial_prompt = description.clone();
291 suggestion_range = snapshot.anchor_in_excerpt(excerpt_id, range.start)?
292 ..snapshot.anchor_in_excerpt(excerpt_id, range.end)?;
293 }
294 WorkflowSuggestion::CreateFile { description } => {
295 initial_prompt = description.clone();
296 suggestion_range = editor::Anchor::min()..editor::Anchor::min();
297 }
298 WorkflowSuggestion::InsertSiblingBefore {
299 position,
300 description,
301 } => {
302 let position = snapshot.anchor_in_excerpt(excerpt_id, *position)?;
303 initial_prompt = description.clone();
304 suggestion_range = buffer.update(cx, |buffer, cx| {
305 buffer.start_transaction(cx);
306 let line_start = buffer.insert_empty_line(position, true, true, cx);
307 initial_transaction_id = buffer.end_transaction(cx);
308 buffer.refresh_preview(cx);
309
310 let line_start = buffer.read(cx).anchor_before(line_start);
311 line_start..line_start
312 });
313 }
314 WorkflowSuggestion::InsertSiblingAfter {
315 position,
316 description,
317 } => {
318 let position = snapshot.anchor_in_excerpt(excerpt_id, *position)?;
319 initial_prompt = description.clone();
320 suggestion_range = buffer.update(cx, |buffer, cx| {
321 buffer.start_transaction(cx);
322 let line_start = buffer.insert_empty_line(position, true, true, cx);
323 initial_transaction_id = buffer.end_transaction(cx);
324 buffer.refresh_preview(cx);
325
326 let line_start = buffer.read(cx).anchor_before(line_start);
327 line_start..line_start
328 });
329 }
330 WorkflowSuggestion::PrependChild {
331 position,
332 description,
333 } => {
334 let position = snapshot.anchor_in_excerpt(excerpt_id, *position)?;
335 initial_prompt = description.clone();
336 suggestion_range = buffer.update(cx, |buffer, cx| {
337 buffer.start_transaction(cx);
338 let line_start = buffer.insert_empty_line(position, false, true, cx);
339 initial_transaction_id = buffer.end_transaction(cx);
340 buffer.refresh_preview(cx);
341
342 let line_start = buffer.read(cx).anchor_before(line_start);
343 line_start..line_start
344 });
345 }
346 WorkflowSuggestion::AppendChild {
347 position,
348 description,
349 } => {
350 let position = snapshot.anchor_in_excerpt(excerpt_id, *position)?;
351 initial_prompt = description.clone();
352 suggestion_range = buffer.update(cx, |buffer, cx| {
353 buffer.start_transaction(cx);
354 let line_start = buffer.insert_empty_line(position, true, false, cx);
355 initial_transaction_id = buffer.end_transaction(cx);
356 buffer.refresh_preview(cx);
357
358 let line_start = buffer.read(cx).anchor_before(line_start);
359 line_start..line_start
360 });
361 }
362 WorkflowSuggestion::Delete { range } => {
363 initial_prompt = "Delete".to_string();
364 suggestion_range = snapshot.anchor_in_excerpt(excerpt_id, range.start)?
365 ..snapshot.anchor_in_excerpt(excerpt_id, range.end)?;
366 }
367 }
368
369 InlineAssistant::update_global(cx, |inline_assistant, cx| {
370 Some(inline_assistant.suggest_assist(
371 editor,
372 suggestion_range,
373 initial_prompt,
374 initial_transaction_id,
375 Some(workspace.clone()),
376 Some(assistant_panel),
377 cx,
378 ))
379 })
380 }
381}
382
383pub mod tool {
384 use std::path::Path;
385
386 use super::*;
387 use anyhow::Context as _;
388 use gpui::AsyncAppContext;
389 use language::ParseStatus;
390 use language_model::LanguageModelTool;
391 use project::ProjectPath;
392 use schemars::JsonSchema;
393
394 #[derive(Debug, Serialize, Deserialize, JsonSchema)]
395 pub struct WorkflowStepResolution {
396 /// An extremely short title for the edit step represented by these operations.
397 pub step_title: String,
398 /// A sequence of operations to apply to the codebase.
399 /// When multiple operations are required for a step, be sure to include multiple operations in this list.
400 pub suggestions: Vec<WorkflowSuggestion>,
401 }
402
403 impl LanguageModelTool for WorkflowStepResolution {
404 fn name() -> String {
405 "edit".into()
406 }
407
408 fn description() -> String {
409 "suggest edits to one or more locations in the codebase".into()
410 }
411 }
412
413 /// A description of an operation to apply to one location in the codebase.
414 ///
415 /// This object represents a single edit operation that can be performed on a specific file
416 /// in the codebase. It encapsulates both the location (file path) and the nature of the
417 /// edit to be made.
418 ///
419 /// # Fields
420 ///
421 /// * `path`: A string representing the file path where the edit operation should be applied.
422 /// This path is relative to the root of the project or repository.
423 ///
424 /// * `kind`: An enum representing the specific type of edit operation to be performed.
425 ///
426 /// # Usage
427 ///
428 /// `EditOperation` is used within a code editor to represent and apply
429 /// programmatic changes to source code. It provides a structured way to describe
430 /// edits for features like refactoring tools or AI-assisted coding suggestions.
431 #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
432 pub struct WorkflowSuggestion {
433 /// The path to the file containing the relevant operation
434 pub path: String,
435 #[serde(flatten)]
436 pub kind: WorkflowSuggestionKind,
437 }
438
439 impl WorkflowSuggestion {
440 pub(super) async fn resolve(
441 &self,
442 project: Model<Project>,
443 mut cx: AsyncAppContext,
444 ) -> Result<(Model<Buffer>, super::WorkflowSuggestion)> {
445 let path = self.path.clone();
446 let kind = self.kind.clone();
447 let buffer = project
448 .update(&mut cx, |project, cx| {
449 let project_path = project
450 .find_project_path(Path::new(&path), cx)
451 .or_else(|| {
452 // If we couldn't find a project path for it, put it in the active worktree
453 // so that when we create the buffer, it can be saved.
454 let worktree = project
455 .active_entry()
456 .and_then(|entry_id| project.worktree_for_entry(entry_id, cx))
457 .or_else(|| project.worktrees(cx).next())?;
458 let worktree = worktree.read(cx);
459
460 Some(ProjectPath {
461 worktree_id: worktree.id(),
462 path: Arc::from(Path::new(&path)),
463 })
464 })
465 .with_context(|| format!("worktree not found for {:?}", path))?;
466 anyhow::Ok(project.open_buffer(project_path, cx))
467 })??
468 .await?;
469
470 let mut parse_status = buffer.read_with(&cx, |buffer, _cx| buffer.parse_status())?;
471 while *parse_status.borrow() != ParseStatus::Idle {
472 parse_status.changed().await?;
473 }
474
475 let snapshot = buffer.update(&mut cx, |buffer, _| buffer.snapshot())?;
476 let outline = snapshot.outline(None).context("no outline for buffer")?;
477
478 let suggestion;
479 match kind {
480 WorkflowSuggestionKind::Update {
481 symbol,
482 description,
483 } => {
484 let symbol = outline
485 .find_most_similar(&symbol)
486 .with_context(|| format!("symbol not found: {:?}", symbol))?
487 .to_point(&snapshot);
488 let start = symbol
489 .annotation_range
490 .map_or(symbol.range.start, |range| range.start);
491 let start = Point::new(start.row, 0);
492 let end = Point::new(
493 symbol.range.end.row,
494 snapshot.line_len(symbol.range.end.row),
495 );
496 let range = snapshot.anchor_before(start)..snapshot.anchor_after(end);
497 suggestion = super::WorkflowSuggestion::Update { range, description };
498 }
499 WorkflowSuggestionKind::Create { description } => {
500 suggestion = super::WorkflowSuggestion::CreateFile { description };
501 }
502 WorkflowSuggestionKind::InsertSiblingBefore {
503 symbol,
504 description,
505 } => {
506 let symbol = outline
507 .find_most_similar(&symbol)
508 .with_context(|| format!("symbol not found: {:?}", symbol))?
509 .to_point(&snapshot);
510 let position = snapshot.anchor_before(
511 symbol
512 .annotation_range
513 .map_or(symbol.range.start, |annotation_range| {
514 annotation_range.start
515 }),
516 );
517 suggestion = super::WorkflowSuggestion::InsertSiblingBefore {
518 position,
519 description,
520 };
521 }
522 WorkflowSuggestionKind::InsertSiblingAfter {
523 symbol,
524 description,
525 } => {
526 let symbol = outline
527 .find_most_similar(&symbol)
528 .with_context(|| format!("symbol not found: {:?}", symbol))?
529 .to_point(&snapshot);
530 let position = snapshot.anchor_after(symbol.range.end);
531 suggestion = super::WorkflowSuggestion::InsertSiblingAfter {
532 position,
533 description,
534 };
535 }
536 WorkflowSuggestionKind::PrependChild {
537 symbol,
538 description,
539 } => {
540 if let Some(symbol) = symbol {
541 let symbol = outline
542 .find_most_similar(&symbol)
543 .with_context(|| format!("symbol not found: {:?}", symbol))?
544 .to_point(&snapshot);
545
546 let position = snapshot.anchor_after(
547 symbol
548 .body_range
549 .map_or(symbol.range.start, |body_range| body_range.start),
550 );
551 suggestion = super::WorkflowSuggestion::PrependChild {
552 position,
553 description,
554 };
555 } else {
556 suggestion = super::WorkflowSuggestion::PrependChild {
557 position: language::Anchor::MIN,
558 description,
559 };
560 }
561 }
562 WorkflowSuggestionKind::AppendChild {
563 symbol,
564 description,
565 } => {
566 if let Some(symbol) = symbol {
567 let symbol = outline
568 .find_most_similar(&symbol)
569 .with_context(|| format!("symbol not found: {:?}", symbol))?
570 .to_point(&snapshot);
571
572 let position = snapshot.anchor_before(
573 symbol
574 .body_range
575 .map_or(symbol.range.end, |body_range| body_range.end),
576 );
577 suggestion = super::WorkflowSuggestion::AppendChild {
578 position,
579 description,
580 };
581 } else {
582 suggestion = super::WorkflowSuggestion::PrependChild {
583 position: language::Anchor::MAX,
584 description,
585 };
586 }
587 }
588 WorkflowSuggestionKind::Delete { symbol } => {
589 let symbol = outline
590 .find_most_similar(&symbol)
591 .with_context(|| format!("symbol not found: {:?}", symbol))?
592 .to_point(&snapshot);
593 let start = symbol
594 .annotation_range
595 .map_or(symbol.range.start, |range| range.start);
596 let start = Point::new(start.row, 0);
597 let end = Point::new(
598 symbol.range.end.row,
599 snapshot.line_len(symbol.range.end.row),
600 );
601 let range = snapshot.anchor_before(start)..snapshot.anchor_after(end);
602 suggestion = super::WorkflowSuggestion::Delete { range };
603 }
604 }
605
606 Ok((buffer, suggestion))
607 }
608 }
609
610 #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
611 #[serde(tag = "kind")]
612 pub enum WorkflowSuggestionKind {
613 /// Rewrites the specified symbol entirely based on the given description.
614 /// This operation completely replaces the existing symbol with new content.
615 Update {
616 /// A fully-qualified reference to the symbol, e.g. `mod foo impl Bar pub fn baz` instead of just `fn baz`.
617 /// The path should uniquely identify the symbol within the containing file.
618 symbol: String,
619 /// A brief description of the transformation to apply to the symbol.
620 description: String,
621 },
622 /// Creates a new file with the given path based on the provided description.
623 /// This operation adds a new file to the codebase.
624 Create {
625 /// A brief description of the file to be created.
626 description: String,
627 },
628 /// Inserts a new symbol based on the given description before the specified symbol.
629 /// This operation adds new content immediately preceding an existing symbol.
630 InsertSiblingBefore {
631 /// A fully-qualified reference to the symbol, e.g. `mod foo impl Bar pub fn baz` instead of just `fn baz`.
632 /// The new content will be inserted immediately before this symbol.
633 symbol: String,
634 /// A brief description of the new symbol to be inserted.
635 description: String,
636 },
637 /// Inserts a new symbol based on the given description after the specified symbol.
638 /// This operation adds new content immediately following an existing symbol.
639 InsertSiblingAfter {
640 /// A fully-qualified reference to the symbol, e.g. `mod foo impl Bar pub fn baz` instead of just `fn baz`.
641 /// The new content will be inserted immediately after this symbol.
642 symbol: String,
643 /// A brief description of the new symbol to be inserted.
644 description: String,
645 },
646 /// Inserts a new symbol as a child of the specified symbol at the start.
647 /// This operation adds new content as the first child of an existing symbol (or file if no symbol is provided).
648 PrependChild {
649 /// An optional fully-qualified reference to the symbol after the code you want to insert, e.g. `mod foo impl Bar pub fn baz` instead of just `fn baz`.
650 /// If provided, the new content will be inserted as the first child of this symbol.
651 /// If not provided, the new content will be inserted at the top of the file.
652 symbol: Option<String>,
653 /// A brief description of the new symbol to be inserted.
654 description: String,
655 },
656 /// Inserts a new symbol as a child of the specified symbol at the end.
657 /// This operation adds new content as the last child of an existing symbol (or file if no symbol is provided).
658 AppendChild {
659 /// An optional fully-qualified reference to the symbol before the code you want to insert, e.g. `mod foo impl Bar pub fn baz` instead of just `fn baz`.
660 /// If provided, the new content will be inserted as the last child of this symbol.
661 /// If not provided, the new content will be applied at the bottom of the file.
662 symbol: Option<String>,
663 /// A brief description of the new symbol to be inserted.
664 description: String,
665 },
666 /// Deletes the specified symbol from the containing file.
667 Delete {
668 /// An fully-qualified reference to the symbol to be deleted, e.g. `mod foo impl Bar pub fn baz` instead of just `fn baz`.
669 symbol: String,
670 },
671 }
672}